Struct argument to kernel has no function attributes

[oliver-batchelor]

A struct passed into a kernel as an argument doesn't have function attributes as a normal struct does, an example showing the difference below.

import taichi as ti
ti.init()

@ti.dataclass
class Foo:
    x: ti.f32
    y: ti.f32

    @ti.func
    def add(self, other:ti.template()):
        return Foo(self.x + other.x, self.y + other.y )

@ti.kernel
def foo(x:Foo):
  print(x.add(x).x)

@ti.kernel
def bar():
  x = Foo(1, 2)
  print(x.add(x).x) #Works

bar() # Works
foo(Foo(1, 2))   # Fails - Struct2 has no attribute 'add'

[Taichi] version 1.6.0, llvm 15.0.4, commit cc4cacfa, linux, python 3.8.16
[Taichi] Starting on arch=x64
Traceback (most recent call last):
  File "test_types.py", line 23, in <module>
    bar() # Works
  File "/home/oliver/mambaforge/envs/conan_ros/lib/python3.8/site-packages/taichi/lang/kernel_impl.py", line 1038, in wrapped
    raise type(e)('\n' + str(e)) from None
taichi.lang.exception.TaichiTypeError: 
File "test_types.py", line 20, in bar:
  print((x + x).x) #Works
         ^^^^^
unsupported operand type(s) for +: 'Struct3' and 'Struct3'

[bobcao3]

Ah... this is something funny. Struct arguments are compiled into IR struct type which does not have encapsulation. We might need to support UFCS to make such things like member methods work in Taichi. @ailzhang @lin-hitonami

[lin-hitonami]

Maybe we can implement this feature in the AST transformer.