[Matrix][routines] Add Column-Major Memory Layout Support Across All Matrix Routines

.github/workflows/run_tests.yaml

@@ -53,3 +53,4 @@ jobs:
         run: |
           magic install
           magic run mojo test tests -I .
+          magic run mojo test tests/core/test_matrix.mojo -I . -D F_CONTIGUOUS

numojo/routines/linalg/decompositions.mojo

@@ -220,8 +220,8 @@ fn lu_decomposition[
     var n = A.shape[0]
 
     # Initiate upper and lower triangular matrices
-    var U = Matrix.full[dtype](shape=(n, n))
-    var L = Matrix.full[dtype](shape=(n, n))
+    var U = Matrix.full[dtype](shape=(n, n), order=A.order())
+    var L = Matrix.full[dtype](shape=(n, n), order=A.order())
 
     # Fill in L and U
     for i in range(0, n):
@@ -305,6 +305,8 @@ fn partial_pivoting[
     """
     var n = A.shape[0]
     var P = Matrix.identity[dtype](n)
+    if A.flags.F_CONTIGUOUS:
+        A = A.reorder_layout()
     var s: Int = 0  # Number of exchanges, for determinant
     for col in range(n):
         var max_p = abs(A[col, col])
@@ -318,6 +320,9 @@ fn partial_pivoting[
 
         if max_p_row != col:
             s = s + 1
+    if A.flags.F_CONTIGUOUS:
+        A = A.reorder_layout()
+        P = P.reorder_layout()
 
     return Tuple(A^, P^, s)
 

numojo/routines/linalg/products.mojo

@@ -381,25 +381,67 @@ fn matmul[
             )
         )
 
-    var C: Matrix[dtype] = Matrix.zeros[dtype](shape=(A.shape[0], B.shape[1]))
+    var C: Matrix[dtype]
 
-    @parameter
-    fn calculate_CC(m: Int):
-        for k in range(A.shape[1]):
+    if A.flags.C_CONTIGUOUS and B.flags.C_CONTIGUOUS:
+        C = Matrix.zeros[dtype](shape=(A.shape[0], B.shape[1]), order=B.order())
 
-            @parameter
-            fn dot[simd_width: Int](n: Int):
-                C._store[simd_width](
-                    m,
-                    n,
-                    C._load[simd_width](m, n)
-                    + A._load(m, k) * B._load[simd_width](k, n),
-                )
+        @parameter
+        fn calculate_CC(m: Int):
+            for k in range(A.shape[1]):
 
-            vectorize[dot, width](B.shape[1])
+                @parameter
+                fn dot[simd_width: Int](n: Int):
+                    C._store[simd_width](
+                        m,
+                        n,
+                        C._load[simd_width](m, n)
+                        + A._load(m, k) * B._load[simd_width](k, n),
+                    )
+
+                vectorize[dot, width](B.shape[1])
 
-    parallelize[calculate_CC](A.shape[0], A.shape[0])
+        parallelize[calculate_CC](A.shape[0], A.shape[0])
+    elif A.flags.F_CONTIGUOUS and B.flags.F_CONTIGUOUS:
+        C = Matrix.zeros[dtype](shape=(A.shape[0], B.shape[1]), order=B.order())
+
+        @parameter
+        fn calculate_FF(n: Int):
+            for k in range(A.shape[1]):
 
+                @parameter
+                fn dot_F[simd_width: Int](m: Int):
+                    C._store[simd_width](
+                        m,
+                        n,
+                        C._load[simd_width](m, n)
+                        + A._load[simd_width](m, k) * B._load(k, n),
+                    )
+
+                vectorize[dot_F, width](A.shape[0])
+
+        parallelize[calculate_FF](B.shape[1], B.shape[1])
+    elif A.flags.C_CONTIGUOUS and B.flags.F_CONTIGUOUS:
+        C = Matrix.zeros[dtype](shape=(A.shape[0], B.shape[1]), order=B.order())
+
+        @parameter
+        fn calculate_CF(m: Int):
+            for n in range(B.shape[1]):
+                var sum: Scalar[dtype] = 0.0
+
+                @parameter
+                fn dot_product[simd_width: Int](k: Int):
+                    sum += (
+                        A._load[simd_width](m, k) * B._load[simd_width](k, n)
+                    ).reduce_add()
+
+                vectorize[dot_product, width](A.shape[1])
+                C._store(m, n, sum)
+
+        parallelize[calculate_CF](A.shape[0], A.shape[0])
+
+    else:
+        C = matmul(A.reorder_layout(), B)
     var _A = A
     var _B = B
 

numojo/routines/linalg/solving.mojo

@@ -123,8 +123,11 @@ fn inv[dtype: DType](A: Matrix[dtype]) raises -> Matrix[dtype]:
         raise Error(
             String("{}x{} matrix is not square.").format(A.shape[0], A.shape[1])
         )
+    var order = "F"
+    if A.flags.C_CONTIGUOUS:
+        order = "C"
 
-    var I = Matrix.identity[dtype](A.shape[0])
+    var I = Matrix.identity[dtype](A.shape[0], order=order)
     var B = solve(A, I)
 
     return B^
@@ -364,16 +367,20 @@ fn solve[
     """
     Solve `AX = Y` using LUP decomposition.
     """
+    if A.flags.C_CONTIGUOUS != Y.flags.C_CONTIGUOUS:
+        raise Error("Input matrices A and Y must have the same memory layout")
+
     var U: Matrix[dtype]
     var L: Matrix[dtype]
+
     A_pivoted, P, _ = partial_pivoting(A)
     L, U = lu_decomposition[dtype](A_pivoted)
 
     var m = A.shape[0]
     var n = Y.shape[1]
 
-    var Z = Matrix.full[dtype]((m, n))
-    var X = Matrix.full[dtype]((m, n))
+    var Z = Matrix.full[dtype]((m, n), order=A.order())
+    var X = Matrix.full[dtype]((m, n), order=A.order())
 
     var PY = P @ Y
 

numojo/routines/manipulation.mojo

@@ -283,8 +283,11 @@ fn transpose[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]:
     """
     Transpose of matrix.
     """
+    var order = "F"
+    if A.flags.C_CONTIGUOUS:
+        order = "C"
 
-    var B = Matrix[dtype](Tuple(A.shape[1], A.shape[0]))
+    var B = Matrix[dtype](Tuple(A.shape[1], A.shape[0]), order=order)
 
     if A.shape[0] == 1 or A.shape[1] == 1:
         memcpy(B._buf.ptr, A._buf.ptr, A.size)
@@ -395,7 +398,9 @@ fn broadcast_to[
 
 fn broadcast_to[
     dtype: DType
-](A: Matrix[dtype], shape: Tuple[Int, Int]) raises -> Matrix[dtype]:
+](
+    A: Matrix[dtype], shape: Tuple[Int, Int], override_order: String = ""
+) raises -> Matrix[dtype]:
     """
     Broadcasts the vector to the given shape.
 
@@ -426,12 +431,17 @@ fn broadcast_to[
     Unhandled exception caught during execution: Cannot broadcast shape 2x2 to shape 4x2!
     ```
     """
+    var ord: String
+    if override_order == "":
+        ord = A.order()
+    else:
+        ord = override_order
 
-    var B = Matrix[dtype](shape)
+    var B = Matrix[dtype](shape, order=ord)
     if (A.shape[0] == shape[0]) and (A.shape[1] == shape[1]):
-        B = A
+        return A
     elif (A.shape[0] == 1) and (A.shape[1] == 1):
-        B = Matrix.full[dtype](shape, A[0, 0])
+        B = Matrix.full[dtype](shape, A[0, 0], order=ord)
     elif (A.shape[0] == 1) and (A.shape[1] == shape[1]):
         for i in range(shape[0]):
             memcpy(
@@ -453,13 +463,15 @@ fn broadcast_to[
 
 fn broadcast_to[
     dtype: DType
-](A: Scalar[dtype], shape: Tuple[Int, Int]) raises -> Matrix[dtype]:
+](A: Scalar[dtype], shape: Tuple[Int, Int], order: String) raises -> Matrix[
+    dtype
+]:
     """
     Broadcasts the scalar to the given shape.
     """
 
     var B = Matrix[dtype](shape)
-    B = Matrix.full[dtype](shape, A)
+    B = Matrix.full[dtype](shape, A, order=order)
     return B^