Default pool() option (#14636)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14636

Add a default CPU option for the pool()

Reviewed By: andrewwdye

Differential Revision: D13281367

fbshipit-source-id: 92dbfce89c900a41731b6d1ff62bb97886c40f77

diff --git a/caffe2/core/net_async_base.cc b/caffe2/core/net_async_base.cc
index e0ad1bc..6d9f741 100644 (file)
--- a/caffe2/core/net_async_base.cc
+++ b/caffe2/core/net_async_base.cc
@@ -161,6 +161,13 @@ TaskThreadPoolBase* AsyncNetBase::poolGetter(
   return pool.get();
 }
 
+TaskThreadPoolBase* AsyncNetBase::pool() {
+  // By default using a non-pinned CPU option
+  DeviceOption dev;
+  dev.set_device_type(PROTO_CPU);
+  return pool(dev);
+}
+
 TaskThreadPoolBase* AsyncNetBase::pool(const DeviceOption& device_option) {
   if (options_.use_single_pool_) {
     return poolGetter(cpu_pools_, PROTO_CPU, -1, num_workers_);

diff --git a/caffe2/core/net_async_base.h b/caffe2/core/net_async_base.h
index 5577f0c..8792c04 100644 (file)
--- a/caffe2/core/net_async_base.h
+++ b/caffe2/core/net_async_base.h
@@ -108,6 +108,7 @@ class CAFFE2_API AsyncNetBase : public NetBase {
   bool run(int task_id, int stream_id);
   int stream(int task_id);
   TaskThreadPoolBase* pool(const DeviceOption& device_option);
+  TaskThreadPoolBase* pool();
 
   void finishTasks(const std::unordered_set<int>& task_ids);
   void finalizeEvents();