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,14 +125,11 @@ } } -/// 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 @@ } } +/// 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 --- a/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir +++ b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir @@ -41,3 +41,157 @@ 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 +}