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
|
//===- TestTopologicalSort.cpp - Pass to test topological sort analysis ---===//
//
// 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/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/TopologicalSortUtils.h"
using namespace mlir;
namespace {
struct TestTopologicalSortAnalysisPass
: public PassWrapper<TestTopologicalSortAnalysisPass,
OperationPass<ModuleOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTopologicalSortAnalysisPass)
StringRef getArgument() const final {
return "test-topological-sort-analysis";
}
StringRef getDescription() const final {
return "Test topological sorting of ops";
}
void runOnOperation() override {
Operation *op = getOperation();
OpBuilder builder(op->getContext());
WalkResult result = op->walk([&](Operation *root) {
if (!root->hasAttr("root"))
return WalkResult::advance();
SmallVector<Operation *> selectedOps;
root->walk([&](Operation *selected) {
if (!selected->hasAttr("selected"))
return WalkResult::advance();
if (root->hasAttr("ordered")) {
// If the root has an "ordered" attribute, we fill the selectedOps
// vector in a certain order.
int64_t pos =
cast<IntegerAttr>(selected->getAttr("selected")).getInt();
if (pos >= static_cast<int64_t>(selectedOps.size()))
selectedOps.append(pos + 1 - selectedOps.size(), nullptr);
selectedOps[pos] = selected;
} else {
selectedOps.push_back(selected);
}
return WalkResult::advance();
});
if (llvm::find(selectedOps, nullptr) != selectedOps.end()) {
root->emitError("invalid test case: some indices are missing among the "
"selected ops");
return WalkResult::skip();
}
if (!computeTopologicalSorting(selectedOps)) {
root->emitError("could not schedule all ops");
return WalkResult::skip();
}
for (const auto &it : llvm::enumerate(selectedOps))
it.value()->setAttr("pos", builder.getIndexAttr(it.index()));
return WalkResult::advance();
});
if (result.wasSkipped())
signalPassFailure();
}
};
} // namespace
namespace mlir {
namespace test {
void registerTestTopologicalSortAnalysisPass() {
PassRegistration<TestTopologicalSortAnalysisPass>();
}
} // namespace test
} // namespace mlir
|
