add the gpu implementation of the logsoftmax operator

src/targets/gpu/CMakeLists.txt

@@ -48,6 +48,7 @@ add_library(migraphx_gpu
     pooling.cpp
     convolution.cpp
     softmax.cpp
+    logsoftmax.cpp
     contiguous.cpp
     concat.cpp
     relu.cpp

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_LOGSOFTMAX_HPP
+#define MIGRAPHX_GUARD_RTGLIB_LOGSOFTMAX_HPP
+
+#include <migraphx/gpu/lowering.hpp>
+#include <migraphx/manage_ptr.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/operators.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/shape_for_each.hpp>
+#include <migraphx/config.hpp>
+#include <migraphx/gpu/miopen.hpp>
+#include <migraphx/gpu/hip.hpp>
+#include <migraphx/dfor.hpp>
+#include <migraphx/gpu/device/contiguous.hpp>
+#include <migraphx/gpu/device/add.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/gpu/rocblas.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <utility>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+struct miopen_logsoftmax
+{
+    op::logsoftmax op;
+    std::string name() const { return "gpu::logsoftmax"; }
+    shape compute_shape(const std::vector<shape>& inputs) const;
+    argument
+    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    int output_alias(const std::vector<shape>& shapes) const { return shapes.size() - 1; }
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/logsoftmax.cpp

+#include <migraphx/gpu/logsoftmax.hpp>
+#include <migraphx/gpu/device/log.hpp>
+#include <migraphx/operators.hpp>
+#include <migraphx/manage_ptr.hpp>
+#include <migraphx/gpu/miopen.hpp>
+#include <utility>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+shape miopen_logsoftmax::compute_shape(const std::vector<shape>& inputs) const
+{
+    check_shapes{inputs, *this}.has(2).standard();
+    return op.compute_shape({inputs.at(0)});
+}
+
+argument miopen_logsoftmax::compute(context& ctx,
+                                 const shape& output_shape,
+                                 const std::vector<argument>& args) const
+{
+    float alpha = 1;
+    float beta  = 0;
+    // temporarily reshape the input to a(0)...a(axis-1)
+    // and a(axis)....a(n)
+    auto lens = output_shape.lens();
+    std::size_t batch_size = std::accumulate(
+                            lens.begin(), lens.begin() + op.axis, 
+                            std::size_t{1}, std::multiplies<std::size_t>());
+    std::size_t n_dims = std::accumulate(lens.begin() + op.axis,
+                            lens.end(), std::size_t{1}, std::multiplies<std::size_t>());
+    migraphx::shape comp_shape{output_shape.type(), {batch_size, n_dims, 1, 1}};
+    auto x_desc = make_tensor(args[0].get_shape());
+    auto y_desc = make_tensor(output_shape);
+
+    miopenSoftmaxForward(ctx.get_stream().get_miopen(),
+                         &alpha,
+                         x_desc.get(),
+                         args[0].implicit(),
+                         &beta,
+                         y_desc.get(),
+                         args[1].implicit());
+
+    // call the device::log function to perform the log operation
+    device::log(ctx.get_stream().get(), args[1], args[0]);
+
+    return args[1];
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/lowering.cpp

@@ -21,6 +21,7 @@
 #include <migraphx/gpu/leaky_relu.hpp>
 #include <migraphx/gpu/elu.hpp>
 #include <migraphx/gpu/softmax.hpp>
+#include <migraphx/gpu/logsoftmax.hpp>
 #include <migraphx/gpu/add.hpp>
 #include <migraphx/gpu/sub.hpp>
 #include <migraphx/gpu/exp.hpp>
@@ -97,6 +98,7 @@ struct miopen_apply
         add_extend_op<miopen_contiguous, op::contiguous>("contiguous");
         add_extend_op<hip_concat, op::concat>("concat");
         add_extend_op<miopen_softmax, op::softmax>("softmax");
+        add_extend_op<miopen_logsoftmax, op::logsoftmax>("logsoftmax");
         add_extend_op<hip_gather, op::gather>("gather");
         add_extend_op<hip_pad, op::pad>("pad");