sbmv implementation

src/level2/__init__.mojo

@@ -3,4 +3,5 @@ from .ger_device import *
 from .syr_device import *
 from .syr2_device import *
 from .gbmv_device import *
+from .sbmv_device import *
 from .trsv_device import *

src/level2/sbmv_device.mojo

@@ -0,0 +1,137 @@
+from gpu import thread_idx, block_idx, block_dim, grid_dim
+from gpu.host import DeviceContext
+from math import ceildiv
+
+comptime TBsize = 512
+
+# level2.sbmv
+# Performs symmetric band matrix-vector multiplication
+#    y := alpha*A*x + beta*y,
+# where A is an n by n symmetric band matrix with k off-diagonals.
+
+fn ssbmv_device(
+    uplo: Int,
+    n: Int,
+    k: Int,
+    alpha: Float32,
+    A: UnsafePointer[Float32, ImmutAnyOrigin],
+    lda: Int,
+    x: UnsafePointer[Float32, ImmutAnyOrigin],
+    incx: Int,
+    beta: Float32,
+    y: UnsafePointer[Float32, MutAnyOrigin],
+    incy: Int,
+):
+    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 = Scalar[DType.float32](0)
+
+        var j_start = max(0, i - k)
+        var j_end = min(n - 1, i + k)
+
+        for j in range(j_start, j_end + 1):
+            var val: Float32
+
+            if uplo: # upper
+                if j >= i:
+                    val = A[i * lda + (j - i)]
+                else:
+                    val = A[j * lda + (i - j)]
+            else: # lower
+                if j <= i:
+                    val = A[i * lda + (i - j)]
+                else:
+                    val = A[j * lda + (j - i)]
+
+            sum += val * x[j * incx]
+
+        y[i * incy] = alpha * sum + beta * y[i * incy]
+
+
+fn dsbmv_device(
+    uplo: Int,
+    n: Int,
+    k: Int,
+    alpha: Float64,
+    A: UnsafePointer[Float64, ImmutAnyOrigin],
+    lda: Int,
+    x: UnsafePointer[Float64, ImmutAnyOrigin],
+    incx: Int,
+    beta: Float64,
+    y: UnsafePointer[Float64, MutAnyOrigin],
+    incy: Int,
+):
+    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 = Scalar[DType.float64](0)
+
+        var j_start = max(0, i - k)
+        var j_end = min(n - 1, i + k)
+
+        for j in range(j_start, j_end + 1):
+            var val: Float64
+
+            if uplo: # upper
+                if j >= i:
+                    val = A[i * lda + (j - i)]
+                else:
+                    val = A[j * lda + (i - j)]
+            else: # lower
+                if j <= i:
+                    val = A[i * lda + (i - j)]
+                else:
+                    val = A[j * lda + (j - i)]
+
+            sum += val * x[j * incx]
+
+        y[i * incy] = alpha * sum + beta * y[i * incy]
+
+
+fn blas_sbmv[dtype: DType](
+    uplo: Int,
+    n: Int,
+    k: Int,
+    alpha: Scalar[dtype],
+    d_A: UnsafePointer[Scalar[dtype], ImmutAnyOrigin],
+    lda: Int,
+    d_x: UnsafePointer[Scalar[dtype], ImmutAnyOrigin],
+    incx: Int,
+    beta: Scalar[dtype],
+    d_y: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    incy: Int,
+    ctx: DeviceContext,
+) raises:
+
+    # TODO:
+    # check n > 0
+    # check k >= 0
+    # check lda >= k + 1
+    # check incx, incy > 0
+
+    @parameter
+    if dtype == DType.float32:
+        ctx.enqueue_function[ssbmv_device, ssbmv_device](
+            uplo, n, k,
+            alpha, d_A, lda,
+            d_x, incx,
+            beta, d_y, incy,
+            grid_dim=ceildiv(n, TBsize),
+            block_dim=TBsize,
+        )
+    elif dtype == DType.float64:
+        ctx.enqueue_function[dsbmv_device, dsbmv_device](
+            uplo, n, k,
+            alpha, d_A, lda,
+            d_x, incx,
+            beta, d_y, incy,
+            grid_dim=ceildiv(n, TBsize),
+            block_dim=TBsize,
+        )
+    else:
+        raise Error("blas_sbmv: Unsupported type")
+
+    ctx.synchronize()

src/testing_utils/testing_utils.mojo

@@ -155,3 +155,103 @@ fn dense_to_band[dtype: DType](
                 B[i * band_width + band_col] = A[i * n + j]
             else:
                 A[i * n + j] = Scalar[dtype](0)
+
+# Packs row-major dense matrix to row-major band buffer
+fn dense_to_sym_band_rm[dtype: DType](
+    A_dense: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    A_band: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    n: Int,
+    k: Int,
+    upper: Bool,
+):
+    var lda = k + 1
+
+    # zero initialize
+    for i in range(n):
+        for b in range(lda):
+            A_band[i * lda + b] = 0
+
+    if upper:
+        # store A[i,j] for j >= i
+        for i in range(n):
+            var j_end = (i + k) if (i + k < n - 1) else (n - 1)
+            for j in range(i, j_end + 1):
+                var band_col = j - i
+                A_band[i * lda + band_col] = A_dense[i * n + j]
+    else:
+        # store A[i,j] for j <= i
+        for i in range(n):
+            var j_start = (i - k) if (i - k > 0) else 0
+            for j in range(j_start, i + 1):
+                var band_col = i - j
+                A_band[i * lda + band_col] = A_dense[i * n + j]
+
+    # Overwrite original matrix with band reconstruction
+    for i in range(n):
+        for j in range(n):
+            var val = Scalar[dtype](0)
+
+            if upper:
+                if j >= i and j <= i + k:
+                    val = A_band[i * lda + (j - i)]
+                elif i >= j and i <= j + k:
+                    # symmetric mirror
+                    val = A_band[j * lda + (i - j)]
+            else:
+                if j <= i and j >= i - k:
+                    val = A_band[i * lda + (i - j)]
+                elif j >= i and j <= i + k:
+                    # symmetric mirror
+                    val = A_band[j * lda + (j - i)]
+
+            A_dense[i * n + j] = val
+
+# Packs row-major dense matrix to column-major band buffer
+fn dense_to_sym_band_cm[dtype: DType](
+    A_dense: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    A_band: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    n: Int,
+    k: Int,
+    upper: Bool,
+):
+    var lda = k + 1
+
+    # zero initialize
+    for i in range(n):
+        for b in range(lda):
+            A_band[i * lda + b] = 0
+
+    if upper:
+        # store A[i,j] for j >= i
+        for j in range(n):
+            var m = k - j
+            var i_start = j - k if (j - k > 0) else 0
+            for i in range(i_start, j + 1):
+                A_band[j * lda + m + i] = A_dense[j * n + i]
+    else:
+        # store A[i,j] for j <= i
+        for j in range(n):
+            var i_end = (j + k) if (j + k < n - 1) else (n - 1)
+            for i in range(j, i_end + 1):
+                var band_col = i - j
+                A_band[j * lda + band_col] = A_dense[i * n + j]
+
+    # Overwrite original matrix with band reconstruction
+    for i in range(n):
+        for j in range(n):
+            var val = Scalar[dtype](0)
+
+            if upper:
+                if j >= i and j <= i + k:
+                    val = A_band[(k - (j - i)) + j * lda]
+                elif i >= j and i <= j + k:
+                    # symmetric mirror
+                    val = A_band[(k - (i - j)) + i * lda]
+            else:
+                if j <= i and j >= i - k:
+                    val = A_band[(i - j) + j * lda]
+                elif j >= i and j <= i + k:
+                    # symmetric mirror
+                    val = A_band[(j - i) + i * lda]
+
+            A_dense[i * n + j] = val