[UX] Changes in new Relax Parser

[yongwww]

More and more UX issues/feature requests were created in the past few weeks, some of them are related to Relax Parser. The new relax parser is under active development (thanks to @junrushao1994 @Hzfengsy @cyx-6 @yelite @YuchenJin), more details about the new parser/printer please take a look at TVMScript Unified Printer and metaprogramming. We would like to summarize the changes we are going to make in the new parser, any comments or suggestions are very welcome!

• C0: The body of relax function could be a single Return, although in the majority of cases it is a SeqExpr (list of BindingBlocks plus an Expr).

# func_1's body is a function
@R.function
def func_1(x: Tensor((m, n), "float32")):
    @R.function
    def local_func(y: Tensor((m, n), "float32")) -> Tensor((m, n), "float32"):
        s = relax.add(x, y)
        return s
    return local_func

@R.function
def func_2(x: Tensor((m,), "float32")):
    return relax.call_tir("my_extern", (x,), (tir.cast("int32", m),), dtype="float32")

@R.function
def func_3(x: Tensor((m,), "float32")):
    return x # or relax.const(1)

• C1: Deprecate () in favor [] for Callable (e.g. Callable[["float32"], "float32]) to make it more pythonic.
• C2: Support both [] and () for shape of Tensor (e.g. Tensor((2, 3), "float32"), Tensor([2, 3], "float32")). Currently the return type of relax Function could be something like Tensor(None, "float32", ndim=2) even the return shape is able to be deduced, we should use the deduced shape (e.g.Tensor((m, n), "float32")) for this case.
• C3: We should parse the IRModule to figure out the function declaration first, then parse the IRModule again. [Bug] [Normalizer] checked_type for function call is not defined #205

# Example by Prakalp
def test_function_call():
    @tvm.script.ir_module
    class InputMod:
        @R.function
        def foo(x: Tensor((2, 3), "float32")) -> Tensor:
            add0 = R.add(x, x)
            return add0

        @R.function
        def main(x: Tensor((2, 3), "float32")) -> Tensor:
            output_foo = foo(x)
            add_again = R.add(x, output_foo)
            return add_again

    mod_before = InputMod
    mod_before.show()

    print("Well Formed: ", relax.analysis.well_formed(mod_before))
    R.parser.pretty_print(mod_before["foo"].body.body.checked_type)
    R.parser.pretty_print(mod_before["main"].body.blocks[0].bindings[0].value.checked_type)

• C4: Add type/shape deduction for standalone relax Function. Currently the Normalize pass is applied for IRModule, but not for Function. [Bug] left-hand side variable type not deducted automatically with call_tir. #209

# test case by Jiawei
@R.function
def bug(x: Tensor((32, 32), "float32")) -> Tensor:
    with R.dataflow():
        lv0 = R.call_tir(my_relu, (x,), (32, 32), dtype="float32")
        R.output(lv0)
    return lv0

bug.body.blocks[0].bindings[0].var.checked_type

• C5: Subscript the tuple literal, and support empty tuple [UX][TVMScript] Feature requests #202

@R.function
def return_unit() -> Tuple():
    return ()

@R.function
def tuple_subscript():
    v0 = R.const(0)
    v1 = R.const(1)
    t = (v0, v1)
    return t[0] # t(0) is supported already

• C6: If a variable was bound to a call which doesn’t have a return value (e.g. callExternFunc with side effects), empty tuple will be used as the type of the var. [UX][TVMScript] Feature requests #202 any other suggestions for this scenario?
• C7: Some literals should be parsed into scalars. Integer literals in shape definition should be parsed into IntImm (e.g., 32 in Tensor((32, 32), "float32")).

# Example by Steven
@R.function
def return_int() -> Tensor((), "int32"): # or "float32"
    return 1.0 # 1 was parsed as an IntImm, but it should be a constant literal

[psrivas2]

Thanks @yongwww for summarizing the parser requirements in one place.

Can we also add #197 to the list as well. It is more of a sugar support than a bug in current parser, but an important one IMO. You can use the very short example below:

Support R.emit_te sugar

@tvm.script.ir_module
class Before:    
    @R.function
    def main(x: Tensor((10, 20), "float32")) -> Tensor:
        with R.dataflow():
            gv = R.emit_te(topi.add, x, x)
            R.output(gv)
        return gv

[psrivas2]

My 2 cents on C3

Most imperative languages rely on user to provide function signature than deducing it. In Relax we would rely on shape/type deduction to deduce the function signature. This is an odd choice given that we cannot in a single pass deduce any function signature (because of function calls). For example how many passes would it take to deduce the return type of foo as Tensor((10, 20), "float32") if let's say we always process functions in the order foo->bar->baz->qux in the following example.

@tvm.script.ir_module
class InputMod:
    @R.function
    def foo(x: Tensor((10, 20), "float32")):
        return bar(x)
    
    @R.function
    def bar(x: Tensor((10, 20), "float32")):
        return baz(x)

    @R.function
    def baz(x: Tensor((10, 20), "float32")):
        return qux(x)

    @R.function
    def qux(x: Tensor((10, 20), "float32")):
        return x

If we assume that the user provides all function signatures, we can break this long chain and ensure that all AST node shape/types are known in two passes: (1) add all function signatures to list (2) deduce types for all values within a function.

Not sure if we support recursion, but if we do, then that would also require user provided function signatures.

[slyubomirsky]

Would we not be able to use a unification-style type inference algorithm (like Hindley-Milner Algorithm W, which was what we based Relay's type inference on) to handle cases like that, including with recursion? OCaml includes imperative features and is still able to infer types in all cases. We would have to be careful with our type system, but I think inference would be possible. (This is the sort of thing a full language spec might help us to reason about haha.)

The way this would work would be to assign types with type variables where they're not known (so a function with two arguments initially has type 'a -> 'b -> 'c, where the letters are type variables), then unify those type variables with concrete types as we encounter them, giving an error when one of these unifications is incompatible. It would require another pass, so it wouldn't work with our current style of inferring types in-place. It depends on how much of a burden we want to put on users to annotate types. (FWIW, we already expect users to provide type annotations for extern function calls.)

[slyubomirsky]

On C1, regarding whether to use square brackets or parentheses for notation, I would support in principle following Python's conventions wherever possible to avoid surprises. Square brackets look a little odd, but I think deviations from Python should be carefully considered (e.g., we do deviate from Python with regard to lexical scoping, and I think there are good reasons for that)

[slyubomirsky]

On C7, I wonder if we should have some kind of conversion (implicit or explicit) between tensors that are scalars (zero-rank tensors) and IntImms/FloatImms etc. PackedFuncs can use either. Can TIR PrimFuncs take IntImms/FloatImms as arguments? What about doing computations with shapes?

[slyubomirsky]

Also, I would really be in favor of not requiring functions to have a return! We could have an implicit return () or even have a None value or something like that (what is returned by a PackedFunc with Void return type? relax.print() prints "None" for that)

[psrivas2]

Would we not be able to use a unification-style type inference algorithm (like Hindley-Milner Algorithm W, which was what we based Relay's type inference on) to handle cases like that, including with recursion

Thanks for the pointer to Relay's type inference system. I'll take your word for it that it is doable :) and someone familiar with the algorithm could work towards a prototype for Relax type inference system.

It depends on how much of a burden we want to put on users to annotate types. (FWIW, we already expect users to provide type annotations for extern function calls.)

It might be my bias from LLVM family of IRs, but I believe annotating function signatures is not much of a burden (we ask users to annotate parameters already, so it is just the return type). It seems pretty straightforward to me that with function signatures annotated, we can infer type of all sub expressions and relieve the Relax type system of type inferencing function calls.

Also, I would really be in favor of not requiring functions to have a return! We could have an implicit return () or even have a None value or something like that (what is returned by a PackedFunc with Void return type? relax.print() prints "None" for that)

To be clear, you are advocating for skipping the return statement only when there is nothing to return?

[slyubomirsky]

I also do not think that annotating types is that much of a burden, but those who like Python may not want to do it. (You make the good point that we already expect arg types to be annotated, so requiring more expectations is not a large additional burden.) Just tossing ideas out there, something like C++'s auto might be a middle ground to consider. But I think type inference is the best way to maintain a Python-like interface--if it's feasible! Doing something like Algorithm W depends on other details of the language (for example, OCaml's parametric polymorphism makes it possible to infer the type 'a -> 'a for let f x = x. If it didn't have parametric polymorphism, then there's no type you could assign it that would work for all possible x).

To be clear, you are advocating for skipping the return statement only when there is nothing to return?

Yes, I think we should treat the lack of a return statement as implicitly returning None. This may be the case for calling PackedFuncs that have side effects. Similarly, function calls whose results are not bound to a return should be parsed as binding to a new, never-used variable. (These are Expr statements in Python.)

[slyubomirsky]

Another request: We should be able to parse Python if expressions: v1 if cond else v2. Currently this is treated as an error.

[slyubomirsky]

Here is another case that is not handled that might be reasonable to parse too:

@R.function
def branch(x: Tensor((), "int32")):
    if x:
        return x
    else:
        return x

The parser currently complains that the body needs to end with a returned expression. It might be hard to formulate a general rule here; it may be worth spec'ing out in more detail. Do we require both branches of a conditional to have a return if either one does? Do we create an implicit else branch if there is no else branch provided? Etc.

[psrivas2]

The current parser_v1 assumes that all of the bindings within a dataflow block are either ast.Assign or ast.UnassignedCall. See code here. However, dataflow block bindings could also have function in RHS which is not handled by the current parser.

Here is a unit test that should pass but would fail with the current parser:

@tvm.script.ir_module
    class Mod:
        @R.function
        def main(x: Tensor((10, 5), "float32")) -> Tensor((10, 5), "float32"):
            with R.dataflow():
                # fails to parse the binding below
                @R.function
                def inner(y: Tensor((10, 5), "float32")) -> Tensor((10, 5), "float32"):
                    s: Tensor((10, 5), "float32") = relax.add(y, y)
                    return s
                gv1 = inner(x)
                R.output(gv1)

            gv2: Tensor((10, 5), "float32") = relax.add(gv1, gv1)
            return gv2

[tqchen]

Closing for now and let us open new issues for the new parser