C++ Eigen/Mex function to perform two pagewise convolutions on a 3D matrix with a simple kernel

Question

I am trying to accelerate an algorithm where I am computing a 2D variance of many complex matrices stored as pages of a 3D array in Matlab. My implementation uses a Mex function with the C++ Matlab Data API. I perform Toeplitz convolution along both dimensions of the matrix with a simple kernel of ones to compute a new matrix of equal size. I then multiply the original matrix by the computed 2D variance matrix and store the sum of the result in my output. This "2D variance" is normalized by the sum of the magnitudes of the corresponding values. Here is a figure illustrating the concept of how the variance is calculated:

There are a few things I think could be improved to make this algorithm and/or code more efficient that I haven't been able to figure out:

1. Instead of Toeplitz convolution with matrices Tband and Rband for each column and row of each page respectively, I could try some sort of fourier convolution. I haven't been able to accurately implement this where the outputs from both codes match...
2. Is the use of temporaries with const auto and limiting evaluation until the end the most efficient way to implement the inside of the parallelized for loop? I tried preallocating some matrices that stored enough data for the number of threads of the given machine (measured at runtime), but I wasn't able to get better performance that way...
3. Is there a way to perform these evaluations in Eigen without a for loop?
4. Are there any other obvious optimizations I'm missing here?

Here's the code:

#include <iostream>
#include "mex.hpp"
#include "mexAdapter.hpp"
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <MatlabDataArray.hpp>
#include <cmath>
#include <math.h>
#include <omp.h>

class MexFunction : public matlab::mex::Function {
public:
    
    // Pointer to MATLAB engine to call fprintf
    std::shared_ptr<matlab::engine::MATLABEngine> matlabPtr = getEngine();

    // Factory to create MATLAB data arrays
    matlab::data::ArrayFactory factory;
    
    // Guarantee Alignment (not sure if this helps)
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    
    void operator()(matlab::mex::ArgumentList outputs, matlab::mex::ArgumentList inputs) {
        checkArguments(outputs, inputs);
        
        // Initialize input parameters
        long int numEl = inputs[0][0];
        long int na = inputs[0][1];
        long int nPoints = inputs[0][2];
        long int szZ = inputs[0][3];
        long int szX = inputs[0][4];
        
        // Initialize input idxtMTX data
        auto ptr = getDataPtr<std::complex<double>>(inputs[1]);
        Eigen::Map< const Eigen::MatrixXcd > idxt( ptr, numEl*na, nPoints );
        
        // Initialize Toeplitz band matrix inputs
        auto ptr2 = getDataPtr<double>(inputs[2]);
        Eigen::Map< const Eigen::MatrixXd > Rband( ptr2, numEl, numEl );
        
        auto ptr3 = getDataPtr<double>(inputs[3]);
        Eigen::Map< const Eigen::MatrixXd > Tband( ptr3, na, na);
        
        outputs[0] = factory.createArray<std::complex<double>>({static_cast<size_t>(szZ),static_cast<size_t>(szX)});
        
        auto ptrRecon = getOutDataPtr<std::complex<double>>(outputs[0]);
        Eigen::Map<Eigen::MatrixXcd> Recon(ptrRecon,szZ,szX);

        #pragma omp parallel for
        for (int i = 0; i < nPoints; i++) {

            // Define temporary pages of complex and absolute data
            const auto idxtCol = idxt.col(i).reshaped(numEl,na);
            const auto idxt2 = idxt.col(i).reshaped(numEl,na).cwiseAbs();
            
            // Compute Numerator and Denominator respectively
            const auto s1 = (Rband*idxt2).array()*(idxt2*Tband).array();
            const auto s2 = ((Rband*idxtCol).array()*(idxtCol*Tband).conjugate().array()).abs();
            
            // Compute result for only values where the denominator is nonzero. Otherwise set to zero. Then multiply by weights
            const auto weight = (s1 != 0).select(s2/s1,0.0);
            const auto result = idxtCol.array()*weight;
            
            // Sum page and assign to output
            Recon(i%szZ,i/szZ) = result.sum();
        }
    }

    void checkArguments(matlab::mex::ArgumentList outputs, matlab::mex::ArgumentList inputs) {
        std::shared_ptr<matlab::engine::MATLABEngine> matlabPtr = getEngine();
        matlab::data::ArrayFactory factory;

        if (inputs.size() != 4) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Five inputs required") }));
        }

        if (inputs[0].getNumberOfElements() != 5) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Need 3 input parameters") }));
        }
        
        if (inputs[0].getType() != matlab::data::ArrayType::INT32) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input parameter must be integer") }));
        }

        if (inputs[1].getType() == matlab::data::ArrayType::DOUBLE ||
            inputs[1].getType() != matlab::data::ArrayType::COMPLEX_DOUBLE) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input idxtMTX must be type complex double") }));
        }

        if (inputs[1].getDimensions().size() != 2) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input must be m-by-n dimension") }));
        }
        
        if (inputs[2].getType() != matlab::data::ArrayType::DOUBLE ||
            inputs[2].getType() == matlab::data::ArrayType::COMPLEX_DOUBLE) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input Rband must be type double") }));
        }

        if (inputs[2].getDimensions().size() != 2) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input must be m-by-n dimension") }));
        }
        
        if (inputs[3].getType() != matlab::data::ArrayType::DOUBLE ||
            inputs[3].getType() == matlab::data::ArrayType::COMPLEX_DOUBLE) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input Tband must be type double") }));
        }

        if (inputs[3].getDimensions().size() != 2) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Input must be m-by-n dimension") }));
        }

    }
    
    template <typename T>
    const T* getDataPtr(matlab::data::Array arr) {
        const matlab::data::TypedArray<T> arr_t = arr;
        matlab::data::TypedIterator<const T> it(arr_t.begin());
        return it.operator->();
    }
    
    template <typename T>
    T* getOutDataPtr(matlab::data::Array& arr) {
      auto range = matlab::data::getWritableElements<T>(arr);
      return range.begin().operator->();
    }

};

Answer 1

This is a great question in most respects, but since the code sample doesn’t compile and can’t be tested, I can only give you a vague and speculative answer.

I’ll skip over your Item 1, because I have no idea whether it would work for your use case.

Array Indices are Never int

Everyone does this, but it’s wrong and dangerous. On many targets, int is too narrow to hold an array index and and could overflow. On some targets, it can index no more than 32,768 elements, but more commonly, objects can be eight billion times larger than an int can address. Since it is a signed type, overflow is undefined behavior. Compilers in 2023 take that as permission to break your program in arbitrary ways, including inserting a buffer-overrun exploit, because it was your responsibility to use the right type.

Even if you’re sure nothing will ever go wrong, you meant to use size_t.

Static Single Assignments Work Great—in C

In fact, (most?) optimizers in 2023 will convert your code to either SSA or an equivalent form, such as continuation-passing style (CPS) or administrative normal form (ANF). Modern compilers are great at optimizing static single assignments on Plain Old Data, like scalars and C-style arrays.

When you write const auto foo in C++, it gets more complicated, and things might be going on behind your back:

• C++ constructors and destructors can have arbitrary side-effects that you do not want inside a high-performance loop, such as allocating memory from the heap.
• C++ functions can return data by reference, in which case you want to declare const auto& foo so you do not copy it.
• It might be more efficient to move data than to copy it. In that case, you want auto foo = bar(); baz(std::move(foo));, or just baz(bar());. However, most classes with move semantics use dynamic memory, so you’d want to avoid creating and destroying them inside the loop body anyway.

I’m not familiar enough with Eigen to tell you if that is or isn’t here (although the fact you’re initializing a dynamic allocator for Eigen makes me suspicious), but in general:

• You might be able to supply your own buffers or objects
• You might be able do declare variables outside the body of the loop and mark them private with OpenMP
• You might be able to supply your own allocator that does not use the global heap
• Where possible, create a slice, subrange or view of your data, instead of copying it. For example, column access can be implemented with a std::ranges::stride_view.

In terms of coding style, I’d probably understand this code better if more of the types were explicit. I think auto is a good idea when it’s obvious to the maintainer what the type is, and maybe for someone with more knowledge of Eigen, that would be the case. You can often get the same flexibility with changing your types through using directives or template parameters.

Recon(i%szZ,i/szZ) is a Code Smell

This doesn’t look to me like a pattern of array access that the optimizer can recognize and transform into something more efficient. It’s possible that the compiler might inline aggressively and be smart enough to see the pattern, if the library you’re calling into is written the right way, but I wouldn’t bet on it. You probably want a nested loop like:

#pragma omp for collapse(2) // Along with other directives.
for (size_t i = 0; i < npoints/szZ; ++i)
    for (size_t j = 0; j < szZ; ++j) {
        const size_t elem_id = i*szZ + j; // Replace the original index with this.
        /* ... */
        Recon(i,j) = result.sum();
}

OpenMP definitely knows how to handle that.

Investigate Whether You can Use SIMD

I have gotten #pragma omp parallel for simd schedule(static) to work on code that is kinda-sorta like this if you squint real hard. I don’t know if Eigen, or a subset of Eigen, is written so that that would work.

Alternatives to a for Loop

In C++20, you can refactor most parallel for loops as a combination of <algorithm> and <ranges>, and give the algorithm a parallel execution policy. For example, many operations on a row of one matrix and a column of another matrix can be written as a std::transform_reduce over a std::zip_view (perhaps over a subrange representing the row and a stride_view representing the column). C++23 will additionally standardize std::mdspan, to view a linear range as a multidimensional array.

As of 2023, compilers don’t fully support these yet, however.

Use a Good Compiler and Libraries

For X64, I’ve gotten much better results with Intel ICX/ICPX than any other compiler. Try the flags -std=c++20 -O3 -march=native -fiopenmp -ipo -fp-model=fast. And of course, turn on warnings.