Revert find_reshape_cont guard relaxation from PR#4858

CHANGELOG.md

@@ -76,6 +76,7 @@ Full documentation for MIGraphX is available at
 * Fixed `QLinearConv` parsing for models with a bias and per-tensor weight quantization, which previously threw `same_dims: dequantizelinear: Dimensions do not match` (e.g. `resnet50_int8`); the bias scale is now broadcast to the bias shape before dequantizing.
 * Fixed the GPU problem cache failing to find entries after reload for pooling operator, resulting in redundant re-benchmarking when using a saved `MIGRAPHX_PROBLEM_CACHE`.
 * Fixed `slice_concat_gather` matcher and interaction between same table and cross table gather fusions(#5038).
+* Fixed system reboot (TDR) when running certain customer model usecase on MIGraphX EP by restoring safety guards in `find_reshape_cont` that prevented incorrect optimization of non-standard tensor layouts (#5052).
 ### Optimized
 * Reduced tuning time by scaling the per-candidate benchmark bundle to the candidate's op count (#4989).
 * Enabled tensor vectorization for GPU fused `argmin` and `argmax` (`gpu::arg_reduce`) (#4790).

src/simplify_reshapes.cpp

@@ -47,6 +47,7 @@
 #include <migraphx/instruction_traversal.hpp>
 #include <migraphx/output_iterator.hpp>
 #include <migraphx/par.hpp>
+#include <migraphx/reshape_dims.hpp>
 
 #include <array>
 #include <functional>
@@ -1291,7 +1292,7 @@ struct find_gather
             return;
 
         const std::size_t axis_index = tune_axis(dlens.size(), gather_op.axis, gather_op.name());
-        const auto axis_len = dlens.at(axis_index);
+        const auto axis_len          = dlens.at(axis_index);
         if(axis_len == 0)
             return;
 
@@ -1619,6 +1620,13 @@ struct find_reshape_cont
         if(ins->get_shape().ndim() > cont_input->get_shape().ndim())
             return;
 
+        auto rdims_sz = std::vector<std::size_t>(dims.begin(), dims.end());
+        if(not std::all_of(ins->inputs().begin(), ins->inputs().end(), [&](auto in) {
+                return in == in_ins or
+                    reshape_dims(in->get_shape(), rdims_sz, {.lazy = true}).has_value();
+            }))
+            return;
+
         auto out_lens = ins->get_shape().lens();
         std::vector<int64_t> out_dims(out_lens.begin(), out_lens.end());
         std::vector<instruction_ref> inputs;

test/simplify_reshapes_test.cpp

@@ -4016,6 +4016,7 @@ TEST_CASE(transpose_contiguous_reshape_binary_packed)
         m1.add_instruction(pass_op{}, add_ins);
     }
     run_pass(m1);
+
     migraphx::module m2;
     {
         auto x  = m2.add_parameter("x", {migraphx::shape::float_type, {2, 128, 28, 28}});
@@ -4034,16 +4035,14 @@ TEST_CASE(transpose_contiguous_reshape_binary_packed)
             conv1,
             w2); // (2, 512, 14, 14)
 
-        auto conv2_rsp = m2.add_instruction(
+        auto conv2_rsp1 = m2.add_instruction(
             migraphx::make_op("reshape", {{"dims", {2, 2, 2, 128, 14, 14}}}), conv2);
         auto conv2_trans = m2.add_instruction(
-            migraphx::make_op("transpose", {{"permutation", {0, 3, 4, 1, 5, 2}}}), conv2_rsp);
-        auto x_rsp =
-            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 128, 14, 2, 14, 2}}}), x);
-        auto add_ins = m2.add_instruction(migraphx::make_op("add"), conv2_trans, x_rsp);
-        auto add_rsp =
-            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 128, 28, 28}}}), add_ins);
-        m2.add_instruction(pass_op{}, add_rsp);
+            migraphx::make_op("transpose", {{"permutation", {0, 3, 4, 1, 5, 2}}}), conv2_rsp1);
+        auto conv2_rsp2 = m2.add_instruction(
+            migraphx::make_op("reshape", {{"dims", {2, 128, 28, 28}}}), conv2_trans);
+        auto add_ins = m2.add_instruction(migraphx::make_op("add"), conv2_rsp2, x);
+        m2.add_instruction(pass_op{}, add_ins);
     }
     EXPECT(m1 == m2);
 }
@@ -4067,6 +4066,7 @@ TEST_CASE(transpose_contiguous_reshape_binary_broadcast)
         m1.add_return({r});
     }
     run_pass(m1);
+
     migraphx::module m2;
     {
         migraphx::shape sx{migraphx::shape::float_type, {4}};
@@ -4075,13 +4075,14 @@ TEST_CASE(transpose_contiguous_reshape_binary_broadcast)
         auto x = m2.add_parameter("x", sx);
         auto y = m2.add_parameter("y", sy);
         auto y_trans =
-            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), y);
+            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}),
+            y);
         auto x_rsp   = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 2}}}), x);
         auto x_brcst = m2.add_instruction(
             migraphx::make_op("broadcast", {{"axis", 1}, {"out_lens", {2, 2, 2, 6}}}), x_rsp);
         auto add_ins = m2.add_instruction(migraphx::make_op("add"), y_trans, x_brcst);
-        auto r = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 4, 6}}}), add_ins);
-        m2.add_return({r});
+        auto r = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {2, 4, 6}}}),
+        add_ins); m2.add_return({r});
     }
     EXPECT(m1 == m2);
 }
@@ -5814,4 +5815,23 @@ TEST_CASE(broadcast_nop_reduce_mean)
     EXPECT(m1.sort() == m2.sort());
 }
 
+TEST_CASE(reshape_cont_nonstandard_groupnorm)
+{
+    migraphx::module m1;
+    {
+        migraphx::shape sx{migraphx::shape::float_type, {1, 8, 4, 4}, {128, 1, 32, 8}};
+        auto x = m1.add_parameter("x", sx);
+        auto y = m1.add_parameter("y", {migraphx::shape::float_type, {1, 2, 1}});
+        auto y_bcast =
+            m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 2, 64}}}), y);
+        // Group norm reshape: {1,8,4,4} → {1,2,64}
+        auto rsp     = m1.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2, 64}}}), x);
+        auto sub_ins = m1.add_instruction(migraphx::make_op("sub"), rsp, y_bcast);
+        m1.add_return({sub_ins});
+    }
+    migraphx::module m2 = m1;
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }