feat: optimize CuTe boundary check, implement TileLang copy kernel, use domestic CUTLASS mirror

ops/copy/nvidia/cute_copy.cuh

@@ -19,14 +19,8 @@ __global__ void copy_contiguous_cute_l40_kernel(T *dst, const T *src, int64_t n)
     auto block_shape = make_shape(Int<elements_per_block>{});
     auto block_coord = make_coord(blockIdx.x);
 
-    Tensor coords = make_identity_tensor(shape(src_tensor));
-    Tensor predicates = cute::lazy::transform(coords, [&](auto coord) {
-        return elem_less(coord, shape(src_tensor));
-    });
-
     Tensor src_tile = local_tile(src_tensor, block_shape, block_coord);
     Tensor dst_tile = local_tile(dst_tensor, block_shape, block_coord);
-    Tensor pred_tile = local_tile(predicates, block_shape, block_coord);
 
     Layout thread_layout = make_layout(make_shape(Int<Threads>{}));
     Layout value_layout = make_layout(make_shape(Int<ElementsPerAccess>{}));
@@ -40,7 +34,16 @@ __global__ void copy_contiguous_cute_l40_kernel(T *dst, const T *src, int64_t n)
 
     Tensor thread_src = thread_copy.partition_S(src_tile);
     Tensor thread_dst = thread_copy.partition_D(dst_tile);
-    Tensor thread_pred = thread_copy.partition_S(pred_tile);
+
+    Tensor coords = make_identity_tensor(shape(src_tensor));
+    Tensor coords_tile = local_tile(coords, block_shape, block_coord);
+    Tensor thread_coords = thread_copy.partition_S(coords_tile);
+    Tensor thread_pred = make_fragment_like<bool>(thread_src);
+    CUTE_UNROLL
+    for (int i = 0; i < size(thread_pred); ++i) {
+        thread_pred(i) = elem_less(thread_coords(i), shape(src_tensor));
+    }
+
     Tensor fragment = make_fragment_like(thread_src);
 
     copy_if(tiled_copy, thread_pred, thread_src, fragment);

ops/copy/tilelang/kernel.py

@@ -15,5 +15,11 @@ def copy_kernel(src, BLOCK_N: int, dtype):
     # 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.
+    # with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
+    with T.Kernel(T.ceildiv(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 + 1) * BLOCK_N]
+        )
 
     return out

scripts/build_nvidia.sh

@@ -12,7 +12,8 @@ CAMP_FORCE_RECONFIGURE="${CAMP_FORCE_RECONFIGURE:-0}"
 MODE="${1:-build}"
 
 CUTLASS_VERSION="${CUTLASS_VERSION:-v3.7.0}"
-CUTLASS_REPO="${CUTLASS_REPO:-https://github.com/NVIDIA/cutlass.git}"
+#CUTLASS_REPO="${CUTLASS_REPO:-https://github.com/NVIDIA/cutlass.git}"
+CUTLASS_REPO="${CUTLASS_REPO:-https://gitee.com/semidrive_1/cutlass.git}"
 
 ensure_third_party() {
   local tp_dir="${ROOT}/third_party"