"vscode:/vscode.git/clone" did not exist on "325d0909a13d2edd0a10c9974e4309cfd6a50438"
Commit 2a2f82c7 authored by Ahsan Saghir's avatar Ahsan Saghir
Browse files

Address review: used contains to check fp8_types set

parent 15df0bc8
...@@ -26,28 +26,24 @@ ...@@ -26,28 +26,24 @@
#include <migraphx/program.hpp> #include <migraphx/program.hpp>
#include <migraphx/instruction.hpp> #include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp> #include <migraphx/iterator_for.hpp>
#include <migraphx/ranges.hpp>
namespace migraphx { namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS { inline namespace MIGRAPHX_INLINE_NS {
bool autocast_fp8_pass::is_fp8_type(const shape::type_t& s) const
{
return fp8_types.find(s) != fp8_types.end();
}
void autocast_fp8_pass::apply(module& m) const void autocast_fp8_pass::apply(module& m) const
{ {
std::vector<instruction_ref> remove_parameters; std::vector<instruction_ref> remove_parameters;
for(auto ins : iterator_for(m)) for(auto ins : iterator_for(m))
{ {
const auto& ins_name = ins->name(); const auto& ins_name = ins->name();
if(ins_name == "@param" and is_fp8_type(ins->get_shape().type())) if(ins_name == "@param" and contains(fp8_types, ins->get_shape().type()))
{ {
shape::type_t fp8_type = ins->get_shape().type(); shape::type_t fp8_type = ins->get_shape().type();
migraphx::shape new_shape = ins->get_shape().with_type(target_type); migraphx::shape new_shape = ins->get_shape().with_type(target_type);
std::string new_param_name = std::string param_name =
ins->get_operator().to_value()["parameter"].to<std::string>(); ins->get_operator().to_value()["parameter"].to<std::string>();
auto new_param = m.add_parameter(new_param_name, new_shape); auto new_param = m.add_parameter(param_name, new_shape);
auto new_ins = m.insert_instruction( auto new_ins = m.insert_instruction(
ins, ins,
migraphx::make_op("convert", {{"target_type", migraphx::to_value(fp8_type)}}), new_param); migraphx::make_op("convert", {{"target_type", migraphx::to_value(fp8_type)}}), new_param);
...@@ -61,7 +57,7 @@ void autocast_fp8_pass::apply(module& m) const ...@@ -61,7 +57,7 @@ void autocast_fp8_pass::apply(module& m) const
std::vector<instruction_ref> new_inputs; std::vector<instruction_ref> new_inputs;
for(const auto& i : inputs) for(const auto& i : inputs)
{ {
if (is_fp8_type(i->get_shape().type())) if (contains(fp8_types, i->get_shape().type()))
{ {
new_inputs.push_back(m.insert_instruction( new_inputs.push_back(m.insert_instruction(
ins, ins,
......
...@@ -33,14 +33,14 @@ struct program; ...@@ -33,14 +33,14 @@ struct program;
struct module; struct module;
/** /**
This will cast input paramater data types to fp8 and return data types to fp32.*/ This pass will convert model with fp8 input parameter to model with fp32
input parameter and internally add casts to fp8 for those converted params.*/
struct MIGRAPHX_EXPORT autocast_fp8_pass struct MIGRAPHX_EXPORT autocast_fp8_pass
{ {
std::set<shape::type_t> fp8_types = {migraphx::shape::fp8e4m3fnuz_type}; std::set<shape::type_t> fp8_types = {migraphx::shape::fp8e4m3fnuz_type};
shape::type_t target_type = migraphx::shape::float_type; shape::type_t target_type = migraphx::shape::float_type;
std::string name() const { return "autocast_fp8_pass"; } std::string name() const { return "autocast_fp8_pass"; }
void apply(module& m) const; void apply(module& m) const;
bool is_fp8_type(const shape::type_t& s) const;
}; };
} // namespace MIGRAPHX_INLINE_NS } // namespace MIGRAPHX_INLINE_NS
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment