From 29384cfd0ce24666dd3dd82008f33395020c751b Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Wed, 5 Aug 2026 10:40:00 -0700 Subject: [PATCH] Tell ptxas the largest block size a kernel can be launched with Without it ptxas has to assume a block could hold the maximum number of threads, which caps it at 65536/1024 = 64 registers per thread. The kernel in apps/cuda_mat_mul is launched with 16x2 threads and wants 70 registers, so it spilled: 20 bytes of spill stores and 16 of spill loads. Giving ptxas the bound removes the spills, and the app goes from 0.314 ms to 0.233 ms. This has to be a function attribute. The equivalent nvvm.annotations entry, which is how the kernel annotation just above is written, is upgraded to one of these when a module is read from a file, but that never happens to a module we built ourselves, and the backend only looks at the attribute. Co-Authored-By: Claude Opus 5 --- src/CodeGen_PTX_Dev.cpp | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 8dcfc0016aa0..dab0b532aa55 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -126,6 +126,30 @@ Type CodeGen_PTX_Dev::upgrade_type_for_storage(const Type &t) const { return CodeGen_LLVM::upgrade_type_for_storage(t); } +// The largest extent of each of the GPU thread loops, if they are all +// constant. A kernel may contain several thread loops in sequence, so take the +// largest of each. +class BlockSize : public IRVisitor { + using IRVisitor::visit; + + void visit(const For *op) override { + for (int i = 0; i < 3; i++) { + if (ends_with(op->name, gpu_thread_name(i))) { + if (auto e = as_const_int(simplify(op->extent()))) { + extent[i] = std::max(extent[i], (int)*e); + } else { + known = false; + } + } + } + IRVisitor::visit(op); + } + +public: + int extent[3] = {1, 1, 1}; + bool known = true; +}; + void CodeGen_PTX_Dev::add_kernel(Stmt stmt, const std::string &name, const std::vector &args) { @@ -204,6 +228,20 @@ void CodeGen_PTX_Dev::add_kernel(Stmt stmt, module->getOrInsertNamedMetadata("nvvm.annotations")->addOperand(md_node); + // Tell ptxas the most threads a block can have. Without this it assumes + // the maximum, and budgets registers for it. + BlockSize block_size; + stmt.accept(&block_size); + if (block_size.known) { + function->addFnAttr("nvvm.maxntid", + std::to_string(block_size.extent[0]) + "," + + std::to_string(block_size.extent[1]) + "," + + std::to_string(block_size.extent[2])); + debug(2) << "Kernel " << name << " has block size " + << block_size.extent[0] << "x" << block_size.extent[1] + << "x" << block_size.extent[2] << "\n"; + } + // Now verify the function is ok verifyFunction(*function);