[Mlir-commits] [mlir] 6f10752 - [MLIR][Tensor] canonicalize splats with static shapes (#237560)

[llvmlistbot at llvm.org]

Author: Lukas Scheller
Date: 2026-11-02T14:47:56Z
New Revision: 6f107527c5cfaa041ff1e33b003d56e7fbb8ce37

URL: https://github.com/llvm/llvm-project/commit/6f107527c5cfaa041ff1e33b003d56e7fbb8ce37
DIFF: https://github.com/llvm/llvm-project/commit/6f107527c5cfaa041ff1e33b003d56e7fbb8ce37.diff

LOG: [MLIR][Tensor] canonicalize splats with static shapes (#237560)

Adds the `ReplaceSplatStaticShapeDims` canonicalizer pattern to
`tensor.splat`.
This is similar to `tensor.empty`
[`ReplaceEmptyTensorStaticShapeDims`](https://github.com/llvm/llvm-project/blob/33e1c2c1b8a6fea5aaf6067fde636dee648df794/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L1158)
pattern and turns
```mlir
%dim = arith.constant 4 : index
%splat = tensor.splat %arg0[%dim] : tensor<?xf64>
```
into
```mlir
%splat = tensor.splat %arg0 : tensor<4xf64>
%casted = tensor.cast %splat : tensor<4xf64> to tensor<?xf64>
```

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
    mlir/test/Dialect/Tensor/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index 6a1a19306c14c1..00057218448d13 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1820,6 +1820,7 @@ def Tensor_SplatOp : Tensor_Op<"splat", [
 
   let hasFolder = 1;
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 5a9b9358337031..f3a59df715f3f8 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -4132,6 +4132,44 @@ OpFoldResult SplatOp::fold(FoldAdaptor adaptor) {
   return SplatElementsAttr::get(getType(), {constOperand});
 }
 
+namespace {
+/// Change the type of the result of a `tensor.splat` by making the result
+/// type statically sized along dimensions that in the original operation were
+/// defined as dynamic, but the size was defined using a `constant` op. For
+/// example
+///
+///  %c4 = arith.constant 4 : index
+///  %0 = tensor.splat %arg0[%c4] : tensor<?xf32>
+///
+///  to
+///
+///  %0 = tensor.splat %arg0 : tensor<4xf32>
+struct ReplaceSplatStaticShapeDims : OpRewritePattern<SplatOp> {
+  using OpRewritePattern<SplatOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SplatOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> foldedDynamicSizes;
+    RankedTensorType foldedTensorType = foldDynamicToStaticDimSizes(
+        op.getType(), op.getDynamicSizes(), foldedDynamicSizes);
+
+    // Stop here if no dynamic size was promoted to static.
+    if (foldedTensorType == op.getType())
+      return failure();
+
+    auto newOp = SplatOp::create(rewriter, op.getLoc(), op.getInput(),
+                                 foldedTensorType, foldedDynamicSizes);
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
+    return success();
+  }
+};
+} // namespace
+
+void SplatOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                          MLIRContext *context) {
+  results.add<ReplaceSplatStaticShapeDims>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // Common Canonicalizers and Folders.
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 690b08620987a2..ef44904219c7ab 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -3267,3 +3267,34 @@ func.func @select_uniform_unselected_poison_tensor(%p: i1, %a: tensor<4xi32>) ->
   %x = arith.select %m, %a, %bad : tensor<4xi1>, tensor<4xi32>
   return %x : tensor<4xi32>
 }
+
+// -----
+
+// Non-constant size with constant index -> splat with static type + cast
+// CHECK-LABEL:   func.func @splat_dynamic_size_constant(
+// CHECK-SAME:      %[[ARG0:.*]]: f32) -> tensor<?xf32> {
+// CHECK:           %[[SPLAT:.*]] = tensor.splat %[[ARG0]] : tensor<4xf32>
+// CHECK:           %[[CAST:.*]] = tensor.cast %[[SPLAT]] : tensor<4xf32> to tensor<?xf32>
+// CHECK:           return %[[CAST]] : tensor<?xf32>
+// CHECK:         }
+func.func @splat_dynamic_size_constant(%v: f32) -> tensor<?xf32> {
+  %c4 = arith.constant 4 : index
+  %0 = tensor.splat %v[%c4] : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Mixed sizes -> splat with semi-static type + cast
+// CHECK-LABEL:   func.func @splat_dynamic_size_mixed(
+// CHECK-SAME:      %[[ARG0:.*]]: f32,
+// CHECK-SAME:      %[[ARG1:.*]]: index) -> tensor<?x?xf32> {
+// CHECK:           %[[SPLAT:.*]] = tensor.splat %[[ARG0]]{{\[}}%[[ARG1]]] : tensor<?x4xf32>
+// CHECK:           %[[CAST:.*]] = tensor.cast %[[SPLAT]] : tensor<?x4xf32> to tensor<?x?xf32>
+// CHECK:           return %[[CAST]] : tensor<?x?xf32>
+// CHECK:         }
+func.func @splat_dynamic_size_mixed(%v: f32, %m: index) -> tensor<?x?xf32> {
+  %c4 = arith.constant 4 : index
+  %0 = tensor.splat %v[%m, %c4] : tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}