[mlir][LLVM] Append call ops alias scopes to inlined operations

[zero9178]

Summary

Currently when inlining, any alias scope information previously attached to the call op is lost. This leads to a loss of information that could be used by alias analysis to determine that two memory access operations do not alias.

This patch fixes this issue by also taking any alias scopes of the call operation into account. These can then simply be appended onto any inlined operations.

This is analogous to the following code in LLVM: https://github.com/llvm/llvm-project/blob/1768c4597e70477af2d69f576f33400181a5f945/llvm/lib/Transforms/Utils/InlineFunction.cpp#L940

[definelicht]

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

190

"Requires" is a bit strong: If both are null, this just returns null, so I would write that explicitly.

200

You already checked that lhs it not null, so no need for the ternary operator I think?

[zero9178]

Address review comments

[definelicht]

LGTM!

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> inlinedBlocks) {
+/// 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
+deepCloneAliasScopes(iterator_range<Region::iterator> inlinedBlocks) {
   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;
 
   // Perform the deep clones while visiting. Builders create a distinct
for (Operation &op : block) {
 
         noAliasScope.setScopeAttr(cast<LLVM::AliasScopeAttr>(
             mapping.lookup(noAliasScope.getScopeAttr())));
       }
     }
   }
 }
 
+/// Creates a new ArrayAttr by concatenating `lhs` with `rhs`.
+/// Returns null if both parameters are null. Otherwise, if either attribute is
+/// null, the other is returned.
+static ArrayAttr concatArrayAttr(ArrayAttr lhs, ArrayAttr rhs) {
+  if (!lhs)
+    return rhs;
+  if (!rhs)
+    return lhs;
+
+  SmallVector<Attribute> result;
+  llvm::append_range(result, lhs);
+  llvm::append_range(result, rhs);
+  return ArrayAttr::get(lhs.getContext(), result);
+}
+
+/// Appends any alias scopes of the call operation to any inlined memory
+/// operation.
+static void
+appendCallOpAliasScopes(Operation *call,
+                        iterator_range<Region::iterator> inlinedBlocks) {
+  auto callAliasInterface = dyn_cast<LLVM::AliasAnalysisOpInterface>(call);
+  if (!callAliasInterface)
+    return;
+
+  ArrayAttr aliasScopes = callAliasInterface.getAliasScopesOrNull();
+  ArrayAttr noAliasScopes = callAliasInterface.getNoAliasScopesOrNull();
+  // If the call has neither alias scopes or noalias scopes we have nothing to
+  // do here.
+  if (!aliasScopes && !noAliasScopes)
+    return;
+
+  // Simply append the call op's alias and noalias scopes to any operation
+  // implementing AliasAnalysisOpInterface.
+  for (Block &block : inlinedBlocks) {
+    for (auto aliasInterface : block.getOps<LLVM::AliasAnalysisOpInterface>()) {
+      if (aliasScopes)
+        aliasInterface.setAliasScopes(concatArrayAttr(
+            aliasInterface.getAliasScopesOrNull(), aliasScopes));
+
+      if (noAliasScopes)
+        aliasInterface.setNoAliasScopes(concatArrayAttr(
+            aliasInterface.getNoAliasScopesOrNull(), noAliasScopes));
+    }
+  }
+}
+
+/// Handles all interactions with alias scopes during inlining.
+static void handleAliasScopes(Operation *call,
+                              iterator_range<Region::iterator> inlinedBlocks) {
+  deepCloneAliasScopes(inlinedBlocks);
+  appendCallOpAliasScopes(call, inlinedBlocks);
+}
+
 /// 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)

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

 // 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
 }
