diff --git a/build_tools/autogen_ltc_backend.py b/build_tools/autogen_ltc_backend.py index 8012a57be..55ec80f6d 100644 --- a/build_tools/autogen_ltc_backend.py +++ b/build_tools/autogen_ltc_backend.py @@ -24,8 +24,9 @@ from torchgen.gen import get_grouped_native_functions, parse_native_yaml from torchgen.gen_backend_stubs import parse_backend_yaml TORCH_DIR = Path(importlib.util.find_spec("torch").origin).resolve().parent.parent -if TORCH_DIR.joinpath("torch", "include").is_dir(): - TORCH_DIR = TORCH_DIR.joinpath("torch", "include") +TORCH_INCLUDE_DIR = TORCH_DIR.joinpath("torch", "include") +if not TORCH_INCLUDE_DIR.is_dir(): + TORCH_INCLUDE_DIR = TORCH_DIR TORCHGEN_DIR = Path(torchgen.__path__[0]).resolve() TORCH_MLIR_DIR = Path(__file__).resolve().parent.parent @@ -167,6 +168,9 @@ class GenTorchMlirLTC: ts_native_yaml = None if ts_native_yaml_path.exists(): ts_native_yaml = yaml.load(ts_native_yaml_path.read_text(), yaml.CLoader) + else: + logging.warning(f"Could not find `ts_native_functions.yaml` at {ts_native_yaml_path}") + parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path) self.native_functions = parsed_yaml.native_functions @@ -290,6 +294,7 @@ class GenTorchMlirLTC: if ts_native_yaml: ts_full_codegen = set(ts_native_yaml["full_codegen"]) + ts_supported = set(ts_native_yaml["supported"]) mlir_full_codegen = set(self.ops) if ts_full_codegen - mlir_full_codegen: @@ -308,6 +313,22 @@ class GenTorchMlirLTC: ) ) + if ts_supported - supported: + logging.debug( + "Ops supported by the TorchScript backend " + "but not by the Torch-MLIR backend:\n {}".format( + "\n ".join(sorted(ts_supported - supported)) + ) + ) + + if supported - ts_supported: + logging.debug( + "Ops supported by the Torch-MLIR backend " + "but not by the TorchScript backend:\n {}".format( + "\n ".join(sorted(supported - ts_supported)) + ) + ) + def generate_shape_inference(self): parsed_backend_yaml = parse_backend_yaml( self.source_yaml, @@ -367,7 +388,7 @@ class GenTorchMlirLTC: ) assert len(shape_inference_decls) > 0 upstream_shape_inference_decls = extract_signatures( - TORCH_DIR.joinpath( + TORCH_INCLUDE_DIR.joinpath( "torch", "csrc", "lazy", "core", "shape_inference.h" ).read_text() ) diff --git a/build_tools/autogen_ltc_backend.yaml b/build_tools/autogen_ltc_backend.yaml index 505701eaa..dde1f0f01 100644 --- a/build_tools/autogen_ltc_backend.yaml +++ b/build_tools/autogen_ltc_backend.yaml @@ -19,6 +19,7 @@ blacklist: - new_empty - rsub - slice.Tensor # Disabled in favour of slice_copy.Tensor +- zeros # Disabled in favour of functionalized alternatives - _reshape_alias @@ -59,14 +60,12 @@ supported: # but their implementations call view operators (which we need to functionalize away). - block_diag - new_empty_strided -- narrow_copy - pixel_shuffle - pixel_unshuffle - select_backward - slice_backward - diagonal_backward - _trilinear -- linalg_inv_ex - linalg_pinv.atol_rtol_tensor - logsumexp.out diff --git a/e2e_testing/torchscript/xfail_sets.py b/e2e_testing/torchscript/xfail_sets.py index 8904d7faa..f106c9749 100644 --- a/e2e_testing/torchscript/xfail_sets.py +++ b/e2e_testing/torchscript/xfail_sets.py @@ -435,12 +435,6 @@ LTC_XFAIL_SET = { "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 ec72c37ab..32cba4fdf 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,12 @@ at::Tensor LazyNativeFunctions::_to_copy( }; at::Tensor LazyNativeFunctions::empty( - at::IntArrayRef size, c10::optional dtype, + at::SymIntArrayRef sym_size, c10::optional dtype, c10::optional layout, c10::optional device, c10::optional pin_memory, c10::optional 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)) @@ -331,7 +333,9 @@ at::Tensor LazyNativeFunctions::empty_strided( c10::optional dtype, c10::optional layout, c10::optional device, c10::optional 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 +354,7 @@ 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. @@ -385,11 +389,6 @@ at::Tensor LazyNativeFunctions::new_empty_strided( self, size, stride, dtype, layout, device, pin_memory); } -at::Tensor LazyNativeFunctions::narrow_copy( - const at::Tensor& self, int64_t dim, int64_t start, int64_t length) { - return at::functionalization::functionalize_aten_op::call(self, dim, start, length); -} at::Tensor LazyNativeFunctions::pixel_shuffle( const at::Tensor& self, int64_t upscale_factor) { return at::functionalization::functionalize_aten_op:: call(i1, i2, i3, expand1, expand2, expand3, sumdim, unroll_dim); } -::std::tuple -LazyNativeFunctions::linalg_inv_ex(const at::Tensor& self, bool check_errors) { - return at::functionalization::functionalize_aten_op::call(self, check_errors); -} at::Tensor LazyNativeFunctions::linalg_pinv( const at::Tensor& self, const c10::optional& atol, const c10::optional& rtol, bool hermitian) {