[RFC] Transformer-based ti.func dispatching

[yuanming-hu]

(Continuing the discussions in #1338 (review))

Recently there have been recurring issues centered around calling Taichi-scope methods in Python-scope. For example, #1051 and #1338, #1379. Just to name a few.

Here I provide an AST transform-based approach, which I've been considering for a while (but not until today did I get a chance to write it down here....)

Note that we are already using the AST transform-based approach for @ti.kernel and @ti.func in Taichi-scope. This proposal simply extends that approach to handle calling @ti.func in Python-scope.

What prevents us from calling ti.func's in Python-scope?

It may seem that many ti.funcs are directly callable in Python. This is indeed true for simple functions (e.g. no branching, no mutable variables). However, this may not be the case for more complex functions, such as Matrix.__matmul__.

As a compiled language, Taichi does have subtle differences from raw Python. For example, its local variables are statically typed:

@ti.func
def foo():
    a = int(0)
    a = 0.4
    print(a) # 0, instead of 0.4

Another difference would be that Taichi uses lexical scoping yet Python uses function scoping:

@ti.func
def bar():
    if True:
        a = 1
    print(a) # SyntaxError: a is no longer available here

How to fix these incompatibilities?

For Taichi-scope kernels and functions, we have already successfully fixed these issues using an AST transformer (in transformer.py): the bar function above will be transformed into

def foo():
  import taichi as ti
  if 1:
    __cond = True
    ti.core.begin_frontend_if(ti.Expr(__cond).ptr)
    ti.core.begin_frontend_if_true()
    a = ti.expr_init(1)
    del a # Note: ensure lexical scoping
    ti.core.pop_scope()
    ti.core.begin_frontend_if_false()
    ti.core.pop_scope()
  ti.ti_print(a)

Similarly, we will be able to directly call @ti.func in Python-scope, if we have a PythonTransformer that translates @ti.func into an AST that is directly runnable by the Python interpreter.

For example,

@ti.func
def foo(x, y)
    a = 0
    if x > 0:
        b = y * 2
        a = b + 1.5
    else:
        a = x
    return a

Should be transformed (by PythonTransformer) into

def foo(x, y)
    a = variable(0) # integer local var
    if x > 0:
        b = variable(y * 2)
        a.set_val(b.val + 1.5) # convert to integer
        del b # Ensure Taich lexical scoping
    else:
        a.set_val(x)
    return a.val

The transformed function is directly executable by the Python interpreter. This enables an interpretation mode of @ti.func. As long as we implement the PythonTransformer, then all the @ti.func can run in both Python- and Taichi-scope.

I think many existing attempts (e.g., invoking functions that get dispatched to different implementations depending on if we are in Taichi- or Python- scope) are already close to this approach (i.e., dispatching functions to Taichi/Python-scope AST transformers.)

Transformer-based dispatching is the ultimate solution.

[archibate]

What's the point of running ti.func in Python-scope? I feel this is a kind of overkill. IMHO when ti.func is called from Python-scope, it should behave like a kernel, if with type-hints.