Merge branch 'develop' into rocblas_fp8

src/onnx/parse_unique.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/onnx/op_parser.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/tune_axis.hpp>
+#include <optional>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+// generate unique output stream y, given input stream x;
+//
+// case unsorted:
+// input x: [2, 1, 1, 3, 4, 3], attr_sorted = 0;
+// output(s):
+// y: [2, 1, 3, 4]   --- the unique output
+// y_indices: [0, 1, 3, 4]  --- first incidence, in terms of indices of x
+// x_rev_indices: [0, 1, 1, 2, 3, 2] --- x seen in terms of indices of y
+// y_count: [1, 2, 2, 1] -- count at each y_index. sum = len(x)
+//
+// case sorted:
+// input x: [2, 1, 1, 3, 4, 3], attr_sorted = 1;
+// output(s):
+// y: [1, 2, 3, 4]   --- the unique output
+// y_indices: [1, 0, 3, 4]  --- first incidence, in terms of indices of x
+// x_rev_indices: [1, 0, 0, 2, 3, 2] --- x seen in terms of indices of y
+// y_count: [2, 1, 2, 1] -- count at each y_index. sum = len(x)
+
+struct parse_unique : op_parser<parse_unique>
+{
+
+    std::vector<op_desc> operators() const { return {{"Unique"}}; }
+
+    std::vector<instruction_ref> parse(const op_desc& opd,
+                                       const onnx_parser& parser,
+                                       const onnx_parser::node_info& info,
+                                       std::vector<instruction_ref> args) const
+    {
+        int64_t sorted = 1; // default = sorted.
+
+        if(contains(info.attributes, "sorted"))
+            sorted = parser.parse_value(info.attributes.at("sorted")).at<int>();
+
+        std::optional<int64_t> axis;
+        if(contains(info.attributes, "axis"))
+        {
+            auto n_dim = args[0]->get_shape().ndim();
+            axis       = parser.parse_value(info.attributes.at("axis")).at<int>();
+            axis       = tune_axis(n_dim, *axis, opd.op_name);
+        }
+        migraphx::argument data_arg = args.back()->eval();
+
+        auto opr     = axis ? migraphx::make_op("unique", {{"axis", *axis}, {"sorted", sorted}})
+                            : migraphx::make_op("unique", {{"sorted", sorted}});
+        auto u_opr   = info.add_instruction(opr, args.at(0));
+        auto i_y     = info.add_instruction(make_op("get_tuple_elem", {{"index", 0}}), u_opr);
+        auto i_y_idx = info.add_instruction(make_op("get_tuple_elem", {{"index", 1}}), u_opr);
+        auto i_x_idx = info.add_instruction(make_op("get_tuple_elem", {{"index", 2}}), u_opr);
+        auto i_count = info.add_instruction(make_op("get_tuple_elem", {{"index", 3}}), u_opr);
+
+        return {i_y, i_y_idx, i_x_idx, i_count};
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/onnx/pooling.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/onnx/pooling.hpp>
+#include <migraphx/onnx/checks.hpp>
+#include <migraphx/onnx/padding.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/op/pooling.hpp>
+#include <migraphx/op/pad.hpp>
+#include <migraphx/ranges.hpp>
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+value handle_pooling_values(const op_desc& opd,
+                            onnx_parser::node_info info,
+                            const shape& in_shape,
+                            value values)
+{
+    auto kdims = in_shape.ndim() - 2;
+    if(starts_with(opd.onnx_name, "Global") or starts_with(opd.onnx_name, "QLinearGlobal"))
+    {
+        // 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");
+    }
+
+    if(contains(info.attributes, "dilations"))
+    {
+        values["dilations"].clear();
+        copy(info.attributes["dilations"].ints(), std::back_inserter(values["dilations"]));
+        check_attr_sizes(
+            kdims, values["dilations"].size(), "PARSE_POOLING: inconsistent dilations");
+    }
+
+    // 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;
+}
+
+instruction_ref add_pooling_op(const op_desc& opd, onnx_parser::node_info info, instruction_ref l0)
+{
+    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 in_shape = l0->get_shape();
+    assert(in_shape.ndim() > 2);
+    auto kdims = in_shape.ndim() - 2;
+
+    values = handle_pooling_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);
+    }
+
+    if(values["dilations"].size() != kdims)
+    {
+        values["dilations"].resize(kdims);
+        std::fill_n(values["dilations"].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>(),
+                                  values["dilations"].to_vector<std::size_t>(),
+                                  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;
+}
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

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/rewrite_pooling.cpp

@@ -35,25 +35,14 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-void rewrite_pooling::apply(module& m) const
+static void replace_with_reduce(module& m, instruction_ref ins)
 {
-    for(auto ins : iterator_for(m))
-    {
-        if(ins->name() != "pooling")
-            continue;
-        if(ins->inputs().empty())
-            continue;
     auto&& s  = ins->inputs().front()->get_shape();
     auto&& op = any_cast<op::pooling>(ins->get_operator());
-        if(not std::all_of(op.padding.begin(), op.padding.end(), [](auto i) { return i == 0; }))
-            continue;
-        if(not std::all_of(op.stride.begin(), op.stride.end(), [](auto i) { return i == 1; }))
-            continue;
     auto lens = s.lens();
-        if(not std::equal(lens.begin() + 2, lens.end(), op.lengths.begin(), op.lengths.end()))
-            continue;
     std::vector<std::int64_t> axes(lens.size() - 2);
     std::iota(axes.begin(), axes.end(), 2);
+
     // average pooling
     if(op.mode == op::pooling_mode::average)
     {
@@ -64,6 +53,131 @@ void rewrite_pooling::apply(module& m) const
     {
         m.replace_instruction(ins, make_op("reduce_max", {{"axes", axes}}), ins->inputs());
     }
+}
+
+static void replace_dilations_with_gather_pooling(module& m, instruction_ref ins)
+{
+    // TODO remove this when MIOpen supports dilated pooling
+    auto&& s  = ins->inputs().front()->get_shape();
+    auto&& op = any_cast<op::pooling>(ins->get_operator());
+    // Ignore N, C axes
+    std::vector<size_t> dims = {s.lens().cbegin() + 2, s.lens().cend()};
+
+    bool default_padding =
+        std::all_of(op.padding.cbegin(), op.padding.cend(), [](auto i) { return i == 0; });
+
+    if(not default_padding)
+    {
+        for(size_t idx{0}; idx < op.padding.size(); ++idx)
+        {
+            // We need to pad both ends
+            dims[idx] += op.padding.at(idx) * 2;
+        }
+    }
+    std::vector<size_t> kernels   = op.lengths;
+    std::vector<size_t> strides   = op.stride;
+    std::vector<size_t> dilations = op.dilations;
+
+    std::vector<std::vector<int>> axis_indices;
+    axis_indices.resize(dims.size());
+
+    for(auto idx{0}; idx < dims.size(); ++idx)
+    {
+        // Only consider if iw fits into the window
+        for(size_t stride{0}; stride < dims.at(idx) - dilations.at(idx) * (kernels.at(idx) - 1);
+            stride += strides.at(idx))
+        {
+            for(size_t step{0}; step < kernels.at(idx); ++step)
+            {
+                axis_indices.at(idx).push_back(stride + dilations.at(idx) * step);
+            }
+        }
+    }
+
+    auto elements = ins->inputs().front();
+    if(not default_padding)
+    {
+        // Pad supports asym, we need to provide both ends
+        std::vector<size_t> padding(2 * s.lens().size(), 0);
+        // Format will be e.g {N, C, P1, P2, N, C, P1, P2}
+        for(size_t idx{0}; idx < op.padding.size(); ++idx)
+        {
+            // Ignore N, C axes
+            padding.at(2 + idx)                   = op.padding.at(idx);
+            padding.at(2 + idx + s.lens().size()) = op.padding.at(idx);
+        }
+
+        // Default value needed for Max pooling
+        elements = m.insert_instruction(
+            ins,
+            make_op("pad", {{"pads", padding}, {"value", std::numeric_limits<float>::lowest()}}),
+            elements);
+    }
+
+    for(auto idx{0}; idx < axis_indices.size(); ++idx)
+    {
+        migraphx::shape s_indices{migraphx::shape::int32_type, {axis_indices.at(idx).size()}};
+        auto indices = m.add_literal(migraphx::literal{s_indices, axis_indices.at(idx)});
+        elements     = m.insert_instruction(
+            ins, make_op("gather", {{"axis", idx + 2 /*ignore N,C*/}}), elements, indices);
+    }
+
+    // Ignore padding
+    std::vector<size_t> new_padding(kernels.size(), 0);
+    // The kernel window elements are places next to each other. E.g. {x1, y1, x2, y2, ...}
+    // We need to skip them to not overlap
+    std::vector<size_t> new_strides(kernels);
+    // Ignore dilations
+    std::vector<size_t> new_dilations(kernels.size(), 1);
+    m.replace_instruction(ins,
+                          make_op("pooling",
+                                  {{"mode", op.mode},
+                                   {"padding", new_padding},
+                                   {"stride", new_strides},
+                                   {"lengths", kernels},
+                                   {"dilations", new_dilations}}),
+                          elements);
+}
+
+void rewrite_pooling::apply(module& m) const
+{
+    for(auto ins : iterator_for(m))
+    {
+        if(ins->name() != "pooling")
+            continue;
+        if(ins->inputs().empty())
+            continue;
+        auto&& s                  = ins->inputs().front()->get_shape();
+        auto&& op                 = any_cast<op::pooling>(ins->get_operator());
+        bool same_kernel_as_shape = std::equal(
+            s.lens().cbegin() + 2, s.lens().cend(), op.lengths.cbegin(), op.lengths.cend());
+        bool default_strides =
+            std::all_of(op.stride.cbegin(), op.stride.cend(), [](auto i) { return i == 1; });
+        bool default_padding =
+            std::all_of(op.padding.cbegin(), op.padding.cend(), [](auto i) { return i == 0; });
+        bool default_dilations =
+            std::all_of(op.dilations.cbegin(), op.dilations.cend(), [](auto i) { return i == 1; });
+        if(same_kernel_as_shape and default_strides and default_padding and default_dilations)
+        {
+            replace_with_reduce(m, ins);
+        }
+        else if(not default_dilations)
+        {
+            // Dilated AvgPool with padding is not supported
+            if(not default_padding and op.mode == op::pooling_mode::average)
+            {
+                continue;
+            }
+            auto size =
+                std::accumulate(s.lens().cbegin(), s.lens().cend(), 1, std::multiplies<size_t>());
+            // Can't handle too much size because of literal size
+            if(size > 100000)
+            {
+                continue;
+            }
+
+            replace_dilations_with_gather_pooling(m, ins);
+        }
     }
 }
 

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);

src/simplify_dyn_ops.cpp

@@ -22,6 +22,7 @@
  * THE SOFTWARE.
  */
 #include <migraphx/simplify_dyn_ops.hpp>
+#include <migraphx/op/slice.hpp>
 #include <migraphx/matcher.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/literal.hpp>
@@ -33,6 +34,10 @@ inline namespace MIGRAPHX_INLINE_NS {
  * Convert 2 input static shape broadcast/multibroadcast into 1 input version.
  * Some compiler passes (ex. simplify_algebra) only support the 1 input versions
  * of the broadcasting operators.
+ * From:
+ * broadcast_op(argument_with_static_shape, argument_with_static_shape)
+ * To:
+ * broadcast_op(argument_with_static_shape); broadcast_op.out_lens = constant_output_dims
  */
 struct find_static_2in_broadcasts
 {
@@ -61,8 +66,65 @@ struct find_static_2in_broadcasts
 };
 
 /**
- * Simplify slice with variable `starts` and `ends` to the constant version if
- * the `input_starts` and `input_ends` inputs are constant.
+ * Simplify slice with 2 inputs to the 1 input version if inputs[1] is constant.
+ * From:
+ * slice(data, constant_input); two attributes set
+ * To:
+ * slice(data); slice.starts, slice.ends. slice.axes set
+ */
+struct find_const_2in_slice
+{
+    auto matcher() const
+    {
+        return match::name("slice")(match::nargs(2), match::arg(1)(match::is_constant()));
+    }
+
+    void apply(module& m, const match::matcher_result& mr) const
+    {
+        auto ins       = mr.result;
+        auto inputs    = ins->inputs();
+        auto slice_op  = any_cast<op::slice>(ins->get_operator());
+        auto set_attrs = slice_op.get_set_attributes();
+        std::vector<int64_t> starts_vec;
+        std::vector<int64_t> ends_vec;
+        std::vector<int64_t> axes_vec;
+        if(set_attrs == op::slice::ends_axes)
+        {
+            // slice(data, starts)
+            inputs.at(1)->eval().visit(
+                [&](auto output) { starts_vec.assign(output.begin(), output.end()); });
+            ends_vec = slice_op.ends;
+            axes_vec = slice_op.axes;
+        }
+        else if(set_attrs == op::slice::starts_axes)
+        {
+            // slice(data, ends)
+            inputs.at(1)->eval().visit(
+                [&](auto output) { ends_vec.assign(output.begin(), output.end()); });
+            starts_vec = slice_op.starts;
+            axes_vec   = slice_op.axes;
+        }
+        else
+        {
+            // slice(data, axes)
+            inputs.at(1)->eval().visit(
+                [&](auto output) { axes_vec.assign(output.begin(), output.end()); });
+            starts_vec = slice_op.starts;
+            ends_vec   = slice_op.ends;
+        }
+        m.replace_instruction(
+            ins,
+            make_op("slice", {{"starts", starts_vec}, {"ends", ends_vec}, {"axes", axes_vec}}),
+            inputs.at(0));
+    }
+};
+
+/**
+ * Simplify slice with 3 inputs to the 1 input version if inputs[1:2] are constant.
+ * From:
+ * slice(data, constant_input1, constant_input2); one attribute set
+ * To:
+ * slice(data); slice.starts, slice.ends. slice.axes set
  */
 struct find_const_3in_slice
 {
@@ -77,27 +139,51 @@ struct find_const_3in_slice
     {
         auto ins            = mr.result;
         auto inputs         = ins->inputs();
-        argument starts_arg = inputs.at(1)->eval();
-        argument ends_arg   = inputs.at(2)->eval();
-        if(not starts_arg.empty() and not ends_arg.empty())
-        {
+        auto slice_op       = any_cast<op::slice>(ins->get_operator());
+        auto set_attrs      = slice_op.get_set_attributes();
         std::vector<int64_t> starts_vec;
         std::vector<int64_t> ends_vec;
-            starts_arg.visit([&](auto output) { starts_vec.assign(output.begin(), output.end()); });
-            ends_arg.visit([&](auto output) { ends_vec.assign(output.begin(), output.end()); });
-            auto slice_val = ins->get_operator().to_value();
-            auto axes_vec  = slice_val.at("axes").to_vector<int64_t>();
+        std::vector<int64_t> axes_vec;
+        if(set_attrs == op::slice::axes_only)
+        {
+            // slice(data, starts, ends)
+            inputs.at(1)->eval().visit(
+                [&](auto output) { starts_vec.assign(output.begin(), output.end()); });
+            inputs.at(2)->eval().visit(
+                [&](auto output) { ends_vec.assign(output.begin(), output.end()); });
+            axes_vec = slice_op.axes;
+        }
+        else if(set_attrs == op::slice::ends_only)
+        {
+            // slice(data, starts, axes)
+            inputs.at(1)->eval().visit(
+                [&](auto output) { starts_vec.assign(output.begin(), output.end()); });
+            inputs.at(2)->eval().visit(
+                [&](auto output) { axes_vec.assign(output.begin(), output.end()); });
+            ends_vec = slice_op.ends;
+        }
+        else
+        {
+            // slice(data, ends, axes)
+            inputs.at(1)->eval().visit(
+                [&](auto output) { ends_vec.assign(output.begin(), output.end()); });
+            inputs.at(2)->eval().visit(
+                [&](auto output) { axes_vec.assign(output.begin(), output.end()); });
+            starts_vec = slice_op.starts;
+        }
         m.replace_instruction(
             ins,
             make_op("slice", {{"starts", starts_vec}, {"ends", ends_vec}, {"axes", axes_vec}}),
             inputs.at(0));
     }
-    }
 };
 
 /**
- * Simplify slice with variable `starts`, `ends`, and `input_axes` to the constant version if
- * the `input_starts`, `input_ends`, and `input_axes` inputs are constant.
+ * Simplify slice with 4 inputs to the 1 input version if inputs[1:3] are constant.
+ * From:
+ * slice(data, constant_starts, constant_ends, constant_axes)
+ * To:
+ * slice(data); slice.starts, slice.ends. slice.axes set
  */
 struct find_const_4in_slice
 {
@@ -113,9 +199,9 @@ struct find_const_4in_slice
     {
         auto ins            = mr.result;
         auto inputs         = ins->inputs();
-        argument starts_arg = inputs.at(1)->eval();
-        argument ends_arg   = inputs.at(2)->eval();
-        argument axes_arg   = inputs.at(3)->eval();
+        argument starts_arg = inputs.at(1)->eval(false);
+        argument ends_arg   = inputs.at(2)->eval(false);
+        argument axes_arg   = inputs.at(3)->eval(false);
         if(not starts_arg.empty() and not ends_arg.empty() and not axes_arg.empty())
         {
             std::vector<int64_t> starts_vec;
@@ -172,13 +258,76 @@ struct find_static_dimensions_of
     }
 };
 
+/**
+ * Simplify allocate into 2 argument reshape that has constant output dimensions into a static 1
+ * argument reshape. Intended to simplify what ONNX parse_reshape creates for dynamic reshapes.
+ * This matcher can be generalized to matching reshape(data, static_shape_output_tensor).
+ * From:
+ * x = allocate(constant_output_dims) -> reshape(data, x)
+ * To:
+ * reshape(data); reshape.dims = constant_output_dims
+ */
+struct find_const_alloc_reshapes
+{
+    auto matcher() const
+    {
+        return match::name("reshape")(match::nargs(2),
+                                      match::arg(1)(match::name("allocate")(match::is_constant())));
+    }
+
+    void apply(module& m, const match::matcher_result& mr) const
+    {
+        auto reshape_ins         = mr.result;
+        auto reshape_inputs      = reshape_ins->inputs();
+        auto alloc_ins           = reshape_inputs.at(1);
+        argument output_dims_arg = alloc_ins->inputs().at(0)->eval(false);
+        std::vector<int64_t> output_dims_vec;
+        output_dims_arg.visit(
+            [&](auto output) { output_dims_vec.assign(output.begin(), output.end()); });
+        m.replace_instruction(
+            reshape_ins, make_op("reshape", {{"dims", output_dims_vec}}), reshape_inputs.at(0));
+        // have dead_code_elimination remove the previous allocate
+    }
+};
+
+/**
+ * Simplify allocate into fill operator that has constant output dimensions and constant value.
+ * The allocate into fill instructions is what is produced when parsing the ONNX
+ * ConstantOfShape operator. This replacement could be handled with propagate_constant, but
+ * would rather have the simplification happen earlier during compiling.
+ * This matcher can be generalized to matching fill(constant_value, static_shape_output_tensor).
+ * From:
+ * x = allocate(constant_ouptut_dims) -> fill(constant_value, x)
+ * To:
+ * literal
+ */
+struct find_const_alloc_fill
+{
+    auto matcher() const
+    {
+        return match::name("fill")(match::arg(0)(match::is_constant()),
+                                   match::arg(1)(match::name("allocate")(match::is_constant())));
+    }
+
+    void apply(module& m, const match::matcher_result& mr) const
+    {
+        auto fill_ins = mr.result;
+        auto fill_arg = fill_ins->eval(false);
+        auto l        = m.add_literal(fill_arg.get_shape(), fill_arg.data());
+        m.replace_instruction(fill_ins, l);
+    }
+};
+
 void simplify_dyn_ops::apply(module& m) const
 {
     match::find_matches(m,
-                        find_static_2in_broadcasts{},
                         find_static_dimensions_of{},
+                        find_const_alloc_reshapes{},
+                        find_static_2in_broadcasts{},
+                        find_const_2in_slice{},
                         find_const_3in_slice{},
-                        find_const_4in_slice{});
+                        find_const_4in_slice{},
+                        find_const_alloc_fill{});
 }
 
 } // namespace MIGRAPHX_INLINE_NS

src/targets/cpu/CMakeLists.txt

@@ -74,21 +74,27 @@ if(MIGRAPHX_ENABLE_ZENDNN)
     target_link_libraries(migraphx_cpu PRIVATE ${BLIS_LIB})
     target_link_libraries(migraphx_cpu PRIVATE ${ZENDNN_LIB})
 else()
-    target_link_libraries(migraphx_cpu PRIVATE DNNL::dnnl)
+    target_link_libraries(migraphx_cpu PUBLIC DNNL::dnnl)
 endif()
 target_link_libraries(migraphx_cpu PRIVATE migraphx)
 
 migraphx_generate_export_header(migraphx_cpu)
 
 find_package(OpenMP)
-target_link_libraries(migraphx_cpu PUBLIC OpenMP::OpenMP_CXX)
-# Add library path to rpath to workaround issues with our broken packages
-foreach(LIBRARY ${OpenMP_CXX_LIBRARIES})
+if(WIN32)
+    target_link_libraries(migraphx_cpu PUBLIC libomp)
+    target_include_directories(migraphx_cpu PUBLIC ${OpenMP_CXX_INCLUDE_DIRS})
+    target_compile_options(migraphx_cpu PUBLIC ${OpenMP_CXX_FLAGS})
+else()
+    target_link_libraries(migraphx_cpu PUBLIC OpenMP::OpenMP_CXX)
+    # Add library path to rpath to workaround issues with our broken packages
+    foreach(LIBRARY ${OpenMP_CXX_LIBRARIES})
         if(LIBRARY MATCHES "libomp")
             get_filename_component(LIBRARY_PATH "${LIBRARY}" PATH)
             target_link_libraries(migraphx_cpu PUBLIC -Wl,-rpath=${LIBRARY_PATH} -Wl,-rpath-link=${LIBRARY_PATH})
         endif()
-endforeach()
+    endforeach()
+endif()
 
 rocm_install_targets(
   TARGETS migraphx_cpu

src/targets/cpu/dnnl.cpp

@@ -67,8 +67,8 @@ dnnl::memory::data_type to_dnnl_memory_data_type(shape::type_t t)
     case st::float_type: return dt::f32;
     case st::int32_type: return dt::s32;
     case st::int8_type: return dt::s8;
-    case st::uint8_type:
-    case st::fp8e4m3fnuz_type: return dt::u8;
+    case st::uint8_type: return dt::u8;
+    case st::fp8e4m3fnuz_type: MIGRAPHX_THROW("fp8e4m3fnuz unsupported in DNNL");
     default: MIGRAPHX_THROW("Unsupported data type");
     }
 }

src/targets/cpu/lowering.cpp

@@ -340,7 +340,6 @@ struct cpu_apply
                               {"reduce_min", "reduction_min"},
                               {"reduce_sum", "reduction_sum"},
                           });
