Skip to content

Commit 36615aa

Browse files
committed
Integrate small bidiag multiplication into general loops
1 parent df7fef1 commit 36615aa

File tree

1 file changed

+26
-15
lines changed

1 file changed

+26
-15
lines changed

src/bidiag.jl

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -629,35 +629,46 @@ function _bibimul!(C, A, B, _add::MulAddMul)
629629
end
630630
function _bibimul_nonzeroalpha!(C, A, B, _add)
631631
n = size(A,1)
632-
if n <= 3
632+
if n == 1
633633
# naive multiplication
634-
for I in CartesianIndices(C)
635-
C[I] += _add(sum(A[I[1], k] * B[k, I[2]] for k in axes(A,2)))
636-
end
634+
@inbounds C[1,1] += _add(A[1,1] * B[1,1])
637635
return C
638636
end
639637
@inbounds begin
640638
# first column of C
641639
C[1,1] += _add(A[1,1]*B[1,1] + A[1, 2]*B[2,1])
642640
C[2,1] += _add(A[2,1]*B[1,1] + A[2,2]*B[2,1])
643-
C[3,1] += _add(A[3,2]*B[2,1])
641+
if n >= 3
642+
C[3,1] += _add(A[3,2]*B[2,1])
643+
end
644644
# second column of C
645645
C[1,2] += _add(A[1,1]*B[1,2] + A[1,2]*B[2,2])
646-
C[2,2] += _add(A[2,1]*B[1,2] + A[2,2]*B[2,2] + A[2,3]*B[3,2])
647-
C[3,2] += _add(A[3,2]*B[2,2] + A[3,3]*B[3,2])
648-
C[4,2] += _add(A[4,3]*B[3,2])
646+
C22 = A[2,1]*B[1,2] + A[2,2]*B[2,2]
647+
if n >= 3
648+
C[2,2] += _add(C22 + A[2,3]*B[3,2])
649+
C[3,2] += _add(A[3,2]*B[2,2] + A[3,3]*B[3,2])
650+
if n >= 4
651+
C[4,2] += _add(A[4,3]*B[3,2])
652+
end
653+
else
654+
C[2,2] += _add(C22)
655+
end
649656
end # inbounds
650657
# middle columns
651658
__bibimul!(C, A, B, _add)
652659
@inbounds begin
653-
C[n-3,n-1] += _add(A[n-3,n-2]*B[n-2,n-1])
654-
C[n-2,n-1] += _add(A[n-2,n-2]*B[n-2,n-1] + A[n-2,n-1]*B[n-1,n-1])
655-
C[n-1,n-1] += _add(A[n-1,n-2]*B[n-2,n-1] + A[n-1,n-1]*B[n-1,n-1] + A[n-1,n]*B[n,n-1])
656-
C[n, n-1] += _add(A[n,n-1]*B[n-1,n-1] + A[n,n]*B[n,n-1])
660+
if n >= 4
661+
C[n-3,n-1] += _add(A[n-3,n-2]*B[n-2,n-1])
662+
C[n-2,n-1] += _add(A[n-2,n-2]*B[n-2,n-1] + A[n-2,n-1]*B[n-1,n-1])
663+
C[n-1,n-1] += _add(A[n-1,n-2]*B[n-2,n-1] + A[n-1,n-1]*B[n-1,n-1] + A[n-1,n]*B[n,n-1])
664+
C[n, n-1] += _add(A[n,n-1]*B[n-1,n-1] + A[n,n]*B[n,n-1])
665+
end
657666
# last column of C
658-
C[n-2, n] += _add(A[n-2,n-1]*B[n-1,n])
659-
C[n-1, n] += _add(A[n-1,n-1]*B[n-1,n ] + A[n-1,n]*B[n,n ])
660-
C[n, n] += _add(A[n,n-1]*B[n-1,n ] + A[n,n]*B[n,n ])
667+
if n >= 3
668+
C[n-2, n] += _add(A[n-2,n-1]*B[n-1,n])
669+
C[n-1, n] += _add(A[n-1,n-1]*B[n-1,n ] + A[n-1,n]*B[n,n ])
670+
C[n, n] += _add(A[n,n-1]*B[n-1,n ] + A[n,n]*B[n,n ])
671+
end
661672
end # inbounds
662673
C
663674
end

0 commit comments

Comments
 (0)