Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
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
97 changes: 90 additions & 7 deletions clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,8 @@ 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();
Expand All @@ -135,6 +137,8 @@ struct LoweringPreparePass : public LoweringPrepareBase<LoweringPreparePass> {

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

///
/// AST related
Expand Down Expand Up @@ -1261,8 +1265,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 @@ -1409,6 +1412,81 @@ std::optional<FuncOp> LoweringPreparePass::buildHIPModuleDtor() {
return dtor;
}

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 varNameStr = makeConstantString(deviceSideName);
mlir::Value varNameValue =
builder.createBitcast(builder.createGetGlobal(varNameStr), voidPtrTy);

auto globalVarValue =
builder.createBitcast(builder.createGetGlobal(global), 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, globalVarValue, varNameValue,
varNameValue, isExtern, varSize, isConstant, zero});
}
}

std::optional<FuncOp> LoweringPreparePass::buildCUDAModuleDtor() {
if (!theModule->getAttr(CIRDialect::getCUDABinaryHandleAttrName()))
return {};
Expand All @@ -1431,9 +1509,9 @@ std::optional<FuncOp> LoweringPreparePass::buildCUDAModuleDtor() {

// void __cuda_module_dtor();
// Despite the name, OG doesn't treat it as a destructor, so it shouldn't be
// put into globalDtorList. If it were a real dtor, then it would cause double
// free above CUDA 9.2. The way to use it is to manually call atexit() at end
// of module ctor.
// put into globalDtorList. If it were a real dtor, then it would cause
// double free above CUDA 9.2. The way to use it is to manually call
// atexit() at end of module ctor.
std::string dtorName = addUnderscoredPrefix(prefix, "_module_dtor");
FuncOp dtor =
buildRuntimeFunction(builder, dtorName, loc, FuncType::get({}, voidTy),
Expand Down Expand Up @@ -1721,8 +1799,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
70 changes: 70 additions & 0 deletions clang/test/CIR/CodeGen/CUDA/registration.cu
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,13 @@
// RUN: %s -o %t.ll
// RUN: FileCheck --check-prefix=LLVM-HOST --input-file=%t.ll %s

// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu \
// RUN: -x cuda -emit-llvm -target-sdk-version=12.3 \
// RUN: -fcuda-include-gpubinary %t.fatbin \
// RUN: %s -o %t.ll
// RUN: FileCheck --check-prefix=OGCG-HOST --input-file=%t.ll %s


// CIR-HOST: module @"{{.*}}" attributes {
// CIR-HOST: cir.cu.binary_handle = #cir.cu.binary_handle<{{.*}}.fatbin>,
// CIR-HOST: cir.global_ctors = [#cir.global_ctor<"__cuda_module_ctor", {{[0-9]+}}>]
Expand Down Expand Up @@ -50,6 +57,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 +73,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 %[[#T3]]
// CIR-HOST: %[[#T4:]] = cir.get_global @a
// CIR-HOST: %[[#Host:]] = cir.cast bitcast %[[#T4]]
// 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 +93,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 @a, 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 Expand Up @@ -110,3 +132,51 @@ __global__ void fn() {}
// LLVM-HOST: call void @__cudaRegisterFatBinaryEnd
// LLVM-HOST: call i32 @atexit(ptr @__cuda_module_dtor)
// LLVM-HOST: }

// OGCG-HOST: @a = internal global i32 undef, align 4
// OGCG-HOST: @0 = private unnamed_addr constant [7 x i8] c"_Z2fnv\00", align 1
// OGCG-HOST: @1 = private unnamed_addr constant [2 x i8] c"a\00", align 1
// OGCG-HOST: @2 = private constant [14 x i8] c"sample fatbin\0A", section ".nv_fatbin", align 8
// OGCG-HOST: @__cuda_fatbin_wrapper = internal constant { i32, i32, ptr, ptr } { i32 1180844977, i32 1, ptr @2, ptr null }, section ".nvFatBinSegment", align 8
// OGCG-HOST: @__cuda_gpubin_handle = internal global ptr null, align 8
// OGCG-HOST: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @__cuda_module_ctor, ptr null }]

// OGCG-HOST: define internal void @__cuda_register_globals(ptr %[[#HANDLE:]]) {
// OGCG-HOST: entry:
// OGCG-HOST: %1 = call i32 @__cudaRegisterFunction(ptr %[[#HANDLE]],
// OGCG-HOST-SAME: ptr @_Z17__device_stub__fnv,
// OGCG-HOST-SAME: ptr @0,
// OGCG-HOST-SAME: ptr @0,
// OGCG-HOST-SAME: i32 -1,
// OGCG-HOST-SAME: ptr null,
// OGCG-HOST-SAME: ptr null,
// OGCG-HOST-SAME: ptr null,
// OGCG-HOST-SAME: ptr null,
// OGCG-HOST-SAME: ptr null)
// OGCG-HOST: call void @__cudaRegisterVar(ptr %[[#HANDLE]],
// OGCG-HOST-SAME: ptr @a,
// OGCG-HOST-SAME: ptr @1,
// OGCG-HOST-SAME: ptr @1,
// OGCG-HOST-SAME: i32 0,
// OGCG-HOST-SAME: i64 4,
// OGCG-HOST-SAME: i32 0,
// OGCG-HOST-SAME: i32 0)
// OGCG-HOST: ret void
// OGCG-HOST: }

// OGCG-HOST: define internal void @__cuda_module_ctor() {
// OGCG-HOST: entry:
// OGCG-HOST: %[[#WRAPADDR:]] = call ptr @__cudaRegisterFatBinary(ptr @__cuda_fatbin_wrapper)
// OGCG-HOST: store ptr %[[#WRAPADDR]], ptr @__cuda_gpubin_handle, align 8
// OGCG-HOST: call void @__cuda_register_globals(ptr %[[#WRAPADDR]])
// OGCG-HOST: call void @__cudaRegisterFatBinaryEnd(ptr %[[#WRAPADDR]])
// OGCG-HOST: %1 = call i32 @atexit(ptr @__cuda_module_dtor)
// OGCG-HOST: ret void
// OGCG-HOST: }

// OGCG-HOST: define internal void @__cuda_module_dtor() {
// OGCG-HOST: entry:
// OGCG-HOST: %[[#HANDLE:]] = load ptr, ptr @__cuda_gpubin_handle, align 8
// OGCG-HOST: call void @__cudaUnregisterFatBinary(ptr %[[#HANDLE]])
// OGCG-HOST: ret void
// OGCG-HOST: }
46 changes: 44 additions & 2 deletions clang/test/CIR/CodeGen/HIP/registration.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,14 @@
// CIR-HOST: cir.global_ctors = [#cir.global_ctor<"__hip_module_ctor", {{[0-9]+}}>]
// CIR-HOST: }

// LLVM-HOST: @.stra0 = private constant [2 x i8] c"a\00"
// LLVM-HOST: @.str_Z2fnv = private constant [7 x i8] c"_Z2fnv\00"
// LLVM-HOST: @__hip_fatbin_str = private constant [14 x i8] c"sample fatbin\0A", section ".hip_fatbin"
// LLVM-HOST: @__hip_fatbin_wrapper = internal constant {
// LLVM-HOST: i32 1212764230, i32 1, ptr @__hip_fatbin_str, ptr null
// LLVM-HOST: }, section ".hipFatBinSegment"
// LLVM-HOST: @_Z2fnv = constant ptr @_Z17__device_stub__fnv, align 8
// LLVM-HOST: @a = global i32 undef, align 4
// LLVM-HOST: @llvm.global_ctors = {{.*}}ptr @__hip_module_ctor

// CIR-HOST: cir.func internal private @__hip_module_dtor() {
Expand Down Expand Up @@ -66,6 +68,9 @@

__global__ void fn() {}


__device__ int a;

// CIR-HOST: cir.func internal private @__hip_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 @@ -80,6 +85,23 @@ __global__ void fn() {}
// CIR-HOST-SAME: %[[#DeviceFn]],
// CIR-HOST-SAME: %[[#MinusOne]],
// CIR-HOST-SAME: %[[#NULL]], %[[#NULL]], %[[#NULL]], %[[#NULL]], %[[#NULL]])
// CIR-HOST: %[[#GVARNAME:]] = cir.get_global @".stra0"
// CIR-HOST: %[[#GVARNAMEPTR:]] = cir.cast bitcast %[[#GVARNAME]]
// CIR-HOST: %[[#GVAR:]] = cir.get_global @a
// CIR-HOST: %[[#GVARPTR:]] = cir.cast bitcast %[[#GVAR]]
// CIR-HOST: %[[#ZERO:]] = cir.const #cir.int<0>
// CIR-HOST: %[[#FOUR:]] = cir.const #cir.int<4>
// CIR-HOST: %[[#ZERON:]] = cir.const #cir.int<0>
// CIR-HOST: %[[#ZERONN:]] = cir.const #cir.int<0>
// CIR-HOST: cir.call @__hipRegisterVar(%[[FatbinHandle]],
// CIR-HOST-SAME: %[[#GVARPTR]],
// CIR-HOST-SAME: %[[#GVARNAMEPTR]],
// CIR-HOST-SAME: %[[#GVARNAMEPTR]],
// CIR-HOST-SAME: %[[#ZERO]],
// CIR-HOST-SAME: %[[#FOUR:]],
// CIR-HOST-SAME: %[[#ZERON]],
// CIR-HOST-SAME: %[[#ZERONN]])
// CIR-HOST: cir.return loc(#loc)
// CIR-HOST: }

// LLVM-HOST: define internal void @__hip_register_globals(ptr %[[#LLVMFatbin:]]) {
Expand All @@ -90,6 +112,15 @@ __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 @__hipRegisterVar(
// LLVM-HOST-SAME: ptr %[[#LLVMFatbin]],
// LLVM-HOST-SAME: ptr @a,
// LLVM-HOST-SAME: ptr @.stra0,
// LLVM-HOST-SAME: ptr @.stra0,
// LLVM-HOST-SAME: i32 0,
// LLVM-HOST-SAME: i64 4,
// LLVM-HOST-SAME: i32 0,
// LLVM-HOST-SAME: i32 0)
// LLVM-HOST: }

// The content in const array should be the same as echoed above,
Expand Down Expand Up @@ -145,9 +176,11 @@ __global__ void fn() {}
// LLVM-HOST: ret void

// OGCG-HOST: @_Z2fnv = constant ptr @_Z17__device_stub__fnv, align 8
// OGCG-HOST: @a = internal global i32 undef, align 4
// OGCG-HOST: @0 = private unnamed_addr constant [7 x i8] c"_Z2fnv\00", align 1
// OGCG-HOST: @1 = private constant [14 x i8] c"sample fatbin\0A", section ".hip_fatbin", align 4096
// OGCG-HOST: @__hip_fatbin_wrapper = internal constant { i32, i32, ptr, ptr } { i32 1212764230, i32 1, ptr @1, ptr null }, section ".hipFatBinSegment", align 8
// OGCG-HOST: @1 = private unnamed_addr constant [2 x i8] c"a\00", align 1
// OGCG-HOST: @2 = private constant [14 x i8] c"sample fatbin\0A", section ".hip_fatbin", align 4096
// OGCG-HOST: @__hip_fatbin_wrapper = internal constant { i32, i32, ptr, ptr } { i32 1212764230, i32 1, ptr @2, ptr null }, section ".hipFatBinSegment", align 8
// OGCG-HOST: @__hip_gpubin_handle = internal global ptr null, align 8
// OGCG-HOST: @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @__hip_module_ctor, ptr null }]

Expand All @@ -160,6 +193,15 @@ __global__ void fn() {}
// OGCG-HOST-SAME: ptr @0,
// OGCG-HOST-SAME: i32 -1,
// OGCG-HOST-SAME: ptr null, ptr null, ptr null, ptr null, ptr null)
// OGCG-HOST: call void @__hipRegisterVar(
// OGCG-HOST-SAME: ptr %[[#LLVMFatbin]],
// OGCG-HOST-SAME: ptr @a,
// OGCG-HOST-SAME: ptr @1,
// OGCG-HOST-SAME: ptr @1,
// OGCG-HOST-SAME: i32 0,
// OGCG-HOST-SAME: i64 4,
// OGCG-HOST-SAME: i32 0, i32 0)
// OGCG-HOST: ret void
// OGCG-HOST: }

// OGCG-HOST: define internal void @__hip_module_ctor() {
Expand Down
Loading