taichi-dev / taichi

Productive, portable, and performant GPU programming in Python.
https://taichi-lang.org
Apache License 2.0
25.05k stars 2.26k forks source link

ndarray struct for support in data_oriented class #7076

Open ailzhang opened 1 year ago

ailzhang commented 1 year ago

Users from ti.field might easily write something like below but taichi will complain about missing arg annotation for self.a such that struct-for will fail. I guess we could either provide a best practice for users about how to work around this (like passing in self.a as an arg with proper annotation), or think about a way how users can annotate self properly so that they don't have to annotate member ndarrays in each kernel one by one. This is kinda similar to the idea what users ask for as argpack IMHO.

import taichi as ti
ti.init(ti.cpu)

@ti.data_oriented
class ExampleModule:
    def __init__(self, n: int):
        self.a = ti.ndarray(ti.f32, (n, n))

    @ti.kernel
    def some_kernel(self):
        for i, j in self.a:
            self.a[i, j] = 1.0
        pass

m = ExampleModule(8)
m.some_kernel()
print(m.a)
oliver-batchelor commented 1 year ago

This is also very nice if you're doing any "metaprogramming" with ti.template(), for example if I pass the ExampleModule to some other data_oriented class as a ti.template() - it can then call back my some_kernel() which can access a.

I don't think there's any other way of doing this right now (and it works very nicely if a is a field)

feisuzhu commented 1 year ago

@jim19930609