Simplify several usages of attributes now that they always have a type and, transitively, access to the context.

    This also fixes a bug where FunctionAttrs were not being remapped for function and function argument attributes.

--

PiperOrigin-RevId: 246876924

2 files changed, 21 insertions, 10 deletions

diff --git a/mlir/lib/Transforms/DialectConversion.cpp b/mlir/lib/Transforms/DialectConversion.cpp
index 66c3b2dbd5e..831a68aa5bb 100644
--- a/mlir/lib/Transforms/DialectConversion.cpp
+++ b/mlir/lib/Transforms/DialectConversion.cpp
@@ -316,8 +316,8 @@ LogicalResult impl::FunctionConversion::run(Module *module) {
     if (!converted)
       return failure();
 
-    auto origFuncAttr = FunctionAttr::get(func, context);
-    auto convertedFuncAttr = FunctionAttr::get(converted, context);
+    auto origFuncAttr = FunctionAttr::get(func);
+    auto convertedFuncAttr = FunctionAttr::get(converted);
     convertedFuncs.push_back(converted);
     functionAttrRemapping.insert({origFuncAttr, convertedFuncAttr});
   }
diff --git a/mlir/lib/Transforms/Utils/Utils.cpp b/mlir/lib/Transforms/Utils/Utils.cpp
index 422d6b136ab..1ab821a9366 100644
--- a/mlir/lib/Transforms/Utils/Utils.cpp
+++ b/mlir/lib/Transforms/Utils/Utils.cpp
@@ -290,33 +290,44 @@ void mlir::createAffineComputationSlice(
   }
 }
 
-void mlir::remapFunctionAttrs(
-    Operation &op, const DenseMap<Attribute, FunctionAttr> &remappingTable) {
-  for (auto attr : op.getAttrs()) {
+static void
+remapFunctionAttrs(NamedAttributeList &attrs,
+                   const DenseMap<Attribute, FunctionAttr> &remappingTable) {
+  for (auto attr : attrs.getAttrs()) {
     // Do the remapping, if we got the same thing back, then it must contain
     // functions that aren't getting remapped.
-    auto newVal =
-        attr.second.remapFunctionAttrs(remappingTable, op.getContext());
+    auto newVal = attr.second.remapFunctionAttrs(remappingTable);
     if (newVal == attr.second)
       continue;
 
     // Otherwise, replace the existing attribute with the new one.  It is safe
     // to mutate the attribute list while we walk it because underlying
     // attribute lists are uniqued and immortal.
-    op.setAttr(attr.first, newVal);
+    attrs.set(attr.first, newVal);
   }
 }
 
 void mlir::remapFunctionAttrs(
+    Operation &op, const DenseMap<Attribute, FunctionAttr> &remappingTable) {
+  ::remapFunctionAttrs(op.getAttrList(), remappingTable);
+}
+
+void mlir::remapFunctionAttrs(
     Function &fn, const DenseMap<Attribute, FunctionAttr> &remappingTable) {
 
+  // Remap the attributes of the function.
+  ::remapFunctionAttrs(fn.getAttrList(), remappingTable);
+
+  // Remap the attributes of the arguments of this function.
+  for (auto &attrList : fn.getAllArgAttrs())
+    ::remapFunctionAttrs(attrList, remappingTable);
+
   // Look at all operations in a Function.
   fn.walk([&](Operation *op) { remapFunctionAttrs(*op, remappingTable); });
 }
 
 void mlir::remapFunctionAttrs(
     Module &module, const DenseMap<Attribute, FunctionAttr> &remappingTable) {
-  for (auto &fn : module) {
+  for (auto &fn : module)
     remapFunctionAttrs(fn, remappingTable);
-  }
 }