liyang/hand-motion-pipeline
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update to 2026-09-17 pipeline snapshot; add weights, L20 assets and recording via Git LFS

Browse Source 
Source: RGB-D -> Dyn-HaMR -> L20 retargeting -> FoundationPose -> reference repair -> SPIDER,
documented in docs/PIPELINE_LATEST.md and docs/SETUP_AND_WEIGHTS.md. Adds FoundationPose and
nvdiffrast upstream snapshots, requirements/pipeline_venv.txt and the FoundationPose weight
manifest/downloader.

Assets (Git LFS): weights/ (WiLoR detector, HandFlow denoiser, UniDepth-L), FoundationPose
checkpoints, HaMeR checkpoint, Dyn-HaMR HMP model and BMC constraints, L20 URDF/meshes, the
20260915_171525 D405 recording and the two box CADs. MANO models are not redistributed
(third_party/hamer/_DATA/data/mano/README.txt). Environments, caches and run outputs excluded.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
liyang
2026-09-17 11:27:46 +08:00
parent 7e4ef6f98b
commit ae28d55f81
761 changed files with 55598 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/nvdiffrast/csrc/torch/torch_antialias.cpp
Vendored
+243
View File
@@ -0,0 +1,243 @@
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
#include "torch_common.inl"
#include "torch_types.h"
#include "../common/common.h"
#include "../common/antialias.h"
//------------------------------------------------------------------------
// Kernel prototypes.
void AntialiasFwdMeshKernel (const AntialiasKernelParams p);
void AntialiasFwdDiscontinuityKernel(const AntialiasKernelParams p);
void AntialiasFwdAnalysisKernel (const AntialiasKernelParams p);
void AntialiasGradKernel (const AntialiasKernelParams p);
//------------------------------------------------------------------------
// Topology hash construction.
TopologyHashWrapper antialias_construct_topology_hash(torch::Tensor tri)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(tri));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AntialiasKernelParams p = {}; // Initialize all fields to zero.
// Check inputs.
NVDR_CHECK_DEVICE(tri);
NVDR_CHECK_CONTIGUOUS(tri);
NVDR_CHECK_I32(tri);
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
// Fill in kernel parameters.
p.numTriangles = tri.size(0);
p.numVertices = 0x7fffffff; // Let's not require vertex positions just to enable an error check.
p.tri = tri.data_ptr<int>();
// Kernel parameters.
p.allocTriangles = 64;
while (p.allocTriangles < p.numTriangles)
p.allocTriangles <<= 1; // Must be power of two.
// Construct the hash tensor and get pointer.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA);
torch::Tensor ev_hash = torch::zeros({(uint64_t)p.allocTriangles * AA_HASH_ELEMENTS_PER_TRIANGLE(p.allocTriangles) * 4}, opts);
p.evHash = (uint4*)(ev_hash.data_ptr<int>());
// Check alignment.
NVDR_CHECK(!((uintptr_t)p.evHash & 15), "ev_hash internal tensor not aligned to int4");
// Populate the hash.
void* args[] = {&p};
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasFwdMeshKernel, (p.numTriangles - 1) / AA_MESH_KERNEL_THREADS_PER_BLOCK + 1, AA_MESH_KERNEL_THREADS_PER_BLOCK, args, 0, stream));
// Return.
TopologyHashWrapper hash_wrap;
hash_wrap.ev_hash = ev_hash;
return hash_wrap;
}
//------------------------------------------------------------------------
// Forward op.
std::tuple<torch::Tensor, torch::Tensor> antialias_fwd(torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, TopologyHashWrapper topology_hash_wrap)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(color));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AntialiasKernelParams p = {}; // Initialize all fields to zero.
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
torch::Tensor& topology_hash = topology_hash_wrap.ev_hash; // Unwrap.
// Check inputs.
NVDR_CHECK_DEVICE(color, rast, pos, tri, topology_hash);
NVDR_CHECK_CONTIGUOUS(color, rast, pos, tri, topology_hash);
NVDR_CHECK_F32(color, rast, pos);
NVDR_CHECK_I32(tri, topology_hash);
// Sanity checks.
NVDR_CHECK(color.sizes().size() == 4 && color.size(0) > 0 && color.size(1) > 0 && color.size(2) > 0 && color.size(3) > 0, "color must have shape[>0, >0, >0, >0]");
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "rast must have shape[>0, >0, >0, 4]");
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
NVDR_CHECK(color.size(1) == rast.size(1) && color.size(2) == rast.size(2), "color and rast inputs must have same spatial dimensions");
if (p.instance_mode)
{
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) > 0 && pos.size(1) > 0 && pos.size(2) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
NVDR_CHECK(rast.size(0) == color.size(0) && pos.size(0) == color.size(0), "minibatch size mismatch between inputs color, rast, pos");
}
else
{
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
NVDR_CHECK(rast.size(0) == color.size(0), "minibatch size mismatch between inputs color, rast");
}
// Extract input dimensions.
p.numVertices = pos.size(p.instance_mode ? 1 : 0);
p.numTriangles = tri.size(0);
p.n = color.size(0);
p.height = color.size(1);
p.width = color.size(2);
p.channels = color.size(3);
// Get input pointers.
p.color = color.data_ptr<float>();
p.rasterOut = rast.data_ptr<float>();
p.tri = tri.data_ptr<int>();
p.pos = pos.data_ptr<float>();
p.evHash = (uint4*)(topology_hash.data_ptr<int>());
// Misc parameters.
p.xh = .5f * (float)p.width;
p.yh = .5f * (float)p.height;
// Determine hash allocation size.
p.allocTriangles = 64;
while (p.allocTriangles < p.numTriangles)
p.allocTriangles <<= 1; // Must be power of two.
// Allocate output tensors.
torch::Tensor out = color.detach().clone(); // Use color as base.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
torch::Tensor work_buffer = torch::empty({p.n * p.width * p.height * 8 + 4}, opts); // 8 int for a maximum of two work items per pixel.
p.output = out.data_ptr<float>();
p.workBuffer = (int4*)(work_buffer.data_ptr<float>());
// Clear the work counters.
NVDR_CHECK_CUDA_ERROR(cudaMemsetAsync(p.workBuffer, 0, sizeof(int4), stream));
// Verify that buffers are aligned to allow float2/float4 operations.
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.rasterOut & 7), "raster_out input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.workBuffer & 15), "work_buffer internal tensor not aligned to int4");
NVDR_CHECK(!((uintptr_t)p.evHash & 15), "topology_hash internal tensor not aligned to int4");
// Choose launch parameters for the discontinuity finder kernel and launch.
void* args[] = {&p};
dim3 blockSize(AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH, AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT, 1);
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.n);
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasFwdDiscontinuityKernel, gridSize, blockSize, args, 0, stream));
// Determine optimum block size for the persistent analysis kernel and launch.
int device = 0;
int numCTA = 0;
int numSM = 0;
NVDR_CHECK_CUDA_ERROR(cudaGetDevice(&device));
NVDR_CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numCTA, (void*)AntialiasFwdAnalysisKernel, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK, 0));
NVDR_CHECK_CUDA_ERROR(cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, device));
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasFwdAnalysisKernel, numCTA * numSM, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK, args, 0, stream));
// Return results.
return std::tuple<torch::Tensor, torch::Tensor>(out, work_buffer);
}
//------------------------------------------------------------------------
// Gradient op.
std::tuple<torch::Tensor, torch::Tensor> antialias_grad(torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, torch::Tensor dy, torch::Tensor work_buffer)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(color));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AntialiasKernelParams p = {}; // Initialize all fields to zero.
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
// Check inputs.
NVDR_CHECK_DEVICE(color, rast, pos, tri, dy, work_buffer);
NVDR_CHECK_CONTIGUOUS(color, rast, pos, tri, work_buffer);
NVDR_CHECK_F32(color, rast, pos, dy, work_buffer);
NVDR_CHECK_I32(tri);
// Sanity checks.
NVDR_CHECK(dy.sizes().size() == 4 && dy.size(0) > 0 && dy.size(1) > 0 && dy.size(2) > 0 && dy.size(3) > 0, "dy must have shape[>0, >0, >0, >0]");
NVDR_CHECK(color.sizes().size() == 4 && color.size(0) > 0 && color.size(1) > 0 && color.size(2) > 0 && color.size(3) > 0, "color must have shape[>0, >0, >0, >0]");
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "raster_out must have shape[>0, >0, >0, 4]");
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
NVDR_CHECK(color.size(1) == rast.size(1) && color.size(2) == rast.size(2), "color and raster_out inputs must have same spatial dimensions");
NVDR_CHECK(color.size(1) == dy.size(1) && color.size(2) == dy.size(2) && color.size(3) == dy.size(3), "color and dy inputs must have same dimensions");
if (p.instance_mode)
{
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) > 0 && pos.size(1) > 0 && pos.size(2) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
NVDR_CHECK(rast.size(0) == color.size(0) && pos.size(0) == color.size(0), "minibatch size mismatch between inputs color, raster_out, pos");
NVDR_CHECK(dy.size(0) == color.size(0) && rast.size(0) == color.size(0) && pos.size(0) ==color.size(0), "minibatch size mismatch between inputs dy, color, raster_out, pos");
}
else
{
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
NVDR_CHECK(rast.size(0) == color.size(0), "minibatch size mismatch between inputs color, raster_out");
NVDR_CHECK(dy.size(0) == color.size(0) && rast.size(0) == color.size(0), "minibatch size mismatch between inputs dy, color, raster_out");
}
// Extract input dimensions.
p.numVertices = pos.size(p.instance_mode ? 1 : 0);
p.numTriangles = tri.size(0);
p.n = color.size(0);
p.height = color.size(1);
p.width = color.size(2);
p.channels = color.size(3);
// Ensure dy is contiguous.
torch::Tensor dy_ = dy.contiguous();
// Get input pointers.
p.color = color.data_ptr<float>();
p.rasterOut = rast.data_ptr<float>();
p.tri = tri.data_ptr<int>();
p.pos = pos.data_ptr<float>();
p.dy = dy_.data_ptr<float>();
p.workBuffer = (int4*)(work_buffer.data_ptr<float>());
// Misc parameters.
p.xh = .5f * (float)p.width;
p.yh = .5f * (float)p.height;
// Allocate output tensors.
torch::Tensor grad_color = dy_.detach().clone(); // Use dy as base.
torch::Tensor grad_pos = torch::zeros_like(pos);
p.gradColor = grad_color.data_ptr<float>();
p.gradPos = grad_pos.data_ptr<float>();
// Clear gradient kernel work counter.
NVDR_CHECK_CUDA_ERROR(cudaMemsetAsync(&p.workBuffer[0].y, 0, sizeof(int), stream));
// Verify that buffers are aligned to allow float2/float4 operations.
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.workBuffer & 15), "work_buffer internal tensor not aligned to int4");
// Determine optimum block size for the gradient kernel and launch.
void* args[] = {&p};
int device = 0;
int numCTA = 0;
int numSM = 0;
NVDR_CHECK_CUDA_ERROR(cudaGetDevice(&device));
NVDR_CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numCTA, (void*)AntialiasGradKernel, AA_GRAD_KERNEL_THREADS_PER_BLOCK, 0));
NVDR_CHECK_CUDA_ERROR(cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, device));
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasGradKernel, numCTA * numSM, AA_GRAD_KERNEL_THREADS_PER_BLOCK, args, 0, stream));
// Return results.
return std::tuple<torch::Tensor, torch::Tensor>(grad_color, grad_pos);
}
//------------------------------------------------------------------------
third_party/nvdiffrast/csrc/torch/torch_bindings.cpp
Vendored
+73
View File
@@ -0,0 +1,73 @@
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
#include "torch_common.inl"
#include "torch_types.h"
#include <tuple>
//------------------------------------------------------------------------
// Op prototypes. Return type macros for readability.
#define OP_RETURN_T torch::Tensor
#define OP_RETURN_TT std::tuple<torch::Tensor, torch::Tensor>
#define OP_RETURN_TTT std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
#define OP_RETURN_TTTT std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
#define OP_RETURN_TTV std::tuple<torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >
#define OP_RETURN_TTTTV std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >
OP_RETURN_TT rasterize_fwd_cuda (RasterizeCRStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple<int, int> resolution, torch::Tensor ranges, int peeling_idx);
OP_RETURN_T rasterize_grad (torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy);
OP_RETURN_T rasterize_grad_db (torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy, torch::Tensor ddb);
OP_RETURN_TT interpolate_fwd (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri);
OP_RETURN_TT interpolate_fwd_da (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor rast_db, bool diff_attrs_all, std::vector<int>& diff_attrs_vec);
OP_RETURN_TT interpolate_grad (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy);
OP_RETURN_TTT interpolate_grad_da (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy, torch::Tensor rast_db, torch::Tensor dda, bool diff_attrs_all, std::vector<int>& diff_attrs_vec);
TextureMipWrapper texture_construct_mip (torch::Tensor tex, int max_mip_level, bool cube_mode);
OP_RETURN_T texture_fwd (torch::Tensor tex, torch::Tensor uv, int filter_mode, int boundary_mode);
OP_RETURN_T texture_fwd_mip (torch::Tensor tex, torch::Tensor uv, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode);
OP_RETURN_T texture_grad_nearest (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode);
OP_RETURN_TT texture_grad_linear (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode);
OP_RETURN_TTV texture_grad_linear_mipmap_nearest (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode);
OP_RETURN_TTTTV texture_grad_linear_mipmap_linear (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode);
TopologyHashWrapper antialias_construct_topology_hash (torch::Tensor tri);
OP_RETURN_TT antialias_fwd (torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, TopologyHashWrapper topology_hash);
OP_RETURN_TT antialias_grad (torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, torch::Tensor dy, torch::Tensor work_buffer);
//------------------------------------------------------------------------
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// State classes.
pybind11::class_<RasterizeCRStateWrapper>(m, "RasterizeCRStateWrapper").def(pybind11::init<int>());
pybind11::class_<TextureMipWrapper>(m, "TextureMipWrapper").def(pybind11::init<>());
pybind11::class_<TopologyHashWrapper>(m, "TopologyHashWrapper");
// Plumbing to torch/c10 logging system.
m.def("get_log_level", [](void) { return FLAGS_caffe2_log_level; }, "get log level");
m.def("set_log_level", [](int level){ FLAGS_caffe2_log_level = level; }, "set log level");
// Ops.
m.def("rasterize_fwd_cuda", &rasterize_fwd_cuda, "rasterize forward op (cuda)");
m.def("rasterize_grad", &rasterize_grad, "rasterize gradient op ignoring db gradients");
m.def("rasterize_grad_db", &rasterize_grad_db, "rasterize gradient op with db gradients");
m.def("interpolate_fwd", &interpolate_fwd, "interpolate forward op with attribute derivatives");
m.def("interpolate_fwd_da", &interpolate_fwd_da, "interpolate forward op without attribute derivatives");
m.def("interpolate_grad", &interpolate_grad, "interpolate gradient op with attribute derivatives");
m.def("interpolate_grad_da", &interpolate_grad_da, "interpolate gradient op without attribute derivatives");
m.def("texture_construct_mip", &texture_construct_mip, "texture mipmap construction");
m.def("texture_fwd", &texture_fwd, "texture forward op without mipmapping");
m.def("texture_fwd_mip", &texture_fwd_mip, "texture forward op with mipmapping");
m.def("texture_grad_nearest", &texture_grad_nearest, "texture gradient op in nearest mode");
m.def("texture_grad_linear", &texture_grad_linear, "texture gradient op in linear mode");
m.def("texture_grad_linear_mipmap_nearest", &texture_grad_linear_mipmap_nearest, "texture gradient op in linear-mipmap-nearest mode");
m.def("texture_grad_linear_mipmap_linear", &texture_grad_linear_mipmap_linear, "texture gradient op in linear-mipmap-linear mode");
m.def("antialias_construct_topology_hash", &antialias_construct_topology_hash, "antialias topology hash construction");
m.def("antialias_fwd", &antialias_fwd, "antialias forward op");
m.def("antialias_grad", &antialias_grad, "antialias gradient op");
}
//------------------------------------------------------------------------
third_party/nvdiffrast/csrc/torch/torch_interpolate.cpp
Vendored
+250
View File
@@ -0,0 +1,250 @@
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
#include "torch_common.inl"
#include "../common/common.h"
#include "../common/interpolate.h"
//------------------------------------------------------------------------
// Kernel prototypes.
void InterpolateFwdKernel (const InterpolateKernelParams p);
void InterpolateFwdKernelDa (const InterpolateKernelParams p);
void InterpolateGradKernel (const InterpolateKernelParams p);
void InterpolateGradKernelDa(const InterpolateKernelParams p);
//------------------------------------------------------------------------
// Helper
static void set_diff_attrs(InterpolateKernelParams& p, bool diff_attrs_all, std::vector<int>& diff_attrs_vec)
{
if (diff_attrs_all)
{
p.numDiffAttr = p.numAttr;
p.diff_attrs_all = 1;
}
else
{
NVDR_CHECK(diff_attrs_vec.size() <= IP_MAX_DIFF_ATTRS, "too many entries in diff_attrs list (increase IP_MAX_DIFF_ATTRS)");
p.numDiffAttr = diff_attrs_vec.size();
memcpy(p.diffAttrs, &diff_attrs_vec[0], diff_attrs_vec.size()*sizeof(int));
}
}
//------------------------------------------------------------------------
// Forward op.
std::tuple<torch::Tensor, torch::Tensor> interpolate_fwd_da(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor rast_db, bool diff_attrs_all, std::vector<int>& diff_attrs_vec)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(attr));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
InterpolateKernelParams p = {}; // Initialize all fields to zero.
bool enable_da = (rast_db.defined()) && (diff_attrs_all || !diff_attrs_vec.empty());
p.instance_mode = (attr.sizes().size() > 2) ? 1 : 0;
// Check inputs.
if (enable_da)
{
NVDR_CHECK_DEVICE(attr, rast, tri, rast_db);
NVDR_CHECK_CONTIGUOUS(attr, rast, tri, rast_db);
NVDR_CHECK_F32(attr, rast, rast_db);
NVDR_CHECK_I32(tri);
}
else
{
NVDR_CHECK_DEVICE(attr, rast, tri);
NVDR_CHECK_CONTIGUOUS(attr, rast, tri);
NVDR_CHECK_F32(attr, rast);
NVDR_CHECK_I32(tri);
}
// Sanity checks.
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "rast must have shape[>0, >0, >0, 4]");
NVDR_CHECK( tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
NVDR_CHECK((attr.sizes().size() == 2 || attr.sizes().size() == 3) && attr.size(0) > 0 && attr.size(1) > 0 && (attr.sizes().size() == 2 || attr.size(2) > 0), "attr must have shape [>0, >0, >0] or [>0, >0]");
if (p.instance_mode)
NVDR_CHECK(attr.size(0) == rast.size(0) || attr.size(0) == 1, "minibatch size mismatch between inputs rast, attr");
if (enable_da)
{
NVDR_CHECK(rast_db.sizes().size() == 4 && rast_db.size(0) > 0 && rast_db.size(1) > 0 && rast_db.size(2) > 0 && rast_db.size(3) == 4, "rast_db must have shape[>0, >0, >0, 4]");
NVDR_CHECK(rast_db.size(1) == rast.size(1) && rast_db.size(2) == rast.size(2), "spatial size mismatch between inputs rast and rast_db");
NVDR_CHECK(rast_db.size(0) == rast.size(0), "minibatch size mismatch between inputs rast, rast_db");
}
// Extract input dimensions.
p.numVertices = attr.size(p.instance_mode ? 1 : 0);
p.numAttr = attr.size(p.instance_mode ? 2 : 1);
p.numTriangles = tri.size(0);
p.height = rast.size(1);
p.width = rast.size(2);
p.depth = rast.size(0);
// Set attribute pixel differential info if enabled, otherwise leave as zero.
if (enable_da)
set_diff_attrs(p, diff_attrs_all, diff_attrs_vec);
else
p.numDiffAttr = 0;
// Get input pointers.
p.attr = attr.data_ptr<float>();
p.rast = rast.data_ptr<float>();
p.tri = tri.data_ptr<int>();
p.rastDB = enable_da ? rast_db.data_ptr<float>() : NULL;
p.attrBC = (p.instance_mode && attr.size(0) == 1) ? 1 : 0;
// Allocate output tensors.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
torch::Tensor out = torch::empty({p.depth, p.height, p.width, p.numAttr}, opts);
torch::Tensor out_da = torch::empty({p.depth, p.height, p.width, p.numDiffAttr * 2}, opts);
p.out = out.data_ptr<float>();
p.outDA = enable_da ? out_da.data_ptr<float>() : NULL;
// Verify that buffers are aligned to allow float2/float4 operations.
NVDR_CHECK(!((uintptr_t)p.rast & 15), "rast input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.rastDB & 15), "rast_db input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.outDA & 7), "out_da output tensor not aligned to float2");
// Choose launch parameters.
dim3 blockSize = getLaunchBlockSize(IP_FWD_MAX_KERNEL_BLOCK_WIDTH, IP_FWD_MAX_KERNEL_BLOCK_HEIGHT, p.width, p.height);
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.depth);
// Launch CUDA kernel.
void* args[] = {&p};
void* func = enable_da ? (void*)InterpolateFwdKernelDa : (void*)InterpolateFwdKernel;
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream));
// Return results.
return std::tuple<torch::Tensor, torch::Tensor>(out, out_da);
}
// Version without derivatives.
std::tuple<torch::Tensor, torch::Tensor> interpolate_fwd(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri)
{
std::vector<int> empty_vec;
torch::Tensor empty_tensor;
return interpolate_fwd_da(attr, rast, tri, empty_tensor, false, empty_vec);
}
//------------------------------------------------------------------------
// Gradient op.
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> interpolate_grad_da(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy, torch::Tensor rast_db, torch::Tensor dda, bool diff_attrs_all, std::vector<int>& diff_attrs_vec)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(attr));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
InterpolateKernelParams p = {}; // Initialize all fields to zero.
bool enable_da = (rast_db.defined()) && (diff_attrs_all || !diff_attrs_vec.empty());
p.instance_mode = (attr.sizes().size() > 2) ? 1 : 0;
// Check inputs.
if (enable_da)
{
NVDR_CHECK_DEVICE(attr, rast, tri, dy, rast_db, dda);
NVDR_CHECK_CONTIGUOUS(attr, rast, tri, rast_db);
NVDR_CHECK_F32(attr, rast, dy, rast_db, dda);
NVDR_CHECK_I32(tri);
}
else
{
NVDR_CHECK_DEVICE(attr, rast, tri, dy);
NVDR_CHECK_CONTIGUOUS(attr, rast, tri);
NVDR_CHECK_F32(attr, rast, dy);
NVDR_CHECK_I32(tri);
}
// Depth of attributes.
int attr_depth = p.instance_mode ? (attr.sizes().size() > 1 ? attr.size(0) : 0) : 1;
// Sanity checks.
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "rast must have shape[>0, >0, >0, 4]");
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
NVDR_CHECK((attr.sizes().size() == 2 || attr.sizes().size() == 3) && attr.size(0) > 0 && attr.size(1) > 0 && (attr.sizes().size() == 2 || attr.size(2) > 0), "attr must have shape [>0, >0, >0] or [>0, >0]");
NVDR_CHECK(dy.sizes().size() == 4 && dy.size(0) > 0 && dy.size(1) == rast.size(1) && dy.size(2) == rast.size(2) && dy.size(3) > 0, "dy must have shape [>0, height, width, >0]");
NVDR_CHECK(dy.size(3) == attr.size(attr.sizes().size() - 1), "argument count mismatch between inputs dy, attr");
NVDR_CHECK((attr_depth == rast.size(0) || attr_depth == 1) && dy.size(0) == rast.size(0), "minibatch size mismatch between inputs rast, dy, attr");
if (enable_da)
{
NVDR_CHECK(dda.sizes().size() == 4 && dda.size(0) > 0 && dda.size(1) == rast.size(1) && dda.size(2) == rast.size(2), "dda must have shape [>0, height, width, ?]");
NVDR_CHECK(dda.size(0) == rast.size(0), "minibatch size mismatch between rast, dda");
NVDR_CHECK(rast_db.sizes().size() == 4 && rast_db.size(0) > 0 && rast_db.size(1) > 0 && rast_db.size(2) > 0 && rast_db.size(3) == 4, "rast_db must have shape[>0, >0, >0, 4]");
NVDR_CHECK(rast_db.size(1) == rast.size(1) && rast_db.size(2) == rast.size(2), "spatial size mismatch between inputs rast and rast_db");
NVDR_CHECK(rast_db.size(0) == rast.size(0), "minibatch size mismatch between inputs rast, rast_db");
}
// Extract input dimensions.
p.numVertices = attr.size(p.instance_mode ? 1 : 0);
p.numAttr = attr.size(p.instance_mode ? 2 : 1);
p.numTriangles = tri.size(0);
p.height = rast.size(1);
p.width = rast.size(2);
p.depth = rast.size(0);
// Ensure gradients are contiguous.
torch::Tensor dy_ = dy.contiguous();
torch::Tensor dda_;
if (enable_da)
dda_ = dda.contiguous();
// Set attribute pixel differential info if enabled, otherwise leave as zero.
if (enable_da)
set_diff_attrs(p, diff_attrs_all, diff_attrs_vec);
else
p.numDiffAttr = 0;
// Get input pointers.
p.attr = attr.data_ptr<float>();
p.rast = rast.data_ptr<float>();
p.tri = tri.data_ptr<int>();
p.dy = dy_.data_ptr<float>();
p.rastDB = enable_da ? rast_db.data_ptr<float>() : NULL;
p.dda = enable_da ? dda_.data_ptr<float>() : NULL;
p.attrBC = (p.instance_mode && attr_depth < p.depth) ? 1 : 0;
// Allocate output tensors.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
torch::Tensor gradAttr = torch::zeros_like(attr);
torch::Tensor gradRaster = torch::empty_like(rast);
torch::Tensor gradRasterDB;
if (enable_da)
gradRasterDB = torch::empty_like(rast_db);
p.gradAttr = gradAttr.data_ptr<float>();
p.gradRaster = gradRaster.data_ptr<float>();
p.gradRasterDB = enable_da ? gradRasterDB.data_ptr<float>() : NULL;
// Verify that buffers are aligned to allow float2/float4 operations.
NVDR_CHECK(!((uintptr_t)p.rast & 15), "rast input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.rastDB & 15), "rast_db input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.dda & 7), "dda input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.gradRaster & 15), "grad_rast output tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.gradRasterDB & 15), "grad_rast_db output tensor not aligned to float4");
// Choose launch parameters.
dim3 blockSize = getLaunchBlockSize(IP_GRAD_MAX_KERNEL_BLOCK_WIDTH, IP_GRAD_MAX_KERNEL_BLOCK_HEIGHT, p.width, p.height);
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.depth);
// Launch CUDA kernel.
void* args[] = {&p};
void* func = enable_da ? (void*)InterpolateGradKernelDa : (void*)InterpolateGradKernel;
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream));
// Return results.
return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>(gradAttr, gradRaster, gradRasterDB);
}
// Version without derivatives.
std::tuple<torch::Tensor, torch::Tensor> interpolate_grad(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy)
{
std::vector<int> empty_vec;
torch::Tensor empty_tensor;
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> result = interpolate_grad_da(attr, rast, tri, dy, empty_tensor, empty_tensor, false, empty_vec);
return std::tuple<torch::Tensor, torch::Tensor>(std::get<0>(result), std::get<1>(result));
}
//------------------------------------------------------------------------
third_party/nvdiffrast/csrc/torch/torch_rasterize.cpp
Vendored
+265
View File
@@ -0,0 +1,265 @@
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
#include "torch_common.inl"
#include "torch_types.h"
#include "../common/common.h"
#include "../common/rasterize.h"
#include "../common/cudaraster/CudaRaster.hpp"
#include "../common/cudaraster/impl/Constants.hpp"
#include <tuple>
//------------------------------------------------------------------------
// Kernel prototypes.
void RasterizeCudaFwdShaderKernel(const RasterizeCudaFwdShaderParams p);
void RasterizeGradKernel(const RasterizeGradParams p);
void RasterizeGradKernelDb(const RasterizeGradParams p);
//------------------------------------------------------------------------
// Python CudaRaster state wrapper methods.
RasterizeCRStateWrapper::RasterizeCRStateWrapper(int cudaDeviceIdx_)
{
const at::cuda::OptionalCUDAGuard device_guard(cudaDeviceIdx_);
cudaDeviceIdx = cudaDeviceIdx_;
cr = new CR::CudaRaster();
}
RasterizeCRStateWrapper::~RasterizeCRStateWrapper(void)
{
const at::cuda::OptionalCUDAGuard device_guard(cudaDeviceIdx);
delete cr;
}
//------------------------------------------------------------------------
// Forward op (Cuda).
std::tuple<torch::Tensor, torch::Tensor> rasterize_fwd_cuda(RasterizeCRStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple<int, int> resolution, torch::Tensor ranges, int peeling_idx)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(pos));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
CR::CudaRaster* cr = stateWrapper.cr;
// Check inputs.
NVDR_CHECK_DEVICE(pos, tri);
NVDR_CHECK_CPU(ranges);
NVDR_CHECK_CONTIGUOUS(pos, tri, ranges);
NVDR_CHECK_F32(pos);
NVDR_CHECK_I32(tri, ranges);
// Check that CudaRaster context was created for the correct GPU.
NVDR_CHECK(pos.get_device() == stateWrapper.cudaDeviceIdx, "CudaRaster context must must reside on the same device as input tensors");
// Determine instance mode and check input dimensions.
bool instance_mode = pos.sizes().size() > 2;
if (instance_mode)
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) > 0 && pos.size(1) > 0 && pos.size(2) == 4, "instance mode - pos must have shape [>0, >0, 4]");
else
{
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "range mode - pos must have shape [>0, 4]");
NVDR_CHECK(ranges.sizes().size() == 2 && ranges.size(0) > 0 && ranges.size(1) == 2, "range mode - ranges must have shape [>0, 2]");
}
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
// Get output shape.
int height_out = std::get<0>(resolution);
int width_out = std::get<1>(resolution);
int depth = instance_mode ? pos.size(0) : ranges.size(0); // Depth of tensor, not related to depth buffering.
NVDR_CHECK(height_out > 0 && width_out > 0, "resolution must be [>0, >0]");
// Round internal resolution up to tile size.
int height = (height_out + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
int width = (width_out + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
// Get position and triangle buffer sizes in vertices / triangles.
int posCount = instance_mode ? pos.size(1) : pos.size(0);
int triCount = tri.size(0);
// Set up CudaRaster buffers.
const float* posPtr = pos.data_ptr<float>();
const int32_t* rangesPtr = instance_mode ? 0 : ranges.data_ptr<int32_t>(); // This is in CPU memory.
const int32_t* triPtr = tri.data_ptr<int32_t>();
cr->setVertexBuffer((void*)posPtr, posCount);
cr->setIndexBuffer((void*)triPtr, triCount);
cr->setBufferSize(width_out, height_out, depth);
// Enable depth peeling?
bool enablePeel = (peeling_idx > 0);
cr->setRenderModeFlags(enablePeel ? CR::CudaRaster::RenderModeFlag_EnableDepthPeeling : 0); // No backface culling.
if (enablePeel)
cr->swapDepthAndPeel(); // Use previous depth buffer as peeling depth input.
// Determine viewport tiling.
int tileCountX = (width + CR_MAXVIEWPORT_SIZE - 1) / CR_MAXVIEWPORT_SIZE;
int tileCountY = (height + CR_MAXVIEWPORT_SIZE - 1) / CR_MAXVIEWPORT_SIZE;
int tileSizeX = ((width + tileCountX - 1) / tileCountX + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
int tileSizeY = ((height + tileCountY - 1) / tileCountY + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
TORCH_CHECK(tileCountX > 0 && tileCountY > 0 && tileSizeX > 0 && tileSizeY > 0, "internal error in tile size calculation: count or size is zero");
TORCH_CHECK(tileSizeX <= CR_MAXVIEWPORT_SIZE && tileSizeY <= CR_MAXVIEWPORT_SIZE, "internal error in tile size calculation: tile larger than allowed");
TORCH_CHECK((tileSizeX & (CR_TILE_SIZE - 1)) == 0 && (tileSizeY & (CR_TILE_SIZE - 1)) == 0, "internal error in tile size calculation: tile not divisible by ", CR_TILE_SIZE);
TORCH_CHECK(tileCountX * tileSizeX >= width && tileCountY * tileSizeY >= height, "internal error in tile size calculation: tiles do not cover viewport");
// Rasterize in tiles.
for (int tileY = 0; tileY < tileCountY; tileY++)
for (int tileX = 0; tileX < tileCountX; tileX++)
{
// Set CudaRaster viewport according to tile.
int offsetX = tileX * tileSizeX;
int offsetY = tileY * tileSizeY;
int sizeX = (width_out - offsetX) < tileSizeX ? (width_out - offsetX) : tileSizeX;
int sizeY = (height_out - offsetY) < tileSizeY ? (height_out - offsetY) : tileSizeY;
cr->setViewport(sizeX, sizeY, offsetX, offsetY);
// Run all triangles in one batch. In case of error, the workload could be split into smaller batches - maybe do that in the future.
// Only enable peeling-specific optimizations to skip first stages when image fits in one tile. Those are not valid otherwise.
cr->deferredClear(0u);
bool success = cr->drawTriangles(rangesPtr, enablePeel && (tileCountX == 1 && tileCountY == 1), stream);
NVDR_CHECK(success, "subtriangle count overflow");
}
// Allocate output tensors.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
torch::Tensor out = torch::empty({depth, height_out, width_out, 4}, opts);
torch::Tensor out_db = torch::empty({depth, height_out, width_out, 4}, opts);
// Populate pixel shader kernel parameters.
RasterizeCudaFwdShaderParams p;
p.pos = posPtr;
p.tri = triPtr;
p.in_idx = (const int*)cr->getColorBuffer();
p.out = out.data_ptr<float>();
p.out_db = out_db.data_ptr<float>();
p.numTriangles = triCount;
p.numVertices = posCount;
p.width_in = width;
p.height_in = height;
p.width_out = width_out;
p.height_out = height_out;
p.depth = depth;
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
p.xs = 2.f / (float)width_out;
p.xo = 1.f / (float)width_out - 1.f;
p.ys = 2.f / (float)height_out;
p.yo = 1.f / (float)height_out - 1.f;
// Verify that buffers are aligned to allow float2/float4 operations.
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.out & 15), "out output tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.out_db & 15), "out_db output tensor not aligned to float4");
// Choose launch parameters.
dim3 blockSize = getLaunchBlockSize(RAST_CUDA_FWD_SHADER_KERNEL_BLOCK_WIDTH, RAST_CUDA_FWD_SHADER_KERNEL_BLOCK_HEIGHT, p.width_out, p.height_out);
dim3 gridSize = getLaunchGridSize(blockSize, p.width_out, p.height_out, p.depth);
// Launch CUDA kernel.
void* args[] = {&p};
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)RasterizeCudaFwdShaderKernel, gridSize, blockSize, args, 0, stream));
// Return.
return std::tuple<torch::Tensor, torch::Tensor>(out, out_db);
}
//------------------------------------------------------------------------
// Gradient op.
torch::Tensor rasterize_grad_db(torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy, torch::Tensor ddb)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(pos));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
RasterizeGradParams p;
bool enable_db = ddb.defined();
// Check inputs.
if (enable_db)
{
NVDR_CHECK_DEVICE(pos, tri, out, dy, ddb);
NVDR_CHECK_CONTIGUOUS(pos, tri, out);
NVDR_CHECK_F32(pos, out, dy, ddb);
NVDR_CHECK_I32(tri);
}
else
{
NVDR_CHECK_DEVICE(pos, tri, out, dy);
NVDR_CHECK_CONTIGUOUS(pos, tri, out);
NVDR_CHECK_F32(pos, out, dy);
NVDR_CHECK_I32(tri);
}
// Determine instance mode.
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
// Shape is taken from the rasterizer output tensor.
NVDR_CHECK(out.sizes().size() == 4, "tensor out must be rank-4");
p.depth = out.size(0);
p.height = out.size(1);
p.width = out.size(2);
NVDR_CHECK(p.depth > 0 && p.height > 0 && p.width > 0, "resolution must be [>0, >0, >0]");
// Check other shapes.
if (p.instance_mode)
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) == p.depth && pos.size(1) > 0 && pos.size(2) == 4, "pos must have shape [depth, >0, 4]");
else
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "pos must have shape [>0, 4]");
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
NVDR_CHECK(out.sizes().size() == 4 && out.size(0) == p.depth && out.size(1) == p.height && out.size(2) == p.width && out.size(3) == 4, "out must have shape [depth, height, width, 4]");
NVDR_CHECK( dy.sizes().size() == 4 && dy.size(0) == p.depth && dy.size(1) == p.height && dy.size(2) == p.width && dy.size(3) == 4, "dy must have shape [depth, height, width, 4]");
if (enable_db)
NVDR_CHECK(ddb.sizes().size() == 4 && ddb.size(0) == p.depth && ddb.size(1) == p.height && ddb.size(2) == p.width && ddb.size(3) == 4, "ddb must have shape [depth, height, width, 4]");
// Ensure gradients are contiguous.
torch::Tensor dy_ = dy.contiguous();
torch::Tensor ddb_;
if (enable_db)
ddb_ = ddb.contiguous();
// Populate parameters.
p.numTriangles = tri.size(0);
p.numVertices = p.instance_mode ? pos.size(1) : pos.size(0);
p.pos = pos.data_ptr<float>();
p.tri = tri.data_ptr<int>();
p.out = out.data_ptr<float>();
p.dy = dy_.data_ptr<float>();
p.ddb = enable_db ? ddb_.data_ptr<float>() : NULL;
// Set up pixel position to clip space x, y transform.
p.xs = 2.f / (float)p.width;
p.xo = 1.f / (float)p.width - 1.f;
p.ys = 2.f / (float)p.height;
p.yo = 1.f / (float)p.height - 1.f;
// Allocate output tensor for position gradients.
torch::Tensor grad = torch::zeros_like(pos);
p.grad = grad.data_ptr<float>();
// Verify that buffers are aligned to allow float2/float4 operations.
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.dy & 7), "dy input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.ddb & 15), "ddb input tensor not aligned to float4");
// Choose launch parameters.
dim3 blockSize = getLaunchBlockSize(RAST_GRAD_MAX_KERNEL_BLOCK_WIDTH, RAST_GRAD_MAX_KERNEL_BLOCK_HEIGHT, p.width, p.height);
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.depth);
// Launch CUDA kernel.
void* args[] = {&p};
void* func = enable_db ? (void*)RasterizeGradKernelDb : (void*)RasterizeGradKernel;
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream));
// Return the gradients.
return grad;
}
// Version without derivatives.
torch::Tensor rasterize_grad(torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy)
{
torch::Tensor empty_tensor;
return rasterize_grad_db(pos, tri, out, dy, empty_tensor);
}
//------------------------------------------------------------------------
third_party/nvdiffrast/csrc/torch/torch_texture.cpp
Vendored
+718
View File
@@ -0,0 +1,718 @@
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
#include "torch_common.inl"
#include "torch_types.h"
#include "../common/common.h"
#include "../common/texture.h"
#include <cuda_runtime.h>
//------------------------------------------------------------------------
// Kernel prototypes.
void MipBuildKernel1 (const TextureKernelParams p);
void MipBuildKernel2 (const TextureKernelParams p);
void MipBuildKernel4 (const TextureKernelParams p);
void TextureFwdKernelNearest1 (const TextureKernelParams p);
void TextureFwdKernelNearest2 (const TextureKernelParams p);
void TextureFwdKernelNearest4 (const TextureKernelParams p);
void TextureFwdKernelLinear1 (const TextureKernelParams p);
void TextureFwdKernelLinear2 (const TextureKernelParams p);
void TextureFwdKernelLinear4 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapNearest1 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapNearest2 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapNearest4 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapLinear1 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapLinear2 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapLinear4 (const TextureKernelParams p);
void TextureFwdKernelCubeNearest1 (const TextureKernelParams p);
void TextureFwdKernelCubeNearest2 (const TextureKernelParams p);
void TextureFwdKernelCubeNearest4 (const TextureKernelParams p);
void TextureFwdKernelCubeLinear1 (const TextureKernelParams p);
void TextureFwdKernelCubeLinear2 (const TextureKernelParams p);
void TextureFwdKernelCubeLinear4 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapNearest1 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapNearest2 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapNearest4 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapLinear1 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapLinear2 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapLinear4 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapNearestBO1 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapNearestBO2 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapNearestBO4 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapLinearBO1 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapLinearBO2 (const TextureKernelParams p);
void TextureFwdKernelLinearMipmapLinearBO4 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapNearestBO1 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapNearestBO2 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapNearestBO4 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapLinearBO1 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapLinearBO2 (const TextureKernelParams p);
void TextureFwdKernelCubeLinearMipmapLinearBO4 (const TextureKernelParams p);
void MipGradKernel1 (const TextureKernelParams p);
void MipGradKernel2 (const TextureKernelParams p);
void MipGradKernel4 (const TextureKernelParams p);
void TextureGradKernelNearest (const TextureKernelParams p);
void TextureGradKernelLinear (const TextureKernelParams p);
void TextureGradKernelLinearMipmapNearest (const TextureKernelParams p);
void TextureGradKernelLinearMipmapLinear (const TextureKernelParams p);
void TextureGradKernelCubeNearest (const TextureKernelParams p);
void TextureGradKernelCubeLinear (const TextureKernelParams p);
void TextureGradKernelCubeLinearMipmapNearest (const TextureKernelParams p);
void TextureGradKernelCubeLinearMipmapLinear (const TextureKernelParams p);
void TextureGradKernelLinearMipmapNearestBO (const TextureKernelParams p);
void TextureGradKernelLinearMipmapLinearBO (const TextureKernelParams p);
void TextureGradKernelCubeLinearMipmapNearestBO (const TextureKernelParams p);
void TextureGradKernelCubeLinearMipmapLinearBO (const TextureKernelParams p);
//------------------------------------------------------------------------
// Modeselektor.
static void set_modes(TextureKernelParams& p, int filter_mode, int boundary_mode, int max_mip_level)
{
// Mip and filter modes.
p.filterMode = filter_mode;
NVDR_CHECK(p.filterMode >= 0 && p.filterMode < TEX_MODE_COUNT, "filter_mode unsupported");
p.enableMip = (p.filterMode == TEX_MODE_LINEAR_MIPMAP_NEAREST || p.filterMode == TEX_MODE_LINEAR_MIPMAP_LINEAR);
// Mip level clamp.
if (p.enableMip)
{
p.mipLevelLimit = max_mip_level;
NVDR_CHECK(p.mipLevelLimit >= -1, "invalid max_mip_level");
}
// Boundary mode.
p.boundaryMode = boundary_mode;
NVDR_CHECK(p.boundaryMode >= 0 && p.boundaryMode < TEX_BOUNDARY_MODE_COUNT, "boundary_mode unsupported");
}
//------------------------------------------------------------------------
// Mipmap construction.
TextureMipWrapper texture_construct_mip(torch::Tensor tex, int max_mip_level, bool cube_mode)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(tex));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
TextureKernelParams p = {}; // Initialize all fields to zero.
p.mipLevelLimit = max_mip_level;
p.boundaryMode = cube_mode ? TEX_BOUNDARY_MODE_CUBE : TEX_BOUNDARY_MODE_WRAP;
NVDR_CHECK(p.mipLevelLimit >= -1, "invalid max_mip_level");
// Check inputs.
NVDR_CHECK_DEVICE(tex);
NVDR_CHECK_CONTIGUOUS(tex);
NVDR_CHECK_F32(tex);
// Populate parameters and sanity check tex shape.
if (!cube_mode)
{
NVDR_CHECK(tex.sizes().size() == 4 && tex.size(0) > 0 && tex.size(1) > 0 && tex.size(2) > 0 && tex.size(3) > 0, "tex must have shape[>0, >0, >0, >0]");
}
else
{
NVDR_CHECK(tex.sizes().size() == 5 && tex.size(0) > 0 && tex.size(1) == 6 && tex.size(2) > 0 && tex.size(3) > 0 && tex.size(4) > 0, "tex must have shape[>0, 6, >0, >0, >0] in cube map mode");
NVDR_CHECK(tex.size(2) == tex.size(3), "texture shape must be square in cube map mode");
}
p.texDepth = tex.size(0);
p.texHeight = tex.size(cube_mode ? 2 : 1);
p.texWidth = tex.size(cube_mode ? 3 : 2);
p.channels = tex.size(cube_mode ? 4 : 3);
// Set texture pointer.
p.tex[0] = tex.data_ptr<float>();
// Generate mip offsets and calculate total size.
int mipOffsets[TEX_MAX_MIP_LEVEL];
int mipTotal = calculateMipInfo(p, mipOffsets);
// Allocate and set mip tensor.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
torch::Tensor mip = torch::empty({mipTotal}, opts);
float* pmip = mip.data_ptr<float>();
for (int i=1; i <= p.mipLevelMax; i++)
p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels.
// Choose kernel variants based on channel count.
void* args[] = {&p};
int channel_div_idx = 0;
if (!(p.channels & 3))
channel_div_idx = 2; // Channel count divisible by 4.
else if (!(p.channels & 1))
channel_div_idx = 1; // Channel count divisible by 2.
// Build mip levels.
for (int i=1; i <= p.mipLevelMax; i++)
{
int2 ms = mipLevelSize(p, i);
int3 sz = make_int3(ms.x, ms.y, p.texDepth);
dim3 blockSize = getLaunchBlockSize(TEX_FWD_MAX_MIP_KERNEL_BLOCK_WIDTH, TEX_FWD_MAX_MIP_KERNEL_BLOCK_HEIGHT, sz.x, sz.y);
dim3 gridSize = getLaunchGridSize(blockSize, sz.x, sz.y, sz.z * (cube_mode ? 6 : 1));
p.mipLevelOut = i;
void* build_func_tbl[3] = { (void*)MipBuildKernel1, (void*)MipBuildKernel2, (void*)MipBuildKernel4 };
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(build_func_tbl[channel_div_idx], gridSize, blockSize, args, 0, stream));
}
// Return the mip tensor in a wrapper.
TextureMipWrapper mip_wrapper;
mip_wrapper.mip = mip;
mip_wrapper.max_mip_level = max_mip_level;
mip_wrapper.texture_size = tex.sizes().vec();
mip_wrapper.cube_mode = cube_mode;
return mip_wrapper;
}
//------------------------------------------------------------------------
// Forward op.
torch::Tensor texture_fwd_mip(torch::Tensor tex, torch::Tensor uv, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(tex));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
TextureKernelParams p = {}; // Initialize all fields to zero.
bool has_mip_stack = (mip_stack.size() > 0);
torch::Tensor& mip_w = mip_wrapper.mip; // Unwrap.
int max_mip_level = has_mip_stack ? mip_stack.size() : mip_wrapper.max_mip_level;
set_modes(p, filter_mode, boundary_mode, max_mip_level);
// See if we have these tensors or not.
bool has_uv_da = uv_da.defined() && uv_da.nbytes();
bool has_mip_level_bias = mip_level_bias.defined() && mip_level_bias.nbytes();
if (p.enableMip)
{
NVDR_CHECK(has_uv_da || has_mip_level_bias, "mipmapping filter mode requires uv_da and/or mip_level_bias input");
NVDR_CHECK(has_mip_stack || mip_w.defined(), "mipmapping filter mode requires mip wrapper or mip stack input");
}
// Check inputs.
NVDR_CHECK_DEVICE(tex, uv);
NVDR_CHECK_CONTIGUOUS(tex, uv);
NVDR_CHECK_F32(tex, uv);
if (p.enableMip)
{
if (has_mip_stack)
{
TORCH_CHECK(at::cuda::check_device(mip_stack), __func__, "(): Mip stack inputs must reside on the correct GPU device");
nvdr_check_contiguous(mip_stack, __func__, "(): Mip stack inputs must be contiguous tensors");
nvdr_check_f32(mip_stack, __func__, "(): Mip stack inputs must be float32 tensors");
}
else
{
NVDR_CHECK_DEVICE(mip_w);
NVDR_CHECK_CONTIGUOUS(mip_w);
NVDR_CHECK_F32(mip_w);
}
if (has_uv_da)
{
NVDR_CHECK_DEVICE(uv_da);
NVDR_CHECK_CONTIGUOUS(uv_da);
NVDR_CHECK_F32(uv_da);
}
if (has_mip_level_bias)
{
NVDR_CHECK_DEVICE(mip_level_bias);
NVDR_CHECK_CONTIGUOUS(mip_level_bias);
NVDR_CHECK_F32(mip_level_bias);
}
}
// Sanity checks and state setters.
bool cube_mode = (boundary_mode == TEX_BOUNDARY_MODE_CUBE);
if (!cube_mode)
{
NVDR_CHECK(tex.sizes().size() == 4 && tex.size(0) > 0 && tex.size(1) > 0 && tex.size(2) > 0 && tex.size(3) > 0, "tex must have shape[>0, >0, >0, >0]");
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 2, "uv must have shape [>0, >0, >0, 2]");
p.texHeight = tex.size(1);
p.texWidth = tex.size(2);
p.channels = tex.size(3);
}
else
{
NVDR_CHECK(tex.sizes().size() == 5 && tex.size(0) > 0 && tex.size(1) == 6 && tex.size(2) > 0 && tex.size(3) > 0 && tex.size(4) > 0, "tex must have shape[>0, 6, >0, >0, >0] in cube map mode");
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 3, "uv must have shape [>0, >0, >0, 3] in cube map mode");
NVDR_CHECK(tex.size(2) == tex.size(3), "texture shape must be square in cube map mode");
p.texHeight = tex.size(2);
p.texWidth = tex.size(3);
p.channels = tex.size(4);
}
NVDR_CHECK(tex.size(0) == 1 || tex.size(0) == uv.size(0), "minibatch size mismatch between inputs tex, uv");
NVDR_CHECK(p.texWidth <= (1 << TEX_MAX_MIP_LEVEL) && p.texHeight <= (1 << TEX_MAX_MIP_LEVEL), "texture size too large");
p.n = uv.size(0);
p.imgHeight = uv.size(1);
p.imgWidth = uv.size(2);
p.texDepth = tex.size(0);
if (p.enableMip)
{
if (has_uv_da)
{
if (!cube_mode)
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 4, "uv_da must have shape [minibatch_size, height, width, 4]");
else
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 6, "uv_da must have shape [minibatch_size, height, width, 6] in cube map mode");
}
if (has_mip_level_bias)
NVDR_CHECK(mip_level_bias.sizes().size() == 3 && mip_level_bias.size(0) == p.n && mip_level_bias.size(1) == p.imgHeight && mip_level_bias.size(2) == p.imgWidth, "mip_level_bias must have shape [minibatch_size, height, width]");
}
// Get input pointers.
p.tex[0] = tex.data_ptr<float>();
p.uv = uv.data_ptr<float>();
p.uvDA = (p.enableMip && has_uv_da) ? uv_da.data_ptr<float>() : NULL;
p.mipLevelBias = (p.enableMip && has_mip_level_bias) ? mip_level_bias.data_ptr<float>() : NULL;
// Allocate output tensor.
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
torch::Tensor out = torch::empty({p.n, p.imgHeight, p.imgWidth, p.channels}, opts);
p.out = out.data_ptr<float>();
// Choose kernel variants based on channel count.
void* args[] = {&p};
int channel_div_idx = 0;
if (!(p.channels & 3))
channel_div_idx = 2; // Channel count divisible by 4.
else if (!(p.channels & 1))
channel_div_idx = 1; // Channel count divisible by 2.
// Mip-related setup.
float* pmip = 0;
if (p.enableMip)
{
if (has_mip_stack)
{
// Custom mip stack supplied. Check that sizes match and assign.
p.mipLevelMax = max_mip_level;
for (int i=1; i <= p.mipLevelMax; i++)
{
torch::Tensor& t = mip_stack[i-1];
int2 sz = mipLevelSize(p, i);
if (!cube_mode)
NVDR_CHECK(t.sizes().size() == 4 && t.size(0) == tex.size(0) && t.size(1) == sz.y && t.size(2) == sz.x && t.size(3) == p.channels, "mip level size mismatch in custom mip stack");
else
NVDR_CHECK(t.sizes().size() == 5 && t.size(0) == tex.size(0) && t.size(1) == 6 && t.size(2) == sz.y && t.size(3) == sz.x && t.size(4) == p.channels, "mip level size mismatch in mip stack");
if (sz.x == 1 && sz.y == 1)
NVDR_CHECK(i == p.mipLevelMax, "mip level size mismatch in mip stack");
p.tex[i] = t.data_ptr<float>();
}
}
else
{
// Generate mip offsets, check mipmap size, and set mip data pointer.
int mipOffsets[TEX_MAX_MIP_LEVEL];
int mipTotal = calculateMipInfo(p, mipOffsets);
NVDR_CHECK(tex.sizes() == mip_wrapper.texture_size && cube_mode == mip_wrapper.cube_mode, "mip does not match texture size");
NVDR_CHECK(mip_w.sizes().size() == 1 && mip_w.size(0) == mipTotal, "wrapped mip tensor size mismatch");
pmip = mip_w.data_ptr<float>();
for (int i=1; i <= p.mipLevelMax; i++)
p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels.
}
}
// Verify that buffers are aligned to allow float2/float4 operations. Unused pointers are zero so always aligned.
if (!cube_mode)
NVDR_CHECK(!((uintptr_t)p.uv & 7), "uv input tensor not aligned to float2");
if ((p.channels & 3) == 0)
{
for (int i=0; i <= p.mipLevelMax; i++)
NVDR_CHECK(!((uintptr_t)p.tex[i] & 15), "tex or mip input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.out & 15), "out output tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)pmip & 15), "mip input tensor not aligned to float4");
}
if ((p.channels & 1) == 0)
{
for (int i=0; i <= p.mipLevelMax; i++)
NVDR_CHECK(!((uintptr_t)p.tex[i] & 7), "tex or mip input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.out & 7), "out output tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)pmip & 7), "mip input tensor not aligned to float2");
}
if (!cube_mode)
NVDR_CHECK(!((uintptr_t)p.uvDA & 15), "uv_da input tensor not aligned to float4");
else
NVDR_CHECK(!((uintptr_t)p.uvDA & 7), "uv_da input tensor not aligned to float2");
// Choose launch parameters for texture lookup kernel.
dim3 blockSize = getLaunchBlockSize(TEX_FWD_MAX_KERNEL_BLOCK_WIDTH, TEX_FWD_MAX_KERNEL_BLOCK_HEIGHT, p.imgWidth, p.imgHeight);
dim3 gridSize = getLaunchGridSize(blockSize, p.imgWidth, p.imgHeight, p.n);
// Choose kernel based on filter mode, cube mode, bias-only mode, and datatype.
void* func_tbl[TEX_MODE_COUNT * 2 * 2 * 3] = {
(void*)TextureFwdKernelNearest1,
(void*)TextureFwdKernelNearest2,
(void*)TextureFwdKernelNearest4,
(void*)TextureFwdKernelLinear1,
(void*)TextureFwdKernelLinear2,
(void*)TextureFwdKernelLinear4,
(void*)TextureFwdKernelLinearMipmapNearest1,
(void*)TextureFwdKernelLinearMipmapNearest2,
(void*)TextureFwdKernelLinearMipmapNearest4,
(void*)TextureFwdKernelLinearMipmapLinear1,
(void*)TextureFwdKernelLinearMipmapLinear2,
(void*)TextureFwdKernelLinearMipmapLinear4,
(void*)TextureFwdKernelCubeNearest1,
(void*)TextureFwdKernelCubeNearest2,
(void*)TextureFwdKernelCubeNearest4,
(void*)TextureFwdKernelCubeLinear1,
(void*)TextureFwdKernelCubeLinear2,
(void*)TextureFwdKernelCubeLinear4,
(void*)TextureFwdKernelCubeLinearMipmapNearest1,
(void*)TextureFwdKernelCubeLinearMipmapNearest2,
(void*)TextureFwdKernelCubeLinearMipmapNearest4,
(void*)TextureFwdKernelCubeLinearMipmapLinear1,
(void*)TextureFwdKernelCubeLinearMipmapLinear2,
(void*)TextureFwdKernelCubeLinearMipmapLinear4,
NULL,
NULL,
NULL,
NULL,
NULL,
NULL,
(void*)TextureFwdKernelLinearMipmapNearestBO1,
(void*)TextureFwdKernelLinearMipmapNearestBO2,
(void*)TextureFwdKernelLinearMipmapNearestBO4,
(void*)TextureFwdKernelLinearMipmapLinearBO1,
(void*)TextureFwdKernelLinearMipmapLinearBO2,
(void*)TextureFwdKernelLinearMipmapLinearBO4,
NULL,
NULL,
NULL,
NULL,
NULL,
NULL,
(void*)TextureFwdKernelCubeLinearMipmapNearestBO1,
(void*)TextureFwdKernelCubeLinearMipmapNearestBO2,
(void*)TextureFwdKernelCubeLinearMipmapNearestBO4,
(void*)TextureFwdKernelCubeLinearMipmapLinearBO1,
(void*)TextureFwdKernelCubeLinearMipmapLinearBO2,
(void*)TextureFwdKernelCubeLinearMipmapLinearBO4,
};
// Function index.
int func_idx = p.filterMode;
if (cube_mode)
func_idx += TEX_MODE_COUNT; // Cube variant.
if (p.enableMip && !has_uv_da)
func_idx += TEX_MODE_COUNT * 2; // Bias-only variant.
func_idx = func_idx * 3 + channel_div_idx; // Choose vector size.
// Launch kernel.
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func_tbl[func_idx], gridSize, blockSize, args, 0, stream));
// Return output tensor.
return out;
}
// Version without mipmaps.
torch::Tensor texture_fwd(torch::Tensor tex, torch::Tensor uv, int filter_mode, int boundary_mode)
{
torch::Tensor empty_tensor;
std::vector<torch::Tensor> empty_vector;
return texture_fwd_mip(tex, uv, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode);
}
//------------------------------------------------------------------------
// Gradient op.
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > texture_grad_linear_mipmap_linear(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode)
{
const at::cuda::OptionalCUDAGuard device_guard(device_of(tex));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
TextureKernelParams p = {}; // Initialize all fields to zero.
bool has_mip_stack = (mip_stack.size() > 0);
torch::Tensor& mip_w = mip_wrapper.mip; // Unwrap.
int max_mip_level = has_mip_stack ? mip_stack.size() : mip_wrapper.max_mip_level;
set_modes(p, filter_mode, boundary_mode, max_mip_level);
// See if we have these tensors or not.
bool has_uv_da = uv_da.defined() && uv_da.nbytes();
bool has_mip_level_bias = mip_level_bias.defined() && mip_level_bias.nbytes();
if (p.enableMip)
{
NVDR_CHECK(has_uv_da || has_mip_level_bias, "mipmapping filter mode requires uv_da and/or mip_level_bias input");
NVDR_CHECK(has_mip_stack || mip_w.defined(), "mipmapping filter mode requires mip wrapper or mip stack input");
}
// Check inputs.
NVDR_CHECK_DEVICE(tex, uv);
NVDR_CHECK_CONTIGUOUS(tex, uv);
NVDR_CHECK_F32(tex, uv);
if (p.enableMip)
{
if (has_mip_stack)
{
TORCH_CHECK(at::cuda::check_device(mip_stack), __func__, "(): Mip stack inputs must reside on the correct GPU device");
nvdr_check_contiguous(mip_stack, __func__, "(): Mip stack inputs must be contiguous tensors");
nvdr_check_f32(mip_stack, __func__, "(): Mip stack inputs must be float32 tensors");
}
else
{
NVDR_CHECK_DEVICE(mip_w);
NVDR_CHECK_CONTIGUOUS(mip_w);
NVDR_CHECK_F32(mip_w);
}
if (has_uv_da)
{
NVDR_CHECK_DEVICE(uv_da);
NVDR_CHECK_CONTIGUOUS(uv_da);
NVDR_CHECK_F32(uv_da);
}
if (has_mip_level_bias)
{
NVDR_CHECK_DEVICE(mip_level_bias);
NVDR_CHECK_CONTIGUOUS(mip_level_bias);
NVDR_CHECK_F32(mip_level_bias);
}
}
// Sanity checks and state setters.
bool cube_mode = (boundary_mode == TEX_BOUNDARY_MODE_CUBE);
if (!cube_mode)
{
NVDR_CHECK(tex.sizes().size() == 4 && tex.size(0) > 0 && tex.size(1) > 0 && tex.size(2) > 0 && tex.size(3) > 0, "tex must have shape[>0, >0, >0, >0]");
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 2, "uv must have shape [>0, >0, >0, 2]");
p.texHeight = tex.size(1);
p.texWidth = tex.size(2);
p.channels = tex.size(3);
}
else
{
NVDR_CHECK(tex.sizes().size() == 5 && tex.size(0) > 0 && tex.size(1) == 6 && tex.size(2) > 0 && tex.size(3) > 0 && tex.size(4) > 0, "tex must have shape[>0, 6, >0, >0, >0] in cube map mode");
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 3, "uv must have shape [>0, >0, >0, 3] in cube map mode");
NVDR_CHECK(tex.size(2) == tex.size(3), "texture shape must be square in cube map mode");
p.texHeight = tex.size(2);
p.texWidth = tex.size(3);
p.channels = tex.size(4);
}
NVDR_CHECK(tex.size(0) == 1 || tex.size(0) == uv.size(0), "minibatch size mismatch between inputs tex, uv");
NVDR_CHECK(p.texWidth <= (1 << TEX_MAX_MIP_LEVEL) && p.texHeight <= (1 << TEX_MAX_MIP_LEVEL), "texture size too large");
p.n = uv.size(0);
p.imgHeight = uv.size(1);
p.imgWidth = uv.size(2);
p.texDepth = tex.size(0);
if (p.enableMip)
{
if (has_uv_da)
{
if (!cube_mode)
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 4, "uv_da must have shape [minibatch_size, height, width, 4]");
else
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 6, "uv_da must have shape [minibatch_size, height, width, 6] in cube map mode");
}
if (has_mip_level_bias)
NVDR_CHECK(mip_level_bias.sizes().size() == 3 && mip_level_bias.size(0) == p.n && mip_level_bias.size(1) == p.imgHeight && mip_level_bias.size(2) == p.imgWidth, "mip_level_bias must have shape [minibatch_size, height, width]");
}
NVDR_CHECK(dy.sizes().size() == 4 && dy.size(0) == p.n && dy.size(1) == p.imgHeight && dy.size(2) == p.imgWidth && dy.size(3) == p.channels, "dy must have shape [minibatch_size, height, width, channels]");
// Get contiguous version of dy.
torch::Tensor dy_ = dy.contiguous();
// Get input pointers.
p.tex[0] = tex.data_ptr<float>();
p.uv = uv.data_ptr<float>();
p.dy = dy_.data_ptr<float>();
p.uvDA = (p.enableMip && has_uv_da) ? uv_da.data_ptr<float>() : NULL;
p.mipLevelBias = (p.enableMip && has_mip_level_bias) ? mip_level_bias.data_ptr<float>() : NULL;
// Allocate output tensor for tex gradient.
torch::Tensor grad_tex = torch::zeros_like(tex);
p.gradTex[0] = grad_tex.data_ptr<float>();
// Allocate output tensor for uv gradient.
torch::Tensor grad_uv;
torch::Tensor grad_uv_da;
torch::Tensor grad_mip_level_bias;
if (p.filterMode != TEX_MODE_NEAREST)
{
grad_uv = torch::empty_like(uv);
p.gradUV = grad_uv.data_ptr<float>();
// Gradients for things affecting mip level.
if (p.filterMode == TEX_MODE_LINEAR_MIPMAP_LINEAR)
{
// Allocate output tensor for uv_da gradient.
if (has_uv_da)
{
grad_uv_da = torch::empty_like(uv_da);
p.gradUVDA = grad_uv_da.data_ptr<float>();
}
// Allocate output tensor for mip_level_bias gradient.
if (has_mip_level_bias)
{
grad_mip_level_bias = torch::empty_like(mip_level_bias);
p.gradMipLevelBias = grad_mip_level_bias.data_ptr<float>();
}
}
}
// Choose kernel variants based on channel count.
int channel_div_idx = 0;
if (!(p.channels & 3))
channel_div_idx = 2; // Channel count divisible by 4.
else if (!(p.channels & 1))
channel_div_idx = 1; // Channel count divisible by 2.
// Mip-related setup.
torch::Tensor grad_mip;
std::vector<torch::Tensor> grad_mip_stack;
float* pmip = 0;
float* pgradMip = 0;
if (p.enableMip)
{
if (has_mip_stack)
{
// Custom mip stack supplied. Check that sizes match, assign, construct gradient tensors.
p.mipLevelMax = max_mip_level;
for (int i=1; i <= p.mipLevelMax; i++)
{
torch::Tensor& t = mip_stack[i-1];
int2 sz = mipLevelSize(p, i);
if (!cube_mode)
NVDR_CHECK(t.sizes().size() == 4 && t.size(0) == tex.size(0) && t.size(1) == sz.y && t.size(2) == sz.x && t.size(3) == p.channels, "mip level size mismatch in mip stack");
else
NVDR_CHECK(t.sizes().size() == 5 && t.size(0) == tex.size(0) && t.size(1) == 6 && t.size(2) == sz.y && t.size(3) == sz.x && t.size(4) == p.channels, "mip level size mismatch in mip stack");
if (sz.x == 1 && sz.y == 1)
NVDR_CHECK(i == p.mipLevelMax, "mip level size mismatch in mip stack");
torch::Tensor g = torch::zeros_like(t);
grad_mip_stack.push_back(g);
p.tex[i] = t.data_ptr<float>();
p.gradTex[i] = g.data_ptr<float>();
}
}
else
{
// Generate mip offsets and get space for temporary mip gradients.
int mipOffsets[TEX_MAX_MIP_LEVEL];
int mipTotal = calculateMipInfo(p, mipOffsets);
NVDR_CHECK(tex.sizes() == mip_wrapper.texture_size && cube_mode == mip_wrapper.cube_mode, "mip does not match texture size");
NVDR_CHECK(mip_w.sizes().size() == 1 && mip_w.size(0) == mipTotal, "mip tensor size mismatch");
grad_mip = torch::zeros_like(mip_w);
pmip = (float*)mip_w.data_ptr<float>();
pgradMip = grad_mip.data_ptr<float>();
for (int i=1; i <= p.mipLevelMax; i++)
{
p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels.
p.gradTex[i] = pgradMip + mipOffsets[i]; // Pointers to mip gradients.
}
}
}
// Verify that buffers are aligned to allow float2/float4 operations. Unused pointers are zero so always aligned.
if (!cube_mode)
{
NVDR_CHECK(!((uintptr_t)p.uv & 7), "uv input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.gradUV & 7), "grad_uv output tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.uvDA & 15), "uv_da input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.gradUVDA & 15), "grad_uv_da output tensor not aligned to float4");
}
else
{
NVDR_CHECK(!((uintptr_t)p.uvDA & 7), "uv_da input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.gradUVDA & 7), "grad_uv_da output tensor not aligned to float2");
}
if ((p.channels & 3) == 0)
{
for (int i=0; i <= p.mipLevelMax; i++)
{
NVDR_CHECK(!((uintptr_t)p.tex[i] & 15), "tex or mip input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)p.gradTex[i] & 15), "grad_tex output tensor not aligned to float4");
}
NVDR_CHECK(!((uintptr_t)p.dy & 15), "dy input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)pmip & 15), "mip input tensor not aligned to float4");
NVDR_CHECK(!((uintptr_t)pgradMip & 15), "internal mip gradient tensor not aligned to float4");
}
if ((p.channels & 1) == 0)
{
for (int i=0; i <= p.mipLevelMax; i++)
{
NVDR_CHECK(!((uintptr_t)p.tex[i] & 7), "tex or mip input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)p.gradTex[i] & 7), "grad_tex output tensor not aligned to float2");
}
NVDR_CHECK(!((uintptr_t)p.dy & 7), "dy output tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)pmip & 7), "mip input tensor not aligned to float2");
NVDR_CHECK(!((uintptr_t)pgradMip & 7), "internal mip gradient tensor not aligned to float2");
}
// Choose launch parameters for main gradient kernel.
void* args[] = {&p};
dim3 blockSize = getLaunchBlockSize(TEX_GRAD_MAX_KERNEL_BLOCK_WIDTH, TEX_GRAD_MAX_KERNEL_BLOCK_HEIGHT, p.imgWidth, p.imgHeight);
dim3 gridSize = getLaunchGridSize(blockSize, p.imgWidth, p.imgHeight, p.n);
void* func_tbl[TEX_MODE_COUNT * 2 * 2] = {
(void*)TextureGradKernelNearest,
(void*)TextureGradKernelLinear,
(void*)TextureGradKernelLinearMipmapNearest,
(void*)TextureGradKernelLinearMipmapLinear,
(void*)TextureGradKernelCubeNearest,
(void*)TextureGradKernelCubeLinear,
(void*)TextureGradKernelCubeLinearMipmapNearest,
(void*)TextureGradKernelCubeLinearMipmapLinear,
NULL,
NULL,
(void*)TextureGradKernelLinearMipmapNearestBO,
(void*)TextureGradKernelLinearMipmapLinearBO,
NULL,
NULL,
(void*)TextureGradKernelCubeLinearMipmapNearestBO,
(void*)TextureGradKernelCubeLinearMipmapLinearBO,
};
// Function index.
int func_idx = p.filterMode;
if (cube_mode)
func_idx += TEX_MODE_COUNT; // Cube variant.
if (p.enableMip && !has_uv_da)
func_idx += TEX_MODE_COUNT * 2; // Bias-only variant.
// Launch main gradient kernel.
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func_tbl[func_idx], gridSize, blockSize, args, 0, stream));
// Launch kernel to pull gradients from mip levels. Don't do this if mip stack was supplied - individual level gradients are already there.
if (p.enableMip && !has_mip_stack)
{
dim3 blockSize = getLaunchBlockSize(TEX_GRAD_MAX_MIP_KERNEL_BLOCK_WIDTH, TEX_GRAD_MAX_MIP_KERNEL_BLOCK_HEIGHT, p.texWidth, p.texHeight);
dim3 gridSize = getLaunchGridSize(blockSize, p.texWidth, p.texHeight, p.texDepth * (cube_mode ? 6 : 1));
int sharedBytes = blockSize.x * blockSize.y * p.channels * sizeof(float);
void* mip_grad_func_tbl[3] = { (void*)MipGradKernel1, (void*)MipGradKernel2, (void*)MipGradKernel4 };
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(mip_grad_func_tbl[channel_div_idx], gridSize, blockSize, args, sharedBytes, stream));
}
// Return output tensors.
return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >(grad_tex, grad_uv, grad_uv_da, grad_mip_level_bias, grad_mip_stack);
}
// Version for nearest filter mode.
torch::Tensor texture_grad_nearest(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode)
{
torch::Tensor empty_tensor;
std::vector<torch::Tensor> empty_vector;
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > result = texture_grad_linear_mipmap_linear(tex, uv, dy, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode);
return std::get<0>(result);
}
// Version for linear filter mode.
std::tuple<torch::Tensor, torch::Tensor> texture_grad_linear(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode)
{
torch::Tensor empty_tensor;
std::vector<torch::Tensor> empty_vector;
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > result = texture_grad_linear_mipmap_linear(tex, uv, dy, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode);
return std::tuple<torch::Tensor, torch::Tensor>(std::get<0>(result), std::get<1>(result));
}
// Version for linear-mipmap-nearest mode.
std::tuple<torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > texture_grad_linear_mipmap_nearest(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode)
{
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > result = texture_grad_linear_mipmap_linear(tex, uv, dy, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode, boundary_mode);
return std::tuple<torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >(std::get<0>(result), std::get<1>(result), std::get<4>(result));
}
//------------------------------------------------------------------------
third_party/nvdiffrast/csrc/torch/torch_types.h
Vendored
+47
View File
@@ -0,0 +1,47 @@
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
#include "torch_common.inl"
//------------------------------------------------------------------------
// Python CudaRaster state wrapper.
namespace CR { class CudaRaster; }
class RasterizeCRStateWrapper
{
public:
RasterizeCRStateWrapper (int cudaDeviceIdx);
~RasterizeCRStateWrapper (void);
CR::CudaRaster* cr;
int cudaDeviceIdx;
};
//------------------------------------------------------------------------
// Mipmap wrapper to prevent intrusion from Python side.
class TextureMipWrapper
{
public:
torch::Tensor mip;
int max_mip_level;
std::vector<int64_t> texture_size; // For error checking.
bool cube_mode; // For error checking.
};
//------------------------------------------------------------------------
// Antialias topology hash wrapper to prevent intrusion from Python side.
class TopologyHashWrapper
{
public:
torch::Tensor ev_hash;
};
//------------------------------------------------------------------------