[MLIR][LLVM] Relax the LLVM dialect's inliner assuming UCF #93514

Dinistro · 2024-05-28T09:06:08Z

This commit changes the LLVM dialect's inliner interface to stop assuming that the inlined function only contained unstructured control flow. This is not necessarily true, and it lead to not properly propagating the noalias information.

llvmbot · 2024-05-28T09:06:42Z

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Christian Ulmann (Dinistro)

Changes

This commit changes the LLVM dialect's inliner interface to stop assuming that the inlined function only contained unstructured control flow. This is not necessarily true, and it lead to not properly propagating the noalias information.

Full diff: https://github.com/llvm/llvm-project/pull/93514.diff

2 Files Affected:

(modified) mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp (+10-13)
(modified) mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir (+52-19)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 4a6154ea6d300..5552dc5e244b8 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -187,7 +187,7 @@ deepCloneAliasScopes(iterator_range<Region::iterator> inlinedBlocks) {
   };
 
   for (Block &block : inlinedBlocks) {
-    for (Operation &op : block) {
+    block.walk([&](Operation *op) {
       if (auto aliasInterface = dyn_cast<LLVM::AliasAnalysisOpInterface>(op)) {
         aliasInterface.setAliasScopes(
             convertScopeList(aliasInterface.getAliasScopesOrNull()));
@@ -202,7 +202,7 @@ deepCloneAliasScopes(iterator_range<Region::iterator> inlinedBlocks) {
         noAliasScope.setScopeAttr(cast<LLVM::AliasScopeAttr>(
             mapping.lookup(noAliasScope.getScopeAttr())));
       }
-    }
+    });
   }
 }
 
@@ -357,9 +357,7 @@ static void createNewAliasScopesFromNoAliasParameter(
   // Go through every instruction and attempt to find which noalias parameters
   // it is definitely based on and definitely not based on.
   for (Block &inlinedBlock : inlinedBlocks) {
-    for (auto aliasInterface :
-         inlinedBlock.getOps<LLVM::AliasAnalysisOpInterface>()) {
-
+    inlinedBlock.walk([&](LLVM::AliasAnalysisOpInterface aliasInterface) {
       // Collect the pointer arguments affected by the alias scopes.
       SmallVector<Value> pointerArgs = aliasInterface.getAccessedOperands();
 
@@ -395,7 +393,7 @@ static void createNewAliasScopesFromNoAliasParameter(
             }
             return true;
           }))
-        continue;
+        return;
 
       // Add all noalias parameter scopes to the noalias scope list that we are
       // not based on.
@@ -438,7 +436,7 @@ static void createNewAliasScopesFromNoAliasParameter(
       // arguments.
       if (aliasesOtherKnownObject ||
           isa<LLVM::CallOp>(aliasInterface.getOperation()))
-        continue;
+        return;
 
       SmallVector<Attribute> aliasScopes;
       for (LLVM::SSACopyOp noAlias : noAliasParams)
@@ -449,7 +447,7 @@ static void createNewAliasScopesFromNoAliasParameter(
         aliasInterface.setAliasScopes(
             concatArrayAttr(aliasInterface.getAliasScopesOrNull(),
                             ArrayAttr::get(call->getContext(), aliasScopes)));
-    }
+    });
   }
 }
 
@@ -472,7 +470,7 @@ appendCallOpAliasScopes(Operation *call,
   // 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>()) {
+    block.walk([&](LLVM::AliasAnalysisOpInterface aliasInterface) {
       if (aliasScopes)
         aliasInterface.setAliasScopes(concatArrayAttr(
             aliasInterface.getAliasScopesOrNull(), aliasScopes));
@@ -480,7 +478,7 @@ appendCallOpAliasScopes(Operation *call,
       if (noAliasScopes)
         aliasInterface.setNoAliasScopes(concatArrayAttr(
             aliasInterface.getNoAliasScopesOrNull(), noAliasScopes));
-    }
+    });
   }
 }
 
@@ -667,7 +665,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
       LLVM_DEBUG(llvm::dbgs() << "Cannot inline: callable is variadic\n");
       return false;
     }
