mlir/lib/Dialect/Func/Utils/Utils.cpp - llvm-project - Git at Google

 //===- 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/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/DebugLog.h"

 #define DEBUG_TYPE "func-utils"

 using namespace mlir;

 /// This method creates an inverse mapping of the provided map `oldToNew`.
 /// Given an array where `oldIdxToNewIdx[i] = j` means old index `i` maps
 /// to new index `j`,
 /// This method returns a vector where `result[j]` contains all old indices
 /// that map to new index `j`.
 ///
 /// Example:
 /// ```
 /// oldIdxToNewIdx = [0, 1, 2, 2, 3]
 /// getInverseMapping(oldIdxToNewIdx) = [[0], [1], [2, 3], [4]]
 /// ```
 ///
 static llvm::SmallVector<llvm::SmallVector<int>>
 getInverseMapping(ArrayRef<int> oldIdxToNewIdx) {
   int numOfNewIdxs = 0;
   if (!oldIdxToNewIdx.empty())
     numOfNewIdxs = 1 + *llvm::max_element(oldIdxToNewIdx);
   llvm::SmallVector<llvm::SmallVector<int>> newToOldIdxs(numOfNewIdxs);
   for (auto [oldIdx, newIdx] : llvm::enumerate(oldIdxToNewIdx))
     newToOldIdxs[newIdx].push_back(oldIdx);
   return newToOldIdxs;
 }

 /// This method returns a new vector of elements that are mapped from the
 /// `origElements` based on the `newIdxToOldIdxs` mapping. This function assumes
 /// that the `newIdxToOldIdxs` mapping is valid, i.e. for each new index, there
 /// is at least one old index that maps to it. Also, It assumes that mapping to
 /// the same old index has the same element in the `origElements` vector.
 template <typename Element>
 static SmallVector<Element> getMappedElements(
     ArrayRef<Element> origElements,
     const llvm::SmallVector<llvm::SmallVector<int>> &newIdxToOldIdxs) {
   SmallVector<Element> newElements;
   for (const auto &oldIdxs : newIdxToOldIdxs) {
     assert(llvm::all_of(oldIdxs,
                         [&origElements](int idx) -> bool {
                           return idx >= 0 &&
                                  static_cast<size_t>(idx) < origElements.size();
                         }) &&
            "idx must be less than the number of elements in the original "
            "elements");
     assert(!oldIdxs.empty() && "oldIdx must not be empty");
     Element origTypeToCheck = origElements[oldIdxs.front()];
     assert(llvm::all_of(oldIdxs,
                         [&](int idx) -> bool {
                           return origElements[idx] == origTypeToCheck;
                         }) &&
            "all oldIdxs must be equal");
     newElements.push_back(origTypeToCheck);
   }
   return newElements;
 }

 FailureOr<func::FuncOp>
 func::replaceFuncWithNewMapping(RewriterBase &rewriter, func::FuncOp funcOp,
                                 ArrayRef<int> oldArgIdxToNewArgIdx,
                                 ArrayRef<int> oldResIdxToNewResIdx) {
   // Generate an empty new function operation with the same name as the
   // original.
   assert(funcOp.getNumArguments() == oldArgIdxToNewArgIdx.size() &&
          "oldArgIdxToNewArgIdx must match the number of arguments in the "
          "function");
   assert(
       funcOp.getNumResults() == oldResIdxToNewResIdx.size() &&
       "oldResIdxToNewResIdx must match the number of results in the function");

   if (!funcOp.getBody().hasOneBlock())
     return rewriter.notifyMatchFailure(
         funcOp, "expected function to have exactly one block");

   // We may have some duplicate arguments in the old function, i.e.
   // in the mapping `newArgIdxToOldArgIdxs` for some new argument index
   // there may be multiple old argument indices.
   llvm::SmallVector<llvm::SmallVector<int>> newArgIdxToOldArgIdxs =
       getInverseMapping(oldArgIdxToNewArgIdx);
   SmallVector<Type> newInputTypes = getMappedElements(
       funcOp.getFunctionType().getInputs(), newArgIdxToOldArgIdxs);

   SmallVector<Location> locs;
   for (const auto &oldArgIdxs : newArgIdxToOldArgIdxs)
     locs.push_back(funcOp.getArgument(oldArgIdxs.front()).getLoc());

   llvm::SmallVector<llvm::SmallVector<int>> newResToOldResIdxs =
       getInverseMapping(oldResIdxToNewResIdx);
   SmallVector<Type> newOutputTypes = getMappedElements(
       funcOp.getFunctionType().getResults(), newResToOldResIdxs);

   rewriter.setInsertionPoint(funcOp);
   auto newFuncOp = func::FuncOp::create(
       rewriter, funcOp.getLoc(), funcOp.getName(),
       rewriter.getFunctionType(newInputTypes, newOutputTypes));

   Region &newRegion = newFuncOp.getBody();
   rewriter.createBlock(&newRegion, newRegion.begin(), newInputTypes, locs);
   newFuncOp.setVisibility(funcOp.getVisibility());

   // 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 (auto [oldArgIdx, newArgIdx] : llvm::enumerate(oldArgIdxToNewArgIdx))
     operandMapper.map(funcOp.getArgument(oldArgIdx),
                       newFuncOp.getArgument(newArgIdx));
   for (auto [newArgIdx, oldArgIdx] : llvm::enumerate(newArgIdxToOldArgIdxs))
     newFuncOp.setArgAttrs(newArgIdx, argAttrs[oldArgIdx.front()]);

   funcOp.getAllResultAttrs(resultAttrs);
   for (auto [newResIdx, oldResIdx] : llvm::enumerate(newResToOldResIdxs))
     newFuncOp.setResultAttrs(newResIdx, resultAttrs[oldResIdx.front()]);

   // 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 (const auto &oldResIdxs : newResToOldResIdxs)
     newReturnValues.push_back(returnOp.getOperand(oldResIdxs.front()));

   rewriter.setInsertionPoint(returnOp);
   func::ReturnOp::create(rewriter, newFuncOp.getLoc(), newReturnValues);
   rewriter.eraseOp(returnOp);

   rewriter.eraseOp(funcOp);

   return newFuncOp;
 }

 func::CallOp
 func::replaceCallOpWithNewMapping(RewriterBase &rewriter, func::CallOp callOp,
                                   ArrayRef<int> oldArgIdxToNewArgIdx,
                                   ArrayRef<int> oldResIdxToNewResIdx) {
   assert(callOp.getNumOperands() == oldArgIdxToNewArgIdx.size() &&
          "oldArgIdxToNewArgIdx must match the number of operands in the call "
          "operation");
   assert(callOp.getNumResults() == oldResIdxToNewResIdx.size() &&
          "oldResIdxToNewResIdx must match the number of results in the call "
          "operation");

   SmallVector<Value> origOperands = callOp.getOperands();
   SmallVector<llvm::SmallVector<int>> newArgIdxToOldArgIdxs =
       getInverseMapping(oldArgIdxToNewArgIdx);
   SmallVector<Value> newOperandsValues =
       getMappedElements<Value>(origOperands, newArgIdxToOldArgIdxs);
   SmallVector<llvm::SmallVector<int>> newResToOldResIdxs =
       getInverseMapping(oldResIdxToNewResIdx);
   SmallVector<Type> origResultTypes = llvm::to_vector(callOp.getResultTypes());
   SmallVector<Type> newResultTypes =
       getMappedElements<Type>(origResultTypes, newResToOldResIdxs);

   // Replace the kernel call operation with a new one that has the
   // mapped arguments.
   rewriter.setInsertionPoint(callOp);
   auto newCallOp =
       func::CallOp::create(rewriter, callOp.getLoc(), callOp.getCallee(),
                            newResultTypes, newOperandsValues);
   newCallOp.setNoInlineAttr(callOp.getNoInlineAttr());
   for (auto &&[oldResIdx, newResIdx] : llvm::enumerate(oldResIdxToNewResIdx))
     rewriter.replaceAllUsesWith(callOp.getResult(oldResIdx),
                                 newCallOp.getResult(newResIdx));
   rewriter.eraseOp(callOp);

   return newCallOp;
 }

 FailureOr<std::pair<func::FuncOp, func::CallOp>>
 func::deduplicateArgsOfFuncOp(RewriterBase &rewriter, func::FuncOp funcOp,
                               ModuleOp moduleOp) {
   SmallVector<func::CallOp> callOps;
   auto traversalResult = moduleOp.walk([&](func::CallOp callOp) {
     if (callOp.getCallee() == funcOp.getSymName()) {
       if (!callOps.empty())
         // Only support one callOp for now
         return WalkResult::interrupt();
       callOps.push_back(callOp);
     }
     return WalkResult::advance();
   });

   if (traversalResult.wasInterrupted()) {
     LDBG() << "function " << funcOp.getName() << " has more than one callOp";
     return failure();
   }

   if (callOps.empty()) {
     LDBG() << "function " << funcOp.getName() << " does not have any callOp";
     return failure();
   }

   func::CallOp callOp = callOps.front();

   // Create mapping for arguments (deduplicate operands)
   SmallVector<int> oldArgIdxToNewArgIdx(callOp.getNumOperands());
   llvm::DenseMap<Value, int> valueToNewArgIdx;
   for (auto [operandIdx, operand] : llvm::enumerate(callOp.getOperands())) {
     auto [iterator, inserted] = valueToNewArgIdx.insert(
         {operand, static_cast<int>(valueToNewArgIdx.size())});
     // Reduce the duplicate operands and maintain the original order.
     oldArgIdxToNewArgIdx[operandIdx] = iterator->second;
   }

   bool hasDuplicateOperands =
       valueToNewArgIdx.size() != callOp.getNumOperands();
   if (!hasDuplicateOperands) {
     LDBG() << "function " << funcOp.getName()
            << " does not have duplicate operands";
     return failure();
   }

   // Create identity mapping for results (no deduplication needed)
   SmallVector<int> oldResIdxToNewResIdx(callOp.getNumResults());
   for (int resultIdx : llvm::seq<int>(0, callOp.getNumResults()))
     oldResIdxToNewResIdx[resultIdx] = resultIdx;

   // Apply the transformation to create new function and call operations
   FailureOr<func::FuncOp> newFuncOpOrFailure = replaceFuncWithNewMapping(
       rewriter, funcOp, oldArgIdxToNewArgIdx, oldResIdxToNewResIdx);
   if (failed(newFuncOpOrFailure)) {
     LDBG() << "failed to replace function signature with name "
            << funcOp.getName() << " with new order";
     return failure();
   }

   func::CallOp newCallOp = replaceCallOpWithNewMapping(
       rewriter, callOp, oldArgIdxToNewArgIdx, oldResIdxToNewResIdx);

   return std::make_pair(*newFuncOpOrFailure, newCallOp);
 }
	//===- 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/Dialect/Func/IR/FuncOps.h"
	#include "mlir/IR/IRMapping.h"
	#include "mlir/IR/PatternMatch.h"
	#include "llvm/ADT/STLExtras.h"
	#include "llvm/ADT/SmallVector.h"
	#include "llvm/Support/DebugLog.h"

	#define DEBUG_TYPE "func-utils"

	using namespace mlir;

	/// This method creates an inverse mapping of the provided map `oldToNew`.
	/// Given an array where `oldIdxToNewIdx[i] = j` means old index `i` maps
	/// to new index `j`,
	/// This method returns a vector where `result[j]` contains all old indices
	/// that map to new index `j`.
	///
	/// Example:
	/// ```
	/// oldIdxToNewIdx = [0, 1, 2, 2, 3]
	/// getInverseMapping(oldIdxToNewIdx) = [[0], [1], [2, 3], [4]]
	/// ```
	///
	static llvm::SmallVector<llvm::SmallVector<int>>
	getInverseMapping(ArrayRef<int> oldIdxToNewIdx) {
	int numOfNewIdxs = 0;
	if (!oldIdxToNewIdx.empty())
	numOfNewIdxs = 1 + *llvm::max_element(oldIdxToNewIdx);
	llvm::SmallVector<llvm::SmallVector<int>> newToOldIdxs(numOfNewIdxs);
	for (auto [oldIdx, newIdx] : llvm::enumerate(oldIdxToNewIdx))
	newToOldIdxs[newIdx].push_back(oldIdx);
	return newToOldIdxs;
	}

	/// This method returns a new vector of elements that are mapped from the
	/// `origElements` based on the `newIdxToOldIdxs` mapping. This function assumes
	/// that the `newIdxToOldIdxs` mapping is valid, i.e. for each new index, there
	/// is at least one old index that maps to it. Also, It assumes that mapping to
	/// the same old index has the same element in the `origElements` vector.
	template <typename Element>
	static SmallVector<Element> getMappedElements(
	ArrayRef<Element> origElements,
	const llvm::SmallVector<llvm::SmallVector<int>> &newIdxToOldIdxs) {
	SmallVector<Element> newElements;
	for (const auto &oldIdxs : newIdxToOldIdxs) {
	assert(llvm::all_of(oldIdxs,
	[&origElements](int idx) -> bool {
	return idx >= 0 &&
	static_cast<size_t>(idx) < origElements.size();
	}) &&
	"idx must be less than the number of elements in the original "
	"elements");
	assert(!oldIdxs.empty() && "oldIdx must not be empty");
	Element origTypeToCheck = origElements[oldIdxs.front()];
	assert(llvm::all_of(oldIdxs,
	[&](int idx) -> bool {
	return origElements[idx] == origTypeToCheck;
	}) &&
	"all oldIdxs must be equal");
	newElements.push_back(origTypeToCheck);
	}
	return newElements;
	}

	FailureOr<func::FuncOp>
	func::replaceFuncWithNewMapping(RewriterBase &rewriter, func::FuncOp funcOp,
	ArrayRef<int> oldArgIdxToNewArgIdx,
	ArrayRef<int> oldResIdxToNewResIdx) {
	// Generate an empty new function operation with the same name as the
	// original.
	assert(funcOp.getNumArguments() == oldArgIdxToNewArgIdx.size() &&
	"oldArgIdxToNewArgIdx must match the number of arguments in the "
	"function");
	assert(
	funcOp.getNumResults() == oldResIdxToNewResIdx.size() &&
	"oldResIdxToNewResIdx must match the number of results in the function");

	if (!funcOp.getBody().hasOneBlock())
	return rewriter.notifyMatchFailure(
	funcOp, "expected function to have exactly one block");

	// We may have some duplicate arguments in the old function, i.e.
	// in the mapping `newArgIdxToOldArgIdxs` for some new argument index
	// there may be multiple old argument indices.
	llvm::SmallVector<llvm::SmallVector<int>> newArgIdxToOldArgIdxs =
	getInverseMapping(oldArgIdxToNewArgIdx);
	SmallVector<Type> newInputTypes = getMappedElements(
	funcOp.getFunctionType().getInputs(), newArgIdxToOldArgIdxs);

	SmallVector<Location> locs;
	for (const auto &oldArgIdxs : newArgIdxToOldArgIdxs)
	locs.push_back(funcOp.getArgument(oldArgIdxs.front()).getLoc());

	llvm::SmallVector<llvm::SmallVector<int>> newResToOldResIdxs =
	getInverseMapping(oldResIdxToNewResIdx);
	SmallVector<Type> newOutputTypes = getMappedElements(
	funcOp.getFunctionType().getResults(), newResToOldResIdxs);

	rewriter.setInsertionPoint(funcOp);
	auto newFuncOp = func::FuncOp::create(
	rewriter, funcOp.getLoc(), funcOp.getName(),
	rewriter.getFunctionType(newInputTypes, newOutputTypes));

	Region &newRegion = newFuncOp.getBody();
	rewriter.createBlock(&newRegion, newRegion.begin(), newInputTypes, locs);
	newFuncOp.setVisibility(funcOp.getVisibility());

	// 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 (auto [oldArgIdx, newArgIdx] : llvm::enumerate(oldArgIdxToNewArgIdx))
	operandMapper.map(funcOp.getArgument(oldArgIdx),
	newFuncOp.getArgument(newArgIdx));
	for (auto [newArgIdx, oldArgIdx] : llvm::enumerate(newArgIdxToOldArgIdxs))
	newFuncOp.setArgAttrs(newArgIdx, argAttrs[oldArgIdx.front()]);

	funcOp.getAllResultAttrs(resultAttrs);
	for (auto [newResIdx, oldResIdx] : llvm::enumerate(newResToOldResIdxs))
	newFuncOp.setResultAttrs(newResIdx, resultAttrs[oldResIdx.front()]);

	// 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 (const auto &oldResIdxs : newResToOldResIdxs)
	newReturnValues.push_back(returnOp.getOperand(oldResIdxs.front()));

	rewriter.setInsertionPoint(returnOp);
	func::ReturnOp::create(rewriter, newFuncOp.getLoc(), newReturnValues);
	rewriter.eraseOp(returnOp);

	rewriter.eraseOp(funcOp);

	return newFuncOp;
	}

	func::CallOp
	func::replaceCallOpWithNewMapping(RewriterBase &rewriter, func::CallOp callOp,
	ArrayRef<int> oldArgIdxToNewArgIdx,
	ArrayRef<int> oldResIdxToNewResIdx) {
	assert(callOp.getNumOperands() == oldArgIdxToNewArgIdx.size() &&
	"oldArgIdxToNewArgIdx must match the number of operands in the call "
	"operation");
	assert(callOp.getNumResults() == oldResIdxToNewResIdx.size() &&
	"oldResIdxToNewResIdx must match the number of results in the call "
	"operation");

	SmallVector<Value> origOperands = callOp.getOperands();
	SmallVector<llvm::SmallVector<int>> newArgIdxToOldArgIdxs =
	getInverseMapping(oldArgIdxToNewArgIdx);
	SmallVector<Value> newOperandsValues =
	getMappedElements<Value>(origOperands, newArgIdxToOldArgIdxs);
	SmallVector<llvm::SmallVector<int>> newResToOldResIdxs =
	getInverseMapping(oldResIdxToNewResIdx);
	SmallVector<Type> origResultTypes = llvm::to_vector(callOp.getResultTypes());
	SmallVector<Type> newResultTypes =
	getMappedElements<Type>(origResultTypes, newResToOldResIdxs);

	// Replace the kernel call operation with a new one that has the
	// mapped arguments.
	rewriter.setInsertionPoint(callOp);
	auto newCallOp =
	func::CallOp::create(rewriter, callOp.getLoc(), callOp.getCallee(),
	newResultTypes, newOperandsValues);
	newCallOp.setNoInlineAttr(callOp.getNoInlineAttr());
	for (auto &&[oldResIdx, newResIdx] : llvm::enumerate(oldResIdxToNewResIdx))
	rewriter.replaceAllUsesWith(callOp.getResult(oldResIdx),
	newCallOp.getResult(newResIdx));
	rewriter.eraseOp(callOp);

	return newCallOp;
	}

	FailureOr<std::pair<func::FuncOp, func::CallOp>>
	func::deduplicateArgsOfFuncOp(RewriterBase &rewriter, func::FuncOp funcOp,
	ModuleOp moduleOp) {
	SmallVector<func::CallOp> callOps;
	auto traversalResult = moduleOp.walk([&](func::CallOp callOp) {
	if (callOp.getCallee() == funcOp.getSymName()) {
	if (!callOps.empty())
	// Only support one callOp for now
	return WalkResult::interrupt();
	callOps.push_back(callOp);
	}
	return WalkResult::advance();
	});

	if (traversalResult.wasInterrupted()) {
	LDBG() << "function " << funcOp.getName() << " has more than one callOp";
	return failure();
	}

	if (callOps.empty()) {
	LDBG() << "function " << funcOp.getName() << " does not have any callOp";
	return failure();
	}

	func::CallOp callOp = callOps.front();

	// Create mapping for arguments (deduplicate operands)
	SmallVector<int> oldArgIdxToNewArgIdx(callOp.getNumOperands());
	llvm::DenseMap<Value, int> valueToNewArgIdx;
	for (auto [operandIdx, operand] : llvm::enumerate(callOp.getOperands())) {
	auto [iterator, inserted] = valueToNewArgIdx.insert(
	{operand, static_cast<int>(valueToNewArgIdx.size())});
	// Reduce the duplicate operands and maintain the original order.
	oldArgIdxToNewArgIdx[operandIdx] = iterator->second;
	}

	bool hasDuplicateOperands =
	valueToNewArgIdx.size() != callOp.getNumOperands();
	if (!hasDuplicateOperands) {
	LDBG() << "function " << funcOp.getName()
	<< " does not have duplicate operands";
	return failure();
	}

	// Create identity mapping for results (no deduplication needed)
	SmallVector<int> oldResIdxToNewResIdx(callOp.getNumResults());
	for (int resultIdx : llvm::seq<int>(0, callOp.getNumResults()))
	oldResIdxToNewResIdx[resultIdx] = resultIdx;

	// Apply the transformation to create new function and call operations
	FailureOr<func::FuncOp> newFuncOpOrFailure = replaceFuncWithNewMapping(
	rewriter, funcOp, oldArgIdxToNewArgIdx, oldResIdxToNewResIdx);
	if (failed(newFuncOpOrFailure)) {
	LDBG() << "failed to replace function signature with name "
	<< funcOp.getName() << " with new order";
	return failure();
	}

	func::CallOp newCallOp = replaceCallOpWithNewMapping(
	rewriter, callOp, oldArgIdxToNewArgIdx, oldResIdxToNewResIdx);

	return std::make_pair(*newFuncOpOrFailure, newCallOp);
	}