mlir/lib/IR/FunctionSupport.cpp - llvm-project - Git at Google

 //===- FunctionSupport.cpp - Utility types for function-like ops ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//

 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/BitVector.h"

 using namespace mlir;

 /// Helper to call a callback once on each index in the range
 /// [0, `totalIndices`), *except* for the indices given in `indices`.
 /// `indices` is allowed to have duplicates and can be in any order.
 inline void iterateIndicesExcept(unsigned totalIndices,
                                  ArrayRef<unsigned> indices,
                                  function_ref<void(unsigned)> callback) {
   llvm::BitVector skipIndices(totalIndices);
   for (unsigned i : indices)
     skipIndices.set(i);

   for (unsigned i = 0; i < totalIndices; ++i)
     if (!skipIndices.test(i))
       callback(i);
 }

 //===----------------------------------------------------------------------===//
 // Function Arguments and Results.
 //===----------------------------------------------------------------------===//

 static bool isEmptyAttrDict(Attribute attr) {
   return attr.cast<DictionaryAttr>().empty();
 }

 DictionaryAttr mlir::function_like_impl::getArgAttrDict(Operation *op,
                                                         unsigned index) {
   ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
   DictionaryAttr argAttrs =
       attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
   return argAttrs;
 }

 DictionaryAttr mlir::function_like_impl::getResultAttrDict(Operation *op,
                                                            unsigned index) {
   ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
   DictionaryAttr resAttrs =
       attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
   return resAttrs;
 }

 void mlir::function_like_impl::detail::setArgResAttrDict(
     Operation *op, StringRef attrName, unsigned numTotalIndices, unsigned index,
     DictionaryAttr attrs) {
   ArrayAttr allAttrs = op->getAttrOfType<ArrayAttr>(attrName);
   if (!allAttrs) {
     if (attrs.empty())
       return;

     // If this attribute is not empty, we need to create a new attribute array.
     SmallVector<Attribute, 8> newAttrs(numTotalIndices,
                                        DictionaryAttr::get(op->getContext()));
     newAttrs[index] = attrs;
     op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
     return;
   }
   // Check to see if the attribute is different from what we already have.
   if (allAttrs[index] == attrs)
     return;

   // If it is, check to see if the attribute array would now contain only empty
   // dictionaries.
   ArrayRef<Attribute> rawAttrArray = allAttrs.getValue();
   if (attrs.empty() &&
       llvm::all_of(rawAttrArray.take_front(index), isEmptyAttrDict) &&
       llvm::all_of(rawAttrArray.drop_front(index + 1), isEmptyAttrDict)) {
     op->removeAttr(attrName);
     return;
   }

   // Otherwise, create a new attribute array with the updated dictionary.
   SmallVector<Attribute, 8> newAttrs(rawAttrArray.begin(), rawAttrArray.end());
   newAttrs[index] = attrs;
   op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
 }

 /// Set all of the argument or result attribute dictionaries for a function.
 static void setAllArgResAttrDicts(Operation *op, StringRef attrName,
                                   ArrayRef<Attribute> attrs) {
   if (llvm::all_of(attrs, isEmptyAttrDict))
     op->removeAttr(attrName);
   else
     op->setAttr(attrName, ArrayAttr::get(op->getContext(), attrs));
 }

 void mlir::function_like_impl::setAllArgAttrDicts(
     Operation *op, ArrayRef<DictionaryAttr> attrs) {
   setAllArgAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
 }
 void mlir::function_like_impl::setAllArgAttrDicts(Operation *op,
                                                   ArrayRef<Attribute> attrs) {
   auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
     return !attr ? DictionaryAttr::get(op->getContext()) : attr;
   });
   setAllArgResAttrDicts(op, getArgDictAttrName(),
                         llvm::to_vector<8>(wrappedAttrs));
 }

 void mlir::function_like_impl::setAllResultAttrDicts(
     Operation *op, ArrayRef<DictionaryAttr> attrs) {
   setAllResultAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
 }
 void mlir::function_like_impl::setAllResultAttrDicts(
     Operation *op, ArrayRef<Attribute> attrs) {
   auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
     return !attr ? DictionaryAttr::get(op->getContext()) : attr;
   });
   setAllArgResAttrDicts(op, getResultDictAttrName(),
                         llvm::to_vector<8>(wrappedAttrs));
 }

 void mlir::function_like_impl::insertFunctionArguments(
     Operation *op, ArrayRef<unsigned> argIndices, TypeRange argTypes,
     ArrayRef<DictionaryAttr> argAttrs, ArrayRef<Optional<Location>> argLocs,
     unsigned originalNumArgs, Type newType) {
   assert(argIndices.size() == argTypes.size());
   assert(argIndices.size() == argAttrs.size() || argAttrs.empty());
   assert(argIndices.size() == argLocs.size() || argLocs.empty());
   if (argIndices.empty())
     return;

   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
   // - Block arguments of entry block.
   Block &entry = op->getRegion(0).front();

   // Update the argument attributes of the function.
   auto oldArgAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
   if (oldArgAttrs || !argAttrs.empty()) {
     SmallVector<DictionaryAttr, 4> newArgAttrs;
     newArgAttrs.reserve(originalNumArgs + argIndices.size());
     unsigned oldIdx = 0;
     auto migrate = [&](unsigned untilIdx) {
       if (!oldArgAttrs) {
         newArgAttrs.resize(newArgAttrs.size() + untilIdx - oldIdx);
       } else {
         auto oldArgAttrRange = oldArgAttrs.getAsRange<DictionaryAttr>();
         newArgAttrs.append(oldArgAttrRange.begin() + oldIdx,
                            oldArgAttrRange.begin() + untilIdx);
       }
       oldIdx = untilIdx;
     };
     for (unsigned i = 0, e = argIndices.size(); i < e; ++i) {
       migrate(argIndices[i]);
       newArgAttrs.push_back(argAttrs.empty() ? DictionaryAttr{} : argAttrs[i]);
     }
     migrate(originalNumArgs);
     setAllArgAttrDicts(op, newArgAttrs);
   }

   // Update the function type and any entry block arguments.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
   for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
     entry.insertArgument(argIndices[i] + i, argTypes[i],
                          argLocs.empty() ? Optional<Location>{} : argLocs[i]);
 }

 void mlir::function_like_impl::insertFunctionResults(
     Operation *op, ArrayRef<unsigned> resultIndices, TypeRange resultTypes,
     ArrayRef<DictionaryAttr> resultAttrs, unsigned originalNumResults,
     Type newType) {
   assert(resultIndices.size() == resultTypes.size());
   assert(resultIndices.size() == resultAttrs.size() || resultAttrs.empty());
   if (resultIndices.empty())
     return;

   // There are 2 things that need to be updated:
   // - Function type.
   // - Result attrs.

   // Update the result attributes of the function.
   auto oldResultAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
   if (oldResultAttrs || !resultAttrs.empty()) {
     SmallVector<DictionaryAttr, 4> newResultAttrs;
     newResultAttrs.reserve(originalNumResults + resultIndices.size());
     unsigned oldIdx = 0;
     auto migrate = [&](unsigned untilIdx) {
       if (!oldResultAttrs) {
         newResultAttrs.resize(newResultAttrs.size() + untilIdx - oldIdx);
       } else {
         auto oldResultAttrsRange = oldResultAttrs.getAsRange<DictionaryAttr>();
         newResultAttrs.append(oldResultAttrsRange.begin() + oldIdx,
                               oldResultAttrsRange.begin() + untilIdx);
       }
       oldIdx = untilIdx;
     };
     for (unsigned i = 0, e = resultIndices.size(); i < e; ++i) {
       migrate(resultIndices[i]);
       newResultAttrs.push_back(resultAttrs.empty() ? DictionaryAttr{}
                                                    : resultAttrs[i]);
     }
     migrate(originalNumResults);
     setAllResultAttrDicts(op, newResultAttrs);
   }

   // Update the function type.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
 }

 void mlir::function_like_impl::eraseFunctionArguments(
     Operation *op, ArrayRef<unsigned> argIndices, unsigned originalNumArgs,
     Type newType) {
   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
   // - Block arguments of entry block.
   Block &entry = op->getRegion(0).front();

   // Update the argument attributes of the function.
   if (auto argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName())) {
     SmallVector<DictionaryAttr, 4> newArgAttrs;
     newArgAttrs.reserve(argAttrs.size());
     iterateIndicesExcept(originalNumArgs, argIndices, [&](unsigned i) {
       newArgAttrs.emplace_back(argAttrs[i].cast<DictionaryAttr>());
     });
     setAllArgAttrDicts(op, newArgAttrs);
   }

   // Update the function type and any entry block arguments.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
   entry.eraseArguments(argIndices);
 }

 void mlir::function_like_impl::eraseFunctionResults(
     Operation *op, ArrayRef<unsigned> resultIndices,
     unsigned originalNumResults, Type newType) {
   // There are 2 things that need to be updated:
   // - Function type.
   // - Result attrs.

   // Update the result attributes of the function.
   if (auto resAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName())) {
     SmallVector<DictionaryAttr, 4> newResultAttrs;
     newResultAttrs.reserve(resAttrs.size());
     iterateIndicesExcept(originalNumResults, resultIndices, [&](unsigned i) {
       newResultAttrs.emplace_back(resAttrs[i].cast<DictionaryAttr>());
     });
     setAllResultAttrDicts(op, newResultAttrs);
   }

   // Update the function type.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
 }

 //===----------------------------------------------------------------------===//
 // Function type signature.
 //===----------------------------------------------------------------------===//

 FunctionType mlir::function_like_impl::getFunctionType(Operation *op) {
   assert(op->hasTrait<OpTrait::FunctionLike>());
   return op->getAttrOfType<TypeAttr>(getTypeAttrName())
       .getValue()
       .cast<FunctionType>();
 }

 void mlir::function_like_impl::setFunctionType(Operation *op,
                                                FunctionType newType) {
   assert(op->hasTrait<OpTrait::FunctionLike>());
   FunctionType oldType = getFunctionType(op);
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));

   // Functor used to update the argument and result attributes of the function.
   auto updateAttrFn = [&](StringRef attrName, unsigned oldCount,
                           unsigned newCount, auto setAttrFn) {
     if (oldCount == newCount)
       return;
     // The new type has no arguments/results, just drop the attribute.
     if (newCount == 0) {
       op->removeAttr(attrName);
       return;
     }
     ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(attrName);
     if (!attrs)
       return;

     // The new type has less arguments/results, take the first N attributes.
     if (newCount < oldCount)
       return setAttrFn(op, attrs.getValue().take_front(newCount));

     // Otherwise, the new type has more arguments/results. Initialize the new
     // arguments/results with empty attributes.
     SmallVector<Attribute> newAttrs(attrs.begin(), attrs.end());
     newAttrs.resize(newCount);
     setAttrFn(op, newAttrs);
   };

   // Update the argument and result attributes.
   updateAttrFn(function_like_impl::getArgDictAttrName(), oldType.getNumInputs(),
                newType.getNumInputs(), [&](Operation *op, auto &&attrs) {
                  setAllArgAttrDicts(op, attrs);
                });
   updateAttrFn(
       function_like_impl::getResultDictAttrName(), oldType.getNumResults(),
       newType.getNumResults(),
       [&](Operation *op, auto &&attrs) { setAllResultAttrDicts(op, attrs); });
 }

 //===----------------------------------------------------------------------===//
 // Function body.
 //===----------------------------------------------------------------------===//

 Region &mlir::function_like_impl::getFunctionBody(Operation *op) {
   assert(op->hasTrait<OpTrait::FunctionLike>());
   return op->getRegion(0);
 }
	//===- FunctionSupport.cpp - Utility types for function-like ops ----------===//
	//
	// 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
	//
	//===----------------------------------------------------------------------===//

	#include "mlir/IR/FunctionSupport.h"
	#include "mlir/Support/LLVM.h"
	#include "llvm/ADT/BitVector.h"

	using namespace mlir;

	/// Helper to call a callback once on each index in the range
	/// [0, `totalIndices`), except for the indices given in `indices`.
	/// `indices` is allowed to have duplicates and can be in any order.
	inline void iterateIndicesExcept(unsigned totalIndices,
	ArrayRef<unsigned> indices,
	function_ref<void(unsigned)> callback) {
	llvm::BitVector skipIndices(totalIndices);
	for (unsigned i : indices)
	skipIndices.set(i);

	for (unsigned i = 0; i < totalIndices; ++i)
	if (!skipIndices.test(i))
	callback(i);
	}

	//===----------------------------------------------------------------------===//
	// Function Arguments and Results.
	//===----------------------------------------------------------------------===//

	static bool isEmptyAttrDict(Attribute attr) {
	return attr.cast<DictionaryAttr>().empty();
	}

	DictionaryAttr mlir::function_like_impl::getArgAttrDict(Operation *op,
	unsigned index) {
	ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
	DictionaryAttr argAttrs =
	attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
	return argAttrs;
	}

	DictionaryAttr mlir::function_like_impl::getResultAttrDict(Operation *op,
	unsigned index) {
	ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
	DictionaryAttr resAttrs =
	attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
	return resAttrs;
	}

	void mlir::function_like_impl::detail::setArgResAttrDict(
	Operation *op, StringRef attrName, unsigned numTotalIndices, unsigned index,
	DictionaryAttr attrs) {
	ArrayAttr allAttrs = op->getAttrOfType<ArrayAttr>(attrName);
	if (!allAttrs) {
	if (attrs.empty())
	return;

	// If this attribute is not empty, we need to create a new attribute array.
	SmallVector<Attribute, 8> newAttrs(numTotalIndices,
	DictionaryAttr::get(op->getContext()));
	newAttrs[index] = attrs;
	op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
	return;
	}
	// Check to see if the attribute is different from what we already have.
	if (allAttrs[index] == attrs)
	return;

	// If it is, check to see if the attribute array would now contain only empty
	// dictionaries.
	ArrayRef<Attribute> rawAttrArray = allAttrs.getValue();
	if (attrs.empty() &&
	llvm::all_of(rawAttrArray.take_front(index), isEmptyAttrDict) &&
	llvm::all_of(rawAttrArray.drop_front(index + 1), isEmptyAttrDict)) {
	op->removeAttr(attrName);
	return;
	}

	// Otherwise, create a new attribute array with the updated dictionary.
	SmallVector<Attribute, 8> newAttrs(rawAttrArray.begin(), rawAttrArray.end());
	newAttrs[index] = attrs;
	op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
	}

	/// Set all of the argument or result attribute dictionaries for a function.
	static void setAllArgResAttrDicts(Operation *op, StringRef attrName,
	ArrayRef<Attribute> attrs) {
	if (llvm::all_of(attrs, isEmptyAttrDict))
	op->removeAttr(attrName);
	else
	op->setAttr(attrName, ArrayAttr::get(op->getContext(), attrs));
	}

	void mlir::function_like_impl::setAllArgAttrDicts(
	Operation *op, ArrayRef<DictionaryAttr> attrs) {
	setAllArgAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
	}
	void mlir::function_like_impl::setAllArgAttrDicts(Operation *op,
	ArrayRef<Attribute> attrs) {
	auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
	return !attr ? DictionaryAttr::get(op->getContext()) : attr;
	});
	setAllArgResAttrDicts(op, getArgDictAttrName(),
	llvm::to_vector<8>(wrappedAttrs));
	}

	void mlir::function_like_impl::setAllResultAttrDicts(
	Operation *op, ArrayRef<DictionaryAttr> attrs) {
	setAllResultAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
	}
	void mlir::function_like_impl::setAllResultAttrDicts(
	Operation *op, ArrayRef<Attribute> attrs) {
	auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
	return !attr ? DictionaryAttr::get(op->getContext()) : attr;
	});
	setAllArgResAttrDicts(op, getResultDictAttrName(),
	llvm::to_vector<8>(wrappedAttrs));
	}

	void mlir::function_like_impl::insertFunctionArguments(
	Operation *op, ArrayRef<unsigned> argIndices, TypeRange argTypes,
	ArrayRef<DictionaryAttr> argAttrs, ArrayRef<Optional<Location>> argLocs,
	unsigned originalNumArgs, Type newType) {
	assert(argIndices.size() == argTypes.size());
	assert(argIndices.size() == argAttrs.size() \|\| argAttrs.empty());
	assert(argIndices.size() == argLocs.size() \|\| argLocs.empty());
	if (argIndices.empty())
	return;

	// There are 3 things that need to be updated:
	// - Function type.
	// - Arg attrs.
	// - Block arguments of entry block.
	Block &entry = op->getRegion(0).front();

	// Update the argument attributes of the function.
	auto oldArgAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
	if (oldArgAttrs \|\| !argAttrs.empty()) {
	SmallVector<DictionaryAttr, 4> newArgAttrs;
	newArgAttrs.reserve(originalNumArgs + argIndices.size());
	unsigned oldIdx = 0;
	auto migrate = [&](unsigned untilIdx) {
	if (!oldArgAttrs) {
	newArgAttrs.resize(newArgAttrs.size() + untilIdx - oldIdx);
	} else {
	auto oldArgAttrRange = oldArgAttrs.getAsRange<DictionaryAttr>();
	newArgAttrs.append(oldArgAttrRange.begin() + oldIdx,
	oldArgAttrRange.begin() + untilIdx);
	}
	oldIdx = untilIdx;
	};
	for (unsigned i = 0, e = argIndices.size(); i < e; ++i) {
	migrate(argIndices[i]);
	newArgAttrs.push_back(argAttrs.empty() ? DictionaryAttr{} : argAttrs[i]);
	}
	migrate(originalNumArgs);
	setAllArgAttrDicts(op, newArgAttrs);
	}

	// Update the function type and any entry block arguments.
	op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
	for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
	entry.insertArgument(argIndices[i] + i, argTypes[i],
	argLocs.empty() ? Optional<Location>{} : argLocs[i]);
	}

	void mlir::function_like_impl::insertFunctionResults(
	Operation *op, ArrayRef<unsigned> resultIndices, TypeRange resultTypes,
	ArrayRef<DictionaryAttr> resultAttrs, unsigned originalNumResults,
	Type newType) {
	assert(resultIndices.size() == resultTypes.size());
	assert(resultIndices.size() == resultAttrs.size() \|\| resultAttrs.empty());
	if (resultIndices.empty())
	return;

	// There are 2 things that need to be updated:
	// - Function type.
	// - Result attrs.

	// Update the result attributes of the function.
	auto oldResultAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
	if (oldResultAttrs \|\| !resultAttrs.empty()) {
	SmallVector<DictionaryAttr, 4> newResultAttrs;
	newResultAttrs.reserve(originalNumResults + resultIndices.size());
	unsigned oldIdx = 0;
	auto migrate = [&](unsigned untilIdx) {
	if (!oldResultAttrs) {
	newResultAttrs.resize(newResultAttrs.size() + untilIdx - oldIdx);
	} else {
	auto oldResultAttrsRange = oldResultAttrs.getAsRange<DictionaryAttr>();
	newResultAttrs.append(oldResultAttrsRange.begin() + oldIdx,
	oldResultAttrsRange.begin() + untilIdx);
	}
	oldIdx = untilIdx;
	};
	for (unsigned i = 0, e = resultIndices.size(); i < e; ++i) {
	migrate(resultIndices[i]);
	newResultAttrs.push_back(resultAttrs.empty() ? DictionaryAttr{}
	: resultAttrs[i]);
	}
	migrate(originalNumResults);
	setAllResultAttrDicts(op, newResultAttrs);
	}

	// Update the function type.
	op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
	}

	void mlir::function_like_impl::eraseFunctionArguments(
	Operation *op, ArrayRef<unsigned> argIndices, unsigned originalNumArgs,
	Type newType) {
	// There are 3 things that need to be updated:
	// - Function type.
	// - Arg attrs.
	// - Block arguments of entry block.
	Block &entry = op->getRegion(0).front();

	// Update the argument attributes of the function.
	if (auto argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName())) {
	SmallVector<DictionaryAttr, 4> newArgAttrs;
	newArgAttrs.reserve(argAttrs.size());
	iterateIndicesExcept(originalNumArgs, argIndices, [&](unsigned i) {
	newArgAttrs.emplace_back(argAttrs[i].cast<DictionaryAttr>());
	});
	setAllArgAttrDicts(op, newArgAttrs);
	}

	// Update the function type and any entry block arguments.
	op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
	entry.eraseArguments(argIndices);
	}

	void mlir::function_like_impl::eraseFunctionResults(
	Operation *op, ArrayRef<unsigned> resultIndices,
	unsigned originalNumResults, Type newType) {
	// There are 2 things that need to be updated:
	// - Function type.
	// - Result attrs.

	// Update the result attributes of the function.
	if (auto resAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName())) {
	SmallVector<DictionaryAttr, 4> newResultAttrs;
	newResultAttrs.reserve(resAttrs.size());
	iterateIndicesExcept(originalNumResults, resultIndices, [&](unsigned i) {
	newResultAttrs.emplace_back(resAttrs[i].cast<DictionaryAttr>());
	});
	setAllResultAttrDicts(op, newResultAttrs);
	}

	// Update the function type.
	op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
	}

	//===----------------------------------------------------------------------===//
	// Function type signature.
	//===----------------------------------------------------------------------===//

	FunctionType mlir::function_like_impl::getFunctionType(Operation *op) {
	assert(op->hasTrait<OpTrait::FunctionLike>());
	return op->getAttrOfType<TypeAttr>(getTypeAttrName())
	.getValue()
	.cast<FunctionType>();
	}

	void mlir::function_like_impl::setFunctionType(Operation *op,
	FunctionType newType) {
	assert(op->hasTrait<OpTrait::FunctionLike>());
	FunctionType oldType = getFunctionType(op);
	op->setAttr(getTypeAttrName(), TypeAttr::get(newType));

	// Functor used to update the argument and result attributes of the function.
	auto updateAttrFn = [&](StringRef attrName, unsigned oldCount,
	unsigned newCount, auto setAttrFn) {
	if (oldCount == newCount)
	return;
	// The new type has no arguments/results, just drop the attribute.
	if (newCount == 0) {
	op->removeAttr(attrName);
	return;
	}
	ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(attrName);
	if (!attrs)
	return;

	// The new type has less arguments/results, take the first N attributes.
	if (newCount < oldCount)
	return setAttrFn(op, attrs.getValue().take_front(newCount));

	// Otherwise, the new type has more arguments/results. Initialize the new
	// arguments/results with empty attributes.
	SmallVector<Attribute> newAttrs(attrs.begin(), attrs.end());
	newAttrs.resize(newCount);
	setAttrFn(op, newAttrs);
	};

	// Update the argument and result attributes.
	updateAttrFn(function_like_impl::getArgDictAttrName(), oldType.getNumInputs(),
	newType.getNumInputs(), [&](Operation *op, auto &&attrs) {
	setAllArgAttrDicts(op, attrs);
	});
	updateAttrFn(
	function_like_impl::getResultDictAttrName(), oldType.getNumResults(),
	newType.getNumResults(),
	[&](Operation *op, auto &&attrs) { setAllResultAttrDicts(op, attrs); });
	}

	//===----------------------------------------------------------------------===//
	// Function body.
	//===----------------------------------------------------------------------===//

	Region &mlir::function_like_impl::getFunctionBody(Operation *op) {
	assert(op->hasTrait<OpTrait::FunctionLike>());
	return op->getRegion(0);
	}