---
# System prepended metadata

title: AD Scheduling
tags: [Enzyme, Paper-Proposals]

---

AD Scheduling
===

###### tags: `Paper-Proposals` `Enzyme`


- Optimize for parallellism
- Spatial locality
- Temporary locality
- Optimize for underlying hardware
- Maximize for outer loop parallelism


Advantages of JAX:
- JAX HLO layer
    - Schedules away the unoptimal code
> Take cues from the HLO scheduler


Random notes:
- Parallel code, Reverse mode, read turns to write, write turns to read etc.
- Good schedulers, just like **polyhedral** schedulers
    - Polly
    - Polygeist -> polyhedral optimization
- Drill up Polly to work with RL optimization
- 


> Collect paper references to help Martin and be more focussed!


1. Some scheduler API needs to be plugged in
2. Integration of a scheduler into Enzyme
3. Want to give high-level information into the schedule, to avoid recovery
    - Integrate with Polly?
        1. In the first round do the opposite of what you would easily want to do
    - Integrate with an existing RL scheduler for a related problem
4. Go down the MLIR route
    - Polygeist
    - Extend Enzyme itself to work on the MLIR level
    - Integrate with any of the other systems, e.g. ETH paper, Google papers, FB papers


Martin:
- Seeks to work with LLVM IR
- Polly-level optimizations
    - Polly -> Enzyme -> Polly
    - Modify the first Polly call to act counter-intuitively to act in a way suited for **reverse-mode AD**
- ProGraML representation of the LLVM IR as graphs to work on top of for the RL
- Set up the Polly-Bench to start off with
    1. CPU-versions (serial)
    2. (GPU-versions)
    3. (OpenMP)
- Ravinia is another benchmark (roginia)
- Pull in the Tensorflow Dataset from the CompilerGym benchmark

> LLVM does not have a good representation of `parallel for`, only do auto-parallelization on stage 2 for now

- Use the manual optimization pass infrastructure