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 _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
