Update MKL exporter to IDEEP ops (#8228)

IDEEP exporter support

3 files changed, 31 insertions, 20 deletions

diff --git a/caffe2/ideep/operators/operator_fallback_ideep.cc b/caffe2/ideep/operators/operator_fallback_ideep.cc
index 216633e38d..1f3ad1eb98 100644
--- a/caffe2/ideep/operators/operator_fallback_ideep.cc
+++ b/caffe2/ideep/operators/operator_fallback_ideep.cc
@@ -8,6 +8,7 @@
 #include <caffe2/operators/filler_op.h>
 #include <caffe2/operators/flatten_op.h>
 #include <caffe2/operators/given_tensor_fill_op.h>
+#include <caffe2/operators/load_save_op.h>
 #include <caffe2/operators/loss_op.h>
 #include <caffe2/operators/reshape_op.h>
 #include <caffe2/operators/softmax_op.h>
@@ -53,5 +54,7 @@ REGISTER_IDEEP_OPERATOR(
 REGISTER_IDEEP_OPERATOR(
     GivenTensorFill,
     IDEEPFallbackOp<GivenTensorFillOp<float, CPUContext>>);
+REGISTER_IDEEP_OPERATOR(Load, IDEEPFallbackOp<LoadOp<CPUContext>>);
+REGISTER_IDEEP_OPERATOR(Save, IDEEPFallbackOp<SaveOp<CPUContext>>);
 
 } // namespace caffe2
diff --git a/caffe2/python/mkl/rewrite_graph.py b/caffe2/python/mkl/rewrite_graph.py
index 45baa59aed..ccf800e440 100644
--- a/caffe2/python/mkl/rewrite_graph.py
+++ b/caffe2/python/mkl/rewrite_graph.py
@@ -8,9 +8,10 @@ from caffe2.proto import caffe2_pb2
 from caffe2.python import core
 
 
-def rewrite_init_net_simple(net):
+def rewrite_init_net_simple(net, ideep=True):
+    device = caffe2_pb2.IDEEP if ideep else caffe2_pb2.MKLDNN
     for op in net.op:
-        op.device_option.device_type = caffe2_pb2.MKLDNN
+        op.device_option.device_type = device
 
 def last_producer(ops, blob):
     for (i, op) in reversed(list(enumerate(ops))):
@@ -19,7 +20,7 @@ def last_producer(ops, blob):
     raise ValueError("Failed to find last producer of blob, %s", blob)
 
 
-def rewrite_run_net_simple(net):
+def rewrite_run_net_simple(net, ideep=True):
     # Simple rewrite for now - assume entire graph can be executed
     # with MKL, so just insert copy ops for external_input[0] and
     # external_output[0]
@@ -32,12 +33,14 @@ def rewrite_run_net_simple(net):
             "Input blob: {} is not consumed by first op: {}".format(
                 input_blob, net.op[0]))
     # Modify input/outputs to point to copied MKL blobs.
+    from_cpu = "CopyCPUToIDEEP" if ideep else "CopyCPUToMKL"
+    to_cpu = "CopyIDEEPToCPU" if ideep else "CopyMKLToCPU"
     copy_input_op = core.CreateOperator(
-        "CopyCPUToMKL", input_blob, mkl_tmp(input_blob))
+        from_cpu, input_blob, mkl_tmp(input_blob))
     net.op[0].input[0] = mkl_tmp(input_blob)
 
     copy_output_ops = [
-        core.CreateOperator("CopyMKLToCPU", mkl_tmp(output_blob), output_blob)
+        core.CreateOperator(to_cpu, mkl_tmp(output_blob), output_blob)
         for output_blob in net.external_output]
 
     for output_blob in net.external_output:
@@ -54,15 +57,16 @@ def rewrite_run_net_simple(net):
     ops = [copy_input_op] + net.op[:] + copy_output_ops
     del net.op[:]
     net.op.extend(ops)
+    device = caffe2_pb2.IDEEP if ideep else caffe2_pb2.MKLDNN
     for op in net.op:
         op.device_option.MergeFrom(
-            core.DeviceOption(device_type=caffe2_pb2.MKLDNN))
+            core.DeviceOption(device_type=device))
         op.engine = ""
 
 
-def rewrite_model_helper_simple(model):
+def rewrite_model_helper_simple(model, ideep=True):
     model = copy.deepcopy(model)
     # All parameter initialization should run on MKL
-    rewrite_init_net_simple(model.param_init_net.Proto())
-    rewrite_run_net_simple(model.net.Proto())
+    rewrite_init_net_simple(model.param_init_net.Proto(), ideep)
+    rewrite_run_net_simple(model.net.Proto(), ideep)
     return model
diff --git a/caffe2/python/mkl/rewrite_graph_test.py b/caffe2/python/mkl/rewrite_graph_test.py
index 65bf7fd53d..87b08d5283 100644
--- a/caffe2/python/mkl/rewrite_graph_test.py
+++ b/caffe2/python/mkl/rewrite_graph_test.py
@@ -180,15 +180,16 @@ def complex_resnet():
     return model, [(1, 1, 224, 224)]
 
 
-@unittest.skipIf(not workspace.C.has_mkldnn,
-                 "Skipping as we do not have mkldnn.")
+@unittest.skipIf(not workspace.C.has_mkldnn or not workspace.C.use_ideep,
+                 "Skipping as we do not have MKLDNN and IDEEP.")
 class MKLRewriteTest(hu.HypothesisTestCase):
     @given(gen=st.sampled_from([simple_relu, simple_fc,
-                                simple_mlp, simple_cnn]))
-    def test_mkl_simple_rewrite(self, gen):
+                                simple_mlp, simple_cnn]),
+           ideep=st.booleans())
+    def test_mkl_simple_rewrite(self, gen, ideep):
         cpu_model, (shape,) = gen()
         cpu_model = deterministic_io(cpu_model)
-        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
+        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model, ideep)
         X = np.random.randn(*shape).astype(np.float32)
 
         def run(model):
@@ -200,10 +201,11 @@ class MKLRewriteTest(hu.HypothesisTestCase):
         np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                    atol=1e-4, rtol=1e-4)
 
-    def test_mkl_resnet_rewrite(self):
+    @given(ideep=st.booleans())
+    def test_mkl_resnet_rewrite(self, ideep):
         cpu_model, (shape,) = complex_resnet()
         cpu_model = deterministic_io(cpu_model)
-        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
+        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model, ideep)
         np.random.seed(1701)
         X = np.random.randn(*shape).astype(np.float32)
 
@@ -215,10 +217,11 @@ class MKLRewriteTest(hu.HypothesisTestCase):
         np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                    atol=1e-4, rtol=1e-4)
 
-    def test_mkl_multi_output_rewrite(self):
+    @given(ideep=st.booleans())
+    def test_mkl_multi_output_rewrite(self, ideep):
         cpu_model, shapes = double_matmul()
         cpu_model = deterministic_io(cpu_model)
-        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
+        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model, ideep)
         np.random.seed(1701)
         Xs = [np.random.randn(*shape).astype(np.float32) for shape in shapes]
 
@@ -236,10 +239,11 @@ class MKLRewriteTest(hu.HypothesisTestCase):
         np.testing.assert_allclose(run(cpu_model), run(mkl_model),
                                    atol=1e-4, rtol=1e-4)
 
-    def test_mkl_alexnet_rewrite(self):
+    @given(ideep=st.booleans())
+    def test_mkl_alexnet_rewrite(self, ideep):
         cpu_model, (shape,) = alexnet()
         cpu_model = deterministic_io(cpu_model)
-        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model)
+        mkl_model = rewrite_graph.rewrite_model_helper_simple(cpu_model, ideep)
         np.random.seed(1701)
         X = np.random.randn(*shape).astype(np.float32)