Skip to content

Fix version number in Bazel downgrade warning - #24114

Closed
byronyi wants to merge 1 commit into
tensorflow:masterfrom
byronyi:patch-1
Closed

byronyi wants to merge 1 commit into
tensorflow:masterfrom
byronyi:patch-1

Conversation

@byronyi

@byronyi byronyi commented Dec 3, 2018 •

Copy link
Copy Markdown
Contributor

I have upgraded to Bazel 0.20:

$ bazel version
INFO: Invocation ID: 00114105-2567-451d-b22d-bb21e2d14b11
Build label: 0.20.0
Build target: bazel-out/k8-opt/bin/src/main/java/com/google/devtools/build/lib/bazel/BazelServer_deploy.jar
Build time: Fri Nov 30 14:39:01 2018 (1543588741)
Build timestamp: 1543588741
Build timestamp as int: 1543588741

Bazel version check is introduced in e7a123f.
Before this patch:

$ ./configure
WARNING: Running Bazel server needs to be killed, because the startup options are different.
WARNING: --batch mode is deprecated. Please instead explicitly shut down your Bazel server using the command "bazel shutdown".
INFO: Invocation ID: b6a7b9a2-db26-4777-a4bb-f5a668e52e82
You have bazel 0.20.0 installed.
Please downgrade your bazel installation to version 0.15.0 or lower to build TensorFlow!
Configuration finished

After this patch:

$ ./configure
WARNING: --batch mode is deprecated. Please instead explicitly shut down your Bazel server using the command "bazel shutdown".
INFO: Invocation ID: 34b75ed8-1359-4fa7-ba71-b67fd887ee3a
You have bazel 0.20.0 installed.
Please downgrade your bazel installation to version 0.19.2 or lower to build TensorFlow!
Configuration finished

@Harshini-Gadige Harshini-Gadige self-assigned this Dec 3, 2018
@Harshini-Gadige Harshini-Gadige added the awaiting review Pull request awaiting review label Dec 3, 2018

@mihaimaruseac mihaimaruseac left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you. Submitted an internal CL for this too early today

@mihaimaruseac mihaimaruseac removed the awaiting review Pull request awaiting review label Dec 3, 2018
@Harshini-Gadige Harshini-Gadige added the kokoro:force-run Tests on submitted change label Dec 3, 2018
@kokoro-team kokoro-team removed the kokoro:force-run Tests on submitted change label Dec 3, 2018
@Harshini-Gadige Harshini-Gadige added ready to pull PR ready for merge process and removed ready to pull PR ready for merge process labels Dec 4, 2018
@byronyi

byronyi commented Dec 7, 2018

Copy link
Copy Markdown
Contributor Author

@Harshini-Gadige Mind to pull?

@Harshini-Gadige Harshini-Gadige added ready to pull PR ready for merge process and removed ready to pull PR ready for merge process labels Dec 7, 2018
@Harshini-Gadige

Copy link
Copy Markdown

@Harshini-Gadige Mind to pull?

Sure. Will keep you posted.

@byronyi

byronyi commented Dec 10, 2018

Copy link
Copy Markdown
Contributor Author

Seems fixed by 3360c72. Closing for now.

@byronyi byronyi closed this Dec 10, 2018
copybara-service Bot pushed a commit that referenced this pull request Mar 27, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

Merging this change closes #24114

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#24114 from kasper0406:kn/fp8-conversion-fix 07ae307879eff24ad2f85607e94503deda1074e4
PiperOrigin-RevId: 741162069
copybara-service Bot pushed a commit that referenced this pull request Mar 27, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

Merging this change closes #24114

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#24114 from kasper0406:kn/fp8-conversion-fix 07ae307879eff24ad2f85607e94503deda1074e4
PiperOrigin-RevId: 741162069
copybara-service Bot pushed a commit that referenced this pull request Mar 27, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

Merging this change closes #24114

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#24114 from kasper0406:kn/fp8-conversion-fix 07ae307879eff24ad2f85607e94503deda1074e4
PiperOrigin-RevId: 741162069
copybara-service Bot pushed a commit that referenced this pull request Mar 27, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

Merging this change closes #24114

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#24114 from kasper0406:kn/fp8-conversion-fix 07ae307879eff24ad2f85607e94503deda1074e4
PiperOrigin-RevId: 741162069
copybara-service Bot pushed a commit that referenced this pull request Mar 27, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

Merging this change closes #24114

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#24114 from kasper0406:kn/fp8-conversion-fix 07ae307879eff24ad2f85607e94503deda1074e4
PiperOrigin-RevId: 741162069
copybara-service Bot pushed a commit that referenced this pull request Mar 28, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

--
fe967ff94ffc5f34f07bff142b5d10d81d5e4dce by Kasper Nielsen <[email protected]>:

Fix support conversion tests

Merging this change closes #24114

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#24114 from kasper0406:kn/fp8-conversion-fix fe967ff94ffc5f34f07bff142b5d10d81d5e4dce
PiperOrigin-RevId: 741162069
copybara-service Bot pushed a commit that referenced this pull request Mar 28, 2025
Imported from GitHub PR openxla/xla#24114

Converting FP8 <-> FP8 fails because the Triton compiler does not support it.
The proposed fix will make the conversion go through FP16.

Two questions:
1) Are there any better approaches of solving this?
2) I could not find a place to put unit tests for this, and in the code there is a comment saying:
    ```
        // TODO(b/266862493): Add end-to-end test once FP8 support lands in XLA as
        // we can't test the code below without patching the feature.
    ```
    Wondering if there is a place where I can add a test?

### Details
When converting FP8 types, the XLA compiler emits a `fp_to_fp` Triton instruction. If the source type is FP8, no rounding strategy is specified.

Concretely, this causes the following Triton to be emitted:
<details>
<summary>
<code>%24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN></code>
</summary>

```
module {
  tt.func @gemm_fusion_dot_320_impl(%arg0: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f8E5M2> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f8E4M3FN> {tt.divisibility = 16 : i32}) {
    %cst = arith.constant dense<0.000000e+00> : tensor<64x64xf8E4M3FN>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<32x64xf8E4M3FN>
    %c90_i32 = arith.constant 90 : i32
    %c32000_i64 = arith.constant 32000 : i64
    %c64_i32 = arith.constant 64 : i32
    %c90_i64 = arith.constant 90 : i64
    %c768_i64 = arith.constant 768 : i64
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %c32_i32 = arith.constant 32 : i32
    %c24_i32 = arith.constant 24 : i32
    %c8_i32 = arith.constant 8 : i32
    %c4000_i32 = arith.constant 4000 : i32
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<32x64xf32>
    %0 = tt.get_program_id x : i32
    %1 = arith.divsi %0, %c4000_i32 : i32
    %2 = arith.muli %1, %c8_i32 : i32
    %3 = arith.subi %c24_i32, %2 : i32
    %4 = arith.cmpi slt, %3, %c8_i32 : i32
    %5 = arith.select %4, %3, %c8_i32 : i32
    %6 = arith.remsi %0, %5 : i32
    %7 = arith.addi %2, %6 : i32
    %8 = arith.remsi %0, %c4000_i32 : i32
    %9 = arith.divsi %8, %5 : i32
    %10 = arith.muli %7, %c32_i32 : i32
    %11 = tt.make_tensor_ptr %arg1, [%c768_i64, %c90_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E5M2>>
    %12 = tt.advance %11, [%10, %c0_i32] : <tensor<32x64xf8E5M2>>
    %13 = arith.muli %9, %c64_i32 : i32
    %14 = tt.make_tensor_ptr %arg0, [%c90_i64, %c32000_i64], [%c1_i64, %c90_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<64x64xf8E4M3FN>>
    %15 = tt.advance %14, [%c0_i32, %13] : <tensor<64x64xf8E4M3FN>>
    %16:3 = scf.for %arg3 = %c0_i32 to %c90_i32 step %c64_i32 iter_args(%arg4 = %12, %arg5 = %15, %arg6 = %cst_1) -> (!tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>)  : i32 {
      %20 = tt.load %arg4 {boundaryCheck = array<i32: 1>, padding = 1 : i32} : !tt.ptr<tensor<32x64xf8E5M2>>
      %21 = tt.advance %arg4, [%c0_i32, %c64_i32] : <tensor<32x64xf8E5M2>>
      %22 = tt.load %arg5 {boundaryCheck = array<i32: 0>, padding = 1 : i32} : !tt.ptr<tensor<64x64xf8E4M3FN>>
      %23 = tt.advance %arg5, [%c64_i32, %c0_i32] : <tensor<64x64xf8E4M3FN>>
      %24 = tt.fp_to_fp %20 : tensor<32x64xf8E5M2> -> tensor<32x64xf8E4M3FN>
      %25 = arith.subi %c90_i32, %arg3 : i32
      %26 = arith.cmpi slt, %25, %c64_i32 : i32
      %27 = scf.if %26 -> (tensor<32x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
        %32 = tt.splat %25 : i32 -> tensor<1x64xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<1x64xi32>
        %34 = tt.broadcast %33 : tensor<1x64xi1> -> tensor<32x64xi1>
        %35 = arith.select %34, %24, %cst_0 : tensor<32x64xi1>, tensor<32x64xf8E4M3FN>
        scf.yield %35 : tensor<32x64xf8E4M3FN>
      } else {
        scf.yield %24 : tensor<32x64xf8E4M3FN>
      }
      %28 = scf.if %26 -> (tensor<64x64xf8E4M3FN>) {
        %30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
        %31 = tt.expand_dims %30 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
        %32 = tt.splat %25 : i32 -> tensor<64x1xi32>
        %33 = arith.cmpi slt, %31, %32 : tensor<64x1xi32>
        %34 = tt.broadcast %33 : tensor<64x1xi1> -> tensor<64x64xi1>
        %35 = arith.select %34, %22, %cst : tensor<64x64xi1>, tensor<64x64xf8E4M3FN>
        scf.yield %35 : tensor<64x64xf8E4M3FN>
      } else {
        scf.yield %22 : tensor<64x64xf8E4M3FN>
      }
      %29 = tt.dot %27, %28, %arg6, inputPrecision = tf32 {maxNumImpreciseAcc = 2147483647 : i32} : tensor<32x64xf8E4M3FN> * tensor<64x64xf8E4M3FN> -> tensor<32x64xf32>
      scf.yield %21, %23, %29 : !tt.ptr<tensor<32x64xf8E5M2>>, !tt.ptr<tensor<64x64xf8E4M3FN>>, tensor<32x64xf32>
    }
    %17 = tt.fp_to_fp %16#2, rounding = rtne : tensor<32x64xf32> -> tensor<32x64xf8E4M3FN>
    %18 = tt.make_tensor_ptr %arg2, [%c768_i64, %c32000_i64], [%c1_i64, %c768_i64], [%c0_i32, %c0_i32] {order = array<i32: 1, 0>} : <tensor<32x64xf8E4M3FN>>
    %19 = tt.advance %18, [%10, %13] : <tensor<32x64xf8E4M3FN>>
    tt.store %19, %17 : !tt.ptr<tensor<32x64xf8E4M3FN>>
    tt.return
  }
}
```
</details>

