trsv implementation

src/level1/__init__.mojo

@@ -5,6 +5,7 @@ from .copy_device import *
 from .rot_device import *
 from .rotg import *
 from .rotm_device import *
+from .rotmg import *
 from .swap_device import *
 from .dot_device import *
 from .dotc_device import *

src/level1/rotmg.mojo

@@ -0,0 +1,150 @@
+fn blas_rotmg[dtype: DType](
+    d1: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    d2: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    x1: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    y1: UnsafePointer[Scalar[dtype], MutAnyOrigin],
+    mut param: List[Scalar[dtype]]
+):
+
+    var flag: Scalar[dtype]
+    var h11: Scalar[dtype] = 0
+    var h12: Scalar[dtype] = 0
+    var h21: Scalar[dtype] = 0
+    var h22: Scalar[dtype] = 0
+    var p1: Scalar[dtype]
+    var p2: Scalar[dtype]
+    var q1: Scalar[dtype]
+    var q2: Scalar[dtype]
+    var temp: Scalar[dtype]
+    var u: Scalar[dtype]
+    var gam: Scalar[dtype] = 4096.0
+    var gamsq: Scalar[dtype] = 16777216.0
+    var rgamsq: Scalar[dtype] = 5.9604645e-8
+
+    if (d1[] < 0):
+        # GO 0-H-D-AND-x1..
+        flag = -1
+        h11 = 0
+        h12 = 0
+        h21 = 0
+        h22 = 0
+
+        d1[] = 0
+        d2[] = 0
+        x1[] = 0
+    else:
+        # CASE-d1-NONNEGATIVE
+        p2 = d2[]*y1[]
+        if (p2 == 0):
+            flag = -2
+            param[0] = flag
+            return
+
+        # REGULAR-CASE..
+        p1 = d1[]*x1[]
+        q2 = p2*y1[]
+        q1 = p1*x1[]
+        #
+        if (abs(q1) > abs(q2)):
+            h21 = -y1[]/x1[]
+            h12 = p2/p1
+            #
+            u = 1 - h12*h21
+            #
+            if (u > 0):
+                flag = 0
+                d1[] /= u
+                d2[] /= u
+                x1[] *= u
+            else:
+                # This code path if here for safety. We do not expect this
+                # condition to ever hold except in edge cases with rounding
+                # errors. See DOI: 10.1145/355841.355847
+                flag = -1
+                h11 = 0
+                h12 = 0
+                h21 = 0
+                h22 = 0
+                #
+                d1[] = 0
+                d2[] = 0
+                x1[] = 0
+        else:
+            if (q2 < 0):
+                # GO 0-H-D-AND-x1..
+                flag = -1
+                h11 = 0
+                h12 = 0
+                h21 = 0
+                h22 = 0
+                #
+                d1[] = 0
+                d2[] = 0
+                x1[] = 0
+            else:
+                flag = 1
+                h11 = p1/p2
+                h22 = x1[]/y1[]
+                u = 1 + h11*h22
+                temp = d2[]/u
+                d2[] = d1[]/u
+                d1[] = temp
+                x1[] = y1[]*u
+
+        # PROCEDURE..SCALE-CHECK
+        if (d1[] != 0):
+            while ((d1[] <= rgamsq)  or  (d1[] >= gamsq)):
+                if (flag == 0):
+                    h11 = 1
+                    h22 = 1
+                    flag = -1
+                else:
+                    h21 = -1
+                    h12 = 1
+                    flag = -1
+
+                if (d1[] <= rgamsq):
+                    d1[] *= gam**2
+                    x1[] /= gam
+                    h11 /= gam
+                    h12 /= gam
+                else:
+                    d1[] /= gam**2
+                    x1[] *= gam
+                    h11 *= gam
+                    h12 *= gam
+
+        if (d2[] != 0):
+            while ( (abs(d2[]) <= rgamsq) or (abs(d2[]) >= gamsq) ):
+                if (flag == 0):
+                    h11 = 1
+                    h22 = 1
+                    flag = -1
+                else:
+                    h21 = -1
+                    h12 = 1
+                    flag = -1
+
+                if (abs(d2[]) <= rgamsq):
+                    d2[] *= gam**2
+                    h21 /= gam
+                    h22 /= gam
+                else:
+                    d2[] /= gam**2
+                    h21 *= gam
+                    h22 *= gam
+
+    if (flag < 0):
+        param[1] = h11
+        param[2] = h21
+        param[3] = h12
+        param[4] = h22
+    elif (flag == 0):
+        param[2] = h21
+        param[3] = h12
+    else:
+        param[1] = h11
+        param[4] = h22
+
+    param[0] = flag
+    return

src/level2/__init__.mojo

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

src/level2/trsv_device.mojo

@@ -0,0 +1,165 @@
+from gpu import thread_idx, block_idx, block_dim, grid_dim
+from gpu.host import DeviceContext
+from math import ceildiv
+
+# level2.trsv
+# Triangular matrix-vector solve
+# Solves one of the systems of equations
+#    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 matrix.
+#
+# uplo:  0 = lower triangular, 1 = upper triangular
+# trans: 0 = no transpose,     1 = transpose
+# diag:  0 = non-unit diagonal, 1 = unit diagonal
+fn strsv_device(
+    uplo: Int,
+    trans: Int,
+    diag: Int,
+    n: 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: forward substitution
+            for i in range(n):
+                var temp = x[i * incx]
+                for j in range(i):
+                    temp -= A[i * lda + j] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+        else:
+            # Upper triangular: backward substitution
+            for k in range(n):
+                var i = n - 1 - k
+                var temp = x[i * incx]
+                for j in range(i + 1, n):
+                    temp -= A[i * lda + j] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+    else:
+        if not uplo:
+            # Transpose lower triangular (A^T is upper): backward substitution
+            for k in range(n):
+                var i = n - 1 - k
+                var temp = x[i * incx]
+                for j in range(i + 1, n):
+                    temp -= A[j * lda + i] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+        else:
+            # Transpose upper triangular (A^T is lower): forward substitution
+            for i in range(n):
+                var temp = x[i * incx]
+                for j in range(i):
+                    temp -= A[j * lda + i] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+
+
+fn dtrsv_device(
+    uplo: Int,
+    trans: Int,
+    diag: Int,
+    n: 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: forward substitution
+            for i in range(n):
+                var temp = x[i * incx]
+                for j in range(i):
+                    temp -= A[i * lda + j] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+        else:
+            # Upper triangular: backward substitution
+            for k in range(n):
+                var i = n - 1 - k
+                var temp = x[i * incx]
+                for j in range(i + 1, n):
+                    temp -= A[i * lda + j] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+    else:
+        if not uplo:
+            # Transpose lower triangular (A^T is upper): backward substitution
+            for k in range(n):
+                var i = n - 1 - k
+                var temp = x[i * incx]
+                for j in range(i + 1, n):
+                    temp -= A[j * lda + i] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+        else:
+            # Transpose upper triangular (A^T is lower): forward substitution
+            for i in range(n):
+                var temp = x[i * incx]
+                for j in range(i):
+                    temp -= A[j * lda + i] * x[j * incx]
+                if not diag:
+                    temp /= A[i * lda + i]
+                x[i * incx] = temp
+
+
+fn blas_trsv[dtype: DType](
+    uplo: Int,
+    trans: Bool,
+    diag: Int,
+    n: 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 incx > 0
+
+    var trans_i = 1 if trans else 0
+
+    @parameter
+    if dtype == DType.float32:
+        ctx.enqueue_function[strsv_device, strsv_device](
+            uplo, trans_i, diag,
+            n, d_A, lda,
+            d_x, incx,
+            grid_dim=1,
+            block_dim=1,
+        )
+    elif dtype == DType.float64:
+        ctx.enqueue_function[dtrsv_device, dtrsv_device](
+            uplo, trans_i, diag,
+            n, d_A, lda,
+            d_x, incx,
+            grid_dim=1,
+            block_dim=1,
+        )
+    else:
+        raise Error("blas_trsv: Unsupported type")
+
+    ctx.synchronize()

src/testing_utils/testing_utils.mojo

@@ -37,7 +37,7 @@ def generate_random_scalar[
 
 
 # Error check following BLAS++ check_gemm:
-# NOTE: can't get epsilon value in Mojo; using tol32, tol64
+# NOTE: might use this for dot, gemv, ger, geru, gemm, symv, hemv, symm, trmv, trsv?, trmm, trsm?
 fn check_gemm_error[dtype: DType](
     m: Int, n: Int, k: Int,
     alpha: Scalar[dtype],