[Fix][S-TIR][DLight] Fall back from non-affine reduction write-back

[Junius-Wynn]

Fixes #20048.

Motivation

A valid Relax conv2d can fail during relax.build(..., target="cuda") when the DLight GPU Reduction rule selects its inner-spatial schedule. Instead of falling through to another schedule rule, the pass raises ScheduleError from bind and aborts compilation.

The Reduction rule first fuses all spatial loops. However, _sch_inner_spatial chooses its threadIdx.x extent using only the original innermost spatial extent. For the reported output shape, the non-unit spatial extents are (2, 2, 20) and the selected thread extent is 10. After rfactor and reverse_compute_at, the write-back block must recover the original axes from the remaining fused loop. This produces bindings that are not recognized as quasi-affine, so bind rejects the block's compact dataflow.

Changes

This PR adds an applicability guard for the known non-affine write-back geometry. Unsupported inner-spatial reductions now return None, allowing GeneralReduction or Fallback to produce a legal schedule. Affine cases continue to use the dedicated Reduction schedule.

The thread-extent selection is moved into a shared helper so that the guard and the scheduling path use the same value. Symbolic spatial extents are handled conservatively.

The default rule-selection loop now catches ScheduleError from an individual rule and continues to the next rule. This prevents an applicability mistake in one schedule rule from becoming a hard compilation failure.

The regression tests cover the reported convolution-like access, a plain reduction without mixed spatial/reduction indices, an affine write-back shape, and transform-level fallback after ScheduleError. They also directly verify that Reduction declines known unsafe shapes instead of relying on the transform-level exception handler.

A strict xfail records a dominant-read spatial-ordering edge case where the shape guard still uses the original block-iterator order. The rule-selection fallback prevents this case from aborting the default schedule chain, while the strict marker ensures that a future guard refinement cannot pass unnoticed.

Testing

The changes were tested in a Python 3.12.13 environment with CUDA 13.0 .

python -m pytest \
  tests/python/s_tir/dlight/test_gpu_reduction.py \
  tests/python/s_tir/dlight/test_gpu_fallback.py -q

Result: 24 passed, 1 xfailed.

The Relax reproduction from #20048 now builds and runs on CUDA. Its result was compared against a NumPy convolution reference:

maximum absolute error: 1.1920928955078125e-07

The previously uncovered relax.op.sum case with input shape (2, 2, 3, 20) and reduction axis 2 also builds successfully through the fallback schedule.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.