Description
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
tracking: