Skip to content

Commit af8ebfd

Browse files
authored
[NVPTX] Allow the ctor/dtor lowering pass to emit kernels (llvm#71549)
Summary: This pass emits the new "nvptx$device$init" and "nvptx$device$fini" kernels that are callable by the device. This intends to mimic the method of lowering for AMDGPU where we emit `amdgcn.device.init` and `amdgcn.device.fini` respectively. These kernels simply iterate a symbol called `__init_array_start/stop` and `__fini_array_start/stop`. Normally, the linker provides these symbols automatically. In the AMDGPU case we only need call the kernel and we call the ctors / dtors. However, for NVPTX we require the user initializes these variables to the associated globals that we already emit as a part of this pass. The motivation behind this change is to move away from OpenMP's handling of ctors / dtors. I would much prefer that the backend / runtime handles this. That allows us to handle ctors / dtors in a language agnostic way, This approach requires that the runtime initializes the associated globals. They are marked `weak` so we can emit this per-TU. The kernel itself is `weak_odr` as it is copied exactly. One downside is that any module containing these kernels elicitis the "stack size cannot be statically determined warning" every time from `nvlink` which is annoying but inconsequential for functionality. It would be nice if there were a way to silence this warning however.
1 parent 3a9cc17 commit af8ebfd

File tree

2 files changed

+248
-7
lines changed

2 files changed

+248
-7
lines changed

llvm/lib/Target/NVPTX/NVPTXCtorDtorLowering.cpp

Lines changed: 186 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
//===----------------------------------------------------------------------===//
1212

1313
#include "NVPTXCtorDtorLowering.h"
14+
#include "MCTargetDesc/NVPTXBaseInfo.h"
1415
#include "NVPTX.h"
1516
#include "llvm/ADT/StringExtras.h"
1617
#include "llvm/IR/Constants.h"
@@ -32,6 +33,11 @@ static cl::opt<std::string>
3233
cl::desc("Override unique ID of ctor/dtor globals."),
3334
cl::init(""), cl::Hidden);
3435

36+
static cl::opt<bool>
37+
CreateKernels("nvptx-emit-init-fini-kernel",
38+
cl::desc("Emit kernels to call ctor/dtor globals."),
39+
cl::init(true), cl::Hidden);
40+
3541
namespace {
3642

3743
static std::string getHash(StringRef Str) {
@@ -42,11 +48,163 @@ static std::string getHash(StringRef Str) {
4248
return llvm::utohexstr(Hash.low(), /*LowerCase=*/true);
4349
}
4450

45-
static bool createInitOrFiniGlobls(Module &M, StringRef GlobalName,
46-
bool IsCtor) {
47-
GlobalVariable *GV = M.getGlobalVariable(GlobalName);
48-
if (!GV || !GV->hasInitializer())
49-
return false;
51+
static void addKernelMetadata(Module &M, GlobalValue *GV) {
52+
llvm::LLVMContext &Ctx = M.getContext();
53+
54+
// Get "nvvm.annotations" metadata node.
55+
llvm::NamedMDNode *MD = M.getOrInsertNamedMetadata("nvvm.annotations");
56+
57+
llvm::Metadata *KernelMDVals[] = {
58+
llvm::ConstantAsMetadata::get(GV), llvm::MDString::get(Ctx, "kernel"),
59+
llvm::ConstantAsMetadata::get(
60+
llvm::ConstantInt::get(llvm::Type::getInt32Ty(Ctx), 1))};
61+
62+
// This kernel is only to be called single-threaded.
63+
llvm::Metadata *ThreadXMDVals[] = {
64+
llvm::ConstantAsMetadata::get(GV), llvm::MDString::get(Ctx, "maxntidx"),
65+
llvm::ConstantAsMetadata::get(
66+
llvm::ConstantInt::get(llvm::Type::getInt32Ty(Ctx), 1))};
67+
llvm::Metadata *ThreadYMDVals[] = {
68+
llvm::ConstantAsMetadata::get(GV), llvm::MDString::get(Ctx, "maxntidy"),
69+
llvm::ConstantAsMetadata::get(
70+
llvm::ConstantInt::get(llvm::Type::getInt32Ty(Ctx), 1))};
71+
llvm::Metadata *ThreadZMDVals[] = {
72+
llvm::ConstantAsMetadata::get(GV), llvm::MDString::get(Ctx, "maxntidz"),
73+
llvm::ConstantAsMetadata::get(
74+
llvm::ConstantInt::get(llvm::Type::getInt32Ty(Ctx), 1))};
75+
76+
llvm::Metadata *BlockMDVals[] = {
77+
llvm::ConstantAsMetadata::get(GV),
78+
llvm::MDString::get(Ctx, "maxclusterrank"),
79+
llvm::ConstantAsMetadata::get(
80+
llvm::ConstantInt::get(llvm::Type::getInt32Ty(Ctx), 1))};
81+
82+
// Append metadata to nvvm.annotations.
83+
MD->addOperand(llvm::MDNode::get(Ctx, KernelMDVals));
84+
MD->addOperand(llvm::MDNode::get(Ctx, ThreadXMDVals));
85+
MD->addOperand(llvm::MDNode::get(Ctx, ThreadYMDVals));
86+
MD->addOperand(llvm::MDNode::get(Ctx, ThreadZMDVals));
87+
MD->addOperand(llvm::MDNode::get(Ctx, BlockMDVals));
88+
}
89+
90+
static Function *createInitOrFiniKernelFunction(Module &M, bool IsCtor) {
91+
StringRef InitOrFiniKernelName =
92+
IsCtor ? "nvptx$device$init" : "nvptx$device$fini";
93+
if (M.getFunction(InitOrFiniKernelName))
94+
return nullptr;
95+
96+
Function *InitOrFiniKernel = Function::createWithDefaultAttr(
97+
FunctionType::get(Type::getVoidTy(M.getContext()), false),
98+
GlobalValue::WeakODRLinkage, 0, InitOrFiniKernelName, &M);
99+
addKernelMetadata(M, InitOrFiniKernel);
100+
101+
return InitOrFiniKernel;
102+
}
103+
104+
// We create the IR required to call each callback in this section. This is
105+
// equivalent to the following code. Normally, the linker would provide us with
106+
// the definitions of the init and fini array sections. The 'nvlink' linker does
107+
// not do this so initializing these values is done by the runtime.
108+
//
109+
// extern "C" void **__init_array_start = nullptr;
110+
// extern "C" void **__init_array_end = nullptr;
111+
// extern "C" void **__fini_array_start = nullptr;
112+
// extern "C" void **__fini_array_end = nullptr;
113+
//
114+
// using InitCallback = void();
115+
// using FiniCallback = void();
116+
//
117+
// void call_init_array_callbacks() {
118+
// for (auto start = __init_array_start; start != __init_array_end; ++start)
119+
// reinterpret_cast<InitCallback *>(*start)();
120+
// }
121+
//
122+
// void call_init_array_callbacks() {
123+
// size_t fini_array_size = __fini_array_end - __fini_array_start;
124+
// for (size_t i = fini_array_size; i > 0; --i)
125+
// reinterpret_cast<FiniCallback *>(__fini_array_start[i - 1])();
126+
// }
127+
static void createInitOrFiniCalls(Function &F, bool IsCtor) {
128+
Module &M = *F.getParent();
129+
LLVMContext &C = M.getContext();
130+
131+
IRBuilder<> IRB(BasicBlock::Create(C, "entry", &F));
132+
auto *LoopBB = BasicBlock::Create(C, "while.entry", &F);
133+
auto *ExitBB = BasicBlock::Create(C, "while.end", &F);
134+
Type *PtrTy = IRB.getPtrTy(llvm::ADDRESS_SPACE_GLOBAL);
135+
136+
auto *Begin = M.getOrInsertGlobal(
137+
IsCtor ? "__init_array_start" : "__fini_array_start",
138+
PointerType::get(C, 0), [&]() {
139+
auto *GV = new GlobalVariable(
140+
M, PointerType::get(C, 0),
141+
/*isConstant=*/false, GlobalValue::WeakAnyLinkage,
142+
Constant::getNullValue(PointerType::get(C, 0)),
143+
IsCtor ? "__init_array_start" : "__fini_array_start",
144+
/*InsertBefore=*/nullptr, GlobalVariable::NotThreadLocal,
145+
/*AddressSpace=*/llvm::ADDRESS_SPACE_GLOBAL);
146+
GV->setVisibility(GlobalVariable::ProtectedVisibility);
147+
return GV;
148+
});
149+
auto *End = M.getOrInsertGlobal(
150+
IsCtor ? "__init_array_end" : "__fini_array_end", PointerType::get(C, 0),
151+
[&]() {
152+
auto *GV = new GlobalVariable(
153+
M, PointerType::get(C, 0),
154+
/*isConstant=*/false, GlobalValue::WeakAnyLinkage,
155+
Constant::getNullValue(PointerType::get(C, 0)),
156+
IsCtor ? "__init_array_end" : "__fini_array_end",
157+
/*InsertBefore=*/nullptr, GlobalVariable::NotThreadLocal,
158+
/*AddressSpace=*/llvm::ADDRESS_SPACE_GLOBAL);
159+
GV->setVisibility(GlobalVariable::ProtectedVisibility);
160+
return GV;
161+
});
162+
163+
// The constructor type is suppoed to allow using the argument vectors, but
164+
// for now we just call them with no arguments.
165+
auto *CallBackTy = FunctionType::get(IRB.getVoidTy(), {});
166+
167+
// The destructor array must be called in reverse order. Get an expression to
168+
// the end of the array and iterate backwards in that case.
169+
Value *BeginVal = IRB.CreateLoad(Begin->getType(), Begin, "begin");
170+
Value *EndVal = IRB.CreateLoad(Begin->getType(), End, "stop");
171+
if (!IsCtor) {
172+
auto *BeginInt = IRB.CreatePtrToInt(BeginVal, IntegerType::getInt64Ty(C));
173+
auto *EndInt = IRB.CreatePtrToInt(EndVal, IntegerType::getInt64Ty(C));
174+
auto *SubInst = IRB.CreateSub(EndInt, BeginInt);
175+
auto *Offset = IRB.CreateAShr(
176+
SubInst, ConstantInt::get(IntegerType::getInt64Ty(C), 3), "offset",
177+
/*IsExact=*/true);
178+
auto *ValuePtr = IRB.CreateGEP(PointerType::get(C, 0), BeginVal,
179+
ArrayRef<Value *>({Offset}));
180+
EndVal = BeginVal;
181+
BeginVal = IRB.CreateInBoundsGEP(
182+
PointerType::get(C, 0), ValuePtr,
183+
ArrayRef<Value *>(ConstantInt::get(IntegerType::getInt64Ty(C), -1)),
184+
"start");
185+
}
186+
IRB.CreateCondBr(
187+
IRB.CreateCmp(IsCtor ? ICmpInst::ICMP_NE : ICmpInst::ICMP_UGT, BeginVal,
188+
EndVal),
189+
LoopBB, ExitBB);
190+
IRB.SetInsertPoint(LoopBB);
191+
auto *CallBackPHI = IRB.CreatePHI(PtrTy, 2, "ptr");
192+
auto *CallBack = IRB.CreateLoad(CallBackTy->getPointerTo(F.getAddressSpace()),
193+
CallBackPHI, "callback");
194+
IRB.CreateCall(CallBackTy, CallBack);
195+
auto *NewCallBack =
196+
IRB.CreateConstGEP1_64(PtrTy, CallBackPHI, IsCtor ? 1 : -1, "next");
197+
auto *EndCmp = IRB.CreateCmp(IsCtor ? ICmpInst::ICMP_EQ : ICmpInst::ICMP_ULT,
198+
NewCallBack, EndVal, "end");
199+
CallBackPHI->addIncoming(BeginVal, &F.getEntryBlock());
200+
CallBackPHI->addIncoming(NewCallBack, LoopBB);
201+
IRB.CreateCondBr(EndCmp, ExitBB, LoopBB);
202+
IRB.SetInsertPoint(ExitBB);
203+
IRB.CreateRetVoid();
204+
}
205+
206+
static bool createInitOrFiniGlobals(Module &M, GlobalVariable *GV,
207+
bool IsCtor) {
50208
ConstantArray *GA = dyn_cast<ConstantArray>(GV->getInitializer());
51209
if (!GA || GA->getNumOperands() == 0)
52210
return false;
@@ -81,14 +239,35 @@ static bool createInitOrFiniGlobls(Module &M, StringRef GlobalName,
81239
appendToUsed(M, {GV});
82240
}
83241

242+
return true;
243+
}
244+
245+
static bool createInitOrFiniKernel(Module &M, StringRef GlobalName,
246+
bool IsCtor) {
247+
GlobalVariable *GV = M.getGlobalVariable(GlobalName);
248+
if (!GV || !GV->hasInitializer())
249+
return false;
250+
251+
if (!createInitOrFiniGlobals(M, GV, IsCtor))
252+
return false;
253+
254+
if (!CreateKernels)
255+
return true;
256+
257+
Function *InitOrFiniKernel = createInitOrFiniKernelFunction(M, IsCtor);
258+
if (!InitOrFiniKernel)
259+
return false;
260+
261+
createInitOrFiniCalls(*InitOrFiniKernel, IsCtor);
262+
84263
GV->eraseFromParent();
85264
return true;
86265
}
87266

88267
static bool lowerCtorsAndDtors(Module &M) {
89268
bool Modified = false;
90-
Modified |= createInitOrFiniGlobls(M, "llvm.global_ctors", /*IsCtor =*/true);
91-
Modified |= createInitOrFiniGlobls(M, "llvm.global_dtors", /*IsCtor =*/false);
269+
Modified |= createInitOrFiniKernel(M, "llvm.global_ctors", /*IsCtor =*/true);
270+
Modified |= createInitOrFiniKernel(M, "llvm.global_dtors", /*IsCtor =*/false);
92271
return Modified;
93272
}
94273

llvm/test/CodeGen/NVPTX/lower-ctor-dtor.ll

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --include-generated-funcs --version 3
12
; RUN: opt -S -mtriple=nvptx64-- -nvptx-lower-ctor-dtor < %s | FileCheck %s
23
; RUN: opt -S -mtriple=nvptx64-- -passes=nvptx-lower-ctor-dtor < %s | FileCheck %s
34
; RUN: opt -S -mtriple=nvptx64-- -passes=nvptx-lower-ctor-dtor \
45
; RUN: -nvptx-lower-global-ctor-dtor-id=unique_id < %s | FileCheck %s --check-prefix=GLOBAL
6+
; RUN: opt -S -mtriple=nvptx64-- -passes=nvptx-lower-ctor-dtor \
7+
; RUN: -nvptx-emit-init-fini-kernel=false < %s | FileCheck %s --check-prefix=KERNEL
58

69
; Make sure we get the same result if we run multiple times
710
; RUN: opt -S -mtriple=nvptx64-- -passes=nvptx-lower-ctor-dtor,nvptx-lower-ctor-dtor < %s | FileCheck %s
@@ -10,15 +13,24 @@
1013
@llvm.global_ctors = appending addrspace(1) global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 1, ptr @foo, ptr null }]
1114
@llvm.global_dtors = appending addrspace(1) global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 1, ptr @bar, ptr null }]
1215

