Add dynamic shared memory allocation

michael-kenzel · michael-kenzel · commit 00bb7bda4828 · 2023-08-25T03:09:28.000+02:00
diff --git a/src/thorin/be/c/c.cpp b/src/thorin/be/c/c.cpp
@@ -510,6 +510,9 @@ void CCodeGen::emit_module() {
             stream_.fmt("__device__ inline int blockDim_{}() {{ return blockDim.{}; }}\n", x, x);
             stream_.fmt("__device__ inline int gridDim_{}() {{ return gridDim.{}; }}\n", x, x);
         }
+
+        stream_.fmt("\n"
+                    "extern __shared__ unsigned char __dynamic_smem[];\n");
     }
 
     stream_.endl() << func_impls_.str();
@@ -737,7 +740,10 @@ void CCodeGen::emit_epilogue(Continuation* cont) {
         bb.tail.fmt("goto {};", label_name(callee));
     } else if (auto callee = body->callee()->isa_nom<Continuation>(); callee && callee->is_intrinsic()) {
         if (callee->intrinsic() == Intrinsic::Reserve) {
+            assert(body->num_args() == 3 && "incorrect number of arguments");
+
             emit_unsafe(body->arg(0));
+
             if (!body->arg(1)->isa<PrimLit>())
                 world().edef(body->arg(1), "reserve_shared: couldn't extract memory size");
 
@@ -753,6 +759,16 @@ void CCodeGen::emit_epilogue(Continuation* cont) {
             }
             bb.tail.fmt("p_{} = {}_reserved;\n", ret_cont->param(1)->unique_name(), cont->unique_name());
             bb.tail.fmt("goto {};", label_name(ret_cont));
+        } else if (callee->intrinsic() == Intrinsic::LocalMemory) {
+            if (lang_ == Lang::HLS)
+                world().edef(body, "local_memory not supported for HLS");
+            assert(body->num_args() == 2 && "incorrect number of arguments");
+
+            emit_unsafe(body->arg(0));
+
+            auto ret_cont = body->arg(1)->as_nom<Continuation>();
+            bb.tail.fmt("p_{} = __dynamic_smem;\n", ret_cont->param(1)->unique_name());
+            bb.tail.fmt("goto {};", label_name(ret_cont));
         } else if (callee->intrinsic() == Intrinsic::Pipeline) {
             assert((lang_ == Lang::OpenCL || lang_ == Lang::HLS) && "pipelining not supported on this backend");
 
@@ -1435,6 +1451,12 @@ std::string CCodeGen::emit_fun_head(Continuation* cont, bool is_proto) {
         }
         needs_comma = true;
     }
+
+    if (cont->is_exported() && lang_ == Lang::OpenCL) {
+        if (needs_comma) s.fmt(", ");
+        s.fmt("__local unsigned char* __dynamic_smem");
+    }
+
     s << ")";
     return s.str();
 }
diff --git a/src/thorin/be/codegen.h b/src/thorin/be/codegen.h
@@ -32,6 +32,7 @@ struct LaunchArgs {
         Device,
         Space,
         Config,
+        LocalMem,
         Body,
         Return,
         Num
diff --git a/src/thorin/be/llvm/amdgpu.cpp b/src/thorin/be/llvm/amdgpu.cpp
@@ -90,4 +90,8 @@ Continuation* AMDGPUCodeGen::emit_reserve(llvm::IRBuilder<>& irbuilder, const Co
     return emit_reserve_shared(irbuilder, continuation, true);
 }
 
+Continuation* AMDGPUCodeGen::emit_local_memory(llvm::IRBuilder<>& irbuilder, const Continuation* continuation) {
+    return emit_local_memory_base_ptr(irbuilder, continuation);
+}
+
 }
diff --git a/src/thorin/be/llvm/amdgpu.h b/src/thorin/be/llvm/amdgpu.h
@@ -23,6 +23,7 @@ class AMDGPUCodeGen : public CodeGen {
     llvm::Value* emit_global(const Global*) override;
     llvm::Value* emit_mathop(llvm::IRBuilder<>&, const MathOp*) override;
     Continuation* emit_reserve(llvm::IRBuilder<>&, const Continuation*) override;
+    Continuation* emit_local_memory(llvm::IRBuilder<>&, const Continuation*) override;
     std::string get_alloc_name() const override { return "malloc"; }
 
     const Cont2Config& kernel_config_;
diff --git a/src/thorin/be/llvm/llvm.cpp b/src/thorin/be/llvm/llvm.cpp
@@ -1182,6 +1182,7 @@ Continuation* CodeGen::emit_intrinsic(llvm::IRBuilder<>& irbuilder, Continuation
         case Intrinsic::CmpXchgWeak: return emit_cmpxchg(irbuilder, continuation, true);
         case Intrinsic::Fence:       return emit_fence(irbuilder, continuation);
         case Intrinsic::Reserve:     return emit_reserve(irbuilder, continuation);
+        case Intrinsic::LocalMemory: return emit_local_memory(irbuilder, continuation);
         case Intrinsic::CUDA:        return runtime_->emit_host_code(*this, irbuilder, Runtime::CUDA_PLATFORM,   ".cu",     continuation);
         case Intrinsic::NVVM:        return runtime_->emit_host_code(*this, irbuilder, Runtime::CUDA_PLATFORM,   ".nvvm",   continuation);
         case Intrinsic::OpenCL:      return runtime_->emit_host_code(*this, irbuilder, Runtime::OPENCL_PLATFORM, ".cl",     continuation);
@@ -1309,7 +1310,7 @@ Continuation* CodeGen::emit_reserve(llvm::IRBuilder<>&, const Continuation* cont
 Continuation* CodeGen::emit_reserve_shared(llvm::IRBuilder<>& irbuilder, const Continuation* continuation, bool init_undef) {
     assert(continuation->has_body());
     auto body = continuation->body();
-    assert(body->num_args() == 3 && "required arguments are missing");
+    assert(body->num_args() == 3 && "incorrect number of arguments");
     if (!body->arg(1)->isa<PrimLit>())
         world().edef(body->arg(1), "reserve_shared: couldn't extract memory size");
     auto num_elems = body->arg(1)->as<PrimLit>()->ps32_value();
@@ -1327,6 +1328,33 @@ Continuation* CodeGen::emit_reserve_shared(llvm::IRBuilder<>& irbuilder, const C
     return cont;
 }
 
+Continuation* CodeGen::emit_local_memory(llvm::IRBuilder<>&, const Continuation* continuation) {
+    world().edef(continuation, "local_memory: only allowed in device code");
+    THORIN_UNREACHABLE;
+}
+
+Continuation* CodeGen::emit_local_memory_base_ptr(llvm::IRBuilder<>& irbuilder, const Continuation* continuation) {
+    static constexpr auto name = "__dynamic_smem";
+
+    assert(continuation->has_body());
+    auto body = continuation->body();
+    assert(body->num_args() == 2 && "incorrect number of arguments");
+    auto cont = body->arg(1)->as_nom<Continuation>();
+
+    auto global = [&] {
+        if (auto found = module().getGlobalVariable(name))
+            return found;
+
+        auto type = llvm::ArrayType::get(llvm::Type::getInt8Ty(context()), 0);
+        auto global = new llvm::GlobalVariable(module(), type, false, llvm::GlobalValue::ExternalLinkage, nullptr, name, nullptr, llvm::GlobalVariable::NotThreadLocal, 3);
+        global->setAlignment(llvm::Align(16));
+        return global;
+    }();
+
+    emit_phi_arg(irbuilder, cont->param(1), global);
+    return cont;
+}
+
 /*
  * backend-specific stuff
  */
diff --git a/src/thorin/be/llvm/llvm.h b/src/thorin/be/llvm/llvm.h
@@ -83,6 +83,8 @@ class CodeGen : public thorin::CodeGen, public thorin::Emitter<llvm::Value*, llv
 
     virtual Continuation* emit_reserve(llvm::IRBuilder<>&, const Continuation*);
     Continuation* emit_reserve_shared(llvm::IRBuilder<>&, const Continuation*, bool=false);
+    virtual Continuation* emit_local_memory(llvm::IRBuilder<>&, const Continuation*);
+    Continuation* emit_local_memory_base_ptr(llvm::IRBuilder<>& irbuilder, const Continuation* continuation);
 
     virtual std::string get_alloc_name() const = 0;
     llvm::BasicBlock* cont2bb(Continuation* cont) { return cont2bb_[cont].first; }
diff --git a/src/thorin/be/llvm/nvvm.cpp b/src/thorin/be/llvm/nvvm.cpp
@@ -246,6 +246,10 @@ Continuation* NVVMCodeGen::emit_reserve(llvm::IRBuilder<>& irbuilder, const Cont
     return emit_reserve_shared(irbuilder, continuation);
 }
 
+Continuation* NVVMCodeGen::emit_local_memory(llvm::IRBuilder<>& irbuilder, const Continuation* continuation) {
+    return emit_local_memory_base_ptr(irbuilder, continuation);
+}
+
 llvm::Value* NVVMCodeGen::emit_mathop(llvm::IRBuilder<>& irbuilder, const MathOp* mathop) {
     auto make_key = [] (MathOpTag tag, unsigned bitwidth) { return (static_cast<unsigned>(tag) << 16) | bitwidth; };
     static const std::unordered_map<unsigned, std::string> libdevice_functions = {
diff --git a/src/thorin/be/llvm/nvvm.h b/src/thorin/be/llvm/nvvm.h
@@ -29,6 +29,7 @@ class NVVMCodeGen : public CodeGen {
     llvm::Value* emit_mathop(llvm::IRBuilder<>&, const MathOp*) override;
 
     Continuation* emit_reserve(llvm::IRBuilder<>&, const Continuation*) override;
+    Continuation* emit_local_memory(llvm::IRBuilder<>&, const Continuation*) override;
 
     llvm::Value* emit_global(const Global*) override;
 
diff --git a/src/thorin/be/llvm/runtime.cpp b/src/thorin/be/llvm/runtime.cpp
@@ -64,21 +64,24 @@ Continuation* Runtime::emit_host_code(CodeGen& code_gen, llvm::IRBuilder<>& buil
     assert(continuation->has_body());
     auto body = continuation->body();
     // to-target is the desired kernel call
-    // target(mem, device, (dim.x, dim.y, dim.z), (block.x, block.y, block.z), body, return, free_vars)
+    // target(mem, device, (dim.x, dim.y, dim.z), (block.x, block.y, block.z), lmem, body, return, free_vars)
     auto target = body->callee()->as_nom<Continuation>();
     assert_unused(target->is_intrinsic());
     assert(body->num_args() >= LaunchArgs::Num && "required arguments are missing");
 
+    auto& world = continuation->world();
+
     // arguments
     auto target_device_id = code_gen.emit(body->arg(LaunchArgs::Device));
     auto target_platform = builder.getInt32(platform);
     auto target_device = builder.CreateOr(target_platform, builder.CreateShl(target_device_id, builder.getInt32(4)));
 
     auto it_space = body->arg(LaunchArgs::Space);
     auto it_config = body->arg(LaunchArgs::Config);
-    auto kernel = body->arg(LaunchArgs::Body)->as<Global>()->init()->as<Continuation>();
 
-    auto& world = continuation->world();
+    auto lmem = code_gen.emit(body->arg(LaunchArgs::LocalMem));
+
+    auto kernel = body->arg(LaunchArgs::Body)->as<Global>()->init()->as<Continuation>();
     auto kernel_name = builder.CreateGlobalStringPtr(kernel->name() == "hls_top" ? kernel->name() : kernel->unique_name());
     auto file_name = builder.CreateGlobalStringPtr(world.name() + ext);
     const size_t num_kernel_args = body->num_args() - LaunchArgs::Num;
@@ -181,6 +184,7 @@ Continuation* Runtime::emit_host_code(CodeGen& code_gen, llvm::IRBuilder<>& buil
     launch_kernel(code_gen, builder, target_device,
                   file_name, kernel_name,
                   grid_size, block_size,
+                  lmem,
                   args, sizes, aligns, allocs, types,
                   builder.getInt32(num_kernel_args));
 
@@ -191,10 +195,11 @@ llvm::Value* Runtime::launch_kernel(
     CodeGen& code_gen, llvm::IRBuilder<>& builder, llvm::Value* device,
     llvm::Value* file, llvm::Value* kernel,
     llvm::Value* grid, llvm::Value* block,
+    llvm::Value* lmem,
     llvm::Value* args, llvm::Value* sizes, llvm::Value* aligns, llvm::Value* allocs, llvm::Value* types,
     llvm::Value* num_args)
 {
-    llvm::Value* launch_args[] = { device, file, kernel, grid, block, args, sizes, aligns, allocs, types, num_args };
+    llvm::Value* launch_args[] = { device, file, kernel, grid, block, lmem, args, sizes, aligns, allocs, types, num_args };
     return builder.CreateCall(get(code_gen, "anydsl_launch_kernel"), launch_args);
 }
 
diff --git a/src/thorin/be/llvm/runtime.h b/src/thorin/be/llvm/runtime.h
@@ -30,6 +30,7 @@ class Runtime {
         CodeGen&, llvm::IRBuilder<>&, llvm::Value* device,
         llvm::Value* file, llvm::Value* kernel,
         llvm::Value* grid, llvm::Value* block,
+        llvm::Value* lmem,
         llvm::Value* args, llvm::Value* sizes, llvm::Value* aligns, llvm::Value* allocs, llvm::Value* types,
         llvm::Value* num_args);
 
diff --git a/src/thorin/be/llvm/runtime.inc b/src/thorin/be/llvm/runtime.inc
@@ -6,7 +6,7 @@ namespace thorin {
         declare noalias ptr @anydsl_alloc(i32, i64);
         declare noalias ptr @anydsl_alloc_unified(i32, i64);
         declare void @anydsl_release(i32, ptr);
-        declare void @anydsl_launch_kernel(i32, ptr, ptr, ptr, ptr, ptr, ptr, ptr, ptr, ptr, i32);
+        declare void @anydsl_launch_kernel(i32, ptr, ptr, ptr, ptr, i32, ptr, ptr, ptr, ptr, ptr, i32);
         declare void @anydsl_parallel_for(i32, i32, i32, ptr, ptr);
         declare void @anydsl_fibers_spawn(i32, i32, i32, ptr, ptr);
         declare i32  @anydsl_spawn_thread(ptr, ptr);
diff --git a/src/thorin/continuation.cpp b/src/thorin/continuation.cpp
@@ -224,6 +224,7 @@ void Continuation::set_intrinsic() {
     else if (name() == "pe_info")        attributes().intrinsic = Intrinsic::PeInfo;
     else if (name() == "pipeline")       attributes().intrinsic = Intrinsic::Pipeline;
     else if (name() == "reserve_shared") attributes().intrinsic = Intrinsic::Reserve;
+    else if (name() == "local_memory")   attributes().intrinsic = Intrinsic::LocalMemory;
     else if (name() == "atomic")         attributes().intrinsic = Intrinsic::Atomic;
     else if (name() == "atomic_load")    attributes().intrinsic = Intrinsic::AtomicLoad;
     else if (name() == "atomic_store")   attributes().intrinsic = Intrinsic::AtomicStore;
diff --git a/src/thorin/continuation.h b/src/thorin/continuation.h
@@ -99,6 +99,7 @@ enum class Intrinsic : uint8_t {
     Vectorize,                  ///< External vectorizer.
     AcceleratorEnd,
     Reserve = AcceleratorEnd,   ///< Intrinsic memory reserve function
+    LocalMemory,                ///< Intrinsic get local memory base pointer
     Atomic,                     ///< Intrinsic atomic function
     AtomicLoad,                 ///< Intrinsic atomic load function
     AtomicStore,                ///< Intrinsic atomic store function

-Original file line number
+Diff line change
         Device,
         Space,
         Config,
 +        LocalMem,
         Body,
         Return,
         Num
Original file line number	Diff line number	Diff line change
`@@ -90,4 +90,8 @@ Continuation* AMDGPUCodeGen::emit_reserve(llvm::IRBuilder<>& irbuilder, const Co`
`90`	`90`	`return emit_reserve_shared(irbuilder, continuation, true);`
`91`	`91`	`}`
`92`	`92`
	`93`	`+Continuation* AMDGPUCodeGen::emit_local_memory(llvm::IRBuilder<>& irbuilder, const Continuation* continuation) {`
	`94`	`+ return emit_local_memory_base_ptr(irbuilder, continuation);`
	`95`	`+}`
	`96`	`+`
`93`	`97`	`}`