Taichi 加速 Python 计算实战：从素数计数到动态规划的性能跃迁-北京尧图网络

Taichi 加速 Python 计算实战从素数计数到动态规划的性能跃迁【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi太极是一种内嵌于 Python 的领域特定语言DSL它允许你用几乎原生的 Python 语法编写计算密集型程序再由编译器将其 JIT 编译为 CPU 或 GPU 上的高性能原生指令从而让 Python 程序获得接近 C/C 乃至 CUDA 的性能。本文基于 accelerate_python.md 的核心内容通过素数计数与最长公共子序列LCS两个完整可运行的示例带你掌握ti.kernel、ti.func、ti.field、ti.loop_config等核心用法并给出对应的仓库源码佐证。读完本文你将能够独立把计算密集的 Python 循环改写成 Taichi 程序并在 CPU 与 GPU 后端之间一键切换。为什么需要加速Python 嵌套循环的性能瓶颈大型或深层嵌套的 Pythonfor循环几乎必然带来糟糕的运行时性能这是因为 Python 的解释执行与动态类型检查在每次迭代中都有巨额开销。Taichi 的解法很直接把循环体搬进ti.kernel装饰的 Python 函数中Taichi 编译器会将其编译到不同设备CPU 或 GPU上执行。从仓库源码看kernel_impl.py 中kernel()的 docstring 明确写道被ti.kernel标记的函数会由 TaichiJIT 编译成原生 CPU/GPU 指令例如一系列 CUDA kernel且顶层for循环会被自动并行化分配到 CPU 线程池或大规模并行的 GPU 上——这正是加速的根本来源。与之配套的ti.func见 kernel_impl.py则把函数标记为可在 Taichi 作用域内调用编译器同样会将其内联编译进调用它的 kernel 中。示例一统计素数count_primes1. 纯 Python 版本与基线性能下面的程序统计[1, n]范围内的素数个数涉及两层嵌套循环外层遍历2到n内层is_prime()遍历2到sqrt(n)检查整除性。Count the prime numbers in the range [1, n] # Checks if a positive integer is a prime number def is_prime(n: int): result True # Traverses the range between 2 and sqrt(n) # - Returns False if n can be divided by one of them; # - otherwise, returns True for k in range(2, int(n ** 0.5) 1): if n % k 0: result False break return result # Traverses the range between 2 and n # Counts the primes according to the return of is_prime() def count_primes(n: int) - int: count 0 for k in range(2, n): if is_prime(k): count 1 return count print(count_primes(1000000))将代码保存为count_prime.py在终端运行并计时time python count_primes.py屏幕上会同时打印素数个数与执行时间。原文档在参考机器上测得基线如下实际结果因机器而异78498 real 0m2.235s user 0m2.235s sys 0m0.000s2. 引入 Taichi三步完成改写第一步导入 Taichi 并初始化先用 CPU 后端import taichi as ti ti.init(archti.cpu)第二步分别用ti.func和ti.kernel装饰is_prime()与count_primes()ti.func def is_prime(n: int): result True for k in range(2, int(n ** 0.5) 1): if n % k 0: result False break return result ti.kernel def count_primes(n: int) - int: count 0 for k in range(2, n): if is_prime(k): count 1 return countTaichi 编译器会把ti.kernel装饰的 Python 代码编译到不同设备如 CPU、GPU上做高性能计算关于 kernel 与 function 这两个核心概念的详细说明参见 Kernels Functions。第三步重新运行time python count_primes.py。原文档在同一台参考机器上测得计算速度提升约6 倍2.235s / 0.363s78498 real 0m0.363s user 0m0.546s sys 0m0.179s3. 扩大规模与切换 GPU 后端将N扩大十倍至10,000,000再运行原文档测得纯 Python 约 55sTaichi 约 0.8s加速约70 倍将后端从 CPU 切换为 GPUti.init(archti.gpu)此时原文档测得Taichi 约 0.45s vs 纯 Python 55s加速进一步来到约120 倍。从源码层面看ti.init()定义在 misc.pyarch参数用于选择全局后端最终通过cfg.arch adaptive_arch_select(arch, enable_fallback)生效代码中还支持cpu_max_num_threadsCPU 线程池线程数、debug开启边界检查等调试行为、offline_cache编译缓存默认开启等常用配置。此外ti.init也遵循TI_ARCH环境变量例如TI_ARCHcuda便于在脚本外控制后端。示例二动态规划——最长公共子序列LCS动态规划的核心思想是用存储空间换取执行时间把中间结果缓存下来避免重复计算。下面实现计算两条序列最长公共子序列Longest Common Subsequence长度的经典 DP 算法。例如序列 a [0,1, 0, 2,4,3, 1,2, 1] 与 b [4,0,1,4, 5,3, 1,2] 的 LCS 为 [0, 1, 4, 3, 1, 2]长度为 6。1. 定义状态表用ti.field表示 DP 矩阵导入 NumPy 与 Taichi 并初始化import taichi as ti import numpy as np ti.init(archti.cpu)生成两条 15,000 长的随机整数序列取值范围[0, 100]N 15000 a_numpy np.random.randint(0, 100, N, dtypenp.int32) b_numpy np.random.randint(0, 100, N, dtypenp.int32)定义一个N×N的 Taichi fieldf用其[i, j]号元素表示序列 a 的前 i 个元素与序列 b 的前 j 个元素的 LCS 长度f ti.field(dtypeti.i32, shape(N 1, N 1))ti.field在仓库中的定义见 impl.pydtype声明元素类型这里是ti.i32shape声明形状f[i, j]的访问方式对用户透明地映射到后端内存布局上。2. DP 转移方程与ti.loop_config(serializeTrue)DP 的转移方程即对 fieldf的遍历f[i, j] ti.max(f[i - 1, j - 1] (a[i - 1] b[j - 1]), ti.max(f[i - 1, j], f[i, j - 1]))将其封装进 kernelcompute_lcs()接收两条序列作为参数ti.kernel def compute_lcs(a: ti.types.ndarray(), b: ti.types.ndarray()) - ti.i32: len_a, len_b a.shape[0], b.shape[0] ti.loop_config(serializeTrue) # Disable auto-parallelism in Taichi for i in range(1, len_a 1): for j in range(1, len_b 1): f[i, j] ti.max(f[i - 1, j - 1] (a[i - 1] b[j - 1]), ti.max(f[i - 1, j], f[i, j - 1])) return f[len_a, len_b]NumPy 数组在 Taichi 中按 ndarray 存储kernel 参数用ti.types.ndarray()标注即可直接传入务必设置ti.loop_config(serializeTrue)关闭 Taichi 的自动并行。这里每个循环迭代的计算都依赖前序迭代的结果f[i-1, j-1]、f[i-1, j]、f[i, j-1]迭代之间不应并行执行。ti.loop_config()的实现位于 misc.py它只作用于紧随其后的下一个循环支持的指令包括参数类型含义block_dimintGPU 上一个 block 的线程数serializebool是否让循环串行执行serializeTrue等价于parallelize1parallelizeintCPU 上使用的线程数block_dim_adaptivebool是否允许后端自适应设置 block_dim默认开启bit_vectorizebool是否对 quant_array 上的 struct for 启用位向量化3. 完整可运行程序将以下完整代码保存为lcs.pybenchmark False时使用文档中的小样例便于核对 LCS 长度 6True时进行大规模基准测试import taichi as ti import numpy as np ti.init(archti.cpu) benchmark True N 15000 f ti.field(dtypeti.i32, shape(N 1, N 1)) if benchmark: a_numpy np.random.randint(0, 100, N, dtypenp.int32) b_numpy np.random.randint(0, 100, N, dtypenp.int32) else: a_numpy np.array([0, 1, 0, 2, 4, 3, 1, 2, 1], dtypenp.int32) b_numpy np.array([4, 0, 1, 4, 5, 3, 1, 2], dtypenp.int32) ti.kernel def compute_lcs(a: ti.types.ndarray(), b: ti.types.ndarray()) - ti.i32: len_a, len_b a.shape[0], b.shape[0] ti.loop_config(serializeTrue) # Disable auto-parallelism in Taichi for i in range(1, len_a 1): for j in range(1, len_b 1): f[i, j] ti.max(f[i - 1, j - 1] (a[i - 1] b[j - 1]), ti.max(f[i - 1, j], f[i, j - 1])) return f[len_a, len_b] print(compute_lcs(a_numpy, b_numpy))运行并计时time python lcs.py系统会打印 LCS 长度与执行时间。原文档在参考机器上的输出示例2721 real 0m1.409s user 0m1.112s sys 0m0.549s原文档提供的同一算法在纯 Python 与 Taichi 下的对照结果是纯 Python 计算两条 15,000 长随机序列的 LCS 需约 476s而 Taichi 仅需约 0.9s加速最高可达约 500 倍。注意实际执行时间会因机器配置而异但可以预期你能观察到与文档相当量级的性能提升。本文引用的所有计时数据均来自原文档在特定参考机器上的实测仅用于说明相对提升量级不应视为普适基准。加速背后的原理小结两个示例揭示了 Taichi 加速 Python 的三条核心机制均能在仓库源码中得到印证JIT 编译到原生指令ti.kernel与ti.funckernel_impl.py把 Python 源码经 AST 转换后编译为 Taichi IR再降级为 CPU/GPU 原生代码消除了解释器逐行执行的开销自动并行化kernel 的顶层for循环被自动并行CPU 上分派到线程池、GPU 上映射为大规模并行线程对于存在数据依赖的循环则用ti.loop_config(serializeTrue)misc.py显式关闭并行以保证正确性数据布局抽象ti.fieldimpl.py与 ndarray 参数让数组数据在 Python 与原生代码间零拷贝传递并允许按后端自动选择最优内存布局。在此基础上你还可以进一步阅读 hello_world.md 了解第一个 Taichi 程序或参考仓库中 tests/python 下的大量测试用例例如 test_loops.py来验证各类循环与并行语义的行为。【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

Taichi 加速 Python 计算实战：从素数计数到动态规划的性能跃迁

相关推荐

Python面向对象编程进阶：多继承、多态与组合实战

Claude Code可视化工具：提升调试效率的动态执行追踪

Paperxie全功能板块硬核梳理｜一篇吃透所有核心模块，毕设全流程全覆盖

更多资讯

DPWM不连续脉宽调制详解：从SVPWM到三相逆变器Simulink仿真对比

MTProxy动态IP高可用架构设计：智能重连与DNS解析解决方案

Mermaid Radar Chart（radar-beta）完全指南：用多维雷达图做技能评估与性能对比

计及多能耦合的电气热能流计算方法与Matlab实现

Resume-Matcher 前端性能进阶：Next.js 15 服务端渲染优化三板斧（React.cache 去重 / 最小化客户端数据 / after() 非阻塞任务）

OpenMontage 中的 Remotion 过渡动画移植指南：把 @remotion/transitions 翻译为 HyperFrames 交叉淡化与 shader-transitions

今日更新

CopilotKit Vue Demo 与 React Demo 的 Parity 实现指南：路由、运行时端点与交互语义的全框架对齐

Harbor RBAC 实战：db_auth 模式下非管理员用户管理项目成员的完整验证与权限模型解析

STM32平衡车纯平衡程序从拆包到参数整定的完整调试指南

本周精选

超人会飞不算本事：系统稳定依赖清晰规则与边界设计

超人VS蜘蛛侠：拆解超级IP的影响力与传播方法论

基于CNN的调制信号识别：MATLAB实现时频图分类实战

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？