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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
//===- TopologicalSortUtils.h - Topological sort utilities ------*- C++ -*-===//
//
// 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/Transforms/TopologicalSortUtils.h"
#include "mlir/IR/OpDefinition.h"
using namespace mlir;
/// Return `true` if the given operation is ready to be scheduled.
static bool isOpReady(Operation *op, DenseSet<Operation *> &unscheduledOps,
function_ref<bool(Value, Operation *)> isOperandReady) {
// An operation is ready to be scheduled if all its operands are ready. An
// operation is ready if:
const auto isReady = [&](Value value) {
// - the user-provided callback marks it as ready,
if (isOperandReady && isOperandReady(value, op))
return true;
Operation *parent = value.getDefiningOp();
// - it is a block argument,
if (!parent)
return true;
// - or it is not defined by an unscheduled op (and also not nested within
// an unscheduled op).
do {
// Stop traversal when op under examination is reached.
if (parent == op)
return true;
if (unscheduledOps.contains(parent))
return false;
} while ((parent = parent->getParentOp()));
// No unscheduled op found.
return true;
};
// An operation is recursively ready to be scheduled of it and its nested
// operations are ready.
WalkResult readyToSchedule = op->walk([&](Operation *nestedOp) {
return llvm::all_of(nestedOp->getOperands(),
[&](Value operand) { return isReady(operand); })
? WalkResult::advance()
: WalkResult::interrupt();
});
return !readyToSchedule.wasInterrupted();
}
bool mlir::sortTopologically(
Block *block, llvm::iterator_range<Block::iterator> ops,
function_ref<bool(Value, Operation *)> isOperandReady) {
if (ops.empty())
return true;
// The set of operations that have not yet been scheduled.
DenseSet<Operation *> unscheduledOps;
// Mark all operations as unscheduled.
for (Operation &op : ops)
unscheduledOps.insert(&op);
Block::iterator nextScheduledOp = ops.begin();
Block::iterator end = ops.end();
bool allOpsScheduled = true;
while (!unscheduledOps.empty()) {
bool scheduledAtLeastOnce = false;
// Loop over the ops that are not sorted yet, try to find the ones "ready",
// i.e. the ones for which there aren't any operand produced by an op in the
// set, and "schedule" it (move it before the `nextScheduledOp`).
for (Operation &op :
llvm::make_early_inc_range(llvm::make_range(nextScheduledOp, end))) {
if (!isOpReady(&op, unscheduledOps, isOperandReady))
continue;
// Schedule the operation by moving it to the start.
unscheduledOps.erase(&op);
op.moveBefore(block, nextScheduledOp);
scheduledAtLeastOnce = true;
// Move the iterator forward if we schedule the operation at the front.
if (&op == &*nextScheduledOp)
++nextScheduledOp;
}
// If no operations were scheduled, give up and advance the iterator.
if (!scheduledAtLeastOnce) {
allOpsScheduled = false;
unscheduledOps.erase(&*nextScheduledOp);
++nextScheduledOp;
}
}
return allOpsScheduled;
}
bool mlir::sortTopologically(
Block *block, function_ref<bool(Value, Operation *)> isOperandReady) {
if (block->empty())
return true;
if (block->back().hasTrait<OpTrait::IsTerminator>())
return sortTopologically(block, block->without_terminator(),
isOperandReady);
return sortTopologically(block, *block, isOperandReady);
}
bool mlir::computeTopologicalSorting(
MutableArrayRef<Operation *> ops,
function_ref<bool(Value, Operation *)> isOperandReady) {
if (ops.empty())
return true;
// The set of operations that have not yet been scheduled.
DenseSet<Operation *> unscheduledOps;
// Mark all operations as unscheduled.
for (Operation *op : ops)
unscheduledOps.insert(op);
unsigned nextScheduledOp = 0;
bool allOpsScheduled = true;
while (!unscheduledOps.empty()) {
bool scheduledAtLeastOnce = false;
// Loop over the ops that are not sorted yet, try to find the ones "ready",
// i.e. the ones for which there aren't any operand produced by an op in the
// set, and "schedule" it (swap it with the op at `nextScheduledOp`).
for (unsigned i = nextScheduledOp; i < ops.size(); ++i) {
if (!isOpReady(ops[i], unscheduledOps, isOperandReady))
continue;
// Schedule the operation by moving it to the start.
unscheduledOps.erase(ops[i]);
std::swap(ops[i], ops[nextScheduledOp]);
scheduledAtLeastOnce = true;
++nextScheduledOp;
}
// If no operations were scheduled, just schedule the first op and continue.
if (!scheduledAtLeastOnce) {
allOpsScheduled = false;
unscheduledOps.erase(ops[nextScheduledOp++]);
}
}
return allOpsScheduled;
}
|
