القائمة
TileLang CUDA 内存访问优化策略，包括 T.alloc_shared/fragment 使用、数据布局优化、合并访存和 Bank Conflict 避免技巧。适用于内存带宽受限、需要优化数据搬运效率的 TileLang 内核性能优化场景
يُحدَّد Prompt الذي يراجع المصدر أولًا بشكل افتراضي. يمكنك التبديل إلى أمر مباشر أو تنزيل نسخة محلية.
اقرأ SKILL.md وأي ملفات مرافقة يعرضها SkillsMP قبل أن تقرر التثبيت.
التثبيت باستخدام Codex أو Claude انسخ هذا Prompt والصقه في Codex أو Claude أو مساعد آخر ليراجع صفحة Skill ويثبّتها لك.
يتجاوز الأمر المباشر Prompt المخصّص للمراجعة. افحص المصدر قبل تشغيله.
npx skills add https://github.com/mindspore-ai/akg --skill tilelang-cuda-memoryيبقى الأمر في سطر واحد. مرّر أفقيًا لمراجعته كاملًا قبل النسخ.
تفضّل نسخة محلية؟ نزّل الملفات المتاحة حاليًا لدى SkillsMP.
عرض SKILL.md
استنادا إلى تصنيف SOC المهني
| name | tilelang-cuda-memory |
| description | TileLang CUDA 内存访问优化策略，包括 T.alloc_shared/fragment 使用、数据布局优化、合并访存和 Bank Conflict 避免技巧。适用于内存带宽受限、需要优化数据搬运效率的 TileLang 内核性能优化场景 |
| category | implementation |
| version | 1.0.0 |
| metadata | {"backend":"cuda","dsl":"tilelang_cuda"} |
内存访问是 GPU 性能的关键瓶颈。本文档提供 TileLang CUDA 的内存访问优化策略。
| 内存类型 | 带宽 (A100) | 延迟 | 容量 | TileLang API |
|---|---|---|---|---|
| 寄存器 | ~19 TB/s | 1 cycle | 256 KB/SM | T.alloc_fragment / T.alloc_local |
| 共享内存 | ~19 TB/s | ~20 cycles | 164 KB/SM | T.alloc_shared |
| L2 缓存 | ~5 TB/s | ~100 cycles | 40 MB | T.use_swizzle 优化 |
| 全局内存 (HBM) | ~2 TB/s | ~400 cycles | 40/80 GB | T.Tensor |
T.copy 进行高效数据传输T.use_swizzle 优化数据局部性# 共享内存用于缓存从全局内存加载的数据块
A_shared = T.alloc_shared((block_M, block_K), "float16")
B_shared = T.alloc_shared((block_K, block_N), "float16")
# 高效数据加载
T.copy(A[by * block_M, ko * block_K], A_shared)
适用场景：
# 寄存器用于累加和局部计算
C_local = T.alloc_fragment((block_M, block_N), "float")
T.clear(C_local)
# 累加操作
T.gemm(A_shared, B_shared, C_local)
适用场景：
# 本地内存用于线程私有变量
C_reg = T.alloc_local((1,), "float")
T.clear(C_reg)
适用场景：
# ✅ 推荐：使用 T.copy 自动合并访问
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)
# ✅ 结果写回
T.copy(C_local, C[by * block_M, bx * block_N])
# 并行数据复制
for k, j in T.Parallel(block_K, block_N):
B_shared[k, j] = B[ko * block_K + k, bx * block_N + j]
# 向量化加载以提高带宽利用率
for k in T.vectorized(TILE_K):
A_local[k] = A[bk * BLOCK_K + tk * TILE_K + k]
B_local[k] = B[bn * BLOCK_N + tn, bk * BLOCK_K + tk * TILE_K + k]
# 使用 T.Pipelined 实现软件流水线
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
# 数据加载和计算自动重叠
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local)
| num_stages | 共享内存使用 | 性能 | 适用场景 |
|---|---|---|---|
| 2 | 最少 | 基础 | 共享内存紧张时 |
| 3 | 中等 | 通常最优 | 默认推荐 |
| 4 | 较多 | 大矩阵时更优 | 共享内存充足时 |
| 5+ | 很多 | 可能下降 | 需要测试验证 |
通过 T.use_swizzle 改善 L2 缓存局部性，特别适用于矩阵乘法等 2D 计算：
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
# 启用 swizzle 以提高 L2 缓存命中率
T.use_swizzle(panel_size=10, enable=True)
# ... 计算逻辑 ...
from tilelang.intrinsics import make_mma_swizzle_layout
# 使用布局注解优化共享内存访问模式
T.annotate_layout({
A_shared: make_mma_swizzle_layout(A_shared),
B_shared: make_mma_swizzle_layout(B_shared),
})
| 算子类型 | block 大小 | 线程数 | 说明 |
|---|---|---|---|
| Element-wise | block=256~1024 | 128-256 | 一维并行 |
| GEMM | M=128, N=128, K=32 | 128 | 二维分块 |
| GEMV | N=64 | 128 | 一维 + 串行 |
| Reduce | block=256~512 | 128-256 | 归约维度分块 |
| LayerNorm | block=256 | 256 | 按行处理 |
import tilelang
import tilelang.language as T
from tilelang.intrinsics import make_mma_swizzle_layout
@tilelang.jit(out_idx=[-1])
def optimized_matmul(M, N, K, block_M, block_N, block_K):
@T.prim_func
def main(A: T.Tensor((M, K), "float16"),
B: T.Tensor((K, N), "float16"),
C: T.Tensor((M, N), "float16")):
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
# 1. 内存分配
A_shared = T.alloc_shared((block_M, block_K), "float16")
B_shared = T.alloc_shared((block_K, block_N), "float16")
C_local = T.alloc_fragment((block_M, block_N), "float")
# 2. 布局优化
T.annotate_layout({
A_shared: make_mma_swizzle_layout(A_shared),
B_shared: make_mma_swizzle_layout(B_shared),
})
# 3. L2 缓存优化
T.use_swizzle(panel_size=10, enable=True)
T.clear(C_local)
# 4. 软件流水线
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local)
T.copy(C_local, C[by * block_M, bx * block_N])
main