more sort tests; rename sort; more navi geg tests; slight optimization of moving sort into branch for multi-out case instead of outside the conditional

Author: bdevorem

src/include/migraphx/module.hpp

@@ -339,7 +339,7 @@ struct MIGRAPHX_EXPORT module
     ins_dep_map calc_implicit_deps() const;
 
     void repeat_while_changes(std::size_t n, const std::function<void()>& f);
-    void localized_sort(instruction_ref start_ins, instruction_ref end_ins);
+    void hoist_external_inputs(instruction_ref start_ins, instruction_ref end_ins);
 
     MIGRAPHX_EXPORT friend std::ostream& operator<<(std::ostream& os, const module& m);
     MIGRAPHX_EXPORT friend bool operator==(const module& x, const module& y);

src/module.cpp

@@ -1600,9 +1600,9 @@ void module::repeat_while_changes(std::size_t n, const std::function<void()>& f)
     }
 }
 
-// For topologically sorting a region in a module, canonically, such that the
-// dependent chain between the two input instructions is last
-void module::localized_sort(instruction_ref start_ins, instruction_ref end_ins)
+// Hoists external inputs (instructions not in the dependency chain between start_ins and end_ins)
+// to before start_ins, while preserving topological order
+void module::hoist_external_inputs(instruction_ref start_ins, instruction_ref end_ins)
 {
     // get the chain of instructions between start_ins and end_ins, inclusive
     auto fusion_ins = find_instructions_between(start_ins, end_ins, this);

src/targets/gpu/fuse_mlir.cpp

@@ -913,17 +913,16 @@ struct find_mlir_fused_geg_ops
         mm->add_return(return_vals);
         auto inputs = find_inputs(map_ins, &mpm.get_module(), mm);
 
-        // sort fusion section of module such that any external inputs are moved before the fusion
-        // so that we can safely place the fused mod in the multi-out case at the beginning of the
-        // chain
-        mpm.get_module().localized_sort(first_gemm_ins, second_gemm_ins);
-        auto fused_ins =
-            mpm.get_module().insert_instruction(first_gemm_ins,
-                                                mlir_op{second_gemm_ins->get_operator()},
-                                                mlir_contiguous(mpm, inputs),
-                                                {mm});
         if(first_gemm_has_multi_outs or elemwise_has_multi_outs)
         {
+            // hoist external inputs before the fusion so that we can safely place the fused mod
+            // in the multi-out case at the beginning of the chain
+            mpm.get_module().hoist_external_inputs(first_gemm_ins, second_gemm_ins);
+            auto fused_ins =
+                mpm.get_module().insert_instruction(first_gemm_ins,
+                                                    mlir_op{second_gemm_ins->get_operator()},
+                                                    mlir_contiguous(mpm, inputs),
+                                                    {mm});
             std::size_t output_idx = 0;
             if(elemwise_has_multi_outs)
             {
@@ -946,6 +945,11 @@ struct find_mlir_fused_geg_ops
         else
         {
             // simple single output case
+            auto fused_ins =
+                mpm.get_module().insert_instruction(first_gemm_ins,
+                                                    mlir_op{second_gemm_ins->get_operator()},
+                                                    mlir_contiguous(mpm, inputs),
+                                                    {mm});
             mpm.get_module().replace_instruction(second_gemm_ins, fused_ins);
         }
     }

test/gpu/fuse_mlir.cpp

@@ -24,13 +24,13 @@
 #include <migraphx/generate.hpp>
 #include <migraphx/dead_code_elimination.hpp>
 #include <migraphx/gpu/fuse_mlir.hpp>
+#include <migraphx/gpu/context.hpp>
 #include <migraphx/instruction.hpp>
 #include <migraphx/pass_manager.hpp>
 #include <migraphx/op/common.hpp>
 #include <migraphx/program.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/param_utils.hpp>
-#include <migraphx/load_save.hpp>
 #include <basic_ops.hpp>
 #include <group.hpp>
 #include <test.hpp>
@@ -55,6 +55,8 @@ struct non_mlir_op
 
 static void run_pass(migraphx::program& p, migraphx::gpu::fuse_mlir fm = {})
 {
+    static migraphx::gpu::context ctx;
+    fm.ctx = &ctx;
     fm.enable_extra = true;
     migraphx::run_passes(p, {fm, migraphx::dead_code_elimination{}});
 }
@@ -1816,13 +1818,13 @@ TEST_CASE(dot_add_dot)
     EXPECT(p1.sort() == p2.sort());
 }
 
-TEST_CASE(dot_add_dot_abc)
+TEST_CASE(dot_add_dot_abc_f32)
 // MLIR currently only supports (A*B)*C GEG patterns
 {
-    migraphx::shape s1{migraphx::shape::half_type, {2, 3}};
-    migraphx::shape s2{migraphx::shape::half_type, {3, 4}};
-    migraphx::shape s3{migraphx::shape::half_type, {2, 4}};
-    migraphx::shape s4{migraphx::shape::half_type, {4, 2}};
+    migraphx::shape s1{migraphx::shape::float_type, {2, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4}};
+    migraphx::shape s3{migraphx::shape::float_type, {2, 4}};
+    migraphx::shape s4{migraphx::shape::float_type, {4, 2}};
     migraphx::program p1;
     {
         auto* mm  = p1.get_main_module();
@@ -1836,15 +1838,14 @@ TEST_CASE(dot_add_dot_abc)
         auto dot2 = mm->add_instruction(migraphx::make_op("dot"), add, y); // {2, 4}*{4, 2} = {2, 2}
         mm->add_return({dot2});
     }
-    save(p1, "testest.mgx");
     run_pass(p1);
     // ensure "geg" is present. Earlier tests ensure the fusion is correct. This is just to ensure
     // it happens.
     std::stringstream ss;
     ss << p1;
     std::string program_str = ss.str();
 
-    // regardless if the matmul is correctly oriented, geg should not happen on navi
+    // regardless if the matmul is correctly oriented, f32 geg should not happen on navi
     auto device_name = migraphx::gpu::get_device_name();
     bool is_navi =
         migraphx::starts_with(device_name, "gfx11") or migraphx::starts_with(device_name, "gfx12");
@@ -1855,6 +1856,39 @@ TEST_CASE(dot_add_dot_abc)
         EXPECT(program_str.find("geg") != std::string::npos);
 }
 
+TEST_CASE(dot_add_dot_abc_fp16)
+// MLIR currently only supports (A*B)*C GEG patterns
+{
+    migraphx::shape s1{migraphx::shape::half_type, {2, 3}};
+    migraphx::shape s2{migraphx::shape::half_type, {3, 4}};
+    migraphx::shape s3{migraphx::shape::half_type, {2, 4}};
+    migraphx::shape s4{migraphx::shape::half_type, {4, 2}};
+    migraphx::program p1;
+    {
+        auto* mm  = p1.get_main_module();
+        auto a    = mm->add_parameter("a", s1);
+        auto b    = mm->add_parameter("b", s2);
+        auto x    = mm->add_parameter("x", s3);
+        auto y    = mm->add_parameter("y", s4);
+        auto dot1 = mm->add_instruction(migraphx::make_op("dot"), a, b); // {2,4}
+        auto add =
+            add_pointwise(p1, "main:pointwise0", {dot1, x}, single_pointwise("add")); // {2, 4}
+        auto dot2 = mm->add_instruction(migraphx::make_op("dot"), add, y); // {2, 4}*{4, 2} = {2, 2}
+        mm->add_return({dot2});
+    }
+    run_pass(p1);
+    std::stringstream ss;
+    ss << p1;
+    std::string program_str = ss.str();
+
+    // ensure "geg" is present if the fusion flag is enabled; type is fp16 so it should
+    // run regardless of if navi
+    if(not migraphx::enabled(MIGRAPHX_ENABLE_MLIR_GEG_FUSION{}))
+        EXPECT(program_str.find("geg") == std::string::npos);
+    else
+        EXPECT(program_str.find("geg") != std::string::npos);
+}
+
 TEST_CASE(dot_add_dot_cab)
 // MLIR currently does not support C*(A*B) GEG patterns
 {

test/module_test.cpp

@@ -1369,7 +1369,7 @@ TEST_CASE(pathological_dfs_graph_sort)
     EXPECT(is_sorted(m));
 }
 
-TEST_CASE(localized_sort)
+TEST_CASE(hoist_external_inputs)
 {
     migraphx::shape s1{migraphx::shape::float_type, {2, 3}};
     migraphx::shape s2{migraphx::shape::float_type, {3, 4}};
@@ -1394,8 +1394,8 @@ TEST_CASE(localized_sort)
         mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 0}}}), dot2);
     mm->add_return({add, transpose});
 
-    // Perform localized sort between dot1 and dot2
-    mm->localized_sort(dot1, dot2);
+    // Hoist external inputs between dot1 and dot2
+    mm->hoist_external_inputs(dot1, dot2);
 
     // Verify the module is still topologically sorted overall
     EXPECT(is_sorted(*mm));
@@ -1415,4 +1415,165 @@ TEST_CASE(localized_sort)
     EXPECT(std::distance(mm->begin(), dot2) < std::distance(mm->begin(), transpose));
 }
 
+TEST_CASE(hoist_external_inputs_no_movement_needed)
+{
+    // Test where external dependencies are already before the fusion chain
+    migraphx::shape s1{migraphx::shape::float_type, {2, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4}};
+    migraphx::shape s3{migraphx::shape::float_type, {2, 4}};
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    auto a = mm->add_parameter("a", s1);
+    auto b = mm->add_parameter("b", s2);
+    auto c = mm->add_parameter("c", s3);
+
+    // External operations already positioned before the fusion chain
+    auto external1 = add_pointwise(p, "main:pointwise0", {a}, single_pointwise("relu"));
+    auto external2 = add_pointwise(p, "main:pointwise1", {b}, single_pointwise("tanh"));
+    auto external3 = add_pointwise(p, "main:pointwise2", {c}, single_pointwise("tanh"));
+
+    // Fusion chain
+    auto dot1 = mm->add_instruction(migraphx::make_op("dot"), external1, external2);
+    auto add = add_pointwise(p, "main:pointwise3", {dot1, external3}, single_pointwise("add"));
+
+    mm->add_return({add});
+
+    // Record positions before hoist_external_inputs
+    auto dot1_pos_before = std::distance(mm->begin(), dot1);
+    auto add_pos_before = std::distance(mm->begin(), add);
+
+    mm->hoist_external_inputs(dot1, add);
+
+    // Verify positions haven't changed (nothing needed to move)
+    EXPECT(std::distance(mm->begin(), dot1) == dot1_pos_before);
+    EXPECT(std::distance(mm->begin(), add) == add_pos_before);
+    EXPECT(is_sorted(*mm));
+}
+
+TEST_CASE(hoist_external_inputs_multiple_external_branches)
+{
+    // Test with multiple independent external branches that need to be moved
+    migraphx::shape s1{migraphx::shape::float_type, {2, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4}};
+    migraphx::shape s3{migraphx::shape::float_type, {4, 3}};
+    migraphx::shape s4{migraphx::shape::float_type, {2, 4}};
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    auto a = mm->add_parameter("a", s1);
+    auto b = mm->add_parameter("b", s2);
+    auto c = mm->add_parameter("c", s4);
+    auto d = mm->add_parameter("d", s3);
+
+    // Start of fusion chain
+    auto dot1 = mm->add_instruction(migraphx::make_op("dot"), a, b); // {2, 4}
+
+    // External branch 1
+    auto ext1 = add_pointwise(p, "main:pointwise0", {c}, single_pointwise("relu")); // {2, 4}
+    auto ext2 = add_pointwise(p, "main:pointwise1", {ext1}, single_pointwise("neg")); // {2, 4}
+
+    // External branch 2
+    auto ext3 = add_pointwise(p, "main:pointwise2", {d}, single_pointwise("tanh")); // {4, 3}
+    auto ext4 = add_pointwise(p, "main:pointwise3", {ext3}, single_pointwise("abs")); // {4, 3}
+
+    // Continue fusion chain using external branches
+    auto add1 = add_pointwise(p, "main:pointwise4", {dot1, ext2}, single_pointwise("add")); // {2, 4}
+    auto dot2 = mm->add_instruction(migraphx::make_op("dot"), add1, ext4); // {2, 4} x {4, 3} = {2, 3}
+
+    mm->add_return({dot2});
+
+    mm->hoist_external_inputs(dot1, dot2);
+
+    EXPECT(is_sorted(*mm));
+
+    // Verify all external operations moved before dot1
+    EXPECT(std::distance(mm->begin(), ext1) < std::distance(mm->begin(), dot1));
+    EXPECT(std::distance(mm->begin(), ext2) < std::distance(mm->begin(), dot1));
+    EXPECT(std::distance(mm->begin(), ext3) < std::distance(mm->begin(), dot1));
+    EXPECT(std::distance(mm->begin(), ext4) < std::distance(mm->begin(), dot1));
+
+    // Verify fusion chain order preserved
+    EXPECT(std::distance(mm->begin(), dot1) < std::distance(mm->begin(), add1));
+    EXPECT(std::distance(mm->begin(), add1) < std::distance(mm->begin(), dot2));
+}
+
+TEST_CASE(hoist_external_inputs_long_fusion_chain)
+{
+    // Test with a longer fusion chain
+    migraphx::shape s{migraphx::shape::float_type, {4, 4}};
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    auto a = mm->add_parameter("a", s);
+    auto b = mm->add_parameter("b", s);
+    auto c = mm->add_parameter("c", s);
+    auto d = mm->add_parameter("d", s);
+
+    // Long fusion chain
+    auto op1 = mm->add_instruction(migraphx::make_op("dot"), a, b);
+
+    // External operations interspersed
+    auto ext1 = add_pointwise(p, "main:pointwise0", {d}, single_pointwise("exp"));
+
+    auto op2 = add_pointwise(p, "main:pointwise1", {op1, c}, single_pointwise("add"));
+
+    auto ext2 = add_pointwise(p, "main:pointwise2", {ext1}, single_pointwise("log"));
+
+    auto op3 = add_pointwise(p, "main:pointwise3", {op2}, single_pointwise("relu"));
+
+    auto ext3 = add_pointwise(p, "main:pointwise4", {ext2}, single_pointwise("abs"));
+
+    auto op4 = add_pointwise(p, "main:pointwise5", {op3}, single_pointwise("tanh"));
+    auto op5 = mm->add_instruction(migraphx::make_op("dot"), op4, ext3);
+
+    mm->add_return({op5});
+
+    mm->hoist_external_inputs(op1, op5);
+
+    EXPECT(is_sorted(*mm));
+
+    // All external operations moved before op1
+    EXPECT(std::distance(mm->begin(), ext1) < std::distance(mm->begin(), op1));
+    EXPECT(std::distance(mm->begin(), ext2) < std::distance(mm->begin(), op1));
+    EXPECT(std::distance(mm->begin(), ext3) < std::distance(mm->begin(), op1));
+
+    // Fusion chain order preserved
+    EXPECT(std::distance(mm->begin(), op1) < std::distance(mm->begin(), op2));
+    EXPECT(std::distance(mm->begin(), op2) < std::distance(mm->begin(), op3));
+    EXPECT(std::distance(mm->begin(), op3) < std::distance(mm->begin(), op4));
+    EXPECT(std::distance(mm->begin(), op4) < std::distance(mm->begin(), op5));
+}
+
+TEST_CASE(hoist_external_inputs_adjacent_instructions)
+{
+    // Test edge case where start and end are adjacent
+    migraphx::shape s1{migraphx::shape::float_type, {2, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4}};
+    migraphx::shape s3{migraphx::shape::float_type, {4, 5}};
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    auto a = mm->add_parameter("a", s1);  // {2, 3}
+    auto b = mm->add_parameter("b", s2);  // {3, 4}
+    auto c = mm->add_parameter("c", s3);  // {4, 5}
+
+    auto dot1 = mm->add_instruction(migraphx::make_op("dot"), a, b);  // {2, 3} x {3, 4} = {2, 4}
+
+    // External operation between dot1 and dot2
+    auto external = add_pointwise(p, "main:pointwise0", {c}, single_pointwise("relu"));  // {4, 5}
+
+    auto dot2 = mm->add_instruction(migraphx::make_op("dot"), dot1, external);  // {2, 4} x {4, 5} = {2, 5}
+
+    mm->add_return({dot2});
+
+    mm->hoist_external_inputs(dot1, dot2);
+
+    EXPECT(is_sorted(*mm));
+
+    // External should be moved before dot1
+    EXPECT(std::distance(mm->begin(), external) < std::distance(mm->begin(), dot1));
+    EXPECT(std::distance(mm->begin(), dot1) < std::distance(mm->begin(), dot2));
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }