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 13 using namespace mlir; 14 15 namespace { 16 /// A simple structure which is used for testing as an underlying location in 17 /// OpaqueLoc. 18 struct MyLocation { 19 MyLocation() = default; MyLocation__anon26dd931e0111::MyLocation20 MyLocation(int id) : id(id) {} getId__anon26dd931e0111::MyLocation21 int getId() { return id; } 22 23 int id{42}; 24 }; 25 } // namespace 26 27 MLIR_DECLARE_EXPLICIT_TYPE_ID(MyLocation *) 28 MLIR_DEFINE_EXPLICIT_TYPE_ID(MyLocation *) 29 30 namespace { 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. 35 struct TestOpaqueLoc 36 : public PassWrapper<TestOpaqueLoc, OperationPass<ModuleOp>> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon26dd931e0211::TestOpaqueLoc37 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestOpaqueLoc) 38 39 StringRef getArgument() const final { return "test-opaque-loc"; } getDescription__anon26dd931e0211::TestOpaqueLoc40 StringRef getDescription() const final { 41 return "Changes all leaf locations to opaque locations"; 42 } 43 runOnOperation__anon26dd931e0211::TestOpaqueLoc44 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(std::make_unique<MyLocation>(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>(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); 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); 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 95 namespace mlir { 96 namespace test { registerTestOpaqueLoc()97void registerTestOpaqueLoc() { PassRegistration<TestOpaqueLoc>(); } 98 } // namespace test 99 } // namespace mlir 100