Junius-Wynn opened a new pull request, #20057:
URL: https://github.com/apache/tvm/pull/20057
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 .
```bash
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:
```text
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
