[mlir][Linalg] Properly order extract_slice traversal in comprehensive bufferization

[nicolasvasilache]

Summary

This revision fixes the traversal order of extract_slice during the inplace analysis.
It was previously thought that such ops could be analyzed at the very end.
This is unfortunately not true as the AliasInfo for dependents of these ops need to be updated.

This change allows the aliases introduced by the bufferization of extract_slice to be properly propagated.

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

 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 
-#define DEBUG_TYPE "comprehensive-func-bufferize"
+#define DEBUG_TYPE "comprehensive-module-bufferize"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 // Forward declarations.
-static std::string printOperationInfo(Operation *);
-static std::string printValueInfo(Value);
+static std::string printOperationInfo(Operation *, bool prefix = true);
+static std::string printValueInfo(Value, bool prefix = true);
 
 //===----------------------------------------------------------------------===//
 // Generic helpers.
 //===----------------------------------------------------------------------===//
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// Return the FuncOp called by `callOp`.
static void printTensorOrBufferInfo(std::string prefix, Value value,
   value.printAsOperand(os, state);
   os << " : " << value.getType();
   if (getInPlace(value) == InPlaceSpec::None)
     return;
   os << " [InPlace=" << stringify(getInPlace(value)) << "]";
 }
 
 /// Print the operation name and bufferization information.
