| 14 | #if MGB_CUDA |
| 15 | |
| 16 | const CudaRuntimeArgs get_cuda_runtime_args(const RuntimeArgs& rt_args) { |
| 17 | mgb_assert( |
| 18 | rt_args.device().enumv() == DeviceEnum::cuda, |
| 19 | "devive type should be cuda."); |
| 20 | const CompNodeEnv& env = |
| 21 | CompNodeEnv::from_comp_node(to_builtin<CompNode, Device>(rt_args.device())); |
| 22 | const CompNodeEnv::CudaEnv& cuda_env = env.cuda_env(); |
| 23 | return {cuda_env.device, cuda_env.stream}; |
| 24 | } |
| 25 | |
| 26 | int get_cuda_device_id(Device device) { |
| 27 | auto cn = to_builtin<CompNode>(device); |