Fuse reduction + transpose of shared input (save full memory pass)

[eellison]

Problem

60 repros emit 2-3 kernels for the pattern:

permute_default = permute(x, [1, 0])     # weight grad [N, C] → [C, N]
sum_dim_int_list = sum(x, dim=[0])        # bias grad [N, C] → [C]

Both read the same large input. Current approach reads it twice (once per kernel). Fused kernel reads once, writes both outputs.

Verified Speedup

1.69x — fused Triton kernel hits 99% of memcopy SOL (0.103ms vs 0.174ms for torch.compile on f32[25216, 3072]).

Prototype: bench_fused_transpose_sum.py — tiles [64,64], uses tl.trans() for coalesced transpose write + tl.atomic_add for partial sum accumulation.

Root Cause in Inductor

The scheduler blocker is tiling_prevents_reduction_fusion at simd.py:1458-1470. select_tiling is called on the permute as a purely pointwise op (reduction_numel=1), returning a 2-tuple. The check expects a 3-tuple with r0_: 1 suffix.

However, initial attempt to add (numel2, rnumel2) to the accepted set didn't reduce kernel count — there may be an earlier rejection. Needs debug tracing through can_fuse to find the actual blocker.

Key Insight

The transpose store can happen INSIDE the reduction loop body. Each tile loaded for the reduction also gets written transposed. fits_in_main_body (simd.py:1493) already supports pointwise stores inside reduction loops when node_numel == numel * rnumel. The scheduler just needs to recognize the permute qualifies.

Impact

60 repros, tensors up to 6.6GB. Saves one full memory pass per repro.

Repros from #26

85 repros from the finalized #26 checklist.

