[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 <type_traits>
 #include <unordered_map>
+#include <variant>
 
 #include <migraphx/config.hpp>
 
@@ -40,8 +43,48 @@ struct value;
 
 namespace sym {
 
+// Scalar value held by literal expressions and interval bounds. Wraps a
+// variant so that integer-literal initialization is unambiguous on stricter
+// libstdc++ versions.
+struct scalar
+{
+    std::variant<int64_t, double> value;
+
+    constexpr scalar() = default;
+
+    template <class T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
+    constexpr scalar(T v) : value{int64_t{v}} // NOLINT(google-explicit-constructor)
+    {
+    }
+
+    template <class T, std::enable_if_t<std::is_floating_point_v<T>, int> = 0>
+    constexpr scalar(T v) : value{double{v}} // NOLINT(google-explicit-constructor)
+    {
+    }
+
+    friend bool operator==(const scalar& a, const scalar& b) { return a.value == b.value; }
+    friend bool operator!=(const scalar& a, const scalar& b) { return not(a == b); }
+};
+
+template <class To>
+To to(const scalar& v)
+{
+    return std::visit([](auto x) -> To { return x; }, v.value);
+}
+
+struct interval
+{
+    scalar min = int64_t{0};
+    scalar max = int64_t{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, std::set<int64_t> optimals = {});
 MIGRAPHX_EXPORT expr lit(int64_t n);
 MIGRAPHX_EXPORT expr parse(const std::string& s);
 
@@ -50,11 +93,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 +109,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 +126,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);