Fix CUDA interop buffer race and add targeted CUDA↔graphics sync (#929) - #934
Open
jhelferty-nv wants to merge 4 commits into
Open
Fix CUDA interop buffer race and add targeted CUDA↔graphics sync (#929)#934jhelferty-nv wants to merge 4 commits into
jhelferty-nv wants to merge 4 commits into
Conversation
…slang#929) Two bugs caused CUDA_ERROR_ILLEGAL_ADDRESS when dispatching diff_pair(None, grad): 1. Vulkan/D3D12 interop: the zeroed primal buffer created by create_zeroed_interop_buffer was destroyed before async CUDA work (memset_device_async) completed, because the buffer ref was only kept alive when copy-back was needed. Now always store interop buffer refs in read_back to extend their lifetime past the dispatch. 2. CUDA direct path: no zeroed buffer was created at all — a null device address was written and the shader dereferenced it. Now allocate a zeroed buffer, store it in read_back, and wait on the submit fence before read_back goes out of scope (CUDA's RHI doesn't track raw device address refs for deferred deletion). Also fix write_torch_tensor_fields to use the interop buffer's device address in the TensorView path, which was previously ignoring the interop_buffer parameter entirely. Made-with: Cursor
Marshalls that perform CUDA work on shared interop buffers (copy_to_buffer, memset_device_async) now record PyTorch's current CUDA stream on the CallContext. exec() passes this stream to submit_command_buffer, which triggers the existing sync_to_cuda/sync_to_device fence mechanisms. This replaces the previous always-sync approach — sync only fires when a marshall actually performed CUDA work on shared buffers, avoiding unnecessary overhead when no interop is involved. Made-with: Cursor
📝 WalkthroughWalkthroughAdds CUDA interop stream tracking and synchronization, ensures temporary interop buffers (including zeroed buffers for null primals) are kept alive until GPU work completes, updates submit handling to record submit IDs and wait for completion when needed, and adds a grad-only diff_pair test. Changes
Sequence DiagramsequenceDiagram
participant PyTorch as PyTorch Tensor
participant SlangPy as SlangPy Interop
participant Device as GPU Device/CUDA
participant CmdSys as Command Submit System
PyTorch->>SlangPy: dispatch diff_pair(primal=None, grad)
SlangPy->>Device: allocate zeroed interop buffer (if needed)
Device-->>SlangPy: device_ptr
SlangPy->>Device: cuda::memset_device_async(device_ptr, 0, stream)
SlangPy->>SlangPy: mark_interop_cuda_stream(stream)
SlangPy->>CmdSys: submit command buffer (use opts.cuda_stream || interop_stream)
CmdSys-->>SlangPy: submit_id
CmdSys->>Device: execute dispatch
alt read_back or interop async ops present on CUDA
SlangPy->>Device: wait_for_submit(submit_id)
Device->>Device: ensure async memset and submits complete
end
SlangPy->>SlangPy: keepalive interop buffers in read_back until safe to destroy
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/slangpy_ext/utils/slangpy.cpp (1)
978-1007:⚠️ Potential issue | 🟠 MajorUse ASCII in the new sync comments.
This block is tripping
check-ascii-source(↔/—), so the build stays red until those characters are replaced with ASCII equivalents.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/slangpy_ext/utils/slangpy.cpp` around lines 978 - 1007, The comments in the block around submit_command_buffer (see variables submit_id, submit_cuda_stream, m_device->submit_command_buffer and the final comment about CUDA devices) contain non-ASCII characters (e.g. ↔ and —) which trigger check-ascii-source; replace them with ASCII equivalents (for example use "CUDA<->graphics" or "CUDA<->graphics synchronization" and use a simple hyphen "--" or "-" instead of an em-dash) so all comment text in that area is pure ASCII.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@slangpy/tests/slangpy_tests/test_torchintegration.py`:
- Around line 736-753: The test test_diffpair_get_shape_grad_only currently only
invokes the compiled function once; call the same function a second time to
reproduce the regression triggered by async CUDA work: after creating func via
helpers.create_function_from_module and constructing pair with diff_pair(None,
grad), invoke func(pair) twice and assert both results are non-None (i.e.,
ensure the first and second calls to func(pair) succeed), keeping the same
grad/device setup and existing assertions.
In `@src/sgl/utils/slangpy.h`:
- Around line 430-441: The comment contains a non-ASCII glyph ("↔") causing
check-ascii-source to fail; update the comment text in the block around
mark_interop_cuda_stream and interop_cuda_stream to use only ASCII (e.g. replace
"CUDA↔graphics" with "CUDA-graphics" or "CUDA <-> graphics") so the header
compiles under ASCII-only checks while leaving the logic in
mark_interop_cuda_stream and the interop_cuda_stream accessor unchanged.
In `@src/slangpy_ext/utils/slangpytorchtensor.cpp`:
- Around line 579-582: Replace non-ASCII characters in the comment that mentions
CUDA↔graphics and the em dash: change "CUDA↔graphics" to an ASCII form like
"CUDA<->graphics" or "CUDA-to-graphics" and replace the long dash "—" with two
hyphens "--" (or the word "only") in the comment that references exec() and
submit_command_buffer so the comment is fully ASCII; update the comment around
the record of PyTorch's current CUDA stream (the block mentioning exec() and
submit_command_buffer and calls like copy_to_buffer and memset_device_async)
accordingly.
- Around line 579-592: Replace non-ASCII punctuation in the comment with plain
ASCII and unify CUDA stream handling by resolving PyTorch's current CUDA stream
once and reusing it for both the memset and for marking the interop stream: in
the lambda mark_interop_stream and in create_zeroed_interop_buffer (and the
inline memset block), call
TorchBridge::instance().get_current_cuda_stream(device_index) once into a local
stream_ptr, pass that same stream_ptr to cuda::memset_device_async (instead of
context->cuda_stream() or nullptr), and then call
context->mark_interop_cuda_stream(...) using
NativeHandle(rhi::NativeHandleType::CUstream,
reinterpret_cast<uint64_t>(stream_ptr)); also ensure the inline memset (lines
~417–420) invokes the same resolution and mark so submit_command_buffer(...,
context->interop_cuda_stream()) synchronizes against the correct stream.
---
Outside diff comments:
In `@src/slangpy_ext/utils/slangpy.cpp`:
- Around line 978-1007: The comments in the block around submit_command_buffer
(see variables submit_id, submit_cuda_stream, m_device->submit_command_buffer
and the final comment about CUDA devices) contain non-ASCII characters (e.g. ↔
and —) which trigger check-ascii-source; replace them with ASCII equivalents
(for example use "CUDA<->graphics" or "CUDA<->graphics synchronization" and use
a simple hyphen "--" or "-" instead of an em-dash) so all comment text in that
area is pure ASCII.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 34669203-5fef-4454-8817-f22e3e9d8c34
📒 Files selected for processing (4)
slangpy/tests/slangpy_tests/test_torchintegration.pysrc/sgl/utils/slangpy.hsrc/slangpy_ext/utils/slangpy.cppsrc/slangpy_ext/utils/slangpytorchtensor.cpp
Replace non-ASCII glyphs (arrows, em-dashes) with ASCII equivalents to pass check-ascii-source pre-commit hook. Call func(pair) twice in test_diffpair_get_shape_grad_only to exercise the second-dispatch regression path from shader-slang#929. Apply clang-format. Made-with: Cursor
Made-with: Cursor # Conflicts: # src/slangpy_ext/utils/slangpytorchtensor.cpp
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/slangpy_ext/utils/slangpy.cpp (1)
1008-1014: Track submit-wait need explicitly, not viaread_backsize.
read_backis the generic post-dispatch container in this file, so this fence now fires for any CUDA call that appends regular readback data, not just the raw-address keepalives this fix is targeting. A dedicatedneeds_submit_waitflag would keep the lifetime fix without widening host-side synchronization.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/slangpy_ext/utils/slangpy.cpp` around lines 1008 - 1014, Introduce a dedicated boolean flag (e.g. needs_submit_wait) and use it instead of nb::len(read_back) to gate the CUDA submit wait; set needs_submit_wait = true at the specific code sites that append raw-device-address keepalive buffers (the places that currently add the temporary zeroed primal/other raw-address lifetimes), keep existing checks for submit_id and DeviceType::cuda, and call m_device->wait_for_submit(submit_id) only when needs_submit_wait is true to avoid widening host-side synchronization caused by general read_back entries.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/slangpy_ext/utils/slangpytorchtensor.cpp`:
- Around line 494-506: The interop branch is incorrectly treating
TensorView/DiffTensorView as supported on non-CUDA backends by writing
interop_buffer->device_address() into TensorViewData (tvd); add a guard that
detects when the bound field's type is TensorView or DiffTensorView (or
otherwise CUDA-only) and if the current backend is not CUDA, fail fast instead
of using interop_buffer. In practice, check the field/type or binding metadata
where interop_buffer is handled, and if type is TensorView/DiffTensorView and
backend != CUDA, return an error/abort the binding (do not set tvd.data or
strides). Keep the existing contiguous stride computation for supported CUDA
paths only.
- Around line 582-595: mark_interop_stream currently records
TorchBridge::instance().get_current_cuda_stream(device_index) at marshalling
time which can differ from the stream used by create_zeroed_interop_buffer's
memset_device_async (context->cuda_stream()), causing wrong-stream sync races;
fix by recording the actual stream used for the zero-fill: change
mark_interop_stream (or its call site) to use context->cuda_stream() (or accept
the stream pointer returned/used by memset_device_async) instead of querying
TorchBridge::get_current_cuda_stream, and ensure create_zeroed_interop_buffer
and the interop sync path both reference the same NativeHandle constructed from
context->cuda_stream(); update any other mark_interop_stream uses (mentioned
lines ~631-635) similarly so the recorded interop stream always matches the
memset stream.
---
Nitpick comments:
In `@src/slangpy_ext/utils/slangpy.cpp`:
- Around line 1008-1014: Introduce a dedicated boolean flag (e.g.
needs_submit_wait) and use it instead of nb::len(read_back) to gate the CUDA
submit wait; set needs_submit_wait = true at the specific code sites that append
raw-device-address keepalive buffers (the places that currently add the
temporary zeroed primal/other raw-address lifetimes), keep existing checks for
submit_id and DeviceType::cuda, and call m_device->wait_for_submit(submit_id)
only when needs_submit_wait is true to avoid widening host-side synchronization
caused by general read_back entries.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b9f3b2e1-1e15-48fb-a5fc-3c6895ab9820
📒 Files selected for processing (4)
slangpy/tests/slangpy_tests/test_torchintegration.pysrc/sgl/utils/slangpy.hsrc/slangpy_ext/utils/slangpy.cppsrc/slangpy_ext/utils/slangpytorchtensor.cpp
Comment on lines
+494
to
+506
| if (interop_buffer) { | ||
| tvd.data = static_cast<uint64_t>(interop_buffer->device_address()); | ||
| // Recalculate strides as contiguous for the interop buffer copy | ||
| Shape contiguous_strides = make_contiguous_strides(shape, info.element_size); | ||
| contiguous_strides = apply_broadcast_stride_zeroing( | ||
| contiguous_strides, | ||
| shape, | ||
| binding->transform(), | ||
| context->call_shape() | ||
| ); | ||
| for (int i = 0; i < info.ndim && i < kSlangPyTensorViewMaxDim; i++) | ||
| tvd.strides[i] = static_cast<uint32_t>(contiguous_strides[i] * info.element_size); | ||
| } |
There was a problem hiding this comment.
Reject TensorView interop on non-CUDA backends.
This branch now feeds interop_buffer->device_address() into TensorViewData, which makes TensorView/DiffTensorView appear to work through the graphics interop path. Those types are CUDA-only, so this should fail fast instead of binding a non-CUDA address model here.
Proposed guard
if (offsets.is_tensorview) {
+ SGL_CHECK(
+ interop_buffer == nullptr || context->device()->type() == DeviceType::cuda,
+ "TensorView/DiffTensorView are CUDA-only and do not support graphics interop backends."
+ );
TensorViewData tvd = populate_tensorview_data(info, shape, strides);
if (interop_buffer) {
tvd.data = static_cast<uint64_t>(interop_buffer->device_address());Based on learnings, TensorView and DiffTensorView types in SlangPy are CUDA-only and do not support non-CUDA backends (D3D12/Vulkan); the interop buffer path should not be used with TensorView fields.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/slangpy_ext/utils/slangpytorchtensor.cpp` around lines 494 - 506, The
interop branch is incorrectly treating TensorView/DiffTensorView as supported on
non-CUDA backends by writing interop_buffer->device_address() into
TensorViewData (tvd); add a guard that detects when the bound field's type is
TensorView or DiffTensorView (or otherwise CUDA-only) and if the current backend
is not CUDA, fail fast instead of using interop_buffer. In practice, check the
field/type or binding metadata where interop_buffer is handled, and if type is
TensorView/DiffTensorView and backend != CUDA, return an error/abort the binding
(do not set tvd.data or strides). Keep the existing contiguous stride
computation for supported CUDA paths only.
Comment on lines
+582
to
+595
| // Record PyTorch's current CUDA stream on the context so that exec() can | ||
| // pass it to submit_command_buffer for CUDA<->graphics synchronization. | ||
| // Called after any CUDA work on shared interop buffers (copy_to_buffer, | ||
| // memset_device_async). Only needed once per dispatch -- subsequent calls | ||
| // are no-ops if the stream is already recorded. | ||
| auto mark_interop_stream = [&](int32_t device_index) | ||
| { | ||
| if (context->interop_cuda_stream().is_valid()) | ||
| return; | ||
| void* stream_ptr = TorchBridge::instance().get_current_cuda_stream(device_index); | ||
| context->mark_interop_cuda_stream( | ||
| NativeHandle(rhi::NativeHandle(rhi::NativeHandleType::CUstream, reinterpret_cast<uint64_t>(stream_ptr))) | ||
| ); | ||
| }; |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "Inspect the stream recorded for interop sync:"
sed -n '582,595p' src/slangpy_ext/utils/slangpytorchtensor.cpp
echo
echo "Inspect the stream used by the zeroed-buffer memset:"
sed -n '629,635p' src/slangpy_ext/utils/slangpytorchtensor.cpp
echo
echo "Find all related stream-selection sites:"
rg -n -C2 'get_current_cuda_stream|mark_interop_stream|memset_device_async|context->cuda_stream\(\)' src/slangpy_ext/utils/slangpytorchtensor.cppRepository: shader-slang/slangpy
Length of output: 3706
🏁 Script executed:
# Find definition of context->cuda_stream()
rg -n 'cuda_stream\(\)' src/slangpy_ext/utils/slangpytorchtensor.cpp -B2 -A2Repository: shader-slang/slangpy
Length of output: 1125
🏁 Script executed:
# Find the context type and its cuda_stream() method definition
fd -e h -e hpp | xargs rg -l 'cuda_stream' | head -10Repository: shader-slang/slangpy
Length of output: 337
🏁 Script executed:
# Look for TorchBridge::get_current_cuda_stream definition
rg -n 'get_current_cuda_stream' src/slangpy_ext/ -A5Repository: shader-slang/slangpy
Length of output: 4363
🏁 Script executed:
# Check if context->cuda_stream() is set explicitly or matches PyTorch's stream
rg -n 'mark_interop_cuda_stream|set.*cuda_stream|cuda_stream.*=' src/slangpy_ext/utils/slangpytorchtensor.cpp -B1 -A1Repository: shader-slang/slangpy
Length of output: 325
🏁 Script executed:
# Find the type of 'context' variable and its cuda_stream() method
rg -n 'context.*=' src/slangpy_ext/utils/slangpytorchtensor.cpp | head -20Repository: shader-slang/slangpy
Length of output: 139
🏁 Script executed:
# Search for cuda_stream() method definition across headers
rg -n 'cuda_stream\(\)' src/slangpy_ext/ -B3 -A3 --type h --type hpp | head -50Repository: shader-slang/slangpy
Length of output: 91
🏁 Script executed:
# Find Device class or similar that has cuda_stream() method
rg -n 'class.*Device|struct.*Device' src/slangpy_ext/ --type h --type hpp | head -10Repository: shader-slang/slangpy
Length of output: 91
🏁 Script executed:
# Find function signature and context parameter
sed -n '560,610p' src/slangpy_ext/utils/slangpytorchtensor.cpp | head -100Repository: shader-slang/slangpy
Length of output: 2218
🏁 Script executed:
# Search for cuda_stream in header files
rg -n 'cuda_stream' src/slangpy_ext/ -t h | head -30Repository: shader-slang/slangpy
Length of output: 1814
🏁 Script executed:
# Look for context initialization in the marshalling function
rg -n 'const auto context|ref<.*Context>' src/slangpy_ext/utils/slangpytorchtensor.cpp | head -20Repository: shader-slang/slangpy
Length of output: 46
🏁 Script executed:
# Find CallContext class definition and cuda_stream() method
rg -n 'class CallContext|struct CallContext' src/slangpy_ext/ -A30 | head -60Repository: shader-slang/slangpy
Length of output: 46
🏁 Script executed:
# Search for cuda_stream() implementation in CallContext
rg -n 'cuda_stream\(\)' src/sgl/ -B2 -A2 --type h | head -50Repository: shader-slang/slangpy
Length of output: 861
🏁 Script executed:
# Check how context's cuda_stream is set/initialized
rg -n 'CallContext.*cuda_stream|cuda_stream.*=' src/sgl/ -B2 -A2 --type h | head -50Repository: shader-slang/slangpy
Length of output: 2476
🏁 Script executed:
# Find where CallContext is created and how cuda_stream is passed
rg -n 'set_call_shape_and_cuda_stream|CallContext' src/slangpy_ext/utils/slangpyfunction.cpp -B3 -A3 | head -80Repository: shader-slang/slangpy
Length of output: 46
🏁 Script executed:
# Look for the full dispatch pipeline to see how cuda_stream in options is set
sed -n '40,100p' src/slangpy_ext/utils/slangpyfunction.cppRepository: shader-slang/slangpy
Length of output: 2557
🏁 Script executed:
# Check if there's explicit synchronization between dispatch-time stream and current stream
rg -n 'cudaStreamSynchronize|sync_to_cuda|sync_to_device' src/slangpy_ext/utils/slangpytorchtensor.cppRepository: shader-slang/slangpy
Length of output: 46
🏁 Script executed:
# Verify the exact flow: when is apply_torch_cuda_stream called vs when is write_shader_cursor_with_interop called
rg -n 'apply_torch_cuda_stream|write_shader_cursor_with_interop|invoke' src/slangpy_ext/utils/slangpyfunction.cpp | head -20Repository: shader-slang/slangpy
Length of output: 645
🏁 Script executed:
# Check the exact line numbers of apply_torch_cuda_stream to understand dispatch timing
sed -n '55,75p' src/slangpy_ext/utils/slangpyfunction.cppRepository: shader-slang/slangpy
Length of output: 955
🏁 Script executed:
# Verify the dispatch flow: when are options passed to context and when are marshalls called
rg -n 'set_call_shape_and_cuda_stream|write_shader_cursor_with_interop' src/slangpy_ext/ -B2 -A2 | head -60Repository: shader-slang/slangpy
Length of output: 1294
🏁 Script executed:
# Check what happens between dispatch invocation and marshalling - any stream changes?
rg -n 'class NativeFunctionNode|invoke.*{' src/slangpy_ext/utils/slangpyfunction.cpp -A30 | head -80Repository: shader-slang/slangpy
Length of output: 195
🏁 Script executed:
# Verify that the memset/mark_interop_stream are truly in the dispatch path
sed -n '605,640p' src/slangpy_ext/utils/slangpytorchtensor.cppRepository: shader-slang/slangpy
Length of output: 1859
🏁 Script executed:
# Final verification: check if there's any explicit stream switching between setup and marshalling
sed -n '100,110p' src/slangpy_ext/utils/slangpyfunction.cppRepository: shader-slang/slangpy
Length of output: 537
🏁 Script executed:
# Verify copy_to_buffer also calls mark_interop_stream to ensure consistency
sed -n '605,615p' src/slangpy_ext/utils/slangpytorchtensor.cppRepository: shader-slang/slangpy
Length of output: 654
🏁 Script executed:
# Verify when marshalling happens relative to stream setup - are they in same dispatch?
sed -n '100,120p' src/slangpy_ext/utils/slangpyfunction.cppRepository: shader-slang/slangpy
Length of output: 749
🏁 Script executed:
# Check if there's any possibility of stream change between setup and marshalling
# by looking at what happens in call_data->call()
rg -n 'class NativeCallData' src/ -A20 | head -40Repository: shader-slang/slangpy
Length of output: 2647
🏁 Script executed:
# Verify the exact synchronization guarantee: does context->cuda_stream() match the one recorded at setup?
rg -n 'mark_interop_cuda_stream\|interop_cuda_stream' src/sgl/utils/slangpy.h -B2 -A2Repository: shader-slang/slangpy
Length of output: 46
🏁 Script executed:
# Search for mark_interop_cuda_stream definition
rg -n 'mark_interop_cuda_stream' src/sgl/ -B2 -A2Repository: shader-slang/slangpy
Length of output: 473
🏁 Script executed:
# Check CallContext to see if cuda_stream and interop_cuda_stream are independent
rg -n 'class CallContext' src/sgl/ -A30 | head -60Repository: shader-slang/slangpy
Length of output: 46
Use the same CUDA stream for both memset_device_async() and interop sync tracking.
mark_interop_stream() records PyTorch's current stream at marshalling time, but create_zeroed_interop_buffer() launches the zero-fill on context->cuda_stream(), which is captured at dispatch setup. If PyTorch's current stream changes between dispatch and marshalling, submit_command_buffer(..., context->interop_cuda_stream()) syncs against the wrong stream, reopening the zeroed-buffer race condition for custom stream scenarios.
Suggested fix
auto mark_interop_stream = [&](int32_t device_index)
{
if (context->interop_cuda_stream().is_valid())
return;
void* stream_ptr = TorchBridge::instance().get_current_cuda_stream(device_index);
context->mark_interop_cuda_stream(
NativeHandle(rhi::NativeHandle(rhi::NativeHandleType::CUstream, reinterpret_cast<uint64_t>(stream_ptr)))
);
};
+
+ auto resolve_interop_stream = [&](int32_t device_index) -> CUstream
+ {
+ void* stream_ptr = TorchBridge::instance().get_current_cuda_stream(device_index);
+ context->mark_interop_cuda_stream(
+ NativeHandle(rhi::NativeHandle(rhi::NativeHandleType::CUstream, reinterpret_cast<uint64_t>(stream_ptr)))
+ );
+ return reinterpret_cast<CUstream>(stream_ptr);
+ };
auto create_zeroed_interop_buffer = [&](const TensorBridgeInfo& info) -> ref<Buffer>
{
size_t buffer_size = static_cast<size_t>(info.numel) * static_cast<size_t>(info.element_size);
if (buffer_size == 0)
@@
});
void* cuda_ptr = interop_buffer->cuda_memory();
if (cuda_ptr && buffer_size > 0) {
- CUstream stream = context->cuda_stream().is_valid()
- ? reinterpret_cast<CUstream>(context->cuda_stream().value())
- : nullptr;
+ CUstream stream = resolve_interop_stream(info.device_index);
cuda::memset_device_async(static_cast<uint8_t*>(cuda_ptr), 0, buffer_size, stream);
- mark_interop_stream(info.device_index);
}
return interop_buffer;
};Also applies to: 631-635
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/slangpy_ext/utils/slangpytorchtensor.cpp` around lines 582 - 595,
mark_interop_stream currently records
TorchBridge::instance().get_current_cuda_stream(device_index) at marshalling
time which can differ from the stream used by create_zeroed_interop_buffer's
memset_device_async (context->cuda_stream()), causing wrong-stream sync races;
fix by recording the actual stream used for the zero-fill: change
mark_interop_stream (or its call site) to use context->cuda_stream() (or accept
the stream pointer returned/used by memset_device_async) instead of querying
TorchBridge::get_current_cuda_stream, and ensure create_zeroed_interop_buffer
and the interop sync path both reference the same NativeHandle constructed from
context->cuda_stream(); update any other mark_interop_stream uses (mentioned
lines ~631-635) similarly so the recorded interop stream always matches the
memset stream.
jhelferty-nv
enabled auto-merge (squash)
April 10, 2026 21:17
jhelferty-nv
added a commit
to jhelferty-nv/slangpy
that referenced
this pull request
Apr 13, 2026
Back out C++ interop buffer changes (slangpytorchtensor.cpp, slangpy.cpp, slangpy.h) per discussion - these need a different approach in shader-slang#934. Re-restrict test_diffpair_get_shape_grad_only to CUDA only since the interop race condition (shader-slang#929) is no longer fixed in this branch. Also fixes: - test_torch_dtypes.py: guard torch import for macOS - Non-ASCII em dashes in test docstrings Made-with: Cursor
ccummingsNV
requested changes
Apr 30, 2026
ccummingsNV
left a comment
ccummingsNV
left a comment
Contributor
There was a problem hiding this comment.
Quite a bit of work to go on this one I think - I don't think we've fixed the correct bits with the correct approaches. Have a read through my comments and see what you think. We can discuss on Teams if need be
| // diff_pair(None, grad)) are referenced by raw device address -- the RHI has | ||
| // no ref to defer their deletion. Wait for the submit to finish so the GPU | ||
| // dispatch is complete before read_back goes out of scope and frees them. | ||
| if (submit_id && m_device->type() == DeviceType::cuda && nb::len(read_back) > 0) { |
Contributor
There was a problem hiding this comment.
If this logic is correct, there is a serious problem deeper in the RHI - it should be impossible for resources to be freed before the gpu has finished with them. We can't add a cpu/gpu sync to solve this.
Contributor
There was a problem hiding this comment.
If the RHI is working correctly, then the requirement is that you should not allow the buffer object's destructor to be called until after the dispatch has been submitted. If this is done, the fence value will be tracked, and the actual device memory will not be freed until that submit is complete.
| if (offsets.is_tensorview) { | ||
| // TensorView path: build TensorViewData struct and write via set_data() | ||
| TensorViewData tvd = populate_tensorview_data(info, shape, strides); | ||
| if (interop_buffer) { |
Contributor
There was a problem hiding this comment.
We made the concious decision to say TensorView should not need to support interop, as slang treats it as a cuda only construct
| // devices before read_back goes out of scope, ensuring the GPU dispatch | ||
| // is complete before the buffer is freed. | ||
| ref<Buffer> primal_zeroed_buffer; | ||
| if (primal_info.data_ptr == nullptr && primal_info.numel > 0) { |
Contributor
There was a problem hiding this comment.
Is there no way for us to support a null binding and handle this. It doesn't feel right to me that if a user has, eg, a 1GB tensor that we allocate 1GB of zeros to avoid passing nulls.
| // Called after any CUDA work on shared interop buffers (copy_to_buffer, | ||
| // memset_device_async). Only needed once per dispatch -- subsequent calls | ||
| // are no-ops if the stream is already recorded. | ||
| auto mark_interop_stream = [&](int32_t device_index) |
Contributor
There was a problem hiding this comment.
This mechanism seems a bit shonky and has already suffered from errors relating to missmatched streams (note the one coderabbit highlights).
Is the correct fix not to simply add a device->sync_to_cuda after the copy, passing in the stream on which the copy was performed?
ccummingsNV
added a commit
that referenced
this pull request
Jul 31, 2026
* Add tests for sliced PyTorch tensors as fixed-size array parameters (#761) Test that PyTorch tensor views created by slicing can be correctly marshalled to Slang fixed-size array parameters (float[N]), covering both basic acceptance and gradient parity with pure PyTorch. Made-with: Cursor * Add stride coverage tests for vector, RWTensor, and Tensor<float,N> params Cover the high-risk gaps in non-trivial tensor view handling: - float3 vector params with prefix/offset/strided sliced tensors - RWTensor write-back correctness with sliced output tensors - Tensor<float,2> with transposed (non-contiguous) input tensors - Gradient parity for float3 params with sliced tensor views Made-with: Cursor * Add medium-risk stride coverage: transpose, Tensor<float,2> views, float[5] - Transpose tests for float[3] and float3 params (non-unit stride in the trailing dimension that maps to array/vector components) - float[5] with prefix/offset/strided slices (verifies array marshalling generalizes beyond the float[3] case) - Tensor<float,2> with column-prefix, column-offset, and column-strided views (extends the existing transpose-only coverage) - Gradient parity tests for both float[3] and float3 with transposed inputs Made-with: Cursor * Add TensorView, DiffTensorView, and numpy stride coverage tests Cover remaining non-contiguous memory layout gaps: - TensorView: sliced input reads, sliced output write-back, sliced add - DiffTensorView: sliced input/output, backward pass with sliced inputs - WTensor: transpose write-back - numpy: non-contiguous arrays for float and float3 params - Remove redundant test_array5_parameter_slice Made-with: Cursor * Add flip, diagonal, permute, and broadcast stride tests Cover remaining medium+ risk non-contiguous memory patterns: - Negative strides (flip): all 1D and 2D param test files - Diagonal views (stride = sum of dims): 1D param tests - 3D permute (both dims non-unit stride): Tensor<float,2> tests - Zero-stride broadcast (expand): float[3], float3, Tensor<float,2> - Numpy flip and broadcast cases Made-with: Cursor * Remove flip and numpy non-contiguous tests torch.flip() always returns a contiguous copy (PyTorch does not support negative strides), so flip cases were just re-testing the contiguous path. The C++ numpy marshalling layer explicitly rejects non-contiguous arrays, so those tests were targeting a missing feature. Made-with: Cursor * Fix scalar DiffPair backward pass codegen The generated Slang struct for DiffPairMarshall had two bugs preventing scalar backward passes from working: 1. Used bare `load`/`store` method names instead of `__slangpy_load`/ `__slangpy_store` (the convention used by all SlangPy Slang types). 2. Generated `__slangpy_load` that output `DifferentialPair<T>` instead of `T`, causing a type mismatch with the trampoline temporaries. The methods were also not `[Differentiable]`, so `bwd_diff` could not propagate gradients through them. The fix rewrites `generate_differential_pair` to follow the same pattern as `DiffTensor`: the struct's load/store methods operate on the primal type and, when gradients are needed, use `[Differentiable]` methods with `[BackwardDerivative]` custom backward functions that explicitly read/write the derivative buffer. Removes the "Awaiting auto-diff changes" skip markers from test_call_with_diff_scalars and test_call_with_diff_pairs, which now pass on both Vulkan and CUDA. Made-with: Cursor * Add NDBuffer and DiffPair type test coverage (#782) - test_ndbuffer.py: 21 tests covering NDBuffer construction (empty, zeros, from_numpy, element_count, empty_like, zeros_like), data transfer (copy_from_numpy, to_numpy, clear), views, broadcasting, validation errors, and end-to-end Slang function dispatch with NDBufferMarshall (1D/2D add, scalar broadcast, read-only pass, StructuredBuffer pass, return type, int dtype). - test_diffpair_type.py: 19 tests covering DiffPair data class construction (defaults, None handling, needs_grad, int types), get/set by PrimType, slangpy_signature, and both factory functions (diffPair, floatDiffPair). Made-with: Cursor * Add mixed torch+NDBuffer and build_shader_object tests (#782) - test_torch_mixed.py: 7 CUDA tests covering mixed torch.Tensor + NDBuffer argument combinations, torch+scalar broadcast, torch return values, 2D torch dispatch, and pack() with torch.Tensor (exercises TorchTensorMarshall.build_shader_object path). Made-with: Cursor * Add coverage tests for torch integration marshalls and bridge fallback (#782) Tests for bridge_fallback.py error/guard paths, torchtensormarshall.py DiffPair factory and error paths, builtin/tensor.py grad validation and properties, and builtin/value.py return types and reduce_type via torch pipelines. Made-with: Cursor * Add Tensor/TensorView/DiffTensorView coverage tests and fix flip_y bug (#782) - Fix non-contiguous array bug in load_buffer_data_from_image when flip_y=True (np.flipud returns a view; add contiguity guard) - Add test_load_image.py: covers Tensor.load_from_image channel variants, flip_y, scale/offset, greyscale, and linearize flags - Add test_torch_dtypes.py: covers create_output branches for int32, int64, float64 in slangpytorchtensor.cpp - Add test_descriptor_marshall.py: covers DescriptorHandle low-level API and DescriptorMarshall error paths - Extend test_tensor.py: covers __str__, detach(), deprecated Tensor.numpy(), deprecated element_count param, and empty() error paths - Extend test_tensorview.py: covers slangpy.Tensor through TensorView and scalar-to-vector TensorView<float2> resolution (tensorcommon.py) - Extend test_difftensorview.py: covers slangpy.Tensor through DiffTensorView forward path (tensorcommon.py) Made-with: Cursor * Add coverage tests for dispatchdata, valueref, function, and structuredbuffer (#782) - test_raw_dispatch: error paths (return value, out param, no params, wrong param name), writer tuple via .write(), dispatch cache hit - test_return_types: inout scalar/vector/matrix ValueRef exercising padding/unpadding paths in valueref.py - test_modules: function.py error paths (set, write, return_type, as_struct) - test_buffers: ByteAddressBuffer read-only and RWByteAddressBuffer write Made-with: Cursor * Add C++ coverage tests for repr, texture shape, benchmarks, and torch scalar types (#782) Made-with: Cursor * Redistribute C++ coverage tests into feature-aligned test files (#782) Move tests from test_cpp_coverage_tier1.py and test_cpp_coverage_torch.py into existing files by topic: repr tests to test_tostring, texture shape to test_textures, Shape/SignatureBuilder to test_shape, buffer view indexing to test_buffer_views, torch dtypes to test_torch_dtypes, raw dispatch to test_raw_dispatch, DiffPair to test_torch_marshall_gaps, and benchmarks to a new standalone test_torch_benchmarks.py. Made-with: Cursor * Rewrite DescriptorMarshall error tests to use functional API (#782) Replace direct DescriptorMarshall construction with indirect triggers: reduce_type error via func.map(tex=(0,)), init error by passing a DescriptorHandle to a function whose module lacks slangpy import. Made-with: Cursor * Rewrite TensorMarshall tests to use functional API (#782) Replace direct TensorMarshall construction with indirect triggers: - Grad dtype mismatch errors via Tensor.with_grads() + function call - Writable check via read-only tensor with grad_in + function call - Properties test via pack() instead of direct construction - Rename create_tensor_marshall test for consistency Made-with: Cursor * Replace direct NativeTorchTensorDiffPair usage with diff_pair() and backward pass (#782) - Use diff_pair() wrapper in 6 tests that used NativeTorchTensorDiffPair with is_input=True (the default) - Rewrite is_input=False test as actual backward pass test that triggers output diff pairs naturally through autograd - Rewrite grad-only factory test to use diff_pair(None, grad) - Add comment to test_slangtype_repr_no_reflection explaining why NativeSlangType is constructed directly Made-with: Cursor * Replace BindContext/CallMode with pack() in build_shader_object test (#782) Use pack(module, diff_pair(primal, grad)) which internally constructs a BindContext and calls build_shader_object, instead of constructing these native objects directly. Made-with: Cursor * Drop build_shader_object gradient stub test (#782) The NotImplementedError in TorchTensorMarshall.build_shader_object is an intentional stub — the C++ autograd path handles gradients directly and bypasses build_shader_object entirely. No value in covering a single raise line for a path that was never needed. Made-with: Cursor * Remove source line number references from test docstrings and comments (#782) Line numbers in the target source files will drift over time, making these references misleading rather than helpful. Made-with: Cursor * Add C++ Tier 3 coverage: StridedBufferView error guards, contiguity, PackedArg properties (#782) - point_to shape/element_stride mismatch errors - Indexing: illegal argument type, too many indices - copy_from_numpy non-contiguous source rejection - clear() with explicit command encoder - is_contiguous singleton dimension skip - maybe_pad_data float3 padding (skips: stride=12 on current hw) - PackedArg .python/.shader_object/.python_object property access Made-with: Cursor * Redistribute torch marshall gap tests into feature-aligned files (#782) Move tests from test_torch_marshall_gaps.py into test_torchintegration.py (factory paths, properties, error guards, nanobind bindings) and test_torch_autograd_workflows.py (backward pass DiffPair output). Broaden parametrization from CUDA-only to DEVICE_TYPES and switch dispatch tests to get_torch_device() for Vulkan interop support. Made-with: Cursor * Remove test numbering from autograd workflow section headers Made-with: Cursor * Apply black formatting to 8 files flagged by pre-commit CI Made-with: Cursor * Add CUDA availability guards to tests that hardcode device="cuda" Tests that directly create torch tensors on CUDA now skip when CUDA is not available. Fixes failures on non-CUDA runners (e.g. macOS). Made-with: Cursor * Use raw strings for regex patterns in pytest.raises match= arguments Prefix match patterns containing regex metacharacters with r to satisfy Ruff RUF043 and make regex intent explicit. Made-with: Cursor * Mark descriptor marshall tests as xfail pending #922 DescriptorHandle type resolution is broken (resolves to StructuredBuffer<Unknown> instead of the actual descriptor type), causing all four tests to fail with ResolveException or wrong error messages. Mark as strict xfail so they surface when #922 is fixed. Made-with: Cursor * Fix test_tensor_marshall_gaps.py CI failures - Remove grad_out pack() call from test_tensor_marshall_properties_via_pack (hits known "primal" field bug; first half already covers properties) - Mark three grad validation tests as strict xfail: they crash with nanobind uninitialized instance error instead of raising ValueError - Make existing test_pack_tensor_with_grads xfail strict Made-with: Cursor * Restrict test_diffpair_get_shape_grad_only to CUDA device only Dispatching diff_pair(None, grad) through Vulkan/D3D12 interop triggers a cleanup race in create_zeroed_interop_buffer where an async CUDA memset outlives the interop buffer, corrupting the CUDA context and crashing subsequent tests. Filed as #929. Running on CUDA alone still covers the intended get_shape fallback path (both native and fallback bridge modes). Made-with: Cursor * Fix interop buffer lifetime race and CUDA null-pointer crash (#929) Two bugs caused CUDA_ERROR_ILLEGAL_ADDRESS when dispatching diff_pair(None, grad): 1. Vulkan/D3D12 interop: the zeroed primal buffer created by create_zeroed_interop_buffer was destroyed before async CUDA work (memset_device_async) completed, because the buffer ref was only kept alive when copy-back was needed. Now always store interop buffer refs in read_back to extend their lifetime past the dispatch. 2. CUDA direct path: no zeroed buffer was created at all — a null device address was written and the shader dereferenced it. Now allocate a zeroed buffer, store it in read_back, and wait on the submit fence before read_back goes out of scope (CUDA's RHI doesn't track raw device address refs for deferred deletion). Also fix write_torch_tensor_fields to use the interop buffer's device address in the TensorView path, which was previously ignoring the interop_buffer parameter entirely. Made-with: Cursor * Add targeted CUDA↔graphics sync for torch interop buffers Marshalls that perform CUDA work on shared interop buffers (copy_to_buffer, memset_device_async) now record PyTorch's current CUDA stream on the CallContext. exec() passes this stream to submit_command_buffer, which triggers the existing sync_to_cuda/sync_to_device fence mechanisms. This replaces the previous always-sync approach — sync only fires when a marshall actually performed CUDA work on shared buffers, avoiding unnecessary overhead when no interop is involved. Made-with: Cursor * Fix non-ASCII characters and invoke test reproducer twice Replace non-ASCII glyphs (arrows, em-dashes) with ASCII equivalents to pass check-ascii-source pre-commit hook. Call func(pair) twice in test_diffpair_get_shape_grad_only to exercise the second-dispatch regression path from #929. Apply clang-format. Made-with: Cursor * Add type annotations and CUDA availability guard in test files Made-with: Cursor * Fix pre-commit formatting issues Made-with: Cursor * Revert fix-929 driver changes and fix CI issues Back out C++ interop buffer changes (slangpytorchtensor.cpp, slangpy.cpp, slangpy.h) per discussion - these need a different approach in #934. Re-restrict test_diffpair_get_shape_grad_only to CUDA only since the interop race condition (#929) is no longer fixed in this branch. Also fixes: - test_torch_dtypes.py: guard torch import for macOS - Non-ASCII em dashes in test docstrings Made-with: Cursor * Remove accidentally committed venv symlink Made-with: Cursor * Move flip_y fix to separate PR and xfail test_load_flip_y (#937) The contiguity fix for load_buffer_data_from_image is now in PR #937. Mark test_load_flip_y as xfail until that fix lands. Made-with: Cursor * Remove test_ndbuffer.py — NDBuffer is being removed (#548, #582) Coverage gains were entirely in builtin/ndbuffer.py and types/buffer.py, both of which are being replaced by Tensor. Made-with: Cursor * Fix black formatting in test_torch_dtypes.py Made-with: Cursor * Convert nanobind uninitialized instance xfails to skip These three grad validation tests crash with a nanobind uninitialized instance error during the dispatch pipeline. Skip instead of xfail to avoid potential state contamination from partially constructed C++ objects. Made-with: Cursor * Skip test_diffpair_get_shape_grad_only entirely (#929) Even the CUDA-only + fallback bridge mode variant triggers the interop buffer lifetime race. Convert from @requires_cuda to @pytest.mark.skip to prevent worker crashes and CUDA context poisoning. Issue #929 updated with re-enablement instructions. Made-with: Cursor * Address review feedback: fix vacuous tests, pytest practices, issue refs Section 1 - Fix vacuous/weak tests: - test_shape: add element_count assertion, fix negative indexing, test copy independence - test_torchintegration: add grad/output assertions to null_grad tests - test_value_torch_pipeline: verify matrix values, packed arg repr Section 2 - Add issue references to all skipped/xfailed tests: - #938: TensorMarshall nanobind lifecycle (grad validation crashes) - #939: DiffTensor/PrimalTensor sample not implemented - #940: Slang compiler bugs (generic vector crash, polynomial derivatives) Section 3 - Pytest best practices: - Remove allow_module_level=True from inside test functions (test_buffer_views) - Fix mutable default arguments in load_test_module (test_raw_dispatch) - Replace bare return with pytest.skip for 3D texture arrays (test_textures) - Move global PYTHON_TYPES mutation to fixture with teardown (test_return_types) Made-with: Cursor * Resolve commented-out / incomplete code from review - test_torchintegration: remove "Should this work??" backward call on non-differentiable RWTensor output (#574). Other tests already cover backward passes with proper differentiable outputs. - test_textures: remove dead transpose code commented out in PR #48 when dimension ordering was fixed upstream. - test_tensor: add assert to test_existential_type_bug compile-only regression test for clarity. Made-with: Cursor * Fix test_element_count_simple: Shape.element_count not exposed to Python The C++ Shape class has element_count but it's not bound via nanobind. Verify element count indirectly through contiguous strides instead. Made-with: Cursor * Skip test_element_count_simple: Shape.element_count not bound to Python The C++ Shape::element_count() method exists but was never registered via nanobind. Skip the test with an explanation rather than removing it, so it can be re-enabled once the binding is added. Made-with: Cursor * Update merged tests for current SlangPy APIs * Update tests for Torch vector signatures --------- Co-authored-by: James Helferty <jhelferty@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
diff_pair(None, grad)causedCUDA_ERROR_ILLEGAL_ADDRESSbecause zeroed interop buffers were destroyed before async CUDA work completed. Now always store interop buffer refs inread_backto extend lifetime past the dispatch.diff_pair(None, grad)wrote a null device address into the shader. Now allocates a zeroed buffer and waits on the submit fence before freeing it.copy_to_buffer/memset_device_async) now record PyTorch's current CUDA stream on theCallContext.exec()passes this stream tosubmit_command_buffer, activating the existingsync_to_cuda/sync_to_devicefence mechanisms only when needed.Closes #929
Changes
slangpytorchtensor.cppread_back; zeroed buffer for CUDA direct path;mark_interop_streamto record PyTorch's CUDA stream; fixTensorViewpath to use interop buffer addressslangpy.cppsubmit_id;wait_for_submiton CUDA whenread_backnon-empty; pass interop CUDA stream tosubmit_command_bufferslangpy.hmark_interop_cuda_stream()/interop_cuda_stream()toCallContexttest_torchintegration.pytest_diffpair_get_shape_grad_onlyreproducerTest plan
test_diffpair_get_shape_grad_only— all 4 variants (native/fallback x vulkan/cuda) passMade with Cursor
Summary by CodeRabbit
New Features
Bug Fixes
Tests