Add solution_id to convolution (#589)

* initial progres

* formatting

* remove comment

* update reflect and error msg

* formatting

* remove header

* move and rename function

* formatting

* fix tidy and remove extra function
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/targets/gpu/convolution.cpp

@@ -36,29 +36,27 @@ argument miopen_convolution::compute(context& ctx,
     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;
-    auto status = miopenConvolutionForward(ctx.get_stream().get_miopen(),
-                                           &alpha,
-                                           x_desc.get(),
-                                           args[0].implicit(),
-                                           w_desc.get(),
-                                           args[1].implicit(),
-                                           cd.get(),
-                                           algo,
-                                           &beta,
-                                           y_desc.get(),
-                                           args[3].implicit(),
-                                           args[2].implicit(),
-                                           args[2].get_shape().bytes());
+    if(solution_id == 0)
+        MIGRAPHX_THROW("MIOpen Convolution: invalid solution ID");
+
+    auto status = miopenConvolutionForwardImmediate(ctx.get_stream().get_miopen(),
+                                                    w_desc.get(),
+                                                    args[1].implicit(),
+                                                    x_desc.get(),
+                                                    args[0].implicit(),
+                                                    cd.get(),
+                                                    y_desc.get(),
+                                                    args[3].implicit(),
+                                                    args[2].implicit(),
+                                                    args[2].get_shape().bytes(),
+                                                    solution_id);
+
     if(status != miopenStatusSuccess)
-        MIGRAPHX_THROW("Running convolution failed");
+        MIGRAPHX_THROW("MIOpen Convolution: running convolution failed");
     return args[3];
 }
 
-shape miopen_convolution::compile(context& ctx,
-                                  const shape& output_shape,
-                                  std::vector<shape> inputs)
+shape miopen_convolution::find(context& ctx, const shape& output_shape, std::vector<shape> inputs)
 {
     shape workspace_shape{};
 
@@ -97,9 +95,32 @@ shape miopen_convolution::compile(context& ctx,
                                                         workspace_size,
                                                         false);
     if(status != miopenStatusSuccess)
-        MIGRAPHX_THROW("Find convolution failed");
+        MIGRAPHX_THROW("MIOpen Convolution: find convolution failed");
     handle = ctx.get_stream().get_miopen();
     algo   = perf.fwd_algo;
+
+    size_t solution_count;
+
+    status = miopenConvolutionForwardGetSolutionCount(
+        handle, w_desc.get(), x_desc.get(), cd.get(), y_desc.get(), &solution_count);
+    if(status != miopenStatusSuccess)
+        MIGRAPHX_THROW("MIOpen Convolution: get solution count failed");
+
+    std::vector<miopenConvSolution_t> solutions(solution_count);
+
+    status = miopenConvolutionForwardGetSolution(handle,
+                                                 w_desc.get(),
+                                                 x_desc.get(),
+                                                 cd.get(),
+                                                 y_desc.get(),
+                                                 solution_count,
+                                                 &solution_count,
+                                                 solutions.data());
+    if(status != miopenStatusSuccess)
+        MIGRAPHX_THROW("MIOpen Convolution: get solution failed");
+
+    solution_id = solutions.front().solution_id;
+
     return shape{shape::int8_type, {perf.memory}};
 }
 
@@ -109,11 +130,24 @@ void miopen_convolution::finalize(context& ctx,
 {
     if(handle == ctx.get_stream().get_miopen())
         return;
-    // Check that workspace hasn't changed
-    auto size = inputs.at(2).bytes();
-    auto ws   = compile(ctx, output_shape, std::move(inputs));
-    if(ws.bytes() > size)
-        MIGRAPHX_THROW("Workspace has changed during finalization.");
+
+    if(solution_id == 0)
+    {
+        // Check that workspace hasn't changed
+        auto size = inputs.at(2).bytes();
+        auto ws   = find(ctx, output_shape, inputs);
+        if(ws.bytes() > size)
+            MIGRAPHX_THROW("MIOpen Convolution: workspace has changed during finalization.");
+    }
+
+    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));
+
+    auto status = miopenConvolutionForwardCompileSolution(
+        handle, w_desc.get(), x_desc.get(), cd.get(), y_desc.get(), solution_id);
+    if(status != miopenStatusSuccess)
+        MIGRAPHX_THROW("MIOpen Convolution: compile solution failed");
 }
 
 } // namespace gpu

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

@@ -17,19 +17,24 @@ struct miopen_convolution
     shared<convolution_descriptor> cd;
     miopenConvFwdAlgorithm_t algo{};
     miopenHandle_t handle = nullptr;
+    uint64_t solution_id  = 0;
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
-        // TODO: Add algo
-        return op::convolution::reflect(self.op, f);
+        return pack(f(self.op.padding, "padding"),
+                    f(self.op.stride, "stride"),
+                    f(self.op.dilation, "dilation"),
+                    f(self.op.group, "group"),
+                    f(self.op.padding_mode, "padding_mode"),
+                    f(self.solution_id, "solution_id"));
     }
 
     std::string name() const { return "gpu::convolution"; }
     shape compute_shape(const std::vector<shape>& inputs) const;
     argument
     compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
-    shape compile(context& ctx, const shape& output_shape, std::vector<shape> inputs);
+    shape find(context& ctx, const shape& output_shape, std::vector<shape> inputs);
     void finalize(context& ctx, const shape& output_shape, std::vector<shape> inputs);
     std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
     {

src/targets/gpu/lowering.cpp

@@ -282,7 +282,7 @@ struct miopen_apply
             auto&& op = any_cast<op::convolution>(ins->get_operator());
 
             auto conv = miopen_convolution{op, make_conv(op)};
-            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
+            auto ws   = conv.find(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
 
             auto workspace = insert_allocation(ins, ws, "workspace");
             auto output    = insert_allocation(ins, ins->get_shape());