Avoiding allocations during argument attribute conversion.

PiperOrigin-RevId: 239144675

3 files changed, 28 insertions, 19 deletions

diff --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index 426281bab67..84792d0582b 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -190,9 +190,12 @@ protected:
   /// The default behavior of this function is to call convertType on individual
   /// function operands and results, and then create a new MLIR function type
   /// from those.
-  virtual std::pair<FunctionType, std::vector<NamedAttributeList>>
-  convertFunctionSignatureType(FunctionType t,
-                               ArrayRef<NamedAttributeList> argAttrs);
+  ///
+  /// Post-condition: if the returned optional attribute list does not have
+  /// a value then no argument attribute conversion happened.
+  virtual FunctionType convertFunctionSignatureType(
+      FunctionType t, ArrayRef<NamedAttributeList> argAttrs,
+      SmallVectorImpl<NamedAttributeList> &convertedArgAttrs);
 };
 
 } // end namespace mlir
diff --git a/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp b/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp
index f6d2f74194d..e9605fa2bfe 100644
--- a/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp
+++ b/mlir/lib/LLVMIR/Transforms/ConvertToLLVMDialect.cpp
@@ -1118,11 +1118,14 @@ protected:
   }
 
   // Convert function signatures using the stored LLVM IR module.
-  std::pair<FunctionType, std::vector<NamedAttributeList>>
-  convertFunctionSignatureType(FunctionType t,
-                               ArrayRef<NamedAttributeList> argAttrs) override {
-    auto convertedType = TypeConverter::convertFunctionSignature(t, *module);
-    return std::make_pair(convertedType, argAttrs.vec());
+  FunctionType convertFunctionSignatureType(
+      FunctionType t, ArrayRef<NamedAttributeList> argAttrs,
+      SmallVectorImpl<NamedAttributeList> &convertedArgAttrs) override {
+
+    convertedArgAttrs.reserve(argAttrs.size());
+    for (auto attr : argAttrs)
+      convertedArgAttrs.push_back(attr);
+    return TypeConverter::convertFunctionSignature(t, *module);
   }
 
   // Make argument-taking successors of each block distinct.  PHI nodes in LLVM
diff --git a/mlir/lib/Transforms/DialectConversion.cpp b/mlir/lib/Transforms/DialectConversion.cpp
index 72a378b7d9a..b033dadfe51 100644
--- a/mlir/lib/Transforms/DialectConversion.cpp
+++ b/mlir/lib/Transforms/DialectConversion.cpp
@@ -215,13 +215,11 @@ Function *impl::FunctionConversion::convertFunction(Function *f) {
     return nullptr;
   };
 
-  // Create a new function with argument types and result types converted.  Wrap
+  // Create a new function with argument types and result types converted. Wrap
   // it into a unique_ptr to make sure it is cleaned up in case of error.
-  auto newFunctionSig = dialectConversion->convertFunctionSignatureType(
-      f->getType(), f->getAllArgAttrs());
-  Type newFunctionType = newFunctionSig.first;
-  std::vector<NamedAttributeList> newFunctionArgAttrs = newFunctionSig.second;
-
+  SmallVector<NamedAttributeList, 4> newFunctionArgAttrs;
+  Type newFunctionType = dialectConversion->convertFunctionSignatureType(
+      f->getType(), f->getAllArgAttrs(), newFunctionArgAttrs);
   if (!newFunctionType)
     return emitError("could not convert function type");
   auto newFunction = llvm::make_unique<Function>(
@@ -310,9 +308,9 @@ LogicalResult impl::FunctionConversion::run(Module *module) {
 
 // Create a function type with arguments and results converted, and argument
 // attributes passed through.
-std::pair<FunctionType, std::vector<NamedAttributeList>>
-DialectConversion::convertFunctionSignatureType(
-    FunctionType type, ArrayRef<NamedAttributeList> argAttrs) {
+FunctionType DialectConversion::convertFunctionSignatureType(
+    FunctionType type, ArrayRef<NamedAttributeList> argAttrs,
+    SmallVectorImpl<NamedAttributeList> &convertedArgAttrs) {
   SmallVector<Type, 8> arguments;
   SmallVector<Type, 4> results;
 
@@ -324,8 +322,13 @@ DialectConversion::convertFunctionSignatureType(
   for (auto t : type.getResults())
     results.push_back(convertType(t));
 
-  return std::make_pair(
-      FunctionType::get(arguments, results, type.getContext()), argAttrs.vec());
+  // Note this will cause an extra allocation only if we need
+  // to grow the caller-provided resulting attribute vector.
+  convertedArgAttrs.reserve(arguments.size());
+  for (auto attr : argAttrs)
+    convertedArgAttrs.push_back(attr);
+
+  return FunctionType::get(arguments, results, type.getContext());
 }
 
 LogicalResult DialectConversion::convert(Module *m) {