Merge pull request #169 from ROCmSoftwarePlatform/lrn

Lrn (initial support)

src/include/migraphx/operators.hpp

@@ -60,6 +60,30 @@ struct batch_norm_inference
     }
 };
 
+struct lrn
+{
+    float alpha = 0.0001;
+    float beta  = 0.75;
+    float bias  = 1.0;
+    int size    = 1;
+    std::string name() const { return "lrn"; }
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.alpha, "alpha"),
+                    f(self.beta, "beta"),
+                    f(self.bias, "bias"),
+                    f(self.size, "size"));
+    }
+
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(1);
+        return inputs.front();
+    }
+};
+
 struct convolution
 {
     std::array<std::size_t, 2> padding  = {{0, 0}};

src/onnx/onnx.cpp

@@ -64,6 +64,7 @@ struct onnx_parser
         add_variadic_op("Max", op::max{});
         add_variadic_op("Min", op::min{});
 
+        add_mem_op("LRN", &onnx_parser::parse_lrn);
         add_mem_op("ImageScaler", &onnx_parser::parse_imagescaler);
         add_mem_op("LeakyRelu", &onnx_parser::parse_leaky_relu);
         add_mem_op("Elu", &onnx_parser::parse_elu);
@@ -536,6 +537,25 @@ struct onnx_parser
         return prog.add_instruction(op, args.front());
     }
 
+    instruction_ref
+    parse_lrn(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
+    {
+        float alpha = 0.0001;
+        float beta  = 0.75;
+        float bias  = 1.0;
+        int size    = 1;
+        if(contains(attributes, "alpha"))
+            alpha = parse_value(attributes.at("alpha")).at<float>();
+        if(contains(attributes, "beta"))
+            beta = parse_value(attributes.at("beta")).at<float>();
+        if(contains(attributes, "bias"))
+            bias = parse_value(attributes.at("bias")).at<float>();
+        if(contains(attributes, "size"))
+            size = parse_value(attributes.at("size")).at<int>();
+        op::lrn op{alpha, beta, bias, size};
+        return prog.add_instruction(op, args.front());
+    }
+
     instruction_ref parse_imagescaler(const std::string&,
                                       attribute_map attributes,
                                       std::vector<instruction_ref> args)

src/targets/cpu/lowering.cpp

@@ -103,6 +103,43 @@ struct cpu_batch_norm_inference
     }
 };
 
+struct cpu_lrn
+{
+    op::lrn op;
+
+    std::string name() const { return "cpu::lrn"; }
+    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
+    argument compute(context&, shape output_shape, std::vector<argument> args) const
+    {
+        argument result{output_shape};
+        visit_all(result, args[0])([&](auto output, auto input) {
+            int n_batch         = output_shape.lens()[0];
+            int channels        = output_shape.lens()[1];
+            int height          = output_shape.lens()[2];
+            int width           = output_shape.lens()[3];
+            float alphaoverarea = op.alpha / op.size;
+            int radius          = (op.size - 1) / 2;
+
+            par_dfor(n_batch, height, width)([&](int b, int h, int w) {
+                float scale = 0;
+                dfor(channels)([&](int c) {
+                    auto start = (c - radius) < 0 ? 0 : (c - radius);
+                    auto end   = (c + radius) > channels ? channels : (c + radius);
+                    for(auto k = start; k < end; ++k)
+                    {
+                        scale += std::pow(input(b, k, h, w), 2);
+                    }
+                    scale *= alphaoverarea;
+                    scale += op.bias;
+                    scale              = std::pow(scale, -op.beta);
+                    output(b, c, h, w) = input(b, c, h, w) * scale;
+                });
+            });
+        });
+        return result;
+    }
+};
+
 struct cpu_convolution
 {
     op::convolution op;
@@ -681,6 +718,7 @@ struct cpu_apply
         apply_map["dot"]         = extend_op<cpu_gemm, op::dot>();
         apply_map["batch_norm_inference"] =
             extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
+        apply_map["lrn"]        = extend_op<cpu_lrn, op::lrn>();
         apply_map["contiguous"] = extend_op<cpu_contiguous, op::contiguous>();
         apply_map["pad"]        = extend_op<cpu_pad, op::pad>();
         apply_map["concat"]     = extend_op<cpu_concat, op::concat>();

src/targets/gpu/CMakeLists.txt

@@ -61,6 +61,7 @@ add_library(migraphx_gpu
     elu.cpp
     pad.cpp
     gather.cpp
+    lrn.cpp
 )
 set_target_properties(migraphx_gpu PROPERTIES EXPORT_NAME gpu)
 rocm_clang_tidy_check(migraphx_gpu)

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_LRN_HPP
+#define MIGRAPHX_GUARD_RTGLIB_LRN_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_lrn
+{
+    shared<lrn_descriptor> ldesc;
+    std::string name() const { return "gpu::lrn"; }
+    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/include/migraphx/gpu/miopen.hpp

@@ -23,6 +23,8 @@ using fusion_plan_descriptor = MIGRAPHX_MANAGE_PTR(miopenFusionPlanDescriptor_t,
                                                    miopenDestroyFusionPlan);
 using fused_operator_args    = MIGRAPHX_MANAGE_PTR(miopenOperatorArgs_t, miopenDestroyOperatorArgs);
 
+using lrn_descriptor = MIGRAPHX_MANAGE_PTR(miopenLRNDescriptor_t, miopenDestroyLRNDescriptor);
+
 template <class Result, class F, class... Ts>
 Result make_obj(F f, Ts... xs)
 {
@@ -89,6 +91,13 @@ inline pooling_descriptor make_pooling(const migraphx::op::pooling& op)
     return p;
 }
 
+inline lrn_descriptor make_lrn(const migraphx::op::lrn& op)
+{
+    auto ldesc = make_obj<lrn_descriptor>(&miopenCreateLRNDescriptor);
+    miopenSetLRNDescriptor(ldesc.get(), miopenLRNCrossChannel, op.size, op.alpha, op.beta, op.bias);
+    return ldesc;
+}
+
 inline activation_descriptor make_relu()
 {
     auto ad = make_obj<activation_descriptor>(&miopenCreateActivationDescriptor);

src/targets/gpu/lowering.cpp

@@ -43,6 +43,7 @@
 #include <migraphx/gpu/concat.hpp>
 #include <migraphx/gpu/pad.hpp>
 #include <migraphx/gpu/gather.hpp>
+#include <migraphx/gpu/lrn.hpp>
 #include <utility>
 #include <functional>
 #include <algorithm>
@@ -99,6 +100,7 @@ struct miopen_apply
         add_extend_op<hip_gather, op::gather>("gather");
         add_extend_op<hip_pad, op::pad>("pad");
 
+        add_lrn_op();
         add_convolution_op();
         add_pooling_op();
         add_batch_norm_inference_op();
@@ -159,6 +161,17 @@ struct miopen_apply
         });
     }
 
