It reports a bug when trying to access a value in static method from ti.kernel

[Morcki]

Describe the bug In static method Test.test, why can't I access the value of v[0].

To Reproduce

# sample code here
import taichi as ti

ti.init(arch=ti.cuda)

class Test:
    @staticmethod
    def test(v):
        print(v[0])
    
@ti.kernel
def test_rotate():
    X = ti.Vector([1.0, 0.0, 0.0])
    Test.test(X)

test_rotate()

After running the python script, it shows the error information as blew:

File "test_3.py", line 9, in test
    print(v[0])
AssertionError: __getitem__ cannot be called in Taichi-scope

[strongoier]

Hi @Morcki. Functions which are called in ti.kernel should be decorated with ti.func:

class Test:
    @staticmethod
    @ti.func
    def test(v):
        print(v[0])

[Morcki]

What if I add this static method to class Matrix(TaichiOperations)? Maybe like this:

@_gen_swizzles
class Matrix(TaichiOperations):
    def __init__(self, ...):
        pass
    ...
    @staticmethod
    def test(v):
        print(v[0])

Here, I want to use Matrix.test both in taichi scope and python scope, but failed when trying to access value of ti.Vector/ti.Matrix.

[Morcki]

Hi, @strongoier, I implement some built-in python method, the link is https://github.com/Morcki/taichi/commit/8309cd5c067b1d89a6bfd8d8e908466ae7fda90f

However, it runs too slowly in taichi scope(slower even than python), may you help to give some suggestions where I should do to run faster.

[strongoier]

Hi @Morcki. Could you show your speed comparison as well?

[Morcki]

Here is the pytest result, it takes 345.05s to run the 6 tests, which I think it shouldn't take that too long. Because there isn't such complex caculations. The math caculations in https://github.com/Morcki/taichi/commit/8309cd5c067b1d89a6bfd8d8e908466ae7fda90f only include sin, cos and @(matrix multiplication).

platform linux -- Python 3.8.10, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/taichigraphics/WorkSpace/dev/taichi, configfile: pyproject.toml
collected 6 items                                                                                                                                                                                             

test_matrix.py ......                                                                                                                                                                                   [100%]

======================================================================================== 6 passed in 345.05s (0:05:45) ========================================================================================

[bobcao3]

This is quite a large linear test case. It takes a long time to JIT compile large kernels like this due to the complexity of simplification optimization step. Also when running on GPU, you should expect a linear kernel performing much worse than the CPU. You should split up the tests into many small kernels, with each kernel only testing one or two operations

[Morcki]

Thanks a lot! I will try your suggestions

Yes, it indeed works, and reduce from 345s to just 17s.