Added test case for DLA device serialization

Signed-off-by: Anurag Dixit <[email protected]>

Author: Anurag Dixit

core/runtime/TRTEngine.cpp

@@ -11,8 +11,6 @@ namespace trtorch {
 namespace core {
 namespace runtime {
 
-const std::string empty_string = std::string();
-
 std::string slugify(std::string s) {
   std::replace(s.begin(), s.end(), '.', '_');
   return s;
@@ -24,7 +22,7 @@ TRTEngine::TRTEngine(std::string serialized_engine)
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
   std::string _name = "deserialized_trt";
-  new (this) TRTEngine(_name, serialized_engine, empty_string);
+  new (this) TRTEngine(_name, serialized_engine, std::string());
 }
 
 TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
@@ -42,7 +40,7 @@ TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
 TRTEngine::TRTEngine(
     std::string mod_name,
     std::string serialized_engine,
-    std::string serialized_device_info = empty_string)
+    std::string serialized_device_info = std::string())
     : logger(
           std::string("[") + mod_name + std::string("_engine] - "),
           util::logging::get_logger().get_reportable_severity(),

tests/modules/BUILD

@@ -28,7 +28,8 @@ test_suite(
        ":test_modules_as_engines",
        ":test_compiled_modules",
        ":test_multiple_registered_engines",
-       ":test_serialization"
+       ":test_serialization",
+       ":test_dla_serialization"
    ]
 )
 
@@ -43,6 +44,17 @@ cc_test(
     ]
 )
 
+cc_test(
+    name = "test_dla_serialization",
+    srcs = ["test_dla_serialization.cpp"],
+    deps = [
+        ":module_test",
+    ],
+    data = [
+        ":jit_models"
+    ]
+)
+
 cc_test(
     name = "test_multiple_registered_engines",
     srcs = ["test_multiple_registered_engines.cpp"],

tests/modules/test_dla_serialization.cpp

@@ -0,0 +1,102 @@
+#include "module_test.h"
+
+std::vector<trtorch::CompileSpec::InputRange> toInputRangesDynamic(std::vector<std::vector<int64_t>> opts) {
+  std::vector<trtorch::CompileSpec::InputRange> a;
+
+  for (auto opt : opts) {
+    std::vector<int64_t> min_range(opt);
+    std::vector<int64_t> max_range(opt);
+
+    min_range[3] = ceil(opt[3] / 2.0);
+    max_range[3] = 2 * opt[3];
+    min_range[2] = ceil(opt[2] / 2.0);
+    max_range[2] = 2 * opt[2];
+
+    a.push_back(trtorch::CompileSpec::InputRange(min_range, opt, max_range));
+  }
+
+  return std::move(a);
+}
+
+TEST_P(ModuleTests, SerializedModuleIsStillCorrect) {
+  trtorch::set_device(0);
+  std::vector<torch::jit::IValue> post_serialized_inputs_ivalues;
+  std::vector<torch::jit::IValue> pre_serialized_inputs_ivalues;
+  for (auto in_shape : input_shapes) {
+    auto in = at::randint(5, in_shape, {at::kCUDA}).to(torch::kF16);
+    post_serialized_inputs_ivalues.push_back(in.clone());
+    pre_serialized_inputs_ivalues.push_back(in.clone());
+  }
+
+  auto compile_spec = trtorch::CompileSpec(toInputRangesDynamic(input_shapes));
+  compile_spec.op_precision = torch::kF16;
+  compile_spec.device.device_type = trtorch::CompileSpec::Device::DeviceType::kDLA;
+  compile_spec.device.gpu_id = 0;
+  compile_spec.device.dla_core = 1;
+  compile_spec.device.allow_gpu_fallback = true;
+  compile_spec.workspace_size = 1 << 28;
+
+  auto pre_serialized_mod = trtorch::CompileGraph(mod, compile_spec);
+  torch::jit::IValue pre_serialized_results_ivalues =
+      trtorch::tests::util::RunModuleForward(pre_serialized_mod, pre_serialized_inputs_ivalues);
+  std::vector<at::Tensor> pre_serialized_results;
+  pre_serialized_results.push_back(pre_serialized_results_ivalues.toTensor());
+
+  pre_serialized_mod.save("test_serialization_mod.ts");
+  auto post_serialized_mod = torch::jit::load("test_serialization_mod.ts");
+
+  torch::jit::IValue post_serialized_results_ivalues =
+      trtorch::tests::util::RunModuleForward(post_serialized_mod, post_serialized_inputs_ivalues);
+  std::vector<at::Tensor> post_serialized_results;
+  post_serialized_results.push_back(post_serialized_results_ivalues.toTensor());
+
+  for (size_t i = 0; i < pre_serialized_results.size(); i++) {
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(
+        post_serialized_results[i], pre_serialized_results[i].reshape_as(post_serialized_results[i]), 2e-5));
+  }
+}
+
+TEST_P(ModuleTests, SerializedDynamicModuleIsStillCorrect) {
+  trtorch::set_device(0);
+  std::vector<torch::jit::IValue> post_serialized_inputs_ivalues;
+  std::vector<torch::jit::IValue> pre_serialized_inputs_ivalues;
+  for (auto in_shape : input_shapes) {
+    auto in = at::randint(5, in_shape, {at::kCUDA}).to(torch::kF16);
+    post_serialized_inputs_ivalues.push_back(in.clone());
+    pre_serialized_inputs_ivalues.push_back(in.clone());
+  }
+
+  auto compile_spec = trtorch::CompileSpec(toInputRangesDynamic(input_shapes));
+  compile_spec.op_precision = torch::kF16;
+  compile_spec.device.device_type = trtorch::CompileSpec::Device::DeviceType::kDLA;
+  compile_spec.device.gpu_id = 0;
+  compile_spec.device.dla_core = 1;
+  compile_spec.device.allow_gpu_fallback = true;
+  compile_spec.workspace_size = 1 << 28;
+
+  auto pre_serialized_mod = trtorch::CompileGraph(mod, compile_spec);
+  torch::jit::IValue pre_serialized_results_ivalues =
+      trtorch::tests::util::RunModuleForward(pre_serialized_mod, pre_serialized_inputs_ivalues);
+  std::vector<at::Tensor> pre_serialized_results;
+  pre_serialized_results.push_back(pre_serialized_results_ivalues.toTensor());
+
+  pre_serialized_mod.save("test_serialization_mod.ts");
+  auto post_serialized_mod = torch::jit::load("test_serialization_mod.ts");
+
+  torch::jit::IValue post_serialized_results_ivalues =
+      trtorch::tests::util::RunModuleForward(post_serialized_mod, post_serialized_inputs_ivalues);
+  std::vector<at::Tensor> post_serialized_results;
+  post_serialized_results.push_back(post_serialized_results_ivalues.toTensor());
+
+  for (size_t i = 0; i < pre_serialized_results.size(); i++) {
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(
+        post_serialized_results[i], pre_serialized_results[i].reshape_as(post_serialized_results[i]), 2e-5));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CompiledModuleForwardIsCloseSuite,
+    ModuleTests,
+    testing::Values(
+        PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}}),
+        PathAndInSize({"tests/modules/pooling_traced.jit.pt", {{1, 3, 10, 10}}})));