Skip to content

feat(cubecl): wire logical any/all reductions#5051

Open
zhan-wei-919 wants to merge 1 commit into
tracel-ai:mainfrom
zhan-wei-919:feat/cubecl-any-all
Open

feat(cubecl): wire logical any/all reductions#5051
zhan-wei-919 wants to merge 1 commit into
tracel-ai:mainfrom
zhan-wei-919:feat/cubecl-any-all

Conversation

@zhan-wei-919
Copy link
Copy Markdown
Contributor

@zhan-wei-919 zhan-wei-919 commented Jun 6, 2026

Checklist

  • Confirmed that cargo run-checks command has been executed.
  • Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Depends on tracel-ai/cubek#302
Related with #5046

Changes

This PR wires Burn's CubeCL backend to cubek's logical any / all reductions.

Changes include:

  • update cubek to the cubek PR commit from feat(reduce): add logical any/all reductions cubek#302
  • update cubecl, cubecl-common, and cubecl-zspace to the matching cubek dependency revision
  • implement CubeCL backend overrides for:
    • bool_any / bool_all
    • float_any / float_all
    • int_any / int_all
  • route logical reductions through cubek ReduceOperationConfig::Any / All
  • add large-dimension float tests for any_dim / all_dim
  • adapt GlobalMemoryStrategy initialization to the newer cubek interpolate API

This makes Burn's contains_nan() path use a real logical reduction through is_nan().any() on CubeCL instead of relying on sum based behavior.

Testing

cargo check -p burn-cubecl
cargo test -p burn-backend-tests --features flex test_any_dim_large -- --nocapture
cargo test -p burn-backend-tests --features flex test_all_dim_large -- --nocapture
cargo test -p burn-backend-tests --features flex test_any -- --nocapture
cargo test -p burn-backend-tests --features flex test_all -- --nocapture
cargo test -p burn-backend-tests --features flex contains_nan -- --nocapture
cargo fmt --all -- --check
git diff --check

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