[MLIR][LLVM] Improve the noalias propagation during inlining (#104750)

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".

Author: Dinistro

mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp

@@ -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(
@@ -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
@@ -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;
@@ -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(

mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir

@@ -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]]{{(,.*)?}}>