-mpmd-pipeline-scheduler

Reorders the fragments to obtain a given pipeline schedule.

Reorders fragments according to a pipeline schedule. The scheduling algorithm relies on a happens before function that takes two fragments f1 and f2 as arguments and checks if f1 must be scheduled before f2. This function requires that:

• f1 and f2 are assigned to the same mesh,
• f1 and f2 are both scheduling units (i.e., user defined fragments with call_counters defined), and
• f1 does not depend on f2 and f2 does not depend on f1. And it should be enough to express many state-of-the-art pipeline schedules.

For each fragment f1 that must be scheduled before f2, the scheduler pass creates a control-dependency from f1 to f2. Then, it applies a topological sort on the module to guarantee that all dependencies are respected (and the program is in a valid SSA form). Finally, the pass removes from the graph any control-dependency introduced.

Options

-must-happen-before          : A comparator that determines whether a fragment must be scheduled before another. Can be parsed from a built-in `PipelineSchedule` as follows: `builtin:<schedule-as-string>`.
-remove-control-dependencies : Whether to remove control dependencies at the end of the pass.

-mpmd-remat-fragment

Rematerializes fragments.

Finds pairs of fragments (forward+backward) that need to be rematerialized and clones every forward fragment before its backward user, replace all backward uses of values produced by the forward fragment with the cloned counterparts. This can be used for activation rematerialization in pipeline parallelism.

When merge_remat_fragments is true, then we merge the remat fragments into their consumer fragments.

Options

-merge-remat-fragments : Whether to merge the remat fragments into their consumer fragments.

-mpmd-rule-based-schedule

Reorders fragments based on user-defined rules.

Reorders fragments by adding control dependencies based on a list of rules. Each rule specifies a sequence of fragments in the desired order of execution, and control dependencies are added to enforce that sequence. For example, for a rule with fragments [A, B, C], control dependencies are added for A->B and B->C. The pass fails if the rules create a cyclic dependency. If a specified fragment in a pair of fragments within a rule is not found, that pair is not scheduled.

Options

-rules                       : A list of fragment schedule rules. Each rule is a list of `FragmentInfo`s that specifies the order in which the fragments should be executed.
-remove-control-dependencies : Whether to remove control dependencies at the end of the pass.