File size: 9,130 Bytes
de4d6ae 1f9a369 de4d6ae 39a1f38 ddbb236 39a1f38 ddbb236 39a1f38 ddbb236 39a1f38 94f765e de4d6ae 94f765e ddbb236 39a1f38 ddbb236 de4d6ae 94f765e 39a1f38 de4d6ae 94f765e de4d6ae 94f765e ddbb236 39a1f38 ddbb236 de4d6ae 94f765e 39a1f38 ddbb236 39a1f38 ddbb236 39a1f38 de4d6ae 1f9a369 de4d6ae 1f9a369 de4d6ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | // Local development binding (torch.utils.cpp_extension JIT; Windows/MSVC).
// Mirrors the torch-ext binding's argument order exactly.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "geometry_grad.h"
#include "pathtracer_launch.h"
namespace {
const float* fptr(const torch::Tensor& t) {
return t.numel() ? t.const_data_ptr<float>() : nullptr;
}
const int* iptr(const torch::Tensor& t) {
return t.numel() ? t.const_data_ptr<int>() : nullptr;
}
PtdSceneArgs pack_args(
const torch::Tensor& tris, const torch::Tensor& mat_ids,
const torch::Tensor& uvs, const torch::Tensor& nodes_f,
const torch::Tensor& nodes_i, const torch::Tensor& light_faces,
const torch::Tensor& light_cdf, double total_light_area,
const torch::Tensor& tex, const torch::Tensor& tex_hdr,
const torch::Tensor& emi_tex, const torch::Tensor& emi_hdr,
const torch::Tensor& mat_type, const torch::Tensor& mat_rough,
const torch::Tensor& mat_ior, const torch::Tensor& med_sa,
const torch::Tensor& med_ss, double med_sbar, const torch::Tensor& env,
const torch::Tensor& env_cdf_m, const torch::Tensor& env_cdf_c,
const torch::Tensor& env_pdf, int64_t env_w, int64_t env_h,
int64_t max_bounces) {
TORCH_CHECK(tris.is_cuda() && tris.is_contiguous(), "tris");
TORCH_CHECK(uvs.numel() == tris.size(0) * 6, "uvs [F, 3, 2]");
TORCH_CHECK(tex_hdr.size(0) <= 64, "at most 64 materials");
TORCH_CHECK(max_bounces >= 1 && max_bounces <= 16, "bounces in [1,16]");
PtdSceneArgs a;
a.tris = tris.const_data_ptr<float>();
a.mat_ids = mat_ids.const_data_ptr<int>();
a.uvs = uvs.const_data_ptr<float>();
a.n_faces = (int)tris.size(0);
a.nodes_f = nodes_f.const_data_ptr<float>();
a.nodes_i = nodes_i.const_data_ptr<int>();
a.n_nodes = (int)nodes_f.size(0);
a.light_faces = iptr(light_faces);
a.light_cdf = fptr(light_cdf);
a.n_lights = (int)light_faces.numel();
a.total_light_area = (float)total_light_area;
a.tex = tex.const_data_ptr<float>();
a.tex_hdr = tex_hdr.const_data_ptr<int>();
a.n_texels = (int)tex.size(0);
a.emi_tex = emi_tex.const_data_ptr<float>();
a.emi_hdr = emi_hdr.const_data_ptr<int>();
a.n_emi_texels = (int)emi_tex.size(0);
a.mat_type = mat_type.const_data_ptr<int>();
a.mat_rough = mat_rough.const_data_ptr<float>();
a.mat_ior = mat_ior.const_data_ptr<float>();
a.n_mats = (int)tex_hdr.size(0);
a.med_sa = fptr(med_sa);
a.med_ss = fptr(med_ss);
a.med_sbar = (float)med_sbar;
a.has_med = med_sa.numel() ? 1 : 0;
a.env = fptr(env);
a.env_w = (int)env_w;
a.env_h = (int)env_h;
a.env_cdf_m = fptr(env_cdf_m);
a.env_cdf_c = fptr(env_cdf_c);
a.env_pdf = fptr(env_pdf);
return a;
}
void pt_forward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
torch::Tensor nodes_f, torch::Tensor nodes_i,
torch::Tensor light_faces, torch::Tensor light_cdf,
double total_light_area, torch::Tensor tex,
torch::Tensor tex_hdr, torch::Tensor emi_tex,
torch::Tensor emi_hdr, torch::Tensor mat_type,
torch::Tensor mat_rough, torch::Tensor mat_ior,
torch::Tensor med_sa, torch::Tensor med_ss, double med_sbar,
torch::Tensor env, torch::Tensor env_cdf_m,
torch::Tensor env_cdf_c, torch::Tensor env_pdf, int64_t env_w,
int64_t env_h, torch::Tensor cam, int64_t spp,
int64_t max_bounces, int64_t mode, int64_t seed,
torch::Tensor image) {
PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
light_cdf, total_light_area, tex, tex_hdr,
emi_tex, emi_hdr, mat_type, mat_rough, mat_ior,
med_sa, med_ss, med_sbar, env, env_cdf_m,
env_cdf_c, env_pdf, env_w, env_h, max_bounces);
const at::cuda::CUDAGuard guard(tris.device());
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
ptd_forward_launch(&a, cam_h.const_data_ptr<float>(), (int)image.size(0),
(int)image.size(1), (int)spp, (int)max_bounces,
(int)mode, (long long)seed, image.data_ptr<float>(),
stream);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
torch::Tensor nodes_f, torch::Tensor nodes_i,
torch::Tensor light_faces, torch::Tensor light_cdf,
double total_light_area, torch::Tensor tex,
torch::Tensor tex_hdr, torch::Tensor emi_tex,
torch::Tensor emi_hdr, torch::Tensor mat_type,
torch::Tensor mat_rough, torch::Tensor mat_ior,
torch::Tensor med_sa, torch::Tensor med_ss, double med_sbar,
torch::Tensor env, torch::Tensor env_cdf_m,
torch::Tensor env_cdf_c, torch::Tensor env_pdf, int64_t env_w,
int64_t env_h, torch::Tensor cam, int64_t spp,
int64_t max_bounces, int64_t mode, int64_t seed,
torch::Tensor grad_image, torch::Tensor grad_tex,
torch::Tensor grad_emi_tex, torch::Tensor grad_env,
torch::Tensor grad_med) {
PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
light_cdf, total_light_area, tex, tex_hdr,
emi_tex, emi_hdr, mat_type, mat_rough, mat_ior,
med_sa, med_ss, med_sbar, env, env_cdf_m,
env_cdf_c, env_pdf, env_w, env_h, max_bounces);
const at::cuda::CUDAGuard guard(tris.device());
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
ptd_backward_launch(&a, cam_h.const_data_ptr<float>(),
(int)grad_image.size(0), (int)grad_image.size(1),
(int)spp, (int)max_bounces, (int)mode, (long long)seed,
grad_image.const_data_ptr<float>(),
grad_tex.data_ptr<float>(),
grad_emi_tex.data_ptr<float>(),
grad_env.numel() ? grad_env.data_ptr<float>() : nullptr,
grad_med.numel() ? grad_med.data_ptr<float>() : nullptr,
stream);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
void pt_geometry_grad(torch::Tensor tris, torch::Tensor mat_ids,
torch::Tensor uvs, torch::Tensor nodes_f,
torch::Tensor nodes_i, torch::Tensor light_faces,
torch::Tensor light_cdf, double total_light_area,
torch::Tensor tex, torch::Tensor tex_hdr,
torch::Tensor emi_tex, torch::Tensor emi_hdr,
torch::Tensor mat_type, torch::Tensor mat_rough,
torch::Tensor mat_ior, torch::Tensor face_verts,
torch::Tensor edges, torch::Tensor edge_cdf,
torch::Tensor cam, int64_t spp, int64_t edge_samples,
int64_t seed, torch::Tensor grad_image,
torch::Tensor grad_verts) {
torch::Tensor z = torch::zeros(0, tris.options());
PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i,
light_faces, light_cdf, total_light_area, tex,
tex_hdr, emi_tex, emi_hdr, mat_type, mat_rough,
mat_ior, z, z, 0.0, z, z, z, z, 0, 0, 4);
const at::cuda::CUDAGuard guard(tris.device());
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
int H = (int)grad_image.size(0), W = (int)grad_image.size(1);
ptd_geo_interior_launch(&a, face_verts.const_data_ptr<int>(),
cam_h.const_data_ptr<float>(), H, W, (int)spp,
(long long)seed,
grad_image.const_data_ptr<float>(),
grad_verts.data_ptr<float>(), stream);
if (edges.numel())
ptd_geo_boundary_launch(&a, face_verts.const_data_ptr<int>(),
edges.const_data_ptr<int>(),
(int)edges.size(0),
edge_cdf.const_data_ptr<float>(),
cam_h.const_data_ptr<float>(), H, W,
(int)edge_samples, (long long)seed,
grad_image.const_data_ptr<float>(),
grad_verts.data_ptr<float>(), stream);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
} // namespace
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("pt_forward", &pt_forward, "pathtracer-diff forward");
m.def("pt_backward", &pt_backward, "pathtracer-diff backward");
m.def("pt_geometry_grad", &pt_geometry_grad,
"pathtracer-diff geometry gradients (direct lighting)");
}
|