| 100 | // The kernel parameters 'Ts' must be constructible from the arguments 'Args'. |
| 101 | template <typename... Ts, typename... Args> |
| 102 | Status GpuLaunchKernel(void (*function)(Ts...), dim3 grid_dim, dim3 block_dim, |
| 103 | size_t shared_memory_size_bytes, gpuStream_t stream, |
| 104 | Args... arguments) { |
| 105 | static_assert(detail::NoneIsReference<Ts...>(), |
| 106 | "Kernels with reference arguments have undefined behaviour."); |
| 107 | #if GOOGLE_CUDA |
| 108 | auto func_ptr = absl::bit_cast<const void*>(function); |
| 109 | // Cast arguments and forward them as an array of pointers. |
| 110 | auto args_tuple = std::tuple<Ts...>(arguments...); |
| 111 | auto arg_ptrs = detail::GetArrayOfElementPointers(&args_tuple); |
| 112 | auto result = cudaLaunchKernel(func_ptr, grid_dim, block_dim, arg_ptrs.data(), |
| 113 | shared_memory_size_bytes, stream); |
| 114 | if (result != cudaSuccess) { |
| 115 | return errors::Internal(cudaGetErrorString(result)); |
| 116 | } |
| 117 | #elif TENSORFLOW_USE_ROCM |
| 118 | hipLaunchKernelGGL(function, grid_dim, block_dim, shared_memory_size_bytes, |
| 119 | stream, std::forward<Args>(arguments)...); |
| 120 | #endif |
| 121 | return Status::OK(); |
| 122 | } |
| 123 | |
| 124 | // Perfect forwarding to make CudaLaunchKernel available to both ROCm and CUDA |
| 125 | // builds |