-    // TODO: Generate aliasing metadata from noalias argument/result attributes.
+    // TODO: Generate aliasing metadata from noalias result attributes.
     if (auto attrs = funcOp.getArgAttrs()) {
       for (DictionaryAttr attrDict : attrs->getAsRange<DictionaryAttr>()) {
         if (attrDict.contains(LLVM::LLVMDialect::getInAllocaAttrName())) {
@@ -755,8 +753,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
       return handleByValArgument(builder, callable, argument, elementType,
                                  requestedAlignment);
     }
-    if ([[maybe_unused]] std::optional<NamedAttribute> attr =
-            argumentAttrs.getNamed(LLVM::LLVMDialect::getNoAliasAttrName())) {
+    if (argumentAttrs.contains(LLVM::LLVMDialect::getNoAliasAttrName())) {
       if (argument.use_empty())
         return argument;
 
diff --git a/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
index 29450833bee59..38561bec116f6 100644
--- a/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir
@@ -20,16 +20,18 @@
 // CHECK: llvm.store
 // CHECK-SAME: alias_scopes = [#[[$FOO_STORE]]]
 // CHECK-SAME: noalias_scopes = [#[[$FOO_LOAD]]]
-llvm.func @foo(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+llvm.func @foo(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %c: i1) {
   %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
+  scf.if %c {
+    %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-LABEL: llvm.func @clone_alias_scopes
 // CHECK: llvm.intr.experimental.noalias.scope.decl #[[$BAR_LOAD]]
 // CHECK: llvm.load
 // CHECK-SAME: alias_scopes = [#[[$BAR_LOAD]]]
@@ -37,8 +39,8 @@ llvm.func @foo(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
 // 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.func @clone_alias_scopes(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: i1) {
+  llvm.call @foo(%arg0, %arg1, %arg2) : (!llvm.ptr, !llvm.ptr, i1) -> ()
   llvm.return
 }
 
@@ -78,7 +80,7 @@ llvm.func @bar(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
 // 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) {
+llvm.func @callee_with_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %c: i1) {
   %0 = llvm.mlir.constant(5 : i64) : i64
   %1 = llvm.mlir.constant(8 : i64) : i64
   %2 = llvm.mlir.constant(7 : i64) : i64
@@ -87,16 +89,18 @@ llvm.func @callee_with_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm
   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
+  scf.if %c {
+    %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) {
+llvm.func @callee_without_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %c: i1) {
   %0 = llvm.mlir.constant(5 : i64) : i64
   %1 = llvm.mlir.constant(8 : i64) : i64
   %2 = llvm.mlir.constant(7 : i64) : i64
@@ -105,9 +109,11 @@ llvm.func @callee_without_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !l
   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
+  scf.if %c {
+    %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
 }
 
@@ -187,12 +193,12 @@ llvm.func @callee_without_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !l
 // CHECK-SAME: alias_scopes = [#[[$CALL_DOMAIN_SCOPE]]]
 // CHECK-NOT: noalias_scopes
 
-llvm.func @caller(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
+llvm.func @caller(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i1) {
   %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.call @callee_with_metadata(%arg0, %arg1, %0, %arg3) {noalias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i1) -> ()
+  llvm.call @callee_with_metadata(%arg1, %arg1, %arg0, %arg3) {alias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i1) -> ()
+  llvm.call @callee_without_metadata(%arg0, %arg1, %0, %arg3) {noalias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i1) -> ()
+  llvm.call @callee_without_metadata(%arg1, %arg1, %arg0, %arg3) {alias_scopes = [#alias_scope2]} : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i1) -> ()
   llvm.return
 }
 
@@ -394,3 +400,30 @@ llvm.func @bar(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
   llvm.call @supported_operations(%arg0, %arg2) : (!llvm.ptr, !llvm.ptr) -> ()
   llvm.return
 }
+
+// -----
+
+// CHECK-DAG: #[[DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
+// CHECK-DAG: #[[$ARG_SCOPE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[DOMAIN]]{{(,.*)?}}>
+
+llvm.func @foo(%arg: i32)
+
+llvm.func @region(%arg0: !llvm.ptr {llvm.noalias}, %c: i1) {
+  scf.if %c {
+    %1 = llvm.load %arg0 : !llvm.ptr -> i32
+    llvm.call @foo(%1) : (i32) -> ()
+    scf.yield
+  }
+  llvm.return
+}
+
+// CHECK-LABEL: llvm.func @bar
+// CHECK: llvm.load
+// CHECK-SAME: alias_scopes = [#[[$ARG_SCOPE]]]
+// CHECK: llvm.call
+// CHECK-NOT: alias_scopes
+// CHECK-SAME: noalias_scopes = [#[[$ARG_SCOPE]]]
+llvm.func @noalias_with_region(%arg0: !llvm.ptr, %arg1: i1) {
+  llvm.call @region(%arg0, %arg1) : (!llvm.ptr, i1) -> ()
+  llvm.return
+}

gysit

LGTM modulo nit.

mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir

Dinistro requested review from metaflow, definelicht and zero9178 May 28, 2024 09:06

llvmbot added mlir:llvm mlir labels May 28, 2024

Dinistro added 2 commits May 28, 2024 10:50

fix test

ebf5dbd

replace scf with test dialect

6b0d045

gysit approved these changes May 28, 2024

View reviewed changes

mlir/test/Dialect/LLVMIR/inlining-alias-scopes.mlir Outdated Show resolved Hide resolved

fix test change remnant

4f0f7c8

Dinistro merged commit 6a3982f into main May 30, 2024
7 checks passed

Dinistro deleted the users/dinistro/relax-llvm-dialect-inlining-assumptions branch May 30, 2024 05:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][LLVM] Relax the LLVM dialect's inliner assuming UCF #93514

[MLIR][LLVM] Relax the LLVM dialect's inliner assuming UCF #93514

Uh oh!

Dinistro commented May 28, 2024

Uh oh!

llvmbot commented May 28, 2024 •

edited

Loading

Uh oh!

gysit left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[MLIR][LLVM] Relax the LLVM dialect's inliner assuming UCF #93514

[MLIR][LLVM] Relax the LLVM dialect's inliner assuming UCF #93514

Uh oh!

Conversation

Dinistro commented May 28, 2024

Uh oh!

llvmbot commented May 28, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gysit left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

llvmbot commented May 28, 2024 •

edited

Loading