Dispatch to use fp32 distance computation in NN Descent depending on data dimensions

[jinsolp]

Closes #1370
Closes #195

This PR adds an option to use fp32 distance computation.

(Outdated) From heuristics, chose dim=16 as the threshold for dispatching to a fp32 distance kernel.
We do manual computation, but since we only target small dimensions, fp32 dispatching ends up being slightly faster end to end with much better recall for small dimensions.

All number below are run on L40 machine and AMD EPYC CPU with 128 cores. Perf and recall is averaged over 5 runs and all time is in seconds. Baseline knn graph is computed using sklearn.neighbors.NearestNeighbors brute for method.

Max iters=20

For larger dimensions there is an inherent issue with the NN Descent algorithm itself that makes the recall low. This can be improved slightly with more iterations.
Also notice that the e2e time taken is similar or slightly less for using fp32.

Max iters=100

Notice how the blue part, the recall doesn't get better compared to the table above even with more iterations (i.e. why we need the fp32 appraoch for this part)

Perf impact on different architectures

H100

L40

[jinsolp]

Maybe we might need to have fp16 vs fp32 as an option (and default to fp16) instead of depending on data dimensions. WIP investigating

[jinsolp]

This is taken out as a separate function from the original kernel because it's commonly used for fp32 and fp16

[jinsolp]

This matmul part is different from the fp16 kernel

[divyegala]

I would like to see a unit test, if possible, where you can demonstrably prove distance computation improvements in fp32 over fp16.

[jinsolp]

Thanks for the feedback @divyegala, changing this to target 26.02 for now because Corey suggested we do further investigation. Force-pushing after rebasing

[jinsolp]

dispatch logic based on dimensions

[jinsolp]

Blue is better (e2e time or kernel time)

• FP32 kernel is usually faster for small dimensions
• FP32 e2e time is sometimes slightly longer for small dimensions despite shorter kernel time. Believe this small gap comes from more data transfers (4bytes vs 2bytes per data element)

L40 GPU

H100 GPU

[jinsolp]

@divyegala based on the latest benchmarks (and confirming from HDBSCAN that we don't have quality issues from using fp16 for larger dimensions), the mechanism is back to dispatching based on the number of dimensions
(chose 16 as a threshold).

I would like to see a unit test, if possible, where you can demonstrably prove distance computation improvements in fp32 over fp16.

Not sure how to add this as a test. Any suggestions?

[divyegala]

Not sure how to add this as a test. Any suggestions?

@jinsolp is it possible to generate a dataset that shows degraded recall for fp16 and good recall for fp32? The way I envision the assertion is ASSERT(fp16_recall < <upperbound_value> && fp32_recall >= <lowerbound_value>).

[jinsolp]

We can (e.g. the blobs data with dim=2 which shows 0.4 recall for fp16 vs 0.99 recall for fp32)
, but if we run through the current code, it will always take the fp32 path.

Should I hardwire the results of fp16 and add that as a vector to the test?

[divyegala]

@jinsolp we should allow the user an override in-case the heuristic fails. Let users choose fp32 explicitly if they want, and exporting that option will also let you test the separate paths.

[jinsolp]

Oh okay, so let the user choose an option, and if none is given fall back to the heuristic dim=16 threshold?

[divyegala]

The other way around. Use the heuristic eagerly, but let user override.

[jinsolp]

@divyegala we now have a dist_comp_dtype option defaulting to auto. can override by using fp32 or fp16.
Added a test as well.

[divyegala]

Please default initialize the index creation with this new parameter.

[jinsolp]

@divyegala It already defaults to AUTO here, which would be taken in nn_descent.cpp

[divyegala]

That file is the source of C API, and this default is for the C++ API.

[jinsolp]

I believe all the defaults of cpp are passed when we call cuvsNNDescentIndexParamsCreate

cuvs/c/src/neighbors/nn_descent.cpp

Lines 167 to 182 in 1959a0d

	extern "C" cuvsError_t cuvsNNDescentIndexParamsCreate(cuvsNNDescentIndexParams_t* params)
	{
	return cuvs::core::translate_exceptions([=] {
	// get defaults from cpp parameters struct
	cuvs::neighbors::nn_descent::index_params cpp_params;
	*params = new cuvsNNDescentIndexParams{
	.metric = static_cast<cuvsDistanceType>((int)cpp_params.metric),
	.metric_arg = cpp_params.metric_arg,
	.graph_degree = cpp_params.graph_degree,
	.intermediate_graph_degree = cpp_params.intermediate_graph_degree,
	.max_iterations = cpp_params.max_iterations,
	.termination_threshold = cpp_params.termination_threshold,
	.return_distances = cpp_params.return_distances};
	});
	}

C doesn't support default struct init and so nothing is initialized in the c struct:

cuvs/c/include/cuvs/neighbors/nn_descent.h

Lines 38 to 46 in 1959a0d

	struct cuvsNNDescentIndexParams {
	cuvsDistanceType metric;
	float metric_arg;
	size_t graph_degree;
	size_t intermediate_graph_degree;
	size_t max_iterations;
	float termination_threshold;
	bool return_distances;
	};

[divyegala]

I haven't kept up with the changes to the C API, my apologies. We used to explicitly default initialize in the create function previously.

[jinsolp]

/merge