Skip to content

Commit 6a3982f

Browse files
DinistroDinistro
authored
[MLIR][LLVM] Relax the LLVM dialect's inliner assuming UCF (#93514)
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.
1 parent 498da62 commit 6a3982f

File tree

2 files changed

+58
-24
lines changed

2 files changed

+58
-24
lines changed

‎mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ deepCloneAliasScopes(iterator_range<Region::iterator> inlinedBlocks) {
187187
};
188188

189189
for (Block &block : inlinedBlocks) {
190-
for (Operation &op : block) {
190+
block.walk([&](Operation *op) {
191191
if (auto aliasInterface = dyn_cast<LLVM::AliasAnalysisOpInterface>(op)) {
192192
aliasInterface.setAliasScopes(
193193
convertScopeList(aliasInterface.getAliasScopesOrNull()));
@@ -202,7 +202,7 @@ deepCloneAliasScopes(iterator_range<Region::iterator> inlinedBlocks) {
202202
noAliasScope.setScopeAttr(cast<LLVM::AliasScopeAttr>(
203203
mapping.lookup(noAliasScope.getScopeAttr())));
204204
}
205-
}
205+
});
206206
}
207207
}
208208

@@ -357,9 +357,7 @@ static void createNewAliasScopesFromNoAliasParameter(
357357
// Go through every instruction and attempt to find which noalias parameters
358358
// it is definitely based on and definitely not based on.
359359
for (Block &inlinedBlock : inlinedBlocks) {
360-
for (auto aliasInterface :
361-
inlinedBlock.getOps<LLVM::AliasAnalysisOpInterface>()) {
362-
360+
inlinedBlock.walk([&](LLVM::AliasAnalysisOpInterface aliasInterface) {
363361
// Collect the pointer arguments affected by the alias scopes.
364362
SmallVector<Value> pointerArgs = aliasInterface.getAccessedOperands();
365363

@@ -395,7 +393,7 @@ static void createNewAliasScopesFromNoAliasParameter(
395393
}
396394
return true;
397395
}))
398-
continue;
396+
return;
399397

400398
// Add all noalias parameter scopes to the noalias scope list that we are
401399
// not based on.
@@ -438,7 +436,7 @@ static void createNewAliasScopesFromNoAliasParameter(
438436
// arguments.
439437
if (aliasesOtherKnownObject ||
440438
isa<LLVM::CallOp>(aliasInterface.getOperation()))
441-
continue;
439+
return;
442440

443441
SmallVector<Attribute> aliasScopes;
444442
for (LLVM::SSACopyOp noAlias : noAliasParams)
@@ -449,7 +447,7 @@ static void createNewAliasScopesFromNoAliasParameter(
449447
aliasInterface.setAliasScopes(
450448
concatArrayAttr(aliasInterface.getAliasScopesOrNull(),
451449
ArrayAttr::get(call->getContext(), aliasScopes)));
452-
}
450+
});
453451
}
454452
}
455453

@@ -472,15 +470,15 @@ appendCallOpAliasScopes(Operation *call,
472470
// Simply append the call op's alias and noalias scopes to any operation
473471
// implementing AliasAnalysisOpInterface.
474472
for (Block &block : inlinedBlocks) {
475-
for (auto aliasInterface : block.getOps<LLVM::AliasAnalysisOpInterface>()) {
473+
block.walk([&](LLVM::AliasAnalysisOpInterface aliasInterface) {
476474
if (aliasScopes)
477475
aliasInterface.setAliasScopes(concatArrayAttr(
478476
aliasInterface.getAliasScopesOrNull(), aliasScopes));
479477

480478
if (noAliasScopes)
481479
aliasInterface.setNoAliasScopes(concatArrayAttr(
482480
aliasInterface.getNoAliasScopesOrNull(), noAliasScopes));
483-
}
481+
});
484482
}
485483
}
486484

