-
Notifications
You must be signed in to change notification settings - Fork 189
feat: extend intrinsic matmul
#951
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from 12 commits
f06f556
fed4d73
27911ae
a7f645c
cc77dee
3958018
35a5a28
ebf92d7
5f5c5a9
06ce735
e709f83
cf5f030
b6d07e6
5e3b588
7d2130a
61851fc
ee7da8d
195c57e
e71b9bb
00c4461
5c2bbc5
79113da
72dc641
cee5bba
a052599
24c5787
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,3 @@ | ||
ADD_EXAMPLE(sum) | ||
ADD_EXAMPLE(dot_product) | ||
ADD_EXAMPLE(dot_product) | ||
ADD_EXAMPLE(matmul) |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,19 @@ | ||
program example_matmul | ||
use stdlib_intrinsics, only: stdlib_matmul | ||
complex :: x(2, 2), y(2, 2) | ||
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 | ||
|
||
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() | ||
call random_number(r1) | ||
call random_number(r2) | ||
call random_number(r3) | ||
|
||
res = stdlib_matmul(r1, r2, r3) ! 50x50 matrix | ||
print *, shape(res) | ||
end program example_matmul |
Original file line number | Diff line number | Diff line change | ||||||
---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,281 @@ | ||||||||
#:include "common.fypp" | ||||||||
#:set I_KINDS_TYPES = list(zip(INT_KINDS, INT_TYPES, INT_KINDS)) | ||||||||
#:set R_KINDS_TYPES = list(zip(REAL_KINDS, REAL_TYPES, REAL_SUFFIX)) | ||||||||
#: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_linalg_state, only: linalg_state_type, linalg_error_handling, LINALG_VALUE_ERROR | ||||||||
use stdlib_constants | ||||||||
implicit none | ||||||||
|
||||||||
character(len=*), parameter :: this = "stdlib_matmul" | ||||||||
|
||||||||
contains | ||||||||
|
||||||||
! Algorithm for the optimal parenthesization of matrices | ||||||||
! Reference: Cormen, "Introduction to Algorithms", 4ed, ch-14, section-2 | ||||||||
! Internal use only! | ||||||||
pure function matmul_chain_order(p) result(s) | ||||||||
integer, intent(in) :: p(:) | ||||||||
integer :: s(1:size(p) - 2, 2:size(p) - 1), m(1:size(p) - 1, 1:size(p) - 1) | ||||||||
integer :: n, l, i, j, k, q | ||||||||
n = size(p) - 1 | ||||||||
m(:,:) = 0 | ||||||||
s(:,:) = 0 | ||||||||
|
||||||||
do l = 2, n | ||||||||
do i = 1, n - l + 1 | ||||||||
j = i + l - 1 | ||||||||
m(i,j) = huge(1) | ||||||||
|
||||||||
do k = i, j - 1 | ||||||||
q = m(i,k) + m(k+1,j) + p(i)*p(k+1)*p(j+1) | ||||||||
|
||||||||
if (q < m(i, j)) then | ||||||||
m(i,j) = q | ||||||||
s(i,j) = k | ||||||||
end if | ||||||||
end do | ||||||||
end do | ||||||||
end do | ||||||||
end function matmul_chain_order | ||||||||
|
||||||||
#:for k, t, s in R_KINDS_TYPES + C_KINDS_TYPES | ||||||||
|
||||||||
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:), 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) | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Very good progress @wassup05, thank you! Imho this PR is almost ready to be merged. As you suggest, it would be good to have a nice wrapper for Here I suggest two possible APIs, and I will ask @jalvesz @jvdp1 @loiseaujc to discuss that together: The first would be similar to ! API Similar to gemm
${t1}$ function stdlib_matmul(A, Astate, B, Bstate) result(C)
${t1}$, intent(in) :: A(:,:), B(:,:)
character, intent(in), optional :: Astate, Bstate and could use the matrix state definitions already in use for the sparse operations stdlib/src/stdlib_sparse_constants.fypp Lines 16 to 18 in 5c64ee6
The second would be more ambitious and essentially zero-overhead, it would wrap the operation in a derived type: (to be templated of course) type :: matrix_state_rdp
real(dp), pointer :: A(:,:) => null()
character(1) :: Astate = 'N'
end type matrix_state_rdp
interface transposed
module procedure transposed_new_rdp
...
end interface transposed
type(matrix_state_rdp) function transposed_new_rdp(A) result(AT)
real(dp), intent(inout), target :: A(:,:)
AT%A => A
AT%Astate = 'T'
end function Then we could define a templated base interface
So the user writing code would have it clear:
we could even make it an operator:
without it triggering any actual data movement. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
How would that work in the written code? Would There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
I also had the same reaction at first, after thinking about it longer I saw that it is actually a pretty clever solution, the user would declare the matrices as regular dense matrices. It is the internal interface that would make the distinction. This would imply though implementing internally several versions to account for the combinations (dense,dense) / (dense,type) / (type,dense). This looks interesting but I wonder if it should be pursued at this stage. The first proposal by @perazz looks easier and totally valid but I would propose it with a slight modification: ${t1}$ function stdlib_matmul(A, B, op_a, op_b) result(C)
${t1}$, intent(in) :: A(:,:), B(:,:)
character, intent(in), optional :: op_a, op_b to let all optional arguments of a procedure at the end of the signature. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes @jalvesz @loiseaujc, I don't know how to write it better, but it is outlined in the previous post.
Only function 4. is actually implemented, and is a There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Oh I see. I quite like that indeed. I still believe though that, as a starting point, restricting ourselves to standard stuff might be easier. Introducing new derived types to represent matrices definitely is something I'm looking forward to but, in line with the discussion here, it might require a broader discussion to have a well-designed set of derived types. I also prefer this signature
for the reasons you've mentioned. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Well, here I would define these "convenience" derived types, because they're only used to make the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I like the first proposal quite a lot too! the 2nd proposal seems convenient too with the |
||||||||
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, p) result(r) | ||||||||
${t}$, intent(in) :: m1(:,:), m2(:,:), m3(:,:), 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 | ||||||||
|
||||||||
end function matmul_chain_mult_${s}$_4 | ||||||||
|
||||||||
pure module subroutine stdlib_matmul_sub_${s}$ (res, m1, m2, m3, m4, m5, err) | ||||||||
${t}$, intent(out), allocatable :: res(:,:) | ||||||||
${t}$, intent(in) :: m1(:,:), m2(:,:) | ||||||||
${t}$, intent(in), optional :: m3(:,:), m4(:,:), m5(:,:) | ||||||||
type(linalg_state_type), intent(out), optional :: err | ||||||||
${t}$, allocatable :: temp(:,:), temp1(:,:) | ||||||||
integer :: p(6), num_present, m, n, k | ||||||||
integer, allocatable :: s(:,:) | ||||||||
|
||||||||
type(linalg_state_type) :: err0 | ||||||||
|
||||||||
p(1) = size(m1, 1) | ||||||||
p(2) = size(m2, 1) | ||||||||
p(3) = size(m2, 2) | ||||||||
|
||||||||
if (size(m1, 2) /= p(2)) then | ||||||||
err0 = linalg_state_type(this, LINALG_VALUE_ERROR, 'matrices m1, m2 not of compatible sizes') | ||||||||
perazz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
call linalg_error_handling(err0, err) | ||||||||
allocate(res(0, 0)) | ||||||||
return | ||||||||
end if | ||||||||
|
||||||||
num_present = 2 | ||||||||
if (present(m3)) then | ||||||||
|
||||||||
if (size(m3, 1) /= p(3)) then | ||||||||
err0 = linalg_state_type(this, LINALG_VALUE_ERROR, 'matrices m2, m3 not of compatible sizes') | ||||||||
perazz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
call linalg_error_handling(err0, err) | ||||||||
allocate(res(0, 0)) | ||||||||
return | ||||||||
end if | ||||||||
|
||||||||
p(3) = size(m3, 1) | ||||||||
wassup05 marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
p(4) = size(m3, 2) | ||||||||
num_present = num_present + 1 | ||||||||
end if | ||||||||
if (present(m4)) then | ||||||||
|
||||||||
if (size(m4, 1) /= p(4)) then | ||||||||
err0 = linalg_state_type(this, LINALG_VALUE_ERROR, 'matrices m3, m4 not of compatible sizes') | ||||||||
perazz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
call linalg_error_handling(err0, err) | ||||||||
allocate(res(0, 0)) | ||||||||
return | ||||||||
end if | ||||||||
|
||||||||
p(4) = size(m4, 1) | ||||||||
p(5) = size(m4, 2) | ||||||||
num_present = num_present + 1 | ||||||||
end if | ||||||||
if (present(m5)) then | ||||||||
|
||||||||
if (size(m5, 1) /= p(5)) then | ||||||||
err0 = linalg_state_type(this, LINALG_VALUE_ERROR, 'matrices m4, m5 not of compatible sizes') | ||||||||
perazz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
call linalg_error_handling(err0, err) | ||||||||
allocate(res(0, 0)) | ||||||||
return | ||||||||
end if | ||||||||
|
||||||||
p(5) = size(m5, 1) | ||||||||
p(6) = size(m5, 2) | ||||||||
num_present = num_present + 1 | ||||||||
end if | ||||||||
|
||||||||
allocate(res(p(1), p(num_present + 1))) | ||||||||
|
||||||||
if (num_present == 2) then | ||||||||
m = p(1) | ||||||||
n = p(3) | ||||||||
k = p(2) | ||||||||
call gemm('N', 'N', m, n, k, one_${s}$, m1, m, m2, k, zero_${s}$, res, m) | ||||||||
return | ||||||||
end if | ||||||||
|
||||||||
! Now num_present >= 3 | ||||||||
allocate(s(1:num_present - 1, 2:num_present)) | ||||||||
|
||||||||
s = matmul_chain_order(p(1: num_present + 1)) | ||||||||
|
||||||||
if (num_present == 3) then | ||||||||
res = matmul_chain_mult_${s}$_3(m1, m2, m3, 1, s, p(1:4)) | ||||||||
return | ||||||||
else if (num_present == 4) then | ||||||||
res = 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) | ||||||||
! 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}$, res, m) | ||||||||
case (2) | ||||||||
! (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}$, res, m) | ||||||||
case (3) | ||||||||
! (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}$, res, m) | ||||||||
case (4) | ||||||||
! (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}$, res, m) | ||||||||
case default | ||||||||
error stop "stdlib_matmul: internal error: unexpected s(i,j)" | ||||||||
perazz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
end select | ||||||||
|
||||||||
end subroutine stdlib_matmul_sub_${s}$ | ||||||||
|
||||||||
pure module function stdlib_matmul_pure_${s}$ (m1, m2, m3, m4, m5) result(r) | ||||||||
${t}$, intent(in) :: m1(:,:), m2(:,:) | ||||||||
${t}$, intent(in), optional :: m3(:,:), m4(:,:), m5(:,:) | ||||||||
${t}$, allocatable :: r(:,:) | ||||||||
|
||||||||
call stdlib_matmul_sub(r, m1, m2, m3, m4, m5) | ||||||||
end function stdlib_matmul_pure_${s}$ | ||||||||
|
||||||||
module function stdlib_matmul_${s}$ (m1, m2, m3, m4, m5, err) result(r) | ||||||||
${t}$, intent(in) :: m1(:,:), m2(:,:) | ||||||||
${t}$, intent(in), optional :: m3(:,:), m4(:,:), m5(:,:) | ||||||||
type(linalg_state_type), intent(out) :: err | ||||||||
${t}$, allocatable :: r(:,:) | ||||||||
|
||||||||
call stdlib_matmul_sub(r, m1, m2, m3, m4, m5, err=err) | ||||||||
end function stdlib_matmul_${s}$ | ||||||||
|
||||||||
#:endfor | ||||||||
end submodule stdlib_intrinsics_matmul |
Uh oh!
There was an error while loading. Please reload this page.