Repository navigation
adj_tile_dot is missing a block-wide broadcast of adj_ret #1374
Description
adj_tile_dot (warp/native/tile_reduce.h:729-756) uses each thread's local adj_ret directly, but in reverse mode adj_ret is only non-zero on the thread(s) that consumed the forward scalar. When the idiomatic "one thread writes the result" pattern is used, only that thread receives a gradient, and the tile elements owned by every other thread get zero.
Repro
import warp as wp
BLOCK_DIM = 64
N = 16
@wp.kernel
def compute(a_in: wp.array(dtype=float),
b_in: wp.array(dtype=float),
out: wp.array(dtype=float)):
i, j = wp.tid()
a = wp.tile_load(a_in, shape=N, offset=0, storage="register")
b = wp.tile_load(b_in, shape=N, offset=0, storage="register")
result = wp.tile_dot(a, b)
if j == 0:
out[i] = resultWith a = [1..16], b = [16..1], after tape.backward():
a_in.grad: [16. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
expected: [16. 15. 14. 13. 12. 11. 10. 9. 8. 7. 6. 5. 4. 3. 2. 1.]
Forward value is correct (816.0); only the gradient is wrong.
Why the existing test_tile_dot_grad* tests don't catch it
They use 1-D i = wp.tid() in a tiled launch. launch_coord divides the linear index by block_dim before unraveling, so wp.tid() returns the tile index, which is the same for every thread in the block. With dim=[1], i == 0 is true on all 64 threads, every thread executes adj_array_store, and every thread pulls adj_out[0] into its own adj_ret. The gate never narrows execution to a single thread, so the tests only validate the uniform-adj_ret case.
Why this is structural to tile_dot and not other reductions
adj_tile_sum(whole-tile,tile_reduce.h:627) already broadcastsadj_retfrom thread 0 through shared memory.adj_tile_sum<Axis>(tile_reduce.h:805) broadcasts viatile_broadcast+tile_add_inplace.adj_tile_reduce_impl,adj_tile_max,adj_tile_min,adj_tile_arg{max,min}are empty stubs.adj_tile_axpyhandles the symmetric scalar-input case via a single-thread accumulate guard onadj_alpha.
tile_dot is the only tile primitive whose forward returns a raw scalar broadcast to every thread (via __shared__ ScalarT scratch at tile_reduce.h:696 and :713-717). The forward path has a broadcast; the reverse path does not.
Context
Originally flagged during review of MR !2216 (the MR that added wp.tile_dot / wp.tile_axpy â GH-1363, GH-1364) but deferred to a follow-up so the primitives could land.
Activity
wp.tile_dot()for fused element-wise multiply and reduce #1364Fix missing block-wide broadcast in adj_tile_dot (NVIDIAGH-1374)
Make tile_dot return a single-element tile instead of a scalar (NVIDIâĶ
Add tile_dot gradient tests for single-thread write pattern (NVIDIAGHâĶ
Make tile_dot return a single-element tile instead of a scalar (NVIDIâĶ
shi-eric commentedon Apr 20, 2026
Fixed by b9af381.
Rather than patching
adj_tile_dotto broadcastadj_ret(as the issue title suggested), the fix restructurestile_dotitself to match thetile_sum/tile_reduceconvention, which side-steps the "which thread holds the canonical gradient" question entirely.Forward change:
tile_dotnow returns a single-element register tile (tile(dtype=T, shape=(1,))) instead of a raw scalar broadcast to all threads. The shared-memory broadcast that previously happened insidetile_dotis gone; only the reducing thread writes tooutput.data[0].Adjoint change:
adj_tile_dotnow takes a 1-element tile foradj_retand reuses the same__shared__ scratchbroadcast pattern asadj_tile_sum, so every thread sees the correct gradient regardless of which thread wrote the forward result.User-facing API change (vs. the originally-merged
wp.tile_dot): Becausetile_dothad not shipped in a release yet (still inUnreleased), we took the opportunity to change the return type rather than keep the scalar-returning signature. Callers must now index the result:The documented equivalent long-form also simplifies, from
wp.tile_extract(wp.tile_sum(wp.tile_map(wp.tensordot, a, b)), 0)towp.tile_sum(wp.tile_map(wp.tensordot, a, b)).tile_axpy, which internally callstile_dot, was updated to readdot_result.data[0].New regression tests
test_tile_dot_grad_extractandtest_tile_dot_grad_extract_sharedcover the single-thread-writes-the-scalar pattern from the repro in this issue.