Select module op (#1569)

Creates the select_module operator that selects one of the submodules passed to it to run based on the submodule parameters.  The submodule is selected by having the exact same static shapes for the arguments to select_module as the parameters in the submodule

src/CMakeLists.txt

@@ -199,6 +199,7 @@ register_migraphx_ops(
     scatternd_add
     scatternd_mul
     scatternd_none
+    select_module
     sigmoid
     sign
     sinh

src/include/migraphx/module.hpp

@@ -54,6 +54,10 @@ using ins_dep_map   = std::unordered_map<instruction_ref, std::unordered_set<ins
  */
 struct module
 {
+    // used by replace_allocate pass
+    // allocate memory in this module rather than using output parmaeters
+    bool use_local_alloc = false;
+
     module(const std::string& name = "");
 
     // move constructor

src/include/migraphx/op/select_module.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_OPERATORS_SELECT_MODULE_HPP
+#define MIGRAPHX_GUARD_OPERATORS_SELECT_MODULE_HPP
+
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/module.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct select_module
+{
+    shape output_dyn_shapes;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.output_dyn_shapes, "output_dyn_shapes"));
+    }
+
+    std::string name() const { return "select_module"; }
+
+    shape compute_shape(const std::vector<shape>&, const std::vector<module_ref>&) const
+    {
+        return shape{output_dyn_shapes};
+    }
+
+    argument compute(const shape&,
+                     const std::vector<argument>& args,
+                     const std::vector<module_ref>& submodule_list,
+                     const std::function<std::vector<argument>(
+                         module_ref&, const std::unordered_map<std::string, argument>&)>& run) const
+    {
+        // find submodule with input parameter shapes exactly the same as the input arguments
+        // assuming arguments are in the same order as the input parameters
+        auto module_iter =
+            std::find_if(submodule_list.cbegin(), submodule_list.cend(), [&](module_ref mr) {
+                auto param_names = mr->get_parameter_names();
+                assert(param_names.size() <= args.size());
+                return std::equal(param_names.cbegin(),
+                                  param_names.cend(),
+                                  args.cbegin(),
+                                  [&](auto p_name, auto a) {
+                                      return a.get_shape() == mr->get_parameter_shape(p_name);
+                                  });
+            });
+
+        if(module_iter == submodule_list.end())
+        {
+            MIGRAPHX_THROW("SELECT_MODULE: no compatible submodules found for given input shapes");
+        }
+        auto* module_to_run = *module_iter;
+        std::unordered_map<std::string, argument> params;
+
+        // add input parameters
+        auto param_names = module_to_run->get_parameter_names();
+        assert(param_names.size() <= args.size());
+        std::transform(param_names.begin(),
+                       param_names.end(),
+                       args.begin(),
+                       std::inserter(params, params.end()),
+                       [](auto&& name, auto&& a) { return std::make_pair(name, a); });
+
+        auto results = run(module_to_run, params);
+        return argument{results};
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/operation.hpp

@@ -140,9 +140,9 @@ template <class T>
 auto compute_shape_op(rank<2>, const T& x, const std::vector<shape>& inputs)
     -> decltype(x.normalize_compute_shape(inputs))
 {
-    dependent_type<operation, T> y = x;
     if(inputs.empty())
         MIGRAPHX_THROW("At least one input is required for " + x.name());
+    dependent_type<operation, T> y = x;
     normalize_attributes(y, inputs[0].max_lens());
     return any_cast<T>(y).normalize_compute_shape(inputs);
 }

src/include/migraphx/shape.hpp

@@ -243,6 +243,9 @@ struct shape
     /// Return true if the shape is dynamic
     bool dynamic() const;
 
+    /// Return true if this shape or any of the sub_shapes are dynamic
+    bool any_of_dynamic() const;
+
     shape normalize_standard() const;
 
     shape with_lens(type_t t, const std::vector<std::size_t>& l) const;

src/program.cpp

@@ -379,7 +379,7 @@ std::vector<argument> generic_eval(const module* mod,
                             }));
         }
         assert(results.find(ins) != results.end());
-        if(not ins->get_shape().dynamic())
+        if(not ins->get_shape().any_of_dynamic())
         {
             assert(results.at(ins).get_shape() == ins->get_shape());
         }

src/replace_allocate.cpp

@@ -104,19 +104,17 @@ void replace_allocate::apply(module& m) const
             continue;
 
         auto s = ins->get_shape();
-
-        if(not main_offload_copy and model.needs_out_params() and contains(mod_output_names, ins))
+        if(not(main_offload_copy or m.use_local_alloc) and model.needs_out_params() and
+           contains(mod_output_names, ins))
         {
-
             auto out_param = m.add_parameter(mod_output_names[ins], s);
             m.replace_instruction(ins, out_param);
-            continue;
         }
-
-        m.replace_instruction(
-            ins,
-            m.insert_instruction(ins,
-                                 make_op(model.name(), migraphx::value{{"shape", to_value(s)}})));
+        else
+        {
+            m.replace_instruction(ins,
+                                  make_op(model.name(), migraphx::value{{"shape", to_value(s)}}));
+        }
     }
 }
 

src/shape.cpp

@@ -483,6 +483,17 @@ std::string shape::type_string() const { return name(this->type()); }
 
 bool shape::dynamic() const { return not impl->m_dyn_dims.empty(); }
 
+bool shape::any_of_dynamic() const
+{
+    if(this->dynamic())
+    {
+        return true;
+    }
+    return std::any_of(this->sub_shapes().cbegin(), this->sub_shapes().cend(), [](auto s) {
+        return s.any_of_dynamic();
+    });
+}
+
 const std::vector<shape::dynamic_dimension>& shape::dyn_dims() const { return impl->m_dyn_dims; }
 
 std::vector<std::size_t> shape::min_lens() const

src/targets/gpu/lowering.cpp

@@ -111,6 +111,7 @@ struct miopen_apply
         add_loop_op();
         add_neg_op();
         add_nms_op();
+        add_select_module_op();
     }
 
     void copy_params() const
@@ -358,6 +359,33 @@ struct miopen_apply
             return mod->replace_instruction(ins, gpu_out);
         });
     }
+
+    /**
+     * Turns on use_local_alloc in the select_module submodules.
+     * Changes the submodule returns to a hip::sync_stream.
+     */
+    void add_select_module_op()
+    {
+        apply_map.emplace("select_module", [=](instruction_ref ins) {
+            std::vector<instruction_ref> inputs = ins->inputs();
+            auto mod_args                       = ins->module_inputs();
+            for(auto* smod : mod_args)
+            {
+                smod->use_local_alloc = true;
+                auto last_ins         = std::prev(smod->end());
+                if(last_ins->name() == "@return")
+                {
+                    for(auto out_ins : last_ins->inputs())
+                    {
+                        auto sync_out = smod->insert_instruction(
+                            last_ins, make_op("hip::sync_stream"), out_ins);
+                        smod->replace_return({sync_out});
+                    }
+                }
+            }
+            return ins;
+        });
+    }
 };
 
 void lowering::apply(module& m) const { miopen_apply{&m, this}.apply(); }

test/op_shape_test.cpp

@@ -2361,6 +2361,18 @@ TEST_CASE(rnn)
     }
 }
 
+TEST_CASE(select_module_dyn)
+{
+    migraphx::shape input{migraphx::shape::float_type, {{1, 4}, {3, 3}, {255, 255}, {255, 255}}};
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {1000, 1000}}});
+    migraphx::shape out_attr = migraphx::shape{sub_shapes};
+    expect_shape(
+        out_attr,
+        migraphx::make_op("select_module", {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+        input);
+}
+
 TEST_CASE(slice_shape)
 {
     migraphx::shape input{migraphx::shape::int32_type, {2, 2, 3}};

test/ref_ops_test.cpp

@@ -7277,6 +7277,188 @@ TEST_CASE(scatternd_reduction_test)
     }
 }
 
+TEST_CASE(select_module_add_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+    auto literal_ins = mm->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 = p.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* batch1 = create_submodule(1, "batch_1");
+    auto* batch2 = create_submodule(2, "batch_2");
+    auto* batch3 = create_submodule(3, "batch_3");
+    auto* batch4 = create_submodule(4, "batch_4");
+
+    migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+    auto input                              = mm->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              = mm->add_instruction(
+        migraphx::make_op("select_module", {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+        {input},
+        {batch1, batch2, batch3, batch4});
+    auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+    mm->add_return({ret});
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data{-4, 8, -1, 4, -1, 8, 8, -4};
+    migraphx::parameter_map params;
+    migraphx::shape input_fixed_shape{migraphx::shape::float_type, {2, 4}};
+    params["data"] = migraphx::argument(input_fixed_shape, input_data.data());
+    auto result    = p.eval(params).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{2, 14, 5, 10, 5, 14, 14, 2};
+    EXPECT(migraphx::verify_range(results_vector, gold));
+}
+
+TEST_CASE(select_module_reduce_test0)
+{
+    migraphx::program p;
+
+    // create batch submodules
+    auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+        auto* submod = p.create_module(module_name);
+        migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 2, 2}};
+        auto sm_input = submod->add_parameter("data", sm_shape);
+        auto reduce_ins =
+            submod->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {1}}}), sm_input);
+        auto squeeze_ins =
+            submod->add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), reduce_ins);
+        submod->add_return({squeeze_ins});
+        return submod;
+    };
+    auto* batch1 = create_submodule(1, "batch_1");
+    auto* batch2 = create_submodule(2, "batch_2");
+    auto* batch3 = create_submodule(3, "batch_3");
+    auto* batch4 = create_submodule(4, "batch_4");
+
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {2, 2}, {2, 2}}};
+    auto input                              = mm->add_parameter("data", s);
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {2, 2}}});
+    migraphx::shape out_attr = migraphx::shape{sub_shapes};
+    auto sm_ins              = mm->add_instruction(
+        migraphx::make_op("select_module", {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+        {input},
+        {batch1, batch2, batch3, batch4});
+    auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+    mm->add_return({ret});
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data{-4, 8, -1, 4, -1, 8, 8, -4};
+    migraphx::parameter_map params;
+    migraphx::shape input_fixed_shape{migraphx::shape::float_type, {2, 2, 2}};
+    params["data"] = migraphx::argument(input_fixed_shape, input_data.data());
+    auto result    = p.eval(params).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{-5, 12, 7, 4};
+    EXPECT(migraphx::verify_range(results_vector, gold));
+}
+
+TEST_CASE(select_module_reduce_test1)
+{
+    migraphx::program p;
+
+    // create batch submodules
+    auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+        auto* submod = p.create_module(module_name);
+        migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 2, 2}};
+        auto sm_input = submod->add_parameter("data", sm_shape);
+        auto reduce_ins =
+            submod->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {1}}}), sm_input);
+        auto squeeze_ins =
+            submod->add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), reduce_ins);
+        submod->add_return({squeeze_ins});
+        return submod;
+    };
+    auto* batch1 = create_submodule(1, "batch_1");
+    auto* batch2 = create_submodule(2, "batch_2");
+    auto* batch3 = create_submodule(3, "batch_3");
+    auto* batch4 = create_submodule(4, "batch_4");
+
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {2, 2}, {2, 2}}};
+    auto input                              = mm->add_parameter("data", s);
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {2, 2}}});
+    migraphx::shape out_attr = migraphx::shape{sub_shapes};
+    auto sm_ins              = mm->add_instruction(
+        migraphx::make_op("select_module", {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+        {input},
+        {batch1, batch2, batch3, batch4});
+    auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+    mm->add_return({ret});
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data{-4, 8, -1, 4, -1, 8, 8, -4, -4, 8, -1, 4, -1, 8, 8, -4};
+    migraphx::parameter_map params;
+    migraphx::shape input_fixed_shape{migraphx::shape::float_type, {4, 2, 2}};
+    params["data"] = migraphx::argument(input_fixed_shape, input_data.data());
+    auto result    = p.eval(params).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{-5, 12, 7, 4, -5, 12, 7, 4};
+    EXPECT(migraphx::verify_range(results_vector, gold));
+}
+
+TEST_CASE(select_module_not_found_error)
+{
+    migraphx::program p;
+
+    // create batch submodules
+    auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+        auto* submod = p.create_module(module_name);
+        migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 2, 2}};
+        auto sm_input = submod->add_parameter("data", sm_shape);
+        auto reduce_ins =
+            submod->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {1}}}), sm_input);
+        auto squeeze_ins =
+            submod->add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), reduce_ins);
+        submod->add_return({squeeze_ins});
+        return submod;
+    };
+    auto* batch1 = create_submodule(1, "batch_1");
+    auto* batch2 = create_submodule(2, "batch_2");
+    auto* batch3 = create_submodule(3, "batch_3");
+    auto* batch4 = create_submodule(4, "batch_4");
+
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {2, 2}, {2, 2}}};
+    auto input                              = mm->add_parameter("data", s);
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {2, 2}}});
+    migraphx::shape out_attr = migraphx::shape{sub_shapes};
+    auto sm_ins              = mm->add_instruction(
+        migraphx::make_op("select_module", {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+        {input},
+        {batch1, batch2, batch3, batch4});
+    auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+    mm->add_return({ret});
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data{-4, 8, -1, 4, -1, 8,  8,  -4, -4, 8,
+                                  -1, 4, -1, 8, 8,  -4, -1, 8,  8,  -4};
+    migraphx::parameter_map params;
+    migraphx::shape input_fixed_shape{migraphx::shape::float_type, {5, 2, 2}};
+    params["data"] = migraphx::argument(input_fixed_shape, input_data.data());
+    EXPECT(test::throws([&] { p.eval(params).back(); }));
+}
+
 TEST_CASE(scatternd_reduction_dyn_test)
 {
     // reduction = add, with dynamic input shapes

test/shape_test.cpp

@@ -238,6 +238,30 @@ TEST_CASE(test_shape_dynamic_serialize)
     EXPECT(s3 != s4);
 }
 
+TEST_CASE(any_of_dynamic_true)
+{
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {3, 4, 5}});
+    migraphx::shape s0{sub_shapes};
+    EXPECT(s0.any_of_dynamic());
+
+    sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 1}, {4, 4}}});
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {3, 4, 5}});
+    migraphx::shape s1{sub_shapes};
+    EXPECT(s1.any_of_dynamic());
+}
+
+TEST_CASE(any_of_dynamic_false)
+{
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {1, 4}});
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {3, 4, 5}});
+    migraphx::shape s{sub_shapes};
+    EXPECT(not s.any_of_dynamic());
+}
+
 TEST_CASE(test_shape_packed)
 {
     migraphx::shape s{migraphx::shape::float_type, {2, 2}, {2, 1}};

test/verify/run_verify.cpp

@@ -185,7 +185,16 @@ void run_verify::verify(const std::string& name, const migraphx::program& p) con
         migraphx::parameter_map m;
         for(auto&& x : p.get_parameter_shapes())
         {
-            m[x.first] = migraphx::generate_argument(x.second, get_hash(x.first));
+            if(x.second.dynamic())
+            {
+                // create static shape using maximum dimensions
+                migraphx::shape static_shape{x.second.type(), x.second.max_lens()};
+                m[x.first] = migraphx::generate_argument(static_shape, get_hash(x.first));
+            }
+            else
+            {
+                m[x.first] = migraphx::generate_argument(x.second, get_hash(x.first));
+            }
         }
 
         auto gold_f = detach_async([=] { return run_ref(p, m); });

test/verify/test_select_module_add.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 "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_select_module_add : verify_program<test_select_module_add>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+        auto literal_ins = mm->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 = p.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* batch1 = create_submodule(1, "batch_1");
+        auto* batch2 = create_submodule(2, "batch_2");
+        auto* batch3 = create_submodule(3, "batch_3");
+        auto* batch4 = create_submodule(4, "batch_4");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+        auto input                              = mm->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              = mm->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input},
+            {batch1, batch2, batch3, batch4});
+        auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm->add_return({ret});
+
+        return p;
+    }
+};

