[Unity][Transform] Handle relax.Var as call_tir args when lowering

[Lunderberg]

Prior to this commit, several transforms assumed that the arguments passed to a call_tir builtin were provided as in-line relax::Tuple objects. Because it would be equally valid for the arguments to instead be a relax::Var instance that had previously been bound to a relax::Tuple object, or had been passed as an input parameter with relax::TupleStructInfo, this assumption shouldn't be made. This PR updates the CallTIRRewrite, FoldConstant, FuseOps, and RewriteDataflowReshape passes to handle variables providing the arguments.

[Lunderberg]

Rebased onto unity head to ensure the CI tests haven't broken in the meantime.

[slyubomirsky]

I think you may need to update the StructInfo inference for call_tir_inplace like in #15971, since (without modification) that assumes the argument is a tuple literal. The test cases here don't try that case, hence why that doesn't result in an error (I'd recommend adding a test case).

[slyubomirsky]

Not that I particularly mind it, but is it preferable to use this construction with a lambda as opposed to just assigning args in different branches?

Array<Expr> args;
if (case1) {
    args = ...
} else if (case2) {
    args = ...
}
// etc.

Does it avoid an allocation or something?

[Lunderberg]

I tend to use this construction for a couple of reasons.

• Avoid ever having a partially-initialized variable.
• Limit the scope of temporary variables that should only be used in the initialization.
• Simplify nested if/else cases with early returns.

Effectively, the immediately-invoked lambda expression acts as a block scope with a return type, similar to Rust's braces (e.g. The value of {let i = 5; i+1} is 6) or relax's SeqExpr.

[slyubomirsky]

You should do this for the other call_tir variants in this file too.

[Lunderberg]

Good call, and updated.

[tqchen]

I think we should instead enforce call_tir to be a more restricted form

[tqchen]

Thanks for the PR, I know this is indeed a generalization and there are some tradeoffs to be considered here. Specifically, we should consider the following alternative:

• C0: We enforce call_tir intrinsic to only take explicit Tuple and enforce in well-form check
  • This is what we can afford to do for special intrinsics, there are other similar things like call_packed etc which follows the same spirit

This indeed limit the flexibility in terms of what is valid, but would greatly simplify the logic that leverages CallTIR. And such simplicity helps us in a lot of cases since passes are simpler and more passes that depends on call_tir pattern matching would enjoy that simplicity.

From expressiveness pov, there is nothing to be lost and we only need passes that generate call_tir to be able to explicitly unpack tuples.

[slyubomirsky]

See also this issue. Allowing call_tir to be more general would allow us to avoid special handling for call_tir (and its variants) in that case. It's a question of which we'd prefer.

FWIW, I think our utilities for looking up bindings make it simple enough to deal with cases like not having a tuple literal in call_tir.

[Lunderberg]

This does increase the complexity of passes that interact with call_tir, but it also decreases the complexity of all other passes. If we require an in-line tuple for call_tir, then every other pass must be aware of this restriction and preserve it. Any ExprMutator class that implements Expr VisitExpr_(const TupleNode*) override must also check the context of the Tuple in order to know what mutations are allowed, and similar checks would be required in VisitBinding. This is complexity that is inherent to interactions with a call_tir node, and a IR restriction not exposed through the C++ type system moves that complexity to passes that shouldn't be aware of call_tir at all.

I think there are ways we could ensure that all call_tir has immediate access to its arguments, working within the C++ type system instead of running a well-formed check

• Provide an explicit IR node for CallTIR and variants. Instead of R.call_tir(tir_gvar, [arg1, ..., argN]) expanding into R.Call(tvm.ir.Op("relax.call_tir"), [tir_gvar, R.tuple(arg1, ..., argN)]), it would expand into R.CallTIR(tir_gvar, [arg1, ..., argN]).

• Move the GlobalVar target into a member variable of the call_tir operator. Instead of R.call_tir(tir_gvar, [arg1, ..., argN]) expanding into R.Call(tvm.ir.Op("relax.call_tir"), [tir_gvar, R.tuple(arg1, ..., argN)]), it would expand into R.Call(R.TIRFunc(tir_gvar), [arg1, ..., argN]). This would be analogous to how R.ExternFunc is currently handled.longer require a tuple, so it couldn't be indirect.

With either of these long-term options, the arguments would be immediately accessible by passes that interact with call_tir, but this would be enforced at the type system, so other passes wouldn't be required to explicitly check for it at runtime.

[Lunderberg]

@slyubomirsky I pulled in the infer struct info improvements and the unit test updates from PR#15971, so this should now include all fixes from both sets of changes.

[slyubomirsky]

My concerns were addressed, but we should have agreement on the design question @tqchen posed before we merge (if we will merge).

[slyubomirsky]

Making a new IR node for call_tir would entail a lot of revisions, but it might be worth it since it is very important to the language. If that's something we would consider, we should give thought to other changes to Call node potentially (several have come up before in community meetings), since making changes to the AST would require careful thought.

Since we have the StructInfo system as a type system, I don't see a reason why call_tir needs to be a special case--a tuple is a tuple and that's the beauty of the type system.

[tqchen]

The principle here is we make common cases(and their optimizations easy), while placing burdens on less common cases.

As for the pass writing patterns, most of our current relax passes starts with mutating leaf-node e.g. Call where we can explicit unpack the patterns. They are structured as follows:

• Override Call, pattern match on the op, which allows us to enable structural intrinsics.
• Directly unpack the args since most op likely will entail the structure there, return a new value.

Mutating TupleExpr directly have the issue of needing to consider possible recursions, given our normal form allows nested tuple, exactly for the ease of representing structural information in some special intrinsics. It can become really hard to reason for most cases due to the existence of recursion. Likely there is really a limited set of passes that involves this directly mutation (as a matter of fact, I only recall perhaps structural info deduction, which is consistent), especially when they are tied to key optimization needs, and they can be restructured to the above pattern.

Taking that rationale and given a lot of instrinsics also follow that pattern(e.g. CallPacked and other possible low level variant (e.g. an intrinsic that explicitly marks read and writes). Supporting intrinsics to allow special structures in the signature helps to simplify passes. It is already part of our design, solving it for CallTIR alone likely won't resolve issues on other calls, and future instrincis that we might need (like explicit read write grouping for multi-stream analysis).

Given the CallTIR and other structured intrinsics are central part of a lot of our analysis. I think the tradeoff is worth spending. Such restriction can be checked through well-formness check so they won't propagate into runtime.

[slyubomirsky]

If we intend to have special cases like call_tir where one argument must be a tuple literal (i.e., not following the normal rule of the type system that any member of the type should be permitted to stand in for any other), we should probably have some centralized place or an API where these rules are defined so that passes know to handle them. I'm not sure what form this should take, but such special cases make it trickier to implement passes like common subexpression elimination: the passes have to know about which operators are special and have special handling for them (as @Lunderberg noted).

I would argue that it doesn't bring that many advantages to require the argument to call_tir be a tuple literal, since the only other possibility (in normal form) is for the argument to a variable, and it is easy to look up the definitions of variables (if they're from the same function) or to use TupleGetItem if necessary (like the GetTupleIndex method in this PR).

[Lunderberg]

My concerns were addressed, but we should have agreement on the design question @tqchen posed before we merge (if we will merge).

Oh, absolutely. The comment was to ensure that we had a clean implementation for the design comparison.

I'm not sure what form this should take, but such special cases make it trickier to implement passes like common subexpression elimination: the passes have to know about which operators are special and have special handling for them

Thank you for this phrasing, as I was trying to put this sentiment into words. The existence of special requirements that must be preserved makes generic implementations of other passes be much more difficult.

[tqchen]

What we are discussing is closely related to our original consideration of what forms a normal form and the general needs. Let me elaborate a bit more

The purpose of normal form is to restrict the set of possibilities in representing the same program thus helping reduce the assumption that one might take when writing pass. Of course it imposes extra demand for passes to ensure things go back to the normal form(the wellform-ness check). Currently, the relax normal form requires that we normalize all non-tuple nesting, but expand all tuple constructions.

That means we encourage forms like

def func(x, y):
    lv0 = call_tir(mm, (x, y))

But do not allow the following in normal form

def func(x, y):
    lv0 = (x, y)
    lv1 = call_tir(mm, lv0)

The reason we do that is to observe a general category of need:

• N0: Structural intrinsic argument

When constructing intrinsics, we would like intrinsics to present certain structures in them, for example call_tir, call_packed. If we want to analyze read write relations low-level, maybe we would like to have intrinsics of
lv0 = call_func_with_read_write(func, read_tuple, write_tuple)

Pattern match and rewriting of these structural intrinsic is a first-class need in our case because of following reasons:

• R0: they are common needs in machine learning passes, i can easily write a customized pass that matches and replaces the call_tir, such needs are also unbounded.
• R1: We would like to make that customization happen at a third party, aka the developers who do not necessarily contribute to upstream can do it as simply as possible without worrying about unpacking tuples.

Such a normal form might bring a bit of restrictions, but won’t be too much to other possible passes that are “generic”. This also depends on the past writing patterns we do. Let us consider the following patterns

