[mlir][Linalg] Fold init_tensor -> linalg.tensor_reshape.

[mravishankar]

Summary

Reshaping an init_tensor can be folded to a init_tensor op of the
final type.

Depends On D93748

[mravishankar]

Rebase

[hanchung]

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

755

This looks weird, how about renaming it to shape?

767

ditto

782–783

Do we really want to crash here? I feel we should just return nullptr and stop rewriting the pattern. Although it's really wrong to have more than one dynamic dims, I don't think it is the pass' responsibility to check it. Otherwise, we could leave many similar check in many passes. I think maybe we can check it in reshape op.

[hanchung]

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

782–783

I saw a similar check in https://reviews.llvm.org/D93724

With the patch, I think we don't need to check it here. We can have a follow cleanup for this.

[mravishankar]

Rebase

[mravishankar]

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

782–783

Thats right. I added it to the verifier too. I just returned nullptr. There might be a path to relaxing this restriction with shape dialect. So I added an assert to say the pattern doesnt handle it. But making it return with nullptr is better. THanks!

[mravishankar]

Address comments

[mravishankar]

Fix build errors

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

LogicalResult matchAndRewrite(DimOp dimOp,
     } else {
       rewriter.replaceOp(dimOp, initTensorOp.getDynamicSize(index));
     }
     return success();
   }
 };
 } // namespace
 
