1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
|
//===- ReconcileUnrealizedCasts.cpp - Eliminate noop unrealized casts -----===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
#define GEN_PASS_DEF_RECONCILEUNREALIZEDCASTS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
using namespace mlir;
namespace {
/// Pass to simplify and eliminate unrealized conversion casts.
///
/// This pass processes unrealized_conversion_cast ops in a worklist-driven
/// fashion. For each matched cast op, if the chain of input casts eventually
/// reaches a cast op where the input types match the output types of the
/// matched op, replace the matched op with the inputs.
///
/// Example:
/// %1 = unrealized_conversion_cast %0 : !A to !B
/// %2 = unrealized_conversion_cast %1 : !B to !C
/// %3 = unrealized_conversion_cast %2 : !C to !A
///
/// In the above example, %0 can be used instead of %3 and all cast ops are
/// folded away.
struct ReconcileUnrealizedCasts
: public impl::ReconcileUnrealizedCastsBase<ReconcileUnrealizedCasts> {
ReconcileUnrealizedCasts() = default;
void runOnOperation() override {
// Gather all unrealized_conversion_cast ops.
SetVector<UnrealizedConversionCastOp> worklist;
getOperation()->walk(
[&](UnrealizedConversionCastOp castOp) { worklist.insert(castOp); });
// Helper function that adds all operands to the worklist that are an
// unrealized_conversion_cast op result.
auto enqueueOperands = [&](UnrealizedConversionCastOp castOp) {
for (Value v : castOp.getInputs())
if (auto inputCastOp = v.getDefiningOp<UnrealizedConversionCastOp>())
worklist.insert(inputCastOp);
};
// Helper function that return the unrealized_conversion_cast op that
// defines all inputs of the given op (in the same order). Return "nullptr"
// if there is no such op.
auto getInputCast =
[](UnrealizedConversionCastOp castOp) -> UnrealizedConversionCastOp {
if (castOp.getInputs().empty())
return {};
auto inputCastOp = castOp.getInputs()
.front()
.getDefiningOp<UnrealizedConversionCastOp>();
if (!inputCastOp)
return {};
if (inputCastOp.getOutputs() != castOp.getInputs())
return {};
return inputCastOp;
};
// Process ops in the worklist bottom-to-top.
while (!worklist.empty()) {
UnrealizedConversionCastOp castOp = worklist.pop_back_val();
if (castOp->use_empty()) {
// DCE: If the op has no users, erase it. Add the operands to the
// worklist to find additional DCE opportunities.
enqueueOperands(castOp);
castOp->erase();
continue;
}
// Traverse the chain of input cast ops to see if an op with the same
// input types can be found.
UnrealizedConversionCastOp nextCast = castOp;
while (nextCast) {
if (nextCast.getInputs().getTypes() == castOp.getResultTypes()) {
// Found a cast where the input types match the output types of the
// matched op. We can directly use those inputs and the matched op can
// be removed.
enqueueOperands(castOp);
castOp.replaceAllUsesWith(nextCast.getInputs());
castOp->erase();
break;
}
nextCast = getInputCast(nextCast);
}
}
}
};
} // namespace
std::unique_ptr<Pass> mlir::createReconcileUnrealizedCastsPass() {
return std::make_unique<ReconcileUnrealizedCasts>();
}
|
