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,77 @@
+/*
+ * 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/make_op.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/iterator_for.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void autocast_fp8_pass::apply(module& m) const
+{
+    for(auto ins : iterator_for(m))
+    {
+        const auto& ins_name = ins->name();
+        if(ins_name == "@param" && ins->get_shape().type() == fp8_type)
+        {
+            migraphx::shape new_shape = ins->get_shape().with_type(fp32_type);
+            std::string new_param_name =
+                ins->get_operator().to_value()["parameter"].to<std::string>() + "_fp32";
+            auto new_param = m.add_parameter(new_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);
+        }
+
+        if(ins_name == "@return")
+        {
+            std::vector<instruction_ref> orig_inputs = ins->inputs();
+            std::vector<instruction_ref> new_inputs;
+            bool is_input_fp8 = false;
+            for(const auto& i : orig_inputs)
+            {
+                if (i->get_shape().type() == fp8_type)
+                {
+                    is_input_fp8 = true;
+                    new_inputs.push_back(m.insert_instruction(
+                        ins,
+                        migraphx::make_op("convert", {{"target_type", migraphx::to_value(fp32_type)}}),
+                        i));
+                }
+
+            }
+            if(is_input_fp8)
+            {
+                auto new_ins = m.insert_instruction(ins, ins->get_operator(), {new_inputs});
+                m.replace_instruction(ins, new_ins);
+            }
+        }
+    }
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/dead_code_elimination.cpp

@@ -72,8 +72,7 @@ void dead_code_elimination::apply(module& m) const
                 leaf->clear_arguments();
                 assert(std::distance(m.begin(), leaf) < std::distance(m.begin(), last));
                 assert(leaf != ins);
-                if(leaf->name() != "@param")
-                    m.move_instruction(leaf, m.end());
+                m.move_instruction(leaf, m.end());
                 for(auto arg : args)
                     self(arg);
             }

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 will cast input paramater data types to fp8 and return data types to fp32.*/
+struct MIGRAPHX_EXPORT autocast_fp8_pass
+{
+    //std::vector<shape::type_t> orig_fp8_type;
+    shape::type_t fp8_type = migraphx::shape::fp8e4m3fnuz_type;
+    shape::type_t fp32_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/module.cpp

@@ -262,7 +262,6 @@ instruction_ref module::insert_instruction(instruction_ref ins,
                                            std::vector<instruction_ref> args)
 {
     assert(has_instruction(ins) or is_end(ins, this->end()));
-    assert(not starts_with(op.name(), "@"));
     shape r     = compute_shape(op, args);
     auto result = impl->insert(ins, {op, r, std::move(args)});
     instruction::backreference(result);

test/autocast_fp8.cpp

@@ -0,0 +1,142 @@
+/*
+ * 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/dead_code_elimination.hpp>
+#include <migraphx/eliminate_identity.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/permutation.hpp>
+#include <migraphx/ranges.hpp>
+#include <test.hpp>
+
+
+void run_pass(migraphx::module& m)
+{
+    migraphx::run_passes(m, {migraphx::autocast_fp8_pass{}, migraphx::dead_code_elimination{}, 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_fp32", {migraphx::shape::float_type, {1}});
+        auto x_fp32   = m2.add_parameter("x_fp32", {migraphx::shape::float_type, {1}});
+        auto y_fp8    = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(migraphx::shape::fp8e4m3fnuz_type)}}), y_fp32);
+        auto x_fp8    = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(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::to_value(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_fp32", {migraphx::shape::float_type, {1}});
+        auto x_fp32 = m2.add_parameter("x_fp32", {migraphx::shape::float_type, {1}});
+        auto y_fp8  = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(migraphx::shape::fp8e4m3fnuz_type)}}), y_fp32);
+        auto x_fp8  = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(migraphx::shape::fp8e4m3fnuz_type)}}), x_fp32);
+        m2.add_instruction(migraphx::make_op("sub"), x_fp8, y_fp8);
+    }
+    EXPECT(m1 == m2);
+}
+
+// multiple inputs 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);
+        auto result = m1.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), sum, diff);
+        m1.add_return({result});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto y_fp32      = m2.add_parameter("y_fp32", {migraphx::shape::float_type, {1}});
+        auto x_fp32      = m2.add_parameter("x_fp32", {migraphx::shape::float_type, {1}});
+        auto y_fp8       = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(migraphx::shape::fp8e4m3fnuz_type)}}), y_fp32);
+        auto x_fp8       = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(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 concat_fp8  = m2.add_instruction(migraphx::make_op("concat", {{"axis", 0}}), sum_fp8, diff_fp8);
+        auto result_fp32 = m2.add_instruction(migraphx::make_op("convert", {{"target_type", migraphx::to_value(migraphx::shape::float_type)}}), concat_fp8);
+        m2.add_return({result_fp32});
+    }
+    EXPECT(m1 == m2);
+}
+
+// autocast pass does not do any changes
+TEST_CASE(autocast_fp8_4)
+{
+    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});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto x   = m2.add_parameter("x", {migraphx::shape::float_type, {1}});
+        auto y   = m2.add_parameter("y", {migraphx::shape::float_type, {1}});
+        auto sum = m2.add_instruction(migraphx::make_op("add"), x, y);
+        m2.add_return({sum});
+
+    }
+    EXPECT(m1 == m2);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }