deploying pre-compiled kernels (C++)

[lfbuchmann]

Intro

I would like to use mitsuba3 / drJit as a part of an integrated C++ pipeline. The aim is to infer properties of distinct objects that are captured in identical settings (lights, sensors, positions). The algorithm can be split up in two steps:

1. Compile the differentiable rendering of the scene with an initial guess using Mitsuba 3 in Python, resulting in a CUDA/LLVM kernel.
2. Run a learning cycle in C++ for a number of steps with a given reference image.

the key is to be able to apply step 2 repeatedly with new reference images, without having to re-compile the kernel and without invoking Python.

Minimal working example

Consider the following script:

Import necessary packages:

import mitsuba as mi
mi.set_variant('cuda_ad_rgb')

import numpy as np
import drjit as dr

Create the scene:

diffuseTexturedSphereDict = {
    'type' : 'sphere',
    'material':{'type':'diffuse',
             'reflectance':{'type':'bitmap', 'bitmap':initBitmap, 'raw':True}}
    }

floorDict = {
    'type': 'rectangle',
    'to_world':mi.ScalarTransform4f.scale([10,10,1]).translate([0,0,-1]),
    'material': {
        'type': 'diffuse',
            'reflectance': 0.9
        }
    }

pointLightDict = {'type':'point',
             'position' : [5,7,7],
            'intensity':{
                        'type':'rgb',
                        'value':100.0
            }}

pointLightDict2 = {'type':'point',
             'position' : [7,5,7],
            'intensity':{
                        'type':'rgb',
                        'value':100.0
            }}

pointLightDict3 = {'type':'point',
             'position' : [0,0,15],
            'intensity':{
                        'type':'rgb',
                        'value':100.0
            }}

sensorDict = {'type':'perspective',
              'to_world':mi.ScalarTransform4f.look_at(
                origin = [5,5,7.5],
                target = [0,0,0],
                up=[0,0,1]
              ),
              'film':{'type':'hdrfilm',
                      'width':640,
                      'height':480},
              #        'sampler':{'type':'independent'}
              'sampler':{'type':'ldsampler',
                         'sample_count':64}
}

integratorDict = {'type': 'prb','max_depth':2}

sceneDict = {'type':'scene'}

sceneDict['floor'] = floorDict
sceneDict['light1'] = pointLightDict
sceneDict['light2'] = pointLightDict2
sceneDict['light3'] = pointLightDict3
sceneDict['sensor']=sensorDict
sceneDict['integrator'] = integratorDict
sceneDict['sphere'] = diffuseTexturedSphereDict
scene = mi.load_dict(sceneDict)

Load a reference image, saved on disk as a numpy array :

ar = np.load('referenceArray.npy')
pic = mi.TensorXf(ar)

The image contained in the array is the same scene, but with a uniformly red sphere. It can be obtained by rendering the scene with sceneDict['sphere'] replaced by

diffuseSphereDict = {
    'type' : 'sphere',
    'material':{'type':'diffuse',
             'reflectance':{'type':'rgb', 'value':[0.8,0,0]}}
    }

Now we set up a learning cycle to learn the red color of the sphere as a texture:

params = mi.traverse(scene)
key = 'sphere.bsdf.reflectance.data'
params.keep(key)
mi.util.convert_to_bitmap(params[key]).write('initialTexture.png')
initPic = mi.render(scene, params, spp = 256)
mi.util.convert_to_bitmap(initPic).write('initialPic.png')
opt = mi.ad.Adam(lr = 5e-2, mask_updates = True)
opt[key] = params[key]
params.update(opt)
steps = 100

def mse(render):
    return dr.mean(dr.sqr(render - pic))

for it in range(steps):
    newPic = mi.render(scene, params, spp = 16)
    loss = mse(newPic)
    dr.backward(loss)
    opt.step()
    dr.clamp(opt[key], 0, 1)
    params.update(opt)
    print(f"Iteration {it:02d}, loss: {loss[0]:5.2e}", end='\n')

pic = mi.render(scene, params, spp=256)
mi.util.convert_to_bitmap(pic).write('finalPic.png')
mi.util.convert_to_bitmap(params[key]).write('finalTexture.png')

The output consist of the initial texture:

the initial rendering:

the texture after 100 iterations:

and the resulting rendering:

Checking the folder (AppData\Temp\drjit on windows) reveals 7 files with suffixes .cuda.bin and 3 "optix7cache.db*" files.

I now would like to write a C++ program that deploys the compiled kernels with a different choice of referenceArray without having to invoke python and recompile the kernels.

Is this possible?

[njroussel]

Hi @lfbuchmann

For anyone coming across this thread, please also refer to the end of this discussion.

As I briefly mentioned in the thread above, I believe that this is reasonably achievable. It is definitely possible out-of-the box, you will need to understand and modify some drjit-core internals.

Here are some pointers/tips:

• By increasing the log level to INFO (dr.set_log_level(dr.LogLevel.Info)) you should see logs regarding the launches of kernels. What we're interested in here are the kernel hashes (identifiers). You'll also see some information regarding whether or not there was a cache hit/miss for the kernel.
• The jitc_run function is where kernels are launched. This is a good starting point to see how the cache is used to launch an existing kernel again. You essentially want to reproduce this behaviour.
• This PR is a fantastic starting point. Although it most likely doesn't work anymore (it was implemented a while ago), the core idea is still valid. One important information is that kernels re-order their input arguments for efficiency, you'll see how that is dealt with in the PR.

We have a need/desire for a generalized solution to this problem of re-launching kernels, albeit it not being our top priority at the moment.