[Mlir-commits] [mlir] [mlir][AMDGPU] Diagnose unsupported packed fp8 conversions (PR #210565)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: qyingwu

<details>
<summary>Changes</summary>

##Summary
Fixes #<!-- -->210266.

`amdgpu.packed_trunc_2xfp8` can crash during AMDGPU-to-ROCDL conversion when the result element type does not match the selected chipset. For example, `gfx942` expects FNUZ fp8 types, but the reproducer uses the OCP `f8E4M3FN` type. The lowering then failed to create an intrinsic result and later tried to bitcast a null value.

This patch diagnoses unsupported packed fp8 conversions instead of crashing. It also adds the same guard for `amdgpu.packed_stoch_round_fp8`, which had the same unchecked result path.

##Tests
  ```
ninja -C build mlir-opt FileCheck count not
build/bin/llvm-lit -v mlir/test/Conversion/AMDGPUToROCDL/packed-trunc-invalid.mlir
build/bin/llvm-lit -v mlir/test/Conversion/AMDGPUToROCDL/8-bit-floats.mlir mlir/test/
Conversion/AMDGPUToROCDL/8-bit-floats-ocp.mlir mlir/test/Conversion/AMDGPUToROCDL/packed-trunc-invalid.mlir
build/bin/llvm-lit -v mlir/test/Conversion/AMDGPUToROCDL
```

---
Full diff: https://github.com/llvm/llvm-project/pull/210565.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp (+7) 
- (added) mlir/test/Conversion/AMDGPUToROCDL/packed-trunc-invalid.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 90e15b1680446..5c4195778be5f 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -3038,6 +3038,9 @@ LogicalResult PackedTrunc2xFp8OpLowering::matchAndRewrite(
   else if (typeIsExpectedFp8ForChipset(chipset, resultElemType))
     result = ROCDL::CvtPkFp8F32Op::create(rewriter, loc, i32, sourceA, sourceB,
                                           existing, op.getWordIndex());
+  else
+    return op.emitOpError(
+        "no intrinsic matching packed trunc conversion on the given chipset");
 
   result = rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(
       op, getTypeConverter()->convertType(resultType), result);
@@ -3072,6 +3075,10 @@ LogicalResult PackedStochRoundFp8OpLowering::matchAndRewrite(
   else if (typeIsExpectedFp8ForChipset(chipset, resultElemType))
     result = ROCDL::CvtSrFp8F32Op::create(rewriter, loc, i32, source, stoch,
                                           existing, op.getStoreIndex());
+  else
+    return op.emitOpError(
+        "no intrinsic matching packed stochastic round conversion on the given "
+        "chipset");
 
   result = rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(
       op, getTypeConverter()->convertType(resultType), result);
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/packed-trunc-invalid.mlir b/mlir/test/Conversion/AMDGPUToROCDL/packed-trunc-invalid.mlir
new file mode 100644
index 0000000000000..874fcae6555a6
--- /dev/null
+++ b/mlir/test/Conversion/AMDGPUToROCDL/packed-trunc-invalid.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt %s --convert-amdgpu-to-rocdl=chipset=gfx942 --split-input-file --verify-diagnostics
+
+func.func @packed_trunc_ocp_type_requires_ocp_chipset(%arg0: f32) {
+  // expected-error at below {{'amdgpu.packed_trunc_2xfp8' op no intrinsic matching packed trunc conversion on the given chipset}}
+  // expected-error at below {{failed to legalize operation 'amdgpu.packed_trunc_2xfp8'}}
+  %0 = amdgpu.packed_trunc_2xfp8 %arg0, undef into undef[word 0] : f32 to vector<4xf8E4M3FN>
+  return
+}
+
+// -----
+
+func.func @packed_stoch_round_ocp_type_requires_ocp_chipset(%arg0: f32,
+                                                            %arg1: i32) {
+  // expected-error at below {{'amdgpu.packed_stoch_round_fp8' op no intrinsic matching packed stochastic round conversion on the given chipset}}
+  // expected-error at below {{failed to legalize operation 'amdgpu.packed_stoch_round_fp8'}}
+  %0 = amdgpu.packed_stoch_round_fp8 %arg0 + %arg1 into undef[0] : f32 to vector<4xf8E4M3FN>
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/210565