diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp @@ -125,6 +125,69 @@ } } +/// 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 inlinedBlocks) { + DenseMap 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(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(op)) { + aliasInterface.setAliasScopes( + convertScopeList(aliasInterface.getAliasScopesOrNull())); + aliasInterface.setNoAliasScopes( + convertScopeList(aliasInterface.getNoAliasScopesOrNull())); + } + + if (auto noAliasScope = dyn_cast(op)) { + // Create the deep clones if necessary. + walker.walk(noAliasScope.getScopeAttr()); + + noAliasScope.setScopeAttr(cast( + 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 @@ // 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(op)) { - if (iface.getAliasScopesOrNull() || iface.getNoAliasScopesOrNull()) { - LLVM_DEBUG(llvm::dbgs() - << "Cannot inline: unhandled alias analysis metadata\n"); - return false; - } - } if (auto iface = dyn_cast(op)) { if (iface.getAccessGroupsOrNull()) { LLVM_DEBUG(llvm::dbgs() @@ -353,6 +409,7 @@ LLVM::MemcpyOp, LLVM::MemmoveOp, LLVM::MemsetOp, + LLVM::NoAliasScopeDeclOp, LLVM::StackRestoreOp, LLVM::StackSaveOp, LLVM::StoreOp, @@ -417,6 +474,7 @@ Operation *call, iterator_range 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 --- /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, description = "foo"> +#alias_scope = #llvm.alias_scope, domain = #alias_scope_domain, description = "foo load"> +#alias_scope1 = #llvm.alias_scope, domain = #alias_scope_domain, description = "foo store"> + +// CHECK-DAG: #[[FOO_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}> +// CHECK-DAG: #[[$FOO_LOAD:.*]] = #llvm.alias_scope +// CHECK-DAG: #[[$FOO_STORE:.*]] = #llvm.alias_scope + +// CHECK-DAG: #[[BAR_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}> +// CHECK-DAG: #[[$BAR_LOAD:.*]] = #llvm.alias_scope +// CHECK-DAG: #[[$BAR_STORE:.*]] = #llvm.alias_scope + +// 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 +}