Support Repeat_M in distribution

include/ck_tile/ops/layernorm2d/kernel/layernorm2d_fwd_shape.hpp

@@ -15,9 +15,9 @@ namespace ck_tile {
         +<----------------------< Repeat_N(2)>--------------------->+
         |                                                           |
         +<--    <WarpPerBlock_N(2)>  -->+
-            Warp_M
+            Warp_N
         +--------------+--------------+--------------+--------------+----+----------------+
- Warp_N | wrap_0       | wrap_1       |                             |    ^                ^
+ Warp_M | wrap_0       | wrap_1       |                             |    ^                ^
         +--------------+--------------+                             |   <WarpPerBlock_M(2)> |
         | wrap_2       | wrap_3       |                             |    v
         +--------------+--------------+--------------+--------------+----+           Block_M

include/ck_tile/ops/layernorm2d/pipeline/layernorm2d_fwd_pipeline_default_policy.hpp

@@ -19,12 +19,12 @@ struct Layernorm2dFwdPipelineDefaultPolicy
         return make_static_tile_distribution(
             tile_distribution_encoding<
                 sequence<>,
-                tuple<sequence<S::WarpPerBlock_M, S::ThreadPerWarp_M, S::Vector_M>,
+                tuple<sequence<S::Repeat_M, S::WarpPerBlock_M, S::ThreadPerWarp_M, S::Vector_M>,
                       sequence<S::Repeat_N, S::WarpPerBlock_N, S::ThreadPerWarp_N, S::Vector_N>>,
                 tuple<sequence<1, 2>, sequence<1, 2>>,
-                tuple<sequence<0, 1>, sequence<1, 2>>,
-                sequence<1, 2, 2>,
-                sequence<2, 0, 3>>{});
+                tuple<sequence<1, 1>, sequence<2, 2>>,
+                sequence<1, 1, 2, 2>,
+                sequence<0, 3, 0, 3>>{});
     }
     template <typename Problem>
     CK_TILE_DEVICE static constexpr auto MakeGammaBetaBlockTileDistribution()