[Xe] [Reorder] Support broadcasting reorders (#589)

Adds support for broadcasting in reorders, where a single src value maps
to multiple dst values.

This is useful for reordering scales/zero points during dequantization.

Author: petercad

include/cute/algorithm/reorder.hpp

@@ -165,21 +165,41 @@ reorder_impl(ReorderAtom               const& atom,
   constexpr int values = size(SLayout{}) / size<0>(SLayout{});
   constexpr int vchunk = sizeof_bits_v<RegistersSrc> / sizeof_bits_v<SType>;
 
-  // Calculate mapping from src val -> dst val on a chunk-by-chunk basis. Unlike a plain copy, there is no intrinsic
-  //   correspondence of src/dst values for subgroup reorders.
-  auto rlayout = coalesce(composition(right_inverse(dlayout), slayout));                 // src index -> dst index
-  auto vrlayout = composition(composition(Layout<Shape<_SG, Int<values>>, Stride<_0, _1>>{},
-                                          rlayout),
-                              Layout<Shape<_1, Int<values>>, Stride<_0, _SG>>{});        // src val -> dst val
-
-  CUTE_UNROLL
-  for (int sv = 0; sv < values; sv += vchunk) {
-    auto pS = recast_ptr<RegTypeSrc>(src.data() + sv);
-    auto pD = recast_ptr<RegTypeDst>(dst.data() + vrlayout(sv));
-
-    detail::explode(detail::CallReorder<ReorderAtom>{},
-                    pS, make_int_sequence<RegNumSrc>{},
-                    pD, make_int_sequence<RegNumDst>{});
+  static constexpr bool has_broadcast = (size(DLayoutWI{}) > size(SLayoutWI{}));
+
+  if (!has_broadcast) {
+    // Calculate mapping from src val -> dst val on a chunk-by-chunk basis. Unlike a plain copy, there is no intrinsic
+    //   correspondence of src/dst values for subgroup reorders.
+    auto rlayout = coalesce(composition(right_inverse(dlayout), slayout));                 // src index -> dst index
+    auto vrlayout = composition(composition(Layout<Shape<_SG, Int<values>>, Stride<_0, _1>>{},
+                                            rlayout),
+                                Layout<Shape<_1, Int<values>>, Stride<_0, _SG>>{});        // src val -> dst val
+
+    CUTE_UNROLL
+    for (int sv = 0; sv < values; sv += vchunk) {
+      auto pS = recast_ptr<RegTypeSrc>(src.data() + sv);
+      auto pD = recast_ptr<RegTypeDst>(dst.data() + vrlayout(sv));
+
+      detail::explode(detail::CallReorder<ReorderAtom>{},
+                      pS, make_int_sequence<RegNumSrc>{},
+                      pD, make_int_sequence<RegNumDst>{});
+    }
+  } else {
+    // If there is broadcast happening, then we need to loop over dst values instead.
+    auto rlayout = coalesce(composition(right_inverse(slayout), dlayout));                 // dst index -> src index
+    auto vrlayout = composition(composition(Layout<Shape<_SG, Int<values>>, Stride<_0, _1>>{},
+                                            rlayout),
+                                Layout<Shape<_1, Int<values>>, Stride<_0, _SG>>{});        // dst val -> src val
+
+    CUTE_UNROLL
+    for (int dv = 0; dv < values; dv += vchunk) {
+      auto pS = recast_ptr<RegTypeSrc>(src.data() + vrlayout(dv));
+      auto pD = recast_ptr<RegTypeDst>(dst.data() + dv);
+
+      detail::explode(detail::CallReorder<ReorderAtom>{},
+                      pS, make_int_sequence<RegNumSrc>{},
+                      pD, make_int_sequence<RegNumDst>{});
+    }
   }
 }
 

include/cute/atom/reorder_atom_xe.hpp

@@ -190,6 +190,7 @@ reorder_impl(ReorderDispatchXeGeneric  const&,
   static constexpr int elems_per_grf = 64 / sizeof(SrcType);
   static constexpr int ds_vl = cute::min(32, cute::min(shape<0>(rlayout), elems_per_grf / stride<0>(rlayout)));
   static constexpr int ss_vl = cute::min(32, cute::min(shape<0>(ilayout), elems_per_grf / stride<0>(ilayout)));
+  static constexpr bool has_broadcast = (size(DLayoutWI{}) > size(SLayoutWI{}));
 
   // Make dst live, to prevent compiler from inserting its own initialization.
 #ifdef __SYCL_DEVICE_ONLY__
@@ -202,7 +203,7 @@ reorder_impl(ReorderDispatchXeGeneric  const&,
   }
 #endif
 
-  if constexpr (ss_vl >= ds_vl) {
+  if constexpr (ss_vl >= ds_vl || has_broadcast) {
     // Stride on src. For simplicity, take 1 GRF at a time.
     for_each(make_seq<size(SLayout{}) / ss_vl>{}, [&](auto i) {
       constexpr auto didx = i * ss_vl;