Dynamic shape support in prefix scan (#1921)

* Add dynamic input to prefix_scan_op

* Added a shape test.  The op should return the same dynamic shape as the input.

* add 2d shape test for prefix_scan

src/include/migraphx/op/prefix_scan_op.hpp

 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 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
@@ -21,6 +21,7 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  */
+
 #ifndef MIGRAPHX_GUARD_OPERATORS_SCAN_OP_HPP
 #define MIGRAPHX_GUARD_OPERATORS_SCAN_OP_HPP
 
@@ -37,6 +38,12 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
+/**
+ * Parent struct for prefix scan operations.  A prefix scan is equivalent to the C++
+ * std::exclusive_scan or std::inclusive_scan.  Given a list of numbers, a prefix scan
+ * sum op returns an equal size list of running totals of the values.  Other operations
+ * besides addition can be supported by their own child ops.
+ */
 template <class Derived>
 struct prefix_scan_op : op_name<Derived>
 {
@@ -60,9 +67,13 @@ struct prefix_scan_op : op_name<Derived>
 
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
+        check_shapes{inputs, *this, true}.has(1);
         auto s = inputs.front();
-        if(s.broadcasted())
+        if(s.dynamic())
+        {
+            return s;
+        }
+        else if(s.broadcasted())
         {
             return {s.type(), s.lens()};
         }
@@ -72,8 +83,9 @@ struct prefix_scan_op : op_name<Derived>
         }
     }
 
-    argument compute(const shape& output_shape, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
+        shape output_shape(dyn_out.computed_shape);
         argument result{output_shape};
         auto s = args[0].get_shape();
         if(s == output_shape)

test/op_shape_test.cpp

@@ -2190,6 +2190,32 @@ TEST_CASE(prefix_scan_sum)
     }
 }
 
+TEST_CASE(prefix_scan_sum_dyn)
+{
+    {
+        std::vector<migraphx::shape::dynamic_dimension> dd{{5, 8}};
+        migraphx::shape s{migraphx::shape::float_type, dd};
+
+        expect_shape(
+            s,
+            migraphx::make_op("prefix_scan_sum", {{"axis", 0}, {"exclusive", 0}, {"reverse", 0}}),
+            s);
+    }
+}
+
+TEST_CASE(prefix_scan_sum_dyn_2d)
+{
+    {
+        std::vector<migraphx::shape::dynamic_dimension> dd{{5, 8}, {3, 7}};
+        migraphx::shape s{migraphx::shape::float_type, dd};
+
+        expect_shape(
+            s,
+            migraphx::make_op("prefix_scan_sum", {{"axis", 1}, {"exclusive", 0}, {"reverse", 0}}),
+            s);
+    }
+}
+
 TEST_CASE(quant_convolution_shape)
 {
     migraphx::shape output{migraphx::shape::int32_type, {4, 4, 1, 1}};

test/ref_ops_test.cpp

@@ -5886,6 +5886,29 @@ TEST_CASE(prefix_scan_sum_1d)
     EXPECT(results_vector == gold);
 }
 
+TEST_CASE(prefix_scan_sum_dyn_1d)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    std::vector<migraphx::shape::dynamic_dimension> dd{{5, 8}};
+    migraphx::shape s{migraphx::shape::float_type, dd};
+    auto input = mm->add_parameter("X", s);
+    mm->add_instruction(migraphx::make_op("prefix_scan_sum", {{"axis", 0}, {"exclusive", false}}),
+                        input);
+    p.compile(migraphx::make_target("ref"));
+
+    std::vector<float> a = {1, 2, 3, 4, 5, 6};
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {6}};
+    migraphx::parameter_map params0;
+    params0["X"] = migraphx::argument(input_fixed_shape0, a.data());
+
+    auto result = p.eval(params0).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold{1.0, 3.0, 6.0, 10.0, 15.0, 21.0};
+    EXPECT(results_vector == gold);
+}
+
 TEST_CASE(prefix_scan_sum_2d)
 {
     {