When changing the type of a Function, also update the type of the respective...

--

PiperOrigin-RevId: 249137558

diff --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h
index 00730a9..acc5e72 100644 (file)
--- a/mlir/include/mlir/IR/Attributes.h
+++ b/mlir/include/mlir/IR/Attributes.h
@@ -410,6 +410,10 @@ public:
   /// This function is used by the internals of the Function class to null out
   /// attributes referring to functions that are about to be deleted.
   static void dropFunctionReference(Function *value);
+
+  /// This function is used by the internals of the Function class to update the
+  /// type of the function attribute for 'value'.
+  static void resetType(Function *value);
 };
 
 /// A base attribute that represents a reference to a vector or tensor constant.

diff --git a/mlir/include/mlir/IR/Function.h b/mlir/include/mlir/IR/Function.h
index 9aefd79..1a7332d 100644 (file)
--- a/mlir/include/mlir/IR/Function.h
+++ b/mlir/include/mlir/IR/Function.h
@@ -72,6 +72,7 @@ public:
   void setType(FunctionType newType) {
     type = newType;
     argAttrs.resize(type.getNumInputs());
+    FunctionAttr::resetType(this);
   }
 
   MLIRContext *getContext();

diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 8085418..b657bbd 100644 (file)
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -318,6 +318,9 @@ struct FunctionAttributeStorage : public AttributeStorage {
     value = nullptr;
   }
 
+  /// Reset the type of this attribute to the type of the held function.
+  void resetType() { setType(value->getType()); }
+
   Function *value;
 };
 

diff --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index b958221..5bb4554 100644 (file)
--- a/mlir/lib/IR/Attributes.cpp
+++ b/mlir/lib/IR/Attributes.cpp
@@ -303,6 +303,12 @@ void FunctionAttr::dropFunctionReference(Function *value) {
                                         StandardAttributes::Function, value);
 }
 
+/// This function is used by the internals of the Function class to update the
+/// type of the attribute for 'value'.
+void FunctionAttr::resetType(Function *value) {
+  FunctionAttr::get(value).getImpl()->resetType();
+}
+
 Function *FunctionAttr::getValue() const { return getImpl()->value; }
 
 FunctionType FunctionAttr::getType() const {