[mlir][bufferization] Generalize getBufferType

[springerm]

Summary

This change generalizes getBufferType. This function can be used to predict the buffer type of any tensor value (not just BlockArguments) without changing any IR. It also subsumes getMemorySpace. This is useful for loop bufferization, where the precise buffer type of an iter_arg cannot be known without examining the loop body.

[bkramer]

looks good

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 bool allocationDoesNotEscape(OpResult opResult);
 
 /// Lookup the buffer for the given value. If the value was not bufferized
 /// yet, wrap it in a ToMemrefOp. Otherwise, it is the result of a ToTensorOp,
 /// from which the memref operand is returned.
 FailureOr<Value> getBuffer(RewriterBase &rewriter, Value value,
                            const BufferizationOptions &options);
 
-/// Return the buffer type for a given Value (tensor) after bufferization.
+/// Return the buffer type for a given Value (tensor) after bufferization
+/// without bufferizing any IR.
 ///
-/// Note: Op implementations should preferrably call `getBuffer()->getType()`.
-/// This function should only be used if `getBuffer` cannot be used.
+/// Note: It should be sufficient to call `getBuffer()->getType()` in most
+/// cases. However, when a buffer type should be predicted without modifying any
+/// IR, this function can be used.
+///
+/// This function is a wrapper around BufferizableOpInterface::getBufferType.
 FailureOr<BaseMemRefType> getBufferType(Value value,
                                         const BufferizationOptions &options);
 
 /// Replace an op with replacement values. The op is deleted. Tensor OpResults
 /// must be replaced with memref values.
 void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
                                    ValueRange values);
 

 BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
                                                    unsigned memorySpace = 0);
 
 /// Return a MemRef type with a static identity layout (i.e., no layout map). If
 /// the given tensor type is unranked, return an unranked MemRef type.
 BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
                                                      unsigned memorySpace = 0);
 
+/// Return the owner of the given value. In case of a BlockArgument that is the
+/// owner of the block. In case of an OpResult that is the defining op.
+Operation *getOwnerOfValue(Value value);
+
+namespace detail {
+/// This is the default implementation of
+/// BufferizableOpInterface::getBufferType. Should not be called from other
+/// places.
+FailureOr<BaseMemRefType>
+defaultGetBufferType(Value value, const BufferizationOptions &options);
+} // namespace detail
+
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
         /*args=*/(ins "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return success();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
-          Return the bufferized type of the given tensor block argument. The
-          block argument is guaranteed to belong to a block of this op.
+          Return the bufferized type of the given tensor value (without
+          bufferizing the IR). The value is either a BlockArgument of a block
+          that belongs to this op or an OpResult of the given op.
+
+          This method is useful when the bufferized type of value must be
+          predicted before modifying any IR.
         }],
         /*retType=*/"FailureOr<BaseMemRefType>",
         /*methodName=*/"getBufferType",
