Integrator failures on cuda with pytorch integrations

[lopsided]

Summary

I am working on shape optimisation with a mesh built in pytorch. I want to experiment with some of the other integrators and options, but most of them fail to run on the gpu.

System configuration

System information:

CPU: 13th Gen Intel(R) Core(TM) i9-13950HX
GPU: NVIDIA GeForce RTX 4090 Laptop GPU
Python: 3.9.18 | packaged by conda-forge | (main, Dec 23 2023, 16:33:10) [GCC 12.3.0]
NVidia driver: 545.29.06
LLVM: 16.0.6

Dr.Jit: 0.4.4
Mitsuba: 3.5.0
Is custom build? False
Compiled with: GNU 10.2.1
Variants:
scalar_rgb
scalar_spectral
cuda_ad_rgb
llvm_ad_rgb

Description

My example here #1086 is working OK, but my actual problem is harder and is struggling to optimise well. (I think possibly because I'm using a roughdielectric bsdf - it works much better with a diffuse material). I want to experiment with some of the other integrators to see if they might give better results, but most of these just crash. I've included a full test script below and included the success/fail state of all combinations of mitsuba/pytorch optimisation, the different integrators, and whether it's using the cpu/gpu.

To summarise my findings:

• The prbvolpath integrator fails in all cases.
• Using mitsuba without pytorch works fine for all integrators (apart from prbvolpath) and either on CPU or GPU.
• Using the pytorch integration on the CPU everything continues to work fine.
• Using the pytorch integration on the GPU causes errors in most cases.
• For translucent shapes, the direct_* integrators don't work well for roughdielectric bsdfs (to be expected I think..).
• On GPU, the *_projective integrators all require 'sppi': 0 otherwise they fail.

Questions:

• For shape optimisation, do I need to use either direct_projective or prb_projective? Some of the integrators say they aren't suitable (eg, prb) but others don't mention it (eg, path).
• For roughdielectric bsdfs, I can't use direct_*?
• Is the sppi parameter important when using roughdielectric bsdfs?

Thanks!

Steps to reproduce

from typing import Tuple

import drjit as dr
import matplotlib.pyplot as plt
import mitsuba as mi
import torch
from pytorch3d.utils import ico_sphere
from torch import nn

USE_CUDA = 1

if USE_CUDA:
    if 'cuda_ad_rgb' not in mi.variants():
        raise RuntimeError('No CUDA variant found.')
    mi.set_variant('cuda_ad_rgb')
    device = torch.device('cuda')

else:
    mi.set_variant('llvm_ad_rgb')
    device = torch.device('cpu')

from mitsuba import ScalarTransform4f as T

SUZANNE_PATH = 'tmp' / 'meshes' / 'suzanne.ply'
SHAPE_NAME = 'sphere'
VERTEX_KEY = SHAPE_NAME + '.vertex_positions'
FACES_KEY = SHAPE_NAME + '.faces'
BSDF_KEY = SHAPE_NAME + '.bsdf'


