[TVM Unity] Open Dev Meeting March 28th

psrivas2 · March 22, 2023, 8:43pm

Hi all,

Our next meeting will take place at 9:00 AM PDT on March 28th (12:00 AM China, 9:30 PM India, 4:00 PM GMT).

Please add topics to the agenda that you would like to discuss

You can find previous meeting notes and recordings here: Apache TVM Relax open development meeting - Google Docs , and join the meeting at https://octoml-ai.zoom.us/j/84450540157?pwd=L2h2SStjWTc1bGU0YS9UZk9KVnptUT09 .

Looking forward to seeing you!

psrivas2 · March 27, 2023, 1:14pm

Current Items on Agenda:

Announcements

Vertical focussed Dev Meetings

Purity Tracking Design Concerns - Steven Lyubomirsky

github.com/apache/tvm

[Unity][IR] Purity Tracking

unity ← slyubomirsky:purity-tracking

opened 03:38AM - 24 Mar 23 UTC

slyubomirsky

+1578 -337

In this PR, I am beginning to implement the tracking of function purity as part …of the `StructInfo` system. This will allow the compiler to enforce that no impure function (one that can possibly have visible side effects) can be called in a `DataflowBlock`. Tracking this requires noting which operators are pure or impure (I am presently doing this with an operator attribute called `FPurity`, a simple boolean), though dealing with calls to other Relax functions means that this information must also be tracked via the `StructInfo` system. Additionally, it is difficult to infer the purity of a function in the general case (when there are calls to other Relax functions), so this change does require users to annotate impure functions using a new field on functions, `is_pure` (in TVMScript, this can be done using `R.is_pure()` or `R.is_impure()`). Since most Relax functions are likely to be pure and purity is the default assumption, this will hopefully not be a large imposition on users. We can consider inferring purity in the easier cases, since those are likely to be common. Note that `PackedFunc`s are conservatively treated as impure. However, in situations where they are needed inside a dataflow block, a call to a `PackedFunc` that is, in reality, pure can be done via the new operator `call_pure_packed` or the existing operator `call_dps_packed` (it is assumed that any `PackedFunc` used with it will be pure). (Similarly, a new operator `invoke_pure_closure` is introduced as a counterpart to `invoke_closure` for dealing with closure objects, though this really should only come up with the `LambdaLifting` pass.) As an "escape hatch" to the purity system, one can use the attribute `relax.force_pure`, which indicates to the compiler to treat the entire function as pure even if it contains an impure call. Additionally, even though `PackedFunc`s are normally treated as impure, a user can use `call_pure_packed` or `call_dps_packed` to call `PackedFunc`s in dataflow blocks when appropriate. These can be used to deal with the following situations: 1. A function does side effects but only on a value that will not be exposed anywhere else or on a new value that will be returned. Even though the individual actions are "impure," the overall function fulfills the definition of being pure. `relax.force_pure` would be useful here. 2. A `PackedFunc` is, in reality, pure. `call_pure_packed` or `call_dps_packed` are useful in this situation. Changes include: * Enforcing that impure functions are not used in `DataflowBlock`s in the well-formed check. * Enforcing that functions that are not labeled impure do not contain impure calls (unless `relax.force_pure` is set). * Implementing the `call_pure` operator Still to be done: - [x] Label purity for all operators - [x] Address certain tricky bugs (Also address the design concerns below) ---- The process of implementing purity tracking has revealed a few issues that may require some further design discussion. ## Labeling purity or impurity Using function attributes to label purity/impurity seems very messy. It might be worth making this part of the `Function` node's AST to avoid having to wrangle attributes in many places in the codebase. ## ~~The treatment of `call_pure`~~ ~~`call_pure` presents a dilemma because many passes look for calls to certain operators, but with `call_pure`, these would be "wrapped" like so: `Call(Op("relay.call_pure"), [inner_op, arg1, ..., argn], attrs, sinfo_args)`. In the posted PR, many passes needed to be revised to look for instances of `call_pure` so that they could be treated the same as calls to ordinary operators. It might cause less disruption to passes to turn `call_pure` into a specialized AST node that literally wraps a call node. This way visitors or mutators could have an easy default case and the lower-level code generation passes could simply ignore the `call_pure` node and deal with the wrapped call. As painful as introducing a new AST node would be, it would likely be more maintainable than having every pass have to have a special case for the `call_pure` operator.~~ The more restricted `call_pure_packed`/`invoke_pure_closure` are less likely to come up in passes and so are less likely to need special handling. ## Staging Related to the above issue, having purity tracking in the well-formedness analysis means that even low-level passes like `VMLowerBuiltin`, which replaces some Relax operators with `PackedFunc`s (treated as impure by default, as they are in principle dangerous black boxes), have to insert calls to `call_pure` to avoid triggering purity errors. Perhaps it might be useful to ignore purity by the time we reach lower stages of compilation, much as the `ToNonDataflow` pass is used in the VM build process to eliminate `DataflowBlock`s. **Update**: Based on the TVM Unity Community Meeting on Mar. 28, 2023, I've provisionally adopted a variant of this approach thanks to a suggestion by @tqchen. Namely, there is a pass called `RemovePurityChecking` that simply adds the `ForcePure` attribute to all pure functions and unwraps all invocations of `call_pure_packed`, `call_pure_dps_packed`, and `invoke_pure_closure`. This allows lower-level passes to remain unchanged and not to have to consider purity at all. It is inserted after `ToNonDataflow` in the `build()` function in `vm_build.py`. **Update 2**: Thanks to suggestions by @tqchen, I've made `is_pure` and `force_pure` into fields on functions rather than attributes and added syntactic sugar in TVMScript to make these easy to set. **Update 3**: Per @tqchen's suggestion `force_pure` will be kept as an attribute instead. This is because it acts more as a direction to the compiler rather than an inherent property of a function.

Having a call_pure op vs making it an AST node
Would it make sense to stop checking purity past a certain point in compilation?

Looking forward to seeing everyone tomorrow!

psrivas2 · March 28, 2023, 4:49pm

Recording: https://octoml-ai.zoom.us/rec/share/jo6MWLrF1ng6Z1YTPPCOlut03nQ1fePX2qnH6nePBVp0JlkpyMdsOlg6XIj95CKG.Ha6fMdazJWRpsdqj

Passcode: T=U?n2Qn