replace all matmul's by gemm

wassup05 · wassup05 · commit e709f838aeb1 · 2025-03-21T01:58:17.000+05:30
diff --git a/example/intrinsics/example_matmul.f90 b/example/intrinsics/example_matmul.f90
@@ -4,9 +4,9 @@ program example_matmul
     real :: r1(50, 100), r2(100, 40), r3(40, 50)
     real, allocatable :: res(:, :)
     x = reshape([(0, 0), (1, 0), (1, 0), (0, 0)], [2, 2])
-    y = reshape([(0, 0), (0, -1), (0, 1), (0, 0)], [2, 2]) ! pauli y-matrix
+    y = reshape([(0, 0), (0, 1), (0, -1), (0, 0)], [2, 2]) ! pauli y-matrix
 
-    print *, stdlib_matmul(y, y, y, y, y) ! should be y
+    print *, stdlib_matmul(y, y, y) ! should be y
     print *, stdlib_matmul(x, x, y, x) ! should be -i x sigma_z
 
     call random_seed()
diff --git a/src/stdlib_intrinsics.fypp b/src/stdlib_intrinsics.fypp
@@ -158,10 +158,10 @@ module stdlib_intrinsics
         !!
         !! matrix multiply more than two matrices with a single function call
         !! the multiplication with the optimal parenthesization for efficiency of computation is done automatically
-        !! Supported data types are `real`, `integer` and `complex`.
+        !! Supported data types are `real` and `complex`.
         !!
         !! Note: The matrices must be of compatible shapes to be multiplied
-        #:for k, t, s in I_KINDS_TYPES + R_KINDS_TYPES + C_KINDS_TYPES
+        #:for k, t, s in R_KINDS_TYPES + C_KINDS_TYPES
             pure module function stdlib_matmul_${s}$ (m1, m2, m3, m4, m5) result(r)
                 ${t}$, intent(in) :: m1(:,:), m2(:,:)
                 ${t}$, intent(in), optional :: m3(:,:), m4(:,:), m5(:,:)
diff --git a/src/stdlib_intrinsics_matmul.fypp b/src/stdlib_intrinsics_matmul.fypp
@@ -4,6 +4,8 @@
 #:set C_KINDS_TYPES = list(zip(CMPLX_KINDS, CMPLX_TYPES, CMPLX_SUFFIX))
 
 submodule (stdlib_intrinsics) stdlib_intrinsics_matmul
+    use stdlib_linalg_blas, only: gemm
+    use stdlib_constants
     implicit none
 
 contains
@@ -36,38 +38,84 @@ contains
         end do
     end function matmul_chain_order
 
-#:for k, t, s in I_KINDS_TYPES + R_KINDS_TYPES + C_KINDS_TYPES
+#:for k, t, s in R_KINDS_TYPES + C_KINDS_TYPES
     
-    pure function matmul_chain_mult_${s}$_3 (m1, m2, m3, start, s) result(r)
+    pure function matmul_chain_mult_${s}$_3 (m1, m2, m3, start, s, p) result(r)
         ${t}$, intent(in) :: m1(:,:), m2(:,:), m3(:,:)
-        integer, intent(in) :: start, s(:,2:)
-        ${t}$, allocatable :: r(:,:)
-        integer :: tmp
-        tmp = s(start, start + 2)
-
-        if (tmp == start) then
-            r = matmul(m1, matmul(m2, m3))
-        else if (tmp == start + 1) then
-            r = matmul(matmul(m1, m2), m3)
+        integer, intent(in) :: start, s(:,2:), p(:)
+        ${t}$, allocatable :: r(:,:), temp(:,:)
+        integer :: ord, m, n, k
+        ord = s(start, start + 2)
+        allocate(r(p(start), p(start + 3)))
+
+        if (ord == start) then
+            ! m1*(m2*m3)
+            m = p(start + 1)
+            n = p(start + 3)
+            k = p(start + 2)
+            allocate(temp(m,n))
+            call gemm('N', 'N', m, n, k, one_${s}$, m2, m, m3, k, zero_${s}$, temp, m)
+            m = p(start)
+            n = p(start + 3)
+            k = p(start + 1)
+            call gemm('N', 'N', m, n, k, one_${s}$, m1, m, temp, k, zero_${s}$, r, m)
+        else if (ord == start + 1) then
+            ! (m1*m2)*m3
+            m = p(start)
+            n = p(start + 2)
+            k = p(start + 1)
+            allocate(temp(m, n))
+            call gemm('N', 'N', m, n, k, one_${s}$, m1, m, m2, k, zero_${s}$, temp, m)
+            m = p(start)
+            n = p(start + 3)
+            k = p(start + 1)
+            call gemm('N', 'N', m, n, k, one_${s}$, temp, m, m3, k, zero_${s}$, r, m)
         else
             error stop "stdlib_matmul: error: unexpected s(i,j)"
         end if
 
     end function matmul_chain_mult_${s}$_3
 
-    pure function matmul_chain_mult_${s}$_4 (m1, m2, m3, m4, start, s) result(r)
+    pure function matmul_chain_mult_${s}$_4 (m1, m2, m3, m4, start, s, p) result(r)
         ${t}$, intent(in) :: m1(:,:), m2(:,:), m3(:,:), m4(:,:)
-        integer, intent(in) :: start, s(:,2:)
-        ${t}$, allocatable :: r(:,:)
-        integer :: tmp
-        tmp = s(start, start + 3)
-
-        if (tmp == start) then
-            r = matmul(m1, matmul_chain_mult_${s}$_3(m2, m3, m4, start + 1, s))
-        else if (tmp == start + 1) then
-            r = matmul(matmul(m1, m2), matmul(m3, m4))
-        else if (tmp == start + 2) then
-            r = matmul(matmul_chain_mult_${s}$_3(m1, m2, m3, start, s), m4)
+        integer, intent(in) :: start, s(:,2:), p(:)
+        ${t}$, allocatable :: r(:,:), temp(:,:), temp1(:,:)
+        integer :: ord, m, n, k
+        ord = s(start, start + 3)
+        allocate(r(p(start), p(start + 4)))
+
+        if (ord == start) then
+            ! m1*(m2*m3*m4)
+            temp = matmul_chain_mult_${s}$_3(m2, m3, m4, start + 1, s, p)
+            m = p(start)
+            n = p(start + 4)
+            k = p(start + 1)
+            call gemm('N', 'N', m, n, k, one_${s}$, m1, m, temp, k, zero_${s}$, r, m)
+        else if (ord == start + 1) then
+            ! (m1*m2)*(m3*m4)
+            m = p(start)
+            n = p(start + 2)
+            k = p(start + 1)
+            allocate(temp(m,n))
+            call gemm('N', 'N', m, n, k, one_${s}$, m1, m, m2, k, zero_${s}$, temp, m)
+
+            m = p(start + 2)
+            n = p(start + 4)
+            k = p(start + 3)
+            allocate(temp1(m,n))
+            call gemm('N', 'N', m, n, k, one_${s}$, m3, m, m4, k, zero_${s}$, temp1, m)
+
+            m = p(start)
+            n = p(start + 4)
+            k = p(start + 2)
+            call gemm('N', 'N', m, n, k, one_${s}$, temp, m, temp1, k, zero_${s}$, r, m)
+        else if (ord == start + 2) then
+            ! (m1*m2*m3)*m4
+            temp = matmul_chain_mult_${s}$_3(m1, m2, m3, start, s, p)
+            m = p(start)
+            n = p(start + 4)
+            k = p(start + 3)
+            call gemm('N', 'N', m, n, k, one_${s}$, temp, m, m4, k, zero_${s}$, r, m)
         else
             error stop "stdlib_matmul: error: unexpected s(i,j)"
         end if
@@ -77,8 +125,8 @@ contains
     pure module function stdlib_matmul_${s}$ (m1, m2, m3, m4, m5) result(r)
         ${t}$, intent(in) :: m1(:,:), m2(:,:)
         ${t}$, intent(in), optional :: m3(:,:), m4(:,:), m5(:,:)
