[Xe] 4-bit unit stride -> VNNI reorders (#593)

Adds support for unit -> VNNI reorders for 4-bit types, to enable
row-major int4 B.

Note that this is an expensive reorder sequence (9 cycles/register), and
it significantly impacts performance for current mainloops. We'll need
to enable an SLM-based mainloop to reduce the cost of this reorder.

Author: petercad

examples/cute/tutorial/xe_gemm.cpp

@@ -414,6 +414,7 @@ int main(int argc, char** argv)
   test_case<int4_t, uint8_t, int32_t, 'R', 'C'>(Q, m, n, k);
 
   test_case<uint4_t, uint4_t, uint32_t, 'R', 'C'>(Q, m, n, k);
+  test_case<uint4_t, uint4_t, uint32_t, 'R', 'R'>(Q, m, n, k);
 
   // Upconversion cases
   test_case<half_t, float_e5m2_t, float, 'R', 'R'>(Q, m, n, k);

include/cute/arch/reorder_xe.hpp

@@ -1264,5 +1264,48 @@ struct Xe_Reorder<ReorderKind::UU, float_ue8m0_t, float>
   }
 };
 
+template <>
+struct Xe_Reorder<ReorderKind::UV, uint4_t, uint4_t>
+{
+  using SRegisters = intel::uchar4[1];
+  using DRegisters = intel::uchar4[1];
+
+  CUTE_HOST_DEVICE static void
+  reorder(intel::uchar4 const& src0, intel::uchar4& dst0)
+  {
+#if defined(CUTE_ARCH_COPY_XE_ENABLED)
+    const uint32_t lshifts = 0x00000004;
+    const uint32_t rshifts = 0x00040000;
+    asm (     /* 9 cycles/output register */
+      "{\n"
+      ".decl IN_UB  v_type=G type=UB num_elts=64 alias=<%1,0>\n"
+      ".decl OUT_UB v_type=G type=UB num_elts=64 alias=<%0,0>\n"
+      ".decl OUT_UW v_type=G type=UW num_elts=32 alias=<%0,0>\n"
+      ".decl LSHIFTS v_type=G type=UW num_elts=2 alias=<%2,0>\n"
+      ".decl RSHIFTS v_type=G type=UW num_elts=2 alias=<%3,0>\n"
+      ".decl TMP_UB v_type=G type=UB num_elts=64 align=64\n"
+      ".decl TMP_UW v_type=G type=UW num_elts=32 alias=<TMP_UB,0>\n"
+      "shr (M1_NM, 16) OUT_UB(0,0)<4> IN_UB(0, 0)<1;2,0> RSHIFTS(0,0)<0;2,1>\n"
+      "shr (M1_NM, 16) OUT_UB(0,1)<4> IN_UB(0,16)<1;2,0> RSHIFTS(0,0)<0;2,1>\n"
+      "shr (M1_NM, 16) OUT_UB(0,2)<4> IN_UB(0,32)<1;2,0> RSHIFTS(0,0)<0;2,1>\n"
+      "shr (M1_NM, 16) OUT_UB(0,3)<4> IN_UB(0,48)<1;2,0> RSHIFTS(0,0)<0;2,1>\n"
+      "shl (M1_NM, 16) TMP_UB(0,0)<4> IN_UB(0, 8)<1;2,0> LSHIFTS(0,0)<0;2,1>\n"
+      "shl (M1_NM, 16) TMP_UB(0,1)<4> IN_UB(0,24)<1;2,0> LSHIFTS(0,0)<0;2,1>\n"
+      "shl (M1_NM, 16) TMP_UB(0,2)<4> IN_UB(0,40)<1;2,0> LSHIFTS(0,0)<0;2,1>\n"
+      "shl (M1_NM, 16) TMP_UB(0,3)<4> IN_UB(0,56)<1;2,0> LSHIFTS(0,0)<0;2,1>\n"
+      "bfn.xCA (M1_NM, 32) OUT_UW(0,0)<1> OUT_UW(0,0)<1;1,0> TMP_UW(0,0)<1;1,0> 0xF0F0:uw\n"
+      "}\n"
+      : "=rw"(dst0)
+      : "rw"(src0), "rw.u"(lshifts), "rw.u"(rshifts)
+    );
+#else
+  CUTE_INVALID_CONTROL_PATH("Not Xe");
+#endif
+  }
+};
+
+template <> struct Xe_Reorder<ReorderKind::UV, int4_t, int4_t>             : Xe_Reorder<ReorderKind::UV, uint4_t, uint4_t> {};
+template <> struct Xe_Reorder<ReorderKind::UV, float_e2m1_t, float_e2m1_t> : Xe_Reorder<ReorderKind::UV, uint4_t, uint4_t> {};
+
 
 } // end namespace cute