[Relax] Default purity as unknown, instead of pure

[Lunderberg]

Prior to this commit, the constructors for relax.Function and relax.FuncStructInfo provided a default annotation as a pure function. While this could be overridden, it

This commit updates the IR to store purity as Optional<Bool> instead of bool, and changes the default to NullOpt.

A function's purity can have three possible values:

• Bool(true): The function is known to be pure.
• Bool(false): The function is known to be impure.
• NullOpt: The function's purity is unknown.

In most cases Bool(false) and NullOpt should be treated equivalently, as they both indicate a function that may contain side effects. However, in the future, inference of purity may occur when the purity is NullOpt, based on analysis of a function's body. A function with purity of Bool(false) is known to be impure, and further analysis is unnecessary.

[slyubomirsky]

I think this idea is very promising, though we should be certain of some corners of the design:

1. Is there any invariant we want with unknown purity? Do we expect all unknown purities to be resolved by a pass?
2. Is unknown purity permitted in a DF block? We could defer checking it until after (say) a purity inference pass. However, deferring the check now means that the purity inference pass is another phase ordering dependency (not a big deal, but users need to remember to use it after any pass that adds new functions to the module).
3. (More speculative.) Would we want PrimFunc/external function calls to default to unknown purity instead of impure? There is no way to infer for external functions, however, so I would assume the answer is probably no for that case. Inferring for PrimFuncs is doable.

[Lunderberg]

2. Is unknown purity permitted in a DF block? We could defer checking it until after (say) a purity inference pass.

I think the key point is at relax.transform.ToNonDataflow, because that is the last point where we could check whether a model meets the assumptions that had been initially provided by the use of a dataflow block.

I'd propose the following: Unknown purity is always allowed in non-dataflow blocks, regardless of where they occur in lowering, because non-dataflow blocks have no purity requirements. Impure functions are never allowed in dataflow blocks (current behavior). Unknown purity is allowed in a dataflow block, but must be resolved

1. Is there any invariant we want with unknown purity? Do we expect all unknown purities to be resolved by a pass?

I think the non-local purity inference pass can be run as often as desired. It is not required to resolve every unknown purity, and may output functions with unknown purity, if they contain calls to other functions of unknown purity. Then, during ToNonDataflow, all expressions within a dataflow block must be known to be pure.

This would be similar to the way we handle shape inference for R.add(A, B), where A: R.Tensor([sym_var_1]) and B: R.Tensor([sym_var_2]), producing R.Tensor(ndim=1) as the inferred shape. The operation may be legal to perform, and so no error is produced at that point. So long as the discrepancy is resolved before it is required (during relax.transform.LegalizeOps), there is no error required.

However, deferring the check now means that the purity inference pass is another phase ordering dependency (not a big deal, but users need to remember to use it after any pass that adds new functions to the module).

I think I'd have the local analysis of a relax::Function be inspected on construction, as that covers the vast majority of cases. A purity inference pass would only be required to handle cases where the local analysis results in unknown purity.

3. (More speculative.) Would we want PrimFunc/external function calls to default to unknown purity instead of impure? There is no way to infer for external functions, however, so I would assume the answer is probably no for that case. Inferring for PrimFuncs is doable.

Hmm. I'd been defaulting to unknown, in case there is additional information provided at a later stage. The main thing I was picturing was replacing an R.ExternFunc with a known implementation. (e.g. Replacing an R.ExternFunc("initialize_something") with a no-op function in cases that do not require initialization.) However, that case would be better handled by accepting a initializer: R.Callable as a function argument, would could then be replaced using BindParams. For user-written R.ExternFunc instances, I think you're right that

I think for PrimFuncs it would be better to default to unknown, since as you mentioned they could be inferred at some point in the future.

[Lunderberg]

That was a lot of breakage across different unit tests, and good to sort it out. I think it's a point in favor of the local-inference of function purity that #16687 didn't run into any of that type of failure.

[slyubomirsky]

Using ToNonDataflow makes sense as a final threshold for making sure dataflow blocks are correctly formed, so I like that.