Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Catch invalid broadcasts in pointwise reduce fusion #3659

Merged
merged 8 commits into from
Nov 29, 2024
Merged

Catch invalid broadcasts in pointwise reduce fusion #3659

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2307be7
small test passing verify
shivadbhavsar Nov 26, 2024
ed7e04b
update condition
shivadbhavsar Nov 26, 2024
b418b44
update logic and existing tests
shivadbhavsar Nov 26, 2024
57b5ebd
tidy and format
shivadbhavsar Nov 26, 2024
a384375
add check to matchers
shivadbhavsar Nov 26, 2024
71cad8b
revert tests
shivadbhavsar Nov 26, 2024
301cd71
add unit tests
shivadbhavsar Nov 26, 2024
4433e9b
add verify test
shivadbhavsar Nov 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 42 additions & 1 deletion src/fuse_reduce.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,12 +174,53 @@ static auto any_input(Ms... ms)
return match::any_of[match::inputs()](match::any(ms...).bind("input"));
}

bool is_valid_broadcast(const instruction_ref b, const std::vector<size_t>& reduce_axes)
{
std::vector<size_t> broadcast_axes;
auto bstrides = b->get_shape().strides();

for(size_t i = 0; i < bstrides.size(); ++i)
{
if(bstrides.at(i) == 0)
broadcast_axes.push_back(i);
}

return broadcast_axes == reduce_axes;
}

template <class M>
static auto match_broadcast_axes(M m)
{
return match::make_basic_fun_matcher(
[=](match::matcher_context& ctx, instruction_ref ins) -> optional<instruction_ref> {
optional<instruction_ref> result = m.match(ctx, ins);
if(contains(ctx.instructions, "broadcast"))
{
instruction_ref reduce;
if(ins->get_operator().name() == "fused_reduce")
{
reduce = ins;
}
else
{
assert(contains(ctx.instructions, "reduce"));
reduce = ctx.instructions["reduce"];
}
auto axes = reduce->get_operator().to_value().at("axes").to_vector<size_t>();
auto broadcast = ctx.instructions["broadcast"];
if(not is_valid_broadcast(broadcast, axes))
return nullopt;
}
return result;
});
}

static auto match_broadcastable_input(const std::string& op, const std::string& name)
{
auto match_op = match::name(op)(used_once_except_broadcast()).bind(name);
auto match_op_input = any_input(match_op, match::used_once());
auto broadcast_match_op_input = any_input(match_broadcast(match_op), match::used_once());
return match::any_of(match_op_input, broadcast_match_op_input);
return match::any_of(match_op_input, match_broadcast_axes(broadcast_match_op_input));
}

static void finalize_reduce_module(module_ref m)
Expand Down
126 changes: 126 additions & 0 deletions test/fuse_reduce.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,44 @@ TEST_CASE(pointwise_reduce)
EXPECT(p1 == p2);
}

TEST_CASE(pointwise_reduce_unfusable_broadcast)
{
migraphx::shape s{migraphx::shape::float_type, {2, 1, 3}};
migraphx::program p1;
{
auto* mm = p1.get_main_module();
auto x = mm->add_parameter("x", s);
auto y = mm->add_parameter("y", s);
auto add = add_pointwise(p1, "main:pointwise0", {x, y}, single_pointwise("add"));
auto addb = mm->add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), add);
auto rsum = mm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {2}}}), addb);
mm->add_return({rsum});
}
run_pass(p1);

migraphx::program p2;
{
auto* mm = p2.get_main_module();
auto x = mm->add_parameter("x", s);
auto y = mm->add_parameter("y", s);
auto add = add_pointwise(p2, "main:pointwise0", {x, y}, single_pointwise("add"));
auto addb = mm->add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), add);
auto rsum =
add_reduce(p2,
"main:reduce_sum0",
{addb},
{2},
[&](auto* rm, const auto& inputs, const auto& axes) {
return rm->add_instruction(
migraphx::make_op("reduce_sum", {{"axes", axes}}), inputs[0]);
});
mm->add_return({rsum});
}
EXPECT(p1 == p2);
}

TEST_CASE(scalar_multibroadcast)
{
// Matches the find_pointwise_reduce matcher, but input x has a (scalar) shape
Expand Down Expand Up @@ -283,6 +321,43 @@ TEST_CASE(reduce_pointwise)
EXPECT(p1 == p2);
}

TEST_CASE(reduce_pointwise_unfusable_broadcast)
{
migraphx::shape s{migraphx::shape::float_type, {2, 1, 3}};
migraphx::program p1;
{
auto* mm = p1.get_main_module();
auto x = mm->add_parameter("x", s);
auto y = mm->add_parameter("y", s);
auto rsum = mm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {2}}}), x);
auto rsumb = mm->add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), rsum);
auto yb =
mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), y);
auto add = add_pointwise(p1, "main:pointwise0", {rsumb, yb}, single_pointwise("add"));
mm->add_return({add});
}
run_pass(p1);
migraphx::program p2;
{
auto* mm = p2.get_main_module();
auto x = mm->add_parameter("x", s);
auto y = mm->add_parameter("y", s);
auto rsum = add_reduce(
p2, "main:reduce_sum0", {x}, {2}, [&](auto* rm, const auto& inputs, const auto& axes) {
return rm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", axes}}),
inputs[0]);
});
auto rsumb = mm->add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), rsum);
auto yb =
mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), y);
auto add = add_pointwise(p2, "main:pointwise0", {rsumb, yb}, single_pointwise("add"));
mm->add_return({add});
}
EXPECT(p1 == p2);
}

TEST_CASE(reduce_reduce)
{
migraphx::shape s{migraphx::shape::float_type, {2, 3}};
Expand Down Expand Up @@ -325,6 +400,57 @@ TEST_CASE(reduce_reduce)
EXPECT(p1 == p2);
}

TEST_CASE(reduce_reduce_unfusable_broadcast)
{
migraphx::shape s{migraphx::shape::float_type, {2, 1, 3}};
migraphx::program p1;
{
auto* mm = p1.get_main_module();
auto x = mm->add_parameter("x", s);
auto rsum = mm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {2}}}), x);
auto rsumb = mm->add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), rsum);
auto xb =
mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), x);
auto rsumdiff = add_pointwise(p1, "main:pointwise0", {rsumb, xb}, single_pointwise("sub"));
auto rsum2 =
mm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {2}}}), rsumdiff);
auto sqrt = add_pointwise(p1, "main:pointwise1", {rsum2}, single_pointwise("sqrt"));
mm->add_return({sqrt});
}
run_pass(p1);
migraphx::program p2;
{
auto* mm = p2.get_main_module();
auto x = mm->add_parameter("x", s);
auto rsum = add_reduce(
p2, "main:reduce_sum0", {x}, {2}, [&](auto* rm, const auto& inputs, const auto& axes) {
return rm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", axes}}),
inputs[0]);
});

auto rsumb = mm->add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), rsum);
auto xb =
mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {2, 4, 3}}}), x);

auto sqrt = add_reduce(
p2,
"main:pointwise0:main:reduce_sum1:main:pointwise1",
{rsumb, xb},
{2},
[&](auto* rm, const auto& inputs, const auto& axes) {
auto rsumdiff = add_pointwise(
p2, rm, "main:pointwise0", {inputs[0], inputs[1]}, single_pointwise("sub"));
auto rsum2 = rm->add_instruction(migraphx::make_op("reduce_sum", {{"axes", axes}}),
rsumdiff);
return add_pointwise(p2, rm, "main:pointwise1", {rsum2}, single_pointwise("sqrt"));
});
mm->add_return({sqrt});
}
EXPECT(p1 == p2);
}

TEST_CASE(parallel_reduce_reduce)
{
migraphx::shape s{migraphx::shape::float_type, {2, 3}};
Expand Down
40 changes: 40 additions & 0 deletions test/include/layernorm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,3 +62,43 @@ inline migraphx::instruction_ref add_layernorm(migraphx::module& m,
m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), bias);
return m.add_instruction(migraphx::make_op("add"), mul, bias_mbcast);
}

inline migraphx::instruction_ref add_pointwise_layernorm(migraphx::module& m,
migraphx::instruction_ref x,
const std::vector<size_t>& dims,
float eps = 1e-12f)
{
auto mgx_type = x->get_shape().type();
auto scale = m.add_parameter("scale", migraphx::shape{mgx_type, {1, 1, dims.back()}});
auto bias = m.add_parameter("bias", migraphx::shape{mgx_type, {1, 1, dims.back()}});

auto epsilon = m.add_literal(migraphx::literal{migraphx::shape{mgx_type}, {eps}});
auto one = m.add_literal(migraphx::literal{migraphx::shape{mgx_type}, {1}});

auto mean = m.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {2}}}), x);
auto mean_mbcast =
m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), mean);
auto x_minus_mean = m.add_instruction(migraphx::make_op("sub"), x, mean_mbcast);
auto sqdiff = m.add_instruction(migraphx::make_op("sqdiff"), x, mean_mbcast);
auto var = m.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {2}}}), sqdiff);

auto epsilon_mbcast = m.add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", var->get_shape().lens()}}), epsilon);
auto var_stable = m.add_instruction(migraphx::make_op("add"), var, epsilon_mbcast);
auto inv_stddev_x = m.add_instruction(migraphx::make_op("rsqrt"), var_stable);
auto inv_stddev_x_mbcast =
m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), inv_stddev_x);
auto norm = m.add_instruction(migraphx::make_op("mul"), x_minus_mean, inv_stddev_x_mbcast);

auto one_mbcast = m.add_instruction(
migraphx::make_op("multibroadcast", {{"out_lens", scale->get_shape().lens()}}), one);
auto add_scale = m.add_instruction(migraphx::make_op("add"), scale, one_mbcast);
auto scale_mbcast =
m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), add_scale);
auto scale_norm = m.add_instruction(migraphx::make_op("mul"), norm, scale_mbcast);

auto bias_mbcast =
m.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), bias);

return m.add_instruction(migraphx::make_op("add"), scale_norm, bias_mbcast);
}
13 changes: 13 additions & 0 deletions test/verify/test_layernorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,3 +185,16 @@ struct test_add_layernorm_add_gemm_nonstd : verify_program<test_add_layernorm_ad
}
std::string section() const { return "gemm"; }
};

struct test_pw_layernorm : verify_program<test_pw_layernorm>
{
migraphx::program create_program() const
{
migraphx::program p;
auto* mm = p.get_main_module();
std::vector<size_t> dims = {1, 9, 6};
auto x = mm->add_parameter("x", migraphx::shape{migraphx::shape::float_type, dims});
add_pointwise_layernorm(*mm, x, dims);
return p;
}
};
Loading