[mlir][linalg][bufferize][NFC] Remove buffer equivalence from bufferize

[springerm]

Summary

Remove all function calls related to buffer equivalence from bufferize implementations.

Depends On D114926

[springerm]

update

[springerm]

update

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h

 namespace mlir {
 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
+/// Equivalence analysis for scf.for. Raise an error if iter_args are not
+/// equivalent to their corresponding loop yield values.
+struct AssertDestinationPassingStyle : public PostAnalysisStep {
+  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+                    SmallVector<Operation *> &newOps) override;
+};
+
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace scf_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_SCF_INTERFACE_IMPL_H

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

LogicalResult bufferize(Operation *op, OpBuilder &b,
     if (!moduleOp)
       return constantOp.emitError(
           "cannot bufferize constants not within builtin.module op");
 
     GlobalCreator globalCreator(moduleOp);
     auto globalMemref = globalCreator.getGlobalFor(constantOp);
     Value memref = b.create<memref::GetGlobalOp>(
         constantOp.getLoc(), globalMemref.type(), globalMemref.getName());
-    state.aliasInfo.insertNewBufferEquivalence(memref, constantOp.getResult());
     state.mapBuffer(constantOp, memref);
 
     return success();
   }
 
   bool isWritable(Operation *op, Value value) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
     assert(value.isa<OpResult>());

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

 
 /// Specify that the value is known to bufferize to writable memory.
 void BufferizationAliasInfo::setBufferizesToWritableMemory(Value v) {
   bufferizeToWritableMemory.insert(v);
 }
 
 /// Return `true` if a value was marked as in-place bufferized.
 bool BufferizationAliasInfo::isInPlace(OpResult opResult) const {
-  bool inplace = inplaceBufferized.contains(opResult);
-#ifndef NDEBUG
-  if (inplace) {
-    auto bufferizableOp =
-        dyn_cast<BufferizableOpInterface>(opResult.getDefiningOp());
-    assert(bufferizableOp &&
-           "expected that in-place bufferized op is bufferizable");
-    SmallVector<OpOperand *> operands =
-        bufferizableOp.getAliasingOpOperand(opResult);
-    for (OpOperand *operand : operands)
-      assert(areAliasingBufferizedValues(operand->get(), opResult) &&
-             "expected that in-place bufferized OpResult aliases with "
-             "aliasing OpOperand");
-  }
-#endif // NDEBUG
-  return inplace;
+  return inplaceBufferized.contains(opResult);
 }
 
 /// Set the inPlace bufferization spec to true.
 void BufferizationAliasInfo::bufferizeInPlace(OpResult result,
                                               OpOperand &operand) {
   LLVM_DEBUG(llvm::dbgs() << "bufferizeInPlace: ");
   LLVM_DEBUG(result.print(llvm::dbgs()));
 
Value mlir::linalg::comprehensive_bufferize::BufferizationState::
   Optional<Value> allocated =
       options.allocationFns->allocationFn(b, loc, allocMemRefType, dynShape);
   // TODO: For now just assert the value is returned. Eventually need to
   // error-propagate.
   assert(allocated && "allocation failed");
   Value casted = allocated.getValue();
   if (memRefType && memRefType != allocMemRefType) {
     casted = b.create<memref::CastOp>(loc, memRefType, allocated.getValue());
-    aliasInfo.insertNewBufferEquivalence(casted, allocated.getValue());
   }
 
   // 2. Create memory deallocation.
   b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
   options.allocationFns->deallocationFn(b, loc, allocated.getValue());
   return casted;
 }
 

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

