[mlir][bufferization] Use Type instead of Value in unknown conversion #144658
Conversation
andrey-golubev
requested review from
aartbik,
PeimingLiu,
yinying-lisa-li and
matthias-springer
as code owners
llvmbot
added
mlir:sparse
|
@llvm/pr-subscribers-mlir-bufferization @llvm/pr-subscribers-mlir-sparse Author: Andrei Golubev (andrey-golubev) ChangesGenerally, bufferization should be able to create a memref from a tensor without needing to know more than just a mlir::Type. Thus, change BufferizationOptions::UnknownTypeConverterFn to accept just a type (mlir::TensorType for now) instead of mlir::Value. Additionally, apply the same rationale to getMemRefType() helper function. Both changes are prerequisites to enable custom types support in one-shot bufferization. Full diff: https://github.com/llvm/llvm-project/pull/144658.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index adccbef754ec5..2fb795f16ae2c 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -265,9 +265,9 @@ struct BufferizationOptions {
std::function<BaseMemRefType(TensorType, Attribute memorySpace,
func::FuncOp, const BufferizationOptions &)>;
/// Tensor -> MemRef type converter.
- /// Parameters: Value, memory space, bufferization options
+ /// Parameters: tensor type, memory space, bufferization options
using UnknownTypeConverterFn = std::function<BaseMemRefType(
- Value, Attribute memorySpace, const BufferizationOptions &)>;
+ TensorType, Attribute memorySpace, const BufferizationOptions &)>;
// Produce a MemorySpace attribute from a tensor type
using DefaultMemorySpaceFn =
std::function<std::optional<Attribute>(TensorType t)>;
@@ -655,7 +655,7 @@ OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
return newOp;
}
-/// Return a MemRefType to which the type of the given value can be bufferized.
+/// Return a MemRefType to which the TensorType can be bufferized.
///
/// If possible, op bufferization implementations should not use this function
/// and instead infer precise memref types for tensor results by themselves.
@@ -667,7 +667,8 @@ OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
/// Note: Canonicalization patterns could clean up layout maps and infer more
/// precise layout maps after bufferization. However, many possible
/// canonicalizations are currently not implemented.
-BaseMemRefType getMemRefType(Value value, const BufferizationOptions &options,
+BaseMemRefType getMemRefType(TensorType tensorType,
+ const BufferizationOptions &options,
MemRefLayoutAttrInterface layout = {},
Attribute memorySpace = nullptr);
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 1d6e1bdaf80f5..dd43647682ea2 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -345,10 +345,9 @@ defaultFunctionArgTypeConverter(TensorType type, Attribute memorySpace,
}
/// Default unknown type converter: Use a fully dynamic layout map.
BaseMemRefType
-defaultUnknownTypeConverter(Value value, Attribute memorySpace,
+defaultUnknownTypeConverter(TensorType tensorType, Attribute memorySpace,
const BufferizationOptions &options) {
- return getMemRefTypeWithFullyDynamicLayout(
- llvm::cast<TensorType>(value.getType()), memorySpace);
+ return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
}
} // namespace
@@ -724,7 +723,8 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
if (!memSpace.has_value())
return op->emitError("could not infer memory space");
- return getMemRefType(value, options, /*layout=*/{}, *memSpace);
+ return getMemRefType(cast<TensorType>(value.getType()), options,
+ /*layout=*/{}, *memSpace);
}
bool bufferization::hasTensorSemantics(Operation *op) {
@@ -797,12 +797,10 @@ LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
// Bufferization-specific IRMapping support with debugging.
//===----------------------------------------------------------------------===//
-BaseMemRefType bufferization::getMemRefType(Value value,
+BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
const BufferizationOptions &options,
MemRefLayoutAttrInterface layout,
Attribute memorySpace) {
- auto tensorType = llvm::cast<TensorType>(value.getType());
-
// Case 1: Unranked memref type.
if (auto unrankedTensorType =
llvm::dyn_cast<UnrankedTensorType>(tensorType)) {
@@ -819,7 +817,7 @@ BaseMemRefType bufferization::getMemRefType(Value value,
memorySpace);
}
- return options.unknownTypeConverterFn(value, memorySpace, options);
+ return options.unknownTypeConverterFn(tensorType, memorySpace, options);
}
BaseMemRefType
@@ -955,10 +953,11 @@ FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
const BufferizationState &bufferizationState,
SmallVector<Value> &invocationStack) {
assert(llvm::isa<TensorType>(value.getType()) && "expected tensor type");
+ auto tensorType = cast<TensorType>(value.getType());
// No further analysis is possible for a block argument.
if (llvm::isa<BlockArgument>(value))
- return bufferization::getMemRefType(value, options);
+ return bufferization::getMemRefType(tensorType, options);
// Value is an OpResult.
Operation *op = getOwnerOfValue(value);
@@ -981,7 +980,7 @@ FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
if (!memSpace.has_value())
return op->emitError("could not infer memory space");
- return getMemRefType(value, options, /*layout=*/{}, *memSpace);
+ return getMemRefType(tensorType, options, /*layout=*/{}, *memSpace);
}
bool bufferization::detail::defaultIsRepetitiveRegion(
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index c7681d309a4af..7e9b9119ce949 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -109,9 +109,9 @@ struct OneShotBufferizePass
"'unknown-type-conversion'");
return signalPassFailure();
}
- opt.unknownTypeConverterFn = [=](Value value, Attribute memorySpace,
+ opt.unknownTypeConverterFn = [=](TensorType tensorType,
+ Attribute memorySpace,
const BufferizationOptions &options) {
- auto tensorType = cast<TensorType>(value.getType());
if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
return bufferization::getMemRefTypeWithStaticIdentityLayout(
tensorType, memorySpace);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index a3ab53d818115..15e5102462ad7 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -223,10 +223,10 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
OneShotBufferizationOptions options;
options.bufferizeFunctionBoundaries = true;
options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
- options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
+ options.unknownTypeConverterFn = [](TensorType tensorType,
+ Attribute memorySpace,
const BufferizationOptions &options) {
- return getMemRefTypeWithStaticIdentityLayout(
- cast<TensorType>(value.getType()), memorySpace);
+ return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
};
if (analysisOnly) {
options.testAnalysisOnly = true;
|
|
@llvm/pr-subscribers-mlir Author: Andrei Golubev (andrey-golubev) ChangesGenerally, bufferization should be able to create a memref from a tensor without needing to know more than just a mlir::Type. Thus, change BufferizationOptions::UnknownTypeConverterFn to accept just a type (mlir::TensorType for now) instead of mlir::Value. Additionally, apply the same rationale to getMemRefType() helper function. Both changes are prerequisites to enable custom types support in one-shot bufferization. Full diff: https://github.com/llvm/llvm-project/pull/144658.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index adccbef754ec5..2fb795f16ae2c 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -265,9 +265,9 @@ struct BufferizationOptions {
std::function<BaseMemRefType(TensorType, Attribute memorySpace,
func::FuncOp, const BufferizationOptions &)>;
/// Tensor -> MemRef type converter.
- /// Parameters: Value, memory space, bufferization options
+ /// Parameters: tensor type, memory space, bufferization options
using UnknownTypeConverterFn = std::function<BaseMemRefType(
- Value, Attribute memorySpace, const BufferizationOptions &)>;
+ TensorType, Attribute memorySpace, const BufferizationOptions &)>;
// Produce a MemorySpace attribute from a tensor type
using DefaultMemorySpaceFn =
std::function<std::optional<Attribute>(TensorType t)>;
@@ -655,7 +655,7 @@ OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
return newOp;
}
-/// Return a MemRefType to which the type of the given value can be bufferized.
+/// Return a MemRefType to which the TensorType can be bufferized.
///
/// If possible, op bufferization implementations should not use this function
/// and instead infer precise memref types for tensor results by themselves.
@@ -667,7 +667,8 @@ OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
/// Note: Canonicalization patterns could clean up layout maps and infer more
/// precise layout maps after bufferization. However, many possible
/// canonicalizations are currently not implemented.
-BaseMemRefType getMemRefType(Value value, const BufferizationOptions &options,
+BaseMemRefType getMemRefType(TensorType tensorType,
+ const BufferizationOptions &options,
MemRefLayoutAttrInterface layout = {},
Attribute memorySpace = nullptr);
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 1d6e1bdaf80f5..dd43647682ea2 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -345,10 +345,9 @@ defaultFunctionArgTypeConverter(TensorType type, Attribute memorySpace,
}
/// Default unknown type converter: Use a fully dynamic layout map.
BaseMemRefType
-defaultUnknownTypeConverter(Value value, Attribute memorySpace,
+defaultUnknownTypeConverter(TensorType tensorType, Attribute memorySpace,
const BufferizationOptions &options) {
- return getMemRefTypeWithFullyDynamicLayout(
- llvm::cast<TensorType>(value.getType()), memorySpace);
+ return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
}
} // namespace
@@ -724,7 +723,8 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
if (!memSpace.has_value())
return op->emitError("could not infer memory space");
- return getMemRefType(value, options, /*layout=*/{}, *memSpace);
+ return getMemRefType(cast<TensorType>(value.getType()), options,
+ /*layout=*/{}, *memSpace);
}
bool bufferization::hasTensorSemantics(Operation *op) {
@@ -797,12 +797,10 @@ LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
// Bufferization-specific IRMapping support with debugging.
//===----------------------------------------------------------------------===//
-BaseMemRefType bufferization::getMemRefType(Value value,
+BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
const BufferizationOptions &options,
MemRefLayoutAttrInterface layout,
Attribute memorySpace) {
- auto tensorType = llvm::cast<TensorType>(value.getType());
-
// Case 1: Unranked memref type.
if (auto unrankedTensorType =
llvm::dyn_cast<UnrankedTensorType>(tensorType)) {
@@ -819,7 +817,7 @@ BaseMemRefType bufferization::getMemRefType(Value value,
memorySpace);
}
- return options.unknownTypeConverterFn(value, memorySpace, options);
+ return options.unknownTypeConverterFn(tensorType, memorySpace, options);
}
BaseMemRefType
@@ -955,10 +953,11 @@ FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
const BufferizationState &bufferizationState,
SmallVector<Value> &invocationStack) {
assert(llvm::isa<TensorType>(value.getType()) && "expected tensor type");
+ auto tensorType = cast<TensorType>(value.getType());
// No further analysis is possible for a block argument.
if (llvm::isa<BlockArgument>(value))
- return bufferization::getMemRefType(value, options);
+ return bufferization::getMemRefType(tensorType, options);
// Value is an OpResult.
Operation *op = getOwnerOfValue(value);
@@ -981,7 +980,7 @@ FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
if (!memSpace.has_value())
return op->emitError("could not infer memory space");
- return getMemRefType(value, options, /*layout=*/{}, *memSpace);
+ return getMemRefType(tensorType, options, /*layout=*/{}, *memSpace);
}
bool bufferization::detail::defaultIsRepetitiveRegion(
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index c7681d309a4af..7e9b9119ce949 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -109,9 +109,9 @@ struct OneShotBufferizePass
"'unknown-type-conversion'");
return signalPassFailure();
}
- opt.unknownTypeConverterFn = [=](Value value, Attribute memorySpace,
+ opt.unknownTypeConverterFn = [=](TensorType tensorType,
+ Attribute memorySpace,
const BufferizationOptions &options) {
- auto tensorType = cast<TensorType>(value.getType());
if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
return bufferization::getMemRefTypeWithStaticIdentityLayout(
tensorType, memorySpace);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index a3ab53d818115..15e5102462ad7 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -223,10 +223,10 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
OneShotBufferizationOptions options;
options.bufferizeFunctionBoundaries = true;
options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
- options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
+ options.unknownTypeConverterFn = [](TensorType tensorType,
+ Attribute memorySpace,
const BufferizationOptions &options) {
- return getMemRefTypeWithStaticIdentityLayout(
- cast<TensorType>(value.getType()), memorySpace);
+ return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
};
if (analysisOnly) {
options.testAnalysisOnly = true;
|
| @@ -265,9 +265,9 @@ struct BufferizationOptions { | |||
| std::function<BaseMemRefType(TensorType, Attribute memorySpace, | |||
| func::FuncOp, const BufferizationOptions &)>; | |||
| /// Tensor -> MemRef type converter. | |||
| /// Parameters: Value, memory space, bufferization options | |||
| /// Parameters: tensor type, memory space, bufferization options | |||
| using UnknownTypeConverterFn = std::function<BaseMemRefType( | |||
There was a problem hiding this comment.
as a bonus: this is now aligned with FunctionArgTypeConverterFn
Generally, bufferization should be able to create a memref from a tensor without needing to know more than just a mlir::Type. Thus, change BufferizationOptions::UnknownTypeConverterFn to accept just a type (mlir::TensorType for now) instead of mlir::Value. Additionally, apply the same rationale to getMemRefType() helper function. Both changes are prerequisites to enable custom types support in one-shot bufferization.
|
@matthias-springer I guess if this looks reasonable, could you merge as well? (I don't have the rights). |
…llvm#144658) Generally, bufferization should be able to create a memref from a tensor without needing to know more than just a mlir::Type. Thus, change BufferizationOptions::UnknownTypeConverterFn to accept just a type (mlir::TensorType for now) instead of mlir::Value. Additionally, apply the same rationale to getMemRefType() helper function. Both changes are prerequisites to enable custom types support in one-shot bufferization.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
…nversion (llvm#144658)" This reverts commit a1c2a71.
Is it true? In IREE, we have special logic for constants. I don't remember all the details, my guess is that we'd like to use private memory for small constants. I can try to make our project happy, but the change itself looks off to me. We name the function as The lit test failure in IREE is that we always expect identity layout for constants. Without passing the value and check if it is constant, we'll create a memref type with fully dynamic layout, while the constant is known static identity layout. Is there a way to recover the behavior? Original output: %cst_0 = arith.constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32>
%0 = bufferization.to_buffer %cst_0 : tensor<5xi32> to memref<5xi32>With the change, we always create dynamic layout for constants: %cst_0 = arith.constant dense<[1, 2, 3, 4, 5]> : tensor<5xi32>
%0 = bufferization.to_buffer %cst_0 : tensor<5xi32> to memref<5xi32, strided<[?], offset: ?>> |
|
FYI, I have a potential fix in our downstream project: iree-org/iree#21243, but it relaxes the checks. It does not only apply to constants, but all types that have static shapes. EDIT: the change is landed to IREE now; we dropped the local revert. |
I would imagine so. I mean, at least this makes sense: you get a type in, you get a type out. It is a type conversion, not a value-to-type conversion.
Thus far what we've come up with Matthias is: TensorLike + BufferLike give us custom type support, while options serve the builtin tensor -> builtin memref conversion. I guess this makes sense (I haven't seen issues but I'm only in the middle of the process with these changes) - unknown type conversion is kind of a last-mile fallback (for builtins?). Supposedly, if you're inside a custom type already (via TensorLike), you wouldn't need it?
Honestly, I am completely lost in all these layout peculiarities upstream. Our downstream does it slightly different: we generally have strides only for "subview" like operations (e.g. in tiling) and all other IR assumes "dense" (if I may) buffers. But, actually, what we do for constants is we strip the strides manually (via canonicalizer - because we have our own constant operation). Now that I think of this, perhaps, this is exactly your problem also? I mean maybe the issue is the default behaviour/implementation? I plan to look at builtin tensor -> memref conversion as well and make sure tensor encoding gets correctly mapped to memref layout. Perhaps it makes sense to revisit what should be done w.r.t. dynamic layouts to solve this issue for both of us? |
The idea seems okay to me now. I don't have a fresh memory about bufferization now, but what you said and the comments in the codebase makes sense.
Yeah, I think the main difference is that you have your own constant in your project, and you can define your own canonicalization patterns to achieve this with your assumptions. IREE uses upstream dialects, e.g., arith/spir-v/etc, and it is very stable. Maybe the issue is in the default behavior in upstream. Or maybe IREE should evolve to the next phase. Again, it has been stable for a long time and I don't have much bandwidth to review it as it touches many components and few different backends, so it might not happen in the near future. If you identify something in the default behavior, I'm happy to learn about it. Thanks for your sharing! 🙂 |
Generally, bufferization should be able to create a memref from a tensor without needing to know more than just a mlir::Type. Thus, change BufferizationOptions::UnknownTypeConverterFn to accept just a type (mlir::TensorType for now) instead of mlir::Value. Additionally, apply the same rationale to getMemRefType() helper function.
Both changes are prerequisites to enable custom types support in one-shot bufferization.