第一次提交

.cnb.yml

@@ -6,4 +6,38 @@ $:
         tags: cnb:arch:amd64:gpu
       services:
         - vscode
-        - docker
+        - 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

.github/workflows/build.yml

@@ -0,0 +1,75 @@
+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

ops/copy/nvidia/kernel.cuh

@@ -14,6 +14,17 @@ __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

ops/copy/tilelang/kernel.py

@@ -16,4 +16,12 @@ 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

ops/reduce_sum/nvidia/kernel.cuh

@@ -14,6 +14,41 @@ __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

ops/reduce_sum/tilelang/kernel.py

@@ -26,4 +26,25 @@ 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

ops/softmax/nvidia/kernel.cuh

@@ -17,6 +17,55 @@ __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

ops/softmax/tilelang/kernel.py

@@ -33,4 +33,78 @@ 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

ops/vector_add/nvidia/kernel.cuh

@@ -29,6 +29,13 @@ __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

ops/vector_add/tilelang/kernel.py

@@ -18,4 +18,9 @@ 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

tests/run_ops.py

@@ -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", choices=[*ops, "all"], default="all")
+    parser.add_argument("--op", 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,7 +118,10 @@ def main() -> int:
         print("CUDA is required for benchmark", file=sys.stderr)
         return 2
 
-    selected_ops = ops if args.op == "all" else (args.op,)
+    if args.op == "all":
+        selected_ops = ops
+    else:
+        selected_ops = tuple(op.strip() for op in args.op.split(",") if op.strip())
     rows: list[list[str]] = []
     bench_rows = []
     failed = False