llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT--> @llvm/pr-subscribers-mlir-core Author: Matthias Springer (matthias-springer) <details> <summary>Changes</summary> This commit adds a new function `ConversionPatternRewriter::replaceOpWithMultiple`. This function is similar to `replaceOp`, but it accepts multiple `ValueRange` replacements, one per op result. Note: This function is not an overload of `replaceOp` because of ambiguous overload resolution that would make the API difficult to use. This commit aligns "block signature conversions" with "op replacements": both support 1:N replacements now. Due to incomplete 1:N support in the dialect conversion driver, an argument materialization is inserted when an SSA value is replaced with multiple values; same as block signature conversions already work around the problem. These argument materializations are going to be removed in a subsequent commit that adds full 1:N support. The purpose of this PR is to add missing features gradually in small increments. This commit also updates two MLIR transformations that have their custom workarounds around missing 1:N support. These can already start using `replaceOpWithMultiple`. Depends on #<!-- -->114940. --- Patch is 29.03 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/115816.diff 7 Files Affected: - (modified) mlir/include/mlir/IR/Builders.h (+1-1) - (modified) mlir/include/mlir/Transforms/DialectConversion.h (+22-2) - (modified) mlir/lib/Dialect/Func/Transforms/DecomposeCallGraphTypes.cpp (+12-28) - (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp (+25-23) - (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp (+13-8) - (modified) mlir/lib/IR/Builders.cpp (+11-5) - (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+102-51) ``````````diff diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h index 7ef03b87179523..78729376507208 100644 --- a/mlir/include/mlir/IR/Builders.h +++ b/mlir/include/mlir/IR/Builders.h @@ -353,7 +353,7 @@ class OpBuilder : public Builder { /// selected insertion point. (E.g., because they are defined in a nested /// region or because they are not visible in an IsolatedFromAbove region.) static InsertPoint after(ArrayRef<Value> values, - const PostDominanceInfo &domInfo); + const PostDominanceInfo *domInfo = nullptr); /// Returns true if this insert point is set. bool isSet() const { return (block != nullptr); } diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h index 5e5957170e646c..e461b7d11602a0 100644 --- a/mlir/include/mlir/Transforms/DialectConversion.h +++ b/mlir/include/mlir/Transforms/DialectConversion.h @@ -795,12 +795,32 @@ class ConversionPatternRewriter final : public PatternRewriter { /// patterns even if a failure is encountered during the rewrite step. bool canRecoverFromRewriteFailure() const override { return true; } - /// PatternRewriter hook for replacing an operation. + /// Replace the given operation with the new values. The number of op results + /// and replacement values must match. The types may differ: the dialect + /// conversion driver will reconcile any surviving type mismatches at the end + /// of the conversion process with source materializations. The given + /// operation is erased. void replaceOp(Operation *op, ValueRange newValues) override; - /// PatternRewriter hook for replacing an operation. + /// Replace the given operation with the results of the new op. The number of + /// op results must match. The types may differ: the dialect conversion + /// driver will reconcile any surviving type mismatches at the end of the + /// conversion process with source materializations. The original operation + /// is erased. void replaceOp(Operation *op, Operation *newOp) override; + /// Replace the given operation with the new value range. The number of op + /// results and value ranges must match. If an original SSA value is replaced + /// by multiple SSA values (i.e., value range has more than 1 element), the + /// conversion driver will insert an argument materialization to convert the + /// N SSA values back into 1 SSA value of the original type. The given + /// operation is erased. + /// + /// Note: The argument materialization is a workaround until we have full 1:N + /// support in the dialect conversion. (It is going to disappear from both + /// `replaceOpWithMultiple` and `applySignatureConversion`.) + void replaceOpWithMultiple(Operation *op, ArrayRef<ValueRange> newValues); + /// PatternRewriter hook for erasing a dead operation. The uses of this /// operation *must* be made dead by the end of the conversion process, /// otherwise an assert will be issued. diff --git a/mlir/lib/Dialect/Func/Transforms/DecomposeCallGraphTypes.cpp b/mlir/lib/Dialect/Func/Transforms/DecomposeCallGraphTypes.cpp index de4aba2ed327db..a08764326a80b6 100644 --- a/mlir/lib/Dialect/Func/Transforms/DecomposeCallGraphTypes.cpp +++ b/mlir/lib/Dialect/Func/Transforms/DecomposeCallGraphTypes.cpp @@ -141,47 +141,31 @@ struct DecomposeCallGraphTypesForCallOp : public OpConversionPattern<CallOp> { getTypeConverter())); } - // Create the new result types for the new `CallOp` and track the indices in - // the new call op's results that correspond to the old call op's results. - // - // expandedResultIndices[i] = "list of new result indices that old result i - // expanded to". + // Create the new result types for the new `CallOp` and track the number of + // replacement types for each original op result. SmallVector<Type, 2> newResultTypes; - SmallVector<SmallVector<unsigned, 2>, 4> expandedResultIndices; + SmallVector<unsigned> expandedResultSizes; for (Type resultType : op.getResultTypes()) { unsigned oldSize = newResultTypes.size(); if (failed(typeConverter->convertType(resultType, newResultTypes))) return failure(); - auto &resultMapping = expandedResultIndices.emplace_back(); - for (unsigned i = oldSize, e = newResultTypes.size(); i < e; i++) - resultMapping.push_back(i); + expandedResultSizes.push_back(newResultTypes.size() - oldSize); } CallOp newCallOp = rewriter.create<CallOp>(op.getLoc(), op.getCalleeAttr(), newResultTypes, newOperands); - // Build a replacement value for each result to replace its uses. If a - // result has multiple mapping values, it needs to be materialized as a - // single value. - SmallVector<Value, 2> replacedValues; + // Build a replacement value for each result to replace its uses. + SmallVector<ValueRange> replacedValues; replacedValues.reserve(op.getNumResults()); + unsigned startIdx = 0; for (unsigned i = 0, e = op.getNumResults(); i < e; ++i) { - auto decomposedValues = llvm::to_vector<6>( - llvm::map_range(expandedResultIndices[i], - [&](unsigned i) { return newCallOp.getResult(i); })); - if (decomposedValues.empty()) { - // No replacement is required. - replacedValues.push_back(nullptr); - } else if (decomposedValues.size() == 1) { - replacedValues.push_back(decomposedValues.front()); - } else { - // Materialize a single Value to replace the original Value. - Value materialized = getTypeConverter()->materializeArgumentConversion( - rewriter, op.getLoc(), op.getType(i), decomposedValues); - replacedValues.push_back(materialized); - } + ValueRange repl = + newCallOp.getResults().slice(startIdx, expandedResultSizes[i]); + replacedValues.push_back(repl); + startIdx += expandedResultSizes[i]; } - rewriter.replaceOp(op, replacedValues); + rewriter.replaceOpWithMultiple(op, replacedValues); return success(); } }; diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp index 062a0ea6cc47cb..09509278d7749a 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp @@ -600,8 +600,8 @@ class SparseCallConverter : public OpConversionPattern<func::CallOp> { flattenOperands(adaptor.getOperands(), flattened); auto newCall = rewriter.create<func::CallOp>(loc, op.getCallee(), finalRetTy, flattened); - // (2) Create cast operation for sparse tensor returns. - SmallVector<Value> castedRet; + // (2) Gather sparse tensor returns. + SmallVector<SmallVector<Value>> packedResultVals; // Tracks the offset of current return value (of the original call) // relative to the new call (after sparse tensor flattening); unsigned retOffset = 0; @@ -618,21 +618,27 @@ class SparseCallConverter : public OpConversionPattern<func::CallOp> { assert(!sparseFlat.empty()); if (sparseFlat.size() > 1) { auto flatSize = sparseFlat.size(); - ValueRange fields(iterator_range<ResultRange::iterator>( - newCall.result_begin() + retOffset, - newCall.result_begin() + retOffset + flatSize)); - castedRet.push_back(genTuple(rewriter, loc, retType, fields)); + packedResultVals.push_back(SmallVector<Value>()); + llvm::append_range(packedResultVals.back(), + iterator_range<ResultRange::iterator>( + newCall.result_begin() + retOffset, + newCall.result_begin() + retOffset + flatSize)); retOffset += flatSize; } else { // If this is an 1:1 conversion, no need for casting. - castedRet.push_back(newCall.getResult(retOffset)); + packedResultVals.push_back(SmallVector<Value>()); + packedResultVals.back().push_back(newCall.getResult(retOffset)); retOffset++; } sparseFlat.clear(); } - assert(castedRet.size() == op.getNumResults()); - rewriter.replaceOp(op, castedRet); + assert(packedResultVals.size() == op.getNumResults()); + SmallVector<ValueRange> ranges; + ranges.reserve(packedResultVals.size()); + for (const SmallVector<Value> &vec : packedResultVals) + ranges.push_back(ValueRange(vec)); + rewriter.replaceOpWithMultiple(op, ranges); return success(); } }; @@ -776,7 +782,7 @@ class SparseTensorAllocConverter // Reuses specifier. fields.push_back(desc.getSpecifier()); assert(fields.size() == desc.getNumFields()); - rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); + rewriter.replaceOpWithMultiple(op, {fields}); return success(); } @@ -796,7 +802,7 @@ class SparseTensorAllocConverter sizeHint, lvlSizesValues, fields); // Replace operation with resulting memrefs. - rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); + rewriter.replaceOpWithMultiple(op, {fields}); return success(); } @@ -837,7 +843,7 @@ class SparseTensorEmptyConverter : public OpConversionPattern<tensor::EmptyOp> { sizeHint, lvlSizesValues, fields); // Replace operation with resulting memrefs. - rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); + rewriter.replaceOpWithMultiple(op, {fields}); return success(); } @@ -893,7 +899,7 @@ class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> { if (op.getHasInserts()) genEndInsert(rewriter, op.getLoc(), desc); // Replace operation with resulting memrefs. - rewriter.replaceOp(op, genTuple(rewriter, op.getLoc(), desc)); + rewriter.replaceOpWithMultiple(op, {desc.getFields()}); return success(); } }; @@ -1006,7 +1012,6 @@ class SparseCompressConverter : public OpConversionPattern<CompressOp> { rewriter.create<scf::YieldOp>(loc, insertRet); rewriter.setInsertionPointAfter(loop); - Value result = genTuple(rewriter, loc, dstType, loop->getResults()); // Deallocate the buffers on exit of the full loop nest. Operation *parent = getTop(op); rewriter.setInsertionPointAfter(parent); @@ -1014,7 +1019,7 @@ class SparseCompressConverter : public OpConversionPattern<CompressOp> { rewriter.create<memref::DeallocOp>(loc, filled); rewriter.create<memref::DeallocOp>(loc, added); // Replace operation with resulting memrefs. - rewriter.replaceOp(op, result); + rewriter.replaceOpWithMultiple(op, {loop->getResults()}); return success(); } }; @@ -1041,8 +1046,7 @@ class SparseInsertConverter : public OpConversionPattern<tensor::InsertOp> { params, /*genCall=*/true); SmallVector<Value> ret = insertGen.genCallOrInline(rewriter, loc); // Replace operation with resulting memrefs. - rewriter.replaceOp(op, - genTuple(rewriter, loc, op.getDest().getType(), ret)); + rewriter.replaceOpWithMultiple(op, {ret}); return success(); } }; @@ -1215,8 +1219,7 @@ class SparseConvertConverter : public OpConversionPattern<ConvertOp> { return true; }); - rewriter.replaceOp( - op, genTuple(rewriter, loc, op.getResult().getType(), fields)); + rewriter.replaceOpWithMultiple(op, {fields}); return success(); } }; @@ -1271,8 +1274,7 @@ class SparseExtractSliceConverter // NOTE: we can not generate tuples directly from descriptor here, as the // descriptor is holding the original type, yet we want the slice type // here (they shared every memref but with an updated specifier). - rewriter.replaceOp(op, genTuple(rewriter, loc, op.getResult().getType(), - desc.getFields())); + rewriter.replaceOpWithMultiple(op, {desc.getFields()}); return success(); } }; @@ -1403,7 +1405,7 @@ struct SparseAssembleOpConverter : public OpConversionPattern<AssembleOp> { } desc.setValMemSize(rewriter, loc, memSize); - rewriter.replaceOp(op, genTuple(rewriter, loc, desc)); + rewriter.replaceOpWithMultiple(op, {desc.getFields()}); return success(); } }; @@ -1577,7 +1579,7 @@ struct SparseNewConverter : public OpConversionPattern<NewOp> { EmitCInterface::Off); // Replace operation with resulting memrefs. - rewriter.replaceOp(op, genTuple(rewriter, loc, dstTp, fields)); + rewriter.replaceOpWithMultiple(op, {fields}); return success(); } }; diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp index a3db50573c2720..834e3634cc130d 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorDescriptor.cpp @@ -54,19 +54,24 @@ convertSparseTensorType(RankedTensorType rtp, SmallVectorImpl<Type> &fields) { // The sparse tensor type converter (defined in Passes.h). //===----------------------------------------------------------------------===// +static Value materializeTuple(OpBuilder &builder, RankedTensorType tp, + ValueRange inputs, Location loc) { + if (!getSparseTensorEncoding(tp)) + // Not a sparse tensor. + return Value(); + // Sparsifier knows how to cancel out these casts. + return genTuple(builder, loc, tp, inputs); +} + SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() { addConversion([](Type type) { return type; }); addConversion(convertSparseTensorType); // Required by scf.for 1:N type conversion. - addSourceMaterialization([](OpBuilder &builder, RankedTensorType tp, - ValueRange inputs, Location loc) -> Value { - if (!getSparseTensorEncoding(tp)) - // Not a sparse tensor. - return Value(); - // Sparsifier knows how to cancel out these casts. - return genTuple(builder, loc, tp, inputs); - }); + addSourceMaterialization(materializeTuple); + + // Required as a workaround until we have full 1:N support. + addArgumentMaterialization(materializeTuple); } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp index 4714c3cace6c78..e85a86e94282ec 100644 --- a/mlir/lib/IR/Builders.cpp +++ b/mlir/lib/IR/Builders.cpp @@ -645,7 +645,7 @@ void OpBuilder::cloneRegionBefore(Region ®ion, Block *before) { OpBuilder::InsertPoint OpBuilder::InsertPoint::after(ArrayRef<Value> values, - const PostDominanceInfo &domInfo) { + const PostDominanceInfo *domInfo) { // Helper function that computes the point after v's definition. auto computeAfterIp = [](Value v) -> std::pair<Block *, Block::iterator> { if (auto blockArg = dyn_cast<BlockArgument>(v)) @@ -658,12 +658,18 @@ OpBuilder::InsertPoint::after(ArrayRef<Value> values, assert(!values.empty() && "expected at least one Value"); auto [block, blockIt] = computeAfterIp(values.front()); + if (values.size() == 1) { + // Fast path: There is only one value. + return InsertPoint(block, blockIt); + } + // Check the other values one-by-one and update the insertion point if // needed. + assert(domInfo && "domInfo expected if >1 values"); for (Value v : values.drop_front()) { auto [candidateBlock, candidateBlockIt] = computeAfterIp(v); - if (domInfo.postDominantes(candidateBlock, candidateBlockIt, block, - blockIt)) { + if (domInfo->postDominantes(candidateBlock, candidateBlockIt, block, + blockIt)) { // The point after v's definition post-dominates the current (and all // previous) insertion points. Note: Post-dominance is transitive. block = candidateBlock; @@ -671,8 +677,8 @@ OpBuilder::InsertPoint::after(ArrayRef<Value> values, continue; } - if (!domInfo.postDominantes(block, blockIt, candidateBlock, - candidateBlockIt)) { + if (!domInfo->postDominantes(block, blockIt, candidateBlock, + candidateBlockIt)) { // The point after v's definition and the current insertion point do not // post-dominate each other. Therefore, there is no insertion point that // post-dominates all values. diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp index 0a62628b9ad240..2f6c0a1ab0bd3b 100644 --- a/mlir/lib/Transforms/Utils/DialectConversion.cpp +++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp @@ -11,6 +11,7 @@ #include "mlir/IR/Block.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dominance.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/Iterators.h" #include "mlir/Interfaces/FunctionInterfaces.h" @@ -53,20 +54,14 @@ static void logFailure(llvm::ScopedPrinter &os, StringRef fmt, Args &&...args) { }); } -/// Helper function that computes an insertion point where the given value is -/// defined and can be used without a dominance violation. -static OpBuilder::InsertPoint computeInsertPoint(Value value) { - Block *insertBlock = value.getParentBlock(); - Block::iterator insertPt = insertBlock->begin(); - if (OpResult inputRes = dyn_cast<OpResult>(value)) - insertPt = ++inputRes.getOwner()->getIterator(); - return OpBuilder::InsertPoint(insertBlock, insertPt); -} - //===----------------------------------------------------------------------===// // ConversionValueMapping //===----------------------------------------------------------------------===// +/// A list of replacement SSA values. Optimized for the common case of a single +/// SSA value. +using ReplacementValues = SmallVector<Value, 1>; + namespace { /// This class wraps a IRMapping to provide recursive lookup /// functionality, i.e. we will traverse if the mapped value also has a mapping. @@ -818,6 +813,22 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { Type originalType, const TypeConverter *converter); + /// Build an N:1 materialization for the given original value that was + /// replaced with the given replacement values. + /// + /// This is a workaround around incomplete 1:N support in the dialect + /// conversion driver. The conversion mapping can store only 1:1 replacements + /// and the conversion patterns only support single Value replacements in the + /// adaptor, so N values must be converted back to a single value. This + /// function will be deleted when full 1:N support has been added. + /// + /// This function inserts an argument materialization back to the original + /// type, followed by a target materialization to the legalized type (if + /// applicable). + void insertNTo1Materialization(OpBuilder::InsertPoint ip, Location loc, + ValueRange replacements, Value originalValue, + const TypeConverter *converter); + //===--------------------------------------------------------------------===// // Rewriter Notification Hooks //===--------------------------------------------------------------------===// @@ -827,7 +838,7 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener { OpBuilder::InsertPoint previous) override; /// Notifies that an op is about to be replaced with the given values. - void notifyOpReplaced(Operation *op, ValueRange newValues); + void notifyOpReplaced(Operation *op, ArrayRef<ReplacementValues> newValues); /// Notifies that a block is about to be erased. void notifyBlockIsBeingErased(Block *block); @@ -1147,8 +1158,9 @@ LogicalResult ConversionPatternRewriterImpl::remapValues( // that the value was replaced with a value of different type and no // ... [truncated] `````````` </details> https://github.com/llvm/llvm-project/pull/115816 _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
