diff --git a/.gitignore b/.gitignore index 0fb68ac7e..29de49a46 100644 --- a/.gitignore +++ b/.gitignore @@ -23,15 +23,10 @@ __pycache__ # Bazel bazel-* -# Libraries -*.so -*.a - # Autogenerated files /generated_native_functions.yaml /generated_backend.hash -/python/torch_mlir/csrc/base_lazy_backend/LazyIr.h -/python/torch_mlir/csrc/base_lazy_backend/LazyNativeFunctions.cpp -/python/torch_mlir/csrc/base_lazy_backend/LazyNativeFunctions.h -/python/torch_mlir/csrc/base_lazy_backend/GenLazyShapeInference.cpp -/python/torch_mlir/csrc/base_lazy_backend/RegisterLazy.cpp +/python/torch_mlir/csrc/base_lazy_backend/generated + +# Example backend +examples/ltc_backend/ltc_backend/_EXAMPLE_MLIR_BACKEND.cpython-37m-x86_64-linux-gnu.so diff --git a/build_tools/autogen_ltc_backend.py b/build_tools/autogen_ltc_backend.py index 623cf476d..9755790e7 100644 --- a/build_tools/autogen_ltc_backend.py +++ b/build_tools/autogen_ltc_backend.py @@ -12,20 +12,20 @@ from textwrap import dedent import yaml TORCH_MLIR_DIR = Path(__file__).parent.parent.resolve() -TORCH_DIR = TORCH_MLIR_DIR.parent.joinpath("pytorch") +TORCH_DIR = TORCH_MLIR_DIR.joinpath("externals", "pytorch") -sys.path.append(str(TORCH_DIR.joinpath("tools"))) +sys.path.append(str(TORCH_DIR)) # PyTorch's LTC backend autogen script -import codegen.dest.lazy_ir -import codegen.gen_lazy_tensor -from codegen.api.lazy import LazyIrSchema -from codegen.gen import get_grouped_native_functions, parse_native_yaml -from codegen.model import NativeFunctionsGroup +import torchgen.dest.lazy_ir +import torchgen.gen_lazy_tensor +from torchgen.api.lazy import LazyIrSchema +from torchgen.gen import get_grouped_native_functions, parse_native_yaml +from torchgen.model import NativeFunctionsGroup def isOptionalCType(arg): - return str(type(arg)) == "" + return str(type(arg)) == "" def generate_native_functions( @@ -33,11 +33,11 @@ def generate_native_functions( ): print("Generating Native Functions Yaml") - native_yaml_path = TORCH_DIR.joinpath( - "aten", "src", "ATen", "native", "native_functions.yaml" - ) + native_path = TORCH_DIR.joinpath("aten", "src", "ATen", "native") + native_yaml_path = native_path.joinpath("native_functions.yaml") + tags_yaml_path = native_path.joinpath("tags.yaml") - parsed_yaml = parse_native_yaml(native_yaml_path) + parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path) native_functions = parsed_yaml.native_functions grouped_native_functions = get_grouped_native_functions(native_functions) @@ -57,6 +57,9 @@ def generate_native_functions( # primarily view ops supported = config.get("supported", []) + # List of non-native ops to do IR codegen for + non_native = config.get("non_native", []) + if which("rg") is not None: # use ripgrep if available as its much faster cmd = ["rg", "-o", "-N", r"aten::[0-9a-zA-Z_\.]+"] else: @@ -105,6 +108,7 @@ def generate_native_functions( "cpp_namespace": "torch::lazy", "full_codegen": opnames, "supported": sorted(supported_ops), + "non_native": non_native, }, f, default_flow_style=False, @@ -123,13 +127,13 @@ def generate_native_functions( @dataclass(frozen=True) -class MlirLazyIr(codegen.gen_lazy_tensor.dest.LazyIR): +class GenMlirLazyIr(torchgen.dest.GenLazyIR): - def lowering_function(self, f): - func = ( - f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func - ) - schema = LazyIrSchema(func) + def lowering_function(self, schema, declaration_only=True): + signature = "TorchMlirOpVector Lower(TorchMlirFunction function, TorchMlirLoweringContext* loctx) const override" + + if declaration_only: + return f"{signature};" emplace_arguments = [] for arg in schema.positional_args: @@ -149,7 +153,7 @@ class MlirLazyIr(codegen.gen_lazy_tensor.dest.LazyIR): [f"kwarguments.emplace_back({a});" for a in emplace_kwarg_values + emplace_kwarg_scalars]) return f""" - TorchMlirOpVector Lower(TorchMlirFunction function, TorchMlirLoweringContext* loctx) const override {{ + {signature} {{ PRINT_FUNCTION(); std::vector arguments; std::vector kwarguments; @@ -159,7 +163,7 @@ class MlirLazyIr(codegen.gen_lazy_tensor.dest.LazyIR): {emplace_arguments_str} {emplace_kwarguments} torch::lazy::TorchMlirOpVector {schema.aten_name}_out = torch::lazy::LowerTorchMlirBuiltin(function, op().op, shapes(), arguments, kwarguments); - CHECK_EQ({schema.aten_name}_out.size(), {len(func.returns)}); + CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)}); return {schema.aten_name}_out; }} @@ -178,13 +182,13 @@ def generate_backend( def gen_fallback_code(*args, **kwargs): return "" - codegen.dest.lazy_ir.gen_fallback_code = gen_fallback_code + torchgen.dest.lazy_ir.gen_fallback_code = gen_fallback_code - codegen.gen_lazy_tensor.run_gen_lazy_tensor( + torchgen.gen_lazy_tensor.run_gen_lazy_tensor( backend_name="TorchMlir", aten_path=str(TORCH_DIR.joinpath("aten", "src", "ATen")), source_yaml=str(source_yaml), - output_dir=str(backend_path), + output_dir=str(backend_path.joinpath("generated")), dry_run=False, impl_path=str(backend_path.joinpath("mlir_native_functions.cpp")), node_base="torch::lazy::TorchMlirNode", @@ -192,7 +196,7 @@ def generate_backend( tensor_class="torch::lazy::LazyTensor", tensor_class_hdr="torch/csrc/lazy/core/tensor.h", shape_inference_hdr=str(backend_path.joinpath("LazyShapeInference.h")), - lazy_ir_cls=MlirLazyIr, + lazy_ir_generator=GenMlirLazyIr, ) # Remove lazy_tensor_core imports @@ -201,7 +205,7 @@ def generate_backend( "sed", "-i", "/lazy_tensor_core/d", - str(backend_path.joinpath("LazyNativeFunctions.cpp")), + str(backend_path.joinpath("generated", "LazyNativeFunctions.cpp")), ] ) @@ -240,14 +244,14 @@ def generate_backend( - shape_inference_defs ) if missing_defs: - backend_path.joinpath("GenLazyShapeInference.cpp").write_text( + backend_path.joinpath("generated", "GenLazyShapeInference.cpp").write_text( dedent( """ // This file contains autogenerated Lazy Shape Inference placeholders // for ops that dont have a corresponding structured kernel or shape definition - #include "LazyShapeInference.h" - #include "../utils/exception.h" + #include "../LazyShapeInference.h" + #include "../../utils/exception.h" namespace torch {{ namespace lazy {{ {} diff --git a/build_tools/autogen_ltc_backend.yaml b/build_tools/autogen_ltc_backend.yaml index f7298748a..d3f04dab4 100644 --- a/build_tools/autogen_ltc_backend.yaml +++ b/build_tools/autogen_ltc_backend.yaml @@ -38,8 +38,7 @@ supported: - empty - expand - fill_ -- native_batch_norm -# - native_batch_norm_backward +- native_batch_norm_backward - permute - squeeze - t @@ -50,4 +49,39 @@ additional_ops: # Additional ops to support that are not supported by Torch-MLIR explicitly - _copy_from - _copy_from_and_resize -- native_batch_norm_backward +# - native_batch_norm_backward + +# List of non native ops that we only want to do IR node class generation for +non_native: + - func: device_data(std::shared_ptr data) -> Tensor + opkind: ltc_device_data + cache_shape: false + - func: scalar(at::Scalar value, at::ScalarType type) -> Tensor + opkind: at::prim::Constant + cache_shape: false + - func: expand(Tensor input, std::vector size, bool is_scalar_expand) -> Tensor + - func: view(Tensor input, std::vector output_size) -> Tensor + cache_shape: false + - func: cast(Tensor input, at::ScalarType dtype, optional stype) -> Tensor + opkind: ltc_cast + cache_shape: false + + # View ops only required until proper functionalization pass is introduced into LTC + - func: as_strided_view_update(Tensor target, Tensor input, std::vector size, std::vector stride, int64_t storage_offset) -> Tensor + opkind: ltc_as_strided_view_update + - func: as_strided(Tensor input, std::vector size, std::vector stride, int64_t storage_offset) -> Tensor + - func: diagonal_view_update(Tensor target, Tensor input, int64_t offset, int64_t dim1, int64_t dim2) -> Tensor + opkind: ltc_diagonal_view_update + cache_shape: false + - func: diagonal(Tensor input, int64_t offset, int64_t dim1, int64_t dim2) -> Tensor + - func: narrow_view_update(Tensor input, Tensor source, std::vector base_indices) -> Tensor + opkind: ltc_narrow_view_update + - func: narrow(Tensor input, std::vector base_indices, std::vector sizes) -> Tensor + - func: permute(Tensor input, std::vector dims) -> Tensor + - func: resize(Tensor input, std::vector size) -> Tensor + - func: select_view_update(Tensor target, Tensor source, int64_t dim, int64_t start, int64_t end, int64_t stride) -> Tensor + opkind: ltc_select_view_update + cache_shape: false + - func: select(Tensor input, int64_t dim, int64_t start, int64_t end, int64_t stride) -> Tensor + - func: squeeze(Tensor input, int dim) -> Tensor + - func: unsqueeze(Tensor input, int dim) -> Tensor diff --git a/e2e_testing/torchscript/main.pyc b/e2e_testing/torchscript/main.pyc new file mode 100644 index 000000000..71b020ae3 Binary files /dev/null and b/e2e_testing/torchscript/main.pyc differ diff --git a/e2e_testing/torchscript/xfail_sets.pyc b/e2e_testing/torchscript/xfail_sets.pyc new file mode 100644 index 000000000..ea6c60b97 Binary files /dev/null and b/e2e_testing/torchscript/xfail_sets.pyc differ diff --git a/examples/ltc_backend/ltc_backend/csrc/backend/backend_impl.cpp b/examples/ltc_backend/ltc_backend/csrc/backend/backend_impl.cpp index c05d93207..a8172bc95 100644 --- a/examples/ltc_backend/ltc_backend/csrc/backend/backend_impl.cpp +++ b/examples/ltc_backend/ltc_backend/csrc/backend/backend_impl.cpp @@ -12,7 +12,7 @@ #include #include -#include +#include #include #include #include diff --git a/externals/pytorch b/externals/pytorch new file mode 160000 index 000000000..9f3d6a00a --- /dev/null +++ b/externals/pytorch @@ -0,0 +1 @@ +Subproject commit 9f3d6a00a76c567d7c046eabc60ae7a578f7bbde diff --git a/python/test/eager_mode/framework.pyc b/python/test/eager_mode/framework.pyc new file mode 100644 index 000000000..9dfe995ef Binary files /dev/null and b/python/test/eager_mode/framework.pyc differ diff --git a/python/torch_mlir/csrc/CMakeLists.txt b/python/torch_mlir/csrc/CMakeLists.txt index f42248cbd..2a20ce93c 100644 --- a/python/torch_mlir/csrc/CMakeLists.txt +++ b/python/torch_mlir/csrc/CMakeLists.txt @@ -21,14 +21,14 @@ link_directories("${TORCH_INSTALL_PREFIX}/lib") add_library(torch_mlir_ltc_backend SHARED base_lazy_backend/backend_impl.cpp - base_lazy_backend/LazyNativeFunctions.cpp + base_lazy_backend/generated/LazyNativeFunctions.cpp + base_lazy_backend/generated/GenLazyShapeInference.cpp + base_lazy_backend/generated/RegisterLazy.cpp base_lazy_backend/LazyShapeInference.cpp - base_lazy_backend/GenLazyShapeInference.cpp base_lazy_backend/mlir_lowering_context.cpp base_lazy_backend/mlir_native_functions.cpp base_lazy_backend/mlir_node.cpp base_lazy_backend/mlir_node_lowering.cpp - base_lazy_backend/RegisterLazy.cpp ) add_dependencies(torch_mlir_ltc_backend diff --git a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp b/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp index 33a065177..f36692297 100644 --- a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp +++ b/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.cpp @@ -13,5 +13,25 @@ namespace torch { namespace lazy { +std::vector compute_shape_native_batch_norm( + const at::Tensor& input, const c10::optional& weight, + const c10::optional& bias, + const c10::optional& running_mean, + const c10::optional& running_var, bool training, + double momentum, double eps) { + std::vector shapes; + shapes.reserve(3); + shapes.emplace_back(input.scalar_type(), input.sizes().vec()); + if (running_mean.has_value()) { + shapes.emplace_back( + running_mean.value().scalar_type(), running_mean.value().sizes().vec()); + if (running_var.has_value()) { + shapes.emplace_back( + running_var.value().scalar_type(), running_var.value().sizes().vec()); + } + } + return shapes; +} + } // namespace lazy } // namespace torch diff --git a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h b/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h index 743d406f5..bbcafa1b2 100644 --- a/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h +++ b/python/torch_mlir/csrc/base_lazy_backend/LazyShapeInference.h @@ -23,6 +23,7 @@ namespace lazy { // clang-format off TORCH_API std::vector compute_shape___and__(const at::Tensor & self, const at::Tensor & other); +TORCH_API std::vector compute_shape__reshape_alias(const at::Tensor & self, at::IntArrayRef size, at::IntArrayRef stride); TORCH_API std::vector compute_shape__shape_as_tensor(const at::Tensor & self); TORCH_API std::vector compute_shape__unsafe_view(const at::Tensor & self, at::IntArrayRef size); TORCH_API std::vector compute_shape_abs(const at::Tensor & self); @@ -37,6 +38,8 @@ TORCH_API std::vector compute_shape_bincount(const at::Tensor & self, con TORCH_API std::vector compute_shape_broadcast_to(const at::Tensor & self, at::IntArrayRef size); TORCH_API std::vector compute_shape_bucketize(const at::Tensor & self, const at::Tensor & boundaries, bool out_int32, bool right); TORCH_API std::vector compute_shape_constant_pad_nd(const at::Tensor & self, at::IntArrayRef pad, const at::Scalar & value); +TORCH_API std::vector compute_shape_convolution(const at::Tensor & input, const at::Tensor & weight, const c10::optional & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups); +TORCH_API std::vector compute_shape_convolution_overrideable(const at::Tensor & input, const at::Tensor & weight, const c10::optional & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool transposed, at::IntArrayRef output_padding, int64_t groups); TORCH_API std::vector compute_shape_conv2d(const at::Tensor & input, const at::Tensor & weight, const c10::optional & bias, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, int64_t groups); TORCH_API std::vector compute_shape_div(const at::Tensor & self, const at::Scalar & other); TORCH_API std::vector compute_shape_div_(at::Tensor & self, const at::Scalar & other); @@ -63,6 +66,7 @@ TORCH_API std::vector compute_shape_mean(const at::Tensor & self, c10::op TORCH_API std::vector compute_shape_mul(const at::Tensor & self, const at::Scalar & other); TORCH_API std::vector compute_shape_mul_(at::Tensor & self, const at::Scalar & other); TORCH_API std::vector compute_shape_native_layer_norm(const at::Tensor & input, at::IntArrayRef normalized_shape, const c10::optional & weight, const c10::optional & bias, double eps); +TORCH_API std::vector compute_shape_new_empty(const at::Tensor & self, at::IntArrayRef size, c10::optional dtype, c10::optional layout, c10::optional device, c10::optional pin_memory); TORCH_API std::vector compute_shape_new_ones(const at::Tensor & self, at::IntArrayRef size, c10::optional dtype, c10::optional layout, c10::optional device, c10::optional pin_memory); TORCH_API std::vector compute_shape_new_zeros(const at::Tensor & self, at::IntArrayRef size, c10::optional dtype, c10::optional layout, c10::optional device, c10::optional pin_memory); TORCH_API std::vector compute_shape_rand_like(const at::Tensor & self, c10::optional dtype, c10::optional layout, c10::optional device, c10::optional pin_memory, c10::optional memory_format); @@ -82,6 +86,9 @@ TORCH_API std::vector compute_shape_sum(const at::Tensor & self, c10::opt TORCH_API std::vector compute_shape_transpose(const at::Tensor & self, int64_t dim0, int64_t dim1); TORCH_API std::vector compute_shape_type_as(const at::Tensor & self, const at::Tensor & other); TORCH_API std::vector compute_shape_var(const at::Tensor & self, bool unbiased); +TORCH_API std::vector compute_shape_zero_(at::Tensor & self); + +TORCH_API std::vector compute_shape_native_batch_norm(const at::Tensor & input, const c10::optional & weight, const c10::optional & bias, const c10::optional & running_mean, const c10::optional & running_var, bool training, double momentum, double eps); // clang-format on diff --git a/python/torch_mlir/csrc/base_lazy_backend/backend_impl.cpp b/python/torch_mlir/csrc/base_lazy_backend/backend_impl.cpp index 8a8b121e2..e71b2f0e2 100644 --- a/python/torch_mlir/csrc/base_lazy_backend/backend_impl.cpp +++ b/python/torch_mlir/csrc/base_lazy_backend/backend_impl.cpp @@ -18,6 +18,7 @@ #include "../utils/debug.h" #include "../utils/exception.h" #include "backend_impl.h" +#include "ir_builder.h" #include "mlir_lowering_context.h" namespace torch { @@ -72,6 +73,15 @@ TorchMlirBackendData::Info* TorchMlirBackendData::mlir_info() const { * */ void TorchMlirBackendImpl::PrepareToExit() const {} +/** + * IR Tracing + * */ + +const IrBuilder* TorchMlirBackendImpl::GetIrBuilder() const { + static const IrBuilder* builder = new TorchMlirIrBuilder(); + return builder; +} + /** * Data Transfer * */ @@ -95,6 +105,16 @@ BackendDataPtr TorchMlirBackendImpl::CreateDataPlaceholder( return std::make_shared(device, shape); } +BackendDataPtr +TorchMlirBackendImpl::GetComputationDataFromNode(Node* node) const { + PRINT_FUNCTION(); + auto* device_data_node = dynamic_cast(node); + if (!device_data_node) { + return nullptr; + } + return device_data_node->data; +} + at::Tensor TorchMlirBackendImpl::MakeTensorFromComputationData( const BackendDataPtr data, c10::optional logical_scalar_type) const { diff --git a/python/torch_mlir/csrc/base_lazy_backend/backend_impl.h b/python/torch_mlir/csrc/base_lazy_backend/backend_impl.h index 685c6a799..cae7266f2 100644 --- a/python/torch_mlir/csrc/base_lazy_backend/backend_impl.h +++ b/python/torch_mlir/csrc/base_lazy_backend/backend_impl.h @@ -65,6 +65,12 @@ public: * */ virtual void PrepareToExit() const override; + /** + * IR Tracing + * */ + + const IrBuilder* GetIrBuilder() const override; + /** * Configuration * */ @@ -84,6 +90,10 @@ public: virtual BackendDataPtr CreateDataPlaceholder( const BackendDevice& device, const Shape& shape) const override; + // Gets backend data if the node is a device data node. Otherwise returns + // nullptr. + virtual BackendDataPtr GetComputationDataFromNode(Node*) const override; + virtual at::Tensor MakeTensorFromComputationData( const BackendDataPtr data, c10::optional logical_scalar_type) const override; diff --git a/python/torch_mlir/csrc/base_lazy_backend/dynamic_ir.cpp b/python/torch_mlir/csrc/base_lazy_backend/dynamic_ir.cpp new file mode 100644 index 000000000..b1de53c60 --- /dev/null +++ b/python/torch_mlir/csrc/base_lazy_backend/dynamic_ir.cpp @@ -0,0 +1,74 @@ +//===- dynamic_ir.cpp -----------------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// +// This file is adapted from pytorch/pytorch +// https://github.com/pytorch/pytorch/blob/master/torch/csrc/lazy/ts_backend/dynamic_ir.cpp +//===----------------------------------------------------------------------===// + +#include + +namespace torch { +namespace lazy { + +DimensionNode::DimensionNode(OpKind op, OpList operands, hash_t hash_seed) + : TorchMlirNode( + op, operands, /*num_outputs=*/1, + /* hash_seed */ HashCombine(op.hash(), hash_seed)) {} + +std::string DimensionNode::ToString() const { return "DimensionNode"; } + +SizeNode::SizeNode(Value input, size_t dim) + : DimensionNode( + OpKind{c10::Symbol::fromQualString("aten::size")}, {input}, + MHash(dim)), + dim_(dim){}; + +int64_t SizeNode::getStaticValue() const { + return dynamic_cast(operand(0).node) + ->shape(0) + .size(dim_); +} + +std::string SizeNode::ToString() const { return "SizeNode"; } + +SizeAdd::SizeAdd(Value a, Value b) + : DimensionNode(OpKind{c10::Symbol::fromQualString("aten::add")}, {a, b}){}; + +int64_t SizeAdd::getStaticValue() const { + return dynamic_cast(operand(0).node)->getStaticValue() + + dynamic_cast(operand(1).node)->getStaticValue(); +} + +std::string SizeAdd::ToString() const { return "SizeAdd"; } + +SizeMul::SizeMul(Value a, Value b) + : DimensionNode(OpKind{c10::Symbol::fromQualString("aten::mul")}, {a, b}){}; + +int64_t SizeMul::getStaticValue() const { + return dynamic_cast(operand(0).node)->getStaticValue() * + dynamic_cast(operand(1).node)->getStaticValue(); +} + +std::string SizeMul::ToString() const { return "SizeMul"; } + +SizeDiv::SizeDiv(Value a, Value b) + : DimensionNode(OpKind{c10::Symbol::fromQualString("aten::div")}, {a, b}){}; + +int64_t SizeDiv::getStaticValue() const { + TORCH_CHECK( + dynamic_cast(operand(1).node)->getStaticValue() != + 0, + "Can't divide a dimension by zero"); + return dynamic_cast(operand(0).node)->getStaticValue() / + dynamic_cast(operand(1).node)->getStaticValue(); +} + +std::string SizeDiv::ToString() const { return "SizeDiv"; } + +} // namespace lazy +} // namespace torch diff --git a/python/torch_mlir/csrc/base_lazy_backend/dynamic_ir.h b/python/torch_mlir/csrc/base_lazy_backend/dynamic_ir.h new file mode 100644 index 000000000..97bad3c15 --- /dev/null +++ b/python/torch_mlir/csrc/base_lazy_backend/dynamic_ir.h @@ -0,0 +1,99 @@ +//===- dynamic_ir.h -------------------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// +// This file is adapted from pytorch/pytorch +// https://github.com/pytorch/pytorch/blob/master/torch/csrc/lazy/ts_backend/dynamic_ir.h +//===----------------------------------------------------------------------===// + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mlir_node.h" +#include +#include +#include +#include +#include + +C10_DECLARE_bool(ltc_enable_dynamic_shapes); + +namespace torch { +namespace lazy { + +/** + * The goal of "dynamic" Nodes is to patch a hole in our tracing. + * Previously, if a user called `sizes` on a Tensor, it would leak out + * of our tracing system, as `sizes` returns a torch.Size or an int. To + * prevent this from happening, we introduce DimensionNode, a new type + * of Node that abstracts the operation of getting the dimensions of a + * Tensor. + * + * Consider the following example: + * ``` + * numel = x.shape()[0] * x.shape()[1] + * ``` + * + * Here, `x.shape()[i]` will be a SizeNode (subclass of DimensionNode), + * and the multiplication of the two SizeNodes will be represented by + * a SizeMul (also a subclass of DimensionNode). Through this, we can + * prevent `numel` from being represented as a Python int and thus + * burned into the Graph. + */ + +class TORCH_API DimensionNode : public lazy::TorchMlirNode { +public: + DimensionNode(OpKind op, OpList operands, hash_t hash_seed = kHashSeed); + bool isDynamic() { return false; } + + std::string ToString() const override; + + virtual int64_t getStaticValue() const = 0; +}; + +// Represents the result of calling `size` on a Tensor +class TORCH_API SizeNode : public DimensionNode { +public: + SizeNode(Value input, size_t dim); + int64_t getStaticValue() const override; + std::string ToString() const override; + size_t dim_ = 0; +}; + +class TORCH_API SizeAdd : public DimensionNode { +public: + SizeAdd(Value a, Value b); + int64_t getStaticValue() const override; + std::string ToString() const override; +}; + +class TORCH_API SizeMul : public DimensionNode { +public: + SizeMul(Value a, Value b); + int64_t getStaticValue() const override; + std::string ToString() const override; +}; + +class TORCH_API SizeDiv : public DimensionNode { +public: + SizeDiv(Value a, Value b); + int64_t getStaticValue() const override; + std::string ToString() const override; +}; + +} // namespace lazy +} // namespace torch diff --git a/python/torch_mlir/csrc/base_lazy_backend/ir_builder.h b/python/torch_mlir/csrc/base_lazy_backend/ir_builder.h new file mode 100644 index 000000000..96e8e7050 --- /dev/null +++ b/python/torch_mlir/csrc/base_lazy_backend/ir_builder.h @@ -0,0 +1,65 @@ +//===- ir_builder.h -------------------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// Also available under a BSD-style license. See LICENSE. +// +//===----------------------------------------------------------------------===// +// This file is adapted from pytorch/pytorch +// https://github.com/pytorch/pytorch/blob/master/torch/csrc/lazy/ts_backend/ir_builder.h +//===----------------------------------------------------------------------===// + +#pragma once + +#include +#include +#include +#include + +#include "dynamic_ir.h" +#include "generated/LazyNonNativeIr.h" +#include "mlir_node.h" +#include "ops/generic.h" + +// This file contains the TorchMlir IrBuilder + +namespace torch { +namespace lazy { + +// clang-format off + +struct TorchMlirIrBuilder : IrBuilder { + NodePtr MakeDeviceData(const std::shared_ptr& data) const override { return MakeNode(data); } + NodePtr MakeScalar(const at::Scalar& value, const at::ScalarType& type) const override { return MakeNode(value, type); } + NodePtr MakeExpand(const Value& input0, const std::vector& size, const bool& is_scalar_expand) const override { return MakeNode(input0, size, is_scalar_expand); } + NodePtr MakeView(const Value& input0, const std::vector& output_size) const override { return MakeNode(input0, output_size); } + NodePtr MakeCast(const Value& input0, const at::ScalarType& dtype, const c10::optional& stype = c10::nullopt) const override { return MakeNode(input0, dtype, stype); } + NodePtr MakeTensorList(const OpList& inputs) const override { return MakeNode(inputs); } + NodePtr MakeGeneric(const OpKind& op, const OpList& operands, const Shape& shape, const size_t& num_outputs = 1, const hash_t& hash_seed = static_cast(0x5a2d296e9)) const override { return MakeNode(op, operands, shape, num_outputs, hash_seed); } + + // view ops + NodePtr MakeAsStridedViewUpdate(const Value& input0, const Value& input1, const std::vector& size, const std::vector& stride, const int64_t& storage_offset) const override { return MakeNode(input0, input1, size, stride, storage_offset); } + NodePtr MakeAsStrided(const Value& input0, const std::vector& size, const std::vector& stride, const int64_t& storage_offset) const override { return MakeNode(input0, size, stride, storage_offset); } + NodePtr MakeDiagonalViewUpdate(const Value& input0, const Value& input1, const int64_t& offset, const int64_t& dim1, const int64_t& dim2) const override { return MakeNode(input0, input1, offset, dim1, dim2); } + NodePtr MakeDiagonal(const Value& input0, const int64_t& offset, const int64_t& dim1, const int64_t& dim2) const override { return MakeNode(input0, offset, dim1, dim2); } + NodePtr MakeNarrowViewUpdate(const Value& input0, const Value& input1, const std::vector& base_indices) const override { return MakeNode(input0, input1, base_indices); } + NodePtr MakeNarrow(const Value& input0, const std::vector& base_indices, const std::vector& sizes) const override { return MakeNode(input0, base_indices, sizes); } + NodePtr MakePermute(const Value& input0, const std::vector& dims) const override { return MakeNode(input0, dims); } + NodePtr MakeResize(const Value& input0, const std::vector& size) const override { return MakeNode(input0, size); } + NodePtr MakeSelectViewUpdate(const Value& input0, const Value& input1, const int64_t& dim, const int64_t& start, const int64_t& end, const int64_t& stride) const override { return MakeNode(input0, input1, dim, start, end, stride); } + NodePtr MakeSelect(const Value& input0, const int64_t& dim, const int64_t& start, const int64_t& end, const int64_t& stride) const override { return MakeNode