cudastereo: TMA path-aggregation kernels for StereoSGM on sm_100+

[DmytroHilei]

cudastereo: TMA path-aggregation kernels for StereoSGM on sm_100+

Motivation

The path-aggregation stage of cv::cuda::StereoSGM::compute runs eight
1-D DP recurrences (vertical U2D/D2U, horizontal L2R/R2L, four oblique
directions), one per direction, each on its own stream. Per-direction
ncu on the current legacy kernels (1920×1080, MD=128, RTX 5060 Laptop
sm_120) shows all of them DRAM-bound with the warp stalled on
long_scoreboard after synchronous LDGs — the block-wide
LDG + __syncthreads pattern pins compute order so the scheduler can't
hide the load.

Hopper introduced TMA (cp.async.bulk.shared::cta.global +
mbarrier.try_wait.parity) precisely for this shape of workload:
asynchronously stream bytes DRAM → smem while the warp keeps computing.
Blackwell (sm_100+) accepts the .shared::cta completion form used
here; a Hopper (.shared::cluster) port is a separate follow-up.

What this PR does

1. Safe fixes (carried forward from PR
   cudastereo-sgm-consistency-fix / commit b590fde):

   • check_consistency launch bounds: dim3(width/16, height/16) →
     dim3(divUp(width, 16), divUp(height, 16)) + in-kernel x >= width || y >= height early-return. Previously the right/bottom strip up
     to 15 px wide was never visited on non-multiple-of-16 images; on
     1920×1080 (1080 % 16 = 8) this affected the bottom 8 rows.
   • store_uint8_vector<2u>: fix uint8x2.y = ptr[0] → ptr[1]
     typo. Unreachable in current code (DP_BLOCK_SIZE ∈ {8, 16}) but
     latent for any future 2-element variant.
   • Adds regression test StereoSGM_NonAligned in test_stereo.cpp
     that verifies the strip is now touched by check_consistency.

2. TMA path-aggregation kernels, one per direction group:

   • aggregate_horizontal_path_kernel_tma — chunked prefetch,
     CHUNK_X = 128 pixels/phase, staging = 8 rows × 132 int32s ×
     2 (L+R) × 2 (double-buffer) ≈ 16.5 KiB smem/block. 16 bulk copies
     per phase transition (PATHS_PER_BLOCK threads issue L + R for
     their own row).
   • aggregate_vertical_path_kernel_tma — per-row ping-pong, one
     bulk copy each for the per-iter right slice + left slice, issued
     by thread 0. Two mbarriers, try_wait.parity gate. Existing
     swizzled right_buffer layout preserved on the consumer side.
   • aggregate_oblique_path_kernel_tma — same as vertical but with
     the diagonal right_x0 ± 1 per iter alignment quirk: per iter
     round the source pointer DOWN by d = ((right_x0 % 4) + 4) % 4
     ints and shift the consumer read index UP by d. Staging sized
     with ALIGN_SLACK = 4 for the slack.

3. Runtime dispatch, added to the vertical + horizontal launcher
   wrappers (4 of the 8) via a shared helper instead of a repeated
   if/else at each call site:

   const bool useTma = canUseTma() && (min_disp % 4 == 0);
   launchAggregationKernel(useTma, gdim, bdim, stream,
       aggregate_horizontal_path_kernel_tma<1, MAX_DISPARITY>,
       aggregate_horizontal_path_kernel<1, MAX_DISPARITY>,
       left, right, dest, left.cols, left.rows, p1, p2, min_disp);

   canUseTma() caches `cudaDeviceGetAttribute(cudaDevAttrComputeCapabilityMajor)

   = 10on first call. sm_9x and below always take the legacy path (no ABI change, no test regression).min_disp % 4 == 0is the TMA source-pointer alignment requirement; unaligned callers fall through to legacy.launchAggregationKernelis a small variadic template (pickstmaKernelorlegacyKernelbased onuseTma`,
   both launched with the same arg pack) so the dispatch decision
   lives in one place rather than repeated per call site.

   The 4 oblique launcher wrappers deliberately do not use this
   dispatch — they always call the legacy kernel. aggregate_oblique_path_kernel_tma
   is implemented, bit-exact, and a real isolated per-kernel win, but
   oblique only runs in HH8 mode where it contends with 7 other
   concurrent streams; under that contention its tail latency roughly
   triples relative to legacy (see Performance below), which flips
   HH8's end-to-end result from a win into a regression if oblique TMA
   is dispatched. Leaving oblique on the legacy path and keeping
   vertical + horizontal TMA on turns HH8 back into a (small) net win.

4. Compile-time guards: each *_tma kernel body is wrapped in
   #if __CUDA_ARCH__ >= 1000 so ptxas accepts the file when compiled
   for older archs (empty body on those archs; runtime never dispatches
   to it).

5. Shared device helpers: the mbarrier ping-pong init
   (tma_init_ping_pong_barriers) and the per-phase wait-then-flip-parity
   step (tma_wait_and_flip) are each defined once and reused by all
   three *_tma kernels, instead of being duplicated inline in each.

Correctness

TMA kernels are bit-exact with their legacy counterparts on the
1920×1080 real-image cost volume (0 differing bytes / 265 MB per
direction; verified externally via the benchmark harness at
benchmarks/cuda_stereosgm/, not shipped with this PR).

Since bit-exactness is designed in, no TMA-specific test is added — the
existing CudaStereo_StereoSGM.regression test (which compares
sgm.compute() output against a reference disparity image) already
covers this. Any TMA regression fails that test.

Cross-arch build verification

The *_tma kernel bodies use cp.async.bulk.shared::cta /
mbarrier.try_wait.parity, which ptxas only accepts for sm_100+. The
#if __CUDA_ARCH__ >= 1000 guards must therefore fully exclude those
bodies when compiling for older archs, or the build breaks on exactly
the targets CI uses (CI does not build sm_120). Verified by compiling
stereosgm.cu with the same nvcc flags for each arch:

target	result
sm_75 (Turing)	compiles clean
sm_86 (Ampere)	compiles clean
sm_90 (Hopper)	compiles clean
fat binary 75/86/90 + PTX (CI-style multi-gencode)	compiles clean

The guard boundary was confirmed at the SASS level, not just by exit
code: the sm_90 / sm_86 objects contain zero bulk-copy
instructions (the Blackwell-only PTX is fully excluded), while the
sm_120 object contains 48 UBLKCP instructions and 12 *_path_kernel_tma
symbols — so the guarded code is emitted where it should be and absent
everywhere else.

Performance

Numbers below are RTX 5060 Laptop (sm_120, 26 SMs), 1920×1080,
MAX_DISPARITY = 128 (cudaEvent elapsed-time), warmup = 400, bench =
3000 for every kernel and the e2e pipeline (bumped from an earlier
200/1000 pass specifically to pin down whether the HH8 e2e result below
was noise — it wasn't). CSVs in data/.

Isolated per-kernel

kernel	direction	legacy median (ms)	tma median (ms)	delta
vertical	U2D	1.526	1.441	−5.6%
vertical	D2U	1.545	1.449	−6.2%
oblique	UL→DR	1.559	1.506	−3.4%
oblique	UR→DL	1.567	1.527	−2.6%
oblique	DR→UL	1.604	1.560	−2.8%
oblique	DL→UR	1.609	1.505	−6.5%
horizontal	L2R	1.523	1.344	−11.7%
horizontal	R2L	1.513	1.367	−9.7%

Aggregate per-direction wins (mean of the direction rows above):

direction group	legacy (ms)	tma (ms)	delta
horizontal	1.518	1.356	−10.7% — biggest per-kernel win in the PR
vertical	1.536	1.445	−5.9%
oblique	1.585	1.524	−3.8% — real, but see Performance/HH8 below

CHUNK_X (horizontal's TMA staging window, currently 128) was swept
64/128/192 — 128 is the validated sweet spot (dominates on L2R
specifically, −19.7% isolated at that setting); 192 regresses
(−14.3% aggregate) and 64 is a noise-driven false peak. 256 overflows
the sm_120 static smem budget. Not an arbitrary first guess.

Isolated per-kernel: ncu (SpeedOfLight + MemoryWorkloadAnalysis)

Oblique's ncu is in the branch report (§6); horizontal's below
(both directions, results/tma/horizontal/ncu/ in the benchmark
harness) is new this round:

metric	legacy L2R	tma L2R	legacy R2L	tma R2L
DRAM Throughput	41.4%	45.2%	42.3%	46.6%
Compute (SM) Throughput	64.3%	58.3%	62.2%	62.8%
L1/TEX Hit Rate	29.1%	0.4%	29.1%	1.8%
Warp Cycles / Issued Inst	6.42	5.65	7.06	5.36
Achieved Occupancy	38.9%	34.4%	40.1%	34.2%

Horizontal's TMA win runs through the opposite mechanism from
oblique's. Oblique's L1/TEX hit rate rises with TMA (5.9%→42%,
better cache locality). Horizontal's collapses (29.1%→0.4-1.8%):
cp.async.bulk routes the bulk read traffic DRAM/L2→smem directly
through the TMA hardware unit, bypassing the L1TEX-mediated LDG path
that gave legacy its real 29% hit rate (a per-pixel register-shift
reuse trick). DRAM Throughput still rises even as L1TEX traffic
vanishes, and Warp Cycles/Issued Instruction drops — both confirm TMA
decouples the load from the warp's critical path rather than
improving what L1TEX sees. Achieved Occupancy drops ~5pp (ping-pong
staging smem cost) — the one thing this design spends to buy the rest.

End-to-end (sgm.compute pipeline)

tma below = vertical + oblique TMA (no horizontal). tma_full = all
three. tma_no_obl (HH8 only) = vertical + horizontal TMA, oblique
left on the legacy path — this is what the PR actually ships.

mode	variant	median (ms)	std (ms)	delta vs legacy median
HH4	legacy	9.854	0.306	—
HH4	tma (vert+obl only)	9.814	0.269	−0.4%
HH4	tma_full (+ horiz)	9.629	0.297	−2.3%
HH8	legacy	17.674	0.502	—
HH8	tma (vert+obl only)	18.251	0.503	+3.3%
HH8	tma_full (+ horiz)	18.136	0.506	+2.9%
HH8	tma_no_obl (shipped)	17.501	0.541	−1.0%

Bit-exact vs legacy in every variant (0 / 2,073,600 differing pixels).
HH4 never invokes oblique, so it only has 3 rows.

HH8 with oblique TMA dispatched is a real, reproducible regression —
not noise. First pass at 200/1000 warmup/bench looked flat-within-1σ
(legacy std was 0.75ms, wide enough to plausibly cover the gap); at
400/3000 the std tightened to ~0.5ms and the same +2.7-3.3% gap held
in two independent runs. Every isolated oblique/vertical/horizontal
median improves with TMA, so a per-kernel view alone predicts an HH8
win that doesn't show up.

Root cause, confirmed via nsys --report cuda_gpu_kern_sum (50
profiled HH8 iterations/variant, isolated with
cudaProfilerStart/Stop + --capture-range=cudaProfilerApi): under
HH8's 8 concurrent streams, TMA kernels' per-launch duration variance
is ~3× higher than legacy's (coefficient of variation ~20-32% for
TMA vs ~6-12% for legacy, oblique the worst offender at ~29-32%).
Tellingly, the legacy horizontal kernel — unmodified, unchanged code
— shows the same inflated variance (~20-25% CV) when co-scheduled
alongside TMA vertical+oblique, versus ~7-11% when it's the only thing
contending with other legacy kernels. That rules out "TMA kernels are
just noisier code" and points at real cross-stream contention: sharing
SMs/L2/DRAM with concurrent TMA streams fattens everyone's tail, TMA's
tail more than legacy's. WTA can't start until all 8 streams finish,
so HH8 e2e tracks max() across those 8 durations, not the mean — a
fatter right tail raises that max even when the mean improved
(order-statistics behavior, not a contradiction).

Fix: since oblique showed the worst tail inflation and is HH8's
only TMA-exclusive contributor (vertical + horizontal both help in
HH4 too, where there's no 8-way contention), we tested dropping
oblique back to the legacy kernel while leaving vertical + horizontal
TMA on. That's the tma_no_obl row above: HH8 flips from +2.9%
regression to −1.0%, a real if modest win, consistent with HH4.
This PR ships that configuration — the 4 oblique launcher wrappers
always call the legacy kernel; aggregate_oblique_path_kernel_tma
stays in the file, bit-exact and functional, as a documented base for
a future contention fix (e.g. CUDA stream priorities) rather than
being deleted.

• The per-kernel horizontal win (−10.7% aggregate, −11.7% on L2R) is
  the largest and most reproducible signal in the PR, and — unlike
  oblique — it holds up end-to-end in both HH4 and HH8 because
  horizontal only has 2 directions, not 4, and isn't exclusive to the
  8-stream mode.
• Vertical's −5.9% isolated win also survives contention (tma_no_obl
  includes it and beats legacy at HH8), unlike oblique's −3.8%, which
  doesn't.

Cross-device check: RTX 5090 (170 SMs vs the laptop's 26)

Reproduced via run_pr_ab.sh on a rented RTX 5090 (also sm_120), same
image/params, same 400/3000 warmup/bench. Bit-exact in every variant
(0 differing pixels, confirmed in run.log). CSVs in
data/*_rtx5090.csv.

	5060 Laptop (26 SMs)	RTX 5090 (170 SMs)
horizontal L2R legacy→v1	−11.7%	−37.1%
horizontal R2L legacy→v1	−9.7%	−48.5%
horizontal v2 (tensor TMA) vs v1	v2 loses by ~12%	v2 wins (−8.7%/−3.9% vs v1)
HH4 e2e legacy→tma_full	−2.3%	−5.1%
HH8 e2e legacy→tma_full (all TMA incl. oblique)	+2.9% (regression)	−1.8% (win)
HH8 e2e legacy→tma_no_obl (shipped)	−1.0%	−0.4%

Horizontal's win scales up sharply on the bigger chip — expected, more
SMs means more concurrent DRAM requests in flight, which is exactly
what cp.async.bulk's decoupled prefetch is designed to exploit.

The oblique finding doesn't fully transfer. On the 5090, tma_full
(oblique TMA dispatched) beats tma_no_obl at HH8 — the opposite of
the laptop result that justified shipping oblique on the legacy path.
170 SMs apparently gives HH8's 8 concurrent streams enough headroom
that oblique's contended tail latency (Performance/HH8 above) doesn't
dominate the way it does with 26 SMs. The shipped dispatch (no oblique
TMA) is not wrong here — it's still a real win on both devices
(−1.0% / −0.4%) — but it leaves ~1.4 points of additional HH8 win on
the table on high-SM-count Blackwell parts. Not making the dispatch
SM-count-aware in this PR: two data points isn't enough to pick a sane
threshold, and a wrong guess risks a device-dependent perf cliff worse
than the modest performance left on the table by staying conservative.

Compatibility

• sm_100+ (Blackwell): TMA kernels dispatched.
• sm_9x (Hopper): legacy path. TMA .shared::cta isn't accepted
  on sm_90 by ptxas; a .shared::cluster port with different
  completion semantics would unlock Hopper — follow-up, not this PR.
• sm_8x (Ampere / Ada / older): legacy path.
• No ABI or public API change. No new build-system dependencies.

Files touched

modules/cudastereo/src/cuda/stereosgm.cu   +~800 lines
modules/cudastereo/test/test_stereo.cpp    +~92 lines (safe-fixes test)