Merge remote-tracking branch 'origin/v0.01' into v0.01

# Conflicts:
#	ops/vector_add/tilelang/kernel.py

ops/copy/tilelang/kernel.py

@@ -5,15 +5,81 @@ import tilelang.language as T
 
 
 @tilelang.jit
+# 按readme，环境配完了，来实现第一个 tilelang kernel - copy_kernel
+# 基本概念：
+# 1.@tilelang.jit 装饰器：将函数编译为 GPU 代码
+# tilelang 库：提供 TileLang 语言的 API
+# 包括所有 TileLang “原语”，用于定义和编译 kernel
+# TileLang 原语：TileLang 语言的基本操作，如 T.copy、T.empty、T.const 等
+# TileLang kernel：TileLang 语言的函数，用于在 GPU 上并行执行数据操作
+# Tile：GPU 处理的基本数据块，类似 CUDA 的 thread block
+# 每个 tile 包含 BLOCK_N 个数据元素
+# 2.函数参数：src, BLOCK_N, dtype
+# 3.函数返回值：out
+
 def copy_kernel(src, BLOCK_N: int, dtype):
+    """
+    TileLang copy kernel - 将数据从 src 拷贝到 out
+    
+    参数说明：
+    - src: 输入张量，形状为 (N,), 类型为 dtype
+    - BLOCK_N: 每个 tile 的大小（编译时常量）
+    - dtype: 数据类型（如 float32）
+    - N: 全局编译时常量，通过 T.const("N") 获取，表示张量总长度
+    
+    功能：将 src 的内容逐 tile 拷贝到 out
+    """
     N = T.const("N")
     src: T.Tensor((N,), dtype)
     out = T.empty((N,), dtype)
 
-    # TODO: implement a tile-wise copy kernel.
+    # 张量操作：
+    # # 1. 创建张量
+    # out = T.empty((N,), dtype)      # 创建空张量
+    # out = T.fill(frag, 0.0)      # 创建全零张量
+    # out = T.fill(frag, value)       # 创建全 value 的张量
+    # out = T.alloc_fragment(shape, dtype)  # 创建空 fragment，分配内存，值未定义
+    # # 2. 张量类型声明
+    # src: T.Tensor((N,), dtype)     # 声明张量形状和类型
+
+    # DONE: implement a tile-wise copy kernel.
     #
     # Suggested steps:
     # 1. Launch one TileLang kernel over the N // BLOCK_N tiles.
     # 2. Use T.copy to move one tile from src to out.
 
-    return out
+    # 计算 tile 数量
+    num_tiles = N // BLOCK_N
+    
+    # 并行遍历所有 tiles
+    """
+    两种循环：
+    # 并行循环（每个 iteration 在不同 GPU 线程执行）
+    for idx in T.Parallel(num_tiles):
+        # 每个线程独立执行
+        T.copy(src[idx*BLOCK_N:(idx+1)*BLOCK_N], 
+               out[idx*BLOCK_N:(idx+1)*BLOCK_N])
+    
+    # 串行循环（在同一个线程内顺序执行）
+    for i in T.Serial(10):
+        # 顺序执行 10 次
+        ...
+    """
+
+    for tile_idx in T.Parallel(num_tiles):
+        # 计算当前 tile 的起始偏移
+        offset = tile_idx * BLOCK_N
+        # 使用 T.copy 搬运一个 tile 的数据
+        # # 3. 张量切片
+        srcTile = src[offset:offset+BLOCK_N]  # 提取一个 tile
+        outTile = out[offset:offset+BLOCK_N]  # 提取一个 tile
+        T.copy(srcTile, outTile)  # 搬运一个 tile 的数据
+        # T.copy(srcTile, outTile)  # 搬运一个 tile 的数据，等价于 T.copy(src[offset:offset+BLOCK_N], out[offset:offset+BLOCK_N])
+        # # 基础用法：将 src 的一部分拷贝到 out
+        # T.copy(src[start:end], out[start:end])
+        #
+        # # 示例：拷贝一个 tile
+        # T.copy(src[offset:offset+BLOCK_N],
+        #        out[offset:offset+BLOCK_N])
+
+    return out

ops/reduce_sum/tilelang/kernel.py

@@ -15,7 +15,7 @@ def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):
     src: T.Tensor((N, M), dtype)
     out = T.empty((N,), dtype)
 
-    # TODO: implement a tiled row-wise reduce_sum kernel.
+    # DONE: implement a tiled row-wise reduce_sum kernel.
     #
     # Suggested steps:
     # 1. Launch one TileLang kernel over row tiles.
@@ -26,4 +26,32 @@ def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):
     # 6. Call T.reduce_sum on the fragment and accumulate into the output fragment.
     # 7. Copy the final output fragment back to global memory.
 
+    # 计算行和列的 tile 数量
+    num_row_tiles = N // BLOCK_N
+    num_col_tiles = M // BLOCK_M
+
+    for row_idx in T.Parallel(num_row_tiles):
+        # 计算当前 tile 的起始偏移
+        row_offset = row_idx * BLOCK_N
+
+        in_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
+        out_frag = T.alloc_fragment((BLOCK_N,), dtype)
+        row_sum_frag = T.alloc_fragment((BLOCK_N,), dtype)
+
+        # 将 out_frag 初始化为 0
+        out_frag = T.fill(out_frag, 0.0)
+
+        for col_idx in T.Serial(num_col_tiles):
+            col_offset = col_idx * BLOCK_M
+
+            # 将一个列 tile 加载到共享内存
+            T.copy(src[row_offset:row_offset + BLOCK_N, col_offset:col_offset + BLOCK_M], in_frag)
+
+            T.reduce_sum(in_frag, row_sum_frag)  # 将每行的和写入 row_sum_frag
+
+            # 累加到输出 fragment
+            out_frag = out_frag + row_sum_frag
+            # 将最终结果写回全局内存
+        T.copy(out_frag, out[row_offset:row_offset + BLOCK_N])
+
     return out

ops/softmax/tilelang/kernel.py

@@ -16,7 +16,7 @@ def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):
     src: T.Tensor((N, M), dtype)
     out = T.empty((N, M), dtype)
 
-    # TODO: implement a tiled row-wise softmax kernel.
+    # DONE: implement a tiled row-wise softmax kernel.
     #
     # Suggested steps:
     # 1. Launch one TileLang kernel over row tiles.
@@ -33,4 +33,72 @@ def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):
     #    - normalize with the final lse
     #    - copy the result tile to global memory
 
+    row_tiles = N // BLOCK_N
+    col_tiles = M // BLOCK_M
+
+    # Step 2: Allocate fragments
+    src_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
+    out_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
+    exp_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
+
+    # Running state fragments
+    running_max_frag = T.alloc_fragment((BLOCK_N,), dtype)  # 跨tile累积的行最大值
+    running_sum_frag = T.alloc_fragment((BLOCK_N,), dtype)  # 跨tile累积的指数和
+
+    # Per-tile fragments
+    tile_max_frag = T.alloc_fragment((BLOCK_N,), dtype)
+    tile_sum_frag = T.alloc_fragment((BLOCK_N,), dtype)
+
+    # Step 1: Launch parallel kernel over row tiles
+    for row_idx in T.Parallel(row_tiles):
+        row_offset = row_idx * BLOCK_N
+
+        # Step 3: Initialize running log-sum-exp state
+        running_max_frag = T.fill(running_max_frag, float('-inf'))
+        running_sum_frag = T.fill(running_sum_frag, 0.0)
+
+        # Step 4: First pass - compute running lse across column tiles
+        for col_idx in T.Serial(col_tiles):
+            col_offset = col_idx * BLOCK_M
+
+            # 4.1 Copy input tile
+            T.copy(src[row_offset:row_offset + BLOCK_N, col_offset:col_offset + BLOCK_M], src_frag)
+
+            # 4.2 Compute tile max
+            T.reduce_max(tile_max_frag, src_frag)
+
+            # 4.3 Compute exp2(x - tile_max) for numerical stability
+            exp_frag = T.exp2((src_frag - tile_max_frag) * log2_e)
+
+            # 4.4 Compute tile sum of exp values
+            T.reduce_sum(tile_sum_frag, exp_frag)
+
+            # 4.5 Update running max and sum
+            # 找到两个max中的较大值
+            max_val = T.max(running_max_frag, tile_max_frag)  # 假设 T.max 存在
+            # 使用数学公式合并sum
+            term1 = running_sum_frag * T.exp2((running_max_frag - max_val) * log2_e)
+            term2 = tile_sum_frag * T.exp2((tile_max_frag - max_val) * log2_e)
+            sum_val = term1 + term2
+
+            running_max_frag = max_val
+            running_sum_frag = sum_val
+
+        # Compute final log-sum-exp: lse = max + log(sum)
+        final_lse_frag = running_max_frag + (T.log2(running_sum_frag) / log2_e)
+
+        # Step 5: Second pass - normalize each tile with final lse
+        for col_idx in T.Serial(col_tiles):
+            col_offset = col_idx * BLOCK_M
+
+            # 5.1 Copy input tile again
+            T.copy(src[row_offset:row_offset + BLOCK_N, col_offset:col_offset + BLOCK_M], src_frag)
+
+            # 5.2 Normalize: softmax(x) = exp(x - lse)
+            out_frag = T.exp2((src_frag - final_lse_frag) * log2_e)
+
+            # 5.3 Write result to global memory
+            T.copy(out_frag, out[row_offset:row_offset + BLOCK_N, col_offset:col_offset + BLOCK_M])
+
+
     return out

ops/vector_add/tilelang/kernel-1.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import tilelang
+import tilelang.language as T
+
+
+@tilelang.jit
+def vector_add_kernel(a, b, BLOCK_N: int, dtype):
+    N = T.const("N")
+    a: T.Tensor((N,), dtype)
+    b: T.Tensor((N,), dtype)
+    out = T.empty((N,), dtype)
+
+    # DONE: implement a tile-wise vector add kernel.
+    #
+    # Suggested steps:
+    # 1. Launch one TileLang kernel over the N // BLOCK_N tiles.
+    # 2. Compute the tile base offset.
+    # 3. Use T.Parallel(BLOCK_N) to fill out[base + i] = a[base + i] + b[base + i].
+    num_tiles = N // BLOCK_N
+
+    for tile_idx in T.Parallel(num_tiles):
+        # 计算当前 tile 的起始偏移
+        offset = tile_idx * BLOCK_N
+        # 使用 T.copy 搬运一个 tile 的数据
+        aTile = a[offset:offset+BLOCK_N]  # 提取一个 tile
+        bTile = b[offset:offset+BLOCK_N]  # 提取一个 tile
+        outTile = out[offset:offset+BLOCK_N]  # 提取一个 tile
+        T.copy(aTile+bTile, outTile)  # 搬运一个 tile 的数据
+    return out
+
+    # 或者也可以像上面那样注释一样写：
+    # for i in T.Parallel(BLOCK_N):  # 内层并行
+    #     out[offset + i] = a[offset + i] + b[offset + i]