initial work on lrn

src/include/migraph/operators.hpp

@@ -51,6 +51,27 @@ struct batch_norm_inference
     }
 };
 
+struct lrn
+{
+    float alpha = 0.0001;
+    float beta = 0.75;
+    float bias = 1.0;
+    int size;
+    std::string name() const { return "lrn"; }
+
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(1);
+        return inputs.front();
+    }
+
+    friend std::ostream& operator<<(std::ostream& os, const lrn& op)
+    {
+        os << op.name() << ":" << op.alpha << ":" << op.beta << ":" << op.bias << ":" << op.size;
+        return os;
+    }
+};
+
 struct convolution
 {
     std::array<std::size_t, 2> padding  = {{0, 0}};

src/targets/cpu/lowering.cpp

@@ -94,6 +94,46 @@ 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];
+            auto alphaoverarea = op.alpha / op.size;
+            auto radius = (op.size - 1) / 2;
+
+            dfor(n_batch, height, width)([&](int b, int h, int w){
+                double 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, c, 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;
@@ -596,6 +636,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["concat"]     = extend_op<cpu_concat, op::concat>();
         apply_map["leaky_relu"] = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();

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

+#ifndef MIGRAPH_GUARD_RTGLIB_LRN_HPP
+#define MIGRAPH_GUARD_RTGLIB_LRN_HPP
+
+#include <migraph/gpu/lowering.hpp>
+#include <migraph/manage_ptr.hpp>
+#include <migraph/instruction.hpp>
+#include <migraph/operators.hpp>
+#include <migraph/generate.hpp>
+#include <migraph/shape_for_each.hpp>
+#include <migraph/config.hpp>
+#include <migraph/gpu/miopen.hpp>
+#include <migraph/gpu/hip.hpp>
+#include <migraph/dfor.hpp>
+#include <migraph/gpu/device/contiguous.hpp>
+#include <migraph/gpu/device/add.hpp>
+#include <migraph/iterator_for.hpp>
+#include <migraph/gpu/rocblas.hpp>
+#include <migraph/gpu/context.hpp>
+#include <utility>
+
+namespace migraph {
+inline namespace MIGRAPH_INLINE_NS {
+namespace gpu {
+
+struct miopen_lrn
+{
+    shared<lrn_descriptor> ld;
+    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 MIGRAPH_INLINE_NS
+} // namespace migraph
+
+#endif

src/targets/gpu/include/migraph/gpu/miopen.hpp

@@ -22,6 +22,7 @@ using activation_descriptor  = MIGRAPH_MANAGE_PTR(miopenActivationDescriptor_t,
 using fusion_plan_descriptor = MIGRAPH_MANAGE_PTR(miopenFusionPlanDescriptor_t,
                                                   miopenDestroyFusionPlan);
 using fused_operator_args    = MIGRAPH_MANAGE_PTR(miopenOperatorArgs_t, miopenDestroyOperatorArgs);
+using lrn_descriptor         = MIGRAPH_MANAGE_PTR(miopenLRNDescriptor_t, miopenDestroyLRNDescriptor);
 
 template <class Result, class F, class... Ts>
 Result make_obj(F f, Ts... xs)