[mlir][LLVM] Perform deep clone of alias scopes during inlining

[zero9178]

Summary

This is the first and most basic and important step for inlining memory operations with alias scopes.

For correctness, it is required that any alias scopes of inlined operations are replaced with deep copies. This is necessary as otherwise the same function could be inlined twice in one function, and suddenly the alias scopes extended.
A simple example would be foo(a, b); foo(a2, b2). a and a2 may alias. If foo is inlined in both instances, the store and load operations from foo may suddenly claim that a and a2 do not alias if we were to keep the original alias scopes.

This is analogous to the following class/code in LLVM: https://github.com/llvm/llvm-project/blob/4eef2e30d6f89328a16d4f1d6b37f1c79afe850c/llvm/lib/Transforms/Utils/InlineFunction.cpp#L985

[gysit]

LGTM module nit suggestion.

mlir/test/Dialect/LLVMIR/inlining.mlir

588

Would it make sense to have a separate inlining-aliascopes.mlir since I assume there will be quite a few additional tests.

[definelicht]

LGTM modulo ultra-nits!

mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp

135

This is the convention used throughout the file 🙂

138–152

It's very non-obvious where the actual cloning/copying happens, please add a comment for the uninitiated 🤠

mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp

for (auto &[allocaOp, arraySize, shouldInsertLifetime] : allocasToMove) {
       if (shouldInsertLifetime)
         builder.create<LLVM::LifetimeEndOp>(
             allocaOp.getLoc(), arraySize.getValue().getLimitedValue(),
             allocaOp.getResult());
     }
   }
 }
 
+/// Handles all interactions with alias scopes during inlining.
+/// Currently:
+/// - Maps all alias scopes in the inlined operations to deep clones of the
+///   scopes and domain. This is required for code such as
+///   `foo(a, b); foo(a2, b2);` to not incorrectly return `noalias` for e.g.
+///   operations on `a` and `a2`.
+static void handleAliasScopes(Operation *call,
+                              iterator_range<Region::iterator> inlineBlocks) {
+  DenseMap<Attribute, Attribute> mapping;
+
+  // Register handles in the walker to create the deep clones.
+  // The walker ensures that an attribute is only ever walked once and does a
+  // post-order walk, ensuring the domain is visited prior to the scope.
+  AttrTypeWalker walker;
+
+  walker.addWalk([&](LLVM::AliasScopeDomainAttr domainAttr) {
+    mapping[domainAttr] = LLVM::AliasScopeDomainAttr::get(
+        domainAttr.getContext(), domainAttr.getDescription());
+  });
+
+  walker.addWalk([&](LLVM::AliasScopeAttr scopeAttr) {
+    mapping[scopeAttr] = LLVM::AliasScopeAttr::get(
+        cast<LLVM::AliasScopeDomainAttr>(mapping.lookup(scopeAttr.getDomain())),
+        scopeAttr.getDescription());
+  });
+
+  // Map an array of scopes to an array of deep clones.
+  auto convertScopeList = [&](ArrayAttr arrayAttr) -> ArrayAttr {
+    if (!arrayAttr)
+      return nullptr;
+
+    // Create the deep clones if necessary.
+    walker.walk(arrayAttr);
+
+    return ArrayAttr::get(arrayAttr.getContext(),
+                          llvm::map_to_vector(arrayAttr, [&](Attribute attr) {
+                            return mapping.lookup(attr);
+                          }));
+  };
+
+  for (Block &block : inlineBlocks) {
+    for (Operation &op : block) {
+      if (auto aliasInterface = dyn_cast<LLVM::AliasAnalysisOpInterface>(op)) {
+        aliasInterface.setAliasScopes(
+            convertScopeList(aliasInterface.getAliasScopesOrNull()));
+        aliasInterface.setNoAliasScopes(
+            convertScopeList(aliasInterface.getNoAliasScopesOrNull()));
+      }
+
+      if (auto noAliasScope = dyn_cast<LLVM::NoAliasScopeDeclOp>(op)) {
+        // Create the deep clones if necessary.
+        walker.walk(noAliasScope.getScopeAttr());
+
+        noAliasScope.setScopeAttr(cast<LLVM::AliasScopeAttr>(
+            mapping.lookup(noAliasScope.getScopeAttr())));
+      }
+    }
+  }
+}
+
 /// If `requestedAlignment` is higher than the alignment specified on `alloca`,
 /// realigns `alloca` if this does not exceed the natural stack alignment.
 /// Returns the post-alignment of `alloca`, whether it was realigned or not.
 static unsigned tryToEnforceAllocaAlignment(LLVM::AllocaOp alloca,
                                             unsigned requestedAlignment,
                                             DataLayout const &dataLayout) {
   unsigned allocaAlignment = alloca.getAlignment().value_or(1);
   if (requestedAlignment <= allocaAlignment)
struct LLVMInlinerInterface : public DialectInlinerInterface {
 
   /// Conservative allowlist of operations supported so far.
   bool isLegalToInline(Operation *op, Region *, bool, IRMapping &) const final {
     if (isPure(op))
       return true;
     // Some attributes on memory operations require handling during
     // inlining. Since this is not yet implemented, refuse to inline memory
     // operations that have any of these attributes.
-    if (auto iface = dyn_cast<LLVM::AliasAnalysisOpInterface>(op)) {
-      if (iface.getAliasScopesOrNull() || iface.getNoAliasScopesOrNull()) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "Cannot inline: unhandled alias analysis metadata\n");
-        return false;
-      }
-    }
     if (auto iface = dyn_cast<LLVM::AccessGroupOpInterface>(op)) {
       if (iface.getAccessGroupsOrNull()) {
         LLVM_DEBUG(llvm::dbgs()
                    << "Cannot inline: unhandled access group metadata\n");
         return false;
       }
     }
     // clang-format off
     if (isa<LLVM::AllocaOp,
             LLVM::AssumeOp,
             LLVM::AtomicRMWOp,
             LLVM::AtomicCmpXchgOp,
             LLVM::CallOp,
             LLVM::DbgDeclareOp,
             LLVM::DbgValueOp,
             LLVM::FenceOp,
             LLVM::InlineAsmOp,
             LLVM::LifetimeEndOp,
             LLVM::LifetimeStartOp,
             LLVM::LoadOp,
             LLVM::MemcpyOp,
             LLVM::MemmoveOp,
             LLVM::MemsetOp,
+            LLVM::NoAliasScopeDeclOp,
             LLVM::StackRestoreOp,
             LLVM::StackSaveOp,
             LLVM::StoreOp,
             LLVM::UnreachableOp>(op))
       return true;
     // clang-format on
     LLVM_DEBUG(llvm::dbgs()
                << "Cannot inline: unhandled side effecting operation \""
Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
     }
     return argument;
   }
 
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const override {
     handleInlinedAllocas(call, inlinedBlocks);
+    handleAliasScopes(call, inlinedBlocks);
   }
 
   // Keeping this (immutable) state on the interface allows us to look up
   // StringAttrs instead of looking up strings, since StringAttrs are bound to
   // the current context and thus cannot be initialized as static fields.
   const DenseSet<StringAttr> disallowedFunctionAttrs;
 };
 
 } // end anonymous namespace
 
 void LLVM::detail::addLLVMInlinerInterface(LLVM::LLVMDialect *dialect) {
   dialect->addInterfaces<LLVMInlinerInterface>();
 }