for (Value oldOutputTensor : oldOutputs) {
 
       const OpResult &opResult = tiledLoopOp->getResult(resultIndex);
       OpOperand &yieldOperand = yieldOp->getOpOperand(resultIndex);
       Value resultBuffer = state.getResultBuffer(opResult);
       if (!resultBuffer)
         return failure();
 
       // Insert mapping and aliasing info.
-      state.aliasInfo.createAliasInfoEntry(resultBuffer);
-      state.aliasInfo.insertNewBufferEquivalence(opResult, resultBuffer);
       state.mapBuffer(opResult, resultBuffer);
 
       // Insert new operand and bbArg.
       tiledLoopOp->insertOperands(nextOutputOperandIndex, resultBuffer);
       BlockArgument newBufferBBArg =
           body->insertArgument(nextOutputBBArgIndex, resultBuffer.getType());
       BlockArgument oldTensorBBArg = body->getArgument(oldOutputBBArgIndex);
       // Insert mapping and aliasing info.
-      state.aliasInfo.createAliasInfoEntry(newBufferBBArg);
-      state.aliasInfo.insertNewBufferEquivalence(oldTensorBBArg,
-                                                 newBufferBBArg);
       state.mapBuffer(oldTensorBBArg, newBufferBBArg);
 
       // Set operand of `linalg.yield` to the bbArg so it just canonicalizes
       // away later.
       yieldOperand.set(oldTensorBBArg);
 
       // Increment indices.
       ++numNewOutputBuffers;
for (Value oldInputTensor : oldInputs) {
 
       // Insert new operand and bbArg.
       tiledLoopOp->insertOperands(nextInputOperandIndex, inputBuffer);
       BlockArgument newBufferBBArg =
           body->insertArgument(nextInputBBArgIndex, inputBuffer.getType());
       BlockArgument oldTensorBBArg = body->getArgument(oldInputBBArgIndex);
 
       // Insert mapping and aliasing info.
-      state.aliasInfo.createAliasInfoEntry(newBufferBBArg);
-      state.aliasInfo.insertNewBufferEquivalence(oldTensorBBArg,
-                                                 newBufferBBArg);
       state.mapBuffer(oldTensorBBArg, newBufferBBArg);
 
       // Increment indices.
       ++numNewInputBuffers;
       ++oldInputBBArgIndex;
       ++nextInputBBArgIndex;
       ++nextInputOperandIndex;
     }

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

 ///
 /// Note: Returning a memref currently fails bufferization. If such memrefs
 /// originate from an op with an Alloc effect, they could be hoisted in the
 /// future.
 static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
                                              BufferizationState &state) {
   LLVM_DEBUG(DBGS() << "Begin bufferizeFuncOpBoundary:\n" << funcOp << "\n");
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
-  BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   // If nothing to do then we are done.
   if (!llvm::any_of(funcOp.getType().getInputs(), isaTensor) &&
       !llvm::any_of(funcOp.getType().getResults(), isaTensor))
     return success();
 
   // Get the bufferized FunctionType for funcOp or construct it if not yet
   // available.
for (unsigned idx = 0; idx < numArgs; ++idx) {
     b.setInsertionPointToStart(&frontBlock);
     // Replace all uses of bbArg through a ToMemRefOp by a memref::CastOp.
     for (auto &use : llvm::make_early_inc_range(bbArg.getUses())) {
       if (auto toMemrefOp =
           dyn_cast<bufferization::ToMemrefOp>(use.getOwner())) {
         auto castOp = b.create<memref::CastOp>(
             funcOp.getLoc(), toMemrefOp.memref().getType(), memref);
         toMemrefOp.memref().replaceAllUsesWith(castOp);
-        aliasInfo.insertNewBufferEquivalence(castOp.dest(),
-                                             toMemrefOp.memref());
       }
     }
     // Replace all remaining uses by a to_tensor.
     if (!bbArg.use_empty()) {
       auto toTensorOp =
           b.create<bufferization::ToTensorOp>(funcOp.getLoc(), memref);
-      aliasInfo.insertNewBufferEquivalence(toTensorOp, bbArg);
       bbArg.replaceAllUsesWith(toTensorOp);
     }
     frontBlock.eraseArgument(0);
     // TODO: add support to erase aliasInfo entries if deemed necessary.
   }
 
   // 4. Rewrite the FuncOp type to buffer form.
   funcOp.setType(bufferizedFuncType);
if (funcOp.body().empty()) {
                 returnOperand.getOperandNumber())) {
           int64_t idx =
               moduleState
                   .equivalentFuncArgs[funcOp][returnOperand.getOperandNumber()];
           Value oldRes = callOp->getResult(returnOperand.getOperandNumber());
           Value buffer = state.lookupBuffer(callOp->getOperand(idx));
           // Add CallOp operand/result equivalence: this is interprocedural
           // info.
-          state.aliasInfo.insertNewBufferEquivalence(oldRes, buffer);
           state.mapBuffer(oldRes, buffer);
           // Add a ToTensorOp to kill all uses of the CallOp return.
           // Replace all uses of the CallOp results so we can erase the CallOp.
           // This ToTensorOp must fold/DCE away or bufferization should be
           // considered failed.
           Value toTensorOp =
               b.create<bufferization::ToTensorOp>(callOp.getLoc(), buffer);
           oldRes.replaceAllUsesWith(toTensorOp);
           // Add new op equivalence info.
-          state.aliasInfo.insertNewBufferEquivalence(toTensorOp, buffer);
           state.mapBuffer(toTensorOp, buffer);
           continue;
         }
 
         resultTypes.push_back(returnType);
       }
     }
 
for (OpOperand &opOperand : callOp->getOpOperands()) {
       // conservatively turn tensors into more dynamic memref than necessary.
       // If the memref type of the callee fails, introduce an extra memref.cast
       // that will either canonicalize away or fail compilation until we can do
       // something better.
       if (buffer.getType() != memRefType) {
         Value castBuffer =
             b.create<memref::CastOp>(callOp.getLoc(), memRefType, buffer);
         // Add new op equivalence info.
-        state.aliasInfo.insertNewBufferEquivalence(castBuffer, buffer);
         state.mapBuffer(tensorOperand, castBuffer);
         buffer = castBuffer;
       }
       newOperands.push_back(buffer);
     }
 
     // 4. Create the new CallOp.
     Operation *newCallOp = b.create<CallOp>(callOp.getLoc(), funcOp.sym_name(),
LogicalResult bufferize(Operation *op, OpBuilder &b,
     for (OpOperand &operand : returnOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;
       Value v = state.lookupBuffer(operand.get());
       Value returnTensor = b.create<bufferization::ToTensorOp>(
           returnOp.getLoc(), v);
       operand.set(returnTensor);
-      state.aliasInfo.insertNewBufferEquivalence(returnTensor, v);
       state.mapBuffer(returnTensor, v);
     }
     return success();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
for (auto bbArg : funcOp.getArguments()) {
       auto rankedTensorType = tensorType.dyn_cast<RankedTensorType>();
       // Cast the tensor to the most dynamic buffer possible. Further
       // canonicalizations will clean up.
       Type memRefType = rankedTensorType
                             ? getDynamicMemRefType(rankedTensorType)
                             : getContiguousOrUnrankedMemRefType(tensorType);
       Value bufferCast = b.create<bufferization::ToMemrefOp>(funcOp.getLoc(),
                                                              memRefType, bbArg);
-      state.aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
       state.mapBuffer(bbArg, bufferCast);
     }
 
     // Bufferize function body.
     return comprehensive_bufferize::bufferize(&funcOp.body(), state);
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }

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

for (OpResult opResult : ifOp->getResults()) {
       // alloc an UnrankedMemRefType + its underlying ranked MemRefType.
       assert(opResult.getType().isa<RankedTensorType>() &&
              "unsupported unranked tensor");
 
       Value resultBuffer = state.getResultBuffer(opResult);
       if (!resultBuffer)
         return failure();
 
-      state.aliasInfo.createAliasInfoEntry(resultBuffer);
       state.mapBuffer(opResult, resultBuffer);
     }
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo) const {
for (OpResult opResult : forOp->getResults()) {
 
       // TODO: More general: Matching bbArg does not bufferize to a read.
       Value resultBuffer = state.getResultBuffer(opResult);
       if (!resultBuffer)
         return failure();
 
       OpOperand &opOperand = forOp.getOpOperandForResult(opResult);
       BlockArgument bbArg = forOp.getRegionIterArgForOpOperand(opOperand);
-      state.aliasInfo.createAliasInfoEntry(resultBuffer);
-      state.aliasInfo.insertNewBufferEquivalence(bbArg, resultBuffer);
       state.mapBuffer(bbArg, resultBuffer);
       state.mapBuffer(opResult, resultBuffer);
     }
 
     // Bufferize loop body.
     if (failed(comprehensive_bufferize::bufferize(&forOp.region(), state)))
       return failure();
 
     // Finish bufferizing scf::ForOp.
     auto yieldOp = cast<scf::YieldOp>(&forOp.region().front().back());
     for (OpOperand &operand : yieldOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;
 
       OpOperand &forOperand = forOp.getOpOperandForResult(
           forOp->getResult(operand.getOperandNumber()));
       auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
+
+      // Buffers are equivalent so the work is already done and we just yield
+      // the bbArg so that it later canonicalizes away.
+      operand.set(bbArg);
+    }
+    return success();
+  }
+};
+
+LogicalResult mlir::linalg::comprehensive_bufferize::scf_ext::
+    AssertDestinationPassingStyle::run(FuncOp funcOp, BufferizationState &state,
+                                       SmallVector<Operation *> &newOps) {
+  LogicalResult status = success();
+  funcOp->walk([&](scf::YieldOp yieldOp) {
+    auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
+    if (!forOp)
+      return WalkResult::advance();
+
+    for (OpOperand &operand : yieldOp->getOpOperands()) {
+      auto tensorType = operand.get().getType().dyn_cast<TensorType>();
+      if (!tensorType)
+        continue;
+
+      OpOperand &forOperand = forOp.getOpOperandForResult(
+          forOp->getResult(operand.getOperandNumber()));
+      auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
       if (!state.aliasInfo.areEquivalentBufferizedValues(operand.get(),
                                                          bbArg)) {
         // TODO: this could get resolved with copies but it can also turn into
         // swaps so we need to be careful about order of copies.
-        return yieldOp->emitError()
-               << "Yield operand #" << operand.getOperandNumber()
-               << " does not bufferize to an equivalent buffer to the matching"
-               << " enclosing scf::for operand";
+        status =
+            yieldOp->emitError()
+        return WalkResult::interrupt();
       }
-
-      // Buffers are equivalent so the work is already done and we just yield
-      // the bbArg so that it later canonicalizes away.
-      operand.set(bbArg);
     }
-    return success();
-  }
-};
+
+    return WalkResult::advance();
+  });
+  return status;
 
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     scf::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 

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

