From cde72af3dbb68dd3cc3b1012225551ccf3879713 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Markus=20B=C3=B6ck?= Date: Tue, 18 Jul 2023 15:23:27 +0200 Subject: [PATCH] [mlir][LLVM] Append call ops alias scopes to inlined operations 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 | 66 +++++++-- .../test/Dialect/LLVMIR/inlining-alias-scopes.mlir | 154 +++++++++++++++++++++ 2 files changed, 212 insertions(+), 8 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp index edffac5..61f13b3c 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp @@ -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 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 inlinedBlocks) { DenseMap 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 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 inlinedBlocks) { + auto callAliasInterface = dyn_cast(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()) { + 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 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. diff --git a/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir index 00c24ca..dd40804 100644 --- a/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir +++ b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir @@ -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, description = "hello2"> +#alias_scope_domain1 = #llvm.alias_scope_domain, description = "hello"> +#alias_scope = #llvm.alias_scope, domain = #alias_scope_domain, description = "hello2: %a"> +#alias_scope1 = #llvm.alias_scope, domain = #alias_scope_domain, description = "hello2: %b"> +#alias_scope2 = #llvm.alias_scope, domain = #alias_scope_domain1, description = "hello: %a"> + +// CHECK-DAG: #[[WITH_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}> +// CHECK-DAG: #[[$WITH_DOMAIN_SCOPE1:.*]] = #llvm.alias_scope +// CHECK-DAG: #[[$WITH_DOMAIN_SCOPE2:.*]] = #llvm.alias_scope + +// CHECK-DAG: #[[CALL_DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}> +// CHECK-DAG: #[[$CALL_DOMAIN_SCOPE:.*]] = #llvm.alias_scope + +// CHECK-DAG: #[[WITH_DOMAIN_NO_ALIAS:.*]] = #llvm.alias_scope_domain<{{.*}}> +// CHECK-DAG: #[[$WITH_DOMAIN_NO_ALIAS_SCOPE1:.*]] = #llvm.alias_scope +// CHECK-DAG: #[[$WITH_DOMAIN_NO_ALIAS_SCOPE2:.*]] = #llvm.alias_scope + +// CHECK-DAG: #[[WITH_DOMAIN_ALIAS:.*]] = #llvm.alias_scope_domain<{{.*}}> +// CHECK-DAG: #[[$WITH_DOMAIN_ALIAS_SCOPE1:.*]] = #llvm.alias_scope +// CHECK-DAG: #[[$WITH_DOMAIN_ALIAS_SCOPE2:.*]] = #llvm.alias_scope + +// 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 +} -- 2.7.4