[OpenMP][OpenMPIRBuilder] Make outlined function parameters i64 and ptr

[jsjodin]

Summary

This patch ensures that all outlined functions parameters are i64 and ptr when compiling for a target device, which is what the OpenMP runtime expects. The values are then cast to the correct type inside the kernel.

[jdoerfert]

I would strongly recommend to make sure the arguments are 64 bit without i2p/p2i casts. This has all sorts of bad effects and it is a reason our clang kernels perform often poorly.

[jsjodin]

In D155628#4512570, @jdoerfert wrote:

I would strongly recommend to make sure the arguments are 64 bit without i2p/p2i casts. This has all sorts of bad effects and it is a reason our clang kernels perform often poorly.

If the size of a pointer is 64-bit can we use the pointers directly, unless pointers are 32-bit and then fall back on the i2p/p2i casts?

[arsenm]

In D155628#4514892, @jsjodin wrote:

In D155628#4512570, @jdoerfert wrote:

I would strongly recommend to make sure the arguments are 64 bit without i2p/p2i casts. This has all sorts of bad effects and it is a reason our clang kernels perform often poorly.

If the size of a pointer is 64-bit can we use the pointers directly, unless pointers are 32-bit and then fall back on the i2p/p2i casts?

clang ABI lowering adds padding arguments and/or structs to account for this, for the 32-bit case it would emit something like { ptr, i32 }

[jdoerfert]

In D155628#4514981, @arsenm wrote:

In D155628#4514892, @jsjodin wrote:

In D155628#4512570, @jdoerfert wrote:

I would strongly recommend to make sure the arguments are 64 bit without i2p/p2i casts. This has all sorts of bad effects and it is a reason our clang kernels perform often poorly.

If the size of a pointer is 64-bit can we use the pointers directly, unless pointers are 32-bit and then fall back on the i2p/p2i casts?

Sure. IIRC, we only support 64 bit pointers right now for offload.

[jsjodin]

Change the code to use both ptr and i64 to pass values.

[arsenm]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

4299

Assuming 64-bit pointer

[jsjodin]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

4299

Assuming 64-bit pointer

Yes, should there be a comment?

[arsenm]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

4299

Yes, I was also more thinking it's assuming a 64-bit address space specifically

[jsjodin]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

4299

Yes, both the given address space and the generic. Does LLVM allow different pointer sizes for different address spaces?

[jsjodin]

Update comment.

[jdoerfert]

LG

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

if (BodyGenCB) {
     } else {
       BeginThenGen(AllocaIP, Builder.saveIP());
     }
   }
 
   return Builder.saveIP();
 }
 