+
+// -----
+
+#alias_scope_domain = #llvm.alias_scope_domain<id = distinct[0]<>, description = "hello2">
+#alias_scope_domain1 = #llvm.alias_scope_domain<id = distinct[1]<>, description = "hello">
+#alias_scope = #llvm.alias_scope<id = distinct[2]<>, domain = #alias_scope_domain, description = "hello2: %a">
+#alias_scope1 = #llvm.alias_scope<id = distinct[3]<>, domain = #alias_scope_domain, description = "hello2: %b">
+#alias_scope2 = #llvm.alias_scope<id = distinct[4]<>, domain = #alias_scope_domain1, description = "hello: %a">
+
+// CHECK-DAG: #[[WITH_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$WITH_DOMAIN_SCOPE1:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[WITH_DOMAIN]], description = {{.*}}>
+// CHECK-DAG: #[[$WITH_DOMAIN_SCOPE2:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[WITH_DOMAIN]], description = {{.*}}>
+
+// CHECK-DAG: #[[CALL_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$CALL_DOMAIN_SCOPE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[CALL_DOMAIN]], description = {{.*}}>
+
+// CHECK-DAG: #[[WITH_DOMAIN_NO_ALIAS:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$WITH_DOMAIN_NO_ALIAS_SCOPE1:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[WITH_DOMAIN_NO_ALIAS]], description = {{.*}}>
+// CHECK-DAG: #[[$WITH_DOMAIN_NO_ALIAS_SCOPE2:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[WITH_DOMAIN_NO_ALIAS]], description = {{.*}}>
+
+// CHECK-DAG: #[[WITH_DOMAIN_ALIAS:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$WITH_DOMAIN_ALIAS_SCOPE1:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[WITH_DOMAIN_ALIAS]], description = {{.*}}>
+// CHECK-DAG: #[[$WITH_DOMAIN_ALIAS_SCOPE2:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[WITH_DOMAIN_ALIAS]], description = {{.*}}>
+
+// CHECK-LABEL: llvm.func @callee_with_metadata(
+// CHECK: llvm.load
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_SCOPE1]], #[[$WITH_DOMAIN_SCOPE2]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$WITH_DOMAIN_SCOPE1]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_SCOPE2]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$WITH_DOMAIN_SCOPE2]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_SCOPE1]]]
+// CHECK: llvm.load
+// CHECK-NOT: {{(no)?}}alias_scopes =
+// CHECK: llvm.store
+// CHECK-NOT: {{(no)?}}alias_scopes =
+llvm.func @callee_with_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+  %0 = llvm.mlir.constant(5 : i64) : i64
+  %1 = llvm.mlir.constant(8 : i64) : i64
+  %2 = llvm.mlir.constant(7 : i64) : i64
+  %3 = llvm.load %arg2 {alignment = 4 : i64, noalias_scopes = [#alias_scope, #alias_scope1]} : !llvm.ptr -> f32
+  %4 = llvm.getelementptr inbounds %arg0[%0] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %3, %4 {alias_scopes = [#alias_scope], alignment = 4 : i64, noalias_scopes = [#alias_scope1]} : f32, !llvm.ptr
+  %5 = llvm.getelementptr inbounds %arg1[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %3, %5 {alias_scopes = [#alias_scope1], alignment = 4 : i64, noalias_scopes = [#alias_scope]} : f32, !llvm.ptr
+  %6 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
+  %7 = llvm.getelementptr inbounds %arg0[%2] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %6, %7 {alignment = 4 : i64} : f32, !llvm.ptr
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @callee_without_metadata(
+// CHECK-NOT: {{(no)?}}alias_scopes =
+
+llvm.func @callee_without_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+  %0 = llvm.mlir.constant(5 : i64) : i64
+  %1 = llvm.mlir.constant(8 : i64) : i64
+  %2 = llvm.mlir.constant(7 : i64) : i64
+  %3 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
+  %4 = llvm.getelementptr inbounds %arg0[%0] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %3, %4 {alignment = 4 : i64} : f32, !llvm.ptr
+  %5 = llvm.getelementptr inbounds %arg1[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %3, %5 {alignment = 4 : i64} : f32, !llvm.ptr
+  %6 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
+  %7 = llvm.getelementptr inbounds %arg0[%2] : (!llvm.ptr, i64) -> !llvm.ptr, f32
+  llvm.store %6, %7 {alignment = 4 : i64} : f32, !llvm.ptr
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @caller(
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+
+// Inlining @callee_with_metadata with noalias_scopes.
+
+// CHECK: llvm.load
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_NO_ALIAS_SCOPE1]], #[[$WITH_DOMAIN_NO_ALIAS_SCOPE2]], #[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$WITH_DOMAIN_NO_ALIAS_SCOPE1]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_NO_ALIAS_SCOPE2]], #[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$WITH_DOMAIN_NO_ALIAS_SCOPE2]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_NO_ALIAS_SCOPE1]], #[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.load
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.store
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+
+// Inlining @callee_with_metadata with alias_scopes.
+
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_ALIAS_SCOPE1]], #[[$WITH_DOMAIN_ALIAS_SCOPE2]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$WITH_DOMAIN_ALIAS_SCOPE1]], #[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_ALIAS_SCOPE2]]]
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$WITH_DOMAIN_ALIAS_SCOPE2]], #[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-SAME: noalias_scopes = [#[[$WITH_DOMAIN_ALIAS_SCOPE1]]]
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+
+// Inlining @callee_without_metadata with noalias_scopes.
+
+// CHECK: llvm.load
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.store
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.store
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.load
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK: llvm.store
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+
+// Inlining @callee_without_metadata with alias_scopes.
+
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+// CHECK: llvm.store
+// CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
+// CHECK-NOT: noalias_scopes
+
+llvm.func @caller(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+  %0 = llvm.load %arg2 {alias_scopes = [#alias_scope2], alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
+  llvm.call @callee_with_metadata(%arg0, %arg1, %0) {noalias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
+  llvm.call @callee_with_metadata(%arg1, %arg1, %arg0) {alias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
+  llvm.call @callee_without_metadata(%arg0, %arg1, %0) {noalias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
+  llvm.call @callee_without_metadata(%arg1, %arg1, %arg0) {alias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr) -> ()
+  llvm.return
+}