forked from ccf-ai-infra/Intro-ops
Compare commits
No commits in common. "v0.03" and "main" have entirely different histories.
|
|
@ -1,5 +0,0 @@
|
|||
# 默认忽略的文件
|
||||
/shelf/
|
||||
/workspace.xml
|
||||
# 基于编辑器的 HTTP 客户端请求
|
||||
/httpRequests/
|
||||
|
|
@ -1,15 +0,0 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<module type="PYTHON_MODULE" version="4">
|
||||
<component name="NewModuleRootManager">
|
||||
<content url="file://$MODULE_DIR$" />
|
||||
<orderEntry type="jdk" jdkName="C:\Users\creat\miniconda3" jdkType="Python SDK" />
|
||||
<orderEntry type="sourceFolder" forTests="false" />
|
||||
</component>
|
||||
<component name="PyDocumentationSettings">
|
||||
<option name="format" value="PLAIN" />
|
||||
<option name="myDocStringFormat" value="Plain" />
|
||||
</component>
|
||||
<component name="TestRunnerService">
|
||||
<option name="PROJECT_TEST_RUNNER" value="py.test" />
|
||||
</component>
|
||||
</module>
|
||||
|
|
@ -1,24 +0,0 @@
|
|||
<component name="InspectionProjectProfileManager">
|
||||
<profile version="1.0">
|
||||
<option name="myName" value="Project Default" />
|
||||
<inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
|
||||
<option name="ourVersions">
|
||||
<value>
|
||||
<list size="4">
|
||||
<item index="0" class="java.lang.String" itemvalue="3.14" />
|
||||
<item index="1" class="java.lang.String" itemvalue="3.13" />
|
||||
<item index="2" class="java.lang.String" itemvalue="3.12" />
|
||||
<item index="3" class="java.lang.String" itemvalue="3.11" />
|
||||
</list>
|
||||
</value>
|
||||
</option>
|
||||
</inspection_tool>
|
||||
<inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
|
||||
<option name="ignoredErrors">
|
||||
<list>
|
||||
<option value="N806" />
|
||||
</list>
|
||||
</option>
|
||||
</inspection_tool>
|
||||
</profile>
|
||||
</component>
|
||||
|
|
@ -1,6 +0,0 @@
|
|||
<component name="InspectionProjectProfileManager">
|
||||
<settings>
|
||||
<option name="USE_PROJECT_PROFILE" value="false" />
|
||||
<version value="1.0" />
|
||||
</settings>
|
||||
</component>
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="Black">
|
||||
<option name="sdkName" value="C:\Users\creat\miniconda3" />
|
||||
</component>
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="C:\Users\creat\miniconda3" project-jdk-type="Python SDK" />
|
||||
</project>
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectModuleManager">
|
||||
<modules>
|
||||
<module fileurl="file://$PROJECT_DIR$/.idea/Intro-ops.iml" filepath="$PROJECT_DIR$/.idea/Intro-ops.iml" />
|
||||
</modules>
|
||||
</component>
|
||||
</project>
|
||||
|
|
@ -1,6 +0,0 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="" vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
|
|
@ -5,81 +5,15 @@ 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)
|
||||
|
||||
# 张量操作：
|
||||
# # 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.
|
||||
# TODO: 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.
|
||||
|
||||
# 计算 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
|
||||
return out
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
# DONE: implement a tiled row-wise reduce_sum kernel.
|
||||
# TODO: implement a tiled row-wise reduce_sum kernel.
|
||||
#
|
||||
# Suggested steps:
|
||||
# 1. Launch one TileLang kernel over row tiles.
|
||||
|
|
@ -26,32 +26,4 @@ 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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
# DONE: implement a tiled row-wise softmax kernel.
|
||||
# TODO: implement a tiled row-wise softmax kernel.
|
||||
#
|
||||
# Suggested steps:
|
||||
# 1. Launch one TileLang kernel over row tiles.
|
||||
|
|
@ -33,72 +33,4 @@ 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
|
||||
|
|
|
|||
|
|
@ -1,34 +0,0 @@
|
|||
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]
|
||||
|
|
@ -1,44 +1,16 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
|
||||
|
||||
@tilelang.jit
|
||||
def vector_add_kernel(a, b, BLOCK_N: int, dtype):
|
||||
"""
|
||||
TileLang vector add kernel - 使用 T.copy 实现向量加法
|
||||
|
||||
核心逻辑：
|
||||
1. 将全局内存中的数据块（tile）加载到片上 fragment
|
||||
2. 在片上执行加法运算
|
||||
3. 将结果写回全局内存
|
||||
"""
|
||||
N = T.const("N")
|
||||
a: T.Tensor((N,), dtype)
|
||||
b: T.Tensor((N,), dtype)
|
||||
out = T.empty((N,), dtype)
|
||||
|
||||
# 分配片上内存（fragment）
|
||||
a_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
||||
b_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
||||
out_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
||||
|
||||
num_tiles = N // BLOCK_N
|
||||
|
||||
# 并行处理每个 tile
|
||||
for tile_idx in T.Parallel(num_tiles):
|
||||
offset = tile_idx * BLOCK_N
|
||||
|
||||
# ========== T.copy 核心用法 ==========
|
||||
# 1. 从全局内存 -> 片上内存（加载）
|
||||
T.copy(a[offset:offset + BLOCK_N], a_frag)
|
||||
T.copy(b[offset:offset + BLOCK_N], b_frag)
|
||||
|
||||
# 2. 片上计算（高效）
|
||||
out_frag = a_frag + b_frag
|
||||
|
||||
# 3. 从片上内存 -> 全局内存（存储）
|
||||
T.copy(out_frag, out[offset:offset + BLOCK_N])
|
||||
|
||||
# TODO: implement a tile-wise vector add kernel.
|
||||
#
|
||||
# Suggested steps:
|
||||
|
|
@ -46,4 +18,4 @@ def vector_add_kernel(a, b, BLOCK_N: int, dtype):
|
|||
# 2. Compute the tile base offset.
|
||||
# 3. Use T.Parallel(BLOCK_N) to fill out[base + i] = a[base + i] + b[base + i].
|
||||
|
||||
return out
|
||||
return out
|
||||
|
|
|
|||
Loading…
Reference in New Issue