[Mlir-commits] [mlir] [NFC][mlir][bufferization] Align alloc/memcpy/cast options hooks to the rest (PR #206966)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 1 05:41:43 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-mlprogram
Author: Andrei Golubev (andrey-golubev)
<details>
<summary>Changes</summary>
There are three hooks that have a different, more complicated API than the other ones: allocationFn, memCpyFn, castFn. All three are called via "helper" functions that check whether the hooks are set to call them and fall back to default implementations otherwise.
Other hooks (e.g. unknownTypeConverterFn) in the bufferization options have a different "API": they are always set to some default implementation and can be overwritten by the user. This is a simpler API overall and seems it can be universal.
---
Full diff: https://github.com/llvm/llvm-project/pull/206966.diff
8 Files Affected:
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+9-17)
- (modified) mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp (+2-4)
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+31-43)
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+11-12)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+2-2)
- (modified) mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp (+1-1)
- (modified) mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp (+1-1)
- (modified) mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp (+4-4)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 9f49a31c1583b..203680796a289 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -303,23 +303,6 @@ struct BufferizationOptions {
/// Return `true` if the given op should be bufferized.
bool isOpAllowed(Operation *op) const;
- /// Helper functions for allocation and memory copying.
- std::optional<AllocationFn> allocationFn;
- std::optional<MemCpyFn> memCpyFn;
- std::optional<CastFn> castFn;
-
- /// Create a memref allocation with the given type and dynamic extents.
- FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
- ValueRange dynShape) const;
-
- /// Creates a memcpy between two given buffers.
- LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from,
- Value to) const;
-
- /// Creates a cast function from a buffer value to a new type.
- FailureOr<Value> createCast(OpBuilder &b, Location loc, Type dest,
- Value value) const;
-
/// Specifies whether not bufferizable ops are allowed in the input. If so,
/// bufferization.to_buffer and bufferization.to_tensor ops are inserted at
/// the boundaries.
@@ -353,6 +336,15 @@ struct BufferizationOptions {
/// predictable.
void setFunctionBoundaryTypeConversion(LayoutMapOption layoutMapOption);
+ /// Create a memref allocation with the given type and dynamic extents.
+ AllocationFn allocationFn = nullptr;
+
+ /// Creates a memcpy between two given buffers.
+ MemCpyFn memCpyFn = nullptr;
+
+ /// Creates a cast function from a buffer value to a new type.
+ CastFn castFn = nullptr;
+
/// Type conversion from tensors to buffers. This type conversion is used to
/// determine bufferized function argument and result types.
///
diff --git a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
index f810000dd082f..0fe7be82959a3 100644
--- a/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -173,12 +173,10 @@ struct SelectOpInterface
"incompatible buffer types on true/false operands");
}
if (trueBuffer.getType() != *targetType) {
- trueBuffer =
- *options.createCast(rewriter, loc, *targetType, trueBuffer);
+ trueBuffer = *options.castFn(rewriter, loc, *targetType, trueBuffer);
}
if (falseBuffer.getType() != *targetType) {
- falseBuffer =
- *options.createCast(rewriter, loc, *targetType, falseBuffer);
+ falseBuffer = *options.castFn(rewriter, loc, *targetType, falseBuffer);
}
}
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 7f656164f3a15..e39b20492e5ed 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -358,6 +358,34 @@ BaseMemRefType getMemRefTypeWithFullyDynamicLayout(ArrayRef<int64_t> shape,
return MemRefType::get(shape, elementType, stridedLayout, memorySpace);
}
+/// Create a memref allocation with the given type and dynamic extents.
+FailureOr<Value> defaultCreateAlloc(OpBuilder &b, Location loc, MemRefType type,
+ ValueRange dynShape,
+ unsigned int bufferAlignment) {
+ // Default buffer allocation via AllocOp.
+ if (bufferAlignment != 0)
+ return memref::AllocOp::create(b, loc, type, dynShape,
+ b.getI64IntegerAttr(bufferAlignment))
+ .getResult();
+ return memref::AllocOp::create(b, loc, type, dynShape).getResult();
+}
+
+/// Create a memory copy between two memref buffers.
+LogicalResult defaultCreateMemCpy(OpBuilder &b, Location loc, Value from,
+ Value to) {
+ memref::CopyOp::create(b, loc, from, to);
+ return success();
+}
+
+FailureOr<Value> defaultCreateCast(OpBuilder &b, Location loc, Type dest,
+ Value value) {
+ assert(isa<BaseMemRefType>(dest) && "expected BaseMemRefType");
+ assert(isa<BaseMemRefType>(value.getType()) && "expected BaseMemRefType");
+ assert(memref::CastOp::areCastCompatible(value.getType(), dest) &&
+ "cast incompatible");
+ return memref::CastOp::create(b, loc, dest, value).getResult();
+}
+
/// Default function arg type converter: Use a fully dynamic layout map.
BufferLikeType
defaultFunctionArgTypeConverter(TensorLikeType type, Attribute memorySpace,
@@ -405,7 +433,9 @@ defaultReconcileBufferTypeMismatch(BufferLikeType x, BufferLikeType y,
// Default constructor for BufferizationOptions.
BufferizationOptions::BufferizationOptions()
- : functionArgTypeConverterFn(defaultFunctionArgTypeConverter),
+ : allocationFn(defaultCreateAlloc), memCpyFn(defaultCreateMemCpy),
+ castFn(defaultCreateCast),
+ functionArgTypeConverterFn(defaultFunctionArgTypeConverter),
unknownTypeConverterFn(defaultUnknownTypeConverter),
reconcileBufferTypeMismatchFn(defaultReconcileBufferTypeMismatch) {}
@@ -811,48 +841,6 @@ void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
rewriter.replaceOp(op, replacements);
}
-//===----------------------------------------------------------------------===//
-// Bufferization-specific scoped alloc insertion support.
-//===----------------------------------------------------------------------===//
-
-/// Create a memref allocation with the given type and dynamic extents.
-FailureOr<Value> BufferizationOptions::createAlloc(OpBuilder &b, Location loc,
- MemRefType type,
- ValueRange dynShape) const {
- if (allocationFn)
- return (*allocationFn)(b, loc, type, dynShape, bufferAlignment);
-
- // Default bufferallocation via AllocOp.
- if (bufferAlignment != 0)
- return memref::AllocOp::create(b, loc, type, dynShape,
- b.getI64IntegerAttr(bufferAlignment))
- .getResult();
- return memref::AllocOp::create(b, loc, type, dynShape).getResult();
-}
-
-/// Create a memory copy between two memref buffers.
-LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
- Value from, Value to) const {
- if (memCpyFn)
- return (*memCpyFn)(b, loc, from, to);
-
- memref::CopyOp::create(b, loc, from, to);
- return success();
-}
-
-FailureOr<Value> BufferizationOptions::createCast(OpBuilder &b, Location loc,
- Type dest,
- Value value) const {
- if (castFn)
- return (*castFn)(b, loc, dest, value);
-
- assert(isa<BaseMemRefType>(dest) && "expected BaseMemRefType");
- assert(isa<BaseMemRefType>(value.getType()) && "expected BaseMemRefType");
- assert(memref::CastOp::areCastCompatible(value.getType(), dest) &&
- "cast incompatible");
- return memref::CastOp::create(b, loc, dest, value).getResult();
-}
-
//===----------------------------------------------------------------------===//
// Bufferization-specific IRMapping support with debugging.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 320f426985750..45f5b3eaa5aea 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -60,7 +60,7 @@ FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
// a fix extra conditions in `isGuaranteedCastCompatible`.
if (memref::CastOp::areCastCompatible(srcType, destType) &&
isGuaranteedCastCompatible(srcType, destType)) {
- Value casted = *options.createCast(b, value.getLoc(), destType, value);
+ Value casted = *options.castFn(b, value.getLoc(), destType, value);
return casted;
}
@@ -73,11 +73,11 @@ FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
dynamicOperands.push_back(size);
}
- FailureOr<Value> copy =
- options.createAlloc(b, loc, destType, dynamicOperands);
+ FailureOr<Value> copy = options.allocationFn(
+ b, loc, destType, dynamicOperands, options.bufferAlignment);
if (failed(copy))
return failure();
- if (failed(options.createMemCpy(b, loc, value, *copy)))
+ if (failed(options.memCpyFn(b, loc, value, *copy)))
return failure();
return copy;
}
@@ -103,11 +103,9 @@ LogicalResult mlir::bufferization::foldToBufferToTensorPair(
if (!llvm::isa<BaseMemRefType>(srcType) ||
!llvm::isa<BaseMemRefType>(destType)) {
// Non-builtin case: the best is to try the user-provided cast.
- assert(options.castFn.has_value() &&
- "user-provided cast is required for non-builtin types");
auto replacement =
- options.createCast(rewriter, bufferToTensor.getBuffer().getLoc(),
- destType, bufferToTensor.getBuffer());
+ options.castFn(rewriter, bufferToTensor.getBuffer().getLoc(), destType,
+ bufferToTensor.getBuffer());
if (failed(replacement))
return failure();
rewriter.replaceOp(toBuffer, *replacement);
@@ -194,14 +192,15 @@ LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
}
- FailureOr<Value> alloc = options.createAlloc(
- rewriter, loc, llvm::cast<MemRefType>(*allocType), dynamicDims);
+ FailureOr<Value> alloc =
+ options.allocationFn(rewriter, loc, llvm::cast<MemRefType>(*allocType),
+ dynamicDims, options.bufferAlignment);
if (failed(alloc))
return failure();
// Create memory copy (if any).
if (getCopy()) {
- if (failed(options.createMemCpy(rewriter, loc, copyBuffer, *alloc)))
+ if (failed(options.memCpyFn(rewriter, loc, copyBuffer, *alloc)))
return failure();
}
@@ -619,7 +618,7 @@ MaterializeInDestinationOp::bufferize(RewriterBase &rewriter,
auto srcBuffer = getBuffer(rewriter, getSource(), options, state);
if (failed(srcBuffer))
return failure();
- if (failed(options.createMemCpy(rewriter, getLoc(), *srcBuffer, buffer)))
+ if (failed(options.memCpyFn(rewriter, getLoc(), *srcBuffer, buffer)))
return failure();
replaceOpWithBufferizedValues(rewriter, getOperation(),
tensorDest ? ValueRange(buffer) : ValueRange());
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index e0262abf313b8..d6e0b12733976 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -477,8 +477,8 @@ bufferization::bufferizeBlockSignature(Block *block, RewriterBase &rewriter,
// A cast is needed if the operand and the block argument have different
// bufferized types.
if (type != *operandBufferType) {
- bufferizedOperand = *options.createCast(rewriter, operand.getLoc(),
- type, bufferizedOperand);
+ bufferizedOperand = *options.castFn(rewriter, operand.getLoc(), type,
+ bufferizedOperand);
}
newOperands.push_back(bufferizedOperand);
}
diff --git a/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
index c2cbc95e94343..a7c5934132b35 100644
--- a/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -146,7 +146,7 @@ struct GlobalStoreOpInterface
}
auto memcpy =
- options.createMemCpy(rewriter, loc, sourceMemref.value(), targetMemref);
+ options.memCpyFn(rewriter, loc, sourceMemref.value(), targetMemref);
if (failed(memcpy)) {
return failure();
}
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index 401be7440f5a2..f8302126a4a4e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -37,7 +37,7 @@ static Value castBuffer(OpBuilder &b, Value buffer, Type type,
if (buffer.getType() == type)
return buffer;
- return *options.createCast(b, buffer.getLoc(), type, buffer);
+ return *options.castFn(b, buffer.getLoc(), type, buffer);
}
/// Helper function for loop bufferization. Return "true" if the given value
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 41a3d96b832d1..525d57341cee5 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -724,7 +724,7 @@ struct InsertSliceOpInterface
getBuffer(rewriter, insertSliceOp.getSource(), options, state);
if (failed(srcMemref))
return failure();
- if (failed(options.createMemCpy(rewriter, loc, *srcMemref, subView)))
+ if (failed(options.memCpyFn(rewriter, loc, *srcMemref, subView)))
return failure();
replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
@@ -1002,8 +1002,8 @@ struct ParallelInsertSliceOpInterface
parallelInsertSliceOp.getMixedStrides());
// This memcpy will fold away if everything bufferizes in-place.
- if (failed(options.createMemCpy(rewriter, parallelInsertSliceOp.getLoc(),
- *srcBuffer, subview)))
+ if (failed(options.memCpyFn(rewriter, parallelInsertSliceOp.getLoc(),
+ *srcBuffer, subview)))
return failure();
// In case the source was allocated in the same block, make sure that the
@@ -1175,7 +1175,7 @@ struct ConcatOpInterface
rewriter, loc, subviewMemRefType, dstBuffer, offsets, sizes, strides);
// Copy the source buffer into the destination subview.
- if (failed(options.createMemCpy(rewriter, loc, *srcBuffer, subview)))
+ if (failed(options.memCpyFn(rewriter, loc, *srcBuffer, subview)))
return failure();
concatDimOffset = sum(concatDimOffset, concatDimSize);
``````````
</details>
https://github.com/llvm/llvm-project/pull/206966
More information about the Mlir-commits
mailing list