• PT0: Only overload visit call and unpack that during visits
• PT1: Possibly overload visit tuple

I would claim that PT0 is actually more desirable and handles almost all the pass needs, and it won’t cause any issues in the structural intrinsics, since all arguments are transformed and pattern matched together.

I think most of the possible ideas of difficulty mentioned are related to passes that follow PT1. Under the current normal form however, we naturally need to be very careful about PT1, because we do not want to lift common values in tuple construction and instead encourage tuples to be always unpacked. A CSE pass that tries to detect sub-tuples and lift common values via PT1 actually would violate the normal form requirement. A correct CSE pass under the current normal form should not overload PT1

A rule of thumb under the current normal form is that we almost always won’t do PT1, unless one passes that populate tuple annotation(aka deduction). If there is another pass that does so, we should visit the assumption and check carefully if the result can violate the normal form assumption.

If there is really a need to do tuple replacement in a pass, and there is a concern that the result might violate the structural intrincis requirment. One simple approach is to have a common pass(like convertSSA in tir) that goes over the Call, always unpack tuple struct info to their bottom. This pass can be inserted after whatever pass so the assumption is preserved, and we only need to register a property(require structured tuple) in the Op.

Finally, one should note that the passes might touch PT1 are usually developed in repo, as a result, the extra consideration of renormalization won’t be hurt, as our experience is sufficient to tackle these considerations. Burdens of R0/R1 are unbounded, a lot of our designs(e.g. Dataflow block) goes into simplifying them for developers who may not have a deep background.

Of course all of the above roots back to our rationale of the normal form. One can argue we should take other normal forms(e.g. Tuple value should always be bound to a variable), in which case the consideration would be different and there are other tradeoffs. We considered these tradeoffs before arriving at the current one.

[Lunderberg]

I think the major issue is the inconsistency of normalization. From the code itself, there is no clear expectation of this being the normal form for Relax, and the majority of the handling in the codebase

Behavior suggesting that the de facto normal form allows Tuple variables

• Normalization does not suppress tuple variables when passed to BlockBuilder::Emit.
• Normalization does not inline tuple variables when used in a call_tir expression.
• Normalization does not unwrap TupleGetItem calls to a suppressed tuple variable.
• A relax.Call node may return a relax.Tuple, which is bound to a variable, and is then immediately used as the argument tuple of call_tir.
• Tuples are allowed to be on the RHS of an expression, bound to a variable.
• User-written TVMScript can define my_var = (x,y).
• Pattern-matching, which acts on a is allowed to replace a relax.Call with a relax.Tuple. (e.g. replacing the output of R.split with a tuple of R.strided_split calls.) Pattern-matching operates on a relax.Expr level.

Behavior suggesting that the de facto normal form requires in-line Tuple variables

• Normalization keeps relax.Tuple arguments in-place, unlike relax.Call arguments which are extracted out into separate bindings.

In order for in-line tuples be the normalization used in practice, and not just in principle, the Normalizer would need to have a number of changes made to it. It would need to check for relax::Tuple in BlockBuilder::Emit, and suppress the binding made for it. Whenever the relax::Var representing the suppressed tuple is used, the Normalizer would need to replace it with the tuple itself, including when used in nested tuple expressions. This seems much more fragile and easy to break than having the normalization of tuples follow the same rules as the normalization of other relax expressions.

[slyubomirsky]

Great work in making the full comparison list, Eric. Are we presently using inline tuples to simplify things significantly? I really don't think it's that big a deal to check the type of a variable and determine if you need to use TupleGetItem, etc.

[Lunderberg]

Are we presently using inline tuples to simplify things significantly?

I haven't found any locations that would require anything more than a binding lookup. Of the 19 instances of Downcast<Tuple> that I found using grep, none of them would be particularly difficult to handle.

• 8 cases which are removed by this PR
• 8 cases that use Downcast<Tuple> inside as part of an IRModule transform, and which could use the UnwrapBindings utility in this PR.
• 2 cases that occur in VM codegen, to handle "relax.call_builtin_with_ctx".
• 1 case that occurs in a helper function, which is then used in IRModule transforms.

The most complicated case involving relax::Tuple that I've found is in the dataflow pattern matching, but that already handles the unwrapping of tuple nodes here with the TryGetVarOfVal line.

[tqchen]

The simpler approach, is actually to just enable well-form check to detect the related issues and require tuple to be inlined.

Normalizer serves as a way to create new bound variables when composite expression occur, which is the common case, but do not necessarily have to handle all the inlining cases. Where pass writer can simply do the items if needed once detected by wellform check.

