【Datawhale2609】算子开发实战 task02-TileLang Add 与 NineToothed Vector Add

系列文章目录

随着人工智能（Artificial Intelligence，AI）技术的迅猛发展，深度学习模型的复杂度和规模正呈指数级增长，这给底层计算效率带来了前所未有的挑战。

数学抽象的主线是：

从处理单个点，到处理整体曲线，再到变换整体，最后到不同数学体系之间做结构翻译。

函数、泛函、算子、函子，就是这四次视角升级。

一门并行编程语言及其编译器。

Triton 最初由 Philippe Tillet 开发，2020 年起由 OpenAI 接管并持续演进。它是一个开源的、基于 Python 的领域特定语言（DSL）和编译器，专门用于编写 GPU 上的深度学习 kernel
Triton的设计宗旨是提升AI模型训练过程中GPU编程的易用性与效率。
它允许开发者通过 Python 语言编写自定义 GPU 内核，实现与专家级 CUDA 代码相当的性能表现，同时无需掌握底层 CUDA 专业知识。
核心设计 ：基于 Tile 的编程模型
Triton 与 CUDA 的根本区别在于抽象层级：CUDA 以单个线程 为编程单位，Triton 以数据块 （tile/block）为编程单位。开发者用近似 NumPy 的 Python 代码描述计算逻辑，编译器自动处理线程调度、内存层级映射和硬件优化。

TileLang（Tile Language ）是一款基于张量分块（Tiling）抽象的高性能 AI 算子编程语言，属于 领域特定语言（DSL）。
它采用声明式语法与类 Python 前端，使开发者能够以接近数学公式的形式描述计算意图，并由编译器自动完成循环优化、内存调度与代码生成。

T.ceildiv(N, BLOCK_N)：计算需要启动的 Block 数量，向上取整：确保即使 N 不能被 BLOCK_N 整除时，也能覆盖所有元素。
T.Kernel(num_blocks) as pid：定义 Grid 级别的并行，pid 相当于 CUDA 中的 blockIdx.x。
T.Parallel(BLOCK_N) as i：定义 Block 内部的线程级并行，i 相当于 CUDA 中的 threadIdx.x。
if index < N:：边界保护（Tail handling），防止在最后一个 Block 中由于 N 不是 BLOCK_N 的整数倍而导致的全局内存越界访问。
return C：将分配并计算完成的 Tensor C 返回。

python 复制代码

def arrangement(lhs, rhs, output):
    return lhs.tile((BLOCK_SIZE,)), rhs.tile((BLOCK_SIZE,)), output.tile((BLOCK_SIZE,))

这里的 lhs, rhs, output 是符号张量，代表完整的 1D 数组。
.tile((BLOCK_SIZE,))：将一维的长数组，在逻辑上切割成一个个大小为 BLOCK_SIZE 的 Tile。
因为 tile 方法需要知道每个维度的切割大小。对于 1D 张量，必须传入包含一个元素的元组。如果是 2D 矩阵，就会是 (BLOCK_M, BLOCK_N)。

python 复制代码

def application(lhs, rhs, output):
    output = lhs + rhs