[SLM][Bugfix] Output debug functions as impure

[Lunderberg]

Prior to this commit, debug functions were generated with relax.call_pure_packed. This resulted in unexpected behavior, as nn.op.print_ can be optimized away as a pure function.

This commit updates debug functions to be generated as impure functions. This requires removing the with bb.dataflow() blocks in the SLM-to-relax conversions, as impure functions may not be used in a dataflow block. To restore dataflow blocks when legal, the ConvertToDataflow pass is applied.

[Lunderberg]

This commit is currently marked as a draft, as it depends on a related bugfix in #16684.

[Lunderberg]

The prerequisite PR#16684 has landed, and this PR is now ready for review.

[slyubomirsky]

Thank you for ensuring that function purity is correctly used.

Regarding the changes to labeling function purity, inferring purity is a welcome change, but I think there are cases that could prove problematic. One would be recursive or mutually recursive functions, as they were the original reason that annotating purity was mandatory rather than optional and inferred. We have generally wanted to avoid needing multiple passes to properly infer StructInfo in Relax--it may be necessary to detect those cases and warn the user to manually label purity for them. Recursive local functions could also be a problematic case for a similar reason.

[slyubomirsky]

Suggested change

	* \note `ret_struct_info` and `is_pure` are required, since it can not deduced by the body.
	* \note `ret_struct_info` and `is_pure` are required, since they cannot be deduced from the body.

Might as well correct the typo too

[Lunderberg]

Good catch, and corrected.

[slyubomirsky]

The docs for is_pure should mention that it's inferred.

[Lunderberg]

And updated.

[slyubomirsky]

How would this behave in the recursive or mutually recursive case? That was the reason for not inferring purity in the first place. Detecting those cases and warning the user that they need to be explicitly annotated would be a reasonable approach

[Lunderberg]

Currently, the struct info for a GlobalVar is determined only by the forward declaration, and not by its body. The default for this is determined here (here), where a function is pure unless explicitly annotated otherwise.

We have a similar problem with return values, where the inferred return type of a function may not omitted, resulting in incorrect struct inference in the calling scope. I think the long-term solution to both is the same: To represent the lack of information while parsing, and to infer full information as a post-proc.

[Lunderberg]

Good points regarding default values. I think I would describe the current state a bit differently, as currently we do not require function purity to be explicitly annotated. At the C++ level, the Python API, the Relax block builder, and in TVMScript, we provide default arguments stating that a function is pure. This can be overridden in most cases, but in some places such as the relax block builder, there isn't a way to override this default.

So, rather than viewing this PR as changing from explicit purity annotations to implicit purity inference, I'd argue that it is only changing the method of purity inference. On main, the purity inference always returns true, where this PR improves the purity inference by checking the immediate body of the function.

[Lunderberg]

Long-term, I'd like to move some of the purity inference out of the parser and into a later transformation pass. The sequence of PRs would look something like the following:

1. Update FuncStructInfo::purity to have type Optional<Bool> rather than bool. This would allow it to express a function whose purity is unknown, similar to the unknown shape/dtype options allowed in TensorStructInfo.
2. Update the defaults to produce NullOpt if the user hasn't specified the function's purity. (Mostly done in this PR.) When inspecting a function body, calls to a subroutine with NullOpt purity mean that we cannot mark the function as pure. No further local analysis would be done at that point.
3. Introduce a separate pass for handling the recursive/mutually-recursive cases.

[slyubomirsky]

I would be in favor of having an extra pass to infer purity, since it clearly seems to be in reach. However, when this subject was raised in design discussions before, the suggestion was turned down due to a preference of having only one pass (normalization) handle StructInfo, of which purity is a part. We could reopen the discussion, of course.

[Lunderberg]

However, when this subject was raised in design discussions before, the suggestion was turned down due to a preference of having only one pass (normalization) handle StructInfo, of which purity is a part. We could reopen the discussion, of course.

I'd be interested in hearing the discussion on it. I wouldn't mind having all StructInfo handling done within a single pass, but I don't think the API of the existing normalization would allow it. The tvm::relax::Normalizer may receive any partially-constructed fragment of relax IR, not necessarily a full IRModule. The fragments are sufficient for performing local analysis, but full struct inference or purity analysis would require a similar inference of the called function.

[Lunderberg]

I think the key next step would be to avoid producing incorrect information. If the FuncStructInfo defaulted to "unknown purity" rather than "pure function", then that would allow for incremental improvement in purity inference. As it is, we're putting incorrect information into the annotations.

[slyubomirsky]

Oh my mistake, a year ago we actually said that trying to infer it is what we would want. See meeting notes from Feb. 28, 2023 (I can't put in a link to a specific heading). This would be consistent with what we talked about so we should give consideration to the approach of marking unknown purity. I think it's a good idea and certainly worth discussing, though it would mean that we might have to rerun the pass at different points.

[Lunderberg]

I like it, and I've made a sibling PR #16744 that would allow FuncStructInfo to describe a function with unknown purity. There's a couple of common elements between this PR and #16744 in the TVMScript handling, but it should be mostly independently.