File size: 990 Bytes
dda32b8 bd11236 dda32b8 |
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 |
# /// script
# requires-python = ">=3.13"
# dependencies = [
# "kernels",
# "numpy",
# "torch",
# ]
# ///
import platform
from pathlib import Path
import kernels
import torch
# Load the locally built kernel
kernel = kernels.get_local_kernel(Path("build"), "__KERNEL_NAME_NORMALIZED__")
# Select device
if platform.system() == "Darwin":
device = torch.device("mps")
elif hasattr(torch, "xpu") and torch.xpu.is_available():
device = torch.device("xpu")
elif torch.version.cuda is not None and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
print(f"Using device: {device}")
# Create input tensor
x = torch.tensor([1.0, 2.0, 3.0], device=device)
print(f"Input: {x}")
# Run kernel (adds 1 to each element)
result = kernel.__KERNEL_NAME_NORMALIZED__(x)
print(f"Output: {result}")
# Verify result
expected = x + 1.0
assert torch.allclose(result, expected), "Kernel output doesn't match expected!"
print("Success!")
|