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

[Hugo Trachino]

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

>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 1/4] [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

>From d9d23dbb6aac024374ed5fb3cf44712e04eca049 Mon Sep 17 00:00:00 2001
From: Hugo Trachino <hugo.trachino at huawei.com>
Date: Thu, 9 Jul 2026 15:57:34 +0100
Subject: [PATCH 2/4] formatting

---
 .../Transforms/DecomposeGenericByUnfoldingPermutation.cpp    | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
index 450cfb2641cb5..7ee89c0f06f56 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
@@ -209,8 +209,9 @@ LogicalResult DecomposeProjectedPermutation::matchAndRewrite(
     // Does it require broadcast?
     if (!broadcastedDims.empty()) {
       assert(!broadcastedDims.empty() && "should have non size broadcast");
-      Value emptyTensor = tensor::EmptyOp::create(rewriter, loc, outputShape,
-                                                  elType);
+      Value emptyTensor =
+          tensor::EmptyOp::create(rewriter, loc, outputShape, elType);
+
 
       auto broadcastOp = linalg::BroadcastOp::create(
           rewriter, loc, newInitValues[i], emptyTensor, broadcastedDims);

>From 91f3300bf5b6aaa948a10f2cf8b77438841d8871 Mon Sep 17 00:00:00 2001
From: Hugo Trachino <hugo.trachino at huawei.com>
Date: Thu, 9 Jul 2026 16:42:32 +0100
Subject: [PATCH 3/4] Update DecomposeGenericByUnfoldingPermutation.cpp

---
 .../Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
index 7ee89c0f06f56..8ab0f67d9b7de 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
@@ -212,7 +212,6 @@ LogicalResult DecomposeProjectedPermutation::matchAndRewrite(
       Value emptyTensor =
           tensor::EmptyOp::create(rewriter, loc, outputShape, elType);
 
-
       auto broadcastOp = linalg::BroadcastOp::create(
           rewriter, loc, newInitValues[i], emptyTensor, broadcastedDims);
 

>From 3b22e2ce1b7c0ed5f1a951d17365924fbbefb5e2 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Mon, 13 Jul 2026 20:07:37 +0800
Subject: [PATCH 4/4] Simplify test

---
 .../Linalg/specialize-generic-ops-fail.mlir        | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
index bb2c2edfa20f2..f0792d7a696de 100644
--- a/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
+++ b/mlir/test/Dialect/Linalg/specialize-generic-ops-fail.mlir
@@ -54,13 +54,11 @@ func.func @not_copy(%input: tensor<8xi32>, %init: tensor<8xi32>) -> tensor<8xi32
 // 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
+func.func @scalar_input(%arg0: tensor<128xf32>, %arg1: f32) -> tensor<128xf32> {
+  %1 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<128xf32>, f32) outs(%arg0 : tensor<128xf32>) {
+  ^bb0(%in: f32, %in_1: f32, %out: f32):
+    %2 = arith.addf %in_1, %in_1 : f32
+    linalg.yield %2 : f32
   } -> tensor<128xf32>
   return %1 : tensor<128xf32>
-}
\ No newline at end of file
+}