mat: optimize trimul
MODIFY
mat/triangular.go
mat/triangular.go
@@ -452,10 +452,33 @@ func (t *TriDense) MulTri(a, b Triangular) {
defer restore()
}
- // TODO(btracey): Improve the set of fast-paths.
+ // Inspect types here, helps keep the loops later clean(er).
+ _, aDiag := aU.(Diagonal)
+ _, bDiag := bU.(Diagonal)
+ // If they are both diagonal only need 1 loop.
+ // All diagonal matrices are Upper.
+ // TODO: Add fast paths for DiagDense.
+ if aDiag && bDiag {
+ t.Zero()
+ for i := 0; i < n; i++ {
+ t.SetTri(i, i, a.At(i, i)*b.At(i, i))
+ }
+ return
+ }
+
+ // Now we know at least one matrix is non-diagonal.
+ // And all diagonal matrices are all Upper.
+ // The both-diagonal case is handled above.
+ // TODO: Add fast paths for Dense variants.
if kind == Upper {
for i := 0; i < n; i++ {
for j := i; j < n; j++ {
+ switch {
+ case aDiag:
+ t.SetTri(i, j, a.At(i, i)*b.At(i, j))
+ case bDiag:
+ t.SetTri(i, j, a.At(i, j)*b.At(j, j))
+ default:
var v float64
for k := i; k <= j; k++ {
v += a.At(i, k) * b.At(k, j)
@@ -463,6 +486,7 @@ func (t *TriDense) MulTri(a, b Triangular) {
t.SetTri(i, j, v)
}
}
+ }
return
}
for i := 0; i < n; i++ {

mat: optimize trimul