[BUG] Unable to convert output dtype from FP32 to BF16 in Group GEMM epilogue

[sanchitintel]

Which component has the problem?

CUTLASS C++

Bug Report

Describe the bug
With BF16 A, B matrices, try computing Group GEMM. Use output dtype as BFloat16 (convert in epilogue).
Did not work with epilogue created directly with cutlass::epilogue::collective::CollectiveEpilogue.

Steps/Code to reproduce bug
Please apply this small diff & compile Group GEMM example (the same file)

diff --git a/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp b/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
index bdda0536..860cbf0c 100644
--- a/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
+++ b/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
@@ -92,7 +92,7 @@ using ElementAccumulator = float; // <- data type of accumulator
using ElementComputeEpilogue = float; // <- data type of epilogue operations
using ElementA = bfloat16_t; // <- data type of elements in input matrix A
using ElementB = bfloat16_t; // <- data type of elements in input matrix B
-using ElementOutput = float; // <- data type of elements in output matrix D
+using ElementOutput = bfloat16_t; // <- data type of elements in output matrix D

///////////////////////////////////////////////////////////////////////////////////////////////////

@@ -198,8 +198,8 @@ struct ExampleRunner {
using LayoutD = typename Gemm::LayoutD;

using CollectiveEpilogue = typename Gemm::CollectiveEpilogue;

• using ElementOutput = typename CollectiveEpilogue::ElementOutput;
• using ElementAccumulator = ElementOutput;

• using ElementOutput = bfloat16_t;

• using ElementAccumulator = float_t;

  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
  @@ -361,7 +361,7 @@ void initialize(const Options &options) {
  std::vector<ElementA *> ptr_A_host(options.groups);
  std::vector<ElementB *> ptr_B_host(options.groups);
  std::vector<ElementC *> ptr_C_host(options.groups);

• std::vector<ElementC *> ptr_D_host(options.groups);

• std::vector<ElementOutput *> ptr_D_host(options.groups);
  std::vector<ElementAccumulator *> ptr_alpha_host(options.groups);
  std::vector<ElementAccumulator *> ptr_beta_host(options.groups);

@@ -599,7 +599,7 @@ int main(int argc, const char** argv)
FusionCallBacks,
XE_2D_U32x8x16_LD_N,
void, void,

•      XE_2D_U32x8x16_ST_N,
  

•      XE_2D_U16x8x16_ST_N,
       void, void>;
  

Expected behavior
Dtype conversion of MMA output should be supported in epilogue

Environment details (please complete the following information):
PVC GPU

Additional context
Main branch

[sanchitintel]

This resolution of this issue requires modifying one line in include/cutlass/epilogue/collective/builders/xe_builder.inl.

Alternatively, a small change can also be made to an example using this API, but the change in include/cutlass/epilogue/collective/builders/xe_builder.inl seems to make more sense, since a UT using that API is using IntelXeXMX16 instead of IntelXeXMX16Group, which doesn't seem to be a good idea from usage perspective, as it's non-intuitive:

https://github.com/intel/cutlass-sycl/blob/e83f147263dd8ca3589b34d76ce6fbec58bac048/test/unit/gemm/device/default_gemm_group_configuration.hpp#L93-L94

I'll submit a PR with an example/UT with a similar diff as below:

diff --git a/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp b/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
diff --git a/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp b/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
index bdda0536..02ad8c86 100644
--- a/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
+++ b/examples/sycl/04_bmg_grouped_gemm/04_bmg_grouped_gemm.cpp
@@ -69,6 +69,7 @@
 #include "cutlass/gemm/group_array_problem_shape.hpp"
 #include "cutlass/gemm/device/gemm_universal.h"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/collective/collective_mma.hpp"
 #include "cutlass/util/GPU_Clock.hpp"

@@ -92,7 +93,7 @@ using ElementAccumulator = float;     // <- data type of accumulator
 using ElementComputeEpilogue = float; // <- data type of epilogue operations
 using ElementA = bfloat16_t;          // <- data type of elements in input matrix A
 using ElementB = bfloat16_t;          // <- data type of elements in input matrix B
-using ElementOutput = float;          // <- data type of elements in output matrix D
+using ElementOutput = bfloat16_t;          // <- data type of elements in output matrix D

 ///////////////////////////////////////////////////////////////////////////////////////////////////

@@ -198,8 +199,8 @@ struct ExampleRunner {
   using LayoutD = typename Gemm::LayoutD;

   using CollectiveEpilogue = typename Gemm::CollectiveEpilogue;
-  using ElementOutput = typename CollectiveEpilogue::ElementOutput;
-  using ElementAccumulator = ElementOutput;
+  using ElementOutput = bfloat16_t;
+  using ElementAccumulator = float_t;

   using StrideA = typename Gemm::GemmKernel::InternalStrideA;
   using StrideB = typename Gemm::GemmKernel::InternalStrideB;
@@ -361,7 +362,7 @@ void initialize(const Options &options) {
   std::vector<ElementA *> ptr_A_host(options.groups);
   std::vector<ElementB *> ptr_B_host(options.groups);
   std::vector<ElementC *> ptr_C_host(options.groups);
-  std::vector<ElementC *> ptr_D_host(options.groups);
+  std::vector<ElementOutput *> ptr_D_host(options.groups);
   std::vector<ElementAccumulator *> ptr_alpha_host(options.groups);
   std::vector<ElementAccumulator *> ptr_beta_host(options.groups);

@@ -584,23 +585,25 @@ int main(int argc, const char** argv)
   using GEMMDispatchPolicy = cutlass::gemm::MainloopIntelXeXMX16Group<PipelineStages>;
   using EpilogueDispatchPolicy = cutlass::epilogue::IntelXeXMX16Group;

-  using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<ElementOutput, ElementComputeEpilogue,
-          ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
+  using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<float_t, float_t>;

-  using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<EpilogueDispatchPolicy, EpilogueOp, TileShape,
-          decltype(tile_shape(TiledMma()))>;
-  using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue<
-          EpilogueDispatchPolicy,
-          TileShape,
-          ElementAccumulator,
-          cutlass::gemm::TagToStrideC_t<LayoutC*>,
-          ElementOutput,
-          cutlass::gemm::TagToStrideC_t<LayoutD*>,
-          FusionCallBacks,
-          XE_2D_U32x8x16_LD_N,
-          void, void,
-          XE_2D_U32x8x16_ST_N,
-          void, void>;
+  using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks<
+    cutlass::epilogue::IntelXeXMX16Group,
+    EpilogueOp,
+    TileShape,
+    decltype(tile_shape(TiledMma()))
+  >;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      cutlass::arch::IntelXe, cutlass::arch::OpClassTensorOp,
+      TileShape, Shape<_1, _1, _1>,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      float, float,
+      float, LayoutC, 1,
+      ElementOutput, LayoutC, 1,
+      EpilogueDispatchPolicy,
+      EpilogueOp
+    >::CollectiveOp;

 // Mainloop
   using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
diff --git a/include/cutlass/epilogue/collective/builders/xe_builder.inl b/include/cutlass/epilogue/collective/builders/xe_builder.inl
index 809cede6..b4b33580 100644
--- a/include/cutlass/epilogue/collective/builders/xe_builder.inl
+++ b/include/cutlass/epilogue/collective/builders/xe_builder.inl
@@ -193,7 +193,7 @@ template <

       //TODO(Codeplay): Should FusionCallbacks use DispatchPolicy IntelXeGroupEpilogue for group gemm? That does not work.
       using FusionCallbacks = typename detail::FusionOpInfo<FusionOpOrCallbacks>::template FusionCallbacks<
-                                  IntelXeXMX16, TileShape_MNK, TileShape_MNK, CopyOpG2R>;
+                                  std::conditional_t<IsGroup, IntelXeXMX16Group, IntelXeXMX16>, TileShape_MNK, TileShape_MNK, CopyOpG2R>;

       using CollectiveOp = cutlass::epilogue::collective::CollectiveEpilogue<
             DispatchPolicy,