[BUG] CUDA crash with tile_reduce in debug mode

[Kenny-Vilella]

Bug Description

Please see below for a repro.
The issue appears somewhere between 1.9.0.dev20250812 and 1.9.0.dev20250813.
I checked that it works well for 1.9.0.dev20250812 and earlier but not with 1.9.0.dev20250813 and after (including 1.10).
The issue disappears if I do not import mujoco_warp or If I do not call tile_reduce.
I added enable_backward = False to make sure we are using exactly the same kernel with and without importing mujoco_warp but this option does not seem to make a difference.

It may be an option problem with set by mujoco_warp ?

Repro:

import warp as wp
# Removing this stop the issue
import mujoco_warp as mjwarp

def _repro_tile_reduce(nv: int):
  @wp.kernel
  def repro_tile_reduce(
    qvel_in: wp.array2d(dtype=float),
    Mqvel: wp.array2d(dtype=float),
  ):
    worldid = wp.tid()

    qvel_tile = wp.tile_load(qvel_in[worldid], shape=wp.static(nv))
    Mqvel_tile = wp.tile_load(Mqvel[worldid], shape=wp.static(nv))
    qvelMqvel_tile = wp.tile_map(wp.mul, qvel_tile, Mqvel_tile)
    # Below is causing issue
    quadratic_tile = wp.tile_reduce(wp.add, qvelMqvel_tile)

  return repro_tile_reduce


wp.config.mode = "debug"
wp.config.enable_backward = False
wp.init()

qvel=wp.zeros((1, 27), dtype=float)
mv = wp.zeros((1, 27), dtype=float)

wp.launch_tiled(
  _repro_tile_reduce(27),
  dim=(1),
  inputs=[qvel, mv],
  block_dim=32,
)

System Information

Run on Horde:

Warp 1.9.0.dev20250813 initialized:
Git commit: f38f51c
CUDA Toolkit 12.8, Driver 12.4
Devices:
"cpu" : "x86_64"
"cuda:0" : "NVIDIA L40" (47 GiB, sm_89, mempool enabled)

[Kenny-Vilella]

This is the crash message:

Warp 1.9.0.dev20250813 initialized:
Git commit: f38f51c
CUDA Toolkit 12.8, Driver 12.4
Devices:
"cpu" : "x86_64"
"cuda:0" : "NVIDIA L40" (47 GiB, sm_89, mempool enabled)
Kernel cache:
/home/horde/.cache/warp/1.9.0.dev20250813
Module main 56649fd load on device 'cuda:0' took 6.42 ms (cached)
Warp CUDA error 715: an illegal instruction was encountered (in function wp_free_device_async, /builds/omniverse/warp/warp/native/warp.cu:659)
Warp CUDA error 715: an illegal instruction was encountered (in function wp_free_device_async, /builds/omniverse/warp/warp/native/warp.cu:659)

[shi-eric]

I don't see a change in behavior with or without the mujoco_warp import. Thanks for cleverly using the nightly builds to bracket the behavioral change. It changed with my commit c4d0c70. which passes the --device-debug aka -G flag to NVRTC when we runtime-compile modules on Linux (this was done so we can place breakpoints inside CUDA kernels). Before thie change, we did not compile device kernels with debug information.

I didn't see the message that @daedalus5 saw:

inlinable function call in a function with debug info must have a !dbg location
  %call = call float @_ZN2wp3mulEff(float %tmp3, float %tmp4)

I don't think this is a LLVM issue since LLVM is not used in this code path...

[shi-eric]

Trying to figure out where this sm70 instruction came from. I am using cuda-gdb with the sample you provided.

Warp 1.10.0.dev0 initialized:
   CUDA Toolkit 13.0, Driver 13.0
   Devices:
     "cpu"      : "x86_64"
     "cuda:0"   : "NVIDIA L40S" (47 GiB, sm_89, mempool enabled)
   Kernel cache:
     /home/horde/.cache/warp/1.10.0.dev0
[New Thread 0x7fffe15f96c0 (LWP 347303)]
Module __main__ 568795a load on device 'cuda:0' took 226.19 ms  (compiled)

CUDA Exception: Warp Illegal Instruction
The exception was triggered at PC 0x7fffa7282480  __cuda_sm70_shflsync_down_p

Thread 1 "python" received signal CUDA_EXCEPTION_4, Warp Illegal Instruction.
[Switching focus to CUDA kernel 0, grid 1, block (0,0,0), thread (0,0,0), device 0, sm 0, warp 0, lane 0]
0x00007fffa72824a0 in __cuda_sm70_shflsync_down_p ()

[Kenny-Vilella]

It's probably not important.
But I tried and confirm that the mujoco_warp import is affecting the results on Horde.
I created a fresh instance and did this exact steps:

sudo apt-get update
sudo apt install python3.10-venv
python -m venv env
source env/bin/activate
pip install --upgrade pip
python -m  pip install -U --pre warp-lang --extra-index-url=https://pypi.nvidia.com/
# Below in running fine without import mujoco_warp
python -m  pip install mujoco --pre --upgrade -f https://py.mujoco.org/
python -m pip install git+https://github.com/google-deepmind/mujoco_warp.git@main
# Below is crashing with import mujoco_warp
python repro.py