• sum_sum_sum_d06bf12e10d0 — 28 kernels, 1.72x gap, 2617MB
• sum_sum_sum_55426f9a4493 — 28 kernels, 1.72x gap, 2483MB
• sum_sum_4bd81dea302d — 9 kernels, 3.30x gap, 126MB
• sum_sum_sum_8d668b8fcd9e — 7 kernels, 3.14x gap, 1617MB
• sum_sum_sum_1d1edd983280 — 7 kernels, 3.12x gap, 1617MB
• sum_sum_sum_83ddb84cc830 — 7 kernels, 2.88x gap, 403MB
• sum_sum_sum_9e7a546b859f — 7 kernels, 1.92x gap, 101MB
• sum_sum_sum_64f701d26f0a — 6 kernels, 2.98x gap, 388MB
• sum_sum_sum_e2388f04f7c2 — 6 kernels, 2.68x gap, 264MB
• sum_sum_sum_4fdc0c9d691f — 6 kernels, 2.68x gap, 19MB
• sum_c2895e89552b — 6 kernels, 2.66x gap, 138MB
• sum_12b455c3b1b0 — 6 kernels, 2.61x gap, 25MB
• sum_sum_sum_5c7c5e63becb — 6 kernels, 2.42x gap, 89MB
• sum_sum_sum_a9940aaceefe — 6 kernels, 2.11x gap, 220MB
• sum_5eec4697d8ff — 5 kernels, 2.60x gap, 72MB
• sum_sum_sum_9814418db3a9 — 5 kernels, 2.58x gap, 206MB
• sum_79c6d0673ca2 — 5 kernels, 2.48x gap, 41MB
• sum_sum_sum_580d823755af — 5 kernels, 2.47x gap, 206MB
• sum_1457ecad8f5c — 5 kernels, 2.37x gap, 189MB
• sum_sum_sum_6107a2f54029 — 4 kernels, 5.86x gap, 824MB
• sum_sum_sum_dc96c4651516 — 4 kernels, 5.76x gap, 824MB
• sum_sum_sum_48033561d121 — 4 kernels, 3.44x gap, 227MB
• sum_sum_sum_1b02daac4062 — 4 kernels, 3.35x gap, 103MB
• sum_sum_sum_da376d09d2cb — 4 kernels, 3.34x gap, 103MB
• sum_sum_sum_81b1543349e5 — 4 kernels, 3.33x gap, 103MB
• sum_sum_sum_3213336f4c0b — 4 kernels, 3.22x gap, 157MB
• sum_fd44d96c0b10 — 4 kernels, 3.11x gap, 164MB
• sum_86b5dc92d54f — 4 kernels, 2.66x gap, 41MB
• sum_sum_sum_f2c6042682bc — 3 kernels, 4.25x gap, 76MB
• sum_sum_sum_1d10e6314ef6 — 3 kernels, 3.22x gap, 107MB
• sum_sum_sum_cb22a50c0d9f — 3 kernels, 3.04x gap, 214MB
• sum_sum_sum_559e754b2ff4 — 3 kernels, 2.74x gap, 15MB
• sum_sum_sum_e2d4961f3571 — 3 kernels, 2.73x gap, 202MB
• sum_sum_sum_69414585b76b — 3 kernels, 2.66x gap, 214MB
• sum_sum_sum_97999676281e — 3 kernels, 2.14x gap, 312MB
• sum_sum_sum_1c9995bc990f — 3 kernels, 2.09x gap, 214MB
• sum_sum_sum_728df9734dbe — 3 kernels, 2.02x gap, 6MB
• sum_13195092a57b — 3 kernels, 2.01x gap, 37MB
• sum_sum_sum_baf315cfc5f0 — 3 kernels, 2.00x gap, 201MB
• sum_sum_sum_59e022113eeb — 3 kernels, 2.00x gap, 403MB
• sum_66b92a2b30bb — 3 kernels, 1.98x gap, 37MB
• sum_sum_sum_766837223ded — 3 kernels, 1.98x gap, 2MB
• sum_cfe874779dd6 — 3 kernels, 1.93x gap, 3294MB
• sum_sum_75ee5fc5ea19 — 3 kernels, 1.91x gap, 16003MB
• sum_cb19b2c26400 — 3 kernels, 1.74x gap, 13MB
• sum_046be72c7ad9 — 3 kernels, 1.72x gap, 6589MB
• sum_60071d2e71f2 — 3 kernels, 1.70x gap, 6MB
• sum_d657e5a1c5df — 3 kernels, 1.70x gap, 14MB
• sum_0568508d372d — 3 kernels, 1.67x gap, 50MB
• sum_a630880e528a — 3 kernels, 1.67x gap, 780KB
• sum_2821c003f795 — 3 kernels, 1.64x gap, 42MB
• sum_27f326f22f76 — 3 kernels, 1.62x gap, 129KB
• sum_df75f19c9c7a — 3 kernels, 1.61x gap, 188KB
• sum_31bf563cdf96 — 3 kernels, 1.59x gap, 465MB
• sum_f160d1f03c1f — 3 kernels, 1.58x gap, 465MB
• sum_sum_sum_b42d335bef54 — 3 kernels, 1.55x gap, 82MB
• sum_587e52500b08 — 3 kernels, 1.55x gap, 467MB
• sum_558522245da0 — 3 kernels, 1.55x gap, 467MB
• sum_51676b1433c8 — 3 kernels, 1.52x gap, 403MB
• sum_7f12bfc45e1a — 3 kernels, 1.52x gap, 403MB
• sum_7c50e5b80bfe — 3 kernels, 1.52x gap, 8KB
• sum_7953a8b0bbba — 3 kernels, 1.49x gap, 12MB
• sum_2bcc7099936f — 3 kernels, 1.47x gap, 302MB
• sum_6080543764bd — 3 kernels, 1.44x gap, 4MB
• sum_bd2fed913c26 — 3 kernels, 1.44x gap, 101MB
• sum_33b494a50932 — 3 kernels, 1.43x gap, 302MB
• sum_ab7572b27d2c — 3 kernels, 1.43x gap, 101MB
• sum_214ee1a84aa5 — 3 kernels, 1.43x gap, 101MB
• sum_b8db5e701976 — 3 kernels, 1.41x gap, 101MB
• sum_426d5bebb63f — 3 kernels, 1.40x gap, 930MB
• sum_196d0dc5c868 — 3 kernels, 1.39x gap, 101MB
• sum_62dca72fabd1 — 3 kernels, 1.35x gap, 50MB
• sum_f95857722916 — 3 kernels, 1.33x gap, 654MB
• sum_6295c187c71d — 3 kernels, 1.31x gap, 50MB
• sum_fe2660cedc0a — 3 kernels, 1.30x gap, 214MB
• sum_80634958a126 — 3 kernels, 1.24x gap, 336MB
• sum_sum_7af057233a52 — 2 kernels, 2.78x gap, 78MB
• sum_sum_f4d29f9ee6ad — 2 kernels, 1.70x gap, 3933MB
• sum_sum_666d023699ba — 2 kernels, 1.51x gap, 1475MB
• sum_096e9f1570b5 — 2 kernels, 1.46x gap, 13MB
• sum_sum_cb1e4e3e8236 — 2 kernels, 1.44x gap, 2857MB
• sum_sum_c769ee8f5293 — 2 kernels, 1.41x gap, 692MB
• sum_b57af4bf1f43 — 2 kernels, 1.39x gap, 94MB
• sum_sum_sum_34c857ab7db3 — 2 kernels, 1.18x gap, 3758MB
• sum_34445852895e — 2 kernels, 0.94x gap, 67MB

[eellison]

Progress

Branch: eellison/pytorch:reduction-transpose-fusion

Fix: Modified _try_reindex_pointwise_for_reduction in scheduler.py to use apply_new_loop_order (preserves buffer strides) instead of apply_loop_reindexing (creates flat index that breaks strides).

Result: 3 kernels → 1 kernel, correctness verified.

Performance:

• Hand-written optimal Triton kernel: 0.104ms (hits SOL, 6x faster than current)
• Inductor with fix (1 fused kernel): 0.232ms (slower than 3-kernel baseline of 0.175ms)
• The fix enables fusion but inductor's auto-tiling chooses a bad config for the fused kernel

