在本系列之前的文章[1][2]中，我们了解到 Taichi 和 Torch 分别服务于不同的应用场景：Taichi 能够更精细地控制并行和每个元素 (element) 的操作，极大地提升了用户操作的灵活性。而 Torch 则将这些细节抽象成张量 (Tensor) 级别的操作，使得用户能聚焦于机器学习的模型结构。作为机器学习、计算机图形学领域炙手可热的框架和编程语言，Torch 和 Taichi 能否各取所长，结合使用呢？答案是肯定的。在本篇文章中，作者将通过两个简单的例子演示：如何使用 Taichi Kernel 来实现 PyTorch 程序中特殊的数据预处理和自定义的算子，告别手写 CUDA，用轻巧便捷的方式提升机器学习模型算法的开发效率和灵活性。

案例1：数据预处理
边缘填充（Padding）是机器学习中常用的预处理方法。如在对图像执行卷积操作时，用户需要对图像边缘进行填充，以保证图像输入输出前后的尺寸不变。一般来说，填充的方法有零填充或 torch.nn.functional.pad 提供的重复填充、循环填充等其他预设模式。但有时候我们想要在边缘上填充某个特殊的纹理或者模式，却并没有一个精心优化过的 PyTorch 算子能够适配这种场景。

解决方案有两个：

使用 PyTorch 或者 Python 逐个操作矩阵元素；
手写 C++ 或 CUDA 代码并接入PyTorch。
前者的计算效率非常低，会拖累神经网络的训练速度；后者学习曲线陡峭，实操非常麻烦，开发流程冗长。那么，有没有更好的方案呢？

接下来我们将通过一个例子，带大家体验如何用 Taichi 做一个砖墙纹理的边缘填充。

用 Taichi 给 PyTorch “添砖加瓦”！
第一步，我们在 PyTorch 中创建一个如下图所示的「砖块」。为了更好地观察填充的规律，我们给这块「砖」填充上了渐变的颜色：


填充的基本单元...
第二步，我们想要在 x 轴上错位重复这个「砖」，也就是如下所示的效果：


由于 PyTorch 中没有为这样的填充提供原生的算子，为了提高运算效率，需要将 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)
这一系列的矩阵操作并不是特别直观，而且需要在 GPU 内存中保存多个中间结果矩阵。一个较为明显的缺点是显存比较小的卡上可能就跑不起来了。



而如果使用 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()
这段代码逻辑非常简单：遍历输出图片的每个像素，计算当前像素对应到输入的「砖块」图片中的位置，最后复制该位置的颜色到这个像素。虽然看起来是在逐个写入每个像素，但 Taichi 会将 kernel 的顶层 for-loop 编译成高度并行的 GPU 代码。同时，上一段代码中我们直接把两个 PyTorch 的 Tensor 传给了 Taichi 函数 ti_pad ，Taichi 会直接使用 PyTorch 分配好的内存，不会因为两个框架间的数据交互而产生额外开销。

最后，实际的运算性能是：在 RTX3090 GPU 上运行时，PyTorch (v1.12.1) 耗费了 30.392 ms[1]，而 Taichi 版本的 Kernel 耗时仅 0.267 ms[2]，Taichi 相对 PyTorch 的加速比超过了 100 倍。

事实上，上述的 PyTorch 底层实现需要启动 58 个 CUDA Kernel，而本例中 Taichi 将全部运算编译成了 1 个 CUDA Kernel。更少的 Kernel 减少了 GPU 函数启动的开销，且相比 PyTorch 实现，Taichi 节省了大量冗余的内存操作。在 GPU 上内存操作远比运算操作开销更「昂贵」，这也是非常夸张的加速比的来源。在 Taichi 的设计过程中，我们遵循了 「Megakernel」的设计准则：使用单个大的 Kernel 去完成尽可能多的运算逻辑，这与机器学习系统设计中常见的 「算子融合优化」是一样的道理。

在数据预处理问题上，一方面 taichi拥有更精细的操作颗粒度，能灵活适配研究人员不同的需求，另一方面 Taichi 能达到更高的计算性能，显著提升预处理部分的运行速度。当然，预处理仅仅是机器学习训练和推理过程中的一小步，对于机器学习领域的研究人员来说，有大量时间花费在模型前向和反向的计算算子中。那么对于定制高性能 ML 算子，Taichi 有什么好办法？taichi https://taichi-lang.cn/