Merge pull request #2290 from j2kun:extract_slice_of_splat

PiperOrigin-RevId: 814850947

Author: copybara-github

lib/Transforms/FoldConstantTensors/FoldConstantTensors.cpp

@@ -223,15 +223,39 @@ struct CollapseEmptyTensor
   }
 };
 
+struct ExtractSliceOfSplat
+    : public OpRewritePattern<mlir::tensor::ExtractSliceOp> {
+ public:
+  ExtractSliceOfSplat(MLIRContext* context)
+      : OpRewritePattern<mlir::tensor::ExtractSliceOp>(context) {}
+
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::tensor::ExtractSliceOp op,
+                                PatternRewriter& rewriter) const override {
+    auto splatOp =
+        dyn_cast_or_null<tensor::SplatOp>(op.getSource().getDefiningOp());
+    if (!splatOp) return failure();
+
+    auto resultTy = op.getResult().getType();
+    auto newSplat = tensor::SplatOp::create(
+        rewriter, op.getLoc(), splatOp.getInput(), resultTy.getShape());
+    rewriter.replaceOp(op, newSplat);
+    return success();
+  }
+};
+
 struct FoldConstantTensors
     : public impl::FoldConstantTensorsBase<FoldConstantTensors> {
   void runOnOperation() override {
     MLIRContext* context = &getContext();
     auto* module = getOperation();
 
     RewritePatternSet patterns(context);
-    patterns.add<InsertAfterConstant, CollapseShapeAfterConstant,
-                 CollapseEmptyTensor, InsertIntoFromElements>(context);
+    patterns
+        .add<InsertAfterConstant, CollapseShapeAfterConstant,
+             CollapseEmptyTensor, InsertIntoFromElements, ExtractSliceOfSplat>(
+            context);
 
     // Run pattern matching and conversion
     if (failed(applyPatternsGreedily(module, std::move(patterns)))) {

lib/Transforms/FoldConstantTensors/FoldConstantTensors.td

@@ -15,7 +15,8 @@ def FoldConstantTensors : Pass<"fold-constant-tensors"> {
 
     The following folders are supported:
     * `tensor.insert` of a constant tensor
-    * `tensor.collapse_shape` of a constant tensor
+    * `tensor.collapse_shape` of a constant or empty tensor
+    * `tensor.extract_slice` of a splat to a splat of the new shape
   }];
 }
 

tests/Transforms/fold_constant_tensors/extract_slice_of_splat.mlir

@@ -0,0 +1,11 @@
+// RUN: heir-opt --fold-constant-tensors %s | FileCheck %s
+
+// CHECK: func @extract_slice_of_splat
+func.func @extract_slice_of_splat(%arg0: i32) -> (tensor<1024xi32>) {
+  // Fold a collapse shape of a constant
+  // CHECK-NEXT: %[[splat:.+]] = tensor.splat
+  // CHECK-NEXT: return %[[splat]]
+  %0 = tensor.splat %arg0 : tensor<10x1024xi32>
+  %slice = tensor.extract_slice %0[1, 0] [1, 1024] [1, 1] : tensor<10x1024xi32> to tensor<1024xi32>
+  return %slice : tensor<1024xi32>
+}