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

[llvmlistbot at llvm.org]

Author: Qixi
Date: 2026-07-06T10:33:31-04:00
New Revision: 4109893350feeea69214d938864e200ca5eec7b5

URL: https://github.com/llvm/llvm-project/commit/4109893350feeea69214d938864e200ca5eec7b5
DIFF: https://github.com/llvm/llvm-project/commit/4109893350feeea69214d938864e200ca5eec7b5.diff

LOG: [mlir][ArithToAMDGPU] Skip tensor scaling ops (#207327)

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.

Co-authored-by: tongjinxuan <tongjinxuan at longcheer.com>

Added: 
    mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir

Modified: 
    mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp b/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
index 3372faf4b16cd..cf44c2a0033ac 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..d22f35a6d07f1
--- /dev/null
+++ b/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -convert-arith-to-amdgpu=chipset=gfx950 | FileCheck %s
+
+// 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>
+}