diff --git a/examples/00_bmg_gemm/00_bmg_gemm.cpp b/examples/00_bmg_gemm/00_bmg_gemm.cpp index 7e9291227e..0f7c665820 100644 --- a/examples/00_bmg_gemm/00_bmg_gemm.cpp +++ b/examples/00_bmg_gemm/00_bmg_gemm.cpp @@ -345,30 +345,33 @@ int main(int argc, const char** argv) using LayoutC = cutlass::layout::RowMajor; using LayoutD = cutlass::layout::RowMajor; - // The 2D block copy operations used for the A and B matrices - using GmemTiledCopyA = XE_2D_U16x32x32_LD_N; - using GmemTiledCopyB = XE_2D_U16x32x32_LD_V; + // [New Copy Atom] When left unspecified (void), make_block_2d_copy_* automatically selects + // appropriate 2D block copy operations for matrices A and B. Alternatively, you can + // explicitly specify new copy atom operations such as XE_LOAD_2D, XE_LOAD_2D_VNNI + // (applicable only to matrix B), or XE_LOAD_2D_TRANSPOSE. + // Refer https://github.com/intel/sycl-tla/blob/petercad/rearchitecture/media/docs/cpp/xe_rearchitecture.md + using GmemTiledCopyA = void; //XE_LOAD_2D<16, 32, 32>; + using GmemTiledCopyB = void; //XE_LOAD_2D_VNNI<16, 32, 32>; // Workgroup-level tile using TileShape = Shape<_256, _256, _32>; - + // A TiledMMA struct defines a tiling of an MMA atom over M, N and K, combining both additional // hardware (sub-groups for Intel BMG) and iterations by each sub-group. // - // The TiledMMAHelper struct defines a specific TiledMMA for a given MMA atom - // (XE_8x16x16_F32BF16BF16F32_TT), TileShape (<256, 256, 32>) and sub-group layout (8x4x1). The - // TiledMMA constructed using TiledMMAHelper has the property that each sub-group operates on a + // The TiledMMAHelper struct defines a specific TiledMMA for a given MMA atom. This example uses + // the XE_DPAS_TT<8, float, cute::bfloat16_t> atom, which represents an 8x16x16 DPAS operation with float32 accumulation and bfloat16 inputs, TileShape (<256, 256, 32>) and sub-group layout (8x4x1). + // The TiledMMA constructed using TiledMMAHelper has the property that each sub-group operates on a // single contiguous chunk of the work-group TileShape. For this configuration, this implies that // each sub-group operates on a contiguous 32x64x32 chunk (4x4x2 iterations). See // 0t_mma_atom.md#TiledMMAs for more info. Sub-groups are arranged row-major (stride 4,1,0) for // performance reasons. - using TiledMma = // M=8,N=16,K=16, D=f32,A=bf16,B=bf16,C=f32 - typename TiledMMAHelper, Layout, - Layout, Stride<_4, _1, _0>>>::TiledMMA; + using TiledMma = typename TiledMMAHelper>, Layout, Layout, Stride<_4, _1, _0>>>::TiledMMA; // For Intel BMG, PipelineStages defines how many k-blocks ahead to prefetch from A and B. constexpr int PipelineStages = 2; - using GEMMDispatchPolicy = cutlass::gemm::MainloopIntelXeXMX16; + // For older version of copy/mma atom, use cutlass::gemm::MainloopIntelXeXMX16 as dispatch policy + using GEMMDispatchPolicy = cutlass::gemm::MainloopXeL1Staged; using EpilogueDispatchPolicy = cutlass::epilogue::IntelXeXMX16; // This is the 'default' epilogue operation (Linear Combination) which performs everything in: diff --git a/examples/00_bmg_gemm/legacy/00_bmg_gemm.cpp b/examples/00_bmg_gemm/legacy/00_bmg_gemm.cpp new file mode 100644 index 0000000000..91139cf0d6 --- /dev/null +++ b/examples/00_bmg_gemm/legacy/00_bmg_gemm.cpp @@ -0,0 +1,429 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2024 Codeplay Software Ltd. All rights reserved. + * Copyright (C) 2025 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief CUTLASS Intel BMG Gemm Example. + + This example constructs and executes a simple CUTLASS GEMM kernel on Intel BMG hardware, and + verifies its correctness with a reference implementation + (cutlass::reference::device::GemmComplex). The example also provides a performance measurement + for the GEMM in TFLOPS. + + This example makes use of BMGs subgroup cooperative 2d-block copy operations and DPAS instructions. + + The shapes of the A and B matrix are defined at runtime by `options.m`, `.n` and `.k`, and the + batch size is defined by `options.l`. The tile shape, which defines how much work is executed by + a single work-group, is defined at compile time by: + ``` + using TileShape = Shape<_256, _256, _32>; + ``` + That is, each work-group processes a tile of M=256, N=256, and iterates over `options.k` in + blocks of K=32. + + Performance of GEMM on BMG is heavily dependent on prefetching the A and B matrices. That is, + executing Intel specific prefetch instructions for future iterations to ensure that the required + blocks of A and B are resident in cache before they are needed. + + To build & run this example (from your build dir): + + $ ninja 00_bmg_gemm + $ ./examples/sycl/00_bmg_gemm/00_bmg_gemm + + Call with `--help` for information about available options +*/ + +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/collective/xe_epilogue.hpp" +#include "cutlass/epilogue/fusion/xe_callbacks.hpp" +#include "cutlass/gemm/device/gemm_universal.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/collective/collective_mma.hpp" +#include "cutlass/util/GPU_Clock.hpp" + +#include +#include + +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/device/tensor_compare.h" +#include "sycl_common.hpp" +#include "helper.h" + +using namespace cute; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// Command line options parsing +struct Options { + + bool help; + bool error; + + int m, n, k, l, iterations; + float alpha, beta; + + Options(): + help(false), + error(false), + m(5120), n(4096), k(4096), l(1), iterations(20), + alpha(1.f), beta(0.f) + { } + + // Parses the command line + void parse(int argc, char const **args) { + cutlass::CommandLine cmd(argc, args); + + if (cmd.check_cmd_line_flag("help")) { + help = true; + return; + } + + cmd.get_cmd_line_argument("m", m, 5120); + cmd.get_cmd_line_argument("n", n, 4096); + cmd.get_cmd_line_argument("k", k, 4096); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("alpha", alpha, 1.f); + cmd.get_cmd_line_argument("beta", beta, 0.f); + cmd.get_cmd_line_argument("iterations", iterations, 100); + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "BMG GEMM Example\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement\n\n" + << " --m= Sets the M extent of the GEMM\n" + << " --n= Sets the N extent of the GEMM\n" + << " --k= Sets the K extent of the GEMM\n" + << " --l= Sets the L extent (batch count) of the GEMM\n" + << " --alpha= Epilogue scalar alpha\n" + << " --beta= Epilogue scalar beta\n\n" + << " --iterations= Iterations\n\n"; + + return out; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class Gemm +> +struct ExampleRunner { + + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + using LayoutA = typename Gemm::LayoutA; + using LayoutB = typename Gemm::LayoutB; + using LayoutC = typename Gemm::LayoutC; + using LayoutD = typename Gemm::LayoutD; + + using ElementA = typename Gemm::ElementA; + using ElementB = typename Gemm::ElementB; + using ElementAcc = typename Gemm::ElementAccumulator; + + using CollectiveEpilogue = typename Gemm::CollectiveEpilogue; + using ElementC = typename Gemm::ElementC; + using ElementOutput = typename CollectiveEpilogue::ElementOutput; + using ElementCompute = typename CollectiveEpilogue::ElementCompute; + using ElementAccumulator = typename CollectiveEpilogue::ElementAccumulator; + + using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape; + + // + // Data members + // + + /// Initialization + StrideA stride_A; + StrideB stride_B; + StrideC stride_C; + StrideD stride_D; + uint64_t seed = 0; + + cutlass::DeviceAllocation block_A; + cutlass::DeviceAllocation block_B; + cutlass::DeviceAllocation block_C; + cutlass::DeviceAllocation block_D; + cutlass::DeviceAllocation block_ref_D; // Reference GEMM result for verification + + // + // Methods + // + + bool verify(const ProblemShapeType& problem_size, ElementCompute alpha, ElementCompute beta) { + auto [M, N, K, L] = problem_size; + + cutlass::TensorRef ref_A(block_A.get(), LayoutA::packed({M, K})); + cutlass::TensorRef ref_B(block_B.get(), LayoutB::packed({K, N})); + cutlass::TensorRef ref_C(block_C.get(), LayoutC::packed({M, N})); + cutlass::TensorRef ref_D(block_ref_D.get(), LayoutD::packed({M, N})); + + cutlass::reference::device::GemmComplex( + {M, N, K}, + alpha, + ref_A, + cutlass::ComplexTransform::kNone, + ref_B, + cutlass::ComplexTransform::kNone, + beta, + ref_C, + ref_D, + ElementAccumulator(0), + L, // batch_count + M * K, // batch_stride_A + K * N, // batch_stride_B + M * N, // batch_stride_C + M * N // batch_stride_D + ); + + // CUTLASS on SYCL uses the compatibility library compat for e.g. default in-order queue + compat::wait(); + + // Check if output from CUTLASS kernel and reference kernel are equal or not + bool passed = cutlass::reference::device::BlockCompareEqual( + block_ref_D.get(), block_D.get(), block_D.size()); + + return passed; + } + + /// Initialize operands to be used in the GEMM and reference GEMM + void initialize(const ProblemShapeType& problem_size) { + auto problem_shape_MNKL = cute::append<4>(problem_size, 1); + auto [M, N, K, L] = problem_shape_MNKL; + + // Complete the stride by combining static layout info (StrideA) with runtime size info (M,K,L) + stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, L)); + stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, L)); + stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L)); + stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L)); + + block_A.reset(static_cast(M) * K * L); + block_B.reset(static_cast(K) * N * L); + block_C.reset(static_cast(M) * N * L); + block_D.reset(static_cast(M) * N * L); + block_ref_D.reset(static_cast(M) * N * L); + + initialize_block(block_A, seed + 2023); + initialize_block(block_B, seed + 2022); + initialize_block(block_C, seed + 2021); + } + + cutlass::Status run(const Options& options, const cutlass::KernelHardwareInfo& hw_info) { + ProblemShapeType problem_size = ProblemShapeType{options.m, options.n, options.k, options.l}; + + initialize(problem_size); + + typename Gemm::GemmKernel::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + problem_size, + {block_A.get(), stride_A, block_B.get(), stride_B}, + {{options.alpha, options.beta}, block_C.get(), stride_C, block_D.get(), stride_D}, + hw_info + }; + + Gemm gemm_op; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + if (gemm_op.can_implement(arguments) != cutlass::Status::kSuccess){ + std::cout << "Invalid Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; + std::exit(1); + } + + CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get())); + + // Run the GEMM + CUTLASS_CHECK(gemm_op.run()); + + compat::wait(); + + // Verify that the result is correct + bool passed = verify(problem_size, options.alpha, options.beta); + std::cout << "Disposition: " << (passed ? "Passed" : "Failed") << std::endl; + + if(!passed) return cutlass::Status::kErrorInternal; + + if (options.iterations > 0) { + GPU_Clock timer; + timer.start(); + for (int i = 0; i < options.iterations; ++i) { + gemm_op.run(); + } + compat::wait(); + + float cute_time = timer.seconds() / options.iterations; + double tflops = (2.0 * options.m * options.n * options.k * options.l) * 1e-12; + std::cout << "Problem Size: " << options.m << 'x' << options.n << 'x' << options.k << 'x' << options.l << std::endl; + printf("Cutlass GEMM Performance: [%4.3f]TFlop/s (%6.4f)ms\n", tflops / cute_time, cute_time*1000); + } + + return cutlass::Status::kSuccess; + } + +}; + +int main(int argc, const char** argv) +{ + // + // Parse options + // + + Options options; + + options.parse(argc, argv); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.error) { + std::cerr << "Aborting execution." << std::endl; + return -1; + } + + // + // Run examples + // + + // The KernelHardwareInfo struct holds the number of EUs on the GPU with a given device ID. This + // information is used by the underlying kernel. + cutlass::KernelHardwareInfo hw_info; + + // Change device_id to another value if you are running on a machine with multiple GPUs and wish + // to use a GPU other than that with device ID 0. + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + bool passed; + + // The code section below describes datatype for input, output matrices and computation between + // elements in input matrices. + using ElementAccumulator = float; // <- data type of accumulator + using ElementComputeEpilogue = float; // <- data type of epilogue operations + using ElementInputA = bfloat16_t; // <- data type of elements in input matrix A + using ElementInputB = bfloat16_t; // <- data type of elements in input matrix B + using ElementOutput = float; // <- data type of elements in output matrix D + + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + using LayoutC = cutlass::layout::RowMajor; + using LayoutD = cutlass::layout::RowMajor; + + // The 2D block copy operations used for the A and B matrices + using GmemTiledCopyA = XE_2D_U16x32x32_LD_N; + using GmemTiledCopyB = XE_2D_U16x32x32_LD_V; + + // Workgroup-level tile + using TileShape = Shape<_256, _256, _32>; + + // A TiledMMA struct defines a tiling of an MMA atom over M, N and K, combining both additional + // hardware (sub-groups for Intel BMG) and iterations by each sub-group. + // + // The TiledMMAHelper struct defines a specific TiledMMA for a given MMA atom + // (XE_8x16x16_F32BF16BF16F32_TT), TileShape (<256, 256, 32>) and sub-group layout (8x4x1). The + // TiledMMA constructed using TiledMMAHelper has the property that each sub-group operates on a + // single contiguous chunk of the work-group TileShape. For this configuration, this implies that + // each sub-group operates on a contiguous 32x64x32 chunk (4x4x2 iterations). See + // 0t_mma_atom.md#TiledMMAs for more info. Sub-groups are arranged row-major (stride 4,1,0) for + // performance reasons. + using TiledMma = // M=8,N=16,K=16, D=f32,A=bf16,B=bf16,C=f32 + typename TiledMMAHelper, Layout, + Layout, Stride<_4, _1, _0>>>::TiledMMA; + + // For Intel BMG, PipelineStages defines how many k-blocks ahead to prefetch from A and B. + constexpr int PipelineStages = 2; + using GEMMDispatchPolicy = cutlass::gemm::MainloopIntelXeXMX16; + using EpilogueDispatchPolicy = cutlass::epilogue::IntelXeXMX16; + + // This is the 'default' epilogue operation (Linear Combination) which performs everything in: + // (D = alpha * (A*B) + beta * C) + // aside from the (A*B), which is handled by the GEMM. See 05_bmg_gemm_with_epilogues for more + // complex epilogue examples. + using EpilogueOp = cutlass::epilogue::fusion::LinearCombination; + + // FusionCallbacks ties the EpilogueOp to an implementation (based on the dispatch + // policy/architecture) and defines the epilogue arguments. + using FusionCallBacks = cutlass::epilogue::fusion::FusionCallbacks; + // GEMM Epilogue - loads & stores C/D matrices, performs epilogue operations & load/stores any + // auxiliary data required + using CollectiveEpilogue = cutlass::epilogue::collective::CollectiveEpilogue< + EpilogueDispatchPolicy, + TileShape, + ElementAccumulator, + cutlass::gemm::TagToStrideC_t, // Converts CUTLASS 2.x to CUTLASS 3.x representation + ElementOutput, + cutlass::gemm::TagToStrideC_t, // Converts CUTLASS 2.x to CUTLASS 3.x representation + FusionCallBacks, + XE_2D_U32x8x16_LD_N, // The copy atom used to load matrix C + void, void, + XE_2D_U32x8x16_ST_N, // The copy atom used to store matrix D + void, void>; + + // GEMM Mainloop - iteration over blocks in K dimension + using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma< + GEMMDispatchPolicy, + TileShape, + ElementInputA, + cutlass::gemm::TagToStrideA_t, // Converts CUTLASS 2.x to CUTLASS 3.x representation + ElementInputB, + cutlass::gemm::TagToStrideB_t, // Converts CUTLASS 2.x to CUTLASS 3.x representation + TiledMma, + GmemTiledCopyA, void, void, cute::identity, // A + GmemTiledCopyB, void, void, cute::identity // B + >; + + // Define the whole kernel (mainloop and epilogue) + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, // Defer global problem shape definition to runtime + CollectiveMainloop, + CollectiveEpilogue + >; + + // The GemmUniversalAdapter wraps the defined GEMM kernel and handles the launch, and e.g. + // persistent scratch memory if required. + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + ExampleRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + + return 0; +} \ No newline at end of file diff --git a/examples/00_bmg_gemm/legacy/CMakeLists.txt b/examples/00_bmg_gemm/legacy/CMakeLists.txt new file mode 100644 index 0000000000..1d40199221 --- /dev/null +++ b/examples/00_bmg_gemm/legacy/CMakeLists.txt @@ -0,0 +1,41 @@ +# Copyright (c) 2024 - 2025 Codeplay Software Ltd. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +set(TEST_BATCHES --l=2) +set(TEST_LARGE "--l=513 --m=8 --n=16384 --k=512") # B matrix capacity > uint32_max +set(TEST_SMALL_SHAPE --m=4 --n=8 --k=8 --l=2) + + +cutlass_example_add_executable( + 00_bmg_gemm_legacy + 00_bmg_gemm.cpp + TEST_COMMAND_OPTIONS + TEST_BATCHES + TEST_LARGE + TEST_SMALL_SHAPE +) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d141f5b7de..38c6b34b75 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -101,6 +101,7 @@ if(CUTLASS_ENABLE_SYCL) message(STATUS "Building examples for Intel GPU targets") foreach(EXAMPLE 00_bmg_gemm + 00_bmg_gemm/legacy 01_bmg_gemm_with_collective_builder 02_bmg_gemm_mixed_dtype 03_bmg_gemm_streamk diff --git a/include/cute/atom/copy_traits_xe_2d.hpp b/include/cute/atom/copy_traits_xe_2d.hpp index 2df8ae0a38..e6cbe664c8 100644 --- a/include/cute/atom/copy_traits_xe_2d.hpp +++ b/include/cute/atom/copy_traits_xe_2d.hpp @@ -756,6 +756,31 @@ make_block_2d_copy_X(CopyOp const& op, // Copy operation return make_block_2d_copy(op, gstride, x_mode, y_mode, atom_shape, sv_layout_t); } +// Helper trait to detect new XE copy ops +template +struct is_new_xe_atom : cute::false_type {}; + +// Helper trait specifically for XE_LOAD_2D_VNNI (for copy B) +template +struct is_new_xe_atom_vnni : cute::false_type {}; + +// Helper trait specifically for XE_STORE_2D (for copy C) +template +struct is_new_xe_atom_store : cute::false_type {}; + +// Check if T is an instantiation of XE_LOAD_2D +template +struct is_new_xe_atom> : cute::true_type {}; + +// Check if T is an instantiation of XE_LOAD_2D_TRANSPOSE +template +struct is_new_xe_atom> : cute::true_type {}; + +// Check if T is an instantiation of XE_LOAD_2D_VNNI +template +struct is_new_xe_atom_vnni> : cute::true_type {}; + + // MMA-focused TiledCopy creation functions. template CUTE_HOST_DEVICE @@ -774,6 +799,12 @@ make_block_2d_copy_A(CopyOp const& op, // Copy operation TiledMMA const& mma, // TiledMMA instance Tensor const& gmem) // Global tensor { + // This will pass for new atoms like XE_LOAD_2D<16, 32, 32> + // and fail for old atoms like XE_2D_U16x32x32_LD_N + static_assert(is_new_xe_atom::value, + "Legacy XE copy atom ops not compatible with make_block_2d_copy_A. " + "Please use the new templated atoms: XE_LOAD_2D or XE_LOAD_2D_TRANSPOSE. " + "Examples: XE_2D_U16x32x32_LD_N -> XE_LOAD_2D<16, 32, 32>, XE_2D_U16x32x32_LD_V -> XE_LOAD_2D_TRANSPOSE<16, 32, 32>"); using ValType = typename GEngine::value_type; return make_block_2d_copy_A(op, mma, gmem.stride()).with(gmem); } @@ -846,6 +877,11 @@ make_block_2d_copy_B(CopyOp const& op, // Copy operation TiledMMA const& mma, // TiledMMA instance Tensor const& gmem) // Global tensor { + // Only accept XE_LOAD_2D_VNNI for copy B + static_assert(is_new_xe_atom_vnni::value, + "Legacy XE copy atom ops not compatible with make_block_2d_copy_B. " + "Please use the new templated atom: XE_LOAD_2D_VNNI. " + "Examples: XE_2D_U16x32x32_LD_V -> XE_LOAD_2D_VNNI<16, 32, 32, 32>"); using ValType = typename GEngine::value_type; return make_block_2d_copy_B(op, mma, gmem.stride()).with(gmem); } diff --git a/include/cutlass/gemm/collective/collective_mma.hpp b/include/cutlass/gemm/collective/collective_mma.hpp index b54776ece5..bb5e2017d8 100644 --- a/include/cutlass/gemm/collective/collective_mma.hpp +++ b/include/cutlass/gemm/collective/collective_mma.hpp @@ -78,6 +78,7 @@ #if defined(SYCL_INTEL_TARGET) #include "cutlass/gemm/collective/xe_mma.hpp" +#include "cutlass/gemm/collective/xe_mma_legacy.hpp" #include "cutlass/gemm/collective/xe_array_mma.hpp" #include "cutlass/gemm/collective/xe_array_mma_fp8.hpp" #include "cutlass/gemm/collective/xe_mma_mixed_input.hpp" diff --git a/include/cutlass/gemm/collective/xe_mma.hpp b/include/cutlass/gemm/collective/xe_mma.hpp index 7d75825a3e..bbd6add94e 100644 --- a/include/cutlass/gemm/collective/xe_mma.hpp +++ b/include/cutlass/gemm/collective/xe_mma.hpp @@ -47,13 +47,13 @@ using namespace cute; template -struct CollectiveMma, TileShape_, ElementA_, StrideA_, ElementB_, StrideB_, TiledMma_, +struct CollectiveMma, TileShape_, ElementA_, StrideA_, ElementB_, StrideB_, TiledMma_, GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_, TransformA_, GmemTiledCopyB_, SmemLayoutAtomB_, SmemCopyAtomB_, TransformB_> { // // Type Aliases // - using DispatchPolicy = MainloopIntelXeXMX16; + using DispatchPolicy = MainloopXeL1Staged; using WorkgroupTileShape = TileShape_; using ElementA = ElementA_; using StrideA = StrideA_; @@ -71,7 +71,7 @@ struct CollectiveMma, TileShape_, Element using TransformB = TransformB_; using ArchTag = typename DispatchPolicy::ArchTag; - static_assert(platform::is_same::value, "MainloopIntelXeXMX16 requires that A and B have same type."); + static_assert(platform::is_same::value, "MainloopXeL1Staged requires that A and B have same type."); static_assert(std::is_same_v, "Transformation for A is not currently supported on Intel PVC"); static_assert(std::is_same_v, "Transformation for B is not currently supported on Intel PVC"); @@ -100,8 +100,36 @@ struct CollectiveMma, TileShape_, Element static constexpr auto Num_SGs = ATOM_N * ATOM_M * ATOM_K; static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}); - using Copy_A = typename Copy_Traits::template DefaultTiledCopy; - using Copy_B = typename Copy_Traits::template DefaultTiledCopy; + // Helper struct to deduce CopyOpA type + template + struct CopyOpAHelper { + static auto get() { + auto tmp = make_tensor(make_gmem_ptr(static_cast(nullptr)), + make_layout(make_shape(Int{}, Int{}, Int<1>{}), SA{})); + if constexpr (!std::is_void_v) { + return make_block_2d_copy_A(GTCA{}, TM{}, tmp(_,_,0)); + } else { + return make_block_2d_copy_A(TM{}, tmp(_,_,0)); + } + } + }; + + // Helper struct to deduce CopyOpB type + template + struct CopyOpBHelper { + static auto get() { + auto tmp = make_tensor(make_gmem_ptr(static_cast(nullptr)), + make_layout(make_shape(Int{}, Int{}, Int<1>{}), SB{})); + if constexpr (!std::is_void_v) { + return make_block_2d_copy_B(GTCB{}, TM{}, tmp(_,_,0)); + } else { + return make_block_2d_copy_B(TM{}, tmp(_,_,0)); + } + } + }; + + using CopyOpA = decltype(CopyOpAHelper::get()); + using CopyOpB = decltype(CopyOpBHelper::get()); // Host side kernel arguments struct Arguments { @@ -112,8 +140,12 @@ struct CollectiveMma, TileShape_, Element }; struct Params { - Copy_A tiled_copy_a; - Copy_B tiled_copy_b; + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + CopyOpA copy_a; + CopyOpB copy_b; }; // @@ -129,12 +161,32 @@ struct CollectiveMma, TileShape_, Element auto [M,N,K,L] = problem_shape; - auto mA_mkl = make_tensor(make_gmem_ptr(args.ptr_A), make_layout(make_shape(M, K, L), args.dA)); - auto mB_nkl = make_tensor(make_gmem_ptr(args.ptr_B), make_layout(make_shape(N, K, L), args.dB)); - Copy_A tiled_copy_a{Copy_A{}.with(mA_mkl)}; - Copy_B tiled_copy_b{Copy_B{}.with(mB_nkl)}; - - return Params{tiled_copy_a, tiled_copy_b}; + auto mA_mkl = make_tensor(make_gmem_ptr(args.ptr_A), + make_layout(make_shape(M, K, L), args.dA)); + auto mB_nkl = make_tensor(make_gmem_ptr(args.ptr_B), + make_layout(make_shape(N, K, L), args.dB)); + + CopyOpA copy_a = [&]() { + if constexpr (!std::is_void_v) { + return make_block_2d_copy_A(GmemTiledCopyA{}, TiledMma{}, mA_mkl(_,_,0)); + } else { + return make_block_2d_copy_A(TiledMma{}, mA_mkl(_,_,0)); + } + }(); + + CopyOpB copy_b = [&]() { + if constexpr (!std::is_void_v) { + return make_block_2d_copy_B(GmemTiledCopyB{}, TiledMma{}, mB_nkl(_,_,0)); + } else { + return make_block_2d_copy_B(TiledMma{}, mB_nkl(_,_,0)); + } + }(); + + return Params{args.ptr_A, + args.dA, + args.ptr_B, + args.dB, + copy_a, copy_b}; } template @@ -177,38 +229,33 @@ struct CollectiveMma, TileShape_, Element static_assert(is_rmem::value, "D tensor must be rmem resident."); static_assert(is_rmem::value, "C tensor must be rmem resident."); - auto thr_copy_A = mainloop.tiled_copy_a.get_slice(thread_idx); - auto thr_copy_B = mainloop.tiled_copy_b.get_slice(thread_idx); + auto thr_copy_a = mainloop.copy_a.get_slice(thread_idx); + auto thr_copy_b = mainloop.copy_b.get_slice(thread_idx); // Instantiate the MMA object and get thread slice TiledMma tiled_mma; - // TODO(Codeplay): see if we can make this nicer - // To make all work items in a subgroup have the same global tensors pass in the index of work item 0 in each subgroup - auto sg = compat::get_nd_item<1>().get_sub_group(); - auto first_thread_in_sg_idx = sg.get_group_linear_id() * DispatchPolicy::SubgroupSize; - auto thr_mma = tiled_mma.get_slice(first_thread_in_sg_idx); - - // Partition global counting tensors for MMA - Tensor tCgA = thr_mma.partition_A(gA); - Tensor tCgB = thr_mma.partition_B(gB); - - Tensor tCrA = make_tensor(make_fragment_layout(mainloop.tiled_copy_a, tCgA(_,_,_,0).shape())); - Tensor tCrB = make_tensor(make_fragment_layout(mainloop.tiled_copy_b, tCgB(_,_,_,0).shape())); - - // Retile registers for copies - Tensor tArA = thr_copy_A.retile_D(tCrA); - Tensor tBrB = thr_copy_B.retile_D(tCrB); - - // Retile global counting tensors for copies - Tensor tAgA = thr_copy_A.retile_S(tCgA); - Tensor tBgB = thr_copy_B.retile_S(tCgB); - - auto tiled_prefetch_a = cute::prefetch_selector,Int>, Num_SGs>(mainloop.tiled_copy_a); - auto tiled_prefetch_b = cute::prefetch_selector,Int>, Num_SGs>(mainloop.tiled_copy_b); - auto thr_prefetch_A = tiled_prefetch_a.get_slice(thread_idx); - auto thr_prefetch_B = tiled_prefetch_b.get_slice(thread_idx); + auto thr_mma = tiled_mma.get_slice(thread_idx); + + /* Register fragments for MMA */ + auto tCrA = thr_mma.partition_sg_fragment_A(gA(_,_,0)); + auto tCrB = thr_mma.partition_sg_fragment_B(gB(_,_,0)); + + /* Register fragments for copies */ + auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_,_,0)); + auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_,_,0)); + + /* Partition global tensor (proxies) for copies */ + Tensor tAgA = thr_copy_a.partition_S(gA); + Tensor tBgB = thr_copy_b.partition_S(gB); - // Partition global tile for prefetch + /* Create prefetch TiledCopy instances */ + auto prefetch_a = make_block_2d_prefetch(mainloop.copy_a); + auto prefetch_b = make_block_2d_prefetch(mainloop.copy_b); + + auto thr_prefetch_A = prefetch_a.get_slice(thread_idx); + auto thr_prefetch_B = prefetch_b.get_slice(thread_idx); + + /* Partition global tensor (proxies) for prefetch */ auto pAgA = thr_prefetch_A.partition_S(gA); auto pBgB = thr_prefetch_B.partition_S(gB); @@ -216,20 +263,18 @@ struct CollectiveMma, TileShape_, Element #define PRINT(x) print(#x ": "); print(x); print("\n"); if (cute::thread(LOG_THREAD, LOG_GROUP)) { print("======================= A: \n"); - PRINT(tCgA); PRINT(tAgA); PRINT(tCrA); PRINT(tArA); - PRINT(mainloop.tiled_copy_a); + PRINT(mainloop.copy_a); print("======================= B: \n"); - PRINT(tCgB); PRINT(tBgB); PRINT(tCrB); PRINT(tBrB); - PRINT(mainloop.tiled_copy_b); + PRINT(mainloop.copy_b); } #undef PRINT #endif @@ -243,21 +288,25 @@ struct CollectiveMma, TileShape_, Element CUTLASS_PRAGMA_UNROLL for (; prefetch_k < DispatchPolicy::Stages; prefetch_k++) { - prefetch(tiled_prefetch_a, pAgA(_, _, _, prefetch_k)); - prefetch(tiled_prefetch_b, pBgB(_, _, _, prefetch_k)); + prefetch(prefetch_a, pAgA(_, _, _, prefetch_k)); + prefetch(prefetch_b, pBgB(_, _, _, prefetch_k)); } for (int k_tile = k_start_idx; k_tile < k_tile_count + k_start_idx; k_tile++, prefetch_k++) { barrier_arrive(barrier_scope); // Copy gmem to rmem for the first k_tile - copy(mainloop.tiled_copy_a, tAgA(_,_,_,k_tile), tArA); - copy(mainloop.tiled_copy_b, tBgB(_,_,_,k_tile), tBrB); + copy(mainloop.copy_a, tAgA(_,_,_,k_tile), tArA); + copy(mainloop.copy_b, tBgB(_,_,_,k_tile), tBrB); if (prefetch_k < k_tile_count) { - prefetch(tiled_prefetch_a, pAgA(_, _, _, prefetch_k)); - prefetch(tiled_prefetch_b, pBgB(_, _, _, prefetch_k)); + prefetch(prefetch_a, pAgA(_, _, _, prefetch_k)); + prefetch(prefetch_b, pBgB(_, _, _, prefetch_k)); } + /* Shuffle data from copy fragments to MMA fragments */ + reorder(tArA, tCrA); + reorder(tBrB, tCrB); + cute::gemm(tiled_mma, tCrA, tCrB, accum); barrier_wait(barrier_scope); } @@ -267,3 +316,4 @@ struct CollectiveMma, TileShape_, Element } // namespace cutlass::gemm::collective ///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/include/cutlass/gemm/collective/xe_mma_legacy.hpp b/include/cutlass/gemm/collective/xe_mma_legacy.hpp new file mode 100644 index 0000000000..ac7fe83884 --- /dev/null +++ b/include/cutlass/gemm/collective/xe_mma_legacy.hpp @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2024 Codeplay Software Ltd. All rights reserved. + * Copyright (C) 2025 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct CollectiveMma, TileShape_, ElementA_, StrideA_, ElementB_, StrideB_, TiledMma_, + GmemTiledCopyA_, SmemLayoutAtomA_, SmemCopyAtomA_, TransformA_, GmemTiledCopyB_, SmemLayoutAtomB_, + SmemCopyAtomB_, TransformB_> { + // + // Type Aliases + // + using DispatchPolicy = MainloopIntelXeXMX16; + using WorkgroupTileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(platform::is_same::value, "MainloopIntelXeXMX16 requires that A and B have same type."); + static_assert(std::is_same_v, "Transformation for A is not currently supported on Intel PVC"); + static_assert(std::is_same_v, "Transformation for B is not currently supported on Intel PVC"); + + static constexpr int SubgroupSize = DispatchPolicy::SubgroupSize; + + using MmaAtomShape = typename TiledMma::AtomShape_MNK; + + static constexpr int BLK_M = get<0>(WorkgroupTileShape{}); + static constexpr int BLK_N = get<1>(WorkgroupTileShape{}); + static constexpr int BLK_K = get<2>(WorkgroupTileShape{}); + + static constexpr int ATOM_M = get<1>(typename TiledMma::ThrLayoutVMNK{}.shape()); + static constexpr int ATOM_N = get<2>(typename TiledMma::ThrLayoutVMNK{}.shape()); + static constexpr int ATOM_K = get<3>(typename TiledMma::ThrLayoutVMNK{}.shape()); + + static_assert(BLK_M % TiledMma{}.template tile_size_mnk<0>() == 0, "TiledMma permutation size must match block size."); + static_assert(BLK_N % TiledMma{}.template tile_size_mnk<1>() == 0, "TiledMma permutation size must match block size."); + static_assert(BLK_K % TiledMma{}.template tile_size_mnk<2>() == 0, "TiledMma permutation size must match block size."); + + static constexpr int SG_M = ceil_div(BLK_M, ATOM_M); + static constexpr int SG_N = ceil_div(BLK_N, ATOM_N); + static constexpr int SG_K = ceil_div(BLK_K, ATOM_K); + using SubgroupTileShape = Shape, C, C>; + + // 32 + static constexpr auto Num_SGs = ATOM_N * ATOM_M * ATOM_K; + static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}); + + using Copy_A = typename Copy_Traits::template DefaultTiledCopy; + using Copy_B = typename Copy_Traits::template DefaultTiledCopy; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + struct Params { + Copy_A tiled_copy_a; + Copy_B tiled_copy_b; + }; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + (void) workspace; + + auto [M,N,K,L] = problem_shape; + + auto mA_mkl = make_tensor(make_gmem_ptr(args.ptr_A), make_layout(make_shape(M, K, L), args.dA)); + auto mB_nkl = make_tensor(make_gmem_ptr(args.ptr_B), make_layout(make_shape(N, K, L), args.dB)); + Copy_A tiled_copy_a{Copy_A{}.with(mA_mkl)}; + Copy_B tiled_copy_b{Copy_B{}.with(mB_nkl)}; + + return Params{tiled_copy_a, tiled_copy_b}; + } + + template + static bool + can_implement( + ProblemShape problem_shapes, + Arguments const& args) { + constexpr int copy_alignment_bits = 128; + constexpr int batch_alignment_bits = 512; + auto problem_shape_MNKL = append<4>(problem_shapes, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + bool implementable = true; + + constexpr int min_aligned_elements_A = copy_alignment_bits / sizeof_bits::value; + implementable &= cutlass::detail::check_alignment(cute::make_shape(M,K,L), args.dA); + constexpr int min_aligned_elements_B = copy_alignment_bits / sizeof_bits::value; + implementable &= cutlass::detail::check_alignment(cute::make_shape(N,K,L), args.dB); + + if (L > 1) { + constexpr int min_batch_aligned_elements_A = batch_alignment_bits / sizeof_bits::value; + implementable &= get<2>(args.dA) % min_batch_aligned_elements_A == 0; + constexpr int min_batch_aligned_elements_B = batch_alignment_bits / sizeof_bits::value; + implementable &= get<2>(args.dB) % min_batch_aligned_elements_B == 0; + } + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for XE 2D copy.\n"); + } + + return implementable; + } + + /// Perform a subgroup-scoped matrix multiply-accumulate + template + CUTLASS_DEVICE void operator()(FrgTensorD &accum, TensorA gA, TensorB gB, FrgTensorC const &src_accum, + KTileIterator k_tile_iter, int k_tile_count, BlkCoord const &blk_coord, int const &K_start, int thread_idx, + Params const &mainloop) { + (void)blk_coord; + static_assert(is_rmem::value, "D tensor must be rmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + + auto thr_copy_A = mainloop.tiled_copy_a.get_slice(thread_idx); + auto thr_copy_B = mainloop.tiled_copy_b.get_slice(thread_idx); + + // Instantiate the MMA object and get thread slice + TiledMma tiled_mma; + // TODO(Codeplay): see if we can make this nicer + // To make all work items in a subgroup have the same global tensors pass in the index of work item 0 in each subgroup + auto sg = compat::get_nd_item<1>().get_sub_group(); + auto first_thread_in_sg_idx = sg.get_group_linear_id() * DispatchPolicy::SubgroupSize; + auto thr_mma = tiled_mma.get_slice(first_thread_in_sg_idx); + + // Partition global counting tensors for MMA + Tensor tCgA = thr_mma.partition_A(gA); + Tensor tCgB = thr_mma.partition_B(gB); + + Tensor tCrA = make_tensor(make_fragment_layout(mainloop.tiled_copy_a, tCgA(_,_,_,0).shape())); + Tensor tCrB = make_tensor(make_fragment_layout(mainloop.tiled_copy_b, tCgB(_,_,_,0).shape())); + + // Retile registers for copies + Tensor tArA = thr_copy_A.retile_D(tCrA); + Tensor tBrB = thr_copy_B.retile_D(tCrB); + + // Retile global counting tensors for copies + Tensor tAgA = thr_copy_A.retile_S(tCgA); + Tensor tBgB = thr_copy_B.retile_S(tCgB); + + auto tiled_prefetch_a = cute::prefetch_selector,Int>, Num_SGs>(mainloop.tiled_copy_a); + auto tiled_prefetch_b = cute::prefetch_selector,Int>, Num_SGs>(mainloop.tiled_copy_b); + auto thr_prefetch_A = tiled_prefetch_a.get_slice(thread_idx); + auto thr_prefetch_B = tiled_prefetch_b.get_slice(thread_idx); + + // Partition global tile for prefetch + auto pAgA = thr_prefetch_A.partition_S(gA); + auto pBgB = thr_prefetch_B.partition_S(gB); + +#if CUTLASS_ENABLE_DEBUG_PRINTS +#define PRINT(x) print(#x ": "); print(x); print("\n"); + if (cute::thread(LOG_THREAD, LOG_GROUP)) { + print("======================= A: \n"); + PRINT(tCgA); + PRINT(tAgA); + + PRINT(tCrA); + PRINT(tArA); + PRINT(mainloop.tiled_copy_a); + + print("======================= B: \n"); + PRINT(tCgB); + PRINT(tBgB); + + PRINT(tCrB); + PRINT(tBrB); + PRINT(mainloop.tiled_copy_b); + } +#undef PRINT +#endif + + // + // Mainloop + // + const auto k_start_idx = crd2idx((*k_tile_iter), make_shape(K_start)); + constexpr int barrier_scope = 2; + int prefetch_k = k_start_idx; + + CUTLASS_PRAGMA_UNROLL + for (; prefetch_k < DispatchPolicy::Stages; prefetch_k++) { + prefetch(tiled_prefetch_a, pAgA(_, _, _, prefetch_k)); + prefetch(tiled_prefetch_b, pBgB(_, _, _, prefetch_k)); + } + + for (int k_tile = k_start_idx; k_tile < k_tile_count + k_start_idx; k_tile++, prefetch_k++) { + barrier_arrive(barrier_scope); + // Copy gmem to rmem for the first k_tile + copy(mainloop.tiled_copy_a, tAgA(_,_,_,k_tile), tArA); + copy(mainloop.tiled_copy_b, tBgB(_,_,_,k_tile), tBrB); + + if (prefetch_k < k_tile_count) { + prefetch(tiled_prefetch_a, pAgA(_, _, _, prefetch_k)); + prefetch(tiled_prefetch_b, pBgB(_, _, _, prefetch_k)); + } + + cute::gemm(tiled_mma, tCrA, tCrB, accum); + barrier_wait(barrier_scope); + } + } +}; + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// \ No newline at end of file diff --git a/include/cutlass/gemm/dispatch_policy.hpp b/include/cutlass/gemm/dispatch_policy.hpp index b742dfd76f..423e14d10f 100644 --- a/include/cutlass/gemm/dispatch_policy.hpp +++ b/include/cutlass/gemm/dispatch_policy.hpp @@ -1247,6 +1247,20 @@ struct MainloopDeviceAgnostic { using Schedule = KernelMultistage; }; #endif + +#if defined(CUTLASS_ENABLE_SYCL) +// Note: This dispatch policy is specifically added for CollectiveMma to support +// the integration of new MMA atoms (XE_DPAS_TT) and copy atoms for Intel XE architecture +template +struct MainloopXeL1Staged { + constexpr static int Stages = Stages_; + constexpr static int SubgroupSize = 16; + using ArchTag = arch::IntelXe; + using Schedule = KernelSchedule; + using ClusterShape = Shape<_1,_1,_1>; +}; +#endif + // n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule template< int LoadABPipelineStageCount_, diff --git a/include/cutlass/gemm/kernel/xe_gemm.hpp b/include/cutlass/gemm/kernel/xe_gemm.hpp index 1b1bebcc18..7a13d0fc1f 100644 --- a/include/cutlass/gemm/kernel/xe_gemm.hpp +++ b/include/cutlass/gemm/kernel/xe_gemm.hpp @@ -238,11 +238,11 @@ class GemmUniversal< constexpr auto workgroup_shape = WorkgroupTileShape{}; // (SUB_M,SUB_N,SUB_K) constexpr auto subgroup_shape = SubgroupTileShape{}; - Tensor mA_mkl = cute::get_xe_tensor(make_shape(M,K,L)); //(m,k,l) - Tensor mB_nkl = cute::get_xe_tensor(make_shape(N,K,L)); //(n,k,l) + Tensor cA = make_identity_tensor(make_shape(M,K,L)); // (M,K,L) + Tensor cB = make_identity_tensor(make_shape(N,K,L)); // (N,K,L) - Tensor gA = local_tile(mA_mkl, select<0,2>(blk_shape), make_coord(m_coord,_,l_coord)); - Tensor gB = local_tile(mB_nkl, select<1,2>(blk_shape), make_coord(n_coord,_,l_coord)); + Tensor gA = local_tile(cA, select<0,2>(blk_shape), make_coord(m_coord,_,l_coord)); + Tensor gB = local_tile(cB, select<1,2>(blk_shape), make_coord(n_coord,_,l_coord)); // Allocate the tiled_mma and the accumulators for the (M,N) subgroup_shape TiledMma tiled_mma;