Added tpmv

bench-level2.mojo

@@ -103,7 +103,7 @@ def parse_args(mut params: RunParams) -> Bool:
     if len(params.routines) == 0:
         params.routines = ["gemv", "gbmv", "ger", "sbmv", "symv", "spmv",
                            "syr", "syr2", "spr", "spr2", "tbmv", "tbsv",
-                           "trmv", "trsv"]
+                           "trmv", "trsv", "tpmv"]
 
     return True
 
@@ -504,6 +504,38 @@ def bench_tbsv[dtype: DType](n: Int, iters: Int, ctx: DeviceContext):
           "," + String(min_max_mean[0] * 1e-9) + "," + String(min_max_mean[1] * 1e-9) +
           "," + String(min_max_mean[2] * 1e-9) + "," + String(bw_gbs))
 
+def bench_tpmv[dtype: DType](n: Int, iters: Int, ctx: DeviceContext):
+    var ap_size = n * (n + 1) // 2
+    x_h = ctx.enqueue_create_host_buffer[dtype](n)
+    AP_h = ctx.enqueue_create_host_buffer[dtype](ap_size)
+    generate_random_arr[dtype](n, x_h.unsafe_ptr(), -1, 1)
+    generate_random_arr[dtype](ap_size, AP_h.unsafe_ptr(), -1, 1)
+    x_d = ctx.enqueue_create_buffer[dtype](n)
+    AP_d = ctx.enqueue_create_buffer[dtype](ap_size)
+    ctx.enqueue_copy(x_d, x_h)
+    ctx.enqueue_copy(AP_d, AP_h)
+    ctx.synchronize()
+
+    for _ in range(WARMUP):
+        blas_tpmv[dtype](1, False, 0, n, AP_d.unsafe_ptr(), x_d.unsafe_ptr(), 1, ctx)
+
+    var timings = List[Float32](length=iters, fill=0.0)
+
+    for i in range(iters):
+        start = monotonic()
+        blas_tpmv[dtype](1, False, 0, n, AP_d.unsafe_ptr(), x_d.unsafe_ptr(), 1, ctx)
+        end = monotonic()
+        timings[i] = Float32(end - start)
+
+    var min_max_mean = arr_min_max_mean(timings)
+
+    # bandwidth: read AP (ap_size) + read/write x (2n) = ap_size + 2n
+    var bw_gbs = Float32((ap_size + 2 * n) * bytes_per_elem(dtype)) / min_max_mean[2]
+
+    print("tpmv," + ctx.name() + "," + String(dtype) + "," + String(n) + "," + String(iters) +
+          "," + String(min_max_mean[0] * 1e-9) + "," + String(min_max_mean[1] * 1e-9) +
+          "," + String(min_max_mean[2] * 1e-9) + "," + String(bw_gbs))
+
 def bench_trmv[dtype: DType](n: Int, iters: Int, ctx: DeviceContext):
     print("trmv: not yet implemented")
 
