Feat/graph logical op debug repr

oneflow/api/python/framework/dtype.cpp

@@ -38,7 +38,10 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
           [](int t) {  // __setstate__
             return CHECK_JUST(DType::Get(DataType(t)));
           }))
-      .def_property_readonly("bytes", [](const Symbol<DType>& dtype) { return dtype->bytes(); });
+      .def_property_readonly("bytes", [](const Symbol<DType>& dtype) { return dtype->bytes(); })
+      .def("get", [](const int data_type_enum) {
+        return CHECK_JUST(DType::Get(static_cast<DataType>(data_type_enum)));
+      });
 
   m.attr("bool") = &CHECK_JUST(DType::Get(DataType::kBool));
   m.attr("char") = &CHECK_JUST(DType::Get(DataType::kChar));

oneflow/api/python/framework/nn_graph.cpp

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
 #include <pybind11/stl.h>
 #include <memory>
 #include <string>
@@ -41,6 +42,11 @@ Maybe<py::object> APINNGraphAdditionalVarTensors(const std::shared_ptr<NNGraph>&
   py::list tensor_list = py::cast(tensors);
   return py::cast<py::object>(tensor_list);
 }
+
+Maybe<py::bytes> APINNGraphGetCompiledSerializedJob(const std::shared_ptr<NNGraph>& graph) {
+  const auto job = *JUST(graph->GetCompiledJob());
+  return py::bytes(job.SerializeAsString());
+}
 }  // namespace
 
 ONEFLOW_API_PYBIND11_MODULE("nn.graph.", m) {
@@ -75,7 +81,8 @@ ONEFLOW_API_PYBIND11_MODULE("nn.graph.", m) {
            &NNGraph::RegisterAdditionalVarOpNamesAndTensorsToBeLoaded)
       .def_property_readonly("additional_var_names", &APINNGraphAdditionalVarNames)
       .def_property_readonly("additional_var_tensors", &APINNGraphAdditionalVarTensors)
-      .def("complie_and_init_runtime", &NNGraph::CompileAndInitRuntime);
+      .def("complie_and_init_runtime", &NNGraph::CompileAndInitRuntime)
+      .def("get_compiled_job_str", &APINNGraphGetCompiledSerializedJob);
 
   m.def("RunLazyNNGraph", &RunLazyNNGraph);
   m.def("SoftSyncNNGraphBuffers", &SoftSyncNNGraphBuffers);

oneflow/api/python/symbol/placement_symbol.cpp

@@ -17,6 +17,7 @@ limitations under the License.
 #include <pybind11/stl.h>
 #include <pybind11/operators.h>
 
+#include "oneflow/core/common/maybe.h"
 #include "oneflow/extension/python/numpy.h"
 #include "oneflow/api/python/framework/size.h"
 #include "oneflow/api/python/of_api_registry.h"
@@ -63,6 +64,18 @@ struct PlacementSymbolExportUtil {
     return parallel_desc;
   }
 
+  static Maybe<ParallelDesc> CreateParallelDesc(const std::string& proto_str) {
+    ParallelConf parallel_conf;
+    CHECK_OR_RETURN(TxtString2PbMessage(proto_str, &parallel_conf));
+    std::shared_ptr<ParallelDesc> parallel_desc;
+    JUST(PhysicalRun([&parallel_desc, &parallel_conf](InstructionsBuilder* builder) -> Maybe<void> {
+      parallel_desc = JUST(builder->GetParallelDescSymbol(parallel_conf));
+      return Maybe<void>::Ok();
+    }));
+
+    return parallel_desc;
+  }
+
   static Maybe<std::vector<std::string>> ParseAndFormatRanks(const py::dict& device_ids) {
     std::vector<std::pair<int64_t, int64_t>> machine_device_id_vec;
     for (const auto& pair : device_ids) {
@@ -137,6 +150,10 @@ struct PlacementSymbolExportUtil {
     return SymbolOf(*JUST(CreateParallelDesc(type, *formated_machine_device_ids, shape)));
   }
 
+  static Maybe<Symbol<ParallelDesc>> CreateParallelDescSymbol(const std::string& proto_str) {
+    return SymbolOf(*JUST(CreateParallelDesc(proto_str)));
+  }
+
   static Maybe<Symbol<ParallelDesc>> AllDevicePlacement(const std::string& type) {
     static thread_local HashMap<std::string, Symbol<ParallelDesc>> device_tag2placement;
     CHECK_NOTNULL((Global<ResourceDesc, ForEnv>::Get()));
@@ -213,6 +230,10 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
              return PlacementSymbolExportUtil::CreateParallelDescSymbol(type, ranks).GetOrThrow();
            }),
            py::arg("type"), py::arg("ranks"))
+      .def(py::init([](const std::string& proto_str) {
+             return PlacementSymbolExportUtil::CreateParallelDescSymbol(proto_str).GetOrThrow();
+           }),
+           py::arg("proto_str"))
       .def_property_readonly(
           "device_type",
           [](Symbol<ParallelDesc> p) {

oneflow/core/framework/nn_graph.h

@@ -72,6 +72,7 @@ class NNGraph final : public NNGraphIf {
   Maybe<std::vector<std::string>> GetAdditionalVarOpNames() const;
   Maybe<std::vector<std::shared_ptr<one::Tensor>>> GetAdditionalVarOpTensors() const;
   Maybe<void> CompileAndInitRuntime();
+  Maybe<Job> GetCompiledJob() { return job_; }
   Maybe<void> Close();
 
  private:

oneflow/core/job/job_build_and_infer_ctx.cpp

@@ -196,7 +196,7 @@ void JobBuildAndInferCtx::AddOpAndUpdateJobParallelViewConf(const OperatorConf&
       (*module_name2module_conf)[module_name].set_name(scope.scope_proto().module_name());
     }
 
-    (*module_name2module_conf)[module_name].add_ops()->CopyFrom(operator_conf);
+    *((*module_name2module_conf)[module_name].add_ops()) = operator_conf.name();
   }
 }
 
@@ -999,7 +999,7 @@ Maybe<void> LazyJobBuildAndInferCtx::Complete() {
   int32_t pass_cnt = 0;
   const int64_t prev_v = FLAGS_v;
   auto DoPass = [&](const std::string& pass_name, int32_t cnt = 0) -> Maybe<void> {
-    VLOG(1) << job_name << " is compiling with pass"
+    VLOG(1) << job_name << " start compiling with pass"
             << " pass_cnt_" + std::to_string(pass_cnt) + "-" + pass_name
             << (cnt > 0 ? std::to_string(cnt) : "");
     if (unlikely(NeedLogJob(pass_name))) {
@@ -1013,6 +1013,9 @@ Maybe<void> LazyJobBuildAndInferCtx::Complete() {
       std::string cnt_str = cnt > 0 ? std::to_string(cnt) : "";
       LogJob("pass_cnt_" + std::to_string(pass_cnt) + "-" + pass_name + cnt_str + "-after");
     }
+    VLOG(1) << job_name << " finish compiling with pass"
+            << " pass_cnt_" + std::to_string(pass_cnt) + "-" + pass_name
+            << (cnt > 0 ? std::to_string(cnt) : "");
     ++pass_cnt;
     return Maybe<void>::Ok();
   };

oneflow/core/job/job_builder.cpp

@@ -19,7 +19,10 @@ limitations under the License.
 #include "oneflow/core/common/container_util.h"
 #include "oneflow/core/job/job.pb.h"
 #include "oneflow/core/job/sbp_parallel.pb.h"
+#include "oneflow/core/operator/op_conf.pb.h"
 #include "oneflow/core/operator/operator.h"
+#include "oneflow/core/vm/symbol_storage.h"
+#include "oneflow/core/framework/scope_util.h"
 
 namespace oneflow {
 
@@ -170,6 +173,7 @@ Maybe<void> JobBuilder::AddOp(const ParallelConf& parallel_conf, const OperatorC
   OperatorConf* mut_op_conf = job_->mutable_net()->add_op();
   *mut_op_conf = op_conf;
   CHECK_OR_RETURN(op_name2op_conf_.emplace(op_conf.name(), mut_op_conf).second);
+  AddOpToModuleConf(op_conf);
   AddOpNamesToPlacementGroup({op_conf.name()}, parallel_conf);
   return Maybe<void>::Ok();
 }
@@ -185,10 +189,35 @@ void JobBuilder::AddOps(const ParallelConf& parallel_conf,
     *mut_op_conf = op_conf;
     CHECK(op_name2op_conf_.emplace(op_conf.name(), mut_op_conf).second);
     op_names.emplace_back(op_conf.name());
+    AddOpToModuleConf(op_conf);
   }
   AddOpNamesToPlacementGroup(op_names, parallel_conf);
 }
 
+void JobBuilder::AddOpToModuleConf(const OperatorConf& op_conf) {
+  // set up the module config
+  if (Global<symbol::Storage<Scope>>::Get()->Has(op_conf.scope_symbol_id())) {
+    const auto& scope = Global<symbol::Storage<Scope>>::Get()->Get(op_conf.scope_symbol_id());
+    if (scope.scope_proto().has_module_name()) {
+      const auto& module_name = scope.scope_proto().module_name();
+      auto* module_name2module_conf = job_->mutable_module_name2module_conf();
+      if (!(*module_name2module_conf)[module_name].has_name()) {
+        (*module_name2module_conf)[module_name].set_name(scope.scope_proto().module_name());
+      }
+
+      *((*module_name2module_conf)[module_name].add_ops()) = op_conf.name();
+      return;
+    }
+  }
+  const auto& module_name = job_->job_conf().job_name();
+  auto* module_name2module_conf = job_->mutable_module_name2module_conf();
+  if (!(*module_name2module_conf)[module_name].has_name()) {
+    (*module_name2module_conf)[module_name].set_name(module_name);
+  }
+
+  *((*module_name2module_conf)[module_name].add_ops()) = op_conf.name();
+}
+
 void JobBuilder::AddOpNamesToPlacementGroup(const std::vector<std::string>& op_names,
                                             const ParallelConf& parallel_conf) {
   PlacementGroup* placement_group = nullptr;
@@ -230,6 +259,21 @@ void JobBuilder::RemoveOpByName(const std::unordered_set<std::string>& removing_
   for (const OperatorConf& op_conf : net.op()) {
     if (removing_names.count(op_conf.name()) == 0) { *(job_->mutable_net()->add_op()) = op_conf; }
   }
+  // Update module conf
+  auto module_confs_map = job_->module_name2module_conf();
+  job_->clear_module_name2module_conf();
+  for (const auto& module_conf_pair : module_confs_map) {
+    const auto& module_name = module_conf_pair.first;
+    auto* module_name2module_conf = job_->mutable_module_name2module_conf();
+    if (!(*module_name2module_conf)[module_name].has_name()) {
+      (*module_name2module_conf)[module_name].set_name(module_name);
+    }
+    for (const auto& op_name : module_conf_pair.second.ops()) {
+      if (removing_names.count(op_name) == 0) {
+        *((*module_name2module_conf)[module_name].add_ops()) = op_name;
+      }
+    }
+  }
   // Update placement
   auto placement_group = job_->placement().placement_group();
   job_->mutable_placement()->clear_placement_group();

oneflow/core/job/job_builder.h

@@ -81,6 +81,7 @@ class JobBuilder final {
  private:
   void AddOpNamesToPlacementGroup(const std::vector<std::string>& op_names,
                                   const ParallelConf& parallel_conf);
+  void AddOpToModuleConf(const OperatorConf& op_conf);
 
   Job* job_;
   HashMap<std::string, OperatorConf*> op_name2op_conf_;

oneflow/core/job/module_conf.proto

@@ -1,9 +1,7 @@
 syntax = "proto2";
 package oneflow;
 
-import "oneflow/core/operator/op_conf.proto";
-
 message ModuleConf {
   required string name = 1;
-  repeated OperatorConf ops = 2;
+  repeated string ops = 2;
 }

oneflow/core/job/plan_util.cpp

@@ -861,8 +861,9 @@ namespace {
 struct MemBlockMemoryInfo {
   int64_t mem_block_id;
   int64_t mem_block_mem_size;
+  bool is_reused;
   std::vector<std::string> ordered_op_names;
-  MemBlockMemoryInfo() : mem_block_id(-1), mem_block_mem_size(-1) {}
+  MemBlockMemoryInfo() : mem_block_id(-1), mem_block_mem_size(-1), is_reused(false) {}
 };
 
 struct ChunkMemoryInfo {
@@ -924,7 +925,10 @@ void PlanUtil::PlanMemoryLog(Plan* plan, const std::string& plan_name) {
     if (mem_block.mem_case().has_device_cuda_mem()) {
       if (mem_block.has_chunk_id()) {
         rank_memory_info.chunk_info.mem_block_ids.push_back(mem_block_id);
+        info.is_reused = true;
       } else {
+        rank_memory_info.chunk_info.mem_block_ids.push_back(mem_block_id);
+        info.is_reused = false;
         rank_memory_info.not_reused_mem_size += mem_block.mem_size();
         rank_memory_info.total_mem_size += mem_block.mem_size();
         if (mem_block.has_variable_op_name()) {
@@ -968,25 +972,26 @@ void PlanUtil::PlanMemoryLog(Plan* plan, const std::string& plan_name) {
               << B2MiB(rank_memory_info.eager_variable_total_mem_size) << " MiB ].";
   }
 
-  if (IsInDebugMode()) {
-    for (const auto& rank_memory_info : rank_device_memory_infos) {
-      int64_t chunk_id = rank_memory_info.chunk_info.chunk_id;
-      VLOG(2) << " For detail: Chunk id: " << chunk_id << " has "
-              << rank_memory_info.chunk_info.mem_block_ids.size() << " MemBlocks.";
-      for (int64_t mem_block_id : rank_memory_info.chunk_info.mem_block_ids) {
-        CHECK(mem_block_id2info.find(mem_block_id) != mem_block_id2info.end());
-        const auto& mem_block_info = mem_block_id2info.at(mem_block_id);
-        VLOG(2) << " In Chunk id: " << chunk_id << " MemBlock id: " << mem_block_id
-                << " has num = " << mem_block_info.ordered_op_names.size()
-                << " ops with mem size = " << B2MiB(mem_block_info.mem_block_mem_size);
-      }
-      for (int64_t mem_block_id : rank_memory_info.chunk_info.mem_block_ids) {
-        CHECK(mem_block_id2info.find(mem_block_id) != mem_block_id2info.end());
-        const auto& mem_block_info = mem_block_id2info.at(mem_block_id);
-        for (int64_t i = 0; i < mem_block_info.ordered_op_names.size(); ++i) {
-          VLOG(3) << " In Chunk id: " << chunk_id << " MemBlock id: " << mem_block_id
-                  << " order: " << i << " op_name: " << mem_block_info.ordered_op_names.at(i);
-        }
+  for (const auto& rank_memory_info : rank_device_memory_infos) {
+    int64_t chunk_id = rank_memory_info.chunk_info.chunk_id;
+    int64_t device_id = rank_memory_info.device_id;
+    int64_t not_reuse_size = rank_memory_info.not_reused_mem_size;
+    VLOG(2) << " For detail: Chunk id: " << chunk_id << " has "
+            << rank_memory_info.chunk_info.mem_block_ids.size() << " MemBlocks"
+            << " not reused size = " << B2MiB(not_reuse_size);
+    for (int64_t mem_block_id : rank_memory_info.chunk_info.mem_block_ids) {
+      CHECK(mem_block_id2info.find(mem_block_id) != mem_block_id2info.end());
+      const auto& mem_block_info = mem_block_id2info.at(mem_block_id);
+      VLOG(2) << "     In Device: " << device_id << " Chunk id: " << chunk_id
+              << " MemBlock id: " << mem_block_id
+              << " has num = " << mem_block_info.ordered_op_names.size()
+              << " ops with mem size = " << B2MiB(mem_block_info.mem_block_mem_size)
+              << " is reused " << mem_block_info.is_reused;
+      for (int64_t i = 0; i < mem_block_info.ordered_op_names.size(); ++i) {
+        VLOG(3) << "         In Device: " << device_id << " Chunk id: " << chunk_id
+                << " In MemBlock id: " << mem_block_id << " order: " << i << " is reused "
+                << mem_block_info.is_reused
+                << " op_name: " << mem_block_info.ordered_op_names.at(i);
       }
     }
   }

python/oneflow/framework/graph_build_util.py

@@ -27,7 +27,6 @@
 import oneflow.framework.scope_util as scope_util
 import oneflow.framework.session_context as session_context
 from oneflow.framework.tensor import Tensor
-
 import oneflow._oneflow_internal._C as _C
 
 lazy_mode = oneflow._oneflow_internal.lazy_mode
@@ -42,9 +41,11 @@ def graph_build_context(config_proto, session):
         config_proto_str, oneflow.placement("cpu", [0]), False,  # is_mirrored
     )
 
+    graph_scope = _make_new_graph_scope(new_scope, config_proto.job_name)
+
     with lazy_mode.guard(True):
         with JobBuildAndInferCtx(config_proto):
-            with BlockScopeContext(prev_scope, new_scope):
+            with BlockScopeContext(prev_scope, graph_scope):
                 yield
 
 
@@ -118,6 +119,36 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         )
 
 
+def _make_new_scope(prev_scope, scope_proto_str_setter):
+    new_scope = None
+
+    def build_scope(builder):
+        nonlocal new_scope
+        new_scope = builder.BuildScopeByProtoStrSetter(
+            prev_scope, scope_proto_str_setter
+        )
+        assert new_scope is not None
+
+    oneflow._oneflow_internal.deprecated.PhysicalRun(build_scope)
+    oneflow._oneflow_internal.eager.Sync()
+    return new_scope
+
+
+def _make_new_graph_scope(prev_scope, graph_name):
+    assert prev_scope is not None
+    attr_dict = dict()
+    name2default = session_context.GetDefaultSession().scope_attr_name2default_val
+
+    def scope_proto_str_setter(serialized_scope_proto: str):
+        scope_proto = text_format.Parse(
+            serialized_scope_proto, scope_pb2_util.ScopeProto()
+        )
+        scope_proto.module_name = graph_name
+        return str(text_format.MessageToString(scope_proto))
+
+    return _make_new_scope(prev_scope, scope_proto_str_setter)
+
+
 def make_new_block_scope(prev_scope, block):
     assert prev_scope is not None
     assert block is not None
@@ -147,21 +178,9 @@ def scope_proto_str_setter(serialized_scope_proto: str):
         # set module name
         if isinstance(block, oneflow.nn.graph.block.ModuleBlock):
             scope_proto.module_name = block.name_prefix + block.name
-
         return str(text_format.MessageToString(scope_proto))
 
-    new_scope = None
-
-    def build_scope(builder):
-        nonlocal new_scope
-        new_scope = builder.BuildScopeByProtoStrSetter(
-            prev_scope, scope_proto_str_setter
-        )
-        assert new_scope is not None
-
-    oneflow._oneflow_internal.deprecated.PhysicalRun(build_scope)
-    oneflow._oneflow_internal.eager.Sync()
-    return new_scope
+    return _make_new_scope(prev_scope, scope_proto_str_setter)
 
 
 def scope_to_proto(scope):

python/oneflow/nn/graph/block.py

@@ -75,6 +75,7 @@ def __init__(
         self._origin = None
         self._scope = None
         self._prev_scope = None
+        assert belonged_graph is None or isinstance(belonged_graph, weakref.ProxyTypes)
         self._belonged_graph = belonged_graph
         self.config = BlockConfig()
 
@@ -563,11 +564,13 @@ def _ops_repr(self):
         )
 
         if self._belonged_graph.is_compiled:
-            module_conf = self._belonged_graph._graph_proto.module_name2module_conf[
-                self.name_prefix + self.name
-            ]
-
-            return operators_repr(module_conf.ops)
+            if self._belonged_graph._compiled_graph_proto is not None:
+                module_conf = self._belonged_graph._compiled_graph_proto.module_name2module_conf[
+                    self.name_prefix + self.name
+                ]
+                return operators_repr(
+                    module_conf.ops, self._belonged_graph._compiled_graph_proto
+                )
 
         return []