[mlir][LLVM] Append call ops alias scopes to inlined operations
authorMarkus Böck <markus.bock+llvm@nextsilicon.com>
Tue, 18 Jul 2023 13:23:27 +0000 (15:23 +0200)
committerMarkus Böck <markus.bock+llvm@nextsilicon.com>
Tue, 18 Jul 2023 14:26:17 +0000 (16:26 +0200)
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

Differential Revision: https://reviews.llvm.org/D155595

mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir

index edffac5..61f13b3 100644 (file)
@@ -125,14 +125,11 @@ handleInlinedAllocas(Operation *call,
   }
 }
 
-/// 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.
@@ -188,6 +185,59 @@ static void handleAliasScopes(Operation *call,
   }
 }
 
+/// Creates a new ArrayAttr by concatenating `lhs` with `rhs`.
+/// Returns null if both parameters are null. If only one attribute is null,
+/// return the other.
+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.
index 00c24ca..dd40804 100644 (file)
@@ -41,3 +41,157 @@ 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
+}