initial

src/include/migraphx/program.hpp

@@ -130,6 +130,7 @@ struct program
 
     // module related api
     module* create_module(const std::string& name);
+    module* copy_module(const std::string& in_name, const std::string& out_name);
     module* get_module(const std::string& name);
     const module* get_module(const std::string& name) const;
 

src/include/migraphx/split_dynamic_batch.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_SPLIT_DYNAMIC_BATCH_HPP
+#define MIGRAPHX_GUARD_RTGLIB_SPLIT_DYNAMIC_BATCH_HPP
+
+#include <string>
+#include <migraphx/program.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+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.
+ */
+struct split_dynamic_batch
+{
+    std::string name() const { return "split_dynamic_batch"; }
+    void apply(program& p) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/split_dynamic_batch.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/split_dynamic_batch.hpp>
+#include <migraphx/functional.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+bool has_one_dyn_dim(std::unordered_map<std::string, shape> param_shapes,
+                     std::string& dyn_param_str,
+                     unsigned int& max_batches)
+{
+    // 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
+    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;
+    for(auto ps : param_shapes)
+    {
+        if(ps.second.dynamic())
+        {
+            num_dynamic += 1;
+            if(num_dynamic > 1)
+            {
+                return false;
+            }
+            int num_nf = 0;
+            for(auto dd : ps.second.dyn_dims())
+            {
+                if(not dd.is_fixed())
+                {
+                    num_nf += 1;
+                    tmp_batches = dd.max;
+                }
+            }
+            if(num_nf == 1)
+            {
+                out_str = ps.first;
+            }
+            else
+            {
+                return false;
+            }
+        }
+    }
+    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
+{
+    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))
+    {
+        // 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)
+        {
+            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;
+        }
+        // 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
+    }
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx