Python API update for dynamic batch (#1723)

Python API with documentation updates

doc/src/reference/py.rst

@@ -6,9 +6,9 @@ Python Reference
 shape
 -----
 
-.. py:class:: shape(type, lens, strides=None)
+.. py:class:: shape(type, lens, strides=None, dyn_dims)
 
-    Describes the shape of a tensor. This includes size, layout, and data type/
+    Describes the shape of a tensor. This includes size, layout, and data type. Can be a dynamic shape by using dyn_dims.
 
 .. py:method:: type()
 
@@ -34,6 +34,12 @@ shape
 
     :rtype: int
 
+.. py:method:: dyn_dims()
+
+    The dynamic dimensions of the shape.
+
+    :rtype: list[dynamic_dimension]
+
 .. py:method:: bytes()
 
     The number of bytes the shape uses.
@@ -46,6 +52,12 @@ shape
 
     :rtype: int
 
+.. py:method:: ndim()
+
+    The number of dimensions for the shape.
+
+    :rtype: int
+
 .. py:method:: packed()
 
     Returns true if the shape is packed.
@@ -64,6 +76,12 @@ shape
 
     :rtype: bool
 
+.. py:method:: dynamic()
+
+    Returns true if the shape is dynamic.
+
+    :rtype: bool
+
 .. py:method:: standard()
 
     Returns true if the shape is a standard shape. That is, the shape is both packed and not transposed.
@@ -76,6 +94,18 @@ shape
 
     :rtype: bool
 
+dynamic_dimension
+--------
+
+.. py:class:: dynamic_dimension(min, max, optimals)
+
+    Construct a dynamic_dimension from a minimum, a maximum, and optionally a set of optimals.
+
+.. py:method:: is_fixed()
+    
+    Returns true if the dynamic_dimension is fixed.
+
+    :rtype : int
 
 argument
 --------
@@ -292,8 +322,10 @@ parse_onnx
     Load and parse an onnx file.
 
     :param str filename: Path to file.
-    :param str default_dim_value: default batch size to use (if not specified in onnx file).
+    :param str default_dim_value: default dimension to use (if not specified in onnx file).
+    :param dynamic_dimension default_dyn_dim_value: default dynamic_dimension value to use.
     :param str map_input_dims: Explicitly specify the dims of an input.
+    :param list[dynamic_dimension] map_dyn_input_dims: Explicitly specify the dynamic_dimensions of an input.
     :param str skip_unknown_operators: Continue parsing onnx file if an unknown operator is found.
     :param str print_program_on_error: Print program if an error occurs.
     :param int max_loop_iterations: Maximum iteration number for the loop operator.

src/py/migraphx_py.cpp

@@ -95,6 +95,10 @@ void visit_py(T x, F f)
     {
         f(x.template cast<std::string>());
     }
+    else if(py::isinstance<migraphx::shape::dynamic_dimension>(x))
+    {
+        f(migraphx::to_value(x.template cast<migraphx::shape::dynamic_dimension>()));
+    }
     else
     {
         MIGRAPHX_THROW("VISIT_PY: Unsupported data type!");
@@ -165,7 +169,10 @@ template <class T>
 py::buffer_info to_buffer_info(T& x)
 {
     migraphx::shape s = x.get_shape();
-    auto strides      = s.strides();
+    assert(s.type() != migraphx::shape::tuple_type);
+    if(s.dynamic())
+        MIGRAPHX_THROW("MIGRAPHX PYTHON: dynamic shape argument passed to to_buffer_info");
+    auto strides = s.strides();
     std::transform(
         strides.begin(), strides.end(), strides.begin(), [&](auto i) { return i * s.type_size(); });
     py::buffer_info b;
@@ -177,7 +184,7 @@ py::buffer_info to_buffer_info(T& x)
             b = py::buffer_info(x.data(),
                                 as.size(),
                                 py::format_descriptor<bool>::format(),
-                                s.lens().size(),
+                                s.ndim(),
                                 s.lens(),
                                 strides);
         }
@@ -186,7 +193,7 @@ py::buffer_info to_buffer_info(T& x)
             b = py::buffer_info(x.data(),
                                 as.size(),
                                 py::format_descriptor<decltype(as())>::format(),
-                                s.lens().size(),
+                                s.ndim(),
                                 s.lens(),
                                 strides);
         }
@@ -239,26 +246,39 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
     py::class_<migraphx::shape> shape_cls(m, "shape");
     shape_cls
         .def(py::init([](py::kwargs kwargs) {
-            auto v    = migraphx::to_value(kwargs);
-            auto t    = migraphx::shape::parse_type(v.get("type", "float"));
-            auto lens = v.get<std::size_t>("lens", {1});
-            if(v.contains("strides"))
-                return migraphx::shape(t, lens, v.at("strides").to_vector<std::size_t>());
+            auto v = migraphx::to_value(kwargs);
+            auto t = migraphx::shape::parse_type(v.get("type", "float"));
+            if(v.contains("dyn_dims"))
+            {
+                auto dyn_dims =
+                    migraphx::from_value<std::vector<migraphx::shape::dynamic_dimension>>(
+                        v.at("dyn_dims"));
+                return migraphx::shape(t, dyn_dims);
+            }
             else
-                return migraphx::shape(t, lens);
+            {
+                auto lens = v.get<std::size_t>("lens", {1});
+                if(v.contains("strides"))
+                    return migraphx::shape(t, lens, v.at("strides").to_vector<std::size_t>());
+                else
+                    return migraphx::shape(t, lens);
+            }
         }))
         .def("type", &migraphx::shape::type)
         .def("lens", &migraphx::shape::lens)
         .def("strides", &migraphx::shape::strides)
+        .def("ndim", &migraphx::shape::ndim)
         .def("elements", &migraphx::shape::elements)
         .def("bytes", &migraphx::shape::bytes)
         .def("type_string", &migraphx::shape::type_string)
         .def("type_size", &migraphx::shape::type_size)
+        .def("dyn_dims", &migraphx::shape::dyn_dims)
         .def("packed", &migraphx::shape::packed)
         .def("transposed", &migraphx::shape::transposed)
         .def("broadcasted", &migraphx::shape::broadcasted)
         .def("standard", &migraphx::shape::standard)
         .def("scalar", &migraphx::shape::scalar)
+        .def("dynamic", &migraphx::shape::dynamic)
         .def("__eq__", std::equal_to<migraphx::shape>{})
         .def("__ne__", std::not_equal_to<migraphx::shape>{})
         .def("__repr__", [](const migraphx::shape& s) { return migraphx::to_string(s); });
@@ -266,6 +286,15 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
     py::enum_<migraphx::shape::type_t>(shape_cls, "type_t")
         MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_PYTHON_GENERATE_SHAPE_ENUM);
 
