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

[Lukas Scheller]

https://github.com/Schottkyc137 created https://github.com/llvm/llvm-project/pull/237560

Adds the `ReplaceSplatStaticShapeDims` canonicalizer pattern to `tensor.splat`.
This is similar to `tensor.empty` `ReplaceEmptyTensorStaticShapeDims` 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>
```

I think this is ready for review, but I'd like to have some pointers concerning design decisions:

1) This pattern can potentially regress the [`SelectSplatCondition`](https://github.com/llvm/llvm-project/blob/33e1c2c1b8a6fea5aaf6067fde636dee648df794/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L4205) pattern (and potentially others too):
```mlir
%c = arith.constant 4 : index
%m = tensor.splat %p[%c] : tensor<?xi1>
%x = arith.select %m, %a, %b : tensor<?xi1>, tensor<?xi32>
```
Using top-down canonicalisation, this pattern now inserts a cast before the `select`. Previously, this would have canonicalised to using `p` directly.
I had considered having the `SelectSplatCondition` look through cast ops, but this must be repeated for every other op as well. Maybe a canonicalisation pattern to push casts through an `arith::SelectOp` is correct?

2) The new pattern is very similar to the [`ReplaceEmptyTensorStaticShapeDims`](https://github.com/llvm/llvm-project/blob/33e1c2c1b8a6fea5aaf6067fde636dee648df794/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L1158) pattern.
I'm happy to refactor the two if this is desired (either here or in a follow-up PR)

>From 5262decf264a131f20483079fb3d1593820e32a4 Mon Sep 17 00:00:00 2001
From: Lukas Scheller <lukas2121 at icloud.com>
Date: Sat, 31 Oct 2026 13:45:29 +0100
Subject: [PATCH] [MLIR][Tensor] canonicalize splats with static shapes

---
 .../mlir/Dialect/Tensor/IR/TensorOps.td       |  1 +
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp      | 38 +++++++++++++++++++
 mlir/test/Dialect/Tensor/canonicalize.mlir    | 31 +++++++++++++++
 3 files changed, 70 insertions(+)

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>
+}