[Feat][Pool] Implement adaptive 2D pooling ops - #1809
Open
RMLYC wants to merge 18 commits into
Open
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
4 tasks
RMLYC
marked this pull request as draft
July 30, 2026 08:09
RMLYC
marked this pull request as ready for review
July 31, 2026 06:15
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
RMLYC
force-pushed
the
feat/pool/adaptive-pool2d
branch
from
July 31, 2026 06:33
67c1b2e to
d8b083d
Compare
Mirror the manifest formula update: per-axis scan counts include the one-row/col overlap of adjacent adaptive bins on non-divisible grids. Divisible cases (incl. the existing snapshot) are unchanged.
Per the torch 2.13 reference docs (output_size: int | None | tuple[int | None, int | None]). Scalar None normalizes to (None, None), i.e. output equals the input spatial size. Tests map the reference call to (None, None) because the torch 2.10 runtime rejects scalar None.
lcy-seso
force-pushed
the
feat/pool/adaptive-pool2d
branch
3 times, most recently
from
August 1, 2026 04:13
815f565 to
bd419c4
Compare
Rebased onto a main that moved input construction into workloads/. Git reported no conflict because the adaptive classes are new, but two gates did: the benchmark and the test each defined gen_inputs for the same NCHW tensor, and the test spelled it gen_inputs(*shape), which is not the WorkloadBase contract. workloads/pool.py now owns AdaptivePool2dWorkload; both stages compose it and keep only their own oracle. The scalar-None normalisation torch 2.10 requires was written out in both ref_programs and is now one property on the workload. The bin arithmetic was identical in three kernel bodies, comment included, as was the static bin extent. Both move to pool/common.py. They are evaluated while tracing, so the emitted TIR is unchanged -- 259 tests and 48 benchmarks pass with latencies matching the previous run. Drops the Adaptive 2D pool protocol section from ops-design.md, leaving that file untouched by this branch. The base class docstring already carried the flow, the playbook is otherwise generic with no per-family section, and both the rule it stated and the pointer to where family-base process lives were already there: Family-Base Refactoring says a family base must not normalize genuine per-op differences, and Out of Scope already routes family protocol variables to the appendix.
lcy-seso
force-pushed
the
feat/pool/adaptive-pool2d
branch
from
August 1, 2026 04:21
bd419c4 to
20465ba
Compare
… int32 The three kernel classes were the same 61 lines three times over -- dtype gate, config space, argument marshalling -- differing only in which builder and which registered op they call. They now subclass one base and bind those two callables. The three torch.library registrations stay separate: they need distinct names and are what torch.compile resolves against. With one copy left, the indices kernel's cost shows. It carried the flat index as int64 through the reduction, so every recorded maximum dragged 64-bit div/mod into the update path. The plane is h_in * w_in, known while tracing, so the internal state is int32 whenever that fits and only the stored index is int64, which is what PyTorch returns. indices, tileops latency in ms before after global 1x1 0.0138 0.0137 SPP 6x6 0.0204 0.0158 non-divisible 7x7 0.0174 0.0164 Matching the prediction: the global case constant-folds the index expression and has nothing to save, the other two do. 259 tests pass, so the indices still agree with PyTorch element for element.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1799
Adds
AdaptiveAvgPool2dFwdOp,AdaptiveMaxPool2dFwdOpand thereturn_indices=TruevariantAdaptiveMaxPool2dIndicesFwdOpin TileLang, matching PyTorch for CHW/NCHW, integer or two-elementoutput_size, per-dimensionNone, non-divisible bins, expanded outputs, max ties andint64indices. The three manifest entries move toimplemented.Rebased onto the main that moved input construction into
workloads/, so the new code follows that shape:workloads/pool.pyownsAdaptivePool2dWorkloadand both stages compose it; the scalar-Nonenormalisation is one property there rather than one perref_program; the bin arithmetic and static bin extent, previously identical in three kernel bodies, live inpool/common.py.Benchmark
H200 @ 1500 MHz, torch 2.10.0+cu129. Latency in ms, global 1x1 / SPP 6x6 / non-divisible 7x7 BF16.
The indices variant is behind PyTorch on the two non-global shapes.
Validation
pytest tests/ops/test_pool.py benchmarks/ops/bench_pool.py307 passed, adaptive latencies as above;validate_manifest.py --strictand per-op; the four structural gates;pre-commit run --all-files;ruff check .. The extracted bin extent was checked against the original expression over all 16384 axis pairs up to 128.