Merge pull request #203 from ROCmSoftwarePlatform/tf_pb_py

tf protobuf pybind support

CMakeLists.txt

@@ -39,6 +39,8 @@ else()
     set(MIGRAPHX_ENABLE_GPU Off CACHE BOOL "")
 endif()
 
+set(MIGRAPHX_ENABLE_TF Off CACHE BOOL "")
+
 add_compile_options(-std=c++14)
 
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake)

src/include/migraphx/onnx.hpp

@@ -7,24 +7,6 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-struct unknown
-{
-    std::string op;
-    std::string name() const { return "unknown:" + op; }
-    shape compute_shape(std::vector<shape> input) const
-    {
-        if(input.empty())
-            return {};
-        else
-            return input.front();
-    }
-    friend std::ostream& operator<<(std::ostream& os, const unknown& x)
-    {
-        os << x.name();
-        return os;
-    }
-};
-
 /// Create a program from an onnx file
 program parse_onnx(const std::string& name);
 

src/include/migraphx/operators.hpp

@@ -1366,6 +1366,25 @@ struct undefined
     argument compute(const shape&, const std::vector<argument>&) const { return {{}, nullptr}; }
 };
 
+struct unknown
+{
+    std::string op;
+    std::string name() const { return "unknown:" + op; }
+    shape compute_shape(std::vector<shape> input) const
+    {
+        if(input.empty())
+            return {};
+        else
+            return input.front();
+    }
+
+    friend std::ostream& operator<<(std::ostream& os, const unknown& x)
+    {
+        os << x.name();
+        return os;
+    }
+};
+
 } // namespace op
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/include/migraphx/tf.hpp

@@ -7,25 +7,7 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-struct unknown
-{
-    std::string op;
-    std::string name() const { return "unknown:" + op; }
-    shape compute_shape(std::vector<shape> input) const
-    {
-        if(input.empty())
-            return {};
-        else
-            return input.front();
-    }
-    friend std::ostream& operator<<(std::ostream& os, const unknown& x)
-    {
-        os << x.name();
-        return os;
-    }
-};
-
-/// Create a program from an onnx file
+/// Create a program from a tf pb file (default is nhwc format)
 program parse_tf(const std::string& name, bool is_nhwc);
 
 } // namespace MIGRAPHX_INLINE_NS

src/onnx/onnx.cpp

@@ -1206,7 +1206,7 @@ struct onnx_parser
             std::vector<instruction_ref> result;
             if(ops.count(node.op_type()) == 0)
             {
-                result.push_back(prog.add_instruction(unknown{node.op_type()}, args));
+                result.push_back(prog.add_instruction(op::unknown{node.op_type()}, args));
             }
             else
             {

src/py/CMakeLists.txt

@@ -12,7 +12,12 @@ if(MIGRAPHX_ENABLE_PYTHON)
         C_VISIBILITY_PRESET hidden
         CXX_VISIBILITY_PRESET hidden
     )
-    target_link_libraries(migraphx_py PRIVATE migraphx migraphx_onnx migraphx_cpu)
+    if(MIGRAPHX_ENABLE_TF)
+        target_link_libraries(migraphx_py PRIVATE migraphx migraphx_tf migraphx_cpu)
+        target_compile_definitions(migraphx_py PRIVATE -DENABLE_TF)
+    else()
+        target_link_libraries(migraphx_py PRIVATE migraphx migraphx_onnx migraphx_cpu)
+    endif()
     if(MIGRAPHX_ENABLE_GPU)
         target_link_libraries(migraphx_py PRIVATE migraphx_gpu)
         target_compile_definitions(migraphx_py PRIVATE -DHAVE_GPU)

src/py/migraphx_py.cpp

@@ -4,8 +4,13 @@
 #include <migraphx/program.hpp>
 #include <migraphx/generate.hpp>
 #include <migraphx/cpu/target.hpp>
-#include <migraphx/onnx.hpp>
 #include <migraphx/stringutils.hpp>
+#ifdef ENABLE_TF
+#include <migraphx/tf.hpp>
+#else
+#include <migraphx/onnx.hpp>
+#endif
+
 #ifdef HAVE_GPU
 #include <migraphx/gpu/target.hpp>
 #include <migraphx/gpu/hip.hpp>
@@ -155,8 +160,16 @@ PYBIND11_MODULE(migraphx, m)
         .def("__ne__", std::not_equal_to<migraphx::program>{})
         .def("__repr__", [](const migraphx::program& p) { return migraphx::to_string(p); });
 
+#ifdef ENABLE_TF
+    m.def("parse_tf",
+          &migraphx::parse_tf,
+          "Parse tf protobuf (default format is nhwc)",
+          py::arg("filename"),
+          py::arg("is_nhwc") = true);
+#else
     m.def("parse_onnx", &migraphx::parse_onnx);
 
+#endif
     m.def("get_target", [](const std::string& name) -> migraphx::target {
         if(name == "cpu")
             return migraphx::cpu::target{};

src/tf/tf.cpp

@@ -525,7 +525,7 @@ struct tf_parser
             }
             if(ops.count(node.op()) == 0)
             {
-                instructions[name] = prog.add_instruction(unknown{node.op()}, args);
+                instructions[name] = prog.add_instruction(op::unknown{node.op()}, args);
             }
             else
             {

test/onnx/onnx_test.cpp

@@ -392,8 +392,8 @@ TEST_CASE(unknown_test)
     migraphx::program p;
     auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}});
     auto l1 = p.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {3, 4}});
-    auto l2 = p.add_instruction(migraphx::unknown{"Unknown"}, l0, l1);
-    p.add_instruction(migraphx::unknown{"Unknown"}, l2);
+    auto l2 = p.add_instruction(migraphx::op::unknown{"Unknown"}, l0, l1);
+    p.add_instruction(migraphx::op::unknown{"Unknown"}, l2);
     auto prog = migraphx::parse_onnx("unknown_test.onnx");
 
     EXPECT(p == prog);