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
Vendored
+20
@@ -0,0 +1,20 @@
|
||||
# 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.
|
||||
|
||||
# Note: you should be able to use a newer image here:
|
||||
FROM nvcr.io/nvidia/pytorch:25.06-py3
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
RUN pip install ninja imageio imageio-ffmpeg
|
||||
|
||||
COPY csrc /tmp/nvdiffrast/csrc
|
||||
COPY nvdiffrast /tmp/nvdiffrast/nvdiffrast
|
||||
COPY *.py *.toml *.md *.txt /tmp/nvdiffrast
|
||||
RUN TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0" pip install /tmp/nvdiffrast --no-build-isolation
|
||||
Vendored
+97
@@ -0,0 +1,97 @@
|
||||
Copyright (c) 2020, NVIDIA Corporation. All rights reserved.
|
||||
|
||||
|
||||
Nvidia Source Code License (1-Way Commercial)
|
||||
|
||||
=======================================================================
|
||||
|
||||
1. Definitions
|
||||
|
||||
"Licensor" means any person or entity that distributes its Work.
|
||||
|
||||
"Software" means the original work of authorship made available under
|
||||
this License.
|
||||
|
||||
"Work" means the Software and any additions to or derivative works of
|
||||
the Software that are made available under this License.
|
||||
|
||||
The terms "reproduce," "reproduction," "derivative works," and
|
||||
"distribution" have the meaning as provided under U.S. copyright law;
|
||||
provided, however, that for the purposes of this License, derivative
|
||||
works shall not include works that remain separable from, or merely
|
||||
link (or bind by name) to the interfaces of, the Work.
|
||||
|
||||
Works, including the Software, are "made available" under this License
|
||||
by including in or with the Work either (a) a copyright notice
|
||||
referencing the applicability of this License to the Work, or (b) a
|
||||
copy of this License.
|
||||
|
||||
2. License Grants
|
||||
|
||||
2.1 Copyright Grant. Subject to the terms and conditions of this
|
||||
License, each Licensor grants to you a perpetual, worldwide,
|
||||
non-exclusive, royalty-free, copyright license to reproduce,
|
||||
prepare derivative works of, publicly display, publicly perform,
|
||||
sublicense and distribute its Work and any resulting derivative
|
||||
works in any form.
|
||||
|
||||
3. Limitations
|
||||
|
||||
3.1 Redistribution. You may reproduce or distribute the Work only
|
||||
if (a) you do so under this License, (b) you include a complete
|
||||
copy of this License with your distribution, and (c) you retain
|
||||
without modification any copyright, patent, trademark, or
|
||||
attribution notices that are present in the Work.
|
||||
|
||||
3.2 Derivative Works. You may specify that additional or different
|
||||
terms apply to the use, reproduction, and distribution of your
|
||||
derivative works of the Work ("Your Terms") only if (a) Your Terms
|
||||
provide that the use limitation in Section 3.3 applies to your
|
||||
derivative works, and (b) you identify the specific derivative
|
||||
works that are subject to Your Terms. Notwithstanding Your Terms,
|
||||
this License (including the redistribution requirements in Section
|
||||
3.1) will continue to apply to the Work itself.
|
||||
|
||||
3.3 Use Limitation. The Work and any derivative works thereof only
|
||||
may be used or intended for use non-commercially. The Work or
|
||||
derivative works thereof may be used or intended for use by Nvidia
|
||||
or its affiliates commercially or non-commercially. As used herein,
|
||||
"non-commercially" means for research or evaluation purposes only
|
||||
and not for any direct or indirect monetary gain.
|
||||
|
||||
3.4 Patent Claims. If you bring or threaten to bring a patent claim
|
||||
against any Licensor (including any claim, cross-claim or
|
||||
counterclaim in a lawsuit) to enforce any patents that you allege
|
||||
are infringed by any Work, then your rights under this License from
|
||||
such Licensor (including the grant in Section 2.1) will terminate
|
||||
immediately.
|
||||
|
||||
3.5 Trademarks. This License does not grant any rights to use any
|
||||
Licensor's or its affiliates' names, logos, or trademarks, except
|
||||
as necessary to reproduce the notices described in this License.
|
||||
|
||||
3.6 Termination. If you violate any term of this License, then your
|
||||
rights under this License (including the grant in Section 2.1) will
|
||||
terminate immediately.
|
||||
|
||||
4. Disclaimer of Warranty.
|
||||
|
||||
THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
|
||||
NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
|
||||
THIS LICENSE.
|
||||
|
||||
5. Limitation of Liability.
|
||||
|
||||
EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
|
||||
THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
|
||||
SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
|
||||
INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
|
||||
OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
|
||||
(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
|
||||
LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
|
||||
COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
|
||||
THE POSSIBILITY OF SUCH DAMAGES.
|
||||
|
||||
=======================================================================
|
||||
Vendored
+21
@@ -0,0 +1,21 @@
|
||||
# Include standard project files
|
||||
include README.md
|
||||
include LICENSE.txt
|
||||
include pyproject.toml
|
||||
include setup.py
|
||||
|
||||
# Include documentation
|
||||
recursive-include docs *
|
||||
|
||||
# Include CUDA/C++ sources and Python files in the source package.
|
||||
recursive-include nvdiffrast *.py
|
||||
recursive-include csrc *.h *.inl *.cu *.cpp *.hpp
|
||||
|
||||
# Exclude compiled files and caches
|
||||
global-exclude *~ *.py[cod] *.so
|
||||
global-exclude __pycache__
|
||||
global-exclude .git*
|
||||
global-exclude .DS_Store
|
||||
global-exclude *.egg-info
|
||||
global-exclude build
|
||||
global-exclude dist
|
||||
Vendored
+49
@@ -0,0 +1,49 @@
|
||||
## Nvdiffrast – Modular Primitives for High-Performance Differentiable Rendering
|
||||
|
||||
|
||||
|
||||
**Modular Primitives for High-Performance Differentiable Rendering**<br>
|
||||
Samuli Laine, Janne Hellsten, Tero Karras, Yeongho Seol, Jaakko Lehtinen, Timo Aila<br>
|
||||
[http://arxiv.org/abs/2011.03277](http://arxiv.org/abs/2011.03277)
|
||||
|
||||
Nvdiffrast is a PyTorch library that provides high-performance primitive operations for rasterization-based differentiable rendering.
|
||||
|
||||
To install:
|
||||
```
|
||||
pip install setuptools wheel ninja
|
||||
pip install git+https://github.com/NVlabs/nvdiffrast.git --no-build-isolation
|
||||
```
|
||||
|
||||
See ☞☞ [nvdiffrast documentation](https://nvlabs.github.io/nvdiffrast) ☜☜ for more information.
|
||||
|
||||
## Licenses
|
||||
|
||||
Copyright © 2020–2025, NVIDIA Corporation. All rights reserved.
|
||||
|
||||
This work is made available under the [Nvidia Source Code License](https://github.com/NVlabs/nvdiffrast/blob/main/LICENSE.txt).
|
||||
|
||||
For business inquiries, please visit our website and submit the form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/)
|
||||
|
||||
We do not currently accept outside code contributions in the form of pull requests.
|
||||
|
||||
Environment map stored as part of `samples/data/envphong.npz` is derived from a Wave Engine
|
||||
[sample material](https://github.com/WaveEngine/Samples-2.5/tree/master/Materials/EnvironmentMap/Content/Assets/CubeMap.cubemap)
|
||||
originally shared under
|
||||
[MIT License](https://github.com/WaveEngine/Samples-2.5/blob/master/LICENSE.md).
|
||||
Mesh and texture stored as part of `samples/data/earth.npz` are derived from
|
||||
[3D Earth Photorealistic 2K](https://www.turbosquid.com/3d-models/3d-realistic-earth-photorealistic-2k-1279125)
|
||||
model originally made available under
|
||||
[TurboSquid 3D Model License](https://blog.turbosquid.com/turbosquid-3d-model-license/#3d-model-license).
|
||||
|
||||
## Citation
|
||||
|
||||
```
|
||||
@article{Laine2020diffrast,
|
||||
title = {Modular Primitives for High-Performance Differentiable Rendering},
|
||||
author = {Samuli Laine and Janne Hellsten and Tero Karras and Yeongho Seol and Jaakko Lehtinen and Timo Aila},
|
||||
journal = {ACM Transactions on Graphics},
|
||||
year = {2020},
|
||||
volume = {39},
|
||||
number = {6}
|
||||
}
|
||||
```
|
||||
+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;
|
||||
};
|
||||
|
||||
//------------------------------------------------------------------------
|
||||
+1033
File diff suppressed because it is too large
Load Diff
+11
@@ -0,0 +1,11 @@
|
||||
# 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.
|
||||
|
||||
from importlib.metadata import version
|
||||
|
||||
__version__ = version(__package__ or 'nvdiffrast')
|
||||
@@ -0,0 +1,10 @@
|
||||
# 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.
|
||||
|
||||
from .ops import RasterizeCudaContext, get_log_level, set_log_level, rasterize, DepthPeeler, interpolate, texture, texture_construct_mip, antialias, antialias_construct_topology_hash, RasterizeGLContext
|
||||
__all__ = ["RasterizeCudaContext", "get_log_level", "set_log_level", "rasterize", "DepthPeeler", "interpolate", "texture", "texture_construct_mip", "antialias", "antialias_construct_topology_hash", "RasterizeGLContext"]
|
||||
+561
@@ -0,0 +1,561 @@
|
||||
# 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.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import warnings
|
||||
import _nvdiffrast_c
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Log level.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def get_log_level():
|
||||
'''Get current log level.
|
||||
|
||||
Returns:
|
||||
Current log level in nvdiffrast. See `set_log_level()` for possible values.
|
||||
'''
|
||||
return _nvdiffrast_c.get_log_level()
|
||||
|
||||
def set_log_level(level):
|
||||
'''Set log level.
|
||||
|
||||
Log levels follow the convention on the C++ side of Torch:
|
||||
0 = Info,
|
||||
1 = Warning,
|
||||
2 = Error,
|
||||
3 = Fatal.
|
||||
The default log level is 1.
|
||||
|
||||
Args:
|
||||
level: New log level as integer. Internal nvdiffrast messages of this
|
||||
severity or higher will be printed, while messages of lower
|
||||
severity will be silent.
|
||||
'''
|
||||
_nvdiffrast_c.set_log_level(level)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# CudaRaster state wrapper.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
class RasterizeCudaContext:
|
||||
def __init__(self, device=None):
|
||||
'''Create a new Cuda rasterizer context.
|
||||
|
||||
The context is deleted and internal storage is released when the object is
|
||||
destroyed.
|
||||
|
||||
Args:
|
||||
device (Optional): Cuda device on which the context is created. Type can be
|
||||
`torch.device`, string (e.g., `'cuda:1'`), or int. If not
|
||||
specified, context will be created on currently active Cuda
|
||||
device.
|
||||
Returns:
|
||||
The newly created Cuda rasterizer context.
|
||||
'''
|
||||
if device is None:
|
||||
cuda_device_idx = torch.cuda.current_device()
|
||||
else:
|
||||
with torch.cuda.device(device):
|
||||
cuda_device_idx = torch.cuda.current_device()
|
||||
self.cpp_wrapper = _nvdiffrast_c.RasterizeCRStateWrapper(cuda_device_idx)
|
||||
self.active_depth_peeler = None
|
||||
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Rasterize.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
class _rasterize_func(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, raster_ctx, pos, tri, resolution, ranges, grad_db, peeling_idx):
|
||||
out, out_db = _nvdiffrast_c.rasterize_fwd_cuda(raster_ctx.cpp_wrapper, pos, tri, resolution, ranges, peeling_idx)
|
||||
ctx.save_for_backward(pos, tri, out)
|
||||
ctx.saved_grad_db = grad_db
|
||||
return out, out_db
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy, ddb):
|
||||
pos, tri, out = ctx.saved_tensors
|
||||
if ctx.saved_grad_db:
|
||||
g_pos = _nvdiffrast_c.rasterize_grad_db(pos, tri, out, dy, ddb)
|
||||
else:
|
||||
g_pos = _nvdiffrast_c.rasterize_grad(pos, tri, out, dy)
|
||||
return None, g_pos, None, None, None, None, None
|
||||
|
||||
# Op wrapper.
|
||||
def rasterize(glctx, pos, tri, resolution, ranges=None, grad_db=True):
|
||||
'''Rasterize triangles.
|
||||
|
||||
All input tensors must be contiguous and reside in GPU memory except for
|
||||
the `ranges` tensor that, if specified, has to reside in CPU memory. The
|
||||
output tensors will be contiguous and reside in GPU memory.
|
||||
|
||||
Args:
|
||||
glctx: Rasterizer context of type `RasterizeCudaContext`.
|
||||
pos: Vertex position tensor with dtype `torch.float32`. To enable range
|
||||
mode, this tensor should have a 2D shape [num_vertices, 4]. To enable
|
||||
instanced mode, use a 3D shape [minibatch_size, num_vertices, 4].
|
||||
tri: Triangle tensor with shape [num_triangles, 3] and dtype `torch.int32`.
|
||||
resolution: Output resolution as integer tuple (height, width).
|
||||
ranges: In range mode, tensor with shape [minibatch_size, 2] and dtype
|
||||
`torch.int32`, specifying start indices and counts into `tri`.
|
||||
Ignored in instanced mode.
|
||||
grad_db: Propagate gradients of image-space derivatives of barycentrics
|
||||
into `pos` in backward pass.
|
||||
|
||||
Returns:
|
||||
A tuple of two tensors, both with shape [minibatch_size, height, width, 4].
|
||||
The first output tensor contains the main rasterizer output in order
|
||||
(u, v, z/w, triangle_id). The second output tensor contains image-space
|
||||
derivatives of barycentrics in order (du/dX, du/dY, dv/dX, dv/dY).
|
||||
'''
|
||||
assert isinstance(glctx, RasterizeCudaContext)
|
||||
assert grad_db is True or grad_db is False
|
||||
|
||||
# Sanitize inputs.
|
||||
assert isinstance(pos, torch.Tensor) and isinstance(tri, torch.Tensor)
|
||||
resolution = tuple(resolution)
|
||||
if ranges is None:
|
||||
ranges = torch.empty(size=(0, 2), dtype=torch.int32, device='cpu')
|
||||
else:
|
||||
assert isinstance(ranges, torch.Tensor)
|
||||
|
||||
# Check that context is not currently reserved for depth peeling.
|
||||
if glctx.active_depth_peeler is not None:
|
||||
return RuntimeError("Cannot call rasterize() during depth peeling operation, use rasterize_next_layer() instead")
|
||||
|
||||
# Instantiate the function.
|
||||
return _rasterize_func.apply(glctx, pos, tri, resolution, ranges, grad_db, -1)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Depth peeler context manager for rasterizing multiple depth layers.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
class DepthPeeler:
|
||||
def __init__(self, glctx, pos, tri, resolution, ranges=None, grad_db=True):
|
||||
'''Create a depth peeler object for rasterizing multiple depth layers.
|
||||
|
||||
Arguments are the same as in `rasterize()`.
|
||||
|
||||
Returns:
|
||||
The newly created depth peeler.
|
||||
'''
|
||||
assert isinstance(glctx, RasterizeCudaContext)
|
||||
assert grad_db is True or grad_db is False
|
||||
|
||||
# Sanitize inputs as usual.
|
||||
assert isinstance(pos, torch.Tensor) and isinstance(tri, torch.Tensor)
|
||||
resolution = tuple(resolution)
|
||||
if ranges is None:
|
||||
ranges = torch.empty(size=(0, 2), dtype=torch.int32, device='cpu')
|
||||
else:
|
||||
assert isinstance(ranges, torch.Tensor)
|
||||
|
||||
# Store all the parameters.
|
||||
self.raster_ctx = glctx
|
||||
self.pos = pos
|
||||
self.tri = tri
|
||||
self.resolution = resolution
|
||||
self.ranges = ranges
|
||||
self.grad_db = grad_db
|
||||
self.peeling_idx = None
|
||||
|
||||
def __enter__(self):
|
||||
if self.raster_ctx is None:
|
||||
raise RuntimeError("Cannot re-enter a terminated depth peeling operation")
|
||||
if self.raster_ctx.active_depth_peeler is not None:
|
||||
raise RuntimeError("Cannot have multiple depth peelers active simultaneously in a rasterization context")
|
||||
self.raster_ctx.active_depth_peeler = self
|
||||
self.peeling_idx = 0
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
assert self.raster_ctx.active_depth_peeler is self
|
||||
self.raster_ctx.active_depth_peeler = None
|
||||
self.raster_ctx = None # Remove all references to input tensor so they're not left dangling.
|
||||
self.pos = None
|
||||
self.tri = None
|
||||
self.resolution = None
|
||||
self.ranges = None
|
||||
self.grad_db = None
|
||||
self.peeling_idx = None
|
||||
return None
|
||||
|
||||
def rasterize_next_layer(self):
|
||||
'''Rasterize next depth layer.
|
||||
|
||||
Operation is equivalent to `rasterize()` except that previously reported
|
||||
surface points are culled away.
|
||||
|
||||
Returns:
|
||||
A tuple of two tensors as in `rasterize()`.
|
||||
'''
|
||||
assert self.raster_ctx.active_depth_peeler is self
|
||||
assert self.peeling_idx >= 0
|
||||
result = _rasterize_func.apply(self.raster_ctx, self.pos, self.tri, self.resolution, self.ranges, self.grad_db, self.peeling_idx)
|
||||
self.peeling_idx += 1
|
||||
return result
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Interpolate.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
# Output pixel differentials for at least some attributes.
|
||||
class _interpolate_func_da(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list):
|
||||
out, out_da = _nvdiffrast_c.interpolate_fwd_da(attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list)
|
||||
ctx.save_for_backward(attr, rast, tri, rast_db)
|
||||
ctx.saved_misc = diff_attrs_all, diff_attrs_list
|
||||
return out, out_da
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy, dda):
|
||||
attr, rast, tri, rast_db = ctx.saved_tensors
|
||||
diff_attrs_all, diff_attrs_list = ctx.saved_misc
|
||||
g_attr, g_rast, g_rast_db = _nvdiffrast_c.interpolate_grad_da(attr, rast, tri, dy, rast_db, dda, diff_attrs_all, diff_attrs_list)
|
||||
return g_attr, g_rast, None, g_rast_db, None, None
|
||||
|
||||
# No pixel differential for any attribute.
|
||||
class _interpolate_func(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, attr, rast, tri):
|
||||
out, out_da = _nvdiffrast_c.interpolate_fwd(attr, rast, tri)
|
||||
ctx.save_for_backward(attr, rast, tri)
|
||||
return out, out_da
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy, _):
|
||||
attr, rast, tri = ctx.saved_tensors
|
||||
g_attr, g_rast = _nvdiffrast_c.interpolate_grad(attr, rast, tri, dy)
|
||||
return g_attr, g_rast, None
|
||||
|
||||
# Op wrapper.
|
||||
def interpolate(attr, rast, tri, rast_db=None, diff_attrs=None):
|
||||
"""Interpolate vertex attributes.
|
||||
|
||||
All input tensors must be contiguous and reside in GPU memory. The output tensors
|
||||
will be contiguous and reside in GPU memory.
|
||||
|
||||
Args:
|
||||
attr: Attribute tensor with dtype `torch.float32`.
|
||||
Shape is [num_vertices, num_attributes] in range mode, or
|
||||
[minibatch_size, num_vertices, num_attributes] in instanced mode.
|
||||
Broadcasting is supported along the minibatch axis.
|
||||
rast: Main output tensor from `rasterize()`.
|
||||
tri: Triangle tensor with shape [num_triangles, 3] and dtype `torch.int32`.
|
||||
rast_db: (Optional) Tensor containing image-space derivatives of barycentrics,
|
||||
i.e., the second output tensor from `rasterize()`. Enables computing
|
||||
image-space derivatives of attributes.
|
||||
diff_attrs: (Optional) List of attribute indices for which image-space
|
||||
derivatives are to be computed. Special value 'all' is equivalent
|
||||
to list [0, 1, ..., num_attributes - 1].
|
||||
|
||||
Returns:
|
||||
A tuple of two tensors. The first output tensor contains interpolated
|
||||
attributes and has shape [minibatch_size, height, width, num_attributes].
|
||||
If `rast_db` and `diff_attrs` were specified, the second output tensor contains
|
||||
the image-space derivatives of the selected attributes and has shape
|
||||
[minibatch_size, height, width, 2 * len(diff_attrs)]. The derivatives of the
|
||||
first selected attribute A will be on channels 0 and 1 as (dA/dX, dA/dY), etc.
|
||||
Otherwise, the second output tensor will be an empty tensor with shape
|
||||
[minibatch_size, height, width, 0].
|
||||
"""
|
||||
# Sanitize the list of pixel differential attributes.
|
||||
if diff_attrs is None:
|
||||
diff_attrs = []
|
||||
elif diff_attrs != 'all':
|
||||
diff_attrs = np.asarray(diff_attrs, np.int32)
|
||||
assert len(diff_attrs.shape) == 1
|
||||
diff_attrs = diff_attrs.tolist()
|
||||
|
||||
diff_attrs_all = int(diff_attrs == 'all')
|
||||
diff_attrs_list = [] if diff_attrs_all else diff_attrs
|
||||
|
||||
# Check inputs.
|
||||
assert all(isinstance(x, torch.Tensor) for x in (attr, rast, tri))
|
||||
if diff_attrs:
|
||||
assert isinstance(rast_db, torch.Tensor)
|
||||
|
||||
# Choose stub.
|
||||
if diff_attrs:
|
||||
return _interpolate_func_da.apply(attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list)
|
||||
else:
|
||||
return _interpolate_func.apply(attr, rast, tri)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Texture
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
# Linear-mipmap-linear and linear-mipmap-nearest: Mipmaps enabled.
|
||||
class _texture_func_mip(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, filter_mode, tex, uv, uv_da, mip_level_bias, mip_wrapper, filter_mode_enum, boundary_mode_enum, *mip_stack):
|
||||
empty = torch.tensor([])
|
||||
if uv_da is None:
|
||||
uv_da = empty
|
||||
if mip_level_bias is None:
|
||||
mip_level_bias = empty
|
||||
if mip_wrapper is None:
|
||||
mip_wrapper = _nvdiffrast_c.TextureMipWrapper()
|
||||
out = _nvdiffrast_c.texture_fwd_mip(tex, uv, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode_enum, boundary_mode_enum)
|
||||
ctx.save_for_backward(tex, uv, uv_da, mip_level_bias, *mip_stack)
|
||||
ctx.saved_misc = filter_mode, mip_wrapper, filter_mode_enum, boundary_mode_enum
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy):
|
||||
tex, uv, uv_da, mip_level_bias, *mip_stack = ctx.saved_tensors
|
||||
filter_mode, mip_wrapper, filter_mode_enum, boundary_mode_enum = ctx.saved_misc
|
||||
if filter_mode == 'linear-mipmap-linear':
|
||||
g_tex, g_uv, g_uv_da, g_mip_level_bias, g_mip_stack = _nvdiffrast_c.texture_grad_linear_mipmap_linear(tex, uv, dy, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode_enum, boundary_mode_enum)
|
||||
return (None, g_tex, g_uv, g_uv_da, g_mip_level_bias, None, None, None) + tuple(g_mip_stack)
|
||||
else: # linear-mipmap-nearest
|
||||
g_tex, g_uv, g_mip_stack = _nvdiffrast_c.texture_grad_linear_mipmap_nearest(tex, uv, dy, uv_da, mip_level_bias, mip_wrapper, mip_stack, filter_mode_enum, boundary_mode_enum)
|
||||
return (None, g_tex, g_uv, None, None, None, None, None) + tuple(g_mip_stack)
|
||||
|
||||
# Linear and nearest: Mipmaps disabled.
|
||||
class _texture_func(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, filter_mode, tex, uv, filter_mode_enum, boundary_mode_enum):
|
||||
out = _nvdiffrast_c.texture_fwd(tex, uv, filter_mode_enum, boundary_mode_enum)
|
||||
ctx.save_for_backward(tex, uv)
|
||||
ctx.saved_misc = filter_mode, filter_mode_enum, boundary_mode_enum
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy):
|
||||
tex, uv = ctx.saved_tensors
|
||||
filter_mode, filter_mode_enum, boundary_mode_enum = ctx.saved_misc
|
||||
if filter_mode == 'linear':
|
||||
g_tex, g_uv = _nvdiffrast_c.texture_grad_linear(tex, uv, dy, filter_mode_enum, boundary_mode_enum)
|
||||
return None, g_tex, g_uv, None, None
|
||||
else: # nearest
|
||||
g_tex = _nvdiffrast_c.texture_grad_nearest(tex, uv, dy, filter_mode_enum, boundary_mode_enum)
|
||||
return None, g_tex, None, None, None
|
||||
|
||||
# Op wrapper.
|
||||
def texture(tex, uv, uv_da=None, mip_level_bias=None, mip=None, filter_mode='auto', boundary_mode='wrap', max_mip_level=None):
|
||||
"""Perform texture sampling.
|
||||
|
||||
All input tensors must be contiguous and reside in GPU memory. The output tensor
|
||||
will be contiguous and reside in GPU memory.
|
||||
|
||||
Args:
|
||||
tex: Texture tensor with dtype `torch.float32`. For 2D textures, must have shape
|
||||
[minibatch_size, tex_height, tex_width, tex_channels]. For cube map textures,
|
||||
must have shape [minibatch_size, 6, tex_height, tex_width, tex_channels] where
|
||||
tex_width and tex_height are equal. Note that `boundary_mode` must also be set
|
||||
to 'cube' to enable cube map mode. Broadcasting is supported along the minibatch axis.
|
||||
uv: Tensor containing per-pixel texture coordinates. When sampling a 2D texture,
|
||||
must have shape [minibatch_size, height, width, 2]. When sampling a cube map
|
||||
texture, must have shape [minibatch_size, height, width, 3].
|
||||
uv_da: (Optional) Tensor containing image-space derivatives of texture coordinates.
|
||||
Must have same shape as `uv` except for the last dimension that is to be twice
|
||||
as long.
|
||||
mip_level_bias: (Optional) Per-pixel bias for mip level selection. If `uv_da` is omitted,
|
||||
determines mip level directly. Must have shape [minibatch_size, height, width].
|
||||
mip: (Optional) Preconstructed mipmap stack from a `texture_construct_mip()` call, or a list
|
||||
of tensors specifying a custom mipmap stack. When specifying a custom mipmap stack,
|
||||
the tensors in the list must follow the same format as `tex` except for width and
|
||||
height that must follow the usual rules for mipmap sizes. The base level texture
|
||||
is still supplied in `tex` and must not be included in the list. Gradients of a
|
||||
custom mipmap stack are not automatically propagated to base texture but the mipmap
|
||||
tensors will receive gradients of their own. If a mipmap stack is not specified
|
||||
but the chosen filter mode requires it, the mipmap stack is constructed internally
|
||||
and discarded afterwards.
|
||||
filter_mode: Texture filtering mode to be used. Valid values are 'auto', 'nearest',
|
||||
'linear', 'linear-mipmap-nearest', and 'linear-mipmap-linear'. Mode 'auto'
|
||||
selects 'linear' if neither `uv_da` or `mip_level_bias` is specified, and
|
||||
'linear-mipmap-linear' when at least one of them is specified, these being
|
||||
the highest-quality modes possible depending on the availability of the
|
||||
image-space derivatives of the texture coordinates or direct mip level information.
|
||||
boundary_mode: Valid values are 'wrap', 'clamp', 'zero', and 'cube'. If `tex` defines a
|
||||
cube map, this must be set to 'cube'. The default mode 'wrap' takes fractional
|
||||
part of texture coordinates. Mode 'clamp' clamps texture coordinates to the
|
||||
centers of the boundary texels. Mode 'zero' virtually extends the texture with
|
||||
all-zero values in all directions.
|
||||
max_mip_level: If specified, limits the number of mipmaps constructed and used in mipmap-based
|
||||
filter modes.
|
||||
|
||||
Returns:
|
||||
A tensor containing the results of the texture sampling with shape
|
||||
[minibatch_size, height, width, tex_channels]. Cube map fetches with invalid uv coordinates
|
||||
(e.g., zero vectors) output all zeros and do not propagate gradients.
|
||||
"""
|
||||
|
||||
# Default filter mode.
|
||||
if filter_mode == 'auto':
|
||||
filter_mode = 'linear-mipmap-linear' if (uv_da is not None or mip_level_bias is not None) else 'linear'
|
||||
|
||||
# Sanitize inputs.
|
||||
if max_mip_level is None:
|
||||
max_mip_level = -1
|
||||
else:
|
||||
max_mip_level = int(max_mip_level)
|
||||
assert max_mip_level >= 0
|
||||
|
||||
# Check inputs.
|
||||
assert isinstance(tex, torch.Tensor) and isinstance(uv, torch.Tensor)
|
||||
if 'mipmap' in filter_mode:
|
||||
assert isinstance(uv_da, torch.Tensor) or isinstance(mip_level_bias, torch.Tensor)
|
||||
|
||||
# If mipping disabled via max level=0, we may as well use simpler filtering internally.
|
||||
if max_mip_level == 0 and filter_mode in ['linear-mipmap-nearest', 'linear-mipmap-linear']:
|
||||
filter_mode = 'linear'
|
||||
|
||||
# Convert filter mode to internal enumeration.
|
||||
filter_mode_dict = {'nearest': 0, 'linear': 1, 'linear-mipmap-nearest': 2, 'linear-mipmap-linear': 3}
|
||||
filter_mode_enum = filter_mode_dict[filter_mode]
|
||||
|
||||
# Convert boundary mode to internal enumeration.
|
||||
boundary_mode_dict = {'cube': 0, 'wrap': 1, 'clamp': 2, 'zero': 3}
|
||||
boundary_mode_enum = boundary_mode_dict[boundary_mode]
|
||||
|
||||
# Construct a mipmap if necessary.
|
||||
if 'mipmap' in filter_mode:
|
||||
mip_wrapper, mip_stack = None, []
|
||||
if mip is not None:
|
||||
assert isinstance(mip, (_nvdiffrast_c.TextureMipWrapper, list))
|
||||
if isinstance(mip, list):
|
||||
assert all(isinstance(x, torch.Tensor) for x in mip)
|
||||
mip_stack = mip
|
||||
else:
|
||||
mip_wrapper = mip
|
||||
else:
|
||||
mip_wrapper = _nvdiffrast_c.texture_construct_mip(tex, max_mip_level, boundary_mode == 'cube')
|
||||
|
||||
# Choose stub.
|
||||
if filter_mode == 'linear-mipmap-linear' or filter_mode == 'linear-mipmap-nearest':
|
||||
return _texture_func_mip.apply(filter_mode, tex, uv, uv_da, mip_level_bias, mip_wrapper, filter_mode_enum, boundary_mode_enum, *mip_stack)
|
||||
else:
|
||||
return _texture_func.apply(filter_mode, tex, uv, filter_mode_enum, boundary_mode_enum)
|
||||
|
||||
# Mipmap precalculation for cases where the texture stays constant.
|
||||
def texture_construct_mip(tex, max_mip_level=None, cube_mode=False):
|
||||
"""Construct a mipmap stack for a texture.
|
||||
|
||||
This function can be used for constructing a mipmap stack for a texture that is known to remain
|
||||
constant. This avoids reconstructing it every time `texture()` is called.
|
||||
|
||||
Args:
|
||||
tex: Texture tensor with the same constraints as in `texture()`.
|
||||
max_mip_level: If specified, limits the number of mipmaps constructed.
|
||||
cube_mode: Must be set to True if `tex` specifies a cube map texture.
|
||||
|
||||
Returns:
|
||||
An opaque object containing the mipmap stack. This can be supplied in a call to `texture()`
|
||||
in the `mip` argument.
|
||||
"""
|
||||
|
||||
assert isinstance(tex, torch.Tensor)
|
||||
assert cube_mode is True or cube_mode is False
|
||||
if max_mip_level is None:
|
||||
max_mip_level = -1
|
||||
else:
|
||||
max_mip_level = int(max_mip_level)
|
||||
assert max_mip_level >= 0
|
||||
return _nvdiffrast_c.texture_construct_mip(tex, max_mip_level, cube_mode)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Antialias.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
class _antialias_func(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, color, rast, pos, tri, topology_hash, pos_gradient_boost):
|
||||
out, work_buffer = _nvdiffrast_c.antialias_fwd(color, rast, pos, tri, topology_hash)
|
||||
ctx.save_for_backward(color, rast, pos, tri)
|
||||
ctx.saved_misc = pos_gradient_boost, work_buffer
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dy):
|
||||
color, rast, pos, tri = ctx.saved_tensors
|
||||
pos_gradient_boost, work_buffer = ctx.saved_misc
|
||||
g_color, g_pos = _nvdiffrast_c.antialias_grad(color, rast, pos, tri, dy, work_buffer)
|
||||
if pos_gradient_boost != 1.0:
|
||||
g_pos = g_pos * pos_gradient_boost
|
||||
return g_color, None, g_pos, None, None, None
|
||||
|
||||
# Op wrapper.
|
||||
def antialias(color, rast, pos, tri, topology_hash=None, pos_gradient_boost=1.0):
|
||||
"""Perform antialiasing.
|
||||
|
||||
All input tensors must be contiguous and reside in GPU memory. The output tensor
|
||||
will be contiguous and reside in GPU memory.
|
||||
|
||||
Note that silhouette edge determination is based on vertex indices in the triangle
|
||||
tensor. For it to work properly, a vertex belonging to multiple triangles must be
|
||||
referred to using the same vertex index in each triangle. Otherwise, nvdiffrast will always
|
||||
classify the adjacent edges as silhouette edges, which leads to bad performance and
|
||||
potentially incorrect gradients. If you are unsure whether your data is good, check
|
||||
which pixels are modified by the antialias operation and compare to the example in the
|
||||
documentation.
|
||||
|
||||
Args:
|
||||
color: Input image to antialias with shape [minibatch_size, height, width, num_channels].
|
||||
rast: Main output tensor from `rasterize()`.
|
||||
pos: Vertex position tensor used in the rasterization operation.
|
||||
tri: Triangle tensor used in the rasterization operation.
|
||||
topology_hash: (Optional) Preconstructed topology hash for the triangle tensor. If not
|
||||
specified, the topology hash is constructed internally and discarded afterwards.
|
||||
pos_gradient_boost: (Optional) Multiplier for gradients propagated to `pos`.
|
||||
|
||||
Returns:
|
||||
A tensor containing the antialiased image with the same shape as `color` input tensor.
|
||||
"""
|
||||
|
||||
# Check inputs.
|
||||
assert all(isinstance(x, torch.Tensor) for x in (color, rast, pos, tri))
|
||||
|
||||
# Construct topology hash unless provided by user.
|
||||
if topology_hash is not None:
|
||||
assert isinstance(topology_hash, _nvdiffrast_c.TopologyHashWrapper)
|
||||
else:
|
||||
topology_hash = _nvdiffrast_c.antialias_construct_topology_hash(tri)
|
||||
|
||||
# Instantiate the function.
|
||||
return _antialias_func.apply(color, rast, pos, tri, topology_hash, pos_gradient_boost)
|
||||
|
||||
# Topology hash precalculation for cases where the triangle array stays constant.
|
||||
def antialias_construct_topology_hash(tri):
|
||||
"""Construct a topology hash for a triangle tensor.
|
||||
|
||||
This function can be used for constructing a topology hash for a triangle tensor that is
|
||||
known to remain constant. This avoids reconstructing it every time `antialias()` is called.
|
||||
|
||||
Args:
|
||||
tri: Triangle tensor with shape [num_triangles, 3]. Must be contiguous and reside in
|
||||
GPU memory.
|
||||
|
||||
Returns:
|
||||
An opaque object containing the topology hash. This can be supplied in a call to
|
||||
`antialias()` in the `topology_hash` argument.
|
||||
"""
|
||||
assert isinstance(tri, torch.Tensor)
|
||||
return _nvdiffrast_c.antialias_construct_topology_hash(tri)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Legacy OpenGL context stub for backwards compatibility.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
class RasterizeGLContext(RasterizeCudaContext):
|
||||
def __init__(self, output_db=True, mode='automatic', device=None):
|
||||
warnings.warn("RasterizeGLContext has been deprecated and uses RasterizeCudaContext internally", DeprecationWarning, stacklevel=2)
|
||||
super().__init__(device=device)
|
||||
|
||||
def set_context(self):
|
||||
pass
|
||||
|
||||
def release_context(self):
|
||||
pass
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+31
@@ -0,0 +1,31 @@
|
||||
[build-system]
|
||||
requires = [
|
||||
"setuptools>=64",
|
||||
"wheel",
|
||||
"ninja",
|
||||
]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "nvdiffrast"
|
||||
version = "0.4.0"
|
||||
description = "nvdiffrast - modular primitives for high-performance differentiable rendering"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.6"
|
||||
#license = "LicenseRef-NSCL"
|
||||
#license-files = ["LICENSE.txt"]
|
||||
authors = [{name = "Samuli Laine", email = "slaine@nvidia.com"}]
|
||||
classifiers = [
|
||||
"Programming Language :: Python :: 3",
|
||||
"Operating System :: OS Independent",
|
||||
]
|
||||
dependencies = ["numpy"]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/NVlabs/nvdiffrast"
|
||||
Repository = "https://github.com/NVlabs/nvdiffrast"
|
||||
Documentation = "https://nvlabs.github.io/nvdiffrast/"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
include = ["nvdiffrast", "nvdiffrast.torch"]
|
||||
exclude = ["nvdiffrast.common*"]
|
||||
+225
@@ -0,0 +1,225 @@
|
||||
|
||||
Environment map stored as part of samples/data/envphong.npz is derived from a Wave Engine sample material originally shared under MIT License that is reproduced below.
|
||||
Original material: https://github.com/WaveEngine/Samples/tree/master/Materials/EnvironmentMap/Content/Assets/CubeMap.cubemap
|
||||
Original license: https://github.com/WaveEngine/Samples/blob/master/LICENSE.md
|
||||
|
||||
Mesh and texture stored as part of samples/data/earth.npz are derived from "3D Earth Photorealistic 2K" model originally made available under TurboSquid 3D Model License that is reproduced below.
|
||||
Original material: https://www.turbosquid.com/3d-models/3d-realistic-earth-photorealistic-2k-1279125
|
||||
Original license: https://blog.turbosquid.com/turbosquid-3d-model-license/#3d-model-license
|
||||
|
||||
|
||||
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2016 Wave Coorporation
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
|
||||
|
||||
|
||||
TurboSquid 3D Model License
|
||||
|
||||
This is a legally binding agreement between licensee ("you"), and TurboSquid regarding your rights to use 3D Models from the Site under this license. "You" refers to the purchasing entity, whether that is a natural person who must be at least 18 years of age, or a corporate entity. The rights granted in this agreement are granted to the purchasing entity, its parent company, and its majority owned affiliates on a "royalty free" basis, which means that after a Purchase, there are no future royalties or payments that are required. This agreement incorporates by reference the Terms of Use as well as the Site's policies and procedures as such.
|
||||
I. Introduction & Definitions
|
||||
|
||||
Definitions
|
||||
|
||||
This agreement is intended to be easy to understand, and to provide clarity for using 3D Models in the work you create ("Creations"). Over the years, TurboSquid has been asked many questions about how 3D Models may be used in Creations, and we have attempted to answer those questions in this agreement.
|
||||
|
||||
Some words in this agreement are given specific meanings. Words that appear initially in quotations, such as "you" and "Creations", are defined in the text preceding the word. Other capitalized words are defined below:
|
||||
|
||||
"3D Model" is the collection of one or more digital files, packaged in the form of a product on the Site that can be identified by a 3D Model ID, and that is made available to you for Purchase on the Site. A 3D Model may include 3D Model files, geometry, texture maps, materials, motion captures, renderings and other constituent files related to the 3D Model data and its representation.
|
||||
|
||||
"Site" refers to the TurboSquid websites, API's, software applications or any approved means or utility either currently in existence or in the future; the software and source code used by TurboSquid to provide such services; user interface layouts, designs, images, text, knowledgebase articles, program offers; site information provided in reports (such as popular keyword searches); and all other intellectual property protected under copyright, trademark, patent, publicity, or any other proprietary right.
|
||||
|
||||
"Purchase" is the acquisition of a 3D Model by you from the Site under this agreement, whether as a purchase of 3D Model made available at a price of greater than $0, or a download of 3D Model made available at no charge.
|
||||
|
||||
"TurboSquid" includes TurboSquid, Inc. and all licensed affiliates and partners that distribute 3D Models on behalf of TurboSquid, Inc.
|
||||
|
||||
"Product Page" is the product page or interface that displays 3D Models available for Purchase on the Site.
|
||||
|
||||
"Computer Game" is a type of Creation that includes digital games, computer-based games, handheld electronic games, mobile games, online games, web-games, social games, game mods, and console-based games.
|
||||
|
||||
"Imagery" is a Creation made of any single image or sequence of images.
|
||||
|
||||
"Depicted Intellectual Property" means any intellectual property depicted in the 3D Model, including any copyright, trademark, trade dress, right of publicity, or any other proprietary right throughout the world that may apply. For purposes of clarity, this does not refer to the copyrights owned by the creator of the 3D Model that are licensed in this agreement.
|
||||
|
||||
To make reading this agreement easier and less repetitive, the following constructions are used:
|
||||
|
||||
"Include," including," and "such as" are considered to be followed with "but not limited to." Examples are used in this agreement to illustrate, rather than limit, the scope of the terms.
|
||||
|
||||
"The following restrictions", "the foregoing restrictions", and "subject to the restrictions" are considered to be followed with "in addition to all other restrictions applicable within this agreement."
|
||||
II. License Rights
|
||||
|
||||
1. Ownership. TurboSquid does not grant title or ownership in 3D Models. All rights in 3D Models not expressly granted in this agreement are reserved by TurboSquid for itself and its licensors.
|
||||
|
||||
2. Rights Granted. For 3D Models, TurboSquid grants to you a non-exclusive, perpetual, worldwide right and license to copy, distribute, reproduce, adapt, publicly display, publicly perform, digitally perform, transmit, broadcast, telecast, advertise, create derivative works, and market 3D Models within Creations in the uses authorized in this agreement. You may request authorization for a use not covered by this agreement ("New Use") by writing use@turbosquid.com. TurboSquid is authorized to approve a New Use if TurboSquid finds in its sole judgment that the New Use is substantially similar to another established use in this agreement and authorizes the New Use in writing.
|
||||
|
||||
3. Rights Granted When Sharing 3D Models. If you Purchase as an employee of a corporate entity, sharing Purchased 3D Models with other employees of your corporate entity is allowed. Examples of allowed sharing include storing files on a networked hard drive, and aggregating 3D Models for later use in future Creations. You are responsible for any downstream distribution, use, or misuse by a recipient of a shared 3D Models. In all cases, sharing 3D Models with external people or entities is only allowed in the following situations, and with the following restrictions:
|
||||
|
||||
a. In the production of a Creation owned by you, if you are working in collaboration with external parties, and there is a need to share 3D Models for the development and production of your Creation, sharing 3D Models with those external parties is allowed. Any external party that receives 3D Models may only use 3D Models on your Creations and must take reasonable care to secure and limit access to 3D Models to that purpose.
|
||||
|
||||
b. In the production of a Creation owned by another entity ("your Client"), if you are working as a contractor and need to share 3D Models with your Client, or any external parties working with your Client, sharing 3D Models is allowed, subject to the restriction that all parties may use 3D Models only for your Client's particular Creation, and for successive versions of your Client's Creation, such as sequel Computer Games or movies that utilize the same 3D Models. All parties must take reasonable care to secure and limit access to 3D Models to the parties working on your Client's Creation. For all other use by any party, 3D Models must be Purchased again to create a new license agreement governing that use
|
||||
|
||||
4. Editorial Use Restriction for Some 3D Models. The following restrictions apply to any 3D Model with an "Editorial Uses Only" label on its Product Page. Permitted use of Depicted Intellectual Property in such 3D Models is limited to news reporting in Creations of some cultural, editorial, journalistic, or otherwise newsworthy value, including news reporting on television and the internet. A second permitted use is use within an academic setting, limited to teaching, scholarship, and research. This restriction does not apply if you have the needed authorization to use the Depicted Intellectual Property for your Creation, such as if you are owner of the Depicted Intellectual Property, or the advertising team, hired party, or licensee of the Depicted Intellectual Property owner.
|
||||
|
||||
5. Depicted Intellectual Property. TurboSquid does not own or license any Depicted Intellectual Property. TurboSquid does not in any way make any representations or warranties about Depicted Intellectual Property associated with 3D Models. You are solely responsible for determining the need for and, if appropriate, obtaining any needed clearance, consent, or release to use any Depicted Intellectual Property in your Creations.
|
||||
|
||||
6. Creations of Imagery.
|
||||
|
||||
Permitted Uses of Creations of Imagery. Subject to the following restrictions, you may use Creations of Imagery within news, film, movies, television programs, video projects, multi-media projects, theatrical display, software user interfaces; architectural renderings, Computer Games, virtual worlds, simulation and training environments; corporate communications, marketing collateral, tradeshow promotional items, booth decorations and presentations; pre-visualizations, product prototyping and research; mobile, web, print, television, and billboard advertising; online and electronic publications of blogs, literature, social media, and email campaigns; website designs and layouts, desktop and mobile wallpapers, screensavers, toolbar skins; books, magazines, posters, greeting cards; apparel items, brochures, framed or printed artwork, household items, office items, lenticular prints, product packaging and manufactured products.
|
||||
|
||||
Restrictions on Permitted Uses of Creations of Imagery.
|
||||
|
||||
a. Stock Media Clearinghouse. You may NOT publish or distribute Creations of Imagery through another stock media clearinghouse, for example as part of an online marketplace for photography, clip art, or design templates.
|
||||
|
||||
b. Promotional Images. Images displayed for the promotion a 3D Model on its Product Page ("Promotional Images") may be used in Creations of Imagery, provided that the 3D Model itself has been Purchased and subject to the following restrictions:
|
||||
|
||||
i. You may NOT use a Promotional Image that has any added element which is not included as part of the 3D Model. An example of this type of restricted use is if the 3D Model contains an airplane, and there is a Promotional Image of that airplane rendered over a blue sky; however, the blue sky image is not included as part of the 3D Model. Other prohibited examples include use of Promotional Images from movies or advertisements that may have used 3D Model.
|
||||
|
||||
ii. You may NOT use any Promotional Image that has a logo, mark, watermark, attribution, copyright or other notice superimposed on the image without prior approval from TurboSquid Support.
|
||||
|
||||
c. Business Logos. You may NOT use Imagery in any Creation that is a trademark, servicemark, or business logo. This restriction is included because the owners of these types of Creations typically seek exclusivity on the use of the imagery in their Creation, which is incompatible with the non-exclusive license granted to you under this agreement.
|
||||
|
||||
|
||||
7. Creations of Computer Games and Software
|
||||
|
||||
Permitted Uses in Creations of Computer Games and Software. Subject to the following restrictions, you may incorporate 3D Models in Creations of Computer Games, virtual worlds, simulation and training environments; mobile, desktop and web applications; and interactive electronic publications of literature such as e-books and electronic textbooks.
|
||||
|
||||
Restrictions on Permitted Uses of 3D Models in Creations of Games and Software.
|
||||
|
||||
a. Interactivity. Your inclusion of 3D Models within any such Creation is limited to uses where 3D Model is contained in an interactive experience for the user and not made available outside of the interactive experience. Such a permitted example of this use would be to include a 3D Model of human anatomy in a medical training application in a way that the 3D Model or its environment may be manipulated or interacted with.
|
||||
|
||||
b. Access to 3D Models. You must take all reasonable and industry standard measures to incorporate 3D Models within Creations to prevent other parties from gaining access to 3D Models. 3D Models must be contained in proprietary formats so that they cannot be opened or imported in a publicly available software application or framework, or extracted without reverse engineering. WebGL exports from Unity, Unreal, and Lumberyard are permitted. Any other open format or format encrypted with decryptable open standards (such as an encrypted compression archive or other WebGL programs not listed here) are prohibited from using 3D Models. If your Creation uses WebGL and you are not sure if it qualifies, please contact use@turbosquid.com and describe your Creation in detail.
|
||||
|
||||
c. Open Systems. You typically may NOT include 3D Models in Creations that have the general functionality for importing and/or exporting 3D Models. Please contact use@turbosquid.com and describe your Creation in detail if this is your desired use. An example of such a prohibited use is to include 3D Models as a starter library within a standard retail Software Creation that allows users to generally work with 3D Models, even if the 3D Model itself is somehow protected and is not capable of being exported. An allowed use is for custom or enterprise software in certain circumstances.
|
||||
|
||||
d. Virtual Good Sales. You may NOT import, upload, reproduce, make available, publish, transmit, distribute, or sublicense 3D Models in Creations of virtual goods or worlds for any 3D community ("Virtual World"), unless you or your Client owns the Virtual World platform and it complies with the previous restrictions.
|
||||
|
||||
|
||||
8. Creations of Physical Form.
|
||||
|
||||
Permitted Uses in Creations of Physical Form. Subject to the following restrictions, you may use 3D Models to make Physical Creations such as 3D printed works, articles of manufacture, custom vehicles, furniture, jewelry, sculptural artwork, toys, and physical entertainment goods ("Creations of Physical Form").
|
||||
|
||||
Restrictions on Permitted Uses in Creations of Physical Form.
|
||||
|
||||
a. Substantially Similar Creations. Permitted use of any Creation of Physical Form in which a 3D Model is untransformed or substantially similar to the 3D Model is limited to personal use, gifts, or charitable donations, with a maximum of 5 instances of such Creation per Purchase; unless the 3D Model is a small part of a much larger array of other physical objects in the Creation. For example, if you are creating a real-world, physical human skeleton for manufacture for sale, it is permitted to add a 3D printed human head that exactly resembles the Purchased 3D Model, but it is not permitted to sell the 3D printed head by itself. Another permitted example of a 3D Model being a small part of a larger array is using a 3D Model that ends up within an automobile as a part of the automobile.
|
||||
|
||||
b. No Depicted Intellectual Property. You may NOT reproduce Depicted Intellectual Property in any Creation of Physical Form for any purpose. For example, you may NOT make Physical Form Creations of a copyrighted character (Spiderman, Elsa, Slimer), or branded technology (Apple, Toshiba, Samsung).
|
||||
|
||||
9. 3D Industry Promotional Use. If TurboSquid has granted you, as a hardware or software partner, access to priced 3D Models on a free-of-charge basis, your use of 3D Models is restricted to internal testing for your 3D software or hardware products, and to the promotion of your software or hardware products with Creations of Imagery provided that an attribution of the artist's name and the Site are included. You agree that should any 3D Models be used outside of these purposes in ways that are normally allowed after a Purchase, that you will notify TurboSquid and promptly Purchase the 3D Models and otherwise comply with the terms herein.
|
||||
|
||||
10. Unauthorized Use. If you use 3D Models in an unauthorized way, TurboSquid may terminate your account and pursue other penalties, damages, losses, and profits TurboSquid is entitled to under this agreement or at law or equity. The following are unauthorized uses that are explicitly prohibited:
|
||||
|
||||
a. Competition. You may NOT use 3D Models in a way that competes with the Site, including distributing through 3D Model Clearinghouses. You may NOT publish, distribute, or make 3D Models available through any online clearinghouse infrastructure. You may not redistribute 3D Models as part of any design template, After Effects template, stock photography, video or clip art for distribution or licensing through any online stock media clearinghouse whatever.
|
||||
|
||||
b. Re-Distribution. You may NOT re-distribute, publish, or make 3D Models available to any third party except in the form of a permitted Creation, or shared as authorized in this agreement.
|
||||
|
||||
c. Group Buying. You may NOT aggregate funds to Purchase 3D Models with one or more other parties. An example of this prohibited use is a website membership where members pool their money to make a single Purchase that is shared by the members of the group. Each such member must Purchase individually.
|
||||
|
||||
d. No Obscene or Unlawful Use. You may NOT use 3D Models for any defamatory, harassing, pornographic, obscene, or racist purpose, or to infringe any party's Depicted Intellectual Property rights.
|
||||
|
||||
e. False Attribution. You may NOT misrepresent yourself as the creator of 3D Models.
|
||||
|
||||
11. Resellers. The license granted herein is wholly transferable by an authorized reseller ("Reseller") to another party ("Transferee"). Each transferred license must be transferred entirely and all transferred 3D Models must be permanently deleted from the Reseller's systems after the transfer. When transferring the license, Reseller represents and warrants that the Reseller has the authority to bind the Transferee to these terms. The Reseller is jointly and severally responsible with any Transferee and each are liable for the transferee's use and compliance with TurboSquid's Terms of Use and Site's policies and procedures as well as any financial obligations hereunder.
|
||||
III. License Term & Termination
|
||||
|
||||
1. Term. Your right and license to 3D Models is perpetual, unless terminated as described herein.
|
||||
|
||||
2. Termination. Your license grant is terminated immediately and without notice in the cases below. In such termination, you and any recipients of 3D Models must cease use, distribution, and destroy all copies of 3D Models.
|
||||
|
||||
a. Reversal of Purchase. Your right and license to 3D Models are contingent on your Purchase of 3D Models. Any payment reversal of a Purchase for any reason immediately terminates all rights granted under this agreement. Potential Reasons for a payment reversal include:
|
||||
|
||||
i. TurboSquid reverses your Purchase at your request.
|
||||
|
||||
ii. TurboSquid receives a charge back or other notice from your bank or credit card cancelling your Purchase and/or withdrawing the funds used for your Purchase.
|
||||
|
||||
iii. TurboSquid determines in its sole discretion that your Purchase was fraudulent.
|
||||
|
||||
iv. When you are granted delayed payment terms, and fail to make payments such that TurboSquid sends you notice and terminates your account.
|
||||
|
||||
b. Failure to Abide by the License Grant. Material failure to abide by the terms of this agreement immediately terminates your right and license to 3D Models. If you detect a violation of the license grant by you or any recipient of shared 3D Models, and promptly report the violation to agent@turbosquid.com, TurboSquid will make a good faith effort to find an appropriate remedy to preserve your license grant.
|
||||
IV. Warranties
|
||||
|
||||
You covenant, represent, and warrant to TurboSquid that:
|
||||
|
||||
You have full right, power, legal capacity, and authority to enter into and perform this agreement, have obtained any third-party consent needed to do so, and, prior to any Purchase, had an opportunity to seek independent legal counsel.
|
||||
You will not use 3D Models except pursuant to the terms of this agreement. Should you use 3D Models in an unauthorized way, you agree to any reasonable fee or penalty exercised by TurboSquid under this agreement or applicable law.
|
||||
You will, prior to Purchase, determine the need for and, if appropriate, obtain any needed third-party clearance, consent, or release to use Depicted Intellectual Property shown in the digital rendering of 3D Models, and shall not use 3D Models to infringe any party's Depicted Intellectual Property rights.
|
||||
You will immediately notify TurboSquid of any legal claim or challenge against your use of 3D Models or any other rights issue, before disclosing such issue to any third-party.
|
||||
|
||||
V. Limitation of Liability
|
||||
|
||||
1. 3D Models are provided on an "as is", "as available", and "with all faults" basis. TurboSquid makes no representations, warranties, conditions, or guarantees as to the usefulness, quality, suitability, truth, fitness for a particular purpose, non-infringement, merchantability, or cosmetic attributes of 3D Models, and does not guarantee the accuracy or completeness of specifications associated with 3D Models, including measurements, weight, durability, strength, materials, general physical properties, regulatory compliance, other engineering or construction attributes.
|
||||
|
||||
2. TurboSquid disclaims all express or implied conditions, representations, and warranties of any kind regarding 3D Models, including any implied warranty or condition of merchantability. TurboSquid allows your Purchase to be refunded under certain reasonable time frames and conditions, subject to the Site's policies.
|
||||
|
||||
3. You assume all risk for any damage to your computer systems and network for any damage to your computer system by obtaining 3D Models, including any damages resulting from computer viruses.
|
||||
|
||||
4. To the fullest extent permitted by law, TurboSquid shall not be liable for (A) any direct, indirect, punitive, special, incidental, consequential, or exemplary damages (including loss of business, revenue, profits, goodwill, use, data, electronically transmitted orders, or other economic advantage) arising out of or in connection with 3D Models, even if TurboSquid has previously been advised of, or reasonably could have foreseen, the possibility of such damages, however they arise, whether in breach of contract or in tort (including negligence) or (B) any damages in excess of $1,000. To the extent that any jurisdiction does not allow the exclusion or limitation of direct, incidental, or consequential damages, portions of the preceding limitation or exclusion may not apply, but should be construed to the greatest extent applicable in such jurisdictions. Notwithstanding anything to the contrary herein, the TurboSquid indemnification obligation set forth below shall be limited to the following depending on the licensing tier:
|
||||
|
||||
Tier 0: 3D Models acquired at free-of-charge are not indemnified.
|
||||
|
||||
Tier 1: Standard License indemnity limitation is ten thousand ($10,000) dollars for all 3D Models acquired with payment. This indemnity is in aggregate for all 3D Models acquired under the Standard License.
|
||||
|
||||
Tier 2: Small Business License indemnity limitation is two hundred and fifty thousand ($250,000) dollars for any 3D Model. This indemnity is in aggregate for all 3D Models acquired under the Small Business License.
|
||||
|
||||
Tier 3: Enterprise License indemnity limitation is one million ($1,000,000) dollars for any 3D Model. This indemnity is in aggregate for all 3D Models acquired under the Enterprise License.
|
||||
|
||||
For any 3D Model labeled Editorial, the above indemnities shall only apply if the model is properly used within the editorial license set forth herein (i.e. for news and editorial purposes in association with newsworthy media.) For use outside the Editorial scope, no indemnification from TurboSquid shall apply.
|
||||
|
||||
5. You agree to indemnify and hold TurboSquid and its subsidiaries, affiliates, shareholders, officers, directors, agents, licensors, licensee, suppliers, alliance members, other partners, employees and representatives ("TurboSquid Parties") harmless from any claim or demand, including reasonable attorneys' fees, made by any third party due to, or arising out of your use of 3D Models or Creations.
|
||||
|
||||
6. Subject to sections 4 and 5 above, TurboSquid shall indemnify, defend, and hold you harmless from and against any claim or demand, including reasonable attorneys' fees made by any third party for copyright or trademark infringement due to or arising out of your use of the 3D Models in accordance with these Terms, but excluding any modifications made by You, if such infringement was caused by the modification. This indemnity shall not apply to any 3D Model labeled for Editorial Use or a brand name, logo, or other Depicted Intellectual Property prior identified in a 3D Model.
|
||||
|
||||
7. In the event of an indemnification claim by You, you agree to provide notice to TurboSquid within thirty days' of receiving any claim and allowing TurboSquid to fully control such claim, including but not limited to, selection of counsel, reasonable diligence into the claim, and if necessary litigation and/or settlement. Notice must be given via email to: agent@turbosquid.com. Notice is not considered made until it is acknowledged in writing by TurboSquid.
|
||||
VI. Other Terms
|
||||
|
||||
1. Entire Agreement. This agreement constitutes the entire agreement between you and TurboSquid relating to your Purchase, unless you have a corporate license agreement with TurboSquid. Corporate licenses are available with additional protections for additional fees. Please contact enterprise@turbosquid.com if your organization requires a corporate license. TurboSquid does not otherwise offer any other changes, additions, variations, or additional signed forms related to this agreement. No modification to this agreement will be binding, unless in writing and signed by an authorized TurboSquid representative.
|
||||
|
||||
2. Material Breach and Injunction.
|
||||
|
||||
Your rights hereunder vary by licensing tier as follows:
|
||||
|
||||
For the Standard License, you agree that any material breach of these Terms will result in irreparable harm to TurboSquid for which damages would be an inadequate remedy and, therefore, in addition to its rights and remedies otherwise available at law, TurboSquid will be entitled to equitable relief, including both a preliminary and permanent injunction, if such a breach occurs. You waive any requirement for the posting of a bond or other security if TurboSquid seeks such an injunction.
|
||||
|
||||
For the Enterprise License, TurboSquid may not seek injunctive relief hereunder for any 3D Model. It hereby waives all right to equitable and injunctive relief and its damages shall be limited to monetary damages.
|
||||
|
||||
Notwithstanding anything to the contrary herein, TurboSquid would be irreparably harmed and shall be entitled to equitable relief including injunctive relief for any hacking, theft, or misuse of the Site.
|
||||
|
||||
3. Import/Export Regulations. 3D Models may be subject to the U.S. export laws and the export or import laws of other countries. You agree to comply strictly with all such laws and, in particular, shall with 3D Models: (a) obtain any export, re-export, or import authorizations required by U.S. or Your local laws; (b) not design, develop or produce missile, chemical/biological, or nuclear weaponry; and (c) not provide 3D Models to prohibited countries and entities identified in the U.S. export regulations.
|
||||
|
||||
4. Governing Law. This agreement is governed by New York law, excluding conflict of law principles. Any action or proceeding arising out of or related to this agreement must be brought in a state or federal court located in New York, New York, and both parties irrevocably submit to the exclusive jurisdiction of such courts. All notices, requests and other communications under this agreement must be in writing (e-mail messages shall be deemed writings).
|
||||
|
||||
5. LIMITED INTERNAL USER ARBITRATION. You acknowledge and agree that TurboSquid may, in its sole discretion, arbitrate disputes between TurboSquid users involving 3D Models (including any purchaser or supplier of 3D Models), and such findings shall be final and non-appealable. Either party may request that TurboSquid arbitrate the dispute, or TurboSquid may elect, at its option, to arbitrate the dispute. After TurboSquid elects to arbitrate any dispute hereunder, TurboSquid will waive any rights to a commission from both the Purchase and arbitration, and the parties must keep the results and process confidential and may not disclose anything related to the dispute to any other party (whether by oral, written, or other type of disclosure). To resolve disputes, TurboSquid may decide to terminate or suspend users, revoke the license, offer replacement 3D Models, reestablish the licensee, or surrender or reallocate fees (whether by refund, charitable donation, or otherwise). TurboSquid may award up to 3X the Purchase price to either party depending on the circumstances. YOU UNDERSTAND, ACKNOWLEDGE, AND AGREE THAT ACCEPTING THIS ARBITRATION PROVISION WAIVES RIGHTS TO JUDICIAL RESOLUTION, TRIAL BY JURY AND RIGHTS YOU WOULD OTHERWISE HAVE IF YOU HAD NOT AGREED TO THIS ARBITRATION PROVISION.
|
||||
|
||||
6. Notice. Any notice under this agreement shall be via email to agent@turbosquid.com, provided that you receive an acknowledgement email from a TurboSquid representative within 5 business days. If no such acknowledgement email is received, notice must be in writing and delivered by mail to the following address.
|
||||
|
||||
TurboSquid, Inc.
|
||||
c/o TurboSquid Support
|
||||
935 Gravier St., Suite 1600
|
||||
New Orleans, LA 70112
|
||||
|
||||
7. Assignment. TurboSquid may not assign its rights under this agreement without providing you notice, except in the case of a bankruptcy, merger, acquisition, sale of all or substantially all of TurboSquid's assets to a subsequent owner or operator, or similar event.
|
||||
|
||||
Your assignment rights vary based on the licensing tier of your purchase:
|
||||
|
||||
For the Standard License, you may not assign your rights under this agreement without the prior written consent of TurboSquid.
|
||||
|
||||
For Small Business or Enterprise Licenses, you may assign your rights under this agreement without the notice and consent of TurboSquid.
|
||||
|
||||
8. English. This agreement may be translated into other languages, but English is the official language of this agreement and in any conflict between the English language version and any other version, the English language version shall control.
|
||||
|
||||
9. Publicity. The following advertising, marketing, and publicity rights are granted to TurboSquid for each licensing tier:
|
||||
|
||||
Standard License purchases may be fully publicized by TurboSquid and you hereby grant TurboSquid the right to use you and your company's name, logo, and project name on the TurboSquid website and in its related marketing and advertising materials.
|
||||
|
||||
Small Business and Enterprise License purchase may not be publicized by TurboSquid in any way without prior written permission of the purchaser.
|
||||
|
||||
10. Time limitations on any claim hereunder. Any claim by you hereunder, including without limitation a claim for indemnification under section V must be made within two years of purchasing the 3D Model.
|
||||
|
||||
This 3D Model License is effective for use with 3D Models for use on or after June 17, 2020.
|
||||
+202
@@ -0,0 +1,202 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import numpy as np
|
||||
import torch
|
||||
import imageio
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
# Transform vertex positions to clip space
|
||||
def transform_pos(mtx, pos):
|
||||
t_mtx = torch.from_numpy(mtx).cuda() if isinstance(mtx, np.ndarray) else mtx
|
||||
# (x,y,z) -> (x,y,z,1)
|
||||
posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
|
||||
return torch.matmul(posw, t_mtx.t())[None, ...]
|
||||
|
||||
def render(glctx, mtx, pos, pos_idx, vtx_col, col_idx, resolution: int):
|
||||
pos_clip = transform_pos(mtx, pos)
|
||||
rast_out, _ = dr.rasterize(glctx, pos_clip, pos_idx, resolution=[resolution, resolution])
|
||||
color, _ = dr.interpolate(vtx_col[None, ...], rast_out, col_idx)
|
||||
color = dr.antialias(color, rast_out, pos_clip, pos_idx)
|
||||
return color
|
||||
|
||||
def make_grid(arr, ncols=2):
|
||||
n, height, width, nc = arr.shape
|
||||
nrows = n//ncols
|
||||
assert n == nrows*ncols
|
||||
return arr.reshape(nrows, ncols, height, width, nc).swapaxes(1,2).reshape(height*nrows, width*ncols, nc)
|
||||
|
||||
def fit_cube(max_iter = 5000,
|
||||
resolution = 4,
|
||||
discontinuous = False,
|
||||
repeats = 1,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 512,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
mp4save_interval = None,
|
||||
mp4save_fn = None):
|
||||
|
||||
log_file = None
|
||||
writer = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(f'{out_dir}/{log_fn}', 'wt')
|
||||
if mp4save_interval != 0:
|
||||
writer = imageio.get_writer(f'{out_dir}/{mp4save_fn}', mode='I', fps=30, codec='libx264', bitrate='16M')
|
||||
else:
|
||||
mp4save_interval = None
|
||||
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
fn = 'cube_%s.npz' % ('d' if discontinuous else 'c')
|
||||
with np.load(f'{datadir}/{fn}') as f:
|
||||
pos_idx, vtxp, col_idx, vtxc = f.values()
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], vtxp.shape[0]))
|
||||
|
||||
# Create position/triangle index tensors
|
||||
pos_idx = torch.from_numpy(pos_idx.astype(np.int32)).cuda()
|
||||
col_idx = torch.from_numpy(col_idx.astype(np.int32)).cuda()
|
||||
vtx_pos = torch.from_numpy(vtxp.astype(np.float32)).cuda()
|
||||
vtx_col = torch.from_numpy(vtxc.astype(np.float32)).cuda()
|
||||
|
||||
# Rasterizer context
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
# Repeats.
|
||||
for rep in range(repeats):
|
||||
|
||||
ang = 0.0
|
||||
gl_avg = []
|
||||
|
||||
vtx_pos_rand = np.random.uniform(-0.5, 0.5, size=vtxp.shape) + vtxp
|
||||
vtx_col_rand = np.random.uniform(0.0, 1.0, size=vtxc.shape)
|
||||
vtx_pos_opt = torch.tensor(vtx_pos_rand, dtype=torch.float32, device='cuda', requires_grad=True)
|
||||
vtx_col_opt = torch.tensor(vtx_col_rand, dtype=torch.float32, device='cuda', requires_grad=True)
|
||||
|
||||
# Adam optimizer for vertex position and color with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([vtx_pos_opt, vtx_col_opt], lr=1e-2)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: max(0.01, 10**(-x*0.0005)))
|
||||
|
||||
for it in range(max_iter + 1):
|
||||
# Random rotation/translation matrix for optimization.
|
||||
r_rot = util.random_rotation_translation(0.25)
|
||||
|
||||
# Smooth rotation for display.
|
||||
a_rot = np.matmul(util.rotate_x(0.4), util.rotate_y(ang))
|
||||
|
||||
# Modelview and modelview + projection matrices.
|
||||
proj = util.projection(x=0.4)
|
||||
r_mv = np.matmul(util.translate(0, 0, -3.5), r_rot)
|
||||
r_mvp = np.matmul(proj, r_mv).astype(np.float32)
|
||||
a_mv = np.matmul(util.translate(0, 0, -3.5), a_rot)
|
||||
a_mvp = np.matmul(proj, a_mv).astype(np.float32)
|
||||
|
||||
# Compute geometric error for logging.
|
||||
with torch.no_grad():
|
||||
geom_loss = torch.mean(torch.sum((torch.abs(vtx_pos_opt) - .5)**2, dim=1)**0.5)
|
||||
gl_avg.append(float(geom_loss))
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
gl_val = np.mean(np.asarray(gl_avg))
|
||||
gl_avg = []
|
||||
s = ("rep=%d," % rep) if repeats > 1 else ""
|
||||
s += "iter=%d,err=%f" % (it, gl_val)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + "\n")
|
||||
|
||||
color = render(glctx, r_mvp, vtx_pos, pos_idx, vtx_col, col_idx, resolution)
|
||||
color_opt = render(glctx, r_mvp, vtx_pos_opt, pos_idx, vtx_col_opt, col_idx, resolution)
|
||||
|
||||
# Compute loss and train.
|
||||
loss = torch.mean((color - color_opt)**2) # L2 pixel loss.
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Show/save image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_mp4 = mp4save_interval and (it % mp4save_interval == 0)
|
||||
|
||||
if display_image or save_mp4:
|
||||
ang = ang + 0.01
|
||||
|
||||
img_b = color[0].cpu().numpy()[::-1]
|
||||
img_o = color_opt[0].detach().cpu().numpy()[::-1]
|
||||
img_d = render(glctx, a_mvp, vtx_pos_opt, pos_idx, vtx_col_opt, col_idx, display_res)[0]
|
||||
img_r = render(glctx, a_mvp, vtx_pos, pos_idx, vtx_col, col_idx, display_res)[0]
|
||||
|
||||
scl = display_res // img_o.shape[0]
|
||||
img_b = np.repeat(np.repeat(img_b, scl, axis=0), scl, axis=1)
|
||||
img_o = np.repeat(np.repeat(img_o, scl, axis=0), scl, axis=1)
|
||||
result_image = make_grid(np.stack([img_o, img_b, img_d.detach().cpu().numpy()[::-1], img_r.cpu().numpy()[::-1]]))
|
||||
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='%d / %d' % (it, max_iter))
|
||||
if save_mp4:
|
||||
writer.append_data(np.clip(np.rint(result_image*255.0), 0, 255).astype(np.uint8))
|
||||
|
||||
# Done.
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Cube fit example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--discontinuous', action='store_true', default=False)
|
||||
parser.add_argument('--resolution', type=int, default=0, required=True)
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--mp4save-interval', type=int, default=100)
|
||||
parser.add_argument('--max-iter', type=int, default=1000)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
ds = 'd' if args.discontinuous else 'c'
|
||||
out_dir = f'{args.outdir}/cube_{ds}_{args.resolution}'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_cube(
|
||||
max_iter=args.max_iter,
|
||||
resolution=args.resolution,
|
||||
discontinuous=args.discontinuous,
|
||||
log_interval=10,
|
||||
display_interval=args.display_interval,
|
||||
out_dir=out_dir,
|
||||
log_fn='log.txt',
|
||||
mp4save_interval=args.mp4save_interval,
|
||||
mp4save_fn='progress.mp4',
|
||||
)
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+205
@@ -0,0 +1,205 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Helpers.
|
||||
|
||||
def transform_pos(mtx, pos):
|
||||
t_mtx = torch.from_numpy(mtx).cuda() if isinstance(mtx, np.ndarray) else mtx
|
||||
posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
|
||||
return torch.matmul(posw, t_mtx.t())[None, ...]
|
||||
|
||||
def render(glctx, mtx, pos, pos_idx, uv, uv_idx, tex, resolution, enable_mip, max_mip_level):
|
||||
pos_clip = transform_pos(mtx, pos)
|
||||
rast_out, rast_out_db = dr.rasterize(glctx, pos_clip, pos_idx, resolution=[resolution, resolution])
|
||||
|
||||
if enable_mip:
|
||||
texc, texd = dr.interpolate(uv[None, ...], rast_out, uv_idx, rast_db=rast_out_db, diff_attrs='all')
|
||||
color = dr.texture(tex[None, ...], texc, texd, filter_mode='linear-mipmap-linear', max_mip_level=max_mip_level)
|
||||
else:
|
||||
texc, _ = dr.interpolate(uv[None, ...], rast_out, uv_idx)
|
||||
color = dr.texture(tex[None, ...], texc, filter_mode='linear')
|
||||
|
||||
color = color * torch.clamp(rast_out[..., -1:], 0, 1) # Mask out background.
|
||||
return color
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def fit_earth(max_iter = 20000,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 1024,
|
||||
enable_mip = True,
|
||||
res = 512,
|
||||
ref_res = 2048, # Dropped from 4096 to 2048 to allow using the Cuda rasterizer.
|
||||
lr_base = 1e-2,
|
||||
lr_ramp = 0.1,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
texsave_interval = None,
|
||||
texsave_fn = None,
|
||||
imgsave_interval = None,
|
||||
imgsave_fn = None):
|
||||
|
||||
log_file = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(out_dir + '/' + log_fn, 'wt')
|
||||
else:
|
||||
imgsave_interval, texsave_interval = None, None
|
||||
|
||||
# Mesh and texture adapted from "3D Earth Photorealistic 2K" model at
|
||||
# https://www.turbosquid.com/3d-models/3d-realistic-earth-photorealistic-2k-1279125
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
with np.load(f'{datadir}/earth.npz') as f:
|
||||
pos_idx, pos, uv_idx, uv, tex = f.values()
|
||||
tex = tex.astype(np.float32)/255.0
|
||||
max_mip_level = 9 # Texture is a 4x3 atlas of 512x512 maps.
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], pos.shape[0]))
|
||||
|
||||
# Some input geometry contains vertex positions in (N, 4) (with v[:,3]==1). Drop
|
||||
# the last column in that case.
|
||||
if pos.shape[1] == 4: pos = pos[:, 0:3]
|
||||
|
||||
# Create position/triangle index tensors
|
||||
pos_idx = torch.from_numpy(pos_idx.astype(np.int32)).cuda()
|
||||
vtx_pos = torch.from_numpy(pos.astype(np.float32)).cuda()
|
||||
uv_idx = torch.from_numpy(uv_idx.astype(np.int32)).cuda()
|
||||
vtx_uv = torch.from_numpy(uv.astype(np.float32)).cuda()
|
||||
|
||||
tex = torch.from_numpy(tex.astype(np.float32)).cuda()
|
||||
tex_opt = torch.full(tex.shape, 0.2, device='cuda', requires_grad=True)
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
ang = 0.0
|
||||
|
||||
# Adam optimizer for texture with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([tex_opt], lr=lr_base)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: lr_ramp**(float(x)/float(max_iter)))
|
||||
|
||||
# Render.
|
||||
ang = 0.0
|
||||
texloss_avg = []
|
||||
for it in range(max_iter + 1):
|
||||
# Random rotation/translation matrix for optimization.
|
||||
r_rot = util.random_rotation_translation(0.25)
|
||||
|
||||
# Smooth rotation for display.
|
||||
a_rot = np.matmul(util.rotate_x(0.4), util.rotate_y(ang))
|
||||
dist = np.random.uniform(0.0, 48.5)
|
||||
|
||||
# Modelview and modelview + projection matrices.
|
||||
proj = util.projection(x=0.4, n=1.0, f=200.0)
|
||||
r_mv = np.matmul(util.translate(0, 0, -1.5-dist), r_rot)
|
||||
r_mvp = np.matmul(proj, r_mv).astype(np.float32)
|
||||
a_mv = np.matmul(util.translate(0, 0, -3.5), a_rot)
|
||||
a_mvp = np.matmul(proj, a_mv).astype(np.float32)
|
||||
|
||||
# Measure texture-space RMSE loss
|
||||
with torch.no_grad():
|
||||
texmask = torch.zeros_like(tex)
|
||||
tr = tex.shape[1]//4
|
||||
texmask[tr+13:2*tr-13, 25:-25, :] += 1.0
|
||||
texmask[25:-25, tr+13:2*tr-13, :] += 1.0
|
||||
# Measure only relevant portions of texture when calculating texture
|
||||
# PSNR.
|
||||
texloss = (torch.sum(texmask * (tex - tex_opt)**2)/torch.sum(texmask))**0.5 # RMSE within masked area.
|
||||
texloss_avg.append(float(texloss))
|
||||
|
||||
# Render reference and optimized frames. Always enable mipmapping for reference.
|
||||
color = render(glctx, r_mvp, vtx_pos, pos_idx, vtx_uv, uv_idx, tex, ref_res, True, max_mip_level)
|
||||
color_opt = render(glctx, r_mvp, vtx_pos, pos_idx, vtx_uv, uv_idx, tex_opt, res, enable_mip, max_mip_level)
|
||||
|
||||
# Reduce the reference to correct size.
|
||||
while color.shape[1] > res:
|
||||
color = util.bilinear_downsample(color)
|
||||
|
||||
# Compute loss and perform a training step.
|
||||
loss = torch.mean((color - color_opt)**2) # L2 pixel loss.
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
texloss_val = np.mean(np.asarray(texloss_avg))
|
||||
texloss_avg = []
|
||||
psnr = -10.0 * np.log10(texloss_val**2) # PSNR based on average RMSE.
|
||||
s = "iter=%d,loss=%f,psnr=%f" % (it, texloss_val, psnr)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + '\n')
|
||||
|
||||
# Show/save image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_image = imgsave_interval and (it % imgsave_interval == 0)
|
||||
save_texture = texsave_interval and (it % texsave_interval) == 0
|
||||
|
||||
if display_image or save_image:
|
||||
ang = ang + 0.1
|
||||
|
||||
with torch.no_grad():
|
||||
result_image = render(glctx, a_mvp, vtx_pos, pos_idx, vtx_uv, uv_idx, tex_opt, res, enable_mip, max_mip_level)[0].cpu().numpy()[::-1]
|
||||
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='%d / %d' % (it, max_iter))
|
||||
if save_image:
|
||||
util.save_image(out_dir + '/' + (imgsave_fn % it), result_image)
|
||||
|
||||
if save_texture:
|
||||
texture = tex_opt.cpu().numpy()[::-1]
|
||||
util.save_image(out_dir + '/' + (texsave_fn % it), texture)
|
||||
|
||||
|
||||
# Done.
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Earth texture fitting example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--mip', help='enable mipmapping', action='store_true', default=False)
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--max-iter', type=int, default=10000)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
ms = 'mip' if args.mip else 'nomip'
|
||||
out_dir = f'{args.outdir}/earth_{ms}'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_earth(max_iter=args.max_iter, log_interval=10, display_interval=args.display_interval, enable_mip=args.mip, out_dir=out_dir, log_fn='log.txt', texsave_interval=1000, texsave_fn='tex_%06d.png', imgsave_interval=1000, imgsave_fn='img_%06d.png')
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+227
@@ -0,0 +1,227 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import numpy as np
|
||||
import torch
|
||||
import os
|
||||
import pathlib
|
||||
import imageio
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Environment map and Phong BRDF learning.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def fit_env_phong(max_iter = 1000,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 1024,
|
||||
res = 1024,
|
||||
lr_base = 1e-2,
|
||||
lr_ramp = 1.0,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
mp4save_interval = None,
|
||||
mp4save_fn = None):
|
||||
|
||||
log_file = None
|
||||
writer = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(out_dir + '/' + log_fn, 'wt')
|
||||
if mp4save_interval != 0:
|
||||
writer = imageio.get_writer(f'{out_dir}/{mp4save_fn}', mode='I', fps=30, codec='libx264', bitrate='16M')
|
||||
else:
|
||||
mp4save_interval = None
|
||||
|
||||
# Texture adapted from https://github.com/WaveEngine/Samples/tree/master/Materials/EnvironmentMap/Content/Assets/CubeMap.cubemap
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
with np.load(f'{datadir}/envphong.npz') as f:
|
||||
pos_idx, pos, normals, env = f.values()
|
||||
env = env.astype(np.float32)/255.0
|
||||
env = np.stack(env)[:, ::-1].copy()
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], pos.shape[0]))
|
||||
|
||||
# Move all the stuff to GPU.
|
||||
pos_idx = torch.as_tensor(pos_idx, dtype=torch.int32, device='cuda')
|
||||
pos = torch.as_tensor(pos, dtype=torch.float32, device='cuda')
|
||||
normals = torch.as_tensor(normals, dtype=torch.float32, device='cuda')
|
||||
env = torch.as_tensor(env, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Target Phong parameters.
|
||||
phong_rgb = np.asarray([1.0, 0.8, 0.6], np.float32)
|
||||
phong_exp = 25.0
|
||||
phong_rgb_t = torch.as_tensor(phong_rgb, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Learned variables: environment maps, phong color, phong exponent.
|
||||
env_var = torch.ones_like(env) * .5
|
||||
env_var.requires_grad_()
|
||||
phong_var_raw = torch.as_tensor(np.random.uniform(size=[4]), dtype=torch.float32, device='cuda')
|
||||
phong_var_raw.requires_grad_()
|
||||
phong_var_mul = torch.as_tensor([1.0, 1.0, 1.0, 10.0], dtype=torch.float32, device='cuda')
|
||||
|
||||
# Render.
|
||||
ang = 0.0
|
||||
imgloss_avg, phong_avg = [], []
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
zero_tensor = torch.as_tensor(0.0, dtype=torch.float32, device='cuda')
|
||||
one_tensor = torch.as_tensor(1.0, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Adam optimizer for environment map and phong with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([env_var, phong_var_raw], lr=lr_base)
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: lr_ramp**(float(x)/float(max_iter)))
|
||||
|
||||
for it in range(max_iter + 1):
|
||||
phong_var = phong_var_raw * phong_var_mul
|
||||
|
||||
# Random rotation/translation matrix for optimization.
|
||||
r_rot = util.random_rotation_translation(0.25)
|
||||
|
||||
# Smooth rotation for display.
|
||||
ang = ang + 0.01
|
||||
a_rot = np.matmul(util.rotate_x(0.4), util.rotate_y(ang))
|
||||
|
||||
# Modelview and modelview + projection matrices.
|
||||
proj = util.projection(x=0.4, n=1.0, f=200.0)
|
||||
r_mv = np.matmul(util.translate(0, 0, -3.5), r_rot)
|
||||
r_mvp = np.matmul(proj, r_mv).astype(np.float32)
|
||||
a_mv = np.matmul(util.translate(0, 0, -3.5), a_rot)
|
||||
a_mvp = np.matmul(proj, a_mv).astype(np.float32)
|
||||
a_mvc = a_mvp
|
||||
r_mvp = torch.as_tensor(r_mvp, dtype=torch.float32, device='cuda')
|
||||
a_mvp = torch.as_tensor(a_mvp, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Solve camera positions.
|
||||
a_campos = torch.as_tensor(np.linalg.inv(a_mv)[:3, 3], dtype=torch.float32, device='cuda')
|
||||
r_campos = torch.as_tensor(np.linalg.inv(r_mv)[:3, 3], dtype=torch.float32, device='cuda')
|
||||
|
||||
# Random light direction.
|
||||
lightdir = np.random.normal(size=[3])
|
||||
lightdir /= np.linalg.norm(lightdir) + 1e-8
|
||||
lightdir = torch.as_tensor(lightdir, dtype=torch.float32, device='cuda')
|
||||
|
||||
def render_refl(ldir, cpos, mvp):
|
||||
# Transform and rasterize.
|
||||
viewvec = pos[..., :3] - cpos[np.newaxis, np.newaxis, :] # View vectors at vertices.
|
||||
reflvec = viewvec - 2.0 * normals[np.newaxis, ...] * torch.sum(normals[np.newaxis, ...] * viewvec, -1, keepdim=True) # Reflection vectors at vertices.
|
||||
reflvec = reflvec / torch.sum(reflvec**2, -1, keepdim=True)**0.5 # Normalize.
|
||||
pos_clip = torch.matmul(pos, mvp.t())[np.newaxis, ...]
|
||||
rast_out, rast_out_db = dr.rasterize(glctx, pos_clip, pos_idx, [res, res])
|
||||
refl, refld = dr.interpolate(reflvec, rast_out, pos_idx, rast_db=rast_out_db, diff_attrs='all') # Interpolated reflection vectors.
|
||||
|
||||
# Phong light.
|
||||
refl = refl / (torch.sum(refl**2, -1, keepdim=True) + 1e-8)**0.5 # Normalize.
|
||||
ldotr = torch.sum(-ldir * refl, -1, keepdim=True) # L dot R.
|
||||
|
||||
# Return
|
||||
return refl, refld, ldotr, (rast_out[..., -1:] == 0)
|
||||
|
||||
# Render the reflections.
|
||||
refl, refld, ldotr, mask = render_refl(lightdir, r_campos, r_mvp)
|
||||
|
||||
# Reference color. No need for AA because we are not learning geometry.
|
||||
color = dr.texture(env[np.newaxis, ...], refl, uv_da=refld, filter_mode='linear-mipmap-linear', boundary_mode='cube')
|
||||
color = color + phong_rgb_t * torch.max(zero_tensor, ldotr) ** phong_exp # Phong.
|
||||
color = torch.where(mask, one_tensor, color) # White background.
|
||||
|
||||
# Candidate rendering same up to this point, but uses learned texture and Phong parameters instead.
|
||||
color_opt = dr.texture(env_var[np.newaxis, ...], refl, uv_da=refld, filter_mode='linear-mipmap-linear', boundary_mode='cube')
|
||||
color_opt = color_opt + phong_var[:3] * torch.max(zero_tensor, ldotr) ** phong_var[3] # Phong.
|
||||
color_opt = torch.where(mask, one_tensor, color_opt) # White background.
|
||||
|
||||
# Compute loss and train.
|
||||
loss = torch.mean((color - color_opt)**2) # L2 pixel loss.
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
|
||||
# Collect losses.
|
||||
imgloss_avg.append(loss.detach().cpu().numpy())
|
||||
phong_avg.append(phong_var.detach().cpu().numpy())
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
imgloss_val, imgloss_avg = np.mean(np.asarray(imgloss_avg, np.float32)), []
|
||||
phong_val, phong_avg = np.mean(np.asarray(phong_avg, np.float32), axis=0), []
|
||||
phong_rgb_rmse = np.mean((phong_val[:3] - phong_rgb)**2)**0.5
|
||||
phong_exp_rel_err = np.abs(phong_val[3] - phong_exp)/phong_exp
|
||||
s = "iter=%d,phong_rgb_rmse=%f,phong_exp_rel_err=%f,img_rmse=%f" % (it, phong_rgb_rmse, phong_exp_rel_err, imgloss_val)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + '\n')
|
||||
|
||||
# Show/save result image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_mp4 = mp4save_interval and (it % mp4save_interval == 0)
|
||||
|
||||
if display_image or save_mp4:
|
||||
lightdir = np.asarray([.8, -1., .5, 0.0])
|
||||
lightdir = np.matmul(a_mvc, lightdir)[:3]
|
||||
lightdir /= np.linalg.norm(lightdir)
|
||||
lightdir = torch.as_tensor(lightdir, dtype=torch.float32, device='cuda')
|
||||
refl, refld, ldotr, mask = render_refl(lightdir, a_campos, a_mvp)
|
||||
color_opt = dr.texture(env_var[np.newaxis, ...], refl, uv_da=refld, filter_mode='linear-mipmap-linear', boundary_mode='cube')
|
||||
color_opt = color_opt + phong_var[:3] * torch.max(zero_tensor, ldotr) ** phong_var[3]
|
||||
color_opt = torch.where(mask, one_tensor, color_opt)
|
||||
result_image = color_opt.detach()[0].cpu().numpy()[::-1]
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='%d / %d' % (it, max_iter))
|
||||
if save_mp4:
|
||||
writer.append_data(np.clip(np.rint(result_image*255.0), 0, 255).astype(np.uint8))
|
||||
|
||||
# Done.
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Main function.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Environment map fitting example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--mp4save-interval', type=int, default=10)
|
||||
parser.add_argument('--max-iter', type=int, default=5000)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
out_dir = f'{args.outdir}/env_phong'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_env_phong(
|
||||
max_iter=args.max_iter,
|
||||
log_interval=100,
|
||||
display_interval=args.display_interval,
|
||||
out_dir=out_dir,
|
||||
mp4save_interval=args.mp4save_interval,
|
||||
mp4save_fn='progress.mp4',
|
||||
)
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
+290
@@ -0,0 +1,290 @@
|
||||
# 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.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pathlib
|
||||
import numpy as np
|
||||
import torch
|
||||
import imageio
|
||||
|
||||
import util
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Quaternion math.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
# Unit quaternion.
|
||||
def q_unit():
|
||||
return np.asarray([1, 0, 0, 0], np.float32)
|
||||
|
||||
# Get a random normalized quaternion.
|
||||
def q_rnd():
|
||||
u, v, w = np.random.uniform(0.0, 1.0, size=[3])
|
||||
v *= 2.0 * np.pi
|
||||
w *= 2.0 * np.pi
|
||||
return np.asarray([(1.0-u)**0.5 * np.sin(v), (1.0-u)**0.5 * np.cos(v), u**0.5 * np.sin(w), u**0.5 * np.cos(w)], np.float32)
|
||||
|
||||
# Get a random quaternion from the octahedral symmetric group S_4.
|
||||
_r2 = 0.5**0.5
|
||||
_q_S4 = [[ 1.0, 0.0, 0.0, 0.0], [ 0.0, 1.0, 0.0, 0.0], [ 0.0, 0.0, 1.0, 0.0], [ 0.0, 0.0, 0.0, 1.0],
|
||||
[-0.5, 0.5, 0.5, 0.5], [-0.5,-0.5,-0.5, 0.5], [ 0.5,-0.5, 0.5, 0.5], [ 0.5, 0.5,-0.5, 0.5],
|
||||
[ 0.5, 0.5, 0.5, 0.5], [-0.5, 0.5,-0.5, 0.5], [ 0.5,-0.5,-0.5, 0.5], [-0.5,-0.5, 0.5, 0.5],
|
||||
[ _r2,-_r2, 0.0, 0.0], [ _r2, _r2, 0.0, 0.0], [ 0.0, 0.0, _r2, _r2], [ 0.0, 0.0,-_r2, _r2],
|
||||
[ 0.0, _r2, _r2, 0.0], [ _r2, 0.0, 0.0,-_r2], [ _r2, 0.0, 0.0, _r2], [ 0.0,-_r2, _r2, 0.0],
|
||||
[ _r2, 0.0, _r2, 0.0], [ 0.0, _r2, 0.0, _r2], [ _r2, 0.0,-_r2, 0.0], [ 0.0,-_r2, 0.0, _r2]]
|
||||
def q_rnd_S4():
|
||||
return np.asarray(_q_S4[np.random.randint(24)], np.float32)
|
||||
|
||||
# Quaternion slerp.
|
||||
def q_slerp(p, q, t):
|
||||
d = np.dot(p, q)
|
||||
if d < 0.0:
|
||||
q = -q
|
||||
d = -d
|
||||
if d > 0.999:
|
||||
a = p + t * (q-p)
|
||||
return a / np.linalg.norm(a)
|
||||
t0 = np.arccos(d)
|
||||
tt = t0 * t
|
||||
st = np.sin(tt)
|
||||
st0 = np.sin(t0)
|
||||
s1 = st / st0
|
||||
s0 = np.cos(tt) - d*s1
|
||||
return s0*p + s1*q
|
||||
|
||||
# Quaterion scale (slerp vs. identity quaternion).
|
||||
def q_scale(q, scl):
|
||||
return q_slerp(q_unit(), q, scl)
|
||||
|
||||
# Quaternion product.
|
||||
def q_mul(p, q):
|
||||
s1, V1 = p[0], p[1:]
|
||||
s2, V2 = q[0], q[1:]
|
||||
s = s1*s2 - np.dot(V1, V2)
|
||||
V = s1*V2 + s2*V1 + np.cross(V1, V2)
|
||||
return np.asarray([s, V[0], V[1], V[2]], np.float32)
|
||||
|
||||
# Angular difference between two quaternions in degrees.
|
||||
def q_angle_deg(p, q):
|
||||
p = p.detach().cpu().numpy()
|
||||
q = q.detach().cpu().numpy()
|
||||
d = np.abs(np.dot(p, q))
|
||||
d = min(d, 1.0)
|
||||
return np.degrees(2.0 * np.arccos(d))
|
||||
|
||||
# Quaternion product
|
||||
def q_mul_torch(p, q):
|
||||
a = p[0]*q[0] - p[1]*q[1] - p[2]*q[2] - p[3]*q[3]
|
||||
b = p[0]*q[1] + p[1]*q[0] + p[2]*q[3] - p[3]*q[2]
|
||||
c = p[0]*q[2] + p[2]*q[0] + p[3]*q[1] - p[1]*q[3]
|
||||
d = p[0]*q[3] + p[3]*q[0] + p[1]*q[2] - p[2]*q[1]
|
||||
return torch.stack([a, b, c, d])
|
||||
|
||||
# Convert quaternion to 4x4 rotation matrix.
|
||||
def q_to_mtx(q):
|
||||
r0 = torch.stack([1.0-2.0*q[1]**2 - 2.0*q[2]**2, 2.0*q[0]*q[1] - 2.0*q[2]*q[3], 2.0*q[0]*q[2] + 2.0*q[1]*q[3]])
|
||||
r1 = torch.stack([2.0*q[0]*q[1] + 2.0*q[2]*q[3], 1.0 - 2.0*q[0]**2 - 2.0*q[2]**2, 2.0*q[1]*q[2] - 2.0*q[0]*q[3]])
|
||||
r2 = torch.stack([2.0*q[0]*q[2] - 2.0*q[1]*q[3], 2.0*q[1]*q[2] + 2.0*q[0]*q[3], 1.0 - 2.0*q[0]**2 - 2.0*q[1]**2])
|
||||
rr = torch.transpose(torch.stack([r0, r1, r2]), 1, 0)
|
||||
rr = torch.cat([rr, torch.tensor([[0], [0], [0]], dtype=torch.float32).cuda()], dim=1) # Pad right column.
|
||||
rr = torch.cat([rr, torch.tensor([[0, 0, 0, 1]], dtype=torch.float32).cuda()], dim=0) # Pad bottom row.
|
||||
return rr
|
||||
|
||||
# Transform vertex positions to clip space
|
||||
def transform_pos(mtx, pos):
|
||||
t_mtx = torch.from_numpy(mtx).cuda() if isinstance(mtx, np.ndarray) else mtx
|
||||
# (x,y,z) -> (x,y,z,1)
|
||||
posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()], axis=1)
|
||||
return torch.matmul(posw, t_mtx.t())[None, ...]
|
||||
|
||||
def render(glctx, mtx, pos, pos_idx, col, col_idx, resolution: int):
|
||||
# Setup TF graph for reference.
|
||||
pos_clip = transform_pos(mtx, pos)
|
||||
rast_out, _ = dr.rasterize(glctx, pos_clip, pos_idx, resolution=[resolution, resolution])
|
||||
color , _ = dr.interpolate(col[None, ...], rast_out, col_idx)
|
||||
color = dr.antialias(color, rast_out, pos_clip, pos_idx)
|
||||
return color
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Cube pose fitter.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def fit_pose(max_iter = 10000,
|
||||
repeats = 1,
|
||||
log_interval = 10,
|
||||
display_interval = None,
|
||||
display_res = 512,
|
||||
lr_base = 0.01,
|
||||
lr_falloff = 1.0,
|
||||
nr_base = 1.0,
|
||||
nr_falloff = 1e-4,
|
||||
grad_phase_start = 0.5,
|
||||
resolution = 256,
|
||||
out_dir = None,
|
||||
log_fn = None,
|
||||
mp4save_interval = None,
|
||||
mp4save_fn = None):
|
||||
|
||||
log_file = None
|
||||
writer = None
|
||||
if out_dir:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
if log_fn:
|
||||
log_file = open(out_dir + '/' + log_fn, 'wt')
|
||||
if mp4save_interval != 0:
|
||||
writer = imageio.get_writer(f'{out_dir}/{mp4save_fn}', mode='I', fps=30, codec='libx264', bitrate='16M')
|
||||
else:
|
||||
mp4save_interval = None
|
||||
|
||||
datadir = f'{pathlib.Path(__file__).absolute().parents[1]}/data'
|
||||
with np.load(f'{datadir}/cube_p.npz') as f:
|
||||
pos_idx, pos, col_idx, col = f.values()
|
||||
print("Mesh has %d triangles and %d vertices." % (pos_idx.shape[0], pos.shape[0]))
|
||||
|
||||
# Some input geometry contains vertex positions in (N, 4) (with v[:,3]==1). Drop
|
||||
# the last column in that case.
|
||||
if pos.shape[1] == 4: pos = pos[:, 0:3]
|
||||
|
||||
# Create position/triangle index tensors
|
||||
pos_idx = torch.from_numpy(pos_idx.astype(np.int32)).cuda()
|
||||
vtx_pos = torch.from_numpy(pos.astype(np.float32)).cuda()
|
||||
col_idx = torch.from_numpy(col_idx.astype(np.int32)).cuda()
|
||||
vtx_col = torch.from_numpy(col.astype(np.float32)).cuda()
|
||||
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
for rep in range(repeats):
|
||||
pose_target = torch.tensor(q_rnd(), device='cuda')
|
||||
pose_init = q_rnd()
|
||||
pose_opt = torch.tensor(pose_init / np.sum(pose_init**2)**0.5, dtype=torch.float32, device='cuda', requires_grad=True)
|
||||
|
||||
loss_best = np.inf
|
||||
pose_best = pose_opt.detach().clone()
|
||||
|
||||
# Modelview + projection matrix.
|
||||
mvp = torch.tensor(np.matmul(util.projection(x=0.4), util.translate(0, 0, -3.5)).astype(np.float32), device='cuda')
|
||||
|
||||
# Adam optimizer for texture with a learning rate ramp.
|
||||
optimizer = torch.optim.Adam([pose_opt], betas=(0.9, 0.999), lr=lr_base)
|
||||
# Render.
|
||||
for it in range(max_iter + 1):
|
||||
# Set learning rate.
|
||||
itf = 1.0 * it / max_iter
|
||||
nr = nr_base * nr_falloff**itf
|
||||
lr = lr_base * lr_falloff**itf
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
# Noise input.
|
||||
if itf >= grad_phase_start:
|
||||
noise = q_unit()
|
||||
else:
|
||||
noise = q_scale(q_rnd(), nr)
|
||||
noise = q_mul(noise, q_rnd_S4()) # Orientation noise.
|
||||
|
||||
# Render.
|
||||
color = render(glctx, torch.matmul(mvp, q_to_mtx(pose_target)), vtx_pos, pos_idx, vtx_col, col_idx, resolution)
|
||||
pose_total_opt = q_mul_torch(pose_opt, noise)
|
||||
mtx_total_opt = torch.matmul(mvp, q_to_mtx(pose_total_opt))
|
||||
color_opt = render(glctx, mtx_total_opt, vtx_pos, pos_idx, vtx_col, col_idx, resolution)
|
||||
|
||||
# Image-space loss.
|
||||
diff = (color_opt - color)**2 # L2 norm.
|
||||
diff = torch.tanh(5.0 * torch.max(diff, dim=-1)[0])
|
||||
loss = torch.mean(diff)
|
||||
|
||||
# Measure image-space loss and update best found pose.
|
||||
loss_val = float(loss)
|
||||
if (loss_val < loss_best) and (loss_val > 0.0):
|
||||
pose_best = pose_total_opt.detach().clone()
|
||||
loss_best = loss_val
|
||||
if itf < grad_phase_start:
|
||||
with torch.no_grad(): pose_opt[:] = pose_best
|
||||
|
||||
# Print/save log.
|
||||
if log_interval and (it % log_interval == 0):
|
||||
err = q_angle_deg(pose_opt, pose_target)
|
||||
ebest = q_angle_deg(pose_best, pose_target)
|
||||
s = "rep=%d,iter=%d,err=%f,err_best=%f,loss=%f,loss_best=%f,lr=%f,nr=%f" % (rep, it, err, ebest, loss_val, loss_best, lr, nr)
|
||||
print(s)
|
||||
if log_file:
|
||||
log_file.write(s + "\n")
|
||||
|
||||
# Run gradient training step.
|
||||
if itf >= grad_phase_start:
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
with torch.no_grad():
|
||||
pose_opt /= torch.sum(pose_opt**2)**0.5
|
||||
|
||||
# Show/save image.
|
||||
display_image = display_interval and (it % display_interval == 0)
|
||||
save_mp4 = mp4save_interval and (it % mp4save_interval == 0)
|
||||
|
||||
if display_image or save_mp4:
|
||||
img_ref = color[0].detach().cpu().numpy()
|
||||
img_opt = color_opt[0].detach().cpu().numpy()
|
||||
img_best = render(glctx, torch.matmul(mvp, q_to_mtx(pose_best)), vtx_pos, pos_idx, vtx_col, col_idx, resolution)[0].detach().cpu().numpy()
|
||||
result_image = np.concatenate([img_ref, img_best, img_opt], axis=1)[::-1]
|
||||
|
||||
if display_image:
|
||||
util.display_image(result_image, size=display_res, title='(%d) %d / %d' % (rep, it, max_iter))
|
||||
if save_mp4:
|
||||
writer.append_data(np.clip(np.rint(result_image*255.0), 0, 255).astype(np.uint8))
|
||||
|
||||
# Done.
|
||||
if writer is not None:
|
||||
writer.close()
|
||||
if log_file:
|
||||
log_file.close()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Cube pose fitting example')
|
||||
parser.add_argument('--outdir', help='specify output directory', default='')
|
||||
parser.add_argument('--display-interval', type=int, default=0)
|
||||
parser.add_argument('--mp4save-interval', type=int, default=10)
|
||||
parser.add_argument('--max-iter', type=int, default=1000)
|
||||
parser.add_argument('--repeats', type=int, default=1)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set up logging.
|
||||
if args.outdir:
|
||||
out_dir = f'{args.outdir}/pose'
|
||||
print (f'Saving results under {out_dir}')
|
||||
else:
|
||||
out_dir = None
|
||||
print ('No output directory specified, not saving log or images')
|
||||
|
||||
# Run.
|
||||
fit_pose(
|
||||
max_iter=args.max_iter,
|
||||
repeats=args.repeats,
|
||||
log_interval=100,
|
||||
display_interval=args.display_interval,
|
||||
out_dir=out_dir,
|
||||
log_fn='log.txt',
|
||||
mp4save_interval=args.mp4save_interval,
|
||||
mp4save_fn='progress.mp4',
|
||||
)
|
||||
|
||||
# Done.
|
||||
print("Done.")
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,30 @@
|
||||
# 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.
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
def tensor(*args, **kwargs):
|
||||
return torch.tensor(*args, device='cuda', **kwargs)
|
||||
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
|
||||
pos = tensor([[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]], dtype=torch.float32)
|
||||
col = tensor([[[1, 0, 0], [0, 1, 0], [0, 0, 1]]], dtype=torch.float32)
|
||||
tri = tensor([[0, 1, 2]], dtype=torch.int32)
|
||||
|
||||
rast, _ = dr.rasterize(glctx, pos, tri, resolution=[256, 256])
|
||||
out, _ = dr.interpolate(col, rast, tri)
|
||||
|
||||
img = out.cpu().numpy()[0, ::-1, :, :] # Flip vertically.
|
||||
img = np.clip(np.rint(img * 255), 0, 255).astype(np.uint8) # Quantize to np.uint8
|
||||
|
||||
print("Saving to 'tri.png'.")
|
||||
imageio.imsave('tri.png', img)
|
||||
+120
@@ -0,0 +1,120 @@
|
||||
# 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.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Projection and transformation matrix helpers.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def projection(x=0.1, n=1.0, f=50.0):
|
||||
return np.array([[n/x, 0, 0, 0],
|
||||
[ 0, n/x, 0, 0],
|
||||
[ 0, 0, -(f+n)/(f-n), -(2*f*n)/(f-n)],
|
||||
[ 0, 0, -1, 0]]).astype(np.float32)
|
||||
|
||||
def translate(x, y, z):
|
||||
return np.array([[1, 0, 0, x],
|
||||
[0, 1, 0, y],
|
||||
[0, 0, 1, z],
|
||||
[0, 0, 0, 1]]).astype(np.float32)
|
||||
|
||||
def rotate_x(a):
|
||||
s, c = np.sin(a), np.cos(a)
|
||||
return np.array([[1, 0, 0, 0],
|
||||
[0, c, -s, 0],
|
||||
[0, s, c, 0],
|
||||
[0, 0, 0, 1]]).astype(np.float32)
|
||||
|
||||
def rotate_y(a):
|
||||
s, c = np.sin(a), np.cos(a)
|
||||
return np.array([[ c, 0, s, 0],
|
||||
[ 0, 1, 0, 0],
|
||||
[-s, 0, c, 0],
|
||||
[ 0, 0, 0, 1]]).astype(np.float32)
|
||||
|
||||
def random_rotation_translation(t):
|
||||
m = np.random.normal(size=[3, 3])
|
||||
m[1] = np.cross(m[0], m[2])
|
||||
m[2] = np.cross(m[0], m[1])
|
||||
m = m / np.linalg.norm(m, axis=1, keepdims=True)
|
||||
m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
|
||||
m[3, 3] = 1.0
|
||||
m[:3, 3] = np.random.uniform(-t, t, size=[3])
|
||||
return m
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Bilinear downsample by 2x.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def bilinear_downsample(x):
|
||||
w = torch.tensor([[1, 3, 3, 1], [3, 9, 9, 3], [3, 9, 9, 3], [1, 3, 3, 1]], dtype=torch.float32, device=x.device) / 64.0
|
||||
w = w.expand(x.shape[-1], 1, 4, 4)
|
||||
x = torch.nn.functional.conv2d(x.permute(0, 3, 1, 2), w, padding=1, stride=2, groups=x.shape[-1])
|
||||
return x.permute(0, 2, 3, 1)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Image display function using OpenGL.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
_glfw_window = None
|
||||
def display_image(image, zoom=None, size=None, title=None): # HWC
|
||||
# Import OpenGL and glfw.
|
||||
import OpenGL.GL as gl
|
||||
import glfw
|
||||
|
||||
# Zoom image if requested.
|
||||
image = np.asarray(image)
|
||||
if size is not None:
|
||||
assert zoom is None
|
||||
zoom = max(1, size // image.shape[0])
|
||||
if zoom is not None:
|
||||
image = image.repeat(zoom, axis=0).repeat(zoom, axis=1)
|
||||
height, width, channels = image.shape
|
||||
|
||||
# Initialize window.
|
||||
if title is None:
|
||||
title = 'Debug window'
|
||||
global _glfw_window
|
||||
if _glfw_window is None:
|
||||
glfw.init()
|
||||
_glfw_window = glfw.create_window(width, height, title, None, None)
|
||||
glfw.make_context_current(_glfw_window)
|
||||
glfw.show_window(_glfw_window)
|
||||
glfw.swap_interval(0)
|
||||
else:
|
||||
glfw.make_context_current(_glfw_window)
|
||||
glfw.set_window_title(_glfw_window, title)
|
||||
glfw.set_window_size(_glfw_window, width, height)
|
||||
|
||||
# Update window.
|
||||
glfw.poll_events()
|
||||
gl.glClearColor(0, 0, 0, 1)
|
||||
gl.glClear(gl.GL_COLOR_BUFFER_BIT)
|
||||
gl.glWindowPos2f(0, 0)
|
||||
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
|
||||
gl_format = {3: gl.GL_RGB, 2: gl.GL_RG, 1: gl.GL_LUMINANCE}[channels]
|
||||
gl_dtype = {'uint8': gl.GL_UNSIGNED_BYTE, 'float32': gl.GL_FLOAT}[image.dtype.name]
|
||||
gl.glDrawPixels(width, height, gl_format, gl_dtype, image[::-1])
|
||||
glfw.swap_buffers(_glfw_window)
|
||||
if glfw.window_should_close(_glfw_window):
|
||||
return False
|
||||
return True
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Image save helper.
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def save_image(fn, x):
|
||||
import imageio
|
||||
x = np.rint(x * 255.0)
|
||||
x = np.clip(x, 0, 255).astype(np.uint8)
|
||||
imageio.imsave(fn, x)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
Vendored
+55
@@ -0,0 +1,55 @@
|
||||
# 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.
|
||||
|
||||
import setuptools
|
||||
import os
|
||||
|
||||
# Print an error message if there's no PyTorch installed.
|
||||
try:
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
except ImportError:
|
||||
# This happens if the user runs 'pip install' with default build isolation
|
||||
# OR if they simply don't have torch installed at all.
|
||||
print("\n\n" + "*" * 70)
|
||||
print("ERROR! Cannot compile nvdiffrast CUDA extension. Please ensure that:\n")
|
||||
print("1. You have PyTorch installed")
|
||||
print("2. You run 'pip install' with --no-build-isolation flag")
|
||||
print("*" * 70 + "\n\n")
|
||||
exit(1)
|
||||
|
||||
setuptools.setup(
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
"_nvdiffrast_c",
|
||||
sources=[
|
||||
"csrc/common/antialias.cu",
|
||||
"csrc/common/common.cpp",
|
||||
"csrc/common/cudaraster/impl/Buffer.cpp",
|
||||
"csrc/common/cudaraster/impl/CudaRaster.cpp",
|
||||
"csrc/common/cudaraster/impl/RasterImpl.cpp",
|
||||
"csrc/common/cudaraster/impl/RasterImpl_kernel.cu",
|
||||
"csrc/common/interpolate.cu",
|
||||
"csrc/common/rasterize.cu",
|
||||
"csrc/common/texture.cpp",
|
||||
"csrc/common/texture_kernel.cu",
|
||||
"csrc/torch/torch_antialias.cpp",
|
||||
"csrc/torch/torch_bindings.cpp",
|
||||
"csrc/torch/torch_interpolate.cpp",
|
||||
"csrc/torch/torch_rasterize.cpp",
|
||||
"csrc/torch/torch_texture.cpp",
|
||||
],
|
||||
extra_compile_args={
|
||||
"cxx": ["-DNVDR_TORCH"]
|
||||
# Disable warnings in torch headers.
|
||||
+ (["/wd4067", "/wd4624", "/wd4996"] if os.name == "nt" else []),
|
||||
"nvcc": ["-DNVDR_TORCH", "-lineinfo"],
|
||||
},
|
||||
)
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
Reference in New Issue
Block a user