Search code examples
c++matrixeigen

User defined Eigen function using twice as much memory as expected


I have defined the following function (MWE)

(Note the formulation is an adaption of this formulation: How to Build a Distance Matrix without a Loop (Vectorization)?, as well as http://nonconditional.com/2014/04/on-the-trick-for-computing-the-squared-euclidian-distances-between-two-sets-of-vectors/)

#include <stdlib.h>
#include <chrono>
#include <Eigen/Dense>
#include <iostream>

using MyMatrix = Eigen::MatrixXd;
using MyMatrix1D = Eigen::VectorXd;

//Calculates e^(scale * ||x-y||_2^2), where ||x-y|| is euclidean distatnce
MyMatrix get_kernel_matrix(const Eigen::Ref<const MyMatrix> x, const Eigen::Ref<const MyMatrix> y)
  {
    const double scale = 0.017;
    const MyMatrix1D XX = x.array().square().rowwise().sum().matrix();
    const MyMatrix1D YY = y.array().square().rowwise().sum().matrix();
     return (((((-2*x)*y.transpose()).colwise() + XX).rowwise() + YY.transpose()).array() * scale).exp().matrix();
  }

int main(int argc, char** argv) {
  const int num_x = 2500;
  const int num_y = 2500;

  const MyMatrix X = MyMatrix::Random(num_x, 2);
  const MyMatrix Y = MyMatrix::Random(num_y, 2);

  const auto t_b_gen = std::chrono::high_resolution_clock::now();
  const MyMatrix k_xp_x(std::move(get_kernel_matrix(X, Y)));
  const auto t_a_gen = std::chrono::high_resolution_clock::now();
  long t_gen = std::chrono::duration_cast<std::chrono::nanoseconds>(t_a_gen - t_b_gen).count();
  std::cout << "Time: " << t_gen << std::endl;
}

which one would expect would take 2500*2500*8bytes = 50MB of memory. However, running /usr/bin/time -v kern_double reports: Maximum resident set size (kbytes): 103288.

Running the program through Massif indicates that the 50MB block is allocated twice, once in the functin call, and once Eigen::internal::cal_dense_assignment. I have attempted with and without std::move to try to force copy elision, however I have not been able to reduce the memory footprint.

What am I doing incorrectly and how can I fix this to use only the required memory rather than double?


Solution

  • This is because the matrix product x*y.transpose() is, by default, evaluated within a temporary to make it more efficient. You can still reuse this temporary by splitting yourself the last expression as follows:

    MyMatrix tmp = -2*x*y.transpose();
    tmp = ((((tmp).colwise() + XX).rowwise() + YY.transpose()).array() * scale).exp();
    return tmp;
    

    Note that neither .matrix(), nor std::move are required here.