FP8 Convolutions in XLA

[philipphack]

Enables scaled convolutions of the form

(X, W, x_scale, w_scale, y_scale) -> Y,

where the input X, the filter W and the output Y are based on the F8E4M3FN and F8E5M2 data types and x_scale, w_scale and y_scale are their scaling factors.

[philipphack]

CC @reedwm, @nluehr.

[reedwm]

Thanks for adding FP8 support to convolutions! Sorry for taking so long to review this.

Also the comments are in a weird order since I went back and forth between the files a lot when reviewing this. I would view the comments in the "Files changed" tab.

[reedwm]

I misunderstood what you were doing before, when I commented

Instead of having this pattern_level concept, can we just iterate over the users again and check for the convert+clamp?

I understand now that you're trying to match a convert of a clamp, which requires looking at the user of the user. I'm still not a fan of the pattern_level concept though, since it's confusing. The best way to address this would be to return a GraphString without the convert+clamp. Then create a FuseConvertToF8 function that is called from CudnnFusedConvRewriter::Run, where you find the existing instruction with a graph string and append the conversion to F8.

If you'd rather handle everything in this function, you can directly get the user of the user to see if that matches, instead of using recursion and pattern_level. E.g. you can do:

if (user->user_count() == 1 && Match(user->users()[0], m::Convert(...))) {...}

Granted, when we eventually match amax calculations, we might need to go deeper, e.g. checking the user of the user of the user. I think doing so directly is better than recursion.

[philipphack]

I think in the context of having multiple users, the recursive approach is advantageous and more straightforward. Can we revisit this issue after you've seen the Amax case?

[reedwm]

Sure, I'll bring this up again when you add amax support if I still think the non-recursive approach is better.

[reedwm]

Can we check both custom_call_string and serialized_graph_string on pre-Hopper, using only the two passes (GpuConvRewriter and CudnnFusedConvRewriter) instead of all the passes? In both the Hopper and non-Hopper case, we can call RunAndFilecheckHloRewrite, and on Hopper only, we can call RunAndCompare. In the Hopper case, you can also call RunFileCheck but only need to do a simple sanity check, such as that the graph string is correct, since other passes may modify things like layout which would make the custom_call_string not match.

This is similar to what we do in gemm_rewrite_test, where we only call RunAndCompare on Hopper

[philipphack]

Isn't that the solution we arrived at? As I understand it, you don't want to verify the final layout even on Hopper systems which in my opinion renders the test somewhat incomplete. I don't think this is directly comparable to the GEMM case where layout conversions play less of a role.

[reedwm]

I'm still unsure why layout is important here, compared to the gemm case. Running GpuConvRewriter and CudnnFusedConvRewriter on pre_hlo_string is still causing the FP8 rewrite to happen even if layout assignment doesn't run, right?

Even layout assignment is important, maybe see what transformations it is doing to pre_hlo_string, and just putting the resulting layouts directly in the HLO strings in the test, instead of relying on layout assignment to run.

[philipphack]

My perspective is that the XLA unit tests are usually based on running the compiler pipeline. When that's not possible, we can do at partial testing by running only the relevant pass in some artificial setting that we can't easily extend to the full pipeline. It's less clear to me though why we'd want to deviate from the normal approach and restrict the testing in cases where we don't have to as well.

[reedwm]

The reason to deviate from the normal approach is that we want to test as much as possible on non-Hopper, and right now, the PR only tests the graph string and not things like the custom_call_target, the dim_labels, etc.

But I'll accept only testing the graph string for now, we can reconsider if the tests get broken later due to a lack of Hopper CI.

[philipphack]

We can discard the first part of the custom_call FileCheck string on non-Hopper systems and still compare the configuration of the Custom Call.

[reedwm]

What is the first part of the custom_call FileCheck string? Is this the f8e4m3fn[1,6,6,16]{3,2,1,0}, u8[{{.*}}]{0}) part and is that part different if you don't run the rest of the passes?

[philipphack]

Yes, the order of the dimensions changes.

[reedwm]

State that copying a cudnn tensor requires cudnn 8.8 in the error.

Also in cudnn_fused_conv_rewriter.cc, you should check CUDNN_VERSION in addition to CUDA_VERSION to avoid this error.

[philipphack]

I think this might be supported but the functionality is only used when we require cuDNN version of at least 8.9 and I can't easily test it. One option would be to give the clone overload of CreateCudnnTensor its own 8.9 version guard instead of sharing it with the existing overload and remove this.

[reedwm]

I'm fine either keeping it as-is or adding an 8.9 version check inside thedefinition of CreateCudnnTensor.

Either way, check CUDNN_VERSION in cudnn_fused_conv_rewriter.cc though.

[reedwm]

Sure, I'll bring this up again when you add amax support if I still think the non-recursive approach is better.

[reedwm]

I'm still unsure why layout is important here, compared to the gemm case. Running GpuConvRewriter and CudnnFusedConvRewriter on pre_hlo_string is still causing the FP8 rewrite to happen even if layout assignment doesn't run, right?

Even layout assignment is important, maybe see what transformations it is doing to pre_hlo_string, and just putting the resulting layouts directly in the HLO strings in the test, instead of relying on layout assignment to run.

[reedwm]

I'm fine either keeping it as-is or adding an 8.9 version check inside thedefinition of CreateCudnnTensor.

Either way, check CUDNN_VERSION in cudnn_fused_conv_rewriter.cc though.

[reedwm]

Can you resolve conflicts?