[FEA] Support for NVIDIA_TF32_OVERRIDE environment variable + handle

[ahendriksen]

Is your feature request related to a problem? Please describe.
I have recently run some brute force KNN benchmarks with @tfeher. Here, we looked at the impact of using 1 x tf32 versus 3 x tf32 performance of brute force knn. On a representative benchmark, using 1 x tf32 resulted in a 2.5x speedup (5 seconds -> 2 seconds). This can be significant for certain workloads (but can also not be set as the default due to unknown effects of reduced numerical accuracy).

We ran into the problem how to support this use case in our current pairwise distance API. We already have two distance types for the L2 distance (expanded and unexpanded). Adding variants for every possible way of speeding up the computation could become prohibitive. CuBLAS supports the NVIDIA_TF32_OVERRIDE environment variable that can force fp32 computations to be performed in tfloat32 precision.

Describe the solution you'd like
Add support for the NVIDIA_TF32_OVERRIDE environment in the RAFT handle. This way, algorithms can interrogate this option without having to continously inspect the environment.

In addition, make it possible to set the tf32 override programmatically. For instance, PyTorch supports the following:

# The flag below controls whether to allow TF32 on matmul. 
torch.backends.cuda.matmul.allow_tf32 = True

Describe alternatives you've considered
Adding another L2 distance type, which I think is unwise (and would not help in the case of cosine distance). Also, adding boolean flags to the pairwise distance API is going to be a mess.

[benfred]

One idea that @cjnolet and I have talked about, is expanding the 'metric_arg' to be a structure - and have it include information like whether to use an expanded/unexpanded distance, in addition to the floating point value for LP distance that it is currently. With this change we could combine the L2Expanded and L2Unexpanded distance metric types into a single DistanceType::L2 - with providing options on how to compute in the metric_arg argument (and potentially even add information about whether to take the final sqrt , meaning we could also get rid of L2SqrtExpanded/L2SqrtUnexpanded).

The rationale is that the DistanceType enum should only really be indicating what the distance is, and that things like expanded/unexpanded are just on how its computed.

If we made this change, I think that passing that flags like allow_tf32 would also belong in this MetricArgs struct - with potentially the default value being picked up from the NVIDIA_TF32_OVERRIDE environment variable.

[ahendriksen]

expanding the 'metric_arg' to be a structure

That is a really great idea! It is more powerful than interpreting the "global" handle state.

With this change we could combine the L2Expanded and L2Unexpanded distance metric types [.. snip ..] (and potentially even add information about whether to take the final sqrt , [.. snip ..]).

Yes, very nice. May I propose that we name the "expanded" form something like "fast" (to describe the goal rather than the implementation). I would also prefer we use a "squared" flag instead of a "sqrt" flag, as the square root is part of the norm (strictly speaking). Then the L2 norm would be called the L2 norm, and the squared L2 norm would be L2 norm with the squared option.

If we made this change, I think that passing that flags like allow_tf32 would also belong in this MetricArgs struct - with potentially the default value being picked up from the NVIDIA_TF32_OVERRIDE environment variable.

100% on board.

One question that I have though, if the metric_arg is going to be a value type is how we make sure that we can add to it in a backward compatible way. It's unfortunate if every change to the metric_arg type requires a breaking change. Do you have any thought on this?