1//===- TestOpaqueLoc.cpp - Pass to test opaque locations ------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir/IR/Builders.h"
10#include "mlir/IR/BuiltinOps.h"
11#include "mlir/Pass/Pass.h"
12
13using namespace mlir;
14
15namespace {
16/// A simple structure which is used for testing as an underlying location in
17/// OpaqueLoc.
18struct MyLocation {
19 MyLocation() = default;
20 MyLocation(int id) : id(id) {}
21 int getId() { return id; }
22
23 int id{42};
24};
25} // namespace
26
27MLIR_DECLARE_EXPLICIT_TYPE_ID(MyLocation *)
28MLIR_DEFINE_EXPLICIT_TYPE_ID(MyLocation *)
29
30namespace {
31/// Pass that changes locations to opaque locations for each operation.
32/// It also takes all operations that are not function operations or
33/// terminators and clones them with opaque locations which store the initial
34/// locations.
35struct TestOpaqueLoc
36 : public PassWrapper<TestOpaqueLoc, OperationPass<ModuleOp>> {
37 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpaqueLoc)
38
39 StringRef getArgument() const final { return "test-opaque-loc"; }
40 StringRef getDescription() const final {
41 return "Changes all leaf locations to opaque locations";
42 }
43
44 void runOnOperation() override {
45 std::vector<std::unique_ptr<MyLocation>> myLocs;
46 int lastIt = 0;
47
48 getOperation().getBody()->walk([&](Operation *op) {
49 myLocs.push_back(x: std::make_unique<MyLocation>(args: lastIt++));
50
51 Location loc = op->getLoc();
52
53 /// Set opaque location without fallback location to test the
54 /// corresponding get method.
55 op->setLoc(
56 OpaqueLoc::get<MyLocation *>(myLocs.back().get(), &getContext()));
57
58 if (isa<ModuleOp>(Val: op->getParentOp()) ||
59 op->hasTrait<OpTrait::IsTerminator>())
60 return;
61
62 OpBuilder builder(op);
63
64 /// Add the same operation but with fallback location to test the
65 /// corresponding get method and serialization.
66 Operation *opCloned1 = builder.clone(op&: *op);
67 opCloned1->setLoc(OpaqueLoc::get<MyLocation *>(myLocs.back().get(), loc));
68
69 /// Add the same operation but with void* instead of MyLocation* to test
70 /// getUnderlyingLocationOrNull method.
71 Operation *opCloned2 = builder.clone(op&: *op);
72 opCloned2->setLoc(OpaqueLoc::get<void *>(nullptr, loc));
73 });
74
75 ScopedDiagnosticHandler diagHandler(&getContext(), [](Diagnostic &diag) {
76 auto &os = llvm::outs();
77 if (isa<OpaqueLoc>(diag.getLocation())) {
78 MyLocation *loc = OpaqueLoc::getUnderlyingLocationOrNull<MyLocation *>(
79 diag.getLocation());
80 if (loc)
81 os << "MyLocation: " << loc->id;
82 else
83 os << "nullptr";
84 }
85 os << ": " << diag << '\n';
86 os.flush();
87 });
88
89 getOperation().walk([&](Operation *op) { op->emitOpError(); });
90 }
91};
92
93} // namespace
94
95namespace mlir {
96namespace test {
97void registerTestOpaqueLoc() { PassRegistration<TestOpaqueLoc>(); }
98} // namespace test
99} // namespace mlir
100

source code of mlir/test/lib/IR/TestOpaqueLoc.cpp