Not sure that it is relevant though.

[shi-eric]

In case this helps for the future, this was what I was seeing. Maybe it's a NVRTC compiler bug.

(cuda-gdb) bt
#0  0x00007fffb12824a0 in __cuda_sm70_shflsync_down_p ()
#1  0x00007fffb12819a0 in __shfl_down_sync(unsigned int, unsigned int, unsigned int, int) ()
#2  0x00007fffb1280b20 in wp::warp_shuffle_down<float> (val=0, offset=16, 
    mask=134217727)

[shi-eric]

It's probably not important. But I tried and confirm that the mujoco_warp import is affecting the results on Horde. I created a fresh instance and did this exact steps:

sudo apt-get update
sudo apt install python3.10-venv
python -m venv env
source env/bin/activate
pip install --upgrade pip
python -m  pip install -U --pre warp-lang --extra-index-url=https://pypi.nvidia.com/
# Below in running fine without import mujoco_warp
python -m  pip install mujoco --pre --upgrade -f https://py.mujoco.org/
python -m pip install git+https://github.com/google-deepmind/mujoco_warp.git@main
# Below is crashing with import mujoco_warp
python repro.py

Not sure that it is relevant though.

Just adding a little more info... I switched to another Horde instance that had a 12.4 driver and was able to observe this behavior where removing the MJWarp import resulted in no errors being reported, but if I add a wp.synchronize_device() to the end of your repro script, I see Warp CUDA error 715: an illegal instruction was encountered (in function wp_cuda_context_synchronize, /builds/omniverse/warp/warp/native/warp.cu:2072) regardless if MJWarp is imported or not.

[Kenny-Vilella]

Yeap you are right and compute sanitizer gives much more information:

========= Illegal instruction
=========     at __cuda_sm70_shflsync_down_p+0x80
=========     by thread (0,0,0) in block (0,0,0)
=========         Device Frame: __shfl_down_sync(unsigned int, unsigned int, unsigned int, int)+0x190
=========         Device Frame: T1 wp::warp_shuffle_down<float>(T1, int, int)+0x610 in tile_reduce.h:72
=========         Device Frame: T1 wp::warp_reduce<float, _repro_tile_reduce__locals__repro_tile_reduce_b774bff3_cuda_kernel_forward::[lambda(T1, T2) (instance 2)]>(T1, T2, unsigned int)+0xb90 in tile_reduce.h:123
=========         Device Frame: auto wp::tile_reduce_impl<wp::tile_register_t<float, wp::tile_layout_register_t<wp::tile_tuple_t<(int)27>>>, _repro_tile_reduce__locals__repro_tile_reduce_b774bff3_cuda_kernel_forward::[lambda(T1, T2) (instance 2)]>(T2, T1 &)+0x1290 in tile_reduce.h:212
=========         Device Frame: _repro_tile_reduce__locals__repro_tile_reduce_b774bff3_cuda_kernel_forward+0x4a20 in repro.py:17
=========     Saved host backtrace up to driver entry point at kernel launch time
=========         Host Frame: wp_cuda_launch_kernel [0x7e808e] in warp.so
=========         Host Frame: launch in context.py:6279
=========         Host Frame: launch_tiled in context.py:6369
=========         Host Frame: <module> in repro.py:29

[daedalus5]

I don't think this is a LLVM issue since LLVM is not used in this code path...

Right... via the NVRTC path this shouldn't be an issue. Not sure why I'm seeing this error though?

[daedalus5]

Apparently the compute-sanitizer is sensitive to warp synchronization methods in debug mode. I can get around the issue if I replace our shuffle down method with a slower implementation:

#if defined(_DEBUG)
    // Debug mode workaround: use shared memory instead of warp shuffle
    // to avoid compute sanitizer "illegal instruction" errors
    __shared__ T warp_shuffle_temp[WP_TILE_BLOCK_DIM]; // Use block-wide shared memory
    
    const int lane_id = threadIdx.x % WP_TILE_WARP_SIZE;
    const int warp_id = threadIdx.x / WP_TILE_WARP_SIZE;
    const int warp_offset = warp_id * WP_TILE_WARP_SIZE;
    
    // Store value in shared memory at this thread's position
    warp_shuffle_temp[threadIdx.x] = val;
    __syncthreads(); // Use block-wide sync to avoid sanitizer issues
    
    // Read from the source lane (current lane + offset)
    int src_lane = lane_id + offset;
    T result = val; // default to current value
    
    // Only read if source lane is valid and part of the mask
    if (src_lane < WP_TILE_WARP_SIZE && (mask & (1U << src_lane)) != 0) {
        result = warp_shuffle_temp[warp_offset + src_lane];
    }
    
    __syncthreads(); // Ensure all reads complete
    return result;
#endif

This uses shared memory and block-wide thread synchronization. Not great having a separate method in debug mode however. We should discuss with the compiler team first.