ND convolution GPU support (#550)

* initial progress

* formatting

* add pooling changes

* formatting

* change eliminate_pad

* formatting

* rename var

* fomratting

* update op shape test and compute

* formatting

* revert conv constructor

* formatting

* change initializer

* formatting

* fix tidy

* change quant conv and shape check

* add tests and fixes

* formatting

* fix type

* fix conv test

* formatting

* add pooling and bn tests

* formatting

* add inconsistent attr tests

* fix padding issue

* formatting

* progress on 1d to 2d

* formatting

* change compute and compile functions

* formatting

* fix duplicate

* fix conflict

* fix issue with 1d conv

* formatting

* add check for 3d limit

* rename function

* formatting

* rename functions

* formatting

* add op_shape test

* change functions

* formatting

* change to copy_backward

* formatting

src/include/migraphx/check_shapes.hpp

@@ -71,6 +71,19 @@ struct check_shapes
         return *this;
     }
 
+    const check_shapes& max_ndims(std::size_t n) const
+    {
+        assert(begin != nullptr);
+        assert(end != nullptr);
+        if(begin != end)
+        {
+            if(begin->lens().size() > n)
+                MIGRAPHX_THROW(prefix() + "Shape must have at most " + std::to_string(n) +
+                               " dimensions");
+        }
+        return *this;
+    }
+
     const check_shapes& min_ndims(std::size_t n) const
     {
         assert(begin != nullptr);

src/include/migraphx/op/convolution.hpp

@@ -37,13 +37,19 @@ struct convolution
     }
 
     std::string name() const { return "convolution"; }
-    shape compute_shape(std::vector<shape> inputs) const
+
+    void check_attribute_size() const
     {
-        check_shapes{inputs, *this}.has(2).same_type().same_ndims().min_ndims(3);
         if(not(padding.size() == stride.size() and padding.size() == dilation.size()))
         {
             MIGRAPHX_THROW("convolution: inconsistent attribute sizes");
         }
+    }
+
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(2).same_type().same_ndims().min_ndims(3);
+        check_attribute_size();
 
         const shape& input   = inputs.at(0);
         const shape& weights = inputs.at(1);
@@ -67,6 +73,12 @@ struct convolution
 
         return {t, output_lens};
     }
+
+    size_t kdims() const
+    {
+        check_attribute_size();
+        return padding.size();
+    }
 };
 
 } // namespace op

src/include/migraphx/op/quant_convolution.hpp

@@ -37,13 +37,19 @@ struct quant_convolution
     }
 
     std::string name() const { return "quant_convolution"; }
-    shape compute_shape(std::vector<shape> inputs) const
+
+    void check_attribute_size() const
     {
-        check_shapes{inputs, *this}.has(2).same_type().same_ndims().min_ndims(3);
         if(not(padding.size() == stride.size() and padding.size() == dilation.size()))
         {
             MIGRAPHX_THROW("quant_convolution: inconsistent attribute sizes");
         }
+    }
+
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(2).same_type().same_ndims().min_ndims(3);
+        check_attribute_size();
 
         const shape& input   = inputs.at(0);
         const shape& weights = inputs.at(1);
@@ -71,6 +77,12 @@ struct quant_convolution
 
         return {t, output_lens};
     }
+
+    size_t kdims() const
+    {
+        check_attribute_size();
+        return padding.size();
+    }
 };
 
 } // namespace op

src/targets/gpu/convolution.cpp

@@ -9,15 +9,32 @@ namespace gpu {
 shape miopen_convolution::compute_shape(const std::vector<shape>& inputs) const
 {
     check_shapes{inputs, *this}.has(4).standard();
-    return op.compute_shape({inputs.at(0), inputs.at(1)});
+    std::vector<shape> conv_inputs(inputs.begin(), inputs.begin() + 2);
+    check_shapes{conv_inputs, *this}.max_ndims(5);
+    return op.compute_shape(conv_inputs);
 }
+
+inline shape recompute_shape_to_2d(const shape& input)
+{
+    shape new_shape{input};
+    auto dims = new_shape.lens();
+
+    if(dims.size() == 3)
+    {
+        std::vector<size_t> new_dims = dims;
+        new_dims.insert(new_dims.begin() + 2, 1);
+        new_shape = shape{input.type(), new_dims};
+    }
+    return new_shape;
+}
+
 argument miopen_convolution::compute(context& ctx,
                                      const shape& output_shape,
                                      const std::vector<argument>& args) const
 {
-    auto x_desc = make_tensor(args[0].get_shape());
-    auto w_desc = make_tensor(args[1].get_shape());
-    auto y_desc = make_tensor(output_shape);
+    auto x_desc = make_tensor(recompute_shape_to_2d(args[0].get_shape()));
+    auto w_desc = make_tensor(recompute_shape_to_2d(args[1].get_shape()));
+    auto y_desc = make_tensor(recompute_shape_to_2d(output_shape));
 
     float alpha = 1;
     float beta  = 0;
@@ -44,9 +61,10 @@ shape miopen_convolution::compile(context& ctx,
                                   std::vector<shape> inputs)
 {
     shape workspace_shape{};
-    auto x_desc = make_tensor(inputs[0]);
-    auto w_desc = make_tensor(inputs[1]);
-    auto y_desc = make_tensor(output_shape);
+
+    auto x_desc = make_tensor(recompute_shape_to_2d(inputs[0]));
+    auto w_desc = make_tensor(recompute_shape_to_2d(inputs[1]));
+    auto y_desc = make_tensor(recompute_shape_to_2d(output_shape));
 
     std::size_t workspace_size = 0;
     miopenConvolutionForwardGetWorkSpaceSize(ctx.get_stream().get_miopen(),

src/targets/gpu/include/migraphx/gpu/miopen.hpp

@@ -82,14 +82,18 @@ inline convolution_descriptor make_conv(const T& op)
     miopenConvolutionMode_t c_mode = miopenConvolution;
     if(op.group > 1)
         c_mode = miopenGroupConv;
-    miopenInitConvolutionDescriptor(c.get(),
-                                    c_mode,
-                                    op.padding[0],
-                                    op.padding[1],
-                                    op.stride[0],
-                                    op.stride[1],
-                                    op.dilation[0],
-                                    op.dilation[1]);
+
+    int kdims = op.kdims();
+    std::vector<int> padding(std::max(2, kdims), 0);
+    std::vector<int> stride(std::max(2, kdims), 1);
+    std::vector<int> dilation(std::max(2, kdims), 1);
+
+    std::copy_backward(op.padding.begin(), op.padding.end(), padding.end());
+    std::copy_backward(op.stride.begin(), op.stride.end(), stride.end());
+    std::copy_backward(op.dilation.begin(), op.dilation.end(), dilation.end());
+
+    miopenInitConvolutionNdDescriptor(
+        c.get(), padding.size(), padding.data(), stride.data(), dilation.data(), c_mode);
     if(op.group > 1)
         miopenSetConvolutionGroupCount(c.get(), op.group);
     return c;

test/op_shape_test.cpp

@@ -96,6 +96,7 @@ TEST_CASE(quant_convolution_shape)
     migraphx::shape weights{migraphx::shape::int8_type, {4, 3, 3, 3}};
     expect_shape(output, migraphx::op::quant_convolution{}, input, weights);
     throws_shape(migraphx::op::quant_convolution{}, input);
+    throws_shape(migraphx::op::quant_convolution{{0}, {1, 1}, {1, 1}}, input, weights);
 
     migraphx::shape input2{migraphx::shape::int32_type, {3, 3}};
     migraphx::shape weights2{migraphx::shape::float_type, {3, 3}};