diff --git a/bench-level2.mojo b/bench-level2.mojo index c3b87ba..4ae5c5e 100644 --- a/bench-level2.mojo +++ b/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) diff --git a/src/level2/__init__.mojo b/src/level2/__init__.mojo index 9a144f2..3f336ad 100644 --- a/src/level2/__init__.mojo +++ b/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 * diff --git a/src/level2/tpmv_device.mojo b/src/level2/tpmv_device.mojo new file mode 100644 index 0000000..728584a --- /dev/null +++ b/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() diff --git a/src/testing_utils/testing_utils.mojo b/src/testing_utils/testing_utils.mojo index 293bea0..8637105 100644 --- a/src/testing_utils/testing_utils.mojo +++ b/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) \ No newline at end of file + 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 \ No newline at end of file diff --git a/test-level2.mojo b/test-level2.mojo index a7e742b..2cd24ba 100644 --- a/test-level2.mojo +++ b/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()