+// Copy input from pointer or i64 to the expected argument type.
+static Value *copyInput(IRBuilderBase &Builder, unsigned AddrSpace,
+                        Value *Input, Argument &Arg) {
+  auto Addr = Builder.CreateAlloca(Arg.getType()->isPointerTy()
+                                       ? Arg.getType()
+                                       : Type::getInt64Ty(Builder.getContext()),
+                                   AddrSpace);
+  auto AddrAscast =
+      Builder.CreatePointerBitCastOrAddrSpaceCast(Addr, Input->getType());
+  Builder.CreateStore(&Arg, AddrAscast);
+  auto Copy = Builder.CreateLoad(Arg.getType(), AddrAscast);
+
+  return Copy;
+}
+
 static Function *
 createOutlinedFunction(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
                        StringRef FuncName, SmallVectorImpl<Value *> &Inputs,
                        OpenMPIRBuilder::TargetBodyGenCallbackTy &CBFunc) {
   SmallVector<Type *> ParameterTypes;
+  if (OMPBuilder.Config.isTargetDevice()) {
+    // All parameters to target devices are passed as pointers
+    // or i64. This assumes 64-bit address spaces/pointers.
+    for (auto &Arg : Inputs)
+      ParameterTypes.push_back(Arg->getType()->isPointerTy()
+                                   ? Arg->getType()
+                                   : Type::getInt64Ty(Builder.getContext()));
+  } else {
     for (auto &Arg : Inputs)
       ParameterTypes.push_back(Arg->getType());
+  }
 
   auto FuncType = FunctionType::get(Builder.getVoidTy(), ParameterTypes,
                                     /*isVarArg*/ false);
   auto Func = Function::Create(FuncType, GlobalValue::InternalLinkage, FuncName,
                                Builder.GetInsertBlock()->getModule());
 
   // Save insert point.
   auto OldInsertPoint = Builder.saveIP();
 
   // Generate the region into the function.
   BasicBlock *EntryBB = BasicBlock::Create(Builder.getContext(), "entry", Func);
   Builder.SetInsertPoint(EntryBB);
 
   // Insert target init call in the device compilation pass.
   if (OMPBuilder.Config.isTargetDevice())
     Builder.restoreIP(OMPBuilder.createTargetInit(Builder, /*IsSPMD*/ false));
 
-  Builder.restoreIP(CBFunc(Builder.saveIP(), Builder.saveIP()));
+  BasicBlock *UserCodeEntryBB = Builder.GetInsertBlock();
 
   // Insert target deinit call in the device compilation pass.
+  Builder.restoreIP(CBFunc(Builder.saveIP(), Builder.saveIP()));
   if (OMPBuilder.Config.isTargetDevice())
     OMPBuilder.createTargetDeinit(Builder, /*IsSPMD*/ false);
 
   // Insert return instruction.
   Builder.CreateRetVoid();
 
   // Rewrite uses of input valus to parameters.
+  Builder.SetInsertPoint(UserCodeEntryBB->getFirstNonPHIOrDbg());
   for (auto InArg : zip(Inputs, Func->args())) {
     Value *Input = std::get<0>(InArg);
     Argument &Arg = std::get<1>(InArg);
 
+    Value *InputCopy =
+        OMPBuilder.Config.isTargetDevice()
+            ? copyInput(Builder,
+                        OMPBuilder.M.getDataLayout().getAllocaAddrSpace(),
+                        Input, Arg)
+            : &Arg;
+
     // Collect all the instructions
     for (User *User : make_early_inc_range(Input->users()))
       if (auto Instr = dyn_cast<Instruction>(User))
         if (Instr->getFunction() == Func)
-          Instr->replaceUsesOfWith(Input, &Arg);
+          Instr->replaceUsesOfWith(Input, InputCopy);
   }
 
   // Restore insert point.
   Builder.restoreIP(OldInsertPoint);
 
   return Func;
 }
 

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, TargetRegionDevice) {
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.setConfig(OpenMPIRBuilderConfig(true, false, false, false));
   OMPBuilder.initialize();
 
   F->setName("func");
   IRBuilder<> Builder(BB);
   OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
 
+  LoadInst *Value = nullptr;
   StoreInst *TargetStore = nullptr;
   llvm::SmallVector<llvm::Value *, 2> CapturedArgs = {
-      Constant::getIntegerValue(Type::getInt32Ty(Ctx), APInt(32, 0)),
+      Constant::getNullValue(Type::getInt32PtrTy(Ctx)),
       Constant::getNullValue(Type::getInt32PtrTy(Ctx))};
 
   auto BodyGenCB = [&](OpenMPIRBuilder::InsertPointTy AllocaIP,
                        OpenMPIRBuilder::InsertPointTy CodeGenIP)
       -> OpenMPIRBuilder::InsertPointTy {
     Builder.restoreIP(CodeGenIP);
-    TargetStore = Builder.CreateStore(CapturedArgs[0], CapturedArgs[1]);
+    Value = Builder.CreateLoad(Type::getInt32Ty(Ctx), CapturedArgs[0]);
+    TargetStore = Builder.CreateStore(Value, CapturedArgs[1]);
     return Builder.saveIP();
   };
 
   IRBuilder<>::InsertPoint EntryIP(&F->getEntryBlock(),
                                    F->getEntryBlock().getFirstInsertionPt());
   TargetRegionEntryInfo EntryInfo("parent", /*DeviceID=*/1, /*FileID=*/2,
                                   /*Line=*/3, /*Count=*/0);
 
   Builder.restoreIP(
       OMPBuilder.createTarget(Loc, EntryIP, EntryInfo, /*NumTeams=*/-1,
                               /*NumThreads=*/-1, CapturedArgs, BodyGenCB));
   Builder.CreateRetVoid();
   OMPBuilder.finalize();
 
   // Check outlined function
   EXPECT_FALSE(verifyModule(*M, &errs()));
   EXPECT_NE(TargetStore, nullptr);
   Function *OutlinedFn = TargetStore->getFunction();
   EXPECT_NE(F, OutlinedFn);
 
   EXPECT_TRUE(OutlinedFn->hasWeakODRLinkage());
   EXPECT_EQ(OutlinedFn->arg_size(), 2U);
   EXPECT_EQ(OutlinedFn->getName(), "__omp_offloading_1_2_parent_l3");
-  EXPECT_TRUE(OutlinedFn->getArg(0)->getType()->isIntegerTy(32));
+  EXPECT_TRUE(OutlinedFn->getArg(0)->getType()->isPointerTy());
   EXPECT_TRUE(OutlinedFn->getArg(1)->getType()->isPointerTy());
 
   // Check entry block
   auto &EntryBlock = OutlinedFn->getEntryBlock();
   Instruction *Init = EntryBlock.getFirstNonPHI();
   EXPECT_NE(Init, nullptr);
 
   auto *InitCall = dyn_cast<CallInst>(Init);
   EXPECT_NE(InitCall, nullptr);
   EXPECT_EQ(InitCall->getCalledFunction()->getName(), "__kmpc_target_init");
   EXPECT_EQ(InitCall->arg_size(), 3U);
   EXPECT_TRUE(isa<GlobalVariable>(InitCall->getArgOperand(0)));
   EXPECT_EQ(InitCall->getArgOperand(1),
             ConstantInt::get(Type::getInt8Ty(Ctx), OMP_TGT_EXEC_MODE_GENERIC));
   EXPECT_EQ(InitCall->getArgOperand(2),
             ConstantInt::get(Type::getInt1Ty(Ctx), true));
 
   auto *EntryBlockBranch = EntryBlock.getTerminator();
   EXPECT_NE(EntryBlockBranch, nullptr);
   EXPECT_EQ(EntryBlockBranch->getNumSuccessors(), 2U);
 
   // Check user code block
   auto *UserCodeBlock = EntryBlockBranch->getSuccessor(0);
   EXPECT_EQ(UserCodeBlock->getName(), "user_code.entry");
-  EXPECT_EQ(UserCodeBlock->getFirstNonPHI(), TargetStore);
-
+  auto *Alloca1 = UserCodeBlock->getFirstNonPHI();
+  EXPECT_TRUE(isa<AllocaInst>(Alloca1));
+  auto *Store1 = Alloca1->getNextNode();
+  EXPECT_TRUE(isa<StoreInst>(Store1));
+  auto *Load1 = Store1->getNextNode();
+  EXPECT_TRUE(isa<LoadInst>(Load1));
+  auto *Alloca2 = Load1->getNextNode();
+  EXPECT_TRUE(isa<AllocaInst>(Alloca2));
+  auto *Store2 = Alloca2->getNextNode();
+  EXPECT_TRUE(isa<StoreInst>(Store2));
+  auto *Load2 = Store2->getNextNode();
+  EXPECT_TRUE(isa<LoadInst>(Load2));
+
+  auto *Value1 = Load2->getNextNode();
+  EXPECT_EQ(Value1, Value);
+  EXPECT_EQ(Value1->getNextNode(), TargetStore);
   auto *Deinit = TargetStore->getNextNode();
   EXPECT_NE(Deinit, nullptr);
 
   auto *DeinitCall = dyn_cast<CallInst>(Deinit);
   EXPECT_NE(DeinitCall, nullptr);
   EXPECT_EQ(DeinitCall->getCalledFunction()->getName(), "__kmpc_target_deinit");
   EXPECT_EQ(DeinitCall->arg_size(), 2U);
   EXPECT_TRUE(isa<GlobalVariable>(DeinitCall->getArgOperand(0)));

mlir/test/Target/LLVMIR/omptarget-region-device-llvm.mlir

 
 // CHECK:      @[[SRC_LOC:.*]] = private unnamed_addr constant [23 x i8] c"{{[^"]*}}", align 1
 // CHECK:      @[[IDENT:.*]] = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @[[SRC_LOC]] }, align 8
 // CHECK:      define weak_odr protected void @__omp_offloading_{{[^_]+}}_{{[^_]+}}_omp_target_region__l{{[0-9]+}}(ptr %[[ADDR_A:.*]], ptr %[[ADDR_B:.*]], ptr %[[ADDR_C:.*]])
 // CHECK:        %[[INIT:.*]] = call i32 @__kmpc_target_init(ptr @[[IDENT]], i8 1, i1 true)
 // CHECK-NEXT:   %[[CMP:.*]] = icmp eq i32 %3, -1
 // CHECK-NEXT:   br i1 %[[CMP]], label %[[LABEL_ENTRY:.*]], label %[[LABEL_EXIT:.*]]
 // CHECK:        [[LABEL_ENTRY]]:
+// CHECK:        %[[TMP_A:.*]] = alloca ptr, align 8
+// CHECK:        store ptr %[[ADDR_A]], ptr %[[TMP_A]], align 8
+// CHECK:        %[[PTR_A:.*]] = load ptr, ptr %[[TMP_A]], align 8
+// CHECK:        %[[TMP_B:.*]] = alloca ptr, align 8
+// CHECK:        store ptr %[[ADDR_B]], ptr %[[TMP_B]], align 8
+// CHECK:        %[[PTR_B:.*]] = load ptr, ptr %[[TMP_B]], align 8
+// CHECK:        %[[TMP_C:.*]] = alloca ptr, align 8
+// CHECK:        store ptr %[[ADDR_C]], ptr %[[TMP_C]], align 8
+// CHECK:        %[[PTR_C:.*]] = load ptr, ptr %[[TMP_C]], align 8
 // CHECK-NEXT:   br label %[[LABEL_TARGET:.*]]
 // CHECK:        [[LABEL_TARGET]]:
-// CHECK:        %[[A:.*]] = load i32, ptr %[[ADDR_A]], align 4
-// CHECK:        %[[B:.*]] = load i32, ptr %[[ADDR_B]], align 4
+// CHECK:        %[[A:.*]] = load i32, ptr %[[PTR_A]], align 4
+// CHECK:        %[[B:.*]] = load i32, ptr %[[PTR_B]], align 4
 // CHECK:        %[[C:.*]] = add i32 %[[A]], %[[B]]
-// CHECK:        store i32 %[[C]], ptr %[[ADDR_C]], align 4
+// CHECK:        store i32 %[[C]], ptr %[[PTR_C]], align 4
 // CHECK:        br label %[[LABEL_DEINIT:.*]]
 // CHECK:        [[LABEL_DEINIT]]:
 // CHECK-NEXT:   call void @__kmpc_target_deinit(ptr @[[IDENT]], i8 1)
 // CHECK-NEXT:   ret void
 // CHECK:        [[LABEL_EXIT]]:
 // CHECK-NEXT:   ret void