[clang] [CIR][CUDA] Register only defined stubs (PR #231390)

[via cfe-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-clang

Author: David Rivera (RiverDave)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/231390.diff


3 Files Affected:

- (modified) clang/lib/CIR/CodeGen/CIRGenCUDANV.cpp (+7-4) 
- (modified) clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp (+2-1) 
- (added) clang/test/CIR/CodeGenCUDA/kernel-declaration.cu (+72) 


``````````diff
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()

``````````

</details>


https://github.com/llvm/llvm-project/pull/231390