| #include <torch/library.h> |
|
|
| #include "registration.h" |
| #include "torch_binding.h" |
|
|
| TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| ops.def("nsw_rhs(Tensor h, Tensor hu, Tensor hv, Tensor b, float dx, float dy) " |
| "-> Tensor[]"); |
| ops.def("sph_forces(Tensor pos, Tensor vel, Tensor rho, Tensor type, " |
| "Tensor cell_start, Tensor cell_end, int gx, int gy, int gz, " |
| "float ox, float oy, float oz, float csz, float h, float mass, " |
| "float rho0, float gamma, float bmod, float c0, float alpha, " |
| "float delta, float eps_xsph, float grav) -> Tensor[]"); |
|
|
| #if defined(CUDA_KERNEL) || defined(ROCM_KERNEL) |
| ops.impl("nsw_rhs", torch::kCUDA, &nsw_rhs_cuda); |
| ops.impl("sph_forces", torch::kCUDA, &sph_forces_cuda); |
| #endif |
| } |
|
|
| REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
|
|