class Sphere(nn.Module):
    def __init__(self, scale: float = 1.0):
        """
        Create a sphere with the given scale.
        """
        super().__init__()
        self.scale = nn.Parameter(torch.tensor(scale, dtype=torch.float32), requires_grad=True)

    def build_mesh(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Create a sphere mesh with the given scale.
        """
        sphere = ico_sphere(level=1, device=self.scale.device)
        vertices = sphere.verts_packed()
        faces = sphere.faces_packed()
        vertices = vertices * self.scale
        return vertices, faces


def build_mitsuba_mesh(shape: Sphere, bsdf: dict) -> mi.Mesh:
    """
    Convert a Sphere object into a Mitsuba mesh.
    """
    # Build the mesh in pytorch and convert the parameters to Mitsuba format
    vertices, faces = shape.build_mesh()
    nv, nf = len(vertices), len(faces)
    vertices = mi.TensorXf(vertices)
    faces = mi.TensorXi64(faces)

    # Set up the material properties
    props = mi.Properties()
    props['bsdf'] = mi.load_dict(bsdf)

    # Construct the mitsuba mesh and set the vertex positions and faces
    mesh = mi.Mesh(
        SHAPE_NAME,
        vertex_count=nv,
        face_count=nf,
        has_vertex_normals=False,
        has_vertex_texcoords=False,
        props=props
    )
    mesh_params = mi.traverse(mesh)
    mesh_params['vertex_positions'] = dr.ravel(vertices)
    mesh_params['faces'] = dr.ravel(faces)

    return mesh


def create_scene(shape, spp, res, integrator, backlight=False):
    return mi.load_dict({
        'type': 'scene',
        'integrator': integrator,
        'sensor': {
            'type': 'perspective',
            'fov': 45,
            'to_world': T.look_at(target=[0, 0, 0], origin=[0, 0, 10], up=[0, 1, 0]),
            'sampler': {
                'type': 'independent',
                'sample_count': spp
            },
            'film': {
                'type': 'hdrfilm',
                'width': res,
                'height': res,
                'filter': {'type': 'gaussian'},
                'sample_border': True,
            },
        },
        'light': {
            'type': 'point',
            'to_world': T.look_at(origin=[0, 0, 10], target=[0, 0, 0], up=[0, 1, 0]),
            'intensity': {
                'type': 'spectrum',
                'value': 300.0,
            }
        },
        SHAPE_NAME: shape
    })


def shape_optimisation(
        opt_with: str = 'mitsuba',
        integrator: dict = {'type': 'path'}
):
    print(f'----- Running optimisation for {opt_with} and integrator {integrator} ------')

    # Parameters
    spp = 64
    res = 100
    lr = 0.1
    n_iterations = 2

    @dr.wrap_ad(source='torch', target='drjit')
    def render_image_wrapper(vertices, faces, seed=1):
        params[VERTEX_KEY] = dr.ravel(vertices)
        params[FACES_KEY] = dr.ravel(faces)
        params.update()
        return mi.render(scene, params, seed=seed)

    # Target shape
    bsdf = {'type': 'diffuse'}
    target_shape = {
        'type': 'ply',
        'filename': str(SUZANNE_PATH),
        'bsdf': bsdf
    }

    # Optimisable shape
    sphere = Sphere(scale=2)
    sphere.to(device)

    # Make scenes
    scene_target = create_scene(target_shape, spp, res, integrator)
    scene = create_scene(build_mitsuba_mesh(sphere, bsdf), spp, res, integrator)
    params = mi.traverse(scene)
    img_target = mi.render(scene_target)
    img_init = mi.render(scene)

    # Optimise
    if opt_with == 'mitsuba':
        opt = mi.ad.Adam(lr=lr)
        opt[VERTEX_KEY] = params[VERTEX_KEY]
    else:
        img_target = img_target.torch()
        img_init = img_init.torch()
        opt = torch.optim.Adam(sphere.parameters(), lr=lr, weight_decay=0)

    for i in range(n_iterations):
        if opt_with == 'mitsuba':
            params.update(opt)
            img_i = mi.render(scene, params, seed=i)
            loss = dr.mean(dr.sqr(img_i - img_target))
            loss_val = float(loss[0])
            dr.backward(loss)
            opt.step()
        else:
            opt.zero_grad()
            v, f = sphere.build_mesh()
            img_i = render_image_wrapper(v, f, seed=i)
            loss = ((img_i - img_target)**2).mean()
            loss_val = loss.item()
            loss.backward()
            opt.step()
        print(f'Iteration {1 + i}: Loss = {loss_val:6f}')

    # Plot
    fig, axes = plt.subplots(1, 3, figsize=(4, 2))
    for ax, img in zip(axes, [img_target, img_init, img_i]):
        if isinstance(img, torch.Tensor):
            img = img.detach().cpu().numpy()
        else:
            img = img.numpy()
        img = img.clip(0, 1)
        ax.imshow(img)
        ax.axis('off')
    fig.tight_layout()
    plt.show()
    plt.close(fig)

    print(f'---- Optimisation complete for {opt_with} and integrator {integrator} ----')
    print('')


if __name__ == '__main__':
    # ---- CPU + Mitsuba
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'direct'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'path'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'volpath'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'volpathmis'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'prb'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'direct_projective'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'prb_projective'})
    shape_optimisation(opt_with='mitsuba', integrator={'type': 'prbvolpath'}) # E1

    # ---- CPU + PyTorch
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'direct'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'path'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'volpath'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'volpathmis'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'prb'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'direct_projective'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'direct_projective', 'sppi': 0})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'prb_projective'})
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'prb_projective', 'sppi': 0})
    shape_optimisation(opt_with='pytorch', integrator={'type': 'prbvolpath'}) # E1

    # ---- CUDA + Mitsuba
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'direct'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'path'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'volpath'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'volpathmis'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'prb'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'direct_projective'})
    # shape_optimisation(opt_with='mitsuba', integrator={'type': 'prb_projective'})
    shape_optimisation(opt_with='mitsuba', integrator={'type': 'prbvolpath'}) # E1

    # ---- CUDA + PyTorch
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'direct'})
    shape_optimisation(opt_with='pytorch', integrator={'type': 'path'})  # E2
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'volpath'})
    shape_optimisation(opt_with='pytorch', integrator={'type': 'volpathmis'}) # E2
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'prb'})
    shape_optimisation(opt_with='pytorch', integrator={'type': 'direct_projective'})  # E0
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'direct_projective', 'sppi': 0})
    shape_optimisation(opt_with='pytorch', integrator={'type': 'prb_projective'}) # E0
    # shape_optimisation(opt_with='pytorch', integrator={'type': 'prb_projective', 'sppi': 0})
    shape_optimisation(opt_with='pytorch', integrator={'type': 'prbvolpath'}) # E1

'''
E0: Process finished with exit code 139 (interrupted by signal 11:SIGSEGV)
E1: drjit.Exception: loop_process_state(): one of the supplied loop state variables of type Float is attached to the AD graph (i.e., grad_enabled(..) is true). However, propagating derivatives through multiple iterations of a recorded loop is not supported (and never will be). Please see the documentation on differentiating loops for details and suggested alternatives.
E2: Critical Dr.Jit compiler failure: cuda_check(): API error 0700 (CUDA_ERROR_ILLEGAL_ADDRESS): "an illegal memory access was encountered" in /project/ext/drjit-core/src/init.cpp:454.
'''