尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

Taichi Hello World 入门指南:用 Python 编写 GPU 并行计算与 Julia 分形

Taichi Hello World 入门指南:用 Python 编写 GPU 并行计算与 Julia 分形 Taichi Hello World 入门指南用 Python 编写 GPU 并行计算与 Julia 分形【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi导读本文以 Taichi 官方入门文档《Hello, World!》为核心骨架带领你从零写出并运行第一个 Taichi 程序——Julia 分形动画。你将掌握 Taichi 的三大核心概念ti.kernel与ti.func装饰器如何触发 JIT 即时编译、ti.field数据容器的定义与遍历、以及外层for循环的自动并行化机制同时理解后端arch的选择逻辑与 GUI 渲染流程。读完本文你将具备独立编写高性能并行计算程序的基础能力。Taichi嵌入 Python 的领域专用语言Taichi 是一种面向高性能并行计算的领域专用语言DSL内嵌于 Python 中。当编写计算密集型任务时只需遵循一套额外的规则并使用两个装饰器ti.func与ti.kernel即可让 Taichi 接管计算任务将装饰的函数通过**即时编译JIT**编译为机器码。此后对这些函数的调用将运行在多核 CPU 或 GPU 上相比原生 Python 代码可实现约 50x100x 的加速。除了 JIT 之外Taichi 还提供**预先编译AOT**系统可将代码导出为二进制/着色器文件脱离 Python 环境运行参见 Tutorial: Run Taichi programs in C application。前置条件与安装在开始之前请确认你的环境满足以下要求Python3.7 / 3.8 / 3.9 / 3.1064 位操作系统Windows、macOS、Linux64 位Taichi 以 PyPI 包的形式分发安装只需一条命令pip install taichiHello, worldJulia 分形程序一个基础的 Julia 分形Julia set动画是理解 Taichi 编程语言核心概念的最佳起点。将下面的代码保存为fractal.pyimport taichi as ti import taichi.math as tm ti.init(archti.gpu) n 320 pixels ti.field(dtypefloat, shape(n * 2, n)) ti.func def complex_sqr(z): # complex square of a 2D vector return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1]) ti.kernel def paint(t: float): for i, j in pixels: # Parallelized over all pixels c tm.vec2(-0.8, tm.cos(t) * 0.2) z tm.vec2(i / n - 1, j / n - 0.5) * 2 iterations 0 while z.norm() 20 and iterations 50: z complex_sqr(z) c iterations 1 pixels[i, j] 1 - iterations * 0.02 gui ti.GUI(Julia Set, res(n * 2, n)) i 0 while gui.running: paint(i * 0.03) gui.set_image(pixels) gui.show() i 1运行方式如果你不使用 IDE直接在终端进入脚本所在目录执行python3 fractal.py将filename替换为你的脚本名务必保留.py扩展名即可看到分形动画。该示例对应的官方实现位于 python/taichi/examples/simulation/fractal.py运行ti example -s fractal即可将其保存到当前工作目录两者逻辑一致可互为对照。下面我们逐段剖析这个看似简单的程序。导入 Taichi 与 math 模块import taichi as ti import taichi.math as tm前两行导入了 Taichi 及其math模块。math模块内置了小型向量与矩阵类型例如 2D 实向量vec2、3×3 实矩阵mat3等。更多用法参见 Math Module。初始化 Taichi选择后端ti.init(archti.gpu)ti.init()初始化 Taichi 运行环境并根据传入的可选参数定制运行时。其中arch参数指定执行编译后代码的后端backend可为ti.cpu或ti.gpu指定ti.gpu时Taichi 会按ti.cuda→ti.vulkan→ti.opengl/ti.Metal的顺序尝试使用 GPU 后端若 GPU 架构均不可用则回退使用 CPU 作为后端也可以直接指定具体后端例如archti.cuda。若指定的架构不可用Taichi 会直接抛出错误。从源码看ti.init()定义于 python/taichi/lang/misc.py其完整签名为init(archNone, default_fpNone, default_ipNone, _test_modeFalse, enable_fallbackTrue, require_versionNone, **kwargs)核心职责是设置整个程序使用的后端。所有受支持的架构在 taichi/inc/archs.inc.h 中集中声明包括x64、arm64、cuda、metal、opengl、dx11、dx12、opencl、amdgpu、vulkan、gles等其中部分处于 WIP开发中或 N/A不可用状态。除了archti.init()还接受其他常用参数来自 python/taichi/lang/misc.py 的 docstring 整理参数类型作用arch枚举计算后端通常为ti.cpu或ti.gpudefault_fp类型默认浮点类型如ti.f32、ti.f64default_ip类型默认整型类型require_version字符串版本约束字符串不满足则报错cpu_max_num_threadsintCPU 线程池使用的线程数debugbool开启调试模式此时 Taichi 会进行越界检查等额外工作print_irbool打印 Taichi 内核的 CHI IR中间表示offline_cachebool启用已编译内核的离线缓存默认开启可加速后续调用random_seedint随机数生成器种子默认 0此外TI_DEFAULT_FP、TI_DEFAULT_IP等环境变量也可用于配置默认浮点/整型类型且ti.init()的显式参数优先级更高见 python/taichi/lang/misc.py。更完整的全局设置说明见 Global Settings。定义 Taichi field数据容器n 320 pixels ti.field(dtypefloat, shape(n * 2, n))ti.field(dtype, shape)定义一个形状为shape、元素类型为dtype的 Taichi field。Field 是 Taichi 中最基础也最常用的数据结构可以类比 NumPy 的ndarray或 PyTorch 的tensor但更加灵活field 可以是空间稀疏的也可以方便地切换不同的数据布局。对本例而言只需知道pixels是一个稠密二维数组即可。从源码看ti.field()定义于 python/taichi/lang/impl.py支持如下参数摘自其 docstring参数说明dtype元素数据类型可为标量类型也可为向量/矩阵类型shapefield 形状可为int或tuple[int]order形状在内存中的排列顺序如ij、jinamefield 名称offsetfield 定义域的偏移量needs_grad是否参与反向模式自动微分需要伴生梯度 fieldneeds_dual是否参与前向模式自动微分需要对偶 field实现上ti.field(dtypefloat, shape(640, 320))等价于声明x ti.field(ti.f32)后通过ti.root.dense(ti.ij, shape(640, 320)).place(x)放置数据见 python/taichi/lang/impl.py 中的等价示例。向量/矩阵类型如ti.math.vec3的 dtype 会被转发到Vector.field/Matrix.field处理。内核kernel与 Taichi 函数functi.func def complex_sqr(z): # complex square of a 2D vector return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1]) ti.kernel def paint(t: float): for i, j in pixels: # Parallelized over all pixels c tm.vec2(-0.8, tm.cos(t) * 0.2) z tm.vec2(i / n - 1, j / n - 0.5) * 2 iterations 0 while z.norm() 20 and iterations 50: z complex_sqr(z) c iterations 1 pixels[i, j] 1 - iterations * 0.02上述代码定义了两个函数分别用ti.func与ti.kernel装饰前者称为Taichi 函数Taichi function后者称为内核kernel。二者都不会被 Python 解释器执行而是由 Taichi 的 JIT 编译器接管并部署到由ti.init()的arch参数决定的并行多核 CPU 或 GPU 上。Taichi 函数与内核的主要区别如下调用入口不同kernel 是 Taichi 接管执行的入口可在程序任意位置调用而 Taichi 函数只能在内核或其他 Taichi 函数内部调用。例如本例中 Taichi 函数complex_sqr是在内核paint内被调用的。类型标注要求不同kernel 的参数与返回值必须显式标注类型而 Taichi 函数不需要。例如内核paint的参数t: float标注了类型而 Taichi 函数complex_sqr的参数z没有。嵌套与递归限制Taichi 支持嵌套函数但不支持嵌套内核同时 Taichi 函数内部不支持递归调用。对于熟悉 CUDA 的读者Taichi 的ti.func相当于 CUDA 的__device__ti.kernel相当于 CUDA 的__global__。 对于熟悉 OpenGL 的读者ti.func可比作 GLSL 中的普通函数ti.kernel可比作 compute shader计算着色器。在本例中complex_sqr计算二维向量的复数平方实部z0² - z1²、虚部2·z0·z1而paint对每个像素迭代求解 Julia 集迭代公式z z² c迭代次数决定像素灰度值。并行 for 循环Taichi 高性能的关键ti.kernel def paint(t: float): for i, j in pixels: # Parallelized over all pixelsTaichi 高性能的关键在于高效的迭代方式——利用并行化循环处理数据。内核最外层作用域中的for循环会被自动并行化无需手动分配线程、回收资源或管理内存本例中for循环同时对i、j两个索引迭代各迭代可并发执行pixelsfield 被当作迭代器i、j分别为取值0到2*n-1、0到n-1的整数索引(i, j)组合按(0,0), (0,1), …, (0,n-1), (1,0), (1,1), …, (2*n-1, n-1)的顺序被并行遍历。需要特别注意的是嵌套在其他结构如if/else或其他循环内部的for循环不会被自动并行化而是按顺序串行执行ti.kernel def fill(): total 0 for i in range(10): # Parallelized for j in range(5): # Serialized in each parallel thread total i * j if total 10: for k in range(5): # Not parallelized because it is not at the outermost scope若想将最外层作用域的for循环改为串行可使用ti.loop_config(serializeTrue)详见 Serialize a specified parallel for loop。⚠️ 警告break语句不支持出现在并行化的循环中ti.kernel def foo(): for i in x: ... break # Error! ti.kernel def foo(): for i in x: for j in range(10): ... break # OK!内层串行循环中使用break是允许的。注意本例内核paint中使用的while循环位于并行for循环体内属于每线程内部的串行控制流因此不违反上述限制。使用 GUI 显示结果gui ti.GUI(Julia Set, res(n * 2, n)) # Sets the window title and the resolution i 0 while gui.running: paint(i * 0.03) gui.set_image(pixels) gui.show() i 1为了将结果渲染到屏幕上Taichi 提供了内置的 GUI System。gui.set_image()用于设置窗口内容gui.show()用于显示更新后的图像。Taichi 的 GUI 系统使用标准笛卡尔坐标系定义像素坐标原点位于屏幕左下角pixels中的(0, 0)元素映射到窗口左下角(639, 319)元素映射到窗口右上角。从实现角度看set_image接受numpy.ndarray、MatrixField、Field 或 Texture 等输入见 python/taichi/ui/canvas.py。官方 fractal 示例中采用了gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT)的循环方式来响应退出事件见 python/taichi/examples/simulation/fractal.py两种写法均可实现动画循环。关键要点回顾通过上述示例你已经掌握了 Taichi 最重要的特性Taichi 在指定后端上编译并执行 Taichi 函数与内核内核最外层作用域的for循环会被自动并行化Taichi 提供了名为field的灵活数据容器并可通过索引对其进行迭代。更多 Taichi 示例Julia 分形是 Taichi 收录的特色示例之一。要查看 Taichi Gallery 中更多精选示例运行ti gallery屏幕会弹出一个新的交互式窗口供你选择并运行示例。要查看完整示例列表运行ti example。以下是一些实用的命令行ti example -p fractal或ti example -P fractal在终端打印 fractal 示例的源代码ti example -s fractal将示例保存到当前工作目录。这些 CLI 命令由 python/taichi/_main.py 中的gallery与example方法实现其中example会以彩色表格列出全部可用示例并按名称运行或导出。示例源文件位于 python/taichi/examples 目录涵盖渲染如fractal3d_ggui.py、模拟如mpm88.py、game_of_life.py、算法如mgpcg.py等众多场景。支持的系统与后端下表概览了 Taichi 支持的操作系统及与之兼容的后端platformCPUCUDAOpenGLMetalVulkanWindows✅✅✅N/A✅Linux✅✅✅N/A✅macOS✅N/AN/A✅✅✅支持N/A不可用。注macOS 上没有 CUDA 与 OpenGL 后端对应的是 Apple Metal 后端该表与 taichi/inc/archs.inc.h 中声明的架构清单相互印证。对于机器上存在多个可用后端的情况ti.init(archti.gpu)会按 CUDA → Vulkan → OpenGL/Metal 的优先级自动选择。下一步现在你已经准备好编写自己的 Taichi 程序了。以下文档介绍了 Taichi 在典型应用场景中的使用方式Accelerate Python with TaichiConduct Physical SimulationAccelerate PyTorch with Taichi【免费下载链接】taichiProductive, portable, and performant GPU programming in Python.项目地址: https://gitcode.com/GitHub_Trending/ta/taichi创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表