Add dynamic shared memory allocation

Author: michael-kenzel

platforms/artic/intrinsics_thorin.impala

@@ -12,10 +12,10 @@
 #[import(cc = "thorin")] fn cmpxchg_weak[T](_addr: &mut T, _cmp: T, _new: T, _success_order: u32, _failure_order: u32, _scope: &[u8]) -> (T, bool); // only for integer data types
 #[import(cc = "thorin")] fn fence(_order: u32, _scope: &[u8]) -> ();
 #[import(cc = "thorin")] fn pe_info[T](_src: &[u8], _val: T) -> ();
-#[import(cc = "thorin")] fn cuda(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _body: fn() -> ()) -> ();
-#[import(cc = "thorin")] fn nvvm(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _body: fn() -> ()) -> ();
-#[import(cc = "thorin")] fn opencl(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _body: fn() -> ()) -> ();
-#[import(cc = "thorin")] fn amdgpu(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _body: fn() -> ()) -> ();
+#[import(cc = "thorin", name = "cuda")] fn cuda_with_shared(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _smem: i32, _body: fn() -> ()) -> ();
+#[import(cc = "thorin", name = "nvvm")] fn nvvm_with_shared(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _smem: i32, _body: fn() -> ()) -> ();
+#[import(cc = "thorin", name = "opencl")] fn opencl_with_shared(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _smem: i32, _body: fn() -> ()) -> ();
+#[import(cc = "thorin", name = "amdgpu")] fn amdgpu_with_shared(_dev: i32, _grid: (i32, i32, i32), _block: (i32, i32, i32), _smem: i32, _body: fn() -> ()) -> ();
 #[import(cc = "thorin")] fn reserve_shared[T](_size: i32) -> &mut addrspace(3)[T];
 #[import(cc = "thorin")] fn hls(_dev: i32, _body: fn() -> ()) -> ();
 #[import(cc = "thorin", name = "pipeline")] fn thorin_pipeline(_initiation_interval: i32, _lower: i32, _upper: i32, _body: fn(i32) -> ()) -> (); // only for HLS/OpenCL backend
@@ -35,6 +35,11 @@
 #[import(cc = "thorin", name = "cmpxchg_weak")] fn cmpxchg_weak_p1[T](_addr: &mut addrspace(1)T, _cmp: T, _new: T, _success_order: u32, _failure_order: u32, _scope: &[u8]) -> (T, bool);
 #[import(cc = "thorin", name = "cmpxchg_weak")] fn cmpxchg_weak_p3[T](_addr: &mut addrspace(3)T, _cmp: T, _new: T, _success_order: u32, _failure_order: u32, _scope: &[u8]) -> (T, bool);
 
+fn @cuda(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) = cuda_with_shared(dev, grid, block, 0, body);
+fn @nvvm(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) = nvvm_with_shared(dev, grid, block, 0, body);
+fn @opencl(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) = opencl_with_shared(dev, grid, block, 0, body);
+fn @amdgpu(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) = amdgpu_with_shared(dev, grid, block, 0, body);
+
 fn @pipeline(body: fn(i32) -> ()) = @|initiation_interval: i32, lower: i32, upper: i32| thorin_pipeline(initiation_interval, lower, upper, body);
 fn @parallel(body: fn(i32) -> ()) = @|num_threads: i32, lower: i32, upper: i32| thorin_parallel(num_threads, lower, upper, body);
 fn @spawn(body: fn() -> ()) = @|| thorin_spawn(body);

platforms/impala/intrinsics_thorin.impala

@@ -10,10 +10,10 @@ extern "thorin" {
     fn insert[T, U](T, i32, U) -> T;
     //fn shuffle[T](T, T, T) -> T;
 
-    fn cuda(i32, (i32, i32, i32), (i32, i32, i32), fn() -> ()) -> ();
-    fn nvvm(i32, (i32, i32, i32), (i32, i32, i32), fn() -> ()) -> ();
-    fn opencl(i32, (i32, i32, i32), (i32, i32, i32), fn() -> ()) -> ();
-    fn amdgpu(i32, (i32, i32, i32), (i32, i32, i32), fn() -> ()) -> ();
+    fn "cuda" cuda_with_shared(i32, (i32, i32, i32), (i32, i32, i32), i32, fn() -> ()) -> ();
+    fn "nvvm" nvvm_with_shared(i32, (i32, i32, i32), (i32, i32, i32), i32, fn() -> ()) -> ();
+    fn "opencl" opencl_with_shared(i32, (i32, i32, i32), (i32, i32, i32), i32, fn() -> ()) -> ();
+    fn "amdgpu" amdgpu_with_shared(i32, (i32, i32, i32), (i32, i32, i32), i32, fn() -> ()) -> ();
     fn reserve_shared[T](i32) -> &mut[3][T];
 
     fn hls(dev: i32, body: fn() -> ()) -> ();
@@ -42,3 +42,8 @@ extern "thorin" {
 
     fn vectorize(vector_length: i32, body: fn(i32) -> ()) -> ();
 }
+
+fn @@cuda(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) { cuda_with_shared(dev, grid, block, 0, body) }
+fn @@nvvm(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) { nvvm_with_shared(dev, grid, block, 0, body) }
+fn @@opencl(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) { opencl_with_shared(dev, grid, block, 0, body) }
+fn @@amdgpu(dev: i32, grid: (i32, i32, i32), block: (i32, i32, i32), body: fn() -> ()) { amdgpu_with_shared(dev, grid, block, 0, body) }

src/anydsl_runtime.cpp

@@ -117,6 +117,7 @@ void anydsl_copy(
 void anydsl_launch_kernel(
     int32_t mask, const char* file_name, const char* kernel_name,
     const uint32_t* grid, const uint32_t* block,
+    uint32_t smem,
     void** arg_data,
     const uint32_t* arg_sizes,
     const uint32_t* arg_aligns,
@@ -128,6 +129,7 @@ void anydsl_launch_kernel(
         kernel_name,
         grid,
         block,
+        smem,
         {
             arg_data,
             arg_sizes,

src/anydsl_runtime.h

@@ -36,6 +36,7 @@ AnyDSL_runtime_API void anydsl_copy(int32_t, const void*, int64_t, int32_t, void
 AnyDSL_runtime_API void anydsl_launch_kernel(
     int32_t, const char*, const char*,
     const uint32_t*, const uint32_t*,
+    uint32_t,
     void**, const uint32_t*, const uint32_t*, const uint32_t*, const uint8_t*,
     uint32_t);
 AnyDSL_runtime_API void anydsl_synchronize(int32_t);

src/cuda_platform.cpp

@@ -225,7 +225,9 @@ void CudaPlatform::launch_kernel(DeviceId dev, const LaunchParams& launch_params
         launch_params.grid[1] / launch_params.block[1],
         launch_params.grid[2] / launch_params.block[2],
         launch_params.block[0], launch_params.block[1], launch_params.block[2],
-        0, nullptr, launch_params.args.data, nullptr);
+        launch_params.smem,
+        nullptr,
+        launch_params.args.data, nullptr);
     CHECK_CUDA(err, "cuLaunchKernel()");
 
     if (runtime_->profiling_enabled()) {

src/hsa_platform.cpp

@@ -327,6 +327,9 @@ void HSAPlatform::launch_kernel(DeviceId dev, const LaunchParams& launch_params)
     if (!queue)
         error("The selected HSA device '%' cannot execute kernels", dev);
 
+    if (launch_params.smem != 0)
+        error("Dynamic shared memory not yet supported on this platform");
+
     auto& kernel_info = load_kernel(dev, launch_params.file_name, launch_params.kernel_name);
 
     auto align_up = [&] (unsigned int start, unsigned int align) -> unsigned int {

src/opencl_platform.cpp

@@ -363,6 +363,9 @@ void OpenCLPlatform::launch_kernel(DeviceId dev, const LaunchParams& launch_para
         return;
     }
 
+    if (launch_params.smem != 0)
+        error("Dynamic shared memory not yet supported on this platform");
+
     auto kernel = load_kernel(dev, launch_params.file_name, launch_params.kernel_name);
 
     // set up arguments

src/runtime.h

@@ -27,6 +27,7 @@ struct LaunchParams {
     const char* kernel_name;
     const uint32_t* grid;
     const uint32_t* block;
+    uint32_t smem;
     struct {
         void** data;
         const uint32_t* sizes;