Add launch bounds

NVIDIA · Jan 3, 2025 · c2c10b7 · c2c10b7
1 parent a82c3fe
commit c2c10b7
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/csrc/codegen.cpp b/csrc/codegen.cpp
@@ -274,6 +274,11 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
   // Generates the kernel function declaration
   void genDeclaration(const std::string& kernel_name) {
     code_ << "__global__ void ";
+    if (kernel_->hasManaged("warp_specialized_num_registers")) {
+      constexpr int64_t threads_per_cta = 384;
+      code_ << "__launch_bounds__(/*MAX_THREADS_PER_BLOCK=*/" << threads_per_cta
+            << ") ";
+    }
     if (kernel_->hasManaged("cluster_dims")) {
       auto cluster_dims =
           kernel_->getManaged<std::tuple<int64_t, int64_t, int64_t>>(