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
+558
@@ -0,0 +1,558 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "antialias.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Helpers.
|
||||
|
||||
#define F32_MAX (3.402823466e+38f)
|
||||
static __forceinline__ __device__ bool same_sign(float a, float b) { return (__float_as_int(a) ^ __float_as_int(b)) >= 0; }
|
||||
static __forceinline__ __device__ bool rational_gt(float n0, float n1, float d0, float d1) { return (n0*d1 > n1*d0) == same_sign(d0, d1); }
|
||||
static __forceinline__ __device__ int max_idx3(float n0, float n1, float n2, float d0, float d1, float d2)
|
||||
{
|
||||
bool g10 = rational_gt(n1, n0, d1, d0);
|
||||
bool g20 = rational_gt(n2, n0, d2, d0);
|
||||
bool g21 = rational_gt(n2, n1, d2, d1);
|
||||
if (g20 && g21) return 2;
|
||||
if (g10) return 1;
|
||||
return 0;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Format of antialiasing work items stored in work buffer. Usually accessed directly as int4.
|
||||
|
||||
struct AAWorkItem
|
||||
{
|
||||
enum
|
||||
{
|
||||
EDGE_MASK = 3, // Edge index in lowest bits.
|
||||
FLAG_DOWN_BIT = 2, // Down instead of right.
|
||||
FLAG_TRI1_BIT = 3, // Edge is from other pixel's triangle.
|
||||
};
|
||||
|
||||
int px, py; // Pixel x, y.
|
||||
unsigned int pz_flags; // High 16 bits = pixel z, low 16 bits = edge index and flags.
|
||||
float alpha; // Antialiasing alpha value. Zero if no AA.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Hash functions. Adapted from public-domain code at http://www.burtleburtle.net/bob/hash/doobs.html
|
||||
|
||||
#define JENKINS_MAGIC (0x9e3779b9u)
|
||||
static __device__ __forceinline__ void jenkins_mix(unsigned int& a, unsigned int& b, unsigned int& c)
|
||||
{
|
||||
a -= b; a -= c; a ^= (c>>13);
|
||||
b -= c; b -= a; b ^= (a<<8);
|
||||
c -= a; c -= b; c ^= (b>>13);
|
||||
a -= b; a -= c; a ^= (c>>12);
|
||||
b -= c; b -= a; b ^= (a<<16);
|
||||
c -= a; c -= b; c ^= (b>>5);
|
||||
a -= b; a -= c; a ^= (c>>3);
|
||||
b -= c; b -= a; b ^= (a<<10);
|
||||
c -= a; c -= b; c ^= (b>>15);
|
||||
}
|
||||
|
||||
// Helper class for hash index iteration. Implements simple odd-skip linear probing with a key-dependent skip.
|
||||
class HashIndex
|
||||
{
|
||||
public:
|
||||
__device__ __forceinline__ HashIndex(const AntialiasKernelParams& p, uint64_t key)
|
||||
{
|
||||
m_mask = (p.allocTriangles << AA_LOG_HASH_ELEMENTS_PER_TRIANGLE(p.allocTriangles)) - 1; // This should work until triangle count exceeds 1073741824.
|
||||
m_idx = (uint32_t)(key & 0xffffffffu);
|
||||
m_skip = (uint32_t)(key >> 32);
|
||||
uint32_t dummy = JENKINS_MAGIC;
|
||||
jenkins_mix(m_idx, m_skip, dummy);
|
||||
m_idx &= m_mask;
|
||||
m_skip &= m_mask;
|
||||
m_skip |= 1;
|
||||
}
|
||||
__device__ __forceinline__ int get(void) const { return m_idx; }
|
||||
__device__ __forceinline__ void next(void) { m_idx = (m_idx + m_skip) & m_mask; }
|
||||
private:
|
||||
uint32_t m_idx, m_skip, m_mask;
|
||||
};
|
||||
|
||||
static __device__ __forceinline__ void hash_insert(const AntialiasKernelParams& p, uint64_t key, int v)
|
||||
{
|
||||
HashIndex idx(p, key);
|
||||
while(1)
|
||||
{
|
||||
uint64_t prev = atomicCAS((unsigned long long*)&p.evHash[idx.get()], 0, (unsigned long long)key);
|
||||
if (prev == 0 || prev == key)
|
||||
break;
|
||||
idx.next();
|
||||
}
|
||||
int* q = (int*)&p.evHash[idx.get()];
|
||||
int a = atomicCAS(q+2, 0, v);
|
||||
if (a != 0 && a != v)
|
||||
atomicCAS(q+3, 0, v);
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ int2 hash_find(const AntialiasKernelParams& p, uint64_t key)
|
||||
{
|
||||
HashIndex idx(p, key);
|
||||
while(1)
|
||||
{
|
||||
uint4 entry = p.evHash[idx.get()];
|
||||
uint64_t k = ((uint64_t)entry.x) | (((uint64_t)entry.y) << 32);
|
||||
if (k == key || k == 0)
|
||||
return make_int2((int)entry.z, (int)entry.w);
|
||||
idx.next();
|
||||
}
|
||||
}
|
||||
|
||||
static __device__ __forceinline__ void evhash_insert_vertex(const AntialiasKernelParams& p, int va, int vb, int vn)
|
||||
{
|
||||
if (va == vb)
|
||||
return;
|
||||
|
||||
uint64_t v0 = (uint32_t)min(va, vb) + 1; // canonical vertex order
|
||||
uint64_t v1 = (uint32_t)max(va, vb) + 1;
|
||||
uint64_t vk = v0 | (v1 << 32); // hash key
|
||||
hash_insert(p, vk, vn + 1);
|
||||
}
|
||||
|
||||
static __forceinline__ __device__ int evhash_find_vertex(const AntialiasKernelParams& p, int va, int vb, int vr)
|
||||
{
|
||||
if (va == vb)
|
||||
return -1;
|
||||
|
||||
uint64_t v0 = (uint32_t)min(va, vb) + 1; // canonical vertex order
|
||||
uint64_t v1 = (uint32_t)max(va, vb) + 1;
|
||||
uint64_t vk = v0 | (v1 << 32); // hash key
|
||||
int2 vn = hash_find(p, vk) - 1;
|
||||
if (vn.x == vr) return vn.y;
|
||||
if (vn.y == vr) return vn.x;
|
||||
return -1;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Mesh analysis kernel.
|
||||
|
||||
__global__ void AntialiasFwdMeshKernel(const AntialiasKernelParams p)
|
||||
{
|
||||
int idx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
if (idx >= p.numTriangles)
|
||||
return;
|
||||
|
||||
int v0 = p.tri[idx * 3 + 0];
|
||||
int v1 = p.tri[idx * 3 + 1];
|
||||
int v2 = p.tri[idx * 3 + 2];
|
||||
|
||||
if (v0 < 0 || v0 >= p.numVertices ||
|
||||
v1 < 0 || v1 >= p.numVertices ||
|
||||
v2 < 0 || v2 >= p.numVertices)
|
||||
return;
|
||||
|
||||
if (v0 == v1 || v1 == v2 || v2 == v0)
|
||||
return;
|
||||
|
||||
evhash_insert_vertex(p, v1, v2, v0);
|
||||
evhash_insert_vertex(p, v2, v0, v1);
|
||||
evhash_insert_vertex(p, v0, v1, v2);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Discontinuity finder kernel.
|
||||
|
||||
__global__ void AntialiasFwdDiscontinuityKernel(const AntialiasKernelParams p)
|
||||
{
|
||||
// Calculate pixel position.
|
||||
int px = blockIdx.x * AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH + threadIdx.x;
|
||||
int py = blockIdx.y * AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT + threadIdx.y;
|
||||
int pz = blockIdx.z;
|
||||
if (px >= p.width || py >= p.height || pz >= p.n)
|
||||
return;
|
||||
|
||||
// Pointer to our TriIdx and fetch.
|
||||
int pidx0 = ((px + p.width * (py + p.height * pz)) << 2) + 3;
|
||||
float tri0 = p.rasterOut[pidx0]; // These can stay as float, as we only compare them against each other.
|
||||
|
||||
// Look right, clamp at edge.
|
||||
int pidx1 = pidx0;
|
||||
if (px < p.width - 1)
|
||||
pidx1 += 4;
|
||||
float tri1 = p.rasterOut[pidx1];
|
||||
|
||||
// Look down, clamp at edge.
|
||||
int pidx2 = pidx0;
|
||||
if (py < p.height - 1)
|
||||
pidx2 += p.width << 2;
|
||||
float tri2 = p.rasterOut[pidx2];
|
||||
|
||||
// Determine amount of work.
|
||||
int count = 0;
|
||||
if (tri1 != tri0) count = 1;
|
||||
if (tri2 != tri0) count += 1;
|
||||
if (!count)
|
||||
return; // Exit warp.
|
||||
|
||||
// Coalesce work counter update to once per CTA.
|
||||
__shared__ int s_temp;
|
||||
s_temp = 0;
|
||||
__syncthreads();
|
||||
int idx = atomicAdd(&s_temp, count);
|
||||
__syncthreads();
|
||||
if (idx == 0)
|
||||
{
|
||||
int base = atomicAdd(&p.workBuffer[0].x, s_temp);
|
||||
s_temp = base + 1; // don't clobber the counters in first slot.
|
||||
}
|
||||
__syncthreads();
|
||||
idx += s_temp;
|
||||
|
||||
// Write to memory.
|
||||
if (tri1 != tri0) p.workBuffer[idx++] = make_int4(px, py, (pz << 16), 0);
|
||||
if (tri2 != tri0) p.workBuffer[idx] = make_int4(px, py, (pz << 16) + (1 << AAWorkItem::FLAG_DOWN_BIT), 0);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Forward analysis kernel.
|
||||
|
||||
__global__ void AntialiasFwdAnalysisKernel(const AntialiasKernelParams p)
|
||||
{
|
||||
__shared__ int s_base;
|
||||
int workCount = p.workBuffer[0].x;
|
||||
for(;;)
|
||||
{
|
||||
// Persistent threads work fetcher.
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0)
|
||||
s_base = atomicAdd(&p.workBuffer[0].y, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK);
|
||||
__syncthreads();
|
||||
int thread_idx = s_base + threadIdx.x;
|
||||
if (thread_idx >= workCount)
|
||||
return;
|
||||
|
||||
int4* pItem = p.workBuffer + thread_idx + 1;
|
||||
int4 item = *pItem;
|
||||
int px = item.x;
|
||||
int py = item.y;
|
||||
int pz = (int)(((unsigned int)item.z) >> 16);
|
||||
int d = (item.z >> AAWorkItem::FLAG_DOWN_BIT) & 1;
|
||||
|
||||
int pixel0 = px + p.width * (py + p.height * pz);
|
||||
int pixel1 = pixel0 + (d ? p.width : 1);
|
||||
float2 zt0 = ((float2*)p.rasterOut)[(pixel0 << 1) + 1];
|
||||
float2 zt1 = ((float2*)p.rasterOut)[(pixel1 << 1) + 1];
|
||||
int tri0 = float_to_triidx(zt0.y) - 1;
|
||||
int tri1 = float_to_triidx(zt1.y) - 1;
|
||||
|
||||
// Select triangle based on background / depth.
|
||||
int tri = (tri0 >= 0) ? tri0 : tri1;
|
||||
if (tri0 >= 0 && tri1 >= 0)
|
||||
tri = (zt0.x < zt1.x) ? tri0 : tri1;
|
||||
if (tri == tri1)
|
||||
{
|
||||
// Calculate with respect to neighbor pixel if chose that triangle.
|
||||
px += 1 - d;
|
||||
py += d;
|
||||
}
|
||||
|
||||
// Bail out if triangle index is corrupt.
|
||||
if (tri < 0 || tri >= p.numTriangles)
|
||||
continue;
|
||||
|
||||
// Fetch vertex indices.
|
||||
int vi0 = p.tri[tri * 3 + 0];
|
||||
int vi1 = p.tri[tri * 3 + 1];
|
||||
int vi2 = p.tri[tri * 3 + 2];
|
||||
|
||||
// Bail out if vertex indices are corrupt.
|
||||
if (vi0 < 0 || vi0 >= p.numVertices ||
|
||||
vi1 < 0 || vi1 >= p.numVertices ||
|
||||
vi2 < 0 || vi2 >= p.numVertices)
|
||||
continue;
|
||||
|
||||
// Fetch opposite vertex indices. Use vertex itself (always silhouette) if no opposite vertex exists.
|
||||
int op0 = evhash_find_vertex(p, vi2, vi1, vi0);
|
||||
int op1 = evhash_find_vertex(p, vi0, vi2, vi1);
|
||||
int op2 = evhash_find_vertex(p, vi1, vi0, vi2);
|
||||
|
||||
// Instance mode: Adjust vertex indices based on minibatch index.
|
||||
if (p.instance_mode)
|
||||
{
|
||||
int vbase = pz * p.numVertices;
|
||||
vi0 += vbase;
|
||||
vi1 += vbase;
|
||||
vi2 += vbase;
|
||||
if (op0 >= 0) op0 += vbase;
|
||||
if (op1 >= 0) op1 += vbase;
|
||||
if (op2 >= 0) op2 += vbase;
|
||||
}
|
||||
|
||||
// Fetch vertex positions.
|
||||
float4 p0 = ((float4*)p.pos)[vi0];
|
||||
float4 p1 = ((float4*)p.pos)[vi1];
|
||||
float4 p2 = ((float4*)p.pos)[vi2];
|
||||
float4 o0 = (op0 < 0) ? p0 : ((float4*)p.pos)[op0];
|
||||
float4 o1 = (op1 < 0) ? p1 : ((float4*)p.pos)[op1];
|
||||
float4 o2 = (op2 < 0) ? p2 : ((float4*)p.pos)[op2];
|
||||
|
||||
// Project vertices to pixel space.
|
||||
float w0 = 1.f / p0.w;
|
||||
float w1 = 1.f / p1.w;
|
||||
float w2 = 1.f / p2.w;
|
||||
float ow0 = 1.f / o0.w;
|
||||
float ow1 = 1.f / o1.w;
|
||||
float ow2 = 1.f / o2.w;
|
||||
float fx = (float)px + .5f - p.xh;
|
||||
float fy = (float)py + .5f - p.yh;
|
||||
float x0 = p0.x * w0 * p.xh - fx;
|
||||
float y0 = p0.y * w0 * p.yh - fy;
|
||||
float x1 = p1.x * w1 * p.xh - fx;
|
||||
float y1 = p1.y * w1 * p.yh - fy;
|
||||
float x2 = p2.x * w2 * p.xh - fx;
|
||||
float y2 = p2.y * w2 * p.yh - fy;
|
||||
float ox0 = o0.x * ow0 * p.xh - fx;
|
||||
float oy0 = o0.y * ow0 * p.yh - fy;
|
||||
float ox1 = o1.x * ow1 * p.xh - fx;
|
||||
float oy1 = o1.y * ow1 * p.yh - fy;
|
||||
float ox2 = o2.x * ow2 * p.xh - fx;
|
||||
float oy2 = o2.y * ow2 * p.yh - fy;
|
||||
|
||||
// Signs to kill non-silhouette edges.
|
||||
float bb = (x1-x0)*(y2-y0) - (x2-x0)*(y1-y0); // Triangle itself.
|
||||
float a0 = (x1-ox0)*(y2-oy0) - (x2-ox0)*(y1-oy0); // Wings.
|
||||
float a1 = (x2-ox1)*(y0-oy1) - (x0-ox1)*(y2-oy1);
|
||||
float a2 = (x0-ox2)*(y1-oy2) - (x1-ox2)*(y0-oy2);
|
||||
|
||||
// If no matching signs anywhere, skip the rest.
|
||||
if (same_sign(a0, bb) || same_sign(a1, bb) || same_sign(a2, bb))
|
||||
{
|
||||
// XY flip for horizontal edges.
|
||||
if (d)
|
||||
{
|
||||
swap(x0, y0);
|
||||
swap(x1, y1);
|
||||
swap(x2, y2);
|
||||
}
|
||||
|
||||
float dx0 = x2 - x1;
|
||||
float dx1 = x0 - x2;
|
||||
float dx2 = x1 - x0;
|
||||
float dy0 = y2 - y1;
|
||||
float dy1 = y0 - y2;
|
||||
float dy2 = y1 - y0;
|
||||
|
||||
// Check if an edge crosses between us and the neighbor pixel.
|
||||
float dc = -F32_MAX;
|
||||
float ds = (tri == tri0) ? 1.f : -1.f;
|
||||
float d0 = ds * (x1*dy0 - y1*dx0);
|
||||
float d1 = ds * (x2*dy1 - y2*dx1);
|
||||
float d2 = ds * (x0*dy2 - y0*dx2);
|
||||
|
||||
if (same_sign(y1, y2)) d0 = -F32_MAX, dy0 = 1.f;
|
||||
if (same_sign(y2, y0)) d1 = -F32_MAX, dy1 = 1.f;
|
||||
if (same_sign(y0, y1)) d2 = -F32_MAX, dy2 = 1.f;
|
||||
|
||||
int di = max_idx3(d0, d1, d2, dy0, dy1, dy2);
|
||||
if (di == 0 && same_sign(a0, bb) && fabsf(dy0) >= fabsf(dx0)) dc = d0 / dy0;
|
||||
if (di == 1 && same_sign(a1, bb) && fabsf(dy1) >= fabsf(dx1)) dc = d1 / dy1;
|
||||
if (di == 2 && same_sign(a2, bb) && fabsf(dy2) >= fabsf(dx2)) dc = d2 / dy2;
|
||||
float eps = .0625f; // Expect no more than 1/16 pixel inaccuracy.
|
||||
|
||||
// Adjust output image if a suitable edge was found.
|
||||
if (dc > -eps && dc < 1.f + eps)
|
||||
{
|
||||
dc = fminf(fmaxf(dc, 0.f), 1.f);
|
||||
float alpha = ds * (.5f - dc);
|
||||
const float* pColor0 = p.color + pixel0 * p.channels;
|
||||
const float* pColor1 = p.color + pixel1 * p.channels;
|
||||
float* pOutput = p.output + (alpha > 0.f ? pixel0 : pixel1) * p.channels;
|
||||
for (int i=0; i < p.channels; i++)
|
||||
atomicAdd(&pOutput[i], alpha * (pColor1[i] - pColor0[i]));
|
||||
|
||||
// Rewrite the work item's flags and alpha. Keep original px, py.
|
||||
unsigned int flags = pz << 16;
|
||||
flags |= di;
|
||||
flags |= d << AAWorkItem::FLAG_DOWN_BIT;
|
||||
flags |= (__float_as_uint(ds) >> 31) << AAWorkItem::FLAG_TRI1_BIT;
|
||||
((int2*)pItem)[1] = make_int2(flags, __float_as_int(alpha));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient kernel.
|
||||
|
||||
__global__ void AntialiasGradKernel(const AntialiasKernelParams p)
|
||||
{
|
||||
// Temporary space for coalesced atomics.
|
||||
CA_DECLARE_TEMP(AA_GRAD_KERNEL_THREADS_PER_BLOCK);
|
||||
__shared__ int s_base; // Work counter communication across entire CTA.
|
||||
|
||||
int workCount = p.workBuffer[0].x;
|
||||
|
||||
for(;;)
|
||||
{
|
||||
// Persistent threads work fetcher.
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0)
|
||||
s_base = atomicAdd(&p.workBuffer[0].y, AA_GRAD_KERNEL_THREADS_PER_BLOCK);
|
||||
__syncthreads();
|
||||
int thread_idx = s_base + threadIdx.x;
|
||||
if (thread_idx >= workCount)
|
||||
return;
|
||||
|
||||
// Read work item filled out by forward kernel.
|
||||
int4 item = p.workBuffer[thread_idx + 1];
|
||||
unsigned int amask = __ballot_sync(0xffffffffu, item.w);
|
||||
if (item.w == 0)
|
||||
continue; // No effect.
|
||||
|
||||
// Unpack work item and replicate setup from forward analysis kernel.
|
||||
int px = item.x;
|
||||
int py = item.y;
|
||||
int pz = (int)(((unsigned int)item.z) >> 16);
|
||||
int d = (item.z >> AAWorkItem::FLAG_DOWN_BIT) & 1;
|
||||
float alpha = __int_as_float(item.w);
|
||||
int tri1 = (item.z >> AAWorkItem::FLAG_TRI1_BIT) & 1;
|
||||
int di = item.z & AAWorkItem::EDGE_MASK;
|
||||
float ds = __int_as_float(__float_as_int(1.0) | (tri1 << 31));
|
||||
int pixel0 = px + p.width * (py + p.height * pz);
|
||||
int pixel1 = pixel0 + (d ? p.width : 1);
|
||||
int tri = float_to_triidx(p.rasterOut[((tri1 ? pixel1 : pixel0) << 2) + 3]) - 1;
|
||||
if (tri1)
|
||||
{
|
||||
px += 1 - d;
|
||||
py += d;
|
||||
}
|
||||
|
||||
// Bail out if triangle index is corrupt.
|
||||
bool triFail = (tri < 0 || tri >= p.numTriangles);
|
||||
amask = __ballot_sync(amask, !triFail);
|
||||
if (triFail)
|
||||
continue;
|
||||
|
||||
// Outgoing color gradients.
|
||||
float* pGrad0 = p.gradColor + pixel0 * p.channels;
|
||||
float* pGrad1 = p.gradColor + pixel1 * p.channels;
|
||||
|
||||
// Incoming color gradients.
|
||||
const float* pDy = p.dy + (alpha > 0.f ? pixel0 : pixel1) * p.channels;
|
||||
|
||||
// Position gradient weight based on colors and incoming gradients.
|
||||
float dd = 0.f;
|
||||
const float* pColor0 = p.color + pixel0 * p.channels;
|
||||
const float* pColor1 = p.color + pixel1 * p.channels;
|
||||
|
||||
// Loop over channels and accumulate.
|
||||
for (int i=0; i < p.channels; i++)
|
||||
{
|
||||
float dy = pDy[i];
|
||||
if (dy != 0.f)
|
||||
{
|
||||
// Update position gradient weight.
|
||||
dd += dy * (pColor1[i] - pColor0[i]);
|
||||
|
||||
// Update color gradients. No coalescing because all have different targets.
|
||||
float v = alpha * dy;
|
||||
atomicAdd(&pGrad0[i], -v);
|
||||
atomicAdd(&pGrad1[i], v);
|
||||
}
|
||||
}
|
||||
|
||||
// If position weight is zero, skip the rest.
|
||||
bool noGrad = (dd == 0.f);
|
||||
amask = __ballot_sync(amask, !noGrad);
|
||||
if (noGrad)
|
||||
continue;
|
||||
|
||||
// Fetch vertex indices of the active edge and their positions.
|
||||
int i1 = (di < 2) ? (di + 1) : 0;
|
||||
int i2 = (i1 < 2) ? (i1 + 1) : 0;
|
||||
int vi1 = p.tri[3 * tri + i1];
|
||||
int vi2 = p.tri[3 * tri + i2];
|
||||
|
||||
// Bail out if vertex indices are corrupt.
|
||||
bool vtxFail = (vi1 < 0 || vi1 >= p.numVertices || vi2 < 0 || vi2 >= p.numVertices);
|
||||
amask = __ballot_sync(amask, !vtxFail);
|
||||
if (vtxFail)
|
||||
continue;
|
||||
|
||||
// Instance mode: Adjust vertex indices based on minibatch index.
|
||||
if (p.instance_mode)
|
||||
{
|
||||
vi1 += pz * p.numVertices;
|
||||
vi2 += pz * p.numVertices;
|
||||
}
|
||||
|
||||
// Fetch vertex positions.
|
||||
float4 p1 = ((float4*)p.pos)[vi1];
|
||||
float4 p2 = ((float4*)p.pos)[vi2];
|
||||
|
||||
// Project vertices to pixel space.
|
||||
float pxh = p.xh;
|
||||
float pyh = p.yh;
|
||||
float fx = (float)px + .5f - pxh;
|
||||
float fy = (float)py + .5f - pyh;
|
||||
|
||||
// XY flip for horizontal edges.
|
||||
if (d)
|
||||
{
|
||||
swap(p1.x, p1.y);
|
||||
swap(p2.x, p2.y);
|
||||
swap(pxh, pyh);
|
||||
swap(fx, fy);
|
||||
}
|
||||
|
||||
// Gradient calculation setup.
|
||||
float w1 = 1.f / p1.w;
|
||||
float w2 = 1.f / p2.w;
|
||||
float x1 = p1.x * w1 * pxh - fx;
|
||||
float y1 = p1.y * w1 * pyh - fy;
|
||||
float x2 = p2.x * w2 * pxh - fx;
|
||||
float y2 = p2.y * w2 * pyh - fy;
|
||||
float dx = x2 - x1;
|
||||
float dy = y2 - y1;
|
||||
float db = x1*dy - y1*dx;
|
||||
|
||||
// Compute inverse delta-y with epsilon.
|
||||
float ep = copysignf(1e-3f, dy); // ~1/1000 pixel.
|
||||
float iy = 1.f / (dy + ep);
|
||||
|
||||
// Compute position gradients.
|
||||
float dby = db * iy;
|
||||
float iw1 = -w1 * iy * dd;
|
||||
float iw2 = w2 * iy * dd;
|
||||
float gp1x = iw1 * pxh * y2;
|
||||
float gp2x = iw2 * pxh * y1;
|
||||
float gp1y = iw1 * pyh * (dby - x2);
|
||||
float gp2y = iw2 * pyh * (dby - x1);
|
||||
float gp1w = -(p1.x * gp1x + p1.y * gp1y) * w1;
|
||||
float gp2w = -(p2.x * gp2x + p2.y * gp2y) * w2;
|
||||
|
||||
// XY flip the gradients.
|
||||
if (d)
|
||||
{
|
||||
swap(gp1x, gp1y);
|
||||
swap(gp2x, gp2y);
|
||||
}
|
||||
|
||||
// Kill position gradients if alpha was saturated.
|
||||
if (fabsf(alpha) >= 0.5f)
|
||||
{
|
||||
gp1x = gp1y = gp1w = 0.f;
|
||||
gp2x = gp2y = gp2w = 0.f;
|
||||
}
|
||||
|
||||
// Initialize coalesced atomics. Match both triangle ID and edge index.
|
||||
// Also note that some threads may be inactive.
|
||||
CA_SET_GROUP_MASK(tri ^ (di << 30), amask);
|
||||
|
||||
// Accumulate gradients.
|
||||
caAtomicAdd3_xyw(p.gradPos + 4 * vi1, gp1x, gp1y, gp1w);
|
||||
caAtomicAdd3_xyw(p.gradPos + 4 * vi2, gp2x, gp2y, gp2w);
|
||||
}
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+50
@@ -0,0 +1,50 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#pragma once
|
||||
#include "common.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Constants and helpers.
|
||||
|
||||
#define AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH 32
|
||||
#define AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT 8
|
||||
#define AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK 256
|
||||
#define AA_MESH_KERNEL_THREADS_PER_BLOCK 256
|
||||
#define AA_HASH_ELEMENTS_PER_TRIANGLE(alloc) ((alloc) >= (2 << 25) ? 4 : 8) // With more than 16777216 triangles (alloc >= 33554432) use smallest possible value of 4 to conserve memory, otherwise use 8 for fewer collisions.
|
||||
#define AA_LOG_HASH_ELEMENTS_PER_TRIANGLE(alloc) ((alloc) >= (2 << 25) ? 2 : 3)
|
||||
#define AA_GRAD_KERNEL_THREADS_PER_BLOCK 256
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// CUDA kernel params.
|
||||
|
||||
struct AntialiasKernelParams
|
||||
{
|
||||
const float* color; // Incoming color buffer.
|
||||
const float* rasterOut; // Incoming rasterizer output buffer.
|
||||
const int* tri; // Incoming triangle buffer.
|
||||
const float* pos; // Incoming position buffer.
|
||||
float* output; // Output buffer of forward kernel.
|
||||
const float* dy; // Incoming gradients.
|
||||
float* gradColor; // Output buffer, color gradient.
|
||||
float* gradPos; // Output buffer, position gradient.
|
||||
int4* workBuffer; // Buffer for storing intermediate work items. First item reserved for counters.
|
||||
uint4* evHash; // Edge-vertex hash.
|
||||
int allocTriangles; // Number of triangles accommodated by evHash. Always power of two.
|
||||
int numTriangles; // Number of triangles.
|
||||
int numVertices; // Number of vertices.
|
||||
int width; // Input width.
|
||||
int height; // Input height.
|
||||
int n; // Minibatch size.
|
||||
int channels; // Channel count in color input.
|
||||
float xh, yh; // Transfer to pixel space.
|
||||
int instance_mode; // 0=normal, 1=instance mode.
|
||||
int tri_const; // 1 if triangle array is known to be constant.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Block and grid size calculators for kernel launches.
|
||||
|
||||
dim3 getLaunchBlockSize(int maxWidth, int maxHeight, int width, int height)
|
||||
{
|
||||
int maxThreads = maxWidth * maxHeight;
|
||||
if (maxThreads <= 1 || (width * height) <= 1)
|
||||
return dim3(1, 1, 1); // Degenerate.
|
||||
|
||||
// Start from max size.
|
||||
int bw = maxWidth;
|
||||
int bh = maxHeight;
|
||||
|
||||
// Optimizations for weirdly sized buffers.
|
||||
if (width < bw)
|
||||
{
|
||||
// Decrease block width to smallest power of two that covers the buffer width.
|
||||
while ((bw >> 1) >= width)
|
||||
bw >>= 1;
|
||||
|
||||
// Maximize height.
|
||||
bh = maxThreads / bw;
|
||||
if (bh > height)
|
||||
bh = height;
|
||||
}
|
||||
else if (height < bh)
|
||||
{
|
||||
// Halve height and double width until fits completely inside buffer vertically.
|
||||
while (bh > height)
|
||||
{
|
||||
bh >>= 1;
|
||||
if (bw < width)
|
||||
bw <<= 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Done.
|
||||
return dim3(bw, bh, 1);
|
||||
}
|
||||
|
||||
dim3 getLaunchGridSize(dim3 blockSize, int width, int height, int depth)
|
||||
{
|
||||
dim3 gridSize;
|
||||
gridSize.x = (width - 1) / blockSize.x + 1;
|
||||
gridSize.y = (height - 1) / blockSize.y + 1;
|
||||
gridSize.z = (depth - 1) / blockSize.z + 1;
|
||||
return gridSize;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+263
@@ -0,0 +1,263 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#pragma once
|
||||
#include <cuda.h>
|
||||
#include <stdint.h>
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// C++ helper function prototypes.
|
||||
|
||||
dim3 getLaunchBlockSize(int maxWidth, int maxHeight, int width, int height);
|
||||
dim3 getLaunchGridSize(dim3 blockSize, int width, int height, int depth);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// The rest is CUDA device code specific stuff.
|
||||
|
||||
#ifdef __CUDACC__
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Helpers for CUDA vector types.
|
||||
|
||||
static __device__ __forceinline__ float2& operator*= (float2& a, const float2& b) { a.x *= b.x; a.y *= b.y; return a; }
|
||||
static __device__ __forceinline__ float2& operator+= (float2& a, const float2& b) { a.x += b.x; a.y += b.y; return a; }
|
||||
static __device__ __forceinline__ float2& operator-= (float2& a, const float2& b) { a.x -= b.x; a.y -= b.y; return a; }
|
||||
static __device__ __forceinline__ float2& operator*= (float2& a, float b) { a.x *= b; a.y *= b; return a; }
|
||||
static __device__ __forceinline__ float2& operator+= (float2& a, float b) { a.x += b; a.y += b; return a; }
|
||||
static __device__ __forceinline__ float2& operator-= (float2& a, float b) { a.x -= b; a.y -= b; return a; }
|
||||
static __device__ __forceinline__ float2 operator* (const float2& a, const float2& b) { return make_float2(a.x * b.x, a.y * b.y); }
|
||||
static __device__ __forceinline__ float2 operator+ (const float2& a, const float2& b) { return make_float2(a.x + b.x, a.y + b.y); }
|
||||
static __device__ __forceinline__ float2 operator- (const float2& a, const float2& b) { return make_float2(a.x - b.x, a.y - b.y); }
|
||||
static __device__ __forceinline__ float2 operator* (const float2& a, float b) { return make_float2(a.x * b, a.y * b); }
|
||||
static __device__ __forceinline__ float2 operator+ (const float2& a, float b) { return make_float2(a.x + b, a.y + b); }
|
||||
static __device__ __forceinline__ float2 operator- (const float2& a, float b) { return make_float2(a.x - b, a.y - b); }
|
||||
static __device__ __forceinline__ float2 operator* (float a, const float2& b) { return make_float2(a * b.x, a * b.y); }
|
||||
static __device__ __forceinline__ float2 operator+ (float a, const float2& b) { return make_float2(a + b.x, a + b.y); }
|
||||
static __device__ __forceinline__ float2 operator- (float a, const float2& b) { return make_float2(a - b.x, a - b.y); }
|
||||
static __device__ __forceinline__ float2 operator- (const float2& a) { return make_float2(-a.x, -a.y); }
|
||||
static __device__ __forceinline__ float3& operator*= (float3& a, const float3& b) { a.x *= b.x; a.y *= b.y; a.z *= b.z; return a; }
|
||||
static __device__ __forceinline__ float3& operator+= (float3& a, const float3& b) { a.x += b.x; a.y += b.y; a.z += b.z; return a; }
|
||||
static __device__ __forceinline__ float3& operator-= (float3& a, const float3& b) { a.x -= b.x; a.y -= b.y; a.z -= b.z; return a; }
|
||||
static __device__ __forceinline__ float3& operator*= (float3& a, float b) { a.x *= b; a.y *= b; a.z *= b; return a; }
|
||||
static __device__ __forceinline__ float3& operator+= (float3& a, float b) { a.x += b; a.y += b; a.z += b; return a; }
|
||||
static __device__ __forceinline__ float3& operator-= (float3& a, float b) { a.x -= b; a.y -= b; a.z -= b; return a; }
|
||||
static __device__ __forceinline__ float3 operator* (const float3& a, const float3& b) { return make_float3(a.x * b.x, a.y * b.y, a.z * b.z); }
|
||||
static __device__ __forceinline__ float3 operator+ (const float3& a, const float3& b) { return make_float3(a.x + b.x, a.y + b.y, a.z + b.z); }
|
||||
static __device__ __forceinline__ float3 operator- (const float3& a, const float3& b) { return make_float3(a.x - b.x, a.y - b.y, a.z - b.z); }
|
||||
static __device__ __forceinline__ float3 operator* (const float3& a, float b) { return make_float3(a.x * b, a.y * b, a.z * b); }
|
||||
static __device__ __forceinline__ float3 operator+ (const float3& a, float b) { return make_float3(a.x + b, a.y + b, a.z + b); }
|
||||
static __device__ __forceinline__ float3 operator- (const float3& a, float b) { return make_float3(a.x - b, a.y - b, a.z - b); }
|
||||
static __device__ __forceinline__ float3 operator* (float a, const float3& b) { return make_float3(a * b.x, a * b.y, a * b.z); }
|
||||
static __device__ __forceinline__ float3 operator+ (float a, const float3& b) { return make_float3(a + b.x, a + b.y, a + b.z); }
|
||||
static __device__ __forceinline__ float3 operator- (float a, const float3& b) { return make_float3(a - b.x, a - b.y, a - b.z); }
|
||||
static __device__ __forceinline__ float3 operator- (const float3& a) { return make_float3(-a.x, -a.y, -a.z); }
|
||||
static __device__ __forceinline__ float4& operator*= (float4& a, const float4& b) { a.x *= b.x; a.y *= b.y; a.z *= b.z; a.w *= b.w; return a; }
|
||||
static __device__ __forceinline__ float4& operator+= (float4& a, const float4& b) { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; return a; }
|
||||
static __device__ __forceinline__ float4& operator-= (float4& a, const float4& b) { a.x -= b.x; a.y -= b.y; a.z -= b.z; a.w -= b.w; return a; }
|
||||
static __device__ __forceinline__ float4& operator*= (float4& a, float b) { a.x *= b; a.y *= b; a.z *= b; a.w *= b; return a; }
|
||||
static __device__ __forceinline__ float4& operator+= (float4& a, float b) { a.x += b; a.y += b; a.z += b; a.w += b; return a; }
|
||||
static __device__ __forceinline__ float4& operator-= (float4& a, float b) { a.x -= b; a.y -= b; a.z -= b; a.w -= b; return a; }
|
||||
static __device__ __forceinline__ float4 operator* (const float4& a, const float4& b) { return make_float4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w); }
|
||||
static __device__ __forceinline__ float4 operator+ (const float4& a, const float4& b) { return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
|
||||
static __device__ __forceinline__ float4 operator- (const float4& a, const float4& b) { return make_float4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w); }
|
||||
static __device__ __forceinline__ float4 operator* (const float4& a, float b) { return make_float4(a.x * b, a.y * b, a.z * b, a.w * b); }
|
||||
static __device__ __forceinline__ float4 operator+ (const float4& a, float b) { return make_float4(a.x + b, a.y + b, a.z + b, a.w + b); }
|
||||
static __device__ __forceinline__ float4 operator- (const float4& a, float b) { return make_float4(a.x - b, a.y - b, a.z - b, a.w - b); }
|
||||
static __device__ __forceinline__ float4 operator* (float a, const float4& b) { return make_float4(a * b.x, a * b.y, a * b.z, a * b.w); }
|
||||
static __device__ __forceinline__ float4 operator+ (float a, const float4& b) { return make_float4(a + b.x, a + b.y, a + b.z, a + b.w); }
|
||||
static __device__ __forceinline__ float4 operator- (float a, const float4& b) { return make_float4(a - b.x, a - b.y, a - b.z, a - b.w); }
|
||||
static __device__ __forceinline__ float4 operator- (const float4& a) { return make_float4(-a.x, -a.y, -a.z, -a.w); }
|
||||
static __device__ __forceinline__ int2& operator*= (int2& a, const int2& b) { a.x *= b.x; a.y *= b.y; return a; }
|
||||
static __device__ __forceinline__ int2& operator+= (int2& a, const int2& b) { a.x += b.x; a.y += b.y; return a; }
|
||||
static __device__ __forceinline__ int2& operator-= (int2& a, const int2& b) { a.x -= b.x; a.y -= b.y; return a; }
|
||||
static __device__ __forceinline__ int2& operator*= (int2& a, int b) { a.x *= b; a.y *= b; return a; }
|
||||
static __device__ __forceinline__ int2& operator+= (int2& a, int b) { a.x += b; a.y += b; return a; }
|
||||
static __device__ __forceinline__ int2& operator-= (int2& a, int b) { a.x -= b; a.y -= b; return a; }
|
||||
static __device__ __forceinline__ int2 operator* (const int2& a, const int2& b) { return make_int2(a.x * b.x, a.y * b.y); }
|
||||
static __device__ __forceinline__ int2 operator+ (const int2& a, const int2& b) { return make_int2(a.x + b.x, a.y + b.y); }
|
||||
static __device__ __forceinline__ int2 operator- (const int2& a, const int2& b) { return make_int2(a.x - b.x, a.y - b.y); }
|
||||
static __device__ __forceinline__ int2 operator* (const int2& a, int b) { return make_int2(a.x * b, a.y * b); }
|
||||
static __device__ __forceinline__ int2 operator+ (const int2& a, int b) { return make_int2(a.x + b, a.y + b); }
|
||||
static __device__ __forceinline__ int2 operator- (const int2& a, int b) { return make_int2(a.x - b, a.y - b); }
|
||||
static __device__ __forceinline__ int2 operator* (int a, const int2& b) { return make_int2(a * b.x, a * b.y); }
|
||||
static __device__ __forceinline__ int2 operator+ (int a, const int2& b) { return make_int2(a + b.x, a + b.y); }
|
||||
static __device__ __forceinline__ int2 operator- (int a, const int2& b) { return make_int2(a - b.x, a - b.y); }
|
||||
static __device__ __forceinline__ int2 operator- (const int2& a) { return make_int2(-a.x, -a.y); }
|
||||
static __device__ __forceinline__ int3& operator*= (int3& a, const int3& b) { a.x *= b.x; a.y *= b.y; a.z *= b.z; return a; }
|
||||
static __device__ __forceinline__ int3& operator+= (int3& a, const int3& b) { a.x += b.x; a.y += b.y; a.z += b.z; return a; }
|
||||
static __device__ __forceinline__ int3& operator-= (int3& a, const int3& b) { a.x -= b.x; a.y -= b.y; a.z -= b.z; return a; }
|
||||
static __device__ __forceinline__ int3& operator*= (int3& a, int b) { a.x *= b; a.y *= b; a.z *= b; return a; }
|
||||
static __device__ __forceinline__ int3& operator+= (int3& a, int b) { a.x += b; a.y += b; a.z += b; return a; }
|
||||
static __device__ __forceinline__ int3& operator-= (int3& a, int b) { a.x -= b; a.y -= b; a.z -= b; return a; }
|
||||
static __device__ __forceinline__ int3 operator* (const int3& a, const int3& b) { return make_int3(a.x * b.x, a.y * b.y, a.z * b.z); }
|
||||
static __device__ __forceinline__ int3 operator+ (const int3& a, const int3& b) { return make_int3(a.x + b.x, a.y + b.y, a.z + b.z); }
|
||||
static __device__ __forceinline__ int3 operator- (const int3& a, const int3& b) { return make_int3(a.x - b.x, a.y - b.y, a.z - b.z); }
|
||||
static __device__ __forceinline__ int3 operator* (const int3& a, int b) { return make_int3(a.x * b, a.y * b, a.z * b); }
|
||||
static __device__ __forceinline__ int3 operator+ (const int3& a, int b) { return make_int3(a.x + b, a.y + b, a.z + b); }
|
||||
static __device__ __forceinline__ int3 operator- (const int3& a, int b) { return make_int3(a.x - b, a.y - b, a.z - b); }
|
||||
static __device__ __forceinline__ int3 operator* (int a, const int3& b) { return make_int3(a * b.x, a * b.y, a * b.z); }
|
||||
static __device__ __forceinline__ int3 operator+ (int a, const int3& b) { return make_int3(a + b.x, a + b.y, a + b.z); }
|
||||
static __device__ __forceinline__ int3 operator- (int a, const int3& b) { return make_int3(a - b.x, a - b.y, a - b.z); }
|
||||
static __device__ __forceinline__ int3 operator- (const int3& a) { return make_int3(-a.x, -a.y, -a.z); }
|
||||
static __device__ __forceinline__ int4& operator*= (int4& a, const int4& b) { a.x *= b.x; a.y *= b.y; a.z *= b.z; a.w *= b.w; return a; }
|
||||
static __device__ __forceinline__ int4& operator+= (int4& a, const int4& b) { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; return a; }
|
||||
static __device__ __forceinline__ int4& operator-= (int4& a, const int4& b) { a.x -= b.x; a.y -= b.y; a.z -= b.z; a.w -= b.w; return a; }
|
||||
static __device__ __forceinline__ int4& operator*= (int4& a, int b) { a.x *= b; a.y *= b; a.z *= b; a.w *= b; return a; }
|
||||
static __device__ __forceinline__ int4& operator+= (int4& a, int b) { a.x += b; a.y += b; a.z += b; a.w += b; return a; }
|
||||
static __device__ __forceinline__ int4& operator-= (int4& a, int b) { a.x -= b; a.y -= b; a.z -= b; a.w -= b; return a; }
|
||||
static __device__ __forceinline__ int4 operator* (const int4& a, const int4& b) { return make_int4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w); }
|
||||
static __device__ __forceinline__ int4 operator+ (const int4& a, const int4& b) { return make_int4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
|
||||
static __device__ __forceinline__ int4 operator- (const int4& a, const int4& b) { return make_int4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w); }
|
||||
static __device__ __forceinline__ int4 operator* (const int4& a, int b) { return make_int4(a.x * b, a.y * b, a.z * b, a.w * b); }
|
||||
static __device__ __forceinline__ int4 operator+ (const int4& a, int b) { return make_int4(a.x + b, a.y + b, a.z + b, a.w + b); }
|
||||
static __device__ __forceinline__ int4 operator- (const int4& a, int b) { return make_int4(a.x - b, a.y - b, a.z - b, a.w - b); }
|
||||
static __device__ __forceinline__ int4 operator* (int a, const int4& b) { return make_int4(a * b.x, a * b.y, a * b.z, a * b.w); }
|
||||
static __device__ __forceinline__ int4 operator+ (int a, const int4& b) { return make_int4(a + b.x, a + b.y, a + b.z, a + b.w); }
|
||||
static __device__ __forceinline__ int4 operator- (int a, const int4& b) { return make_int4(a - b.x, a - b.y, a - b.z, a - b.w); }
|
||||
static __device__ __forceinline__ int4 operator- (const int4& a) { return make_int4(-a.x, -a.y, -a.z, -a.w); }
|
||||
static __device__ __forceinline__ uint2& operator*= (uint2& a, const uint2& b) { a.x *= b.x; a.y *= b.y; return a; }
|
||||
static __device__ __forceinline__ uint2& operator+= (uint2& a, const uint2& b) { a.x += b.x; a.y += b.y; return a; }
|
||||
static __device__ __forceinline__ uint2& operator-= (uint2& a, const uint2& b) { a.x -= b.x; a.y -= b.y; return a; }
|
||||
static __device__ __forceinline__ uint2& operator*= (uint2& a, unsigned int b) { a.x *= b; a.y *= b; return a; }
|
||||
static __device__ __forceinline__ uint2& operator+= (uint2& a, unsigned int b) { a.x += b; a.y += b; return a; }
|
||||
static __device__ __forceinline__ uint2& operator-= (uint2& a, unsigned int b) { a.x -= b; a.y -= b; return a; }
|
||||
static __device__ __forceinline__ uint2 operator* (const uint2& a, const uint2& b) { return make_uint2(a.x * b.x, a.y * b.y); }
|
||||
static __device__ __forceinline__ uint2 operator+ (const uint2& a, const uint2& b) { return make_uint2(a.x + b.x, a.y + b.y); }
|
||||
static __device__ __forceinline__ uint2 operator- (const uint2& a, const uint2& b) { return make_uint2(a.x - b.x, a.y - b.y); }
|
||||
static __device__ __forceinline__ uint2 operator* (const uint2& a, unsigned int b) { return make_uint2(a.x * b, a.y * b); }
|
||||
static __device__ __forceinline__ uint2 operator+ (const uint2& a, unsigned int b) { return make_uint2(a.x + b, a.y + b); }
|
||||
static __device__ __forceinline__ uint2 operator- (const uint2& a, unsigned int b) { return make_uint2(a.x - b, a.y - b); }
|
||||
static __device__ __forceinline__ uint2 operator* (unsigned int a, const uint2& b) { return make_uint2(a * b.x, a * b.y); }
|
||||
static __device__ __forceinline__ uint2 operator+ (unsigned int a, const uint2& b) { return make_uint2(a + b.x, a + b.y); }
|
||||
static __device__ __forceinline__ uint2 operator- (unsigned int a, const uint2& b) { return make_uint2(a - b.x, a - b.y); }
|
||||
static __device__ __forceinline__ uint3& operator*= (uint3& a, const uint3& b) { a.x *= b.x; a.y *= b.y; a.z *= b.z; return a; }
|
||||
static __device__ __forceinline__ uint3& operator+= (uint3& a, const uint3& b) { a.x += b.x; a.y += b.y; a.z += b.z; return a; }
|
||||
static __device__ __forceinline__ uint3& operator-= (uint3& a, const uint3& b) { a.x -= b.x; a.y -= b.y; a.z -= b.z; return a; }
|
||||
static __device__ __forceinline__ uint3& operator*= (uint3& a, unsigned int b) { a.x *= b; a.y *= b; a.z *= b; return a; }
|
||||
static __device__ __forceinline__ uint3& operator+= (uint3& a, unsigned int b) { a.x += b; a.y += b; a.z += b; return a; }
|
||||
static __device__ __forceinline__ uint3& operator-= (uint3& a, unsigned int b) { a.x -= b; a.y -= b; a.z -= b; return a; }
|
||||
static __device__ __forceinline__ uint3 operator* (const uint3& a, const uint3& b) { return make_uint3(a.x * b.x, a.y * b.y, a.z * b.z); }
|
||||
static __device__ __forceinline__ uint3 operator+ (const uint3& a, const uint3& b) { return make_uint3(a.x + b.x, a.y + b.y, a.z + b.z); }
|
||||
static __device__ __forceinline__ uint3 operator- (const uint3& a, const uint3& b) { return make_uint3(a.x - b.x, a.y - b.y, a.z - b.z); }
|
||||
static __device__ __forceinline__ uint3 operator* (const uint3& a, unsigned int b) { return make_uint3(a.x * b, a.y * b, a.z * b); }
|
||||
static __device__ __forceinline__ uint3 operator+ (const uint3& a, unsigned int b) { return make_uint3(a.x + b, a.y + b, a.z + b); }
|
||||
static __device__ __forceinline__ uint3 operator- (const uint3& a, unsigned int b) { return make_uint3(a.x - b, a.y - b, a.z - b); }
|
||||
static __device__ __forceinline__ uint3 operator* (unsigned int a, const uint3& b) { return make_uint3(a * b.x, a * b.y, a * b.z); }
|
||||
static __device__ __forceinline__ uint3 operator+ (unsigned int a, const uint3& b) { return make_uint3(a + b.x, a + b.y, a + b.z); }
|
||||
static __device__ __forceinline__ uint3 operator- (unsigned int a, const uint3& b) { return make_uint3(a - b.x, a - b.y, a - b.z); }
|
||||
static __device__ __forceinline__ uint4& operator*= (uint4& a, const uint4& b) { a.x *= b.x; a.y *= b.y; a.z *= b.z; a.w *= b.w; return a; }
|
||||
static __device__ __forceinline__ uint4& operator+= (uint4& a, const uint4& b) { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; return a; }
|
||||
static __device__ __forceinline__ uint4& operator-= (uint4& a, const uint4& b) { a.x -= b.x; a.y -= b.y; a.z -= b.z; a.w -= b.w; return a; }
|
||||
static __device__ __forceinline__ uint4& operator*= (uint4& a, unsigned int b) { a.x *= b; a.y *= b; a.z *= b; a.w *= b; return a; }
|
||||
static __device__ __forceinline__ uint4& operator+= (uint4& a, unsigned int b) { a.x += b; a.y += b; a.z += b; a.w += b; return a; }
|
||||
static __device__ __forceinline__ uint4& operator-= (uint4& a, unsigned int b) { a.x -= b; a.y -= b; a.z -= b; a.w -= b; return a; }
|
||||
static __device__ __forceinline__ uint4 operator* (const uint4& a, const uint4& b) { return make_uint4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w); }
|
||||
static __device__ __forceinline__ uint4 operator+ (const uint4& a, const uint4& b) { return make_uint4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
|
||||
static __device__ __forceinline__ uint4 operator- (const uint4& a, const uint4& b) { return make_uint4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w); }
|
||||
static __device__ __forceinline__ uint4 operator* (const uint4& a, unsigned int b) { return make_uint4(a.x * b, a.y * b, a.z * b, a.w * b); }
|
||||
static __device__ __forceinline__ uint4 operator+ (const uint4& a, unsigned int b) { return make_uint4(a.x + b, a.y + b, a.z + b, a.w + b); }
|
||||
static __device__ __forceinline__ uint4 operator- (const uint4& a, unsigned int b) { return make_uint4(a.x - b, a.y - b, a.z - b, a.w - b); }
|
||||
static __device__ __forceinline__ uint4 operator* (unsigned int a, const uint4& b) { return make_uint4(a * b.x, a * b.y, a * b.z, a * b.w); }
|
||||
static __device__ __forceinline__ uint4 operator+ (unsigned int a, const uint4& b) { return make_uint4(a + b.x, a + b.y, a + b.z, a + b.w); }
|
||||
static __device__ __forceinline__ uint4 operator- (unsigned int a, const uint4& b) { return make_uint4(a - b.x, a - b.y, a - b.z, a - b.w); }
|
||||
|
||||
template<class T> static __device__ __forceinline__ T zero_value(void);
|
||||
template<> __device__ __forceinline__ float zero_value<float> (void) { return 0.f; }
|
||||
template<> __device__ __forceinline__ float2 zero_value<float2>(void) { return make_float2(0.f, 0.f); }
|
||||
template<> __device__ __forceinline__ float4 zero_value<float4>(void) { return make_float4(0.f, 0.f, 0.f, 0.f); }
|
||||
static __device__ __forceinline__ float3 make_float3(const float2& a, float b) { return make_float3(a.x, a.y, b); }
|
||||
static __device__ __forceinline__ float4 make_float4(const float3& a, float b) { return make_float4(a.x, a.y, a.z, b); }
|
||||
static __device__ __forceinline__ float4 make_float4(const float2& a, const float2& b) { return make_float4(a.x, a.y, b.x, b.y); }
|
||||
static __device__ __forceinline__ int3 make_int3(const int2& a, int b) { return make_int3(a.x, a.y, b); }
|
||||
static __device__ __forceinline__ int4 make_int4(const int3& a, int b) { return make_int4(a.x, a.y, a.z, b); }
|
||||
static __device__ __forceinline__ int4 make_int4(const int2& a, const int2& b) { return make_int4(a.x, a.y, b.x, b.y); }
|
||||
static __device__ __forceinline__ uint3 make_uint3(const uint2& a, unsigned int b) { return make_uint3(a.x, a.y, b); }
|
||||
static __device__ __forceinline__ uint4 make_uint4(const uint3& a, unsigned int b) { return make_uint4(a.x, a.y, a.z, b); }
|
||||
static __device__ __forceinline__ uint4 make_uint4(const uint2& a, const uint2& b) { return make_uint4(a.x, a.y, b.x, b.y); }
|
||||
|
||||
template<class T> static __device__ __forceinline__ void swap(T& a, T& b) { T temp = a; a = b; b = temp; }
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Triangle ID <-> float32 conversion functions to support very large triangle IDs.
|
||||
//
|
||||
// Values up to and including 16777216 (also, negative values) are converted trivially and retain
|
||||
// compatibility with previous versions. Larger values are mapped to unique float32 that are not equal to
|
||||
// the ID. The largest value that converts to float32 and back without generating inf or nan is 889192447.
|
||||
|
||||
static __device__ __forceinline__ int float_to_triidx(float x) { if (x <= 16777216.f) return (int)x; return __float_as_int(x) - 0x4a800000; }
|
||||
static __device__ __forceinline__ float triidx_to_float(int x) { if (x <= 0x01000000) return (float)x; return __int_as_float(0x4a800000 + x); }
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Coalesced atomics. These are all done via macros.
|
||||
|
||||
#if __CUDA_ARCH__ >= 700 // Warp match instruction __match_any_sync() is only available on compute capability 7.x and higher
|
||||
|
||||
#define CA_TEMP _ca_temp
|
||||
#define CA_TEMP_PARAM float* CA_TEMP
|
||||
#define CA_DECLARE_TEMP(threads_per_block) \
|
||||
__shared__ float CA_TEMP[(threads_per_block)]
|
||||
|
||||
#define CA_SET_GROUP_MASK(group, thread_mask) \
|
||||
bool _ca_leader; \
|
||||
float* _ca_ptr; \
|
||||
do { \
|
||||
int tidx = threadIdx.x + blockDim.x * threadIdx.y; \
|
||||
int lane = tidx & 31; \
|
||||
int warp = tidx >> 5; \
|
||||
int tmask = __match_any_sync((thread_mask), (group)); \
|
||||
int leader = __ffs(tmask) - 1; \
|
||||
_ca_leader = (leader == lane); \
|
||||
_ca_ptr = &_ca_temp[((warp << 5) + leader)]; \
|
||||
} while(0)
|
||||
|
||||
#define CA_SET_GROUP(group) \
|
||||
CA_SET_GROUP_MASK((group), 0xffffffffu)
|
||||
|
||||
#define caAtomicAdd(ptr, value) \
|
||||
do { \
|
||||
if (_ca_leader) \
|
||||
*_ca_ptr = 0.f; \
|
||||
atomicAdd(_ca_ptr, (value)); \
|
||||
if (_ca_leader) \
|
||||
atomicAdd((ptr), *_ca_ptr); \
|
||||
} while(0)
|
||||
|
||||
#define caAtomicAdd3_xyw(ptr, x, y, w) \
|
||||
do { \
|
||||
caAtomicAdd((ptr), (x)); \
|
||||
caAtomicAdd((ptr)+1, (y)); \
|
||||
caAtomicAdd((ptr)+3, (w)); \
|
||||
} while(0)
|
||||
|
||||
#define caAtomicAddTexture(ptr, level, idx, value) \
|
||||
do { \
|
||||
CA_SET_GROUP((idx) ^ ((level) << 27)); \
|
||||
caAtomicAdd((ptr)+(idx), (value)); \
|
||||
} while(0)
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Disable atomic coalescing for compute capability lower than 7.x
|
||||
|
||||
#else // __CUDA_ARCH__ >= 700
|
||||
#define CA_TEMP _ca_temp
|
||||
#define CA_TEMP_PARAM float CA_TEMP
|
||||
#define CA_DECLARE_TEMP(threads_per_block) CA_TEMP_PARAM
|
||||
#define CA_SET_GROUP_MASK(group, thread_mask)
|
||||
#define CA_SET_GROUP(group)
|
||||
#define caAtomicAdd(ptr, value) atomicAdd((ptr), (value))
|
||||
#define caAtomicAdd3_xyw(ptr, x, y, w) \
|
||||
do { \
|
||||
atomicAdd((ptr), (x)); \
|
||||
atomicAdd((ptr)+1, (y)); \
|
||||
atomicAdd((ptr)+3, (w)); \
|
||||
} while(0)
|
||||
#define caAtomicAddTexture(ptr, level, idx, value) atomicAdd((ptr)+(idx), (value))
|
||||
#endif // __CUDA_ARCH__ >= 700
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
#endif // __CUDACC__
|
||||
@@ -0,0 +1,63 @@
|
||||
// Copyright (c) 2009-2022, 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
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// This is a slimmed-down and modernized version of the original
|
||||
// CudaRaster codebase that accompanied the HPG 2011 paper
|
||||
// "High-Performance Software Rasterization on GPUs" by Laine and Karras.
|
||||
// Modifications have been made to accommodate post-Volta execution model
|
||||
// with warp divergence. Support for shading, blending, quad rendering,
|
||||
// and supersampling have been removed as unnecessary for nvdiffrast.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
namespace CR
|
||||
{
|
||||
|
||||
class RasterImpl;
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Interface class to isolate user from implementation details.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
class CudaRaster
|
||||
{
|
||||
public:
|
||||
enum
|
||||
{
|
||||
RenderModeFlag_EnableBackfaceCulling = 1 << 0, // Enable backface culling.
|
||||
RenderModeFlag_EnableDepthPeeling = 1 << 1, // Enable depth peeling. Must have a peel buffer set.
|
||||
};
|
||||
|
||||
public:
|
||||
CudaRaster (void);
|
||||
~CudaRaster (void);
|
||||
|
||||
void setBufferSize (int width, int height, int numImages); // Width and height are internally rounded up to multiples of tile size (8x8) for buffer sizes.
|
||||
void setViewport (int width, int height, int offsetX, int offsetY); // Tiled rendering viewport setup.
|
||||
void setRenderModeFlags (unsigned int renderModeFlags); // Affects all subsequent calls to drawTriangles(). Defaults to zero.
|
||||
void deferredClear (unsigned int clearColor); // Clears color and depth buffers during next call to drawTriangles().
|
||||
void setVertexBuffer (void* vertices, int numVertices); // GPU pointer managed by caller. Vertex positions in clip space as float4 (x, y, z, w).
|
||||
void setIndexBuffer (void* indices, int numTriangles); // GPU pointer managed by caller. Triangle index+color quadruplets as uint4 (idx0, idx1, idx2, color).
|
||||
bool drawTriangles (const int* ranges, bool peel, cudaStream_t stream); // Ranges (offsets and counts) as #triangles entries, not as bytes. If NULL, draw all triangles. Returns false in case of internal overflow.
|
||||
void* getColorBuffer (void); // GPU pointer managed by CudaRaster.
|
||||
void* getDepthBuffer (void); // GPU pointer managed by CudaRaster.
|
||||
void swapDepthAndPeel (void); // Swap depth and peeling buffers.
|
||||
|
||||
private:
|
||||
CudaRaster (const CudaRaster&); // forbidden
|
||||
CudaRaster& operator= (const CudaRaster&); // forbidden
|
||||
|
||||
private:
|
||||
RasterImpl* m_impl; // Opaque pointer to implementation.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
} // namespace CR
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
// Copyright (c) 2009-2022, 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 "../../framework.h"
|
||||
#include "Buffer.hpp"
|
||||
|
||||
using namespace CR;
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// GPU buffer.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
Buffer::Buffer(void)
|
||||
: m_gpuPtr(NULL),
|
||||
m_bytes (0)
|
||||
{
|
||||
// empty
|
||||
}
|
||||
|
||||
Buffer::~Buffer(void)
|
||||
{
|
||||
if (m_gpuPtr)
|
||||
cudaFree(m_gpuPtr); // Don't throw an exception.
|
||||
}
|
||||
|
||||
void Buffer::reset(size_t bytes)
|
||||
{
|
||||
if (bytes == m_bytes)
|
||||
return;
|
||||
|
||||
if (m_gpuPtr)
|
||||
{
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFree(m_gpuPtr));
|
||||
m_gpuPtr = NULL;
|
||||
}
|
||||
|
||||
if (bytes > 0)
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMalloc(&m_gpuPtr, bytes));
|
||||
|
||||
m_bytes = bytes;
|
||||
}
|
||||
|
||||
void Buffer::grow(size_t bytes)
|
||||
{
|
||||
if (bytes > m_bytes)
|
||||
reset(bytes);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Host buffer with page-locked memory.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
HostBuffer::HostBuffer(void)
|
||||
: m_hostPtr(NULL),
|
||||
m_bytes (0)
|
||||
{
|
||||
// empty
|
||||
}
|
||||
|
||||
HostBuffer::~HostBuffer(void)
|
||||
{
|
||||
if (m_hostPtr)
|
||||
cudaFreeHost(m_hostPtr); // Don't throw an exception.
|
||||
}
|
||||
|
||||
void HostBuffer::reset(size_t bytes)
|
||||
{
|
||||
if (bytes == m_bytes)
|
||||
return;
|
||||
|
||||
if (m_hostPtr)
|
||||
{
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFreeHost(m_hostPtr));
|
||||
m_hostPtr = NULL;
|
||||
}
|
||||
|
||||
if (bytes > 0)
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMallocHost(&m_hostPtr, bytes));
|
||||
|
||||
m_bytes = bytes;
|
||||
}
|
||||
|
||||
void HostBuffer::grow(size_t bytes)
|
||||
{
|
||||
if (bytes > m_bytes)
|
||||
reset(bytes);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,55 @@
|
||||
// Copyright (c) 2009-2022, 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 "Defs.hpp"
|
||||
|
||||
namespace CR
|
||||
{
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
class Buffer
|
||||
{
|
||||
public:
|
||||
Buffer (void);
|
||||
~Buffer (void);
|
||||
|
||||
void reset (size_t bytes);
|
||||
void grow (size_t bytes);
|
||||
void* getPtr (size_t offset = 0) { return (void*)(((uintptr_t)m_gpuPtr) + offset); }
|
||||
size_t getSize (void) const { return m_bytes; }
|
||||
|
||||
void setPtr (void* ptr) { m_gpuPtr = ptr; }
|
||||
|
||||
private:
|
||||
void* m_gpuPtr;
|
||||
size_t m_bytes;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
class HostBuffer
|
||||
{
|
||||
public:
|
||||
HostBuffer (void);
|
||||
~HostBuffer (void);
|
||||
|
||||
void reset (size_t bytes);
|
||||
void grow (size_t bytes);
|
||||
void* getPtr (void) { return m_hostPtr; }
|
||||
size_t getSize (void) const { return m_bytes; }
|
||||
|
||||
void setPtr (void* ptr) { m_hostPtr = ptr; }
|
||||
|
||||
private:
|
||||
void* m_hostPtr;
|
||||
size_t m_bytes;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
// Copyright (c) 2009-2022, 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
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
#define CR_MAXVIEWPORT_LOG2 11 // ViewportSize / PixelSize.
|
||||
#define CR_SUBPIXEL_LOG2 4 // PixelSize / SubpixelSize.
|
||||
|
||||
#define CR_MAXBINS_LOG2 4 // ViewportSize / BinSize.
|
||||
#define CR_BIN_LOG2 4 // BinSize / TileSize.
|
||||
#define CR_TILE_LOG2 3 // TileSize / PixelSize.
|
||||
|
||||
#define CR_COVER8X8_LUT_SIZE 768 // 64-bit entries.
|
||||
#define CR_FLIPBIT_FLIP_Y 2
|
||||
#define CR_FLIPBIT_FLIP_X 3
|
||||
#define CR_FLIPBIT_SWAP_XY 4
|
||||
#define CR_FLIPBIT_COMPL 5
|
||||
|
||||
#define CR_BIN_STREAMS_LOG2 4
|
||||
#define CR_BIN_SEG_LOG2 9 // 32-bit entries.
|
||||
#define CR_TILE_SEG_LOG2 5 // 32-bit entries.
|
||||
|
||||
#define CR_MAXSUBTRIS_LOG2 24 // Triangle structs. Dictated by CoarseRaster.
|
||||
#define CR_COARSE_QUEUE_LOG2 10 // Triangles.
|
||||
|
||||
#define CR_SETUP_WARPS 2
|
||||
#define CR_SETUP_OPT_BLOCKS 8
|
||||
#define CR_BIN_WARPS 16
|
||||
#define CR_COARSE_WARPS 16 // Must be a power of two.
|
||||
#define CR_FINE_MAX_WARPS 20
|
||||
|
||||
#define CR_EMBED_IMAGE_PARAMS 32 // Number of per-image parameter structs embedded in kernel launch parameter block.
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
#define CR_MAXVIEWPORT_SIZE (1 << CR_MAXVIEWPORT_LOG2)
|
||||
#define CR_SUBPIXEL_SIZE (1 << CR_SUBPIXEL_LOG2)
|
||||
#define CR_SUBPIXEL_SQR (1 << (CR_SUBPIXEL_LOG2 * 2))
|
||||
|
||||
#define CR_MAXBINS_SIZE (1 << CR_MAXBINS_LOG2)
|
||||
#define CR_MAXBINS_SQR (1 << (CR_MAXBINS_LOG2 * 2))
|
||||
#define CR_BIN_SIZE (1 << CR_BIN_LOG2)
|
||||
#define CR_BIN_SQR (1 << (CR_BIN_LOG2 * 2))
|
||||
|
||||
#define CR_MAXTILES_LOG2 (CR_MAXBINS_LOG2 + CR_BIN_LOG2)
|
||||
#define CR_MAXTILES_SIZE (1 << CR_MAXTILES_LOG2)
|
||||
#define CR_MAXTILES_SQR (1 << (CR_MAXTILES_LOG2 * 2))
|
||||
#define CR_TILE_SIZE (1 << CR_TILE_LOG2)
|
||||
#define CR_TILE_SQR (1 << (CR_TILE_LOG2 * 2))
|
||||
|
||||
#define CR_BIN_STREAMS_SIZE (1 << CR_BIN_STREAMS_LOG2)
|
||||
#define CR_BIN_SEG_SIZE (1 << CR_BIN_SEG_LOG2)
|
||||
#define CR_TILE_SEG_SIZE (1 << CR_TILE_SEG_LOG2)
|
||||
|
||||
#define CR_MAXSUBTRIS_SIZE (1 << CR_MAXSUBTRIS_LOG2)
|
||||
#define CR_COARSE_QUEUE_SIZE (1 << CR_COARSE_QUEUE_LOG2)
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// When evaluating interpolated Z pixel centers, we may introduce an error
|
||||
// of (+-CR_LERP_ERROR) ULPs.
|
||||
|
||||
#define CR_LERP_ERROR(SAMPLES_LOG2) (2200u << (SAMPLES_LOG2))
|
||||
#define CR_DEPTH_MIN CR_LERP_ERROR(3)
|
||||
#define CR_DEPTH_MAX (CR_U32_MAX - CR_LERP_ERROR(3))
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,79 @@
|
||||
// Copyright (c) 2009-2022, 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 "Defs.hpp"
|
||||
#include "../CudaRaster.hpp"
|
||||
#include "RasterImpl.hpp"
|
||||
|
||||
using namespace CR;
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Stub interface implementation.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
CudaRaster::CudaRaster()
|
||||
{
|
||||
m_impl = new RasterImpl();
|
||||
}
|
||||
|
||||
CudaRaster::~CudaRaster()
|
||||
{
|
||||
delete m_impl;
|
||||
}
|
||||
|
||||
void CudaRaster::setBufferSize(int width, int height, int numImages)
|
||||
{
|
||||
m_impl->setBufferSize(Vec3i(width, height, numImages));
|
||||
}
|
||||
|
||||
void CudaRaster::setViewport(int width, int height, int offsetX, int offsetY)
|
||||
{
|
||||
m_impl->setViewport(Vec2i(width, height), Vec2i(offsetX, offsetY));
|
||||
}
|
||||
|
||||
void CudaRaster::setRenderModeFlags(U32 flags)
|
||||
{
|
||||
m_impl->setRenderModeFlags(flags);
|
||||
}
|
||||
|
||||
void CudaRaster::deferredClear(U32 clearColor)
|
||||
{
|
||||
m_impl->deferredClear(clearColor);
|
||||
}
|
||||
|
||||
void CudaRaster::setVertexBuffer(void* vertices, int numVertices)
|
||||
{
|
||||
m_impl->setVertexBuffer(vertices, numVertices);
|
||||
}
|
||||
|
||||
void CudaRaster::setIndexBuffer(void* indices, int numTriangles)
|
||||
{
|
||||
m_impl->setIndexBuffer(indices, numTriangles);
|
||||
}
|
||||
|
||||
bool CudaRaster::drawTriangles(const int* ranges, bool peel, cudaStream_t stream)
|
||||
{
|
||||
return m_impl->drawTriangles((const Vec2i*)ranges, peel, stream);
|
||||
}
|
||||
|
||||
void* CudaRaster::getColorBuffer(void)
|
||||
{
|
||||
return m_impl->getColorBuffer();
|
||||
}
|
||||
|
||||
void* CudaRaster::getDepthBuffer(void)
|
||||
{
|
||||
return m_impl->getDepthBuffer();
|
||||
}
|
||||
|
||||
void CudaRaster::swapDepthAndPeel(void)
|
||||
{
|
||||
m_impl->swapDepthAndPeel();
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,90 @@
|
||||
// Copyright (c) 2009-2022, 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 <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace CR
|
||||
{
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
#ifndef NULL
|
||||
# define NULL 0
|
||||
#endif
|
||||
|
||||
#ifdef __CUDACC__
|
||||
# define CR_CUDA 1
|
||||
#else
|
||||
# define CR_CUDA 0
|
||||
#endif
|
||||
|
||||
#if CR_CUDA
|
||||
# define CR_CUDA_FUNC __device__ __inline__
|
||||
# define CR_CUDA_CONST __constant__
|
||||
#else
|
||||
# define CR_CUDA_FUNC inline
|
||||
# define CR_CUDA_CONST static const
|
||||
#endif
|
||||
|
||||
#define CR_UNREF(X) ((void)(X))
|
||||
#define CR_ARRAY_SIZE(X) ((int)(sizeof(X) / sizeof((X)[0])))
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
typedef uint8_t U8;
|
||||
typedef uint16_t U16;
|
||||
typedef uint32_t U32;
|
||||
typedef uint64_t U64;
|
||||
typedef int8_t S8;
|
||||
typedef int16_t S16;
|
||||
typedef int32_t S32;
|
||||
typedef int64_t S64;
|
||||
typedef float F32;
|
||||
typedef double F64;
|
||||
typedef void (*FuncPtr)(void);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
#define CR_U32_MAX (0xFFFFFFFFu)
|
||||
#define CR_S32_MIN (~0x7FFFFFFF)
|
||||
#define CR_S32_MAX (0x7FFFFFFF)
|
||||
#define CR_U64_MAX ((U64)(S64)-1)
|
||||
#define CR_S64_MIN ((S64)-1 << 63)
|
||||
#define CR_S64_MAX (~((S64)-1 << 63))
|
||||
#define CR_F32_MIN (1.175494351e-38f)
|
||||
#define CR_F32_MAX (3.402823466e+38f)
|
||||
#define CR_F64_MIN (2.2250738585072014e-308)
|
||||
#define CR_F64_MAX (1.7976931348623158e+308)
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Misc types.
|
||||
|
||||
class Vec2i
|
||||
{
|
||||
public:
|
||||
Vec2i(int x_, int y_) : x(x_), y(y_) {}
|
||||
int x, y;
|
||||
};
|
||||
|
||||
class Vec3i
|
||||
{
|
||||
public:
|
||||
Vec3i(int x_, int y_, int z_) : x(x_), y(y_), z(z_) {}
|
||||
int x, y, z;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// CUDA utilities.
|
||||
|
||||
#if CR_CUDA
|
||||
# define globalThreadIdx (threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * (blockIdx.x + gridDim.x * blockIdx.y)))
|
||||
#endif
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
} // namespace CR
|
||||
@@ -0,0 +1,153 @@
|
||||
// Copyright (c) 2009-2022, 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 "Defs.hpp"
|
||||
#include "Constants.hpp"
|
||||
|
||||
namespace CR
|
||||
{
|
||||
//------------------------------------------------------------------------
|
||||
// Projected triangle.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
struct CRTriangleHeader
|
||||
{
|
||||
S16 v0x; // Subpixels relative to viewport center. Valid if triSubtris = 1.
|
||||
S16 v0y;
|
||||
S16 v1x;
|
||||
S16 v1y;
|
||||
S16 v2x;
|
||||
S16 v2y;
|
||||
|
||||
U32 misc; // triSubtris=1: (zmin:20, f01:4, f12:4, f20:4), triSubtris>=2: (subtriBase)
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
struct CRTriangleData
|
||||
{
|
||||
U32 zx; // zx * sampleX + zy * sampleY + zb = lerp(CR_DEPTH_MIN, CR_DEPTH_MAX, (clipZ / clipW + 1) / 2)
|
||||
U32 zy;
|
||||
U32 zb;
|
||||
U32 id; // Triangle id.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Device-side structures.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
struct CRAtomics
|
||||
{
|
||||
// Setup.
|
||||
S32 numSubtris; // = numTris
|
||||
|
||||
// Bin.
|
||||
S32 binCounter; // = 0
|
||||
S32 numBinSegs; // = 0
|
||||
|
||||
// Coarse.
|
||||
S32 coarseCounter; // = 0
|
||||
S32 numTileSegs; // = 0
|
||||
S32 numActiveTiles; // = 0
|
||||
|
||||
// Fine.
|
||||
S32 fineCounter; // = 0
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
struct CRImageParams
|
||||
{
|
||||
S32 triOffset; // First triangle index to draw.
|
||||
S32 triCount; // Number of triangles to draw.
|
||||
S32 binBatchSize; // Number of triangles per batch.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
struct CRParams
|
||||
{
|
||||
// Common.
|
||||
|
||||
CRAtomics* atomics; // Work counters. Per-image.
|
||||
S32 numImages; // Batch size.
|
||||
S32 totalCount; // In range mode, total number of triangles to render.
|
||||
S32 instanceMode; // 0 = range mode, 1 = instance mode.
|
||||
|
||||
S32 numVertices; // Number of vertices in input buffer, not counting multiples in instance mode.
|
||||
S32 numTriangles; // Number of triangles in input buffer.
|
||||
void* vertexBuffer; // numVertices * float4(x, y, z, w)
|
||||
void* indexBuffer; // numTriangles * int3(vi0, vi1, vi2)
|
||||
|
||||
S32 widthPixels; // Render buffer size in pixels. Must be multiple of tile size (8x8).
|
||||
S32 heightPixels;
|
||||
S32 widthPixelsVp; // Viewport size in pixels.
|
||||
S32 heightPixelsVp;
|
||||
S32 widthBins; // widthPixels / CR_BIN_SIZE
|
||||
S32 heightBins; // heightPixels / CR_BIN_SIZE
|
||||
S32 numBins; // widthBins * heightBins
|
||||
|
||||
F32 xs; // Vertex position adjustments for tiled rendering.
|
||||
F32 ys;
|
||||
F32 xo;
|
||||
F32 yo;
|
||||
|
||||
S32 widthTiles; // widthPixels / CR_TILE_SIZE
|
||||
S32 heightTiles; // heightPixels / CR_TILE_SIZE
|
||||
S32 numTiles; // widthTiles * heightTiles
|
||||
|
||||
U32 renderModeFlags;
|
||||
S32 deferredClear; // 1 = Clear framebuffer before rendering triangles.
|
||||
U32 clearColor;
|
||||
U32 clearDepth;
|
||||
|
||||
// These are uniform across batch.
|
||||
|
||||
S32 maxSubtris;
|
||||
S32 maxBinSegs;
|
||||
S32 maxTileSegs;
|
||||
|
||||
// Setup output / bin input.
|
||||
|
||||
void* triSubtris; // maxSubtris * U8
|
||||
void* triHeader; // maxSubtris * CRTriangleHeader
|
||||
void* triData; // maxSubtris * CRTriangleData
|
||||
|
||||
// Bin output / coarse input.
|
||||
|
||||
void* binSegData; // maxBinSegs * CR_BIN_SEG_SIZE * S32
|
||||
void* binSegNext; // maxBinSegs * S32
|
||||
void* binSegCount; // maxBinSegs * S32
|
||||
void* binFirstSeg; // CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * (S32 segIdx), -1 = none
|
||||
void* binTotal; // CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * (S32 numTris)
|
||||
|
||||
// Coarse output / fine input.
|
||||
|
||||
void* tileSegData; // maxTileSegs * CR_TILE_SEG_SIZE * S32
|
||||
void* tileSegNext; // maxTileSegs * S32
|
||||
void* tileSegCount; // maxTileSegs * S32
|
||||
void* activeTiles; // CR_MAXTILES_SQR * (S32 tileIdx)
|
||||
void* tileFirstSeg; // CR_MAXTILES_SQR * (S32 segIdx), -1 = none
|
||||
|
||||
// Surface buffers. Outer tile offset is baked into pointers.
|
||||
|
||||
void* colorBuffer; // sizePixels.x * sizePixels.y * numImages * U32
|
||||
void* depthBuffer; // sizePixels.x * sizePixels.y * numImages * U32
|
||||
void* peelBuffer; // sizePixels.x * sizePixels.y * numImages * U32, only if peeling enabled.
|
||||
S32 strideX; // horizontal size in pixels
|
||||
S32 strideY; // vertical stride in pixels
|
||||
|
||||
// Per-image parameters for first images are embedded here to avoid extra memcpy for small batches.
|
||||
|
||||
CRImageParams imageParamsFirst[CR_EMBED_IMAGE_PARAMS];
|
||||
const CRImageParams* imageParamsExtra; // After CR_EMBED_IMAGE_PARAMS.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
// Copyright (c) 2009-2022, 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 "../../framework.h"
|
||||
#include "PrivateDefs.hpp"
|
||||
#include "Constants.hpp"
|
||||
#include "RasterImpl.hpp"
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
using namespace CR;
|
||||
using std::min;
|
||||
using std::max;
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Kernel prototypes and variables.
|
||||
|
||||
void triangleSetupKernel (const CRParams p);
|
||||
void binRasterKernel (const CRParams p);
|
||||
void coarseRasterKernel (const CRParams p);
|
||||
void fineRasterKernel (const CRParams p);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
RasterImpl::RasterImpl(void)
|
||||
: m_renderModeFlags (0),
|
||||
m_deferredClear (false),
|
||||
m_clearColor (0),
|
||||
m_vertexPtr (NULL),
|
||||
m_indexPtr (NULL),
|
||||
m_numVertices (0),
|
||||
m_numTriangles (0),
|
||||
m_bufferSizesReported (0),
|
||||
|
||||
m_numImages (0),
|
||||
m_bufferSizePixels (0, 0),
|
||||
m_bufferSizeVp (0, 0),
|
||||
m_sizePixels (0, 0),
|
||||
m_sizeVp (0, 0),
|
||||
m_offsetPixels (0, 0),
|
||||
m_sizeBins (0, 0),
|
||||
m_numBins (0),
|
||||
m_sizeTiles (0, 0),
|
||||
m_numTiles (0),
|
||||
|
||||
m_numSMs (1),
|
||||
m_numCoarseBlocksPerSM (1),
|
||||
m_numFineBlocksPerSM (1),
|
||||
m_numFineWarpsPerBlock (1),
|
||||
|
||||
m_maxSubtris (1),
|
||||
m_maxBinSegs (1),
|
||||
m_maxTileSegs (1)
|
||||
{
|
||||
// Query relevant device attributes.
|
||||
|
||||
int currentDevice = 0;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaGetDevice(¤tDevice));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaDeviceGetAttribute(&m_numSMs, cudaDevAttrMultiProcessorCount, currentDevice));
|
||||
cudaFuncAttributes attr;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFuncGetAttributes(&attr, (void*)fineRasterKernel));
|
||||
m_numFineWarpsPerBlock = min(attr.maxThreadsPerBlock / 32, CR_FINE_MAX_WARPS);
|
||||
NVDR_CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&m_numCoarseBlocksPerSM, (void*)coarseRasterKernel, 32 * CR_COARSE_WARPS, 0));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&m_numFineBlocksPerSM, (void*)fineRasterKernel, 32 * m_numFineWarpsPerBlock, 0));
|
||||
|
||||
// Setup functions.
|
||||
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFuncSetCacheConfig((void*)triangleSetupKernel, cudaFuncCachePreferShared));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFuncSetCacheConfig((void*)binRasterKernel, cudaFuncCachePreferShared));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFuncSetCacheConfig((void*)coarseRasterKernel, cudaFuncCachePreferShared));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaFuncSetCacheConfig((void*)fineRasterKernel, cudaFuncCachePreferShared));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
RasterImpl::~RasterImpl(void)
|
||||
{
|
||||
// Empty.
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
void RasterImpl::setBufferSize(Vec3i size)
|
||||
{
|
||||
// Internal buffer width and height must be divisible by tile size.
|
||||
int w = (size.x + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
int h = (size.y + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
|
||||
m_bufferSizePixels = Vec2i(w, h);
|
||||
m_bufferSizeVp = Vec2i(size.x, size.y);
|
||||
m_numImages = size.z;
|
||||
|
||||
m_colorBuffer.reset(w * h * size.z * sizeof(U32));
|
||||
m_depthBuffer.reset(w * h * size.z * sizeof(U32));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
void RasterImpl::setViewport(Vec2i size, Vec2i offset)
|
||||
{
|
||||
// Offset must be divisible by tile size.
|
||||
NVDR_CHECK((offset.x & (CR_TILE_SIZE - 1)) == 0 && (offset.y & (CR_TILE_SIZE - 1)) == 0, "invalid viewport offset");
|
||||
|
||||
// Round internal viewport size to multiples of tile size.
|
||||
int w = (size.x + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
int h = (size.y + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
|
||||
m_sizePixels = Vec2i(w, h);
|
||||
m_offsetPixels = offset;
|
||||
m_sizeVp = Vec2i(size.x, size.y);
|
||||
m_sizeTiles.x = m_sizePixels.x >> CR_TILE_LOG2;
|
||||
m_sizeTiles.y = m_sizePixels.y >> CR_TILE_LOG2;
|
||||
m_numTiles = m_sizeTiles.x * m_sizeTiles.y;
|
||||
m_sizeBins.x = (m_sizeTiles.x + CR_BIN_SIZE - 1) >> CR_BIN_LOG2;
|
||||
m_sizeBins.y = (m_sizeTiles.y + CR_BIN_SIZE - 1) >> CR_BIN_LOG2;
|
||||
m_numBins = m_sizeBins.x * m_sizeBins.y;
|
||||
}
|
||||
|
||||
void RasterImpl::swapDepthAndPeel(void)
|
||||
{
|
||||
m_peelBuffer.reset(m_depthBuffer.getSize()); // Ensure equal size and valid pointer.
|
||||
|
||||
void* tmp = m_depthBuffer.getPtr();
|
||||
m_depthBuffer.setPtr(m_peelBuffer.getPtr());
|
||||
m_peelBuffer.setPtr(tmp);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
bool RasterImpl::drawTriangles(const Vec2i* ranges, bool peel, cudaStream_t stream)
|
||||
{
|
||||
bool instanceMode = (!ranges);
|
||||
|
||||
int maxSubtrisSlack = 4096; // x 81B = 324KB
|
||||
int maxBinSegsSlack = 256; // x 2137B = 534KB
|
||||
int maxTileSegsSlack = 4096; // x 136B = 544KB
|
||||
|
||||
// Resize atomics as needed.
|
||||
m_crAtomics .grow(m_numImages * sizeof(CRAtomics));
|
||||
m_crAtomicsHost.grow(m_numImages * sizeof(CRAtomics));
|
||||
|
||||
// Size of these buffers doesn't depend on input.
|
||||
m_binFirstSeg .grow(m_numImages * CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * sizeof(S32));
|
||||
m_binTotal .grow(m_numImages * CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * sizeof(S32));
|
||||
m_activeTiles .grow(m_numImages * CR_MAXTILES_SQR * sizeof(S32));
|
||||
m_tileFirstSeg .grow(m_numImages * CR_MAXTILES_SQR * sizeof(S32));
|
||||
|
||||
// Construct per-image parameters and determine worst-case buffer sizes.
|
||||
m_crImageParamsHost.grow(m_numImages * sizeof(CRImageParams));
|
||||
CRImageParams* imageParams = (CRImageParams*)m_crImageParamsHost.getPtr();
|
||||
for (int i=0; i < m_numImages; i++)
|
||||
{
|
||||
CRImageParams& ip = imageParams[i];
|
||||
|
||||
int roundSize = CR_BIN_WARPS * 32;
|
||||
int minBatches = CR_BIN_STREAMS_SIZE * 2;
|
||||
int maxRounds = 32;
|
||||
|
||||
ip.triOffset = instanceMode ? 0 : ranges[i].x;
|
||||
ip.triCount = instanceMode ? m_numTriangles : ranges[i].y;
|
||||
ip.binBatchSize = min(max(ip.triCount / (roundSize * minBatches), 1), maxRounds) * roundSize;
|
||||
|
||||
m_maxSubtris = max(m_maxSubtris, min(ip.triCount + maxSubtrisSlack, CR_MAXSUBTRIS_SIZE));
|
||||
m_maxBinSegs = max(m_maxBinSegs, max(m_numBins * CR_BIN_STREAMS_SIZE, (ip.triCount - 1) / CR_BIN_SEG_SIZE + 1) + maxBinSegsSlack);
|
||||
m_maxTileSegs = max(m_maxTileSegs, max(m_numTiles, (ip.triCount - 1) / CR_TILE_SEG_SIZE + 1) + maxTileSegsSlack);
|
||||
}
|
||||
|
||||
// Retry until successful.
|
||||
|
||||
for (;;)
|
||||
{
|
||||
// Allocate buffers.
|
||||
m_triSubtris.reset(m_numImages * m_maxSubtris * sizeof(U8));
|
||||
m_triHeader .reset(m_numImages * m_maxSubtris * sizeof(CRTriangleHeader));
|
||||
m_triData .reset(m_numImages * m_maxSubtris * sizeof(CRTriangleData));
|
||||
|
||||
m_binSegData .reset(m_numImages * m_maxBinSegs * CR_BIN_SEG_SIZE * sizeof(S32));
|
||||
m_binSegNext .reset(m_numImages * m_maxBinSegs * sizeof(S32));
|
||||
m_binSegCount.reset(m_numImages * m_maxBinSegs * sizeof(S32));
|
||||
|
||||
m_tileSegData .reset(m_numImages * m_maxTileSegs * CR_TILE_SEG_SIZE * sizeof(S32));
|
||||
m_tileSegNext .reset(m_numImages * m_maxTileSegs * sizeof(S32));
|
||||
m_tileSegCount.reset(m_numImages * m_maxTileSegs * sizeof(S32));
|
||||
|
||||
// Report if buffers grow from last time.
|
||||
size_t sizesTotal = getTotalBufferSizes();
|
||||
if (sizesTotal > m_bufferSizesReported)
|
||||
{
|
||||
size_t sizesMB = ((sizesTotal - 1) >> 20) + 1; // Round up.
|
||||
sizesMB = ((sizesMB + 9) / 10) * 10; // 10MB granularity enough in this day and age.
|
||||
LOG(INFO) << "Internal buffers grown to " << sizesMB << " MB";
|
||||
m_bufferSizesReported = sizesMB << 20;
|
||||
}
|
||||
|
||||
// Launch stages. Blocks until everything is done.
|
||||
launchStages(instanceMode, peel, stream);
|
||||
|
||||
// Peeling iteration cannot fail, so no point checking things further.
|
||||
if (peel)
|
||||
break;
|
||||
|
||||
// Atomics after coarse stage are now available.
|
||||
CRAtomics* atomics = (CRAtomics*)m_crAtomicsHost.getPtr();
|
||||
|
||||
// Success?
|
||||
bool failed = false;
|
||||
for (int i=0; i < m_numImages; i++)
|
||||
{
|
||||
const CRAtomics& a = atomics[i];
|
||||
failed = failed || (a.numSubtris > m_maxSubtris) || (a.numBinSegs > m_maxBinSegs) || (a.numTileSegs > m_maxTileSegs);
|
||||
}
|
||||
if (!failed)
|
||||
break; // Success!
|
||||
|
||||
// If we were already at maximum capacity, no can do.
|
||||
if (m_maxSubtris == CR_MAXSUBTRIS_SIZE)
|
||||
return false;
|
||||
|
||||
// Enlarge buffers and try again.
|
||||
for (int i=0; i < m_numImages; i++)
|
||||
{
|
||||
const CRAtomics& a = atomics[i];
|
||||
m_maxSubtris = max(m_maxSubtris, min(a.numSubtris + maxSubtrisSlack, CR_MAXSUBTRIS_SIZE));
|
||||
m_maxBinSegs = max(m_maxBinSegs, a.numBinSegs + maxBinSegsSlack);
|
||||
m_maxTileSegs = max(m_maxTileSegs, a.numTileSegs + maxTileSegsSlack);
|
||||
}
|
||||
}
|
||||
|
||||
m_deferredClear = false;
|
||||
return true; // Success.
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
size_t RasterImpl::getTotalBufferSizes(void) const
|
||||
{
|
||||
return
|
||||
m_colorBuffer.getSize() + m_depthBuffer.getSize() + // Don't include atomics and image params.
|
||||
m_triSubtris.getSize() + m_triHeader.getSize() + m_triData.getSize() +
|
||||
m_binFirstSeg.getSize() + m_binTotal.getSize() + m_binSegData.getSize() + m_binSegNext.getSize() + m_binSegCount.getSize() +
|
||||
m_activeTiles.getSize() + m_tileFirstSeg.getSize() + m_tileSegData.getSize() + m_tileSegNext.getSize() + m_tileSegCount.getSize();
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
void RasterImpl::launchStages(bool instanceMode, bool peel, cudaStream_t stream)
|
||||
{
|
||||
CRImageParams* imageParams = (CRImageParams*)m_crImageParamsHost.getPtr();
|
||||
|
||||
// Unless peeling, initialize atomics to mostly zero.
|
||||
CRAtomics* atomics = (CRAtomics*)m_crAtomicsHost.getPtr();
|
||||
if (!peel)
|
||||
{
|
||||
memset(atomics, 0, m_numImages * sizeof(CRAtomics));
|
||||
for (int i=0; i < m_numImages; i++)
|
||||
atomics[i].numSubtris = imageParams[i].triCount;
|
||||
}
|
||||
|
||||
// Copy to device. If peeling, this is the state after coarse raster launch on first iteration.
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMemcpyAsync(m_crAtomics.getPtr(), atomics, m_numImages * sizeof(CRAtomics), cudaMemcpyHostToDevice, stream));
|
||||
|
||||
// Copy per-image parameters if there are more than fits in launch parameter block and we haven't done it already.
|
||||
if (!peel && m_numImages > CR_EMBED_IMAGE_PARAMS)
|
||||
{
|
||||
int numImageParamsExtra = m_numImages - CR_EMBED_IMAGE_PARAMS;
|
||||
m_crImageParamsExtra.grow(numImageParamsExtra * sizeof(CRImageParams));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMemcpyAsync(m_crImageParamsExtra.getPtr(), imageParams + CR_EMBED_IMAGE_PARAMS, numImageParamsExtra * sizeof(CRImageParams), cudaMemcpyHostToDevice, stream));
|
||||
}
|
||||
|
||||
// Set global parameters.
|
||||
CRParams p;
|
||||
{
|
||||
p.atomics = (CRAtomics*)m_crAtomics.getPtr();
|
||||
p.numImages = m_numImages;
|
||||
p.totalCount = 0; // Only relevant in range mode.
|
||||
p.instanceMode = instanceMode ? 1 : 0;
|
||||
|
||||
p.numVertices = m_numVertices;
|
||||
p.numTriangles = m_numTriangles;
|
||||
p.vertexBuffer = m_vertexPtr;
|
||||
p.indexBuffer = m_indexPtr;
|
||||
|
||||
p.widthPixels = m_sizePixels.x;
|
||||
p.heightPixels = m_sizePixels.y;
|
||||
p.widthPixelsVp = m_sizeVp.x;
|
||||
p.heightPixelsVp = m_sizeVp.y;
|
||||
p.widthBins = m_sizeBins.x;
|
||||
p.heightBins = m_sizeBins.y;
|
||||
p.numBins = m_numBins;
|
||||
|
||||
p.xs = (float)m_bufferSizeVp.x / (float)m_sizeVp.x;
|
||||
p.ys = (float)m_bufferSizeVp.y / (float)m_sizeVp.y;
|
||||
p.xo = (float)(m_bufferSizeVp.x - m_sizeVp.x - 2 * m_offsetPixels.x) / (float)m_sizeVp.x;
|
||||
p.yo = (float)(m_bufferSizeVp.y - m_sizeVp.y - 2 * m_offsetPixels.y) / (float)m_sizeVp.y;
|
||||
|
||||
p.widthTiles = m_sizeTiles.x;
|
||||
p.heightTiles = m_sizeTiles.y;
|
||||
p.numTiles = m_numTiles;
|
||||
|
||||
p.renderModeFlags = m_renderModeFlags;
|
||||
p.deferredClear = m_deferredClear ? 1 : 0;
|
||||
p.clearColor = m_clearColor;
|
||||
p.clearDepth = CR_DEPTH_MAX;
|
||||
|
||||
p.maxSubtris = m_maxSubtris;
|
||||
p.maxBinSegs = m_maxBinSegs;
|
||||
p.maxTileSegs = m_maxTileSegs;
|
||||
|
||||
p.triSubtris = m_triSubtris.getPtr();
|
||||
p.triHeader = m_triHeader.getPtr();
|
||||
p.triData = m_triData.getPtr();
|
||||
p.binSegData = m_binSegData.getPtr();
|
||||
p.binSegNext = m_binSegNext.getPtr();
|
||||
p.binSegCount = m_binSegCount.getPtr();
|
||||
p.binFirstSeg = m_binFirstSeg.getPtr();
|
||||
p.binTotal = m_binTotal.getPtr();
|
||||
p.tileSegData = m_tileSegData.getPtr();
|
||||
p.tileSegNext = m_tileSegNext.getPtr();
|
||||
p.tileSegCount = m_tileSegCount.getPtr();
|
||||
p.activeTiles = m_activeTiles.getPtr();
|
||||
p.tileFirstSeg = m_tileFirstSeg.getPtr();
|
||||
|
||||
size_t byteOffset = ((size_t)m_offsetPixels.x + (size_t)m_offsetPixels.y * (size_t)m_bufferSizePixels.x) * sizeof(U32);
|
||||
p.colorBuffer = m_colorBuffer.getPtr(byteOffset);
|
||||
p.depthBuffer = m_depthBuffer.getPtr(byteOffset);
|
||||
p.peelBuffer = (m_renderModeFlags & CudaRaster::RenderModeFlag_EnableDepthPeeling) ? m_peelBuffer.getPtr(byteOffset) : 0;
|
||||
p.strideX = m_bufferSizePixels.x;
|
||||
p.strideY = m_bufferSizePixels.y;
|
||||
|
||||
memcpy(&p.imageParamsFirst, imageParams, min(m_numImages, CR_EMBED_IMAGE_PARAMS) * sizeof(CRImageParams));
|
||||
p.imageParamsExtra = (CRImageParams*)m_crImageParamsExtra.getPtr();
|
||||
}
|
||||
|
||||
// Setup block sizes.
|
||||
|
||||
dim3 brBlock(32, CR_BIN_WARPS);
|
||||
dim3 crBlock(32, CR_COARSE_WARPS);
|
||||
dim3 frBlock(32, m_numFineWarpsPerBlock);
|
||||
void* args[] = {&p};
|
||||
|
||||
// Launch stages from setup to coarse and copy atomics to host only if this is not a single-tile peeling iteration.
|
||||
if (!peel)
|
||||
{
|
||||
if (instanceMode)
|
||||
{
|
||||
int setupBlocks = (m_numTriangles - 1) / (32 * CR_SETUP_WARPS) + 1;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)triangleSetupKernel, dim3(setupBlocks, 1, m_numImages), dim3(32, CR_SETUP_WARPS), args, 0, stream));
|
||||
}
|
||||
else
|
||||
{
|
||||
for (int i=0; i < m_numImages; i++)
|
||||
p.totalCount += imageParams[i].triCount;
|
||||
int setupBlocks = (p.totalCount - 1) / (32 * CR_SETUP_WARPS) + 1;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)triangleSetupKernel, dim3(setupBlocks, 1, 1), dim3(32, CR_SETUP_WARPS), args, 0, stream));
|
||||
}
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)binRasterKernel, dim3(CR_BIN_STREAMS_SIZE, 1, m_numImages), brBlock, args, 0, stream));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)coarseRasterKernel, dim3(m_numSMs * m_numCoarseBlocksPerSM, 1, m_numImages), crBlock, args, 0, stream));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMemcpyAsync(m_crAtomicsHost.getPtr(), m_crAtomics.getPtr(), sizeof(CRAtomics) * m_numImages, cudaMemcpyDeviceToHost, stream));
|
||||
}
|
||||
|
||||
// Fine rasterizer is launched always.
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)fineRasterKernel, dim3(m_numSMs * m_numFineBlocksPerSM, 1, m_numImages), frBlock, args, 0, stream));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,102 @@
|
||||
// Copyright (c) 2009-2022, 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 "PrivateDefs.hpp"
|
||||
#include "Buffer.hpp"
|
||||
#include "../CudaRaster.hpp"
|
||||
|
||||
namespace CR
|
||||
{
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
class RasterImpl
|
||||
{
|
||||
public:
|
||||
RasterImpl (void);
|
||||
~RasterImpl (void);
|
||||
|
||||
void setBufferSize (Vec3i size);
|
||||
void setViewport (Vec2i size, Vec2i offset);
|
||||
void setRenderModeFlags (U32 flags) { m_renderModeFlags = flags; }
|
||||
void deferredClear (U32 color) { m_deferredClear = true; m_clearColor = color; }
|
||||
void setVertexBuffer (void* ptr, int numVertices) { m_vertexPtr = ptr; m_numVertices = numVertices; } // GPU pointer.
|
||||
void setIndexBuffer (void* ptr, int numTriangles) { m_indexPtr = ptr; m_numTriangles = numTriangles; } // GPU pointer.
|
||||
bool drawTriangles (const Vec2i* ranges, bool peel, cudaStream_t stream);
|
||||
void* getColorBuffer (void) { return m_colorBuffer.getPtr(); } // GPU pointer.
|
||||
void* getDepthBuffer (void) { return m_depthBuffer.getPtr(); } // GPU pointer.
|
||||
void swapDepthAndPeel (void);
|
||||
size_t getTotalBufferSizes (void) const;
|
||||
|
||||
private:
|
||||
void launchStages (bool instanceMode, bool peel, cudaStream_t stream);
|
||||
|
||||
// State.
|
||||
|
||||
unsigned int m_renderModeFlags;
|
||||
bool m_deferredClear;
|
||||
unsigned int m_clearColor;
|
||||
void* m_vertexPtr;
|
||||
void* m_indexPtr;
|
||||
int m_numVertices; // Input buffer size.
|
||||
int m_numTriangles; // Input buffer size.
|
||||
size_t m_bufferSizesReported; // Previously reported buffer sizes.
|
||||
|
||||
// Surfaces.
|
||||
|
||||
Buffer m_colorBuffer;
|
||||
Buffer m_depthBuffer;
|
||||
Buffer m_peelBuffer;
|
||||
int m_numImages;
|
||||
Vec2i m_bufferSizePixels; // Internal buffer size.
|
||||
Vec2i m_bufferSizeVp; // Total viewport size.
|
||||
Vec2i m_sizePixels; // Internal size at which all computation is done, buffers reserved, etc.
|
||||
Vec2i m_sizeVp; // Size to which output will be cropped outside, determines viewport size.
|
||||
Vec2i m_offsetPixels; // Viewport offset for tiled rendering.
|
||||
Vec2i m_sizeBins;
|
||||
S32 m_numBins;
|
||||
Vec2i m_sizeTiles;
|
||||
S32 m_numTiles;
|
||||
|
||||
// Launch sizes etc.
|
||||
|
||||
S32 m_numSMs;
|
||||
S32 m_numCoarseBlocksPerSM;
|
||||
S32 m_numFineBlocksPerSM;
|
||||
S32 m_numFineWarpsPerBlock;
|
||||
|
||||
// Global intermediate buffers. Individual images have offsets to these.
|
||||
|
||||
Buffer m_crAtomics;
|
||||
HostBuffer m_crAtomicsHost;
|
||||
HostBuffer m_crImageParamsHost;
|
||||
Buffer m_crImageParamsExtra;
|
||||
Buffer m_triSubtris;
|
||||
Buffer m_triHeader;
|
||||
Buffer m_triData;
|
||||
Buffer m_binFirstSeg;
|
||||
Buffer m_binTotal;
|
||||
Buffer m_binSegData;
|
||||
Buffer m_binSegNext;
|
||||
Buffer m_binSegCount;
|
||||
Buffer m_activeTiles;
|
||||
Buffer m_tileFirstSeg;
|
||||
Buffer m_tileSegData;
|
||||
Buffer m_tileSegNext;
|
||||
Buffer m_tileSegCount;
|
||||
|
||||
// Actual buffer sizes.
|
||||
|
||||
S32 m_maxSubtris;
|
||||
S32 m_maxBinSegs;
|
||||
S32 m_maxTileSegs;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
} // namespace CR
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) 2009-2022, 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 "../CudaRaster.hpp"
|
||||
#include "PrivateDefs.hpp"
|
||||
#include "Constants.hpp"
|
||||
#include "Util.inl"
|
||||
|
||||
namespace CR
|
||||
{
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Stage implementations.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
#include "TriangleSetup.inl"
|
||||
#include "BinRaster.inl"
|
||||
#include "CoarseRaster.inl"
|
||||
#include "FineRaster.inl"
|
||||
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Stage entry points.
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
__global__ void __launch_bounds__(CR_SETUP_WARPS * 32, CR_SETUP_OPT_BLOCKS) triangleSetupKernel (const CR::CRParams p) { CR::triangleSetupImpl(p); }
|
||||
__global__ void __launch_bounds__(CR_BIN_WARPS * 32, 1) binRasterKernel (const CR::CRParams p) { CR::binRasterImpl(p); }
|
||||
__global__ void __launch_bounds__(CR_COARSE_WARPS * 32, 1) coarseRasterKernel (const CR::CRParams p) { CR::coarseRasterImpl(p); }
|
||||
__global__ void __launch_bounds__(CR_FINE_MAX_WARPS * 32, 1) fineRasterKernel (const CR::CRParams p) { CR::fineRasterImpl(p); }
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#pragma once
|
||||
|
||||
// Framework-specific macros to enable code sharing.
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// PyTorch.
|
||||
|
||||
#ifdef NVDR_TORCH
|
||||
#ifndef __CUDACC__
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/CUDAUtils.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <pybind11/numpy.h>
|
||||
#endif
|
||||
#define NVDR_CHECK(COND, ERR) do { TORCH_CHECK(COND, ERR) } while(0)
|
||||
#define NVDR_CHECK_CUDA_ERROR(CUDA_CALL) do { cudaError_t err = CUDA_CALL; TORCH_CHECK(!err, "Cuda error: ", cudaGetLastError(), "[", #CUDA_CALL, ";]"); } while(0)
|
||||
#endif
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+276
@@ -0,0 +1,276 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "common.h"
|
||||
#include "interpolate.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Forward kernel.
|
||||
|
||||
template <bool ENABLE_DA>
|
||||
static __forceinline__ __device__ void InterpolateFwdKernelTemplate(const InterpolateKernelParams p)
|
||||
{
|
||||
// Calculate pixel position.
|
||||
int px = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int py = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int pz = blockIdx.z;
|
||||
if (px >= p.width || py >= p.height || pz >= p.depth)
|
||||
return;
|
||||
|
||||
// Pixel index.
|
||||
int pidx = px + p.width * (py + p.height * pz);
|
||||
|
||||
// Output ptrs.
|
||||
float* out = p.out + pidx * p.numAttr;
|
||||
float2* outDA = ENABLE_DA ? (((float2*)p.outDA) + pidx * p.numDiffAttr) : 0;
|
||||
|
||||
// Fetch rasterizer output.
|
||||
float4 r = ((float4*)p.rast)[pidx];
|
||||
int triIdx = float_to_triidx(r.w) - 1;
|
||||
bool triValid = (triIdx >= 0 && triIdx < p.numTriangles);
|
||||
|
||||
// If no geometry in entire warp, zero the output and exit.
|
||||
// Otherwise force barys to zero and output with live threads.
|
||||
if (__all_sync(0xffffffffu, !triValid))
|
||||
{
|
||||
for (int i=0; i < p.numAttr; i++)
|
||||
out[i] = 0.f;
|
||||
if (ENABLE_DA)
|
||||
for (int i=0; i < p.numDiffAttr; i++)
|
||||
outDA[i] = make_float2(0.f, 0.f);
|
||||
return;
|
||||
}
|
||||
|
||||
// Fetch vertex indices.
|
||||
int vi0 = triValid ? p.tri[triIdx * 3 + 0] : 0;
|
||||
int vi1 = triValid ? p.tri[triIdx * 3 + 1] : 0;
|
||||
int vi2 = triValid ? p.tri[triIdx * 3 + 2] : 0;
|
||||
|
||||
// Bail out if corrupt indices.
|
||||
if (vi0 < 0 || vi0 >= p.numVertices ||
|
||||
vi1 < 0 || vi1 >= p.numVertices ||
|
||||
vi2 < 0 || vi2 >= p.numVertices)
|
||||
return;
|
||||
|
||||
// In instance mode, adjust vertex indices by minibatch index unless broadcasting.
|
||||
if (p.instance_mode && !p.attrBC)
|
||||
{
|
||||
vi0 += pz * p.numVertices;
|
||||
vi1 += pz * p.numVertices;
|
||||
vi2 += pz * p.numVertices;
|
||||
}
|
||||
|
||||
// Pointers to attributes.
|
||||
const float* a0 = p.attr + vi0 * p.numAttr;
|
||||
const float* a1 = p.attr + vi1 * p.numAttr;
|
||||
const float* a2 = p.attr + vi2 * p.numAttr;
|
||||
|
||||
// Barys. If no triangle, force all to zero -> output is zero.
|
||||
float b0 = triValid ? r.x : 0.f;
|
||||
float b1 = triValid ? r.y : 0.f;
|
||||
float b2 = triValid ? (1.f - r.x - r.y) : 0.f;
|
||||
|
||||
// Interpolate and write attributes.
|
||||
for (int i=0; i < p.numAttr; i++)
|
||||
out[i] = b0*a0[i] + b1*a1[i] + b2*a2[i];
|
||||
|
||||
// No diff attrs? Exit.
|
||||
if (!ENABLE_DA)
|
||||
return;
|
||||
|
||||
// Read bary pixel differentials if we have a triangle.
|
||||
float4 db = make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
if (triValid)
|
||||
db = ((float4*)p.rastDB)[pidx];
|
||||
|
||||
// Unpack a bit.
|
||||
float dudx = db.x;
|
||||
float dudy = db.y;
|
||||
float dvdx = db.z;
|
||||
float dvdy = db.w;
|
||||
|
||||
// Calculate the pixel differentials of chosen attributes.
|
||||
for (int i=0; i < p.numDiffAttr; i++)
|
||||
{
|
||||
// Input attribute index.
|
||||
int j = p.diff_attrs_all ? i : p.diffAttrs[i];
|
||||
if (j < 0)
|
||||
j += p.numAttr; // Python-style negative indices.
|
||||
|
||||
// Zero output if invalid index.
|
||||
float dsdx = 0.f;
|
||||
float dsdy = 0.f;
|
||||
if (j >= 0 && j < p.numAttr)
|
||||
{
|
||||
float s0 = a0[j];
|
||||
float s1 = a1[j];
|
||||
float s2 = a2[j];
|
||||
float dsdu = s0 - s2;
|
||||
float dsdv = s1 - s2;
|
||||
dsdx = dudx*dsdu + dvdx*dsdv;
|
||||
dsdy = dudy*dsdu + dvdy*dsdv;
|
||||
}
|
||||
|
||||
// Write.
|
||||
outDA[i] = make_float2(dsdx, dsdy);
|
||||
}
|
||||
}
|
||||
|
||||
// Template specializations.
|
||||
__global__ void InterpolateFwdKernel (const InterpolateKernelParams p) { InterpolateFwdKernelTemplate<false>(p); }
|
||||
__global__ void InterpolateFwdKernelDa(const InterpolateKernelParams p) { InterpolateFwdKernelTemplate<true>(p); }
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient kernel.
|
||||
|
||||
template <bool ENABLE_DA>
|
||||
static __forceinline__ __device__ void InterpolateGradKernelTemplate(const InterpolateKernelParams p)
|
||||
{
|
||||
// Temporary space for coalesced atomics.
|
||||
CA_DECLARE_TEMP(IP_GRAD_MAX_KERNEL_BLOCK_WIDTH * IP_GRAD_MAX_KERNEL_BLOCK_HEIGHT);
|
||||
|
||||
// Calculate pixel position.
|
||||
int px = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int py = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int pz = blockIdx.z;
|
||||
if (px >= p.width || py >= p.height || pz >= p.depth)
|
||||
return;
|
||||
|
||||
// Pixel index.
|
||||
int pidx = px + p.width * (py + p.height * pz);
|
||||
|
||||
// Fetch triangle ID. If none, output zero bary/db gradients and exit.
|
||||
float4 r = ((float4*)p.rast)[pidx];
|
||||
int triIdx = float_to_triidx(r.w) - 1;
|
||||
if (triIdx < 0 || triIdx >= p.numTriangles)
|
||||
{
|
||||
((float4*)p.gradRaster)[pidx] = make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
if (ENABLE_DA)
|
||||
((float4*)p.gradRasterDB)[pidx] = make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
return;
|
||||
}
|
||||
|
||||
// Fetch vertex indices.
|
||||
int vi0 = p.tri[triIdx * 3 + 0];
|
||||
int vi1 = p.tri[triIdx * 3 + 1];
|
||||
int vi2 = p.tri[triIdx * 3 + 2];
|
||||
|
||||
// Bail out if corrupt indices.
|
||||
if (vi0 < 0 || vi0 >= p.numVertices ||
|
||||
vi1 < 0 || vi1 >= p.numVertices ||
|
||||
vi2 < 0 || vi2 >= p.numVertices)
|
||||
return;
|
||||
|
||||
// In instance mode, adjust vertex indices by minibatch index unless broadcasting.
|
||||
if (p.instance_mode && !p.attrBC)
|
||||
{
|
||||
vi0 += pz * p.numVertices;
|
||||
vi1 += pz * p.numVertices;
|
||||
vi2 += pz * p.numVertices;
|
||||
}
|
||||
|
||||
// Initialize coalesced atomics.
|
||||
CA_SET_GROUP(triIdx);
|
||||
|
||||
// Pointers to inputs.
|
||||
const float* a0 = p.attr + vi0 * p.numAttr;
|
||||
const float* a1 = p.attr + vi1 * p.numAttr;
|
||||
const float* a2 = p.attr + vi2 * p.numAttr;
|
||||
const float* pdy = p.dy + pidx * p.numAttr;
|
||||
|
||||
// Pointers to outputs.
|
||||
float* ga0 = p.gradAttr + vi0 * p.numAttr;
|
||||
float* ga1 = p.gradAttr + vi1 * p.numAttr;
|
||||
float* ga2 = p.gradAttr + vi2 * p.numAttr;
|
||||
|
||||
// Barys and bary gradient accumulators.
|
||||
float b0 = r.x;
|
||||
float b1 = r.y;
|
||||
float b2 = 1.f - r.x - r.y;
|
||||
float gb0 = 0.f;
|
||||
float gb1 = 0.f;
|
||||
|
||||
// Loop over attributes and accumulate attribute gradients.
|
||||
for (int i=0; i < p.numAttr; i++)
|
||||
{
|
||||
float y = pdy[i];
|
||||
float s0 = a0[i];
|
||||
float s1 = a1[i];
|
||||
float s2 = a2[i];
|
||||
gb0 += y * (s0 - s2);
|
||||
gb1 += y * (s1 - s2);
|
||||
caAtomicAdd(ga0 + i, b0 * y);
|
||||
caAtomicAdd(ga1 + i, b1 * y);
|
||||
caAtomicAdd(ga2 + i, b2 * y);
|
||||
}
|
||||
|
||||
// Write the bary gradients.
|
||||
((float4*)p.gradRaster)[pidx] = make_float4(gb0, gb1, 0.f, 0.f);
|
||||
|
||||
// If pixel differentials disabled, we're done.
|
||||
if (!ENABLE_DA)
|
||||
return;
|
||||
|
||||
// Calculate gradients based on attribute pixel differentials.
|
||||
const float2* dda = ((float2*)p.dda) + pidx * p.numDiffAttr;
|
||||
float gdudx = 0.f;
|
||||
float gdudy = 0.f;
|
||||
float gdvdx = 0.f;
|
||||
float gdvdy = 0.f;
|
||||
|
||||
// Read bary pixel differentials.
|
||||
float4 db = ((float4*)p.rastDB)[pidx];
|
||||
float dudx = db.x;
|
||||
float dudy = db.y;
|
||||
float dvdx = db.z;
|
||||
float dvdy = db.w;
|
||||
|
||||
for (int i=0; i < p.numDiffAttr; i++)
|
||||
{
|
||||
// Input attribute index.
|
||||
int j = p.diff_attrs_all ? i : p.diffAttrs[i];
|
||||
if (j < 0)
|
||||
j += p.numAttr; // Python-style negative indices.
|
||||
|
||||
// Check that index is valid.
|
||||
if (j >= 0 && j < p.numAttr)
|
||||
{
|
||||
float2 dsdxy = dda[i];
|
||||
float dsdx = dsdxy.x;
|
||||
float dsdy = dsdxy.y;
|
||||
|
||||
float s0 = a0[j];
|
||||
float s1 = a1[j];
|
||||
float s2 = a2[j];
|
||||
|
||||
// Gradients of db.
|
||||
float dsdu = s0 - s2;
|
||||
float dsdv = s1 - s2;
|
||||
gdudx += dsdu * dsdx;
|
||||
gdudy += dsdu * dsdy;
|
||||
gdvdx += dsdv * dsdx;
|
||||
gdvdy += dsdv * dsdy;
|
||||
|
||||
// Gradients of attributes.
|
||||
float du = dsdx*dudx + dsdy*dudy;
|
||||
float dv = dsdx*dvdx + dsdy*dvdy;
|
||||
caAtomicAdd(ga0 + j, du);
|
||||
caAtomicAdd(ga1 + j, dv);
|
||||
caAtomicAdd(ga2 + j, -du - dv);
|
||||
}
|
||||
}
|
||||
|
||||
// Write.
|
||||
((float4*)p.gradRasterDB)[pidx] = make_float4(gdudx, gdudy, gdvdx, gdvdy);
|
||||
}
|
||||
|
||||
// Template specializations.
|
||||
__global__ void InterpolateGradKernel (const InterpolateKernelParams p) { InterpolateGradKernelTemplate<false>(p); }
|
||||
__global__ void InterpolateGradKernelDa(const InterpolateKernelParams p) { InterpolateGradKernelTemplate<true>(p); }
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,49 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#pragma once
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Constants and helpers.
|
||||
|
||||
#define IP_FWD_MAX_KERNEL_BLOCK_WIDTH 8
|
||||
#define IP_FWD_MAX_KERNEL_BLOCK_HEIGHT 8
|
||||
#define IP_GRAD_MAX_KERNEL_BLOCK_WIDTH 8
|
||||
#define IP_GRAD_MAX_KERNEL_BLOCK_HEIGHT 8
|
||||
#define IP_MAX_DIFF_ATTRS 32
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// CUDA kernel params.
|
||||
|
||||
struct InterpolateKernelParams
|
||||
{
|
||||
const int* tri; // Incoming triangle buffer.
|
||||
const float* attr; // Incoming attribute buffer.
|
||||
const float* rast; // Incoming rasterizer output buffer.
|
||||
const float* rastDB; // Incoming rasterizer output buffer for bary derivatives.
|
||||
const float* dy; // Incoming attribute gradients.
|
||||
const float* dda; // Incoming attr diff gradients.
|
||||
float* out; // Outgoing interpolated attributes.
|
||||
float* outDA; // Outgoing texcoord major axis lengths.
|
||||
float* gradAttr; // Outgoing attribute gradients.
|
||||
float* gradRaster; // Outgoing rasterizer gradients.
|
||||
float* gradRasterDB; // Outgoing rasterizer bary diff gradients.
|
||||
int numTriangles; // Number of triangles.
|
||||
int numVertices; // Number of vertices.
|
||||
int numAttr; // Number of total vertex attributes.
|
||||
int numDiffAttr; // Number of attributes to differentiate.
|
||||
int width; // Image width.
|
||||
int height; // Image height.
|
||||
int depth; // Minibatch size.
|
||||
int attrBC; // 0=normal, 1=attr is broadcast.
|
||||
int instance_mode; // 0=normal, 1=instance mode.
|
||||
int diff_attrs_all; // 0=normal, 1=produce pixel differentials for all attributes.
|
||||
int diffAttrs[IP_MAX_DIFF_ATTRS]; // List of attributes to differentiate.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+279
@@ -0,0 +1,279 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "common.h"
|
||||
#include "rasterize.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Cuda forward rasterizer pixel shader kernel.
|
||||
|
||||
__global__ void RasterizeCudaFwdShaderKernel(const RasterizeCudaFwdShaderParams p)
|
||||
{
|
||||
// Calculate pixel position.
|
||||
int px = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int py = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int pz = blockIdx.z;
|
||||
if (px >= p.width_out || py >= p.height_out || pz >= p.depth)
|
||||
return;
|
||||
|
||||
// Pixel indices.
|
||||
int pidx_in = px + p.width_in * (py + p.height_in * pz);
|
||||
int pidx_out = px + p.width_out * (py + p.height_out * pz);
|
||||
|
||||
// Fetch triangle idx.
|
||||
int triIdx = p.in_idx[pidx_in] - 1;
|
||||
if (triIdx < 0 || triIdx >= p.numTriangles)
|
||||
{
|
||||
// No or corrupt triangle.
|
||||
((float4*)p.out)[pidx_out] = make_float4(0.0, 0.0, 0.0, 0.0); // Clear out.
|
||||
((float4*)p.out_db)[pidx_out] = make_float4(0.0, 0.0, 0.0, 0.0); // Clear out_db.
|
||||
return;
|
||||
}
|
||||
|
||||
// Fetch vertex indices.
|
||||
int vi0 = p.tri[triIdx * 3 + 0];
|
||||
int vi1 = p.tri[triIdx * 3 + 1];
|
||||
int vi2 = p.tri[triIdx * 3 + 2];
|
||||
|
||||
// Bail out if vertex indices are corrupt.
|
||||
if (vi0 < 0 || vi0 >= p.numVertices ||
|
||||
vi1 < 0 || vi1 >= p.numVertices ||
|
||||
vi2 < 0 || vi2 >= p.numVertices)
|
||||
return;
|
||||
|
||||
// In instance mode, adjust vertex indices by minibatch index.
|
||||
if (p.instance_mode)
|
||||
{
|
||||
vi0 += pz * p.numVertices;
|
||||
vi1 += pz * p.numVertices;
|
||||
vi2 += pz * p.numVertices;
|
||||
}
|
||||
|
||||
// Fetch vertex positions.
|
||||
float4 p0 = ((float4*)p.pos)[vi0];
|
||||
float4 p1 = ((float4*)p.pos)[vi1];
|
||||
float4 p2 = ((float4*)p.pos)[vi2];
|
||||
|
||||
// Evaluate edge functions.
|
||||
float fx = p.xs * (float)px + p.xo;
|
||||
float fy = p.ys * (float)py + p.yo;
|
||||
float p0x = p0.x - fx * p0.w;
|
||||
float p0y = p0.y - fy * p0.w;
|
||||
float p1x = p1.x - fx * p1.w;
|
||||
float p1y = p1.y - fy * p1.w;
|
||||
float p2x = p2.x - fx * p2.w;
|
||||
float p2y = p2.y - fy * p2.w;
|
||||
float a0 = p1x*p2y - p1y*p2x;
|
||||
float a1 = p2x*p0y - p2y*p0x;
|
||||
float a2 = p0x*p1y - p0y*p1x;
|
||||
|
||||
// Perspective correct, normalized barycentrics.
|
||||
float iw = 1.f / (a0 + a1 + a2);
|
||||
float b0 = a0 * iw;
|
||||
float b1 = a1 * iw;
|
||||
|
||||
// Compute z/w for depth buffer.
|
||||
float z = p0.z * a0 + p1.z * a1 + p2.z * a2;
|
||||
float w = p0.w * a0 + p1.w * a1 + p2.w * a2;
|
||||
float zw = z / w;
|
||||
|
||||
// Clamps to avoid NaNs.
|
||||
b0 = __saturatef(b0); // Clamp to [+0.0, 1.0].
|
||||
b1 = __saturatef(b1); // Clamp to [+0.0, 1.0].
|
||||
float bs = 1.f / fmaxf(b0 + b1, 1.f);
|
||||
b0 *= bs; // Clamp sum to 1.0 diagonally.
|
||||
b1 *= bs;
|
||||
zw = fmaxf(fminf(zw, 1.f), -1.f);
|
||||
|
||||
// Emit output.
|
||||
((float4*)p.out)[pidx_out] = make_float4(b0, b1, zw, triidx_to_float(triIdx + 1));
|
||||
|
||||
// Calculate bary pixel differentials.
|
||||
float dfxdx = p.xs * iw;
|
||||
float dfydy = p.ys * iw;
|
||||
float da0dx = p2.y*p1.w - p1.y*p2.w;
|
||||
float da0dy = p1.x*p2.w - p2.x*p1.w;
|
||||
float da1dx = p0.y*p2.w - p2.y*p0.w;
|
||||
float da1dy = p2.x*p0.w - p0.x*p2.w;
|
||||
float da2dx = p1.y*p0.w - p0.y*p1.w;
|
||||
float da2dy = p0.x*p1.w - p1.x*p0.w;
|
||||
float datdx = da0dx + da1dx + da2dx;
|
||||
float datdy = da0dy + da1dy + da2dy;
|
||||
float dudx = dfxdx * (b0 * datdx - da0dx);
|
||||
float dudy = dfydy * (b0 * datdy - da0dy);
|
||||
float dvdx = dfxdx * (b1 * datdx - da1dx);
|
||||
float dvdy = dfydy * (b1 * datdy - da1dy);
|
||||
|
||||
// Emit bary pixel differentials.
|
||||
((float4*)p.out_db)[pidx_out] = make_float4(dudx, dudy, dvdx, dvdy);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient Cuda kernel.
|
||||
|
||||
template <bool ENABLE_DB>
|
||||
static __forceinline__ __device__ void RasterizeGradKernelTemplate(const RasterizeGradParams p)
|
||||
{
|
||||
// Temporary space for coalesced atomics.
|
||||
CA_DECLARE_TEMP(RAST_GRAD_MAX_KERNEL_BLOCK_WIDTH * RAST_GRAD_MAX_KERNEL_BLOCK_HEIGHT);
|
||||
|
||||
// Calculate pixel position.
|
||||
int px = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int py = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int pz = blockIdx.z;
|
||||
if (px >= p.width || py >= p.height || pz >= p.depth)
|
||||
return;
|
||||
|
||||
// Pixel index.
|
||||
int pidx = px + p.width * (py + p.height * pz);
|
||||
|
||||
// Read triangle idx and dy.
|
||||
float2 dy = ((float2*)p.dy)[pidx * 2];
|
||||
float4 ddb = ENABLE_DB ? ((float4*)p.ddb)[pidx] : make_float4(0.f, 0.f, 0.f, 0.f);
|
||||
int triIdx = float_to_triidx(((float*)p.out)[pidx * 4 + 3]) - 1;
|
||||
|
||||
// Exit if nothing to do.
|
||||
if (triIdx < 0 || triIdx >= p.numTriangles)
|
||||
return; // No or corrupt triangle.
|
||||
int grad_all_dy = __float_as_int(dy.x) | __float_as_int(dy.y); // Bitwise OR of all incoming gradients.
|
||||
int grad_all_ddb = 0;
|
||||
if (ENABLE_DB)
|
||||
grad_all_ddb = __float_as_int(ddb.x) | __float_as_int(ddb.y) | __float_as_int(ddb.z) | __float_as_int(ddb.w);
|
||||
if (((grad_all_dy | grad_all_ddb) << 1) == 0)
|
||||
return; // All incoming gradients are +0/-0.
|
||||
|
||||
// Fetch vertex indices.
|
||||
int vi0 = p.tri[triIdx * 3 + 0];
|
||||
int vi1 = p.tri[triIdx * 3 + 1];
|
||||
int vi2 = p.tri[triIdx * 3 + 2];
|
||||
|
||||
// Bail out if vertex indices are corrupt.
|
||||
if (vi0 < 0 || vi0 >= p.numVertices ||
|
||||
vi1 < 0 || vi1 >= p.numVertices ||
|
||||
vi2 < 0 || vi2 >= p.numVertices)
|
||||
return;
|
||||
|
||||
// In instance mode, adjust vertex indices by minibatch index.
|
||||
if (p.instance_mode)
|
||||
{
|
||||
vi0 += pz * p.numVertices;
|
||||
vi1 += pz * p.numVertices;
|
||||
vi2 += pz * p.numVertices;
|
||||
}
|
||||
|
||||
// Initialize coalesced atomics.
|
||||
CA_SET_GROUP(triIdx);
|
||||
|
||||
// Fetch vertex positions.
|
||||
float4 p0 = ((float4*)p.pos)[vi0];
|
||||
float4 p1 = ((float4*)p.pos)[vi1];
|
||||
float4 p2 = ((float4*)p.pos)[vi2];
|
||||
|
||||
// Evaluate edge functions.
|
||||
float fx = p.xs * (float)px + p.xo;
|
||||
float fy = p.ys * (float)py + p.yo;
|
||||
float p0x = p0.x - fx * p0.w;
|
||||
float p0y = p0.y - fy * p0.w;
|
||||
float p1x = p1.x - fx * p1.w;
|
||||
float p1y = p1.y - fy * p1.w;
|
||||
float p2x = p2.x - fx * p2.w;
|
||||
float p2y = p2.y - fy * p2.w;
|
||||
float a0 = p1x*p2y - p1y*p2x;
|
||||
float a1 = p2x*p0y - p2y*p0x;
|
||||
float a2 = p0x*p1y - p0y*p1x;
|
||||
|
||||
// Compute inverse area with epsilon.
|
||||
float at = a0 + a1 + a2;
|
||||
float ep = copysignf(1e-6f, at); // ~1 pixel in 1k x 1k image.
|
||||
float iw = 1.f / (at + ep);
|
||||
|
||||
// Perspective correct, normalized barycentrics.
|
||||
float b0 = a0 * iw;
|
||||
float b1 = a1 * iw;
|
||||
|
||||
// Position gradients.
|
||||
float gb0 = dy.x * iw;
|
||||
float gb1 = dy.y * iw;
|
||||
float gbb = gb0 * b0 + gb1 * b1;
|
||||
float gp0x = gbb * (p2y - p1y) - gb1 * p2y;
|
||||
float gp1x = gbb * (p0y - p2y) + gb0 * p2y;
|
||||
float gp2x = gbb * (p1y - p0y) - gb0 * p1y + gb1 * p0y;
|
||||
float gp0y = gbb * (p1x - p2x) + gb1 * p2x;
|
||||
float gp1y = gbb * (p2x - p0x) - gb0 * p2x;
|
||||
float gp2y = gbb * (p0x - p1x) + gb0 * p1x - gb1 * p0x;
|
||||
float gp0w = -fx * gp0x - fy * gp0y;
|
||||
float gp1w = -fx * gp1x - fy * gp1y;
|
||||
float gp2w = -fx * gp2x - fy * gp2y;
|
||||
|
||||
// Bary differential gradients.
|
||||
if (ENABLE_DB && ((grad_all_ddb) << 1) != 0)
|
||||
{
|
||||
float dfxdX = p.xs * iw;
|
||||
float dfydY = p.ys * iw;
|
||||
ddb.x *= dfxdX;
|
||||
ddb.y *= dfydY;
|
||||
ddb.z *= dfxdX;
|
||||
ddb.w *= dfydY;
|
||||
|
||||
float da0dX = p1.y * p2.w - p2.y * p1.w;
|
||||
float da1dX = p2.y * p0.w - p0.y * p2.w;
|
||||
float da2dX = p0.y * p1.w - p1.y * p0.w;
|
||||
float da0dY = p2.x * p1.w - p1.x * p2.w;
|
||||
float da1dY = p0.x * p2.w - p2.x * p0.w;
|
||||
float da2dY = p1.x * p0.w - p0.x * p1.w;
|
||||
float datdX = da0dX + da1dX + da2dX;
|
||||
float datdY = da0dY + da1dY + da2dY;
|
||||
|
||||
float x01 = p0.x - p1.x;
|
||||
float x12 = p1.x - p2.x;
|
||||
float x20 = p2.x - p0.x;
|
||||
float y01 = p0.y - p1.y;
|
||||
float y12 = p1.y - p2.y;
|
||||
float y20 = p2.y - p0.y;
|
||||
float w01 = p0.w - p1.w;
|
||||
float w12 = p1.w - p2.w;
|
||||
float w20 = p2.w - p0.w;
|
||||
|
||||
float a0p1 = fy * p2.x - fx * p2.y;
|
||||
float a0p2 = fx * p1.y - fy * p1.x;
|
||||
float a1p0 = fx * p2.y - fy * p2.x;
|
||||
float a1p2 = fy * p0.x - fx * p0.y;
|
||||
|
||||
float wdudX = 2.f * b0 * datdX - da0dX;
|
||||
float wdudY = 2.f * b0 * datdY - da0dY;
|
||||
float wdvdX = 2.f * b1 * datdX - da1dX;
|
||||
float wdvdY = 2.f * b1 * datdY - da1dY;
|
||||
|
||||
float c0 = iw * (ddb.x * wdudX + ddb.y * wdudY + ddb.z * wdvdX + ddb.w * wdvdY);
|
||||
float cx = c0 * fx - ddb.x * b0 - ddb.z * b1;
|
||||
float cy = c0 * fy - ddb.y * b0 - ddb.w * b1;
|
||||
float cxy = iw * (ddb.x * datdX + ddb.y * datdY);
|
||||
float czw = iw * (ddb.z * datdX + ddb.w * datdY);
|
||||
|
||||
gp0x += c0 * y12 - cy * w12 + czw * p2y + ddb.w * p2.w;
|
||||
gp1x += c0 * y20 - cy * w20 - cxy * p2y - ddb.y * p2.w;
|
||||
gp2x += c0 * y01 - cy * w01 + cxy * p1y - czw * p0y + ddb.y * p1.w - ddb.w * p0.w;
|
||||
gp0y += cx * w12 - c0 * x12 - czw * p2x - ddb.z * p2.w;
|
||||
gp1y += cx * w20 - c0 * x20 + cxy * p2x + ddb.x * p2.w;
|
||||
gp2y += cx * w01 - c0 * x01 - cxy * p1x + czw * p0x - ddb.x * p1.w + ddb.z * p0.w;
|
||||
gp0w += cy * x12 - cx * y12 - czw * a1p0 + ddb.z * p2.y - ddb.w * p2.x;
|
||||
gp1w += cy * x20 - cx * y20 - cxy * a0p1 - ddb.x * p2.y + ddb.y * p2.x;
|
||||
gp2w += cy * x01 - cx * y01 - cxy * a0p2 - czw * a1p2 + ddb.x * p1.y - ddb.y * p1.x - ddb.z * p0.y + ddb.w * p0.x;
|
||||
}
|
||||
|
||||
// Accumulate using coalesced atomics.
|
||||
caAtomicAdd3_xyw(p.grad + 4 * vi0, gp0x, gp0y, gp0w);
|
||||
caAtomicAdd3_xyw(p.grad + 4 * vi1, gp1x, gp1y, gp1w);
|
||||
caAtomicAdd3_xyw(p.grad + 4 * vi2, gp2x, gp2y, gp2w);
|
||||
}
|
||||
|
||||
// Template specializations.
|
||||
__global__ void RasterizeGradKernel (const RasterizeGradParams p) { RasterizeGradKernelTemplate<false>(p); }
|
||||
__global__ void RasterizeGradKernelDb(const RasterizeGradParams p) { RasterizeGradKernelTemplate<true>(p); }
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#pragma once
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Constants and helpers.
|
||||
|
||||
#define RAST_CUDA_FWD_SHADER_KERNEL_BLOCK_WIDTH 8
|
||||
#define RAST_CUDA_FWD_SHADER_KERNEL_BLOCK_HEIGHT 8
|
||||
#define RAST_GRAD_MAX_KERNEL_BLOCK_WIDTH 8
|
||||
#define RAST_GRAD_MAX_KERNEL_BLOCK_HEIGHT 8
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// CUDA forward rasterizer shader kernel params.
|
||||
|
||||
struct RasterizeCudaFwdShaderParams
|
||||
{
|
||||
const float* pos; // Vertex positions.
|
||||
const int* tri; // Triangle indices.
|
||||
const int* in_idx; // Triangle idx buffer from rasterizer.
|
||||
float* out; // Main output buffer.
|
||||
float* out_db; // Bary pixel gradient output buffer.
|
||||
int numTriangles; // Number of triangles.
|
||||
int numVertices; // Number of vertices.
|
||||
int width_in; // Input image width.
|
||||
int height_in; // Input image height.
|
||||
int width_out; // Output image width.
|
||||
int height_out; // Output image height.
|
||||
int depth; // Size of minibatch.
|
||||
int instance_mode; // 1 if in instance rendering mode.
|
||||
float xs, xo, ys, yo; // Pixel position to clip-space x, y transform.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient CUDA kernel params.
|
||||
|
||||
struct RasterizeGradParams
|
||||
{
|
||||
const float* pos; // Incoming position buffer.
|
||||
const int* tri; // Incoming triangle buffer.
|
||||
const float* out; // Rasterizer output buffer.
|
||||
const float* dy; // Incoming gradients of rasterizer output buffer.
|
||||
const float* ddb; // Incoming gradients of bary diff output buffer.
|
||||
float* grad; // Outgoing position gradients.
|
||||
int numTriangles; // Number of triangles.
|
||||
int numVertices; // Number of vertices.
|
||||
int width; // Image width.
|
||||
int height; // Image height.
|
||||
int depth; // Size of minibatch.
|
||||
int instance_mode; // 1 if in instance rendering mode.
|
||||
float xs, xo, ys, yo; // Pixel position to clip-space x, y transform.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "framework.h"
|
||||
#include "texture.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Mip stack construction and access helpers.
|
||||
|
||||
void raiseMipSizeError(const TextureKernelParams& p)
|
||||
{
|
||||
char buf[1024];
|
||||
int bufsz = 1024;
|
||||
|
||||
std::string msg = "Mip-map size error - cannot downsample an odd extent greater than 1. Resize the texture so that both spatial extents are powers of two, or limit the number of mip maps using max_mip_level argument.\n";
|
||||
|
||||
int w = p.texWidth;
|
||||
int h = p.texHeight;
|
||||
bool ew = false;
|
||||
bool eh = false;
|
||||
|
||||
msg += "Attempted mip stack construction:\n";
|
||||
msg += "level width height\n";
|
||||
msg += "----- ----- ------\n";
|
||||
snprintf(buf, bufsz, "base %5d %5d\n", w, h);
|
||||
msg += buf;
|
||||
|
||||
int mipTotal = 0;
|
||||
int level = 0;
|
||||
while ((w|h) > 1 && !(ew || eh)) // Stop at first impossible size.
|
||||
{
|
||||
// Current level.
|
||||
level += 1;
|
||||
|
||||
// Determine if downsampling fails.
|
||||
ew = ew || (w > 1 && (w & 1));
|
||||
eh = eh || (h > 1 && (h & 1));
|
||||
|
||||
// Downsample.
|
||||
if (w > 1) w >>= 1;
|
||||
if (h > 1) h >>= 1;
|
||||
|
||||
// Append level size to error message.
|
||||
snprintf(buf, bufsz, "mip %-2d ", level);
|
||||
msg += buf;
|
||||
if (ew) snprintf(buf, bufsz, " err ");
|
||||
else snprintf(buf, bufsz, "%5d ", w);
|
||||
msg += buf;
|
||||
if (eh) snprintf(buf, bufsz, " err\n");
|
||||
else snprintf(buf, bufsz, "%5d\n", h);
|
||||
msg += buf;
|
||||
}
|
||||
|
||||
NVDR_CHECK(0, msg);
|
||||
}
|
||||
|
||||
int calculateMipInfo(TextureKernelParams& p, int* mipOffsets)
|
||||
{
|
||||
// No levels at all?
|
||||
if (p.mipLevelLimit == 0)
|
||||
{
|
||||
p.mipLevelMax = 0;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Current level size.
|
||||
int w = p.texWidth;
|
||||
int h = p.texHeight;
|
||||
|
||||
int mipTotal = 0;
|
||||
int level = 0;
|
||||
int c = (p.boundaryMode == TEX_BOUNDARY_MODE_CUBE) ? (p.channels * 6) : p.channels;
|
||||
mipOffsets[0] = 0;
|
||||
while ((w|h) > 1)
|
||||
{
|
||||
// Current level.
|
||||
level += 1;
|
||||
|
||||
// Quit if cannot downsample.
|
||||
if ((w > 1 && (w & 1)) || (h > 1 && (h & 1)))
|
||||
raiseMipSizeError(p);
|
||||
|
||||
// Downsample.
|
||||
if (w > 1) w >>= 1;
|
||||
if (h > 1) h >>= 1;
|
||||
|
||||
mipOffsets[level] = mipTotal; // Store the mip offset (#floats).
|
||||
mipTotal += w * h * p.texDepth * c;
|
||||
|
||||
// Hit the level limit?
|
||||
if (p.mipLevelLimit >= 0 && level == p.mipLevelLimit)
|
||||
break;
|
||||
}
|
||||
|
||||
p.mipLevelMax = level;
|
||||
return mipTotal;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#pragma once
|
||||
#include "framework.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Constants.
|
||||
|
||||
#define TEX_DEBUG_MIP_RETAIN_VARIANCE 0 // For debugging
|
||||
#define TEX_FWD_MAX_KERNEL_BLOCK_WIDTH 8
|
||||
#define TEX_FWD_MAX_KERNEL_BLOCK_HEIGHT 8
|
||||
#define TEX_FWD_MAX_MIP_KERNEL_BLOCK_WIDTH 8
|
||||
#define TEX_FWD_MAX_MIP_KERNEL_BLOCK_HEIGHT 8
|
||||
#define TEX_GRAD_MAX_KERNEL_BLOCK_WIDTH 8
|
||||
#define TEX_GRAD_MAX_KERNEL_BLOCK_HEIGHT 8
|
||||
#define TEX_GRAD_MAX_MIP_KERNEL_BLOCK_WIDTH 8
|
||||
#define TEX_GRAD_MAX_MIP_KERNEL_BLOCK_HEIGHT 8
|
||||
#define TEX_MAX_MIP_LEVEL 16 // Currently a texture cannot be larger than 2 GB because we use 32-bit indices everywhere.
|
||||
#define TEX_MODE_NEAREST 0 // Nearest on base level.
|
||||
#define TEX_MODE_LINEAR 1 // Bilinear on base level.
|
||||
#define TEX_MODE_LINEAR_MIPMAP_NEAREST 2 // Bilinear on nearest mip level.
|
||||
#define TEX_MODE_LINEAR_MIPMAP_LINEAR 3 // Trilinear.
|
||||
#define TEX_MODE_COUNT 4
|
||||
#define TEX_BOUNDARY_MODE_CUBE 0 // Cube map mode.
|
||||
#define TEX_BOUNDARY_MODE_WRAP 1 // Wrap (u, v).
|
||||
#define TEX_BOUNDARY_MODE_CLAMP 2 // Clamp (u, v).
|
||||
#define TEX_BOUNDARY_MODE_ZERO 3 // Pad with zeros.
|
||||
#define TEX_BOUNDARY_MODE_COUNT 4
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// CUDA kernel params.
|
||||
|
||||
struct TextureKernelParams
|
||||
{
|
||||
const float* tex[TEX_MAX_MIP_LEVEL]; // Incoming texture buffer with mip levels.
|
||||
const float* uv; // Incoming texcoord buffer.
|
||||
const float* uvDA; // Incoming uv pixel diffs or NULL.
|
||||
const float* mipLevelBias; // Incoming mip level bias or NULL.
|
||||
const float* dy; // Incoming output gradient.
|
||||
float* out; // Outgoing texture data.
|
||||
float* gradTex[TEX_MAX_MIP_LEVEL]; // Outgoing texture gradients with mip levels.
|
||||
float* gradUV; // Outgoing texcoord gradient.
|
||||
float* gradUVDA; // Outgoing texcoord pixel differential gradient.
|
||||
float* gradMipLevelBias; // Outgoing mip level bias gradient.
|
||||
int enableMip; // If true, we have uv_da and/or mip_level_bias input(s), and a mip tensor.
|
||||
int filterMode; // One of the TEX_MODE_ constants.
|
||||
int boundaryMode; // One of the TEX_BOUNDARY_MODE_ contants.
|
||||
int texConst; // If true, texture is known to be constant.
|
||||
int mipLevelLimit; // Mip level limit coming from the op.
|
||||
int channels; // Number of texture channels.
|
||||
int imgWidth; // Image width.
|
||||
int imgHeight; // Image height.
|
||||
int texWidth; // Texture width.
|
||||
int texHeight; // Texture height.
|
||||
int texDepth; // Texture depth.
|
||||
int n; // Minibatch size.
|
||||
int mipLevelMax; // Maximum mip level index. Zero if mips disabled.
|
||||
int mipLevelOut; // Mip level being calculated in builder kernel.
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// C++ helper function prototypes.
|
||||
|
||||
void raiseMipSizeError(const TextureKernelParams& p);
|
||||
int calculateMipInfo(TextureKernelParams& p, int* mipOffsets);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Macros.
|
||||
|
||||
#define mipLevelSize(p, i) make_int2(((p).texWidth >> (i)) > 1 ? ((p).texWidth >> (i)) : 1, ((p).texHeight >> (i)) > 1 ? ((p).texHeight >> (i)) : 1)
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+1156
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,243 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "torch_common.inl"
|
||||
#include "torch_types.h"
|
||||
#include "../common/common.h"
|
||||
#include "../common/antialias.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Kernel prototypes.
|
||||
|
||||
void AntialiasFwdMeshKernel (const AntialiasKernelParams p);
|
||||
void AntialiasFwdDiscontinuityKernel(const AntialiasKernelParams p);
|
||||
void AntialiasFwdAnalysisKernel (const AntialiasKernelParams p);
|
||||
void AntialiasGradKernel (const AntialiasKernelParams p);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Topology hash construction.
|
||||
|
||||
TopologyHashWrapper antialias_construct_topology_hash(torch::Tensor tri)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(tri));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
AntialiasKernelParams p = {}; // Initialize all fields to zero.
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(tri);
|
||||
NVDR_CHECK_CONTIGUOUS(tri);
|
||||
NVDR_CHECK_I32(tri);
|
||||
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
|
||||
// Fill in kernel parameters.
|
||||
p.numTriangles = tri.size(0);
|
||||
p.numVertices = 0x7fffffff; // Let's not require vertex positions just to enable an error check.
|
||||
p.tri = tri.data_ptr<int>();
|
||||
|
||||
// Kernel parameters.
|
||||
p.allocTriangles = 64;
|
||||
while (p.allocTriangles < p.numTriangles)
|
||||
p.allocTriangles <<= 1; // Must be power of two.
|
||||
|
||||
// Construct the hash tensor and get pointer.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA);
|
||||
torch::Tensor ev_hash = torch::zeros({(uint64_t)p.allocTriangles * AA_HASH_ELEMENTS_PER_TRIANGLE(p.allocTriangles) * 4}, opts);
|
||||
p.evHash = (uint4*)(ev_hash.data_ptr<int>());
|
||||
|
||||
// Check alignment.
|
||||
NVDR_CHECK(!((uintptr_t)p.evHash & 15), "ev_hash internal tensor not aligned to int4");
|
||||
|
||||
// Populate the hash.
|
||||
void* args[] = {&p};
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasFwdMeshKernel, (p.numTriangles - 1) / AA_MESH_KERNEL_THREADS_PER_BLOCK + 1, AA_MESH_KERNEL_THREADS_PER_BLOCK, args, 0, stream));
|
||||
|
||||
// Return.
|
||||
TopologyHashWrapper hash_wrap;
|
||||
hash_wrap.ev_hash = ev_hash;
|
||||
return hash_wrap;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Forward op.
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> antialias_fwd(torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, TopologyHashWrapper topology_hash_wrap)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(color));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
AntialiasKernelParams p = {}; // Initialize all fields to zero.
|
||||
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
|
||||
torch::Tensor& topology_hash = topology_hash_wrap.ev_hash; // Unwrap.
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(color, rast, pos, tri, topology_hash);
|
||||
NVDR_CHECK_CONTIGUOUS(color, rast, pos, tri, topology_hash);
|
||||
NVDR_CHECK_F32(color, rast, pos);
|
||||
NVDR_CHECK_I32(tri, topology_hash);
|
||||
|
||||
// Sanity checks.
|
||||
NVDR_CHECK(color.sizes().size() == 4 && color.size(0) > 0 && color.size(1) > 0 && color.size(2) > 0 && color.size(3) > 0, "color must have shape[>0, >0, >0, >0]");
|
||||
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "rast must have shape[>0, >0, >0, 4]");
|
||||
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
NVDR_CHECK(color.size(1) == rast.size(1) && color.size(2) == rast.size(2), "color and rast inputs must have same spatial dimensions");
|
||||
if (p.instance_mode)
|
||||
{
|
||||
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) > 0 && pos.size(1) > 0 && pos.size(2) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
|
||||
NVDR_CHECK(rast.size(0) == color.size(0) && pos.size(0) == color.size(0), "minibatch size mismatch between inputs color, rast, pos");
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
|
||||
NVDR_CHECK(rast.size(0) == color.size(0), "minibatch size mismatch between inputs color, rast");
|
||||
}
|
||||
|
||||
// Extract input dimensions.
|
||||
p.numVertices = pos.size(p.instance_mode ? 1 : 0);
|
||||
p.numTriangles = tri.size(0);
|
||||
p.n = color.size(0);
|
||||
p.height = color.size(1);
|
||||
p.width = color.size(2);
|
||||
p.channels = color.size(3);
|
||||
|
||||
// Get input pointers.
|
||||
p.color = color.data_ptr<float>();
|
||||
p.rasterOut = rast.data_ptr<float>();
|
||||
p.tri = tri.data_ptr<int>();
|
||||
p.pos = pos.data_ptr<float>();
|
||||
p.evHash = (uint4*)(topology_hash.data_ptr<int>());
|
||||
|
||||
// Misc parameters.
|
||||
p.xh = .5f * (float)p.width;
|
||||
p.yh = .5f * (float)p.height;
|
||||
|
||||
// Determine hash allocation size.
|
||||
p.allocTriangles = 64;
|
||||
while (p.allocTriangles < p.numTriangles)
|
||||
p.allocTriangles <<= 1; // Must be power of two.
|
||||
|
||||
// Allocate output tensors.
|
||||
torch::Tensor out = color.detach().clone(); // Use color as base.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
torch::Tensor work_buffer = torch::empty({p.n * p.width * p.height * 8 + 4}, opts); // 8 int for a maximum of two work items per pixel.
|
||||
p.output = out.data_ptr<float>();
|
||||
p.workBuffer = (int4*)(work_buffer.data_ptr<float>());
|
||||
|
||||
// Clear the work counters.
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMemsetAsync(p.workBuffer, 0, sizeof(int4), stream));
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations.
|
||||
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.rasterOut & 7), "raster_out input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.workBuffer & 15), "work_buffer internal tensor not aligned to int4");
|
||||
NVDR_CHECK(!((uintptr_t)p.evHash & 15), "topology_hash internal tensor not aligned to int4");
|
||||
|
||||
// Choose launch parameters for the discontinuity finder kernel and launch.
|
||||
void* args[] = {&p};
|
||||
dim3 blockSize(AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH, AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT, 1);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.n);
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasFwdDiscontinuityKernel, gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Determine optimum block size for the persistent analysis kernel and launch.
|
||||
int device = 0;
|
||||
int numCTA = 0;
|
||||
int numSM = 0;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaGetDevice(&device));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numCTA, (void*)AntialiasFwdAnalysisKernel, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK, 0));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, device));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasFwdAnalysisKernel, numCTA * numSM, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK, args, 0, stream));
|
||||
|
||||
// Return results.
|
||||
return std::tuple<torch::Tensor, torch::Tensor>(out, work_buffer);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient op.
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> antialias_grad(torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, torch::Tensor dy, torch::Tensor work_buffer)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(color));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
AntialiasKernelParams p = {}; // Initialize all fields to zero.
|
||||
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(color, rast, pos, tri, dy, work_buffer);
|
||||
NVDR_CHECK_CONTIGUOUS(color, rast, pos, tri, work_buffer);
|
||||
NVDR_CHECK_F32(color, rast, pos, dy, work_buffer);
|
||||
NVDR_CHECK_I32(tri);
|
||||
|
||||
// Sanity checks.
|
||||
NVDR_CHECK(dy.sizes().size() == 4 && dy.size(0) > 0 && dy.size(1) > 0 && dy.size(2) > 0 && dy.size(3) > 0, "dy must have shape[>0, >0, >0, >0]");
|
||||
NVDR_CHECK(color.sizes().size() == 4 && color.size(0) > 0 && color.size(1) > 0 && color.size(2) > 0 && color.size(3) > 0, "color must have shape[>0, >0, >0, >0]");
|
||||
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "raster_out must have shape[>0, >0, >0, 4]");
|
||||
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
NVDR_CHECK(color.size(1) == rast.size(1) && color.size(2) == rast.size(2), "color and raster_out inputs must have same spatial dimensions");
|
||||
NVDR_CHECK(color.size(1) == dy.size(1) && color.size(2) == dy.size(2) && color.size(3) == dy.size(3), "color and dy inputs must have same dimensions");
|
||||
if (p.instance_mode)
|
||||
{
|
||||
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) > 0 && pos.size(1) > 0 && pos.size(2) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
|
||||
NVDR_CHECK(rast.size(0) == color.size(0) && pos.size(0) == color.size(0), "minibatch size mismatch between inputs color, raster_out, pos");
|
||||
NVDR_CHECK(dy.size(0) == color.size(0) && rast.size(0) == color.size(0) && pos.size(0) ==color.size(0), "minibatch size mismatch between inputs dy, color, raster_out, pos");
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "pos must have shape [>0, >0, 4] or [>0, 4]");
|
||||
NVDR_CHECK(rast.size(0) == color.size(0), "minibatch size mismatch between inputs color, raster_out");
|
||||
NVDR_CHECK(dy.size(0) == color.size(0) && rast.size(0) == color.size(0), "minibatch size mismatch between inputs dy, color, raster_out");
|
||||
}
|
||||
|
||||
// Extract input dimensions.
|
||||
p.numVertices = pos.size(p.instance_mode ? 1 : 0);
|
||||
p.numTriangles = tri.size(0);
|
||||
p.n = color.size(0);
|
||||
p.height = color.size(1);
|
||||
p.width = color.size(2);
|
||||
p.channels = color.size(3);
|
||||
|
||||
// Ensure dy is contiguous.
|
||||
torch::Tensor dy_ = dy.contiguous();
|
||||
|
||||
// Get input pointers.
|
||||
p.color = color.data_ptr<float>();
|
||||
p.rasterOut = rast.data_ptr<float>();
|
||||
p.tri = tri.data_ptr<int>();
|
||||
p.pos = pos.data_ptr<float>();
|
||||
p.dy = dy_.data_ptr<float>();
|
||||
p.workBuffer = (int4*)(work_buffer.data_ptr<float>());
|
||||
|
||||
// Misc parameters.
|
||||
p.xh = .5f * (float)p.width;
|
||||
p.yh = .5f * (float)p.height;
|
||||
|
||||
// Allocate output tensors.
|
||||
torch::Tensor grad_color = dy_.detach().clone(); // Use dy as base.
|
||||
torch::Tensor grad_pos = torch::zeros_like(pos);
|
||||
p.gradColor = grad_color.data_ptr<float>();
|
||||
p.gradPos = grad_pos.data_ptr<float>();
|
||||
|
||||
// Clear gradient kernel work counter.
|
||||
NVDR_CHECK_CUDA_ERROR(cudaMemsetAsync(&p.workBuffer[0].y, 0, sizeof(int), stream));
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations.
|
||||
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.workBuffer & 15), "work_buffer internal tensor not aligned to int4");
|
||||
|
||||
// Determine optimum block size for the gradient kernel and launch.
|
||||
void* args[] = {&p};
|
||||
int device = 0;
|
||||
int numCTA = 0;
|
||||
int numSM = 0;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaGetDevice(&device));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numCTA, (void*)AntialiasGradKernel, AA_GRAD_KERNEL_THREADS_PER_BLOCK, 0));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, device));
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)AntialiasGradKernel, numCTA * numSM, AA_GRAD_KERNEL_THREADS_PER_BLOCK, args, 0, stream));
|
||||
|
||||
// Return results.
|
||||
return std::tuple<torch::Tensor, torch::Tensor>(grad_color, grad_pos);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,73 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "torch_common.inl"
|
||||
#include "torch_types.h"
|
||||
#include <tuple>
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Op prototypes. Return type macros for readability.
|
||||
|
||||
#define OP_RETURN_T torch::Tensor
|
||||
#define OP_RETURN_TT std::tuple<torch::Tensor, torch::Tensor>
|
||||
#define OP_RETURN_TTT std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
|
||||
#define OP_RETURN_TTTT std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
|
||||
#define OP_RETURN_TTV std::tuple<torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >
|
||||
#define OP_RETURN_TTTTV std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >
|
||||
|
||||
OP_RETURN_TT rasterize_fwd_cuda (RasterizeCRStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple<int, int> resolution, torch::Tensor ranges, int peeling_idx);
|
||||
OP_RETURN_T rasterize_grad (torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy);
|
||||
OP_RETURN_T rasterize_grad_db (torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy, torch::Tensor ddb);
|
||||
OP_RETURN_TT interpolate_fwd (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri);
|
||||
OP_RETURN_TT interpolate_fwd_da (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor rast_db, bool diff_attrs_all, std::vector<int>& diff_attrs_vec);
|
||||
OP_RETURN_TT interpolate_grad (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy);
|
||||
OP_RETURN_TTT interpolate_grad_da (torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy, torch::Tensor rast_db, torch::Tensor dda, bool diff_attrs_all, std::vector<int>& diff_attrs_vec);
|
||||
TextureMipWrapper texture_construct_mip (torch::Tensor tex, int max_mip_level, bool cube_mode);
|
||||
OP_RETURN_T texture_fwd (torch::Tensor tex, torch::Tensor uv, int filter_mode, int boundary_mode);
|
||||
OP_RETURN_T texture_fwd_mip (torch::Tensor tex, torch::Tensor uv, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode);
|
||||
OP_RETURN_T texture_grad_nearest (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode);
|
||||
OP_RETURN_TT texture_grad_linear (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode);
|
||||
OP_RETURN_TTV texture_grad_linear_mipmap_nearest (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode);
|
||||
OP_RETURN_TTTTV texture_grad_linear_mipmap_linear (torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode);
|
||||
TopologyHashWrapper antialias_construct_topology_hash (torch::Tensor tri);
|
||||
OP_RETURN_TT antialias_fwd (torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, TopologyHashWrapper topology_hash);
|
||||
OP_RETURN_TT antialias_grad (torch::Tensor color, torch::Tensor rast, torch::Tensor pos, torch::Tensor tri, torch::Tensor dy, torch::Tensor work_buffer);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
// State classes.
|
||||
pybind11::class_<RasterizeCRStateWrapper>(m, "RasterizeCRStateWrapper").def(pybind11::init<int>());
|
||||
pybind11::class_<TextureMipWrapper>(m, "TextureMipWrapper").def(pybind11::init<>());
|
||||
pybind11::class_<TopologyHashWrapper>(m, "TopologyHashWrapper");
|
||||
|
||||
// Plumbing to torch/c10 logging system.
|
||||
m.def("get_log_level", [](void) { return FLAGS_caffe2_log_level; }, "get log level");
|
||||
m.def("set_log_level", [](int level){ FLAGS_caffe2_log_level = level; }, "set log level");
|
||||
|
||||
// Ops.
|
||||
m.def("rasterize_fwd_cuda", &rasterize_fwd_cuda, "rasterize forward op (cuda)");
|
||||
m.def("rasterize_grad", &rasterize_grad, "rasterize gradient op ignoring db gradients");
|
||||
m.def("rasterize_grad_db", &rasterize_grad_db, "rasterize gradient op with db gradients");
|
||||
m.def("interpolate_fwd", &interpolate_fwd, "interpolate forward op with attribute derivatives");
|
||||
m.def("interpolate_fwd_da", &interpolate_fwd_da, "interpolate forward op without attribute derivatives");
|
||||
m.def("interpolate_grad", &interpolate_grad, "interpolate gradient op with attribute derivatives");
|
||||
m.def("interpolate_grad_da", &interpolate_grad_da, "interpolate gradient op without attribute derivatives");
|
||||
m.def("texture_construct_mip", &texture_construct_mip, "texture mipmap construction");
|
||||
m.def("texture_fwd", &texture_fwd, "texture forward op without mipmapping");
|
||||
m.def("texture_fwd_mip", &texture_fwd_mip, "texture forward op with mipmapping");
|
||||
m.def("texture_grad_nearest", &texture_grad_nearest, "texture gradient op in nearest mode");
|
||||
m.def("texture_grad_linear", &texture_grad_linear, "texture gradient op in linear mode");
|
||||
m.def("texture_grad_linear_mipmap_nearest", &texture_grad_linear_mipmap_nearest, "texture gradient op in linear-mipmap-nearest mode");
|
||||
m.def("texture_grad_linear_mipmap_linear", &texture_grad_linear_mipmap_linear, "texture gradient op in linear-mipmap-linear mode");
|
||||
m.def("antialias_construct_topology_hash", &antialias_construct_topology_hash, "antialias topology hash construction");
|
||||
m.def("antialias_fwd", &antialias_fwd, "antialias forward op");
|
||||
m.def("antialias_grad", &antialias_grad, "antialias gradient op");
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,250 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "torch_common.inl"
|
||||
#include "../common/common.h"
|
||||
#include "../common/interpolate.h"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Kernel prototypes.
|
||||
|
||||
void InterpolateFwdKernel (const InterpolateKernelParams p);
|
||||
void InterpolateFwdKernelDa (const InterpolateKernelParams p);
|
||||
void InterpolateGradKernel (const InterpolateKernelParams p);
|
||||
void InterpolateGradKernelDa(const InterpolateKernelParams p);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Helper
|
||||
|
||||
static void set_diff_attrs(InterpolateKernelParams& p, bool diff_attrs_all, std::vector<int>& diff_attrs_vec)
|
||||
{
|
||||
if (diff_attrs_all)
|
||||
{
|
||||
p.numDiffAttr = p.numAttr;
|
||||
p.diff_attrs_all = 1;
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(diff_attrs_vec.size() <= IP_MAX_DIFF_ATTRS, "too many entries in diff_attrs list (increase IP_MAX_DIFF_ATTRS)");
|
||||
p.numDiffAttr = diff_attrs_vec.size();
|
||||
memcpy(p.diffAttrs, &diff_attrs_vec[0], diff_attrs_vec.size()*sizeof(int));
|
||||
}
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Forward op.
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> interpolate_fwd_da(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor rast_db, bool diff_attrs_all, std::vector<int>& diff_attrs_vec)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(attr));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
InterpolateKernelParams p = {}; // Initialize all fields to zero.
|
||||
bool enable_da = (rast_db.defined()) && (diff_attrs_all || !diff_attrs_vec.empty());
|
||||
p.instance_mode = (attr.sizes().size() > 2) ? 1 : 0;
|
||||
|
||||
// Check inputs.
|
||||
if (enable_da)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(attr, rast, tri, rast_db);
|
||||
NVDR_CHECK_CONTIGUOUS(attr, rast, tri, rast_db);
|
||||
NVDR_CHECK_F32(attr, rast, rast_db);
|
||||
NVDR_CHECK_I32(tri);
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK_DEVICE(attr, rast, tri);
|
||||
NVDR_CHECK_CONTIGUOUS(attr, rast, tri);
|
||||
NVDR_CHECK_F32(attr, rast);
|
||||
NVDR_CHECK_I32(tri);
|
||||
}
|
||||
|
||||
// Sanity checks.
|
||||
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "rast must have shape[>0, >0, >0, 4]");
|
||||
NVDR_CHECK( tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
NVDR_CHECK((attr.sizes().size() == 2 || attr.sizes().size() == 3) && attr.size(0) > 0 && attr.size(1) > 0 && (attr.sizes().size() == 2 || attr.size(2) > 0), "attr must have shape [>0, >0, >0] or [>0, >0]");
|
||||
if (p.instance_mode)
|
||||
NVDR_CHECK(attr.size(0) == rast.size(0) || attr.size(0) == 1, "minibatch size mismatch between inputs rast, attr");
|
||||
if (enable_da)
|
||||
{
|
||||
NVDR_CHECK(rast_db.sizes().size() == 4 && rast_db.size(0) > 0 && rast_db.size(1) > 0 && rast_db.size(2) > 0 && rast_db.size(3) == 4, "rast_db must have shape[>0, >0, >0, 4]");
|
||||
NVDR_CHECK(rast_db.size(1) == rast.size(1) && rast_db.size(2) == rast.size(2), "spatial size mismatch between inputs rast and rast_db");
|
||||
NVDR_CHECK(rast_db.size(0) == rast.size(0), "minibatch size mismatch between inputs rast, rast_db");
|
||||
}
|
||||
|
||||
// Extract input dimensions.
|
||||
p.numVertices = attr.size(p.instance_mode ? 1 : 0);
|
||||
p.numAttr = attr.size(p.instance_mode ? 2 : 1);
|
||||
p.numTriangles = tri.size(0);
|
||||
p.height = rast.size(1);
|
||||
p.width = rast.size(2);
|
||||
p.depth = rast.size(0);
|
||||
|
||||
// Set attribute pixel differential info if enabled, otherwise leave as zero.
|
||||
if (enable_da)
|
||||
set_diff_attrs(p, diff_attrs_all, diff_attrs_vec);
|
||||
else
|
||||
p.numDiffAttr = 0;
|
||||
|
||||
// Get input pointers.
|
||||
p.attr = attr.data_ptr<float>();
|
||||
p.rast = rast.data_ptr<float>();
|
||||
p.tri = tri.data_ptr<int>();
|
||||
p.rastDB = enable_da ? rast_db.data_ptr<float>() : NULL;
|
||||
p.attrBC = (p.instance_mode && attr.size(0) == 1) ? 1 : 0;
|
||||
|
||||
// Allocate output tensors.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
torch::Tensor out = torch::empty({p.depth, p.height, p.width, p.numAttr}, opts);
|
||||
torch::Tensor out_da = torch::empty({p.depth, p.height, p.width, p.numDiffAttr * 2}, opts);
|
||||
|
||||
p.out = out.data_ptr<float>();
|
||||
p.outDA = enable_da ? out_da.data_ptr<float>() : NULL;
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations.
|
||||
NVDR_CHECK(!((uintptr_t)p.rast & 15), "rast input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.rastDB & 15), "rast_db input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.outDA & 7), "out_da output tensor not aligned to float2");
|
||||
|
||||
// Choose launch parameters.
|
||||
dim3 blockSize = getLaunchBlockSize(IP_FWD_MAX_KERNEL_BLOCK_WIDTH, IP_FWD_MAX_KERNEL_BLOCK_HEIGHT, p.width, p.height);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.depth);
|
||||
|
||||
// Launch CUDA kernel.
|
||||
void* args[] = {&p};
|
||||
void* func = enable_da ? (void*)InterpolateFwdKernelDa : (void*)InterpolateFwdKernel;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Return results.
|
||||
return std::tuple<torch::Tensor, torch::Tensor>(out, out_da);
|
||||
}
|
||||
|
||||
// Version without derivatives.
|
||||
std::tuple<torch::Tensor, torch::Tensor> interpolate_fwd(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri)
|
||||
{
|
||||
std::vector<int> empty_vec;
|
||||
torch::Tensor empty_tensor;
|
||||
return interpolate_fwd_da(attr, rast, tri, empty_tensor, false, empty_vec);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient op.
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> interpolate_grad_da(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy, torch::Tensor rast_db, torch::Tensor dda, bool diff_attrs_all, std::vector<int>& diff_attrs_vec)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(attr));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
InterpolateKernelParams p = {}; // Initialize all fields to zero.
|
||||
bool enable_da = (rast_db.defined()) && (diff_attrs_all || !diff_attrs_vec.empty());
|
||||
p.instance_mode = (attr.sizes().size() > 2) ? 1 : 0;
|
||||
|
||||
// Check inputs.
|
||||
if (enable_da)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(attr, rast, tri, dy, rast_db, dda);
|
||||
NVDR_CHECK_CONTIGUOUS(attr, rast, tri, rast_db);
|
||||
NVDR_CHECK_F32(attr, rast, dy, rast_db, dda);
|
||||
NVDR_CHECK_I32(tri);
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK_DEVICE(attr, rast, tri, dy);
|
||||
NVDR_CHECK_CONTIGUOUS(attr, rast, tri);
|
||||
NVDR_CHECK_F32(attr, rast, dy);
|
||||
NVDR_CHECK_I32(tri);
|
||||
}
|
||||
|
||||
// Depth of attributes.
|
||||
int attr_depth = p.instance_mode ? (attr.sizes().size() > 1 ? attr.size(0) : 0) : 1;
|
||||
|
||||
// Sanity checks.
|
||||
NVDR_CHECK(rast.sizes().size() == 4 && rast.size(0) > 0 && rast.size(1) > 0 && rast.size(2) > 0 && rast.size(3) == 4, "rast must have shape[>0, >0, >0, 4]");
|
||||
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
NVDR_CHECK((attr.sizes().size() == 2 || attr.sizes().size() == 3) && attr.size(0) > 0 && attr.size(1) > 0 && (attr.sizes().size() == 2 || attr.size(2) > 0), "attr must have shape [>0, >0, >0] or [>0, >0]");
|
||||
NVDR_CHECK(dy.sizes().size() == 4 && dy.size(0) > 0 && dy.size(1) == rast.size(1) && dy.size(2) == rast.size(2) && dy.size(3) > 0, "dy must have shape [>0, height, width, >0]");
|
||||
NVDR_CHECK(dy.size(3) == attr.size(attr.sizes().size() - 1), "argument count mismatch between inputs dy, attr");
|
||||
NVDR_CHECK((attr_depth == rast.size(0) || attr_depth == 1) && dy.size(0) == rast.size(0), "minibatch size mismatch between inputs rast, dy, attr");
|
||||
if (enable_da)
|
||||
{
|
||||
NVDR_CHECK(dda.sizes().size() == 4 && dda.size(0) > 0 && dda.size(1) == rast.size(1) && dda.size(2) == rast.size(2), "dda must have shape [>0, height, width, ?]");
|
||||
NVDR_CHECK(dda.size(0) == rast.size(0), "minibatch size mismatch between rast, dda");
|
||||
NVDR_CHECK(rast_db.sizes().size() == 4 && rast_db.size(0) > 0 && rast_db.size(1) > 0 && rast_db.size(2) > 0 && rast_db.size(3) == 4, "rast_db must have shape[>0, >0, >0, 4]");
|
||||
NVDR_CHECK(rast_db.size(1) == rast.size(1) && rast_db.size(2) == rast.size(2), "spatial size mismatch between inputs rast and rast_db");
|
||||
NVDR_CHECK(rast_db.size(0) == rast.size(0), "minibatch size mismatch between inputs rast, rast_db");
|
||||
}
|
||||
|
||||
// Extract input dimensions.
|
||||
p.numVertices = attr.size(p.instance_mode ? 1 : 0);
|
||||
p.numAttr = attr.size(p.instance_mode ? 2 : 1);
|
||||
p.numTriangles = tri.size(0);
|
||||
p.height = rast.size(1);
|
||||
p.width = rast.size(2);
|
||||
p.depth = rast.size(0);
|
||||
|
||||
// Ensure gradients are contiguous.
|
||||
torch::Tensor dy_ = dy.contiguous();
|
||||
torch::Tensor dda_;
|
||||
if (enable_da)
|
||||
dda_ = dda.contiguous();
|
||||
|
||||
// Set attribute pixel differential info if enabled, otherwise leave as zero.
|
||||
if (enable_da)
|
||||
set_diff_attrs(p, diff_attrs_all, diff_attrs_vec);
|
||||
else
|
||||
p.numDiffAttr = 0;
|
||||
|
||||
// Get input pointers.
|
||||
p.attr = attr.data_ptr<float>();
|
||||
p.rast = rast.data_ptr<float>();
|
||||
p.tri = tri.data_ptr<int>();
|
||||
p.dy = dy_.data_ptr<float>();
|
||||
p.rastDB = enable_da ? rast_db.data_ptr<float>() : NULL;
|
||||
p.dda = enable_da ? dda_.data_ptr<float>() : NULL;
|
||||
p.attrBC = (p.instance_mode && attr_depth < p.depth) ? 1 : 0;
|
||||
|
||||
// Allocate output tensors.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
torch::Tensor gradAttr = torch::zeros_like(attr);
|
||||
torch::Tensor gradRaster = torch::empty_like(rast);
|
||||
torch::Tensor gradRasterDB;
|
||||
if (enable_da)
|
||||
gradRasterDB = torch::empty_like(rast_db);
|
||||
|
||||
p.gradAttr = gradAttr.data_ptr<float>();
|
||||
p.gradRaster = gradRaster.data_ptr<float>();
|
||||
p.gradRasterDB = enable_da ? gradRasterDB.data_ptr<float>() : NULL;
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations.
|
||||
NVDR_CHECK(!((uintptr_t)p.rast & 15), "rast input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.rastDB & 15), "rast_db input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.dda & 7), "dda input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradRaster & 15), "grad_rast output tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradRasterDB & 15), "grad_rast_db output tensor not aligned to float4");
|
||||
|
||||
// Choose launch parameters.
|
||||
dim3 blockSize = getLaunchBlockSize(IP_GRAD_MAX_KERNEL_BLOCK_WIDTH, IP_GRAD_MAX_KERNEL_BLOCK_HEIGHT, p.width, p.height);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.depth);
|
||||
|
||||
// Launch CUDA kernel.
|
||||
void* args[] = {&p};
|
||||
void* func = enable_da ? (void*)InterpolateGradKernelDa : (void*)InterpolateGradKernel;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Return results.
|
||||
return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>(gradAttr, gradRaster, gradRasterDB);
|
||||
}
|
||||
|
||||
// Version without derivatives.
|
||||
std::tuple<torch::Tensor, torch::Tensor> interpolate_grad(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy)
|
||||
{
|
||||
std::vector<int> empty_vec;
|
||||
torch::Tensor empty_tensor;
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> result = interpolate_grad_da(attr, rast, tri, dy, empty_tensor, empty_tensor, false, empty_vec);
|
||||
return std::tuple<torch::Tensor, torch::Tensor>(std::get<0>(result), std::get<1>(result));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
@@ -0,0 +1,265 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "torch_common.inl"
|
||||
#include "torch_types.h"
|
||||
#include "../common/common.h"
|
||||
#include "../common/rasterize.h"
|
||||
#include "../common/cudaraster/CudaRaster.hpp"
|
||||
#include "../common/cudaraster/impl/Constants.hpp"
|
||||
#include <tuple>
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Kernel prototypes.
|
||||
|
||||
void RasterizeCudaFwdShaderKernel(const RasterizeCudaFwdShaderParams p);
|
||||
void RasterizeGradKernel(const RasterizeGradParams p);
|
||||
void RasterizeGradKernelDb(const RasterizeGradParams p);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Python CudaRaster state wrapper methods.
|
||||
|
||||
RasterizeCRStateWrapper::RasterizeCRStateWrapper(int cudaDeviceIdx_)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(cudaDeviceIdx_);
|
||||
cudaDeviceIdx = cudaDeviceIdx_;
|
||||
cr = new CR::CudaRaster();
|
||||
}
|
||||
|
||||
RasterizeCRStateWrapper::~RasterizeCRStateWrapper(void)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(cudaDeviceIdx);
|
||||
delete cr;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Forward op (Cuda).
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> rasterize_fwd_cuda(RasterizeCRStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple<int, int> resolution, torch::Tensor ranges, int peeling_idx)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(pos));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
CR::CudaRaster* cr = stateWrapper.cr;
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(pos, tri);
|
||||
NVDR_CHECK_CPU(ranges);
|
||||
NVDR_CHECK_CONTIGUOUS(pos, tri, ranges);
|
||||
NVDR_CHECK_F32(pos);
|
||||
NVDR_CHECK_I32(tri, ranges);
|
||||
|
||||
// Check that CudaRaster context was created for the correct GPU.
|
||||
NVDR_CHECK(pos.get_device() == stateWrapper.cudaDeviceIdx, "CudaRaster context must must reside on the same device as input tensors");
|
||||
|
||||
// Determine instance mode and check input dimensions.
|
||||
bool instance_mode = pos.sizes().size() > 2;
|
||||
if (instance_mode)
|
||||
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) > 0 && pos.size(1) > 0 && pos.size(2) == 4, "instance mode - pos must have shape [>0, >0, 4]");
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "range mode - pos must have shape [>0, 4]");
|
||||
NVDR_CHECK(ranges.sizes().size() == 2 && ranges.size(0) > 0 && ranges.size(1) == 2, "range mode - ranges must have shape [>0, 2]");
|
||||
}
|
||||
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
|
||||
// Get output shape.
|
||||
int height_out = std::get<0>(resolution);
|
||||
int width_out = std::get<1>(resolution);
|
||||
int depth = instance_mode ? pos.size(0) : ranges.size(0); // Depth of tensor, not related to depth buffering.
|
||||
NVDR_CHECK(height_out > 0 && width_out > 0, "resolution must be [>0, >0]");
|
||||
|
||||
// Round internal resolution up to tile size.
|
||||
int height = (height_out + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
int width = (width_out + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
|
||||
// Get position and triangle buffer sizes in vertices / triangles.
|
||||
int posCount = instance_mode ? pos.size(1) : pos.size(0);
|
||||
int triCount = tri.size(0);
|
||||
|
||||
// Set up CudaRaster buffers.
|
||||
const float* posPtr = pos.data_ptr<float>();
|
||||
const int32_t* rangesPtr = instance_mode ? 0 : ranges.data_ptr<int32_t>(); // This is in CPU memory.
|
||||
const int32_t* triPtr = tri.data_ptr<int32_t>();
|
||||
cr->setVertexBuffer((void*)posPtr, posCount);
|
||||
cr->setIndexBuffer((void*)triPtr, triCount);
|
||||
cr->setBufferSize(width_out, height_out, depth);
|
||||
|
||||
// Enable depth peeling?
|
||||
bool enablePeel = (peeling_idx > 0);
|
||||
cr->setRenderModeFlags(enablePeel ? CR::CudaRaster::RenderModeFlag_EnableDepthPeeling : 0); // No backface culling.
|
||||
if (enablePeel)
|
||||
cr->swapDepthAndPeel(); // Use previous depth buffer as peeling depth input.
|
||||
|
||||
// Determine viewport tiling.
|
||||
int tileCountX = (width + CR_MAXVIEWPORT_SIZE - 1) / CR_MAXVIEWPORT_SIZE;
|
||||
int tileCountY = (height + CR_MAXVIEWPORT_SIZE - 1) / CR_MAXVIEWPORT_SIZE;
|
||||
int tileSizeX = ((width + tileCountX - 1) / tileCountX + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
int tileSizeY = ((height + tileCountY - 1) / tileCountY + CR_TILE_SIZE - 1) & (-CR_TILE_SIZE);
|
||||
TORCH_CHECK(tileCountX > 0 && tileCountY > 0 && tileSizeX > 0 && tileSizeY > 0, "internal error in tile size calculation: count or size is zero");
|
||||
TORCH_CHECK(tileSizeX <= CR_MAXVIEWPORT_SIZE && tileSizeY <= CR_MAXVIEWPORT_SIZE, "internal error in tile size calculation: tile larger than allowed");
|
||||
TORCH_CHECK((tileSizeX & (CR_TILE_SIZE - 1)) == 0 && (tileSizeY & (CR_TILE_SIZE - 1)) == 0, "internal error in tile size calculation: tile not divisible by ", CR_TILE_SIZE);
|
||||
TORCH_CHECK(tileCountX * tileSizeX >= width && tileCountY * tileSizeY >= height, "internal error in tile size calculation: tiles do not cover viewport");
|
||||
|
||||
// Rasterize in tiles.
|
||||
for (int tileY = 0; tileY < tileCountY; tileY++)
|
||||
for (int tileX = 0; tileX < tileCountX; tileX++)
|
||||
{
|
||||
// Set CudaRaster viewport according to tile.
|
||||
int offsetX = tileX * tileSizeX;
|
||||
int offsetY = tileY * tileSizeY;
|
||||
int sizeX = (width_out - offsetX) < tileSizeX ? (width_out - offsetX) : tileSizeX;
|
||||
int sizeY = (height_out - offsetY) < tileSizeY ? (height_out - offsetY) : tileSizeY;
|
||||
cr->setViewport(sizeX, sizeY, offsetX, offsetY);
|
||||
|
||||
// Run all triangles in one batch. In case of error, the workload could be split into smaller batches - maybe do that in the future.
|
||||
// Only enable peeling-specific optimizations to skip first stages when image fits in one tile. Those are not valid otherwise.
|
||||
cr->deferredClear(0u);
|
||||
bool success = cr->drawTriangles(rangesPtr, enablePeel && (tileCountX == 1 && tileCountY == 1), stream);
|
||||
NVDR_CHECK(success, "subtriangle count overflow");
|
||||
}
|
||||
|
||||
// Allocate output tensors.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
torch::Tensor out = torch::empty({depth, height_out, width_out, 4}, opts);
|
||||
torch::Tensor out_db = torch::empty({depth, height_out, width_out, 4}, opts);
|
||||
|
||||
// Populate pixel shader kernel parameters.
|
||||
RasterizeCudaFwdShaderParams p;
|
||||
p.pos = posPtr;
|
||||
p.tri = triPtr;
|
||||
p.in_idx = (const int*)cr->getColorBuffer();
|
||||
p.out = out.data_ptr<float>();
|
||||
p.out_db = out_db.data_ptr<float>();
|
||||
p.numTriangles = triCount;
|
||||
p.numVertices = posCount;
|
||||
p.width_in = width;
|
||||
p.height_in = height;
|
||||
p.width_out = width_out;
|
||||
p.height_out = height_out;
|
||||
p.depth = depth;
|
||||
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
|
||||
p.xs = 2.f / (float)width_out;
|
||||
p.xo = 1.f / (float)width_out - 1.f;
|
||||
p.ys = 2.f / (float)height_out;
|
||||
p.yo = 1.f / (float)height_out - 1.f;
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations.
|
||||
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.out & 15), "out output tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.out_db & 15), "out_db output tensor not aligned to float4");
|
||||
|
||||
// Choose launch parameters.
|
||||
dim3 blockSize = getLaunchBlockSize(RAST_CUDA_FWD_SHADER_KERNEL_BLOCK_WIDTH, RAST_CUDA_FWD_SHADER_KERNEL_BLOCK_HEIGHT, p.width_out, p.height_out);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.width_out, p.height_out, p.depth);
|
||||
|
||||
// Launch CUDA kernel.
|
||||
void* args[] = {&p};
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((void*)RasterizeCudaFwdShaderKernel, gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Return.
|
||||
return std::tuple<torch::Tensor, torch::Tensor>(out, out_db);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient op.
|
||||
|
||||
torch::Tensor rasterize_grad_db(torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy, torch::Tensor ddb)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(pos));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
RasterizeGradParams p;
|
||||
bool enable_db = ddb.defined();
|
||||
|
||||
// Check inputs.
|
||||
if (enable_db)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(pos, tri, out, dy, ddb);
|
||||
NVDR_CHECK_CONTIGUOUS(pos, tri, out);
|
||||
NVDR_CHECK_F32(pos, out, dy, ddb);
|
||||
NVDR_CHECK_I32(tri);
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK_DEVICE(pos, tri, out, dy);
|
||||
NVDR_CHECK_CONTIGUOUS(pos, tri, out);
|
||||
NVDR_CHECK_F32(pos, out, dy);
|
||||
NVDR_CHECK_I32(tri);
|
||||
}
|
||||
|
||||
// Determine instance mode.
|
||||
p.instance_mode = (pos.sizes().size() > 2) ? 1 : 0;
|
||||
|
||||
// Shape is taken from the rasterizer output tensor.
|
||||
NVDR_CHECK(out.sizes().size() == 4, "tensor out must be rank-4");
|
||||
p.depth = out.size(0);
|
||||
p.height = out.size(1);
|
||||
p.width = out.size(2);
|
||||
NVDR_CHECK(p.depth > 0 && p.height > 0 && p.width > 0, "resolution must be [>0, >0, >0]");
|
||||
|
||||
// Check other shapes.
|
||||
if (p.instance_mode)
|
||||
NVDR_CHECK(pos.sizes().size() == 3 && pos.size(0) == p.depth && pos.size(1) > 0 && pos.size(2) == 4, "pos must have shape [depth, >0, 4]");
|
||||
else
|
||||
NVDR_CHECK(pos.sizes().size() == 2 && pos.size(0) > 0 && pos.size(1) == 4, "pos must have shape [>0, 4]");
|
||||
NVDR_CHECK(tri.sizes().size() == 2 && tri.size(0) > 0 && tri.size(1) == 3, "tri must have shape [>0, 3]");
|
||||
NVDR_CHECK(out.sizes().size() == 4 && out.size(0) == p.depth && out.size(1) == p.height && out.size(2) == p.width && out.size(3) == 4, "out must have shape [depth, height, width, 4]");
|
||||
NVDR_CHECK( dy.sizes().size() == 4 && dy.size(0) == p.depth && dy.size(1) == p.height && dy.size(2) == p.width && dy.size(3) == 4, "dy must have shape [depth, height, width, 4]");
|
||||
if (enable_db)
|
||||
NVDR_CHECK(ddb.sizes().size() == 4 && ddb.size(0) == p.depth && ddb.size(1) == p.height && ddb.size(2) == p.width && ddb.size(3) == 4, "ddb must have shape [depth, height, width, 4]");
|
||||
|
||||
// Ensure gradients are contiguous.
|
||||
torch::Tensor dy_ = dy.contiguous();
|
||||
torch::Tensor ddb_;
|
||||
if (enable_db)
|
||||
ddb_ = ddb.contiguous();
|
||||
|
||||
// Populate parameters.
|
||||
p.numTriangles = tri.size(0);
|
||||
p.numVertices = p.instance_mode ? pos.size(1) : pos.size(0);
|
||||
p.pos = pos.data_ptr<float>();
|
||||
p.tri = tri.data_ptr<int>();
|
||||
p.out = out.data_ptr<float>();
|
||||
p.dy = dy_.data_ptr<float>();
|
||||
p.ddb = enable_db ? ddb_.data_ptr<float>() : NULL;
|
||||
|
||||
// Set up pixel position to clip space x, y transform.
|
||||
p.xs = 2.f / (float)p.width;
|
||||
p.xo = 1.f / (float)p.width - 1.f;
|
||||
p.ys = 2.f / (float)p.height;
|
||||
p.yo = 1.f / (float)p.height - 1.f;
|
||||
|
||||
// Allocate output tensor for position gradients.
|
||||
torch::Tensor grad = torch::zeros_like(pos);
|
||||
p.grad = grad.data_ptr<float>();
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations.
|
||||
NVDR_CHECK(!((uintptr_t)p.pos & 15), "pos input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.dy & 7), "dy input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.ddb & 15), "ddb input tensor not aligned to float4");
|
||||
|
||||
// Choose launch parameters.
|
||||
dim3 blockSize = getLaunchBlockSize(RAST_GRAD_MAX_KERNEL_BLOCK_WIDTH, RAST_GRAD_MAX_KERNEL_BLOCK_HEIGHT, p.width, p.height);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.depth);
|
||||
|
||||
// Launch CUDA kernel.
|
||||
void* args[] = {&p};
|
||||
void* func = enable_db ? (void*)RasterizeGradKernelDb : (void*)RasterizeGradKernel;
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Return the gradients.
|
||||
return grad;
|
||||
}
|
||||
|
||||
// Version without derivatives.
|
||||
torch::Tensor rasterize_grad(torch::Tensor pos, torch::Tensor tri, torch::Tensor out, torch::Tensor dy)
|
||||
{
|
||||
torch::Tensor empty_tensor;
|
||||
return rasterize_grad_db(pos, tri, out, dy, empty_tensor);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+718
@@ -0,0 +1,718 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "torch_common.inl"
|
||||
#include "torch_types.h"
|
||||
#include "../common/common.h"
|
||||
#include "../common/texture.h"
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Kernel prototypes.
|
||||
|
||||
void MipBuildKernel1 (const TextureKernelParams p);
|
||||
void MipBuildKernel2 (const TextureKernelParams p);
|
||||
void MipBuildKernel4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelNearest1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelNearest2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelNearest4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinear1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinear2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinear4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapNearest1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapNearest2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapNearest4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapLinear1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapLinear2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapLinear4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeNearest1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeNearest2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeNearest4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinear1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinear2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinear4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapNearest1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapNearest2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapNearest4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapLinear1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapLinear2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapLinear4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapNearestBO1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapNearestBO2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapNearestBO4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapLinearBO1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapLinearBO2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelLinearMipmapLinearBO4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapNearestBO1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapNearestBO2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapNearestBO4 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapLinearBO1 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapLinearBO2 (const TextureKernelParams p);
|
||||
void TextureFwdKernelCubeLinearMipmapLinearBO4 (const TextureKernelParams p);
|
||||
void MipGradKernel1 (const TextureKernelParams p);
|
||||
void MipGradKernel2 (const TextureKernelParams p);
|
||||
void MipGradKernel4 (const TextureKernelParams p);
|
||||
void TextureGradKernelNearest (const TextureKernelParams p);
|
||||
void TextureGradKernelLinear (const TextureKernelParams p);
|
||||
void TextureGradKernelLinearMipmapNearest (const TextureKernelParams p);
|
||||
void TextureGradKernelLinearMipmapLinear (const TextureKernelParams p);
|
||||
void TextureGradKernelCubeNearest (const TextureKernelParams p);
|
||||
void TextureGradKernelCubeLinear (const TextureKernelParams p);
|
||||
void TextureGradKernelCubeLinearMipmapNearest (const TextureKernelParams p);
|
||||
void TextureGradKernelCubeLinearMipmapLinear (const TextureKernelParams p);
|
||||
void TextureGradKernelLinearMipmapNearestBO (const TextureKernelParams p);
|
||||
void TextureGradKernelLinearMipmapLinearBO (const TextureKernelParams p);
|
||||
void TextureGradKernelCubeLinearMipmapNearestBO (const TextureKernelParams p);
|
||||
void TextureGradKernelCubeLinearMipmapLinearBO (const TextureKernelParams p);
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Modeselektor.
|
||||
|
||||
static void set_modes(TextureKernelParams& p, int filter_mode, int boundary_mode, int max_mip_level)
|
||||
{
|
||||
// Mip and filter modes.
|
||||
p.filterMode = filter_mode;
|
||||
NVDR_CHECK(p.filterMode >= 0 && p.filterMode < TEX_MODE_COUNT, "filter_mode unsupported");
|
||||
p.enableMip = (p.filterMode == TEX_MODE_LINEAR_MIPMAP_NEAREST || p.filterMode == TEX_MODE_LINEAR_MIPMAP_LINEAR);
|
||||
|
||||
// Mip level clamp.
|
||||
if (p.enableMip)
|
||||
{
|
||||
p.mipLevelLimit = max_mip_level;
|
||||
NVDR_CHECK(p.mipLevelLimit >= -1, "invalid max_mip_level");
|
||||
}
|
||||
|
||||
// Boundary mode.
|
||||
p.boundaryMode = boundary_mode;
|
||||
NVDR_CHECK(p.boundaryMode >= 0 && p.boundaryMode < TEX_BOUNDARY_MODE_COUNT, "boundary_mode unsupported");
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Mipmap construction.
|
||||
|
||||
TextureMipWrapper texture_construct_mip(torch::Tensor tex, int max_mip_level, bool cube_mode)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(tex));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
TextureKernelParams p = {}; // Initialize all fields to zero.
|
||||
p.mipLevelLimit = max_mip_level;
|
||||
p.boundaryMode = cube_mode ? TEX_BOUNDARY_MODE_CUBE : TEX_BOUNDARY_MODE_WRAP;
|
||||
NVDR_CHECK(p.mipLevelLimit >= -1, "invalid max_mip_level");
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(tex);
|
||||
NVDR_CHECK_CONTIGUOUS(tex);
|
||||
NVDR_CHECK_F32(tex);
|
||||
|
||||
// Populate parameters and sanity check tex shape.
|
||||
if (!cube_mode)
|
||||
{
|
||||
NVDR_CHECK(tex.sizes().size() == 4 && tex.size(0) > 0 && tex.size(1) > 0 && tex.size(2) > 0 && tex.size(3) > 0, "tex must have shape[>0, >0, >0, >0]");
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(tex.sizes().size() == 5 && tex.size(0) > 0 && tex.size(1) == 6 && tex.size(2) > 0 && tex.size(3) > 0 && tex.size(4) > 0, "tex must have shape[>0, 6, >0, >0, >0] in cube map mode");
|
||||
NVDR_CHECK(tex.size(2) == tex.size(3), "texture shape must be square in cube map mode");
|
||||
}
|
||||
p.texDepth = tex.size(0);
|
||||
p.texHeight = tex.size(cube_mode ? 2 : 1);
|
||||
p.texWidth = tex.size(cube_mode ? 3 : 2);
|
||||
p.channels = tex.size(cube_mode ? 4 : 3);
|
||||
|
||||
// Set texture pointer.
|
||||
p.tex[0] = tex.data_ptr<float>();
|
||||
|
||||
// Generate mip offsets and calculate total size.
|
||||
int mipOffsets[TEX_MAX_MIP_LEVEL];
|
||||
int mipTotal = calculateMipInfo(p, mipOffsets);
|
||||
|
||||
// Allocate and set mip tensor.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
torch::Tensor mip = torch::empty({mipTotal}, opts);
|
||||
float* pmip = mip.data_ptr<float>();
|
||||
for (int i=1; i <= p.mipLevelMax; i++)
|
||||
p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels.
|
||||
|
||||
// Choose kernel variants based on channel count.
|
||||
void* args[] = {&p};
|
||||
int channel_div_idx = 0;
|
||||
if (!(p.channels & 3))
|
||||
channel_div_idx = 2; // Channel count divisible by 4.
|
||||
else if (!(p.channels & 1))
|
||||
channel_div_idx = 1; // Channel count divisible by 2.
|
||||
|
||||
// Build mip levels.
|
||||
for (int i=1; i <= p.mipLevelMax; i++)
|
||||
{
|
||||
int2 ms = mipLevelSize(p, i);
|
||||
int3 sz = make_int3(ms.x, ms.y, p.texDepth);
|
||||
dim3 blockSize = getLaunchBlockSize(TEX_FWD_MAX_MIP_KERNEL_BLOCK_WIDTH, TEX_FWD_MAX_MIP_KERNEL_BLOCK_HEIGHT, sz.x, sz.y);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, sz.x, sz.y, sz.z * (cube_mode ? 6 : 1));
|
||||
p.mipLevelOut = i;
|
||||
|
||||
void* build_func_tbl[3] = { (void*)MipBuildKernel1, (void*)MipBuildKernel2, (void*)MipBuildKernel4 };
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(build_func_tbl[channel_div_idx], gridSize, blockSize, args, 0, stream));
|
||||
}
|
||||
|
||||
// Return the mip tensor in a wrapper.
|
||||
TextureMipWrapper mip_wrapper;
|
||||
mip_wrapper.mip = mip;
|
||||
mip_wrapper.max_mip_level = max_mip_level;
|
||||
mip_wrapper.texture_size = tex.sizes().vec();
|
||||
mip_wrapper.cube_mode = cube_mode;
|
||||
return mip_wrapper;
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Forward op.
|
||||
|
||||
torch::Tensor texture_fwd_mip(torch::Tensor tex, torch::Tensor uv, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(tex));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
TextureKernelParams p = {}; // Initialize all fields to zero.
|
||||
bool has_mip_stack = (mip_stack.size() > 0);
|
||||
torch::Tensor& mip_w = mip_wrapper.mip; // Unwrap.
|
||||
int max_mip_level = has_mip_stack ? mip_stack.size() : mip_wrapper.max_mip_level;
|
||||
set_modes(p, filter_mode, boundary_mode, max_mip_level);
|
||||
|
||||
// See if we have these tensors or not.
|
||||
bool has_uv_da = uv_da.defined() && uv_da.nbytes();
|
||||
bool has_mip_level_bias = mip_level_bias.defined() && mip_level_bias.nbytes();
|
||||
|
||||
if (p.enableMip)
|
||||
{
|
||||
NVDR_CHECK(has_uv_da || has_mip_level_bias, "mipmapping filter mode requires uv_da and/or mip_level_bias input");
|
||||
NVDR_CHECK(has_mip_stack || mip_w.defined(), "mipmapping filter mode requires mip wrapper or mip stack input");
|
||||
}
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(tex, uv);
|
||||
NVDR_CHECK_CONTIGUOUS(tex, uv);
|
||||
NVDR_CHECK_F32(tex, uv);
|
||||
if (p.enableMip)
|
||||
{
|
||||
if (has_mip_stack)
|
||||
{
|
||||
TORCH_CHECK(at::cuda::check_device(mip_stack), __func__, "(): Mip stack inputs must reside on the correct GPU device");
|
||||
nvdr_check_contiguous(mip_stack, __func__, "(): Mip stack inputs must be contiguous tensors");
|
||||
nvdr_check_f32(mip_stack, __func__, "(): Mip stack inputs must be float32 tensors");
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK_DEVICE(mip_w);
|
||||
NVDR_CHECK_CONTIGUOUS(mip_w);
|
||||
NVDR_CHECK_F32(mip_w);
|
||||
}
|
||||
if (has_uv_da)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(uv_da);
|
||||
NVDR_CHECK_CONTIGUOUS(uv_da);
|
||||
NVDR_CHECK_F32(uv_da);
|
||||
}
|
||||
if (has_mip_level_bias)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(mip_level_bias);
|
||||
NVDR_CHECK_CONTIGUOUS(mip_level_bias);
|
||||
NVDR_CHECK_F32(mip_level_bias);
|
||||
}
|
||||
}
|
||||
|
||||
// Sanity checks and state setters.
|
||||
bool cube_mode = (boundary_mode == TEX_BOUNDARY_MODE_CUBE);
|
||||
if (!cube_mode)
|
||||
{
|
||||
NVDR_CHECK(tex.sizes().size() == 4 && tex.size(0) > 0 && tex.size(1) > 0 && tex.size(2) > 0 && tex.size(3) > 0, "tex must have shape[>0, >0, >0, >0]");
|
||||
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 2, "uv must have shape [>0, >0, >0, 2]");
|
||||
p.texHeight = tex.size(1);
|
||||
p.texWidth = tex.size(2);
|
||||
p.channels = tex.size(3);
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(tex.sizes().size() == 5 && tex.size(0) > 0 && tex.size(1) == 6 && tex.size(2) > 0 && tex.size(3) > 0 && tex.size(4) > 0, "tex must have shape[>0, 6, >0, >0, >0] in cube map mode");
|
||||
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 3, "uv must have shape [>0, >0, >0, 3] in cube map mode");
|
||||
NVDR_CHECK(tex.size(2) == tex.size(3), "texture shape must be square in cube map mode");
|
||||
p.texHeight = tex.size(2);
|
||||
p.texWidth = tex.size(3);
|
||||
p.channels = tex.size(4);
|
||||
}
|
||||
NVDR_CHECK(tex.size(0) == 1 || tex.size(0) == uv.size(0), "minibatch size mismatch between inputs tex, uv");
|
||||
NVDR_CHECK(p.texWidth <= (1 << TEX_MAX_MIP_LEVEL) && p.texHeight <= (1 << TEX_MAX_MIP_LEVEL), "texture size too large");
|
||||
p.n = uv.size(0);
|
||||
p.imgHeight = uv.size(1);
|
||||
p.imgWidth = uv.size(2);
|
||||
p.texDepth = tex.size(0);
|
||||
if (p.enableMip)
|
||||
{
|
||||
if (has_uv_da)
|
||||
{
|
||||
if (!cube_mode)
|
||||
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 4, "uv_da must have shape [minibatch_size, height, width, 4]");
|
||||
else
|
||||
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 6, "uv_da must have shape [minibatch_size, height, width, 6] in cube map mode");
|
||||
}
|
||||
if (has_mip_level_bias)
|
||||
NVDR_CHECK(mip_level_bias.sizes().size() == 3 && mip_level_bias.size(0) == p.n && mip_level_bias.size(1) == p.imgHeight && mip_level_bias.size(2) == p.imgWidth, "mip_level_bias must have shape [minibatch_size, height, width]");
|
||||
}
|
||||
|
||||
// Get input pointers.
|
||||
p.tex[0] = tex.data_ptr<float>();
|
||||
p.uv = uv.data_ptr<float>();
|
||||
p.uvDA = (p.enableMip && has_uv_da) ? uv_da.data_ptr<float>() : NULL;
|
||||
p.mipLevelBias = (p.enableMip && has_mip_level_bias) ? mip_level_bias.data_ptr<float>() : NULL;
|
||||
|
||||
// Allocate output tensor.
|
||||
torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
torch::Tensor out = torch::empty({p.n, p.imgHeight, p.imgWidth, p.channels}, opts);
|
||||
p.out = out.data_ptr<float>();
|
||||
|
||||
// Choose kernel variants based on channel count.
|
||||
void* args[] = {&p};
|
||||
int channel_div_idx = 0;
|
||||
if (!(p.channels & 3))
|
||||
channel_div_idx = 2; // Channel count divisible by 4.
|
||||
else if (!(p.channels & 1))
|
||||
channel_div_idx = 1; // Channel count divisible by 2.
|
||||
|
||||
// Mip-related setup.
|
||||
float* pmip = 0;
|
||||
if (p.enableMip)
|
||||
{
|
||||
if (has_mip_stack)
|
||||
{
|
||||
// Custom mip stack supplied. Check that sizes match and assign.
|
||||
p.mipLevelMax = max_mip_level;
|
||||
for (int i=1; i <= p.mipLevelMax; i++)
|
||||
{
|
||||
torch::Tensor& t = mip_stack[i-1];
|
||||
int2 sz = mipLevelSize(p, i);
|
||||
if (!cube_mode)
|
||||
NVDR_CHECK(t.sizes().size() == 4 && t.size(0) == tex.size(0) && t.size(1) == sz.y && t.size(2) == sz.x && t.size(3) == p.channels, "mip level size mismatch in custom mip stack");
|
||||
else
|
||||
NVDR_CHECK(t.sizes().size() == 5 && t.size(0) == tex.size(0) && t.size(1) == 6 && t.size(2) == sz.y && t.size(3) == sz.x && t.size(4) == p.channels, "mip level size mismatch in mip stack");
|
||||
if (sz.x == 1 && sz.y == 1)
|
||||
NVDR_CHECK(i == p.mipLevelMax, "mip level size mismatch in mip stack");
|
||||
p.tex[i] = t.data_ptr<float>();
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Generate mip offsets, check mipmap size, and set mip data pointer.
|
||||
int mipOffsets[TEX_MAX_MIP_LEVEL];
|
||||
int mipTotal = calculateMipInfo(p, mipOffsets);
|
||||
NVDR_CHECK(tex.sizes() == mip_wrapper.texture_size && cube_mode == mip_wrapper.cube_mode, "mip does not match texture size");
|
||||
NVDR_CHECK(mip_w.sizes().size() == 1 && mip_w.size(0) == mipTotal, "wrapped mip tensor size mismatch");
|
||||
pmip = mip_w.data_ptr<float>();
|
||||
for (int i=1; i <= p.mipLevelMax; i++)
|
||||
p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels.
|
||||
}
|
||||
}
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations. Unused pointers are zero so always aligned.
|
||||
if (!cube_mode)
|
||||
NVDR_CHECK(!((uintptr_t)p.uv & 7), "uv input tensor not aligned to float2");
|
||||
if ((p.channels & 3) == 0)
|
||||
{
|
||||
for (int i=0; i <= p.mipLevelMax; i++)
|
||||
NVDR_CHECK(!((uintptr_t)p.tex[i] & 15), "tex or mip input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.out & 15), "out output tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)pmip & 15), "mip input tensor not aligned to float4");
|
||||
}
|
||||
if ((p.channels & 1) == 0)
|
||||
{
|
||||
for (int i=0; i <= p.mipLevelMax; i++)
|
||||
NVDR_CHECK(!((uintptr_t)p.tex[i] & 7), "tex or mip input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.out & 7), "out output tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)pmip & 7), "mip input tensor not aligned to float2");
|
||||
}
|
||||
if (!cube_mode)
|
||||
NVDR_CHECK(!((uintptr_t)p.uvDA & 15), "uv_da input tensor not aligned to float4");
|
||||
else
|
||||
NVDR_CHECK(!((uintptr_t)p.uvDA & 7), "uv_da input tensor not aligned to float2");
|
||||
|
||||
// Choose launch parameters for texture lookup kernel.
|
||||
dim3 blockSize = getLaunchBlockSize(TEX_FWD_MAX_KERNEL_BLOCK_WIDTH, TEX_FWD_MAX_KERNEL_BLOCK_HEIGHT, p.imgWidth, p.imgHeight);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.imgWidth, p.imgHeight, p.n);
|
||||
|
||||
// Choose kernel based on filter mode, cube mode, bias-only mode, and datatype.
|
||||
void* func_tbl[TEX_MODE_COUNT * 2 * 2 * 3] = {
|
||||
(void*)TextureFwdKernelNearest1,
|
||||
(void*)TextureFwdKernelNearest2,
|
||||
(void*)TextureFwdKernelNearest4,
|
||||
(void*)TextureFwdKernelLinear1,
|
||||
(void*)TextureFwdKernelLinear2,
|
||||
(void*)TextureFwdKernelLinear4,
|
||||
(void*)TextureFwdKernelLinearMipmapNearest1,
|
||||
(void*)TextureFwdKernelLinearMipmapNearest2,
|
||||
(void*)TextureFwdKernelLinearMipmapNearest4,
|
||||
(void*)TextureFwdKernelLinearMipmapLinear1,
|
||||
(void*)TextureFwdKernelLinearMipmapLinear2,
|
||||
(void*)TextureFwdKernelLinearMipmapLinear4,
|
||||
(void*)TextureFwdKernelCubeNearest1,
|
||||
(void*)TextureFwdKernelCubeNearest2,
|
||||
(void*)TextureFwdKernelCubeNearest4,
|
||||
(void*)TextureFwdKernelCubeLinear1,
|
||||
(void*)TextureFwdKernelCubeLinear2,
|
||||
(void*)TextureFwdKernelCubeLinear4,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapNearest1,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapNearest2,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapNearest4,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapLinear1,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapLinear2,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapLinear4,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
(void*)TextureFwdKernelLinearMipmapNearestBO1,
|
||||
(void*)TextureFwdKernelLinearMipmapNearestBO2,
|
||||
(void*)TextureFwdKernelLinearMipmapNearestBO4,
|
||||
(void*)TextureFwdKernelLinearMipmapLinearBO1,
|
||||
(void*)TextureFwdKernelLinearMipmapLinearBO2,
|
||||
(void*)TextureFwdKernelLinearMipmapLinearBO4,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapNearestBO1,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapNearestBO2,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapNearestBO4,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapLinearBO1,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapLinearBO2,
|
||||
(void*)TextureFwdKernelCubeLinearMipmapLinearBO4,
|
||||
};
|
||||
|
||||
// Function index.
|
||||
int func_idx = p.filterMode;
|
||||
if (cube_mode)
|
||||
func_idx += TEX_MODE_COUNT; // Cube variant.
|
||||
if (p.enableMip && !has_uv_da)
|
||||
func_idx += TEX_MODE_COUNT * 2; // Bias-only variant.
|
||||
func_idx = func_idx * 3 + channel_div_idx; // Choose vector size.
|
||||
|
||||
// Launch kernel.
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func_tbl[func_idx], gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Return output tensor.
|
||||
return out;
|
||||
}
|
||||
|
||||
// Version without mipmaps.
|
||||
torch::Tensor texture_fwd(torch::Tensor tex, torch::Tensor uv, int filter_mode, int boundary_mode)
|
||||
{
|
||||
torch::Tensor empty_tensor;
|
||||
std::vector<torch::Tensor> empty_vector;
|
||||
return texture_fwd_mip(tex, uv, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode);
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Gradient op.
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > texture_grad_linear_mipmap_linear(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode)
|
||||
{
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(tex));
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
TextureKernelParams p = {}; // Initialize all fields to zero.
|
||||
bool has_mip_stack = (mip_stack.size() > 0);
|
||||
torch::Tensor& mip_w = mip_wrapper.mip; // Unwrap.
|
||||
int max_mip_level = has_mip_stack ? mip_stack.size() : mip_wrapper.max_mip_level;
|
||||
set_modes(p, filter_mode, boundary_mode, max_mip_level);
|
||||
|
||||
// See if we have these tensors or not.
|
||||
bool has_uv_da = uv_da.defined() && uv_da.nbytes();
|
||||
bool has_mip_level_bias = mip_level_bias.defined() && mip_level_bias.nbytes();
|
||||
|
||||
if (p.enableMip)
|
||||
{
|
||||
NVDR_CHECK(has_uv_da || has_mip_level_bias, "mipmapping filter mode requires uv_da and/or mip_level_bias input");
|
||||
NVDR_CHECK(has_mip_stack || mip_w.defined(), "mipmapping filter mode requires mip wrapper or mip stack input");
|
||||
}
|
||||
|
||||
// Check inputs.
|
||||
NVDR_CHECK_DEVICE(tex, uv);
|
||||
NVDR_CHECK_CONTIGUOUS(tex, uv);
|
||||
NVDR_CHECK_F32(tex, uv);
|
||||
if (p.enableMip)
|
||||
{
|
||||
if (has_mip_stack)
|
||||
{
|
||||
TORCH_CHECK(at::cuda::check_device(mip_stack), __func__, "(): Mip stack inputs must reside on the correct GPU device");
|
||||
nvdr_check_contiguous(mip_stack, __func__, "(): Mip stack inputs must be contiguous tensors");
|
||||
nvdr_check_f32(mip_stack, __func__, "(): Mip stack inputs must be float32 tensors");
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK_DEVICE(mip_w);
|
||||
NVDR_CHECK_CONTIGUOUS(mip_w);
|
||||
NVDR_CHECK_F32(mip_w);
|
||||
}
|
||||
if (has_uv_da)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(uv_da);
|
||||
NVDR_CHECK_CONTIGUOUS(uv_da);
|
||||
NVDR_CHECK_F32(uv_da);
|
||||
}
|
||||
if (has_mip_level_bias)
|
||||
{
|
||||
NVDR_CHECK_DEVICE(mip_level_bias);
|
||||
NVDR_CHECK_CONTIGUOUS(mip_level_bias);
|
||||
NVDR_CHECK_F32(mip_level_bias);
|
||||
}
|
||||
}
|
||||
|
||||
// Sanity checks and state setters.
|
||||
bool cube_mode = (boundary_mode == TEX_BOUNDARY_MODE_CUBE);
|
||||
if (!cube_mode)
|
||||
{
|
||||
NVDR_CHECK(tex.sizes().size() == 4 && tex.size(0) > 0 && tex.size(1) > 0 && tex.size(2) > 0 && tex.size(3) > 0, "tex must have shape[>0, >0, >0, >0]");
|
||||
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 2, "uv must have shape [>0, >0, >0, 2]");
|
||||
p.texHeight = tex.size(1);
|
||||
p.texWidth = tex.size(2);
|
||||
p.channels = tex.size(3);
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(tex.sizes().size() == 5 && tex.size(0) > 0 && tex.size(1) == 6 && tex.size(2) > 0 && tex.size(3) > 0 && tex.size(4) > 0, "tex must have shape[>0, 6, >0, >0, >0] in cube map mode");
|
||||
NVDR_CHECK(uv.sizes().size() == 4 && uv.size(0) > 0 && uv.size(1) > 0 && uv.size(2) > 0 && uv.size(3) == 3, "uv must have shape [>0, >0, >0, 3] in cube map mode");
|
||||
NVDR_CHECK(tex.size(2) == tex.size(3), "texture shape must be square in cube map mode");
|
||||
p.texHeight = tex.size(2);
|
||||
p.texWidth = tex.size(3);
|
||||
p.channels = tex.size(4);
|
||||
}
|
||||
NVDR_CHECK(tex.size(0) == 1 || tex.size(0) == uv.size(0), "minibatch size mismatch between inputs tex, uv");
|
||||
NVDR_CHECK(p.texWidth <= (1 << TEX_MAX_MIP_LEVEL) && p.texHeight <= (1 << TEX_MAX_MIP_LEVEL), "texture size too large");
|
||||
p.n = uv.size(0);
|
||||
p.imgHeight = uv.size(1);
|
||||
p.imgWidth = uv.size(2);
|
||||
p.texDepth = tex.size(0);
|
||||
if (p.enableMip)
|
||||
{
|
||||
if (has_uv_da)
|
||||
{
|
||||
if (!cube_mode)
|
||||
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 4, "uv_da must have shape [minibatch_size, height, width, 4]");
|
||||
else
|
||||
NVDR_CHECK(uv_da.sizes().size() == 4 && uv_da.size(0) == p.n && uv_da.size(1) == p.imgHeight && uv_da.size(2) == p.imgWidth && uv_da.size(3) == 6, "uv_da must have shape [minibatch_size, height, width, 6] in cube map mode");
|
||||
}
|
||||
if (has_mip_level_bias)
|
||||
NVDR_CHECK(mip_level_bias.sizes().size() == 3 && mip_level_bias.size(0) == p.n && mip_level_bias.size(1) == p.imgHeight && mip_level_bias.size(2) == p.imgWidth, "mip_level_bias must have shape [minibatch_size, height, width]");
|
||||
}
|
||||
NVDR_CHECK(dy.sizes().size() == 4 && dy.size(0) == p.n && dy.size(1) == p.imgHeight && dy.size(2) == p.imgWidth && dy.size(3) == p.channels, "dy must have shape [minibatch_size, height, width, channels]");
|
||||
|
||||
// Get contiguous version of dy.
|
||||
torch::Tensor dy_ = dy.contiguous();
|
||||
|
||||
// Get input pointers.
|
||||
p.tex[0] = tex.data_ptr<float>();
|
||||
p.uv = uv.data_ptr<float>();
|
||||
p.dy = dy_.data_ptr<float>();
|
||||
p.uvDA = (p.enableMip && has_uv_da) ? uv_da.data_ptr<float>() : NULL;
|
||||
p.mipLevelBias = (p.enableMip && has_mip_level_bias) ? mip_level_bias.data_ptr<float>() : NULL;
|
||||
|
||||
// Allocate output tensor for tex gradient.
|
||||
torch::Tensor grad_tex = torch::zeros_like(tex);
|
||||
p.gradTex[0] = grad_tex.data_ptr<float>();
|
||||
|
||||
// Allocate output tensor for uv gradient.
|
||||
torch::Tensor grad_uv;
|
||||
torch::Tensor grad_uv_da;
|
||||
torch::Tensor grad_mip_level_bias;
|
||||
if (p.filterMode != TEX_MODE_NEAREST)
|
||||
{
|
||||
grad_uv = torch::empty_like(uv);
|
||||
p.gradUV = grad_uv.data_ptr<float>();
|
||||
|
||||
// Gradients for things affecting mip level.
|
||||
if (p.filterMode == TEX_MODE_LINEAR_MIPMAP_LINEAR)
|
||||
{
|
||||
// Allocate output tensor for uv_da gradient.
|
||||
if (has_uv_da)
|
||||
{
|
||||
grad_uv_da = torch::empty_like(uv_da);
|
||||
p.gradUVDA = grad_uv_da.data_ptr<float>();
|
||||
}
|
||||
|
||||
// Allocate output tensor for mip_level_bias gradient.
|
||||
if (has_mip_level_bias)
|
||||
{
|
||||
grad_mip_level_bias = torch::empty_like(mip_level_bias);
|
||||
p.gradMipLevelBias = grad_mip_level_bias.data_ptr<float>();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Choose kernel variants based on channel count.
|
||||
int channel_div_idx = 0;
|
||||
if (!(p.channels & 3))
|
||||
channel_div_idx = 2; // Channel count divisible by 4.
|
||||
else if (!(p.channels & 1))
|
||||
channel_div_idx = 1; // Channel count divisible by 2.
|
||||
|
||||
// Mip-related setup.
|
||||
torch::Tensor grad_mip;
|
||||
std::vector<torch::Tensor> grad_mip_stack;
|
||||
float* pmip = 0;
|
||||
float* pgradMip = 0;
|
||||
if (p.enableMip)
|
||||
{
|
||||
if (has_mip_stack)
|
||||
{
|
||||
// Custom mip stack supplied. Check that sizes match, assign, construct gradient tensors.
|
||||
p.mipLevelMax = max_mip_level;
|
||||
for (int i=1; i <= p.mipLevelMax; i++)
|
||||
{
|
||||
torch::Tensor& t = mip_stack[i-1];
|
||||
int2 sz = mipLevelSize(p, i);
|
||||
if (!cube_mode)
|
||||
NVDR_CHECK(t.sizes().size() == 4 && t.size(0) == tex.size(0) && t.size(1) == sz.y && t.size(2) == sz.x && t.size(3) == p.channels, "mip level size mismatch in mip stack");
|
||||
else
|
||||
NVDR_CHECK(t.sizes().size() == 5 && t.size(0) == tex.size(0) && t.size(1) == 6 && t.size(2) == sz.y && t.size(3) == sz.x && t.size(4) == p.channels, "mip level size mismatch in mip stack");
|
||||
if (sz.x == 1 && sz.y == 1)
|
||||
NVDR_CHECK(i == p.mipLevelMax, "mip level size mismatch in mip stack");
|
||||
|
||||
torch::Tensor g = torch::zeros_like(t);
|
||||
grad_mip_stack.push_back(g);
|
||||
|
||||
p.tex[i] = t.data_ptr<float>();
|
||||
p.gradTex[i] = g.data_ptr<float>();
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Generate mip offsets and get space for temporary mip gradients.
|
||||
int mipOffsets[TEX_MAX_MIP_LEVEL];
|
||||
int mipTotal = calculateMipInfo(p, mipOffsets);
|
||||
NVDR_CHECK(tex.sizes() == mip_wrapper.texture_size && cube_mode == mip_wrapper.cube_mode, "mip does not match texture size");
|
||||
NVDR_CHECK(mip_w.sizes().size() == 1 && mip_w.size(0) == mipTotal, "mip tensor size mismatch");
|
||||
grad_mip = torch::zeros_like(mip_w);
|
||||
pmip = (float*)mip_w.data_ptr<float>();
|
||||
pgradMip = grad_mip.data_ptr<float>();
|
||||
for (int i=1; i <= p.mipLevelMax; i++)
|
||||
{
|
||||
p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels.
|
||||
p.gradTex[i] = pgradMip + mipOffsets[i]; // Pointers to mip gradients.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Verify that buffers are aligned to allow float2/float4 operations. Unused pointers are zero so always aligned.
|
||||
if (!cube_mode)
|
||||
{
|
||||
NVDR_CHECK(!((uintptr_t)p.uv & 7), "uv input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradUV & 7), "grad_uv output tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.uvDA & 15), "uv_da input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradUVDA & 15), "grad_uv_da output tensor not aligned to float4");
|
||||
}
|
||||
else
|
||||
{
|
||||
NVDR_CHECK(!((uintptr_t)p.uvDA & 7), "uv_da input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradUVDA & 7), "grad_uv_da output tensor not aligned to float2");
|
||||
}
|
||||
if ((p.channels & 3) == 0)
|
||||
{
|
||||
for (int i=0; i <= p.mipLevelMax; i++)
|
||||
{
|
||||
NVDR_CHECK(!((uintptr_t)p.tex[i] & 15), "tex or mip input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradTex[i] & 15), "grad_tex output tensor not aligned to float4");
|
||||
}
|
||||
NVDR_CHECK(!((uintptr_t)p.dy & 15), "dy input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)pmip & 15), "mip input tensor not aligned to float4");
|
||||
NVDR_CHECK(!((uintptr_t)pgradMip & 15), "internal mip gradient tensor not aligned to float4");
|
||||
}
|
||||
if ((p.channels & 1) == 0)
|
||||
{
|
||||
for (int i=0; i <= p.mipLevelMax; i++)
|
||||
{
|
||||
NVDR_CHECK(!((uintptr_t)p.tex[i] & 7), "tex or mip input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)p.gradTex[i] & 7), "grad_tex output tensor not aligned to float2");
|
||||
}
|
||||
NVDR_CHECK(!((uintptr_t)p.dy & 7), "dy output tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)pmip & 7), "mip input tensor not aligned to float2");
|
||||
NVDR_CHECK(!((uintptr_t)pgradMip & 7), "internal mip gradient tensor not aligned to float2");
|
||||
}
|
||||
|
||||
// Choose launch parameters for main gradient kernel.
|
||||
void* args[] = {&p};
|
||||
dim3 blockSize = getLaunchBlockSize(TEX_GRAD_MAX_KERNEL_BLOCK_WIDTH, TEX_GRAD_MAX_KERNEL_BLOCK_HEIGHT, p.imgWidth, p.imgHeight);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.imgWidth, p.imgHeight, p.n);
|
||||
|
||||
void* func_tbl[TEX_MODE_COUNT * 2 * 2] = {
|
||||
(void*)TextureGradKernelNearest,
|
||||
(void*)TextureGradKernelLinear,
|
||||
(void*)TextureGradKernelLinearMipmapNearest,
|
||||
(void*)TextureGradKernelLinearMipmapLinear,
|
||||
(void*)TextureGradKernelCubeNearest,
|
||||
(void*)TextureGradKernelCubeLinear,
|
||||
(void*)TextureGradKernelCubeLinearMipmapNearest,
|
||||
(void*)TextureGradKernelCubeLinearMipmapLinear,
|
||||
NULL,
|
||||
NULL,
|
||||
(void*)TextureGradKernelLinearMipmapNearestBO,
|
||||
(void*)TextureGradKernelLinearMipmapLinearBO,
|
||||
NULL,
|
||||
NULL,
|
||||
(void*)TextureGradKernelCubeLinearMipmapNearestBO,
|
||||
(void*)TextureGradKernelCubeLinearMipmapLinearBO,
|
||||
};
|
||||
|
||||
// Function index.
|
||||
int func_idx = p.filterMode;
|
||||
if (cube_mode)
|
||||
func_idx += TEX_MODE_COUNT; // Cube variant.
|
||||
if (p.enableMip && !has_uv_da)
|
||||
func_idx += TEX_MODE_COUNT * 2; // Bias-only variant.
|
||||
|
||||
// Launch main gradient kernel.
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(func_tbl[func_idx], gridSize, blockSize, args, 0, stream));
|
||||
|
||||
// Launch kernel to pull gradients from mip levels. Don't do this if mip stack was supplied - individual level gradients are already there.
|
||||
if (p.enableMip && !has_mip_stack)
|
||||
{
|
||||
dim3 blockSize = getLaunchBlockSize(TEX_GRAD_MAX_MIP_KERNEL_BLOCK_WIDTH, TEX_GRAD_MAX_MIP_KERNEL_BLOCK_HEIGHT, p.texWidth, p.texHeight);
|
||||
dim3 gridSize = getLaunchGridSize(blockSize, p.texWidth, p.texHeight, p.texDepth * (cube_mode ? 6 : 1));
|
||||
int sharedBytes = blockSize.x * blockSize.y * p.channels * sizeof(float);
|
||||
|
||||
void* mip_grad_func_tbl[3] = { (void*)MipGradKernel1, (void*)MipGradKernel2, (void*)MipGradKernel4 };
|
||||
NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel(mip_grad_func_tbl[channel_div_idx], gridSize, blockSize, args, sharedBytes, stream));
|
||||
}
|
||||
|
||||
// Return output tensors.
|
||||
return std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >(grad_tex, grad_uv, grad_uv_da, grad_mip_level_bias, grad_mip_stack);
|
||||
}
|
||||
|
||||
// Version for nearest filter mode.
|
||||
torch::Tensor texture_grad_nearest(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode)
|
||||
{
|
||||
torch::Tensor empty_tensor;
|
||||
std::vector<torch::Tensor> empty_vector;
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > result = texture_grad_linear_mipmap_linear(tex, uv, dy, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode);
|
||||
return std::get<0>(result);
|
||||
}
|
||||
|
||||
// Version for linear filter mode.
|
||||
std::tuple<torch::Tensor, torch::Tensor> texture_grad_linear(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode)
|
||||
{
|
||||
torch::Tensor empty_tensor;
|
||||
std::vector<torch::Tensor> empty_vector;
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > result = texture_grad_linear_mipmap_linear(tex, uv, dy, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode);
|
||||
return std::tuple<torch::Tensor, torch::Tensor>(std::get<0>(result), std::get<1>(result));
|
||||
}
|
||||
|
||||
// Version for linear-mipmap-nearest mode.
|
||||
std::tuple<torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > texture_grad_linear_mipmap_nearest(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, torch::Tensor uv_da, torch::Tensor mip_level_bias, TextureMipWrapper mip_wrapper, std::vector<torch::Tensor> mip_stack, int filter_mode, int boundary_mode)
|
||||
{
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::vector<torch::Tensor> > result = texture_grad_linear_mipmap_linear(tex, uv, dy, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode, boundary_mode);
|
||||
return std::tuple<torch::Tensor, torch::Tensor, std::vector<torch::Tensor> >(std::get<0>(result), std::get<1>(result), std::get<4>(result));
|
||||
}
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
// Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
//
|
||||
// NVIDIA CORPORATION and its licensors retain all intellectual property
|
||||
// and proprietary rights in and to this software, related documentation
|
||||
// and any modifications thereto. Any use, reproduction, disclosure or
|
||||
// distribution of this software and related documentation without an express
|
||||
// license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
#include "torch_common.inl"
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Python CudaRaster state wrapper.
|
||||
|
||||
namespace CR { class CudaRaster; }
|
||||
class RasterizeCRStateWrapper
|
||||
{
|
||||
public:
|
||||
RasterizeCRStateWrapper (int cudaDeviceIdx);
|
||||
~RasterizeCRStateWrapper (void);
|
||||
|
||||
CR::CudaRaster* cr;
|
||||
int cudaDeviceIdx;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Mipmap wrapper to prevent intrusion from Python side.
|
||||
|
||||
class TextureMipWrapper
|
||||
{
|
||||
public:
|
||||
torch::Tensor mip;
|
||||
int max_mip_level;
|
||||
std::vector<int64_t> texture_size; // For error checking.
|
||||
bool cube_mode; // For error checking.
|
||||
};
|
||||
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
// Antialias topology hash wrapper to prevent intrusion from Python side.
|
||||
|
||||
class TopologyHashWrapper
|
||||
{
|
||||
public:
|
||||
torch::Tensor ev_hash;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
Reference in New Issue
Block a user