100 lines · cpp
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-exception6//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 in17/// OpaqueLoc.18struct MyLocation {19 MyLocation() = default;20 MyLocation(int id) : id(id) {}21 int getId() { return id; }22 23 int id{42};24};25} // namespace26 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 or33/// terminators and clones them with opaque locations which store the initial34/// locations.35struct TestOpaqueLoc36 : 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(std::make_unique<MyLocation>(lastIt++));50 51 Location loc = op->getLoc();52 53 /// Set opaque location without fallback location to test the54 /// 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 the65 /// 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 test70 /// 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 else83 os << "nullptr";84 }85 os << ": " << diag << '\n';86 os.flush();87 });88 89 getOperation().walk([&](Operation *op) { op->emitOpError(); });90 }91};92 93} // namespace94 95namespace mlir {96namespace test {97void registerTestOpaqueLoc() { PassRegistration<TestOpaqueLoc>(); }98} // namespace test99} // namespace mlir100