-
         extend_op("concat", "dnnl::concat");
         extend_op("contiguous", "dnnl::reorder");
         extend_op("convolution", "dnnl::convolution");
@@ -376,6 +375,12 @@ struct cpu_apply
         // Apply these operators first so the inputs can be const folded
         for(auto it : iterator_for(*modl))
         {
+            // skip lowering if input has fp8 as one of the inputs since oneDNN doesn't have fp8
+            // supported yet.
+            if(std::any_of(it->inputs().begin(), it->inputs().end(), [](const auto& i) {
+                   return i->get_shape().type() == migraphx::shape::fp8e4m3fnuz_type;
+               }))
+                continue;
             if(it->name() == "pow")
             {
                 apply_pow(it);
@@ -383,6 +388,12 @@ struct cpu_apply
         }
         for(auto it : iterator_for(*modl))
         {
+            // skip lowering if input has fp8 as one of the inputs since oneDNN doesn't have fp8
+            // supported yet.
+            if(std::any_of(it->inputs().begin(), it->inputs().end(), [](const auto& i) {
+                   return i->get_shape().type() == migraphx::shape::fp8e4m3fnuz_type;
+               }))
+                continue;
             if(it->name() == "pooling")
             {
                 apply_pooling(it);

src/targets/cpu/pooling.cpp

@@ -34,23 +34,32 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace cpu {
 
-struct dnnl_pooling : dnnl_extend_op<dnnl_pooling, dnnl::pooling_forward, op::pooling>
+struct dnnl_pooling : dnnl_extend_op<dnnl_pooling, dnnl::pooling_v2_forward, op::pooling>
 {
     std::vector<int> arg_map(int) const { return {MIGRAPHX_DNNL_PREFIX(ARG_SRC)}; }
 
-    dnnl::pooling_forward::desc get_desc(const std::unordered_map<int, dnnl::memory::desc>& m) const
+    dnnl::pooling_v2_forward::desc
+    get_desc(const std::unordered_map<int, dnnl::memory::desc>& m) const
     {
         auto algo  = op.mode == op::pooling_mode::max ? dnnl::algorithm::pooling_max
                                                       : dnnl::algorithm::pooling_avg;
         auto kdims = op.kdims();
         std::vector<size_t> padding_l(op.padding.begin(), op.padding.begin() + kdims);
         std::vector<size_t> padding_r(op.padding.begin() + kdims, op.padding.end());
+        // Note: It is not documented, but the default dilation seems to be 0 instead of 1.
+        //       We need to offset dilations with -1.
+        std::vector<size_t> dilations;
+        std::transform(op.dilations.cbegin(),
+                       op.dilations.cend(),
+                       std::back_inserter(dilations),
+                       [](size_t d) { return d - 1; });
         return {dnnl::prop_kind::forward_inference,
                 algo,
                 m.at(MIGRAPHX_DNNL_PREFIX(ARG_SRC)),
                 m.at(MIGRAPHX_DNNL_PREFIX(ARG_DST)),
                 to_dnnl_dims(op.stride),
                 to_dnnl_dims(op.lengths),
+                to_dnnl_dims(dilations),
                 to_dnnl_dims(padding_l),
                 to_dnnl_dims(padding_r)};
     }

src/targets/gpu/CMakeLists.txt

@@ -22,20 +22,20 @@
 # THE SOFTWARE.
 # ####################################################################################
 
-list(APPEND CMAKE_PREFIX_PATH /opt/rocm)
-find_package(hip)
+find_package(hip REQUIRED)
 if(NOT GPU_TARGETS)
-    message(FATAL_ERROR "HIP package is broken and has no GPU_TARGETS, please pass -DGPU_TARGETS=$(/opt/rocm/bin/rocminfo | grep -o -m1 'gfx.*') to cmake to build for your gpu.")
+    set(fatal_msg "HIP package is broken and has no GPU_TARGETS. Please pass GPU_TARGETS to cmake.")
+    if(NOT WIN32)
+        set(fatal_msg "${fatal_msg}\nUse -DGPU_TARGETS=$(/opt/rocm/bin/rocminfo | grep -o -m1 'gfx.*') to build for your GPU.")
+    endif()
+    message(FATAL_ERROR ${fatal_msg})
 endif()
-find_package(miopen)
+find_package(miopen REQUIRED)
+message(STATUS "MIGraphX is using MIOpen")
 
 # rocblas
-find_package(rocblas REQUIRED PATHS /opt/rocm)
-message(STATUS "Build with rocblas")
-
-if(NOT TARGET MIOpen)
-    message(SEND_ERROR "Cant find miopen")
-endif()
+find_package(rocblas REQUIRED)
+message(STATUS "MIGraphX build with rocBLAS")
 
 if(MIGRAPHX_USE_COMPOSABLEKERNEL)
     find_package(composable_kernel 1.0.0 REQUIRED COMPONENTS jit_library)
@@ -49,7 +49,6 @@ endif()
 
 file(GLOB KERNEL_FILES CONFIGURE_DEPENDS
     ${CMAKE_CURRENT_SOURCE_DIR}/kernels/include/migraphx/kernels/*.hpp)
-message(STATUS "KERNEL_FILES: ${KERNEL_FILES}")
 
 if(NOT MIGRAPHX_USE_COMPOSABLEKERNEL)
     list(REMOVE_ITEM KERNEL_FILES
@@ -66,8 +65,10 @@ file(GLOB DEVICE_GPU_SRCS CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/device/*
 add_library(migraphx_device ${DEVICE_GPU_SRCS})
 
 add_library(compile_for_gpu INTERFACE)
-target_compile_options(compile_for_gpu INTERFACE -std=c++17 -fno-gpu-rdc -Wno-cuda-compat -Wno-unused-command-line-argument -Xclang -fallow-half-arguments-and-returns)
-target_link_libraries(compile_for_gpu INTERFACE hip::device -fno-gpu-rdc -Wno-invalid-command-line-argument -Wno-unused-command-line-argument -Wno-option-ignored)
+target_compile_features(compile_for_gpu INTERFACE cxx_std_17)
+target_compile_options(compile_for_gpu INTERFACE -fno-gpu-rdc -Wno-cuda-compat -Wno-unused-command-line-argument -Xclang -fallow-half-arguments-and-returns)
+target_link_options(compile_for_gpu INTERFACE  -fno-gpu-rdc -Wno-invalid-command-line-argument -Wno-unused-command-line-argument -Wno-option-ignored)
+target_link_libraries(compile_for_gpu INTERFACE hip::device)
 check_cxx_compiler_flag("--cuda-host-only -fhip-lambda-host-device -x hip" HAS_HIP_LAMBDA_HOST_DEVICE)
 
 if(HAS_HIP_LAMBDA_HOST_DEVICE)
@@ -211,8 +212,10 @@ if(MIGRAPHX_ENABLE_MLIR)
 endif()
 
 if(MIGRAPHX_USE_HIPRTC)
+    find_package(hiprtc REQUIRED)
     message(STATUS "MIGraphX is using hipRTC")
     target_compile_definitions(migraphx_gpu PRIVATE -DMIGRAPHX_USE_HIPRTC=1)
+    target_link_libraries(migraphx_gpu PUBLIC hiprtc::hiprtc)
 else()
     message(STATUS "MIGraphX is using HIP Clang")
 
@@ -221,12 +224,15 @@ else()
     target_flags(HIP_COMPILER_FLAGS hip::device)
 
     # Remove cuda arch flags
-    string(REGEX REPLACE --cuda-gpu-arch=[a-z0-9]+ "" HIP_COMPILER_FLAGS "${HIP_COMPILER_FLAGS}")
-    string(REGEX REPLACE --offload-arch=[a-z0-9:+-]+ "" HIP_COMPILER_FLAGS "${HIP_COMPILER_FLAGS}")
+    string(REGEX REPLACE "--cuda-gpu-arch=[a-z0-9]+ ?" "" HIP_COMPILER_FLAGS "${HIP_COMPILER_FLAGS}")
+    string(REGEX REPLACE "--offload-arch=[a-z0-9:+-]+ ?" "" HIP_COMPILER_FLAGS "${HIP_COMPILER_FLAGS}")
 
     # Skip library paths since hip will incorrectly treat it as a source file
     string(APPEND HIP_COMPILER_FLAGS " ")
 
+    if(WIN32)
+        string(REPLACE "\\" "/" HIP_COMPILER_FLAGS "${HIP_COMPILER_FLAGS}")
+    endif()
     foreach(_unused RANGE 2)
         string(REGEX REPLACE " /[^ ]+\\.(a|so) " " " HIP_COMPILER_FLAGS "${HIP_COMPILER_FLAGS}")
     endforeach()
@@ -238,7 +244,11 @@ else()
     )
 
     if(DEFINED CMAKE_CXX_COMPILER_LAUNCHER)
+        if(WIN32)
+            execute_process(COMMAND where ${CMAKE_CXX_COMPILER_LAUNCHER} OUTPUT_VARIABLE MIGRAPHX_HIP_COMPILER_LAUNCHER)
+        else()
             execute_process(COMMAND which ${CMAKE_CXX_COMPILER_LAUNCHER} OUTPUT_VARIABLE MIGRAPHX_HIP_COMPILER_LAUNCHER)
+        endif()
         string(STRIP "${MIGRAPHX_HIP_COMPILER_LAUNCHER}" MIGRAPHX_HIP_COMPILER_LAUNCHER)
         target_compile_definitions(migraphx_gpu PRIVATE -DMIGRAPHX_HIP_COMPILER_LAUNCHER="${MIGRAPHX_HIP_COMPILER_LAUNCHER}")
     endif()

src/targets/gpu/compile_hip.cpp

@@ -194,7 +194,7 @@ struct hiprtc_program
 };
 
 std::vector<std::vector<char>> compile_hip_src_with_hiprtc(std::vector<hiprtc_src_file> srcs,
-                                                           std::string params,
+                                                           const std::string& params,
                                                            const std::string& arch)
 {
     hiprtc_program prog(std::move(srcs));
@@ -238,8 +238,9 @@ bool hip_has_flags(const std::vector<std::string>& flags)
     }
 }
 
-std::vector<std::vector<char>>
-compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std::string& arch)
+std::vector<std::vector<char>> compile_hip_src(const std::vector<src_file>& srcs,
+                                               const std::string& params,
+                                               const std::string& arch)
 {
     std::vector<hiprtc_src_file> hsrcs{srcs.begin(), srcs.end()};
     if(enabled(MIGRAPHX_GPU_DUMP_SRC{}))
@@ -281,13 +282,13 @@ compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std
         if(fs::exists(out))
             return {read_buffer(out.string())};
     }
-    return compile_hip_src_with_hiprtc(std::move(hsrcs), std::move(params), arch);
+    return compile_hip_src_with_hiprtc(std::move(hsrcs), params, arch);
 }
 
 #else // MIGRAPHX_USE_HIPRTC
 
 std::vector<std::vector<char>> compile_hip_src_with_hiprtc(std::vector<hiprtc_src_file>, // NOLINT
-                                                           std::string,                  // NOLINT
+                                                           const std::string&,           // NOLINT
                                                            const std::string&)
 {
     MIGRAPHX_THROW("Not using hiprtc");
@@ -316,29 +317,15 @@ src_compiler assemble(src_compiler compiler)
     return compiler;
 }
 
-std::vector<std::vector<char>>
-compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std::string& arch)
+std::vector<std::vector<char>> compile_hip_src(const std::vector<src_file>& srcs,
+                                               const std::string& params,
+                                               const std::string& arch)
 {
     assert(not srcs.empty());
+
     if(not is_hip_clang_compiler())
         MIGRAPHX_THROW("Unknown hip compiler: " MIGRAPHX_HIP_COMPILER);
 
-    if(params.find("-std=") == std::string::npos)
-        params += " --std=c++17";
-    params += " -fno-gpu-rdc";
-    if(enabled(MIGRAPHX_GPU_DEBUG_SYM{}))
-        params += " -g";
-    params += " -c";
-    params += " --offload-arch=" + arch;
-    params += " --cuda-device-only";
-    params += " -O" + string_value_of(MIGRAPHX_GPU_OPTIMIZE{}, "3") + " ";
-
-    if(enabled(MIGRAPHX_GPU_DEBUG{}))
-        params += " -DMIGRAPHX_DEBUG";
-
-    params += " -Wno-unused-command-line-argument -Wno-cuda-compat ";
-    params += MIGRAPHX_HIP_COMPILER_FLAGS;
-
     src_compiler compiler;
     compiler.flags    = params;
     compiler.compiler = MIGRAPHX_HIP_COMPILER;
@@ -346,6 +333,23 @@ compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std
     if(has_compiler_launcher())
         compiler.launcher = MIGRAPHX_HIP_COMPILER_LAUNCHER;
 #endif
+
+    if(params.find("-std=") == std::string::npos)
+        compiler.flags += " --std=c++17";
+    compiler.flags += " -fno-gpu-rdc";
+    if(enabled(MIGRAPHX_GPU_DEBUG_SYM{}))
+        compiler.flags += " -g";
+    compiler.flags += " -c";
+    compiler.flags += " --offload-arch=" + arch;
+    compiler.flags += " --cuda-device-only";
+    compiler.flags += " -O" + string_value_of(MIGRAPHX_GPU_OPTIMIZE{}, "3") + " ";
+
+    if(enabled(MIGRAPHX_GPU_DEBUG{}))
+        compiler.flags += " -DMIGRAPHX_DEBUG";
+
+    compiler.flags += " -Wno-unused-command-line-argument -Wno-cuda-compat ";
+    compiler.flags += MIGRAPHX_HIP_COMPILER_FLAGS;
+
     if(enabled(MIGRAPHX_GPU_DUMP_SRC{}))
     {
         for(const auto& src : srcs)

src/targets/gpu/compile_hip_code_object.cpp

@@ -200,7 +200,7 @@ operation compile_hip_code_object(const std::string& content, hip_compile_option
     options.params += " " + join_strings(compiler_warnings(), " ");
     options.params += " -ftemplate-backtrace-limit=0";
     options.params += " -Werror";
-    auto cos = compile_hip_src(srcs, std::move(options.params), get_device_name());
+    auto cos = compile_hip_src(srcs, options.params, get_device_name());
     if(cos.size() != 1)
         MIGRAPHX_THROW("No code object");
     return code_object_op{value::binary{cos.front()},

src/targets/gpu/device/include/migraphx/gpu/device/scan.hpp

@@ -43,24 +43,32 @@ template <index_int N,
 __device__ void block_scan(index idx, Op op, T init, ForStride fs, Input input, Output output)
 {
     using type = decltype(input(deduce_for_stride(fs)));
-    MIGRAPHX_DEVICE_SHARED type buffer[N];
+    MIGRAPHX_DEVICE_SHARED type buffer[2][N];
     type x = init;
     fs([&](auto i) {
+        index_int iout = 0;
+        index_int iin  = 1;
         if(idx.local == 0)
-            buffer[idx.local] = op(input(i), x);
+            buffer[iout][idx.local] = op(input(i), x);
         else
-            buffer[idx.local] = input(i);
+            buffer[iout][idx.local] = input(i);
         __syncthreads();
         for(index_int s = 1; s < idx.nlocal(); s *= 2)
         {
-            if(idx.local + s < idx.nlocal())
+            iout = 1 - iout;
+            iin  = 1 - iin;
+            if(idx.local >= s)
             {
-                buffer[idx.local + s] = op(buffer[idx.local], buffer[idx.local + s]);
+                buffer[iout][idx.local] = op(buffer[iin][idx.local], buffer[iin][idx.local - s]);
+            }
+            else
+            {
+                buffer[iout][idx.local] = buffer[iin][idx.local];
             }
             __syncthreads();
         }
-        x = buffer[idx.nlocal() - 1];
-        output(i, buffer[idx.local]);
+        x = buffer[iout][idx.nlocal() - 1];
+        output(i, buffer[iout][idx.local]);
     });
 }
 

src/targets/gpu/driver/precompile_op.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/gpu/driver/action.hpp>
+#include <migraphx/gpu/time_op.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/gpu/lowering.hpp>
+#include <migraphx/gpu/compile_ops.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/instruction.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace driver {
+
+struct precompile_op : action<precompile_op>
+{
+    static program create_preop_program(const operation& preop, std::vector<shape> inputs)
+    {
+        program p;
+        auto* mm = p.get_main_module();
+        std::vector<instruction_ref> args;
+        inputs.pop_back();
+        transform(inputs, range(inputs.size()), std::back_inserter(args), [&](auto input, auto i) {
+            return mm->add_parameter("x" + std::to_string(i), input);
+        });
+        mm->add_instruction(preop, args);
+        return p;
+    }
+
+    static operation get_code_object(const program& p)
+    {
+        MIGRAPHX_TIDY_CONST auto* mm = p.get_main_module();
+        auto it                      = std::find_if(mm->begin(), mm->end(), [](const auto& ins) {
+            return (ins.name() == "gpu::code_object");
+        });
+        if(it == mm->end())
+            MIGRAPHX_THROW("Failed to create code object");
+        return it->get_operator();
+    }
+    static void apply(const parser& p, const value& v)
+    {
+        context ctx;
+        auto inputs = p.parse_shapes(v.at("inputs"));
+        auto name   = v.at("name").to<std::string>();
+        auto preop  = make_op(name);
+        if(v.contains("fields"))
+            preop.from_value(v.at("fields"));
+        bool exhaustive = v.get("exhaustive", false);
+        auto prog       = create_preop_program(preop, inputs);
+        run_passes(prog, {lowering{}, compile_ops{&ctx, exhaustive}});
+        auto op = get_code_object(prog);
+        auto t  = time_op(ctx, op, inputs, p.get(v, "iterations", 100));
+        std::cout << preop << ": " << t << "ms" << std::endl;
+    }
+};
+
+} // namespace driver
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/include/migraphx/gpu/compile_hip.hpp

@@ -58,10 +58,10 @@ struct hiprtc_src_file
 MIGRAPHX_GPU_EXPORT bool hip_has_flags(const std::vector<std::string>& flags);
 
 MIGRAPHX_GPU_EXPORT std::vector<std::vector<char>> compile_hip_src_with_hiprtc(
-    std::vector<hiprtc_src_file> srcs, std::string params, const std::string& arch);
+    std::vector<hiprtc_src_file> srcs, const std::string& params, const std::string& arch);
 
-MIGRAPHX_GPU_EXPORT std::vector<std::vector<char>>
-compile_hip_src(const std::vector<src_file>& srcs, std::string params, const std::string& arch);
+MIGRAPHX_GPU_EXPORT std::vector<std::vector<char>> compile_hip_src(
+    const std::vector<src_file>& srcs, const std::string& params, const std::string& arch);
 
 MIGRAPHX_GPU_EXPORT std::string enum_params(std::size_t count, std::string param);
 

src/targets/gpu/include/migraphx/gpu/fuse_mlir.hpp

@@ -34,6 +34,7 @@ struct module_pass_manager;
 namespace gpu {
 
 MIGRAPHX_GPU_EXPORT bool mlir_enabled();
+MIGRAPHX_GPU_EXPORT bool mlir_attention_enabled();
 
 struct MIGRAPHX_GPU_EXPORT fuse_mlir
 {

src/targets/gpu/include/migraphx/gpu/gemm_softmax_gemm.hpp

@@ -66,6 +66,10 @@ struct gemm_softmax_gemm
     }
 
     static bool is_ck_supported_type(shape::type_t t) { return contains({shape::half_type}, t); }
+    static bool is_mlir_supported_type(shape::type_t t)
+    {
+        return contains({shape::type_t::float_type, shape::half_type}, t);
+    }
 };
 
 } // namespace gpu

src/targets/gpu/include/migraphx/gpu/miopen.hpp

@@ -211,6 +211,12 @@ inline pooling_descriptor make_pooling(const migraphx::op::pooling& op)
         ss << op.mode;
         MIGRAPHX_THROW(ss.str());
     }
+    if(not std::all_of(
+           op.dilations.cbegin(), op.dilations.cend(), [](std::size_t d) { return d == 1; }))
+    {
+        MIGRAPHX_THROW("Unsupported dilations for pooling: [" + to_string_range(op.dilations) +
+                       "]");
+    }
     auto p = make_obj<pooling_descriptor>(&miopenCreatePoolingDescriptor);
 
     int kdims = op.kdims();