test/verify/test_select_module_conv.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 "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_select_module_conv : verify_program<test_select_module_conv>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 3, 4, 4}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            migraphx::shape weights_shape{migraphx::shape::float_type, {2, 3, 3, 3}};
+            std::vector<float> weights_data(2 * 3 * 3 * 3, 2.0);
+            auto weights  = submod->add_literal(migraphx::literal{weights_shape, weights_data});
+            auto conv_ins = submod->add_instruction(
+                migraphx::make_op("convolution", {{"padding", {1, 1}}}), sm_input, weights);
+            submod->add_return({conv_ins});
+            return submod;
+        };
+        auto* batch1 = create_submodule(1, "batch_1");
+        auto* batch2 = create_submodule(2, "batch_2");
+        auto* batch3 = create_submodule(3, "batch_3");
+        auto* batch4 = create_submodule(4, "batch_4");
+
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {3, 3}, {4, 4}, {4, 4}}};
+        auto input                              = mm->add_parameter("data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(
+            migraphx::shape{migraphx::shape::float_type, {{1, 4}, {2, 2}, {4, 4}, {4, 4}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input},
+            {batch1, batch2, batch3, batch4});
+        auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm->add_return({ret});
+
+        return p;
+    }
+};

test/verify/test_select_module_reduce.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 "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_select_module_reduce : verify_program<test_select_module_reduce>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+
+        // create batch submodules
+        auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+            auto* submod = p.create_module(module_name);
+            migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 2, 2}};
+            auto sm_input = submod->add_parameter("data", sm_shape);
+            auto reduce_ins =
+                submod->add_instruction(migraphx::make_op("reduce_sum", {{"axes", {1}}}), sm_input);
+            auto squeeze_ins =
+                submod->add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), reduce_ins);
+            submod->add_return({squeeze_ins});
+            return submod;
+        };
+        auto* batch1 = create_submodule(1, "batch_1");
+        auto* batch2 = create_submodule(2, "batch_2");
+        auto* batch3 = create_submodule(3, "batch_3");
+        auto* batch4 = create_submodule(4, "batch_4");
+
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {2, 2}, {2, 2}}};
+        auto input                              = mm->add_parameter("data", s);
+        std::vector<migraphx::shape> sub_shapes = {};
+        sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {2, 2}}});
+        migraphx::shape out_attr = migraphx::shape{sub_shapes};
+        auto sm_ins              = mm->add_instruction(
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+            {input},
+            {batch1, batch2, batch3, batch4});
+        auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+        mm->add_return({ret});
+
+        return p;
+    }
+};