Skip to content

implement the std::batching feature #140907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZuseZ4 opened this issue May 10, 2025 · 1 comment
Open

implement the std::batching feature #140907

ZuseZ4 opened this issue May 10, 2025 · 1 comment
Assignees
Labels
C-tracking-issue Category: An issue tracking the progress of sth. like the implementation of an RFC F-batching `#![feature(batching)]` T-compiler Relevant to the compiler team, which will review and decide on the PR/issue.

Comments

@ZuseZ4
Copy link
Member

ZuseZ4 commented May 10, 2025

The batching feature is already exposed as part of the autodiff macro, but if someone isn't interested in computing derivatives, then just using the autodiff macro for batching is a bit verbose.

For users it would therefore be good to have a standalone batching macro.
It would be brought into scope with

use std::batching::batching;

similar to one that currently introduces autodiff. My autodiff frontend PR shows how such a macro can be introduced: https://github.com/rust-lang/rust/pull/129458/files#diff-ea30f57e78f3b861b784ae315b3ac31358b7cea2aa14a668084ad412de12e586 Especially the code under library is relevant.

The std::autodiff macro is getting glowered into a rustc_autodiff internal attribute. There are tests in tests/pretty/autodiff showing how this looks like (with the relevant code again being in the frontend PR).
For the std::batching macro, we should do something similar and introduce a rustc_batching attribute.
It is totally fine to just do one step at a time, and merge such a frontend PR first.

The macro itself accepts a positive integer as the first argument, that's the batch width (e.g. 4, 8, 32).
Then it accepts one activity per function argument. Unlike autodiff, there are only two allowed activities.
We could follow Enzyme and call them batch and leaf for now, but I'm open to better names. I think that's enough for the first PR.

In the second step, we would then lower the batching attribute to llvm-ir, like we already do it for the autodiff attribute. We should in some way make sure that we don't duplicate much code here, since most things will behave identical. As an outlook, here are some batching tests in Enzyme core: https://github.com/EnzymeAD/Enzyme/tree/main/enzyme/test/Enzyme/BatchMode

cc @TheDeveloper101

tracking:

@ZuseZ4 ZuseZ4 added the F-batching `#![feature(batching)]` label May 10, 2025
@rustbot rustbot added the needs-triage This issue may need triage. Remove it if it has been sufficiently triaged. label May 10, 2025
@ZuseZ4 ZuseZ4 removed the needs-triage This issue may need triage. Remove it if it has been sufficiently triaged. label May 10, 2025
@TheDeveloper101
Copy link

@rustbot claim

@fmease fmease added T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. C-tracking-issue Category: An issue tracking the progress of sth. like the implementation of an RFC labels May 11, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
C-tracking-issue Category: An issue tracking the progress of sth. like the implementation of an RFC F-batching `#![feature(batching)]` T-compiler Relevant to the compiler team, which will review and decide on the PR/issue.
Projects
None yet
Development

No branches or pull requests

4 participants