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

[llvmlistbot at llvm.org]

https://github.com/Qixi-1 created https://github.com/llvm/llvm-project/pull/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.

>From 14840d36550b5d9c7fe64e5a0fb41d50dfa274f1 Mon Sep 17 00:00:00 2001
From: tongjinxuan <tongjinxuan at longcheer.com>
Date: Fri, 3 Jul 2026 14:08:52 +0800
Subject: [PATCH] [mlir][ArithToAMDGPU] Skip tensor scaling ops

---
 .../ArithToAMDGPU/ArithToAMDGPU.cpp           | 10 ++++++++++
 .../ArithToAMDGPU/scaling-truncf-tensor.mlir  | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+)
 create mode 100644 mlir/test/Conversion/ArithToAMDGPU/scaling-truncf-tensor.mlir

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