-        /*args=*/(ins "BlockArgument":$bbArg,
+        /*args=*/(ins "Value":$value,
                       "const BufferizationOptions &":$options),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
-          assert(bbArg.getOwner()->getParentOp() == $_op &&
-                 "bbArg must belong to this op");
-          assert(bbArg.getType().isa<TensorType>() &&
-                 "expected tensor type");
-          return bufferization::getMemRefType(bbArg, options);
-        }]
-      >,
-      InterfaceMethod<
-        /*desc=*/[{
-          Return the memory space of the given tensor OpResult if specified on
-          this op. If not specified, return `failure`.
-
-          This method will never be called with OpResults that do not bufferize
-          to a memory allocation.
-        }],
-        /*retType=*/"FailureOr<unsigned>",
-        /*methodName=*/"getMemorySpace",
-        /*args=*/(ins "OpResult":$opResult),
-        /*methodBody=*/"",
-        /*defaultImplementation=*/[{
-          assert(cast<BufferizableOpInterface>($_op.getOperation())
-                     .bufferizesToAllocation(opResult)
-                 && "expected allocation");
-          return failure();
+          assert(getOwnerOfValue(value) == $_op.getOperation() &&
+                 "expected that value belongs to this op");
+          return bufferization::detail::defaultGetBufferType(value, options);
         }]
       >,
   ];
 
   let extraClassDeclaration = [{
     /// Resolve out-of-place tensor OpOperands with explicit allocations in the
     /// form of `bufferization.alloc_tensor` ops.
     LogicalResult resolveTensorOpOperandConflicts(

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
 
     bool isMemoryWrite(OpResult opResult, const AnalysisState &state);
 
     bool bufferizesToAllocation(OpResult opResult) { return true; }
 
-    FailureOr<unsigned> getMemorySpace(OpResult opResult) {
-      if (getMemorySpace().has_value())
-        return static_cast<unsigned>(*getMemorySpace());
-      return failure();
-    }
-
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
     SmallVector<OpResult> getAliasingOpResult(
         OpOperand &opOperand, const AnalysisState &state);
 
+    FailureOr<BaseMemRefType> getBufferType(
+        Value value, const BufferizationOptions &options);
+
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicDim(unsigned idx) {
       return getType().isDynamicDim(idx);
     }
LogicalResult bufferize(RewriterBase &rewriter,
       // FinalizingBufferize.
       return success();
     }
 
     bool isWritable(Value value, const AnalysisState &state) const {
       // It is unknown whether the memref operand is writable or not.
       return false;
     }
+
+    FailureOr<BaseMemRefType> getBufferType(
+        Value value, const BufferizationOptions &options) {
+      return getMemref().getType().cast<BaseMemRefType>();
+    }
   }];
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

 
 #define DEBUG_TYPE "bufferizable-op-interface"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << (X))
 
 using namespace mlir;
 using namespace bufferization;
 
-/// Return the owner of the given value.
-static Operation *getOwnerOfValue(Value value) {
+Operation *bufferization::getOwnerOfValue(Value value) {
   if (auto opResult = value.dyn_cast<OpResult>())
     return opResult.getDefiningOp();
   return value.cast<BlockArgument>().getOwner()->getParentOp();
 }
 
 bool bufferization::allocationDoesNotEscape(OpResult opResult) {
 #ifndef NDEBUG
   auto bufferizableOp = opResult.getDefiningOp<BufferizableOpInterface>();
#endif // NDEBUG
   if (failed(memrefType))
     return failure();
   ensureToMemrefOpIsValid(value, *memrefType);
   return rewriter
       .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
+FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
+    Value value, const BufferizationOptions &options) {
+  assert(value.getType().isa<TensorType>() && "expected tensor type");
+
+  // No further analysis is possible for a block argument.
+  if (value.isa<BlockArgument>())
+    return bufferization::getMemRefType(value, options);
+
+  // Value is an OpResult.
+  Operation *op = getOwnerOfValue(value);
+  auto opResult = value.cast<OpResult>();
+  auto bufferizableOp = cast<BufferizableOpInterface>(op);
+  AnalysisState state(options);
+  auto aliasingOperands = bufferizableOp.getAliasingOpOperand(opResult, state);
+  if (!aliasingOperands.empty() &&
+      bufferizableOp.bufferRelation(opResult, state) ==
+          BufferRelation::Equivalent) {
+    // If the OpResult has an equivalent OpOperand, both OpResult and
+    // OpOperand bufferize to the exact same buffer type.
+    Value equivalentOperand = aliasingOperands.front()->get();
+    return getBufferType(equivalentOperand, options);
+  }
+
+  // If we do not know the memory space and there is no default memory space,
+  // report a failure.
+  if (!options.defaultMemorySpace.has_value())
+    return op->emitError("could not infer memory space");
+
+  return getMemRefType(value, options, /*layout=*/{},
+                       *options.defaultMemorySpace);
+}
+
 /// Return the buffer type for a given Value (tensor) after bufferization.
 FailureOr<BaseMemRefType>
 bufferization::getBufferType(Value value, const BufferizationOptions &options) {
   assert(value.getType().isa<TensorType>() && "unexpected non-tensor type");
   Operation *op = getOwnerOfValue(value);
+  auto bufferizableOp = options.dynCastBufferizableOp(op);
+  if (bufferizableOp)
+    return bufferizableOp.getBufferType(value, options);
 
-  // ToTensorOp: Take buffer type directly from the op.
-  if (auto toTensorOp = value.getDefiningOp<bufferization::ToTensorOp>())
-    return toTensorOp.getMemref().getType().cast<BaseMemRefType>();
-
-  // If value is a bbArg of a bufferizable op: query op interface.
-  if (auto bbArg = value.dyn_cast<BlockArgument>())
-    if (auto bufferizableOp =
-            options.dynCastBufferizableOp(bbArg.getOwner()->getParentOp()))
-      return bufferizableOp.getBufferType(bbArg, options);
-
-  // Check value is a new buffer allocation with a memory space attribute. In
-  // that case we can at least infer the memory space.
-  Optional<unsigned> memorySpace;
-  if (auto opResult = value.dyn_cast<OpResult>()) {
-    if (auto bufferizableOp =
-            options.dynCastBufferizableOp(opResult.getDefiningOp())) {
-      if (bufferizableOp.bufferizesToAllocation(opResult)) {
-        FailureOr<unsigned> queriedMemorySpace =
-            bufferizableOp.getMemorySpace(opResult);
-        if (!failed(queriedMemorySpace))
-          memorySpace = *queriedMemorySpace;
-      }
-    }
-  }
-
-  // If we still do not know the memory space, use the default memory space (if
-  // any).
-  if (!memorySpace.has_value())
-    memorySpace = options.defaultMemorySpace;
-
-  // If we still do not know the memory space, report a failure.
-  if (!memorySpace.has_value())
+  // Op is not bufferizable.
+  if (!options.defaultMemorySpace.has_value())
     return op->emitError("could not infer memory space");
 
-  return getMemRefType(value, options, /*layout=*/{}, *memorySpace);
+  return getMemRefType(value, options, /*layout=*/{},
+                       *options.defaultMemorySpace);
 }
 
 void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
                                                   Operation *op,
                                                   ValueRange values) {
   assert(values.size() == op->getNumResults() &&
          "expected one value per OpResult");
   OpBuilder::InsertionGuard g(rewriter);

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

 
 //===----------------------------------------------------------------------===//
 // AllocTensorOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
                                        const BufferizationOptions &options) {
   OpBuilder::InsertionGuard g(rewriter);
-  Operation *op = this->getOperation();
   Location loc = getLoc();
 
   // Nothing to do for dead AllocTensorOps.
   if (getOperation()->getUses().empty()) {
     rewriter.eraseOp(getOperation());
     return success();
   }
 
   // Get "copy" buffer.
   Value copyBuffer;
   if (getCopy()) {
     FailureOr<Value> maybeCopyBuffer = getBuffer(rewriter, getCopy(), options);
     if (failed(maybeCopyBuffer))
       return failure();
     copyBuffer = *maybeCopyBuffer;
   }
 
-  // Compute memory space of this allocation.
-  unsigned memorySpace;
-  if (getMemorySpace().has_value()) {
-    memorySpace = *getMemorySpace();
-  } else if (getCopy()) {
-    memorySpace =
-        copyBuffer.getType().cast<BaseMemRefType>().getMemorySpaceAsInt();
-  } else if (options.defaultMemorySpace.has_value()) {
-    memorySpace = *options.defaultMemorySpace;
-  } else {
-    return op->emitError("could not infer memory space");
-  }
-
   // Create memory allocation.
-  auto allocType =
-      MemRefType::get(getType().getShape(), getType().getElementType(),
-                      AffineMap(), memorySpace);
+  auto allocType = getBufferType(getResult(), options);
+  if (failed(allocType))
+    return failure();
   SmallVector<Value> dynamicDims = getDynamicSizes();
   if (getCopy()) {
     assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
     populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
   }
-  FailureOr<Value> alloc =
-      options.createAlloc(rewriter, loc, allocType, dynamicDims);
+  FailureOr<Value> alloc = options.createAlloc(
+      rewriter, loc, allocType->cast<MemRefType>(), dynamicDims);
   if (failed(alloc))
     return failure();
 
   // Create memory copy (if any).
   if (getCopy()) {
     if (failed(options.createMemCpy(rewriter, loc, copyBuffer, *alloc)))
       return failure();
   }

 
 SmallVector<OpResult>
 AllocTensorOp::getAliasingOpResult(OpOperand &opOperand,
                                    const AnalysisState &state) {
   // This is a new allocation. It does not alias with any other buffer.
   return {};
 }
 
+FailureOr<BaseMemRefType>
+AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options) {
+  assert(value == getResult() && "invalid value");
+
+  // Compute memory space of this allocation.
+  unsigned memorySpace;
+  if (getMemorySpace().has_value()) {
+    memorySpace = *getMemorySpace();
+  } else if (getCopy()) {
+    auto copyBufferType = bufferization::getBufferType(getCopy(), options);
+    if (failed(copyBufferType))
+      return failure();
+    memorySpace = copyBufferType->getMemorySpaceAsInt();
+  } else if (options.defaultMemorySpace.has_value()) {
+    memorySpace = *options.defaultMemorySpace;
+  } else {
+    return getOperation()->emitError("could not infer memory space");
+  }
+
+  return getMemRefTypeWithStaticIdentityLayout(getType(), memorySpace);
+}
+
 LogicalResult AllocTensorOp::verify() {
   if (getCopy() && !getDynamicSizes().empty())
     return emitError("dynamic sizes not needed when copying a tensor");
   if (!getCopy() && getType().getNumDynamicDims() !=
                         static_cast<int64_t>(getDynamicSizes().size()))
     return emitError("expected ")
            << getType().getNumDynamicDims() << " dynamic sizes";
   if (getCopy() && getCopy().getType() != getType())

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

LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     }
 
     rewriter.updateRootInPlace(
         yieldOp, [&]() { yieldOp.getResultsMutable().assign(yieldValues); });
     return success();
   }
 
   FailureOr<BaseMemRefType>
