diff --git a/src/fuse_horizontal.cpp b/src/fuse_horizontal.cpp index 60f9c41c2dc..62cb02a171b 100644 --- a/src/fuse_horizontal.cpp +++ b/src/fuse_horizontal.cpp @@ -158,13 +158,79 @@ static std::vector slice_gather_rows(module& m, return results; } +// Fuse a group of gathers whose indices may have *different* shapes. Each index is +// flattened to 1-D and concatenated, a single batched gather runs over `table`, then each +// element range is sliced back out and reshaped to the original gather's output shape +// (index dims + embedding dim). indices[i] is the (already offset-adjusted) index for +// gathers[i]; table is the gather data operand (a single, possibly concatenated table). +static std::vector +fuse_gathers_flattened(module& m, + const std::vector& gathers, + const std::vector& indices, + instruction_ref table, + instruction_ref insert_pt) +{ + std::vector flat_inputs(indices.size()); + std::transform(indices.begin(), indices.end(), flat_inputs.begin(), [&](instruction_ref idx) { + if(idx->get_shape().lens().size() == 1) + return idx; + std::int64_t n = idx->get_shape().elements(); + return m.insert_instruction(insert_pt, make_op("reshape", {{"dims", {n}}}), idx); + }); + + auto big_idx = m.insert_instruction(insert_pt, make_op("concat", {{"axis", 0}}), flat_inputs); + auto batched_gather = + m.insert_instruction(insert_pt, make_op("gather", {{"axis", 0}}), table, big_idx); + + // Inclusive prefix sum of element counts gives each range's end; shift right and + // prepend 0 for the (exclusive) start offsets into the batched gather. + std::vector slice_ends(gathers.size()); + transform_partial_sum( + gathers.begin(), gathers.end(), slice_ends.begin(), std::plus<>{}, [](auto g) { + return g->inputs().at(1)->get_shape().elements(); + }); + std::vector slice_starts(gathers.size()); + slice_starts[0] = 0; + std::copy(slice_ends.begin(), std::prev(slice_ends.end()), slice_starts.begin() + 1); + + const std::int64_t emb_dim = table->get_shape().lens().back(); + std::vector results(gathers.size()); + std::transform(gathers.begin(), + gathers.end(), + slice_starts.begin(), + results.begin(), + [&](instruction_ref g, std::size_t start) -> instruction_ref { + const auto& idx_lens = g->inputs().at(1)->get_shape().lens(); + std::size_t n = g->inputs().at(1)->get_shape().elements(); + auto sliced = m.insert_instruction( + insert_pt, + make_op("slice", + {{"axes", {0}}, + {"starts", {static_cast(start)}}, + {"ends", {static_cast(start + n)}}}), + batched_gather); + // A 1-D index already yields {n, emb_dim}; only multi-dim indices need a + // reshape back to (index dims + embedding dim). + if(idx_lens.size() == 1) + return sliced; + std::vector out_dims(idx_lens.begin(), idx_lens.end()); + out_dims.push_back(emb_dim); + return m.insert_instruction( + insert_pt, make_op("reshape", {{"dims", out_dims}}), sliced); + }); + + return results; +} + // --------------------------------------------------------------------------- // Same-table gather horizontal fusion // // Candidates: gather(axis=0) with 2D constant embedding table and a non-scalar // index whose first dim is >= min_index_batch (worthwhile batch) // Grouping: by (table instruction, index type, index trailing dims) so only -// gathers reading the *same* table are merged +// gathers reading the *same* table are merged. With merge_mixed_lengths +// the trailing dims are dropped, so same-table gathers with differing +// index shapes are merged too (via the flattened path). // Fusion: concatenate the indices, single batched gather, slice rows back. // No index offset adjustment is needed since the table is shared. // --------------------------------------------------------------------------- @@ -174,6 +240,11 @@ struct same_table_gather_horizontal_fusion // Minimum first index dim for the batched gather to be worthwhile. static constexpr std::size_t min_index_batch = 4; + // When set, merge same-table gathers whose indices differ only in shape by flattening + // each index (rather than requiring matching trailing dims). Because the table is + // shared, no index offset adjustment is introduced. + bool merge_mixed_lengths = false; + std::size_t min_group_size() const { return 2; } bool is_candidate(instruction_ref ins) const @@ -211,9 +282,12 @@ struct same_table_gather_horizontal_fusion auto idx_type = idx->get_shape().type(); const auto& lens = idx->get_shape().lens(); assert(not lens.empty()); - // Trailing index dims (all except first) — must match for concat on axis 0. + // Trailing index dims (all except first) — must match for concat on axis 0. When + // merging mixed lengths the indices are flattened, so trailing dims are ignored. // Keying on the data instruction itself restricts grouping to one table. - std::vector trailing(lens.begin() + 1, lens.end()); + std::vector trailing; + if(not merge_mixed_lengths) + trailing.assign(lens.begin() + 1, lens.end()); return std::make_tuple(data, idx_type, std::move(trailing)); } @@ -224,11 +298,25 @@ struct same_table_gather_horizontal_fusion auto data = gathers.front()->inputs().at(0); assert(data->get_shape().lens().size() == 2); - // Concatenate the per-gather indices into a single index tensor. + // Collect the per-gather indices (the table is shared, so no offset adjustment). std::vector idx_inputs(gathers.size()); std::transform(gathers.begin(), gathers.end(), idx_inputs.begin(), [](auto g) { return g->inputs().at(1); }); + + // Mixed index shapes cannot be concatenated on axis 0; flatten each instead. Only + // take the flattened path when the group is actually non-uniform, so uniform-shape + // groups produce the same (cheaper, reshape-free) IR regardless of the flag. + if(merge_mixed_lengths) + { + const auto& first_lens = gathers.front()->inputs().at(1)->get_shape().lens(); + bool uniform = std::all_of(gathers.begin(), gathers.end(), [&](auto g) { + return g->inputs().at(1)->get_shape().lens() == first_lens; + }); + if(not uniform) + return fuse_gathers_flattened(m, gathers, idx_inputs, data, insert_pt); + } + auto concat_idx = m.insert_instruction(insert_pt, make_op("concat", {{"axis", 0}}), idx_inputs); @@ -246,8 +334,9 @@ struct same_table_gather_horizontal_fusion // Candidates: gather(axis=0) with 2D constant embedding table, static shapes, // non-scalar index // Grouping: by (embedding dimension, index type, index trailing dims) -// Fusion: concatenate embedding tables, adjust indices with offsets, -// single batched gather, slice results back +// Fusion: concatenate the *distinct* embedding tables (shared tables are kept +// once), adjust indices with per-table offsets, single batched gather, +// slice results back // --------------------------------------------------------------------------- struct gather_horizontal_fusion @@ -296,51 +385,50 @@ struct gather_horizontal_fusion { auto idx_type = gathers.front()->inputs().at(1)->get_shape().type(); - // Concatenate all embedding tables - std::vector emb_inputs(gathers.size()); - std::transform(gathers.begin(), gathers.end(), emb_inputs.begin(), [](auto g) { - return g->inputs().at(0); - }); + // Deduplicate shared tables. Several gathers in a group may read the *same* + // table instruction (e.g. after same-table siblings are folded in). Emitting + // one table copy per gather would replicate that data, bloating both the + // concatenated table literal and the batched gather. Instead keep one copy per + // distinct table (in first-appearance order) and record the row offset of each. + std::vector unique_tables; + std::unordered_map table_offset; + std::size_t running_offset = 0; + for(auto g : gathers) + { + auto data = g->inputs().at(0); + if(table_offset.emplace(data, running_offset).second) + { + unique_tables.push_back(data); + running_offset += data->get_shape().lens().front(); + } + } + + // Concatenate the distinct tables (skip the concat when only one remains). auto concat_emb = - m.insert_instruction(insert_pt, make_op("concat", {{"axis", 0}}), emb_inputs); - - // Compute cumulative embedding offsets using transform_partial_sum. - // Inclusive partial sum gives end offsets; shift right and prepend 0 - // to get start (exclusive) offsets. - std::vector cum_sizes(gathers.size()); - transform_partial_sum( - gathers.begin(), gathers.end(), cum_sizes.begin(), std::plus<>{}, [](auto g) { - return g->inputs().at(0)->get_shape().lens().front(); - }); + unique_tables.size() == 1 + ? unique_tables.front() + : m.insert_instruction(insert_pt, make_op("concat", {{"axis", 0}}), unique_tables); - // Exclusive offsets: [0, cum_sizes[0], cum_sizes[1], ...] - std::vector emb_offsets(gathers.size()); - emb_offsets[0] = 0; - std::copy(cum_sizes.begin(), std::prev(cum_sizes.end()), emb_offsets.begin() + 1); - - // Build adjusted indices (add offset to shift into concatenated table) + // Build adjusted indices (add each gather's table offset to shift into the + // concatenated table). Gathers whose table lands at offset 0 need no adjustment. std::vector adjusted_idx_inputs; adjusted_idx_inputs.reserve(gathers.size()); - - migraphx::for_each( - gathers.begin(), gathers.end(), emb_offsets.begin(), [&](auto g, auto offset) { - auto idx = g->inputs().at(1); - if(offset == 0) - { - adjusted_idx_inputs.push_back(idx); - } - else - { - auto offset_scalar = m.add_literal(literal{shape{idx_type}, {offset}}); - auto offset_broadcast = m.insert_instruction( - insert_pt, - make_op("multibroadcast", {{"out_lens", idx->get_shape().lens()}}), - offset_scalar); - auto adjusted_idx = - m.insert_instruction(insert_pt, make_op("add"), idx, offset_broadcast); - adjusted_idx_inputs.push_back(adjusted_idx); - } - }); + std::transform(gathers.begin(), + gathers.end(), + std::back_inserter(adjusted_idx_inputs), + [&](auto g) -> instruction_ref { + auto idx = g->inputs().at(1); + auto offset = table_offset.at(g->inputs().at(0)); + if(offset == 0) + return idx; + auto offset_scalar = m.add_literal(literal{shape{idx_type}, {offset}}); + auto offset_broadcast = m.insert_instruction( + insert_pt, + make_op("multibroadcast", {{"out_lens", idx->get_shape().lens()}}), + offset_scalar); + return m.insert_instruction( + insert_pt, make_op("add"), idx, offset_broadcast); + }); // Concatenate adjusted indices auto concat_idx = @@ -372,9 +460,19 @@ void fuse_horizontal::apply(module_pass_manager& mpm) const { auto& m = mpm.get_module(); - // Collapse gathers that share the same table first; any sibling gathers left - // across *different* tables then fall through to cross-table fusion. - fuse_horizontal_ops(m, same_table_gather_horizontal_fusion{}, gather_horizontal_fusion{}); + // Run cross-embedding fusion first so a group of gathers sharing an embedding + // dimension is bundled together even when it spans several tables (table dedup + // keeps each distinct table once, so shared tables are not replicated). Running + // same-table fusion first would greedily collapse the same-table subset and strand + // the remaining siblings below cross-table fusion's group-size threshold. Same-table + // fusion then mops up the smaller same-instance groups (size 2-3) that cross-table + // fusion's larger threshold skips. merge_mixed_lengths only affects same-table fusion: + // it lets same-table gathers with differing index shapes merge (via a flattened index) + // at no offset-add cost, without over-merging unrelated tables. + fuse_horizontal_ops( + m, + gather_horizontal_fusion{}, + same_table_gather_horizontal_fusion{.merge_mixed_lengths = merge_mixed_lengths}); } } // namespace MIGRAPHX_INLINE_NS diff --git a/src/include/migraphx/fuse_horizontal.hpp b/src/include/migraphx/fuse_horizontal.hpp index fcac7b3151b..be5af30cc49 100644 --- a/src/include/migraphx/fuse_horizontal.hpp +++ b/src/include/migraphx/fuse_horizontal.hpp @@ -39,9 +39,16 @@ inline namespace MIGRAPHX_INLINE_NS { * embedding dimension and index layout, then fuses them into a single gather * over a concatenated embedding table with offset-adjusted indices. * The batched result is sliced back to produce the original outputs. + * + * When merge_mixed_lengths is set, same-table fusion also merges gathers on one + * table whose indices differ only in shape: each index is flattened to 1-D before + * the batched gather and each slice is reshaped back to its original output shape. + * Because the table is shared this introduces no index offset adjustment, and it + * never merges across distinct tables. Left off by default. */ struct MIGRAPHX_EXPORT fuse_horizontal { + bool merge_mixed_lengths = false; std::string name() const { return "fuse_horizontal"; } void apply(module_pass_manager& mpm) const; }; diff --git a/src/include/migraphx/simplify_reshapes.hpp b/src/include/migraphx/simplify_reshapes.hpp index 3ecfada3f36..c1056fe79d0 100644 --- a/src/include/migraphx/simplify_reshapes.hpp +++ b/src/include/migraphx/simplify_reshapes.hpp @@ -41,6 +41,10 @@ struct MIGRAPHX_EXPORT simplify_reshapes size_t depth = 4; bool enable_op_shape_transform_op = false; bool enable_gather_rewrite = false; + // Controls find_gather_slice_concat (rewrite of concat(slice(gather)) patterns). On by + // default; disabled only in the pre-fuse_horizontal GPU pass so the rewrite cannot reshape + // the gather/index graph that horizontal gather fusion groups on. + bool enable_gather_slice_concat = true; std::string name() const { return "simplify_reshapes"; } void apply(module& m) const; }; diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp index d5920a01814..a91467ffd4d 100644 --- a/src/simplify_reshapes.cpp +++ b/src/simplify_reshapes.cpp @@ -2092,6 +2092,7 @@ struct find_slice_squeeze m.replace_instruction(op_ins, new_sq); } }; + } // namespace void simplify_reshapes::apply(module& m) const @@ -2101,13 +2102,18 @@ void simplify_reshapes::apply(module& m) const if(enable_gather_rewrite) match::find_matches(m, find_gather{}); m.repeat_while_changes(depth, [&] { + // find_gather_slice_concat rewrites concat(slice(gather)) patterns into a gather with + // reordered indices. It is on by default but disabled in the pre-fuse_horizontal GPU + // pass, so it cannot reshape the gather/index graph that horizontal gather fusion groups + // on (running it before fusion can change which fusion path a lookup takes). + if(enable_gather_slice_concat) + match::find_matches(m, find_gather_slice_concat{}); match::find_matches(m, find_nop_reshapes{}, find_flatten{}, find_reshape_cont{}, find_slice_shape_transforms{}, find_nested_shape_transforms{}, - find_gather_slice_concat{}, find_concat_slice{}, find_concat_transpose{}, find_concat_reshape{}, diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp index 19ead8af09c..d01d5ea61bc 100644 --- a/src/targets/gpu/target.cpp +++ b/src/targets/gpu/target.cpp @@ -138,7 +138,7 @@ struct pipeline_factory eliminate_data_type_for_gpu{.disable_64bit = options.fast_math, .ctx = get_context()}, rewrite_resize{.affine_only = true}, dead_code_elimination{}, - simplify_reshapes{.enable_gather_rewrite = true}, + simplify_reshapes{.enable_gather_rewrite = true, .enable_gather_slice_concat = false}, eliminate_identity{}, eliminate_pad{}, dead_code_elimination{}, @@ -159,7 +159,8 @@ struct pipeline_factory optimize_module{}, layout_convolution{.channels_last = enabled(MIGRAPHX_ENABLE_NHWC{})}, dead_code_elimination{}, - enable_pass(disabled(MIGRAPHX_ENABLE_FULL_DYNAMIC{}), fuse_horizontal{}), + enable_pass(disabled(MIGRAPHX_ENABLE_FULL_DYNAMIC{}), + fuse_horizontal{.merge_mixed_lengths = true}), dead_code_elimination{}, prefuse_ops{get_context()}, dead_code_elimination{}, diff --git a/test/fuse_horizontal_test.cpp b/test/fuse_horizontal_test.cpp index 64816e23324..daa8b479e2d 100644 --- a/test/fuse_horizontal_test.cpp +++ b/test/fuse_horizontal_test.cpp @@ -476,6 +476,206 @@ TEST_CASE(gather_horiz_fusion_shared_index) EXPECT(m1 == m2); } +// Cross-table fusion where every gather reads the *same* constant table. +// Index batch (2) is below same_table's min_index_batch (4), so the group falls +// through to cross-table fusion. Table dedup must keep a single table copy (no +// concat / no offset adjustment) rather than replicating it once per gather. +TEST_CASE(gather_horiz_fusion_dedup_single_shared_table) +{ + migraphx::module m1; + { + auto emb = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {6, 2}}, 0)); + + auto idx1 = m1.add_parameter("idx1", {migraphx::shape::int32_type, {2}}); + auto idx2 = m1.add_parameter("idx2", {migraphx::shape::int32_type, {2}}); + auto idx3 = m1.add_parameter("idx3", {migraphx::shape::int32_type, {2}}); + auto idx4 = m1.add_parameter("idx4", {migraphx::shape::int32_type, {2}}); + + auto g1 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx1); + auto g2 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx2); + auto g3 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx3); + auto g4 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx4); + + m1.add_return({g1, g2, g3, g4}); + } + run_pass(m1); + + migraphx::module m2; + { + auto emb = + m2.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {6, 2}}, 0)); + + auto idx1 = m2.add_parameter("idx1", {migraphx::shape::int32_type, {2}}); + auto idx2 = m2.add_parameter("idx2", {migraphx::shape::int32_type, {2}}); + auto idx3 = m2.add_parameter("idx3", {migraphx::shape::int32_type, {2}}); + auto idx4 = m2.add_parameter("idx4", {migraphx::shape::int32_type, {2}}); + + // Single shared table → indices simply concatenated, gathered once, sliced back. + auto concat_idx = + m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), + std::vector{idx1, idx2, idx3, idx4}); + + auto bg = m2.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, concat_idx); + + auto s1 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}), bg); + auto s2 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {2}}, {"ends", {4}}}), bg); + auto s3 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {4}}, {"ends", {6}}}), bg); + auto s4 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {6}}, {"ends", {8}}}), bg); + + m2.add_return({s1, s2, s3, s4}); + } + EXPECT(m1.sort() == m2.sort()); +} + +// Cross-table fusion with two distinct tables, each read by two gathers. Dedup +// concatenates each table once ({3,2}+{4,2} = {7,2}); only the second table's +// gathers get a +3 offset. Without dedup the table would be replicated to {14,2}. +TEST_CASE(gather_horiz_fusion_dedup_two_shared_tables) +{ + migraphx::module m1; + { + auto emb_a = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {3, 2}}, 0)); + auto emb_b = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 2}}, 1)); + + auto idx_a1 = m1.add_parameter("idx_a1", {migraphx::shape::int32_type, {2}}); + auto idx_b1 = m1.add_parameter("idx_b1", {migraphx::shape::int32_type, {2}}); + auto idx_a2 = m1.add_parameter("idx_a2", {migraphx::shape::int32_type, {2}}); + auto idx_b2 = m1.add_parameter("idx_b2", {migraphx::shape::int32_type, {2}}); + + auto ga1 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb_a, idx_a1); + auto gb1 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb_b, idx_b1); + auto ga2 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb_a, idx_a2); + auto gb2 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb_b, idx_b2); + + m1.add_return({ga1, gb1, ga2, gb2}); + } + run_pass(m1); + + migraphx::module m2; + { + auto emb_a = + m2.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {3, 2}}, 0)); + auto emb_b = + m2.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {4, 2}}, 1)); + + auto idx_a1 = m2.add_parameter("idx_a1", {migraphx::shape::int32_type, {2}}); + auto idx_b1 = m2.add_parameter("idx_b1", {migraphx::shape::int32_type, {2}}); + auto idx_a2 = m2.add_parameter("idx_a2", {migraphx::shape::int32_type, {2}}); + auto idx_b2 = m2.add_parameter("idx_b2", {migraphx::shape::int32_type, {2}}); + + // Distinct tables concatenated once: [3+4, 2] = [7, 2] + auto concat_emb = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), + std::vector{emb_a, emb_b}); + + // Table B lands at offset 3; each of its gathers is shifted independently. + auto offb1 = m2.add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int32_type}, {std::size_t(3)}}); + auto bcb1 = + m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2}}}), offb1); + auto adj_b1 = m2.add_instruction(migraphx::make_op("add"), idx_b1, bcb1); + + auto offb2 = m2.add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int32_type}, {std::size_t(3)}}); + auto bcb2 = + m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2}}}), offb2); + auto adj_b2 = m2.add_instruction(migraphx::make_op("add"), idx_b2, bcb2); + + auto concat_idx = m2.add_instruction( + migraphx::make_op("concat", {{"axis", 0}}), + std::vector{idx_a1, adj_b1, idx_a2, adj_b2}); + + auto bg = + m2.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), concat_emb, concat_idx); + + auto s1 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}), bg); + auto s2 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {2}}, {"ends", {4}}}), bg); + auto s3 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {4}}, {"ends", {6}}}), bg); + auto s4 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {6}}, {"ends", {8}}}), bg); + + m2.add_return({s1, s2, s3, s4}); + } + EXPECT(m1.sort() == m2.sort()); +} + +static void run_pass_mixed(migraphx::module& m) +{ + migraphx::run_passes(m, + {migraphx::fuse_horizontal{.merge_mixed_lengths = true}, + migraphx::dead_code_elimination{}}); +} + +// merge_mixed_lengths: four gathers on the *same* table with different index shapes (so the +// default same-table finder, which needs matching trailing dims, would not group them all) +// are merged by flattening each index to 1-D, one batched gather, then slice + reshape back. +// Index first dims are >= min_index_batch (4) so they qualify for same-table fusion. +TEST_CASE(gather_horiz_fusion_mixed_index_lengths) +{ + migraphx::module m1; + { + auto emb = + m1.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {6, 2}}, 0)); + auto idx1 = m1.add_parameter("idx1", {migraphx::shape::int32_type, {4, 3}}); + auto idx2 = m1.add_parameter("idx2", {migraphx::shape::int32_type, {5, 2}}); + auto idx3 = m1.add_parameter("idx3", {migraphx::shape::int32_type, {4, 4}}); + auto idx4 = m1.add_parameter("idx4", {migraphx::shape::int32_type, {6, 2}}); + + auto g1 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx1); + auto g2 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx2); + auto g3 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx3); + auto g4 = m1.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, idx4); + + m1.add_return({g1, g2, g3, g4}); + } + run_pass_mixed(m1); + + migraphx::module m2; + { + auto emb = + m2.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {6, 2}}, 0)); + auto idx1 = m2.add_parameter("idx1", {migraphx::shape::int32_type, {4, 3}}); + auto idx2 = m2.add_parameter("idx2", {migraphx::shape::int32_type, {5, 2}}); + auto idx3 = m2.add_parameter("idx3", {migraphx::shape::int32_type, {4, 4}}); + auto idx4 = m2.add_parameter("idx4", {migraphx::shape::int32_type, {6, 2}}); + + // Each 2-D index flattened to 1-D (element counts 12, 10, 16, 12). + auto f1 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {12}}}), idx1); + auto f2 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {10}}}), idx2); + auto f3 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {16}}}), idx3); + auto f4 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {12}}}), idx4); + + auto big_idx = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), + std::vector{f1, f2, f3, f4}); + auto bg = m2.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb, big_idx); + + auto s1 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {12}}}), bg); + auto r1 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {4, 3, 2}}}), s1); + auto s2 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {12}}, {"ends", {22}}}), bg); + auto r2 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {5, 2, 2}}}), s2); + auto s3 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {22}}, {"ends", {38}}}), bg); + auto r3 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {4, 4, 2}}}), s3); + auto s4 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {38}}, {"ends", {50}}}), bg); + auto r4 = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {6, 2, 2}}}), s4); + + m2.add_return({r1, r2, r3, r4}); + } + EXPECT(m1.sort() == m2.sort()); +} + // Dependent gathers: g2 depends on g1's output → only independent ones fuse // Since g1→g2 dependency exists, group_by won't group them together. // With only 3 remaining independent gathers, below min_group_size=4, no fusion. @@ -848,31 +1048,41 @@ TEST_CASE(same_table_gathers_multiple_tables) auto idx_b1 = m2.add_parameter("idx_b1", {migraphx::shape::int32_type, {4}}); auto idx_b2 = m2.add_parameter("idx_b2", {migraphx::shape::int32_type, {5}}); - // Table A is fused first (anchor sees ga1) - auto concat_idx_a = - m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), - std::vector{idx_a1, idx_a2}); + // Cross-embedding fusion runs first and, because both tables share the same + // embedding dim, bundles all four gathers into one batched gather. Table dedup + // keeps each table once: concat is [3+4, 2] = [7, 2], and only table B's indices + // are shifted by +3 (table A's offset is 0). + auto concat_emb = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), + std::vector{emb_a, emb_b}); - auto bg_a = - m2.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb_a, concat_idx_a); + auto offb1 = m2.add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int32_type}, {std::size_t(3)}}); + auto bcb1 = + m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {4}}}), offb1); + auto adj_b1 = m2.add_instruction(migraphx::make_op("add"), idx_b1, bcb1); - auto sa1 = m2.add_instruction( - migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {4}}}), bg_a); - auto sa2 = m2.add_instruction( - migraphx::make_op("slice", {{"axes", {0}}, {"starts", {4}}, {"ends", {9}}}), bg_a); + auto offb2 = m2.add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int32_type}, {std::size_t(3)}}); + auto bcb2 = + m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {5}}}), offb2); + auto adj_b2 = m2.add_instruction(migraphx::make_op("add"), idx_b2, bcb2); - // Table B is fused next - auto concat_idx_b = - m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), - std::vector{idx_b1, idx_b2}); + // Indices concatenated in gather (position) order: ga1, gb1, ga2, gb2. + auto concat_idx = m2.add_instruction( + migraphx::make_op("concat", {{"axis", 0}}), + std::vector{idx_a1, adj_b1, idx_a2, adj_b2}); - auto bg_b = - m2.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), emb_b, concat_idx_b); + auto bg = + m2.add_instruction(migraphx::make_op("gather", {{"axis", 0}}), concat_emb, concat_idx); + auto sa1 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {4}}}), bg); auto sb1 = m2.add_instruction( - migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {4}}}), bg_b); + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {4}}, {"ends", {8}}}), bg); + auto sa2 = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {8}}, {"ends", {13}}}), bg); auto sb2 = m2.add_instruction( - migraphx::make_op("slice", {{"axes", {0}}, {"starts", {4}}, {"ends", {9}}}), bg_b); + migraphx::make_op("slice", {{"axes", {0}}, {"starts", {13}}, {"ends", {18}}}), bg); m2.add_return({sa1, sb1, sa2, sb2}); }