fix tests and error cases; allow std::optional for streamutils (#2362)

src/include/migraphx/op/allocate.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
@@ -41,16 +41,16 @@ namespace op {
  * Dynamic allocate:
  * One input: `allocate(output_dims)`
  * `output_dims` are the output buffer dimensions and has a static shape.
- * Either `this.s` or `this.buf_type` must be set to calculate the dynamic output shape at compute
- * time. If `this.buf_type` is set, the compute_shape() of allocate at compile time will have
- * dynamic_dimensions from {0, max_int} with rank = output_dims.ndim(). If `this.s` is set then the
- * compute_shape() will output `this.s`; `this.s` should be a dynamic shape.
+ * Either `this.s` or `this.buf_type` (but not both) must be set to calculate the dynamic output
+ * shape at compute time. If `this.buf_type` is set, the compute_shape() of allocate at compile time
+ * will have dynamic_dimensions from {0, max_int} with rank = output_dims.ndim(). If `this.s` is set
+ * then the compute_shape() will output `this.s`; `this.s` should be a dynamic shape.
  */
 struct allocate
 {
-    shape s{};
+    optional<shape> s;
     // for dynamic allocate to set the buffer type
-    shape::type_t buf_type = shape::half_type;
+    optional<shape::type_t> buf_type;
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
@@ -62,26 +62,38 @@ struct allocate
 
     shape compute_shape(const std::vector<shape>& inputs) const
     {
-        if(s != shape())
+        if(s.has_value())
         {
+            if(buf_type.has_value())
+            {
+                MIGRAPHX_THROW("ALLOCATE: shape and buf_type attributes both set");
+            }
             if(inputs.size() == 1)
             {
                 migraphx::check_shapes{inputs, *this, false}.only_dims(1);
             }
             else
             {
+                if(s->dynamic())
+                {
+                    MIGRAPHX_THROW("ALLOCATE: dynamic shape attribute and no input");
+                }
                 migraphx::check_shapes{inputs, *this, false}.has(0);
             }
-            return s;
+            return s.value();
         }
         else
         {
+            if(not buf_type.has_value())
+            {
+                MIGRAPHX_THROW("ALLOCATE: shape and buf_type attributes both not set");
+            }
             migraphx::check_shapes{inputs, *this, false}.has(1).only_dims(1);
             const auto& out_dims = inputs.at(0);
             std::size_t max_val = std::numeric_limits<std::size_t>::max();
             std::vector<shape::dynamic_dimension> dyn_dims(out_dims.lens().at(0),
                                                            shape::dynamic_dimension{0, max_val});
-            return {buf_type, dyn_dims};
+            return {buf_type.value(), dyn_dims};
         }
     }
     argument compute(const shape& output_shape, const std::vector<argument>& args) const
@@ -94,7 +106,11 @@ struct allocate
         {
             std::vector<std::size_t> output_dims(output_shape.ndim());
             args.at(0).visit([&](auto a) { output_dims.assign(a.begin(), a.end()); });
-            return argument{shape{buf_type, output_dims}};
+            if(s)
+            {
+                return argument{shape{s->type(), output_dims}};
+            }
+            return argument{shape{buf_type.value(), output_dims}};
         }
     }
 };

src/include/migraphx/streamutils.hpp

@@ -30,6 +30,7 @@
 #include <migraphx/rank.hpp>
 #include <migraphx/requires.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/optional.hpp>
 #include <vector>
 
 namespace migraphx {
@@ -68,6 +69,19 @@ auto stream_write_value_impl(rank<1>, std::ostream& os, const T& x) -> decltype(
     os << x;
 }
 
+template <class T>
+auto stream_write_value_impl(rank<1>, std::ostream& os, const optional<T>& x)
+{
+    if(x.has_value())
+    {
+        os << *x;
+    }
+    else
+    {
+        os << "nullopt";
+    }
+}
+
 template <class T>
 void stream_write_value_impl(rank<1>, std::ostream& os, const std::vector<T>& r)
 {

test/op_shape_test.cpp

@@ -88,7 +88,7 @@ TEST_CASE(allocate_static)
     expect_shape(out_shape, migraphx::make_op("allocate", {{"shape", to_value(out_shape)}}));
 }
 
-TEST_CASE(allocate_static_input_error)
+TEST_CASE(allocate_static_input)
 {
     migraphx::shape input{migraphx::shape::int64_type, {3}};
     migraphx::shape out_shape{migraphx::shape::float_type, {2, 3, 4}};
@@ -120,8 +120,22 @@ TEST_CASE(allocate_dyn_no_input_error)
 {
     migraphx::shape shape_attr{migraphx::shape::float_type,
                                {{1, 4}, {3, 3}, {4, 8, {4, 6}}, {4, 8}, {4, 6}}};
-    expect_shape(shape_attr,
-                 migraphx::make_op("allocate", {{"shape", migraphx::to_value(shape_attr)}}));
+    throws_shape(migraphx::make_op("allocate", {{"shape", migraphx::to_value(shape_attr)}}));
+}
+
+TEST_CASE(allocate_shape_and_buf_type_error)
+{
+    migraphx::shape shape_attr{migraphx::shape::float_type,
+                               {{1, 4}, {3, 3}, {4, 8, {4, 6}}, {4, 8}, {4, 6}}};
+    throws_shape(migraphx::make_op(
+        "allocate",
+        {{"shape", migraphx::to_value(shape_attr)}, {"buf_type", migraphx::shape::half_type}}));
+}
+
+TEST_CASE(allocate_no_attr_error)
+{
+    migraphx::shape input{migraphx::shape::int64_type, {4}};
+    throws_shape(migraphx::make_op("allocate"), input);
 }
 
 TEST_CASE(argmax_axis0)

test/ref/allocate.cpp

@@ -30,7 +30,7 @@
 
 #include <test.hpp>
 
-TEST_CASE(allocate_dyn)
+TEST_CASE(allocate_dyn0)
 {
     migraphx::program p;
     auto* mm = p.get_main_module();
@@ -47,3 +47,21 @@ TEST_CASE(allocate_dyn)
     migraphx::shape sresult{migraphx::shape::float_type, {2, 3, 4, 4}};
     result.visit([&](auto output) { EXPECT(output.get_shape() == sresult); });
 }
+
+TEST_CASE(allocate_dyn1)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::int64_type, {4}};
+    migraphx::shape out_shape{migraphx::shape::float_type, {2, 3, 4, 4}};
+    auto out_dims = mm->add_parameter("out_dims", s);
+    mm->add_instruction(migraphx::make_op("allocate", {{"shape", migraphx::to_value(out_shape)}}),
+                        out_dims);
+    p.compile(migraphx::make_target("ref"));
+
+    migraphx::parameter_map params;
+    std::vector<int64_t> data = {2, 3, 4, 4};
+    params["out_dims"]        = migraphx::argument(s, data.data());
+    auto result               = p.eval(params).back();
+    result.visit([&](auto output) { EXPECT(output.get_shape() == out_shape); });
+}