+    void add_lrn_op()
+    {
+        apply_map.emplace("lrn", [=](instruction_ref ins) {
+            auto&& op   = any_cast<op::lrn>(ins->get_operator());
+            auto ldesc  = make_lrn(op);
+            auto output = insert_allocation(ins, ins->get_shape());
+            return prog->replace_instruction(
+                ins, miopen_lrn{std::move(ldesc)}, ins->inputs().at(0), output);
+        });
+    }
+
     template <class T>
     void add_generic_op(std::string name)
     {

src/targets/gpu/lrn.cpp

+#include <migraphx/gpu/lrn.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_lrn::compute_shape(const std::vector<shape>& inputs) const
+{
+    check_shapes{inputs, *this}.has(2).not_broadcasted();
+    return inputs.at(1);
+}
+
+argument miopen_lrn::compute(context& ctx,
+                             const shape& output_shape,
+                             const std::vector<argument>& args) const
+{
+    float alpha = 1;
+    float beta  = 0;
+    auto x_desc = make_tensor(args[0].get_shape());
+    auto y_desc = make_tensor(output_shape);
+    miopenLRNForward(ctx.get_stream().get_miopen(),
+                     ldesc.get(),
+                     &alpha,
+                     x_desc.get(),
+                     args[0].implicit(),
+                     &beta,
+                     y_desc.get(),
+                     args[1].implicit(),
+                     false,
+                     nullptr);
+
+    return args[1];
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

test/cpu_ops_test.cpp

@@ -732,6 +732,20 @@ TEST_CASE(leaky_relu_test)
     EXPECT(migraphx::verify_range(results_vector, gold));
 }
 
+TEST_CASE(lrn_test)
+{
+    migraphx::program p;
+    migraphx::shape s{migraphx::shape::float_type, {1, 5, 1, 1}};
+    auto l = p.add_literal(migraphx::literal{s, {-2.0f, 1.0f, 0.f, 1.0f, 2.0f}});
+    p.add_instruction(migraphx::op::lrn{0.0001, 0.75, 1, 5}, l);
+    p.compile(migraphx::cpu::target{});
+    auto result = p.eval({});
+    std::vector<float> results_vector(5);
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold = {-2 / 1.000075, 1 / 1.00009, 0 / 1.000145, 1 / 1.00009, 2 / 1.000075};
+    EXPECT(migraphx::verify_range(results_vector, gold));
+}
+
 TEST_CASE(imagescaler_test)
 {
     migraphx::program p;

test/gpu/miopen.cpp

@@ -669,6 +669,18 @@ struct test_elu
     }
 };
 
+struct test_relu_lrn
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto x = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {1, 5, 2, 2}});
+        auto y = p.add_instruction(migraphx::op::relu{}, x);
+        p.add_instruction(migraphx::op::lrn{0.0001, 0.75, 1.0, 5}, y);
+        return p;
+    }
+};
+
 struct test_conv_pooling
 {
     migraphx::program create_program() const
@@ -1506,6 +1518,7 @@ struct test_rnn_bi_3args
 
 int main()
 {
+    verify_program<test_relu_lrn>();
     verify_program<test_pooling_autopad>();
     verify_program<test_abs>();
     verify_program<test_concat>();

test/onnx/onnx_test.cpp

@@ -774,4 +774,18 @@ TEST_CASE(pad_test)
     migraphx::parse_onnx("pad_test.onnx");
 }
 
+TEST_CASE(lrn_test)
+{
+    migraphx::program p;
+
+    auto l0 = p.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 28, 24, 24}});
+    migraphx::op::lrn op;
+    op.size  = 5;
+    op.alpha = 0.0001;
+    op.beta  = 0.75;
+    op.bias  = 1.0;
+    p.add_instruction(op, l0);
+    migraphx::parse_onnx("lrn_test.onnx");
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }