From 329dc0eb69585713eb895194b5d3e6ddfd2247f7 Mon Sep 17 00:00:00 2001 From: Antonio Kim Date: Thu, 21 Jul 2022 15:04:50 -0700 Subject: [PATCH] Add support for lift_fresh op --- build_tools/autogen_ltc_backend.yaml | 1 + .../csrc/base_lazy_backend/mlir_native_functions.cpp | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/build_tools/autogen_ltc_backend.yaml b/build_tools/autogen_ltc_backend.yaml index b8a8905e06c..ff6c7f5eeab 100644 --- 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 index 13522c605da..eb793615b86 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 @@ -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,