Fixing trD compute type in the Xe Epilogue

Author: joyalbin

include/cutlass/epilogue/collective/xe_epilogue.hpp

@@ -389,7 +389,10 @@ class CollectiveEpilogue<
     cst_callbacks.begin();
 
     auto acc_frag = recast<Array<ElementCompute, FragmentSize>>(accumulators);
-    auto trD_compute_frag = recast<Array<ElementCompute, FragmentSize>>(trD_compute);
+    using FragmentVisit = decltype(cst_callbacks.visit(acc_frag(0), 0, 0, 0));
+    constexpr bool IsDirectR2S = cute::is_same_v<FragmentVisit, Array<ElementD, FragmentSize>>;
+    using RegisterElementD = cute::conditional_t<!IsDirectR2S, ElementCompute, ElementD>;
+    auto trD_compute_frag = recast<Array<RegisterElementD, FragmentSize>>(trD_compute);
 
     Tensor trD = make_tensor<ElementOutput>(Shape<Int<FragmentSize>>{});
     auto trD_frag = recast<Array<ElementOutput, FragmentSize>>(trD);
@@ -423,7 +426,7 @@ class CollectiveEpilogue<
         if constexpr (is_destination_supported) {
           CUTLASS_PRAGMA_UNROLL
           for (int i = 0; i < size(trD_compute_frag); ++i) {
-            trD_frag(i) = cutlass::NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize>{}(trD_compute_frag(i));
+            trD_frag(i) = cutlass::NumericArrayConverter<ElementOutput, RegisterElementD, FragmentSize>{}(trD_compute_frag(i));
           }
           copy(params.xe_store_d, trD, tCgD(_, epi_m, epi_n));
         }