From 0e0ac8ec81a0ba4215ea45d2957eaa3a4b31e325 Mon Sep 17 00:00:00 2001 From: Jackson Mowry Date: Tue, 14 Apr 2026 14:19:45 +0000 Subject: [PATCH 1/2] multi-threaded tbmv impl --- src/level2/__init__.mojo | 1 + src/level2/tbmv_device.mojo | 197 +++++++++++++++++++++++++++ src/testing_utils/testing_utils.mojo | 23 ++++ test-level2.mojo | 118 ++++++++++++++++ 4 files changed, 339 insertions(+) create mode 100644 src/level2/tbmv_device.mojo diff --git a/src/level2/__init__.mojo b/src/level2/__init__.mojo index 59539ef..d256a33 100644 --- a/src/level2/__init__.mojo +++ b/src/level2/__init__.mojo @@ -5,4 +5,5 @@ from .syr2_device import * from .gbmv_device import * from .sbmv_device import * from .trsv_device import * +from .tbmv_device import * from .symv_device import * diff --git a/src/level2/tbmv_device.mojo b/src/level2/tbmv_device.mojo new file mode 100644 index 0000000..18b1fcc --- /dev/null +++ b/src/level2/tbmv_device.mojo @@ -0,0 +1,197 @@ +from gpu import thread_idx, block_idx, block_dim, grid_dim +from gpu.host import DeviceContext +from math import ceildiv +from ..level1.copy_device import blas_copy + +comptime TBsize = 512 + +# level2.tbmv +# Performs triangular band matrix-vector multiplication +# x := A*x, or x := A**T*x, +# where x is an n element vector and A is an n by n unit, or non-unit, +# upper or lower triangular band matrix with k super/sub-diagonals. +# +# Row-major band storage: +# Upper: A[i,j] stored at A_band[i*lda + (j-i)], for j in [i, min(n-1, i+k)] +# Lower: A[i,j] stored at A_band[i*lda + (i-j)], for j in [max(0, i-k), i] +# +# uplo: 0 = lower triangular, 1 = upper triangular +# trans: 0 = no transpose, 1 = transpose +# diag: 0 = non-unit diagonal, 1 = unit diagonal (diagonal not accessed) + +fn stbmv_device( + uplo: Int, + trans: Int, + diag: Int, + n: Int, + k: Int, + A: UnsafePointer[Float32, ImmutAnyOrigin], + lda: Int, + x: UnsafePointer[Float32, ImmutAnyOrigin], + incx: Int, + temp: UnsafePointer[Float32, MutAnyOrigin], +): + var global_i = block_dim.x * block_idx.x + thread_idx.x + var n_threads = grid_dim.x * block_dim.x + + if not trans: + if uplo: + # Upper triangular, no-transpose: temp[i] = sum_{j=i}^{min(n-1,i+k)} A[i,j]*x[j] + for i in range(global_i, n, n_threads): + var sum = Float32(0) + var j_end = min(n - 1, i + k) + for j in range(i, j_end + 1): + if j == i and diag: + sum += x[j * incx] + else: + sum += A[i * lda + (j - i)] * x[j * incx] + temp[i] = sum + else: + # Lower triangular, no-transpose: temp[i] = sum_{j=max(0,i-k)}^{i} A[i,j]*x[j] + for i in range(global_i, n, n_threads): + var sum = Float32(0) + var j_start = max(0, i - k) + for j in range(j_start, i + 1): + if j == i and diag: + sum += x[j * incx] + else: + sum += A[i * lda + (i - j)] * x[j * incx] + temp[i] = sum + else: + if uplo: + # Upper triangular, transpose: temp[j] = sum_{i=max(0,j-k)}^{j} A[i,j]*x[i] + for j in range(global_i, n, n_threads): + var sum = Float32(0) + var i_start = max(0, j - k) + for i in range(i_start, j + 1): + if i == j and diag: + sum += x[i * incx] + else: + sum += A[i * lda + (j - i)] * x[i * incx] + temp[j] = sum + else: + # Lower triangular, transpose: temp[j] = sum_{i=j}^{min(n-1,j+k)} A[i,j]*x[i] + for j in range(global_i, n, n_threads): + var sum = Float32(0) + var i_end = min(n - 1, j + k) + for i in range(j, i_end + 1): + if i == j and diag: + sum += x[i * incx] + else: + sum += A[i * lda + (i - j)] * x[i * incx] + temp[j] = sum + + +fn dtbmv_device( + uplo: Int, + trans: Int, + diag: Int, + n: Int, + k: Int, + A: UnsafePointer[Float64, ImmutAnyOrigin], + lda: Int, + x: UnsafePointer[Float64, ImmutAnyOrigin], + incx: Int, + temp: UnsafePointer[Float64, MutAnyOrigin], +): + var global_i = block_dim.x * block_idx.x + thread_idx.x + var n_threads = grid_dim.x * block_dim.x + + if not trans: + if uplo: + # Upper triangular, no-transpose: temp[i] = sum_{j=i}^{min(n-1,i+k)} A[i,j]*x[j] + for i in range(global_i, n, n_threads): + var sum = Float64(0) + var j_end = min(n - 1, i + k) + for j in range(i, j_end + 1): + if j == i and diag: + sum += x[j * incx] + else: + sum += A[i * lda + (j - i)] * x[j * incx] + temp[i] = sum + else: + # Lower triangular, no-transpose: temp[i] = sum_{j=max(0,i-k)}^{i} A[i,j]*x[j] + for i in range(global_i, n, n_threads): + var sum = Float64(0) + var j_start = max(0, i - k) + for j in range(j_start, i + 1): + if j == i and diag: + sum += x[j * incx] + else: + sum += A[i * lda + (i - j)] * x[j * incx] + temp[i] = sum + else: + if uplo: + # Upper triangular, transpose: temp[j] = sum_{i=max(0,j-k)}^{j} A[i,j]*x[i] + for j in range(global_i, n, n_threads): + var sum = Float64(0) + var i_start = max(0, j - k) + for i in range(i_start, j + 1): + if i == j and diag: + sum += x[i * incx] + else: + sum += A[i * lda + (j - i)] * x[i * incx] + temp[j] = sum + else: + # Lower triangular, transpose: temp[j] = sum_{i=j}^{min(n-1,j+k)} A[i,j]*x[i] + for j in range(global_i, n, n_threads): + var sum = Float64(0) + var i_end = min(n - 1, j + k) + for i in range(j, i_end + 1): + if i == j and diag: + sum += x[i * incx] + else: + sum += A[i * lda + (i - j)] * x[i * incx] + temp[j] = sum + + +fn blas_tbmv[dtype: DType]( + uplo: Int, + trans: Bool, + diag: Int, + n: Int, + k: Int, + d_A: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], + lda: Int, + d_x: UnsafePointer[Scalar[dtype], MutAnyOrigin], + incx: Int, + d_temp: UnsafePointer[Scalar[dtype], MutAnyOrigin], + ctx: DeviceContext, +) raises: + + # NOTE: add error checking here? + # check n > 0 + # check k >= 0 + # check lda >= k + 1 + # check incx > 0 + + var trans_i = 1 if trans else 0 + + # Use the caller-supplied temp buffer when provided; otherwise allocate one + # automatically for this call. owned_temp is only valid when d_temp is null. + var owned_temp = ctx.enqueue_create_buffer[dtype](0 if d_temp else n) + var work_temp = d_temp if d_temp else owned_temp.unsafe_ptr() + + @parameter + if dtype == DType.float32: + ctx.enqueue_function[stbmv_device, stbmv_device]( + uplo, trans_i, diag, + n, k, d_A, lda, + d_x, incx, work_temp, + grid_dim=ceildiv(n, TBsize), + block_dim=TBsize, + ) + elif dtype == DType.float64: + ctx.enqueue_function[dtbmv_device, dtbmv_device]( + uplo, trans_i, diag, + n, k, d_A, lda, + d_x, incx, work_temp, + grid_dim=ceildiv(n, TBsize), + block_dim=TBsize, + ) + else: + raise Error("blas_tbmv: Unsupported type") + + # Copy accumulated results from temp back into x using the level-1 + # copy routine (temp has stride 1; x has stride incx). + blas_copy[dtype](n, work_temp, 1, d_x, incx, ctx) diff --git a/src/testing_utils/testing_utils.mojo b/src/testing_utils/testing_utils.mojo index 9f0e5ef..071b0fe 100644 --- a/src/testing_utils/testing_utils.mojo +++ b/src/testing_utils/testing_utils.mojo @@ -135,6 +135,29 @@ fn frobenius_norm_symmetric[dtype: DType]( return sqrt(sum) +fn dense_to_tri_band_rm[dtype: DType]( + A_dense: UnsafePointer[Scalar[dtype], MutAnyOrigin], + A_band: UnsafePointer[Scalar[dtype], MutAnyOrigin], + n: Int, + k: Int, + upper: Int, +): + var lda = k + 1 + + for i in range(n * lda): + A_band[i] = 0 + + if upper: + for i in range(n): + var j_end = min(n - 1, i + k) + for j in range(i, j_end + 1): + A_band[i * lda + (j - i)] = A_dense[i * n + j] + else: + for i in range(n): + var j_start = max(0, i - k) + for j in range(j_start, i + 1): + A_band[i * lda + (i - j)] = A_dense[i * n + j] + fn dense_to_band[dtype: DType]( A: UnsafePointer[Scalar[dtype], MutAnyOrigin], B: UnsafePointer[Scalar[dtype], MutAnyOrigin], diff --git a/test-level2.mojo b/test-level2.mojo index 9f44af8..d318bf0 100644 --- a/test-level2.mojo +++ b/test-level2.mojo @@ -764,6 +764,96 @@ def trsv_test[ ) assert_true(ok) +def tbmv_test[ + dtype: DType, + n: Int, + k: Int, + uplo: Int, + trans: Bool, + diag: Int, +](): + comptime lda = k + 1 + + with DeviceContext() as ctx: + A_dense = ctx.enqueue_create_host_buffer[dtype](n * n) + A_band = ctx.enqueue_create_host_buffer[dtype](lda * n) + A_d = ctx.enqueue_create_buffer[dtype](lda * n) + x = ctx.enqueue_create_host_buffer[dtype](n) + x_d = ctx.enqueue_create_buffer[dtype](n) + + generate_random_arr[dtype](n * n, A_dense.unsafe_ptr(), -1, 1) + generate_random_arr[dtype](n, x.unsafe_ptr(), -100, 100) + + # Zero out elements outside the triangle and outside the band + for i in range(n): + for j in range(n): + if uplo: # upper triangular + if j < i or j > i + k: + A_dense[i * n + j] = 0 + else: # lower triangular + if j > i or j < i - k: + A_dense[i * n + j] = 0 + # Set unit diagonal to 1 so the numpy reference is consistent + if diag: + A_dense[i * n + i] = 1 + + # Pack dense triangular band into row-major band storage + dense_to_tri_band_rm(A_dense.unsafe_ptr(), A_band.unsafe_ptr(), n, k, uplo) + + ctx.enqueue_copy(A_d, A_band) + ctx.enqueue_copy(x_d, x) + ctx.synchronize() + + var norm_A = frobenius_norm[dtype](A_dense.unsafe_ptr(), n * n) + var norm_x = frobenius_norm[dtype](x.unsafe_ptr(), n) + + blas_tbmv[dtype]( + uplo, trans, diag, + n, k, + A_d.unsafe_ptr(), lda, + x_d.unsafe_ptr(), 1, + UnsafePointer[Scalar[dtype], MutAnyOrigin](), # Emtpy temp (forces auto allocation) + ctx, + ) + + np = Python.import_module("numpy") + + py_A = Python.list() + py_x = Python.list() + for i in range(n * n): + py_A.append(A_dense[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).reshape(n, n) + np_x = np.array(py_x, dtype=np.float32) + sp_res = np.dot(np_A.T, np_x) if trans else np.dot(np_A, np_x) + elif dtype == DType.float64: + np_A = np.array(py_A, dtype=np.float64).reshape(n, n) + np_x = np.array(py_x, dtype=np.float64) + sp_res = np.dot(np_A.T, np_x) if trans else np.dot(np_A, np_x) + 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]) + norm_diff += diff * diff + norm_diff = sqrt(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]() @@ -848,6 +938,33 @@ def test_trsv(): trsv_test[DType.float64, 64, False, 0, 1]() trsv_test[DType.float64, 64, False, 1, 1]() +def test_tbmv(): + # uplo: 0=lower, 1=upper trans: False/True diag: 0=non-unit, 1=unit + tbmv_test[DType.float32, 64, 1, 1, False, 0]() + tbmv_test[DType.float32, 64, 1, 0, False, 0]() + tbmv_test[DType.float32, 64, 2, 1, True, 0]() + tbmv_test[DType.float32, 64, 2, 0, True, 0]() + tbmv_test[DType.float32, 64, 4, 1, False, 1]() + tbmv_test[DType.float32, 64, 4, 0, False, 1]() + tbmv_test[DType.float32, 64, 4, 1, True, 1]() + tbmv_test[DType.float32, 64, 4, 0, True, 1]() + tbmv_test[DType.float64, 64, 1, 1, False, 0]() + tbmv_test[DType.float64, 64, 1, 0, False, 0]() + tbmv_test[DType.float64, 64, 2, 1, True, 0]() + tbmv_test[DType.float64, 64, 2, 0, True, 0]() + tbmv_test[DType.float64, 64, 4, 1, False, 1]() + tbmv_test[DType.float64, 64, 4, 0, False, 1]() + tbmv_test[DType.float64, 64, 4, 1, True, 1]() + tbmv_test[DType.float64, 64, 4, 0, True, 1]() + tbmv_test[DType.float32, 512, 16, 1, False, 0]() + tbmv_test[DType.float32, 512, 16, 0, False, 0]() + tbmv_test[DType.float32, 512, 16, 1, True, 0]() + tbmv_test[DType.float32, 512, 16, 0, True, 0]() + tbmv_test[DType.float64, 512, 16, 1, False, 0]() + tbmv_test[DType.float64, 512, 16, 0, False, 0]() + tbmv_test[DType.float64, 512, 16, 1, True, 0]() + tbmv_test[DType.float64, 512, 16, 0, True, 0]() + def main(): print("--- MojoBLAS Level 2 routines testing ---") var args = argv() @@ -864,6 +981,7 @@ def main(): elif args[i] == "syr2": suite.test[test_syr2]() elif args[i] == "gbmv": suite.test[test_gbmv]() elif args[i] == "trsv": suite.test[test_trsv]() + elif args[i] == "tbmv": suite.test[test_tbmv]() elif args[i] == "symv": suite.test[test_symv]() else: print("unknown routine:", args[i]) suite^.run() From 7b1b5b34d8e6880b93155706d6d7bbc4ad141470 Mon Sep 17 00:00:00 2001 From: Jackson Mowry Date: Tue, 14 Apr 2026 14:48:56 +0000 Subject: [PATCH 2/2] single-threaded tbsv implementation --- src/level2/__init__.mojo | 1 + src/level2/tbsv_device.mojo | 188 ++++++++++++++++++++++++++++++++++++ test-level2.mojo | 135 ++++++++++++++++++++++++++ 3 files changed, 324 insertions(+) create mode 100644 src/level2/tbsv_device.mojo diff --git a/src/level2/__init__.mojo b/src/level2/__init__.mojo index d256a33..809b0f3 100644 --- a/src/level2/__init__.mojo +++ b/src/level2/__init__.mojo @@ -6,4 +6,5 @@ from .gbmv_device import * from .sbmv_device import * from .trsv_device import * from .tbmv_device import * +from .tbsv_device import * from .symv_device import * diff --git a/src/level2/tbsv_device.mojo b/src/level2/tbsv_device.mojo new file mode 100644 index 0000000..f9bc217 --- /dev/null +++ b/src/level2/tbsv_device.mojo @@ -0,0 +1,188 @@ +from gpu import thread_idx, block_idx, block_dim, grid_dim +from gpu.host import DeviceContext + +# level2.tbsv +# Solves a triangular band system +# A*x = b, or A**T*x = b, +# where b and x are n element vectors and A is an n by n unit, or non-unit, +# upper or lower triangular band matrix with k super/sub-diagonals. +# +# Row-major band storage (same layout as tbmv): +# Upper: A[i,j] stored at A_band[i*lda + (j-i)], for j in [i, min(n-1, i+k)] +# Lower: A[i,j] stored at A_band[i*lda + (i-j)], for j in [max(0, i-k), i] +# Diagonal A[i,i] is always at A_band[i*lda + 0] for both upper and lower. +# +# uplo: 0 = lower triangular, 1 = upper triangular +# trans: 0 = no transpose, 1 = transpose +# diag: 0 = non-unit diagonal, 1 = unit diagonal (diagonal not accessed) + +fn stbsv_device( + uplo: Int, + trans: Int, + diag: Int, + n: Int, + k: Int, + A: UnsafePointer[Float32, ImmutAnyOrigin], + lda: Int, + x: UnsafePointer[Float32, MutAnyOrigin], + incx: Int, +): + var tid = block_dim.x * block_idx.x + thread_idx.x + if tid != 0: + return + + if not trans: + if not uplo: + # Lower triangular, no-transpose: forward substitution + # A[i,j] at A_band[i*lda + (i-j)]; diagonal at A_band[i*lda + 0] + for i in range(n): + var temp = x[i * incx] + var j_start = max(0, i - k) + for j in range(j_start, i): + temp -= A[i * lda + (i - j)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + else: + # Upper triangular, no-transpose: backward substitution + # A[i,j] at A_band[i*lda + (j-i)]; diagonal at A_band[i*lda + 0] + for idx in range(n): + var i = n - 1 - idx + var temp = x[i * incx] + var j_end = min(n - 1, i + k) + for j in range(i + 1, j_end + 1): + temp -= A[i * lda + (j - i)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + else: + if not uplo: + # Transpose lower triangular (A^T is upper): backward substitution + # A^T[i,j] = A[j,i]; for lower A, A[j,i] at A_band[j*lda + (j-i)] + # where j in [i+1, min(n-1, i+k)] + for idx in range(n): + var i = n - 1 - idx + var temp = x[i * incx] + var j_end = min(n - 1, i + k) + for j in range(i + 1, j_end + 1): + temp -= A[j * lda + (j - i)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + else: + # Transpose upper triangular (A^T is lower): forward substitution + # A^T[i,j] = A[j,i]; for upper A, A[j,i] at A_band[j*lda + (i-j)] + # where j in [max(0, i-k), i-1] + for i in range(n): + var temp = x[i * incx] + var j_start = max(0, i - k) + for j in range(j_start, i): + temp -= A[j * lda + (i - j)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + + +fn dtbsv_device( + uplo: Int, + trans: Int, + diag: Int, + n: Int, + k: Int, + A: UnsafePointer[Float64, ImmutAnyOrigin], + lda: Int, + x: UnsafePointer[Float64, MutAnyOrigin], + incx: Int, +): + var tid = block_dim.x * block_idx.x + thread_idx.x + if tid != 0: + return + + if not trans: + if not uplo: + # Lower triangular, no-transpose: forward substitution + for i in range(n): + var temp = x[i * incx] + var j_start = max(0, i - k) + for j in range(j_start, i): + temp -= A[i * lda + (i - j)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + else: + # Upper triangular, no-transpose: backward substitution + for idx in range(n): + var i = n - 1 - idx + var temp = x[i * incx] + var j_end = min(n - 1, i + k) + for j in range(i + 1, j_end + 1): + temp -= A[i * lda + (j - i)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + else: + if not uplo: + # Transpose lower triangular (A^T is upper): backward substitution + for idx in range(n): + var i = n - 1 - idx + var temp = x[i * incx] + var j_end = min(n - 1, i + k) + for j in range(i + 1, j_end + 1): + temp -= A[j * lda + (j - i)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + else: + # Transpose upper triangular (A^T is lower): forward substitution + for i in range(n): + var temp = x[i * incx] + var j_start = max(0, i - k) + for j in range(j_start, i): + temp -= A[j * lda + (i - j)] * x[j * incx] + if not diag: + temp /= A[i * lda] + x[i * incx] = temp + + +fn blas_tbsv[dtype: DType]( + uplo: Int, + trans: Bool, + diag: Int, + n: Int, + k: Int, + d_A: UnsafePointer[Scalar[dtype], ImmutAnyOrigin], + lda: Int, + d_x: UnsafePointer[Scalar[dtype], MutAnyOrigin], + incx: Int, + ctx: DeviceContext, +) raises: + + # NOTE: add error checking here? + # check n > 0 + # check k >= 0 + # check lda >= k + 1 + # check incx > 0 + + var trans_i = 1 if trans else 0 + + @parameter + if dtype == DType.float32: + ctx.enqueue_function[stbsv_device, stbsv_device]( + uplo, trans_i, diag, + n, k, d_A, lda, + d_x, incx, + grid_dim=1, + block_dim=1, + ) + elif dtype == DType.float64: + ctx.enqueue_function[dtbsv_device, dtbsv_device]( + uplo, trans_i, diag, + n, k, d_A, lda, + d_x, incx, + grid_dim=1, + block_dim=1, + ) + else: + raise Error("blas_tbsv: Unsupported type") + + ctx.synchronize() diff --git a/test-level2.mojo b/test-level2.mojo index d318bf0..e029195 100644 --- a/test-level2.mojo +++ b/test-level2.mojo @@ -854,6 +854,113 @@ def tbmv_test[ ) assert_true(ok) +def tbsv_test[ + dtype: DType, + n: Int, + k: Int, + uplo: Int, + trans: Bool, + diag: Int, +](): + comptime lda = k + 1 + + with DeviceContext() as ctx: + A_dense = ctx.enqueue_create_host_buffer[dtype](n * n) + A_band = ctx.enqueue_create_host_buffer[dtype](lda * n) + A_d = ctx.enqueue_create_buffer[dtype](lda * n) + b = ctx.enqueue_create_host_buffer[dtype](n) + x_d = ctx.enqueue_create_buffer[dtype](n) + + generate_random_arr[dtype](n * n, A_dense.unsafe_ptr(), -1, 1) + generate_random_arr[dtype](n, b.unsafe_ptr(), -1, 1) + + # Zero out elements outside the triangle and band; ensure diagonal + # dominance for a well-conditioned system + for i in range(n): + for j in range(n): + if uplo: # upper triangular + if j < i or j > i + k: + A_dense[i * n + j] = 0 + else: # lower triangular + if j > i or j < i - k: + A_dense[i * n + j] = 0 + if diag: + A_dense[i * n + i] = 1 + else: + # Make diagonally dominant to keep the system well-conditioned + A_dense[i * n + i] += 1000 + + # Pack dense triangular band into row-major band storage + dense_to_tri_band_rm(A_dense.unsafe_ptr(), A_band.unsafe_ptr(), n, k, uplo) + + ctx.enqueue_copy(A_d, A_band) + ctx.enqueue_copy(x_d, b) + ctx.synchronize() + + var norm_A = frobenius_norm[dtype](A_dense.unsafe_ptr(), n * n) + var norm_b = frobenius_norm[dtype](b.unsafe_ptr(), n) + + blas_tbsv[dtype]( + uplo, trans, diag, + n, k, + A_d.unsafe_ptr(), lda, + x_d.unsafe_ptr(), 1, + ctx, + ) + + sp = Python.import_module("scipy") + np = Python.import_module("numpy") + + py_A = Python.list() + py_b = Python.list() + for i in range(n * n): + py_A.append(A_dense[i]) + for i in range(n): + py_b.append(b[i]) + + var sp_res: PythonObject + + var is_lower = 0 if uplo else 1 + var trans_mode = 1 if trans else 0 + var unit_diag = True if diag else False + + if dtype == DType.float32: + np_A = np.array(py_A, dtype=np.float32).reshape(n, n) + np_b = np.array(py_b, dtype=np.float32) + sp_res = sp.linalg.solve_triangular( + np_A, np_b, + lower=is_lower, + trans=trans_mode, + unit_diagonal=unit_diag, + ) + elif dtype == DType.float64: + np_A = np.array(py_A, dtype=np.float64).reshape(n, n) + np_b = np.array(py_b, dtype=np.float64) + sp_res = sp.linalg.solve_triangular( + np_A, np_b, + lower=is_lower, + trans=trans_mode, + unit_diagonal=unit_diag, + ) + 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]) + norm_diff += diff * diff + norm_diff = sqrt(norm_diff) + + var ok = check_gemm_error[dtype]( + 1, n, n, + Scalar[dtype](1), Scalar[dtype](0), + norm_A, norm_b, 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]() @@ -965,6 +1072,33 @@ def test_tbmv(): tbmv_test[DType.float64, 512, 16, 1, True, 0]() tbmv_test[DType.float64, 512, 16, 0, True, 0]() +def test_tbsv(): + # uplo: 0=lower, 1=upper trans: False/True diag: 0=non-unit, 1=unit + tbsv_test[DType.float32, 64, 1, 1, False, 0]() + tbsv_test[DType.float32, 64, 1, 0, False, 0]() + tbsv_test[DType.float32, 64, 2, 1, True, 0]() + tbsv_test[DType.float32, 64, 2, 0, True, 0]() + tbsv_test[DType.float32, 64, 4, 1, False, 1]() + tbsv_test[DType.float32, 64, 4, 0, False, 1]() + tbsv_test[DType.float32, 64, 4, 1, True, 1]() + tbsv_test[DType.float32, 64, 4, 0, True, 1]() + tbsv_test[DType.float64, 64, 1, 1, False, 0]() + tbsv_test[DType.float64, 64, 1, 0, False, 0]() + tbsv_test[DType.float64, 64, 2, 1, True, 0]() + tbsv_test[DType.float64, 64, 2, 0, True, 0]() + tbsv_test[DType.float64, 64, 4, 1, False, 1]() + tbsv_test[DType.float64, 64, 4, 0, False, 1]() + tbsv_test[DType.float64, 64, 4, 1, True, 1]() + tbsv_test[DType.float64, 64, 4, 0, True, 1]() + tbsv_test[DType.float32, 512, 16, 1, False, 0]() + tbsv_test[DType.float32, 512, 16, 0, False, 0]() + tbsv_test[DType.float32, 512, 16, 1, True, 0]() + tbsv_test[DType.float32, 512, 16, 0, True, 0]() + tbsv_test[DType.float64, 512, 16, 1, False, 0]() + tbsv_test[DType.float64, 512, 16, 0, False, 0]() + tbsv_test[DType.float64, 512, 16, 1, True, 0]() + tbsv_test[DType.float64, 512, 16, 0, True, 0]() + def main(): print("--- MojoBLAS Level 2 routines testing ---") var args = argv() @@ -982,6 +1116,7 @@ def main(): elif args[i] == "gbmv": suite.test[test_gbmv]() 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] == "symv": suite.test[test_symv]() else: print("unknown routine:", args[i]) suite^.run()