4
$\begingroup$

Consider the following simple matrix operation ($a_j, b_j, z, w$ are matrices)

$$ w=\sum_j a_j\cdot z\cdot b_j$$

For performance reasons I am compiling the code

downShift = 
 Compile[{{z, _Complex, 2}, { a, _Real, 3}, {b, _Real, 3}}, 
  Module[{nI, nF, ns, j},
    {ns, nF, nI} = Dimensions[a];
    Sum[a[[j]].z.b[[j]], {j, ns}]
   ]]

It brings no speed up. Trying to investigate the issue reveals the following huge code

Needs["CompiledFunctionTools`"]
CompilePrint[downShift]

3 arguments
        1 Boolean register
        19 Integer registers
        1 Complex register
        11 Tensor registers
        Underflow checking off
        Overflow checking off
        Integer overflow checking on
        RuntimeAttributes -> {}

    T(C2)0 = A1
    T(R3)1 = A2
    T(R3)2 = A3
    I17 = 0
    I6 = 4
    C0 = 0. + 0. I
    I3 = 2
    I7 = 12
    I2 = 1
    I1 = 3
    Result = T(C2)8

1   T(I1)3 = Dimensions[ T(R3)1]]
2   I0 = Length[ T(I1)3]
3   B0 = I0 == I1
4   B0 = ! B0
5   if[ !B0] goto 8
6   Return Error
7   goto 8
8   I0 = GetElement[ T(I1)3, I2]
9   I4 = GetElement[ T(I1)3, I3]
10  I5 = GetElement[ T(I1)3, I1]
11  T(R2)8 = Part[ T(R3)1, I0]
12  T(C2)5 = CoerceTensor[ I6, T(R2)8]]
13  T(C2)4 = Dot[ T(C2)5, T(C2)0, I7]]
14  T(R2)8 = Part[ T(R3)2, I0]
15  T(C2)5 = CoerceTensor[ I6, T(R2)8]]
16  T(C2)9 = Dot[ T(C2)4, T(C2)5, I7]]
17  T(C3)4 = {T(C2)9}
18  T(C2)9 = Part[ T(C3)4, I2]
19  I12 = Length[ T(C2)9]
20  T(R2)9 = Part[ T(R3)1, I0]
21  T(C2)4 = CoerceTensor[ I6, T(R2)9]]
22  T(C2)8 = Dot[ T(C2)4, T(C2)0, I7]]
23  T(R2)9 = Part[ T(R3)2, I0]
24  T(C2)4 = CoerceTensor[ I6, T(R2)9]]
25  T(C2)5 = Dot[ T(C2)8, T(C2)4, I7]]
26  T(C3)8 = {T(C2)5}
27  T(C2)5 = Part[ T(C3)8, I2]
28  T(C1)8 = Part[ T(C2)5, I2]
29  I16 = Length[ T(C1)8]
30  I14 = I17
31  T(C2)8 = Table[ I12, I16]
32  I15 = I17
33  goto 38
34  I11 = I17
35  goto 37
36  Element[ T(C2)8, I14] = C0
37  if[ ++ I11 <= I16] goto 36
38  if[ ++ I15 <= I12] goto 34
39  I8 = I0
40  I9 = I17
41  goto 50
42  T(R2)5 = Part[ T(R3)1, I9]
43  T(C2)9 = CoerceTensor[ I6, T(R2)5]]
44  T(C2)4 = Dot[ T(C2)9, T(C2)0, I7]]
45  T(R2)5 = Part[ T(R3)2, I9]
46  T(C2)9 = CoerceTensor[ I6, T(R2)5]]
47  T(C2)10 = Dot[ T(C2)4, T(C2)9, I7]]
48  T(C2)4 = T(C2)8 + T(C2)10
49  T(C2)8 = CopyTensor[ T(C2)4]]
50  if[ ++ I9 <= I8] goto 42
51  Return

Why the code is so complicated if it can be realized with just a couple of calls to BLAS subroutines? How to speed up the code?

There is an important update to the question: $$\text{ns} \ll \text{nI} \approx \text{nF}.$$

Typical dimensions are $\text{ns}=12$, $\text{nI}=\text{nF}=600$. The code has to be executed repeatedly inside NDSolve.

$\endgroup$
3
  • $\begingroup$ Please post the dimension of z, a, b` at which you target. BLAS is really slow with many low-dimensional matrices; so in this case one wants to write out the dot-products exicitly (as George Varnavides already did). This can be significantly improved if the dimensions of z and the last two dimensions of a and b are already known at comple time. $\endgroup$ Feb 27, 2023 at 6:33
  • $\begingroup$ The call to CoerceTensor must have to do with the fact that a and b are real tensors, while z is complex. Since BLAS can only do "pure" dot products in a single type, chunks of a and b are converted first. It might be a good idea to treat Re[z] and Im[z] independently first, and to add the results only in the end. This would use only real artihmetic, skip the conversions, and save half of the flops. $\endgroup$ Feb 27, 2023 at 6:38
  • $\begingroup$ @HenrikSchumacher All your remarks are very relevant. I updated the post. To separately treat real and imaginary parts is an excellent idea, however, it still involves an overhead related to the copy. Unfortunately, BLAS cannot do DGEMM with storage spacing between elements equal to 2. $\endgroup$
    – yarchik
    Feb 27, 2023 at 8:55

3 Answers 3

6
$\begingroup$

You probably don't want to see it, but here is a LibraryLink implementation. After all, I was curious as to how fast one could get this in principle.

While the code is basically platform independent, I rely here on Apple's Accelerate framework. I have also marked some places that you would have to modify to make this work on other platforms...

I tried also OpenBLAS -- and found out that it would be about 4 times slower. (Maybe because OpenBLAS cannot source on Apple's the secret matrix processing unit.)

Here is the code:

Needs["CCompilerDriver`"]

Quiet[LibraryFunctionUnload[downShift6]];
ClearAll[downShift6];
downShift6 = Module[{lib,file,name},
    name = "downShift6";
    file=Export[
        FileNameJoin[{$TemporaryDirectory,name<>".cpp"}],
        StringJoin["
            #include \"WolframLibrary.h\"
            
            #include <algorithm>

            #include <Accelerate/Accelerate.h>
            //#include \"openblas.h\"

            EXTERN_C DLLEXPORT int "<>name<>"(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res)
            {
                MTensor a_ = MArgument_getMTensor(Args[0]); 
                MTensor z_ = MArgument_getMTensor(Args[1]);
                MTensor b_ = MArgument_getMTensor(Args[2]); 
                
                const mint n  = libData->MTensor_getDimensions(a_)[0];  
                const mint n0 = libData->MTensor_getDimensions(a_)[1];  
                const mint n1 = libData->MTensor_getDimensions(a_)[2];
                const mint n2 = libData->MTensor_getDimensions(b_)[1];  
                const mint n3 = libData->MTensor_getDimensions(b_)[2];  
                
                // Create MTensor for the result.
                MTensor c_;
                
                const mint dims [2] {n0, n3};
                
                (void)libData->MTensor_new(MType_Complex, 2, dims, &c_);
                
                const mreal * const a     = libData->MTensor_getRealData(a_);
                const mreal * const b     = libData->MTensor_getRealData(b_);
                const mreal * const z_buf = reinterpret_cast<mreal*>(libData->MTensor_getComplexData(z_));
                      mreal * const c_buf = reinterpret_cast<mreal*>(libData->MTensor_getComplexData(c_));
                
                // Buffers for real and imaginary parts of result.
                mreal * c [2] = { new mreal [n0 * n3], new mreal [n0 * n3]};

                // Some scratch space.
                mreal * const z = new mreal [n1 * n2];
                mreal * const t = new mreal [n1 * n3];
                    
                for( mint k = 0; k < 2; ++k )
                {
                    std::fill( &c[k][0], &c[k][n0 * n3], static_cast<mreal>(0) );

                    // Read only real/imaginary parts of z_buf, strided by 2.;
                    cblas_dcopy( n1 * n2, &z_buf[k], 2, z, 1 );
                    
                    for( mint i = 0; i < n; ++i )
                    {   
                        // Compute t = z * b;
                        cblas_dgemm( CblasRowMajor, CblasNoTrans, CblasNoTrans, 
                            n1, n3, n2, 1., z, n2, &b[n2*n3*i], n2, 0., t, n3 
                        );
                    
                        // Compute c += a * t;
                        cblas_dgemm( CblasRowMajor, CblasNoTrans, CblasNoTrans, 
                            n0, n3, n1, 1., &a[n0*n1*i], n1, t, n3, 1., c[k], n3 
                        );
                    }
                }

                // Copy real and imaginary parts into c_buf.

                // Write real parts to c_buf, strided by 2.
                cblas_dcopy( n0 * n3, c[0], 1, &c_buf[0], 2 );

                // Write imaginary parts to c_buf, strided by 2.
                cblas_dcopy( n0 * n3, c[1], 1, &c_buf[1], 2 );

                // Free the scratch space.
                delete[] t;
                delete[] z;
                delete[] c[0];
                delete[] c[1];
                
                MArgument_setMTensor(Res, c_);
                
                return LIBRARY_NO_ERROR;
            }"
        ],
        "Text"
    ];
    
    lib=CreateLibrary[{file},name,      
        "CompileOptions"->{
            " -Wall"
            ,"-Wextra"
            ,"-Wno-unused-parameter"
            ,"-std=c++11"
            ,"-Ofast"
            ,"-flto"
            ,"-framework Accelerate"
            (*,"-lopenblas"*)
        }
        ,"IncludeDirectories" -> {(*Put the path to OpenBLAS' openblas.h header here.*)}
        ,"LibraryDirectories" -> {(*Put the path to OpenBLAS' library here.*)}
        ,"ShellOutputFunction"->Print
    ];
    
    LibraryFunctionLoad[lib,name, {{Real,3,"Constant"},{Complex,2,"Constant"},{Real,3,"Constant"}}, {Complex,2}]
];

Here a usage example:

nn = 12;
n0 = n1 = n2 = n3 = 600;
a = RandomReal[{-1, 1}, {nn, n0, n1}];
z = RandomComplex[{-1 - I, 1 + I}, {n1, n2}];
b = RandomReal[{-1, 1}, {nn, n2, n3}];


c0 = Activate[TensorContract[Inactive[TensorProduct][a, z, b], {{1, 6}, {3, 4}, {5, 7}}]]; // RepeatedTiming // First

c6 = downShift6[a, z, b]; // RepeatedTiming // First

Max[Abs[c0 - c6]]/Max[Abs[c0]]

0.238828

0.0428804

7.13858*10^-15

What struck me was that TensorContract required quite exactly as long as my OpenBLAS version. Maybe that is because Mathematica is not linked against the Accelerate framework? Would be great to hear about your experiences.

$\endgroup$
10
  • 2
    $\begingroup$ Impressive! I will test it in a real-life situation and give feedback. $\endgroup$
    – yarchik
    Feb 27, 2023 at 18:47
  • 1
    $\begingroup$ I tested your code and I am quite impressed with the performance. On my old MacBook air 2015 I get only 2-fold speed up. On M1 pro, I get results similar to yours, i.e., 5-6 times speed up. I notice your dcopy tricks with the stride 2. I also use this when programming in fortran, although I have been heavily criticized for that. So I am very happy that I am not alone. I am aware though that it is a bit dangerous. $\endgroup$
    – yarchik
    Feb 27, 2023 at 20:14
  • $\begingroup$ Great work! I don't think it's a good idea to write the source code into a file and then reading it back in. The documentation describes as the first usage to use a string as first argument to CreateLibrary. $\endgroup$
    – Roman
    Feb 27, 2023 at 20:21
  • 1
    $\begingroup$ @yarchik stride-2 tricks with complex numbers are standard practice for interfacing S/D-type with C/Z-type BLAS functions. I'm surprised that you were criticised for it! $\endgroup$
    – Roman
    Feb 27, 2023 at 20:26
  • 1
    $\begingroup$ @HenrikSchumacher you'll have to specify "Language" -> "C++" because the string interface assumes C by default. $\endgroup$
    – Roman
    Feb 27, 2023 at 20:34
2
$\begingroup$

Haven't looked at why Compile is so slow here, presumably writing the dot products out using Sum and Table indices could help.

Does this work for you though?

TensorContract[
    Inactive[TensorProduct][a, z, b],
    {{1, 6}, {3, 4}, {5, 7}}] // Activate; // AbsoluteTiming

Not sure how large your arrays are, but this is fairly fast on my machine for {ns,nF,nI} = {1000,100,100}.

$\endgroup$
2
  • $\begingroup$ My parameters are such that ns$\ll$nI, nF. I updated the post with relevant information. $\endgroup$
    – yarchik
    Feb 27, 2023 at 8:57
  • $\begingroup$ Indeed, I see a small improvement for a realistic usage case (+1). Unfortunately, it is still too slow to be practically useful. $\endgroup$
    – yarchik
    Feb 27, 2023 at 9:26
2
$\begingroup$

This is about a factor of two faster than the compiled code, and could probably be translated to BLAS calls in an external C function without much trouble:

ns = 12;
nI = 600;
nF = 610;
a = RandomVariate[NormalDistribution[], {ns, nI, nF}];
z = RandomVariate[NormalDistribution[], {nF, nF, 2}] . {1, I};
b = RandomVariate[NormalDistribution[], {ns, nF, nI}];

Straight sum:

w1 = Sum[a[[j]] . z . b[[j]], {j, ns}]; // AbsoluteTiming // First
(*    0.234662    *)

Compiled code:

w2 = downShift[z, a, b]; // AbsoluteTiming // First
(*    0.231027    *)

Real-valued explicit tensor contraction:

w3 = ArrayReshape[Transpose[a . Re[z]], {nI, ns*nF}] .
       ArrayReshape[b, {ns*nF, nI}] +
     I*ArrayReshape[Transpose[a . Im[z]], {nI, ns*nF}] .
       ArrayReshape[b, {ns*nF, nI}]; // AbsoluteTiming // First
(*    0.107337    *)
$\endgroup$
4
  • $\begingroup$ Yes, it offers stable performance improvement (+). I was hoping I can compile your solution, but ArrayReshape forces the main evaluation. Just to be consistent with OP you may exchange nI and nF. $\endgroup$
    – yarchik
    Feb 27, 2023 at 13:22
  • $\begingroup$ You can compile it by using the full form of the size lists: {1, nI, ns*nF} and {ns*nF, nI, 1}. But compiling does not speed up the evaluation. The speed seems limited by the actual calculation. $\endgroup$
    – Roman
    Feb 27, 2023 at 14:56
  • $\begingroup$ Are you going to use the matrix $w$ as-is or are you going to contract it further? If you are going to contract it further, then it would be good to see the full calculation here and propose more global optimizations. $\endgroup$
    – Roman
    Feb 27, 2023 at 15:42
  • $\begingroup$ No, I do not contract it further, only flatten, but this is not the bottleneck and it can even be avoided. So the function in the OP is essentially what I need. $\endgroup$
    – yarchik
    Feb 27, 2023 at 18:46

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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