[mlir][Linalg] Add canonicalization for init_tensor -> subtensor op.

[asaadaldien]

LGTM!

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

909

should it be a failure in case of not zero offset and or not unit strides ?

913

Is that guaranteed to have static sizes ?

[mravishankar]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

909

I think it doesnt matter. The size is all that is needed. Thats the semantics of init_tensor.

913

It will return -1 when it is static.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (hasDynamic) {
       dynamicShapes.push_back(dynamicDim);
     }
   }
   return builder.create<InitTensorOp>(loc, dynamicShapes, staticShapes,
                                       srcType.getElementType());
 }
 
 namespace {
-struct FoldWithTensorReshapeOp : public OpRewritePattern<TensorReshapeOp> {
+/// Since `init_tensor` operation creates a tensor needed only for its shape, a
+/// subtensor of this is also needed only for its shape. The result can be
+/// replaced by a new init_tensor operation of the same size as the subtensor
+/// op.
+struct FoldInitTensorWithSubTensorOp : public OpRewritePattern<SubTensorOp> {
+  using OpRewritePattern<SubTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SubTensorOp subtensorOp,
+                                PatternRewriter &rewriter) const override {
+    if (!subtensorOp.source().getDefiningOp<linalg::InitTensorOp>())
+      return failure();
+    rewriter.replaceOpWithNewOp<linalg::InitTensorOp>(
+        subtensorOp, subtensorOp.sizes(),
+        llvm::to_vector<4>(llvm::map_range(
+            subtensorOp.static_sizes(),
+            [](Attribute attr) { return attr.cast<IntegerAttr>().getInt(); })),
+        subtensorOp.getSourceType().getElementType());
+    return success();
+  }
+};
+
+struct FoldInitTensorWithTensorReshapeOp
+    : public OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     if (!reshapeOp.src().getDefiningOp<InitTensorOp>())
       return failure();
     RankedTensorType collapsedType = reshapeOp.getSrcType();
     RankedTensorType expandedType = reshapeOp.getResultType();
     bool isCollapsed = expandedType.getRank() < collapsedType.getRank();
     if (isCollapsed)
       std::swap(collapsedType, expandedType);
     Value initTensorOp = isCollapsed
                              ? getCollapsedInitTensor(rewriter, reshapeOp)
                              : getExpandedInitTensor(rewriter, reshapeOp);
     if (!initTensorOp)
       return failure();
     rewriter.replaceOp(reshapeOp, initTensorOp);
     return success();
   }
 };
 } // namespace
 
 void InitTensorOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results.insert<FoldWithTensorReshapeOp, ReplaceDimOfInitTensorOp,
-                 ReplaceStaticShapeDims>(context);
+  results
+      .insert<FoldInitTensorWithSubTensorOp, FoldInitTensorWithTensorReshapeOp,
+              ReplaceDimOfInitTensorOp, ReplaceStaticShapeDims>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // PadTensorOp
 //===----------------------------------------------------------------------===//
 
 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {

mlir/test/Dialect/Linalg/canonicalize.mlir

%3:2 = linalg.generic
     ^bb0(%arg3: f32, %arg4 : f32, %arg5 : f32, %arg6 : f32):
       linalg.yield %arg2, %arg4 : f32, f32
     } -> tensor<?x?xf32>, tensor<?x?xf32>
   return %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
 }
 // CHECK-LABEL: func @keep_not_noop
 //       CHECK:   %[[RESULT:.+]]:2 = linalg.generic
 //       CHECK:   return %[[RESULT]]#0, %[[RESULT]]#1
+
+// -----
+
+func @fold_init_tensor_with_subtensor
+  (%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>
+{
+  %0 = linalg.init_tensor[%arg0, 10, 40] : tensor<?x10x40xf32>
+  %1 = subtensor %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]
+    : tensor<?x10x40xf32> to tensor<5x?x20xf32>
+  return %1 : tensor<5x?x20xf32>
+}
+//      CHECK: func @fold_init_tensor_with_subtensor
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
+//      CHECK:   %[[T0:.+]] = linalg.init_tensor [5, %[[ARG1]], 20]
+//      CHECK:   return %[[T0]]