Skip to content

Commit

Permalink
Add Epilogue Pipeline for PVC using EVT (#80)
Browse files Browse the repository at this point in the history
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]>
  • Loading branch information
muhammad-tanvir-1211 and aacostadiaz authored Jul 15, 2024
1 parent 8ff7158 commit 1cb2841
Show file tree
Hide file tree
Showing 10 changed files with 502 additions and 26 deletions.
34 changes: 20 additions & 14 deletions examples/sycl/pvc/pvc_bfloat_dpas_gemm_cute.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)) );
});
}

Expand Down Expand Up @@ -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>,
Expand Down
2 changes: 1 addition & 1 deletion include/cute/arch/copy_xe.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
4 changes: 4 additions & 0 deletions include/cutlass/epilogue/collective/collective_epilogue.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
/////////////////////////////////////////////////////////////////////////////////////////////////
Loading

0 comments on commit 1cb2841

Please sign in to comment.