Skip to content

Commit 53fa92d

Browse files
NimishMishraNimishMishra
authored
[mlir][llvm][OpenMP] Hoist __atomic_load alloca (#132888)
Current implementation of `__atomic_compare_exchange` uses an alloca for `__atomic_load`, leading to issues like #120724. This PR hoists this alloca to `AllocaIP`. Fixes: #120724
1 parent 712c213 commit 53fa92d

File tree

8 files changed

+33
-18
lines changed

8 files changed

+33
-18
lines changed

‎flang/test/Integration/OpenMP/atomic-capture-complex.f90

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,14 @@
99
!RUN: %if x86-registered-target %{ %flang_fc1 -triple x86_64-unknown-linux-gnu -emit-llvm -fopenmp %s -o - | FileCheck --check-prefixes=CHECK,X86 %s %}
1010
!RUN: %if aarch64-registerd-target %{ %flang_fc1 -triple aarch64-unknown-linux-gnu -emit-llvm -fopenmp %s -o - | FileCheck --check-prefixes=CHECK,AARCH64 %s %}
1111

12+
!CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
1213
!CHECK: %[[X_NEW_VAL:.*]] = alloca { float, float }, align 8
1314
!CHECK: %[[VAL_1:.*]] = alloca { float, float }, i64 1, align 8
1415
!CHECK: %[[ORIG_VAL:.*]] = alloca { float, float }, i64 1, align 8
1516
!CHECK: store { float, float } { float 2.000000e+00, float 2.000000e+00 }, ptr %[[ORIG_VAL]], align 4
1617
!CHECK: br label %entry
1718

1819
!CHECK: entry:
19-
!CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
2020
!CHECK: call void @__atomic_load(i64 8, ptr %[[ORIG_VAL]], ptr %[[ATOMIC_TEMP_LOAD]], i32 0)
2121
!CHECK: %[[PHI_NODE_ENTRY_1:.*]] = load { float, float }, ptr %[[ATOMIC_TEMP_LOAD]], align 8
2222
!CHECK: br label %.atomic.cont

‎llvm/include/llvm/Frontend/Atomic/Atomic.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,15 @@ class AtomicInfo {
2222
Align AtomicAlign;
2323
Align ValueAlign;
2424
bool UseLibcall;
25+
IRBuilderBase::InsertPoint AllocaIP;
2526

2627
public:
2728
AtomicInfo(IRBuilderBase *Builder, Type *Ty, uint64_t AtomicSizeInBits,
2829
uint64_t ValueSizeInBits, Align AtomicAlign, Align ValueAlign,
29-
bool UseLibcall)
30+
bool UseLibcall, IRBuilderBase::InsertPoint AllocaIP)
3031
: Builder(Builder), Ty(Ty), AtomicSizeInBits(AtomicSizeInBits),
3132
ValueSizeInBits(ValueSizeInBits), AtomicAlign(AtomicAlign),
32-
ValueAlign(ValueAlign), UseLibcall(UseLibcall) {}
33+
ValueAlign(ValueAlign), UseLibcall(UseLibcall), AllocaIP(AllocaIP) {}
3334

3435
virtual ~AtomicInfo() = default;
3536

‎llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -489,9 +489,10 @@ class OpenMPIRBuilder {
489489
public:
490490
AtomicInfo(IRBuilder<> *Builder, llvm::Type *Ty, uint64_t AtomicSizeInBits,
491491
uint64_t ValueSizeInBits, llvm::Align AtomicAlign,
492-
llvm::Align ValueAlign, bool UseLibcall, llvm::Value *AtomicVar)
492+
llvm::Align ValueAlign, bool UseLibcall,
493+
IRBuilderBase::InsertPoint AllocaIP, llvm::Value *AtomicVar)
493494
: llvm::AtomicInfo(Builder, Ty, AtomicSizeInBits, ValueSizeInBits,
494-
AtomicAlign, ValueAlign, UseLibcall),
495+
AtomicAlign, ValueAlign, UseLibcall, AllocaIP),
495496
AtomicVar(AtomicVar) {}
496497

497498
llvm::Value *getAtomicPointer() const override { return AtomicVar; }
@@ -3270,11 +3271,12 @@ class OpenMPIRBuilder {
32703271
/// value
32713272
/// \param AO Atomic ordering of the generated atomic
32723273
/// instructions.
3273-
///
3274+
/// \param AllocaIP Insert point for allocas
3275+
//
32743276
/// \return Insertion point after generated atomic read IR.
32753277
InsertPointTy createAtomicRead(const LocationDescription &Loc,
32763278
AtomicOpValue &X, AtomicOpValue &V,
3277-
AtomicOrdering AO);
3279+
AtomicOrdering AO, InsertPointTy AllocaIP);
32783280

32793281
/// Emit atomic write for : X = Expr --- Only Scalar data types.
32803282
///

‎llvm/lib/Frontend/Atomic/Atomic.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,12 @@ AtomicInfo::EmitAtomicLoadLibcall(AtomicOrdering AO) {
118118
Value *PtrVal = getAtomicPointer();
119119
PtrVal = Builder->CreateAddrSpaceCast(PtrVal, PointerType::getUnqual(Ctx));
120120
Args.push_back(PtrVal);
121+
122+
auto CurrentIP = Builder->saveIP();
123+
Builder->restoreIP(AllocaIP);
121124
AllocaInst *AllocaResult =
122125
CreateAlloca(Ty, getAtomicPointer()->getName() + "atomic.temp.load");
126+
Builder->restoreIP(CurrentIP);
123127
const Align AllocaAlignment = DL.getPrefTypeAlign(SizedIntTy);
124128
AllocaResult->setAlignment(AllocaAlignment);
125129
Args.push_back(AllocaResult);

‎llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8631,7 +8631,7 @@ bool OpenMPIRBuilder::checkAndEmitFlushAfterAtomic(
86318631
OpenMPIRBuilder::InsertPointTy
86328632
OpenMPIRBuilder::createAtomicRead(const LocationDescription &Loc,
86338633
AtomicOpValue &X, AtomicOpValue &V,
8634-
AtomicOrdering AO) {
8634+
AtomicOrdering AO, InsertPointTy AllocaIP) {
86358635
if (!updateToLocation(Loc))
86368636
return Loc.IP;
86378637

@@ -8659,7 +8659,7 @@ OpenMPIRBuilder::createAtomicRead(const LocationDescription &Loc,
86598659
LoadDL.getTypeStoreSize(OldVal->getPointerOperand()->getType());
86608660
OpenMPIRBuilder::AtomicInfo atomicInfo(
86618661
&Builder, XElemTy, LoadSize * 8, LoadSize * 8, OldVal->getAlign(),
8662-
OldVal->getAlign(), true /* UseLibcall */, X.Var);
8662+
OldVal->getAlign(), true /* UseLibcall */, AllocaIP, X.Var);
86638663
auto AtomicLoadRes = atomicInfo.EmitAtomicLoadLibcall(AO);
86648664
XRead = AtomicLoadRes.first;
86658665
OldVal->eraseFromParent();
@@ -8824,7 +8824,7 @@ Expected<std::pair<Value *, Value *>> OpenMPIRBuilder::emitAtomicUpdate(
88248824

88258825
OpenMPIRBuilder::AtomicInfo atomicInfo(
88268826
&Builder, XElemTy, LoadSize * 8, LoadSize * 8, OldVal->getAlign(),
8827-
OldVal->getAlign(), true /* UseLibcall */, X);
8827+
OldVal->getAlign(), true /* UseLibcall */, AllocaIP, X);
88288828
auto AtomicLoadRes = atomicInfo.EmitAtomicLoadLibcall(AO);
88298829
BasicBlock *CurBB = Builder.GetInsertBlock();
88308830
Instruction *CurBBTI = CurBB->getTerminator();

‎llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3781,6 +3781,9 @@ TEST_F(OpenMPIRBuilderTest, OMPAtomicReadFlt) {
37813781
IRBuilder<> Builder(BB);
37823782

37833783
OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
3784+
BasicBlock *EntryBB = BB;
3785+
OpenMPIRBuilder::InsertPointTy AllocaIP(EntryBB,
3786+
EntryBB->getFirstInsertionPt());
37843787

37853788
Type *Float32 = Type::getFloatTy(M->getContext());
37863789
AllocaInst *XVal = Builder.CreateAlloca(Float32);
@@ -3791,7 +3794,7 @@ TEST_F(OpenMPIRBuilderTest, OMPAtomicReadFlt) {
37913794
OpenMPIRBuilder::AtomicOpValue X = {XVal, Float32, false, false};
37923795
OpenMPIRBuilder::AtomicOpValue V = {VVal, Float32, false, false};
37933796

3794-
Builder.restoreIP(OMPBuilder.createAtomicRead(Loc, X, V, AO));
3797+
Builder.restoreIP(OMPBuilder.createAtomicRead(Loc, X, V, AO, AllocaIP));
37953798

37963799
IntegerType *IntCastTy =
37973800
IntegerType::get(M->getContext(), Float32->getScalarSizeInBits());
@@ -3821,6 +3824,9 @@ TEST_F(OpenMPIRBuilderTest, OMPAtomicReadInt) {
38213824
IRBuilder<> Builder(BB);
38223825

38233826
OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
3827+
BasicBlock *EntryBB = BB;
3828+
OpenMPIRBuilder::InsertPointTy AllocaIP(EntryBB,
3829+
EntryBB->getFirstInsertionPt());
38243830

38253831
IntegerType *Int32 = Type::getInt32Ty(M->getContext());
38263832
AllocaInst *XVal = Builder.CreateAlloca(Int32);
@@ -3831,9 +3837,8 @@ TEST_F(OpenMPIRBuilderTest, OMPAtomicReadInt) {
38313837
OpenMPIRBuilder::AtomicOpValue X = {XVal, Int32, false, false};
38323838
OpenMPIRBuilder::AtomicOpValue V = {VVal, Int32, false, false};
38333839

3834-
BasicBlock *EntryBB = BB;
3840+
Builder.restoreIP(OMPBuilder.createAtomicRead(Loc, X, V, AO, AllocaIP));
38353841

3836-
Builder.restoreIP(OMPBuilder.createAtomicRead(Loc, X, V, AO));
38373842
LoadInst *AtomicLoad = nullptr;
38383843
StoreInst *StoreofAtomic = nullptr;
38393844

‎mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2761,6 +2761,8 @@ convertOmpAtomicRead(Operation &opInst, llvm::IRBuilderBase &builder,
27612761
return failure();
27622762

27632763
llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
2764+
llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
2765+
findAllocaInsertPoint(builder, moduleTranslation);
27642766

27652767
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
27662768

@@ -2773,7 +2775,7 @@ convertOmpAtomicRead(Operation &opInst, llvm::IRBuilderBase &builder,
27732775

27742776
llvm::OpenMPIRBuilder::AtomicOpValue V = {v, elementType, false, false};
27752777
llvm::OpenMPIRBuilder::AtomicOpValue X = {x, elementType, false, false};
2776-
builder.restoreIP(ompBuilder->createAtomicRead(ompLoc, X, V, AO));
2778+
builder.restoreIP(ompBuilder->createAtomicRead(ompLoc, X, V, AO, allocaIP));
27772779
return success();
27782780
}
27792781

‎mlir/test/Target/LLVMIR/openmp-llvm.mlir

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1368,6 +1368,7 @@ llvm.func @omp_atomic_read(%arg0 : !llvm.ptr, %arg1 : !llvm.ptr) -> () {
13681368

13691369
// CHECK-LABEL: @omp_atomic_read_implicit_cast
13701370
llvm.func @omp_atomic_read_implicit_cast () {
1371+
//CHECK: %[[ATOMIC_LOAD_TEMP:.*]] = alloca { float, float }, align 8
13711372
//CHECK: %[[Z:.*]] = alloca float, i64 1, align 4
13721373
//CHECK: %[[Y:.*]] = alloca double, i64 1, align 8
13731374
//CHECK: %[[X:.*]] = alloca [2 x { float, float }], i64 1, align 8
@@ -1392,7 +1393,7 @@ llvm.func @omp_atomic_read_implicit_cast () {
13921393
%16 = llvm.mul %10, %9 overflow<nsw> : i64
13931394
%17 = llvm.getelementptr %5[%15] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.struct<(f32, f32)>
13941395

1395-
//CHECK: %[[ATOMIC_LOAD_TEMP:.*]] = alloca { float, float }, align 8
1396+
13961397
//CHECK: call void @__atomic_load(i64 8, ptr %[[X_ELEMENT]], ptr %[[ATOMIC_LOAD_TEMP]], i32 0)
13971398
//CHECK: %[[LOAD:.*]] = load { float, float }, ptr %[[ATOMIC_LOAD_TEMP]], align 8
13981399
//CHECK: %[[EXT:.*]] = extractvalue { float, float } %[[LOAD]], 0
@@ -1480,14 +1481,14 @@ llvm.func @omp_atomic_update(%x:!llvm.ptr, %expr: i32, %xbool: !llvm.ptr, %exprb
14801481

14811482
// -----
14821483

1484+
//CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
14831485
//CHECK: %[[X_NEW_VAL:.*]] = alloca { float, float }, align 8
14841486
//CHECK: {{.*}} = alloca { float, float }, i64 1, align 8
14851487
//CHECK: %[[ORIG_VAL:.*]] = alloca { float, float }, i64 1, align 8
14861488

14871489
//CHECK: br label %entry
14881490

14891491
//CHECK: entry:
1490-
//CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
14911492
//CHECK: call void @__atomic_load(i64 8, ptr %[[ORIG_VAL]], ptr %[[ATOMIC_TEMP_LOAD]], i32 0)
14921493
//CHECK: %[[PHI_NODE_ENTRY_1:.*]] = load { float, float }, ptr %[[ATOMIC_TEMP_LOAD]], align 8
14931494
//CHECK: br label %.atomic.cont
@@ -1532,14 +1533,14 @@ llvm.func @_QPomp_atomic_update_complex() {
15321533

15331534
// -----
15341535

1536+
//CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
15351537
//CHECK: %[[X_NEW_VAL:.*]] = alloca { float, float }, align 8
15361538
//CHECK: %[[VAL_1:.*]] = alloca { float, float }, i64 1, align 8
15371539
//CHECK: %[[ORIG_VAL:.*]] = alloca { float, float }, i64 1, align 8
15381540
//CHECK: store { float, float } { float 2.000000e+00, float 2.000000e+00 }, ptr %[[ORIG_VAL]], align 4
15391541
//CHECK: br label %entry
15401542

15411543
//CHECK: entry: ; preds = %0
1542-
//CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
15431544
//CHECK: call void @__atomic_load(i64 8, ptr %[[ORIG_VAL]], ptr %[[ATOMIC_TEMP_LOAD]], i32 0)
15441545
//CHECK: %[[PHI_NODE_ENTRY_1:.*]] = load { float, float }, ptr %[[ATOMIC_TEMP_LOAD]], align 8
15451546
//CHECK: br label %.atomic.cont
@@ -1597,9 +1598,9 @@ llvm.func @_QPomp_atomic_capture_complex() {
15971598
// CHECK-LABEL: define void @omp_atomic_read_complex() {
15981599
llvm.func @omp_atomic_read_complex(){
15991600

1601+
// CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
16001602
// CHECK: %[[a:.*]] = alloca { float, float }, i64 1, align 8
16011603
// CHECK: %[[b:.*]] = alloca { float, float }, i64 1, align 8
1602-
// CHECK: %[[ATOMIC_TEMP_LOAD:.*]] = alloca { float, float }, align 8
16031604
// CHECK: call void @__atomic_load(i64 8, ptr %[[b]], ptr %[[ATOMIC_TEMP_LOAD]], i32 0)
16041605
// CHECK: %[[LOADED_VAL:.*]] = load { float, float }, ptr %[[ATOMIC_TEMP_LOAD]], align 8
16051606
// CHECK: store { float, float } %[[LOADED_VAL]], ptr %[[a]], align 4

0 commit comments

Comments
 (0)