16+
1317
; CHECK-NOT: @llvm.global_ctors
1418
; CHECK-NOT: @llvm.global_dtors
1519

1620
; CHECK: @__init_array_object_foo_[[HASH:[0-9a-f]+]]_1 = protected addrspace(4) constant ptr @foo, section ".init_array.1"
1721
; CHECK: @__fini_array_object_bar_[[HASH:[0-9a-f]+]]_1 = protected addrspace(4) constant ptr @bar, section ".fini_array.1"
1822
; CHECK: @llvm.used = appending global [2 x ptr] [ptr addrspacecast (ptr addrspace(4) @__init_array_object_foo_[[HASH]]_1 to ptr), ptr addrspacecast (ptr addrspace(4) @__fini_array_object_bar_[[HASH]]_1 to ptr)], section "llvm.metadata"
23+
; CHECK: @__fini_array_start = weak protected addrspace(1) global ptr null
24+
; CHECK: @__fini_array_end = weak protected addrspace(1) global ptr null
25+
1926
; GLOBAL: @__init_array_object_foo_unique_id_1 = protected addrspace(4) constant ptr @foo, section ".init_array.1"
2027
; GLOBAL: @__fini_array_object_bar_unique_id_1 = protected addrspace(4) constant ptr @bar, section ".fini_array.1"
2128
; GLOBAL: @llvm.used = appending global [2 x ptr] [ptr addrspacecast (ptr addrspace(4) @__init_array_object_foo_unique_id_1 to ptr), ptr addrspacecast (ptr addrspace(4) @__fini_array_object_bar_unique_id_1 to ptr)], section "llvm.metadata"
29+
; GLOBAL: @__fini_array_start = weak protected addrspace(1) global ptr null
30+
; GLOBAL: @__fini_array_end = weak protected addrspace(1) global ptr null
31+
32+
; KERNEL: @__init_array_object_foo_[[HASH:[0-9a-f]+]]_1 = protected addrspace(4) constant ptr @foo, section ".init_array.1"
33+
; KERNEL: @__fini_array_object_bar_[[HASH:[0-9a-f]+]]_1 = protected addrspace(4) constant ptr @bar, section ".fini_array.1"
2234

