Taichi Metaprogramming and Object-Oriented Programming

引言

太极图形课会分成「太极」语言与计算机「图形学」两部分。前半部分主讲太极编程语言的基础语法,高级用法,以及调试和优化太极程序的方式。后半部分侧重于图形学知识,我们将以实践者的角度来聊一聊基于太极语言的渲染和仿真。

元编程用于提升代码复用性,面向对象编程主要用来提供代码的可扩展性和可维护性。

元编程

Meta

Metaprogramming

Metaprogramming 是凌驾于 programming 之上的一种编程方式,传统的程序是直接产生程序输入数据产生结果的;而元编程则是让你编写一部分代码,然后让这部分代码产生真正可以运行的代码,然后让运行的代码来得到你的结果。

举个例子,造车是你最后的结果,那我们为了造车需要搭建一条流水线(也就是计算过程)。你把流水线上放上数据,数据一路推过流水线完成计算之后就变成了一辆车。

当然同一条流水线可以生产出很多很多的同样的汽车。而这时候你老板跑过来说想要不同的汽车,那常理来说就要加钱,做三条流水线:

但老板又说了,最近公司日子有点紧,没办法加钱,但这三条流水线你还得造出来。那怎么办呢?我们可以不直接造流水线,直接去生产一个图纸:

我们的制作内容从制作一个流水线变成了画一个流水线图纸。这个图纸也许可以引申一些额外的参数,存在一些细微的改动,最后实例化生成三条不同的流水线。

Metaprogramming in Taichi

  • 太极很多时候被用来写一个高效的并行计算的,可能会写一些模版类处理 2D/3D 数据
  • 给一些编译器理解的语句,编译器就会生成和你编写时完全不同的代码,只需编写少量代码就可以达到比较高的效率

Copy a Taichi field to another

def copy_4(src, dst):
for i in range(4):
dst[i] = src[i]

a = ti.field(ti.f32, 4)
b = ti.field(ti.f32, 4)

copy_4(a, b)

现在有两个类型为 Taichi field 的四维数组,我想做一次拷贝。如果在 Python- scope 直接让 b = a 的话会指向同一块内存,上面是一个比较天真的拷贝实现。

def copy(src, ddst, size):
for i in range(size):
dst[i] = src[i]

a = ti.field(ti.f32, 4)
b = ti.field(ti.f32, 4)
c = ti.field(ti.f32, 12)
d = ti.filed(ti.f32, 12)

copy(a, b, 4)
copy(c, d, 12)

有人会说之前的代码复用性差,我们可以把长度也作为一个参数传进来完成拷贝。

@ti.kernel
def copy(src: ti.template(), dst: ti.template(), size: ti.i32):
for i in range(size):
dst[i] = src[i]

a = ti.field(ti.f32, 4)
b = ti.field(ti.f32, 4)
c = ti.field(ti.f32, 12)
d = ti.field(ti.f32, 12)

copy(a, b, 4)
copy(c, d, 12)

我们可以再给力一些,拷贝这个工作相对独立,可以并行执行。我们自然而然请出了 @ti.kernel,最外层自动并行,但这里多出了 ti.template(),最最直观的解释是没这东西传不进去 ti.field,之前解释过 Python-scope 向 Taichi-scope 只能传最多八个标量,而且它们必须是强类型的。

ti.template() 支持任何类型,你可以丢 field进去,还可以丢别的东西进去。

ti.template()

实际上 ti.template() 都会被 Taichi 认为是一个模版,这个模版会在具体调用的时候才会被生成。当你传的参数是一个 vector 时它会专门生成一个只传 vector 的模版,或者说当你传的参数是 field 的时候它会生成一个只传 field 的模版。

Passing arguments using ti.template()

@ti.kernel
def foo(x: ti.template()):
print(x[0], x[1])

a = [42, 3.14]
foo(a) # NOT allowed


@ti.kernel
def foo(x: ti.template()):
print(x[0], x[1])

a = ti.Vector([42, 3.14])
foo(a) # 42, 3.14

它可以传递任何东西,这里的星号代表必须是 Taichi 认识的东西,例如 ti.f32ti.i32ti.f64……上图中 a 如果是普通的 python list 产生太极 kernel 时会直接报错。

vec = ti.Vector([0.0, 0.0])

@ti.kernel
def my_kernel(x: ti.template()):
vec2 = x
vec2[0] = 2.0
print(vec2) # [2.0, 0.0]

my_kernel(vec)


vec = ti.Vector([0.0, 0.0])

@ti.kernel
def my_kernel(x: ti.template()):
x[0] = 2.0 # bad assignment, x is in the Python scope
print(x)

my_kernel(vec)