LogicalResult bufferize(Operation *op, OpBuilder &b,
     MemRefLayoutAttrInterface layout =
         rankedMemRefType && tensorType.isa<RankedTensorType>()
             ? rankedMemRefType.getLayout()
             : MemRefLayoutAttrInterface();
     Type memRefType = getContiguousOrUnrankedMemRefType(
         castOp.getResult().getType(), layout, memorySpace);
     Value res =
         b.create<memref::CastOp>(castOp.getLoc(), memRefType, resultBuffer);
-    state.aliasInfo.insertNewBufferEquivalence(res, castOp.getResult());
     state.mapBuffer(castOp.getResult(), res);
     return success();
   }
 };
 
 struct DimOpInterface
     : public BufferizableOpInterface::ExternalModel<DimOpInterface,
                                                     tensor::DimOp> {
struct InsertOpInterface
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto insertOp = cast<tensor::InsertOp>(op);
     Location loc = insertOp.getLoc();
     Value destMemref = state.getResultBuffer(insertOp->getOpResult(0));
     b.create<memref::StoreOp>(loc, insertOp.scalar(), destMemref,
                               insertOp.indices());
     state.mapBuffer(insertOp, destMemref);
-    state.aliasInfo.insertNewBufferAlias(insertOp, destMemref);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo) const {
     return BufferRelation::Equivalent;
   }
 };
