llvm-project/mlir/lib/Transforms/NormalizeMemRefs.cpp

524 lines
24 KiB
C++

//===- NormalizeMemRefs.cpp -----------------------------------------------===//
//
// 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 an interprocedural pass to normalize memrefs to have
// identity layout maps.
//
//===----------------------------------------------------------------------===//
#include "PassDetail.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/Utils.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/Support/Debug.h"
#define DEBUG_TYPE "normalize-memrefs"
using namespace mlir;
namespace {
/// All memrefs passed across functions with non-trivial layout maps are
/// converted to ones with trivial identity layout ones.
/// If all the memref types/uses in a function are normalizable, we treat
/// such functions as normalizable. Also, if a normalizable function is known
/// to call a non-normalizable function, we treat that function as
/// non-normalizable as well. We assume external functions to be normalizable.
struct NormalizeMemRefs : public NormalizeMemRefsBase<NormalizeMemRefs> {
void runOnOperation() override;
void normalizeFuncOpMemRefs(FuncOp funcOp, ModuleOp moduleOp);
bool areMemRefsNormalizable(FuncOp funcOp);
void updateFunctionSignature(FuncOp funcOp, ModuleOp moduleOp);
void setCalleesAndCallersNonNormalizable(FuncOp funcOp, ModuleOp moduleOp,
DenseSet<FuncOp> &normalizableFuncs);
Operation *createOpResultsNormalized(FuncOp funcOp, Operation *oldOp);
};
} // end anonymous namespace
std::unique_ptr<OperationPass<ModuleOp>> mlir::createNormalizeMemRefsPass() {
return std::make_unique<NormalizeMemRefs>();
}
void NormalizeMemRefs::runOnOperation() {
LLVM_DEBUG(llvm::dbgs() << "Normalizing Memrefs...\n");
ModuleOp moduleOp = getOperation();
// We maintain all normalizable FuncOps in a DenseSet. It is initialized
// with all the functions within a module and then functions which are not
// normalizable are removed from this set.
// TODO: Change this to work on FuncLikeOp once there is an operation
// interface for it.
DenseSet<FuncOp> normalizableFuncs;
// Initialize `normalizableFuncs` with all the functions within a module.
moduleOp.walk([&](FuncOp funcOp) { normalizableFuncs.insert(funcOp); });
// Traverse through all the functions applying a filter which determines
// whether that function is normalizable or not. All callers/callees of
// a non-normalizable function will also become non-normalizable even if
// they aren't passing any or specific non-normalizable memrefs. So,
// functions which calls or get called by a non-normalizable becomes non-
// normalizable functions themselves.
moduleOp.walk([&](FuncOp funcOp) {
if (normalizableFuncs.contains(funcOp)) {
if (!areMemRefsNormalizable(funcOp)) {
LLVM_DEBUG(llvm::dbgs()
<< "@" << funcOp.getName()
<< " contains ops that cannot normalize MemRefs\n");
// Since this function is not normalizable, we set all the caller
// functions and the callees of this function as not normalizable.
// TODO: Drop this conservative assumption in the future.
setCalleesAndCallersNonNormalizable(funcOp, moduleOp,
normalizableFuncs);
}
}
});
LLVM_DEBUG(llvm::dbgs() << "Normalizing " << normalizableFuncs.size()
<< " functions\n");
// Those functions which can be normalized are subjected to normalization.
for (FuncOp &funcOp : normalizableFuncs)
normalizeFuncOpMemRefs(funcOp, moduleOp);
}
/// Check whether all the uses of oldMemRef are either dereferencing uses or the
/// op is of type : DeallocOp, CallOp or ReturnOp. Only if these constraints
/// are satisfied will the value become a candidate for replacement.
/// TODO: Extend this for DimOps.
static bool isMemRefNormalizable(Value::user_range opUsers) {
if (llvm::any_of(opUsers, [](Operation *op) {
if (op->hasTrait<OpTrait::MemRefsNormalizable>())
return false;
return true;
}))
return false;
return true;
}
/// Set all the calling functions and the callees of the function as not
/// normalizable.
void NormalizeMemRefs::setCalleesAndCallersNonNormalizable(
FuncOp funcOp, ModuleOp moduleOp, DenseSet<FuncOp> &normalizableFuncs) {
if (!normalizableFuncs.contains(funcOp))
return;
LLVM_DEBUG(
llvm::dbgs() << "@" << funcOp.getName()
<< " calls or is called by non-normalizable function\n");
normalizableFuncs.erase(funcOp);
// Caller of the function.
Optional<SymbolTable::UseRange> symbolUses = funcOp.getSymbolUses(moduleOp);
for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
// TODO: Extend this for ops that are FunctionLike. This would require
// creating an OpInterface for FunctionLike ops.
FuncOp parentFuncOp = symbolUse.getUser()->getParentOfType<FuncOp>();
for (FuncOp &funcOp : normalizableFuncs) {
if (parentFuncOp == funcOp) {
setCalleesAndCallersNonNormalizable(funcOp, moduleOp,
normalizableFuncs);
break;
}
}
}
// Functions called by this function.
funcOp.walk([&](CallOp callOp) {
StringAttr callee = callOp.getCalleeAttr();
for (FuncOp &funcOp : normalizableFuncs) {
// We compare FuncOp and callee's name.
if (callee == funcOp.getNameAttr()) {
setCalleesAndCallersNonNormalizable(funcOp, moduleOp,
normalizableFuncs);
break;
}
}
});
}
/// Check whether all the uses of AllocOps, CallOps and function arguments of a
/// function are either of dereferencing type or are uses in: DeallocOp, CallOp
/// or ReturnOp. Only if these constraints are satisfied will the function
/// become a candidate for normalization. We follow a conservative approach here
/// wherein even if the non-normalizable memref is not a part of the function's
/// argument or return type, we still label the entire function as
/// non-normalizable. We assume external functions to be normalizable.
bool NormalizeMemRefs::areMemRefsNormalizable(FuncOp funcOp) {
// We assume external functions to be normalizable.
if (funcOp.isExternal())
return true;
if (funcOp
.walk([&](memref::AllocOp allocOp) -> WalkResult {
Value oldMemRef = allocOp.getResult();
if (!isMemRefNormalizable(oldMemRef.getUsers()))
return WalkResult::interrupt();
return WalkResult::advance();
})
.wasInterrupted())
return false;
if (funcOp
.walk([&](CallOp callOp) -> WalkResult {
for (unsigned resIndex :
llvm::seq<unsigned>(0, callOp.getNumResults())) {
Value oldMemRef = callOp.getResult(resIndex);
if (oldMemRef.getType().isa<MemRefType>())
if (!isMemRefNormalizable(oldMemRef.getUsers()))
return WalkResult::interrupt();
}
return WalkResult::advance();
})
.wasInterrupted())
return false;
for (unsigned argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
BlockArgument oldMemRef = funcOp.getArgument(argIndex);
if (oldMemRef.getType().isa<MemRefType>())
if (!isMemRefNormalizable(oldMemRef.getUsers()))
return false;
}
return true;
}
/// Fetch the updated argument list and result of the function and update the
/// function signature. This updates the function's return type at the caller
/// site and in case the return type is a normalized memref then it updates
/// the calling function's signature.
/// TODO: An update to the calling function signature is required only if the
/// returned value is in turn used in ReturnOp of the calling function.
void NormalizeMemRefs::updateFunctionSignature(FuncOp funcOp,
ModuleOp moduleOp) {
FunctionType functionType = funcOp.getType();
SmallVector<Type, 4> resultTypes;
FunctionType newFuncType;
resultTypes = llvm::to_vector<4>(functionType.getResults());
// External function's signature was already updated in
// 'normalizeFuncOpMemRefs()'.
if (!funcOp.isExternal()) {
SmallVector<Type, 8> argTypes;
for (const auto &argEn : llvm::enumerate(funcOp.getArguments()))
argTypes.push_back(argEn.value().getType());
// Traverse ReturnOps to check if an update to the return type in the
// function signature is required.
funcOp.walk([&](ReturnOp returnOp) {
for (const auto &operandEn : llvm::enumerate(returnOp.getOperands())) {
Type opType = operandEn.value().getType();
MemRefType memrefType = opType.dyn_cast<MemRefType>();
// If type is not memref or if the memref type is same as that in
// function's return signature then no update is required.
if (!memrefType || memrefType == resultTypes[operandEn.index()])
continue;
// Update function's return type signature.
// Return type gets normalized either as a result of function argument
// normalization, AllocOp normalization or an update made at CallOp.
// There can be many call flows inside a function and an update to a
// specific ReturnOp has not yet been made. So we check that the result
// memref type is normalized.
// TODO: When selective normalization is implemented, handle multiple
// results case where some are normalized, some aren't.
if (memrefType.getAffineMaps().empty())
resultTypes[operandEn.index()] = memrefType;
}
});
// We create a new function type and modify the function signature with this
// new type.
newFuncType = FunctionType::get(&getContext(), /*inputs=*/argTypes,
/*results=*/resultTypes);
}
// Since we update the function signature, it might affect the result types at
// the caller site. Since this result might even be used by the caller
// function in ReturnOps, the caller function's signature will also change.
// Hence we record the caller function in 'funcOpsToUpdate' to update their
// signature as well.
llvm::SmallDenseSet<FuncOp, 8> funcOpsToUpdate;
// We iterate over all symbolic uses of the function and update the return
// type at the caller site.
Optional<SymbolTable::UseRange> symbolUses = funcOp.getSymbolUses(moduleOp);
for (SymbolTable::SymbolUse symbolUse : *symbolUses) {
Operation *userOp = symbolUse.getUser();
OpBuilder builder(userOp);
// When `userOp` can not be casted to `CallOp`, it is skipped. This assumes
// that the non-CallOp has no memrefs to be replaced.
// TODO: Handle cases where a non-CallOp symbol use of a function deals with
// memrefs.
auto callOp = dyn_cast<CallOp>(userOp);
if (!callOp)
continue;
Operation *newCallOp =
builder.create<CallOp>(userOp->getLoc(), callOp.getCalleeAttr(),
resultTypes, userOp->getOperands());
bool replacingMemRefUsesFailed = false;
bool returnTypeChanged = false;
for (unsigned resIndex : llvm::seq<unsigned>(0, userOp->getNumResults())) {
OpResult oldResult = userOp->getResult(resIndex);
OpResult newResult = newCallOp->getResult(resIndex);
// This condition ensures that if the result is not of type memref or if
// the resulting memref was already having a trivial map layout then we
// need not perform any use replacement here.
if (oldResult.getType() == newResult.getType())
continue;
AffineMap layoutMap =
oldResult.getType().dyn_cast<MemRefType>().getAffineMaps().front();
if (failed(replaceAllMemRefUsesWith(oldResult, /*newMemRef=*/newResult,
/*extraIndices=*/{},
/*indexRemap=*/layoutMap,
/*extraOperands=*/{},
/*symbolOperands=*/{},
/*domInstFilter=*/nullptr,
/*postDomInstFilter=*/nullptr,
/*allowDereferencingOps=*/true,
/*replaceInDeallocOp=*/true))) {
// If it failed (due to escapes for example), bail out.
// It should never hit this part of the code because it is called by
// only those functions which are normalizable.
newCallOp->erase();
replacingMemRefUsesFailed = true;
break;
}
returnTypeChanged = true;
}
if (replacingMemRefUsesFailed)
continue;
// Replace all uses for other non-memref result types.
userOp->replaceAllUsesWith(newCallOp);
userOp->erase();
if (returnTypeChanged) {
// Since the return type changed it might lead to a change in function's
// signature.
// TODO: If funcOp doesn't return any memref type then no need to update
// signature.
// TODO: Further optimization - Check if the memref is indeed part of
// ReturnOp at the parentFuncOp and only then updation of signature is
// required.
// TODO: Extend this for ops that are FunctionLike. This would require
// creating an OpInterface for FunctionLike ops.
FuncOp parentFuncOp = newCallOp->getParentOfType<FuncOp>();
funcOpsToUpdate.insert(parentFuncOp);
}
}
// Because external function's signature is already updated in
// 'normalizeFuncOpMemRefs()', we don't need to update it here again.
if (!funcOp.isExternal())
funcOp.setType(newFuncType);
// Updating the signature type of those functions which call the current
// function. Only if the return type of the current function has a normalized
// memref will the caller function become a candidate for signature update.
for (FuncOp parentFuncOp : funcOpsToUpdate)
updateFunctionSignature(parentFuncOp, moduleOp);
}
/// Normalizes the memrefs within a function which includes those arising as a
/// result of AllocOps, CallOps and function's argument. The ModuleOp argument
/// is used to help update function's signature after normalization.
void NormalizeMemRefs::normalizeFuncOpMemRefs(FuncOp funcOp,
ModuleOp moduleOp) {
// Turn memrefs' non-identity layouts maps into ones with identity. Collect
// alloc ops first and then process since normalizeMemRef replaces/erases ops
// during memref rewriting.
SmallVector<memref::AllocOp, 4> allocOps;
funcOp.walk([&](memref::AllocOp op) { allocOps.push_back(op); });
for (memref::AllocOp allocOp : allocOps)
(void)normalizeMemRef(&allocOp);
// We use this OpBuilder to create new memref layout later.
OpBuilder b(funcOp);
FunctionType functionType = funcOp.getType();
SmallVector<Type, 8> inputTypes;
// Walk over each argument of a function to perform memref normalization (if
for (unsigned argIndex :
llvm::seq<unsigned>(0, functionType.getNumInputs())) {
Type argType = functionType.getInput(argIndex);
MemRefType memrefType = argType.dyn_cast<MemRefType>();
// Check whether argument is of MemRef type. Any other argument type can
// simply be part of the final function signature.
if (!memrefType) {
inputTypes.push_back(argType);
continue;
}
// Fetch a new memref type after normalizing the old memref to have an
// identity map layout.
MemRefType newMemRefType = normalizeMemRefType(memrefType, b,
/*numSymbolicOperands=*/0);
if (newMemRefType == memrefType || funcOp.isExternal()) {
// Either memrefType already had an identity map or the map couldn't be
// transformed to an identity map.
inputTypes.push_back(newMemRefType);
continue;
}
// Insert a new temporary argument with the new memref type.
BlockArgument newMemRef =
funcOp.front().insertArgument(argIndex, newMemRefType);
BlockArgument oldMemRef = funcOp.getArgument(argIndex + 1);
AffineMap layoutMap = memrefType.getAffineMaps().front();
// Replace all uses of the old memref.
if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newMemRef,
/*extraIndices=*/{},
/*indexRemap=*/layoutMap,
/*extraOperands=*/{},
/*symbolOperands=*/{},
/*domInstFilter=*/nullptr,
/*postDomInstFilter=*/nullptr,
/*allowNonDereferencingOps=*/true,
/*replaceInDeallocOp=*/true))) {
// If it failed (due to escapes for example), bail out. Removing the
// temporary argument inserted previously.
funcOp.front().eraseArgument(argIndex);
continue;
}
// All uses for the argument with old memref type were replaced
// successfully. So we remove the old argument now.
funcOp.front().eraseArgument(argIndex + 1);
}
// Walk over normalizable operations to normalize memrefs of the operation
// results. When `op` has memrefs with affine map in the operation results,
// new operation containin normalized memrefs is created. Then, the memrefs
// are replaced. `CallOp` is skipped here because it is handled in
// `updateFunctionSignature()`.
funcOp.walk([&](Operation *op) {
if (op->hasTrait<OpTrait::MemRefsNormalizable>() &&
op->getNumResults() > 0 && !isa<CallOp>(op) && !funcOp.isExternal()) {
// Create newOp containing normalized memref in the operation result.
Operation *newOp = createOpResultsNormalized(funcOp, op);
// When all of the operation results have no memrefs or memrefs without
// affine map, `newOp` is the same with `op` and following process is
// skipped.
if (op != newOp) {
bool replacingMemRefUsesFailed = false;
for (unsigned resIndex : llvm::seq<unsigned>(0, op->getNumResults())) {
// Replace all uses of the old memrefs.
Value oldMemRef = op->getResult(resIndex);
Value newMemRef = newOp->getResult(resIndex);
MemRefType oldMemRefType = oldMemRef.getType().dyn_cast<MemRefType>();
// Check whether the operation result is MemRef type.
if (!oldMemRefType)
continue;
MemRefType newMemRefType = newMemRef.getType().cast<MemRefType>();
if (oldMemRefType == newMemRefType)
continue;
// TODO: Assume single layout map. Multiple maps not supported.
AffineMap layoutMap = oldMemRefType.getAffineMaps().front();
if (failed(replaceAllMemRefUsesWith(oldMemRef,
/*newMemRef=*/newMemRef,
/*extraIndices=*/{},
/*indexRemap=*/layoutMap,
/*extraOperands=*/{},
/*symbolOperands=*/{},
/*domInstFilter=*/nullptr,
/*postDomInstFilter=*/nullptr,
/*allowDereferencingOps=*/true,
/*replaceInDeallocOp=*/true))) {
newOp->erase();
replacingMemRefUsesFailed = true;
continue;
}
}
if (!replacingMemRefUsesFailed) {
// Replace other ops with new op and delete the old op when the
// replacement succeeded.
op->replaceAllUsesWith(newOp);
op->erase();
}
}
}
});
// In a normal function, memrefs in the return type signature gets normalized
// as a result of normalization of functions arguments, AllocOps or CallOps'
// result types. Since an external function doesn't have a body, memrefs in
// the return type signature can only get normalized by iterating over the
// individual return types.
if (funcOp.isExternal()) {
SmallVector<Type, 4> resultTypes;
for (unsigned resIndex :
llvm::seq<unsigned>(0, functionType.getNumResults())) {
Type resType = functionType.getResult(resIndex);
MemRefType memrefType = resType.dyn_cast<MemRefType>();
// Check whether result is of MemRef type. Any other argument type can
// simply be part of the final function signature.
if (!memrefType) {
resultTypes.push_back(resType);
continue;
}
// Computing a new memref type after normalizing the old memref to have an
// identity map layout.
MemRefType newMemRefType = normalizeMemRefType(memrefType, b,
/*numSymbolicOperands=*/0);
resultTypes.push_back(newMemRefType);
continue;
}
FunctionType newFuncType =
FunctionType::get(&getContext(), /*inputs=*/inputTypes,
/*results=*/resultTypes);
// Setting the new function signature for this external function.
funcOp.setType(newFuncType);
}
updateFunctionSignature(funcOp, moduleOp);
}
/// Create an operation containing normalized memrefs in the operation results.
/// When the results of `oldOp` have memrefs with affine map, the memrefs are
/// normalized, and new operation containing them in the operation results is
/// returned. If all of the results of `oldOp` have no memrefs or memrefs
/// without affine map, `oldOp` is returned without modification.
Operation *NormalizeMemRefs::createOpResultsNormalized(FuncOp funcOp,
Operation *oldOp) {
// Prepare OperationState to create newOp containing normalized memref in
// the operation results.
OperationState result(oldOp->getLoc(), oldOp->getName());
result.addOperands(oldOp->getOperands());
result.addAttributes(oldOp->getAttrs());
// Add normalized MemRefType to the OperationState.
SmallVector<Type, 4> resultTypes;
OpBuilder b(funcOp);
bool resultTypeNormalized = false;
for (unsigned resIndex : llvm::seq<unsigned>(0, oldOp->getNumResults())) {
auto resultType = oldOp->getResult(resIndex).getType();
MemRefType memrefType = resultType.dyn_cast<MemRefType>();
// Check whether the operation result is MemRef type.
if (!memrefType) {
resultTypes.push_back(resultType);
continue;
}
// Fetch a new memref type after normalizing the old memref.
MemRefType newMemRefType = normalizeMemRefType(memrefType, b,
/*numSymbolicOperands=*/0);
if (newMemRefType == memrefType) {
// Either memrefType already had an identity map or the map couldn't
// be transformed to an identity map.
resultTypes.push_back(memrefType);
continue;
}
resultTypes.push_back(newMemRefType);
resultTypeNormalized = true;
}
result.addTypes(resultTypes);
// When all of the results of `oldOp` have no memrefs or memrefs without
// affine map, `oldOp` is returned without modification.
if (resultTypeNormalized) {
OpBuilder bb(oldOp);
for (auto &oldRegion : oldOp->getRegions()) {
Region *newRegion = result.addRegion();
newRegion->takeBody(oldRegion);
}
return bb.createOperation(result);
} else
return oldOp;
}