Skip to content

[CIR][CUDA] Register global variables #1978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
+108 −4
Closed

[CIR][CUDA] Register global variables #1978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,10 @@ class CIRDataLayout {
mlir::Type getCharType(mlir::MLIRContext *ctx) const {
return typeSizeInfo.getCharType(ctx);
}

mlir::Type getSizeType(mlir::MLIRContext *ctx) const {
return typeSizeInfo.getSizeType(ctx);
}
};

/// Used to lazily calculate structure layout information for a target machine,
Expand Down
93 changes: 89 additions & 4 deletions clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,13 +127,17 @@ struct LoweringPreparePass : public LoweringPrepareBase<LoweringPreparePass> {

// Maps CUDA kernel name to device stub function.
llvm::StringMap<FuncOp> cudaKernelMap;
// Maps CUDA device-side variable name to host-side (shadow) GlobalOp.
llvm::StringMap<GlobalOp> cudaVarMap;

void buildCUDAModuleCtor();
std::optional<FuncOp> buildCUDAModuleDtor();
std::optional<FuncOp> buildCUDARegisterGlobals();

void buildCUDARegisterGlobalFunctions(cir::CIRBaseBuilderTy &builder,
FuncOp regGlobalFunc);
void buildCUDARegisterVars(cir::CIRBaseBuilderTy &builder,
FuncOp regGlobalFunc);

///
/// AST related
Expand Down Expand Up @@ -1198,8 +1202,7 @@ std::optional<FuncOp> LoweringPreparePass::buildCUDARegisterGlobals() {
builder.setInsertionPointToStart(regGlobalFunc.addEntryBlock());

buildCUDARegisterGlobalFunctions(builder, regGlobalFunc);

// TODO(cir): registration for global variables.
buildCUDARegisterVars(builder, regGlobalFunc);

ReturnOp::create(builder, loc);
return regGlobalFunc;
Expand Down Expand Up @@ -1273,6 +1276,83 @@ void LoweringPreparePass::buildCUDARegisterGlobalFunctions(
}
}

void LoweringPreparePass::buildCUDARegisterVars(cir::CIRBaseBuilderTy &builder,
FuncOp regGlobalFunc) {
auto loc = theModule.getLoc();
auto cudaPrefix = getCUDAPrefix(astCtx);

auto voidTy = VoidType::get(&getContext());
auto voidPtrTy = PointerType::get(voidTy);
auto voidPtrPtrTy = PointerType::get(voidPtrTy);
auto intTy = datalayout->getIntType(&getContext());
auto charTy = datalayout->getCharType(&getContext());
auto sizeTy = datalayout->getSizeType(&getContext());

// Extract the GPU binary handle argument.
mlir::Value fatbinHandle = *regGlobalFunc.args_begin();

cir::CIRBaseBuilderTy globalBuilder(getContext());
globalBuilder.setInsertionPointToStart(theModule.getBody());

// Declare CUDA internal function:
// void __cudaRegisterVar(
// void **fatbinHandle,
// char *hostVarName,
// char *deviceVarName,
// const char *deviceVarName,
// int isExtern, size_t varSize,
// int isConstant, int zero
// );
// Similar to the registration of global functions, OG does not care about
// pointer types. They will generate the same IR anyway.

FuncOp cudaRegisterVar = buildRuntimeFunction(
globalBuilder, addUnderscoredPrefix(cudaPrefix, "RegisterVar"), loc,
FuncType::get({voidPtrPtrTy, voidPtrTy, voidPtrTy, voidPtrTy, intTy,
sizeTy, intTy, intTy},
voidTy));

unsigned int count = 0;
auto makeConstantString = [&](llvm::StringRef str) -> GlobalOp {
auto strType = ArrayType::get(&getContext(), charTy, 1 + str.size());

auto tmpString = GlobalOp::create(
globalBuilder, loc, (".str" + str + std::to_string(count++)).str(),
strType, /*isConstant=*/true,
/*linkage=*/cir::GlobalLinkageKind::PrivateLinkage);

// We must make the string zero-terminated.
tmpString.setInitialValueAttr(ConstArrayAttr::get(
strType, StringAttr::get(&getContext(), str + "\0")));
tmpString.setPrivate();
return tmpString;
};

for (auto &[deviceSideName, global] : cudaVarMap) {
GlobalOp deviceNameStr = makeConstantString(deviceSideName);
mlir::Value deviceNameValue = builder.createBitcast(
builder.createGetGlobal(deviceNameStr), voidPtrTy);

GlobalOp hostNameStr = makeConstantString(global.getName());
mlir::Value hostNameValue =
builder.createBitcast(builder.createGetGlobal(hostNameStr), voidPtrTy);

// Every device variable that has a shadow on host will not be extern.
// See CIRGenModule::emitGlobalVarDefinition.
auto isExtern = ConstantOp::create(builder, loc, IntAttr::get(intTy, 0));
llvm::TypeSize size = datalayout->getTypeSizeInBits(global.getSymType());
auto varSize = ConstantOp::create(
builder, loc, IntAttr::get(sizeTy, size.getFixedValue() / 8));
auto isConstant = ConstantOp::create(
builder, loc, IntAttr::get(intTy, global.getConstant()));
auto zero = ConstantOp::create(builder, loc, IntAttr::get(intTy, 0));
builder.createCallOp(loc, cudaRegisterVar,
{fatbinHandle, hostNameValue, deviceNameValue,
deviceNameValue, isExtern, varSize, isConstant,
zero});
}
}

std::optional<FuncOp> LoweringPreparePass::buildCUDAModuleDtor() {
if (!theModule->getAttr(CIRDialect::getCUDABinaryHandleAttrName()))
return {};
Expand Down Expand Up @@ -1585,8 +1665,13 @@ void LoweringPreparePass::runOnOp(Operation *op) {
lowerVAArgOp(vaArgOp);
} else if (auto deleteArrayOp = dyn_cast<DeleteArrayOp>(op)) {
lowerDeleteArrayOp(deleteArrayOp);
} else if (auto getGlobal = dyn_cast<GlobalOp>(op)) {
lowerGlobalOp(getGlobal);
} else if (auto global = dyn_cast<GlobalOp>(op)) {
lowerGlobalOp(global);
if (auto attr = op->getAttr(cir::CUDAShadowNameAttr::getMnemonic())) {
auto shadowNameAttr = dyn_cast<CUDAShadowNameAttr>(attr);
std::string deviceSideName = shadowNameAttr.getDeviceSideName();
cudaVarMap[deviceSideName] = global;
}
} else if (auto dynamicCast = dyn_cast<DynamicCastOp>(op)) {
lowerDynamicCastOp(dynamicCast);
} else if (auto stdFind = dyn_cast<StdFindOp>(op)) {
Expand Down
15 changes: 15 additions & 0 deletions clang/test/CIR/CodeGen/CUDA/registration.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@

__global__ void fn() {}

__device__ int a;

// CIR-HOST: cir.func internal private @__cuda_register_globals(%[[FatbinHandle:[a-zA-Z0-9]+]]{{.*}}) {
// CIR-HOST: %[[#NULL:]] = cir.const #cir.ptr<null>
// CIR-HOST: %[[#T1:]] = cir.get_global @".str_Z2fnv"
Expand All @@ -64,6 +66,16 @@ __global__ void fn() {}
// CIR-HOST-SAME: %[[#DeviceFn]],
// CIR-HOST-SAME: %[[#MinusOne]],
// CIR-HOST-SAME: %[[#NULL]], %[[#NULL]], %[[#NULL]], %[[#NULL]], %[[#NULL]])
// CIR-HOST: %[[#T3:]] = cir.get_global @".stra0"
// CIR-HOST: %[[#Device:]] = cir.cast bitcast %7
// CIR-HOST: %[[#T4:]] = cir.get_global @".stra1"
// CIR-HOST: %[[#Host:]] = cir.cast bitcast %9
// CIR-HOST: %[[#Ext:]] = cir.const #cir.int<0>
// CIR-HOST: %[[#Sz:]] = cir.const #cir.int<4>
// CIR-HOST: %[[#Const:]] = cir.const #cir.int<0>
// CIR-HOST: %[[#Zero:]] = cir.const #cir.int<0>
// CIR-HOST: cir.call @__cudaRegisterVar(%arg0, %[[#Host]], %[[#Device]], %[[#Device]],
// CIR-HOST-SAME: %[[#Ext]], %[[#Sz]], %[[#Const]], %[[#Zero]])
// CIR-HOST: }

// LLVM-HOST: define internal void @__cuda_register_globals(ptr %[[#LLVMFatbin:]]) {
Expand All @@ -74,6 +86,9 @@ __global__ void fn() {}
// LLVM-HOST-SAME: ptr @.str_Z2fnv,
// LLVM-HOST-SAME: i32 -1,
// LLVM-HOST-SAME: ptr null, ptr null, ptr null, ptr null, ptr null)
// LLVM-HOST: call void @__cudaRegisterVar(
// LLVM-HOST-SAME: ptr %0, ptr @.stra1, ptr @.stra0, ptr @.stra0,
// LLVM-HOST-SAME: i32 0, i64 4, i32 0, i32 0)
// LLVM-HOST: }

// The content in const array should be the same as echoed above,
Expand Down
Loading