promote_literals pass (#1593)

Adds the promote_literals compiler pass that moves literals from the submodules to the main module.
With the eliminate_common_subexpression pass, it will remove copies of literals created during split_single_dyn_dim.
Pass is enabled with the split_single_dyn_dim compile option.

src/CMakeLists.txt

@@ -73,6 +73,7 @@ add_library(migraphx
     process.cpp
     program.cpp
     propagate_constant.cpp
+    promote_literals.cpp
     quantization.cpp
     quantize_fp16.cpp
     quantize_int8.cpp

src/include/migraphx/pass_manager.hpp

@@ -39,6 +39,7 @@ struct module_pass_manager
     virtual module& get_module()                           = 0;
     virtual module* create_module(const std::string& name) = 0;
     virtual module* get_common_parent()                    = 0;
+    virtual module* get_root_module()                      = 0;
     virtual void run_pass(const pass& p)                   = 0;
 
     protected:

src/include/migraphx/promote_literals.hpp

+/*
+ * 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_RTGLIB_PROMOTE_LITERALS_HPP
+#define MIGRAPHX_GUARD_RTGLIB_PROMOTE_LITERALS_HPP
+
+#include <string>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+/**
+ * Replace literals in submodules with literals in the root module.
+ * Intended to allow for reuse of the literals between submodules.
+ */
+struct promote_literals
+{
+    std::string name() const { return "promote_literals"; }
+    void apply(module_pass_manager&) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/split_single_dyn_dim.hpp

@@ -25,7 +25,7 @@
 #define MIGRAPHX_GUARD_RTGLIB_SPLIT_SINGLE_DYN_DIM_HPP
 
 #include <string>
-#include <migraphx/program.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/instruction_ref.hpp>
 #include <migraphx/config.hpp>
 

src/pass_manager.cpp

@@ -86,12 +86,21 @@ struct module_pm : module_pass_manager
         assert(mod);
         return *mod;
     }
+
     virtual module* create_module(const std::string& name) override
     {
         assert(prog);
         return prog->create_module(name);
     }
+
     virtual module* get_common_parent() override { return common_parent; }
+
+    virtual module* get_root_module() override
+    {
+        assert(prog);
+        return prog->get_main_module();
+    }
+
     virtual void run_pass(const pass& p) override
     {
         assert(mod);

src/promote_literals.cpp

+/*
+ * 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/promote_literals.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/module.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void promote_literals::apply(module_pass_manager& mpm) const
+{
+    module& m              = mpm.get_module();
+    module_ref root_module = mpm.get_root_module();
+    if(m.name() == "main")
+        return;
+
+    for(auto ins : iterator_for(m))
+    {
+        if(ins->name() == "@literal")
+        {
+            auto new_lit = root_module->add_literal(ins->get_literal());
+            for(auto out_ins : ins->outputs())
+            {
+                out_ins->replace_argument(out_ins, ins, new_lit);
+            }
+        }
+    }
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/target.cpp

@@ -26,7 +26,6 @@
 #include <migraphx/check_context.hpp>
 #include <migraphx/dead_code_elimination.hpp>
 #include <migraphx/eliminate_allocation.hpp>
-#include <migraphx/eliminate_common_subexpression.hpp>
 #include <migraphx/eliminate_concat.hpp>
 #include <migraphx/eliminate_contiguous.hpp>
 #include <migraphx/eliminate_data_type.hpp>
@@ -40,7 +39,7 @@
 #include <migraphx/normalize_ops.hpp>
 #include <migraphx/optimize_module.hpp>
 #include <migraphx/preallocate_param.hpp>
-#include <migraphx/propagate_constant.hpp>
+#include <migraphx/promote_literals.hpp>
 #include <migraphx/register_target.hpp>
 #include <migraphx/replace_allocate.hpp>
 #include <migraphx/rewrite_gelu.hpp>
@@ -48,7 +47,6 @@
 #include <migraphx/rewrite_quantization.hpp>
 #include <migraphx/rewrite_rnn.hpp>
 #include <migraphx/schedule.hpp>
-#include <migraphx/simplify_algebra.hpp>
 #include <migraphx/simplify_qdq.hpp>
 #include <migraphx/simplify_reshapes.hpp>
 #include <migraphx/split_single_dyn_dim.hpp>
@@ -150,6 +148,8 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
         dead_code_elimination{},
         compile_ops{&ctx},
         dead_code_elimination{},
+        promote_literals{},
+        dead_code_elimination{},
         write_literals{&ctx},
         schedule{gpu::schedule_model{ctx.get_current_device().nstreams()}, not enabled(MIGRAPHX_DISABLE_SCHEDULE_PASS{})},
         memory_coloring{"hip::allocate"},

test/promote_literals_test.cpp

+/*
+ * 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/promote_literals.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/eliminate_common_subexpression.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/serialize.hpp>
+#include <test.hpp>
+
+void run_promote(migraphx::program& p)
+{
+    migraphx::run_passes(p, {migraphx::promote_literals{}, migraphx::dead_code_elimination{}});
+}
+
+void run_promote_and_ecs(migraphx::program& p)
+{
+    migraphx::run_passes(p,
+                         {migraphx::promote_literals{},
+                          migraphx::dead_code_elimination{},
+                          migraphx::eliminate_common_subexpression{},
+                          migraphx::dead_code_elimination{}});
+}
+
+TEST_CASE(promote_only)
+{
+    migraphx::program p0;
+    {
+        auto* mm0 = p0.get_main_module();
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p0.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+            auto literal_ins = submod->add_literal(migraphx::literal{lit_s, {6}});
+            auto broadcast_lit =
+                submod->add_instruction(migraphx::make_op("multibroadcast"), literal_ins, sm_input);
+            auto add_ins =
+                submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit);
+            submod->add_return({add_ins});
+            return submod;
+        };
+        auto* dim1 = create_submodule(1, "dim_1");
+        auto* dim2 = create_submodule(2, "dim_2");
+        auto* dim3 = create_submodule(3, "dim_3");
+        auto* dim4 = create_submodule(4, "dim_4");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input0                             = mm0->add_parameter("data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm0->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input0},
+            {dim1, dim2, dim3, dim4});
+        auto ret =
+            mm0->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm0->add_return({ret});
+    }
+    run_promote(p0);
+
+    migraphx::program p1;
+    {
+        auto* mm1 = p1.get_main_module();
+        migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+        auto literal_ins3 = mm1->add_literal(migraphx::literal{lit_s, {6}});
+        auto literal_ins2 = mm1->add_literal(migraphx::literal{lit_s, {6}});
+        auto literal_ins1 = mm1->add_literal(migraphx::literal{lit_s, {6}});
+        auto literal_ins0 = mm1->add_literal(migraphx::literal{lit_s, {6}});
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size,
+                                    migraphx::instruction_ref lit,
+                                    const std::string& module_name) {
+            auto* submod = p1.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            auto broadcast_lit =
+                submod->add_instruction(migraphx::make_op("multibroadcast"), lit, sm_input);
+            auto add_ins =
+                submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit);
+            submod->add_return({add_ins});
+            return submod;
+        };
+        auto* dim1 = create_submodule(1, literal_ins0, "dim_1");
+        auto* dim2 = create_submodule(2, literal_ins1, "dim_2");
+        auto* dim3 = create_submodule(3, literal_ins2, "dim_3");
+        auto* dim4 = create_submodule(4, literal_ins3, "dim_4");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input0 = mm1->insert_parameter(std::next(literal_ins3), "data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm1->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input0},
+            {dim1, dim2, dim3, dim4});
+        auto ret =
+            mm1->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm1->add_return({ret});
+    }
+
+    EXPECT(p0 == p1);
+}
+
+TEST_CASE(promote_and_ecs0)
+{
+    migraphx::program p0;
+    {
+        auto* mm0 = p0.get_main_module();
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p0.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+            auto literal_ins = submod->add_literal(migraphx::literal{lit_s, {6}});
+            auto broadcast_lit =
+                submod->add_instruction(migraphx::make_op("multibroadcast"), literal_ins, sm_input);
+            auto add_ins =
+                submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit);
+            submod->add_return({add_ins});
+            return submod;
+        };
+        auto* dim1 = create_submodule(1, "dim_1");
+        auto* dim2 = create_submodule(2, "dim_2");
+        auto* dim3 = create_submodule(3, "dim_3");
+        auto* dim4 = create_submodule(4, "dim_4");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input0                             = mm0->add_parameter("data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm0->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input0},
+            {dim1, dim2, dim3, dim4});
+        auto ret =
+            mm0->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm0->add_return({ret});
+    }
+    run_promote_and_ecs(p0);
+
+    migraphx::program p1;
+    {
+        auto* mm1 = p1.get_main_module();
+        migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+        auto literal_ins = mm1->add_literal(migraphx::literal{lit_s, {6}});
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p1.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            auto broadcast_lit =
+                submod->add_instruction(migraphx::make_op("multibroadcast"), literal_ins, sm_input);
+            auto add_ins =
+                submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit);
+            submod->add_return({add_ins});
+            return submod;
+        };
+        auto* dim1 = create_submodule(1, "dim_1");
+        auto* dim2 = create_submodule(2, "dim_2");
+        auto* dim3 = create_submodule(3, "dim_3");
+        auto* dim4 = create_submodule(4, "dim_4");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input0 = mm1->insert_parameter(std::next(literal_ins), "data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm1->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input0},
+            {dim1, dim2, dim3, dim4});
+        auto ret =
+            mm1->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm1->add_return({ret});
+    }
+
+    EXPECT(p0 == p1);
+}
+
+TEST_CASE(promote_and_ecs1)
+{
+    migraphx::program p0;
+    {
+        auto* mm0 = p0.get_main_module();
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p0.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+            auto literal_ins0   = submod->add_literal(migraphx::literal{lit_s, {6}});
+            auto literal_ins1   = submod->add_literal(migraphx::literal{lit_s, {2}});
+            auto broadcast_lit0 = submod->add_instruction(
+                migraphx::make_op("multibroadcast"), literal_ins0, sm_input);
+            auto broadcast_lit1 = submod->add_instruction(
+                migraphx::make_op("multibroadcast"), literal_ins1, sm_input);
+            auto add_ins =
+                submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit0);
+            auto mul_ins =
+                submod->add_instruction(migraphx::make_op("mul"), add_ins, broadcast_lit1);
+            submod->add_return({mul_ins});
+            return submod;
+        };
+        auto* dim1 = create_submodule(1, "dim_1");
+        auto* dim2 = create_submodule(2, "dim_2");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input0                             = mm0->add_parameter("data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm0->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input0},
+            {dim1, dim2});
+        auto ret =
+            mm0->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm0->add_return({ret});
+    }
+    run_promote_and_ecs(p0);
+
+    migraphx::program p1;
+    {
+        auto* mm1 = p1.get_main_module();
+        migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+        auto literal_ins1 = mm1->add_literal(migraphx::literal{lit_s, {2}});
+        auto literal_ins0 = mm1->add_literal(migraphx::literal{lit_s, {6}});
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p1.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+            auto sm_input       = submod->add_parameter("data", sm_shape);
+            auto broadcast_lit0 = submod->add_instruction(
+                migraphx::make_op("multibroadcast"), literal_ins0, sm_input);
+            auto broadcast_lit1 = submod->add_instruction(
+                migraphx::make_op("multibroadcast"), literal_ins1, sm_input);
+            auto add_ins =
+                submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit0);
+            auto mul_ins =
+                submod->add_instruction(migraphx::make_op("mul"), add_ins, broadcast_lit1);
+            submod->add_return({mul_ins});
+            return submod;
+        };
+        auto* dim1 = create_submodule(1, "dim_1");
+        auto* dim2 = create_submodule(2, "dim_2");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input0 = mm1->insert_parameter(std::next(literal_ins1), "data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm1->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input0},
+            {dim1, dim2});
+        auto ret =
+            mm1->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm1->add_return({ret});
+    }
+
+    EXPECT(p0 == p1);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }