[llvm-branch-commits] [mlir] [mlir-c] Add 1:N TypeConverter conversion and materialization bindings (PR #208935)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Fri Jul 17 14:03:39 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Maksim Levental (makslevental)
<details>
<summary>Changes</summary>
Builds on the source/target materialization C bindings (#<!-- -->208934) to expose the 1:N dialect-conversion functionality through the MLIR C API. Continues the buildout of the dialect-conversion C bindings (follows #<!-- -->206146 and #<!-- -->206161).
Assisted by: Claude
---
Patch is 55.01 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/208935.diff
4 Files Affected:
- (modified) mlir/include/mlir-c/Rewrite.h (+87-5)
- (modified) mlir/lib/Bindings/Python/Rewrite.cpp (+3-1)
- (modified) mlir/lib/CAPI/Transforms/Rewrite.cpp (+141-14)
- (modified) mlir/test/CAPI/rewrite.c (+734-4)
``````````diff
diff --git a/mlir/include/mlir-c/Rewrite.h b/mlir/include/mlir-c/Rewrite.h
index e074fc96f2aeb..968690f79c321 100644
--- a/mlir/include/mlir-c/Rewrite.h
+++ b/mlir/include/mlir-c/Rewrite.h
@@ -501,6 +501,15 @@ mlirConversionPatternRewriterConvertRegionTypes(
MlirConversionPatternRewriter rewriter, MlirRegion region,
MlirTypeConverter typeConverter);
+/// Replace the given operation with multiple value ranges -- one range per
+/// result of `op` -- and erase it. `nRanges` must equal the number of results
+/// of `op`. `rangeSizes[i]` is the number of values in the i-th range, and
+/// `values` is the flat concatenation of all ranges (its length is the sum of
+/// `rangeSizes[0..nRanges)`).
+MLIR_CAPI_EXPORTED void mlirConversionPatternRewriterReplaceOpWithMultiple(
+ MlirConversionPatternRewriter rewriter, MlirOperation op, intptr_t nRanges,
+ intptr_t *rangeSizes, MlirValue *values);
+
//===----------------------------------------------------------------------===//
/// ConversionTarget API
//===----------------------------------------------------------------------===//
@@ -601,32 +610,93 @@ mlirTypeConverterAddConversion(MlirTypeConverter typeConverter,
MlirTypeConverterConversionCallback convertType,
void *userData);
-/// Convert the given type using the given TypeConverter.
+/// Opaque accumulator for the result types of a 1:N type conversion. It is
+/// passed to a MlirTypeConverter1ToNConversionCallback, which appends converted
+/// types to it via mlirTypeConverterConversionResultsAppend.
+typedef struct MlirTypeConverterConversionResults {
+ void *ptr;
+} MlirTypeConverterConversionResults;
+
+/// Append a converted result type to the given 1:N conversion result
+/// accumulator.
+MLIR_CAPI_EXPORTED void mlirTypeConverterConversionResultsAppend(
+ MlirTypeConverterConversionResults results, MlirType type);
+
+/// Callback type for 1:N type conversion functions. For the given `type`, the
+/// callback appends zero or more converted result types to `results` (via
+/// mlirTypeConverterConversionResultsAppend) and returns success. Returning
+/// failure leaves the type unconverted and allows another conversion function
+/// to be tried; any types appended before returning failure are discarded.
+/// Appending a single type is a 1:1 conversion; appending several is a 1:N
+/// conversion; appending none (on success) erases the type.
+typedef MlirLogicalResult (*MlirTypeConverter1ToNConversionCallback)(
+ MlirType type, MlirTypeConverterConversionResults results, void *userData);
+
+/// Add a 1:N type conversion function to the given TypeConverter.
+MLIR_CAPI_EXPORTED void mlirTypeConverterAdd1ToNConversion(
+ MlirTypeConverter typeConverter,
+ MlirTypeConverter1ToNConversionCallback convertType, void *userData);
+
+/// Convert the given type using the given TypeConverter. This is the 1:1
+/// convenience form: it returns the single converted type, or a null MlirType
+/// on failure or if the type converts to anything other than exactly one type
+/// (e.g. a 1:N conversion registered via mlirTypeConverterAdd1ToNConversion, or
+/// an erasure to zero types).
MLIR_CAPI_EXPORTED MlirType
mlirTypeConverterConvertType(MlirTypeConverter typeConverter, MlirType type);
-/// Callback type for type materializations. Given a builder (passed as a
+/// Callback type for source materializations. Given a builder (passed as a
/// rewriter), the desired output type, the input values, and a location, the
/// callback must build a cast-like operation that produces a single value of
/// `outputType` and return it. Returning a null MlirValue indicates failure, in
/// which case another registered materialization may be attempted.
-typedef MlirValue (*MlirTypeConverterMaterializationCallback)(
+typedef MlirValue (*MlirTypeConverterSourceMaterializationCallback)(
MlirRewriterBase rewriter, MlirType outputType, intptr_t nInputs,
MlirValue *inputs, MlirLocation loc, void *userData);
+/// Callback type for 1:1 target materializations. Behaves like
+/// MlirTypeConverterSourceMaterializationCallback, but additionally receives
+/// `originalType`: the original type of the SSA value being materialized.
+///
+/// Note: This callback is single-output. For the 1:N (multiple-output) form,
+/// use MlirTypeConverter1ToNTargetMaterializationCallback.
+typedef MlirValue (*MlirTypeConverterTargetMaterializationCallback)(
+ MlirRewriterBase rewriter, MlirType outputType, intptr_t nInputs,
+ MlirValue *inputs, MlirLocation loc, MlirType originalType, void *userData);
+
/// Register a source materialization with the given TypeConverter. This is
/// invoked when a replacement value must be converted back to its original
/// source type because some uses persist beyond the main conversion.
MLIR_CAPI_EXPORTED void mlirTypeConverterAddSourceMaterialization(
MlirTypeConverter typeConverter,
- MlirTypeConverterMaterializationCallback callback, void *userData);
+ MlirTypeConverterSourceMaterializationCallback callback, void *userData);
/// Register a target materialization with the given TypeConverter. This is
/// invoked when a value must be converted to a target type according to a
/// pattern's type converter.
MLIR_CAPI_EXPORTED void mlirTypeConverterAddTargetMaterialization(
MlirTypeConverter typeConverter,
- MlirTypeConverterMaterializationCallback callback, void *userData);
+ MlirTypeConverterTargetMaterializationCallback callback, void *userData);
+
+/// Callback type for 1:N target materializations. Like
+/// MlirTypeConverterTargetMaterializationCallback, but produces a value for
+/// each of the `nOutputTypes` requested output types instead of a single value.
+/// The callback must fill `outputs` -- a caller-allocated array of length
+/// `nOutputTypes` -- with that many values and return success. Returning
+/// failure signals that this materialization declined (so another may be
+/// attempted); in that case `outputs` is ignored. A success that leaves any of
+/// `outputs` null is likewise treated as a decline. `originalType` carries the
+/// original type of the value being materialized and may be a null MlirType.
+typedef MlirLogicalResult (*MlirTypeConverter1ToNTargetMaterializationCallback)(
+ MlirRewriterBase rewriter, intptr_t nOutputTypes, MlirType *outputTypes,
+ intptr_t nInputs, MlirValue *inputs, MlirLocation loc,
+ MlirType originalType, MlirValue *outputs, void *userData);
+
+/// Register a 1:N target materialization with the given TypeConverter.
+MLIR_CAPI_EXPORTED void mlirTypeConverterAdd1ToNTargetMaterialization(
+ MlirTypeConverter typeConverter,
+ MlirTypeConverter1ToNTargetMaterializationCallback callback,
+ void *userData);
//===----------------------------------------------------------------------===//
/// ConversionPattern API
@@ -647,6 +717,18 @@ typedef struct {
MlirValue *operands,
MlirConversionPatternRewriter rewriter,
void *userData);
+ /// Optional callback corresponding to the 1:N
+ /// ConversionPattern::matchAndRewrite(Operation *, ArrayRef<ValueRange>, ...)
+ /// overload, used when one or more operands are remapped to several values
+ /// (e.g. under a 1:N type conversion). `operands` is the flat concatenation
+ /// of all operand ranges; there are `nRanges` ranges (one per original
+ /// operand) and `rangeSizes[i]` is the number of values in the i-th range.
+ /// When this is non-null it takes precedence; when null, the driver falls
+ /// back to the 1:1 `matchAndRewrite` above.
+ MlirLogicalResult (*matchAndRewrite1ToN)(
+ MlirConversionPattern pattern, MlirOperation op, intptr_t nRanges,
+ intptr_t *rangeSizes, intptr_t nOperands, MlirValue *operands,
+ MlirConversionPatternRewriter rewriter, void *userData);
} MlirConversionPatternCallbacks;
/// Create a conversion pattern that matches the operation with the given
diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp
index 750c7b1e4b1d8..2a9b0ea1c218d 100644
--- a/mlir/lib/Bindings/Python/Rewrite.cpp
+++ b/mlir/lib/Bindings/Python/Rewrite.cpp
@@ -226,7 +226,9 @@ void PyRewritePatternSet::addConversion(nb::handle root,
std::string opName = operationNameFromObject(root);
MlirStringRef rootName = mlirStringRefCreate(opName.data(), opName.size());
- MlirConversionPatternCallbacks callbacks;
+ // Value-initialize so optional callbacks (e.g. matchAndRewrite1ToN) default
+ // to null rather than an indeterminate pointer.
+ MlirConversionPatternCallbacks callbacks{};
callbacks.construct = [](void *userData) {
nb::handle(static_cast<PyObject *>(userData)).inc_ref();
};
diff --git a/mlir/lib/CAPI/Transforms/Rewrite.cpp b/mlir/lib/CAPI/Transforms/Rewrite.cpp
index 92456d6d8a435..5f0da4776e7bd 100644
--- a/mlir/lib/CAPI/Transforms/Rewrite.cpp
+++ b/mlir/lib/CAPI/Transforms/Rewrite.cpp
@@ -543,6 +543,22 @@ MlirLogicalResult mlirConversionPatternRewriterConvertRegionTypes(
*unwrap(typeConverter)));
}
+void mlirConversionPatternRewriterReplaceOpWithMultiple(
+ MlirConversionPatternRewriter rewriter, MlirOperation op, intptr_t nRanges,
+ intptr_t *rangeSizes, MlirValue *values) {
+ SmallVector<SmallVector<Value>> ranges;
+ ranges.reserve(nRanges);
+ MlirValue *cur = values;
+ for (intptr_t i = 0; i < nRanges; ++i) {
+ SmallVector<Value> range;
+ range.reserve(rangeSizes[i]);
+ for (intptr_t j = 0; j < rangeSizes[i]; ++j, ++cur)
+ range.push_back(unwrap(*cur));
+ ranges.push_back(std::move(range));
+ }
+ unwrap(rewriter)->replaceOpWithMultiple(unwrap(op), std::move(ranges));
+}
+
//===----------------------------------------------------------------------===//
/// ConversionTarget API
//===----------------------------------------------------------------------===//
@@ -664,25 +680,53 @@ void mlirTypeConverterAddConversion(
});
}
+void mlirTypeConverterConversionResultsAppend(
+ MlirTypeConverterConversionResults results, MlirType type) {
+ static_cast<SmallVectorImpl<Type> *>(results.ptr)->push_back(unwrap(type));
+}
+
+void mlirTypeConverterAdd1ToNConversion(
+ MlirTypeConverter typeConverter,
+ MlirTypeConverter1ToNConversionCallback convertType, void *userData) {
+ unwrap(typeConverter)
+ ->addConversion(
+ [convertType, userData](Type type, SmallVectorImpl<Type> &results)
+ -> std::optional<LogicalResult> {
+ size_t numPriorResults = results.size();
+ MlirTypeConverterConversionResults wrappedResults{&results};
+ MlirLogicalResult result =
+ convertType(wrap(type), wrappedResults, userData);
+ if (mlirLogicalResultIsFailure(result)) {
+ // The callback declined. Restore any types it appended so the
+ // driver's "try the next conversion" invariant holds (a declining
+ // conversion function must not mutate `results`).
+ results.truncate(numPriorResults);
+ return std::nullopt;
+ }
+ return success();
+ });
+}
+
MlirType mlirTypeConverterConvertType(MlirTypeConverter typeConverter,
MlirType type) {
return wrap(unwrap(typeConverter)->convertType(unwrap(type)));
}
namespace {
-/// Wraps a C materialization callback as a C++ materialization callback of the
-/// form `Value(OpBuilder &, Type, ValueRange, Location)`, shared by both source
-/// and target materializations. The builder is always a RewriterBase in the
-/// conversion driver, so it is safe to expose it as an MlirRewriterBase.
+SmallVector<MlirValue> wrapInputs(ValueRange inputs) {
+ SmallVector<MlirValue> wrappedInputs;
+ wrappedInputs.reserve(inputs.size());
+ for (Value v : inputs)
+ wrappedInputs.push_back(wrap(v));
+ return wrappedInputs;
+}
+
std::function<Value(OpBuilder &, Type, ValueRange, Location)>
-wrapMaterializationCallback(MlirTypeConverterMaterializationCallback callback,
- void *userData) {
+wrapSourceMaterializationCallback(
+ MlirTypeConverterSourceMaterializationCallback callback, void *userData) {
return [callback, userData](OpBuilder &builder, Type type, ValueRange inputs,
Location loc) -> Value {
- SmallVector<MlirValue> wrappedInputs;
- wrappedInputs.reserve(inputs.size());
- for (Value v : inputs)
- wrappedInputs.push_back(wrap(v));
+ SmallVector<MlirValue> wrappedInputs = wrapInputs(inputs);
MlirValue result =
callback(wrap(static_cast<RewriterBase *>(&builder)), wrap(type),
static_cast<intptr_t>(wrappedInputs.size()),
@@ -690,24 +734,85 @@ wrapMaterializationCallback(MlirTypeConverterMaterializationCallback callback,
return mlirValueIsNull(result) ? Value() : unwrap(result);
};
}
+
+std::function<Value(OpBuilder &, Type, ValueRange, Location, Type)>
+wrapTargetMaterializationCallback(
+ MlirTypeConverterTargetMaterializationCallback callback, void *userData) {
+ return [callback, userData](OpBuilder &builder, Type type, ValueRange inputs,
+ Location loc, Type originalType) -> Value {
+ SmallVector<MlirValue> wrappedInputs = wrapInputs(inputs);
+ MlirValue result =
+ callback(wrap(static_cast<RewriterBase *>(&builder)), wrap(type),
+ static_cast<intptr_t>(wrappedInputs.size()),
+ wrappedInputs.data(), wrap(loc), wrap(originalType), userData);
+ return mlirValueIsNull(result) ? Value() : unwrap(result);
+ };
+}
+
+std::function<SmallVector<Value>(OpBuilder &, TypeRange, ValueRange, Location,
+ Type)>
+wrap1ToNTargetMaterializationCallback(
+ MlirTypeConverter1ToNTargetMaterializationCallback callback,
+ void *userData) {
+ return [callback, userData](OpBuilder &builder, TypeRange outputTypes,
+ ValueRange inputs, Location loc,
+ Type originalType) -> SmallVector<Value> {
+ SmallVector<MlirType> wrappedOutputTypes;
+ wrappedOutputTypes.reserve(outputTypes.size());
+ for (Type t : outputTypes)
+ wrappedOutputTypes.push_back(wrap(t));
+ SmallVector<MlirValue> wrappedInputs = wrapInputs(inputs);
+ SmallVector<MlirValue> wrappedOutputs(outputTypes.size(),
+ MlirValue{nullptr});
+ MlirLogicalResult result = callback(
+ wrap(static_cast<RewriterBase *>(&builder)),
+ static_cast<intptr_t>(wrappedOutputTypes.size()),
+ wrappedOutputTypes.data(), static_cast<intptr_t>(wrappedInputs.size()),
+ wrappedInputs.data(), wrap(loc), wrap(originalType),
+ wrappedOutputs.data(), userData);
+ if (mlirLogicalResultIsFailure(result))
+ return {}; // declined; another materialization may be attempted
+ SmallVector<Value> outputs;
+ outputs.reserve(wrappedOutputs.size());
+ for (MlirValue v : wrappedOutputs) {
+ // A success that left any output unfilled is an incomplete
+ // materialization. Treat it as a decline rather than returning a
+ // null-containing result (which would trip an assert in the driver).
+ if (mlirValueIsNull(v))
+ return {};
+ outputs.push_back(unwrap(v));
+ }
+ return outputs;
+ };
+}
} // namespace
void mlirTypeConverterAddSourceMaterialization(
MlirTypeConverter typeConverter,
- MlirTypeConverterMaterializationCallback callback, void *userData) {
+ MlirTypeConverterSourceMaterializationCallback callback, void *userData) {
assert(callback && "expected non-null materialization callback");
unwrap(typeConverter)
->addSourceMaterialization(
- wrapMaterializationCallback(callback, userData));
+ wrapSourceMaterializationCallback(callback, userData));
}
void mlirTypeConverterAddTargetMaterialization(
MlirTypeConverter typeConverter,
- MlirTypeConverterMaterializationCallback callback, void *userData) {
+ MlirTypeConverterTargetMaterializationCallback callback, void *userData) {
+ assert(callback && "expected non-null materialization callback");
+ unwrap(typeConverter)
+ ->addTargetMaterialization(
+ wrapTargetMaterializationCallback(callback, userData));
+}
+
+void mlirTypeConverterAdd1ToNTargetMaterialization(
+ MlirTypeConverter typeConverter,
+ MlirTypeConverter1ToNTargetMaterializationCallback callback,
+ void *userData) {
assert(callback && "expected non-null materialization callback");
unwrap(typeConverter)
->addTargetMaterialization(
- wrapMaterializationCallback(callback, userData));
+ wrap1ToNTargetMaterializationCallback(callback, userData));
}
//===----------------------------------------------------------------------===//
@@ -747,6 +852,28 @@ class ExternalConversionPattern : public mlir::ConversionPattern {
userData));
}
+ LogicalResult
+ matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Without a 1:N callback, defer to the default behavior, which dispatches
+ // to the 1:1 matchAndRewrite above or fails to match on a 1:N mapping.
+ if (!callbacks.matchAndRewrite1ToN)
+ return dispatchTo1To1(*this, op, operands, rewriter);
+ SmallVector<intptr_t> rangeSizes;
+ rangeSizes.reserve(operands.size());
+ std::vector<MlirValue> wrappedOperands;
+ for (ValueRange range : operands) {
+ rangeSizes.push_back(static_cast<intptr_t>(range.size()));
+ for (Value val : range)
+ wrappedOperands.push_back(wrap(val));
+ }
+ return unwrap(callbacks.matchAndRewrite1ToN(
+ wrap(static_cast<const mlir::ConversionPattern *>(this)), wrap(op),
+ static_cast<intptr_t>(rangeSizes.size()), rangeSizes.data(),
+ static_cast<intptr_t>(wrappedOperands.size()), wrappedOperands.data(),
+ wrap(&rewriter), userData));
+ }
+
private:
MlirConversionPatternCallbacks callbacks;
void *userData;
diff --git a/mlir/test/CAPI/rewrite.c b/mlir/test/CAPI/rewrite.c
index b10afbd8b812c..16e40753c59dc 100644
--- a/mlir/test/CAPI/rewrite.c
+++ b/mlir/test/CAPI/rewrite.c
@@ -833,7 +833,178 @@ static MlirValue buildCastMaterialization(MlirRewriterBase rewriter,
return mlirOperationGetResult(castOp, 0);
}
-// Conversion pattern for `test.source`: replaces it with a `test.source_i64`
+// Source materialization callback that always declines (returns a null value)
+// and records that it was consulted. Used to exercise the "this materialization
+// declined, try the next one" fallback path.
+static MlirValue
+declineSourceMaterialization(MlirRewriterBase rewriter, MlirType outputType,
+ intptr_t nInputs, MlirValue *inputs,
+ MlirLocation loc, void *userData) {
+ (void)rewriter;
+ (void)outputType;
+ (void)nInputs;
+ (void)inputs;
+ (void)loc;
+ intptr_t *declined = (intptr_t *)userData;
+ if (declined)
+ ++(*declined);
+ return (MlirValue){NULL};
+}
+
+// Source materialization callback that records the number of inputs it was
+// invoked with (into userData) before building the `test.cast`. Used to verify
+// that a 1:N replacement drives a source materialization with nInputs > 1.
+static MlirValue buildSourceCastRecordInputs(MlirRewriterBase rewriter,
+ MlirType outputType,
+ intptr_t nInputs,
+ MlirValue *inputs,
+ MlirLocation loc, void *userData) {
+ intptr_t *observedInputs = (intptr_t *)userData;
+ if (observedInputs)
+ *observedInputs = nInputs;
+ MlirOperationState state =
+ mlirOperationStateGet(mlirStringRefCreateFromCString("test.cast"), loc);
+ mlirOperationStateAddOperands(&state, nInputs, inputs);
+ mlirOperationStateAddResults(&state, 1, &outputType);
+ MlirOperation castOp = mlirOperationCreate(&state);
+ mlirRewriterBaseInsert(rewriter, castOp);
+ return mlirOperationGetResult(castOp, 0);
+}
+
+// 1:1 target materialization callback. Builds a `test.cast` like
+// ...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/208935
More information about the llvm-branch-commits
mailing list