-static std::string printOperationInfo(Operation *op) {
+static std::string printOperationInfo(Operation *op, bool prefix) {
   std::string result;
   llvm::raw_string_ostream os(result);
   AsmState state(op->getParentOfType<mlir::FuncOp>());
-  os << op->getName();
+  StringRef tab = prefix ? "\n[" DEBUG_TYPE "]\t" : "";
+  os << tab << op->getName();
   SmallVector<Value> shapedOperands;
   for (OpOperand &opOperand : op->getOpOperands()) {
     std::string prefix =
-        llvm::formatv("\n\t-> #{0} ", opOperand.getOperandNumber());
+        llvm::formatv("{0}  -> #{1} ", tab, opOperand.getOperandNumber());
     printTensorOrBufferInfo(prefix, opOperand.get(), state, os);
   }
   for (OpResult opResult : op->getOpResults()) {
     std::string prefix =
-        llvm::formatv("\n\t<- #{0} ", opResult.getResultNumber());
+        llvm::formatv("{0}  <- #{1} ", tab, opResult.getResultNumber());
     printTensorOrBufferInfo(prefix, opResult, state, os);
   }
   return result;
 }
 
 /// Print the bufferization information for the defining op or block argument.
-static std::string printValueInfo(Value value) {
+static std::string printValueInfo(Value value, bool prefix) {
   auto *op = value.getDefiningOp();
   if (op)
-    return printOperationInfo(op);
+    return printOperationInfo(op, prefix);
   // Print the block argument bufferization information.
   std::string result;
   llvm::raw_string_ostream os(result);
   AsmState state(value.getParentRegion()->getParentOfType<mlir::FuncOp>());
   os << value;
   printTensorOrBufferInfo("\n\t - ", value, state, os);
   return result;
 }
return TypeSwitch<Operation *, OpOperand *>(result.getDefiningOp())
       .Case([&](ConstantOp op) { return &op->getOpOperand(0); })
       .Case([&](ExtractSliceOp op) { return &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) {
         return &op.getIterOpOperands()[result.getResultNumber()];
       })
+      .Case([&](InitTensorOp op) { return nullptr; })
       .Case([&](InsertSliceOp op) { return &op->getOpOperand(1); })
       .Case([&](LinalgOp op) {
         return op.getOutputTensorOperands()[result.getResultNumber()];
       })
       .Case([&](TiledLoopOp op) {
         // TODO: TiledLoopOp helper method to avoid leaking impl details.
         return &op->getOpOperand(op.getNumControlOperands() +
                                  op.getNumInputs() + result.getResultNumber());

 /// bufferization aliasing behavior, which indicates that the op must allocate
 /// all of its tensor results.
 /// TODO: in the future this may need to evolve towards a list of OpResult.
 static Optional<OpResult> getAliasingOpResult(OpOperand &opOperand) {
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
     return None;
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
       // These terminators legitimately have no result.
-      .Case<ReturnOp, linalg::YieldOp, scf::YieldOp>(
+      .Case<ReturnOp, linalg::InitTensorOp, linalg::YieldOp, scf::YieldOp>(
           [&](auto op) { return OpResult(); })
       // DimOp has no tensor result.
       .Case<tensor::DimOp>([&](auto op) { return None; })
       // ConstantOp is never inplaceable.
       .Case([&](ConstantOp op) { return op->getResult(0); })
       // ExtractSliceOp is different: its result is not inplaceable on op.source
       // but when bufferized inplace, the result is an aliasing subregion of
       // op.source.
public:
   /// equivalent ExtractSliceOp.
   bool isSourceEquivalentToAMatchingInplaceExtractSliceOp(
       InsertSliceOp insertSliceOp) const;
 
   /// Apply `fun` to all the members of the equivalence class of `v`.
   void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const;
 
   /// Print to `os`.
-  void print(raw_ostream &os) const;
+  void printAliases(raw_ostream &os) const;
+  void printEquivalences(raw_ostream &os) const;
 
   /// Print to `errs()`.
-  void dump() const { print(llvm::errs()); }
+  void dumpAliases() const { printAliases(llvm::errs()); }
+  void dumpEquivalences() const { printEquivalences(llvm::errs()); }
 
 private:
   /// Check that aliasInfo for `v` exists and return a reference to it.
   DenseSet<Value> &getAliasInfoRef(Value v);
 
   const DenseSet<Value> &getAliasInfoRef(Value v) const {
     return const_cast<BufferizationAliasInfo *>(this)->getAliasInfoRef(v);
   }

 /// Set the inPlace bufferization spec to true.
 /// Merge result's and operand's aliasing sets and iterates to a fixed point.
 void BufferizationAliasInfo::bufferizeInPlace(OpResult result,
                                               OpOperand &operand,
                                               BufferRelation bufferRelation) {
   setInPlaceOpResult(result, InPlaceSpec::True);
   if (mergeAliases(result, operand.get()))
     mergeAliasesToFixedPoint();
+  // Dump the updated alias analysis.
+  LLVM_DEBUG(dumpAliases());
   if (bufferRelation == BufferRelation::Equivalent)
     equivalentInfo.unionSets(result, operand.get());
-  // Dump the updated analysis.
-  LLVM_DEBUG(dump());
+  // Dump the updated equivalence analysis.
+  LLVM_DEBUG(dumpEquivalences());
 }
 
 /// Set the inPlace bufferization spec to false.
 void BufferizationAliasInfo::bufferizeOutOfPlace(OpResult result) {
   setInPlaceOpResult(result, InPlaceSpec::False);
 }
 
 /// Return true if it is possible to find an inplace write W among the uses of
bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
     OpResult result, const DominanceInfo &domInfo) const {
   Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(result);
   if (!maybeAliasingOperand)
     return false;
 
   Operation *opToBufferize = result.getDefiningOp();
   Value root = (*maybeAliasingOperand)->get();
   LDBG("----Start wouldCreateReadAfterWriteInterference\n");
-  LDBG("--------rootValue: " << printValueInfo(root) << "\n");
+  LDBG("--------aliasing rootValue: " << printValueInfo(root) << "\n");
 
   // Collect:
   //   1. all the inplace write uses of some alias of `root`.
   //   2. all the write uses that belong to `opToBufferize`.
   // opToBufferize is not yet inplace, we want to determine if it can be inplace
   // so we also consider all its write uses, not just the inplace ones.
   DenseSet<OpOperand *> usesWrite;
   for (Value vWrite : getAliasInfoRef(root)) {
for (OpOperand *uWrite : usesWrite) {
         continue;
       // If aliasingReadOp properly dominates aliasingWriteOp, the read cannot
       // be affected by the write: there is no interference.
       if (domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp))
         continue;
       // At this point, aliasingWriteOp properly dominates aliasingReadOp or
       // there is no clear dominance and we need to be conservative.
       LDBG("---->found RaW interference\n");
-      LDBG("     Interfering  read -> #" << uRead->getOperandNumber() << ":\n"
+      LDBG("     Interfering  read -> #" << uRead->getOperandNumber() << ":"
                                          << printOperationInfo(aliasingReadOp)
                                          << '\n');
-      LDBG("     Interfering write -> #" << uWrite->getOperandNumber() << ":\n"
+      LDBG("     Interfering write -> #" << uWrite->getOperandNumber() << ":"
                                          << printOperationInfo(aliasingWriteOp)
                                          << '\n');
       LDBG("---->opportunity to clobber RaW interference\n");
       if (isClobberedWriteBeforeRead(opToBufferize, *uRead, *uWrite, domInfo)) {
         LDBG("---->clobbered! -> skip\n");
         continue;
       }
 
void BufferizationAliasInfo::applyOnEquivalenceClass(
     Value v, function_ref<void(Value)> fun) const {
   auto leaderIt = equivalentInfo.findLeader(v);
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
     fun(mit->v);
   }
 }
 
-void BufferizationAliasInfo::print(raw_ostream &os) const {
+void BufferizationAliasInfo::printAliases(raw_ostream &os) const {
   os << "\n/========================== AliasInfo "
         "==========================\n";
   for (auto it : aliasInfo) {
-    os << "|\n| -- source: " << printValueInfo(it.getFirst()) << '\n';
+    os << "|\n| -- source: " << printValueInfo(it.getFirst(), /*prefix=*/false)
+       << '\n';
     for (auto v : it.getSecond())
-      os << "| ---- target: " << printValueInfo(v) << '\n';
+      os << "| ---- target: " << printValueInfo(v, /*prefix=*/false) << '\n';
   }
   os << "|\n\\====================== End AliasInfo "
         "======================\n\n";
+}
+
+void BufferizationAliasInfo::printEquivalences(raw_ostream &os) const {
   os << "\n/********************* Equivalent Buffers *********************\n";
   for (auto it = equivalentInfo.begin(), eit = equivalentInfo.end(); it != eit;
        ++it) {
     if (!it->isLeader())
       continue;
     Value leader = it->getData();
-    os << "|\n| -- leader: " << printValueInfo(leader) << '\n';
+    os << "|\n| -- leader: " << printValueInfo(leader, /*prefix=*/false)
+       << '\n';
     for (auto mit = equivalentInfo.member_begin(it),
               meit = equivalentInfo.member_end();
          mit != meit; ++mit) {
       Value v = static_cast<Value>(*mit);
-      os << "| ---- equivalent member: " << printValueInfo(v) << '\n';
+      os << "| ---- equivalent member: " << printValueInfo(v, /*prefix=*/false)
+         << '\n';
     }
   }
   os << "|\n\\***************** End Equivalent Buffers *****************\n\n";
 }
 
 DenseSet<Value> &BufferizationAliasInfo::getAliasInfoRef(Value v) {
   auto it = aliasInfo.find(v);
   if (it == aliasInfo.end())
bool BufferizationAliasInfo::existsInterleavedValueClobber(
   Operation *aliasingWriteOp = aliasingWrite.getOwner();
   assert(!domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp) &&
          "Unexpected aliasingReadOp properly dominates aliasingWriteOp");
 
   for (Value valueToClobber : {aliasingRead.get(), aliasingWrite.get()}) {
     auto leaderIt = equivalentInfo.findLeader(valueToClobber);
     for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
          ++mit) {
-      /// Note: the "would write to memory after bufferization" condition is
-      /// verified by `candidateOp` since it would produce a value that
-      /// bufferizes to an equivalent buffer.
       Operation *candidateOp = mit->v.getDefiningOp();
       if (!candidateOp)
         continue;
+      auto maybeAliasingOperand = getAliasingOpOperand(mit->v.cast<OpResult>());
+      if (!maybeAliasingOperand || !*maybeAliasingOperand ||
+          !bufferizesToMemoryWrite(**maybeAliasingOperand))
+        continue;
       LDBG("---->clobbering candidate: " << printOperationInfo(candidateOp)
                                          << '\n');
       if (domInfo.properlyDominates(aliasingWriteOp, candidateOp) &&
           domInfo.properlyDominates(candidateOp, aliasingReadOp))
         return true;
     }
   }
   return false;
else
     aliasInfo.bufferizeInPlace(
         result, operand, BufferizationAliasInfo::BufferRelation::Equivalent);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
-/// Analyze the `funcOp` body to determine which OpResults are inplaceable.
+/// Analyze the `funcOp` body to determine which OpResults are inplaceable:
+///   1. First, analyze InsertSliceOp greedily: we almost never want to
+///      bufferize the tensor "inserted into" to become out-of-place.
+///   2. Walk the other ops in reverse. This is a good starter heuristic.
+///      ExtractSliceOps are interleaved with other ops in traversal order.
+///
 static LogicalResult
 inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
                           const DominanceInfo &domInfo) {
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
   assert(funcOp && funcOp->getNumRegions() > 0 && !funcOp.body().empty() &&
          "expected a funcOp definition with a body");
 
   // Collect ops so we can build our own traversal.
-  SmallVector<ExtractSliceOp> extractSliceOps;
+  SmallVector<Operation *> otherOps;
   SmallVector<InsertSliceOp> insertSliceOps;
-  SmallVector<Operation *> nonSliceOps;
   funcOp.walk([&](Operation *op) {
-    if (auto extractSliceOp = dyn_cast<ExtractSliceOp>(op))
-      return extractSliceOps.push_back(extractSliceOp);
     if (auto insertSliceOp = dyn_cast<InsertSliceOp>(op))
       return insertSliceOps.push_back(insertSliceOp);
     // No tensors => no buffers.
     if (none_of(op->getOperandTypes(), isaTensor) &&
         none_of(op->getResultTypes(), isaTensor))
       return;
-    nonSliceOps.push_back(op);
+    otherOps.push_back(op);
   });
 
-  // Bufferize InsertSliceOp greedily: we almost never want to bufferize
+  // First, analyze InsertSliceOp greedily: we almost never want to bufferize
   // the tensor "inserted into" to become out-of-place. This implementation
   // does not distinguish between different InsertSliceOp. If we want
   // finer-grained behavior, we could order the InsertSliceOp with some metric.
-  // Walk InsertSliceOp in reverse for better interference behavior.
   for (InsertSliceOp insertSliceOp : reverse(insertSliceOps)) {
     OpOperand &destOpOperand = insertSliceOp->getOpOperand(1);
     if (failed(bufferizableInPlaceAnalysis(
             destOpOperand, getInplaceableOpResult(destOpOperand), aliasInfo,
             domInfo)))
       return failure();
   }
 
-  // Analyze all ops that return a tensors, except ExtractSliceOp and
-  // InsertSliceOp which are handled separately.
-  // Walk other ops in reverse for better interference behavior.
-  for (Operation *op : reverse(nonSliceOps))
-    for (OpOperand &opOperand : op->getOpOperands())
+  // Walk ops in reverse for better interference analysis.
+  for (Operation *op : reverse(otherOps)) {
+    for (OpOperand &opOperand : op->getOpOperands()) {
       if (OpResult result = getInplaceableOpResult(opOperand))
         if (result.getType().isa<TensorType>() &&
             failed(bufferizableInPlaceAnalysis(opOperand, result, aliasInfo,
                                                domInfo)))
           return failure();
-
-  // Finally, bufferize ExtractSliceOp.
-  // Walk ExtractSliceOps in reverse for better clobbering behavior: it is
-  // easier to detect clobbers of smaller slices before larger ones.
-  for (ExtractSliceOp extractSliceOp : reverse(extractSliceOps))
-    if (failed(bufferizableInPlaceAnalysis(extractSliceOp, aliasInfo, domInfo)))
+    }
+    // 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)) {
+      if (failed(
+              bufferizableInPlaceAnalysis(extractSliceOp, aliasInfo, domInfo)))
         return failure();
+      continue;
+    }
+  }
 
   LDBG("End InPlaceAnalysisFuncOpInternals:\n" << funcOp << '\n');
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for functions.

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

 
 // CHECK-LABEL: func @extract_slice_to_linalg_write_use
 func @extract_slice_to_linalg_write_use(
     %A : tensor<4x4xf32>,
     %B : tensor<?x?xf32>,
     %C : tensor<?x?xf32> {linalg.inplaceable = true})
   ->  (tensor<4x4xf32>, tensor<4x4xf32>)
 {
-  // Step 3. %sB forward propagates to a write in %D but it is not inplace.
+  // Step 4. %sB forward propagates to a write in %D but it is not inplace.
   // So this is only ever read and can bufferize inplace.
   //     CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %sB = tensor.extract_slice %B[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
 
-  // Step 2. %sB has a read interference in %E, it does not bufferize inplace.
+  // Step 3. %sB has a read interference in %E, it does not bufferize inplace.
   //     CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
   %D = linalg.matmul  ins(%B, %C: tensor<?x?xf32>, tensor<?x?xf32>)
                      outs(%sB: tensor<4x4xf32>)
     -> tensor<4x4xf32>
 
-  // Step 4. %sC forward propagates to an inplace write in %E.
+  // Step 2. %sC forward propagates to an inplace write in %E.
   // %sC backward propagates to %C which is inplaceable.
   // As a consequence this is bufferized inplace.
   //     CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %sC = tensor.extract_slice %C[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
 
   // Step 1. %sC backprops to the tensor.extract_slice producer which is not
   // considered an interference. This bufferizes inplace.

 {
   // Step 4. %sB forward propagates to an inplace write in %D.
   // %sB backward propagates to %B which is not inplaceable.
   // As a consequence this is bufferized out of place.
   //     CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
   %sB = tensor.extract_slice %B[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
 
-  // Step 1. %sB backprops to the tensor.extract_slice producer which is not
+  // Step 3. %sB backprops to the tensor.extract_slice producer which is not
   // considered an interference. This bufferizes inplace.
   //     CHECK: linalg.matmul
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %D = linalg.matmul  ins(%B, %C: tensor<?x?xf32>, tensor<?x?xf32>)
                      outs(%sB: tensor<4x4xf32>)
     -> tensor<4x4xf32>
 
-  // Step 3. %sC forward propagates to an inplace write in %E.
+  // Step 2. %sC forward propagates to an inplace write in %E.
   // %sC backward propagates to %C which is inplaceable.
   // As a consequence this is bufferized inplace.
   //     CHECK: tensor.extract_slice
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %sC = tensor.extract_slice %C[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
 
   // Step 1. %sC backprops to the tensor.extract_slice producer which is not
   // considered an interference. This bufferizes inplace.

 func private @some_use(tensor<?xf32>) -> ()
 
 // CHECK-LABEL: func @scf_for_deps
 func @scf_for_deps(%A : tensor<?xf32> {linalg.inplaceable = true},
                    %B : tensor<?xf32> {linalg.inplaceable = true},
                    %lb : index, %ub : index, %step : index)
   -> (tensor<?xf32>, tensor<?xf32>)
 {
-  // %r0 must be out of place because one use of %t in the subsequent production 
+  // %r0 must be out of place because one use of %t in the subsequent production
   // of %r1 is read.
   //      CHECK: scf.for
   // CHECK-NEXT: call
   // CHECK-NEXT: scf.yield
   // CHECK-NEXT: {__inplace_results_attr__ = ["false"]}
   %r0 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
     call @some_use(%t) : (tensor<?xf32>) -> ()
     scf.yield %t : tensor<?xf32>
   }
 
   // %r1 bufferizes inplace fine.
   //      CHECK: scf.for
   // CHECK-NEXT: call
   // CHECK-NEXT: scf.yield
   // CHECK-NEXT: {__inplace_results_attr__ = ["true"]}
   %r1 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
     call @some_use(%t) : (tensor<?xf32>) -> ()
     scf.yield %t : tensor<?xf32>
   }
 
-  // %r2 must be out of place because one use of %t in the subsequent production 
+  // %r2 must be out of place because one use of %t in the subsequent production
   // of %r3 is read.
   //      CHECK: linalg.tiled_loop
   // CHECK-NEXT: call
   // CHECK-NEXT: linalg.yield
   // CHECK-NEXT: {__inplace_results_attr__ = ["false"]}
   %r2 = linalg.tiled_loop (%i) = (%lb) to (%ub) step (%step)
         ins()
         outs(%t = %B: tensor<?xf32>) {
func @read_dependence_through_scf_and_call(
   //     CHECK: fill
   // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
   %B2 = linalg.fill(%f2, %I2) : f32, tensor<64xf32> -> tensor<64xf32>
 
   call @bar(%A2) : (tensor<64xf32>) -> ()
   call @bar(%B2) : (tensor<64xf32>) -> ()
   return
 }
+
+//===----------------------------------------------------------------------===//
+// Transitive cases through extract_slice.
+//===----------------------------------------------------------------------===//
+
+builtin.func @matmul_on_tensors(
+    %arg0: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
+    %arg1: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
+    %arg2: tensor<256x256xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = true})
+    -> tensor<256x256xf32>
+{
+  %c0 = constant 0 : index
+  %cst_0 = constant 0.000000e+00 : f32
+  %cst_1 = constant 1.000000e+00 : f32
+
+  %7 = linalg.init_tensor [256, 256] : tensor<256x256xf32>
+
+  //      CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  //      CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]}
+  %8 = linalg.fill(%cst_0, %7) : f32, tensor<256x256xf32> -> tensor<256x256xf32>
+  %11 = linalg.fill(%cst_1, %7) : f32, tensor<256x256xf32> -> tensor<256x256xf32>
+
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  //      CHECK: linalg.matmul
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %sA = tensor.extract_slice %8[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32>
+  %sB = tensor.extract_slice %11[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32>
+  %r = linalg.matmul
+         ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>)
+        outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
+
+  return %r : tensor<256x256xf32>
+}
+
+// -----
+
+builtin.func @matmul_on_tensors(
+    %arg0: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
+    %arg1: tensor<518x518xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = false},
+    %arg2: tensor<256x256xf32> {linalg.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, linalg.inplaceable = true})
+    -> tensor<256x256xf32>
+{
+  %c0 = constant 0 : index
+  %cst_0 = constant 0.000000e+00 : f32
+  %cst_1 = constant 1.000000e+00 : f32
+
+  %7 = linalg.init_tensor [256, 256] : tensor<256x256xf32>
+
+  //     CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  //      CHECK: vector.transfer_write
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]
+  %8 = linalg.fill(%cst_0, %7) : f32, tensor<256x256xf32> -> tensor<256x256xf32>
+  %9 = vector.transfer_read %arg0[%c0, %c0], %cst_0 {in_bounds = [false, true]} : tensor<518x518xf32>, vector<256x256xf32>
+  %10 = vector.transfer_write %9, %8[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, tensor<256x256xf32>
+
+  //      CHECK: linalg.fill
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  //      CHECK: vector.transfer_write
+  // CHECK-SAME: {__inplace_results_attr__ = ["false"]
+  %11 = linalg.fill(%cst_1, %7) : f32, tensor<256x256xf32> -> tensor<256x256xf32>
+  %12 = vector.transfer_read %arg1[%c0, %c0], %cst_0 {in_bounds = [false, true]} : tensor<518x518xf32>, vector<256x256xf32>
+  %13 = vector.transfer_write %12, %11[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, tensor<256x256xf32>
+
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  //      CHECK: tensor.extract_slice
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  //      CHECK: linalg.matmul
+  // CHECK-SAME: {__inplace_results_attr__ = ["true"]}
+  %sA = tensor.extract_slice %10[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32>
+  %sB = tensor.extract_slice %13[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32>
+  %r = linalg.matmul
+         ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>)
+        outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
+
+  return %r : tensor<256x256xf32>
+}