跳转至主要内容
Version: v1.1.3

面向数据对象式编程

Taichi 是一种 面向数据 编程 (DOP) 的语言。 但是,单纯的 DOP 会使模块化变得困难。

为了允许代码模块化,Taichi 从面向对象编程 (OOP) 中借鉴了一些概念。

为了方便起见,我们将称此混合方案为面向数据对象式编程 (ODOP)。

面向数据类

简介

如果您需要定义一个 Taichi kernel 作为Python 类成员函数,请使用 @ti.data_oriented 来修饰该类。 然后您就可以在您的 面向数据 Python 类中定义 ti.kernels 和 ti.funcs。

note

The first argument of the function should be the class instance ("self"), unless you are defining a @staticmethod.

一个简单的例子:

@ti.data_oriented
class TiArray:
def __init__(self, n):
self.x = ti.field(dtype=ti.i32, shape=n)

@ti.kernel
def inc(self):
for i in self.x:
self.x[i] += 1

a = TiArray(32)
a.inc()

Taichi field 不仅可以在 init 函数中定义,也可以在 Python 作用域内任一面向数据类的函数内定义。 例如,

import taichi as ti

ti.init()

@ti.data_oriented
class MyClass:
@ti.kernel
def inc(self, temp: ti.template()):
for I in ti.grouped(temp):
temp[I] += 1

def call_inc(self):
self.inc(self.temp)

def allocate_temp(self, n):
self.temp = ti.field(dtype = ti.i32, shape=n)


a = MyClass()
# a.call_inc() cannot be called, since a.temp has not been allocated at this point
a.allocate_temp(4)
a.call_inc()
a.call_inc()
print(a.temp) # [2 2 2 2]
a.allocate_temp(8)
a.call_inc()
print(a.temp) # [1 1 1 1 1 1 1 1]

另一个内存回收利用的例子:

import taichi as ti

ti.init()

@ti.data_oriented
class Calc:
def __init__(self):
self.x = ti.field(dtype=ti.f32, shape=16)
self.y = ti.field(dtype=ti.f32, shape=4)

@ti.kernel
def func(self, temp: ti.template()):
for i in range(8):
temp[i] = self.x[i * 2] + self.x[i * 2 + 1]

for i in range(4):
self.y[i] = max(temp[i * 2], temp[i * 2 + 1])

def call_func(self):
fb = ti.FieldsBuilder()
temp = ti.field(dtype=ti.f32)
fb.dense(ti.i, 8).place(temp)
tree = fb.finalize()
self.func(temp)
tree.destroy()


a = Calc()
for i in range(16):
a.x[i] = i
a.call_func()
print(a.y) # [ 5. 13. 21. 29.]

面向数据类的继承

data-oriented 属性将会被自动地传到所继承的Python类之外。 这意味着 Taichi Kernel 可以被任何经由 @ti.data_oriented 修饰过的父类所调用。

一个示例:

import taichi as ti

ti.init(arch=ti.cuda)

class BaseClass:
def __init__(self):
self.n = 10
self.num = ti.field(dtype=ti.i32, shape=(self.n, ))

@ti.kernel
def count(self) -> ti.i32:
ret = 0
for i in range(self.n):
ret += self.num[i]
return ret

@ti.kernel
def add(self, d: ti.i32):
for i in range(self.n):
self.num[i] += d


@ti.data_oriented
class DataOrientedClass(BaseClass):
pass

class DeviatedClass(DataOrientedClass):
@ti.kernel
def sub(self, d: ti.i32):
for i in range(self.n):
self.num[i] -= d


a = DeviatedClass()
a.add(1)
a.sub(1)
print(a.count()) # 0


b = DataOrientedClass()
b.add(2)
print(b.count()) # 1

c = BaseClass()
# c.add(3)
# print(c.count())
# The two lines above will trigger a kernel define error, since class c is not decorated by @ti.data_oriented

Python 内置修饰器

在 Python 中修饰符通常会被预编译, @staticmethod1@classmethod2面向数据 类中可以修饰 Taichi kernel

staticmethod 示例 :

import taichi as ti

ti.init()

@ti.data_oriented
class Array2D:
def __init__(self, n, m, increment):
self.n = n
self.m = m
self.val = ti.field(ti.f32)
self.total = ti.field(ti.f32)
self.increment = float(increment)
ti.root.dense(ti.ij, (self.n, self.m)).place(self.val)
ti.root.place(self.total)

@staticmethod
@ti.func
def clamp(x): # Clamp to [0, 1)
return max(0., min(1 - 1e-6, x))

