TBB task graph (#15041)

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

Adding an alternative implementation of a task graph based on TBB

Reviewed By: dmudiger

Differential Revision: D13412517

fbshipit-source-id: f5efedd680bbe0072bf38d504e5682ab51dd630f

diff --git a/caffe2/core/net.cc b/caffe2/core/net.cc
index 4093472..54ca39d 100644 (file)
--- a/caffe2/core/net.cc
+++ b/caffe2/core/net.cc
@@ -182,6 +182,14 @@ TaskThreadPoolBase* ExecutorHelper::GetPool(
   CAFFE_THROW("Not implemented");
 }
 
+std::vector<OperatorBase*> ExecutorHelper::GetOperators() const {
+  CAFFE_THROW("Not implemented");
+}
+
+int ExecutorHelper::GetNumWorkers() const {
+  CAFFE_THROW("Not implemented");
+}
+
 std::vector<float> NetBase::TEST_Benchmark(
     const int warmup_runs,
     const int main_runs,

diff --git a/caffe2/core/net.h b/caffe2/core/net.h
index 7b89f8b..e1c62dd 100644 (file)
--- a/caffe2/core/net.h
+++ b/caffe2/core/net.h
@@ -131,6 +131,8 @@ class CAFFE2_API ExecutorHelper {
  public:
   ExecutorHelper() {}
   virtual TaskThreadPoolBase* GetPool(const DeviceOption& option) const;
+  virtual std::vector<OperatorBase*> GetOperators() const;
+  virtual int GetNumWorkers() const;
   virtual ~ExecutorHelper() {}
 };
 

diff --git a/caffe2/core/net_parallel.h b/caffe2/core/net_parallel.h
index df5b6d8..756637f 100644 (file)
--- a/caffe2/core/net_parallel.h
+++ b/caffe2/core/net_parallel.h
@@ -68,6 +68,14 @@ class ParallelNetExecutorHelper : public ExecutorHelper {
     return net_->Pool(option);
   }
 
+  std::vector<OperatorBase*> GetOperators() const override {
+    return net_->GetOperators();
+  }
+
+  int GetNumWorkers() const override {
+    return net_->num_workers_;
+  }
+
  private:
   ParallelNet* net_;
 };