+static Value getCollapsedInitTensor(OpBuilder &builder,
+                                    TensorReshapeOp reshapeOp) {
+  Location loc = reshapeOp.getLoc();
+  SmallVector<Value, 4> dynamicShapes;
+  SmallVector<int64_t, 4> staticShapes;
+  auto reassociation = reshapeOp.getReassociationMaps();
+  Value src = reshapeOp.src();
+  RankedTensorType srcType = reshapeOp.getSrcType();
+  ArrayRef<int64_t> srcShape = srcType.getShape();
+  for (auto map : reassociation) {
+    Value linearizedDynamicDim = nullptr;
+    int64_t linearizedStaticDim = 1;
+    for (unsigned i : llvm::map_range(map.getResults(), [](AffineExpr e) {
+           return e.cast<AffineDimExpr>().getPosition();
+         })) {
+      if (ShapedType::isDynamic(srcShape[i])) {
+        Value shapeVal = builder.create<DimOp>(loc, src, i);
+        if (linearizedDynamicDim) {
+          linearizedDynamicDim =
+              builder.create<MulIOp>(loc, linearizedDynamicDim, shapeVal);
+        } else {
+          linearizedDynamicDim = shapeVal;
+        }
+      } else {
+        linearizedStaticDim *= srcShape[i];
+      }
+    }
+    if (linearizedDynamicDim) {
+      if (linearizedStaticDim != 1) {
+        linearizedDynamicDim = builder.create<MulIOp>(
+            loc, linearizedDynamicDim,
+            builder.create<ConstantIndexOp>(loc, linearizedStaticDim));
+      }
+      dynamicShapes.push_back(linearizedDynamicDim);
+      staticShapes.push_back(ShapedType::kDynamicSize);
+    } else {
+      staticShapes.push_back(linearizedStaticDim);
+    }
+  }
+  return builder.create<InitTensorOp>(loc, dynamicShapes, staticShapes,
+                                      srcType.getElementType());
+}
+
+static Value getExpandedInitTensor(OpBuilder &builder,
+                                   TensorReshapeOp reshapeOp) {
+  SmallVector<Value, 4> dynamicShapes;
+  SmallVector<int64_t, 4> staticShapes;
+  auto reassociation = reshapeOp.getReassociationMaps();
+  Value src = reshapeOp.src();
+  RankedTensorType srcType = reshapeOp.getSrcType();
+  ArrayRef<int64_t> srcShape = srcType.getShape();
+  ArrayRef<int64_t> dstShape = reshapeOp.getResultType().getShape();
+  Location loc = reshapeOp.getLoc();
+  for (auto map : enumerate(reassociation)) {
+    int64_t linearizedStaticDim = 1;
+    bool hasDynamic = false;
+    for (unsigned i :
+         llvm::map_range(map.value().getResults(), [](AffineExpr e) {
+           return e.cast<AffineDimExpr>().getPosition();
+         })) {
+      if (ShapedType::isDynamic(dstShape[i])) {
+        if (hasDynamic)
+          return nullptr;
+        hasDynamic = true;
+        staticShapes.push_back(ShapedType::kDynamicSize);
+        continue;
+      }
+      staticShapes.push_back(dstShape[i]);
+      linearizedStaticDim *= dstShape[i];
+    }
+    if (hasDynamic) {
+      assert(ShapedType::isDynamic(srcShape[map.index()]) &&
+             "expected shape in collapsed type to be static as well");
+      Value dynamicDim = builder.create<DimOp>(loc, src, map.index());
+      if (linearizedStaticDim != 1) {
+        dynamicDim = builder.create<UnsignedDivIOp>(
+            loc, dynamicDim,
+            builder.create<ConstantIndexOp>(loc, linearizedStaticDim));
+      }
+      dynamicShapes.push_back(dynamicDim);
+    }
+  }
+  return builder.create<InitTensorOp>(loc, dynamicShapes, staticShapes,
+                                      srcType.getElementType());
+}
+
+namespace {
+struct FoldWithTensorReshapeOp : 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<ReplaceDimOfInitTensorOp, ReplaceStaticShapeDims>(context);
+  results.insert<FoldWithTensorReshapeOp, ReplaceDimOfInitTensorOp,
+                 ReplaceStaticShapeDims>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 /// Collapse reassociation maps that are used in pair of reshape ops where one
 /// is a producer and other is the consumer. Only valid to use this method when

 template <typename OpTy>
 static LogicalResult verifyReshapeLikeShapes(OpTy op, ShapedType collapsedType,
                                              ShapedType expandedType,
                                              bool isExpandingReshape) {
   ArrayRef<int64_t> collapsedShape = collapsedType.getShape();
   ArrayRef<int64_t> expandedShape = expandedType.getShape();
   unsigned expandedDimStart = 0;
   for (auto map : llvm::enumerate(op.getReassociationMaps())) {
-    Optional<int64_t> dynamicDims;
+    Optional<int64_t> dynamicShape;
     int64_t linearizedStaticShape = 1;
     for (auto dim : llvm::enumerate(expandedShape.slice(
              expandedDimStart, map.value().getNumResults()))) {
       if (ShapedType::isDynamic(dim.value())) {
-        if (isExpandingReshape && dynamicDims) {
+        if (isExpandingReshape && dynamicShape) {
           return op->emitOpError("invalid to have a single dimension (")
                  << map.index() << ") expanded into multiple dynamic dims ("
-                 << expandedDimStart + dynamicDims.getValue() << ","
+                 << expandedDimStart + dynamicShape.getValue() << ","
                  << expandedDimStart + dim.index() << ")";
         }
-        dynamicDims = dim.index();
+        dynamicShape = dim.index();
       } else {
         linearizedStaticShape *= dim.value();
       }
     }
-    if (dynamicDims) {
+    if (dynamicShape) {
       if (!ShapedType::isDynamic(collapsedShape[map.index()])) {
         return op->emitOpError("expected dimension ")
                << map.index()
                << " of collapsed type to be dynamic since one or more of the "
                   "corresponding dimensions in the expanded type is dynamic";
       }
     } else {
       if (collapsedShape[map.index()] != linearizedStaticShape) {

mlir/test/Dialect/Linalg/canonicalize.mlir

func @init_tensor_dim_of_linalg_result(%arg_0 : tensor<?xf32>,
   %num_elem_1 = dim %1, %c0 : tensor<?xf32>
   %result_1 = linalg.init_tensor [%num_elem_1] : tensor<?xf32>
   return %result_0, %result_1 : tensor<?xf32>, tensor<?xf32>
 }
 // CHECK-LABEL: func @init_tensor_dim_of_linalg_result(
 // CHECK-SAME: [[ARG_0:%.*]]: tensor<?xf32>, [[ARG_1:%.*]]: tensor<?xf32>)
 // CHECK: dim [[ARG_0]]
 // CHECK: dim [[ARG_1]]
+
+// -----
+
+func @init_tensor_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
+  %0 = linalg.init_tensor [6, 5, %arg0] : tensor<6x5x?xf32>
+  %1 = linalg.tensor_reshape %0
+    [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
+     affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
+     affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>] :
+     tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
+  return %1 : tensor<2x3x5x4x?x7xf32>
+}
+//      CHECK: func @init_tensor_reshape_expansion
+// CHECK-SAME:   %[[ARG0:.+]]: index
+//      CHECK:   %[[C28:.+]] = constant 28 : index
+//      CHECK:   %[[T0:.+]] = divi_unsigned %[[ARG0]], %[[C28]]
+//      CHECK:   %[[T1:.+]] = linalg.init_tensor [2, 3, 5, 4, %[[T0]], 7]
+//      CHECK:   return %[[T1]]
+
+// -----
+
+func @init_tensor_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
+  %0 = linalg.init_tensor [2, 3, 5, 4, %arg0, 7] : tensor<2x3x5x4x?x7xf32>
+  %1 = linalg.tensor_reshape %0
+    [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>,
+     affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>,
+     affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>] :
+    tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>
+  return %1 : tensor<6x5x?xf32>
+}
+//      CHECK: func @init_tensor_reshape_collapse
+// CHECK-SAME:   %[[ARG0:.+]]: index
+//      CHECK:   %[[C28:.+]] = constant 28 : index
+//      CHECK:   %[[T0:.+]] = muli %[[ARG0]], %[[C28]]
+//      CHECK:   %[[T1:.+]] = linalg.init_tensor [6, 5, %[[T0]]]
+//      CHECK:   return %[[T1]]