企业官网建设流程全解析

用 Taichi 加速 PyTorch：数据预处理与自定义 ML 算子的高性能实践

【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi

Taichi 与 PyTorch 的应用场景不同，却可以形成互补：PyTorch 将细节抽象为类似 LEGO 积木的张量级算子，让开发者专注构建机器学习模型；Taichi 则提供对并行化的精细控制，支持更"细粒度"（元素级）的操作，赋予开发者更大的灵活性。本文基于仓库文档 docs/lang/articles/get-started/accelerate_pytorch.md，通过"数据预处理（padding）"与"自定义高性能 ML 算子（RWKV 深度卷积）"两个完整示例，展示如何用 Taichi kernel 实现高性能算子、直接消费 PyTorch 张量，并对比 Python、PyTorch、CUDA 与 Taichi 四种写法的可读性与性能差异。读完本文，你将掌握用 Taichi 替代"低效 Python 循环"与"C++/CUDA 算子扩展"两条传统路径的实战方案。

为什么需要 "Taichi + PyTorch" 组合

PyTorch 擅长承载机器学习中绝大部分的张量计算，但存在两类短板：

• 缺乏特定算子的原生支持：对于按自定义图案进行 padding、一维深度卷积这类"小众"算子，PyTorch 没有专门的原生实现；
• 运行时性能不理想：将元素级逻辑强行改写为张量级矩阵运算，会产生大量中间结果与内核启动开销。

面对这两类场景，传统上有两条路可走：用 Python 或 PyTorch 逐元素迭代（效率极低，可能拖垮整个训练流程）；或编写 C++/CUDA 算子并通过 Python 自定义算子扩展接入 PyTorch（需要大量底层硬件架构知识，上手周期长）。Taichi 提供了第三条路——用接近 Python 的语法编写自动并行化、性能可媲美 CUDA 的 kernel，并让 kernel 直接接收 PyTorch 张量。

示例一：数据预处理——用 Taichi 实现"砖墙"padding

Padding 是机器学习中常用的数据预处理手段，例如防止卷积操作改变输入图像尺寸。但没有任何 PyTorch 算子专门针对"特定自定义图案"的 padding。本节以"砖墙"拼接为例：先创建一块"砖"并用渐变色填充，再以固定偏移水平重复拼接，形成交错布局。

用 PyTorch 实现 padding 的痛点

下面的torch_pad()将 padding 过程改写为一系列原生 PyTorch 矩阵运算以提升效率：

def torch_pad(arr, tile, y): # image_pixel_to_coord arr[:, :, 0] = image_height - 1 + ph - arr[:, :, 0] arr[:, :, 1] -= pw arr1 = torch.flip(arr, (2, )) # map_coord v = torch.floor(arr1[:, :, 1] / tile_height).to(torch.int) u = torch.floor((arr1[:, :, 0] - v * shift_y[0]) / tile_width).to(torch.int) uu = torch.stack((u, u), axis=2) vv = torch.stack((v, v), axis=2) arr2 = arr1 - uu * shift_x - vv * shift_y # coord_to_tile_pixel arr2[:, :, 1] = tile_height - 1 - arr2[:, :, 1] table = torch.flip(arr2, (2, )) table = table.view(-1, 2).to(torch.float) inds = table.mv(y) gathered = torch.index_select(tile.view(-1), 0, inds.to(torch.long)) return gathered with Timer(): gathered = torch_pad(coords, tile, y) torch.cuda.synchronize(device=device)

这种"矩阵化"改写存在两个突出问题：

1. 不直观：坐标映射、翻转、堆叠、索引收集等逻辑层层嵌套，可读性差、难以维护；
2. 显存开销大：过程中需要保存大量中间结果，显存较小的旧 GPU 甚至无法负担。

用 Taichi 实现 padding

Taichi 版本的ti_pad()思路更直接：遍历输出图像的每个像素，反推其在"砖"中的对应位置，再将该位置的 RGB 颜色填入输出。顶层 for 循环由 Taichi 自动并行化，矩阵运算写法也远比张量拼接直观：