2335
; VISIBILITY: .visible .const .align 8 .u64 __init_array_object_foo_[[HASH:[0-9a-f]+]]_1 = foo;
2436
; VISIBILITY: .visible .const .align 8 .u64 __fini_array_object_bar_[[HASH:[0-9a-f]+]]_1 = bar;
@@ -30,3 +42,53 @@ define internal void @foo() {
3042
define internal void @bar() {
3143
ret void
3244
}
45+
46+
; CHECK-LABEL: define weak_odr void @"nvptx$device$init"() {
47+
; CHECK-NEXT: entry:
48+
; CHECK-NEXT: [[BEGIN:%.*]] = load ptr addrspace(1), ptr addrspace(1) @__init_array_start, align 8
49+
; CHECK-NEXT: [[STOP:%.*]] = load ptr addrspace(1), ptr addrspace(1) @__init_array_end, align 8
50+
; CHECK-NEXT: [[TMP0:%.*]] = icmp ne ptr addrspace(1) [[BEGIN]], [[STOP]]
51+
; CHECK-NEXT: br i1 [[TMP0]], label [[WHILE_ENTRY:%.*]], label [[WHILE_END:%.*]]
52+
; CHECK: while.entry:
53+
; CHECK-NEXT: [[PTR:%.*]] = phi ptr addrspace(1) [ [[BEGIN]], [[ENTRY:%.*]] ], [ [[NEXT:%.*]], [[WHILE_ENTRY]] ]
54+
; CHECK-NEXT: [[CALLBACK:%.*]] = load ptr, ptr addrspace(1) [[PTR]], align 8
55+
; CHECK-NEXT: call void [[CALLBACK]]()
56+
; CHECK-NEXT: [[NEXT]] = getelementptr ptr addrspace(1), ptr addrspace(1) [[PTR]], i64 1
57+
; CHECK-NEXT: [[END:%.*]] = icmp eq ptr addrspace(1) [[NEXT]], [[STOP]]
58+
; CHECK-NEXT: br i1 [[END]], label [[WHILE_END]], label [[WHILE_ENTRY]]
59+
; CHECK: while.end:
60+
; CHECK-NEXT: ret void
61+
;
62+
;
63+
; CHECK-LABEL: define weak_odr void @"nvptx$device$fini"() {
64+
; CHECK-NEXT: entry:
65+
; CHECK-NEXT: [[BEGIN:%.*]] = load ptr addrspace(1), ptr addrspace(1) @__fini_array_start, align 8
66+
; CHECK-NEXT: [[STOP:%.*]] = load ptr addrspace(1), ptr addrspace(1) @__fini_array_end, align 8
67+
; CHECK-NEXT: [[TMP0:%.*]] = ptrtoint ptr addrspace(1) [[BEGIN]] to i64
68+
; CHECK-NEXT: [[TMP1:%.*]] = ptrtoint ptr addrspace(1) [[STOP]] to i64
69+
; CHECK-NEXT: [[TMP2:%.*]] = sub i64 [[TMP1]], [[TMP0]]
70+
; CHECK-NEXT: [[OFFSET:%.*]] = ashr exact i64 [[TMP2]], 3
71+
; CHECK-NEXT: [[TMP3:%.*]] = getelementptr ptr, ptr addrspace(1) [[BEGIN]], i64 [[OFFSET]]
72+
; CHECK-NEXT: [[START:%.*]] = getelementptr inbounds ptr, ptr addrspace(1) [[TMP3]], i64 -1
73+
; CHECK-NEXT: [[TMP4:%.*]] = icmp ugt ptr addrspace(1) [[START]], [[BEGIN]]
74+
; CHECK-NEXT: br i1 [[TMP4]], label [[WHILE_ENTRY:%.*]], label [[WHILE_END:%.*]]
75+
; CHECK: while.entry:
76+
; CHECK-NEXT: [[PTR:%.*]] = phi ptr addrspace(1) [ [[START]], [[ENTRY:%.*]] ], [ [[NEXT:%.*]], [[WHILE_ENTRY]] ]
77+
; CHECK-NEXT: [[CALLBACK:%.*]] = load ptr, ptr addrspace(1) [[PTR]], align 8
78+
; CHECK-NEXT: call void [[CALLBACK]]()
79+
; CHECK-NEXT: [[NEXT]] = getelementptr ptr addrspace(1), ptr addrspace(1) [[PTR]], i64 -1
80+
; CHECK-NEXT: [[END:%.*]] = icmp ult ptr addrspace(1) [[NEXT]], [[BEGIN]]
81+
; CHECK-NEXT: br i1 [[END]], label [[WHILE_END]], label [[WHILE_ENTRY]]
82+
; CHECK: while.end:
83+
; CHECK-NEXT: ret void
84+
85+
; CHECK: [[META0:![0-9]+]] = !{ptr @"nvptx$device$init", !"kernel", i32 1}
86+
; CHECK: [[META1:![0-9]+]] = !{ptr @"nvptx$device$init", !"maxntidx", i32 1}
87+
; CHECK: [[META2:![0-9]+]] = !{ptr @"nvptx$device$init", !"maxntidy", i32 1}
88+
; CHECK: [[META3:![0-9]+]] = !{ptr @"nvptx$device$init", !"maxntidz", i32 1}
89+
; CHECK: [[META4:![0-9]+]] = !{ptr @"nvptx$device$init", !"maxclusterrank", i32 1}
90+
; CHECK: [[META5:![0-9]+]] = !{ptr @"nvptx$device$fini", !"kernel", i32 1}
91+
; CHECK: [[META6:![0-9]+]] = !{ptr @"nvptx$device$fini", !"maxntidx", i32 1}
92+
; CHECK: [[META7:![0-9]+]] = !{ptr @"nvptx$device$fini", !"maxntidy", i32 1}
93+
; CHECK: [[META8:![0-9]+]] = !{ptr @"nvptx$device$fini", !"maxntidz", i32 1}
94+
; CHECK: [[META9:![0-9]+]] = !{ptr @"nvptx$device$fini", !"maxclusterrank", i32 1}

0 commit comments

Comments
 (0)