Skip to content

Conversation

@copybara-service
Copy link

@copybara-service copybara-service bot commented Sep 30, 2025

[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.

@copybara-service copybara-service bot changed the title [mpmd] Initial custom pipeline API [mpmd] Add custom pipeline API for user-defined scheduling and merging Oct 1, 2025
@copybara-service copybara-service bot force-pushed the test_808549034 branch 2 times, most recently from d9214ad to 3eb4859 Compare October 2, 2025 10:53
@copybara-service copybara-service bot force-pushed the test_808549034 branch 6 times, most recently from e0e5eac to 0e7d8e5 Compare November 6, 2025 10:27
@copybara-service copybara-service bot changed the title [mpmd] Add custom pipeline API for user-defined scheduling and merging [mpmd] Integrate PipelineSchedule API into jit lowering flow Nov 14, 2025
@copybara-service copybara-service bot force-pushed the test_808549034 branch 4 times, most recently from 931260b to 4f53d94 Compare November 17, 2025 17:00
@copybara-service copybara-service bot force-pushed the test_808549034 branch 3 times, most recently from 7bca49a to 8e81972 Compare December 18, 2025 18:56
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: 808549034
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant