Repository: taichi-dev/taichi (language: python)

Issue:
Adding support for larger SharedArray
**Concisely describe the proposed feature**

SharedArray helps store the intermediate results so we don't need to write back to global memory, which really increases the performance. 

However, as quoted from CUDA doc: [link](https://docs.nvidia.com/cuda/cuda-c-programming-guide/#shared-memory-8-x)

> Devices of compute capability 8.0 allow a single thread block to address up to 163 KB of shared memory, while devices of compute capabilities 8.6 and 8.9 allow up to 99 KB of shared memory. Kernels relying on shared memory allocations over 48 KB per block are architecture-specific, and must use dynamic shared memory rather than statically sized shared memory arrays. These kernels require an explicit opt-in by using cudaFuncSetAttribute() to set the cudaFuncAttributeMaxDynamicSharedMemorySize; see [Shared Memory](https://docs.nvidia.com/cuda/cuda-c-programming-guide/topics/compute-capabilities.html#shared-memory-7.x) for the Volta architecture.

Right now, we can only allocate 48KB of shared memory (SharedArray), even with new high-end GPUs which already have 99KB or more shared memory. So, is this possible for Taichi to support dynamic shared memory in CUDA that allows us to allocate more shared memory?




Discussion / hints:
Thank you for bring this up!

We need to generate code for different device capabilities. We currently pin the codegen at sm75 which only supports up to 64KB. It's a little complex at the time being for Taichi. Let me think about this carefully.
Thanks!  This is a great feature for Taichi.



Fix patch:
diff --git a/taichi/codegen/cuda/codegen_cuda.cpp b/taichi/codegen/cuda/codegen_cuda.cpp
index 2f6708b0cb6dd..c61ea0c4748c3 100644
--- a/taichi/codegen/cuda/codegen_cuda.cpp
+++ b/taichi/codegen/cuda/codegen_cuda.cpp
@@ -39,6 +39,7 @@ static bool is_half2(DataType dt) {
 class TaskCodeGenCUDA : public TaskCodeGenLLVM {
  public:
   using IRVisitor::visit;
+  size_t dynamic_shared_array_bytes{0};
 
   explicit TaskCodeGenCUDA(const CompileConfig &config,
                            TaichiLLVMContext &tlctx,
@@ -161,6 +162,44 @@ class TaskCodeGenCUDA : public TaskCodeGenLLVM {
     llvm_val[stmt] = create_print(formats, types, values);
   }
 
+  void visit(AllocaStmt *stmt) override {
+    // Override shared memory codegen logic for large shared memory
+    if (stmt->ret_type->is<TensorType>() && stmt->is_shared) {
+      auto tensor_type = stmt->ret_type->cast<TensorType>();
+      size_t shared_array_bytes =
+          tensor_type->get_num_elements() *
+          data_type_size(tensor_type->get_element_type());
+      if (shared_array_bytes > cuda_dynamic_shared_array_threshold_bytes) {
+        if (dynamic_shared_array_bytes > 0) {
+          /* Current version only allows one dynamic shared array allocation,
+           * otherwise the results could be wrong.
+           * However, we should be able to collect multiple user allocations
+           * and transparently apply a proper offset.
+           *
+           * TODO: remove the limits.
+           */
+          TI_ERROR(
+              "Only one single large shared array instance is allowed in "
+              "current version.")
+        }
+        // Clear tensor shape for dynamic shared memory.
+        tensor_type->set_shape(std::vector<int>({0}));
+        dynamic_shared_array_bytes += shared_array_bytes;
+      }
+
+      auto type = tlctx->get_data_type(tensor_type);
+      auto base = new llvm::GlobalVariable(
+          *module, type, false, llvm::GlobalValue::ExternalLinkage, nullptr,
+          fmt::format("shared_array_{}", stmt->id), nullptr,
+          llvm::GlobalVariable::NotThreadLocal, 3 /*addrspace=shared*/);
+      base->setAlignment(llvm::MaybeAlign(8));
+      auto ptr_type = llvm::PointerType::get(type, 0);
+      llvm_val[stmt] = builder->CreatePointerCast(base, ptr_type);
+    } else {
+      TaskCodeGenLLVM::visit(stmt);
+    }
+  }
+
   void emit_extra_unary(UnaryOpStmt *stmt) override {
     // functions from libdevice
     auto input = llvm_val[stmt->operand];
@@ -595,6 +634,7 @@ class TaskCodeGenCUDA : public TaskCodeGenLLVM {
         current_task->grid_dim = num_SMs * query_max_block_per_sm;
       }
       current_task->block_dim = stmt->block_dim;
+      current_task->dynamic_shared_array_bytes = dynamic_shared_array_bytes;
       TI_ASSERT(current_task->grid_dim != 0);
       TI_ASSERT(current_task->block_dim != 0);
       offloaded_tasks.push_back(*current_task);
@@ -797,7 +837,8 @@ FunctionType CUDAModuleToFunctionConverter::convert(
     for (auto task : offloaded_tasks) {
       TI_TRACE("Launching kernel {}<<<{}, {}>>>", task.name, task.grid_dim,
                task.block_dim);
-      cuda_module->launch(task.name, task.grid_dim, task.block_dim, 0,
+      cuda_module->launch(task.name, task.grid_dim, task.block_dim,
+                          task.dynamic_shared_array_bytes,
                           {&context.get_context()}, {});
     }
     if (context.arg_buffer_size > 0) {
diff --git a/taichi/codegen/llvm/llvm_compiled_data.h b/taichi/codegen/llvm/llvm_compiled_data.h
index 72f6dd877618c..4f27eaa632acf 100644
--- a/taichi/codegen/llvm/llvm_compiled_data.h
+++ b/taichi/codegen/llvm/llvm_compiled_data.h
@@ -13,12 +13,17 @@ class OffloadedTask {
   std::string name;
   int block_dim{0};
   int grid_dim{0};
+  int dynamic_shared_array_bytes{0};
 
   explicit OffloadedTask(const std::string &name = "",
                          int block_dim = 0,
-                         int grid_dim = 0)
-      : name(name), block_dim(block_dim), grid_dim(grid_dim){};
-  TI_IO_DEF(name, block_dim, grid_dim);
+                         int grid_dim = 0,
+                         int dynamic_shared_array_bytes = 0)
+      : name(name),
+        block_dim(block_dim),
+        grid_dim(grid_dim),
+        dynamic_shared_array_bytes(dynamic_shared_array_bytes){};
+  TI_IO_DEF(name, block_dim, grid_dim, dynamic_shared_array_bytes);
 };
 
 struct LLVMCompiledTask {
diff --git a/taichi/inc/constants.h b/taichi/inc/constants.h
index 275d892336daa..c7162e975cc10 100644
--- a/taichi/inc/constants.h
+++ b/taichi/inc/constants.h
@@ -27,6 +27,14 @@ constexpr std::size_t taichi_result_buffer_runtime_query_id = 31;
 
 constexpr int taichi_listgen_max_element_size = 1024;
 
+// By default, CUDA could allocate up to 48KB static shared arrays.
+// It requires dynamic shared memory to allocate a larger array.
+// Therefore, when one shared array request for size greater than 48KB,
+// we switch it to dynamic allocation.
+// In current version, only one dynamic instance is allowed.
+// TODO: remove the limit.
+constexpr std::size_t cuda_dynamic_shared_array_threshold_bytes = 49152;
+
 // use for auto mesh_local to determine shared-mem size per block (in bytes)
 // TODO: get this at runtime
 constexpr std::size_t default_shared_mem_size = 65536;
diff --git a/taichi/ir/type.h b/taichi/ir/type.h
index 267822c631576..76458a77540d4 100644
--- a/taichi/ir/type.h
+++ b/taichi/ir/type.h
@@ -230,6 +230,10 @@ class TI_DLL_EXPORT TensorType : public Type {
     return shape_;
   }
 
+  void set_shape(const std::vector<int> &shape) {
+    shape_ = shape;
+  }
+
   Type *get_compute_type() override {
     return this;
   }
diff --git a/taichi/rhi/cuda/cuda_context.cpp b/taichi/rhi/cuda/cuda_context.cpp
index 9d3415b5e27b0..120a1debd3ebb 100644
--- a/taichi/rhi/cuda/cuda_context.cpp
+++ b/taichi/rhi/cuda/cuda_context.cpp
@@ -57,13 +57,14 @@ CUDAContext::CUDAContext()
 
   compute_capability_ = cc_major * 10 + cc_minor;
 
-  if (compute_capability_ > 75) {
-    // The NVPTX backend of LLVM 10.0.0 does not seem to support
-    // compute_capability > 75 yet. See
-    // llvm-10.0.0.src/build/lib/Target/NVPTX/NVPTXGenSubtargetInfo.inc
-    compute_capability_ = 75;
+  if (compute_capability_ > 86) {
+    compute_capability_ = 86;
   }
 
+  driver_.device_get_attribute(
+      &max_shared_memory_bytes_,
+      CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device_);
+
   mcpu_ = fmt::format("sm_{}", compute_capability_);
 
   TI_TRACE("Emitting CUDA code for {}", mcpu_);
@@ -130,6 +131,17 @@ void CUDAContext::launch(void *func,
 
   if (grid_dim > 0) {
     std::lock_guard<std::mutex> _(lock_);
+    if (dynamic_shared_mem_bytes > 0) {
+      if (dynamic_shared_mem_bytes > max_shared_memory_bytes_) {
+        TI_ERROR(
+            "Requested dynamic shared memory size of {} bytes, but the device "
+            "supports max capacity of {} bytes.",
+            dynamic_shared_mem_bytes, max_shared_memory_bytes_);
+      }
+      driver_.kernel_set_attribute(
+          func, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+          dynamic_shared_mem_bytes);
+    }
     driver_.launch_kernel(func, grid_dim, 1, 1, block_dim, 1, 1,
                           dynamic_shared_mem_bytes, nullptr,
                           arg_pointers.data(), nullptr);
diff --git a/taichi/rhi/cuda/cuda_context.h b/taichi/rhi/cuda/cuda_context.h
index 313981210ab77..e912cca7aa0c1 100644
--- a/taichi/rhi/cuda/cuda_context.h
+++ b/taichi/rhi/cuda/cuda_context.h
@@ -26,6 +26,7 @@ class CUDAContext {
   std::mutex lock_;
   KernelProfilerBase *profiler_;
   CUDADriver &driver_;
+  int max_shared_memory_bytes_;
   bool debug_;
   bool supports_mem_pool_;
 
diff --git a/taichi/rhi/cuda/cuda_driver.h b/taichi/rhi/cuda/cuda_driver.h
index 809128a936419..07f0b6d1ad952 100644
--- a/taichi/rhi/cuda/cuda_driver.h
+++ b/taichi/rhi/cuda/cuda_driver.h
@@ -31,6 +31,7 @@ constexpr uint32 CU_STREAM_NON_BLOCKING = 0x1;
 constexpr uint32 CU_MEM_ATTACH_GLOBAL = 0x1;
 constexpr uint32 CU_MEM_ADVISE_SET_PREFERRED_LOCATION = 3;
 constexpr uint32 CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X = 2;
+constexpr uint32 CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 97;
 constexpr uint32 CU_DEVICE_ATTRIBUTE_MAX_BLOCKS_PER_MULTIPROCESSOR = 106;
 constexpr uint32 CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT = 16;
 constexpr uint32 CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR = 75;
diff --git a/taichi/rhi/cuda/cuda_driver_functions.inc.h b/taichi/rhi/cuda/cuda_driver_functions.inc.h
index 3c0999f3b59d1..25b3c7958eda1 100644
--- a/taichi/rhi/cuda/cuda_driver_functions.inc.h
+++ b/taichi/rhi/cuda/cuda_driver_functions.inc.h
@@ -47,6 +47,8 @@ PER_CUDA_FUNCTION(launch_kernel, cuLaunchKernel, void *, uint32, uint32, uint32,
                   uint32, uint32, uint32, uint32, void *, void **, void **);
 PER_CUDA_FUNCTION(kernel_get_attribute, cuFuncGetAttribute, int *, uint32, void *);
 PER_CUDA_FUNCTION(kernel_get_occupancy, cuOccupancyMaxActiveBlocksPerMultiprocessor, int *, void *, int, size_t);
+PER_CUDA_FUNCTION(kernel_set_attribute, cuFuncSetAttribute, void *, CUfunction_attribute_enum, int);
+
 
 // Stream management
 PER_CUDA_FUNCTION(stream_synchronize, cuStreamSynchronize, void *);


Test patch:
diff --git a/tests/python/test_shared_array.py b/tests/python/test_shared_array.py
index f4133a181a993..665b188e7e29b 100644
--- a/tests/python/test_shared_array.py
+++ b/tests/python/test_shared_array.py
@@ -4,8 +4,11 @@
 from tests import test_utils
 
 
-@test_utils.test(arch=[ti.cuda, ti.vulkan, ti.amdgpu])
-def test_shared_array_nested_loop():
+@test_utils.test(arch=[ti.cuda])
+def test_large_shared_array():
+    # Skip the GPUs prior to Volta which doesn't have dynamical shared memory.
+    if ti.lang.impl.get_cuda_compute_capability() < 70:
+        return
     block_dim = 128
     nBlocks = 64
     N = nBlocks * block_dim
@@ -31,7 +34,7 @@ def calc_shared_array(v: ti.types.ndarray(ndim=1),
         ti.loop_config(block_dim=block_dim)
         for i in range(nBlocks * block_dim):
             tid = i % block_dim
-            pad = ti.simt.block.SharedArray((block_dim, ), ti.f32)
+            pad = ti.simt.block.SharedArray((65536 // 4, ), ti.f32)
             acc = 0.0
             v_val = v[i]
             for k in range(nBlocks):
@@ -47,6 +50,58 @@ def calc_shared_array(v: ti.types.ndarray(ndim=1),
     assert np.allclose(reference, a_arr)
 
 
+@test_utils.test(arch=[ti.cuda, ti.vulkan, ti.amdgpu])
+def test_multiple_shared_array():
+    block_dim = 128
+    nBlocks = 64
+    N = nBlocks * block_dim * 4
+    v_arr = np.random.randn(N).astype(np.float32)
+    d_arr = np.random.randn(N).astype(np.float32)
+    a_arr = np.zeros(N).astype(np.float32)
+    reference = np.zeros(N).astype(np.float32)
+
+    @ti.kernel
+    def calc(v: ti.types.ndarray(ndim=1), d: ti.types.ndarray(ndim=1),
+             a: ti.types.ndarray(ndim=1)):
+        for i in range(N):
+            acc = 0.0
+            v_val = v[i]
+            for j in range(N):
+                acc += v_val * d[j]
+            a[i] = acc
+
+    @ti.kernel
+    def calc_shared_array(v: ti.types.ndarray(ndim=1),
+                          d: ti.types.ndarray(ndim=1),
+                          a: ti.types.ndarray(ndim=1)):
+        ti.loop_config(block_dim=block_dim)
+        for i in range(nBlocks * block_dim * 4):
+            tid = i % block_dim
+            pad0 = ti.simt.block.SharedArray((block_dim, ), ti.f32)
+            pad1 = ti.simt.block.SharedArray((block_dim, ), ti.f32)
+            pad2 = ti.simt.block.SharedArray((block_dim, ), ti.f32)
+            pad3 = ti.simt.block.SharedArray((block_dim, ), ti.f32)
+            acc = 0.0
+            v_val = v[i]
+            for k in range(nBlocks):
+                pad0[tid] = d[k * block_dim * 4 + tid]
+                pad1[tid] = d[k * block_dim * 4 + block_dim + tid]
+                pad2[tid] = d[k * block_dim * 4 + 2 * block_dim + tid]
+                pad3[tid] = d[k * block_dim * 4 + 3 * block_dim + tid]
+                ti.simt.block.sync()
+                for j in range(block_dim):
+                    acc += v_val * pad0[j]
+                    acc += v_val * pad1[j]
+                    acc += v_val * pad2[j]
+                    acc += v_val * pad3[j]
+                ti.simt.block.sync()
+            a[i] = acc
+
+    calc(v_arr, d_arr, reference)
+    calc_shared_array(v_arr, d_arr, a_arr)
+    assert np.allclose(reference, a_arr, rtol=1e-4)
+
+
 @test_utils.test(arch=[ti.cuda, ti.vulkan, ti.amdgpu])
 def test_shared_array_atomics():
     N = 256