Update to 2026-09-17 pipeline snapshot; add weights, L20 assets and recording via Git LFS
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
@@ -0,0 +1,20 @@
|
||||
/*
|
||||
* Copyright (c) 2023, 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/extension.h>
|
||||
#include "common.h"
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("sampleRaysUniformOccupiedVoxels", &sampleRaysUniformOccupiedVoxels);
|
||||
m.def("postprocessOctreeRayTracing", &postprocessOctreeRayTracing);
|
||||
m.def("rayColorToTextureImageCUDA", &rayColorToTextureImageCUDA);
|
||||
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
/*
|
||||
* Copyright (c) 2023, 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 <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <bits/stdc++.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/torch.h>
|
||||
#include <stdint.h>
|
||||
#include <cstdio>
|
||||
#include <stdint.h>
|
||||
#include <stdexcept>
|
||||
#include <limits>
|
||||
#include <set>
|
||||
#include "common.h"
|
||||
#include "Eigen/Dense"
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* @brief
|
||||
*
|
||||
* @tparam scalar_t
|
||||
* @param z_sampled
|
||||
* @param z_in_out
|
||||
* @param z_vals
|
||||
* @return __global__
|
||||
*/
|
||||
template <typename scalar_t>
|
||||
__global__ void sample_rays_uniform_occupied_voxels_kernel(const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> z_sampled, const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> z_in_out, torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> z_vals)
|
||||
{
|
||||
const int i_ray = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int i_sample = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (i_ray>=z_sampled.size(0)) return;
|
||||
if (i_sample>=z_sampled.size(1)) return;
|
||||
|
||||
int i_box = 0;
|
||||
float z_remain = z_sampled[i_ray][i_sample];
|
||||
auto z_in_out_cur_ray = z_in_out[i_ray];
|
||||
const float eps = 1e-4;
|
||||
const int max_n_box = z_in_out.size(1);
|
||||
|
||||
if (z_in_out_cur_ray[0][0]==0) return;
|
||||
|
||||
while (1)
|
||||
{
|
||||
if (i_box>=max_n_box)
|
||||
{
|
||||
if (z_remain<=eps)
|
||||
{
|
||||
z_vals[i_ray][i_sample] = z_in_out_cur_ray[max_n_box-1][1];
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("ERROR sample_rays_uniform_occupied_voxels_kernel: z_remain=%f, i_ray=%d, i_sample=%d, i_box=%d, z_in_out_cur_ray=(%f,%f)\n",z_remain,i_ray,i_sample,i_box,z_in_out_cur_ray[i_box][0],z_in_out_cur_ray[i_box][1]);
|
||||
for (int i=0;i<z_in_out.size(1);i++)
|
||||
{
|
||||
printf("z_in_out_cur_ray[%d]=(%f,%f)\n",i,z_in_out_cur_ray[i][0],z_in_out_cur_ray[i][1]);
|
||||
}
|
||||
while (1){};
|
||||
}
|
||||
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (z_in_out_cur_ray[i_box][0]==0)
|
||||
{
|
||||
if (z_remain<=eps && i_box>=1)
|
||||
{
|
||||
z_vals[i_ray][i_sample] = z_in_out_cur_ray[i_box-1][1];
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("ERROR sample_rays_uniform_occupied_voxels_kernel: z_remain=%f, i_ray=%d, i_sample=%d, i_box=%d, z_in_out_cur_ray=(%f,%f)\n",z_remain,i_ray,i_sample,i_box,z_in_out_cur_ray[i_box][0],z_in_out_cur_ray[i_box][1]);
|
||||
for (int i=0;i<z_in_out.size(1);i++)
|
||||
{
|
||||
printf("z_in_out_cur_ray[%d]=(%f,%f)\n",i,z_in_out_cur_ray[i][0],z_in_out_cur_ray[i][1]);
|
||||
}
|
||||
while (1){};
|
||||
}
|
||||
}
|
||||
|
||||
float box_len = z_in_out_cur_ray[i_box][1]-z_in_out_cur_ray[i_box][0];
|
||||
if (z_remain<=box_len)
|
||||
{
|
||||
z_vals[i_ray][i_sample] = z_in_out_cur_ray[i_box][0] + z_remain;
|
||||
return;
|
||||
}
|
||||
z_remain -= box_len;
|
||||
i_box++;
|
||||
}
|
||||
}
|
||||
|
||||
at::Tensor sampleRaysUniformOccupiedVoxels(const at::Tensor z_in_out, const at::Tensor z_sampled, at::Tensor z_vals)
|
||||
{
|
||||
CHECK_INPUT(z_in_out);
|
||||
CHECK_INPUT(z_sampled);
|
||||
CHECK_INPUT(z_vals);
|
||||
AT_ASSERTM(z_vals.sizes()==z_sampled.sizes());
|
||||
|
||||
const int N_rays = z_sampled.sizes()[0];
|
||||
const int N_samples = z_sampled.sizes()[1];
|
||||
const int threadx = 32;
|
||||
const int thready = 32;
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(z_in_out.type(), "sample_rays_uniform_occupied_voxels_kernel", ([&]
|
||||
{
|
||||
sample_rays_uniform_occupied_voxels_kernel<scalar_t><<<{divCeil(N_rays,threadx),divCeil(N_samples,thready)}, {threadx,thready}>>>(z_sampled.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),z_in_out.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),z_vals.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
|
||||
}));
|
||||
|
||||
return z_vals;
|
||||
}
|
||||
|
||||
template<class scalar_t>
|
||||
__global__ void postprocessOctreeRayTracingKernel(const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> ray_index, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> depth_in_out, const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> unique_intersect_ray_ids, const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> start_poss, torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> depths_in_out_padded)
|
||||
{
|
||||
const int unique_id_pos = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (unique_id_pos>=unique_intersect_ray_ids.size(0)) return;
|
||||
const int i_ray = unique_intersect_ray_ids[unique_id_pos];
|
||||
|
||||
int i_intersect = 0;
|
||||
auto cur_depths_in_out_padded = depths_in_out_padded[i_ray];
|
||||
for (int i=start_poss[unique_id_pos];i<ray_index.size(0);i++)
|
||||
{
|
||||
if (ray_index[i]!=i_ray) break;
|
||||
if (depth_in_out[i][0]==0 || depth_in_out[i][1]==0) break;
|
||||
if (depth_in_out[i][0]>depth_in_out[i][1]) continue;
|
||||
if (abs(depth_in_out[i][1]-depth_in_out[i][0])<1e-4) continue;
|
||||
|
||||
cur_depths_in_out_padded[i_intersect][0] = depth_in_out[i][0];
|
||||
cur_depths_in_out_padded[i_intersect][1] = depth_in_out[i][1];
|
||||
|
||||
i_intersect++;
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
at::Tensor postprocessOctreeRayTracing(const at::Tensor ray_index, const at::Tensor depth_in_out, const at::Tensor unique_intersect_ray_ids, const at::Tensor start_poss, const int max_intersections, const int N_rays)
|
||||
{
|
||||
CHECK_INPUT(ray_index);
|
||||
CHECK_INPUT(depth_in_out);
|
||||
CHECK_INPUT(start_poss);
|
||||
|
||||
const int n_unique_ids = unique_intersect_ray_ids.sizes()[0];
|
||||
at::Tensor depths_in_out_padded = at::zeros({N_rays,max_intersections,2}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, 0).requires_grad(false));
|
||||
dim3 threads = {256};
|
||||
dim3 blocks = {divCeil(n_unique_ids,threads.x)};
|
||||
AT_DISPATCH_FLOATING_TYPES(depth_in_out.type(), "postprocessOctreeRayTracingKernel", ([&]
|
||||
{
|
||||
postprocessOctreeRayTracingKernel<scalar_t><<<blocks,threads>>>(ray_index.packed_accessor32<long,1,torch::RestrictPtrTraits>(), depth_in_out.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), unique_intersect_ray_ids.packed_accessor32<long,1,torch::RestrictPtrTraits>(), start_poss.packed_accessor32<long,1,torch::RestrictPtrTraits>(), depths_in_out_padded.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
|
||||
}));
|
||||
|
||||
return depths_in_out_padded;
|
||||
}
|
||||
|
||||
|
||||
at::Tensor calculateBarycentricCoordinate3D(const at::Tensor &triangle, const at::Tensor &p)
|
||||
{
|
||||
auto vector_A = triangle[1]-triangle[2];
|
||||
auto vector_B = triangle[1]-triangle[0];
|
||||
auto normal = vector_A.cross(vector_B);
|
||||
auto areaABC = (normal * at::cross(triangle[1]-triangle[0], triangle[2]-triangle[0])).sum();
|
||||
auto areaPBC = (normal * at::cross(triangle[1]-p, triangle[2]-p)).sum();
|
||||
auto areaPCA = (normal * at::cross(triangle[2]-p, triangle[0]-p)).sum();
|
||||
at::Tensor ws = at::zeros({3}).to(torch::kFloat32);
|
||||
ws[0] = areaPBC / areaABC;
|
||||
ws[1] = areaPCA / areaABC;
|
||||
ws[2] = 1-ws[0]-ws[1];
|
||||
return ws;
|
||||
}
|
||||
|
||||
|
||||
__device__ Eigen::Vector3f calculateBarycentricCoordinate3DKernel(const Eigen::Matrix<float,3,3> &triangle, const Eigen::Vector3f &p)
|
||||
{
|
||||
Eigen::Vector3f vector_A = triangle.row(1)-triangle.row(2);
|
||||
Eigen::Vector3f vector_B = triangle.row(1)-triangle.row(0);
|
||||
Eigen::Vector3f normal = vector_A.cross(vector_B);
|
||||
float areaABC = (normal.array() * (triangle.row(1)-triangle.row(0)).cross(triangle.row(2)-triangle.row(0)).transpose().array()).sum();
|
||||
float areaPBC = (normal.array() * (triangle.row(1).transpose()-p).cross(triangle.row(2).transpose()-p).array()).sum();
|
||||
float areaPCA = (normal.array() * (triangle.row(2).transpose()-p).cross(triangle.row(0).transpose()-p).array()).sum();
|
||||
Eigen::Vector3f w;
|
||||
w[0] = areaPBC / areaABC;
|
||||
w[1] = areaPCA / areaABC;
|
||||
w[2] = 1-w[0]-w[1];
|
||||
return w;
|
||||
|
||||
}
|
||||
|
||||
|
||||
__device__ void calculateBarycentricCoordinate2DKernel(const Eigen::Matrix<float,3,2> &triangle, const Eigen::Vector2f &p, Eigen::Vector3f &w)
|
||||
{
|
||||
Eigen::Vector2f CA = triangle.row(0)-triangle.row(2);
|
||||
Eigen::Vector2f AC = -CA;
|
||||
Eigen::Vector2f CP = p-triangle.row(2).transpose();
|
||||
Eigen::Vector2f AB = triangle.row(1)-triangle.row(0);
|
||||
Eigen::Vector2f AP = p-triangle.row(0).transpose();
|
||||
Eigen::Matrix2f numerator, denominator;
|
||||
numerator << CA, CP;
|
||||
denominator<<AB, AC;
|
||||
w(1) = numerator.determinant()/denominator.determinant();
|
||||
numerator << AB, AP;
|
||||
w(2) = numerator.determinant()/denominator.determinant();
|
||||
w(0) = 1-w(1)-w(2);
|
||||
}
|
||||
|
||||
|
||||
|
||||
template<class scalar_t>
|
||||
__global__ void rayColorToTextureImageKernel(const torch::PackedTensorAccessor32<long,2,torch::RestrictPtrTraits> F, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> V, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> hit_locations, const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> hit_face_ids, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> uvs_tex, torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> uvs)
|
||||
{
|
||||
const int i_hit = blockIdx.x*blockDim.x + threadIdx.x;
|
||||
if (i_hit>=hit_locations.size(0)) return;
|
||||
|
||||
auto face = F[hit_face_ids[i_hit]];
|
||||
Eigen::Matrix3f tri_v = Eigen::Matrix3f::Zero();
|
||||
for (int r=0;r<3;r++)
|
||||
{
|
||||
for (int c=0;c<3;c++)
|
||||
{
|
||||
tri_v(r,c) = V[face[r]][c];
|
||||
}
|
||||
}
|
||||
auto location = hit_locations[i_hit];
|
||||
Eigen::Vector3f p(location[0], location[1], location[2]);
|
||||
auto w = calculateBarycentricCoordinate3DKernel(tri_v, p);
|
||||
|
||||
Eigen::Matrix<float,3,2> tri_uvs;
|
||||
for (int i=0;i<3;i++)
|
||||
{
|
||||
for (int j=0;j<2;j++)
|
||||
{
|
||||
tri_uvs(i,j) = uvs_tex[face[i]][j];
|
||||
}
|
||||
}
|
||||
|
||||
Eigen::Vector2f cur_uv = tri_uvs.transpose() * w;
|
||||
uvs[i_hit][0] = cur_uv(0);
|
||||
uvs[i_hit][1] = cur_uv(1);
|
||||
}
|
||||
|
||||
|
||||
|
||||
void rayColorToTextureImageCUDA(const at::Tensor &F, const at::Tensor &V, const at::Tensor &hit_locations, const at::Tensor &hit_face_ids, const at::Tensor &uvs_tex, at::Tensor &uvs)
|
||||
{
|
||||
CHECK_CONTIGUOUS(F);
|
||||
CHECK_CONTIGUOUS(V);
|
||||
CHECK_CONTIGUOUS(hit_locations);
|
||||
CHECK_CONTIGUOUS(hit_face_ids);
|
||||
CHECK_CONTIGUOUS(uvs_tex);
|
||||
|
||||
dim3 threads = {512};
|
||||
dim3 blocks = {divCeil(int(hit_locations.sizes()[0]),threads.x)};
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(V.type(), "rayColorToTextureImageKernel", ([&]
|
||||
{
|
||||
rayColorToTextureImageKernel<scalar_t><<<blocks,threads>>>(F.packed_accessor32<long,2,torch::RestrictPtrTraits>(), V.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), hit_locations.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), hit_face_ids.packed_accessor32<long,1,torch::RestrictPtrTraits>(), uvs_tex.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), uvs.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
|
||||
}));
|
||||
}
|
||||
|
||||
|
||||
|
||||
__device__ bool isLineIntersectLine(Eigen::Vector2f l1p1, Eigen::Vector2f l1p2, Eigen::Vector2f l2p1, Eigen::Vector2f l2p2)
|
||||
{
|
||||
float q = (l1p1(1) - l2p1(1)) * (l2p2(0) - l2p1(0)) - (l1p1(0) - l2p1(0)) * (l2p2(1) - l2p1(1));
|
||||
float d = (l1p2(0) - l1p1(0)) * (l2p2(1) - l2p1(1)) - (l1p2(1) - l1p1(1)) * (l2p2(0) - l2p1(0));
|
||||
|
||||
if ( d == 0 )
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
float r = q / d;
|
||||
|
||||
q = (l1p1(1) - l2p1(1)) * (l1p2(0) - l1p1(0)) - (l1p1(0) - l2p1(0)) * (l1p2(1) - l1p1(1));
|
||||
float s = q / d;
|
||||
|
||||
if( r < 0 || r > 1 || s < 0 || s > 1 )
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
__device__ bool isLineIntersectSquare(const Eigen::Vector2f &la, const Eigen::Vector2f &lb, const Eigen::Vector2f &p0, const Eigen::Vector2f &p1, const Eigen::Vector2f &p2, const Eigen::Vector2f &p3)
|
||||
{
|
||||
return isLineIntersectLine(la, lb, p0, p1) || isLineIntersectLine(la, lb, p1, p2) || isLineIntersectLine(la, lb, p2, p3) || isLineIntersectLine(la, lb, p3, p0);
|
||||
}
|
||||
|
||||
|
||||
__device__ bool isPixelInsideTriangle(const Eigen::Matrix<float,3,2> &triangle, Eigen::Vector2f p, Eigen::Vector3f &w)
|
||||
{
|
||||
calculateBarycentricCoordinate2DKernel(triangle, p, w);
|
||||
for (int j=0;j<3;j++)
|
||||
{
|
||||
if (w(j)<0) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
/*
|
||||
* Copyright (c) 2023, 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.
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
|
||||
|
||||
|
||||
inline int divCeil(int a, int b)
|
||||
{
|
||||
return (a+b-1)/b;
|
||||
};
|
||||
|
||||
|
||||
at::Tensor sampleRaysUniformOccupiedVoxels(const at::Tensor z_in_out, const at::Tensor z_sampled, at::Tensor z_vals);
|
||||
at::Tensor postprocessOctreeRayTracing(const at::Tensor ray_index, const at::Tensor depth_in_out, const at::Tensor unique_intersect_ray_ids, const at::Tensor start_poss, const int max_intersections, const int N_rays);
|
||||
void rayColorToTextureImageCUDA(const at::Tensor &F, const at::Tensor &V, const at::Tensor &hit_locations, const at::Tensor &hit_face_ids, const at::Tensor &uvs_tex, at::Tensor &uvs);
|
||||
@@ -0,0 +1,41 @@
|
||||
# Copyright (c) 2023, 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.
|
||||
|
||||
|
||||
from setuptools import setup
|
||||
import os,sys
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
|
||||
nvcc_flags = ['-Xcompiler', '-O3', '-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__']
|
||||
c_flags = ['-O3', '-std=c++14']
|
||||
|
||||
setup(
|
||||
name='common',
|
||||
extra_cflags=c_flags,
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
ext_modules=[
|
||||
CUDAExtension('common', [
|
||||
'bindings.cpp',
|
||||
'common.cu',
|
||||
],extra_compile_args={'gcc': c_flags, 'nvcc': nvcc_flags}),
|
||||
CUDAExtension('gridencoder', [
|
||||
f"{code_dir}/torch_ngp_grid_encoder/gridencoder.cu",
|
||||
f"{code_dir}/torch_ngp_grid_encoder/bindings.cpp",
|
||||
],extra_compile_args={'gcc': c_flags, 'nvcc': nvcc_flags}),
|
||||
],
|
||||
include_dirs=[
|
||||
"/usr/local/include/eigen3",
|
||||
"/usr/include/eigen3",
|
||||
],
|
||||
cmdclass={
|
||||
'build_ext': BuildExtension
|
||||
})
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2022 hawkey
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include "gridencoder.h"
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("grid_encode_forward", &grid_encode_forward, "grid_encode_forward (CUDA)");
|
||||
m.def("grid_encode_backward", &grid_encode_backward, "grid_encode_backward (CUDA)");
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
import numpy as np
|
||||
import os,sys,pdb
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(code_dir)
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.autograd import Function
|
||||
from torch.autograd.function import once_differentiable
|
||||
from torch.cuda.amp import custom_bwd, custom_fwd
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
from torch.utils.cpp_extension import load
|
||||
import os,sys
|
||||
import gridencoder
|
||||
|
||||
|
||||
_gridtype_to_id = {
|
||||
'hash': 0,
|
||||
'tiled': 1,
|
||||
}
|
||||
|
||||
class _grid_encode(Function):
|
||||
@staticmethod
|
||||
@custom_fwd
|
||||
def forward(ctx, inputs, embeddings, offsets, per_level_scale, base_resolution, calc_grad_inputs=False, gridtype=0, align_corners=False):
|
||||
# inputs: [B, D], float in [0, 1]
|
||||
# embeddings: [sO, C], float
|
||||
# offsets: [L + 1], int
|
||||
# RETURN: [B, F], float
|
||||
|
||||
inputs = inputs.contiguous()
|
||||
|
||||
B, D = inputs.shape
|
||||
L = offsets.shape[0] - 1
|
||||
C = embeddings.shape[1]
|
||||
S = np.log2(per_level_scale)
|
||||
H = base_resolution
|
||||
|
||||
if torch.is_autocast_enabled() and C % 2 == 0:
|
||||
embeddings = embeddings.to(torch.half)
|
||||
|
||||
outputs = torch.empty(L, B, C, device=inputs.device, dtype=embeddings.dtype)
|
||||
|
||||
if calc_grad_inputs:
|
||||
dy_dx = torch.empty(B, L * D * C, device=inputs.device, dtype=embeddings.dtype)
|
||||
else:
|
||||
dy_dx = torch.empty(1, device=inputs.device, dtype=embeddings.dtype) # placeholder... TODO: a better way?
|
||||
|
||||
gridencoder.grid_encode_forward(inputs, embeddings, offsets, outputs, B, D, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners)
|
||||
|
||||
outputs = outputs.permute(1, 0, 2).reshape(B, L * C)
|
||||
|
||||
ctx.save_for_backward(inputs, embeddings, offsets, dy_dx)
|
||||
ctx.dims = [B, D, C, L, S, H, gridtype]
|
||||
ctx.calc_grad_inputs = calc_grad_inputs
|
||||
ctx.align_corners = align_corners
|
||||
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
@custom_bwd
|
||||
def backward(ctx, grad):
|
||||
|
||||
inputs, embeddings, offsets, dy_dx = ctx.saved_tensors
|
||||
B, D, C, L, S, H, gridtype = ctx.dims
|
||||
calc_grad_inputs = ctx.calc_grad_inputs
|
||||
align_corners = ctx.align_corners
|
||||
|
||||
# grad: [B, L * C] --> [L, B, C]
|
||||
grad = grad.view(B, L, C).permute(1, 0, 2).contiguous()
|
||||
|
||||
grad_embeddings = torch.zeros_like(embeddings)
|
||||
|
||||
if calc_grad_inputs:
|
||||
grad_inputs = torch.zeros_like(inputs, dtype=embeddings.dtype)
|
||||
else:
|
||||
grad_inputs = torch.zeros(1, device=inputs.device, dtype=embeddings.dtype)
|
||||
|
||||
gridencoder.grid_encode_backward(grad, inputs, embeddings, offsets, grad_embeddings, B, D, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners)
|
||||
|
||||
if calc_grad_inputs:
|
||||
grad_inputs = grad_inputs.to(inputs.dtype)
|
||||
return grad_inputs, grad_embeddings, None, None, None, None, None, None
|
||||
else:
|
||||
return None, grad_embeddings, None, None, None, None, None, None
|
||||
|
||||
|
||||
grid_encode = _grid_encode.apply
|
||||
|
||||
|
||||
# https://github.com/ashawkey/torch-ngp/blob/main/gridencoder/grid.py
|
||||
class GridEncoder(nn.Module):
|
||||
def __init__(self, input_dim=3, n_levels=16, level_dim=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=None, gridtype='hash', align_corners=False):
|
||||
super().__init__()
|
||||
|
||||
per_level_scale = np.exp2(np.log2(desired_resolution / base_resolution) / (n_levels - 1))
|
||||
|
||||
self.input_dim = input_dim # coord dims, 2 or 3
|
||||
self.n_levels = n_levels # num levels, each level multiply resolution by 2
|
||||
self.level_dim = level_dim # encode channels per level
|
||||
self.per_level_scale = per_level_scale # multiply resolution by this scale at each level.
|
||||
self.log2_hashmap_size = log2_hashmap_size
|
||||
self.base_resolution = base_resolution
|
||||
self.out_dim = n_levels * level_dim
|
||||
self.gridtype = gridtype
|
||||
self.gridtype_id = _gridtype_to_id[gridtype] # "tiled" or "hash"
|
||||
self.align_corners = align_corners
|
||||
|
||||
# allocate parameters
|
||||
offsets = []
|
||||
offset = 0
|
||||
self.max_params = 2 ** log2_hashmap_size
|
||||
for i in range(n_levels):
|
||||
resolution = int(np.ceil(base_resolution * per_level_scale ** i))
|
||||
params_in_level = min(self.max_params, (resolution if align_corners else resolution + 1) ** input_dim) # limit max number
|
||||
params_in_level = int(np.ceil(params_in_level / 8) * 8) # make divisible
|
||||
print(f"level {i}, resolution: {resolution}")
|
||||
offsets.append(offset)
|
||||
offset += params_in_level
|
||||
offsets.append(offset)
|
||||
offsets = torch.from_numpy(np.array(offsets, dtype=np.int32))
|
||||
self.register_buffer('offsets', offsets)
|
||||
|
||||
self.n_params = offsets[-1] * level_dim
|
||||
|
||||
# parameters
|
||||
self.embeddings = nn.Parameter(torch.empty(offset, level_dim))
|
||||
# self.embeddings = nn.Embedding(offset, level_dim, sparse=True)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
std = 1e-4
|
||||
self.embeddings.data.uniform_(-std, std)
|
||||
# self.embeddings.weight.data.uniform_(-std, std)
|
||||
|
||||
|
||||
def __repr__(self):
|
||||
return f"GridEncoder: input_dim={self.input_dim} n_levels={self.n_levels} level_dim={self.level_dim} resolution={self.base_resolution} -> {int(round(self.base_resolution * self.per_level_scale ** (self.n_levels - 1)))} per_level_scale={self.per_level_scale:.4f} params={tuple(self.embeddings.shape)} gridtype={self.gridtype} align_corners={self.align_corners}"
|
||||
|
||||
|
||||
def forward(self, inputs, bound=1):
|
||||
# inputs: [..., input_dim], normalized real world positions in [-bound, bound]
|
||||
# return: [..., num_levels * level_dim]
|
||||
|
||||
inputs = (inputs + bound) / (2 * bound) # map to [0, 1]
|
||||
|
||||
#print('inputs', inputs.shape, inputs.dtype, inputs.min().item(), inputs.max().item())
|
||||
|
||||
prefix_shape = list(inputs.shape[:-1])
|
||||
inputs = inputs.view(-1, self.input_dim)
|
||||
|
||||
outputs = grid_encode(inputs, self.embeddings, self.offsets, self.per_level_scale, self.base_resolution, inputs.requires_grad, self.gridtype_id, self.align_corners)
|
||||
outputs = outputs.view(prefix_shape + [self.out_dim])
|
||||
|
||||
#print('outputs', outputs.shape, outputs.dtype, outputs.min().item(), outputs.max().item())
|
||||
|
||||
return outputs
|
||||
+500
@@ -0,0 +1,500 @@
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <stdexcept>
|
||||
|
||||
#include <stdint.h>
|
||||
#include <cstdio>
|
||||
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor")
|
||||
#define CHECK_IS_INT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor")
|
||||
#define CHECK_IS_FLOATING(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float || x.scalar_type() == at::ScalarType::Half || x.scalar_type() == at::ScalarType::Double, #x " must be a floating tensor")
|
||||
|
||||
|
||||
// just for compatability of half precision in AT_DISPATCH_FLOATING_TYPES_AND_HALF...
|
||||
static inline __device__ at::Half atomicAdd(at::Half *address, at::Half val) {
|
||||
// requires CUDA >= 10 and ARCH >= 70
|
||||
// this is very slow compared to float or __half2, and never used.
|
||||
//return atomicAdd(reinterpret_cast<__half*>(address), val);
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static inline __host__ __device__ T div_round_up(T val, T divisor) {
|
||||
return (val + divisor - 1) / divisor;
|
||||
}
|
||||
|
||||
|
||||
template <uint32_t D>
|
||||
__device__ uint32_t fast_hash(const uint32_t pos_grid[D]) {
|
||||
static_assert(D <= 7, "fast_hash can only hash up to 7 dimensions.");
|
||||
|
||||
// While 1 is technically not a good prime for hashing (or a prime at all), it helps memory coherence
|
||||
// and is sufficient for our use case of obtaining a uniformly colliding index from high-dimensional
|
||||
// coordinates.
|
||||
constexpr uint32_t primes[7] = { 1, 2654435761, 805459861, 3674653429, 2097192037, 1434869437, 2165219737 };
|
||||
|
||||
uint32_t result = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < D; ++i) {
|
||||
result ^= pos_grid[i] * primes[i];
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
template <uint32_t D, uint32_t C>
|
||||
__device__ uint32_t get_grid_index(const uint32_t gridtype, const bool align_corners, const uint32_t ch, const uint32_t hashmap_size, const uint32_t resolution, const uint32_t pos_grid[D]) {
|
||||
uint32_t stride = 1;
|
||||
uint32_t index = 0;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D && stride <= hashmap_size; d++) {
|
||||
index += pos_grid[d] * stride;
|
||||
stride *= align_corners ? resolution: (resolution + 1);
|
||||
}
|
||||
|
||||
// NOTE: for NeRF, the hash is in fact not necessary. Check https://github.com/NVlabs/instant-ngp/issues/97.
|
||||
// gridtype: 0 == hash, 1 == tiled
|
||||
if (gridtype == 0 && stride > hashmap_size) {
|
||||
index = fast_hash<D>(pos_grid);
|
||||
}
|
||||
|
||||
return (index % hashmap_size) * C + ch;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief
|
||||
*
|
||||
* @tparam scalar_t
|
||||
* @tparam D : point dimension usually 3D
|
||||
* @tparam C : embedding dim 1/2/4
|
||||
* @param inputs
|
||||
* @param grid
|
||||
* @param offsets
|
||||
* @param outputs
|
||||
* @param B
|
||||
* @param L
|
||||
* @param S
|
||||
* @param H
|
||||
* @param calc_grad_inputs
|
||||
* @param dy_dx
|
||||
* @param gridtype
|
||||
* @param align_corners
|
||||
* @return __global__
|
||||
*/
|
||||
template <typename scalar_t, uint32_t D, uint32_t C>
|
||||
__global__ void kernel_grid(
|
||||
const float * __restrict__ inputs,
|
||||
const scalar_t * __restrict__ grid,
|
||||
const int * __restrict__ offsets,
|
||||
scalar_t * __restrict__ outputs,
|
||||
const uint32_t B, const uint32_t L, const float S, const uint32_t H,
|
||||
const bool calc_grad_inputs,
|
||||
scalar_t * __restrict__ dy_dx,
|
||||
const uint32_t gridtype,
|
||||
const bool align_corners
|
||||
) {
|
||||
const uint32_t b = blockIdx.x * blockDim.x + threadIdx.x; // batch id
|
||||
|
||||
if (b >= B) return;
|
||||
|
||||
const uint32_t level = blockIdx.y;
|
||||
|
||||
// locate
|
||||
grid += (uint32_t)offsets[level] * C;
|
||||
inputs += b * D;
|
||||
outputs += level * B * C + b * C;
|
||||
|
||||
// check input range (should be in [0, 1])
|
||||
bool flag_oob = false;
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if (inputs[d] < 0 || inputs[d] > 1) {
|
||||
flag_oob = true;
|
||||
}
|
||||
}
|
||||
// if input out of bound, just set output to 0
|
||||
if (flag_oob) {
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
outputs[ch] = 0;
|
||||
}
|
||||
if (calc_grad_inputs) {
|
||||
dy_dx += b * D * L * C + level * D * C; // B L D C
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
dy_dx[d * C + ch] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const uint32_t hashmap_size = offsets[level + 1] - offsets[level];
|
||||
const float scale = exp2f(level * S) * H - 1.0f;
|
||||
const uint32_t resolution = (uint32_t)ceil(scale) + 1;
|
||||
|
||||
// calculate coordinate
|
||||
float pos[D];
|
||||
uint32_t pos_grid[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
pos[d] = inputs[d] * scale + (align_corners ? 0.0f : 0.5f);
|
||||
pos_grid[d] = floorf(pos[d]);
|
||||
pos[d] -= (float)pos_grid[d];
|
||||
}
|
||||
|
||||
//printf("[b=%d, l=%d] pos=(%f, %f)+(%d, %d)\n", b, level, pos[0], pos[1], pos_grid[0], pos_grid[1]);
|
||||
|
||||
// interpolate
|
||||
scalar_t results[C] = {0}; // temp results in register
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t idx = 0; idx < (1 << D); idx++) {
|
||||
float w = 1;
|
||||
uint32_t pos_grid_local[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if ((idx & (1 << d)) == 0) {
|
||||
w *= 1 - pos[d];
|
||||
pos_grid_local[d] = pos_grid[d];
|
||||
} else {
|
||||
w *= pos[d];
|
||||
pos_grid_local[d] = pos_grid[d] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t index = get_grid_index<D, C>(gridtype, align_corners, 0, hashmap_size, resolution, pos_grid_local); // Hashing
|
||||
|
||||
// writing to register (fast)
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
results[ch] += w * grid[index + ch];
|
||||
}
|
||||
|
||||
//printf("[b=%d, l=%d] int %d, idx %d, w %f, val %f\n", b, level, idx, index, w, grid[index]);
|
||||
}
|
||||
|
||||
// writing to global memory (slow)
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
outputs[ch] = results[ch];
|
||||
}
|
||||
|
||||
// prepare dy_dx for calc_grad_inputs
|
||||
// differentiable (soft) indexing: https://discuss.pytorch.org/t/differentiable-indexing/17647/9
|
||||
if (calc_grad_inputs) {
|
||||
|
||||
dy_dx += b * D * L * C + level * D * C; // B L D C
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t gd = 0; gd < D; gd++) {
|
||||
|
||||
scalar_t results_grad[C] = {0};
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t idx = 0; idx < (1 << (D - 1)); idx++) {
|
||||
float w = scale;
|
||||
uint32_t pos_grid_local[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t nd = 0; nd < D - 1; nd++) {
|
||||
const uint32_t d = (nd >= gd) ? (nd + 1) : nd;
|
||||
|
||||
if ((idx & (1 << nd)) == 0) {
|
||||
w *= 1 - pos[d];
|
||||
pos_grid_local[d] = pos_grid[d];
|
||||
} else {
|
||||
w *= pos[d];
|
||||
pos_grid_local[d] = pos_grid[d] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
pos_grid_local[gd] = pos_grid[gd];
|
||||
uint32_t index_left = get_grid_index<D, C>(gridtype, align_corners, 0, hashmap_size, resolution, pos_grid_local);
|
||||
pos_grid_local[gd] = pos_grid[gd] + 1;
|
||||
uint32_t index_right = get_grid_index<D, C>(gridtype, align_corners, 0, hashmap_size, resolution, pos_grid_local);
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
results_grad[ch] += w * (grid[index_right + ch] - grid[index_left + ch]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
dy_dx[gd * C + ch] = results_grad[ch];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename scalar_t, uint32_t D, uint32_t C, uint32_t N_C>
|
||||
__global__ void kernel_grid_backward(
|
||||
const scalar_t * __restrict__ grad,
|
||||
const float * __restrict__ inputs,
|
||||
const scalar_t * __restrict__ grid,
|
||||
const int * __restrict__ offsets,
|
||||
scalar_t * __restrict__ grad_grid,
|
||||
const uint32_t B, const uint32_t L, const float S, const uint32_t H,
|
||||
const uint32_t gridtype,
|
||||
const bool align_corners
|
||||
) {
|
||||
const uint32_t b = (blockIdx.x * blockDim.x + threadIdx.x) * N_C / C;
|
||||
if (b >= B) return;
|
||||
|
||||
const uint32_t level = blockIdx.y;
|
||||
const uint32_t ch = (blockIdx.x * blockDim.x + threadIdx.x) * N_C - b * C;
|
||||
|
||||
// locate
|
||||
grad_grid += offsets[level] * C;
|
||||
inputs += b * D;
|
||||
grad += level * B * C + b * C + ch; // L, B, C
|
||||
|
||||
const uint32_t hashmap_size = offsets[level + 1] - offsets[level];
|
||||
const float scale = exp2f(level * S) * H - 1.0f;
|
||||
const uint32_t resolution = (uint32_t)ceil(scale) + 1;
|
||||
|
||||
// check input range (should be in [0, 1])
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if (inputs[d] < 0 || inputs[d] > 1) {
|
||||
return; // grad is init as 0, so we simply return.
|
||||
}
|
||||
}
|
||||
|
||||
// calculate coordinate
|
||||
float pos[D];
|
||||
uint32_t pos_grid[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
pos[d] = inputs[d] * scale + (align_corners ? 0.0f : 0.5f);
|
||||
pos_grid[d] = floorf(pos[d]);
|
||||
pos[d] -= (float)pos_grid[d];
|
||||
}
|
||||
|
||||
scalar_t grad_cur[N_C] = {0}; // fetch to register
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < N_C; c++) {
|
||||
grad_cur[c] = grad[c];
|
||||
}
|
||||
|
||||
// interpolate
|
||||
#pragma unroll
|
||||
for (uint32_t idx = 0; idx < (1 << D); idx++) {
|
||||
float w = 1;
|
||||
uint32_t pos_grid_local[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if ((idx & (1 << d)) == 0) {
|
||||
w *= 1 - pos[d];
|
||||
pos_grid_local[d] = pos_grid[d];
|
||||
} else {
|
||||
w *= pos[d];
|
||||
pos_grid_local[d] = pos_grid[d] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t index = get_grid_index<D, C>(gridtype, align_corners, ch, hashmap_size, resolution, pos_grid_local);
|
||||
|
||||
// atomicAdd for __half is slow (especially for large values), so we use __half2 if N_C % 2 == 0
|
||||
// TODO: use float which is better than __half, if N_C % 2 != 0
|
||||
if (std::is_same<scalar_t, at::Half>::value && N_C % 2 == 0) {
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < N_C; c += 2) {
|
||||
// process two __half at once (by interpreting as a __half2)
|
||||
__half2 v = {(__half)(w * grad_cur[c]), (__half)(w * grad_cur[c + 1])};
|
||||
atomicAdd((__half2*)&grad_grid[index + c], v);
|
||||
}
|
||||
// float, or __half when N_C % 2 != 0 (which means C == 1)
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < N_C; c++) {
|
||||
atomicAdd(&grad_grid[index + c], w * grad_cur[c]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename scalar_t, uint32_t D, uint32_t C>
|
||||
__global__ void kernel_input_backward(
|
||||
const scalar_t * __restrict__ grad,
|
||||
const scalar_t * __restrict__ dy_dx,
|
||||
scalar_t * __restrict__ grad_inputs,
|
||||
uint32_t B, uint32_t L
|
||||
) {
|
||||
const uint32_t t = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
if (t >= B * D) return;
|
||||
|
||||
const uint32_t b = t / D;
|
||||
const uint32_t d = t - b * D;
|
||||
|
||||
dy_dx += b * L * D * C;
|
||||
|
||||
scalar_t result = 0;
|
||||
|
||||
# pragma unroll
|
||||
for (int l = 0; l < L; l++) {
|
||||
# pragma unroll
|
||||
for (int ch = 0; ch < C; ch++) {
|
||||
result += grad[l * B * C + b * C + ch] * dy_dx[l * D * C + d * C + ch];
|
||||
}
|
||||
}
|
||||
|
||||
grad_inputs[t] = result;
|
||||
}
|
||||
|
||||
|
||||
template <typename scalar_t, uint32_t D>
|
||||
void kernel_grid_wrapper(const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *outputs, const uint32_t B, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, const uint32_t gridtype, const bool align_corners) {
|
||||
static constexpr uint32_t N_THREAD = 512;
|
||||
const dim3 blocks_hashgrid = { div_round_up(B, N_THREAD), L, 1 };
|
||||
switch (C) {
|
||||
case 1: kernel_grid<scalar_t, D, 1><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 2: kernel_grid<scalar_t, D, 2><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 4: kernel_grid<scalar_t, D, 4><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 8: kernel_grid<scalar_t, D, 8><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."};
|
||||
}
|
||||
}
|
||||
|
||||
// inputs: [B, D], float, in [0, 1]
|
||||
// embeddings: [sO, C], float
|
||||
// offsets: [L + 1], uint32_t
|
||||
// outputs: [L, B, C], float (L first, so only one level of hashmap needs to fit into cache at a time.)
|
||||
// H: base resolution
|
||||
// dy_dx: [B, L * D * C]
|
||||
template <typename scalar_t>
|
||||
void grid_encode_forward_cuda(const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *outputs, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, const uint32_t gridtype, const bool align_corners) {
|
||||
switch (D) {
|
||||
case 1: kernel_grid_wrapper<scalar_t, 1>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 2: kernel_grid_wrapper<scalar_t, 2>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 3: kernel_grid_wrapper<scalar_t, 3>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 4: kernel_grid_wrapper<scalar_t, 4>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 5: kernel_grid_wrapper<scalar_t, 5>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
default: throw std::runtime_error{"GridEncoding: D must be 1, 2, 3, 4, or 5."};
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
template <typename scalar_t, uint32_t D>
|
||||
void kernel_grid_backward_wrapper(const scalar_t *grad, const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *grad_embeddings, const uint32_t B, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, scalar_t *grad_inputs, const uint32_t gridtype, const bool align_corners) {
|
||||
static constexpr uint32_t N_THREAD = 256;
|
||||
const uint32_t N_C = std::min(2u, C); // n_features_per_thread
|
||||
const dim3 blocks_hashgrid = { div_round_up(B * C / N_C, N_THREAD), L, 1 };
|
||||
switch (C) {
|
||||
case 1:
|
||||
kernel_grid_backward<scalar_t, D, 1, 1><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 1><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
case 2:
|
||||
kernel_grid_backward<scalar_t, D, 2, 2><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 2><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
case 4:
|
||||
kernel_grid_backward<scalar_t, D, 4, 2><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 4><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
case 8:
|
||||
kernel_grid_backward<scalar_t, D, 8, 2><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 8><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."};
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// grad: [L, B, C], float
|
||||
// inputs: [B, D], float, in [0, 1]
|
||||
// embeddings: [sO, C], float
|
||||
// offsets: [L + 1], uint32_t
|
||||
// grad_embeddings: [sO, C]
|
||||
// H: base resolution
|
||||
template <typename scalar_t>
|
||||
void grid_encode_backward_cuda(const scalar_t *grad, const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *grad_embeddings, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, scalar_t *grad_inputs, const uint32_t gridtype, const bool align_corners) {
|
||||
switch (D) {
|
||||
case 1: kernel_grid_backward_wrapper<scalar_t, 1>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 2: kernel_grid_backward_wrapper<scalar_t, 2>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 3: kernel_grid_backward_wrapper<scalar_t, 3>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 4: kernel_grid_backward_wrapper<scalar_t, 4>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 5: kernel_grid_backward_wrapper<scalar_t, 5>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
default: throw std::runtime_error{"GridEncoding: D must be 1, 2, 3, 4, or 5."};
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
void grid_encode_forward(const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor outputs, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, at::Tensor dy_dx, const uint32_t gridtype, const bool align_corners) {
|
||||
CHECK_CUDA(inputs);
|
||||
CHECK_CUDA(embeddings);
|
||||
CHECK_CUDA(offsets);
|
||||
CHECK_CUDA(outputs);
|
||||
CHECK_CUDA(dy_dx);
|
||||
|
||||
CHECK_CONTIGUOUS(inputs);
|
||||
CHECK_CONTIGUOUS(embeddings);
|
||||
CHECK_CONTIGUOUS(offsets);
|
||||
CHECK_CONTIGUOUS(outputs);
|
||||
CHECK_CONTIGUOUS(dy_dx);
|
||||
|
||||
CHECK_IS_FLOATING(inputs);
|
||||
CHECK_IS_FLOATING(embeddings);
|
||||
CHECK_IS_INT(offsets);
|
||||
CHECK_IS_FLOATING(outputs);
|
||||
CHECK_IS_FLOATING(dy_dx);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
|
||||
embeddings.scalar_type(), "grid_encode_forward", ([&] {
|
||||
grid_encode_forward_cuda<scalar_t>(inputs.data_ptr<float>(), embeddings.data_ptr<scalar_t>(), offsets.data_ptr<int>(), outputs.data_ptr<scalar_t>(), B, D, C, L, S, H, calc_grad_inputs, dy_dx.data_ptr<scalar_t>(), gridtype, align_corners);
|
||||
}));
|
||||
}
|
||||
|
||||
void grid_encode_backward(const at::Tensor grad, const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor grad_embeddings, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, const at::Tensor dy_dx, at::Tensor grad_inputs, const uint32_t gridtype, const bool align_corners) {
|
||||
CHECK_CUDA(grad);
|
||||
CHECK_CUDA(inputs);
|
||||
CHECK_CUDA(embeddings);
|
||||
CHECK_CUDA(offsets);
|
||||
CHECK_CUDA(grad_embeddings);
|
||||
CHECK_CUDA(dy_dx);
|
||||
CHECK_CUDA(grad_inputs);
|
||||
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
CHECK_CONTIGUOUS(inputs);
|
||||
CHECK_CONTIGUOUS(embeddings);
|
||||
CHECK_CONTIGUOUS(offsets);
|
||||
CHECK_CONTIGUOUS(grad_embeddings);
|
||||
CHECK_CONTIGUOUS(dy_dx);
|
||||
CHECK_CONTIGUOUS(grad_inputs);
|
||||
|
||||
CHECK_IS_FLOATING(grad);
|
||||
CHECK_IS_FLOATING(inputs);
|
||||
CHECK_IS_FLOATING(embeddings);
|
||||
CHECK_IS_INT(offsets);
|
||||
CHECK_IS_FLOATING(grad_embeddings);
|
||||
CHECK_IS_FLOATING(dy_dx);
|
||||
CHECK_IS_FLOATING(grad_inputs);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
|
||||
grad.scalar_type(), "grid_encode_backward", ([&] {
|
||||
grid_encode_backward_cuda<scalar_t>(grad.data_ptr<scalar_t>(), inputs.data_ptr<float>(), embeddings.data_ptr<scalar_t>(), offsets.data_ptr<int>(), grad_embeddings.data_ptr<scalar_t>(), B, D, C, L, S, H, calc_grad_inputs, dy_dx.data_ptr<scalar_t>(), grad_inputs.data_ptr<scalar_t>(), gridtype, align_corners);
|
||||
}));
|
||||
|
||||
}
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
#ifndef _HASH_ENCODE_H
|
||||
#define _HASH_ENCODE_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
// inputs: [B, D], float, in [0, 1]
|
||||
// embeddings: [sO, C], float
|
||||
// offsets: [L + 1], uint32_t
|
||||
// outputs: [B, L * C], float
|
||||
// H: base resolution
|
||||
void grid_encode_forward(const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor outputs, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, at::Tensor dy_dx, const uint32_t gridtype, const bool align_corners);
|
||||
void grid_encode_backward(const at::Tensor grad, const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor grad_embeddings, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, const at::Tensor dy_dx, at::Tensor grad_inputs, const uint32_t gridtype, const bool align_corners);
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user