mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Simplify pthreadpool implementation on top of Caffe2 thread pool (#7666)
Remove one layer of pointer dereference when calling the thread pool.
This commit is contained in:
@ -11,7 +11,6 @@
|
||||
#include "caffe2/operators/conv_pool_op_base.h"
|
||||
|
||||
#include "caffe2/utils/math.h"
|
||||
#include "caffe2/utils/threadpool/pthreadpool_impl.h"
|
||||
#include "nnpack.h"
|
||||
|
||||
CAFFE2_DEFINE_bool(caffe2_profile_nnpack, false, "");
|
||||
@ -196,7 +195,7 @@ bool NNPACKConvOp::RunOnDeviceWithOrderNCHW() {
|
||||
const nnp_size output_subsample = {.width = static_cast<size_t>(stride_w()),
|
||||
.height = static_cast<size_t>(stride_h())};
|
||||
initNNPACK();
|
||||
pthreadpool pool(ws_->GetThreadPool());
|
||||
pthreadpool_t pool = reinterpret_cast<pthreadpool_t>(ws_->GetThreadPool());
|
||||
|
||||
runWithSharedBuffer<CPUContext>(ws_, [&](Tensor<CPUContext>* buffer) {
|
||||
if (transformStrategy_ == nnp_convolution_transform_strategy_precompute) {
|
||||
@ -220,7 +219,7 @@ bool NNPACKConvOp::RunOnDeviceWithOrderNCHW() {
|
||||
&transformedFilterSize,
|
||||
nnp_activation_identity,
|
||||
nullptr /* activation parameter */,
|
||||
&pool,
|
||||
pool,
|
||||
nullptr /* profile */);
|
||||
if (status == nnp_status_success) {
|
||||
/* For these convolution parameters filter transforms can be
|
||||
@ -257,7 +256,7 @@ bool NNPACKConvOp::RunOnDeviceWithOrderNCHW() {
|
||||
&transformedFilterSize,
|
||||
nnp_activation_identity,
|
||||
nullptr /* activation parameter */,
|
||||
&pool,
|
||||
pool,
|
||||
nullptr /* profile */);
|
||||
CAFFE_ENFORCE(
|
||||
nnp_status_success == status,
|
||||
@ -320,7 +319,7 @@ bool NNPACKConvOp::RunOnDeviceWithOrderNCHW() {
|
||||
&workspaceSize,
|
||||
activation_,
|
||||
nullptr /* activation parameter */,
|
||||
&pool,
|
||||
pool,
|
||||
FLAGS_caffe2_profile_nnpack ? &profile : nullptr);
|
||||
if (status == nnp_status_insufficient_buffer) {
|
||||
/* Query required workspace size, increase buffer, and try again */
|
||||
@ -341,7 +340,7 @@ bool NNPACKConvOp::RunOnDeviceWithOrderNCHW() {
|
||||
&workspaceSize,
|
||||
activation_,
|
||||
nullptr /* activation parameter */,
|
||||
&pool,
|
||||
pool,
|
||||
nullptr /* profile */);
|
||||
if (status == nnp_status_success) {
|
||||
/* Division with rounding up, in case size is not multiple of
|
||||
@ -374,7 +373,7 @@ bool NNPACKConvOp::RunOnDeviceWithOrderNCHW() {
|
||||
&workspaceSize,
|
||||
activation_,
|
||||
nullptr /* activation parameter */,
|
||||
&pool,
|
||||
pool,
|
||||
FLAGS_caffe2_profile_nnpack ? &profile : nullptr);
|
||||
}
|
||||
}
|
||||
|
@ -1,5 +1,4 @@
|
||||
#include "caffe2/utils/threadpool/pthreadpool.h"
|
||||
#include "caffe2/utils/threadpool/pthreadpool_impl.h"
|
||||
#include "caffe2/utils/threadpool/ThreadPool.h"
|
||||
|
||||
|
||||
@ -7,17 +6,19 @@
|
||||
// External API
|
||||
//
|
||||
|
||||
void pthreadpool_compute_1d(struct pthreadpool* threadpool,
|
||||
pthreadpool_function_1d_t function,
|
||||
void* argument,
|
||||
size_t range) {
|
||||
threadpool->pool_->run(
|
||||
[function, argument](int threadId, size_t workId) {
|
||||
function(argument, workId);
|
||||
},
|
||||
range);
|
||||
void pthreadpool_compute_1d(
|
||||
pthreadpool_t threadpool,
|
||||
pthreadpool_function_1d_t function,
|
||||
void* argument,
|
||||
size_t range) {
|
||||
reinterpret_cast<caffe2::ThreadPool*>(threadpool)
|
||||
->run(
|
||||
[function, argument](int threadId, size_t workId) {
|
||||
function(argument, workId);
|
||||
},
|
||||
range);
|
||||
}
|
||||
|
||||
size_t pthreadpool_get_threads_count(struct pthreadpool* threadpool) {
|
||||
return threadpool->pool_->getNumThreads();
|
||||
size_t pthreadpool_get_threads_count(pthreadpool_t threadpool) {
|
||||
return reinterpret_cast<caffe2::ThreadPool*>(threadpool)->getNumThreads();
|
||||
}
|
||||
|
@ -1,23 +0,0 @@
|
||||
#ifndef CAFFE2_UTILS_PTHREADPOOL_IMPL_H_
|
||||
#define CAFFE2_UTILS_PTHREADPOOL_IMPL_H_
|
||||
|
||||
#include "ThreadPoolCommon.h"
|
||||
|
||||
|
||||
namespace caffe2 {
|
||||
|
||||
class ThreadPool;
|
||||
|
||||
} // namespace caffe2
|
||||
|
||||
extern "C" {
|
||||
|
||||
// Wrapper for the caffe2 threadpool for the usage of NNPACK
|
||||
struct pthreadpool {
|
||||
pthreadpool(caffe2::ThreadPool* pool) : pool_(pool) {}
|
||||
caffe2::ThreadPool* pool_;
|
||||
};
|
||||
|
||||
} // extern "C"
|
||||
|
||||
#endif // CAFFE2_UTILS_PTHREADPOOL_IMPL_H_
|
Reference in New Issue
Block a user