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 PackedFuncs 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 PackedFuncs are normally treated as impure, a user can use call_pure_packed or call_dps_packed to call PackedFuncs 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 DataflowBlocks 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:

• Label purity for all operators
• 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 PackedFuncs (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 DataflowBlocks.

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.