jax
5a3b1cc0 - [mpmd] Integrate PipelineSchedule API into jit lowering flow

Commit
51 days ago
[mpmd] Integrate PipelineSchedule API into jit lowering flow Lowering Flow: When a `PipelineSchedule` is provided in `MpmdConfig`, the `mpmd.jit` lowering process first runs the IMPORT pass only to generate initial fragments. The rule builders defined within the `PipelineSchedule` are then invoked to generate scheduling and merging rules based on these fragments. These rules are used during the subsequent passes, specifically OPTIMIZE. User defined partitioning options are checked to ensure they don't conflict with options required by `PipelineSchedule`. This new API replaces the direct usage of `fragment_merge_rules` and `fragment_schedule_rules` in `MpmdConfig`. PiperOrigin-RevId: 869242735
Parents
Loading