[Mlir-commits] [mlir] [mlir][func]: Introduce ReplaceFuncSignature tranform operation (PR #143381)

[Aviad Cohen]

================
@@ -0,0 +1,103 @@
+//===- Utils.cpp - Utilities to support the Func dialect ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements utilities for the Func dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/Utils/Utils.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+
+func::FuncOp func::replaceFuncWithNewOrder(func::FuncOp funcOp,
+                                           ArrayRef<int> newArgsOrder,
+                                           ArrayRef<int> newResultsOrder) {
+  // Generate an empty new function operation with the same name as the
+  // original.
+  assert(funcOp.getNumArguments() == newArgsOrder.size());
+  assert(funcOp.getNumResults() == newResultsOrder.size());
+  auto origInputTypes = funcOp.getFunctionType().getInputs();
+  auto origOutputTypes = funcOp.getFunctionType().getResults();
+  SmallVector<Type> newInputTypes, newOutputTypes;
+  for (unsigned int i = 0; i < origInputTypes.size(); ++i)
+    newInputTypes.push_back(origInputTypes[newArgsOrder[i]]);
+  for (unsigned int i = 0; i < origOutputTypes.size(); ++i)
+    newOutputTypes.push_back(origOutputTypes[newResultsOrder[i]]);
+  IRRewriter rewriter(funcOp);
+  rewriter.setInsertionPoint(funcOp);
+  auto newFuncOp = rewriter.create<func::FuncOp>(
+      funcOp.getLoc(), funcOp.getName(),
+      rewriter.getFunctionType(newInputTypes, newOutputTypes));
+  newFuncOp.addEntryBlock();
+  newFuncOp.setVisibility(funcOp.getVisibility());
+  newFuncOp->setDiscardableAttrs(funcOp->getDiscardableAttrDictionary());
+
+  // Map the arguments of the original function to the new function in
+  // the new order and adjust the attributes accordingly.
+  IRMapping operandMapper;
+  SmallVector<DictionaryAttr> argAttrs, resultAttrs;
+  funcOp.getAllArgAttrs(argAttrs);
+  for (unsigned int i = 0; i < newArgsOrder.size(); ++i) {
+    operandMapper.map(funcOp.getArgument(newArgsOrder[i]),
+                      newFuncOp.getArgument(i));
+    newFuncOp.setArgAttrs(i, argAttrs[newArgsOrder[i]]);
+  }
+  funcOp.getAllResultAttrs(resultAttrs);
+  for (unsigned int i = 0; i < newResultsOrder.size(); ++i)
+    newFuncOp.setResultAttrs(i, resultAttrs[newResultsOrder[i]]);
+
+  // Clone the operations from the original function to the new function.
+  rewriter.setInsertionPointToStart(&newFuncOp.getBody().front());
+  for (Operation &op : funcOp.getOps())
+    rewriter.clone(op, operandMapper);
+
+  // Handle the return operation.
+  auto returnOp = cast<func::ReturnOp>(
+      newFuncOp.getFunctionBody().begin()->getTerminator());
----------------
AviadCo wrote:

Ack, limited the util to one block.

https://github.com/llvm/llvm-project/pull/143381