The purpose of inline tuple is to remove the need of bound lookup, and make the structural information in the arguments like call_tir to be much more clear(as te structural belongs to the function itself.

Syntactically, that means explicit values like below.

def main(x, y):
    lv0 = (x, y)
    lv1 = call_tir(mm, lv0)

The goal is to be aligned with the rationale to support N0: Structural intrinsic argument

Again the issue is that the use cases of strutural matching of intrinsics can be unbounded (as many passes are pattern rewriting of structural intrinsics), we would like to simplify these cases both in terms of syntax, as well as the pattern matching code. Aka we would like to eliminate the extra binding lookup for both syntax reason, as well as reduction of metral overhead for having to think about optional binding lookup or polymorphism.

This complexity would even get worse if we start to look into nested structure, which we might need in future, it also grows with number of structural instrinsics we add, which goes beyond call_tir itself.

In the mean time, the complexity of keeping things normalized(e.g. have util postproc that automatically insert tuple getitem for cases that needs to be unpacked for passes that might generate tuple argument passing), is managable with the clear wellform-ness check. And also likely simplifies most of our needs here.

[Lunderberg]

The simpler approach, is actually to just enable well-form check to detect the related issues and require tuple to be inlined.

Oh, certainly agreed that the well-formed check is the simpler one to introduce in the short-term. I think it adds complexity in the long-term, though.

• Not evident to readers of the code. A well-formed check can only detect ill-formed IR once it has been generated. Encoding it in the C++ types used in the IR informs a reader what structures are allowed.

• Not evident to pass writers. A well-formed check can only detect ill-formed IR after a pass runs. Encoding it in the C++ types of the IR provides feedback from the C++ compiler, because ill-formed structures would cause a compilation error.

• Not comprehensive. A well-formed check can only detect whether the currently tested input to a pass produces ill-formed IR. Encoding a requirement in the C++ types of the IR ensures that a pass cannot produce ill-formed IR.

Where pass writer can simply do the items if needed once detected by wellform check.

This is the case I want to avoid, because this is too late of a warning. It occurs after a pass has been designed and written, often only when more comprehensive inputs are used in CI. Adding a new design constraint that late in development often requires re-design of the implementation in order to follow the new constraint, resulting in significant repetition of effort.

Aka we would like to eliminate the extra binding lookup for both syntax reason, as well as reduction of metral overhead for having to think about optional binding lookup or polymorphism.

I agree with the reasoning of reducing the mental overhead, but think that comes from having as many constraints exposed to the C++ type system as possible. (See earlier comment for ways this could be done.) By making illegal states be impossible to represent, the mental overhead of avoiding illegal states is removed. If illegal states can be represented, we increase the mental overhead of avoiding those illegal states.

[tqchen]

as many constraints exposed to the C++ type system as possible

This is indeed also a tradeoff here in terms of IR design. We cannot encode all constraints, for example the ANF itself is not encoded directly in the c++ data structure, but still enforced through well-form checkers.

For this particular case, because we can possibly allow unbounded usecases to support N0: Structural intrinsic argument. The possible usecases go beyond CallTIR, that also include call packed, possible nested structures etc. So enabling it in the normal form is a good tradeoff

It occurs after a pass has been designed and written

A pass that follows the common pattern(rewrites Call and not rewrites Tuple) will not face such issues. A pass that needs to do tuple replacement already have other complications to consider. And if the pass writer is not willing to do so, having a followup util function(like ConvertSSA) will simply unpack the tuple values and recovert things back to the normal form, without need to redesign the original pass.

Indeed, this extra burden is when we that develops pass that involves tuple remapping, such pass in nature would require extra set of considerations. And the extra overhead of insertin a re-normalization is not high. They are less common than the passes that involves pattern match and rewrite structural intrinsic argument. That is also why the design prioritizes these cases.

[Lunderberg]

@slyubomirsky @tqchen Can you take a look at PRs #16067 and #16068? The former introduces new functionality through a FNormalize attribute, and the latter applies that functionality to relax.op.call_tir.

1. Downstream passes can make assumptions about the specific form of the AST. For call_tir, the current use of Downcast<Tuple> be guaranteed to have an in-line tuple.
2. If the operator can be automatically normalized, then upstream passes do not need to be aware of the operator-specific requirements. The normalization is applied as part of every ExprMutator usage, the same as any other normalization step.
3. If automatic normalization is impossible, an error is raised during the pass that introduced the ill-formed IR, rather than at the downstream pass that tried to use it.

Properties (2) and (3) fulfill my goals of avoiding an increased mental burden when writing an upstream pass, and of making all IR requirements be explicit within the code.

[Lunderberg]

Closing this PR, as it is superseded by #16067 and #16068.