Upload 288 files

.gitattributes

@@ -33,3 +33,38 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/chair_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/chair_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/chair_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_4.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_5.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_6.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_7.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/flower_8.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/monkey_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/monkey_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/monkey_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/monkey_4.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_4.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_5.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_6.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_7.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/paopao_8.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/puppet_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/puppet_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/puppet_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/robot_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/robot_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/SpongeBob_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/SpongeBob_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/SpongeBob_3.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/SpongeBob_4.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/toolcar_1.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/toolcar_2.png filter=lfs diff=lfs merge=lfs -text
+assets/example_multi_image/toolcar_3.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,16 @@
 ---
 title: ReconViaGen
-emoji: 💻
-colorFrom: green
-colorTo: purple
+emoji: 🖥️
+colorFrom: indigo
+colorTo: blue
 sdk: gradio
-sdk_version: 5.44.1
+sdk_version: 5.34.2
 app_file: app.py
 pinned: false
-license: apache-2.0
+license: mit
+short_description: High-fidelity 3D Geometry Generation from single view image
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+Project Page: https://jiahao620.github.io/reconviagen/

app.py

@@ -0,0 +1,392 @@
+import gradio as gr
+import spaces
+from gradio_litmodel3d import LitModel3D
+
+import os
+import shutil
+os.environ['SPCONV_ALGO'] = 'native'
+from typing import *
+import torch
+import numpy as np
+import imageio
+from easydict import EasyDict as edict
+from PIL import Image
+from trellis.pipelines import TrellisVGGTTo3DPipeline
+from trellis.representations import Gaussian, MeshExtractResult
+from trellis.utils import render_utils, postprocessing_utils
+
+
+
+MAX_SEED = np.iinfo(np.int32).max
+# TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+TMP_DIR = "tmp/Trellis-demo"
+os.environ['GRADIO_TEMP_DIR'] = 'tmp'
+os.makedirs(TMP_DIR, exist_ok=True)
+
+def start_session(req: gr.Request):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    os.makedirs(user_dir, exist_ok=True)
+    
+    
+def end_session(req: gr.Request):
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    shutil.rmtree(user_dir)
+def preprocess_image(image: Image.Image) -> Image.Image:
+    """
+    Preprocess the input image for 3D generation.
+    
+    This function is called when a user uploads an image or selects an example.
+    It applies background removal and other preprocessing steps necessary for
+    optimal 3D model generation.
+
+    Args:
+        image (Image.Image): The input image from the user
+
+    Returns:
+        Image.Image: The preprocessed image ready for 3D generation
+    """
+    processed_image = pipeline.preprocess_image(image)
+    return processed_image
+
+
+def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
+    """
+    Preprocess a list of input images for multi-image 3D generation.
+    
+    This function is called when users upload multiple images in the gallery.
+    It processes each image to prepare them for the multi-image 3D generation pipeline.
+    
+    Args:
+        images (List[Tuple[Image.Image, str]]): The input images from the gallery
+        
+    Returns:
+        List[Image.Image]: The preprocessed images ready for 3D generation
+    """
+    images = [image[0] for image in images]
+    processed_images = [pipeline.preprocess_image(image) for image in images]
+    return processed_images
+
+
+def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
+    return {
+        'gaussian': {
+            **gs.init_params,
+            '_xyz': gs._xyz.cpu().numpy(),
+            '_features_dc': gs._features_dc.cpu().numpy(),
+            '_scaling': gs._scaling.cpu().numpy(),
+            '_rotation': gs._rotation.cpu().numpy(),
+            '_opacity': gs._opacity.cpu().numpy(),
+        },
+        'mesh': {
+            'vertices': mesh.vertices.cpu().numpy(),
+            'faces': mesh.faces.cpu().numpy(),
+        },
+    }
+    
+    
+def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
+    gs = Gaussian(
+        aabb=state['gaussian']['aabb'],
+        sh_degree=state['gaussian']['sh_degree'],
+        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
+        scaling_bias=state['gaussian']['scaling_bias'],
+        opacity_bias=state['gaussian']['opacity_bias'],
+        scaling_activation=state['gaussian']['scaling_activation'],
+    )
+    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
+    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
+    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
+    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
+    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+    
+    mesh = edict(
+        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+    )
+    
+    return gs, mesh
+
+
+def get_seed(randomize_seed: bool, seed: int) -> int:
+    """
+    Get the random seed for generation.
+    
+    This function is called by the generate button to determine whether to use
+    a random seed or the user-specified seed value.
+    
+    Args:
+        randomize_seed (bool): Whether to generate a random seed
+        seed (int): The user-specified seed value
+        
+    Returns:
+        int: The seed to use for generation
+    """
+    return np.random.randint(0, MAX_SEED) if randomize_seed else seed
+
+
+@spaces.GPU(duration=120)
+def generate_and_extract_glb(
+    multiimages: List[Tuple[Image.Image, str]],
+    seed: int,
+    ss_guidance_strength: float,
+    ss_sampling_steps: int,
+    slat_guidance_strength: float,
+    slat_sampling_steps: int,
+    multiimage_algo: Literal["multidiffusion", "stochastic"],
+    mesh_simplify: float,
+    texture_size: int,
+    req: gr.Request,
+) -> Tuple[dict, str, str, str]:
+    """
+    Convert an image to a 3D model and extract GLB file.
+
+    Args:
+        image (Image.Image): The input image.
+        multiimages (List[Tuple[Image.Image, str]]): The input images in multi-image mode.
+        is_multiimage (bool): Whether is in multi-image mode.
+        seed (int): The random seed.
+        ss_guidance_strength (float): The guidance strength for sparse structure generation.
+        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
+        slat_guidance_strength (float): The guidance strength for structured latent generation.
+        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
+        multiimage_algo (Literal["multidiffusion", "stochastic"]): The algorithm for multi-image generation.
+        mesh_simplify (float): The mesh simplification factor.
+        texture_size (int): The texture resolution.
+
+    Returns:
+        dict: The information of the generated 3D model.
+        str: The path to the video of the 3D model.
+        str: The path to the extracted GLB file.
+        str: The path to the extracted GLB file (for download).
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    image_files = [image[0] for image in multiimages]
+
+    # Generate 3D model
+    outputs = pipeline.run(
+        image=image_files,
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+        mode=multiimage_algo,
+    )
+
+    # Render video
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    video_path = os.path.join(user_dir, 'sample.mp4')
+    imageio.mimsave(video_path, video, fps=15)
+    
+    # Extract GLB
+    gs = outputs['gaussian'][0]
+    mesh = outputs['mesh'][0]
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+    glb_path = os.path.join(user_dir, 'sample.glb')
+    glb.export(glb_path)
+    
+    # Pack state for optional Gaussian extraction
+    state = pack_state(gs, mesh)
+    
+    torch.cuda.empty_cache()
+    return state, video_path, glb_path, glb_path
+
+
+@spaces.GPU
+def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
+    """
+    Extract a Gaussian splatting file from the generated 3D model.
+    
+    This function is called when the user clicks "Extract Gaussian" button.
+    It converts the 3D model state into a .ply file format containing
+    Gaussian splatting data for advanced 3D applications.
+
+    Args:
+        state (dict): The state of the generated 3D model containing Gaussian data
+        req (gr.Request): Gradio request object for session management
+
+    Returns:
+        Tuple[str, str]: Paths to the extracted Gaussian file (for display and download)
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, _ = unpack_state(state)
+    gaussian_path = os.path.join(user_dir, 'sample.ply')
+    gs.save_ply(gaussian_path)
+    torch.cuda.empty_cache()
+    return gaussian_path, gaussian_path
+
+
+def prepare_multi_example() -> List[Image.Image]:
+    multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
+    images = []
+    for case in multi_case:
+        _images = []
+        for i in range(1, 9):
+            if os.path.exists(f'assets/example_multi_image/{case}_{i}.png'):
+                img = Image.open(f'assets/example_multi_image/{case}_{i}.png')
+                W, H = img.size
+                img = img.resize((int(W / H * 512), 512))
+                _images.append(np.array(img))
+        if len(_images) > 0:
+            images.append(Image.fromarray(np.concatenate(_images, axis=1)))
+    return images
+
+
+def split_image(image: Image.Image) -> List[Image.Image]:
+    """
+    Split a multi-view image into separate view images.
+    
+    This function is called when users select multi-image examples that contain
+    multiple views in a single concatenated image. It automatically splits them
+    based on alpha channel boundaries and preprocesses each view.
+    
+    Args:
+        image (Image.Image): A concatenated image containing multiple views
+        
+    Returns:
+        List[Image.Image]: List of individual preprocessed view images
+    """
+    image = np.array(image)
+    alpha = image[..., 3]
+    alpha = np.any(alpha>0, axis=0)
+    start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist()
+    end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist()
+    images = []
+    for s, e in zip(start_pos, end_pos):
+        images.append(Image.fromarray(image[:, s:e+1]))
+    return [preprocess_image(image) for image in images]
+
+
+with gr.Blocks(delete_cache=(600, 600)) as demo:
+    gr.Markdown("""
+    ## Multi-view images to 3D Asset with [ReconViaGen](https://jiahao620.github.io/reconviagen/)
+    * Upload an image and click "Generate & Extract GLB" to create a 3D asset and automatically extract the GLB file.
+    * If you want the Gaussian file as well, click "Extract Gaussian" after generation.
+    * If the image has alpha channel, it will be used as the mask. Otherwise, we use `rembg` to remove the background.
+    
+    ✨This demo is partial. We will release the whole model later. Stay tuned!✨
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            with gr.Tabs() as input_tabs:
+                with gr.Tab(label="Multiple Images", id=0) as multiimage_input_tab:
+                    image_prompt = gr.Image(label="Image Prompt", format="png", visible=False, image_mode="RGBA", type="pil", height=300)
+                    multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3)
+                    gr.Markdown("""
+                        Input different views of the object in separate images. 
+                        
+                        *NOTE: this is an experimental algorithm without training a specialized model. It may not produce the best results for all images, especially those having different poses or inconsistent details.*
+                    """)
+        
+            with gr.Accordion(label="Generation Settings", open=False):
+                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+                gr.Markdown("Stage 1: Sparse Structure Generation")
+                with gr.Row():
+                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
+                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=30, step=1)
+                gr.Markdown("Stage 2: Structured Latent Generation")
+                with gr.Row():
+                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
+                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
+                multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="multidiffusion")
+            
+            with gr.Accordion(label="GLB Extraction Settings", open=False):
+                mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
+                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
+
+            generate_btn = gr.Button("Generate & Extract GLB", variant="primary")
+            extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+            gr.Markdown("""
+                        *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*
+                        """)
+
+        with gr.Column():
+            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
+            model_output = LitModel3D(label="Extracted GLB/Gaussian", exposure=10.0, height=300)
+            
+            with gr.Row():
+                download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
+                download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
+    
+    output_buf = gr.State()
+
+    # Example images at the bottom of the page
+    with gr.Row() as multiimage_example:
+        examples_multi = gr.Examples(
+            examples=prepare_multi_example(),
+            inputs=[image_prompt],
+            fn=split_image,
+            outputs=[multiimage_prompt],
+            run_on_click=True,
+            examples_per_page=8,
+        )
+
+    # Handlers
+    demo.load(start_session)
+    demo.unload(end_session)
+
+
+    multiimage_prompt.upload(
+        preprocess_images,
+        inputs=[multiimage_prompt],
+        outputs=[multiimage_prompt],
+    )
+
+    generate_btn.click(
+        get_seed,
+        inputs=[randomize_seed, seed],
+        outputs=[seed],
+    ).then(
+        lambda: [None, None, None, None],  # 先清空 video_output
+        inputs=[],
+        outputs=[video_output, model_output, download_glb, download_gs],
+    ).then(
+        generate_and_extract_glb,
+        inputs=[multiimage_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo, mesh_simplify, texture_size],
+        outputs=[output_buf, video_output, model_output, download_glb],
+    ).then(
+        lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
+        outputs=[extract_gs_btn, download_glb],
+    )
+
+    video_output.clear(
+        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False), gr.Button(interactive=False)]),
+        outputs=[extract_gs_btn, download_glb, download_gs],
+    )
+    
+    extract_gs_btn.click(
+        extract_gaussian,
+        inputs=[output_buf],
+        outputs=[model_output, download_gs],
+    ).then(
+        lambda: gr.Button(interactive=True),
+        outputs=[download_gs],
+    )
+
+    model_output.clear(
+        lambda: tuple([gr.Button(interactive=False), gr.Button(interactive=False)]),
+        outputs=[download_glb, download_gs],
+    )
+    
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    pipeline = TrellisVGGTTo3DPipeline.from_pretrained("weights/trellis-vggt-v0-1")
+    # pipeline = TrellisVGGTTo3DPipeline.from_pretrained("Stable-X/trellis-vggt-v0-1")
+    pipeline.cuda()
+    pipeline.VGGT_model.cuda()
+    try:
+        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
+    except:
+        pass
+    demo.launch()

extensions/nvdiffrast/LICENSE.txt

@@ -0,0 +1,97 @@
+Copyright (c) 2020, NVIDIA Corporation. All rights reserved.
+
+
+Nvidia Source Code License (1-Way Commercial)
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. The Work or
+    derivative works thereof may be used or intended for use by Nvidia
+    or its affiliates commercially or non-commercially. As used herein,
+    "non-commercially" means for research or evaluation purposes only
+    and not for any direct or indirect monetary gain.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor's or its affiliates' names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

extensions/nvdiffrast/README.md

@@ -0,0 +1,42 @@
+## Nvdiffrast – Modular Primitives for High-Performance Differentiable Rendering
+
+![Teaser image](./docs/img/teaser.png)
+
+**Modular Primitives for High-Performance Differentiable Rendering**<br>
+Samuli Laine, Janne Hellsten, Tero Karras, Yeongho Seol, Jaakko Lehtinen, Timo Aila<br>
+[http://arxiv.org/abs/2011.03277](http://arxiv.org/abs/2011.03277)
+
+Nvdiffrast is a PyTorch/TensorFlow library that provides high-performance primitive operations for rasterization-based differentiable rendering.
+Please refer to &#x261E;&#x261E; [nvdiffrast documentation](https://nvlabs.github.io/nvdiffrast) &#x261C;&#x261C; for more information.
+
+## Licenses
+
+Copyright &copy; 2020&ndash;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}
+}
+```

extensions/nvdiffrast/nvdiffrast/__init__.py

@@ -0,0 +1,9 @@
+# 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'

extensions/nvdiffrast/nvdiffrast/common/antialias.cu

@@ -0,0 +1,558 @@
+// Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include "antialias.h"
+
+//------------------------------------------------------------------------
+// Helpers.
+
+#define F32_MAX (3.402823466e+38f)
+static __forceinline__ __device__ bool same_sign(float a, float b) { return (__float_as_int(a) ^ __float_as_int(b)) >= 0; }
+static __forceinline__ __device__ bool rational_gt(float n0, float n1, float d0, float d1) { return (n0*d1 > n1*d0) == same_sign(d0, d1); }
+static __forceinline__ __device__ int max_idx3(float n0, float n1, float n2, float d0, float d1, float d2)
+{
+    bool g10 = rational_gt(n1, n0, d1, d0);
+    bool g20 = rational_gt(n2, n0, d2, d0);
+    bool g21 = rational_gt(n2, n1, d2, d1);
+    if (g20 && g21) return 2;
+    if (g10) return 1;
+    return 0;
+}
+
+//------------------------------------------------------------------------
+// Format of antialiasing work items stored in work buffer. Usually accessed directly as int4.
+
+struct AAWorkItem
+{
+    enum
+    {
+        EDGE_MASK       = 3,    // Edge index in lowest bits.
+        FLAG_DOWN_BIT   = 2,    // Down instead of right.
+        FLAG_TRI1_BIT   = 3,    // Edge is from other pixel's triangle.
+    };
+
+    int             px, py;         // Pixel x, y.
+    unsigned int    pz_flags;       // High 16 bits = pixel z, low 16 bits = edge index and flags.
+    float           alpha;          // Antialiasing alpha value. Zero if no AA.
+};
+
+//------------------------------------------------------------------------
+// Hash functions. Adapted from public-domain code at http://www.burtleburtle.net/bob/hash/doobs.html
+
+#define JENKINS_MAGIC (0x9e3779b9u)
+static __device__ __forceinline__ void jenkins_mix(unsigned int& a, unsigned int& b, unsigned int& c)
+{
+    a -= b; a -= c; a ^= (c>>13);
+    b -= c; b -= a; b ^= (a<<8);
+    c -= a; c -= b; c ^= (b>>13);
+    a -= b; a -= c; a ^= (c>>12);
+    b -= c; b -= a; b ^= (a<<16);
+    c -= a; c -= b; c ^= (b>>5);
+    a -= b; a -= c; a ^= (c>>3);
+    b -= c; b -= a; b ^= (a<<10);
+    c -= a; c -= b; c ^= (b>>15);
+}
+
+// Helper class for hash index iteration. Implements simple odd-skip linear probing with a key-dependent skip.
+class HashIndex
+{
+public:
+    __device__ __forceinline__ HashIndex(const AntialiasKernelParams& p, uint64_t key)
+    {
+        m_mask = (p.allocTriangles << AA_LOG_HASH_ELEMENTS_PER_TRIANGLE(p.allocTriangles)) - 1; // This should work until triangle count exceeds 1073741824.
+        m_idx  = (uint32_t)(key & 0xffffffffu);
+        m_skip = (uint32_t)(key >> 32);
+        uint32_t dummy = JENKINS_MAGIC;
+        jenkins_mix(m_idx, m_skip, dummy);
+        m_idx &= m_mask;
+        m_skip &= m_mask;
+        m_skip |= 1;
+    }
+    __device__ __forceinline__ int get(void) const { return m_idx; }
+    __device__ __forceinline__ void next(void) { m_idx = (m_idx + m_skip) & m_mask; }
+private:
+    uint32_t m_idx, m_skip, m_mask;
+};
+
+static __device__ __forceinline__ void hash_insert(const AntialiasKernelParams& p, uint64_t key, int v)
+{
+    HashIndex idx(p, key);
+    while(1)
+    {
+        uint64_t prev = atomicCAS((unsigned long long*)&p.evHash[idx.get()], 0, (unsigned long long)key);
+        if (prev == 0 || prev == key)
+            break;
+        idx.next();
+    }
+    int* q = (int*)&p.evHash[idx.get()];
+    int a = atomicCAS(q+2, 0, v);
+    if (a != 0 && a != v)
+        atomicCAS(q+3, 0, v);
+}
+
+static __device__ __forceinline__ int2 hash_find(const AntialiasKernelParams& p, uint64_t key)
+{
+    HashIndex idx(p, key);
+    while(1)
+    {
+        uint4 entry = p.evHash[idx.get()];
+        uint64_t k = ((uint64_t)entry.x) | (((uint64_t)entry.y) << 32);
+        if (k == key || k == 0)
+            return make_int2((int)entry.z, (int)entry.w);
+        idx.next();
+    }
+}
+
+static __device__ __forceinline__ void evhash_insert_vertex(const AntialiasKernelParams& p, int va, int vb, int vn)
+{
+    if (va == vb)
+        return;
+    
+    uint64_t v0 = (uint32_t)min(va, vb) + 1; // canonical vertex order
+    uint64_t v1 = (uint32_t)max(va, vb) + 1;
+    uint64_t vk = v0 | (v1 << 32); // hash key
+    hash_insert(p, vk, vn + 1);
+}
+
+static __forceinline__ __device__ int evhash_find_vertex(const AntialiasKernelParams& p, int va, int vb, int vr)
+{
+    if (va == vb)
+        return -1;
+
+    uint64_t v0 = (uint32_t)min(va, vb) + 1; // canonical vertex order
+    uint64_t v1 = (uint32_t)max(va, vb) + 1;
+    uint64_t vk = v0 | (v1 << 32); // hash key
+    int2 vn = hash_find(p, vk) - 1;
+    if (vn.x == vr) return vn.y;
+    if (vn.y == vr) return vn.x;
+    return -1;
+}
+
+//------------------------------------------------------------------------
+// Mesh analysis kernel.
+
+__global__ void AntialiasFwdMeshKernel(const AntialiasKernelParams p)
+{
+    int idx = threadIdx.x + blockIdx.x * blockDim.x;
+    if (idx >= p.numTriangles)
+        return;
+
+    int v0 = p.tri[idx * 3 + 0];
+    int v1 = p.tri[idx * 3 + 1];
+    int v2 = p.tri[idx * 3 + 2];
+
+    if (v0 < 0 || v0 >= p.numVertices ||
+        v1 < 0 || v1 >= p.numVertices ||
+        v2 < 0 || v2 >= p.numVertices)
+        return;
+
+    if (v0 == v1 || v1 == v2 || v2 == v0)
+        return;
+
+    evhash_insert_vertex(p, v1, v2, v0);
+    evhash_insert_vertex(p, v2, v0, v1);
+    evhash_insert_vertex(p, v0, v1, v2);
+}
+
+//------------------------------------------------------------------------
+// Discontinuity finder kernel.
+
+__global__ void AntialiasFwdDiscontinuityKernel(const AntialiasKernelParams p)
+{
+    // Calculate pixel position.
+    int px = blockIdx.x * AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH + threadIdx.x;
+    int py = blockIdx.y * AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT + threadIdx.y;
+    int pz = blockIdx.z;
+    if (px >= p.width || py >= p.height || pz >= p.n)
+        return;
+
+    // Pointer to our TriIdx and fetch.
+    int pidx0 = ((px + p.width * (py + p.height * pz)) << 2) + 3;
+    float tri0 = p.rasterOut[pidx0]; // These can stay as float, as we only compare them against each other.
+
+    // Look right, clamp at edge.
+    int pidx1 = pidx0;
+    if (px < p.width - 1)
+        pidx1 += 4;
+    float tri1 = p.rasterOut[pidx1];
+
+    // Look down, clamp at edge.
+    int pidx2 = pidx0;
+    if (py < p.height - 1)
+        pidx2 += p.width << 2;
+    float tri2 = p.rasterOut[pidx2];
+
+    // Determine amount of work.
+    int count = 0;
+    if (tri1 != tri0) count  = 1;
+    if (tri2 != tri0) count += 1;
+    if (!count)
+        return; // Exit warp.
+
+    // Coalesce work counter update to once per CTA.
+    __shared__ int s_temp;
+    s_temp = 0;
+    __syncthreads();
+    int idx = atomicAdd(&s_temp, count);
+    __syncthreads();
+    if (idx == 0)
+    {
+        int base = atomicAdd(&p.workBuffer[0].x, s_temp);
+        s_temp = base + 1; // don't clobber the counters in first slot.
+    }
+    __syncthreads();
+    idx += s_temp;
+
+    // Write to memory.
+    if (tri1 != tri0) p.workBuffer[idx++] = make_int4(px, py, (pz << 16), 0);
+    if (tri2 != tri0) p.workBuffer[idx]   = make_int4(px, py, (pz << 16) + (1 << AAWorkItem::FLAG_DOWN_BIT), 0);
+}
+
+//------------------------------------------------------------------------
+// Forward analysis kernel.
+
+__global__ void AntialiasFwdAnalysisKernel(const AntialiasKernelParams p)
+{
+    __shared__ int s_base;
+    int workCount = p.workBuffer[0].x;
+    for(;;)
+    {
+        // Persistent threads work fetcher.
+        __syncthreads();
+        if (threadIdx.x == 0)
+            s_base = atomicAdd(&p.workBuffer[0].y, AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK);
+        __syncthreads();
+        int thread_idx = s_base + threadIdx.x;
+        if (thread_idx >= workCount)
+            return;
+
+        int4* pItem = p.workBuffer + thread_idx + 1;
+        int4 item = *pItem;
+        int px = item.x;
+        int py = item.y;
+        int pz = (int)(((unsigned int)item.z) >> 16);
+        int d  = (item.z >> AAWorkItem::FLAG_DOWN_BIT) & 1;
+
+        int pixel0 = px + p.width * (py + p.height * pz);
+        int pixel1 = pixel0 + (d ? p.width : 1);
+        float2 zt0 = ((float2*)p.rasterOut)[(pixel0 << 1) + 1];
+        float2 zt1 = ((float2*)p.rasterOut)[(pixel1 << 1) + 1];
+        int tri0 = float_to_triidx(zt0.y) - 1;
+        int tri1 = float_to_triidx(zt1.y) - 1;
+
+        // Select triangle based on background / depth.
+        int tri = (tri0 >= 0) ? tri0 : tri1;
+        if (tri0 >= 0 && tri1 >= 0)
+            tri = (zt0.x < zt1.x) ? tri0 : tri1;
+        if (tri == tri1)
+        {
+            // Calculate with respect to neighbor pixel if chose that triangle.
+            px += 1 - d;
+            py += d;
+        }
+
+        // Bail out if triangle index is corrupt.
+        if (tri < 0 || tri >= p.numTriangles)
+            continue;
+
+        // Fetch vertex indices.
+        int vi0 = p.tri[tri * 3 + 0];
+        int vi1 = p.tri[tri * 3 + 1];
+        int vi2 = p.tri[tri * 3 + 2];
+
+        // Bail out if vertex indices are corrupt.
+        if (vi0 < 0 || vi0 >= p.numVertices ||
+            vi1 < 0 || vi1 >= p.numVertices ||
+            vi2 < 0 || vi2 >= p.numVertices)
+            continue;
+
+        // Fetch opposite vertex indices. Use vertex itself (always silhouette) if no opposite vertex exists.
+        int op0 = evhash_find_vertex(p, vi2, vi1, vi0);
+        int op1 = evhash_find_vertex(p, vi0, vi2, vi1);
+        int op2 = evhash_find_vertex(p, vi1, vi0, vi2);
+
+        // Instance mode: Adjust vertex indices based on minibatch index.
+        if (p.instance_mode)
+        {
+            int vbase = pz * p.numVertices;
+            vi0 += vbase;
+            vi1 += vbase;
+            vi2 += vbase;
+            if (op0 >= 0) op0 += vbase;
+            if (op1 >= 0) op1 += vbase;
+            if (op2 >= 0) op2 += vbase;
+        }
+
+        // Fetch vertex positions.
+        float4 p0 = ((float4*)p.pos)[vi0];
+        float4 p1 = ((float4*)p.pos)[vi1];
+        float4 p2 = ((float4*)p.pos)[vi2];
+        float4 o0 = (op0 < 0) ? p0 : ((float4*)p.pos)[op0];
+        float4 o1 = (op1 < 0) ? p1 : ((float4*)p.pos)[op1];
+        float4 o2 = (op2 < 0) ? p2 : ((float4*)p.pos)[op2];
+
+        // Project vertices to pixel space.
+        float w0  = 1.f / p0.w;
+        float w1  = 1.f / p1.w;
+        float w2  = 1.f / p2.w;
+        float ow0 = 1.f / o0.w;
+        float ow1 = 1.f / o1.w;
+        float ow2 = 1.f / o2.w;
+        float fx  = (float)px + .5f - p.xh;
+        float fy  = (float)py + .5f - p.yh;
+        float x0  = p0.x * w0 * p.xh - fx;
+        float y0  = p0.y * w0 * p.yh - fy;
+        float x1  = p1.x * w1 * p.xh - fx;
+        float y1  = p1.y * w1 * p.yh - fy;
+        float x2  = p2.x * w2 * p.xh - fx;
+        float y2  = p2.y * w2 * p.yh - fy;
+        float ox0 = o0.x * ow0 * p.xh - fx;
+        float oy0 = o0.y * ow0 * p.yh - fy;
+        float ox1 = o1.x * ow1 * p.xh - fx;
+        float oy1 = o1.y * ow1 * p.yh - fy;
+        float ox2 = o2.x * ow2 * p.xh - fx;
+        float oy2 = o2.y * ow2 * p.yh - fy;
+
+        // Signs to kill non-silhouette edges.
+        float bb = (x1-x0)*(y2-y0) - (x2-x0)*(y1-y0); // Triangle itself.
+        float a0 = (x1-ox0)*(y2-oy0) - (x2-ox0)*(y1-oy0); // Wings.
+        float a1 = (x2-ox1)*(y0-oy1) - (x0-ox1)*(y2-oy1);
+        float a2 = (x0-ox2)*(y1-oy2) - (x1-ox2)*(y0-oy2);
+
+        // If no matching signs anywhere, skip the rest.
+        if (same_sign(a0, bb) || same_sign(a1, bb) || same_sign(a2, bb))
+        {
+            // XY flip for horizontal edges.
+            if (d)
+            {
+                swap(x0, y0);
+                swap(x1, y1);
+                swap(x2, y2);
+            }
+
+            float dx0 = x2 - x1;
+            float dx1 = x0 - x2;
+            float dx2 = x1 - x0;
+            float dy0 = y2 - y1;
+            float dy1 = y0 - y2;
+            float dy2 = y1 - y0;
+
+            // Check if an edge crosses between us and the neighbor pixel.
+            float dc = -F32_MAX;
+            float ds = (tri == tri0) ? 1.f : -1.f;
+            float d0 = ds * (x1*dy0 - y1*dx0);
+            float d1 = ds * (x2*dy1 - y2*dx1);
+            float d2 = ds * (x0*dy2 - y0*dx2);
+
+            if (same_sign(y1, y2)) d0 = -F32_MAX, dy0 = 1.f;
+            if (same_sign(y2, y0)) d1 = -F32_MAX, dy1 = 1.f;
+            if (same_sign(y0, y1)) d2 = -F32_MAX, dy2 = 1.f;
+
+            int di = max_idx3(d0, d1, d2, dy0, dy1, dy2);
+            if (di == 0 && same_sign(a0, bb) && fabsf(dy0) >= fabsf(dx0)) dc = d0 / dy0;
+            if (di == 1 && same_sign(a1, bb) && fabsf(dy1) >= fabsf(dx1)) dc = d1 / dy1;
+            if (di == 2 && same_sign(a2, bb) && fabsf(dy2) >= fabsf(dx2)) dc = d2 / dy2;
+            float eps = .0625f; // Expect no more than 1/16 pixel inaccuracy.
+
+            // Adjust output image if a suitable edge was found.
+            if (dc > -eps && dc < 1.f + eps)
+            {
+                dc = fminf(fmaxf(dc, 0.f), 1.f);
+                float alpha = ds * (.5f - dc);
+                const float* pColor0 = p.color + pixel0 * p.channels;
+                const float* pColor1 = p.color + pixel1 * p.channels;
+                float* pOutput = p.output + (alpha > 0.f ? pixel0 : pixel1) * p.channels;
+                for (int i=0; i < p.channels; i++)
+                    atomicAdd(&pOutput[i], alpha * (pColor1[i] - pColor0[i]));
+
+                // Rewrite the work item's flags and alpha. Keep original px, py.
+                unsigned int flags = pz << 16;
+                flags |= di;
+                flags |= d << AAWorkItem::FLAG_DOWN_BIT;
+                flags |= (__float_as_uint(ds) >> 31) << AAWorkItem::FLAG_TRI1_BIT;
+                ((int2*)pItem)[1] = make_int2(flags, __float_as_int(alpha));
+            }
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// Gradient kernel.
+
+__global__ void AntialiasGradKernel(const AntialiasKernelParams p)
+{
+    // Temporary space for coalesced atomics.
+    CA_DECLARE_TEMP(AA_GRAD_KERNEL_THREADS_PER_BLOCK);
+    __shared__ int s_base; // Work counter communication across entire CTA.
+
+    int workCount = p.workBuffer[0].x;
+
+    for(;;)
+    {
+        // Persistent threads work fetcher.
+        __syncthreads();
+        if (threadIdx.x == 0)
+            s_base = atomicAdd(&p.workBuffer[0].y, AA_GRAD_KERNEL_THREADS_PER_BLOCK);
+        __syncthreads();
+        int thread_idx = s_base + threadIdx.x;
+        if (thread_idx >= workCount)
+            return;
+
+        // Read work item filled out by forward kernel.
+        int4 item = p.workBuffer[thread_idx + 1];
+        unsigned int amask = __ballot_sync(0xffffffffu, item.w);
+        if (item.w == 0)
+            continue; // No effect.
+
+        // Unpack work item and replicate setup from forward analysis kernel.
+        int px = item.x;
+        int py = item.y;
+        int pz = (int)(((unsigned int)item.z) >> 16);
+        int d = (item.z >> AAWorkItem::FLAG_DOWN_BIT) & 1;
+        float alpha = __int_as_float(item.w);
+        int tri1 = (item.z >> AAWorkItem::FLAG_TRI1_BIT) & 1;
+        int di = item.z & AAWorkItem::EDGE_MASK;
+        float ds = __int_as_float(__float_as_int(1.0) | (tri1 << 31));
+        int pixel0 = px + p.width * (py + p.height * pz);
+        int pixel1 = pixel0 + (d ? p.width : 1);
+        int tri = float_to_triidx(p.rasterOut[((tri1 ? pixel1 : pixel0) << 2) + 3]) - 1;
+        if (tri1)
+        {
+            px += 1 - d;
+            py += d;
+        }
+
+        // Bail out if triangle index is corrupt.
+        bool triFail = (tri < 0 || tri >= p.numTriangles);
+        amask = __ballot_sync(amask, !triFail);
+        if (triFail)
+            continue;
+
+        // Outgoing color gradients.
+        float* pGrad0 = p.gradColor + pixel0 * p.channels;
+        float* pGrad1 = p.gradColor + pixel1 * p.channels;
+
+        // Incoming color gradients.
+        const float* pDy = p.dy + (alpha > 0.f ? pixel0 : pixel1) * p.channels;
+
+        // Position gradient weight based on colors and incoming gradients.
+        float dd = 0.f;
+        const float* pColor0 = p.color + pixel0 * p.channels;
+        const float* pColor1 = p.color + pixel1 * p.channels;
+
+        // Loop over channels and accumulate.
+        for (int i=0; i < p.channels; i++)
+        {
+            float dy = pDy[i];
+            if (dy != 0.f)
+            {
+                // Update position gradient weight.
+                dd += dy * (pColor1[i] - pColor0[i]);
+
+                // Update color gradients. No coalescing because all have different targets.
+                float v = alpha * dy;
+                atomicAdd(&pGrad0[i], -v);
+                atomicAdd(&pGrad1[i], v);
+            }
+        }
+
+        // If position weight is zero, skip the rest.
+        bool noGrad = (dd == 0.f);
+        amask = __ballot_sync(amask, !noGrad);
+        if (noGrad)
+            continue;
+
+        // Fetch vertex indices of the active edge and their positions.
+        int i1 = (di < 2) ? (di + 1) : 0;
+        int i2 = (i1 < 2) ? (i1 + 1) : 0;
+        int vi1 = p.tri[3 * tri + i1];
+        int vi2 = p.tri[3 * tri + i2];
+
+        // Bail out if vertex indices are corrupt.
+        bool vtxFail = (vi1 < 0 || vi1 >= p.numVertices || vi2 < 0 || vi2 >= p.numVertices);
+        amask = __ballot_sync(amask, !vtxFail);
+        if (vtxFail)
+            continue;
+
+        // Instance mode: Adjust vertex indices based on minibatch index.
+        if (p.instance_mode)
+        {
+            vi1 += pz * p.numVertices;
+            vi2 += pz * p.numVertices;
+        }
+
+        // Fetch vertex positions.
+        float4 p1 = ((float4*)p.pos)[vi1];
+        float4 p2 = ((float4*)p.pos)[vi2];
+
+        // Project vertices to pixel space.
+        float pxh = p.xh;
+        float pyh = p.yh;
+        float fx = (float)px + .5f - pxh;
+        float fy = (float)py + .5f - pyh;
+
+        // XY flip for horizontal edges.
+        if (d)
+        {
+            swap(p1.x, p1.y);
+            swap(p2.x, p2.y);
+            swap(pxh, pyh);
+            swap(fx, fy);
+        }
+
+        // Gradient calculation setup.
+        float w1 = 1.f / p1.w;
+        float w2 = 1.f / p2.w;
+        float x1 = p1.x * w1 * pxh - fx;
+        float y1 = p1.y * w1 * pyh - fy;
+        float x2 = p2.x * w2 * pxh - fx;
+        float y2 = p2.y * w2 * pyh - fy;
+        float dx = x2 - x1;
+        float dy = y2 - y1;
+        float db = x1*dy - y1*dx;
+
+        // Compute inverse delta-y with epsilon.
+        float ep = copysignf(1e-3f, dy); // ~1/1000 pixel.
+        float iy = 1.f / (dy + ep);
+
+        // Compute position gradients.
+        float dby = db * iy;
+        float iw1 = -w1 * iy * dd;
+        float iw2 =  w2 * iy * dd;
+        float gp1x = iw1 * pxh * y2;
+        float gp2x = iw2 * pxh * y1;
+        float gp1y = iw1 * pyh * (dby - x2);
+        float gp2y = iw2 * pyh * (dby - x1);
+        float gp1w = -(p1.x * gp1x + p1.y * gp1y) * w1;
+        float gp2w = -(p2.x * gp2x + p2.y * gp2y) * w2;
+
+        // XY flip the gradients.
+        if (d)
+        {
+            swap(gp1x, gp1y);
+            swap(gp2x, gp2y);
+        }
+
+        // Kill position gradients if alpha was saturated.
+        if (fabsf(alpha) >= 0.5f)
+        {
+            gp1x = gp1y = gp1w = 0.f;
+            gp2x = gp2y = gp2w = 0.f;
+        }
+
+        // Initialize coalesced atomics. Match both triangle ID and edge index.
+        // Also note that some threads may be inactive.
+        CA_SET_GROUP_MASK(tri ^ (di << 30), amask);
+
+        // Accumulate gradients.
+        caAtomicAdd3_xyw(p.gradPos + 4 * vi1, gp1x, gp1y, gp1w);
+        caAtomicAdd3_xyw(p.gradPos + 4 * vi2, gp2x, gp2y, gp2w);
+    }
+}
+
+//------------------------------------------------------------------------

extensions/nvdiffrast/nvdiffrast/common/antialias.h

@@ -0,0 +1,50 @@
+// Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+#include "common.h"
+
+//------------------------------------------------------------------------
+// Constants and helpers.
+
+#define AA_DISCONTINUITY_KERNEL_BLOCK_WIDTH         32
+#define AA_DISCONTINUITY_KERNEL_BLOCK_HEIGHT        8
+#define AA_ANALYSIS_KERNEL_THREADS_PER_BLOCK        256
+#define AA_MESH_KERNEL_THREADS_PER_BLOCK            256
+#define AA_HASH_ELEMENTS_PER_TRIANGLE(alloc)        ((alloc) >= (2 << 25) ? 4 : 8) // With more than 16777216 triangles (alloc >= 33554432) use smallest possible value of 4 to conserve memory, otherwise use 8 for fewer collisions.
+#define AA_LOG_HASH_ELEMENTS_PER_TRIANGLE(alloc)    ((alloc) >= (2 << 25) ? 2 : 3)
+#define AA_GRAD_KERNEL_THREADS_PER_BLOCK            256
+
+//------------------------------------------------------------------------
+// CUDA kernel params.
+
+struct AntialiasKernelParams
+{
+    const float*    color;          // Incoming color buffer.
+    const float*    rasterOut;      // Incoming rasterizer output buffer.
+    const int*      tri;            // Incoming triangle buffer.
+    const float*    pos;            // Incoming position buffer.
+    float*          output;         // Output buffer of forward kernel.
+    const float*    dy;             // Incoming gradients.
+    float*          gradColor;      // Output buffer, color gradient.
+    float*          gradPos;        // Output buffer, position gradient.
+    int4*           workBuffer;     // Buffer for storing intermediate work items. First item reserved for counters.
+    uint4*          evHash;         // Edge-vertex hash.
+    int             allocTriangles; // Number of triangles accommodated by evHash. Always power of two.
+    int             numTriangles;   // Number of triangles.
+    int             numVertices;    // Number of vertices.
+    int             width;          // Input width.
+    int             height;         // Input height.
+    int             n;              // Minibatch size.
+    int             channels;       // Channel count in color input.
+    float           xh, yh;         // Transfer to pixel space.
+    int             instance_mode;  // 0=normal, 1=instance mode.
+    int             tri_const;      // 1 if triangle array is known to be constant.
+};
+
+//------------------------------------------------------------------------

extensions/nvdiffrast/nvdiffrast/common/common.cpp

@@ -0,0 +1,60 @@
+// Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <cuda_runtime.h>
+
+//------------------------------------------------------------------------
+// Block and grid size calculators for kernel launches.
+
+dim3 getLaunchBlockSize(int maxWidth, int maxHeight, int width, int height)
+{
+    int maxThreads = maxWidth * maxHeight;
+    if (maxThreads <= 1 || (width * height) <= 1)
+        return dim3(1, 1, 1); // Degenerate.
+
+    // Start from max size.
+    int bw = maxWidth;
+    int bh = maxHeight;
+
+    // Optimizations for weirdly sized buffers.
+    if (width < bw)
+    {
+        // Decrease block width to smallest power of two that covers the buffer width.
+        while ((bw >> 1) >= width)
+            bw >>= 1;
+
+        // Maximize height.
+        bh = maxThreads / bw;
+        if (bh > height)
+            bh = height;
+    }
+    else if (height < bh)
+    {
+        // Halve height and double width until fits completely inside buffer vertically.
+        while (bh > height)
+        {
+            bh >>= 1;
+            if (bw < width)
+                bw <<= 1;
+        }
+    }
+
+    // Done.
+    return dim3(bw, bh, 1);
+}
+
+dim3 getLaunchGridSize(dim3 blockSize, int width, int height, int depth)
+{
+    dim3 gridSize;
+    gridSize.x = (width  - 1) / blockSize.x + 1;
+    gridSize.y = (height - 1) / blockSize.y + 1;
+    gridSize.z = (depth  - 1) / blockSize.z + 1;
+    return gridSize;
+}
+
+//------------------------------------------------------------------------

extensions/nvdiffrast/nvdiffrast/common/common.h

@@ -0,0 +1,263 @@
+// Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+#include <cuda.h>
+#include <stdint.h>
+
+//------------------------------------------------------------------------
+// C++ helper function prototypes.
+
+dim3 getLaunchBlockSize(int maxWidth, int maxHeight, int width, int height);
+dim3 getLaunchGridSize(dim3 blockSize, int width, int height, int depth);
+
+//------------------------------------------------------------------------
+// The rest is CUDA device code specific stuff.
+
+#ifdef __CUDACC__
+
+//------------------------------------------------------------------------
+// Helpers for CUDA vector types.
+
+static __device__ __forceinline__ float2&   operator*=  (float2& a, const float2& b)       { a.x *= b.x; a.y *= b.y; return a; }
+static __device__ __forceinline__ float2&   operator+=  (float2& a, const float2& b)       { a.x += b.x; a.y += b.y; return a; }
+static __device__ __forceinline__ float2&   operator-=  (float2& a, const float2& b)       { a.x -= b.x; a.y -= b.y; return a; }
+static __device__ __forceinline__ float2&   operator*=  (float2& a, float b)               { a.x *= b; a.y *= b; return a; }
+static __device__ __forceinline__ float2&   operator+=  (float2& a, float b)               { a.x += b; a.y += b; return a; }
+static __device__ __forceinline__ float2&   operator-=  (float2& a, float b)               { a.x -= b; a.y -= b; return a; }
+static __device__ __forceinline__ float2    operator*   (const float2& a, const float2& b) { return make_float2(a.x * b.x, a.y * b.y); }
+static __device__ __forceinline__ float2    operator+   (const float2& a, const float2& b) { return make_float2(a.x + b.x, a.y + b.y); }
+static __device__ __forceinline__ float2    operator-   (const float2& a, const float2& b) { return make_float2(a.x - b.x, a.y - b.y); }
+static __device__ __forceinline__ float2    operator*   (const float2& a, float b)         { return make_float2(a.x * b, a.y * b); }
+static __device__ __forceinline__ float2    operator+   (const float2& a, float b)         { return make_float2(a.x + b, a.y + b); }
+static __device__ __forceinline__ float2    operator-   (const float2& a, float b)         { return make_float2(a.x - b, a.y - b); }
+static __device__ __forceinline__ float2    operator*   (float a, const float2& b)         { return make_float2(a * b.x, a * b.y); }
+static __device__ __forceinline__ float2    operator+   (float a, const float2& b)         { return make_float2(a + b.x, a + b.y); }
+static __device__ __forceinline__ float2    operator-   (float a, const float2& b)         { return make_float2(a - b.x, a - b.y); }
+static __device__ __forceinline__ float2    operator-   (const float2& a)                  { return make_float2(-a.x, -a.y); }
+static __device__ __forceinline__ float3&   operator*=  (float3& a, const float3& b)       { a.x *= b.x; a.y *= b.y; a.z *= b.z; return a; }
+static __device__ __forceinline__ float3&   operator+=  (float3& a, const float3& b)       { a.x += b.x; a.y += b.y; a.z += b.z; return a; }
+static __device__ __forceinline__ float3&   operator-=  (float3& a, const float3& b)       { a.x -= b.x; a.y -= b.y; a.z -= b.z; return a; }
+static __device__ __forceinline__ float3&   operator*=  (float3& a, float b)               { a.x *= b; a.y *= b; a.z *= b; return a; }
+static __device__ __forceinline__ float3&   operator+=  (float3& a, float b)               { a.x += b; a.y += b; a.z += b; return a; }
+static __device__ __forceinline__ float3&   operator-=  (float3& a, float b)               { a.x -= b; a.y -= b; a.z -= b; return a; }
+static __device__ __forceinline__ float3    operator*   (const float3& a, const float3& b) { return make_float3(a.x * b.x, a.y * b.y, a.z * b.z); }
+static __device__ __forceinline__ float3    operator+   (const float3& a, const float3& b) { return make_float3(a.x + b.x, a.y + b.y, a.z + b.z); }
+static __device__ __forceinline__ float3    operator-   (const float3& a, const float3& b) { return make_float3(a.x - b.x, a.y - b.y, a.z - b.z); }
+static __device__ __forceinline__ float3    operator*   (const float3& a, float b)         { return make_float3(a.x * b, a.y * b, a.z * b); }
+static __device__ __forceinline__ float3    operator+   (const float3& a, float b)         { return make_float3(a.x + b, a.y + b, a.z + b); }
+static __device__ __forceinline__ float3    operator-   (const float3& a, float b)         { return make_float3(a.x - b, a.y - b, a.z - b); }
+static __device__ __forceinline__ float3    operator*   (float a, const float3& b)         { return make_float3(a * b.x, a * b.y, a * b.z); }
+static __device__ __forceinline__ float3    operator+   (float a, const float3& b)         { return make_float3(a + b.x, a + b.y, a + b.z); }
+static __device__ __forceinline__ float3    operator-   (float a, const float3& b)         { return make_float3(a - b.x, a - b.y, a - b.z); }
+static __device__ __forceinline__ float3    operator-   (const float3& a)                  { return make_float3(-a.x, -a.y, -a.z); }
+static __device__ __forceinline__ float4&   operator*=  (float4& a, const float4& b)       { a.x *= b.x; a.y *= b.y; a.z *= b.z; a.w *= b.w; return a; }
+static __device__ __forceinline__ float4&   operator+=  (float4& a, const float4& b)       { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; return a; }
+static __device__ __forceinline__ float4&   operator-=  (float4& a, const float4& b)       { a.x -= b.x; a.y -= b.y; a.z -= b.z; a.w -= b.w; return a; }
+static __device__ __forceinline__ float4&   operator*=  (float4& a, float b)               { a.x *= b; a.y *= b; a.z *= b; a.w *= b; return a; }
+static __device__ __forceinline__ float4&   operator+=  (float4& a, float b)               { a.x += b; a.y += b; a.z += b; a.w += b; return a; }
+static __device__ __forceinline__ float4&   operator-=  (float4& a, float b)               { a.x -= b; a.y -= b; a.z -= b; a.w -= b; return a; }
+static __device__ __forceinline__ float4    operator*   (const float4& a, const float4& b) { return make_float4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w); }
+static __device__ __forceinline__ float4    operator+   (const float4& a, const float4& b) { return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
+static __device__ __forceinline__ float4    operator-   (const float4& a, const float4& b) { return make_float4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w); }
+static __device__ __forceinline__ float4    operator*   (const float4& a, float b)         { return make_float4(a.x * b, a.y * b, a.z * b, a.w * b); }
+static __device__ __forceinline__ float4    operator+   (const float4& a, float b)         { return make_float4(a.x + b, a.y + b, a.z + b, a.w + b); }
+static __device__ __forceinline__ float4    operator-   (const float4& a, float b)         { return make_float4(a.x - b, a.y - b, a.z - b, a.w - b); }
+static __device__ __forceinline__ float4    operator*   (float a, const float4& b)         { return make_float4(a * b.x, a * b.y, a * b.z, a * b.w); }
+static __device__ __forceinline__ float4    operator+   (float a, const float4& b)         { return make_float4(a + b.x, a + b.y, a + b.z, a + b.w); }
+static __device__ __forceinline__ float4    operator-   (float a, const float4& b)         { return make_float4(a - b.x, a - b.y, a - b.z, a - b.w); }
+static __device__ __forceinline__ float4    operator-   (const float4& a)                  { return make_float4(-a.x, -a.y, -a.z, -a.w); }
+static __device__ __forceinline__ int2&     operator*=  (int2& a, const int2& b)           { a.x *= b.x; a.y *= b.y; return a; }
+static __device__ __forceinline__ int2&     operator+=  (int2& a, const int2& b)           { a.x += b.x; a.y += b.y; return a; }
+static __device__ __forceinline__ int2&     operator-=  (int2& a, const int2& b)           { a.x -= b.x; a.y -= b.y; return a; }
+static __device__ __forceinline__ int2&     operator*=  (int2& a, int b)                   { a.x *= b; a.y *= b; return a; }
+static __device__ __forceinline__ int2&     operator+=  (int2& a, int b)                   { a.x += b; a.y += b; return a; }
+static __device__ __forceinline__ int2&     operator-=  (int2& a, int b)                   { a.x -= b; a.y -= b; return a; }
+static __device__ __forceinline__ int2      operator*   (const int2& a, const int2& b)     { return make_int2(a.x * b.x, a.y * b.y); }
+static __device__ __forceinline__ int2      operator+   (const int2& a, const int2& b)     { return make_int2(a.x + b.x, a.y + b.y); }
+static __device__ __forceinline__ int2      operator-   (const int2& a, const int2& b)     { return make_int2(a.x - b.x, a.y - b.y); }
+static __device__ __forceinline__ int2      operator*   (const int2& a, int b)             { return make_int2(a.x * b, a.y * b); }
+static __device__ __forceinline__ int2      operator+   (const int2& a, int b)             { return make_int2(a.x + b, a.y + b); }
+static __device__ __forceinline__ int2      operator-   (const int2& a, int b)             { return make_int2(a.x - b, a.y - b); }
+static __device__ __forceinline__ int2      operator*   (int a, const int2& b)             { return make_int2(a * b.x, a * b.y); }
+static __device__ __forceinline__ int2      operator+   (int a, const int2& b)             { return make_int2(a + b.x, a + b.y); }
+static __device__ __forceinline__ int2      operator-   (int a, const int2& b)             { return make_int2(a - b.x, a - b.y); }
+static __device__ __forceinline__ int2      operator-   (const int2& a)                    { return make_int2(-a.x, -a.y); }
+static __device__ __forceinline__ int3&     operator*=  (int3& a, const int3& b)           { a.x *= b.x; a.y *= b.y; a.z *= b.z; return a; }
+static __device__ __forceinline__ int3&     operator+=  (int3& a, const int3& b)           { a.x += b.x; a.y += b.y; a.z += b.z; return a; }
+static __device__ __forceinline__ int3&     operator-=  (int3& a, const int3& b)           { a.x -= b.x; a.y -= b.y; a.z -= b.z; return a; }
+static __device__ __forceinline__ int3&     operator*=  (int3& a, int b)                   { a.x *= b; a.y *= b; a.z *= b; return a; }
+static __device__ __forceinline__ int3&     operator+=  (int3& a, int b)                   { a.x += b; a.y += b; a.z += b; return a; }
+static __device__ __forceinline__ int3&     operator-=  (int3& a, int b)                   { a.x -= b; a.y -= b; a.z -= b; return a; }
+static __device__ __forceinline__ int3      operator*   (const int3& a, const int3& b)     { return make_int3(a.x * b.x, a.y * b.y, a.z * b.z); }
+static __device__ __forceinline__ int3      operator+   (const int3& a, const int3& b)     { return make_int3(a.x + b.x, a.y + b.y, a.z + b.z); }
+static __device__ __forceinline__ int3      operator-   (const int3& a, const int3& b)     { return make_int3(a.x - b.x, a.y - b.y, a.z - b.z); }
+static __device__ __forceinline__ int3      operator*   (const int3& a, int b)             { return make_int3(a.x * b, a.y * b, a.z * b); }
+static __device__ __forceinline__ int3      operator+   (const int3& a, int b)             { return make_int3(a.x + b, a.y + b, a.z + b); }
+static __device__ __forceinline__ int3      operator-   (const int3& a, int b)             { return make_int3(a.x - b, a.y - b, a.z - b); }
+static __device__ __forceinline__ int3      operator*   (int a, const int3& b)             { return make_int3(a * b.x, a * b.y, a * b.z); }
+static __device__ __forceinline__ int3      operator+   (int a, const int3& b)             { return make_int3(a + b.x, a + b.y, a + b.z); }
+static __device__ __forceinline__ int3      operator-   (int a, const int3& b)             { return make_int3(a - b.x, a - b.y, a - b.z); }
+static __device__ __forceinline__ int3      operator-   (const int3& a)                    { return make_int3(-a.x, -a.y, -a.z); }
+static __device__ __forceinline__ int4&     operator*=  (int4& a, const int4& b)           { a.x *= b.x; a.y *= b.y; a.z *= b.z; a.w *= b.w; return a; }
+static __device__ __forceinline__ int4&     operator+=  (int4& a, const int4& b)           { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; return a; }
+static __device__ __forceinline__ int4&     operator-=  (int4& a, const int4& b)           { a.x -= b.x; a.y -= b.y; a.z -= b.z; a.w -= b.w; return a; }
+static __device__ __forceinline__ int4&     operator*=  (int4& a, int b)                   { a.x *= b; a.y *= b; a.z *= b; a.w *= b; return a; }
+static __device__ __forceinline__ int4&     operator+=  (int4& a, int b)                   { a.x += b; a.y += b; a.z += b; a.w += b; return a; }
+static __device__ __forceinline__ int4&     operator-=  (int4& a, int b)                   { a.x -= b; a.y -= b; a.z -= b; a.w -= b; return a; }
+static __device__ __forceinline__ int4      operator*   (const int4& a, const int4& b)     { return make_int4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w); }
+static __device__ __forceinline__ int4      operator+   (const int4& a, const int4& b)     { return make_int4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
+static __device__ __forceinline__ int4      operator-   (const int4& a, const int4& b)     { return make_int4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w); }
+static __device__ __forceinline__ int4      operator*   (const int4& a, int b)             { return make_int4(a.x * b, a.y * b, a.z * b, a.w * b); }
+static __device__ __forceinline__ int4      operator+   (const int4& a, int b)             { return make_int4(a.x + b, a.y + b, a.z + b, a.w + b); }
+static __device__ __forceinline__ int4      operator-   (const int4& a, int b)             { return make_int4(a.x - b, a.y - b, a.z - b, a.w - b); }
+static __device__ __forceinline__ int4      operator*   (int a, const int4& b)             { return make_int4(a * b.x, a * b.y, a * b.z, a * b.w); }
+static __device__ __forceinline__ int4      operator+   (int a, const int4& b)             { return make_int4(a + b.x, a + b.y, a + b.z, a + b.w); }
+static __device__ __forceinline__ int4      operator-   (int a, const int4& b)             { return make_int4(a - b.x, a - b.y, a - b.z, a - b.w); }
+static __device__ __forceinline__ int4      operator-   (const int4& a)                    { return make_int4(-a.x, -a.y, -a.z, -a.w); }
+static __device__ __forceinline__ uint2&    operator*=  (uint2& a, const uint2& b)         { a.x *= b.x; a.y *= b.y; return a; }
+static __device__ __forceinline__ uint2&    operator+=  (uint2& a, const uint2& b)         { a.x += b.x; a.y += b.y; return a; }
+static __device__ __forceinline__ uint2&    operator-=  (uint2& a, const uint2& b)         { a.x -= b.x; a.y -= b.y; return a; }
+static __device__ __forceinline__ uint2&    operator*=  (uint2& a, unsigned int b)         { a.x *= b; a.y *= b; return a; }
+static __device__ __forceinline__ uint2&    operator+=  (uint2& a, unsigned int b)         { a.x += b; a.y += b; return a; }
+static __device__ __forceinline__ uint2&    operator-=  (uint2& a, unsigned int b)         { a.x -= b; a.y -= b; return a; }
+static __device__ __forceinline__ uint2     operator*   (const uint2& a, const uint2& b)   { return make_uint2(a.x * b.x, a.y * b.y); }
+static __device__ __forceinline__ uint2     operator+   (const uint2& a, const uint2& b)   { return make_uint2(a.x + b.x, a.y + b.y); }
+static __device__ __forceinline__ uint2     operator-   (const uint2& a, const uint2& b)   { return make_uint2(a.x - b.x, a.y - b.y); }
+static __device__ __forceinline__ uint2     operator*   (const uint2& a, unsigned int b)   { return make_uint2(a.x * b, a.y * b); }
+static __device__ __forceinline__ uint2     operator+   (const uint2& a, unsigned int b)   { return make_uint2(a.x + b, a.y + b); }
+static __device__ __forceinline__ uint2     operator-   (const uint2& a, unsigned int b)   { return make_uint2(a.x - b, a.y - b); }
+static __device__ __forceinline__ uint2     operator*   (unsigned int a, const uint2& b)   { return make_uint2(a * b.x, a * b.y); }
+static __device__ __forceinline__ uint2     operator+   (unsigned int a, const uint2& b)   { return make_uint2(a + b.x, a + b.y); }
+static __device__ __forceinline__ uint2     operator-   (unsigned int a, const uint2& b)   { return make_uint2(a - b.x, a - b.y); }
+static __device__ __forceinline__ uint3&    operator*=  (uint3& a, const uint3& b)         { a.x *= b.x; a.y *= b.y; a.z *= b.z; return a; }
+static __device__ __forceinline__ uint3&    operator+=  (uint3& a, const uint3& b)         { a.x += b.x; a.y += b.y; a.z += b.z; return a; }
+static __device__ __forceinline__ uint3&    operator-=  (uint3& a, const uint3& b)         { a.x -= b.x; a.y -= b.y; a.z -= b.z; return a; }
+static __device__ __forceinline__ uint3&    operator*=  (uint3& a, unsigned int b)         { a.x *= b; a.y *= b; a.z *= b; return a; }
+static __device__ __forceinline__ uint3&    operator+=  (uint3& a, unsigned int b)         { a.x += b; a.y += b; a.z += b; return a; }
+static __device__ __forceinline__ uint3&    operator-=  (uint3& a, unsigned int b)         { a.x -= b; a.y -= b; a.z -= b; return a; }
+static __device__ __forceinline__ uint3     operator*   (const uint3& a, const uint3& b)   { return make_uint3(a.x * b.x, a.y * b.y, a.z * b.z); }
+static __device__ __forceinline__ uint3     operator+   (const uint3& a, const uint3& b)   { return make_uint3(a.x + b.x, a.y + b.y, a.z + b.z); }
+static __device__ __forceinline__ uint3     operator-   (const uint3& a, const uint3& b)   { return make_uint3(a.x - b.x, a.y - b.y, a.z - b.z); }
+static __device__ __forceinline__ uint3     operator*   (const uint3& a, unsigned int b)   { return make_uint3(a.x * b, a.y * b, a.z * b); }
+static __device__ __forceinline__ uint3     operator+   (const uint3& a, unsigned int b)   { return make_uint3(a.x + b, a.y + b, a.z + b); }
+static __device__ __forceinline__ uint3     operator-   (const uint3& a, unsigned int b)   { return make_uint3(a.x - b, a.y - b, a.z - b); }
+static __device__ __forceinline__ uint3     operator*   (unsigned int a, const uint3& b)   { return make_uint3(a * b.x, a * b.y, a * b.z); }
+static __device__ __forceinline__ uint3     operator+   (unsigned int a, const uint3& b)   { return make_uint3(a + b.x, a + b.y, a + b.z); }
+static __device__ __forceinline__ uint3     operator-   (unsigned int a, const uint3& b)   { return make_uint3(a - b.x, a - b.y, a - b.z); }
+static __device__ __forceinline__ uint4&    operator*=  (uint4& a, const uint4& b)         { a.x *= b.x; a.y *= b.y; a.z *= b.z; a.w *= b.w; return a; }
+static __device__ __forceinline__ uint4&    operator+=  (uint4& a, const uint4& b)         { a.x += b.x; a.y += b.y; a.z += b.z; a.w += b.w; return a; }
+static __device__ __forceinline__ uint4&    operator-=  (uint4& a, const uint4& b)         { a.x -= b.x; a.y -= b.y; a.z -= b.z; a.w -= b.w; return a; }
+static __device__ __forceinline__ uint4&    operator*=  (uint4& a, unsigned int b)         { a.x *= b; a.y *= b; a.z *= b; a.w *= b; return a; }
+static __device__ __forceinline__ uint4&    operator+=  (uint4& a, unsigned int b)         { a.x += b; a.y += b; a.z += b; a.w += b; return a; }
+static __device__ __forceinline__ uint4&    operator-=  (uint4& a, unsigned int b)         { a.x -= b; a.y -= b; a.z -= b; a.w -= b; return a; }
+static __device__ __forceinline__ uint4     operator*   (const uint4& a, const uint4& b)   { return make_uint4(a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w); }
+static __device__ __forceinline__ uint4     operator+   (const uint4& a, const uint4& b)   { return make_uint4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); }
+static __device__ __forceinline__ uint4     operator-   (const uint4& a, const uint4& b)   { return make_uint4(a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w); }
+static __device__ __forceinline__ uint4     operator*   (const uint4& a, unsigned int b)   { return make_uint4(a.x * b, a.y * b, a.z * b, a.w * b); }
+static __device__ __forceinline__ uint4     operator+   (const uint4& a, unsigned int b)   { return make_uint4(a.x + b, a.y + b, a.z + b, a.w + b); }
+static __device__ __forceinline__ uint4     operator-   (const uint4& a, unsigned int b)   { return make_uint4(a.x - b, a.y - b, a.z - b, a.w - b); }
+static __device__ __forceinline__ uint4     operator*   (unsigned int a, const uint4& b)   { return make_uint4(a * b.x, a * b.y, a * b.z, a * b.w); }
+static __device__ __forceinline__ uint4     operator+   (unsigned int a, const uint4& b)   { return make_uint4(a + b.x, a + b.y, a + b.z, a + b.w); }
+static __device__ __forceinline__ uint4     operator-   (unsigned int a, const uint4& b)   { return make_uint4(a - b.x, a - b.y, a - b.z, a - b.w); }
+
+template<class T> static __device__ __forceinline__ T zero_value(void);
+template<> __device__ __forceinline__ float  zero_value<float> (void)                      { return 0.f; }
+template<> __device__ __forceinline__ float2 zero_value<float2>(void)                      { return make_float2(0.f, 0.f); }
+template<> __device__ __forceinline__ float4 zero_value<float4>(void)                      { return make_float4(0.f, 0.f, 0.f, 0.f); }
+static __device__ __forceinline__ float3 make_float3(const float2& a, float b)             { return make_float3(a.x, a.y, b); }
+static __device__ __forceinline__ float4 make_float4(const float3& a, float b)             { return make_float4(a.x, a.y, a.z, b); }
+static __device__ __forceinline__ float4 make_float4(const float2& a, const float2& b)     { return make_float4(a.x, a.y, b.x, b.y); }
+static __device__ __forceinline__ int3 make_int3(const int2& a, int b)                     { return make_int3(a.x, a.y, b); }
+static __device__ __forceinline__ int4 make_int4(const int3& a, int b)                     { return make_int4(a.x, a.y, a.z, b); }
+static __device__ __forceinline__ int4 make_int4(const int2& a, const int2& b)             { return make_int4(a.x, a.y, b.x, b.y); }
+static __device__ __forceinline__ uint3 make_uint3(const uint2& a, unsigned int b)         { return make_uint3(a.x, a.y, b); }
+static __device__ __forceinline__ uint4 make_uint4(const uint3& a, unsigned int b)         { return make_uint4(a.x, a.y, a.z, b); }
+static __device__ __forceinline__ uint4 make_uint4(const uint2& a, const uint2& b)         { return make_uint4(a.x, a.y, b.x, b.y); }
+
+template<class T> static __device__ __forceinline__ void swap(T& a, T& b)                  { T temp = a; a = b; b = temp; }
+
+//------------------------------------------------------------------------
+// Triangle ID <-> float32 conversion functions to support very large triangle IDs.
+//
+// Values up to and including 16777216 (also, negative values) are converted trivially and retain
+// compatibility with previous versions. Larger values are mapped to unique float32 that are not equal to
+// the ID. The largest value that converts to float32 and back without generating inf or nan is 889192447.
+
+static __device__ __forceinline__ int   float_to_triidx(float x) { if (x <= 16777216.f) return (int)x;   return __float_as_int(x) - 0x4a800000; }
+static __device__ __forceinline__ float triidx_to_float(int x)   { if (x <= 0x01000000) return (float)x; return __int_as_float(0x4a800000 + x); }
+
+//------------------------------------------------------------------------
+// Coalesced atomics. These are all done via macros.
+
+#if __CUDA_ARCH__ >= 700 // Warp match instruction __match_any_sync() is only available on compute capability 7.x and higher
+
+#define CA_TEMP       _ca_temp
+#define CA_TEMP_PARAM float* CA_TEMP
+#define CA_DECLARE_TEMP(threads_per_block) \
+    __shared__ float CA_TEMP[(threads_per_block)]
+
+#define CA_SET_GROUP_MASK(group, thread_mask)                   \
+    bool   _ca_leader;                                          \
+    float* _ca_ptr;                                             \
+    do {                                                        \
+        int tidx   = threadIdx.x + blockDim.x * threadIdx.y;    \
+        int lane   = tidx & 31;                                 \
+        int warp   = tidx >> 5;                                 \
+        int tmask  = __match_any_sync((thread_mask), (group));  \
+        int leader = __ffs(tmask) - 1;                          \
+        _ca_leader = (leader == lane);                          \
+        _ca_ptr    = &_ca_temp[((warp << 5) + leader)];         \
+    } while(0)
+
+#define CA_SET_GROUP(group) \
+    CA_SET_GROUP_MASK((group), 0xffffffffu)
+
+#define caAtomicAdd(ptr, value)         \
+    do {                                \
+        if (_ca_leader)                 \
+            *_ca_ptr = 0.f;             \
+        atomicAdd(_ca_ptr, (value));    \
+        if (_ca_leader)                 \
+            atomicAdd((ptr), *_ca_ptr); \
+    } while(0)
+
+#define caAtomicAdd3_xyw(ptr, x, y, w)  \
+    do {                                \
+        caAtomicAdd((ptr), (x));        \
+        caAtomicAdd((ptr)+1, (y));      \
+        caAtomicAdd((ptr)+3, (w));      \
+    } while(0)
+
+#define caAtomicAddTexture(ptr, level, idx, value)  \
+    do {                                            \
+        CA_SET_GROUP((idx) ^ ((level) << 27));      \
+        caAtomicAdd((ptr)+(idx), (value));          \
+    } while(0)
+
+//------------------------------------------------------------------------
+// Disable atomic coalescing for compute capability lower than 7.x
+
+#else // __CUDA_ARCH__ >= 700
+#define CA_TEMP _ca_temp
+#define CA_TEMP_PARAM float CA_TEMP
+#define CA_DECLARE_TEMP(threads_per_block) CA_TEMP_PARAM
+#define CA_SET_GROUP_MASK(group, thread_mask)
+#define CA_SET_GROUP(group)
+#define caAtomicAdd(ptr, value) atomicAdd((ptr), (value))
+#define caAtomicAdd3_xyw(ptr, x, y, w)  \
+    do {                                \
+        atomicAdd((ptr), (x));          \
+        atomicAdd((ptr)+1, (y));        \
+        atomicAdd((ptr)+3, (w));        \
+    } while(0)
+#define caAtomicAddTexture(ptr, level, idx, value) atomicAdd((ptr)+(idx), (value))
+#endif // __CUDA_ARCH__ >= 700
+
+//------------------------------------------------------------------------
+#endif // __CUDACC__

extensions/nvdiffrast/nvdiffrast/common/cudaraster/CudaRaster.hpp

@@ -0,0 +1,63 @@
+// Copyright (c) 2009-2022, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+
+//------------------------------------------------------------------------
+// This is a slimmed-down and modernized version of the original
+// CudaRaster codebase that accompanied the HPG 2011 paper
+// "High-Performance Software Rasterization on GPUs" by Laine and Karras.
+// Modifications have been made to accommodate post-Volta execution model
+// with warp divergence. Support for shading, blending, quad rendering,
+// and supersampling have been removed as unnecessary for nvdiffrast.
+//------------------------------------------------------------------------
+
+namespace CR
+{
+
+class RasterImpl;
+
+//------------------------------------------------------------------------
+// Interface class to isolate user from implementation details.
+//------------------------------------------------------------------------
+
+class CudaRaster
+{
+public:
+    enum
+    {
+        RenderModeFlag_EnableBackfaceCulling = 1 << 0,   // Enable backface culling.
+        RenderModeFlag_EnableDepthPeeling    = 1 << 1,   // Enable depth peeling. Must have a peel buffer set.
+    };
+
+public:
+					        CudaRaster				(void);
+					        ~CudaRaster				(void);
+
+    void                    setBufferSize           (int width, int height, int numImages);              // Width and height are internally rounded up to multiples of tile size (8x8) for buffer sizes.
+    void                    setViewport             (int width, int height, int offsetX, int offsetY);   // Tiled rendering viewport setup.
+    void                    setRenderModeFlags      (unsigned int renderModeFlags);                      // Affects all subsequent calls to drawTriangles(). Defaults to zero.
+    void                    deferredClear           (unsigned int clearColor);                           // Clears color and depth buffers during next call to drawTriangles().
+    void                    setVertexBuffer         (void* vertices, int numVertices);                   // GPU pointer managed by caller. Vertex positions in clip space as float4 (x, y, z, w).
+    void                    setIndexBuffer          (void* indices, int numTriangles);                   // GPU pointer managed by caller. Triangle index+color quadruplets as uint4 (idx0, idx1, idx2, color).
+    bool                    drawTriangles           (const int* ranges, bool peel, cudaStream_t stream); // Ranges (offsets and counts) as #triangles entries, not as bytes. If NULL, draw all triangles. Returns false in case of internal overflow.
+    void*                   getColorBuffer          (void);                                              // GPU pointer managed by CudaRaster.
+    void*                   getDepthBuffer          (void);                                              // GPU pointer managed by CudaRaster.
+    void                    swapDepthAndPeel        (void);                                              // Swap depth and peeling buffers.
+
+private:
+					        CudaRaster           	(const CudaRaster&); // forbidden
+	CudaRaster&             operator=           	(const CudaRaster&); // forbidden
+
+private:
+    RasterImpl*             m_impl;                 // Opaque pointer to implementation.
+};
+
+//------------------------------------------------------------------------
+} // namespace CR
+

extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/BinRaster.inl

@@ -0,0 +1,423 @@
+// 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<<threadIdx.x))
+                    {
+                        int idx = __popc(outMask & getLaneMaskLt());
+                        if (threadIdx.y > 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];
+}
+
+//------------------------------------------------------------------------

extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.cpp

@@ -0,0 +1,94 @@
+// Copyright (c) 2009-2022, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include "../../framework.h"
+#include "Buffer.hpp"
+
+using namespace CR;
+
+//------------------------------------------------------------------------
+// GPU buffer.
+//------------------------------------------------------------------------
+
+Buffer::Buffer(void)
+:   m_gpuPtr(NULL),
+    m_bytes (0)
+{
+    // empty
+}
+
+Buffer::~Buffer(void)
+{
+    if (m_gpuPtr)
+        cudaFree(m_gpuPtr); // Don't throw an exception.
+}
+
+void Buffer::reset(size_t bytes)
+{
+    if (bytes == m_bytes)
+        return;
+
+    if (m_gpuPtr)
+    {
+        NVDR_CHECK_CUDA_ERROR(cudaFree(m_gpuPtr));
+        m_gpuPtr = NULL;
+    }
+
+    if (bytes > 0)
+        NVDR_CHECK_CUDA_ERROR(cudaMalloc(&m_gpuPtr, bytes));
+
+    m_bytes = bytes;
+}
+
+void Buffer::grow(size_t bytes)
+{
+    if (bytes > m_bytes)
+        reset(bytes);
+}
+
+//------------------------------------------------------------------------
+// Host buffer with page-locked memory.
+//------------------------------------------------------------------------
+
+HostBuffer::HostBuffer(void)
+:   m_hostPtr(NULL),
+    m_bytes  (0)
+{
+    // empty
+}
+
+HostBuffer::~HostBuffer(void)
+{
+    if (m_hostPtr)
+        cudaFreeHost(m_hostPtr); // Don't throw an exception.
+}
+
+void HostBuffer::reset(size_t bytes)
+{
+    if (bytes == m_bytes)
+        return;
+
+    if (m_hostPtr)
+    {
+        NVDR_CHECK_CUDA_ERROR(cudaFreeHost(m_hostPtr));
+        m_hostPtr = NULL;
+    }
+
+    if (bytes > 0)
+        NVDR_CHECK_CUDA_ERROR(cudaMallocHost(&m_hostPtr, bytes));
+
+    m_bytes = bytes;
+}
+
+void HostBuffer::grow(size_t bytes)
+{
+    if (bytes > m_bytes)
+        reset(bytes);
+}
+
+//------------------------------------------------------------------------

extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/Buffer.hpp

@@ -0,0 +1,55 @@
+// Copyright (c) 2009-2022, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA CORPORATION and its licensors retain all intellectual property
+// and proprietary rights in and to this software, related documentation
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#pragma once
+#include "Defs.hpp"
+
+namespace CR
+{
+//------------------------------------------------------------------------
+
+class Buffer
+{
+public:
+                    Buffer      (void);
+                    ~Buffer     (void);
+
+    void            reset       (size_t bytes);
+    void            grow        (size_t bytes);
+    void*           getPtr      (size_t offset = 0) { return (void*)(((uintptr_t)m_gpuPtr) + offset); }
+    size_t          getSize     (void) const { return m_bytes; }
+
+    void            setPtr      (void* ptr) { m_gpuPtr = ptr; }
+
+private:
+    void*           m_gpuPtr;
+    size_t          m_bytes;
+};
+
+//------------------------------------------------------------------------
+
+class HostBuffer
+{
+public:
+                    HostBuffer  (void);
+                    ~HostBuffer (void);
+
+    void            reset       (size_t bytes);
+    void            grow        (size_t bytes);
+    void*           getPtr      (void) { return m_hostPtr; }
+    size_t          getSize     (void) const { return m_bytes; }
+
+    void            setPtr      (void* ptr) { m_hostPtr = ptr; }
+
+private:
+    void*           m_hostPtr;
+    size_t          m_bytes;
+};
+
+//------------------------------------------------------------------------
+}

extensions/nvdiffrast/nvdiffrast/common/cudaraster/impl/CoarseRaster.inl

@@ -0,0 +1,730 @@
+// 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);
+            }
+        }
+    }
+}
+
+//------------------------------------------------------------------------