Skip to content

Make bound-variable recursion and tensor marshalls safe for reference-typed fields - #1045

Open
szihs wants to merge 5 commits into
mainfrom
haaggarwal/tensor-marshall-reference-aware
Open

Make bound-variable recursion and tensor marshalls safe for reference-typed fields#1045
szihs wants to merge 5 commits into
mainfrom
haaggarwal/tensor-marshall-reference-aware

Conversation

@szihs

@szihs szihs commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Slang's CUDA target is gaining a by-reference ABI for entry-point uniform structs that carry descriptor tables — fixed-size arrays of resources or pointer-backed tensors (shader-slang/slang#11774): such a parameter reflects as an implicit ParameterBlock sub-object occupying 8 bytes of kernel-argument space, with the payload in device global memory.

NativeBoundVariableRuntime::write_shader_cursor_pre_dispatch navigates into a bound struct argument with cursor[name] and passes the resulting cursor to children. For a reference-typed field, sgl field lookups auto-dereference into the sub-object (the shader-cursor "Just Work" path), so a child marshall extracts sub-object-relative offsets — while cursor.shader_object() still returns the parent object. A cached-offset writer like TensorMarshall (reserve_data(field_offset, field_size) + raw writes) would then silently corrupt the parent object's memory.

Proposed solution

  • One fix point, at the recursion site (slangpy.cpp): if the child field cursor is_reference(), dereference it before recursing — children always receive a cursor whose shader_object() owns the offsets they extract. This is the same pattern the ParameterBlock<CallData> fallback path already uses for call_data, and covers every leaf marshall at once.
  • Reference-aware NativeValueMarshall (slangpyvalue.cpp/h): the vectorized-array path binds a generated Array1DValueType<T,N> struct wrapping the array; when T carries descriptors (tensors, buffers) the parameter converts, and ensure_cached's cursor[name]["value"] navigation crossed the reference boundary — cached a sub-object-relative offset, then wrote through the parent object. This was the actual heap corruption observed on L40S in test_vectorize_struct_with_tensor_array / test_vectorize_struct_with_resource_array / test_2d_mapped_vectorize_struct_with_tensor_array (CUDA-only; Vulkan never converts). The marshall now dereferences explicitly, caches the field index, and targets the sub-object's ShaderObject at write time.
  • Loud guards in the cached-offset tensor writers (slangpytensor.cpp, slangpytorchtensor.cpp): SGL_CHECK(!field.is_reference(), ...) in the offset-caching path. Tensor types themselves are never passed by reference, so the check never fires on supported shapes — it converts any unexpected future shape or compiler/host version skew from silent memory corruption into a hard error.

The change is compatible with current Slang (every is_reference() is simply false today, so behavior is unchanged) and must land before slangpy bumps to a Slang containing the new ABI.

What was deliberately not changed

estimate_entrypoint_arguments_size still counts a convertible parameter at its by-value size. Correcting it to 8 bytes would require re-implementing the compiler's conversion predicate host-side — a second source of truth that can drift. The error direction is safe: over-estimation only pushes toward the ParameterBlock<CallData> fallback path.

Follow-up (separate PR, after the slang bump)

Bump the pinned Slang, verify test_benchmark_tensor.py::test_tensor_sum_indirect reaches parity with the forced-PB numbers (≈0.026/≈0.195 ms count 16/32 on L40S), then retire the #1030 CUDA routing heuristic.

🤖 Generated with Claude Code

…-typed fields

Slang's CUDA target is gaining a by-reference ABI for entry-point uniform
structs that carry descriptor tables (fixed-size arrays of resources or
pointer-backed tensors): such a parameter reflects as an implicit
ParameterBlock sub-object (shader-slang/slang#11774).

The bound-variable recursion navigated into such a field with
cursor[name] and passed the resulting cursor to children. Field lookups
on a reference-typed cursor auto-dereference into the sub-object, so a
child marshall caching those (sub-object-relative) offsets while writing
through cursor.shader_object() - still the parent object - would silently
corrupt memory. Dereference once at the recursion site so children always
receive a cursor whose shader_object() owns the offsets they extract,
matching what the ParameterBlock<CallData> fallback path already does for
call_data.

TensorMarshall and NativeTorchTensorMarshall additionally fail loudly
(SGL_CHECK) if their own bound field is ever reference-typed, converting
what would be silent corruption under a compiler/host version skew into a
hard error. Tensor types themselves are never passed by reference, so the
check never fires on supported shapes.

This change is compatible with both current and by-reference Slang
compilers (is_reference() is simply false everywhere today), and must land
before slangpy bumps to a Slang containing the new ABI.
@coderabbitai

coderabbitai Bot commented Jul 3, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

This change addresses reference-typed shader field handling to prevent silent offset corruption. The bound variable runtime now dereferences reference-typed child cursors before recursion. Tensor marshalls add validation rejecting reference-typed fields in cached-offset paths. The value marshall caches reference state and dereferences to select correct write targets.

Changes

Reference-typed field handling

Layer / File(s) Summary
Child cursor dereferencing
src/slangpy_ext/utils/slangpy.cpp
Dereferences reference-typed child cursors before recursive shader cursor writes to prevent offset corruption in children.
Tensor marshall reference guards
src/slangpy_ext/utils/slangpytensor.cpp, src/slangpy_ext/utils/slangpytorchtensor.cpp
Adds runtime checks that throw via SGL_CHECK when cached tensor fields are reference-typed, since the cached-offset fast path does not support this case.
Value marshall reference-aware caching and write targeting
src/slangpy_ext/utils/slangpyvalue.h, src/slangpy_ext/utils/slangpyvalue.cpp
Adds field_is_reference and field_index to cached write state; ensure_cached resolves and dereferences reference-typed fields before caching offsets, and write_shader_cursor_pre_dispatch targets the correct sub-object ShaderObject when the cached field is reference-typed.

Sequence Diagram(s)

sequenceDiagram
  participant NativeValueMarshall
  participant ShaderCursor
  participant ShaderObject

  NativeValueMarshall->>ShaderCursor: ensure_cached looks up field via variable_name
  ShaderCursor-->>NativeValueMarshall: field (possibly reference-typed)
  NativeValueMarshall->>NativeValueMarshall: cache field_is_reference, field_index
  NativeValueMarshall->>ShaderCursor: write_shader_cursor_pre_dispatch
  NativeValueMarshall->>ShaderObject: select target_object via cached field_index if reference
  NativeValueMarshall->>ShaderObject: write value_cursor into target_object
Loading

Estimated code review effort: High

Related issues: None found in the provided context.

Related PRs: None found in the provided context.

Suggested labels: bug, correctness, native-extension

Suggested reviewers: None specified in the provided context.

Poem
A rabbit hopped through cursors deep,
Where reference fields could offsets creep,
Dereferenced now, each sub-object found,
No silent corruption, safe and sound,
Hop hop hooray, the bindings keep!

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main change: making bound-variable recursion and tensor marshalling safe for reference-typed fields.
Description check ✅ Passed The description is directly related to the changeset and explains the motivation, fix, and follow-up clearly.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@szihs

szihs commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator Author

/format

@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Formatting requested - I'm on it!

@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Applied formatting and pushed a commit to this PR.

szihs added 2 commits July 3, 2026 07:17
The vectorized-array path binds a generated Array1DValueType<T,N> struct
wrapping the array. When T carries descriptors (tensors, buffers), Slang's
CUDA by-reference ABI reflects that parameter as a ParameterBlock
sub-object. NativeValueMarshall::ensure_cached navigated it with
cursor[name]["value"]: the nested lookup auto-dereferences into the
sub-object (making the cached offset sub-object-relative), while the write
constructed a cursor on cursor.shader_object() - the parent object -
producing out-of-bounds raw writes and heap corruption
(test_vectorize_struct_with_tensor_array and friends on CUDA).

Dereference explicitly when the bound field is reference-typed, cache the
field index, and target the sub-object's ShaderObject at write time -
the same pattern the call_data ParameterBlock fallback and the
bound-variable recursion already use. No behavior change with compilers
that pass everything by value (is_reference() is false everywhere).
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Pre-commit failed for this PR.

Workflow run: https://github.com/shader-slang/slangpy/actions/runs/28666521739

GitHub used:

pre-commit 4.3.0

You can check your local version with:

pre-commit --version

You can fix this locally by running:

pre-commit run --all-files

To install the hooks locally for future commits:

pre-commit install

Or ask GitHub to apply the formatting here by adding a PR comment containing:

/format

@szihs

szihs commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator Author

/format

@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Formatting requested - I'm on it!

@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Applied formatting and pushed a commit to this PR.

szihs added a commit to shader-slang/slang that referenced this pull request Jul 3, 2026
…d at parameter binding (#11941)

> **Stacked on #11939** (base branch is
`haaggarwal/cuda-param-dynamic-index-floor`); the diff shown here is
only this PR's commit. Fixes the compiler-side root cause of #11774.

## Motivation

The local-copy legalization in #11939 is a floor, not the fix: the host
still ships the full descriptor blob by value on every dispatch, every
thread pays a prologue copy, and the placement remains wrong for what a
descriptor array *is*. The measured regression (#11774: SlangPy
`test_tensor_sum_indirect`, 24×/13× slower than Vulkan on L40S) is:

```slang
struct TensorList { RWTensor<float,2> tensors[4]; }  // elements are pointer-backed structs
void main(..., uniform TensorList tensor_list, ...)  // by value in .param → serial ld.param chain
{ ... tensor_list.tensors[tensor_indices[i]] ... }
```

On every other target, indexed resource collections live in
dynamically-indexable descriptor memory (descriptor heaps) — which is
why Vulkan/D3D12 never regressed. On CUDA the equivalent placement is
*payload in global memory behind one pointer*: exactly what
`ParameterBlock` already produces, and measured **faster than Vulkan**
(0.173 ms vs 0.199 ms on RTX 5090 for the forced-PB variant of this
workload).

## Proposed solution

Model descriptor-table-carrying entry-point uniforms as what they
semantically are — parameter blocks — with **one decision site, at
parameter binding**, where reflection is computed:

1. **Layout (the only predicate evaluation)** — in
`computeEntryPointParameterTypeLayout`: on CUDA compute entry points, a
`uniform` struct parameter transitively containing a fixed-size array
whose elements carry descriptors (resources or pointers;
`typeContainsDescriptorBearingFixedArray`) is wrapped: `paramType =
astBuilder->getParameterBlockType(paramType)`. The standard layout
machinery then yields the parameter-group layout — an 8-byte pointer
container, element layout preserving all binding ranges — which
reflection reports truthfully.
2. **IR consumes the decision** — new post-link pass
`reconcileCUDAByRefEntryPointParams`: any entry-point parameter whose
recorded `IRVarLayout` is a parameter-group layout while the parameter
is still value-typed is retyped to `IRParameterBlockType(T)` and its
body value-uses rewritten to loads-through-addresses (via
`rewriteValueUsesToAddrUses`, hoisted from the constref pass into a
shared utility). The pass evaluates **no predicate**, so the emitted
kernel signature and the reflected layout cannot disagree — the
ABI-desync failure class (`CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES`) that
sank the emit-stage attempts (#11747, and the emit-side half of #11830)
is impossible by construction. This follows the #9869 precedent
(entry-point parameter-passing conventions rewritten by a dedicated
post-link IR pass).
3. **Zero emit changes, zero slang-rhi changes** —
`IRParameterBlockType` entry-point params already emit as `T*`, and
slang-rhi's CUDA backend deliberately supports parameter-block
sub-objects whose parent is an entry point (*"Sub-objects are always
written to global memory, even if the parent represents an
entry-point"*, cuda-shader-object.cpp), allocating the payload and
patching the device pointer at the reflected kernel-argument offset.
4. **Escape hatch** — new `-cuda-entry-point-params-by-value` option
(`CompilerOptionName::CudaEntryPointParamsByValue`, appended before
`CountOf`) restores the legacy all-by-value ABI; it gates the binding
decision, so layout and IR flip together.

Emitted result for the motivating shape:

```cpp
extern "C" __global__ void computeMain(TensorList_0* list_0, ...)
{ ... *((&list_0->tensors_0[_S3])->data_0 + _S1) ... }  // O(1) global load
```

Reflection reports the parameter as `kind: parameterBlock`, 8 bytes at
its kernel-argument offset, element preserving the array (count N,
stride 16); the next parameter starts after the pointer.

## Change summary

| File | Change |
|---|---|
| `source/slang/slang-parameter-binding.cpp` | predicate + gate + wrap
at the single decision site |
| `source/slang/slang-ir-cuda-byref-entry-point-params.{h,cpp}` |
layout-consuming IR reconciliation pass |
| `source/slang/slang-ir-transform-params-to-constref.{h,cpp}` | hoist
`rewriteValueUsesToAddrUses` to a shared free function (pure move) |
| `source/slang/slang-emit.cpp` | invoke the pass in the CUDA arm after
`collectOptiXEntryPointUniformParams` |
| `include/slang.h`, `source/slang/slang-options.cpp` | the opt-out
option |
| `tests/cuda/entry-point-uniform-descriptor-array-byref.slang` | emit +
reflection-JSON lock for the motivating shape |
| `tests/cuda/entry-point-uniform-resource-array-byref.slang` | true
resource array in a struct converts too |
| `tests/cuda/entry-point-uniform-byref-negative.slang` | negatives:
single tensor struct, plain-data struct, source-written PB (no
double-wrap) |
| `tests/cuda/entry-point-uniform-byref-opt-out.slang` | the flag
restores the legacy ABI |
| `tests/cuda/kernel-param-dynamic-index-nested-struct.slang` |
repurposed: floor behavior under the legacy ABI |
| `docs/cuda-target.md`, `docs/user-guide/09-targets.md` | document the
actual ABI rule |

## Concepts and vocabulary

- **Descriptor table** — an indexed collection of resource descriptors;
every binding model places these in dynamically-indexable memory.
Distinct from **root constants** (small plain data), whose canonical
CUDA home genuinely is `.param` — the predicate preserves that
distinction, which is why entry-point args remain meaningful (and fast)
after this change.
- **Layout-consuming pass** — the IR pass reads `IRLayoutDecoration` and
reconciles the IR to it; it re-derives nothing. Single source of truth
for the ABI.
- **E5 shape** — a source-written `uniform ParameterBlock<S>`
entry-point parameter, which compiles, reflects, and binds correctly
today; this PR makes binding *produce* that existing canonical form,
adding no new representation to the system.

## Process report

- **Why binding, not emit:** reflection is finalized at
parameter-binding time; any later ABI decision is invisible to hosts —
the #11747 crash class. The wrap happens where reflection is computed;
the IR pass consumes the recorded result. There is no second predicate
kept in sync by discipline (the flaw in #11830's split).
- **Why struct-typed parameters only (staged scope):** a bare descriptor
array would reflect as `ParameterBlock<T[N]>`, a group shape slang-rhi's
`_unwrapParameterGroups` cannot represent (silently collapses to 1 slot
on the sub-object path). Struct-wrapped groups take its `default:` path
and bind N slots on **already-shipped** slang-rhi binaries. Bare arrays
keep the by-value layout (covered by #11939's local copy); this covers
100% of the measured regression — including the issue's hand-written
reproducer, which also nests its array in a struct. The end state for
bare arrays (synthesized wrapper struct vs slang-rhi group-of-array
support) is left as an explicit maintainer decision rather than decided
implicitly here.
- **Cascading issue found during implementation:** the first
"struct-typed" check (`as<DeclRefType>` + `as<StructDecl>`) also matched
bare arrays — `T[N]` *is* a `DeclRefType` of the core-module `Array`
struct — converting `uniform RWStructuredBuffer<float> t[32]` into the
unsupported `ParameterBlock<array>` shape and crashing emit downstream.
Fixed with an explicit `as<ArrayExpressionType>` rejection (commented);
the underlying emit crash is fixed independently in #11940 since
source-written `ParameterBlock<float[4]>` crashes master today.
- **Input-shape check (per the methodology):** `uniform TensorList<N>`
living by value in `.param` was never a chosen representation — CUDA
compute is the arm where `moveEntryPointUniformParamsToGlobalScope` is
skipped and no placement pass runs; the placement was the *absence* of a
decision (`collectOptiXEntryPointUniformParams` returns immediately for
compute). This PR adds the decision at the layer that owns parameter
placement. Each exclusion names the component that owns that shape
instead: RT stages → the OptiX SBT pass;
torch/`[AutoPyBindCUDA]`/`[CudaKernel]` → their generated host launch
code; explicit parameter groups → already by-reference.
- **Predicate scope:** descriptor-bearing arrays only, *not* "any array"
— on CUDA every SlangPy tensor contains `uint[D]` shape/stride arrays
and the functional API emits `int[2]` metadata parameters; a broad
predicate would convert essentially every parameter into a sub-object
allocation, destroying the root-constant fast path that entry-point
parameters exist to provide. Plain-data dynamic indexing is served
ABI-invisibly by #11939.
- `rewriteValueUsesToAddrUses` hoisting is a pure move (the context
method now delegates). `isLoadFromImmutableAddress` already treats
`ParameterBlockType` as an immutable address root, so the downstream
`transformParamsToConstRef` passes these addresses into device functions
without re-materializing copies — the whole-array copy cannot be
silently reintroduced.
- Verified locally: emit FileCheck + reflection JSON across all shapes
and the opt-out; all `tests/cuda` pass; full `slang-test` suite clean vs
a pre-change build. **Needs GPU CI / manual L40S verification:** slangpy
`test_benchmark_tensor.py::test_tensor_sum_indirect` parity with the
forced-PB numbers (≈0.026/≈0.195 ms for count 16/32) with the
slangpy#1030 heuristic disabled, the `test_array.py` CUDA suite (the
#11747 crash canaries), and `ptxas -v` no-spill + no dynamic-address
`ld.param` chain. Landing order: shader-slang/slangpy#1045 must merge
before slangpy bumps to a Slang containing this change.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
@szihs
szihs marked this pull request as ready for review July 3, 2026 16:13
@szihs
szihs requested a review from a team as a code owner July 3, 2026 16:13
@szihs
szihs requested review from bmillsNV and ccummingsNV and removed request for a team July 3, 2026 16:13

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1


ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: c6de7480-4a4f-4e7e-8e14-c20040ce7d83

📥 Commits

Reviewing files that changed from the base of the PR and between 20c0170 and 9b7ffe1.

📒 Files selected for processing (5)
  • src/slangpy_ext/utils/slangpy.cpp
  • src/slangpy_ext/utils/slangpytensor.cpp
  • src/slangpy_ext/utils/slangpytorchtensor.cpp
  • src/slangpy_ext/utils/slangpyvalue.cpp
  • src/slangpy_ext/utils/slangpyvalue.h

Comment thread src/slangpy_ext/utils/slangpy.cpp
szihs added a commit to szihs/slangpy that referenced this pull request Jul 7, 2026
…ly on current shader-slang#1045)

Re-applies the DNI workflow (clone floor branch #11939+#11941, merge #11940,
SGL_LOCAL_SLANG, --benchmark-save bypassing Mongo, pwsh on Windows, always-unlock)
on top of the current host-safety shader-slang#1045 head + the cherry-picked param-array benchmark.
szihs added a commit to shader-slang/slang that referenced this pull request Jul 8, 2026
…d at parameter binding

A CUDA compute entry-point uniform parameter that carries a descriptor
table - a fixed-size array of resources or pointer-backed structs - is now
laid out and reflected as an implicit ParameterBlock: the kernel receives
one 8-byte device pointer in kernel-argument space and the payload lives in
device global memory. A runtime index into the array becomes an ordinary
dynamically-addressed global load instead of the serial .param-space load
chain that made such kernels 10-30x slower than Vulkan/D3D12 (issue #11774).

The decision is made exactly once, in computeEntryPointParameterTypeLayout
at parameter-binding time, where reflection is computed; a post-link IR
pass (reconcileCUDAByRefEntryPointParams) then retypes the parameter to
match the recorded layout and rewrites its value uses into loads through
addresses, reusing the constref pass's rewriteValueUsesToAddrUses (exported
as a shared utility). The IR pass consumes the layout decision rather than
re-deriving a predicate, so the emitted kernel signature and the reflected
layout cannot disagree - the failure mode that sank earlier emit-stage
attempts (#11747).

Restricted to struct-typed parameters; bare descriptor arrays keep the
by-value layout and are covered by the dynamic-index local-copy
legalization this PR is stacked on. Ray-tracing stages, torch/
[AutoPyBindCUDA]/[CudaKernel] kernels, and source-written parameter groups
are excluded. -cuda-entry-point-params-by-value restores the legacy ABI.

Stacked on the dynamic-index floor pass (this repo's other CUDA PR);
requires shader-slang/slangpy#1045 before slangpy bumps its Slang pin.
@jhelferty-nv
jhelferty-nv removed the request for review from bmillsNV July 8, 2026 19:19
@ccummingsNV

Copy link
Copy Markdown
Contributor

Is this still needed if we aren't making the 'auto-promote-to-param-block' change?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants