Skip to content

Split std::batching out of std::autodiff to enable AoS and SoA vectorization. #145900

@ZuseZ4

Description

@ZuseZ4

The std::autodiff feature allows batching of calls. So instead of calling the function f(x: f32) -> f32 four times with four different arguments, one can call batched_f(x1: f32, x2: f32, x3: f32, x4: f32) -> [f32;4].

In the Python world, this feature is known from JAX as .vmap().
std::batching should allow Array-of-struct and Struct-of-array vectorization. It is a common autodiff feature since e.g. ML frameworks allow batching to improve training performance, but we have enough interested users who are not interested in autodiff/derivatives/gradients, so it would be good to also offer it as standalone feature.

The first step is to add a std::batching intrinsic, and lower it to an __enzyme_batch call on llvm-ir level.
Luckily, this has been done for autodiff as part of the latest GSOC project, so you can copy most of it: #142640. In this first MVP part, you just need to focus about getting the infrastructure up, even handling a single testcase is enough. The parsing logic will be a strict subset of the parsing logic for std::autodiff, so it should be straightforward.

In a second PR we should then talk with @ Jubilee and some other people from the std::simd project to improve the actual types which we generate. For example, it is likely that we'd want to generate Rust's SIMD types in most cases instead of just generating arrays. Enzyme has a lot of configurations, and we should talk about how to map them nicely to Rust types. When implementing std::autodiff I exposed these in one way, but we likely want to re-evaluate my choices. Any improvements can later be fed back into the std::autodiff macro.

Please ask any questions you might have here and ping @Sa4dUs for questions about introducing an intrinsic, and me about the Enzyme/autodiff/batching logic or lowering to LLVM.

Metadata

Metadata

Assignees

No one assigned

    Labels

    A-frontendArea: Compiler frontend (errors, parsing and HIR)C-enhancementCategory: An issue proposing an enhancement or a PR with one.F-autodiff`#![feature(autodiff)]`F-batching`#![feature(batching)]`

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions