import locomp
import numpy as np

@locomp.kernel
def vector_add(X: locomp.Tensor, Y: locomp.Tensor, O: locomp.Tensor, N: locomp.constexpr):
    i = locomp.program_id(0)
    locomp.store(O + i, locomp.load(X + i) + locomp.load(Y + i))

# Apple Silicon (Metal)
x = locomp.tensor(np.ones(4, dtype=np.float32))
y = locomp.tensor(np.ones(4, dtype=np.float32) * 2)
o = locomp.empty(4)
vector_add[(4,)](x, y, o, N=4)
print(o.numpy())  # [3. 3. 3. 3.]

# NVIDIA GPU (CUDA)
x = locomp.tensor(np.ones(4, dtype=np.float32), backend="cuda")
y = locomp.tensor(np.ones(4, dtype=np.float32) * 2, backend="cuda")
o = locomp.empty(4, backend="cuda")
vector_add[(4,)](x, y, o, N=4)
print(o.numpy())  # [3. 3. 3. 3.]

import locomp
import numpy as np

@locomp.kernel
def rms_norm(X: locomp.Tensor, W: locomp.Tensor, O: locomp.Tensor,
             N: locomp.constexpr, EPS: locomp.constexpr):
    row = locomp.program_id(0)
    tid = locomp.local_id(0)
    smem = locomp.shared_memory(32)

    local_sum = 0.0
    for i in range(tid, N, 128):
        val = locomp.load(X + row * N + i)
        local_sum = local_sum + val * val
    local_sum = locomp.simd_sum(local_sum)

    if locomp.simd_lane_id() == 0:
        locomp.shared_store(smem, locomp.simd_group_id(), local_sum)
    locomp.barrier()

    if tid == 0:
        total = 0.0
        for g in range(4):
            total = total + locomp.shared_load(smem, g)
        locomp.shared_store(smem, 0, locomp.rsqrt(total / N + EPS))
    locomp.barrier()

    rms = locomp.shared_load(smem, 0)
    for i in range(tid, N, 128):
        val = locomp.load(X + row * N + i)
        locomp.store(O + row * N + i, val * rms * locomp.load(W + i))

# Dispatch: 1 threadgroup per row, 128 threads per group
rows, dim = 128, 4096
x = locomp.tensor(np.random.randn(rows, dim).astype(np.float32))
w = locomp.tensor(np.ones(dim, dtype=np.float32))
o = locomp.empty((rows, dim))
rms_norm[(rows,), (128,)](x, w, o, N=dim, EPS=1e-5)

import locomp

@locomp.kernel
def fused_rope(Q: locomp.Tensor, Cos: locomp.Tensor, Sin: locomp.Tensor,
               N: locomp.constexpr, D: locomp.constexpr):
    bh = locomp.program_id(0)
    t  = locomp.program_id(1)
    for d in range(0, D // 2):
        qi = locomp.load(Q + bh * N * D + t * D + d)
        qj = locomp.load(Q + bh * N * D + t * D + d + D // 2)
        c  = locomp.load(Cos + t * (D // 2) + d)
        s  = locomp.load(Sin + t * (D // 2) + d)
        locomp.store(Q + bh * N * D + t * D + d,         qi * c - qj * s)
        locomp.store(Q + bh * N * D + t * D + d + D // 2, qi * s + qj * c)

class InferenceServer:
    def __init__(self, backend: str = "auto"):
        self.backend = backend  # "metal" | "cuda" | "riscv" | "auto"

    def run_rope(self, q_np, cos_np, sin_np, heads, seq, dim):
        q   = locomp.tensor(q_np,   backend=self.backend)
        cos = locomp.tensor(cos_np, backend=self.backend)
        sin = locomp.tensor(sin_np, backend=self.backend)
        # compiled once, cached forever — near-zero overhead on all subsequent calls
        fused_rope[(heads, seq)](q, cos, sin, N=seq, D=dim)
        return q.numpy()

@locomp.kernel
def gelu(X: locomp.Tensor, O: locomp.Tensor, N: locomp.constexpr):
    i = locomp.program_id(0)
    x = locomp.load(X + i)
    inner = locomp.clamp(0.7978845608 * (x + 0.044715 * x * x * x), -10.0, 10.0)
    locomp.store(O + i, 0.5 * x * (1.0 + locomp.tanh(inner)))

@locomp.kernel
def dequant_matvec(W_packed: locomp.Tensor, X: locomp.Tensor, Scales: locomp.Tensor,
                   O: locomp.Tensor, K: locomp.constexpr, GROUP: locomp.constexpr):
    row = locomp.program_id(0)
    acc = 0.0
    for k in range(K // 2):
        packed = locomp.cast(locomp.load(W_packed + row * (K // 2) + k), locomp.UInt8)
        w0 = locomp.cast(packed & 0xF, locomp.Float32) - 8.0
        w1 = locomp.cast((packed >> 4) & 0xF, locomp.Float32) - 8.0
        scale = locomp.load(Scales + row * (K // GROUP) + (2 * k) // GROUP)
        acc = acc + (w0 * scale) * locomp.load(X + 2 * k)
        acc = acc + (w1 * scale) * locomp.load(X + 2 * k + 1)
    locomp.store(O + row, acc)

from locomp import autotune, Config

@autotune(
    configs=[
        Config(BLOCK=16, num_warps=4),
        Config(BLOCK=32, num_warps=8),
        Config(BLOCK=64, num_warps=16),
    ],
    key=["N", "D"],
)
@locomp.kernel
def softmax(X: locomp.Tensor, O: locomp.Tensor,
            N: locomp.constexpr, D: locomp.constexpr, BLOCK: locomp.constexpr):
    ...
# Benchmarks all configs on first call · caches winner to ~/.cache/locomp/autotune.json

kernel[(N,)](args...)                          # 1D grid, 1 thread/group
kernel[(N,), (T,)](args...)                    # 1D grid, T threads/group
kernel[(gx, gy)](args...)                      # 2D grid
kernel[(gx, gy, gz), (tx, ty)](args...)        # 3D grid, 2D threadgroup

t = locomp.tensor(np_array)                   # numpy → GPU
t = locomp.tensor(np_array, backend="cuda")   # to specific backend
t = locomp.empty(shape)                        # uninitialized
t = locomp.zeros(shape)                        # zero-filled
t.numpy()                                      # GPU → numpy (auto-syncs)
t.reshape(new_shape)                           # zero-copy reshape
t.transpose(d0, d1)                            # swap dimensions
t.permute(*dims)                               # arbitrary reorder
t[slices]                                      # numpy-style slicing

a = locomp.ag.tensor(np.random.randn(N), requires_grad=True)
b = locomp.ag.tensor(np.random.randn(N), requires_grad=True)
loss = locomp.ag.sum(locomp.ag.mul(a, b))
locomp.ag.backward(loss)
print(a.grad)  # dL/da = b

ga = locomp.gpu_ag
x  = ga.tensor(np.random.randn(N), requires_grad=True)
y  = ga.relu(ga.exp(x))
ga.backward(ga.sum(y))
print(x.grad.numpy())   # GPU gradient, read back to CPU