diff --git a/extensions/nvdiffrast/LICENSE.txt b/extensions/nvdiffrast/LICENSE.txt deleted file mode 100644 index 9ef3a30911b6676b2fefa11e4e6044beb5b0bc8e..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/LICENSE.txt +++ /dev/null @@ -1,97 +0,0 @@ -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. - -======================================================================= diff --git a/extensions/nvdiffrast/README.md b/extensions/nvdiffrast/README.md deleted file mode 100644 index c17baacc4209f34d3d74efdbd7cbb080d0451ca9..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/README.md +++ /dev/null @@ -1,42 +0,0 @@ -## Nvdiffrast – Modular Primitives for High-Performance Differentiable Rendering - -![Teaser image](./docs/img/teaser.png) - -**Modular Primitives for High-Performance Differentiable Rendering**
-Samuli Laine, Janne Hellsten, Tero Karras, Yeongho Seol, Jaakko Lehtinen, Timo Aila
-[http://arxiv.org/abs/2011.03277](http://arxiv.org/abs/2011.03277) - -Nvdiffrast is a PyTorch/TensorFlow library that provides high-performance primitive operations for rasterization-based differentiable rendering. -Please refer to ☞☞ [nvdiffrast documentation](https://nvlabs.github.io/nvdiffrast) ☜☜ for more information. - -## Licenses - -Copyright © 2020–2024, 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} -} -``` diff --git a/extensions/nvdiffrast/nvdiffrast/__init__.py b/extensions/nvdiffrast/nvdiffrast/__init__.py deleted file mode 100644 index f3d077dd95e58d52169d5699b6d3ede78a309c21..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# 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. - -__version__ = '0.3.3' diff --git a/extensions/nvdiffrast/nvdiffrast/common/antialias.cu b/extensions/nvdiffrast/nvdiffrast/common/antialias.cu deleted file mode 100644 index a62c1b34bc63788bf2e8f960f8bb2d7d2c9d1d45..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/antialias.cu +++ /dev/null @@ -1,558 +0,0 @@ -// 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); - } -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/antialias.h b/extensions/nvdiffrast/nvdiffrast/common/antialias.h deleted file mode 100644 index bc2fd480e63ca6ad12fc5f4e15f1857880d145c3..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/antialias.h +++ /dev/null @@ -1,50 +0,0 @@ -// 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. -}; - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/common.cpp b/extensions/nvdiffrast/nvdiffrast/common/common.cpp deleted file mode 100644 index 24e59ba6242ea362eac349e5127833908a4f251c..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/common.cpp +++ /dev/null @@ -1,60 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// 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; -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/common.h b/extensions/nvdiffrast/nvdiffrast/common/common.h deleted file mode 100644 index df1e30d40ceb32c247516b68882c678b6dba215a..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/common.h +++ /dev/null @@ -1,263 +0,0 @@ -// 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 -#include - -//------------------------------------------------------------------------ -// 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 static __device__ __forceinline__ T zero_value(void); -template<> __device__ __forceinline__ float zero_value (void) { return 0.f; } -template<> __device__ __forceinline__ float2 zero_value(void) { return make_float2(0.f, 0.f); } -template<> __device__ __forceinline__ float4 zero_value(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 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__ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/CudaRaster.hpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/CudaRaster.hpp deleted file mode 100644 index 9261f995af645f50d25f9de5267e692576093cf1..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/CudaRaster.hpp +++ /dev/null @@ -1,63 +0,0 @@ -// 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 - diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/BinRaster.inl b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/BinRaster.inl deleted file mode 100644 index 0cfb2c356727a6ee9d3b1063c28fd4b6817e093c..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/BinRaster.inl +++ /dev/null @@ -1,423 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ - -__device__ __inline__ void binRasterImpl(const CRParams p) -{ - __shared__ volatile U32 s_broadcast [CR_BIN_WARPS + 16]; - __shared__ volatile S32 s_outOfs [CR_MAXBINS_SQR]; - __shared__ volatile S32 s_outTotal [CR_MAXBINS_SQR]; - __shared__ volatile S32 s_overIndex [CR_MAXBINS_SQR]; - __shared__ volatile S32 s_outMask [CR_BIN_WARPS][CR_MAXBINS_SQR + 1]; // +1 to avoid bank collisions - __shared__ volatile S32 s_outCount [CR_BIN_WARPS][CR_MAXBINS_SQR + 1]; // +1 to avoid bank collisions - __shared__ volatile S32 s_triBuf [CR_BIN_WARPS*32*4]; // triangle ring buffer - __shared__ volatile U32 s_batchPos; - __shared__ volatile U32 s_bufCount; - __shared__ volatile U32 s_overTotal; - __shared__ volatile U32 s_allocBase; - - const CRImageParams& ip = getImageParams(p, blockIdx.z); - CRAtomics& atomics = p.atomics[blockIdx.z]; - const U8* triSubtris = (const U8*)p.triSubtris + p.maxSubtris * blockIdx.z; - const CRTriangleHeader* triHeader = (const CRTriangleHeader*)p.triHeader + p.maxSubtris * blockIdx.z; - - S32* binFirstSeg = (S32*)p.binFirstSeg + CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * blockIdx.z; - S32* binTotal = (S32*)p.binTotal + CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * blockIdx.z; - S32* binSegData = (S32*)p.binSegData + p.maxBinSegs * CR_BIN_SEG_SIZE * blockIdx.z; - S32* binSegNext = (S32*)p.binSegNext + p.maxBinSegs * blockIdx.z; - S32* binSegCount = (S32*)p.binSegCount + p.maxBinSegs * blockIdx.z; - - if (atomics.numSubtris > p.maxSubtris) - return; - - // per-thread state - int thrInBlock = threadIdx.x + threadIdx.y * 32; - int batchPos = 0; - - // first 16 elements of s_broadcast are always zero - if (thrInBlock < 16) - s_broadcast[thrInBlock] = 0; - - // initialize output linked lists and offsets - if (thrInBlock < p.numBins) - { - binFirstSeg[(thrInBlock << CR_BIN_STREAMS_LOG2) + blockIdx.x] = -1; - s_outOfs[thrInBlock] = -CR_BIN_SEG_SIZE; - s_outTotal[thrInBlock] = 0; - } - - // repeat until done - for(;;) - { - // get batch - if (thrInBlock == 0) - s_batchPos = atomicAdd(&atomics.binCounter, ip.binBatchSize); - __syncthreads(); - batchPos = s_batchPos; - - // all batches done? - if (batchPos >= ip.triCount) - break; - - // per-thread state - int bufIndex = 0; - int bufCount = 0; - int batchEnd = min(batchPos + ip.binBatchSize, ip.triCount); - - // loop over batch as long as we have triangles in it - do - { - // read more triangles - while (bufCount < CR_BIN_WARPS*32 && batchPos < batchEnd) - { - // get subtriangle count - - int triIdx = batchPos + thrInBlock; - int num = 0; - if (triIdx < batchEnd) - num = triSubtris[triIdx]; - - // cumulative sum of subtriangles within each warp - U32 myIdx = __popc(__ballot_sync(~0u, num & 1) & getLaneMaskLt()); - if (__any_sync(~0u, num > 1)) - { - myIdx += __popc(__ballot_sync(~0u, num & 2) & getLaneMaskLt()) * 2; - myIdx += __popc(__ballot_sync(~0u, num & 4) & getLaneMaskLt()) * 4; - } - if (threadIdx.x == 31) // Do not assume that last thread in warp wins the write. - s_broadcast[threadIdx.y + 16] = myIdx + num; - __syncthreads(); - - // cumulative sum of per-warp subtriangle counts - // Note: cannot have more than 32 warps or this needs to sync between each step. - bool act = (thrInBlock < CR_BIN_WARPS); - U32 actMask = __ballot_sync(~0u, act); - if (threadIdx.y == 0 && act) - { - volatile U32* ptr = &s_broadcast[thrInBlock + 16]; - U32 val = *ptr; - #if (CR_BIN_WARPS > 1) - val += ptr[-1]; __syncwarp(actMask); - *ptr = val; __syncwarp(actMask); - #endif - #if (CR_BIN_WARPS > 2) - val += ptr[-2]; __syncwarp(actMask); - *ptr = val; __syncwarp(actMask); - #endif - #if (CR_BIN_WARPS > 4) - val += ptr[-4]; __syncwarp(actMask); - *ptr = val; __syncwarp(actMask); - #endif - #if (CR_BIN_WARPS > 8) - val += ptr[-8]; __syncwarp(actMask); - *ptr = val; __syncwarp(actMask); - #endif - #if (CR_BIN_WARPS > 16) - val += ptr[-16]; __syncwarp(actMask); - *ptr = val; __syncwarp(actMask); - #endif - - // initially assume that we consume everything - // only last active thread does the writes - if (threadIdx.x == CR_BIN_WARPS - 1) - { - s_batchPos = batchPos + CR_BIN_WARPS * 32; - s_bufCount = bufCount + val; - } - } - __syncthreads(); - - // skip if no subtriangles - if (num) - { - // calculate write position for first subtriangle - U32 pos = bufCount + myIdx + s_broadcast[threadIdx.y + 16 - 1]; - - // only write if entire triangle fits - if (pos + num <= CR_ARRAY_SIZE(s_triBuf)) - { - pos += bufIndex; // adjust for current start position - pos &= CR_ARRAY_SIZE(s_triBuf)-1; - if (num == 1) - s_triBuf[pos] = triIdx * 8 + 7; // single triangle - else - { - for (int i=0; i < num; i++) - { - s_triBuf[pos] = triIdx * 8 + i; - pos++; - pos &= CR_ARRAY_SIZE(s_triBuf)-1; - } - } - } else if (pos <= CR_ARRAY_SIZE(s_triBuf)) - { - // this triangle is the first that failed, overwrite total count and triangle count - s_batchPos = batchPos + thrInBlock; - s_bufCount = pos; - } - } - - // update triangle counts - __syncthreads(); - batchPos = s_batchPos; - bufCount = s_bufCount; - } - - // make every warp clear its output buffers - for (int i=threadIdx.x; i < p.numBins; i += 32) - s_outMask[threadIdx.y][i] = 0; - __syncwarp(); - - // choose our triangle - uint4 triData = make_uint4(0, 0, 0, 0); - if (thrInBlock < bufCount) - { - U32 triPos = bufIndex + thrInBlock; - triPos &= CR_ARRAY_SIZE(s_triBuf)-1; - - // find triangle - int triIdx = s_triBuf[triPos]; - int dataIdx = triIdx >> 3; - int subtriIdx = triIdx & 7; - if (subtriIdx != 7) - dataIdx = triHeader[dataIdx].misc + subtriIdx; - - // read triangle - - triData = *(((const uint4*)triHeader) + dataIdx); - } - - // setup bounding box and edge functions, and rasterize - S32 lox, loy, hix, hiy; - bool hasTri = (thrInBlock < bufCount); - U32 hasTriMask = __ballot_sync(~0u, hasTri); - if (hasTri) - { - S32 v0x = add_s16lo_s16lo(triData.x, p.widthPixelsVp * (CR_SUBPIXEL_SIZE >> 1)); - S32 v0y = add_s16hi_s16lo(triData.x, p.heightPixelsVp * (CR_SUBPIXEL_SIZE >> 1)); - S32 d01x = sub_s16lo_s16lo(triData.y, triData.x); - S32 d01y = sub_s16hi_s16hi(triData.y, triData.x); - S32 d02x = sub_s16lo_s16lo(triData.z, triData.x); - S32 d02y = sub_s16hi_s16hi(triData.z, triData.x); - int binLog = CR_BIN_LOG2 + CR_TILE_LOG2 + CR_SUBPIXEL_LOG2; - lox = add_clamp_0_x((v0x + min_min(d01x, 0, d02x)) >> binLog, 0, p.widthBins - 1); - loy = add_clamp_0_x((v0y + min_min(d01y, 0, d02y)) >> binLog, 0, p.heightBins - 1); - hix = add_clamp_0_x((v0x + max_max(d01x, 0, d02x)) >> binLog, 0, p.widthBins - 1); - hiy = add_clamp_0_x((v0y + max_max(d01y, 0, d02y)) >> binLog, 0, p.heightBins - 1); - - U32 bit = 1 << threadIdx.x; -#if __CUDA_ARCH__ >= 700 - bool multi = (hix != lox || hiy != loy); - if (!__any_sync(hasTriMask, multi)) - { - int binIdx = lox + p.widthBins * loy; - U32 mask = __match_any_sync(hasTriMask, binIdx); - s_outMask[threadIdx.y][binIdx] = mask; - __syncwarp(hasTriMask); - } else -#endif - { - bool complex = (hix > lox+1 || hiy > loy+1); - if (!__any_sync(hasTriMask, complex)) - { - int binIdx = lox + p.widthBins * loy; - atomicOr((U32*)&s_outMask[threadIdx.y][binIdx], bit); - if (hix > lox) atomicOr((U32*)&s_outMask[threadIdx.y][binIdx + 1], bit); - if (hiy > loy) atomicOr((U32*)&s_outMask[threadIdx.y][binIdx + p.widthBins], bit); - if (hix > lox && hiy > loy) atomicOr((U32*)&s_outMask[threadIdx.y][binIdx + p.widthBins + 1], bit); - } else - { - S32 d12x = d02x - d01x, d12y = d02y - d01y; - v0x -= lox << binLog, v0y -= loy << binLog; - - S32 t01 = v0x * d01y - v0y * d01x; - S32 t02 = v0y * d02x - v0x * d02y; - S32 t12 = d01x * d12y - d01y * d12x - t01 - t02; - S32 b01 = add_sub(t01 >> binLog, max(d01x, 0), min(d01y, 0)); - S32 b02 = add_sub(t02 >> binLog, max(d02y, 0), min(d02x, 0)); - S32 b12 = add_sub(t12 >> binLog, max(d12x, 0), min(d12y, 0)); - - int width = hix - lox + 1; - d01x += width * d01y; - d02x += width * d02y; - d12x += width * d12y; - - U8* currPtr = (U8*)&s_outMask[threadIdx.y][lox + loy * p.widthBins]; - U8* skipPtr = (U8*)&s_outMask[threadIdx.y][(hix + 1) + loy * p.widthBins]; - U8* endPtr = (U8*)&s_outMask[threadIdx.y][lox + (hiy + 1) * p.widthBins]; - int stride = p.widthBins * 4; - int ptrYInc = stride - width * 4; - - do - { - if (b01 >= 0 && b02 >= 0 && b12 >= 0) - atomicOr((U32*)currPtr, bit); - currPtr += 4, b01 -= d01y, b02 += d02y, b12 -= d12y; - if (currPtr == skipPtr) - currPtr += ptrYInc, b01 += d01x, b02 -= d02x, b12 += d12x, skipPtr += stride; - } - while (currPtr != endPtr); - } - } - } - - // count per-bin contributions - if (thrInBlock == 0) - s_overTotal = 0; // overflow counter - - // ensure that out masks are done - __syncthreads(); - - int overIndex = -1; - bool act = (thrInBlock < p.numBins); - U32 actMask = __ballot_sync(~0u, act); - if (act) - { - U8* srcPtr = (U8*)&s_outMask[0][thrInBlock]; - U8* dstPtr = (U8*)&s_outCount[0][thrInBlock]; - int total = 0; - for (int i = 0; i < CR_BIN_WARPS; i++) - { - total += __popc(*(U32*)srcPtr); - *(U32*)dstPtr = total; - srcPtr += (CR_MAXBINS_SQR + 1) * 4; - dstPtr += (CR_MAXBINS_SQR + 1) * 4; - } - - // overflow => request a new segment - int ofs = s_outOfs[thrInBlock]; - bool ovr = (((ofs - 1) >> CR_BIN_SEG_LOG2) != (((ofs - 1) + total) >> CR_BIN_SEG_LOG2)); - U32 ovrMask = __ballot_sync(actMask, ovr); - if (ovr) - { - overIndex = __popc(ovrMask & getLaneMaskLt()); - if (overIndex == 0) - s_broadcast[threadIdx.y + 16] = atomicAdd((U32*)&s_overTotal, __popc(ovrMask)); - __syncwarp(ovrMask); - overIndex += s_broadcast[threadIdx.y + 16]; - s_overIndex[thrInBlock] = overIndex; - } - } - - // sync after overTotal is ready - __syncthreads(); - - // at least one segment overflowed => allocate segments - U32 overTotal = s_overTotal; - U32 allocBase = 0; - if (overTotal > 0) - { - // allocate memory - if (thrInBlock == 0) - { - U32 allocBase = atomicAdd(&atomics.numBinSegs, overTotal); - s_allocBase = (allocBase + overTotal <= p.maxBinSegs) ? allocBase : 0; - } - __syncthreads(); - allocBase = s_allocBase; - - // did my bin overflow? - if (overIndex != -1) - { - // calculate new segment index - int segIdx = allocBase + overIndex; - - // add to linked list - if (s_outOfs[thrInBlock] < 0) - binFirstSeg[(thrInBlock << CR_BIN_STREAMS_LOG2) + blockIdx.x] = segIdx; - else - binSegNext[(s_outOfs[thrInBlock] - 1) >> CR_BIN_SEG_LOG2] = segIdx; - - // defaults - binSegNext [segIdx] = -1; - binSegCount[segIdx] = CR_BIN_SEG_SIZE; - } - } - - // concurrent emission -- each warp handles its own triangle - if (thrInBlock < bufCount) - { - int triPos = (bufIndex + thrInBlock) & (CR_ARRAY_SIZE(s_triBuf) - 1); - int currBin = lox + loy * p.widthBins; - int skipBin = (hix + 1) + loy * p.widthBins; - int endBin = lox + (hiy + 1) * p.widthBins; - int binYInc = p.widthBins - (hix - lox + 1); - - // loop over triangle's bins - do - { - U32 outMask = s_outMask[threadIdx.y][currBin]; - if (outMask & (1< 0) - idx += s_outCount[threadIdx.y-1][currBin]; - - int base = s_outOfs[currBin]; - int free = (-base) & (CR_BIN_SEG_SIZE - 1); - if (idx >= free) - idx += ((allocBase + s_overIndex[currBin]) << CR_BIN_SEG_LOG2) - free; - else - idx += base; - - binSegData[idx] = s_triBuf[triPos]; - } - - currBin++; - if (currBin == skipBin) - currBin += binYInc, skipBin += p.widthBins; - } - while (currBin != endBin); - } - - // wait all triangles to finish, then replace overflown segment offsets - __syncthreads(); - if (thrInBlock < p.numBins) - { - U32 total = s_outCount[CR_BIN_WARPS - 1][thrInBlock]; - U32 oldOfs = s_outOfs[thrInBlock]; - if (overIndex == -1) - s_outOfs[thrInBlock] = oldOfs + total; - else - { - int addr = oldOfs + total; - addr = ((addr - 1) & (CR_BIN_SEG_SIZE - 1)) + 1; - addr += (allocBase + overIndex) << CR_BIN_SEG_LOG2; - s_outOfs[thrInBlock] = addr; - } - s_outTotal[thrInBlock] += total; - } - - // these triangles are now done - int count = ::min(bufCount, CR_BIN_WARPS * 32); - bufCount -= count; - bufIndex += count; - bufIndex &= CR_ARRAY_SIZE(s_triBuf)-1; - } - while (bufCount > 0 || batchPos < batchEnd); - - // flush all bins - if (thrInBlock < p.numBins) - { - int ofs = s_outOfs[thrInBlock]; - if (ofs & (CR_BIN_SEG_SIZE-1)) - { - int seg = ofs >> CR_BIN_SEG_LOG2; - binSegCount[seg] = ofs & (CR_BIN_SEG_SIZE-1); - s_outOfs[thrInBlock] = (ofs + CR_BIN_SEG_SIZE - 1) & -CR_BIN_SEG_SIZE; - } - } - } - - // output totals - if (thrInBlock < p.numBins) - binTotal[(thrInBlock << CR_BIN_STREAMS_LOG2) + blockIdx.x] = s_outTotal[thrInBlock]; -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.cpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.cpp deleted file mode 100644 index acd42262ec68e324657bb7cbd7449150c136c6fc..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.cpp +++ /dev/null @@ -1,94 +0,0 @@ -// 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); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.hpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.hpp deleted file mode 100644 index 1063a68d0c024f8580a99efefeeb0ed56ec3f239..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.hpp +++ /dev/null @@ -1,55 +0,0 @@ -// 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; -}; - -//------------------------------------------------------------------------ -} diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CoarseRaster.inl b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CoarseRaster.inl deleted file mode 100644 index dc39da88d341da0d9d704f9cc3cdff1c793cca54..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CoarseRaster.inl +++ /dev/null @@ -1,730 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ - -__device__ __inline__ int globalTileIdx(int tileInBin, int widthTiles) -{ - int tileX = tileInBin & (CR_BIN_SIZE - 1); - int tileY = tileInBin >> CR_BIN_LOG2; - return tileX + tileY * widthTiles; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ void coarseRasterImpl(const CRParams p) -{ - // Common. - - __shared__ volatile U32 s_workCounter; - __shared__ volatile U32 s_scanTemp [CR_COARSE_WARPS][48]; // 3KB - - // Input. - - __shared__ volatile U32 s_binOrder [CR_MAXBINS_SQR]; // 1KB - __shared__ volatile S32 s_binStreamCurrSeg [CR_BIN_STREAMS_SIZE]; // 0KB - __shared__ volatile S32 s_binStreamFirstTri [CR_BIN_STREAMS_SIZE]; // 0KB - __shared__ volatile S32 s_triQueue [CR_COARSE_QUEUE_SIZE]; // 4KB - __shared__ volatile S32 s_triQueueWritePos; - __shared__ volatile U32 s_binStreamSelectedOfs; - __shared__ volatile U32 s_binStreamSelectedSize; - - // Output. - - __shared__ volatile U32 s_warpEmitMask [CR_COARSE_WARPS][CR_BIN_SQR + 1]; // 16KB, +1 to avoid bank collisions - __shared__ volatile U32 s_warpEmitPrefixSum [CR_COARSE_WARPS][CR_BIN_SQR + 1]; // 16KB, +1 to avoid bank collisions - __shared__ volatile U32 s_tileEmitPrefixSum [CR_BIN_SQR + 1]; // 1KB, zero at the beginning - __shared__ volatile U32 s_tileAllocPrefixSum[CR_BIN_SQR + 1]; // 1KB, zero at the beginning - __shared__ volatile S32 s_tileStreamCurrOfs [CR_BIN_SQR]; // 1KB - __shared__ volatile U32 s_firstAllocSeg; - __shared__ volatile U32 s_firstActiveIdx; - - // Pointers and constants. - - CRAtomics& atomics = p.atomics[blockIdx.z]; - const CRTriangleHeader* triHeader = (const CRTriangleHeader*)p.triHeader + p.maxSubtris * blockIdx.z; - const S32* binFirstSeg = (const S32*)p.binFirstSeg + CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * blockIdx.z; - const S32* binTotal = (const S32*)p.binTotal + CR_MAXBINS_SQR * CR_BIN_STREAMS_SIZE * blockIdx.z; - const S32* binSegData = (const S32*)p.binSegData + p.maxBinSegs * CR_BIN_SEG_SIZE * blockIdx.z; - const S32* binSegNext = (const S32*)p.binSegNext + p.maxBinSegs * blockIdx.z; - const S32* binSegCount = (const S32*)p.binSegCount + p.maxBinSegs * blockIdx.z; - S32* activeTiles = (S32*)p.activeTiles + CR_MAXTILES_SQR * blockIdx.z; - S32* tileFirstSeg = (S32*)p.tileFirstSeg + CR_MAXTILES_SQR * blockIdx.z; - S32* tileSegData = (S32*)p.tileSegData + p.maxTileSegs * CR_TILE_SEG_SIZE * blockIdx.z; - S32* tileSegNext = (S32*)p.tileSegNext + p.maxTileSegs * blockIdx.z; - S32* tileSegCount = (S32*)p.tileSegCount + p.maxTileSegs * blockIdx.z; - - int tileLog = CR_TILE_LOG2 + CR_SUBPIXEL_LOG2; - int thrInBlock = threadIdx.x + threadIdx.y * 32; - int emitShift = CR_BIN_LOG2 * 2 + 5; // We scan ((numEmits << emitShift) | numAllocs) over tiles. - - if (atomics.numSubtris > p.maxSubtris || atomics.numBinSegs > p.maxBinSegs) - return; - - // Initialize sharedmem arrays. - - if (thrInBlock == 0) - { - s_tileEmitPrefixSum[0] = 0; - s_tileAllocPrefixSum[0] = 0; - } - s_scanTemp[threadIdx.y][threadIdx.x] = 0; - - // Sort bins in descending order of triangle count. - - for (int binIdx = thrInBlock; binIdx < p.numBins; binIdx += CR_COARSE_WARPS * 32) - { - int count = 0; - for (int i = 0; i < CR_BIN_STREAMS_SIZE; i++) - count += binTotal[(binIdx << CR_BIN_STREAMS_LOG2) + i]; - s_binOrder[binIdx] = (~count << (CR_MAXBINS_LOG2 * 2)) | binIdx; - } - - __syncthreads(); - sortShared(s_binOrder, p.numBins); - - // Process each bin by one block. - - for (;;) - { - // Pick a bin for the block. - - if (thrInBlock == 0) - s_workCounter = atomicAdd(&atomics.coarseCounter, 1); - __syncthreads(); - - int workCounter = s_workCounter; - if (workCounter >= p.numBins) - break; - - U32 binOrder = s_binOrder[workCounter]; - bool binEmpty = ((~binOrder >> (CR_MAXBINS_LOG2 * 2)) == 0); - if (binEmpty && !p.deferredClear) - break; - - int binIdx = binOrder & (CR_MAXBINS_SQR - 1); - - // Initialize input/output streams. - - int triQueueWritePos = 0; - int triQueueReadPos = 0; - - if (thrInBlock < CR_BIN_STREAMS_SIZE) - { - int segIdx = binFirstSeg[(binIdx << CR_BIN_STREAMS_LOG2) + thrInBlock]; - s_binStreamCurrSeg[thrInBlock] = segIdx; - s_binStreamFirstTri[thrInBlock] = (segIdx == -1) ? ~0u : binSegData[segIdx << CR_BIN_SEG_LOG2]; - } - - for (int tileInBin = CR_COARSE_WARPS * 32 - 1 - thrInBlock; tileInBin < CR_BIN_SQR; tileInBin += CR_COARSE_WARPS * 32) - s_tileStreamCurrOfs[tileInBin] = -CR_TILE_SEG_SIZE; - - // Initialize per-bin state. - - int binY = idiv_fast(binIdx, p.widthBins); - int binX = binIdx - binY * p.widthBins; - int originX = (binX << (CR_BIN_LOG2 + tileLog)) - (p.widthPixelsVp << (CR_SUBPIXEL_LOG2 - 1)); - int originY = (binY << (CR_BIN_LOG2 + tileLog)) - (p.heightPixelsVp << (CR_SUBPIXEL_LOG2 - 1)); - int maxTileXInBin = ::min(p.widthTiles - (binX << CR_BIN_LOG2), CR_BIN_SIZE) - 1; - int maxTileYInBin = ::min(p.heightTiles - (binY << CR_BIN_LOG2), CR_BIN_SIZE) - 1; - int binTileIdx = (binX + binY * p.widthTiles) << CR_BIN_LOG2; - - // Entire block: Merge input streams and process triangles. - - if (!binEmpty) - do - { - //------------------------------------------------------------------------ - // Merge. - //------------------------------------------------------------------------ - - // Entire block: Not enough triangles => merge and queue segments. - // NOTE: The bin exit criterion assumes that we queue more triangles than we actually need. - - while (triQueueWritePos - triQueueReadPos <= CR_COARSE_WARPS * 32) - { - // First warp: Choose the segment with the lowest initial triangle index. - - bool hasStream = (thrInBlock < CR_BIN_STREAMS_SIZE); - U32 hasStreamMask = __ballot_sync(~0u, hasStream); - if (hasStream) - { - // Find the stream with the lowest triangle index. - - U32 firstTri = s_binStreamFirstTri[thrInBlock]; - U32 t = firstTri; - volatile U32* v = &s_scanTemp[0][thrInBlock + 16]; - - #if (CR_BIN_STREAMS_SIZE > 1) - v[0] = t; __syncwarp(hasStreamMask); t = ::min(t, v[-1]); __syncwarp(hasStreamMask); - #endif - #if (CR_BIN_STREAMS_SIZE > 2) - v[0] = t; __syncwarp(hasStreamMask); t = ::min(t, v[-2]); __syncwarp(hasStreamMask); - #endif - #if (CR_BIN_STREAMS_SIZE > 4) - v[0] = t; __syncwarp(hasStreamMask); t = ::min(t, v[-4]); __syncwarp(hasStreamMask); - #endif - #if (CR_BIN_STREAMS_SIZE > 8) - v[0] = t; __syncwarp(hasStreamMask); t = ::min(t, v[-8]); __syncwarp(hasStreamMask); - #endif - #if (CR_BIN_STREAMS_SIZE > 16) - v[0] = t; __syncwarp(hasStreamMask); t = ::min(t, v[-16]); __syncwarp(hasStreamMask); - #endif - v[0] = t; __syncwarp(hasStreamMask); - - // Consume and broadcast. - - bool first = (s_scanTemp[0][CR_BIN_STREAMS_SIZE - 1 + 16] == firstTri); - U32 firstMask = __ballot_sync(hasStreamMask, first); - if (first && (firstMask >> threadIdx.x) == 1u) - { - int segIdx = s_binStreamCurrSeg[thrInBlock]; - s_binStreamSelectedOfs = segIdx << CR_BIN_SEG_LOG2; - if (segIdx != -1) - { - int segSize = binSegCount[segIdx]; - int segNext = binSegNext[segIdx]; - s_binStreamSelectedSize = segSize; - s_triQueueWritePos = triQueueWritePos + segSize; - s_binStreamCurrSeg[thrInBlock] = segNext; - s_binStreamFirstTri[thrInBlock] = (segNext == -1) ? ~0u : binSegData[segNext << CR_BIN_SEG_LOG2]; - } - } - } - - // No more segments => break. - - __syncthreads(); - triQueueWritePos = s_triQueueWritePos; - int segOfs = s_binStreamSelectedOfs; - if (segOfs < 0) - break; - - int segSize = s_binStreamSelectedSize; - __syncthreads(); - - // Fetch triangles into the queue. - - for (int idxInSeg = CR_COARSE_WARPS * 32 - 1 - thrInBlock; idxInSeg < segSize; idxInSeg += CR_COARSE_WARPS * 32) - { - S32 triIdx = binSegData[segOfs + idxInSeg]; - s_triQueue[(triQueueWritePos - segSize + idxInSeg) & (CR_COARSE_QUEUE_SIZE - 1)] = triIdx; - } - } - - // All threads: Clear emit masks. - - for (int maskIdx = thrInBlock; maskIdx < CR_COARSE_WARPS * CR_BIN_SQR; maskIdx += CR_COARSE_WARPS * 32) - s_warpEmitMask[maskIdx >> (CR_BIN_LOG2 * 2)][maskIdx & (CR_BIN_SQR - 1)] = 0; - - __syncthreads(); - - //------------------------------------------------------------------------ - // Raster. - //------------------------------------------------------------------------ - - // Triangle per thread: Read from the queue. - - int triIdx = -1; - if (triQueueReadPos + thrInBlock < triQueueWritePos) - triIdx = s_triQueue[(triQueueReadPos + thrInBlock) & (CR_COARSE_QUEUE_SIZE - 1)]; - - uint4 triData = make_uint4(0, 0, 0, 0); - if (triIdx != -1) - { - int dataIdx = triIdx >> 3; - int subtriIdx = triIdx & 7; - if (subtriIdx != 7) - dataIdx = triHeader[dataIdx].misc + subtriIdx; - triData = *((uint4*)triHeader + dataIdx); - } - - // 32 triangles per warp: Record emits (= tile intersections). - - if (__any_sync(~0u, triIdx != -1)) - { - S32 v0x = sub_s16lo_s16lo(triData.x, originX); - S32 v0y = sub_s16hi_s16lo(triData.x, originY); - S32 d01x = sub_s16lo_s16lo(triData.y, triData.x); - S32 d01y = sub_s16hi_s16hi(triData.y, triData.x); - S32 d02x = sub_s16lo_s16lo(triData.z, triData.x); - S32 d02y = sub_s16hi_s16hi(triData.z, triData.x); - - // Compute tile-based AABB. - - int lox = add_clamp_0_x((v0x + min_min(d01x, 0, d02x)) >> tileLog, 0, maxTileXInBin); - int loy = add_clamp_0_x((v0y + min_min(d01y, 0, d02y)) >> tileLog, 0, maxTileYInBin); - int hix = add_clamp_0_x((v0x + max_max(d01x, 0, d02x)) >> tileLog, 0, maxTileXInBin); - int hiy = add_clamp_0_x((v0y + max_max(d01y, 0, d02y)) >> tileLog, 0, maxTileYInBin); - int sizex = add_sub(hix, 1, lox); - int sizey = add_sub(hiy, 1, loy); - int area = sizex * sizey; - - // Miscellaneous init. - - U8* currPtr = (U8*)&s_warpEmitMask[threadIdx.y][lox + (loy << CR_BIN_LOG2)]; - int ptrYInc = CR_BIN_SIZE * 4 - (sizex << 2); - U32 maskBit = 1 << threadIdx.x; - - // Case A: All AABBs are small => record the full AABB using atomics. - - if (__all_sync(~0u, sizex <= 2 && sizey <= 2)) - { - if (triIdx != -1) - { - atomicOr((U32*)currPtr, maskBit); - if (sizex == 2) atomicOr((U32*)(currPtr + 4), maskBit); - if (sizey == 2) atomicOr((U32*)(currPtr + CR_BIN_SIZE * 4), maskBit); - if (sizex == 2 && sizey == 2) atomicOr((U32*)(currPtr + 4 + CR_BIN_SIZE * 4), maskBit); - } - } - else - { - // Compute warp-AABB (scan-32). - - U32 aabbMask = add_sub(2 << hix, 0x20000 << hiy, 1 << lox) - (0x10000 << loy); - if (triIdx == -1) - aabbMask = 0; - - volatile U32* v = &s_scanTemp[threadIdx.y][threadIdx.x + 16]; - v[0] = aabbMask; __syncwarp(); aabbMask |= v[-1]; __syncwarp(); - v[0] = aabbMask; __syncwarp(); aabbMask |= v[-2]; __syncwarp(); - v[0] = aabbMask; __syncwarp(); aabbMask |= v[-4]; __syncwarp(); - v[0] = aabbMask; __syncwarp(); aabbMask |= v[-8]; __syncwarp(); - v[0] = aabbMask; __syncwarp(); aabbMask |= v[-16]; __syncwarp(); - v[0] = aabbMask; __syncwarp(); aabbMask = s_scanTemp[threadIdx.y][47]; - - U32 maskX = aabbMask & 0xFFFF; - U32 maskY = aabbMask >> 16; - int wlox = findLeadingOne(maskX ^ (maskX - 1)); - int wloy = findLeadingOne(maskY ^ (maskY - 1)); - int whix = findLeadingOne(maskX); - int whiy = findLeadingOne(maskY); - int warea = (add_sub(whix, 1, wlox)) * (add_sub(whiy, 1, wloy)); - - // Initialize edge functions. - - S32 d12x = d02x - d01x; - S32 d12y = d02y - d01y; - v0x -= lox << tileLog; - v0y -= loy << tileLog; - - S32 t01 = v0x * d01y - v0y * d01x; - S32 t02 = v0y * d02x - v0x * d02y; - S32 t12 = d01x * d12y - d01y * d12x - t01 - t02; - S32 b01 = add_sub(t01 >> tileLog, ::max(d01x, 0), ::min(d01y, 0)); - S32 b02 = add_sub(t02 >> tileLog, ::max(d02y, 0), ::min(d02x, 0)); - S32 b12 = add_sub(t12 >> tileLog, ::max(d12x, 0), ::min(d12y, 0)); - - d01x += sizex * d01y; - d02x += sizex * d02y; - d12x += sizex * d12y; - - // Case B: Warp-AABB is not much larger than largest AABB => Check tiles in warp-AABB, record using ballots. - if (__any_sync(~0u, warea * 4 <= area * 8)) - { - // Not sure if this is any faster than Case C after all the post-Volta ballot mask tracking. - bool act = (triIdx != -1); - U32 actMask = __ballot_sync(~0u, act); - if (act) - { - for (int y = wloy; y <= whiy; y++) - { - bool yIn = (y >= loy && y <= hiy); - U32 yMask = __ballot_sync(actMask, yIn); - if (yIn) - { - for (int x = wlox; x <= whix; x++) - { - bool xyIn = (x >= lox && x <= hix); - U32 xyMask = __ballot_sync(yMask, xyIn); - if (xyIn) - { - U32 res = __ballot_sync(xyMask, b01 >= 0 && b02 >= 0 && b12 >= 0); - if (threadIdx.x == 31 - __clz(xyMask)) - *(U32*)currPtr = res; - currPtr += 4, b01 -= d01y, b02 += d02y, b12 -= d12y; - } - } - currPtr += ptrYInc, b01 += d01x, b02 -= d02x, b12 += d12x; - } - } - } - } - - // Case C: General case => Check tiles in AABB, record using atomics. - - else - { - if (triIdx != -1) - { - U8* skipPtr = currPtr + (sizex << 2); - U8* endPtr = currPtr + (sizey << (CR_BIN_LOG2 + 2)); - do - { - if (b01 >= 0 && b02 >= 0 && b12 >= 0) - atomicOr((U32*)currPtr, maskBit); - currPtr += 4, b01 -= d01y, b02 += d02y, b12 -= d12y; - if (currPtr == skipPtr) - currPtr += ptrYInc, b01 += d01x, b02 -= d02x, b12 += d12x, skipPtr += CR_BIN_SIZE * 4; - } - while (currPtr != endPtr); - } - } - } - } - - __syncthreads(); - - //------------------------------------------------------------------------ - // Count. - //------------------------------------------------------------------------ - - // Tile per thread: Initialize prefix sums. - - for (int tileInBin_base = 0; tileInBin_base < CR_BIN_SQR; tileInBin_base += CR_COARSE_WARPS * 32) - { - int tileInBin = tileInBin_base + thrInBlock; - bool act = (tileInBin < CR_BIN_SQR); - U32 actMask = __ballot_sync(~0u, act); - if (act) - { - // Compute prefix sum of emits over warps. - - U8* srcPtr = (U8*)&s_warpEmitMask[0][tileInBin]; - U8* dstPtr = (U8*)&s_warpEmitPrefixSum[0][tileInBin]; - int tileEmits = 0; - for (int i = 0; i < CR_COARSE_WARPS; i++) - { - tileEmits += __popc(*(U32*)srcPtr); - *(U32*)dstPtr = tileEmits; - srcPtr += (CR_BIN_SQR + 1) * 4; - dstPtr += (CR_BIN_SQR + 1) * 4; - } - - // Determine the number of segments to allocate. - - int spaceLeft = -s_tileStreamCurrOfs[tileInBin] & (CR_TILE_SEG_SIZE - 1); - int tileAllocs = (tileEmits - spaceLeft + CR_TILE_SEG_SIZE - 1) >> CR_TILE_SEG_LOG2; - volatile U32* v = &s_tileEmitPrefixSum[tileInBin + 1]; - - // All counters within the warp are small => compute prefix sum using ballot. - - if (!__any_sync(actMask, tileEmits >= 2)) - { - U32 m = getLaneMaskLe(); - *v = (__popc(__ballot_sync(actMask, tileEmits & 1) & m) << emitShift) | __popc(__ballot_sync(actMask, tileAllocs & 1) & m); - } - - // Otherwise => scan-32 within the warp. - - else - { - U32 sum = (tileEmits << emitShift) | tileAllocs; - *v = sum; __syncwarp(actMask); if (threadIdx.x >= 1) sum += v[-1]; __syncwarp(actMask); - *v = sum; __syncwarp(actMask); if (threadIdx.x >= 2) sum += v[-2]; __syncwarp(actMask); - *v = sum; __syncwarp(actMask); if (threadIdx.x >= 4) sum += v[-4]; __syncwarp(actMask); - *v = sum; __syncwarp(actMask); if (threadIdx.x >= 8) sum += v[-8]; __syncwarp(actMask); - *v = sum; __syncwarp(actMask); if (threadIdx.x >= 16) sum += v[-16]; __syncwarp(actMask); - *v = sum; __syncwarp(actMask); - } - } - } - - // First warp: Scan-8. - - __syncthreads(); - - bool scan8 = (thrInBlock < CR_BIN_SQR / 32); - U32 scan8Mask = __ballot_sync(~0u, scan8); - if (scan8) - { - int sum = s_tileEmitPrefixSum[(thrInBlock << 5) + 32]; - volatile U32* v = &s_scanTemp[0][thrInBlock + 16]; - v[0] = sum; __syncwarp(scan8Mask); - #if (CR_BIN_SQR > 1 * 32) - sum += v[-1]; __syncwarp(scan8Mask); v[0] = sum; __syncwarp(scan8Mask); - #endif - #if (CR_BIN_SQR > 2 * 32) - sum += v[-2]; __syncwarp(scan8Mask); v[0] = sum; __syncwarp(scan8Mask); - #endif - #if (CR_BIN_SQR > 4 * 32) - sum += v[-4]; __syncwarp(scan8Mask); v[0] = sum; __syncwarp(scan8Mask); - #endif - } - - __syncthreads(); - - // Tile per thread: Finalize prefix sums. - // Single thread: Allocate segments. - - for (int tileInBin = thrInBlock; tileInBin < CR_BIN_SQR; tileInBin += CR_COARSE_WARPS * 32) - { - int sum = s_tileEmitPrefixSum[tileInBin + 1] + s_scanTemp[0][(tileInBin >> 5) + 15]; - int numEmits = sum >> emitShift; - int numAllocs = sum & ((1 << emitShift) - 1); - s_tileEmitPrefixSum[tileInBin + 1] = numEmits; - s_tileAllocPrefixSum[tileInBin + 1] = numAllocs; - - if (tileInBin == CR_BIN_SQR - 1 && numAllocs != 0) - { - int t = atomicAdd(&atomics.numTileSegs, numAllocs); - s_firstAllocSeg = (t + numAllocs <= p.maxTileSegs) ? t : 0; - } - } - - __syncthreads(); - int firstAllocSeg = s_firstAllocSeg; - int totalEmits = s_tileEmitPrefixSum[CR_BIN_SQR]; - int totalAllocs = s_tileAllocPrefixSum[CR_BIN_SQR]; - - //------------------------------------------------------------------------ - // Emit. - //------------------------------------------------------------------------ - - // Emit per thread: Write triangle index to globalmem. - - for (int emitInBin = thrInBlock; emitInBin < totalEmits; emitInBin += CR_COARSE_WARPS * 32) - { - // Find tile in bin. - - U8* tileBase = (U8*)&s_tileEmitPrefixSum[0]; - U8* tilePtr = tileBase; - U8* ptr; - - #if (CR_BIN_SQR > 128) - ptr = tilePtr + 0x80 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 64) - ptr = tilePtr + 0x40 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 32) - ptr = tilePtr + 0x20 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 16) - ptr = tilePtr + 0x10 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 8) - ptr = tilePtr + 0x08 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 4) - ptr = tilePtr + 0x04 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 2) - ptr = tilePtr + 0x02 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - #if (CR_BIN_SQR > 1) - ptr = tilePtr + 0x01 * 4; if (emitInBin >= *(U32*)ptr) tilePtr = ptr; - #endif - - int tileInBin = (tilePtr - tileBase) >> 2; - int emitInTile = emitInBin - *(U32*)tilePtr; - - // Find warp in tile. - - int warpStep = (CR_BIN_SQR + 1) * 4; - U8* warpBase = (U8*)&s_warpEmitPrefixSum[0][tileInBin] - warpStep; - U8* warpPtr = warpBase; - - #if (CR_COARSE_WARPS > 8) - ptr = warpPtr + 0x08 * warpStep; if (emitInTile >= *(U32*)ptr) warpPtr = ptr; - #endif - #if (CR_COARSE_WARPS > 4) - ptr = warpPtr + 0x04 * warpStep; if (emitInTile >= *(U32*)ptr) warpPtr = ptr; - #endif - #if (CR_COARSE_WARPS > 2) - ptr = warpPtr + 0x02 * warpStep; if (emitInTile >= *(U32*)ptr) warpPtr = ptr; - #endif - #if (CR_COARSE_WARPS > 1) - ptr = warpPtr + 0x01 * warpStep; if (emitInTile >= *(U32*)ptr) warpPtr = ptr; - #endif - - int warpInTile = (warpPtr - warpBase) >> (CR_BIN_LOG2 * 2 + 2); - U32 emitMask = *(U32*)(warpPtr + warpStep + ((U8*)s_warpEmitMask - (U8*)s_warpEmitPrefixSum)); - int emitInWarp = emitInTile - *(U32*)(warpPtr + warpStep) + __popc(emitMask); - - // Find thread in warp. - - int threadInWarp = 0; - int pop = __popc(emitMask & 0xFFFF); - bool pred = (emitInWarp >= pop); - if (pred) emitInWarp -= pop; - if (pred) emitMask >>= 0x10; - if (pred) threadInWarp += 0x10; - - pop = __popc(emitMask & 0xFF); - pred = (emitInWarp >= pop); - if (pred) emitInWarp -= pop; - if (pred) emitMask >>= 0x08; - if (pred) threadInWarp += 0x08; - - pop = __popc(emitMask & 0xF); - pred = (emitInWarp >= pop); - if (pred) emitInWarp -= pop; - if (pred) emitMask >>= 0x04; - if (pred) threadInWarp += 0x04; - - pop = __popc(emitMask & 0x3); - pred = (emitInWarp >= pop); - if (pred) emitInWarp -= pop; - if (pred) emitMask >>= 0x02; - if (pred) threadInWarp += 0x02; - - if (emitInWarp >= (emitMask & 1)) - threadInWarp++; - - // Figure out where to write. - - int currOfs = s_tileStreamCurrOfs[tileInBin]; - int spaceLeft = -currOfs & (CR_TILE_SEG_SIZE - 1); - int outOfs = emitInTile; - - if (outOfs < spaceLeft) - outOfs += currOfs; - else - { - int allocLo = firstAllocSeg + s_tileAllocPrefixSum[tileInBin]; - outOfs += (allocLo << CR_TILE_SEG_LOG2) - spaceLeft; - } - - // Write. - - int queueIdx = warpInTile * 32 + threadInWarp; - int triIdx = s_triQueue[(triQueueReadPos + queueIdx) & (CR_COARSE_QUEUE_SIZE - 1)]; - - tileSegData[outOfs] = triIdx; - } - - //------------------------------------------------------------------------ - // Patch. - //------------------------------------------------------------------------ - - // Allocated segment per thread: Initialize next-pointer and count. - - for (int i = CR_COARSE_WARPS * 32 - 1 - thrInBlock; i < totalAllocs; i += CR_COARSE_WARPS * 32) - { - int segIdx = firstAllocSeg + i; - tileSegNext[segIdx] = segIdx + 1; - tileSegCount[segIdx] = CR_TILE_SEG_SIZE; - } - - // Tile per thread: Fix previous segment's next-pointer and update s_tileStreamCurrOfs. - - __syncthreads(); - for (int tileInBin = CR_COARSE_WARPS * 32 - 1 - thrInBlock; tileInBin < CR_BIN_SQR; tileInBin += CR_COARSE_WARPS * 32) - { - int oldOfs = s_tileStreamCurrOfs[tileInBin]; - int newOfs = oldOfs + s_warpEmitPrefixSum[CR_COARSE_WARPS - 1][tileInBin]; - int allocLo = s_tileAllocPrefixSum[tileInBin]; - int allocHi = s_tileAllocPrefixSum[tileInBin + 1]; - - if (allocLo != allocHi) - { - S32* nextPtr = &tileSegNext[(oldOfs - 1) >> CR_TILE_SEG_LOG2]; - if (oldOfs < 0) - nextPtr = &tileFirstSeg[binTileIdx + globalTileIdx(tileInBin, p.widthTiles)]; - *nextPtr = firstAllocSeg + allocLo; - - newOfs--; - newOfs &= CR_TILE_SEG_SIZE - 1; - newOfs |= (firstAllocSeg + allocHi - 1) << CR_TILE_SEG_LOG2; - newOfs++; - } - s_tileStreamCurrOfs[tileInBin] = newOfs; - } - - // Advance queue read pointer. - // Queue became empty => bin done. - - triQueueReadPos += CR_COARSE_WARPS * 32; - } - while (triQueueReadPos < triQueueWritePos); - - // Tile per thread: Fix next-pointer and count of the last segment. - // 32 tiles per warp: Count active tiles. - - __syncthreads(); - - for (int tileInBin_base = 0; tileInBin_base < CR_BIN_SQR; tileInBin_base += CR_COARSE_WARPS * 32) - { - int tileInBin = tileInBin_base + thrInBlock; - bool act = (tileInBin < CR_BIN_SQR); - U32 actMask = __ballot_sync(~0u, act); - if (act) - { - int tileX = tileInBin & (CR_BIN_SIZE - 1); - int tileY = tileInBin >> CR_BIN_LOG2; - bool force = (p.deferredClear & tileX <= maxTileXInBin & tileY <= maxTileYInBin); - - int ofs = s_tileStreamCurrOfs[tileInBin]; - int segIdx = (ofs - 1) >> CR_TILE_SEG_LOG2; - int segCount = ofs & (CR_TILE_SEG_SIZE - 1); - - if (ofs >= 0) - tileSegNext[segIdx] = -1; - else if (force) - { - s_tileStreamCurrOfs[tileInBin] = 0; - tileFirstSeg[binTileIdx + tileX + tileY * p.widthTiles] = -1; - } - - if (segCount != 0) - tileSegCount[segIdx] = segCount; - - U32 res = __ballot_sync(actMask, ofs >= 0 | force); - if (threadIdx.x == 0) - s_scanTemp[0][(tileInBin >> 5) + 16] = __popc(res); - } - } - - // First warp: Scan-8. - // One thread: Allocate space for active tiles. - - __syncthreads(); - - bool scan8 = (thrInBlock < CR_BIN_SQR / 32); - U32 scan8Mask = __ballot_sync(~0u, scan8); - if (scan8) - { - volatile U32* v = &s_scanTemp[0][thrInBlock + 16]; - U32 sum = v[0]; - #if (CR_BIN_SQR > 1 * 32) - sum += v[-1]; __syncwarp(scan8Mask); v[0] = sum; __syncwarp(scan8Mask); - #endif - #if (CR_BIN_SQR > 2 * 32) - sum += v[-2]; __syncwarp(scan8Mask); v[0] = sum; __syncwarp(scan8Mask); - #endif - #if (CR_BIN_SQR > 4 * 32) - sum += v[-4]; __syncwarp(scan8Mask); v[0] = sum; __syncwarp(scan8Mask); - #endif - - if (thrInBlock == CR_BIN_SQR / 32 - 1) - s_firstActiveIdx = atomicAdd(&atomics.numActiveTiles, sum); - } - - // Tile per thread: Output active tiles. - - __syncthreads(); - - for (int tileInBin_base = 0; tileInBin_base < CR_BIN_SQR; tileInBin_base += CR_COARSE_WARPS * 32) - { - int tileInBin = tileInBin_base + thrInBlock; - bool act = (tileInBin < CR_BIN_SQR) && (s_tileStreamCurrOfs[tileInBin] >= 0); - U32 actMask = __ballot_sync(~0u, act); - if (act) - { - int activeIdx = s_firstActiveIdx; - activeIdx += s_scanTemp[0][(tileInBin >> 5) + 15]; - activeIdx += __popc(actMask & getLaneMaskLt()); - activeTiles[activeIdx] = binTileIdx + globalTileIdx(tileInBin, p.widthTiles); - } - } - } -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Constants.hpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Constants.hpp deleted file mode 100644 index 42d757fbba82e75c1358caee7419a4be009d7464..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Constants.hpp +++ /dev/null @@ -1,73 +0,0 @@ -// 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)) - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CudaRaster.cpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CudaRaster.cpp deleted file mode 100644 index 21f9a484324c7e5409b8e10796e7b6d12ffdbe2f..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CudaRaster.cpp +++ /dev/null @@ -1,79 +0,0 @@ -// 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(); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Defs.hpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Defs.hpp deleted file mode 100644 index 5831664434b613c8443f2e1ece0a0ede0505ca80..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Defs.hpp +++ /dev/null @@ -1,90 +0,0 @@ -// 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 -#include - -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 diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/FineRaster.inl b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/FineRaster.inl deleted file mode 100644 index 8704cc8fd5473c67cb9f97bfb42d76ff0539beb8..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/FineRaster.inl +++ /dev/null @@ -1,385 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Utility funcs. -//------------------------------------------------------------------------ - -__device__ __inline__ void initTileZMax(U32& tileZMax, bool& tileZUpd, volatile U32* tileDepth) -{ - tileZMax = CR_DEPTH_MAX; - tileZUpd = (::min(tileDepth[threadIdx.x], tileDepth[threadIdx.x + 32]) < tileZMax); -} - -__device__ __inline__ void updateTileZMax(U32& tileZMax, bool& tileZUpd, volatile U32* tileDepth, volatile U32* temp) -{ - // Entry is warp-coherent. - if (__any_sync(~0u, tileZUpd)) - { - U32 z = ::max(tileDepth[threadIdx.x], tileDepth[threadIdx.x + 32]); __syncwarp(); - temp[threadIdx.x + 16] = z; __syncwarp(); - z = ::max(z, temp[threadIdx.x + 16 - 1]); __syncwarp(); temp[threadIdx.x + 16] = z; __syncwarp(); - z = ::max(z, temp[threadIdx.x + 16 - 2]); __syncwarp(); temp[threadIdx.x + 16] = z; __syncwarp(); - z = ::max(z, temp[threadIdx.x + 16 - 4]); __syncwarp(); temp[threadIdx.x + 16] = z; __syncwarp(); - z = ::max(z, temp[threadIdx.x + 16 - 8]); __syncwarp(); temp[threadIdx.x + 16] = z; __syncwarp(); - z = ::max(z, temp[threadIdx.x + 16 - 16]); __syncwarp(); temp[threadIdx.x + 16] = z; __syncwarp(); - tileZMax = temp[47]; - tileZUpd = false; - } -} - -//------------------------------------------------------------------------ - -__device__ __inline__ void getTriangle(const CRParams& p, S32& triIdx, S32& dataIdx, uint4& triHeader, S32& segment) -{ - const CRTriangleHeader* triHeaderPtr = (const CRTriangleHeader*)p.triHeader + blockIdx.z * p.maxSubtris;; - const S32* tileSegData = (const S32*)p.tileSegData + p.maxTileSegs * CR_TILE_SEG_SIZE * blockIdx.z; - const S32* tileSegNext = (const S32*)p.tileSegNext + p.maxTileSegs * blockIdx.z; - const S32* tileSegCount = (const S32*)p.tileSegCount + p.maxTileSegs * blockIdx.z; - - if (threadIdx.x >= tileSegCount[segment]) - { - triIdx = -1; - dataIdx = -1; - } - else - { - int subtriIdx = tileSegData[segment * CR_TILE_SEG_SIZE + threadIdx.x]; - triIdx = subtriIdx >> 3; - dataIdx = triIdx; - subtriIdx &= 7; - if (subtriIdx != 7) - dataIdx = triHeaderPtr[triIdx].misc + subtriIdx; - triHeader = *((uint4*)triHeaderPtr + dataIdx); - } - - // advance to next segment - segment = tileSegNext[segment]; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ bool earlyZCull(uint4 triHeader, U32 tileZMax) -{ - U32 zmin = triHeader.w & 0xFFFFF000; - return (zmin > tileZMax); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U64 trianglePixelCoverage(const CRParams& p, const uint4& triHeader, int tileX, int tileY, volatile U64* s_cover8x8_lut) -{ - int baseX = (tileX << (CR_TILE_LOG2 + CR_SUBPIXEL_LOG2)) - ((p.widthPixelsVp - 1) << (CR_SUBPIXEL_LOG2 - 1)); - int baseY = (tileY << (CR_TILE_LOG2 + CR_SUBPIXEL_LOG2)) - ((p.heightPixelsVp - 1) << (CR_SUBPIXEL_LOG2 - 1)); - - // extract S16 vertex positions while subtracting tile coordinates - S32 v0x = sub_s16lo_s16lo(triHeader.x, baseX); - S32 v0y = sub_s16hi_s16lo(triHeader.x, baseY); - S32 v01x = sub_s16lo_s16lo(triHeader.y, triHeader.x); - S32 v01y = sub_s16hi_s16hi(triHeader.y, triHeader.x); - S32 v20x = sub_s16lo_s16lo(triHeader.x, triHeader.z); - S32 v20y = sub_s16hi_s16hi(triHeader.x, triHeader.z); - - // extract flipbits - U32 f01 = (triHeader.w >> 6) & 0x3C; - U32 f12 = (triHeader.w >> 2) & 0x3C; - U32 f20 = (triHeader.w << 2) & 0x3C; - - // compute per-edge coverage masks - U64 c01, c12, c20; - c01 = cover8x8_exact_fast(v0x, v0y, v01x, v01y, f01, s_cover8x8_lut); - c12 = cover8x8_exact_fast(v0x + v01x, v0y + v01y, -v01x - v20x, -v01y - v20y, f12, s_cover8x8_lut); - c20 = cover8x8_exact_fast(v0x, v0y, v20x, v20y, f20, s_cover8x8_lut); - - // combine masks - return c01 & c12 & c20; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U32 scan32_value(U32 value, volatile U32* temp) -{ - __syncwarp(); - temp[threadIdx.x + 16] = value; __syncwarp(); - value += temp[threadIdx.x + 16 - 1]; __syncwarp(); temp[threadIdx.x + 16] = value; __syncwarp(); - value += temp[threadIdx.x + 16 - 2]; __syncwarp(); temp[threadIdx.x + 16] = value; __syncwarp(); - value += temp[threadIdx.x + 16 - 4]; __syncwarp(); temp[threadIdx.x + 16] = value; __syncwarp(); - value += temp[threadIdx.x + 16 - 8]; __syncwarp(); temp[threadIdx.x + 16] = value; __syncwarp(); - value += temp[threadIdx.x + 16 - 16]; __syncwarp(); temp[threadIdx.x + 16] = value; __syncwarp(); - return value; -} - -__device__ __inline__ volatile const U32& scan32_total(volatile U32* temp) -{ - return temp[47]; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ S32 findBit(U64 mask, int idx) -{ - U32 x = getLo(mask); - int pop = __popc(x); - bool p = (pop <= idx); - if (p) x = getHi(mask); - if (p) idx -= pop; - int bit = p ? 32 : 0; - - pop = __popc(x & 0x0000ffffu); - p = (pop <= idx); - if (p) x >>= 16; - if (p) bit += 16; - if (p) idx -= pop; - - U32 tmp = x & 0x000000ffu; - pop = __popc(tmp); - p = (pop <= idx); - if (p) tmp = x & 0x0000ff00u; - if (p) idx -= pop; - - return findLeadingOne(tmp) + bit - idx; -} - -//------------------------------------------------------------------------ -// Single-sample implementation. -//------------------------------------------------------------------------ - -__device__ __inline__ void executeROP(U32 color, U32 depth, volatile U32* pColor, volatile U32* pDepth, U32 ropMask) -{ - atomicMin((U32*)pDepth, depth); - __syncwarp(ropMask); - bool act = (depth == *pDepth); - __syncwarp(ropMask); - U32 actMask = __ballot_sync(ropMask, act); - if (act) - { - *pDepth = 0; - __syncwarp(actMask); - atomicMax((U32*)pDepth, threadIdx.x); - __syncwarp(actMask); - if (*pDepth == threadIdx.x) - { - *pDepth = depth; - *pColor = color; - } - __syncwarp(actMask); - } -} - -//------------------------------------------------------------------------ - -__device__ __inline__ void fineRasterImpl(const CRParams p) -{ - // for 20 warps: - __shared__ volatile U64 s_cover8x8_lut[CR_COVER8X8_LUT_SIZE]; // 6KB - __shared__ volatile U32 s_tileColor [CR_FINE_MAX_WARPS][CR_TILE_SQR]; // 5KB - __shared__ volatile U32 s_tileDepth [CR_FINE_MAX_WARPS][CR_TILE_SQR]; // 5KB - __shared__ volatile U32 s_tilePeel [CR_FINE_MAX_WARPS][CR_TILE_SQR]; // 5KB - __shared__ volatile U32 s_triDataIdx [CR_FINE_MAX_WARPS][64]; // 5KB CRTriangleData index - __shared__ volatile U64 s_triangleCov [CR_FINE_MAX_WARPS][64]; // 10KB coverage mask - __shared__ volatile U32 s_triangleFrag[CR_FINE_MAX_WARPS][64]; // 5KB fragment index - __shared__ volatile U32 s_temp [CR_FINE_MAX_WARPS][80]; // 6.25KB - // = 47.25KB total - - CRAtomics& atomics = p.atomics[blockIdx.z]; - const CRTriangleData* triData = (const CRTriangleData*)p.triData + blockIdx.z * p.maxSubtris; - - const S32* activeTiles = (const S32*)p.activeTiles + CR_MAXTILES_SQR * blockIdx.z; - const S32* tileFirstSeg = (const S32*)p.tileFirstSeg + CR_MAXTILES_SQR * blockIdx.z; - - volatile U32* tileColor = s_tileColor[threadIdx.y]; - volatile U32* tileDepth = s_tileDepth[threadIdx.y]; - volatile U32* tilePeel = s_tilePeel[threadIdx.y]; - volatile U32* triDataIdx = s_triDataIdx[threadIdx.y]; - volatile U64* triangleCov = s_triangleCov[threadIdx.y]; - volatile U32* triangleFrag = s_triangleFrag[threadIdx.y]; - volatile U32* temp = s_temp[threadIdx.y]; - - if (atomics.numSubtris > p.maxSubtris || atomics.numBinSegs > p.maxBinSegs || atomics.numTileSegs > p.maxTileSegs) - return; - - temp[threadIdx.x] = 0; // first 16 elements of temp are always zero - cover8x8_setupLUT(s_cover8x8_lut); - __syncthreads(); - - // loop over tiles - for (;;) - { - // pick a tile - if (threadIdx.x == 0) - temp[16] = atomicAdd(&atomics.fineCounter, 1); - __syncwarp(); - int activeIdx = temp[16]; - if (activeIdx >= atomics.numActiveTiles) - break; - - int tileIdx = activeTiles[activeIdx]; - S32 segment = tileFirstSeg[tileIdx]; - int tileY = tileIdx / p.widthTiles; - int tileX = tileIdx - tileY * p.widthTiles; - int px = (tileX << CR_TILE_LOG2) + (threadIdx.x & (CR_TILE_SIZE - 1)); - int py = (tileY << CR_TILE_LOG2) + (threadIdx.x >> CR_TILE_LOG2); - - // initialize per-tile state - int triRead = 0, triWrite = 0; - int fragRead = 0, fragWrite = 0; - if (threadIdx.x == 0) - triangleFrag[63] = 0; // "previous triangle" - - // deferred clear => clear tile - if (p.deferredClear) - { - tileColor[threadIdx.x] = p.clearColor; - tileDepth[threadIdx.x] = p.clearDepth; - tileColor[threadIdx.x + 32] = p.clearColor; - tileDepth[threadIdx.x + 32] = p.clearDepth; - } - else // otherwise => read tile from framebuffer - { - U32* pColor = (U32*)p.colorBuffer + p.strideX * p.strideY * blockIdx.z; - U32* pDepth = (U32*)p.depthBuffer + p.strideX * p.strideY * blockIdx.z; - tileColor[threadIdx.x] = pColor[px + p.strideX * py]; - tileDepth[threadIdx.x] = pDepth[px + p.strideX * py]; - tileColor[threadIdx.x + 32] = pColor[px + p.strideX * (py + 4)]; - tileDepth[threadIdx.x + 32] = pDepth[px + p.strideX * (py + 4)]; - } - - // read peeling inputs if enabled - if (p.renderModeFlags & CudaRaster::RenderModeFlag_EnableDepthPeeling) - { - U32* pPeel = (U32*)p.peelBuffer + p.strideX * p.strideY * blockIdx.z; - tilePeel[threadIdx.x] = pPeel[px + p.strideX * py]; - tilePeel[threadIdx.x + 32] = pPeel[px + p.strideX * (py + 4)]; - } - - U32 tileZMax; - bool tileZUpd; - initTileZMax(tileZMax, tileZUpd, tileDepth); - - // process fragments - for(;;) - { - // need to queue more fragments? - if (fragWrite - fragRead < 32 && segment >= 0) - { - // update tile z - coherent over warp - updateTileZMax(tileZMax, tileZUpd, tileDepth, temp); - - // read triangles - do - { - // read triangle index and data, advance to next segment - S32 triIdx, dataIdx; - uint4 triHeader; - getTriangle(p, triIdx, dataIdx, triHeader, segment); - - // early z cull - if (triIdx >= 0 && earlyZCull(triHeader, tileZMax)) - triIdx = -1; - - // determine coverage - U64 coverage = trianglePixelCoverage(p, triHeader, tileX, tileY, s_cover8x8_lut); - S32 pop = (triIdx == -1) ? 0 : __popcll(coverage); - - // fragment count scan - U32 frag = scan32_value(pop, temp); - frag += fragWrite; // frag now holds cumulative fragment count - fragWrite += scan32_total(temp); - - // queue non-empty triangles - U32 goodMask = __ballot_sync(~0u, pop != 0); - if (pop != 0) - { - int idx = (triWrite + __popc(goodMask & getLaneMaskLt())) & 63; - triDataIdx [idx] = dataIdx; - triangleFrag[idx] = frag; - triangleCov [idx] = coverage; - } - triWrite += __popc(goodMask); - } - while (fragWrite - fragRead < 32 && segment >= 0); - } - __syncwarp(); - - // end of segment? - if (fragRead == fragWrite) - break; - - // clear triangle boundaries - temp[threadIdx.x + 16] = 0; - __syncwarp(); - - // tag triangle boundaries - if (triRead + threadIdx.x < triWrite) - { - int idx = triangleFrag[(triRead + threadIdx.x) & 63] - fragRead; - if (idx <= 32) - temp[idx + 16 - 1] = 1; - } - __syncwarp(); - - int ropLaneIdx = threadIdx.x; - U32 boundaryMask = __ballot_sync(~0u, temp[ropLaneIdx + 16]); - - // distribute fragments - bool hasFragment = (ropLaneIdx < fragWrite - fragRead); - U32 fragmentMask = __ballot_sync(~0u, hasFragment); - if (hasFragment) - { - int triBufIdx = (triRead + __popc(boundaryMask & getLaneMaskLt())) & 63; - int fragIdx = add_sub(fragRead, ropLaneIdx, triangleFrag[(triBufIdx - 1) & 63]); - U64 coverage = triangleCov[triBufIdx]; - int pixelInTile = findBit(coverage, fragIdx); - int dataIdx = triDataIdx[triBufIdx]; - - // determine pixel position - U32 pixelX = (tileX << CR_TILE_LOG2) + (pixelInTile & 7); - U32 pixelY = (tileY << CR_TILE_LOG2) + (pixelInTile >> 3); - - // depth test - U32 depth = 0; - uint4 td = *((uint4*)triData + dataIdx * (sizeof(CRTriangleData) >> 4)); - - depth = td.x * pixelX + td.y * pixelY + td.z; - bool zkill = (p.renderModeFlags & CudaRaster::RenderModeFlag_EnableDepthPeeling) && (depth <= tilePeel[pixelInTile]); - if (!zkill) - { - U32 oldDepth = tileDepth[pixelInTile]; - if (depth > oldDepth) - zkill = true; - else if (oldDepth == tileZMax) - tileZUpd = true; // we are replacing previous zmax => need to update - } - - U32 ropMask = __ballot_sync(fragmentMask, !zkill); - if (!zkill) - executeROP(td.w, depth, &tileColor[pixelInTile], &tileDepth[pixelInTile], ropMask); - } - // no need to sync, as next up is updateTileZMax that does internal warp sync - - // update counters - fragRead = ::min(fragRead + 32, fragWrite); - triRead += __popc(boundaryMask); - } - - // Write tile back to the framebuffer. - if (true) - { - int px = (tileX << CR_TILE_LOG2) + (threadIdx.x & (CR_TILE_SIZE - 1)); - int py = (tileY << CR_TILE_LOG2) + (threadIdx.x >> CR_TILE_LOG2); - U32* pColor = (U32*)p.colorBuffer + p.strideX * p.strideY * blockIdx.z; - U32* pDepth = (U32*)p.depthBuffer + p.strideX * p.strideY * blockIdx.z; - pColor[px + p.strideX * py] = tileColor[threadIdx.x]; - pDepth[px + p.strideX * py] = tileDepth[threadIdx.x]; - pColor[px + p.strideX * (py + 4)] = tileColor[threadIdx.x + 32]; - pDepth[px + p.strideX * (py + 4)] = tileDepth[threadIdx.x + 32]; - } - } -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/PrivateDefs.hpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/PrivateDefs.hpp deleted file mode 100644 index fda59339b310f51253c29cac42d47086ef35c28a..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/PrivateDefs.hpp +++ /dev/null @@ -1,153 +0,0 @@ -// 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. -}; - -//------------------------------------------------------------------------ -} diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl.cpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl.cpp deleted file mode 100644 index 48643b1e59ff00b4b2c44b3e7aa280bc4f486925..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl.cpp +++ /dev/null @@ -1,370 +0,0 @@ -// 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 - -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)p.strideX) * 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)); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl.hpp b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl.hpp deleted file mode 100644 index 321c7b30787ef50ac0d066f07ad5113c7079f499..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl.hpp +++ /dev/null @@ -1,102 +0,0 @@ -// 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 - diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl_.cu b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl_.cu deleted file mode 100644 index 69c137fdbfaf98f958ed3fe7afebc0d2debf05f1..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/RasterImpl_.cu +++ /dev/null @@ -1,37 +0,0 @@ -// 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); } - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/TriangleSetup.inl b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/TriangleSetup.inl deleted file mode 100644 index 716cb0e7c3d9b4f3a8d639f8ad0e23d425fefd91..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/TriangleSetup.inl +++ /dev/null @@ -1,402 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ - -__device__ __inline__ void snapTriangle( - const CRParams& p, - float4 v0, float4 v1, float4 v2, - int2& p0, int2& p1, int2& p2, float3& rcpW, int2& lo, int2& hi) -{ - F32 viewScaleX = (F32)(p.widthPixelsVp << (CR_SUBPIXEL_LOG2 - 1)); - F32 viewScaleY = (F32)(p.heightPixelsVp << (CR_SUBPIXEL_LOG2 - 1)); - rcpW = make_float3(1.0f / v0.w, 1.0f / v1.w, 1.0f / v2.w); - p0 = make_int2(f32_to_s32_sat(v0.x * rcpW.x * viewScaleX), f32_to_s32_sat(v0.y * rcpW.x * viewScaleY)); - p1 = make_int2(f32_to_s32_sat(v1.x * rcpW.y * viewScaleX), f32_to_s32_sat(v1.y * rcpW.y * viewScaleY)); - p2 = make_int2(f32_to_s32_sat(v2.x * rcpW.z * viewScaleX), f32_to_s32_sat(v2.y * rcpW.z * viewScaleY)); - lo = make_int2(min_min(p0.x, p1.x, p2.x), min_min(p0.y, p1.y, p2.y)); - hi = make_int2(max_max(p0.x, p1.x, p2.x), max_max(p0.y, p1.y, p2.y)); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U32 cover8x8_selectFlips(S32 dx, S32 dy) // 10 instr -{ - U32 flips = 0; - if (dy > 0 || (dy == 0 && dx <= 0)) - flips ^= (1 << CR_FLIPBIT_FLIP_X) ^ (1 << CR_FLIPBIT_FLIP_Y) ^ (1 << CR_FLIPBIT_COMPL); - if (dx > 0) - flips ^= (1 << CR_FLIPBIT_FLIP_X) ^ (1 << CR_FLIPBIT_FLIP_Y); - if (::abs(dx) < ::abs(dy)) - flips ^= (1 << CR_FLIPBIT_SWAP_XY) ^ (1 << CR_FLIPBIT_FLIP_Y); - return flips; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ bool prepareTriangle( - const CRParams& p, - int2 p0, int2 p1, int2 p2, int2 lo, int2 hi, - int2& d1, int2& d2, S32& area) -{ - // Backfacing or degenerate => cull. - - d1 = make_int2(p1.x - p0.x, p1.y - p0.y); - d2 = make_int2(p2.x - p0.x, p2.y - p0.y); - area = d1.x * d2.y - d1.y * d2.x; - - if (area == 0) - return false; // Degenerate. - - if (area < 0 && (p.renderModeFlags & CudaRaster::RenderModeFlag_EnableBackfaceCulling) != 0) - return false; // Backfacing. - - // AABB falls between samples => cull. - - int sampleSize = 1 << CR_SUBPIXEL_LOG2; - int biasX = (p.widthPixelsVp << (CR_SUBPIXEL_LOG2 - 1)) - (sampleSize >> 1); - int biasY = (p.heightPixelsVp << (CR_SUBPIXEL_LOG2 - 1)) - (sampleSize >> 1); - int lox = (int)add_add(lo.x, sampleSize - 1, biasX) & -sampleSize; - int loy = (int)add_add(lo.y, sampleSize - 1, biasY) & -sampleSize; - int hix = (hi.x + biasX) & -sampleSize; - int hiy = (hi.y + biasY) & -sampleSize; - - if (lox > hix || loy > hiy) - return false; // Between pixels. - - // AABB covers 1 or 2 samples => cull if they are not covered. - - int diff = add_sub(hix, hiy, lox) - loy; - if (diff <= sampleSize) - { - int2 t0 = make_int2(add_sub(p0.x, biasX, lox), add_sub(p0.y, biasY, loy)); - int2 t1 = make_int2(add_sub(p1.x, biasX, lox), add_sub(p1.y, biasY, loy)); - int2 t2 = make_int2(add_sub(p2.x, biasX, lox), add_sub(p2.y, biasY, loy)); - S32 e0 = t0.x * t1.y - t0.y * t1.x; - S32 e1 = t1.x * t2.y - t1.y * t2.x; - S32 e2 = t2.x * t0.y - t2.y * t0.x; - if (area < 0) - { - e0 = -e0; - e1 = -e1; - e2 = -e2; - } - - if (e0 < 0 || e1 < 0 || e2 < 0) - { - if (diff == 0) - return false; // Between pixels. - - t0 = make_int2(add_sub(p0.x, biasX, hix), add_sub(p0.y, biasY, hiy)); - t1 = make_int2(add_sub(p1.x, biasX, hix), add_sub(p1.y, biasY, hiy)); - t2 = make_int2(add_sub(p2.x, biasX, hix), add_sub(p2.y, biasY, hiy)); - e0 = t0.x * t1.y - t0.y * t1.x; - e1 = t1.x * t2.y - t1.y * t2.x; - e2 = t2.x * t0.y - t2.y * t0.x; - if (area < 0) - { - e0 = -e0; - e1 = -e1; - e2 = -e2; - } - - if (e0 < 0 || e1 < 0 || e2 < 0) - return false; // Between pixels. - } - } - - // Otherwise => proceed to output the triangle. - - return true; // Visible. -} - -//------------------------------------------------------------------------ - -__device__ __inline__ void setupTriangle( - const CRParams& p, - CRTriangleHeader* th, CRTriangleData* td, int triId, - float v0z, float v1z, float v2z, - int2 p0, int2 p1, int2 p2, float3 rcpW, - int2 d1, int2 d2, S32 area) -{ - // Swap vertices 1 and 2 if area is negative. Only executed if backface culling is - // disabled (if it is enabled, we never come here with area < 0). - - if (area < 0) - { - swap(d1, d2); - swap(p1, p2); - swap(v1z, v2z); - swap(rcpW.y, rcpW.z); - area = -area; - } - - int2 wv0; - wv0.x = p0.x + (p.widthPixelsVp << (CR_SUBPIXEL_LOG2 - 1)); - wv0.y = p0.y + (p.heightPixelsVp << (CR_SUBPIXEL_LOG2 - 1)); - - // Setup depth plane equation. - - F32 zcoef = (F32)(CR_DEPTH_MAX - CR_DEPTH_MIN) * 0.5f; - F32 zbias = (F32)(CR_DEPTH_MAX + CR_DEPTH_MIN) * 0.5f; - float3 zvert = make_float3( - (v0z * zcoef) * rcpW.x + zbias, - (v1z * zcoef) * rcpW.y + zbias, - (v2z * zcoef) * rcpW.z + zbias - ); - int2 zv0 = make_int2( - wv0.x - (1 << (CR_SUBPIXEL_LOG2 - 1)), - wv0.y - (1 << (CR_SUBPIXEL_LOG2 - 1)) - ); - uint3 zpleq = setupPleq(zvert, zv0, d1, d2, 1.0f / (F32)area); - - U32 zmin = f32_to_u32_sat(fminf(fminf(zvert.x, zvert.y), zvert.z) - (F32)CR_LERP_ERROR(0)); - - // Write CRTriangleData. - - *(uint4*)td = make_uint4(zpleq.x, zpleq.y, zpleq.z, triId); - - // Determine flipbits. - - U32 f01 = cover8x8_selectFlips(d1.x, d1.y); - U32 f12 = cover8x8_selectFlips(d2.x - d1.x, d2.y - d1.y); - U32 f20 = cover8x8_selectFlips(-d2.x, -d2.y); - - // Write CRTriangleHeader. - - *(uint4*)th = make_uint4( - prmt(p0.x, p0.y, 0x5410), - prmt(p1.x, p1.y, 0x5410), - prmt(p2.x, p2.y, 0x5410), - (zmin & 0xfffff000u) | (f01 << 6) | (f12 << 2) | (f20 >> 2)); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ void triangleSetupImpl(const CRParams p) -{ - __shared__ F32 s_bary[CR_SETUP_WARPS * 32][18]; - F32* bary = s_bary[threadIdx.x + threadIdx.y * 32]; - - // Compute task and image indices. - - int taskIdx = threadIdx.x + 32 * (threadIdx.y + CR_SETUP_WARPS * blockIdx.x); - int imageIdx = 0; - if (p.instanceMode) - { - imageIdx = blockIdx.z; - if (taskIdx >= p.numTriangles) - return; - } - else - { - while (imageIdx < p.numImages) - { - int count = getImageParams(p, imageIdx).triCount; - if (taskIdx < count) - break; - taskIdx -= count; - imageIdx += 1; - } - if (imageIdx == p.numImages) - return; - } - - // Per-image data structures. - - const CRImageParams& ip = getImageParams(p, imageIdx); - CRAtomics& atomics = p.atomics[imageIdx]; - - const int* indexBuffer = (const int*)p.indexBuffer; - U8* triSubtris = (U8*)p.triSubtris + imageIdx * p.maxSubtris; - CRTriangleHeader* triHeader = (CRTriangleHeader*)p.triHeader + imageIdx * p.maxSubtris; - CRTriangleData* triData = (CRTriangleData*)p.triData + imageIdx * p.maxSubtris; - - // Determine triangle index. - - int triIdx = taskIdx; - if (!p.instanceMode) - triIdx += ip.triOffset; - - // Read vertex indices. - - if ((U32)triIdx >= (U32)p.numTriangles) - { - // Bad triangle index. - triSubtris[taskIdx] = 0; - return; - } - - uint4 vidx; - vidx.x = indexBuffer[triIdx * 3 + 0]; - vidx.y = indexBuffer[triIdx * 3 + 1]; - vidx.z = indexBuffer[triIdx * 3 + 2]; - vidx.w = triIdx + 1; // Triangle index. - - if (vidx.x >= (U32)p.numVertices || - vidx.y >= (U32)p.numVertices || - vidx.z >= (U32)p.numVertices) - { - // Bad vertex index. - triSubtris[taskIdx] = 0; - return; - } - - // Read vertex positions. - - const float4* vertexBuffer = (const float4*)p.vertexBuffer; - if (p.instanceMode) - vertexBuffer += p.numVertices * imageIdx; // Instance offset. - - float4 v0 = vertexBuffer[vidx.x]; - float4 v1 = vertexBuffer[vidx.y]; - float4 v2 = vertexBuffer[vidx.z]; - - // Adjust vertex positions according to current viewport size and offset. - - v0.x = v0.x * p.xs + v0.w * p.xo; - v0.y = v0.y * p.ys + v0.w * p.yo; - v1.x = v1.x * p.xs + v1.w * p.xo; - v1.y = v1.y * p.ys + v1.w * p.yo; - v2.x = v2.x * p.xs + v2.w * p.xo; - v2.y = v2.y * p.ys + v2.w * p.yo; - - // Outside view frustum => cull. - - if (v0.w < fabsf(v0.x) | v0.w < fabsf(v0.y) | v0.w < fabsf(v0.z)) - { - if ((v0.w < +v0.x & v1.w < +v1.x & v2.w < +v2.x) | - (v0.w < -v0.x & v1.w < -v1.x & v2.w < -v2.x) | - (v0.w < +v0.y & v1.w < +v1.y & v2.w < +v2.y) | - (v0.w < -v0.y & v1.w < -v1.y & v2.w < -v2.y) | - (v0.w < +v0.z & v1.w < +v1.z & v2.w < +v2.z) | - (v0.w < -v0.z & v1.w < -v1.z & v2.w < -v2.z)) - { - triSubtris[taskIdx] = 0; - return; - } - } - - // Inside depth range => try to snap vertices. - - if (v0.w >= fabsf(v0.z) & v1.w >= fabsf(v1.z) & v2.w >= fabsf(v2.z)) - { - // Inside S16 range and small enough => fast path. - // Note: aabbLimit comes from the fact that cover8x8 - // does not support guardband with maximal viewport. - - int2 p0, p1, p2, lo, hi; - float3 rcpW; - - snapTriangle(p, v0, v1, v2, p0, p1, p2, rcpW, lo, hi); - S32 loxy = ::min(lo.x, lo.y); - S32 hixy = ::max(hi.x, hi.y); - S32 aabbLimit = (1 << (CR_MAXVIEWPORT_LOG2 + CR_SUBPIXEL_LOG2)) - 1; - - if (loxy >= -32768 && hixy <= 32767 && hixy - loxy <= aabbLimit) - { - int2 d1, d2; - S32 area; - bool res = prepareTriangle(p, p0, p1, p2, lo, hi, d1, d2, area); - triSubtris[taskIdx] = res ? 1 : 0; - - if (res) - setupTriangle( - p, - &triHeader[taskIdx], &triData[taskIdx], vidx.w, - v0.z, v1.z, v2.z, - p0, p1, p2, rcpW, - d1, d2, area); - - return; - } - } - - // Clip to view frustum. - - float4 ov0 = v0; - float4 od1 = make_float4(v1.x - v0.x, v1.y - v0.y, v1.z - v0.z, v1.w - v0.w); - float4 od2 = make_float4(v2.x - v0.x, v2.y - v0.y, v2.z - v0.z, v2.w - v0.w); - int numVerts = clipTriangleWithFrustum(bary, &ov0.x, &v1.x, &v2.x, &od1.x, &od2.x); - - // Count non-culled subtriangles. - - v0.x = ov0.x + od1.x * bary[0] + od2.x * bary[1]; - v0.y = ov0.y + od1.y * bary[0] + od2.y * bary[1]; - v0.z = ov0.z + od1.z * bary[0] + od2.z * bary[1]; - v0.w = ov0.w + od1.w * bary[0] + od2.w * bary[1]; - v1.x = ov0.x + od1.x * bary[2] + od2.x * bary[3]; - v1.y = ov0.y + od1.y * bary[2] + od2.y * bary[3]; - v1.z = ov0.z + od1.z * bary[2] + od2.z * bary[3]; - v1.w = ov0.w + od1.w * bary[2] + od2.w * bary[3]; - float4 tv1 = v1; - - int numSubtris = 0; - for (int i = 2; i < numVerts; i++) - { - v2.x = ov0.x + od1.x * bary[i * 2 + 0] + od2.x * bary[i * 2 + 1]; - v2.y = ov0.y + od1.y * bary[i * 2 + 0] + od2.y * bary[i * 2 + 1]; - v2.z = ov0.z + od1.z * bary[i * 2 + 0] + od2.z * bary[i * 2 + 1]; - v2.w = ov0.w + od1.w * bary[i * 2 + 0] + od2.w * bary[i * 2 + 1]; - - int2 p0, p1, p2, lo, hi, d1, d2; - float3 rcpW; - S32 area; - - snapTriangle(p, v0, v1, v2, p0, p1, p2, rcpW, lo, hi); - if (prepareTriangle(p, p0, p1, p2, lo, hi, d1, d2, area)) - numSubtris++; - - v1 = v2; - } - - triSubtris[taskIdx] = numSubtris; - - // Multiple subtriangles => allocate. - - int subtriBase = taskIdx; - if (numSubtris > 1) - { - subtriBase = atomicAdd(&atomics.numSubtris, numSubtris); - triHeader[taskIdx].misc = subtriBase; - if (subtriBase + numSubtris > p.maxSubtris) - numVerts = 0; - } - - // Setup subtriangles. - - v1 = tv1; - for (int i = 2; i < numVerts; i++) - { - v2.x = ov0.x + od1.x * bary[i * 2 + 0] + od2.x * bary[i * 2 + 1]; - v2.y = ov0.y + od1.y * bary[i * 2 + 0] + od2.y * bary[i * 2 + 1]; - v2.z = ov0.z + od1.z * bary[i * 2 + 0] + od2.z * bary[i * 2 + 1]; - v2.w = ov0.w + od1.w * bary[i * 2 + 0] + od2.w * bary[i * 2 + 1]; - - int2 p0, p1, p2, lo, hi, d1, d2; - float3 rcpW; - S32 area; - - snapTriangle(p, v0, v1, v2, p0, p1, p2, rcpW, lo, hi); - if (prepareTriangle(p, p0, p1, p2, lo, hi, d1, d2, area)) - { - setupTriangle( - p, - &triHeader[subtriBase], &triData[subtriBase], vidx.w, - v0.z, v1.z, v2.z, - p0, p1, p2, rcpW, - d1, d2, area); - - subtriBase++; - } - - v1 = v2; - } -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Util.inl b/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Util.inl deleted file mode 100644 index cf5bb79f9606e3b81e904c2815910e3a0626598e..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Util.inl +++ /dev/null @@ -1,452 +0,0 @@ -// 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 "PrivateDefs.hpp" - -namespace CR -{ -//------------------------------------------------------------------------ - -template __device__ __inline__ void swap(T& a, T& b) { T t = a; a = b; b = t; } - -__device__ __inline__ U32 getLo (U64 a) { return __double2loint(__longlong_as_double(a)); } -__device__ __inline__ S32 getLo (S64 a) { return __double2loint(__longlong_as_double(a)); } -__device__ __inline__ U32 getHi (U64 a) { return __double2hiint(__longlong_as_double(a)); } -__device__ __inline__ S32 getHi (S64 a) { return __double2hiint(__longlong_as_double(a)); } -__device__ __inline__ U64 combineLoHi (U32 lo, U32 hi) { return __double_as_longlong(__hiloint2double(hi, lo)); } -__device__ __inline__ S64 combineLoHi (S32 lo, S32 hi) { return __double_as_longlong(__hiloint2double(hi, lo)); } -__device__ __inline__ U32 getLaneMaskLt (void) { U32 r; asm("mov.u32 %0, %lanemask_lt;" : "=r"(r)); return r; } -__device__ __inline__ U32 getLaneMaskLe (void) { U32 r; asm("mov.u32 %0, %lanemask_le;" : "=r"(r)); return r; } -__device__ __inline__ U32 getLaneMaskGt (void) { U32 r; asm("mov.u32 %0, %lanemask_gt;" : "=r"(r)); return r; } -__device__ __inline__ U32 getLaneMaskGe (void) { U32 r; asm("mov.u32 %0, %lanemask_ge;" : "=r"(r)); return r; } -__device__ __inline__ int findLeadingOne (U32 v) { U32 r; asm("bfind.u32 %0, %1;" : "=r"(r) : "r"(v)); return r; } -__device__ __inline__ bool singleLane (void) { return ((::__ballot_sync(~0u, true) & getLaneMaskLt()) == 0); } - -__device__ __inline__ void add_add_carry (U32& rlo, U32 alo, U32 blo, U32& rhi, U32 ahi, U32 bhi) { U64 r = combineLoHi(alo, ahi) + combineLoHi(blo, bhi); rlo = getLo(r); rhi = getHi(r); } -__device__ __inline__ S32 f32_to_s32_sat (F32 a) { S32 v; asm("cvt.rni.sat.s32.f32 %0, %1;" : "=r"(v) : "f"(a)); return v; } -__device__ __inline__ U32 f32_to_u32_sat (F32 a) { U32 v; asm("cvt.rni.sat.u32.f32 %0, %1;" : "=r"(v) : "f"(a)); return v; } -__device__ __inline__ U32 f32_to_u32_sat_rmi (F32 a) { U32 v; asm("cvt.rmi.sat.u32.f32 %0, %1;" : "=r"(v) : "f"(a)); return v; } -__device__ __inline__ U32 f32_to_u8_sat (F32 a) { U32 v; asm("cvt.rni.sat.u8.f32 %0, %1;" : "=r"(v) : "f"(a)); return v; } -__device__ __inline__ S64 f32_to_s64 (F32 a) { S64 v; asm("cvt.rni.s64.f32 %0, %1;" : "=l"(v) : "f"(a)); return v; } -__device__ __inline__ S32 add_s16lo_s16lo (S32 a, S32 b) { S32 v; asm("vadd.s32.s32.s32 %0, %1.h0, %2.h0;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 add_s16hi_s16lo (S32 a, S32 b) { S32 v; asm("vadd.s32.s32.s32 %0, %1.h1, %2.h0;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 add_s16lo_s16hi (S32 a, S32 b) { S32 v; asm("vadd.s32.s32.s32 %0, %1.h0, %2.h1;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 add_s16hi_s16hi (S32 a, S32 b) { S32 v; asm("vadd.s32.s32.s32 %0, %1.h1, %2.h1;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_s16lo_s16lo (S32 a, S32 b) { S32 v; asm("vsub.s32.s32.s32 %0, %1.h0, %2.h0;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_s16hi_s16lo (S32 a, S32 b) { S32 v; asm("vsub.s32.s32.s32 %0, %1.h1, %2.h0;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_s16lo_s16hi (S32 a, S32 b) { S32 v; asm("vsub.s32.s32.s32 %0, %1.h0, %2.h1;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_s16hi_s16hi (S32 a, S32 b) { S32 v; asm("vsub.s32.s32.s32 %0, %1.h1, %2.h1;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_u16lo_u16lo (U32 a, U32 b) { S32 v; asm("vsub.s32.u32.u32 %0, %1.h0, %2.h0;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_u16hi_u16lo (U32 a, U32 b) { S32 v; asm("vsub.s32.u32.u32 %0, %1.h1, %2.h0;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_u16lo_u16hi (U32 a, U32 b) { S32 v; asm("vsub.s32.u32.u32 %0, %1.h0, %2.h1;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ S32 sub_u16hi_u16hi (U32 a, U32 b) { S32 v; asm("vsub.s32.u32.u32 %0, %1.h1, %2.h1;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ U32 add_b0 (U32 a, U32 b) { U32 v; asm("vadd.u32.u32.u32 %0, %1.b0, %2;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ U32 add_b1 (U32 a, U32 b) { U32 v; asm("vadd.u32.u32.u32 %0, %1.b1, %2;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ U32 add_b2 (U32 a, U32 b) { U32 v; asm("vadd.u32.u32.u32 %0, %1.b2, %2;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ U32 add_b3 (U32 a, U32 b) { U32 v; asm("vadd.u32.u32.u32 %0, %1.b3, %2;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ U32 vmad_b0 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b0, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b1 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b2 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b2, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b3 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b3, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b0_b3 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b0, %2.b3, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b1_b3 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b1, %2.b3, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b2_b3 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b2, %2.b3, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 vmad_b3_b3 (U32 a, U32 b, U32 c) { U32 v; asm("vmad.u32.u32.u32 %0, %1.b3, %2.b3, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 add_mask8 (U32 a, U32 b) { U32 v; U32 z=0; asm("vadd.u32.u32.u32 %0.b0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(z)); return v; } -__device__ __inline__ U32 sub_mask8 (U32 a, U32 b) { U32 v; U32 z=0; asm("vsub.u32.u32.u32 %0.b0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(z)); return v; } -__device__ __inline__ S32 max_max (S32 a, S32 b, S32 c) { S32 v; asm("vmax.s32.s32.s32.max %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 min_min (S32 a, S32 b, S32 c) { S32 v; asm("vmin.s32.s32.s32.min %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 max_add (S32 a, S32 b, S32 c) { S32 v; asm("vmax.s32.s32.s32.add %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 min_add (S32 a, S32 b, S32 c) { S32 v; asm("vmin.s32.s32.s32.add %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 add_add (U32 a, U32 b, U32 c) { U32 v; asm("vadd.u32.u32.u32.add %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 sub_add (U32 a, U32 b, U32 c) { U32 v; asm("vsub.u32.u32.u32.add %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 add_sub (U32 a, U32 b, U32 c) { U32 v; asm("vsub.u32.u32.u32.add %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(c), "r"(b)); return v; } -__device__ __inline__ S32 add_clamp_0_x (S32 a, S32 b, S32 c) { S32 v; asm("vadd.u32.s32.s32.sat.min %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 add_clamp_b0 (S32 a, S32 b, S32 c) { S32 v; asm("vadd.u32.s32.s32.sat %0.b0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 add_clamp_b2 (S32 a, S32 b, S32 c) { S32 v; asm("vadd.u32.s32.s32.sat %0.b2, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ U32 prmt (U32 a, U32 b, U32 c) { U32 v; asm("prmt.b32 %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 u32lo_sext (U32 a) { U32 v; asm("cvt.s16.u32 %0, %1;" : "=r"(v) : "r"(a)); return v; } -__device__ __inline__ U32 slct (U32 a, U32 b, S32 c) { U32 v; asm("slct.u32.s32 %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ S32 slct (S32 a, S32 b, S32 c) { S32 v; asm("slct.s32.s32 %0, %1, %2, %3;" : "=r"(v) : "r"(a), "r"(b), "r"(c)); return v; } -__device__ __inline__ F32 slct (F32 a, F32 b, S32 c) { F32 v; asm("slct.f32.s32 %0, %1, %2, %3;" : "=f"(v) : "f"(a), "f"(b), "r"(c)); return v; } -__device__ __inline__ U32 isetge (S32 a, S32 b) { U32 v; asm("set.ge.u32.s32 %0, %1, %2;" : "=r"(v) : "r"(a), "r"(b)); return v; } -__device__ __inline__ F64 rcp_approx (F64 a) { F64 v; asm("rcp.approx.ftz.f64 %0, %1;" : "=d"(v) : "d"(a)); return v; } -__device__ __inline__ F32 fma_rm (F32 a, F32 b, F32 c) { F32 v; asm("fma.rm.f32 %0, %1, %2, %3;" : "=f"(v) : "f"(a), "f"(b), "f"(c)); return v; } -__device__ __inline__ U32 idiv_fast (U32 a, U32 b); - -__device__ __inline__ uint3 setupPleq (float3 values, int2 v0, int2 d1, int2 d2, F32 areaRcp); - -__device__ __inline__ void cover8x8_setupLUT (volatile U64* lut); -__device__ __inline__ U64 cover8x8_exact_fast (S32 ox, S32 oy, S32 dx, S32 dy, U32 flips, volatile const U64* lut); // Assumes viewport <= 2^11, subpixels <= 2^4, no guardband. -__device__ __inline__ U64 cover8x8_lookupMask (S64 yinit, U32 yinc, U32 flips, volatile const U64* lut); - -__device__ __inline__ U64 cover8x8_exact_noLUT (S32 ox, S32 oy, S32 dx, S32 dy); // optimized reference implementation, does not require look-up table -__device__ __inline__ U64 cover8x8_conservative_noLUT (S32 ox, S32 oy, S32 dx, S32 dy); -__device__ __inline__ U64 cover8x8_generateMask_noLUT (S32 curr, S32 dx, S32 dy); - -template __device__ __inline__ void sortShared(T* ptr, int numItems); // Assumes that numItems <= threadsInBlock. Must sync before & after the call. - -__device__ __inline__ const CRImageParams& getImageParams(const CRParams& p, int idx) -{ - return (idx < CR_EMBED_IMAGE_PARAMS) ? p.imageParamsFirst[idx] : p.imageParamsExtra[idx - CR_EMBED_IMAGE_PARAMS]; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ int clipPolygonWithPlane(F32* baryOut, const F32* baryIn, int numIn, F32 v0, F32 v1, F32 v2) -{ - int numOut = 0; - if (numIn >= 3) - { - int ai = (numIn - 1) * 2; - F32 av = v0 + v1 * baryIn[ai + 0] + v2 * baryIn[ai + 1]; - for (int bi = 0; bi < numIn * 2; bi += 2) - { - F32 bv = v0 + v1 * baryIn[bi + 0] + v2 * baryIn[bi + 1]; - if (av * bv < 0.0f) - { - F32 bc = av / (av - bv); - F32 ac = 1.0f - bc; - baryOut[numOut + 0] = baryIn[ai + 0] * ac + baryIn[bi + 0] * bc; - baryOut[numOut + 1] = baryIn[ai + 1] * ac + baryIn[bi + 1] * bc; - numOut += 2; - } - if (bv >= 0.0f) - { - baryOut[numOut + 0] = baryIn[bi + 0]; - baryOut[numOut + 1] = baryIn[bi + 1]; - numOut += 2; - } - ai = bi; - av = bv; - } - } - return (numOut >> 1); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ int clipTriangleWithFrustum(F32* bary, const F32* v0, const F32* v1, const F32* v2, const F32* d1, const F32* d2) -{ - int num = 3; - bary[0] = 0.0f, bary[1] = 0.0f; - bary[2] = 1.0f, bary[3] = 0.0f; - bary[4] = 0.0f, bary[5] = 1.0f; - - if ((v0[3] < fabsf(v0[0])) | (v1[3] < fabsf(v1[0])) | (v2[3] < fabsf(v2[0]))) - { - F32 temp[18]; - num = clipPolygonWithPlane(temp, bary, num, v0[3] + v0[0], d1[3] + d1[0], d2[3] + d2[0]); - num = clipPolygonWithPlane(bary, temp, num, v0[3] - v0[0], d1[3] - d1[0], d2[3] - d2[0]); - } - if ((v0[3] < fabsf(v0[1])) | (v1[3] < fabsf(v1[1])) | (v2[3] < fabsf(v2[1]))) - { - F32 temp[18]; - num = clipPolygonWithPlane(temp, bary, num, v0[3] + v0[1], d1[3] + d1[1], d2[3] + d2[1]); - num = clipPolygonWithPlane(bary, temp, num, v0[3] - v0[1], d1[3] - d1[1], d2[3] - d2[1]); - } - if ((v0[3] < fabsf(v0[2])) | (v1[3] < fabsf(v1[2])) | (v2[3] < fabsf(v2[2]))) - { - F32 temp[18]; - num = clipPolygonWithPlane(temp, bary, num, v0[3] + v0[2], d1[3] + d1[2], d2[3] + d2[2]); - num = clipPolygonWithPlane(bary, temp, num, v0[3] - v0[2], d1[3] - d1[2], d2[3] - d2[2]); - } - return num; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U32 idiv_fast(U32 a, U32 b) -{ - return f32_to_u32_sat_rmi(((F32)a + 0.5f) / (F32)b); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U32 toABGR(float4 color) -{ - // 11 instructions: 4*FFMA, 4*F2I, 3*PRMT - U32 x = f32_to_u32_sat_rmi(fma_rm(color.x, (1 << 24) * 255.0f, (1 << 24) * 0.5f)); - U32 y = f32_to_u32_sat_rmi(fma_rm(color.y, (1 << 24) * 255.0f, (1 << 24) * 0.5f)); - U32 z = f32_to_u32_sat_rmi(fma_rm(color.z, (1 << 24) * 255.0f, (1 << 24) * 0.5f)); - U32 w = f32_to_u32_sat_rmi(fma_rm(color.w, (1 << 24) * 255.0f, (1 << 24) * 0.5f)); - return prmt(prmt(x, y, 0x0073), prmt(z, w, 0x0073), 0x5410); -} - -//------------------------------------------------------------------------ -// v0 = subpixels relative to the bottom-left sampling point - -__device__ __inline__ uint3 setupPleq(float3 values, int2 v0, int2 d1, int2 d2, F32 areaRcp) -{ - F32 mx = fmaxf(fmaxf(values.x, values.y), values.z); - int sh = ::min(::max((__float_as_int(mx) >> 23) - (127 + 22), 0), 8); - S32 t0 = (U32)values.x >> sh; - S32 t1 = ((U32)values.y >> sh) - t0; - S32 t2 = ((U32)values.z >> sh) - t0; - - U32 rcpMant = (__float_as_int(areaRcp) & 0x007FFFFF) | 0x00800000; - int rcpShift = (23 + 127) - (__float_as_int(areaRcp) >> 23); - - uint3 pleq; - S64 xc = ((S64)t1 * d2.y - (S64)t2 * d1.y) * rcpMant; - S64 yc = ((S64)t2 * d1.x - (S64)t1 * d2.x) * rcpMant; - pleq.x = (U32)(xc >> (rcpShift - (sh + CR_SUBPIXEL_LOG2))); - pleq.y = (U32)(yc >> (rcpShift - (sh + CR_SUBPIXEL_LOG2))); - - S32 centerX = (v0.x * 2 + min_min(d1.x, d2.x, 0) + max_max(d1.x, d2.x, 0)) >> (CR_SUBPIXEL_LOG2 + 1); - S32 centerY = (v0.y * 2 + min_min(d1.y, d2.y, 0) + max_max(d1.y, d2.y, 0)) >> (CR_SUBPIXEL_LOG2 + 1); - S32 vcx = v0.x - (centerX << CR_SUBPIXEL_LOG2); - S32 vcy = v0.y - (centerY << CR_SUBPIXEL_LOG2); - - pleq.z = t0 << sh; - pleq.z -= (U32)(((xc >> 13) * vcx + (yc >> 13) * vcy) >> (rcpShift - (sh + 13))); - pleq.z -= pleq.x * centerX + pleq.y * centerY; - return pleq; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ void cover8x8_setupLUT(volatile U64* lut) -{ - for (S32 lutIdx = threadIdx.x + blockDim.x * threadIdx.y; lutIdx < CR_COVER8X8_LUT_SIZE; lutIdx += blockDim.x * blockDim.y) - { - int half = (lutIdx < (12 << 5)) ? 0 : 1; - int yint = (lutIdx >> 5) - half * 12 - 3; - U32 shape = ((lutIdx >> 2) & 7) << (31 - 2); - S32 slctSwapXY = lutIdx << (31 - 1); - S32 slctNegX = lutIdx << (31 - 0); - S32 slctCompl = slctSwapXY ^ slctNegX; - - U64 mask = 0; - int xlo = half * 4; - int xhi = xlo + 4; - for (int x = xlo; x < xhi; x++) - { - int ylo = slct(0, ::max(yint, 0), slctCompl); - int yhi = slct(::min(yint, 8), 8, slctCompl); - for (int y = ylo; y < yhi; y++) - { - int xx = slct(x, y, slctSwapXY); - int yy = slct(y, x, slctSwapXY); - xx = slct(xx, 7 - xx, slctNegX); - mask |= (U64)1 << (xx + yy * 8); - } - yint += shape >> 31; - shape <<= 1; - } - lut[lutIdx] = mask; - } -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U64 cover8x8_exact_fast(S32 ox, S32 oy, S32 dx, S32 dy, U32 flips, volatile const U64* lut) // 52 instr -{ - F32 yinitBias = (F32)(1 << (31 - CR_MAXVIEWPORT_LOG2 - CR_SUBPIXEL_LOG2 * 2)); - F32 yinitScale = (F32)(1 << (32 - CR_SUBPIXEL_LOG2)); - F32 yincScale = 65536.0f * 65536.0f; - - S32 slctFlipY = flips << (31 - CR_FLIPBIT_FLIP_Y); - S32 slctFlipX = flips << (31 - CR_FLIPBIT_FLIP_X); - S32 slctSwapXY = flips << (31 - CR_FLIPBIT_SWAP_XY); - - // Evaluate cross product. - - S32 t = ox * dy - oy * dx; - F32 det = (F32)slct(t, t - dy * (7 << CR_SUBPIXEL_LOG2), slctFlipX); - if (flips >= (1 << CR_FLIPBIT_COMPL)) - det = -det; - - // Represent Y as a function of X. - - F32 xrcp = 1.0f / (F32)::abs(slct(dx, dy, slctSwapXY)); - F32 yzero = det * yinitScale * xrcp + yinitBias; - S64 yinit = f32_to_s64(slct(yzero, -yzero, slctFlipY)); - U32 yinc = f32_to_u32_sat((F32)::abs(slct(dy, dx, slctSwapXY)) * xrcp * yincScale); - - // Lookup. - - return cover8x8_lookupMask(yinit, yinc, flips, lut); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U64 cover8x8_lookupMask(S64 yinit, U32 yinc, U32 flips, volatile const U64* lut) -{ - // First half. - - U32 yfrac = getLo(yinit); - U32 shape = add_clamp_0_x(getHi(yinit) + 4, 0, 11); - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - int oct = flips & ((1 << CR_FLIPBIT_FLIP_X) | (1 << CR_FLIPBIT_SWAP_XY)); - U64 mask = *(U64*)((U8*)lut + oct + (shape << 5)); - - // Second half. - - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - shape = add_clamp_0_x(getHi(yinit) + 4, __popc(shape & 15), 11); - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - add_add_carry(yfrac, yfrac, yinc, shape, shape, shape); - mask |= *(U64*)((U8*)lut + oct + (shape << 5) + (12 << 8)); - return (flips >= (1 << CR_FLIPBIT_COMPL)) ? ~mask : mask; -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U64 cover8x8_exact_noLUT(S32 ox, S32 oy, S32 dx, S32 dy) -{ - S32 curr = ox * dy - oy * dx; - if (dy > 0 || (dy == 0 && dx <= 0)) curr--; // exclusive - return cover8x8_generateMask_noLUT(curr, dx, dy); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U64 cover8x8_conservative_noLUT(S32 ox, S32 oy, S32 dx, S32 dy) -{ - S32 curr = ox * dy - oy * dx; - if (dy > 0 || (dy == 0 && dx <= 0)) curr--; // exclusive - curr += (::abs(dx) + ::abs(dy)) << (CR_SUBPIXEL_LOG2 - 1); - return cover8x8_generateMask_noLUT(curr, dx, dy); -} - -//------------------------------------------------------------------------ - -__device__ __inline__ U64 cover8x8_generateMask_noLUT(S32 curr, S32 dx, S32 dy) -{ - curr += (dx - dy) * (7 << CR_SUBPIXEL_LOG2); - S32 stepX = dy << (CR_SUBPIXEL_LOG2 + 1); - S32 stepYorig = -dx - dy * 7; - S32 stepY = stepYorig << (CR_SUBPIXEL_LOG2 + 1); - - U32 hi = isetge(curr, 0); - U32 frac = curr + curr; - for (int i = 62; i >= 32; i--) - add_add_carry(frac, frac, ((i & 7) == 7) ? stepY : stepX, hi, hi, hi); - - U32 lo = 0; - for (int i = 31; i >= 0; i--) - add_add_carry(frac, frac, ((i & 7) == 7) ? stepY : stepX, lo, lo, lo); - - lo ^= lo >> 1, hi ^= hi >> 1; - lo ^= lo >> 2, hi ^= hi >> 2; - lo ^= lo >> 4, hi ^= hi >> 4; - lo ^= lo >> 8, hi ^= hi >> 8; - lo ^= lo >> 16, hi ^= hi >> 16; - - if (dy < 0) - { - lo ^= 0x55AA55AA; - hi ^= 0x55AA55AA; - } - if (stepYorig < 0) - { - lo ^= 0xFF00FF00; - hi ^= 0x00FF00FF; - } - if ((hi & 1) != 0) - lo = ~lo; - - return combineLoHi(lo, hi); -} - -//------------------------------------------------------------------------ - -template __device__ __inline__ void sortShared(T* ptr, int numItems) -{ - int thrInBlock = threadIdx.x + threadIdx.y * blockDim.x; - int range = 16; - - // Use transposition sort within each 16-wide subrange. - - int base = thrInBlock * 2; - bool act = (base < numItems - 1); - U32 actMask = __ballot_sync(~0u, act); - if (act) - { - bool tryOdd = (base < numItems - 2 && (~base & (range - 2)) != 0); - T mid = ptr[base + 1]; - - for (int iter = 0; iter < range; iter += 2) - { - // Evens. - - T tmp = ptr[base + 0]; - if (tmp > mid) - { - ptr[base + 0] = mid; - mid = tmp; - } - __syncwarp(actMask); - - // Odds. - - if (tryOdd) - { - tmp = ptr[base + 2]; - if (mid > tmp) - { - ptr[base + 2] = mid; - mid = tmp; - } - } - __syncwarp(actMask); - } - ptr[base + 1] = mid; - } - - // Multiple subranges => Merge hierarchically. - - for (; range < numItems; range <<= 1) - { - // Assuming that we would insert the current item into the other - // subrange, use binary search to find the appropriate slot. - - __syncthreads(); - - T item; - int slot; - if (thrInBlock < numItems) - { - item = ptr[thrInBlock]; - slot = (thrInBlock & -range) ^ range; - if (slot < numItems) - { - T tmp = ptr[slot]; - bool inclusive = ((thrInBlock & range) != 0); - if (tmp < item || (inclusive && tmp == item)) - { - for (int step = (range >> 1); step != 0; step >>= 1) - { - int probe = slot + step; - if (probe < numItems) - { - tmp = ptr[probe]; - if (tmp < item || (inclusive && tmp == item)) - slot = probe; - } - } - slot++; - } - } - } - - // Store the item at an appropriate place. - - __syncthreads(); - - if (thrInBlock < numItems) - ptr[slot + (thrInBlock & (range * 2 - 1)) - range] = item; - } -} - -//------------------------------------------------------------------------ -} diff --git a/extensions/nvdiffrast/nvdiffrast/common/framework.h b/extensions/nvdiffrast/nvdiffrast/common/framework.h deleted file mode 100644 index 3b1b8e993f1c21bd9f0129e14645db9e773cec71..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/framework.h +++ /dev/null @@ -1,49 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Tensorflow. - -#ifdef NVDR_TENSORFLOW -#define EIGEN_USE_GPU -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/shape_inference.h" -#include "tensorflow/core/platform/default/logging.h" -using namespace tensorflow; -using namespace tensorflow::shape_inference; -#define NVDR_CTX_ARGS OpKernelContext* _nvdr_ctx -#define NVDR_CTX_PARAMS _nvdr_ctx -#define NVDR_CHECK(COND, ERR) OP_REQUIRES(_nvdr_ctx, COND, errors::Internal(ERR)) -#define NVDR_CHECK_CUDA_ERROR(CUDA_CALL) OP_CHECK_CUDA_ERROR(_nvdr_ctx, CUDA_CALL) -#define NVDR_CHECK_GL_ERROR(GL_CALL) OP_CHECK_GL_ERROR(_nvdr_ctx, GL_CALL) -#endif - -//------------------------------------------------------------------------ -// PyTorch. - -#ifdef NVDR_TORCH -#ifndef __CUDACC__ -#include -#include -#include -#include -#include -#endif -#define NVDR_CTX_ARGS int _nvdr_ctx_dummy -#define NVDR_CTX_PARAMS 0 -#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) -#define NVDR_CHECK_GL_ERROR(GL_CALL) do { GL_CALL; GLenum err = glGetError(); TORCH_CHECK(err == GL_NO_ERROR, "OpenGL error: ", getGLErrorString(err), "[", #GL_CALL, ";]"); } while(0) -#endif - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/glutil.cpp b/extensions/nvdiffrast/nvdiffrast/common/glutil.cpp deleted file mode 100644 index bea4d86bf8ae95bc958a0cc8dd18238e4b663324..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/glutil.cpp +++ /dev/null @@ -1,403 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Common. -//------------------------------------------------------------------------ - -#include "framework.h" -#include "glutil.h" -#include -#include - -// Create the function pointers. -#define GLUTIL_EXT(return_type, name, ...) return_type (GLAPIENTRY* name)(__VA_ARGS__) = 0; -#include "glutil_extlist.h" -#undef GLUTIL_EXT - -// Track initialization status. -static volatile bool s_glExtInitialized = false; - -// Error strings. -const char* getGLErrorString(GLenum err) -{ - switch(err) - { - case GL_NO_ERROR: return "GL_NO_ERROR"; - case GL_INVALID_ENUM: return "GL_INVALID_ENUM"; - case GL_INVALID_VALUE: return "GL_INVALID_VALUE"; - case GL_INVALID_OPERATION: return "GL_INVALID_OPERATION"; - case GL_STACK_OVERFLOW: return "GL_STACK_OVERFLOW"; - case GL_STACK_UNDERFLOW: return "GL_STACK_UNDERFLOW"; - case GL_OUT_OF_MEMORY: return "GL_OUT_OF_MEMORY"; - case GL_INVALID_FRAMEBUFFER_OPERATION: return "GL_INVALID_FRAMEBUFFER_OPERATION"; - case GL_TABLE_TOO_LARGE: return "GL_TABLE_TOO_LARGE"; - case GL_CONTEXT_LOST: return "GL_CONTEXT_LOST"; - } - return "Unknown error"; -} - -//------------------------------------------------------------------------ -// Windows. -//------------------------------------------------------------------------ - -#ifdef _WIN32 - -static CRITICAL_SECTION getInitializedCriticalSection(void) -{ - CRITICAL_SECTION cs; - InitializeCriticalSection(&cs); - return cs; -} - -static CRITICAL_SECTION s_getProcAddressMutex = getInitializedCriticalSection(); - -static void safeGetProcAddress(const char* name, PROC* pfn) -{ - PROC result = wglGetProcAddress(name); - if (!result) - { - LeaveCriticalSection(&s_getProcAddressMutex); // Prepare for thread exit. - LOG(FATAL) << "wglGetProcAddress() failed for '" << name << "'"; - exit(1); // Should never get here but make sure we exit. - } - *pfn = result; -} - -static void initializeGLExtensions(void) -{ - // Use critical section for thread safety. - EnterCriticalSection(&s_getProcAddressMutex); - - // Only dig function pointers if not done already. - if (!s_glExtInitialized) - { - // Generate code to populate the function pointers. -#define GLUTIL_EXT(return_type, name, ...) safeGetProcAddress(#name, (PROC*)&name); -#include "glutil_extlist.h" -#undef GLUTIL_EXT - - // Mark as initialized. - s_glExtInitialized = true; - } - - // Done. - LeaveCriticalSection(&s_getProcAddressMutex); - return; -} - -void setGLContext(GLContext& glctx) -{ - if (!glctx.hglrc) - LOG(FATAL) << "setGLContext() called with null gltcx"; - if (!wglMakeCurrent(glctx.hdc, glctx.hglrc)) - LOG(FATAL) << "wglMakeCurrent() failed when setting GL context"; - - if (glctx.extInitialized) - return; - initializeGLExtensions(); - glctx.extInitialized = 1; -} - -void releaseGLContext(void) -{ - if (!wglMakeCurrent(NULL, NULL)) - LOG(FATAL) << "wglMakeCurrent() failed when releasing GL context"; -} - -extern "C" int set_gpu(const char*); // In setgpu.lib -GLContext createGLContext(int cudaDeviceIdx) -{ - if (cudaDeviceIdx >= 0) - { - char pciBusId[256] = ""; - LOG(INFO) << "Creating GL context for Cuda device " << cudaDeviceIdx; - if (cudaDeviceGetPCIBusId(pciBusId, 255, cudaDeviceIdx)) - { - LOG(INFO) << "PCI bus id query failed"; - } - else - { - int res = set_gpu(pciBusId); - LOG(INFO) << "Selecting device with PCI bus id " << pciBusId << " - " << (res ? "failed, expect crash or major slowdown" : "success"); - } - } - - HINSTANCE hInstance = GetModuleHandle(NULL); - WNDCLASS wc = {}; - wc.style = CS_OWNDC; - wc.lpfnWndProc = DefWindowProc; - wc.hInstance = hInstance; - wc.lpszClassName = "__DummyGLClassCPP"; - int res = RegisterClass(&wc); - - HWND hwnd = CreateWindow( - "__DummyGLClassCPP", // lpClassName - "__DummyGLWindowCPP", // lpWindowName - WS_OVERLAPPEDWINDOW, // dwStyle - CW_USEDEFAULT, // x - CW_USEDEFAULT, // y - 0, 0, // nWidth, nHeight - NULL, NULL, // hWndParent, hMenu - hInstance, // hInstance - NULL // lpParam - ); - - PIXELFORMATDESCRIPTOR pfd = {}; - pfd.dwFlags = PFD_SUPPORT_OPENGL; - pfd.iPixelType = PFD_TYPE_RGBA; - pfd.iLayerType = PFD_MAIN_PLANE; - pfd.cColorBits = 32; - pfd.cDepthBits = 24; - pfd.cStencilBits = 8; - - HDC hdc = GetDC(hwnd); - int pixelformat = ChoosePixelFormat(hdc, &pfd); - SetPixelFormat(hdc, pixelformat, &pfd); - - HGLRC hglrc = wglCreateContext(hdc); - LOG(INFO) << std::hex << std::setfill('0') - << "WGL OpenGL context created (hdc: 0x" << std::setw(8) << (uint32_t)(uintptr_t)hdc - << ", hglrc: 0x" << std::setw(8) << (uint32_t)(uintptr_t)hglrc << ")"; - - GLContext glctx = {hdc, hglrc, 0}; - return glctx; -} - -void destroyGLContext(GLContext& glctx) -{ - if (!glctx.hglrc) - LOG(FATAL) << "destroyGLContext() called with null gltcx"; - - // If this is the current context, release it. - if (wglGetCurrentContext() == glctx.hglrc) - releaseGLContext(); - - HWND hwnd = WindowFromDC(glctx.hdc); - if (!hwnd) - LOG(FATAL) << "WindowFromDC() failed"; - if (!ReleaseDC(hwnd, glctx.hdc)) - LOG(FATAL) << "ReleaseDC() failed"; - if (!wglDeleteContext(glctx.hglrc)) - LOG(FATAL) << "wglDeleteContext() failed"; - if (!DestroyWindow(hwnd)) - LOG(FATAL) << "DestroyWindow() failed"; - - LOG(INFO) << std::hex << std::setfill('0') - << "WGL OpenGL context destroyed (hdc: 0x" << std::setw(8) << (uint32_t)(uintptr_t)glctx.hdc - << ", hglrc: 0x" << std::setw(8) << (uint32_t)(uintptr_t)glctx.hglrc << ")"; - - memset(&glctx, 0, sizeof(GLContext)); -} - -#endif // _WIN32 - -//------------------------------------------------------------------------ -// Linux. -//------------------------------------------------------------------------ - -#ifdef __linux__ - -static pthread_mutex_t s_getProcAddressMutex; - -typedef void (*PROCFN)(); - -static void safeGetProcAddress(const char* name, PROCFN* pfn) -{ - PROCFN result = eglGetProcAddress(name); - if (!result) - { - pthread_mutex_unlock(&s_getProcAddressMutex); // Prepare for thread exit. - LOG(FATAL) << "wglGetProcAddress() failed for '" << name << "'"; - exit(1); // Should never get here but make sure we exit. - } - *pfn = result; -} - -static void initializeGLExtensions(void) -{ - pthread_mutex_lock(&s_getProcAddressMutex); - - // Only dig function pointers if not done already. - if (!s_glExtInitialized) - { - // Generate code to populate the function pointers. -#define GLUTIL_EXT(return_type, name, ...) safeGetProcAddress(#name, (PROCFN*)&name); -#include "glutil_extlist.h" -#undef GLUTIL_EXT - - // Mark as initialized. - s_glExtInitialized = true; - } - - pthread_mutex_unlock(&s_getProcAddressMutex); - return; -} - -void setGLContext(GLContext& glctx) -{ - if (!glctx.context) - LOG(FATAL) << "setGLContext() called with null gltcx"; - - if (!eglMakeCurrent(glctx.display, EGL_NO_SURFACE, EGL_NO_SURFACE, glctx.context)) - LOG(ERROR) << "eglMakeCurrent() failed when setting GL context"; - - if (glctx.extInitialized) - return; - initializeGLExtensions(); - glctx.extInitialized = 1; -} - -void releaseGLContext(void) -{ - EGLDisplay display = eglGetCurrentDisplay(); - if (display == EGL_NO_DISPLAY) - LOG(WARNING) << "releaseGLContext() called with no active display"; - if (!eglMakeCurrent(display, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT)) - LOG(FATAL) << "eglMakeCurrent() failed when releasing GL context"; -} - -static EGLDisplay getCudaDisplay(int cudaDeviceIdx) -{ - typedef EGLBoolean (*eglQueryDevicesEXT_t)(EGLint, EGLDeviceEXT, EGLint*); - typedef EGLBoolean (*eglQueryDeviceAttribEXT_t)(EGLDeviceEXT, EGLint, EGLAttrib*); - typedef EGLDisplay (*eglGetPlatformDisplayEXT_t)(EGLenum, void*, const EGLint*); - - eglQueryDevicesEXT_t eglQueryDevicesEXT = (eglQueryDevicesEXT_t)eglGetProcAddress("eglQueryDevicesEXT"); - if (!eglQueryDevicesEXT) - { - LOG(INFO) << "eglGetProcAddress(\"eglQueryDevicesEXT\") failed"; - return 0; - } - - eglQueryDeviceAttribEXT_t eglQueryDeviceAttribEXT = (eglQueryDeviceAttribEXT_t)eglGetProcAddress("eglQueryDeviceAttribEXT"); - if (!eglQueryDeviceAttribEXT) - { - LOG(INFO) << "eglGetProcAddress(\"eglQueryDeviceAttribEXT\") failed"; - return 0; - } - - eglGetPlatformDisplayEXT_t eglGetPlatformDisplayEXT = (eglGetPlatformDisplayEXT_t)eglGetProcAddress("eglGetPlatformDisplayEXT"); - if (!eglGetPlatformDisplayEXT) - { - LOG(INFO) << "eglGetProcAddress(\"eglGetPlatformDisplayEXT\") failed"; - return 0; - } - - int num_devices = 0; - eglQueryDevicesEXT(0, 0, &num_devices); - if (!num_devices) - return 0; - - EGLDisplay display = 0; - EGLDeviceEXT* devices = (EGLDeviceEXT*)malloc(num_devices * sizeof(void*)); - eglQueryDevicesEXT(num_devices, devices, &num_devices); - for (int i=0; i < num_devices; i++) - { - EGLDeviceEXT device = devices[i]; - intptr_t value = -1; - if (eglQueryDeviceAttribEXT(device, EGL_CUDA_DEVICE_NV, &value) && value == cudaDeviceIdx) - { - display = eglGetPlatformDisplayEXT(EGL_PLATFORM_DEVICE_EXT, device, 0); - break; - } - } - - free(devices); - return display; -} - -GLContext createGLContext(int cudaDeviceIdx) -{ - EGLDisplay display = 0; - - if (cudaDeviceIdx >= 0) - { - char pciBusId[256] = ""; - LOG(INFO) << "Creating GL context for Cuda device " << cudaDeviceIdx; - display = getCudaDisplay(cudaDeviceIdx); - if (!display) - LOG(INFO) << "Failed, falling back to default display"; - } - - if (!display) - { - display = eglGetDisplay(EGL_DEFAULT_DISPLAY); - if (display == EGL_NO_DISPLAY) - LOG(FATAL) << "eglGetDisplay() failed"; - } - - EGLint major; - EGLint minor; - if (!eglInitialize(display, &major, &minor)) - LOG(FATAL) << "eglInitialize() failed"; - - // Choose configuration. - - const EGLint context_attribs[] = { - EGL_RED_SIZE, 8, - EGL_GREEN_SIZE, 8, - EGL_BLUE_SIZE, 8, - EGL_ALPHA_SIZE, 8, - EGL_DEPTH_SIZE, 24, - EGL_STENCIL_SIZE, 8, - EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT, - EGL_SURFACE_TYPE, EGL_PBUFFER_BIT, - EGL_NONE - }; - - EGLConfig config; - EGLint num_config; - if (!eglChooseConfig(display, context_attribs, &config, 1, &num_config)) - LOG(FATAL) << "eglChooseConfig() failed"; - - // Create GL context. - - if (!eglBindAPI(EGL_OPENGL_API)) - LOG(FATAL) << "eglBindAPI() failed"; - - EGLContext context = eglCreateContext(display, config, EGL_NO_CONTEXT, NULL); - if (context == EGL_NO_CONTEXT) - LOG(FATAL) << "eglCreateContext() failed"; - - // Done. - - LOG(INFO) << "EGL " << (int)minor << "." << (int)major << " OpenGL context created (disp: 0x" - << std::hex << std::setfill('0') - << std::setw(16) << (uintptr_t)display - << ", ctx: 0x" << std::setw(16) << (uintptr_t)context << ")"; - - GLContext glctx = {display, context, 0}; - return glctx; -} - -void destroyGLContext(GLContext& glctx) -{ - if (!glctx.context) - LOG(FATAL) << "destroyGLContext() called with null gltcx"; - - // If this is the current context, release it. - if (eglGetCurrentContext() == glctx.context) - releaseGLContext(); - - if (!eglDestroyContext(glctx.display, glctx.context)) - LOG(ERROR) << "eglDestroyContext() failed"; - - LOG(INFO) << "EGL OpenGL context destroyed (disp: 0x" - << std::hex << std::setfill('0') - << std::setw(16) << (uintptr_t)glctx.display - << ", ctx: 0x" << std::setw(16) << (uintptr_t)glctx.context << ")"; - - memset(&glctx, 0, sizeof(GLContext)); -} - -//------------------------------------------------------------------------ - -#endif // __linux__ - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/glutil.h b/extensions/nvdiffrast/nvdiffrast/common/glutil.h deleted file mode 100644 index 6e4c384626324b5d36e0cc9ac002eacbb10c65f2..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/glutil.h +++ /dev/null @@ -1,113 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// Windows-specific headers and types. -//------------------------------------------------------------------------ - -#ifdef _WIN32 -#define NOMINMAX -#include // Required by gl.h in Windows. -#define GLAPIENTRY APIENTRY - -struct GLContext -{ - HDC hdc; - HGLRC hglrc; - int extInitialized; -}; - -#endif // _WIN32 - -//------------------------------------------------------------------------ -// Linux-specific headers and types. -//------------------------------------------------------------------------ - -#ifdef __linux__ -#define EGL_NO_X11 // X11/Xlib.h has "#define Status int" which breaks Tensorflow. Avoid it. -#define MESA_EGL_NO_X11_HEADERS -#include -#include -#define GLAPIENTRY - -struct GLContext -{ - EGLDisplay display; - EGLContext context; - int extInitialized; -}; - -#endif // __linux__ - -//------------------------------------------------------------------------ -// OpenGL, CUDA interop, GL extensions. -//------------------------------------------------------------------------ -#define GL_GLEXT_LEGACY -#include -#include - -// Constants. -#ifndef GL_VERSION_1_2 -#define GL_CLAMP_TO_EDGE 0x812F -#define GL_TEXTURE_3D 0x806F -#endif -#ifndef GL_VERSION_1_5 -#define GL_ARRAY_BUFFER 0x8892 -#define GL_DYNAMIC_DRAW 0x88E8 -#define GL_ELEMENT_ARRAY_BUFFER 0x8893 -#endif -#ifndef GL_VERSION_2_0 -#define GL_FRAGMENT_SHADER 0x8B30 -#define GL_INFO_LOG_LENGTH 0x8B84 -#define GL_LINK_STATUS 0x8B82 -#define GL_VERTEX_SHADER 0x8B31 -#endif -#ifndef GL_VERSION_3_0 -#define GL_MAJOR_VERSION 0x821B -#define GL_MINOR_VERSION 0x821C -#define GL_RGBA32F 0x8814 -#define GL_TEXTURE_2D_ARRAY 0x8C1A -#endif -#ifndef GL_VERSION_3_2 -#define GL_GEOMETRY_SHADER 0x8DD9 -#endif -#ifndef GL_ARB_framebuffer_object -#define GL_COLOR_ATTACHMENT0 0x8CE0 -#define GL_COLOR_ATTACHMENT1 0x8CE1 -#define GL_DEPTH_STENCIL 0x84F9 -#define GL_DEPTH_STENCIL_ATTACHMENT 0x821A -#define GL_DEPTH24_STENCIL8 0x88F0 -#define GL_FRAMEBUFFER 0x8D40 -#define GL_INVALID_FRAMEBUFFER_OPERATION 0x0506 -#define GL_UNSIGNED_INT_24_8 0x84FA -#endif -#ifndef GL_ARB_imaging -#define GL_TABLE_TOO_LARGE 0x8031 -#endif -#ifndef GL_KHR_robustness -#define GL_CONTEXT_LOST 0x0507 -#endif - -// Declare function pointers to OpenGL extension functions. -#define GLUTIL_EXT(return_type, name, ...) extern return_type (GLAPIENTRY* name)(__VA_ARGS__); -#include "glutil_extlist.h" -#undef GLUTIL_EXT - -//------------------------------------------------------------------------ -// Common functions. -//------------------------------------------------------------------------ - -void setGLContext (GLContext& glctx); -void releaseGLContext (void); -GLContext createGLContext (int cudaDeviceIdx); -void destroyGLContext (GLContext& glctx); -const char* getGLErrorString (GLenum err); - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/glutil_extlist.h b/extensions/nvdiffrast/nvdiffrast/common/glutil_extlist.h deleted file mode 100644 index 6b0fdc2c7820ee4ead0c3e68ad8bafce033c482f..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/glutil_extlist.h +++ /dev/null @@ -1,48 +0,0 @@ -// 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. - -#ifndef GL_VERSION_1_2 -GLUTIL_EXT(void, glTexImage3D, GLenum target, GLint level, GLint internalFormat, GLsizei width, GLsizei height, GLsizei depth, GLint border, GLenum format, GLenum type, const void *pixels); -#endif -#ifndef GL_VERSION_1_5 -GLUTIL_EXT(void, glBindBuffer, GLenum target, GLuint buffer); -GLUTIL_EXT(void, glBufferData, GLenum target, ptrdiff_t size, const void* data, GLenum usage); -GLUTIL_EXT(void, glGenBuffers, GLsizei n, GLuint* buffers); -#endif -#ifndef GL_VERSION_2_0 -GLUTIL_EXT(void, glAttachShader, GLuint program, GLuint shader); -GLUTIL_EXT(void, glCompileShader, GLuint shader); -GLUTIL_EXT(GLuint, glCreateProgram, void); -GLUTIL_EXT(GLuint, glCreateShader, GLenum type); -GLUTIL_EXT(void, glDrawBuffers, GLsizei n, const GLenum* bufs); -GLUTIL_EXT(void, glEnableVertexAttribArray, GLuint index); -GLUTIL_EXT(void, glGetProgramInfoLog, GLuint program, GLsizei bufSize, GLsizei* length, char* infoLog); -GLUTIL_EXT(void, glGetProgramiv, GLuint program, GLenum pname, GLint* param); -GLUTIL_EXT(void, glLinkProgram, GLuint program); -GLUTIL_EXT(void, glShaderSource, GLuint shader, GLsizei count, const char *const* string, const GLint* length); -GLUTIL_EXT(void, glUniform1f, GLint location, GLfloat v0); -GLUTIL_EXT(void, glUniform2f, GLint location, GLfloat v0, GLfloat v1); -GLUTIL_EXT(void, glUseProgram, GLuint program); -GLUTIL_EXT(void, glVertexAttribPointer, GLuint index, GLint size, GLenum type, GLboolean normalized, GLsizei stride, const void* pointer); -#endif -#ifndef GL_VERSION_3_2 -GLUTIL_EXT(void, glFramebufferTexture, GLenum target, GLenum attachment, GLuint texture, GLint level); -#endif -#ifndef GL_ARB_framebuffer_object -GLUTIL_EXT(void, glBindFramebuffer, GLenum target, GLuint framebuffer); -GLUTIL_EXT(void, glGenFramebuffers, GLsizei n, GLuint* framebuffers); -#endif -#ifndef GL_ARB_vertex_array_object -GLUTIL_EXT(void, glBindVertexArray, GLuint array); -GLUTIL_EXT(void, glGenVertexArrays, GLsizei n, GLuint* arrays); -#endif -#ifndef GL_ARB_multi_draw_indirect -GLUTIL_EXT(void, glMultiDrawElementsIndirect, GLenum mode, GLenum type, const void *indirect, GLsizei primcount, GLsizei stride); -#endif - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/interpolate.cu b/extensions/nvdiffrast/nvdiffrast/common/interpolate.cu deleted file mode 100644 index f75f5f1f4ffa39bee40c01dc6b62843e658b83c6..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/interpolate.cu +++ /dev/null @@ -1,276 +0,0 @@ -// 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 -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(p); } -__global__ void InterpolateFwdKernelDa(const InterpolateKernelParams p) { InterpolateFwdKernelTemplate(p); } - -//------------------------------------------------------------------------ -// Gradient kernel. - -template -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(p); } -__global__ void InterpolateGradKernelDa(const InterpolateKernelParams p) { InterpolateGradKernelTemplate(p); } - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/interpolate.h b/extensions/nvdiffrast/nvdiffrast/common/interpolate.h deleted file mode 100644 index 335bb5a8baf53aa96f0cf2f51d0a0fc91ba49527..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/interpolate.h +++ /dev/null @@ -1,49 +0,0 @@ -// 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. -}; - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/rasterize.cu b/extensions/nvdiffrast/nvdiffrast/common/rasterize.cu deleted file mode 100644 index bdc518c9bec3ef45b74dcd75bc8429107b3c9896..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/rasterize.cu +++ /dev/null @@ -1,276 +0,0 @@ -// 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]. - 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 -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(p); } -__global__ void RasterizeGradKernelDb(const RasterizeGradParams p) { RasterizeGradKernelTemplate(p); } - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/rasterize.h b/extensions/nvdiffrast/nvdiffrast/common/rasterize.h deleted file mode 100644 index 30c5e4951239891c4d9c1f28bc0012a39c4a6b06..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/rasterize.h +++ /dev/null @@ -1,60 +0,0 @@ -// 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. -}; - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/rasterize_gl.cpp b/extensions/nvdiffrast/nvdiffrast/common/rasterize_gl.cpp deleted file mode 100644 index 404dcf221bb1182a6755e503927127be113453a8..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/rasterize_gl.cpp +++ /dev/null @@ -1,644 +0,0 @@ -// 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 "rasterize_gl.h" -#include "glutil.h" -#include -#define STRINGIFY_SHADER_SOURCE(x) #x - -//------------------------------------------------------------------------ -// Helpers. - -#define ROUND_UP(x, y) ((((x) + ((y) - 1)) / (y)) * (y)) -static int ROUND_UP_BITS(uint32_t x, uint32_t y) -{ - // Round x up so that it has at most y bits of mantissa. - if (x < (1u << y)) - return x; - uint32_t m = 0; - while (x & ~m) - m = (m << 1) | 1u; - m >>= y; - if (!(x & m)) - return x; - return (x | m) + 1u; -} - -//------------------------------------------------------------------------ -// Draw command struct used by rasterizer. - -struct GLDrawCmd -{ - uint32_t count; - uint32_t instanceCount; - uint32_t firstIndex; - uint32_t baseVertex; - uint32_t baseInstance; -}; - -//------------------------------------------------------------------------ -// GL helpers. - -static void compileGLShader(NVDR_CTX_ARGS, const RasterizeGLState& s, GLuint* pShader, GLenum shaderType, const char* src_buf) -{ - std::string src(src_buf); - - // Set preprocessor directives. - int n = src.find('\n') + 1; // After first line containing #version directive. - if (s.enableZModify) - src.insert(n, "#define IF_ZMODIFY(x) x\n"); - else - src.insert(n, "#define IF_ZMODIFY(x)\n"); - - const char *cstr = src.c_str(); - *pShader = 0; - NVDR_CHECK_GL_ERROR(*pShader = glCreateShader(shaderType)); - NVDR_CHECK_GL_ERROR(glShaderSource(*pShader, 1, &cstr, 0)); - NVDR_CHECK_GL_ERROR(glCompileShader(*pShader)); -} - -static void constructGLProgram(NVDR_CTX_ARGS, GLuint* pProgram, GLuint glVertexShader, GLuint glGeometryShader, GLuint glFragmentShader) -{ - *pProgram = 0; - - GLuint glProgram = 0; - NVDR_CHECK_GL_ERROR(glProgram = glCreateProgram()); - NVDR_CHECK_GL_ERROR(glAttachShader(glProgram, glVertexShader)); - NVDR_CHECK_GL_ERROR(glAttachShader(glProgram, glGeometryShader)); - NVDR_CHECK_GL_ERROR(glAttachShader(glProgram, glFragmentShader)); - NVDR_CHECK_GL_ERROR(glLinkProgram(glProgram)); - - GLint linkStatus = 0; - NVDR_CHECK_GL_ERROR(glGetProgramiv(glProgram, GL_LINK_STATUS, &linkStatus)); - if (!linkStatus) - { - GLint infoLen = 0; - NVDR_CHECK_GL_ERROR(glGetProgramiv(glProgram, GL_INFO_LOG_LENGTH, &infoLen)); - if (infoLen) - { - const char* hdr = "glLinkProgram() failed:\n"; - std::vector info(strlen(hdr) + infoLen); - strcpy(&info[0], hdr); - NVDR_CHECK_GL_ERROR(glGetProgramInfoLog(glProgram, infoLen, &infoLen, &info[strlen(hdr)])); - NVDR_CHECK(0, &info[0]); - } - NVDR_CHECK(0, "glLinkProgram() failed"); - } - - *pProgram = glProgram; -} - -//------------------------------------------------------------------------ -// Shared C++ functions. - -void rasterizeInitGLContext(NVDR_CTX_ARGS, RasterizeGLState& s, int cudaDeviceIdx) -{ - // Create GL context and set it current. - s.glctx = createGLContext(cudaDeviceIdx); - setGLContext(s.glctx); - - // Version check. - GLint vMajor = 0; - GLint vMinor = 0; - glGetIntegerv(GL_MAJOR_VERSION, &vMajor); - glGetIntegerv(GL_MINOR_VERSION, &vMinor); - glGetError(); // Clear possible GL_INVALID_ENUM error in version query. - LOG(INFO) << "OpenGL version reported as " << vMajor << "." << vMinor; - NVDR_CHECK((vMajor == 4 && vMinor >= 4) || vMajor > 4, "OpenGL 4.4 or later is required"); - - // Enable depth modification workaround on A100 and later. - int capMajor = 0; - NVDR_CHECK_CUDA_ERROR(cudaDeviceGetAttribute(&capMajor, cudaDevAttrComputeCapabilityMajor, cudaDeviceIdx)); - s.enableZModify = (capMajor >= 8); - - // Number of output buffers. - int num_outputs = s.enableDB ? 2 : 1; - - // Set up vertex shader. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glVertexShader, GL_VERTEX_SHADER, - "#version 330\n" - "#extension GL_ARB_shader_draw_parameters : enable\n" - STRINGIFY_SHADER_SOURCE( - layout(location = 0) in vec4 in_pos; - out int v_layer; - out int v_offset; - void main() - { - int layer = gl_DrawIDARB; - gl_Position = in_pos; - v_layer = layer; - v_offset = gl_BaseInstanceARB; // Sneak in TriID offset here. - } - ) - ); - - // Geometry and fragment shaders depend on if bary differential output is enabled or not. - if (s.enableDB) - { - // Set up geometry shader. Calculation of per-pixel bary differentials is based on: - // u = (u/w) / (1/w) - // --> du/dX = d((u/w) / (1/w))/dX - // --> du/dX = [d(u/w)/dX - u*d(1/w)/dX] * w - // and we know both d(u/w)/dX and d(1/w)/dX are constant over triangle. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glGeometryShader, GL_GEOMETRY_SHADER, - "#version 430\n" - STRINGIFY_SHADER_SOURCE( - layout(triangles) in; - layout(triangle_strip, max_vertices=3) out; - layout(location = 0) uniform vec2 vp_scale; - in int v_layer[]; - in int v_offset[]; - out vec4 var_uvzw; - out vec4 var_db; - void main() - { - // Plane equations for bary differentials. - float w0 = gl_in[0].gl_Position.w; - float w1 = gl_in[1].gl_Position.w; - float w2 = gl_in[2].gl_Position.w; - vec2 p0 = gl_in[0].gl_Position.xy; - vec2 p1 = gl_in[1].gl_Position.xy; - vec2 p2 = gl_in[2].gl_Position.xy; - vec2 e0 = p0*w2 - p2*w0; - vec2 e1 = p1*w2 - p2*w1; - float a = e0.x*e1.y - e0.y*e1.x; - - // Clamp area to an epsilon to avoid arbitrarily high bary differentials. - float eps = 1e-6f; // ~1 pixel in 1k x 1k image. - float ca = (abs(a) >= eps) ? a : (a < 0.f) ? -eps : eps; // Clamp with sign. - float ia = 1.f / ca; // Inverse area. - - vec2 ascl = ia * vp_scale; - float dudx = e1.y * ascl.x; - float dudy = -e1.x * ascl.y; - float dvdx = -e0.y * ascl.x; - float dvdy = e0.x * ascl.y; - - float duwdx = w2 * dudx; - float dvwdx = w2 * dvdx; - float duvdx = w0 * dudx + w1 * dvdx; - float duwdy = w2 * dudy; - float dvwdy = w2 * dvdy; - float duvdy = w0 * dudy + w1 * dvdy; - - vec4 db0 = vec4(duvdx - dvwdx, duvdy - dvwdy, dvwdx, dvwdy); - vec4 db1 = vec4(duwdx, duwdy, duvdx - duwdx, duvdy - duwdy); - vec4 db2 = vec4(duwdx, duwdy, dvwdx, dvwdy); - - int layer_id = v_layer[0]; - int prim_id = gl_PrimitiveIDIn + v_offset[0]; - - gl_Layer = layer_id; gl_PrimitiveID = prim_id; gl_Position = vec4(gl_in[0].gl_Position.x, gl_in[0].gl_Position.y, gl_in[0].gl_Position.z, gl_in[0].gl_Position.w); var_uvzw = vec4(1.f, 0.f, gl_in[0].gl_Position.z, gl_in[0].gl_Position.w); var_db = db0; EmitVertex(); - gl_Layer = layer_id; gl_PrimitiveID = prim_id; gl_Position = vec4(gl_in[1].gl_Position.x, gl_in[1].gl_Position.y, gl_in[1].gl_Position.z, gl_in[1].gl_Position.w); var_uvzw = vec4(0.f, 1.f, gl_in[1].gl_Position.z, gl_in[1].gl_Position.w); var_db = db1; EmitVertex(); - gl_Layer = layer_id; gl_PrimitiveID = prim_id; gl_Position = vec4(gl_in[2].gl_Position.x, gl_in[2].gl_Position.y, gl_in[2].gl_Position.z, gl_in[2].gl_Position.w); var_uvzw = vec4(0.f, 0.f, gl_in[2].gl_Position.z, gl_in[2].gl_Position.w); var_db = db2; EmitVertex(); - } - ) - ); - - // Set up fragment shader. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glFragmentShader, GL_FRAGMENT_SHADER, - "#version 430\n" - STRINGIFY_SHADER_SOURCE( - in vec4 var_uvzw; - in vec4 var_db; - layout(location = 0) out vec4 out_raster; - layout(location = 1) out vec4 out_db; - IF_ZMODIFY( - layout(location = 1) uniform float in_dummy; - ) - void main() - { - int id_int = gl_PrimitiveID + 1; - float id_float = (id_int <= 0x01000000) ? float(id_int) : intBitsToFloat(0x4a800000 + id_int); - - out_raster = vec4(var_uvzw.x, var_uvzw.y, var_uvzw.z / var_uvzw.w, id_float); - out_db = var_db * var_uvzw.w; - IF_ZMODIFY(gl_FragDepth = gl_FragCoord.z + in_dummy;) - } - ) - ); - - // Set up fragment shader for depth peeling. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glFragmentShaderDP, GL_FRAGMENT_SHADER, - "#version 430\n" - STRINGIFY_SHADER_SOURCE( - in vec4 var_uvzw; - in vec4 var_db; - layout(binding = 0) uniform sampler2DArray out_prev; - layout(location = 0) out vec4 out_raster; - layout(location = 1) out vec4 out_db; - IF_ZMODIFY( - layout(location = 1) uniform float in_dummy; - ) - void main() - { - int id_int = gl_PrimitiveID + 1; - float id_float = (id_int <= 0x01000000) ? float(id_int) : intBitsToFloat(0x4a800000 + id_int); - - vec4 prev = texelFetch(out_prev, ivec3(gl_FragCoord.x, gl_FragCoord.y, gl_Layer), 0); - float depth_new = var_uvzw.z / var_uvzw.w; - if (prev.w == 0 || depth_new <= prev.z) - discard; - out_raster = vec4(var_uvzw.x, var_uvzw.y, depth_new, id_float); - out_db = var_db * var_uvzw.w; - IF_ZMODIFY(gl_FragDepth = gl_FragCoord.z + in_dummy;) - } - ) - ); - } - else - { - // Geometry shader without bary differential output. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glGeometryShader, GL_GEOMETRY_SHADER, - "#version 330\n" - STRINGIFY_SHADER_SOURCE( - layout(triangles) in; - layout(triangle_strip, max_vertices=3) out; - in int v_layer[]; - in int v_offset[]; - out vec4 var_uvzw; - void main() - { - int layer_id = v_layer[0]; - int prim_id = gl_PrimitiveIDIn + v_offset[0]; - - gl_Layer = layer_id; gl_PrimitiveID = prim_id; gl_Position = vec4(gl_in[0].gl_Position.x, gl_in[0].gl_Position.y, gl_in[0].gl_Position.z, gl_in[0].gl_Position.w); var_uvzw = vec4(1.f, 0.f, gl_in[0].gl_Position.z, gl_in[0].gl_Position.w); EmitVertex(); - gl_Layer = layer_id; gl_PrimitiveID = prim_id; gl_Position = vec4(gl_in[1].gl_Position.x, gl_in[1].gl_Position.y, gl_in[1].gl_Position.z, gl_in[1].gl_Position.w); var_uvzw = vec4(0.f, 1.f, gl_in[1].gl_Position.z, gl_in[1].gl_Position.w); EmitVertex(); - gl_Layer = layer_id; gl_PrimitiveID = prim_id; gl_Position = vec4(gl_in[2].gl_Position.x, gl_in[2].gl_Position.y, gl_in[2].gl_Position.z, gl_in[2].gl_Position.w); var_uvzw = vec4(0.f, 0.f, gl_in[2].gl_Position.z, gl_in[2].gl_Position.w); EmitVertex(); - } - ) - ); - - // Fragment shader without bary differential output. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glFragmentShader, GL_FRAGMENT_SHADER, - "#version 430\n" - STRINGIFY_SHADER_SOURCE( - in vec4 var_uvzw; - layout(location = 0) out vec4 out_raster; - IF_ZMODIFY( - layout(location = 1) uniform float in_dummy; - ) - void main() - { - int id_int = gl_PrimitiveID + 1; - float id_float = (id_int <= 0x01000000) ? float(id_int) : intBitsToFloat(0x4a800000 + id_int); - - out_raster = vec4(var_uvzw.x, var_uvzw.y, var_uvzw.z / var_uvzw.w, id_float); - IF_ZMODIFY(gl_FragDepth = gl_FragCoord.z + in_dummy;) - } - ) - ); - - // Depth peeling variant of fragment shader. - compileGLShader(NVDR_CTX_PARAMS, s, &s.glFragmentShaderDP, GL_FRAGMENT_SHADER, - "#version 430\n" - STRINGIFY_SHADER_SOURCE( - in vec4 var_uvzw; - layout(binding = 0) uniform sampler2DArray out_prev; - layout(location = 0) out vec4 out_raster; - IF_ZMODIFY( - layout(location = 1) uniform float in_dummy; - ) - void main() - { - int id_int = gl_PrimitiveID + 1; - float id_float = (id_int <= 0x01000000) ? float(id_int) : intBitsToFloat(0x4a800000 + id_int); - - vec4 prev = texelFetch(out_prev, ivec3(gl_FragCoord.x, gl_FragCoord.y, gl_Layer), 0); - float depth_new = var_uvzw.z / var_uvzw.w; - if (prev.w == 0 || depth_new <= prev.z) - discard; - out_raster = vec4(var_uvzw.x, var_uvzw.y, var_uvzw.z / var_uvzw.w, id_float); - IF_ZMODIFY(gl_FragDepth = gl_FragCoord.z + in_dummy;) - } - ) - ); - } - - // Finalize programs. - constructGLProgram(NVDR_CTX_PARAMS, &s.glProgram, s.glVertexShader, s.glGeometryShader, s.glFragmentShader); - constructGLProgram(NVDR_CTX_PARAMS, &s.glProgramDP, s.glVertexShader, s.glGeometryShader, s.glFragmentShaderDP); - - // Construct main fbo and bind permanently. - NVDR_CHECK_GL_ERROR(glGenFramebuffers(1, &s.glFBO)); - NVDR_CHECK_GL_ERROR(glBindFramebuffer(GL_FRAMEBUFFER, s.glFBO)); - - // Enable two color attachments. - GLenum draw_buffers[2] = { GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT1 }; - NVDR_CHECK_GL_ERROR(glDrawBuffers(num_outputs, draw_buffers)); - - // Construct vertex array object. - NVDR_CHECK_GL_ERROR(glGenVertexArrays(1, &s.glVAO)); - NVDR_CHECK_GL_ERROR(glBindVertexArray(s.glVAO)); - - // Construct position buffer, bind permanently, enable, set ptr. - NVDR_CHECK_GL_ERROR(glGenBuffers(1, &s.glPosBuffer)); - NVDR_CHECK_GL_ERROR(glBindBuffer(GL_ARRAY_BUFFER, s.glPosBuffer)); - NVDR_CHECK_GL_ERROR(glEnableVertexAttribArray(0)); - NVDR_CHECK_GL_ERROR(glVertexAttribPointer(0, 4, GL_FLOAT, GL_FALSE, 0, 0)); - - // Construct index buffer and bind permanently. - NVDR_CHECK_GL_ERROR(glGenBuffers(1, &s.glTriBuffer)); - NVDR_CHECK_GL_ERROR(glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, s.glTriBuffer)); - - // Set up depth test. - NVDR_CHECK_GL_ERROR(glEnable(GL_DEPTH_TEST)); - NVDR_CHECK_GL_ERROR(glDepthFunc(GL_LESS)); - NVDR_CHECK_GL_ERROR(glClearDepth(1.0)); - - // Create and bind output buffers. Storage is allocated later. - NVDR_CHECK_GL_ERROR(glGenTextures(num_outputs, s.glColorBuffer)); - for (int i=0; i < num_outputs; i++) - { - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glColorBuffer[i])); - NVDR_CHECK_GL_ERROR(glFramebufferTexture(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0 + i, s.glColorBuffer[i], 0)); - } - - // Create and bind depth/stencil buffer. Storage is allocated later. - NVDR_CHECK_GL_ERROR(glGenTextures(1, &s.glDepthStencilBuffer)); - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glDepthStencilBuffer)); - NVDR_CHECK_GL_ERROR(glFramebufferTexture(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, s.glDepthStencilBuffer, 0)); - - // Create texture name for previous output buffer (depth peeling). - NVDR_CHECK_GL_ERROR(glGenTextures(1, &s.glPrevOutBuffer)); -} - -void rasterizeResizeBuffers(NVDR_CTX_ARGS, RasterizeGLState& s, bool& changes, int posCount, int triCount, int width, int height, int depth) -{ - changes = false; - - // Resize vertex buffer? - if (posCount > s.posCount) - { - if (s.cudaPosBuffer) - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaPosBuffer)); - s.posCount = (posCount > 64) ? ROUND_UP_BITS(posCount, 2) : 64; - LOG(INFO) << "Increasing position buffer size to " << s.posCount << " float32"; - NVDR_CHECK_GL_ERROR(glBufferData(GL_ARRAY_BUFFER, s.posCount * sizeof(float), NULL, GL_DYNAMIC_DRAW)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsGLRegisterBuffer(&s.cudaPosBuffer, s.glPosBuffer, cudaGraphicsRegisterFlagsWriteDiscard)); - changes = true; - } - - // Resize triangle buffer? - if (triCount > s.triCount) - { - if (s.cudaTriBuffer) - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaTriBuffer)); - s.triCount = (triCount > 64) ? ROUND_UP_BITS(triCount, 2) : 64; - LOG(INFO) << "Increasing triangle buffer size to " << s.triCount << " int32"; - NVDR_CHECK_GL_ERROR(glBufferData(GL_ELEMENT_ARRAY_BUFFER, s.triCount * sizeof(int32_t), NULL, GL_DYNAMIC_DRAW)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsGLRegisterBuffer(&s.cudaTriBuffer, s.glTriBuffer, cudaGraphicsRegisterFlagsWriteDiscard)); - changes = true; - } - - // Resize framebuffer? - if (width > s.width || height > s.height || depth > s.depth) - { - int num_outputs = s.enableDB ? 2 : 1; - if (s.cudaColorBuffer[0]) - for (int i=0; i < num_outputs; i++) - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaColorBuffer[i])); - - if (s.cudaPrevOutBuffer) - { - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaPrevOutBuffer)); - s.cudaPrevOutBuffer = 0; - } - - // New framebuffer size. - s.width = (width > s.width) ? width : s.width; - s.height = (height > s.height) ? height : s.height; - s.depth = (depth > s.depth) ? depth : s.depth; - s.width = ROUND_UP(s.width, 32); - s.height = ROUND_UP(s.height, 32); - LOG(INFO) << "Increasing frame buffer size to (width, height, depth) = (" << s.width << ", " << s.height << ", " << s.depth << ")"; - - // Allocate color buffers. - for (int i=0; i < num_outputs; i++) - { - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glColorBuffer[i])); - NVDR_CHECK_GL_ERROR(glTexImage3D(GL_TEXTURE_2D_ARRAY, 0, GL_RGBA32F, s.width, s.height, s.depth, 0, GL_RGBA, GL_UNSIGNED_BYTE, 0)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MAG_FILTER, GL_NEAREST)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MIN_FILTER, GL_NEAREST)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)); - } - - // Allocate depth/stencil buffer. - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glDepthStencilBuffer)); - NVDR_CHECK_GL_ERROR(glTexImage3D(GL_TEXTURE_2D_ARRAY, 0, GL_DEPTH24_STENCIL8, s.width, s.height, s.depth, 0, GL_DEPTH_STENCIL, GL_UNSIGNED_INT_24_8, 0)); - - // (Re-)register all GL buffers into Cuda. - for (int i=0; i < num_outputs; i++) - NVDR_CHECK_CUDA_ERROR(cudaGraphicsGLRegisterImage(&s.cudaColorBuffer[i], s.glColorBuffer[i], GL_TEXTURE_3D, cudaGraphicsRegisterFlagsReadOnly)); - - changes = true; - } -} - -void rasterizeRender(NVDR_CTX_ARGS, RasterizeGLState& s, cudaStream_t stream, const float* posPtr, int posCount, int vtxPerInstance, const int32_t* triPtr, int triCount, const int32_t* rangesPtr, int width, int height, int depth, int peeling_idx) -{ - // Only copy inputs if we are on first iteration of depth peeling or not doing it at all. - if (peeling_idx < 1) - { - if (triPtr) - { - // Copy both position and triangle buffers. - void* glPosPtr = NULL; - void* glTriPtr = NULL; - size_t posBytes = 0; - size_t triBytes = 0; - NVDR_CHECK_CUDA_ERROR(cudaGraphicsMapResources(2, &s.cudaPosBuffer, stream)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsResourceGetMappedPointer(&glPosPtr, &posBytes, s.cudaPosBuffer)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsResourceGetMappedPointer(&glTriPtr, &triBytes, s.cudaTriBuffer)); - NVDR_CHECK(posBytes >= posCount * sizeof(float), "mapped GL position buffer size mismatch"); - NVDR_CHECK(triBytes >= triCount * sizeof(int32_t), "mapped GL triangle buffer size mismatch"); - NVDR_CHECK_CUDA_ERROR(cudaMemcpyAsync(glPosPtr, posPtr, posCount * sizeof(float), cudaMemcpyDeviceToDevice, stream)); - NVDR_CHECK_CUDA_ERROR(cudaMemcpyAsync(glTriPtr, triPtr, triCount * sizeof(int32_t), cudaMemcpyDeviceToDevice, stream)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnmapResources(2, &s.cudaPosBuffer, stream)); - } - else - { - // Copy position buffer only. Triangles are already copied and known to be constant. - void* glPosPtr = NULL; - size_t posBytes = 0; - NVDR_CHECK_CUDA_ERROR(cudaGraphicsMapResources(1, &s.cudaPosBuffer, stream)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsResourceGetMappedPointer(&glPosPtr, &posBytes, s.cudaPosBuffer)); - NVDR_CHECK(posBytes >= posCount * sizeof(float), "mapped GL position buffer size mismatch"); - NVDR_CHECK_CUDA_ERROR(cudaMemcpyAsync(glPosPtr, posPtr, posCount * sizeof(float), cudaMemcpyDeviceToDevice, stream)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnmapResources(1, &s.cudaPosBuffer, stream)); - } - } - - // Select program based on whether we have a depth peeling input or not. - if (peeling_idx < 1) - { - // Normal case: No peeling, or peeling disabled. - NVDR_CHECK_GL_ERROR(glUseProgram(s.glProgram)); - } - else - { - // If we don't have a third buffer yet, create one. - if (!s.cudaPrevOutBuffer) - { - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glPrevOutBuffer)); - NVDR_CHECK_GL_ERROR(glTexImage3D(GL_TEXTURE_2D_ARRAY, 0, GL_RGBA32F, s.width, s.height, s.depth, 0, GL_RGBA, GL_UNSIGNED_BYTE, 0)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MAG_FILTER, GL_NEAREST)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MIN_FILTER, GL_NEAREST)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)); - NVDR_CHECK_GL_ERROR(glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)); - NVDR_CHECK_CUDA_ERROR(cudaGraphicsGLRegisterImage(&s.cudaPrevOutBuffer, s.glPrevOutBuffer, GL_TEXTURE_3D, cudaGraphicsRegisterFlagsReadOnly)); - } - - // Swap the GL buffers. - GLuint glTempBuffer = s.glPrevOutBuffer; - s.glPrevOutBuffer = s.glColorBuffer[0]; - s.glColorBuffer[0] = glTempBuffer; - - // Swap the Cuda buffers. - cudaGraphicsResource_t cudaTempBuffer = s.cudaPrevOutBuffer; - s.cudaPrevOutBuffer = s.cudaColorBuffer[0]; - s.cudaColorBuffer[0] = cudaTempBuffer; - - // Bind the new output buffer. - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glColorBuffer[0])); - NVDR_CHECK_GL_ERROR(glFramebufferTexture(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, s.glColorBuffer[0], 0)); - - // Bind old buffer as the input texture. - NVDR_CHECK_GL_ERROR(glBindTexture(GL_TEXTURE_2D_ARRAY, s.glPrevOutBuffer)); - - // Activate the correct program. - NVDR_CHECK_GL_ERROR(glUseProgram(s.glProgramDP)); - } - - // Set viewport, clear color buffer(s) and depth/stencil buffer. - NVDR_CHECK_GL_ERROR(glViewport(0, 0, width, height)); - NVDR_CHECK_GL_ERROR(glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT | GL_STENCIL_BUFFER_BIT)); - - // If outputting bary differentials, set resolution uniform - if (s.enableDB) - NVDR_CHECK_GL_ERROR(glUniform2f(0, 2.f / (float)width, 2.f / (float)height)); - - // Set the dummy uniform if depth modification workaround is active. - if (s.enableZModify) - NVDR_CHECK_GL_ERROR(glUniform1f(1, 0.f)); - - // Render the meshes. - if (depth == 1 && !rangesPtr) - { - // Trivial case. - NVDR_CHECK_GL_ERROR(glDrawElements(GL_TRIANGLES, triCount, GL_UNSIGNED_INT, 0)); - } - else - { - // Populate a buffer for draw commands and execute it. - std::vector drawCmdBuffer(depth); - - if (!rangesPtr) - { - // Fill in range array to instantiate the same triangles for each output layer. - // Triangle IDs starts at zero (i.e., one) for each layer, so they correspond to - // the first dimension in addressing the triangle array. - for (int i=0; i < depth; i++) - { - GLDrawCmd& cmd = drawCmdBuffer[i]; - cmd.firstIndex = 0; - cmd.count = triCount; - cmd.baseVertex = vtxPerInstance * i; - cmd.baseInstance = 0; - cmd.instanceCount = 1; - } - } - else - { - // Fill in the range array according to user-given ranges. Triangle IDs point - // to the input triangle array, NOT index within range, so they correspond to - // the first dimension in addressing the triangle array. - for (int i=0, j=0; i < depth; i++) - { - GLDrawCmd& cmd = drawCmdBuffer[i]; - int first = rangesPtr[j++]; - int count = rangesPtr[j++]; - NVDR_CHECK(first >= 0 && count >= 0, "range contains negative values"); - NVDR_CHECK((first + count) * 3 <= triCount, "range extends beyond end of triangle buffer"); - cmd.firstIndex = first * 3; - cmd.count = count * 3; - cmd.baseVertex = 0; - cmd.baseInstance = first; - cmd.instanceCount = 1; - } - } - - // Draw! - NVDR_CHECK_GL_ERROR(glMultiDrawElementsIndirect(GL_TRIANGLES, GL_UNSIGNED_INT, &drawCmdBuffer[0], depth, sizeof(GLDrawCmd))); - } -} - -void rasterizeCopyResults(NVDR_CTX_ARGS, RasterizeGLState& s, cudaStream_t stream, float** outputPtr, int width, int height, int depth) -{ - // Copy color buffers to output tensors. - cudaArray_t array = 0; - cudaChannelFormatDesc arrayDesc = {}; // For error checking. - cudaExtent arrayExt = {}; // For error checking. - int num_outputs = s.enableDB ? 2 : 1; - NVDR_CHECK_CUDA_ERROR(cudaGraphicsMapResources(num_outputs, s.cudaColorBuffer, stream)); - for (int i=0; i < num_outputs; i++) - { - NVDR_CHECK_CUDA_ERROR(cudaGraphicsSubResourceGetMappedArray(&array, s.cudaColorBuffer[i], 0, 0)); - NVDR_CHECK_CUDA_ERROR(cudaArrayGetInfo(&arrayDesc, &arrayExt, NULL, array)); - NVDR_CHECK(arrayDesc.f == cudaChannelFormatKindFloat, "CUDA mapped array data kind mismatch"); - NVDR_CHECK(arrayDesc.x == 32 && arrayDesc.y == 32 && arrayDesc.z == 32 && arrayDesc.w == 32, "CUDA mapped array data width mismatch"); - NVDR_CHECK(arrayExt.width >= width && arrayExt.height >= height && arrayExt.depth >= depth, "CUDA mapped array extent mismatch"); - cudaMemcpy3DParms p = {0}; - p.srcArray = array; - p.dstPtr.ptr = outputPtr[i]; - p.dstPtr.pitch = width * 4 * sizeof(float); - p.dstPtr.xsize = width; - p.dstPtr.ysize = height; - p.extent.width = width; - p.extent.height = height; - p.extent.depth = depth; - p.kind = cudaMemcpyDeviceToDevice; - NVDR_CHECK_CUDA_ERROR(cudaMemcpy3DAsync(&p, stream)); - } - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnmapResources(num_outputs, s.cudaColorBuffer, stream)); -} - -void rasterizeReleaseBuffers(NVDR_CTX_ARGS, RasterizeGLState& s) -{ - int num_outputs = s.enableDB ? 2 : 1; - - if (s.cudaPosBuffer) - { - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaPosBuffer)); - s.cudaPosBuffer = 0; - } - - if (s.cudaTriBuffer) - { - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaTriBuffer)); - s.cudaTriBuffer = 0; - } - - for (int i=0; i < num_outputs; i++) - { - if (s.cudaColorBuffer[i]) - { - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaColorBuffer[i])); - s.cudaColorBuffer[i] = 0; - } - } - - if (s.cudaPrevOutBuffer) - { - NVDR_CHECK_CUDA_ERROR(cudaGraphicsUnregisterResource(s.cudaPrevOutBuffer)); - s.cudaPrevOutBuffer = 0; - } -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/rasterize_gl.h b/extensions/nvdiffrast/nvdiffrast/common/rasterize_gl.h deleted file mode 100644 index ed0f75036a837cf5345a7aa27abcf39bf0ddbd40..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/rasterize_gl.h +++ /dev/null @@ -1,60 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// Do not try to include OpenGL stuff when compiling CUDA kernels for torch. - -#if !(defined(NVDR_TORCH) && defined(__CUDACC__)) -#include "framework.h" -#include "glutil.h" - -//------------------------------------------------------------------------ -// OpenGL-related persistent state for forward op. - -struct RasterizeGLState // Must be initializable by memset to zero. -{ - int width; // Allocated frame buffer width. - int height; // Allocated frame buffer height. - int depth; // Allocated frame buffer depth. - int posCount; // Allocated position buffer in floats. - int triCount; // Allocated triangle buffer in ints. - GLContext glctx; - GLuint glFBO; - GLuint glColorBuffer[2]; - GLuint glPrevOutBuffer; - GLuint glDepthStencilBuffer; - GLuint glVAO; - GLuint glTriBuffer; - GLuint glPosBuffer; - GLuint glProgram; - GLuint glProgramDP; - GLuint glVertexShader; - GLuint glGeometryShader; - GLuint glFragmentShader; - GLuint glFragmentShaderDP; - cudaGraphicsResource_t cudaColorBuffer[2]; - cudaGraphicsResource_t cudaPrevOutBuffer; - cudaGraphicsResource_t cudaPosBuffer; - cudaGraphicsResource_t cudaTriBuffer; - int enableDB; - int enableZModify; // Modify depth in shader, workaround for a rasterization issue on A100. -}; - -//------------------------------------------------------------------------ -// Shared C++ code prototypes. - -void rasterizeInitGLContext(NVDR_CTX_ARGS, RasterizeGLState& s, int cudaDeviceIdx); -void rasterizeResizeBuffers(NVDR_CTX_ARGS, RasterizeGLState& s, bool& changes, int posCount, int triCount, int width, int height, int depth); -void rasterizeRender(NVDR_CTX_ARGS, RasterizeGLState& s, cudaStream_t stream, const float* posPtr, int posCount, int vtxPerInstance, const int32_t* triPtr, int triCount, const int32_t* rangesPtr, int width, int height, int depth, int peeling_idx); -void rasterizeCopyResults(NVDR_CTX_ARGS, RasterizeGLState& s, cudaStream_t stream, float** outputPtr, int width, int height, int depth); -void rasterizeReleaseBuffers(NVDR_CTX_ARGS, RasterizeGLState& s); - -//------------------------------------------------------------------------ -#endif // !(defined(NVDR_TORCH) && defined(__CUDACC__)) diff --git a/extensions/nvdiffrast/nvdiffrast/common/texture.cpp b/extensions/nvdiffrast/nvdiffrast/common/texture.cpp deleted file mode 100644 index 8c38fb5a2d8fe8075719d24b25162c8c34347c02..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/texture.cpp +++ /dev/null @@ -1,104 +0,0 @@ -// 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(NVDR_CTX_ARGS, 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(NVDR_CTX_ARGS, 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(NVDR_CTX_PARAMS, 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; -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/texture.h b/extensions/nvdiffrast/nvdiffrast/common/texture.h deleted file mode 100644 index a8d5ee1b6aa5b19ead3460afd24b9f8e64f456dd..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/texture.h +++ /dev/null @@ -1,78 +0,0 @@ -// 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(NVDR_CTX_ARGS, const TextureKernelParams& p); -int calculateMipInfo(NVDR_CTX_ARGS, 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) - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/common/texture_.cu b/extensions/nvdiffrast/nvdiffrast/common/texture_.cu deleted file mode 100644 index 40bd62974b76a7cd5fa594bb74a1d61330857011..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/common/texture_.cu +++ /dev/null @@ -1,1156 +0,0 @@ -// 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 "texture.h" - -//------------------------------------------------------------------------ -// Memory access and math helpers. - -static __device__ __forceinline__ void accum_from_mem(float* a, int s, float b, float c) { a[0] += b * c; } -static __device__ __forceinline__ void accum_from_mem(float* a, int s, float2 b, float c) { a[0] += b.x * c; a[s] += b.y * c; } -static __device__ __forceinline__ void accum_from_mem(float* a, int s, float4 b, float c) { a[0] += b.x * c; a[s] += b.y * c; a[2*s] += b.z * c; a[3*s] += b.w * c; } -static __device__ __forceinline__ void accum_to_mem(float& a, float* b, int s) { a += b[0]; } -static __device__ __forceinline__ void accum_to_mem(float2& a, float* b, int s) { float2 v = a; v.x += b[0]; v.y += b[s]; a = v; } -static __device__ __forceinline__ void accum_to_mem(float4& a, float* b, int s) { float4 v = a; v.x += b[0]; v.y += b[s]; v.z += b[2*s]; v.w += b[3*s]; a = v; } -static __device__ __forceinline__ bool isfinite_vec3(const float3& a) { return isfinite(a.x) && isfinite(a.y) && isfinite(a.z); } -static __device__ __forceinline__ bool isfinite_vec4(const float4& a) { return isfinite(a.x) && isfinite(a.y) && isfinite(a.z) && isfinite(a.w); } -template static __device__ __forceinline__ T lerp (const T& a, const T& b, float c) { return a + c * (b - a); } -template static __device__ __forceinline__ T bilerp(const T& a, const T& b, const T& c, const T& d, const float2& e) { return lerp(lerp(a, b, e.x), lerp(c, d, e.x), e.y); } - -//------------------------------------------------------------------------ -// Cube map wrapping for smooth filtering across edges and corners. At corners, -// one of the texture coordinates will be negative. For correct interpolation, -// the missing texel must take the average color of the other three. - -static __constant__ uint32_t c_cubeWrapMask1[48] = -{ - 0x1530a440, 0x1133a550, 0x6103a110, 0x1515aa44, 0x6161aa11, 0x40154a04, 0x44115a05, 0x04611a01, - 0x2630a440, 0x2233a550, 0x5203a110, 0x2626aa44, 0x5252aa11, 0x40264a04, 0x44225a05, 0x04521a01, - 0x32608064, 0x3366a055, 0x13062091, 0x32328866, 0x13132299, 0x50320846, 0x55330a55, 0x05130219, - 0x42508064, 0x4455a055, 0x14052091, 0x42428866, 0x14142299, 0x60420846, 0x66440a55, 0x06140219, - 0x5230a044, 0x5533a055, 0x1503a011, 0x5252aa44, 0x1515aa11, 0x40520a44, 0x44550a55, 0x04150a11, - 0x6130a044, 0x6633a055, 0x2603a011, 0x6161aa44, 0x2626aa11, 0x40610a44, 0x44660a55, 0x04260a11, -}; - -static __constant__ uint8_t c_cubeWrapMask2[48] = -{ - 0x26, 0x33, 0x11, 0x05, 0x00, 0x09, 0x0c, 0x04, 0x04, 0x00, 0x00, 0x05, 0x00, 0x81, 0xc0, 0x40, - 0x02, 0x03, 0x09, 0x00, 0x0a, 0x00, 0x00, 0x02, 0x64, 0x30, 0x90, 0x55, 0xa0, 0x99, 0xcc, 0x64, - 0x24, 0x30, 0x10, 0x05, 0x00, 0x01, 0x00, 0x00, 0x06, 0x03, 0x01, 0x05, 0x00, 0x89, 0xcc, 0x44, -}; - -static __device__ __forceinline__ int4 wrapCubeMap(int face, int ix0, int ix1, int iy0, int iy1, int w) -{ - // Calculate case number. - int cx = (ix0 < 0) ? 0 : (ix1 >= w) ? 2 : 1; - int cy = (iy0 < 0) ? 0 : (iy1 >= w) ? 6 : 3; - int c = cx + cy; - if (c >= 5) - c--; - c = (face << 3) + c; - - // Compute coordinates and faces. - unsigned int m = c_cubeWrapMask1[c]; - int x0 = (m >> 0) & 3; x0 = (x0 == 0) ? 0 : (x0 == 1) ? ix0 : iy0; - int x1 = (m >> 2) & 3; x1 = (x1 == 0) ? 0 : (x1 == 1) ? ix1 : iy0; - int x2 = (m >> 4) & 3; x2 = (x2 == 0) ? 0 : (x2 == 1) ? ix0 : iy1; - int x3 = (m >> 6) & 3; x3 = (x3 == 0) ? 0 : (x3 == 1) ? ix1 : iy1; - int y0 = (m >> 8) & 3; y0 = (y0 == 0) ? 0 : (y0 == 1) ? ix0 : iy0; - int y1 = (m >> 10) & 3; y1 = (y1 == 0) ? 0 : (y1 == 1) ? ix1 : iy0; - int y2 = (m >> 12) & 3; y2 = (y2 == 0) ? 0 : (y2 == 1) ? ix0 : iy1; - int y3 = (m >> 14) & 3; y3 = (y3 == 0) ? 0 : (y3 == 1) ? ix1 : iy1; - int f0 = ((m >> 16) & 15) - 1; - int f1 = ((m >> 20) & 15) - 1; - int f2 = ((m >> 24) & 15) - 1; - int f3 = ((m >> 28) ) - 1; - - // Flips. - unsigned int f = c_cubeWrapMask2[c]; - int w1 = w - 1; - if (f & 0x01) x0 = w1 - x0; - if (f & 0x02) x1 = w1 - x1; - if (f & 0x04) x2 = w1 - x2; - if (f & 0x08) x3 = w1 - x3; - if (f & 0x10) y0 = w1 - y0; - if (f & 0x20) y1 = w1 - y1; - if (f & 0x40) y2 = w1 - y2; - if (f & 0x80) y3 = w1 - y3; - - // Done. - int4 tcOut; - tcOut.x = x0 + (y0 + f0 * w) * w; - tcOut.y = x1 + (y1 + f1 * w) * w; - tcOut.z = x2 + (y2 + f2 * w) * w; - tcOut.w = x3 + (y3 + f3 * w) * w; - return tcOut; -} - -//------------------------------------------------------------------------ -// Cube map indexing and gradient functions. - -// Map a 3D lookup vector into an (s,t) face coordinates (returned in first . -// two parameters) and face index. -static __device__ __forceinline__ int indexCubeMap(float& x, float& y, float z) -{ - float ax = fabsf(x); - float ay = fabsf(y); - float az = fabsf(z); - int idx; - float c; - if (az > fmaxf(ax, ay)) { idx = 4; c = z; } - else if (ay > ax) { idx = 2; c = y; y = z; } - else { idx = 0; c = x; x = z; } - if (c < 0.f) idx += 1; - float m = __frcp_rz(fabsf(c)) * .5; - float m0 = __uint_as_float(__float_as_uint(m) ^ ((0x21u >> idx) << 31)); - float m1 = (idx != 2) ? -m : m; - x = x * m0 + .5; - y = y * m1 + .5; - if (!isfinite(x) || !isfinite(y)) - return -1; // Invalid uv. - x = fminf(fmaxf(x, 0.f), 1.f); - y = fminf(fmaxf(y, 0.f), 1.f); - return idx; -} - -// Based on dA/d{s,t}, compute dA/d{x,y,z} at a given 3D lookup vector. -static __device__ __forceinline__ float3 indexCubeMapGrad(float3 uv, float gu, float gv) -{ - float ax = fabsf(uv.x); - float ay = fabsf(uv.y); - float az = fabsf(uv.z); - int idx; - float c; - float c0 = gu; - float c1 = gv; - if (az > fmaxf(ax, ay)) { idx = 0x10; c = uv.z; c0 *= uv.x; c1 *= uv.y; } - else if (ay > ax) { idx = 0x04; c = uv.y; c0 *= uv.x; c1 *= uv.z; } - else { idx = 0x01; c = uv.x; c0 *= uv.z; c1 *= uv.y; } - if (c < 0.f) idx += idx; - float m = __frcp_rz(fabsf(c)); - c0 = (idx & 0x34) ? -c0 : c0; - c1 = (idx & 0x2e) ? -c1 : c1; - float gl = (c0 + c1) * m; - float gx = (idx & 0x03) ? gl : (idx & 0x20) ? -gu : gu; - float gy = (idx & 0x0c) ? gl : -gv; - float gz = (idx & 0x30) ? gl : (idx & 0x03) ? gu : gv; - gz = (idx & 0x09) ? -gz : gz; - float3 res = make_float3(gx, gy, gz) * (m * .5f); - if (!isfinite_vec3(res)) - return make_float3(0.f, 0.f, 0.f); // Invalid uv. - return res; -} - -// Based on dL/d(d{s,t}/s{X,Y}), compute dL/d(d{x,y,z}/d{X,Y}). This is just two -// indexCubeMapGrad() functions rolled together. -static __device__ __forceinline__ void indexCubeMapGrad4(float3 uv, float4 dw, float3& g0, float3& g1) -{ - float ax = fabsf(uv.x); - float ay = fabsf(uv.y); - float az = fabsf(uv.z); - int idx; - float c, c0, c1; - if (az > fmaxf(ax, ay)) { idx = 0x10; c = uv.z; c0 = uv.x; c1 = uv.y; } - else if (ay > ax) { idx = 0x04; c = uv.y; c0 = uv.x; c1 = uv.z; } - else { idx = 0x01; c = uv.x; c0 = uv.z; c1 = uv.y; } - if (c < 0.f) idx += idx; - float m = __frcp_rz(fabsf(c)); - c0 = (idx & 0x34) ? -c0 : c0; - c1 = (idx & 0x2e) ? -c1 : c1; - float gl0 = (dw.x * c0 + dw.z * c1) * m; - float gl1 = (dw.y * c0 + dw.w * c1) * m; - float gx0 = (idx & 0x03) ? gl0 : (idx & 0x20) ? -dw.x : dw.x; - float gx1 = (idx & 0x03) ? gl1 : (idx & 0x20) ? -dw.y : dw.y; - float gy0 = (idx & 0x0c) ? gl0 : -dw.z; - float gy1 = (idx & 0x0c) ? gl1 : -dw.w; - float gz0 = (idx & 0x30) ? gl0 : (idx & 0x03) ? dw.x : dw.z; - float gz1 = (idx & 0x30) ? gl1 : (idx & 0x03) ? dw.y : dw.w; - if (idx & 0x09) - { - gz0 = -gz0; - gz1 = -gz1; - } - g0 = make_float3(gx0, gy0, gz0) * (m * .5f); - g1 = make_float3(gx1, gy1, gz1) * (m * .5f); - if (!isfinite_vec3(g0) || !isfinite_vec3(g1)) - { - g0 = make_float3(0.f, 0.f, 0.f); // Invalid uv. - g1 = make_float3(0.f, 0.f, 0.f); - } -} - -// Compute d{s,t}/d{X,Y} based on d{x,y,z}/d{X,Y} at a given 3D lookup vector. -// Result is (ds/dX, ds/dY, dt/dX, dt/dY). -static __device__ __forceinline__ float4 indexCubeMapGradST(float3 uv, float3 dvdX, float3 dvdY) -{ - float ax = fabsf(uv.x); - float ay = fabsf(uv.y); - float az = fabsf(uv.z); - int idx; - float c, gu, gv; - if (az > fmaxf(ax, ay)) { idx = 0x10; c = uv.z; gu = uv.x; gv = uv.y; } - else if (ay > ax) { idx = 0x04; c = uv.y; gu = uv.x; gv = uv.z; } - else { idx = 0x01; c = uv.x; gu = uv.z; gv = uv.y; } - if (c < 0.f) idx += idx; - if (idx & 0x09) - { - dvdX.z = -dvdX.z; - dvdY.z = -dvdY.z; - } - float m = __frcp_rz(fabsf(c)); - float dm = m * .5f; - float mm = m * dm; - gu *= (idx & 0x34) ? -mm : mm; - gv *= (idx & 0x2e) ? -mm : mm; - - float4 res; - if (idx & 0x03) - { - res = make_float4(gu * dvdX.x + dm * dvdX.z, - gu * dvdY.x + dm * dvdY.z, - gv * dvdX.x - dm * dvdX.y, - gv * dvdY.x - dm * dvdY.y); - } - else if (idx & 0x0c) - { - res = make_float4(gu * dvdX.y + dm * dvdX.x, - gu * dvdY.y + dm * dvdY.x, - gv * dvdX.y + dm * dvdX.z, - gv * dvdY.y + dm * dvdY.z); - } - else // (idx & 0x30) - { - res = make_float4(gu * dvdX.z + copysignf(dm, c) * dvdX.x, - gu * dvdY.z + copysignf(dm, c) * dvdY.x, - gv * dvdX.z - dm * dvdX.y, - gv * dvdY.z - dm * dvdY.y); - } - - if (!isfinite_vec4(res)) - return make_float4(0.f, 0.f, 0.f, 0.f); - - return res; -} - -// Compute d(d{s,t}/d{X,Y})/d{x,y,z}, i.e., how the pixel derivatives of 2D face -// coordinates change w.r.t. 3D texture coordinate vector, returned as follows: -// | d(ds/dX)/dx d(ds/dY)/dx d(dt/dX)/dx d(dt/dY)/dx | -// | d(ds/dX)/dy d(ds/dY)/dy d(dt/dX)/dy d(dt/dY)/dy | -// | d(ds/dX)/dz d(ds/dY)/dz d(dt/dX)/dz d(dt/dY)/dz | -static __device__ __forceinline__ void indexCubeMapGrad2(float3 uv, float3 dvdX, float3 dvdY, float4& dx, float4& dy, float4& dz) -{ - float ax = fabsf(uv.x); - float ay = fabsf(uv.y); - float az = fabsf(uv.z); - int idx; - float c, gu, gv; - if (az > fmaxf(ax, ay)) { idx = 0x10; c = uv.z; gu = uv.x; gv = uv.y; } - else if (ay > ax) { idx = 0x04; c = uv.y; gu = uv.x; gv = uv.z; } - else { idx = 0x01; c = uv.x; gu = uv.z; gv = uv.y; } - if (c < 0.f) idx += idx; - - if (idx & 0x09) - { - dvdX.z = -dvdX.z; - dvdY.z = -dvdY.z; - } - - float m = __frcp_rz(c); - float dm = -m * fabsf(m) * .5; - float mm = m * m * .5; - float mu = (idx & 0x34) ? -mm : mm; - float mv = (idx & 0x2e) ? -mm : mm; - gu *= -2.0 * m * mu; - gv *= -2.0 * m * mv; - - if (idx & 0x03) - { - dx.x = gu * dvdX.x + dm * dvdX.z; - dx.y = gu * dvdY.x + dm * dvdY.z; - dx.z = gv * dvdX.x - dm * dvdX.y; - dx.w = gv * dvdY.x - dm * dvdY.y; - dy.x = 0.f; - dy.y = 0.f; - dy.z = mv * dvdX.x; - dy.w = mv * dvdY.x; - dz.x = mu * dvdX.x; - dz.y = mu * dvdY.x; - dz.z = 0.f; - dz.w = 0.f; - } - else if (idx & 0x0c) - { - dx.x = mu * dvdX.y; - dx.y = mu * dvdY.y; - dx.z = 0.f; - dx.w = 0.f; - dy.x = gu * dvdX.y + dm * dvdX.x; - dy.y = gu * dvdY.y + dm * dvdY.x; - dy.z = gv * dvdX.y + dm * dvdX.z; - dy.w = gv * dvdY.y + dm * dvdY.z; - dz.x = 0.f; - dz.y = 0.f; - dz.z = mv * dvdX.y; - dz.w = mv * dvdY.y; - } - else // (idx & 0x30) - { - dx.x = mu * dvdX.z; - dx.y = mu * dvdY.z; - dx.z = 0.f; - dx.w = 0.f; - dy.x = 0.f; - dy.y = 0.f; - dy.z = mv * dvdX.z; - dy.w = mv * dvdY.z; - dz.x = gu * dvdX.z - fabsf(dm) * dvdX.x; - dz.y = gu * dvdY.z - fabsf(dm) * dvdY.x; - dz.z = gv * dvdX.z - dm * dvdX.y; - dz.w = gv * dvdY.z - dm * dvdY.y; - } -} - -//------------------------------------------------------------------------ -// General texture indexing. - -template -static __device__ __forceinline__ int indexTextureNearest(const TextureKernelParams& p, float3 uv, int tz) -{ - int w = p.texWidth; - int h = p.texHeight; - float u = uv.x; - float v = uv.y; - - // Cube map indexing. - if (CUBE_MODE) - { - // No wrap. Fold face index into tz right away. - int idx = indexCubeMap(u, v, uv.z); // Rewrites u, v. - if (idx < 0) - return -1; // Invalid uv. - tz = 6 * tz + idx; - } - else - { - // Handle boundary. - if (p.boundaryMode == TEX_BOUNDARY_MODE_WRAP) - { - u = u - (float)__float2int_rd(u); - v = v - (float)__float2int_rd(v); - } - } - - u = u * (float)w; - v = v * (float)h; - - int iu = __float2int_rd(u); - int iv = __float2int_rd(v); - - // In zero boundary mode, return texture address -1. - if (!CUBE_MODE && p.boundaryMode == TEX_BOUNDARY_MODE_ZERO) - { - if (iu < 0 || iu >= w || iv < 0 || iv >= h) - return -1; - } - - // Otherwise clamp and calculate the coordinate properly. - iu = min(max(iu, 0), w-1); - iv = min(max(iv, 0), h-1); - return iu + w * (iv + tz * h); -} - -template -static __device__ __forceinline__ float2 indexTextureLinear(const TextureKernelParams& p, float3 uv, int tz, int4& tcOut, int level) -{ - // Mip level size. - int2 sz = mipLevelSize(p, level); - int w = sz.x; - int h = sz.y; - - // Compute texture-space u, v. - float u = uv.x; - float v = uv.y; - bool clampU = false; - bool clampV = false; - - // Cube map indexing. - int face = 0; - if (CUBE_MODE) - { - // Neither clamp or wrap. - face = indexCubeMap(u, v, uv.z); // Rewrites u, v. - if (face < 0) - { - tcOut.x = tcOut.y = tcOut.z = tcOut.w = -1; // Invalid uv. - return make_float2(0.f, 0.f); - } - u = u * (float)w - 0.5f; - v = v * (float)h - 0.5f; - } - else - { - if (p.boundaryMode == TEX_BOUNDARY_MODE_WRAP) - { - // Wrap. - u = u - (float)__float2int_rd(u); - v = v - (float)__float2int_rd(v); - } - - // Move to texel space. - u = u * (float)w - 0.5f; - v = v * (float)h - 0.5f; - - if (p.boundaryMode == TEX_BOUNDARY_MODE_CLAMP) - { - // Clamp to center of edge texels. - u = fminf(fmaxf(u, 0.f), w - 1.f); - v = fminf(fmaxf(v, 0.f), h - 1.f); - clampU = (u == 0.f || u == w - 1.f); - clampV = (v == 0.f || v == h - 1.f); - } - } - - // Compute texel coordinates and weights. - int iu0 = __float2int_rd(u); - int iv0 = __float2int_rd(v); - int iu1 = iu0 + (clampU ? 0 : 1); // Ensure zero u/v gradients with clamped. - int iv1 = iv0 + (clampV ? 0 : 1); - u -= (float)iu0; - v -= (float)iv0; - - // Cube map wrapping. - bool cubeWrap = CUBE_MODE && (iu0 < 0 || iv0 < 0 || iu1 >= w || iv1 >= h); - if (cubeWrap) - { - tcOut = wrapCubeMap(face, iu0, iu1, iv0, iv1, w); - tcOut += 6 * tz * w * h; // Bring in tz. - return make_float2(u, v); // Done. - } - - // Fold cube map face into tz. - if (CUBE_MODE) - tz = 6 * tz + face; - - // Wrap overflowing texel indices. - if (!CUBE_MODE && p.boundaryMode == TEX_BOUNDARY_MODE_WRAP) - { - if (iu0 < 0) iu0 += w; - if (iv0 < 0) iv0 += h; - if (iu1 >= w) iu1 -= w; - if (iv1 >= h) iv1 -= h; - } - - // Coordinates with tz folded in. - int iu0z = iu0 + tz * w * h; - int iu1z = iu1 + tz * w * h; - tcOut.x = iu0z + w * iv0; - tcOut.y = iu1z + w * iv0; - tcOut.z = iu0z + w * iv1; - tcOut.w = iu1z + w * iv1; - - // Invalidate texture addresses outside unit square if we are in zero mode. - if (!CUBE_MODE && p.boundaryMode == TEX_BOUNDARY_MODE_ZERO) - { - bool iu0_out = (iu0 < 0 || iu0 >= w); - bool iu1_out = (iu1 < 0 || iu1 >= w); - bool iv0_out = (iv0 < 0 || iv0 >= h); - bool iv1_out = (iv1 < 0 || iv1 >= h); - if (iu0_out || iv0_out) tcOut.x = -1; - if (iu1_out || iv0_out) tcOut.y = -1; - if (iu0_out || iv1_out) tcOut.z = -1; - if (iu1_out || iv1_out) tcOut.w = -1; - } - - // All done. - return make_float2(u, v); -} - -//------------------------------------------------------------------------ -// Mip level calculation. - -template -static __device__ __forceinline__ void calculateMipLevel(int& level0, int& level1, float& flevel, const TextureKernelParams& p, int pidx, float3 uv, float4* pdw, float3* pdfdv) -{ - // Do nothing if mips not in use. - if (FILTER_MODE == TEX_MODE_NEAREST || FILTER_MODE == TEX_MODE_LINEAR) - return; - - // Determine mip level based on UV pixel derivatives. If no derivatives are given (mip level bias only), leave as zero. - if (!BIAS_ONLY) - { - // Get pixel derivatives of texture coordinates. - float4 uvDA; - float3 dvdX, dvdY; // Gradients use these later. - if (CUBE_MODE) - { - // Fetch. - float2 d0 = ((const float2*)p.uvDA)[3 * pidx + 0]; - float2 d1 = ((const float2*)p.uvDA)[3 * pidx + 1]; - float2 d2 = ((const float2*)p.uvDA)[3 * pidx + 2]; - - // Map d{x,y,z}/d{X,Y} into d{s,t}/d{X,Y}. - dvdX = make_float3(d0.x, d1.x, d2.x); // d{x,y,z}/dX - dvdY = make_float3(d0.y, d1.y, d2.y); // d{x,y,z}/dY - uvDA = indexCubeMapGradST(uv, dvdX, dvdY); // d{s,t}/d{X,Y} - } - else - { - // Fetch. - uvDA = ((const float4*)p.uvDA)[pidx]; - } - - // Scaling factors. - float uscl = p.texWidth; - float vscl = p.texHeight; - - // d[s,t]/d[X,Y]. - float dsdx = uvDA.x * uscl; - float dsdy = uvDA.y * uscl; - float dtdx = uvDA.z * vscl; - float dtdy = uvDA.w * vscl; - - // Calculate footprint axis lengths. - float A = dsdx*dsdx + dtdx*dtdx; - float B = dsdy*dsdy + dtdy*dtdy; - float C = dsdx*dsdy + dtdx*dtdy; - float l2b = 0.5 * (A + B); - float l2n = 0.25 * (A-B)*(A-B) + C*C; - float l2a = sqrt(l2n); - float lenMinorSqr = fmaxf(0.0, l2b - l2a); - float lenMajorSqr = l2b + l2a; - - // Footprint vs. mip level gradient. - if (pdw && FILTER_MODE == TEX_MODE_LINEAR_MIPMAP_LINEAR) - { - float dw = 0.72134752f / (l2n + l2a * l2b); // Constant is 0.5/ln(2). - float AB = dw * .5f * (A - B); - float Cw = dw * C; - float l2aw = dw * l2a; - float d_f_ddsdX = uscl * (dsdx * (l2aw + AB) + dsdy * Cw); - float d_f_ddsdY = uscl * (dsdy * (l2aw - AB) + dsdx * Cw); - float d_f_ddtdX = vscl * (dtdx * (l2aw + AB) + dtdy * Cw); - float d_f_ddtdY = vscl * (dtdy * (l2aw - AB) + dtdx * Cw); - - float4 d_f_dw = make_float4(d_f_ddsdX, d_f_ddsdY, d_f_ddtdX, d_f_ddtdY); - if (!CUBE_MODE) - *pdw = isfinite_vec4(d_f_dw) ? d_f_dw : make_float4(0.f, 0.f, 0.f, 0.f); - - // In cube maps, there is also a texture coordinate vs. mip level gradient. - // Only output nonzero vectors if both are free of inf/Nan garbage. - if (CUBE_MODE) - { - float4 dx, dy, dz; - indexCubeMapGrad2(uv, dvdX, dvdY, dx, dy, dz); - float3 d_dsdX_dv = make_float3(dx.x, dy.x, dz.x); - float3 d_dsdY_dv = make_float3(dx.y, dy.y, dz.y); - float3 d_dtdX_dv = make_float3(dx.z, dy.z, dz.z); - float3 d_dtdY_dv = make_float3(dx.w, dy.w, dz.w); - - float3 d_f_dv = make_float3(0.f, 0.f, 0.f); - d_f_dv += d_dsdX_dv * d_f_ddsdX; - d_f_dv += d_dsdY_dv * d_f_ddsdY; - d_f_dv += d_dtdX_dv * d_f_ddtdX; - d_f_dv += d_dtdY_dv * d_f_ddtdY; - - bool finite = isfinite_vec4(d_f_dw) && isfinite_vec3(d_f_dv); - *pdw = finite ? d_f_dw : make_float4(0.f, 0.f, 0.f, 0.f); - *pdfdv = finite ? d_f_dv : make_float3(0.f, 0.f, 0.f); - } - } - - // Finally, calculate mip level. - flevel = .5f * __log2f(lenMajorSqr); // May be inf/NaN, but clamp fixes it. - } - - // Bias the mip level and clamp. - if (p.mipLevelBias) - flevel += p.mipLevelBias[pidx]; - flevel = fminf(fmaxf(flevel, 0.f), (float)p.mipLevelMax); - - // Calculate levels depending on filter mode. - level0 = __float2int_rd(flevel); - - // Leave everything else at zero if flevel == 0 (magnification) or when in linear-mipmap-nearest mode. - if (FILTER_MODE == TEX_MODE_LINEAR_MIPMAP_LINEAR && flevel > 0.f) - { - level1 = min(level0 + 1, p.mipLevelMax); - flevel -= level0; // Fractional part. Zero if clamped on last level. - } -} - -//------------------------------------------------------------------------ -// Texel fetch and accumulator helpers that understand cube map corners. - -template -static __device__ __forceinline__ void fetchQuad(T& a00, T& a10, T& a01, T& a11, const float* pIn, int4 tc, bool corner) -{ - // For invalid cube map uv, tc will be all negative, and all texel values will be zero. - if (corner) - { - T avg = zero_value(); - if (tc.x >= 0) avg += (a00 = *((const T*)&pIn[tc.x])); - if (tc.y >= 0) avg += (a10 = *((const T*)&pIn[tc.y])); - if (tc.z >= 0) avg += (a01 = *((const T*)&pIn[tc.z])); - if (tc.w >= 0) avg += (a11 = *((const T*)&pIn[tc.w])); - avg *= 0.33333333f; - if (tc.x < 0) a00 = avg; - if (tc.y < 0) a10 = avg; - if (tc.z < 0) a01 = avg; - if (tc.w < 0) a11 = avg; - } - else - { - a00 = (tc.x >= 0) ? *((const T*)&pIn[tc.x]) : zero_value(); - a10 = (tc.y >= 0) ? *((const T*)&pIn[tc.y]) : zero_value(); - a01 = (tc.z >= 0) ? *((const T*)&pIn[tc.z]) : zero_value(); - a11 = (tc.w >= 0) ? *((const T*)&pIn[tc.w]) : zero_value(); - } -} - -static __device__ __forceinline__ void accumQuad(float4 c, float* pOut, int level, int4 tc, bool corner, CA_TEMP_PARAM) -{ - // For invalid cube map uv, tc will be all negative, and no accumulation will take place. - if (corner) - { - float cb; - if (tc.x < 0) cb = c.x; - if (tc.y < 0) cb = c.y; - if (tc.z < 0) cb = c.z; - if (tc.w < 0) cb = c.w; - cb *= 0.33333333f; - if (tc.x >= 0) caAtomicAddTexture(pOut, level, tc.x, c.x + cb); - if (tc.y >= 0) caAtomicAddTexture(pOut, level, tc.y, c.y + cb); - if (tc.z >= 0) caAtomicAddTexture(pOut, level, tc.z, c.z + cb); - if (tc.w >= 0) caAtomicAddTexture(pOut, level, tc.w, c.w + cb); - } - else - { - if (tc.x >= 0) caAtomicAddTexture(pOut, level, tc.x, c.x); - if (tc.y >= 0) caAtomicAddTexture(pOut, level, tc.y, c.y); - if (tc.z >= 0) caAtomicAddTexture(pOut, level, tc.z, c.z); - if (tc.w >= 0) caAtomicAddTexture(pOut, level, tc.w, c.w); - } -} - -//------------------------------------------------------------------------ -// Mip builder kernel. - -template -static __forceinline__ __device__ void MipBuildKernelTemplate(const TextureKernelParams p) -{ - // Sizes. - int2 sz_in = mipLevelSize(p, p.mipLevelOut - 1); - int2 sz_out = mipLevelSize(p, p.mipLevelOut); - - // 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 >= sz_out.x || py >= sz_out.y) - return; - - // Pixel indices. - int pidx_in0 = p.channels * (((px + sz_in.x * py) << 1) + (pz * sz_in.x * sz_in.y)); - int pidx_in1 = pidx_in0 + p.channels * sz_in.x; // Next pixel down. - int pidx_out = p.channels * (px + sz_out.x * (py + sz_out.y * pz)); - - // Input and output pointers. - const float* pin = p.tex[p.mipLevelOut - 1]; - float* pout = (float*)p.tex[p.mipLevelOut]; - - // Special case: Input texture height or width is 1. - if (sz_in.x == 1 || sz_in.y == 1) - { - if (sz_in.y == 1) - pidx_in1 = pidx_in0 + p.channels; // Next pixel on the right. - - for (int i=0; i < p.channels; i += C) - { - T v0 = *((const T*)&pin[pidx_in0 + i]); - T v1 = *((const T*)&pin[pidx_in1 + i]); - T avg = .5f * (v0 + v1); -#if TEX_DEBUG_MIP_RETAIN_VARIANCE - avg = (avg - .5f) * 1.41421356f + .5f; -#endif - *((T*)&pout[pidx_out + i]) = avg; - } - - return; - } - - for (int i=0; i < p.channels; i += C) - { - T v0 = *((const T*)&pin[pidx_in0 + i]); - T v1 = *((const T*)&pin[pidx_in0 + i + p.channels]); - T v2 = *((const T*)&pin[pidx_in1 + i]); - T v3 = *((const T*)&pin[pidx_in1 + i + p.channels]); - T avg = .25f * (v0 + v1 + v2 + v3); -#if TEX_DEBUG_MIP_RETAIN_VARIANCE - avg = (avg - .5f) * 2.f + .5f; -#endif - *((T*)&pout[pidx_out + i]) = avg; - } -} - -// Template specializations. -__global__ void MipBuildKernel1(const TextureKernelParams p) { MipBuildKernelTemplate(p); } -__global__ void MipBuildKernel2(const TextureKernelParams p) { MipBuildKernelTemplate(p); } -__global__ void MipBuildKernel4(const TextureKernelParams p) { MipBuildKernelTemplate(p); } - -//------------------------------------------------------------------------ -// Forward kernel. - -template -static __forceinline__ __device__ void TextureFwdKernelTemplate(const TextureKernelParams p) -{ - // Calculate pixel position. - int px = blockIdx.x * blockDim.x + threadIdx.x; - int py = blockIdx.y * blockDim.y + threadIdx.y; - int pz = blockIdx.z; - int tz = (p.texDepth == 1) ? 0 : pz; - if (px >= p.imgWidth || py >= p.imgHeight || pz >= p.n) - return; - - // Pixel index. - int pidx = px + p.imgWidth * (py + p.imgHeight * pz); - - // Output ptr. - float* pOut = p.out + pidx * p.channels; - - // Get UV. - float3 uv; - if (CUBE_MODE) - uv = ((const float3*)p.uv)[pidx]; - else - uv = make_float3(((const float2*)p.uv)[pidx], 0.f); - - // Nearest mode. - if (FILTER_MODE == TEX_MODE_NEAREST) - { - int tc = indexTextureNearest(p, uv, tz); - tc *= p.channels; - const float* pIn = p.tex[0]; - - // Copy if valid tc, otherwise output zero. - for (int i=0; i < p.channels; i += C) - *((T*)&pOut[i]) = (tc >= 0) ? *((const T*)&pIn[tc + i]) : zero_value(); - - return; // Exit. - } - - // Calculate mip level. In 'linear' mode these will all stay zero. - float flevel = 0.f; // Fractional level. - int level0 = 0; // Discrete level 0. - int level1 = 0; // Discrete level 1. - calculateMipLevel(level0, level1, flevel, p, pidx, uv, 0, 0); - - // Get texel indices and pointer for level 0. - int4 tc0 = make_int4(0, 0, 0, 0); - float2 uv0 = indexTextureLinear(p, uv, tz, tc0, level0); - const float* pIn0 = p.tex[level0]; - bool corner0 = CUBE_MODE && ((tc0.x | tc0.y | tc0.z | tc0.w) < 0); - tc0 *= p.channels; - - // Bilinear fetch. - if (FILTER_MODE == TEX_MODE_LINEAR || FILTER_MODE == TEX_MODE_LINEAR_MIPMAP_NEAREST) - { - // Interpolate. - for (int i=0; i < p.channels; i += C, tc0 += C) - { - T a00, a10, a01, a11; - fetchQuad(a00, a10, a01, a11, pIn0, tc0, corner0); - *((T*)&pOut[i]) = bilerp(a00, a10, a01, a11, uv0); - } - return; // Exit. - } - - // Get texel indices and pointer for level 1. - int4 tc1 = make_int4(0, 0, 0, 0); - float2 uv1 = indexTextureLinear(p, uv, tz, tc1, level1); - const float* pIn1 = p.tex[level1]; - bool corner1 = CUBE_MODE && ((tc1.x | tc1.y | tc1.z | tc1.w) < 0); - tc1 *= p.channels; - - // Trilinear fetch. - for (int i=0; i < p.channels; i += C, tc0 += C, tc1 += C) - { - // First level. - T a00, a10, a01, a11; - fetchQuad(a00, a10, a01, a11, pIn0, tc0, corner0); - T a = bilerp(a00, a10, a01, a11, uv0); - - // Second level unless in magnification mode. - if (flevel > 0.f) - { - T b00, b10, b01, b11; - fetchQuad(b00, b10, b01, b11, pIn1, tc1, corner1); - T b = bilerp(b00, b10, b01, b11, uv1); - a = lerp(a, b, flevel); // Interpolate between levels. - } - - // Write. - *((T*)&pOut[i]) = a; - } -} - -// Template specializations. -__global__ void TextureFwdKernelNearest1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelNearest2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelNearest4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinear1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinear2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinear4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapNearest1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapNearest2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapNearest4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapLinear1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapLinear2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapLinear4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeNearest1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeNearest2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeNearest4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinear1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinear2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinear4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapNearest1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapNearest2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapNearest4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapLinear1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapLinear2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapLinear4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapNearestBO1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapNearestBO2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapNearestBO4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapLinearBO1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapLinearBO2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelLinearMipmapLinearBO4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapNearestBO1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapNearestBO2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapNearestBO4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapLinearBO1 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapLinearBO2 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } -__global__ void TextureFwdKernelCubeLinearMipmapLinearBO4 (const TextureKernelParams p) { TextureFwdKernelTemplate(p); } - -//------------------------------------------------------------------------ -// Gradient mip puller kernel. - -template -static __forceinline__ __device__ void MipGradKernelTemplate(const TextureKernelParams 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.texWidth || py >= p.texHeight) - return; - - // Number of wide elements. - int c = p.channels; - if (C == 2) c >>= 1; - if (C == 4) c >>= 2; - - // Dynamically allocated shared memory for holding a texel. - extern __shared__ float s_texelAccum[]; - int sharedOfs = threadIdx.x + threadIdx.y * blockDim.x; - int sharedStride = blockDim.x * blockDim.y; -# define TEXEL_ACCUM(_i) (s_texelAccum + (sharedOfs + (_i) * sharedStride)) - - // Clear the texel. - for (int i=0; i < p.channels; i++) - *TEXEL_ACCUM(i) = 0.f; - - // Track texel position and accumulation weight over the mip stack. - int x = px; - int y = py; - float w = 1.f; - - // Pull gradients from all levels. - int2 sz = mipLevelSize(p, 0); // Previous level size. - for (int level=1; level <= p.mipLevelMax; level++) - { - // Weight decay depends on previous level size. - if (sz.x > 1) w *= .5f; - if (sz.y > 1) w *= .5f; - - // Current level size and coordinates. - sz = mipLevelSize(p, level); - x >>= 1; - y >>= 1; - - T* pIn = (T*)(p.gradTex[level] + (x + sz.x * (y + sz.y * pz)) * p.channels); - for (int i=0; i < c; i++) - accum_from_mem(TEXEL_ACCUM(i * C), sharedStride, pIn[i], w); - } - - // Add to main texture gradients. - T* pOut = (T*)(p.gradTex[0] + (px + p.texWidth * (py + p.texHeight * pz)) * p.channels); - for (int i=0; i < c; i++) - accum_to_mem(pOut[i], TEXEL_ACCUM(i * C), sharedStride); -} - -// Template specializations. -__global__ void MipGradKernel1(const TextureKernelParams p) { MipGradKernelTemplate(p); } -__global__ void MipGradKernel2(const TextureKernelParams p) { MipGradKernelTemplate(p); } -__global__ void MipGradKernel4(const TextureKernelParams p) { MipGradKernelTemplate(p); } - -//------------------------------------------------------------------------ -// Gradient kernel. - -template -static __forceinline__ __device__ void TextureGradKernelTemplate(const TextureKernelParams p) -{ - // Temporary space for coalesced atomics. - CA_DECLARE_TEMP(TEX_GRAD_MAX_KERNEL_BLOCK_WIDTH * TEX_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; - int tz = (p.texDepth == 1) ? 0 : pz; - if (px >= p.imgWidth || py >= p.imgHeight || pz >= p.n) - return; - - // Pixel index. - int pidx = px + p.imgWidth * (py + p.imgHeight * pz); - - // Early exit if output gradients are zero. - const float* pDy = p.dy + pidx * p.channels; - unsigned int dmax = 0u; - if ((p.channels & 3) == 0) - { - for (int i=0; i < p.channels; i += 4) - { - uint4 dy = *((const uint4*)&pDy[i]); - dmax |= (dy.x | dy.y | dy.z | dy.w); - } - } - else - { - for (int i=0; i < p.channels; i++) - dmax |= __float_as_uint(pDy[i]); - } - - // Store zeros and exit. - if (__uint_as_float(dmax) == 0.f) - { - if (CUBE_MODE) - { - if (FILTER_MODE != TEX_MODE_NEAREST) - ((float3*)p.gradUV)[pidx] = make_float3(0.f, 0.f, 0.f); - if (FILTER_MODE == TEX_MODE_LINEAR_MIPMAP_LINEAR) - { - if (p.gradUVDA) - { - ((float2*)p.gradUVDA)[3 * pidx + 0] = make_float2(0.f, 0.f); - ((float2*)p.gradUVDA)[3 * pidx + 1] = make_float2(0.f, 0.f); - ((float2*)p.gradUVDA)[3 * pidx + 2] = make_float2(0.f, 0.f); - } - if (p.gradMipLevelBias) - p.gradMipLevelBias[pidx] = 0.f; - } - } - else - { - if (FILTER_MODE != TEX_MODE_NEAREST) - ((float2*)p.gradUV)[pidx] = make_float2(0.f, 0.f); - if (FILTER_MODE == TEX_MODE_LINEAR_MIPMAP_LINEAR) - { - if (p.gradUVDA) - ((float4*)p.gradUVDA)[pidx] = make_float4(0.f, 0.f, 0.f, 0.f); - if (p.gradMipLevelBias) - p.gradMipLevelBias[pidx] = 0.f; - } - } - return; - } - - // Get UV. - float3 uv; - if (CUBE_MODE) - uv = ((const float3*)p.uv)[pidx]; - else - uv = make_float3(((const float2*)p.uv)[pidx], 0.f); - - // Nearest mode - texture gradients only. - if (FILTER_MODE == TEX_MODE_NEAREST) - { - int tc = indexTextureNearest(p, uv, tz); - if (tc < 0) - return; // Outside texture. - - tc *= p.channels; - float* pOut = p.gradTex[0]; - - // Accumulate texture gradients. - for (int i=0; i < p.channels; i++) - caAtomicAddTexture(pOut, 0, tc + i, pDy[i]); - - return; // Exit. - } - - // Calculate mip level. In 'linear' mode these will all stay zero. - float4 dw = make_float4(0.f, 0.f, 0.f, 0.f); - float3 dfdv = make_float3(0.f, 0.f, 0.f); - float flevel = 0.f; // Fractional level. - int level0 = 0; // Discrete level 0. - int level1 = 0; // Discrete level 1. - calculateMipLevel(level0, level1, flevel, p, pidx, uv, &dw, &dfdv); - - // UV gradient accumulators. - float gu = 0.f; - float gv = 0.f; - - // Get texel indices and pointers for level 0. - int4 tc0 = make_int4(0, 0, 0, 0); - float2 uv0 = indexTextureLinear(p, uv, tz, tc0, level0); - const float* pIn0 = p.tex[level0]; - float* pOut0 = p.gradTex[level0]; - bool corner0 = CUBE_MODE && ((tc0.x | tc0.y | tc0.z | tc0.w) < 0); - tc0 *= p.channels; - - // Texel weights. - float uv011 = uv0.x * uv0.y; - float uv010 = uv0.x - uv011; - float uv001 = uv0.y - uv011; - float uv000 = 1.f - uv0.x - uv001; - float4 tw0 = make_float4(uv000, uv010, uv001, uv011); - - // Attribute weights. - int2 sz0 = mipLevelSize(p, level0); - float sclu0 = (float)sz0.x; - float sclv0 = (float)sz0.y; - - // Bilinear mode - texture and uv gradients. - if (FILTER_MODE == TEX_MODE_LINEAR || FILTER_MODE == TEX_MODE_LINEAR_MIPMAP_NEAREST) - { - for (int i=0; i < p.channels; i++, tc0 += 1) - { - float dy = pDy[i]; - accumQuad(tw0 * dy, pOut0, level0, tc0, corner0, CA_TEMP); - - float a00, a10, a01, a11; - fetchQuad(a00, a10, a01, a11, pIn0, tc0, corner0); - float ad = (a11 + a00 - a10 - a01); - gu += dy * ((a10 - a00) + uv0.y * ad) * sclu0; - gv += dy * ((a01 - a00) + uv0.x * ad) * sclv0; - } - - // Store UV gradients and exit. - if (CUBE_MODE) - ((float3*)p.gradUV)[pidx] = indexCubeMapGrad(uv, gu, gv); - else - ((float2*)p.gradUV)[pidx] = make_float2(gu, gv); - - return; - } - - // Accumulate fractional mip level gradient. - float df = 0; // dL/df. - - // Get texel indices and pointers for level 1. - int4 tc1 = make_int4(0, 0, 0, 0); - float2 uv1 = indexTextureLinear(p, uv, tz, tc1, level1); - const float* pIn1 = p.tex[level1]; - float* pOut1 = p.gradTex[level1]; - bool corner1 = CUBE_MODE && ((tc1.x | tc1.y | tc1.z | tc1.w) < 0); - tc1 *= p.channels; - - // Texel weights. - float uv111 = uv1.x * uv1.y; - float uv110 = uv1.x - uv111; - float uv101 = uv1.y - uv111; - float uv100 = 1.f - uv1.x - uv101; - float4 tw1 = make_float4(uv100, uv110, uv101, uv111); - - // Attribute weights. - int2 sz1 = mipLevelSize(p, level1); - float sclu1 = (float)sz1.x; - float sclv1 = (float)sz1.y; - - // Trilinear mode. - for (int i=0; i < p.channels; i++, tc0 += 1, tc1 += 1) - { - float dy = pDy[i]; - float dy0 = (1.f - flevel) * dy; - accumQuad(tw0 * dy0, pOut0, level0, tc0, corner0, CA_TEMP); - - // UV gradients for first level. - float a00, a10, a01, a11; - fetchQuad(a00, a10, a01, a11, pIn0, tc0, corner0); - float ad = (a11 + a00 - a10 - a01); - gu += dy0 * ((a10 - a00) + uv0.y * ad) * sclu0; - gv += dy0 * ((a01 - a00) + uv0.x * ad) * sclv0; - - // Second level unless in magnification mode. - if (flevel > 0.f) - { - // Texture gradients for second level. - float dy1 = flevel * dy; - accumQuad(tw1 * dy1, pOut1, level1, tc1, corner1, CA_TEMP); - - // UV gradients for second level. - float b00, b10, b01, b11; - fetchQuad(b00, b10, b01, b11, pIn1, tc1, corner1); - float bd = (b11 + b00 - b10 - b01); - gu += dy1 * ((b10 - b00) + uv1.y * bd) * sclu1; - gv += dy1 * ((b01 - b00) + uv1.x * bd) * sclv1; - - // Mip level gradient. - float a = bilerp(a00, a10, a01, a11, uv0); - float b = bilerp(b00, b10, b01, b11, uv1); - df += (b-a) * dy; - } - } - - // Store UV gradients. - if (CUBE_MODE) - ((float3*)p.gradUV)[pidx] = indexCubeMapGrad(uv, gu, gv) + (dfdv * df); - else - ((float2*)p.gradUV)[pidx] = make_float2(gu, gv); - - // Store mip level bias gradient. - if (p.gradMipLevelBias) - p.gradMipLevelBias[pidx] = df; - - // Store UV pixel differential gradients. - if (!BIAS_ONLY) - { - // Final gradients. - dw *= df; // dL/(d{s,y}/d{X,Y}) = df/(d{s,y}/d{X,Y}) * dL/df. - - // Store them. - if (CUBE_MODE) - { - // Remap from dL/(d{s,t}/s{X,Y}) to dL/(d{x,y,z}/d{X,Y}). - float3 g0, g1; - indexCubeMapGrad4(uv, dw, g0, g1); - ((float2*)p.gradUVDA)[3 * pidx + 0] = make_float2(g0.x, g1.x); - ((float2*)p.gradUVDA)[3 * pidx + 1] = make_float2(g0.y, g1.y); - ((float2*)p.gradUVDA)[3 * pidx + 2] = make_float2(g0.z, g1.z); - } - else - ((float4*)p.gradUVDA)[pidx] = dw; - } -} - -// Template specializations. -__global__ void TextureGradKernelNearest (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelLinear (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelLinearMipmapNearest (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelLinearMipmapLinear (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelCubeNearest (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelCubeLinear (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelCubeLinearMipmapNearest (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelCubeLinearMipmapLinear (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelLinearMipmapNearestBO (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelLinearMipmapLinearBO (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelCubeLinearMipmapNearestBO (const TextureKernelParams p) { TextureGradKernelTemplate(p); } -__global__ void TextureGradKernelCubeLinearMipmapLinearBO (const TextureKernelParams p) { TextureGradKernelTemplate(p); } - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/lib/setgpu.lib b/extensions/nvdiffrast/nvdiffrast/lib/setgpu.lib deleted file mode 100644 index add9a0c4f631cb56dbee31a05ed97339930301e2..0000000000000000000000000000000000000000 Binary files a/extensions/nvdiffrast/nvdiffrast/lib/setgpu.lib and /dev/null differ diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/__init__.py b/extensions/nvdiffrast/nvdiffrast/tensorflow/__init__.py deleted file mode 100644 index c2ac67a712d7df7ecc40cc2d1dfd67e5f738be9f..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# 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 rasterize, interpolate, texture, antialias -from .plugin_loader import set_cache_dir - -__all__ = ["rasterize", "interpolate", "texture", "antialias", "set_cache_dir"] diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/ops.py b/extensions/nvdiffrast/nvdiffrast/tensorflow/ops.py deleted file mode 100644 index 130e2e195e2dce3e007aaebc55cc03837cf77c02..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/ops.py +++ /dev/null @@ -1,303 +0,0 @@ -# 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 tensorflow as tf -import numpy as np -import os -from . import plugin_loader - -#---------------------------------------------------------------------------- -# Helpers. -#---------------------------------------------------------------------------- - -# OpenGL-related linker options depending on platform. -def _get_gl_opts(): - libs = { - 'posix': ['GL', 'EGL'], - 'nt': ['gdi32', 'opengl32', 'user32', 'setgpu'], - } - return ['-l' + x for x in libs[os.name]] - -# Load the cpp plugin. -def _get_plugin(): - fn = os.path.join(os.path.dirname(__file__), 'tf_all.cu') - return plugin_loader.get_plugin(fn, extra_nvcc_options=_get_gl_opts() + ['-DNVDR_TENSORFLOW']) - -# Convert parameter to a numpy array if possible. -def _get_constant(x, dtype): - try: - return np.asarray(x, dtype=dtype) - except (TypeError, ValueError): - return None - -# Tests for a construction-time constantness instead of tf.constant node because -# the latter can be overridden in Session.run() feed_dict at evaluation time. -def _is_constant(x, dtype): - if isinstance(x, np.ndarray): - return np.can_cast(x.dtype, dtype, 'unsafe') - else: - return _get_constant(x, dtype) is not None - -#---------------------------------------------------------------------------- -# Rasterize. -#---------------------------------------------------------------------------- - -def rasterize(pos, tri, resolution, ranges=None, tri_const=False, output_db=True, grad_db=True): - assert tri_const is True or tri_const is False - assert output_db is True or output_db is False - - # Known constant resolution? - resolution_c = _get_constant(resolution, np.int32) - - # Known constant triangles? - tri_const = tri_const or _is_constant(tri, np.int32) - - # Convert all inputs to tensors / base types. - tri_const = 1 if tri_const else 0 - tri = tf.convert_to_tensor(tri, dtype=tf.int32) - pos = tf.convert_to_tensor(pos, dtype=tf.float32) - resolution = tf.convert_to_tensor(resolution, dtype=tf.int32) - if ranges is None: - ranges = tf.convert_to_tensor(np.zeros(shape=[0, 2], dtype=np.int32)) # Empty tensor. - else: - ranges = tf.convert_to_tensor(ranges, dtype=tf.int32) # Convert input to tensor. - - # Infer as much about the output shape as possible. - out_shape = [None, None, None, 4] - if pos.shape.rank == 3: # Instanced mode. - out_shape[0] = pos.shape[0].value - elif pos.shape.rank == 2: # Range mode. - if ranges.shape.rank not in [None, 0]: - out_shape[0] = ranges.shape[0].value - if resolution_c is not None: - assert resolution_c.shape == (2,) - out_shape[1], out_shape[2] = resolution_c - - # Output pixel differentials. - @tf.custom_gradient - def func_db(pos): - out, out_db = _get_plugin().rasterize_fwd(pos, tri, resolution, ranges, 1, tri_const) - out.set_shape(out_shape) - out_db.set_shape(out_shape) - def grad(dy, ddb): - if grad_db: - return _get_plugin().rasterize_grad_db(pos, tri, out, dy, ddb) - else: - return _get_plugin().rasterize_grad(pos, tri, out, dy) - return (out, out_db), grad - - # Do not output pixel differentials. - @tf.custom_gradient - def func(pos): - out, out_db = _get_plugin().rasterize_fwd(pos, tri, resolution, ranges, 0, tri_const) - out.set_shape(out_shape) - out_db.set_shape(out_shape[:-1] + [0]) # Zero channels in out_db. - def grad(dy, _): - return _get_plugin().rasterize_grad(pos, tri, out, dy) - return (out, out_db), grad - - # Choose stub. - if output_db: - return func_db(pos) - else: - return func(pos) - -#---------------------------------------------------------------------------- -# Interpolate. -#---------------------------------------------------------------------------- - -def interpolate(attr, rast, tri, rast_db=None, diff_attrs=None): - # Sanitize the list of pixel differential attributes. - if diff_attrs is None: - diff_attrs = [] - elif diff_attrs != 'all': - diff_attrs = _get_constant(diff_attrs, np.int32) - assert (diff_attrs is not None) and len(diff_attrs.shape) == 1 - diff_attrs = diff_attrs.tolist() - - # Convert all inputs to tensors. - attr = tf.convert_to_tensor(attr, dtype=tf.float32) - rast = tf.convert_to_tensor(rast, dtype=tf.float32) - tri = tf.convert_to_tensor(tri, dtype=tf.int32) - if diff_attrs: - rast_db = tf.convert_to_tensor(rast_db, dtype=tf.float32) - - # Infer output shape. - out_shape = [None, None, None, None] - if rast.shape.rank is not None: - out_shape = [rast.shape[0].value, rast.shape[1].value, rast.shape[2].value, None] - if attr.shape.rank in [2, 3]: - out_shape[3] = attr.shape[-1].value - - # Output pixel differentials for at least some attributes. - @tf.custom_gradient - def func_da(attr, rast, rast_db): - diff_attrs_all = int(diff_attrs == 'all') - diff_attrs_list = [] if diff_attrs_all else diff_attrs - out, out_da = _get_plugin().interpolate_fwd_da(attr, rast, tri, rast_db, diff_attrs_all, diff_attrs_list) - - # Infer number of channels in out_da. - if not diff_attrs_all: - da_channels = 2 * len(diff_attrs) - if (attr.shape.rank in [2, 3]) and (attr.shape[-1].value is not None): - da_channels = 2 * attr.shape[-1].value - else: - da_channels = None - - # Set output shapes. - out.set_shape(out_shape) - out_da.set_shape([out_shape[0], out_shape[1], out_shape[2], da_channels]) - - def grad(dy, dda): - return _get_plugin().interpolate_grad_da(attr, rast, tri, dy, rast_db, dda, diff_attrs_all, diff_attrs_list) - return (out, out_da), grad - - # No pixel differentials for any attribute. - @tf.custom_gradient - def func(attr, rast): - out, out_da = _get_plugin().interpolate_fwd(attr, rast, tri) - out.set_shape(out_shape) - out_da.set_shape(out_shape[:-1] + [0]) # Zero channels in out_da. - def grad(dy, _): - return _get_plugin().interpolate_grad(attr, rast, tri, dy) - return (out, out_da), grad - - # Choose stub. - if diff_attrs: - return func_da(attr, rast, rast_db) - else: - return func(attr, rast) - -#---------------------------------------------------------------------------- -# Texture. -#---------------------------------------------------------------------------- - -def texture(tex, uv, uv_da=None, filter_mode='auto', boundary_mode='wrap', tex_const=False, max_mip_level=None): - assert tex_const is True or tex_const is False - - # Default filter mode. - if filter_mode == 'auto': - filter_mode = 'linear-mipmap-linear' if (uv_da is not None) else 'linear' - - # Known constant texture? - tex_const = tex_const or _is_constant(tex, np.float32) - - # Sanitize inputs. - tex_const = 1 if tex_const else 0 - if max_mip_level is None: - max_mip_level = -1 - else: - max_mip_level = int(max_mip_level) - assert max_mip_level >= 0 - - # Convert inputs to tensors. - tex = tf.convert_to_tensor(tex, dtype=tf.float32) - uv = tf.convert_to_tensor(uv, dtype=tf.float32) - if 'mipmap' in filter_mode: - uv_da = tf.convert_to_tensor(uv_da, dtype=tf.float32) - - # Infer output shape. - out_shape = [None, None, None, None] - if uv.shape.rank is not None: - assert uv.shape.rank == 4 - out_shape = [uv.shape[0].value, uv.shape[1].value, uv.shape[2].value, None] - if tex.shape.rank is not None: - assert tex.shape.rank == (5 if boundary_mode == 'cube' else 4) - out_shape[-1] = tex.shape[-1].value - - # 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] - - # Linear-mipmap-linear: Mipmaps enabled, all gradients active. - @tf.custom_gradient - def func_linear_mipmap_linear(tex, uv, uv_da): - out, mip = _get_plugin().texture_fwd_mip(tex, uv, uv_da, filter_mode_enum, boundary_mode_enum, tex_const, max_mip_level) - out.set_shape(out_shape) - def grad(dy): - return _get_plugin().texture_grad_linear_mipmap_linear(tex, uv, dy, uv_da, mip, filter_mode_enum, boundary_mode_enum, max_mip_level) - return out, grad - - # Linear-mipmap-nearest: Mipmaps enabled, no gradients to uv_da. - @tf.custom_gradient - def func_linear_mipmap_nearest(tex, uv): - out, mip = _get_plugin().texture_fwd_mip(tex, uv, uv_da, filter_mode_enum, boundary_mode_enum, tex_const, max_mip_level) - out.set_shape(out_shape) - def grad(dy): - return _get_plugin().texture_grad_linear_mipmap_nearest(tex, uv, dy, uv_da, mip, filter_mode_enum, boundary_mode_enum, max_mip_level) - return out, grad - - # Linear: Mipmaps disabled, no uv_da, no gradients to uv_da. - @tf.custom_gradient - def func_linear(tex, uv): - out = _get_plugin().texture_fwd(tex, uv, filter_mode_enum, boundary_mode_enum) - out.set_shape(out_shape) - def grad(dy): - return _get_plugin().texture_grad_linear(tex, uv, dy, filter_mode_enum, boundary_mode_enum) - return out, grad - - # Nearest: Mipmaps disabled, no uv_da, no gradients to uv_da or uv. - @tf.custom_gradient - def func_nearest(tex): - out = _get_plugin().texture_fwd(tex, uv, filter_mode_enum, boundary_mode_enum) - out.set_shape(out_shape) - def grad(dy): - return _get_plugin().texture_grad_nearest(tex, uv, dy, filter_mode_enum, boundary_mode_enum) - return out, grad - - # Choose stub. - if filter_mode == 'linear-mipmap-linear': - return func_linear_mipmap_linear(tex, uv, uv_da) - elif filter_mode == 'linear-mipmap-nearest': - return func_linear_mipmap_nearest(tex, uv) - elif filter_mode == 'linear': - return func_linear(tex, uv) - elif filter_mode == 'nearest': - return func_nearest(tex) - -#---------------------------------------------------------------------------- -# Antialias. -#---------------------------------------------------------------------------- - -def antialias(color, rast, pos, tri, tri_const=False, pos_gradient_boost=1.0): - assert tri_const is True or tri_const is False - - # Known constant triangles? - tri_const = tri_const or _is_constant(tri, np.int32) - - # Convert inputs to tensors. - color = tf.convert_to_tensor(color, dtype=tf.float32) - rast = tf.convert_to_tensor(rast, dtype=tf.float32) - pos = tf.convert_to_tensor(pos, dtype=tf.float32) - tri = tf.convert_to_tensor(tri, dtype=tf.int32) - - # Sanitize inputs. - tri_const = 1 if tri_const else 0 - - @tf.custom_gradient - def func(color, pos): - color_out, work_buffer = _get_plugin().antialias_fwd(color, rast, pos, tri, tri_const) - color_out.set_shape(color.shape) - def grad(dy): - grad_color, grad_pos = _get_plugin().antialias_grad(color, rast, pos, tri, dy, work_buffer) - if pos_gradient_boost != 1.0: - grad_pos = grad_pos * pos_gradient_boost - return grad_color, grad_pos - return color_out, grad - - return func(color, pos) - -#---------------------------------------------------------------------------- diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/plugin_loader.py b/extensions/nvdiffrast/nvdiffrast/tensorflow/plugin_loader.py deleted file mode 100644 index 1347df8a5a276298dbeff81e7b06f62e6765a3ef..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/plugin_loader.py +++ /dev/null @@ -1,219 +0,0 @@ -# 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 glob -import os -import re -import uuid -import hashlib -import tempfile -import shutil -import tensorflow as tf -from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module - -#---------------------------------------------------------------------------- -# Global options. - -_nvdiffrast_cache_dir = None - -def set_cache_dir(path: str) -> None: - '''Set CUDA kernel compilation temp dir. - - If `set_cache_dir` is not called, the cache directory will default to - one of the below: - - - Value of NVDIFFRAST_CACHE_DIR env var, if set - - $HOME/.cache/nvdiffrast if HOME env var is set - - $USERPROFILE/.cache/nvdiffrast if USERPROFILE is set. - - Args: - path: Where to save CUDA kernel build temporaries - ''' - global _nvdiffrast_cache_dir - _nvdiffrast_cache_dir = path - -def make_cache_dir_path(*paths: str) -> str: - if _nvdiffrast_cache_dir is not None: - return os.path.join(_nvdiffrast_cache_dir, *paths) - if 'NVDIFFRAST_CACHE_DIR' in os.environ: - return os.path.join(os.environ['NVDIFFRAST_CACHE_DIR'], *paths) - if 'HOME' in os.environ: - return os.path.join(os.environ['HOME'], '.cache', 'nvdiffrast', *paths) - if 'USERPROFILE' in os.environ: - return os.path.join(os.environ['USERPROFILE'], '.cache', 'nvdiffrast', *paths) - return os.path.join(tempfile.gettempdir(), '.cache', 'nvdiffrast', *paths) - -cuda_cache_version_tag = 'v1' -do_not_hash_included_headers = False # Speed up compilation by assuming that headers included by the CUDA code never change. Unsafe! -verbose = True # Print status messages to stdout. - -#---------------------------------------------------------------------------- -# Internal helper funcs. - -def _find_compiler_bindir(): - hostx64_paths = sorted(glob.glob('C:/Program Files/Microsoft Visual Studio/*/Enterprise/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Enterprise/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - hostx64_paths = sorted(glob.glob('C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64'), reverse=True) - if hostx64_paths != []: - return hostx64_paths[0] - vc_bin_dir = 'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin' - if os.path.isdir(vc_bin_dir): - return vc_bin_dir - return None - -def _get_compute_cap(device): - caps_str = device.physical_device_desc - m = re.search('compute capability: (\\d+).(\\d+)', caps_str) - major = m.group(1) - minor = m.group(2) - return (major, minor) - -def _get_cuda_gpu_arch_string(): - gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU'] - if len(gpus) == 0: - raise RuntimeError('No GPU devices found') - (major, minor) = _get_compute_cap(gpus[0]) - return 'sm_%s%s' % (major, minor) - -def _run_cmd(cmd): - with os.popen(cmd) as pipe: - output = pipe.read() - status = pipe.close() - if status is not None: - raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output)) - -def _prepare_nvcc_cli(opts): - cmd = 'nvcc ' + opts.strip() - cmd += ' --disable-warnings' - cmd += ' --include-path "%s"' % tf.sysconfig.get_include() - cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src') - cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl') - cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive') - - compiler_bindir = _find_compiler_bindir() - if compiler_bindir is None: - # Require that _find_compiler_bindir succeeds on Windows. Allow - # nvcc to use whatever is the default on Linux. - if os.name == 'nt': - raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__) - else: - cmd += ' --compiler-bindir "%s"' % compiler_bindir - cmd += ' 2>&1' - return cmd - -#---------------------------------------------------------------------------- -# Main entry point. - -_plugin_cache = dict() - -def get_plugin(cuda_file, extra_nvcc_options=[]): - cuda_file_base = os.path.basename(cuda_file) - cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base) - - # Already in cache? - if cuda_file in _plugin_cache: - return _plugin_cache[cuda_file] - - # Setup plugin. - if verbose: - print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True) - try: - # Hash CUDA source. - md5 = hashlib.md5() - with open(cuda_file, 'rb') as f: - md5.update(f.read()) - md5.update(b'\n') - - # Hash headers included by the CUDA code by running it through the preprocessor. - if not do_not_hash_included_headers: - if verbose: - print('Preprocessing... ', end='', flush=True) - with tempfile.TemporaryDirectory() as tmp_dir: - tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext) - _run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))) - with open(tmp_file, 'rb') as f: - bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros - good_file_str = ('"' + cuda_file_base + '"').encode('utf-8') - for ln in f: - if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas - ln = ln.replace(bad_file_str, good_file_str) - md5.update(ln) - md5.update(b'\n') - - # Select compiler options. - compile_opts = '' - if os.name == 'nt': - compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib') - compile_opts += ' --library-path="%s"' % (os.path.dirname(__file__) + r"\..\lib") # Find libraries during compilation. - elif os.name == 'posix': - compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.so') - compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\'' - else: - assert False # not Windows or Linux, w00t? - compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string() - compile_opts += ' --use_fast_math' - for opt in extra_nvcc_options: - compile_opts += ' ' + opt - nvcc_cmd = _prepare_nvcc_cli(compile_opts) - - # Hash build configuration. - md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n') - md5.update(('tf.VERSION: ' + tf.VERSION).encode('utf-8') + b'\n') - md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n') - - # Compile if not already compiled. - bin_file_ext = '.dll' if os.name == 'nt' else '.so' - cuda_cache_path = make_cache_dir_path() - bin_file = os.path.join(make_cache_dir_path(), cuda_file_name + '_' + md5.hexdigest() + bin_file_ext) - if not os.path.isfile(bin_file): - if verbose: - print('Compiling... ', end='', flush=True) - with tempfile.TemporaryDirectory() as tmp_dir: - tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext) - _run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)) - os.makedirs(cuda_cache_path, exist_ok=True) - intermediate_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext) - shutil.copyfile(tmp_file, intermediate_file) - os.rename(intermediate_file, bin_file) # atomic - - # Load. - if verbose: - print('Loading... ', end='', flush=True) - plugin = tf.load_op_library(bin_file) - - # Add to cache. - _plugin_cache[cuda_file] = plugin - if verbose: - print('Done.', flush=True) - return plugin - - except: - if verbose: - print('Failed!', flush=True) - raise - -#---------------------------------------------------------------------------- diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_all.cu b/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_all.cu deleted file mode 100644 index 21b32399c8a2429faa89e3356a34837d0d7e3138..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_all.cu +++ /dev/null @@ -1,36 +0,0 @@ -// 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. - -// TF-specific helpers. - -#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal("Cuda error: ", cudaGetErrorName(err), "[", #CUDA_CALL, ";]")); } while (0) -#define OP_CHECK_GL_ERROR(CTX, GL_CALL) do { GL_CALL; GLenum err = glGetError(); OP_REQUIRES(CTX, err == GL_NO_ERROR, errors::Internal("OpenGL error: ", getGLErrorString(err), "[", #GL_CALL, ";]")); } while (0) - -// Cuda kernels and CPP all together. What an absolute compilation unit. - -#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__ -#include "../common/framework.h" -#include "../common/glutil.cpp" - -#include "../common/common.h" -#include "../common/common.cpp" - -#include "../common/rasterize.h" -#include "../common/rasterize_gl.cpp" -#include "../common/rasterize.cu" -#include "tf_rasterize.cu" - -#include "../common/interpolate.cu" -#include "tf_interpolate.cu" - -#include "../common/texture.cpp" -#include "../common/texture.cu" -#include "tf_texture.cu" - -#include "../common/antialias.cu" -#include "tf_antialias.cu" diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_antialias.cu b/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_antialias.cu deleted file mode 100644 index a8f1b5930a1c4145bf63f925168e01728c36a28c..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_antialias.cu +++ /dev/null @@ -1,278 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Forward TensorFlow op. - -struct AntialiasFwdOp : public OpKernel -{ - AntialiasKernelParams m_attribs; - - AntialiasFwdOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - OP_REQUIRES_OK(ctx, ctx->GetAttr("tri_const", &m_attribs.tri_const)); - } - - void Compute(OpKernelContext* ctx) - { - AntialiasKernelParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - - // Get input. - const Tensor& color = ctx->input(0); - const Tensor& rasterOut = ctx->input(1); - const Tensor& pos = ctx->input(2); - const Tensor& tri = ctx->input(3); - - // Instance rendering mode? - p.instance_mode = pos.dims() > 2; - - // Extract input dimensions. - if (p.instance_mode) - p.numVertices = (pos.dims() > 1) ? pos.dim_size(1) : 0; - else - p.numVertices = (pos.dims() > 0) ? pos.dim_size(0) : 0; - p.numTriangles = (tri.dims() > 0) ? tri.dim_size(0) : 0; - p.n = (color.dims() > 0) ? color.dim_size(0) : 0; - p.height = (color.dims() > 1) ? color.dim_size(1) : 0; - p.width = (color.dims() > 2) ? color.dim_size(2) : 0; - p.channels = (color.dims() > 3) ? color.dim_size(3) : 0; - - // Sanity checks. - OP_REQUIRES(ctx, color.dims() == 4 && color.dim_size(0) > 0 && color.dim_size(1) > 0 && color.dim_size(2) > 0 && color.dim_size(3) > 0, errors::InvalidArgument("color must have shape[>0, >0, >0, >0]")); - OP_REQUIRES(ctx, rasterOut.dims() == 4 && rasterOut.dim_size(0) > 0 && rasterOut.dim_size(1) > 0 && rasterOut.dim_size(2) > 0 && rasterOut.dim_size(3) == 4, errors::InvalidArgument("raster_out must have shape[>0, >0, >0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, color.dim_size(1) == rasterOut.dim_size(1) && color.dim_size(2) == rasterOut.dim_size(2), errors::InvalidArgument("color and raster_out inputs must have same spatial dimensions")); - if (p.instance_mode) - { - OP_REQUIRES(ctx, pos.dims() == 3 && pos.dim_size(0) > 0 && pos.dim_size(1) > 0 && pos.dim_size(2) == 4, errors::InvalidArgument("pos must have shape [>0, >0, 4] or [>0, 4]")); - OP_REQUIRES(ctx, rasterOut.dim_size(0) == p.n && pos.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs color, raster_out, pos")); - } - else - { - OP_REQUIRES(ctx, pos.dims() == 2 && pos.dim_size(0) > 0 && pos.dim_size(1) == 4, errors::InvalidArgument("pos must have shape [>0, >0, 4] or [>0, 4]")); - OP_REQUIRES(ctx, rasterOut.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs color, raster_out")); - } - - // Get input pointers. - p.color = color.flat().data(); - p.rasterOut = rasterOut.flat().data(); - p.tri = tri.flat().data(); - p.pos = pos.flat().data(); - - // Misc parameters. - p.xh = .5f * (float)p.width; - p.yh = .5f * (float)p.height; - - // Allocate output tensor. - Tensor* outputTensor = NULL; - TensorShape outputShape; - outputShape.AddDim(p.n); - outputShape.AddDim(p.height); - outputShape.AddDim(p.width); - outputShape.AddDim(p.channels); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, outputShape, &outputTensor)); - p.output = outputTensor->flat().data(); - - // Allocate work buffer. One extra int4 for storing counters. - Tensor* workTensor = NULL; - TensorShape workShape; - workShape.AddDim(p.n * p.width * p.height * 8 + 4); // 8 int for a maximum of two work items per pixel. - OP_REQUIRES_OK(ctx, ctx->allocate_output(1, workShape, &workTensor)); - p.workBuffer = (int4*)(workTensor->flat().data()); - - // Clear the work counters. - OP_CHECK_CUDA_ERROR(ctx, cudaMemsetAsync(p.workBuffer, 0, sizeof(int4), stream)); - - // Verify that buffers are aligned to allow float2/float4 operations. - OP_REQUIRES(ctx, !((uintptr_t)p.pos & 15), errors::Internal("pos input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.rasterOut & 7), errors::Internal("raster_out input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.workBuffer & 15), errors::Internal("work_buffer internal tensor not aligned to int4")); - - // Kernel parameters. - void* args[] = {&p}; - - // (Re-)calculate opposite vertex hash. - if (!p.evHash || !p.tri_const) - { - if (p.allocTriangles < p.numTriangles) - { - p.allocTriangles = max(p.allocTriangles, 64); - while (p.allocTriangles < p.numTriangles) - p.allocTriangles <<= 1; // Must be power of two. - - // (Re-)allocate memory for the hash. - OP_CHECK_CUDA_ERROR(ctx, cudaFree(p.evHash)); - OP_CHECK_CUDA_ERROR(ctx, cudaMalloc(&p.evHash, p.allocTriangles * AA_HASH_ELEMENTS_PER_TRIANGLE(p.allocTriangles) * sizeof(uint4))); - LOG(INFO) << "Increasing topology hash size to accommodate " << p.allocTriangles << " triangles"; - } - - // Clear the hash and launch the mesh kernel to populate it. - OP_CHECK_CUDA_ERROR(ctx, cudaMemsetAsync(p.evHash, 0, p.allocTriangles * AA_HASH_ELEMENTS_PER_TRIANGLE(p.allocTriangles) * sizeof(uint4), stream)); - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)AntialiasFwdMeshKernel, (p.numTriangles - 1) / AA_MESH_KERNEL_THREADS_PER_BLOCK + 1, AA_MESH_KERNEL_THREADS_PER_BLOCK, args, 0, stream)); - } - - // Copy input to output as a baseline. - OP_CHECK_CUDA_ERROR(ctx, cudaMemcpyAsync(p.output, p.color, p.n * p.height * p.width * p.channels * sizeof(float), cudaMemcpyDeviceToDevice, stream)); - - // Choose launch parameters for the discontinuity finder kernel and launch. - dim3 blockSize(AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH, AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT, 1); - dim3 gridSize = getLaunchGridSize(blockSize, p.width, p.height, p.n); - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)AntialiasFwdDiscontinuityKernel, gridSize, blockSize, args, 0, stream)); - - // Determine optimum block size for the persistent analysis kernel. - int device = 0; - int numCTA = 0; - int numSM = 0; - OP_CHECK_CUDA_ERROR(ctx, cudaGetDevice(&device)); - OP_CHECK_CUDA_ERROR(ctx, cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numCTA, (void*)AntialiasFwdAnalysisKernel, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK, 0)); - OP_CHECK_CUDA_ERROR(ctx, cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, device)); - - // Launch analysis kernel. - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)AntialiasFwdAnalysisKernel, numCTA * numSM, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK, args, 0, stream)); - } -}; - -REGISTER_OP("AntialiasFwd") - .Input ("color: float") - .Input ("raster_out: float") - .Input ("pos: float") - .Input ("tri: int32") - .Output ("output: float") - .Output ("work_buffer: int32") - .Attr ("tri_const: int"); - -REGISTER_KERNEL_BUILDER(Name("AntialiasFwd").Device(DEVICE_GPU), AntialiasFwdOp); - -//------------------------------------------------------------------------ -// Gradient TensorFlow op. - -struct AntialiasGradOp : public OpKernel -{ - AntialiasKernelParams m_attribs; - - AntialiasGradOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - } - - void Compute(OpKernelContext* ctx) - { - AntialiasKernelParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - - // Get input. - const Tensor& color = ctx->input(0); - const Tensor& rasterOut = ctx->input(1); - const Tensor& pos = ctx->input(2); - const Tensor& tri = ctx->input(3); - const Tensor& dy = ctx->input(4); - const Tensor& workBuffer = ctx->input(5); - - // Instance rendering mode? - p.instance_mode = pos.dims() > 2; - - // Extract input dimensions. - if (p.instance_mode) - p.numVertices = (pos.dims() > 1) ? pos.dim_size(1) : 0; - else - p.numVertices = (pos.dims() > 0) ? pos.dim_size(0) : 0; - p.numTriangles = (tri.dims() > 0) ? tri.dim_size(0) : 0; - p.n = (color.dims() > 0) ? color.dim_size(0) : 0; - p.height = (color.dims() > 1) ? color.dim_size(1) : 0; - p.width = (color.dims() > 2) ? color.dim_size(2) : 0; - p.channels = (color.dims() > 3) ? color.dim_size(3) : 0; - - // Sanity checks. - OP_REQUIRES(ctx, dy.dims() == 4 && dy.dim_size(0) > 0 && dy.dim_size(1) > 0 && dy.dim_size(2) > 0 && dy.dim_size(3) > 0, errors::InvalidArgument("dy must have shape[>0, >0, >0, >0]")); - OP_REQUIRES(ctx, color.dims() == 4 && color.dim_size(0) > 0 && color.dim_size(1) > 0 && color.dim_size(2) > 0 && color.dim_size(3) > 0, errors::InvalidArgument("color must have shape[>0, >0, >0, >0]")); - OP_REQUIRES(ctx, rasterOut.dims() == 4 && rasterOut.dim_size(0) > 0 && rasterOut.dim_size(1) > 0 && rasterOut.dim_size(2) > 0 && rasterOut.dim_size(3) == 4, errors::InvalidArgument("raster_out must have shape[>0, >0, >0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, color.dim_size(1) == rasterOut.dim_size(1) && color.dim_size(2) == rasterOut.dim_size(2), errors::InvalidArgument("color and raster_out inputs must have same spatial dimensions")); - OP_REQUIRES(ctx, color.dim_size(1) == dy.dim_size(1) && color.dim_size(2) == dy.dim_size(2) && color.dim_size(3) == dy.dim_size(3), errors::InvalidArgument("color and dy inputs must have same dimensions")); - if (p.instance_mode) - { - OP_REQUIRES(ctx, pos.dims() == 3 && pos.dim_size(0) > 0 && pos.dim_size(1) > 0 && pos.dim_size(2) == 4, errors::InvalidArgument("pos must have shape [>0, >0, 4] or [>0, 4]")); - OP_REQUIRES(ctx, rasterOut.dim_size(0) == p.n && pos.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs color, raster_out, pos")); - OP_REQUIRES(ctx, dy.dim_size(0) == p.n && rasterOut.dim_size(0) == p.n && pos.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs dy, color, raster_out, pos")); - } - else - { - OP_REQUIRES(ctx, pos.dims() == 2 && pos.dim_size(0) > 0 && pos.dim_size(1) == 4, errors::InvalidArgument("pos must have shape [>0, >0, 4] or [>0, 4]")); - OP_REQUIRES(ctx, rasterOut.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs color, raster_out")); - OP_REQUIRES(ctx, dy.dim_size(0) == p.n && rasterOut.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs dy, color, raster_out")); - } - - // Get input pointers. - p.dy = dy.flat().data(); - p.color = color.flat().data(); - p.rasterOut = rasterOut.flat().data(); - p.tri = tri.flat().data(); - p.pos = pos.flat().data(); - p.workBuffer = (int4*)(workBuffer.flat().data()); - - // Misc parameters. - p.xh = .5f * (float)p.width; - p.yh = .5f * (float)p.height; - - // Allocate color gradient output tensor. - Tensor* gradColor = NULL; - TensorShape gradColorShape; - gradColorShape.AddDim(p.n); - gradColorShape.AddDim(p.height); - gradColorShape.AddDim(p.width); - gradColorShape.AddDim(p.channels); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, gradColorShape, &gradColor)); - p.gradColor = gradColor->flat().data(); - - // Allocate position gradient output tensor. - Tensor* gradPos = NULL; - TensorShape gradPosShape; - if (p.instance_mode) - gradPosShape.AddDim(p.n); - gradPosShape.AddDim(p.numVertices); - gradPosShape.AddDim(4); - OP_REQUIRES_OK(ctx, ctx->allocate_output(1, gradPosShape, &gradPos)); - p.gradPos = gradPos->flat().data(); - - // Initialize all the stuff. - OP_CHECK_CUDA_ERROR(ctx, cudaMemsetAsync(&p.workBuffer[0].y, 0, sizeof(int), stream)); // Gradient kernel work counter. - OP_CHECK_CUDA_ERROR(ctx, cudaMemcpyAsync(p.gradColor, p.dy, p.n * p.height * p.width * p.channels * sizeof(float), cudaMemcpyDeviceToDevice, stream)); - OP_CHECK_CUDA_ERROR(ctx, cudaMemsetAsync(p.gradPos, 0, (p.instance_mode ? p.n : 1) * p.numVertices * 4 * sizeof(float), stream)); - - // Verify that buffers are aligned to allow float2/float4 operations. - OP_REQUIRES(ctx, !((uintptr_t)p.pos & 15), errors::Internal("pos input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.workBuffer & 15), errors::Internal("work_buffer internal tensor not aligned to int4")); - - // Launch the gradient kernel. - void* args[] = {&p}; - - int device = 0; - int numCTA = 0; - int numSM = 0; - OP_CHECK_CUDA_ERROR(ctx, cudaGetDevice(&device)); - OP_CHECK_CUDA_ERROR(ctx, cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numCTA, (void*)AntialiasGradKernel, AA_GRAD_KERNEL_THREADS_PER_BLOCK, 0)); - OP_CHECK_CUDA_ERROR(ctx, cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, device)); - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)AntialiasGradKernel, numCTA * numSM, AA_GRAD_KERNEL_THREADS_PER_BLOCK, args, 0, stream)); - } -}; - -REGISTER_OP("AntialiasGrad") - .Input ("color: float") - .Input ("raster_out: float") - .Input ("pos: float") - .Input ("tri: int32") - .Input ("dy: float") - .Input ("work_buffer: int32") - .Output ("grad_color: float") - .Output ("grad_pos: float"); - -REGISTER_KERNEL_BUILDER(Name("AntialiasGrad").Device(DEVICE_GPU), AntialiasGradOp); - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_interpolate.cu b/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_interpolate.cu deleted file mode 100644 index 6b44165cab026937014362893751e5c26be3fac3..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_interpolate.cu +++ /dev/null @@ -1,301 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Common op attribute parser. - -static __host__ void interpolateParseOpAttributes(OpKernelConstruction* ctx, InterpolateKernelParams& p, bool enableDA) -{ - if (enableDA) - { - OP_REQUIRES_OK(ctx, ctx->GetAttr("diff_attrs_all", &p.diff_attrs_all)); - if (!p.diff_attrs_all) - { - std::vector diff_attrs_vec; - OP_REQUIRES_OK(ctx, ctx->GetAttr("diff_attrs", &diff_attrs_vec)); - OP_REQUIRES(ctx, diff_attrs_vec.size() > 0, errors::InvalidArgument("differentiation enabled with empty diff_attrs list")); - OP_REQUIRES(ctx, diff_attrs_vec.size() <= IP_MAX_DIFF_ATTRS, errors::InvalidArgument("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 TensorFlow op. - -template -struct InterpolateFwdOp : public OpKernel -{ - InterpolateKernelParams m_attribs; - - InterpolateFwdOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - interpolateParseOpAttributes(ctx, m_attribs, ENABLE_DA); - } - - void Compute(OpKernelContext* ctx) - { - InterpolateKernelParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - - // Get input. - const Tensor& attr = ctx->input(0); - const Tensor& rast = ctx->input(1); - const Tensor& tri = ctx->input(2); - const Tensor& rast_db = ctx->input(ENABLE_DA ? 3 : 2); - - // Instance rendering mode? - p.instance_mode = attr.dims() > 2; - - // Extract input dimensions. - if (p.instance_mode) - { - p.numVertices = (attr.dims() > 1) ? attr.dim_size(1) : 0; - p.numAttr = (attr.dims() > 2) ? attr.dim_size(2) : 0; - } - else - { - p.numVertices = (attr.dims() > 0) ? attr.dim_size(0) : 0; - p.numAttr = (attr.dims() > 1) ? attr.dim_size(1) : 0; - } - p.numTriangles = (tri.dims() > 0) ? tri.dim_size(0) : 0; - p.height = (rast.dims() > 1) ? rast.dim_size(1) : 0; - p.width = (rast.dims() > 2) ? rast.dim_size(2) : 0; - p.depth = (rast.dims() > 0) ? rast.dim_size(0) : 0; - - // Sanity checks. - OP_REQUIRES(ctx, rast.dims() == 4 && rast.dim_size(0) > 0 && rast.dim_size(1) > 0 && rast.dim_size(2) > 0 && rast.dim_size(3) == 4, errors::InvalidArgument("rast must have shape[>0, >0, >0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, (attr.dims() == 2 || attr.dims() == 3) && attr.dim_size(0) > 0 && attr.dim_size(1) > 0 && (attr.dims() == 2 || attr.dim_size(2) > 0), errors::InvalidArgument("attr must have shape [>0, >0, >0] or [>0, >0]")); - if (p.instance_mode) - OP_REQUIRES(ctx, attr.dim_size(0) == p.depth || attr.dim_size(0) == 1, errors::InvalidArgument("minibatch size mismatch between inputs rast, attr")); - if (ENABLE_DA) - { - OP_REQUIRES(ctx, rast_db.dims() == 4 && rast_db.dim_size(0) > 0 && rast_db.dim_size(1) > 0 && rast_db.dim_size(2) > 0 && rast_db.dim_size(3) == 4, errors::InvalidArgument("rast_db must have shape[>0, >0, >0, 4]")); - OP_REQUIRES(ctx, rast_db.dim_size(1) == rast.dim_size(1) && rast_db.dim_size(2) == rast.dim_size(2), errors::InvalidArgument("spatial size mismatch between inputs rast and rast_db")); - OP_REQUIRES(ctx, rast_db.dim_size(0) == p.depth, errors::InvalidArgument("minibatch size mismatch between inputs rast, rast_db")); - } - - // All diff attrs mode. - if (p.diff_attrs_all) - p.numDiffAttr = p.numAttr; - - // Get input pointers. - p.attr = attr.flat().data(); - p.rast = rast.flat().data(); - p.tri = tri.flat().data(); - p.attrBC = (p.instance_mode && attr.dim_size(0) == 1) ? 1 : 0; - p.rastDB = ENABLE_DA ? rast_db.flat().data() : 0; - - // Allocate main output tensor. - Tensor* out_tensor = NULL; - TensorShape out_shape; - out_shape.AddDim(p.depth); - out_shape.AddDim(p.height); - out_shape.AddDim(p.width); - out_shape.AddDim(p.numAttr); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, out_shape, &out_tensor)); - p.out = out_tensor->flat().data(); - - // Allocate pixel differential output tensor. - Tensor* out_da_tensor = NULL; - out_shape.set_dim(3, p.numDiffAttr * 2); - OP_REQUIRES_OK(ctx, ctx->allocate_output(1, out_shape, &out_da_tensor)); - p.outDA = ENABLE_DA ? out_da_tensor->flat().data() : 0; - - // Verify that buffers are aligned to allow float2/float4 operations. - OP_REQUIRES(ctx, !((uintptr_t)p.rast & 15), errors::Internal("rast input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.rastDB & 15), errors::Internal("rast_db input tensor not aligned to float4")); - if (ENABLE_DA) - OP_REQUIRES(ctx, !((uintptr_t)p.outDA & 7), errors::Internal("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; - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream)); - } -}; - -REGISTER_OP("InterpolateFwd") - .Input ("attr: float") - .Input ("rast: float") - .Input ("tri: int32") - .Output ("out: float") - .Output ("out_da: float"); - -REGISTER_OP("InterpolateFwdDa") - .Input ("attr: float") - .Input ("rast: float") - .Input ("tri: int32") - .Input ("rast_db: float") - .Output ("out: float") - .Output ("out_da: float") - .Attr ("diff_attrs_all: int") - .Attr ("diff_attrs: list(int)"); - -REGISTER_KERNEL_BUILDER(Name("InterpolateFwd") .Device(DEVICE_GPU), InterpolateFwdOp); -REGISTER_KERNEL_BUILDER(Name("InterpolateFwdDa").Device(DEVICE_GPU), InterpolateFwdOp); - -//------------------------------------------------------------------------ -// Gradient TensorFlow op. - -template -struct InterpolateGradOp : public OpKernel -{ - InterpolateKernelParams m_attribs; - - InterpolateGradOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - interpolateParseOpAttributes(ctx, m_attribs, ENABLE_DA); - } - - void Compute(OpKernelContext* ctx) - { - InterpolateKernelParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - - // Get input. - const Tensor& attr = ctx->input(0); - const Tensor& rast = ctx->input(1); - const Tensor& tri = ctx->input(2); - const Tensor& dy = ctx->input(3); - const Tensor& rast_db = ctx->input(ENABLE_DA ? 4 : 3); - const Tensor& dda = ctx->input(ENABLE_DA ? 5 : 3); - - // Instance rendering mode? - p.instance_mode = attr.dims() > 2; - - // Extract input dimensions. - if (p.instance_mode) - { - p.numVertices = (attr.dims() > 1) ? attr.dim_size(1) : 0; - p.numAttr = (attr.dims() > 2) ? attr.dim_size(2) : 0; - } - else - { - p.numVertices = (attr.dims() > 0) ? attr.dim_size(0) : 0; - p.numAttr = (attr.dims() > 1) ? attr.dim_size(1) : 0; - } - p.numTriangles = (tri.dims() > 0) ? tri.dim_size(0) : 0; - p.depth = (rast.dims() > 0) ? rast.dim_size(0) : 0; - p.height = (rast.dims() > 1) ? rast.dim_size(1) : 0; - p.width = (rast.dims() > 2) ? rast.dim_size(2) : 0; - int attr_depth = p.instance_mode ? (attr.dims() > 1 ? attr.dim_size(0) : 0) : 1; - - // Sanity checks. - OP_REQUIRES(ctx, rast.dims() == 4 && rast.dim_size(0) > 0 && rast.dim_size(1) > 0 && rast.dim_size(2) > 0 && rast.dim_size(3) == 4, errors::InvalidArgument("rast must have shape[>0, >0, >0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, (attr.dims() == 2 || attr.dims() == 3) && attr.dim_size(0) > 0 && attr.dim_size(1) > 0 && (attr.dims() == 2 || attr.dim_size(2) > 0), errors::InvalidArgument("attr must have shape [>0, >0, >0] or [>0, >0]")); - OP_REQUIRES(ctx, dy.dims() == 4 && dy.dim_size(0) > 0 && dy.dim_size(1) == p.height && dy.dim_size(2) == p.width && dy.dim_size(3) > 0, errors::InvalidArgument("dy must have shape [>0, height, width, >0]")); - OP_REQUIRES(ctx, dy.dim_size(3) == p.numAttr, errors::InvalidArgument("argument count mismatch between inputs dy, attr")); - OP_REQUIRES(ctx, (attr_depth == p.depth || attr_depth == 1) && dy.dim_size(0) == p.depth, errors::InvalidArgument("minibatch size mismatch between inputs rast, dy, attr")); - if (ENABLE_DA) - { - OP_REQUIRES(ctx, dda.dims() == 4 && dda.dim_size(0) > 0 && dda.dim_size(1) == p.height && dda.dim_size(2) == p.width, errors::InvalidArgument("dda must have shape [>0, height, width, ?]")); - OP_REQUIRES(ctx, dda.dim_size(0) == p.depth, errors::InvalidArgument("minibatch size mismatch between rast, dda")); - } - - // All diff attrs mode. - if (p.diff_attrs_all) - p.numDiffAttr = p.numAttr; - - // Get input pointers. - p.attr = attr.flat().data(); - p.rast = rast.flat().data(); - p.tri = tri.flat().data(); - p.dy = dy.flat().data(); - p.rastDB = ENABLE_DA ? rast_db.flat().data() : 0; - p.dda = ENABLE_DA ? dda.flat().data() : 0; - p.attrBC = (p.instance_mode && attr_depth < p.depth) ? 1 : 0; - - // Allocate attribute gradient output tensor. - Tensor* grad_attr_tensor = NULL; - TensorShape grad_attr_shape; - if (p.instance_mode) - grad_attr_shape.AddDim(attr_depth); - grad_attr_shape.AddDim(p.numVertices); - grad_attr_shape.AddDim(p.numAttr); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, grad_attr_shape, &grad_attr_tensor)); - p.gradAttr = grad_attr_tensor->flat().data(); - - // Allocate bary gradient output tensor. - Tensor* grad_rast_tensor = NULL; - TensorShape grad_rast_shape; - grad_rast_shape.AddDim(p.depth); - grad_rast_shape.AddDim(p.height); - grad_rast_shape.AddDim(p.width); - grad_rast_shape.AddDim(4); - OP_REQUIRES_OK(ctx, ctx->allocate_output(1, grad_rast_shape, &grad_rast_tensor)); - p.gradRaster = grad_rast_tensor->flat().data(); - - // Allocate bary pixel diff gradient output tensor. - if (ENABLE_DA) - { - Tensor* grad_rast_db_tensor = NULL; - OP_REQUIRES_OK(ctx, ctx->allocate_output(2, grad_rast_shape, &grad_rast_db_tensor)); - p.gradRasterDB = grad_rast_db_tensor->flat().data(); - } - - // Clear attribute gradients. - cudaMemsetAsync(p.gradAttr, 0, attr_depth * p.numVertices * p.numAttr * sizeof(float), stream); - - // Verify that buffers are aligned to allow float2/float4 operations. - OP_REQUIRES(ctx, !((uintptr_t)p.rast & 15), errors::Internal("rast input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradRaster & 15), errors::Internal("grad_rast output tensor not aligned to float4")); - if (ENABLE_DA) - { - OP_REQUIRES(ctx, !((uintptr_t)p.dda & 7), errors::Internal("dda input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.rastDB & 15), errors::Internal("rast_db input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradRasterDB & 15), errors::Internal("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; - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream)); - } -}; - -REGISTER_OP("InterpolateGrad") - .Input ("attr: float") - .Input ("rast: float") - .Input ("tri: int32") - .Input ("dy: float") - .Output ("grad_attr: float") - .Output ("grad_rast: float") - ; - -REGISTER_OP("InterpolateGradDa") - .Input ("attr: float") - .Input ("rast: float") - .Input ("tri: int32") - .Input ("dy: float") - .Input ("rast_db: float") - .Input ("dda: float") - .Output ("grad_attr: float") - .Output ("grad_rast: float") - .Output ("grad_rast_db: float") - .Attr ("diff_attrs_all: int") - .Attr ("diff_attrs: list(int)"); - ; - -REGISTER_KERNEL_BUILDER(Name("InterpolateGrad") .Device(DEVICE_GPU), InterpolateGradOp); -REGISTER_KERNEL_BUILDER(Name("InterpolateGradDa").Device(DEVICE_GPU), InterpolateGradOp); - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_rasterize.cu b/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_rasterize.cu deleted file mode 100644 index de7e1e540e1154705418a348c722698923f4fca4..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_rasterize.cu +++ /dev/null @@ -1,242 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Forward TensorFlow op. - -struct RasterizeFwdOp : public OpKernel -{ - RasterizeGLState m_glState; // OpenGL-related persistent state. - int m_tri_const; // 1 if triangle array is known to be constant. - - RasterizeFwdOp(OpKernelConstruction* ctx): - OpKernel(ctx) - { - memset(&m_glState, 0, sizeof(RasterizeGLState)); - OP_REQUIRES_OK(ctx, ctx->GetAttr("enable_db", &m_glState.enableDB)); - OP_REQUIRES_OK(ctx, ctx->GetAttr("tri_const", &m_tri_const)); - } - - void Compute(OpKernelContext* ctx) - { - cudaStream_t stream = ctx->eigen_device().stream(); - - // Check that input shapes are correct. - const Tensor& pos = ctx->input(0); - const Tensor& tri = ctx->input(1); - const Tensor& resolution = ctx->input(2); - const Tensor& ranges = ctx->input(3); - - // Determine number of outputs - int num_outputs = m_glState.enableDB ? 2 : 1; - - // Determine instance mode and check input dimensions. - bool instance_mode = pos.dims() > 2; - if (instance_mode) - { - OP_REQUIRES(ctx, pos.dims() == 3 && pos.dim_size(0) > 0 && pos.dim_size(1) > 0 && pos.dim_size(2) == 4, errors::InvalidArgument("instance mode - pos must have shape [>0, >0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, resolution.dims() == 1 && resolution.dim_size(0) == 2, errors::InvalidArgument("resolution must have shape [2]")); - } - else - { - OP_REQUIRES(ctx, pos.dims() == 2 && pos.dim_size(0) > 0 && pos.dim_size(1) == 4, errors::InvalidArgument("range mode - pos must have shape [>0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, resolution.dims() == 1 && resolution.dim_size(0) == 2, errors::InvalidArgument("resolution must have shape [2]")); - OP_REQUIRES(ctx, ranges.dims() == 2 && ranges.dim_size(0) > 0 && ranges.dim_size(1) == 2, errors::InvalidArgument("range mode - ranges must have shape [>0, 2]")); - } - - // Get output shape. - const int32_t* res_in = resolution.flat().data(); // This is in CPU memory. - int height = res_in[0]; - int width = res_in[1]; - int depth = instance_mode ? pos.dim_size(0) : ranges.dim_size(0); - OP_REQUIRES(ctx, height > 0 && width > 0, errors::InvalidArgument("resolution must be [>0, >0]")); - - // Get position and triangle buffer sizes in int32/float32. - int posCount = 4 * pos.dim_size(0) * (instance_mode ? pos.dim_size(1) : 1); - int triCount = 3 * tri.dim_size(0); - - // Init context and GL? - bool initCtx = !m_glState.glFBO; - if (initCtx) - { - const DeviceBase::GpuDeviceInfo* g = ctx->device()->tensorflow_gpu_device_info(); - int cudaDeviceIdx = g ? g->gpu_id : -1; - rasterizeInitGLContext(ctx, m_glState, cudaDeviceIdx); // In common/rasterize.cpp - } - else - setGLContext(m_glState.glctx); // (Re-)Activate GL context. - - // Resize all buffers. - bool changes = false; - rasterizeResizeBuffers(ctx, m_glState, changes, posCount, triCount, width, height, depth); // In common/rasterize_gl.cpp - if (changes) - { -#ifdef _WIN32 - // Workaround for occasional blank first frame on Windows. - releaseGLContext(); - setGLContext(m_glState.glctx); -#endif - } - - // Copy input data to GL and render. - const float* posPtr = pos.flat().data(); - const int32_t* rangesPtr = instance_mode ? 0 : ranges.flat().data(); // This is in CPU memory. - const int32_t* triPtr = (initCtx || !m_tri_const) ? tri.flat().data() : NULL; // Copy triangles only if needed. - int vtxPerInstance = instance_mode ? pos.dim_size(1) : 0; - rasterizeRender(ctx, m_glState, stream, posPtr, posCount, vtxPerInstance, triPtr, triCount, rangesPtr, width, height, depth, -1); - - // Allocate output tensors. - TensorShape output_shape; - output_shape.AddDim(depth); - output_shape.AddDim(height); - output_shape.AddDim(width); - output_shape.AddDim(4); - float* outputPtr[2]; - for (int i=0; i < 2; i++) - { - if (i >= num_outputs) - output_shape.set_dim(3, 0); // Zero channels for unwanted out_db tensor. - Tensor* output_tensor = NULL; - OP_REQUIRES_OK(ctx, ctx->allocate_output(i, output_shape, &output_tensor)); - if (i < num_outputs) - outputPtr[i] = output_tensor->flat().data(); - } - - // Copy rasterized results into CUDA buffers. - rasterizeCopyResults(ctx, m_glState, stream, outputPtr, width, height, depth); - - // Done. Release GL context. - releaseGLContext(); - } -}; - -REGISTER_OP("RasterizeFwd") - .Input ("pos: float") - .Input ("tri: int32") - .Input ("resolution: int32") - .Input ("ranges: int32") - .Output ("out: float") - .Output ("out_db: float") - .Attr ("enable_db: int") - .Attr ("tri_const: int"); - -REGISTER_KERNEL_BUILDER(Name("RasterizeFwd").Device(DEVICE_GPU).HostMemory("resolution").HostMemory("ranges"), RasterizeFwdOp); - -//------------------------------------------------------------------------ -// Gradient TensorFlow op. - -template -struct RasterizeGradOp : public OpKernel -{ - RasterizeGradParams m_attribs; - - RasterizeGradOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - } - - void Compute(OpKernelContext* ctx) - { - RasterizeGradParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - - // Input tensors. - const Tensor& pos = ctx->input(0); - const Tensor& tri = ctx->input(1); - const Tensor& out = ctx->input(2); - const Tensor& dy = ctx->input(3); - const Tensor& ddb = ctx->input(ENABLE_DB ? 4 : 3); - - // Determine instance mode. - p.instance_mode = (pos.dims() > 2) ? 1 : 0; - - // Shape is taken from the rasterizer output tensor. - OP_REQUIRES(ctx, out.dims() == 4, errors::InvalidArgument("out must be rank-4")); - p.depth = out.dim_size(0); - p.height = out.dim_size(1); - p.width = out.dim_size(2); - OP_REQUIRES(ctx, p.depth > 0 && p.height > 0 && p.width > 0, errors::InvalidArgument("resolution must be [>0, >0, >0]")); - - // Check other shapes. - if (p.instance_mode) - OP_REQUIRES(ctx, pos.dims() == 3 && pos.dim_size(0) == p.depth && pos.dim_size(1) > 0 && pos.dim_size(2) == 4, errors::InvalidArgument("pos must have shape [depth, >0, 4]")); - else - OP_REQUIRES(ctx, pos.dims() == 2 && pos.dim_size(0) > 0 && pos.dim_size(1) == 4, errors::InvalidArgument("pos must have shape [>0, 4]")); - OP_REQUIRES(ctx, tri.dims() == 2 && tri.dim_size(0) > 0 && tri.dim_size(1) == 3, errors::InvalidArgument("tri must have shape [>0, 3]")); - OP_REQUIRES(ctx, out.dims() == 4 && out.dim_size(0) == p.depth && out.dim_size(1) == p.height && out.dim_size(2) == p.width && out.dim_size(3) == 4, errors::InvalidArgument("out must have shape [depth, height, width, 4]")); - OP_REQUIRES(ctx, dy.dims() == 4 && dy.dim_size(0) == p.depth && dy.dim_size(1) == p.height && dy.dim_size(2) == p.width && dy.dim_size(3) == 4, errors::InvalidArgument("dy must have shape [depth, height, width, 4]")); - if (ENABLE_DB) - OP_REQUIRES(ctx, ddb.dims() == 4 && ddb.dim_size(0) == p.depth && ddb.dim_size(1) == p.height && ddb.dim_size(2) == p.width && ddb.dim_size(3) == 4, errors::InvalidArgument("ddb must have shape [depth, height, width, 4]")); - - // Populate parameters. - p.numTriangles = tri.dim_size(0); - p.numVertices = p.instance_mode ? pos.dim_size(1) : pos.dim_size(0); - p.pos = pos.flat().data(); - p.tri = tri.flat().data(); - p.out = out.flat().data(); - p.dy = dy.flat().data(); - p.ddb = ENABLE_DB ? ddb.flat().data() : 0; - - // 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. - Tensor* grad_tensor = NULL; - TensorShape grad_shape; - if (p.instance_mode) - grad_shape.AddDim(p.depth); - grad_shape.AddDim(p.numVertices); - grad_shape.AddDim(4); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, grad_shape, &grad_tensor)); - p.grad = grad_tensor->flat().data(); - - // Clear the output buffers. - size_t gradBytes = (p.instance_mode ? p.depth : 1) * p.numVertices * 4 * sizeof(float); - cudaMemsetAsync(p.grad, 0, gradBytes, stream); - - // Verify that buffers are aligned to allow float2/float4 operations. - OP_REQUIRES(ctx, !((uintptr_t)p.pos & 15), errors::Internal("pos input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.dy & 7), errors::Internal("dy input tensor not aligned to float2")); - if (ENABLE_DB) - OP_REQUIRES(ctx, !((uintptr_t)p.ddb & 15), errors::Internal("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; - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(func, gridSize, blockSize, args, 0, stream)); - } -}; - -REGISTER_OP("RasterizeGrad") - .Input ("pos: float") - .Input ("tri: int32") - .Input ("out: float") - .Input ("dy: float") - .Output ("grad: float"); - -REGISTER_OP("RasterizeGradDb") - .Input ("pos: float") - .Input ("tri: int32") - .Input ("out: float") - .Input ("dy: float") - .Input ("ddb: float") - .Output ("grad: float"); - -REGISTER_KERNEL_BUILDER(Name("RasterizeGrad") .Device(DEVICE_GPU), RasterizeGradOp); -REGISTER_KERNEL_BUILDER(Name("RasterizeGradDb").Device(DEVICE_GPU), RasterizeGradOp); - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_texture.cu b/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_texture.cu deleted file mode 100644 index a43db1e20e2edc34f535afa548b166d823fe087f..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/tensorflow/tf_texture.cu +++ /dev/null @@ -1,525 +0,0 @@ -// 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. - -//------------------------------------------------------------------------ -// Common op attribute parser. - -static __host__ void parseOpAttributes(OpKernelConstruction* ctx, TextureKernelParams& p) -{ - // Mip and filter modes. - OP_REQUIRES_OK(ctx, ctx->GetAttr("filter_mode", &p.filterMode)); - OP_REQUIRES(ctx, p.filterMode >= 0 && p.filterMode < TEX_MODE_COUNT, errors::InvalidArgument("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) - { - OP_REQUIRES_OK(ctx, ctx->GetAttr("max_mip_level", &p.mipLevelLimit)); - OP_REQUIRES(ctx, p.mipLevelLimit >= -1, errors::InvalidArgument("invalid max_mip_level")); - ctx->GetAttr("tex_const", &p.texConst); // Only available in forward op. - } - - // Boundary mode. - OP_REQUIRES_OK(ctx, ctx->GetAttr("boundary_mode", &p.boundaryMode)); - OP_REQUIRES(ctx, p.boundaryMode >= 0 && p.boundaryMode < TEX_BOUNDARY_MODE_COUNT, errors::InvalidArgument("boundary_mode unsupported")); -} - -//------------------------------------------------------------------------ -// Forward TensorFlow op. - -struct TextureFwdOp : public OpKernel -{ - TextureKernelParams m_attribs; - PersistentTensor m_persistentMipTensor; // Used if texture is constant and mips are enabled. - bool m_persistentMipTensorInitialized; - - TextureFwdOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - m_persistentMipTensorInitialized = false; - parseOpAttributes(ctx, m_attribs); - } - - void Compute(OpKernelContext* ctx) - { - TextureKernelParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - bool cube_mode = (p.boundaryMode == TEX_BOUNDARY_MODE_CUBE); - - // Get input. - const Tensor& tex = ctx->input(0); - const Tensor& uv = ctx->input(1); - const Tensor& uv_da = ctx->input(p.enableMip ? 2 : 1); - - // Extract input dimensions. - p.n = (uv.dims() > 0) ? uv.dim_size(0) : 0; - p.imgHeight = (uv.dims() > 1) ? uv.dim_size(1) : 0; - p.imgWidth = (uv.dims() > 2) ? uv.dim_size(2) : 0; - p.texDepth = (tex.dims() > 0) ? tex.dim_size(0) : 0; - if (!cube_mode) - { - p.texHeight = (tex.dims() > 1) ? tex.dim_size(1) : 0; - p.texWidth = (tex.dims() > 2) ? tex.dim_size(2) : 0; - p.channels = (tex.dims() > 3) ? tex.dim_size(3) : 0; - } - else - { - p.texHeight = (tex.dims() > 2) ? tex.dim_size(2) : 0; - p.texWidth = (tex.dims() > 3) ? tex.dim_size(3) : 0; - p.channels = (tex.dims() > 4) ? tex.dim_size(4) : 0; - } - - // Sanity checks. - if (!cube_mode) - { - OP_REQUIRES(ctx, tex.dims() == 4 && tex.dim_size(0) > 0 && tex.dim_size(1) > 0 && tex.dim_size(2) > 0 && tex.dim_size(3) > 0, errors::InvalidArgument("tex must have shape[>0, >0, >0, >0]")); - OP_REQUIRES(ctx, uv.dims() == 4 && uv.dim_size(0) > 0 && uv.dim_size(1) > 0 && uv.dim_size(2) > 0 && uv.dim_size(3) == 2, errors::InvalidArgument("uv must have shape [>0, >0, >0, 2]")); - } - else - { - OP_REQUIRES(ctx, tex.dims() == 5 && tex.dim_size(0) > 0 && tex.dim_size(1) == 6 && tex.dim_size(2) > 0 && tex.dim_size(3) > 0 && tex.dim_size(4) > 0, errors::InvalidArgument("tex must have shape[>0, 6, >0, >0, >0] in cube map mode")); - OP_REQUIRES(ctx, uv.dims() == 4 && uv.dim_size(0) > 0 && uv.dim_size(1) > 0 && uv.dim_size(2) > 0 && uv.dim_size(3) == 3, errors::InvalidArgument("uv must have shape [>0, >0, >0, 3] in cube map mode")); - OP_REQUIRES(ctx, tex.dim_size(2) == tex.dim_size(3), errors::InvalidArgument("texture shape must be square in cube map mode")); - } - OP_REQUIRES(ctx, tex.dim_size(0) == 1 || tex.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs tex, uv")); - OP_REQUIRES(ctx, p.texWidth <= (1 << TEX_MAX_MIP_LEVEL) && p.texHeight <= (1 << TEX_MAX_MIP_LEVEL), errors::InvalidArgument("texture size too large")); - if (p.enableMip) - { - if (!cube_mode) - OP_REQUIRES(ctx, uv_da.dims() == 4 && uv_da.dim_size(0) == p.n && uv_da.dim_size(1) == p.imgHeight && uv_da.dim_size(2) == p.imgWidth && uv_da.dim_size(3) == 4, errors::InvalidArgument("uv_da must have shape [minibatch_size, height, width, 4]")); - else - OP_REQUIRES(ctx, uv_da.dims() == 4 && uv_da.dim_size(0) == p.n && uv_da.dim_size(1) == p.imgHeight && uv_da.dim_size(2) == p.imgWidth && uv_da.dim_size(3) == 6, errors::InvalidArgument("uv_da must have shape [minibatch_size, height, width, 6] in cube map mode")); - } - - // Get input pointers. - p.tex[0] = tex.flat().data(); - p.uv = uv.flat().data(); - p.uvDA = p.enableMip ? uv_da.flat().data() : 0; - - // Allocate output tensor. - Tensor* out_tensor = NULL; - TensorShape out_shape; - out_shape.AddDim(p.n); - out_shape.AddDim(p.imgHeight); - out_shape.AddDim(p.imgWidth); - out_shape.AddDim(p.channels); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, out_shape, &out_tensor)); - p.out = out_tensor->flat().data(); - - // 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) - { - // Generate mip offsets. - int mipOffsets[TEX_MAX_MIP_LEVEL]; - int mipTotal = calculateMipInfo(ctx, p, mipOffsets); - - // Mip output tensor. - Tensor* mip_tensor = NULL; - TensorShape mip_shape; - mip_shape.AddDim(mipTotal); - - // If texture is constant, calculate mip stack only once. - bool computeMip = true; - if (p.texConst) - { - // First execution? - if (!m_persistentMipTensorInitialized) - { - // Allocate a persistent mip tensor. - OP_REQUIRES_OK(ctx, ctx->allocate_persistent(DT_FLOAT, mip_shape, &m_persistentMipTensor, &mip_tensor)); - m_persistentMipTensorInitialized = true; - } - else - { - // Reuse the persistent tensor, do not recompute mip levels. - mip_tensor = m_persistentMipTensor.AccessTensor(ctx); - computeMip = false; - } - - // Set as output tensor as well. - ctx->set_output(1, *mip_tensor); - } - else - { - // Allocate an output tensor as usual. - OP_REQUIRES_OK(ctx, ctx->allocate_output(1, mip_shape, &mip_tensor)); - } - - pmip = mip_tensor->flat().data(); // Pointer to data. - for (int i=1; i <= p.mipLevelMax; i++) - p.tex[i] = pmip + mipOffsets[i]; // Pointers to mip levels. - - // Build mip levels if needed. - if (computeMip) - { - 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 }; - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(build_func_tbl[channel_div_idx], gridSize, blockSize, args, 0, stream)); - } - } - } - - // Verify that buffers are aligned to allow float2/float4 operations. Unused pointers are zero so always aligned. - if (!cube_mode) - OP_REQUIRES(ctx, !((uintptr_t)p.uv & 7), errors::Internal("uv input tensor not aligned to float2")); - if ((p.channels & 3) == 0) - { - OP_REQUIRES(ctx, !((uintptr_t)p.tex[0] & 15), errors::Internal("tex input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.out & 15), errors::Internal("out output tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)pmip & 15), errors::Internal("mip output tensor not aligned to float4")); - } - if ((p.channels & 1) == 0) - { - OP_REQUIRES(ctx, !((uintptr_t)p.tex[0] & 7), errors::Internal("tex input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.out & 7), errors::Internal("out output tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)pmip & 7), errors::Internal("mip output tensor not aligned to float2")); - } - if (!cube_mode) - OP_REQUIRES(ctx, !((uintptr_t)p.uvDA & 15), errors::Internal("uv_da input tensor not aligned to float4")); - else - OP_REQUIRES(ctx, !((uintptr_t)p.uvDA & 7), errors::Internal("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, and datatype. - void* func_tbl[TEX_MODE_COUNT * 3 * 2] = { - (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, - }; - - // Function index. - int func_idx = p.filterMode; - if (cube_mode) - func_idx += TEX_MODE_COUNT; - func_idx = func_idx * 3 + channel_div_idx; - - // Launch kernel. - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(func_tbl[func_idx], gridSize, blockSize, args, 0, stream)); - } -}; - -REGISTER_OP("TextureFwd") - .Input ("tex: float") - .Input ("uv: float") - .Output ("out: float") - .Attr ("filter_mode: int") - .Attr ("boundary_mode: int"); - -REGISTER_OP("TextureFwdMip") - .Input ("tex: float") - .Input ("uv: float") - .Input ("uv_da: float") - .Output ("out: float") - .Output ("mip: float") - .Attr ("filter_mode: int") - .Attr ("boundary_mode: int") - .Attr ("tex_const: int") - .Attr ("max_mip_level: int"); - -REGISTER_KERNEL_BUILDER(Name("TextureFwd") .Device(DEVICE_GPU), TextureFwdOp); -REGISTER_KERNEL_BUILDER(Name("TextureFwdMip").Device(DEVICE_GPU), TextureFwdOp); - -//------------------------------------------------------------------------ -// Gradient TensorFlow op. - -struct TextureGradOp : public OpKernel -{ - TextureKernelParams m_attribs; - - TextureGradOp(OpKernelConstruction* ctx): OpKernel(ctx) - { - memset(&m_attribs, 0, sizeof(m_attribs)); - parseOpAttributes(ctx, m_attribs); - } - - void Compute(OpKernelContext* ctx) - { - TextureKernelParams& p = m_attribs; - cudaStream_t stream = ctx->eigen_device().stream(); - bool cube_mode = (p.boundaryMode == TEX_BOUNDARY_MODE_CUBE); - - // Get input. - const Tensor& tex = ctx->input(0); - const Tensor& uv = ctx->input(1); - const Tensor& dy = ctx->input(2); - const Tensor& uv_da = ctx->input(p.enableMip ? 3 : 2); - const Tensor& mip = ctx->input(p.enableMip ? 4 : 2); - - // Extract input dimensions. - p.n = (uv.dims() > 0) ? uv.dim_size(0) : 0; - p.imgHeight = (uv.dims() > 1) ? uv.dim_size(1) : 0; - p.imgWidth = (uv.dims() > 2) ? uv.dim_size(2) : 0; - p.texDepth = (tex.dims() > 0) ? tex.dim_size(0) : 0; - if (!cube_mode) - { - p.texHeight = (tex.dims() > 1) ? tex.dim_size(1) : 0; - p.texWidth = (tex.dims() > 2) ? tex.dim_size(2) : 0; - p.channels = (tex.dims() > 3) ? tex.dim_size(3) : 0; - } - else - { - p.texHeight = (tex.dims() > 2) ? tex.dim_size(2) : 0; - p.texWidth = (tex.dims() > 3) ? tex.dim_size(3) : 0; - p.channels = (tex.dims() > 4) ? tex.dim_size(4) : 0; - } - - // Sanity checks. - if (!cube_mode) - { - OP_REQUIRES(ctx, tex.dims() == 4 && tex.dim_size(0) > 0 && tex.dim_size(1) > 0 && tex.dim_size(2) > 0 && tex.dim_size(3) > 0, errors::InvalidArgument("tex must have shape[>0, >0, >0, >0]")); - OP_REQUIRES(ctx, uv.dims() == 4 && uv.dim_size(0) > 0 && uv.dim_size(1) > 0 && uv.dim_size(2) > 0 && uv.dim_size(3) == 2, errors::InvalidArgument("uv must have shape [>0, >0, >0, 2]")); - } - else - { - OP_REQUIRES(ctx, tex.dims() == 5 && tex.dim_size(0) > 0 && tex.dim_size(1) == 6 && tex.dim_size(2) > 0 && tex.dim_size(3) > 0 && tex.dim_size(4) > 0, errors::InvalidArgument("tex must have shape[>0, 6, >0, >0, >0] in cube map mode")); - OP_REQUIRES(ctx, uv.dims() == 4 && uv.dim_size(0) > 0 && uv.dim_size(1) > 0 && uv.dim_size(2) > 0 && uv.dim_size(3) == 3, errors::InvalidArgument("uv must have shape [>0, >0, >0, 3] in cube map mode")); - OP_REQUIRES(ctx, tex.dim_size(2) == tex.dim_size(3), errors::InvalidArgument("texture shape must be square in cube map mode")); - } - OP_REQUIRES(ctx, tex.dim_size(0) == 1 || tex.dim_size(0) == p.n, errors::InvalidArgument("minibatch size mismatch between inputs tex, uv")); - OP_REQUIRES(ctx, dy.dims() == 4 && dy.dim_size(0) == p.n && dy.dim_size(1) == p.imgHeight && dy.dim_size(2) == p.imgWidth && dy.dim_size(3) == p.channels, errors::InvalidArgument("dy must have shape [minibatch_size, height, width, channels]")); - if (p.enableMip) - { - if (!cube_mode) - OP_REQUIRES(ctx, uv_da.dims() == 4 && uv_da.dim_size(0) == p.n && uv_da.dim_size(1) == p.imgHeight && uv_da.dim_size(2) == p.imgWidth && uv_da.dim_size(3) == 4, errors::InvalidArgument("uv_da must have shape [minibatch_size, height, width, 4]")); - else - OP_REQUIRES(ctx, uv_da.dims() == 4 && uv_da.dim_size(0) == p.n && uv_da.dim_size(1) == p.imgHeight && uv_da.dim_size(2) == p.imgWidth && uv_da.dim_size(3) == 6, errors::InvalidArgument("uv_da must have shape [minibatch_size, height, width, 6] in cube map mode")); - } - - // Get input pointers. - p.tex[0] = tex.flat().data(); - p.uv = uv.flat().data(); - p.dy = dy.flat().data(); - p.uvDA = p.enableMip ? uv_da.flat().data() : 0; - float* pmip = p.enableMip ? (float*)mip.flat().data() : 0; - - // Allocate output tensor for tex gradient. - Tensor* grad_tex_tensor = NULL; - TensorShape grad_tex_shape; - grad_tex_shape.AddDim(p.texDepth); - if (cube_mode) - grad_tex_shape.AddDim(6); - grad_tex_shape.AddDim(p.texHeight); - grad_tex_shape.AddDim(p.texWidth); - grad_tex_shape.AddDim(p.channels); - OP_REQUIRES_OK(ctx, ctx->allocate_output(0, grad_tex_shape, &grad_tex_tensor)); - p.gradTex[0] = grad_tex_tensor->flat().data(); - - // Allocate output tensor for uv gradient. - if (p.filterMode != TEX_MODE_NEAREST) - { - TensorShape grad_uv_shape; - Tensor* grad_uv_tensor = NULL; - grad_uv_shape.AddDim(p.n); - grad_uv_shape.AddDim(p.imgHeight); - grad_uv_shape.AddDim(p.imgWidth); - grad_uv_shape.AddDim(uv.dim_size(3)); - OP_REQUIRES_OK(ctx, ctx->allocate_output(1, grad_uv_shape, &grad_uv_tensor)); - p.gradUV = grad_uv_tensor->flat().data(); - - // Allocate output tensor for uv_da gradient. - if (p.filterMode == TEX_MODE_LINEAR_MIPMAP_LINEAR) - { - Tensor* grad_uv_da_tensor = NULL; - grad_uv_shape.set_dim(3, uv_da.dim_size(3)); - OP_REQUIRES_OK(ctx, ctx->allocate_output(2, grad_uv_shape, &grad_uv_da_tensor)); - p.gradUVDA = grad_uv_da_tensor->flat().data(); - } - } - - // 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. - Tensor grad_mip_tensor; - float* pgradMip = 0; - if (p.enableMip) - { - // Generate mip offsets. - int mipOffsets[TEX_MAX_MIP_LEVEL]; - int mipTotal = calculateMipInfo(ctx, p, mipOffsets); - - // Get space for temporary mip gradients. - TensorShape grad_mip_shape; - grad_mip_shape.AddDim(mipTotal); - ctx->allocate_temp(DT_FLOAT, grad_mip_shape, &grad_mip_tensor); - pgradMip = grad_mip_tensor.flat().data(); - 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. - } - - // Clear mip gradients. - OP_CHECK_CUDA_ERROR(ctx, cudaMemsetAsync(pgradMip, 0, mipTotal * sizeof(float), stream)); - } - - // Initialize texture gradients to zero. - int texBytes = p.texHeight * p.texWidth * p.texDepth * p.channels * sizeof(float); - if (cube_mode) - texBytes *= 6; - OP_CHECK_CUDA_ERROR(ctx, cudaMemsetAsync(p.gradTex[0], 0, texBytes, stream)); - - // Verify that buffers are aligned to allow float2/float4 operations. Unused pointers are zero so always aligned. - if (!cube_mode) - { - OP_REQUIRES(ctx, !((uintptr_t)p.uv & 7), errors::Internal("uv input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradUV & 7), errors::Internal("grad_uv output tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.uvDA & 15), errors::Internal("uv_da input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradUVDA & 15), errors::Internal("grad_uv_da output tensor not aligned to float4")); - } - else - { - OP_REQUIRES(ctx, !((uintptr_t)p.uvDA & 7), errors::Internal("uv_da input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradUVDA & 7), errors::Internal("grad_uv_da output tensor not aligned to float2")); - } - if ((p.channels & 3) == 0) - { - OP_REQUIRES(ctx, !((uintptr_t)p.tex[0] & 15), errors::Internal("tex input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradTex[0] & 15), errors::Internal("grad_tex output tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)p.dy & 15), errors::Internal("dy input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)pmip & 15), errors::Internal("mip input tensor not aligned to float4")); - OP_REQUIRES(ctx, !((uintptr_t)pgradMip & 15), errors::Internal("internal mip gradient tensor not aligned to float4")); - } - if ((p.channels & 1) == 0) - { - OP_REQUIRES(ctx, !((uintptr_t)p.tex[0] & 7), errors::Internal("tex input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.gradTex[0] & 7), errors::Internal("grad_tex output tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)p.dy & 7), errors::Internal("dy output tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)pmip & 7), errors::Internal("mip input tensor not aligned to float2")); - OP_REQUIRES(ctx, !((uintptr_t)pgradMip & 7), errors::Internal("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] = { - (void*)TextureGradKernelNearest, - (void*)TextureGradKernelLinear, - (void*)TextureGradKernelLinearMipmapNearest, - (void*)TextureGradKernelLinearMipmapLinear, - (void*)TextureGradKernelCubeNearest, - (void*)TextureGradKernelCubeLinear, - (void*)TextureGradKernelCubeLinearMipmapNearest, - (void*)TextureGradKernelCubeLinearMipmapLinear, - }; - - // Function index. - int func_idx = p.filterMode; - if (cube_mode) - func_idx += TEX_MODE_COUNT; - - // Launch main gradient kernel. - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(func_tbl[func_idx], gridSize, blockSize, args, 0, stream)); - - // Launch kernel to pull gradients from mip levels. - if (p.enableMip) - { - 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 }; - OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel(mip_grad_func_tbl[channel_div_idx], gridSize, blockSize, args, sharedBytes, stream)); - } - } -}; - -REGISTER_OP("TextureGradNearest") - .Input ("tex: float") - .Input ("uv: float") - .Input ("dy: float") - .Output ("grad_tex: float") - .Attr ("filter_mode: int") - .Attr ("boundary_mode: int"); - -REGISTER_OP("TextureGradLinear") - .Input ("tex: float") - .Input ("uv: float") - .Input ("dy: float") - .Output ("grad_tex: float") - .Output ("grad_uv: float") - .Attr ("filter_mode: int") - .Attr ("boundary_mode: int"); - -REGISTER_OP("TextureGradLinearMipmapNearest") - .Input ("tex: float") - .Input ("uv: float") - .Input ("dy: float") - .Input ("uv_da: float") - .Input ("mip: float") - .Output ("grad_tex: float") - .Output ("grad_uv: float") - .Attr ("filter_mode: int") - .Attr ("boundary_mode: int") - .Attr ("max_mip_level: int"); - -REGISTER_OP("TextureGradLinearMipmapLinear") - .Input ("tex: float") - .Input ("uv: float") - .Input ("dy: float") - .Input ("uv_da: float") - .Input ("mip: float") - .Output ("grad_tex: float") - .Output ("grad_uv: float") - .Output ("grad_uv_da: float") - .Attr ("filter_mode: int") - .Attr ("boundary_mode: int") - .Attr ("max_mip_level: int"); - -REGISTER_KERNEL_BUILDER(Name("TextureGradNearest") .Device(DEVICE_GPU), TextureGradOp); -REGISTER_KERNEL_BUILDER(Name("TextureGradLinear") .Device(DEVICE_GPU), TextureGradOp); -REGISTER_KERNEL_BUILDER(Name("TextureGradLinearMipmapNearest").Device(DEVICE_GPU), TextureGradOp); -REGISTER_KERNEL_BUILDER(Name("TextureGradLinearMipmapLinear") .Device(DEVICE_GPU), TextureGradOp); - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/__init__.py b/extensions/nvdiffrast/nvdiffrast/torch/__init__.py deleted file mode 100644 index c12ddfffee1da1f71a5cb34bbb79455f180263c5..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# 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, RasterizeGLContext, get_log_level, set_log_level, rasterize, DepthPeeler, interpolate, texture, texture_construct_mip, antialias, antialias_construct_topology_hash -__all__ = ["RasterizeCudaContext", "RasterizeGLContext", "get_log_level", "set_log_level", "rasterize", "DepthPeeler", "interpolate", "texture", "texture_construct_mip", "antialias", "antialias_construct_topology_hash"] diff --git a/extensions/nvdiffrast/nvdiffrast/torch/ops.py b/extensions/nvdiffrast/nvdiffrast/torch/ops.py deleted file mode 100644 index 5514bbb5254e5cd99f7ae3870069c7e4e81ec2dd..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/ops.py +++ /dev/null @@ -1,734 +0,0 @@ -# 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 importlib -import logging -import numpy as np -import os -import torch -import torch.utils.cpp_extension -from . import _C - -#---------------------------------------------------------------------------- -# C++/Cuda plugin compiler/loader. - -_cached_plugin = {} -def _get_plugin(gl=False): - assert isinstance(gl, bool) - - # Modified with precompiled torch CUDA extension - if not gl: - return _C - - # Return cached plugin if already loaded. - if _cached_plugin.get(gl, None) is not None: - return _cached_plugin[gl] - - # Make sure we can find the necessary compiler and libary binaries. - if os.name == 'nt': - lib_dir = os.path.dirname(__file__) + r"\..\lib" - def find_cl_path(): - import glob - def get_sort_key(x): - # Primary criterion is VS version, secondary is edition, third is internal MSVC version. - x = x.split('\\')[3:] - x[1] = {'BuildTools': '~0', 'Community': '~1', 'Pro': '~2', 'Professional': '~3', 'Enterprise': '~4'}.get(x[1], x[1]) - return x - vs_relative_path = r"\Microsoft Visual Studio\*\*\VC\Tools\MSVC\*\bin\Hostx64\x64" - paths = glob.glob(r"C:\Program Files" + vs_relative_path) - paths += glob.glob(r"C:\Program Files (x86)" + vs_relative_path) - if paths: - return sorted(paths, key=get_sort_key)[-1] - - # If cl.exe is not on path, try to find it. - if os.system("where cl.exe >nul 2>nul") != 0: - cl_path = find_cl_path() - if cl_path is None: - raise RuntimeError("Could not locate a supported Microsoft Visual C++ installation") - os.environ['PATH'] += ';' + cl_path - - # Compiler options. - common_opts = ['-DNVDR_TORCH'] - cc_opts = [] - if os.name == 'nt': - cc_opts += ['/wd4067', '/wd4624'] # Disable warnings in torch headers. - - # Linker options for the GL-interfacing plugin. - ldflags = [] - if gl: - if os.name == 'posix': - ldflags = ['-lGL', '-lEGL'] - elif os.name == 'nt': - libs = ['gdi32', 'opengl32', 'user32', 'setgpu'] - ldflags = ['/LIBPATH:' + lib_dir] + ['/DEFAULTLIB:' + x for x in libs] - - # List of source files. - if gl: - source_files = [ - '../common/common.cpp', - '../common/glutil.cpp', - '../common/rasterize_gl.cpp', - 'torch_bindings_gl.cpp', - 'torch_rasterize_gl.cpp', - ] - else: - source_files = [ - '../common/cudaraster/impl/Buffer.cpp', - '../common/cudaraster/impl/CudaRaster.cpp', - '../common/cudaraster/impl/RasterImpl.cu', - '../common/cudaraster/impl/RasterImpl.cpp', - '../common/common.cpp', - '../common/rasterize.cu', - '../common/interpolate.cu', - '../common/texture.cu', - '../common/texture.cpp', - '../common/antialias.cu', - 'torch_bindings.cpp', - 'torch_rasterize.cpp', - 'torch_interpolate.cpp', - 'torch_texture.cpp', - 'torch_antialias.cpp', - ] - - # Some containers set this to contain old architectures that won't compile. We only need the one installed in the machine. - os.environ['TORCH_CUDA_ARCH_LIST'] = '' - - # On Linux, show a warning if GLEW is being forcibly loaded when compiling the GL plugin. - if gl and (os.name == 'posix') and ('libGLEW' in os.environ.get('LD_PRELOAD', '')): - logging.getLogger('nvdiffrast').warning("Warning: libGLEW is being loaded via LD_PRELOAD, and will probably conflict with the OpenGL plugin") - - # Try to detect if a stray lock file is left in cache directory and show a warning. This sometimes happens on Windows if the build is interrupted at just the right moment. - plugin_name = 'nvdiffrast_plugin' + ('_gl' if gl else '') - try: - lock_fn = os.path.join(torch.utils.cpp_extension._get_build_directory(plugin_name, False), 'lock') - if os.path.exists(lock_fn): - logging.getLogger('nvdiffrast').warning("Lock file exists in build directory: '%s'" % lock_fn) - except: - pass - - # Speed up compilation on Windows. - if os.name == 'nt': - # Skip telemetry sending step in vcvarsall.bat - os.environ['VSCMD_SKIP_SENDTELEMETRY'] = '1' - - # Opportunistically patch distutils to cache MSVC environments. - try: - import distutils._msvccompiler - import functools - if not hasattr(distutils._msvccompiler._get_vc_env, '__wrapped__'): - distutils._msvccompiler._get_vc_env = functools.lru_cache()(distutils._msvccompiler._get_vc_env) - except: - pass - - # Compile and load. - source_paths = [os.path.join(os.path.dirname(__file__), fn) for fn in source_files] - torch.utils.cpp_extension.load(name=plugin_name, sources=source_paths, extra_cflags=common_opts+cc_opts, extra_cuda_cflags=common_opts+['-lineinfo'], extra_ldflags=ldflags, with_cuda=True, verbose=False) - - # Import, cache, and return the compiled module. - _cached_plugin[gl] = importlib.import_module(plugin_name) - return _cached_plugin[gl] - -#---------------------------------------------------------------------------- -# Log level. -#---------------------------------------------------------------------------- - -def get_log_level(): - '''Get current log level. - - Returns: - Current log level in nvdiffrast. See `set_log_level()` for possible values. - ''' - return _get_plugin().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. - ''' - _get_plugin().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 = _get_plugin().RasterizeCRStateWrapper(cuda_device_idx) - self.output_db = True - self.active_depth_peeler = None - -#---------------------------------------------------------------------------- -# GL state wrapper. -#---------------------------------------------------------------------------- - -class RasterizeGLContext: - def __init__(self, output_db=True, mode='automatic', device=None): - '''Create a new OpenGL rasterizer context. - - Creating an OpenGL context is a slow operation so you should usually reuse the same - context in all calls to `rasterize()` on the same CPU thread. The OpenGL context - is deleted when the object is destroyed. - - Side note: When using the OpenGL context in a rasterization operation, the - context's internal framebuffer object is automatically enlarged to accommodate the - rasterization operation's output shape, but it is never shrunk in size until the - context is destroyed. Thus, if you need to rasterize, say, deep low-resolution - tensors and also shallow high-resolution tensors, you can conserve GPU memory by - creating two separate OpenGL contexts for these tasks. In this scenario, using the - same OpenGL context for both tasks would end up reserving GPU memory for a deep, - high-resolution output tensor. - - Args: - output_db (bool): Compute and output image-space derivates of barycentrics. - mode: OpenGL context handling mode. Valid values are 'manual' and 'automatic'. - 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 OpenGL rasterizer context. - ''' - assert output_db is True or output_db is False - assert mode in ['automatic', 'manual'] - self.output_db = output_db - self.mode = mode - 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 = _get_plugin(gl=True).RasterizeGLStateWrapper(output_db, mode == 'automatic', cuda_device_idx) - self.active_depth_peeler = None # For error checking only. - - def set_context(self): - '''Set (activate) OpenGL context in the current CPU thread. - Only available if context was created in manual mode. - ''' - assert self.mode == 'manual' - self.cpp_wrapper.set_context() - - def release_context(self): - '''Release (deactivate) currently active OpenGL context. - Only available if context was created in manual mode. - ''' - assert self.mode == 'manual' - self.cpp_wrapper.release_context() - -#---------------------------------------------------------------------------- -# Rasterize. -#---------------------------------------------------------------------------- - -class _rasterize_func(torch.autograd.Function): - @staticmethod - def forward(ctx, raster_ctx, pos, tri, resolution, ranges, grad_db, peeling_idx): - if isinstance(raster_ctx, RasterizeGLContext): - out, out_db = _get_plugin(gl=True).rasterize_fwd_gl(raster_ctx.cpp_wrapper, pos, tri, resolution, ranges, peeling_idx) - else: - out, out_db = _get_plugin().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 = _get_plugin().rasterize_grad_db(pos, tri, out, dy, ddb) - else: - g_pos = _get_plugin().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 `RasterizeGLContext` or `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. Ignored if using an OpenGL context that - was not configured to output image-space derivatives. - - Returns: - A tuple of two tensors. The first output tensor has shape [minibatch_size, - height, width, 4] and contains the main rasterizer output in order (u, v, z/w, - triangle_id). If the OpenGL context was configured to output image-space - derivatives of barycentrics, the second output tensor will also have shape - [minibatch_size, height, width, 4] and contain said derivatives in order - (du/dX, du/dY, dv/dX, dv/dY). Otherwise it will be an empty tensor with shape - [minibatch_size, height, width, 0]. - ''' - assert isinstance(glctx, (RasterizeGLContext, RasterizeCudaContext)) - assert grad_db is True or grad_db is False - grad_db = grad_db and glctx.output_db - - # 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, (RasterizeGLContext, RasterizeCudaContext)) - assert grad_db is True or grad_db is False - grad_db = grad_db and glctx.output_db - - # 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 = _get_plugin().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 = _get_plugin().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 = _get_plugin().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 = _get_plugin().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 = _get_plugin().TextureMipWrapper() - out = _get_plugin().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 = _get_plugin().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 = _get_plugin().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 = _get_plugin().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 = _get_plugin().texture_grad_linear(tex, uv, dy, filter_mode_enum, boundary_mode_enum) - return None, g_tex, g_uv, None, None - else: # nearest - g_tex = _get_plugin().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, (_get_plugin().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 = _get_plugin().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 _get_plugin().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 = _get_plugin().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 = _get_plugin().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, _get_plugin().TopologyHashWrapper) - else: - topology_hash = _get_plugin().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 _get_plugin().antialias_construct_topology_hash(tri) - -#---------------------------------------------------------------------------- diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_antialias.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_antialias.cpp deleted file mode 100644 index 9e52f957daa2d7bffe6e6ec52e1435bbf1ea2f55..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_antialias.cpp +++ /dev/null @@ -1,243 +0,0 @@ -// 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(); - - // 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()); - - // 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 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(); - p.rasterOut = rast.data_ptr(); - p.tri = tri.data_ptr(); - p.pos = pos.data_ptr(); - p.evHash = (uint4*)(topology_hash.data_ptr()); - - // 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(); - p.workBuffer = (int4*)(work_buffer.data_ptr()); - - // 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(out, work_buffer); -} - -//------------------------------------------------------------------------ -// Gradient op. - -std::tuple 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(); - p.rasterOut = rast.data_ptr(); - p.tri = tri.data_ptr(); - p.pos = pos.data_ptr(); - p.dy = dy_.data_ptr(); - p.workBuffer = (int4*)(work_buffer.data_ptr()); - - // 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(); - p.gradPos = grad_pos.data_ptr(); - - // 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(grad_color, grad_pos); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_bindings.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_bindings.cpp deleted file mode 100644 index 431b68c629504510e474d92909395c4d3d06c34c..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_bindings.cpp +++ /dev/null @@ -1,73 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// Op prototypes. Return type macros for readability. - -#define OP_RETURN_T torch::Tensor -#define OP_RETURN_TT std::tuple -#define OP_RETURN_TTT std::tuple -#define OP_RETURN_TTTT std::tuple -#define OP_RETURN_TTV std::tuple > -#define OP_RETURN_TTTTV std::tuple > - -OP_RETURN_TT rasterize_fwd_cuda (RasterizeCRStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple 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& 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& 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 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 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 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_(m, "RasterizeCRStateWrapper").def(pybind11::init()); - pybind11::class_(m, "TextureMipWrapper").def(pybind11::init<>()); - pybind11::class_(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"); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_bindings_gl.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_bindings_gl.cpp deleted file mode 100644 index f403d73cf3149531ccedd277e6b8b26afa5cdda4..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_bindings_gl.cpp +++ /dev/null @@ -1,30 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// Op prototypes. - -std::tuple rasterize_fwd_gl(RasterizeGLStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple resolution, torch::Tensor ranges, int peeling_idx); - -//------------------------------------------------------------------------ - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - // State classes. - pybind11::class_(m, "RasterizeGLStateWrapper").def(pybind11::init()) - .def("set_context", &RasterizeGLStateWrapper::setContext) - .def("release_context", &RasterizeGLStateWrapper::releaseContext); - - // Ops. - m.def("rasterize_fwd_gl", &rasterize_fwd_gl, "rasterize forward op (opengl)"); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_common.inl b/extensions/nvdiffrast/nvdiffrast/torch/torch_common.inl deleted file mode 100644 index 3809ecef8b091fa911d1fdf89c5b84dc1038e9a1..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_common.inl +++ /dev/null @@ -1,29 +0,0 @@ -// 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/framework.h" - -//------------------------------------------------------------------------ -// Input check helpers. -//------------------------------------------------------------------------ - -#ifdef _MSC_VER -#define __func__ __FUNCTION__ -#endif - -#define NVDR_CHECK_DEVICE(...) do { TORCH_CHECK(at::cuda::check_device({__VA_ARGS__}), __func__, "(): Inputs " #__VA_ARGS__ " must reside on the same GPU device") } while(0) -#define NVDR_CHECK_CPU(...) do { nvdr_check_cpu({__VA_ARGS__}, __func__, "(): Inputs " #__VA_ARGS__ " must reside on CPU"); } while(0) -#define NVDR_CHECK_CONTIGUOUS(...) do { nvdr_check_contiguous({__VA_ARGS__}, __func__, "(): Inputs " #__VA_ARGS__ " must be contiguous tensors"); } while(0) -#define NVDR_CHECK_F32(...) do { nvdr_check_f32({__VA_ARGS__}, __func__, "(): Inputs " #__VA_ARGS__ " must be float32 tensors"); } while(0) -#define NVDR_CHECK_I32(...) do { nvdr_check_i32({__VA_ARGS__}, __func__, "(): Inputs " #__VA_ARGS__ " must be int32 tensors"); } while(0) -inline void nvdr_check_cpu(at::ArrayRef ts, const char* func, const char* err_msg) { for (const at::Tensor& t : ts) TORCH_CHECK(t.device().type() == c10::DeviceType::CPU, func, err_msg); } -inline void nvdr_check_contiguous(at::ArrayRef ts, const char* func, const char* err_msg) { for (const at::Tensor& t : ts) TORCH_CHECK(t.is_contiguous(), func, err_msg); } -inline void nvdr_check_f32(at::ArrayRef ts, const char* func, const char* err_msg) { for (const at::Tensor& t : ts) TORCH_CHECK(t.dtype() == torch::kFloat32, func, err_msg); } -inline void nvdr_check_i32(at::ArrayRef ts, const char* func, const char* err_msg) { for (const at::Tensor& t : ts) TORCH_CHECK(t.dtype() == torch::kInt32, func, err_msg); } -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_interpolate.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_interpolate.cpp deleted file mode 100644 index c3f303a2ec9ae02a1db81d326966415efca6568e..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_interpolate.cpp +++ /dev/null @@ -1,250 +0,0 @@ -// 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& 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 interpolate_fwd_da(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor rast_db, bool diff_attrs_all, std::vector& 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(); - p.rast = rast.data_ptr(); - p.tri = tri.data_ptr(); - p.rastDB = enable_da ? rast_db.data_ptr() : 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(); - p.outDA = enable_da ? out_da.data_ptr() : 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(out, out_da); -} - -// Version without derivatives. -std::tuple interpolate_fwd(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri) -{ - std::vector empty_vec; - torch::Tensor empty_tensor; - return interpolate_fwd_da(attr, rast, tri, empty_tensor, false, empty_vec); -} - -//------------------------------------------------------------------------ -// Gradient op. - -std::tuple 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& 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(); - p.rast = rast.data_ptr(); - p.tri = tri.data_ptr(); - p.dy = dy_.data_ptr(); - p.rastDB = enable_da ? rast_db.data_ptr() : NULL; - p.dda = enable_da ? dda_.data_ptr() : 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(); - p.gradRaster = gradRaster.data_ptr(); - p.gradRasterDB = enable_da ? gradRasterDB.data_ptr() : 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(gradAttr, gradRaster, gradRasterDB); -} - -// Version without derivatives. -std::tuple interpolate_grad(torch::Tensor attr, torch::Tensor rast, torch::Tensor tri, torch::Tensor dy) -{ - std::vector empty_vec; - torch::Tensor empty_tensor; - std::tuple result = interpolate_grad_da(attr, rast, tri, dy, empty_tensor, empty_tensor, false, empty_vec); - return std::tuple(std::get<0>(result), std::get<1>(result)); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_rasterize.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_rasterize.cpp deleted file mode 100644 index 6f2ef697aa9bd409c82674bd2c1d921140a691d7..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_rasterize.cpp +++ /dev/null @@ -1,265 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// 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 rasterize_fwd_cuda(RasterizeCRStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple 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(); - const int32_t* rangesPtr = instance_mode ? 0 : ranges.data_ptr(); // This is in CPU memory. - const int32_t* triPtr = tri.data_ptr(); - 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(); - p.out_db = out_db.data_ptr(); - 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(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(); - p.tri = tri.data_ptr(); - p.out = out.data_ptr(); - p.dy = dy_.data_ptr(); - p.ddb = enable_db ? ddb_.data_ptr() : 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(); - - // 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); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_rasterize_gl.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_rasterize_gl.cpp deleted file mode 100644 index 76f875d6adaa273107ec7e44f7724cf59e90de48..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_rasterize_gl.cpp +++ /dev/null @@ -1,132 +0,0 @@ -// 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_gl.h" -#include - -//------------------------------------------------------------------------ -// Python GL state wrapper methods. - -RasterizeGLStateWrapper::RasterizeGLStateWrapper(bool enableDB, bool automatic_, int cudaDeviceIdx_) -{ - pState = new RasterizeGLState(); - automatic = automatic_; - cudaDeviceIdx = cudaDeviceIdx_; - memset(pState, 0, sizeof(RasterizeGLState)); - pState->enableDB = enableDB ? 1 : 0; - rasterizeInitGLContext(NVDR_CTX_PARAMS, *pState, cudaDeviceIdx_); - releaseGLContext(); -} - -RasterizeGLStateWrapper::~RasterizeGLStateWrapper(void) -{ - setGLContext(pState->glctx); - rasterizeReleaseBuffers(NVDR_CTX_PARAMS, *pState); - releaseGLContext(); - destroyGLContext(pState->glctx); - delete pState; -} - -void RasterizeGLStateWrapper::setContext(void) -{ - setGLContext(pState->glctx); -} - -void RasterizeGLStateWrapper::releaseContext(void) -{ - releaseGLContext(); -} - -//------------------------------------------------------------------------ -// Forward op (OpenGL). - -std::tuple rasterize_fwd_gl(RasterizeGLStateWrapper& stateWrapper, torch::Tensor pos, torch::Tensor tri, std::tuple resolution, torch::Tensor ranges, int peeling_idx) -{ - const at::cuda::OptionalCUDAGuard device_guard(device_of(pos)); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - RasterizeGLState& s = *stateWrapper.pState; - - // 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 GL context was created for the correct GPU. - NVDR_CHECK(pos.get_device() == stateWrapper.cudaDeviceIdx, "GL context must must reside on the same device as input tensors"); - - // Determine number of outputs - int num_outputs = s.enableDB ? 2 : 1; - - // 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 = std::get<0>(resolution); - int width = std::get<1>(resolution); - int depth = instance_mode ? pos.size(0) : ranges.size(0); - NVDR_CHECK(height > 0 && width > 0, "resolution must be [>0, >0]"); - - // Get position and triangle buffer sizes in int32/float32. - int posCount = 4 * pos.size(0) * (instance_mode ? pos.size(1) : 1); - int triCount = 3 * tri.size(0); - - // Set the GL context unless manual context. - if (stateWrapper.automatic) - setGLContext(s.glctx); - - // Resize all buffers. - bool changes = false; - rasterizeResizeBuffers(NVDR_CTX_PARAMS, s, changes, posCount, triCount, width, height, depth); - if (changes) - { -#ifdef _WIN32 - // Workaround for occasional blank first frame on Windows. - releaseGLContext(); - setGLContext(s.glctx); -#endif - } - - // Copy input data to GL and render. - const float* posPtr = pos.data_ptr(); - const int32_t* rangesPtr = instance_mode ? 0 : ranges.data_ptr(); // This is in CPU memory. - const int32_t* triPtr = tri.data_ptr(); - int vtxPerInstance = instance_mode ? pos.size(1) : 0; - rasterizeRender(NVDR_CTX_PARAMS, s, stream, posPtr, posCount, vtxPerInstance, triPtr, triCount, rangesPtr, width, height, depth, peeling_idx); - - // Allocate output tensors. - torch::TensorOptions opts = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - torch::Tensor out = torch::empty({depth, height, width, 4}, opts); - torch::Tensor out_db = torch::empty({depth, height, width, s.enableDB ? 4 : 0}, opts); - float* outputPtr[2]; - outputPtr[0] = out.data_ptr(); - outputPtr[1] = s.enableDB ? out_db.data_ptr() : NULL; - - // Copy rasterized results into CUDA buffers. - rasterizeCopyResults(NVDR_CTX_PARAMS, s, stream, outputPtr, width, height, depth); - - // Done. Release GL context and return. - if (stateWrapper.automatic) - releaseGLContext(); - - return std::tuple(out, out_db); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_texture.cpp b/extensions/nvdiffrast/nvdiffrast/torch/torch_texture.cpp deleted file mode 100644 index 6e3693eedd9e17d862d74ca08ee9636dcca290cc..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_texture.cpp +++ /dev/null @@ -1,718 +0,0 @@ -// 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 - -//------------------------------------------------------------------------ -// 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(); - - // Generate mip offsets and calculate total size. - int mipOffsets[TEX_MAX_MIP_LEVEL]; - int mipTotal = calculateMipInfo(NVDR_CTX_PARAMS, 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(); - 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 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(); - p.uv = uv.data_ptr(); - p.uvDA = (p.enableMip && has_uv_da) ? uv_da.data_ptr() : NULL; - p.mipLevelBias = (p.enableMip && has_mip_level_bias) ? mip_level_bias.data_ptr() : 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(); - - // 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(); - } - } - else - { - // Generate mip offsets, check mipmap size, and set mip data pointer. - int mipOffsets[TEX_MAX_MIP_LEVEL]; - int mipTotal = calculateMipInfo(NVDR_CTX_PARAMS, 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(); - 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 empty_vector; - return texture_fwd_mip(tex, uv, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode); -} - -//------------------------------------------------------------------------ -// Gradient op. - -std::tuple > 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 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(); - p.uv = uv.data_ptr(); - p.dy = dy_.data_ptr(); - p.uvDA = (p.enableMip && has_uv_da) ? uv_da.data_ptr() : NULL; - p.mipLevelBias = (p.enableMip && has_mip_level_bias) ? mip_level_bias.data_ptr() : NULL; - - // Allocate output tensor for tex gradient. - torch::Tensor grad_tex = torch::zeros_like(tex); - p.gradTex[0] = grad_tex.data_ptr(); - - // 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(); - - // 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(); - } - - // 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(); - } - } - } - - // 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 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(); - p.gradTex[i] = g.data_ptr(); - } - } - else - { - // Generate mip offsets and get space for temporary mip gradients. - int mipOffsets[TEX_MAX_MIP_LEVEL]; - int mipTotal = calculateMipInfo(NVDR_CTX_PARAMS, 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(); - pgradMip = grad_mip.data_ptr(); - 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 >(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 empty_vector; - std::tuple > 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 texture_grad_linear(torch::Tensor tex, torch::Tensor uv, torch::Tensor dy, int filter_mode, int boundary_mode) -{ - torch::Tensor empty_tensor; - std::vector empty_vector; - std::tuple > result = texture_grad_linear_mipmap_linear(tex, uv, dy, empty_tensor, empty_tensor, TextureMipWrapper(), empty_vector, filter_mode, boundary_mode); - return std::tuple(std::get<0>(result), std::get<1>(result)); -} - -// Version for linear-mipmap-nearest mode. -std::tuple > 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 mip_stack, int filter_mode, int boundary_mode) -{ - std::tuple > 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 >(std::get<0>(result), std::get<1>(result), std::get<4>(result)); -} - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/nvdiffrast/torch/torch_types.h b/extensions/nvdiffrast/nvdiffrast/torch/torch_types.h deleted file mode 100644 index 9ea4bdea7c011866cd80b72d7f0a593ee091d070..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/nvdiffrast/torch/torch_types.h +++ /dev/null @@ -1,65 +0,0 @@ -// 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 GL state wrapper. - -class RasterizeGLState; -class RasterizeGLStateWrapper -{ -public: - RasterizeGLStateWrapper (bool enableDB, bool automatic, int cudaDeviceIdx); - ~RasterizeGLStateWrapper (void); - - void setContext (void); - void releaseContext (void); - - RasterizeGLState* pState; - bool automatic; - int cudaDeviceIdx; -}; - -//------------------------------------------------------------------------ -// 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 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; -}; - -//------------------------------------------------------------------------ diff --git a/extensions/nvdiffrast/run_sample.sh b/extensions/nvdiffrast/run_sample.sh deleted file mode 100644 index 08b4104d6e8aef18b0ff8b9491f74586df0cd550..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/run_sample.sh +++ /dev/null @@ -1,52 +0,0 @@ -#!/bin/bash - -# 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. - -function print_help { - echo "Usage: `basename $0` [--build-container] " - echo "" - echo "Option --build-container will build the Docker container based on" - echo "docker/Dockerfile and tag the image with gltorch:latest." - echo "" - echo "Example: `basename $0` samples/torch/envphong.py" -} - -build_container=0 -sample="" -while [[ "$#" -gt 0 ]]; do - case $1 in - --build-container) build_container=1;; - -h|--help) print_help; exit 0 ;; - --*) echo "Unknown parameter passed: $1"; exit 1 ;; - *) sample="$1"; shift; break; - esac - shift -done - -rest=$@ - -# Build the docker container -if [ "$build_container" = "1" ]; then - docker build --tag gltorch:latest -f docker/Dockerfile . -fi - -if [ ! -f "$sample" ]; then - echo - echo "No python sample given or file '$sample' not found. Exiting." - exit 1 -fi - -image="gltorch:latest" - -echo "Using container image: $image" -echo "Running command: $sample $rest" - -# Run a sample with docker -docker run --rm -it --gpus all --user $(id -u):$(id -g) \ - -v `pwd`:/app --workdir /app -e TORCH_EXTENSIONS_DIR=/app/tmp $image python3 $sample $rest diff --git a/extensions/nvdiffrast/setup copy.py b/extensions/nvdiffrast/setup copy.py deleted file mode 100644 index 07541da2a746eb9cd5f6ae61cb34475b890439cd..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/setup copy.py +++ /dev/null @@ -1,51 +0,0 @@ -# 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 nvdiffrast -import setuptools -import os - -with open("README.md", "r") as fh: - long_description = fh.read() - -setuptools.setup( - name="nvdiffrast", - version=nvdiffrast.__version__, - author="Samuli Laine", - author_email="slaine@nvidia.com", - description="nvdiffrast - modular primitives for high-performance differentiable rendering", - long_description=long_description, - long_description_content_type="text/markdown", - url="https://github.com/NVlabs/nvdiffrast", - packages=setuptools.find_packages(), - package_data={ - 'nvdiffrast': [ - 'common/*.h', - 'common/*.inl', - 'common/*.cu', - 'common/*.cpp', - 'common/cudaraster/*.hpp', - 'common/cudaraster/impl/*.cpp', - 'common/cudaraster/impl/*.hpp', - 'common/cudaraster/impl/*.inl', - 'common/cudaraster/impl/*.cu', - 'lib/*.h', - 'torch/*.h', - 'torch/*.inl', - 'torch/*.cpp', - 'tensorflow/*.cu', - ] + (['lib/*.lib'] if os.name == 'nt' else []) - }, - include_package_data=True, - install_requires=['numpy'], # note: can't require torch here as it will install torch even for a TensorFlow container - classifiers=[ - "Programming Language :: Python :: 3", - "Operating System :: OS Independent", - ], - python_requires='>=3.6', -) diff --git a/extensions/nvdiffrast/setup.py b/extensions/nvdiffrast/setup.py deleted file mode 100644 index b335971c46c41edee3352c882c6d23823d6b5ad1..0000000000000000000000000000000000000000 --- a/extensions/nvdiffrast/setup.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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 nvdiffrast -import setuptools -import os -from torch.utils.cpp_extension import CUDAExtension, BuildExtension - - -with open("README.md", "r") as fh: - long_description = fh.read() - -setuptools.setup( - name="nvdiffrast", - version=nvdiffrast.__version__, - author="Samuli Laine", - author_email="slaine@nvidia.com", - description="nvdiffrast - modular primitives for high-performance differentiable rendering", - long_description=long_description, - long_description_content_type="text/markdown", - url="https://github.com/NVlabs/nvdiffrast", - packages=setuptools.find_packages(), - # package_data={ - # 'nvdiffrast': [ - # 'common/*.h', - # 'common/*.inl', - # 'common/*.cu', - # 'common/*.cpp', - # 'common/cudaraster/*.hpp', - # 'common/cudaraster/impl/*.cpp', - # 'common/cudaraster/impl/*.hpp', - # 'common/cudaraster/impl/*.inl', - # 'common/cudaraster/impl/*.cu', - # 'lib/*.h', - # 'torch/*.h', - # 'torch/*.inl', - # 'torch/*.cpp', - # 'tensorflow/*.cu', - # ] + (['lib/*.lib'] if os.name == 'nt' else []) - # }, - # include_package_data=True, - ext_modules=[ - CUDAExtension( - name="nvdiffrast.torch._C", - sources=[ - 'nvdiffrast/common/cudaraster/impl/Buffer.cpp', - 'nvdiffrast/common/cudaraster/impl/CudaRaster.cpp', - 'nvdiffrast/common/cudaraster/impl/RasterImpl_.cu', - 'nvdiffrast/common/cudaraster/impl/RasterImpl.cpp', - 'nvdiffrast/common/common.cpp', - 'nvdiffrast/common/rasterize.cu', - 'nvdiffrast/common/interpolate.cu', - 'nvdiffrast/common/texture_.cu', - 'nvdiffrast/common/texture.cpp', - 'nvdiffrast/common/antialias.cu', - 'nvdiffrast/torch/torch_bindings.cpp', - 'nvdiffrast/torch/torch_rasterize.cpp', - 'nvdiffrast/torch/torch_interpolate.cpp', - 'nvdiffrast/torch/torch_texture.cpp', - 'nvdiffrast/torch/torch_antialias.cpp', - ], - extra_compile_args={ - 'cxx': ['-DNVDR_TORCH'], - 'nvcc': ['-DNVDR_TORCH', '-lineinfo'], - }, - ) - ], - cmdclass={ - 'build_ext': BuildExtension - }, - install_requires=['numpy'], # note: can't require torch here as it will install torch even for a TensorFlow container - classifiers=[ - "Programming Language :: Python :: 3", - "Operating System :: OS Independent", - ], - python_requires='>=3.6', -)