pathtracer-diff / torch-ext /torch_binding.cpp

v3: geometry gradients (dual-number interior + shadow and camera silhouette edge sampling)

1f9a369 verified 20 days ago

13.9 kB

	#include <torch/library.h>
	#include <torch/all.h>

	#include <ATen/cuda/CUDAContext.h>
	#include <c10/cuda/CUDAGuard.h>

	#include "registration.h"
	#include "torch_binding.h"
	#include "../pathtracer_diff_cuda/geometry_grad.h"
	#include "../pathtracer_diff_cuda/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;
	}

	void chk(bool ok, const char* msg) { TORCH_CHECK(ok, msg); }

	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 spp,
	int64_t max_bounces, int64_t mode) {
	chk(tris.is_cuda() && tris.is_contiguous() &&
	tris.dtype() == torch::kFloat32 && tris.dim() == 2 &&
	tris.size(1) == 9,
	"tris must be contiguous CUDA f32 [F, 9]");
	chk(mat_ids.dtype() == torch::kInt32 && mat_ids.numel() == tris.size(0),
	"mat_ids must be i32 [F]");
	chk(uvs.numel() == tris.size(0) * 6, "uvs must be [F, 3, 2]");
	chk(nodes_f.dim() == 2 && nodes_f.size(1) == 6, "nodes_f [N, 6]");
	chk(nodes_i.dim() == 2 && nodes_i.size(1) == 3 &&
	nodes_i.size(0) == nodes_f.size(0),
	"nodes_i [N, 3]");
	chk(light_cdf.numel() == light_faces.numel(), "light list mismatch");
	chk(tex.dim() == 2 && tex.size(1) == 3, "tex [T, 3]");
	chk(tex_hdr.dim() == 2 && tex_hdr.size(1) == 3, "tex_hdr [M, 3]");
	int64_t M = tex_hdr.size(0);
	chk(M <= 64, "at most 64 materials");
	chk(emi_tex.dim() == 2 && emi_tex.size(1) == 3, "emi_tex [Te, 3]");
	chk(emi_hdr.sizes() == tex_hdr.sizes(), "emi_hdr [M, 3]");
	chk(mat_type.numel() == M && mat_rough.numel() == M && mat_ior.numel() == M,
	"per-material arrays must have M entries");
	chk(med_sa.numel() == med_ss.numel() &&
	(med_sa.numel() == 0 \|\| med_sa.numel() == 3),
	"medium sigmas must be [3] or empty");
	if (env.numel()) {
	chk(env.numel() == env_w * env_h * 3 && env_cdf_m.numel() == env_h &&
	env_cdf_c.numel() == env_w * env_h &&
	env_pdf.numel() == env_w * env_h,
	"env tables mismatch");
	}
	chk(spp >= 1, "spp must be >= 1");
	chk(max_bounces >= 1 && max_bounces <= 16, "max_bounces in [1, 16]");
	chk(mode >= 0 && mode <= 2, "mode must be 0\|1\|2");

	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)M;
	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;
	}

	} // namespace

	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, spp,
	max_bounces, mode);
	TORCH_CHECK(image.is_cuda() && image.is_contiguous() &&
	image.dtype() == torch::kFloat32 && image.dim() == 3 &&
	image.size(2) == 3,
	"image must be contiguous CUDA f32 [H, W, 3]");
	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, spp,
	max_bounces, mode);
	TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous() &&
	grad_image.dim() == 3 && grad_image.size(2) == 3,
	"grad_image must be contiguous CUDA f32 [H, W, 3]");
	TORCH_CHECK(grad_tex.sizes() == tex.sizes(), "grad_tex must match tex");
	TORCH_CHECK(grad_emi_tex.sizes() == emi_tex.sizes(),
	"grad_emi_tex must match emi_tex");
	TORCH_CHECK(grad_env.sizes() == env.sizes(), "grad_env must match env");
	TORCH_CHECK(grad_med.numel() == (med_sa.numel() ? 6 : 0),
	"grad_med must be [6] with a medium, else empty");
	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,
	/spp=/1, /max_bounces=/4, /mode=/0);
	TORCH_CHECK(face_verts.dtype() == torch::kInt32 &&
	face_verts.numel() == tris.size(0) * 3,
	"face_verts must be i32 [F, 3]");
	TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous() &&
	grad_image.dim() == 3 && grad_image.size(2) == 3,
	"grad_image must be contiguous CUDA f32 [H, W, 3]");
	TORCH_CHECK(grad_verts.dim() == 2 && grad_verts.size(1) == 3,
	"grad_verts must be [V, 3]");
	TORCH_CHECK(edges.numel() == 0 \|\|
	(edges.dtype() == torch::kInt32 && edges.dim() == 2 &&
	edges.size(1) == 4 &&
	edge_cdf.numel() == edges.size(0) + 1),
	"edges [E, 4] i32 with edge_cdf [E + 1]");
	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();
	}

	TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
	ops.def(
	"pt_forward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
	" Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
	" float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emi_tex,"
	" Tensor emi_hdr, Tensor mat_type, Tensor mat_rough, Tensor mat_ior,"
	" Tensor med_sa, Tensor med_ss, float med_sbar, Tensor env,"
	" Tensor env_cdf_m, Tensor env_cdf_c, Tensor env_pdf, int env_w,"
	" int env_h, Tensor cam, int spp, int max_bounces, int mode, int seed,"
	" Tensor! image) -> ()");
	ops.impl("pt_forward", torch::kCUDA, &pt_forward);

	ops.def(
	"pt_backward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
	" Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
	" float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emi_tex,"
	" Tensor emi_hdr, Tensor mat_type, Tensor mat_rough, Tensor mat_ior,"
	" Tensor med_sa, Tensor med_ss, float med_sbar, Tensor env,"
	" Tensor env_cdf_m, Tensor env_cdf_c, Tensor env_pdf, int env_w,"
	" int env_h, Tensor cam, int spp, int max_bounces, int mode, int seed,"
	" Tensor grad_image, Tensor! grad_tex, Tensor! grad_emi_tex,"
	" Tensor! grad_env, Tensor! grad_med) -> ()");
	ops.impl("pt_backward", torch::kCUDA, &pt_backward);

	ops.def(
	"pt_geometry_grad(Tensor tris, Tensor mat_ids, Tensor uvs,"
	" Tensor nodes_f, Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
	" float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emi_tex,"
	" Tensor emi_hdr, Tensor mat_type, Tensor mat_rough, Tensor mat_ior,"
	" Tensor face_verts, Tensor edges, Tensor edge_cdf, Tensor cam,"
	" int spp, int edge_samples, int seed, Tensor grad_image,"
	" Tensor! grad_verts) -> ()");
	ops.impl("pt_geometry_grad", torch::kCUDA, &pt_geometry_grad);
	}

	REGISTER_EXTENSION(TORCH_EXTENSION_NAME)