Petition to keep ti.pyfunc

[AmesingFlank]

I understand that the plan is to remove ti.pyfunc in 1.6.0, but I found this to be a very important feature which I use all the time, so I'd like to make a case for keeping it in taichi.

In #4414, @lin-hitonami explained that even without ti.pyfunc, you can still create functions that work in both taichi scope and python scope using the following pattern:

def fmod_python(a, b):
    return a-ti.floor(a/b)*b

fmod_taichi = ti.func(fmod_python)

While this does work for the simplest scenarios, there are many use cases that this pattern cannot support:

• Multiple functions that calls each other. For example:

xs = ti.field(float, shape=100)

@ti.pyfunc
def f(i):
    return xs[i] + 1

@ti.pyfunc
def g(i):
    return f(i) + xs[i] + 1

@ti.kernel
def k():
    print(g(3))

k()

The solution proposed in #4414 does not work here:

xs = ti.field(float, shape=100)

def f(i):
    return xs[i] + 1

def g(i):
    return f(i) + xs[i] + 1

@ti.kernel
def k():
    print(ti.func(g)(3))

k()

because in ti.func(g), the f function that it calls is not transformed by ti.func, so this triggers an error:

  File "c:\users\ldfra\code\taichi\taichi\python\taichi\lang\util.py", line 308, in wrapped
    assert in_python_scope(), \
AssertionError: __getitem__ cannot be called in Taichi-scope

• member functions of @ti.data_oriented classes:
  Consider for example:

@ti.data_oriented
class C:
    @ti.pyfunc
    def f(self, i):
        return i + 1

c = C()

print(c.f(1))

@ti.kernel
def k():
    print(c.f(1))

k()

this would be very awkward without ti.pyfunc. The only solution I see is by doing

@ti.data_oriented
class C:
    def f(self, i):
        return i + 1

@ti.kernel
def k():
    print(ti.func(C.f)(c, 1))

which is quite bad in terms of readability

• member functions struct types (@ti.dataclass). This is the same as with ti.data_oriented classes, because members of struct types are also callable from both taichi and python scope.

These are all very common use cases, and not having ti.pyfunc really makes writing large-scale taichi programs a lot more awkward. So please consider keeping ti.pyfunc. Alternatively, I would also greatly appreciate if you could help me identify other patterns that could solve all of the three aforementioned problems!

cc @lin-hitonami @strongoier @k-ye @ailzhang @yuanming-hu

[ailzhang]

Hey @AmesingFlank thanks for bringing this up! We discussed this offline with @strongoier and @lin-hitonami and totally agreed this is a valid use case! We'll keep pyfunc and raise better error message when an invalid pyfunc is written/called instead. Cheers!

[AmesingFlank]

Hey @AmesingFlank thanks for bringing this up! We discussed this offline with @strongoier and @lin-hitonami and totally agreed this is a valid use case! We'll keep pyfunc and raise better error message when an invalid pyfunc is written/called instead. Cheers!

That's awesome, thank you!