Spaces:
Running
Running
| """Candidate sources shared by tests (ReLU task).""" | |
| HEADER = """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import triton | |
| import triton.language as tl | |
| """ | |
| KERNEL = """ | |
| @triton.jit | |
| def relu_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr): | |
| offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) | |
| mask = offs < n | |
| x = tl.load(x_ptr + offs, mask=mask) | |
| tl.store(y_ptr + offs, tl.maximum(x, 0.0), mask=mask) | |
| """ | |
| GOOD = HEADER + KERNEL + """ | |
| class ModelNew(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| def forward(self, x): | |
| x = x.contiguous() | |
| y = torch.empty_like(x) | |
| n = x.numel() | |
| relu_kernel[(triton.cdiv(n, 1024),)](x, y, n, BLOCK=1024) | |
| return y | |
| """ | |
| def with_forward(body: str, extra: str = "", header: str = HEADER) -> str: | |
| """ReLU kernel module whose ModelNew.forward body is replaced (4-space indented lines).""" | |
| lines = "\n".join(" " + ln for ln in body.strip().splitlines()) | |
| return header + extra + KERNEL + f""" | |
| class ModelNew(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.act = nn.ReLU() | |
| def forward(self, x): | |
| {lines} | |
| """ | |
| WRONG = GOOD.replace("tl.maximum(x, 0.0)", "x") # identity, not relu | |
| ALIAS = with_forward("relu_kernel[(triton.cdiv(x.numel(), 1024),)](x, x, x.numel(), BLOCK=1024)\nreturn x") | |
| HANG = with_forward("while True:\n pass") | |
| CUDA_GOOD = """ | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.cpp_extension import load_inline | |
| src = ''' | |
| __global__ void relu_k(const float* x, float* y, int n) { | |
| int i = blockIdx.x * blockDim.x + threadIdx.x; | |
| if (i < n) y[i] = x[i] > 0.f ? x[i] : 0.f; | |
| } | |
| torch::Tensor relu(torch::Tensor x) { | |
| auto y = torch::empty_like(x); | |
| int n = x.numel(); | |
| relu_k<<<(n + 255) / 256, 256>>>(x.data_ptr<float>(), y.data_ptr<float>(), n); | |
| return y; | |
| } | |
| ''' | |
| mod = load_inline(name="relu_ext", cpp_sources="torch::Tensor relu(torch::Tensor x);", | |
| cuda_sources=src, functions=["relu"]) | |
| class ModelNew(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| def forward(self, x): | |
| return mod.relu(x.contiguous()) | |
| """ | |
| MATMUL = """ | |
| import torch | |
| import torch.nn as nn | |
| import triton | |
| import triton.language as tl | |
| @triton.jit | |
| def mm_kernel(a, b, c, M, N, K, sam, sak, sbk, sbn, scm, scn, | |
| BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): | |
| rm = tl.program_id(0) * BM + tl.arange(0, BM) | |
| rn = tl.program_id(1) * BN + tl.arange(0, BN) | |
| rk = tl.arange(0, BK) | |
| acc = tl.zeros((BM, BN), dtype=tl.float32) | |
| for k in range(0, K, BK): | |
| x = tl.load(a + rm[:, None] * sam + (k + rk)[None, :] * sak, | |
| mask=(rm[:, None] < M) & ((k + rk)[None, :] < K), other=0.0) | |
| y = tl.load(b + (k + rk)[:, None] * sbk + rn[None, :] * sbn, | |
| mask=((k + rk)[:, None] < K) & (rn[None, :] < N), other=0.0) | |
| acc += tl.dot(x, y) | |
| tl.store(c + rm[:, None] * scm + rn[None, :] * scn, acc, mask=(rm[:, None] < M) & (rn[None, :] < N)) | |
| class ModelNew(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| def forward(self, A, B): | |
| M, K = A.shape | |
| N = B.shape[1] | |
| C = torch.empty((M, N), device=A.device, dtype=torch.float32) | |
| grid = (triton.cdiv(M, 128), triton.cdiv(N, 128)) | |
| mm_kernel[grid](A, B, C, M, N, K, A.stride(0), A.stride(1), B.stride(0), B.stride(1), | |
| C.stride(0), C.stride(1), BM=128, BN=128, BK=32) | |
| return C | |
| """ | |
| # correct but ~seconds per call: 100k sequential launches | |
| SLOW = with_forward("""y = torch.empty_like(x) | |
| for _ in range(100000): | |
| relu_kernel[(triton.cdiv(x.numel(), 1024),)](x, y, x.numel(), BLOCK=1024) | |
| return y""") | |