Merge pull request #278 from ROCmSoftwarePlatform/device-refactor

Refactor device

src/targets/gpu/include/migraphx/gpu/device/logsoftmax.hpp

@@ -10,10 +10,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 namespace device {
 
-argument logsoftmax(hipStream_t stream,
-                    const migraphx::shape& output_shape,
-                    std::vector<migraphx::argument> args,
-                    int axis);
+argument logsoftmax(hipStream_t stream, argument result, argument arg, int axis);
 
 } // namespace device
 } // namespace gpu

src/targets/gpu/include/migraphx/gpu/device/softmax.hpp

@@ -10,10 +10,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 namespace device {
 
-argument softmax(hipStream_t stream,
-                 const migraphx::shape& output_shape,
-                 std::vector<migraphx::argument> args,
-                 int axis);
+argument softmax(hipStream_t stream, argument result, argument arg, int axis);
 
 } // namespace device
 } // namespace gpu

src/targets/gpu/logsoftmax.cpp

@@ -15,11 +15,10 @@ shape hip_logsoftmax::compute_shape(const std::vector<shape>& inputs) const
     return op.compute_shape({inputs.at(0)});
 }
 
-argument hip_logsoftmax::compute(context& ctx,
-                                 const shape& output_shape,
-                                 const std::vector<argument>& args) const
+argument
+hip_logsoftmax::compute(context& ctx, const shape&, const std::vector<argument>& args) const
 {
-    return device::logsoftmax(ctx.get_stream().get(), output_shape, args, op.axis);
+    return device::logsoftmax(ctx.get_stream().get(), args[1], args[0], op.axis);
 }
 
 } // namespace gpu

src/targets/gpu/softmax.cpp

@@ -37,11 +37,9 @@ shape hip_softmax::compute_shape(const std::vector<shape>& inputs) const
     return op.compute_shape({inputs.at(0)});
 }
 
-argument hip_softmax::compute(context& ctx,
-                              const shape& output_shape,
-                              const std::vector<argument>& args) const
+argument hip_softmax::compute(context& ctx, const shape&, const std::vector<argument>& args) const
 {
-    return device::softmax(ctx.get_stream().get(), output_shape, args, op.axis);
+    return device::softmax(ctx.get_stream().get(), args[1], args[0], op.axis);
 }
 
 } // namespace gpu