https://github.com/RiverDave created https://github.com/llvm/llvm-project/pull/231390
For some very odd reason, launching a kernel defined in another TU would cause us to embed two binaries. This patch makes sure we only register stubs with a body. >From 8dcb2526f4ce375aa27e20d5220154b5149511a1 Mon Sep 17 00:00:00 2001 From: David Rivera <[email protected]> Date: Tue, 13 Oct 2026 03:05:11 -0400 Subject: [PATCH] [CIR][CUDA] Register only defined stubs --- clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp | 11 +-- .../Dialect/Transforms/LoweringPrepare.cpp | 3 +- .../CIR/CodeGenCUDA/kernel-declaration.cu | 72 +++++++++++++++++++ 3 files changed, 81 insertions(+), 5 deletions(-) create mode 100644 clang/test/CIR/CodeGenCUDA/kernel-declaration.cu diff --git a/clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp b/clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp index 79153ca788756..1c6d049163dad 100644 --- a/clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp +++ b/clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp @@ -524,15 +524,18 @@ void CIRGenNVCUDARuntime::handleGlobalReplace(cir::GlobalOp oldGV, } /// Whether this translation unit has anything for the CUDA runtime to register. -/// These are the same two attributes LoweringPrepare collects to decide whether -/// to build a module ctor, so both sides answer the question from one source. +/// Like LoweringPrepare, count only kernel stub definitions and device variables +/// carrying registration metadata, not kernel declarations or launch-site calls. static bool hasEntitiesToRegister(mlir::ModuleOp module) { // A walk, not a scan of the module body, so this stays in agreement with the // recursive walk LoweringPrepare collects them with. return module ->walk([](mlir::Operation *op) { - if (op->hasAttr(cir::CUDAKernelNameAttr::getMnemonic()) || - op->hasAttr(cir::CUDAVarRegistrationInfoAttr::getMnemonic())) + if (auto fn = mlir::dyn_cast<cir::FuncOp>(op); + fn && !fn.isDeclaration() && + fn->hasAttr(cir::CUDAKernelNameAttr::getMnemonic())) + return mlir::WalkResult::interrupt(); + if (op->hasAttr(cir::CUDAVarRegistrationInfoAttr::getMnemonic())) return mlir::WalkResult::interrupt(); return mlir::WalkResult::advance(); }) diff --git a/clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp b/clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp index 23fb81e2f706d..ec8f36d93811f 100644 --- a/clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp +++ b/clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp @@ -2467,7 +2467,8 @@ void LoweringPreparePass::runOnOp(mlir::Operation *op) { globalDtorList.emplace_back(fnOp.getName(), globalDtor.value()); if (mlir::Attribute attr = - fnOp->getAttr(cir::CUDAKernelNameAttr::getMnemonic())) { + fnOp->getAttr(cir::CUDAKernelNameAttr::getMnemonic()); + attr && !fnOp.isDeclaration()) { auto kernelNameAttr = dyn_cast<CUDAKernelNameAttr>(attr); llvm::StringRef kernelName = kernelNameAttr.getKernelName(); cudaKernelMap[kernelName] = fnOp; diff --git a/clang/test/CIR/CodeGenCUDA/kernel-declaration.cu b/clang/test/CIR/CodeGenCUDA/kernel-declaration.cu new file mode 100644 index 0000000000000..46bfc85f9af66 --- /dev/null +++ b/clang/test/CIR/CodeGenCUDA/kernel-declaration.cu @@ -0,0 +1,72 @@ +// A launch of a kernel defined in another TU must not register that kernel or +// embed this TU's fat binary. Match classic codegen for both declaration-only +// TUs and TUs that also define a different kernel. + +// RUN: echo -n "GPU binary would be here." > %t.fatbin + +// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-cir %s -x cuda \ +// RUN: -target-sdk-version=12.3 -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefix=CIR --implicit-check-not=fatbin \ +// RUN: --implicit-check-not=Register --implicit-check-not=module_ctor + +// With no entities to register, the binary must not even be read. +// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-cir %s -x cuda \ +// RUN: -target-sdk-version=12.3 -fcuda-include-gpubinary %t.nonexistent -o - \ +// RUN: | FileCheck %s --check-prefix=CIR --implicit-check-not=fatbin \ +// RUN: --implicit-check-not=Register --implicit-check-not=module_ctor + +// RUN: %clang_cc1 -triple x86_64-linux-gnu -fclangir -emit-llvm %s -x cuda \ +// RUN: -target-sdk-version=12.3 -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefix=LAUNCH --implicit-check-not=fatbin \ +// RUN: --implicit-check-not=Register --implicit-check-not=module_ctor \ +// RUN: --implicit-check-not=llvm.global_ctors +// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-llvm %s -x cuda \ +// RUN: -target-sdk-version=12.3 -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefix=LAUNCH --implicit-check-not=fatbin \ +// RUN: --implicit-check-not=Register --implicit-check-not=module_ctor \ +// RUN: --implicit-check-not=llvm.global_ctors + +// RUN: %clang_cc1 -triple x86_64-linux-gnu -fclangir -emit-llvm %s -x hip \ +// RUN: -fhip-new-launch-api -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefix=LAUNCH --implicit-check-not=fatbin \ +// RUN: --implicit-check-not=Register --implicit-check-not=module_ctor \ +// RUN: --implicit-check-not=llvm.global_ctors + +// RUN: %clang_cc1 -triple x86_64-linux-gnu -fclangir -emit-llvm %s -x cuda \ +// RUN: -DDEFINE_KERNEL -target-sdk-version=12.3 -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefixes=LAUNCH,MIXED \ +// RUN: --implicit-check-not='call{{.*}} @__{{cuda|hip}}RegisterFunction({{.*}}ptr @{{[^,]*}}k0' +// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-llvm %s -x cuda \ +// RUN: -DDEFINE_KERNEL -target-sdk-version=12.3 -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefixes=LAUNCH,MIXED \ +// RUN: --implicit-check-not='call{{.*}} @__{{cuda|hip}}RegisterFunction({{.*}}ptr @{{[^,]*}}k0' +// RUN: %clang_cc1 -triple x86_64-linux-gnu -fclangir -emit-llvm %s -x hip \ +// RUN: -DDEFINE_KERNEL -fhip-new-launch-api -fcuda-include-gpubinary %t.fatbin -o - \ +// RUN: | FileCheck %s --check-prefixes=LAUNCH,MIXED \ +// RUN: --implicit-check-not='call{{.*}} @__{{cuda|hip}}RegisterFunction({{.*}}ptr @{{[^,]*}}k0' + +#include "Inputs/cuda.h" + +__global__ void k0(int); +#ifdef DEFINE_KERNEL +__global__ void k1(int); +#endif + +void launch() { + k0<<<1, 1>>>(42); +#ifdef DEFINE_KERNEL + k1<<<1, 1>>>(42); +#endif +} + +// Define after the launch to exercise upgrading an existing stub declaration. +#ifdef DEFINE_KERNEL +__global__ void k1(int) {} +#endif + +// CIR: cir.call @_Z17__device_stub__k0i( +// LAUNCH-DAG: call void @_Z17__device_stub__k0i(i32{{.*}} 42) +// LAUNCH-DAG: declare void @_Z17__device_stub__k0i( +// MIXED-DAG: define{{.*}} void @_Z17__device_stub__k1i( +// MIXED-DAG: call{{.*}} @__{{cuda|hip}}RegisterFunction({{.*}}ptr @{{[^,]*}}k1{{[^,]*}}, +// MIXED-DAG: define internal void @__{{cuda|hip}}_module_ctor() _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
