iree/compiler/Codegen/Common/IREEComprehensiveBufferizePass.cpp
/// Stitch comprehensive bufferization inside of IREE by proceeding as follows:
///   1. a. Bufferizes InterfaceBindingSubspanOp optimistically
///      b. Insert a memref::TensorLoad to serve as glue between the buffer and
///         tensor worlds.
///      c. Record aliasInfo of memref::TensorLoad manually
///      d. Record inplaceability of memref::TensorLoad manually
///      e. Record the bufferization of memref::TensorLoad manually
///   2. Rewrite all Flow::Dispatch::TensorLoad ops as Tensor::ExtractSliceOp
///      that comprehensive bufferization understands.
///   3. Specifically select the ops we want to bufferize / skip. In the future,
///      this may be better specified with a BufferizationOpInterface.
///   4. Perform analysis and bufferization on the ops.
Stitch comprehensive bufferization inside of IREE by proceeding as follows:
1. a. Bufferizes InterfaceBindingSubspanOp optimistically
b. Insert a memref::TensorLoad to serve as glue between the buffer and
tensor worlds.
c. Record aliasInfo of memref::TensorLoad manually
d. Record inplaceability of memref::TensorLoad manually
e. Record the bufferization of memref::TensorLoad manually
2.
Stitch comprehensive bufferization inside of IREE by proceeding as follows : 1 . a . Bufferizes InterfaceBindingSubspanOp optimistically b . Insert a memref :: TensorLoad to serve as glue between the buffer and tensor worlds . c . Record aliasInfo of memref :: TensorLoad manually d . Record inplaceability of memref :: TensorLoad manually e . Record the bufferization of memref :: TensorLoad manually 2 .
void IREEComprehensiveBufferizePass::runOnOperation() {
  ModuleOp moduleOp = getOperation();
  MLIRContext *context = &getContext();
  for (auto funcOp : moduleOp.getOps<FuncOp>()) {
    OpBuilder b(context);
    BlockAndValueMapping bvm, tensorLoads;
    linalg::BufferizationAliasInfo aliasInfo(funcOp);
    SmallVector<Operation *> toEraseLate;
    auto scopeGuard = llvm::make_scope_exit([&]() {
      for (Operation *op : llvm::reverse(toEraseLate)) op->erase();
    });
    funcOp.walk([&](IREE::HAL::InterfaceBindingSubspanOp op) {
      auto shapedType =
          op.getResult().getType().dyn_cast<IREE::Flow::DispatchTensorType>();
      if (!shapedType || !shapedType.hasRank()) return;
      OpBuilder::InsertionGuard g(b);
      b.setInsertionPoint(op);
      auto memRefType = getMemrefTypeForTensor(shapedType);
      auto baseBuffer = b.create<IREE::HAL::InterfaceBindingSubspanOp>(
          op->getLoc(), memRefType, op.binding(), op.byte_offset(),
          op.byte_length(), op.dynamic_dims());
      bvm.map(op, baseBuffer);
      Value tensor = b.create<memref::TensorLoadOp>(op->getLoc(), baseBuffer);
      aliasInfo.createAliasInfoEntry(op.result());
      aliasInfo.createAliasInfoEntry(tensor);
      tensorLoads.map(op.result(), tensor);
      if (shapedType.getAccess() != IREE::Flow::TensorAccess::ReadOnly) {
        aliasInfo.setBufferizesToWritableMemory(tensor);
      }
      bvm.map(tensor, baseBuffer);
      toEraseLate.push_back(op);
    });
    funcOp.walk<WalkOrder::PostOrder>([&](IREE::Flow::DispatchTensorLoadOp op) {
      OpBuilder b(op);
      Value v = b.create<tensor::ExtractSliceOp>(
          op->getLoc(), op.result().getType().cast<RankedTensorType>(),
          tensorLoads.lookup(op.source()), op.getMixedOffsets(),
          op.getMixedSizes(), op.getMixedStrides());
      aliasInfo.createAliasInfoEntry(v);
      op.result().replaceAllUsesWith(v);
      toEraseLate.push_back(op);
    });
    funcOp.walk<WalkOrder::PostOrder>(
        [&](IREE::Flow::DispatchTensorStoreOp op) {
          OpBuilder b(op);
          Value v = b.create<tensor::InsertSliceOp>(
              op->getLoc(), op.value(), tensorLoads.lookup(op.target()),
              op.getMixedOffsets(), op.getMixedSizes(), op.getMixedStrides());
          aliasInfo.createAliasInfoEntry(v);
          toEraseLate.push_back(op);
        });
    DominanceInfo domInfo(funcOp);
    SmallVector<Operation *> ops;
    ops.reserve(funcOp.body().front().getOperations().size());
    WalkResult opsSelected =
        funcOp.body().walk([&](Operation *op) -> WalkResult {
          if (isa<IREE::HAL::InterfaceBindingSubspanOp,
                  IREE::Flow::DispatchTensorLoadOp,
                  IREE::Flow::DispatchTensorStoreOp>(op)) {
            return WalkResult::advance();
          }
          if (llvm::none_of(op->getOperandTypes(), isaTensor) &&
              llvm::none_of(op->getResultTypes(), isaTensor)) {
            return WalkResult::advance();
          }
          if (op->getParentOfType<linalg::LinalgOp>())
            return WalkResult::advance();
          if (isa<CallOpInterface>(op)) {
            return static_cast<LogicalResult>(op->emitError(
                "CallOpInterface bufferization not supported in IREE"));
          }
          ops.push_back(op);
          return WalkResult::advance();
        });
    if (opsSelected.wasInterrupted() ||
        failed(linalg::inPlaceAnalysis(ops, aliasInfo, domInfo))) {
      return signalPassFailure();
    }
    for (Operation *op : ops) {
      if (failed(linalg::bufferizeOp(op, bvm, aliasInfo, allocationFn,
                                     nullptr))) {
        return signalPassFailure();
      }
    }
  }
}
void IREEComprehensiveBufferizePass :: runOnOperation ( ) { ModuleOp moduleOp = getOperation ( ) ; MLIRContext * context = & getContext ( ) ; for ( auto funcOp : moduleOp . getOps < FuncOp > ( ) ) { OpBuilder b ( context ) ; BlockAndValueMapping bvm , tensorLoads ; linalg :: BufferizationAliasInfo aliasInfo ( funcOp ) ; SmallVector < Operation * > toEraseLate ; auto scopeGuard = llvm :: make_scope_exit ( [ & ] ( ) { for ( Operation * op : llvm :: reverse ( toEraseLate ) ) op -> erase ( ) ; } ) ; funcOp . walk ( [ & ] ( IREE :: HAL :: InterfaceBindingSubspanOp op ) { auto shapedType = op . getResult ( ) . getType ( ) . dyn_cast < IREE :: Flow :: DispatchTensorType > ( ) ; if ( ! shapedType || ! shapedType . hasRank ( ) ) return ; OpBuilder :: InsertionGuard g ( b ) ; b . setInsertionPoint ( op ) ; auto memRefType = getMemrefTypeForTensor ( shapedType ) ; auto baseBuffer = b . create < IREE :: HAL :: InterfaceBindingSubspanOp > ( op -> getLoc ( ) , memRefType , op . binding ( ) , op . byte_offset ( ) , op . byte_length ( ) , op . dynamic_dims ( ) ) ; bvm . map ( op , baseBuffer ) ; Value tensor = b . create < memref :: TensorLoadOp > ( op -> getLoc ( ) , baseBuffer ) ; aliasInfo . createAliasInfoEntry ( op . result ( ) ) ; aliasInfo . createAliasInfoEntry ( tensor ) ; tensorLoads . map ( op . result ( ) , tensor ) ; if ( shapedType . getAccess ( ) != IREE :: Flow :: TensorAccess :: ReadOnly ) { aliasInfo . setBufferizesToWritableMemory ( tensor ) ; } bvm . map ( tensor , baseBuffer ) ; toEraseLate . push_back ( op ) ; } ) ; funcOp . walk < WalkOrder :: PostOrder > ( [ & ] ( IREE :: Flow :: DispatchTensorLoadOp op ) { OpBuilder b ( op ) ; Value v = b . create < tensor :: ExtractSliceOp > ( op -> getLoc ( ) , op . result ( ) . getType ( ) . cast < RankedTensorType > ( ) , tensorLoads . lookup ( op . source ( ) ) , op . getMixedOffsets ( ) , op . getMixedSizes ( ) , op . getMixedStrides ( ) ) ; aliasInfo . createAliasInfoEntry ( v ) ; op . result ( ) . replaceAllUsesWith ( v ) ; toEraseLate . push_back ( op ) ; } ) ; funcOp . walk < WalkOrder :: PostOrder > ( [ & ] ( IREE :: Flow :: DispatchTensorStoreOp op ) { OpBuilder b ( op ) ; Value v = b . create < tensor :: InsertSliceOp > ( op -> getLoc ( ) , op . value ( ) , tensorLoads . lookup ( op . target ( ) ) , op . getMixedOffsets ( ) , op . getMixedSizes ( ) , op . getMixedStrides ( ) ) ; aliasInfo . createAliasInfoEntry ( v ) ; toEraseLate . push_back ( op ) ; } ) ; DominanceInfo domInfo ( funcOp ) ; SmallVector < Operation * > ops ; ops . reserve ( funcOp . body ( ) . front ( ) . getOperations ( ) . size ( ) ) ; WalkResult opsSelected = funcOp . body ( ) . walk ( [ & ] ( Operation * op ) -> WalkResult { if ( isa < IREE :: HAL :: InterfaceBindingSubspanOp , IREE :: Flow :: DispatchTensorLoadOp , IREE :: Flow :: DispatchTensorStoreOp > ( op ) ) { return WalkResult :: advance ( ) ; } if ( llvm :: none_of ( op -> getOperandTypes ( ) , isaTensor ) && llvm :: none_of ( op -> getResultTypes ( ) , isaTensor ) ) { return WalkResult :: advance ( ) ; } if ( op -> getParentOfType < linalg :: LinalgOp > ( ) ) return WalkResult :: advance ( ) ; if ( isa < CallOpInterface > ( op ) ) { return static_cast < LogicalResult > ( op -> emitError ( " " ) ) ; } ops . push_back ( op ) ; return WalkResult :: advance ( ) ; } ) ; if ( opsSelected . wasInterrupted ( ) || failed ( linalg :: inPlaceAnalysis ( ops , aliasInfo , domInfo ) ) ) { return signalPassFailure ( ) ; } for ( Operation * op : ops ) { if ( failed ( linalg :: bufferizeOp ( op , bvm , aliasInfo , allocationFn , nullptr ) ) ) { return signalPassFailure ( ) ; } } } }
Stitch comprehensive bufferization inside of IREE by proceeding as follows:
1. a. Bufferizes InterfaceBindingSubspanOp optimistically
b. Insert a memref::TensorLoad to serve as glue between the buffer and
tensor worlds.
Stitch comprehensive bufferization inside of IREE by proceeding as follows : 1 . a . Bufferizes InterfaceBindingSubspanOp optimistically b . Insert a memref :: TensorLoad to serve as glue between the buffer and tensor worlds .
// 1. First go over all hal.interface.binding.subspan ops and create // counterparts working with memrefs. // These are used until late, erase on scoped exit. // 1.a. Just change the result type of the InterfaceBindingSubspanOp to // from the base buffer. // This op does not operate on core tensor types and has half-side // effecting semantics. It cannot be added to BufferizationAliasInfo. // Instead: // 1.b. Insert a memref::TensorLoad to serve as glue between the buffer // and tensor worlds. // 1.c. Insert a new entry manually into the existing aliasInfo. // 1.d. Mark tensors that bufferize to writeable memory as such. // 1.e. Save tensor -> baseBuffer into BVM. // Drop the original op that is now bufferized. // 2. Rewrite all Flow::Dispatch::TensorLoad ops as Tensor::ExtractSliceOp. // Insert a new entry manually into the existing aliasInfo. // Insert a new entry manually into the existing aliasInfo. // TODO: Visit all the operations that return `tensor`s that are not handled // by comprehensive bufferize. // 3. Specifically select the ops we want to bufferize / skip. In the // future, this may be better specified with a BufferizationOpInterface. // TODO: if we want to bufferize function calls, we need FuncOp // and to pass a proper bufferizedFunctionTypes. // 4. Perform inplaceability analysis of `ops`. // 5. Perform bufferization. /*bufferizedFunctionTypes=*/