[mlir] Implement the SubElement interfaces for the builtin locations

This enables find/replace of nested components for location attributes.

Differential Revision: https://reviews.llvm.org/D136408

diff --git a/mlir/include/mlir/IR/BuiltinLocationAttributes.td b/mlir/include/mlir/IR/BuiltinLocationAttributes.td
index cc483cd..ca96fb9 100644 (file)
--- a/mlir/include/mlir/IR/BuiltinLocationAttributes.td
+++ b/mlir/include/mlir/IR/BuiltinLocationAttributes.td
@@ -15,10 +15,11 @@
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinDialect.td"
+include "mlir/IR/SubElementInterfaces.td"
 
 // Base class for Builtin dialect location attributes.
-class Builtin_LocationAttr<string name>
-    : AttrDef<Builtin_Dialect, name, [], "::mlir::LocationAttr"> {
+class Builtin_LocationAttr<string name, list<Trait> traits = []>
+    : AttrDef<Builtin_Dialect, name, traits, "::mlir::LocationAttr"> {
   let cppClassName = name;
   let mnemonic = ?;
 }
@@ -27,7 +28,9 @@ class Builtin_LocationAttr<string name>
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
-def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc"> {
+def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A callsite source location";
   let description = [{
     Syntax:
@@ -104,7 +107,9 @@ def FileLineColLoc : Builtin_LocationAttr<"FileLineColLoc"> {
 // FusedLoc
 //===----------------------------------------------------------------------===//
 
-def FusedLoc : Builtin_LocationAttr<"FusedLoc"> {
+def FusedLoc : Builtin_LocationAttr<"FusedLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A tuple of other source locations";
   let description = [{
     Syntax:
@@ -143,7 +148,9 @@ def FusedLoc : Builtin_LocationAttr<"FusedLoc"> {
 // NameLoc
 //===----------------------------------------------------------------------===//
 
-def NameLoc : Builtin_LocationAttr<"NameLoc"> {
+def NameLoc : Builtin_LocationAttr<"NameLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A named source location";
   let description = [{
     Syntax:
@@ -180,7 +187,9 @@ def NameLoc : Builtin_LocationAttr<"NameLoc"> {
 // OpaqueLoc
 //===----------------------------------------------------------------------===//
 
-def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc"> {
+def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "An opaque source location";
   let description = [{
     An instance of this location essentially contains a pointer to some data

diff --git a/mlir/include/mlir/IR/Location.h b/mlir/include/mlir/IR/Location.h
index 424d93d..fc3ee12 100644 (file)
--- a/mlir/include/mlir/IR/Location.h
+++ b/mlir/include/mlir/IR/Location.h
@@ -15,6 +15,7 @@
 #define MLIR_IR_LOCATION_H
 
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/SubElementInterfaces.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {

diff --git a/mlir/lib/IR/Location.cpp b/mlir/lib/IR/Location.cpp
index ce88b24..cd19d59 100644 (file)
--- a/mlir/lib/IR/Location.cpp
+++ b/mlir/lib/IR/Location.cpp
@@ -80,6 +80,20 @@ CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames) {
   return CallSiteLoc::get(name, caller);
 }
 
+void CallSiteLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getCallee());
+  walkAttrsFn(getCaller());
+}
+
+Attribute
+CallSiteLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                         ArrayRef<Type> replTypes) const {
+  return get(replAttrs[0].cast<LocationAttr>(),
+             replAttrs[1].cast<LocationAttr>());
+}
+
 //===----------------------------------------------------------------------===//
 // FusedLoc
 //===----------------------------------------------------------------------===//
@@ -121,3 +135,55 @@ Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
 
   return Base::get(context, locs, metadata);
 }
+
+void FusedLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  for (Attribute attr : getLocations())
+    walkAttrsFn(attr);
+  walkAttrsFn(getMetadata());
+}
+
+Attribute
+FusedLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                      ArrayRef<Type> replTypes) const {
+  SmallVector<Location> newLocs;
+  newLocs.reserve(replAttrs.size() - 1);
+  for (Attribute attr : replAttrs.drop_back())
+    newLocs.push_back(attr.cast<LocationAttr>());
+  return get(getContext(), newLocs, replAttrs.back());
+}
+
+//===----------------------------------------------------------------------===//
+// NameLoc
+//===----------------------------------------------------------------------===//
+
+void NameLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getName());
+  walkAttrsFn(getChildLoc());
+}
+
+Attribute NameLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                               ArrayRef<Type> replTypes) const {
+  return get(replAttrs[0].cast<StringAttr>(),
+             replAttrs[1].cast<LocationAttr>());
+}
+
+//===----------------------------------------------------------------------===//
+// OpaqueLoc
+//===----------------------------------------------------------------------===//
+
+void OpaqueLoc::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getFallbackLocation());
+}
+
+Attribute
+OpaqueLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                       ArrayRef<Type> replTypes) const {
+  return get(getUnderlyingLocation(), getUnderlyingTypeID(),
+             replAttrs[0].cast<LocationAttr>());
+}