Taichi Metaprogramming and Object-Oriented Programming

Posted on In Views: Word count in article: 7.9k Reading time ≈ 13 mins.

引言

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

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

元编程

Meta

Metaprogramming

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

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

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

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

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

Metaprogramming in Taichi

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

Copy a Taichi field to another

1
2
3
4
5
6
7
8
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 的话会指向同一块内存，上面是一个比较天真的拷贝实现。

1
2
3
4
5
6
7
8
9
10
11
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)

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

1
2
3
4
5
6
7
8
9
10
11
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()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
@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.f32、ti.i32、ti.f64……上图中 a 如果是普通的 python list 产生太极 kernel 时会直接报错。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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 全部编译为常量。

1
2
3
4
@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 中被修改，所以我们可以尽情修改 dst 和 src 中的所有成员。

1
2
3
4
5
6
7
8
9
10
11
@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

1
2
3
4
5
6
7
8
9
10
11
12
@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: 这种形式专门处理 c 到 d 的拷贝，这种情况下也许需要写好几个函数。

Dimension independent programming

1
2
3
4
5
6
7
8
9
10
11
12
13
14
@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 循环针对不同的维度进行调整。

1
2
3
4
5
6
7
8
@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

1
2
3
4
5
6
7
8
9
10
11
12
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

1
2
3
4
5
6
7
8
@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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
@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()

1
2
3
4
5
6
7
8
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() 来修饰它进行一些代码加速。

1
2
3
4
5
6
7
8
9
10
11
12
@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() 手动拆开。

1
2
3
4
5
6
7
8
# 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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
@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

1
2
3
4
5
6
7
8
9
10
11
12
13
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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
@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

1
2
3
4
5
6
@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

1
2
3
4
5
6
7
8
9
@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

1
2
3
4
5
6
7
8
9
10
11
12
13
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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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