mlir/lib/IR/Location.cpp

   1 //===- Location.cpp - MLIR Location Classes -------------------------------===//
   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/Location.h"
  10 #include "mlir/IR/BuiltinDialect.h"
  11 #include "mlir/IR/Visitors.h"
  12 #include "llvm/ADT/SetVector.h"
  13 #include "llvm/ADT/TypeSwitch.h"
  14
  15 using namespace mlir;
  16 using namespace mlir::detail;
  17
  18 //===----------------------------------------------------------------------===//
  19 /// Tablegen Attribute Definitions
  20 //===----------------------------------------------------------------------===//
  21
  22 #define GET_ATTRDEF_CLASSES
  23 #include "mlir/IR/BuiltinLocationAttributes.cpp.inc"
  24
  25 //===----------------------------------------------------------------------===//
  26 // BuiltinDialect
  27 //===----------------------------------------------------------------------===//
  28
  29 void BuiltinDialect::registerLocationAttributes() {
  30   addAttributes<
  31 #define GET_ATTRDEF_LIST
  32 #include "mlir/IR/BuiltinLocationAttributes.cpp.inc"
  33       >();
  34 }
  35
  36 //===----------------------------------------------------------------------===//
  37 // LocationAttr
  38 //===----------------------------------------------------------------------===//
  39
  40 WalkResult LocationAttr::walk(function_ref<WalkResult(Location)> walkFn) {
  41   AttrTypeWalker walker;
  42   // Walk locations, but skip any other attribute.
  43   walker.addWalk([&](Attribute attr) {
  44     if (auto loc = llvm::dyn_cast<LocationAttr>(attr))
  45       return walkFn(loc);
  46
  47     return WalkResult::skip();
  48   });
  49   return walker.walk<WalkOrder::PreOrder>(*this);
  50 }
  51
  52 /// Methods for support type inquiry through isa, cast, and dyn_cast.
  53 bool LocationAttr::classof(Attribute attr) {
  54   return attr.hasTrait<AttributeTrait::IsLocation>();
  55 }
  56
  57 //===----------------------------------------------------------------------===//
  58 // CallSiteLoc
  59 //===----------------------------------------------------------------------===//
  60
  61 CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames) {
  62   assert(!frames.empty() && "required at least 1 call frame");
  63   Location caller = frames.back();
  64   for (auto frame : llvm::reverse(frames.drop_back()))
  65     caller = CallSiteLoc::get(frame, caller);
  66   return CallSiteLoc::get(name, caller);
  67 }
  68
  69 //===----------------------------------------------------------------------===//
  70 // FusedLoc
  71 //===----------------------------------------------------------------------===//
  72
  73 Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
  74                        MLIRContext *context) {
  75   // Unique the set of locations to be fused.
  76   llvm::SmallSetVector<Location, 4> decomposedLocs;
  77   for (auto loc : locs) {
  78     // If the location is a fused location we decompose it if it has no
  79     // metadata or the metadata is the same as the top level metadata.
  80     if (auto fusedLoc = llvm::dyn_cast<FusedLoc>(loc)) {
  81       if (fusedLoc.getMetadata() == metadata) {
  82         // UnknownLoc's have already been removed from FusedLocs so we can
  83         // simply add all of the internal locations.
  84         decomposedLocs.insert(fusedLoc.getLocations().begin(),
  85                               fusedLoc.getLocations().end());
  86         continue;
  87       }
  88     }
  89     // Otherwise, only add known locations to the set.
  90     if (!llvm::isa<UnknownLoc>(loc))
  91       decomposedLocs.insert(loc);
  92   }
  93   locs = decomposedLocs.getArrayRef();
  94
  95   // Handle the simple cases of less than two locations. Ensure the metadata (if
  96   // provided) is not dropped.
  97   if (locs.empty()) {
  98     if (!metadata)
  99       return UnknownLoc::get(context);
 100     // TODO: Investigate ASAN failure when using implicit conversion from
 101     // Location to ArrayRef<Location> below.
 102     return Base::get(context, ArrayRef<Location>{UnknownLoc::get(context)},
 103                      metadata);
 104   }
 105   if (locs.size() == 1 && !metadata)
 106     return locs.front();
 107
 108   return Base::get(context, locs, metadata);
 109 }