[lang] Allow non portable dynamic shared memory size

Signed-off-by: Asher Mancinelli <amancinelli@nvidia.com>

Author: ashermancinelli

cext/cuda_loader.h

@@ -26,6 +26,7 @@
     X(cuCtxGetDevice, 2000) \
     X(cuCtxGetId, 12000) \
     X(cuDeviceGet, 2000) \
+    X(cuDeviceGetCount, 2000) \
     X(cuDeviceGetAttribute, 2000) \
     X(cuDevicePrimaryCtxRetain, 7000) \
     X(cuDriverGetVersion, 2020) \
@@ -34,6 +35,8 @@
     X(cuEventQuery, 2000) \
     X(cuEventRecord, 2000) \
     X(cuKernelGetFunction, 12000) \
+    X(cuKernelGetAttribute, 12000) \
+    X(cuKernelSetAttribute, 12000) \
     X(cuMemAlloc, 3020) \
     X(cuMemAllocHost, 3020) \
     X(cuMemFree, 3020) \

cext/tile_kernel.cpp

@@ -280,6 +280,66 @@ struct CudaKernel {
     CUkernel kernel;
 };
 
+static Status enable_maximum_dynamic_shared_memory(const DriverApi *driver,
+                                                   const CUkernel kernel,
+                                                   const char *func_name) {
+    int device_count;
+    CUresult res = driver->cuDeviceGetCount(&device_count);
+    if (res != CUDA_SUCCESS) {
+        return raise(PyExc_RuntimeError, "Failed to get device count: %s",
+                     get_cuda_error(driver, res));
+    }
+
+    for (int device_ordinal = 0; device_ordinal < device_count;
+         device_ordinal++) {
+        CUdevice device;
+        res = driver->cuDeviceGet(&device, device_ordinal);
+        if (res != CUDA_SUCCESS) {
+            return raise(PyExc_RuntimeError, "Failed to get device %d: %s",
+                         device_ordinal, get_cuda_error(driver, res));
+        }
+
+        int max_smem;
+        res = driver->cuDeviceGetAttribute(
+            &max_smem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
+            device);
+        if (res != CUDA_SUCCESS) {
+            return raise(PyExc_RuntimeError,
+                         "Failed to get maximum shared memory for device %d: %s",
+                         device_ordinal, get_cuda_error(driver, res));
+        }
+
+        int static_smem;
+        res = driver->cuKernelGetAttribute(
+            &static_smem, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, kernel, device);
+        if (res != CUDA_SUCCESS) {
+            return raise(PyExc_RuntimeError,
+                         "Failed to get static shared memory for kernel %s: %s",
+                         func_name, get_cuda_error(driver, res));
+        }
+
+        if (max_smem < static_smem) {
+            // If the user's program uses more static shared memory than the
+            // current device has available, then we cannot request enough
+            // shared memory. If the user has another device capable of running
+            // their program, they must run on that device and errors will be
+            // reported at launch time.
+            continue;
+        }
+        int largest_possible_dynamic_smem = max_smem - static_smem;
+        res = driver->cuKernelSetAttribute(
+            CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+            largest_possible_dynamic_smem, kernel, device);
+        if (res != CUDA_SUCCESS) {
+            return raise(PyExc_RuntimeError,
+                         "Failed to set dynamic shared memory for kernel %s: %s",
+                         func_name, get_cuda_error(driver, res));
+        }
+    }
+
+    return OK;
+}
+
 static Result<CudaKernel> load_cuda_kernel(const DriverApi* driver,
                                            const char* cubin_data,
                                            size_t cubin_size,
@@ -291,11 +351,12 @@ static Result<CudaKernel> load_cuda_kernel(const DriverApi* driver,
 
     CUkernel kernel;
     CUresult res = driver->cuLibraryGetKernel(&kernel, lib->get(), func_name);
-    if (res == CUDA_SUCCESS)
-        return CudaKernel{std::move(*lib), kernel};
+    if (res != CUDA_SUCCESS) {
+        return raise(PyExc_RuntimeError, "Failed to get kernel %s from library: %s",
+                     func_name, get_cuda_error(driver, res));
+    }
 
-    return raise(PyExc_RuntimeError, "Failed to get kernel %s from library: %s",
-                 func_name, get_cuda_error(driver, res));
+    return CudaKernel{std::move(*lib), kernel};
 }
 
 
@@ -2156,6 +2217,12 @@ static Result<TileKernel> compile(const DriverApi* driver,
     Result<CudaKernel> cukernel = load_cuda_kernel(driver, cubin_data, cubin_size, cufunc_name);
     if (!cukernel.is_ok()) return ErrorRaised;
 
+    if (py_dyn_smem_size_prog != Py_None) {
+        Status status = enable_maximum_dynamic_shared_memory(
+            driver, cukernel->kernel, cufunc_name);
+        if (!status) return ErrorRaised;
+    }
+
     Result<HostProgram> dyn_smem_size_prog = host_program_parse(py_dyn_smem_size_prog, 1);
     if (!dyn_smem_size_prog.is_ok()) return ErrorRaised;
 

experimental/cuda-lang/test/test_dynamic_shared_mem.py

@@ -172,10 +172,8 @@ def kern(n):
 def test_dynamic_shared_alignment_runtime_round_up_launch():
     @cl.kernel
     def kern(x, n):
-        smem = cl.shared_array(
-            shape=(n,), dtype=cl.uint8, dynamic=True, alignment=128
-        )
-        values = cl.shared_array(shape=(1,), dtype=cl.int32, dynamic=True)
+        smem = cl.shared_array(n, cl.uint8, dynamic=True, alignment=128)
+        values = cl.shared_array(1, cl.int32, dynamic=True)
         smem[0] = cl.uint8(0)
         values[0] = 42
         x[0] = values[0] + cl.int32(smem[0])
@@ -188,6 +186,48 @@ def kern(x, n):
     assert spy.get_dynamic_smem_size() == 128 + 4
 
 
+def test_dynamic_shared_memory_over_device_limit():
+    properties = torch.cuda.get_device_properties(torch.cuda.current_device())
+    max_smem_bytes = properties.shared_memory_per_block_optin
+
+    @cl.kernel
+    def kern():
+        smem = cl.shared_array(max_smem_bytes + 1, cl.int8, dynamic=True)
+        if cl.thread_index(0) == 0:
+            smem[max_smem_bytes] = cl.int8(0)
+
+    with pytest.raises(RuntimeError, match="Failed to launch"):
+        cl.launch(torch.cuda.current_stream(), (1,), (32,), kern, ())
+
+
+def test_max_shared_memory_with_static_allocation():
+    properties = torch.cuda.get_device_properties(torch.cuda.current_device())
+    max_smem_bytes = properties.shared_memory_per_block_optin
+    static_smem_bytes = 1024
+    dynamic_smem_bytes = max_smem_bytes - static_smem_bytes
+
+    @cl.kernel
+    def kern(output):
+        static_smem = cl.shared_array(
+            static_smem_bytes, cl.int8, alignment=1024
+        )
+        dynamic_smem = cl.shared_array(
+            dynamic_smem_bytes, cl.int8, dynamic=True
+        )
+        if cl.thread_index(0) == 0:
+            static_smem[static_smem_bytes - 1] = cl.int8(3)
+            dynamic_smem[dynamic_smem_bytes - 1] = cl.int8(4)
+            output[0] = (
+                static_smem[static_smem_bytes - 1]
+                + dynamic_smem[dynamic_smem_bytes - 1]
+            )
+
+    output = torch.zeros(1, dtype=torch.int8, device="cuda")
+    cl.launch(torch.cuda.current_stream(), (1,), (32,), kern, (output,))
+
+    assert output.item() == 7
+
+
 def test_dynamic_1d_array_and_static_1d_array():
     @cl.kernel
     def kern(x, n):