Merge

src/onnx/parse_pooling.cpp

@@ -22,14 +22,8 @@
  * THE SOFTWARE.
  */
 #include <migraphx/onnx/op_parser.hpp>
-#include <migraphx/onnx/checks.hpp>
-#include <migraphx/onnx/padding.hpp>
-#include <migraphx/op/pad.hpp>
-#include <migraphx/op/pooling.hpp>
+#include <migraphx/onnx/pooling.hpp>
 #include <migraphx/instruction.hpp>
-#include <migraphx/ranges.hpp>
-#include <migraphx/stringutils.hpp>
-#include <migraphx/make_op.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -39,68 +33,14 @@ struct parse_pooling : op_parser<parse_pooling>
 {
     std::vector<op_desc> operators() const
     {
-        return {{"AveragePool", "average"},
+        return {
+            {"AveragePool", "average"},
             {"GlobalAveragePool", "average"},
             {"GlobalMaxPool", "max"},
             {"MaxPool", "max"},
             {"LpPool", "lpnorm"},
-                {"GlobalLpPool", "lpnorm"}};
-    }
-
-    value handle_values(const op_desc& opd,
-                        onnx_parser::node_info info,
-                        const shape& in_shape,
-                        value values) const
-    {
-        auto kdims = in_shape.ndim() - 2;
-        if(starts_with(opd.onnx_name, "Global"))
-        {
-            // if spatial dimensions are dynamic use dyn_global flag
-            if(in_shape.dynamic() and std::any_of(in_shape.dyn_dims().cbegin() + 2,
-                                                  in_shape.dyn_dims().cend(),
-                                                  [](auto dd) { return not dd.is_fixed(); }))
-            {
-                values["dyn_global"] = true;
-                values["lengths"]    = std::vector<size_t>();
-            }
-            else
-            {
-                // works with static and fixed dynamic shape
-                auto m_lens       = in_shape.max_lens();
-                values["lengths"] = std::vector<size_t>(m_lens.begin() + 2, m_lens.end());
-            }
-        }
-
-        if(contains(info.attributes, "ceil_mode"))
-        {
-            values["ceil_mode"] = static_cast<bool>(info.attributes.at("ceil_mode").i());
-        }
-
-        if(contains(info.attributes, "strides"))
-        {
-            values["stride"].clear();
-            copy(info.attributes["strides"].ints(), std::back_inserter(values["stride"]));
-            check_attr_sizes(kdims, values["stride"].size(), "PARSE_POOLING: inconsistent strides");
-        }
-
-        if(contains(info.attributes, "kernel_shape"))
-        {
-            values["lengths"].clear();
-            copy(info.attributes["kernel_shape"].ints(), std::back_inserter(values["lengths"]));
-            check_attr_sizes(
-                kdims, values["lengths"].size(), "PARSE_POOLING: inconsistent lengths");
-        }
-
-        // lp_order attribute
-        if(contains(info.attributes, "p"))
-        {
-            values["lp_order"] = info.attributes.at("p").i();
-        }
-
-        // ensure pads available only when auto_pad is "NOT_SET"
-        check_padding_mode(info, "POOLING");
-
-        return values;
+            {"GlobalLpPool", "lpnorm"},
+        };
     }
 
     instruction_ref parse(const op_desc& opd,
@@ -108,144 +48,8 @@ struct parse_pooling : op_parser<parse_pooling>
                           onnx_parser::node_info info,
                           std::vector<instruction_ref> args) const
     {
-        std::string mode                                                 = opd.op_name;
-        const std::unordered_map<std::string, op::pooling_mode> mode_map = {
-            {"max", op::pooling_mode::max},
-            {"average", op::pooling_mode::average},
-            {"lpnorm", op::pooling_mode::lpnorm}};
-        if(not contains(mode_map, mode))
-        {
-            MIGRAPHX_THROW(
-                "PARSE_POOLING: onnx pooling mode must be [\"max\", \"average\", \"lpnorm\"]");
-        }
-        operation op  = make_op("pooling", {{"mode", mode_map.at(mode)}});
-        value values  = op.to_value();
-        auto l0       = args[0];
-        auto in_shape = l0->get_shape();
-        assert(in_shape.ndim() > 2);
-        auto kdims = in_shape.ndim() - 2;
-
-        values = handle_values(opd, info, in_shape, values);
-
-        // count include padding, if count include pad is 1, we always use
-        // explicit pad
-        int count_include_pad = 0;
-        if(contains(info.attributes, "count_include_pad"))
-        {
-            if(in_shape.dynamic())
-            {
-                MIGRAPHX_THROW("PARSE_POOLING: count_include_pad attribute is not supported for "
-                               "dynamic input shape");
-            }
-            count_include_pad = info.attributes.at("count_include_pad").i();
-        }
-
-        std::vector<int64_t> paddings;
-        float pad_val = ((mode == "max") ? std::numeric_limits<float>::lowest() : 0.0f);
-
-        if(contains(info.attributes, "pads"))
-        {
-            values["padding"].clear();
-            copy(info.attributes["pads"].ints(), std::back_inserter(paddings));
-            check_attr_sizes(
-                kdims, paddings.size() / 2, "PARSE_POOLING: inconsistent explicit paddings");
-        }
-
-        if(paddings.size() != 2 * kdims)
-        {
-            paddings.resize(kdims * 2);
-            std::fill_n(paddings.begin(), 2 * kdims, 0);
-        }
-
-        if(values["padding"].size() != kdims)
-        {
-            values["padding"].resize(kdims);
-            std::fill_n(values["padding"].begin(), kdims, 0);
-        }
-
-        if(values["stride"].size() != kdims)
-        {
-            values["stride"].resize(kdims);
-            std::fill_n(values["stride"].begin(), kdims, 1);
-        }
-
-        // used to calculate the supposed output shape
-        std::vector<int64_t> orig_padding = paddings;
-
-        // TODO:  add parsing for dilations
-        if(contains(info.attributes, "auto_pad") and
-           to_upper(info.attributes["auto_pad"].s()) != "NOTSET")
-        {
-            auto auto_pad = to_upper(info.attributes["auto_pad"].s());
-            // don't use the given padding sizes, if any
-            // values["padding"].clear();
-            if(in_shape.dynamic())
-            {
-                // set padding_mode to trigger auto padding at runtime
-                bool is_same_upper     = (auto_pad.find("SAME_UPPER") != std::string::npos);
-                values["padding_mode"] = is_same_upper ? to_value(op::padding_mode_t::same_upper)
-                                                       : to_value(op::padding_mode_t::same_lower);
-            }
-            else
-            {
-                // Calculate auto padding
-                // dilations (argument 4) not supported; default to all 1's
-                cal_auto_padding_size(info,
-                                      values,
-                                      values["lengths"].to_vector<std::size_t>(),
-                                      std::vector<size_t>(in_shape.ndim() - 2, 1),
-                                      in_shape.lens(),
-                                      paddings);
-                values["padding"] = paddings;
-                // default padding_mode indicates that padding sizes are not calculated dynamically
-                values["padding_mode"] = migraphx::op::padding_mode_t::default_;
-            }
-        }
-
-        std::vector<int64_t> slice_start;
-        std::vector<int64_t> slice_end;
-        tune_padding_size(values, paddings, count_include_pad, slice_start);
-
-        if(not slice_start.empty())
-        {
-            if(in_shape.dynamic())
-            {
-                MIGRAPHX_THROW(
-                    "PARSE_POOLING: asymmetric padding not supported for dynamic input shape");
-            }
-            // calculate expected output shape
-            orig_padding.insert(orig_padding.begin() + kdims, 2, 0);
-            orig_padding.insert(orig_padding.begin(), 2, 0);
-            op::pad pad{orig_padding, 0.0f};
-            shape padded_shape = pad.compute_shape({l0->get_shape()});
-
-            // make an op just to get its output shape
-            auto out_lens = make_op("pooling", values).compute_shape({padded_shape}).lens();
-            // compute slice_end information
-            slice_end.resize(slice_start.size());
-            std::transform(out_lens.begin() + 2,
-                           out_lens.end(),
-                           slice_start.begin(),
-                           slice_end.begin(),
-                           [](auto i, auto j) { return i + j; });
-        }
-        values["padding"] = std::vector<size_t>(paddings.begin(), paddings.end());
-
-        check_asym_padding(info, l0, paddings, values, count_include_pad, pad_val);
-        op.from_value(values);
-
-        auto l1 = info.add_instruction(op, l0);
-        if(not slice_start.empty())
-        {
-            std::vector<int64_t> axes(kdims);
-            std::iota(axes.begin(), axes.end(), 2);
-            l1 = info.add_instruction(
-                make_op("slice", {{"axes", axes}, {"starts", slice_start}, {"ends", slice_end}}),
-                l1);
-        }
-
-        return l1;
-    }
+        return add_pooling_op(opd, std::move(info), args[0]);
+    };
 };
 
 } // namespace onnx

src/onnx/parse_scatternd.cpp

@@ -39,15 +39,17 @@ struct parse_scatternd : op_parser<parse_scatternd>
                           const onnx_parser::node_info& info,
                           std::vector<instruction_ref>& args) const
     {
+        std::string reduction = "none";
         if(contains(info.attributes, "reduction"))
         {
-            if(info.attributes.at("reduction").s() == "add")
-                return info.add_instruction(migraphx::make_op("scatternd_add"), args);
-            if(info.attributes.at("reduction").s() == "mul")
-                return info.add_instruction(migraphx::make_op("scatternd_mul"), args);
+            reduction = info.attributes.at("reduction").s();
+            if(not contains({"none", "add", "mul", "min", "max"}, reduction))
+            {
+                MIGRAPHX_THROW("PARSE_SCATTERND: unsupported reduction mode " + reduction);
+            }
         }
 
-        return info.add_instruction(migraphx::make_op("scatternd_none"), args);
+        return info.add_instruction(migraphx::make_op("scatternd_" + reduction), args);
     }
 };
 

src/py/migraphx_py.cpp

@@ -40,7 +40,7 @@
 #include <migraphx/json.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/op/common.hpp>
-
+#include <migraphx/float8.hpp>
 #ifdef HAVE_GPU
 #include <migraphx/gpu/hip.hpp>
 #endif
@@ -144,6 +144,18 @@ struct npy_format_descriptor<half>
     static constexpr auto name() { return _("half"); }
 };
 
+template <>
+struct npy_format_descriptor<migraphx::fp8::fp8e4m3fnuz>
+{
+    static std::string format()
+    {
+        // following: https://docs.python.org/3/library/struct.html#format-characters
+        // TODO: need to figure out correct encoding
+        return "z";
+    }
+    static constexpr auto name() { return _("fp8e4m3fnuz"); }
+};
+
 } // namespace detail
 } // namespace pybind11
 

src/register_target.cpp

@@ -56,7 +56,11 @@ target make_target(const std::string& name)
 {
     if(not contains(target_map(), name))
     {
+#ifdef _WIN32
+        std::string target_name = "migraphx_" + name + ".dll";
+#else
         std::string target_name = "libmigraphx_" + name + ".so";
+#endif
         store_target_lib(dynamic_loader(target_name));
     }
     const auto it = target_map().find(name);

src/schedule.cpp

@@ -27,7 +27,7 @@
 #include <migraphx/iterator_for.hpp>
 #include <migraphx/iterator.hpp>
 #include <migraphx/dfor.hpp>
-#include <migraphx/par_for.hpp>
+#include <migraphx/simple_par_for.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/dom_info.hpp>
@@ -461,7 +461,7 @@ struct stream_info
                        std::back_inserter(index_to_ins),
                        [](auto&& it) { return it.first; });
 
-        par_for(concur_ins.size(), [&](auto ins_index, auto tid) {
+        simple_par_for(concur_ins.size(), [&](auto ins_index, auto tid) {
             auto merge_first = index_to_ins[ins_index];
             assert(concur_ins.count(merge_first) > 0);
             auto& merge_second = concur_ins.at(merge_first);