diff --git a/src/include/migraphx/dim_like.hpp b/src/include/migraphx/dim_like.hpp index d3933fe2a40..cb8a008e02d 100644 --- a/src/include/migraphx/dim_like.hpp +++ b/src/include/migraphx/dim_like.hpp @@ -24,14 +24,17 @@ #ifndef MIGRAPHX_GUARD_MIGRAPHLIB_DIM_LIKE_HPP #define MIGRAPHX_GUARD_MIGRAPHLIB_DIM_LIKE_HPP +#include #include #include #include +#include #include #include #include #include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { @@ -59,6 +62,36 @@ inline std::ostream& operator<<(std::ostream& os, const dim_like& d) return os; } +// Extracts the concrete int64_t from each entry; throws (via std::get) if any +// entry holds a symbolic dynamic_dimension. +inline std::vector to_ints(const std::vector& dims) +{ + std::vector result(dims.size()); + std::transform(dims.begin(), dims.end(), result.begin(), [](const dim_like& d) { + return std::get(d); + }); + return result; +} + +inline std::vector to_sym_exprs(const std::vector& dims) +{ + std::vector result(dims.size()); + std::transform(dims.begin(), dims.end(), result.begin(), [](const dim_like& d) -> sym::expr { + if(std::holds_alternative(d)) + return std::get(d).sym_expr; + return sym::lit(std::get(d)); + }); + return result; +} + +// Check if any of the dim_like are a shape::dynamic_dimension. +inline bool any_sym(const std::vector& dims) +{ + return std::any_of(dims.begin(), dims.end(), [](const dim_like& d) { + return std::holds_alternative(d); + }); +} + MIGRAPHX_EXPORT void migraphx_to_value(value& v, const dim_like& d); MIGRAPHX_EXPORT void migraphx_from_value(const value& v, dim_like& d); diff --git a/src/include/migraphx/enum.hpp b/src/include/migraphx/enum.hpp index a2a59e7a42d..19bc7171a4c 100644 --- a/src/include/migraphx/enum.hpp +++ b/src/include/migraphx/enum.hpp @@ -26,6 +26,7 @@ #include #include +#include #include #include #include @@ -141,6 +142,76 @@ struct is_named_enum +struct is_bit_flag : std::false_type +{ +}; + +template +struct is_bit_flag()))>> + : std::true_type +{ +}; + +// Type-safe bitwise operators for enums declared with MIGRAPHX_BIT_FLAG_ENUM. They operate on the +// enum's underlying integer and return the enum type, so combining values never leaks to a raw +// integer and different flag enums cannot be mixed. +template {})> +constexpr E operator|(E lhs, E rhs) +{ + using U = std::underlying_type_t; + return static_cast(static_cast(lhs) | static_cast(rhs)); +} + +template {})> +constexpr E operator&(E lhs, E rhs) +{ + using U = std::underlying_type_t; + return static_cast(static_cast(lhs) & static_cast(rhs)); +} + +template {})> +constexpr E operator^(E lhs, E rhs) +{ + using U = std::underlying_type_t; + return static_cast(static_cast(lhs) ^ static_cast(rhs)); +} + +template {})> +constexpr E operator~(E val) +{ + using U = std::underlying_type_t; + return static_cast(~static_cast(val)); +} + +template {})> +constexpr E& operator|=(E& lhs, E rhs) +{ + return lhs = lhs | rhs; +} + +template {})> +constexpr E& operator&=(E& lhs, E rhs) +{ + return lhs = lhs & rhs; +} + +template {})> +constexpr E& operator^=(E& lhs, E rhs) +{ + return lhs = lhs ^ rhs; +} + +// Returns true when every bit set in flag is also set in val. +template {})> +constexpr bool has_flag(E val, E flag) +{ + using U = std::underlying_type_t; + return (static_cast(val) & static_cast(flag)) == static_cast(flag); +} + // Returns the array of enumerator values for an enum declared with MIGRAPHX_ENUM. template {})> auto enum_entries() @@ -285,4 +356,43 @@ Enum from_string(const std::string& name) MIGRAPHX_DETAIL_ENUM_HELPERS( \ friend, name, MIGRAPHX_DETAIL_ENUM_CLASS_CAPTURE, using enum_scope = name;, __VA_ARGS__) +// Emits the ADL hook that marks an enum as a bit flag, which is_bit_flag detects to enable the +// bitwise operators. `linkage` is `inline` at namespace scope or `friend` at class scope, mirroring +// MIGRAPHX_DETAIL_ENUM_HELPERS. +// NOLINTNEXTLINE(cppcoreguidelines-macro-usage) +#define MIGRAPHX_DETAIL_BIT_FLAG_ENUM(linkage, name) \ + linkage constexpr bool migraphx_is_bit_flag(name) { return true; } + +// Declares a scoped enum (enum class) with the given underlying type and enables the type-safe +// bitwise operators |, &, ^, ~, |=, &=, ^= and has_flag() on it. Use it at namespace scope: +// +// MIGRAPHX_BIT_FLAG_ENUM(access, std::uint8_t, +// none = 0, +// read = 1 << 0, +// write = 1 << 1) +// +// auto rw = access::read | access::write; +// if(has_flag(rw, access::read)) { /* ... */ } +// +// Unlike MIGRAPHX_ENUM_CLASS, no to_string/from_string helpers are generated, so enumerator values +// should be self-contained bit masks. +// NOLINTNEXTLINE(cppcoreguidelines-macro-usage) +#define MIGRAPHX_BIT_FLAG_ENUM(name, type, ...) \ + enum class name : type \ + { \ + __VA_ARGS__ \ + }; \ + MIGRAPHX_DETAIL_BIT_FLAG_ENUM(inline, name) + +// Like MIGRAPHX_BIT_FLAG_ENUM, but for a scoped enum declared inside a class or struct; the hook is +// generated as a hidden friend so argument-dependent lookup still finds it. Use it inside the +// class/struct body. +// NOLINTNEXTLINE(cppcoreguidelines-macro-usage) +#define MIGRAPHX_NESTED_BIT_FLAG_ENUM(name, type, ...) \ + enum class name : type \ + { \ + __VA_ARGS__ \ + }; \ + MIGRAPHX_DETAIL_BIT_FLAG_ENUM(friend, name) + #endif // MIGRAPHX_GUARD_MIGRAPHX_ENUM_HPP diff --git a/src/include/migraphx/op/slice.hpp b/src/include/migraphx/op/slice.hpp index 2654561c42b..34dc867f544 100644 --- a/src/include/migraphx/op/slice.hpp +++ b/src/include/migraphx/op/slice.hpp @@ -27,72 +27,73 @@ #include #include #include +#include #include #include #include #include +#include #include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { namespace op { -/** - * Slice operator that accepts variable axes, starts and ends. - * All of `starts`, `ends`, and `axes` must be supplied by either - * their attribute or an input (but not both). - * - * Valid calls: - * slice(input); axes, starts, ends set - * slice(input, starts); axes, ends set - * slice(input, ends); starts, axes set - * slice(input, axes); starts, ends set - * slice(input, starts, ends); axes set - * slice(input, starts, axes); ends set - * slice(input, ends, axes); starts set - * slice(input, start, ends, axes); none set - * - * Attributes: - * axes: constant axes to slice over (optional) - * starts: constant slice starting indices (optional) - * ends: constant slice ending indices (optional) - * - * Parameters: - * data: the input tensor to slice (dynamic or static shape) - * input_starts: starting indices of slice (optional, static shape) - * input_ends: ending indices of slice (optional, static shape) - * input_axes: axes to slice over (optional, static shape) - */ +/// Slice operator that accepts variable axes, starts and ends. +/// All of `starts`, `ends`, and `axes` attributes must be supplied. +/// +/// `mode` specifies what the inputs to slice are: +/// one_input: slice(input); +/// starts_input: slice(input, starts); +/// ends_input: slice(input, ends); +/// axes_input: slice(input, axes); +/// starts_ends_input: slice(input, starts, ends); +/// starts_axes_input: slice(input, starts, axes); +/// ends_axes_input: slice(input, ends, axes); +/// starts_ends_axes_input: slice(input, start, ends, axes); +/// +/// Attributes: +/// axes: axes to slice over +/// starts: slice starting indices +/// ends: slice ending indices +/// +/// Parameters: +/// data: the input tensor to slice (dynamic or static shape) +/// starts_input: starting indices of slice (optional, static shape) +/// ends_input: ending indices of slice (optional, static shape) +/// axes_input: axes to slice over (optional, static shape) struct slice { - std::vector axes{}; - std::vector starts{}; - std::vector ends{}; + MIGRAPHX_NESTED_ENUM_CLASS(slice_mode, + one_input, + starts_input, + ends_input, + axes_input, + starts_ends_input, + starts_axes_input, + ends_axes_input, + starts_ends_axes_input); + + friend std::ostream& operator<<(std::ostream& os, slice_mode v) { return os << to_string(v); } - /** - * Named arrays for the set attribute possibilities. - */ - static constexpr std::array all_set = {true, true, true}; - static constexpr std::array ends_axes = {false, true, true}; - static constexpr std::array starts_axes = {true, false, true}; - static constexpr std::array starts_ends = {true, true, false}; - static constexpr std::array axes_only = {false, false, true}; - static constexpr std::array ends_only = {false, true, false}; - static constexpr std::array starts_only = {true, false, false}; - static constexpr std::array none_set = {false, false, false}; + std::vector axes{}; + std::vector starts{}; + std::vector ends{}; + slice_mode mode = slice_mode::one_input; template static auto reflect(Self& self, F f) { - return pack(f(self.axes, "axes"), f(self.starts, "starts"), f(self.ends, "ends")); + return pack(f(self.axes, "axes"), + f(self.starts, "starts"), + f(self.ends, "ends"), + f(self.mode, "mode")); } - /** - * Ensure that attribute axes is within limits. - * Will attempt to normalize starts and ends; but will use the dynamic_dimension.max - * values for dynamic shapes. This makes it so you have to renormalize for - * non-fixed dynamic_dimensions. - */ + /// Ensure that attribute axes is within limits. + /// Will attempt to normalize starts and ends; but will use the dynamic_dimension.max + /// values for dynamic shapes. This makes it so you have to renormalize for + /// non-fixed dynamic_dimensions. value attributes() const { value normalize_axes = value::object{}; @@ -112,12 +113,10 @@ struct slice std::string name() const { return "slice"; } - /** - * Computes the slice output shape dimensions for given starts, ends,and axes. - * Templated to also handle tensor views. - * Possibly different type between [in_starts, in_ends] and [in_axes] if in_axes is this - * object's axes attribute. Assumes in_starts and in_ends are normalized; in_axes are valid. - */ + /// Computes the slice output shape dimensions for given starts, ends,and axes. + /// Templated to also handle tensor views. + /// Possibly different type between [in_starts, in_ends] and [in_axes] if in_axes is this + /// object's axes attribute. Assumes in_starts and in_ends are normalized; in_axes are valid. template std::vector lens_calc(const std::vector& lengths, A in_starts, A in_ends, B in_axes) const @@ -131,192 +130,195 @@ struct slice return new_lens; } - /// Get the attributes that are non-empty - std::array get_set_attributes() const + /// Check that the inputs, attributes, and mode are valid. + void check_inputs_and_attributes(std::vector inputs) const { - std::array, 3> attrs = {this->starts, this->ends, this->axes}; - std::array bool_vec; - std::transform(attrs.cbegin(), attrs.cend(), bool_vec.begin(), [](const auto& a) { - return not a.empty(); - }); - return bool_vec; - } - - /// Helper function for normalize_compute_shape() - shape compute_two_or_more(std::vector inputs) const - { - auto input_shape = inputs[0]; - auto set_attributes = get_set_attributes(); - // check that inputs [1, end) are all 1D, have the same - // dimension, and are static + auto input_shape = inputs[0]; + // All set (non-empty) bound attributes must agree on the number of sliced axes. + // A variable (input-provided) bound leaves its attribute empty, so empty attrs are skipped. + std::size_t attr_size = 0; + for(auto s : {axes.size(), starts.size(), ends.size()}) + { + if(s == 0) + continue; + if(attr_size == 0) + attr_size = s; + else if(s != attr_size) + MIGRAPHX_THROW("SLICE: set starts/ends/axes attributes must have the same length"); + } + if(inputs.size() == 1) + { + if(any_sym(starts) or any_sym(ends)) + MIGRAPHX_THROW("SLICE: Invalid attributes: symbolic in attribute for 1 input slice"); + if(mode != slice_mode::one_input) + MIGRAPHX_THROW("SLICE: Invalid mode for 1 input"); + return; + } + // Check that inputs [1, end) are all 1D, have the same dimension, and are static shape. check_shapes{inputs.begin() + 1, inputs.end(), - std::string("SLICE: inputs (starts, ends, and input_axes)"), + std::string("SLICE: inputs (starts_input, ends_input, axes_input)"), false} .only_dims(1) .same_dims(); - auto dds = input_shape.to_dynamic().dyn_dims(); if(inputs.size() == 2) { - if(set_attributes == ends_axes) + std::vector two_input_modes_not_axes = {slice_mode::starts_input, slice_mode::ends_input}; + if(contains(two_input_modes_not_axes, mode)) { - // attr ends and axes set; inputs are (data, input_starts) - if(inputs[1].lens().at(0) != axes.size()) - { - MIGRAPHX_THROW("SLICE: 2 input and attributes mismatch: input_starts length (" + - to_string(inputs[1].lens().at(0)) + ") != number of axes (" + - to_string(axes.size()) + ")"); - } - std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) { - dds.at(axis) = {0, dds.at(axis).get_interval().max}; - }); + if(inputs[1].lens()[0] != axes.size()) + MIGRAPHX_THROW("SLICE: input length (" + migraphx::to_string(inputs[1].lens()[0]) + ") does not match attribute length (" + migraphx::to_string(axes.size()) + ")"); } - else if(set_attributes == starts_axes) + else if(mode == slice_mode::axes_input) { - // attr starts and axes set; inputs are (data, input_ends) - if(inputs[1].lens().at(0) != axes.size()) - { - MIGRAPHX_THROW("SLICE: 2 input and attributes mismatch: input_ends length (" + - to_string(inputs[1].lens().at(0)) + ") != number of axes (" + - to_string(axes.size()) + ")"); - } - std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) { - dds.at(axis) = {0, dds.at(axis).get_interval().max}; - }); - } - else if(set_attributes == starts_ends) - { - // attr starts and ends set; inputs are (data, input_axes) - if(inputs[1].lens().at(0) != starts.size()) - { - MIGRAPHX_THROW("SLICE: 2 input and attributes mismatch: input_axes length (" + - to_string(inputs[1].lens().at(0)) + ") != number of starts (" + - to_string(starts.size()) + ")"); - } - std::transform(dds.begin(), dds.end(), dds.begin(), [](const auto& dd) { - return shape::dynamic_dimension{0, dd.get_interval().max}; - }); + if(inputs[1].lens()[0] != starts.size()) + MIGRAPHX_THROW("SLICE: input length (" + migraphx::to_string(inputs[1].lens()[0]) + ") does not match attribute length (" + migraphx::to_string(starts.size()) + ")"); } else { - MIGRAPHX_THROW("SLICE: Invalid 2 input and attributes configuration"); + MIGRAPHX_THROW("SLICE: Invalid mode for 2 inputs"); } } else if(inputs.size() == 3) { - if(set_attributes == axes_only) + if(mode == slice_mode::starts_ends_input) { - // attr axes set; inputs are (data, input_starts, input_ends) - if(inputs[1].lens().at(0) != axes.size()) - { - MIGRAPHX_THROW("SLICE: 3 input and attributes mismatch: input_starts length (" + - to_string(inputs[1].lens().at(0)) + ") != number of axes (" + - to_string(axes.size()) + ")"); - } - std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) { - dds.at(axis) = {0, dds.at(axis).get_interval().max}; - }); + if(inputs[1].lens()[0] != axes.size()) + MIGRAPHX_THROW("SLICE: input length (" + migraphx::to_string(inputs[1].lens()[0]) + ") does not match attribute length (" + migraphx::to_string(axes.size()) + ")"); } - else if(set_attributes == ends_only) + else if(mode == slice_mode::starts_axes_input) { - // attr ends set; inputs are (data, input_starts, input_axes) - if(inputs[1].lens().at(0) != ends.size()) - { - MIGRAPHX_THROW("SLICE: 3 input and attributes mismatch: input_starts length (" + - to_string(inputs[1].lens().at(0)) + ") != number of ends (" + - to_string(ends.size()) + ")"); - } - std::transform(dds.begin(), dds.end(), dds.begin(), [](const auto& dd) { - return shape::dynamic_dimension{0, dd.get_interval().max}; - }); + if(inputs[1].lens()[0] != ends.size()) + MIGRAPHX_THROW("SLICE: input length (" + migraphx::to_string(inputs[1].lens()[0]) + ") does not match attribute length (" + migraphx::to_string(ends.size()) + ")"); } - else if(set_attributes == starts_only) - + else if(mode == slice_mode::ends_axes_input) { - // attr starts set; inputs are (data, input_ends, input_axes) - if(inputs[1].lens().at(0) != starts.size()) - { - MIGRAPHX_THROW("SLICE: 3 input and attributes mismatch: input_ends length (" + - to_string(inputs[1].lens().at(0)) + ") != number of starts (" + - to_string(starts.size()) + ")"); - } - std::transform(dds.begin(), dds.end(), dds.begin(), [](const auto& dd) { - return shape::dynamic_dimension{0, dd.get_interval().max}; - }); + if(inputs[1].lens()[0] != starts.size()) + MIGRAPHX_THROW("SLICE: input length (" + migraphx::to_string(inputs[1].lens()[0]) + ") does not match attribute length (" + migraphx::to_string(starts.size()) + ")"); } else { - MIGRAPHX_THROW("Invalid 3 input and attributes configuration"); + MIGRAPHX_THROW("SLICE: Invalid mode for 3 inputs"); + } + } + else + { + if(mode != slice_mode::starts_ends_axes_input) + { + MIGRAPHX_THROW("SLICE: Invalid mode for 4 inputs"); } } + return; + } + + // TODO: remove this once range-based dynamic shapes are deprecated + bool use_range_based_logic() const + { + if(mode == slice_mode::starts_input and starts.empty()) + return true; + else if(mode == slice_mode::ends_input and ends.empty()) + return true; + else if(mode == slice_mode::axes_input and axes.empty()) + return true; + else if(mode == slice_mode::starts_ends_input and starts.empty() and ends.empty()) + return true; + else if(mode == slice_mode::starts_axes_input and starts.empty() and axes.empty()) + return true; + else if(mode == slice_mode::ends_axes_input and ends.empty() and axes.empty()) + return true; + else if(mode == slice_mode::starts_ends_axes_input and starts.empty() and ends.empty() and axes.empty()) + return true; else + return false; + } + + // For when there is a variable input and the associated attribute is not set. + // ex: slice(data, starts) starts = {}, ends = {2, 3}, axes = {0, 1} + // TODO: remove this once range-based dynamic shapes are deprecated + shape range_based_compute_shape_for_two_or_more(shape input_shape) const + { + auto dds = input_shape.to_dynamic().dyn_dims(); + static std::vector has_axes_input = { + slice_mode::axes_input, + slice_mode::starts_axes_input, + slice_mode::ends_axes_input, + slice_mode::starts_ends_axes_input + }; + if(contains(has_axes_input, mode)) { - // all 4 inputs (data, inputs_starts, input_ends, input_axes) std::transform(dds.begin(), dds.end(), dds.begin(), [](const auto& dd) { return shape::dynamic_dimension{0, dd.get_interval().max}; }); } + + std::for_each(axes.cbegin(), axes.cend(), [&](const auto& axis) { + dds.at(axis) = {0, dds.at(axis).get_interval().max}; + }); return shape{input_shape.type(), dds}; } + + // Static and symbolic inputs share this path; the result is demoted back to + // static when fully fixed (slice is a view). + shape symbolic_compute_shape(const shape& s) const + { + if(starts.size() != axes.size() or ends.size() != axes.size()) + MIGRAPHX_THROW("SLICE: Attribute sizes do not match for symbolic_compute_shape()"); + auto sym_in = s.to_symbolic(); + auto dds = sym_in.dyn_dims(); + auto start_exprs = to_sym_exprs(starts); + auto end_exprs = to_sym_exprs(ends); + for(std::size_t i = 0; i < axes.size(); ++i) + dds[axes[i]] = shape::dynamic_dimension{end_exprs[i] - start_exprs[i]}; + shape result{s.type(), std::move(dds), sym_in.dyn_strides()}; + if(not s.symbolic() and result.is_fixed()) + return result.to_static(); + return result; + } + // uses the normalize_axes flag to normalize axes, starts, and ends shape normalize_compute_shape(std::vector inputs) const { check_shapes{inputs, *this, true}.has(1, 2, 3, 4); - if(inputs.size() != 1) - return compute_two_or_more(inputs); - + check_inputs_and_attributes(inputs); auto input_shape = inputs[0]; - auto set_attributes = get_set_attributes(); - if(set_attributes != all_set) - MIGRAPHX_THROW("SLICE 1_arg: Invalid 1 input and attributes configuration"); - - // TODO: support slicing non-fixed symbolic dims (output dim would be - // a sym::expr derived from starts/ends and the symbolic axis bound). - if(input_shape.dynamic() and std::any_of(axes.begin(), axes.end(), [&](auto axis) { - return not input_shape.dyn_dims()[axis].is_fixed(); - })) + if(inputs.size() == 1 and input_shape.dynamic() and not input_shape.symbolic()) { - if(input_shape.symbolic()) - { - MIGRAPHX_THROW( - "SLICE 1_arg: slicing is not allowed on non-fixed symbolic input axis "); - } - // Attributes are not normalized for this case, so they can be negative or - // out-of-bounds. Using a relaxed dimension bound for now instead of calculating the - // tightest possible bound. - auto dds = input_shape.dyn_dims(); - for(auto axis : this->axes) + // Fallback for range-based dynamic shapes. + // Non-fixed sliced axis: bounds aren't normalized (can be negative or + // out-of-bounds), so use a relaxed [0, max] bound. (#5015) + // TODO: remove this once range-based dynamic shapes are deprecated + if(std::any_of(axes.begin(), axes.end(), [&](auto axis) { + return not input_shape.dyn_dims()[axis].is_fixed(); + })) { - dds[axis] = {0, dds[axis].get_interval().max}; + auto dds = input_shape.dyn_dims(); + for(auto axis : axes) + dds[axis] = {0, dds[axis].get_interval().max}; + return shape{input_shape.type(), dds}; } + + auto new_lens = lens_calc(input_shape.max_lens(), to_ints(starts), to_ints(ends), axes); + auto dds = input_shape.dyn_dims(); + for(auto axis : axes) + dds[axis] = shape::dynamic_dimension{new_lens[axis], new_lens[axis]}; return shape{input_shape.type(), dds}; } - - auto new_lens = lens_calc(input_shape.max_lens(), this->starts, this->ends, this->axes); - - if(not input_shape.dynamic()) - return shape{input_shape.type(), new_lens, input_shape.strides()}; - - auto dds = input_shape.dyn_dims(); - for(auto axis : this->axes) + else if(inputs.size() > 1 and use_range_based_logic()) { - dds[axis] = input_shape.symbolic() - ? shape::dynamic_dimension{sym::lit(new_lens[axis])} - : shape::dynamic_dimension{new_lens[axis], new_lens[axis]}; + // TODO: remove this once range-based dynamic shapes are deprecated + return range_based_compute_shape_for_two_or_more(input_shape); + } + else + { + return symbolic_compute_shape(input_shape); } - - if(input_shape.symbolic()) - return shape{input_shape.type(), dds, input_shape.dyn_strides()}; - return shape{input_shape.type(), dds}; } - /** - * Calculates the starting offset for the sliced tensor. - * Used in compute when only data input and all other information are in the attributes. - * - * \param s static input shape - */ + /// Calculates the starting offset for the sliced tensor. + /// Used in compute when only data input and all other information are in the attributes. + /// + /// s: static input shape auto compute_offset(const shape& s) const { const std::vector& lens = s.lens(); @@ -327,85 +329,85 @@ struct slice for(std::size_t i = 0; i < axes.size(); i++) { auto axis = axes[i]; - offset += starts[i] * strides[axis]; + offset += std::get(starts[i]) * strides[axis]; } } else { for(std::size_t axis = 0; axis < lens.size(); axis++) { - offset += starts[axis] * strides[axis]; + offset += std::get(starts[axis]) * strides[axis]; } } return offset * s.type_size(); } - /** - * Calculates the starting offset for the sliced tensor (for aliasing). - * Used for 2-4 inputs to `slice. - * - * \param s static input shape - * \param input_starts starting indices of slice - * \param ax_vec axes to slice on - */ + /// Calculates the starting offset for the sliced tensor (for aliasing). + /// Used for 2-4 inputs to `slice. + /// + /// s: static input shape + /// starts_input: starting indices of slice + /// ax_vec: axes to slice on template - auto compute_offset(const shape& s, const T& input_starts, const T& ax_vec) const + auto compute_offset(const shape& s, const T& starts_input, const T& ax_vec) const { auto ret = 0; for(std::size_t i = 0; i < ax_vec.size(); ++i) { auto axis = ax_vec[i]; - ret += input_starts[i] * s.strides().at(axis); + ret += starts_input[i] * s.strides().at(axis); } return ret * s.type_size(); } - /** - * If given, normalize the inputs. Otherwise get from operator attributes. - * Return the values in a map. - * - * Parameters - * input_shape: static shape of the input - * input_starts: optional - * input_ends: optional - * input_ends: optional - */ + /// Used to normalize starts/ends/axes at runtime. + /// If given, normalize the starts/ends/axes inputs. Otherwise get from operator attributes. + /// If starts_input or ends_input is not given, assuming starts/ends attributes are only + /// integers. If starts/ends have symbolics, they should go through the starts_input and + /// ends_input instead. `axes` attribute should always be correctly normalized at compile-time + /// because shapes with dynamic rank are not supported. + /// + /// input_shape: static shape of the input + /// starts_input: optional + /// ends_input: optional + /// axes_input: optional std::unordered_map> normalize_starts_ends_axes(shape input_shape, - const optional>& input_starts, - const optional>& input_ends, - const optional>& input_axes) const + const optional>& starts_input, + const optional>& ends_input, + const optional>& axes_input) const { + assert(not input_shape.dynamic()); auto axes_attrs = this->attributes().at("normalize_axes"); std::vector norm_starts; std::vector norm_ends; std::vector norm_axes; - if(input_axes) + if(axes_input) { - norm_axes = normalize_axes(input_axes.value(), + norm_axes = normalize_axes(axes_input.value(), input_shape, axes_attrs.at("axes"), - "Slice variable input_axes"); + "Slice variable axes_input"); } else { norm_axes = this->axes; } - if(input_starts) + if(starts_input) { - norm_starts = normalize_indices(input_starts.value(), + norm_starts = normalize_indices(starts_input.value(), norm_axes, input_shape, axes_attrs.at("starts"), - "Slice variable input_starts"); + "Slice variable starts_input"); } else { - norm_starts = this->starts; + norm_starts = to_ints(this->starts); } - if(input_ends) + if(ends_input) { - norm_ends = normalize_indices(input_ends.value(), + norm_ends = normalize_indices(ends_input.value(), norm_axes, input_shape, axes_attrs.at("ends"), @@ -413,7 +415,7 @@ struct slice } else { - norm_ends = this->ends; + norm_ends = to_ints(this->ends); } return {{"norm_starts", norm_starts}, {"norm_ends", norm_ends}, {"norm_axes", norm_axes}}; } @@ -429,87 +431,80 @@ struct slice } else { - // Note that we re-normalize both the attributes and inputs because of the non-fixed - // dynamic input shape case. It's possible to only re-normalize if slicing over - // non-fixed dynamic_dimensions. - auto set_attributes = get_set_attributes(); std::unordered_map> norm_inputs; - if(set_attributes == ends_axes) + // Attribute-provided dims are passed as their concrete int values (via to_ints) so they + // are re-normalized/clamped against the runtime input shape, just like the runtime + // inputs. Symbolic bounds only ever appear on the input-provided dims. + if(mode == slice_mode::starts_input) { - // attr ends and axes set; inputs are (data, input_starts) - args[1].visit([&](auto input_starts) { + args[1].visit([&](auto starts_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - input_starts.template to_vector(), - this->ends, + starts_input.template to_vector(), + to_ints(this->ends), this->axes); }); } - else if(set_attributes == starts_axes) + else if(mode == slice_mode::ends_input) { - // attr starts and axes set; inputs are (data, input_ends) - args[1].visit([&](auto input_ends) { + args[1].visit([&](auto ends_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - this->starts, - input_ends.template to_vector(), + to_ints(this->starts), + ends_input.template to_vector(), this->axes); }); } - else if(set_attributes == starts_ends) + else if(mode == slice_mode::axes_input) { - // attr starts and ends set; inputs are (data, input_axes) - args[1].visit([&](auto input_axes) { + args[1].visit([&](auto axes_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - this->starts, - this->ends, - input_axes.template to_vector()); + to_ints(this->starts), + to_ints(this->ends), + axes_input.template to_vector()); }); } - else if(set_attributes == axes_only) + else if(mode == slice_mode::starts_ends_input) { - // attr axes set; inputs are (data, input_starts, input_ends) - visit_all(args[1], args[2])([&](auto input_starts, auto input_ends) { + // attr axes set; inputs are (data, starts_input, ends_input) + visit_all(args[1], args[2])([&](auto starts_input, auto ends_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - input_starts.template to_vector(), - input_ends.template to_vector(), + starts_input.template to_vector(), + ends_input.template to_vector(), this->axes); }); } - else if(set_attributes == ends_only) + else if(mode == slice_mode::starts_axes_input) { - // attr ends set; inputs are (data, input_starts, input_axes) - visit_all(args[1], args[2])([&](auto input_starts, auto input_axes) { + visit_all(args[1], args[2])([&](auto starts_input, auto axes_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - input_starts.template to_vector(), - this->ends, - input_axes.template to_vector()); + starts_input.template to_vector(), + to_ints(this->ends), + axes_input.template to_vector()); }); } - else if(set_attributes == starts_only) + else if(mode == slice_mode::ends_axes_input) { - // attr starts set; inputs are (data, input_ends, input_axes) - visit_all(args[1], args[2])([&](auto input_ends, auto input_axes) { + visit_all(args[1], args[2])([&](auto ends_input, auto axes_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - this->starts, - input_ends.template to_vector(), - input_axes.template to_vector()); + to_ints(this->starts), + ends_input.template to_vector(), + axes_input.template to_vector()); }); } - else + else // mode == slice_mode::starts_ends_axes_input { - // no attr set, all inputs visit_all(args[1], args[2], args[3])( - [&](auto input_starts, auto input_ends, auto input_axes) { + [&](auto starts_input, auto ends_input, auto axes_input) { norm_inputs = normalize_starts_ends_axes(input_shape, - input_starts.template to_vector(), - input_ends.template to_vector(), - input_axes.template to_vector()); + starts_input.template to_vector(), + ends_input.template to_vector(), + axes_input.template to_vector()); }); } auto offset = compute_offset( diff --git a/src/include/migraphx/op/topk.hpp b/src/include/migraphx/op/topk.hpp index 521004b170d..e2877f89561 100644 --- a/src/include/migraphx/op/topk.hpp +++ b/src/include/migraphx/op/topk.hpp @@ -33,19 +33,22 @@ #include #include #include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { namespace op { /** - * TopK with constant `k` value. Significantly different from ONNX spec's TopK. + * TopK with additional input compared to ONNX spec's TopK. + * Valid calls topk(x, k), topk(x, k, indexing); * arg[0]: input data - * arg[1]: optional indexing information used for rewrite_topk + * arg[1]: k value input + * arg[2]: optional indexing information used for rewrite_topk */ struct topk { - int64_t k = 1; + dim_like k = 1; int64_t axis = 0; bool largest = true; @@ -66,16 +69,17 @@ struct topk shape normalize_compute_shape(std::vector inputs) const { - check_shapes{inputs, *this, true}.has(1, 2); - auto type = inputs.at(0).type(); + check_shapes{inputs, *this, true}.has(2, 3); + auto type = inputs.at(0).type(); + auto k_val = std::get(k); if(inputs.at(0).dynamic()) { auto dyn_dims = inputs.at(0).dyn_dims(); auto min_lens_vec = inputs.at(0).min_lens(); auto max_lens_vec = inputs.at(0).max_lens(); - auto min_kk = std::min(k, min_lens_vec[axis]); - auto max_kk = std::min(k, max_lens_vec[axis]); + auto min_kk = std::min(k_val, min_lens_vec[axis]); + auto max_kk = std::min(k_val, max_lens_vec[axis]); dyn_dims[axis] = {min_kk, max_kk}; shape s_val{type, dyn_dims}; @@ -85,7 +89,7 @@ struct topk else { auto lens = inputs.at(0).lens(); - auto kk = std::min(k, lens[axis]); + auto kk = std::min(k_val, lens[axis]); lens[axis] = kk; shape s_val{type, lens}; @@ -114,14 +118,14 @@ struct topk argument res_ind{vec_ss.back()}; auto in_val = args.front(); auto relements = in_val.get_shape().lens()[axis]; - auto actual_k = std::min(k, relements); + auto actual_k = std::min(std::get(k), relements); auto make_indices = [&](const auto& m_idx) { return [&](int64_t i) { - if(args.size() < 2) + if(args.size() < 3) return i; auto j = m_idx; j[axis] = i; - return args[1].at(j); + return args.back().at(j); }; }; auto outer_lens = in_val.get_shape().lens(); diff --git a/src/normalize_attributes.cpp b/src/normalize_attributes.cpp index 48804c9034f..7f34340e9c2 100644 --- a/src/normalize_attributes.cpp +++ b/src/normalize_attributes.cpp @@ -25,11 +25,79 @@ #include #include #include +#include +#include +#include #include #include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { +// min/max that fold to one operand when the ordering is provable via intervals, +// and only fall back to a symbolic min/max node when it is indeterminate. +static sym::expr fold_min(const sym::expr& a, const sym::expr& b) +{ + auto lt = sym::strict_less(a, b); + if(lt.has_value()) + return *lt ? a : b; + return sym::min(a, b); +} +static sym::expr fold_max(const sym::expr& a, const sym::expr& b) +{ + auto lt = sym::strict_less(a, b); + if(lt.has_value()) + return *lt ? b : a; + return sym::max(a, b); +} + +static sym::expr axis_len_expr(const shape& s, int64_t axis) +{ + if(not s.dynamic()) + return sym::lit(static_cast(s.lens().at(axis))); + const auto& dd = s.dyn_dims().at(axis); + if(dd.is_symbolic()) + return dd.sym_expr; + if(dd.is_fixed()) + return sym::lit(static_cast(dd.get_interval().max)); + MIGRAPHX_THROW("normalize_attributes: cannot normalize a symbolic bound on a non-fixed axis"); +} + +static dim_like to_dim_like(const sym::expr& e) +{ + if(e.name() == "literal") + return sym::to(e.eval({})); + return shape::dynamic_dimension{e}; +} + +// Symbolic analog of tune_attribute for dim_like bounds (clip_min/clip_max + +// use_len). Applies the ONNX clamp norm(v) = clamp(v < 0 ? v + D : v, 0, D) +// symbolically, folding against the interval bounds where provable. +template +static value tune_attribute_sym(const std::vector& dims, + const std::vector& axes, + const std::vector& attrs, + const shape& input_shape, + Message m) +{ + if(not contains(attrs, op::normalize_attribute::use_len)) + MIGRAPHX_THROW(m() + "symbolic bounds are only supported with use_len normalization"); + if(axes.size() != dims.size()) + MIGRAPHX_THROW(m() + "symbolic bounds require one axis per bound"); + auto zero = sym::lit(std::int64_t{0}); + auto exprs = to_sym_exprs(dims); + std::vector result(dims.size()); + std::transform( + exprs.begin(), exprs.end(), axes.begin(), result.begin(), [&](const auto& v, auto axis) { + auto len = axis_len_expr(input_shape, axis); + auto neg = sym::strict_less(v, zero); // from-the-end (negative) index? + if(not neg.has_value()) + MIGRAPHX_THROW(m() + "symbolic bound of indeterminate sign cannot be normalized"); + auto abs_v = *neg ? v + len : v; + return to_dim_like(fold_min(fold_max(abs_v, zero), len)); + }); + return migraphx::to_value(result); +} + /** * Parameters: * vec: the vector attribute to normalize @@ -235,6 +303,22 @@ bool normalize_attributes(operation& op, const shape& input_shape) { axes = val.at("axes").without_key().to_vector(); } + // Symbolic (dim_like) bounds serialize as objects; clamp them + // symbolically rather than against a compile-time length. + if(std::any_of(vv.begin(), vv.end(), [](const auto& e) { return e.is_object(); })) + { + auto dims = migraphx::from_value>(vv); + val[key] = + tune_attribute_sym(dims, + axes, + rv.without_key().to_vector(), + input_shape, + message); + op.from_value(val); + val = op.to_value(); + tuned = true; + continue; + } auto vec = vv.to_vector(); auto result = tune_attribute(vec, axes, rv.without_key(), input_shape, message); val[key] = result; diff --git a/src/onnx/parse_nonmaxsuppression.cpp b/src/onnx/parse_nonmaxsuppression.cpp index eb422329f67..6387f855af4 100644 --- a/src/onnx/parse_nonmaxsuppression.cpp +++ b/src/onnx/parse_nonmaxsuppression.cpp @@ -24,9 +24,7 @@ #include #include #include -#include - -MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_USE_DYNAMIC_NMS) +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { @@ -45,18 +43,18 @@ struct parse_nonmaxsuppression : op_parser auto nms_ins = info.add_instruction(op, args); // slice with variable ends to handle dynamic shape output. auto indices = info.add_instruction(make_op("get_tuple_elem", {{"index", 0}}), nms_ins); - if(enabled(MIGRAPHX_USE_DYNAMIC_NMS{})) - { - // TODO: planning to make this the default behavior and removing the env var. - auto num_selected = - info.add_instruction(make_op("get_tuple_elem", {{"index", 1}}), nms_ins); - return info.add_instruction( - make_op("slice", {{"axes", {0}}, {"starts", {0}}}), indices, num_selected); - } - else - { - return indices; - } + auto num_selected = + info.add_instruction(make_op("get_tuple_elem", {{"index", 1}}), nms_ins); + auto num_selected_var = shape::dynamic_dimension{sym::var(info.name)}; + return info.add_instruction( + make_op( + "slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", value::array{to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + indices, + num_selected); } }; diff --git a/src/onnx/parse_slice.cpp b/src/onnx/parse_slice.cpp index 57e546694b2..72f071fb54b 100644 --- a/src/onnx/parse_slice.cpp +++ b/src/onnx/parse_slice.cpp @@ -1,7 +1,7 @@ /* * The MIT License (MIT) * - * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved. + * Copyright (c) 2015-2026 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 @@ -27,11 +27,20 @@ #include #include #include +#include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { namespace onnx { +MIGRAPHX_BIT_FLAG_ENUM(slice_input_flags, + std::uint8_t, + none = 0, + starts_input = 1 << 0, + ends_input = 1 << 1, + axes_input = 1 << 2) + struct parse_slice : op_parser { @@ -39,17 +48,18 @@ struct parse_slice : op_parser struct slice_desc { - op::slice op; std::vector op_args; + std::vector axes; + std::vector starts; + std::vector ends; std::vector steps; std::vector raxes; + slice_input_flags flags = slice_input_flags::none; void always_insert(instruction_ref arg) { op_args.insert(op_args.begin(), arg); } - /** - * Either insert argument into `this->op_args` or return the constant value of the argument - */ - std::vector insert(instruction_ref arg) + // Either insert argument into `this->op_args` or return the constant value of the argument + std::vector try_insert(instruction_ref arg) { std::vector result; migraphx::argument arg_value = arg->eval(); @@ -63,6 +73,41 @@ struct parse_slice : op_parser } return result; } + + op::slice::slice_mode get_slice_mode(slice_input_flags slice_flags) + { + switch (slice_flags) + { + case slice_input_flags::none: + return op::slice::slice_mode::one_input; + case slice_input_flags::starts_input: + return op::slice::slice_mode::starts_input; + case slice_input_flags::ends_input: + return op::slice::slice_mode::ends_input; + case slice_input_flags::axes_input: + return op::slice::slice_mode::axes_input; + case (slice_input_flags::starts_input | slice_input_flags::ends_input): + return op::slice::slice_mode::starts_ends_input; + case (slice_input_flags::starts_input | slice_input_flags::axes_input): + return op::slice::slice_mode::starts_axes_input; + case (slice_input_flags::ends_input | slice_input_flags::axes_input): + return op::slice::slice_mode::ends_axes_input; + case (slice_input_flags::starts_input | slice_input_flags::ends_input | slice_input_flags::axes_input): + return op::slice::slice_mode::starts_ends_axes_input; + default: + MIGRAPHX_THROW("PARSE_SLICE: invalid slice_mode"); + } + } + + op::slice create_slice_operator() + { + op::slice slice_op; + slice_op.axes = axes; + slice_op.starts = starts; + slice_op.ends = ends; + slice_op.mode = get_slice_mode(flags); + return slice_op; + } }; instruction_ref parse(const op_desc& /*opd*/, @@ -71,7 +116,7 @@ struct parse_slice : op_parser const std::vector& args) const { auto sd = construct_slice_desc(parser, info, args); - auto ins = info.add_instruction(sd.op, sd.op_args); + auto ins = info.add_instruction(sd.create_slice_operator(), sd.op_args); if(not sd.raxes.empty()) { ins = info.add_instruction(make_op("reverse", {{"axes", sd.raxes}}), ins); @@ -85,7 +130,7 @@ struct parse_slice : op_parser std::back_inserter(nsteps), [](auto s) { return std::abs(s); }); return ins = info.add_instruction( - make_op("step", {{"axes", sd.op.axes}, {"steps", nsteps}}), ins); + make_op("step", {{"axes", sd.axes}, {"steps", nsteps}}), ins); } else return ins; @@ -97,8 +142,8 @@ struct parse_slice : op_parser { slice_desc sd; - // slice can have up to 5 inputs, we first check the 5th one - // to decide whether MIGRAPHX can handle this slice. + // ONNX Slice can have up to 5 inputs, we first check the 5th one + // to decide whether MIGX can handle this slice. if(args.size() == 5) { migraphx::argument step_arg = args.back()->eval(); @@ -108,51 +153,69 @@ struct parse_slice : op_parser if(args.size() >= 4) { - sd.op.axes = sd.insert(args.at(3)); + auto _axes = sd.try_insert(args.at(3)); + if(_axes.empty()) + sd.flags |= slice_input_flags::axes_input; + sd.axes = _axes; } else if(contains(info.attributes, "axes")) { literal s = parser.parse_value(info.attributes.at("axes")); - s.visit([&](auto v) { copy(v, std::back_inserter(sd.op.axes)); }); + s.visit([&](auto v) { copy(v, std::back_inserter(sd.axes)); }); } + // NOTE: goes through range-based dynamic shapes pathway only if(args.size() >= 3) { - sd.op.ends = sd.insert(args.at(2)); + auto _ends = sd.try_insert(args.at(2)); + if(_ends.empty()) + sd.flags |= slice_input_flags::ends_input; + sd.ends.assign(_ends.begin(), _ends.end()); } else if(contains(info.attributes, "ends")) { literal s = parser.parse_value(info.attributes.at("ends")); - s.visit([&](auto v) { copy(v, std::back_inserter(sd.op.ends)); }); + s.visit([&](auto v) { + std::transform(v.begin(), v.end(), std::back_inserter(sd.ends), [](auto e) { + return static_cast(e); + }); + }); } if(args.size() >= 2) { - sd.op.starts = sd.insert(args.at(1)); + auto _starts = sd.try_insert(args.at(1)); + if(_starts.empty()) + sd.flags |= slice_input_flags::starts_input; + sd.starts.assign(_starts.begin(), _starts.end()); } else if(contains(info.attributes, "starts")) { literal s = parser.parse_value(info.attributes.at("starts")); - s.visit([&](auto v) { copy(v, std::back_inserter(sd.op.starts)); }); + s.visit([&](auto v) { + std::transform(v.begin(), v.end(), std::back_inserter(sd.starts), [](auto e) { + return static_cast(e); + }); + }); } // data input argument sd.always_insert(args.at(0)); // If axes arg is not given, the default is all of them. - if(sd.op.axes.empty() and sd.op_args.size() <= 3) + if(sd.axes.empty() and sd.op_args.size() <= 3) { std::vector axes(args[0]->get_shape().ndim()); std::iota(axes.begin(), axes.end(), int64_t{0}); - sd.op.axes = axes; + sd.axes = axes; } if(std::any_of(sd.steps.begin(), sd.steps.end(), [](auto s) { return s != 1; })) { - if(sd.op.starts.empty() or sd.op.ends.empty()) + if(sd.starts.empty() or sd.ends.empty()) MIGRAPHX_THROW( "PARSE_SLICE: steps and variable starts and/or ends is not supported"); - if(sd.op.axes.empty()) + if(sd.axes.empty()) MIGRAPHX_THROW("PARSE_SLICE: steps and variable axes is not supported"); } @@ -161,12 +224,13 @@ struct parse_slice : op_parser { if(sd.steps[i] >= 0) continue; - sd.op.starts[i] += 1; - if(sd.op.starts[i] == 0) - sd.op.starts[i] = INT_MAX; - sd.op.ends[i] += 1; - sd.raxes.push_back(sd.op.axes[i]); - std::swap(sd.op.starts[i], sd.op.ends[i]); + auto start = std::get(sd.starts[i]) + 1; + if(start == 0) + start = INT_MAX; + sd.starts[i] = start; + sd.ends[i] = std::get(sd.ends[i]) + 1; + sd.raxes.push_back(sd.axes[i]); + std::swap(sd.starts[i], sd.ends[i]); } return sd; } diff --git a/src/onnx/parse_topk.cpp b/src/onnx/parse_topk.cpp index 7481ddcb5e7..277c61981f1 100644 --- a/src/onnx/parse_topk.cpp +++ b/src/onnx/parse_topk.cpp @@ -27,6 +27,7 @@ #include #include #include +#include namespace migraphx { inline namespace MIGRAPHX_INLINE_NS { @@ -53,47 +54,66 @@ struct parse_topk : op_parser axis = parser.parse_value(info.attributes.at("axis")).at(); } - bool var_k = false; - int64_t k = 0; - if(args.size() == 2) + // opset-1 form: `k` is an attribute. Synthesize a constant `k` input so the topk + // operator always has (x, k) inputs. + if(args.size() == 1) { - auto arg_k = args.at(1)->eval(); - if(not arg_k.empty()) + int64_t k = 0; + if(contains(info.attributes, "k")) { - k = arg_k.at(); + k = info.attributes.at("k").i(); } - else - { - var_k = true; - } - } - else if(contains(info.attributes, "k")) - { - k = info.attributes.at("k").i(); + const shape k_shape{shape::int64_type, {1}}; + auto k_lit = info.add_literal(literal{k_shape, {k}}); + auto topk_ret = info.add_instruction( + make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0), k_lit); + + auto ret_val = info.add_instruction(make_op("get_tuple_elem", {{"index", 0}}), topk_ret); + auto ret_ind = info.add_instruction(make_op("get_tuple_elem", {{"index", 1}}), topk_ret); + return {ret_val, ret_ind}; } - if(var_k) + // opset-10+ form: `k` is a runtime input. + auto arg_k = args.at(1)->eval(); + if(not arg_k.empty()) { - // set `k` to axis dimension - auto input_shape = args.at(0)->get_shape(); - auto norm_axis = axis < 0 ? axis + input_shape.ndim() : axis; - k = input_shape.max_lens().at(norm_axis); + // Constant `k`: use its value for the attribute; topk output is already the exact size. + int64_t k = arg_k.at(); + auto topk_ret = info.add_instruction( + make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0), args.at(1)); + + auto ret_val = info.add_instruction(make_op("get_tuple_elem", {{"index", 0}}), topk_ret); + auto ret_ind = info.add_instruction(make_op("get_tuple_elem", {{"index", 1}}), topk_ret); + return {ret_val, ret_ind}; } + // Variable (data-dependent) `k`: run topk over the whole axis dimension, then slice the + // outputs down to the runtime `k` using a symbolic dimension. + auto input_shape = args.at(0)->get_shape(); + auto norm_axis = axis < 0 ? axis + input_shape.ndim() : axis; + int64_t k = input_shape.max_lens().at(norm_axis); + auto topk_ret = info.add_instruction( - make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0)); + make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0), args.at(1)); auto ret_val = info.add_instruction(make_op("get_tuple_elem", {{"index", 0}}), topk_ret); auto ret_ind = info.add_instruction(make_op("get_tuple_elem", {{"index", 1}}), topk_ret); - if(var_k) - { - // dynamic slice on outputs of `topk` - ret_val = info.add_instruction( - make_op("slice", {{"starts", {0}}, {"axes", {axis}}}), ret_val, args.at(1)); - ret_ind = info.add_instruction( - make_op("slice", {{"starts", {0}}, {"axes", {axis}}}), ret_ind, args.at(1)); - } + auto k_var = shape::dynamic_dimension{sym::var(info.name)}; + ret_val = info.add_instruction(make_op("slice", + {{"axes", {axis}}, + {"starts", {0}}, + {"ends", value::array{to_value(k_var)}}, + {"mode", "ends_input"}}), + ret_val, + args.at(1)); + ret_ind = info.add_instruction(make_op("slice", + {{"axes", {axis}}, + {"starts", {0}}, + {"ends", value::array{to_value(k_var)}}, + {"mode", "ends_input"}}), + ret_ind, + args.at(1)); return {ret_val, ret_ind}; } diff --git a/src/rewrite_topk.cpp b/src/rewrite_topk.cpp index 8597d55c331..dfca3db0b6d 100644 --- a/src/rewrite_topk.cpp +++ b/src/rewrite_topk.cpp @@ -43,6 +43,7 @@ struct find_large_topk { auto ins = r.result; auto input = ins->inputs().front(); + auto k_ins = ins->inputs().at(1); auto op = ins->get_operator().to_value(); auto axis = op["axis"].to(); auto dims = input->get_shape().lens(); @@ -75,10 +76,10 @@ struct find_large_topk ins, make_op("broadcast", {{"axis", axis}, {"out_lens", dims}}), indices_lit); auto gindices = m.insert_instruction(ins, make_op("reshape", {{"dims", gdims}}), indices); auto ginput = m.insert_instruction(ins, make_op("reshape", {{"dims", gdims}}), input); - auto topk1 = m.insert_instruction(ins, make_op("topk", op), ginput, gindices); + auto topk1 = m.insert_instruction(ins, make_op("topk", op), ginput, k_ins, gindices); auto finput = insert_final(topk1, 0); auto findices = insert_final(topk1, 1); - m.replace_instruction(ins, ins->get_operator(), finput, findices); + m.replace_instruction(ins, ins->get_operator(), finput, k_ins, findices); } }; diff --git a/src/simplify_algebra.cpp b/src/simplify_algebra.cpp index b054d007edf..6f0696a8fe8 100644 --- a/src/simplify_algebra.cpp +++ b/src/simplify_algebra.cpp @@ -205,8 +205,10 @@ struct find_mul_slice_conv auto sop = any_cast(i->get_operator()); if(sop.axes != slice_op.axes) return true; - if(std::max(sop.starts.front(), slice_op.starts.front()) < - std::min(sop.ends.front(), slice_op.ends.front())) + if(std::max(std::get(sop.starts.front()), + std::get(slice_op.starts.front())) < + std::min(std::get(sop.ends.front()), + std::get(slice_op.ends.front()))) return true; return false; })) @@ -223,17 +225,19 @@ struct find_mul_slice_conv auto new_mul = m.insert_instruction(ins, make_op("mul"), new_a, slice_w_ins); std::vector sliced_weights; - if(slice_op.starts.front() != 0) + if(std::get(slice_op.starts.front()) != 0) sliced_weights.push_back(m.insert_instruction( ins, - make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", slice_op.starts}}), + make_op("slice", + {{"axes", {0}}, {"starts", {0}}, {"ends", to_ints(slice_op.starts)}}), w_ins)); sliced_weights.push_back(new_mul); int64_t end_axis = w_ins->get_shape().lens().at(0); - if(slice_op.ends.front() != end_axis) + if(std::get(slice_op.ends.front()) != end_axis) sliced_weights.push_back(m.insert_instruction( ins, - make_op("slice", {{"axes", {0}}, {"starts", slice_op.ends}, {"ends", {end_axis}}}), + make_op("slice", + {{"axes", {0}}, {"starts", to_ints(slice_op.ends)}, {"ends", {end_axis}}}), w_ins)); auto new_weights = @@ -1388,11 +1392,12 @@ static std::vector get_splits(instruction_ref ins) auto get_end = [&](auto& i) -> auto& { return get_slice(i).ends; }; // Sort the "slice" instructions in order of starts - std::sort( - result.begin(), result.end(), [&](auto x, auto y) { return get_start(x) < get_start(y); }); - if(std::any_of(get_start(result.front()).begin(), get_start(result.front()).end(), [&](auto i) { - return i != 0; - })) + std::sort(result.begin(), result.end(), [&](auto x, auto y) { + return to_ints(get_start(x)) < to_ints(get_start(y)); + }); + if(std::any_of(get_start(result.front()).begin(), + get_start(result.front()).end(), + [&](const auto& i) { return std::get(i) != 0; })) return {}; // one slice must "start" where the last slice "end" @@ -1403,7 +1408,7 @@ static std::vector get_splits(instruction_ref ins) for(std::size_t i = 0; i < axes.size(); i++) { auto axis = axes[i]; - if(ins->get_shape().lens()[axis] != get_slice(result.back()).ends[i]) + if(ins->get_shape().lens()[axis] != std::get(get_slice(result.back()).ends[i])) return {}; } return result; @@ -1778,7 +1783,11 @@ struct find_split_concat if(not std::is_sorted(it, it + splits.size(), [](instruction_ref x, instruction_ref y) { auto xop = any_cast(x->get_operator()); auto yop = any_cast(y->get_operator()); - return std::tie(xop.starts, xop.ends) < std::tie(yop.starts, yop.ends); + auto xs = to_ints(xop.starts); + auto xe = to_ints(xop.ends); + auto ys = to_ints(yop.starts); + auto ye = to_ints(yop.ends); + return std::tie(xs, xe) < std::tie(ys, ye); })) return; @@ -2496,8 +2505,8 @@ struct find_split_transpose for(auto in : split_outputs) { auto oper = any_cast(in->get_operator()); - auto starts = oper.starts; - auto ends = oper.ends; + auto starts = to_ints(oper.starts); + auto ends = to_ints(oper.ends); auto tr_orig = in->outputs().front(); m.replace_instruction( tr_orig, diff --git a/src/simplify_dyn_ops.cpp b/src/simplify_dyn_ops.cpp index 6a90d97fe70..ea196740370 100644 --- a/src/simplify_dyn_ops.cpp +++ b/src/simplify_dyn_ops.cpp @@ -154,6 +154,18 @@ struct find_static_2in_broadcasts : match::supports_dynamic_shapes } }; +// Matches a slice whose set starts/ends bounds are all concrete ints (no +// symbolic dim_like bounds), so the const-input rewrites can read them as ints. +MIGRAPHX_PRED_MATCHER(slice_concrete_bounds, instruction_ref ins) +{ + if(ins->name() != "slice") + return false; + auto slice_op = any_cast(ins->get_operator()); + auto is_int = [](const dim_like& d) { return std::holds_alternative(d); }; + return std::all_of(slice_op.starts.begin(), slice_op.starts.end(), is_int) and + std::all_of(slice_op.ends.begin(), slice_op.ends.end(), is_int); +} + /** * Simplify slice with 2 inputs to the 1 input version if inputs[1] is constant. * From: @@ -165,32 +177,32 @@ struct find_const_2in_slice : match::supports_dynamic_shapes { auto matcher() const { - return match::name("slice")(match::nargs(2), match::arg(1)(match::is_constant())); + return match::name("slice")( + match::nargs(2), match::arg(1)(match::is_constant()), slice_concrete_bounds()); } void apply(module& m, const match::matcher_result& mr) const { - auto ins = mr.result; - auto inputs = ins->inputs(); - auto slice_op = any_cast(ins->get_operator()); - auto set_attrs = slice_op.get_set_attributes(); + auto ins = mr.result; + auto inputs = ins->inputs(); + auto slice_op = any_cast(ins->get_operator()); std::vector starts_vec; std::vector ends_vec; std::vector axes_vec; - if(set_attrs == op::slice::ends_axes) + if(slice_op.mode == op::slice::slice_mode::starts_input) { // slice(data, starts) inputs.at(1)->eval().visit( [&](auto output) { starts_vec.assign(output.begin(), output.end()); }); - ends_vec = slice_op.ends; + ends_vec = to_ints(slice_op.ends); axes_vec = slice_op.axes; } - else if(set_attrs == op::slice::starts_axes) + else if(slice_op.mode == op::slice::slice_mode::ends_input) { // slice(data, ends) inputs.at(1)->eval().visit( [&](auto output) { ends_vec.assign(output.begin(), output.end()); }); - starts_vec = slice_op.starts; + starts_vec = to_ints(slice_op.starts); axes_vec = slice_op.axes; } else @@ -198,8 +210,8 @@ struct find_const_2in_slice : match::supports_dynamic_shapes // 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; + starts_vec = to_ints(slice_op.starts); + ends_vec = to_ints(slice_op.ends); } m.replace_instruction( ins, @@ -221,19 +233,19 @@ struct find_const_3in_slice : match::supports_dynamic_shapes { return match::name("slice")(match::nargs(3), match::arg(1)(match::is_constant()), - match::arg(2)(match::is_constant())); + match::arg(2)(match::is_constant()), + slice_concrete_bounds()); } void apply(module& m, const match::matcher_result& mr) const { - auto ins = mr.result; - auto inputs = ins->inputs(); - auto slice_op = any_cast(ins->get_operator()); - auto set_attrs = slice_op.get_set_attributes(); + auto ins = mr.result; + auto inputs = ins->inputs(); + auto slice_op = any_cast(ins->get_operator()); std::vector starts_vec; std::vector ends_vec; std::vector axes_vec; - if(set_attrs == op::slice::axes_only) + if(slice_op.mode == op::slice::slice_mode::starts_ends_input) { // slice(data, starts, ends) inputs.at(1)->eval().visit( @@ -242,14 +254,14 @@ struct find_const_3in_slice : match::supports_dynamic_shapes [&](auto output) { ends_vec.assign(output.begin(), output.end()); }); axes_vec = slice_op.axes; } - else if(set_attrs == op::slice::ends_only) + else if(slice_op.mode == op::slice::slice_mode::starts_axes_input) { // 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; + ends_vec = to_ints(slice_op.ends); } else { @@ -258,7 +270,7 @@ struct find_const_3in_slice : match::supports_dynamic_shapes [&](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; + starts_vec = to_ints(slice_op.starts); } m.replace_instruction( ins, diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp index f31bce1d3b3..310dd6dedd8 100644 --- a/src/simplify_reshapes.cpp +++ b/src/simplify_reshapes.cpp @@ -663,7 +663,8 @@ struct find_nested_slice auto op = any_cast(ins->get_operator()); for(std::size_t i = 0; i < op.axes.size(); i++) { - result[op.axes[i]] = std::make_pair(op.starts[i], op.ends[i]); + result[op.axes[i]] = + std::make_pair(std::get(op.starts[i]), std::get(op.ends[i])); } return result; } @@ -861,8 +862,8 @@ struct find_concat_slice for(const auto& sins : slice_candidates) { auto sop = any_cast(sins->get_operator()); - size_t slice_start = sop.starts.front(); - size_t slice_len = sop.ends.front() - slice_start; + size_t slice_start = std::get(sop.starts.front()); + size_t slice_len = std::get(sop.ends.front()) - slice_start; auto fii = std::find_if(prefix_scan.begin(), prefix_scan.end(), [&](const auto& j) { return j == slice_start; }); @@ -1291,7 +1292,7 @@ struct find_gather return; const std::size_t axis_index = tune_axis(dlens.size(), gather_op.axis, gather_op.name()); - const auto axis_len = dlens.at(axis_index); + const auto axis_len = dlens.at(axis_index); if(axis_len == 0) return; @@ -1820,8 +1821,13 @@ struct find_transpose_slice { assert(op.starts.size() == op.ends.size()); std::vector result(op.starts.size()); - std::transform( - op.ends.begin(), op.ends.end(), op.starts.begin(), result.begin(), std::minus<>{}); + std::transform(op.ends.begin(), + op.ends.end(), + op.starts.begin(), + result.begin(), + [](const auto& e, const auto& s) { + return std::get(e) - std::get(s); + }); return result; } @@ -1846,14 +1852,15 @@ struct find_transpose_slice sdistance.begin(), 0, std::plus<>{}, - [&](auto x, auto d) -> uint64_t { + [&](const auto& x, auto d) -> uint64_t { if(d == 0) return 1; return f(x) % d; }); }; + auto get_int = [](const auto& x) { return std::get(x); }; if(mod_by_distance(slice.axes, [&](auto x) { return ins->get_shape().lens()[x]; }) != 0 or - mod_by_distance(slice.starts, id{}) != 0 or mod_by_distance(slice.ends, id{}) != 0) + mod_by_distance(slice.starts, get_int) != 0 or mod_by_distance(slice.ends, get_int) != 0) return; // TODO: Handle multiple axes if(sdistance.size() != 1) @@ -1891,8 +1898,8 @@ struct find_transpose_slice { auto op = any_cast(s->get_operator()); op.axes = {0}; - op.starts = {op.starts.front() / sdistance.front()}; - op.ends = {op.ends.front() / sdistance.front()}; + op.starts = {std::get(op.starts.front()) / sdistance.front()}; + op.ends = {std::get(op.ends.front()) / sdistance.front()}; auto slice_ins = m.insert_instruction(ins, op, transpose); auto squeeze = m.insert_instruction(ins, make_op("squeeze", {{"axes", {0}}}), slice_ins); diff --git a/src/targets/gpu/kernels/include/migraphx/kernels/topk.hpp b/src/targets/gpu/kernels/include/migraphx/kernels/topk.hpp index d771fab6b65..fac609e26b3 100644 --- a/src/targets/gpu/kernels/include/migraphx/kernels/topk.hpp +++ b/src/targets/gpu/kernels/include/migraphx/kernels/topk.hpp @@ -209,7 +209,10 @@ topk_impl(index idx, Compare compare, T init, Y y, YIndex y_idx, X x, XIndices.. template __device__ auto topk(Compare compare, T init) { - return [=](auto output, auto out_indices, auto input, auto... in_indices) { + // The `k` input is consumed positionally and ignored here: compile_ops bakes a constant + // `k` attribute into the output shape, so the trailing `in_indices` (if any) is the + // optional indexing input used by rewrite_topk. + return [=](auto output, auto out_indices, auto input, auto, auto... in_indices) { auto idx = make_index(); slice_schedule(idx, slice_axes())(output, out_indices, input, in_indices...)( diff --git a/src/targets/gpu/topk.cpp b/src/targets/gpu/topk.cpp index 2e799c650af..05466a8aca1 100644 --- a/src/targets/gpu/topk.cpp +++ b/src/targets/gpu/topk.cpp @@ -31,23 +31,25 @@ namespace gpu { shape hip_topk::compute_shape(std::vector inputs) const { - return op.normalize_compute_shape({inputs.front()}); + // Drop the trailing output-allocation input; topk needs the (x, k, [indexing]) inputs. + return op.normalize_compute_shape({inputs.begin(), inputs.end() - 1}); } argument hip_topk::compute(context& ctx, const shape&, const std::vector& args) const { auto outputs = args.back().get_sub_objects(); + auto k_val = std::get(op.k); return op.largest ? device::topk_largest(ctx.get_stream().get(), outputs.front(), outputs.back(), args[0], - op.k, + k_val, op.axis) : device::topk_smallest(ctx.get_stream().get(), outputs.front(), outputs.back(), args[0], - op.k, + k_val, op.axis); } diff --git a/test/enum.cpp b/test/enum.cpp index 14c9f313c3d..74a455442de 100644 --- a/test/enum.cpp +++ b/test/enum.cpp @@ -148,6 +148,14 @@ namespace migraphx { // check that the generated to_string(status) is not ambiguous with migraphx::to_string(const T&). // NOLINTNEXTLINE(misc-use-internal-linkage) MIGRAPHX_ENUM(status, ok, busy = 4, done) + +// Bit-flag enums. Declared in the migraphx namespace so that argument-dependent lookup finds the +// operators (which also live in migraphx) from the global-scope test cases below. Two widths +// exercise the selectable underlying type. +// NOLINTNEXTLINE(misc-use-internal-linkage) +MIGRAPHX_BIT_FLAG_ENUM(access, std::uint8_t, none = 0, read = 1 << 0, write = 1 << 1, exec = 1 << 2) +// NOLINTNEXTLINE(misc-use-internal-linkage) +MIGRAPHX_BIT_FLAG_ENUM(wide_flag, std::uint16_t, none = 0, lo = 1 << 0, hi = 1 << 15) } // namespace migraphx TEST_CASE(underlying_values) @@ -307,6 +315,56 @@ TEST_CASE(is_named_enum_trait) EXPECT(not migraphx::is_named_enum{}); } +TEST_CASE(bit_flag_operators) +{ + auto rw = migraphx::access::read | migraphx::access::write; + EXPECT(has_flag(rw, migraphx::access::read)); + EXPECT(has_flag(rw, migraphx::access::write)); + EXPECT(not has_flag(rw, migraphx::access::exec)); + + // & masks, and the result stays the enum type. + EXPECT((rw & migraphx::access::read) == migraphx::access::read); + EXPECT((rw & migraphx::access::exec) == migraphx::access::none); + + rw |= migraphx::access::exec; + EXPECT(has_flag(rw, migraphx::access::exec)); + + // ~ and &= clear a bit. + rw &= ~migraphx::access::write; + EXPECT(not has_flag(rw, migraphx::access::write)); + EXPECT(has_flag(rw, migraphx::access::read)); + + // ^ and ^= toggle a bit. + EXPECT((migraphx::access::read ^ migraphx::access::read) == migraphx::access::none); + auto toggled = migraphx::access::none; + toggled ^= migraphx::access::write; + EXPECT(toggled == migraphx::access::write); +} + +TEST_CASE(bit_flag_width_and_scoped) +{ + // The underlying type is honored, so the storage width is selectable. + EXPECT(sizeof(migraphx::access) == 1); + EXPECT(sizeof(migraphx::wide_flag) == 2); + // A bit that only fits in the wider type round-trips through the underlying value. + EXPECT(static_cast(migraphx::wide_flag::hi) == (1 << 15)); + // It is still a real scoped enum: no implicit conversion to its underlying type. + EXPECT(not std::is_convertible{}); +} + +TEST_CASE(is_bit_flag_trait) +{ + EXPECT(migraphx::is_bit_flag{}); + EXPECT(migraphx::is_bit_flag{}); + // False for named enums, plain enums, and non-enum types: the operators do not leak to them. + EXPECT(not migraphx::is_bit_flag{}); + EXPECT(not migraphx::is_bit_flag{}); + EXPECT(not migraphx::is_bit_flag{}); + EXPECT(not migraphx::is_bit_flag{}); + // And a bit-flag enum does not get the named-enum string helpers. + EXPECT(not migraphx::is_named_enum{}); +} + TEST_CASE(value_stores_named_enum_as_string) { migraphx::value v = green; diff --git a/test/gpu/dyn_slice_lowering.cpp b/test/gpu/dyn_slice_lowering.cpp index 11d743012af..32c55950635 100644 --- a/test/gpu/dyn_slice_lowering.cpp +++ b/test/gpu/dyn_slice_lowering.cpp @@ -49,8 +49,11 @@ TEST_CASE(dyn_slice_lowering_runtime_inputs) auto data = m1.add_parameter("data", data_s); auto starts = m1.add_parameter("starts", idx_s); auto ends = m1.add_parameter("ends", idx_s); - auto sl = - m1.add_instruction(migraphx::make_op("slice", {{"axes", {2}}}), data, starts, ends); + auto sl = m1.add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"mode", "starts_ends_input"}}), + data, + starts, + ends); m1.add_return({sl}); } run_lowering(m1); @@ -64,8 +67,11 @@ TEST_CASE(dyn_slice_lowering_runtime_inputs) auto copy_ends = m2.add_instruction(migraphx::make_op("hip::copy_from_gpu"), ends); auto sync = m2.add_instruction(migraphx::make_op("hip::sync_stream"), copy_starts, copy_ends); - auto sl = - m2.add_instruction(migraphx::make_op("slice", {{"axes", {2}}}), data, sync, copy_ends); + auto sl = m2.add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"mode", "starts_ends_input"}}), + data, + sync, + copy_ends); m2.add_return({sl}); } EXPECT(m1 == m2); diff --git a/test/multi_target/multitarget_test.cpp b/test/multi_target/multitarget_test.cpp index 04ba53901d3..d944b37b582 100644 --- a/test/multi_target/multitarget_test.cpp +++ b/test/multi_target/multitarget_test.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -230,8 +231,16 @@ TEST_CASE(single_target_multi_compile) score_threshold); auto idx = gpu_mod->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); auto cnt = gpu_mod->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), nms); - auto r = gpu_mod->add_instruction( - migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}}), idx, cnt); + auto num_selected_var = + migraphx::shape::dynamic_dimension{migraphx::sym::var("nms_num_selected")}; + auto r = gpu_mod->add_instruction( + migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + idx, + cnt); gpu_mod->add_return({r}); auto run_on_gpu = mm->add_instruction( diff --git a/test/onnx/parse/nms_dynamic_batch_test.cpp b/test/onnx/parse/nms_dynamic_batch_test.cpp index f9ac10fa4aa..98f9964820c 100644 --- a/test/onnx/parse/nms_dynamic_batch_test.cpp +++ b/test/onnx/parse/nms_dynamic_batch_test.cpp @@ -40,7 +40,18 @@ TEST_CASE(nms_dynamic_batch_test) auto st = mm->add_parameter("score_threshold", sst); auto nms = mm->add_instruction( migraphx::make_op("nonmaxsuppression", {{"center_point_box", true}}), b, s, mo, iou, st); - auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto idx = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto cnt = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), nms); + auto num_selected_var = + migraphx::shape::dynamic_dimension{migraphx::sym::var("NonMaxSuppression_5")}; + auto ret = mm->add_instruction( + migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + idx, + cnt); mm->add_return({ret}); migraphx::onnx_options options; diff --git a/test/onnx/parse/nms_dynamic_boxes_test.cpp b/test/onnx/parse/nms_dynamic_boxes_test.cpp index 2b11265d00c..17d5de4d674 100644 --- a/test/onnx/parse/nms_dynamic_boxes_test.cpp +++ b/test/onnx/parse/nms_dynamic_boxes_test.cpp @@ -39,7 +39,18 @@ TEST_CASE(nms_dynamic_boxes_test) migraphx::shape sst{migraphx::shape::float_type, {1}}; auto st = mm->add_parameter("score_threshold", sst); auto nms = mm->add_instruction(migraphx::make_op("nonmaxsuppression"), b, s, mo, iou, st); - auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto idx = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto cnt = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), nms); + auto num_selected_var = + migraphx::shape::dynamic_dimension{migraphx::sym::var("NonMaxSuppression_5")}; + auto ret = mm->add_instruction( + migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + idx, + cnt); mm->add_return({ret}); migraphx::onnx_options options; diff --git a/test/onnx/parse/nms_dynamic_classes_test.cpp b/test/onnx/parse/nms_dynamic_classes_test.cpp index 8f8a3abd9d2..adaf189baec 100644 --- a/test/onnx/parse/nms_dynamic_classes_test.cpp +++ b/test/onnx/parse/nms_dynamic_classes_test.cpp @@ -39,7 +39,18 @@ TEST_CASE(nms_dynamic_classes_test) migraphx::shape sst{migraphx::shape::float_type, {1}}; auto st = mm->add_parameter("score_threshold", sst); auto nms = mm->add_instruction(migraphx::make_op("nonmaxsuppression"), b, s, mo, iou, st); - auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto idx = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto cnt = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), nms); + auto num_selected_var = + migraphx::shape::dynamic_dimension{migraphx::sym::var("NonMaxSuppression_5")}; + auto ret = mm->add_instruction( + migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + idx, + cnt); mm->add_return({ret}); migraphx::onnx_options options; diff --git a/test/onnx/parse/nms_test.cpp b/test/onnx/parse/nms_test.cpp index 3dbf522b504..beefd297652 100644 --- a/test/onnx/parse/nms_test.cpp +++ b/test/onnx/parse/nms_test.cpp @@ -45,7 +45,18 @@ TEST_CASE(nms_test) auto nms = mm->add_instruction( migraphx::make_op("nonmaxsuppression", {{"center_point_box", true}}), b, s, mo, iou, st); - auto ret = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto idx = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); + auto cnt = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), nms); + auto num_selected_var = + migraphx::shape::dynamic_dimension{migraphx::sym::var("NonMaxSuppression_5")}; + auto ret = mm->add_instruction( + migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + idx, + cnt); mm->add_return({ret}); auto prog = read_onnx("nms_test.onnx"); diff --git a/test/onnx/parse/slice_var_input_default_steps.cpp b/test/onnx/parse/slice_var_input_default_steps.cpp index e432b124f27..d3caa52e22e 100644 --- a/test/onnx/parse/slice_var_input_default_steps.cpp +++ b/test/onnx/parse/slice_var_input_default_steps.cpp @@ -34,7 +34,8 @@ TEST_CASE(slice_var_input_default_steps) auto ends = mm->add_parameter("ends", migraphx::shape{migraphx::shape::int64_type, {2}}); auto axes = mm->add_parameter("axes", migraphx::shape{migraphx::shape::int64_type, {2}}); mm->add_literal({{migraphx::shape::int64_type, {2}}, {1, 1}}); - auto ret = mm->add_instruction(migraphx::make_op("slice"), data, starts, ends, axes); + auto ret = mm->add_instruction( + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), data, starts, ends, axes); mm->add_return({ret}); migraphx::onnx_options options; diff --git a/test/onnx/parse/slice_var_input_dyn0.cpp b/test/onnx/parse/slice_var_input_dyn0.cpp index ae4ab5739f8..11f64ef960a 100644 --- a/test/onnx/parse/slice_var_input_dyn0.cpp +++ b/test/onnx/parse/slice_var_input_dyn0.cpp @@ -32,8 +32,11 @@ TEST_CASE(slice_var_input_dyn0) mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {{3, 8}, {2, 2}}}); auto starts = mm->add_parameter("starts", migraphx::shape{migraphx::shape::int32_type, {2}}); auto ends = mm->add_parameter("ends", migraphx::shape{migraphx::shape::int32_type, {2}}); - auto ret = - mm->add_instruction(migraphx::make_op("slice", {{"axes", {0, 1}}}), data, starts, ends); + auto ret = mm->add_instruction( + migraphx::make_op("slice", {{"axes", {0, 1}}, {"mode", "starts_ends_input"}}), + data, + starts, + ends); mm->add_return({ret}); migraphx::onnx_options options; diff --git a/test/onnx/parse/slice_var_input_dyn1.cpp b/test/onnx/parse/slice_var_input_dyn1.cpp index 1fe4681cb57..c660924a622 100644 --- a/test/onnx/parse/slice_var_input_dyn1.cpp +++ b/test/onnx/parse/slice_var_input_dyn1.cpp @@ -33,7 +33,8 @@ TEST_CASE(slice_var_input_dyn1) auto starts = mm->add_parameter("starts", migraphx::shape{migraphx::shape::int32_type, {2}}); auto ends = mm->add_parameter("ends", migraphx::shape{migraphx::shape::int32_type, {2}}); auto axes = mm->add_parameter("axes", migraphx::shape{migraphx::shape::int32_type, {2}}); - auto ret = mm->add_instruction(migraphx::make_op("slice"), data, starts, ends, axes); + auto ret = mm->add_instruction( + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), data, starts, ends, axes); mm->add_return({ret}); migraphx::onnx_options options; diff --git a/test/onnx/parse/slice_var_input_static0.cpp b/test/onnx/parse/slice_var_input_static0.cpp index 2bea481379f..650a0c5783c 100644 --- a/test/onnx/parse/slice_var_input_static0.cpp +++ b/test/onnx/parse/slice_var_input_static0.cpp @@ -31,7 +31,11 @@ TEST_CASE(slice_var_input_static0) auto data = mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {3, 2}}); auto starts = mm->add_parameter("starts", migraphx::shape{migraphx::shape::int32_type, {2}}); auto ends = mm->add_parameter("ends", migraphx::shape{migraphx::shape::int32_type, {2}}); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {0, 1}}}), data, starts, ends); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {0, 1}}, {"mode", "starts_ends_input"}}), + data, + starts, + ends); auto prog = optimize_onnx("slice_var_input_static0.onnx"); EXPECT(p == prog); diff --git a/test/onnx/parse/slice_var_input_static1.cpp b/test/onnx/parse/slice_var_input_static1.cpp index d98a1429499..f9024cecaf6 100644 --- a/test/onnx/parse/slice_var_input_static1.cpp +++ b/test/onnx/parse/slice_var_input_static1.cpp @@ -32,7 +32,8 @@ TEST_CASE(slice_var_input_static1) auto starts = mm->add_parameter("starts", migraphx::shape{migraphx::shape::int64_type, {2}}); auto ends = mm->add_parameter("ends", migraphx::shape{migraphx::shape::int64_type, {2}}); auto axes = mm->add_parameter("axes", migraphx::shape{migraphx::shape::int64_type, {2}}); - mm->add_instruction(migraphx::make_op("slice"), data, starts, ends, axes); + mm->add_instruction( + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), data, starts, ends, axes); auto prog = optimize_onnx("slice_var_input_static1.onnx"); EXPECT(p == prog); diff --git a/test/onnx/parse/topk_attrk_test.cpp b/test/onnx/parse/topk_attrk_test.cpp index 6f8f6f3b9fc..d8a3fd62cd9 100644 --- a/test/onnx/parse/topk_attrk_test.cpp +++ b/test/onnx/parse/topk_attrk_test.cpp @@ -30,7 +30,9 @@ TEST_CASE(topk_attrk_test) auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {2, 5, 3, 2}}; auto data = mm->add_parameter("data", s); - auto out = mm->add_instruction(migraphx::make_op("topk", {{"k", 2}, {"axis", -1}}), data); + auto k = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {2}}); + auto out = mm->add_instruction(migraphx::make_op("topk", {{"k", 2}, {"axis", -1}}), data, k); auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); mm->add_return({val, ind}); diff --git a/test/onnx/parse/topk_neg_axis_test.cpp b/test/onnx/parse/topk_neg_axis_test.cpp index 882bf43ecca..99ef3ab93ab 100644 --- a/test/onnx/parse/topk_neg_axis_test.cpp +++ b/test/onnx/parse/topk_neg_axis_test.cpp @@ -29,11 +29,11 @@ TEST_CASE(topk_neg_axis_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape sk{migraphx::shape::int64_type, {1}}; - mm->add_literal(migraphx::literal(sk, {3})); + auto k = mm->add_literal(migraphx::literal(sk, {3})); migraphx::shape s{migraphx::shape::float_type, {3, 4, 5, 6}}; auto data = mm->add_parameter("data", s); auto out = mm->add_instruction( - migraphx::make_op("topk", {{"k", 3}, {"axis", -2}, {"largest", 1}}), data); + migraphx::make_op("topk", {{"k", 3}, {"axis", -2}, {"largest", 1}}), data, k); auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); mm->add_return({val, ind}); diff --git a/test/onnx/parse/topk_test.cpp b/test/onnx/parse/topk_test.cpp index 4ac45c73df2..9af389a8d78 100644 --- a/test/onnx/parse/topk_test.cpp +++ b/test/onnx/parse/topk_test.cpp @@ -29,11 +29,11 @@ TEST_CASE(topk_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape sk{migraphx::shape::int64_type, {1}}; - mm->add_literal(migraphx::literal(sk, {4})); + auto k = mm->add_literal(migraphx::literal(sk, {4})); migraphx::shape s{migraphx::shape::float_type, {2, 5, 3, 2}}; auto data = mm->add_parameter("data", s); auto out = mm->add_instruction( - migraphx::make_op("topk", {{"k", 4}, {"axis", 1}, {"largest", 0}}), data); + migraphx::make_op("topk", {{"k", 4}, {"axis", 1}, {"largest", 0}}), data, k); auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); mm->add_return({val, ind}); diff --git a/test/onnx/parse/topk_var_k_test.cpp b/test/onnx/parse/topk_var_k_test.cpp index 2338236abe0..a7c7f7f9cc1 100644 --- a/test/onnx/parse/topk_var_k_test.cpp +++ b/test/onnx/parse/topk_var_k_test.cpp @@ -25,8 +25,8 @@ #include // `k` is a runtime input (graph input, not an initializer), so the parser takes the var_k -// path: topk runs with k set to the axis dimension, then the outputs are sliced down to the -// runtime `k`. +// path: topk runs with k set to the axis dimension's max length, then the outputs are sliced +// down to the runtime `k` using a symbolic dimension. TEST_CASE(topk_var_k_test) { migraphx::program p; @@ -34,11 +34,24 @@ TEST_CASE(topk_var_k_test) auto data = mm->add_parameter("data", {migraphx::shape::float_type, {2, 4}}); auto k = mm->add_parameter("k", {migraphx::shape::int64_type, {1}}); auto out = mm->add_instruction( - migraphx::make_op("topk", {{"k", 4}, {"axis", 1}, {"largest", 1}}), data); - auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); - auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); - val = mm->add_instruction(migraphx::make_op("slice", {{"starts", {0}}, {"axes", {1}}}), val, k); - ind = mm->add_instruction(migraphx::make_op("slice", {{"starts", {0}}, {"axes", {1}}}), ind, k); + migraphx::make_op("topk", {{"k", 4}, {"axis", 1}, {"largest", 1}}), data, k); + auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); + auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); + auto k_var = migraphx::shape::dynamic_dimension{migraphx::sym::var("TopK_2")}; + val = mm->add_instruction(migraphx::make_op("slice", + {{"axes", {1}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(k_var)}}, + {"mode", "ends_input"}}), + val, + k); + ind = mm->add_instruction(migraphx::make_op("slice", + {{"axes", {1}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(k_var)}}, + {"mode", "ends_input"}}), + ind, + k); mm->add_return({val, ind}); auto prog = read_onnx("topk_var_k_test.onnx"); @@ -55,11 +68,24 @@ TEST_CASE(topk_var_k_dynamic_test) auto data = mm->add_parameter("data", {migraphx::shape::float_type, {{1, 4}, {2, 4}}}); auto k = mm->add_parameter("k", {migraphx::shape::int64_type, {1}}); auto out = mm->add_instruction( - migraphx::make_op("topk", {{"k", 4}, {"axis", 1}, {"largest", 1}}), data); - auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); - auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); - val = mm->add_instruction(migraphx::make_op("slice", {{"starts", {0}}, {"axes", {1}}}), val, k); - ind = mm->add_instruction(migraphx::make_op("slice", {{"starts", {0}}, {"axes", {1}}}), ind, k); + migraphx::make_op("topk", {{"k", 4}, {"axis", 1}, {"largest", 1}}), data, k); + auto val = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); + auto ind = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); + auto k_var = migraphx::shape::dynamic_dimension{migraphx::sym::var("TopK_2")}; + val = mm->add_instruction(migraphx::make_op("slice", + {{"axes", {1}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(k_var)}}, + {"mode", "ends_input"}}), + val, + k); + ind = mm->add_instruction(migraphx::make_op("slice", + {{"axes", {1}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(k_var)}}, + {"mode", "ends_input"}}), + ind, + k); mm->add_return({val, ind}); migraphx::onnx_options options; diff --git a/test/op/builder/torch_kit_test.cpp b/test/op/builder/torch_kit_test.cpp index ea207e86f0b..fec515c9ab9 100644 --- a/test/op/builder/torch_kit_test.cpp +++ b/test/op/builder/torch_kit_test.cpp @@ -278,7 +278,8 @@ TEST_CASE(torch_kit_ops_op_builder_test) EXPECT(check_plain_op("softmax", {{"axis", 1}}, {{"a", {f, {4, 6}}}})); EXPECT(check_plain_op("squeeze", {{"axes", {0}}}, {{"a", {f, {1, 4, 6}}}})); EXPECT(check_plain_op("step", {{"axes", {0}}, {"steps", {2}}}, {{"a", {f, {4, 6}}}})); - EXPECT(check_plain_op("topk", {{"k", 2}, {"axis", 0}}, {{"a", {f, {4, 6}}}})); + EXPECT( + check_plain_op("topk", {{"k", 2}, {"axis", 0}}, {{"a", {f, {4, 6}}}, {"k", {i64, {1}}}})); EXPECT(check_plain_op("transpose", {{"permutation", {1, 0}}}, {{"a", {f, {4, 6}}}})); EXPECT(check_plain_op("undefined", obj, {})); EXPECT(check_plain_op("unsqueeze", {{"axes", {0}}}, {{"a", {f, {4, 6}}}})); diff --git a/test/op_shape_test.cpp b/test/op_shape_test.cpp index 12369334843..387fceddeaf 100644 --- a/test/op_shape_test.cpp +++ b/test/op_shape_test.cpp @@ -4906,18 +4906,11 @@ TEST_CASE(slice_var_inputs_static_shape0) // attr ends and axes set; inputs are (data, input_starts) migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{3, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice", {{"ends", {2, 3}}, {"axes", {1, 2}}}), - input, - starts); -} - -TEST_CASE(slice_var_inputs_static_mismatch_error0) -{ - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; - migraphx::shape starts{migraphx::shape::int64_type, {2}}; - throws_shape( - migraphx::make_op("slice", {{"ends", {2, 3, 4}}, {"axes", {0, 1, 2}}}), input, starts); + expect_shape( + migraphx::shape{migraphx::shape::float_type, {{3, 3}, {0, 4}, {0, 4}}}, + migraphx::make_op("slice", {{"ends", {2, 3}}, {"axes", {1, 2}}, {"mode", "starts_input"}}), + input, + starts); } TEST_CASE(slice_var_inputs_static_shape1) @@ -4925,18 +4918,11 @@ TEST_CASE(slice_var_inputs_static_shape1) // attr starts and axes set; inputs are (data, input_ends) migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{3, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice", {{"starts", {0, 1}}, {"axes", {1, 2}}}), - input, - ends); -} - -TEST_CASE(slice_var_inputs_static_mismatch_error1) -{ - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; - migraphx::shape ends{migraphx::shape::int64_type, {2}}; - throws_shape( - migraphx::make_op("slice", {{"starts", {0, 1, 2}}, {"axes", {0, 1, 2}}}), input, ends); + expect_shape( + migraphx::shape{migraphx::shape::float_type, {{3, 3}, {0, 4}, {0, 4}}}, + migraphx::make_op("slice", {{"starts", {0, 1}}, {"axes", {1, 2}}, {"mode", "ends_input"}}), + input, + ends); } TEST_CASE(slice_var_inputs_static_shape2) @@ -4944,18 +4930,11 @@ TEST_CASE(slice_var_inputs_static_shape2) // attr starts and ends set; inputs are (data, input_axes) migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice", {{"starts", {0, 1}}, {"ends", {1, 2}}}), - input, - axes); -} - -TEST_CASE(slice_var_inputs_static_mismatch_error2) -{ - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; - migraphx::shape axes{migraphx::shape::int64_type, {2}}; - throws_shape( - migraphx::make_op("slice", {{"starts", {0, 1, 2}}, {"ends", {3, 4, 4}}}), input, axes); + expect_shape( + migraphx::shape{migraphx::shape::float_type, {{0, 3}, {0, 4}, {0, 4}}}, + migraphx::make_op("slice", {{"starts", {0, 1}}, {"ends", {1, 2}}, {"mode", "axes_input"}}), + input, + axes); } TEST_CASE(slice_var_inputs_static_shape3) @@ -4965,20 +4944,12 @@ TEST_CASE(slice_var_inputs_static_shape3) migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; expect_shape(migraphx::shape{migraphx::shape::float_type, {{3, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice", {{"axes", {1, 2}}}), + migraphx::make_op("slice", {{"axes", {1, 2}}, {"mode", "starts_ends_input"}}), input, starts, ends); } -TEST_CASE(slice_var_inputs_static_mismatch_error3) -{ - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; - migraphx::shape starts{migraphx::shape::int64_type, {2}}; - migraphx::shape ends{migraphx::shape::int64_type, {2}}; - throws_shape(migraphx::make_op("slice", {{"axes", {0, 1, 2}}}), input, starts, ends); -} - TEST_CASE(slice_var_inputs_static_shape4) { // attr ends set; inputs are (data, input_starts, input_axes) @@ -4986,20 +4957,12 @@ TEST_CASE(slice_var_inputs_static_shape4) migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice", {{"ends", {3, 4}}}), + migraphx::make_op("slice", {{"ends", {3, 4}}, {"mode", "starts_axes_input"}}), input, starts, axes); } -TEST_CASE(slice_var_inputs_static_mismatch_error4) -{ - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; - migraphx::shape starts{migraphx::shape::int64_type, {2}}; - migraphx::shape axes{migraphx::shape::int64_type, {2}}; - throws_shape(migraphx::make_op("slice", {{"ends", {3, 3, 3}}}), input, starts, axes); -} - TEST_CASE(slice_var_inputs_static_shape5) { // attr starts set; inputs are (data, input_ends, input_axes) @@ -5007,20 +4970,12 @@ TEST_CASE(slice_var_inputs_static_shape5) migraphx::shape ends{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice", {{"starts", {0, 2}}}), + migraphx::make_op("slice", {{"starts", {0, 2}}, {"mode", "ends_axes_input"}}), input, ends, axes); } -TEST_CASE(slice_var_inputs_static_mismatch_error5) -{ - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; - migraphx::shape ends{migraphx::shape::int64_type, {2}}; - migraphx::shape axes{migraphx::shape::int64_type, {2}}; - throws_shape(migraphx::make_op("slice", {{"starts", {0, 1, 2}}}), input, ends, axes); -} - TEST_CASE(slice_var_inputs_static_shape6) { migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; @@ -5028,154 +4983,251 @@ TEST_CASE(slice_var_inputs_static_shape6) migraphx::shape ends{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 3}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice"), + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), input, starts, ends, axes); } -TEST_CASE(slice_var_inputs_static_mismatch_error6) +TEST_CASE(slice_var_inputs_dyn_shape0) { - migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; + // attr ends and axes set; inputs are (data, input_starts) + migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; - migraphx::shape ends{migraphx::shape::int64_type, {2}}; - migraphx::shape axes{migraphx::shape::int64_type, {3}}; - throws_shape(migraphx::make_op("slice"), input, starts, ends, axes); + expect_shape( + migraphx::shape{migraphx::shape::float_type, {{3, 6}, {0, 6}, {0, 6}}}, + migraphx::make_op("slice", {{"ends", {2, 3}}, {"axes", {1, 2}}, {"mode", "starts_input"}}), + input, + starts); } -TEST_CASE(slice_var_inputs_dyn_shape0) +TEST_CASE(slice_var_inputs_dyn_shape1) { - // attr ends and axes set; inputs are (data, input_starts) + // attr starts and axes set; inputs are (data, input_ends) migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; - migraphx::shape starts{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{3, 6}, {0, 6}, {0, 6}}}, - migraphx::make_op("slice", {{"ends", {2, 3}}, {"axes", {1, 2}}}), - input, - starts); + migraphx::shape ends{migraphx::shape::int64_type, {2}}; + expect_shape( + migraphx::shape{migraphx::shape::float_type, {{3, 6}, {0, 6}, {0, 6}}}, + migraphx::make_op("slice", {{"starts", {0, 1}}, {"axes", {1, 2}}, {"mode", "ends_input"}}), + input, + ends); } -TEST_CASE(slice_var_inputs_dyn_mismatch_error0) +TEST_CASE(slice_var_inputs_dyn_shape2) { + // attr starts and ends set; inputs are (data, input_axes) migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; - migraphx::shape starts{migraphx::shape::int64_type, {2}}; - throws_shape( - migraphx::make_op("slice", {{"ends", {2, 3, 4}}, {"axes", {0, 1, 2}}}), input, starts); + migraphx::shape axes{migraphx::shape::int64_type, {2}}; + expect_shape( + migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 6}, {0, 6}}}, + migraphx::make_op("slice", {{"starts", {0, 1}}, {"ends", {8, 8}}, {"mode", "axes_input"}}), + input, + axes); } -TEST_CASE(slice_var_inputs_dyn_shape1) +TEST_CASE(slice_var_inputs_dyn_shape3) { - // attr starts and axes set; inputs are (data, input_ends) + // attr axes set; inputs are (data, input_starts, input_ends) migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; expect_shape(migraphx::shape{migraphx::shape::float_type, {{3, 6}, {0, 6}, {0, 6}}}, - migraphx::make_op("slice", {{"starts", {0, 1}}, {"axes", {1, 2}}}), + migraphx::make_op("slice", {{"axes", {1, 2}}, {"mode", "starts_ends_input"}}), input, + starts, ends); } -TEST_CASE(slice_var_inputs_dyn_mismatch_error1) +TEST_CASE(slice_var_inputs_dyn_shape4) { + // attr ends set; inputs are (data, input_starts, input_axes) migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; - migraphx::shape ends{migraphx::shape::int64_type, {2}}; - throws_shape( - migraphx::make_op("slice", {{"starts", {0, 1, 2}}, {"axes", {0, 1, 2}}}), input, ends); + migraphx::shape starts{migraphx::shape::int64_type, {2}}; + migraphx::shape axes{migraphx::shape::int64_type, {2}}; + expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 6}, {0, 6}}}, + migraphx::make_op("slice", {{"ends", {3, 4}}, {"mode", "starts_axes_input"}}), + input, + starts, + axes); } -TEST_CASE(slice_var_inputs_dyn_shape2) +TEST_CASE(slice_var_inputs_dyn_shape5) { - // attr starts and ends set; inputs are (data, input_axes) + // attr starts set; inputs are (data, input_ends, input_axes) migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape ends{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 6}, {0, 6}}}, - migraphx::make_op("slice", {{"starts", {0, 1}}, {"ends", {8, 8}}}), + migraphx::make_op("slice", {{"starts", {0, 2}}, {"mode", "ends_axes_input"}}), input, + ends, axes); } -TEST_CASE(slice_var_inputs_dyn_mismatch_error2) +TEST_CASE(slice_var_inputs_dyn_shape6) { - migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {2, 4, {2, 4}}, {2, 4, {2, 4}}}}; + migraphx::shape starts{migraphx::shape::int64_type, {2}}; + migraphx::shape ends{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; - throws_shape( - migraphx::make_op("slice", {{"starts", {0, 1, 2}}, {"ends", {3, 4, 4}}}), input, axes); + expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 4}, {0, 4}}}, + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), + input, + starts, + ends, + axes); } -TEST_CASE(slice_var_inputs_dyn_shape3) +TEST_CASE(slice_var_inputs_static_mismatch_error0) { - // attr axes set; inputs are (data, input_starts, input_ends) - migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; + throws_shape( + migraphx::make_op("slice", + {{"ends", {2, 3, 4}}, {"axes", {0, 1, 2}}, {"mode", "starts_input"}}), + input, + starts); +} + +TEST_CASE(slice_var_inputs_static_mismatch_error1) +{ + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{3, 6}, {0, 6}, {0, 6}}}, - migraphx::make_op("slice", {{"axes", {1, 2}}}), - input, - starts, - ends); + throws_shape( + migraphx::make_op("slice", + {{"starts", {0, 1, 2}}, {"axes", {0, 1, 2}}, {"mode", "ends_input"}}), + input, + ends); } -TEST_CASE(slice_var_inputs_dyn_mismatch_error3) +TEST_CASE(slice_var_inputs_static_mismatch_error2) { - migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; + migraphx::shape axes{migraphx::shape::int64_type, {2}}; + throws_shape( + migraphx::make_op("slice", + {{"starts", {0, 1, 2}}, {"ends", {3, 4, 4}}, {"mode", "axes_input"}}), + input, + axes); +} + +TEST_CASE(slice_var_inputs_static_mismatch_error3) +{ + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; - throws_shape(migraphx::make_op("slice", {{"axes", {0, 1, 2}}}), input, starts, ends); + throws_shape( + migraphx::make_op("slice", {{"axes", {0, 1, 2}}, {"mode", "starts_ends_input"}}), + input, + starts, + ends); } -TEST_CASE(slice_var_inputs_dyn_shape4) +TEST_CASE(slice_var_inputs_static_mismatch_error4) { - // attr ends set; inputs are (data, input_starts, input_axes) - migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 6}, {0, 6}}}, - migraphx::make_op("slice", {{"ends", {3, 4}}}), + throws_shape( + migraphx::make_op("slice", {{"ends", {3, 3, 3}}, {"mode", "starts_axes_input"}}), + input, + starts, + axes); +} + +TEST_CASE(slice_var_inputs_static_mismatch_error5) +{ + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; + migraphx::shape ends{migraphx::shape::int64_type, {2}}; + migraphx::shape axes{migraphx::shape::int64_type, {2}}; + throws_shape( + migraphx::make_op("slice", {{"starts", {0, 1, 2}}, {"mode", "ends_axes_input"}}), + input, + ends, + axes); +} + +TEST_CASE(slice_var_inputs_static_mismatch_error6) +{ + migraphx::shape input{migraphx::shape::float_type, {3, 4, 4}}; + migraphx::shape starts{migraphx::shape::int64_type, {2}}; + migraphx::shape ends{migraphx::shape::int64_type, {2}}; + migraphx::shape axes{migraphx::shape::int64_type, {3}}; + throws_shape(migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), input, starts, + ends, axes); } -TEST_CASE(slice_var_inputs_dyn_mismatch_error4) +TEST_CASE(slice_var_inputs_dyn_mismatch_error0) { migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; - migraphx::shape axes{migraphx::shape::int64_type, {2}}; - throws_shape(migraphx::make_op("slice", {{"ends", {3, 3, 3}}}), input, starts, axes); + throws_shape( + migraphx::make_op("slice", + {{"ends", {2, 3, 4}}, {"axes", {0, 1, 2}}, {"mode", "starts_input"}}), + input, + starts); } -TEST_CASE(slice_var_inputs_dyn_shape5) +TEST_CASE(slice_var_inputs_dyn_mismatch_error1) { - // attr starts set; inputs are (data, input_ends, input_axes) migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; + throws_shape( + migraphx::make_op("slice", + {{"starts", {0, 1, 2}}, {"axes", {0, 1, 2}}, {"mode", "ends_input"}}), + input, + ends); +} + +TEST_CASE(slice_var_inputs_dyn_mismatch_error2) +{ + migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 6}, {0, 6}}}, - migraphx::make_op("slice", {{"starts", {0, 2}}}), - input, - ends, - axes); + throws_shape( + migraphx::make_op("slice", + {{"starts", {0, 1, 2}}, {"ends", {3, 4, 4}}, {"mode", "axes_input"}}), + input, + axes); } -TEST_CASE(slice_var_inputs_dyn_mismatch_error5) +TEST_CASE(slice_var_inputs_dyn_mismatch_error3) { migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; + migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; - migraphx::shape axes{migraphx::shape::int64_type, {2}}; - throws_shape(migraphx::make_op("slice", {{"starts", {0, 1, 2}}}), input, ends, axes); + throws_shape( + migraphx::make_op("slice", {{"axes", {0, 1, 2}}, {"mode", "starts_ends_input"}}), + input, + starts, + ends); } -TEST_CASE(slice_var_inputs_dyn_shape6) +TEST_CASE(slice_var_inputs_dyn_mismatch_error4) { - migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {2, 4, {2, 4}}, {2, 4, {2, 4}}}}; + migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; migraphx::shape starts{migraphx::shape::int64_type, {2}}; + migraphx::shape axes{migraphx::shape::int64_type, {2}}; + throws_shape( + migraphx::make_op("slice", {{"ends", {3, 3, 3}}, {"mode", "starts_axes_input"}}), + input, + starts, + axes); +} + +TEST_CASE(slice_var_inputs_dyn_mismatch_error5) +{ + migraphx::shape input{migraphx::shape::float_type, {{3, 6}, {4, 6}, {4, 6}}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {2}}; - expect_shape(migraphx::shape{migraphx::shape::float_type, {{0, 6}, {0, 4}, {0, 4}}}, - migraphx::make_op("slice"), - input, - starts, - ends, - axes); + throws_shape( + migraphx::make_op("slice", {{"starts", {0, 1, 2}}, {"mode", "ends_axes_input"}}), + input, + ends, + axes); } TEST_CASE(slice_var_inputs_dyn_mismatch_error6) @@ -5184,7 +5236,11 @@ TEST_CASE(slice_var_inputs_dyn_mismatch_error6) migraphx::shape starts{migraphx::shape::int64_type, {2}}; migraphx::shape ends{migraphx::shape::int64_type, {2}}; migraphx::shape axes{migraphx::shape::int64_type, {3}}; - throws_shape(migraphx::make_op("slice"), input, starts, ends, axes); + throws_shape(migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), + input, + starts, + ends, + axes); } TEST_CASE(slice_dyn_shape0) @@ -5366,12 +5422,18 @@ TEST_CASE(slice_sym_fixed_bound_var) EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); } -TEST_CASE(slice_sym_non_fixed_throws) +TEST_CASE(slice_sym_non_fixed_axis) { - // Slicing on a non-fixed symbolic axis is rejected (same contract as range). - auto n = var("n", {1, 8}); - migraphx::shape sin{migraphx::shape::float_type, {dd{lit(4)}, dd{n}, dd{lit(8)}}}; - throws_shape(migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {2}}}), sin); + // Slicing a non-fixed symbolic axis is allowed: the bounds are trusted to be valid at + // runtime, so the extent is just ends - starts (no clamping). The other symbol is kept. + auto n = var("n", {1, 8}); + auto m = var("m", {1, 8}); + std::unordered_map sym_map = {{n, 5}, {m, 3}}; + migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{m}}}; + auto op = migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}); + migraphx::shape sout{migraphx::shape::float_type, {dd{lit(2)}, dd{m}}, sin.dyn_strides()}; + expect_shape(sout, op, sin); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); } TEST_CASE(slice_sym_nonstandard_layout) @@ -5387,6 +5449,138 @@ TEST_CASE(slice_sym_nonstandard_layout) EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); } +TEST_CASE(slice_sym_clamped_and_negative_bounds) +{ + // Concrete bounds on a symbolic input are clipped/resolved against the fixed sliced axis + // like the static path; the symbolic axis is untouched. + auto n = var("n", {1, 8}); + std::unordered_map sym_map = {{n, 5}}; + migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(3)}}}; + + { + // end clipped to len 3: dim = 3 - 2 = 1. + auto op = migraphx::make_op("slice", {{"axes", {2}}, {"starts", {2}}, {"ends", {10}}}); + migraphx::shape sout{ + migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(1)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); + } + { + // negative start: -1 -> 2, dim = 3 - 2 = 1. + auto op = migraphx::make_op("slice", {{"axes", {2}}, {"starts", {-1}}, {"ends", {10}}}); + migraphx::shape sout{ + migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(1)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); + } + { + // negative end on axis 1 (len 2): -1 -> 1, dim = 1 - 0 = 1. + auto op = migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {-1}}}); + migraphx::shape sout{ + migraphx::shape::float_type, {dd{n}, dd{lit(1)}, dd{lit(3)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); + } +} + +TEST_CASE(slice_sym_symbolic_end_static_input) +{ + // Static input + symbolic end bound: output is symbolic (not fixed -> not demoted to static). + // Symbolic attributes require 2+ inputs, so the runtime end value is supplied as an input. + auto n = var("n", {1, 16}); + auto op = migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}, + {"mode", "ends_input"}}); + + // end=n is clamped to the axis length 10: dim = min(n, 10). + migraphx::shape sin{migraphx::shape::float_type, {10}}; + migraphx::shape ends_in{migraphx::shape::int64_type, {1}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{migraphx::sym::min(n, lit(10))}}, {lit(1)}}; + expect_shape(sout, op, sin, ends_in); + EXPECT(sout.symbolic()); + EXPECT(not sout.is_fixed()); + EXPECT(sout.to_static({{n, 7}}) == migraphx::shape{migraphx::shape::float_type, {7}, {1}}); + EXPECT(sout.to_static({{n, 10}}) == migraphx::shape{migraphx::shape::float_type, {10}, {1}}); +} + +TEST_CASE(slice_sym_symbolic_bounds) +{ + // The sliced extent (ends - starts) must be non-negative across the whole variable + // range, so each var range is chosen to keep end >= start. + { + // Symbolic end clamped to the axis length 12: dim = min(n, 12) - 2. + auto m = var("m", {1, 16}); + auto n = var("n", {2, 16}); + auto op = migraphx::make_op("slice", + {{"axes", {1}}, + {"starts", {2}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}, + {"mode", "ends_input"}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{m}, dd{lit(12)}}}; + migraphx::shape ends_in{migraphx::shape::int64_type, {1}}; + migraphx::shape sout{migraphx::shape::float_type, + {dd{m}, dd{migraphx::sym::min(n, lit(12)) - lit(2)}}, + sin.dyn_strides()}; + expect_shape(sout, op, sin, ends_in); + EXPECT(sout.symbolic()); + EXPECT(not sout.is_fixed()); + EXPECT(sout.to_static({{m, 4}, {n, 9}}) == + migraphx::shape{migraphx::shape::float_type, {4, 7}, {12, 1}}); + } + { + // Symbolic end provably >= the axis length collapses to the length: extent is concrete. + auto n = var("n", {13, 20}); + auto op = migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {2}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}, + {"mode", "ends_input"}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{lit(12)}, dd{lit(4)}}}; + migraphx::shape ends_in{migraphx::shape::int64_type, {1}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{lit(10)}, dd{lit(4)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, ends_in); + } + { + // Symbolic start: dim = 8 - n (n <= 8 keeps the extent non-negative). + auto n = var("n", {1, 8}); + auto op = migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{n})}}, + {"ends", {8}}, + {"mode", "starts_input"}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{lit(10)}, dd{lit(4)}}}; + migraphx::shape starts_in{migraphx::shape::int64_type, {1}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{lit(8) - n}, dd{lit(4)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, starts_in); + EXPECT(sout.symbolic()); + EXPECT(sout.to_static({{n, 3}}) == + migraphx::shape{migraphx::shape::float_type, {5, 4}, {4, 1}}); + } + { + // Both bounds symbolic: dim = n - m (ranges disjoint so n >= m always). + auto m = var("m", {1, 5}); + auto n = var("n", {5, 16}); + auto op = migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{m})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}, + {"mode", "starts_ends_input"}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{lit(20)}}}; + migraphx::shape starts_in{migraphx::shape::int64_type, {1}}; + migraphx::shape ends_in{migraphx::shape::int64_type, {1}}; + migraphx::shape sout{migraphx::shape::float_type, {dd{n - m}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, starts_in, ends_in); + EXPECT(sout.symbolic()); + EXPECT(sout.to_static({{m, 3}, {n, 9}}) == + migraphx::shape{migraphx::shape::float_type, {6}, {1}}); + } +} + TEST_CASE(test_scan_slice1) { migraphx::shape input{migraphx::shape::float_type, {2, 3, 4}}; diff --git a/test/quantization.cpp b/test/quantization.cpp index a33c88badec..e12c2de44ac 100644 --- a/test/quantization.cpp +++ b/test/quantization.cpp @@ -446,8 +446,10 @@ TEST_CASE(topk) auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto data = mm->add_parameter("data", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 0}, {"k", 3}, {"largest", 0}}), data); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {3}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", 0}, {"k", 3}, {"largest", 0}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({r0, r1}); @@ -460,10 +462,12 @@ TEST_CASE(topk) auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto fdata = mm->add_parameter("data", s); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {3}}); auto hdata = mm->add_instruction( migraphx::make_op("convert", {{"target_type", migraphx::shape::half_type}}), fdata); auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 0}, {"k", 3}, {"largest", 0}}), hdata); + migraphx::make_op("topk", {{"axis", 0}, {"k", 3}, {"largest", 0}}), hdata, kk); auto hr0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto fr0 = mm->add_instruction( migraphx::make_op("convert", {{"target_type", migraphx::shape::float_type}}), hr0); diff --git a/test/ref/nonmaxsuppression.cpp b/test/ref/nonmaxsuppression.cpp index 552f3e2b816..3ca60073fc8 100644 --- a/test/ref/nonmaxsuppression.cpp +++ b/test/ref/nonmaxsuppression.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include #include @@ -35,8 +36,18 @@ static migraphx::instruction_ref add_nms_dynamic_slice(migraphx::module* mm, { auto idx = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), nms); auto cnt = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), nms); + // Slice the padded indices down to the runtime num_selected value using a symbolic + // dimension, matching the IR the NonMaxSuppression ONNX parser emits. + auto num_selected_var = + migraphx::shape::dynamic_dimension{migraphx::sym::var("nms_num_selected")}; return mm->add_instruction( - migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}}), idx, cnt); + migraphx::make_op("slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(num_selected_var)}}, + {"mode", "ends_input"}}), + idx, + cnt); } TEST_CASE(nms_dyn_out_test) diff --git a/test/ref/slice.cpp b/test/ref/slice.cpp index efde93f1ab0..0550434f870 100644 --- a/test/ref/slice.cpp +++ b/test/ref/slice.cpp @@ -87,7 +87,11 @@ TEST_CASE(slice_var_inputs_static0) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto starts = mm->add_parameter("starts", s1); auto ends = mm->add_parameter("ends", s1); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {2}}}), l0, starts, ends); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"mode", "starts_ends_input"}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -113,7 +117,11 @@ TEST_CASE(slice_var_inputs_static1) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto starts = mm->add_parameter("starts", s1); auto ends = mm->add_parameter("ends", s1); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {2}}}), l0, starts, ends); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"mode", "starts_ends_input"}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -140,7 +148,8 @@ TEST_CASE(slice_var_inputs_static2) auto starts = mm->add_parameter("starts", s1); auto ends = mm->add_parameter("ends", s1); auto axes = mm->add_parameter("axes", s1); - mm->add_instruction(migraphx::make_op("slice"), l0, starts, ends, axes); + mm->add_instruction( + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), l0, starts, ends, axes); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -165,7 +174,10 @@ TEST_CASE(slice_var_inputs_dyn0) auto input = mm->add_parameter("input", s0); migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto starts = mm->add_parameter("starts", s1); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {2}}, {"ends", {10}}}), input, starts); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"ends", {10}}, {"mode", "starts_input"}}), + input, + starts); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -190,7 +202,10 @@ TEST_CASE(slice_var_inputs_dyn1) auto input = mm->add_parameter("input", s0); migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto ends = mm->add_parameter("ends", s1); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {2}}, {"starts", {-5}}}), input, ends); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"starts", {-5}}, {"mode", "ends_input"}}), + input, + ends); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -215,7 +230,10 @@ TEST_CASE(slice_var_inputs_dyn2) auto input = mm->add_parameter("input", s0); migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto axes = mm->add_parameter("axes", s1); - mm->add_instruction(migraphx::make_op("slice", {{"starts", {1}}, {"ends", {-1}}}), input, axes); + mm->add_instruction( + migraphx::make_op("slice", {{"starts", {1}}, {"ends", {-1}}, {"mode", "axes_input"}}), + input, + axes); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -241,7 +259,11 @@ TEST_CASE(slice_var_inputs_dyn3) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto starts = mm->add_parameter("starts", s1); auto ends = mm->add_parameter("ends", s1); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {2}}}), input, starts, ends); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"mode", "starts_ends_input"}}), + input, + starts, + ends); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -269,10 +291,12 @@ TEST_CASE(slice_var_inputs_dyn4) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto starts = mm->add_parameter("starts", s1); auto axes = mm->add_parameter("axes", s1); - mm->add_instruction(migraphx::make_op("slice", {{"ends", {std::numeric_limits::max()}}}), - input, - starts, - axes); + mm->add_instruction( + migraphx::make_op( + "slice", {{"ends", {std::numeric_limits::max()}}, {"mode", "starts_axes_input"}}), + input, + starts, + axes); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -300,7 +324,11 @@ TEST_CASE(slice_var_inputs_dyn5) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto ends = mm->add_parameter("ends", s1); auto axes = mm->add_parameter("axes", s1); - mm->add_instruction(migraphx::make_op("slice", {{"starts", {-4}}}), input, ends, axes); + mm->add_instruction( + migraphx::make_op("slice", {{"starts", {-4}}, {"mode", "ends_axes_input"}}), + input, + ends, + axes); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; @@ -328,7 +356,11 @@ TEST_CASE(slice_var_inputs_dyn6) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto starts = mm->add_parameter("starts", s1); auto ends = mm->add_parameter("ends", s1); - mm->add_instruction(migraphx::make_op("slice", {{"axes", {2}}}), input, starts, ends); + mm->add_instruction( + migraphx::make_op("slice", {{"axes", {2}}, {"mode", "starts_ends_input"}}), + input, + starts, + ends); p.compile(migraphx::make_target("ref")); migraphx::parameter_map params; diff --git a/test/ref/topk.cpp b/test/ref/topk.cpp index 3f71f4f8879..535bb6ecd05 100644 --- a/test/ref/topk.cpp +++ b/test/ref/topk.cpp @@ -37,6 +37,9 @@ static auto run_program(const migraphx::value& op, bool custom_idx = false, bool migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto data = mm->add_parameter("data", s); std::vector topk_inputs = {data}; + auto k_val = op.at("k").to(); + topk_inputs.push_back( + mm->add_literal(migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {k_val}})); if(custom_idx) { migraphx::shape is{migraphx::shape::uint16_type, {3, 5}}; @@ -155,9 +158,11 @@ TEST_CASE(topk_k_greater_than_n_dynamic) std::vector dds = {{1, 100}}; migraphx::shape s{migraphx::shape::float_type, dds}; auto data = mm->add_parameter("data", s); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {100}}); // k=100 is the max placeholder from parse time auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 0}, {"k", 100}, {"largest", 1}}), data); + migraphx::make_op("topk", {{"axis", 0}, {"k", 100}, {"largest", 1}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({r0, r1}); @@ -191,9 +196,11 @@ TEST_CASE(topk_k_equals_n) auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto data = mm->add_parameter("data", s); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {5}}); // k=5 equals axis=1 dimension of 5 auto r = mm->add_instruction(migraphx::make_op("topk", {{"axis", 1}, {"k", 5}, {"largest", 0}}), - data); + data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({r0, r1}); diff --git a/test/rewrite_topk.cpp b/test/rewrite_topk.cpp index 4d29aac52b5..22a338a7a3a 100644 --- a/test/rewrite_topk.cpp +++ b/test/rewrite_topk.cpp @@ -39,7 +39,8 @@ TEST_CASE(small_topk) migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::float_type, {300}}); - auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}}), x); + auto k = m1.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); + auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}}), x, k); m1.add_return({r}); } migraphx::module m2 = m1; @@ -52,7 +53,8 @@ TEST_CASE(large_topk_no_split) migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::float_type, {240000}}); - auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 120000}}), x); + auto k = m1.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {120000}}); + auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 120000}}), x, k); m1.add_return({r}); } migraphx::module m2 = m1; @@ -66,7 +68,8 @@ TEST_CASE(split_topk_batch_1) migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::float_type, {n}}); - auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}}), x); + auto k = m1.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); + auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}}), x, k); m1.add_return({r}); } run_pass(m1); @@ -76,6 +79,7 @@ TEST_CASE(split_topk_batch_1) std::vector indices(n); std::iota(indices.begin(), indices.end(), 0); auto x = m2.add_parameter("x", {migraphx::shape::float_type, {n}}); + auto k = m2.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); auto input_idx = m2.add_literal(migraphx::literal{{migraphx::shape::uint32_type, {n}}, indices}); auto input_idxb = m2.add_instruction( @@ -84,15 +88,15 @@ TEST_CASE(split_topk_batch_1) migraphx::make_op("reshape", {{"dims", {group, n / group}}}), input_idxb); auto xr = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {group, n / group}}}), x); - auto r1 = - m2.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), xr, input_idxr); + auto r1 = m2.add_instruction( + migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), xr, k, input_idxr); auto value1 = m2.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r1); auto idx1 = m2.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r1); auto valuer = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {8 * group}}}), value1); auto idxr = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {8 * group}}}), idx1); - auto r2 = - m2.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 0}}), valuer, idxr); + auto r2 = m2.add_instruction( + migraphx::make_op("topk", {{"k", 8}, {"axis", 0}}), valuer, k, idxr); m2.add_return({r2}); } EXPECT(m1.sort() == m2.sort()); @@ -105,7 +109,8 @@ TEST_CASE(split_topk_batch_64) migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::float_type, {batch, n}}); - auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), x); + auto k = m1.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); + auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), x, k); m1.add_return({r}); } run_pass(m1); @@ -115,6 +120,7 @@ TEST_CASE(split_topk_batch_64) std::vector indices(n); std::iota(indices.begin(), indices.end(), 0); auto x = m2.add_parameter("x", {migraphx::shape::float_type, {batch, n}}); + auto k = m2.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); auto input_idx = m2.add_literal(migraphx::literal{{migraphx::shape::uint32_type, {n}}, indices}); auto input_idxb = m2.add_instruction( @@ -123,16 +129,16 @@ TEST_CASE(split_topk_batch_64) migraphx::make_op("reshape", {{"dims", {batch, group, n / group}}}), input_idxb); auto xr = m2.add_instruction( migraphx::make_op("reshape", {{"dims", {batch, group, n / group}}}), x); - auto r1 = - m2.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 2}}), xr, input_idxr); + auto r1 = m2.add_instruction( + migraphx::make_op("topk", {{"k", 8}, {"axis", 2}}), xr, k, input_idxr); auto value1 = m2.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r1); auto idx1 = m2.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r1); auto valuer = m2.add_instruction( migraphx::make_op("reshape", {{"dims", {batch, 8 * group}}}), value1); auto idxr = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {batch, 8 * group}}}), idx1); - auto r2 = - m2.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), valuer, idxr); + auto r2 = m2.add_instruction( + migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), valuer, k, idxr); m2.add_return({r2}); } EXPECT(m1.sort() == m2.sort()); @@ -145,7 +151,8 @@ TEST_CASE(split_topk_batch_64_last) migraphx::module m1; { auto x = m1.add_parameter("x", {migraphx::shape::float_type, {n, batch}}); - auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 0}}), x); + auto k = m1.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); + auto r = m1.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 0}}), x, k); m1.add_return({r}); } run_pass(m1); @@ -155,6 +162,7 @@ TEST_CASE(split_topk_batch_64_last) std::vector indices(n); std::iota(indices.begin(), indices.end(), 0); auto x = m2.add_parameter("x", {migraphx::shape::float_type, {n, batch}}); + auto k = m2.add_literal(migraphx::literal{{migraphx::shape::int64_type, {1}}, {8}}); auto input_idx = m2.add_literal(migraphx::literal{{migraphx::shape::uint32_type, {n}}, indices}); auto input_idxb = m2.add_instruction( @@ -163,16 +171,16 @@ TEST_CASE(split_topk_batch_64_last) migraphx::make_op("reshape", {{"dims", {group, n / group, batch}}}), input_idxb); auto xr = m2.add_instruction( migraphx::make_op("reshape", {{"dims", {group, n / group, batch}}}), x); - auto r1 = - m2.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), xr, input_idxr); + auto r1 = m2.add_instruction( + migraphx::make_op("topk", {{"k", 8}, {"axis", 1}}), xr, k, input_idxr); auto value1 = m2.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r1); auto idx1 = m2.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r1); auto valuer = m2.add_instruction( migraphx::make_op("reshape", {{"dims", {8 * group, batch}}}), value1); auto idxr = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {8 * group, batch}}}), idx1); - auto r2 = - m2.add_instruction(migraphx::make_op("topk", {{"k", 8}, {"axis", 0}}), valuer, idxr); + auto r2 = m2.add_instruction( + migraphx::make_op("topk", {{"k", 8}, {"axis", 0}}), valuer, k, idxr); m2.add_return({r2}); } EXPECT(m1.sort() == m2.sort()); diff --git a/test/simplify_dyn_ops_test.cpp b/test/simplify_dyn_ops_test.cpp index b6a770df194..85d1895a53e 100644 --- a/test/simplify_dyn_ops_test.cpp +++ b/test/simplify_dyn_ops_test.cpp @@ -308,7 +308,9 @@ TEST_CASE(const_slice_2input_ends_axes) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto input_starts = m0.add_literal(migraphx::literal{s1, {0}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"ends", {3}}, {"axes", {0}}}), input, input_starts); + migraphx::make_op("slice", {{"ends", {3}}, {"axes", {0}}, {"mode", "starts_input"}}), + input, + input_starts); m0.add_return({slice_ins}); } run_pass(m0); @@ -333,7 +335,9 @@ TEST_CASE(const_slice_2input_starts_axes) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto input_ends = m0.add_literal(migraphx::literal{s1, {3}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"starts", {0}}, {"axes", {0}}}), input, input_ends); + migraphx::make_op("slice", {{"starts", {0}}, {"axes", {0}}, {"mode", "ends_input"}}), + input, + input_ends); m0.add_return({slice_ins}); } run_pass(m0); @@ -358,7 +362,9 @@ TEST_CASE(const_slice_2input_starts_ends) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto input_axes = m0.add_literal(migraphx::literal{s1, {0}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}}), input, input_axes); + migraphx::make_op("slice", {{"starts", {0}}, {"ends", {3}}, {"mode", "axes_input"}}), + input, + input_axes); m0.add_return({slice_ins}); } run_pass(m0); @@ -384,7 +390,10 @@ TEST_CASE(const_slice_3input_axes_only) auto input_starts = m0.add_literal(migraphx::literal{s1, {0}}); auto input_ends = m0.add_literal(migraphx::literal{s1, {3}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"axes", {0}}}), input, input_starts, input_ends); + migraphx::make_op("slice", {{"axes", {0}}, {"mode", "starts_ends_input"}}), + input, + input_starts, + input_ends); m0.add_return({slice_ins}); } run_pass(m0); @@ -410,7 +419,10 @@ TEST_CASE(const_slice_3input_ends_only) auto input_starts = m0.add_literal(migraphx::literal{s1, {0}}); auto input_axes = m0.add_literal(migraphx::literal{s1, {0}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"ends", {3}}}), input, input_starts, input_axes); + migraphx::make_op("slice", {{"ends", {3}}, {"mode", "starts_axes_input"}}), + input, + input_starts, + input_axes); m0.add_return({slice_ins}); } run_pass(m0); @@ -436,7 +448,10 @@ TEST_CASE(const_slice_3inputs_starts_only) auto input_ends = m0.add_literal(migraphx::literal{s1, {3}}); auto input_axes = m0.add_literal(migraphx::literal{s1, {0}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"starts", {0}}}), input, input_ends, input_axes); + migraphx::make_op("slice", {{"starts", {0}}, {"mode", "ends_axes_input"}}), + input, + input_ends, + input_axes); m0.add_return({slice_ins}); } run_pass(m0); @@ -461,7 +476,9 @@ TEST_CASE(const_slice_2input_ends_axes_dyn) migraphx::shape s1{migraphx::shape::int32_type, {1}}; auto input_starts = m0.add_literal(migraphx::literal{s1, {0}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"ends", {3}}, {"axes", {0}}}), input, input_starts); + migraphx::make_op("slice", {{"ends", {3}}, {"axes", {0}}, {"mode", "starts_input"}}), + input, + input_starts); m0.add_return({slice_ins}); } run_pass(m0); @@ -488,7 +505,10 @@ TEST_CASE(const_slice_3input_dyn) auto input_starts = m0.add_literal(migraphx::literal{s1, {0}}); auto input_ends = m0.add_literal(migraphx::literal{s1, {3}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice", {{"axes", {0}}}), input, input_starts, input_ends); + migraphx::make_op("slice", {{"axes", {0}}, {"mode", "starts_ends_input"}}), + input, + input_starts, + input_ends); m0.add_return({slice_ins}); } run_pass(m0); @@ -515,7 +535,11 @@ TEST_CASE(const_slice_4input) auto input_ends = m0.add_literal(migraphx::literal{s1, {3}}); auto input_axes = m0.add_literal(migraphx::literal{s1, {0}}); auto slice_ins = m0.add_instruction( - migraphx::make_op("slice"), input, input_starts, input_ends, input_axes); + migraphx::make_op("slice", {{"mode", "starts_ends_axes_input"}}), + input, + input_starts, + input_ends, + input_axes); m0.add_return({slice_ins}); } run_pass(m0); diff --git a/test/verify/test_topk.cpp b/test/verify/test_topk.cpp index af001797aa8..92809e2a99a 100644 --- a/test/verify/test_topk.cpp +++ b/test/verify/test_topk.cpp @@ -39,8 +39,10 @@ struct test_topk : verify_program> unsigned int batch = 3; migraphx::shape s{DType, {batch, N}}; auto x1 = mm->add_parameter("x1", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 1}, {"k", K}, {"largest", 1}}), x1); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {K}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", 1}, {"k", K}, {"largest", 1}}), x1, kk); auto values = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto indices = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({values, indices}); diff --git a/test/verify/test_topk_0.cpp b/test/verify/test_topk_0.cpp index 675d67efbcd..ed1cd51354d 100644 --- a/test/verify/test_topk_0.cpp +++ b/test/verify/test_topk_0.cpp @@ -36,8 +36,10 @@ struct test_topk_0 : verify_program> auto* mm = p.get_main_module(); migraphx::shape s{DType, {3, 5}}; auto data = mm->add_parameter("data", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 1}, {"k", 4}, {"largest", 1}}), data); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {4}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", 1}, {"k", 4}, {"largest", 1}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); mm->add_return({r0}); diff --git a/test/verify/test_topk_1.cpp b/test/verify/test_topk_1.cpp index 3e5d0bf331f..9d98131a23d 100644 --- a/test/verify/test_topk_1.cpp +++ b/test/verify/test_topk_1.cpp @@ -35,8 +35,10 @@ struct test_topk_1 : verify_program auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto data = mm->add_parameter("data", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", -2}, {"k", 3}, {"largest", 1}}), data); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {3}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", -2}, {"k", 3}, {"largest", 1}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({r0, r1}); diff --git a/test/verify/test_topk_2.cpp b/test/verify/test_topk_2.cpp index 44f0283579e..ebde18d88b1 100644 --- a/test/verify/test_topk_2.cpp +++ b/test/verify/test_topk_2.cpp @@ -35,8 +35,10 @@ struct test_topk_2 : verify_program auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto data = mm->add_parameter("data", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 1}, {"k", 4}, {"largest", 0}}), data); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {4}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", 1}, {"k", 4}, {"largest", 0}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); mm->add_return({r0}); diff --git a/test/verify/test_topk_3.cpp b/test/verify/test_topk_3.cpp index 44d15f9b22d..26ec829e17a 100644 --- a/test/verify/test_topk_3.cpp +++ b/test/verify/test_topk_3.cpp @@ -35,8 +35,10 @@ struct test_topk_3 : verify_program auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::float_type, {3, 5}}; auto data = mm->add_parameter("data", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", -2}, {"k", 3}, {"largest", 0}}), data); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {3}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", -2}, {"k", 3}, {"largest", 0}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({r0, r1}); diff --git a/test/verify/test_topk_dynamic.cpp b/test/verify/test_topk_dynamic.cpp index 25b0836df00..0bdb4a853dd 100644 --- a/test/verify/test_topk_dynamic.cpp +++ b/test/verify/test_topk_dynamic.cpp @@ -38,8 +38,10 @@ struct test_topk_dynamic : verify_program> std::vector dds = {{1, 100}}; migraphx::shape s{migraphx::shape::float_type, dds}; auto data = mm->add_parameter("data", s); - auto r = mm->add_instruction( - migraphx::make_op("topk", {{"axis", 0}, {"k", 100}, {"largest", 1}}), data); + auto kk = mm->add_literal( + migraphx::literal{migraphx::shape{migraphx::shape::int64_type, {1}}, {100}}); + auto r = mm->add_instruction( + migraphx::make_op("topk", {{"axis", 0}, {"k", 100}, {"largest", 1}}), data, kk); auto r0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), r); auto r1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), r); mm->add_return({r0, r1});