Add support for lift_fresh op

llvm · Jul 22, 2022 · 329dc0e · 329dc0e
1 parent 005748b
commit 329dc0e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/build_tools/autogen_ltc_backend.yaml b/build_tools/autogen_ltc_backend.yaml
@@ -52,6 +52,7 @@ supported:
 
 # ops required for functionalization
 - lift
+- lift_fresh
 # Below are all operators that are "composite" in core,
 # but require us to explicitly re-enable functionalization in order to use them.
 # Why? These operators are all CompositeExplicitAutograd, which mean that they run

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
@@ -417,6 +417,12 @@ at::Tensor LazyNativeFunctions::lift(const at::Tensor& tensor) {
   return at::functionalization::impl::to_functional_tensor(tensor);
 }
 
+at::Tensor LazyNativeFunctions::lift_fresh(const at::Tensor& tensor) {
+  TORCH_INTERNAL_ASSERT(
+      !at::functionalization::impl::isFunctionalTensor(tensor));
+  return at::functionalization::impl::to_functional_tensor(tensor);
+}
+
 // All of the below ops correspond to CompositeExplicitAutograd kernels from core
 // that call into view operators internally.
 // These are all composite ops that LTC can technically re-use / get for free,