spr implementation

src/level2/__init__.mojo

@@ -6,3 +6,4 @@ from .gbmv_device import *
 from .sbmv_device import *
 from .trsv_device import *
 from .symv_device import *
+from .spr_device import *

src/level2/spr_device.mojo

@@ -0,0 +1,100 @@
+from gpu import thread_idx, block_idx, block_dim, grid_dim
+from gpu.host import DeviceContext
+from math import ceildiv
+
+comptime TBsize = 512
+
+# level2.spr
+# Performs symmetric rank-1 update of packed symmetric matrix:
+#   A := alpha*x*x**T + A
+# where A is stored in packed format (column-major).
+# uplo: 0 = upper triangle, 1 = lower triangle
+#
+# Upper triangular packed storage (column-major):
+#   AP[j*(j+1)//2 + i] = A[i,j] for 0 <= i <= j < n
+#
+# Lower triangular packed storage (column-major):
+#   AP[j*(2*n-j+1)//2 + (i-j)] = A[i,j] for 0 <= j <= i < n
+#
+# Each thread handles a unique row i, updating all relevant AP entries
+# for that row with no data race.
+fn sspr_device(
+    uplo: Int,
+    n: Int,
+    alpha: Float32,
+    x: UnsafePointer[Float32, ImmutAnyOrigin],
+    incx: Int,
+    AP: UnsafePointer[Float32, MutAnyOrigin],
+):
+    var global_i = block_dim.x * block_idx.x + thread_idx.x
+    var n_threads = grid_dim.x * block_dim.x
+
+    # upper triangle: AP[j*(j+1)/2 + i] for j in range [i, n)
+    if not uplo:
+        for i in range(global_i, n, n_threads):
+            var xi = alpha * x[i * incx]
+            for j in range(i, n):
+                AP[j * (j + 1) // 2 + i] += xi * x[j * incx]
+    # lower triangle: AP[j*(2*n-j+1)/2 + (i-j)] for j in range [0, i]
+    else:
+        for i in range(global_i, n, n_threads):
+            var xi = alpha * x[i * incx]
+            for j in range(0, i + 1):
+                AP[j * (2 * n - j + 1) // 2 + (i - j)] += xi * x[j * incx]
+
+
+fn dspr_device(
+    uplo: Int,
+    n: Int,
+    alpha: Float64,
+    x: UnsafePointer[Float64, ImmutAnyOrigin],
+    incx: Int,
+    AP: UnsafePointer[Float64, MutAnyOrigin],
+):
+    var global_i = block_dim.x * block_idx.x + thread_idx.x
+    var n_threads = grid_dim.x * block_dim.x
+
+    # upper triangle: AP[j*(j+1)/2 + i] for j in range [i, n)
+    if not uplo:
+        for i in range(global_i, n, n_threads):
+            var xi = alpha * x[i * incx]
+            for j in range(i, n):
+                AP[j * (j + 1) // 2 + i] += xi * x[j * incx]
+    # lower triangle: AP[j*(2*n-j+1)/2 + (i-j)] for j in range [0, i]
+    else:
+        for i in range(global_i, n, n_threads):
+            var xi = alpha * x[i * incx]
+            for j in range(0, i + 1):
+                AP[j * (2 * n - j + 1) // 2 + (i - j)] += xi * x[j * incx]
+
+
+fn blas_spr[dtype: DType](
+    uplo: Int,
+    n: Int,
+    alpha: Scalar[dtype],
+    d_x: UnsafePointer[Scalar[dtype], ImmutAnyOrigin],
+    incx: Int,
+    d_AP: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    ctx: DeviceContext,
+) raises:
+    @parameter
+    if dtype == DType.float32:
+        ctx.enqueue_function[sspr_device, sspr_device](
+            uplo, n,
+            alpha, d_x, incx,
+            d_AP,
+            grid_dim=ceildiv(n, TBsize),
+            block_dim=TBsize,
+        )
+    elif dtype == DType.float64:
+        ctx.enqueue_function[dspr_device, dspr_device](
+            uplo, n,
+            alpha, d_x, incx,
+            d_AP,
+            grid_dim=ceildiv(n, TBsize),
+            block_dim=TBsize,
+        )
+    else:
+        raise Error("blas_spr: Unsupported type")
+
+    ctx.synchronize()

test-level2.mojo

@@ -223,6 +223,59 @@ def syr_test[
             for i in range(n * n):
                 assert_almost_equal(res_mojo[i], Scalar[dtype](py=sp_flat[i]), atol=atol)
 
+def spr_test[
+    dtype: DType,
+    n: Int,
+    uplo: Int,
+]():
+    comptime ap_len = n * (n + 1) // 2
+
+    with DeviceContext() as ctx:
+        AP_d = ctx.enqueue_create_buffer[dtype](ap_len)
+        AP = ctx.enqueue_create_host_buffer[dtype](ap_len)
+        x_d = ctx.enqueue_create_buffer[dtype](n)
+        x = ctx.enqueue_create_host_buffer[dtype](n)
+
+        generate_random_arr[dtype](ap_len, AP.unsafe_ptr(), -100, 100)
+        generate_random_arr[dtype](n, x.unsafe_ptr(), -100, 100)
+
+        ctx.enqueue_copy(AP_d, AP)
+        ctx.enqueue_copy(x_d, x)
+        ctx.synchronize()
+
+        var alpha = generate_random_scalar[dtype](-100, 100)
+
+        blas_spr[dtype](uplo, n, alpha, x_d.unsafe_ptr(), 1, AP_d.unsafe_ptr(), ctx)
+
+        # Import SciPy and numpy
+        sp = Python.import_module("scipy")
+        np = Python.import_module("numpy")
+        sp_blas = sp.linalg.blas
+
+        py_AP = Python.list()
+        py_x = Python.list()
+        for i in range(ap_len):
+            py_AP.append(AP[i])
+        for i in range(n):
+            py_x.append(x[i])
+
+        var sp_res: PythonObject
+        if dtype == DType.float32:
+            np_AP = np.array(py_AP, dtype=np.float32)
+            np_x = np.array(py_x, dtype=np.float32)
+            sp_res = sp_blas.sspr(alpha, np_x, lower=uplo, ap=np_AP, overwrite_ap=False)
+        elif dtype == DType.float64:
+            np_AP = np.array(py_AP, dtype=np.float64)
+            np_x = np.array(py_x, dtype=np.float64)
+            sp_res = sp_blas.dspr(alpha, np_x, lower=uplo, ap=np_AP, overwrite_ap=False)
+        else:
+            print("Unsupported type: ", dtype)
+            return
+
+        with AP_d.map_to_host() as res_mojo:
+            for i in range(ap_len):
+                assert_almost_equal(res_mojo[i], Scalar[dtype](py=sp_res[i]), atol=atol)
+
 def syr2_test[
     dtype: DType,
     n: Int,
@@ -780,6 +833,12 @@ def test_ger():
     ger_test[DType.float64, 64, 64]()
     ger_test[DType.float64, 256, 256]()
 
+def test_spr():
+    spr_test[DType.float32,  256, 1]()
+    spr_test[DType.float32, 1024, 0]()
+    spr_test[DType.float64,  256, 0]()
+    spr_test[DType.float64, 1024, 1]()
+
 def test_syr():
     syr_test[DType.float32,  256, 1]()
     syr_test[DType.float32, 1024, 0]()
@@ -860,6 +919,7 @@ def main():
         if   args[i] == "gemv":  suite.test[test_gemv]()
         elif args[i] == "ger":   suite.test[test_ger]()
         elif args[i] == "sbmv":  suite.test[test_sbmv]()
+        elif args[i] == "spr":   suite.test[test_spr]()
         elif args[i] == "syr":   suite.test[test_syr]()
         elif args[i] == "syr2":  suite.test[test_syr2]()
         elif args[i] == "gbmv":  suite.test[test_gbmv]()