Progress

src/include/migraphx/split_dynamic_batch.hpp → src/include/migraphx/split_single_dyn_dim.hpp

@@ -34,12 +34,12 @@ inline namespace MIGRAPHX_INLINE_NS {
 
 /**
  * Split dynamic batch dimension over submodules if exactly one dimension in the parameter list
- * is dynamic. Applies on program to only work on main module.
+ * is dynamic. Should only run on the main module.
  */
 struct split_dynamic_batch
 {
     std::string name() const { return "split_dynamic_batch"; }
-    void apply(program& p) const;
+    void apply(module_pass_manager& p) const;
 };
 
 } // namespace MIGRAPHX_INLINE_NS

src/split_dynamic_batch.cpp → src/split_single_dyn_dim.cpp

@@ -22,24 +22,32 @@
  * THE SOFTWARE.
  */
 
-#include <migraphx/split_dynamic_batch.hpp>
+#include <migraphx/split_single_dyn_dim.hpp>
 #include <migraphx/functional.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/make_op.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+// TODO code needs cleanup
+
 bool has_one_dyn_dim(std::unordered_map<std::string, shape> param_shapes,
                      std::string& dyn_param_str,
-                     unsigned int& max_batches)
+                     int& dyn_index,
+                     int& min_dim,
+                     int& max_dim)
 {
     // true if exactly one dynamic shape with exactly one non-fixed dynamic_dimension
-    // dyn_param_str is updated to the parameter string with the dynamic_dimension
+    // dyn_param_name is updated to the parameter string with the dynamic_dimension
     if(std::none_of(
            param_shapes.cbegin(), param_shapes.cend(), [](auto ps) { return ps.second.dynamic(); }))
         return false;
     int num_dynamic = 0;
     std::string out_str;
-    int tmp_batches;
+    int tmp_min = -1;
+    int tmp_max = -1;
+    int tmp_ind = -1;
     for(auto ps : param_shapes)
     {
         if(ps.second.dynamic())
@@ -50,12 +58,16 @@ bool has_one_dyn_dim(std::unordered_map<std::string, shape> param_shapes,
                 return false;
             }
             int num_nf = 0;
-            for(auto dd : ps.second.dyn_dims())
+            auto dds   = ps.second.dyn_dims();
+            for(int i = 0; i < dds.size(); ++i)
             {
+                const auto& dd = dds.at(i);
                 if(not dd.is_fixed())
                 {
                     num_nf += 1;
-                    tmp_batches = dd.max;
+                    tmp_min = dd.min;
+                    tmp_max = dd.max;
+                    tmp_ind = i;
                 }
             }
             if(num_nf == 1)
@@ -68,42 +80,59 @@ bool has_one_dyn_dim(std::unordered_map<std::string, shape> param_shapes,
             }
         }
     }
+    min_dim       = tmp_min;
+    max_dim       = tmp_max;
+    dyn_index     = tmp_ind;
     dyn_param_str = out_str;
-    max_batches   = tmp_batches;
     return true;
 }
 
-// don't use program
-// use module_pass_manager
-void split_dynamic_batch::apply(program& p) const
+/**
+ * Make all the batch sizes in the range for now
+ * Probably won't work for `if` and `loop` instructions, depending on how the submodules for those
+ * work create additional submodules for optimal values if not already done insert select_module
+ * instruction to the top, replace return bypassing other instructions. Unused instructions should
+ * be removed by dead_code_elimination
+ */
+void split_dynamic_batch::apply(module_pass_manager& mpm) const
 {
-    auto param_shapes = p.get_parameter_shapes();
-    std::string dyn_param_str;
-    unsigned int max_batches;
-    if(has_one_dyn_dim(param_shapes, dyn_param_str, max_batches))
+    module_ref mm;
+    mm                = &mpm.get_module();
+    auto param_names  = mm->get_parameter_names();
+    auto param_shapes = mm->get_parameter_shapes();
+    std::string dyn_param_name;
+    int dyn_index;
+    int min_dim;
+    int max_dim;
+    if(has_one_dyn_dim(param_shapes, dyn_param_name, dyn_index, min_dim, max_dim))
     {
-        // floor(log_2(max_batches))
-        unsigned int mask = 1U << 31;
-        while(not(max_batches & mask))
-        {
-            mask >>= 1;
-        }
-        // create submodules based on binary base
-        while(mask)
+        std::vector<module_ref> submodules;
+        for(int dim_size = min_dim; dim_size <= max_dim; ++dim_size)
         {
-            auto mod =
-                p.copy_module(p.get_main_module()->name(), "batch_size_" + std::to_string(mask));
-            // change dynamic input parameter shape to static
-            // propagate the shape change through the model
-            // create new parameter_list
-            // use module.add_instructions() with map between old and new static parameter
-            // add a return to the submodule
-            mask >>= 1;
+            auto submod = mpm.create_module("batch_" + std::to_string(dim_size));
+            std::unordered_map<instruction_ref, instruction_ref> map_ins;
+            auto dps                  = mm->get_parameter_shape(dyn_param_name);
+            auto static_lens          = dps.max_lens();
+            static_lens.at(dyn_index) = dim_size;
+            auto static_param =
+                submod->add_parameter(dyn_param_name, migraphx::shape{dps.type(), static_lens});
+            map_ins[mm->get_parameter(dyn_param_name)] = static_param;
+            submod->add_instructions(mm, map_ins);
+            submodules.push_back(submod);
         }
-        // create additional submodules for optimal values if not already done
-        //
-        // delete stuff in main module, insert instruction to the top, replace return bypassing
-        // other instructions unused instructions should be removed by dead_code_elimination
+        // redirect to select_module operator and return;
+        std::vector<instruction_ref> sm_inputs;
+        std::transform(param_names.cbegin(),
+                       param_names.cend(),
+                       std::back_inserter(sm_inputs),
+                       [&](auto pn) { return mm->get_parameter(pn); });
+        auto sm_ins = mm->insert_instruction(
+            mm->begin(),
+            migraphx::make_op("select_module",
+                              {{"output_dyn_shapes", migraphx::to_value(mm->get_output_shapes())}}),
+            sm_inputs,
+            submodules);
+        mm->replace_return({sm_ins});
     }
 }
 

src/targets/gpu/target.cpp

@@ -51,6 +51,7 @@
 #include <migraphx/simplify_algebra.hpp>
 #include <migraphx/simplify_qdq.hpp>
 #include <migraphx/simplify_reshapes.hpp>
+#include <migraphx/split_single_dyn_dim.hpp>
 #include <migraphx/gpu/allocation_model.hpp>
 #include <migraphx/gpu/compile_miopen.hpp>
 #include <migraphx/gpu/compile_ops.hpp>
@@ -102,6 +103,8 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
     // clang-format off
     return
     {
+        split_single_dyn_dim{},
+        dead_code_elimination{},
         normalize_ops{},
         dead_code_elimination{},
         simplify_qdq{},