Create dyn_topk operator for handling variable k value with symbolics

CHANGELOG.md

@@ -41,6 +41,9 @@ 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).
+* 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
@@ -61,6 +64,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
 

src/CMakeLists.txt

@@ -215,6 +215,8 @@ register_migraphx_ops(
     dimensions_of
     div
     dot
+    dyn_slice
+    dyn_topk
     elu
     equal
     erf

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 <algorithm>
 #include <cstdint>
 #include <ostream>
 #include <type_traits>
+#include <vector>
 
 #include <migraphx/config.hpp>
 #include <migraphx/picked_variant.hpp>
 #include <migraphx/requires.hpp>
 #include <migraphx/shape.hpp>
+#include <migraphx/sym.hpp>
 
 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<dim_like>& dims)
+{
+    return std::all_of(dims.begin(), dims.end(), [](const dim_like& d) {
+        return std::holds_alternative<int64_t>(d);
+    });
+}
+
+/// Extracts the concrete int64_t from each entry; throws (via std::get) if any entry holds a
+/// dynamic_dimension.
+inline std::vector<int64_t> to_ints(const std::vector<dim_like>& dims)
+{
+    std::vector<int64_t> result(dims.size());
+    std::transform(dims.begin(), dims.end(), result.begin(), [](const dim_like& d) {
+        return std::get<int64_t>(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<sym::expr> to_sym_exprs(const std::vector<dim_like>& dims)
+{
+    std::vector<sym::expr> result(dims.size());
+    std::transform(dims.begin(), dims.end(), result.begin(), [](const dim_like& d) -> sym::expr {
+        if(std::holds_alternative<int64_t>(d))
+            return sym::lit(std::get<int64_t>(d));
+        if(not is_symbolic(d))
+            MIGRAPHX_THROW("DIM_LIKE: cannot convert a range-based dimension to a symbolic "
+                           "expression");
+        return std::get<shape::dynamic_dimension>(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);
 

src/include/migraphx/op/dyn_slice.hpp

@@ -0,0 +1,194 @@
+/*
+ * 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 <migraphx/algorithm.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/config.hpp>
+#include <migraphx/dim_like.hpp>
+#include <migraphx/dyn_output.hpp>
+#include <migraphx/normalize_attributes.hpp>
+#include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/value.hpp>
+
+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 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
+/// 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<int64_t> axes{};
+    std::vector<dim_like> starts{};
+    std::vector<dim_like> ends{};
+
+    template <class Self, class F>
+    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<shape>& 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<shape> inputs) const
+    {
+        check_shapes{inputs, *this, true}.has(3);
+        check_inputs_and_attributes(inputs);
+        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");
+
+        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<sym::expr> 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<argument> args) const
+    {
+        const auto& input = args.front();
+        auto input_shape  = input.get_shape();
+        auto read         = [](const argument& arg) {
+            std::vector<int64_t> result;
+            arg.visit([&](auto values) { result = values.template to_vector<int64_t>(); });
+            return result;
+        };
+        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");
+
+        // Get end-start for output dimension sizes. Reject if ends before starts (no wrap around).
+        std::vector<std::size_t> 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<std::size_t> 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<std::size_t> output_alias(const std::vector<shape>&) const { return {0}; }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif