[IR] [Refactor] Support return in serial kernels

[yuanming-hu]

Remove ArgStoreStmt

taichi/taichi/ir/ir.h

Line 985 in f5373b1

class ArgStoreStmt : public Stmt {

Currently the only use of ArgStoreStmt is for returning values in a kernel. I personally feel like this is a confusing design (#839 (comment)). I suggest removing this statement and replace it with KernelReturnStmt.

Add KernelReturnStmt

KernelReturnStmt can only work in serial offloaded tasks.

• Set return values;
• Terminate the kernel. Actually, we should enforce that KernelReturnStmt can only be the final statement of a serial kernel. If we add this restriction, then no need to worry about terminating the kernel.

Benefits

1. SNode readers can finally use return
2. We can allow return values of arbitrary kernels, e.g.

@ti.kernel
def compute_sum():
  s = 0.0  
  for i in x:
    s += x[i]
  return s

s = compute_sum()

Discussions

• How about returning a vector/matrix?

[archibate]

Cool! This is so cool because we never use ArgStoreStmt in python-defined kernels. But removing ArgStoreStmt, what if we want multiple return values like the pythonicreturn a, b?

[yuanming-hu]

Good point. KernelReturnStmt should have a member std::vector<Stmt *> return_values to support multiple return values.

[archibate]

Does return make sense in diff-kernels? e.g.:

@ti.kernel
def func(x: ti.i32):
  return x ** 2

@ti.kernel
def func_grad(x: ti.i32):
  return 2 * x  # am I understanding correct?