Add Epilogue Pipeline for PVC using EVT (codeplaysoftware#80)

This PR introduces the Epilogue implementation for PVC using the Epilogue Visitor Tree available for SM90 (and onwards) GPUs for NVIDIA. We only support fusion::LinearCombination operation for PVC i.e. D = alpha * A * B + beta * C through this PR, but it can be extended further to add other fusion operations by partial specialization of the FusionCallBacks struct available in the include/cutlass/epilogue/fusion/intel_pvc_callbacks.hpp file. --------- Co-authored-by: Alejandro Acosta <[email protected]>
aacostadiaz · Jul 15, 2024 · 1cb2841 · 1cb2841
1 parent 8ff7158
commit 1cb2841
Show file tree

Hide file tree

Showing 10 changed files with 502 additions and 26 deletions.
diff --git a/examples/sycl/pvc/pvc_bfloat_dpas_gemm_cute.cpp b/examples/sycl/pvc/pvc_bfloat_dpas_gemm_cute.cpp
@@ -31,6 +31,8 @@
 
 #include "cutlass/gemm/device/gemm.h"
 #include "cutlass/epilogue/collective/default_epilogue.hpp"
+#include "cutlass/epilogue/collective/intel_pvc_epilogue.hpp"
+#include "cutlass/epilogue/fusion/intel_pvc_callbacks.hpp"
 #include "cutlass/gemm/device/gemm_universal.h"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
 #include "cutlass/gemm/collective/collective_mma.hpp"
@@ -49,7 +51,7 @@ template <typename T>
 static void fill_matrix(std::vector<T> &vector)
 {
   std::generate(std::begin(vector), std::end(vector), [&] {
-      return static_cast<T>( (rand() / double(RAND_MAX)) );
+    return static_cast<T>( (rand() / double(RAND_MAX)) );
   });
 }
 
@@ -360,26 +362,30 @@ int main(int argc, const char** argv)
           Layout<Shape<_1,_1,_1>>,
           Tile<_32,_64,_32>>;  // Subgroup level-tile
 
-  using DispatchPolicy = cutlass::gemm::MainloopIntelPVCUnpredicated;
+  using GEMMDispatchPolicy = cutlass::gemm::MainloopIntelPVCUnpredicated;
+  using EpilogueDispatchPolicy = cutlass::epilogue::IntelPVCEpilogue;
 
-  using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
-          ElementOutput,                                     // <- data type of output matrix
-          128 / cutlass::sizeof_bits<ElementOutput>::value,  // <- the number of elements per vectorized
-          // memory access. For a byte, it's 16
-          // elements. This becomes the vector width of
-          // math instructions in the epilogue too
-          ElementAccumulator,                                // <- data type of accumulator
-          ElementComputeEpilogue>;  // <- data type for alpha/beta in linear combination function
+  using EpilogueOp = cutlass::epilogue::fusion::LinearCombination<ElementOutput, ElementComputeEpilogue,
+          ElementAccumulator, ElementAccumulator, cutlass::FloatRoundStyle::round_to_nearest>;
 
-  using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
+  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>,
-          EpilogueOp,
-          cutlass::gemm::EpilogueDefault>;
+          FusionCallBacks,
+          XE_2D_U32x8x16x1x1_LD_N,
+          void, void,
+          XE_2D_U32x8x16x1x1_ST_N,
+          void, void>;
 
 // Mainloop
   using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
-          DispatchPolicy,
+          GEMMDispatchPolicy,
           TileShape,
           ElementInputA,
           cutlass::gemm::TagToStrideA_t<LayoutA>,

diff --git a/include/cute/arch/copy_xe.hpp b/include/cute/arch/copy_xe.hpp
@@ -287,7 +287,7 @@ struct XE_2D_U16x16x16x1x1_V
 struct XE_2D_U32x8x16x1x1_ST_N
 {
   template <class T>
-  CUTE_HOST_DEVICE static void copy(void *baseoffset, int width, int height,
+  CUTE_HOST_DEVICE static void copy(const void *baseoffset, int width, int height,
                                     int pitch, intel::coord_t coord, const T *src) {
     #if defined(SYCL_INTEL_TARGET)
       static_assert(sizeof(T) == 4, "Expected T to have size 4");

diff --git a/include/cutlass/epilogue/collective/collective_epilogue.hpp b/include/cutlass/epilogue/collective/collective_epilogue.hpp
@@ -56,7 +56,11 @@ class CollectiveEpilogue {
 #include "default_epilogue.hpp"
 #include "default_epilogue_array.hpp"
 #include "epilogue_tensor_broadcast.hpp"
+#if defined (SYCL_INTEL_TARGET)
+#include "intel_pvc_epilogue.hpp"
+#else
 #include "sm70_epilogue_vectorized.hpp"
 #include "sm90_epilogue_tma_warpspecialized.hpp"
 #include "sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp"
+#endif
 /////////////////////////////////////////////////////////////////////////////////////////////////