[mlir][linalg][bufferize] Support TensorExpand/CollapseShape

[springerm]

Summary

Depends On D112511

[springerm]

update

[springerm]

rebase

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

1993

can we use something like:

llvm::map_iterator(
  attr.cast<ArrayAttr>().getAsValueRange<IntegerAttr>(),
  [](const APInt &val) { return !val.getInt(); }))

?

1997

dyn_cast + assert is an anti-pattern FWIH, you can just cast

[springerm]

update

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

return
           ConstantOp,
           tensor::DimOp,
           ExtractSliceOp,
           scf::ForOp,
           InsertSliceOp,
           InitTensorOp,
           LinalgOp,
           ReturnOp,
+          TensorCollapseShapeOp,
+          TensorExpandShapeOp,
           TiledLoopOp,
           VectorTransferOpInterface,
           linalg::YieldOp,
           scf::YieldOp>(op)
       // clang-format on
       || (none_of(op->getResultTypes(), isaTensor) &&
           none_of(op->getOperandTypes(), isaTensor));
 }
return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
         // result(s).
         .Case<tensor::CastOp,
               scf::ForOp,
               InsertSliceOp,
               LinalgOp,
               TiledLoopOp,
               VectorTransferOpInterface>(
             [&](auto op) { return getInplaceableOpResult(op, opOperand); })
-        // ExtractSliceOp is special, when bufferized inplace it just returns an
-        // alias to its operand. Its result is never inplaceable on its operand.
-        .Case([&](ExtractSliceOp op) { return OpResult(); })
+        // Some ops just return an alias to an operand when bufferized inplace.
+        // Such OpResults are never inplaceable on an OpOperand.
+        .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+            [] (auto op) { return OpResult(); })
         // CallOpInterface is special, it needs to wait for the callee to be
         // bufferized and needs to inspect the BufferAliasInfo object. It can't
         // make a proper determination by itself and needs to be conservative.
         .Case([&](CallOpInterface op) { return OpResult(); })
         // Other ops.
         .Default([&](Operation *op) { return OpResult(); });
   // clang-format on
 }
 
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place. Note that multiple OpOperands can may potentially alias with an
 /// OpResult. E.g.: std.select in the future.
 static SmallVector<OpOperand *> getAliasingOpOperand(OpResult result) {
   SmallVector<OpOperand *> r;
   // Unknown ops are handled conservatively and never bufferize in-place.
   if (!hasKnownBufferizationAliasingBehavior(result.getDefiningOp()))
     return SmallVector<OpOperand *>();
   TypeSwitch<Operation *>(result.getDefiningOp())
       .Case([&](tensor::CastOp op) { r.push_back(&op->getOpOperand(0)); })
-      .Case([&](ExtractSliceOp op) { r.push_back(&op->getOpOperand(0)); })
+      .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+          [&](auto op) { r.push_back(&op->getOpOperand(0)); })
       // In the case of scf::ForOp, this currently assumes the iter_args / yield
       // are 1-1. This may fail and is verified at the end.
       // TODO: update this.
       .Case([&](scf::ForOp op) {
         r.push_back(&op.getIterOpOperands()[result.getResultNumber()]);
       })
       .Case([&](InsertSliceOp op) { r.push_back(&op->getOpOperand(1)); })
       .Case([&](LinalgOp op) {
TypeSwitch<Operation *>(result.getDefiningOp())
         llvm_unreachable("unexpected defining op");
       });
   return r;
 }
 
 /// If the an ExtractSliceOp is bufferized in-place, the source operand will
 /// alias with the result.
 static OpResult getAliasingOpResult(ExtractSliceOp op, OpOperand &opOperand) {
-  if (op.source() == opOperand.get())
+  if (&op->getOpOperand(0) == &opOperand)
+    return op->getResult(0);
+  return OpResult();
+}
+
+/// If the a TensorExpandShapeOp is bufferized in-place, the source operand will
+/// alias with the result.
+static OpResult getAliasingOpResult(TensorExpandShapeOp op,
+                                    OpOperand &opOperand) {
+  if (&op->getOpOperand(0) == &opOperand)
+    return op->getResult(0);
+  return OpResult();
+}
+
+/// If the a TensorCollapseShapeOp is bufferized in-place, the source operand
+/// will alias with the result.
+static OpResult getAliasingOpResult(TensorCollapseShapeOp op,
+                                    OpOperand &opOperand) {
+  if (&op->getOpOperand(0) == &opOperand)
     return op->getResult(0);
   return OpResult();
 }
 
 /// Determine which OpResult will alias with `opOperand` if the op is bufferized
 /// in place. This is a superset of `getInplaceableOpResult`.
 /// TODO: in the future this may need to evolve towards a list of OpResult.
 static OpResult getAliasingOpResult(OpOperand &opOperand) {
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
-      // ExtractSliceOp is different: its result is not inplaceable on op.source
-      // but when bufferized inplace, the result is an aliasing subregion of
-      // op.source.
-      .Case(
-          [&](ExtractSliceOp op) { return getAliasingOpResult(op, opOperand); })
+      // Some ops are different: Their result is not inplaceable on an OpOperand
+      // but when bufferized inplace, their result is aliasing (a subregion of)
+      // an OpOperand.
+      .Case<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+          [&](auto op) { return getAliasingOpResult(op, opOperand); })
       // All other ops, return the result of `getInplaceableOpResult`.
       .Default(
           [&](Operation *op) { return getInplaceableOpResult(opOperand); });
 }
 
 // Predeclaration of function.
 static bool bufferizesToMemoryRead(OpOperand &opOperand);
 
 /// Return true if the given value is read by an op that bufferizes to a memory
 /// read. Also takes into account ops that create an alias but do not read by
 /// themselves (e.g., ExtractSliceOp).
 static bool isValueRead(Value value) {
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
-    // Skip over all ExtractSliceOps. These do not read by themselves but just
-    // add a new alias.
-    if (auto extractSliceOp =
-            dyn_cast<ExtractSliceOp>(uMaybeReading->getOwner()))
-      for (OpOperand &use : extractSliceOp.result().getUses())
+    // Skip over all ops that create an alias but do not read.
+    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+            uMaybeReading->getOwner()))
+      for (OpOperand &use : uMaybeReading->getOwner()->getResult(0).getUses())
         workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
       return true;
   }
 
   return false;
 }
 
 /// Return true if `opOperand` bufferizes to a memory read.
 static bool bufferizesToMemoryRead(OpOperand &opOperand) {
   // Unknown op that returns a tensor. The inplace analysis does not support
   // it. Conservatively return true.
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
     return true;
-  // ExtractSliceOp alone doesn't bufferize to a memory read, one of its uses
+  // Some ops alone do not bufferize to a memory read, but one of their uses
   // may.
-  if (isa<ExtractSliceOp>(opOperand.getOwner()))
+  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+          opOperand.getOwner()))
     return false;
   // scf::ForOp alone doesn't bufferize to a memory read, one of the uses of its
   // matching bbArg may.
   if (auto forOp = dyn_cast<scf::ForOp>(opOperand.getOwner()))
     return isValueRead(forOp.getRegionIterArgForOpOperand(opOperand));
   // TiledLoop alone doesn't bufferize to a memory read, one of the uses of its
   // matching bbArg may.
   if (auto tiledLoopOp = dyn_cast<TiledLoopOp>(opOperand.getOwner()))
static bool bufferizesToMemoryRead(OpOperand &opOperand) {
   return true;
 }
 
 /// Return true if `opOperand` bufferizes to a memory write.
 static bool bufferizesToMemoryWrite(OpOperand &opOperand) {
   // These terminators are not writes.
   if (isa<ReturnOp, linalg::YieldOp, scf::YieldOp>(opOperand.getOwner()))
     return false;
-  // ExtractSliceOp alone doesn't bufferize to a memory write, one of its uses
+  // Some ops alone do not bufferize to a memory write, but one of their uses
   // may.
-  if (isa<ExtractSliceOp>(opOperand.getOwner()))
+  if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+          opOperand.getOwner()))
     return false;
   // CallOpInterface alone doesn't bufferize to a memory write, one of the uses
   // of the matching bbArg may. It is the responsibility of the caller to
   // inspect bbArgs. In the absence of a BufferizationAliasInfo, we need to be
   // conservative.
   if (auto callOp = dyn_cast<CallOpInterface>(opOperand.getOwner()))
     return true;
   // Unknown op that returns a tensor. The inplace analysis does not support
tiledLoopOp->setAttr(
       b.getI32VectorAttr(
           {numLoops, numLoops, numLoops,
            static_cast<int>(oldInputs.size()) + numNewInputBuffers,
            static_cast<int>(oldOutputs.size()) + numNewOutputBuffers}));
 
   return success();
 }
 
+/// Helper function for TensorCollapseShapeOp and TensorExpandShapeOp
+/// bufferization. Retrieve reassociation indices from an attribute.
+static SmallVector<ReassociationIndices>
+attrToReassociationIndices(ArrayAttr reassocAttr) {
+  SmallVector<ReassociationIndices> reassociation;
+  for (Attribute attr : reassocAttr) {
+    ReassociationIndices indices;
+    auto elem = attr.dyn_cast<ArrayAttr>();
+    assert(elem && "unexpected reassociation attr format");
+    for (Attribute idxAttr : elem) {
+      auto idxIntAttr = idxAttr.dyn_cast<IntegerAttr>();
+      assert(idxIntAttr && "unexpected reassociation attr format");
+      indices.push_back(idxIntAttr.getInt());
+    }
+    reassociation.push_back(indices);
+  }
+  return reassociation;
+}
+
+static LogicalResult bufferize(OpBuilder &b, TensorCollapseShapeOp op,
+                               BlockAndValueMapping &bvm,
+                               BufferizationAliasInfo &aliasInfo) {
+  // Take a guard and get result buffer.
+  OpBuilder::InsertionGuard g(b);
+  Value resultBuffer = getResultBuffer(b, op->getResult(0), bvm, aliasInfo);
+  if (!resultBuffer)
+    return failure();
+
+  // Create memref::CollapseShapeOp.
+  b.setInsertionPoint(op);
+  auto memrefOp = b.create<memref::CollapseShapeOp>(
+      op.getLoc(), resultBuffer,
+      attrToReassociationIndices(op.reassociation()));
+  map(bvm, op.result(), memrefOp.result());
+  return success();
+}
+
+static LogicalResult bufferize(OpBuilder &b, TensorExpandShapeOp op,
+                               BlockAndValueMapping &bvm,
+                               BufferizationAliasInfo &aliasInfo) {
+  // Take a guard and get result buffer.
+  OpBuilder::InsertionGuard g(b);
+  Value resultBuffer = getResultBuffer(b, op->getResult(0), bvm, aliasInfo);
+  if (!resultBuffer)
+    return failure();
+
+  // Create memref::ExpandShapeOp.
+  b.setInsertionPoint(op);
+  auto memrefOp = b.create<memref::ExpandShapeOp>(
+      op.getLoc(), resultBuffer, attrToReassociationIndices(op.reassociation()),
+      op.getResultType().getShape());
+  map(bvm, op.result(), memrefOp.result());
+  return success();
+}
+
 /// Bufferize ExtractSliceOp to subview with optional alloc + copy depending on
 /// whether or not it is marked inplaceable.
 /// Note that `getInplaceableOpResult` on a ExtractSliceOp always returns null.
 /// As consequence a ExtractSliceOp always alloc + copy when taken in
 /// isolation.
 static LogicalResult bufferize(OpBuilder &b, ExtractSliceOp extractSliceOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
#endif // NDEBUG
   else
     aliasInfo.bufferizeInPlace(result, operand);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
+/// This analysis function is used for ops where the first OpOperand aliases
+/// with the first OpResult, without creating a read or write. There are a few
+/// ops besides ExtractSliceOp that have such semantics.
 ///
-/// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace.
-/// ===========================================================
+/// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace:
 ///
-/// When bufferized out of place, a ExtractSlice lowers to alloc + copy. This
+/// When bufferized out of place, an ExtractSliceOp lowers to alloc + copy. This
 /// cannot change the flow of information for either the source or the
 /// result buffers.
 ///
-/// When bufferized inplace, a ExtractSliceOp does not by itself create any read
-/// or write from memory. Instead, it has the effect of merging the alias sets
-/// of the source and the result buffers.
+/// When bufferized inplace, an ExtractSliceOp does not by itself create any
+/// read or write from memory. Instead, it has the effect of merging the alias
+/// sets of the source and the result buffers.
 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult
-bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
+bufferizableInPlaceAnalysisAliasOnlyOp(Operation *op,
                                        BufferizationAliasInfo &aliasInfo,
                                        const DominanceInfo &domInfo) {
-  return bufferizableInPlaceAnalysisImpl(extractSliceOp->getOpOperand(0),
-                                         extractSliceOp->getOpResult(0),
-                                         aliasInfo, domInfo);
+  return bufferizableInPlaceAnalysisImpl(
+      op->getOpOperand(0), op->getOpResult(0), aliasInfo, domInfo);
 }
 
 /// Determine if `operand` can be bufferized in-place with one of the op's
 /// results. If so, set InPlaceSpec::True on the result. Otherwise, set
 /// InPlaceSpec::False on the result.
 static LogicalResult
 bufferizableInPlaceAnalysis(OpOperand &operand,
                             BufferizationAliasInfo &aliasInfo,
LogicalResult mlir::linalg::inPlaceAnalysis(SmallVector<Operation *> &ops,
                                             BufferizationAliasInfo &aliasInfo,
                                             const DominanceInfo &domInfo) {
   // Walk ops in reverse for better interference analysis.
   for (Operation *op : reverse(ops)) {
     for (OpOperand &opOperand : op->getOpOperands())
       if (failed(bufferizableInPlaceAnalysis(opOperand, aliasInfo, domInfo)))
         return failure();
 
-    // Special logic to analyze ExtractSliceOp.
-    // Note that ExtractSliceOp analysis needs to be interleaved with other ops
-    // to properly capture aliases.
-    // Walk ExtractSliceOps in reverse for better clobbering analysis behavior:
-    // it is easier to detect clobbers of smaller slices before larger ones.
-    if (auto extractSliceOp = dyn_cast<ExtractSliceOp>(op))
+    // Special logic to analyze ops who's OpResults are not inplaceable on an
+    // OpOperand but may create an alias.
+    if (isa<ExtractSliceOp, TensorCollapseShapeOp, TensorExpandShapeOp>(op))
       if (failed(
-              bufferizableInPlaceAnalysis(extractSliceOp, aliasInfo, domInfo)))
+              bufferizableInPlaceAnalysisAliasOnlyOp(op, aliasInfo, domInfo)))
         return failure();
   }
 
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable.
 static LogicalResult
LogicalResult mlir::linalg::bufferizeOp(
   OpBuilder b(op->getContext());
   return TypeSwitch<Operation *, LogicalResult>(op)
       // Skip BufferCast and TensorLoad ops.
       .Case<memref::BufferCastOp, memref::TensorLoadOp>(
           [&](auto) { return success(); })
       .Case<tensor::CastOp, tensor::DimOp, ExtractSliceOp, scf::ForOp,
             InitTensorOp, InsertSliceOp, tensor::ExtractOp, LinalgOp, ReturnOp,
             TiledLoopOp, VectorTransferOpInterface, linalg::YieldOp,
-            scf::YieldOp>([&](auto op) {
+            scf::YieldOp, TensorCollapseShapeOp, TensorExpandShapeOp>(
+          [&](auto op) {
             LDBG("Begin bufferize:\n" << op << '\n');
             return bufferize(b, op, bvm, aliasInfo);
           })
       .Case([&](CallOpInterface op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         if (!bufferizedFunctionTypes)
           llvm_unreachable(
               "null bufferizedFunctionTypes when bufferizing CallOpInterface");
         return bufferize(b, op, bvm, aliasInfo, *bufferizedFunctionTypes);
       })
       .Case([&](arith::ConstantOp op) {
auto extractOp = b.create<tensor::ExtractSliceOp>(
         initTensor.getLoc(), insertOp.dest(), insertOp.getMixedOffsets(),
         insertOp.getMixedSizes(), insertOp.getMixedStrides());
     // Uses of the InitTensorOp are replaced here, but the op is not deleted.
     // InitTensorOps without uses are ignored by the bufferization.
     initTensor.replaceAllUsesWith(extractOp.result());
     aliasInfo.createAliasInfoEntry(extractOp.result());
 
     // Run analysis on the ExtractSliceOp.
-    if (failed(bufferizableInPlaceAnalysis(extractOp, aliasInfo, domInfo)))
+    if (failed(bufferizableInPlaceAnalysisAliasOnlyOp(extractOp, aliasInfo,
+                                                      domInfo)))
       return WalkResult::interrupt();
 
     // Advance to the next operation.
     return WalkResult::advance();
   });
 
   return failure(status.wasInterrupted());
 }

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @buffer_forwarding_no_conflict(
 
   // Self-copy canonicalizes away later.
   // CHECK: linalg.copy(%[[T_SUBVIEW]], %[[T_SUBVIEW]])
   %r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
 
   return %r1: tensor<?xf32>
 }
 
+// -----
+
+//       CHECK: #[[$map0:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+//       CHECK: #[[$map1:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+// CHECK-LABEL: func @test_tensor_collapse_shape(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<?x?x?xf32, #[[$map0]]>
+func @test_tensor_collapse_shape(
+    %A : tensor<?x?x?xf32> {linalg.inplaceable = true}) -> vector<5x6xf32> {
+  //      CHECK: memref.collapse_shape %arg0 [
+  // CHECK-SAME:     [0, 1], [2]] : memref<?x?x?xf32, #[[$map0]]> into memref<?x?xf32, #[[$map1]]>
+  %s = linalg.tensor_collapse_shape %A[[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>
+  %cst = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %v = vector.transfer_read %s[%c0, %c0], %cst : tensor<?x?xf32>, vector<5x6xf32>
+  return %v : vector<5x6xf32>
+}
+
+// -----
+
+//       CHECK: #[[$map2:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+//       CHECK: #[[$map3:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
+// CHECK-LABEL: func @test_tensor_expand_shape(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<50x?xf32, #[[$map2]]>
+func @test_tensor_expand_shape(
+  %A : tensor<50x?xf32> {linalg.inplaceable = true}) -> vector<5x6x7xf32> {
+  //      CHECK: memref.expand_shape %arg0 [
+  // CHECK-SAME:     [0, 1], [2]] : memref<50x?xf32, #[[$map2]]> into memref<10x5x?xf32, #[[$map3]]>
+  %s = linalg.tensor_expand_shape %A[[0, 1], [2]] : tensor<50x?xf32> into tensor<10x5x?xf32>
+  %cst = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %v = vector.transfer_read %s[%c0, %c0, %c0], %cst : tensor<10x5x?xf32>, vector<5x6x7xf32>
+  return %v : vector<5x6x7xf32>
+}