Glow model onnx exporters (pytorch#3167)

Summary:
Glow exporter (ONNX) serializes the Glow Function graph from memory into the file according to ONNX standard (extended).
Documentation:
Three layers of classes (similar to Glow loaders architecture). Protobuf writer, CommonOperatorWriter, and ONNXModelWriter.
[Optional Fixes #issue]
Pull Request resolved: pytorch#3167

Test Plan: Please see a detailed explanation of how to fill out the fields in the relevant sections in PULL_REQUEST.md.

Differential Revision: D15998903

fbshipit-source-id: 302b3c8884d501e16d138bce9f6edc68bb4cd677

Author: Yuri Putivsky

include/glow/Exporter/CommonOperatorWriter.h

@@ -0,0 +1,57 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License.");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef GLOW_EXPORTER_COMMONOPERATORWRITER_H
+#define GLOW_EXPORTER_COMMONOPERATORWRITER_H
+
+#include "glow/Exporter/ProtobufWriter.h"
+
+namespace glow {
+/// Declares writer methods for all operators. Every writer method serializes
+/// Glow node into the provided protobuf.
+template <typename Traits> class CommonOperatorWriter : public ProtobufWriter {
+protected:
+  virtual ~CommonOperatorWriter() = default;
+
+  /// Declare pure virtual methods, one per each node kind.
+  /// Derived class must to implement all of it.
+#define DEF_NODE(CLASS, NAME)                                                  \
+  virtual llvm::Error write##NAME(const CLASS *node,                           \
+                                  typename Traits::GraphProto &graph) = 0;
+#include "glow/AutoGenNodes.def"
+
+  /// Function invokes the correspondent virtual method according to \p node
+  /// type to serialize node information into \p graph (protobuf), reports
+  /// visited intermediate nodes through \p reporter, \returns llvm::Error.
+  llvm::Error writeOperator(const Node *node,
+                            typename Traits::GraphProto &graph) {
+    switch (node->getKind()) {
+#define DEF_NODE(CLASS, NAME)                                                  \
+  case glow::Kinded::Kind::CLASS##Kind:                                        \
+    return write##NAME(llvm::cast<CLASS>(node), graph);
+#include "glow/AutoGenNodes.def"
+    default:
+      llvm_unreachable(
+          "Not reachable, values and instructions are not handled here");
+      return llvm::Error::success();
+    }
+  }
+
+  using ProtobufWriter::ProtobufWriter;
+};
+} // namespace glow
+
+#endif // GLOW_EXPORTER_COMMONOPERATORWRITER_H

include/glow/Exporter/ONNXModelWriter.h

@@ -0,0 +1,95 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef GLOW_EXPORTER_ONNXMODELWRITER_H
+#define GLOW_EXPORTER_ONNXMODELWRITER_H
+
+#include "glow/Exporter/CommonOperatorWriter.h"
+#include "glow/Graph/Graph.h"
+
+#include "onnx/onnx_pb.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <string>
+
+/// ONNX traits for protobuf types.
+struct ONNX_TRAITS {
+  using GraphProto = ONNX_NAMESPACE::GraphProto;
+};
+
+namespace glow {
+
+/// Unique set of visited nodes.
+using ReportedNodes = std::unordered_set<const Node *>;
+
+/// Writes ONNX models.
+class ONNXModelWriter : public CommonOperatorWriter<ONNX_TRAITS> {
+  // Declare shorter aliases.
+  using GraphType = typename ONNX_TRAITS::GraphProto;
+  using NodeType = ONNX_NAMESPACE::NodeProto;
+  using TensorType = ONNX_NAMESPACE::TensorProto;
+  using AttrType = ONNX_NAMESPACE::AttributeProto;
+  using ValueInfoType = ONNX_NAMESPACE::ValueInfoProto;
+
+  /// Current version of ONNX standard.
+  size_t opsetVersion_;
+  /// Keeps the track of already visited or processed nodes.
+  ReportedNodes reportedNodes_;
+  /// Converts \p glowType to \p protoType.
+  static typename TensorType::DataType convertType(const Type &glowType);
+  /// Writes Glow tensor \p T to proto output \p out.
+  static void writeTensor(const Tensor &T, TensorType *out);
+  /// Writes tensor shape from placeholder \p PH into protpbuf \p valueProto.
+  static void tensorShapeFromPlaceholder(const Placeholder *PH,
+                                         ValueInfoType *valueProto);
+  /// Writes all inputs and outputs with operator name \p opName from give Node
+  /// \p node into protobuf \p proto.
+  static llvm::Error writeAllWithNode(const std::string &opName,
+                                      const Node *node, NodeType *proto);
+  /// Writes all inputs and outputs with operator name \p opName from give Node
+  /// \p node into created node protobuf using \p graph.
+  static llvm::Error writeAll(const std::string &opName, const Node *node,
+                              GraphType &graph);
+  // Finds if uses of \p node have node with the provided \p kind.
+  static bool hasUsesOfKind(const Node *node, Kinded::Kind kind);
+
+public:
+  /// Creates an ONNX model writer to serialize \p F graph into file
+  /// \p modelFilename, writing \p irVersion and \p opsetVersion.
+  /// If \p errPtr is not null then if an error occurs it will get assigned
+  /// there otherwise if an error occurs it will abort.
+  ONNXModelWriter(const std::string &modelFilename, Function &F,
+                  size_t irVersion, size_t opsetVersion,
+                  llvm::Error *errPtr = nullptr, bool textMode = false);
+
+private:
+  /// \returns error for the unexpected node kind.
+  static llvm::Error writeUnexpectedKind(const Node *node) {
+    RETURN_ERR(strFormat("Glow can not export node %s, unsupported kind: %d.",
+                         node->getName().str().c_str(), node->getKind()));
+  }
+
+  /// Declares the overriden all pure virtual methods, declared in base class.
+#define DEF_NODE(CLASS, NAME)                                                  \
+  llvm::Error write##NAME(const CLASS *, GraphType &) override;
+#include "glow/AutoGenNodes.def"
+};
+
+} // namespace glow
+
+#endif // GLOW_EXPORTER_ONNXMODELWRITER_H

include/glow/Exporter/ProtobufWriter.h

@@ -0,0 +1,49 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef GLOW_EXPORTER_PROTOBUFWRITER_H
+#define GLOW_EXPORTER_PROTOBUFWRITER_H
+
+#include "glow/Graph/Graph.h"
+#include "glow/Support/Error.h"
+
+#include <fstream>
+#include <google/protobuf/text_format.h>
+
+namespace glow {
+/// Writes model: graph and weights.
+class ProtobufWriter {
+protected:
+  /// The graph that we are constructing.
+  Function &G_;
+  /// Output file stream.
+  std::ofstream ff_;
+
+  llvm::Error writeModel(const ::google::protobuf::Message &modelProto,
+                         bool textMode = false);
+
+public:
+  /// Constructs new ProtobufWriter object. It will write protopuf messages into
+  /// \p modelFilename using graph and constants from \p F.
+  /// If \p errPtr is not null then if an error occurs it will get assigned
+  /// there otherwise if an error occurs it will abort.
+  ProtobufWriter(const std::string &modelFilename, Function &F,
+                 llvm::Error *errPtr = nullptr);
+};
+
+} // namespace glow
+
+#endif // GLOW_EXPORTER_PROTOBUFWRITER_H

include/glow/Support/Error.h

@@ -85,6 +85,10 @@ class GlowErr final : public llvm::ErrorInfo<GlowErr> {
     COMPILE_UNSUPPORTED_NODE_AFTER_OPTIMIZE,
     // Compilation error; Compilation context not correctly setup.
     COMPILE_CONTEXT_MALFORMED,
+    // Model writer encountered an invalid file name.
+    MODEL_WRITER_INVALID_FILENAME,
+    // Model writer cannot serialize graph to the file.
+    MODEL_WRITER_SERIALIZATION_ERROR,
   };
 
   /// GlowErr is not convertable to std::error_code. This is included for
@@ -156,6 +160,10 @@ class GlowErr final : public llvm::ErrorInfo<GlowErr> {
       return "COMPILE_UNSUPPORTED_NODE_AFTER_OPTIMIZE";
     case ErrorCode::COMPILE_CONTEXT_MALFORMED:
       return "COMPILE_CONTEXT_MALFORMED";
+    case ErrorCode::MODEL_WRITER_INVALID_FILENAME:
+      return "MODEL_WRITER_INVALID_FILENAME";
+    case ErrorCode::MODEL_WRITER_SERIALIZATION_ERROR:
+      return "MODEL_WRITER_SERIALIZATION_ERROR";
     };
 
     llvm_unreachable("unsupported ErrorCode");

lib/CMakeLists.txt

@@ -7,6 +7,7 @@ add_subdirectory(CodeGen)
 add_subdirectory(Converter)
 add_subdirectory(ExecutionContext)
 add_subdirectory(ExecutionEngine)
+add_subdirectory(Exporter)
 add_subdirectory(Graph)
 add_subdirectory(IR)
 add_subdirectory(Importer)

lib/Exporter/CMakeLists.txt

@@ -0,0 +1,34 @@
+include_directories(${PROTOBUF_INCLUDE_DIRS})
+include_directories(${CMAKE_CURRENT_BINARY_DIR})
+
+add_definitions(-DGOOGLE_PROTOBUF_NO_RTTI)
+
+add_library(Exporter
+              ProtobufWriter.cpp
+              ONNXModelWriter.cpp)
+target_compile_definitions(Exporter
+                           INTERFACE
+                             -DGOOGLE_PROTOBUF_NO_RTTI)
+target_link_libraries(Exporter
+                      PRIVATE
+                        Base
+                        Graph
+                        Importer
+                        LLVMSupport
+                        Support)
+target_link_libraries(Exporter PUBLIC onnx_proto ${PROTOBUF_LIBRARY})
+
+if (MSVC AND LINK_PROTOBUF_AS_DLL)
+  # For protobuf warning when it is build as dll.
+  # Suppresses a warning that is treated as error.
+  # Basically one of the header files has interface class
+  # containing STL string. Which might cause issues
+  # if things are build with different compilers.
+  #
+  # Sets general warning level as 2 for this project.
+  # There are few warnings that are treated as errors that
+  # come from VS include headers
+  target_compile_options(onnx_proto PUBLIC /wd4251)
+  target_compile_options(onnx_proto PUBLIC /W2)
+  target_compile_definitions(onnx_proto PUBLIC -DPROTOBUF_USE_DLLS)
+endif()