net_simple_refcount type to help experimentation with dynamic allocation. (pytorch#13370)

Summary:
Pull Request resolved: pytorch#13370

This diff adds a new net type (simple_refcount) that does one thing: for all
intermediate results produced by a net, it will keep refcount about internal
usage, and when it finishes its consumption, the net will delete the blob
content to mimic the case of dynamic allocation. In fact, this would also be
the behavior when we go functional: anything that is not explicitly marked as
input or output will be up to the executor for lifetime management.

See the comments in net_simple_refcount.cc for details.

Reviewed By: dzhulgakov

Differential Revision: D12855489

fbshipit-source-id: 594a47a786305d595fd505b6700864dd1d9c72aa

Author: Yangqing Jia

caffe2/core/net_simple_refcount.cc

@@ -0,0 +1,81 @@
+#include "caffe2/core/net_simple_refcount.h"
+#include "caffe2/core/net.h"
+
+#include <iostream>
+#include <set>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "caffe2/core/operator.h"
+#include "caffe2/core/static_tracepoint.h"
+#include "caffe2/core/timer.h"
+#include "caffe2/proto/caffe2_pb.h"
+#include "caffe2/utils/proto_utils.h"
+
+namespace caffe2 {
+
+SimpleRefCountNet::SimpleRefCountNet(
+    const std::shared_ptr<const NetDef>& net_def,
+    Workspace* ws)
+    : SimpleNet(net_def, ws) {
+  VLOG(1) << "Constructing SimpleRefCountNet " << net_def->name();
+  // Construct the "to delete" list.
+  delete_list_.resize(net_def->op_size());
+
+  std::map<string, int> last_consumed_at;
+  std::set<string> created_by_me;
+  // For each opeartor
+  for (int idx = 0; idx < net_def->op_size(); ++idx) {
+    const auto& op_def = net_def->op(idx);
+    for (const string& in_name : op_def.input()) {
+      last_consumed_at[in_name] = idx;
+    }
+    for (const string& out_name : op_def.output()) {
+      created_by_me.insert(out_name);
+    }
+  }
+  // We do not delete any operator that is not produced by the net, and
+  // any operator that is marked as external_output. Any blob that is not
+  // consumed won't be in the last_consumed_at map, so we don't need to
+  // do anything special.
+  for (auto& kv : last_consumed_at) {
+    if (!created_by_me.count(kv.first)) {
+      kv.second = -1;
+    }
+  }
+  for (const string& name : net_def->external_output()) {
+    last_consumed_at[name] = -1;
+  }
+  // Set up the delete list.
+  for (auto& kv : last_consumed_at) {
+    if (kv.second > 0) {
+      delete_list_[kv.second].push_back(ws->GetBlob(kv.first));
+      VLOG(1) << "NetSimpleRefCountNet: will delete " << kv.first
+              << " at operator #" << kv.second;
+    }
+  }
+}
+
+bool SimpleRefCountNet::Run() {
+  StartAllObservers();
+  VLOG(1) << "Running net " << name_;
+  for (int op_id = 0; op_id < operators_.size(); ++op_id) {
+    auto& op = operators_[op_id];
+    VLOG(1) << "Running operator " << op->debug_def().name() << "("
+            << op->debug_def().type() << ").";
+    bool res = op->Run();
+    if (!res) {
+      LOG(ERROR) << "Operator failed: " << ProtoDebugString(op->debug_def());
+      return false;
+    }
+    for (Blob* blob : delete_list_[op_id]) {
+      blob->Reset();
+    }
+  }
+  StopAllObservers();
+  return true;
+}
+
+REGISTER_NET(simple_refcount, SimpleRefCountNet);
+
+} // namespace caffe2

caffe2/core/net_simple_refcount.h

@@ -0,0 +1,59 @@
+#ifndef CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_
+#define CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_
+
+#include <vector>
+
+#include "c10/util/Registry.h"
+#include "caffe2/core/common.h"
+#include "caffe2/core/logging.h"
+#include "caffe2/core/net.h"
+#include "caffe2/core/net_simple.h"
+#include "caffe2/core/tensor.h"
+#include "caffe2/core/workspace.h"
+#include "caffe2/proto/caffe2_pb.h"
+
+namespace caffe2 {
+
+// SimpleRefcountNet is an implementation that adds an additional abstraction
+// on top of SimpleRefCountNet: it tracks all the tensors and for those that are
+// considered internal/temporary, delete them once their refcount go to zero.
+// In the context of a simple static run, this can be carried out during
+// construction time: we will do a pass through the network and track what
+// blobs we need to do reset on, after the execution of every op.
+//
+// To identify which blob is considered temporary, we employ the following
+// strategy: any blob that is
+// (1) consumed but not produced by ops in the net, or
+// (2) produced but not consumed by ops in the net, or
+// (3) is marked as external_output in the protobuf
+// will NOT be considered temporary.
+//
+// In the long run, we should design proper functional interfaces so that
+// nets are less imperative and more functional.
+//
+// Also, for now, SimpleRefCountNet should only be used for benchmarking
+// purposes and not product use, since it is not going to provide better
+// performance gain, and is implicitly incompatible with the contract that
+// earlier Nets expose - that all intermediate blobs are visible to the users.
+class SimpleRefCountNet final : public SimpleNet {
+ public:
+  SimpleRefCountNet(
+      const std::shared_ptr<const NetDef>& net_def,
+      Workspace* ws);
+
+ protected:
+  bool Run() override;
+
+  using SimpleNet::operators_;
+
+ private:
+  // The list of blobs to delete when each operator finishes its run.
+  // This will be populated during construction time.
+  vector<vector<Blob*>> delete_list_;
+
+  C10_DISABLE_COPY_AND_ASSIGN(SimpleRefCountNet);
+};
+
+} // namespace caffe2
+
+#endif // CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_

caffe2/core/net_simple_refcount_test.cc

@@ -0,0 +1,71 @@
+#include <gtest/gtest.h>
+#include "c10/util/StringUtil.h"
+#include "caffe2/core/net.h"
+#include "caffe2/core/net_async_scheduling.h"
+#include "caffe2/core/net_dag.h"
+#include "caffe2/core/operator.h"
+#include "caffe2/core/scope_guard.h"
+
+#include <google/protobuf/text_format.h>
+
+namespace caffe2 {
+
+namespace {
+
+// A net test dummy op that does nothing but scaffolding. Here, we
+// inherit from OperatorBase because we instantiate on both CPU and
+// GPU. In general, you want to only inherit from Operator<Context>.
+class NetSimpleRefCountTestOp final : public Operator<CPUContext> {
+ public:
+  NetSimpleRefCountTestOp(const OperatorDef& operator_def, Workspace* ws)
+      : Operator<CPUContext>(operator_def, ws) {}
+  USE_OPERATOR_FUNCTIONS(CPUContext);
+
+  bool RunOnDevice() override {
+    const int32_t& input = OperatorBase::Input<int32_t>(0);
+    int32_t* output = OperatorBase::Output<int32_t>(0);
+    *output = input + 1;
+    return true;
+  }
+};
+
+REGISTER_CPU_OPERATOR(NetSimpleRefCountTest, NetSimpleRefCountTestOp);
+
+OPERATOR_SCHEMA(NetSimpleRefCountTest).NumInputs(1).NumOutputs(1);
+
+TEST(NetSimpleRefCountTest, TestCorrectness) {
+  Workspace ws;
+  *(ws.CreateBlob("a")->GetMutable<int32_t>()) = 1;
+  NetDef net_def;
+  net_def.set_type("simple_refcount");
+  net_def.add_op()->CopyFrom(
+      CreateOperatorDef("NetSimpleRefCountTest", "", {"a"}, {"b"}));
+  net_def.add_op()->CopyFrom(
+      CreateOperatorDef("NetSimpleRefCountTest", "", {"b"}, {"c"}));
+  net_def.add_op()->CopyFrom(
+      CreateOperatorDef("NetSimpleRefCountTest", "", {"b"}, {"d"}));
+  net_def.add_op()->CopyFrom(
+      CreateOperatorDef("NetSimpleRefCountTest", "", {"c"}, {"e"}));
+  // After execution, what should look like is:
+  // a = 1
+  // b = deallocated
+  // c = deallocated
+  // d = 3
+  // e = 4
+  std::unique_ptr<NetBase> net(CreateNet(net_def, &ws));
+  net->Run();
+  // Note on ASSERT vs EXPECT: ASSERT will quit directly if condition not
+  // met, which is why we guard IsType<> calls with ASSERT so that the
+  // subsequent Get() calls do not product an exception.
+  ASSERT_TRUE(ws.GetBlob("a")->IsType<int32_t>());
+  EXPECT_EQ(ws.GetBlob("a")->Get<int32_t>(), 1);
+  EXPECT_EQ(ws.GetBlob("b")->GetRaw(), nullptr);
+  EXPECT_EQ(ws.GetBlob("c")->GetRaw(), nullptr);
+  ASSERT_TRUE(ws.GetBlob("d")->IsType<int32_t>());
+  EXPECT_EQ(ws.GetBlob("d")->Get<int32_t>(), 3);
+  ASSERT_TRUE(ws.GetBlob("e")->IsType<int32_t>());
+  EXPECT_EQ(ws.GetBlob("e")->Get<int32_t>(), 4);
+}
+
+} // namespace
+} // namespace caffe2