submission 117551

- from utils import make_match_reference, DeterministicContext
- import torch
- from task import input_t, output_t
- from torch.utils.cpp_extension import load_inline
-  
-  
- cpp_stub = r"""
- #include <torch/extension.h>
- void vector_add(torch::Tensor a, torch::Tensor b, torch::Tensor out);
- """
-  
- cuda_src = r"""
- #include <torch/extension.h>
- #include <cuda.h>
- #include <cuda_runtime.h>
- #include <cuda_fp16.h>
- #include <ATen/cuda/CUDAContext.h>
-  
- __global__ void vector_add_half2_kernel(const __half* __restrict__ a,
-                                         const __half* __restrict__ b,
-                                         __half* __restrict__ c,
-                                         int n) {
-     int idx2 = blockIdx.x * blockDim.x + threadIdx.x;  // half2 单位索引
-     int i = idx2 * 2;
-  
-     // 主路径：half2
-     if (i + 1 < n) {
-         const __half2* a2 = reinterpret_cast<const __half2*>(a);
-         const __half2* b2 = reinterpret_cast<const __half2*>(b);
-               __half2* c2 = reinterpret_cast<__half2*>(c);
-         __half2 ah = a2[idx2];
-         __half2 bh = b2[idx2];
-         c2[idx2] = __hadd2(ah, bh);
-     }
-  
-     // 尾元素（总元素为奇数时）
-     if ((n & 1) && idx2 == (n >> 1)) {
-         c[n - 1] = __hadd(a[n - 1], b[n - 1]);
-     }
- }
-  
- void vector_add(torch::Tensor a, torch::Tensor b, torch::Tensor out) {
-     TORCH_CHECK(a.is_cuda() && b.is_cuda() && out.is_cuda(), "tensors must be on CUDA");
-     TORCH_CHECK(a.dtype() == torch::kFloat16 && b.dtype() == torch::kFloat16 && out.dtype() == torch::kFloat16,
-                 "only float16 supported");
-     TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && out.is_contiguous(),
-                 "tensors must be contiguous");
-     TORCH_CHECK(a.sizes() == b.sizes() && a.sizes() == out.sizes(), "shape mismatch");
-     TORCH_CHECK(a.dim() == 1 || a.dim() == 2, "expect 1D or 2D tensors");
-  
-     // 线性化处理（1D/2D 通用）
-     int n = a.numel();
-     int threads = 256;
-     int pairs = (n + 1) / 2;  // half2 元素数
-     int blocks = (pairs + threads - 1) / threads;
-  
-     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-     vector_add_half2_kernel<<<blocks, threads, 0, stream>>>(
-         reinterpret_cast<const __half*>(a.data_ptr<at::Half>()),
-         reinterpret_cast<const __half*>(b.data_ptr<at::Half>()),
-         reinterpret_cast<__half*>(out.data_ptr<at::Half>()),
-         n);
-     TORCH_CHECK(cudaGetLastError() == cudaSuccess, "kernel launch failed");
- }
- """
-  
- major, minor = torch.cuda.get_device_capability()
- # major, minor = 8, 9
- arch_flag = f"-arch=sm_{major}{minor}"
-  
- ext = load_inline(
-     name="vecadd_half2_2d_inline",
-     cpp_sources=[cpp_stub],
-     cuda_sources=[cuda_src],
-     functions=["vector_add"],
-     extra_cuda_cflags=[arch_flag],
-     verbose=True,
- )
-  
-  
- # def ref_kernel(data: input_t) -> output_t:
- #     """
- #     Reference implementation of vector addition using PyTorch.
- #     Args:
- #         data: Tuple of tensors [A, B] to be added.
- #     Returns:
- #         Tensor containing element-wise sums.
- #     """
- #     with DeterministicContext():
- #         A, B, output = data
- #         output[...] = A + B
-  
- #         return output
-  
-  
- def custom_kernel(data: input_t) -> output_t:
-     """
-     Reference implementation of vector addition using PyTorch.
-     Args:
-         data: Tuple of tensors [A, B] to be added.
-     Returns:
-         Tensor containing element-wise sums.
-     """
-     with DeterministicContext():
-         A, B, output = data
-         ext.vector_add(A, B, output)
-         return output
-  
-  
- # def generate_input(size: int, seed: int) -> input_t:
- #     """
- #     Generates random input tensors of specified shapes.
- #     Returns:
- #         Tuple of tensors [A, B] to be added.
- #     """
- #     gen = torch.Generator(device="cuda")
- #     gen.manual_seed(seed)
- #     A = torch.randn(
- #         size, size, device="cuda", dtype=torch.float16, generator=gen
- #     ).contiguous()
- #     B = torch.randn(
- #         size, size, device="cuda", dtype=torch.float16, generator=gen
- #     ).contiguous()
- #     C = torch.empty(size, size, device="cuda", dtype=torch.float16).contiguous()
- #     return A, B, C
-  
-  
- # check_implementation = make_match_reference(ref_kernel)
   No newline at end of file
+ from __future__ import annotations
+ 
+ import os
+ 
+ os.system("pip install tilelang")
+ 
+ import importlib
+ import subprocess
+ import sys
+ from typing import Dict, Tuple
+ 
+ import torch
+ 
+ 
+ # def ensure_pkg(pkg, extra_args=None):
+ #     try:
+ #         importlib.import_module(pkg)
+ #     except ImportError:
+ #         args = [sys.executable, "-m", "pip", "install", pkg]
+ #         if extra_args:
+ #             args.extend(extra_args)
+ #         subprocess.check_call(args)
+ 
+ 
+ # ensure_pkg("tilelang", ["-q"])
+ # ensure_pkg("einops", ["-q"])
+ 
+ 
+ import tilelang  # noqa: E402
+ import tilelang.language as T  # noqa: E402
+ 
+ TensorTriplet = Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+ 
+ def _tile_shape(numel: int) -> Tuple[int, int]:
+     return 1, numel
+ 
+ @tilelang.jit
+ def elementwise_add(
+     M,
+     N,
+     dt,
+     block_M,
+     block_N,
+     threads,
+ ):
+ 
+     @T.prim_func
+     def tlk(
+             A: T.Tensor((M, N), dt),
+             B: T.Tensor((M, N), dt),
+             C: T.Tensor((M, N), dt),
+     ):
+         with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
+             start_x = bx * block_N
+             start_y = by * block_M
+ 
+             for (local_y, local_x) in T.Parallel(block_M, block_N):
+                 y = start_y + local_y
+                 x = start_x + local_x
+ 
+                 C[y, x] = A[y, x] + B[y, x]
+ 
+     return tlk
+ 
+ 
+ @tilelang.jit
+ def tl_vecadd(M, N, dtype="float16", block_M=128, block_N=128, threads=128):
+ 
+     @T.prim_func
+     def kernel(
+         A: T.Tensor((M, N), dtype),
+         B: T.Tensor((M, N), dtype),
+         C: T.Tensor((M, N), dtype),
+     ):
+         with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
+             A_shared = T.alloc_shared((block_M, block_N), dtype)
+             B_shared = T.alloc_shared((block_M, block_N), dtype)
+             C_local = T.alloc_fragment((block_M, block_N), dtype)
+             C_shared = T.alloc_shared((block_M, block_N), dtype)
+ 
+             T.copy(A[by * block_M, bx * block_N], A_shared)
+             T.copy(B[by * block_M, bx * block_N], B_shared)
+ 
+             for (ly, lx) in T.Parallel(block_M, block_N):
+                 C_local[ly, lx] = A_shared[ly, lx] + B_shared[ly, lx]
+ 
+             T.copy(C_local, C_shared)
+             T.copy(C_shared, C[by * block_M, bx * block_N])
+ 
+     return kernel
+ 
+ 
+ _kernel_cache: Dict[Tuple[int, int], tilelang.runtime.Module] = {}
+ 
+ 
+ def custom_kernel(data: TensorTriplet) -> torch.Tensor:
+     A, B, output = data
+     rows, cols = A.shape
+     # rows, cols = _tile_shape(A.numel())
+     key = (rows, cols)
+     if key not in _kernel_cache:
+         # import pdb; pdb.set_trace()
+         # _kernel_cache[key] = elementwise_add(rows, cols, "float16", 128, 128, 128) # tl_vecadd(rows, cols)
+         _kernel_cache[key] = tl_vecadd(rows, cols)
+     kernel = _kernel_cache[key]
+     kernel(A, B, output)
+     return output
   No newline at end of file