[Mlir-commits] [mlir] [mlir][ArithToAMDGPU] Skip tensor scaling ops (PR #207327)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Qixi (Qixi-1)

<details>
<summary>Changes</summary>

This fixes a crash in -convert-arith-to-amdgpu=chipset=gfx950 when the input contains tensor arith.scaling_truncf operations.

The scaling rewrite patterns are written for scalar/vector values. Tensor values should not be rewritten by these patterns, so this patch returns failure for tensor operands and leaves those ops unchanged.

A regression test is added for the tensor reproducer from #<!-- -->207311.

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


2 Files Affected:

- (modified) mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp (+10) 
- (added) mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp b/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
index 3372faf4b16cd..d61eeaf3d1e8d 100644
--- a/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
+++ b/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
@@ -464,6 +464,11 @@ ScalingExtFRewritePattern::matchAndRewrite(arith::ScalingExtFOp op,
   if (outVecType && outVecType.isScalable())
     return failure();
 
+  if (isa<RankedTensorType>(out.getType()) ||
+    isa<RankedTensorType>(in.getType()) ||
+    isa<RankedTensorType>(scale.getType()))
+    return failure();
+
   Type scaleF32Type =
       scaleVecType ? VectorType::get(scaleVecType.getShape(), f32) : f32;
   if (scaleType.getIntOrFloatBitWidth() < 32)
@@ -577,6 +582,11 @@ ScalingTruncFRewritePattern::matchAndRewrite(arith::ScalingTruncFOp op,
   if (outVecType && outVecType.isScalable())
     return failure();
 
+  if (isa<RankedTensorType>(out.getType()) ||
+    isa<RankedTensorType>(in.getType()) ||
+    isa<RankedTensorType>(scale.getType()))
+    return failure();
+
   Type scaleF32Type =
       scaleVecType ? VectorType::get(scaleVecType.getShape(), f32) : f32;
   if (scaleType.getIntOrFloatBitWidth() < 32)
diff --git a/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir b/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir
new file mode 100644
index 0000000000000..e92bfafc0babc
--- /dev/null
+++ b/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir
@@ -0,0 +1,19 @@
+// RUN: mlir-opt %s -convert-arith-to-amdgpu=chipset=gfx950 | FileCheck %s
+
+module {
+  // CHECK-LABEL: func.func @m0
+  // CHECK: arith.scaling_truncf
+  func.func @m0(%arg0: tensor<16xf16>, %arg1: tensor<16xf8E8M0FNU>) -> tensor<16xf4E2M1FN> {
+    %0 = arith.scaling_truncf %arg0, %arg1 : tensor<16xf16>, tensor<16xf8E8M0FNU> to tensor<16xf4E2M1FN>
+    return %0 : tensor<16xf4E2M1FN>
+  }
+
+  // CHECK-LABEL: func.func @m1
+  // CHECK: arith.constant
+  // CHECK: arith.scaling_truncf
+  func.func @m1(%arg0: tensor<4xf32>) -> tensor<4xf4E2M1FN> {
+    %cst = arith.constant dense<1.000000e+00> : tensor<4xf8E8M0FNU>
+    %0 = arith.scaling_truncf %arg0, %cst : tensor<4xf32>, tensor<4xf8E8M0FNU> to tensor<4xf4E2M1FN>
+    return %0 : tensor<4xf4E2M1FN>
+  }
+}
\ No newline at end of file

``````````

</details>


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