[QNN-EP] Add reshape handler for split axes reshape

onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc

@@ -2408,7 +2408,7 @@ static bool FinalizeReshapeShape(const std::vector<int64_t>& input_shape,      /
   return true;
 }
 
-bool HandleReshape(HandlerArgs& args) {
+static bool HandleReshapeAsTranspose(HandlerArgs& args) {
   // A Reshape can be logically equivalent to a Transpose if all dims with a value > 1 remain in the same order
   // and do not change size. If so, we can use HandleTransposeImpl to merge them.
   //  e.g. Reshape(input {1, 512, 4, 1}, shape {1, 1, 512, 4}) is equivalent to Transpose with perms { 0, 3, 1, 2 }
@@ -2498,6 +2498,143 @@ bool HandleReshape(HandlerArgs& args) {
   return HandleTransposeImpl(args, perms);
 }
 
+// Push a Transpose past a Reshape when the Reshape *splits* one or more
+// post-transpose axes into contiguous groups of output axes.
+//
+// Example (motivating case):
+//   input {1, 12, 20, 24} -> Transpose(perm=[0,3,1,2]) -> {1, 24, 12, 20}
+//                         -> Reshape({1, 3, 8, 12, 20}) -> {1, 3, 8, 12, 20}
+//
+//   After: input -> Reshape({1, 12, 20, 3, 8}) -> Transpose(perm=[0,3,4,1,2])
+//                                              -> {1, 3, 8, 12, 20}
+//
+// Restrictions enforced:
+//   - Both transpose input shape and Reshape output shape (from shape
+//     inference on the Reshape's value_info) must be fully concrete
+//     (no symbolic / -1 / 0 dims). We drive off the output shape rather
+//     than the shape-input initializer so -1 / 0 have already been resolved
+//     by upstream shape inference.
+//   - Reshape output rank must exceed the transpose rank; equal or smaller
+//     ranks are covered by HandleReshapeAsTranspose.
+//   - The reshape output must partition into exactly one contiguous group per
+//     post-transpose axis, using every output dim.
+static bool HandleReshapeSplit(HandlerArgs& args) {
+  auto transpose_input_shape_opt = args.ctx.graph.GetValueInfo(args.transpose.Inputs()[0])->Shape();
+  if (!transpose_input_shape_opt.has_value()) {
+    return false;
+  }
+  const std::vector<int64_t>& transpose_input_shape = *transpose_input_shape_opt;
+  for (int64_t d : transpose_input_shape) {
+    if (d < 0) {
+      return false;
+    }
+  }
+
+  size_t rank = args.perm.size();
+  if (transpose_input_shape.size() != rank) {
+    return false;
+  }
+  std::vector<int64_t> transposed_shape(rank);
+  for (size_t i = 0; i < rank; ++i) {
+    transposed_shape[i] = transpose_input_shape[gsl::narrow_cast<size_t>(args.perm[i])];
+  }
+
+  auto reshape_output_shape_opt = args.ctx.graph.GetValueInfo(args.node.Outputs()[0])->Shape();
+  if (!reshape_output_shape_opt.has_value()) {
+    return false;
+  }
+  const std::vector<int64_t>& requested_shape = *reshape_output_shape_opt;
+  for (int64_t d : requested_shape) {
+    if (d <= 0) {
+      return false;  // symbolic / -1 / 0 — inference didn't resolve; bail.
+    }
+  }
+
+  // This handler is for pure splits (output rank > post-transpose rank).
+  // Equal or smaller ranks are the domain of HandleReshapeAsTranspose.
+  if (requested_shape.size() <= rank) {
+    return false;
+  }
+
+  // Partition requested_shape into `rank` contiguous groups: group j's product
+  // must equal transposed_shape[j], consuming output dims left-to-right.
+  // Consume at least one output dim per group so that size-1 post-transpose
+  // axes get a unique home rather than an empty group (which would leave the
+  // next group's assignment ambiguous when the requested shape also contains
+  // leading 1s).
+  std::vector<std::pair<size_t, size_t>> groups;  // [start, end) in reshape-output axis space
+  groups.reserve(rank);
+  size_t cursor = 0;
+  for (size_t j = 0; j < rank; ++j) {
+    int64_t target = transposed_shape[j];
+    size_t group_start = cursor;
+    int64_t prod = 1;
+    do {
+      if (cursor >= requested_shape.size()) {
+        return false;
+      }
+      prod *= requested_shape[cursor];
+      ++cursor;
+    } while (prod < target);
+    if (prod != target) {
+      return false;
+    }
+    groups.emplace_back(group_start, cursor);
+  }
+  if (cursor != requested_shape.size()) {
+    return false;
+  }
+
+  // Build the new (pre-Transpose) Reshape shape: iterate pre-transpose axes
+  // in order 0..rank-1. For each pre-transpose axis i, the post-transpose axis
+  // that carries it is args.perm_inv[i]; append that group's dims from
+  // requested_shape. This is the shape the Reshape emits when applied
+  // directly to the pre-transpose tensor. new_perm[k] records the position
+  // that requested_shape[k] ends up at, giving the Transpose that restores
+  // the original Reshape's externally-observed ordering.
+  std::vector<int64_t> new_reshape_shape;
+  new_reshape_shape.reserve(requested_shape.size());
+  std::vector<int64_t> new_perm(requested_shape.size());
+  for (size_t i = 0; i < rank; ++i) {
+    size_t j = gsl::narrow_cast<size_t>(args.perm_inv[i]);
+    const auto& [gs, ge] = groups[j];
+    for (size_t k = gs; k < ge; ++k) {
+      new_perm[k] = gsl::narrow_cast<int64_t>(new_reshape_shape.size());
+      new_reshape_shape.push_back(requested_shape[k]);
+    }
+  }
+
+  // Rewrite the graph:
+  //   1) Point the Reshape at the pre-transpose tensor with the new shape.
+  //   2) Insert a Transpose(new_perm) on the Reshape's output unless it's
+  //      identity (only possible when the original Transpose was a no-op).
+  //   3) Drop the original Transpose if nothing else consumes it.
+  std::string_view transpose_input = args.transpose.Inputs()[0];
+  std::vector<int64_t> shape_initializer_shape{gsl::narrow_cast<int64_t>(new_reshape_shape.size())};
+  std::string_view new_shape_init =
+      AddInitializerInt64(args.ctx.graph, shape_initializer_shape, new_reshape_shape);
+
+  args.node.SetInput(0, transpose_input);
+  args.node.SetInput(1, new_shape_init);
+
+  if (!IsIdentityPerm(new_perm)) {
+    TransposeOutput(args.ctx.graph, args.node, 0, new_perm, InvertPerm(new_perm));
+  }
+
+  if (!args.ctx.graph.HasValueConsumers(args.transpose.Outputs()[0])) {
+    args.ctx.graph.RemoveNode(args.transpose);
+  }
+
+  return true;
+}
+
+bool HandleReshape(HandlerArgs& args) {
+  if (HandleReshapeAsTranspose(args)) {
+    return true;
+  }
+  return HandleReshapeSplit(args);
+}
+
 constexpr HandlerInfo reshape_handler = {&FirstInput, &HandleReshape, /*transposes_outputs*/ false};
 
 // TODO: check binary size of this and replace it with constexpr if large

onnxruntime/test/optimizer/transpose_optimizer_test.cc

@@ -12,6 +12,7 @@
 #include "core/framework/op_node_proto_helper.h"
 #include "core/graph/graph.h"
 #include "core/graph/node_attr_utils.h"
+#include "core/optimizer/initializer.h"
 #include "core/optimizer/transpose_optimization/onnx_transpose_optimization.h"
 #include "core/optimizer/transpose_optimization/optimizer_api.h"
 #include "core/optimizer/transpose_optimization/ort_optimizer_utils.h"
@@ -4139,6 +4140,181 @@ TEST(TransposeOptimizerTests, TestReshapeWithZero) {
                        true);  // allow_zero
 }
 
+// Transpose -> Reshape where the Reshape splits one or more post-transpose axes
+// into contiguous groups of output axes. HandleReshapeSplit should rewrite this
+// as Reshape (on the pre-transpose tensor) followed by a Transpose whose perm
+// restores the original externally-observed ordering. The original Transpose
+// should be removed and the new Transpose's perm captured in `expected_new_perms`.
+static void TestTransposeReshapeSplit(const std::vector<int64_t>& input_shape,
+                                      const std::vector<int64_t>& perms,
+                                      const std::vector<int64_t>& reshape_shape,
+                                      const std::vector<int64_t>& expected_new_perms) {
+  auto build_test_case = [&](ModelTestBuilder& builder) {
+    auto* input_arg = builder.MakeInput<float>(input_shape, 0.0, 1.0);
+    auto* mul_arg1 = builder.MakeInput<float>({1}, 0.0, 1.0);
+    auto* reshape_shape_value =
+        builder.MakeInitializer<int64_t>({int64_t(reshape_shape.size())}, reshape_shape);
+
+    auto* mul_out_0 = builder.MakeOutput();
+    auto* transpose_out_0 = builder.MakeIntermediate();
+    auto* reshape_out_0 = builder.MakeIntermediate();
+    auto* identity_out_0 = builder.MakeOutput();
+
+    builder.AddNode("Mul", {input_arg, mul_arg1}, {mul_out_0});
+
+    auto& transpose_1 = builder.AddNode("Transpose", {mul_out_0}, {transpose_out_0});
+    transpose_1.AddAttribute("perm", perms);
+
+    builder.AddNode("Reshape", {transpose_out_0, reshape_shape_value}, {reshape_out_0});
+    builder.AddNode("Identity", {reshape_out_0}, {identity_out_0});
+  };
+
+  auto check_optimized_graph = [&](InferenceSessionWrapper& session) {
+    const auto& graph = session.GetGraph();
+    std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+
+    ASSERT_EQ(op_to_count["Transpose"], 1);
+    ASSERT_EQ(op_to_count["Reshape"], 1);
+
+    const auto& nodes = graph.Nodes();
+    const Node& transpose = *std::find_if(nodes.begin(), nodes.end(),
+                                          [](const auto& node) { return node.OpType() == "Transpose"; });
+    ProtoHelperNodeContext proto_helper_ctx(transpose);
+    OpNodeProtoHelper<ProtoHelperNodeContext> proto_helper(&proto_helper_ctx);
+    std::vector<int64_t> actual_perms;
+    ASSERT_STATUS_OK(proto_helper.GetAttrs<int64_t>("perm", actual_perms));
+    ASSERT_THAT(actual_perms, testing::ContainerEq(expected_new_perms));
+  };
+
+  TransformerTester(build_test_case,
+                    check_optimized_graph,
+                    TransformerLevel::Default,
+                    TransformerLevel::Level1,
+                    /*opset_version*/ {15, 23});
+}
+
+// Negative-case variant: assert HandleReshapeSplit bails and the graph is
+// unchanged (Transpose still upstream of Reshape).
+static void TestTransposeReshapeSplitUnchanged(const std::vector<int64_t>& input_shape,
+                                               const std::vector<int64_t>& perms,
+                                               const std::vector<int64_t>& reshape_shape) {
+  auto build_test_case = [&](ModelTestBuilder& builder) {
+    auto* input_arg = builder.MakeInput<float>(input_shape, 0.0, 1.0);
+    auto* mul_arg1 = builder.MakeInput<float>({1}, 0.0, 1.0);
+    auto* reshape_shape_value =
+        builder.MakeInitializer<int64_t>({int64_t(reshape_shape.size())}, reshape_shape);
+
+    auto* mul_out_0 = builder.MakeOutput();
+    auto* transpose_out_0 = builder.MakeIntermediate();
+    auto* reshape_out_0 = builder.MakeIntermediate();
+    auto* identity_out_0 = builder.MakeOutput();
+
+    builder.AddNode("Mul", {input_arg, mul_arg1}, {mul_out_0});
+
+    auto& transpose_1 = builder.AddNode("Transpose", {mul_out_0}, {transpose_out_0});
+    transpose_1.AddAttribute("perm", perms);
+
+    builder.AddNode("Reshape", {transpose_out_0, reshape_shape_value}, {reshape_out_0});
+    builder.AddNode("Identity", {reshape_out_0}, {identity_out_0});
+  };
+
+  auto check_optimized_graph = [&](InferenceSessionWrapper& session) {
+    const auto& graph = session.GetGraph();
+    std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+    ASSERT_EQ(op_to_count["Transpose"], 1);
+    ASSERT_EQ(op_to_count["Reshape"], 1);
+
+    // Reshape's data input should still be the Transpose's output. If the
+    // split-handler had fired, the Reshape would be fed by the pre-transpose
+    // tensor instead.
+    const auto& nodes = graph.Nodes();
+    const Node& transpose = *std::find_if(nodes.begin(), nodes.end(),
+                                          [](const auto& node) { return node.OpType() == "Transpose"; });
+    const Node& reshape = *std::find_if(nodes.begin(), nodes.end(),
+                                        [](const auto& node) { return node.OpType() == "Reshape"; });
+    ASSERT_EQ(reshape.InputDefs()[0]->Name(), transpose.OutputDefs()[0]->Name());
+  };
+
+  TransformerTester(build_test_case,
+                    check_optimized_graph,
+                    TransformerLevel::Default,
+                    TransformerLevel::Level1,
+                    /*opset_version*/ {15, 23});
+}
+
+// Transpose -> Reshape(split) is rewritten to Reshape -> Transpose. The
+// motivating case (documented in HandleReshapeSplit):
+//   input {1,12,20,24} -Transpose[0,3,1,2]-> {1,24,12,20}
+//     -Reshape({1,3,8,12,20})-> {1,3,8,12,20}
+//   becomes
+//   input {1,12,20,24} -Reshape({1,12,20,3,8})-> {1,12,20,3,8}
+//     -Transpose[0,3,4,1,2]-> {1,3,8,12,20}
+TEST(TransposeOptimizerTests, TestReshapeSplit) {
+  // Motivating case: split the transposed channel axis into (3, 8).
+  // new_reshape_shape (pre-transpose) = {1, 12, 20, 3, 8}.
+  TestTransposeReshapeSplit(/*input_shape*/ {1, 12, 20, 24},
+                            /*perms*/ {0, 3, 1, 2},
+                            /*reshape_shape*/ {1, 3, 8, 12, 20},
+                            /*expected_new_perms*/ {0, 3, 4, 1, 2});
+
+  // Split multiple post-transpose axes. Transpose (2,4,6) with perm [2,0,1]
+  // gives (6,2,4); splitting 6->(2,3) and 4->(2,2) yields output (2,3,2,2,2).
+  // perm_inv=[1,2,0]. New Reshape emits groups in pre-transpose order:
+  //   i=0 -> groups[1]=[2,3) -> req[2]=2         new_perm[2]=0
+  //   i=1 -> groups[2]=[3,5) -> req[3]=2,req[4]=2 new_perm[3]=1,new_perm[4]=2
+  //   i=2 -> groups[0]=[0,2) -> req[0]=2,req[1]=3 new_perm[0]=3,new_perm[1]=4
+  // => new_reshape_shape={2,2,2,2,3}, new_perm={3,4,0,1,2}.
+  TestTransposeReshapeSplit(/*input_shape*/ {2, 4, 6},
+                            /*perms*/ {2, 0, 1},
+                            /*reshape_shape*/ {2, 3, 2, 2, 2},
+                            /*expected_new_perms*/ {3, 4, 0, 1, 2});
+
+  // Size-1 post-transpose axis gets its own group.
+  //   input {1, 6, 4} -Transpose[2,0,1]-> {4, 1, 6}
+  //   -Reshape({2, 2, 1, 2, 3})- (groups: 4->(2,2), 1->(1), 6->(2,3))
+  // Pre-transpose axes 0,1,2 map to post-transpose axes 1,2,0 (perm_inv=[1,2,0]).
+  // Pre-transpose order picks groups[1]=(1), groups[2]=(2,3), groups[0]=(2,2)
+  //   => new_reshape_shape = {1, 2, 3, 2, 2}
+  //   requested[0..1]=(2,2) sit at new positions 3,4 (from group 0)
+  //   requested[2]=1        sits at new position 0 (from group 1)
+  //   requested[3..4]=(2,3) sit at new positions 1,2 (from group 2)
+  //   new_perm = {3, 4, 0, 1, 2}
+  TestTransposeReshapeSplit(/*input_shape*/ {1, 6, 4},
+                            /*perms*/ {2, 0, 1},
+                            /*reshape_shape*/ {2, 2, 1, 2, 3},
+                            /*expected_new_perms*/ {3, 4, 0, 1, 2});
+}
+
+// Negative cases for HandleReshapeSplit: the split rewrite must bail so the
+// Transpose -> Reshape order is preserved.
+TEST(TransposeOptimizerTests, TestReshapeSplitBails) {
+  // Rank-shrinking reshape (flatten). HandleReshapeSplit only handles pure
+  // splits (output rank > post-transpose rank). Shrinks / same-rank shapes
+  // are out of scope and the graph must be left alone.
+  //   input {2,3,4} -Transpose[1,0,2]-> {3,2,4} -Reshape({6,4})-> {6,4}
+  // requested_shape.size()=2 <= rank=3, so the handler bails immediately.
+  TestTransposeReshapeSplitUnchanged(/*input_shape*/ {2, 3, 4},
+                                     /*perms*/ {1, 0, 2},
+                                     /*reshape_shape*/ {6, 4});
+
+  // Partition impossible: no contiguous prefix of output dims multiplies to
+  // the first transposed dim exactly.
+  //   input {2,3} -Transpose[1,0]-> {3,2} -Reshape({2,3,1})- ...
+  // rank=2, requested rank=3, so the size guard passes. Then target=3:
+  //   consume req[0]=2, prod=2 (<3); consume req[1]=3, prod=6 (!=3) -> bail.
+  TestTransposeReshapeSplitUnchanged(/*input_shape*/ {2, 3},
+                                     /*perms*/ {1, 0},
+                                     /*reshape_shape*/ {2, 3, 1});
+
+  // Split that would require re-ordering across the boundary of a transposed
+  // axis. Transposed shape (6,4); requested {2,12,1} tries to combine part of
+  // axis-0 with all of axis-1, which is not expressible as a per-axis split.
+  //   target=6: prod=2 (<6); prod=2*12=24 (!=6) -> bail.
+  TestTransposeReshapeSplitUnchanged(/*input_shape*/ {4, 6},
+                                     /*perms*/ {1, 0},
+                                     /*reshape_shape*/ {2, 12, 1});
+}
+
 // test Reshape with an inferred dim due to value of -1
 // test valid (inferred dim is 1:1 with existing) and invalid (inferred size differs)
 TEST(TransposeOptimizerTests, TestReshapeWithMinusOne) {