From 2d835e803ae900585d50eeca2f9dfc751d677ee2 Mon Sep 17 00:00:00 2001
From: Henry Tu <henry@henrytu.me>
Date: Fri, 26 Aug 2022 14:01:36 -0400
Subject: [PATCH] Revert "Revert updating mlir_native_functions.cpp signature
 (#1281)"

This reverts commit a1ace0657d1027a82b09a652a157a4c45e9d62d1.
---
 e2e_testing/torchscript/xfail_sets.py         |  6 ----
 .../mlir_native_functions.cpp                 | 29 ++++++++++++++-----
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/e2e_testing/torchscript/xfail_sets.py b/e2e_testing/torchscript/xfail_sets.py
index 8904d7faa91..f106c974910 100644
--- a/e2e_testing/torchscript/xfail_sets.py
+++ b/e2e_testing/torchscript/xfail_sets.py
@@ -435,12 +435,6 @@
     "NewOnesModuleFloat3D_basic",
     "NewOnesModuleInt2D_basic",
     "NewOnesModuleInt3D_basic",
-    "NewZerosModuleDefaultDtype_basic",
-    "NewZerosModuleFalsePinMemory_basic",
-    "NewZerosModuleFloat2D_basic",
-    "NewZerosModuleFloat3D_basic",
-    "NewZerosModuleInt2D_basic",
-    "NewZerosModuleInt3D_basic",
     "OnesLikeModule_defaultDtype",
     "OnesLikeModule_falsePinMemory",
     "OnesLikeModule_float",
diff --git a/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp b/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp
index ec72c37ab50..cfd0cf68e15 100644
--- a/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp
+++ b/python/torch_mlir/csrc/base_lazy_backend/mlir_native_functions.cpp
@@ -302,10 +302,14 @@ at::Tensor LazyNativeFunctions::_to_copy(
 };
 
 at::Tensor LazyNativeFunctions::empty(
-    at::IntArrayRef size, c10::optional<at::ScalarType> dtype,
-    c10::optional<at::Layout> layout, c10::optional<at::Device> device,
+    at::SymIntArrayRef sym_size,
+    c10::optional<at::ScalarType> dtype,
+    c10::optional<at::Layout> layout,
+    c10::optional<at::Device> device,
     c10::optional<bool> pin_memory,
     c10::optional<at::MemoryFormat> memory_format) {
+  // TODO: support this directly
+  auto size = c10::asIntArrayRefSlow(sym_size);
   const auto device_type = torch::lazy::getBackend()->EagerFallbackDeviceType();
   at::TensorOptions options = at::TensorOptions()
                                   .device(c10::Device(device_type))
@@ -317,8 +321,9 @@ at::Tensor LazyNativeFunctions::empty(
   // See Note [Lazy Tensor Functionalization]
   if (c10::impl::tls_local_dispatch_key_set().excluded_.has(
           c10::DispatchKey::Functionalize)) {
-    // Invariant: if the functionalization key is in the exclude set, then we're expected
-    // to return an ordinary tensor, which will be "lifted" into a functional wrapper later.
+    // Invariant: if the functionalization key is in the exclude set, then we're
+    // expected to return an ordinary tensor, which will be "lifted" into a
+    // functional wrapper later.
     return tensor;
   } else {
     auto wrapped = at::functionalization::impl::to_functional_tensor(tensor);
@@ -331,7 +336,13 @@ at::Tensor LazyNativeFunctions::empty_strided(
     c10::optional<at::ScalarType> dtype, c10::optional<at::Layout> layout,
     c10::optional<at::Device> device, c10::optional<bool> pin_memory) {
   TORCH_LAZY_FN_COUNTER("lazy::");
-  at::Tensor t = empty(size, dtype, layout, device, pin_memory, c10::nullopt);
+  at::Tensor t = empty(
+      c10::SymIntArrayRef::fromIntArrayRef(size),
+      dtype,
+      layout,
+      device,
+      pin_memory,
+      c10::nullopt);
   return t.as_strided(size, stride, /*storage_offset=*/0);
 }
 
@@ -350,7 +361,8 @@ LazyNativeFunctions::fill_(at::Tensor& self, const at::Scalar& value) {
 at::Tensor LazyNativeFunctions::_unsafe_view(
     const at::Tensor& self, at::IntArrayRef size) {
   TORCH_LAZY_FN_COUNTER("lazy::");
-  return LazyNativeFunctions::view_copy(self, size);
+  return LazyNativeFunctions::view_copy(
+      self, c10::SymIntArrayRef::fromIntArrayRef(size));
 }
 
 // This is needed by the torch.tensor constructor.
@@ -386,7 +398,10 @@ at::Tensor LazyNativeFunctions::new_empty_strided(
 }
 
 at::Tensor LazyNativeFunctions::narrow_copy(
-    const at::Tensor& self, int64_t dim, int64_t start, int64_t length) {
+    const at::Tensor& self,
+    int64_t dim,
+    c10::SymInt start,
+    c10::SymInt length) {
   return at::functionalization::functionalize_aten_op<ATEN_OP(
       narrow_copy)>::call(self, dim, start, length);
 }