@ti.kernel
def inc(self):
for i, j in self.val:
ti.atomic_add(self.val[i, j], self.increment)

@ti.kernel
def inc2(self, increment: ti.i32):
for i, j in self.val:
ti.atomic_add(self.val[i, j], increment)

@ti.kernel
def reduce(self):
for i, j in self.val:
ti.atomic_add(self.total[None], self.val[i, j] * 4)

arr = Array2D(2, 2, 3)

double_total = ti.field(ti.f32, shape=())

ti.root.lazy_grad()

arr.inc()
arr.inc.grad()
print(arr.val[0, 0]) # 3
arr.inc2(4)
print(arr.val[0, 0]) # 7

with ti.ad.Tape(loss=arr.total):
arr.reduce()

for i in range(arr.n):
for j in range(arr.m):
print(arr.val.grad[i, j]) # 4

@ti.kernel
def double():
double_total[None] = 2 * arr.total[None]

with ti.ad.Tape(loss=double_total):
arr.reduce()
double()

for i in range(arr.n):
for j in range(arr.m):
print(arr.val.grad[i, j]) # 8

classmethod 示例:

import taichi as ti

ti.init(arch=ti.cuda)

@ti.data_oriented
class Counter:
num_ = ti.field(dtype=ti.i32, shape=(32, ))
def __init__(self, data_range):
self.range = data_range
self.add(data_range[0], data_range[1], 1)

@classmethod
@ti.kernel
def add(cls, l: ti.i32, r: ti.i32, d: ti.i32):
for i in range(l, r):
cls.num_[i] += d

@ti.kernel
def num(self) -> ti.i32:
ret = 0
for i in range(self.range[0], self.range[1]):
ret += self.num_[i]
return ret

a = Counter((0, 5))
print(a.num()) # 5
b = Counter((4, 10))
print(a.num()) # 6
print(b.num()) # 7

Taichi dataclasses

Taichi provides custom struct types for developers to associate pieces of data together. 不过,如果能有以下两样往往会更加方便:

  1. 一个更加面向对象的结构体类型的 Python 表示。
  2. 与结构体类型关联的函数。 (C++ 风格结构体)

To achieve these two points, developers can use the @ti.dataclass decorator on a Python class. 这一做法受到 Python 的 dataclass 功能的启发,这一功能使用包含注释的 class field 来创建数据类型。

从 Python 类创建一个结构体

以下是一个从 Python 类创建一个 Taichi 结构体类型的示例:

@ti.dataclass
class Sphere:
center: vec3
radius: ti.f32

这将生成与下面代码生成 完全一致的类型:

Sphere = ti.types.struct(center=vec3, radius=ti.f32)

Using the @ti.dataclass decorator will convert the annotated fields in the Python class to members in the resulting struct type. 上述两个示例都会创建一个相同的结构体 field。

sphere_field = Sphere.field(shape=(n,))

将函数与结构体类型关联

Python 类可以有附属函数,Taichi 结构体类型也可以有。 下面这段代码衍生自上面的两个例子,展示了如何将函数附属在一个结构体上

@ti.dataclass
class Sphere:
center: vec3
radius: ti.f32

@ti.func
def area(self):
# a function to run in taichi scope
return 4 * math.pi * self.radius * self.radius

def is_zero_sized(self):
# a python scope function
return self.radius == 0.0

与结构体相关联的函数与正常函数都遵守相同的 作用域规则, 因为它所在的作用域可以是 Taichi 作用域也可以是 Python 作用域。 现在,Sphere 结构体类型生成的每个实例都会添加上上述函数。 这些函数可以这样被调用:

a_python_struct = Sphere(center=vec3(0.0), radius=1.0)
# calls a python scope function from python
a_python_struct.is_zero_sized() # False

@ti.kernel
def get_area() -> ti.f32:
a_taichi_struct = Sphere(center=vec3(0.0), radius=4.0)
# return the area of the sphere, a taichi scope function
return a_taichi_struct.area()
get_area() # 201.062...

Notes

  • Inheritance of Taichi dataclasses is not implemented.
  • While functions attached to a struct with the @ti.dataclass decorator is convenient and encouraged, it is actually possible to associate a function to structs with the older method of defining structs. 如上所述,两种定义结构体类型的方法最后的生成是完全相同的。 要做到这点,请在调用 ti.types.struct 时使用 __struct_methods 参数。
@ti.func
def area(self):
# a function to run in taichi scope
return 4 * math.pi * self.radius * self.radius

Sphere = ti.types.struct(center=vec3, radius=ti.f32,
__struct_methods={'area': area})