这里的传参是设计为引用,当我们传入两个 ti.field 我们不想进行两次拷贝构造,想去直接修改这些 field 中的值。基于以上理由,ti.template() 会以引用传递的方式进行。

引用传递的限制在于在 Python-scope 中的一些数据没办法在 Taichi-scope 中做修改,比如 Python-scope 中有 a = 0,在 @ti.kernel 中无法修改 a = 0,因为 @ti.kernel 被编译的时候我们已经知道它会把 Python-scope 中的 data 全部编译为常量。

@ti.kernel
def copy(src: ti.template(), dst: ti.tempalte(), size: ti.i32):
for i in range(size):
dst[i] = src[i]

我们改不了 Python-scope 中的数据,但可以修改全局数据。我们知道一个 ti.field 是全局的,它既可以在 Python-scope 中被执行调用,也可以在 Taichi-scope 中被修改,所以我们可以尽情修改 dstsrc 中的所有成员。

@ti.func
def my_func(x: ti.template()):
x += 1 # This line will change the original value of x

@ti.kernel
def my_kernel():
x = 24
my_func(x)
print(x) # 25

my_kernel()

如果我在 @ti.func 中用 ti.template() 做修饰呢?上面的代码如果没有 ti.template(),该 my_func() 会执行一次拷贝,在这个词法作用域里做得改变不会影响到 @ti.kernel 里面。

Copy a Taichi field to another

@ti.kernel
def copy(src: ti.template(), dst: ti.template()):
for i in src:
dst[i] = src[i]

a = ti.field(ti.f32, 4)
b = ti.field(ti.f32, 4)
c = ti.field(ti.f32, 12)
d = ti.field(ti.f32, 12)

copy(a, b)
copy(c, d)

在熟悉 ti.template() 后我们可以把参数 size 甩掉了,因为传进两个 field 就直接可以用 Taichi struct-for 的形式进行循环调用,直接进行拷贝操作。

同理拷贝 ti.Vector.field 等等也可行,只要 shape 长得一样就ok。

如果 shape 不一样,在 Taichi struct-for 需要使用 for i, j in src: 这种形式专门处理 cd 的拷贝,这种情况下也许需要写好几个函数。

Dimension independent programming

@ti.kernel
def copy_1D(x: ti.template(), y: ti.template()):
for i in x:
y[i] = x[i]

@ti.kernel
def copy_2D(x: ti.template(), y: ti.template()):
for i, j in x:
y[i, j] = x[i, j]

@ti.kernel
def copy_3D(x: ti.template(), y: ti.template()):
for i, j, k in x:
y[i, j, k] = x[i, j, k]

在这种情况下你可能需要写好几个函数,它们之间唯一的区别可能只是外面的 ti.template() 传进来的是一维场、二维场、三维场,for 循环针对不同的维度进行调整。

@ti.kernel
def copy(x: ti.template(), y: ti.template()):
for I in ti.grouped(y):
# if y is 0D, then I = ti.Vector([])
# if y is 1D, then I = ti.vector([i])
# if y is 2D, then I = ti.Vector([i, j])
# if y is 3D, then I = ti.Vector([i, j, k])
x[I] = y[I]

可不可以再给力一些?当然可以,我们可以把这些函数 group 成同一个函数,I 作为 ti.grouped() 返回类型相同的 field

Metadata

import taichi as ti
ti.init(arch = ti.cpu, debug = True)

@ti.kernel
def copy(src: ti.template(), dst: ti.template()):
assert src.shape == dst.shape
for i in dst:
dst[i] = src[i]

a = ti.field(ti.f32, 4)
b = ti.field(ti.f32, 100)
copy(a, b)

我们在进行元编程的时候需要一些 Metadata(描述数据的数据),对于 field 调用:

  • ti.dtype
  • ti.shape

@ti.kernel
def foo():
matrix = ti.Matrix([[1, 2], [3, 4], [5, 6]])
print(matrix.n) # 3
print(matrix.m) # 2
vector = ti.Vector([7, 8, 9])
print(vector.n) # 3
print(vector.m) # 1

对于 Taichi 中一些复合类如矩阵、向量也会有 Metadata,主要描述它的行和列:

  • matrix.n
  • matrix.m
  • vector.n
  • vector.m

Use ti.template() with caution

@ti.kernel
def foo(x: ti.template()):
...

a = 1
foo(a)
foo(a)
foo(a) # foo is instantiated once


@ti.kernel
def foo(x: ti.template()):
...

a = 1
b = 2
c = 3
foo(a)
foo(b)
foo(c) # foo is instantiated three times

