用 Taichi 加速 PyTorch数据预处理与自定义 ML 算子的高性能实践【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi 与 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 实现砖墙paddingPadding 是机器学习中常用的数据预处理手段例如防止卷积操作改变输入图像尺寸。但没有任何 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), axis2) vv torch.stack((v, v), axis2) 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(devicedevice)这种矩阵化改写存在两个突出问题不直观坐标映射、翻转、堆叠、索引收集等逻辑层层嵌套可读性差、难以维护显存开销大过程中需要保存大量中间结果显存较小的旧 GPU 甚至无法负担。用 Taichi 实现 paddingTaichi 版本的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、layoutAOS/SOA、needs_grad等参数例如ti.types.ndarray(ndim3)。此外Taichi 还提供to_torch()/from_torch()双向转换接口见 python/taichi/lang/field.py字段与 PyTorch 张量之间的数据交换同样便捷。性能对比58 个 CUDA 内核 vs 1 个 CUDA 内核Kernel 函数平均耗时 (ms)启动的 CUDA 内核数torch_pad()30.39258ti_pad()0.2671测试环境GPU RTX3090PyTorch v1.12.1Taichi 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 版深度卷积算子。该深度卷积算子的功能可概括为三步遍历两个输入张量w和k将w、k对应元素的乘积累加到s将s写入输出张量out。四种实现的可读性与性能对比实现方式可读性性能Python极佳最慢PyTorch差慢CUDA差快Taichi极佳与 CUDA 相当甚至更优Python 参考实现最直观、易于理解但慢到性能曲线几乎画不出来def run_formula_very_slow(w, k, B, C, T, eps): out torch.empty((B, C, T), devicecpu) for b in range(B): for c in range(C): for t in range(T): s eps for u in range(t-T1, t1): s w[c][0][(T-1)-(t-u)] * k[b][c][uT-1] out[b][c][t] s return outPyTorch 实现要把上面的 Python 参考代码翻译成下面这一行必须对F.conv1d、nn.ZeroPad2d等算子的底层逻辑了如指掌out eps F.conv1d(nn.ZeroPad2d((T-1, 0, 0, 0))(k), w.unsqueeze(1), groupsC)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(ndim3), w: ti.types.ndarray(ndim3), k: ti.types.ndarray(ndim3), eps: ti.f32): for b, c, t in out: s eps for u in range(t-T1, t1): s w[c, 0, (T-1)-(t-u)] * k[b, c, uT-1] out[b, c, t] s注意这里的for b, c, t in out直接遍历三维 ndarray 的索引配合ti.types.ndarray(ndim3)的类型标注明确声明了张量维度eps以ti.f32标量参数传入。在相同优化技巧如循环融合下Taichi 在多数场景下性能与 CUDA 相当部分场景甚至更优。加速原理Taichi 为什么能做到结合上文两个示例与仓库源码Taichi 相对 PyTorch/CUDA 的加速与开发效率优势主要来自三点顶层 for 自动并行化ti.ndrange将多维索引空间压平为单个顶层循环由 Taichi 编译器自动映射到并行线程见 python/taichi/lang/_ndrange.py开发者无需手工管理 block/thread单 CUDA 内核编译元素级逻辑全部编译进一个内核大幅减少 GPU 启动开销这正是 padding 示例中 58 个内核降到 1 个、加速超 100 倍的根本原因零拷贝张量互操作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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
用 Taichi 加速 PyTorch：数据预处理与自定义 ML 算子的高性能实践
用 Taichi 加速 PyTorch数据预处理与自定义 ML 算子的高性能实践【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi 与 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 实现砖墙paddingPadding 是机器学习中常用的数据预处理手段例如防止卷积操作改变输入图像尺寸。但没有任何 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), axis2) vv torch.stack((v, v), axis2) 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(devicedevice)这种矩阵化改写存在两个突出问题不直观坐标映射、翻转、堆叠、索引收集等逻辑层层嵌套可读性差、难以维护显存开销大过程中需要保存大量中间结果显存较小的旧 GPU 甚至无法负担。用 Taichi 实现 paddingTaichi 版本的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、layoutAOS/SOA、needs_grad等参数例如ti.types.ndarray(ndim3)。此外Taichi 还提供to_torch()/from_torch()双向转换接口见 python/taichi/lang/field.py字段与 PyTorch 张量之间的数据交换同样便捷。性能对比58 个 CUDA 内核 vs 1 个 CUDA 内核Kernel 函数平均耗时 (ms)启动的 CUDA 内核数torch_pad()30.39258ti_pad()0.2671测试环境GPU RTX3090PyTorch v1.12.1Taichi 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 版深度卷积算子。该深度卷积算子的功能可概括为三步遍历两个输入张量w和k将w、k对应元素的乘积累加到s将s写入输出张量out。四种实现的可读性与性能对比实现方式可读性性能Python极佳最慢PyTorch差慢CUDA差快Taichi极佳与 CUDA 相当甚至更优Python 参考实现最直观、易于理解但慢到性能曲线几乎画不出来def run_formula_very_slow(w, k, B, C, T, eps): out torch.empty((B, C, T), devicecpu) for b in range(B): for c in range(C): for t in range(T): s eps for u in range(t-T1, t1): s w[c][0][(T-1)-(t-u)] * k[b][c][uT-1] out[b][c][t] s return outPyTorch 实现要把上面的 Python 参考代码翻译成下面这一行必须对F.conv1d、nn.ZeroPad2d等算子的底层逻辑了如指掌out eps F.conv1d(nn.ZeroPad2d((T-1, 0, 0, 0))(k), w.unsqueeze(1), groupsC)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(ndim3), w: ti.types.ndarray(ndim3), k: ti.types.ndarray(ndim3), eps: ti.f32): for b, c, t in out: s eps for u in range(t-T1, t1): s w[c, 0, (T-1)-(t-u)] * k[b, c, uT-1] out[b, c, t] s注意这里的for b, c, t in out直接遍历三维 ndarray 的索引配合ti.types.ndarray(ndim3)的类型标注明确声明了张量维度eps以ti.f32标量参数传入。在相同优化技巧如循环融合下Taichi 在多数场景下性能与 CUDA 相当部分场景甚至更优。加速原理Taichi 为什么能做到结合上文两个示例与仓库源码Taichi 相对 PyTorch/CUDA 的加速与开发效率优势主要来自三点顶层 for 自动并行化ti.ndrange将多维索引空间压平为单个顶层循环由 Taichi 编译器自动映射到并行线程见 python/taichi/lang/_ndrange.py开发者无需手工管理 block/thread单 CUDA 内核编译元素级逻辑全部编译进一个内核大幅减少 GPU 启动开销这正是 padding 示例中 58 个内核降到 1 个、加速超 100 倍的根本原因零拷贝张量互操作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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
