Spaces:
Running
Running
File size: 3,721 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """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""")
|