44 lines
1 KiB
Python
44 lines
1 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Minimal Quadrants CUDA kernel smoke run."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import quadrants as qd
|
||
|
|
|
||
|
|
|
||
|
|
@qd.kernel
|
||
|
|
def axpy(n: int, x: qd.types.NDArray[qd.f32, 1], y: qd.types.NDArray[qd.f32, 1]) -> None:
|
||
|
|
for i in range(n):
|
||
|
|
y[i] = 2.0 * x[i] + 1.0
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--arch", default="cuda", choices=("cuda", "gpu", "cpu", "x64"))
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
qd.init(arch=getattr(qd, args.arch))
|
||
|
|
|
||
|
|
x_np = np.arange(8, dtype=np.float32)
|
||
|
|
expected = 2.0 * x_np + 1.0
|
||
|
|
|
||
|
|
x = qd.ndarray(qd.f32, shape=x_np.shape)
|
||
|
|
y = qd.ndarray(qd.f32, shape=x_np.shape)
|
||
|
|
x.from_numpy(x_np)
|
||
|
|
|
||
|
|
axpy(x_np.size, x, y)
|
||
|
|
qd.sync()
|
||
|
|
|
||
|
|
actual = y.to_numpy()
|
||
|
|
if not np.array_equal(actual, expected):
|
||
|
|
raise SystemExit(f"Quadrants hello kernel mismatch: {actual} != {expected}")
|
||
|
|
|
||
|
|
print(f"quadrants hello world ok on arch={args.arch}: {actual}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|