diff options
| author |  Yuefeng Zhou <yuefengz@google.com> | 2018-02-20 16:50:05 -0800 |
|---|---|---|
| committer | Yuefeng Zhou <yuefengz@google.com> | 2018-02-20 16:50:05 -0800 |
| commit | 1eff6cf8a77f1b8699671d31f8f307a6fd9170ea (patch) | |
| tree | f327431b9b918148d548b849cd24a3633bd50274 /unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h | |
| parent | adb134d47e70906ed15d877d7a0fe03f3f665c02 (diff) | |
Use device's allocate function instead of internal::aligned_malloc. This would make it easier to track memory usage in device instances.
Diffstat (limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h')
| -rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h b/unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h index d30cc96ab..6fb69910e 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorContractionThreadPool.h @@ -377,7 +377,7 @@ struct TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgT divup<size_t>(bm_ * bk_ * sizeof(LhsScalar), align) * align; size_t rhs_size = divup<size_t>(bn_ * bk_ * sizeof(RhsScalar), align) * align; - packed_mem_ = static_cast<char*>(internal::aligned_malloc( + packed_mem_ = static_cast<char*>(device_.allocate( (nm0_ * lhs_size + nn0_ * rhs_size) * std::min<size_t>(nk_, P - 1))); char* mem = static_cast<char*>(packed_mem_); for (Index x = 0; x < numext::mini<Index>(nk_, P - 1); x++) { @@ -399,7 +399,7 @@ struct TensorEvaluator<const TensorContractionOp<Indices, LeftArgType, RightArgT for (Index m = 0; m < nm_; m++) delete[] state_kernel_[x][m]; delete[] state_kernel_[x]; } - internal::aligned_free(packed_mem_); + device_.deallocate(packed_mem_); } void run() { |