+    py::class_<migraphx::shape::dynamic_dimension>(shape_cls, "dynamic_dimension")
+        .def(py::init<>())
+        .def(py::init<std::size_t, std::size_t>())
+        .def(py::init<std::size_t, std::size_t, std::set<std::size_t>>())
+        .def_readwrite("min", &migraphx::shape::dynamic_dimension::min)
+        .def_readwrite("max", &migraphx::shape::dynamic_dimension::max)
+        .def_readwrite("optimals", &migraphx::shape::dynamic_dimension::optimals)
+        .def("is_fixed", &migraphx::shape::dynamic_dimension::is_fixed);
+
     py::class_<migraphx::argument>(m, "argument", py::buffer_protocol())
         .def_buffer([](migraphx::argument& x) -> py::buffer_info { return to_buffer_info(x); })
         .def(py::init([](py::buffer b) {
@@ -440,13 +469,18 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
         "parse_onnx",
         [](const std::string& filename,
            unsigned int default_dim_value,
+           migraphx::shape::dynamic_dimension default_dyn_dim_value,
            std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
+           std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>
+               map_dyn_input_dims,
            bool skip_unknown_operators,
            bool print_program_on_error,
            int64_t max_loop_iterations) {
             migraphx::onnx_options options;
             options.default_dim_value      = default_dim_value;
+            options.default_dyn_dim_value  = default_dyn_dim_value;
             options.map_input_dims         = map_input_dims;
+            options.map_dyn_input_dims     = map_dyn_input_dims;
             options.skip_unknown_operators = skip_unknown_operators;
             options.print_program_on_error = print_program_on_error;
             options.max_loop_iterations    = max_loop_iterations;
@@ -454,8 +488,11 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
         },
         "Parse onnx file",
         py::arg("filename"),
-        py::arg("default_dim_value") = 1,
-        py::arg("map_input_dims")    = std::unordered_map<std::string, std::vector<std::size_t>>(),
+        py::arg("default_dim_value")     = 0,
+        py::arg("default_dyn_dim_value") = migraphx::shape::dynamic_dimension{1, 1},
+        py::arg("map_input_dims") = std::unordered_map<std::string, std::vector<std::size_t>>(),
+        py::arg("map_dyn_input_dims") =
+            std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>(),
         py::arg("skip_unknown_operators") = false,
         py::arg("print_program_on_error") = false,
         py::arg("max_loop_iterations")    = 10);
@@ -464,20 +501,28 @@ MIGRAPHX_PYBIND11_MODULE(migraphx, m)
         "parse_onnx_buffer",
         [](const std::string& onnx_buffer,
            unsigned int default_dim_value,
+           migraphx::shape::dynamic_dimension default_dyn_dim_value,
            std::unordered_map<std::string, std::vector<std::size_t>> map_input_dims,
+           std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>
+               map_dyn_input_dims,
            bool skip_unknown_operators,
            bool print_program_on_error) {
             migraphx::onnx_options options;
             options.default_dim_value      = default_dim_value;
+            options.default_dyn_dim_value  = default_dyn_dim_value;
             options.map_input_dims         = map_input_dims;
+            options.map_dyn_input_dims     = map_dyn_input_dims;
             options.skip_unknown_operators = skip_unknown_operators;
             options.print_program_on_error = print_program_on_error;
             return migraphx::parse_onnx_buffer(onnx_buffer, options);
         },
         "Parse onnx file",
         py::arg("filename"),
-        py::arg("default_dim_value") = 1,
-        py::arg("map_input_dims")    = std::unordered_map<std::string, std::vector<std::size_t>>(),
+        py::arg("default_dim_value")     = 0,
+        py::arg("default_dyn_dim_value") = migraphx::shape::dynamic_dimension{1, 1},
+        py::arg("map_input_dims") = std::unordered_map<std::string, std::vector<std::size_t>>(),
+        py::arg("map_dyn_input_dims") =
+            std::unordered_map<std::string, std::vector<migraphx::shape::dynamic_dimension>>(),
         py::arg("skip_unknown_operators") = false,
         py::arg("print_program_on_error") = false);
 

src/shape.cpp

@@ -706,14 +706,10 @@ void migraphx_from_value(const value& v, shape& s)
         {
             auto v_dd = v.at("dynamic_dimensions");
             std::vector<shape::dynamic_dimension> dyn_dims(v.at("dynamic_dimensions").size());
-            std::transform(v_dd.begin(), v_dd.end(), dyn_dims.begin(), [](migraphx::value x) {
-                auto x_min      = x.at("min").template to<size_t>();
-                auto x_max      = x.at("max").template to<size_t>();
-                auto v_optimals = x.at("optimals");
-                std::set<size_t> set_x_optimals =
-                    from_value<std::set<std::size_t>>(x.at("optimals"));
-                return shape::dynamic_dimension{x_min, x_max, set_x_optimals};
-            });
+            std::transform(
+                v_dd.begin(), v_dd.end(), dyn_dims.begin(), [](const migraphx::value& x) {
+                    return from_value<shape::dynamic_dimension>(x);
+                });
 
             s = shape{shape::parse_type(t), dyn_dims};
         }

test/py/test_gpu.py

@@ -86,8 +86,8 @@ def test_nonzero():
     params = {}
 
     shapes = p.get_parameter_shapes()
-    params["data"] = np.array([1, 1, 0, 1]).reshape(
-        shapes["data"].lens()).astype(np.bool)
+    params["data"] = np.array([1, 1, 0,
+                               1]).reshape(shapes["data"].lens()).astype(bool)
 
     r = p.run(params)
     print(r)
@@ -127,15 +127,54 @@ def test_if_pl():
     params["x"] = np.ones(6).reshape(shapes["x"].lens()).astype(np.float32)
     params["y"] = np.array([2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0
                             ]).reshape(shapes["y"].lens()).astype(np.float32)
-    params["cond"] = np.array([1]).reshape(()).astype(np.bool)
+    params["cond"] = np.array([1]).reshape(()).astype(bool)
 
     r = p.run(params)[-1]
     print(r)
 
 
+def test_dyn_batch():
+    a = migraphx.shape.dynamic_dimension(1, 4, {2, 4})
+    b = migraphx.shape.dynamic_dimension(3, 3)
+    c = migraphx.shape.dynamic_dimension(32, 32)
+    dd_map = {"0": [a, b, c, c]}
+    p = migraphx.parse_onnx("conv_relu_maxpool_test.onnx",
+                            map_dyn_input_dims=dd_map)
+    print(p)
+    print("Compiling ...")
+    p.compile(migraphx.get_target("gpu"))
+    print(p)
+
+    def run_prog(batch_size):
+        params = {}
+        for key, value in p.get_parameter_shapes().items():
+            # convert to a static shape
+            if value.dynamic():
+                dds = value.dyn_dims()
+                new_lens = []
+                for dd in dds:
+                    if dd.is_fixed():
+                        new_lens.append(dd.min)
+                    else:
+                        new_lens.append(batch_size)
+                s = migraphx.shape(type=value.type_string(), lens=new_lens)
+            else:
+                s = value
+            print("Parameter {} -> {}".format(key, s))
+            params[key] = migraphx.generate_argument(s)
+        r = p.run(params)
+        print(r)
+
+    run_prog(1)
+    run_prog(2)
+    run_prog(3)
+    run_prog(4)
+
+
 test_conv_relu()
 test_sub_uint64()
 test_neg_int64()
 test_fp16_imagescaler()
 test_if_pl()
 test_nonzero()
+test_dyn_batch()

test/py/test_gpu_offload.py

@@ -23,16 +23,53 @@
 #####################################################################################
 import migraphx
 
-p = migraphx.parse_onnx("conv_relu_maxpool_test.onnx")
-print(p)
-print("Compiling ...")
-p.compile(migraphx.get_target("gpu"), offload_copy=False)
-print(p)
-params = {}
 
-for key, value in p.get_parameter_shapes().items():
-    print("Parameter {} -> {}".format(key, value))
-    params[key] = migraphx.to_gpu(migraphx.generate_argument(value))
+def test_conv_relu():
+    p = migraphx.parse_onnx("conv_relu_maxpool_test.onnx")
+    print(p)
+    print("Compiling ...")
+    p.compile(migraphx.get_target("gpu"), offload_copy=False)
+    print(p)
+    params = {}
 
-r = migraphx.from_gpu(p.run(params)[-1])
-print(r)
+    for key, value in p.get_parameter_shapes().items():
+        print("Parameter {} -> {}".format(key, value))
+        params[key] = migraphx.to_gpu(migraphx.generate_argument(value))
+
+    r = migraphx.from_gpu(p.run(params)[-1])
+    print(r)
+
+
+# TODO: placeholder until tuple shapes and arguments exposed
+#def test_dyn_batch():
+#    a = migraphx.shape.dynamic_dimension(1, 4, {2, 4})
+#    b = migraphx.shape.dynamic_dimension(3, 3)
+#    c = migraphx.shape.dynamic_dimension(32, 32)
+#    dd_map = {"0": [a, b, c, c]}
+#    p = migraphx.parse_onnx("conv_relu_maxpool_test.onnx",
+#                            map_dyn_input_dims=dd_map)
+#    print(p)
+#    print("Compiling ...")
+#    p.compile(migraphx.get_target("gpu"), offload_copy=False)
+#
+#    print(p)
+#
+#    def run_prog(batch_size):
+#        params = {}
+#        for key, value in p.get_parameter_shapes().items():
+#            print("Parameter {} -> {}".format(key, value))
+#            params[key] = migraphx.to_gpu(
+#                migraphx.generate_argument(value.to_static(batch_size)))
+#
+#        print("before_output")
+#        outputs = p.run(params)
+#        print(outputs)
+#        r = migraphx.from_gpu(p.run(params)[-1])
+#        print(r)
+#
+#    run_prog(1)
+#    run_prog(2)
+#    run_prog(3)
+#    run_prog(4)
+
+test_conv_relu()

test/py/test_shape.py

@@ -29,6 +29,7 @@ def test_create_shape():
     assert s.standard()
     assert s.packed()
     assert s.lens() == [1, 64, 3, 3]
+    assert s.ndim() == 4
 
 
 def test_create_shape_broadcast():
@@ -49,6 +50,35 @@ def test_create_shape_type():
     assert s.type_size() == 4
 
 
+def test_create_dyn_dims():
+    a = migraphx.shape.dynamic_dimension()
+    assert a.is_fixed()
+    assert a.min == 0
+    b = migraphx.shape.dynamic_dimension(4, 4)
+    assert b.is_fixed()
+    assert b.max == 4
+    c = migraphx.shape.dynamic_dimension(1, 4, {2, 4})
+    assert not c.is_fixed()
+    assert c.min == 1
+    assert c.max == 4
+    assert c.optimals == {2, 4}
+
+    dyn_dims = [a, b]
+    dyn_dims.append(c)
+    assert dyn_dims[1] == b
+
+
+def test_create_dyn_shape():
+    a = migraphx.shape.dynamic_dimension(1, 4, {2, 4})
+    b = migraphx.shape.dynamic_dimension(4, 4)
+    dds = [a, b]
+    dyn_shape = migraphx.shape(type='float', dyn_dims=dds)
+    assert dyn_shape.dynamic()
+    assert dyn_shape.dyn_dims()[0].min == dds[0].min
+    assert dyn_shape.dyn_dims()[0].max == dds[0].max
+    assert dyn_shape.dyn_dims()[0].optimals == dds[0].optimals
+
+
 def test_type_enum():
     mgx_types = [
         'bool_type', 'double_type', 'float_type', 'half_type', 'int16_type',
@@ -63,3 +93,5 @@ if __name__ == "__main__":
     test_create_shape()
     test_create_shape_broadcast()
     test_create_shape_type()
+    test_create_dyn_dims()
+    test_create_dyn_shape()

test/shape_test.cpp

@@ -201,6 +201,20 @@ TEST_CASE(dynamic_dimension_add_sub_fixed)
     EXPECT((2 + e) == d);
 }
 
+TEST_CASE(dynamic_dimension_serialize)
+{
+    using migraphx::shape;
+    auto a  = shape::dynamic_dimension{2, 5, {2, 3}};
+    auto b  = shape::dynamic_dimension{3, 6, {3}};
+    auto v1 = migraphx::to_value(a);
+    auto v2 = migraphx::to_value(b);
+    EXPECT(v1 != v2);
+    auto c = migraphx::from_value<shape::dynamic_dimension>(v1);
+    EXPECT(a == c);
+    auto d = migraphx::from_value<shape::dynamic_dimension>(v2);
+    EXPECT(b == d);
+}
+
 TEST_CASE(test_shape_dynamic_errors)
 {
     using migraphx::shape;