Add initial splitk reduce

docs/dev/env_vars.rst

@@ -107,6 +107,10 @@ Disables the ``schedule`` pass.
 Set to "1", "enable", "enabled", "yes", or "true" to use.
 Disables the ``fuse_reduce`` pass.
 
+.. envvar:: MIGRAPHX_ENABLE_SPLIT_REDUCE
+Set to "1", "enable", "enabled", "yes", or "true" to use.
+Enable split_reduce.
+
 .. envvar:: MIGRAPHX_ENABLE_NHWC
 
 Set to "1", "enable", "enabled", "yes", or "true" to use.

src/CMakeLists.txt

@@ -77,6 +77,7 @@ add_library(migraphx
     operation.cpp
     optimize_module.cpp
     pad_calc.cpp
+    param_utils.cpp
     pass.cpp
     pass_manager.cpp
     permutation.cpp
@@ -94,6 +95,7 @@ add_library(migraphx
     replace_allocate.cpp
     rewrite_reduce.cpp
     simplify_qdq.cpp
+    split_reduce.cpp
     sqlite.cpp
     rewrite_gelu.cpp
     rewrite_low_precision.cpp

src/fuse_pointwise.cpp

@@ -162,7 +162,7 @@ static std::vector<instruction_ref> append_pointwise_module(instruction_ref ins,
             input_map[input] = map_ins[param];
         }
     }
-    pm->replace_return(pm->insert_instructions(last, xm, map_ins));
+    pm->replace_return(pm->insert_instructions(last, xm, &map_ins));
     return inputs;
 }
 

src/fuse_reduce.cpp

@@ -83,23 +83,6 @@ struct fused_reduce
 };
 MIGRAPHX_REGISTER_OP(fused_reduce);
 
-static std::unordered_map<instruction_ref, instruction_ref>
-get_ins_param_map(const std::vector<instruction_ref>& inputs, const_module_ref sm)
-{
-    std::unordered_map<instruction_ref, instruction_ref> result;
-    auto names = sm->get_parameter_names();
-    std::sort(names.begin(), names.end());
-    assert(names.size() == inputs.size());
-    std::transform(names.begin(),
-                   names.end(),
-                   inputs.begin(),
-                   std::inserter(result, result.end()),
-                   [&](const auto& name, auto input) {
-                       return std::make_pair(input, sm->get_parameter(name));
-                   });
-    return result;
-}
-
 static void insert_params(module_ref sm,
                           const std::vector<instruction_ref>& inputs,
                           std::unordered_map<instruction_ref, instruction_ref>& map_ins)
@@ -119,7 +102,7 @@ static auto insert_ins_in_submodule(module_ref sm,
                                     std::unordered_map<instruction_ref, instruction_ref>& map_ins)
 {
     insert_params(sm, ins->inputs(), map_ins);
-    return sm->add_instructions({ins}, map_ins);
+    return sm->add_instructions({ins}, &map_ins);
 }
 
 static auto insert_ins_in_submodule(module_ref sm, instruction_ref ins)
@@ -136,12 +119,12 @@ insert_module_in_submodule(module_ref sm,
                            module::inserter insert = nullptr)
 {
     insert_params(sm, inputs, map_ins);
-    auto param_map = get_ins_param_map(inputs, m);
+    auto param_map = m->get_ins_param_map(inputs);
     for(auto&& [input, param] : param_map)
     {
         map_ins[param] = map_ins.at(input);
     }
-    return sm->add_instructions(m, map_ins, std::move(insert));
+    return sm->add_instructions(m, &map_ins, std::move(insert));
 }
 
 static auto

src/include/migraphx/liveness.hpp

@@ -0,0 +1,77 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2024 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_MIGRAPHX_LIVENESS_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_LIVENESS_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/ranges.hpp>
+#include <unordered_set>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+// This will do liveness analysis on the module, and it will call the
+// function `f` with the instruction and the set of the other instructions
+// that are live
+template <class F>
+void liveness(const module& m, F f)
+{
+    auto implicit_deps = m.calc_implicit_deps();
+    std::unordered_set<instruction_ref> live_set;
+    auto rp = reverse(m);
+    for(auto rins : iterator_for(rp)) // NOLINT
+    {
+        // The base iterator is one ahead, so we need to use the previous iterator
+        auto ins = std::prev(rins.base());
+        // Add live variables
+        auto add_live_variables = [&](const auto& inputs) {
+            for(auto input : inputs)
+            {
+                auto i = instruction::get_output_alias(input);
+                // Skip if variable comes from parent
+                if(not m.has_instruction(i))
+                    continue;
+                live_set.insert(i);
+            }
+        };
+        add_live_variables(ins->inputs());
+        add_live_variables(implicit_deps[ins]);
+        // Remove last usage
+        auto it = live_set.find(ins);
+        if(it != live_set.end())
+        {
+            live_set.erase(it);
+            f(ins, live_set);
+        }
+    }
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_LIVENESS_HPP

src/include/migraphx/module.hpp

@@ -50,6 +50,8 @@ struct module_impl;
 using parameter_map = std::unordered_map<std::string, argument>;
 using ins_dep_map   = std::unordered_map<instruction_ref, std::unordered_set<instruction_ref>>;
 
+struct module_with_inputs;
+
 /**
  * @brief Stores the instruction stream
  */
@@ -127,38 +129,38 @@ struct MIGRAPHX_EXPORT module
 
     std::vector<instruction_ref>
     add_instructions(const std::vector<instruction_ref>& instructions,
-                     std::unordered_map<instruction_ref, instruction_ref> map_ins = {},
-                     inserter insert                                              = nullptr);
+                     std::unordered_map<instruction_ref, instruction_ref>* map_ins = nullptr,
+                     inserter insert                                               = nullptr);
 
     std::vector<instruction_ref>
     add_instructions(const_module_ref m,
-                     std::unordered_map<instruction_ref, instruction_ref> map_ins = {},
-                     inserter insert                                              = nullptr);
+                     std::unordered_map<instruction_ref, instruction_ref>* map_ins = nullptr,
+                     inserter insert                                               = nullptr);
 
     std::vector<instruction_ref>
     add_instructions(instruction_ref start,
                      instruction_ref last,
-                     std::unordered_map<instruction_ref, instruction_ref> map_ins = {},
-                     inserter insert                                              = nullptr);
+                     std::unordered_map<instruction_ref, instruction_ref>* map_ins = nullptr,
+                     inserter insert                                               = nullptr);
 
     std::vector<instruction_ref>
     insert_instructions(instruction_ref ins,
                         const std::vector<instruction_ref>& instructions,
-                        std::unordered_map<instruction_ref, instruction_ref> map_ins = {},
-                        inserter insert                                              = nullptr);
+                        std::unordered_map<instruction_ref, instruction_ref>* map_ins = nullptr,
+                        inserter insert                                               = nullptr);
 
     std::vector<instruction_ref>
     insert_instructions(instruction_ref ins,
                         const_module_ref m,
-                        std::unordered_map<instruction_ref, instruction_ref> map_ins = {},
-                        inserter insert                                              = nullptr);
+                        std::unordered_map<instruction_ref, instruction_ref>* map_ins = nullptr,
+                        inserter insert                                               = nullptr);
 
     std::vector<instruction_ref>
     insert_instructions(instruction_ref ins,
                         instruction_ref start,
                         instruction_ref last,
-                        std::unordered_map<instruction_ref, instruction_ref> map_ins = {},
-                        inserter insert                                              = nullptr);
+                        std::unordered_map<instruction_ref, instruction_ref>* map_ins = nullptr,
+                        inserter insert                                               = nullptr);
 
     template <class... Ts>
     instruction_ref add_literal(Ts&&... xs)
@@ -186,6 +188,8 @@ struct MIGRAPHX_EXPORT module
 
     instruction_ref get_parameter(std::string name) const;
 
+    std::vector<instruction_ref> get_parameters() const;
+
     void rename_parameter(instruction_ref ins, const std::string& name);
 
     std::unordered_map<std::string, shape> get_parameter_shapes() const;
@@ -205,6 +209,28 @@ struct MIGRAPHX_EXPORT module
 
     void finalize(std::vector<context>& contexts);
 
+    /// Create a mapping from the input instruction to the corresponding
+    /// parameter instruction. Use the `reverse` flag to reverse the lookup
+    /// to be from parameter instruction to input instread.
+    std::unordered_map<instruction_ref, instruction_ref>
+    get_ins_param_map(const std::vector<instruction_ref>& inputs, bool reverse = false) const;
+
+    using with_inputs = module_with_inputs;
+
+    /// This will split the module into two parts at the instruction splits.
+    /// Each split instruction becomes an input parameter in the second
+    /// module. As such the inputs instructions to the second module will use
+    /// the split instructions as input placeholders that can be replaced
+    /// later.
+    std::array<with_inputs, 2> split(const std::vector<instruction_ref>& args,
+                                     const std::vector<instruction_ref>& splits) const;
+
+    /// This will split the module in 3 parts using different split
+    /// instruction for each additional module.
+    std::array<with_inputs, 3> split(const std::vector<instruction_ref>& args,
+                                     const std::vector<instruction_ref>& splits1,
+                                     const std::vector<instruction_ref>& splits2) const;
+
     void debug_print() const;
     void debug_print(instruction_ref ins) const;
     void debug_print(instruction_ref ins,
@@ -266,6 +292,21 @@ struct MIGRAPHX_EXPORT module
     std::unique_ptr<module_impl> impl;
 };
 
+struct module_with_inputs
+{
+    module mod;
+    std::vector<instruction_ref> inputs;
+    /// Replace the instruction in the inputs with rep
+    void replace(instruction_ref ins, instruction_ref rep);
+    /// Replace the input instructions using the map_ins to lookup the replacement
+    void replace(const std::unordered_map<instruction_ref, instruction_ref>& map_ins);
+
+    /// Replace the input instructions of the keys with the instructions
+    /// passed as values. Both vectors should be in the same order.
+    void replace(const std::vector<instruction_ref>& keys,
+                 const std::vector<instruction_ref>& values);
+};
+
 inline module& get_module(module& m) { return m; }
 
 } // namespace MIGRAPHX_INLINE_NS

src/include/migraphx/param_utils.hpp

@@ -0,0 +1,42 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2024 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_MIGRAPHX_PARAM_UTILS_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_PARAM_UTILS_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/instruction_ref.hpp>
+#include <vector>
+#include <string>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+std::string param_name(std::size_t i, const std::string& prefix = "x");
+
+void sort_params(std::vector<instruction_ref>& params);
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_PARAM_UTILS_HPP

src/include/migraphx/pass_manager.hpp

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

src/include/migraphx/split_reduce.hpp

@@ -0,0 +1,51 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2024 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_MIGRAPHX_SPLIT_REDUCE_HPP
+#define MIGRAPHX_GUARD_MIGRAPHX_SPLIT_REDUCE_HPP
+
+#include <migraphx/config.hpp>
+#include <string>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+struct module_pass_manager;
+
+/// For large reductions that are larger than the split_size, this pass will
+/// split the fused_reduce operators so that the reduction will happen across
+/// multiple compute units gaining better occupancy for targets with many
+/// compute units. Since the reduction is split across compute units, any
+/// elementwise operators will be split into separate operators as well due to
+/// needing global synchronization.
+struct MIGRAPHX_EXPORT split_reduce
+{
+    std::size_t split_size = 8192;
+    std::string name() const { return "split_reduce"; }
+    void apply(module_pass_manager& mpm) const;
+};
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_MIGRAPHX_SPLIT_REDUCE_HPP