[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: Fix missing source materialization (PR #97903)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Jul 6 09:09:49 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm
Author: Matthias Springer (matthias-springer)
<details>
<summary>Changes</summary>
This commit fixes a bug in the dialect conversion. During a 1:N signature conversion, the dialect conversion did not insert a cast back to the original block argument type, producing invalid IR.
See `test-block-legalization.mlir`: Without this commit, the operand type of the op changes because an `unrealized_conversion_cast` is missing:
```
"test.consumer_of_complex"(%v) : (!llvm.struct<(f64, f64)>) -> ()
```
To implement this fix, it was necessary to change the meaning of argument materializations. An argument materialization now maps from the new block argument types to the original block argument type. This also addresses a `FIXME` in the code base:
```
// FIXME: The current argument materialization hook expects the original
// output type, even though it doesn't use that as the actual output type
// of the generated IR. The output type is just used as an indicator of
// the type of materialization to do. This behavior is really awkward in
// that it diverges from the behavior of the other hooks, and can be
// easily misunderstood. We should clean up the argument hooks to better
// represent the desired invariants we actually care about.
```
It is no longer necessary to distinguish between the "output type" and the "original output type".
Most type converter are already written according to the new API. (Most implementations use the same conversion functions as for source materializations.) One exception is the MemRef-to-LLVM type converter, which materialized an `!llvm.struct` based on the elements of a memref descriptor. It still does that, but casts the `!llvm.struct` back to the original memref type. The dialect conversion inserts a target materialization (to `!llvm.struct`) which cancels out with the other cast.
This commit also fixes a bug in `computeNecessaryMaterializations`. The implementation did not account for the possibility that a value was replaced multiple times. E.g., replace `a` by `b`, then `b` by `c`.
This commit also adds a transform dialect op to populate SCF-to-CF patterns. This transform op was needed to write a test case. The bug described here appears only during a complex interplay of 1:N signature conversions and op replacements. (I was not able to trigger it with ops and patterns from the `test` dialect without duplicating the `scf.if` pattern.)
---
Full diff: https://github.com/llvm/llvm-project/pull/97903.diff
9 Files Affected:
- (modified) mlir/docs/DialectConversion.md (+2-1)
- (modified) mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td (+11)
- (modified) mlir/include/mlir/Transforms/DialectConversion.h (+2-1)
- (modified) mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp (+21-7)
- (modified) mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp (+7)
- (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+40-43)
- (modified) mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir (+2-2)
- (added) mlir/test/Transforms/test-block-legalization.mlir (+44)
``````````diff
diff --git a/mlir/docs/DialectConversion.md b/mlir/docs/DialectConversion.md
index db26e6477d5fc7..23e74470a835f7 100644
--- a/mlir/docs/DialectConversion.md
+++ b/mlir/docs/DialectConversion.md
@@ -352,7 +352,8 @@ class TypeConverter {
/// This method registers a materialization that will be called when
/// converting (potentially multiple) block arguments that were the result of
- /// a signature conversion of a single block argument, to a single SSA value.
+ /// a signature conversion of a single block argument, to a single SSA value
+ /// with the old argument type.
template <typename FnT,
typename T = typename llvm::function_traits<FnT>::template arg_t<1>>
void addArgumentMaterialization(FnT &&callback) {
diff --git a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
index 7bf914f6456ce1..20880d94a83cac 100644
--- a/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
+++ b/mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td
@@ -38,6 +38,17 @@ def ApplySCFStructuralConversionPatternsOp : Op<Transform_Dialect,
let assemblyFormat = "attr-dict";
}
+def ApplySCFToControlFlowPatternsOp : Op<Transform_Dialect,
+ "apply_conversion_patterns.scf.scf_to_control_flow",
+ [DeclareOpInterfaceMethods<ConversionPatternDescriptorOpInterface>]> {
+ let description = [{
+ Collects patterns that lower structured control flow ops to unstructured
+ control flow.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
def Transform_ScfForOp : Transform_ConcreteOpType<"scf.for">;
def ForallToForOp : Op<Transform_Dialect, "loop.forall_to_for",
diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index a22f198bdf2520..6999af3909010a 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -182,7 +182,8 @@ class TypeConverter {
/// This method registers a materialization that will be called when
/// converting (potentially multiple) block arguments that were the result of
- /// a signature conversion of a single block argument, to a single SSA value.
+ /// a signature conversion of a single block argument, to a single SSA value
+ /// with the old block argument type.
template <typename FnT, typename T = typename llvm::function_traits<
std::decay_t<FnT>>::template arg_t<1>>
void addArgumentMaterialization(FnT &&callback) {
diff --git a/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp b/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
index f5620a6a7cd913..32d02d5e438bdd 100644
--- a/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
+++ b/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
@@ -153,9 +153,11 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
type.isVarArg());
});
- // Materialization for memrefs creates descriptor structs from individual
- // values constituting them, when descriptors are used, i.e. more than one
- // value represents a memref.
+ // Argument materializations convert from the new block argument types
+ // (multiple SSA values that make up a memref descriptor) back to the
+ // original block argument type. The dialect conversion framework will then
+ // insert a target materialization from the original block argument type to
+ // a legal type.
addArgumentMaterialization(
[&](OpBuilder &builder, UnrankedMemRefType resultType, ValueRange inputs,
Location loc) -> std::optional<Value> {
@@ -164,12 +166,18 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
// memref descriptor cannot be built just from a bare pointer.
return std::nullopt;
}
- return UnrankedMemRefDescriptor::pack(builder, loc, *this, resultType,
- inputs);
+ Value desc = UnrankedMemRefDescriptor::pack(builder, loc, *this,
+ resultType, inputs);
+ // An argument materialization must return a value of type
+ // `resultType`, so insert a cast from the memref descriptor type
+ // (!llvm.struct) to the original memref type.
+ return builder.create<UnrealizedConversionCastOp>(loc, resultType, desc)
+ .getResult(0);
});
addArgumentMaterialization([&](OpBuilder &builder, MemRefType resultType,
ValueRange inputs,
Location loc) -> std::optional<Value> {
+ Value desc;
if (inputs.size() == 1) {
// This is a bare pointer. We allow bare pointers only for function entry
// blocks.
@@ -180,10 +188,16 @@ LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
if (!block->isEntryBlock() ||
!isa<FunctionOpInterface>(block->getParentOp()))
return std::nullopt;
- return MemRefDescriptor::fromStaticShape(builder, loc, *this, resultType,
+ desc = MemRefDescriptor::fromStaticShape(builder, loc, *this, resultType,
inputs[0]);
+ } else {
+ desc = MemRefDescriptor::pack(builder, loc, *this, resultType, inputs);
}
- return MemRefDescriptor::pack(builder, loc, *this, resultType, inputs);
+ // An argument materialization must return a value of type `resultType`,
+ // so insert a cast from the memref descriptor type (!llvm.struct) to the
+ // original memref type.
+ return builder.create<UnrealizedConversionCastOp>(loc, resultType, desc)
+ .getResult(0);
});
// Add generic source and target materializations to handle cases where
// non-LLVM types persist after an LLVM conversion.
diff --git a/mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt b/mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt
index 1d6f9ebd153f0b..06bccab80e7d80 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt
+++ b/mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt
@@ -13,6 +13,7 @@ add_mlir_dialect_library(MLIRSCFTransformOps
MLIRIR
MLIRLoopLikeInterface
MLIRSCFDialect
+ MLIRSCFToControlFlow
MLIRSCFTransforms
MLIRSCFUtils
MLIRTransformDialect
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 56ff2709a589ec..9921b09fcba7f4 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -7,6 +7,8 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
+
+#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
@@ -49,6 +51,11 @@ void transform::ApplySCFStructuralConversionPatternsOp::
conversionTarget);
}
+void transform::ApplySCFToControlFlowPatternsOp::populatePatterns(
+TypeConverter &typeConverter, RewritePatternSet &patterns) {
+ populateSCFToControlFlowConversionPatterns(patterns);
+}
+
//===----------------------------------------------------------------------===//
// ForallToForOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index e6c0ee2ab29490..4fad77ff24759f 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -707,10 +707,9 @@ class UnresolvedMaterializationRewrite : public OperationRewrite {
UnresolvedMaterializationRewrite(
ConversionPatternRewriterImpl &rewriterImpl,
UnrealizedConversionCastOp op, const TypeConverter *converter = nullptr,
- MaterializationKind kind = MaterializationKind::Target,
- Type origOutputType = nullptr)
+ MaterializationKind kind = MaterializationKind::Target)
: OperationRewrite(Kind::UnresolvedMaterialization, rewriterImpl, op),
- converterAndKind(converter, kind), origOutputType(origOutputType) {}
+ converterAndKind(converter, kind) {}
static bool classof(const IRRewrite *rewrite) {
return rewrite->getKind() == Kind::UnresolvedMaterialization;
@@ -734,17 +733,11 @@ class UnresolvedMaterializationRewrite : public OperationRewrite {
return converterAndKind.getInt();
}
- /// Return the original illegal output type of the input values.
- Type getOrigOutputType() const { return origOutputType; }
-
private:
/// The corresponding type converter to use when resolving this
/// materialization, and the kind of this materialization.
llvm::PointerIntPair<const TypeConverter *, 1, MaterializationKind>
converterAndKind;
-
- /// The original output type. This is only used for argument conversions.
- Type origOutputType;
};
} // namespace
@@ -860,12 +853,10 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
Block *insertBlock,
Block::iterator insertPt, Location loc,
ValueRange inputs, Type outputType,
- Type origOutputType,
const TypeConverter *converter);
Value buildUnresolvedArgumentMaterialization(Block *block, Location loc,
ValueRange inputs,
- Type origOutputType,
Type outputType,
const TypeConverter *converter);
@@ -1388,20 +1379,24 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion(
if (replArgs.size() == 1 &&
(!converter || replArgs[0].getType() == origArg.getType())) {
newArg = replArgs.front();
+ mapping.map(origArg, newArg);
} else {
- Type origOutputType = origArg.getType();
-
- // Legalize the argument output type.
- Type outputType = origOutputType;
- if (Type legalOutputType = converter->convertType(outputType))
- outputType = legalOutputType;
-
- newArg = buildUnresolvedArgumentMaterialization(
- newBlock, origArg.getLoc(), replArgs, origOutputType, outputType,
- converter);
+ // Build argument materialization: new block arguments -> old block
+ // argument type.
+ Value argMat = buildUnresolvedArgumentMaterialization(
+ newBlock, origArg.getLoc(), replArgs, origArg.getType(), converter);
+ mapping.map(origArg, argMat);
+
+ // Build target materialization: old block argument type -> legal type.
+ if (Type legalOutputType = converter->convertType(origArg.getType())) {
+ newArg = buildUnresolvedTargetMaterialization(
+ origArg.getLoc(), argMat, legalOutputType, converter);
+ mapping.map(argMat, newArg);
+ } else {
+ newArg = argMat;
+ }
}
- mapping.map(origArg, newArg);
appendRewrite<ReplaceBlockArgRewrite>(block, origArg);
argInfo[i] = ConvertedArgInfo(inputMap->inputNo, inputMap->size, newArg);
}
@@ -1424,7 +1419,7 @@ Block *ConversionPatternRewriterImpl::applySignatureConversion(
/// of input operands.
Value ConversionPatternRewriterImpl::buildUnresolvedMaterialization(
MaterializationKind kind, Block *insertBlock, Block::iterator insertPt,
- Location loc, ValueRange inputs, Type outputType, Type origOutputType,
+ Location loc, ValueRange inputs, Type outputType,
const TypeConverter *converter) {
// Avoid materializing an unnecessary cast.
if (inputs.size() == 1 && inputs.front().getType() == outputType)
@@ -1435,16 +1430,15 @@ Value ConversionPatternRewriterImpl::buildUnresolvedMaterialization(
OpBuilder builder(insertBlock, insertPt);
auto convertOp =
builder.create<UnrealizedConversionCastOp>(loc, outputType, inputs);
- appendRewrite<UnresolvedMaterializationRewrite>(convertOp, converter, kind,
- origOutputType);
+ appendRewrite<UnresolvedMaterializationRewrite>(convertOp, converter, kind);
return convertOp.getResult(0);
}
Value ConversionPatternRewriterImpl::buildUnresolvedArgumentMaterialization(
- Block *block, Location loc, ValueRange inputs, Type origOutputType,
- Type outputType, const TypeConverter *converter) {
+ Block *block, Location loc, ValueRange inputs, Type outputType,
+ const TypeConverter *converter) {
return buildUnresolvedMaterialization(MaterializationKind::Argument, block,
block->begin(), loc, inputs, outputType,
- origOutputType, converter);
+ converter);
}
Value ConversionPatternRewriterImpl::buildUnresolvedTargetMaterialization(
Location loc, Value input, Type outputType,
@@ -1456,7 +1450,7 @@ Value ConversionPatternRewriterImpl::buildUnresolvedTargetMaterialization(
return buildUnresolvedMaterialization(MaterializationKind::Target,
insertBlock, insertPt, loc, input,
- outputType, outputType, converter);
+ outputType, converter);
}
//===----------------------------------------------------------------------===//
@@ -2672,6 +2666,9 @@ static void computeNecessaryMaterializations(
ConversionPatternRewriterImpl &rewriterImpl,
DenseMap<Value, SmallVector<Value>> &inverseMapping,
SetVector<UnresolvedMaterializationRewrite *> &necessaryMaterializations) {
+ // Helper function to check if the given value or a not yet materialized
+ // replacement of the given value is live.
+ // Note: `inverseMapping` maps from replaced values to original values.
auto isLive = [&](Value value) {
auto findFn = [&](Operation *user) {
auto matIt = materializationOps.find(user);
@@ -2679,12 +2676,18 @@ static void computeNecessaryMaterializations(
return !necessaryMaterializations.count(matIt->second);
return rewriterImpl.isOpIgnored(user);
};
- // This value may be replacing another value that has a live user.
- for (Value inv : inverseMapping.lookup(value))
- if (llvm::find_if_not(inv.getUsers(), findFn) != inv.user_end())
+ // A worklist is needed because a value may have gone through a chain of
+ // replacements and each of the replaced values may have live users.
+ SmallVector<Value> worklist;
+ worklist.push_back(value);
+ while (!worklist.empty()) {
+ Value next = worklist.pop_back_val();
+ if (llvm::find_if_not(next.getUsers(), findFn) != next.user_end())
return true;
- // Or have live users itself.
- return llvm::find_if_not(value.getUsers(), findFn) != value.user_end();
+ // This value may be replacing another value that has a live user.
+ llvm::append_range(worklist, inverseMapping.lookup(next));
+ }
+ return false;
};
llvm::unique_function<Value(Value, Value, Type)> lookupRemappedValue =
@@ -2844,18 +2847,10 @@ static LogicalResult legalizeUnresolvedMaterialization(
switch (mat.getMaterializationKind()) {
case MaterializationKind::Argument:
// Try to materialize an argument conversion.
- // FIXME: The current argument materialization hook expects the original
- // output type, even though it doesn't use that as the actual output type
- // of the generated IR. The output type is just used as an indicator of
- // the type of materialization to do. This behavior is really awkward in
- // that it diverges from the behavior of the other hooks, and can be
- // easily misunderstood. We should clean up the argument hooks to better
- // represent the desired invariants we actually care about.
newMaterialization = converter->materializeArgumentConversion(
- rewriter, op->getLoc(), mat.getOrigOutputType(), inputOperands);
+ rewriter, op->getLoc(), outputType, inputOperands);
if (newMaterialization)
break;
-
// If an argument materialization failed, fallback to trying a target
// materialization.
[[fallthrough]];
@@ -2865,6 +2860,8 @@ static LogicalResult legalizeUnresolvedMaterialization(
break;
}
if (newMaterialization) {
+ assert(newMaterialization.getType() == opResult.getType() &&
+ "materialization callback produced value of incorrect type");
replaceMaterialization(rewriterImpl, opResult, newMaterialization,
inverseMapping);
return success();
diff --git a/mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir b/mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir
index 91ef571cb3bf71..6b9df32fe02dd3 100644
--- a/mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir
+++ b/mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir
@@ -1,8 +1,8 @@
// RUN: mlir-opt -convert-func-to-llvm -reconcile-unrealized-casts %s | FileCheck %s
-// RUN: mlir-opt -convert-func-to-llvm='use-bare-ptr-memref-call-conv=1' %s | FileCheck %s --check-prefix=BAREPTR
+// RUN: mlir-opt -convert-func-to-llvm='use-bare-ptr-memref-call-conv=1' -reconcile-unrealized-casts %s | FileCheck %s --check-prefix=BAREPTR
-// RUN: mlir-opt -transform-interpreter %s | FileCheck %s --check-prefix=BAREPTR
+// RUN: mlir-opt -transform-interpreter -reconcile-unrealized-casts %s | FileCheck %s --check-prefix=BAREPTR
// These tests were separated from func-memref.mlir because applying
// -reconcile-unrealized-casts resulted in `llvm.extractvalue` ops getting
diff --git a/mlir/test/Transforms/test-block-legalization.mlir b/mlir/test/Transforms/test-block-legalization.mlir
new file mode 100644
index 00000000000000..d739f95a569472
--- /dev/null
+++ b/mlir/test/Transforms/test-block-legalization.mlir
@@ -0,0 +1,44 @@
+// RUN: mlir-opt %s -transform-interpreter | FileCheck %s
+
+// CHECK-LABEL: func @complex_block_signature_conversion(
+// CHECK: %[[cst:.*]] = complex.constant
+// CHECK: %[[complex_llvm:.*]] = builtin.unrealized_conversion_cast %[[cst]] : complex<f64> to !llvm.struct<(f64, f64)>
+// Note: Some blocks are omitted.
+// CHECK: llvm.br ^[[block1:.*]](%[[complex_llvm]]
+// CHECK: ^[[block1]](%[[arg:.*]]: !llvm.struct<(f64, f64)>):
+// CHECK: %[[cast:.*]] = builtin.unrealized_conversion_cast %[[arg]] : !llvm.struct<(f64, f64)> to complex<f64>
+// CHECK: llvm.br ^[[block2:.*]]
+// CHECK: ^[[block2]]:
+// CHECK: "test.consumer_of_complex"(%[[cast]]) : (complex<f64>) -> ()
+func.func @complex_block_signature_conversion() {
+ %cst = complex.constant [0.000000e+00, 0.000000e+00] : complex<f64>
+ %true = arith.constant true
+ %0 = scf.if %true -> complex<f64> {
+ scf.yield %cst : complex<f64>
+ } else {
+ scf.yield %cst : complex<f64>
+ }
+
+ // Regression test to ensure that the a source materialization is inserted.
+ // The operand of "test.consumer_of_complex" must not change.
+ "test.consumer_of_complex"(%0) : (complex<f64>) -> ()
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%toplevel_module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %toplevel_module
+ : (!transform.any_op) -> !transform.any_op
+ transform.apply_conversion_patterns to %func {
+ transform.apply_conversion_patterns.dialect_to_llvm "cf"
+ transform.apply_conversion_patterns.func.func_to_llvm
+ transform.apply_conversion_patterns.scf.scf_to_control_flow
+ } with type_converter {
+ transform.apply_conversion_patterns.memref.memref_to_llvm_type_converter
+ } {
+ legal_dialects = ["llvm"],
+ partial_conversion
+ } : !transform.any_op
+ transform.yield
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/97903
More information about the Mlir-commits
mailing list