Skip to content

Commit

Permalink
[MLIR][LLVM] Improve the noalias propagation during inlining (#104750)
Browse files Browse the repository at this point in the history 
This commit changes the LLVM dialect's inliner interface to properly
propagate noalias information to memory accesses that have different
underlying object. By always introducing an SSACopy intrinsic, it's
possible to understand that specific memory operations are using
unrelated pointers. Previously, the backwards slice walk did continue
beyond the boundary of the original function and failed to reason about
the "underlying objects".
  • Loading branch information
@Dinistro
Dinistro authored Aug 19, 2024
1 parent 63267ca commit 065d2d9
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 35 deletions.
72 changes: 37 additions & 35 deletions mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -271,37 +271,41 @@ getUnderlyingObjectSet(Value pointerValue) {
static void createNewAliasScopesFromNoAliasParameter(
Operation *call, iterator_range<Region::iterator> inlinedBlocks) {

// First collect all noalias parameters. These have been specially marked by
// the `handleArgument` implementation by using the `ssa.copy` intrinsic and
// attaching a `noalias` attribute to it.
// These are only meant to be temporary and should therefore be deleted after
// we're done using them here.
// First, collect all ssa copy operations, which correspond to function
// parameters, and additionally store the noalias parameters. All parameters
// have been marked by the `handleArgument` implementation by using the
// `ssa.copy` intrinsic. Additionally, noalias parameters have an attached
// `noalias` attribute to the intrinsics. These intrinsics are only meant to
// be temporary and should therefore be deleted after we're done using them
// here.
SetVector<LLVM::SSACopyOp> ssaCopies;
SetVector<LLVM::SSACopyOp> noAliasParams;
for (Value argument : cast<LLVM::CallOp>(call).getArgOperands()) {
for (Operation *user : argument.getUsers()) {
auto ssaCopy = llvm::dyn_cast<LLVM::SSACopyOp>(user);
if (!ssaCopy)
continue;
ssaCopies.insert(ssaCopy);

if (!ssaCopy->hasAttr(LLVM::LLVMDialect::getNoAliasAttrName()))
continue;

noAliasParams.insert(ssaCopy);
}
}

// If there were none, we have nothing to do here.
if (noAliasParams.empty())
return;

// Scope exit block to make it impossible to forget to get rid of the
// intrinsics.
auto exit = llvm::make_scope_exit([&] {
for (LLVM::SSACopyOp ssaCopyOp : noAliasParams) {
for (LLVM::SSACopyOp ssaCopyOp : ssaCopies) {
ssaCopyOp.replaceAllUsesWith(ssaCopyOp.getOperand());
ssaCopyOp->erase();
}
});

// If there were no noalias parameters, we have nothing to do here.
if (noAliasParams.empty())
return;

// Create a new domain for this specific inlining and a new scope for every
// noalias parameter.
auto functionDomain = LLVM::AliasScopeDomainAttr::get(
Expand Down Expand Up @@ -335,7 +339,7 @@ static void createNewAliasScopesFromNoAliasParameter(
bool aliasesOtherKnownObject = false;
// Go through the based on pointers and check that they are either:
// * Constants that can be ignored (undef, poison, null pointer).
// * Based on a noalias parameter.
// * Based on a pointer parameter.
// * Other pointers that we know can't alias with our noalias parameter.
//
// Any other value might be a pointer based on any noalias parameter that
Expand All @@ -346,11 +350,13 @@ static void createNewAliasScopesFromNoAliasParameter(
if (matchPattern(object, m_Constant()))
return false;

if (noAliasParams.contains(object.getDefiningOp<LLVM::SSACopyOp>()))
if (auto ssaCopy = object.getDefiningOp<LLVM::SSACopyOp>()) {
// If that value is based on a noalias parameter, it is guaranteed
// to not alias with any other object.
aliasesOtherKnownObject |= !noAliasParams.contains(ssaCopy);
return false;
}

// TODO: This should include other arguments from the inlined
// callable.
if (isa_and_nonnull<LLVM::AllocaOp, LLVM::AddressOfOp>(
object.getDefiningOp())) {
aliasesOtherKnownObject = true;
Expand Down Expand Up @@ -773,29 +779,25 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
return handleByValArgument(builder, callable, argument, elementType,
requestedAlignment);
}
if (argumentAttrs.contains(LLVM::LLVMDialect::getNoAliasAttrName())) {
if (argument.use_empty())
return argument;

// This code is essentially a workaround for deficiencies in the
// inliner interface: We need to transform operations *after* inlined
// based on the argument attributes of the parameters *before* inlining.
// This method runs prior to actual inlining and thus cannot transform the
// post-inlining code, while `processInlinedCallBlocks` does not have
// access to pre-inlining function arguments. Additionally, it is required
// to distinguish which parameter an SSA value originally came from.
// As a workaround until this is changed: Create an ssa.copy intrinsic
// with the noalias attribute that can easily be found, and is extremely
// unlikely to exist in the code prior to inlining, using this to
// communicate between this method and `processInlinedCallBlocks`.
// TODO: Fix this by refactoring the inliner interface.
auto copyOp = builder.create<LLVM::SSACopyOp>(call->getLoc(), argument);

// This code is essentially a workaround for deficiencies in the inliner
// interface: We need to transform operations *after* inlined based on the
// argument attributes of the parameters *before* inlining. This method runs
// prior to actual inlining and thus cannot transform the post-inlining
// code, while `processInlinedCallBlocks` does not have access to
// pre-inlining function arguments. Additionally, it is required to
// distinguish which parameter an SSA value originally came from. As a
// workaround until this is changed: Create an ssa.copy intrinsic with the
// noalias attribute (when it was present before) that can easily be found,
// and is extremely unlikely to exist in the code prior to inlining, using
// this to communicate between this method and `processInlinedCallBlocks`.
// TODO: Fix this by refactoring the inliner interface.
auto copyOp = builder.create<LLVM::SSACopyOp>(call->getLoc(), argument);
if (argumentAttrs.contains(LLVM::LLVMDialect::getNoAliasAttrName()))
copyOp->setDiscardableAttr(
builder.getStringAttr(LLVM::LLVMDialect::getNoAliasAttrName()),
builder.getUnitAttr());
return copyOp;
}
return argument;
return copyOp;
}

void processInlinedCallBlocks(
Expand Down
23 changes: 23 additions & 0 deletions mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,29 @@ llvm.func @bar(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {

// -----

// CHECK-DAG: #[[DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
// CHECK-DAG: #[[$ARG_SCOPE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[DOMAIN]]{{(,.*)?}}>

llvm.func @foo(%arg0: !llvm.ptr {llvm.noalias}, %arg1: !llvm.ptr) {
%0 = llvm.mlir.constant(5 : i64) : i64
%1 = llvm.load %arg0 {alignment = 4 : i64} : !llvm.ptr -> f32
%2 = llvm.getelementptr inbounds %arg1[%0] : (!llvm.ptr, i64) -> !llvm.ptr, f32
llvm.store %1, %2 {alignment = 4 : i64} : f32, !llvm.ptr
llvm.return
}

// CHECK-LABEL: llvm.func @missing_noalias_on_one_ptr
// CHECK: llvm.load
// CHECK-SAME: alias_scopes = [#[[$ARG_SCOPE]]]
// CHECK: llvm.store
// CHECK-SAME: noalias_scopes = [#[[$ARG_SCOPE]]]
llvm.func @missing_noalias_on_one_ptr(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
llvm.call @foo(%arg0, %arg2) : (!llvm.ptr, !llvm.ptr) -> ()
llvm.return
}

// -----

// CHECK-DAG: #[[DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
// CHECK-DAG: #[[$ARG0_SCOPE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[DOMAIN]]{{(,.*)?}}>
// CHECK-DAG: #[[$ARG1_SCOPE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[DOMAIN]]{{(,.*)?}}>
Expand Down

0 comments on commit 065d2d9

Please sign in to comment.