-  getBufferType(Operation *op, BlockArgument bbArg,
+  getBufferType(Operation *op, Value value,
                 const BufferizationOptions &options) const {
     auto forOp = cast<scf::ForOp>(op);
+    // TODO: Only block arguments supported at the moment.
+    if (value.isa<OpResult>())
+      return failure();
+    auto bbArg = value.cast<BlockArgument>();
     return bufferization::getBufferType(
         forOp.getOpOperandForRegionIterArg(bbArg).get(), options);
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto forOp = cast<scf::ForOp>(op);
     Block *oldLoopBody = &forOp.getLoopBody().front();

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
     Location loc = extractSliceOp.getLoc();
 
     // Get source buffer.
     FailureOr<Value> srcMemref =
         getBuffer(rewriter, extractSliceOp.getSource(), options);
     if (failed(srcMemref))
       return failure();
-    auto srcMemrefType = srcMemref->getType().cast<MemRefType>();
 
     // Take a subview of the source buffer.
-    auto subviewMemRefType =
-        memref::SubViewOp::inferRankReducedResultType(
-            extractSliceOp.getType().getShape(), srcMemrefType, mixedOffsets,
-            mixedSizes, mixedStrides)
-            .cast<MemRefType>();
+    auto resultMemrefType =
+        getBufferType(op, extractSliceOp.getResult(), options);
+    if (failed(resultMemrefType))
+      return failure();
     Value subView = rewriter.create<memref::SubViewOp>(
-        loc, subviewMemRefType, *srcMemref, mixedOffsets, mixedSizes,
-        mixedStrides);
+        loc, resultMemrefType->cast<MemRefType>(), *srcMemref, mixedOffsets,
+        mixedSizes, mixedStrides);
 
     replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
+
+  FailureOr<BaseMemRefType>
+  getBufferType(Operation *op, Value value,
+                const BufferizationOptions &options) const {
+    auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
+    assert(value == extractSliceOp.getResult() && "invalid value");
+    auto srcMemrefType =
+        bufferization::getBufferType(extractSliceOp.getSource(), options);
+    if (failed(srcMemrefType))
+      return failure();
+    SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
+    return memref::SubViewOp::inferRankReducedResultType(
+               extractSliceOp.getType().getShape(),
+               srcMemrefType->cast<MemRefType>(), mixedOffsets, mixedSizes,
+               mixedStrides)
+        .cast<BaseMemRefType>();
+  }
 };
 
 /// Bufferization of tensor.extract. Replace with memref.load.
 struct ExtractOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
                                                     tensor::ExtractOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {