[Mlir-commits] [mlir] 1da1141 - [mlir] Implement the SubElement interfaces for the builtin locations

[River Riddle]

Author: River Riddle
Date: 2022-10-21T15:32:36-07:00
New Revision: 1da1141830718353fa7a38beada710d399d760dc

URL: https://github.com/llvm/llvm-project/commit/1da1141830718353fa7a38beada710d399d760dc
DIFF: https://github.com/llvm/llvm-project/commit/1da1141830718353fa7a38beada710d399d760dc.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/IR/BuiltinLocationAttributes.td
    mlir/include/mlir/IR/Location.h
    mlir/lib/IR/Location.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/BuiltinLocationAttributes.td b/mlir/include/mlir/IR/BuiltinLocationAttributes.td
index cc483cd7bed5..ca96fb9e53bb 100644
--- 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 424d93ddf6c6..fc3ee122060a 100644
--- 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 ce88b244a90d..cd19d59f2167 100644
--- 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>());
+}