Autocast_fp8 pass

src/CMakeLists.txt

@@ -35,6 +35,7 @@ add_library(migraphx
     analyze_streams.cpp
     apply_alpha_beta.cpp
     argument.cpp
+    autocast_fp8.cpp
     auto_contiguous.cpp
     common.cpp
     common_dims.cpp

src/autocast_fp8.cpp

@@ -0,0 +1,81 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/autocast_fp8.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/ranges.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void autocast_fp8_pass::apply(module& m) const
+{
+    std::vector<instruction_ref> remove_parameters;
+    for(auto ins : iterator_for(m))
+    {
+        const auto& ins_name = ins->name();
+        if(ins_name == "@param" and contains(fp8_types, ins->get_shape().type()))
+        {
+            shape::type_t fp8_type    = ins->get_shape().type();
+            migraphx::shape new_shape = ins->get_shape().with_type(target_type);
+            std::string param_name = ins->get_operator().to_value()["parameter"].to<std::string>();
+            m.rename_parameter(ins, param_name + "_old");
+            auto new_param = m.add_parameter(param_name, new_shape);
+            auto new_ins   = m.insert_instruction(
+                ins,
+                migraphx::make_op("convert", {{"target_type", migraphx::to_value(fp8_type)}}),
+                new_param);
+            m.replace_instruction(ins, new_ins);
+            remove_parameters.push_back(ins);
+        }
+
+        if(ins_name == "@return")
+        {
+            std::vector<instruction_ref> inputs = ins->inputs();
+            std::vector<instruction_ref> new_inputs;
+            std::transform(
+                inputs.begin(), inputs.end(), std::back_inserter(new_inputs), [&](auto i) {
+                    if(contains(fp8_types, i->get_shape().type()))
+                    {
+                        return m.insert_instruction(
+                            ins,
+                            migraphx::make_op("convert",
+                                              {{"target_type", migraphx::to_value(target_type)}}),
+                            i);
+                    }
+                    else
+                        return i;
+                });
+            m.replace_return({new_inputs});
+        }
+    }
+    // Remove unused parameters with fp8 type
+    for(const auto& i : remove_parameters)
+        m.remove_instruction(i);
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/include/migraphx/autocast_fp8.hpp

@@ -0,0 +1,49 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#ifndef MIGRAPHX_GUARD_AMDMIGRAPHX_AUTOCAST_FP8_HPP
+#define MIGRAPHX_GUARD_AMDMIGRAPHX_AUTOCAST_FP8_HPP
+
+#include <migraphx/shape.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct program;
+struct module;
+
+/**
+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
+{
+    std::set<shape::type_t> fp8_types = {migraphx::shape::fp8e4m3fnuz_type};
+    shape::type_t target_type         = migraphx::shape::float_type;
+    std::string name() const { return "autocast_fp8_pass"; }
+    void apply(module& m) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/module.hpp

@@ -175,6 +175,8 @@ struct MIGRAPHX_EXPORT module
 
     instruction_ref get_parameter(std::string name) const;
 
+    void rename_parameter(instruction_ref ins, const std::string& name);
+
     std::unordered_map<std::string, shape> get_parameter_shapes() const;
 
     bool has_instruction(instruction_ref ins) const;

src/module.cpp

@@ -554,6 +554,17 @@ instruction_ref module::get_parameter(std::string name) const
         return this->end();
 }
 
+void module::rename_parameter(instruction_ref ins, const std::string& name)
+{
+    assert(ins->name() == "@param");
+    auto op      = any_cast<builtin::param>(ins->get_operator());
+    op.parameter = name;
+    auto outputs = ins->outputs();
+    *ins         = instruction{op, ins->get_shape(), {}};
+    for(auto output : outputs)
+        ins->add_output(output);
+}
+
 std::unordered_map<std::string, shape> module::get_parameter_shapes() const
 {
     std::unordered_map<std::string, shape> result;

test/autocast_fp8.cpp

@@ -0,0 +1,175 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <basic_ops.hpp>
+#include <migraphx/autocast_fp8.hpp>
+#include <migraphx/eliminate_identity.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <test.hpp>
+
+void run_pass(migraphx::module& m)
+{
+    migraphx::run_passes(m, {migraphx::autocast_fp8_pass{}, migraphx::eliminate_identity{}});
+}
+
+// with return
+TEST_CASE(autocast_fp8_1)
+{
+    migraphx::module m1;
+    {
+        auto x   = m1.add_parameter("x", {migraphx::shape::fp8e4m3fnuz_type, {1}});
+        auto y   = m1.add_parameter("y", {migraphx::shape::fp8e4m3fnuz_type, {1}});
+        auto sum = m1.add_instruction(migraphx::make_op("add"), x, y);
+        m1.add_return({sum});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto y_fp32 = m2.add_parameter("y", {migraphx::shape::float_type, {1}});
+        auto x_fp32 = m2.add_parameter("x", {migraphx::shape::float_type, {1}});
+        auto y_fp8  = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}),
+            y_fp32);
+        auto x_fp8 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}),
+            x_fp32);
+        auto sum_fp8  = m2.add_instruction(migraphx::make_op("add"), x_fp8, y_fp8);
+        auto sum_fp32 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::float_type}}), sum_fp8);
+        m2.add_return({sum_fp32});
+    }
+    EXPECT(m1 == m2);
+}
+
+// without return
+TEST_CASE(autocast_fp8_2)
+{
+    migraphx::module m1;
+    {
+        auto x = m1.add_parameter("x", {migraphx::shape::fp8e4m3fnuz_type, {1}});
+        auto y = m1.add_parameter("y", {migraphx::shape::fp8e4m3fnuz_type, {1}});
+        m1.add_instruction(migraphx::make_op("sub"), x, y);
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto y_fp32 = m2.add_parameter("y", {migraphx::shape::float_type, {1}});
+        auto x_fp32 = m2.add_parameter("x", {migraphx::shape::float_type, {1}});
+        auto y_fp8  = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}),
+            y_fp32);
+        auto x_fp8 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}),
+            x_fp32);
+        m2.add_instruction(migraphx::make_op("sub"), x_fp8, y_fp8);
+    }
+    EXPECT(m1 == m2);
+}
+
+// multiple inputs (of same type) to return
+TEST_CASE(autocast_fp8_3)
+{
+    migraphx::module m1;
+    {
+        auto x    = m1.add_parameter("x", {migraphx::shape::fp8e4m3fnuz_type, {1}});
+        auto y    = m1.add_parameter("y", {migraphx::shape::fp8e4m3fnuz_type, {1}});
+        auto sum  = m1.add_instruction(migraphx::make_op("add"), x, y);
+        auto diff = m1.add_instruction(migraphx::make_op("sub"), x, y);
+        m1.add_return({sum, diff});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto y_fp32 = m2.add_parameter("y", {migraphx::shape::float_type, {1}});
+        auto x_fp32 = m2.add_parameter("x", {migraphx::shape::float_type, {1}});
+        auto y_fp8  = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}),
+            y_fp32);
+        auto x_fp8 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}),
+            x_fp32);
+        auto sum_fp8  = m2.add_instruction(migraphx::make_op("add"), x_fp8, y_fp8);
+        auto diff_fp8 = m2.add_instruction(migraphx::make_op("sub"), x_fp8, y_fp8);
+        auto sum_fp32 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::float_type}}), sum_fp8);
+        auto diff_fp32 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::float_type}}), diff_fp8);
+        m2.add_return({sum_fp32, diff_fp32});
+    }
+    EXPECT(m1 == m2);
+}
+
+// multiple inputs (of different types) to return
+TEST_CASE(autocast_fp8_4)
+{
+    migraphx::module m1;
+    {
+        auto x1   = m1.add_parameter("x1", {migraphx::shape::fp8e4m3fnuz_type, {3, 4}, {1, 3}});
+        auto y1   = m1.add_parameter("y1", {migraphx::shape::fp8e4m3fnuz_type, {3, 4}, {1, 3}});
+        auto x2   = m1.add_parameter("x2", {migraphx::shape::float_type, {3, 4}, {1, 3}});
+        auto y2   = m1.add_parameter("y2", {migraphx::shape::float_type, {3, 4}, {1, 3}});
+        auto sum1 = m1.add_instruction(migraphx::make_op("add"), x1, y1);
+        auto sum2 = m1.add_instruction(migraphx::make_op("add"), x2, y2);
+        m1.add_return({sum1, sum2});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto x2     = m2.add_parameter("x2", {migraphx::shape::float_type, {3, 4}, {1, 3}});
+        auto y2     = m2.add_parameter("y2", {migraphx::shape::float_type, {3, 4}, {1, 3}});
+        auto y1     = m2.add_parameter("y1", {migraphx::shape::float_type, {3, 4}, {1, 3}});
+        auto x1     = m2.add_parameter("x1", {migraphx::shape::float_type, {3, 4}, {1, 3}});
+        auto y1_fp8 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}), y1);
+        auto x1_fp8 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::fp8e4m3fnuz_type}}), x1);
+        auto sum1        = m2.add_instruction(migraphx::make_op("add"), x1_fp8, y1_fp8);
+        auto sum2        = m2.add_instruction(migraphx::make_op("add"), x2, y2);
+        auto result_sum1 = m2.add_instruction(
+            migraphx::make_op("convert", {{"target_type", migraphx::shape::float_type}}), sum1);
+        m2.add_return({result_sum1, sum2});
+    }
+    EXPECT(m1 == m2);
+}
+
+// autocast pass does not do any changes
+TEST_CASE(autocast_fp8_5)
+{
+    migraphx::module m1;
+    {
+        auto x   = m1.add_parameter("x", {migraphx::shape::float_type, {1}});
+        auto y   = m1.add_parameter("y", {migraphx::shape::float_type, {1}});
+        auto sum = m1.add_instruction(migraphx::make_op("add"), x, y);
+        m1.add_return({sum});
+    }
+    migraphx::module m2 = m1;
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }