[mlir] Implement replacement of SymbolRefAttrs in Dialect attributes using SubElementAttr interface

[zero9178]

Summary

This patch extends the SubElementAttr interface to allow replacing a contained sub attribute. The attribute that should be replaced is identified by an index which denotes the n-th element returned by the accompanying walkImmediateSubElements method.

Using this addition the patch implements replacing SymbolRefAttrs contained within any dialect attributes.

[rriddle]

Awesome! Thanks for taking this on. This has been a TODO for a long time (and more recently since I added the sub element interfaces).

Can you also update the documentation for SymbolRefAttr to remove the restrictions on where it can be held? (https://github.com/llvm/llvm-project/blob/998e067a0a579bd483c2a2d411385e58494ffbca/mlir/include/mlir/IR/BuiltinAttributes.td#L951)

mlir/include/mlir/IR/BuiltinAttributes.td

69

Please wrap tablegen files at 80 characters.

mlir/include/mlir/IR/SubElementInterfaces.td

38–41

Can you add an explanation for the result value?

43

Can you drop the std:: here? We don't do that anywhere else.

46

Do we need a default implementation? Is there a downside to enforcing this method be defined? This would also remove the error handling that would arise from the default implementation.

mlir/lib/IR/BuiltinAttributes.cpp

56

Please run clang-format, and also drop the std:: from size_t.

58

nit: Spell out auto here.

225

Same here.

mlir/lib/IR/SymbolTable.cpp

494–495

You should be able to do walkSubAttrs here.

500

Prefer proper constructor calls.

832

Drop the llvm:: here.

833–835

walkSubAttrs?

841–842

Why not update all attributes and then rebuild once at the end? Otherwise, this may result in a bunch of uniquing that isn't necessary

[zero9178]

Addressed reviewers comments:

• Fixed style issues
• Changed corresponding documentation
• Avoid redundant uniquing by replacing attributes in a batch

[zero9178]

mlir/include/mlir/IR/SubElementInterfaces.td

46

I am not quite sure about this one as well. The cases I was thinking of in particular were types and attributes that do not contain any attributes.
In this case simply calling this operation is invalid. As this is a clear programmer error I chose to use an assert statement for now. As there are even many types in builtin alone that implement SubElementTypeInterface but don't have attributes I felt like this might warrant a default implementation.

mlir/lib/IR/SymbolTable.cpp

833–835

walkSubAttrs recursively walks through the attributes and does so in a depth first order which would yield wrong indices and more attributes than can be replaced by the interface. replaceImmediateSubAttribute is meant as a counterpart to walkImmediateSubElements as they should only be concerned with attributes that are immediately contained within the attribute.

[rriddle]

mlir/include/mlir/IR/SubElementInterfaces.td

47–48

Can you use llvm_unreachable instead?

mlir/lib/IR/SymbolTable.cpp

833–835

I forgot originally, but if we expand support we will also likely need to support references inside of Types as well. That can likely be deferred, but the documentation will need to indicate that holding references inside of Types (including attributes within those types) is not supported.

[zero9178]

Address reviewer comments:

• Document that symbol reference attributes are currently not supported inside of types
• Use llvm_unreachble instead of assert

[zero9178]

Friendly bump :)

[rriddle]

In D111357#3074452, @zero9178 wrote:

Friendly bump :)

Thanks for the ping, appreciate it! Got bogged down with some other work.

Did a scan, LGTM.

[GMNGeoffrey]

I'm getting weird intermittent build errors that I think are associated with this change. So far:

ld.lld: error: undefined symbol: mlir::SubElementAttrInterface::replaceImmediateSubAttribute(llvm::ArrayRef<std::pair<unsigned long, mlir::Attribute> >) const
>>> referenced by SymbolTable.cpp

In file included from /usr/local/google/home/gcmn/src/iree/third_party/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h:13:
/usr/local/google/home/gcmn/src/iree/third_party/llvm-project/mlir/include/mlir/IR/SubElementInterfaces.h:21:10: fatal error: 'mlir/IR/SubElementAttrInterfaces.h.inc' file not found
#include "mlir/IR/SubElementAttrInterfaces.h.inc"

My guess is that this is something missing in the CMake deps, but I actually don't really understand how that can cause a linker error. I'm poking around, but wanted to give a heads up in case others are seeing the same thing

[GMNGeoffrey]

In D111357#3104749, @GMNGeoffrey wrote:

I'm getting weird intermittent build errors that I think are associated with this change.

My theory is now that this was an issue in mlir-hlo with missing deps and somehow that *also* poisoned my ccache. This change probably just tickled something.

mlir/include/mlir/IR/BuiltinAttributes.td

def Builtin_AffineMapAttr : Builtin_Attr<"AffineMap"> {
   let typeBuilder = "IndexType::get($_value.getContext())";
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_ArrayAttr : Builtin_Attr<"Array", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    DeclareAttrInterfaceMethods<SubElementAttrInterface, ["replaceImmediateSubAttribute"]>
   ]> {
   let summary = "A collection of other Attribute values";
   let description = [{
     Syntax:
 
     ```
     array-attribute ::= `[` (attribute-value (`,` attribute-value)*)? `]`
     ```
def Builtin_DenseStringElementsAttr : Builtin_Attr<
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    DeclareAttrInterfaceMethods<SubElementAttrInterface, ["replaceImmediateSubAttribute"]>
   ]> {
   let summary = "An dictionary of named Attribute values";
   let description = [{
     Syntax:
 
     ```
     dictionary-attribute ::= `{` (attribute-entry (`,` attribute-entry)*)? `}`
     ```

mlir/include/mlir/IR/SubElementInterfaces.td

list<InterfaceMethod> methods = [
     InterfaceMethod<
       /*desc=*/[{
         Walk all of the immediately nested sub-attributes and sub-types. This
         method does not recurse into sub elements.
       }], "void", "walkImmediateSubElements",
       (ins "llvm::function_ref<void(mlir::Attribute)>":$walkAttrsFn,
            "llvm::function_ref<void(mlir::Type)>":$walkTypesFn)
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Replace the attribute identified by the index with the specified value.
+        The index is derived from the order of the attributes returned by the
+        attribute callback of `walkImmediateSubElements`. An index of 0 would
+        replace the very first attribute given by `walkImmediateSubElements`
+      }], cppNamespace # "::" # interfaceName, "replaceImmediateSubAttribute",
+      (ins "std::size_t":$index, "::mlir::Attribute":$value),
+      [{}],
+      /*defaultImplementation=*/[{
+        return {};
+      }]
+    >,
   ];
 
   code extraClassDeclaration = [{
     /// Walk all of the held sub-attributes.
     void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
       walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
     }
 

mlir/lib/IR/BuiltinAttributes.cpp

 
 void ArrayAttr::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Attribute attr : getValue())
     walkAttrsFn(attr);
 }
 
+SubElementAttrInterface ArrayAttr::replaceImmediateSubAttribute(std::size_t index,
+                                                  Attribute value) const {
+  auto vector = getValue().vec();
+  vector[index] = value;
+  return get(getContext(), vector);
+}
+
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 /// Helper function that does either an in place sort or sorts from source array
 /// into destination. If inPlace then storage is both the source and the
 /// destination, else value is the source and storage destination. Returns
 /// whether source was sorted.

 
 void DictionaryAttr::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Attribute attr : llvm::make_second_range(getValue()))
     walkAttrsFn(attr);
 }
 
+SubElementAttrInterface DictionaryAttr::replaceImmediateSubAttribute(std::size_t index,
+                                                       Attribute value) const {
+  auto vec = getValue().vec();
+  vec[index].second = value;
+  // The above only modifies the mapped value, but not the key, and therefore
+  // not the order of the elements. It remains sorted
+  return getWithSorted(getContext(), vec);
+}
+
 //===----------------------------------------------------------------------===//
 // StringAttr
 //===----------------------------------------------------------------------===//
 
 StringAttr StringAttr::getEmptyStringAttrUnchecked(MLIRContext *context) {
   return Base::get(context, "", NoneType::get(context));
 }
 

mlir/lib/IR/SymbolTable.cpp

static WalkResult walkSymbolRefs(
     function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
   // Check to see if the operation has any attributes.
   DictionaryAttr attrDict = op->getAttrDictionary();
   if (attrDict.empty())
     return WalkResult::advance();
 
   // A worklist of a container attribute and the current index into the held
   // attribute list.
-  SmallVector<Attribute, 1> attrWorklist(1, attrDict);
+  struct WorklistItem {
+    SubElementAttrInterface container;
+    SmallVector<Attribute> immediateSubElements;
+
+    explicit WorklistItem(SubElementAttrInterface container) {
+      SmallVector<Attribute> subElements;
+      container.walkImmediateSubElements(
+          [&](Attribute attr) { subElements.push_back(attr); }, [](Type) {});
+      immediateSubElements = std::move(subElements);
+    }
+  };
+
+  SmallVector<WorklistItem, 1> attrWorklist(1, WorklistItem{attrDict});
   SmallVector<int, 1> curAccessChain(1, /*Value=*/-1);
 
   // Process the symbol references within the given nested attribute range.
-  auto processAttrs = [&](int &index, auto attrRange) -> WalkResult {
-    for (Attribute attr : llvm::drop_begin(attrRange, index)) {
+  auto processAttrs = [&](int &index,
+                          WorklistItem &worklistItem) -> WalkResult {
+    for (Attribute attr :
+         llvm::drop_begin(worklistItem.immediateSubElements, index)) {
       /// Check for a nested container attribute, these will also need to be
       /// walked.
-      if (attr.isa<ArrayAttr, DictionaryAttr>()) {
-        attrWorklist.push_back(attr);
+      if (auto interface = attr.dyn_cast<SubElementAttrInterface>()) {
+        attrWorklist.emplace_back(interface);
         curAccessChain.push_back(-1);
         return WalkResult::advance();
       }
 
       // Invoke the provided callback if we find a symbol use and check for a
       // requested interrupt.
       if (auto symbolRef = attr.dyn_cast<SymbolRefAttr>())
         if (callback({op, symbolRef}, curAccessChain).wasInterrupted())
           return WalkResult::interrupt();
 
       // Make sure to keep the index counter in sync.
       ++index;
     }
 
     // Pop this container attribute from the worklist.
     attrWorklist.pop_back();
     curAccessChain.pop_back();
     return WalkResult::advance();
   };
 
   WalkResult result = WalkResult::advance();
   do {
-    Attribute attr = attrWorklist.back();
+    WorklistItem &item = attrWorklist.back();
     int &index = curAccessChain.back();
     ++index;
 
     // Process the given attribute, which is guaranteed to be a container.
-    if (auto dict = attr.dyn_cast<DictionaryAttr>())
-      result = processAttrs(index, make_second_range(dict.getValue()));
-    else
-      result = processAttrs(index, attr.cast<ArrayAttr>().getValue());
+    result = processAttrs(index, item);
   } while (!attrWorklist.empty() && !result.wasInterrupted());
   return result;
 }
 
 /// Walk all of the uses, for any symbol, that are nested within the given
 /// regions, invoking the provided callback for each. This does not traverse
 /// into any nested symbol tables.
 static Optional<WalkResult> walkSymbolUses(
bool SymbolTable::symbolKnownUseEmpty(Operation *symbol, Region *from) {
   return symbolKnownUseEmptyImpl(symbol, from);
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolTable::replaceAllSymbolUses
 
 /// Rebuild the given attribute container after replacing all references to a
 /// symbol with the updated attribute in 'accesses'.
-static Attribute rebuildAttrAfterRAUW(
-    Attribute container,
+static SubElementAttrInterface rebuildAttrAfterRAUW(
+    SubElementAttrInterface container,
     ArrayRef<std::pair<SmallVector<int, 1>, SymbolRefAttr>> accesses,
     unsigned depth) {
   // Given a range of Attributes, update the ones referred to by the given
   // access chains to point to the new symbol attribute.
-  auto updateAttrs = [&](auto &&attrRange) {
-    auto attrBegin = std::begin(attrRange);
-    for (unsigned i = 0, e = accesses.size(); i != e;) {
-      ArrayRef<int> access = accesses[i].first;
-      Attribute &attr = *std::next(attrBegin, access[depth]);
+
+  llvm::SmallVector<Attribute> subElements;
+  container.walkImmediateSubElements(
+      [&](Attribute attribute) { subElements.push_back(attribute); },
+      [](Type) {});
 
-      // Check to see if this is a leaf access, i.e. a SymbolRef.
-      if (access.size() == depth + 1) {
-        attr = accesses[i].second;
-        ++i;
-        continue;
-      }
+      container = container.replaceImmediateSubAttribute(access.back(),
+                                                         accesses[i].second);
 
-      // Otherwise, this is a container. Collect all of the accesses for this
-      // index and recurse. The recursion here is bounded by the size of the
-      // largest access array.
-      auto nestedAccesses = accesses.drop_front(i).take_while([&](auto &it) {
-        ArrayRef<int> nextAccess = it.first;
-        return nextAccess.size() > depth + 1 &&
-               nextAccess[depth] == access[depth];
-      });
-      attr = rebuildAttrAfterRAUW(attr, nestedAccesses, depth + 1);
+    auto result = rebuildAttrAfterRAUW(subElements[access[depth]],
+                                       nestedAccesses, depth + 1);
+    container = container.replaceImmediateSubAttribute(access[depth], result);
 
-      // Skip over all of the accesses that refer to the nested container.
-      i += nestedAccesses.size();
-    }
-  };
 
-  if (auto dictAttr = container.dyn_cast<DictionaryAttr>()) {
-    auto newAttrs = llvm::to_vector<4>(dictAttr.getValue());
-    updateAttrs(make_second_range(newAttrs));
-    return DictionaryAttr::get(dictAttr.getContext(), newAttrs);
-  }
-  auto newAttrs = llvm::to_vector<4>(container.cast<ArrayAttr>().getValue());
-  updateAttrs(newAttrs);
-  return ArrayAttr::get(container.getContext(), newAttrs);
+  return container;
 }
 
 /// Generates a new symbol reference attribute with a new leaf reference.
 static SymbolRefAttr generateNewRefAttr(SymbolRefAttr oldAttr,
                                         FlatSymbolRefAttr newLeafAttr) {
   if (oldAttr.isa<FlatSymbolRefAttr>())
     return newLeafAttr;
   auto nestedRefs = llvm::to_vector<2>(oldAttr.getNestedReferences());

mlir/test/IR/test-symbol-rauw.mlir

 // Check that the replacement fails for potentially unknown symbol tables.
 module {
   // CHECK: func private @failed_repl
   func private @failed_repl() attributes {sym.new_name = "replaced_name" }
 
   "foo.possibly_unknown_symbol_table"() ({
   }) : () -> ()
 }
+
+// -----
+
+// Check that replacement works in any implementations of SubElementsAttrInterface
+module {
+    // CHECK: func private @replaced_foo
+    func private @symbol_foo() attributes {sym.new_name = "replaced_foo" }
+
+    // CHECK: func @symbol_bar
+    func @symbol_bar() {
+      // CHECK: foo.op
+      // CHECK-SAME: non_symbol_attr,
+      // CHECK-SAME: use = [#test.sub_elements_access<[@replaced_foo], @symbol_bar, @replaced_foo>],
+      // CHECK-SAME: z_non_symbol_attr_3
+      "foo.op"() {
+        non_symbol_attr,
+        use = [#test.sub_elements_access<[@symbol_foo],@symbol_bar,@symbol_foo>],
+        z_non_symbol_attr_3
+      } : () -> ()
+    }
+}

mlir/test/lib/Dialect/Test/TestAttrDefs.td

 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_ATTRDEFS
 #define TEST_ATTRDEFS
 
 // To get the test dialect definition.
 include "TestOps.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
+include "mlir/IR/SubElementInterfaces.td"
 
 // All of the attributes will extend this class.
 class Test_Attr<string name, list<Trait> traits = []>
     : AttrDef<Test_Dialect, name, traits>;
 
 def SimpleAttrA : Test_Attr<"SimpleA"> {
   let mnemonic = "smpla";
 }
auto value_begin_impl(OverloadToken<llvm::APInt>) const {
       return llvm::map_range(getElements(), [=](uint64_t value) {
         return llvm::APInt(/*numBits=*/64, value);
       }).begin();
     }
   }];
   let genVerifyDecl = 1;
 }
 
+def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface, ["replaceImmediateSubAttribute"]>
+  ]> {
+
+  let mnemonic = "sub_elements_access";
+
+  let parameters = (ins
+    "::mlir::Attribute":$first,
+    "::mlir::Attribute":$second,
+    "::mlir::Attribute":$third
+  );
+}
+
 #endif // TEST_ATTRDEFS

mlir/test/lib/Dialect/Test/TestAttributes.cpp

TestI64ElementsAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   }
   if (type.getRank() != 1 || !type.getElementType().isSignlessInteger(64))
     return emitError() << "expected single rank 64-bit shape type, but got: "
                        << type;
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// TestSubElementsAccessAttr
+//===----------------------------------------------------------------------===//
+
+Attribute TestSubElementsAccessAttr::parse(::mlir::DialectAsmParser &parser,
+                                           ::mlir::Type type) {
+  Attribute first, second, third;
+  if (parser.parseLess() || parser.parseAttribute(first) ||
+      parser.parseComma() || parser.parseAttribute(second) ||
+      parser.parseComma() || parser.parseAttribute(third) ||
+      parser.parseGreater()) {
+    return {};
+  }
+  return get(parser.getContext(), first, second, third);
+}
+
+void TestSubElementsAccessAttr::print(
+    ::mlir::DialectAsmPrinter &printer) const {
+  printer << getMnemonic() << "<" << getFirst() << ", " << getSecond() << ", "
+          << getThird() << ">";
+}
+
+void TestSubElementsAccessAttr::walkImmediateSubElements(
+    llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
+    llvm::function_ref<void(mlir::Type)> walkTypesFn) const {
+  walkAttrsFn(getFirst());
+  walkAttrsFn(getSecond());
+  walkAttrsFn(getThird());
+}
+
+SubElementAttrInterface TestSubElementsAccessAttr::replaceImmediateSubAttribute(
+    std::size_t index, ::mlir::Attribute value) const {
+  switch (index) {
+  case 0:
+    return get(getContext(), value, getSecond(), getThird());
+  case 1:
+    return get(getContext(), getFirst(), value, getThird());
+  case 2:
+    return get(getContext(), getFirst(), getSecond(), value);
+  }
+  LLVM_BUILTIN_UNREACHABLE;
+}
+
+//===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//
 
 #include "TestAttrInterfaces.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.cpp.inc"