Skip to content

feat(pool): add AdaptiveAvgPool3d kernels - #521

Open
jcwal1516 wants to merge 1 commit into
tracel-ai:mainfrom
jcwal1516:branch/adaptive-avg-pool3d-local
Open

feat(pool): add AdaptiveAvgPool3d kernels#521
jcwal1516 wants to merge 1 commit into
tracel-ai:mainfrom
jcwal1516:branch/adaptive-avg-pool3d-local

Conversation

@jcwal1516

@jcwal1516 jcwal1516 commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

Validate your PR with burn.

It is important that you make sure that you don't introduce any bugs in burn.

Instructions

  • Create a new branch or fork of the burn repo
  • Update the main Cargo.toml with this PR hash.
  • Fix any broken tests or compilation errors in burn.
  • Submit a PR in burn with your fixes and link it here.

Burn validation: tracel-ai/burn#5372 at Burn commit f62594eb975f04d2d95bb0c728bc1e47484dc588. It pins this PR's Cubek revision 604d51ac0403ecbe06156c958ac458ca45a550e6.

Related issue: tracel-ai/burn#5288.

Context

This replaces #497, which GitHub automatically closed after #495 merged and its base branch was deleted. The AdaptiveAvgPool3d feature patch is unchanged; the branch is rebased directly on Cubek main at 22b77ade29af26a41dac872afb9f02b5a1bd8010.

Changes

  • Add native channels-last AdaptiveAvgPool3d forward and backward kernels.
  • Gather input-gradient contributions in the backward kernel, avoiding atomic accumulation.
  • Derive pooling geometry, vectorization, and cube launch settings from runtime tensor shapes.
  • Accumulate pooling values in f32 before casting to the output element type.
  • Validate rank, channel count, output size, and gradient shape at the public boundary.
  • Generalize the CPU pooling oracle to three spatial dimensions.
  • Correct the existing 2D backward test shapes so PoolBackwardProblem consistently represents channels-last tensors after generalizing the CPU reference.
  • Add forward and backward coverage for divisible, non-divisible, asymmetric, multichannel, output-one, and larger-output cases.

Testing

Passed against Cubek revision 604d51ac0403ecbe06156c958ac458ca45a550e6:

  • cargo xtask check format
  • cargo xtask check lint
  • cargo xtask check typos
  • cargo xtask check audit
  • cargo clippy --tests --features extended -- -D warnings
  • cargo clippy -p cubek-matmul --no-default-features --features std,tiled,extended,cubecl/cpu --tests -- -D warnings
  • cargo clippy -p cubek-matmul --no-default-features --features std,multi-level,extended,cubecl/cpu --tests -- -D warnings
  • cargo xtask doc build
  • cargo xtask doc tests
  • macOS Metal: cargo test-metal-light -p cubek-pool -- --nocapture — 55 passed, 0 failed, 10 ignored
  • macOS CPU: cargo test-cpu-light -p cubek-pool -- --nocapture — 55 passed, 0 failed, 10 ignored
  • Burn macOS Metal, no fusion — 32 targeted AdaptiveAvgPool3d tests passed
  • Burn macOS Metal with fusion — 32 targeted AdaptiveAvgPool3d tests passed
  • Burn cargo run-checks --backend metal — passed the complete canonical check, including both full Metal backend suites
  • Burn cargo xtask doc build
  • Burn cargo xtask doc tests
  • Burn cargo metadata --format-version 1 --locked --no-deps

Hosted Cubek CI on this exact commit passed code quality, documentation, and both stable and previous Linux jobs.

Hosted Burn CI on the exact linked Burn commit passed every job, including code quality, documentation, macOS Metal, stable and previous Linux jobs, stable Windows jobs, no-std, and coverage upload.

Before this rebase, the byte-identical feature implementation also passed WSL2/Linux CUDA and native Windows CPU/CUDA cubek-pool runs — 55 passed and 10 ignored in each configuration.

Confirmed baseline exception:

  • Local cargo xtask test --ci reaches the unrelated cubek-reduce failure reduce::it::reduce_dim::perpendicular_matrix_large_odd_batch::f32::parallel_vectorization_enabled::full_unit::reduce_dim::test_topk_with_indices_3. The identical isolated failure reproduces on clean Cubek main at 22b77ade29af26a41dac872afb9f02b5a1bd8010. The cubek-pool portion passes 55 tests with 10 ignored, and hosted Linux CI passes on the exact PR commit.

No benchmarks were added or used, and this PR makes no performance claim.

@jcwal1516
jcwal1516 force-pushed the branch/adaptive-avg-pool3d-local branch from 27d9bf4 to 604d51a Compare August 26, 2026 05:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant