From 82f414bed3506a8c55ae1874d811e6e1c50cd7e0 Mon Sep 17 00:00:00 2001
From: powderluv <powderluv@users.noreply.github.com>
Date: Wed, 31 Aug 2022 15:01:59 -0700
Subject: [PATCH] Add shapelib and Torch ODS gen tests (#1318)

---
 build_tools/python_deploy/build_linux_packages.sh | 14 ++++++++++++++
 build_tools/update_shape_lib.sh                   |  9 +++++++++
 build_tools/update_torch_ods.sh                   |  9 +++++++++
 3 files changed, 32 insertions(+)

diff --git a/build_tools/python_deploy/build_linux_packages.sh b/build_tools/python_deploy/build_linux_packages.sh
index 44ab910c0fb5f..f86fa8365c8e3 100755
--- a/build_tools/python_deploy/build_linux_packages.sh
+++ b/build_tools/python_deploy/build_linux_packages.sh
@@ -193,6 +193,20 @@ function test_in_tree() {
   cd /main_checkout/torch-mlir/
   export PYTHONPATH="/main_checkout/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir"
 
+  echo ":::: Run shapelib update tests"
+  if ! ./build_tools/update_shape_lib.sh; then
+    echo Shape Lib is out of date with the installed PyTorch version
+  else
+    echo Shape Lib is up to date
+  fi
+
+  echo ":::: Run torch_ods update tests"
+  if ! ./build_tools/update_torch_ods.sh; then
+    echo Torch ODS is out of date with the installed PyTorch version
+  else
+    echo Torch ODS is up to date
+  fi
+
   echo ":::: Run refbackend e2e integration tests"
   python -m e2e_testing.main --config=refbackend -v
 
diff --git a/build_tools/update_shape_lib.sh b/build_tools/update_shape_lib.sh
index b2d619d3484be..3422440a98419 100755
--- a/build_tools/update_shape_lib.sh
+++ b/build_tools/update_shape_lib.sh
@@ -30,3 +30,12 @@ PYTHONPATH="${pypath}" python \
   -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.shape_lib_gen \
   --pytorch_op_extensions=${ext_module:-""} \
   --torch_transforms_cpp_dir="${torch_transforms_cpp_dir}"
+
+if ! git diff --quiet ./lib/Dialect/Torch/Transforms/ShapeLibrary.cpp &> /dev/null; then
+  echo "#######################################################"
+  echo "Generated files are not up to date (see diff below)"
+  echo ">>> Please verify you are running the latest PyTorch nightly or source build and if required push these changes up for review <<<"
+  echo "#######################################################"
+  git diff --color=always ./lib/Dialect/Torch/Transforms/ShapeLibrary.cpp
+  exit 1
+fi
diff --git a/build_tools/update_torch_ods.sh b/build_tools/update_torch_ods.sh
index 2b30ffe663aeb..e0f407d17ac01 100755
--- a/build_tools/update_torch_ods.sh
+++ b/build_tools/update_torch_ods.sh
@@ -32,3 +32,12 @@ PYTHONPATH="${pypath}" python \
   --torch_ir_include_dir="${torch_ir_include_dir}" \
   --pytorch_op_extensions="${ext_module}" \
   --debug_registry_dump="${torch_ir_include_dir}/JITOperatorRegistryDump.txt"
+
+if ! git diff --quiet ./include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td &> /dev/null; then
+  echo "#######################################################"
+  echo "Generated files are not up to date (see diff below)"
+  echo ">>> Please verify you are running the latest PyTorch nightly or source build and if required push these changes up for review <<<"
+  echo "#######################################################"
+  git diff --color=always ./include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
+  exit 1
+fi