Merge branch 'dyn_unsqueeze' of github.com:ROCmSoftwarePlatform/AMDMIGraphX into dyn_onnx_matmul

src/eliminate_contiguous.cpp

@@ -42,6 +42,13 @@ static bool try_compute_shape(instruction_ref ins,
     try
     {
         shape new_shape = ins->get_operator().compute_shape(inputs, mods);
+
+        // Cannot tell if a dynamic shape will need to be made contiguous
+        if(new_shape.dynamic())
+        {
+            return false;
+        }
+
         // If the output shape is a standard shape, no need to try its output
         if(new_shape.standard())
         {
@@ -133,14 +140,20 @@ static void remove_contiguous(const std::string& op_name, module& m, F f)
         }
     }
 
-    // Perform evaluations in parallel
+    // Perform static contiguous evaluations in parallel
     std::vector<argument> literals(const_instructions.size());
     par_for(const_instructions.size(), 1, [&](const auto i) {
         auto c    = op::contiguous{};
         auto prev = const_instructions[i]->inputs().front();
-        literals[i] = c.compute(c.compute_shape({prev->get_shape()}), {prev->eval()});
+        // compute the output contiguous shape from the previous instruction shape
+        shape computed_shape                   = c.compute_shape({prev->get_shape()});
+        const std::vector<argument>& prev_eval = {prev->eval()};
+        // prev_eval should not be used in make_compute_output_shape() as computed_shape is static
+        auto co_shape = make_compute_output_shape(pack(c, computed_shape, prev_eval));
+        literals[i]   = c.compute(co_shape, prev_eval);
     });
 
+    // Replace static contiguous operations with a literal
     for(size_t i = 0; i < const_instructions.size(); i++)
     {
         auto l = m.add_literal(literals[i].get_shape(), literals[i].data());

src/fuse_pointwise.cpp

@@ -45,7 +45,16 @@ static literal get_scalar(instruction_ref ins)
         return {};
     auto e = ins->eval();
     literal r{};
+    // needed for bool as visit_at invokes as() which promotes bool to int8
+    // Without this we'll break type checks for logical ops that are fused.
+    if(e.get_shape().type() == shape::bool_type)
+    {
+        r = literal{e.at<bool>()};
+    }
+    else
+    {
         e.visit_at([&](auto x) { r = literal{x}; });
+    }
     return r;
 }
 

src/include/migraphx/op/contiguous.hpp

@@ -28,6 +28,7 @@
 #include <migraphx/argument.hpp>
 #include <migraphx/shape_for_each.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -42,19 +43,27 @@ namespace op {
 struct contiguous
 {
     std::string name() const { return "contiguous"; }
+
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
-        if(inputs.front().standard())
-            return inputs.front();
-        auto lens = inputs.at(0).lens();
-        auto t    = inputs.at(0).type();
+        check_shapes{inputs, *this, true}.has(1);
+        auto s0 = inputs.front();
+        if(s0.dynamic() or s0.standard())
+        {
+            return s0;
+        }
+        else
+        {
+            const auto& lens = s0.lens();
+            auto t           = s0.type();
             return {t, lens};
         }
-    argument compute(const shape& output_shape, std::vector<argument> args) const
+    }
+
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        assert(output_shape.standard());
-        argument result{output_shape};
+        assert(dyn_out.computed_shape.standard());
+        argument result{dyn_out.computed_shape};
         visit_all(result, args[0])([&](auto output, auto input) {
             shape_for_each(output.get_shape(), [&](const auto& idx) {
                 output(idx.begin(), idx.end()) = input(idx.begin(), idx.end());

src/include/migraphx/op/squeeze.hpp

@@ -29,6 +29,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -54,14 +55,51 @@ struct squeeze
     std::string name() const { return "squeeze"; }
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
+        check_shapes{inputs, *this, true}.has(1);
         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;
+               }))
+            {
+                MIGRAPHX_THROW(
+                    "SQUEEZE: dynamic axis dimension should be equal to {1, 1, 0} or {1, 1, 1}");
+            }
+            std::vector<shape::dynamic_dimension> dyn_dims = {};
+            if(axes.empty())
+            {
+                for(auto i : range(input_shape.ndim()))
+                {
+                    auto dd = input_shape.dyn_dims()[i];
+                    if(dd != one_dyn_dim)
+                    {
+                        dyn_dims.push_back(dd);
+                    }
+                }
+            }
+            else
+            {
+                for(auto i : range(input_shape.ndim()))
+                {
+                    if(std::find(axes.begin(), axes.end(), i) == axes.end())
+                    {
+                        dyn_dims.push_back(input_shape.dyn_dims()[i]);
+                    }
+                }
+            }
+            return {input_shape.type(), dyn_dims};
+        }
+        else
+        {
             auto type        = input_shape.type();
             auto old_lens    = input_shape.lens();
             auto old_strides = input_shape.strides();
-        if(std::any_of(axes.begin(), axes.end(), [&](auto axis) { return old_lens[axis] != 1; }))
+            if(std::any_of(
+                   axes.begin(), axes.end(), [&](auto axis) { return old_lens[axis] != 1; }))
             {
-            MIGRAPHX_THROW("squeeze axis dimension should be equal to 1");
+                MIGRAPHX_THROW("SQUEEZE: static axis dimension should be equal to 1");
             }
             std::vector<std::size_t> new_lens;
             std::vector<std::size_t> new_strides;
@@ -96,10 +134,11 @@ struct squeeze
                 return shape{type, new_lens, new_strides};
             }
         }
+    }
 
-    argument compute(shape output_shape, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        return args[0].reshape(output_shape);
+        return args[0].reshape(dyn_out.computed_shape);
     }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };

src/include/migraphx/op/unsqueeze.hpp

@@ -29,11 +29,20 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
+/**
+ * Adds dimensions to a tensor based on the axes attribute.
+ * `axes` are based on the number of output shape dimensions and should not contain duplicates.
+ * `steps` are for modifying dimensions added to the middle of the original shape.
+ * Each step must be a factor of the original dimension.
+ * ex: unsqueeze(shape = [3, 4, 10], axes = [2, 4, 5], steps = [2]) -> shape = [3, 4, 2, 5, 1, 1]
+ * Dynamic shape version does not handle `steps`.
+ */
 struct unsqueeze
 {
     std::vector<int64_t> axes;
@@ -56,8 +65,33 @@ struct unsqueeze
     std::string name() const { return "unsqueeze"; }
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
+        check_shapes{inputs, *this, true}.has(1);
         auto input_shape = inputs[0];
+
+        if(input_shape.dynamic())
+        {
+            if(not steps.empty())
+            {
+                MIGRAPHX_THROW("UNSQUEEZE_dyn: nonempty steps attribute");
+            }
+            std::vector<shape::dynamic_dimension> dyn_dims = {};
+            auto new_ndim                                  = input_shape.ndim() + axes.size();
+            std::size_t k                                  = 0;
+            for(auto i : range(new_ndim))
+            {
+                if(std::find(axes.begin(), axes.end(), i) != axes.end())
+                {
+                    dyn_dims.push_back({1, 1, 0});
+                }
+                else
+                {
+                    dyn_dims.push_back(input_shape.dyn_dims().at(k++));
+                }
+            }
+            return {input_shape.type(), dyn_dims};
+        }
+        else
+        {
             auto type        = input_shape.type();
             auto old_lens    = input_shape.lens();
             auto old_strides = input_shape.strides();
@@ -110,9 +144,10 @@ struct unsqueeze
             }
             return shape{type, new_lens, new_strides};
         }
-    argument compute(shape output_shape, std::vector<argument> args) const
+    }
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        return args[0].reshape(output_shape);
+        return args[0].reshape(dyn_out.computed_shape);
     }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };

src/targets/gpu/CMakeLists.txt

@@ -233,11 +233,14 @@ get_target_property(MIOPEN_LOCATION MIOpen LOCATION)
 check_library_exists(MIOpen "miopenHiddenSetConvolutionFindMode" "${MIOPEN_LOCATION}" HAS_FIND_MODE_API)
 check_library_exists(MIOpen "miopenFindSolutions" "${MIOPEN_LOCATION}" HAS_FIND_2_API)
 
-if(HAS_FIND_2_API) 
+# TODO: Set default to HAS_FIND_2_API
+set(MIGRAPHX_USE_FIND_2_API OFF CACHE BOOL "")
+
+if(MIGRAPHX_USE_FIND_2_API) 
     target_compile_definitions(migraphx_gpu PUBLIC -DMIGRAPHX_HAS_FIND_2_API)
     message(STATUS "MIGraphx is using Find-2.0 API of MIOpen")
 else()
-    message(STATUS "MIOpen does not have Find-2.0 API")
+    message(STATUS "MIGraphx is using legacy Find API in MIOpen")
 endif()
 
 if(HAS_FIND_MODE_API)

test/fuse_pointwise.cpp

@@ -272,6 +272,35 @@ TEST_CASE(contiguous_input)
     EXPECT(p1 == p2);
 }
 
+TEST_CASE(contiguous_boolean_input)
+{
+
+    migraphx::shape s{migraphx::shape::bool_type, {2, 3}};
+    migraphx::shape s_lit{migraphx::shape::bool_type, {1}, {0}};
+    migraphx::program p1;
+    {
+        auto* mm = p1.get_main_module();
+        auto x   = mm->add_parameter("x", s);
+        auto one = mm->add_literal(migraphx::literal(s_lit, {1.0}));
+        auto yb =
+            mm->add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s.lens()}}), one);
+        auto y    = mm->add_instruction(migraphx::make_op("contiguous"), yb);
+        auto xor1 = mm->add_instruction(migraphx::make_op("logical_xor"), x, y);
+        mm->add_return({xor1});
+    }
+    run_pass(p1);
+    migraphx::program p2;
+    {
+        auto* mm  = p2.get_main_module();
+        auto x    = mm->add_parameter("x", s);
+        auto xor1 = add_pointwise(p2, "main:pointwise0", {x}, [=](auto* pm, const auto& inputs) {
+            auto y = pm->add_literal(migraphx::literal(s_lit, {1}));
+            return pm->add_instruction(migraphx::make_op("logical_xor"), inputs[0], y);
+        });
+        mm->add_return({xor1});
+    }
+}
+
 TEST_CASE(all_scalar_input)
 {
     migraphx::shape s{migraphx::shape::float_type};

test/onnx/gen_onnx.py

@@ -5984,6 +5984,26 @@ def squeeze_unsqueeze_test():
     return ([node, node2], [x], [y])
 
 
+@onnx_test
+def squeeze_unsqueeze_dyn_test():
+    x = helper.make_tensor_value_info('0', TensorProto.FLOAT,
+                                      [1, None, 1, 1, None, 1])
+    y = helper.make_tensor_value_info('2', TensorProto.FLOAT,
+                                      [1, 1, None, 1, None, 1])
+
+    node = onnx.helper.make_node('Squeeze',
+                                 inputs=['0'],
+                                 axes=[0, 2, 3, 5],
+                                 outputs=['1'])
+
+    node2 = onnx.helper.make_node('Unsqueeze',
+                                  inputs=['1'],
+                                  axes=[0, 1, 3, 5],
+                                  outputs=['2'])
+
+    return ([node, node2], [x], [y])
+
+
 @onnx_test
 def sub_bcast_test():
     arg0 = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])

test/onnx/onnx_test.cpp

@@ -5776,6 +5776,29 @@ TEST_CASE(squeeze_unsqueeze_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(squeeze_unsqueeze_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    std::vector<int64_t> squeeze_axes{0, 2, 3, 5};
+    std::vector<int64_t> unsqueeze_axes{0, 1, 3, 5};
+    auto l0 = mm->add_parameter(
+        "0",
+        migraphx::shape{migraphx::shape::float_type,
+                        {{1, 1, 0}, {1, 4, 0}, {1, 1, 0}, {1, 1, 0}, {1, 4, 0}, {1, 1, 0}}});
+    auto c0  = mm->add_instruction(migraphx::make_op("contiguous"), l0);
+    auto l1  = mm->add_instruction(migraphx::make_op("squeeze", {{"axes", squeeze_axes}}), c0);
+    auto c1  = mm->add_instruction(migraphx::make_op("contiguous"), l1);
+    auto ret = mm->add_instruction(migraphx::make_op("unsqueeze", {{"axes", unsqueeze_axes}}), c1);
+    mm->add_return({ret});
+
+    migraphx::onnx_options options;
+    options.default_dyn_dim_value = {1, 4, 0};
+    auto prog                     = parse_onnx("squeeze_unsqueeze_dyn_test.onnx", options);
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(squeeze_axes_input_test)
 {
     migraphx::program p;

test/op_shape_test.cpp

@@ -365,6 +365,12 @@ TEST_CASE(contiguous_shape)
     expect_shape(single, migraphx::make_op("contiguous"), single);
 }
 
+TEST_CASE(contiguous_dyn_shape)
+{
+    migraphx::shape s0{migraphx::shape::float_type, {{1, 4, 0}, {2, 2, 2}}};
+    expect_shape(s0, migraphx::make_op("contiguous"), s0);
+}
+
 TEST_CASE(contiguous_shape_scalar)
 {
     migraphx::shape output{migraphx::shape::float_type};
@@ -2089,6 +2095,30 @@ TEST_CASE(test_squeeze_all)
     expect_shape(s2, migraphx::make_op("squeeze", {{"axes", {0}}}), s1);
 }
 
+TEST_CASE(test_squeeze_dyn)
+{
+    migraphx::shape s1{migraphx::shape::float_type,
+                       {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}, {3, 3, 0}}};
+    migraphx::shape s2{migraphx::shape::float_type, {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s2, migraphx::make_op("squeeze", {{"axes", {3}}}), s1);
+
+    migraphx::shape s3{migraphx::shape::float_type, {{1, 4, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s3, migraphx::make_op("squeeze"), s1);
+
+    throws_shape(migraphx::make_op("squeeze", {{"axes", {0}}}), s1);
+}
+
+TEST_CASE(test_squeeze_dyn_neg_axes)
+{
+    migraphx::shape s1{migraphx::shape::float_type,
+                       {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}, {3, 3, 0}}};
+    migraphx::shape s2{migraphx::shape::float_type, {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s2, migraphx::make_op("squeeze", {{"axes", {-2}}}), s1);
+
+    migraphx::shape s3{migraphx::shape::float_type, {{1, 4, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s3, migraphx::make_op("squeeze", {{"axes", {-2, -4}}}), s1);
+}
+
 TEST_CASE(test_squeeze_transpose)
 {
     migraphx::shape s1{migraphx::shape::float_type, {4, 4, 1}, {4, 1, 4}};
@@ -2130,6 +2160,30 @@ TEST_CASE(test_unsqueeze)
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}}), s1);
 }
 
+TEST_CASE(test_unsqueeze_dyn)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {{1, 4, 3}, {2, 5, 0}, {3, 3, 0}}};
+    migraphx::shape s2{migraphx::shape::float_type, {{1, 4, 3}, {2, 5, 0}, {1, 1, 0}, {3, 3, 0}}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}}), s1);
+
+    migraphx::shape s3{migraphx::shape::float_type,
+                       {{1, 4, 3}, {2, 5, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}}};
+    expect_shape(s3, migraphx::make_op("unsqueeze", {{"axes", {2, 4}}}), s1);
+
+    throws_shape(migraphx::make_op("unsqueeze", {{"axes", {2, 4}}, {"steps", {2}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_dyn_neg_axes)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {{1, 4, 3}, {2, 5, 0}, {3, 3, 0}}};
+    migraphx::shape s2{migraphx::shape::float_type, {{1, 4, 3}, {2, 5, 0}, {1, 1, 0}, {3, 3, 0}}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {-2}}}), s1);
+
+    migraphx::shape s3{migraphx::shape::float_type,
+                       {{1, 4, 3}, {2, 5, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}}};
+    expect_shape(s3, migraphx::make_op("unsqueeze", {{"axes", {-1, -3}}}), s1);
+}
+
 TEST_CASE(test_unsqueeze_step)
 {
     migraphx::shape s1{migraphx::shape::float_type, {4, 5, 12}};
@@ -2161,13 +2215,27 @@ TEST_CASE(test_unsqueeze_mismatch_step_axis)
     throws_shape(migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2, 3}}}), s1);
 }
 
-TEST_CASE(test_unsqueeze_negative_axis)
+TEST_CASE(test_unsqueeze_negative_axis1)
 {
     migraphx::shape s1{migraphx::shape::float_type, {4, 5, 3}};
     migraphx::shape s2{migraphx::shape::float_type, {4, 5, 1, 3}};
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {-2}}}), s1);
 }
 
+TEST_CASE(test_unsqueeze_negative_axis2)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 5, 3, 1}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {-1}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_negative_axis3)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 1, 5, 3}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {-3}}}), s1);
+}
+
 TEST_CASE(test_unsqueeze_scalar)
 {
     migraphx::shape s1{migraphx::shape::float_type, {1}, {0}};

test/ref_ops_test.cpp

@@ -925,6 +925,33 @@ TEST_CASE(contiguous_test)
     EXPECT(migraphx::verify_range(results_vector, data));
 }
 
+TEST_CASE(contiguous_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape dyn_shape{migraphx::shape::float_type,
+                              {{1, 1, 0}, {2, 6, 0}, {2, 2, 0}, {2, 2, 0}}};
+    auto input = mm->add_parameter("X", dyn_shape);
+    mm->add_instruction(migraphx::make_op("contiguous"), input);
+    p.compile(migraphx::ref::target{});
+
+    migraphx::shape static_shape{migraphx::shape::float_type, {1, 3, 2, 2}, {12, 1, 6, 3}};
+    std::vector<float> data(12);
+    std::iota(data.begin(), data.end(), 0);
+    migraphx::parameter_map params;
+    params["X"] = migraphx::argument(static_shape, data.data());
+    auto result = p.eval(params).back();
+
+    std::vector<size_t> new_strides = {12, 4, 2, 1};
+    EXPECT(result.get_shape().strides() == new_strides);
+
+    std::vector<float> results_vector(12);
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+
+    std::vector<float> gold = {0, 3, 6, 9, 1, 4, 7, 10, 2, 5, 8, 11};
+    EXPECT(migraphx::verify_range(results_vector, gold));
+}
+
 TEST_CASE(conv_dynamic_batch_test)
 {
     migraphx::program p;
@@ -7084,6 +7111,25 @@ TEST_CASE(squeeze_test)
     }
 }
 
+TEST_CASE(squeeze_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s1{migraphx::shape::float_type,
+                       {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}, {3, 3, 0}}};
+    auto p0 = mm->add_parameter("x", s1);
+    mm->add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), p0);
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data(4 * 3 * 3);
+    migraphx::parameter_map params0;
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {4, 1, 3, 1, 3}};
+    params0["x"] = migraphx::argument(input_fixed_shape0, input_data.data());
+    auto result  = p.eval(params0).back();
+    migraphx::shape s2{migraphx::shape::float_type, {4, 3, 1, 3}};
+    EXPECT(result.get_shape() == s2);
+}
+
 TEST_CASE(step_test)
 {
     {
@@ -7361,6 +7407,25 @@ TEST_CASE(unsqueeze_test)
     }
 }
 
+TEST_CASE(unsqueeze_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    migraphx::shape s1{migraphx::shape::float_type, {{1, 4, 0}, {3, 3, 0}, {3, 3, 0}}};
+    auto p0 = mm->add_parameter("x", s1);
+    mm->add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), p0);
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data(4 * 3 * 3);
+    migraphx::parameter_map params0;
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {4, 3, 3}};
+    params0["x"] = migraphx::argument(input_fixed_shape0, input_data.data());
+    auto result  = p.eval(params0).back();
+    migraphx::shape s2{migraphx::shape::float_type, {4, 1, 3, 3}};
+    EXPECT(result.get_shape() == s2);
+}
+
 TEST_CASE(where_test)
 {
     migraphx::program p;