Taichi高级编程：元编程技术详解-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01138/article/details/148360823

Taichi高级编程：元编程技术详解

taichi Productive & portable high-performance programming in Python. 项目地址: https://gitcode.com/gh_mirrors/ta/taichi

前言

在Taichi编程语言中，元编程是一项强大的功能，它允许开发者在编译时而非运行时执行某些操作。本文将深入探讨Taichi中的元编程技术，包括模板编程、维度无关编程、字段元数据访问以及编译时计算等核心概念。

什么是元编程？

元编程是指编写能够生成或操作其他程序的程序。在Taichi中，元编程主要体现为：

延迟实例化：Taichi内核是延迟实例化的模板内核
编译时计算：将部分计算从运行时转移到编译时
代码复用：通过模板实现不同场景下的代码复用

模板编程

模板编程是元编程的重要形式，它允许我们编写通用的、可重用的代码。

基本用法

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

在这个例子中，ti.template()作为参数类型提示，表示可以接受任何Taichi字段或Python对象。这使得同一个内核可以用于不同形状的字段：

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_1D(a, b)  # 适用于4元素字段
copy_1D(c, d)  # 同样适用于12元素字段

注意事项

非Taichi对象的模板参数不能在Taichi内核内重新赋值
模板参数在编译后会内联到生成的内核中

维度无关编程

Taichi提供了ti.grouped语法，可以将循环索引分组为向量，实现维度无关的编程。

传统方式的问题

传统上，我们需要为不同维度的数据编写不同的内核：

# 1D版本
@ti.kernel
def copy_1D(x: ti.template(), y: ti.template()):
    for i in x:
        y[i] = x[i]

# 2D版本        
@ti.kernel
def copy_2d(x: ti.template(), y: ti.template()):
    for i, j in x:
        y[i, j] = x[i, j]

使用ti.grouped统一处理

@ti.kernel
def copy(x: ti.template(), y: ti.template()):
    for I in ti.grouped(y):
        x[I] = y[I]

ti.grouped会根据y的维度自动调整I的形式：

0D: I = ti.Vector([]) (相当于None)
1D: I = ti.Vector([i])
2D: I = ti.Vector([i, j])
3D: I = ti.Vector([i, j, k])

字段元数据访问

Taichi字段有两个重要元数据可以访问：

数据类型: 通过field.dtype访问
形状: 通过field.shape访问

Python作用域访问

x = ti.field(dtype=ti.f32, shape=(3, 3))
print("字段维度:", x.shape)  # 输出 (3, 3)
print("数据类型:", x.dtype)  # 输出 float32

Taichi作用域访问

@ti.kernel
def print_metadata(x: ti.template()):
    print("字段维度:", len(x.shape))
    for i in ti.static(range(len(x.shape))):
        print("维度", i, "大小:", x.shape[i])
    ti.static_print("数据类型:", x.dtype)

注意：对于稀疏字段，返回的是完整域的形状。

矩阵和向量元数据

对于矩阵和向量，可以访问以下元数据：

matrix.n: 行数
matrix.m: 列数
向量被视为单列矩阵，vector.n是元素个数，vector.m总是1

@ti.kernel
def matrix_info():
    mat = ti.Matrix([[1,2], [3,4], [5,6]])
    print(mat.n)  # 3行
    print(mat.m)  # 2列
    
    vec = ti.Vector([7,8,9])
    print(vec.n)  # 3个元素
    print(vec.m)  # 总是1

编译时计算

编译时计算可以将部分计算从运行时转移到编译时，提高运行时性能。

静态作用域

ti.static函数提示编译器在编译时评估其参数，参数的作用域称为静态作用域。

编译时分支

使用ti.static实现编译时分支，类似于C++17的if constexpr：

enable_projection = True

@ti.kernel
def static_example():
    if ti.static(enable_projection):
        x[0] = 1  # 编译时决定是否包含这段代码

注意：static if的两个分支之一会在编译后被丢弃。

循环展开

使用ti.static强制循环展开：

@ti.kernel
def unrolled_loop():
    for i in ti.static(range(4)):
        print(i)
    
    # 等价于：
    print(0)
    print(1)
    print(2)
    print(3)

历史说明：在v1.4.0之前，访问矩阵/向量元素的索引必须是编译时常量，因此涉及索引的循环必须展开。v1.4.0之后，索引可以是运行时变量，但展开仍有助于减少运行时开销。

编译时递归函数

编译时递归函数是指可以在编译时递归内联的函数，递归条件在编译时评估。

@ti.func
def sum_from_one_to(n: ti.template()) -> ti.i32:
    ret = 0
    if ti.static(n > 0):
        ret = n + sum_from_one_to(n - 1)
    return ret

@ti.kernel
def test():
    print(sum_from_one_to(10))  # 输出55

警告：递归深度过大时不建议使用编译时递归，因为会导致编译时代码膨胀，增加编译时间。