[Mlir-commits] [mlir] [mlir][Vector] Add interleave(splat, splat) -> wide splat fold (PR #215255)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: Bryth (matteo-rm)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/215255.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+1) 
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+9) 
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+11) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 24fdb4f0bb258..5f701cfcbd1b7 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -585,6 +585,7 @@ def Vector_InterleaveOp :
     }
   }];
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 class ResultIsHalfSourceVectorType<string result> : TypesMatchWith<
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 9322c11d401d9..3bb9d79934af9 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -8489,6 +8489,15 @@ void InterleaveOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<InterleaveDeinterleaveFolder>(context);
 }
 
+OpFoldResult InterleaveOp::fold(FoldAdaptor adaptor) {
+  // interleave(splat(x), splat(x)) -> widened splat(x)
+  auto splat = dyn_cast_if_present<SplatElementsAttr>(adaptor.getLhs());
+  if (!splat || adaptor.getLhs() != adaptor.getRhs())
+    return {};
+  return SplatElementsAttr::get(getResultVectorType(),
+                                splat.getSplatValue<Attribute>());
+}
+
 std::optional<SmallVector<int64_t, 4>> InterleaveOp::getShapeForUnroll() {
   return llvm::to_vector<4>(getResultVectorType().getShape());
 }
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index ed4c908c6e5f2..99a4165b7cb32 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4422,3 +4422,14 @@ func.func @interleave_deinterleave_fold(%arg0: vector<4xf32>) -> vector<4xf32> {
   %result = vector.interleave %even, %odd : vector<2xf32> -> vector<4xf32>
   return %result : vector<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @interleave_splat
+//       CHECK:   %[[C:.*]] = arith.constant dense<0.000000e+00> : vector<[4]xf32>
+//       CHECK:   return %[[C]]
+func.func @interleave_splat() -> vector<[4]xf32> {
+  %cst = arith.constant dense<0.0> : vector<[2]xf32>
+  %0 = vector.interleave %cst, %cst : vector<[2]xf32> -> vector<[4]xf32>
+  return %0 : vector<[4]xf32>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/215255