implement the std::batching feature #140907
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.
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
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 arustc_autodiff
internal attribute. There are tests intests/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 arustc_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:
The text was updated successfully, but these errors were encountered: