Faster FP32 power calculation on CUDA GPUs

This is related to #8068 but concerns performance rather than accuracy.

@nouiz and I noticed that jax.numpy.power currently lowers to nearly 200 PTX instructions. This is similar to how nvcc would compile powerf without -use_fast_math.

However, the hardware-implemented SP exponential/logarithm functions can compute a**b as exp(b*log(a)) with much better performance and slightly worse accuracy. This corresponds to the __powf CUDA C++ intrinsic or powerf with -use_fast_math.

        lg2.approx.ftz.f32      %f3, %f1;
        mul.ftz.f32     %f4, %f2, %f3;
        ex2.approx.ftz.f32      %f5, %f4;

Is this something worths further investigation?

[hawkinsp]

Out of curiosity, what was the original workload that calls pow() heavily? Do you have a benchmark?

[hawkinsp]

I think it’s a good question. This optimization could certainly be done in XLA and seems worthwhile particularly for neural network users.

The key thing to my mind is we need a better API at both XLA and JAX levels for allowing the user to choose either the “fast” implementation or the “accurate” implementation. If we had that, we would have more flexibility to have both and choose between them. The closest analogue I can think of is the precision flag on Dot and Conv operations in HLO.