[Mlir-commits] [mlir] dd28cc7 - [mlir][LLVM] Perform deep clone of alias scopes during inlining
Markus Böck
llvmlistbot at llvm.org
Tue Jul 18 04:23:45 PDT 2023
Author: Markus Böck
Date: 2023-07-18T13:23:40+02:00
New Revision: dd28cc707c49e90c3a10ac724f6f5f46a21031c3
URL: https://github.com/llvm/llvm-project/commit/dd28cc707c49e90c3a10ac724f6f5f46a21031c3
DIFF: https://github.com/llvm/llvm-project/commit/dd28cc707c49e90c3a10ac724f6f5f46a21031c3.diff
LOG: [mlir][LLVM] Perform deep clone of alias scopes during inlining
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
Differential Revision: https://reviews.llvm.org/D155479
Added:
mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
Modified:
mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index d678c14d63acc9..edffac581b8b39 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -125,6 +125,69 @@ 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) {
+ 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
+ // attribute to make sure that new instances are always created by the
+ // uniquer.
+ 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 : inlinedBlocks) {
+ 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.
@@ -323,13 +386,6 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
// 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()
@@ -353,6 +409,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
LLVM::MemcpyOp,
LLVM::MemmoveOp,
LLVM::MemsetOp,
+ LLVM::NoAliasScopeDeclOp,
LLVM::StackRestoreOp,
LLVM::StackSaveOp,
LLVM::StoreOp,
@@ -417,6 +474,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
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
diff --git a/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
new file mode 100644
index 00000000000000..00c24cad9880db
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
@@ -0,0 +1,43 @@
+// RUN: mlir-opt %s -inline -split-input-file | FileCheck %s
+
+#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
+}
More information about the Mlir-commits
mailing list