Which leads to a failing assertion:
```
#0  0x000073413786d9fc in pthread_kill () from /lib/x86_64-linux-gnu/libc.so.6
#1  0x0000734137819476 in raise () from /lib/x86_64-linux-gnu/libc.so.6
#2  0x00007341377ff7f3 in abort () from /lib/x86_64-linux-gnu/libc.so.6
#3  0x00007341377ff71b in ?? () from /lib/x86_64-linux-gnu/libc.so.6
#4  0x0000734137810e96 in __assert_fail () from /lib/x86_64-linux-gnu/libc.so.6
#5  0x000057d936b1777b in mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion::createDestOps (this=0x733d08425cc0, op=..., adaptor=..., rewriter=..., elemTy=..., operands=..., loc=...)
    at external/triton/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp:500
#6  0x000057d936b17195 in mlir::triton::gpu::ElementwiseOpConversionBase<mlir::triton::FpToFpOp, mlir::triton::gpu::(anonymous namespace)::FpToFpOpConversion>::matchAndRewrite (this=0x733d08425cc0, op=..., adaptor=..., rewriter=...)
    at external/triton/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h:188
[...]
#29 0x000057d93fa6cade in mlir::PassManager::run (this=0x733e80fba158, op=0x733d080bbc20) at external/llvm-project/mlir/lib/Pass/Pass.cpp:885
#30 0x000057d9363f6b1b in xla::gpu::CompileTritonToLLVM (hlo_config=..., hlo_module_name="gemm_fusion_dot.320", device_info=..., block_level_parameters=..., triton_module=..., llvm_module=0x733d0816d6a0, mlir_context=..., is_xla_fusion=true, emit_kernel=true)
    at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1627
#31 0x000057d9363f5a5d in xla::gpu::TritonWrapper (fn_name="gemm_fusion_dot_320_impl", fusion=0x733d080a31c0, cc=std::variant<stream_executor::CudaComputeCapability, stream_executor::RocmComputeCapability> [index 0] = {...}, device_info=..., block_level_parameters=...,
    llvm_module=0x733d0816d6a0, mlir_context=...) at xla/backends/gpu/codegen/triton/fusion_emitter.cc:1531
```

However, this fails Triton compilation:
* First it hits an assertion that the rounding strategy when the destination type is FP8 must be specified
* Adding the rounding strategy, then goes on to another issue, that no methods for converting FP8 <-> FP8 are specified

To work around the above two issues, I propose going through FP16 when both the source and destination types are FP8's.
Copybara import of the project:

--
afd3929099fc4d1045275ca3210e0bc727a2b906 by Kasper Nielsen <[email protected]>:

Fix fused fp8 <-> fp8 conversions

--
66340aa808f58e5dc6ab1c2e06790ceccde95540 by Kasper Nielsen <[email protected]>:

Add unit tests and refactor duplicated code

--
07ae307879eff24ad2f85607e94503deda1074e4 by Kasper Nielsen <[email protected]>:

Run clang-format

--
fe967ff94ffc5f34f07bff142b5d10d81d5e4dce by Kasper Nielsen <[email protected]>:

Fix support conversion tests

Merging this change closes #24114

PiperOrigin-RevId: 741473648
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla: yes ready to pull PR ready for merge process

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants