跳转至主要内容
Version: develop

在 Android 上使用 Ndarray 运行 Taichi 程序

Taichi 的 JIT (即时编译) 模块根据指定的后端架构(即 ti.init()arch 指定的后端)将 Taichi kernel 编译为计算着色器,并在 Taichi 的 JIT 运行时执行这些着色器。 而 AOT(预先编译)模块则构建并保存必要的计算着色器,以便你可以在没有 Python 运行环境的情况下在自己的运行时加载和执行这些着色器。

本篇教程以天体轨道模拟为例,讲解在 Android 上用 Ndarray 运行一个 Taichi 程序的过程。

Taichi 的 AOT 模块 作为一个正在开发的概念验证,今后可能会发生变化。

Ndarray 定义

Taichi 提供一种叫做 Ndarray 的数据容器。 Ndarray 是由相同类型和大小的元素组成的多维容器;Ndarray 中的元素都是标量或张量。

Ndarray 的形状

Ndarray 的形状决定了 Ndarray 的布局;元素的形状决定了元素的布局。 例如:

  • 一个形状为 [2, 1024] 且其中元素形状为 [] 的 Ndarry,是一个由 2×1,024=2,048 个标量组成的数组。
  • 一个形状为 [128, 128] 且其中元素形状为 [2, 4] 的 Ndarry,是一个由 128 x 128 = 16,384 个 2 x 4 矩阵组成的数组。

Ndarray 的维度

这里的维度指 Ndarray 的维度数量。 例如:

  • 一个形状为 [1, 2, 3] 的 Ndarray 是三维的。
  • 一个形状为 [500] 的 Ndarray 是一维的。

Ndarray 的优势

每个 Ndarray 都有一个固定的维度,但是你可以根据维度灵活改变形状。

若 field 的形状发生改变,就需要重写 Taichi 程序并重新编译,但 Ndarray 不同于 field,它的形状可以动态改变而无需重新编译。

以天体轨道模拟为示例,假如你想要将天体的数量翻一番,达到 2,000 个:

  • 使用 Taichi field,你必须编译两次;
  • 使用 Ndarray,你只需更新运行时程序。

在 Android 上使用 Ndarray 运行一个 Taichi 程序

下面的部分描述了在 Android 上使用 Ndarray 运行一个 Taichi 程序的过程。

  1. 生成必要的计算着色器
  2. 解析生成的 JSON 文件
  3. 准备 SSBO 和形状信息
  4. 准备渲染着色器
  5. 执行所有着色器
note

从第 2 步开始,你需要有自己的运行时程序。 我们提供了一个 用于 Android 的 Java 运行时程序示例 以供参考,但你可能需要根据所用平台和编程语言修改这些代码。

生成必要的计算着色器

下面的 Python 脚本定义了一个 Taichi AOT 模块,根据所选后端(OpenGL)生成并保存必要的计算着色器(GLES 着色器)。

Taichi kernel 和计算着色器不是一对一映射。 每个 Taichi kernel 都可以生成多个计算着色器,着色器数量通常和 kernel 中循环的数量相当。

import taichi as ti

ti.init(arch=ti.opengl, use_gles=True, allow_nv_shader_extension=False)

# Define constants for computation
G = 1
PI = 3.141592653
N = 1000
m = 5
galaxy_size = 0.4
planet_radius = 1
init_vel = 120
h = 1e-5
substepping = 10

# Define Taichi kernels
@ti.kernel
def initialize(pos: ti.types.ndarray(element_dim=1), vel: ti.types.ndarray(element_dim=1)):
center=ti.Vector([0.5, 0.5])
for i in pos:
theta = ti.random() * 2 * PI
r = (ti.sqrt(ti.random()) * 0.7 + 0.3) * galaxy_size
offset = r * ti.Vector([ti.cos(theta), ti.sin(theta)])
pos[i] = center+offset
vel[i] = [-offset.y, offset.x]
vel[i] *= init_vel

@ti.kernel
def compute_force(pos: ti.types.ndarray(element_dim=1), vel: ti.types.ndarray(element_dim=1), force: ti.types.ndarray(element_dim=1)):
for i in pos:
force[i] = ti.Vector([0.0, 0.0])
for i in pos:
p = pos[i]
for j in pos:
if i != j:
diff = p-pos[j]
r = diff.norm(1e-5)
f = -G * m * m * (1.0/r)**3 * diff
force[i] += f
dt = h/substepping
for i in pos:
vel[i].atomic_add(dt*force[i]/m)
pos[i].atomic_add(dt*vel[i])

# Define Ndarrays
pos = ti.Vector.ndarray(2, ti.f32, N)
vel = ti.Vector.ndarray(2, ti.f32, N)
force = ti.Vector.ndarray(2, ti.f32, N)

# Run the AOT module builder
def aot():
m = ti.aot.Module(ti.opengl)
m.add_kernel(initialize, (pos, vel))
m.add_kernel(compute_force, (pos, vel, force))

dir_name = 'nbody_aot'
m.save(dir_name, '')
aot()

第 3 行初始化 Taichi:

  1. 设置 use_glesTrue,生成用于 Android 的 GLES 计算着色器。
  2. 设置 allow_nv_shader_extensionFalse,防止生成的 GLES 计算着色器在 Android 上使用 Nvidia 的 GL 扩展。

此设置是因为 Android 支持 GLES API,但 GLES 不支持 NV_SHADER_EXTENSION

第 50-58 行定义和构建了 Taichi AOT模块:

  1. 创建一个 Taichi AOT 模块,指定其后端为 OpenGL:
     m = ti.aot.Module(ti.opengl)
  1. 在模块中添加必需的 kernel:initializecompute_force,每个 kernel 有自己的 Ndarrays:
m.add_kernel(initialize, (pos, vel))

m.add_kernel(compute_force, (pos, vel, force))
  1. 在当前的工作目录下指定一个文件夹来保存模块生成的文件:
dir_name = 'nbody_aot'

m.save(dir_name, '')

必要的计算着色器和一个 JSON 文件会出现在指定目录下。

解析生成的 JSON 文件

note

从本节中,你必须有自己的运行时程序。 我们提供了一个 [用于 Android 的 Java 运行时程序示例](https://github. com/taichi-dev/taichi-aot-demo/blob/master/nbody_ndarray/java_runtime/NbodyNdarray. java) 以供参考,但你可能需要根据所用平台和编程语言修改这些代码。

生成必要的 GLES 计算着色器后,你需要编写运行时程序,将以下 JSON 文件解析为某些数据结构。 该 JSON 文件包含执行计算着色器所需的全部必要信息。 文件由 Taichi kernel 组织,明确了每个 kernel 中的计算着色器和 Ndarray。 让我们进一步分析其中结构。

此处为简洁起见,省略了 kernel compute_force 的 JSON 对象。 完整的 JSON 文件,见 metadata.json

  • 由 Taichi kernel 组织

    • initialize(第 4 行)
    • compute_force(第 51 行)
  • kernel 特定计算着色器

    initialize 为例。该 kernel 生成了一个叫做 initialize_c54_00(第 7 行)的计算着色器和一个叫做 initialize_c54_01(第 13 行)的计算着色器。

  • kernel 相关 args_buff

    initialize kernel 被分到了一个 128 字节的 args_buffer(第 21 行)。 注意 args_buffer 的大小取决于 kernel 接收的 Ndarray(posvel)的数量(见第 19 行的 arg_count)。 initialize kernel,或者更准确地说是每个 kernel,都有一个专属的 args_buffer,用于存储 scalar_args(第 27 行)指定的标量参数和 array_args(第 28-45 行)指定的 Ndarray 形状信息。

    Ndarray 的形状信息是由 JSON 数组 array_args 中的参数索引编排的:0(第 29 行)对应 Ndarray pos1 (第 37 行)对应 Ndarray vel。 参数索引由调用 add_kernel() 时传递 Ndarray 的顺序决定。 请参阅 Python 脚本中的第 53 行。

    args_buffer 中 Ndarray pos 的形状信息在 args_buffer(第 64 行)中有 64 字节的偏移。 根据第 35 行和第 43 行,Ndarray pos 的形状信息占用了 args_buffer 的 96 - 64 = 32 字节。

    :::tip ATTENTION JSON 文件只指定了对应的 Ndarray 的维度(第 30、38 行),你可以在自己的运行时程序中动态更新 Ndarray 的形状。 :::

  • kernel 特定绑定索引

    used.arr_arg_to_bind_idx(第 46 行)将 kernel 中每个 Ndarray 的 SSBO 映射到计算着色器的“更全局”的绑定索引。 例如,"1": 5(第 48 行)将 Ndarray vel 绑定到绑定索引 5

{
"aot_data": {
"kernels": {
"initialize": {
"tasks": [
{
"name": "initialize_c54_00",
"src": "nbody_aot/initialize_c54_00.glsl",
"workgroup_size": 1,
"num_groups": 1
},
{
"name": "initialize_c54_01",
"src": "nbody_aot/initialize_c54_01.glsl",
"workgroup_size": 128,
"num_groups": 256
}
],
"arg_count": 2,
"ret_count": 0,
"args_buf_size": 128,
"ret_buf_size": 0,
"ext_arr_access": {
"0": 2,
"1": 3
},
"scalar_args": {},
"arr_args": {
"0": {
"field_dim": 1,
"is_scalar": false,
"element_shape": [
2
],
"shape_offset_in_bytes_in_args_buf": 64
},
"1": {
"field_dim": 1,
"is_scalar": false,
"element_shape": [
2
],
"shape_offset_in_bytes_in_args_buf": 96
}
},
"used.arr_arg_to_bind_idx": {
"0": 4,
"1": 5
}
},
"compute_force": {...}
},
"kernel_tmpls": {},
"fields": [],
"root_buffer_size": 0
}
}

下面详细描述生成的 JSON 文件中的键:

aot_data:总体的 JSON 对象。

  • kernels:所有的 Taichi kernel。
    • $(kernel_name):特定 Taichi kernel 的名称。
      • tasks:生成的计算着色器组成的 JSON 数组。
        • name:特定计算着色器的名称。
        • src:着色器文件的相对路径。
        • workgroup_size:N/A
        • num_groups:N/A
      • arg_count:Taichi kernel 接收的参数数量。
      • ret_count:Taichi kernel 返回值的数量。
      • args_buf_sizeargs_buf 所占字节大小.
      • ret_buf_sizeret_buf 所占字节大小。
      • scalar_args:kernel 接收的标量参数。
      • arr_args:kernel 中 Ndarray 的形状信息。
        • $(arg_index):Ndarray 的参数索引。
          • field_dim:Ndarray 的维度。
          • is_scalar:Ndarray 的元素是否为标量。
          • element_shape:一个 int 数组,表示 Ndarray 中每个元素的形状。
          • shape_offset_in_bytes_in_args_buf:Ndarray 的形状信息在 args_buf 中的偏移。
      • used.arr_arg_to_bind_idx:映射,明确给定 Ndarray 要绑定的 SSBO。 例如,"1": 5(第 48 行)将 Ndarray vel 绑定到绑定索引 5

希望你没有因为突然接收大量的信息而不知所措。 在下一节中,我们将重新审阅 JSON 文件, 提供表格和图表,帮助说明上文所列的一些概念。

准备 SSBO 和形状信息

在运行时程序执行 GLES 计算着色器之前,你需要准备好所有资源,包括:

  • 将 SSBO 绑定到对应的缓冲区
  • 绑定每一个 Ndarray 的 SSBO
  • 用 Ndarray 的形状信息填充 args_buffer

将 SSBO 绑定到对应的缓冲区

下表列出了 Taichi 程序常用的缓冲区及绑定索引:

缓冲区全局/kernel 相关存储绑定索引
root_buffer全局有固定偏移量和固定尺寸的全部 field0
gtmp_buffer全局全局临时数据1
args_bufferkernel 特定传递给 taichi kernel 的参数
  • 标量参数
  • 每个 Ndarray 的形状信息:
    • Ndarray 的形状
    • 元素的形状
2
  1. 如果你的 Taichi 脚本使用了至少一个 field,你需要为 root_buffer 绑定一个 SSBO。 如果脚本中不涉及 field,请跳过这一步。
  2. 将一个小 SSBO(如一个 1,024 字节的 SSBO)绑定到 gtmp_buffer 的绑定索引 1
  3. 将一个 64 x 5 = 320 字节的 SSBO 绑定到 args_buffer 的绑定参数 2

绑定每一个 Ndarray 的 SSBO

在运行时程序运行特定 kernel 前(如 initialize kernel),须根据 used.arr_arg_to_bind_idx 的值为 kernel 中每一个 Ndarray 绑定大小合适的 SSBO。

以下是上述 JSON 文件第 29-49 行的总结:

NdarrayTaichi kernel维度元素的形状参数索引绑定索引
posinitialize1[2]04
velinitialize1[2]15

如果设置每个 Ndarray 的形状为 [500],元素形状为 [2](意味着每个元素都是一个 2D 向量),那么:

  • 每个 Ndarray 有 500 x 2 = 1000 个数字
  • 因为数字类型是浮点(上文的 Python 脚本中指定),所以每个 Ndarray 的 SSBO 大小为 1000 x 4 = 4000 字节。

因此,你需要:

  • 对于 Ndarray pos,绑定一个 4,000 字节的 SSBO 到绑定索引 4
  • 对于 Ndarray vel,绑定一个 4,000 字节的 SSBO 到绑定索引 5

用 Ndarray 的形状信息填充 args_buffer

在说明 JSON 文件时,我们提到每个 kernel 都有一个专属的 args_buffer,用于存储 scalar_args 指定的标量参数和 array_args 指定的 Ndarray 形状信息。 array_args 没有指定 Ndarray 的形状,因此,准备工作的最后一步是在运行时程序中往 args_buffer 填充每个 Ndarray 的形状信息。

一个 args_buffer 的大小通常是 64 + 64 x 4 字节。 第一个 64 字节为标量参数预留;缓冲区剩余空间即为 64 x 4 字节。 每个 Ndarray 都被分配到 8 x 4 字节,用于存储其形状信息(每个用至多 8 个数字表示形状信息),因此,缓冲区剩余空间可以存储最多 8 个 Ndarray 的形状信息。

  • 如果你的 Ndarray 形状是 [100, 200],元素尺寸为 [3, 2],那么你应在相应位置填写 100、200、3 和 2。
  • 在本文示例中,posvel 的形状均为 [500],元素维度为 [2]。 因此,你应在相应位置填写 500 和 2。

准备渲染着色器

要执行渲染(在示例中也就是绘制天体),你需要遍写一个顶点着色器和一个片段着色器。

执行所有着色器

当在运行时程序中执行着色器时,请确保你在执行 Taichi kernel 前绑定了SSBO,并在完成后取消绑定。

我们的 用于 Android 的 Java 运行时程序示例 完成了以下步骤:

  1. initialize 中运行 GLES 计算着色器一次。
  2. 对于每帧:
    1. compute_force 运行 GLES 计算着色器 10 次。
    2. 为渲染运行顶点和片段着色器一次。

OpenGL相关术语与定义

OpenGL ES(GLES)

OpenGL ES(GLES)是 OpenGL 为嵌入式系统设计的 API。 根据其 说明,OpenGL 的桌面驱动程序支持所有的 GLES API。

OpenGL 着色语言(GLSL)

OpenGL 着色语言(GLSL)是 OpenGL 的主要着色语言。 GLSL 是一种受 OpenGL 直接支持的类 C 语言,无需扩展。

着色器

着色器是在图形处理器的某个阶段进行计算或渲染的用户自定义程序。

SSBO(着色器存储缓冲期对象)

每个 Taichi kernel 都可以生成多个计算着色器,这些着色器使用 SSBO 作为访问数据的缓冲。

有两种类型的 SSBO:一种对应 Taichi 维护的缓冲区,包括 root_buffergtmp_bufferargs_buffer;另一种对应开发者维护的 Ndarray,用于数据分享。

你需要在运行时程序中将生成的着色器绑定到相应的 SSBO。 Ndarray 的 SSBO 的绑定索引从 4 开始。