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

[Oleksandr Alex Zinenko]

================
@@ -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());
+  SmallVector<Value> newReturnValues;
+  for (unsigned int i = 0; i < newResultsOrder.size(); ++i)
+    newReturnValues.push_back(returnOp.getOperand(newResultsOrder[i]));
+  rewriter.setInsertionPoint(returnOp);
+  auto newReturnOp =
+      rewriter.create<func::ReturnOp>(newFuncOp.getLoc(), newReturnValues);
+  newReturnOp->setDiscardableAttrs(returnOp->getDiscardableAttrDictionary());
----------------
ftynse wrote:

Not really. It's a thorny question we've had for a while with no good general solution until now. Most transformations drop discardable attributes mechanically, by not actively trying to preserve them. Proper preservation would involve some sort of abstracted check for whether an attribute makes sense on new operation or operations.

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