Taichi语言学习笔记-1

Taichi语言学习笔记-1

这个语言我在上大学的时候就听说过，以高性能著称，当时一个99行代码渲染冰雪奇缘的视频在b站上斩获了不小的播放量，那个时候我就想来尝试一下这个非常厉害的语言。不过到了今天，我才有充分的“理由”“不得不”学习这个语言，遂写下这篇文章，一方面促进自己学习，另一方面也请广大网友有理有据拷打我，让我借大伙之力进步更快。

这一篇博客主要想写Taichi与python之间的关系与区别，针对Taichi本身的各种内容之后会慢慢学习。

更改日志

2024-03-17 : 修复一点bug,核对了一些案例

What？

这门语言，或者说这个DSL，实际上是嵌入python中的：至少在一部分使用场景下，他的语法和python几乎没什么差别，甚至写都是写在一个文件下的：

import taichi as ti  
ti.init(arch=ti.cpu)  
  
@ti.func  
def inv_square(x): # A Taichi function  
	return 1.0 / (x * x)  
  
@ti.kernel  
def partial_sum(n: int) -> float: # A kernel  
	total = 0.0  
	for i in range(1, n + 1):  
		total += inv_square(n)  
	return total  
  
partial_sum(1000)

这是官方给的案例。可以看到这里taichi被作为一个包引进进来，做了一个简单的初始化，就开始进行了计算。一切看起来都是那么友好，而且和python看起来可以说是完全一样。

Why？

官方在文档中给出的说法是：将taichi内嵌入python，并使用其作为自己的前端，是有原因的：易学易用，不需要繁琐的编译什么的环节，可以借助极为丰富的python生态来加速开发，且python的AST非常适合拿来分析以供下一步运算，这些都使得python成为了当仁不让地首选。

问题当然是有的：那就是python还是慢的，纯python脚本开发出来的taichi应用初始化的时候会有点费劲，需要一些时间。

那到底这玩意是怎么让代码变快的呢？

HOW？

简单来说这玩意实际上跑了一个独立于python runtime的另外一个runtime，taichi在这之中下足了功夫：只要标注了ti.kernel代码块，就会被用一种奇妙的方式塞到taichi runtime里，在那里代码得到了非常良好的优化，这使得最终代码运行效率提升。

所以叫成taichi lang的确是有道理的：taichi runtime对于喂进去的python代码是有比较多要求的，至少比python要求的多，所以涉及到taichi的那一部分就必须用一套新的（虽然其实几乎一样）的规则来去写代码。

那么自然而然地，代码就必须分成两块：taichi scope和python scope。用上述代码做例子的话：

#########python scope#############
import taichi as ti  
ti.init(arch=ti.cpu)  
#########python scope END#########

###########taichi scope###########
@ti.func  
def inv_square(x): # A Taichi function  
	return 1.0 / (x * x)  
###########taichi scope END#######

###########taichi scope###########
@ti.kernel  
def partial_sum(n: int) -> float: # A kernel  
	total = 0.0  
	for i in range(1, n + 1):  
		total += inv_square(n)  
	return total  
###########taichi scope END#######

#########python scope#############
partial_sum(1000)
#########python scope END#########

关于taichi的代码被划分成一部分，剩下的都是python scope。

既然是两个scope，其中的概念被分成两个也是合情合理的。taichi scope中的function实际上和python中的function是分开的。

有时候我更喜欢这样理解这一整套系统：python作为一个胶水语言，他要做的事情就是安排好taichi scope内的代码按照什么顺序、以什么样子运行，以及其他的脏活（读数据）累活（处理数据，输出output），taichi代码就好好做好运算就好。

taichi kernel就是python代码需要管理的运算单元。一个kernel定义了一个运算任务。kernel可以在python scope代码中的任何部分出现，但唯独不能出现在另一个taichi kernel或taichi func里。

@ti.kernel
def hello():
	print("hello.") #可以写原生python代码的，就像这样一样

@ti.kernel
def what():
	hello() # -> pylint不会报错，但运行会报错

if __name__ == "__main__":
	what() # -> correct

上述代码自然是有问题的，报错信息非常明确：

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.11.8
[Taichi] Starting on arch=x64
Traceback (most recent call last):
  File "/home/leviathan/everything/project/taichi-test/test.py", line 17, in <module>
    what() # -> correct
    ^^^^^^
  File "/home/leviathan/.local/lib/python3.11/site-packages/taichi/lang/kernel_impl.py", line 1107, in wrapped
    raise type(e)("\n" + str(e)) from None
taichi.lang.exception.TaichiSyntaxError:
File "/home/leviathan/everything/project/taichi-test/test.py", line 13, in what:
    hello() # -> pylint不会报错，但运行会报错
    ^^^^^^^

Kernels cannot call other kernels. I.e., nested kernels are not allowed. Please check if you have direct/indirect invocation of kernels within kernels. Note that some methods provided by the Taichi standard library may invoke kernels, and please move their invocations to Python-scope.

那么很好理解，ti.func就是针对kernel的function。python scope自然是不可以调用这些func，其需要放在kernel或者其他taichi function里被调用。

不过有个坏消息：ti.func是不允许运行时递归的，因为ti实际上是这样处理函数调用的：直接内联，即将所有func最后合成出一个超级大的func，如果做出运行时递归的话很有可能会导致出现无限长的展开导致出问题。

值得注意的一点是，只要你不调用，写错的代码是不会报错的。

接下来是几个错误例子：

@ti.func
def error():
  error()
# 运行后会告诉你栈溢出。

count = 0
@ti.func
def out_count():
  if count <= 2:
    count += 1
    out_count()
# 实际上count是编译时常量，所以也会爆炸

@ti.func
def error_yield():
  count = 0
  if count == 2:
    count += 1
    yield count
    error_yield()
# taichi貌似还没支持yield，写出来运行后会告诉你yield不被支持。

类型

方法声明

taichi实际上还是比较在意类型的。假设如下代码：

@ti.kernel
def test(a):
  print(a)

if __name__ == "__main__":
  test(12)

@ti.kernel
def test(a: int):
  return a

if __name__ == "__main__":
  test(12)

上述两段代码是无法运行的。Taichi对于kernel的声明并不像Python那样那么的随便，其args内所有的参数都要明确的给出类型，且如果存在返回值，那么也要说明返回类型。

对于@ti.func就没那么多规矩了。

数据类型

Taichi和Python有很多细微上不一样的地方。对于类型而言，Taichi是一个静态类型语言，且一个变量的类型实际上在Taichi的编译期间就已经确定好了。官网例子举得就相当不错：

@ti.kernel  
def test():  
	x = 1 # x is the integer 1  
	x = 3.14 # x is an integer, so the value 3.14 is cast to 3 and x takes the value 3  
	x = ti.Vector([1, 1]) # Error!

这实际上就是静态类型的特点。第一次出现实际上就相当于完成了声明与初始化，之后就是基于这一类型来做事，所以会有上述的执行情况。

至此，基本上就讲述清楚两者之间的关系与区别了。