Nd deconv GPU support (#558)

* initial progress

* formatting

* check existing tests

* formatting

* change for loop to transform

* formatting

* add tests

* formatting

* remove comment

* add more tests

* update gpu miopen calls

* formatting

* fix error msg
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/include/migraphx/op/convolution.hpp

@@ -55,6 +55,10 @@ struct convolution
         const shape& weights = inputs.at(1);
         auto t               = input.type();
         size_t kdims         = input.lens().size() - 2;
+        if(kdims != this->kdims())
+        {
+            MIGRAPHX_THROW("convolution: input k-dims does not match attribute size");
+        }
 
         if(input.lens().at(1) != (weights.lens().at(1) * group))
             MIGRAPHX_THROW("CONVOLUTION: Mismatch channel numbers");

src/include/migraphx/op/deconvolution.hpp

@@ -37,18 +37,27 @@ struct deconvolution
     }
 
     std::string name() const { return "deconvolution"; }
-    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("deconvolution: inconsistent attribute sizes");
         }
+    }
+
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(2).same_type().same_ndims().min_ndims(3);
 
         const shape& input   = inputs.at(0);
         const shape& weights = inputs.at(1);
         auto t               = input.type();
         size_t kdims         = input.lens().size() - 2;
+        if(kdims != this->kdims())
+        {
+            MIGRAPHX_THROW("deconvolution: input k-dims does not match attribute size");
+        }
 
         std::vector<size_t> output_lens{input.lens()[0], weights.lens()[1]};
 
@@ -61,6 +70,12 @@ struct deconvolution
         }
         return {t, output_lens};
     }
+
+    size_t kdims() const
+    {
+        check_attribute_size();
+        return padding.size();
+    }
 };
 
 } // namespace op

src/include/migraphx/op/quant_convolution.hpp

@@ -55,6 +55,10 @@ struct quant_convolution
         const shape& weights = inputs.at(1);
         auto t               = input.type();
         size_t kdims         = input.lens().size() - 2;
+        if(kdims != this->kdims())
+        {
+            MIGRAPHX_THROW("quant_convolution: input k-dims does not match attribute size");
+        }
 
         // all input type must be int8_type and output is float_type
         if(t != shape::int8_type)

src/targets/gpu/convolution.cpp

@@ -14,7 +14,7 @@ shape miopen_convolution::compute_shape(const std::vector<shape>& inputs) const
     return op.compute_shape(conv_inputs);
 }
 
-inline shape recompute_shape_to_2d(const shape& input)
+inline shape reshape_if_1d(const shape& input)
 {
     shape new_shape{input};
     auto dims = new_shape.lens();
@@ -32,9 +32,9 @@ argument miopen_convolution::compute(context& ctx,
                                      const shape& output_shape,
                                      const std::vector<argument>& args) const
 {
-    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));
+    auto x_desc = make_tensor(reshape_if_1d(args[0].get_shape()));
+    auto w_desc = make_tensor(reshape_if_1d(args[1].get_shape()));
+    auto y_desc = make_tensor(reshape_if_1d(output_shape));
 
     float alpha = 1;
     float beta  = 0;
@@ -62,9 +62,9 @@ shape miopen_convolution::compile(context& ctx,
 {
     shape workspace_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));
+    auto x_desc = make_tensor(reshape_if_1d(inputs[0]));
+    auto w_desc = make_tensor(reshape_if_1d(inputs[1]));
+    auto y_desc = make_tensor(reshape_if_1d(output_shape));
 
     std::size_t workspace_size = 0;
     miopenConvolutionForwardGetWorkSpaceSize(ctx.get_stream().get_miopen(),

src/targets/gpu/deconvolution.cpp

@@ -9,15 +9,32 @@ namespace gpu {
 shape miopen_deconvolution::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 reshape_if_1d(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_deconvolution::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(reshape_if_1d(args[0].get_shape()));
+    auto w_desc = make_tensor(reshape_if_1d(args[1].get_shape()));
+    auto y_desc = make_tensor(reshape_if_1d(output_shape));
 
     float alpha = 1;
     float beta  = 0;
@@ -44,9 +61,9 @@ shape miopen_deconvolution::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(reshape_if_1d(inputs[0]));
+    auto w_desc = make_tensor(reshape_if_1d(inputs[1]));
+    auto y_desc = make_tensor(reshape_if_1d(output_shape));
 
     std::size_t workspace_size = 0;
     miopenConvolutionForwardGetWorkSpaceSize(ctx.get_stream().get_miopen(),

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

@@ -104,14 +104,17 @@ inline convolution_descriptor make_deconv(const T& op)
 {
     auto c = make_obj<convolution_descriptor>(&miopenCreateConvolutionDescriptor);
     miopenConvolutionMode_t c_mode = miopenTranspose;
-    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;