rust-lang / rust-lang/rust

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

Open
#145,900 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

A-frontend C-enhancement F-autodiff F-batching
Dominant language
Rust
Stars
119k
Forks
16.1k
PR merge metrics
PR metrics pending

Description

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: https://github.com/rust-lang/rust/pull/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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the linked Rust pull request and the existing std::autodiff intrinsic, parsing, and LLVM lowering implementation. Trace how the autodiff path reaches LLVM IR, then scope the MVP to a std::batching intrinsic that lowers to __enzyme_batch and supports one testcase. Done means the infrastructure builds and that testcase passes; SIMD type design is explicitly deferred.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.