-        ${t}$, allocatable :: r(:,:)
-        integer :: p(6), num_present
+        ${t}$, allocatable :: r(:,:), temp(:,:), temp1(:,:)
+        integer :: p(6), num_present, m, n, k
         integer, allocatable :: s(:,:)
 
         p(1) = size(m1, 1)
@@ -102,8 +150,13 @@ contains
             num_present = num_present + 1
         end if
 
+        allocate(r(p(1), p(num_present + 1)))
+
         if (num_present == 2) then
-            r = matmul(m1, m2)
+            m = p(1)
+            n = p(3)
+            k = p(2)
+            call gemm('N', 'N', m, n, k, one_${s}$, m1, m, m2, k, zero_${s}$, r, m)
             return
         end if
 
@@ -113,24 +166,56 @@ contains
         s = matmul_chain_order(p(1: num_present + 1))
 
         if (num_present == 3) then
-            r = matmul_chain_mult_${s}$_3(m1, m2, m3, 1, s)
+            r = matmul_chain_mult_${s}$_3(m1, m2, m3, 1, s, p(1:4))
             return
         else if (num_present == 4) then
-            r = matmul_chain_mult_${s}$_4(m1, m2, m3, m4, 1, s)
+            r = matmul_chain_mult_${s}$_4(m1, m2, m3, m4, 1, s, p(1:5))
             return
         end if
 
         ! Now num_present is 5
 
         select case (s(1, 5))
             case (1)
-                r = matmul(m1, matmul_chain_mult_${s}$_4(m2, m3, m4, m5, 2, s))
+                ! m1*(m2*m3*m4*m5)
+                temp = matmul_chain_mult_${s}$_4(m2, m3, m4, m5, 2, s, p)
+                m = p(1)
+                n = p(6)
+                k = p(2)
+                call gemm('N', 'N', m, n, k, one_${s}$, m1, m, temp, k, zero_${s}$, r, m)
             case (2)
-                r = matmul(matmul(m1, m2), matmul_chain_mult_${s}$_3(m3, m4, m5, 3, s))
+                ! (m1*m2)*(m3*m4*m5)
+                m = p(1)
+                n = p(3)
+                k = p(2)
+                allocate(temp(m,n))
+                call gemm('N', 'N', m, n, k, one_${s}$, m1, m, m2, k, zero_${s}$, temp, m)
+
+                temp1 = matmul_chain_mult_${s}$_3(m3, m4, m5, 3, s, p)
+
+                k = n
+                n = p(6)
+                call gemm('N', 'N', m, n, k, one_${s}$, temp, m, temp1, k, zero_${s}$, r, m)
             case (3)
-                r = matmul(matmul_chain_mult_${s}$_3(m1, m2, m3, 1, s), matmul(m4, m5))
+                ! (m1*m2*m3)*(m4*m5)
+                temp = matmul_chain_mult_${s}$_3(m1, m2, m3, 3, s, p)
+
+                m = p(4)
+                n = p(6)
+                k = p(5)
+                allocate(temp1(m,n))
+                call gemm('N', 'N', m, n, k, one_${s}$, m4, m, m5, k, zero_${s}$, temp1, m)
+
+                k = m
+                m = p(1)
+                call gemm('N', 'N', m, n, k, one_${s}$, temp, m, temp1, k, zero_${s}$, r, m)
             case (4)
-                r = matmul(matmul_chain_mult_${s}$_4(m1, m2, m3, m4, 1, s), m5)
+                ! (m1*m2*m3*m4)*m5
+                temp = matmul_chain_mult_${s}$_4(m1, m2, m3, m4, 1, s, p)
+                m = p(1)
+                n = p(6)
+                k = p(5)
+                call gemm('N', 'N', m, n, k, one_${s}$, temp, m, m5, k, zero_${s}$, r, m)
             case default
                 error stop "stdlib_matmul: error: unexpected s(i,j)"
         end select