What's needed next: The fused kernel needs the tiled transpose approach (BLOCK_M=64, BLOCK_N=64, tl.trans() for coalesced stores, tl.atomic_add for reduction). Inductor's default tiling doesn't know about the transpose access pattern. Options:

1. Teach select_tiling to recognize the transposed store pattern and pick appropriate tile sizes
2. Add a transpose-aware autotuning config space
3. Template-based approach for the reduction+transpose pattern

The ceiling (0.104ms) is 6x better than current. The scheduler fix is correct — it's the codegen tiling that needs work.

[eellison]

Final result: 1.87x speedup via generic cooperative reduction

No pattern match needed. The fix is two parts:

1. Scheduler: _try_reindex_pointwise_for_reduction with apply_new_loop_order enables the transpose to fuse into the reduction kernel (1 kernel instead of 2-3).

2. Codegen: Cooperative reduction with tuned config (XBLOCK=16, R0_BLOCK=256, relaxed eligibility for atomic cooperative mode) gives coalesced loads and efficient reduction.

Performance

Version	Time
Baseline (2-3 kernels)	0.231ms
Generic cooperative (tuned)	0.123ms
Hand-written optimal	0.104ms

1.87x speedup, 16% gap to ceiling. The remaining gap is non-coalesced transposed stores — closing it requires a fundamentally different grid layout (2D tile-and-reduce with atomic accumulation). Not worth the complexity for 16%.

Changes

• scheduler.py: loop reorder fallback in _try_reindex_pointwise_for_reduction
• choices.py: relaxed cooperative_reduction eligibility for atomic mode when tiling_scores show r-dominant pattern
• triton_heuristics.py: tuned RBLOCK/RSPLIT for atomic cooperative (larger blocks, fewer splits)
• triton.py: should_use_cooperative_reduction extended for transpose patterns
• simd.py: moved tiling_scores before cooperative_reduction init

Validation

• Correctness: transpose exact match, sum error < 0.001 (fp32 atomic accumulation)
• No regressions on 47-repro sample (with tightened heuristic: rhint >= 8192)

[eellison]

Local default-off prototype for this bucket is now packaged as a draft PR against my fork, not upstream PyTorch:

eellison/pytorch#1

Current best evidence is in scratch_issue21_benchmark_validation_20260529/hardened_32m_paired_20260529_130715:

• sum_sum_sum_6107a2f54029: 811.0us -> 180.0us, kernels 4 -> 1
• sum_sum_sum_dc96c4651516: 823.4us -> 184.7us, kernels 4 -> 1
• sum_sum_sum_baf315cfc5f0: 74.7us -> 71.4us, kernels 3 -> 1
• sum_sum_sum_3213336f4c0b now rejects under the 32MiB guard and is neutral: 111.4us -> 110.6us, kernels 4 -> 4

The path is still default-off because the profitability model is intentionally conservative and not ready for default-on.

[eellison]

Simplified the fork PR for this bucket so it is now mixed-order-only and no longer includes the marginal simple-producer path.

PR: eellison/pytorch#1
Head: c774b530fc8

Latest tightened sweep artifact: scratch_issue21_guarded_validation_simplified_64m_20260529

Results from the 64MiB default-off gate sweep:

• sum_sum_sum_6107a2f54029: 0.8357ms -> 0.1793ms, kernels 4 -> 1
• sum_sum_sum_dc96c4651516: 0.8098ms -> 0.1766ms, kernels 4 -> 1
• sum_sum_sum_e2d4961f3571: 0.0788ms -> 0.0739ms, kernels 3 -> 1
• sum_sum_sum_59e022113eeb: 0.1115ms -> 0.1066ms, kernels 3 -> 1
• sum_sum_sum_97999676281e: 0.0933ms -> 0.0907ms, kernels 3 -> 1
• sum_sum_sum_64f701d26f0a: 0.1725ms -> 0.1679ms, kernels 6 -> 4

The earlier 32MiB gate selected one small regression (sum_sum_sum_69414585b76b, 3 -> 1 kernels but 0.085ms -> 0.090ms), so I tightened the default-off min-size gate to 64MiB. The old 32MiB boundary case sum_sum_sum_baf315cfc5f0 now rejects. The 30-candidate 64MiB sweep had 30/30 correct, 6 selected, and no selected regressions.

This is the simpler shape: existing mixed-order producer + extra producer-body saved sum partials; no generic nested-reduction, split-reduction, or simple-producer path.

[eellison]

Closing as part of the bf16 corpus migration cleanup. These predate the corpus flip — numbers/hashes are from the old fp32 corpus, and several may already be addressed by in-flight Inductor changes (tmp_work branch). Perf/codegen items will be re-derived and re-prioritized by aggregate model-level impact via the attribution pipeline, then re-filed with current data. Reopen if still relevant.