Context
There are currently two incompatible dialect conversion drivers is MLIR: the “default” driver and the 1:N dialect conversion driver. The 1:N driver was added some time ago to address limitations in the default driver. E.g., replacing an op that produces 1 SSA value with multiple SSA values is not supported by the default driver. The 1:N driver has uses mainly in downstream projects.
This RFC proposes:
- Extending the default driver with full 1:N conversion support.
- Deleting the 1:N driver.
The purpose of this change is to remove the duplicate dialect conversion infrastructure. Another goal is to clean up the API (in particular, remove argument materializations) before continuing with the One-Shot Dialect Conversion effort.
1:N Conversions in the Default Driver
The default driver has limited support for 1:N conversions: The type converter API (convertType and materializations) supports it. There is also an API to replace a block argument with multiple replacement values (during applySignatureConversion). However, the remaining infrastructure (in particular the ConversionValueMapping and ConversionPattern API) does not support it. To support 1:N replacements for block arguments, the default driver uses “argument materializations” as a workaround: these convert the N replacement values back into a single SSA value that can be stored in the mapping and that can be passed to adaptors.
Implementation Outline
The following changes are required to support 1:N conversion support in the default driver:
- Minor extensions to the
TypeConverterAPI. The argument materializations carry additional information that must be passed to target materializations. For details, see [mlir][Transforms] Dialect conversion: add `originalType` param to materializations by matthias-springer · Pull Request #112128 · llvm/llvm-project · GitHub. - Minor extension to the
OpAdaptorC++ code generation, so that a 1:1 adaptor can be constructed from a 1:N adaptor. For details, see [mlir][tblgen] Add additional constructor to Adaptor class by matthias-springer · Pull Request #112144 · llvm/llvm-project · GitHub. ConversionValueMapping: StoreValue→ValueRangeinstead ofValue→Value. That’s because an SSA value can now be replaced with multiple SSA values.- Provide an additional 1:N
ConversionPatternAPI, where each adaptor method now returns aValueRangeinstead of aValue. - Provide an additional
replaceOpAPI that supports 1:N replacements. - Remove argument materializations entirely. They are no longer needed.
- Remove other workarounds that were needed because of 1:N limitations. A typical workaround is "manually packing/unpacking replacement values in an
unrealized_conversion_cast. (E.g., there’s one in the sparse compiler.)
ConversionValueMapping Data Structure
The ConversionValueMapping keeps track of replacements of SSA values that have not materialized in IR yet. Until now, this essentially used to be an IRMapping. To support 1:N mappings, a data structure such as the following is needed:
struct ConversionValueMapping {
/// Replacement mapping: Value -> ValueRange
DenseMap<Value, SmallVector<Value, 1>> mapping;
/// Materializations: ValueRange -> [ValueRange]
DenseMap<SmallVector<Value, 1>,
SmallVector<SmallVector<Value, 1>>,
SmallVectorMapInfo>
materializations;
};
The materializations map stores unresolved materializations. In the current implementation, it is combined with IRMapping. (Which I believe is incorrect in some cases, because we cannot distinguish between actual replacements and type conversions of replaced values.) Materializations can occur in the following shapes:
1:1: This is the typical case in a 1:1 dialect conversion.1:N: A target materialization could be a 1:N conversion. (1 type is legalized to N types.)N:1: A source materialization could convert a legalized type range (N types) back to a single type.N:M: E.g.: a single value of typetwas replaced with N legalized types. Now a subsequent pattern has a different type converter that legalizestto M different types. Now anN:Mconversion is inserted.
Furthermore, there can be multiple materializations for the same type/type range. That’s why the key of the map is a ValueRange and the value of the map is a list of ValueRanges.
Note: To optimize for the more common case of 1:1 conversions, all ValueRanges are actually SmallVector<Value, 1>.
Note: With the One-Shot Dialect Conversion driver, the ConversionValueMapping will be deleted. However, I cannot do all of that in one step. I have to break things down into smaller refactoring to make it possible for downstream users to adapt to the changes step-by-step.
ConversionPattern API
To support N replacement values, the adaptor must now return ValueRange instead of Value. The auto-generated adaptors already allow for custom value types (via C++ templates). I propose to add a second matchAndRewrite overload that supports 1:N replacements via their adaptors.
class ConversionPattern {
public:
// Current 1:1 API. Will remain unchanged.
virtual LogicalResult
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const;
// New 1:N API.
virtual LogicalResult
matchAndRewrite(Operation *op, ArrayRef<ArrayRef<Value>> operands,
ConversionPatternRewriter &rewriter) const {
// By default, dispatches to 1:1 API.
return matchAndRewrite(op, getOneToOneAdaptorOperands(operands), rewriter);
}
private:
// Unpack 1:1 replacements. Fail if 1:N replacements are found.
SmallVector<Value> getOneToOneAdaptorOperands(
ArrayRef<ArrayRef<Value>> operands) const {
SmallVector<Value> oneToOneOperands;
oneToOneOperands.reserve(operands.size());
for (ArrayRef<Value> operand : operands) {
if (operand.size() != 1)
llvm::report_fatal_error("pattern '" + getDebugName() +
"' does not support 1:N conversion");
oneToOneOperands.push_back(operand.front());
}
return oneToOneOperands;
}
}
The dialect conversion driver always invokes the 1:N matchAndRewrite entry point of patterns. That’s because it fills the adaptor with mappings from the ConversionValueMapping, which are stored as 1:N mappings.
If an existing 1:1 pattern is run in a 1:1 conversion, nothing changes. However, if the type converter returns a 1:N mapping and the pattern does not implement the 1:N entry point a fatal error is triggered.
(Similar changes for OpConversionPattern etc. See prototype for details.)
ConversionPatternRewriter API
The rewriter API remains mostly unchanged. Just one new function is needed:
void ConversionPatternRewriter::replaceOpWithMultiple(
Operation *op, ArrayRef<ValueRange> newValues);
I tried overloading replaceOp but ran into various cases where existing code no longer compiles due to ambiguous overload resolution.
Changes for Downstream Users
If you are a users of the 1:N dialect conversion driver, you have to update your patterns. (The 1:N driver does not use ConversionPattern.) This can be done gradually: the 1:N driver can stay around for a while after 1:N support was added to the default driver.
If you are using argument materializations, you have to incorporate the functionality into target materializations. That’s because argument materializations will be removed.
If you are using a type converter that has 1:N type conversions, you are exercising a code path in the dialect conversion that is incomplete (example). You may have to update some existing 1:1 conversion patterns to 1:N conversion patterns, as those will start receiving 1:N replacements.
Status
I am working a prototype here: [mlir][draft] Support 1:N dialect conversion by matthias-springer · Pull Request #112141 · llvm/llvm-project · GitHub. It is still in a very rough state (not ready for review yet) and not all tests are passing yet. In particular, the MemRef → LLVM conversions are not working yet.
Over the next weeks, I’d like to gather feedback about the design described above. I hope to have a working and polished prototype some time after the LLVM conference.
@ingomueller-net @zero9178 @ftynse @mehdi_amini @River707 @jpienaar