if (needCopy) {
           memref::SubViewOp::inferRankReducedResultType(
               insertSliceOp.getSourceType().getRank(), dstMemrefType,
               insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
               insertSliceOp.getMixedStrides())
               .cast<MemRefType>();
       Value subView = b.create<memref::SubViewOp>(
           loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
           insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
-      // Insert new alias.
-      state.aliasInfo.insertNewBufferAlias(subView, dstMemref);
       // Copy tensor.
       Value srcMemref = state.lookupBuffer(insertSliceOp.source());
       state.options.allocationFns->memCpyFn(b, insertSliceOp.getLoc(),
                                             srcMemref, subView);
     }
 
     state.mapBuffer(insertSliceOp.result(), dstMemref);
     return success();

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

void LinalgComprehensiveModuleBufferize::runOnOperation() {
   options.testAnalysisOnly = testAnalysisOnly;
 
   // Enable InitTensorOp elimination.
   options.addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing tensor
   // dialect ops.
   options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
+  options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, options))) {
     signalPassFailure();
     return;
   }

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

func @reading_scf_for(%t1: tensor<?xf32> {linalg.inplaceable = true},
   // CHECK:      scf.for
   %r, %v3 = scf.for %i = %c0 to %s step %c1 iter_args(%t2 = %t1, %v0 = %v) -> (tensor<?xf32>, vector<5xf32>) {
     // CHECK:      tensor.extract_slice
     // CHECK-SAME: __inplace_results_attr__ = ["true"]
     %e = tensor.extract_slice %t2[%s][%s][1] : tensor<?xf32> to tensor<?xf32>
 
     // Read from %t1 via alias %e.
     %v2 = vector.transfer_read %e[%s], %cst : tensor<?xf32>, vector<5xf32>
-    scf.yield %e, %v2 : tensor<?xf32>, vector<5xf32>
+    scf.yield %t2, %v2 : tensor<?xf32>, vector<5xf32>
   }
   // CHECK: __inplace_results_attr__ = ["true", "false"]
 
   // Use %t3 in some way without reading it, so that it does not get DCE'd.
   // CHECK:      linalg.generic
   // CHECK-SAME: __inplace_results_attr__ = ["true"]
   %o = linalg.generic #trait outs (%t3 : tensor<?xf32>) {
       ^bb(%0: f32) :
func @non_reading_scf_for(%t1: tensor<?xf32> {linalg.inplaceable = true},
   // Write to %t1.
   // CHECK:      vector.transfer_write
   // CHECK-SAME: __inplace_results_attr__ = ["true"]
   %t3 = vector.transfer_write %v, %t1[%s] : vector<5xf32>, tensor<?xf32>
 
   // This loop does not read from %t1. It only writes to it.
   // CHECK:      scf.for
   %r, %v3 = scf.for %i = %c0 to %s step %c1 iter_args(%t2 = %t1, %v0 = %v) -> (tensor<?xf32>, vector<5xf32>) {
-    // CHECK:      tensor.extract_slice
-    // CHECK-SAME: __inplace_results_attr__ = ["true"]
-    %e = tensor.extract_slice %t2[%s][%s][1] : tensor<?xf32> to tensor<?xf32>
-
-    // Write to %t1 via alias. (Overwrite %t3.)
+    // Write to %t1 via %t2. (Overwrite %t3.)
     // CHECK:      linalg.generic
     // CHECK-SAME: __inplace_results_attr__ = ["true"]
-    %o2 = linalg.generic #trait outs (%e : tensor<?xf32>) {
+    %o2 = linalg.generic #trait outs (%t2 : tensor<?xf32>) {
         ^bb(%0: f32) :
           linalg.yield %cst : f32
       } -> (tensor<?xf32>)
 
     // Read overwritten value. This is not a read of %t1.
     %v2 = vector.transfer_read %o2[%s], %cst : tensor<?xf32>, vector<5xf32>
     scf.yield %o2, %v2 : tensor<?xf32>, vector<5xf32>
   }
 
   // Use %t3 in some way without reading it, so that it does not get DCE'd.
-    // CHECK:      linalg.generic
-    // CHECK-SAME: __inplace_results_attr__ = ["true"]
   %o = linalg.generic #trait outs (%t3 : tensor<?xf32>) {
       ^bb(%0: f32) :
         linalg.yield %cst : f32
     } -> (tensor<?xf32>)
 
   //      CHECK: return
   // CHECK-SAME: {__equivalent_func_args__ = [0, -1]}
   return %o, %v3 : tensor<?xf32>, vector<5xf32>