taichi-dev / taichi

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

Structs & @ti.dataclass: variable cannot be assigned inside kernel or func #7579

Open jarmitage opened 1 year ago

jarmitage commented 1 year ago

[Taichi] version 1.4.1, llvm 16.0.0git, commit e67c674e, osx, python 3.10.6

t's not possible to assign to a StructType from within a kernel or a func

import taichi as ti
ti.init()
@ti.dataclass
class Particle:
    pos: ti.math.vec2
    vel: ti.math.vec2
p = Particle(pos=[0.5,0.5],vel=[0.5,0.5])
print(p.pos[0]) # 0.5
p.pos[0] = 0.6
print(p.pos[0]) # 0.6
@ti.kernel
def assign():
    p.pos[0] = 0.7 # error: Variable 'p.pos[0]' cannot be assigned. Maybe it is not a Taichi object?
assign()
print(p.pos[0])

Same error with ti.types.struct instead of @ti.dataclass:

Particle = ti.types.struct(pos=ti.math.vec2, vel=ti.math.vec2)

Same error inside func instead of kernel:

@ti.func
def assignfunc():
    p.pos[0] = 0.7
@ti.kernel
def assignkernel():
    assignfunc()
assignkernel()
BouchardMath commented 1 year ago

I'm not a maintainer, but your issue got me curious since I use ti.dataclass (and ti.data_oriented) successfully in my code.

Typically I initialize a taichi dataclass/struct the following way p = Particle.field(shape=1) and then assign values to it.

This fixes the issue with your code snippet.

In fact, p = Particle(pos=[0.5,0.5],vel=[0.5,0.5]) returns a python dict instead of a taichi.lang.struct.StructField. I'm not sure what the intended behaviour is supposed to be...

BouchardMath commented 1 year ago

Working code

import taichi as ti
ti.init(default_fp=ti.f64)

@ti.dataclass
class Particle:
    pos: ti.math.vec2
    vel: ti.math.vec2

p = Particle.field(shape=1)
p[0].pos.fill(0.5)
p[0].vel.fill(0.5)

print(p[0].pos[0]) # 0.5
p[0].pos[0] = 0.6
print(p[0].pos[0]) # 0.6

@ti.kernel
def assign():
    p[0].pos[0] = 0.7¸
assign()
print(p[0].pos[0]) # 0.7
jarmitage commented 1 year ago

@BouchardMath this is true and I also use that pattern. But according to the docs the other way should be possible too:

https://docs.taichi-lang.org/docs/type#struct-types-and-dataclass

neozhaoliang commented 1 year ago

@jarmitage @BouchardMath Hi, the error is because when you declared p = Particle(pos=[0.5,0.5],vel=[0.5,0.5]) in the Python scope, p will be a Python object (indeed, a wrapper of a Dict), and is treated as a global constant by the kernel. Hence you can read the members of p in a kernel, but cannot modify it. For example, you can read the members of p:

p = Particle(pos=[0.5,0.5],vel=[0.5,0.5])

@ti.kernel
def test():
    print(p.pos.x)

but cannot reassign it a value:

p = Particle(pos=[0.5,0.5],vel=[0.5,0.5])

@ti.kernel
def test():
    p.pos.x = 1.0

One solution is to define p inside the Taichi scope, so that p will be a struct and can be modified inside Taichi.

For example, you can change the code like below:

@ti.kernel
def test():
    p = Particle(pos=[0.5,0.5],vel=[0.5,0.5])
    p.pos.x = 1.0