[OMPIRBuilder] Fix store inst alignment for ordered depend directive

[peixin]

Summary

OpenMP runtime requires depend vec with i64 type and the alignment of
store instruction should be set as 8.

[shraiysh]

Thanks for this patch @peixin. Can you please provide a reference to this requirement? (I was looking for the __kmpc_doacross_post in the RTL reference here but could not find it)

[kiranchandramohan]

In D116300#3221620, @shraiysh wrote:

Thanks for this patch @peixin. Can you please provide a reference to this requirement? (I was looking for the __kmpc_doacross_post in the RTL reference here but could not find it)

I think it is the following.
https://github.com/llvm/llvm-project/blob/bde561c4813952847112600e5efe72d9015556f7/openmp/runtime/src/kmp_csupport.cpp#L4220
https://github.com/llvm/llvm-project/blob/df20599597079d620acb9ec2161fb54f838fbbb4/openmp/runtime/src/kmp.h#L4011

[shraiysh]

Thanks Kiran. The i64 vec change LGTM (I still didn't understand why the alignment has to be 8).

[peixin]

In D116300#3221883, @shraiysh wrote:

Thanks Kiran. The i64 vec change LGTM (I still didn't understand why the alignment has to be 8).

Actually, the 4-byte alignment can also give correct running result. But the value is always stored with the type of integer 64, and the 8-byte alignment is more reasonable.

[shraiysh]

LGTM.

[peixin]

@kiranchandramohan Do you have further comments? If not, I will land this patch.

[kiranchandramohan]

LGTM.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

return EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB,
                               /*Conditional*/ false, /*hasFinalize*/ true);
 }
 
 OpenMPIRBuilder::InsertPointTy
 OpenMPIRBuilder::createOrderedDepend(const LocationDescription &Loc,
                                      InsertPointTy AllocaIP, unsigned NumLoops,
                                      ArrayRef<llvm::Value *> StoreValues,
                                      const Twine &Name, bool IsDependSource) {
+  for (size_t I = 0; I < StoreValues.size(); I++)
+    assert(StoreValues[I]->getType()->isIntegerTy(64) &&
+           "OpenMP runtime requires depend vec with i64 type");
+
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   // Allocate space for vector and generate alloc instruction.
   auto *ArrI64Ty = ArrayType::get(Int64, NumLoops);
   Builder.restoreIP(AllocaIP);
   AllocaInst *ArgsBase = Builder.CreateAlloca(ArrI64Ty, nullptr, Name);
   ArgsBase->setAlignment(Align(8));
   Builder.restoreIP(Loc.IP);
 
   // Store the index value with offset in depend vector.
   for (unsigned I = 0; I < NumLoops; ++I) {
     Value *DependAddrGEPIter = Builder.CreateInBoundsGEP(
         ArrI64Ty, ArgsBase, {Builder.getInt64(0), Builder.getInt64(I)});
-    Builder.CreateStore(StoreValues[I], DependAddrGEPIter);
+    StoreInst *STInst = Builder.CreateStore(StoreValues[I], DependAddrGEPIter);
+    STInst->setAlignment(Align(8));
   }
 
   Value *DependBaseAddrGEP = Builder.CreateInBoundsGEP(
       ArrI64Ty, ArgsBase, {Builder.getInt64(0), Builder.getInt64(0)});
 
   uint32_t SrcLocStrSize;
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
   Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

for (unsigned Iter = 0; Iter < NumLoops; Iter++) {
     ASSERT_NE(SecondIdx, nullptr);
     EXPECT_EQ(FirstIdx->getValue(), 0);
     EXPECT_EQ(SecondIdx->getValue(), Iter);
     StoreInst *StoreValue =
         dyn_cast<StoreInst>(DependAddrGEPIter->getNextNode());
     ASSERT_NE(StoreValue, nullptr);
     EXPECT_EQ(StoreValue->getValueOperand(), StoreValues[Iter]);
     EXPECT_EQ(StoreValue->getPointerOperand(), DependAddrGEPIter);
+    EXPECT_EQ(StoreValue->getAlignment(), 8);
     IterInst = dyn_cast<Instruction>(StoreValue);
   }
 
   GetElementPtrInst *DependBaseAddrGEP =
       dyn_cast<GetElementPtrInst>(IterInst->getNextNode());
   ASSERT_NE(DependBaseAddrGEP, nullptr);
   EXPECT_EQ(DependBaseAddrGEP->getPointerOperand(), AllocInst);
   EXPECT_EQ(DependBaseAddrGEP->getNumIndices(), (unsigned)2);
for (unsigned Iter = 0; Iter < NumLoops; Iter++) {
     ASSERT_NE(SecondIdx, nullptr);
     EXPECT_EQ(FirstIdx->getValue(), 0);
     EXPECT_EQ(SecondIdx->getValue(), Iter);
     StoreInst *StoreValue =
         dyn_cast<StoreInst>(DependAddrGEPIter->getNextNode());
     ASSERT_NE(StoreValue, nullptr);
     EXPECT_EQ(StoreValue->getValueOperand(), StoreValues[Iter]);
     EXPECT_EQ(StoreValue->getPointerOperand(), DependAddrGEPIter);
+    EXPECT_EQ(StoreValue->getAlignment(), 8);
     IterInst = dyn_cast<Instruction>(StoreValue);
   }
 
   GetElementPtrInst *DependBaseAddrGEP =
       dyn_cast<GetElementPtrInst>(IterInst->getNextNode());
   ASSERT_NE(DependBaseAddrGEP, nullptr);
   EXPECT_EQ(DependBaseAddrGEP->getPointerOperand(), AllocInst);
   EXPECT_EQ(DependBaseAddrGEP->getNumIndices(), (unsigned)2);

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

omp.ordered_region  {
       omp.terminator
     // CHECK: call void @__kmpc_end_ordered(%struct.ident_t* @[[GLOB1]], i32 [[OMP_THREAD]])
     }
     omp.yield
   }
 
   omp.wsloop (%arg7) : i32 = (%arg0) to (%arg1) step (%arg2) ordered(1) {
     // CHECK: [[TMP:%.*]] = getelementptr inbounds [1 x i64], [1 x i64]* [[ADDR]], i64 0, i64 0
-    // CHECK: store i64 [[ARG0:%.*]], i64* [[TMP]], align 4
+    // CHECK: store i64 [[ARG0:%.*]], i64* [[TMP]], align 8
     // CHECK: [[TMP2:%.*]] = getelementptr inbounds [1 x i64], [1 x i64]* [[ADDR]], i64 0, i64 0
     // CHECK: [[OMP_THREAD2:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB3:[0-9]+]])
     // CHECK: call void @__kmpc_doacross_wait(%struct.ident_t* @[[GLOB3]], i32 [[OMP_THREAD2]], i64* [[TMP2]])
     omp.ordered depend_type("dependsink") depend_vec(%arg3 : i64) {num_loops_val = 1 : i64}
 
     // CHECK: [[TMP3:%.*]] = getelementptr inbounds [1 x i64], [1 x i64]* [[ADDR3]], i64 0, i64 0
-    // CHECK: store i64 [[ARG0]], i64* [[TMP3]], align 4
+    // CHECK: store i64 [[ARG0]], i64* [[TMP3]], align 8
     // CHECK: [[TMP4:%.*]] = getelementptr inbounds [1 x i64], [1 x i64]* [[ADDR3]], i64 0, i64 0
     // CHECK: [[OMP_THREAD4:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB5:[0-9]+]])
     // CHECK: call void @__kmpc_doacross_post(%struct.ident_t* @[[GLOB5]], i32 [[OMP_THREAD4]], i64* [[TMP4]])
     omp.ordered depend_type("dependsource") depend_vec(%arg3 : i64) {num_loops_val = 1 : i64}
 
     omp.yield
   }
 
   omp.wsloop (%arg7) : i32 = (%arg0) to (%arg1) step (%arg2) ordered(2) {
     // CHECK: [[TMP5:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR5]], i64 0, i64 0
-    // CHECK: store i64 [[ARG0]], i64* [[TMP5]], align 4
+    // CHECK: store i64 [[ARG0]], i64* [[TMP5]], align 8
     // CHECK: [[TMP6:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR5]], i64 0, i64 1
-    // CHECK: store i64 [[ARG1:%.*]], i64* [[TMP6]], align 4
+    // CHECK: store i64 [[ARG1:%.*]], i64* [[TMP6]], align 8
     // CHECK: [[TMP7:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR5]], i64 0, i64 0
     // CHECK: [[OMP_THREAD6:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB7:[0-9]+]])
     // CHECK: call void @__kmpc_doacross_wait(%struct.ident_t* @[[GLOB7]], i32 [[OMP_THREAD6]], i64* [[TMP7]])
     // CHECK: [[TMP8:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR7]], i64 0, i64 0
-    // CHECK: store i64 [[ARG2:%.*]], i64* [[TMP8]], align 4
+    // CHECK: store i64 [[ARG2:%.*]], i64* [[TMP8]], align 8
     // CHECK: [[TMP9:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR7]], i64 0, i64 1
-    // CHECK: store i64 [[ARG3:%.*]], i64* [[TMP9]], align 4
+    // CHECK: store i64 [[ARG3:%.*]], i64* [[TMP9]], align 8
     // CHECK: [[TMP10:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR7]], i64 0, i64 0
     // CHECK: [[OMP_THREAD8:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB7]])
     // CHECK: call void @__kmpc_doacross_wait(%struct.ident_t* @[[GLOB7]], i32 [[OMP_THREAD8]], i64* [[TMP10]])
     omp.ordered depend_type("dependsink") depend_vec(%arg3, %arg4, %arg5, %arg6 : i64, i64, i64, i64) {num_loops_val = 2 : i64}
 
     // CHECK: [[TMP11:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR9]], i64 0, i64 0
-    // CHECK: store i64 [[ARG0]], i64* [[TMP11]], align 4
+    // CHECK: store i64 [[ARG0]], i64* [[TMP11]], align 8
     // CHECK: [[TMP12:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR9]], i64 0, i64 1
-    // CHECK: store i64 [[ARG1]], i64* [[TMP12]], align 4
+    // CHECK: store i64 [[ARG1]], i64* [[TMP12]], align 8
     // CHECK: [[TMP13:%.*]] = getelementptr inbounds [2 x i64], [2 x i64]* [[ADDR9]], i64 0, i64 0
     // CHECK: [[OMP_THREAD10:%.*]] = call i32 @__kmpc_global_thread_num(%struct.ident_t* @[[GLOB9:[0-9]+]])
     // CHECK: call void @__kmpc_doacross_post(%struct.ident_t* @[[GLOB9]], i32 [[OMP_THREAD10]], i64* [[TMP13]])
     omp.ordered depend_type("dependsource") depend_vec(%arg3, %arg4 : i64, i64) {num_loops_val = 2 : i64}
 
     omp.yield
   }