@@ -533,6 +565,7 @@ def run_dtype[
         elif (routine == "spr"):  bench_spr[dtype](n, params.iters, ctx)
         elif (routine == "spr2"): bench_spr2[dtype](n, params.iters, ctx)
         elif (routine == "tbmv"): bench_tbmv[dtype](n, params.iters, ctx)
+        elif (routine == "tpmv"): bench_tpmv[dtype](n, params.iters, ctx)
         elif (routine == "tbsv"): bench_tbsv[dtype](n, params.iters, ctx)
         elif (routine == "trmv"): bench_trmv[dtype](n, params.iters, ctx)
         elif (routine == "trsv"): bench_trsv[dtype](n, params.iters, ctx)

src/level2/__init__.mojo

@@ -9,3 +9,4 @@ from .tbmv_device import *
 from .tbsv_device import *
 from .symv_device import *
 from .spr_device import *
+from .tpmv_device import *

src/level2/tpmv_device.mojo

@@ -0,0 +1,140 @@
+from gpu import thread_idx, block_idx, block_dim, grid_dim
+from gpu.host import DeviceContext
+from math import ceildiv
+
+comptime TBsize = 512
+
+# level2.tpmv
+# Performs matrix-vector multiplication of form
+#    x := A*x,   or   x := A**T*x
+# where x is a vector and A is an n by n unit or non-unit, upper or lower
+# triangular band matrix with (k+1) diagonals.
+fn stpmv_device(
+    uplo: Int,
+    trans: Int,
+    diag: Int,
+    n: Int,
+    AP: UnsafePointer[Float32, ImmutAnyOrigin],
+    x: UnsafePointer[Float32, ImmutAnyOrigin],
+    incx: Int,
+    workspace: UnsafePointer[Float32, MutAnyOrigin],
+):
+    var global_i = block_dim.x * block_idx.x + thread_idx.x
+    var n_threads = grid_dim.x * block_dim.x
+
+    for i in range(global_i, n, n_threads):
+        var sum = x[i * incx]
+        if trans:
+            if uplo:
+                var col_start = i * (i + 1) / 2
+                if not diag:
+                    sum *= AP[col_start + i]
+                for j in range(0, i):
+                    sum += AP[col_start + j] * x[j * incx]
+            else:
+                var col_start = i * n - ((i - 1) * i) / 2
+                if not diag:
+                    sum *= AP[col_start]
+                for j in range(i + 1, n):
+                    sum += AP[col_start + (j - i)] * x[j * incx]
+        else:
+            if uplo:
+                if not diag:
+                    sum *= AP[i * (i + 1) / 2 + i]
+                for j in range(i + 1, n):
+                    var index = j * (j + 1) / 2 + i
+                    sum += AP[index] * x[j * incx]
+            else:
+                if not diag:
+                    sum *= AP[i * n - ((i - 1) * i) / 2]
+                for j in range(0, i):
+                    var index = j * n - ((j - 1) * j) / 2 + (i - j)
+                    sum += AP[index] * x[j * incx]
+        workspace[i * incx] = sum
+
+
+fn dtpmv_device(
+    uplo: Int,
+    trans: Int,
+    diag: Int,
+    n: Int,
+    AP: UnsafePointer[Float64, ImmutAnyOrigin],
+    x: UnsafePointer[Float64, ImmutAnyOrigin],
+    incx: Int,
+    workspace: UnsafePointer[Float64, MutAnyOrigin],
+):
+    var global_i = block_dim.x * block_idx.x + thread_idx.x
+    var n_threads = grid_dim.x * block_dim.x
+
+    for i in range(global_i, n, n_threads):
+        var sum = x[i * incx]
+        if trans:
+            if uplo:
+                var col_start = i * (i + 1) / 2
+                if not diag:
+                    sum *= AP[col_start + i]
+                for j in range(0, i):
+                    sum += AP[col_start + j] * x[j * incx]
+            else:
+                var col_start = i * n - ((i - 1) * i) / 2
+                if not diag:
+                    sum *= AP[col_start]
+                for j in range(i + 1, n):
+                    sum += AP[col_start + (j - i)] * x[j * incx]
+        else:
+            if uplo:
+                if not diag:
+                    sum *= AP[i * (i + 1) / 2 + i]
+                for j in range(i + 1, n):
+                    var index = j * (j + 1) / 2 + i
+                    sum += AP[index] * x[j * incx]
+            else:
+                if not diag:
+                    sum *= AP[i * n - ((i - 1) * i) / 2]
+                for j in range(0, i):
+                    var index = j * n - ((j - 1) * j) / 2 + (i - j)
+                    sum += AP[index] * x[j * incx]
+        workspace[i * incx] = sum
+
+
+fn blas_tpmv[dtype: DType](
+    uplo: Int,
+    trans: Bool,
+    diag: Int,
+    n: Int,
+    d_AP: UnsafePointer[Scalar[dtype], ImmutAnyOrigin],
+    d_x: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    incx: Int,
+    ctx: DeviceContext,
+) raises:
+    # NOTE: add error checking here?
+
+    var trans_i = 1 if trans else 0
+
+    var workspace = ctx.enqueue_create_buffer[dtype](n)
+
+    @parameter
+    if dtype == DType.float32:
+        ctx.enqueue_function[stpmv_device, stpmv_device](
+            uplo, trans_i, diag,
+            n, d_AP,
+            d_x, incx,
+            workspace,
+            grid_dim=ceildiv(n, TBsize),
+            block_dim=TBsize,
+        )
+    elif dtype == DType.float64:
+        ctx.enqueue_function[dtpmv_device, dtpmv_device](
+            uplo, trans_i, diag,
+            n, d_AP,
+            d_x, incx,
+            workspace,
+            grid_dim=ceildiv(n, TBsize),
+            block_dim=TBsize,
+        )
+    else:
+        raise Error("blas_tpmv: Unsupported type")
+
+    ctx.enqueue_copy(d_x, workspace)
+
+    ctx.synchronize()

src/testing_utils/testing_utils.mojo

@@ -294,4 +294,26 @@ def arr_min_max_mean(
             a_max = a
         a_mean += a
     a_mean /= arr.__len__()
-    return (a_min, a_max, a_mean)
+    return (a_min, a_max, a_mean)
+
+
+fn dense_to_tri_packed[dtype: DType](
+    A_dense: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    A_packed: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    n: Int,
+    uplo: Int,
+):
+    var index = 0
+    for j in range(n):
+        if uplo:
+            for i in range(0, j+1):
+                A_packed[index] = A_dense[i * n + j]
+                index += 1
+        else:
+            for i in range(j, n):
+                A_packed[index] = A_dense[i * n + j]
+                index += 1
+
+    var n_packed = n * (n + 1) / 2
+    for i in range(index, n_packed):
+        A_packed[i] = 0

test-level2.mojo

@@ -1014,6 +1014,120 @@ def tbsv_test[
             )
             assert_true(ok)
 
+
+def tpmv_test[
+    dtype: DType,
+    n: Int,
+    trans: Bool,
+    uplo: Int,
+    diag: Int,
+]():
+    with DeviceContext() as ctx:
+        A_dense = ctx.enqueue_create_host_buffer[dtype](n * n)
+
+        AP_d = ctx.enqueue_create_buffer[dtype](n * n)
+        AP = ctx.enqueue_create_host_buffer[dtype](n * n)
+        x_d = ctx.enqueue_create_buffer[dtype](n)
+        x = ctx.enqueue_create_host_buffer[dtype](n)
+
+        generate_random_arr[dtype](n * n, A_dense.unsafe_ptr(), -1, 1)
+        generate_random_arr[dtype](n, x.unsafe_ptr(), -1, 1)
+
+        # Zero out off-triangle elements to make A strictly triangular
+        for i in range(n):
+            for j in range(n):
+                # upper triangular
+                if uplo == 0:
+                    if i > j:
+                        A_dense[i * n + j] = 0
+                # lower triangular
+                else:
+                    if i < j:
+                        A_dense[i * n + j] = 0
+
+        # Handle diagonal based on diag parameter
+        for i in range(n):
+            # unit diagonal
+            if diag == 1:
+                A_dense[i * n + i] = 1
+            # non-unit diagonal: make diagonally dominant to reduce numerical instability
+            else:
+                A_dense[i * n + i] += 1000
+
+        # Pack into column-major packed storage
+        dense_to_tri_packed(A_dense.unsafe_ptr(), AP.unsafe_ptr(), n, uplo)
+
+        ctx.enqueue_copy(AP_d, AP)
+        ctx.enqueue_copy(x_d, x)
+        ctx.synchronize()
+
+        # Compute norms for error checks
+        var norm_A = frobenius_norm[dtype](A_dense.unsafe_ptr(), n * n)
+        var norm_x = frobenius_norm[dtype](x.unsafe_ptr(), n)
+
+        blas_tpmv[dtype](
+            uplo,
+            trans,
+            diag,
+            n,
+            AP_d.unsafe_ptr(),
+            x_d.unsafe_ptr(), 1,
+            ctx,
+        )
+
+        # Import SciPy and numpy
+        sp = Python.import_module("scipy")
+        np = Python.import_module("numpy")
+        sp_blas = sp.linalg.blas
+
+        py_A = Python.list()
+        py_x = Python.list()
+        for i in range(n * n):
+            py_A.append(AP[i])
+        for i in range(n):
+            py_x.append(x[i])
+
+        var sp_res: PythonObject
+
+        if dtype == DType.float32:
+            np_A = np.array(py_A, dtype=np.float32)
+            np_x = np.array(py_x, dtype=np.float32)
+            sp_res = sp_blas.stpmv(n, np_A, np_x,
+                lower=0 if uplo else 1,
+                trans=1 if trans else 0,
+                diag=1 if diag else 0
+            )
+        elif dtype == DType.float64:
+            np_A = np.array(py_A, dtype=np.float64)
+            np_x = np.array(py_x, dtype=np.float64)
+            sp_res = sp_blas.dtpmv(n, np_A, np_x,
+                lower=0 if uplo else 1,
+                trans=1 if trans else 0,
+                diag=1 if diag else 0
+            )
+        else:
+            print("Unsupported type: ", dtype)
+            return
+
+        with x_d.map_to_host() as res_mojo:
+
+            var norm_diff = Scalar[dtype](0)
+            for i in range(n):
+                var diff = res_mojo[i] - Scalar[dtype](py=sp_res[i])
+                print(res_mojo[i], Scalar[dtype](py=sp_res[i]))
+                norm_diff += diff * diff
+            norm_diff = sqrt(norm_diff)
+            print(norm_diff)
+
+            var ok = check_gemm_error[dtype](
+                1, n, n,
+                Scalar[dtype](1), Scalar[dtype](0),
+                norm_A, norm_x, Scalar[dtype](0),
+                norm_diff
+            )
+            assert_true(ok)
+
+
 def test_gemv():
     gemv_test[DType.float32,  64,  64, False]()
     gemv_test[DType.float32,  64,  64, True]()
@@ -1158,6 +1272,24 @@ def test_tbsv():
     tbsv_test[DType.float64, 512, 16, 1, True,  0]()
     tbsv_test[DType.float64, 512, 16, 0, True,  0]()
 
+def test_tpmv():
+    tpmv_test[DType.float32,  64,  True,  0, 0]()
+    tpmv_test[DType.float32,  64,  True,  1, 0]()
+    tpmv_test[DType.float32,  64,  True,  0, 1]()
+    tpmv_test[DType.float32,  64,  True,  1, 1]()
+    tpmv_test[DType.float32,  64,  False, 0, 0]()
+    tpmv_test[DType.float32,  64,  False, 1, 0]()
+    tpmv_test[DType.float32,  64,  False, 0, 1]()
+    tpmv_test[DType.float32,  64,  False, 1, 1]()
+    tpmv_test[DType.float64,  64,  True,  0, 0]()
+    tpmv_test[DType.float64,  64,  True,  1, 0]()
+    tpmv_test[DType.float64,  64,  True,  0, 1]()
+    tpmv_test[DType.float64,  64,  True,  1, 1]()
+    tpmv_test[DType.float64,  64,  False, 0, 0]()
+    tpmv_test[DType.float64,  64,  False, 1, 0]()
+    tpmv_test[DType.float64,  64,  False, 0, 1]()
+    tpmv_test[DType.float64,  64,  False, 1, 1]()
+
 def main():
     print("--- MojoBLAS Level 2 routines testing ---")
     var args = argv()
@@ -1177,6 +1309,7 @@ def main():
         elif args[i] == "trsv":  suite.test[test_trsv]()
         elif args[i] == "tbmv":  suite.test[test_tbmv]()
         elif args[i] == "tbsv":  suite.test[test_tbsv]()
+        elif args[i] == "tpmv":  suite.test[test_tpmv]()
         elif args[i] == "symv":  suite.test[test_symv]()
         else: print("unknown routine:", args[i])
     suite^.run()