[mlir][tensor][bufferize] Implement getBufferType for Expand/CollapseShapeOp

[springerm]

Summary

This function must be implemented for all ops, where the result memref type is different from the input memref type.

[jreiffers]

Thanks, I was about to do just that :-).

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

[{
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>,
 
     // Builder that infers the result layout map. The result shape must be
     // specified. Otherwise, the op may be ambiguous.
     OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
                "ArrayRef<ReassociationIndices>":$reassociation)>
   ];
-  let extraClassDeclaration = commonExtraClassDeclaration;
+
+  let extraClassDeclaration = commonExtraClassDeclaration # [{
+    static FailureOr<MemRefType> computeExpandedType(
+        MemRefType srcType, ArrayRef<int64_t> resultShape,
+        ArrayRef<ReassociationIndices> reassociation);
+  }];
+
   let hasVerifier = 1;
 }
 
 def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape"> {
   let summary = "operation to produce a memref with a smaller rank.";
   let description = [{
     The `memref.collapse_shape` op produces a new view with a smaller rank
     whose sizes are a reassociation of the original `view`. The operation is
let description = [{
     sizes and/or strides. I.e., a result dimension must be dynamic if and only
     if at least one dimension in the corresponding reassociation group is
     dynamic. Similarly, the stride of a result dimension must be dynamic if and
     only if the corresponding start dimension in the source type is dynamic.
 
     Note: This op currently assumes that the inner strides are of the
     source/result layout map are the faster-varying ones.
   }];
+
   let builders = [
     // Builders for a contracting reshape whose result type is computed from
     // `src` and `reassociation`.
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationIndices>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
     OpBuilder<(ins "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
OpBuilder<(ins "Type":$resultType, "Value":$src,
       "ArrayRef<ReassociationExprs>":$reassociation,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
     [{
       auto reassociationMaps =
           convertReassociationMapsToIndices($_builder, reassociation);
       build($_builder, $_state, resultType, src, reassociationMaps, attrs);
     }]>
   ];
+
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     /// Return `true` if this source MemRef type is guaranteed to be collapsible
     /// according to the given reassociation indices. In the presence of dynamic
     /// strides this is usually not the case.
     static bool isGuaranteedCollapsible(
         MemRefType srcType, ArrayRef<ReassociationIndices> reassociation);
+
+    static MemRefType computeCollapsedType(
+        MemRefType srcType, ArrayRef<ReassociationIndices> reassociation);
   }];
 
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

for (unsigned idx = 0, e = reassoc.size(); idx < e; ++idx) {
                                   .asStride();
     }
   }
   auto resultStrides = llvm::to_vector<8>(llvm::reverse(reverseResultStrides));
   resultStrides.resize(resultShape.size(), 1);
   return StridedLayoutAttr::get(srcType.getContext(), srcOffset, resultStrides);
 }
 
-static FailureOr<MemRefType>
-computeExpandedType(MemRefType srcType, ArrayRef<int64_t> resultShape,
-                    ArrayRef<ReassociationIndices> reassociation) {
+FailureOr<MemRefType> ExpandShapeOp::computeExpandedType(
+    MemRefType srcType, ArrayRef<int64_t> resultShape,
   if (srcType.getLayout().isIdentity()) {
     // If the source is contiguous (i.e., no layout map specified), so is the
     // result.
     MemRefLayoutAttrInterface layout;
     return MemRefType::get(resultShape, srcType.getElementType(), layout,
                            srcType.getMemorySpace());
   }
 
   // Source may not be contiguous. Compute the layout map.
   FailureOr<StridedLayoutAttr> computedLayout =
       computeExpandedLayoutMap(srcType, resultShape, reassociation);
   if (failed(computedLayout))
     return failure();
   return MemRefType::get(resultShape, srcType.getElementType(), *computedLayout,
                          srcType.getMemorySpace());
 }
 
 void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
                           ArrayRef<int64_t> resultShape, Value src,
                           ArrayRef<ReassociationIndices> reassociation) {
   // Only ranked memref source values are supported.
   auto srcType = src.getType().cast<MemRefType>();
   FailureOr<MemRefType> resultType =
-      computeExpandedType(srcType, resultShape, reassociation);
+      ExpandShapeOp::computeExpandedType(srcType, resultShape, reassociation);
   // Failure of this assertion usually indicates a problem with the source
   // type, e.g., could not get strides/offset.
   assert(succeeded(resultType) && "could not compute layout");
   build(builder, result, *resultType, src, reassociation);
 }
 
 LogicalResult ExpandShapeOp::verify() {
   MemRefType srcType = getSrcType();
   MemRefType resultType = getResultType();
 
   // Verify result shape.
   if (failed(verifyCollapsedShape(getOperation(), srcType.getShape(),
                                   resultType.getShape(),
                                   getReassociationIndices(),
                                   /*allowMultipleDynamicDimsPerGroup=*/false)))
     return failure();
 
   // Compute expected result type (including layout map).
-  FailureOr<MemRefType> expectedResultType = computeExpandedType(
+  FailureOr<MemRefType> expectedResultType = ExpandShapeOp::computeExpandedType(
       srcType, resultType.getShape(), getReassociationIndices());
   if (failed(expectedResultType))
     return emitOpError("invalid source layout map");
 
   // Check actual result type.
   if (*expectedResultType != resultType)
     return emitOpError("expected expanded type to be ")
            << *expectedResultType << " but found " << resultType;
bool CollapseShapeOp::isGuaranteedCollapsible(
   // MemRefs with identity layout are always collapsible.
   if (srcType.getLayout().isIdentity())
     return true;
 
   return succeeded(computeCollapsedLayoutMap(srcType, reassociation,
                                              /*strict=*/true));
 }
 
-static MemRefType
-computeCollapsedType(MemRefType srcType,
-                     ArrayRef<ReassociationIndices> reassociation) {
+MemRefType CollapseShapeOp::computeCollapsedType(
+    MemRefType srcType, ArrayRef<ReassociationIndices> reassociation) {
   SmallVector<int64_t> resultShape;
   resultShape.reserve(reassociation.size());
   for (const ReassociationIndices &group : reassociation) {
     using saturated_arith::Wrapper;
     auto groupSize = Wrapper::size(1);
     for (int64_t srcDim : group)
       groupSize = groupSize * Wrapper::size(srcType.getDimSize(srcDim));
     resultShape.push_back(groupSize.asSize());
MemRefType CollapseShapeOp::computeCollapsedType(
   return MemRefType::get(resultShape, srcType.getElementType(), *computedLayout,
                          srcType.getMemorySpace());
 }
 
 void CollapseShapeOp::build(OpBuilder &b, OperationState &result, Value src,
                             ArrayRef<ReassociationIndices> reassociation,
                             ArrayRef<NamedAttribute> attrs) {
   auto srcType = src.getType().cast<MemRefType>();
-  MemRefType resultType = computeCollapsedType(srcType, reassociation);
+  MemRefType resultType =
+      CollapseShapeOp::computeCollapsedType(srcType, reassociation);
   build(b, result, resultType, src, attrs);
   result.addAttribute(::mlir::getReassociationAttrName(),
                       getReassociationIndicesAttribute(b, reassociation));
 }
 
 LogicalResult CollapseShapeOp::verify() {
   MemRefType srcType = getSrcType();
   MemRefType resultType = getResultType();
LogicalResult matchAndRewrite(CollapseShapeOp op,
                                 PatternRewriter &rewriter) const override {
     auto cast = op.getOperand().getDefiningOp<CastOp>();
     if (!cast)
       return failure();
 
     if (!CastOp::canFoldIntoConsumerOp(cast))
       return failure();
 
-    Type newResultType =
-        computeCollapsedType(cast.getOperand().getType().cast<MemRefType>(),
-                             op.getReassociationIndices());
+    Type newResultType = CollapseShapeOp::computeCollapsedType(
+        cast.getOperand().getType().cast<MemRefType>(),
 
     if (newResultType == op.getResultType()) {
       rewriter.updateRootInPlace(
           op, [&]() { op.getSrcMutable().assign(cast.getSource()); });
     } else {
       Value newOp = rewriter.create<CollapseShapeOp>(
           op->getLoc(), cast.getSource(), op.getReassociationIndices());
       rewriter.replaceOpWithNewOp<CastOp>(op, op.getType(), newOp);

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+    auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
+    auto maybeSrcBufferType = bufferization::getBufferType(
+        collapseShapeOp.getSrc(), options, fixedTypes);
+    if (failed(maybeSrcBufferType))
+      return failure();
+    auto srcBufferType = maybeSrcBufferType->cast<MemRefType>();
+    bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
+        srcBufferType, collapseShapeOp.getReassociationIndices());
+
+    if (!canBeCollapsed) {
+      // If dims cannot be collapsed, this op bufferizes to a new allocation.
+      RankedTensorType tensorResultType = collapseShapeOp.getResultType();
+      return bufferization::getMemRefTypeWithStaticIdentityLayout(
+          tensorResultType, srcBufferType.getMemorySpaceAsInt());
+    }
+
+    return memref::CollapseShapeOp::computeCollapsedType(
+        srcBufferType, collapseShapeOp.getReassociationIndices());
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     RankedTensorType tensorResultType = collapseShapeOp.getResultType();
     FailureOr<Value> maybeBuffer =
         getBuffer(rewriter, collapseShapeOp.getSrc(), options);
     if (failed(maybeBuffer))
       return failure();
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+    auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
+    auto maybeSrcBufferType = bufferization::getBufferType(
+        expandShapeOp.getSrc(), options, fixedTypes);
+    if (failed(maybeSrcBufferType))
+      return failure();
+    auto srcBufferType = maybeSrcBufferType->cast<MemRefType>();
+    auto maybeResultType = memref::ExpandShapeOp::computeExpandedType(
+        srcBufferType, expandShapeOp.getResultType().getShape(),
+        expandShapeOp.getReassociationIndices());
+    if (failed(maybeResultType))
+      return failure();
+    return *maybeResultType;
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
     auto tensorResultType = expandShapeOp.getResultType();
     FailureOr<Value> buffer =
         getBuffer(rewriter, expandShapeOp.getSrc(), options);
     if (failed(buffer))
       return failure();

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

func.func @non_tensor_for_arg(%A : tensor<?xf32> {bufferization.writable = true})
   %c2 = arith.constant 2.0 : f32
   %c10 = arith.constant 10 : index
   %r1:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%idx = %c1, %t = %A) -> (index, tensor<?xf32>) {
     %t2 = tensor.insert %c2 into %t[%idx] : tensor<?xf32>
     scf.yield %idx, %t2 : index, tensor<?xf32>
   }
   return %r1#1 : tensor<?xf32>
 }
+
+// -----
+
+// This is a regression test. Just check that the IR bufferizes.
+  
+// CHECK-LABEL: func @buffer_type_of_collapse_shape
+func.func @buffer_type_of_collapse_shape(%arg0: tensor<f64>) {
+  %true = arith.constant true
+  %0 = scf.while (%arg1 = %arg0) : (tensor<f64>) -> (tensor<f64>) {
+    scf.condition(%true) %arg1 : tensor<f64>
+  } do {
+  ^bb0(%_: tensor<f64>):
+    %3 = bufferization.alloc_tensor() : tensor<1xf64>
+    %16 = tensor.collapse_shape %3 [] : tensor<1xf64> into tensor<f64>
+    scf.yield %16 : tensor<f64>
+  }
+  return
+}