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

Improvements to propagate_const for int4 #3494

Merged
merged 11 commits into from
Oct 3, 2024
Merged

Improvements to propagate_const for int4 #3494

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
e61af57
Add tests for constants
pfultz2 Sep 25, 2024
1b14c51
Update function
pfultz2 Sep 25, 2024
fe3e7b2
Add pack unpack test
pfultz2 Sep 25, 2024
04220d5
Add test for transpose
pfultz2 Oct 1, 2024
2ac2e14
Format
pfultz2 Oct 1, 2024
8a397bf
Add int4 tests
pfultz2 Oct 1, 2024
62a70cf
Format
pfultz2 Oct 1, 2024
0898b5e
Dont save as non-packed tensor
pfultz2 Oct 1, 2024
23bd0db
Format
pfultz2 Oct 1, 2024
fb71a35
Fix assertions
pfultz2 Oct 2, 2024
8865c6c
Format
pfultz2 Oct 2, 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
24 changes: 18 additions & 6 deletions src/propagate_constant.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,16 @@ MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_PROPAGATE_CONSTANT)

bool skip_propagate(instruction_ref ins)
{
if(ins->name() == "contiguous" or ins->name() == "dequantizelinear")
if(contains({"contiguous", "dequantizelinear", "reshape"}, ins->name()))
return skip_propagate(ins->inputs().front());
if(ins->name() == "unpack_int4")
return true;
auto&& s = ins->get_shape();
if(s.broadcasted() and not s.scalar())
return true;
if(s.scalar() and s.elements() != 1)
if(s.broadcasted() and s.element_space() < s.elements())
return true;
auto alias = instruction::get_output_alias(ins, true);
if(alias != ins)
return skip_propagate(alias);
return false;
}

Expand All @@ -56,6 +57,16 @@ bool is_const_ins(instruction_ref ins, const std::unordered_set<std::string>& sk
skip_ops.find(ins->name()) == skip_ops.end();
}

argument as_packed(const argument& c)
{
if(c.get_shape().packed())
return c;
auto s = c.get_shape().with_lens(c.get_shape().lens());
argument result;
c.visit([&](auto x) { result = literal{s, x.begin(), x.end()}.get_argument(); });
return result;
}

void propagate_constant::apply(module& m) const
{
std::unordered_set<instruction_ref> const_instrs;
Expand Down Expand Up @@ -92,7 +103,7 @@ void propagate_constant::apply(module& m) const
grainsize = const_instrs_vec.size() / n;
#endif
simple_par_for(const_instrs_vec.size(), grainsize, [&](const auto i) {
literals[i] = const_instrs_vec[i]->eval();
literals[i] = as_packed(const_instrs_vec[i]->eval());
});

// Replace instructions in m
Expand All @@ -113,7 +124,8 @@ void propagate_constant::apply(module& m) const
})(const_instrs_vec[i]);
m.debug_print(inss);
}
assert(literals[i].get_shape() == const_instrs_vec[i]->get_shape());
assert(literals[i].get_shape().lens() == const_instrs_vec[i]->get_shape().lens());
assert(literals[i].get_shape().bytes() <= const_instrs_vec[i]->get_shape().bytes());
auto l = m.add_literal(literals[i].get_shape(), literals[i].data());
m.replace_instruction(const_instrs_vec[i], l);
}
Expand Down
Loading
Loading