@ti.kernel def ti_pad(image_pixels: ti.types.ndarray(), tile: ti.types.ndarray()): for row, col in ti.ndrange(image_height, image_width): # image_pixel_to_coord x1, y1 = ti.math.ivec2(col - pw, image_height - 1 - row + ph) # map_coord v: ti.i32 = ti.floor(y1 / tile_height) u: ti.i32 = ti.floor((x1 - v * shift_y[0]) / tile_width) x2, y2 = ti.math.ivec2(x1 - u * shift_x[0] - v * shift_y[0], y1 - u * shift_x[1] - v * shift_y[1]) # coord_to_tile_pixel x, y = ti.math.ivec2(tile_height - 1 - y2, x2) image_pixels[row, col] = tile[x, y] with Timer(): ti_pad(image_pixels, tile) ti.sync()

这段代码中的image_height、image_width、tile_height、tile_width、pw、ph、shift_x、shift_y均为 padding 布局相关的全局常量（图像尺寸、砖块尺寸、偏移量等），在完整示例中按需定义即可。

源码佐证一：顶层 for 与ti.ndrange的并行语义

ti_pad()中for row, col in ti.ndrange(image_height, image_width)是关键——在 Taichi 中只有顶层 for 循环会被自动并行化。从 python/taichi/lang/_ndrange.py 的文档注释可以看到，ti.ndrange返回多维整数索引的笛卡尔积迭代器，文档明确建议：需要并行遍历张量时，不要写嵌套 for 循环，而应使用ndrange将全部条目放进一个顶层循环，例如for row, col, channel in ti.ndrange(image_height, image_width, channels)。每个参数可以是整数n（等价range(0, n)）或(start, end)二元组。

源码佐证二：PyTorch 张量作为"外部数组"零拷贝传入

ti_pad(image_pixels, tile)直接接收 PyTorch 张量，关键在于ti.types.ndarray()类型标注。Taichi 的 kernel 参数系统对 ndarray/外部数组做了专门处理：在 python/taichi/lang/kernel_impl.py 的set_arg_ext_array()中，numpy 数组或 PyTorch 张量通过v.ctypes.data直接取底层内存指针，以set_arg_external_array_with_shape()传入启动上下文，复用 PyTorch 已分配的内存，不产生两框架之间的数据搬运开销。类型定义位于 python/taichi/types/ndarray_type.py，支持ndim、layout（AOS/SOA）、needs_grad等参数，例如ti.types.ndarray(ndim=3)。

此外，Taichi 还提供to_torch()/from_torch()双向转换接口（见 python/taichi/lang/field.py），字段与 PyTorch 张量之间的数据交换同样便捷。

性能对比：58 个 CUDA 内核 vs 1 个 CUDA 内核

• 测试环境：GPU RTX3090；PyTorch v1.12.1，Taichi v1.1.0；
• 实际加速比可能因具体实现与 GPU 配置而异。

如上表所示，torch_pad()需要启动58 个 CUDA 内核，而 Taichi 将全部计算编译进1 个 CUDA 内核（本例加速超过 100 倍）。CUDA 内核数量越少，GPU 启动开销越低；同时 Taichi 版本省去了大量冗余内存操作。GPU 启动开销与冗余内存操作，正是此类场景中潜在的优化与加速来源。

示例二：自定义 ML 算子——RWKV 模型中的一维深度卷积

机器学习研究者通常花大量时间设计模型架构，但新设计或自定义的算子往往得不到 PyTorch 的良好支持，只能去啃 CUDA 做调优。原文档引用了 RWKV-CUDA 项目作为背景：作者用类似一维深度卷积的自定义算子开发 RWKV 语言模型，该算子计算量不大却因缺少 PyTorch 原生支持而运行缓慢，作者用循环融合、Shared Memory 等 CUDA 优化技巧实现后，性能比纯 PyTorch 提升约 20 倍。参照这份 CUDA 代码，原文档作者用同样的优化思路，在 RWKV 模型中定制了 Taichi 版深度卷积算子。

该深度卷积算子的功能可概括为三步：

1. 遍历两个输入张量w和k；
2. 将w、k对应元素的乘积累加到s；
3. 将s写入输出张量out。

四种实现的可读性与性能对比

Python 参考实现

最直观、易于理解，但慢到性能曲线几乎画不出来：

def run_formula_very_slow(w, k, B, C, T, eps): out = torch.empty((B, C, T), device='cpu') for b in range(B): for c in range(C): for t in range(T): s = eps for u in range(t-T+1, t+1): s += w[c][0][(T-1)-(t-u)] * k[b][c][u+T-1] out[b][c][t] = s return out

PyTorch 实现

要把上面的 Python 参考代码"翻译"成下面这一行，必须对F.conv1d、nn.ZeroPad2d等算子的底层逻辑了如指掌：

out = eps + F.conv1d(nn.ZeroPad2d((T-1, 0, 0, 0))(k), w.unsqueeze(1), groups=C)

CUDA 实现

CUDA 参考代码可读性差：最外层循环由线程并行度隐式定义，索引计算复杂，每个元素在矩阵中的位置难以一眼看清；实现更复杂的算法时极易出错。此外，运行 CUDA 代码还需要搭建编译环境，若预编译为动态链接库，还要处理环境配置与 Python API 封装等琐事：

__global__ void kernel_forward(const float* w, const float* k, float* x, const float eps, const int B, const int C, const int T) { const int i = blockIdx.y; const int t = threadIdx.x; float s = eps; const float* www = w + (i % C) * T + (T - 1) - t; const float* kk = k + i * T; for (int u = 0; u <= t; u++){ s += www[u] * kk[u]; } x[i * T + t] = s; }

Taichi 实现

Taichi 版本与 Python 参考实现几乎逐行对应，无需关心并行化、指针偏移等底层细节即可获得与 CUDA 相当的性能：

@ti.kernel def taichi_forward_v0( out: ti.types.ndarray(ndim=3), w: ti.types.ndarray(ndim=3), k: ti.types.ndarray(ndim=3), eps: ti.f32): for b, c, t in out: s = eps for u in range(t-T+1, t+1): s += w[c, 0, (T-1)-(t-u)] * k[b, c, u+T-1] out[b, c, t] = s

注意这里的for b, c, t in out直接遍历三维 ndarray 的索引，配合ti.types.ndarray(ndim=3)的类型标注明确声明了张量维度；eps以ti.f32标量参数传入。在相同优化技巧（如循环融合）下，Taichi 在多数场景下性能与 CUDA 相当，部分场景甚至更优。

加速原理：Taichi 为什么能做到

结合上文两个示例与仓库源码，Taichi 相对 PyTorch/CUDA 的加速与开发效率优势主要来自三点：

1. 顶层 for 自动并行化：ti.ndrange将多维索引空间压平为单个顶层循环，由 Taichi 编译器自动映射到并行线程（见 python/taichi/lang/_ndrange.py），开发者无需手工管理 block/thread；
2. 单 CUDA 内核编译：元素级逻辑全部编译进一个内核，大幅减少 GPU 启动开销，这正是 padding 示例中 58 个内核降到 1 个、加速超 100 倍的根本原因；
3. 零拷贝张量互操作：PyTorch 张量经外部数组机制直接共享内存（见 python/taichi/lang/kernel_impl.py），无需在框架间搬运数据。

总结

PyTorch 能高效处理机器学习中的绝大部分计算任务，但在算子原生支持不足、运行性能不理想等场景存在短板。作为一种嵌入 Python 的高性能编程语言，Taichi 具备：

• 易读性：kernel 写法贴近 Python 参考实现，学习成本低；
• 优化的内存占用：直接复用 PyTorch 张量内存，避免中间结果与数据搬运；
• 媲美 CUDA 的运行时性能：在 RWKV 深度卷积示例中，Taichi 性能与 CUDA 相当甚至更优；
• 良好的可移植性：便于社区间共享、复现代码。

本文的两个示例（数据预处理 padding 与自定义深度卷积算子）展示了 Taichi 与 PyTorch 如何互补，解决真实世界的高性能编程问题：预处理等细粒度元素操作交给 Taichi，模型构建等张量级抽象保留在 PyTorch，二者各取所长。

【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi