Helion_example_blog.txt

txt

import torch
import helion
import helion.language as hl
import time

# ============================================================
# Example 1: Matrix Multiplication
# ============================================================
@helion.kernel()
def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    m, k = x.size()
    k2, n = y.size()
    out = torch.empty([m, n], dtype=x.dtype, device=x.device)
    for tile_m, tile_n in hl.tile([m, n]):
        acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n])
        out[tile_m, tile_n] = acc
    return out


# ============================================================
# Example 2: Vector Addition (elementwise)
# ============================================================
@helion.kernel()
def add_kernel(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    size = x.size(0)
    out = torch.empty_like(x)
    for tile in hl.tile(size):
        out[tile] = x[tile] + y[tile]
    return out


# ============================================================
# Example 3: Softmax (from official Helion examples)
# ============================================================
@helion.kernel()
def softmax_kernel(x: torch.Tensor) -> torch.Tensor:
    n, _m = x.size()
    out = torch.empty_like(x)
    for tile_n in hl.tile(n):
        values = x[tile_n, :]
        amax = torch.amax(values, dim=1, keepdim=True)
        exp = torch.exp(values - amax)
        sum_exp = torch.sum(exp, dim=1, keepdim=True)
        out[tile_n, :] = exp / sum_exp
    return out


if __name__ == "__main__":
    print("=" * 60)
    print("HELION LIVE CODE EXAMPLES ON NVIDIA H200")
    print("=" * 60)

    # --- Example 1: Matrix Multiplication ---
    print("\n[1] MATRIX MULTIPLICATION")
    M, K, N = 512, 256, 512
    x = torch.randn(M, K, device="cuda", dtype=torch.float16)
    y = torch.randn(K, N, device="cuda", dtype=torch.float16)

    result = matmul(x, y)
    expected = torch.mm(x.float(), y.float()).half()

    print(f"    Shapes: ({M}x{K}) @ ({K}x{N}) -> {list(result.shape)}")
    print(f"    Max error: {(result.float() - expected.float()).abs().max().item():.6f}")
    print(f"    Correct: {torch.allclose(result.float(), expected.float(), atol=1e-1, rtol=1e-2)}")

    # --- Example 2: Vector Addition ---
    print("\n[2] VECTOR ADDITION")
    size = 1_000_000
    a = torch.randn(size, device="cuda", dtype=torch.float32)
    b = torch.randn(size, device="cuda", dtype=torch.float32)

    result_add = add_kernel(a, b)
    expected_add = a + b

    print(f"    Size: {size:,}")
    print(f"    Max error: {(result_add - expected_add).abs().max().item():.8f}")
    print(f"    Correct: {torch.allclose(result_add, expected_add)}")

    # --- Example 3: Softmax ---
    print("\n[3] SOFTMAX")
    rows, cols = 256, 1024
    x_soft = torch.randn(rows, cols, device="cuda", dtype=torch.float32)

    result_soft = softmax_kernel(x_soft)
    expected_soft = torch.softmax(x_soft, dim=-1)

    print(f"    Shape: ({rows}x{cols})")
    print(f"    Max error: {(result_soft - expected_soft).abs().max().item():.8f}")
    print(f"    Correct: {torch.allclose(result_soft, expected_soft, atol=1e-3)}")
    print(f"    Row sums (should be ~1.0): {result_soft.sum(dim=-1)[:5].tolist()}")

    print("\n" + "=" * 60)
    print("ALL EXAMPLES COMPLETED SUCCESSFULLY")
    print("=" * 60)