Compare commits
No commits in common. "main" and "dev" have entirely different histories.
36
.cnb.yml
36
.cnb.yml
|
|
@ -6,38 +6,4 @@ $:
|
|||
tags: cnb:arch:amd64:gpu
|
||||
services:
|
||||
- vscode
|
||||
- docker
|
||||
|
||||
push:
|
||||
- docker:
|
||||
image: docker.cnb.cool/yutianyu.yi/image/ncu
|
||||
runner:
|
||||
tags: cnb:arch:amd64:gpu
|
||||
stages:
|
||||
- name: Build & Test NVIDIA
|
||||
script: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
pip install ninja
|
||||
bash scripts/build_nvidia.sh all
|
||||
|
||||
- name: Test TileLang
|
||||
script: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
PYTHONPATH=python:. python tests/run_ops.py --op copy,vector_add,reduce_sum,softmax --backend tilelang --mode all
|
||||
|
||||
pull_request:
|
||||
- docker:
|
||||
image: docker.cnb.cool/yutianyu.yi/image/ncu
|
||||
runner:
|
||||
tags: cnb:arch:amd64:gpu
|
||||
stages:
|
||||
- name: Build & Test NVIDIA
|
||||
script: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
pip install ninja
|
||||
bash scripts/build_nvidia.sh all
|
||||
|
||||
- name: Test TileLang
|
||||
script: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
PYTHONPATH=python:. python tests/run_ops.py --op copy,vector_add,reduce_sum,softmax --backend tilelang --mode all
|
||||
- docker
|
||||
|
|
@ -1,75 +0,0 @@
|
|||
name: Jobs
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main, master]
|
||||
pull_request:
|
||||
branches: [main, master]
|
||||
|
||||
env:
|
||||
PYTHONPATH: python:.
|
||||
|
||||
jobs:
|
||||
# =========================================================================
|
||||
# NVIDIA CUDA backend
|
||||
# =========================================================================
|
||||
build-test-nvidia:
|
||||
runs-on: [self-hosted, nvidia]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
pip install -r requirements.txt
|
||||
|
||||
- name: Build & Test NVIDIA
|
||||
run: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
bash scripts/build_nvidia.sh all
|
||||
|
||||
# =========================================================================
|
||||
# TileLang JIT backend
|
||||
# =========================================================================
|
||||
test-tilelang:
|
||||
runs-on: [self-hosted, nvidia]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
pip install -r requirements.txt
|
||||
|
||||
- name: Run TileLang tests
|
||||
run: |
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate py312
|
||||
PYTHONPATH=python:. python tests/run_ops.py --op copy,vector_add,reduce_sum,softmax --backend tilelang --mode all
|
||||
|
||||
# =========================================================================
|
||||
# Summary
|
||||
# =========================================================================
|
||||
ci-summary:
|
||||
needs: [build-test-nvidia, test-tilelang]
|
||||
if: always()
|
||||
runs-on: [self-hosted, nvidia]
|
||||
steps:
|
||||
- name: Print CI Summary
|
||||
run: |
|
||||
echo "============================================"
|
||||
echo " CI Test Summary"
|
||||
echo "============================================"
|
||||
echo ""
|
||||
echo " NVIDIA: ${{ needs.build-test-nvidia.result }}"
|
||||
echo " TileLang: ${{ needs.test-tilelang.result }}"
|
||||
echo ""
|
||||
echo "============================================"
|
||||
|
||||
- name: Check overall status
|
||||
run: |
|
||||
FAILURES=0
|
||||
[ "${{ needs.build-test-nvidia.result }}" = "success" ] || FAILURES=$((FAILURES+1))
|
||||
[ "${{ needs.test-tilelang.result }}" = "success" ] || FAILURES=$((FAILURES+1))
|
||||
echo ""
|
||||
echo "Total job failures: $FAILURES / 2"
|
||||
[ "$FAILURES" -eq 0 ] || exit 1
|
||||
|
|
@ -14,17 +14,6 @@ __global__ void copy_contiguous_kernel(T *dst, const T *src, int64_t n) {
|
|||
// 2. Compute the grid-wide stride.
|
||||
// 3. Loop over i = idx; i < n; i += stride.
|
||||
// 4. Copy src[i] to dst[i].
|
||||
// 全局线程索引：当前线程在整个 grid 中的唯一 ID
|
||||
int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
// grid-stride：一次 grid 能处理的元素数，用于跨步遍历
|
||||
int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
|
||||
// grid-stride loop：每个线程处理 idx, idx+stride, idx+2*stride, ...
|
||||
// 好处：即使 n >> 线程总数，也能完整覆盖；且对 GPU cache 更友好
|
||||
for (int64_t i = idx; i < n; i += stride) {
|
||||
dst[i] = src[i];
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
} // namespace oprt::copy::nvidia
|
||||
|
|
|
|||
|
|
@ -16,12 +16,4 @@ def copy_kernel(src, BLOCK_N: int, dtype):
|
|||
# 1. Launch one TileLang kernel over the N // BLOCK_N tiles.
|
||||
# 2. Use T.copy to move one tile from src to out.
|
||||
|
||||
with T.Kernel(N // BLOCK_N, threads = 256) as pid_n:
|
||||
T.copy(
|
||||
src[pid_n * BLOCK_N:(pid_n + 1) * BLOCK_N],
|
||||
out[pid_n * BLOCK_N : pid_n * BLOCK_N],
|
||||
)
|
||||
|
||||
|
||||
|
||||
return out
|
||||
|
|
|
|||
|
|
@ -14,41 +14,6 @@ __global__ void reduce_sum_rowwise_kernel(float *out, const float *in, int64_t r
|
|||
// 3. Store the partial sums in shared memory.
|
||||
// 4. Reduce shared memory with a tree reduction.
|
||||
// 5. Let thread 0 write the final row sum to out[row].
|
||||
// 动态分配的共享内存，大小 = threads * sizeof(float)，由调用方传入
|
||||
// 共享内存在同一个 block 的所有线程间共享，速度接近寄存器
|
||||
extern __shared__ float sdata[];
|
||||
|
||||
// 一个 block 处理一行：blockIdx.x 就是行号
|
||||
int64_t row = blockIdx.x;
|
||||
if (row >= rows) return;
|
||||
|
||||
const float *row_in = in + row * cols;
|
||||
|
||||
// 阶段1：每个线程先串行累加自己负责的那部分元素
|
||||
// 当 cols > blockDim.x 时，线程需要跨步处理多个元素
|
||||
float sum = 0.0f;
|
||||
for (int64_t i = threadIdx.x; i < cols; i += blockDim.x) {
|
||||
sum += row_in[i];
|
||||
}
|
||||
// 每个线程的局部和写入共享内存对应位置
|
||||
sdata[threadIdx.x] = sum;
|
||||
// 必须同步：确保所有线程都写完 sdata 后再开始规约
|
||||
__syncthreads();
|
||||
|
||||
// 阶段2：树形规约（tree reduction）
|
||||
// 每轮将活跃元素减半：[0]+[128], [1]+[129], ... → [0]+[64], ... → 最终 sdata[0] 即总和
|
||||
// 时间复杂度 O(log2(threads))，比串行 O(threads) 快得多
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (threadIdx.x < stride) {
|
||||
sdata[threadIdx.x] += sdata[threadIdx.x + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 阶段3：线程0拥有最终结果，写入全局内存
|
||||
if (threadIdx.x == 0) {
|
||||
out[row] = sdata[0];
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace oprt::reduce_sum::nvidia
|
||||
|
|
|
|||
|
|
@ -26,25 +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.
|
||||
|
||||
with T.Kernel(N // BLOCK_N, threads = 256) as pid_n:
|
||||
src_frag = T.alloc_fragment((BLOCK_N,BLOCK_M),dtype)
|
||||
out_frag = T.alloc_fragment((BLOCK_N,),dtype)
|
||||
T.fill(out_frag,0.0)
|
||||
|
||||
for pid_m in T.Serial(M//BLOCK_M):
|
||||
T.copy(
|
||||
src[
|
||||
pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N,
|
||||
pid_m * BLOCK_M : (pid_m + 1) * BLOCK_M,
|
||||
],
|
||||
src_frag,
|
||||
)
|
||||
|
||||
T.reduce_sum(src_frag,out_frag,dim = -1, clear = False)
|
||||
|
||||
T.copy(out_frag , out[pid*BLOCK_N : (pid + 1) * BLOCK_N])
|
||||
|
||||
|
||||
|
||||
|
||||
return out
|
||||
|
|
|
|||
|
|
@ -17,55 +17,6 @@ __global__ void softmax_rowwise_kernel(float *out, const float *in, int64_t rows
|
|||
// and accumulate their sum.
|
||||
// 4. Reduce to get the row sum.
|
||||
// 5. Normalize each output element by row_sum.
|
||||
|
||||
extern __shared__ float sdata[];
|
||||
int64_t row = blockIdx.x;
|
||||
if (row >= rows) return;
|
||||
|
||||
const float *row_in = in + row * cols;
|
||||
float *row_out = out + row * cols;
|
||||
|
||||
// ---- 第1轮规约：求行最大值 row_max ----
|
||||
// 数值稳定技巧：softmax = exp(x_i) / sum(exp(x_j))
|
||||
// 直接算会溢出，必须减去最大值：exp(x_i - max) / sum(exp(x_j - max))
|
||||
float local_max = -FLT_MAX;
|
||||
for (int64_t i = threadIdx.x; i < cols; i += blockDim.x) {
|
||||
local_max = fmaxf(local_max, row_in[i]);
|
||||
}
|
||||
sdata[threadIdx.x] = local_max;
|
||||
__syncthreads();
|
||||
// 用树形规约取 max（与 reduce_sum 的加法规约结构一样，只是运算换成 fmaxf）
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (threadIdx.x < stride) {
|
||||
sdata[threadIdx.x] = fmaxf(sdata[threadIdx.x], sdata[threadIdx.x + stride]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
float row_max = sdata[0];
|
||||
__syncthreads(); // 必须同步，确保所有线程读到 row_max 后再进入下一阶段
|
||||
|
||||
// ---- 第2轮：计算 exp(x_i - row_max)，暂存到 out，同时累加求和 ----
|
||||
float local_sum = 0.0f;
|
||||
for (int64_t i = threadIdx.x; i < cols; i += blockDim.x) {
|
||||
float val = expf(row_in[i] - row_max);
|
||||
row_out[i] = val; // 先把 exp 值写出去，省掉额外的临时数组
|
||||
local_sum += val;
|
||||
}
|
||||
sdata[threadIdx.x] = local_sum;
|
||||
__syncthreads();
|
||||
// 树形规约求 sum（和 reduce_sum 的加法树完全一样）
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (threadIdx.x < stride) {
|
||||
sdata[threadIdx.x] += sdata[threadIdx.x + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
float row_sum = sdata[0];
|
||||
|
||||
// ---- 第3轮：归一化 out[i] /= row_sum ----
|
||||
for (int64_t i = threadIdx.x; i < cols; i += blockDim.x) {
|
||||
row_out[i] /= row_sum;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace oprt::softmax::nvidia
|
||||
|
|
|
|||
|
|
@ -33,78 +33,4 @@ def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):
|
|||
# - normalize with the final lse
|
||||
# - copy the result tile to global memory
|
||||
|
||||
# 每个 block 处理 BLOCK_N 行
|
||||
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
|
||||
# --- 分配寄存器片段 ---
|
||||
src_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) # 输入 tile
|
||||
exp_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) # 存放 exp(x - row_max)
|
||||
row_max = T.alloc_fragment((BLOCK_N,), dtype) # 每行的全局最大值（running）
|
||||
tile_max = T.alloc_fragment((BLOCK_N,), dtype) # 当前 tile 的行最大值
|
||||
row_sum = T.alloc_fragment((BLOCK_N,), dtype) # 每行的 exp 值之和（running）
|
||||
tile_sum = T.alloc_fragment((BLOCK_N,), dtype) # 当前 tile 的 exp 值之和
|
||||
|
||||
# 初始化：row_max 设为极小值，row_sum 设为 0
|
||||
T.fill(row_max, -1e30)
|
||||
T.fill(row_sum, 0.0)
|
||||
|
||||
# ============ 第一遍：Online Softmax ============
|
||||
# 核心：分块计算每行的 softmax 所需的 max 和 sum
|
||||
# 因为列方向可能很长（M 很大），不能一次性处理，
|
||||
# 所以用 online 算法逐 tile 更新 running max 和 running sum
|
||||
for pid_m in T.Serial(M // BLOCK_M):
|
||||
# 1) 从全局内存加载当前 (BLOCK_N, BLOCK_M) tile
|
||||
T.copy(
|
||||
src[
|
||||
pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N,
|
||||
pid_m * BLOCK_M : (pid_m + 1) * BLOCK_M,
|
||||
],
|
||||
src_frag,
|
||||
)
|
||||
# 2) 沿列方向 reduce_max：得到当前 tile 每行的最大值
|
||||
# (BLOCK_N, BLOCK_M) → (BLOCK_N,)
|
||||
T.reduce_max(src_frag, tile_max, dim=-1, clear=True)
|
||||
|
||||
# 3) 更新 running max，并对 running sum 做修正
|
||||
# 当 max 变大时，之前累积的 exp 值都偏大了，需要乘以一个校正因子缩放
|
||||
for i in T.Parallel(BLOCK_N):
|
||||
old_max = row_max[i]
|
||||
new_max = T.max(old_max, tile_max[i])
|
||||
# 校正因子：exp(old_max - new_max) = 2^(log2_e * (old_max - new_max))
|
||||
# 因为 old_max ≤ new_max，所以这个因子 ≤ 1，即缩小之前的 sum
|
||||
row_sum[i] = row_sum[i] * T.exp2(log2_e * (old_max - new_max))
|
||||
row_max[i] = new_max
|
||||
|
||||
# 4) 计算当前 tile 中每个元素的 exp(x - new_max)
|
||||
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
|
||||
exp_frag[i, j] = T.exp2(log2_e * (src_frag[i, j] - row_max[i]))
|
||||
|
||||
# 5) 沿列方向 reduce_sum：得到当前 tile 每行的 exp 值之和
|
||||
T.reduce_sum(exp_frag, tile_sum, dim=-1, clear=True)
|
||||
# 6) 将当前 tile 的部分和累加到 running sum
|
||||
for i in T.Parallel(BLOCK_N):
|
||||
row_sum[i] = row_sum[i] + tile_sum[i]
|
||||
|
||||
# ============ 第二遍：归一化 ============
|
||||
# 此时 row_max[i] 和 row_sum[i] 已经是最终的全局 max 和全局 sum
|
||||
# softmax(x) = exp(x - max) / sum
|
||||
for pid_m in T.Serial(M // BLOCK_M):
|
||||
T.copy(
|
||||
src[
|
||||
pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N,
|
||||
pid_m * BLOCK_M : (pid_m + 1) * BLOCK_M,
|
||||
],
|
||||
src_frag,
|
||||
)
|
||||
# 对每个元素计算 exp(x - row_max) / row_sum
|
||||
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
|
||||
src_frag[i, j] = T.exp2(log2_e * (src_frag[i, j] - row_max[i])) / row_sum[i]
|
||||
# 写回全局内存
|
||||
T.copy(
|
||||
src_frag,
|
||||
out[
|
||||
pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N,
|
||||
pid_m * BLOCK_M : (pid_m + 1) * BLOCK_M,
|
||||
],
|
||||
)
|
||||
|
||||
return out
|
||||
|
|
|
|||
|
|
@ -29,13 +29,6 @@ __global__ void vector_add_contiguous_kernel(T * __restrict__ out,
|
|||
// 2. Compute the grid-wide stride.
|
||||
// 3. Loop over i = idx; i < n; i += stride.
|
||||
// 4. Write out[i] = add_values(a[i], b[i]).
|
||||
|
||||
int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
|
||||
// 同样的 grid-stride loop 模式，通过 add_values 统一处理不同精度
|
||||
for (int64_t i = idx; i < n; i += stride) {
|
||||
out[i] = add_values(a[i], b[i]);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace oprt::vector_add::nvidia
|
||||
|
|
|
|||
|
|
@ -18,9 +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].
|
||||
|
||||
with T.Kernel(N // BLOCK_N, threads = 256) as pid_n:
|
||||
base = pid_n * BLOCK_N
|
||||
for i in T.Parallel(BLOCK_N):
|
||||
out[base + i] = a[base + i] + b[base + i]
|
||||
|
||||
return out
|
||||
|
|
|
|||
|
|
@ -109,7 +109,7 @@ def _run_bench_module(op: str, backend: str):
|
|||
def main() -> int:
|
||||
ops = _discover_ops()
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--op", default="all")
|
||||
parser.add_argument("--op", choices=[*ops, "all"], default="all")
|
||||
parser.add_argument("--backend", choices=["nvidia", "tilelang", "metax"], default="nvidia")
|
||||
parser.add_argument("--mode", choices=["test", "bench", "all"], default="all")
|
||||
args = parser.parse_args()
|
||||
|
|
@ -118,10 +118,7 @@ def main() -> int:
|
|||
print("CUDA is required for benchmark", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
if args.op == "all":
|
||||
selected_ops = ops
|
||||
else:
|
||||
selected_ops = tuple(op.strip() for op in args.op.split(",") if op.strip())
|
||||
selected_ops = ops if args.op == "all" else (args.op,)
|
||||
rows: list[list[str]] = []
|
||||
bench_rows = []
|
||||
failed = False
|
||||
|
|
|
|||
Loading…
Reference in New Issue