From 86195d36e40e90cf1349cf8b4b2fdcad51333da4 Mon Sep 17 00:00:00 2001 From: charlie Date: Tue, 4 Aug 2026 18:32:00 -0500 Subject: [PATCH 1/6] Dynamic slice from symbolic slice in other branch --- CHANGELOG.md | 3 + src/CMakeLists.txt | 1 + src/include/migraphx/dim_like.hpp | 37 ++ src/include/migraphx/op/dyn_slice.hpp | 199 ++++++++++ .../migraphx/op/normalize_attribute.hpp | 7 +- src/include/migraphx/op/slice.hpp | 19 +- src/include/migraphx/operators.hpp | 1 + src/normalize_attributes.cpp | 100 ++++- src/sym.cpp | 8 + src/targets/gpu/lowering.cpp | 6 +- test/gpu/dyn_slice_lowering.cpp | 32 ++ test/normalize_ops_test.cpp | 222 +++++++++++ test/op_shape_test.cpp | 360 +++++++++++++----- test/ref/dyn_slice.cpp | 279 ++++++++++++++ test/sym.cpp | 58 +++ 15 files changed, 1222 insertions(+), 110 deletions(-) create mode 100644 src/include/migraphx/op/dyn_slice.hpp create mode 100644 test/ref/dyn_slice.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index b8f6e00b721..79fa01f196f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,8 @@ Full documentation for MIGraphX is available at * Added mixed length gather fusion in same_table_gather_horizontal_fusion to bundle gather kernels that share the same data (#5044). * Added a verbose terminate handler for exceptions on Windows (#5084). * Added a `--start-from` or `-s` flag to test binaries which resumes from a test name in the list instead of the beginning (#5072). +* Added a `dyn_slice` operator, `dyn_slice(data, starts, ends)`, that describes its bound inputs with symbolic attributes so slicing by a data-dependent bound keeps a symbolic output shape; the axes are an attribute since they must be known when the shape is computed (#5088). +* Added symbolic normalization of operator attributes, selected with the `use_sym` normalize attribute, which clamps a bound against a symbolic axis length instead of leaving it unnormalized (#5088). ### Changed @@ -61,6 +63,7 @@ Full documentation for MIGraphX is available at * Updated python API to allow getting and adding debug symbols from instructions. (#4803) * Allow for 1 arg slicing over a dynamic dimension. (#5015) * Route convolutions and dot operations through rocMLIR when MIOpen or GEMM libraries are disabled at build time (#5059). +* The `slice` operator now rejects symbolic input shapes since its integer bounds cannot express a symbolic output extent; use `dyn_slice` for those (#5088). ### Resolved issues diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 4b18a01803d..bf770276c3b 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -215,6 +215,7 @@ register_migraphx_ops( dimensions_of div dot + dyn_slice elu equal erf diff --git a/src/include/migraphx/dim_like.hpp b/src/include/migraphx/dim_like.hpp index 1a90b6f9ed7..7049592175f 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 { @@ -65,6 +68,40 @@ inline std::ostream& operator<<(std::ostream& os, const dim_like& d) return os; } +inline bool all_ints(const std::vector& dims) +{ + return std::all_of(dims.begin(), dims.end(), [](const dim_like& d) { + return std::holds_alternative(d); + }); +} + +/// Extracts the concrete int64_t from each entry; throws (via std::get) if any entry holds a +/// 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; +} + +/// Converts each entry to a symbolic expression. A range-based dynamic_dimension has no +/// expression to convert, so it is rejected. +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 sym::lit(std::get(d)); + if(not is_symbolic(d)) + MIGRAPHX_THROW("DIM_LIKE: cannot convert a range-based dimension to a symbolic " + "expression"); + return std::get(d).sym_expr; + }); + return result; +} + 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/op/dyn_slice.hpp b/src/include/migraphx/op/dyn_slice.hpp new file mode 100644 index 00000000000..47842fc97c4 --- /dev/null +++ b/src/include/migraphx/op/dyn_slice.hpp @@ -0,0 +1,199 @@ +/* + * The MIT License (MIT) + * + * 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 + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +#ifndef MIGRAPHX_GUARD_OPERATORS_DYN_SLICE_HPP +#define MIGRAPHX_GUARD_OPERATORS_DYN_SLICE_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace migraphx { +inline namespace MIGRAPHX_INLINE_NS { +namespace op { + +/// Slice operator whose bounds are only known at run time. +/// +/// The starts and ends are always supplied as inputs. The attribute of the same name describes +/// that input at compile time: either a concrete value or, when the value is only known +/// symbolically, a symbolic dynamic_dimension whose expression evaluates to what the input will +/// hold at run time. That is what lets the output shape stay symbolic instead of collapsing to a +/// range. The axes have to be known when the shape is computed, so they are an attribute only. +/// +/// Attributes: +/// axes: axes to slice over +/// starts: slice starting indices +/// ends: slice ending indices +/// +/// Parameters: +/// data: the input tensor to slice (static or symbolic shape) +/// starts_input: starting indices of the slice (static shape, 1D) +/// ends_input: ending indices of the slice (static shape, 1D) +struct dyn_slice +{ + std::vector axes{}; + std::vector starts{}; + std::vector ends{}; + + template + static auto reflect(Self& self, F f) + { + return pack(f(self.axes, "axes"), f(self.starts, "starts"), f(self.ends, "ends")); + } + + /// Ensure the axes attribute is within limits, and clip starts and ends to the sliced axis + /// length. Symbolic bounds are clipped symbolically; see tune_attribute_sym(). + value attributes() const + { + value normalize_axes = value::object{}; + normalize_axes["axes"] = value::array{normalize_attribute::include_min}; + normalize_axes["starts"] = value::array{normalize_attribute::clip_max, + normalize_attribute::clip_min, + normalize_attribute::include_max, + normalize_attribute::use_len, + normalize_attribute::include_min, + normalize_attribute::use_sym}; + normalize_axes["ends"] = value::array{normalize_attribute::clip_max, + normalize_attribute::clip_min, + normalize_attribute::include_max, + normalize_attribute::use_len, + normalize_attribute::include_min, + normalize_attribute::use_sym}; + return {{"normalize_axes", normalize_axes}, {"fillcolor", "#FFA500" /* orange */}}; + } + + std::string name() const { return "dyn_slice"; } + + /// Check the attributes against each other and against the inputs. + void check_inputs_and_attributes(const std::vector& inputs) const + { + if(axes.empty() or starts.size() != axes.size() or ends.size() != axes.size()) + { + MIGRAPHX_THROW("DYN_SLICE: axes, starts, and ends attributes must all be set and " + "have the same length"); + } + // The inputs carry the run-time value of the starts and ends attributes, so there is one + // entry per sliced axis. + check_shapes{inputs.begin() + 1, + inputs.end(), + std::string("DYN_SLICE: inputs (starts, ends)"), + false} + .only_dims(1) + .same_dims(); + if(inputs[1].lens().front() != axes.size()) + { + MIGRAPHX_THROW( + "DYN_SLICE: input length (" + migraphx::to_string(inputs[1].lens().front()) + + ") does not match attribute length (" + migraphx::to_string(axes.size()) + ")"); + } + } + + shape normalize_compute_shape(std::vector inputs) const + { + check_shapes{inputs, *this, true}.has(3); + check_inputs_and_attributes(inputs); + auto input_shape = inputs.front(); + if(input_shape.dynamic() and not input_shape.symbolic()) + MIGRAPHX_THROW("DYN_SLICE: data input must have a static or symbolic shape"); + + auto sym_in = input_shape.to_symbolic(); + auto dds = sym_in.dyn_dims(); + auto start_exprs = to_sym_exprs(starts); + auto end_exprs = to_sym_exprs(ends); + std::vector extents(axes.size()); + std::transform(end_exprs.begin(), + end_exprs.end(), + start_exprs.begin(), + extents.begin(), + [](const auto& end, const auto& start) { return end - start; }); + migraphx::for_each(axes.begin(), axes.end(), extents.begin(), [&](auto axis, auto extent) { + dds[axis] = shape::dynamic_dimension{std::move(extent)}; + }); + shape result{input_shape.type(), std::move(dds), sym_in.dyn_strides()}; + // A slice is a view, so a fully concrete result of a static input must stay static. + if(not input_shape.symbolic() and result.is_fixed()) + return result.to_static(); + return result; + } + + argument compute(const dyn_output&, std::vector args) const + { + const auto& input = args.front(); + auto input_shape = input.get_shape(); + auto read = [](const argument& arg) { + std::vector result; + arg.visit([&](auto values) { result = values.template to_vector(); }); + return result; + }; + // The bound attributes are only the compile-time view of the inputs; the inputs hold the + // values to slice with, so they are renormalized against the run-time shape here. The + // axes attribute needs no such handling: shapes with a dynamic rank are not supported, so + // it is already normalized at compile time. + auto axes_attrs = this->attributes().at("normalize_axes"); + auto norm_starts = normalize_indices( + read(args[1]), axes, input_shape, axes_attrs.at("starts"), "DYN_SLICE: starts input"); + auto norm_ends = normalize_indices( + read(args[2]), axes, input_shape, axes_attrs.at("ends"), "DYN_SLICE: ends input"); + + // The compile-time shape asserts a non-negative extent on every sliced axis, so + // inconsistent run-time bounds are rejected instead of wrapping around to a huge length. + std::vector extents(axes.size()); + std::transform(norm_ends.begin(), + norm_ends.end(), + norm_starts.begin(), + extents.begin(), + [](auto end, auto start) { + if(end < start) + MIGRAPHX_THROW("DYN_SLICE: end (" + migraphx::to_string(end) + + ") is before start (" + migraphx::to_string(start) + + ")"); + return std::size_t(end - start); + }); + auto new_lens = input_shape.lens(); + std::vector start_indices(input_shape.ndim(), 0); + migraphx::for_each(axes.begin(), + axes.end(), + norm_starts.begin(), + [&](auto axis, auto start) { start_indices[axis] = start; }); + migraphx::for_each(axes.begin(), axes.end(), extents.begin(), [&](auto axis, auto extent) { + new_lens[axis] = extent; + }); + auto offset = input_shape.index(start_indices) * input_shape.type_size(); + shape output_shape{input_shape.type(), new_lens, input_shape.strides()}; + return {output_shape, [=] { return input.data() + offset; }}; + } + + std::vector output_alias(const std::vector&) const { return {0}; } +}; + +} // namespace op +} // namespace MIGRAPHX_INLINE_NS +} // namespace migraphx + +#endif diff --git a/src/include/migraphx/op/normalize_attribute.hpp b/src/include/migraphx/op/normalize_attribute.hpp index c545c2045ea..7989ca0c3ad 100644 --- a/src/include/migraphx/op/normalize_attribute.hpp +++ b/src/include/migraphx/op/normalize_attribute.hpp @@ -52,6 +52,10 @@ namespace op { * Include or exclude the maximum value/index for range checking and clipping. * 7. `normalize_padding`: * To normalize the padding to `2*(pad ndim)` dimensions. + * 8. `use_sym` vs. `no_sym` (default): + * Whether the attribute can hold a symbolic value, which requires it to be a `dim_like`. + * Such an attribute is normalized symbolically when either the attribute or the axis it is + * normalized against is symbolic. Only meaningful together with `use_len`. */ enum class normalize_attribute { @@ -61,7 +65,8 @@ enum class normalize_attribute clip_min, include_max, include_min, - normalize_padding + normalize_padding, + use_sym }; } // namespace op diff --git a/src/include/migraphx/op/slice.hpp b/src/include/migraphx/op/slice.hpp index 2654561c42b..fcd67c9f6e2 100644 --- a/src/include/migraphx/op/slice.hpp +++ b/src/include/migraphx/op/slice.hpp @@ -58,7 +58,8 @@ namespace op { * ends: constant slice ending indices (optional) * * Parameters: - * data: the input tensor to slice (dynamic or static shape) + * data: the input tensor to slice (static or range-based dynamic shape; a symbolic shape is + * rejected because the output extent cannot be expressed with integer bounds, use dyn_slice) * 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) @@ -263,6 +264,8 @@ struct slice shape normalize_compute_shape(std::vector inputs) const { check_shapes{inputs, *this, true}.has(1, 2, 3, 4); + if(inputs.front().symbolic()) + MIGRAPHX_THROW("SLICE: symbolic input shapes are not supported, use dyn_slice"); if(inputs.size() != 1) return compute_two_or_more(inputs); @@ -271,17 +274,10 @@ struct slice 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(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. @@ -301,13 +297,8 @@ struct slice auto dds = input_shape.dyn_dims(); for(auto axis : this->axes) { - dds[axis] = input_shape.symbolic() - ? shape::dynamic_dimension{sym::lit(new_lens[axis])} - : shape::dynamic_dimension{new_lens[axis], new_lens[axis]}; + dds[axis] = {new_lens[axis], new_lens[axis]}; } - - if(input_shape.symbolic()) - return shape{input_shape.type(), dds, input_shape.dyn_strides()}; return shape{input_shape.type(), dds}; } diff --git a/src/include/migraphx/operators.hpp b/src/include/migraphx/operators.hpp index 1e527c762ff..78b0d28b088 100644 --- a/src/include/migraphx/operators.hpp +++ b/src/include/migraphx/operators.hpp @@ -53,6 +53,7 @@ #include #include #include +#include #include #include #include diff --git a/src/normalize_attributes.cpp b/src/normalize_attributes.cpp index 48804c9034f..c8709b25f94 100644 --- a/src/normalize_attributes.cpp +++ b/src/normalize_attributes.cpp @@ -25,11 +25,85 @@ #include #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. +// 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(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(dd.get_interval().max); + MIGRAPHX_THROW("normalize_attributes: cannot normalize against 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 values that are symbolic or that are normalized against + * a symbolic axis length. Applies the ONNX clamp norm(v) = clamp(v < 0 ? v + D : v, 0, D) + * symbolically, folding against the interval bounds where provable. + * + * Returns the entries as dim_like so a value that folds to a literal is stored as a plain + * integer. + */ +template +static std::vector tune_attribute_sym(const std::vector& exprs, + 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() + "use_sym normalization requires use_len!"); + if(axes.size() != exprs.size()) + MIGRAPHX_THROW(m() + "one axis per value is required to normalize symbolically!"); + auto zero = sym::lit(std::int64_t{0}); + std::vector result(exprs.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() + "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 result; +} + /** * Parameters: * vec: the vector attribute to normalize @@ -64,6 +138,9 @@ static auto tune_attribute(const std::vector& vec, if(contains(vec_attrs, op::normalize_attribute::use_len)) { + // max_vals has one entry per value, and each axis picks the entry to fill. + if(axes.size() > vec.size()) + MIGRAPHX_THROW(m() + "more axes than values to normalize!"); if(input_shape.dynamic()) { // return the unchanged `vec` if the dynamic_dimensions at `axes` are not fixed @@ -230,14 +307,31 @@ bool normalize_attributes(operation& op, const shape& input_shape) auto vv = val.at(key).without_key(); if(vv.is_array()) { + auto attrs = rv.without_key().to_vector(); + // A symbolic value serializes as an object. Normalizing against a symbolic axis + // length can turn a concrete value symbolic too, so a symbolic input shape takes + // the same path even when every value is a plain integer. + bool sym_values = + std::any_of(vv.begin(), vv.end(), [](const auto& e) { return e.is_object(); }); + bool allow_sym = contains(attrs, op::normalize_attribute::use_sym); + if(sym_values and not allow_sym) + MIGRAPHX_THROW(message() + "symbolic values are not supported!"); std::vector axes; if(val.contains("axes")) { axes = val.at("axes").without_key().to_vector(); } - auto vec = vv.to_vector(); - auto result = tune_attribute(vec, axes, rv.without_key(), input_shape, message); - val[key] = result; + if(not vv.empty() and allow_sym and (sym_values or input_shape.symbolic())) + { + auto dims = migraphx::from_value>(vv); + val[key] = migraphx::to_value( + tune_attribute_sym(to_sym_exprs(dims), axes, attrs, input_shape, message)); + } + else + { + auto vec = vv.to_vector(); + val[key] = tune_attribute(vec, axes, rv.without_key(), input_shape, message); + } op.from_value(val); val = op.to_value(); tuned = true; diff --git a/src/sym.cpp b/src/sym.cpp index bf299ad3456..380a3014d8b 100644 --- a/src/sym.cpp +++ b/src/sym.cpp @@ -969,6 +969,14 @@ static const std::vector& get_rewrite_rules() sqrt(_1 / _2) >> sqrt(_1) / sqrt(_2), log(exp(_1)) >> _1, exp(log(_1)) >> _1, + // Clamping against a bound the expression already clamps to only nests a + // redundant node, so repeated clamping (as attribute normalization does on + // every shape computation) keeps a single min/max instead of growing without + // bound. The inner node can hold the bound in either operand. + min(min(_1, _2), _2) >> min(_1, _2), + min(min(_2, _1), _2) >> min(_2, _1), + max(max(_1, _2), _2) >> max(_1, _2), + max(max(_2, _1), _2) >> max(_2, _1), }; }(); return rules; diff --git a/src/targets/gpu/lowering.cpp b/src/targets/gpu/lowering.cpp index a0a6dc67afa..e9e7d7c2b88 100644 --- a/src/targets/gpu/lowering.cpp +++ b/src/targets/gpu/lowering.cpp @@ -693,7 +693,7 @@ struct miopen_apply void add_dyn_slice_op() { - apply_map.emplace("slice", [=](instruction_ref ins) { + auto lower_runtime_bounds = [=](instruction_ref ins) { auto inputs = ins->inputs(); if(inputs.size() > 1) { @@ -715,7 +715,9 @@ struct miopen_apply ins, mod->insert_instruction(ins, ins->get_operator(), inputs)); } return ins; - }); + }; + apply_map.emplace("slice", lower_runtime_bounds); + apply_map.emplace("dyn_slice", lower_runtime_bounds); } // Get the argument's shape dimensions on host and then copy to gpu diff --git a/test/gpu/dyn_slice_lowering.cpp b/test/gpu/dyn_slice_lowering.cpp index 11d743012af..109a273eaeb 100644 --- a/test/gpu/dyn_slice_lowering.cpp +++ b/test/gpu/dyn_slice_lowering.cpp @@ -71,6 +71,38 @@ TEST_CASE(dyn_slice_lowering_runtime_inputs) EXPECT(m1 == m2); } +// dyn_slice always has runtime bound inputs, so both of them are copied to the host. +TEST_CASE(dyn_slice_lowering_dyn_slice_op) +{ + migraphx::shape data_s{migraphx::shape::float_type, {2, 2, 4}}; + migraphx::shape idx_s{migraphx::shape::int64_type, {1}}; + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {2}}}); + + migraphx::module m1; + { + 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(op, data, starts, ends); + m1.add_return({sl}); + } + run_lowering(m1); + + migraphx::module m2; + { + auto data = m2.add_parameter("data", data_s); + auto starts = m2.add_parameter("starts", idx_s); + auto ends = m2.add_parameter("ends", idx_s); + auto copy_starts = m2.add_instruction(migraphx::make_op("hip::copy_from_gpu"), starts); + 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(op, data, sync, copy_ends); + m2.add_return({sl}); + } + EXPECT(m1 == m2); +} + // A slice with only 1 input (all attributes inline) should not be modified // by the dynamic slice lowering. TEST_CASE(dyn_slice_lowering_single_input) diff --git a/test/normalize_ops_test.cpp b/test/normalize_ops_test.cpp index dca9f78125a..dfe1a9bef08 100644 --- a/test/normalize_ops_test.cpp +++ b/test/normalize_ops_test.cpp @@ -26,10 +26,17 @@ #include #include #include +#include #include +#include +#include #include #include +using dd = migraphx::shape::dynamic_dimension; +using migraphx::sym::lit; +using migraphx::sym::var; + struct normalize_test_op { std::vector axes = {}; @@ -61,6 +68,37 @@ struct normalize_test_op } }; +// A bound attribute that can hold a symbolic value has to declare use_sym, otherwise its +// normalized value could not be stored back. This operator deliberately leaves it out. +struct no_use_sym_test_op +{ + std::vector axes = {}; + std::vector bound = {}; + + template + static auto reflect(Self& self, F f) + { + return migraphx::pack(f(self.axes, "axes"), f(self.bound, "bound")); + } + + migraphx::value attributes() const + { + migraphx::value normalize; + normalize["bound"] = migraphx::value::array{migraphx::op::normalize_attribute::clip_max, + migraphx::op::normalize_attribute::clip_min, + migraphx::op::normalize_attribute::include_max, + migraphx::op::normalize_attribute::use_len, + migraphx::op::normalize_attribute::include_min}; + return {{"normalize_axes", normalize}}; + } + + std::string name() const { return "normalize_ops_test::no_use_sym_op"; } + migraphx::shape normalize_compute_shape(std::vector inputs) const + { + return inputs[0]; + } +}; + static void run_pass(migraphx::module& m) { migraphx::run_passes(m, {migraphx::normalize_ops{}, migraphx::dead_code_elimination{}}); @@ -179,6 +217,190 @@ TEST_CASE(slice_test_1) EXPECT(m1 == m2); } +// dyn_slice always takes its bounds as inputs, and its starts/ends attributes describe those +// inputs at compile time. Only the attributes are normalized here. +static migraphx::module create_dyn_slice(const migraphx::shape& data_shape, + const migraphx::value& attributes, + std::size_t nbounds = 1) +{ + migraphx::module m; + migraphx::shape bounds_shape{migraphx::shape::int64_type, {nbounds}}; + auto data = m.add_parameter("data", data_shape); + auto starts = m.add_parameter("starts", bounds_shape); + auto ends = m.add_parameter("ends", bounds_shape); + auto r = m.add_instruction(migraphx::make_op("dyn_slice", attributes), data, starts, ends); + m.add_return({r}); + + return m; +} + +static migraphx::value sym_bound(const migraphx::sym::expr& e) +{ + return migraphx::value::array{migraphx::to_value(dd{e})}; +} + +TEST_CASE(dyn_slice_sym_ends_clamped_test) +{ + // n is not provably ordered against the axis length 5, so the bound clamps to min(n, 5). + auto n = var("n", {1, 8}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + auto m1 = create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(n)}}); + auto m2 = create_dyn_slice( + s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(migraphx::sym::min(n, lit(5)))}}); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_sym_ends_below_len_test) +{ + // n < 5 is provable, so the bound keeps the bare symbol rather than gaining a min wrapper. + auto n = var("n", {1, 4}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + auto m1 = create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(n)}}); + auto m2 = create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(n)}}); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_sym_ends_at_len_test) +{ + // n >= 5 is provable, so the bound collapses to the axis length and demotes to an integer. + auto n = var("n", {6, 9}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + auto m1 = create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(n)}}); + auto m2 = create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", {5}}}); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_sym_starts_clamped_test) +{ + // The starts attribute normalizes the same way. Axis 1 has length 3. + auto n = var("n", {0, 5}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + auto m1 = create_dyn_slice(s, {{"axes", {1}}, {"starts", sym_bound(n)}, {"ends", {3}}}); + auto m2 = create_dyn_slice( + s, {{"axes", {1}}, {"starts", sym_bound(migraphx::sym::min(n, lit(3)))}, {"ends", {3}}}); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_sym_mixed_bounds_test) +{ + // Each entry is clamped against its own axis length (3 for axis 1, 5 for axis 3), and the + // concrete entry stays concrete. + auto n = var("n", {1, 8}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + auto m1 = create_dyn_slice(s, + {{"axes", {1, 3}}, + {"starts", {0, 0}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n}), 9}}}, + 2); + auto m2 = create_dyn_slice( + s, + {{"axes", {1, 3}}, + {"starts", {0, 0}}, + {"ends", + migraphx::value::array{migraphx::to_value(dd{migraphx::sym::min(n, lit(3))}), 5}}}, + 2); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_sym_symbolic_axis_len_test) +{ + // When the sliced axis is itself symbolic, the clamp bound is that axis's symbol instead of + // a compile-time length. + auto k = var("k", {2, 6}); + auto n = var("n", {1, 8}); + migraphx::shape s{migraphx::shape::float_type, {dd{k}, dd{lit(4)}}}; + + auto m1 = create_dyn_slice(s, {{"axes", {0}}, {"starts", {0}}, {"ends", sym_bound(n)}}); + auto m2 = create_dyn_slice( + s, {{"axes", {0}}, {"starts", {0}}, {"ends", sym_bound(migraphx::sym::min(n, k))}}); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_concrete_bounds_symbolic_axis_len_test) +{ + // A concrete bound still normalizes symbolically when the axis it is clamped against is a + // symbol: end 2 is not provably below k, so it becomes min(2, k). + auto k = var("k", {1, 6}); + migraphx::shape s{migraphx::shape::float_type, {dd{k}, dd{lit(4)}}}; + + auto m1 = create_dyn_slice(s, {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}); + auto m2 = create_dyn_slice( + s, {{"axes", {0}}, {"starts", {0}}, {"ends", sym_bound(migraphx::sym::min(lit(2), k))}}); + run_pass(m1); + + EXPECT(m1 == m2); +} + +TEST_CASE(dyn_slice_sym_normalize_idempotent_test) +{ + // Normalizing an already normalized symbolic bound must not nest a second clamp. + auto n = var("n", {1, 8}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + auto m1 = create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(n)}}); + run_pass(m1); + auto once = m1; + run_pass(m1); + + EXPECT(m1 == once); +} + +TEST_CASE(dyn_slice_sym_indeterminate_sign_throws) +{ + // A bound that may be negative cannot be resolved into a from-the-end index. + auto n = var("n", {-2, 2}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + EXPECT(test::throws( + [&] { create_dyn_slice(s, {{"axes", {3}}, {"starts", {0}}, {"ends", sym_bound(n)}}); }, + "bound of indeterminate sign cannot be normalized")); +} + +TEST_CASE(dyn_slice_sym_nonfixed_axis_throws) +{ + // A range-based dynamic axis is neither symbolic nor fixed, so it has no length expression + // to clamp the bound against. + auto n = var("n", {1, 8}); + migraphx::shape s{migraphx::shape::float_type, {{2, 4}, {3, 3}}}; + + EXPECT(test::throws( + [&] { create_dyn_slice(s, {{"axes", {0}}, {"starts", {0}}, {"ends", sym_bound(n)}}); }, + "cannot normalize against a non-fixed axis")); +} + +TEST_CASE(sym_value_without_use_sym_throws) +{ + // Normalizing a symbolic value into an attribute that did not opt in is rejected instead of + // silently leaving the bound unnormalized. + auto n = var("n", {1, 8}); + migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 5}}; + + EXPECT(test::throws( + [&] { + migraphx::module m; + auto data = m.add_parameter("data", s); + m.add_instruction(no_use_sym_test_op{{3}, {dd{n}}}, data); + }, + "symbolic values are not supported")); +} + static migraphx::module create_test_op(const std::vector& axes) { migraphx::module m; diff --git a/test/op_shape_test.cpp b/test/op_shape_test.cpp index ae60d28d96c..d0f351d97ee 100644 --- a/test/op_shape_test.cpp +++ b/test/op_shape_test.cpp @@ -1408,6 +1408,243 @@ TEST_CASE(dot_sym_k_vs_range) expect_shape(expected, migraphx::make_op("dot"), s_a, s_b); } +// dyn_slice takes its bounds as inputs; this is the shape they have for `n` sliced axes. +static migraphx::shape dyn_slice_bounds(std::size_t n) +{ + return migraphx::shape{migraphx::shape::int64_type, {n}}; +} + +TEST_CASE(dyn_slice_static) +{ + // Concrete bounds over a static input stay static: a slice is just a view. + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape expected{migraphx::shape::float_type, {2, 2, 2}, {6, 3, 1}}; + expect_shape(expected, + migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}), + input, + dyn_slice_bounds(1), + dyn_slice_bounds(1)); +} + +TEST_CASE(dyn_slice_static_clamped_bounds) +{ + // Out of range bounds are clipped and negative bounds resolved against the axis length. + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape expected{migraphx::shape::float_type, {2, 2, 2}, {6, 3, 1}}; + expect_shape(expected, + migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {-2}}, {"ends", {10}}}), + input, + dyn_slice_bounds(1), + dyn_slice_bounds(1)); +} + +TEST_CASE(dyn_slice_negative_axis) +{ + // A negative axis attribute is normalized against the input rank. + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape expected{migraphx::shape::float_type, {2, 2, 2}, {6, 3, 1}}; + expect_shape(expected, + migraphx::make_op("dyn_slice", {{"axes", {-1}}, {"starts", {1}}, {"ends", {3}}}), + input, + dyn_slice_bounds(1), + dyn_slice_bounds(1)); +} + +TEST_CASE(dyn_slice_symbolic_end_static_input) +{ + // Static input + symbolic end bound: the output is symbolic, so it is not demoted to static. + auto n = var("n", {1, 16}); + auto op = migraphx::make_op("dyn_slice", + {{"axes", {0}}, + {"starts", {0}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}}); + + // end=n is clamped to the axis length 10: dim = min(n, 10). + migraphx::shape sin{migraphx::shape::float_type, {10}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{migraphx::sym::min(n, lit(10))}}, {lit(1)}}; + expect_shape(sout, op, sin, dyn_slice_bounds(1), dyn_slice_bounds(1)); + 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(dyn_slice_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. + auto bounds = dyn_slice_bounds(1); + { + // 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("dyn_slice", + {{"axes", {1}}, + {"starts", {2}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{m}, dd{lit(12)}}}; + 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, bounds, bounds); + 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("dyn_slice", + {{"axes", {0}}, + {"starts", {2}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{n})}}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{lit(12)}, dd{lit(4)}}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{lit(10)}, dd{lit(4)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, bounds, bounds); + } + { + // Symbolic start: dim = 8 - n (n <= 8 keeps the extent non-negative). + auto n = var("n", {1, 8}); + auto op = migraphx::make_op("dyn_slice", + {{"axes", {0}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{n})}}, + {"ends", {8}}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{lit(10)}, dd{lit(4)}}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{lit(8) - n}, dd{lit(4)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, bounds, bounds); + EXPECT(sout.symbolic()); + } +} + +TEST_CASE(dyn_slice_sym_data_fixed_axis) +{ + // Symbolic input sliced on a fixed axis: the symbol on the other axis is untouched and the + // result matches slicing the equivalent static shape. + auto n = var("n", {1, 8}); + std::unordered_map sym_map = {{n, 5}}; + auto bounds = dyn_slice_bounds(1); + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); + + migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(3)}}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(2)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, bounds, bounds); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map), bounds, bounds})); +} + +TEST_CASE(dyn_slice_sym_data_symbolic_axis) +{ + // Slicing the non-fixed symbolic axis: the concrete end is clamped against the axis symbol + // rather than a compile-time length, so the extent stays symbolic. + auto n = var("n", {1, 8}); + auto m = var("m", {1, 8}); + auto bounds = dyn_slice_bounds(1); + auto op = migraphx::make_op("dyn_slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}); + + migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{m}}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{migraphx::sym::min(lit(2), n)}, dd{m}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, bounds, bounds); + EXPECT(sout.to_static({{n, 5}, {m, 3}}) == + migraphx::shape{migraphx::shape::float_type, {2, 3}, {3, 1}}); + EXPECT(sout.to_static({{n, 1}, {m, 3}}) == + migraphx::shape{migraphx::shape::float_type, {1, 3}, {3, 1}}); +} + +TEST_CASE(dyn_slice_sym_multiple_axes) +{ + // Slice two axes at once; the symbol at the untouched axis survives. + auto n = var("n", {1, 8}); + std::unordered_map sym_map = {{n, 4}}; + auto bounds = dyn_slice_bounds(2); + auto op = + migraphx::make_op("dyn_slice", {{"axes", {0, 2}}, {"starts", {1, 2}}, {"ends", {4, 5}}}); + + migraphx::shape sin{migraphx::shape::float_type, {dd{lit(6)}, dd{n}, dd{lit(8)}}}; + migraphx::shape sout{ + migraphx::shape::float_type, {dd{lit(3)}, dd{n}, dd{lit(3)}}, sin.dyn_strides()}; + expect_shape(sout, op, sin, bounds, bounds); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map), bounds, bounds})); +} + +TEST_CASE(dyn_slice_sym_nonstandard_layout) +{ + // Non-standard symbolic input: the slice must preserve the permutation. + auto n = var("n", {1, 8}); + std::unordered_map sym_map = {{n, 6}}; + auto bounds = dyn_slice_bounds(1); + + auto sin = migraphx::shape::from_permutation( + migraphx::shape::float_type, {dd{n}, dd{lit(3)}, dd{lit(5)}, dd{lit(7)}}, {0, 2, 3, 1}); + auto op = migraphx::make_op("dyn_slice", {{"axes", {3}}, {"starts", {1}}, {"ends", {6}}}); + auto sout = op.compute_shape({sin, bounds, bounds}); + EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map), bounds, bounds})); +} + +TEST_CASE(dyn_slice_wrong_number_of_inputs_error) +{ + // The bounds are the only inputs: the axes are an attribute, not a fourth input. + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + auto bounds = dyn_slice_bounds(1); + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); + throws_shape(op, input, bounds); + throws_shape(op, input, bounds, bounds, bounds); +} + +TEST_CASE(dyn_slice_bounds_input_rank_error) +{ + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape bounds_2d{migraphx::shape::int64_type, {1, 1}}; + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); + throws_shape(op, input, bounds_2d, bounds_2d); +} + +TEST_CASE(dyn_slice_dynamic_bounds_input_error) +{ + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape dyn_bounds{migraphx::shape::int64_type, {dd{1, 4}}}; + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); + throws_shape(op, input, dyn_bounds, dyn_bounds); +} + +TEST_CASE(dyn_slice_bounds_input_length_error) +{ + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); + throws_shape(op, input, dyn_slice_bounds(2), dyn_slice_bounds(2)); +} + +TEST_CASE(dyn_slice_attribute_length_error) +{ + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + auto bounds = dyn_slice_bounds(2); + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {0, 1}}, {"ends", {2, 3}}}); + throws_shape(op, input, bounds, bounds); +} + +TEST_CASE(dyn_slice_missing_attribute_error) +{ + migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + auto bounds = dyn_slice_bounds(1); + throws_shape( + migraphx::make_op("dyn_slice", {{"starts", {1}}, {"ends", {3}}}), input, bounds, bounds); +} + +TEST_CASE(dyn_slice_range_dynamic_data_error) +{ + // A range-based dynamic dimension has no expression to build a symbolic extent from. + migraphx::shape input{migraphx::shape::float_type, {{2, 4}, {3, 3}}}; + auto bounds = dyn_slice_bounds(1); + throws_shape(migraphx::make_op("dyn_slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {2}}}), + input, + bounds, + bounds); +} + TEST_CASE(broadcast_for_dot_static) { migraphx::shape s0{migraphx::shape::float_type, {481, 356}}; @@ -5585,96 +5822,39 @@ TEST_CASE(slice_dyn_nonfixed_keeps_other_optimals) input); } -TEST_CASE(slice_sym) -{ - auto n = var("n", {1, 8}); - auto m = var("m", {1, 16}); - auto k = var("k", {1, 64}); - std::unordered_map sym_map = {{n, 3}, {m, 5}, {k, 7}}; - - auto expect_matches_static = [&](const migraphx::operation& op, - const migraphx::shape& sin, - const migraphx::shape& sym_out) { - EXPECT(sym_out.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); - }; - - { - // Slice axis 0 (first); sym at axis 1. - auto op = migraphx::make_op("slice", {{"axes", {0}}, {"starts", {1}}, {"ends", {3}}}); - migraphx::shape sin{migraphx::shape::float_type, {dd{lit(5)}, dd{n}, dd{lit(4)}}}; - migraphx::shape sout{ - migraphx::shape::float_type, {dd{lit(2)}, dd{n}, dd{lit(4)}}, sin.dyn_strides()}; - expect_shape(sout, op, sin); - expect_matches_static(op, sin, sout); - } - { - // Slice axis 1 (middle); syms at axes 0 and 2. - auto op = migraphx::make_op("slice", {{"axes", {1}}, {"starts", {2}}, {"ends", {6}}}); - migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{lit(8)}, dd{m}}}; - migraphx::shape sout{ - migraphx::shape::float_type, {dd{n}, dd{lit(4)}, dd{m}}, sin.dyn_strides()}; - expect_shape(sout, op, sin); - expect_matches_static(op, sin, sout); - } - { - // Slice axis 3 (last) on a 4D shape; syms at axes 0, 1, 2. - auto op = migraphx::make_op("slice", {{"axes", {3}}, {"starts", {0}}, {"ends", {3}}}); - migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{m}, dd{k}, dd{lit(10)}}}; - migraphx::shape sout{ - migraphx::shape::float_type, {dd{n}, dd{m}, dd{k}, dd{lit(3)}}, sin.dyn_strides()}; - expect_shape(sout, op, sin); - expect_matches_static(op, sin, sout); - } -} - -TEST_CASE(slice_sym_multiple_axes) -{ - // Slice axes 0 and 2 at once; sym at axis 1 is untouched. - auto n = var("n", {1, 8}); - std::unordered_map sym_map = {{n, 4}}; - - auto op = migraphx::make_op("slice", {{"axes", {0, 2}}, {"starts", {1, 2}}, {"ends", {4, 5}}}); - migraphx::shape sin{migraphx::shape::float_type, {dd{lit(6)}, dd{n}, dd{lit(8)}}}; - migraphx::shape sout{ - migraphx::shape::float_type, {dd{lit(3)}, dd{n}, 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_fixed_bound_var) -{ - // var("k", {3, 3}) is fixed (collapsed bound), so slicing the axis is allowed. - auto k = var("k", {3, 3}); - auto n = var("n", {1, 8}); - std::unordered_map sym_map = {{n, 5}}; - - auto op = migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {2}}}); - migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{k}, dd{lit(4)}}}; - migraphx::shape sout{ - migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(4)}}, 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_non_fixed_throws) -{ - // 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); -} - -TEST_CASE(slice_sym_nonstandard_layout) -{ - // Non-standard symbolic input: the slice must preserve the permutation - auto n = var("n", {1, 8}); - std::unordered_map sym_map = {{n, 6}}; - - auto sin = migraphx::shape::from_permutation( - migraphx::shape::float_type, {dd{n}, dd{lit(3)}, dd{lit(5)}, dd{lit(7)}}, {0, 2, 3, 1}); - auto op = migraphx::make_op("slice", {{"axes", {3}}, {"starts", {1}}, {"ends", {6}}}); - auto sout = op.compute_shape({sin}); - EXPECT(sout.to_static(sym_map) == op.compute_shape({sin.to_static(sym_map)})); +TEST_CASE(slice_sym_throws) +{ + // slice has no way to express a symbolic output extent, so a symbolic input is rejected + // whatever the sliced axis looks like. dyn_slice handles these instead. + auto n = var("n", {1, 8}); + auto k = var("k", {3, 3}); + auto op_axis0 = migraphx::make_op("slice", {{"axes", {0}}, {"starts", {1}}, {"ends", {3}}}); + auto op_axis1 = migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {2}}}); + + // Fixed sliced axis, symbol on another axis. + throws_shape(op_axis0, + migraphx::shape{migraphx::shape::float_type, {dd{lit(5)}, dd{n}, dd{lit(4)}}}); + // Non-fixed symbolic sliced axis. + throws_shape(op_axis1, + migraphx::shape{migraphx::shape::float_type, {dd{lit(4)}, dd{n}, dd{lit(8)}}}); + // A collapsed-bound symbol is fixed, but the shape is still symbolic. + throws_shape(op_axis1, + migraphx::shape{migraphx::shape::float_type, {dd{n}, dd{k}, dd{lit(4)}}}); + // Multiple sliced axes. + throws_shape( + migraphx::make_op("slice", {{"axes", {0, 2}}, {"starts", {1, 2}}, {"ends", {4, 5}}}), + migraphx::shape{migraphx::shape::float_type, {dd{lit(6)}, dd{n}, dd{lit(8)}}}); + // Non-standard layout. + throws_shape(migraphx::make_op("slice", {{"axes", {3}}, {"starts", {1}}, {"ends", {6}}}), + migraphx::shape::from_permutation(migraphx::shape::float_type, + {dd{n}, dd{lit(3)}, dd{lit(5)}, dd{lit(7)}}, + {0, 2, 3, 1})); + // Variable bounds over symbolic data, which used to degrade to a range. + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; + throws_shape(migraphx::make_op("slice", {{"axes", {0}}}), + migraphx::shape{migraphx::shape::float_type, {dd{n}, dd{lit(4)}}}, + bounds, + bounds); } TEST_CASE(test_scan_slice1) diff --git a/test/ref/dyn_slice.cpp b/test/ref/dyn_slice.cpp new file mode 100644 index 00000000000..f38154ff8fa --- /dev/null +++ b/test/ref/dyn_slice.cpp @@ -0,0 +1,279 @@ +/* + * The MIT License (MIT) + * + * 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 + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using dd = migraphx::shape::dynamic_dimension; +using migraphx::sym::var; + +// The bounds are always inputs, so every test supplies them as parameters named starts/ends +// with this shape. +static migraphx::shape dyn_slice_bounds_shape(std::size_t n) +{ + return migraphx::shape{migraphx::shape::int64_type, {n}}; +} + +static void add_dyn_slice(migraphx::module& m, + const migraphx::value& attributes, + migraphx::instruction_ref data, + std::size_t nbounds) +{ + auto bounds_shape = dyn_slice_bounds_shape(nbounds); + auto starts = m.add_parameter("starts", bounds_shape); + auto ends = m.add_parameter("ends", bounds_shape); + m.add_instruction(migraphx::make_op("dyn_slice", attributes), data, starts, ends); +} + +static migraphx::parameter_map bounds_params(const std::vector& starts, + const std::vector& ends) +{ + auto bounds_shape = dyn_slice_bounds_shape(starts.size()); + migraphx::parameter_map params; + // Going through a literal so the arguments own their buffers. + params["starts"] = migraphx::literal{bounds_shape, starts}.get_argument(); + params["ends"] = migraphx::literal{bounds_shape, ends}.get_argument(); + return params; +} + +static migraphx::literal iota_literal(const migraphx::shape& s) +{ + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + return migraphx::literal{s, data}; +} + +static std::vector read_ints(const migraphx::argument& arg) +{ + std::vector result; + arg.visit([&](auto output) { result.assign(output.begin(), output.end()); }); + return result; +} + +TEST_CASE(dyn_slice_concrete_bounds_test) +{ + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}, data, 1); + // Every bound is concrete, so the output shape is known at compile time. + EXPECT(p.get_output_shapes().back() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); + p.compile(migraphx::make_target("ref")); + + auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + // The static output shape lets the compiler make the aliased view contiguous. + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {4, 2, 1}}); +} + +TEST_CASE(dyn_slice_sym_ends_test) +{ + // Symbolic `ends` bound with a variable input supplying its runtime value. The same + // compiled program handles both values of the symbol. + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, + {{"axes", {2}}, + {"starts", {1}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 3})})}}}, + data, + 1); + p.compile(migraphx::make_target("ref")); + + auto result0 = p.eval(bounds_params({1}, {3})).back(); + std::vector gold0 = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result0), gold0)); + EXPECT(result0.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); + + auto result1 = p.eval(bounds_params({1}, {2})).back(); + std::vector gold1 = {1, 4, 7, 10}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result1), gold1)); + EXPECT(result1.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 1}, {6, 3, 1}}); +} + +TEST_CASE(dyn_slice_sym_starts_test) +{ + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 2})})}}, + {"ends", {3}}}, + data, + 1); + p.compile(migraphx::make_target("ref")); + + auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); +} + +TEST_CASE(dyn_slice_sym_both_bounds_test) +{ + // Both bounds symbolic, each with its own runtime value. + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 1})})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 3})})}}}, + data, + 1); + p.compile(migraphx::make_target("ref")); + + auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); +} + +TEST_CASE(dyn_slice_sym_data_test) +{ + // Symbolic input shape sliced on a fixed axis: the output is symbolic after compiling and + // resolves once the parameter is bound to a static shape. + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, + {dd{var("n", {1, 4})}, dd{migraphx::sym::lit(2)}, dd{migraphx::sym::lit(3)}}}; + auto data = mm->add_parameter("x", s); + add_dyn_slice(*mm, {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}, data, 1); + p.compile(migraphx::make_target("ref")); + + migraphx::shape input_fixed_shape{migraphx::shape::int32_type, {2, 2, 3}}; + auto data_literal = iota_literal(input_fixed_shape); + auto params = bounds_params({1}, {3}); + params["x"] = data_literal.get_argument(); + + auto result = p.eval(params).back(); + std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); +} + +TEST_CASE(dyn_slice_sym_bounds_multi_axes_test) +{ + // Two symbolic end bounds over two axes, so each bound is clamped against its own axis + // length and more than one axis is sliced. + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, + {{"axes", {1, 2}}, + {"starts", {1, 0}}, + {"ends", + migraphx::value::array{migraphx::to_value(dd{var("n", {1, 2})}), + migraphx::to_value(dd{var("m", {1, 3})})}}}, + data, + 2); + p.compile(migraphx::make_target("ref")); + + auto result = p.eval(bounds_params({1, 0}, {2, 2})).back(); + std::vector gold = {3, 4, 9, 10}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 1, 2}, {6, 3, 1}}); +} + +TEST_CASE(dyn_slice_runtime_bounds_clamped_test) +{ + // A runtime end that is out of range and a negative runtime start are both resolved + // against the axis length when the slice runs. + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 1})})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 8})})}}}, + data, + 1); + p.compile(migraphx::make_target("ref")); + + auto result = p.eval(bounds_params({-2}, {100})).back(); + std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); +} + +TEST_CASE(dyn_slice_negative_axis_test) +{ + // The axes attribute is normalized when the program is compiled, so the runtime bounds are + // applied to axis 2. + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, {{"axes", {-1}}, {"starts", {1}}, {"ends", {3}}}, data, 1); + p.compile(migraphx::make_target("ref")); + + auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; + EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(result.get_shape() == + migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {4, 2, 1}}); +} + +TEST_CASE(dyn_slice_end_before_start_error_test) +{ + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; + auto data = mm->add_literal(iota_literal(s)); + add_dyn_slice(*mm, + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 3})})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {0, 3})})}}}, + data, + 1); + p.compile(migraphx::make_target("ref")); + + EXPECT(test::throws([&] { p.eval(bounds_params({2}, {1})); })); +} diff --git a/test/sym.cpp b/test/sym.cpp index e435b093ab8..f02861f7d5c 100644 --- a/test/sym.cpp +++ b/test/sym.cpp @@ -2647,6 +2647,64 @@ TEST_CASE(builtin_log_exp_nested) EXPECT(log(exp(x)) + log(exp(y)) == x + y); } +TEST_CASE(builtin_min_max_already_clamped) +{ + auto x = var("x"); + auto y = var("y"); + // Clamping against a bound that is already applied folds away, with the bound in either + // operand of the inner node. + EXPECT(min(min(x, y), y) == min(x, y)); + EXPECT(max(max(x, y), y) == max(x, y)); + EXPECT(min(min(y, x), y) == min(y, x)); + EXPECT(max(max(y, x), y) == max(y, x)); +} + +TEST_CASE(builtin_min_max_already_clamped_literal_bound) +{ + auto x = var("x"); + EXPECT(min(min(x, lit(5)), lit(5)) == min(x, lit(5))); + EXPECT(max(max(x, lit(0)), lit(0)) == max(x, lit(0))); + EXPECT(min(min(lit(5), x), lit(5)) == min(lit(5), x)); + EXPECT(max(max(lit(0), x), lit(0)) == max(lit(0), x)); +} + +TEST_CASE(builtin_min_max_clamp_repeated) +{ + // Clamping a clamped expression any number of times keeps a single min/max node, as + // repeated attribute normalization does. + auto n = var("n", interval{int64_t{1}, int64_t{8}}); + auto once = min(n, lit(5)); + EXPECT(min(once, lit(5)) == once); + EXPECT(min(min(once, lit(5)), lit(5)) == once); + auto once_from = max(n, lit(0)); + EXPECT(max(once_from, lit(0)) == once_from); + EXPECT(max(max(once_from, lit(0)), lit(0)) == once_from); +} + +TEST_CASE(builtin_min_max_different_bound_not_folded) +{ + auto x = var("x"); + auto y = var("y"); + auto z = var("z"); + // The inner bound is not the outer bound, so the nesting is meaningful and kept + EXPECT(min(min(x, y), z) != min(x, y)); + EXPECT(min(min(x, y), z).children().front() == min(x, y)); + EXPECT(max(max(x, y), z) != max(x, y)); + EXPECT(max(max(x, y), z).children().front() == max(x, y)); +} + +TEST_CASE(builtin_min_max_already_clamped_eval) +{ + auto x = var("x"); + auto y = var("y"); + auto e = min(min(x, y), y); + EXPECT(e.eval({{var("x"), int64_t{7}}, {var("y"), int64_t{5}}}) == scalar{int64_t{5}}); + EXPECT(e.eval({{var("x"), int64_t{3}}, {var("y"), int64_t{5}}}) == scalar{int64_t{3}}); + auto f = max(max(x, y), y); + EXPECT(f.eval({{var("x"), int64_t{7}}, {var("y"), int64_t{5}}}) == scalar{int64_t{7}}); + EXPECT(f.eval({{var("x"), int64_t{3}}, {var("y"), int64_t{5}}}) == scalar{int64_t{5}}); +} + TEST_CASE(builtin_raw_no_leak) { auto x = var("x"); From 37ce9660d6fec45c55c656b7fdb39c7ba1dc4fe9 Mon Sep 17 00:00:00 2001 From: charlie Date: Wed, 5 Aug 2026 14:25:36 -0500 Subject: [PATCH 2/6] Test and comment cleanup --- src/include/migraphx/op/dyn_slice.hpp | 10 +- test/op_shape_test.cpp | 55 ++--- test/ref/dyn_slice.cpp | 333 +++++++++++++++++--------- 3 files changed, 249 insertions(+), 149 deletions(-) diff --git a/src/include/migraphx/op/dyn_slice.hpp b/src/include/migraphx/op/dyn_slice.hpp index 47842fc97c4..f1d25c60539 100644 --- a/src/include/migraphx/op/dyn_slice.hpp +++ b/src/include/migraphx/op/dyn_slice.hpp @@ -118,7 +118,7 @@ struct dyn_slice { check_shapes{inputs, *this, true}.has(3); check_inputs_and_attributes(inputs); - auto input_shape = inputs.front(); + const auto& input_shape = inputs.front(); if(input_shape.dynamic() and not input_shape.symbolic()) MIGRAPHX_THROW("DYN_SLICE: data input must have a static or symbolic shape"); @@ -151,18 +151,14 @@ struct dyn_slice arg.visit([&](auto values) { result = values.template to_vector(); }); return result; }; - // The bound attributes are only the compile-time view of the inputs; the inputs hold the - // values to slice with, so they are renormalized against the run-time shape here. The - // axes attribute needs no such handling: shapes with a dynamic rank are not supported, so - // it is already normalized at compile time. auto axes_attrs = this->attributes().at("normalize_axes"); + // Only use the starts_input and ends_input for the output slice. Not the attributes. auto norm_starts = normalize_indices( read(args[1]), axes, input_shape, axes_attrs.at("starts"), "DYN_SLICE: starts input"); auto norm_ends = normalize_indices( read(args[2]), axes, input_shape, axes_attrs.at("ends"), "DYN_SLICE: ends input"); - // The compile-time shape asserts a non-negative extent on every sliced axis, so - // inconsistent run-time bounds are rejected instead of wrapping around to a huge length. + // Get end-start for output dimension sizes. Reject if ends before starts (no wrap around). std::vector extents(axes.size()); std::transform(norm_ends.begin(), norm_ends.end(), diff --git a/test/op_shape_test.cpp b/test/op_shape_test.cpp index d0f351d97ee..c2ae18fc667 100644 --- a/test/op_shape_test.cpp +++ b/test/op_shape_test.cpp @@ -1408,46 +1408,44 @@ TEST_CASE(dot_sym_k_vs_range) expect_shape(expected, migraphx::make_op("dot"), s_a, s_b); } -// dyn_slice takes its bounds as inputs; this is the shape they have for `n` sliced axes. -static migraphx::shape dyn_slice_bounds(std::size_t n) -{ - return migraphx::shape{migraphx::shape::int64_type, {n}}; -} - TEST_CASE(dyn_slice_static) { // Concrete bounds over a static input stay static: a slice is just a view. migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + // dyn_slice takes its bounds as inputs, with one element per sliced axis. + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; migraphx::shape expected{migraphx::shape::float_type, {2, 2, 2}, {6, 3, 1}}; expect_shape(expected, migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}), input, - dyn_slice_bounds(1), - dyn_slice_bounds(1)); + bounds, + bounds); } TEST_CASE(dyn_slice_static_clamped_bounds) { // Out of range bounds are clipped and negative bounds resolved against the axis length. migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; migraphx::shape expected{migraphx::shape::float_type, {2, 2, 2}, {6, 3, 1}}; expect_shape(expected, migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {-2}}, {"ends", {10}}}), input, - dyn_slice_bounds(1), - dyn_slice_bounds(1)); + bounds, + bounds); } TEST_CASE(dyn_slice_negative_axis) { // A negative axis attribute is normalized against the input rank. migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; migraphx::shape expected{migraphx::shape::float_type, {2, 2, 2}, {6, 3, 1}}; expect_shape(expected, migraphx::make_op("dyn_slice", {{"axes", {-1}}, {"starts", {1}}, {"ends", {3}}}), input, - dyn_slice_bounds(1), - dyn_slice_bounds(1)); + bounds, + bounds); } TEST_CASE(dyn_slice_symbolic_end_static_input) @@ -1463,7 +1461,8 @@ TEST_CASE(dyn_slice_symbolic_end_static_input) migraphx::shape sin{migraphx::shape::float_type, {10}}; migraphx::shape sout{ migraphx::shape::float_type, {dd{migraphx::sym::min(n, lit(10))}}, {lit(1)}}; - expect_shape(sout, op, sin, dyn_slice_bounds(1), dyn_slice_bounds(1)); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; + expect_shape(sout, op, sin, bounds, bounds); EXPECT(sout.symbolic()); EXPECT(not sout.is_fixed()); EXPECT(sout.to_static({{n, 7}}) == migraphx::shape{migraphx::shape::float_type, {7}, {1}}); @@ -1474,7 +1473,7 @@ TEST_CASE(dyn_slice_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. - auto bounds = dyn_slice_bounds(1); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; { // Symbolic end clamped to the axis length 12: dim = min(n, 12) - 2. auto m = var("m", {1, 16}); @@ -1526,7 +1525,7 @@ TEST_CASE(dyn_slice_sym_data_fixed_axis) // result matches slicing the equivalent static shape. auto n = var("n", {1, 8}); std::unordered_map sym_map = {{n, 5}}; - auto bounds = dyn_slice_bounds(1); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{lit(2)}, dd{lit(3)}}}; @@ -1540,10 +1539,10 @@ TEST_CASE(dyn_slice_sym_data_symbolic_axis) { // Slicing the non-fixed symbolic axis: the concrete end is clamped against the axis symbol // rather than a compile-time length, so the extent stays symbolic. - auto n = var("n", {1, 8}); - auto m = var("m", {1, 8}); - auto bounds = dyn_slice_bounds(1); - auto op = migraphx::make_op("dyn_slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}); + auto n = var("n", {1, 8}); + auto m = var("m", {1, 8}); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; + auto op = migraphx::make_op("dyn_slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {2}}}); migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{m}}}; migraphx::shape sout{ @@ -1560,7 +1559,7 @@ TEST_CASE(dyn_slice_sym_multiple_axes) // Slice two axes at once; the symbol at the untouched axis survives. auto n = var("n", {1, 8}); std::unordered_map sym_map = {{n, 4}}; - auto bounds = dyn_slice_bounds(2); + migraphx::shape bounds{migraphx::shape::int64_type, {2}}; auto op = migraphx::make_op("dyn_slice", {{"axes", {0, 2}}, {"starts", {1, 2}}, {"ends", {4, 5}}}); @@ -1576,7 +1575,7 @@ TEST_CASE(dyn_slice_sym_nonstandard_layout) // Non-standard symbolic input: the slice must preserve the permutation. auto n = var("n", {1, 8}); std::unordered_map sym_map = {{n, 6}}; - auto bounds = dyn_slice_bounds(1); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; auto sin = migraphx::shape::from_permutation( migraphx::shape::float_type, {dd{n}, dd{lit(3)}, dd{lit(5)}, dd{lit(7)}}, {0, 2, 3, 1}); @@ -1589,8 +1588,8 @@ TEST_CASE(dyn_slice_wrong_number_of_inputs_error) { // The bounds are the only inputs: the axes are an attribute, not a fourth input. migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; - auto bounds = dyn_slice_bounds(1); - auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; + auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); throws_shape(op, input, bounds); throws_shape(op, input, bounds, bounds, bounds); } @@ -1614,14 +1613,16 @@ TEST_CASE(dyn_slice_dynamic_bounds_input_error) TEST_CASE(dyn_slice_bounds_input_length_error) { migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; + // One axis is sliced, so a two element bounds input does not match. + migraphx::shape bounds{migraphx::shape::int64_type, {2}}; auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}); - throws_shape(op, input, dyn_slice_bounds(2), dyn_slice_bounds(2)); + throws_shape(op, input, bounds, bounds); } TEST_CASE(dyn_slice_attribute_length_error) { migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; - auto bounds = dyn_slice_bounds(2); + migraphx::shape bounds{migraphx::shape::int64_type, {2}}; auto op = migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {0, 1}}, {"ends", {2, 3}}}); throws_shape(op, input, bounds, bounds); } @@ -1629,7 +1630,7 @@ TEST_CASE(dyn_slice_attribute_length_error) TEST_CASE(dyn_slice_missing_attribute_error) { migraphx::shape input{migraphx::shape::float_type, {2, 2, 3}}; - auto bounds = dyn_slice_bounds(1); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; throws_shape( migraphx::make_op("dyn_slice", {{"starts", {1}}, {"ends", {3}}}), input, bounds, bounds); } @@ -1638,7 +1639,7 @@ TEST_CASE(dyn_slice_range_dynamic_data_error) { // A range-based dynamic dimension has no expression to build a symbolic extent from. migraphx::shape input{migraphx::shape::float_type, {{2, 4}, {3, 3}}}; - auto bounds = dyn_slice_bounds(1); + migraphx::shape bounds{migraphx::shape::int64_type, {1}}; throws_shape(migraphx::make_op("dyn_slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {2}}}), input, bounds, diff --git a/test/ref/dyn_slice.cpp b/test/ref/dyn_slice.cpp index f38154ff8fa..6a8b083e2db 100644 --- a/test/ref/dyn_slice.cpp +++ b/test/ref/dyn_slice.cpp @@ -36,64 +36,38 @@ using dd = migraphx::shape::dynamic_dimension; using migraphx::sym::var; -// The bounds are always inputs, so every test supplies them as parameters named starts/ends -// with this shape. -static migraphx::shape dyn_slice_bounds_shape(std::size_t n) -{ - return migraphx::shape{migraphx::shape::int64_type, {n}}; -} - -static void add_dyn_slice(migraphx::module& m, - const migraphx::value& attributes, - migraphx::instruction_ref data, - std::size_t nbounds) -{ - auto bounds_shape = dyn_slice_bounds_shape(nbounds); - auto starts = m.add_parameter("starts", bounds_shape); - auto ends = m.add_parameter("ends", bounds_shape); - m.add_instruction(migraphx::make_op("dyn_slice", attributes), data, starts, ends); -} - -static migraphx::parameter_map bounds_params(const std::vector& starts, - const std::vector& ends) -{ - auto bounds_shape = dyn_slice_bounds_shape(starts.size()); - migraphx::parameter_map params; - // Going through a literal so the arguments own their buffers. - params["starts"] = migraphx::literal{bounds_shape, starts}.get_argument(); - params["ends"] = migraphx::literal{bounds_shape, ends}.get_argument(); - return params; -} - -static migraphx::literal iota_literal(const migraphx::shape& s) -{ - std::vector data(s.elements()); - std::iota(data.begin(), data.end(), 0); - return migraphx::literal{s, data}; -} - -static std::vector read_ints(const migraphx::argument& arg) -{ - std::vector result; - arg.visit([&](auto output) { result.assign(output.begin(), output.end()); }); - return result; -} - TEST_CASE(dyn_slice_concrete_bounds_test) { migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}, data, 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}), + l0, + starts, + ends); // Every bound is concrete, so the output shape is known at compile time. EXPECT(p.get_output_shapes().back() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); p.compile(migraphx::make_target("ref")); - auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector starts_data = {1}; + std::vector ends_data = {3}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); // The static output shape lets the compiler make the aliased view contiguous. EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {4, 2, 1}}); @@ -106,24 +80,45 @@ TEST_CASE(dyn_slice_sym_ends_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, - {{"axes", {2}}, - {"starts", {1}}, - {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 3})})}}}, - data, - 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op( + "dyn_slice", + {{"axes", {2}}, + {"starts", {1}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 3})})}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - auto result0 = p.eval(bounds_params({1}, {3})).back(); + std::vector starts_data = {1}; + std::vector ends_data0 = {3}; + std::vector ends_data1 = {2}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data0.data()); + + auto result0 = p.eval(params).back(); + std::vector results_vector0; + result0.visit([&](auto output) { results_vector0.assign(output.begin(), output.end()); }); std::vector gold0 = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result0), gold0)); + EXPECT(migraphx::verify::verify_rms_range(results_vector0, gold0)); EXPECT(result0.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); - auto result1 = p.eval(bounds_params({1}, {2})).back(); + params["ends"] = migraphx::argument(bounds_shape, ends_data1.data()); + + auto result1 = p.eval(params).back(); + std::vector results_vector1; + result1.visit([&](auto output) { results_vector1.assign(output.begin(), output.end()); }); std::vector gold1 = {1, 4, 7, 10}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result1), gold1)); + EXPECT(migraphx::verify::verify_rms_range(results_vector1, gold1)); EXPECT(result1.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 1}, {6, 3, 1}}); } @@ -133,18 +128,34 @@ TEST_CASE(dyn_slice_sym_starts_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, - {{"axes", {2}}, - {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 2})})}}, - {"ends", {3}}}, - data, - 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op( + "dyn_slice", + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 2})})}}, + {"ends", {3}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector starts_data = {1}; + std::vector ends_data = {3}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); } @@ -155,18 +166,34 @@ TEST_CASE(dyn_slice_sym_both_bounds_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, - {{"axes", {2}}, - {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 1})})}}, - {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 3})})}}}, - data, - 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op( + "dyn_slice", + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 1})})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 3})})}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector starts_data = {1}; + std::vector ends_data = {3}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); } @@ -179,18 +206,32 @@ TEST_CASE(dyn_slice_sym_data_test) auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {dd{var("n", {1, 4})}, dd{migraphx::sym::lit(2)}, dd{migraphx::sym::lit(3)}}}; - auto data = mm->add_parameter("x", s); - add_dyn_slice(*mm, {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}, data, 1); + auto x = mm->add_parameter("x", s); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op("dyn_slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}), + x, + starts, + ends); p.compile(migraphx::make_target("ref")); migraphx::shape input_fixed_shape{migraphx::shape::int32_type, {2, 2, 3}}; - auto data_literal = iota_literal(input_fixed_shape); - auto params = bounds_params({1}, {3}); - params["x"] = data_literal.get_argument(); + std::vector data(input_fixed_shape.elements()); + std::iota(data.begin(), data.end(), 0); + std::vector starts_data = {1}; + std::vector ends_data = {3}; + migraphx::parameter_map params; + params["x"] = migraphx::argument(input_fixed_shape, data.data()); + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); - auto result = p.eval(params).back(); + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); } @@ -202,20 +243,35 @@ TEST_CASE(dyn_slice_sym_bounds_multi_axes_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, - {{"axes", {1, 2}}, - {"starts", {1, 0}}, - {"ends", - migraphx::value::array{migraphx::to_value(dd{var("n", {1, 2})}), - migraphx::to_value(dd{var("m", {1, 3})})}}}, - data, - 2); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {2}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op("dyn_slice", + {{"axes", {1, 2}}, + {"starts", {1, 0}}, + {"ends", + migraphx::value::array{migraphx::to_value(dd{var("n", {1, 2})}), + migraphx::to_value(dd{var("m", {1, 3})})}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - auto result = p.eval(bounds_params({1, 0}, {2, 2})).back(); + std::vector starts_data = {1, 0}; + std::vector ends_data = {2, 2}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {3, 4, 9, 10}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 1, 2}, {6, 3, 1}}); } @@ -227,18 +283,34 @@ TEST_CASE(dyn_slice_runtime_bounds_clamped_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, - {{"axes", {2}}, - {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 1})})}}, - {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 8})})}}}, - data, - 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op( + "dyn_slice", + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 1})})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {1, 8})})}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - auto result = p.eval(bounds_params({-2}, {100})).back(); + std::vector starts_data = {-2}; + std::vector ends_data = {100}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}}); } @@ -250,13 +322,30 @@ TEST_CASE(dyn_slice_negative_axis_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, {{"axes", {-1}}, {"starts", {1}}, {"ends", {3}}}, data, 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op("dyn_slice", {{"axes", {-1}}, {"starts", {1}}, {"ends", {3}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - auto result = p.eval(bounds_params({1}, {3})).back(); + std::vector starts_data = {1}; + std::vector ends_data = {3}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + auto result = p.eval(params).back(); + std::vector results_vector; + result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); }); std::vector gold = {1, 2, 4, 5, 7, 8, 10, 11}; - EXPECT(migraphx::verify::verify_rms_range(read_ints(result), gold)); + EXPECT(migraphx::verify::verify_rms_range(results_vector, gold)); EXPECT(result.get_shape() == migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {4, 2, 1}}); } @@ -266,14 +355,28 @@ TEST_CASE(dyn_slice_end_before_start_error_test) migraphx::program p; auto* mm = p.get_main_module(); migraphx::shape s{migraphx::shape::int32_type, {2, 2, 3}}; - auto data = mm->add_literal(iota_literal(s)); - add_dyn_slice(*mm, - {{"axes", {2}}, - {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 3})})}}, - {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {0, 3})})}}}, - data, - 1); + std::vector data(s.elements()); + std::iota(data.begin(), data.end(), 0); + auto l0 = mm->add_literal(migraphx::literal{s, data}); + migraphx::shape bounds_shape{migraphx::shape::int64_type, {1}}; + auto starts = mm->add_parameter("starts", bounds_shape); + auto ends = mm->add_parameter("ends", bounds_shape); + mm->add_instruction( + migraphx::make_op( + "dyn_slice", + {{"axes", {2}}, + {"starts", migraphx::value::array{migraphx::to_value(dd{var("m", {0, 3})})}}, + {"ends", migraphx::value::array{migraphx::to_value(dd{var("n", {0, 3})})}}}), + l0, + starts, + ends); p.compile(migraphx::make_target("ref")); - EXPECT(test::throws([&] { p.eval(bounds_params({2}, {1})); })); + std::vector starts_data = {2}; + std::vector ends_data = {1}; + migraphx::parameter_map params; + params["starts"] = migraphx::argument(bounds_shape, starts_data.data()); + params["ends"] = migraphx::argument(bounds_shape, ends_data.data()); + + EXPECT(test::throws([&] { p.eval(params); })); } From d58785c41b610a2b8e7ddcc970900dda6ac85880 Mon Sep 17 00:00:00 2001 From: charlie Date: Wed, 5 Aug 2026 16:24:54 -0500 Subject: [PATCH 3/6] Cleanup normalize_attributes --- src/normalize_attributes.cpp | 396 ++++++++++++++++++----------------- 1 file changed, 202 insertions(+), 194 deletions(-) diff --git a/src/normalize_attributes.cpp b/src/normalize_attributes.cpp index c8709b25f94..9464ffdc719 100644 --- a/src/normalize_attributes.cpp +++ b/src/normalize_attributes.cpp @@ -22,6 +22,7 @@ * THE SOFTWARE. */ #include +#include #include #include #include @@ -104,6 +105,98 @@ static std::vector tune_attribute_sym(const std::vector& ex return result; } +/** + * The maximum that each value is normalized against: the rank of the input, or the length of the + * axis the value applies to when `use_len` is set. There is one entry per value, and each axis + * picks the entry to fill. + * + * Returns nullopt when a dynamic_dimension at `axes` is not fixed, since it has no single length + * to normalize against. + */ +template +static optional> +attribute_max_vals(std::size_t nvals, + const std::vector& axes, + const std::vector& attrs, + const shape& input_shape, + Message m) +{ + int64_t n_rank = input_shape.ndim(); + if(contains(attrs, op::normalize_attribute::use_output)) + { + n_rank = n_rank + nvals; + } + std::vector max_vals(nvals, n_rank); + if(not contains(attrs, op::normalize_attribute::use_len)) + return max_vals; + if(axes.size() > nvals) + MIGRAPHX_THROW(m() + "more axes than values to normalize!"); + if(not input_shape.dynamic()) + { + std::transform(axes.begin(), axes.end(), max_vals.begin(), [&](auto i) { + return input_shape.lens().at(i); + }); + return max_vals; + } + if(std::any_of(axes.begin(), axes.end(), [&](auto ax) { + return not input_shape.dyn_dims().at(ax).is_fixed(); + })) + return nullopt; + std::transform(axes.begin(), axes.end(), max_vals.begin(), [&](auto i) { + return input_shape.dyn_dims().at(i).get_interval().max; + }); + return max_vals; +} + +/// Clips the values above the maximum, or range checks them when clipping is off. +template +static void clip_or_check_max(std::vector& result, + const std::vector& max_vals, + const std::vector& attrs, + Message m) +{ + bool inclusive = contains(attrs, op::normalize_attribute::include_max); + if(contains(attrs, op::normalize_attribute::clip_max)) + { + std::transform( + result.begin(), result.end(), max_vals.begin(), result.begin(), [&](auto v, auto mv) { + auto limit = inclusive ? mv : mv - 1; + return v > limit ? limit : v; + }); + return; + } + bool in_range = + inclusive ? std::equal(result.begin(), result.end(), max_vals.begin(), std::less_equal<>{}) + : std::equal(result.begin(), result.end(), max_vals.begin(), std::less<>{}); + if(not in_range) + MIGRAPHX_THROW(m() + "value out of range!"); +} + +/// Clips the values below the minimum, or range checks them when clipping is off. +template +static void clip_or_check_min(std::vector& result, + const std::vector& min_vals, + const std::vector& attrs, + Message m) +{ + bool inclusive = contains(attrs, op::normalize_attribute::include_min); + if(contains(attrs, op::normalize_attribute::clip_min)) + { + std::transform( + result.begin(), result.end(), min_vals.begin(), result.begin(), [&](auto v, auto mv) { + auto limit = inclusive ? mv : mv + 1; + return v < limit ? limit : v; + }); + return; + } + bool in_range = + inclusive + ? std::equal(min_vals.begin(), min_vals.end(), result.begin(), std::less_equal<>{}) + : std::equal(min_vals.begin(), min_vals.end(), result.begin(), std::less<>{}); + if(not in_range) + MIGRAPHX_THROW(m() + "attribute out of range!"); +} + /** * Parameters: * vec: the vector attribute to normalize @@ -116,131 +209,30 @@ static std::vector tune_attribute_sym(const std::vector& ex * See normalize_attribute.hpp for explaining the options. */ template -static auto tune_attribute(const std::vector& vec, - const std::vector& axes, - const value& val, - const shape& input_shape, - Message m) +static std::vector tune_attribute(const std::vector& vec, + const std::vector& axes, + const value& val, + const shape& input_shape, + Message m) { std::vector result(vec); if(result.empty()) - { return result; - }; - int64_t n_rank = input_shape.ndim(); - std::vector vec_attrs = val.to_vector(); - if(contains(vec_attrs, op::normalize_attribute::use_output)) - { - n_rank = n_rank + vec.size(); - } - - std::vector max_vals(vec.size(), n_rank); - - if(contains(vec_attrs, op::normalize_attribute::use_len)) - { - // max_vals has one entry per value, and each axis picks the entry to fill. - if(axes.size() > vec.size()) - MIGRAPHX_THROW(m() + "more axes than values to normalize!"); - if(input_shape.dynamic()) - { - // return the unchanged `vec` if the dynamic_dimensions at `axes` are not fixed - if(std::any_of(axes.begin(), axes.end(), [&](auto ax) { - return not input_shape.dyn_dims().at(ax).is_fixed(); - })) - { - return vec; - } - std::transform(axes.begin(), axes.end(), max_vals.begin(), [&](auto i) { - return input_shape.dyn_dims().at(i).get_interval().max; - }); - } - else - { - std::transform(axes.begin(), axes.end(), max_vals.begin(), [&](auto i) { - return input_shape.lens().at(i); - }); - } - } - - if(contains(vec_attrs, op::normalize_attribute::clip_max)) - { - if(contains(vec_attrs, op::normalize_attribute::include_max)) - { - std::transform(result.begin(), - result.end(), - max_vals.begin(), - result.begin(), - [](auto v, auto mv) { return v > mv ? mv : v; }); - } - else - { - std::transform(result.begin(), - result.end(), - max_vals.begin(), - result.begin(), - [](auto v, auto mv) { return v >= mv ? mv - 1 : v; }); - } - } - else - { - if(contains(vec_attrs, op::normalize_attribute::include_max)) - { - if(not std::equal(result.begin(), result.end(), max_vals.begin(), std::less_equal<>{})) - { - MIGRAPHX_THROW(m() + "value out of range!"); - } - } - else - { - if(not std::equal(result.begin(), result.end(), max_vals.begin(), std::less<>{})) - { - MIGRAPHX_THROW(m() + "value out of range!"); - } - } - } + auto attrs = val.to_vector(); + auto max_vals = attribute_max_vals(vec.size(), axes, attrs, input_shape, m); + // Without a length to normalize against, the values are returned unchanged. The caller has to + // renormalize once the dimensions are known. + if(not max_vals.has_value()) + return result; + clip_or_check_max(result, *max_vals, attrs, m); - std::vector min_vals = max_vals; - std::transform(min_vals.begin(), min_vals.end(), min_vals.begin(), [](auto v) { return -v; }); - if(contains(vec_attrs, op::normalize_attribute::clip_min)) - { - if(contains(vec_attrs, op::normalize_attribute::include_min)) - { - std::transform(result.begin(), - result.end(), - min_vals.begin(), - result.begin(), - [](auto v, auto mv) { return v < mv ? mv : v; }); - } - else - { - std::transform(result.begin(), - result.end(), - min_vals.begin(), - result.begin(), - [](auto v, auto mv) { return v < mv + 1 ? mv + 1 : v; }); - } - } - else - { - if(contains(vec_attrs, op::normalize_attribute::include_min)) - { - if(not std::equal( - min_vals.begin(), min_vals.end(), result.begin(), std::less_equal<>{})) - { - MIGRAPHX_THROW(m() + "attribute out of range!"); - } - } - else - { - if(not std::equal(result.begin(), result.end(), min_vals.begin(), std::less<>{})) - { - MIGRAPHX_THROW(m() + "attribute out of range!"); - } - } - } + std::vector min_vals(max_vals->size()); + std::transform(max_vals->begin(), max_vals->end(), min_vals.begin(), [](auto v) { return -v; }); + clip_or_check_min(result, min_vals, attrs, m); + // Resolve the from-the-end (negative) values against the maximum. std::transform( - result.begin(), result.end(), max_vals.begin(), result.begin(), [](auto v, auto mv) { + result.begin(), result.end(), max_vals->begin(), result.begin(), [](auto v, auto mv) { return v < 0 ? v + mv : v; }); @@ -257,9 +249,93 @@ static auto tune_pad_attribute(const value& val) return result; } +/** + * Doubles a padding attribute that only gives the padding for one side of each spatial dimension. + * Dimensions to pad start from the third dimension (index 2). Auto padding is left to the target. + * + * Returns whether the padding attribute is normalized. + */ +static bool normalize_padding_attribute(operation& op, + value& val, + const std::string& key, + const shape& input_shape) +{ + bool use_auto_padding = + (val.contains("padding_mode") and + (val.at("padding_mode").to() != migraphx::op::padding_mode_t::default_)); + if(use_auto_padding) + return false; + auto padding = val.at(key); + auto npad = input_shape.ndim() - 2; + if(padding.size() == 2 * npad) + return true; + if(padding.size() != npad) + MIGRAPHX_THROW("normalize_attributes: inconsistent padding vector size "); + val[key] = tune_pad_attribute(padding); + op.from_value(val); + return true; +} + +/** + * Normalizes an array attribute, symbolically when the attribute opts in with `use_sym` and either + * a value or the axis length it is normalized against is symbolic. See tune_attribute_sym(). + */ +template +static value tune_array_attribute(const value& vv, + const std::vector& axes, + const value& opts, + const shape& input_shape, + Message m) +{ + auto norm_attrs = opts.to_vector(); + // A symbolic value serializes as an object. Normalizing against a symbolic axis length can + // turn a concrete value symbolic too, so a symbolic input shape takes the same path even when + // every value is a plain integer. + bool sym_values = + std::any_of(vv.begin(), vv.end(), [](const auto& e) { return e.is_object(); }); + bool allow_sym = contains(norm_attrs, op::normalize_attribute::use_sym); + if(sym_values and not allow_sym) + MIGRAPHX_THROW(m() + "symbolic values are not supported!"); + if(not vv.empty() and allow_sym and (sym_values or input_shape.symbolic())) + { + auto dims = migraphx::from_value>(vv); + return migraphx::to_value( + tune_attribute_sym(to_sym_exprs(dims), axes, norm_attrs, input_shape, m)); + } + return value(tune_attribute(vv.to_vector(), axes, opts, input_shape, m)); +} + +/// Normalizes one entry of the `normalize_axes` map and writes it back into the operator. +static void +normalize_axes_attribute(operation& op, value& val, const value& rv, const shape& input_shape) +{ + const auto& key = rv.get_key(); + if(not val.contains(key)) + MIGRAPHX_THROW("NORMALIZE_ATTR : op " + op.name() + " attribute \"" + key + + "\" not exist!"); + auto message = [&] { return op.name() + ": " + key + ": "; }; + auto opts = rv.without_key(); + auto vv = val.at(key).without_key(); + if(vv.is_array()) + { + std::vector axes; + if(val.contains("axes")) + { + axes = val.at("axes").without_key().to_vector(); + } + val[key] = tune_array_attribute(vv, axes, opts, input_shape, message); + } + else + { + auto num = vv.to(); + val[key] = tune_attribute({num}, {num}, opts, input_shape, message).front(); + } + op.from_value(val); + val = op.to_value(); +} + /** * Assumptions: - * Dimensions to pad start from the third dimension (index 2). * Called by compute_shape_op() with the shape of the first input. */ bool normalize_attributes(operation& op, const shape& input_shape) @@ -269,88 +345,20 @@ bool normalize_attributes(operation& op, const shape& input_shape) auto val = op.to_value(); if(attrs.contains("normalize_padding")) { - bool use_auto_padding = - (val.contains("padding_mode") and - (val.at("padding_mode").to() != migraphx::op::padding_mode_t::default_)); - if(not use_auto_padding) - { - auto padding = val.at(attrs.at("normalize_padding").to()); - auto padding_size = padding.size(); - auto padding_start = 2; - if(padding_size == 2 * (input_shape.ndim() - padding_start)) - tuned = true; - else if(padding_size != (input_shape.ndim() - padding_start)) - { - MIGRAPHX_THROW("normalize_attributes: inconsistent padding vector size "); - } - else - { - auto result = tune_pad_attribute(padding); - val["padding"] = result; - op.from_value(val); - tuned = true; - } - } + tuned = normalize_padding_attribute( + op, val, attrs.at("normalize_padding").to(), input_shape); } if(not attrs.contains("normalize_axes")) { return tuned; } - auto attr_v = attrs.at("normalize_axes").without_key(); - for(const auto& rv : attr_v) + // The keys are normalized in the order the operator declares them, so `axes` is resolved + // before the bounds that are normalized against it. + for(const auto& rv : attrs.at("normalize_axes").without_key()) { - const auto& key = rv.get_key(); - if(val.contains(key)) - { - auto message = [&] { return op.name() + ": " + key + ": "; }; - auto vv = val.at(key).without_key(); - if(vv.is_array()) - { - auto attrs = rv.without_key().to_vector(); - // A symbolic value serializes as an object. Normalizing against a symbolic axis - // length can turn a concrete value symbolic too, so a symbolic input shape takes - // the same path even when every value is a plain integer. - bool sym_values = - std::any_of(vv.begin(), vv.end(), [](const auto& e) { return e.is_object(); }); - bool allow_sym = contains(attrs, op::normalize_attribute::use_sym); - if(sym_values and not allow_sym) - MIGRAPHX_THROW(message() + "symbolic values are not supported!"); - std::vector axes; - if(val.contains("axes")) - { - axes = val.at("axes").without_key().to_vector(); - } - if(not vv.empty() and allow_sym and (sym_values or input_shape.symbolic())) - { - auto dims = migraphx::from_value>(vv); - val[key] = migraphx::to_value( - tune_attribute_sym(to_sym_exprs(dims), axes, attrs, input_shape, message)); - } - else - { - auto vec = vv.to_vector(); - val[key] = tune_attribute(vec, axes, rv.without_key(), input_shape, message); - } - op.from_value(val); - val = op.to_value(); - tuned = true; - } - else - { - auto num = vv.to(); - auto result = tune_attribute({num}, {num}, rv.without_key(), input_shape, message); - val[key] = result.front(); - op.from_value(val); - val = op.to_value(); - tuned = true; - } - } - else - { - MIGRAPHX_THROW("NORMALIZE_ATTR : op " + op.name() + " attribute \"" + key + - "\" not exist!"); - } + normalize_axes_attribute(op, val, rv, input_shape); + tuned = true; } return tuned; From d4784d4631b0264bfdc12e9f31e9db70ef04726f Mon Sep 17 00:00:00 2001 From: charlie Date: Wed, 5 Aug 2026 16:32:37 -0500 Subject: [PATCH 4/6] Fix licensing --- src/include/migraphx/op/normalize_attribute.hpp | 2 +- test/normalize_ops_test.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/include/migraphx/op/normalize_attribute.hpp b/src/include/migraphx/op/normalize_attribute.hpp index 7989ca0c3ad..bae3116c68a 100644 --- a/src/include/migraphx/op/normalize_attribute.hpp +++ b/src/include/migraphx/op/normalize_attribute.hpp @@ -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 diff --git a/test/normalize_ops_test.cpp b/test/normalize_ops_test.cpp index dfe1a9bef08..9b602617005 100644 --- a/test/normalize_ops_test.cpp +++ b/test/normalize_ops_test.cpp @@ -1,7 +1,7 @@ /* * The MIT License (MIT) * - * Copyright (c) 2015-2025 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 From fe49783d58fe79de7f138cdfc78938f521e99591 Mon Sep 17 00:00:00 2001 From: charlie Date: Wed, 5 Aug 2026 16:41:04 -0500 Subject: [PATCH 5/6] Simplify comment --- src/include/migraphx/op/dyn_slice.hpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/include/migraphx/op/dyn_slice.hpp b/src/include/migraphx/op/dyn_slice.hpp index f1d25c60539..1cee67f3427 100644 --- a/src/include/migraphx/op/dyn_slice.hpp +++ b/src/include/migraphx/op/dyn_slice.hpp @@ -41,10 +41,9 @@ namespace op { /// Slice operator whose bounds are only known at run time. /// /// The starts and ends are always supplied as inputs. The attribute of the same name describes -/// that input at compile time: either a concrete value or, when the value is only known -/// symbolically, a symbolic dynamic_dimension whose expression evaluates to what the input will -/// hold at run time. That is what lets the output shape stay symbolic instead of collapsing to a -/// range. The axes have to be known when the shape is computed, so they are an attribute only. +/// that input at compile time: either a concrete value or a symbolic dynamic_dimension whose +/// expression evaluates to what the input will hold at run time. The axes have to be known when +/// the shape is computed, so they are an attribute only. /// /// Attributes: /// axes: axes to slice over From b769ef9ea8a97d84e5a40c88de8d17bb057aa73e Mon Sep 17 00:00:00 2001 From: charlie Date: Thu, 6 Aug 2026 13:33:25 -0500 Subject: [PATCH 6/6] Change over to using other dyn_slice base --- CHANGELOG.md | 1 + src/CMakeLists.txt | 1 + src/include/migraphx/op/dyn_topk.hpp | 119 +++++++++ src/include/migraphx/operators.hpp | 1 + src/onnx/parse_topk.cpp | 69 ++--- src/targets/gpu/CMakeLists.txt | 2 - src/targets/gpu/device/topk.cpp | 239 ------------------ .../gpu/include/migraphx/gpu/device/topk.hpp | 55 ---- src/targets/gpu/include/migraphx/gpu/topk.hpp | 62 ----- src/targets/gpu/lowering.cpp | 27 +- src/targets/gpu/topk.cpp | 56 ---- test/gpu/dyn_topk.cpp | 82 ++++++ test/onnx/parse/topk_var_k_test.cpp | 66 ++--- test/op_shape_test.cpp | 72 ++++++ test/ref/dyn_topk.cpp | 106 ++++++++ 15 files changed, 466 insertions(+), 492 deletions(-) create mode 100644 src/include/migraphx/op/dyn_topk.hpp delete mode 100644 src/targets/gpu/device/topk.cpp delete mode 100644 src/targets/gpu/include/migraphx/gpu/device/topk.hpp delete mode 100644 src/targets/gpu/include/migraphx/gpu/topk.hpp delete mode 100644 src/targets/gpu/topk.cpp create mode 100644 test/gpu/dyn_topk.cpp create mode 100644 test/ref/dyn_topk.cpp diff --git a/CHANGELOG.md b/CHANGELOG.md index 79fa01f196f..2939fe91944 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,7 @@ Full documentation for MIGraphX is available at * Added a `--start-from` or `-s` flag to test binaries which resumes from a test name in the list instead of the beginning (#5072). * Added a `dyn_slice` operator, `dyn_slice(data, starts, ends)`, that describes its bound inputs with symbolic attributes so slicing by a data-dependent bound keeps a symbolic output shape; the axes are an attribute since they must be known when the shape is computed (#5088). * Added symbolic normalization of operator attributes, selected with the `use_sym` normalize attribute, which clamps a bound against a symbolic axis length instead of leaving it unnormalized (#5088). +* Added a `dyn_topk` operator that takes `k` as a runtime input, so ONNX TopK with a data-dependent `K` no longer has to pad to the axis length and slice. ### Changed diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index bf770276c3b..da9964e48d1 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -216,6 +216,7 @@ register_migraphx_ops( div dot dyn_slice + dyn_topk elu equal erf diff --git a/src/include/migraphx/op/dyn_topk.hpp b/src/include/migraphx/op/dyn_topk.hpp new file mode 100644 index 00000000000..10dec420250 --- /dev/null +++ b/src/include/migraphx/op/dyn_topk.hpp @@ -0,0 +1,119 @@ +/* + * The MIT License (MIT) + * + * 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 + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ +#ifndef MIGRAPHX_GUARD_OPERATORS_DYN_TOPK_HPP +#define MIGRAPHX_GUARD_OPERATORS_DYN_TOPK_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace migraphx { +inline namespace MIGRAPHX_INLINE_NS { +namespace op { + +/** + * TopK with a data-dependent `k`, matching the ONNX spec where `K` is a runtime input. + * arg[0]: input data + * arg[1]: k value, a static 1-D single-element tensor + * + * The `k` attribute is the symbol standing for that runtime value, so the output length + * along `axis` can be described as min(k, axis length) at compile time. The exact length + * is only known in compute(). + */ +struct dyn_topk +{ + sym::expr k{}; + int64_t axis = 0; + bool largest = true; + + template + static auto reflect(Self& self, F f) + { + return pack(f(self.k, "k"), f(self.axis, "axis"), f(self.largest, "largest")); + } + + value attributes() const + { + value normalize; + normalize["axis"] = value::array{normalize_attribute::include_min}; + return {{"normalize_axes", normalize}}; + } + + std::string name() const { return "dyn_topk"; } + + shape normalize_compute_shape(std::vector inputs) const + { + check_shapes{inputs, *this, true}.has(2); + check_shapes{inputs.begin() + 1, inputs.end(), std::string("DYN_TOPK: k input"), false} + .only_dims(1) + .elements(1); + if(k.empty() or k.name() != "variable") + MIGRAPHX_THROW("DYN_TOPK: k attribute must be a symbolic variable"); + + const auto& input_shape = inputs.at(0); + auto type = input_shape.type(); + + // A range-based axis has no symbolic length to clamp against, so fall back to the + // widest possible extent. + // TODO: remove this when range-based dynamic shapes are removed + if(input_shape.dynamic() and not input_shape.symbolic()) + { + auto dyn_dims = input_shape.dyn_dims(); + dyn_dims[axis] = {0, input_shape.max_lens().at(axis)}; + return shape({shape{type, dyn_dims}, shape{shape::int64_type, dyn_dims}}); + } + + auto sym_in = input_shape.to_symbolic(); + auto dyn_dims = sym_in.dyn_dims(); + dyn_dims[axis] = shape::dynamic_dimension{sym::min(k, dyn_dims[axis].sym_expr)}; + return shape({shape{type, dyn_dims}, shape{shape::int64_type, dyn_dims}}); + } + + argument compute(const dyn_output&, std::vector args) const + { + auto input_shape = args.front().get_shape(); + std::size_t k_val = 0; + args.at(1).visit([&](auto v) { k_val = v.front(); }); + auto actual_k = std::min(k_val, input_shape.lens().at(axis)); + auto out_lens = input_shape.lens(); + out_lens[axis] = actual_k; + + shape out{{shape{input_shape.type(), out_lens}, shape{shape::int64_type, out_lens}}}; + return topk{static_cast(actual_k), axis, largest}.compute(dyn_output{out, out}, + {args.front()}); + } +}; + +} // namespace op +} // namespace MIGRAPHX_INLINE_NS +} // namespace migraphx + +#endif diff --git a/src/include/migraphx/operators.hpp b/src/include/migraphx/operators.hpp index 78b0d28b088..0023afe7f17 100644 --- a/src/include/migraphx/operators.hpp +++ b/src/include/migraphx/operators.hpp @@ -54,6 +54,7 @@ #include #include #include +#include #include #include #include diff --git a/src/onnx/parse_topk.cpp b/src/onnx/parse_topk.cpp index 7481ddcb5e7..2626aa2a34e 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,48 +54,54 @@ 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(); - } - else - { - var_k = true; + k = info.attributes.at("k").i(); } - } - else if(contains(info.attributes, "k")) - { - k = info.attributes.at("k").i(); + auto topk_ret = info.add_instruction( + make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0)); + 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. A constant `k` gives an exactly sized output, + // so no slicing is needed. + 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); + auto topk_ret = info.add_instruction( + make_op("topk", {{"k", arg_k.at()}, {"axis", axis}, {"largest", largest}}), + args.at(0)); + 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}; } - auto topk_ret = info.add_instruction( - make_op("topk", {{"k", k}, {"axis", axis}, {"largest", largest}}), args.at(0)); + // Variable (data-dependent) `k`: name the runtime value with a symbol and let dyn_topk + // describe its output as min(k, axis length). ONNX requires 1 <= k <= axis length, and + // those bounds are what keep the resulting dimension's interval finite. + // TODO: rewrite_topk should later turn this into topk + slice so it can run on a target. + auto input_shape = args.at(0)->get_shape(); + auto norm_axis = axis < 0 ? axis + input_shape.ndim() : axis; + int64_t k_max = input_shape.max_lens().at(norm_axis); + auto k_var = sym::var(info.name, {1, k_max}); + auto topk_ret = info.add_instruction( + make_op("dyn_topk", {{"k", to_value(k_var)}, {"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)); - } - return {ret_val, ret_ind}; } }; diff --git a/src/targets/gpu/CMakeLists.txt b/src/targets/gpu/CMakeLists.txt index 2dbaac523ea..a367bf3a84b 100644 --- a/src/targets/gpu/CMakeLists.txt +++ b/src/targets/gpu/CMakeLists.txt @@ -243,7 +243,6 @@ add_library(migraphx_gpu sync_device.cpp target.cpp time_op.cpp - topk.cpp write_literals.cpp fuse_mlss.cpp mlss_conv_op.cpp @@ -271,7 +270,6 @@ endfunction() register_migraphx_gpu_ops(hip_ fixed_pad loop - topk ) if (MIGRAPHX_USE_MIOPEN) register_migraphx_gpu_ops(miopen_ diff --git a/src/targets/gpu/device/topk.cpp b/src/targets/gpu/device/topk.cpp deleted file mode 100644 index 0eb8b790160..00000000000 --- a/src/targets/gpu/device/topk.cpp +++ /dev/null @@ -1,239 +0,0 @@ -/* - * The MIT License (MIT) - * - * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ -#include -#include -#include -#include -#include -#include -#include -#include - -namespace migraphx { -inline namespace MIGRAPHX_INLINE_NS { -namespace gpu { -namespace device { - -template -struct hip_heap_vector -{ - MIGRAPHX_DEVICE_CONSTEXPR hip_heap_vector(T* val, index_int n, Index v_idx, Compare comp) - : data(val), size(n), data_index(v_idx), compare(comp) - { - make_heap(size); - } - - MIGRAPHX_DEVICE_CONSTEXPR void try_push(const T val) - { - if(compare(val, data[data_index(0)])) - return; - - pop_heap(size - 1); - data[data_index(size - 1)] = val; - push_heap(size - 1); - } - - MIGRAPHX_DEVICE_CONSTEXPR void sort() { sort_heap(size); } - - private: - MIGRAPHX_DEVICE_CONSTEXPR inline static void swap(T& v1, T& v2) noexcept - { - T v = v1; - v1 = v2; - v2 = v; - } - - MIGRAPHX_DEVICE_CONSTEXPR inline void heapify_down(index_int n, index_int index) - { - while(index < n) - { - auto pre_index = index; - index_int l = 2 * index + 1; - index_int r = 2 * index + 2; - - if(l < n and compare(data[data_index(l)], data[data_index(index)])) - { - index = l; - } - - if(r < n and compare(data[data_index(r)], data[data_index(index)])) - { - index = r; - if(compare(data[data_index(l)], data[data_index(r)])) - { - index = l; - } - } - - if(index == pre_index) - { - break; - } - - swap(data[data_index(index)], data[data_index(pre_index)]); - } - } - - MIGRAPHX_DEVICE_CONSTEXPR inline void heapify_up(index_int index) - { - while(index > 0) - { - auto parent_idx = (index - 1) / 2; - - if(not compare(data[data_index(index)], data[data_index(parent_idx)])) - { - break; - } - - swap(data[data_index(index)], data[data_index(parent_idx)]); - index = parent_idx; - } - } - - MIGRAPHX_DEVICE_CONSTEXPR inline void make_heap(index_int n) - { - for(int j = n / 2 - 1; j >= 0; --j) - { - heapify_down(n, j); - } - } - - MIGRAPHX_DEVICE_CONSTEXPR inline void push_heap(index_int loc) { heapify_up(loc); } - - MIGRAPHX_DEVICE_CONSTEXPR inline void pop_heap(index_int loc) - { - swap(data[data_index(0)], data[data_index(loc)]); - heapify_down(loc, 0); - } - - MIGRAPHX_DEVICE_CONSTEXPR inline void sort_heap(index_int n) - { - for(int j = n - 1; j > 0; --j) - { - swap(data[data_index(0)], data[data_index(j)]); - heapify_down(j, 0); - } - } - - T* data = nullptr; - index_int size; - Index data_index; - Compare compare; -}; - -template -__device__ static hip_heap_vector -make_heap(T* data, index_int n, Index idx, Compare compare) -{ - return {data, n, idx, compare}; -} - -template -static std::vector topk(hipStream_t stream, - const argument& val_res, - const argument& ind_res, - const argument& arg, - int64_t k, - int64_t axis, - Compare compare) -{ - auto in_s = arg.get_shape(); - auto in_lens = in_s.lens(); - auto out_s = val_res.get_shape(); - auto axis_dim = in_s.lens()[axis]; - auto comp_lens = in_lens; - comp_lens[axis] = 1; - shape comp_s{in_s.type(), comp_lens}; - std::size_t elem_num = comp_s.elements(); - - hip_visit_all(val_res, arg, out_s, in_s, comp_s)( - [&](auto out_val, auto input, auto oss, auto iss, auto css) { - auto* data = device_cast(input.data()); - auto* out = device_cast(out_val.data()); - auto* const ind = ind_res.cast(); - gs_launch(stream, elem_num)([=](auto i) __device__ { - auto idx = css.multi(i); - - auto in_idx = [&](int ii) { - auto iidx = idx; - iidx[axis] = ii; - return iss.index(iidx); - }; - - auto out_idx = [&](int ii) { - auto iidx = idx; - iidx[axis] = ii; - return oss.index(iidx); - }; - - auto data_compare = [=](auto ii, auto jj) { - return compare(data[in_idx(ii)], data[in_idx(jj)]); - }; - - for(int j = 0; j < k; ++j) - { - ind[out_idx(j)] = j; - } - - auto hp = make_heap(ind, k, out_idx, data_compare); - for(int j = k; j < axis_dim; ++j) - { - hp.try_push(j); - } - hp.sort(); - - for(int j = 0; j < k; ++j) - { - out[out_idx(j)] = data[in_idx(ind[out_idx(j)])]; - } - }); - }); - - return {val_res, ind_res}; -} - -argument topk_largest(hipStream_t stream, - const argument& val_res, - const argument& ind_res, - const argument& arg, - int64_t k, - int64_t axis) -{ - return {topk(stream, val_res, ind_res, arg, k, axis, std::less<>{})}; -} - -argument topk_smallest(hipStream_t stream, - const argument& val_res, - const argument& ind_res, - const argument& arg, - int64_t k, - int64_t axis) -{ - return {topk(stream, val_res, ind_res, arg, k, axis, std::greater<>{})}; -} - -} // namespace device -} // namespace gpu -} // namespace MIGRAPHX_INLINE_NS -} // namespace migraphx diff --git a/src/targets/gpu/include/migraphx/gpu/device/topk.hpp b/src/targets/gpu/include/migraphx/gpu/device/topk.hpp deleted file mode 100644 index b1fb4e8e2ee..00000000000 --- a/src/targets/gpu/include/migraphx/gpu/device/topk.hpp +++ /dev/null @@ -1,55 +0,0 @@ -/* - * The MIT License (MIT) - * - * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ -#ifndef MIGRAPHX_GUARD_RTGLIB_DEVICE_TOPK_HPP -#define MIGRAPHX_GUARD_RTGLIB_DEVICE_TOPK_HPP - -#include -#include -#include - -namespace migraphx { -inline namespace MIGRAPHX_INLINE_NS { -namespace gpu { -namespace device { - -argument MIGRAPHX_DEVICE_EXPORT topk_smallest(hipStream_t stream, - const argument& val_res, - const argument& ind_res, - const argument& arg, - int64_t k, - int64_t axis); - -argument MIGRAPHX_DEVICE_EXPORT topk_largest(hipStream_t stream, - const argument& val_res, - const argument& ind_res, - const argument& arg, - int64_t k, - int64_t axis); - -} // namespace device -} // namespace gpu -} // namespace MIGRAPHX_INLINE_NS -} // namespace migraphx - -#endif diff --git a/src/targets/gpu/include/migraphx/gpu/topk.hpp b/src/targets/gpu/include/migraphx/gpu/topk.hpp deleted file mode 100644 index e07a7c4c24e..00000000000 --- a/src/targets/gpu/include/migraphx/gpu/topk.hpp +++ /dev/null @@ -1,62 +0,0 @@ -/* - * The MIT License (MIT) - * - * 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 - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ -#ifndef MIGRAPHX_GUARD_RTGLIB_TOPK_HPP -#define MIGRAPHX_GUARD_RTGLIB_TOPK_HPP - -#include -#include -#include -#include - -namespace migraphx { -inline namespace MIGRAPHX_INLINE_NS { -namespace gpu { - -struct context; - -struct hip_topk -{ - op::topk op; - - template - static auto reflect(Self& self, F f) - { - return migraphx::reflect(self.op, f); - } - - std::string name() const { return "gpu::topk"; } - shape compute_shape(std::vector inputs) const; - argument - compute(context& ctx, const shape& output_shape, const std::vector& args) const; - std::vector output_alias(const std::vector& shapes) const - { - return {shapes.size() - 1}; - } -}; - -} // namespace gpu -} // namespace MIGRAPHX_INLINE_NS -} // namespace migraphx - -#endif diff --git a/src/targets/gpu/lowering.cpp b/src/targets/gpu/lowering.cpp index e9e7d7c2b88..9002889eb15 100644 --- a/src/targets/gpu/lowering.cpp +++ b/src/targets/gpu/lowering.cpp @@ -111,6 +111,7 @@ struct miopen_apply add_neg_op(); add_lrn_op(); add_nms_op(); + add_dyn_topk_op(); add_convolution_backwards_op(); add_select_module_op(); add_reshape_lazy_op(); @@ -458,14 +459,14 @@ struct miopen_apply const auto& boxes_s = ins->inputs()[0]->get_shape(); const auto& scores_s = ins->inputs()[1]->get_shape(); if(boxes_s.dynamic() or scores_s.dynamic()) - return lower_nms_to_ref(ins); + return lower_tuple_op_to_ref(ins); const auto num_boxes = boxes_s.lens().at(1); const auto num_bc = boxes_s.lens().at(0) * scores_s.lens().at(1); // Route to ref (CPU) when: // - num_boxes < 2: Single box or no boxes, no sort or IoU comparison needed. // - num_bc > 8192: shared-memory limit on the compact kernel. if(num_boxes < 2 or num_bc > 8192) - return lower_nms_to_ref(ins); + return lower_tuple_op_to_ref(ins); return lower_nms_to_gpu_pipeline(ins); }); } @@ -532,10 +533,10 @@ struct miopen_apply return mod->replace_instruction(ins, compact); } - // Dynamic-shape fallback: run the ref op on the host. The tuple has to be - // split host-side before copy_to_gpu (which is not tuple-aware), and the - // downstream get_tuple_elem consumers are rewritten in place. - instruction_ref lower_nms_to_ref(instruction_ref ins) const + // Fallback for a tuple-returning op the GPU cannot handle: run the ref op on the + // host. The tuple has to be split host-side before copy_to_gpu (which is not + // tuple-aware), and the downstream get_tuple_elem consumers are rewritten in place. + instruction_ref lower_tuple_op_to_ref(instruction_ref ins) const { auto inputs = ins->inputs(); std::vector cpu_inputs; @@ -565,9 +566,9 @@ struct miopen_apply for(auto consumer : consumers) { if(consumer->name() != "get_tuple_elem") - MIGRAPHX_THROW("gpu::add_nms_op: dynamic NMS fallback expects only " - "get_tuple_elem consumers of nonmaxsuppression; got: " + - consumer->name()); + MIGRAPHX_THROW("gpu::lower_tuple_op_to_ref: host fallback expects only " + "get_tuple_elem consumers of " + + ins->name() + "; got: " + consumer->name()); auto idx = consumer->get_operator().to_value().at("index").to(); assert(idx < gpu_subs.size()); mod->replace_instruction(consumer, gpu_subs[idx]); @@ -578,6 +579,14 @@ struct miopen_apply return ins; } + // dyn_topk's output length is data-dependent, so there is no fixed-size kernel to compile. + // Run it on the host until rewrite_topk can turn it into topk + slice. + void add_dyn_topk_op() + { + apply_map.emplace("dyn_topk", + [=](instruction_ref ins) { return lower_tuple_op_to_ref(ins); }); + } + void add_lrn_op() { apply_map.emplace("lrn", [=](instruction_ref ins) { diff --git a/src/targets/gpu/topk.cpp b/src/targets/gpu/topk.cpp deleted file mode 100644 index 2e799c650af..00000000000 --- a/src/targets/gpu/topk.cpp +++ /dev/null @@ -1,56 +0,0 @@ -/* - * The MIT License (MIT) - * - * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved. - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in - * all copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN - * THE SOFTWARE. - */ -#include -#include -#include - -namespace migraphx { -inline namespace MIGRAPHX_INLINE_NS { -namespace gpu { - -shape hip_topk::compute_shape(std::vector inputs) const -{ - return op.normalize_compute_shape({inputs.front()}); -} - -argument hip_topk::compute(context& ctx, const shape&, const std::vector& args) const -{ - auto outputs = args.back().get_sub_objects(); - return op.largest ? device::topk_largest(ctx.get_stream().get(), - outputs.front(), - outputs.back(), - args[0], - op.k, - op.axis) - : device::topk_smallest(ctx.get_stream().get(), - outputs.front(), - outputs.back(), - args[0], - op.k, - op.axis); -} - -} // namespace gpu -} // namespace MIGRAPHX_INLINE_NS -} // namespace migraphx diff --git a/test/gpu/dyn_topk.cpp b/test/gpu/dyn_topk.cpp new file mode 100644 index 00000000000..9686225211a --- /dev/null +++ b/test/gpu/dyn_topk.cpp @@ -0,0 +1,82 @@ +/* + * The MIT License (MIT) + * + * 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 + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using migraphx::sym::var; + +// dyn_topk has no GPU kernel because its output length is data-dependent, so lowering routes it +// to the host ref op. This checks that round trip: inputs copied off the device, the op run on +// the host, and both tuple elements copied back. +TEST_CASE(dyn_topk_gpu_host_fallback) +{ + migraphx::program p; + auto* mm = p.get_main_module(); + migraphx::shape data_s{migraphx::shape::float_type, {2, 4}}; + std::vector data_values = {1, 3, 2, 4, 8, 5, 7, 6}; + auto data = mm->add_literal(migraphx::literal{data_s, data_values}); + migraphx::shape k_s{migraphx::shape::int64_type, {1}}; + auto k = mm->add_parameter("k", k_s); + auto out = mm->add_instruction( + migraphx::make_op( + "dyn_topk", {{"k", migraphx::to_value(var("k", {1, 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); + mm->add_return({val, ind}); + + migraphx::target t = migraphx::make_target("gpu"); + p.compile(t); + + std::vector k_data = {2}; + migraphx::parameter_map params; + for(auto&& x : p.get_parameter_shapes()) + { + if(x.first == "k") + params[x.first] = t.copy_to(migraphx::argument(k_s, k_data.data())); + else + params[x.first] = t.allocate(x.second); + } + + auto results = p.eval(params); + std::vector val_v; + std::vector ind_v; + t.copy_from(results.at(0)).visit([&](auto o) { val_v.assign(o.begin(), o.end()); }); + t.copy_from(results.at(1)).visit([&](auto o) { ind_v.assign(o.begin(), o.end()); }); + + EXPECT(migraphx::verify::verify_rms_range(val_v, std::vector{4, 3, 8, 7})); + EXPECT(ind_v == std::vector{3, 1, 0, 2}); +} + +int main(int argc, const char* argv[]) { test::run(argc, argv); } diff --git a/test/onnx/parse/topk_var_k_test.cpp b/test/onnx/parse/topk_var_k_test.cpp index 2338236abe0..f180890706f 100644 --- a/test/onnx/parse/topk_var_k_test.cpp +++ b/test/onnx/parse/topk_var_k_test.cpp @@ -24,47 +24,37 @@ #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`. -TEST_CASE(topk_var_k_test) +// `k` is a runtime input (graph input, not an initializer), so the parser emits dyn_topk with +// the runtime `k` named by a symbol bounded by the axis length. +static void add_dyn_topk(migraphx::module& m, const std::vector& args) { - migraphx::program p; - auto* mm = p.get_main_module(); - 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); - mm->add_return({val, ind}); - - auto prog = read_onnx("topk_var_k_test.onnx"); - - EXPECT(p == prog); + auto k_var = migraphx::sym::var("TopK_2", {1, 4}); + auto out = m.add_instruction( + migraphx::make_op("dyn_topk", + {{"k", migraphx::to_value(k_var)}, {"axis", 1}, {"largest", 1}}), + args[0], + args[1]); + auto val = m.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), out); + auto ind = m.add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), out); + m.add_return({val, ind}); } -// Same model, but `data` is overridden to a dynamic shape. `k` stays a runtime input, so the -// var_k path still fires and sets the topk `k` to the axis dimension's max length. -TEST_CASE(topk_var_k_dynamic_test) +TEST_CASE(topk_var_k_test) { - migraphx::program p; - auto* mm = p.get_main_module(); - 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); - mm->add_return({val, ind}); - - migraphx::onnx_options options; - options.map_dyn_input_dims["data"] = {{1, 4}, {2, 4}}; - auto prog = read_onnx("topk_var_k_test.onnx", options); + EXPECT(check_parse("topk_var_k_test.onnx", + {{"data", {migraphx::shape::float_type, {2, 4}}}, + {"k", {migraphx::shape::int64_type, {1}}}}, + add_dyn_topk)); +} - EXPECT(p == prog); +// Same model, but `data` is overridden to a symbolic shape. The `k` symbol is still bounded by +// the axis length, which comes from the symbol's upper bound rather than a fixed length. +TEST_CASE(topk_var_k_symbolic_test) +{ + using migraphx::sym::var; + EXPECT(check_parse( + "topk_var_k_test.onnx", + {{"data", {migraphx::shape::float_type, sym_dims({var("n", {1, 4}), var("m", {2, 4})})}}, + {"k", {migraphx::shape::int64_type, {1}}}}, + add_dyn_topk)); } diff --git a/test/op_shape_test.cpp b/test/op_shape_test.cpp index c2ae18fc667..63b0539bdb0 100644 --- a/test/op_shape_test.cpp +++ b/test/op_shape_test.cpp @@ -1646,6 +1646,78 @@ TEST_CASE(dyn_slice_range_dynamic_data_error) bounds); } +static migraphx::shape dyn_topk_shape(const std::vector
& dims) +{ + return migraphx::shape({migraphx::shape{migraphx::shape::float_type, dims}, + migraphx::shape{migraphx::shape::int64_type, dims}}); +} + +TEST_CASE(dyn_topk_static_input) +{ + // The runtime `k` is unknown, so a static input still yields a symbolic output: min(k, 4). + auto k = var("k", {1, 4}); + auto op = + migraphx::make_op("dyn_topk", {{"k", migraphx::to_value(k)}, {"axis", 1}, {"largest", 1}}); + migraphx::shape sin{migraphx::shape::float_type, {2, 4}}; + migraphx::shape kin{migraphx::shape::int64_type, {1}}; + expect_shape(dyn_topk_shape({dd{lit(2)}, dd{migraphx::sym::min(k, lit(4))}}), op, sin, kin); +} + +TEST_CASE(dyn_topk_symbolic_input) +{ + // The axis length is itself a symbol, so the sliced extent is min(k, m). + auto k = var("k", {1, 4}); + auto n = var("n", {1, 4}); + auto m = var("m", {2, 4}); + auto op = + migraphx::make_op("dyn_topk", {{"k", migraphx::to_value(k)}, {"axis", 1}, {"largest", 1}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{n}, dd{m}}}; + migraphx::shape kin{migraphx::shape::int64_type, {1}}; + expect_shape(dyn_topk_shape({dd{n}, dd{migraphx::sym::min(k, m)}}), op, sin, kin); +} + +TEST_CASE(dyn_topk_range_input) +{ + // A range-based axis has no symbol to clamp against, so the extent widens to [0, max]. + auto k = var("k", {1, 4}); + auto op = + migraphx::make_op("dyn_topk", {{"k", migraphx::to_value(k)}, {"axis", 1}, {"largest", 1}}); + migraphx::shape sin{migraphx::shape::float_type, {dd{1, 4}, dd{2, 4}}}; + migraphx::shape kin{migraphx::shape::int64_type, {1}}; + expect_shape(dyn_topk_shape({dd{1, 4}, dd{0, 4}}), op, sin, kin); +} + +TEST_CASE(dyn_topk_bad_inputs) +{ + auto k = var("k", {1, 4}); + auto op = + migraphx::make_op("dyn_topk", {{"k", migraphx::to_value(k)}, {"axis", 1}, {"largest", 1}}); + migraphx::shape sin{migraphx::shape::float_type, {2, 4}}; + migraphx::shape kin{migraphx::shape::int64_type, {1}}; + // `k` input is required and there is no indexing input + throws_shape(op, sin); + throws_shape(op, sin, kin, kin); + // `k` must be a static single-element 1-D tensor + throws_shape(op, sin, migraphx::shape{migraphx::shape::int64_type, {2}}); + throws_shape(op, sin, migraphx::shape{migraphx::shape::int64_type, {1, 1}}); + throws_shape(op, sin, migraphx::shape{migraphx::shape::int64_type, {dd{1, 4}}}); +} + +TEST_CASE(dyn_topk_bad_k_attribute) +{ + migraphx::shape sin{migraphx::shape::float_type, {2, 4}}; + migraphx::shape kin{migraphx::shape::int64_type, {1}}; + // `k` names the runtime value, so it has to be a bare variable + throws_shape(migraphx::make_op("dyn_topk", {{"axis", 1}}), sin, kin); + throws_shape( + migraphx::make_op("dyn_topk", {{"k", migraphx::to_value(lit(3))}, {"axis", 1}}), sin, kin); + throws_shape( + migraphx::make_op("dyn_topk", + {{"k", migraphx::to_value(var("k", {1, 4}) + lit(1))}, {"axis", 1}}), + sin, + kin); +} + TEST_CASE(broadcast_for_dot_static) { migraphx::shape s0{migraphx::shape::float_type, {481, 356}}; diff --git a/test/ref/dyn_topk.cpp b/test/ref/dyn_topk.cpp new file mode 100644 index 00000000000..d64f3af386f --- /dev/null +++ b/test/ref/dyn_topk.cpp @@ -0,0 +1,106 @@ +/* + * The MIT License (MIT) + * + * 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 + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using migraphx::sym::var; + +static const migraphx::shape data_shape{migraphx::shape::float_type, {2, 4}}; +static const std::vector data_values = {1, 3, 2, 4, 8, 5, 7, 6}; + +static migraphx::program make_dyn_topk_program(bool largest) +{ + migraphx::program p; + auto* mm = p.get_main_module(); + auto data = mm->add_literal(migraphx::literal{data_shape, data_values}); + auto k = mm->add_parameter("k", {migraphx::shape::int64_type, {1}}); + auto out = mm->add_instruction( + migraphx::make_op( + "dyn_topk", + {{"k", migraphx::to_value(var("k", {1, 4}))}, {"axis", 1}, {"largest", largest}}), + 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}); + p.compile(migraphx::make_target("ref")); + return p; +} + +static std::pair, std::vector> run_with_k(migraphx::program& p, + int64_t k) +{ + migraphx::shape ks{migraphx::shape::int64_type, {1}}; + std::vector k_data = {k}; + migraphx::parameter_map pp; + pp["k"] = migraphx::argument(ks, k_data.data()); + + auto results = p.eval(pp); + std::vector val; + std::vector ind; + results[0].visit([&](auto o) { val.assign(o.begin(), o.end()); }); + results[1].visit([&](auto o) { ind.assign(o.begin(), o.end()); }); + return {val, ind}; +} + +// The compile-time shape is symbolic, so the output size has to come from the `k` argument at +// eval time. Running one compiled program with two different `k` values proves it does. +TEST_CASE(dyn_topk_follows_runtime_k) +{ + auto p = make_dyn_topk_program(true); + + auto [val2, ind2] = run_with_k(p, 2); + EXPECT(val2 == std::vector{4, 3, 8, 7}); + EXPECT(ind2 == std::vector{3, 1, 0, 2}); + + auto [val3, ind3] = run_with_k(p, 3); + EXPECT(val3 == std::vector{4, 3, 2, 8, 7, 6}); + EXPECT(ind3 == std::vector{3, 1, 2, 0, 2, 3}); +} + +TEST_CASE(dyn_topk_smallest) +{ + auto p = make_dyn_topk_program(false); + auto [val, ind] = run_with_k(p, 2); + EXPECT(val == std::vector{1, 2, 5, 6}); + EXPECT(ind == std::vector{0, 2, 1, 3}); +} + +// `k` at the axis length degenerates to a full sort. +TEST_CASE(dyn_topk_k_equals_axis) +{ + auto p = make_dyn_topk_program(true); + auto [val, ind] = run_with_k(p, 4); + EXPECT(val == std::vector{4, 3, 2, 1, 8, 7, 6, 5}); + EXPECT(ind == std::vector{3, 1, 2, 0, 0, 2, 3, 1}); +}