[cuda] Add argument "gpu_max_reg" to ti.init

taichi/backends/cpu/jit_cpu.cpp

@@ -126,7 +126,8 @@ class JITSessionCPU : public JITSession {
     return DL;
   }
 
-  JITModule *add_module(std::unique_ptr<llvm::Module> M) override {
+  JITModule *add_module(std::unique_ptr<llvm::Module> M, int max_reg) override {
+    TI_ASSERT(max_reg == 0);  // No need to specify max_reg on CPUs
     TI_ASSERT(M);
     global_optimize_module_cpu(M);
     std::lock_guard<std::mutex> _(mut);

taichi/backends/cuda/codegen_cuda.cpp

@@ -41,7 +41,8 @@ class CodeGenLLVMCUDA : public CodeGenLLVM {
     }
 
     auto jit = kernel->program.llvm_context_device->jit.get();
-    auto cuda_module = jit->add_module(std::move(module));
+    auto cuda_module =
+        jit->add_module(std::move(module), kernel->program.config.gpu_max_reg);
 
     return [offloaded_local, cuda_module,
             kernel = this->kernel](Context &context) {

taichi/backends/cuda/cuda_driver.h

@@ -33,6 +33,7 @@ constexpr uint32 CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT = 16;
 constexpr uint32 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR = 75;
 constexpr uint32 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR = 76;
 constexpr uint32 CUDA_ERROR_ASSERT = 710;
+constexpr uint32 CU_JIT_MAX_REGISTERS = 0;
 
 std::string get_cuda_error_message(uint32 err);
 

taichi/backends/cuda/jit_cuda.cpp

@@ -86,7 +86,8 @@ class JITSessionCUDA : public JITSession {
       : data_layout(data_layout) {
   }
 
-  virtual JITModule *add_module(std::unique_ptr<llvm::Module> M) override {
+  virtual JITModule *add_module(std::unique_ptr<llvm::Module> M,
+                                int max_reg) override {
     auto ptx = compile_module_to_ptx(M);
     if (get_current_program().config.print_kernel_nvptx) {
       static FileSequenceWriter writer("taichi_kernel_nvptx_{:04d}.ptx",
@@ -103,8 +104,23 @@ class JITSessionCUDA : public JITSession {
     TI_TRACE("Loading module...");
     [[maybe_unused]] auto &&_ =
         std::move(CUDAContext::get_instance().get_lock_guard());
-    CUDADriver::get_instance().module_load_data_ex(&cuda_module, ptx.c_str(), 0,
-                                                   nullptr, nullptr);
+
+    constexpr int max_num_options = 8;
+    int num_options = 0;
+    uint32 options[max_num_options];
+    void *option_values[max_num_options];
+
+    // Insert options
+    if (max_reg != 0) {
+      options[num_options] = CU_JIT_MAX_REGISTERS;
+      option_values[num_options] = &max_reg;
+      num_options++;
+    }
+
+    TI_ASSERT(num_options <= max_num_options);
+
+    CUDADriver::get_instance().module_load_data_ex(
+        &cuda_module, ptx.c_str(), num_options, options, option_values);
     TI_TRACE("CUDA module load time : {}ms", (Time::get_time() - t) * 1000);
     // cudaModules.push_back(cudaModule);
     modules.push_back(std::make_unique<JITModuleCUDA>(cuda_module));

taichi/jit/jit_session.h

@@ -19,7 +19,8 @@ class JITSession {
   JITSession() {
   }
 
-  virtual JITModule *add_module(std::unique_ptr<llvm::Module> M) = 0;
+  virtual JITModule *add_module(std::unique_ptr<llvm::Module> M,
+                                int max_reg = 0) = 0;
 
   // virtual void remove_module(JITModule *module) = 0;
 

taichi/program/compile_config.cpp

@@ -30,6 +30,7 @@ CompileConfig::CompileConfig() {
   kernel_profiler = false;
   default_cpu_block_dim = 32;
   default_gpu_block_dim = 128;
+  gpu_max_reg = 0;  // 0 means using the default value from the CUDA driver.
   verbose = true;
   fast_math = true;
   async_mode = false;

taichi/program/compile_config.h

@@ -39,6 +39,7 @@ struct CompileConfig {
   std::string extra_flags;
   int default_cpu_block_dim;
   int default_gpu_block_dim;
+  int gpu_max_reg;
   int ad_stack_size;
 
   int saturating_grid_dim;

taichi/python/export_lang.cpp

@@ -137,6 +137,7 @@ void export_lang(py::module &m) {
                      &CompileConfig::default_cpu_block_dim)
       .def_readwrite("default_gpu_block_dim",
                      &CompileConfig::default_gpu_block_dim)
+      .def_readwrite("gpu_max_reg", &CompileConfig::gpu_max_reg)
       .def_readwrite("saturating_grid_dim", &CompileConfig::saturating_grid_dim)
       .def_readwrite("max_block_dim", &CompileConfig::max_block_dim)
       .def_readwrite("cpu_max_num_threads", &CompileConfig::cpu_max_num_threads)