Accelerating computation of a "2D Variance" using C++ and Eigen in a MEX function

Question

I have the following function in Matlab:

function [JCF] = computeJCF(idxtMTX,p,alpha)
    % Compute Weight matrix based on 2D variance of idxtMTX per pixel.
    % idxtMTX - [N,M,nPixels] (stored in p as [p.numEl,p.na,p.nPoints])

    % Preallocate output matrix
    JCF = zeros(size(idxtMTX));
    denomFactor = double(p.numEl*p.na);

    for i = 1:p.nPoints
        % Set aside matrices to operate on
        kMTX = idxtMTX(:,:,i);
        kMTX2 = abs(kMTX).^alpha;

        % Compute weight matrix
        weightMTX = abs(sum(kMTX,1).*sum(kMTX,2)).^alpha ./ ...
          (denomFactor.^(alpha-1)*sum(kMTX2,1).*sum(kMTX2,2));

        % Throw away values that were invalid (i.e. Denominator was equal
        % to zero) and save weightMTX
        weightMTX(isnan(weightMTX)) = 0;
        JCF(:,:,i) = weightMTX;
    end

    % Apply Weight Matrix and sum to get result
    JCF = idxtMTX.*(JCF);
    JCF = reshape(squeeze(sum(JCF,[1,2])/denomFactor),[p.szZ,p.szX]);
end

Using the Matlab Data API for C++ and the Eigen toolbox (3.4.0), I have implemented the code above as the following mex function:

#include <iostream>
#include "mex.hpp"
#include "mexAdapter.hpp"
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <MatlabDataArray.hpp>
#include <cmath>
#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;
    
    void operator()(matlab::mex::ArgumentList outputs, matlab::mex::ArgumentList inputs) {
        checkArguments(outputs, inputs);
        
        // Initialize input parameters
        long int N = inputs[0][0];      // Typically number of elements
        long int M = inputs[0][1];      // Typically number of Plane Waves or Transmits
        long int nPixels = inputs[0][2];
        long int szZ = inputs[0][3];
        long int szX = inputs[0][4];
        long int alpha = inputs[0][5];
        
        // Initialize input idxtMTX data
        auto ptr = getDataPtr<std::complex<double>>(inputs[1]);
        Eigen::Map< const Eigen::MatrixXcd > idxt( ptr, N*M, nPixels );

        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);

        // Get num threads
        int numThreads = omp_get_max_threads();
        int nProc = omp_get_num_procs();
        omp_set_num_threads(nProc);
        
        double denomFactor = static_cast<double>(N*M);
        double denomFactor2 = static_cast<double>(std::pow(N*M,alpha-1));
        #pragma omp parallel for
        for (int i = 0; i < nPixels; i++) {
            // Define Matrix for given pixel
            const auto idxtCol = idxt.col(i).reshaped(N,M);
            const auto idxtPow = idxtCol.cwiseAbs().array().pow(alpha).matrix();
            
            // Compute Numerator and Denominator of weight matrix
            const auto s1 = (idxtCol.rowwise().sum()*idxtCol.colwise().sum()).cwiseAbs().array().pow(alpha);
            const auto s2 = idxtPow.rowwise().sum()*idxtPow.colwise().sum();
            
            // Compute weight matrix and result while excluding points where denom is zero.
            const auto weight = (s2.array() != 0).select(s1.array()/(s2.array()*denomFactor2),0.0);
            const auto result = idxtCol.array()*weight;

            // Store result in output matrix
            Recon(i%szZ,i/szZ) = result.sum()/denomFactor;
        }
    }

    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() != 2) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Two inputs required") }));
        }

        if (inputs[0].getNumberOfElements() != 6) {
            matlabPtr->feval(u"error", 
                0, std::vector<matlab::data::Array>({ factory.createScalar("Need 6 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") }));
        }

    }
    
    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->();
    }

};

This code compiles and works (compilation settings and test cases included below). However, say I replace the parallel for loop with the following for the case of alpha = 2:

double denomFactor = static_cast<double>(numEl*na);
#pragma omp parallel for
for (int i = 0; i < nPoints; i++) {
    const auto idxtCol = idxt.col(i).reshaped(numEl,na);
    
    const auto s1 = (idxtCol.rowwise().sum()*idxtCol.colwise().sum()).cwiseAbs2();
    const auto s2 = idxtCol.cwiseAbs2().rowwise().sum()*idxtCol.cwiseAbs2().colwise().sum();
    
    const auto weight = (s2.array() != 0).select(s1.array()/(s2.array()*denomFactor),0.0);
    const auto result = idxtCol.array()*weight;
    
    Recon(i%szZ,i/szZ) = result.sum()/denomFactor;
}

Compiling and running with this version offers an almost ~20x speedup. The problem appears to be the fact that the pow function does not have a vectorized implementation in Eigen. I've tried versions of this loop for alpha = 4 and alpha = 1 and they are all at least an order of magnitude faster than the general case. However, ideally, I'd like to be able to handle variable inputs of alpha without having to sacrifice so much in performance. I now have to call this function for a range of alphas going beyond 4, but my grasp of C++ is not strong enough to cleanly implement something that handles variable inputs of alpha at high performance.

I can further break this problem down for a few situations of increasing generality:

1. A function that can handle positive integer values of alpha upto some limit (e.g. 20)
2. A function that can handle any positive integer values of alpha.
3. A function that can handle any positive values of alpha.

(note: alpha is always positive and nonzero)

Here's how I compile the function(s) from Matlab (with mingw GCC version 8.3):

currentDir = matlab.desktop.editor.getActiveFilename; 
currentDir = regexp(currentDir, filesep, 'split');
nL = length(currentDir);
ipath = % path to eigen library here
mingwFlags = {'CXXFLAGS="$CXXFLAGS -march=native -std=c++14 -fno-math-errno -ffast-math -fopenmp -DNDEBUG"',...
            'LDFLAGS="$LDFLAGS -fopenmp"','CXXOPTIMFLAGS="-O3"'};

outFileDir = fullfile(currentDir{1:nL-1},'MEX Files');

% Constant Alpha = 2 version
tic; mex(ipath,mingwFlags{:},'JCFproductA2.cpp','-outdir',outFileDir); toc

% Variable Alpha (integer only) version
tic; mex(ipath,mingwFlags{:},'JCFproductExp.cpp','-outdir',outFileDir); toc

% Constant Alpha = 4 version
tic; mex(ipath,mingwFlags{:},'JCFproductA4.cpp','-outdir',outFileDir); toc

% Constant Alpha = 1 version
tic; mex(ipath,mingwFlags{:},'JCFproductA1.cpp','-outdir',outFileDir); toc

Here are some example inputs:

param = int32([192,75,836,38,22]); % [numEl,na,nPixels,szZ,szX]
tic; chk = JCFproductA2(param,idxtMTX); toc
tic; chk = JCFproductA4(param,idxtMTX); toc
tic; chk = JCFproductA1(param,idxtMTX); toc

% alpha = 2 using the general function
param2 = int32([param,2]);
tic; chk5 = JCFproductExp(param2,idxtMTX); toc

In real scenarios, nPixels, szZ, and szX are much larger. Also, nPixels = szZ*szX. However, using these dimensions and rerunning the code 100 times and taking the average timing of all of these instances, I got the following profiled results:

"Exp2"  0.2253
"A2"    0.01481
"Matlab2"   0.3180
"Exp1"  0.2291
"A1"    0.03093
"Matlab1"   0.3204
"Exp4"  0.2271
"A4"    0.01572
"Matlab4"   0.6400

where A1, A2 and A4 are the specialized mex functions, Matlab1, Matlab2 and Matlab4 are the timings of the Matlab function with the corresponding value for alpha, and Exp1, Exp2 and Exp4 are the generalized mex function with the corresponding value of alpha. All values are in seconds.

Answer 1

Maybe try exponentiation by squaring?

#include <iostream>

// computes base.pow(exp) for exp > 0
double pow_int(double base, int exp)
{
    // result already has 1 power of base
    --exp;
    double result = base;

    // answer is (result * term^exp)
    double term = base;

    while (exp > 0)
    {
        if (exp % 2 == 1)
            result = term * result;
        term = term * term;
        exp /= 2;
    }

    return result;
}

int main()
{
    std::cout << "pow_int(2.1, 1) = " << pow_int(2.1, 1) << "\n";
    std::cout << "pow_int(2.1, 2) = " << pow_int(2.1, 2) << "\n";
    std::cout << "pow_int(2.1, 3) = " << pow_int(2.1, 3) << "\n";
}

This could help for your second case, as exponentiation is much easier for integer alpha. You would replace the type of base with the type of idxtCol.rowwise().sum() * idxtCol.colwise().sum() and term * term with cwiseAbs2(). Hope this helps!