From a8755bb5221ae6c2c44cdbadf89a2e2eaa265cc2 Mon Sep 17 00:00:00 2001
From: Sahan Paliskara <sahancpal@gmail.com>
Date: Fri, 13 Jan 2023 17:51:09 -0800
Subject: [PATCH] torch dyanamo in c++ test

ghstack-source-id: 53fa37b8b40ac4c4b4778cef83f9405aba1d6259
Pull Request resolved: https://github.com/pytorch/multipy/pull/299
---
 multipy/runtime/CMakeLists.txt         | 12 +++++
 multipy/runtime/test_deploy.cpp        |  4 ++
 multipy/runtime/test_dynamo_compat.cpp | 68 ++++++++++++++++++++++++++
 3 files changed, 84 insertions(+)
 create mode 100644 multipy/runtime/test_dynamo_compat.cpp

diff --git a/multipy/runtime/CMakeLists.txt b/multipy/runtime/CMakeLists.txt
index 4d50f210..bf18f6fb 100644
--- a/multipy/runtime/CMakeLists.txt
+++ b/multipy/runtime/CMakeLists.txt
@@ -88,6 +88,10 @@ set(INTERPRETER_TEST_SOURCES_GPU
   ${DEPLOY_DIR}/test_deploy_gpu.cpp
 )
 
+set(INTERPRETER_TEST_SOURCES_COMPAT
+  ${DEPLOY_DIR}/test_dynamo_compat.cpp
+)
+
 # TODO: Currently tests can only be done when ABI=1 as the testing infrustructure
 # used by ASSERT_TRUE requires ABI=1 in Github actions, we should fix this!
 
@@ -99,6 +103,14 @@ target_link_libraries(test_deploy
 )
 target_include_directories(test_deploy PRIVATE ${CMAKE_SOURCE_DIR}/../..)
 
+add_executable(test_compat ${INTERPRETER_TEST_SOURCES_COMPAT})
+# target_compile_definitions(test_compat PUBLIC TEST_CUSTOM_LIBRARY)
+target_include_directories(test_compat PRIVATE ${PYTORCH_ROOT}/torch)
+target_link_libraries(test_compat
+  PUBLIC "-Wl,--no-as-needed -rdynamic" gtest dl torch_deploy_interface c10 torch_cpu
+)
+target_include_directories(test_compat PRIVATE ${CMAKE_SOURCE_DIR}/../..)
+
 if(BUILD_CUDA_TESTS)
   LINK_DIRECTORIES("${PYTORCH_ROOT}/torch/lib")
   add_executable(test_deploy_gpu ${INTERPRETER_TEST_SOURCES_GPU})
diff --git a/multipy/runtime/test_deploy.cpp b/multipy/runtime/test_deploy.cpp
index b559eab7..18dfc107 100644
--- a/multipy/runtime/test_deploy.cpp
+++ b/multipy/runtime/test_deploy.cpp
@@ -88,6 +88,10 @@ TEST(TorchpyTest, SimpleModel) {
   compare_torchpy_jit(path("SIMPLE", simple), path("SIMPLE_JIT", simple_jit));
 }
 
+TEST(TorchpyTest, DynamoTest) {
+
+}
+
 #ifdef FBCODE_CAFFE2
 TEST(TorchpyTest, LoadTextAndBinary) {
   torch::deploy::InterpreterManager manager(1);
diff --git a/multipy/runtime/test_dynamo_compat.cpp b/multipy/runtime/test_dynamo_compat.cpp
new file mode 100644
index 00000000..7d24c687
--- /dev/null
+++ b/multipy/runtime/test_dynamo_compat.cpp
@@ -0,0 +1,68 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <ATen/Parallel.h>
+#include <gtest/gtest.h>
+#include <libgen.h>
+#include <cstring>
+
+#include <c10/util/irange.h>
+#include <libgen.h>
+#include <multipy/runtime/deploy.h>
+#include <torch/script.h>
+#include <torch/torch.h>
+
+#include <future>
+#include <iostream>
+#include <string>
+
+void compare_torchpy_jit(const char* model_filename, const char* jit_filename) {
+  // Test
+
+  torch::deploy::InterpreterManager m(2);
+  torch::deploy::Package p = m.loadPackage(model_filename);
+  auto model = p.loadPickle("model", "model.pkl");
+  at::IValue eg;
+  {
+    auto I = p.acquireSession();
+    eg = I.self.attr("load_pickle")({"model", "example.pkl"}).toIValue();
+  }
+  auto I = p.acquireSession();
+  auto cModelObj = I.global("torch", "compile")(model.toObj(&I));
+  auto cModel = m.createMovable(cModelObj, &I);
+  at::Tensor output = cModel(eg.toTupleRef().elements()).toTensor();
+
+  // Reference
+  auto ref_model = torch::jit::load(jit_filename);
+  at::Tensor ref_output =
+      ref_model.forward(eg.toTupleRef().elements()).toTensor();
+
+  ASSERT_TRUE(ref_output.allclose(output, 1e-03, 1e-05));
+}
+
+const char* simple = "multipy/runtime/example/generated/simple";
+const char* simple_jit = "multipy/runtime/example/generated/simple_jit";
+
+const char* path(const char* envname, const char* path) {
+  const char* e = getenv(envname);
+  return e ? e : path;
+}
+
+TEST(TorchpyTest, SimpleModel) {
+  compare_torchpy_jit(path("SIMPLE", simple), path("SIMPLE_JIT", simple_jit));
+}
+
+int main(int argc, char* argv[]) {
+  ::testing::InitGoogleTest(&argc, argv);
+  char tempeh[256];
+  getcwd(tempeh, 256);
+  std::cout << "Current working directory: " << tempeh << std::endl;
+  int rc = RUN_ALL_TESTS();
+  char tmp[256];
+  getcwd(tmp, 256);
+  std::cout << "Current working directory: " << tmp << std::endl;
+  return rc;
+}