mlir/test/Dialect/LLVMIR/inlining.mlir

 }
 
 // CHECK-LABEL: @test_disallow_arg_attr
 // CHECK-NEXT: llvm.call
 llvm.func @test_disallow_arg_attr(%ptr : !llvm.ptr) {
   llvm.call @disallowed_arg_attr(%ptr) : (!llvm.ptr) -> ()
   llvm.return
 }
+
+// -----
+
+#alias_scope_domain = #llvm.alias_scope_domain<id = distinct[0]<>, description = "foo">
+#alias_scope = #llvm.alias_scope<id = distinct[1]<>, domain = #alias_scope_domain, description = "foo load">
+#alias_scope1 = #llvm.alias_scope<id = distinct[2]<>, domain = #alias_scope_domain, description = "foo store">
+
+// CHECK-DAG: #[[FOO_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$FOO_LOAD:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[FOO_DOMAIN]], description = {{.*}}>
+// CHECK-DAG: #[[$FOO_STORE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[FOO_DOMAIN]], description = {{.*}}>
+
+// CHECK-DAG: #[[BAR_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$BAR_LOAD:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[BAR_DOMAIN]], description = {{.*}}>
+// CHECK-DAG: #[[$BAR_STORE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[BAR_DOMAIN]], description = {{.*}}>
+
+// CHECK-LABEL: llvm.func @foo
+// CHECK: llvm.intr.experimental.noalias.scope.decl #[[$FOO_LOAD]]
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$FOO_LOAD]]]
+// CHECK-SAME: noalias_scopes = [#[[$FOO_STORE]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$FOO_STORE]]]
+// CHECK-SAME: noalias_scopes = [#[[$FOO_LOAD]]]
+llvm.func @foo(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+  %0 = llvm.mlir.constant(5 : i64) : i64
+  llvm.intr.experimental.noalias.scope.decl #alias_scope
+  %2 = llvm.load %arg1 {alias_scopes = [#alias_scope], alignment = 4 : i64, noalias_scopes = [#alias_scope1]} : !llvm.ptr -> f32
+  %3 = llvm.getelementptr inbounds %arg0[%0] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %2, %3 {alias_scopes = [#alias_scope1], alignment = 4 : i64, noalias_scopes = [#alias_scope]} : f32, !llvm.ptr
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @bar
+// CHECK: llvm.intr.experimental.noalias.scope.decl #[[$BAR_LOAD]]
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$BAR_LOAD]]]
+// CHECK-SAME: noalias_scopes = [#[[$BAR_STORE]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$BAR_STORE]]]
+// CHECK-SAME: noalias_scopes = [#[[$BAR_LOAD]]]
+llvm.func @bar(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+  llvm.call @foo(%arg0, %arg2) : (!llvm.ptr, !llvm.ptr) -> ()
+  llvm.return
+}