API Reference

@locomp.kernel
def my_kernel(A: locomp.Tensor, B: locomp.Tensor, N: locomp.constexpr):
    i = locomp.program_id(0)
    locomp.store(B + i, locomp.load(A + i) * 2.0)

# Dispatch: grid=(N,), threadgroup size default 1
my_kernel[(N,)](a, b, N=N)

# With explicit threadgroup size
my_kernel[(N,), (128,)](a, b, N=N)

from locomp import autotune, Config

@autotune(
    configs=[Config(BLOCK_M=16), Config(BLOCK_M=32), Config(BLOCK_M=64)],
    key=["M", "N"],
)
@locomp.kernel
def my_kernel(..., BLOCK_M: locomp.constexpr):
    ...

# Context manager — runs on exit
with locomp.graph() as g:
    g.add(rms_norm, (rows,), (128,), x, w, h, N=dim, eps=1e-5)
    g.add(matmul,   (M * N,),        h, w2, out, M=M, N=K, K=dim)

# Manual build + run (re-runnable)
g = locomp.graph()
g.add(kernel_a, (N,), x, tmp, N=N)
g.add(kernel_b, (N,), tmp, out, N=N)
g.run()   # first run
g.run()   # re-run same sequence on same tensors

with locomp.profile() as p:
    rms_norm[(rows,), (128,)](x, w, out, N=dim, eps=1e-5)
    matmul[(M * N,)](a, b, c, M=M, N=N, K=K)

print(p.report())
# Kernel          Grid     ms
# rms_norm        (128,)   0.245
# matmul          (4096,)  1.103

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)  # on GPU
y  = ga.relu(ga.exp(x))
ga.backward(ga.sum(y))
print(x.grad.numpy())

# Supported ops: add sub mul div exp log relu pow
#                sigmoid tanh sum mean matvec matmul softmax cross_entropy

if locomp.cuda_available():
    print(f"{locomp.cuda_device_count()} CUDA GPUs available")
    t = locomp.tensor(np.ones(1024, np.float32), backend="cuda")
    my_kernel[(1024,)](t, out, N=1024)
    print(out.numpy())