Reverse Op (#846)

* init reverseOp branch: ref op + ref test. WIP

* first passing basic test

* cleanup

* additional axis implementation

* additional test

* ref op implementation vec to int for axis

* ref op test change for axis

* initial gpu files and test

* updates to implementation and test

* fixed some issues

* clang format

* cleanup

* formatting

* removing comments

* remove local size, back to default

* update tests: replace with std functions

* multiple axis for reverse op

* fix a build error

* clang format

* more tests

* fix a bug for the reverse device function

* clang format

* fix a bug

* clang format

* ref test updates, multiaxis

* formatting
Co-authored-by: Shucai Xiao <Shucai.Xiao@amd.com>
Co-authored-by: mvermeulen <5479696+mvermeulen@users.noreply.github.com>

src/CMakeLists.txt

@@ -135,6 +135,7 @@ register_migraphx_ops(
     reduce_sum
     relu
     reshape
+    reverse
     rnn
     rnn_last_cell_output
     rnn_last_hs_output

src/include/migraphx/op/reverse.hpp

+#ifndef MIGRAPHX_GUARD_OPERATORS_REVERSE_HPP
+#define MIGRAPHX_GUARD_OPERATORS_REVERSE_HPP
+
+#include <algorithm>
+#include <vector>
+#include <cmath>
+#include <utility>
+#include <migraphx/config.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/par_for.hpp>
+#include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/shape_for_each.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/value.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct reverse
+{
+
+    std::vector<int64_t> axes;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.axes, "axes"));
+    }
+
+    std::string name() const { return "reverse"; }
+
+    value attributes() const
+    {
+        value normalize;
+        normalize["axes"] = value::array{normalize_attribute::include_min};
+        return {{"normalize_axes", normalize}};
+    }
+
+    shape normalize_compute_shape(std::vector<shape> inputs) const
+    {
+        auto lens = inputs[0].lens();
+        auto type = inputs[0].type();
+        return shape{type, lens};
+    }
+
+    argument compute(const shape& s, std::vector<argument> args) const
+    {
+        argument result{s};
+        auto lens = s.lens();
+        visit_all(result, args.front())([&](auto output, auto input) {
+            shape_for_each(s, [&](const auto& out_idx) {
+                auto in_idx = out_idx;
+                for(const auto& axis : axes)
+                {
+                    in_idx[axis] = lens[axis] - 1 - out_idx[axis];
+                }
+                output[s.index(out_idx)] = input[s.index(in_idx)];
+            });
+        });
+
+        return result;
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/operators.hpp

@@ -71,6 +71,7 @@
 #include <migraphx/op/reduce_sum.hpp>
 #include <migraphx/op/relu.hpp>
 #include <migraphx/op/reshape.hpp>
+#include <migraphx/op/reverse.hpp>
 #include <migraphx/op/rnn.hpp>
 #include <migraphx/op/rnn_last_cell_output.hpp>
 #include <migraphx/op/rnn_last_hs_output.hpp>

src/targets/gpu/CMakeLists.txt

@@ -68,6 +68,7 @@ add_library(migraphx_device
     device/reduce_sum.cpp
     device/reduce_prod.cpp
     device/relu.cpp
+    device/reverse.cpp
     device/rnn_variable_seq_lens.cpp
     device/round.cpp
     device/rsqrt.cpp
@@ -140,6 +141,7 @@ add_library(migraphx_gpu
     pooling.cpp
     preallocate_param.cpp
     quant_convolution.cpp
+    reverse.cpp
     rnn_variable_seq_lens.cpp
     rocblas.cpp
     softmax.cpp
@@ -197,6 +199,7 @@ register_migraphx_gpu_ops(hip_
     reduce_prod
     reduce_sum
     relu
+    reverse
     round
     rsqrt
     sigmoid

src/targets/gpu/device/reverse.cpp

+#include "migraphx/gpu/device/visit.hpp"
+#include <migraphx/shape.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/gpu/device/reverse.hpp>
+#include <migraphx/gpu/device/tensor.hpp>
+#include <migraphx/gpu/device/launch.hpp>
+#include <migraphx/gpu/device/types.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace device {
+
+argument
+reverse(hipStream_t stream, argument result, argument arg1, const std::vector<int64_t>& axes)
+{
+    auto s = arg1.get_shape();
+    // auto lens             = s.lens();
+    std::vector<std::size_t> axis_len(axes.begin(), axes.end());
+    shape sa{shape::float_type, axis_len};
+    std::size_t nelements = s.elements();
+    visit_all(result, arg1)([&](auto output1, auto input1) {
+        hip_visit_views(output1, input1, s)([&](auto output, auto input, auto hs) {
+            hip_visit_views(sa)([&](auto daxes) {
+                auto lens = hs.lens;
+                gs_launch(stream, nelements)([=](auto i) __device__ {
+                    auto idx    = hs.multi(i);
+                    auto in_idx = idx;
+                    for(auto axis : daxes.lens)
+                        in_idx[axis] = lens[axis] - 1 - idx[axis];
+                    output[idx] = input[in_idx];
+                });
+            });
+        });
+    });
+
+    return result;
+}
+
+} // namespace device
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_DEVICE_REVERSE_HPP
+#define MIGRAPHX_GUARD_RTGLIB_DEVICE_REVERSE_HPP
+
+#include <migraphx/argument.hpp>
+#include <migraphx/config.hpp>
+#include <hip/hip_runtime_api.h>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+namespace device {
+
+argument
+reverse(hipStream_t stream, argument result, argument arg1, const std::vector<int64_t>& axes);
+
+} // namespace device
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_REVERSE_HPP
+#define MIGRAPHX_GUARD_RTGLIB_REVERSE_HPP
+
+#include <migraphx/argument.hpp>
+#include <migraphx/reflect.hpp>
+#include <migraphx/op/reverse.hpp>
+#include <migraphx/gpu/miopen.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+struct context;
+
+struct hip_reverse
+{
+    op::reverse op;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return migraphx::reflect(self.op, f);
+    }
+
+    std::string name() const { return "gpu::reverse"; }
+    shape compute_shape(std::vector<shape> inputs) const;
+    argument
+    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const;
+    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
+    {
+        return shapes.size() - 1;
+    }
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/lowering.cpp

@@ -160,6 +160,7 @@ struct miopen_apply
         add_extend_op("reduce_min");
         add_extend_op("reduce_prod");
         add_extend_op("reduce_sum");
+        add_extend_op("reverse");
         add_extend_op("rnn_var_sl_last_output");
         add_extend_op("rnn_var_sl_shift_output");
         add_extend_op("rnn_var_sl_shift_sequence");

src/targets/gpu/reverse.cpp

+#include <migraphx/gpu/reverse.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/gpu/device/reverse.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+shape hip_reverse::compute_shape(std::vector<shape> inputs) const
+{
+    inputs.pop_back();
+    return op.normalize_compute_shape(inputs);
+}
+
+argument hip_reverse::compute(context& ctx, const shape&, const std::vector<argument>& args) const
+{
+    return device::reverse(ctx.get_stream().get(), args.back(), args[0], op.axes);
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

test/ref_ops_test.cpp

@@ -3560,6 +3560,66 @@ TEST_CASE(reshape_test)
     }
 }
 
+TEST_CASE(reverse_test_axis0)
+{
+    migraphx::shape in_shape{migraphx::shape::float_type, {2, 16}};
+    std::vector<float> data(32);
+    std::iota(data.begin(), data.end(), 1);
+    migraphx::program p;
+    auto* mm              = p.get_main_module();
+    auto l                = mm->add_literal(migraphx::literal{in_shape, data});
+    std::vector<int> axes = {0};
+    mm->add_instruction(migraphx::make_op("reverse", {{"axes", axes}}), l);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> target_data = data;
+    std::swap_ranges(target_data.begin(), target_data.begin() + 16, target_data.begin() + 16);
+    EXPECT(migraphx::verify_range(results_vector, target_data));
+}
+
+TEST_CASE(reverse_test_axis1)
+{
+    migraphx::shape in_shape{migraphx::shape::float_type, {2, 16}};
+    std::vector<float> data(32);
+    std::iota(data.begin(), data.end(), 1);
+    migraphx::program p;
+    auto* mm              = p.get_main_module();
+    auto l                = mm->add_literal(migraphx::literal{in_shape, data});
+    std::vector<int> axes = {1};
+    mm->add_instruction(migraphx::make_op("reverse", {{"axes", axes}}), l);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> target_data = data;
+    std::reverse(target_data.begin(), target_data.begin() + 16);
+    std::reverse(target_data.end() - 16, target_data.end());
+    EXPECT(migraphx::verify_range(results_vector, target_data));
+}
+
+TEST_CASE(reverse_test_axis10)
+{
+    migraphx::shape in_shape{migraphx::shape::float_type, {2, 16}};
+    std::vector<float> data(32);
+    std::iota(data.begin(), data.end(), 1);
+    migraphx::program p;
+    auto* mm              = p.get_main_module();
+    auto l                = mm->add_literal(migraphx::literal{in_shape, data});
+    std::vector<int> axes = {1, 0};
+    mm->add_instruction(migraphx::make_op("reverse", {{"axes", axes}}), l);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> target_data = data;
+    std::reverse(target_data.begin(), target_data.begin() + 16);
+    std::reverse(target_data.end() - 16, target_data.end());
+    std::swap_ranges(target_data.begin(), target_data.begin() + 16, target_data.begin() + 16);
+    EXPECT(migraphx::verify_range(results_vector, target_data));
+}
+
 TEST_CASE(round_test)
 {
     migraphx::program p;

test/verify/test_reverse.cpp

+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_reverse : verify_program<test_reverse>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::float_type, {4, 16}};
+        auto a0                   = mm->add_parameter("data", s);
+        std::vector<int64_t> axis = {0};
+        mm->add_instruction(migraphx::make_op("reverse", {{"axes", axis}}), a0);
+        return p;
+    }
+};

test/verify/test_reverse_multiaxis.cpp

+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_reverse_multiaxis : verify_program<test_reverse_multiaxis>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::float_type, {4, 16}};
+        auto a0                   = mm->add_parameter("data", s);
+        std::vector<int64_t> axes = {0, 1};
+        mm->add_instruction(migraphx::make_op("reverse", {{"axes", axes}}), a0);
+        return p;
+    }
+};

test/verify/test_reverse_negaxis.cpp

+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct test_reverse_negaxis : verify_program<test_reverse_negaxis>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape s{migraphx::shape::float_type, {4, 16}};
+        auto a0                   = mm->add_parameter("data", s);
+        std::vector<int64_t> axis = {-1};
+        mm->add_instruction(migraphx::make_op("reverse", {{"axes", axis}}), a0);
+        return p;
+    }
+};