Taichi: Frontend kernel/function structure checking

[yuanming-hu]

• Disallow return in ti.kernel
• Enforce exactly one return in ti.func
• Disallow function definition in ti.kernel and ti.func

[archibate]

If I want to add multi-return support, is there any suggestions?

[yuanming-hu]

If I want to add multi-return support, is there any suggestions?

Oh that would be great! We might need to think about types:

@ti.func
def foo(x):
  if x > 0:
    return int(x)
  else:
    return float(x)

What should be the return type in this case? We should do some type checking to prevent the ambiguity from happening. One solution could be using type hints:

@ti.func
def foo(x) -> float:
  if x > 0:
    return int(x)
  else:
    return float(x)

Note that we need an AST transform that translates the function above into

@ti.func
def foo(x) -> float:
  ret = float(0)
  if x > 0:
    ret = int(x)
  else:
    ret = float(x)
  return ret

Note that we probably also need to support ti.Matrixs.

[xumingkuan]

In the following case,

@ti.all_archs
def test_return():

  @ti.kernel
  def kernel():

    @ti.func
    def func():
      return 233

    print(func())
    return 1

  print(kernel())

we will visit both return statements with is_kernel = True, and then visit return 233 with is_kernel = False.
How can we distinguish return 233 in that it is inside @ti.func but not directly in @ti.kernel?

[yuanming-hu]

Please disallow function definition within taichi kernels as well :-)

[xumingkuan]

I think a solution may be first to check if the function has multi-return with different types, and if the answer is yes, raise an error to tell the user to convert these types explicitly.

[xumingkuan]

An example of allowing different return types with implicit conversions (only between subclasses and superclasses, disallowing conversion between basic types like i32 and f32) is https://github.com/xumingkuan/decaf-2019/blob/master/src/main/java/decaf/frontend/tree/Tree.java#L485 (upperBound refers to https://github.com/xumingkuan/decaf-2019/blob/master/src/main/java/decaf/frontend/type/Type.java#L97).
It would be confusing if we allow some implicit conversions (e.g. i32 to f64, since the latter can represent all values of the former) but not some others (e.g. i64 and f32, since we can't determine which type is "subtype" of the other).

[xumingkuan]

Ideally, we can support the following two cases at the same time:

@ti.func
def foo(x) -> float:
  if x > 0:
    return int(x)
  else:
    return float(x)

@ti.func
def foo(x):
  if x > 0:
    return float(int(x))
  else:
    return float(x)

[yuanming-hu]

I would actually challenge you guys to think about a bigger problem: how to support function definition/calls in Taichi IR. Currently, the function (@ti.func) calls are force-inlined, which has a few limitations such as potentially HUGE code size, and lack of recursion support. Ultimately we should consider compiling ti.func into a piece of Taichi IR (instead of a Python function that gets inlined), and then add a ReturnStmt as a new IR instruction.

[archibate]

Does IR have goto statements? If so, we can simulate function call stacks. Also sloved the problem that GLSL doesn't support recursion.(
Also some articles say that we can simulate recursion using trees.

[yuanming-hu]

No, no goto support needed. Simulating function call stacks is pretty advanced :-) For the first step, we should simply consider supporting CPU/CUDA where recursion is supported by the backend compiler.

It's more like a program representation issue within Taichi IR. For example, you might need to implement

• FuncDefStmt which contains a function prototype and function body
• FuncCallStmt with a function pointer (FuncDefStmt) and arguments
• ReturnStmt with a return value

[archibate]

Correct. Let's simply make recursion a limitation for OpenGL backend for now. And FuncDefStmt is something like OffloadStmt?
Btw, do we need @ti.inlinefunc after that?

[yuanming-hu]

Yeah sort of. We may also need a new IRNode (something like TaichiProgram) that wraps all the functions and the main function.

[yuanming-hu]

Btw, do we need @ti.inlinefunc after that?

Yes, actually I suggest that we make functions inline by default unless marked as ti.noinline. Or we can let the compiler make the decision.

[archibate]

I believe main function (ti.kernel) should be treated as a normal function too, and TaichiProgram contains a pointer to that IRNode.

[k-ye]

I guess it's fairly easy to support FuncDef, FuncCall in source-to-source backends. But for LLVM ones, do you need to create your own calling conventions. Or maybe LLVM already handles that..?

[yuanming-hu]

Or maybe LLVM already handles that..?

Exactly. LLVM already handles that. Nothing to worry about in that direction :-)

[github-actions[bot]]

Warning: The issue has been out-of-update for 50 days, marking stale.