PTX contains invalid parameter load

[mworchel]

Summary

I have (quite complex) Python code that nests multiple symbolic loops and on some occasions (more details below) Dr.Jit produces invalid PTX that leads to the following error:

COMPILE ERROR: Invalid PTX input: ptx2llvm-module-003: error: Failed to parse input PTX string
ptx2llvm-module-003, line 6238; error   : State space mismatch between instruction and address in instruction 'ld'
ptx2llvm-module-003, line 6242; error   : State space mismatch between instruction and address in instruction 'ld'
ptx2llvm-module-003, line 6358; error   : State space mismatch between instruction and address in instruction 'ld'
ptx2llvm-module-003, line 6361; error   : State space mismatch between instruction and address in instruction 'ld'
Cannot parse input PTX string

The offending line(s) in PTX read like this:

.visible .func (.param .align 4 .b8 result[36]) __direct_callable__c7d274ed0e138d86d000b2ea00bb0bcc(.reg .u64 data) {
    // Call: mitsuba::Emitter::sample_ray()
    .reg.b8   %b <177>; .reg.b16  %w<177>; .reg.b32 %r<177>;
    .reg.b64  %rd<177>; .reg.f16  %h<177>; .reg.f32 %f<177>;
    .reg.f64  %d <177>; .reg.pred %p<177>;
    [...]
    ld.param.b32 %f22, [params+65535]; // <--------- invalid load
    mov.b32 %f23, 0xbf800000;
    fma.rn.ftz.f32 %f24, %f21, %f22, %f23;
    abs.f32 %f25, %f24;
    ld.param.b32 %f26, [params+65535]; // <--------- invalid load
    [...]

There's a parameter load although it seems no parameters are passed.

Details

The code is intersection logic for a custom Mitsuba shape that is invoked inside the active integrator (I have wrapped the scene and minimally adapted the Python integrators to support the new shape in Python). This works great for prb and other custom integrators, both the primal rendering and forward/backward-mode differentiation. But, there is one custom integrator (basically a Python port of the particle tracer) where the primal rendering works as expected but backward-mode differentiation produces the above PTX code (calling dr.backward_from(dr.sum(img)))

I stepped through the PTX generation and found that a CallInput variable v depends on a variable a that neither has a valid reg_index nor param_offset assigned. Basically, the following happens:

1. a is assigned 0xFFFF (65535) as param_offset at eval.cpp#L351.
2. reg_index of a remains unassigned and therefore it is deemed unused in call.cpp#L480, which leaves param_offset at 0xFFFF.
3. The invalid load is generated at cuda_eval.cpp#L301.

I am not 100% sure at which point the state becomes invalid. Unfortunately, I could not create a self-contained reproducer and I cannot (yet) share the code. Interestingly, before the invalid compilation, the same call launches a kernel with a very similar function that is correctly assembled:

.visible .func (.param .align 4 .b8 result[36]) __direct_callable__7e02206b58dbcab2fbe0b1545125e371(.reg .u64 data, .param .align 4 .b8 params[16]) {
    // Call: mitsuba::Emitter::sample_ray()
    .reg.b8   %b <177>; .reg.b16  %w<177>; .reg.b32 %r<177>;
    .reg.b64  %rd<177>; .reg.f16  %h<177>; .reg.f32 %f<177>;
    .reg.f64  %d <177>; .reg.pred %p<177>;
    [...]
    ld.param.b32 %f22, [params+8];
    mov.b32 %f23, 0xbf800000;
    fma.rn.ftz.f32 %f24, %f21, %f22, %f23;
    abs.f32 %f25, %f24;
    ld.param.b32 %f26, [params+12];
    [...]

I think it would be helpful to get some pointers on where I should keep digging. Thanks!

[merlinND]

Hello @mworchel,

Just a shot in the dark, but could you please try with this fix from @DoeringChristian that was just merged?
2d41937 (PR #129)

[njroussel]

Hi @mworchel

It's maybe worth going through the steps described in this document. Specifically, turning on the sanitizers might catch something.

Also, try turning off OptimizeCalls. It does sound like the bug could be a consequence of that feature.
If that didn't help, then it most likely will require some deep digging in Dr.Jit.

---

For this problem here's how I'd go about it:

• Disable parallel scene loading (mi.load_dict(..., parallel=False))
• Set dr.set_flag(dr.JitFlag.ReuseIndices, False)

This will make JIT indices consistent and easier to track across multiple sessions.

The key is to figure out why the variable a isn't assigned a reg_index during the main assembly in jitc_eval_impl (eval.cpp:676). This is how it should happen:

• Your main code requests to evaluate some variable out
• jitc_eval_impl() will call jitc_var_traverse on out to find all its dependencies.
• During the traversal, it will find a CallOutput - one of the sample_ray outputs.
• When traversing that output, it will go through all the subgraphs for each emitter instance's sample_ray function in your scene (eval.cpp:187).
• Those individual traversals should end up finding the CallInput variable.
• The CallInput variable should add a to list of visited variables (eval.cpp:165) and list of scheduled variables (eval.cpp:221).
• A register index reg_index is assigned to every scheduled variable eval.cpp:302.

If you don't see that last step happening, then it should just be a matter of figuring out what's wrong during the traversal phase. If you do see it happening, something terribly wrong is happening between there and the Call assembly.

[mworchel]

Hi,

@merlinND unfortunately that didn't fix the issue but thanks for the pointer.

@njroussel thanks for the details. I debugged a little more earlier today and I think it aligns nicely with what you suggest:

1. Without OptimizeCalls there is no invalid PTX error but the detail::check_grad_enabled("backward_from", value); fails and results in a crash. This might be another issue so I suspect the main problem is calls being optimized.
2. There is a CallOutput variable that belongs to sample_ray, which is encountered during the traversal. This variable results in the Call being scheduled. But: when the inner (function) variables are traversed, the traversal never reaches any CallInput because (as far as I have seen) all dependent variables "inside" the function are either Literal or Evaluated. This aligns with the observation that a (working) call to the function is assembled earlier. Should the CallOutput also become Evaluated? I am wondering why the traversal enters the function again if any variable the output depends on is either evaluated or literal.

So, basically, the CallInputs are never scheduled.

EDIT: Setting the integrator loop to mode='evaluated' (with OptimizeCalls=False) appears to work and gradients are computed.

[wjakob]

Hi @mworchel, to really help you with this, we will need a reproducer. Ideally pure Dr.Jit -- for example, you can trigger the same virtual function recording mechanism using dr.switch()

[wjakob]

Hi @mworchel -- would you kindly be able to produce a reproducer for this issue (it's fine for it to even be a big one, no need to make a minimal repro).

[mworchel]

Hi @wjakob,

that's probably the easiest way, because I'm not sure what exactly it takes to recreate the issue in pure drjit.

I can send you a reproducer but its currently internal code so I would prefer sending it by mail, if that's ok.

[wjakob]

Ok, that works -- thank you!