Conversation
…le parameter address - Add [Differentiable] and [ForwardDerivative(fwd_getOffset)] to getOffset on all IPointerLikeAddress implementations (BindlessAddress, PointerAddress, TorchTensorViewAddress, Ptr extension) so autodiff can propagate DifferentialPtrPair through address offset computation. - Simplify ILayer.eval from eval(input, weightAddr, biasAddr) to eval(input, parameterAddr). FFLayer.eval now internally computes the bias address via parameterAddr.getOffset(weightCount), keeping address arithmetic inside the differentiable function. - Rename basic-ilayer-ffn-training-test to basic-ilayer-ffn-forward-test and add basic-ilayer-ffn-backward-test that verifies gradients through a 2-layer FFN with getOffset called inside the differentiable function. - Update all test call sites to use the new single-address eval API.
kaizhangNV
requested review from
bmillsNV and
Copilot
and removed request for
a team
📝 WalkthroughWalkthroughConsolidates weights and biases into a single contiguous parameter address for layers and adds forward-mode differentiable support for Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 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 |
There was a problem hiding this comment.
Pull request overview
This PR updates the neural layer parameter addressing model to support autodiff through address arithmetic, by making getOffset differentiable and switching ILayer/FFLayer.eval to accept a single base parameter address for a contiguous (weights-then-bias) parameter block.
Changes:
- Add custom forward derivatives for
getOffsetacrossIPointerLikeAddressimplementations so autodiff can propagateDifferentialPtrPairthrough pointer offsets. - Simplify
ILayer.eval/FFLayer.evaltoeval(input, parameterAddr)and compute bias address internally viaparameterAddr.getOffset(weightCount). - Update/extend neural tests to use the new API and add a backward test that exercises
getOffsetinside the differentiable function.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/gfx-unit-test/neural-tensorview-address.slang | Switch unit test to single-parameter-address eval and route grads through a single DifferentialPtrPair. |
| tests/neural/fflayer-wavetangled-vector-test.slang | Update WaveTangledVector FFLayer tests to the new eval(input, baseAddr) signature. |
| tests/neural/fflayer-two-storage-forward-test.slang | Convert test to contiguous parameter storage and update call sites to the new API. |
| tests/neural/fflayer-no-bias-test.slang | Update no-bias layer tests/wrappers to pass a base parameter address. |
| tests/neural/fflayer-autodiff-backward-test.slang | Update autodiff backward test to exercise getOffset inside eval. |
| tests/neural/basic-ilayer-frontend-smoke-test.slang | Update ILayer-constrained call site to new single-address eval. |
| tests/neural/basic-ilayer-ffn-forward-test.slang | Update 2-layer FFN forward test to use per-layer base offsets. |
| tests/neural/basic-ilayer-ffn-backward-test.slang | New backward test verifying gradients through a 2-layer FFN with internal getOffset. |
| tests/neural/activation-with-fflayer-test.slang | Update activation + FFLayer tests to pass only the base address. |
| source/standard-modules/neural/layers.slang | Change FFLayer.eval signature and compute bias address internally via getOffset. |
| source/standard-modules/neural/ilayer.slang | Change ILayer.eval signature to accept a single contiguous-parameter base address. |
| source/standard-modules/neural/bindless-storage.slang | Add [Differentiable] + [ForwardDerivative] for getOffset on address types. |
| static DifferentialPtrPair<This> fwd_getOffset(DifferentialPtrPair<This> self, int elements) | ||
| { | ||
| return DifferentialPtrPair<This>( | ||
| self.p.getOffset(elements), | ||
| self.d.getOffset(elements)); | ||
| } |
There was a problem hiding this comment.
TorchTensorViewAddress.getOffset is marked [require(cuda)], but its forward derivative helper fwd_getOffset is not. On non-CUDA targets this helper may still be type-checked/compiled and it calls getOffset, which is CUDA-only, potentially causing compilation errors. Mark fwd_getOffset with the same [require(cuda)] (or otherwise ensure it is excluded on non-CUDA targets).
| @@ -1,4 +1,4 @@ | |||
| // Test FFLayer eval() with separate weight/bias address instances. | |||
| // Test FFLayer eval() with contiguous parameter storage. | |||
There was a problem hiding this comment.
This test now verifies contiguous parameter storage rather than separate weight/bias storage, but the file name (fflayer-two-storage-forward-test.slang) still suggests the old behavior. Consider renaming the file to match the updated intent to keep the test suite easier to navigate.
| // Test FFLayer eval() with contiguous parameter storage. | |
| // Test FFLayer eval() with contiguous parameter storage (note: file name | |
| // `fflayer-two-storage-forward-test.slang` is legacy and now refers to this | |
| // contiguous-parameter layout test). |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/neural/basic-ilayer-ffn-forward-test.slang (1)
49-54: Prefer deriving the second layer base fromLayer1.ParameterCount.Line 51 hardcodes the current packing size as
6. Using the layer constant keeps this test aligned if the parameter layout changes.♻️ Suggested cleanup
// Compute base addresses for each layer's parameter block let layer1Addr = baseAddr.getOffset(0); - let layer2Addr = baseAddr.getOffset(6); + let layer2Addr = baseAddr.getOffset(Layer1.ParameterCount);
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 804adb37-6a5c-42a1-9b9c-c186c5912b1d
📒 Files selected for processing (12)
source/standard-modules/neural/bindless-storage.slangsource/standard-modules/neural/ilayer.slangsource/standard-modules/neural/layers.slangtests/neural/activation-with-fflayer-test.slangtests/neural/basic-ilayer-ffn-backward-test.slangtests/neural/basic-ilayer-ffn-forward-test.slangtests/neural/basic-ilayer-frontend-smoke-test.slangtests/neural/fflayer-autodiff-backward-test.slangtests/neural/fflayer-no-bias-test.slangtests/neural/fflayer-two-storage-forward-test.slangtests/neural/fflayer-wavetangled-vector-test.slangtools/gfx-unit-test/neural-tensorview-address.slang
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/standard-modules/neural/bindless-storage.slang (1)
82-89:⚠️ Potential issue | 🟠 MajorKeep
BindlessAddress.getOffset()signed until the final validation/cast.Line 84 stores
baseIndex + elementsintouint, while the othergetOffset()implementations keepelementsin signed pointer arithmetic. BecauseIPointerLikeAddress.getOffset()takesint, this makes negative offsets behave differently forBindlessAddress, andfwd_getOffset()will propagate the same wrapped index into the differential address. Either reject negative offsets consistently at the interface boundary or do the arithmetic in a signed type before the final cast.Possible direction
public This getOffset(int elements) { - uint newBaseIndex = baseIndex + elements; + int newBaseIndex = int(baseIndex) + elements; + // Reject `newBaseIndex < 0` here if bindless addresses are not + // allowed to move backwards. This address = This(handle); - address.baseIndex = newBaseIndex; + address.baseIndex = uint(newBaseIndex); return address; }Also applies to: 91-95
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: c517abeb-f425-4fec-a6c0-9bb675d92729
📒 Files selected for processing (1)
source/standard-modules/neural/bindless-storage.slang
| [Differentiable] | ||
| [ForwardDerivative(fwd_getOffset)] | ||
| public This getOffset(int elements) |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
In current Slang autodiff documentation, can a method marked [Differentiable]legally return a type that implementsIDifferentiablePtrType, or are such types restricted to input parameters only?
💡 Result:
In Slang’s current autodiff docs, a function/method marked [Differentiable] may not return a type that implements IDifferentiablePtrType. Types implementing IDifferentiablePtrType are restricted to input (in) parameters only (i.e., they cannot be used as out, inout, or return types), and attempting to do so is a compile-time error. [1]
getOffset() methods cannot be marked [Differentiable] because they return IDifferentiablePtrType.
In Slang's autodiff system, types implementing IDifferentiablePtrType are restricted to input parameters only and cannot be returned from [Differentiable] methods—this is a compile-time error per the documented rules. The current annotations on these methods violate that restriction.
Also applies to: 141–143, 250–252, 330–332
There was a problem hiding this comment.
PR Review: Add custom derivatives for getOffset
Summary
This PR makes two coupled changes: (1) adds [Differentiable] + [ForwardDerivative(fwd_getOffset)] to getOffset on all four IPointerLikeAddress implementations so autodiff can propagate DifferentialPtrPair through address offset computation, and (2) simplifies ILayer.eval from eval(input, weightAddr, biasAddr) to eval(input, parameterAddr), moving the bias address computation inside the differentiable function. This is a clean design improvement — the old API forced callers to split addresses outside the differentiable function; the new API keeps address arithmetic inside the differentiable scope where it belongs.
What looks good
- Custom derivative pattern is correct. The
fwd_getOffsetimplementations follow the established pattern fromtests/autodiff/diff-ptr-type-custom.slang— constructingDifferentialPtrPair<This>by callinggetOffset(elements)on both.pand.d. no_diff()usage is consistent. Applied toPointerAddressandPtrextension (which do pointer arithmetic the compiler can't differentiate), correctly omitted fromBindlessAddress(uint arithmetic) andTorchTensorViewAddress(body bypassed by custom derivative).- Bias offset calculation is correct.
OutputVector.Size * InputVector.Sizematches theLinearLayoutweight storage convention (weights areOut × Inscalars at offset 0, bias follows immediately). - The
Ptrextension'sfwd_getOffsetusesself.p + elementsinstead of.getOffset(elements), avoiding potential infinite recursion through the derivative chain. Good choice. - New backward test is well-designed. The 2-layer FFN with
getOffset(0)andgetOffset(6)inside the differentiable function directly exercises the custom derivative with non-trivial offsets, and verifies all 9 gradient values. - Test updates are thorough. All existing call sites are updated to the new 2-param API.
Issues
-
[Minor] Missing
[require(cuda)]onTorchTensorViewAddress.fwd_getOffset— see inline comment. All other methods on this type have the attribute. -
[Question]
IPointerLikeAddress.getOffsetinterface declaration lacks[Differentiable]— see inline comment onistorages.slang. This works today due to[ForceInline]onFFLayer.evalcausing specialization before autodiff, but may be fragile for future non-inlined generic usage.
Observations (non-blocking)
- The new backward test only exercises
BindlessAddress.PointerAddresscoverage comes indirectly from the updatedfflayer-autodiff-backward-test.slang(which hasTEST_POINTER=1). ThePtrextension derivative has no runtime test coverage (tests are disabled due to tracked issues #8630, #8631, #8834), which is understandable. - The
(compute, vulkan)category tag on metal/cuda test lines in the new backward test follows the existing convention in the neural test directory — not a copy-paste issue.
Overall this is a well-structured change that both enables a needed autodiff capability and simplifies the public API. The breaking label is appropriate.
|
|
||
| [ForceInline] | ||
| [require(cuda_glsl_hlsl_metal_spirv, sm_6_6)] | ||
| public void atomicAdd(uint index, T value) |
There was a problem hiding this comment.
nit: Missing [require(cuda)] on fwd_getOffset.
The primal getOffset on TorchTensorViewAddress (and the __init, __subscript) are all marked [require(cuda)]. The new fwd_getOffset should be consistent:
| public void atomicAdd(uint index, T value) | |
| [require(cuda)] | |
| static DifferentialPtrPair<This> fwd_getOffset(DifferentialPtrPair<This> self, int elements) |
Without it, a non-CUDA target could theoretically resolve this derivative function even though the primal is CUDA-only, producing a confusing error instead of a clean capability mismatch.
| // | ||
| //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUFFER):-vk -compute -shaderobj -xslang -experimental-feature -output-using-type -emit-spirv-directly | ||
| //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUFFER):-mtl -compute -shaderobj -output-using-type -xslang -experimental-feature | ||
| //TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUFFER):-cuda -compute -shaderobj -output-using-type -capability cuda_sm_7_0 -xslang -experimental-feature |
There was a problem hiding this comment.
suggestion: Consider adding a PointerAddress test variant as well.
This backward test only exercises BindlessAddress. The existing fflayer-autodiff-backward-test.slang does cover PointerAddress (via TEST_POINTER=1), and that test is updated in this PR to call getOffset inside the differentiable function. However, since this new test is the primary one that exercises multi-layer getOffset inside a differentiable context (with non-trivial offset=6), a TEST_POINTER=1 CUDA variant here would strengthen coverage of PointerAddress.fwd_getOffset in the multi-layer scenario.
Not blocking — the existing backward test suite provides reasonable coverage.
Add custom derivatives for getOffset and simplify ILayer.eval to single parameter address
Add [Differentiable] and [ForwardDerivative(fwd_getOffset)] to getOffset on all IPointerLikeAddress implementations (BindlessAddress, PointerAddress, TorchTensorViewAddress, Ptr extension) so autodiff can propagate DifferentialPtrPair through address offset computation.
Simplify ILayer.eval from eval(input, weightAddr, biasAddr) to eval(input, parameterAddr). FFLayer.eval now internally computes the bias address via parameterAddr.getOffset(weightCount), keeping address arithmetic inside the differentiable function.
Rename basic-ilayer-ffn-training-test to basic-ilayer-ffn-forward-test and add basic-ilayer-ffn-backward-test that verifies gradients through a 2-layer FFN with getOffset called inside the differentiable function.
Update all test call sites to use the new single-address eval API.