Fix for void ElementC in epilogue.

Todo: Handle void ElementD
Check if we need to handle assumption that accumulator is same as compute.

Signed-off-by: Chawla, Amit K <[email protected]>

Author: amitchawla1

include/cutlass/epilogue/collective/builders/xe_builder.inl

@@ -195,7 +195,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>;
-
+      static_assert(cute::is_same_v<ElementAccumulator, ElementCompute>, "Assumption is that in Epilogue, ElementAccumulator and ElementCompute are same");
       using CollectiveOp = cutlass::epilogue::collective::CollectiveEpilogue<
             DispatchPolicy,
             TileShape_MNK,

include/cutlass/epilogue/collective/xe_array_epilogue.hpp

@@ -90,7 +90,6 @@ class CollectiveEpilogue<
   using DispatchPolicy = IntelXeXMX16Group;
   using CtaTileMNK = CtaTileMNK_;
   using FusionCallbacks = FusionCallbacks_;
-  using ElementC = ElementC_;
   using StrideC = StrideC_;
   using InternalStrideC = cute::remove_pointer_t<StrideC>;
   using ElementD = ElementD_;
@@ -110,6 +109,7 @@ class CollectiveEpilogue<
   using ElementOutput = ElementD;
   using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
   using ElementAccumulator = ElementCompute;
+  using ElementC = conditional_t<cute::is_void_v<CopyOpG2R>, ElementAccumulator, ElementC_>;
   using ElementSource = typename FusionCallbacks::ElementSource;
   using ElementScalar = typename FusionCallbacks::ElementScalar;
   static constexpr FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest;
@@ -139,7 +139,7 @@ class CollectiveEpilogue<
                                              Layout<CopyThreadShape>{},
                                              make_layout(shape_div(typename Trait_D::BlockShape{}, CopyThreadShape{}))));
 private:
-  constexpr static bool is_source_supported = not cute::is_void_v<ElementC>;
+  constexpr static bool is_source_supported = not cute::is_void_v<ElementC_>;
   constexpr static bool is_destination_supported = not cute::is_void_v<ElementD> && not cute::is_void_v<CopyOpR2G>;
 
 public:

include/cutlass/epilogue/collective/xe_epilogue.hpp

@@ -89,7 +89,6 @@ class CollectiveEpilogue<
   using DispatchPolicy = IntelXeXMX16;
   using CtaTileMNK = CtaTileMNK_;
   using FusionCallbacks = FusionCallbacks_;
-  using ElementC = ElementC_;
   using StrideC = StrideC_;
   using ElementD = ElementD_;
   using StrideD = StrideD_;
@@ -107,6 +106,7 @@ class CollectiveEpilogue<
   using ElementOutput = ElementD;
   using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
   using ElementAccumulator = ElementCompute;
+  using ElementC = conditional_t<cute::is_void_v<CopyOpG2R>, ElementAccumulator, ElementC_>;
   static constexpr int SubgroupSize = DispatchPolicy::SubgroupSize;
 
   static_assert(cute::rank(CtaTileMNK{}) == 3, "CtaTileMNK must be rank-3: [CTA_M, CTA_N, CTA_K]");
@@ -129,7 +129,7 @@ class CollectiveEpilogue<
   using XE_Copy_D = decltype(make_tiled_copy(Copy_Atom<Trait_D, ElementD>{}, Layout<CopyThreadShape>{}, val_layout_store_D{}));
 
 private:
-  constexpr static bool is_source_supported = not cute::is_void_v<ElementC> && not cute::is_void_v<CopyOpG2R>;
+  constexpr static bool is_source_supported = not cute::is_void_v<ElementC_> && not cute::is_void_v<CopyOpG2R>;
   constexpr static bool is_destination_supported = not cute::is_void_v<ElementD> && not cute::is_void_v<CopyOpR2G>;
 
   constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();