devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merged in ezhulenev/eigen-01 (pull request PR-514)

Browse Source 
Add tests for evalShardedByInnerDim contraction + fix bugs
This commit is contained in:
Rasmus Munk Larsen
2018-09-28 18:37:54 +00:00
parent 86ba50be39 524c81f3fa
commit 7c1b47840a
4 changed files with 135 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h
View File

@@ -798,14 +798,15 @@ struct TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgT
auto process_block = [=, &barrier](Scalar* buf, Index first, Index last) {
::memset(buf, 0, m * n * sizeof(Scalar));
TENSOR_CONTRACTION_DISPATCH(
this->template evalGemmPartial, Alignment,
this->template evalGemmPartialWithoutOutputKernel, Alignment,
(buf, first, last, this->m_device.numThreads()));
barrier.Notify();
};
Index start = 0;
for (int blocks_left = num_blocks; blocks_left > 0; --blocks_left) {
// The underlying GEMM kernel assumes that k is a multiple of 8 and
// subtle breakage occurs if this is violated.
// The underlying GEMM kernel assumes that k is a multiple of packet size
// (currently largest packet size is 8) and subtle breakage occurs if
// this is violated.
block_size = 8 * divup<Index>(k - start, 8 * blocks_left);
Scalar* buf;
if (start == 0) {
@@ -830,6 +831,14 @@ struct TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgT
addToBuffer<Alignment>(m * n, buf, result);
this->m_device.deallocate(buf);
}
// Finally call output kernel with finalized output buffer.
typedef internal::blas_data_mapper<Scalar, Index, ColMajor> OutputMapper;
this->m_output_kernel(OutputMapper(result, m),
this->m_tensor_contraction_params,
static_cast<Eigen::Index>(0),
static_cast<Eigen::Index>(0),
m, n);
}
TensorOpCost contractionCostPerInnerDim(Index m, Index n, Index k) const {