[Unity] Implement FNormalize attribute for operators

[Lunderberg]

Some Relax operators have requirements regarding their AST that are stronger than are checked by the C++ types being used. These are similar to checks that are present in the tvm::relax::WellFormed utility, such as checks forbidding the use of undefined variables, which are also stronger than required by the underlying C++ types. However, because every operator may have unique requirements, it would be unreasonable to expect a writer of a relax::ExprMutator to be aware of and to maintain all such requirements.

This PR introduces an operation operator attribute FNormalize. If defined, this function is used to apply an operator-specific normalization. The implementation of FNormalize has the following design decisions.

• If no change is required, FNormalize should return the input argument unmodified.

• FNormalize is only responsible for normalization of the operator itself. The expression it returns may be unnormalized (e.g. contain nested expressions).

• FNormalize receives the BlockBuilder as an argument, to allow context-dependent normalization.

  For example, an operator whose normalization requires in-line expressions may use BlockBuilder::LookupBinding to perform variable replacement.

• FNormalize is applied after FInferStructInfo. FNormalize may assume that the relax::Call passed to FNormalize has well-defined struct info.

  • Corollary: FInferStructInfo may not assume that its relax::Call argument has been passed through FNormalize.

    This is a reasonable requirement, because (1) shape inference should depend only on the struct info of arguments and not the values themselves, and (2) this only impacts operators that use FNormalize.

• FNormalize should not be used to apply simplifications, and should be limited to cases where the same computation may be expressed in multiple manners.

  For example, replacing a by-variable tuple with an in-line tuple in R.call_tir is a form of normalization, but replacing R.add(arg, R.const(0)) with arg is a form of simplification.

  This separation is to ensure that FNormalize has minimal overhead, as some simplifications may have large computational costs, and FNormalize is applied as part of all ExprMutator usage. A later PR will introduce an attribute FSimplify, along with a dedicated pass to apply simplifications.

• Use of FNormalize is suppressed while parsing TVMScript. TVMScript must be able to generate test cases that trigger specific failure modes, and that may include producing un-normalized relax IR. In addition, TVMScript must be stable when passed through a round-trip from IR to text to IR.

• If an IRModule contains any non-normalized operators, the IRModule is ill-formed. That is, all FNormalize operations on a well-formed module are no-ops.

[tqchen]

Thanks for the proposed change. I like how FNormalize can help reducing overhead of creating certain operators and bring them back to normal form.

I only have one comment on the wellform check side. In this case, it is useful to have an intentionally duplicated check that is different from FNormalize , e.g. have a TEnforceExplicitTupleInArgs attribute that enforces the tuple argument being unpacked, and check this condition. This provides extra layer of protection, makes the intention clear and is also more efficient

[Lunderberg]

Thank you, and I like the overall design. I think we still want to keep all the normalization logic in FNormalize, without adding boolean flags for specific cases. The more boolean flags we have, the more difficult it is for a developer to know the rules for all flags. For example, a developer would need to check if the TEnforceExplicitTupleInArgs is used as part of the well-formed check, whether it triggers an assert during normalization, whether it triggers a normalization step during normalization, whether the normalization step applies when parsing TVMScript, etc. By implementing each of these features on top of the same FNormalize functionality, new operator-specific normalization rules can be implemented without adding to a developer's mental overhead. A developer only needs to know that the new normalization is handled the same as all existing integrations.

(Also, see the other comment for performance benchmarking.)

[tqchen]

after thinking a bit more, i now agree that we can reuse FNormalize in wellform check. thanks for proposing the change

[tqchen]

We can use this function https://github.com/apache/tvm/blob/main/include/tvm/ir/op.h#L476

[Lunderberg]

Thank you, and updated! I had checked for a single-parameter .get, and an iterator-style .find, but hadn't found the two-parameter .get.

Updated to use if (auto func_normalize = op_map_normalize_.get(call->op, nullptr); func_normalize != nullptr), here and in well_formed.cc.

[tqchen]

https://github.com/apache/tvm/blob/main/include/tvm/ir/op.h#L476 we can directly use this function to simplofy the logic

[tqchen]

op_map_normalize_.get(call->op, nullptr)

[Lunderberg]

Thank you, and updated to use if (auto func_normalize = op_map_normalize_.get(call->op, nullptr); func_normalize != nullptr).

[Lunderberg]

Thank you, and changes made as suggested!