[TIR] Support Return in TIR

[ZihengJiang]

This PR supports building TIR function directly and allows TIR returns value directly.
Current approach is by adding an intrinsic for return. Example:

def add():
    a = tir.Var("a", "float32") 
    b = tir.Var("b", "float32") 
    c = a + b
    c = tir.call_intrin("float32", "tir.myreturn", c) 
    c = tir.Evaluate(c)
    func = tir.PrimFunc([a, b], c)
    mod = tvm.IRModule({'add': func})
    func = tvm.build(mod['add'])
    out = func(1.0, 2.0)
    print(out)

In the future, with the tvm script, we should be able to write some like:

@tvm.script
def add(a, b):
    tir.myreturn(a + b)

There are two things need to be discussed:

• Because return is a keyword in Python, currently I use myreturn, we should have another name for it.
• Whether we should support multiple return in a single function.

@tqchen @junrushao1994 @areusch

[areusch]

I think we may need to return at least DLTensor for AOT

[tqchen]

One additional comment, overall lgtm. @junrushao1994 @tkonolige @areusch please take another look and https://tvm.apache.org/docs/contribute/code_review.html#approve-and-request-changes-explicitly

[tkonolige]

Please add which data datatypes are supported to this error message

[tqchen]

Thanks @tkonolige @jroesch @ZihengJiang on discussing the issue of ICHECK and comments. While I in general agree that for cases that are user facing, having clear error message would be helpful.

I would be hesitant to enforce error message for every internal invariance check. This decision is more of an engineering tradeoff. For example are many ways to improve code robustness and readability, code comment is one of them. However, it does not necessarily mean we need to comment every line of code, even though it is possible to do so and doing so might improve readability slightly.

Developers may want to spend their attentions in other matters and we also need to balance these concerns, so as long as the code structure is clear, we could focus the code review on more important things -- e.g. documenting the interface clearly, making sure the overall logics are correct.

Similarly, while it is always possible to add error message to every internal invariance checks, I don't think we are required to do so for every care. We can find many similar examples(like the ones below) in other large code base like TensorFlow and PyTorch, doing so doesn't make pytorch or tensorflow not product ready.

• https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/throughput_benchmark.cpp#L25
• https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/replicate_per_replica_nodes_test.cc#L54

Of course this does not mean that we should not add error messages to the checks that could cause user facing errors. We should actually shift our attention to address these checks quickly, while giving flexibility to the developers in other cases that have minimum gain.

In summary, context of the code matters, and I believe code review should focus on understanding the overall logics, trying to think about readability in a broader context(other than error message itself), while encouraging things to move on smoothly once things met a good quality.

In this case we have discussed and I believe the ICHECK on pointer null invariant checking is not necessary. While it is indeed helpful to have better error message in other two cases.

[tkonolige]

I'm a little confused on how we use return in tir? How is the example you give different if we just remove the return?

a = tir.Var("a", "float32") 
b = tir.Var("b", "float32") 
c = a + b
# c = tir.call_intrin("float32", "tir.myreturn", c)
c = tir.Evaluate(c) /# doesn't this return a + b?

[ZihengJiang]

@tkonolige
This will generate a function that only evaluate but not return:

def func(a, b):
    c = a + b
    # return c

Think and have a try by yourself would make you understand other's PR.

[tkonolige]

I'm just trying to understand things better. There's not really any documentation on TIR.

I had thought that the last statement in every TIR expression was implicitly returned. Am I correct in saying that before this PR calling a TIR function would always return nullptr/None? And that now you can actually return a value?

[ZihengJiang]

@tkonolige
Happy to address your comment and hope it will make you understand TIR. Yes, usually a zero value would be returned before this PR and now a value can be explicitly returned.

[tkonolige]

Thanks for the clarification.

Given that return is a useful feature, could you add tir.myreturn (or whatever name you think best) as a function to python/tvm/tir/stmt.py. That way we don't have to call the intrinsic manually. Could you also add documentation with that function. An example would be nice.

(Trying to use the submit review feature correctly here.)

[tqchen]

Thanks @tkonolige , I agree, and perhaps https://github.com/apache/tvm/blob/main/python/tvm/tir/op.py is a better place given other intrinsics are in there

[tqchen]

@ZihengJiang would be great if we can act on the thread and bring it in :)

[tqchen]

Thanks @ZihengJiang please check the comment on control flow termination

[tqchen]

Thanks @ZihengJiang @tkonolige @junrushao1994 @jroesch @areusch