testbed_sdf.cu
/*
* Copyright (c) 2020-2025, 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.
*/
/** @file testbed_sdf.cu
* @author Thomas Müller & Alex Evans, NVIDIA
*/
#include <neural-graphics-primitives/common.h>
#include <neural-graphics-primitives/common_device.cuh>
#include <neural-graphics-primitives/discrete_distribution.h>
#include <neural-graphics-primitives/envmap.cuh>
#include <neural-graphics-primitives/random_val.cuh> // helpers to generate random values, directions
#include <neural-graphics-primitives/render_buffer.h>
#include <neural-graphics-primitives/takikawa_encoding.cuh>
#include <neural-graphics-primitives/testbed.h>
#include <neural-graphics-primitives/tinyobj_loader_wrapper.h>
#include <neural-graphics-primitives/trainable_buffer.cuh>
#include <neural-graphics-primitives/triangle_bvh.cuh>
#include <neural-graphics-primitives/triangle_octree.cuh>
#include <tiny-cuda-nn/encodings/grid.h>
#include <tiny-cuda-nn/gpu_matrix.h>
#include <tiny-cuda-nn/network_with_input_encoding.h>
#include <tiny-cuda-nn/trainer.h>
#include <cmrc/cmrc.hpp>
CMRC_DECLARE(ngp);
#ifdef copysign
# undef copysign
#endif
namespace ngp {
static constexpr uint32_t MARCH_ITER = 10000;
Testbed::NetworkDims Testbed::network_dims_sdf() const {
NetworkDims dims;
dims.n_input = 3;
dims.n_output = 1;
dims.n_pos = 3;
return dims;
}
__device__ inline float square(float x) { return x * x; }
__device__ inline float mix(float a, float b, float t) { return a + (b - a) * t; }
__device__ inline vec3 mix(const vec3& a, const vec3& b, float t) { return a + (b - a) * t; }
__device__ inline float SchlickFresnel(float u) {
float m = __saturatef(1.0 - u);
return square(square(m)) * m;
}
__device__ inline float G1(float NdotH, float a) {
if (a >= 1.0) {
return 1.0 / PI();
}
float a2 = square(a);
float t = 1.0 + (a2 - 1.0) * NdotH * NdotH;
return (a2 - 1.0) / (PI() * log(a2) * t);
}
__device__ inline float G2(float NdotH, float a) {
float a2 = square(a);
float t = 1.0 + (a2 - 1.0) * NdotH * NdotH;
return a2 / (PI() * t * t);
}
__device__ inline float SmithG_GGX(float NdotV, float alphaG) {
float a = alphaG * alphaG;
float b = NdotV * NdotV;
return 1.0 / (NdotV + sqrtf(a + b - a * b));
}
// this function largely based on:
// https://github.com/wdas/brdf/blob/master/src/brdfs/disney.brdf
// http://blog.selfshadow.com/publications/s2012-shading-course/burley/s2012_pbs_disney_brdf_notes_v3.pdf
__device__ vec3 evaluate_shading(
const vec3& base_color,
const vec3& ambient_color, // :)
const vec3& light_color, // :)
float metallic,
float subsurface,
float specular,
float roughness,
float specular_tint,
float sheen,
float sheen_tint,
float clearcoat,
float clearcoat_gloss,
vec3 L,
vec3 V,
vec3 N
) {
float NdotL = dot(N, L);
float NdotV = dot(N, V);
vec3 H = normalize(L + V);
float NdotH = dot(N, H);
float LdotH = dot(L, H);
// Diffuse fresnel - go from 1 at normal incidence to .5 at grazing
// and mix in diffuse retro-reflection based on roughness
float FL = SchlickFresnel(NdotL), FV = SchlickFresnel(NdotV);
vec3 amb = (ambient_color * mix(0.2f, FV, metallic));
amb *= base_color;
if (NdotL < 0.f || NdotV < 0.f) {
return amb;
}
float luminance = dot(base_color, vec3{0.3f, 0.6f, 0.1f});
// normalize luminance to isolate hue and saturation components
vec3 Ctint = base_color * (1.f / (luminance + 0.00001f));
vec3 Cspec0 = mix(mix(vec3(1.0f), Ctint, specular_tint) * specular * 0.08f, base_color, metallic);
vec3 Csheen = mix(vec3(1.0f), Ctint, sheen_tint);
float Fd90 = 0.5f + 2.0f * LdotH * LdotH * roughness;
float Fd = mix(1, Fd90, FL) * mix(1.f, Fd90, FV);
// Based on Hanrahan-Krueger BRDF approximation of isotropic BSSRDF
// 1.25 scale is used to (roughly) preserve albedo
// Fss90 used to "flatten" retroreflection based on roughness
float Fss90 = LdotH * LdotH * roughness;
float Fss = mix(1.0f, Fss90, FL) * mix(1.0f, Fss90, FV);
float ss = 1.25f * (Fss * (1.f / (NdotL + NdotV) - 0.5f) + 0.5f);
// Specular
float a = std::max(0.001f, square(roughness));
float Ds = G2(NdotH, a);
float FH = SchlickFresnel(LdotH);
vec3 Fs = mix(Cspec0, vec3(1.0f), FH);
float Gs = SmithG_GGX(NdotL, a) * SmithG_GGX(NdotV, a);
// sheen
vec3 Fsheen = FH * sheen * Csheen;
// clearcoat (ior = 1.5 -> F0 = 0.04)
float Dr = G1(NdotH, mix(0.1f, 0.001f, clearcoat_gloss));
float Fr = mix(0.04f, 1.0f, FH);
float Gr = SmithG_GGX(NdotL, 0.25f) * SmithG_GGX(NdotV, 0.25f);
float CCs = 0.25f * clearcoat * Gr * Fr * Dr;
vec3 brdf = (float(1.0f / PI()) * mix(Fd, ss, subsurface) * base_color + Fsheen) * (1.0f - metallic) + Gs * Fs * Ds + vec3{CCs, CCs, CCs};
return vec3(brdf * light_color) * NdotL + amb;
}
__global__ void advance_pos_kernel_sdf(
const uint32_t n_elements,
const float zero_offset,
vec3* __restrict__ positions,
float* __restrict__ distances,
SdfPayload* __restrict__ payloads,
BoundingBox aabb,
float floor_y,
const TriangleOctreeNode* __restrict__ octree_nodes,
int max_octree_depth,
float distance_scale,
float maximum_distance,
float k,
float* __restrict__ prev_distances,
float* __restrict__ total_distances,
float* __restrict__ min_visibility
) {
const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
if (i >= n_elements) {
return;
}
SdfPayload& payload = payloads[i];
if (!payload.alive) {
return;
}
float distance = distances[i] - zero_offset;
distance *= distance_scale;
// Advance by the predicted distance
vec3 pos = positions[i];
pos += distance * payload.dir;
// Skip over regions not covered by the octree
if (octree_nodes && !contains(octree_nodes, max_octree_depth, pos)) {
float octree_distance = ray_intersect(octree_nodes, max_octree_depth, pos, payload.dir) + 1e-6f;
distance += octree_distance;
pos += octree_distance * payload.dir;
}
if (pos.y < floor_y && payload.dir.y < 0.f) {
float floor_dist = -(pos.y - floor_y) / payload.dir.y;
distance += floor_dist;
pos += floor_dist * payload.dir;
payload.alive = false;
}
positions[i] = pos;
if (total_distances && distance > 0.0f) {
// From https://www.iquilezles.org/www/articles/rmshadows/rmshadows.htm
float total_distance = total_distances[i];
float y = distance * distance / (2.0f * prev_distances[i]);
float d = sqrtf(distance * distance - y * y);
min_visibility[i] = fminf(min_visibility[i], k * d / fmaxf(0.0f, total_distance - y));
prev_distances[i] = distance;
total_distances[i] = total_distance + distance;
}
bool stay_alive = distance > maximum_distance && fabsf(distance / 2) > 3 * maximum_distance;
if (!stay_alive) {
payload.alive = false;
return;
}
if (!aabb.contains(pos)) {
payload.alive = false;
return;
}
++payload.n_steps;
}
__global__ void perturb_sdf_samples(
uint32_t n_elements, const vec3* __restrict__ perturbations, vec3* __restrict__ positions, float* __restrict__ distances
) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_elements) {
return;
}
vec3 perturbation = perturbations[i];
positions[i] += perturbation;
// Small epsilon above 1 to ensure a triangle is always found.
distances[i] = length(perturbation) * 1.001f;
}
__global__ void prepare_shadow_rays(
const uint32_t n_elements,
vec3 sun_dir,
vec3* __restrict__ positions,
vec3* __restrict__ normals,
float* __restrict__ distances,
float* __restrict__ prev_distances,
float* __restrict__ total_distances,
float* __restrict__ min_visibility,
SdfPayload* __restrict__ payloads,
BoundingBox aabb,
const TriangleOctreeNode* __restrict__ octree_nodes,
int max_octree_depth
) {
const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
if (i >= n_elements) {
return;
}
SdfPayload& payload = payloads[i];
// Step back a little along the ray to prevent self-intersection
vec3 view_pos = positions[i] + normalize(faceforward(normals[i], payload.dir, normals[i])) * 1e-3f;
vec3 dir = normalize(sun_dir);
float t = fmaxf(aabb.ray_intersect(view_pos, dir).x + 1e-6f, 0.0f);
view_pos += t * dir;
if (octree_nodes && !contains(octree_nodes, max_octree_depth, view_pos)) {
t = fmaxf(0.0f, ray_intersect(octree_nodes, max_octree_depth, view_pos, dir) + 1e-6f);
view_pos += t * dir;
}
positions[i] = view_pos;
if (!aabb.contains(view_pos)) {
distances[i] = MAX_DEPTH();
payload.alive = false;
min_visibility[i] = 1.0f;
return;
}
distances[i] = MAX_DEPTH();
payload.idx = i;
payload.dir = dir;
payload.n_steps = 0;
payload.alive = true;
if (prev_distances) {
prev_distances[i] = 1e20f;
}
if (total_distances) {
total_distances[i] = 0.0f;
}
if (min_visibility) {
min_visibility[i] = 1.0f;
}
}
__global__ void write_shadow_ray_result(
const uint32_t n_elements,
BoundingBox aabb,
const vec3* __restrict__ positions,
const SdfPayload* __restrict__ shadow_payloads,
const float* __restrict__ min_visibility,
float* __restrict__ shadow_factors
) {
const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
if (i >= n_elements) {
return;
}
shadow_factors[shadow_payloads[i].idx] = aabb.contains(positions[i]) ? 0.0f : min_visibility[i];
}
__global__ void shade_kernel_sdf(
const uint32_t n_elements,
BoundingBox aabb,
float floor_y,
const ERenderMode mode,
const BRDFParams brdf,
vec3 sun_dir,
vec3 up_dir,
mat4x3 camera_matrix,
vec3* __restrict__ positions,
vec3* __restrict__ normals,
float* __restrict__ distances,
SdfPayload* __restrict__ payloads,
vec4* __restrict__ frame_buffer,
float* __restrict__ depth_buffer
) {
const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
if (i >= n_elements) {
return;
}
SdfPayload& payload = payloads[i];
if (!aabb.contains(positions[i])) {
return;
}
// The normal in memory isn't normalized yet
vec3 normal = normalize(normals[i]);
vec3 pos = positions[i];
bool floor = false;
if (pos.y < floor_y + 0.001f && payload.dir.y < 0.f) {
normal = vec3{0.0f, 1.0f, 0.0f};
floor = true;
}
vec3 cam_pos = camera_matrix[3];
vec3 cam_fwd = camera_matrix[2];
float ao = powf(0.92f, payload.n_steps * 0.5f) * (1.f / 0.92f);
vec3 color;
switch (mode) {
case ERenderMode::AO: color = vec3(powf(0.92f, payload.n_steps)); break;
case ERenderMode::Shade: {
float skyam = -dot(normal, up_dir) * 0.5f + 0.5f;
vec3 suncol = vec3{255.f / 255.0f, 225.f / 255.0f, 195.f / 255.0f} * 4.f *
distances[i]; // Distance encodes shadow occlusion. 0=occluded, 1=no shadow
const vec3 skycol = vec3{195.f / 255.0f, 215.f / 255.0f, 255.f / 255.0f} * 4.f * skyam;
float check_size = 8.f / aabb.diag().x;
float check = ((int(floorf(check_size * (pos.x - aabb.min.x))) ^ int(floorf(check_size * (pos.z - aabb.min.z)))) & 1) ? 0.8f :
0.2f;
const vec3 floorcol = vec3{check * check * check, check * check, check};
color = evaluate_shading(
floor ? floorcol : brdf.basecolor * brdf.basecolor,
brdf.ambientcolor * skycol,
suncol,
floor ? 0.f : brdf.metallic,
floor ? 0.f : brdf.subsurface,
floor ? 1.f : brdf.specular,
floor ? 0.5f : brdf.roughness,
0.f,
floor ? 0.f : brdf.sheen,
0.f,
floor ? 0.f : brdf.clearcoat,
brdf.clearcoat_gloss,
sun_dir,
-normalize(payload.dir),
normal
);
} break;
case ERenderMode::Depth: color = vec3(dot(cam_fwd, pos - cam_pos)); break;
case ERenderMode::Positions: {
color = (pos - 0.5f) / 2.0f + 0.5f;
} break;
case ERenderMode::Normals: color = 0.5f * normal + 0.5f; break;
case ERenderMode::Cost: color = vec3((float)payload.n_steps / 30); break;
case ERenderMode::EncodingVis: color = normals[i]; break;
}
frame_buffer[payload.idx] = {color.r, color.g, color.b, 1.0f};
depth_buffer[payload.idx] = dot(cam_fwd, pos - cam_pos);
}
__global__ void compact_kernel_shadow_sdf(
const uint32_t n_elements,
const float zero_offset,
vec3* src_positions,
float* src_distances,
SdfPayload* src_payloads,
float* src_prev_distances,
float* src_total_distances,
float* src_min_visibility,
vec3* dst_positions,
float* dst_distances,
SdfPayload* dst_payloads,
float* dst_prev_distances,
float* dst_total_distances,
float* dst_min_visibility,
vec3* dst_final_positions,
float* dst_final_distances,
SdfPayload* dst_final_payloads,
float* dst_final_prev_distances,
float* dst_final_total_distances,
float* dst_final_min_visibility,
BoundingBox aabb,
uint32_t* counter,
uint32_t* finalCounter
) {
const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
if (i >= n_elements) {
return;
}
SdfPayload& src_payload = src_payloads[i];
if (src_payload.alive) {
uint32_t idx = atomicAdd(counter, 1);
dst_payloads[idx] = src_payload;
dst_positions[idx] = src_positions[i];
dst_distances[idx] = src_distances[i];
dst_prev_distances[idx] = src_prev_distances[i];
dst_total_distances[idx] = src_total_distances[i];
dst_min_visibility[idx] = src_min_visibility[i];
} else { // For shadow rays, collect _all_ final samples to keep track of their partial visibility
uint32_t idx = atomicAdd(finalCounter, 1);
dst_final_payloads[idx] = src_payload;
dst_final_positions[idx] = src_positions[i];
dst_final_distances[idx] = src_distances[i];
dst_final_prev_distances[idx] = src_prev_distances[i];
dst_final_total_distances[idx] = src_total_distances[i];
dst_final_min_visibility[idx] = aabb.contains(src_positions[i]) ? 0.0f : src_min_visibility[i];
}
}
__global__ void compact_kernel_sdf(
const uint32_t n_elements,
const float zero_offset,
vec3* src_positions,
float* src_distances,
SdfPayload* src_payloads,
vec3* dst_positions,
float* dst_distances,
SdfPayload* dst_payloads,
vec3* dst_final_positions,
float* dst_final_distances,
SdfPayload* dst_final_payloads,
BoundingBox aabb,
uint32_t* counter,
uint32_t* finalCounter
) {
const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
if (i >= n_elements) {
return;
}
SdfPayload& src_payload = src_payloads[i];
if (src_payload.alive) {
uint32_t idx = atomicAdd(counter, 1);
dst_payloads[idx] = src_payload;
dst_positions[idx] = src_positions[i];
dst_distances[idx] = src_distances[i];
} else if (aabb.contains(src_positions[i])) {
uint32_t idx = atomicAdd(finalCounter, 1);
dst_final_payloads[idx] = src_payload;
dst_final_positions[idx] = src_positions[i];
dst_final_distances[idx] = 1.0f; // HACK: Distances encode shadowing factor when shading
}
}
__global__ void uniform_octree_sample_kernel(
const uint32_t num_elements,
default_rng_t rng,
const TriangleOctreeNode* __restrict__ octree_nodes,
uint32_t num_nodes,
uint32_t depth,
vec3* __restrict__ samples
) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= num_elements) {
return;
}
rng.advance(i * (1 << 8));
// Samples random nodes until a leaf is picked
uint32_t node;
uint32_t child;
do {
node = umin((uint32_t)(random_val(rng) * num_nodes), num_nodes - 1);
child = umin((uint32_t)(random_val(rng) * 8), 8u - 1);
} while (octree_nodes[node].depth < depth - 2 || octree_nodes[node].children[child] == -1);
// Here it should be guaranteed that any child of the node is -1
float size = scalbnf(1.0f, -depth + 1);
u16vec3 pos = octree_nodes[node].pos * uint16_t(2);
if (child & 1) {
++pos.x;
}
if (child & 2) {
++pos.y;
}
if (child & 4) {
++pos.z;
}
samples[i] = size * (vec3(pos) + samples[i]);
}
__global__ void scale_to_aabb_kernel(uint32_t n_elements, BoundingBox aabb, vec3* __restrict__ inout) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_elements) {
return;
}
inout[i] = aabb.min + inout[i] * aabb.diag();
}
__global__ void compare_signs_kernel(
uint32_t n_elements,
const vec3* positions,
const float* distances_ref,
const float* distances_model,
uint32_t* counters,
const TriangleOctreeNode* octree_nodes,
int max_octree_depth
) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_elements) {
return;
}
bool inside1 = distances_ref[i] <= 0.f;
bool inside2 = distances_model[i] <= 0.f;
if (octree_nodes && !contains(octree_nodes, max_octree_depth, positions[i])) {
inside2 = inside1; // assume, when using the octree, that the model is always correct outside the octree.
atomicAdd(&counters[6], 1); // outside the octree
} else {
atomicAdd(&counters[7], 1); // inside the octree
}
atomicAdd(&counters[inside1 ? 0 : 1], 1);
atomicAdd(&counters[inside2 ? 2 : 3], 1);
if (inside1 && inside2) {
atomicAdd(&counters[4], 1);
}
if (inside1 || inside2) {
atomicAdd(&counters[5], 1);
}
}
__global__ void scale_iou_counters_kernel(uint32_t n_elements, uint32_t* counters, float scale) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_elements) {
return;
}
counters[i] = uint32_t(roundf(counters[i] * scale));
}
__global__ void assign_float(uint32_t n_elements, float value, float* __restrict__ out) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_elements) {
return;
}
out[i] = value;
}
__global__ void init_rays_with_payload_kernel_sdf(
uint32_t sample_index,
vec3* __restrict__ positions,
float* __restrict__ distances,
SdfPayload* __restrict__ payloads,
ivec2 resolution,
vec2 focal_length,
mat4x3 camera_matrix,
vec2 screen_center,
vec3 parallax_shift,
bool snap_to_pixel_centers,
BoundingBox aabb,
float floor_y,
float near_distance,
float plane_z,
float aperture_size,
Foveation foveation,
Buffer2DView<const vec4> envmap,
vec4* __restrict__ frame_buffer,
float* __restrict__ depth_buffer,
Buffer2DView<const uint8_t> hidden_area_mask,
Lens lens,
const TriangleOctreeNode* __restrict__ octree_nodes = nullptr,
int max_octree_depth = 0
) {
uint32_t x = threadIdx.x + blockDim.x * blockIdx.x;
uint32_t y = threadIdx.y + blockDim.y * blockIdx.y;
if (x >= resolution.x || y >= resolution.y) {
return;
}
uint32_t idx = x + resolution.x * y;
if (plane_z < 0) {
aperture_size = 0.0;
}
Ray ray = pixel_to_ray(
sample_index,
{(int)x, (int)y},
resolution,
focal_length,
camera_matrix,
screen_center,
parallax_shift,
snap_to_pixel_centers,
near_distance,
plane_z,
aperture_size,
foveation,
hidden_area_mask,
lens
);
distances[idx] = MAX_DEPTH();
depth_buffer[idx] = MAX_DEPTH();
SdfPayload& payload = payloads[idx];
if (!ray.is_valid()) {
payload.dir = ray.d;
payload.idx = idx;
payload.n_steps = 0;
payload.alive = false;
positions[idx] = ray.o;
return;
}
if (plane_z < 0) {
float n = length(ray.d);
payload.dir = (1.0f / n) * ray.d;
payload.idx = idx;
payload.n_steps = 0;
payload.alive = false;
positions[idx] = ray.o - plane_z * ray.d;
depth_buffer[idx] = -plane_z;
return;
}
ray.d = normalize(ray.d);
float t = max(aabb.ray_intersect(ray.o, ray.d).x, 0.0f);
ray.advance(t + 1e-6f);
if (octree_nodes && !contains(octree_nodes, max_octree_depth, ray.o)) {
t = max(0.0f, ray_intersect(octree_nodes, max_octree_depth, ray.o, ray.d));
if (ray.o.y > floor_y && ray.d.y < 0.f) {
float floor_dist = -(ray.o.y - floor_y) / ray.d.y;
if (floor_dist > 0.f) {
t = min(t, floor_dist);
}
}
ray.advance(t + 1e-6f);
}
positions[idx] = ray.o;
if (envmap) {
frame_buffer[idx] = read_envmap(envmap, ray.d);
}
payload.dir = ray.d;
payload.idx = idx;
payload.n_steps = 0;
payload.alive = aabb.contains(ray.o);
}
__host__ __device__ uint32_t sample_discrete(float uniform_sample, const float* __restrict__ cdf, int length) {
return binary_search(uniform_sample, cdf, length);
}
__global__ void sample_uniform_on_triangle_kernel(
uint32_t n_elements, const float* __restrict__ cdf, uint32_t length, const Triangle* __restrict__ triangles, vec3* __restrict__ sampled_positions
) {
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_elements) {
return;
}
vec3 sample = sampled_positions[i];
uint32_t tri_idx = sample_discrete(sample.x, cdf, length);
sampled_positions[i] = triangles[tri_idx].sample_uniform_position(sample.yz());
}
void Testbed::SphereTracer::init_rays_from_camera(
uint32_t sample_index,
const ivec2& resolution,
const vec2& focal_length,
const mat4x3& camera_matrix,
const vec2& screen_center,
const vec3& parallax_shift,
bool snap_to_pixel_centers,
const BoundingBox& aabb,
float floor_y,
float near_distance,
float plane_z,
float aperture_size,
const Foveation& foveation,
const Buffer2DView<const vec4>& envmap,
vec4* frame_buffer,
float* depth_buffer,
const Buffer2DView<const uint8_t>& hidden_area_mask,
const Lens& lens,
const TriangleOctree* octree,
uint32_t n_octree_levels,
cudaStream_t stream
) {
// Make sure we have enough memory reserved to render at the requested resolution
size_t n_pixels = (size_t)resolution.x * resolution.y;
enlarge(n_pixels, stream);
const dim3 threads = {16, 8, 1};
const dim3 blocks = {div_round_up((uint32_t)resolution.x, threads.x), div_round_up((uint32_t)resolution.y, threads.y), 1};
init_rays_with_payload_kernel_sdf<<<blocks, threads, 0, stream>>>(
sample_index,
m_rays[0].pos,
m_rays[0].distance,
m_rays[0].payload,
resolution,
focal_length,
camera_matrix,
screen_center,
parallax_shift,
snap_to_pixel_centers,
aabb,
floor_y,
near_distance,
plane_z,
aperture_size,
foveation,
envmap,
frame_buffer,
depth_buffer,
hidden_area_mask,
lens,
octree ? octree->nodes_gpu() : nullptr,
octree ? n_octree_levels : 0
);
m_n_rays_initialized = (uint32_t)n_pixels;
m_resolution = resolution;
}
void Testbed::SphereTracer::init_rays_from_data(uint32_t n_elements, const RaysSdfSoa& data, cudaStream_t stream) {
enlarge(n_elements, stream);
m_rays[0].copy_from_other_async(n_elements, data, stream);
m_n_rays_initialized = n_elements;
m_resolution = {(int)m_n_rays_initialized, 1};
}
uint32_t Testbed::SphereTracer::trace_bvh(TriangleBvh* bvh, const Triangle* triangles, cudaStream_t stream) {
uint32_t n_alive = m_n_rays_initialized;
m_n_rays_initialized = 0;
if (!bvh) {
return 0;
}
// Abuse the normal buffer to temporarily hold ray directions
parallel_for_gpu(stream, n_alive, [payloads = m_rays[0].payload, normals = m_rays[0].normal] __device__(size_t i) {
normals[i] = payloads[i].dir;
});
bvh->ray_trace_gpu(n_alive, m_rays[0].pos, m_rays[0].normal, triangles, stream);
return n_alive;
}
uint32_t Testbed::SphereTracer::trace(
const distance_fun_t& distance_function,
const Network<float, network_precision_t>* network,
float zero_offset,
float distance_scale,
float maximum_distance,
const BoundingBox& aabb,
const float floor_y,
const TriangleOctree* octree,
const uint32_t n_octree_levels,
cudaStream_t stream
) {
if (m_n_rays_initialized == 0) {
return 0;
}
CUDA_CHECK_THROW(cudaMemsetAsync(m_hit_counter, 0, sizeof(uint32_t), stream));
const uint32_t STEPS_INBETWEEN_COMPACTION = 4;
uint32_t n_alive = m_n_rays_initialized;
m_n_rays_initialized = 0;
uint32_t i = 1;
uint32_t double_buffer_index = 0;
while (i < MARCH_ITER) {
// Compact more frequently in the first couple of steps
uint32_t step_size = std::min(i, STEPS_INBETWEEN_COMPACTION);
RaysSdfSoa& rays_current = m_rays[(double_buffer_index + 1) % 2];
RaysSdfSoa& rays_tmp = m_rays[double_buffer_index % 2];
++double_buffer_index;
if (m_fused_trace_kernel) {
dim3 threads, blocks;
if (m_resolution.x >= 8 && m_resolution.y >= 16) {
threads = {8, 16, 1};
blocks = {
div_round_up((uint32_t)m_resolution.x, threads.x), div_round_up((uint32_t)m_resolution.y, threads.y), 1
};
} else {
threads = {N_THREADS_LINEAR, 1, 1};
blocks = {n_blocks_linear(n_alive, threads.x), 1, 1};
}
m_fused_trace_kernel->launch(
blocks,
threads,
0,
stream,
m_resolution,
zero_offset,
rays_tmp.pos,
rays_tmp.distance,
rays_tmp.payload,
aabb,
floor_y,
octree ? octree->nodes_gpu() : nullptr,
octree ? n_octree_levels : 0,
distance_scale,
maximum_distance,
m_shadow_sharpness,
m_trace_shadow_rays ? rays_tmp.prev_distance : nullptr,
m_trace_shadow_rays ? rays_tmp.total_distance : nullptr,
m_trace_shadow_rays ? rays_tmp.min_visibility : nullptr,
network->inference_params()
);
}
// Compact rays that did not diverge yet
{
CUDA_CHECK_THROW(cudaMemsetAsync(m_alive_counter, 0, sizeof(uint32_t), stream));
if (m_trace_shadow_rays) {
linear_kernel(
compact_kernel_shadow_sdf,
0,
stream,
n_alive,
zero_offset,
rays_tmp.pos,
rays_tmp.distance,
rays_tmp.payload,
rays_tmp.prev_distance,
rays_tmp.total_distance,
rays_tmp.min_visibility,
rays_current.pos,
rays_current.distance,
rays_current.payload,
rays_current.prev_distance,
rays_current.total_distance,
rays_current.min_visibility,
m_rays_hit.pos,
m_rays_hit.distance,
m_rays_hit.payload,
m_rays_hit.prev_distance,
m_rays_hit.total_distance,
m_rays_hit.min_visibility,
aabb,
m_alive_counter,
m_hit_counter
);
} else {
linear_kernel(
compact_kernel_sdf,
0,
stream,
n_alive,
zero_offset,
rays_tmp.pos,
rays_tmp.distance,
rays_tmp.payload,
rays_current.pos,
rays_current.distance,
rays_current.payload,
m_rays_hit.pos,
m_rays_hit.distance,
m_rays_hit.payload,
aabb,
m_alive_counter,
m_hit_counter
);
}
CUDA_CHECK_THROW(cudaMemcpyAsync(&n_alive, m_alive_counter, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream));
CUDA_CHECK_THROW(cudaStreamSynchronize(stream));
}
if (n_alive == 0) {
break;
}
for (uint32_t j = 0; j < step_size; ++j) {
distance_function(n_alive, rays_current.pos, rays_current.distance, stream);
linear_kernel(
advance_pos_kernel_sdf,
0,
stream,
n_alive,
zero_offset,
rays_current.pos,
rays_current.distance,
rays_current.payload,
aabb,
floor_y,
octree ? octree->nodes_gpu() : nullptr,
octree ? n_octree_levels : 0,
distance_scale,
maximum_distance,
m_shadow_sharpness,
m_trace_shadow_rays ? rays_current.prev_distance : nullptr,
m_trace_shadow_rays ? rays_current.total_distance : nullptr,
m_trace_shadow_rays ? rays_current.min_visibility : nullptr
);
}
i += step_size;
}
uint32_t n_hit;
CUDA_CHECK_THROW(cudaMemcpyAsync(&n_hit, m_hit_counter, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream));
CUDA_CHECK_THROW(cudaStreamSynchronize(stream));
return n_hit;
}
void Testbed::SphereTracer::enlarge(size_t n_elements, cudaStream_t stream) {
n_elements = next_multiple(n_elements, size_t(BATCH_SIZE_GRANULARITY));
auto scratch = allocate_workspace_and_distribute<
vec3,
vec3,
float,
float,
float,
float,
SdfPayload, // m_rays[0]
vec3,
vec3,
float,
float,
float,
float,
SdfPayload, // m_rays[1]
vec3,
vec3,
float,
float,
float,
float,
SdfPayload, // m_rays_hit
uint32_t,
uint32_t>(
stream,
&m_scratch_alloc,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
n_elements,
32, // 2 full cache lines to ensure no overlap
32 // 2 full cache lines to ensure no overlap
);
m_rays[0].set(
std::get<0>(scratch),
std::get<1>(scratch),
std::get<2>(scratch),
std::get<3>(scratch),
std::get<4>(scratch),
std::get<5>(scratch),
std::get<6>(scratch)
);
m_rays[1].set(
std::get<7>(scratch),
std::get<8>(scratch),
std::get<9>(scratch),
std::get<10>(scratch),
std::get<11>(scratch),
std::get<12>(scratch),
std::get<13>(scratch)
);
m_rays_hit.set(
std::get<14>(scratch),
std::get<15>(scratch),
std::get<16>(scratch),
std::get<17>(scratch),
std::get<18>(scratch),
std::get<19>(scratch),
std::get<20>(scratch)
);
m_hit_counter = std::get<21>(scratch);
m_alive_counter = std::get<22>(scratch);
}
void Testbed::FiniteDifferenceNormalsApproximator::enlarge(uint32_t n_elements, cudaStream_t stream) {
n_elements = next_multiple(n_elements, BATCH_SIZE_GRANULARITY);
auto scratch = allocate_workspace_and_distribute<vec3, vec3, vec3, float, float, float, float, float, float>(
stream, &m_scratch_alloc, n_elements, n_elements, n_elements, n_elements, n_elements, n_elements, n_elements, n_elements, n_elements
);
dx = std::get<0>(scratch);
dy = std::get<1>(scratch);
dz = std::get<2>(scratch);
dist_dx_pos = std::get<3>(scratch);
dist_dy_pos = std::get<4>(scratch);
dist_dz_pos = std::get<5>(scratch);
dist_dx_neg = std::get<6>(scratch);
dist_dy_neg = std::get<7>(scratch);
dist_dz_neg = std::get<8>(scratch);
}
void Testbed::FiniteDifferenceNormalsApproximator::normal(
uint32_t n_elements, const distance_fun_t& distance_function, const vec3* pos, vec3* normal, float epsilon, cudaStream_t stream
) {
enlarge(n_elements, stream);
parallel_for_gpu(stream, n_elements, [pos = pos, dx = dx, dy = dy, dz = dz, epsilon] __device__(size_t i) {
vec3 p = pos[i];
dx[i] = vec3{p.x + epsilon, p.y, p.z};
dy[i] = vec3{p.x, p.y + epsilon, p.z};
dz[i] = vec3{p.x, p.y, p.z + epsilon};
});
distance_function(n_elements, dx, dist_dx_pos, stream);
distance_function(n_elements, dy, dist_dy_pos, stream);
distance_function(n_elements, dz, dist_dz_pos, stream);
parallel_for_gpu(stream, n_elements, [pos = pos, dx = dx, dy = dy, dz = dz, epsilon] __device__(size_t i) {
vec3 p = pos[i];
dx[i] = vec3{p.x - epsilon, p.y, p.z};
dy[i] = vec3{p.x, p.y - epsilon, p.z};
dz[i] = vec3{p.x, p.y, p.z - epsilon};
});
distance_function(n_elements, dx, dist_dx_neg, stream);
distance_function(n_elements, dy, dist_dy_neg, stream);
distance_function(n_elements, dz, dist_dz_neg, stream);
parallel_for_gpu(
stream,
n_elements,
[normal = normal,
dist_dx_pos = dist_dx_pos,
dist_dx_neg = dist_dx_neg,
dist_dy_pos = dist_dy_pos,
dist_dy_neg = dist_dy_neg,
dist_dz_pos = dist_dz_pos,
dist_dz_neg = dist_dz_neg] __device__(size_t i) {
normal[i] = {dist_dx_pos[i] - dist_dx_neg[i], dist_dy_pos[i] - dist_dy_neg[i], dist_dz_pos[i] - dist_dz_neg[i]};
}
);
}
void Testbed::render_sdf(
cudaStream_t stream,
CudaDevice& device,
const distance_fun_t& distance_function,
const normals_fun_t& normals_function,
const CudaRenderBufferView& render_buffer,
const vec2& focal_length,
const mat4x3& camera_matrix,
const vec2& screen_center,
const Foveation& foveation,
const Lens& lens,
int visualized_dimension
) {
auto jit_guard = m_network->jit_guard(stream, true);
float plane_z = m_slice_plane_z + m_scale;
if (m_render_mode == ERenderMode::Slice) {
plane_z = -plane_z;
}
auto* octree_ptr = m_sdf.uses_takikawa_encoding || m_sdf.use_triangle_octree ? m_sdf.triangle_octree.get() : nullptr;
SphereTracer tracer;
uint32_t n_octree_levels = octree_ptr ? octree_ptr->depth() : 0;
BoundingBox sdf_bounding_box = m_aabb;
sdf_bounding_box.inflate(m_sdf.zero_offset);
if (m_jit_fusion) {
if (!device.fused_render_kernel()) {
try {
device.set_fused_render_kernel(std::make_unique<CudaRtcKernel>(
"trace_sdf",
fmt::format(
"{}\n#include <neural-graphics-primitives/fused_kernels/trace_sdf.cuh>\n",
m_network->generate_device_function("eval_sdf")
),
all_files(cmrc::ngp::get_filesystem())
));
} catch (const std::runtime_error& e) {
tlog::warning() << e.what();
tlog::warning() << "Disabling JIT fusion.";
m_jit_fusion = false;
}
}
if (device.fused_render_kernel()) {
tracer.set_fused_trace_kernel(device.fused_render_kernel());
}
}
tracer.init_rays_from_camera(
render_buffer.spp,
render_buffer.resolution,
focal_length,
camera_matrix,
screen_center,
m_parallax_shift,
m_snap_to_pixel_centers,
sdf_bounding_box,
get_floor_y(),
m_render_near_distance,
plane_z,
m_aperture_size,
foveation,
m_envmap.inference_view(),
render_buffer.frame_buffer,
render_buffer.depth_buffer,
render_buffer.hidden_area_mask ? render_buffer.hidden_area_mask->const_view() : Buffer2DView<const uint8_t>{},
lens,
octree_ptr,
n_octree_levels,
stream
);
bool gt_raytrace = m_render_ground_truth && m_sdf.groundtruth_mode == ESDFGroundTruthMode::RaytracedMesh;
auto trace = [&](SphereTracer& tracer) {
if (gt_raytrace) {
return tracer.trace_bvh(m_sdf.triangle_bvh.get(), m_sdf.triangles_gpu.data(), stream);
} else {
return tracer.trace(
distance_function,
m_network.get(),
m_sdf.zero_offset,
m_sdf.distance_scale,
m_sdf.maximum_distance,
sdf_bounding_box,
get_floor_y(),
octree_ptr,
n_octree_levels,
stream
);
}
};
uint32_t n_hit;
if (m_render_mode == ERenderMode::Slice) {
n_hit = tracer.n_rays_initialized();
} else {
n_hit = trace(tracer);
}
RaysSdfSoa& rays_hit = m_render_mode == ERenderMode::Slice || gt_raytrace ? tracer.rays_init() : tracer.rays_hit();
if (m_render_mode == ERenderMode::Slice) {
if (visualized_dimension == -1) {
distance_function(n_hit, rays_hit.pos, rays_hit.distance, stream);
extract_dimension_pos_neg_kernel<float><<<n_blocks_linear(n_hit * 3), N_THREADS_LINEAR, 0, stream>>>(
n_hit * 3, 0, 1, 3, rays_hit.distance, CM, (float*)rays_hit.normal
);
} else {
// Store colors in the normal buffer
uint32_t n_elements = next_multiple(n_hit, BATCH_SIZE_GRANULARITY);
GPUMatrix<float> positions_matrix((float*)rays_hit.pos, 3, n_elements);
GPUMatrix<float> colors_matrix((float*)rays_hit.normal, 3, n_elements);
m_network->visualize_activation(stream, m_visualized_layer, visualized_dimension, positions_matrix, colors_matrix);
}
}
ERenderMode render_mode = (visualized_dimension > -1 || m_render_mode == ERenderMode::Slice) ? ERenderMode::EncodingVis : m_render_mode;
if (render_mode == ERenderMode::Shade || render_mode == ERenderMode::Normals) {
if (m_sdf.analytic_normals || gt_raytrace) {
normals_function(n_hit, rays_hit.pos, rays_hit.normal, stream);
} else {
float fd_normals_epsilon = m_sdf.fd_normals_epsilon;
FiniteDifferenceNormalsApproximator fd_normals;
fd_normals.normal(n_hit, distance_function, rays_hit.pos, rays_hit.normal, fd_normals_epsilon, stream);
}
if (render_mode == ERenderMode::Shade && n_hit > 0) {
// Shadow rays towards the sun
SphereTracer shadow_tracer;
shadow_tracer.init_rays_from_data(n_hit, rays_hit, stream);
shadow_tracer.set_fused_trace_kernel(tracer.fused_trace_kernel());
shadow_tracer.set_trace_shadow_rays(true);
shadow_tracer.set_shadow_sharpness(m_sdf.shadow_sharpness);
RaysSdfSoa& shadow_rays_init = shadow_tracer.rays_init();
linear_kernel(
prepare_shadow_rays,
0,
stream,
n_hit,
normalize(m_sun_dir),
shadow_rays_init.pos,
shadow_rays_init.normal,
shadow_rays_init.distance,
shadow_rays_init.prev_distance,
shadow_rays_init.total_distance,
shadow_rays_init.min_visibility,
shadow_rays_init.payload,
sdf_bounding_box,
octree_ptr ? octree_ptr->nodes_gpu() : nullptr,
n_octree_levels
);
uint32_t n_hit_shadow = trace(shadow_tracer);
auto& shadow_rays_hit = gt_raytrace ? shadow_tracer.rays_init() : shadow_tracer.rays_hit();
linear_kernel(
write_shadow_ray_result,
0,
stream,
n_hit_shadow,
sdf_bounding_box,
shadow_rays_hit.pos,
shadow_rays_hit.payload,
shadow_rays_hit.min_visibility,
rays_hit.distance
);
// todo: Reflection rays?
}
} else if (render_mode == ERenderMode::EncodingVis && m_render_mode != ERenderMode::Slice) {
// HACK: Store colors temporarily in the normal buffer
uint32_t n_elements = next_multiple(n_hit, BATCH_SIZE_GRANULARITY);
GPUMatrix<float> positions_matrix((float*)rays_hit.pos, 3, n_elements);
GPUMatrix<float> colors_matrix((float*)rays_hit.normal, 3, n_elements);
m_network->visualize_activation(stream, m_visualized_layer, visualized_dimension, positions_matrix, colors_matrix);
}
linear_kernel(
shade_kernel_sdf,
0,
stream,
n_hit,
m_aabb,
get_floor_y(),
render_mode,
m_sdf.brdf,
normalize(m_sun_dir),
normalize(m_up_dir),
camera_matrix,
rays_hit.pos,
rays_hit.normal,
rays_hit.distance,
rays_hit.payload,
render_buffer.frame_buffer,
render_buffer.depth_buffer
);
if (render_mode == ERenderMode::Cost) {
std::vector<SdfPayload> payloads_final_cpu(n_hit);
CUDA_CHECK_THROW(
cudaMemcpyAsync(payloads_final_cpu.data(), rays_hit.payload, n_hit * sizeof(SdfPayload), cudaMemcpyDeviceToHost, stream)
);
CUDA_CHECK_THROW(cudaStreamSynchronize(stream));
size_t total_n_steps = 0;
for (uint32_t i = 0; i < n_hit; ++i) {
total_n_steps += payloads_final_cpu[i].n_steps;
}
tlog::info() << "Total steps per hit= " << total_n_steps << "/" << n_hit << " = " << ((float)total_n_steps / (float)n_hit);
}
}
std::vector<vec3> load_stl(const fs::path& path) {
std::vector<vec3> vertices;
std::ifstream f{native_string(path), std::ios::in | std::ios::binary};
if (!f) {
throw std::runtime_error{fmt::format("Mesh file '{}' not found", path.str())};
}
uint32_t buf[21] = {};
f.read((char*)buf, 4 * 21);
if (f.gcount() < 4 * 21) {
throw std::runtime_error{fmt::format("Mesh file '{}' too small for STL header", path.str())};
}
uint32_t nfaces = buf[20];
if (memcmp(buf, "solid", 5) == 0 || buf[20] == 0) {
throw std::runtime_error{fmt::format("ASCII STL file '{}' not supported", path.str())};
}
vertices.reserve(nfaces * 3);
for (uint32_t i = 0; i < nfaces; ++i) {
f.read((char*)buf, 50);
if (f.gcount() < 50) {
nfaces = i;
break;
}
vertices.push_back(*(vec3*)(buf + 3));
vertices.push_back(*(vec3*)(buf + 6));
vertices.push_back(*(vec3*)(buf + 9));
}
return vertices;
}
void Testbed::load_mesh(const fs::path& data_path) {
tlog::info() << "Loading mesh from '" << data_path << "'";
auto start = std::chrono::steady_clock::now();
std::vector<vec3> vertices;
if (equals_case_insensitive(data_path.extension(), "obj")) {
vertices = load_obj(data_path.str());
} else if (equals_case_insensitive(data_path.extension(), "stl")) {
vertices = load_stl(data_path.str());
} else {
throw std::runtime_error{"SDF data path must be a mesh in ascii .obj or binary .stl format."};
}
// The expected format is
// [v1.x][v1.y][v1.z][v2.x]...
size_t n_vertices = vertices.size();
size_t n_triangles = n_vertices / 3;
m_raw_aabb.min = vec3(std::numeric_limits<float>::infinity());
m_raw_aabb.max = vec3(-std::numeric_limits<float>::infinity());
for (size_t i = 0; i < n_vertices; ++i) {
m_raw_aabb.enlarge(vertices[i]);
}
// Inflate AABB by 1% to give the network a little wiggle room.
const float inflation = 0.005f;
m_raw_aabb.inflate(length(m_raw_aabb.diag()) * inflation);
m_sdf.mesh_scale = max(m_raw_aabb.diag());
// Normalize vertex coordinates to lie within [0,1]^3.
// This way, none of the constants need to carry around
// bounding box factors.
for (size_t i = 0; i < n_vertices; ++i) {
vertices[i] = (vertices[i] - m_raw_aabb.min - 0.5f * m_raw_aabb.diag()) / m_sdf.mesh_scale + 0.5f;
}
m_aabb = {};
for (size_t i = 0; i < n_vertices; ++i) {
m_aabb.enlarge(vertices[i]);
}
m_aabb.inflate(length(m_aabb.diag()) * inflation);
m_aabb = m_aabb.intersection(BoundingBox{vec3(0.0f), vec3(1.0f)});
m_render_aabb = m_aabb;
m_render_aabb_to_local = mat3::identity();
m_mesh.thresh = 0.f;
m_sdf.triangles_cpu.resize(n_triangles);
for (size_t i = 0; i < n_vertices; i += 3) {
m_sdf.triangles_cpu[i / 3] = {vertices[i + 0], vertices[i + 1], vertices[i + 2]};
}
if (!m_sdf.triangle_bvh) {
m_sdf.triangle_bvh = TriangleBvh::make();
}
m_sdf.triangle_bvh->build(m_sdf.triangles_cpu, 8);
m_sdf.triangles_gpu.resize_and_copy_from_host(m_sdf.triangles_cpu);
m_sdf.triangle_bvh->build_optix(m_sdf.triangles_gpu, m_stream.get());
m_sdf.triangle_octree.reset(new TriangleOctree{});
m_sdf.triangle_octree->build(*m_sdf.triangle_bvh, m_sdf.triangles_cpu, 10);
m_bounding_radius = length(vec3(0.5f));
// Compute discrete probability distribution for later sampling of the mesh's surface
m_sdf.triangle_weights.resize(n_triangles);
for (size_t i = 0; i < n_triangles; ++i) {
m_sdf.triangle_weights[i] = m_sdf.triangles_cpu[i].surface_area();
}
m_sdf.triangle_distribution.build(m_sdf.triangle_weights);
// Move CDF to gpu
m_sdf.triangle_cdf.resize_and_copy_from_host(m_sdf.triangle_distribution.cdf);
// Clear training data as it's no longer representative
// of the previously loaded mesh... but don't clear the network.
// Perhaps it'll look interesting while morphing from one mesh to another.
m_sdf.training.idx = 0;
m_sdf.training.size = 0;
tlog::success() << "Loaded mesh after " << tlog::durationToString(std::chrono::steady_clock::now() - start);
tlog::info() << " n_triangles=" << n_triangles << " aabb=" << m_raw_aabb;
}
void Testbed::generate_training_samples_sdf(vec3* positions, float* distances, uint32_t n_to_generate, cudaStream_t stream, bool uniform_only) {
uint32_t n_to_generate_base = n_to_generate / 8;
const uint32_t n_to_generate_surface_exact = uniform_only ? 0 : n_to_generate_base * 4;
const uint32_t n_to_generate_surface_offset = uniform_only ? 0 : n_to_generate_base * 3;
const uint32_t n_to_generate_uniform = uniform_only ? n_to_generate : n_to_generate_base * 1;
const uint32_t n_to_generate_surface = n_to_generate_surface_exact + n_to_generate_surface_offset;
// Generate uniform 3D samples. Some of these will be transformed to cover the surfaces uniformly. Others will be left as-is.
generate_random_uniform<float>(stream, m_rng, n_to_generate * 3, (float*)positions);
linear_kernel(
sample_uniform_on_triangle_kernel,
0,
stream,
n_to_generate_surface,
m_sdf.triangle_cdf.data(),
(uint32_t)m_sdf.triangle_cdf.size(),
m_sdf.triangles_gpu.data(),
positions
);
// The distances of points on the mesh are zero. Can immediately set.
CUDA_CHECK_THROW(cudaMemsetAsync(distances, 0, n_to_generate_surface_exact * sizeof(float), stream));
// If we have an octree, generate uniform samples within that octree.
// Otherwise, at least confine uniform samples to the AABB.
// (For the uniform_only case, we always use the AABB, then the IoU kernel checks against the octree later)
float stddev = m_bounding_radius / 1024.0f * m_sdf.training.surface_offset_scale;
if (!uniform_only && (m_sdf.uses_takikawa_encoding || m_sdf.use_triangle_octree)) {
linear_kernel(
uniform_octree_sample_kernel,
0,
stream,
n_to_generate_uniform,
m_rng,
m_sdf.triangle_octree->nodes_gpu(),
m_sdf.triangle_octree->n_nodes(),
m_sdf.triangle_octree->depth(),
positions + n_to_generate_surface
);
m_rng.advance();
// If we know the finest discretization of the octree, we can concentrate
// points MUCH closer to the mesh surface
float leaf_size = scalbnf(1.0f, -m_sdf.triangle_octree->depth() + 1);
if (leaf_size < stddev) {
tlog::warning() << "leaf_size < stddev";
stddev = leaf_size;
}
linear_kernel(assign_float, 0, stream, n_to_generate_uniform, length(vec3(leaf_size)) * 1.001f, distances + n_to_generate_surface);
} else {
BoundingBox sdf_aabb = m_aabb;
sdf_aabb.inflate(m_sdf.zero_offset);
linear_kernel(scale_to_aabb_kernel, 0, stream, n_to_generate_uniform, sdf_aabb, positions + n_to_generate_surface);
linear_kernel(assign_float, 0, stream, n_to_generate_uniform, length(sdf_aabb.diag()) * 1.001f, distances + n_to_generate_surface);
}
m_sdf.training.perturbations.enlarge(n_to_generate_surface_offset);
generate_random_logistic<float>(stream, m_rng, n_to_generate_surface_offset * 3, (float*)m_sdf.training.perturbations.data(), 0.0f, stddev);
linear_kernel(
perturb_sdf_samples,
0,
stream,
n_to_generate_surface_offset,
m_sdf.training.perturbations.data(),
positions + n_to_generate_surface_exact,
distances + n_to_generate_surface_exact
);
// The following function expects `distances` to contain an upper bound on the
// true distance. This accelerates lookups.
m_sdf.triangle_bvh->signed_distance_gpu(
n_to_generate_uniform + n_to_generate_surface_offset,
m_sdf.mesh_sdf_mode,
positions + n_to_generate_surface_exact,
distances + n_to_generate_surface_exact,
m_sdf.triangles_gpu.data(),
true,
stream
);
CUDA_CHECK_THROW(cudaStreamSynchronize(stream));
}
__global__ void generate_grid_samples_sdf_uniform(ivec3 res_3d, BoundingBox aabb, const mat3& render_aabb_to_local, vec3* __restrict__ out) {
uint32_t x = threadIdx.x + blockIdx.x * blockDim.x;
uint32_t y = threadIdx.y + blockIdx.y * blockDim.y;
uint32_t z = threadIdx.z + blockIdx.z * blockDim.z;
if (x >= res_3d.x || y >= res_3d.y || z >= res_3d.z) {
return;
}
uint32_t i = x + y * res_3d.x + z * res_3d.x * res_3d.y;
vec3 pos = vec3{(float)x, (float)y, (float)z} * vec3{1.f / res_3d.x, 1.f / res_3d.y, 1.f / res_3d.z};
pos = pos * (aabb.max - aabb.min) + aabb.min;
out[i] = transpose(render_aabb_to_local) * pos;
}
GPUMemory<float> Testbed::get_sdf_gt_on_grid(ivec3 res3d, const BoundingBox& aabb, const mat3& render_aabb_to_local) {
const uint32_t n_elements = (res3d.x * res3d.y * res3d.z);
GPUMemory<float> density(n_elements);
GPUMemoryArena::Allocation alloc;
auto scratch = allocate_workspace_and_distribute<vec3>(m_stream.get(), &alloc, n_elements);
vec3* positions = std::get<0>(scratch);
float* sdf_out = density.data();
const dim3 threads = {16, 8, 1};
const dim3 blocks = {
div_round_up((uint32_t)res3d.x, threads.x), div_round_up((uint32_t)res3d.y, threads.y), div_round_up((uint32_t)res3d.z, threads.z)
};
generate_grid_samples_sdf_uniform<<<blocks, threads, 0, m_stream.get()>>>(res3d, aabb, render_aabb_to_local, positions);
CUDA_CHECK_THROW(cudaStreamSynchronize(m_stream.get()));
m_sdf.triangle_bvh->signed_distance_gpu(
n_elements, m_sdf.mesh_sdf_mode, positions, sdf_out, m_sdf.triangles_gpu.data(), false, m_stream.get()
);
CUDA_CHECK_THROW(cudaStreamSynchronize(m_stream.get()));
/*
std::vector<float> cpudensity(density.size());
std::vector<vec3> cpupositions(n_elements);
density.copy_to_host(cpudensity);
cudaMemcpy(cpupositions.data(),positions,n_elements*12,cudaMemcpyDeviceToHost);
for (int i=0;i<64;++i)
printf("[%0.3f %0.3f %0.3f] -> %0.3f\n", cpupositions[i].x,cpupositions[i].y,cpupositions[i].z,cpudensity[i]);
*/
return density;
}
void Testbed::train_sdf(size_t target_batch_size, bool get_loss_scalar, cudaStream_t stream) {
const uint32_t n_output_dims = 1;
const uint32_t n_input_dims = 3;
if (m_sdf.training.size >= target_batch_size) {
// Auxiliary matrices for training
const uint32_t batch_size = (uint32_t)std::min(m_sdf.training.size, target_batch_size);
// Permute all training records to de-correlate training data
linear_kernel(
shuffle<vec3>,
0,
stream,
m_sdf.training.size,
1,
m_training_step,
m_sdf.training.positions.data(),
m_sdf.training.positions_shuffled.data()
);
linear_kernel(
shuffle<float>,
0,
stream,
m_sdf.training.size,
1,
m_training_step,
m_sdf.training.distances.data(),
m_sdf.training.distances_shuffled.data()
);
GPUMatrix<float> training_target_matrix(m_sdf.training.distances_shuffled.data(), n_output_dims, batch_size);
GPUMatrix<float> training_batch_matrix((float*)(m_sdf.training.positions_shuffled.data()), n_input_dims, batch_size);
auto ctx = m_trainer->training_step(stream, training_batch_matrix, training_target_matrix);
m_training_step++;
if (get_loss_scalar) {
m_loss_scalar.update(m_trainer->loss(stream, *ctx));
}
}
}
void Testbed::training_prep_sdf(uint32_t batch_size, cudaStream_t stream) {
if (m_sdf.training.generate_sdf_data_online) {
m_sdf.training.size = batch_size;
m_sdf.training.positions.enlarge(m_sdf.training.size);
m_sdf.training.positions_shuffled.enlarge(m_sdf.training.size);
m_sdf.training.distances.enlarge(m_sdf.training.size);
m_sdf.training.distances_shuffled.enlarge(m_sdf.training.size);
generate_training_samples_sdf(m_sdf.training.positions.data(), m_sdf.training.distances.data(), batch_size, stream, false);
}
}
// set scale_existing_results_factor=0. to reset any existing results; set it to 1.0 to accumulate more samples onto existing results
// set it to a fraction near 1 to use a sliding EMA
// if blocking is false, then this returns the iou from the *last* call
double Testbed::calculate_iou(uint32_t n_samples, float scale_existing_results_factor, bool blocking, bool force_use_octree) {
cudaStream_t stream = m_stream.get();
uint32_t countercpu[8];
m_sdf.iou_counter.enlarge(8);
if (!blocking) { // when not blocking, returns data from the *last* run, then kicks off work to accumulate some more samples
cudaMemcpy(countercpu, m_sdf.iou_counter.data(), 8 * 4, cudaMemcpyDeviceToHost);
}
if (scale_existing_results_factor < 1.f) {
linear_kernel(scale_iou_counters_kernel, 0, stream, 8, m_sdf.iou_counter.data(), scale_existing_results_factor);
}
while (n_samples > 0) {
uint32_t batch_size = std::min(uint32_t(128 * 128 * 128), n_samples);
m_sdf.training.size = batch_size;
n_samples -= batch_size;
m_sdf.training.positions.enlarge(m_sdf.training.size);
m_sdf.training.distances.enlarge(m_sdf.training.size); // we use this buffer for the GT distances
m_sdf.training.distances_shuffled.enlarge(m_sdf.training.size); // we use the shuffled output for the output of inference
generate_training_samples_sdf(m_sdf.training.positions.data(), m_sdf.training.distances.data(), (uint32_t)(batch_size), stream, true);
GPUMatrix<float> positions_matrix((float*)m_sdf.training.positions.data(), 3, batch_size);
GPUMatrix<float> distances_matrix(m_sdf.training.distances_shuffled.data(), 1, batch_size);
m_network->inference(stream, positions_matrix, distances_matrix);
auto* octree_ptr = (m_sdf.uses_takikawa_encoding || m_sdf.use_triangle_octree || force_use_octree) ? m_sdf.triangle_octree.get() :
nullptr;
linear_kernel(
compare_signs_kernel,
0,
stream,
batch_size,
m_sdf.training.positions.data(),
m_sdf.training.distances.data(), // ref
m_sdf.training.distances_shuffled.data(), // model
m_sdf.iou_counter.data(),
octree_ptr ? octree_ptr->nodes_gpu() : nullptr,
octree_ptr ? octree_ptr->depth() : 0
);
}
if (blocking) {
CUDA_CHECK_THROW(cudaMemcpyAsync(countercpu, m_sdf.iou_counter.data(), 8 * 4, cudaMemcpyDeviceToHost, stream));
CUDA_CHECK_THROW(cudaStreamSynchronize(stream));
}
return countercpu[4] / double(countercpu[5]);
}
} // namespace ngp