当我们生成 Taichi 的模板之后是如何生成对应的代码呢?它会在见到一个新的参数的时候重新为 foo() 实例化一套最后会执行的 kernel,比如上例中左边代码只被实例化一次,而右边会实例化三次。

Metaprogramming in Taichi

  • Unify the development of dimensionality-dependent code, such as 2D/3D physical simulations
  • Improve run-time performance by taking run-time costs to compile time

ti.static()

enable_projection = False

x = ti.field(ti.f32, shape = 10)

@ti.kernel
def static():
if ti.static(enable_projection): # No runtime overhead
x[0] = 1

如果有一些数字在编译的时候就已经知道的,我们可以用 ti.staic() 来修饰它进行一些代码加速。

@ti.kernel
def foo():
for i in ti.static(range(4)):
print(i)

# is equivalent to:
@ti.kernel
def foo():
print(0)
print(1)
print(2)
print(3)

它也能帮助我们做一些循环展开,上面的 for 比如不需要并行,我们可以使用 ti.static() 手动拆开。

# Here we declare a field contains 8 vectors. Each vector contains 3 elements.
x = ti.Vector.field(3, ti.f32, shape=(8))
@ti.kernel
def reset():
for i in x:
for j in ti.static(range(x.n)):
# The inner loop must be unrolled since j is an index for accessing a vector
x[i][j] = 0

Example

@ti.kernel
def computeForce(self, stars: ti.template()):
self.clearForce()
for i in range(self.n):
p = self.pos[i]

for j in range(self.n): if i != j:
diff = self.pos[j] - p
r = diff.norm(1e-2)
self.force[i] += G * self.Mass() * self.Mass() * diff / r**3
for j in range(stars.Number()):
diff = stars.Pos()[j] - p
r = diff.norm(1e-2)
self.force[i] += G * self.Mass() * stars.Mass() * diff / r**3

面向对象编程

If you want to build a car… (POP)

Object-oriented programming (OOP)

Object-oriented programming (OOP) is a programming paradigm based on the concept of “objects”.

Python OOP in a nutshell

class Wheel:
def _init_(self, radius, width, rolling_fric):
self.radius = radius
self.width = width
self.rolling_fric = rolling_fric

def Roll(self):
...

w1 = Wheel(5, 1, 0.1)
w2 = Wheel(5, 1, 0.1)
w3 = Wheel(6, 1.2, 0.15)
w4 = Wheel(6, 1.2, 0.15)

OOP 有以下两个好处:

  • 在轮子类设计时发动机不需要知道轮子是怎么造的
  • 进行类之间的继承,比如在轮胎上加个灯

Python OOP + Taichi DOP = ODOP

@ti.data_oriented
class TaichiWheel:
def _init_(self, radius, width, rolling_fric):
self.radius = radius
self.width = width
self.rolling_fric = rolling_fric
self.pos = ti.Vector.field(3, ti.f32, shape=4)

@ti.kernel
def Roll(self):
...

@ti.func
def foo(self):
...

Taichi 的类和 Python 差不太多,而 ODOP 则是在原有 Python 的 OOP 基础上加入 DOP。

@ti.data_oriented

image-20211102162822518

@ti.data_oriented
class CelestiablObject:
def _init_(self):
self.pos = ti.Vector(2, ti.f32, shape = N)
self.vel = ti.Vector(2, ti.f32, shape = N)
self.force = ti.Vector(2, ti.f32, shape = N)

Use @ti.kernel / @ti.func in your class

@ti.data_oriented
class CelestiablObject:
...

@ti.kernel
def update(self, h: ti.f32):
for i in self.vel:
self.vel[i] += h * self.force[i] / self.Mass()
self.pos[i] += h * self.vel[i]

@ti.data_oriented 的类中同样可以加入 @ti.kernel 加速类的成员函数计算。

PDOP: Use Python scope variables in Taichi scope with caution

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

d = 1

@ti.kernel
def foo():
print("d in Taichi scope =", d)

d += 1 # d = 2
foo() # d in Taichi scope = 2
d += 1 # d = 3
foo() # d in Taichi scope = 2
ODOP: Use Python scope members in Taichi scope with caution

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

@ti.data_oriented
class MyTaichiClass:
def _init_(self):
self.d = 1

def IncreaseD(self):
self.d += 1

@ti.kernel
def PrintD(self):
print("d in Taichi scope =", self.d)

a = MyTaichiClass() # d = 1
a.IncreaseD() # d = 2
a.PrintD() # print: d = 2
a.IncreaseD() # d = 3
a.PrintD() # print: d = 2

上面例子中 self.d 是 Python-scope 中的变量,当我在 @ti.kernel 中调用 a.IncreaseD()seld.d 就当成常量处理。

Encapsulation

Inheritance

Polymorphism