@@ -667,7 +665,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
667665
LLVM_DEBUG(llvm::dbgs() << "Cannot inline: callable is variadic\n");
668666
return false;
669667
}
670-
// TODO: Generate aliasing metadata from noalias argument/result attributes.
668+
// TODO: Generate aliasing metadata from noalias result attributes.
671669
if (auto attrs = funcOp.getArgAttrs()) {
672670
for (DictionaryAttr attrDict : attrs->getAsRange<DictionaryAttr>()) {
673671
if (attrDict.contains(LLVM::LLVMDialect::getInAllocaAttrName())) {
@@ -755,8 +753,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
755753
return handleByValArgument(builder, callable, argument, elementType,
756754
requestedAlignment);
757755
}
758-
if ([[maybe_unused]] std::optional<NamedAttribute> attr =
759-
argumentAttrs.getNamed(LLVM::LLVMDialect::getNoAliasAttrName())) {
756+
if (argumentAttrs.contains(LLVM::LLVMDialect::getNoAliasAttrName())) {
760757
if (argument.use_empty())
761758
return argument;
762759

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

Lines changed: 48 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,21 +24,24 @@ llvm.func @foo(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
2424
%0 = llvm.mlir.constant(5 : i64) : i64
2525
llvm.intr.experimental.noalias.scope.decl #alias_scope
2626
%2 = llvm.load %arg1 {alias_scopes = [#alias_scope], alignment = 4 : i64, noalias_scopes = [#alias_scope1]} : !llvm.ptr -> f32
27-
%3 = llvm.getelementptr inbounds %arg0[%0] : (!llvm.ptr, i64) -> !llvm.ptr, f32
28-
llvm.store %2, %3 {alias_scopes = [#alias_scope1], alignment = 4 : i64, noalias_scopes = [#alias_scope]} : f32, !llvm.ptr
27+
"test.one_region_op"() ({
28+
%3 = llvm.getelementptr inbounds %arg0[%0] : (!llvm.ptr, i64) -> !llvm.ptr, f32
29+
llvm.store %2, %3 {alias_scopes = [#alias_scope1], alignment = 4 : i64, noalias_scopes = [#alias_scope]} : f32, !llvm.ptr
30+
"test.terminator"() : () -> ()
31+
}) : () -> ()
2932
llvm.return
3033
}
3134

32-
// CHECK-LABEL: llvm.func @bar
35+
// CHECK-LABEL: llvm.func @clone_alias_scopes
3336
// CHECK: llvm.intr.experimental.noalias.scope.decl #[[$BAR_LOAD]]
3437
// CHECK: llvm.load
3538
// CHECK-SAME: alias_scopes = [#[[$BAR_LOAD]]]
3639
// CHECK-SAME: noalias_scopes = [#[[$BAR_STORE]]]
3740
// CHECK: llvm.store
3841
// CHECK-SAME: alias_scopes = [#[[$BAR_STORE]]]
3942
// CHECK-SAME: noalias_scopes = [#[[$BAR_LOAD]]]
40-
llvm.func @bar(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
41-
llvm.call @foo(%arg0, %arg2) : (!llvm.ptr, !llvm.ptr) -> ()
43+
llvm.func @clone_alias_scopes(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
44+
llvm.call @foo(%arg0, %arg1) : (!llvm.ptr, !llvm.ptr) -> ()
4245
llvm.return
4346
}
4447

@@ -87,9 +90,12 @@ llvm.func @callee_with_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm
8790
llvm.store %3, %4 {alias_scopes = [#alias_scope], alignment = 4 : i64, noalias_scopes = [#alias_scope1]} : f32, !llvm.ptr
8891
%5 = llvm.getelementptr inbounds %arg1[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
8992
llvm.store %3, %5 {alias_scopes = [#alias_scope1], alignment = 4 : i64, noalias_scopes = [#alias_scope]} : f32, !llvm.ptr
90-
%6 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
91-
%7 = llvm.getelementptr inbounds %arg0[%2] : (!llvm.ptr, i64) -> !llvm.ptr, f32
92-
llvm.store %6, %7 {alignment = 4 : i64} : f32, !llvm.ptr
93+
"test.one_region_op"() ({
94+
%6 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
95+
%7 = llvm.getelementptr inbounds %arg0[%2] : (!llvm.ptr, i64) -> !llvm.ptr, f32
96+
llvm.store %6, %7 {alignment = 4 : i64} : f32, !llvm.ptr
97+
"test.terminator"() : () -> ()
98+
}) : () -> ()
9399
llvm.return
94100
}
95101

@@ -105,9 +111,13 @@ llvm.func @callee_without_metadata(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !l
105111
llvm.store %3, %4 {alignment = 4 : i64} : f32, !llvm.ptr
106112
%5 = llvm.getelementptr inbounds %arg1[%1] : (!llvm.ptr, i64) -> !llvm.ptr, f32
107113
llvm.store %3, %5 {alignment = 4 : i64} : f32, !llvm.ptr
108-
%6 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
109-
%7 = llvm.getelementptr inbounds %arg0[%2] : (!llvm.ptr, i64) -> !llvm.ptr, f32
110-
llvm.store %6, %7 {alignment = 4 : i64} : f32, !llvm.ptr
114+
"test.one_region_op"() ({
115+
%6 = llvm.load %arg2 {alignment = 4 : i64} : !llvm.ptr -> f32
116+
%7 = llvm.getelementptr inbounds %arg0[%2] : (!llvm.ptr, i64) -> !llvm.ptr, f32
117+
llvm.store %6, %7 {alignment = 4 : i64} : f32, !llvm.ptr
118+
"test.terminator"() : () -> ()
119+
}) : () -> ()
120+
111121
llvm.return
112122
}
113123

@@ -394,3 +404,30 @@ llvm.func @bar(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) {
394404
llvm.call @supported_operations(%arg0, %arg2) : (!llvm.ptr, !llvm.ptr) -> ()
395405
llvm.return
396406
}
407+
408+
// -----
409+
410+
// CHECK-DAG: #[[DOMAIN:.*]] = #llvm.alias_scope_domain<{{.*}}>
411+
// CHECK-DAG: #[[$ARG_SCOPE:.*]] = #llvm.alias_scope<id = {{.*}}, domain = #[[DOMAIN]]{{(,.*)?}}>
412+
413+
llvm.func @foo(%arg: i32)
414+
415+
llvm.func @region(%arg0: !llvm.ptr {llvm.noalias}) {
416+
"test.one_region_op"() ({
417+
%1 = llvm.load %arg0 : !llvm.ptr -> i32
418+
llvm.call @foo(%1) : (i32) -> ()
419+
"test.terminator"() : () -> ()
420+
}) : () -> ()
421+
llvm.return
422+
}
423+
424+
// CHECK-LABEL: llvm.func @noalias_with_region
425+
// CHECK: llvm.load
426+
// CHECK-SAME: alias_scopes = [#[[$ARG_SCOPE]]]
427+
// CHECK: llvm.call
428+
// CHECK-NOT: alias_scopes
429+
// CHECK-SAME: noalias_scopes = [#[[$ARG_SCOPE]]]
430+
llvm.func @noalias_with_region(%arg0: !llvm.ptr) {
431+
llvm.call @region(%arg0) : (!llvm.ptr) -> ()
432+
llvm.return
433+
}

0 commit comments

Comments
 (0)