Dot product batch-wise

Stone · June 10, 2018, 2:11am

Yeah you’re right, it uses a single batched kernel. Actually the CPU version is a loop in batch dimension.

pytorch/pytorch/blob/29849e428cba03f4a9e24a157781ce822512c39e/aten/src/TH/generic/THTensorMath.cpp#L2308


  THTensor_(resizeAs)(result, t);
  if (beta != 0.0) {
    THTensor_(copy)(result, t);
  }
}


THTensor *matrix1 = THTensor_(new)();
THTensor *matrix2 = THTensor_(new)();
THTensor *result_matrix = THTensor_(new)();


for (batch = 0; batch < THTensor_(size)(batch1, 0); ++batch) {
  THTensor_(select)(matrix1, batch1, 0, batch);
  THTensor_(select)(matrix2, batch2, 0, batch);
  THTensor_(select)(result_matrix, result, 0, batch);


  THTensor_(addmm)(result_matrix, beta, result_matrix, alpha, matrix1, matrix2);
}


THTensor_(free)(matrix1);
THTensor_(free)(matrix2);
THTensor_(free)(result_matrix);