return vs hcl.return_

[zhenkun]

What is the difference between return and hcl.return_? When to use what?

[seanlatias]

The main difference between return and hcl.return_ is that the former one would not be recognized by the HeteroCL compiler. In other words, the returned value(s) will always be inlined. For hcl.return_, currently it has two use cases. First, you can use it inside hcl.compute API (and its family such as hcl.update and hcl.scalar) for computing the RHS value. Note that currently we do not support nested stages inside those APIs. Second, you can use it inside hc.def_ as a general return statement. Following we show an example:

# Example 1 - return and hcl.return_ have the same results
# In this example, we use hcl.return_ in hcl.compute

# with Python return
def f(arg):
    return arg+1
A = hcl.compute((10,), lambda x: f(x))

# with hcl.return_
def f(arg):
    hcl.return_(arg+1)
A = hcl.compute((10,), lambda x: f(x))

# Equivalent Python code
for x in range(0, 10):
    A[x] = x+1

In the above example, we can see that with Python return, the returned value is directly inlined, which is the same for hcl.return_.

# Example 2 - return and hcl.return_ work differently
# In this example, we still use hcl.compute API. 
# However, since Python cannot recognize HeteroCL API, it doesn't understand the conditional statement.

# with Python return
def f(arg):
    with hcl.if_(arg > 5):
        return arg+2
    with hcl.else_():
        return arg+1
A = hcl.compute((10,), lambda x: f(x))

# Equivalent Python code - The Python program returns when it hits the first return statement
# Thus it never visits the else branch. Also, the if branch is omitted because Python doesn't understand it
for x in range(0, 10):
    A[x] = x+2

# with hcl.return_
def f(arg):
    with hcl.if_(arg > 5):
        hcl.return_(arg+2)
    with hcl.else_():
        hcl.return_(arg+1)
A = hcl.compute((10,), lambda x: f(x))

# Equivalent Python code - This one acts as expected
for x in range(0, 10):
    if x > 5:
        A[x] = x+2
    else:
        A[x] = x+1

As you can see, the Python return doesn't work in this case because: 1) It cannot recognize HeteroCL imperative API, and 2) It simply returns when hitting the return statement.

# Example 3: Differences between using def+reutrn and hcl.def_+hcl.return_

# with Python return + def
def f(arg):
    return arg+1
A = hcl.compute((10,), lambda x: f(x))

# Equivalent Python code - f is inlined
for x in range(0, 10):
    A[x] = x+1

# with hcl.return_ + hcl.def_
@hcl.def_([()])
def f(arg):
    hcl.return_(arg+1)
A = hcl.compute((10,), lambda x: f(x))

# Equivalent Python code - f is not inlined
def f(arg):
    return arg+1

for x in range(0, 10):
    A[x] = f(x)

In this example, the behavior of Python return simply inlines the RHS as we have explained. If we don't want the inlining, we can use hcl.def_ to manually create a module and use hcl.return_ inside it. Note that we must use hcl.return_ inside hcl.def_.