Skip to content

[TIR] Support Return in TIR - #7084

Merged
tqchen merged 21 commits into
apache:mainfrom
ZihengJiang:tir
Jan 16, 2021
Merged

[TIR] Support Return in TIR#7084
tqchen merged 21 commits into
apache:mainfrom
ZihengJiang:tir

Conversation

@ZihengJiang

Copy link
Copy Markdown
Contributor

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

Loading
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants