[PT FE] Support aten::unsafe_chunk op

mvafin · Oct 7, 2024 · d5abde0 · d5abde0
1 parent cf870cd
commit d5abde0
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 7 deletions.
diff --git a/src/frontends/pytorch/src/op/getitem.cpp b/src/frontends/pytorch/src/op/getitem.cpp
@@ -23,10 +23,10 @@ OutputVector translate_getitem(const NodeContext& context) {
     PYTORCH_OP_CONVERSION_CHECK(!idx_type.is<type::Str>(),
                                 "String index in aten::__getitem__ means dict input, this is not supported.");
     if (ov::as_type_ptr<ov::op::util::FrameworkNode>(input.get_node_shared_ptr())) {
-        PYTORCH_OP_CONVERSION_CHECK(!cast_fw_node(input.get_node_shared_ptr(), "aten::split"),
-                                    "special case for aten::__getitem__");
-        PYTORCH_OP_CONVERSION_CHECK(!cast_fw_node(input.get_node_shared_ptr(), "aten::chunk"),
-                                    "special case for aten::__getitem__");
+        PYTORCH_OP_CONVERSION_CHECK(
+            !cast_fw_node(input.get_node_shared_ptr(),
+                          std::vector<std::string>{"aten::split", "aten::chunk", "aten::unsafe_chunk"}),
+            "special case for aten::__getitem__");
         const auto&& list_elems = get_list_as_outputs(input);
         auto getitem_idx = context.const_input<int64_t>(1);
         if (getitem_idx < 0) {

diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp
@@ -663,6 +663,7 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
         // aten::unbind - Supported in limited set of patterns
         {"aten::unflatten", op::translate_unflatten},
         {"aten::unfold", op::translate_unfold},
+        // aten::unsafe_chunk - Supported in limited set of patterns
         {"aten::unsqueeze", op::quantizable_op<op::translate_1to1_match_2_inputs<opset10::Unsqueeze>>},
         {"aten::upsample_bicubic2d", op::translate_upsample_bicubic2d},
         {"aten::upsample_bilinear2d", op::translate_upsample_bilinear2d},

diff --git a/src/frontends/pytorch/src/transforms/aten_getitem_replacer.cpp b/src/frontends/pytorch/src/transforms/aten_getitem_replacer.cpp
@@ -124,7 +124,8 @@ AtenGetItemReplacer::AtenGetItemReplacer() {
                 auto gather = rg.make<v8::Gather>(input_concat, getitem_idx, zero);
                 replace_node(getitem, gather);
             }
-        } else if (auto chunk = cast_fw_node(input_node, "aten::chunk")) {
+        } else if (auto chunk =
+                       cast_fw_node(input_node, std::vector<std::string>{"aten::chunk", "aten::unsafe_chunk"})) {
             auto input_tensor = chunk->get_input_source_output(0);
             auto chunks_i32 = chunk->get_input_source_output(1);
             auto dim_i32 = chunk->get_input_source_output(2);

diff --git a/src/frontends/pytorch/src/transforms/index_loop_getitem_replacer.cpp b/src/frontends/pytorch/src/transforms/index_loop_getitem_replacer.cpp
@@ -52,7 +52,8 @@ IndexLoopGetitemReplacer::IndexLoopGetitemReplacer() {
         size_t chunk_idx = 0;
         auto loop_inputs = loop_op->input_values();
         for (size_t i = 1; i < loop_inputs.size(); i++) {
-            if (cast_fw_node(loop_inputs.at(i).get_node_shared_ptr(), "aten::chunk")) {
+            if (cast_fw_node(loop_inputs.at(i).get_node_shared_ptr(),
+                             std::vector<std::string>{"aten::chunk", "aten::unsafe_chunk"})) {
                 chunk_op = loop_inputs.at(i).get_node_shared_ptr();
                 chunk_idx = i;
                 break;

diff --git a/src/frontends/pytorch/src/transforms/prim_list_unpack_replacer.cpp b/src/frontends/pytorch/src/transforms/prim_list_unpack_replacer.cpp
@@ -88,7 +88,8 @@ PrimListUnpackReplacer::PrimListUnpackReplacer() {
             replace_node(list_unpack, split);
 
             return true;
-        } else if (auto chunk = cast_fw_node(input_node, "aten::chunk")) {
+        } else if (auto chunk =
+                       cast_fw_node(input_node, std::vector<std::string>{"aten::chunk", "aten::unsafe_chunk"})) {
             if (list_unpack->get_output_size() == 1) {
                 list_unpack->output(0).replace(input_node->input_value(0));
                 return true;

diff --git a/src/frontends/pytorch/src/utils.cpp b/src/frontends/pytorch/src/utils.cpp
@@ -378,6 +378,21 @@ std::shared_ptr<ov::op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node>
     return fw_node;
 }
 
+std::shared_ptr<ov::op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node> node,
+                                                          const std::vector<std::string>& types) {
+    auto fw_node = std::dynamic_pointer_cast<ov::op::util::FrameworkNode>(node);
+    if (!fw_node) {
+        return nullptr;
+    }
+    const auto& attrs = fw_node->get_attrs();
+    for (auto type : types) {
+        if (attrs.find(PtFrameworkNode::op_type_key) != attrs.end() && attrs.at(PtFrameworkNode::op_type_key) == type) {
+            return fw_node;
+        }
+    }
+    return nullptr;
+}
+
 std::shared_ptr<ov::Node> make_list_construct(const ov::OutputVector& inputs) {
     auto list_construct = std::make_shared<ov::op::util::FrameworkNode>(inputs, inputs.size());
     ov::op::util::FrameworkNodeAttrs attrs;

diff --git a/src/frontends/pytorch/src/utils.hpp b/src/frontends/pytorch/src/utils.hpp
@@ -80,6 +80,8 @@ OutputVector make_framework_node_ignore_bodies(const NodeContext& context, const
 OutputVector make_framework_node(const NodeContext& context, const std::string& exception);
 
 std::shared_ptr<op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node> node, const std::string& type);
+std::shared_ptr<op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node> node,
+                                                      const std::vector<std::string>& types);
 
 std::shared_ptr<Node> make_list_construct(const ov::OutputVector& inputs);