diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py index 371d3fae131fe2..72c40148cff5b9 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py @@ -112,46 +112,48 @@ def _call(*args): def fx_openvino(subgraph, example_inputs): - try: - executor_parameters = None - inputs_reversed = False - if os.getenv("OPENVINO_TORCH_MODEL_CACHING") is not None: - # Create a hash to be used for caching - model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest() - executor_parameters = {"model_hash_str": model_hash_str} - # Check if the model was fully supported and already cached - example_inputs.reverse() - inputs_reversed = True - maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", get_device(), example_inputs, cache_root_path()) - if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"): - # Model is fully supported and already cached. Run the cached OV model directly. - compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, *example_inputs) - def _call(*args): - res = execute_cached(compiled_model, *args) - return res - return _call - if inputs_reversed: - example_inputs.reverse() - model = make_fx(subgraph)(*example_inputs) - with torch.no_grad(): - model.eval() - partitioner = Partitioner() - compiled_model = partitioner.make_partitions(model) - - if executor_parameters is not None and 'model_hash_str' in executor_parameters: - # Check if the model is fully supported. - fully_supported = partitioner.check_fully_supported(compiled_model) - if fully_supported: - executor_parameters["model_hash_str"] += "_fs" - - def _call(*args): - res = execute(compiled_model, *args, executor="openvino", - executor_parameters=executor_parameters) - return res - return _call - except Exception as e: - log.debug(f"Failed in OpenVINO execution: {e}") - return compile_fx(subgraph, example_inputs) + #try: + print("DEBUG - fx_openvino - A") + executor_parameters = None + inputs_reversed = False + if os.getenv("OPENVINO_TORCH_MODEL_CACHING") is not None: + # Create a hash to be used for caching + model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest() + executor_parameters = {"model_hash_str": model_hash_str} + # Check if the model was fully supported and already cached + example_inputs.reverse() + inputs_reversed = True + maybe_fs_cached_name = cached_model_name(model_hash_str + "_fs", get_device(), example_inputs, cache_root_path()) + if os.path.isfile(maybe_fs_cached_name + ".xml") and os.path.isfile(maybe_fs_cached_name + ".bin"): + # Model is fully supported and already cached. Run the cached OV model directly. + compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, *example_inputs) + def _call(*args): + res = execute_cached(compiled_model, *args) + return res + return _call + if inputs_reversed: + example_inputs.reverse() + model = make_fx(subgraph)(*example_inputs) + with torch.no_grad(): + model.eval() + partitioner = Partitioner() + compiled_model = partitioner.make_partitions(model) + + if executor_parameters is not None and 'model_hash_str' in executor_parameters: + # Check if the model is fully supported. + fully_supported = partitioner.check_fully_supported(compiled_model) + if fully_supported: + executor_parameters["model_hash_str"] += "_fs" + + def _call(*args): + print("DEBUG - fx_openvino - B") + res = execute(compiled_model, *args, executor="openvino", + executor_parameters=executor_parameters) + return res + return _call + #except Exception as e: + # log.debug(f"Failed in OpenVINO execution: {e}") + # return compile_fx(subgraph, example_inputs) def reset(): clear_caches() diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/compile.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/compile.py index df8addbf68d44a..9ea4d6739bcb1d 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/compile.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/compile.py @@ -34,9 +34,15 @@ def cached_model_name(model_hash_str, device, args, cache_root, reversed = False inputs_str = "" for idx, input_data in enumerate(args): if reversed: - inputs_str = "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str + if isinstance(input_data, torch.SymInt): + inputs_str = "_"+str(type(input_data)) + inputs_str + else: + inputs_str = "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + inputs_str else: - inputs_str += "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") + if isinstance(input_data, torch.SymInt): + inputs_str += "_"+str(type(input_data)) + else: + inputs_str += "_" + str(input_data.type()) + str(input_data.size())[11:-1].replace(" ", "") inputs_str = sha256(inputs_str.encode('utf-8')).hexdigest() file_name += inputs_str @@ -101,7 +107,46 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None): input_types = [] for idx, input_data in enumerate(args): input_types.append(input_data.type()) - input_shapes.append(input_data.size()) + if input_data.size() == torch.Size([17, 1, 2, 128]): + input_shapes.append(torch.Size([-1, 1, 2, 128])) + elif input_data.size() == torch.Size([17, 1, 4096]): + input_shapes.append(torch.Size([-1, 1, 4096])) + elif input_data.size() == torch.Size([17, 1, 32, 2]): + input_shapes.append(torch.Size([-1, 1, 32, 2])) + elif input_data.size() == torch.Size([18, 1, 2, 128]): + input_shapes.append(torch.Size([-1, 1, 2, 128])) + elif input_data.size() == torch.Size([18, 1, 4096]): + input_shapes.append(torch.Size([-1, 1, 4096])) + elif input_data.size() == torch.Size([18, 1, 32, 2]): + input_shapes.append(torch.Size([-1, 1, 32, 2])) + elif input_data.size() == torch.Size([48, 1, 2, 128]): + input_shapes.append(torch.Size([-1, 1, 2, 128])) + elif input_data.size() == torch.Size([48, 1, 4096]): + input_shapes.append(torch.Size([-1, 1, 4096])) + elif input_data.size() == torch.Size([48, 1, 32, 2]): + input_shapes.append(torch.Size([-1, 1, 32, 2])) + elif input_data.size() == torch.Size([22, 1, 2, 128]): + input_shapes.append(torch.Size([-1, 1, 2, 128])) + elif input_data.size() == torch.Size([22, 1, 4096]): + input_shapes.append(torch.Size([-1, 1, 4096])) + elif input_data.size() == torch.Size([22, 1, 32, 2]): + input_shapes.append(torch.Size([-1, 1, 32, 2])) + elif input_data.size() == torch.Size([23, 1, 2, 128]): + input_shapes.append(torch.Size([-1, 1, 2, 128])) + elif input_data.size() == torch.Size([23, 1, 4096]): + input_shapes.append(torch.Size([-1, 1, 4096])) + elif input_data.size() == torch.Size([23, 1, 32, 2]): + input_shapes.append(torch.Size([-1, 1, 32, 2])) + elif input_data.size() == torch.Size([24, 1, 2, 128]): + input_shapes.append(torch.Size([-1, 1, 2, 128])) + elif input_data.size() == torch.Size([24, 1, 4096]): + input_shapes.append(torch.Size([-1, 1, 4096])) + elif input_data.size() == torch.Size([24, 1, 32, 2]): + input_shapes.append(torch.Size([-1, 1, 32, 2])) + elif input_data.size() == torch.Size([1, 32, 24, 128]): + input_shapes.append(torch.Size([1, 32, -1, 128])) + else: + input_shapes.append(input_data.size()) decoder = TorchFXPythonDecoder(gm, gm, input_shapes=input_shapes, input_types=input_types) @@ -125,7 +170,46 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None): for idx, input_data in enumerate(args): om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype]) - om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape))) + if input_data.size() == torch.Size([17, 1, 2, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 2, 128])))) + elif input_data.size() == torch.Size([17, 1, 4096]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 4096])))) + elif input_data.size() == torch.Size([17, 1, 32, 2]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 32, 2])))) + elif input_data.size() == torch.Size([18, 1, 2, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 2, 128])))) + elif input_data.size() == torch.Size([18, 1, 4096]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 4096])))) + elif input_data.size() == torch.Size([18, 1, 32, 2]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 32, 2])))) + elif input_data.size() == torch.Size([48, 1, 2, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 2, 128])))) + elif input_data.size() == torch.Size([48, 1, 4096]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 4096])))) + elif input_data.size() == torch.Size([48, 1, 32, 2]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 32, 2])))) + elif input_data.size() == torch.Size([22, 1, 2, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 2, 128])))) + elif input_data.size() == torch.Size([22, 1, 4096]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 4096])))) + elif input_data.size() == torch.Size([22, 1, 32, 2]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 32, 2])))) + elif input_data.size() == torch.Size([23, 1, 2, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 2, 128])))) + elif input_data.size() == torch.Size([23, 1, 4096]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 4096])))) + elif input_data.size() == torch.Size([23, 1, 32, 2]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 32, 2])))) + elif input_data.size() == torch.Size([24, 1, 2, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 2, 128])))) + elif input_data.size() == torch.Size([24, 1, 4096]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 4096])))) + elif input_data.size() == torch.Size([23, 1, 32, 2]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([-1, 1, 32, 2])))) + elif input_data.size() == torch.Size([1, 32, 24, 128]): + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(torch.Size([1, 32, -1, 128])))) + else: + om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape))) om.validate_nodes_and_infer_types() if model_hash_str is not None: diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py index 4b45c1b7a6af5a..f08edf0fed63c4 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py @@ -83,11 +83,20 @@ def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition if not fully_supported: model_hash_str = model_hash_str + "_p" + str(partition_id) - if use_cache and (partition_id in compiled_cache): - compiled = compiled_cache[partition_id] + #if use_cache and (partition_id in compiled_cache): + # compiled = compiled_cache[partition_id] + #else: + # compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str) + # compiled_cache[partition_id] = compiled + + cached_pid = partition_id + #if cached_pid > 1: + # cached_pid = 1 + if use_cache and (cached_pid in compiled_cache): + compiled = compiled_cache[cached_pid] else: compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str) - compiled_cache[partition_id] = compiled + compiled_cache[cached_pid] = compiled flat_args, _ = tree_flatten(args) ov_inputs = [a.detach().cpu().numpy() for a in flat_args] @@ -113,11 +122,12 @@ def __call__(self, *args): if self.perm_fallback: return self.gm(*args) - try: - result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id) - except Exception: - self.perm_fallback = True - return self.gm(*args) + result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id) + #try: + # result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id) + #except Exception: + # self.perm_fallback = True + # return self.gm(*args) return result @@ -154,11 +164,11 @@ def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=Non model_hash_str = executor_parameters.get("model_hash_str", None) signature = str(id(gm)) - for idx, input_data in enumerate(args): - if isinstance(input_data, torch.Tensor): - signature = signature + "_" + str(idx) + ":" + str(input_data.type())[6:] + ":" + str(input_data.size())[11:-1].replace(" ", "") - else: - signature = signature + "_" + str(idx) + ":" + type(input_data).__name__ + ":val(" + str(input_data) + ")" + #for idx, input_data in enumerate(args): + # if isinstance(input_data, torch.Tensor): + # signature = signature + "_" + str(idx) + ":" + str(input_data.type())[6:] + ":" + str(input_data.size())[11:-1].replace(" ", "") + # else: + # signature = signature + "_" + str(idx) + ":" + type(input_data).__name__ + ":val(" + str(input_data) + ")" if signature not in partitioned_modules: partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache, diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py index a6fb4de094d3eb..18802d2bcd5c14 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py @@ -37,12 +37,16 @@ def __init__(self): "torch.ops.aten.add.Tensor": None, "torch.ops.aten.add_.Tensor": None, "torch.ops.aten.addmm.default": None, + "torch.ops.aten.all.default": None, + "torch.ops.aten.any.default": None, "torch.ops.aten.arange.start": None, + "torch.ops.aten.arange.start_step": None, "torch.ops.aten.arange.default": None, "torch.ops.aten.argmax.default": None, "torch.ops.aten.avg_pool2d.default": None, "torch.ops.aten.baddbmm.default": None, "torch.ops.aten.bitwise_and.Tensor": None, + "torch.ops.aten.bitwise_not.default": None, "torch.ops.aten.bmm.default": None, "torch.ops.aten.cat.default": None, "torch.ops.aten.clamp_min.default": None, @@ -52,6 +56,7 @@ def __init__(self): "torch.ops.aten.cos.default": None, "torch.ops.aten.cumsum.default": None, "torch.ops.aten.detach.default": None, + "torch.ops.aten.detach_.default": None, "torch.ops.aten.div.Scalar": None, "torch.ops.aten.div.Tensor": None, "torch.ops.aten.embedding.default": None, @@ -60,6 +65,7 @@ def __init__(self): "torch.ops.aten.eq.Tensor": None, "torch.ops.aten.exp.default": None, "torch.ops.aten.expand.default": None, + "torch.ops.aten.fill_.Tensor": None, "torch.ops.aten.full.default": None, "torch.ops.aten.gather.default": None, "torch.ops.aten.gelu.default": None, @@ -68,12 +74,17 @@ def __init__(self): "torch.ops.aten.hardswish_.default": None, "torch.ops.aten.hardtanh_.default": None, "torch.ops.aten.index.Tensor": None, + "torch.ops.aten.isinf.default": None, + "torch.ops.aten.isnan.default": None, + "torch.ops.aten.le.Scalar": None, "torch.ops.aten.leaky_relu_.default": None, "torch.ops.aten.lift_fresh_copy.default": None, "torch.ops.aten.linalg_vector_norm.default": None, "torch.ops.aten.lt.Tensor": None, "torch.ops.aten.log.default": None, + "torch.ops.aten.logical_not.default": None, "torch.ops.aten.logsumexp.default": None, + "torch.ops.aten.masked_fill.Scalar": None, "torch.ops.aten.masked_fill_.Scalar": None, "torch.ops.aten.masked_fill.Tensor": None, "torch.ops.aten.max_pool2d_with_indices.default": None, @@ -88,11 +99,17 @@ def __init__(self): "torch.ops.aten.native_layer_norm.default": None, "torch.ops.aten.neg.default": None, "torch.ops.aten.new_ones.default": None, + "torch.ops.aten.ones.default": None, "torch.ops.aten.permute.default": None, + "torch.ops.aten.pow.Scalar": None, "torch.ops.aten.pow.Tensor_Scalar": None, + "torch.ops.aten.pow.Tensor_Tensor": None, + "torch.ops.aten.reciprocal.default": None, "torch.ops.aten.relu.default": None, "torch.ops.aten.relu_.default": None, + "torch.ops.aten.rsqrt.default": None, "torch.ops.aten.rsub.Scalar": None, + "torch.ops.aten.scatter.src": None, "torch.ops.aten._scaled_dot_product_flash_attention.default": None, "torch.ops.aten.select.int": None, "torch.ops.aten.sigmoid.default": None, @@ -100,12 +117,18 @@ def __init__(self): "torch.ops.aten.silu_.default": None, "torch.ops.aten.sin.default": None, "torch.ops.aten.slice.Tensor": None, + "torch.ops.aten.sort.default": None, "torch.ops.aten.split.Tensor": None, + "torch.ops.aten.split_with_sizes.default": None, + "torch.ops.aten.stack.default": None, "torch.ops.aten.sub.default": None, "torch.ops.aten.sub.Tensor": None, "torch.ops.aten.t.default": None, "torch.ops.aten.tanh.default": None, + "torch.ops.aten.topk.default": None, "torch.ops.aten.transpose.int": None, + "torch.ops.aten.tril.default": None, + "torch.ops.aten.tril_.default": None, "torch.ops.aten.unsqueeze.default": None, "torch.ops.aten.upsample_nearest2d.default": None, "torch.ops.aten.view.default": None, diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/partition.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/partition.py index cc3381783b00ee..02364c693158b9 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/partition.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/partition.py @@ -19,6 +19,7 @@ import typing as t import logging +import os logger = logging.getLogger(__name__) logger.setLevel(logging.WARNING) @@ -59,7 +60,14 @@ def make_partitions(self, graph_module: GraphModule) -> GraphModule: partitioner = CapabilityBasedPartitioner( graph_module, self.supported_ops, allows_single_node_partition=False) partitions = partitioner.propose_partitions() - self.add_get_attr_inputs(partitions) - fused_graph_module = partitioner.fuse_partitions(partitions) + new_partitions = [] + min_num_nodes = 0 + if os.getenv("OPENVINO_TORCH_MIN_NUM_NODES") is not None: + min_num_nodes = int(os.getenv("OPENVINO_TORCH_MIN_NUM_NODES")) + for part in partitions: + if len(part.nodes) > min_num_nodes: + new_partitions.append(part) + self.add_get_attr_inputs(new_partitions) + fused_graph_module = partitioner.fuse_partitions(new_partitions) return fused_graph_module diff --git a/src/frontends/pytorch/src/input_model.cpp b/src/frontends/pytorch/src/input_model.cpp index bce1937151d92d..765b4b70f3ce89 100644 --- a/src/frontends/pytorch/src/input_model.cpp +++ b/src/frontends/pytorch/src/input_model.cpp @@ -24,7 +24,7 @@ InputModel::InputModel(const std::shared_ptr& model_decoder) : m_m const auto& outputs = m_model_decoder->outputs(); for (size_t i = 0; i < outputs.size(); ++i) { auto out_place = std::make_shared(*this, outputs[i]); - m_name_to_place.emplace(std::to_string(inputs[i]), std::dynamic_pointer_cast(out_place)); + m_name_to_place.emplace(std::to_string(outputs[i]), std::dynamic_pointer_cast(out_place)); for (const auto& name : out_place->get_names()) { m_name_to_place.emplace(name, std::dynamic_pointer_cast(out_place)); } diff --git a/src/frontends/pytorch/src/op/any.cpp b/src/frontends/pytorch/src/op/any.cpp new file mode 100644 index 00000000000000..32f4664e8b3195 --- /dev/null +++ b/src/frontends/pytorch/src/op/any.cpp @@ -0,0 +1,32 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/reduce_logical_or.hpp" +#include "openvino/op/reshape.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/not_equal.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +OutputVector translate_any_fx(const NodeContext& context) { + num_inputs_check(context, 1, 1); + auto x = context.get_input(0); + auto num_inputs = context.get_input_size(); + bool keep_dims = false; + auto const_minus_one = context.mark_node(ov::op::v0::Constant::create(element::i32, Shape{1}, {-1})); + auto flatten_source = context.mark_node(std::make_shared(x, const_minus_one, false)); + auto const_zero = context.mark_node(ov::op::v0::Constant::create(element::i32, Shape{1}, {0})); + auto any = context.mark_node(std::make_shared(flatten_source, const_zero, keep_dims)); + return {any}; +}; + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op/bitwise.cpp b/src/frontends/pytorch/src/op/bitwise.cpp index 84465502969d81..e692284ff85ca4 100644 --- a/src/frontends/pytorch/src/op/bitwise.cpp +++ b/src/frontends/pytorch/src/op/bitwise.cpp @@ -5,8 +5,10 @@ #include "openvino/frontend/pytorch/node_context.hpp" #include "openvino/op/bitwise_and.hpp" #include "openvino/op/bitwise_not.hpp" +#include "openvino/op/logical_not.hpp" #include "openvino/op/bitwise_or.hpp" #include "openvino/op/bitwise_xor.hpp" +#include "openvino/op/convert.hpp" #include "utils.hpp" namespace ov { @@ -24,6 +26,24 @@ OutputVector translate_bitwise_not(const NodeContext& context) { return {not_x}; }; +OutputVector translate_bitwise_not_fx(const NodeContext& context) { + num_inputs_check(context, 1, 2); + auto x = context.get_input(0); + if (x.get_element_type() != element::boolean) { + auto x_bool = context.mark_node(std::make_shared(x, element::boolean)); + auto not_x = context.mark_node(std::make_shared(x_bool)); + if (!context.input_is_none(1)) { + context.mutate_input(1, not_x); + } + return {not_x}; + } + auto not_x = context.mark_node(std::make_shared(x)); + if (!context.input_is_none(1)) { + context.mutate_input(1, not_x); + } + return {not_x}; +}; + OutputVector translate_bitwise_and(const NodeContext& context) { num_inputs_check(context, 2, 3); auto x = context.get_input(0); diff --git a/src/frontends/pytorch/src/op/cat.cpp b/src/frontends/pytorch/src/op/cat.cpp index 9476979a118bd7..156e0ac914a073 100644 --- a/src/frontends/pytorch/src/op/cat.cpp +++ b/src/frontends/pytorch/src/op/cat.cpp @@ -5,6 +5,7 @@ #include "openvino/frontend/pytorch/node_context.hpp" #include "openvino/op/concat.hpp" #include "openvino/op/parameter.hpp" +#include "openvino/op/unsqueeze.hpp" #include "openvino/op/reshape.hpp" #include "openvino/op/scatter_elements_update.hpp" #include "openvino/op/shape_of.hpp" @@ -73,13 +74,22 @@ OutputVector translate_cat(const NodeContext& context) { OutputVector translate_cat_fx(const NodeContext& context) { // This translator is only needed to get axis as constant from external scope + // TODO: This is a temporary solution until the nested input issue is resolved num_inputs_check(context, 2, context.get_input_size()); - std::deque> list_elems; - for (size_t i = 0; i < context.get_input_size() - 1; i++) { - list_elems.push_back(context.get_input(static_cast(i))); + if (context.get_input(context.get_input_size()-1).get_shape().size() == 0) { + std::deque> list_elems; + for (size_t i = 0; i < context.get_input_size() - 1; i++) { + list_elems.push_back(context.get_input(static_cast(i))); + } + auto axis = context.const_input(context.get_input_size() - 1); + return translate_cat_common(context, list_elems, axis, true); + } else { + std::deque> list_elems; + for (size_t i = 0; i < context.get_input_size(); i++) { + list_elems.push_back(context.get_input(static_cast(i))); + } + return translate_cat_common(context, list_elems, 0, true); } - auto axis = context.const_input(context.get_input_size() - 1); - return translate_cat_common(context, list_elems, axis, true); }; OutputVector translate_quantized_cat(const NodeContext& context) { @@ -94,6 +104,18 @@ OutputVector translate_quantized_cat(const NodeContext& context) { list_elems.front())}; }; +OutputVector translate_stack_fx(const NodeContext& context) { + // This translator is only needed to get axis as constant from external scope + num_inputs_check(context, 2, context.get_input_size()); + std::deque> list_elems; + for (size_t i = 0; i < context.get_input_size() - 1; i++) { + auto in_unsqueeze = context.mark_node(std::make_shared(context.get_input(static_cast(i)), context.get_input(context.get_input_size() - 1))); + list_elems.push_back(in_unsqueeze); + } + auto axis = context.const_input(context.get_input_size() - 1); + return translate_cat_common(context, list_elems, axis, true); +}; + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op/expand.cpp b/src/frontends/pytorch/src/op/expand.cpp index 10dd8fdc9605a7..e13b5959671061 100644 --- a/src/frontends/pytorch/src/op/expand.cpp +++ b/src/frontends/pytorch/src/op/expand.cpp @@ -56,10 +56,10 @@ OutputVector translate_expand_fx(const NodeContext& context) { // TODO: This is a temporary solution to optimize out Broadcast if the input and // output shapes are same. This should be removed after a proper optimization is // implemented. - auto sizes_const = context.const_input(1); - if (x.get_shape() == sizes_const) { - return {x}; - } + //auto sizes_const = context.const_input(1); + //if (x.get_shape() == sizes_const) { + // return {x}; + //} auto sizes = context.get_input(1); // TODO: figure out what implicit means FRONT_END_OP_CONVERSION_CHECK(context.input_is_none(2) || context.const_input(2) == false, @@ -70,4 +70,4 @@ OutputVector translate_expand_fx(const NodeContext& context) { } // namespace op } // namespace pytorch } // namespace frontend -} // namespace ov \ No newline at end of file +} // namespace ov diff --git a/src/frontends/pytorch/src/op/index.cpp b/src/frontends/pytorch/src/op/index.cpp index e739476619ff7f..799ae3aa98b469 100644 --- a/src/frontends/pytorch/src/op/index.cpp +++ b/src/frontends/pytorch/src/op/index.cpp @@ -4,8 +4,9 @@ #include "openvino/frontend/pytorch/node_context.hpp" #include "openvino/op/concat.hpp" +#include "openvino/op/constant.hpp" #include "openvino/op/convert.hpp" -#include "openvino/op/gather_nd.hpp" +#include "openvino/op/gather.hpp" #include "utils.hpp" namespace ov { @@ -28,7 +29,8 @@ OutputVector translate_index_fx(const NodeContext& context) { } auto concat = context.mark_node(std::make_shared(OutputVector(list_elems.begin(), list_elems.end()), 0)); - auto gather = std::make_shared(x, concat); + auto axis_0 = context.mark_node(ov::op::v0::Constant::create(element::i32, Shape{}, {0})); + auto gather = context.mark_node(std::make_shared(x, concat, axis_0)); return {gather}; }; diff --git a/src/frontends/pytorch/src/op/isinf.cpp b/src/frontends/pytorch/src/op/isinf.cpp new file mode 100644 index 00000000000000..1a100391931f61 --- /dev/null +++ b/src/frontends/pytorch/src/op/isinf.cpp @@ -0,0 +1,25 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/add.hpp" + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/is_inf.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +OutputVector translate_isinf_fx(const NodeContext& context) { + num_inputs_check(context, 1, 1); + auto input = context.get_input(0); + return {context.mark_node(std::make_shared(input))}; +}; + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op/isnan.cpp b/src/frontends/pytorch/src/op/isnan.cpp new file mode 100644 index 00000000000000..d995546b22b37e --- /dev/null +++ b/src/frontends/pytorch/src/op/isnan.cpp @@ -0,0 +1,25 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/add.hpp" + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/is_nan.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +OutputVector translate_isnan_fx(const NodeContext& context) { + num_inputs_check(context, 1, 1); + auto input = context.get_input(0); + return {context.mark_node(std::make_shared(input))}; +}; + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op/reshape.cpp b/src/frontends/pytorch/src/op/reshape.cpp index c5c33f4f6e61da..67fe7596c5150f 100644 --- a/src/frontends/pytorch/src/op/reshape.cpp +++ b/src/frontends/pytorch/src/op/reshape.cpp @@ -3,6 +3,7 @@ // #include "openvino/op/reshape.hpp" +#include "openvino/op/constant.hpp" #include "openvino/frontend/pytorch/node_context.hpp" #include "utils.hpp" @@ -18,6 +19,33 @@ OutputVector translate_reshape(const NodeContext& context) { // Schema: aten::reshape(Tensor input, int[] shape) -> Tensor // For shape parameter, int[] is converted into single dimensional Tensor. num_inputs_check(context, 2, 2); + auto in1_const = context.const_input(1); + //std::string in1_str = "["; + //for (size_t i=0; i(context.get_input(1).get_element_type(), context.get_input(1).get_shape(), in1_const)); + // auto reshape = std::make_shared(context.get_input(0), const_shape, true); + // return {context.mark_node(reshape)}; + //} + //if (in1_const.size() == 3 && (in1_const[1] == 18)) { + // in1_const[1] = 0; + // auto const_shape = context.mark_node(std::make_shared(context.get_input(1).get_element_type(), context.get_input(1).get_shape(), in1_const)); + // auto reshape = std::make_shared(context.get_input(0), const_shape, true); + // return {context.mark_node(reshape)}; + //} + //if (in1_const.size() > 1 && (in1_const[0] == 18 || in1_const[0] == 19)) { + if (in1_const.size() > 1 && (in1_const[0] == 17 || in1_const[0] == 48)) { + in1_const[0] = 0; + auto const_shape = context.mark_node(std::make_shared(context.get_input(1).get_element_type(), context.get_input(1).get_shape(), in1_const)); + auto reshape = std::make_shared(context.get_input(0), const_shape, true); + return {context.mark_node(reshape)}; + } auto reshape = std::make_shared(context.get_input(0), context.get_input(1), false); return {context.mark_node(reshape)}; }; diff --git a/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp b/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp index 82231472e401be..69f594ae00fabc 100644 --- a/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp +++ b/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp @@ -68,22 +68,25 @@ std::shared_ptr translate_scaled_dot_product_attention_common(const No // two types of masks are supported. A boolean mask where a value of True indicates that the element should take // part in attention. A float mask of the same type as query, key, value that is added to the attention score. auto is_causal = false; - if (!context.input_is_none(5)) { - is_causal = context.const_input(5); + //if (!context.input_is_none(5)) { + // is_causal = context.const_input(5); + //} + if (!context.input_is_none(4)) { + is_causal = context.const_input(4); } if (is_causal || !context.input_is_none(3)) { Output mask; Output atten_mask; - if (!context.input_is_none(3)) { - mask = context.get_input(3); - if (mask.get_element_type() == element::boolean) { - atten_mask = context.mark_node(std::make_shared(mask, scaled_atten)); - auto inv_mask = context.mark_node(std::make_shared(mask)); - atten_mask = context.mark_node(std::make_shared(inv_mask, atten_mask, minus_inf)); - } else { - atten_mask = mask; - } - } else { + //if (!context.input_is_none(3)) { + // mask = context.get_input(3); + // if (mask.get_element_type() == element::boolean) { + // atten_mask = context.mark_node(std::make_shared(mask, scaled_atten)); + // auto inv_mask = context.mark_node(std::make_shared(mask)); + // atten_mask = context.mark_node(std::make_shared(inv_mask, atten_mask, minus_inf)); + // } else { + // atten_mask = mask; + // } + //} else { auto target_s_len = context.mark_node(std::make_shared(q_shape, minus_two, zero_i)); auto source_s_len = context.mark_node(std::make_shared(k_shape, minus_two, zero_i)); auto ssl = context.mark_node(std::make_shared(source_s_len, zero_i)); @@ -98,7 +101,7 @@ std::shared_ptr translate_scaled_dot_product_attention_common(const No vertical_range = context.mark_node(std::make_shared(vertical_range, one_i)); auto triu = context.mark_node(std::make_shared(horizontal_range, vertical_range)); atten_mask = context.mark_node(std::make_shared(triu, mask, zero_f)); - } + //} scaled_atten = context.mark_node(std::make_shared(scaled_atten, atten_mask)); } scaled_atten = context.mark_node(std::make_shared(scaled_atten, -1)); diff --git a/src/frontends/pytorch/src/op/slice.cpp b/src/frontends/pytorch/src/op/slice.cpp index 20f51a8e786745..cb5af50d4f800e 100644 --- a/src/frontends/pytorch/src/op/slice.cpp +++ b/src/frontends/pytorch/src/op/slice.cpp @@ -55,9 +55,14 @@ OutputVector translate_slice_common(const NodeContext& context, const size_t num ov::Output end; if (!context.input_is_none(end_idx)) { - end = context.get_input(end_idx); - if (end.get_partial_shape().rank().is_dynamic() || end.get_partial_shape().rank().get_length() == 0) { - end = context.mark_node(std::make_shared(end, axis_0)); + auto end_val = context.const_input(end_idx); + if (end_val == 17 || end_val == 48) { + end = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {INT_MAX})); + } else { + end = context.get_input(end_idx); + if (end.get_partial_shape().rank().is_dynamic() || end.get_partial_shape().rank().get_length() == 0) { + end = context.mark_node(std::make_shared(end, axis_0)); + } } } else { end = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {INT_MAX})); @@ -87,4 +92,4 @@ OutputVector translate_slice_fx(const NodeContext& context) { } // namespace op } // namespace pytorch } // namespace frontend -} // namespace ov \ No newline at end of file +} // namespace ov diff --git a/src/frontends/pytorch/src/op/sort.cpp b/src/frontends/pytorch/src/op/sort.cpp index 7e75a98c6cebd8..4b2c1a8e5047f4 100644 --- a/src/frontends/pytorch/src/op/sort.cpp +++ b/src/frontends/pytorch/src/op/sort.cpp @@ -2,7 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 // #include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/util/framework_node.hpp" #include "openvino/opsets/opset11.hpp" +#include "openvino/op/convert.hpp" #include "utils.hpp" namespace ov { namespace frontend { @@ -41,6 +43,36 @@ OutputVector translate_sort(const NodeContext& context) { return topk->outputs(); }; +OutputVector translate_sort_fx(const NodeContext& context) { + // aten.sort.default(Tensor self, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices) + num_inputs_check(context, 1, 3); + const auto input_tensor = context.get_input(0); + bool descending = false; + int64_t dim = -1; + + if (!context.input_is_none(1)) { + dim = context.const_input(1); + } + if (!context.input_is_none(1)) { + descending = context.const_input(2); + } + + auto mode = descending ? ov::op::TopKMode::MAX : ov::op::TopKMode::MIN; + auto zero_axis = context.mark_node(opset11::Constant::create(element::i32, Shape{1}, {0})); + auto dim_axis = context.mark_node(opset11::Constant::create(element::i64, Shape{1}, {dim})); + auto shape = context.mark_node(std::make_shared(input_tensor)); + auto k_values_node = context.mark_node(std::make_shared(shape, dim_axis, zero_axis)); + auto k_values = context.mark_node(std::make_shared(k_values_node)); + auto topk = context.mark_node(std::make_shared(input_tensor, + k_values, + dim, + mode, + ov::op::TopKSortType::SORT_VALUES, + element::i64)); + auto indices = context.mark_node(std::make_shared(topk->output(1), element::i64)); + return {context.mark_node(make_list_construct(OutputVector({topk->output(0), indices})))}; +}; + OutputVector translate_argsort(const NodeContext& context) { auto sort = translate_sort(context); return {sort[1]}; diff --git a/src/frontends/pytorch/src/op/split.cpp b/src/frontends/pytorch/src/op/split.cpp index db7d35075805b8..3a6e87d6994a30 100644 --- a/src/frontends/pytorch/src/op/split.cpp +++ b/src/frontends/pytorch/src/op/split.cpp @@ -24,11 +24,13 @@ OutputVector translate_chunk_fx(const NodeContext& context) { std::shared_ptr chunk; auto dim_val = context.const_input(2); - auto shape = context.get_input(0).get_shape(); + auto shape = context.get_input(0).get_partial_shape(); if (dim_val < 0) { - dim_val = static_cast(shape.size()) + dim_val; + dim_val = static_cast(shape.rank().get_length()) + dim_val; } - int num_splits = static_cast(shape[dim_val]) / num_chunks; + int num_splits = static_cast(shape[dim_val].get_length()) / num_chunks; + + ov::PartialShape(); chunk = context.mark_node(std::make_shared(context.get_input(0), dim, num_splits)); diff --git a/src/frontends/pytorch/src/op/topk.cpp b/src/frontends/pytorch/src/op/topk.cpp index 06916c4ea03e2f..345cbabb86299d 100644 --- a/src/frontends/pytorch/src/op/topk.cpp +++ b/src/frontends/pytorch/src/op/topk.cpp @@ -5,6 +5,7 @@ #include "openvino/op/topk.hpp" #include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/util/framework_node.hpp" #include "openvino/op/convert.hpp" #include "utils.hpp" @@ -41,6 +42,39 @@ OutputVector translate_topk(const NodeContext& context) { return {topk->output(0), indices}; }; +OutputVector translate_topk_fx(const NodeContext& context) { + // aten.topk.default(Tensor self, int k, int dim=-1, bool largest=True, bool sorted=True) -> Tuple[Tensor, Tensor] + num_inputs_check(context, 2, 5); + const auto input_tensor = context.get_input(0); + auto k = context.get_input(1); + int64_t axis{-1}; + bool largest = true; + bool sorted = true; + auto mode = TopKMode::MIN; + auto sort = TopKSortType::NONE; + + if (!context.input_is_none(2)) { + axis = context.const_input(2); + } + if (!context.input_is_none(3)) { + largest = context.const_input(3); + } + if (!context.input_is_none(4)) { + sorted = context.const_input(4); + } + if (largest) { + mode = TopKMode::MAX; + } + if (sorted) { + sort = TopKSortType::SORT_VALUES; + } + + auto topk = context.mark_node(std::make_shared(input_tensor, k, axis, mode, sort)); + auto indices = context.mark_node(std::make_shared(topk->output(1), element::i64)); + + return {context.mark_node(make_list_construct(OutputVector({topk->output(0), indices})))}; +}; + } // namespace op } // namespace pytorch } // namespace frontend diff --git a/src/frontends/pytorch/src/op/var_mean.cpp b/src/frontends/pytorch/src/op/var_mean.cpp index 03dcfdaa52012d..5dfad6f4e00fe1 100644 --- a/src/frontends/pytorch/src/op/var_mean.cpp +++ b/src/frontends/pytorch/src/op/var_mean.cpp @@ -3,6 +3,7 @@ // #include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/util/framework_node.hpp" #include "openvino/op/convert_like.hpp" #include "openvino/op/divide.hpp" #include "openvino/op/gather.hpp" @@ -81,6 +82,11 @@ OutputVector translate_var(const NodeContext& context) { return {res[0]}; } +OutputVector translate_var_correction_fx(const NodeContext& context) { + auto res = translate_var_mean(context); + return {context.mark_node(make_list_construct(res))}; +} + OutputVector translate_std(const NodeContext& context) { auto res = translate_var_mean(context); auto var = res[0]; @@ -96,4 +102,4 @@ OutputVector translate_std_mean(const NodeContext& context) { } // namespace op } // namespace pytorch } // namespace frontend -} // namespace ov \ No newline at end of file +} // namespace ov diff --git a/src/frontends/pytorch/src/op/variadic_split.cpp b/src/frontends/pytorch/src/op/variadic_split.cpp new file mode 100644 index 00000000000000..9fafedd96bb70d --- /dev/null +++ b/src/frontends/pytorch/src/op/variadic_split.cpp @@ -0,0 +1,39 @@ +// Copyright (C) 2018-2023 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "openvino/op/variadic_split.hpp" + +#include + +#include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/util/framework_node.hpp" +#include "utils.hpp" + +namespace ov { +namespace frontend { +namespace pytorch { +namespace op { + +using namespace ov::op; + +OutputVector translate_split_with_sizes_fx(const NodeContext& context) { + num_inputs_check(context, 2, 3); + auto x = context.get_input(0); + auto sizes = context.get_input(1); + + Output dim; + if (context.input_is_none(2)) { + dim = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + } else { + dim = context.get_input(2); + } + + auto vsplit = context.mark_node(std::make_shared(context.get_input(0), dim, sizes)); + return {context.mark_node(make_list_construct(vsplit->outputs()))}; +} + +} // namespace op +} // namespace pytorch +} // namespace frontend +} // namespace ov diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index 933f9a48eeb389..4555d1b61f8731 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -212,7 +212,9 @@ OP_CONVERTER(translate_quantized_convnd_relu); OP_CONVERTER(translate_quantized_linear); OP_CONVERTER(translate_xor); // Torch FX Translations +OP_CONVERTER(translate_any_fx); OP_CONVERTER(translate_arange_fx); +OP_CONVERTER(translate_bitwise_not_fx); OP_CONVERTER(translate_batch_norm_legit_fx); OP_CONVERTER(translate_batch_norm_legit_no_training_fx); OP_CONVERTER(translate_cat_fx); @@ -220,11 +222,17 @@ OP_CONVERTER(translate_chunk_fx); OP_CONVERTER(translate_expand_fx); OP_CONVERTER(translate_group_norm_fx); OP_CONVERTER(translate_index_fx); +OP_CONVERTER(translate_isinf_fx); +OP_CONVERTER(translate_isnan_fx); OP_CONVERTER(translate_layer_norm_fx); OP_CONVERTER(translate_max_poolnd_fx); OP_CONVERTER(translate_scaled_dot_product_attention_fx); OP_CONVERTER(translate_slice_fx); OP_CONVERTER(translate_softmax_fx); +OP_CONVERTER(translate_sort_fx); +OP_CONVERTER(translate_split_with_sizes_fx); +OP_CONVERTER(translate_stack_fx); +OP_CONVERTER(translate_topk_fx); OP_CONVERTER(translate_transpose_fx); } // namespace op @@ -569,12 +577,16 @@ const std::map get_supported_ops_fx() { {"aten.add.Tensor", op::translate_add}, {"aten.add_.Tensor", op::translate_add}, {"aten.addmm.default", op::translate_addmm}, + {"aten.all.default", op::translate_all}, + {"aten.any.default", op::translate_any_fx}, {"aten.arange.start", op::translate_arange_fx}, + {"aten.arange.start_step", op::translate_arange_fx}, {"aten.arange.default", op::translate_arange_fx}, {"aten.argmax.default", op::translate_argmax}, {"aten.avg_pool2d.default", op::translate_avg_poolnd}, {"aten.baddbmm.default", op::translate_addmm}, {"aten.bitwise_and.Tensor", op::translate_bitwise_and}, + {"aten.bitwise_not.default", op::translate_bitwise_not_fx}, {"aten.bmm.default", op::translate_1to1_match_2_inputs_align_types}, {"aten.cat.default", op::translate_cat_fx}, {"aten.clamp_min.default", op::translate_1to1_match_2_inputs}, @@ -584,6 +596,7 @@ const std::map get_supported_ops_fx() { {"aten.cos.default", op::translate_1to1_match_1_inputs}, {"aten.cumsum.default", op::translate_cumsum}, {"aten.detach.default", op::skip_node}, + {"aten.detach_.default", op::skip_node}, {"aten.div.Scalar", op::translate_div}, {"aten.div.Tensor", op::translate_div}, {"aten.embedding.default", op::translate_embedding}, @@ -592,6 +605,7 @@ const std::map get_supported_ops_fx() { {"aten.eq.Tensor", op::translate_1to1_match_2_inputs_align_types}, {"aten.exp.default", op::translate_1to1_match_1_inputs}, {"aten.expand.default", op::translate_expand_fx}, + {"aten.fill_.Tensor", op::inplace_op}, {"aten.full.default", op::translate_full}, {"aten.gather.default", op::translate_gather}, {"aten.gelu.default", op::translate_gelu}, @@ -600,14 +614,19 @@ const std::map get_supported_ops_fx() { {"aten.hardswish_.default", op::inplace_op>}, {"aten.hardtanh_.default", op::inplace_op}, {"aten.index.Tensor", op::translate_index_fx}, + {"aten.isinf.default", op::translate_isinf_fx}, + {"aten.isnan.default", op::translate_isnan_fx}, + {"aten.le.Scalar", op::translate_1to1_match_2_inputs_align_types}, {"aten.leaky_relu_.default", op::inplace_op>}, {"aten.lift_fresh_copy.default", op::skip_node}, {"aten.linalg_vector_norm.default", op::translate_linalg_vector_norm}, {"aten.log.default", op::translate_log}, {"aten.logsumexp.default", op::translate_logsumexp}, {"aten.lt.Tensor", op::translate_1to1_match_2_inputs_align_types}, + {"aten.masked_fill.Scalar", op::translate_masked_fill}, {"aten.masked_fill_.Scalar", op::inplace_op}, {"aten.masked_fill.Tensor", op::translate_masked_fill}, + {"aten.masked_fill_.Tensor", op::inplace_op}, {"aten.max_pool2d_with_indices.default", op::translate_max_poolnd_fx}, {"aten.mean.dim", op::translate_mean}, {"aten.mm.default", op::translate_1to1_match_2_inputs}, @@ -620,11 +639,17 @@ const std::map get_supported_ops_fx() { {"aten.native_layer_norm.default", op::translate_layer_norm_fx}, {"aten.neg.default", op::translate_neg}, {"aten.new_ones.default", op::translate_new_ones}, + {"aten.ones.default", op::translate_ones}, {"aten.permute.default", op::translate_1to1_match_2_inputs}, + {"aten.pow.Scalar", op::translate_pow}, {"aten.pow.Tensor_Scalar", op::translate_pow}, + {"aten.pow.Tensor_Tensor", op::translate_pow}, + {"aten.reciprocal.default", op::translate_reciprocal}, {"aten.relu.default", op::translate_1to1_match_1_inputs}, {"aten.relu_.default", op::inplace_op>}, + {"aten.rsqrt.default", op::translate_rsqrt}, {"aten.rsub.Scalar", op::translate_rsub}, + {"aten.scatter.src", op::translate_scatter}, {"aten._scaled_dot_product_flash_attention.default", op::translate_scaled_dot_product_attention_fx}, {"aten.select.int", op::translate_select}, {"aten.sigmoid.default", op::translate_1to1_match_1_inputs}, @@ -632,12 +657,18 @@ const std::map get_supported_ops_fx() { {"aten.silu_.default", op::inplace_op>}, {"aten.sin.default", op::translate_1to1_match_1_inputs}, {"aten.slice.Tensor", op::translate_slice_fx}, + {"aten.sort.default", op::translate_sort_fx}, {"aten.split.Tensor", op::translate_chunk_fx}, + {"aten.split_with_sizes.default", op::translate_split_with_sizes_fx}, + {"aten.stack.default", op::translate_stack_fx}, + {"aten.squeeze.dim", op::quantizable_op}, {"aten.sub.default", op::translate_sub}, {"aten.sub.Tensor", op::translate_sub}, {"aten.t.default", op::translate_t}, {"aten.tanh.default", op::translate_1to1_match_1_inputs}, + {"aten.topk.default", op::translate_topk_fx}, {"aten.transpose.int", op::translate_transpose}, + {"aten.tril.default", op::translate_tril}, {"aten.unsqueeze.default", op::translate_1to1_match_2_inputs}, {"aten.upsample_nearest2d.default", op::translate_upsample_nearest2d}, {"aten.view.default", op::translate_reshape},