[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 {
@@ -112,42 +113,77 @@ 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();
+                if(ival.min < 0 or ival.max < 0)
+                    MIGRAPHX_THROW("dynamic_dimension: symbolic expression has negative bounds");
+                return {static_cast<std::size_t>(ival.min), static_cast<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 +200,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 +242,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 +284,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 +317,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 +436,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,6 +28,7 @@
 #include <cstdint>
 #include <memory>
 #include <ostream>
+#include <set>
 #include <string>
 #include <unordered_map>
 
@@ -40,8 +41,21 @@ struct value;
 
 namespace sym {
 
+struct interval
+{
+    int64_t min = 0;
+    int64_t max = 0;
+    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            = {1, 1},
+                         std::set<int64_t> optimals = {});
 MIGRAPHX_EXPORT expr lit(int64_t n);
 MIGRAPHX_EXPORT expr parse(const std::string& s);
 
@@ -50,11 +64,14 @@ struct MIGRAPHX_EXPORT expr
     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 +80,10 @@ struct MIGRAPHX_EXPORT expr
     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 +97,8 @@ struct MIGRAPHX_EXPORT expr
 
     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,25 @@ 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())
+    {
+        const auto& strides = s.dyn_strides();
+        const auto& dds     = s.dyn_dims();
+        std::stable_sort(result.begin(), result.end(), by(std::greater<>{}, [&](auto x) {
+                             return std::make_tuple(strides[x].eval_interval().max,
+                                                    dds[x].sym_expr.eval_interval().max);
+                         }));
+    }
+    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 +83,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;
     }