[AIMIGRAPHX-835] integrate symbolic expression in dynamic_dimension

src/gemm.cpp

@@ -72,15 +72,18 @@ struct batch_slicer
     batch_slicer(const shape& mat_shape)
     {
         auto n_batch_dims = mat_shape.ndim() - 2;
-        inner_shape       = shape{mat_shape.type(),
-                                  {mat_shape.lens().end() - 2, mat_shape.lens().end()},
-                                  {mat_shape.strides().end() - 2, mat_shape.strides().end()}};
+        inner_shape       = shape{
+            mat_shape.type(),
+            std::vector<std::size_t>{mat_shape.lens().end() - 2, mat_shape.lens().end()},
+            std::vector<std::size_t>{mat_shape.strides().end() - 2, mat_shape.strides().end()}};
         if(n_batch_dims > 0)
         {
             outer_shape =
                 shape{mat_shape.type(),
-                      {mat_shape.lens().begin(), mat_shape.lens().begin() + n_batch_dims},
-                      {mat_shape.strides().begin(), mat_shape.strides().begin() + n_batch_dims}};
+                      std::vector<std::size_t>{mat_shape.lens().begin(),
+                                               mat_shape.lens().begin() + n_batch_dims},
+                      std::vector<std::size_t>{mat_shape.strides().begin(),
+                                               mat_shape.strides().begin() + n_batch_dims}};
         }
     }
 

src/include/migraphx/shape.hpp

@@ -39,6 +39,7 @@
 #include <migraphx/bf16.hpp>
 #include <migraphx/float8.hpp>
 #include <migraphx/serialize.hpp>
+#include <migraphx/sym.hpp>
 #include <migraphx/config.hpp>
 
 namespace migraphx {
@@ -94,6 +95,13 @@ struct MIGRAPHX_EXPORT shape
     {
     };
 
+    // TODO: Deprecate the pure range-based form of dynamic_dimension in favor
+    // of the symbolic form (sym_expr). The current design carries two parallel
+    // notions of bounds -- dynamic_dimension::interval (std::size_t min/max,
+    // here) and sym::interval (int64_t min/max, attached to each sym::var) --
+    // which is a source of confusion. Once all shape-producing paths go through
+    // symbolic expressions, `range`/`optimals` and this nested `interval` can
+    // be removed and bounds will live solely on sym::var.
     struct MIGRAPHX_EXPORT dynamic_dimension
     {
         struct interval
@@ -112,42 +120,76 @@ struct MIGRAPHX_EXPORT shape
             friend bool operator!=(const interval& a, const interval& b) { return not(a == b); }
         };
 
-        interval range = {0, 0};
-        std::set<std::size_t> optimals{};
+        std::optional<interval> range;
+        std::optional<std::set<std::size_t>> optimals;
+        sym::expr sym_expr;
 
         dynamic_dimension() = default;
-        dynamic_dimension(std::size_t min_v, std::size_t max_v) : range{min_v, max_v} {}
+        dynamic_dimension(std::size_t min_v, std::size_t max_v)
+            : range{interval{min_v, max_v}}, optimals{std::set<std::size_t>{}}
+        {
+        }
         dynamic_dimension(std::size_t min_v, std::size_t max_v, std::set<std::size_t> opt)
-            : range{min_v, max_v}, optimals(std::move(opt))
+            : range{interval{min_v, max_v}},
+              optimals(min_v == max_v ? std::set<std::size_t>{} : std::move(opt))
+        {
+        }
+        dynamic_dimension(sym::expr s) : sym_expr(std::move(s))
         {
+            if(sym_expr.empty())
+                MIGRAPHX_THROW(
+                    "dynamic_dimension: cannot construct from an empty symbolic expression");
         }
 
         template <class Self, class F>
         static auto reflect(Self& self, F f)
         {
-            return pack(f(self.range, "range"), f(self.optimals, "optimals"));
+            return pack(
+                f(self.range, "range"), f(self.optimals, "optimals"), f(self.sym_expr, "sym"));
         }
 
-        interval get_interval() const { return range; }
-        std::set<std::size_t> get_optimals() const { return optimals; }
+        interval get_interval() const
+        {
+            if(is_symbolic())
+            {
+                auto ival = sym_expr.eval_interval();
+                assert(sym::to<int64_t>(ival.min) >= 0 and sym::to<int64_t>(ival.max) >= 0);
+                return {sym::to<std::size_t>(ival.min), sym::to<std::size_t>(ival.max)};
+            }
+            return *range;
+        }
+        std::set<std::size_t> get_optimals() const
+        {
+            if(is_symbolic())
+                return sym_expr.eval_optimals();
+            if(optimals.has_value())
+                return *optimals;
+            return {};
+        }
 
         bool is_fixed() const;
+        bool is_symbolic() const { return not sym_expr.empty(); }
         bool has_optimal() const;
 
         /**
          * Return a dynamic_dimension with the intersection of two dynamic_dimension ranges if
-         * possible.
+         * possible. When both dimensions are symbolic, they are compatible only if they
+         * share the same symbolic expression.
          */
         std::optional<dynamic_dimension> intersection(const dynamic_dimension& other) const
         {
+            if(this->is_symbolic() and other.is_symbolic())
+            {
+                if(this->sym_expr == other.sym_expr)
+                    return *this;
+                return nullopt;
+            }
             auto this_interval  = this->get_interval();
             auto other_interval = other.get_interval();
             auto left           = std::max(this_interval.min, other_interval.min);
             auto right          = std::min(this_interval.max, other_interval.max);
             if(left <= right)
-            {
                 return dynamic_dimension{left, right};
-            }
             return nullopt;
         }
 
@@ -164,20 +206,24 @@ struct MIGRAPHX_EXPORT shape
         MIGRAPHX_EXPORT friend bool operator!=(const dynamic_dimension& x, const std::size_t& y);
         MIGRAPHX_EXPORT friend bool operator!=(const std::size_t& x, const dynamic_dimension& y);
 
-        // add, subtract, multiply fixed std::size_t dimension
-        dynamic_dimension& operator+=(const std::size_t& x);
-        dynamic_dimension& operator-=(const std::size_t& x);
-        dynamic_dimension& operator*=(const std::size_t& x);
-        MIGRAPHX_EXPORT friend dynamic_dimension operator+(const dynamic_dimension& x,
-                                                           const std::size_t& y);
-        MIGRAPHX_EXPORT friend dynamic_dimension operator+(const std::size_t& x,
-                                                           const dynamic_dimension& y);
-        MIGRAPHX_EXPORT friend dynamic_dimension operator-(const dynamic_dimension& x,
-                                                           const std::size_t& y);
-        MIGRAPHX_EXPORT friend dynamic_dimension operator*(const dynamic_dimension& x,
-                                                           const std::size_t& y);
-        MIGRAPHX_EXPORT friend dynamic_dimension operator*(const std::size_t& x,
-                                                           const dynamic_dimension& y);
+        // clang-format off
+// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
+#define MIGRAPHX_SHAPE_DYN_DIM_DEFINE_OP(binary_op, assign_op)                                  \
+        dynamic_dimension& operator assign_op(const dynamic_dimension& x);                      \
+        dynamic_dimension& operator assign_op(const std::size_t& x);                            \
+        MIGRAPHX_EXPORT friend dynamic_dimension operator binary_op(                            \
+            const dynamic_dimension& x, const dynamic_dimension& y);                            \
+        MIGRAPHX_EXPORT friend dynamic_dimension operator binary_op(                            \
+            const dynamic_dimension& x, const std::size_t& y);                                  \
+        MIGRAPHX_EXPORT friend dynamic_dimension operator binary_op(                            \
+            const std::size_t& x, const dynamic_dimension& y);
+
+        MIGRAPHX_SHAPE_DYN_DIM_DEFINE_OP(+, +=)
+        MIGRAPHX_SHAPE_DYN_DIM_DEFINE_OP(-, -=)
+        MIGRAPHX_SHAPE_DYN_DIM_DEFINE_OP(*, *=)
+        MIGRAPHX_SHAPE_DYN_DIM_DEFINE_OP(/, /=)
+#undef MIGRAPHX_SHAPE_DYN_DIM_DEFINE_OP
+        // clang-format on
     };
 
     static std::string to_sizes_string(const std::vector<shape>& shapes);
@@ -202,8 +248,10 @@ struct MIGRAPHX_EXPORT shape
     // Force all calls of the format `shape( type_t, { size_t compatibles } )` to map to
     // shape(type_t, std::vector<std::size_t> l)
     shape(type_t t, std::initializer_list<std::size_t> d);
+    shape(type_t t, std::initializer_list<std::size_t> l, std::initializer_list<std::size_t> s);
 
     shape(type_t t, std::vector<dynamic_dimension> dims);
+    shape(type_t t, std::vector<dynamic_dimension> dims, std::vector<sym::expr> dstrides);
 
     // Construct a dynamic shape from vectors of mins, maxes, and optimals.
     // optimals_list is a vector of optimals that corresponds to each min and max.
@@ -242,6 +290,9 @@ struct MIGRAPHX_EXPORT shape
      */
     static shape
     from_permutation(type_t t, const std::vector<std::size_t>& l, const std::vector<int64_t>& perm);
+    static shape from_permutation(type_t t,
+                                  const std::vector<dynamic_dimension>& dds,
+                                  const std::vector<int64_t>& perm);
 
     type_t type() const;
     const std::vector<std::size_t>& lens() const;
@@ -272,6 +323,9 @@ struct MIGRAPHX_EXPORT shape
 
     const std::vector<dynamic_dimension>& dyn_dims() const;
 
+    bool symbolic() const;
+    const std::vector<sym::expr>& dyn_strides() const;
+
     /*!
      * Minimum lengths for dynamic shape.
      * lens() for static shape.
@@ -388,14 +442,17 @@ struct MIGRAPHX_EXPORT shape
 
     shape with_lens(type_t t, const std::vector<std::size_t>& l) const;
     shape with_lens(const std::vector<std::size_t>& l) const;
+    shape with_lens(type_t t, const std::vector<dynamic_dimension>& dds) const;
+    shape with_lens(const std::vector<dynamic_dimension>& dds) const;
 
     shape with_type(type_t t) const;
 
-    // convert the shape to an equivalent dynamic shape with empty optimals
+    // convert the shape to an equivalent dynamic shape with constant symbolic strides
     shape to_dynamic() const;
 
     // convert the shape to a static one setting any non-fixed dynamic_dimensions to x
     shape to_static(std::size_t x) const;
+    shape to_static(const std::unordered_map<sym::expr, std::size_t>& symbol_map) const;
 
     MIGRAPHX_EXPORT friend bool operator==(const shape& x, const shape& y);
     MIGRAPHX_EXPORT friend bool operator!=(const shape& x, const shape& y);

src/include/migraphx/sym.hpp

@@ -28,8 +28,11 @@
 #include <cstdint>
 #include <memory>
 #include <ostream>
+#include <set>
 #include <string>
 #include <unordered_map>
+#include <utility>
+#include <variant>
 
 #include <migraphx/config.hpp>
 
@@ -40,8 +43,36 @@
 
 namespace sym {
 
+using scalar = std::variant<int64_t, double>;
+
+template <class To>
+To to(const scalar& v)
+{
+    return std::visit([](auto x) -> To { return x; }, v);
+}
+
+struct interval
+{
+    scalar min = int64_t{0};
+    scalar max = int64_t{0};
+
+    interval() = default;
+    interval(scalar mn, scalar mx) : min{std::move(mn)}, max{std::move(mx)} {}
+    // Convenience overload so brace-init with bare integer literals (e.g.
+    // `interval{1, 8}` or `var("n", {1, 8})`) resolves unambiguously rather
+    // than triggering the variant converting-ctor's int-vs-double tie that
+    // some libstdc++ versions (notably on SLES) reject.
+    interval(int64_t mn, int64_t mx) : min{mn}, max{mx} {}
+
+    friend bool operator==(const interval& a, const interval& b)
+    {
+        return a.min == b.min and a.max == b.max;
+    }
+    friend bool operator!=(const interval& a, const interval& b) { return not(a == b); }
+};
+
 struct expr;
-MIGRAPHX_EXPORT expr var(const std::string& name);
+MIGRAPHX_EXPORT expr var(const std::string& name, interval bounds, std::set<int64_t> optimals = {});
 MIGRAPHX_EXPORT expr lit(int64_t n);
 MIGRAPHX_EXPORT expr parse(const std::string& s);
 
@@ -50,11 +81,14 @@
     expr();
 
     bool empty() const;
+    bool is_literal() const;
     std::size_t hash() const;
     std::string to_string() const;
     value to_value() const;
     void from_value(const value& v);
     std::size_t eval_uint(const std::unordered_map<expr, std::size_t>& symbol_map) const;
+    interval eval_interval() const;
+    std::set<std::size_t> eval_optimals() const;
     expr subs(const std::unordered_map<expr, expr>& symbol_map) const;
 
     MIGRAPHX_EXPORT friend expr operator+(const expr& a, const expr& b);
@@ -63,6 +97,10 @@
     MIGRAPHX_EXPORT friend expr operator/(const expr& a, const expr& b);
     MIGRAPHX_EXPORT friend bool operator==(const expr& a, const expr& b);
     MIGRAPHX_EXPORT friend bool operator!=(const expr& a, const expr& b);
+    MIGRAPHX_EXPORT friend bool operator<(const expr& a, const expr& b);
+    friend bool operator>(const expr& a, const expr& b) { return b < a; }
+    friend bool operator<=(const expr& a, const expr& b) { return not(b < a); }
+    friend bool operator>=(const expr& a, const expr& b) { return not(a < b); }
     MIGRAPHX_EXPORT friend std::ostream& operator<<(std::ostream& os, const expr& e);
 
     friend expr operator+(const expr& a, int64_t b) { return a + lit(b); }
@@ -76,7 +114,8 @@
 
     struct impl;
 
-    MIGRAPHX_EXPORT friend expr var(const std::string& name);
+    MIGRAPHX_EXPORT friend expr
+    var(const std::string& name, interval bounds, std::set<int64_t> optimals);
     MIGRAPHX_EXPORT friend expr lit(int64_t n);
     MIGRAPHX_EXPORT friend expr parse(const std::string& s);
 

src/permutation.cpp

@@ -25,6 +25,7 @@
 #include <migraphx/permutation.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/algorithm.hpp>
+#include <migraphx/sym.hpp>
 #include <map>
 #include <functional>
 
@@ -33,6 +34,10 @@ inline namespace MIGRAPHX_INLINE_NS {
 
 shape reorder_shape(const shape& s, const std::vector<int64_t>& permutation)
 {
+    if(s.symbolic())
+        return {s.type(),
+                reorder_dims(s.dyn_dims(), permutation),
+                reorder_dims(s.dyn_strides(), permutation)};
     return {s.type(), reorder_dims(s.lens(), permutation), reorder_dims(s.strides(), permutation)};
 }
 
@@ -43,11 +48,50 @@ std::vector<int64_t> invert_permutation(const std::vector<int64_t>& permutation)
 
 std::vector<int64_t> find_permutation(const shape& s)
 {
-    std::vector<std::int64_t> result(s.lens().size());
+    if(s.dynamic() and not s.symbolic())
+        MIGRAPHX_THROW("FIND_PERMUTATION: non-symbolic dynamic shapes not supported");
+    std::vector<std::int64_t> result(s.ndim());
     std::iota(result.begin(), result.end(), 0);
-    std::stable_sort(result.begin(), result.end(), by(std::greater<>{}, [&](auto x) {
-                         return std::make_tuple(s.strides()[x], s.lens()[x]);
-                     }));
+    if(s.symbolic())
+    {
+        // Sort symbolic strides by evaluating at max variable values.
+        // Assumptions (see is_sorted_strides in shape.cpp for details):
+        //  1. Strides are products of dim variables * constant factors (no symbolic divisors)
+        //  2. Strides come from compute_strides() or permutations thereof
+        //  3. Max-eval ordering is consistent with all non-degenerate runtime orderings
+        const auto& strides = s.dyn_strides();
+        const auto& dds     = s.dyn_dims();
+        std::vector<sym::interval> stride_intervals(strides.size());
+        std::transform(strides.begin(), strides.end(), stride_intervals.begin(), [](const auto& e) {
+            return e.eval_interval();
+        });
+        std::vector<int64_t> dim_max(dds.size());
+        std::transform(dds.begin(), dds.end(), dim_max.begin(), [](const auto& dd) {
+            return sym::to<int64_t>(dd.sym_expr.eval_interval().max);
+        });
+        std::stable_sort(result.begin(), result.end(), by(std::greater<>{}, [&](auto x) {
+                             return std::make_tuple(sym::to<int64_t>(stride_intervals[x].max),
+                                                    dim_max[x]);
+                         }));
+        // Assumption 3 guard: when max-eval gives a strict ordering between two
+        // adjacent strides, min-eval must not reverse it. Collapse to equality at
+        // min is expected (e.g. when a dim has min=1), but a sign flip indicates
+        // a symbolic divisor violating assumption 1.
+        if(std::adjacent_find(result.begin(), result.end(), [&](auto a, auto b) {
+               return sym::to<int64_t>(stride_intervals[a].max) >
+                          sym::to<int64_t>(stride_intervals[b].max) and
+                      sym::to<int64_t>(stride_intervals[a].min) <
+                          sym::to<int64_t>(stride_intervals[b].min);
+           }) != result.end())
+            MIGRAPHX_THROW("FIND_PERMUTATION: symbolic stride ordering reversal between "
+                           "max-eval and min-eval. Violation of symbolic stride assumptions.");
+    }
+    else
+    {
+        std::stable_sort(result.begin(), result.end(), by(std::greater<>{}, [&](auto x) {
+                             return std::make_tuple(s.strides()[x], s.lens()[x]);
+                         }));
+    }
     return result;
 }
 
@@ -64,7 +108,7 @@ std::vector<int64_t> find_permutation(const std::vector<shape>& shapes)
     }
     if(count.empty())
     {
-        std::vector<int64_t> r(shapes.front().lens().size());
+        std::vector<int64_t> r(shapes.front().ndim());
         std::iota(r.begin(), r.end(), 0);
         return r;
     }