Skip to content

Commit 73831f9

Browse files
committed
Integrate small bidiag multiplication into general loops
1 parent 881a86f commit 73831f9

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
@@ -640,35 +640,46 @@ function _bibimul!(C, A, B, _add::MulAddMul)
640640
end
641641
function _bibimul_nonzeroalpha!(C, A, B, _add)
642642
n = size(A,1)
643-
if n <= 3
643+
if n == 1
644644
# naive multiplication
645-
for I in CartesianIndices(C)
646-
C[I] += _add(sum(A[I[1], k] * B[k, I[2]] for k in axes(A,2)))
647-
end
645+
@inbounds C[1,1] += _add(A[1,1] * B[1,1])
648646
return C
649647
end
650648
@inbounds begin
651649
# first column of C
652650
C[1,1] += _add(A[1,1]*B[1,1] + A[1, 2]*B[2,1])
653651
C[2,1] += _add(A[2,1]*B[1,1] + A[2,2]*B[2,1])
654-
C[3,1] += _add(A[3,2]*B[2,1])
652+
if n >= 3
653+
C[3,1] += _add(A[3,2]*B[2,1])
654+
end
655655
# second column of C
656656
C[1,2] += _add(A[1,1]*B[1,2] + A[1,2]*B[2,2])
657-
C[2,2] += _add(A[2,1]*B[1,2] + A[2,2]*B[2,2] + A[2,3]*B[3,2])
658-
C[3,2] += _add(A[3,2]*B[2,2] + A[3,3]*B[3,2])
659-
C[4,2] += _add(A[4,3]*B[3,2])
657+
C22 = A[2,1]*B[1,2] + A[2,2]*B[2,2]
658+
if n >= 3
659+
C[2,2] += _add(C22 + A[2,3]*B[3,2])
660+
C[3,2] += _add(A[3,2]*B[2,2] + A[3,3]*B[3,2])
661+
if n >= 4
662+
C[4,2] += _add(A[4,3]*B[3,2])
663+
end
664+
else
665+
C[2,2] += _add(C22)
666+
end
660667
end # inbounds
661668
# middle columns
662669
__bibimul!(C, A, B, _add)
663670
@inbounds begin
664-
C[n-3,n-1] += _add(A[n-3,n-2]*B[n-2,n-1])
665-
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])
666-
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])
667-
C[n, n-1] += _add(A[n,n-1]*B[n-1,n-1] + A[n,n]*B[n,n-1])
671+
if n >= 4
672+
C[n-3,n-1] += _add(A[n-3,n-2]*B[n-2,n-1])
673+
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])
674+
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])
675+
C[n, n-1] += _add(A[n,n-1]*B[n-1,n-1] + A[n,n]*B[n,n-1])
676+
end
668677
# last column of C
669-
C[n-2, n] += _add(A[n-2,n-1]*B[n-1,n])
670-
C[n-1, n] += _add(A[n-1,n-1]*B[n-1,n ] + A[n-1,n]*B[n,n ])
671-
C[n, n] += _add(A[n,n-1]*B[n-1,n ] + A[n,n]*B[n,n ])
678+
if n >= 3
679+
C[n-2, n] += _add(A[n-2,n-1]*B[n-1,n])
680+
C[n-1, n] += _add(A[n-1,n-1]*B[n-1,n ] + A[n-1,n]*B[n,n ])
681+
C[n, n] += _add(A[n,n-1]*B[n-1,n ] + A[n,n]*B[n,n ])
682+
end
672683
end # inbounds
673684
C
674685
end

0 commit comments

Comments
 (0)