[Unity][Tracking Issue] Phase Ordering

[slyubomirsky]

Per the previous discussion of instituting some explicit phase ordering in Relax to account for the fact that some passes already rely on each other, this is a proposal for adopting some of the measures from the discussion.

• P1. Add a comment in transform.h or a README.md in include/tvm/relax that describes the known pass dependencies.
• P2. Include LambdaLifting as a final step after parsing.
• P3. Implement phase tracking during the build.
  • P3.1. Use a module attribute to indicate the phase. The phases will be "ingestion" (0, between parsing and normalization), "optimization" (1, applying high-level optimizations), "lowering" (2, some lower-level optimizations that rely on legalizing operators, eliminating purity tracking, and device-specific properties), and "build" (3, the final lowering of builtins to PackedFuncs and preparing for VM code generation). It will be important to mark transition points accordingly (LambdaLifting as 0->1, LegalizeOps as 1->2, VMBuiltinLower as 2->3).
  • P3.2. Include checks for the expected phase in each pass. This will largely be boilerplate.
  • P3.3. Modify pass implementations to account for phase invariants. For example, after phase 0, passes do not have to handle inner functions (except for BYOC-specific ones, unless we want to change that).
  • P3.4. Enforce certain phase invariants in the well-formed checker, tied to the module attribute for phase.
    • Inner functions should not be permitted after phase 0 (we can make an exception for those used by BYOC or change the implementation of RunCodegen).
    • Dataflow blocks should not permitted after phase 1.
    • All functions must be marked either impure or use the relax.force_pure attribute after phase 2. Additionally, there should be no remaining uses of call_pure_packed or invoke_pure_closure after phase 2. (These are both already implemented by the RemovePurityTracking pass; this would just enforce the invariant.)
    • There should be no remaining uses of call_tir and call_dps_packed after phase 2.
    • We could also consider enforcing that all arithmetic operators will have been legalized after phase 2 (by checking only for a whitelist of builtins that are lowered in phase 3).

This list is meant as a starting point for further discussion, as there were still some unresolved disagreements in the earlier thread. By dividing the implementation into steps, this could help make the discussion more concrete by focusing on implementation.

In terms of implementation, the biggest question centers on the BYOC passes, which rely on inner functions. This is the only major case where the current implementation of passes differs from what was proposed (namely, eliminating inner functions early on). We should also be certain of which passes are responsible for updating the phase attribute (those proposed above under P3.1 are all non-optional passes, so I would argue they are reasonable candidates).

We may also consider trying to enforce some finer-grained ordering constraints (see the "megapasses" in the discussion thread), though there did not seem to be wide support for this; it may suffice to note these constraints in the documentation without enforcing them directly.

cc @quic-sanirudh

[slyubomirsky]

Note from June 13, 2023, community meeting: Documenting pass dependencies might be best to do on the individual passes' docstrings. This would likely be easiest to maintain in the long term as well.