Lower lrn to pooling

src/rewrite_pooling.cpp

@@ -30,6 +30,8 @@
 #include <migraphx/op/reduce_max.hpp>
 #include <migraphx/make_op.hpp>
 
+#include <migraphx/op/lrn.hpp>
+#include <migraphx/permutation.hpp>
 #include <migraphx/program.hpp>
 
 namespace migraphx {
@@ -55,6 +57,75 @@ static void replace_with_reduce(module& m, instruction_ref ins)
     }
 }
 
+static void lower_lrn_to_pooling(module& m, instruction_ref ins)
+{
+    auto v = ins->get_operator().to_value();
+
+    float alpha             = v.at("alpha").to<float>();
+    float beta              = v.at("beta").to<float>();
+    float k                 = v.at("bias").to<float>();
+    int size                = v.at("size").to<int>();
+    const unsigned int axis = 1;
+
+    auto x             = ins->inputs().at(0);
+    const auto& xshape = x->get_shape();
+    auto lens          = xshape.lens();
+    const int64_t rank = static_cast<int64_t>(lens.size());
+    if(rank < 2)
+        return;
+    if(size <= 0)
+        return;
+
+    auto x2 = m.insert_instruction(ins, make_op("mul"), x, x);
+
+    std::vector<int64_t> perm(rank);
+    std::iota(perm.begin(), perm.end(), 0);
+    std::swap(perm[static_cast<std::size_t>(axis)], perm.back());
+    auto moved      = m.insert_instruction(ins, make_op("transpose", {{"permutation", perm}}), x2);
+    auto moved_lens = moved->get_shape().lens();
+    auto b =
+        std::accumulate(moved_lens.begin(), moved_lens.end() - 1, 1, std::multiplies<size_t>());
+    const int64_t c = static_cast<int64_t>(moved_lens.back());
+    auto pooled_in  = m.insert_instruction(
+        ins,
+         make_op("reshape", {{"dims", std::vector<int64_t>{static_cast<int64_t>(b), 1, 1, c}}}),
+         moved);
+
+    auto avg = m.insert_instruction(ins,
+                                    make_op("pooling",
+                                            {{"mode", op::pooling_mode::average},
+                                             {"lengths", std::vector<int64_t>{1, size}},
+                                             {"stride", std::vector<int64_t>{1, 1}},
+                                             {"padding", std::vector<int64_t>{0, size / 2}},
+                                             {"count_include_pad", true}}),
+                                    pooled_in);
+
+    auto moved_shape_back = std::vector<int64_t>(moved_lens.begin(), moved_lens.end());
+    auto avg_moved =
+        m.insert_instruction(ins, make_op("reshape", {{"dims", moved_shape_back}}), avg);
+
+
+    auto invp = invert_permutation(perm);
+    auto avg_ch =
+        m.insert_instruction(ins, make_op("transpose", {{"permutation", invp}}), avg_moved);
+
+    auto k_lit = m.add_literal(k);
+    auto a_lit = m.add_literal(alpha);
+    auto k_mb  = m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), k_lit);
+    auto a_mb  = m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), a_lit);
+    auto alpha_avg = m.insert_instruction(ins, make_op("mul"), a_mb, avg_ch);
+    auto den       = m.insert_instruction(ins, make_op("add"), k_mb, alpha_avg);
+
+    auto b_lit  = m.add_literal(beta);
+    auto b_mb   = m.insert_instruction(ins, make_op("multibroadcast", {{"out_lens", lens}}), b_lit);
+    auto denpow = m.insert_instruction(ins, make_op("pow"), den, b_mb);
+    auto y      = m.insert_instruction(ins, make_op("div"), ins->inputs().front(), denpow);
+
+    m.replace_instruction(ins, y);
+}
+
+
+
 static void replace_dilations_with_gather_pooling(module& m, instruction_ref ins)
 {
     // TODO remove this when MIOpen supports dilated pooling
@@ -143,10 +214,16 @@ void rewrite_pooling::apply(module& m) const
 {
     for(auto ins : iterator_for(m))
     {
-        if(ins->name() != "pooling")
-            continue;
         if(ins->inputs().empty())
+            continue;  
+        if(ins->name() == "lrn")
+        {
+            lower_lrn_to_pooling(m, ins);
             continue;
+        }
+        if(ins->name() != "pooling")
+            continue;
+
         auto&& s                  = ins->inputs().front()->get_shape();
         auto&& op                 = any_cast<op::pooling>(ins->get_operator());
         bool same_kernel_as_shape = std::equal(

test/rewrite_pooling_test.cpp

@@ -34,6 +34,8 @@
 
 #include <migraphx/verify.hpp>
 
+#include <migraphx/iterator.hpp>
+
 static void opt_pooling(migraphx::module& m)
 {
     migraphx::rewrite_pooling rp;
@@ -309,6 +311,64 @@ TEST_CASE(rewrite_pooling_dialtions_test5)
     test_rewrite(migraphx::op::pooling_mode::max);
 }
 
+TEST_CASE(lower_lrn_to_pooling)
+{
+    migraphx::shape s{migraphx::shape::float_type, {1, 64, 55, 55}};
+    migraphx::module m1;
+    {
+        auto x   = m1.add_parameter("x", s);
+        auto lrn = m1.add_instruction(
+            migraphx::make_op("lrn",
+                              {{"alpha", 0.0001}, {"beta", 0.75}, {"bias", 1.0}, {"size", 5}}),
+            x);
+        m1.add_return({lrn});
+    }
+    opt_pooling(m1);
+
+    migraphx::module m2;
+    {
+        auto x = m2.add_parameter("x", s);
+
+        auto x_squared  = m2.add_instruction(migraphx::make_op("mul"), x, x);
+        auto transpose1 = m2.add_instruction(
+            migraphx::make_op("transpose", {{"permutation", {0, 3, 2, 1}}}), x_squared);
+        auto reshape1 = m2.add_instruction(
+            migraphx::make_op("reshape", {{"dims", {3025, 1, 1, 64}}}), transpose1);
+        auto pooling =
+            m2.add_instruction(migraphx::make_op("pooling",
+                                                 {{"mode", migraphx::op::pooling_mode::average},
+                                                  {"lengths", std::vector<int64_t>{1, 5}},
+                                                  {"stride", std::vector<int64_t>{1, 1}},
+                                                  {"padding", std::vector<int64_t>{0, 2}},
+                                                  {"count_include_pad", true}}),
+                               reshape1);
+        auto reshape2 =
+            m2.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 55, 55, 64}}}), pooling);
+        auto transpose2 = m2.add_instruction(
+            migraphx::make_op("transpose", {{"permutation", {0, 3, 2, 1}}}), reshape2);
+
+        auto beta_lit  = m2.add_literal(migraphx::literal{migraphx::shape::float_type, {0.75}});
+        auto alpha_lit = m2.add_literal(migraphx::literal{migraphx::shape::float_type, {0.0001}});
+        auto bias_lit  = m2.add_literal(migraphx::literal{migraphx::shape::float_type, {1.0}});
+
+        auto bias_mb = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 64, 55, 55}}}), bias_lit);
+        auto alpha_mb = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 64, 55, 55}}}), alpha_lit);
+        auto beta_mb = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 64, 55, 55}}}), beta_lit);
+
+        auto alpha_avg   = m2.add_instruction(migraphx::make_op("mul"), alpha_mb, transpose2);
+        auto denominator = m2.add_instruction(migraphx::make_op("add"), bias_mb, alpha_avg);
+        auto powered     = m2.add_instruction(migraphx::make_op("pow"), denominator, beta_mb);
+        auto result      = m2.add_instruction(migraphx::make_op("div"), x, powered); 
+
+        m2.add_return({result});
+    }
+
+    EXPECT(m1.sort() == m2.sort());
+}
+
 TEST_CASE(rewrite_avgpool_rank3_dil_test)
 {
     // 1D case 1, input is 3D