Refactor dynamic_dimension fixed compare (#1470)

Implements the operator==(dynamic_dimension, size_t) functions

src/common.cpp

@@ -77,7 +77,6 @@ std::vector<shape::dynamic_dimension> compute_broadcasted_dyn_dims(shape s0, sha
     }
     auto offset = s1.ndim() - s0.ndim();
     std::vector<shape::dynamic_dimension> out_dims(s1.dyn_dims());
-    shape::dynamic_dimension one_dyn_dim{1, 1, 0};
     std::transform(
         s0.dyn_dims().cbegin(),
         s0.dyn_dims().cend(),
@@ -88,7 +87,7 @@ std::vector<shape::dynamic_dimension> compute_broadcasted_dyn_dims(shape s0, sha
             {
                 return a;
             }
-            else if(a == one_dyn_dim or b == one_dyn_dim)
+            else if(a == 1 or b == 1)
             {
                 // setting opt to 0, may need to be changed
                 return shape::dynamic_dimension{std::max(a.min, b.min), std::max(a.max, b.max), 0};

src/include/migraphx/op/squeeze.hpp

@@ -59,9 +59,8 @@ struct squeeze
         auto input_shape = inputs[0];
         if(input_shape.dynamic())
         {
-            shape::dynamic_dimension one_dyn_dim{1, 1, 0};
             if(std::any_of(axes.begin(), axes.end(), [&](auto axis) {
-                   return input_shape.dyn_dims()[axis] != one_dyn_dim;
+                   return input_shape.dyn_dims()[axis] != 1;
                }))
             {
                 MIGRAPHX_THROW(
@@ -73,7 +72,7 @@ struct squeeze
                 std::copy_if(input_shape.dyn_dims().cbegin(),
                              input_shape.dyn_dims().cend(),
                              std::back_inserter(dyn_dims),
-                             [&](auto dd) { return dd != one_dyn_dim; });
+                             [&](auto dd) { return dd != 1; });
             }
             else
             {

src/include/migraphx/shape.hpp

@@ -101,6 +101,12 @@ struct shape
         friend bool operator==(const dynamic_dimension& x, const dynamic_dimension& y);
         friend bool operator!=(const dynamic_dimension& x, const dynamic_dimension& y);
         friend std::ostream& operator<<(std::ostream& os, const dynamic_dimension& x);
+
+        // compare to fixed std::size_t dimension
+        friend bool operator==(const dynamic_dimension& x, const std::size_t& y);
+        friend bool operator==(const std::size_t& x, const dynamic_dimension& y);
+        friend bool operator!=(const dynamic_dimension& x, const std::size_t& y);
+        friend bool operator!=(const std::size_t& x, const dynamic_dimension& y);
     };
 
     static const std::vector<type_t>& types();

src/shape.cpp

@@ -521,6 +521,14 @@ std::ostream& operator<<(std::ostream& os, const shape::dynamic_dimension& x)
     return os;
 }
 
+bool operator==(const shape::dynamic_dimension& x, const std::size_t& y)
+{
+    return x.min == y and x.max == y;
+}
+bool operator==(const std::size_t& x, const shape::dynamic_dimension& y) { return y == x; }
+bool operator!=(const shape::dynamic_dimension& x, const std::size_t& y) { return not(x == y); }
+bool operator!=(const std::size_t& x, const shape::dynamic_dimension& y) { return not(x == y); }
+
 bool operator==(const shape& x, const shape& y)
 {
     if(x.dynamic() and y.dynamic())

test/shape_test.cpp

@@ -160,6 +160,20 @@ TEST_CASE(test_shape_dynamic_compares)
     EXPECT(ss0.str() != ss3.str());
 }
 
+TEST_CASE(dynamic_dimension_size_t_compares)
+{
+    using migraphx::shape;
+    auto a = shape::dynamic_dimension{2, 2, 2};
+    EXPECT(a == 2);
+    EXPECT(a != 3);
+    EXPECT(static_cast<std::size_t>(2) == a);
+    EXPECT(static_cast<std::size_t>(3) != a);
+
+    auto b = shape::dynamic_dimension{2, 4, 0};
+    EXPECT(b != 2);
+    EXPECT(static_cast<std::size_t>(2) != b);
+}
+
 TEST_CASE(test_shape_dynamic_errors)
 {
     using migraphx::shape;