Use parse_layernorm to un-fuse LayerNormalization op

src/onnx/parse_attention.cpp

@@ -31,7 +31,6 @@ struct parse_attention : op_parser<parse_attention>
         auto input_lens        = input->get_shape().lens();
         auto batch_size        = input_lens.at(0);
         auto sequence_length   = input_lens.at(1);
-        auto input_hidden_size = input_lens.at(2);
 
         // bias shape= (3 * hidden_size)
         auto bias_lens   = bias->get_shape().lens();

src/onnx/parse_layernorm.cpp

@@ -18,7 +18,14 @@ struct parse_layernorm : op_parser<parse_layernorm>
                           onnx_parser::node_info info,
                           const std::vector<instruction_ref>& args) const
     {
-        float epsilon = 1e-3f;
+        // un-fuse layernorm op so migraphx can handle fusion instead
+        
+        auto x        = args.front(); 
+        auto x_type   = x->get_shape().type();
+        auto weights        = args.at(1); 
+        auto bias        = args.at(2); 
+
+        float epsilon = 1e-12f;
         int64_t axis  = -1;
         if(contains(info.attributes, "epsilon"))
         {
@@ -28,17 +35,25 @@ struct parse_layernorm : op_parser<parse_layernorm>
         {
             epsilon = parser.parse_value(info.attributes.at("axis")).at<int64_t>();
         }
+    
+        auto epsilon_lit  = info.add_literal(literal{shape{x_type, {1}}, {epsilon}});
+        auto exponent = info.add_literal(literal{shape{x_type, {1}}, {2.0}});
+        auto dims = x->get_shape().lens();
 
-        auto layernorm = info.add_instruction(
-            make_op("layernorm", {{"epsilon", epsilon}, {"axis", axis}}), args.front());
-
-        if(args.size() == 3)
-        {
-            layernorm = info.add_broadcastable_binary_op("mul", layernorm, args.at(1));
-            layernorm = info.add_broadcastable_binary_op("add", layernorm, args.at(2));
-        }
+        auto mean = info.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {axis}}}), x);
+        auto mean_mbcast =
+            info.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), mean);
+        auto sub = info.add_instruction(migraphx::make_op("sub"), x, mean_mbcast);
+        auto exponent_mbcast =
+            info.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", dims}}), exponent);
+        auto pow = info.add_instruction(migraphx::make_op("pow"), sub, exponent_mbcast);
+        auto var = info.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {axis}}}), pow);
+        auto add_epsilon = info.add_broadcastable_binary_op("add", var, epsilon_lit);
+        auto sqrt = info.add_instruction(migraphx::make_op("sqrt"), add_epsilon);
+        auto div = info.add_broadcastable_binary_op("div", sub, sqrt);
+        auto mul = info.add_broadcastable_binary_op("mul", div, weights);
 
-        return layernorm;
+        return info.add_broadcastable_binary_op("add", mul, bias);
     }
 };