eic-opticks/sysrap/torch.cu

0001 #include "torch.h"
0002
0003 #include <cuda_runtime.h>
0004 #include <curand_kernel.h>
0005
0006 #include <cstddef>
0007
0008 #include "CUDA_CHECK.h"
0009 #include "srng.h"
0010
0011 namespace
0012 {
0013
0014 __global__ void generate_photons_kernel(storch torch, sphoton* photons, unsigned int num_photons, unsigned int seed)
0015 {
0016     unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
0017     if (idx >= num_photons)
0018     {
0019         return;
0020     }
0021
0022     curandStatePhilox4_32_10 rng;
0023     curand_init(seed, idx, 0, &rng);
0024
0025     qtorch  qt{.t = torch};
0026
0027     sphoton photon;
0028     storch::generate(photon, rng, qt.q, idx, 0);
0029     photons[idx] = photon;
0030 }
0031
0032 /**
0033 Small RAII guard for the temporary device photon array used by
0034 generate_photons_gpu. The wrapper is not needed for the successful path, but
0035 keeps cleanup correct when CUDA_CHECK or CUDA_SYNC_CHECK throws after
0036 cudaMalloc. Without a guard, failures in the kernel launch, synchronization, or
0037 device-to-host copy would skip cudaFree and leak the device allocation.
0038 **/
0039 struct device_photon_buffer
0040 {
0041     explicit device_photon_buffer(size_t count)
0042     {
0043         CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&ptr), count * sizeof(sphoton)));
0044     }
0045
0046     ~device_photon_buffer()
0047     {
0048         CUDA_CHECK_NOEXCEPT(cudaFree(ptr));
0049     }
0050
0051     device_photon_buffer(const device_photon_buffer&) = delete;
0052     device_photon_buffer& operator=(const device_photon_buffer&) = delete;
0053
0054     sphoton* ptr = nullptr;
0055 };
0056
0057 } // namespace
0058
0059 std::vector<sphoton> generate_photons_gpu(const storch& torch, unsigned int num_photons, unsigned int seed)
0060 {
0061     if (num_photons == 0)
0062     {
0063         num_photons = torch.numphoton;
0064     }
0065
0066     std::vector<sphoton> photons(num_photons);
0067     if (num_photons == 0)
0068     {
0069         return photons;
0070     }
0071
0072     device_photon_buffer d_photons(photons.size());
0073
0074     constexpr unsigned int threads_per_block = 256;
0075     unsigned int blocks = (num_photons + threads_per_block - 1) / threads_per_block;
0076
0077     generate_photons_kernel<<<blocks, threads_per_block>>>(torch, d_photons.ptr, num_photons, seed);
0078     CUDA_SYNC_CHECK();
0079
0080     CUDA_CHECK(cudaMemcpy(photons.data(), d_photons.ptr, photons.size() * sizeof(sphoton), cudaMemcpyDeviceToHost));
0081
0082     return photons;
0083 }