4
$\begingroup$

I want to optimize the code like this:

matmul[a_, b_] :=
{{vecmul[a[[1, 1]], b[[1, 1]]] + vecmul[a[[1, 2]], b[[2, 1]]], 
  vecmul[a[[1, 1]], b[[1, 2]]] + vecmul[a[[1, 2]], b[[2, 2]]]}, 
 {vecmul[a[[2, 1]], b[[1, 1]]] + vecmul[a[[2, 2]], b[[2, 1]]], 
  vecmul[a[[2, 1]], b[[1, 2]]] + vecmul[a[[2, 2]], b[[2, 2]]]}}

where a and b are matrices each of whose entries are vectors (of dimensions equal to each other but unrelated to dimensions of the matrices), and vecmul operates on pairs of such vectors. Say, each a[[i, j]] and each b[[i, j]] is a 7-dimensional vector.

I've tried

matmul[a_, b_] := Inner[vecmul[#1, #2] &, a, b]

but unfortunately Inner views a and b as tensors with three indices, tries to contract the last index of the first tensor with the first index of the last tensor and fails since they are of different dimensions. The version Inner[f, tensor1, tensor2, Plus, k] contracts kth index of tensor1 with the first index of tensor2 which in my case either gives wrong results (for k equal to 1 or 2) or an error (for k equal to 3, which is the default value in this case).

So what can I do to avoid these ugly endless redundant repetitions in the code?

$\endgroup$
2
$\begingroup$

My first guess is

matmul2[a_, b_] := TensorContract[ Outer[vecmul, a, Transpose[b], 2], {2, 4}]

Surprisingly, the idomatic way

matmul3[a_, b_] := Table[
  Sum[vecmul[a[[i, j]], b[[j, k]]], {j, 1, Dimensions[a][[2]]}]
  , {i, 1, Dimensions[a][[1]]}, {k, 1, Dimensions[b][[2]]}]

seems to be quite a bit faster for large arrays.

But much faster is the following:

With[{dot = vecmul},
  cvecmul = Compile[{{x, _Real, 2}, {y, _Real, 2}},
    Total@Table[dot[x[[j]], y[[j]]], {j, 1, Length[x]}],
    CompilationTarget -> "WVM"
    ]
  ];
matmul4[a_, b_] := Outer[cvecmul, a, Transpose[b], 1];

If you have a C compiler installed then you can also change CompilationTarget -> "WVM" to CompilationTarget -> "C" to get a further speedup.

In the case that vecmul is really Dot and nothing more fancy, you can fully vectorize everything with

matmul5[a_, b_] := Dot[
  Flatten[a, {{1}, {2, 3}}],
  Flatten[Transpose[b, {1, 3, 2}], 1]
 ];

Here is a short test suite:

vecmul = Dot;
n = 30;
m = 20;
nn = 40;
a = RandomReal[{-1, 1}, {n, m, nn}];
b = RandomReal[{-1, 1}, {m, n, nn}];
vecmul = Dot;

aa2 = matmul2[a, b]; // RepeatedTiming
aa3 = matmul3[a, b]; // RepeatedTiming
aa4 = matmul4[a, b]; // RepeatedTiming
aa5 = matmul5[a, b]; // RepeatedTiming

Max[Abs[aa2 - aa3]]
Max[Abs[aa2 - aa4]]
Max[Abs[aa2 - aa5]]

{0.10, Null}

{0.037, Null}

{0.0072, Null}

{0.00029, Null}

0.

1.06581*10^-14

1.77636*10^-14

An explanation for the poor performance of matmul2 might be that Outer has to allocate roughly m^2 = 400 times the amount of memory compared to what the Outer in the matmul3 has to allocate.

Edit

For the function vecmul prescribed by the OP in one of the comments, an integer variant of matmul4 should do the trick.

rg = 7;
ng = 12;
vecmul = With[{r = rg, n = ng},
   Compile[{{x, _Integer, 1}, {y, _Integer, 1}},
    Table[Mod[RotateRight[y, k].x, n], {k, 0, r - 1}],
    CompilationTarget -> "WVM"
    ]
   ];    
With[{dot = vecmul},
  cvecmul = Compile[{{x, _Integer, 2}, {y, _Integer, 2}},
    Total@Table[dot[x[[j]], y[[j]]], {j, 1, Length[x]}],
    CompilationTarget -> "WVM"
    ]
  ];
matmul4[a_, b_] := Outer[cvecmul, a, Transpose[b], 1];
$\endgroup$
  • $\begingroup$ Unfortunately in my case it is not Dot, rather it is the product in the group algebra of a cyclic group over a cyclic ring, vecmul[a_, b_] := Mod[a.NestList[RotateRight, b, r - 1], n] with fixed natural numbers r and n (r is a prime, but n not in general). In particular, it is a vector of the same type again, so that matmul[a, b] is also the matrix of vectors of the same type. $\endgroup$ – მამუკა ჯიბლაძე Nov 22 '17 at 8:47
  • $\begingroup$ Do you mean vecmul[a_, b_] := Mod[NestList[RotateRight, b, r - 1].a, n]? $\endgroup$ – Henrik Schumacher Nov 22 '17 at 8:51
  • $\begingroup$ No it is correct as written. b.NestList[RotateRight, a, r-1] gives the same result (the multiplication is commutative). It is supposed to represent the product of $a_1+a_2x+...+a_rx^{r-1}$ and $b_1+b_2x+...+b_rx^{r-1}$ where $x^r=1$ holds $\endgroup$ – მამუკა ჯიბლაძე Nov 22 '17 at 9:58
  • $\begingroup$ But the dimensions do not fit... $\endgroup$ – Henrik Schumacher Nov 22 '17 at 9:59
  • $\begingroup$ Oops sorry I forgot to say that r is actually equal to the lengths of both a and b! $\endgroup$ – მამუკა ჯიბლაძე Nov 22 '17 at 10:03

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.