Add dynamic input to prefix_scan_op

src/include/migraphx/op/prefix_scan_op.hpp

@@ -60,8 +60,12 @@ 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.dynamic())
+        {
+            return {s.type(), s.max_lens()};
+        }
         if(s.broadcasted())
         {
             return {s.type(), s.lens()};
@@ -72,8 +76,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/ref_ops_test.cpp

@@ -5539,6 +5539,32 @@ 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, 6}};
+    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::parameter_map params0;    
+
+    // auto input = migraphx::literal{s, {1, 2, 3, 4, 5, 6}};
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {6}};
+    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)
 {
     {