[Mlir-commits] [mlir] [MLIR][Linalg] Fix LinalgSpecializeGenericOpsPass for scalar inputs (PR #208466)

[Hugo Trachino]

https://github.com/nujaa created https://github.com/llvm/llvm-project/pull/208466

The test case added would crash because the Ranked Tensor-ness of operands was assumed but not checked. This crash only happens when there are multiple inputs as it must bailout somewhere else.

>From df521a8ba9b7907c65f42787e38a620c71a68406 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Thu, 9 Jul 2026 21:41:33 +0800
Subject: [PATCH] [MLIR][Linalg] Fix LinalgSpecializeGenericOpsPass for scalar
 inputs

---
 .../DecomposeGenericByUnfoldingPermutation.cpp |  6 +++---
 .../Linalg/specialize-generic-ops-fail.mlir    | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
index 9015cbb096f88..450cfb2641cb5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
@@ -164,8 +164,8 @@ LogicalResult DecomposeProjectedPermutation::matchAndRewrite(
   // out which operand can supply that runtime-value (tensor.dim).
   // Leaving it as a future TODO.
   if (llvm::any_of(op->getOpOperands(), [](OpOperand &oper) {
-        auto opType = cast<RankedTensorType>(oper.get().getType());
-        return ShapedType::isDynamicShape(opType.getShape());
+        auto opType = dyn_cast<RankedTensorType>(oper.get().getType());
+        return !opType || ShapedType::isDynamicShape(opType.getShape());
       }))
     return failure();
 
@@ -210,7 +210,7 @@ LogicalResult DecomposeProjectedPermutation::matchAndRewrite(
     if (!broadcastedDims.empty()) {
       assert(!broadcastedDims.empty() && "should have non size broadcast");
       Value emptyTensor = tensor::EmptyOp::create(rewriter, loc, outputShape,
-                                                  inputRTType.getElementType());
+                                                  elType);
 
       auto broadcastOp = linalg::BroadcastOp::create(
           rewriter, loc, newInitValues[i], emptyTensor, broadcastedDims);
diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
index 5d66837fca510..bb2c2edfa20f2 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
@@ -46,3 +46,21 @@ func.func @not_copy(%input: tensor<8xi32>, %init: tensor<8xi32>) -> tensor<8xi32
   } -> tensor<8xi32>
   return %res : tensor<8xi32>
 }
+
+
+// -----
+
+// CHECK-LABEL: @scalar_input
+// CHECK: linalg.generic
+#map = affine_map<(d0) -> (d0)>
+#map1 = affine_map<(d0) -> ()>
+func.func @scalar_input(%arg0: tensor<128xi64>, %arg1: f32) -> tensor<128xf32> {
+  %0 = tensor.empty() : tensor<128xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<128xi64>, f32) outs(%0 : tensor<128xf32>) {
+  ^bb0(%in: i64, %in_1: f32, %out: f32):
+    %2 = arith.sitofp %in : i64 to f32
+    %3 = arith.divf %2, %in_1 : f32
+    linalg.yield %3 : f32
+  } -> tensor<128xf32>
+  return %1 : tensor<128xf32>
+}
\ No newline at end of file