diff --git a/CHANGELOG.md b/CHANGELOG.md index fe5c4d7a..21814c68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,67 @@ # CUTLASS 4.x +## [4.4.0](https://github.com/NVIDIA/cutlass/tree/main) (2026-01-23) + +### CuTe DSL +* New features + - Ahead of Time (AoT) compilation is now available! + + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/cute/export for example usage + - JAX support - you can now use CuTeDSL along with JAX + + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/jax for example usage + - Introduced versioning support in DSL: + + cutlass.__version__ for a string representation of DSL version + + cutlass.CUDA_VERSION for a version class to tell the CUDA version used for DSL + - Added CopyDsmemStoreOp to store data to distributed shared memory with explicit synchronization. + +* Bug fixing and improvements + - Fixed `cute.printf` with f-string + - Fixed an issue that cutlass.cuda.initialize_cuda_context() silently kills python + +* API changes + - Deprecate get_num_tmem_alloc_cols from blackwell_helpers.py. Use the one from tmem_allocator.py instead. + - Deprecate SM100_TMEM_CAPACITY_COLUMNS and SM100_TMEM_MIN_ALLOC_COLUMNS. + - LdMatrix16x16x8bOp and StMatrix16x8x8bOp now require explicit transpose=True when calling __init__, to avoid ambiguity in data transposition. + - LdMatrix16x16x8bOp copy traits updated to be faithful to PTX without permutations. Permuted variant is renamed to LdMatrix16x8x8bOp. + - group_bulk_copy_modes in async bulk copy example is now deprecated, use group_modes directly instead. + - cute.arch.calc_packed_f32x2_op default enable ftz to default disable ftz + +### CUTLASS C++ +* Add Hopper e2m1 to fp32 optimized conversion and e2m1 * TF32 tensor core GEMM. + - Set MmaType to tfloat32_t for FP32 mode. + - TF32 provides FP32 inputs with reduced precision (19-bit vs 32-bit) + - Set TileShapeK=64 for TF32 (K must be multiple of 8) + - Shuffle optimization enabled via `compute_memory_reordering_atom()` + - E2M1 -> FP32 -> TF32 TC path for mixed-precision GEMM + - Enable [example 55](https://github.com/NVIDIA/cutlass/tree/main/examples/55_hopper_mixed_dtype_gemm) with TF32 support +* Add [example 93](https://github.com/NVIDIA/cutlass/tree/main/examples/93_blackwell_low_latency_gqa/) for Blackwell low latency generation phase GQA kernel. + - Kernel design details please check [Readme](https://github.com/NVIDIA/cutlass/tree/main/examples/93_blackwell_low_latency_gqa/readme.md). +* Add [example 94](https://github.com/NVIDIA/cutlass/tree/main/examples/94_ada_fp8_blockwise/) for Ada FP8xFP8 -> BF16 GEMM with blockwise dequantization of input matrices in the MMA loop with FP32 accumulation. + - Generate additional device/kernel/threadblock files in CUTLASS include directory that add functionality to carry the scaling tensors + use them in MMA loop. + - Add gemm_blockwise to include files in [default_mma_core_sm80](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/threadblock/default_mma_core_sm80.h) +* Add Hopper SM90 State Space Decomposition (SSD) kernel in [example 111](https://github.com/NVIDIA/cutlass/tree/main/examples/111_hopper_ssd). +* Add Blackwell SM100 State Space Decomposition (SSD) kernel in [example 112](https://github.com/NVIDIA/cutlass/tree/main/examples/112_blackwell_ssd). +* Add support for arbitrary application-provided strides for block-scale tensors. + - Users and applications now must pass valid block-scale strides in all cases, even when the tensor is packed. +* Support 4x blockscaled public ptx for CUDA 13.1. +* Allow non-static `TmaGbasis` in `AuxTmaParams`. + - Some cases in attention kernel may require non-static `tma_gbasis`. + - Relax the restriction on `TmaGbasis` parameter of `AuxTmaParams` and users are allowed to manually construct a dynamic gbasis. +* Fix some kernel issues: + - Fix MSVC pre process issue. + - Fix a self assign issue in GEMV kernel. + - Fix a TMA descriptor bug where the CUDA driver is not properly setting the OOB address gen mode correctly. + - Fix memory fence for clc scheduler in Blackwell SM120 pingpong kernel. + - Fix missing SMEM alignment in Blackwell SM120 scale factors. +* Fix some profiler issues: + - Refactor L1 functional test generation logic to reduce the L1 test cases to avoid timeout. + - Fix a core dump issue for nvfp4 grouped GEMM kernel. + - Fix inconsistent GEMM verification logic. + - Rework grouped gemm verification logic for different types. +* Fix some broken links under `media/docs`. +* Various improvements and fixes from the community and CUTLASS team. Thanks to everyone who submitted PRs! +* Optimal code generation with CUDA toolkit versions 13.1. + ## [4.3.5](https://github.com/NVIDIA/cutlass/releases/tag/v4.3.5) (2026-01-09) ### CuTe DSL @@ -405,7 +466,7 @@ - Sorting performance results by GFLOPs/second: Users can now sort the final performance report based on GFLOPs/second, making it easier to identify the most efficient kernels. - Exhaustive search for best kernel performance in GFLOPs/second: The profiler now searches for the best-performing kernel across a range of problem sizes, swizzle sizes, rasterization orders, and dynamic cluster configurations to maximize performance. - Performance search under a fixed GEMM shape: Enables exhaustive tuning within a fixed GEMM shape, exploring various kernel parameters to find the best configuration. - - More detailed introductions and examples to leverage this feature can be found in [profiler.md](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/profiler.html#exhaustive-search-mode-and-top-k-output-ranking-according-to-performance-in-gflopss). + - More detailed introductions and examples to leverage this feature can be found in [profiler.md](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/profiler.html#exhaustive-search-mode-and-top-k-output-ranking-according-to-performance-in-gflops-s). * Support `void` as the D element in sm100 kernel epilogues. * Various improvements and fixes from the community and CUTLASS team. Thanks to everyone who submitted PRs! * Optimal code generation with CUDA toolkit versions 12.8U1. @@ -494,7 +555,7 @@ - [An improved mixed input GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/55_hopper_mixed_dtype_gemm/README.md) and a [lookup table implementation](https://github.com/NVIDIA/cutlass/tree/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu) for `INT4`x`FP8` scale-only mode. - [EVT nodes for Top-K selection and softmax](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp) and [GEMM example using those](https://github.com/NVIDIA/cutlass/tree/main/examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu). - [Programmatic Dependent Launch](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/arch/grid_dependency_control.h) (PDL) that leverages a new Hopper feature to speedup two back-to-back kernels, and its corresponding [documentations](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/dependent_kernel_launch.html). -- [A new debugging tool, synclog](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/arch/synclog.hpp), for dumping out all synchronization events from within a kernel to a file. Please see [synclog documentation](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/utilities.html#debugging-asynchronous-kernels-with-cutlasss-built-in-synclog-tool) for details. +- [A new debugging tool, synclog](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/arch/synclog.hpp), for dumping out all synchronization events from within a kernel to a file. Please see [synclog documentation](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/utilities.html#debugging-asynchronous-kernels-with-cutlass-s-built-in-synclog-tool) for details. - A new TMA-enabled [epilogue](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp) for grouped GEMM that brings significant performance improvement, as well as its EVT support. - A SIMT-enabled pointer-array [epilogue](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/collective/sm70_epilogue_vectorized_array.hpp). - A new [Ping-Pong kernel schedule for Grouped GEMM](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp) and some other optimizations. @@ -617,7 +678,7 @@ * [Batched B2B GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/13_two_tensor_op_fusion) now can run multiple Back-to-Back GEMM with the same problem size in parallel. * [Batched Strided GEMV](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemv.cu) support both row major and column major input matrix. * [Permute + GEMM fusion](https://github.com/NVIDIA/cutlass/tree/main/examples/39_gemm_permute) can fuse Permute with following GEMM now. Before, we only support fusing GEMM with Permute in the epilogue. -* [Row Broadcast](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h) can be fused in the epilogue. +* [Row Broadcast](https://github.com/NVIDIA/cutlass/blob/8236f30675bbe98f81d11c05764b77bfcb25b8cc/include/cutlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h) can be fused in the epilogue. * The GitHub branch is renamed from `master` to `main` in this release. * Optimal performance using [**CUDA 12.1**](https://developer.nvidia.com/cuda-downloads) * Updates and bugfixes from the community (thanks!) @@ -632,7 +693,7 @@ * Extensions to CUTLASS profiler to support threadblock cluster shapes in library and profiler tile configurations. * [CUTLASS library integration](https://github.com/NVIDIA/cutlass/tree/main/tools/library/src/gemm_operation_3x.hpp) for 3.x API kernels built through the new `CollectiveBuilder` API, enabling CUTLASS profiler. * Support for [Hopper GEMMs](https://github.com/NVIDIA/cutlass/tree/main/examples/48_hopper_warp_specialized_gemm) through the new 3.0 API with CuTe-based exposure of the Hopper [Tensor Memory Accelerator](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor) and [WGMMA Tensor Core](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-warpgroup-level-matrix-instructions) features. -* Set of examples that demonstrate the usage of the new 3.0 API to easily build GEMM kernels targeting Hopper: examples [48](https://github.com/NVIDIA/cutlass/tree/main/examples/48_hopper_warp_specialized_gemm), [49](https://github.com/NVIDIA/cutlass/tree/main/examples/49_hopper_gemm_schedules_with_collective_builder), and [50](https://github.com/NVIDIA/cutlass/tree/main/examples/50_hopper_gemm_with_epilogue_swizzle). +* Set of examples that demonstrate the usage of the new 3.0 API to easily build GEMM kernels targeting Hopper: examples [48](https://github.com/NVIDIA/cutlass/tree/main/examples/48_hopper_warp_specialized_gemm), [49](https://github.com/NVIDIA/cutlass/tree/main/examples/49_hopper_gemm_with_collective_builder), and [50](https://github.com/NVIDIA/cutlass/tree/main/examples/50_hopper_gemm_with_epilogue_swizzle). # CUTLASS 2.x @@ -662,7 +723,7 @@ * [CUTLASS Python](https://github.com/NVIDIA/cutlass/tree/main/examples/40_cutlass_py) now supports GEMM, CONV, Group GEMM for different data types as well as different epilogue flavours. * Optimizations for CUTLASS's [Grouped GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/24_gemm_grouped/gemm_grouped.cu) kernel. Threadblock scheduling part is improved. Some computation can be moved to the host side if applicable. [Grouped Syr2k](https://github.com/NVIDIA/cutlass/tree/main/examples/38_syr2k_grouped/syr2k_grouped.cu) kernels are added, too. * Optimizations for [GEMM+Softmax](https://github.com/NVIDIA/cutlass/tree/main/examples/35_gemm_softmax). All the reduction computation is fused into the previous GEMM. More template arguments are provided to fine tune the performance. -* [Grouped GEMM for Multihead Attention](https://github.com/NVIDIA/cutlass/tree/main/examples/41_multi_head_attention). This general group gemm based MHA does not require the sequence length of all GEMMs to be the same which makes it most useful for natural language processing. +* [Grouped GEMM for Multihead Attention](https://github.com/NVIDIA/cutlass/tree/main/examples/41_fused_multi_head_attention). This general group gemm based MHA does not require the sequence length of all GEMMs to be the same which makes it most useful for natural language processing. * [GEMM + Layer norm fusion for Ampere](https://github.com/NVIDIA/cutlass/tree/main/examples/37_gemm_layernorm_gemm_fusion/) splits the layernorm into two parts and both of them can be fused into the GEMMs before and after separately. In addition to use square sum to compute variance of layernorm, [Shift-K](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Computing_shifted_data) is provided if square sum raise numerical issues. * [GEMM Epilogue Permutation Fusion](https://github.com/NVIDIA/cutlass/tree/main/examples/39_gemm_permute) can apply user provided permutation layout mapping in the GEMM epilogue. * [Grouped convolution targeting implicit GEMM](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device/group_conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu) introduces the first group convolution implementation to CUTLASS. It is an Analytical implementation, not an Optimized. The restrictions are: 1) input and output channel number should be multiple of group number. 2) split-K is not supported. The implementation has 2 modes: @@ -689,7 +750,7 @@ * [TRMM](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/trmm_f32n_f32t_f32t_tensor_op_fast_f32_ls_sm80.cu) with [emitter](https://github.com/NVIDIA/cutlass/tree/main/python/cutlass_library/trmm_operation.py) * [Unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/testbed_rank_k_universal.h) * [CUTLASS Python](https://github.com/NVIDIA/cutlass/tree/main/examples/40_cutlass_py) demonstrating JIT compilation of CUTLASS kernels and a Python-based runtime using [CUDA Python](https://developer.nvidia.com/cuda-python) - * [Python-based runtime](https://github.com/NVIDIA/cutlass/tree/main/tools/library/scripts/rt.py) interoperable with existing emitters + * [Python-based runtime](https://github.com/NVIDIA/cutlass/blob/d572cc1aabfcbd45944219fb8690f0e49e22b5a3/tools/library/scripts/rt.py) interoperable with existing emitters * [GEMM + Softmax example](https://github.com/NVIDIA/cutlass/tree/main/examples/35_gemm_softmax) * [Gather and Scatter Fusion with GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/36_gather_scatter_fusion) can gather inputs and scatters outputs based on indices vectors in the same GEMM kernel. * It can select random rows in a row major matrix. @@ -811,7 +872,7 @@ ## [2.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.3.0) (2020-09-23) * [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/) * [Sparse Tensor Core GEMM kernels](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16n_f32t_tensor_op_f32_sparse_sm80.cu): - * Direct access to Sparse Tensor Cores and maximum performance via [`mma.sp.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma-and-friends) + * Direct access to Sparse Tensor Cores and maximum performance via [`mma.sp.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions) * Fast SGEMM targeting GeForce RTX 30-series CUDA Cores * Minor Features: * [Activation functions](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/thread/activation.h) such as [GeLU](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/thread/linear_combination_gelu.h) and [Sigmoid](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/epilogue/thread/linear_combination_sigmoid.h) @@ -825,7 +886,7 @@ ## [2.2.0](https://github.com/NVIDIA/cutlass/releases/tag/v2.2.0) (2020-06-08) * [NVIDIA Ampere Architecture features](https://devblogs.nvidia.com/nvidia-ampere-architecture-in-depth/) * Fast Tensor Core operations: - * Maximum performance via [`mma.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma-and-friends) + * Maximum performance via [`mma.sync`](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions) * Tensor Float 32, BFloat16, and double-precision data types * Mixed integer data types (int8, int4, bin1) * Asynchronous copy for deep software pipelines via [`cp.async`](https://docs.nvidia.com/cuda/parallel-thread-execution) @@ -874,9 +935,6 @@ ## [1.3.2](https://github.com/NVIDIA/cutlass/releases/tag/v1.3.2) (2019-07-09) * Performance improvement for Volta Tensor Cores TN and TT layouts. -## [1.3.1](https://github.com/NVIDIA/cutlass/releases/tag/v1.3.1) (2019-04-09) - * Corrected NVRTC unit tests. - ## [1.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v1.3.0) (2019-03-20) * Efficient GEMM kernel targeting Volta Tensor Cores via `mma.sync` instruction added in CUDA 10.1. diff --git a/CMakeLists.txt b/CMakeLists.txt index c941d5aa..0c9c35d3 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -267,6 +267,10 @@ if (WIN32) # Disable warning on Unicode characters list(APPEND CUTLASS_CUDA_NVCC_FLAGS -Xcompiler=/wd4819) + # Disable warning on macro expansion producing 'defined' has undefined behavior + list(APPEND CUTLASS_CUDA_NVCC_FLAGS -Xcompiler=/wd5105) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd5105") + # Disable excess x86 floating point precision that can lead to results being labeled incorrectly list(APPEND CUTLASS_CUDA_NVCC_FLAGS -Xcompiler=/fp:strict) endif(WIN32) @@ -620,8 +624,8 @@ if (MSVC) # # See https://developercommunity.visualstudio.com/t/msvc-incorrectly-defines-cplusplus/139261 - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:__cplusplus") - set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:__cplusplus /Zc:preprocessor") + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus -Xcompiler /Zc:preprocessor") endif() @@ -1134,6 +1138,8 @@ function(cutlass_generate_profiler_tests NAME) endif() endforeach() + message(STATUS "Finished processing ${CUTLASS_PROFILER_REGRESSION_LIST_FILE} to generate profiler-based tests.") + cutlass_add_executable_tests( ${NAME} cutlass_profiler diff --git a/README.md b/README.md index af692c80..b7b32ee7 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,9 @@ ![ALT](./media/images/gemm-hierarchy-with-epilogue-no-labels.png "Complete CUDA GEMM decomposition") # Overview -# CUTLASS 4.3.5 +# CUTLASS 4.4.0 -_CUTLASS 4.3.5 - Jan 2026_ +_CUTLASS 4.4.0 - Jan 2026_ CUTLASS is a collection of abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and related computations at all levels and scales within CUDA. It incorporates strategies for @@ -43,116 +43,64 @@ To get started quickly - please refer : - [CUTLASS C++ Quick Start Guide](https://docs.nvidia.com/cutlass/latest/media/docs/cpp/quickstart.html). - [CuTe DSL Quick Start Guide](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/quick_start.html). -# What's New in CUTLASS 4.3 +# What's New in CUTLASS 4.4 -## CuTe DSL -* New features: - - Supported Apache [TVM-FFI](https://tvm.apache.org/ffi/index.html) for further reduced host runtime overhead for JIT functions, better PyTorch and ML frameworks interopability - - Added fake tensor and stream to decouple compile jit function with "from_dlpack" flow. Now we no longer require users to have real tensor when compile jit function. - - Added FastDivmodDivisor with Python operator overloads, new APIs, Cute dialect integration, and optimized static tile scheduler performance for faster index mapping. - - Added l2 cache evict priority for tma related ops. Users could do fine-grain l2 cache control. - - Added Blackwell SM103 support. - - Multiple dependent DSOs in the wheel have been merged into one single DSO. - - New env var `CUTE_DSL_CACHE_DIR` to specify the path for dumping caches. - - Supported namedtuple and kwargs for JIT function arguments in tvm-ffi. - - Supported variadic tuples for JIT function argument in tvm-ffi. - - Added PDL support along with example [Kernel launch with Programmatic Dependent Launch](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/programmatic_dependent_launch.py) -* Debuggability improvements: - - Supported source location tracking for DSL APIs (Allow tools like ``nsight`` profiling to correlate perf metrics with Python source code) - - Supported dumping PTX and CUBIN code: [Hello World Example](https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/notebooks/hello_world.ipynb) -* More examples and notebooks to get started with CuTe DSL: - - Improved performance of [elementwise example](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/elementwise_apply.py): - + Generalize code to handle list of input tensors - + Generalize TV layout computation to handle different data types - - Improved [Blackwell SM100 persistent dense GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py): - + To demonstrate usage of new Pipeline APIs `PipelineProducer` and `PipelineConsumer` to simplify code without explicit pipeline state management (Exiting APIs are still maintained) - + Separated epilogue code for non-TMA and TMA implementation - - [Tutorial for Blackwell GEMM: Basic Blackwell SM100 GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm) - + [Baseline Blackwell GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py) achieves 84% SOL performance with MNK 8K - + More examples are coming for demo of optimization: `Baseline + X` - - [Tutorial for Async Pipeline API](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/async_pipeline.ipynb) - - Reworked [elementwise add notebook](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/elementwise_add.ipynb) with more details and detailed explanation about TV layout - + Updated implementation to handle general data type and multiple inputs - + Updated explanation for TV layout in simpler language - + Added visualization of TV Layout with 3rd party utils - - [Benchmark and autotune demonstration](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/notebooks/benchmark_autotune.ipynb) -* More examples of authorizing peak-performance kernels: - - [Blackwell SM100 mixed-input GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/mixed_input_gemm.py) - - [Blackwell SM100 persistent blockwise dense GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py) - - [Blackwell SM100 persistent blockwise contiguous grouped dense GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py) - - [Blackwell SM100 persistent blockwise masked grouped dense GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py) - - [Blackwell SM100 fmha bwd](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/fmha_bwd.py) - - [Blackwell SM100 mla](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/mla.py) - - [Hopper SM90 persistent dense GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/hopper/dense_gemm_persistent.py) - - [Blackwell GeForce batched dense GEMM](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell_geforce/dense_gemm.py) - - [Ampere HSTU Attention](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/ampere/hstu_attention.py) -* API updates: - - Please refer to [DSL API changelog](https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_api/changelog.html) for details -* Bug fixings and improvements - - Add mma_tiler_n=64 and mma_tiler_n=192 support in [Blackwell SM100 persistent dense blockscaled GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py). - - Fixed ``TensorSSA.reduce`` to support static value as initial value - - Updated docstring for following APIs to be more concise and easier to understand: - - ``make_layout_tv`` - - ``is_static`` - - ``PipelineAsync`` - - ``SmemAllocator`` - - Fixed documentation for ``pipeline``, ``utils`` and ``cute.math`` - - Added overlapping accumulator optimization for block tile N = 256 case for better epilogue latency hiding in [Blackwell SM100 persistent dense blockscaled GEMM with static scheduling](https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py). - - Fixed TensorSSA.__getitem__ indexing to match CuTe's indexing convention - - Fixed an issue with cutlass.max and cutlass.min - - Fixed an issue with mark_compact_shape_dynamic - - Fixed device reset issue with tvm-ffi - - Fixed tvm-ffi export compiled function - - Fixed an issue of CUDA JitExecutor when unloading kernels - - Fixed an issue of allocating max smem when there's statically allocated smem - - Fixed an issue when JIT function argument with union type annotation for tvm-ffi - - Clearer error message for the case of runtime error cudaErrorInsufficientDriver - - Fixed a frame refcnt issue with cuda graph - - Enhancement for tvm-ffi AoT case for earlier module unload - - Fixed order issue in make_smem_layout_a in utils/hopper_helpers.py - - Fixed the unexpected CPU overhead issue introduced by 4.3.4 -* Update copyright to 2026. +### CuTe DSL +* New features + - Ahead of Time (AoT) compilation is now available! + + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/cute/export for example usage + - JAX support - you can now use CuTeDSL along with JAX + + Refer to files under https://github.com/NVIDIA/cutlass/tree/main/examples/python/CuTeDSL/jax for example usage + - Introduced versioning support in DSL: + + cutlass.__version__ for a string representation of DSL version + + cutlass.CUDA_VERSION for a version class to tell the CUDA version used for DSL + - Added CopyDsmemStoreOp to store data to distributed shared memory with explicit synchronization. -## CUTLASS C++ -* Further enhance Blackwell SM100 Attention kernels in [example 77](https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha/). - - Add softmax skip correction. - - Fix a shared memory allocation bug where it needs to opt in maximum dynamics shared memory explicitly once it exceeds 48KB. - - Fix a dead hang issue caused by early return warp. -* Add support through cmdline argument lists for `batch`, `no_verif`, `cluster_shape` and `cluster_shape_fallback` in [example 89](https://github.com/NVIDIA/cutlass/tree/main/examples/89_sm103_fp4_ultra_gemm/). -* Add Ragged Contiguous Grouped gemm kernel in [example 92](https://github.com/NVIDIA/cutlass/tree/main/examples/92_blackwell_moe_gemm/). - - This kernel uses a TMA 3D load to load the weights matrix and use the tensormap update method to load activations. -* Add 256x128 tile size support for Hopper SM90 deepgemm in [example 67](https://github.com/NVIDIA/cutlass/tree/main/examples/67_hopper_fp8_warp_specialized_gemm_with_blockwise_scaling/). - - Performance is optimized to align with Deepseek implementation. -* Simplification of API for MoE gemms. - - Instead of requiring users to call several cute utilities to set up the stride, API `moe_stride_utils` is introduced to help setup strides in the kernel. - - Instead of requiring users to set vectors like `problem_shapes_device` and `problem_shapes_hosts`, a new problem shape struct called `MoEProblemShape` is introduced which takes in max_m, max_n, max_k and counts vector as input and deduce problem shapes internally whenever required. -* Enable GEMM_K = 0 in grouped gemm. -* Optimize group gemm kernels by enabling async TMA desc update. -* Support Blackwell SM100 convolution stream-K kernel. - - Unit tests: [fprop_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/fprop/sm100_conv3d_fprop_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu), [dgrad_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/dgrad/sm100_conv3d_dgrad_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu), [wgrad_streamK](https://github.com/NVIDIA/cutlass/tree/main/test/unit/conv/device_3x/wgrad/sm100_conv2d_wgrad_implicit_gemm_f16_f16_f16_tensorop_f16_streamk.cu). -* Add Blackwell SM100 sparse gemm compressor unit tests. - - Unit tests: [compressor_fp16](https://github.com/NVIDIA/cutlass/tree/main/test/unit/transform/device/sm100_sparse_gemm_compressor_f16.cu). - - Add sub-bytes and runtime data type support in compressor unit test testbed. -* Add profiler support for: - - Blackwell SM100 and SM120 blockscaled sparse kernels. - - New MoE grouped gemm API. - - Blackwell SM100 cpasync kernel. +* Bug fixing and improvements + - Fixed `cute.printf` with f-string + - Fixed an issue that cutlass.cuda.initialize_cuda_context() silently kills python + +* API changes + - Deprecate get_num_tmem_alloc_cols from blackwell_helpers.py. Use the one from tmem_allocator.py instead. + - Deprecate SM100_TMEM_CAPACITY_COLUMNS and SM100_TMEM_MIN_ALLOC_COLUMNS. + - LdMatrix16x16x8bOp and StMatrix16x8x8bOp now require explicit transpose=True when calling __init__, to avoid ambiguity in data transposition. + - LdMatrix16x16x8bOp copy traits updated to be faithful to PTX without permutations. Permuted variant is renamed to LdMatrix16x8x8bOp. + - group_bulk_copy_modes in async bulk copy example is now deprecated, use group_modes directly instead. + - cute.arch.calc_packed_f32x2_op default enable ftz to default disable ftz + +### CUTLASS C++ +* Add Hopper e2m1 to fp32 optimized conversion and e2m1 * TF32 tensor core GEMM. + - Set MmaType to tfloat32_t for FP32 mode. + - TF32 provides FP32 inputs with reduced precision (19-bit vs 32-bit) + - Set TileShapeK=64 for TF32 (K must be multiple of 8) + - Shuffle optimization enabled via `compute_memory_reordering_atom()` + - E2M1 -> FP32 -> TF32 TC path for mixed-precision GEMM + - Enable [example 55](https://github.com/NVIDIA/cutlass/tree/main/examples/55_hopper_mixed_dtype_gemm) with TF32 support +* Add [example 93](https://github.com/NVIDIA/cutlass/tree/main/examples/93_blackwell_low_latency_gqa/) for Blackwell low latency generation phase GQA kernel. + - Kernel design details please check [Readme](https://github.com/NVIDIA/cutlass/tree/main/examples/93_blackwell_low_latency_gqa/readme.md). +* Add [example 94](https://github.com/NVIDIA/cutlass/tree/main/examples/94_ada_fp8_blockwise/) for Ada FP8xFP8 -> BF16 GEMM with blockwise dequantization of input matrices in the MMA loop with FP32 accumulation. + - Generate additional device/kernel/threadblock files in CUTLASS include directory that add functionality to carry the scaling tensors + use them in MMA loop. + - Add gemm_blockwise to include files in [default_mma_core_sm80](https://github.com/NVIDIA/cutlass/tree/main/include/cutlass/gemm/threadblock/default_mma_core_sm80.h) +* Add Hopper SM90 State Space Decomposition (SSD) kernel in [example 111](https://github.com/NVIDIA/cutlass/tree/main/examples/111_hopper_ssd). +* Add Blackwell SM100 State Space Decomposition (SSD) kernel in [example 112](https://github.com/NVIDIA/cutlass/tree/main/examples/112_blackwell_ssd). +* Add support for arbitrary application-provided strides for block-scale tensors. + - Users and applications now must pass valid block-scale strides in all cases, even when the tensor is packed. +* Support 4x blockscaled public ptx for CUDA 13.1. +* Allow non-static `TmaGbasis` in `AuxTmaParams`. + - Some cases in attention kernel may require non-static `tma_gbasis`. + - Relax the restriction on `TmaGbasis` parameter of `AuxTmaParams` and users are allowed to manually construct a dynamic gbasis. * Fix some kernel issues: - - Fix a race check issue of Blackwell SM103 kernels by adding missing elect one for prefetch barrier initialization. - - Allow user to directly specify the number of stages for Hopper sm90 mixed input gemm. - - Remove warnings caused by cuda vector type alignment setting in CUDA 13. - - Remove problematic `cutlass::int8_t` and replace it with `int8_t`. - - Fix a few bugs in distributed gemm API and examples. - - Fix handling negative zero in sparse compressor. - - Add missing `wait_on_dependent_grids` for PDL use case. - - Work around a driver bug which will cause occasionally errors when executing kernels. - - Use CUDA Driver Get Version Runtime APIs Rather than Driver APIs. + - Fix MSVC pre process issue. + - Fix a self assign issue in GEMV kernel. + - Fix a TMA descriptor bug where the CUDA driver is not properly setting the OOB address gen mode correctly. + - Fix memory fence for clc scheduler in Blackwell SM120 pingpong kernel. + - Fix missing SMEM alignment in Blackwell SM120 scale factors. * Fix some profiler issues: - - Add some missing reference kernels. - - Support VoidC reference kernels. - - Add calculation of scale factor A and B in function `bytes_with_problem_shape` of block scaled profiler. - - Fix an issue when epilogue tile N is not divided by default subtile N. -* Update copyright to 2026. + - Refactor L1 functional test generation logic to reduce the L1 test cases to avoid timeout. + - Fix a core dump issue for nvfp4 grouped GEMM kernel. + - Fix inconsistent GEMM verification logic. + - Rework grouped gemm verification logic for different types. +* Fix some broken links under `media/docs`. Note: CUTLASS 4.x builds are known to be down on Windows platforms for all CUDA toolkits. CUTLASS team is working on a fix. diff --git a/examples/111_hopper_ssd/111_hopper_ssd.cu b/examples/111_hopper_ssd/111_hopper_ssd.cu new file mode 100644 index 00000000..78d77530 --- /dev/null +++ b/examples/111_hopper_ssd/111_hopper_ssd.cu @@ -0,0 +1,849 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + * + **************************************************************************************************/ + +#include +#include "cutlass/util/command_line.h" + +#include "cutlass/cutlass.h" + +#include "cute/tensor.hpp" +#include "cute/layout.hpp" +#include "cutlass/kernel_hardware_info.hpp" + +#include "thrust/universal_vector.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/device/tensor_fill.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +#include "reference/reference_ssd_cumsum.hpp" +#include "reference/reference_ssd.hpp" + +#include "cutlass/transform/device/transform_universal_adapter.hpp" + +#include "device/ssd.hpp" +#include "kernel/sm90_ssd_kernel_builder.hpp" + +using namespace cute; + +// Command line options parsing +struct Options { + + using Element = cutlass::bfloat16_t; + using ElementAcc = float; + using ElementDA = float; + static constexpr bool D_HAS_HDIM = true; + static constexpr bool HAS_D = true; + static constexpr bool HAS_Z = true; + + bool help; + bool error; + + // All static number now + int G = 2; + int B = 3; + int E = 2; + int H = 2; + // Reference kernel doesn't support dynamic C now. + static constexpr auto C = Int<8>{}; + static constexpr auto D = Int<64>{}; + static constexpr auto L = Int<128>{}; + static constexpr auto N = Int<128>{}; + int EH = E * H; + + int iterations; + bool verify; + bool verbose; + + int warmups; + bool measure; + + Options(): + help(false), + error(false), + iterations(1), verify(true), + measure(false), warmups(3) + {} + + // Parses the command line + void parse(int argc, char const **args) { + cutlass::CommandLine cmd(argc, args); + + Options defaults; + + if (cmd.check_cmd_line_flag("help")) { + help = true; + return; + } + + cmd.get_cmd_line_argument("iterations", iterations, defaults.iterations); + cmd.get_cmd_line_argument("G", G, defaults.G); + cmd.get_cmd_line_argument("B", B, defaults.B); + cmd.get_cmd_line_argument("E", E, defaults.E); + cmd.get_cmd_line_argument("H", H, defaults.H); + verbose = cmd.check_cmd_line_flag("verbose"); + verify = !(cmd.check_cmd_line_flag("without_verify")); + + EH = E*H; + + if (iterations > 1) { + measure = true; + verbose = true; + } + + auto problem_shape = cute::make_tuple(G, B, EH, C, L, D, N); + cute::print("problem_shape : "); cute::print(problem_shape); cute::print("\n"); + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "111_hopper_ssd\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement\n\n" + << " --iterations= Benchmarking iterations.\n" + << " --without_verify Don't verify the results.\n" + << " --verbose Print execution time per kernel\n" + << " --G= Group\n" + << " --B= Batch\n" + << " --E= Expanded factor\n" + << " --H= Number of heads\n" + << "\n"; + + return out; + } + + + auto get_problem_shape() const { + return cute::make_tuple(G, B, EH, C, L, D, N); + } + + // acceptable layout by cuDNN + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + + auto layoutX() const { + auto layout = make_layout(make_shape(L, C, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutDelta() const { + auto layout = make_layout(make_shape(L, C, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutDeltaA() const { + auto layout = make_layout(make_shape(L, C, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutB() const { + auto layout = make_layout(make_shape(L, C, N, G, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutC() const { + auto layout = make_layout(make_shape(L, C, N, G, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutY() const { + auto layout = make_layout(make_shape(L, C, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutF() const { + auto layout = make_layout(make_shape(N, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutD() const { + if constexpr (D_HAS_HDIM) { + auto layout = make_layout(make_shape(D, EH)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + else { + auto layout = make_layout(make_shape(Int<1>{}, EH)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + } + + auto layoutZ() const { + auto layout = make_layout(make_shape(L, C, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + // transformed layout for kernel parameters + + auto layoutX_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B)); + return make_layout( + make_shape(D,L,int32_t(C),EH*B), + make_stride( + stride<2>(layout), + stride<0>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutB_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),N,G*B)); + return make_layout( + make_shape(L,N,int32_t(C),G*B), + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutC_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),N,G*B)); + return make_layout( + make_shape(L,N,int32_t(C),G*B), + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutDelta_transformed() const { + return make_layout(make_shape(L,int32_t(C),EH*B)); + } + + auto layoutY_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B)); + return make_layout( + make_shape(L,D,int32_t(C),EH*B), // (M,K,L,...) + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutF_transformed() const { + auto layout = make_layout(make_shape(N,D,EH*B)); + return make_layout( + make_shape(D,N,EH*B), + make_stride( + stride<1>(layout), + stride<0>(layout), + stride<2>(layout) + ) + ); + } + + auto layoutD_transformed() const { + if constexpr (D_HAS_HDIM) { + return make_layout(make_shape(D, EH)); + } + else { + return make_layout(make_shape(Int<1>{}, EH)); + } + } + + auto layoutZ_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B)); + return make_layout( + make_shape(L,D,int32_t(C),EH*B), + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + +}; + +template +static void +initialize_values( + thrust::universal_vector& dst_ptr, + cutlass::Distribution::Kind dist_kind, + uint64_t seed, + Element var = Element(1.f)) { + if (cutlass::Distribution::Uniform == dist_kind) { + int scope = 2; + cutlass::reference::host::BlockFillRandomUniform( + dst_ptr.data().get(), dst_ptr.size(), seed, scope, -scope, 0); + } + else if (cutlass::Distribution::AllZeros == dist_kind) { + cutlass::reference::host::BlockFillRandomUniform( + dst_ptr.data().get(), dst_ptr.size(), seed, 0, 0, 0); + } + else if (cutlass::Distribution::AllOnes == dist_kind) { + cutlass::reference::host::BlockFillRandomUniform( + dst_ptr.data().get(), dst_ptr.size(), seed, 1, 1, 0); + } + else if (cutlass::Distribution::Gaussian == dist_kind) { + cutlass::reference::device::BlockFillRandomGaussian( + dst_ptr.data().get(), dst_ptr.size(), seed, (Element) 0, var); + } + else if (cutlass::Distribution::Sequential == dist_kind) { + cutlass::reference::host::BlockFillSequential(dst_ptr.data().get(), dst_ptr.size()); + } + else { + std::cerr << "Invalid distribution kind!\n."; + exit(1); + } +} + +template < + class Options_ +> +struct TestBed { + using Option = Options_; + using Element = typename Option::Element; + using ElementDA = typename Option::ElementDA; + using ElementAcc = typename Option::ElementAcc; + + thrust::universal_vector tensor_X; + thrust::universal_vector tensor_DeltaA; + thrust::universal_vector tensor_DeltaA_cumsum; + thrust::universal_vector tensor_Delta; + thrust::universal_vector tensor_B; + thrust::universal_vector tensor_C; + thrust::universal_vector tensor_D; + thrust::universal_vector tensor_Y; + thrust::universal_vector tensor_Z; + thrust::universal_vector tensor_Y_ref_0; + thrust::universal_vector tensor_Y_ref_1; + thrust::universal_vector tensor_F; + thrust::universal_vector tensor_F_ref_0; + thrust::universal_vector tensor_F_ref_1; + + cutlass::Distribution::Kind init_X = cutlass::Distribution::Uniform; + cutlass::Distribution::Kind init_DeltaA = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind init_Delta = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform; + cutlass::Distribution::Kind init_C = cutlass::Distribution::Uniform; + + using TileShape = decltype(make_shape(Options::L, Options::D, Options::N)); // (L, D, N) + using SsdOperation = cutlass::ssd::device::SSD< + typename cutlass::ssd::kernel::Sm90SsdBuilder< + Element, ElementDA, ElementAcc, Element, + TileShape, + Option::HAS_D, Option::D_HAS_HDIM, Option::HAS_Z + >::Kernel>; + using CumsumKenrel = cutlass::ssd::kernel::CumsumKernel; + using CumsumOperation = cutlass::transform::device::TransformUniversalAdapter; + + bool initialize(Options const& options, const cutlass::KernelHardwareInfo& hw_info, uint64_t seed = 2023) { + auto [g, b, eh, c, l, d, n] = options.get_problem_shape(); + assert(g == 1 && "Only group size == 1 is supported") ; + + auto size_X = b * eh * c * l * d; + auto size_DeltaA = b * eh * c * l; + auto size_Delta = b * eh * c * l; + auto size_B = g * b * c * n * l; + auto size_C = g * b * c * n * l; + auto size_Y = b * eh * c * l * d; + auto size_F = b * eh * d * n; + + tensor_X .resize(sizeof(Element) * size(options.layoutX())); + tensor_DeltaA .resize(sizeof(Element) * size(options.layoutDeltaA())); + tensor_Delta .resize(sizeof(Element) * size(options.layoutDelta())); + tensor_B .resize(sizeof(Element) * size(options.layoutB())); + tensor_C .resize(sizeof(Element) * size(options.layoutC())); + tensor_D .resize(sizeof(Element) * size(options.layoutD())); + tensor_Z .resize(sizeof(Element) * size(options.layoutZ())); + tensor_Y .resize(sizeof(Element) * size(options.layoutY())); + tensor_Y_ref_0.resize(sizeof(Element) * size(options.layoutY())); + tensor_Y_ref_1.resize(sizeof(Element) * size(options.layoutY())); + tensor_F .resize(sizeof(Element) * size(options.layoutF())); + tensor_F_ref_0.resize(sizeof(Element) * size(options.layoutF())); + tensor_F_ref_1.resize(sizeof(Element) * size(options.layoutF())); + + tensor_DeltaA_cumsum.resize(sizeof(ElementDA) * size(options.layoutDeltaA())); + + // Limit distribution to reduce skew between hosts and devices + initialize_values(tensor_X, init_X, seed); + initialize_values(tensor_DeltaA, init_DeltaA, seed + 1, Element(0.05f)); + initialize_values(tensor_Delta, init_Delta, seed + 3, Element(0.05f)); + initialize_values(tensor_B, init_B, seed + 5); + initialize_values(tensor_C, init_C, seed + 7); + initialize_values(tensor_D, init_C, seed + 9); + initialize_values(tensor_Z, init_X, seed); + + cudaError_t result; + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the Initialization kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + } + + // apply cumsum(device) before kernel launch + typename CumsumOperation::Arguments arguments{ + make_shape(int(b), int(eh), int(c), int(l)), + { + tensor_DeltaA.data().get(), + tensor_DeltaA_cumsum.data().get(), + }, + hw_info + }; + + CumsumOperation op; + + size_t workspace_size = CumsumOperation::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + cutlass::Status status = op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "This kernel is not supported. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + status = op.initialize(arguments, workspace.get()); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + // may be used uninitialized + cudaEvent_t start; + cudaEvent_t end; + cudaEventCreate(&start); + cudaEventCreate(&end); + + // warm up + if (options.measure) { + for (int i = 0; i < options.warmups; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + } + result = cudaEventRecord(start); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + // Run + for (int i = 0; i < options.iterations; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + result = cudaEventRecord(end); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return false; + } + + float runtime_ms = 0; + result = cudaEventElapsedTime(&runtime_ms, start, end); + if (result != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + runtime_ms /= static_cast(options.iterations); + + if (options.verbose) { + printf("[iters = %d, warmups = %d] cumsum kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms); + } + + return true; + } + + bool sufficient() const { + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + int max_smem_size; + result = cudaDeviceGetAttribute(&max_smem_size, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_idx); + if (result != cudaSuccess) { + throw std::runtime_error("cudaDeviceGetAttribute() failed"); + } + + return true; + } + + bool run(Options const& options, const cutlass::KernelHardwareInfo& hw_info) { + if (!sufficient()) { + std::cerr << "Test waived due to insufficient CUDA device.\n"; + return true; + } + + if (!initialize(options, hw_info)) { + std::cerr << "Failed to initialize the test.\n"; + return true; + }; + + auto [g, b, eh, c, l, d, n] = options.get_problem_shape(); + typename SsdOperation::Arguments arguments{ + make_shape(int(g), int(b), int(eh), int(c), int(l), int(d), int(n)), + { + tensor_X.data().get(), + tensor_DeltaA_cumsum.data().get(), + tensor_Delta.data().get(), + tensor_B.data().get(), + tensor_C.data().get(), + options.layoutX_transformed(), + options.layoutB_transformed(), + options.layoutC_transformed(), + options.layoutDelta_transformed() + }, + { + tensor_Y.data().get(), + tensor_F.data().get(), + tensor_D.data().get(), + tensor_Z.data().get(), + options.layoutY_transformed(), + options.layoutF_transformed(), + options.layoutD_transformed(), + options.layoutZ_transformed() + }, + hw_info + }; + + SsdOperation op; + + size_t workspace_size = SsdOperation::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + cutlass::Status status = op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "This kernel is not supported. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + status = op.initialize(arguments, workspace.get()); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + cudaError_t result; + // may be used uninitialized + cudaEvent_t start; + cudaEvent_t end; + cudaEventCreate(&start); + cudaEventCreate(&end); + + // warm up + if (options.measure) { + for (int i = 0; i < options.warmups; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + } + result = cudaEventRecord(start); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + // Run + for (int i = 0; i < options.iterations; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + result = cudaEventRecord(end); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return false; + } + + float runtime_ms = 0; + result = cudaEventElapsedTime(&runtime_ms, start, end); + if (result != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + runtime_ms /= static_cast(options.iterations); + + if (options.verbose) { + printf("[iters = %d, warmups = %d] ssd kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms); + printf("smem size = %d\n", SsdOperation::Kernel::SharedStorageSize); + } + // Matrix + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + auto mY_ref_0 = cute::make_tensor(tensor_Y_ref_0.data().get(), options.layoutY()); + auto mY_ref_1 = cute::make_tensor(tensor_Y_ref_1.data().get(), options.layoutY()); + auto mY_res = cute::make_tensor(tensor_Y.data().get(), options.layoutY()); + auto mF_ref_0 = cute::make_tensor(tensor_F_ref_0.data().get(), options.layoutF()); + auto mF_ref_1 = cute::make_tensor(tensor_F_ref_1.data().get(), options.layoutF()); + auto mF_res = cute::make_tensor(tensor_F.data().get(), options.layoutF()); + auto mX = cute::make_tensor(tensor_X.data().get(), options.layoutX()); + auto mB = cute::make_tensor(tensor_B.data().get(), options.layoutB()); + auto mC = cute::make_tensor(tensor_C.data().get(), options.layoutC()); + auto mD = cute::make_tensor(tensor_D.data().get(), options.layoutD()); + auto mZ = cute::make_tensor(tensor_Z.data().get(), options.layoutZ()); + auto mDelta = cute::make_tensor(tensor_Delta.data().get(), options.layoutDelta()); + auto mDeltaA = cute::make_tensor(tensor_DeltaA.data().get(), options.layoutDeltaA()); + + // Reference Device kernel + if (options.verify) { + ssd_reference( + mY_ref_1, + mF_ref_1, + mX, + mDelta, + mDeltaA, + mB, + mC, + mD, + mZ, + options + ); + } + + bool passed = true; + + if (options.verify) { + printf("[TensorY]verifying...\n"); + passed &= compare_reference<5>(mY_ref_1, mY_res); + printf("[TensorF]verifying...\n"); + passed &= compare_reference<4>(mF_ref_1, mF_res); + } + + return passed; + } + + template< + int TensorDim, + class Engine, class Layout + > + static constexpr bool + compare_reference( + cute::Tensor const& reference, + cute::Tensor const& computed, + float epsilon = 0.05f) { + if (size(reference) != size(computed)) { + return false; + } + + bool passed = true; + if (epsilon == 0.0f) { + // fast refcheck w/o epsilon + for (size_t i = 0; i < size_t(size(reference)); ++i) { + if (reference(i) != computed(i)) { + passed = false; + printf("[%llu] %f, %f\n", static_cast(i), + float(reference(i)), float(computed(i))); + break; + } + } + } + else { + // refcheck with epsilon + for (size_t i = 0; i < size_t(size(reference)); ++i) { + auto ref = static_cast(reference(i)); + auto act = static_cast(computed(i)); + auto abs_error = std::abs(act - ref); + auto rel_error = abs_error / (std::max(std::abs(act), std::abs(ref)) + 0.00001f); + if (std::isnan(abs_error) || std::isnan(rel_error) || + std::min(rel_error, abs_error) > epsilon) { + passed = false; + printf("[%llu] %f, %f\n", static_cast(i), + float(reference(i)), float(computed(i))); + break; + } + } + } + if (not passed) { + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + auto m = cute::shape<2>(reference); + auto n = cute::shape(reference); + printf("reference:\n"); + for (int mi = 0; mi < m; ++mi) { + for (int ni = 0; ni < n; ++ni) { + if constexpr (TensorDim == 5) { + printf("%.4f ", static_cast(reference(0,0,mi,2,ni))); + } + else { + printf("%.4f ", static_cast(reference(0,0,mi,ni))); + } + } + printf("\n"); + } + printf("\n"); + printf("computed:\n"); + for (int mi = 0; mi < m; ++mi) { + for (int ni = 0; ni < n; ++ni) { + if constexpr (TensorDim == 5) { + printf("%.4f ", static_cast(computed(0,0,mi,2,ni))); + } + else { + printf("%.4f ", static_cast(computed(0,0,mi,ni))); + } + } + printf("\n"); + } + printf("\n"); + } + return passed; + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + +int main(int argc, char const **args) { + + cudaDeviceProp props; + + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return -1; + } + + if (__CUDACC_VER_MAJOR__ < 12 || props.major < 9) { + std::cout + << "This example requires a GPU of NVIDIA's Hopper Architecture or " + << "later (compute capability 90 or greater) and CUDA 12.0 or greater.\n"; + return 0; + } + else if (__CUDACC_VER_MAJOR__ < 12 || (props.major != 9 || props.minor != 0)) { + std::cout + << "This example requires a GPU of NVIDIA's Hopper Architecture " + << "(compute capability 90) and CUDA 12.0 or greater.\n"; + return 0; + } + +#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + + // + // Parse options + // + + Options options; + + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.error) { + std::cerr << "Aborting execution." << std::endl; + return -1; + } + + // Execute kernel + + printf("start testing....\n"); + + // The KernelHardwareInfo struct holds the number of SMs 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.device_id = 0; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + // Check Device/Host ref kernel + TestBed testbed{}; + bool passed = testbed.run(options, hw_info); + + if (passed) { + printf("everything is ok.\n"); + } + else { + printf("something is wrong!!!!!\n"); + } + +#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED) + + return 0; +} diff --git a/examples/111_hopper_ssd/CMakeLists.txt b/examples/111_hopper_ssd/CMakeLists.txt new file mode 100644 index 00000000..e1962109 --- /dev/null +++ b/examples/111_hopper_ssd/CMakeLists.txt @@ -0,0 +1,41 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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_property( + SOURCE 111_hopper_ssd.cu + PROPERTY COMPILE_FLAGS "--use_fast_math" + ) + +cutlass_example_add_executable( + 111_hopper_ssd + 111_hopper_ssd.cu + ) + +if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang"))) + +endif() diff --git a/examples/111_hopper_ssd/README.md b/examples/111_hopper_ssd/README.md new file mode 100644 index 00000000..7f5f1aa1 --- /dev/null +++ b/examples/111_hopper_ssd/README.md @@ -0,0 +1,67 @@ +# NVIDIA Hopper SSD (State Space Decomposition) CUDA Example + +## Overview + +This example demonstrates the implementation of State Space Decomposition (SSD) operations on NVIDIA's Hopper GPU architecture. It showcases the use of CUTLASS library components for high-performance tensor computations that efficiently leverage Hopper's advanced hardware capabilities, including TMA (Tensor Memory Accelerator) and warp specialization. + +## System Requirements ++ NVIDIA GPU with Hopper Architecture (compute capability 9.0) ++ CUDA Toolkit 12.0 or newer ++ C++17 compatible compiler + +## Build the example +Follow the cutlass example building. + +## Command Line Options +The example supports the following command line options: +--help: Display the usage statement +--iterations=: Number of iterations for benchmarking (default: 1) +--without_verify: Skip result verification +--verbose: Print execution time per kernel +--G=: Group size (default: 2) +--B=: Batch size (default: 3) +--E=: Expanded factor (default: 2) +--H=: Number of heads (default: 2) + +## Limitation ++ Only support LxDxN = 128x64x128 ++ Uses bfloat16 precision for input/output tensors + +## Performance ++ Utilizes TMA (Tensor Memory Accelerator) for efficient memory access ++ Implements warp-specialized kernels for optimal resource utilization ++ Limited by the SEGSUM (segment sum) computation part ++ ALU bound operation + +# Copyright + +Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. +``` + diff --git a/examples/111_hopper_ssd/collective/common.hpp b/examples/111_hopper_ssd/collective/common.hpp new file mode 100644 index 00000000..a1842e0f --- /dev/null +++ b/examples/111_hopper_ssd/collective/common.hpp @@ -0,0 +1,174 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/kernel_hardware_info.h" +#include "cute/tensor.hpp" + +namespace cutlass::ssd::collective { + +using namespace cute; + +template +CUTE_DEVICE void gemm_reset_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) { + constexpr int rA = decltype(rank(tA))::value; + constexpr int rB = decltype(rank(tB))::value; + constexpr int rC = decltype(rank(tC))::value; + if constexpr (rA == 2 && rB == 2 && rC == 1) { + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<1>(tA); k_block++) { + cute::gemm(atom, tA(_,k_block), tB(_,k_block), tC); + atom.accumulate_ = GMMA::ScaleOut::One; + } + } + else { + static_assert(rA == 3 && rB == 3 && rC == 3); + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tA); k_block++) { + cute::gemm(atom, tA(_,_,k_block), tB(_,_,k_block), tC); + atom.accumulate_ = GMMA::ScaleOut::One; + } + } +} + +template +CUTE_DEVICE void gemm_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) { + atom.accumulate_ = GMMA::ScaleOut::Zero; + gemm_reset_zero_acc(atom, tA, tB, tC); +} + +template class Primitive, cute::GMMA::Major tA, cute::GMMA::Major tB, cute::GMMA::ScaleIn sA, cute::GMMA::ScaleIn sB> +inline auto __device__ constexpr convert_to_gmma_rs(cute::MMA_Atom> const& tiled_mma) { + using Atom = cute::MMA_Atom>; + using ElementA = typename Atom::ValTypeA; + using ElementB = typename Atom::ValTypeB; + using ElementC = typename Atom::ValTypeC; + using Shape_MNK = typename Atom::Shape_MNK; + using RS = decltype(cute::GMMA::rs_op_selector()); + return cute::MMA_Atom{}; +} + +template class Primitive, cute::GMMA::ScaleIn sA, cute::GMMA::ScaleIn sB> +inline auto __device__ constexpr convert_to_gmma_rs(cute::MMA_Atom> const& tiled_mma) { + using Atom = cute::MMA_Atom>; + using ElementA = typename Atom::ValTypeA; + using ElementB = typename Atom::ValTypeB; + using ElementC = typename Atom::ValTypeC; + using Shape_MNK = typename Atom::Shape_MNK; + constexpr auto tA = cute::GMMA::Major::K; + constexpr auto tB = cute::GMMA::Major::K; + using RS = decltype(cute::GMMA::rs_op_selector()); + return cute::MMA_Atom{}; +} + +template +CUTE_DEVICE auto constexpr convert_to_gmma_rs(cute::TiledMMA const& tiled_mma) { + return cute::TiledMMA{}; +} + +template +CUTE_DEVICE auto constexpr convert_c_layout_to_a_layout(CLayout const& c, AValueShape const& a) { + return make_layout( + make_shape(a, shape<1>(c), make_shape(shape<2>(c), size<0>(c) / size(a))), + make_stride(stride<0>(c), stride<1>(c), make_stride(stride<2>(c), size<2>(a) * stride<0,2>(c)))); +} + +template +CUTE_DEVICE constexpr auto unstageSmemLayout(Layout const& layout, Stages stages = {}) { + return composition(layout, make_tuple(_, _, make_layout(stages))); +} + +template +CUTE_DEVICE auto make_acc_into_op(Accumulator const& acc, OperandLayout_TV const& operand_layout_tv) { + Tensor operand = make_fragment_like(convert_c_layout_to_a_layout(acc.layout(), shape<1>(operand_layout_tv))); + Tensor operand_as_acc = make_tensor(operand.data(), acc.layout()); + + cute::copy(acc, operand_as_acc); + + if constexpr (sizeof(Element) == 1) { + + // 00 11 22 33 00 11 22 33 acc layout + // 00 00 11 11 22 22 33 33 operand layout + // BB AA AA BB AA BB BB AA conflict-free exchange pattern + // 16-bit exchange; so process two at a time potentially + int tid = threadIdx.x % 4; + auto values_u32 = recast(operand); + + CUTE_UNROLL + for (int n = 0; n < size<1>(values_u32); n++) { + CUTE_UNROLL + for (int k = 0; k < size<2>(values_u32); k++) { + CUTE_UNROLL + for (int ii = 0; ii < 8; ii += 4) { + + uint32_t values_tmp_0 = values_u32(ii / 2 + 0, n, k); + uint32_t values_tmp_1 = values_u32(ii / 2 + 1, n, k); + + // step A: + // t 1 v 0 -> t 0 v 1 + // t 2 v 0 -> t 1 v 0 + // t 0 v 1 -> t 2 v 0 + // t 3 v 1 -> t 3 v 1 + + int v_to_send = tid == 1 || tid == 2 ? 0 : 1; + int v_to_recv = v_to_send; + int t_to_recv_from = (0x3021 >> (tid * 4)) & 0xF; + + uint32_t values_tmp_a = v_to_send == 0 ? values_tmp_0 : values_tmp_1; + + values_tmp_a = __shfl_sync(0xFFFFFFFF, values_tmp_a, t_to_recv_from, 4); + + // step B: + // t 0 v 0 -> t 0 v 0 + // t 3 v 0 -> t 1 v 1 + // t 1 v 1 -> t 2 v 1 + // t 2 v 1 -> t 3 v 0 + + v_to_send = 1 - v_to_send; + v_to_recv = 1 - v_to_recv; + t_to_recv_from = (0x2130 >> (tid * 4)) & 0xF; + + uint32_t values_tmp_b = v_to_send == 0 ? values_tmp_0 : values_tmp_1; + + values_tmp_b = __shfl_sync(0xFFFFFFFF, values_tmp_b, t_to_recv_from, 4); + + values_u32(ii / 2 + 0, n, k) = __byte_perm(values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x1054 : 0x5410); + values_u32(ii / 2 + 1, n, k) = __byte_perm(values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x3276 : 0x7632); + } + } + } + } + + return operand; +} + +} // namespace cutlass::fmha::collective diff --git a/examples/111_hopper_ssd/collective/sm90_ssd_epilogue.hpp b/examples/111_hopper_ssd/collective/sm90_ssd_epilogue.hpp new file mode 100644 index 00000000..9d2ef017 --- /dev/null +++ b/examples/111_hopper_ssd/collective/sm90_ssd_epilogue.hpp @@ -0,0 +1,838 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/epilogue/collective/collective_builder.hpp" +#include "cutlass/fast_math.h" + +namespace cutlass::ssd::collective { +using namespace cute; + +template< + class ElementAcc_, + class Element_, + class TileShape_, + class EpilogueTile_, + class SmemLayoutX_, + class SmemLayoutY_, + class SmemLayoutPartialY_, + class SmemLayoutP_, + class SmemLayoutZ_, + int StagesD_, + int StagesY_, + int StagesZ_, + bool HAS_D_, + bool D_HAS_HDIM_, + bool HAS_Z_> +struct SsdEpilogue { + + using TileShape = TileShape_; + using ElementAcc = ElementAcc_; + using Element = Element_; + using ElementY = Element_; + using ElementP = Element_; + using ElementD = Element_; + using ElementX = Element_; + using ElementZ = Element_; + using EpilogueTile = EpilogueTile_; + using SmemLayoutX = SmemLayoutX_; + using SmemLayoutY = SmemLayoutY_; + using SmemLayoutPartialY = SmemLayoutPartialY_; + using SmemLayoutP = SmemLayoutP_; + using SmemLayoutZ = SmemLayoutZ_; + + static constexpr int StagesD = StagesD_; + static constexpr int StagesY = StagesY_; + static constexpr int StagesZ = StagesZ_; + static constexpr bool HAS_D = HAS_D_; + static constexpr bool D_HAS_HDIM = D_HAS_HDIM_; + static constexpr bool HAS_Z = HAS_Z_; + + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + static constexpr auto L = get<0>(TileShape{}); + static constexpr auto D = get<1>(TileShape{}); + static constexpr auto N = get<2>(TileShape{}); + + constexpr static size_t SmemAlignmentY = cutlass::detail::alignment_for_swizzle(SmemLayoutY{}); + + struct CollectiveStorage { + alignas(SmemAlignmentY) ArrayEngine> smem_y_partial; + alignas(SmemAlignmentY) ArrayEngine> smem_y; + alignas(SmemAlignmentY) ArrayEngine> smem_z; + }; + + using EpiloadPipelineD = cutlass::PipelineTmaAsync; + using EpiloadPipelineZ = cutlass::PipelineTmaAsync; + static constexpr int kEpiloadDBytes = D_HAS_HDIM ? D * sizeof(ElementD) : 0; + static constexpr int kEpiloadZBytes = HAS_Z ? cosize_v * sizeof(ElementZ) : 0; + + // TMA pipeline for storing D + using StorePipeline = cutlass::PipelineTmaStore; + using StorePipelineState = cutlass::PipelineState; + + // TMA pipeline for storing P + using StorePPipeline = cutlass::PipelineTmaStore<1>; + using StorePPipelineState = cutlass::PipelineState<1>; + + using CooperatePipeline = cutlass::PipelineAsync; + using CooperatePipelineState = cutlass::PipelineState; + + struct SharedStorage { + using TensorStorage = CollectiveStorage; + TensorStorage tensors; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + + using StrideY = cute::tuple<_1, int, int, int>; // (L,D,C,B) + using StrideP = cute::tuple; // (D,N,B) + using StrideZ = cute::tuple<_1, int, int, int>; // (L,D,C,B) + + using LayoutY = decltype(make_layout(make_shape(L, D, int32_t(0), int32_t(0)), + make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,D,C,B) + using LayoutP = decltype(make_layout(make_shape(D, N, int32_t(0)), make_stride(N, _1{}, D*N))); // (D,N,B) + + using LayoutD_2D = decltype(make_layout(make_shape(D, int32_t(0)), make_stride(_1{}, D))); // (D,EH) + using LayoutD_1D = decltype(make_layout(make_shape(_1{}, int32_t(0)), make_stride(_0{}, _1{}))); // (D,EH) + using LayoutD = cute::conditional_t< + D_HAS_HDIM, + LayoutD_2D, + LayoutD_1D + >; + using LayoutZ = decltype(make_layout(make_shape(L, D, int32_t(0), int32_t(0)), + make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,D,C,B) + + struct Arguments { + ElementY* ptr_Y{nullptr}; + ElementP* ptr_P{nullptr}; + const ElementD* ptr_D{nullptr}; + const ElementZ* ptr_Z{nullptr}; + LayoutY layout_Y{}; + LayoutP layout_P{}; + LayoutD layout_D{}; + LayoutZ layout_Z{}; + }; + + using CopyOpS2G = SM90_TMA_STORE; + using CopyOpG2S = SM90_TMA_LOAD; + + struct Params { + using TMA_Y = decltype(make_tma_atom( + CopyOpS2G{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)),LayoutY{}), + take<0,2>(SmemLayoutY{}), + EpilogueTile{})); + using TMA_P = decltype(make_tma_atom( + CopyOpS2G{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)),LayoutP{}), + take<0,2>(SmemLayoutP{}), + make_tile(shape<1>(TileShape{}), shape<2>(TileShape{})))); + using TMA_Z = decltype(make_tma_atom( + CopyOpG2S{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)),LayoutZ{}), + take<0,2>(SmemLayoutZ{}), + make_tile(shape<0>(TileShape{}), shape<1>(TileShape{})))); + using TensorD = decltype(make_tensor( + make_gmem_ptr(static_cast(nullptr)), + LayoutD{})); + + TMA_Y tma_store_y; + TMA_P tma_store_p; + TMA_Z tma_load_z; + TensorD tensor_d; + }; + + template + static Params to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace = nullptr) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + auto tensor_y = make_tensor(make_gmem_ptr(args.ptr_Y), args.layout_Y); + auto tensor_p = make_tensor(make_gmem_ptr(args.ptr_P), args.layout_P); + auto tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), args.layout_D); + auto tensor_z = make_tensor(make_gmem_ptr(args.ptr_Z), args.layout_Z); + + auto tma_store_y = make_tma_atom( + CopyOpS2G{}, + tensor_y, + take<0,2>(SmemLayoutY{}), + EpilogueTile{}); + auto tma_store_p = make_tma_atom( + CopyOpS2G{}, + tensor_p, + take<0,2>(SmemLayoutP{}), + make_tile(shape<1>(TileShape{}), shape<2>(TileShape{}))); + auto tma_load_z = make_tma_atom( + CopyOpG2S{}, + tensor_z, + take<0,2>(SmemLayoutZ{}), + make_tile(shape<0>(TileShape{}), shape<1>(TileShape{}))); + return Params{ + tma_store_y, + tma_store_p, + tma_load_z, + tensor_d + }; + } + + template< + class Params, class ProblemShape, + class EpiloadPipeline, class PipelineState, + class TensorStorage + > + CUTLASS_DEVICE + void load_d( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + EpiloadPipeline& pipeline, PipelineState& pipeline_state, + TensorStorage& shared_tensors) { + + if constexpr (D_HAS_HDIM) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + + auto& gD = params.tensor_d; + + ElementD* ptr_d = shared_tensors.smem_d.data(); + auto smem_layout = make_layout(make_shape(get<1>(TileShape{}), Int{})); // (D,) + auto sD = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_d), smem_layout)); + + auto bulk_atom = Copy_Atom{}; + + int write_stage = pipeline_state.index(); + + // LOCK pipeline_state for _writing_ + pipeline.producer_acquire(pipeline_state); + + using BarrierType = typename EpiloadPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_state); + copy(bulk_atom.with(*tma_barrier), gD(_,blk_coord), sD(_,write_stage)); + + // Advance pipeline_state + ++pipeline_state; + } + } + + } + + template< + class EpiloadPipeline, class PipelineState + > + CUTLASS_DEVICE void + load_d_tail(EpiloadPipeline pipeline, PipelineState pipeline_state) { + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + pipeline.producer_tail(pipeline_state); + } + } + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_z_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mZ_mkl = params.tma_load_z.get_tma_tensor(make_shape(L,D,C,EH*B)); + Tensor gZ_mkl = local_tile(mZ_mkl, TileShape{}, make_coord(_,_,_), Step<_1,_1,X>{}); + + return make_tuple(gZ_mkl); + } + + template< + class Params, class ProblemShape, + class EpiLoadPipeline, class PipelineState, + class GTensor, + class TensorStorage + > + CUTLASS_DEVICE + void load_z( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + EpiLoadPipeline& pipeline, PipelineState& pipeline_state, + cute::tuple const& load_inputs, + TensorStorage& shared_tensors) { + + if constexpr (HAS_Z) { + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + Tensor sZ_ = make_tensor(make_smem_ptr(shared_tensors.smem_z.begin()), SmemLayoutZ{}); + Tensor sZ = as_position_independent_swizzle_tensor(sZ_); + + // + // Prepare the TMA loads for Z + // + + Tensor gZ_mkl = get<0>(load_inputs); + + // Partition the inputs based on the current block coordinates. + Tensor gZ = gZ_mkl(_,_,_0{},_0{},_,blk_coord); + + auto oprands = tma_partition(params.tma_load_z, Int<0>{}, Layout<_1>{}, + group_modes<0,2>(sZ), group_modes<0,2>(gZ)); // (TMA,k) and (TMA,PIPE) + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + auto tZgZ = get<0>(oprands); + auto tZsZ = get<1>(oprands); + + // Disable multicast + uint16_t mcast_mask_z = 0; + + auto chunk = get<3>(problem_size); + // Mainloop + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) { + int write_stage = pipeline_state.index(); + + // LOCK pipeline_state for _writing_ + pipeline.producer_acquire(pipeline_state); + + using BarrierType = typename EpiLoadPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_state); + + copy(params.tma_load_z.with(*tma_barrier, mcast_mask_z), tZgZ(_,chunk_idx), tZsZ(_,write_stage)); + + // Advance pipeline_state + ++pipeline_state; + } + } + } + } + + template< + class EpiLoadPipeline, class PipelineState + > + CUTLASS_DEVICE void + load_z_tail(EpiLoadPipeline pipeline, PipelineState pipeline_state) { + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(pipeline_state); + } + } + + template< + class Params, class ProblemShape, + class CooperatePipeline, class CooperatePipelineState, + class TensorIntra, + class TiledMma, + class TensorStorage + > + CUTLASS_DEVICE + auto store_intra( + int& chunk, int const& blk_coord, Params const& params, ProblemShape const& problem_size, + CooperatePipeline& cooperate_pipeline, CooperatePipelineState& cooperate_pipe_producer_state, + TensorIntra& tIntra, + TiledMma tiled_mma, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mY_mn = params.tma_store_y.get_tma_tensor(make_shape(L,D,C,B*EH)); + Tensor gY_mn = local_tile(mY_mn, take<0,2>(TileShape{}), make_coord(_,_,_))(_,_,_0{},_0{},chunk,blk_coord); + // Apply epilogue subtiling + Tensor gY_epi = flat_divide(gY_mn, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + auto epi_tile_m = size<0>(EpilogueTile{}); + auto epi_tile_n = size<1>(EpilogueTile{}); + auto partial_m = Int<128>{}; + auto partial_n = Int{}; + + auto ptr_sY = shared_tensors.smem_y_partial.begin(); + Tensor sY_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sY), + tile_to_shape(UMMA::Layout_K_SW64_Atom{}, make_shape(partial_m, partial_n, Int<2>{})))); + using CopyAtomC = Copy_Atom; + TiledCopy tiled_copy_C_atom = make_tiled_copy_C_atom(CopyAtomC{}, tiled_mma); + using CopyOpR2S = SM90_U16x8_STSM_T; + TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rIntra = thread_r2s.retile_S(tIntra); // ((R2S,R2S_V),MMA_M,MMA_N) + + // Hard code + static constexpr int FragmentSize = 4; + auto tRS_rY = make_tensor(shape(sY_epi(thread_idx,_,_0{}))); + Tensor tRS_rIntra_frg = recast>(tRS_rIntra); + Tensor tRS_sY_frg = recast>(sY_epi(thread_idx,_,_)); + Tensor tRS_rY_frg = recast>(tRS_rY); + + auto mma_tile_m = size<0>(TileShape{}) / size<1>(tRS_rIntra); + auto mma_tile_n = size<1>(TileShape{}) / size<2>(tRS_rIntra); + + // For each output tile + CUTLASS_PRAGMA_UNROLL + for (int epi_n = 0; epi_n < size<3>(gY_epi); ++epi_n) { + CUTLASS_PRAGMA_UNROLL + for (int epi_m = 0; epi_m < size<2>(gY_epi); ++epi_m) { + + int mma_m = epi_m; + int mma_n = (epi_n * size<1>(EpilogueTile{})) / mma_tile_n; + Tensor tRS_rIntra_frg_mn = tRS_rIntra_frg(_,mma_m,mma_n); + + // Vectorized fragment loop with visitor callback entry point + // Epilogue op + int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n); + int r2s_v = epi_n_in_mma * size(tRS_rY_frg); + + CUTLASS_PRAGMA_UNROLL + for (int epi_v = 0; epi_v < size(tRS_rY_frg); ++epi_v) { + tRS_rY_frg(epi_v) = tRS_rIntra_frg_mn(r2s_v + epi_v); + } + + cooperate_pipeline.producer_acquire(cooperate_pipe_producer_state); + + copy(tRS_rY_frg, tRS_sY_frg(_,cooperate_pipe_producer_state.index())); + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to Cooperate warps + + cooperate_pipeline.producer_commit(cooperate_pipe_producer_state); + ++cooperate_pipe_producer_state; + } + } + } + + template< + class Params, + class EpiloadPipeline, class EpiloadPipelineState, + class TiledMma, + class TensorStorage + > + CUTLASS_DEVICE + auto update_d( + int const& blk_coord, Params params, + bool is_first_iteration, + EpiloadPipeline& epi_load_pipeline, EpiloadPipelineState& epi_load_pipe_consumer_state, + TiledMma tiled_mma, + TensorStorage& shared_tensors + ) { + + int thread_idx = int(threadIdx.x % 128); + + int read_stage = epi_load_pipe_consumer_state.index(); + // Load delta for epilogue + auto row_layout = make_layout(make_shape(get<0>(TileShape{}), get<1>(TileShape{}), Int{}), make_stride(_0{}, _1{}, get<1>(TileShape{}))); + auto c_tv_layout = typename TiledMma::LayoutC_TV{}; + auto tile_mn = make_shape(tile_size<0>(tiled_mma), tile_size<1>(tiled_mma)); + auto d_layout = zipped_divide(row_layout, tile_mn); + auto d_tv_layout = composition(d_layout, make_tuple(c_tv_layout,_)); + + auto sD = make_tensor(make_smem_ptr(shared_tensors.smem_d.data()), d_tv_layout)(make_coord(thread_idx,_), make_coord(_,_,read_stage)); + auto tSR_sD = as_position_independent_swizzle_tensor(sD); + auto tSR_rD = make_tensor(shape(sD)); + auto tD = make_tensor(shape(sD)); + + if constexpr (D_HAS_HDIM) { + if (is_first_iteration) { + epi_load_pipeline.consumer_wait(epi_load_pipe_consumer_state); + } + copy(tSR_sD, tSR_rD); + type_convert(tSR_rD, tD); + } + else if constexpr (HAS_D) { + auto& gD = params.tensor_d; + auto value = static_cast(gD(_0{}, blk_coord)); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tSR_rD); ++ii) { + tD(ii) = value; + } + } + + return make_tuple(tD); + } + + template< + class Params, class ProblemShape, + class StorePipeline, class StorePipelineState, + class CooperatePipeline, class CooperatePipelineState, + class MainloopPipelineX, class PipelineStateX, + class EpiloadPipelineZ, class PipelineStateZ, + class TensorInter, class TensorDeltaA, class TensorD, + class TiledMma, + class TensorStorage, class TensorStorageX + > + CUTLASS_DEVICE + auto store( + int& chunk, int const& blk_coord, Params const& params, ProblemShape const& problem_size, + StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state, + CooperatePipeline& cooperate_pipeline, CooperatePipelineState& cooperate_pipe_consumer_state, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_state_x, + EpiloadPipelineZ& pipeline_z, PipelineStateZ& pipeline_state_z, + TensorInter& tInter, TensorDeltaA& tDeltaA, TensorD& tD, + TiledMma tiled_mma, + TensorStorage& shared_tensors, TensorStorageX& shared_tensors_x) { + + int thread_idx = int(threadIdx.x % 128); + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mY_mn = params.tma_store_y.get_tma_tensor(make_shape(L,D,C,B*EH)); + Tensor gY_mn = local_tile(mY_mn, take<0,2>(TileShape{}), make_coord(_,_,_))(_,_,_0{},_0{},chunk,blk_coord); + + // Apply epilogue subtiling + Tensor gY_epi = flat_divide(gY_mn, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Construct the corresponding pipelined smem tensors + auto ptr_sY = shared_tensors.smem_y.begin(); + Tensor sY_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sY), SmemLayoutY{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + auto ptr_sX = shared_tensors_x.smem_x.data(); + Tensor sX_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sX), SmemLayoutX{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + auto ptr_sZ = shared_tensors.smem_z.begin(); + Tensor sZ_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sZ), SmemLayoutZ{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + auto epi_tile_m = size<0>(EpilogueTile{}); + auto epi_tile_n = size<1>(EpilogueTile{}); + auto partial_m = Int<128>{}; + auto partial_n = Int{}; + + auto ptr_sY_partial = shared_tensors.smem_y_partial.begin(); + Tensor sY_epi_partial = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sY_partial), + tile_to_shape(UMMA::Layout_K_SW64_Atom{}, make_shape(partial_m, partial_n, Int<2>{})))); + + using CopyAtomC = Copy_Atom; + TiledCopy tiled_copy_C_atom = make_tiled_copy_C_atom(CopyAtomC{}, tiled_mma); + + using CopyOpR2S = SM90_U16x8_STSM_T; + TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + + Tensor tRS_rInter = thread_r2s.retile_S(tInter); // ((R2S,R2S_V),MMA_M,MMA_N) + Tensor tRS_rDeltaA = thread_r2s.retile_S(tDeltaA); // ((R2S,R2S_V),MMA_M,MMA_N) + Tensor tRS_rD = thread_r2s.retile_S(tD); // ((R2S,R2S_V),MMA_M,MMA_N) + Tensor tRS_sY = thread_r2s.partition_D(sY_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + + using CopyOpS2R = SM75_U16x8_LDSM_T; + TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sX = thread_s2r.partition_S(flat_divide(sX_epi, EpilogueTile{})); + Layout tSR_rX_layout = make_layout(take<0,3>(shape(thread_s2r.partition_D(flat_divide(sX_epi, EpilogueTile{}))))); + Tensor tSR_rX = make_tensor(tSR_rX_layout); + Tensor tRS_rX = make_tensor(tSR_rX_layout); + + using CopyOpS2R_Z = SM75_U16x8_LDSM_T; + TiledCopy tiled_s2r_z = make_tiled_copy_S(Copy_Atom{}, tiled_copy_C_atom); + ThrCopy thread_s2r_z = tiled_s2r_z.get_slice(thread_idx); + Tensor tSR_sZ = thread_s2r_z.partition_S(flat_divide(sZ_epi, EpilogueTile{})); + Layout tSR_rZ_layout = make_layout(take<0,3>(shape(thread_s2r_z.partition_D(flat_divide(sZ_epi, EpilogueTile{}))))); + Tensor tSR_rZ = make_tensor(tSR_rZ_layout); + Tensor tRS_rZ = make_tensor(tSR_rZ_layout); + + // Hard code + static constexpr int FragmentSize = 4; + // Allocate Y registers + Layout tRS_rY_layout = make_layout(take<0,3>(shape(thread_r2s.partition_S(sY_epi)))); + Tensor tRS_rY = make_tensor(tRS_rY_layout); // (R2S,R2S_M,R2S_N) + Tensor tRS_rY_frg = recast>(tRS_rY); + Tensor tRS_rX_frg = recast>(tRS_rX); + Tensor tRS_rZ_frg = recast>(tRS_rZ); + // Tensor tRS_rIntra_frg = recast>(tRS_rIntra); + Tensor tRS_rInter_frg = recast>(tRS_rInter); + Tensor tRS_rDeltaA_frg = recast>(tRS_rDeltaA); + Tensor tRS_rD_frg = recast>(tRS_rD); + Tensor tRS_rCompute = make_tensor(tRS_rY_layout); // (R2S,R2S_M,R2S_N) + Tensor tRS_rCompute_frg = recast>(tRS_rCompute); + + auto tSR_rY = make_tensor(shape(sY_epi_partial(thread_idx,_,_0{}))); + Tensor tSR_sY_frg = recast>(sY_epi_partial(thread_idx,_,_)); + Tensor tSR_rY_frg = recast>(tSR_rY); + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + auto oprands = tma_partition(params.tma_store_y, Int<0>{}, Layout<_1>{}, + group_modes<0,2>(sY_epi), group_modes<0,2>(gY_epi)); // (TMA,k) and (TMA,PIPE) + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + auto bSG_gY = get<0>(oprands); + auto bSG_sY = get<1>(oprands); + + auto mma_tile_m = size<0>(TileShape{}) / size<1>(tRS_rInter); + auto mma_tile_n = size<1>(TileShape{}) / size<2>(tRS_rInter); + +#if 0 + if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0) && chunk == 0) { + print("tRS_rInter : ");print(tRS_rInter);print("\n"); + print("tRS_sY : ");print(tRS_sY);print("\n"); + print("bSG_sY : ");print(bSG_sY);print("\n"); + print("bSG_gY : ");print(bSG_gY);print("\n"); + + print("tRS_rY_frg : ");print(tRS_rY_frg);print("\n"); + print("tSR_rY_frg : ");print(tSR_rY_frg);print("\n"); + print("tRS_rInter_frg : ");print(tRS_rInter_frg);print("\n"); + print("tRS_rCompute_frg : ");print(tRS_rCompute_frg);print("\n"); + } +#endif + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0; + + int epi_m_prev = 0, epi_n_prev = 0; + + auto tma_store_fn = [&] (int epi_m, int epi_n) { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_y, bSG_sY(_,store_pipe_producer_state.index()), bSG_gY(_,epi_m,epi_n)); + } + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + }; + + static constexpr bool DelayTmaStore = true; + int subtile_idx = -1; + + if constexpr (HAS_Z) { + pipeline_z.consumer_wait(pipeline_state_z); + } + + // For each output tile + CUTLASS_PRAGMA_UNROLL + for (int epi_n = 0; epi_n < size<3>(gY_epi); ++epi_n) { + CUTLASS_PRAGMA_UNROLL + for (int epi_m = 0; epi_m < size<2>(gY_epi); ++epi_m) { + bool is_first_iteration = epi_m == 0 && epi_n == 0; + bool is_last_iteration = epi_m == size<2>(gY_epi)-1 && epi_n == size<3>(gY_epi)-1; + + if (subtile_idx != -1 && (epi_n * static_cast(size<2>(gY_epi)) + epi_m) != subtile_idx) { + continue; + } + + if constexpr (HAS_D) { + // load x + copy(tiled_s2r, tSR_sX(_,_,_,epi_m,epi_n,pipeline_state_x.index()), tSR_rX); + type_convert(tSR_rX, tRS_rX); + } + + if constexpr (HAS_Z) { + // load x + copy(tiled_s2r_z, tSR_sZ(_,_,_,epi_m,epi_n,pipeline_state_z.index()), tSR_rZ); + type_convert(tSR_rZ, tRS_rZ); + // SiLu + cutlass::epilogue::thread::SiLu op; + for (int ii = 0; ii < size(tRS_rZ); ++ii) { + tRS_rZ(ii) = op(tRS_rZ(ii)); + } + } + + cooperate_pipeline.consumer_wait(cooperate_pipe_consumer_state); + + copy(tSR_sY_frg(_,cooperate_pipe_consumer_state.index()), tSR_rY_frg); + + cooperate_pipeline.consumer_release(cooperate_pipe_consumer_state); + ++cooperate_pipe_consumer_state; + + cutlass::arch::fence_view_async_shared(); + synchronize(); + + int mma_m = epi_m; + int mma_n = (epi_n * size<1>(EpilogueTile{})) / mma_tile_n; + // Tensor tRS_rIntra_frg_mn = tRS_rIntra_frg(_,mma_m,mma_n); + Tensor tRS_rInter_frg_mn = tRS_rInter_frg(_,mma_m,mma_n); + Tensor tRS_rDeltaA_frg_mn = tRS_rDeltaA_frg(_,mma_m,mma_n); + Tensor tRS_rD_frg_mn = tRS_rD_frg(_,mma_m,mma_n); + + // Vectorized fragment loop with visitor callback entry point + // Epilogue op + int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n); + int r2s_v = epi_n_in_mma * size(tRS_rCompute_frg); + CUTLASS_PRAGMA_UNROLL + for (int epi_v = 0; epi_v < size(tRS_rCompute_frg); ++epi_v) { + tRS_rCompute_frg(epi_v) = tRS_rDeltaA_frg_mn(r2s_v + epi_v) * tRS_rInter_frg_mn(r2s_v + epi_v) + tSR_rY_frg(epi_v); + if constexpr (HAS_D) { + tRS_rCompute_frg(epi_v) = tRS_rD_frg_mn(r2s_v + epi_v) * tRS_rX_frg(epi_v) + tRS_rCompute_frg(epi_v); + } + if constexpr (HAS_Z) { + tRS_rCompute_frg(epi_v) = tRS_rCompute_frg(epi_v) * tRS_rZ_frg(epi_v); + } + } + + // The latest we can delay the TMA store is right before the smem store of the next iteration + // since the current TMA store needs to be committed before we can acquire the next smem buffer + if constexpr (DelayTmaStore) { + // Issue TMA stores for the previous subtile + if (not is_first_iteration and subtile_idx == -1) { + tma_store_fn(epi_m_prev, epi_n_prev); + } + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tRS_rY_frg); ++i) { + tRS_rY_frg(i) = cutlass::NumericArrayConverter{}(tRS_rCompute_frg(i)); + } + + copy(tiled_r2s, tRS_rY, tRS_sY(_,_,_,store_pipe_producer_state.index())); + + if constexpr (not DelayTmaStore) { + // Issue TMA stores for this subtile + tma_store_fn(epi_m, epi_n); + } + + } // for epi_m + } // for epi_n + + pipeline_x.consumer_release(pipeline_state_x); + ++pipeline_state_x; + pipeline_z.consumer_release(pipeline_state_z); + ++pipeline_state_z; + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the last subtile + tma_store_fn(epi_m_prev, epi_n_prev); + } + + } + + template< + class Params, class ProblemShape, + class StorePipeline, class StorePipelineState, + class TensorStorage + > + CUTLASS_DEVICE + auto store_p( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mP_mn = params.tma_store_p.get_tma_tensor(make_shape(D,N,B*EH)); + Tensor gP_mn = local_tile(mP_mn, take<1,3>(TileShape{}), make_coord(_,_,blk_coord)); + auto gP_epi = gP_mn(_,_,_0{},_0{}); + + // Construct the corresponding pipelined smem tensors + auto ptr_sP = shared_tensors.smem_p.begin(); + Tensor sP_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sP), SmemLayoutP{})); // (EPI_TILE_M,EPI_TILE_N) + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + auto oprands = tma_partition(params.tma_store_p, Int<0>{}, Layout<_1>{}, + group_modes<0,2>(sP_epi), group_modes<0,2>(gP_epi)); // (TMA,k) and (TMA,PIPE) + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + auto bSG_gY = get<0>(oprands); + auto bSG_sY = get<1>(oprands); + +#if 0 + if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("sP_epi : ");print(sP_epi);print("\n"); + print("gP_epi : ");print(gP_epi);print("\n"); + print("bSG_sY : ");print(bSG_sY);print("\n"); + print("bSG_gY : ");print(bSG_gY);print("\n"); + } +#endif + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + // Use the reserved named barrier `streamkbarrier` since this kernel doesn't support streamk + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::StreamkBarrier0); }; + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0; + + auto tma_store_fn = [&] () { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_p, bSG_sY, bSG_gY); + } + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + // don't need pipeline + synchronize(); + }; + + tma_store_fn(); + } + + template< + class ElementSrc, class ElementDst, + class TensorSrc, class TensorDst + > + CUTLASS_DEVICE + auto type_convert( + TensorSrc& tS, + TensorDst& tD) { + + static constexpr int FragmentSize = 2; + NumericArrayConverter converter; + + auto tS_frg = recast>(tS); + auto tD_frg = recast>(tD); + + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tS_frg); ++ii) { + tD_frg(ii) = converter(tS_frg(ii)); + } + } +}; + +} // namespace cutlass::fmha::collective + diff --git a/examples/111_hopper_ssd/collective/sm90_ssd_gemm_tma_warpspecialized.hpp b/examples/111_hopper_ssd/collective/sm90_ssd_gemm_tma_warpspecialized.hpp new file mode 100644 index 00000000..81b37f17 --- /dev/null +++ b/examples/111_hopper_ssd/collective/sm90_ssd_gemm_tma_warpspecialized.hpp @@ -0,0 +1,1146 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/collective/collective_builder.hpp" +#include "../collective/common.hpp" + +namespace cutlass::ssd::collective { + +using namespace cute; + +template< + class Element_, + class ElementDA_, + class ElementAcc_, + class ElementD_, + class TileShape_, + int Stages_ +> +struct SsdMainloopTmaWarpSpecialized { + + using Element = Element_; + using ElementDA = ElementDA_; + using ElementAcc = ElementAcc_; + using TileShape = TileShape_; // (L,D,N) + + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + static constexpr auto L = get<0>(TileShape{}); + static constexpr auto D = get<1>(TileShape{}); + static constexpr auto N = get<2>(TileShape{}); + + using TileShapeIntraBMM1 = decltype(make_shape(L, L, N)); // (L,L,N) + using TileShapeIntraBMM2 = decltype(make_shape(L, D, N)); // (L,D,L) + using TileShapeInterBMM1 = decltype(make_shape(N, D, L)); // (N,D,L) + using TileShapeInterBMM2 = decltype(make_shape(L, D, N)); // (L,D,N) + + using ClusterShape = Shape<_1, _1, _1>; + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using Stages = cutlass::gemm::collective::StageCount; + // IntraMma1 NT + using CollectiveIntraMma1 = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, + Element, cutlass::layout::ColumnMajor, Alignment, + Element, cutlass::layout::RowMajor, Alignment, + ElementAcc, + TileShapeIntraBMM1, ClusterShape, Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + // IntraMma2 TN + using CollectiveIntraMma2 = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, + Element, cutlass::layout::RowMajor, Alignment, + Element, cutlass::layout::ColumnMajor, Alignment, + ElementAcc, + TileShapeIntraBMM2, ClusterShape, Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + // InterMma1 TN + using CollectiveInterMma1 = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, + Element, cutlass::layout::RowMajor, Alignment, + Element, cutlass::layout::ColumnMajor, Alignment, + ElementAcc, + TileShapeInterBMM1, ClusterShape, Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + // InterMma2 NN + using CollectiveInterMma2 = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, + Element, cutlass::layout::ColumnMajor, Alignment, + Element, cutlass::layout::ColumnMajor, Alignment, + ElementAcc, + TileShapeInterBMM2, ClusterShape, Stages, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaIntra1 = typename CollectiveIntraMma1::TiledMma; + using TiledMmaIntra2 = decltype(convert_to_gmma_rs(typename CollectiveIntraMma2::TiledMma{})); + using TiledMmaInter1 = decltype(convert_to_gmma_rs(typename CollectiveInterMma1::TiledMma{})); + using TiledMmaInter2 = typename CollectiveInterMma2::TiledMma; + + using SmemLayoutX = typename CollectiveIntraMma2::SmemLayoutB; + using SmemLayoutC = typename CollectiveIntraMma1::SmemLayoutA; + using SmemLayoutB = typename CollectiveIntraMma1::SmemLayoutB; + + // Use 1 stage buffer since it is immediate result by the previous gemm + using SmemLayoutP_Ref = typename CollectiveInterMma2::SmemLayoutB; + using SmemLayoutP = decltype(take<0,2>(SmemLayoutP_Ref{})); + + using MainloopPipelineX = cutlass::PipelineTmaAsync; + using MainloopPipelineB = cutlass::PipelineTmaAsync; + using MainloopPipelineC = cutlass::PipelineTmaAsync; + using MainloopPipelineDelta = cutlass::PipelineTmaAsync; + + static constexpr int kXLoadBytes = size(SmemLayoutX{}(_,_,_0{})) * sizeof(Element); + static constexpr int kBLoadBytes = size(SmemLayoutB{}(_,_,_0{})) * sizeof(Element); + static constexpr int kCLoadBytes = size(SmemLayoutC{}(_,_,_0{})) * sizeof(Element); + static constexpr int kDeltaLoadBytes = get<0>(TileShape{}) * sizeof(Element); + static constexpr int kDeltaALoadBytes = get<0>(TileShape{}) * sizeof(ElementDA); + + struct SharedStorage : cute::aligned_struct<128, _0> { + cute::array_aligned> smem_x; + cute::array_aligned> smem_b; + cute::array_aligned> smem_c; + cute::array_aligned> smem_p; + cute::array_aligned(TileShape{}) * Stages::value> smem_delta; + cute::array_aligned(TileShape{}) * Stages::value> smem_delta_a; + cute::array_aligned(TileShape{}) * Stages::value> smem_d; + }; + + using LayoutX = decltype(make_layout(make_shape(D, L, int32_t(0), int32_t(0)), make_stride(int32_t(0), _1{}, L, int32_t(0)))); // (D,L,C,B) + using LayoutB = decltype(make_layout(make_shape(L, N, int32_t(0), int32_t(0)), make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,N,C,B) + using LayoutC = decltype(make_layout(make_shape(L, N, int32_t(0), int32_t(0)), make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,N,C,B) + using LayoutDelta = decltype(make_layout(make_shape(L, int32_t(0), int32_t(0)), make_stride(_1{}, L, int32_t(0)))); // (L,C,B) + + // Not support Mcast + using GmemTiledCopyX = cute::SM90_TMA_LOAD; + using GmemTiledCopyB = cute::SM90_TMA_LOAD; + using GmemTiledCopyC = cute::SM90_TMA_LOAD; + + struct Arguments { + const Element* ptr_X{nullptr}; + const ElementDA* ptr_DeltaA{nullptr}; + const Element* ptr_Delta{nullptr}; + const Element* ptr_B{nullptr}; + const Element* ptr_C{nullptr}; + LayoutX layout_X{}; + LayoutB layout_B{}; + LayoutC layout_C{}; + LayoutDelta layout_Delta{}; + }; + + struct Params { + using TMA_X = decltype(make_tma_copy( + GmemTiledCopyX{}, + make_tensor(static_cast(nullptr), LayoutX{}), + SmemLayoutX{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShapeIntraBMM2{}), shape<2>(TileShapeIntraBMM2{})), + _1{})); + using TMA_B = decltype(make_tma_copy( + GmemTiledCopyB{}, + make_tensor(static_cast(nullptr), LayoutB{}), + SmemLayoutB{}(_,_,cute::Int<0>{}), + make_shape(shape<0>(TileShapeIntraBMM1{}), shape<2>(TileShapeIntraBMM1{})), + _1{})); + using TMA_C = decltype(make_tma_copy( + GmemTiledCopyC{}, + make_tensor(static_cast(nullptr), LayoutC{}), + SmemLayoutC{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShapeIntraBMM1{}), shape<2>(TileShapeIntraBMM1{})), + _1{})); + using TensorDelta = decltype(make_tensor( + make_gmem_ptr(static_cast(nullptr)), + LayoutDelta{})); + using TensorDeltaA = decltype(make_tensor( + make_gmem_ptr(static_cast(nullptr)), + LayoutDelta{})); + + TMA_X tma_load_x; + TMA_B tma_load_b; + TMA_C tma_load_c; + TensorDelta tensor_delta; + TensorDeltaA tensor_delta_a; + }; + + template + static bool can_implement(ProblemShape const& problem_size, Arguments const& args) { + return true; + } + + template + static Params to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor tensor_x = make_tensor(make_gmem_ptr(args.ptr_X), args.layout_X); + Tensor tensor_b = make_tensor(make_gmem_ptr(args.ptr_B), args.layout_B); + Tensor tensor_c = make_tensor(make_gmem_ptr(args.ptr_C), args.layout_C); + Tensor tensor_delta = make_tensor(make_gmem_ptr(args.ptr_Delta), args.layout_Delta); + Tensor tensor_delta_a = make_tensor(make_gmem_ptr(args.ptr_DeltaA), args.layout_Delta); + +#if 0 + print("tensor_x : ");print(tensor_x);print("\n"); + print("smem_X : ");print(SmemLayoutX{});print("\n"); + print("tensor_b : ");print(tensor_b);print("\n"); + print("smem_b : ");print(SmemLayoutB{});print("\n"); + print("tensor_c : ");print(tensor_c);print("\n"); + print("smem_c : ");print(SmemLayoutC{});print("\n"); + print("tensor_delta : ");print(tensor_delta);print("\n"); + print("tensor_delta_a: ");print(tensor_delta_a);print("\n"); + print("tile : ");print(TileShape{});print("\n"); + print("tileIntra1 : ");print(TileShapeIntraBMM1{});print("\n"); + print("tileIntra2 : ");print(TileShapeIntraBMM2{});print("\n"); + print("tileInter1 : ");print(TileShapeInterBMM1{});print("\n"); + print("tileInter2 : ");print(TileShapeInterBMM2{});print("\n"); +#endif + + typename Params::TMA_X tma_load_x = make_tma_copy( + GmemTiledCopyX{}, + tensor_x, + SmemLayoutX{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShapeIntraBMM2{}), shape<2>(TileShapeIntraBMM2{})), + _1{}); + typename Params::TMA_B tma_load_b = make_tma_copy( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,cute::Int<0>{}), + make_shape(shape<0>(TileShapeIntraBMM1{}), shape<2>(TileShapeIntraBMM1{})), + _1{}); + typename Params::TMA_C tma_load_c = make_tma_copy( + GmemTiledCopyC{}, + tensor_c, + SmemLayoutC{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShapeIntraBMM1{}), shape<2>(TileShapeIntraBMM1{})), + _1{}); + + return Params{ + tma_load_x, + tma_load_b, + tma_load_c, + tensor_delta, + tensor_delta_a + }; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { + cute::prefetch_tma_descriptor(params.tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_c.get_tma_descriptor()); + } + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_x_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mX_mkl = params.tma_load_x.get_tma_tensor(make_shape(D,L,C,EH*B)); + Tensor gX_mkl = local_tile(mX_mkl, TileShapeIntraBMM2{}, make_coord(_,_,_), Step{}); // (BLK_M,BLK_K,m,k,l) + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("mX_mkl : ");print(mX_mkl);print("\n"); + print("gX_mkl : ");print(gX_mkl);print("\n"); + } +#endif + return make_tuple(gX_mkl); + } + + template< + class Params, class ProblemShape, + class MainloopPipeline, class PipelineState, + class GTensor, + class TensorStorage + > + CUTLASS_DEVICE + void load_x( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + MainloopPipeline& pipeline, PipelineState& pipeline_state, + cute::tuple const& load_inputs, + TensorStorage& shared_tensors) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + Tensor sX_ = make_tensor(make_smem_ptr(shared_tensors.smem_x.data()), SmemLayoutX{}); // (BLK_M,BLK_K,PIPE) + Tensor sX = as_position_independent_swizzle_tensor(sX_); // (BLK_M,BLK_K,PIPE) + + // + // Prepare the TMA loads for X + // + + Tensor gX_mkl = get<0>(load_inputs); + + // Partition the inputs based on the current block coordinates. + Tensor gX = gX_mkl(_,_,_0{},_0{},_,blk_coord); // (BLK_M,BLK_K,k) + + // no multicast + auto tma_load_x = params.tma_load_x.get_slice(_0{}); + // Applies the mapping from block_tma_a + Tensor tXgX = tma_load_x.partition_S(gX); // (TMA,TMA_M,TMA_K,k) + Tensor tXsX = tma_load_x.partition_D(sX); // (TMA,TMA_M,TMA_K,PIPE) + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("tXgX : ");print(tXgX);print("\n"); + print("tXsX : ");print(tXsX);print("\n"); + } +#endif + // Disable multicast + uint16_t mcast_mask_x = 0; + + auto chunk = get<3>(problem_size); + // Mainloop + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) { + int write_stage = pipeline_state.index(); + + // LOCK pipeline_state for _writing_ + pipeline.producer_acquire(pipeline_state); + + // + // Copy gmem to smem for chunk + // + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_state); + + copy(params.tma_load_x.with(*tma_barrier, mcast_mask_x), tXgX(_,_,_,chunk_idx), tXsX(_,_,_,write_stage)); + + // Advance pipeline_state + ++pipeline_state; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + template< + class MainloopPipeline, class PipelineState + > + CUTLASS_DEVICE void + load_x_tail(MainloopPipeline pipeline, PipelineState pipeline_state) { + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(pipeline_state); + } + } + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_b_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mB_mkl = params.tma_load_b.get_tma_tensor(make_shape(L,N,C,G*B)); + Tensor gB_mkl = local_tile(mB_mkl, TileShapeIntraBMM1{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("mB_mkl : ");print(mB_mkl);print("\n"); + print("gB_mkl : ");print(gB_mkl);print("\n"); + } +#endif + return make_tuple(gB_mkl); + } + + template< + class Params, class ProblemShape, + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineC, class PipelineStateC, + class GTensorB, class GTensorC, + class TensorStorage + > + CUTLASS_DEVICE + void load_b_c( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + MainloopPipelineB& pipeline_b, PipelineStateB& pipeline_state_b, + MainloopPipelineC& pipeline_c, PipelineStateC& pipeline_state_c, + cute::tuple const& load_inputs_b, + cute::tuple const& load_inputs_c, + TensorStorage& shared_tensors) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_b.data()), SmemLayoutB{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE) + Tensor sC_ = make_tensor(make_smem_ptr(shared_tensors.smem_c.data()), SmemLayoutC{}); // (BLK_M,BLK_K,PIPE) + Tensor sC = as_position_independent_swizzle_tensor(sC_); // (BLK_M,BLK_K,PIPE) + + // + // Prepare the TMA loads for B + // + + Tensor gB_mkl = get<0>(load_inputs_b); + Tensor gB = gB_mkl(_,_,_0{},_0{},_,blk_coord); // (BLK_M,BLK_K,k) + // no multicast + auto tma_load_b = params.tma_load_b.get_slice(_0{}); + // Applies the mapping from block_tma_a + Tensor tBgB = tma_load_b.partition_S(gB); // (TMA,TMA_M,TMA_K,k) + Tensor tBsB = tma_load_b.partition_D(sB); // (TMA,TMA_M,TMA_K,PIPE) + + // + // Prepare the TMA loads for X + // + + Tensor gC_mkl = get<0>(load_inputs_c); + Tensor gC = gC_mkl(_,_,_0{},_0{},_,blk_coord); // (BLK_M,BLK_K,k) + // no multicast + auto tma_load_c = params.tma_load_c.get_slice(_0{}); + // Applies the mapping from block_tma_a + Tensor tCgC = tma_load_c.partition_S(gC); // (TMA,TMA_M,TMA_K,k) + Tensor tCsC = tma_load_c.partition_D(sC); // (TMA,TMA_M,TMA_K,PIPE) + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("tCgC : ");print(tCgC);print("\n"); + print("tCsC : ");print(tCsC);print("\n"); + print("tBgB : ");print(tBgB);print("\n"); + print("tBsB : ");print(tBsB);print("\n"); + } +#endif + // Disable multicast + uint16_t mcast_mask_b = 0; + uint16_t mcast_mask_c = 0; + + auto chunk = get<3>(problem_size); + // Mainloop + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) { + int write_stage_b = pipeline_state_b.index(); + int write_stage_c = pipeline_state_c.index(); + + // Load B + pipeline_b.producer_acquire(pipeline_state_b); + + using BarrierTypeB = typename MainloopPipelineB::ProducerBarrierType; + BarrierTypeB* tma_barrier_b = pipeline_b.producer_get_barrier(pipeline_state_b); + + copy(params.tma_load_b.with(*tma_barrier_b, mcast_mask_b), tBgB(_,_,_,chunk_idx), tBsB(_,_,_,write_stage_b)); + + // Advance pipeline_state + ++pipeline_state_b; + + // Load C + pipeline_c.producer_acquire(pipeline_state_c); + + using BarrierTypeC = typename MainloopPipelineC::ProducerBarrierType; + BarrierTypeC* tma_barrier_c = pipeline_c.producer_get_barrier(pipeline_state_c); + + copy(params.tma_load_c.with(*tma_barrier_c, mcast_mask_c), tCgC(_,_,_,chunk_idx), tCsC(_,_,_,write_stage_c)); + + // Advance pipeline_state + ++pipeline_state_c; + } + } + } + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_c_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mC_mkl = params.tma_load_c.get_tma_tensor(make_shape(L,N,C,G*B)); + Tensor gC_mkl = local_tile(mC_mkl, TileShapeIntraBMM1{}, make_coord(_,_,_), Step{}); // (BLK_M,BLK_K,n,k,l) + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("mC_mkl : ");print(mC_mkl);print("\n"); + print("gC_mkl : ");print(gC_mkl);print("\n"); + } +#endif + return make_tuple(gC_mkl); + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + template< + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineC, class PipelineStateC + > + CUTLASS_DEVICE void + load_b_c_tail(MainloopPipelineB pipeline_b, PipelineStateB pipeline_state_b, + MainloopPipelineC pipeline_c, PipelineStateC pipeline_state_c) { + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline_b.producer_tail(pipeline_state_b); + pipeline_c.producer_tail(pipeline_state_c); + } + } + + template< + class Params, class ProblemShape, + class MainloopPipeline, class PipelineState, + class TensorStorage + > + CUTLASS_DEVICE + void load_delta( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + MainloopPipeline& pipeline, PipelineState& pipeline_state, + TensorStorage& shared_tensors) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + + auto& gDelta = params.tensor_delta; + auto& gDeltaA = params.tensor_delta_a; + + Element* ptr_delta = shared_tensors.smem_delta.data(); + ElementDA* ptr_delta_a = shared_tensors.smem_delta_a.data(); + auto smem_layout = make_layout(make_shape(get<0>(TileShape{}), Int{})); + auto sDelta = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_delta), smem_layout)); + auto sDeltaA = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_delta_a), smem_layout)); + + auto bulk_atom_dt = Copy_Atom{}; + auto bulk_atom_dA = Copy_Atom{}; + + auto chunk = get<3>(problem_size); + // Mainloop + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) { + int write_stage = pipeline_state.index(); + + // LOCK pipeline_state for _writing_ + pipeline.producer_acquire(pipeline_state); + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_state); + copy(bulk_atom_dt.with(*tma_barrier), gDelta(_,chunk_idx,blk_coord), sDelta(_,write_stage)); + copy(bulk_atom_dA.with(*tma_barrier), gDeltaA(_,chunk_idx,blk_coord), sDeltaA(_,write_stage)); + + // Advance pipeline_state + ++pipeline_state; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + template< + class MainloopPipeline, class PipelineState + > + CUTLASS_DEVICE void + load_delta_tail(MainloopPipeline pipeline, PipelineState pipeline_state) { + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(pipeline_state); + } + } + + template< + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineC, class PipelineStateC, + class TensorStorage + > + CUTLASS_DEVICE + auto mma_intra_1( + int& chunk, + MainloopPipelineB& pipeline_b, PipelineStateB& pipeline_state_b, + MainloopPipelineC& pipeline_c, PipelineStateC& pipeline_state_c, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + + int read_b_stage = pipeline_state_b.index(); + int read_c_stage = pipeline_state_c.index(); + + TiledMmaIntra1 tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + + // Mainloop setup QK + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_b.data()), SmemLayoutB{}); + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_c.data()), SmemLayoutC{}); + + Tensor tCsC = thr_mma.partition_A(sC); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tBsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrC = thr_mma.make_fragment_A(tCsC); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tBrB = thr_mma.make_fragment_B(tBsB); // (MMA,MMA_M,MMA_N,PIPE) + + Tensor accumulator = partition_fragment_C(tiled_mma, take<0, 2>(TileShapeIntraBMM1{})); + + auto barrier_token_b = pipeline_b.consumer_try_wait(pipeline_state_b); + auto barrier_token_c = pipeline_c.consumer_try_wait(pipeline_state_c); + + pipeline_b.consumer_wait(pipeline_state_b, barrier_token_b); + pipeline_c.consumer_wait(pipeline_state_c, barrier_token_c); + +#if 0 + if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0) && chunk == 0) { + print("tBsB : ");print(tBsB);print("\n"); + print("tCsC : ");print(tCsC);print("\n"); + print("tBrB : ");print(tBrB);print("\n"); + print("tCrC : ");print(tCrC);print("\n"); + print("accumulator : ");print(accumulator);print("\n"); + } +#endif + + warpgroup_fence_operand(accumulator); + warpgroup_arrive(); + gemm_zero_acc(tiled_mma, tCrC(_,_,_,read_c_stage), tBrB(_,_,_,read_b_stage), accumulator); + warpgroup_commit_batch(); + + warpgroup_wait<0>(); + + // Maybe delay + pipeline_b.consumer_release(pipeline_state_b); + ++pipeline_state_b; + pipeline_c.consumer_release(pipeline_state_c); + ++pipeline_state_c; + + return make_tuple(accumulator); + } + + template< + class MainloopPipelineDelta, class PipelineStateDelta, + class TensorIntra1, + class TensorStorage + > + CUTLASS_DEVICE + auto pre_intra_2( + int& chunk, + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_state_delta, + TensorIntra1& tIntra1, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + int read_delta_stage = pipeline_state_delta.index(); + + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::TransformBarrier); }; + + TiledMmaIntra1 mma; + + auto ts0 = size<0>(TileShape{}); + + Layout row_layout = make_layout(make_shape ( ts0, ts0, Int{}), + make_stride(_0{}, _1{}, ts0)); + Layout col_layout = make_layout(make_shape ( ts0, ts0, Int{}), + make_stride(_1{}, _0{}, ts0)); + + Tensor sDelta_row = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta.data()), row_layout)); + Tensor sDeltaA_row = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), row_layout)); + Tensor sDeltaA_col = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), col_layout)); + Tensor coord_tensor = make_coord_tensor(make_identity_layout(make_shape(ts0, ts0))); + + ThrMMA thr_mma = mma.get_slice(thread_idx); + + Tensor tDelta_row = thr_mma.partition_C(sDelta_row)(_,_,_,read_delta_stage); + Tensor tDeltaA_row = thr_mma.partition_C(sDeltaA_row)(_,_,_,read_delta_stage); + Tensor tDeltaA_col = thr_mma.partition_C(sDeltaA_col)(_,_,_,read_delta_stage); + Tensor tC = thr_mma.partition_C(coord_tensor); + + auto tSR_Delta_bf16 = make_tensor(shape(tDelta_row)); + auto tSR_DeltaA_row_bf16 = make_tensor(shape(tDeltaA_row)); + auto tSR_DeltaA_col_bf16 = make_tensor(shape(tDeltaA_col)); + + auto tSR_Delta = make_tensor(shape(tDelta_row)); + auto tSR_DeltaA_row = make_tensor(shape(tDeltaA_row)); + auto tSR_DeltaA_col = make_tensor(shape(tDeltaA_col)); + + auto barrier_token = pipeline_delta.consumer_try_wait(pipeline_state_delta); + pipeline_delta.consumer_wait(pipeline_state_delta, barrier_token); + + // Load delta/delta_a + copy(tDelta_row, tSR_Delta_bf16); + copy(tDeltaA_row, tSR_DeltaA_row_bf16); + copy(tDeltaA_col, tSR_DeltaA_col_bf16); + + pipeline_delta.consumer_release(pipeline_state_delta); + ++pipeline_state_delta; + + // Data convert + type_convert(tSR_Delta_bf16, tSR_Delta); + type_convert(tSR_DeltaA_row_bf16, tSR_DeltaA_row); + type_convert(tSR_DeltaA_col_bf16, tSR_DeltaA_col); + + // SegSum + #pragma unroll + for (int ii = 0; ii < size(tIntra1); ++ii) { + ElementAcc tmp(0.f); + auto [m,n] = tC(ii); + if (m >= n) { + tmp = expf(tSR_DeltaA_col(ii) - tSR_DeltaA_row(ii)); + } + tIntra1(ii) = tmp * tSR_Delta(ii) * tIntra1(ii); + } + + synchronize(); + + return make_tuple(tIntra1); + } + + + template< + class MainloopPipelineX, class PipelineStateX, + class TensorIntra1, + class TensorStorage + > + CUTLASS_DEVICE + auto mma_intra_2( + int& chunk, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_state_x, + TensorIntra1& tIntra1, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + + Tensor sX = make_tensor(make_smem_ptr(shared_tensors.smem_x.data()), SmemLayoutX{}); + + TiledMmaIntra2 tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCsX = thr_mma.partition_B(sX); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrX = thr_mma.make_fragment_B(tCsX); // (MMA,MMA_N,MMA_K,PIPE) + + auto tCrIntra = make_acc_into_op(tIntra1, typename TiledMmaIntra2::LayoutA_TV{}); + auto accumulator = partition_fragment_C(tiled_mma, take<0, 2>(TileShapeIntraBMM2{})); + +#if 0 + if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0) && chunk == 0) { + print("tCsX : ");print(tCsX);print("\n"); + print("tCrX : ");print(tCrX);print("\n"); + print("tIntra1 : ");print(tIntra1);print("\n"); + print("tCrIntra : ");print(tCrIntra);print("\n"); + print("accumulator : ");print(accumulator);print("\n"); + } +#endif + + int read_stage = pipeline_state_x.index(); + + auto barrier_token = pipeline_x.consumer_try_wait(pipeline_state_x); + pipeline_x.consumer_wait(pipeline_state_x, barrier_token); + + // Delay the X pipeline release + warpgroup_wait<0>(); + warpgroup_fence_operand(tCrIntra); + warpgroup_fence_operand(accumulator); + + gemm_zero_acc(tiled_mma, tCrIntra, tCrX(_,_,_,read_stage), accumulator); + + warpgroup_commit_batch(); + + // May delay + warpgroup_wait<0>(); + + pipeline_x.consumer_release(pipeline_state_x); + ++pipeline_state_x; + + return make_tuple(accumulator); + } + + template< + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineDelta, class PipelineStateDelta, + class TensorStorage + > + CUTLASS_DEVICE + auto pre_inter_1( + int& chunk, + MainloopPipelineB& pipeline_b, PipelineStateB& pipeline_state_b, + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_state_delta, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + + // Different from the original SmemB + using SmemLayoutB = typename CollectiveInterMma1::SmemLayoutA; + + Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_b.data()), SmemLayoutB{}); + Tensor sB = as_position_independent_swizzle_tensor(sB_); + + int read_delta_stage = pipeline_state_delta.index(); + int read_b_stage = pipeline_state_b.index(); + + auto barrier_token_b = pipeline_b.consumer_try_wait(pipeline_state_b); + auto barrier_token_delta = pipeline_delta.consumer_try_wait(pipeline_state_delta); + + // construct + TiledMmaInter1 tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCrB_mma = thr_mma.partition_fragment_A(sB(_,_,Int<0>{})); + Tensor tCrB_load_bf16 = make_fragment_like(tCrB_mma); + Tensor tCrB_load = make_fragment_like(tCrB_mma); + + // using CopyAtom = Copy_Atom; + using CopyAtom = Copy_Atom; + auto smem_tiled_copy_B = make_tiled_copy_A(CopyAtom{}, tiled_mma); + auto smem_thr_copy_B = smem_tiled_copy_B.get_thread_slice(thread_idx); + Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB_load_bf16); + Tensor tCsB = smem_thr_copy_B.partition_S(sB); + + // load delta/delta_a + auto ts0 = size<0>(TileShape{}); + + Layout row_layout = make_layout(make_shape ( ts0, ts0, Int{}), + make_stride(_0{}, _1{}, ts0)); + + Tensor sDelta = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta.data()), row_layout)); + Tensor sDeltaA = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), row_layout)); + + Tensor tDelta = thr_mma.partition_A(sDelta)(_,_,_,read_delta_stage); + Tensor tDeltaA = thr_mma.partition_A(sDeltaA)(_,_,_,read_delta_stage); + + auto tSR_Delta_bf16 = make_tensor(shape(tDelta)); + auto tSR_DeltaA_bf16 = make_tensor(shape(tDeltaA)); + + auto tSR_Delta = make_tensor(shape(tDelta)); + auto tSR_DeltaA = make_tensor(shape(tDeltaA)); + + pipeline_b.consumer_wait(pipeline_state_b, barrier_token_b); + + // load B + copy(smem_tiled_copy_B, tCsB(_,_,_,read_b_stage), tCrB_copy_view(_,_,_)); + + pipeline_delta.consumer_wait(pipeline_state_delta, barrier_token_delta); + + // Load delta/delta_a + copy(tDelta, tSR_Delta_bf16); + copy(tDeltaA, tSR_DeltaA_bf16); + + // data conver + type_convert(tSR_Delta_bf16, tSR_Delta); + type_convert(tSR_DeltaA_bf16, tSR_DeltaA); + type_convert(tCrB_load_bf16, tCrB_load); + + // xxx : This is specific to GMMA and it depends on the known TV partitioner. + // Only work for 128x64x128 tile as the layout is ((REG_N,REG_M,REG_N_REP),ATOM_M_REP,ATOM_N_REP). + // So we can use ((-1,_,-1),_,-1) to iterate the last column process (here -1 means the last coordinate) + static_assert(cute::rank(tSR_DeltaA) == Int<3>{}, "The rank of Tensor is mismatched"); + static_assert(cute::rank<0>(tSR_DeltaA) == Int<3>{}, "The rank of MMA_ATOM is mismatched"); + static_assert(std::is_same::value, "Only ALayout_64x16 is supported"); + // get last column + bool last_thread_per_row = thread_idx % 4 == 3; + ElementAcc last_column = ElementAcc(0.f); + if (last_thread_per_row) { + auto coord0 = get<0>(shape<0>(tSR_DeltaA)) - _1{}; + auto coord1 = get<2>(shape<0>(tSR_DeltaA)) - _1{}; + auto coord2 = get<2>(shape(tSR_DeltaA)) - _1{}; + auto tSR_DeltaA_slice = tSR_DeltaA(make_coord(coord0,_,coord1),_,coord2); + last_column = tSR_DeltaA_slice(_0{}); + } + + // broadcast the last column + last_column = __shfl_sync(0xFFFFFFFF, last_column, 3, 32); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tSR_DeltaA); ++ii) { + tSR_DeltaA(ii) = expf(last_column - tSR_DeltaA(ii)); + } + + // Delta*DeltaA + auto tCompute = make_tensor(tCrB_mma.layout()); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tCompute); ++ii) { + tCompute(ii) = tSR_DeltaA(ii) * tSR_Delta(ii); + } + + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tCompute); ++ii) { + tCompute(ii) = tCrB_load(ii) * tCompute(ii); + } + + copy(tCompute, tCrB_mma); + + // release B buffer + pipeline_b.consumer_release(pipeline_state_b); + ++pipeline_state_b; + return make_tuple(tCrB_mma, last_column); + } + + template< + class MainloopPipelineX, class PipelineStateX, + class TensorB, + class TensorStorage + > + CUTLASS_DEVICE + auto mma_inter_1( + int& chunk, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_state_x, + TensorB tCrB, + TensorStorage& shared_tensors) { + + Tensor sX = make_tensor(make_smem_ptr(shared_tensors.smem_x.data()), SmemLayoutX{}); + + int thread_idx = int(threadIdx.x % 128); + TiledMmaInter1 tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + + Tensor tCsX = thr_mma.partition_B(sX); + Tensor tCrX = thr_mma.make_fragment_B(tCsX); + + Tensor accumulator = partition_fragment_C(tiled_mma, take<0, 2>(TileShapeInterBMM1{})); + + int read_stage = pipeline_state_x.index(); + + auto barrier_token_x = pipeline_x.consumer_try_wait(pipeline_state_x); + pipeline_x.consumer_wait(pipeline_state_x, barrier_token_x); + + warpgroup_fence_operand(accumulator); + warpgroup_arrive(); + gemm_zero_acc(tiled_mma, tCrB, tCrX(_,_,_,read_stage), accumulator); + warpgroup_commit_batch(); + + warpgroup_wait<0>(); + + return make_tuple(accumulator); + } + + template< + class TensorStorage + > + CUTLASS_DEVICE + auto state_init(TensorStorage& shared_tensors) { + TiledMmaInter1 tiled_mma; + Tensor tensor_state = partition_fragment_C(tiled_mma, take<0, 2>(TileShapeInterBMM1{})); + + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::TransposeBarrier); }; + int thread_idx = int(threadIdx.x % 128); + + Tensor sP_ = make_tensor(make_smem_ptr(shared_tensors.smem_p.data()), SmemLayoutP{}); // (D, N) + // Transpose + auto sP_tmp = make_tensor(sP_.data(), make_layout(reverse(shape(sP_)), reverse(stride(sP_)))); // (N, D) + Tensor sP = as_position_independent_swizzle_tensor(sP_tmp); + + using CopyAtomC = Copy_Atom; + TiledCopy tiled_r2s_atom = make_tiled_copy_C(CopyAtomC{}, tiled_mma); + + using CopyOpR2S = SM90_U16x8_STSM_T; + TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom{}, tiled_r2s_atom); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rP = make_tensor(shape(thread_r2s.partition_S(sP))); + Tensor tRS_sP = thread_r2s.partition_D(sP); + + clear(tRS_rP); + clear(tensor_state); + // Store P + copy(tiled_r2s, tRS_rP, tRS_sP); + + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + + return make_tuple(tensor_state); + } + template< + class TensorInter1, + class TensorState + > + CUTLASS_DEVICE + void pre_inter_2( + ElementAcc& last_column, + TensorInter1& tInter1, + TensorState& tState) { + + // Apply the state + for (int ii = 0; ii < size(tInter1); ++ii) { + tInter1(ii) = tInter1(ii) + expf(last_column) * static_cast(tState(ii)); + } + // Update the state + copy(tInter1, tState); + + } + + template< + class TensorState, + class TensorStorage + > + CUTLASS_DEVICE + void post_inter_2( + TensorState& tState, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::TransposeBarrier); }; + + Tensor sP_ = make_tensor(make_smem_ptr(shared_tensors.smem_p.data()), SmemLayoutP{}); // (D, N) + // Transpose + auto sP_tmp = make_tensor(sP_.data(), make_layout(reverse(shape(sP_)), reverse(stride(sP_)))); // (N, D) + Tensor sP = as_position_independent_swizzle_tensor(sP_tmp); + auto tP = make_tensor(shape(tState)); + + type_convert(tState, tP); + + synchronize(); // ensure all threads have issued their async fence + + using CopyAtomC = Copy_Atom; + TiledMmaInter1 tiled_mma; + TiledCopy tiled_r2s_atom = make_tiled_copy_C(CopyAtomC{}, tiled_mma); + + using CopyOpR2S = SM90_U16x8_STSM_T; + TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom{}, tiled_r2s_atom); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rP = thread_r2s.retile_S(tP); + Tensor tRS_sP = thread_r2s.partition_D(sP); + + // Store P + copy(tiled_r2s, tRS_rP, tRS_sP); + + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + } + + template< + class MainloopPipelineC, class PipelineStateC, + class MainloopPipelineDelta, class PipelineStateDelta, + class TensorStorage + > + CUTLASS_DEVICE + auto mma_inter_2( + int& chunk, + MainloopPipelineC& pipeline_c, PipelineStateC& pipeline_state_c, + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_state_delta, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + + int read_c_stage = pipeline_state_c.index(); + + TiledMmaInter2 tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + + // Mainloop setup QK + Tensor sP = make_tensor(make_smem_ptr(shared_tensors.smem_p.data()), SmemLayoutP{}); + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_c.data()), SmemLayoutC{}); + + Tensor tCsC = thr_mma.partition_A(sC); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tPsP = thr_mma.partition_B(sP); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrC = thr_mma.make_fragment_A(tCsC); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tPrP = thr_mma.make_fragment_B(tPsP); // (MMA,MMA_M,MMA_N,PIPE) + + Tensor accumulator = partition_fragment_C(tiled_mma, take<0, 2>(TileShapeInterBMM2{})); + +#if 0 + if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0) && chunk == 0) { + print("sP : ");print(sP);print("\n"); + print("sC : ");print(sC);print("\n"); + print("tCsC : ");print(tCsC);print("\n"); + print("tPsP : ");print(tPsP);print("\n"); + print("tCrC : ");print(tCrC);print("\n"); + print("tPrP : ");print(tPrP);print("\n"); + print("accumulator : ");print(accumulator);print("\n"); + } +#endif + + auto barrier_token_c = pipeline_c.consumer_try_wait(pipeline_state_c); + pipeline_c.consumer_wait(pipeline_state_c, barrier_token_c); + + warpgroup_fence_operand(accumulator); + warpgroup_arrive(); + gemm_zero_acc(tiled_mma, tCrC(_,_,_,read_c_stage), tPrP, accumulator); + warpgroup_commit_batch(); + + // release C + warpgroup_wait<0>(); + + pipeline_c.consumer_release(pipeline_state_c); + ++pipeline_state_c; + + int read_delta_stage = pipeline_state_delta.index(); + // Load delta for epilogue + auto col_layout = make_layout(make_shape(get<0>(TileShape{}), get<0>(TileShape{}), Int{}), make_stride(_1{}, _0{}, get<0>(TileShape{}))); + auto c_tv_layout = typename TiledMmaInter2::LayoutC_TV{}; + auto tile_mn = make_shape(tile_size<0>(tiled_mma), tile_size<1>(tiled_mma)); + auto delta_col_layout = zipped_divide(col_layout, tile_mn); + auto delta_a_col_tv_layout = composition(delta_col_layout, make_tuple(c_tv_layout,_)); + + auto tDeltaA_col_ = make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), delta_a_col_tv_layout)(make_coord(thread_idx,_), make_coord(_,_,read_delta_stage)); + auto tDeltaA_col = as_position_independent_swizzle_tensor(tDeltaA_col_); + auto tSR_DeltaA_col = make_tensor(shape(tDeltaA_col)); + + copy(tDeltaA_col, tSR_DeltaA_col); + + pipeline_delta.consumer_release(pipeline_state_delta); + ++pipeline_state_delta; + + for (int ii = 0; ii < size(tSR_DeltaA_col); ++ii) { + tSR_DeltaA_col(ii) = expf(tSR_DeltaA_col(ii)); + } + + return make_tuple(accumulator, tSR_DeltaA_col); + } + + template< + class ElementSrc, class ElementDst, + class TensorSrc, class TensorDst + > + CUTLASS_DEVICE + auto type_convert( + TensorSrc& tS, + TensorDst& tD) { + + static constexpr int FragmentSize = 2; + NumericArrayConverter converter; + + auto tS_frg = recast>(tS); + auto tD_frg = recast>(tD); + + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tS_frg); ++ii) { + tD_frg(ii) = converter(tS_frg(ii)); + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/examples/111_hopper_ssd/device/ssd.hpp b/examples/111_hopper_ssd/device/ssd.hpp new file mode 100644 index 00000000..21e9f07a --- /dev/null +++ b/examples/111_hopper_ssd/device/ssd.hpp @@ -0,0 +1,273 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + +// common +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" + +#if !defined(__CUDACC_RTC__) +#include "cutlass/cluster_launch.hpp" +#include "cutlass/trace.h" +#endif // !defined(__CUDACC_RTC__) + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::ssd::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class SSD { +public: + using Kernel = Kernel_; + + static int const kThreadCount = Kernel::MaxThreadsPerBlock; + + /// Argument structure: User API + using Arguments = typename Kernel::Arguments; + /// Argument structure: Kernel API + using Params = typename Kernel::Params; + +private: + + /// Kernel API parameters object + Params params_; + + bool is_initialized(bool set = false) { + static bool initialized = false; + if (set) initialized = true; + return initialized; + } + +public: + + /// Access the Params structure + Params const& params() const { + return params_; + } + + /// Determines whether the GEMM can execute the given problem. + static Status + can_implement(Arguments const& args) { + if (Kernel::can_implement(args)) { + return Status::kSuccess; + } + else { + return Status::kInvalid; + } + } + + /// Gets the workspace size + static size_t + get_workspace_size(Arguments const& args) { + size_t workspace_bytes = 0; + workspace_bytes += Kernel::get_workspace_size(args); + return workspace_bytes; + } + + /// Computes the grid shape + static dim3 + get_grid_shape(Params const& params) { + return Kernel::get_grid_shape(params); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int /* smem_capacity */ = -1) { + CUTLASS_TRACE_HOST("Universal::maximum_active_blocks()"); + int max_active_blocks = -1; + int smem_size = Kernel::SharedStorageSize; + + // first, account for dynamic smem capacity if needed + cudaError_t result; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + result = cudaFuncSetAttribute( + device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " + << cudaGetErrorString(result)); + return -1; + } + } + + // query occupancy after setting smem size + result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + device_kernel, + Kernel::MaxThreadsPerBlock, + smem_size); + + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " + << cudaGetErrorString(result)); + return -1; + } + + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Initializes GEMM state from arguments. + Status + initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("Universal::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Initialize the workspace + Status status = Kernel::initialize_workspace(args, workspace, stream); + if (status != Status::kSuccess) { + return status; + } + + // Initialize the Params structure + params_ = Kernel::to_underlying_arguments(args, workspace); + + if (is_initialized()) return Status::kSuccess; + + // account for dynamic smem capacity if needed + int smem_size = Kernel::SharedStorageSize; + printf("[Usage] smem : %d\n", smem_size); + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + cudaError_t result = cudaFuncSetAttribute( + device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + } + + is_initialized(true); + + return Status::kSuccess; + } + + /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. + Status + update(Arguments const& args, void* workspace = nullptr) { + CUTLASS_TRACE_HOST("Universal()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + if (workspace_bytes > 0 && nullptr == workspace) { + return Status::kErrorWorkspaceNull; + } + + params_ = Kernel::to_underlying_arguments(args, workspace); + return Status::kSuccess; + } + + /// Primary run() entry point API that is static allowing users to create and manage their own params. + /// Supplied params struct must be construct by calling Kernel::to_underling_arguments() + static Status + run(Params& params, cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("Universal::run()"); + dim3 const block = Kernel::get_block_shape(); + dim3 const grid = get_grid_shape(params); + + // configure smem size and carveout + int smem_size = Kernel::SharedStorageSize; + + Status launch_result; + // Use extended launch API only for mainloops that use it + if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) { + dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}), + cute::size<1>(typename Kernel::ClusterShape{}), + cute::size<2>(typename Kernel::ClusterShape{})); + void const* kernel = (void const*) device_kernel; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params); + } + else { + launch_result = Status::kSuccess; + device_kernel<<>>(params); + } + + cudaError_t result = cudaGetLastError(); + if (cudaSuccess == result && Status::kSuccess == launch_result) { + return Status::kSuccess; + } + else { + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } + } + + // + // Non-static launch overloads that first create and set the internal params struct of this kernel handle. + // + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + if (Status::kSuccess == status) { + status = run(params_, stream); + } + return status; + } + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + return run(args, workspace, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + run(cudaStream_t stream = nullptr) { + return run(params_, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + operator()(cudaStream_t stream = nullptr) { + return run(params_, stream); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/examples/111_hopper_ssd/kernel/sm90_ssd_kernel_builder.hpp b/examples/111_hopper_ssd/kernel/sm90_ssd_kernel_builder.hpp new file mode 100644 index 00000000..d285b002 --- /dev/null +++ b/examples/111_hopper_ssd/kernel/sm90_ssd_kernel_builder.hpp @@ -0,0 +1,116 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 "../collective/sm90_ssd_epilogue.hpp" +#include "../collective/sm90_ssd_gemm_tma_warpspecialized.hpp" +#include "../kernel/sm90_ssd_kernel_tma_warpspecialized.hpp" +#include "../kernel/sm90_ssd_tile_scheduler.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +namespace cutlass::ssd::kernel { + +template< + class Element_, + class ElementDA_, + class ElementAcc_, + class ElementY_, + class TileShape_, + bool HAS_D_, + bool D_HAS_HDIM_, + bool HAS_Z_ +> +struct Sm90SsdBuilder { + + using Element = Element_; + using ElementDA = ElementDA_; + using ElementAcc = ElementAcc_; + using ElementY = ElementY_; + using TileShape = TileShape_; + + static constexpr bool HAS_D = HAS_D_; + static constexpr bool D_HAS_HDIM = D_HAS_HDIM_; + static constexpr bool HAS_Z = HAS_Z_; + + static constexpr int StagesY = 2; + static constexpr int StagesX = 2; + static constexpr int StagesZ = 1; // smem size limitation + using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; + using Schedule = cutlass::epilogue::TmaWarpSpecialized; + using EpilogueTile = decltype(cutlass::epilogue::collective::detail::sm90_compute_tile_shape_or_override< + ElementY, EpilogueTileType, Schedule, TileShape>()); + + using SmemLayoutAtomY = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::MN, ElementY, decltype(get<0>(EpilogueTile{})), decltype(get<1>(EpilogueTile{}))>()); + using SmemLayoutY = decltype(tile_to_shape( + SmemLayoutAtomY{}, + make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int{}), + Step<_2,_1,_3>{})); + + using SmemLayoutAtomX = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::MN, Element, decltype(get<0>(TileShape{})), decltype(get<1>(TileShape{}))>()); + using SmemLayoutX = decltype(tile_to_shape( + SmemLayoutAtomY{}, + make_shape(size<0>(TileShape{}), size<1>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + + using SmemLayoutAtomZ = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::MN, Element, decltype(get<0>(TileShape{})), decltype(get<1>(TileShape{}))>()); + using SmemLayoutZ = decltype(tile_to_shape( + SmemLayoutAtomZ{}, + make_shape(size<0>(TileShape{}), size<1>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + + static constexpr auto epi_tile_m = size<0>(EpilogueTile{}); + static constexpr auto epi_tile_n = size<1>(EpilogueTile{}); + static constexpr auto partial_m = Int<128>{}; + static constexpr auto partial_n = Int{}; + + using SmemLayoutAtomPartialY = typename GMMA::Layout_K_SW64_Atom; + using SmemLayoutPartialY = decltype(tile_to_shape( + SmemLayoutAtomPartialY{}, + make_shape(partial_m, partial_n, Int{}))); + + using CollectiveMainloop = cutlass::ssd::collective::SsdMainloopTmaWarpSpecialized; + using CollectiveEpilogue = cutlass::ssd::collective::SsdEpilogue< + ElementAcc, ElementY, TileShape, + EpilogueTile, SmemLayoutX, SmemLayoutY, SmemLayoutPartialY, typename CollectiveMainloop::SmemLayoutP, SmemLayoutZ, + StagesX, StagesY, StagesZ, + HAS_D, D_HAS_HDIM, HAS_Z>; + using TileScheduler = cutlass::ssd::kernel::PersistentTileScheduler; + using Kernel = cutlass::ssd::kernel::SsdKernelTmaWarpSpecialized; + +}; + +} diff --git a/examples/111_hopper_ssd/kernel/sm90_ssd_kernel_tma_warpspecialized.hpp b/examples/111_hopper_ssd/kernel/sm90_ssd_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..0619e8db --- /dev/null +++ b/examples/111_hopper_ssd/kernel/sm90_ssd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,553 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/arch/reg_reconfig.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/arch/arch.h" + +namespace cutlass::ssd::kernel { + +using namespace cute; + +template< + class CollectiveMainloop, + class CollectiveEpilogue, + class TileScheduler +> +struct SsdKernelTmaWarpSpecialized { + + static const int NumLoadWarpGroups = 1; + // hard code + static constexpr int NumMmaWarpGroups = 2; + + // TileShape: LDN + using TileShape = typename CollectiveMainloop::TileShape; + // Force to use 1x1x1 + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + // Pipeline for Tensor X + using MainloopPipelineX = typename CollectiveMainloop::MainloopPipelineX; + using PipelineParamsX = typename MainloopPipelineX::Params; + using PipelineStateX = typename cutlass::PipelineState; + + // Pipeline for Tensor Delta && DeltaA + using MainloopPipelineDelta = typename CollectiveMainloop::MainloopPipelineDelta; + using PipelineParamsDelta = typename MainloopPipelineDelta::Params; + using PipelineStateDelta = typename cutlass::PipelineState; + + // Pipeline for Tensor X + using MainloopPipelineB = typename CollectiveMainloop::MainloopPipelineB; + using PipelineParamsB = typename MainloopPipelineB::Params; + using PipelineStateB = typename cutlass::PipelineState; + + // Pipeline for Tensor X + using MainloopPipelineC = typename CollectiveMainloop::MainloopPipelineC; + using PipelineParamsC = typename MainloopPipelineC::Params; + using PipelineStateC = typename cutlass::PipelineState; + + // Pipeline for cooperate-warps + using CooperatePipeline = typename CollectiveEpilogue::CooperatePipeline; + using PipelineParamsCo = typename CooperatePipeline::Params; + using PipelineStateCo = typename cutlass::PipelineState; + + // Pipeline for Tensor D + using EpiloadPipelineD = typename CollectiveEpilogue::EpiloadPipelineD; + using PipelineParamsD = typename EpiloadPipelineD::Params; + using PipelineStateD = typename cutlass::PipelineState; + + // Pipeline for Tensor Z + using EpiloadPipelineZ = typename CollectiveEpilogue::EpiloadPipelineZ; + using PipelineParamsZ = typename EpiloadPipelineZ::Params; + using PipelineStateZ = typename cutlass::PipelineState; + + struct TensorStorage { + typename CollectiveMainloop::SharedStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + struct SharedStorage { + TensorStorage tensors; + + using PipelineStorageX = typename MainloopPipelineX::SharedStorage; + using PipelineStorageDelta = typename MainloopPipelineDelta::SharedStorage; + using PipelineStorageB = typename MainloopPipelineB::SharedStorage; + using PipelineStorageC = typename MainloopPipelineC::SharedStorage; + using PipelineStorageCo = typename CooperatePipeline::SharedStorage; + using PipelineStorageD = typename EpiloadPipelineD::SharedStorage; + using PipelineStorageZ = typename EpiloadPipelineZ::SharedStorage; + + // pipeline + alignas(16) PipelineStorageX pipeline_storage_x; + alignas(16) PipelineStorageDelta pipeline_storage_delta; + alignas(16) PipelineStorageB pipeline_storage_b; + alignas(16) PipelineStorageC pipeline_storage_c; + alignas(16) PipelineStorageCo pipeline_storage_co; + alignas(16) PipelineStorageD pipeline_storage_d; + alignas(16) PipelineStorageZ pipeline_storage_z; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + // [G, B, EH, C, L, D, N] + using ProblemShape = cute::tuple; + + struct Arguments { + ProblemShape problem_size; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_size; + typename CollectiveMainloop::Params mainloop; + typename CollectiveEpilogue::Params epilogue; + typename TileScheduler::Params tile_scheduler; + }; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = (NumMmaWarpGroups + NumLoadWarpGroups) * cutlass::NumThreadsPerWarpGroup; + using ArchTag = cutlass::arch::Sm90; + + // CTA reconfig (TBD) + static constexpr uint32_t LoadRegisterRequirement = 40 - 2 * 8; + static constexpr uint32_t TotalRegisterSupply = (64*1024 / MaxThreadsPerBlock / MinBlocksPerMultiprocessor / 8) * 8 * MaxThreadsPerBlock / cutlass::NumThreadsPerWarpGroup; + static constexpr uint32_t MmaRegisterRequirement = ((TotalRegisterSupply - LoadRegisterRequirement) / NumMmaWarpGroups / 8) * 8; + // static constexpr uint32_t LoadRegisterRequirement = 40; + // static constexpr uint32_t MmaRegisterRequirement = 232; + + static size_t get_workspace_size(Arguments const& args) { return 0; } + static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t) { + return cutlass::Status::kSuccess; + } + + static bool can_implement(Arguments const& args) { + return CollectiveMainloop::can_implement(args.problem_size, args.mainloop); + } + + static dim3 get_grid_shape(Params const& params) { + return TileScheduler::get_grid_shape(params.tile_scheduler); + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + return Params{ + args.problem_size, + CollectiveMainloop::to_underlying_arguments(args.problem_size, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_size, args.epilogue, workspace), + TileScheduler::to_underlying_arguments(args.problem_size, args.hw_info, ClusterShape{}, TileShape{}) + }; + } + + CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { + + // TBD + enum class WarpGroupRole { + Producer = 0, + Consumer0 = 1, + Consumer1 = 2 + }; + // TBD + enum class ProducerWarpRole { + LoadX = 0, + LoadDelta = 1, + LoadBC = 2, + LoadZ = 3 + }; + + + // Parameters + // [G, B, EH, C, L, D, N] + auto C = get<3>(params.problem_size); + + // Shared memory. + auto& storage = *reinterpret_cast(smem); + + int lane_idx = cutlass::canonical_lane_idx(); + int warp_idx = cutlass::canonical_warp_idx_sync(); + int warp_idx_in_warp_group = warp_idx % cutlass::NumWarpsPerWarpGroup; + int warp_group_idx = cutlass::canonical_warp_group_idx(); + auto warp_group_role = WarpGroupRole(warp_group_idx); + auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group); + int lane_predicate = cute::elect_one_sync(); + uint32_t block_rank_in_cluster = cute::block_rank_in_cluster(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + // Pipeline (TBD) + PipelineParamsX pipeline_params_x; + pipeline_params_x.transaction_bytes = CollectiveMainloop::kXLoadBytes; + pipeline_params_x.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadX); + pipeline_params_x.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups; + pipeline_params_x.initializing_warp = 4; + + PipelineParamsDelta pipeline_params_delta; + pipeline_params_delta.transaction_bytes = CollectiveMainloop::kDeltaLoadBytes + CollectiveMainloop::kDeltaALoadBytes; + pipeline_params_delta.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadDelta); + pipeline_params_delta.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups; + pipeline_params_delta.initializing_warp = 5; + + PipelineParamsB pipeline_params_b; + pipeline_params_b.transaction_bytes = CollectiveMainloop::kBLoadBytes; + pipeline_params_b.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadBC); + pipeline_params_b.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups; + pipeline_params_b.initializing_warp = 6; + + PipelineParamsC pipeline_params_c; + pipeline_params_c.transaction_bytes = CollectiveMainloop::kCLoadBytes; + pipeline_params_c.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadBC); + pipeline_params_c.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups; + pipeline_params_c.initializing_warp = 7; + + PipelineParamsCo pipeline_params_co; + pipeline_params_co.producer_arv_count = cutlass::NumThreadsPerWarpGroup; + pipeline_params_co.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; + pipeline_params_co.initializing_warp = 8; + + PipelineParamsD pipeline_params_d; + pipeline_params_d.transaction_bytes = CollectiveEpilogue::kEpiloadDBytes; + pipeline_params_d.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadDelta); + pipeline_params_d.num_consumers = cutlass::NumThreadsPerWarpGroup; + pipeline_params_d.initializing_warp = 9; + + PipelineParamsZ pipeline_params_z; + pipeline_params_z.transaction_bytes = CollectiveEpilogue::kEpiloadZBytes; + pipeline_params_z.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadZ); + pipeline_params_z.num_consumers = cutlass::NumThreadsPerWarpGroup; + pipeline_params_z.initializing_warp = 10; + + if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadX) { + pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadDelta) { + pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Producer; + pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadBC) { + pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Producer; + pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadZ) { + pipeline_params_z.role = EpiloadPipelineZ::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer; + pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Consumer; + pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer; + pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Consumer; + } + if (warp_group_role == WarpGroupRole::Consumer0) { + pipeline_params_co.role = CooperatePipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer1) { + pipeline_params_co.role = CooperatePipeline::ThreadCategory::Consumer; + pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Consumer; + pipeline_params_z.role = EpiloadPipelineZ::ThreadCategory::Consumer; + } + + MainloopPipelineX pipeline_x(storage.pipeline_storage_x, pipeline_params_x, Shape<_1,_1,_1>{}); + PipelineStateX mainloop_pipe_x_consumer; + PipelineStateX mainloop_pipe_x_producer = cutlass::make_producer_start_state(); + + MainloopPipelineDelta pipeline_delta(storage.pipeline_storage_delta, pipeline_params_delta, Shape<_1,_1,_1>{}); + PipelineStateDelta mainloop_pipe_delta_consumer; + PipelineStateDelta mainloop_pipe_delta_producer = cutlass::make_producer_start_state(); + + MainloopPipelineB pipeline_b(storage.pipeline_storage_b, pipeline_params_b, Shape<_1,_1,_1>{}); + PipelineStateB mainloop_pipe_b_consumer; + PipelineStateB mainloop_pipe_b_producer = cutlass::make_producer_start_state(); + + MainloopPipelineC pipeline_c(storage.pipeline_storage_c, pipeline_params_c, Shape<_1,_1,_1>{}); + PipelineStateC mainloop_pipe_c_consumer; + PipelineStateC mainloop_pipe_c_producer = cutlass::make_producer_start_state(); + + CooperatePipeline pipeline_co(storage.pipeline_storage_co, pipeline_params_co); + PipelineStateCo cooperate_pipe_consumer_state; + PipelineStateCo cooperate_pipe_producer_state = cutlass::make_producer_start_state(); + + EpiloadPipelineD pipeline_d(storage.pipeline_storage_d, pipeline_params_d, Shape<_1,_1,_1>{}); + PipelineStateD epi_load_pipe_d_consumer; + PipelineStateD epi_load_pipe_d_producer = cutlass::make_producer_start_state(); + + EpiloadPipelineZ pipeline_z(storage.pipeline_storage_z, pipeline_params_z, Shape<_1,_1,_1>{}); + PipelineStateZ epi_load_pipe_z_consumer; + PipelineStateZ epi_load_pipe_z_producer = cutlass::make_producer_start_state(); + + // Epilogue Store pipeline + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state(); + + // Epilogue Store P pipeline + using EpiStorePPipeline = typename CollectiveEpilogue::StorePPipeline; + typename EpiStorePPipeline::Params epi_store_p_pipeline_params; + epi_store_p_pipeline_params.always_wait = true; + EpiStorePPipeline epi_store_p_pipeline(epi_store_p_pipeline_params); + + PipelineState epi_store_p_pipe_producer_state = cutlass::make_producer_start_state(); + + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer blocks in the Cluster + // and to finish smem init + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); + } + else { + __syncthreads(); + } + + // Kernel implement(TBD) + + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue; + + if (warp_group_role == WarpGroupRole::Producer) { + // disable reg dealloc to enable print + cutlass::arch::warpgroup_reg_dealloc(); + if (producer_warp_role == ProducerWarpRole::LoadX) { + // use local variable to avoid STL/LDL + TileScheduler tile_scheduler{params.tile_scheduler}; + auto load_input = collective_mainloop.load_x_init(params.mainloop, params.problem_size); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + // Load X + collective_mainloop.load_x( + blk_coord, params.mainloop, params.problem_size, + pipeline_x, mainloop_pipe_x_producer, + load_input, + storage.tensors.mainloop + ); + } + collective_mainloop.load_x_tail( + pipeline_x, mainloop_pipe_x_producer + ); + } + else if (producer_warp_role == ProducerWarpRole::LoadDelta) { + TileScheduler tile_scheduler{params.tile_scheduler}; + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto blk_coord_eh = tile_scheduler.get_block_coord_eh(); + // Epiload + collective_epilogue.load_d( + blk_coord_eh, params.epilogue, params.problem_size, + pipeline_d, epi_load_pipe_d_producer, + storage.tensors.mainloop + ); + // Load Delta + // Load DeltaA + collective_mainloop.load_delta( + blk_coord, params.mainloop, params.problem_size, + pipeline_delta, mainloop_pipe_delta_producer, + storage.tensors.mainloop + ); + } + collective_mainloop.load_delta_tail( + pipeline_delta, mainloop_pipe_delta_producer + ); + } + else if (producer_warp_role == ProducerWarpRole::LoadBC) { + TileScheduler tile_scheduler{params.tile_scheduler}; + auto load_input_b = collective_mainloop.load_b_init(params.mainloop, params.problem_size); + auto load_input_c = collective_mainloop.load_c_init(params.mainloop, params.problem_size); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord_b(); + // Load B + collective_mainloop.load_b_c( + blk_coord, params.mainloop, params.problem_size, + pipeline_b, mainloop_pipe_b_producer, + pipeline_c, mainloop_pipe_c_producer, + load_input_b, + load_input_c, + storage.tensors.mainloop + ); + } + collective_mainloop.load_b_c_tail( + pipeline_b, mainloop_pipe_b_producer, + pipeline_c, mainloop_pipe_c_producer + ); + } + else if (producer_warp_role == ProducerWarpRole::LoadZ) { + // use local variable to avoid STL/LDL + TileScheduler tile_scheduler{params.tile_scheduler}; + auto load_input = collective_epilogue.load_z_init(params.epilogue, params.problem_size); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + // Load X + collective_epilogue.load_z( + blk_coord, params.epilogue, params.problem_size, + pipeline_z, epi_load_pipe_z_producer, + load_input, + storage.tensors.epilogue + ); + } + collective_epilogue.load_z_tail( + pipeline_z, epi_load_pipe_z_producer + ); + } + } + // Warpgroup1 for Intra + // Warpgroup2 for Inter + else if (warp_group_role == WarpGroupRole::Consumer0) { + TileScheduler tile_scheduler{params.tile_scheduler}; + cutlass::arch::warpgroup_reg_alloc(); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + for (int chunk = 0; chunk < C; ++chunk) { + auto blk_coord = tile_scheduler.get_block_coord(); + // IntraBMM1 + // Wait B + // Wait C + auto [tIntra1] = collective_mainloop.mma_intra_1( + chunk, + pipeline_b, mainloop_pipe_b_consumer, + pipeline_c, mainloop_pipe_c_consumer, + storage.tensors.mainloop + ); + // Pre Intra2 + auto [tPreIntra2] = collective_mainloop.pre_intra_2( + chunk, + pipeline_delta, mainloop_pipe_delta_consumer, + tIntra1, + storage.tensors.mainloop + ); + // IntraBMM2 + auto [tIntra2] = collective_mainloop.mma_intra_2( + chunk, + pipeline_x, mainloop_pipe_x_consumer, + tPreIntra2, + storage.tensors.mainloop + ); + collective_epilogue.store_intra( + chunk, blk_coord, params.epilogue, params.problem_size, + pipeline_co, cooperate_pipe_producer_state, + tIntra2, + typename CollectiveMainloop::TiledMmaIntra2{}, + storage.tensors.epilogue + ); + } + } + } + else if (warp_group_role == WarpGroupRole::Consumer1) { + TileScheduler tile_scheduler{params.tile_scheduler}; + cutlass::arch::warpgroup_reg_alloc(); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto [tState] = collective_mainloop.state_init(storage.tensors.mainloop); + auto blk_coord_eh = tile_scheduler.get_block_coord_eh(); + bool is_first_iteration = true; + for (int chunk = 0; chunk < C; ++chunk) { + auto blk_coord = tile_scheduler.get_block_coord(); + // Pre Inter1 + // Wait delta + // Wait deltaA + auto [tPreInter1, last_column] = collective_mainloop.pre_inter_1( + chunk, + pipeline_b, mainloop_pipe_b_consumer, + pipeline_delta, mainloop_pipe_delta_consumer, + storage.tensors.mainloop + ); + // InterBMM1 + // Wait X + auto [tInter1] = collective_mainloop.mma_inter_1( + chunk, + pipeline_x, mainloop_pipe_x_consumer, + tPreInter1, + storage.tensors.mainloop + ); + collective_mainloop.pre_inter_2( + last_column, + tInter1, + tState + ); + // InterBMM2 + auto [tInter2, tDelta] = collective_mainloop.mma_inter_2( + chunk, + pipeline_c, mainloop_pipe_c_consumer, + pipeline_delta, mainloop_pipe_delta_consumer, + storage.tensors.mainloop + ); + // Pre Inter2 + collective_mainloop.post_inter_2( + tState, + storage.tensors.mainloop + ); + // update_d + auto [tD] = collective_epilogue.update_d( + blk_coord_eh, params.epilogue, + is_first_iteration, + pipeline_d, epi_load_pipe_d_consumer, + typename CollectiveMainloop::TiledMmaInter2{}, + storage.tensors.mainloop + ); + is_first_iteration = false; + // Epilogue TensorY store + collective_epilogue.store( + chunk, blk_coord, params.epilogue, params.problem_size, + epi_store_pipeline, epi_store_pipe_producer_state, + pipeline_co, cooperate_pipe_consumer_state, + pipeline_x, mainloop_pipe_x_consumer, + pipeline_z, epi_load_pipe_z_consumer, + tInter2, tDelta, tD, + typename CollectiveMainloop::TiledMmaInter2{}, + storage.tensors.epilogue, storage.tensors.mainloop + ); + } + + if constexpr (CollectiveEpilogue::D_HAS_HDIM) { + // update the barrier + pipeline_d.consumer_release(epi_load_pipe_d_consumer); + ++epi_load_pipe_d_consumer; + } + + auto blk_coord = tile_scheduler.get_block_coord(); + // Epilogue Fstate store + collective_epilogue.store_p( + blk_coord, params.epilogue, params.problem_size, + epi_store_p_pipeline, epi_store_p_pipe_producer_state, + storage.tensors.mainloop + ); + } + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/examples/111_hopper_ssd/kernel/sm90_ssd_tile_scheduler.hpp b/examples/111_hopper_ssd/kernel/sm90_ssd_tile_scheduler.hpp new file mode 100644 index 00000000..20a2e509 --- /dev/null +++ b/examples/111_hopper_ssd/kernel/sm90_ssd_tile_scheduler.hpp @@ -0,0 +1,132 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/fast_math.h" +#include "cutlass/kernel_hardware_info.h" + +namespace cutlass::ssd::kernel { + +//////////////////////////////////////////////////////////////////////////////// + +struct PersistentTileScheduler { + + struct Params { + int num_blocks; + int num_groups; + FastDivmod divmod_eh; + FastDivmod divmod_ngroup_ratio; + + KernelHardwareInfo hw_info; + }; + + int block_idx = 0; + Params params; + + CUTLASS_DEVICE + PersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, TileShape const& tile_shape) + { + using namespace cute; + auto [G, B, EH, C, L, D, N] = problem_size; + + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = hw_info.sm_count; + if (sm_count <= 0) { + CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + } + + CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count); + hw_info.sm_count = sm_count; + + int num_blocks = B * EH; + int ngroup_ratio = EH / G; + + return Params { + num_blocks, + G, + {EH}, + {ngroup_ratio}, + hw_info + }; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1); + return grid; + } + + CUTLASS_DEVICE + bool is_valid() { + return block_idx < params.num_blocks; + } + + CUTLASS_DEVICE + auto get_block_coord() { + return block_idx; + } + + CUTLASS_DEVICE + auto get_block_coord_b() { + using namespace cute; + int eh_idx, b_idx; + int g_idx, rest_idx; + params.divmod_eh(b_idx, eh_idx, block_idx); + params.divmod_ngroup_ratio(g_idx, rest_idx, eh_idx); + return (params.num_groups * b_idx + g_idx); + } + + CUTLASS_DEVICE + auto get_block_coord_eh() { + using namespace cute; + int eh_idx, b_idx; + params.divmod_eh(b_idx, eh_idx, block_idx); + return eh_idx; + } + + CUTLASS_DEVICE + PersistentTileScheduler& operator++() { + block_idx += gridDim.x; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::ssd::kernel diff --git a/examples/111_hopper_ssd/reference/reference_ssd.hpp b/examples/111_hopper_ssd/reference/reference_ssd.hpp new file mode 100644 index 00000000..05a524df --- /dev/null +++ b/examples/111_hopper_ssd/reference/reference_ssd.hpp @@ -0,0 +1,345 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 "cute/tensor.hpp" + +// training or inference phase (not used yet) +// PHASE 0 : training +// PHASE 1 : inference +#define PHASE 0 + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template< + bool transA, + bool transB, + class Element, + class TensorA, + class TensorB, + class TensorC +> +void mma( + TensorA tA, + TensorB tB, + TensorC tC) { + using namespace cute; + + int M = transA ? int(shape<1>(tA)) : int(shape<0>(tA)); + int N = transB ? int(shape<1>(tB)) : int(shape<0>(tB)); + int K = transA ? int(shape<0>(tA)) : int(shape<1>(tA)); + for (int mi = 0; mi < M; ++mi) { + for (int ni = 0; ni < N; ++ni) { + for (int ki = 0; ki < K; ++ki) { + float a = static_cast(Element(transA ? tA(ki, mi) : tA(mi, ki))); + float b = static_cast(Element(transB ? tB(ki, ni) : tB(ni, ki))); + tC(mi, ni) += a * b; + } + } + } + +} + +template< + class Element, + class Tensor +> +auto segsum(Tensor tensor) { + using namespace cute; + auto C = shape<0>(tensor); + auto L = shape<1>(tensor); + auto cum_sum = make_tensor(make_shape(C,L)); + // cum_sum + for (int ci = 0; ci < C; ++ci) { + for (int li = 0; li < L; ++li) { + if (li == 0) { + cum_sum(ci, li) = tensor(ci, li); + } + else { + cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li); + } + } + } + auto seg_sum_out = make_tensor(make_shape(C, L, L)); + // seg_sum + // [ 1, 0, 0] + // [ e^y, 1, 0] + // [e^(y+z), e^z, 1] + CUTLASS_PRAGMA_UNROLL + for (int ci = 0; ci < C; ++ci) { + for (int i = 0; i < L; ++i) { + for (int j = 0; j < L; ++j) { + if (j < i) { + float tmp = static_cast(cum_sum(ci, i)) - static_cast(cum_sum(ci, j)); + seg_sum_out(ci, i, j) = expf(tmp); + } + else if (j == i) { + seg_sum_out(ci, i, j) = 1.f; + } + else { + seg_sum_out(ci, i, j) = 0.f; + } + } + } + } + + return seg_sum_out; +} + +template< + class Element, + class Tensor +> +auto cumsum( + Tensor tensor) { + using namespace cute; + auto C = shape<0>(tensor); + auto L = shape<1>(tensor); + auto cum_sum = make_tensor(make_shape(C,L)); + auto cum_sum_out = make_tensor(make_shape(C,L)); + auto cum_sum_exp_out = make_tensor(make_shape(C, L)); + auto cum_sum_exp_out_last = make_tensor(make_shape(C, L)); + auto last_column = make_tensor(make_shape(C)); + // [x, x+y, x+y+z, ..] + CUTLASS_PRAGMA_UNROLL + for (int ci = 0; ci < C; ++ci) { + for (int li = 0; li < L; ++li) { + if (li == 0) { + cum_sum(ci, li) = tensor(ci, li); + } + else { + cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li); + } + // cum_sum_out(ci, li) = static_cast(cum_sum(ci, li)); + } + } + CUTLASS_PRAGMA_UNROLL + for (int ci = 0; ci < C; ++ci) { + last_column(ci) = static_cast(cum_sum(ci, L-1)); + CUTLASS_PRAGMA_UNROLL + for (int li = 0; li < L; ++li) { + cum_sum_exp_out_last(ci, li) = expf(static_cast(last_column(ci) - cum_sum(ci, li))); + cum_sum_exp_out(ci, li) = expf(static_cast(cum_sum(ci, li))); + } + } + + return make_tuple(cum_sum_exp_out_last, last_column, cum_sum_exp_out); +} + +template< + bool HAS_D, + bool D_HAS_HDIM, + bool HAS_Z, + class TensorY, + class TensorF, + class TensorX, + class TensorDelta, + class TensorDeltaA, + class TensorB, + class TensorC, + class TensorD, + class TensorZ, + class Params +> +void ssd_reference_impl( + TensorY mY, TensorF mF, + TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA, + TensorB mB, TensorC mC, TensorD mD, TensorZ mZ, + Params params) { + + using namespace cute; + using Element = typename Params::Element; + using ElementAcc = typename Params::ElementAcc; + + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + // d [ eh, d] + auto [G, B, EH, C, L, D, N] = params.get_problem_shape(); + int group_ratio = EH / G; + for (int b = 0; b < B; ++b) { + for (int eh = 0; eh < EH; ++eh) { + int g = eh / group_ratio; + auto tY = mY(b,eh,_,_,_); + auto tF = mF(b,eh,_,_); + auto tX = mX(b,eh,_,_,_); + auto tDelta = mDelta(b,eh,_,_); + auto tDeltaA = mDeltaA(b,eh,_,_); + auto tB = mB(b,g,_,_,_); + auto tC = mC(b,g,_,_,_); + auto tD = mD(eh,_); + auto tZ = mZ(b,eh,_,_,_); + // IntraBMM1 BxC, LxLxN, NT + // B: [n, c, l] + // C: [n, c, l] + // O: [c, l, l] + auto tIntraBMM1_out = make_tensor(make_shape(C,L,L)); + for (int ci = 0; ci < C; ++ci) { + mma(tC(_,ci,_), tB(_,ci,_), tIntraBMM1_out(ci,_,_)); + } + // Pre_IntraBMM2 DeltaA_IntraBMM2 x Delta x IntraBMM_out + // DeltaA_xxx : [c, l, l] + // Delta : [c, l, _] + // IntraBMM1_out: [c, l, l] + auto tDeltaA_IntraBMM2 = segsum(tDeltaA); + auto tIntraBMM2_inp = make_tensor(make_shape(C, L, L)); + for (int ci = 0; ci < C; ++ci) { + for (int i = 0; i < L; ++i) { + for (int j = 0; j < L; ++j) { + tIntraBMM2_inp(ci, i, j) = tDeltaA_IntraBMM2(ci, i, j) * tDelta(ci, j) * tIntraBMM1_out(ci, i, j); + } + } + } + // IntraBMM2 IntraBMM2_inp x X, LxDxL, TT + // IntraBMM2_inp: [c, l, l] + // X : [d, c, l] + // IntraBMM2_out: [c, l, d] + auto tIntraBMM2_out = make_tensor(make_shape(C,L,D)); + for (int ci = 0; ci < C; ++ci) { + mma(tIntraBMM2_inp(ci,_,_), tX(_,ci,_), tIntraBMM2_out(ci,_,_)); + } + // Pre_InterBMM1 DeltaA_InterBMM1 x Delta x B + // DeltaA_xxx : [c, l] + // Delta : [c, l] + // IntraBMM1_out: [c, n, l] + auto [tDeltaA_InterBMM1, tLast, tCumsum_exp] = cumsum(tDeltaA); + auto tInterBMM1_inp = make_tensor(make_shape(C, N, L)); + for (int ci = 0; ci < C; ++ci) { + for (int i = 0; i < N; ++i) { + for (int j = 0; j < L; ++j) { + tInterBMM1_inp(ci, i, j) = tDeltaA_InterBMM1(ci, j) * tDelta(ci, j) * tB(i, ci, j); + } + } + } + // InterBMM1 InterBMM1_inp x X, NxDxL, swapAB, TT + // InterBMM1_inp: [c, n, l] + // X : [d, c, l] + // InterBMM1_out: [c, n, d] + auto tInterBMM1_out = make_tensor(make_shape(C,N,D)); + for (int ci = 0; ci < C; ++ci) { + mma(tInterBMM1_inp(ci,_,_), tX(_,ci,_), tInterBMM1_out(ci,_,_)); + } + // Initialize state + // PreInterBMM2 + // InterBMM1_out: [c, n, d] + // Last : [c] + auto tInterBMM2_inp = make_tensor(make_shape(C, N, D)); + for (int ci = 0; ci < C; ++ci) { + for (int ni = 0; ni < N; ++ ni){ + for (int di = 0; di < D; ++di) { + if (ci == 0) { + tInterBMM2_inp(ci, ni, di) = 0; + } + else { + tInterBMM2_inp(ci, ni, di) = tInterBMM1_out(ci - 1, ni, di) + expf(tLast(ci - 1)) * tInterBMM2_inp(ci - 1, ni, di); + } + } + } + } + // InterBMM2 InterBMM2_inp x C, LxDxN, NT + // C : [n, c, l] + // InterBMM2_inp: [c, n, d] + // InterBMM2_out: [c, l, d] + auto tInterBMM2_out = make_tensor(make_shape(C,L,D)); + for (int ci = 0; ci < C; ++ci) { + mma(tC(_,ci,_), tInterBMM2_inp(ci,_,_), tInterBMM2_out(ci,_,_)); + } + // Epilogue Cumsum_exp x InterBMM2_out + IntraBMM2_out + // InterBMM2_out: [c, l, d] + // IntraBMM2_out: [c, l, d] + // Cumsum_exp : [c, l] + for (int ci = 0; ci < C; ++ci) { + for (int li = 0; li < L; ++li) { + for (int di = 0; di < D; ++di) { + float y = tInterBMM2_out(ci, li, di) * tCumsum_exp(ci, li) + tIntraBMM2_out(ci, li, di); + float scale; + if constexpr (D_HAS_HDIM) { + scale = static_cast(tD(di)); + } + else { + scale = static_cast(tD(_0{})); + } + if constexpr (HAS_D) { + y = y + static_cast(tX(di, ci, li)) * scale; + } + else { + y = y; + } + if constexpr (HAS_Z) { + float z = static_cast(tZ(di, ci, li)); + // y = y * z * (1 / (1 + exp(-z))); + y = y * z * (1 / (1 + exp(-z))); + } + tY(di, ci, li) = static_cast(y); + } + } + } + // Epilogue Fstate(last C) + for (int ni = 0; ni < N; ++ ni){ + for (int di = 0; di < D; ++di) { + tF(di, ni) = static_cast(tInterBMM1_out(C - 1, ni, di) + expf(tLast(C - 1)) * tInterBMM2_inp(C - 1, ni, di)); + } + } + } + } +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template< + bool HAS_D, + bool D_HAS_HDIM, + bool HAS_Z, + class TensorY, + class TensorF, + class TensorX, + class TensorDelta, + class TensorDeltaA, + class TensorB, + class TensorC, + class TensorD, + class TensorZ, + class Params +> +void ssd_reference( + TensorY mY, TensorF mF, + TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA, + TensorB mB, TensorC mC, TensorD mD, TensorZ mZ, + Params params) { + ssd_reference_impl(mY, mF, mX, mDelta, mDeltaA, mB, mC, mD, mZ, params); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/111_hopper_ssd/reference/reference_ssd_cumsum.hpp b/examples/111_hopper_ssd/reference/reference_ssd_cumsum.hpp new file mode 100644 index 00000000..124052f8 --- /dev/null +++ b/examples/111_hopper_ssd/reference/reference_ssd_cumsum.hpp @@ -0,0 +1,194 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 +#include + +#include "cutlass/coord.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/tensor_view.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/host/gemm.h" +#include "cutlass/arch/arch.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +#include "cute/int_tuple.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "cute/util/debug.hpp" +#include "cute/config.hpp" + +namespace cutlass::ssd::kernel { + +using namespace cute; + +template< + class Element_, + class ElementD_, + class TileShape_> +struct CumsumKernel { + using Element = Element_; + using ElementD = ElementD_; + using TileShape = TileShape_; // L,D,N + + // Required by `device_kernel` + static constexpr int MaxThreadsPerBlock = 128; + static constexpr int MinBlocksPerMultiprocessor = 1; + using ArchTag = arch::Sm90; + + static constexpr int AlignmentBytes = 16; + + struct SharedStorage { + /* empty, no smem needed */ + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + struct TransformArguments { + const Element* ptr_DeltaA; + ElementD* ptr_Cumsum; + }; + + struct TransformParams { + const Element* ptr_DeltaA; + ElementD* ptr_Cumsum; + }; + + using ProblemShape = cute::tuple; // b, eh, c, l + struct Arguments { + ProblemShape problem_shape{}; + TransformArguments transform{}; + KernelHardwareInfo hw_info{}; + }; + + struct Params { + ProblemShape problem_shape{}; + TransformParams transform{}; + KernelHardwareInfo hw_info{}; + }; + + static Params + to_underlying_arguments(Arguments const& args, void* workspace) { + return Params{ + ProblemShape{args.problem_shape}, + TransformParams{args.transform.ptr_DeltaA, args.transform.ptr_Cumsum}, + KernelHardwareInfo{args.hw_info}}; + } + + static Status + can_implement(Arguments const& args) { + return Status::kSuccess; + } + + static size_t + get_workspace_size(Arguments const& args) { + return size_t(0); + } + + static Status + initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr, + CudaHostAdapter *cuda_adapter = nullptr) { + return Status::kSuccess; + } + + static dim3 + get_grid_shape(Params const& params) { + auto [B, EH, C, L] = params.problem_shape; + return dim3(B*EH, 1, 1); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + CUTE_HOST_DEVICE + void + operator()(Params params, [[maybe_unused]] char* smem_buf = nullptr) { + auto [B, EH, C, L] = params.problem_shape; + + auto layout = make_layout(make_shape(L, C, EH*B)); + + auto mD_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_DeltaA), make_layout(reverse(layout.shape()), reverse(layout.stride()))); + auto mC_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_Cumsum), make_layout(reverse(layout.shape()), reverse(layout.stride()))); + auto cD_bcl = make_identity_tensor(shape(mD_bcl)); + + int blk_idx = blockIdx.x; + int thread_idx = threadIdx.x; + + auto tD = logical_divide(mD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_); + auto tC = logical_divide(mC_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_); + auto cD = logical_divide(cD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_); + + static constexpr int NumPacked = AlignmentBytes / sizeof(ElementD); + using PackedTypeDeltaA = uint_bit_t * NumPacked>; + using PackedTypeCumsum = uint_bit_t * NumPacked>; + +#if 0 + if (thread_idx % 128 == 0 && blk_idx == 0) { + print("tD : ");print(tD);print("\n"); + print("tC : ");print(tC);print("\n"); + print("cD : ");print(cD);print("\n"); + } +#endif + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < shape<0>(tD); ++i) { + float last_element = 0.f; + auto crd = cD(i,_0{}); + auto tD_recast = recast(tD); + auto tC_recast = recast(tC); + if (elem_less(crd, shape(mD_bcl))) { + for (int j = 0; j < shape<1>(tD_recast); ++j) { + auto tD_slice = make_tensor(make_shape(Int{})); + auto tC_slice = make_tensor(make_shape(Int{})); + auto tD_slice_recast = recast(tD_slice); + auto tC_slice_recast = recast(tC_slice); + + tD_slice_recast(_0{}) = tD_recast(i,j); + for (int k = 0; k < NumPacked; ++ k) { + last_element += static_cast(tD_slice(k)); + tC_slice(k) = static_cast(last_element); + } + tC_recast(i,j) = tC_slice_recast(_0{}); + } + } + } + } + +private: + +}; + +} // End namespace cutlass diff --git a/examples/112_blackwell_ssd/112_blackwell_ssd.cu b/examples/112_blackwell_ssd/112_blackwell_ssd.cu new file mode 100644 index 00000000..ba2b1cbb --- /dev/null +++ b/examples/112_blackwell_ssd/112_blackwell_ssd.cu @@ -0,0 +1,850 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + * + **************************************************************************************************/ + +#include +#include "cutlass/util/command_line.h" + +#include "cutlass/cutlass.h" + +#include "cute/tensor.hpp" +#include "cute/layout.hpp" +#include "cutlass/kernel_hardware_info.hpp" + +#include "thrust/universal_vector.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/device/tensor_fill.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +#include "reference/reference_ssd_cumsum.hpp" +#include "reference/reference_ssd.hpp" + +#include "cutlass/transform/device/transform_universal_adapter.hpp" + +#include "device/ssd.hpp" +#include "kernel/sm100_ssd_kernel_builder.hpp" + +using namespace cute; + +// Command line options parsing +struct Options { + + using Element = cutlass::bfloat16_t; + using ElementAcc = float; + using ElementDA = float; + static constexpr bool D_HAS_HDIM = true; + static constexpr bool HAS_D = true; + // Blackwell SSD doesn't support Z now(huge perf drop). + static constexpr bool HAS_Z = false; + + bool help; + bool error; + + // All static number now + int G = 2; + int B = 3; + int E = 2; + int H = 2; + // Reference kernel doesn't support dynamic C now. + static constexpr auto C = Int<8>{}; + static constexpr auto D = Int<64>{}; + static constexpr auto L = Int<128>{}; + static constexpr auto N = Int<128>{}; + int EH = E * H; + + int iterations; + bool verify; + bool verbose; + + int warmups; + bool measure; + + Options(): + help(false), + error(false), + iterations(1), verify(true), + measure(false), warmups(3) + {} + + // Parses the command line + void parse(int argc, char const **args) { + cutlass::CommandLine cmd(argc, args); + + Options defaults; + + if (cmd.check_cmd_line_flag("help")) { + help = true; + return; + } + + cmd.get_cmd_line_argument("iterations", iterations, defaults.iterations); + cmd.get_cmd_line_argument("G", G, defaults.G); + cmd.get_cmd_line_argument("B", B, defaults.B); + cmd.get_cmd_line_argument("E", E, defaults.E); + cmd.get_cmd_line_argument("H", H, defaults.H); + verbose = cmd.check_cmd_line_flag("verbose"); + verify = !(cmd.check_cmd_line_flag("without_verify")); + + EH = E*H; + + if (iterations > 1) { + measure = true; + verbose = true; + } + + auto problem_shape = cute::make_tuple(G, B, EH, C, L, D, N); + cute::print("problem_shape : "); cute::print(problem_shape); cute::print("\n"); + } + + /// Prints the usage statement. + std::ostream & print_usage(std::ostream &out) const { + + out << "112_blackwell_ssd\n\n" + << "Options:\n\n" + << " --help If specified, displays this usage statement\n\n" + << " --iterations= Benchmarking iterations.\n" + << " --without_verify Don't verify the results.\n" + << " --verbose Print execution time per kernel\n" + << " --G= Group\n" + << " --B= Batch\n" + << " --E= Expanded factor\n" + << " --H= Number of heads\n" + << "\n"; + + return out; + } + + + auto get_problem_shape() const { + return cute::make_tuple(G, B, EH, C, L, D, N); + } + + // acceptable layout by cuDNN + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + + auto layoutX() const { + auto layout = make_layout(make_shape(L, C, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutDelta() const { + auto layout = make_layout(make_shape(L, C, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutDeltaA() const { + auto layout = make_layout(make_shape(L, C, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutB() const { + auto layout = make_layout(make_shape(L, C, N, G, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutC() const { + auto layout = make_layout(make_shape(L, C, N, G, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutY() const { + auto layout = make_layout(make_shape(L, C, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutF() const { + auto layout = make_layout(make_shape(N, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + auto layoutD() const { + if constexpr (D_HAS_HDIM) { + auto layout = make_layout(make_shape(D, EH)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + else { + auto layout = make_layout(make_shape(Int<1>{}, EH)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + } + + auto layoutZ() const { + auto layout = make_layout(make_shape(L, C, D, EH, B)); + return make_layout(reverse(layout.shape()), reverse(layout.stride())); + } + + // transformed layout for kernel parameters + + auto layoutX_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B)); + return make_layout( + make_shape(D,L,int32_t(C),EH*B), + make_stride( + stride<2>(layout), + stride<0>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutB_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),N,G*B)); + return make_layout( + make_shape(L,N,int32_t(C),G*B), + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutC_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),N,G*B)); + return make_layout( + make_shape(L,N,int32_t(C),G*B), + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutDelta_transformed() const { + return make_layout(make_shape(L,int32_t(C),EH*B)); + } + + auto layoutY_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B)); + return make_layout( + make_shape(L,D,int32_t(C),EH*B), // (M,K,L,...) + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + + auto layoutF_transformed() const { + auto layout = make_layout(make_shape(N,D,EH*B)); + return make_layout( + make_shape(D,N,EH*B), + make_stride( + stride<1>(layout), + stride<0>(layout), + stride<2>(layout) + ) + ); + } + + auto layoutD_transformed() const { + if constexpr (D_HAS_HDIM) { + return make_layout(make_shape(D, EH)); + } + else { + return make_layout(make_shape(Int<1>{}, EH)); + } + } + + auto layoutZ_transformed() const { + auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B)); + return make_layout( + make_shape(L,D,int32_t(C),EH*B), + make_stride( + stride<0>(layout), + stride<2>(layout), + stride<1>(layout), + stride<3>(layout) + ) + ); + } + +}; + +template +static void +initialize_values( + thrust::universal_vector& dst_ptr, + cutlass::Distribution::Kind dist_kind, + uint64_t seed, + Element var = Element(1.f)) { + if (cutlass::Distribution::Uniform == dist_kind) { + int scope = 2; + cutlass::reference::host::BlockFillRandomUniform( + dst_ptr.data().get(), dst_ptr.size(), seed, scope, -scope, 0); + } + else if (cutlass::Distribution::AllZeros == dist_kind) { + cutlass::reference::host::BlockFillRandomUniform( + dst_ptr.data().get(), dst_ptr.size(), seed, 0, 0, 0); + } + else if (cutlass::Distribution::AllOnes == dist_kind) { + cutlass::reference::host::BlockFillRandomUniform( + dst_ptr.data().get(), dst_ptr.size(), seed, 1, 1, 0); + } + else if (cutlass::Distribution::Gaussian == dist_kind) { + cutlass::reference::device::BlockFillRandomGaussian( + dst_ptr.data().get(), dst_ptr.size(), seed, (Element) 0, var); + } + else if (cutlass::Distribution::Sequential == dist_kind) { + cutlass::reference::host::BlockFillSequential(dst_ptr.data().get(), dst_ptr.size()); + } + else { + std::cerr << "Invalid distribution kind!\n."; + exit(1); + } +} + +template < + class Options_ +> +struct TestBed { + using Option = Options_; + using Element = typename Option::Element; + using ElementDA = typename Option::ElementDA; + using ElementAcc = typename Option::ElementAcc; + + thrust::universal_vector tensor_X; + thrust::universal_vector tensor_DeltaA; + thrust::universal_vector tensor_DeltaA_cumsum; + thrust::universal_vector tensor_Delta; + thrust::universal_vector tensor_B; + thrust::universal_vector tensor_C; + thrust::universal_vector tensor_D; + thrust::universal_vector tensor_Y; + thrust::universal_vector tensor_Z; + thrust::universal_vector tensor_Y_ref_0; + thrust::universal_vector tensor_Y_ref_1; + thrust::universal_vector tensor_F; + thrust::universal_vector tensor_F_ref_0; + thrust::universal_vector tensor_F_ref_1; + + cutlass::Distribution::Kind init_X = cutlass::Distribution::Uniform; + cutlass::Distribution::Kind init_DeltaA = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind init_Delta = cutlass::Distribution::Gaussian; + cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform; + cutlass::Distribution::Kind init_C = cutlass::Distribution::Uniform; + + using TileShape = decltype(make_shape(Options::L, Options::D, Options::N)); // (L, D, N) + using SsdOperation = cutlass::ssd::device::SSD< + typename cutlass::ssd::kernel::Sm100SsdBuilder< + Element, ElementDA, ElementAcc, Element, + TileShape, + Option::HAS_D, Option::D_HAS_HDIM + >::Kernel>; + using CumsumKenrel = cutlass::ssd::kernel::CumsumKernel; + using CumsumOperation = cutlass::transform::device::TransformUniversalAdapter; + + bool initialize(Options const& options, const cutlass::KernelHardwareInfo& hw_info, uint64_t seed = 2024) { + auto [g, b, eh, c, l, d, n] = options.get_problem_shape(); + assert(g == 1 && "Only group size == 1 is supported") ; + + auto size_X = b * eh * c * l * d; + auto size_DeltaA = b * eh * c * l; + auto size_Delta = b * eh * c * l; + auto size_B = g * b * c * n * l; + auto size_C = g * b * c * n * l; + auto size_Y = b * eh * c * l * d; + auto size_F = b * eh * d * n; + + tensor_X .resize(sizeof(Element) * size(options.layoutX())); + tensor_DeltaA .resize(sizeof(Element) * size(options.layoutDeltaA())); + tensor_Delta .resize(sizeof(Element) * size(options.layoutDelta())); + tensor_B .resize(sizeof(Element) * size(options.layoutB())); + tensor_C .resize(sizeof(Element) * size(options.layoutC())); + tensor_D .resize(sizeof(Element) * size(options.layoutD())); + tensor_Z .resize(sizeof(Element) * size(options.layoutZ())); + tensor_Y .resize(sizeof(Element) * size(options.layoutY())); + tensor_Y_ref_0.resize(sizeof(Element) * size(options.layoutY())); + tensor_Y_ref_1.resize(sizeof(Element) * size(options.layoutY())); + tensor_F .resize(sizeof(Element) * size(options.layoutF())); + tensor_F_ref_0.resize(sizeof(Element) * size(options.layoutF())); + tensor_F_ref_1.resize(sizeof(Element) * size(options.layoutF())); + + tensor_DeltaA_cumsum.resize(sizeof(ElementDA) * size(options.layoutDeltaA())); + + // Limit distribution to reduce skew between hosts and devices + initialize_values(tensor_X, init_X, seed); + initialize_values(tensor_DeltaA, init_DeltaA, seed + 1, Element(0.05f)); + initialize_values(tensor_Delta, init_Delta, seed + 3, Element(0.05f)); + initialize_values(tensor_B, init_B, seed + 5); + initialize_values(tensor_C, init_C, seed + 7); + initialize_values(tensor_D, init_C, seed + 9); + initialize_values(tensor_Z, init_X, seed); + + cudaError_t result; + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the Initialization kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + } + + // apply cumsum(device) before kernel launch + typename CumsumOperation::Arguments arguments{ + make_shape(int(b), int(eh), int(c), int(l)), + { + tensor_DeltaA.data().get(), + tensor_DeltaA_cumsum.data().get(), + }, + hw_info + }; + + CumsumOperation op; + + size_t workspace_size = CumsumOperation::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + cutlass::Status status = op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "This kernel is not supported. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + status = op.initialize(arguments, workspace.get()); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + // may be used uninitialized + cudaEvent_t start; + cudaEvent_t end; + cudaEventCreate(&start); + cudaEventCreate(&end); + + // warm up + if (options.measure) { + for (int i = 0; i < options.warmups; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + } + result = cudaEventRecord(start); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + // Run + for (int i = 0; i < options.iterations; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + result = cudaEventRecord(end); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return false; + } + + float runtime_ms = 0; + result = cudaEventElapsedTime(&runtime_ms, start, end); + if (result != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + runtime_ms /= static_cast(options.iterations); + + if (options.verbose) { + printf("[iters = %d, warmups = %d] cumsum kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms); + } + + return true; + } + + bool sufficient() const { + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + int max_smem_size; + result = cudaDeviceGetAttribute(&max_smem_size, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_idx); + if (result != cudaSuccess) { + throw std::runtime_error("cudaDeviceGetAttribute() failed"); + } + + return true; + } + + bool run(Options const& options, const cutlass::KernelHardwareInfo& hw_info) { + if (!sufficient()) { + std::cerr << "Test waived due to insufficient CUDA device.\n"; + return true; + } + + if (!initialize(options, hw_info)) { + std::cerr << "Failed to initialize the test.\n"; + return true; + }; + + auto [g, b, eh, c, l, d, n] = options.get_problem_shape(); + typename SsdOperation::Arguments arguments{ + make_shape(int(g), int(b), int(eh), int(c), int(l), int(d), int(n)), + { + tensor_X.data().get(), + tensor_DeltaA_cumsum.data().get(), + tensor_Delta.data().get(), + tensor_B.data().get(), + tensor_C.data().get(), + options.layoutX_transformed(), + options.layoutB_transformed(), + options.layoutC_transformed(), + options.layoutDelta_transformed() + }, + { + tensor_Y.data().get(), + tensor_F.data().get(), + tensor_D.data().get(), + // tensor_Z.data().get(), + options.layoutY_transformed(), + options.layoutF_transformed(), + options.layoutD_transformed(), + // options.layoutZ_transformed() + }, + hw_info + }; + + SsdOperation op; + + size_t workspace_size = SsdOperation::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + cutlass::Status status = op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "This kernel is not supported. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + status = op.initialize(arguments, workspace.get()); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + + cudaError_t result; + // may be used uninitialized + cudaEvent_t start; + cudaEvent_t end; + cudaEventCreate(&start); + cudaEventCreate(&end); + + // warm up + if (options.measure) { + for (int i = 0; i < options.warmups; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + } + result = cudaEventRecord(start); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + // Run + for (int i = 0; i < options.iterations; i++) { + status = op.run(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + return false; + } + } + result = cudaEventRecord(end); + if (result != cudaSuccess) { + std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(result) << std::endl; + return false; + } + + float runtime_ms = 0; + result = cudaEventElapsedTime(&runtime_ms, start, end); + if (result != cudaSuccess) { + std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl; + return false; + } + runtime_ms /= static_cast(options.iterations); + + if (options.verbose) { + printf("[iters = %d, warmups = %d] ssd kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms); + printf("smem size = %d\n", SsdOperation::Kernel::SharedStorageSize); + } + // Matrix + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + auto mY_ref_0 = cute::make_tensor(tensor_Y_ref_0.data().get(), options.layoutY()); + auto mY_ref_1 = cute::make_tensor(tensor_Y_ref_1.data().get(), options.layoutY()); + auto mY_res = cute::make_tensor(tensor_Y.data().get(), options.layoutY()); + auto mF_ref_0 = cute::make_tensor(tensor_F_ref_0.data().get(), options.layoutF()); + auto mF_ref_1 = cute::make_tensor(tensor_F_ref_1.data().get(), options.layoutF()); + auto mF_res = cute::make_tensor(tensor_F.data().get(), options.layoutF()); + auto mX = cute::make_tensor(tensor_X.data().get(), options.layoutX()); + auto mB = cute::make_tensor(tensor_B.data().get(), options.layoutB()); + auto mC = cute::make_tensor(tensor_C.data().get(), options.layoutC()); + auto mD = cute::make_tensor(tensor_D.data().get(), options.layoutD()); + auto mZ = cute::make_tensor(tensor_Z.data().get(), options.layoutZ()); + auto mDelta = cute::make_tensor(tensor_Delta.data().get(), options.layoutDelta()); + auto mDeltaA = cute::make_tensor(tensor_DeltaA.data().get(), options.layoutDeltaA()); + + // Reference Device kernel + if (options.verify) { + ssd_reference( + mY_ref_1, + mF_ref_1, + mX, + mDelta, + mDeltaA, + mB, + mC, + mD, + mZ, + options + ); + } + + bool passed = true; + + if (options.verify) { + printf("[TensorY]verifying...\n"); + passed &= compare_reference<5>(mY_ref_1, mY_res); + printf("[TensorF]verifying...\n"); + passed &= compare_reference<4>(mF_ref_1, mF_res); + } + + return passed; + } + + template< + int TensorDim, + class Engine, class Layout + > + static constexpr bool + compare_reference( + cute::Tensor const& reference, + cute::Tensor const& computed, + float epsilon = 0.05f) { + if (size(reference) != size(computed)) { + return false; + } + + bool passed = true; + if (epsilon == 0.0f) { + // fast refcheck w/o epsilon + for (size_t i = 0; i < size_t(size(reference)); ++i) { + if (reference(i) != computed(i)) { + passed = false; + printf("[%llu] %f, %f\n", static_cast(i), + float(reference(i)), float(computed(i))); + break; + } + } + } + else { + // refcheck with epsilon + for (size_t i = 0; i < size_t(size(reference)); ++i) { + auto ref = static_cast(reference(i)); + auto act = static_cast(computed(i)); + auto abs_error = std::abs(act - ref); + auto rel_error = abs_error / (std::max(std::abs(act), std::abs(ref)) + 0.00001f); + if (std::isnan(abs_error) || std::isnan(rel_error) || + std::min(rel_error, abs_error) > epsilon) { + passed = false; + printf("[%llu] %f, %f\n", static_cast(i), + float(reference(i)), float(computed(i))); + break; + } + } + } + if (not passed) { + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + auto m = cute::shape<2>(reference); + auto n = cute::shape(reference); + printf("reference:\n"); + for (int mi = 0; mi < m; ++mi) { + for (int ni = 0; ni < n; ++ni) { + if constexpr (TensorDim == 5) { + printf("%.4f ", static_cast(reference(0,0,mi,0,ni))); + } + else { + printf("%.4f ", static_cast(reference(0,0,mi,ni))); + } + } + printf("\n"); + } + printf("\n"); + printf("computed:\n"); + for (int mi = 0; mi < m; ++mi) { + for (int ni = 0; ni < n; ++ni) { + if constexpr (TensorDim == 5) { + printf("%.4f ", static_cast(computed(0,0,mi,0,ni))); + } + else { + printf("%.4f ", static_cast(computed(0,0,mi,ni))); + } + } + printf("\n"); + } + printf("\n"); + } + return passed; + } +}; + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +int main(int argc, char const **args) { + + cudaDeviceProp props; + + cudaError_t error = cudaGetDeviceProperties(&props, 0); + if (error != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl; + return -1; + } + + if (__CUDACC_VER_MAJOR__ < 12 || props.major < 10) { + std::cout + << "This example requires a GPU of NVIDIA's Blackwell Architecture or " + << "later (compute capability 100 or greater) and CUDA 12.0 or greater.\n"; + return 0; + } + else if (__CUDACC_VER_MAJOR__ < 12 || (props.major != 10 || props.minor != 0)) { + std::cout + << "This example requires a GPU of NVIDIA's Blackwell Architecture " + << "(compute capability 100) and CUDA 12.0 or greater.\n"; + return 0; + } + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + + // + // Parse options + // + + Options options; + + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + if (options.error) { + std::cerr << "Aborting execution." << std::endl; + return -1; + } + + // Execute kernel + + printf("start testing....\n"); + + // The KernelHardwareInfo struct holds the number of SMs 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.device_id = 0; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + // Check Device/Host ref kernel + TestBed testbed{}; + bool passed = testbed.run(options, hw_info); + + if (passed) { + printf("everything is ok.\n"); + } + else { + printf("something is wrong!!!!!\n"); + } + +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + + return 0; +} diff --git a/test/utils/device_info.py b/examples/112_blackwell_ssd/CMakeLists.txt similarity index 61% rename from test/utils/device_info.py rename to examples/112_blackwell_ssd/CMakeLists.txt index 00294399..0a924e07 100644 --- a/test/utils/device_info.py +++ b/examples/112_blackwell_ssd/CMakeLists.txt @@ -1,20 +1,20 @@ # Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 @@ -26,34 +26,20 @@ # 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. +if (CUTLASS_NVCC_ARCHS MATCHES 100a) -def _get_device_compute_capability(): - try: - import cuda.bindings.driver as drv - from cuda.bindings.driver import CUdevice_attribute as dev_attr +set_property( + SOURCE 112_blackwell_ssd.cu + PROPERTY COMPILE_FLAGS "--use_fast_math" + ) - def drv_api(api_name, *args): - ret_code, *result = getattr(drv, api_name)(*args) - if ret_code: - raise ValueError(f"CUDA error: {ret_code}") - return result[0] if len(result) == 1 else result +cutlass_example_add_executable( + 112_blackwell_ssd + 112_blackwell_ssd.cu + ) - drv_api("cuInit", 0) - device = drv_api("cuDeviceGet", 0) - major = drv_api( - "cuDeviceGetAttribute", - dev_attr.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, - device, - ) - minor = drv_api( - "cuDeviceGetAttribute", - dev_attr.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, - device, - ) - return f"{major}{minor}" - except Exception as e: - print(f"Failed to get CUDA compute capability: {e}") - return None +if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang"))) +endif() -compute_capability = _get_device_compute_capability() +endif() diff --git a/examples/112_blackwell_ssd/README.md b/examples/112_blackwell_ssd/README.md new file mode 100644 index 00000000..0768503a --- /dev/null +++ b/examples/112_blackwell_ssd/README.md @@ -0,0 +1,65 @@ +# NVIDIA Blackwell SSD (State Space Decomposition) CUDA Example + +## Overview + +This example demonstrates the implementation of State Space Decomposition (SSD) operations on NVIDIA's Blackwell GPU architecture. It showcases the use of CUTLASS library components for high-performance tensor computations that efficiently leverage Blackwell's advanced hardware capabilities. + +## System Requirements ++ NVIDIA GPU with Blackwell Architecture (compute capability 10.0) ++ CUDA Toolkit 12.8 or newer ++ C++17 compatible compiler + +## Build the example +Follow the cutlass example building. + +## Command Line Options +The example supports the following command line options: +--help: Display the usage statement +--iterations=: Number of iterations for benchmarking (default: 1) +--without_verify: Skip result verification +--verbose: Print execution time per kernel +--G=: Group size (default: 2) +--B=: Batch size (default: 3) +--E=: Expanded factor (default: 2) +--H=: Number of heads (default: 2) + +## Limitation ++ Only support LxDxN = 128x64x128 ++ Require all TMEM at once and no more cta on the same SM + +## Performance ++ Limited by the SEGSUM part. ++ Low MMA utils ++ ALU bound. + +# Copyright + +Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. +``` diff --git a/examples/112_blackwell_ssd/collective/sm100_ssd_epilogue.hpp b/examples/112_blackwell_ssd/collective/sm100_ssd_epilogue.hpp new file mode 100644 index 00000000..9bad8faa --- /dev/null +++ b/examples/112_blackwell_ssd/collective/sm100_ssd_epilogue.hpp @@ -0,0 +1,535 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/epilogue/collective/collective_builder.hpp" + +namespace cutlass::ssd::collective { +using namespace cute; + +template< + class ElementAcc_, + class Element_, + class ElementDA_, + class TileShape_, + class EpilogueTile_, + class SmemLayoutY_, + class SmemLayoutP_, + class SmemLayoutX_, + int StagesInput, + int StagesOutput, + bool HasScaleD_, + bool HasBlockScaleD_> +struct SsdEpilogue { + + using TileShape = TileShape_; + using ElementAcc = ElementAcc_; + using Element = Element_; + using ElementD = Element_; + using ElementDA = ElementDA_; + using ElementY = Element_; + using ElementP = Element_; + using EpilogueTile = EpilogueTile_; + using SmemLayoutY = SmemLayoutY_; + using SmemLayoutP = SmemLayoutP_; + using SmemLayoutX = SmemLayoutX_; + + static constexpr bool HasScaleD = HasScaleD_; + static constexpr bool HasBlockScaleD = HasBlockScaleD_; + + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + static constexpr auto L = get<0>(TileShape{}); + static constexpr auto D = get<1>(TileShape{}); + static constexpr auto N = get<2>(TileShape{}); + + constexpr static int ThreadCount = 128; + constexpr static size_t SmemAlignmentY = cutlass::detail::alignment_for_swizzle(SmemLayoutY{}); + + struct CollectiveStorage { + alignas(SmemAlignmentY) ArrayEngine> smem_y; + }; + + // TMA pipeline for storing D + using StorePipeline = cutlass::PipelineTmaStore; + using StorePipelineState = cutlass::PipelineState; + + using StorePPipeline = cutlass::PipelineTmaStore<1>; + using StorePPipelineState = cutlass::PipelineState<1>; + + using CooperatePipeline = cutlass::PipelineAsync; + using CooperatePipelineState = cutlass::PipelineState; + + using EpiloadPipelineD = cutlass::PipelineTmaAsync; + static constexpr int kEpiloadDBytes = HasBlockScaleD ? D * sizeof(ElementD) : 0; + + struct SharedStorage { + using TensorStorage = CollectiveStorage; + TensorStorage tensors; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + + using LayoutY = decltype(make_layout(make_shape(L, D, int32_t(0), int32_t(0)), + make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,D,C,B) + using LayoutP = decltype(make_layout(make_shape(D, N, int32_t(0)), make_stride(N, _1{}, D*N))); // (D,N,B) + using LayoutD_2D = decltype(make_layout(make_shape(D, int32_t(0)), make_stride(_1{}, D))); // (D,EH) + using LayoutD_1D = decltype(make_layout(make_shape(_1{}, int32_t(0)), make_stride(_0{}, _1{}))); // (D,EH) + using LayoutD = cute::conditional_t< + HasBlockScaleD, + LayoutD_2D, + LayoutD_1D + >; + + struct Arguments { + ElementY* ptr_Y{nullptr}; + ElementP* ptr_P{nullptr}; + const ElementD* ptr_D{nullptr}; + LayoutY layout_Y{}; + LayoutP layout_P{}; + LayoutD layout_D{}; + }; + + using StrideY = cute::tuple<_1, int, int, int>; // (L,D,C,B) + using StrideP = cute::tuple; // (D,N,B) + + using CopyOpS2G = SM90_TMA_STORE; + + struct Params { + using TMA_Y = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideY{}, int32_t(0)), StrideY{}), + take<0,2>(SmemLayoutY{}), + EpilogueTile{}, + _1{})); + using TMA_P = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideP{}, int32_t(0)), StrideP{}), + take<0,3>(SmemLayoutP{}), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), // (D,N) + _1{})); + using TensorD = decltype(make_tensor( + make_gmem_ptr(static_cast(nullptr)), + LayoutD{})); + + TMA_Y tma_store_y; + TMA_P tma_store_p; + TensorD tensor_d; + }; + + template + static Params to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace = nullptr) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + auto tensor_y = make_tensor(make_gmem_ptr(args.ptr_Y), args.layout_Y); + auto tensor_p = make_tensor(make_gmem_ptr(args.ptr_P), args.layout_P); + auto tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), args.layout_D); + + auto tma_store_y = make_tma_copy_C_sm90( + CopyOpS2G{}, + tensor_y, + take<0,2>(SmemLayoutY{}), + EpilogueTile{}); + auto tma_store_p = make_tma_copy_C_sm90( + CopyOpS2G{}, + tensor_p, + take<0,2>(SmemLayoutP{}), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}))); + return Params{ + tma_store_y, + tma_store_p, + tensor_d + }; + } + + template< + class Params, class ProblemShape, + class EpiloadPipeline, class PipelineState, + class TensorStorage + > + CUTLASS_DEVICE + void load_d( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + EpiloadPipeline& pipeline, PipelineState& pipeline_d_producer_state, + TensorStorage& shared_tensors) { + + if constexpr (HasBlockScaleD) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + + auto& gD = params.tensor_d; + + ElementD* ptr_d = shared_tensors.smem_d.data(); + auto smem_layout = make_layout(make_shape(get<1>(TileShape{}), Int{})); // (D,) + auto sD = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_d), smem_layout)); + + auto bulk_atom = Copy_Atom{}; + + int write_stage = pipeline_d_producer_state.index(); + + // LOCK pipeline_state for _writing_ + pipeline.producer_acquire(pipeline_d_producer_state); + + using BarrierType = typename EpiloadPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_d_producer_state); + copy(bulk_atom.with(*tma_barrier), gD(_,blk_coord), sD(_,write_stage)); + + // Advance pipeline_state + ++pipeline_d_producer_state; + } + } + + } + + template< + class Params, class ProblemShape, + class MainloopPipelineIntra, class PipelineStateIntra, + class MainloopPipelineAcc, class PipelineStateAcc, + class MainloopPipelineDelta, class PipelineStateDelta, + class MainloopPipelineX, class PipelineStateX, + class EpiloadPipelineD, class PipelineStateD, + class FragmentC_Intra_1, class FragmentC_Intra_2, + class FragmentC_Inter_1, class FragmentC_Inter_2, + class MainloopStorage, class EpilogueStorage + > + CUTLASS_DEVICE + auto store( + int& chunk, int const& blk_coord, int const& blk_coord_eh, + Params const& params, ProblemShape const& problem_size, + MainloopPipelineIntra& pipeline_intra, PipelineStateIntra& pipeline_intra_consumer_state, + MainloopPipelineAcc& pipeline_acc, PipelineStateAcc& pipeline_acc_consumer_state, + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_delta_consumer_state, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_x_consumer_state, + EpiloadPipelineD& pipeline_d, PipelineStateD& pipeline_d_producer_state, + StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state, + cute::tuple& acc_intra, + cute::tuple& acc_inter, + MainloopStorage& mainloop_tensors, EpilogueStorage& epilogue_tensors, + bool is_first_iteration) { + + using CopyOpT2R = SM100_TMEM_LOAD_16dp256b4x; + using CopyOpR2S = SM90_U16x8_STSM_T; + using CopyOpS2R = SM75_U16x8_LDSM_T; + + // Epilogue + + int thread_idx = int(threadIdx.x % 128); + int warp_idx = thread_idx / 32; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mY_mn = params.tma_store_y.get_tma_tensor(make_shape(L,D,C,B*EH)); + Tensor gY_mn = local_tile(mY_mn, take<0,2>(TileShape{}), make_coord(_,_,_))(_,_,_0{},_0{},chunk,blk_coord); + + auto ptr_sY = epilogue_tensors.smem_y.begin(); + Tensor sY_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sY), SmemLayoutY{})); + auto ptr_sX = mainloop_tensors.smem_x.data(); + Tensor sX_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sX), SmemLayoutX{})); + + auto ts0 = size<0>(TileShape{}); + auto ts1 = size<1>(TileShape{}); + Layout col_layout = make_layout(make_shape ( ts0, ts1, Int{}), + make_stride(_1{}, _0{}, ts0)); + Tensor sDeltaA = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(mainloop_tensors.smem_delta_a.data()), col_layout)); + auto tDeltaA = sDeltaA(_,_,pipeline_delta_consumer_state.index()); + Layout row_layout = make_layout(make_shape ( ts0, ts1, Int{}), + make_stride(_0{}, _1{}, ts1)); + Tensor sD = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(mainloop_tensors.smem_d.data()), row_layout)); + auto tD = sD(_,_,pipeline_d_producer_state.index()); + + auto accumulator_intra_2 = get<1>(acc_intra); + auto accumulator_inter_2 = get<1>(acc_inter); + auto tIntra = accumulator_intra_2(make_coord(_,_),_0{},_0{},_0{}); + auto tInter = accumulator_inter_2(make_coord(_,_),_0{},_0{},_0{}); + Tensor tIntra_epi = flat_divide(tIntra, EpilogueTile{}); + Tensor tInter_epi = flat_divide(tInter, EpilogueTile{}); + Tensor tDeltaA_epi = flat_divide(tDeltaA, EpilogueTile{}); + Tensor tD_epi = flat_divide(tD, EpilogueTile{}); + Tensor gY_epi = flat_divide(gY_mn, EpilogueTile{}); + + TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tIntra_epi(_,_,_0{},_0{})); + ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_tIntra = thread_t2r.partition_S(tIntra_epi); + Tensor tTR_tInter = thread_t2r.partition_S(tInter_epi); + Tensor tTR_tDeltaA = thread_t2r.partition_D(tDeltaA_epi); + Tensor tTR_tD = thread_t2r.partition_D(tD_epi); + Tensor tTR_sY = thread_t2r.partition_D(sY_epi(_,_,_0{})); + Tensor tTR_rIntra = make_tensor(shape(thread_t2r.partition_D(gY_epi))); + // Tensor tTR_rInter = make_tensor(shape(thread_t2r.partition_D(gY_epi))); + + // Tensor tTR_rIntra = make_tensor(shape(tTR_sY)); + Tensor tTR_rInter = make_tensor(shape(tTR_sY)); + Tensor tTR_rY = make_tensor(shape(tTR_sY)); + Tensor tTR_rDeltaA = make_tensor(shape(tTR_sY)); + Tensor tTR_rD = make_tensor(shape(tTR_sY)); + Tensor tTR_rD_Acc = make_tensor(shape(tTR_sY)); + Tensor tTR_rX = make_tensor(shape(tTR_sY)); + Tensor tTR_rX_Acc = make_tensor(shape(tTR_sY)); + + // (t)hread-partition for (r)egister to (s)mem copy (tRS_) + TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_sY = thread_r2s.partition_D(sY_epi); + Tensor tRS_rY = thread_r2s.retile_S(tTR_rY); + Tensor tRS_rDeltaA = thread_r2s.retile_S(tTR_rDeltaA); + Tensor tRS_rD = thread_r2s.retile_S(tTR_rD); + Tensor tRS_rCompute = make_tensor(shape(tRS_rY)); + + TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sX = thread_s2r.partition_S(flat_divide(sX_epi, EpilogueTile{})); + Tensor tSR_rX = thread_s2r.retile_D(tTR_rX); + + constexpr int FragmentSize = 4; + // Tensor tTR_rIntra_frg = recast>(coalesce(tTR_rIntra)); + // Tensor tTR_rInter_frg = recast>(coalesce(tTR_rInter)); + Tensor tTR_rY_frg = recast>(coalesce(tTR_rY)); + Tensor tRS_rY_frg = recast>(coalesce(tRS_rY)); + Tensor tRS_rDeltaA_frg = recast>(coalesce(tRS_rDeltaA)); + Tensor tRS_rD_frg = recast>(coalesce(tRS_rD)); + Tensor tRS_rCompute_frg = recast>(coalesce(tRS_rCompute)); + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + ThrCopy thrblk_s2g = params.tma_store_y.get_slice(Int<0>{}); + Tensor bSG_sY = thrblk_s2g.partition_S(sY_epi); + Tensor bSG_gY = thrblk_s2g.partition_D(gY_epi); + + auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; + [[maybe_unused]] bool issue_smem_store = true; + [[maybe_unused]] bool issue_tma_store = warp_idx == 0; + + // The TMA store sequence for one subtile iteration + auto tma_store_fn = [&] (int epi_m, int epi_n) { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_y, bSG_sY(_,_,_,store_pipe_producer_state.index()), bSG_gY(_,_,_,epi_m,epi_n)); + } + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + }; + + // Require Acc + pipeline_intra.consumer_wait(pipeline_intra_consumer_state); + copy(tiled_t2r, tTR_tIntra, tTR_rIntra); + + cutlass::arch::fence_view_async_tmem_load(); + + pipeline_intra.consumer_release(pipeline_intra_consumer_state); + ++pipeline_intra_consumer_state; + + // Require Inter Acc + pipeline_acc.consumer_wait(pipeline_acc_consumer_state); + + CUTLASS_PRAGMA_UNROLL + for (int iter_n = 0; iter_n < size<3>(gY_epi); ++iter_n) { + CUTLASS_PRAGMA_UNROLL + for (int iter_m = 0; iter_m < size<2>(gY_epi); ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + + Tensor tTR_rIntra_mn = tTR_rIntra(_,_,_,epi_m,epi_n); + Tensor tTR_tInter_mn = tTR_tInter(_,_,_,epi_m,epi_n); + Tensor tTR_tDeltaA_mn = tTR_tDeltaA(_,_,_,epi_m,epi_n); + Tensor tTR_tD_mn = tTR_tD(_,_,_,epi_m,epi_n); + + copy(tiled_t2r, tTR_tInter_mn, tTR_rInter); + cutlass::arch::fence_view_async_tmem_load(); + + copy(tTR_tDeltaA_mn, tTR_rDeltaA); + + if constexpr (HasScaleD) { + copy(tiled_s2r, tSR_sX(_,_,_,epi_m,epi_n,pipeline_x_consumer_state.index()), tSR_rX); + type_convert(tTR_rX, tTR_rX_Acc); + } + + if constexpr (HasBlockScaleD) { + if (is_first_iteration) { + pipeline_d.consumer_wait(pipeline_d_producer_state); + } + copy(tTR_tD_mn, tTR_rD); + type_convert(tTR_rD, tTR_rD_Acc); + } + else if constexpr (HasScaleD) { + auto& gD = params.tensor_d; + auto value = static_cast(gD(_0{}, blk_coord_eh)); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tTR_rD_Acc); ++ii) { + tTR_rD_Acc(ii) = value; + } + } + // ? + synchronize(); + + NumericArrayConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tRS_rCompute); ++ii) { + tRS_rCompute(ii) = tTR_rIntra_mn(ii) + tTR_rInter(ii) * expf(tTR_rDeltaA(ii)); + if constexpr (HasScaleD) { + tRS_rCompute(ii) += tTR_rD_Acc(ii) * tTR_rX_Acc(ii); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int epi_v = 0; epi_v < size(tTR_rY_frg); ++epi_v) { + tRS_rY_frg(epi_v) = converter(tRS_rCompute_frg(epi_v)); + } + + copy(tiled_r2s, tRS_rY, tRS_sY(_,_,_,store_pipe_producer_state.index())); + + tma_store_fn(epi_m, epi_n); + } + } + + pipeline_acc.consumer_release(pipeline_acc_consumer_state); + ++pipeline_acc_consumer_state; + pipeline_x.consumer_release_from_threads(pipeline_x_consumer_state); + ++pipeline_x_consumer_state; + + pipeline_delta.consumer_release_from_threads(pipeline_delta_consumer_state); + ++pipeline_delta_consumer_state; + } + + template< + class Params, class ProblemShape, + class StorePipeline, class StorePipelineState, + class TensorStorage + > + CUTLASS_DEVICE + auto store_p( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state, + TensorStorage& shared_tensors) { + + int thread_idx = int(threadIdx.x % 128); + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mP_mn = params.tma_store_p.get_tma_tensor(make_shape(D,N,B*EH)); + Tensor gP_mn = local_tile(mP_mn, take<1,3>(TileShape{}), make_coord(_,_,blk_coord)); + auto gP_epi = gP_mn(_,_,_0{},_0{}); + + // Construct the corresponding pipelined smem tensors + auto ptr_sP = shared_tensors.smem_p.begin(); + Tensor sP_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sP), SmemLayoutP{}))(_,_,_0{}); // (EPI_TILE_M,EPI_TILE_N) + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + auto oprands = tma_partition(params.tma_store_p, Int<0>{}, Layout<_1>{}, + group_modes<0,2>(sP_epi), group_modes<0,2>(gP_epi)); // (TMA,k) and (TMA,PIPE) + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + auto bSG_gY = get<0>(oprands); + auto bSG_sY = get<1>(oprands); + +#if 0 + if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("mP_mn : ");print(mP_mn);print("\n"); + print("gP_mn : ");print(gP_mn);print("\n"); + print("sP_epi : ");print(sP_epi);print("\n"); + print("gP_epi : ");print(gP_epi);print("\n"); + print("bSG_sY : ");print(bSG_sY);print("\n"); + print("bSG_gY : ");print(bSG_gY);print("\n"); + } +#endif + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + // Use the reserved named barrier `streamkbarrier` since this kernel doesn't support streamk + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::StreamkBarrier0); }; + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0; + + auto tma_store_fn = [&] () { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_p, bSG_sY, bSG_gY); + } + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + // don't need pipeline + synchronize(); + }; + + tma_store_fn(); + } + + template< + class ElementSrc, class ElementDst, + class TensorSrc, class TensorDst + > + CUTLASS_DEVICE + auto type_convert( + TensorSrc& tS, + TensorDst& tD) { + + static constexpr int FragmentSize = 2; + NumericArrayConverter converter; + + auto tS_frg = recast>(tS); + auto tD_frg = recast>(tD); + + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tS_frg); ++ii) { + tD_frg(ii) = converter(tS_frg(ii)); + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/examples/112_blackwell_ssd/collective/sm100_ssd_gemm_tma_warpspecialized.hpp b/examples/112_blackwell_ssd/collective/sm100_ssd_gemm_tma_warpspecialized.hpp new file mode 100644 index 00000000..43a11d8d --- /dev/null +++ b/examples/112_blackwell_ssd/collective/sm100_ssd_gemm_tma_warpspecialized.hpp @@ -0,0 +1,1145 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/collective/collective_builder.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/gemm/gemm.h" + +#include "../utils/pipeline.h" + +namespace cutlass::ssd::collective { + +using namespace cute; + +template< + class Element_, + class ElementDA_, + class ElementAcc_, + class ElementD_, + class TileShape_, + int StagesInput_, + int StagesOutput_, + class TiledMmaIntra1_, + class TiledMmaIntra2_, + class TiledMmaInter1_, + class TiledMmaInter2_, + class SmemLayoutX_, + class SmemLayoutB_, + class SmemLayoutC_, + class SmemLayoutP_, + class SmemLayoutBT_, + class SmemLayoutPT_, + class SmemLayoutQ_, + class TmemLayoutB_, + class TmemLayoutQ_ +> +struct SsdMainloopTmaWarpSpecialized { + + + using Element = Element_; + using ElementDA = ElementDA_; + using ElementAcc = ElementAcc_; + static constexpr int StagesInput = StagesInput_; + static constexpr int StagesOutput = StagesOutput_; + + using TileShape = TileShape_; // (L,D,N) + + // avoid "warning #3357-D: capturing structured bindings is a C++20 feature" + static constexpr auto L = get<0>(TileShape{}); + static constexpr auto D = get<1>(TileShape{}); + static constexpr auto N = get<2>(TileShape{}); + + using TileShapeIntraBMM1 = decltype(make_shape(L, L, N)); // (L,L,N) + using TileShapeIntraBMM2 = decltype(make_shape(L, D, N)); // (L,D,L) + using TileShapeInterBMM1 = decltype(make_shape(N, D, L)); // (N,D,L) + using TileShapeInterBMM2 = decltype(make_shape(L, D, N)); // (L,D,N) + + // 16B alignment lets us use TMA + static constexpr int Alignment = 16 / sizeof(Element); + + using TiledMmaIntra1 = TiledMmaIntra1_; + using TiledMmaIntra2 = TiledMmaIntra2_; + using TiledMmaInter1 = TiledMmaInter1_; + using TiledMmaInter2 = TiledMmaInter2_; + using SmemLayoutX = SmemLayoutX_; + using SmemLayoutB = SmemLayoutB_; + using SmemLayoutC = SmemLayoutC_; + using SmemLayoutP = SmemLayoutP_; + using SmemLayoutBT = SmemLayoutBT_; + using SmemLayoutPT = SmemLayoutPT_; + using TmemLayoutB = TmemLayoutB_; + using TmemLayoutQ = TmemLayoutQ_; + using SmemLayoutQ = SmemLayoutQ_; + + using ClusterShape = Shape<_1, _1, _1>; + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout( + ClusterShape{}), + make_tile(typename TiledMmaIntra1::AtomThrID{}))); + + using AtomThrShapeMNK = Shape(typename TiledMmaIntra1::ThrLayoutVMNK{})), _1, _1>; + + using MainloopPipelineX = cutlass::PipelineTmaMultiConsumersAsync< + StagesInput, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineB = cutlass::PipelineTmaMultiConsumersAsync< + StagesInput, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineC = cutlass::PipelineTmaMultiConsumersAsync< + StagesInput, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineDelta = cutlass::PipelineTmaMultiConsumersAsync< + StagesInput, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineIntra = cutlass::PipelineUmmaAsync<1, AtomThrShapeMNK>; + using MainloopPipelineInter = cutlass::PipelineUmmaAsync<1, AtomThrShapeMNK>; + using AccumulatorPipeline = cutlass::PipelineUmmaAsync<1, AtomThrShapeMNK>; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + + static constexpr int kXLoadBytes = size(SmemLayoutX{}(_,_,_,_0{})) * sizeof(Element); + static constexpr int kBLoadBytes = size(SmemLayoutB{}(_,_,_,_0{})) * sizeof(Element); + static constexpr int kCLoadBytes = size(SmemLayoutC{}(_,_,_,_0{})) * sizeof(Element); + static constexpr int kDeltaLoadBytes = get<0>(TileShape{}) * sizeof(Element); + static constexpr int kDeltaALoadBytes = get<0>(TileShape{}) * sizeof(ElementDA); + + struct TensorStorage : cute::aligned_struct<128, _0> { + cute::array_aligned> smem_x; + cute::array_aligned> smem_b; + cute::array_aligned> smem_c; + cute::array_aligned> smem_p; + cute::array_aligned(TileShape{}) * StagesInput> smem_delta; + cute::array_aligned(TileShape{}) * StagesInput> smem_delta_a; + cute::array_aligned(TileShape{}) * StagesInput> smem_d; + }; + + // TMEM management (for LDN = 128x64x128) + static constexpr uint32_t tmem_intra_1_Acc = 0; + static constexpr uint32_t tmem_intra_2_A = tmem_intra_1_Acc + 128; + static constexpr uint32_t tmem_intra_2_Acc = tmem_intra_2_A + 64; + static constexpr uint32_t tmem_inter_1_A = tmem_intra_2_Acc + 64; + static constexpr uint32_t tmem_inter_1_Acc = tmem_inter_1_A + 64; + static constexpr uint32_t tmem_inter_2_Acc = tmem_inter_1_Acc + 64; + + using StrideX = cute::tuple; // (D,L,C,B) + using StrideB = cute::tuple<_1, int, int, int>; // (L,N,C,B) + using StrideC = cute::tuple<_1, int, int, int>; // (L,N,C,B) + using StrideDelta = cute::tuple<_1, int, int>; // (L,C,B) + + using LayoutX = decltype(make_layout(make_shape(D, L, int32_t(0), int32_t(0)), make_stride(int32_t(0), _1{}, L, int32_t(0)))); // (D,L,C,B) + using LayoutB = decltype(make_layout(make_shape(L, N, int32_t(0), int32_t(0)), make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,N,C,B) + using LayoutC = decltype(make_layout(make_shape(L, N, int32_t(0), int32_t(0)), make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,N,C,B) + using LayoutDelta = decltype(make_layout(make_shape(L, int32_t(0), int32_t(0)), make_stride(_1{}, L, int32_t(0)))); // (L,C,B) + + // Not support Mcast + using GmemTiledCopyX = cute::SM90_TMA_LOAD; + using GmemTiledCopyB = cute::SM90_TMA_LOAD; + using GmemTiledCopyC = cute::SM90_TMA_LOAD; + + struct Arguments { + const Element* ptr_X{nullptr}; + const ElementDA* ptr_DeltaA{nullptr}; + const Element* ptr_Delta{nullptr}; + const Element* ptr_B{nullptr}; + const Element* ptr_C{nullptr}; + LayoutX layout_X{}; + LayoutB layout_B{}; + LayoutC layout_C{}; + LayoutDelta layout_Delta{}; + }; + + template + static constexpr auto + get_tma_load_x_instance(TensorX const& tensor_x, ClusterShapeVMNK const& cluster_shape_vmnk) { + return make_tma_atom_B_sm100( + GmemTiledCopyX{}, + tensor_x, + SmemLayoutX{}(_,_,_,cute::Int<0>{}), + TileShapeIntraBMM2{}, + TiledMmaIntra2{}, + cluster_shape_vmnk); + } + + template + static constexpr auto + get_tma_load_b_instance(TensorB const& tensor_b, ClusterShapeVMNK const& cluster_shape_vmnk) { + return make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShapeIntraBMM1{}, + TiledMmaIntra1{}, + cluster_shape_vmnk); + } + + template + static constexpr auto + get_tma_load_c_instance(TensorC const& tensor_c, ClusterShapeVMNK const& cluster_shape_vmnk) { + return make_tma_atom_A_sm100( + GmemTiledCopyC{}, + tensor_c, + SmemLayoutC{}(_,_,_,cute::Int<0>{}), + TileShapeIntraBMM1{}, + TiledMmaIntra1{}, + cluster_shape_vmnk); + } + + struct Params { + using TMA_X = decltype(get_tma_load_x_instance( + make_tensor( + make_gmem_ptr(recast_ptr(nullptr)), + LayoutX{}), + ClusterLayout_VMNK{})); + using TMA_B = decltype(get_tma_load_b_instance( + make_tensor( + make_gmem_ptr(recast_ptr(nullptr)), + LayoutB{}), + ClusterLayout_VMNK{})); + using TMA_C = decltype(get_tma_load_c_instance( + make_tensor( + make_gmem_ptr(recast_ptr(nullptr)), + LayoutC{}), + ClusterLayout_VMNK{})); + using TensorDelta = decltype(make_tensor( + make_gmem_ptr(static_cast(nullptr)), + LayoutDelta{})); + using TensorDeltaA = decltype(make_tensor( + make_gmem_ptr(static_cast(nullptr)), + LayoutDelta{})); + + TMA_X tma_load_x; + TMA_B tma_load_b; + TMA_C tma_load_c; + TensorDelta tensor_delta; + TensorDeltaA tensor_delta_a; + }; + + template + static bool can_implement(ProblemShape const& problem_size, Arguments const& args) { + return true; + } + + template + static Params to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor tensor_x = make_tensor(make_gmem_ptr(args.ptr_X), args.layout_X); + Tensor tensor_b = make_tensor(make_gmem_ptr(args.ptr_B), args.layout_B); + Tensor tensor_c = make_tensor(make_gmem_ptr(args.ptr_C), args.layout_C); + Tensor tensor_delta = make_tensor(make_gmem_ptr(args.ptr_Delta), args.layout_Delta); + Tensor tensor_delta_a = make_tensor(make_gmem_ptr(args.ptr_DeltaA), args.layout_Delta); + +#if 0 + print("tensor_x : ");print(tensor_x);print("\n"); + print("smem_X : ");print(SmemLayoutX{});print("\n"); + print("tensor_b : ");print(tensor_b);print("\n"); + print("smem_b : ");print(SmemLayoutB{});print("\n"); + print("tensor_c : ");print(tensor_c);print("\n"); + print("smem_c : ");print(SmemLayoutC{});print("\n"); + print("tensor_delta : ");print(tensor_delta);print("\n"); + print("tensor_delta_a: ");print(tensor_delta_a);print("\n"); + print("tile : ");print(TileShape{});print("\n"); + print("tileIntra1 : ");print(TileShapeIntraBMM1{});print("\n"); + print("tileIntra2 : ");print(TileShapeIntraBMM2{});print("\n"); + print("tileInter1 : ");print(TileShapeInterBMM1{});print("\n"); + print("tileInter2 : ");print(TileShapeInterBMM2{});print("\n"); +#endif + + typename Params::TMA_X tma_load_x = get_tma_load_x_instance( + tensor_x, + ClusterLayout_VMNK{}); + typename Params::TMA_B tma_load_b = get_tma_load_b_instance( + tensor_b, + ClusterLayout_VMNK{}); + typename Params::TMA_C tma_load_c = get_tma_load_c_instance( + tensor_c, + ClusterLayout_VMNK{}); + + return Params{ + tma_load_x, + tma_load_b, + tma_load_c, + tensor_delta, + tensor_delta_a + }; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { + cute::prefetch_tma_descriptor(params.tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_c.get_tma_descriptor()); + } + + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_x_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mX_mkl = params.tma_load_x.get_tma_tensor(make_shape(D,L,C,EH*B)); + Tensor gX_mkl = local_tile(mX_mkl, TileShapeIntraBMM2{}, make_coord(_,_,_), Step{}); // (BLK_M,BLK_K,m,k,l) + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("mX_mkl : ");print(mX_mkl);print("\n"); + print("gX_mkl : ");print(gX_mkl);print("\n"); + print("tma_desc_ : ");print(params.tma_load_x.tma_desc_);print("\n"); + } +#endif + return make_tuple(gX_mkl); + } + + template< + class Params, class ProblemShape, + class MainloopPipelineX, class PipelineStateX, + class MainloopPipelineDelta, class PipelineStateDelta, + class GTensorX, + class TensorStorage + > + CUTLASS_DEVICE + void load_x_delta( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_x_producer_state, + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_delta_producer_state, + cute::tuple const& load_inputs, + TensorStorage& shared_tensors) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + + auto& gDelta = params.tensor_delta; + auto& gDeltaA = params.tensor_delta_a; + + Element* ptr_delta = shared_tensors.smem_delta.data(); + ElementDA* ptr_delta_a = shared_tensors.smem_delta_a.data(); + auto smem_layout = make_layout(make_shape(get<0>(TileShape{}), Int{})); + auto sDelta = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_delta), smem_layout)); + auto sDeltaA = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_delta_a), smem_layout)); + + auto bulk_atom_dt = Copy_Atom{}; + auto bulk_atom_dA = Copy_Atom{}; + + Tensor sX_ = make_tensor(make_smem_ptr(shared_tensors.smem_x.data()), SmemLayoutX{}); // (BLK_M,BLK_K,PIPE) + Tensor sX = as_position_independent_swizzle_tensor(sX_); // (BLK_M,BLK_K,PIPE) + + ThrMMA cta_mma = TiledMmaIntra2{}.get_slice(blockIdx.x % size(typename TiledMmaIntra2::AtomThrID{})); + + Tensor gX_mkl = get<0>(load_inputs); + Tensor tCgX_mk = cta_mma.partition_B(gX_mkl); + + auto [tXgX_mk, tXsX] = tma_partition(params.tma_load_x, + group_modes<0,3>(sX), group_modes<0,3>(tCgX_mk)); + auto tXgX = tXgX_mk(_,_0{},_0{},_,blk_coord); + + // Disable multicast + uint16_t mcast_mask_x = 0; + + auto chunk = get<3>(problem_size); + // Mainloop + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) { + int write_stage_x = pipeline_x_producer_state.index(); + int write_stage_delta = pipeline_delta_producer_state.index(); + + using BarrierTypeX = typename MainloopPipelineX::ProducerBarrierType; + BarrierTypeX* tma_barrier_x = pipeline_x.producer_get_barrier(pipeline_x_producer_state); + using BarrierTypeDelta = typename MainloopPipelineDelta::ProducerBarrierType; + BarrierTypeDelta* tma_barrier_delta = pipeline_delta.producer_get_barrier(pipeline_delta_producer_state); + + pipeline_x.producer_acquire(pipeline_x_producer_state); + + copy(params.tma_load_x.with(*tma_barrier_x, mcast_mask_x), tXgX(_,chunk_idx), tXsX(_,write_stage_x)); + + pipeline_delta.producer_acquire(pipeline_delta_producer_state); + + copy(bulk_atom_dt.with(*tma_barrier_delta), gDelta(_,chunk_idx,blk_coord), sDelta(_,write_stage_delta)); + copy(bulk_atom_dA.with(*tma_barrier_delta), gDeltaA(_,chunk_idx,blk_coord), sDeltaA(_,write_stage_delta)); + + // Advance pipeline_state + ++pipeline_x_producer_state; + ++pipeline_delta_producer_state; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + template< + class MainloopPipeline, class PipelineState + > + CUTLASS_DEVICE void + load_x_tail(MainloopPipeline pipeline, PipelineState pipeline_state) { + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(pipeline_state); + } + } + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_b_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mB_mkl = params.tma_load_b.get_tma_tensor(make_shape(L,N,C,G*B)); + Tensor gB_mkl = local_tile(mB_mkl, TileShapeIntraBMM1{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) + + return make_tuple(gB_mkl); + } + + template < + class Params, class ProblemShape + > + CUTLASS_DEVICE + auto load_c_init(Params const& params, ProblemShape const& problem_size) { + + using X = Underscore; + auto [G, B, EH, C, L, D, N] = problem_size; + + Tensor mC_mkl = params.tma_load_c.get_tma_tensor(make_shape(L,N,C,G*B)); + Tensor gC_mkl = local_tile(mC_mkl, TileShapeIntraBMM1{}, make_coord(_,_,_), Step{}); // (BLK_M,BLK_K,n,k,l) + + return make_tuple(gC_mkl); + } + + template< + class Params, class ProblemShape, + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineC, class PipelineStateC, + class GTensorB, class GTensorC, + class TensorStorage + > + CUTLASS_DEVICE + void load_b_c( + int const& blk_coord, Params const& params, ProblemShape const& problem_size, + MainloopPipelineB& pipeline_b, PipelineStateB& pipeline_b_producer_state, + MainloopPipelineC& pipeline_c, PipelineStateC& pipeline_c_producer_state, + cute::tuple const& load_inputs_b, cute::tuple const& load_inputs_c, + TensorStorage& shared_tensors) { + + int lane_predicate = cute::elect_one_sync(); + + if (lane_predicate) { + Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_b.data()), SmemLayoutB{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE) + Tensor sC_ = make_tensor(make_smem_ptr(shared_tensors.smem_c.data()), SmemLayoutC{}); // (BLK_M,BLK_K,PIPE) + Tensor sC = as_position_independent_swizzle_tensor(sC_); // (BLK_M,BLK_K,PIPE) + + // + // Prepare the TMA loads for B&C + // + + Tensor gB_mkl = get<0>(load_inputs_b); + Tensor gC_mkl = get<0>(load_inputs_c); + + ThrMMA cta_mma = TiledMmaIntra1{}.get_slice(blockIdx.x % size(typename TiledMmaIntra1::AtomThrID{})); + + Tensor tCgC_mk = cta_mma.partition_A(gC_mkl); + Tensor tCgB_mk = cta_mma.partition_B(gB_mkl); + + auto [tCgC_, tCsC] = tma_partition(params.tma_load_c, + group_modes<0,3>(sC), group_modes<0,3>(tCgC_mk)); + auto tCgC = tCgC_(_,_0{},_0{},_,blk_coord); + + auto [tBgB_, tBsB] = tma_partition(params.tma_load_b, + group_modes<0,3>(sB), group_modes<0,3>(tCgB_mk)); + auto tBgB = tBgB_(_,_0{},_0{},_,blk_coord); + +#if 0 + if (threadIdx.x % 32 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) { + print("gC : ");print(gC);print("\n"); + print("tCgC : ");print(tCgC);print("\n"); + print("tCsC : ");print(tCsC);print("\n"); + } +#endif + // no multicast + // Disable multicast + uint16_t mcast_mask_b = 0; + uint16_t mcast_mask_c = 0; + + auto chunk = get<3>(problem_size); + // Mainloop + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) { + int write_stage_b = pipeline_b_producer_state.index(); + int write_stage_c = pipeline_c_producer_state.index(); + + using BarrierTypeB = typename MainloopPipelineB::ProducerBarrierType; + BarrierTypeB* tma_barrier_b = pipeline_b.producer_get_barrier(pipeline_b_producer_state); + using BarrierTypeC = typename MainloopPipelineC::ProducerBarrierType; + BarrierTypeC* tma_barrier_c = pipeline_c.producer_get_barrier(pipeline_c_producer_state); + + pipeline_b.producer_acquire(pipeline_b_producer_state); + + copy(params.tma_load_b.with(*tma_barrier_b, mcast_mask_b), tBgB(_,chunk_idx), tBsB(_,write_stage_b)); + + pipeline_c.producer_acquire(pipeline_c_producer_state); + + copy(params.tma_load_c.with(*tma_barrier_c, mcast_mask_c), tCgC(_,chunk_idx), tCsC(_,write_stage_c)); + + // Advance pipeline_state + ++pipeline_b_producer_state; + ++pipeline_c_producer_state; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + template< + class MainloopPipeline, class PipelineState + > + CUTLASS_DEVICE void + load_b_c_tail(MainloopPipeline pipeline, PipelineState pipeline_state) { + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(pipeline_state); + } + } + + template< + class TiledMma, + class FragmentA, class FragmentB, + class FrgEngine, class FrgLayout + > + CUTLASS_DEVICE + auto mma( + int read_A_stage, int read_B_stage, + cute::tuple const& mma_inputs, + cute::Tensor& accumulators) { + auto [tiled_mma, tCrA, tCrB] = mma_inputs; + tiled_mma.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block,read_A_stage), tCrB(_,_,k_block,read_B_stage), accumulators); + tiled_mma.accumulate_ = UMMA::ScaleOut::One; + } + } + + CUTLASS_DEVICE auto + get_mma_intra_acc(){ + static constexpr uint32_t tmem_base_ptr = 0; // reserved all tmem columns + auto acc_shape_1 = partition_shape_C(TiledMmaIntra1{}, take<0,2>(TileShapeIntraBMM1{})); + auto accumulator_1 = TiledMmaIntra1::make_fragment_C(append(acc_shape_1, + Int<1>{})); + accumulator_1.data() = tmem_base_ptr + tmem_intra_1_Acc; + + auto acc_shape_2 = partition_shape_C(TiledMmaIntra2{}, take<0,2>(TileShapeIntraBMM2{})); + auto accumulator_2 = TiledMmaIntra2::make_fragment_C(append(acc_shape_2, + Int<1>{})); + accumulator_2.data() = tmem_base_ptr + tmem_intra_2_Acc; + return cute::make_tuple(accumulator_1, accumulator_2); + } + + CUTLASS_DEVICE auto + mma_intra_init(TensorStorage& shared_tensors) { + // Init Intra1 + TiledMmaIntra1 tiled_mma_1; + + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_b.data()), SmemLayoutB{}); + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_c.data()), SmemLayoutC{}); + + Tensor tCrC = tiled_mma_1.make_fragment_A(sC); + Tensor tCrB = tiled_mma_1.make_fragment_B(sB); + + // Init Intra2 + TiledMmaIntra2 tiled_mma_2; + + Tensor sX = make_tensor(make_smem_ptr(shared_tensors.smem_x.data()), SmemLayoutX{}); + Tensor tCrQ = tiled_mma_2.make_fragment_A(shape(TmemLayoutQ{})); + Tensor tCrX = tiled_mma_2.make_fragment_B(sX); + + static constexpr uint32_t tmem_base_ptr = 0; // reserved all tmem columns + tCrQ.data() = tmem_base_ptr + tmem_intra_2_A; +#if 0 + if (threadIdx.x % 32 == 0) { + print("sB : ");print(sB);print("\n"); + print("sC : ");print(sC);print("\n"); + print("sX : ");print(sX);print("\n"); + print("tCrC : ");print(tCrC);print("\n"); + print("tCrB : ");print(tCrB);print("\n"); + print("tCrQ : ");print(tCrQ);print("\n"); + print("tCrX : ");print(tCrX);print("\n"); + } +#endif + return cute::make_tuple( + cute::make_tuple(tiled_mma_1, tCrC, tCrB), // mma intra 1 + cute::make_tuple(tiled_mma_2, tCrQ, tCrX) // mma intra 2 + ); + } + + template< + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineC, class PipelineStateC, + class MainloopPipelineX, class PipelineStateX, + class MainloopPipelineIntra, class PipelineStateIntra, + class TiledMma_1, class FragmentA_1, class FragmentB_1, + class TiledMma_2, class FragmentA_2, class FragmentB_2, + class FragmentC_1, class FragmentC_2 + > + CUTLASS_DEVICE + auto mma_intra( + MainloopPipelineB& pipeline_b, PipelineStateB& pipeline_b_consumer_state, + MainloopPipelineC& pipeline_c, PipelineStateC& pipeline_c_consumer_state, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_x_consumer_state, + MainloopPipelineIntra& pipeline_intra, PipelineStateIntra& pipeline_intra_producer_state, + cute::tuple const& mma_inputs_1, + cute::tuple const& mma_inputs_2, + cute::tuple& acc_intra) { + + auto [acc_intra_1, acc_intra_2] = acc_intra; + // require B&C + auto pipeline_b_token = pipeline_b.consumer_try_wait(pipeline_b_consumer_state); + auto pipeline_c_token = pipeline_c.consumer_try_wait(pipeline_c_consumer_state); + auto pipeline_x_token = pipeline_x.consumer_try_wait(pipeline_x_consumer_state); + pipeline_b.consumer_wait(pipeline_b_consumer_state, pipeline_b_token); + pipeline_c.consumer_wait(pipeline_c_consumer_state, pipeline_c_token); + pipeline_x.consumer_wait(pipeline_x_consumer_state, pipeline_x_token); + // do we need try wait? + pipeline_intra.producer_acquire(pipeline_intra_producer_state); + { + auto read_A_stage = pipeline_c_consumer_state.index(); + auto read_B_stage = pipeline_b_consumer_state.index(); + auto accumulator = acc_intra_1(_,_,_,_0{}); + mma(read_A_stage, read_B_stage, mma_inputs_1, accumulator); + } + // release B&C + pipeline_b.consumer_release_from_umma(pipeline_b_consumer_state); + pipeline_c.consumer_release_from_umma(pipeline_c_consumer_state); + pipeline_intra.producer_commit(pipeline_intra_producer_state); + ++pipeline_b_consumer_state; + ++pipeline_c_consumer_state; + ++pipeline_intra_producer_state; + // MMA Intra2 + pipeline_intra.producer_acquire(pipeline_intra_producer_state); + { + auto read_B_stage = pipeline_x_consumer_state.index(); + auto accumulator = acc_intra_2(_,_,_,_0{}); + mma(_0{}, read_B_stage, mma_inputs_2, accumulator); + } + pipeline_intra.producer_commit(pipeline_intra_producer_state); + pipeline_x.consumer_release_from_umma(pipeline_x_consumer_state); + ++pipeline_x_consumer_state; + ++pipeline_intra_producer_state; + } + + CUTLASS_DEVICE auto + get_mma_inter_acc(){ + static constexpr uint32_t tmem_base_ptr = 0; // reserved all tmem columns + auto acc_shape_1 = partition_shape_C(TiledMmaInter1{}, take<0,2>(TileShapeInterBMM1{})); + auto accumulator_1 = TiledMmaInter1::make_fragment_C(append(acc_shape_1, + Int<1>{})); + accumulator_1.data() = tmem_base_ptr + tmem_inter_1_Acc; + + auto acc_shape_2 = partition_shape_C(TiledMmaInter2{}, take<0,2>(TileShapeInterBMM2{})); + auto accumulator_2 = TiledMmaInter2::make_fragment_C(append(acc_shape_2, + Int<1>{})); + accumulator_2.data() = tmem_base_ptr + tmem_inter_2_Acc; + return cute::make_tuple(accumulator_1, accumulator_2); + } + + CUTLASS_DEVICE auto + mma_inter_init(TensorStorage& shared_tensors) { + // Init Inter1 + TiledMmaInter1 tiled_mma_1; + + Tensor sX = make_tensor(make_smem_ptr(shared_tensors.smem_x.data()), SmemLayoutX{}); + + Tensor tCrB = tiled_mma_1.make_fragment_A(shape(TmemLayoutB{})); + Tensor tCrX = tiled_mma_1.make_fragment_B(sX); + + static constexpr uint32_t tmem_base_ptr = 0; // reserved all tmem columns + tCrB.data() = tmem_base_ptr + tmem_inter_1_A; + + // Init Inter2 + TiledMmaInter2 tiled_mma_2; + + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_c.data()), SmemLayoutC{}); + Tensor sP = make_tensor(make_smem_ptr(shared_tensors.smem_p.data()), SmemLayoutP{}); + + Tensor tCrC = tiled_mma_2.make_fragment_A(sC); + Tensor tCrP = tiled_mma_2.make_fragment_B(sP); + + return cute::make_tuple( + cute::make_tuple(tiled_mma_1, tCrB, tCrX), // mma intra 1 + cute::make_tuple(tiled_mma_2, tCrC, tCrP) // mma intra 2 + ); + } + + template< + class MainloopPipelineC, class PipelineStateC, + class MainloopPipelineX, class PipelineStateX, + class MainloopPipelineInter, class PipelineStateInter, + class MainloppPipelineAcc, class PipelineStateAcc, + class TiledMma_1, class FragmentA_1, class FragmentB_1, + class TiledMma_2, class FragmentA_2, class FragmentB_2, + class FragmentC_1, class FragmentC_2 + > + CUTLASS_DEVICE + auto mma_inter( + MainloopPipelineC& pipeline_c, PipelineStateC& pipeline_c_consumer_state, + MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_x_consumer_state, + MainloopPipelineInter& pipeline_inter, PipelineStateInter& pipeline_inter_producer_state, + MainloppPipelineAcc& pipeline_acc, PipelineStateAcc& pipeline_acc_producer_state, + cute::tuple const& mma_inputs_1, + cute::tuple const& mma_inputs_2, + cute::tuple& acc_inter) { + + auto [acc_inter_1, acc_inter_2] = acc_inter; + + // require C + auto pipeline_x_token = pipeline_x.consumer_try_wait(pipeline_x_consumer_state); + auto pipeline_c_token = pipeline_c.consumer_try_wait(pipeline_c_consumer_state); + + pipeline_x.consumer_wait(pipeline_x_consumer_state, pipeline_x_token); + // MMA inter 1 + pipeline_inter.producer_acquire(pipeline_inter_producer_state); + { + auto read_B_stage = pipeline_x_consumer_state.index(); + auto accumulator = acc_inter_1(_,_,_,_0{}); + mma(_0{}, read_B_stage, mma_inputs_1, accumulator); + } + // release B&C + pipeline_x.consumer_release_from_umma(pipeline_x_consumer_state); + ++pipeline_x_consumer_state; + + pipeline_c.consumer_wait(pipeline_c_consumer_state, pipeline_c_token); + // MMA inter 2 + pipeline_acc.producer_acquire(pipeline_acc_producer_state); + { + auto read_A_stage = pipeline_c_consumer_state.index(); + auto accumulator = acc_inter_2(_,_,_,_0{}); + mma(read_A_stage, _0{}, mma_inputs_2, accumulator); + } + pipeline_c.consumer_release_from_umma(pipeline_c_consumer_state); + pipeline_inter.producer_commit(pipeline_inter_producer_state); + pipeline_acc.producer_commit(pipeline_acc_producer_state); + ++pipeline_c_consumer_state; + ++pipeline_inter_producer_state; + ++pipeline_acc_producer_state; + + } + + template< + class FragmentC_1, class FragmentC_2, + class TensorStorage + > + CUTLASS_DEVICE + auto state_init( + cute::tuple& acc_inter, + TensorStorage& shared_tensors) { + + int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; + using CopyAtomT2R = SM100_TMEM_LOAD_16dp256b8x; + using CopyOpR2S = SM90_U16x8_STSM_T; + + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::TransposeBarrier); }; + + // Pre inter 2 + auto accumulator = get<0>(acc_inter); + auto tIntra = make_tensor(accumulator.data(), get<0>(accumulator.layout())); + Tensor sP = make_tensor(make_smem_ptr(shared_tensors.smem_p.data()), SmemLayoutPT{}); + auto tiled_t2r = make_tmem_copy(CopyAtomT2R{}, tIntra); + auto thr_t2r = tiled_t2r.get_slice(thread_idx); + auto tTR_sP = thr_t2r.partition_D(sP(_,_,_0{})); + auto tTR_rP_compute = make_tensor(shape(tTR_sP)); + auto tTR_rP = make_tensor(shape(tTR_sP)); + + clear(tTR_rP); + clear(tTR_rP_compute); + + auto tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_r2s = tiled_r2s.get_slice(thread_idx); + auto tRS_sP = thr_r2s.partition_D(sP(_,_,_0{})); + auto tRS_rP = thr_r2s.retile_S(tTR_rP); + + copy(tiled_r2s, tRS_rP, tRS_sP); + cutlass::arch::fence_view_async_shared(); + + synchronize(); + return make_tuple(tTR_rP_compute); + } + + template< + class MainloopPipelineB, class PipelineStateB, + class MainloopPipelineDelta, class PipelineStateDelta, + class MainloopPipelineInter, class PipelineStateInter, + class FragmentC_1, class FragmentC_2, + class TensorState, + class TensorStorage + > + CUTLASS_DEVICE + auto pre_inter( + MainloopPipelineB& pipeline_b, PipelineStateB& pipeline_b_consumer_state, + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_delta_consumer_state, + MainloopPipelineInter& pipeline_inter, PipelineStateInter& pipeline_inter_consumer_state, + cute::tuple& acc_inter, + TensorState& tState, + TensorStorage& shared_tensors) { + + static constexpr uint32_t tmem_base_ptr = 0; // reserved all tmem columns + int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; + using CopyAtomR2T = SM100_TMEM_STORE_16dp128b16x; + using CopyAtomT2R = SM100_TMEM_LOAD_16dp256b8x; + using CoptAtomS2R = SM75_U32x4_LDSM_N; + using CopyOpR2S = SM90_U16x8_STSM_T; + + // Pipeline check + auto pipeline_b_token = pipeline_b.consumer_try_wait(pipeline_b_consumer_state); + auto pipeline_delta_token = pipeline_delta.consumer_try_wait(pipeline_delta_consumer_state); + pipeline_b.consumer_wait(pipeline_b_consumer_state, pipeline_b_token); + pipeline_delta.consumer_wait(pipeline_delta_consumer_state, pipeline_delta_token); + + // Pre inter 1 + TiledMmaInter1 tiled_mma_1; + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_b.data()), SmemLayoutBT{}); + Tensor tBtB = tiled_mma_1.make_fragment_A(shape(TmemLayoutB{})); + tBtB.data() = tmem_base_ptr + tmem_inter_1_A; + + auto tiled_s2t = make_tmem_copy(CopyAtomR2T{}, tBtB(_,_,_,0)); + auto thr_s2t = tiled_s2t.get_slice(thread_idx); + auto tiled_s2r = make_tiled_copy_S(Copy_Atom{}, tiled_s2t); + auto thr_s2r = tiled_s2r.get_slice(thread_idx); + + auto tBsB_s2r = thr_s2r.partition_S(sB); + auto tBtB_r2t = thr_s2t.partition_D(tBtB); + auto tBrB_r2t = make_tensor(shape(thr_s2t.partition_S(sB(_,_,_,_0{})))); + auto tBrB_s2r = thr_s2r.retile_D(tBrB_r2t); + + auto ts0 = size<0>(TileShape{}); + Layout row_layout = make_layout(make_shape ( ts0, ts0, Int{}), + make_stride(_0{}, _1{}, ts0)); + + Tensor sDelta = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta.data()), row_layout)); + Tensor sDeltaA = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), row_layout)); + + auto thr_mma = tiled_mma_1.get_thread_slice(thread_idx); + auto tCsDelta = thr_mma.partition_A(sDelta); + auto tCsDeltaA = thr_mma.partition_A(sDeltaA); + + auto tBsDelta = thr_s2t.partition_S(tCsDelta(_,_,_,pipeline_delta_consumer_state.index())); + auto tBrDelta = make_tensor(shape(tBsDelta)); + auto tBsDeltaA = thr_s2t.partition_S(tCsDeltaA(_,_,_,pipeline_delta_consumer_state.index())); + auto tBrDeltaA = make_tensor(shape(tBsDeltaA)); + + Layout linear_layout = make_layout(make_shape ( ts0, Int{}), + make_stride(_1{}, ts0)); + Tensor sDeltaA_last = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), linear_layout)); + +#if 0 + if (threadIdx.x % 128 == 0) { + print("sB : ");print(sB);print("\n"); + print("sB : ");print(SmemLayoutB{});print("\n"); + print("tBtB : ");print(tBtB);print("\n"); + print("tBsB_s2r : ");print(tBsB_s2r);print("\n"); + print("tBtB_r2t : ");print(tBtB_r2t);print("\n"); + print("tBrB_r2t : ");print(tBrB_r2t);print("\n"); + print("tBrB_s2r : ");print(tBrB_s2r);print("\n"); + print("tBsDelta : ");print(tBsDelta);print("\n"); + print("tBsDeltaA : ");print(tBsDeltaA);print("\n"); + print("tBrDelta : ");print(tBrDelta);print("\n"); + print("tBrDeltaA : ");print(tBrDeltaA);print("\n"); + } +#endif + + copy(tiled_s2r, tBsB_s2r(_,_,_,_,pipeline_b_consumer_state.index()), tBrB_s2r); + + copy(tBsDelta, tBrDelta); + copy(tBsDeltaA, tBrDeltaA); + + ElementAcc last_column = ElementAcc(0.f); + last_column = sDeltaA_last(_,pipeline_delta_consumer_state.index())(int(ts0) - 1); + + cutlass::arch::fence_view_async_shared(); + + auto tCompute = make_tensor(shape(tBrB_s2r)); + auto tBrB_Compute = make_tensor(shape(tBrB_s2r)); + auto tBrDelta_Compute = make_tensor(shape(tBrDelta)); + auto tBrDeltaA_Compute = make_tensor(shape(tBrDeltaA)); + + type_convert(tBrB_s2r, tBrB_Compute); + type_convert(tBrDelta, tBrDelta_Compute); + type_convert(tBrDeltaA, tBrDeltaA_Compute); + + auto tBrDeltaA_packed = recast(tBrDeltaA_Compute); + auto tBrDelta_packed = recast(tBrDelta_Compute); + auto tBrB_packed = recast(tBrB_Compute); + auto tCompute_packed = recast(tCompute); + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tCompute_packed); ++ii) { + tBrDeltaA_Compute(2 * ii) = expf(last_column - tBrDeltaA_Compute(2 * ii)); + tBrDeltaA_Compute(2 * ii + 1) = expf(last_column - tBrDeltaA_Compute(2 * ii + 1)); + + cute::mul(tCompute_packed(ii), tBrDelta_packed(ii), tBrDeltaA_packed(ii)); + cute::mul(tCompute_packed(ii), tBrB_packed(ii), tCompute_packed(ii)); + } + + type_convert(tCompute, tBrB_r2t); + + copy(tiled_s2t, tBrB_r2t, tBtB_r2t(_,_,_,_,_0{})); + cutlass::arch::fence_view_async_tmem_store(); + + // Pipeline update + pipeline_b.consumer_release_from_threads(pipeline_b_consumer_state); + pipeline_delta.consumer_release_from_threads(pipeline_delta_consumer_state); + pipeline_inter.consumer_release(pipeline_inter_consumer_state); + ++pipeline_b_consumer_state; + ++pipeline_delta_consumer_state; + ++pipeline_inter_consumer_state; + + // Pipeline check + pipeline_inter.consumer_wait(pipeline_inter_consumer_state); + + // Pre inter 2 + auto accumulator = get<0>(acc_inter); + auto tIntra = make_tensor(accumulator.data(), get<0>(accumulator.layout())); + Tensor sP = make_tensor(make_smem_ptr(shared_tensors.smem_p.data()), SmemLayoutPT{}); + auto tiled_t2r = make_tmem_copy(CopyAtomT2R{}, tIntra); + auto thr_t2r = tiled_t2r.get_slice(thread_idx); + auto tTR_tP = thr_t2r.partition_S(tIntra); + auto tTR_sP = thr_t2r.partition_D(sP(_,_,_0{})); + auto tTR_rP_compute = make_tensor(shape(tTR_sP)); + auto tTR_rP = make_tensor(shape(tTR_sP)); + + auto tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_r2s = tiled_r2s.get_slice(thread_idx); + auto tRS_sP = thr_r2s.partition_D(sP(_,_,_0{})); + auto tRS_rP = thr_r2s.retile_S(tTR_rP); + + copy(tiled_t2r, tTR_tP, tTR_rP_compute); + cutlass::arch::fence_view_async_tmem_load(); + + for (int ii = 0; ii < size(tTR_rP_compute); ++ii) { + tTR_rP_compute(ii) = tTR_rP_compute(ii) + expf(last_column) * static_cast(tState(ii)); + } + + type_convert(tTR_rP_compute,tTR_rP); + // Update the state + copy(tTR_rP_compute, tState); + + copy(tiled_r2s, tRS_rP, tRS_sP); + // Fence for the next iteration `pipeline_inter.consumer_release` + cutlass::arch::fence_view_async_shared(); + +#if 0 + if (threadIdx.x % 128 == 0) { + print("accumulator:");print(accumulator);print("\n"); + print("sP : ");print(sP);print("\n"); + print("sP : ");print(SmemLayoutP{});print("\n"); + print("tTR_tP : ");print(tTR_tP);print("\n"); + print("tTR_sP : ");print(tTR_sP);print("\n"); + print("tTR_rP : ");print(tTR_rP);print("\n"); + print("tRS_sP : ");print(tRS_sP);print("\n"); + print("tRS_rP : ");print(tRS_rP);print("\n"); + } +#endif + + } + + template< + class MainloopPipelineDelta, class PipelineStateDelta, + class MainloopPipelineIntra, class PipelineStateIntra, + class FragmentC_1, class FragmentC_2 + > + CUTLASS_DEVICE + auto pre_intra( + MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_delta_consumer_state, + MainloopPipelineIntra& pipeline_intra, PipelineStateIntra& pipeline_intra_consumer_state, + cute::tuple& acc_intra, + TensorStorage& shared_tensors) { + + static constexpr uint32_t tmem_base_ptr = 0; // reserved all tmem columns + int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; + using CopyAtomT2R = SM100_TMEM_LOAD_16dp256b16x; + using CopyAtomR2T = SM100_TMEM_STORE_16dp128b16x; + + // require B&C + auto pipeline_delta_token = pipeline_delta.consumer_try_wait(pipeline_delta_consumer_state); + pipeline_delta.consumer_wait(pipeline_delta_consumer_state, pipeline_delta_token); + // PreIntra2 + auto pipeline_intra_token = pipeline_intra.consumer_try_wait(pipeline_intra_consumer_state); + pipeline_intra.consumer_wait(pipeline_intra_consumer_state, pipeline_intra_token); + + TiledMmaIntra2 tiled_mma_2; + + auto accumulator = get<0>(acc_intra); + Tensor tQ = tiled_mma_2.make_fragment_A(shape(TmemLayoutQ{})); + // Tensor tQ = tiled_mma_2.make_fragment_A(shape(TmemLayoutQ{})); + tQ.data() = tmem_base_ptr + tmem_intra_2_A; + + auto tIntra = make_tensor(accumulator.data(), get<0>(accumulator.layout())); + auto tiled_t2r = make_tmem_copy(CopyAtomT2R{}, tIntra); + auto thr_t2r = tiled_t2r.get_slice(thread_idx); + auto tTR_tQ = thr_t2r.partition_S(tIntra); + auto tTR_sQ = thr_t2r.partition_D(make_tensor(make_shape(get<0>(TileShape{}), get<0>(TileShape{})))); + auto tTR_rQ = make_tensor(shape(tTR_sQ)); + + auto tiled_r2t = make_tmem_copy(CopyAtomR2T{}, tQ); + auto thr_r2t = tiled_r2t.get_slice(thread_idx); + auto tRT_rQ = make_tensor(shape(thr_r2t.partition_S(tQ))); + auto tRT_tQ = thr_r2t.partition_D(tQ); + auto tCompute = make_tensor(shape(tRT_rQ)); + + auto ts0 = size<0>(TileShape{}); + Layout row_layout = make_layout(make_shape ( ts0, ts0, Int{}), + make_stride(_0{}, _1{}, ts0)); + Layout col_layout = make_layout(make_shape ( ts0, ts0, Int{}), + make_stride(_1{}, _0{}, ts0)); + Tensor coord_tensor = make_coord_tensor(make_identity_layout(make_shape(ts0, ts0))); + + Tensor sDelta = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta.data()), row_layout)); + Tensor sDeltaA_Row = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), row_layout)); + Tensor sDeltaA_Col = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_tensors.smem_delta_a.data()), col_layout)); + + auto tQsDeltaA_Row = thr_t2r.partition_D(sDeltaA_Row)(_,_,_,pipeline_delta_consumer_state.index()); + auto tQsDeltaA_Col = thr_t2r.partition_D(sDeltaA_Col)(_,_,_,pipeline_delta_consumer_state.index()); + auto tQsDelta = thr_t2r.partition_D(sDelta)(_,_,_,pipeline_delta_consumer_state.index()); + auto tCoord = thr_t2r.partition_D(coord_tensor); + + auto tQrDeltaA_Row = make_tensor(shape(tQsDeltaA_Row)); + auto tQrDeltaA_Col = make_tensor(shape(tQsDeltaA_Col)); + auto tQrDelta = make_tensor(shape(tQsDelta)); + + auto tCrDeltaA_Row = make_tensor(shape(tQsDeltaA_Row)); + auto tCrDeltaA_Col = make_tensor(shape(tQsDeltaA_Col)); + auto tCrDelta = make_tensor(shape(tQsDelta)); + +#if 0 + if (threadIdx.x % 128 == 0) { + print("SmemLayoutQ:");print(SmemLayoutQ{});print("\n"); + print("row_layout:");print(row_layout);print("\n"); + print("col_layout:");print(col_layout);print("\n"); + print("tQ : ");print(tQ);print("\n"); + print("accumulator: ");print(accumulator);print("\n"); + print("tTR_tQ : ");print(tTR_tQ);print("\n"); + print("tTR_sQ : ");print(tTR_sQ);print("\n"); + print("tTR_rQ : ");print(tTR_rQ);print("\n"); + print("tRT_rQ : ");print(tRT_rQ);print("\n"); + print("tRT_tQ : ");print(tRT_tQ);print("\n"); + } +#endif + + copy(tiled_t2r, tTR_tQ, tTR_rQ); + cutlass::arch::fence_view_async_tmem_load(); + + copy(tQsDeltaA_Row, tQrDeltaA_Row); + copy(tQsDeltaA_Col, tQrDeltaA_Col); + copy(tQsDelta , tQrDelta); + + // Data convert + type_convert(tQsDelta, tCrDelta); + type_convert(tQrDeltaA_Row, tCrDeltaA_Row); + type_convert(tQrDeltaA_Col, tCrDeltaA_Col); + + // SegSum + #pragma unroll + for (int ii = 0; ii < size(tTR_rQ); ++ii) { + ElementAcc tmp(0); + tmp = tQrDeltaA_Col(ii) - tCrDeltaA_Row(ii); + tCompute(ii) = expf(tmp) * tCrDelta(ii) * tTR_rQ(ii); + } + #pragma unroll + for (int ii = 0; ii < size(tTR_rQ); ++ii) { + auto [m,n] = tCoord(ii); + if (m < n) { + tCompute(ii) = 0; + } + } + + type_convert(tCompute, tRT_rQ); + + copy(tiled_r2t, tRT_rQ, tRT_tQ); + cutlass::arch::fence_view_async_tmem_store(); + + // release B&C + pipeline_intra.consumer_release(pipeline_intra_consumer_state); + + ++pipeline_intra_consumer_state; + } + + template< + class ElementSrc, class ElementDst, + class TensorSrc, class TensorDst + > + CUTLASS_DEVICE + auto type_convert( + TensorSrc& tS, + TensorDst& tD) { + + static constexpr int FragmentSize = 2; + NumericArrayConverter converter; + + auto tS_frg = recast>(tS); + auto tD_frg = recast>(tD); + + CUTLASS_PRAGMA_UNROLL + for (int ii = 0; ii < size(tS_frg); ++ii) { + tD_frg(ii) = converter(tS_frg(ii)); + } + } +}; + + +} // namespace cutlass::ssd::collective diff --git a/examples/112_blackwell_ssd/device/ssd.hpp b/examples/112_blackwell_ssd/device/ssd.hpp new file mode 100644 index 00000000..ec983cb5 --- /dev/null +++ b/examples/112_blackwell_ssd/device/ssd.hpp @@ -0,0 +1,272 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + +// common +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" + +#if !defined(__CUDACC_RTC__) +#include "cutlass/cluster_launch.hpp" +#include "cutlass/trace.h" +#endif // !defined(__CUDACC_RTC__) + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::ssd::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class SSD { +public: + using Kernel = Kernel_; + + static int const kThreadCount = Kernel::MaxThreadsPerBlock; + + /// Argument structure: User API + using Arguments = typename Kernel::Arguments; + /// Argument structure: Kernel API + using Params = typename Kernel::Params; + +private: + + /// Kernel API parameters object + Params params_; + + bool is_initialized(bool set = false) { + static bool initialized = false; + if (set) initialized = true; + return initialized; + } + +public: + + /// Access the Params structure + Params const& params() const { + return params_; + } + + /// Determines whether the GEMM can execute the given problem. + static Status + can_implement(Arguments const& args) { + if (Kernel::can_implement(args)) { + return Status::kSuccess; + } + else { + return Status::kInvalid; + } + } + + /// Gets the workspace size + static size_t + get_workspace_size(Arguments const& args) { + size_t workspace_bytes = 0; + workspace_bytes += Kernel::get_workspace_size(args); + return workspace_bytes; + } + + /// Computes the grid shape + static dim3 + get_grid_shape(Params const& params) { + return Kernel::get_grid_shape(params); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int /* smem_capacity */ = -1) { + CUTLASS_TRACE_HOST("Universal::maximum_active_blocks()"); + int max_active_blocks = -1; + int smem_size = Kernel::SharedStorageSize; + + // first, account for dynamic smem capacity if needed + cudaError_t result; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + result = cudaFuncSetAttribute( + device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " + << cudaGetErrorString(result)); + return -1; + } + } + + // query occupancy after setting smem size + result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + device_kernel, + Kernel::MaxThreadsPerBlock, + smem_size); + + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " + << cudaGetErrorString(result)); + return -1; + } + + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Initializes GEMM state from arguments. + Status + initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("Universal::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Initialize the workspace + Status status = Kernel::initialize_workspace(args, workspace, stream); + if (status != Status::kSuccess) { + return status; + } + + // Initialize the Params structure + params_ = Kernel::to_underlying_arguments(args, workspace); + + if (is_initialized()) return Status::kSuccess; + + // account for dynamic smem capacity if needed + int smem_size = Kernel::SharedStorageSize; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + cudaError_t result = cudaFuncSetAttribute( + device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + } + + is_initialized(true); + + return Status::kSuccess; + } + + /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. + Status + update(Arguments const& args, void* workspace = nullptr) { + CUTLASS_TRACE_HOST("Universal()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + if (workspace_bytes > 0 && nullptr == workspace) { + return Status::kErrorWorkspaceNull; + } + + params_ = Kernel::to_underlying_arguments(args, workspace); + return Status::kSuccess; + } + + /// Primary run() entry point API that is static allowing users to create and manage their own params. + /// Supplied params struct must be construct by calling Kernel::to_underling_arguments() + static Status + run(Params& params, cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("Universal::run()"); + dim3 const block = Kernel::get_block_shape(); + dim3 const grid = get_grid_shape(params); + + // configure smem size and carveout + int smem_size = Kernel::SharedStorageSize; + + Status launch_result; + // Use extended launch API only for mainloops that use it + if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) { + dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}), + cute::size<1>(typename Kernel::ClusterShape{}), + cute::size<2>(typename Kernel::ClusterShape{})); + void const* kernel = (void const*) device_kernel; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params); + } + else { + launch_result = Status::kSuccess; + device_kernel<<>>(params); + } + + cudaError_t result = cudaGetLastError(); + if (cudaSuccess == result && Status::kSuccess == launch_result) { + return Status::kSuccess; + } + else { + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } + } + + // + // Non-static launch overloads that first create and set the internal params struct of this kernel handle. + // + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + if (Status::kSuccess == status) { + status = run(params_, stream); + } + return status; + } + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + return run(args, workspace, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + run(cudaStream_t stream = nullptr) { + return run(params_, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + operator()(cudaStream_t stream = nullptr) { + return run(params_, stream); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_builder.hpp b/examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_builder.hpp new file mode 100644 index 00000000..155abd7a --- /dev/null +++ b/examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_builder.hpp @@ -0,0 +1,259 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 "../collective/sm100_ssd_epilogue.hpp" +#include "../collective/sm100_ssd_gemm_tma_warpspecialized.hpp" +#include "../kernel/sm100_ssd_kernel_tma_warpspecialized.hpp" +#include "../kernel/sm100_ssd_tile_scheduler.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" + +namespace cutlass::ssd::kernel::detail { + +template< + class ElementA, + class ElementB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + UMMA::Major UmmaMajorA, + UMMA::Major UmmaMajorB +> +constexpr auto +sm100_make_ts_tiled_mma() { + return cutlass::gemm::collective::detail::sm100_make_1sm_ts_trivial_tiled_mma< + ElementA, ElementB, ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB>(); +} + +template< + class ElementA, + class ElementB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + UMMA::Major UmmaMajorA, + UMMA::Major UmmaMajorB +> +constexpr auto +sm100_make_ss_tiled_mma() { + return cutlass::gemm::collective::detail::sm100_make_1sm_trivial_tiled_mma< + ElementA, ElementB, ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB>(); +} + +} + +namespace cutlass::ssd::kernel { + +template< + class Element_, + class ElementDA_, + class ElementAcc_, + class ElementY_, + class TileShape_, + bool HAS_D_, + bool D_HAS_HDIM_ +> +struct Sm100SsdBuilder { + + using Element = Element_; + using ElementDA = ElementDA_; + using ElementAcc = ElementAcc_; + using ElementY = ElementY_; + using TileShape = TileShape_; + using ArchTag = cutlass::arch::Sm100; + + // hard-code + using ClusterShape = Shape<_1,_1,_1>; + + static constexpr int StagesInput = 2; + static constexpr int StagesOutput = 2; + + using TileShapeIntraBMM1 = decltype(make_shape(get<0>(TileShape{}), get<0>(TileShape{}), get<2>(TileShape{}))); // (L,L,N) + using TileShapeIntraBMM2 = decltype(make_shape(get<0>(TileShape{}), get<1>(TileShape{}), get<0>(TileShape{}))); // (L,D,L) + using TileShapeInterBMM1 = decltype(make_shape(get<2>(TileShape{}), get<1>(TileShape{}), get<0>(TileShape{}))); // (N,D,L) + using TileShapeInterBMM2 = decltype(make_shape(get<0>(TileShape{}), get<1>(TileShape{}), get<2>(TileShape{}))); // (L,D,N) + + // LxLxN, NT + using TiledMmaIntra1 = decltype(detail::sm100_make_ss_tiled_mma()); + // LxNxL, TN + using TiledMmaIntra2 = decltype(detail::sm100_make_ts_tiled_mma()); + // NxDxL, TN + using TiledMmaInter1 = decltype(detail::sm100_make_ts_tiled_mma()); + // LxDxN, NN + using TiledMmaInter2 = decltype(detail::sm100_make_ss_tiled_mma()); + + // ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K) + using MmaShapeC_MK = decltype(partition_shape_A(TiledMmaIntra1{}, make_shape(cute::size<0>(TileShapeIntraBMM1{}), + cute::size<2>(TileShapeIntraBMM1{})))); + using MmaShapeB_NK = decltype(partition_shape_B(TiledMmaIntra1{}, make_shape(cute::size<1>(TileShapeIntraBMM1{}), + cute::size<2>(TileShapeIntraBMM1{})))); + using MmaShapeQ_MK = decltype(partition_shape_A(TiledMmaIntra2{}, make_shape(cute::size<0>(TileShapeIntraBMM2{}), + cute::size<2>(TileShapeIntraBMM2{})))); + using MmaShapeX_NK = decltype(partition_shape_B(TiledMmaIntra2{}, make_shape(cute::size<1>(TileShapeIntraBMM2{}), + cute::size<2>(TileShapeIntraBMM2{})))); + + using MmaShapeB_MK = decltype(partition_shape_A(TiledMmaInter1{}, make_shape(cute::size<0>(TileShapeInterBMM1{}), + cute::size<2>(TileShapeInterBMM1{})))); + using MmaShapeP_NK = decltype(partition_shape_B(TiledMmaInter2{}, make_shape(cute::size<1>(TileShapeInterBMM2{}), + cute::size<2>(TileShapeInterBMM2{})))); + + using GmemTiledCopyX = cute::SM90_TMA_LOAD; + using GmemTiledCopyB = cute::SM90_TMA_LOAD; + using GmemTiledCopyC = cute::SM90_TMA_LOAD; + + using BlockTileX_N = decltype(cute::size<0,0>(MmaShapeX_NK{}) * cute::size<1>(MmaShapeX_NK{})); + using BlockTileX_K = decltype(cute::size<0,1>(MmaShapeX_NK{}) * cute::size<2>(MmaShapeX_NK{})); + using SmemLayoutAtomX = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, BlockTileX_N, BlockTileX_K>()); + + using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{})); + using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{})); + using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, Element, BlockTileB_N, BlockTileB_K>()); + + using BlockTileBT_M = decltype(cute::size<0,0>(MmaShapeB_MK{}) * cute::size<1>(MmaShapeB_MK{})); + using BlockTileBT_K = decltype(cute::size<0,1>(MmaShapeB_MK{}) * cute::size<2>(MmaShapeB_MK{})); + using SmemLayoutAtomBT = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, BlockTileBT_M, BlockTileBT_K>()); + using TmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, BlockTileBT_M, BlockTileBT_K>()); + + using BlockTileC_M = decltype(cute::size<0,0>(MmaShapeC_MK{}) * cute::size<1>(MmaShapeC_MK{})); + using BlockTileC_K = decltype(cute::size<0,1>(MmaShapeC_MK{}) * cute::size<2>(MmaShapeC_MK{})); + using SmemLayoutAtomC = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, Element, BlockTileC_M, BlockTileC_K>()); + + using BlockTileP_N = decltype(cute::size<0,0>(MmaShapeP_NK{}) * cute::size<1>(MmaShapeP_NK{})); + using BlockTileP_K = decltype(cute::size<0,1>(MmaShapeP_NK{}) * cute::size<2>(MmaShapeP_NK{})); + using SmemLayoutAtomP = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, BlockTileP_N, BlockTileP_K>()); + + using SmemLayoutAtomPT = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, Element, BlockTileP_K, BlockTileP_N>()); + + using BlockTileQ_M = decltype(cute::size<0,0>(MmaShapeQ_MK{}) * cute::size<1>(MmaShapeQ_MK{})); + using BlockTileQ_K = decltype(cute::size<0,1>(MmaShapeQ_MK{}) * cute::size<2>(MmaShapeQ_MK{})); + using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, BlockTileP_N, BlockTileP_K>()); + using TmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, BlockTileQ_M, BlockTileQ_K>()); + + using SmemLayoutX = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomX{}, + append(MmaShapeX_NK{}, Int{}), + Step<_2,_1,_3>{})); + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(MmaShapeB_NK{}, Int{}), + Step<_2,_1,_3>{})); + // Be consistent with SmemLayoutB + using SmemLayoutBT = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomBT{}, + append(MmaShapeB_MK{}, Int{}), + Step<_1,_2,_3>{})); + using TmemLayoutB = decltype(UMMA::tile_to_mma_shape( + TmemLayoutAtomB{}, + append(MmaShapeB_MK{}, Int<1>{}), + Step<_2,_1,_3>{})); + using SmemLayoutC = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomC{}, + append(MmaShapeC_MK{}, Int{}), + Step<_2,_1,_3>{})); + // P only need 1 stage in this case + using SmemLayoutPT = decltype(tile_to_shape( + SmemLayoutAtomPT{}, + append(make_shape(get<2>(TileShape{}), get<1>(TileShape{})), Int<1>{}))); + using SmemLayoutP = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomP{}, + append(MmaShapeP_NK{}, Int<1>{}), + Step<_2,_1,_3>{})); + using TmemLayoutQ = decltype(UMMA::tile_to_mma_shape( + TmemLayoutAtomQ{}, + append(MmaShapeQ_MK{}, Int<1>{}), + Step<_2,_1,_3>{})); + using SmemLayoutQ = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomQ{}, + append(MmaShapeQ_MK{}, Int<2>{}), + Step<_2,_1,_3>{})); + + using SmemLayoutAtomXT = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::MN, Element, decltype(get<0>(TileShape{})), decltype(get<1>(TileShape{}))>()); + using SmemLayoutXT = decltype(tile_to_shape( + SmemLayoutAtomXT{}, + make_shape(size<0>(TileShape{}), size<1>(TileShape{}), Int{}), + Step<_1,_2,_3>{})); + + using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; + using Schedule = cutlass::epilogue::TmaWarpSpecialized; + using EpilogueTile = Shape, Int<32>>; + + using SmemLayoutAtomY = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::MN, ElementY, decltype(get<0>(EpilogueTile{})), decltype(get<1>(EpilogueTile{}))>()); + using SmemLayoutY = decltype(tile_to_shape( + SmemLayoutAtomY{}, + make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int{}), + Step<_2,_1,_3>{})); + using SmemLayoutStoreP = decltype(tile_to_shape( + SmemLayoutAtomP{}, + append(make_shape(get<1>(TileShape{}), get<2>(TileShape{})), Int<1>{}), + Step<_2,_1,_3>{})); + + using CollectiveMainloop = cutlass::ssd::collective::SsdMainloopTmaWarpSpecialized< + Element, ElementDA, ElementAcc, ElementY, TileShape, + StagesInput, StagesOutput, + TiledMmaIntra1, TiledMmaIntra2, + TiledMmaInter1, TiledMmaInter2, + SmemLayoutX, SmemLayoutB, SmemLayoutC, SmemLayoutP, + SmemLayoutBT, SmemLayoutPT, SmemLayoutQ, + TmemLayoutB, TmemLayoutQ>; + using CollectiveEpilogue = cutlass::ssd::collective::SsdEpilogue< + ElementAcc, Element, ElementDA, TileShape, + EpilogueTile, SmemLayoutY, SmemLayoutStoreP, SmemLayoutXT, StagesInput, StagesOutput, + HAS_D_, D_HAS_HDIM_>; + using TileScheduler = cutlass::ssd::kernel::PersistentTileScheduler; + using Kernel = cutlass::ssd::kernel::SsdKernelTmaWarpSpecialized; + +}; + +} diff --git a/examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_tma_warpspecialized.hpp b/examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_tma_warpspecialized.hpp new file mode 100644 index 00000000..c2c428f2 --- /dev/null +++ b/examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,519 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/arch/reg_reconfig.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/arch/arch.h" + +namespace cutlass::ssd::kernel { + +using namespace cute; + +template< + class CollectiveMainloop_, + class CollectiveEpilogue_, + class TileScheduler_ +> +struct SsdKernelTmaWarpSpecialized { + + using CollectiveMainloop = CollectiveMainloop_; + using CollectiveEpilogue = CollectiveEpilogue_; + using TileScheduler = TileScheduler_; + + // TileShape: LDN + using TileShape = typename CollectiveMainloop::TileShape; + // Force to use 1x1x1 + using ClusterShape = typename CollectiveMainloop::ClusterShape; + using ArchTag = cutlass::arch::Sm100; + + static constexpr uint32_t NumMmaThreads = NumThreadsPerWarp; // 1 warp + static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount; + static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp; + + // Pipeline for Tensor X + using MainloopPipelineX = typename CollectiveMainloop::MainloopPipelineX; + using PipelineStateX = typename cutlass::PipelineState; + using PipelineParamsX = typename MainloopPipelineX::Params; + + // Pipeline for Tensor Delta && DeltaA + using MainloopPipelineDelta = typename CollectiveMainloop::MainloopPipelineDelta; + using PipelineStateDelta = typename cutlass::PipelineState; + using PipelineParamsDelta = typename MainloopPipelineDelta::Params; + + // Pipeline for Tensor X + using MainloopPipelineB = typename CollectiveMainloop::MainloopPipelineB; + using PipelineStateB = typename cutlass::PipelineState; + using PipelineParamsB = typename MainloopPipelineB::Params; + + // Pipeline for Tensor X + using MainloopPipelineC = typename CollectiveMainloop::MainloopPipelineC; + using PipelineStateC = typename cutlass::PipelineState; + using PipelineParamsC = typename MainloopPipelineC::Params; + + // Pipeline for Intra Fusion + using MainloopPipelineIntra = typename CollectiveMainloop::MainloopPipelineIntra; + using PipelineStateIntra = typename cutlass::PipelineState; + using PipelineParamsIntra = typename MainloopPipelineIntra::Params; + + // Pipeline for Inter Fusion + using MainloopPipelineInter = typename CollectiveMainloop::MainloopPipelineInter; + using PipelineStateInter = typename cutlass::PipelineState; + using PipelineParamsInter = typename MainloopPipelineInter::Params; + + // Pipeline for Accumulator + using AccumulatorPipeline = typename CollectiveMainloop::AccumulatorPipeline; + using AccumulatorPipelineState = typename CollectiveMainloop::AccumulatorPipelineState; + using PipelineParamsAcc = typename AccumulatorPipeline::Params; + + // Pipeline for Epilgoue store + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + using EpiStorePipelineState = typename CollectiveEpilogue::StorePipelineState; + + // Pipeline for Epilgoue load D + using EpiloadPipelineD = typename CollectiveEpilogue::EpiloadPipelineD; + using PipelineParamsD = typename EpiloadPipelineD::Params; + using PipelineStateD = typename cutlass::PipelineState; + using TmemAllocator = typename cute::TMEM::Allocator1Sm; + + struct SharedStorage { + struct PipelineStorage : cute::aligned_struct<16, _1> { + using PipelineStorageX = typename MainloopPipelineX::SharedStorage; + using PipelineStorageDelta = typename MainloopPipelineDelta::SharedStorage; + using PipelineStorageB = typename MainloopPipelineB::SharedStorage; + using PipelineStorageC = typename MainloopPipelineC::SharedStorage; + using PipelineStorageIntra = typename MainloopPipelineIntra::SharedStorage; + using PipelineStorageInter = typename MainloopPipelineInter::SharedStorage; + using PipelineStorageAcc = typename AccumulatorPipeline::SharedStorage; + using PipelineStorageD = typename EpiloadPipelineD::SharedStorage; + alignas(16) PipelineStorageX pipeline_storage_x; + alignas(16) PipelineStorageDelta pipeline_storage_delta; + alignas(16) PipelineStorageB pipeline_storage_b; + alignas(16) PipelineStorageC pipeline_storage_c; + alignas(16) PipelineStorageD pipeline_storage_d; + alignas(16) PipelineStorageIntra pipeline_storage_intra; + alignas(16) PipelineStorageInter pipeline_storage_inter; + alignas(16) PipelineStorageAcc pipeline_storage_acc; + } pipelines; + + uint32_t tmem_base_ptr; + + struct TensorStorage : cute::aligned_struct<128, _1> { + using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; + using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; + + EpilogueTensorStorage epilogue; + MainloopTensorStorage mainloop; + } tensors; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + // [G, B, EH, C, L, D, N] + using ProblemShape = cute::tuple; + + struct Arguments { + ProblemShape problem_size; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_size; + typename CollectiveMainloop::Params mainloop; + typename CollectiveEpilogue::Params epilogue; + typename TileScheduler::Params tile_scheduler; + }; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = 384; + + static size_t get_workspace_size(Arguments const& args) { return 0; } + static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t) { + return cutlass::Status::kSuccess; + } + + static bool can_implement(Arguments const& args) { + return CollectiveMainloop::can_implement(args.problem_size, args.mainloop); + } + + static dim3 get_grid_shape(Params const& params) { + return TileScheduler::get_grid_shape(params.tile_scheduler); + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + return Params{ + args.problem_size, + CollectiveMainloop::to_underlying_arguments(args.problem_size, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_size, args.epilogue, workspace), + TileScheduler::to_underlying_arguments(args.problem_size, args.hw_info, ClusterShape{}, TileShape{}) + }; + } + + CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { + + using namespace cute; + using X = Underscore; + + // TBD + enum class WarpCategory : int32_t { + MMAInter = 0, + MMAIntra = 1, + DMA0 = 2, // Delta, X + DMA1 = 3, // B, C + PreInter = 4, + PreIntra = 8 + }; + + // Parameters + // [G, B, EH, C, L, D, N] + auto C = get<3>(params.problem_size); + + // Shared memory. + auto& storage = *reinterpret_cast(smem); + + int lane_predicate = cute::elect_one_sync(); + int lane_idx = cutlass::canonical_lane_idx(); + int warp_idx = cutlass::canonical_warp_idx_sync(); + auto warp_category = (WarpCategory(warp_idx) < WarpCategory::PreInter) ? + WarpCategory(warp_idx) : + (WarpCategory(warp_idx) < WarpCategory::PreIntra) ? + WarpCategory::PreInter : WarpCategory::PreIntra; + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + // Pipeline (TBD) + PipelineParamsX pipeline_params_x; + pipeline_params_x.transaction_bytes = CollectiveMainloop::kXLoadBytes; + pipeline_params_x.is_leader = lane_predicate && (warp_category == WarpCategory::DMA0); + pipeline_params_x.num_consumers = 1 + 1 + cutlass::NumThreadsPerWarpGroup; // MMA0 + MMA1 + PreInter + pipeline_params_x.initializing_warp = 4; + + PipelineParamsDelta pipeline_params_delta; + pipeline_params_delta.transaction_bytes = CollectiveMainloop::kDeltaLoadBytes + CollectiveMainloop::kDeltaALoadBytes; + pipeline_params_delta.is_leader = lane_predicate && (warp_category == WarpCategory::DMA0); + pipeline_params_delta.num_consumers = cutlass::NumThreadsPerWarpGroup + cutlass::NumThreadsPerWarpGroup; // PreInter + PreIntra + pipeline_params_delta.initializing_warp = 5; + + PipelineParamsB pipeline_params_b; + pipeline_params_b.transaction_bytes = CollectiveMainloop::kBLoadBytes; + pipeline_params_b.is_leader = lane_predicate && (warp_category == WarpCategory::DMA1); + pipeline_params_b.num_consumers = 1 + cutlass::NumThreadsPerWarpGroup; + pipeline_params_b.initializing_warp = 6; + + PipelineParamsC pipeline_params_c; + pipeline_params_c.transaction_bytes = CollectiveMainloop::kCLoadBytes; + pipeline_params_c.is_leader = lane_predicate && (warp_category == WarpCategory::DMA1); + pipeline_params_c.num_consumers = 1 + 1; + pipeline_params_c.initializing_warp = 7; + + PipelineParamsIntra pipeline_params_intra; + pipeline_params_intra.producer_arv_count = 1; + pipeline_params_intra.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; + pipeline_params_intra.initializing_warp = 8; + + PipelineParamsIntra pipeline_params_inter; + pipeline_params_inter.producer_arv_count = 1; + pipeline_params_inter.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; + pipeline_params_inter.initializing_warp = 9; + + PipelineParamsAcc pipeline_params_acc; + pipeline_params_acc.producer_arv_count = 1; + pipeline_params_acc.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; + pipeline_params_acc.initializing_warp = 10; + + PipelineParamsD pipeline_params_d; + pipeline_params_d.transaction_bytes = CollectiveEpilogue::kEpiloadDBytes; + pipeline_params_d.is_leader = lane_predicate && (warp_category == WarpCategory::DMA0); + pipeline_params_d.num_consumers = cutlass::NumThreadsPerWarpGroup; + pipeline_params_d.initializing_warp = 11; + + + if (warp_category == WarpCategory::DMA0) { + pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Producer; + pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Producer; + pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Producer; + } + if (warp_category == WarpCategory::DMA1) { + pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Producer; + pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Producer; + } + // TBD + if (warp_category == WarpCategory::MMAInter) { + pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer; + // pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer; + pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Consumer; + pipeline_params_inter.role = MainloopPipelineInter::ThreadCategory::Producer; + pipeline_params_acc.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (warp_category == WarpCategory::MMAIntra) { + pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer; + pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer; + pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Consumer; + pipeline_params_intra.role = MainloopPipelineIntra::ThreadCategory::Producer; + } + if (warp_category == WarpCategory::PreInter) { + pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer; + pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Consumer; + pipeline_params_inter.role = MainloopPipelineInter::ThreadCategory::Consumer; + } + if (warp_category == WarpCategory::PreIntra) { + pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer; + pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Consumer; + pipeline_params_acc.role = AccumulatorPipeline::ThreadCategory::Consumer; + pipeline_params_intra.role = MainloopPipelineIntra::ThreadCategory::Consumer; + pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Consumer; + } + + MainloopPipelineX pipeline_x(storage.pipelines.pipeline_storage_x, pipeline_params_x, Shape<_1,_1,_1>{}); + PipelineStateX mainloop_pipe_x_consumer; + PipelineStateX mainloop_pipe_x_producer = cutlass::make_producer_start_state(); + + MainloopPipelineDelta pipeline_delta(storage.pipelines.pipeline_storage_delta, pipeline_params_delta, Shape<_1,_1,_1>{}); + PipelineStateDelta mainloop_pipe_delta_consumer; + PipelineStateDelta mainloop_pipe_delta_producer = cutlass::make_producer_start_state(); + + MainloopPipelineB pipeline_b(storage.pipelines.pipeline_storage_b, pipeline_params_b, Shape<_1,_1,_1>{}); + PipelineStateB mainloop_pipe_b_consumer; + PipelineStateB mainloop_pipe_b_producer = cutlass::make_producer_start_state(); + + MainloopPipelineC pipeline_c(storage.pipelines.pipeline_storage_c, pipeline_params_c, Shape<_1,_1,_1>{}); + PipelineStateC mainloop_pipe_c_consumer; + PipelineStateC mainloop_pipe_c_producer = cutlass::make_producer_start_state(); + + EpiloadPipelineD pipeline_d(storage.pipelines.pipeline_storage_d, pipeline_params_d, Shape<_1,_1,_1>{}); + PipelineStateD epi_load_pipe_d_consumer; + PipelineStateD epi_load_pipe_d_producer = cutlass::make_producer_start_state(); + + MainloopPipelineIntra pipeline_intra(storage.pipelines.pipeline_storage_intra, pipeline_params_intra, Shape<_1,_1,_1>{}); + PipelineStateIntra mainloop_pipe_intra_consumer; + PipelineStateIntra mainloop_pipe_intra_producer = cutlass::make_producer_start_state(); + + MainloopPipelineInter pipeline_inter(storage.pipelines.pipeline_storage_inter, pipeline_params_inter, Shape<_1,_1,_1>{}); + // Opposite pipeline state + PipelineStateInter mainloop_pipe_inter_consumer = cutlass::make_producer_start_state(); + PipelineStateInter mainloop_pipe_inter_producer; + + AccumulatorPipeline pipeline_acc(storage.pipelines.pipeline_storage_acc, pipeline_params_acc, Shape<_1,_1,_1>{}); + AccumulatorPipelineState mainloop_pipe_acc_consumer; + AccumulatorPipelineState mainloop_pipe_acc_producer = cutlass::make_producer_start_state(); + + // Epilogue Store pipeline + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state(); + + // Epilogue Store P pipeline + using EpiStorePPipeline = typename CollectiveEpilogue::StorePPipeline; + typename EpiStorePPipeline::Params epi_store_p_pipeline_params; + epi_store_p_pipeline_params.always_wait = true; + EpiStorePPipeline epi_store_p_pipeline(epi_store_p_pipeline_params); + + PipelineState epi_store_p_pipe_producer_state = cutlass::make_producer_start_state(); + + // Tmem allocator + TmemAllocator tmem_allocator{}; + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer blocks in the Cluster + // and to finish smem init + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); + } + else { + __syncthreads(); + } + + // ignore the tmem alloc + arch::NamedBarrier tmem_allocation_result_barrier(NumMmaThreads + NumMmaThreads + NumEpilogueThreads + NumEpilogueThreads, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + if (warp_category == WarpCategory::MMAIntra) { + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + } + if (warp_category == WarpCategory::MMAInter || warp_category == WarpCategory::PreIntra || warp_category == WarpCategory::PreInter) { + tmem_allocation_result_barrier.arrive_and_wait(); + } + + // Kernel implement(TBD) + + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue; + TileScheduler tile_scheduler{params.tile_scheduler}; + + auto mma_output_intra = collective_mainloop.get_mma_intra_acc(); + auto mma_output_inter = collective_mainloop.get_mma_inter_acc(); + + if (warp_category == WarpCategory::DMA0) { + auto load_input = collective_mainloop.load_x_init(params.mainloop, params.problem_size); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto blk_coord_eh = tile_scheduler.get_block_coord_eh(); + // Load D + collective_epilogue.load_d( + blk_coord_eh, params.epilogue, params.problem_size, + pipeline_d, epi_load_pipe_d_producer, + storage.tensors.mainloop + ); + // Load X + // Load Delta + // Load DeltaA + collective_mainloop.load_x_delta( + blk_coord, params.mainloop, params.problem_size, + pipeline_x, mainloop_pipe_x_producer, + pipeline_delta, mainloop_pipe_delta_producer, + load_input, + storage.tensors.mainloop + ); + } + } + else if (warp_category == WarpCategory::DMA1) { + auto load_input_b = collective_mainloop.load_b_init(params.mainloop, params.problem_size); + auto load_input_c = collective_mainloop.load_c_init(params.mainloop, params.problem_size); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord_b(); + // Load B + // Load C + collective_mainloop.load_b_c( + blk_coord, params.mainloop, params.problem_size, + pipeline_b, mainloop_pipe_b_producer, + pipeline_c, mainloop_pipe_c_producer, + load_input_b, load_input_c, + storage.tensors.mainloop + ); + } + } + else if (warp_category == WarpCategory::MMAIntra) { + auto [mma_inputs_1, mma_inputs_2] = collective_mainloop.mma_intra_init(storage.tensors.mainloop); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + for (int chunk = 0; chunk < C; ++chunk) { + collective_mainloop.mma_intra( + pipeline_b, mainloop_pipe_b_consumer, + pipeline_c, mainloop_pipe_c_consumer, + pipeline_x, mainloop_pipe_x_consumer, + pipeline_intra, mainloop_pipe_intra_producer, + mma_inputs_1, mma_inputs_2, + mma_output_intra + ); + } + } + } + else if (warp_category == WarpCategory::MMAInter) { + auto [mma_inputs_1, mma_inputs_2] = collective_mainloop.mma_inter_init(storage.tensors.mainloop); + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + for (int chunk = 0; chunk < C; ++chunk) { + collective_mainloop.mma_inter( + pipeline_c, mainloop_pipe_c_consumer, + pipeline_x, mainloop_pipe_x_consumer, + pipeline_inter, mainloop_pipe_inter_producer, + pipeline_acc, mainloop_pipe_acc_producer, + mma_inputs_1, mma_inputs_2, + mma_output_inter + ); + } + } + } + else if (warp_category == WarpCategory::PreIntra) { + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto blk_coord_eh = tile_scheduler.get_block_coord_eh(); + bool is_first_iteration = true; + for (int chunk = 0; chunk < C; ++chunk) { + collective_mainloop.pre_intra( + pipeline_delta, mainloop_pipe_delta_consumer, + pipeline_intra, mainloop_pipe_intra_consumer, + mma_output_intra, + storage.tensors.mainloop + ); + collective_epilogue.store( + chunk, blk_coord, blk_coord_eh, params.epilogue, params.problem_size, + pipeline_intra, mainloop_pipe_intra_consumer, + pipeline_acc, mainloop_pipe_acc_consumer, + pipeline_delta, mainloop_pipe_delta_consumer, + pipeline_x, mainloop_pipe_x_consumer, + pipeline_d, epi_load_pipe_d_consumer, + epi_store_pipeline, epi_store_pipe_producer_state, + mma_output_intra, + mma_output_inter, + storage.tensors.mainloop, storage.tensors.epilogue, + is_first_iteration + ); + is_first_iteration = false; + } + if constexpr (CollectiveEpilogue::HasBlockScaleD) { + // update the barrier + pipeline_d.consumer_release(epi_load_pipe_d_consumer); + ++epi_load_pipe_d_consumer; + } + } + uint32_t free_stage_ptr = storage.tmem_base_ptr; + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + else if (warp_category == WarpCategory::PreInter) { + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto [tState] = collective_mainloop.state_init(mma_output_inter, storage.tensors.mainloop); + for (int chunk = 0; chunk < C; ++chunk) { + collective_mainloop.pre_inter( + pipeline_b, mainloop_pipe_b_consumer, + pipeline_delta, mainloop_pipe_delta_consumer, + pipeline_inter, mainloop_pipe_inter_consumer, + mma_output_inter, + tState, + storage.tensors.mainloop + ); + } + auto blk_coord = tile_scheduler.get_block_coord(); + // Epilogue Fstate store + collective_epilogue.store_p( + blk_coord, params.epilogue, params.problem_size, + epi_store_p_pipeline, epi_store_p_pipe_producer_state, + storage.tensors.mainloop + ); + } + } + } +}; + +} // namespace cutlass::ssd::kernel diff --git a/examples/112_blackwell_ssd/kernel/sm100_ssd_tile_scheduler.hpp b/examples/112_blackwell_ssd/kernel/sm100_ssd_tile_scheduler.hpp new file mode 100644 index 00000000..20a2e509 --- /dev/null +++ b/examples/112_blackwell_ssd/kernel/sm100_ssd_tile_scheduler.hpp @@ -0,0 +1,132 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/fast_math.h" +#include "cutlass/kernel_hardware_info.h" + +namespace cutlass::ssd::kernel { + +//////////////////////////////////////////////////////////////////////////////// + +struct PersistentTileScheduler { + + struct Params { + int num_blocks; + int num_groups; + FastDivmod divmod_eh; + FastDivmod divmod_ngroup_ratio; + + KernelHardwareInfo hw_info; + }; + + int block_idx = 0; + Params params; + + CUTLASS_DEVICE + PersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {} + + template + static Params to_underlying_arguments( + ProblemSize const& problem_size, KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, TileShape const& tile_shape) + { + using namespace cute; + auto [G, B, EH, C, L, D, N] = problem_size; + + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = hw_info.sm_count; + if (sm_count <= 0) { + CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + } + + CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count); + hw_info.sm_count = sm_count; + + int num_blocks = B * EH; + int ngroup_ratio = EH / G; + + return Params { + num_blocks, + G, + {EH}, + {ngroup_ratio}, + hw_info + }; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1); + return grid; + } + + CUTLASS_DEVICE + bool is_valid() { + return block_idx < params.num_blocks; + } + + CUTLASS_DEVICE + auto get_block_coord() { + return block_idx; + } + + CUTLASS_DEVICE + auto get_block_coord_b() { + using namespace cute; + int eh_idx, b_idx; + int g_idx, rest_idx; + params.divmod_eh(b_idx, eh_idx, block_idx); + params.divmod_ngroup_ratio(g_idx, rest_idx, eh_idx); + return (params.num_groups * b_idx + g_idx); + } + + CUTLASS_DEVICE + auto get_block_coord_eh() { + using namespace cute; + int eh_idx, b_idx; + params.divmod_eh(b_idx, eh_idx, block_idx); + return eh_idx; + } + + CUTLASS_DEVICE + PersistentTileScheduler& operator++() { + block_idx += gridDim.x; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::ssd::kernel diff --git a/examples/112_blackwell_ssd/reference/reference_ssd.hpp b/examples/112_blackwell_ssd/reference/reference_ssd.hpp new file mode 100644 index 00000000..0cb80e6c --- /dev/null +++ b/examples/112_blackwell_ssd/reference/reference_ssd.hpp @@ -0,0 +1,341 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 "cute/tensor.hpp" + +// training or inference phase (not used yet) +// PHASE 0 : training +// PHASE 1 : inference +#define PHASE 0 + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template< + bool transA, + bool transB, + class Element, + class TensorA, + class TensorB, + class TensorC +> +void mma(TensorA tA, TensorB tB, TensorC tC) { + using namespace cute; + + int M = transA ? int(shape<1>(tA)) : int(shape<0>(tA)); + int N = transB ? int(shape<1>(tB)) : int(shape<0>(tB)); + int K = transA ? int(shape<0>(tA)) : int(shape<1>(tA)); + for (int mi = 0; mi < M; ++mi) { + for (int ni = 0; ni < N; ++ni) { + for (int ki = 0; ki < K; ++ki) { + float a = static_cast(Element(transA ? tA(ki, mi) : tA(mi, ki))); + float b = static_cast(Element(transB ? tB(ki, ni) : tB(ni, ki))); + tC(mi, ni) += a * b; + } + } + } + +} + +template< + class Element, + class Tensor +> +auto segsum(Tensor tensor) { + using namespace cute; + auto C = shape<0>(tensor); + auto L = shape<1>(tensor); + auto cum_sum = make_tensor(make_shape(C,L)); + // cum_sum + for (int ci = 0; ci < C; ++ci) { + for (int li = 0; li < L; ++li) { + if (li == 0) { + cum_sum(ci, li) = tensor(ci, li); + } + else { + cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li); + } + } + } + auto seg_sum_out = make_tensor(make_shape(C, L, L)); + // seg_sum + // [ 1, 0, 0] + // [ e^y, 1, 0] + // [e^(y+z), e^z, 1] + CUTLASS_PRAGMA_UNROLL + for (int ci = 0; ci < C; ++ci) { + for (int i = 0; i < L; ++i) { + for (int j = 0; j < L; ++j) { + if (j < i) { + float tmp = static_cast(cum_sum(ci, i)) - static_cast(cum_sum(ci, j)); + seg_sum_out(ci, i, j) = expf(tmp); + } + else if (j == i) { + seg_sum_out(ci, i, j) = 1.f; + } + else { + seg_sum_out(ci, i, j) = 0.f; + } + } + } + } + + return seg_sum_out; +} + +template< + class Element, + class Tensor +> +auto cumsum(Tensor tensor) { + using namespace cute; + auto C = shape<0>(tensor); + auto L = shape<1>(tensor); + auto cum_sum = make_tensor(make_shape(C,L)); + auto cum_sum_out = make_tensor(make_shape(C,L)); + auto cum_sum_exp_out = make_tensor(make_shape(C, L)); + auto cum_sum_exp_out_last = make_tensor(make_shape(C, L)); + auto last_column = make_tensor(make_shape(C)); + // [x, x+y, x+y+z, ..] + CUTLASS_PRAGMA_UNROLL + for (int ci = 0; ci < C; ++ci) { + for (int li = 0; li < L; ++li) { + if (li == 0) { + cum_sum(ci, li) = tensor(ci, li); + } + else { + cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li); + } + // cum_sum_out(ci, li) = static_cast(cum_sum(ci, li)); + } + } + CUTLASS_PRAGMA_UNROLL + for (int ci = 0; ci < C; ++ci) { + last_column(ci) = static_cast(cum_sum(ci, L-1)); + CUTLASS_PRAGMA_UNROLL + for (int li = 0; li < L; ++li) { + cum_sum_exp_out_last(ci, li) = expf(static_cast(last_column(ci) - cum_sum(ci, li))); + cum_sum_exp_out(ci, li) = expf(static_cast(cum_sum(ci, li))); + } + } + + return make_tuple(cum_sum_exp_out_last, last_column, cum_sum_exp_out); +} + +template< + bool HAS_D, + bool D_HAS_HDIM, + bool HAS_Z, + class TensorY, + class TensorF, + class TensorX, + class TensorDelta, + class TensorDeltaA, + class TensorB, + class TensorC, + class TensorD, + class TensorZ, + class Params +> +void ssd_reference_impl( + TensorY mY, TensorF mF, + TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA, + TensorB mB, TensorC mC, TensorD mD, TensorZ mZ, + Params params) { + + using namespace cute; + using Element = typename Params::Element; + using ElementAcc = typename Params::ElementAcc; + + // x [b, eh, d, c, l] + // delta [b, eh, c, l] + // delta_A [b, eh, c, l] + // B [b, g, n, c, l] + // C [b, g, n, c, l] + // y [b, eh, d, c, l] + // fstate [b, eh, d, n] + // d [ eh, d] + auto [G, B, EH, C, L, D, N] = params.get_problem_shape(); + int group_ratio = EH / G; + for (int b = 0; b < B; ++b) { + for (int eh = 0; eh < EH; ++eh) { + int g = eh / group_ratio; + auto tY = mY(b,eh,_,_,_); + auto tF = mF(b,eh,_,_); + auto tX = mX(b,eh,_,_,_); + auto tDelta = mDelta(b,eh,_,_); + auto tDeltaA = mDeltaA(b,eh,_,_); + auto tB = mB(b,g,_,_,_); + auto tC = mC(b,g,_,_,_); + auto tD = mD(eh,_); + auto tZ = mZ(b,eh,_,_,_); + // IntraBMM1 BxC, LxLxN, NT + // B: [n, c, l] + // C: [n, c, l] + // O: [c, l, l] + auto tIntraBMM1_out = make_tensor(make_shape(C,L,L)); + for (int ci = 0; ci < C; ++ci) { + mma(tC(_,ci,_), tB(_,ci,_), tIntraBMM1_out(ci,_,_)); + } + // Pre_IntraBMM2 DeltaA_IntraBMM2 x Delta x IntraBMM_out + // DeltaA_xxx : [c, l, l] + // Delta : [c, l, _] + // IntraBMM1_out: [c, l, l] + auto tDeltaA_IntraBMM2 = segsum(tDeltaA); + auto tIntraBMM2_inp = make_tensor(make_shape(C, L, L)); + for (int ci = 0; ci < C; ++ci) { + for (int i = 0; i < L; ++i) { + for (int j = 0; j < L; ++j) { + tIntraBMM2_inp(ci, i, j) = tDeltaA_IntraBMM2(ci, i, j) * tDelta(ci, j) * tIntraBMM1_out(ci, i, j); + } + } + } + // IntraBMM2 IntraBMM2_inp x X, LxDxL, TT + // IntraBMM2_inp: [c, l, l] + // X : [d, c, l] + // IntraBMM2_out: [c, l, d] + auto tIntraBMM2_out = make_tensor(make_shape(C,L,D)); + for (int ci = 0; ci < C; ++ci) { + mma(tIntraBMM2_inp(ci,_,_), tX(_,ci,_), tIntraBMM2_out(ci,_,_)); + } + // Pre_InterBMM1 DeltaA_InterBMM1 x Delta x B + // DeltaA_xxx : [c, l] + // Delta : [c, l] + // IntraBMM1_out: [c, n, l] + auto [tDeltaA_InterBMM1, tLast, tCumsum_exp] = cumsum(tDeltaA); + auto tInterBMM1_inp = make_tensor(make_shape(C, N, L)); + for (int ci = 0; ci < C; ++ci) { + for (int i = 0; i < N; ++i) { + for (int j = 0; j < L; ++j) { + tInterBMM1_inp(ci, i, j) = tDeltaA_InterBMM1(ci, j) * tDelta(ci, j) * tB(i, ci, j); + } + } + } + // InterBMM1 InterBMM1_inp x X, NxDxL, swapAB, TT + // InterBMM1_inp: [c, n, l] + // X : [d, c, l] + // InterBMM1_out: [c, n, d] + auto tInterBMM1_out = make_tensor(make_shape(C,N,D)); + for (int ci = 0; ci < C; ++ci) { + mma(tInterBMM1_inp(ci,_,_), tX(_,ci,_), tInterBMM1_out(ci,_,_)); + } + // Initialize state + // PreInterBMM2 + // InterBMM1_out: [c, n, d] + // Last : [c] + auto tInterBMM2_inp = make_tensor(make_shape(C, N, D)); + for (int ci = 0; ci < C; ++ci) { + for (int ni = 0; ni < N; ++ ni){ + for (int di = 0; di < D; ++di) { + if (ci == 0) { + tInterBMM2_inp(ci, ni, di) = 0; + } + else { + tInterBMM2_inp(ci, ni, di) = tInterBMM1_out(ci - 1, ni, di) + expf(tLast(ci - 1)) * tInterBMM2_inp(ci - 1, ni, di); + } + } + } + } + // InterBMM2 InterBMM2_inp x C, LxDxN, NT + // C : [n, c, l] + // InterBMM2_inp: [c, n, d] + // InterBMM2_out: [c, l, d] + auto tInterBMM2_out = make_tensor(make_shape(C,L,D)); + for (int ci = 0; ci < C; ++ci) { + mma(tC(_,ci,_), tInterBMM2_inp(ci,_,_), tInterBMM2_out(ci,_,_)); + } + // Epilogue Cumsum_exp x InterBMM2_out + IntraBMM2_out + // InterBMM2_out: [c, l, d] + // IntraBMM2_out: [c, l, d] + // Cumsum_exp : [c, l] + for (int ci = 0; ci < C; ++ci) { + for (int li = 0; li < L; ++li) { + for (int di = 0; di < D; ++di) { + float y = tInterBMM2_out(ci, li, di) * tCumsum_exp(ci, li) + tIntraBMM2_out(ci, li, di); + float scale; + if constexpr (D_HAS_HDIM) { + scale = static_cast(tD(di)); + } + else { + scale = static_cast(tD(_0{})); + } + if constexpr (HAS_D) { + y = y + static_cast(tX(di, ci, li)) * scale; + } + else { + y = y; + } + if constexpr (HAS_Z) { + float z = static_cast(tZ(di, ci, li)); + // y = y * z * (1 / (1 + exp(-z))); + y = y * z * (1 / (1 + exp(-z))); + } + tY(di, ci, li) = static_cast(y); + } + } + } + // Epilogue Fstate(last C) + for (int ni = 0; ni < N; ++ ni){ + for (int di = 0; di < D; ++di) { + tF(di, ni) = static_cast(tInterBMM1_out(C - 1, ni, di) + expf(tLast(C - 1)) * tInterBMM2_inp(C - 1, ni, di)); + } + } + } + } +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template< + bool HAS_D, + bool D_HAS_HDIM, + bool HAS_Z, + class TensorY, + class TensorF, + class TensorX, + class TensorDelta, + class TensorDeltaA, + class TensorB, + class TensorC, + class TensorD, + class TensorZ, + class Params +> +void ssd_reference( + TensorY mY, TensorF mF, + TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA, + TensorB mB, TensorC mC, TensorD mD, TensorZ mZ, + Params params) { + ssd_reference_impl(mY, mF, mX, mDelta, mDeltaA, mB, mC, mD, mZ, params); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/examples/112_blackwell_ssd/reference/reference_ssd_cumsum.hpp b/examples/112_blackwell_ssd/reference/reference_ssd_cumsum.hpp new file mode 100644 index 00000000..124052f8 --- /dev/null +++ b/examples/112_blackwell_ssd/reference/reference_ssd_cumsum.hpp @@ -0,0 +1,194 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 +#include + +#include "cutlass/coord.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/tensor_view.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/reference/host/gemm.h" +#include "cutlass/arch/arch.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +#include "cute/int_tuple.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "cute/util/debug.hpp" +#include "cute/config.hpp" + +namespace cutlass::ssd::kernel { + +using namespace cute; + +template< + class Element_, + class ElementD_, + class TileShape_> +struct CumsumKernel { + using Element = Element_; + using ElementD = ElementD_; + using TileShape = TileShape_; // L,D,N + + // Required by `device_kernel` + static constexpr int MaxThreadsPerBlock = 128; + static constexpr int MinBlocksPerMultiprocessor = 1; + using ArchTag = arch::Sm90; + + static constexpr int AlignmentBytes = 16; + + struct SharedStorage { + /* empty, no smem needed */ + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + struct TransformArguments { + const Element* ptr_DeltaA; + ElementD* ptr_Cumsum; + }; + + struct TransformParams { + const Element* ptr_DeltaA; + ElementD* ptr_Cumsum; + }; + + using ProblemShape = cute::tuple; // b, eh, c, l + struct Arguments { + ProblemShape problem_shape{}; + TransformArguments transform{}; + KernelHardwareInfo hw_info{}; + }; + + struct Params { + ProblemShape problem_shape{}; + TransformParams transform{}; + KernelHardwareInfo hw_info{}; + }; + + static Params + to_underlying_arguments(Arguments const& args, void* workspace) { + return Params{ + ProblemShape{args.problem_shape}, + TransformParams{args.transform.ptr_DeltaA, args.transform.ptr_Cumsum}, + KernelHardwareInfo{args.hw_info}}; + } + + static Status + can_implement(Arguments const& args) { + return Status::kSuccess; + } + + static size_t + get_workspace_size(Arguments const& args) { + return size_t(0); + } + + static Status + initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr, + CudaHostAdapter *cuda_adapter = nullptr) { + return Status::kSuccess; + } + + static dim3 + get_grid_shape(Params const& params) { + auto [B, EH, C, L] = params.problem_shape; + return dim3(B*EH, 1, 1); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + CUTE_HOST_DEVICE + void + operator()(Params params, [[maybe_unused]] char* smem_buf = nullptr) { + auto [B, EH, C, L] = params.problem_shape; + + auto layout = make_layout(make_shape(L, C, EH*B)); + + auto mD_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_DeltaA), make_layout(reverse(layout.shape()), reverse(layout.stride()))); + auto mC_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_Cumsum), make_layout(reverse(layout.shape()), reverse(layout.stride()))); + auto cD_bcl = make_identity_tensor(shape(mD_bcl)); + + int blk_idx = blockIdx.x; + int thread_idx = threadIdx.x; + + auto tD = logical_divide(mD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_); + auto tC = logical_divide(mC_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_); + auto cD = logical_divide(cD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_); + + static constexpr int NumPacked = AlignmentBytes / sizeof(ElementD); + using PackedTypeDeltaA = uint_bit_t * NumPacked>; + using PackedTypeCumsum = uint_bit_t * NumPacked>; + +#if 0 + if (thread_idx % 128 == 0 && blk_idx == 0) { + print("tD : ");print(tD);print("\n"); + print("tC : ");print(tC);print("\n"); + print("cD : ");print(cD);print("\n"); + } +#endif + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < shape<0>(tD); ++i) { + float last_element = 0.f; + auto crd = cD(i,_0{}); + auto tD_recast = recast(tD); + auto tC_recast = recast(tC); + if (elem_less(crd, shape(mD_bcl))) { + for (int j = 0; j < shape<1>(tD_recast); ++j) { + auto tD_slice = make_tensor(make_shape(Int{})); + auto tC_slice = make_tensor(make_shape(Int{})); + auto tD_slice_recast = recast(tD_slice); + auto tC_slice_recast = recast(tC_slice); + + tD_slice_recast(_0{}) = tD_recast(i,j); + for (int k = 0; k < NumPacked; ++ k) { + last_element += static_cast(tD_slice(k)); + tC_slice(k) = static_cast(last_element); + } + tC_recast(i,j) = tC_slice_recast(_0{}); + } + } + } + } + +private: + +}; + +} // End namespace cutlass diff --git a/examples/112_blackwell_ssd/utils/pipeline.h b/examples/112_blackwell_ssd/utils/pipeline.h new file mode 100644 index 00000000..a26004cf --- /dev/null +++ b/examples/112_blackwell_ssd/utils/pipeline.h @@ -0,0 +1,225 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 "cute/numeric/integral_constant.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cutlass/arch/barrier.h" +#include "cutlass/pipeline/sm90_pipeline.hpp" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { + +using namespace cute; + +// Producer-consumer pipeline implementation +// for TMA producer. In this case, Multi-consumers +// (UMMAs, TransformWarps, ...) +// will arrive at the same empty barrier. +// A naive implement without mcast support. +template , class AtomThrShape_MNK_ = Shape<_1,_1,_1>> +class PipelineTmaMultiConsumersAsync { +public: + static constexpr uint32_t Stages = Stages_; + using AtomThrShape_MNK = AtomThrShape_MNK_; +private: + using Impl = PipelineTmaAsync; +public: + using FullBarrier = typename Impl::FullBarrier; + using EmptyBarrier = typename Impl::EmptyBarrier; + using ProducerBarrierType = typename Impl::ProducerBarrierType; + using ConsumerBarrierType = typename Impl::ConsumerBarrierType; + using PipelineState = typename Impl::PipelineState; + using SharedStorage = typename Impl::SharedStorage; + using ThreadCategory = typename Impl::ThreadCategory; + using Params = typename Impl::Params; + + // Helper function to initialize barriers + static + CUTLASS_DEVICE + void + init_barriers(SharedStorage& storage, Params params, ClusterShape cluster_shape) { + int warp_idx = canonical_warp_idx_sync(); + if (warp_idx == params.initializing_warp) { + constexpr int producer_arv_cnt = 1; + int const consumer_arv_cnt = params.num_consumers; + cutlass::arch::detail::initialize_barrier_array_pair_aligned( + storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, consumer_arv_cnt); + } + cutlass::arch::fence_barrier_init(); + } + + CUTLASS_DEVICE + void init_masks(ClusterShape cluster_shape) { + // Calculate consumer mask + if (params_.role == ThreadCategory::Consumer) { + is_signalling_thread_ = 1; + dst_blockid_ = 0; + } + } + + // Constructor by default initializes barriers and calculates masks. + // These operations can be explicity deferred by specifying InitBarriers and InitMasks. + // If deferred, user code needs to guarantee init_masks and/or init_barriers is/are called. + template + CUTLASS_DEVICE + PipelineTmaMultiConsumersAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape, InitBarriers = {}, InitMasks = {}) + : impl_(storage, params, cluster_shape) + , params_(params) + , empty_barrier_ptr_(&storage.empty_barrier_[0]) + , full_barrier_ptr_(&storage.full_barrier_[0]) { + static_assert(cute::is_same_v || cute::is_same_v); + static_assert(size(cluster_shape) == 1, "PipelineTmaMultiConsumersAsync only supports 1x1x1 cluster shape now"); + if constexpr (cute::is_same_v) { + init_barriers(storage, params_, cluster_shape); + } + + static_assert(cute::is_same_v || cute::is_same_v); + if constexpr (cute::is_same_v) { + init_masks(cluster_shape); + } + } + + //////////////////// + // Producer APIs + //////////////////// + // Four member functions are always used in pairs: + // + // * producer_try_acquire and producer_acquire, and + // * consumer_try_wait and consumer_wait. + // + // The two functions with "try" in their names are called "try" functions, + // and the other two are conceptually "finalize" functions. + // The "try" function in each pair starts the process of waiting on the barrier to flip. + // It opportunistically waits for an implementation-dependent timeout. + // Whether or not the barrier has flipped yet, the try function will return a token. + // If the token indicates that the barrier has not flipped, + // then the token must be passed into the corresponding "finalize" function. + // The finalize function will then block until the barrier has flipped. + // If the token indicates that the barrier _has_ flipped, + // then it is still correct to pass it into the finalize function. + // The finalize function will return immediately in that case. + CUTLASS_DEVICE + ProducerToken producer_try_acquire(PipelineState state, uint32_t skip_wait = false) { + return impl_.producer_try_acquire(state, skip_wait); + } + + CUTLASS_DEVICE + void producer_acquire(PipelineState state, ProducerToken barrier_token = {BarrierStatus::WaitAgain}) { + impl_.producer_acquire(state, barrier_token); + } + + // NOP for TMA based mainloop + CUTLASS_DEVICE + void producer_commit(PipelineState state, uint32_t bytes) { + impl_.producer_commit(state, bytes); + } + + // Prevents early exit of producer blocks in Cluster. + // This should be called once before kernel exits. + CUTLASS_DEVICE + void producer_tail(PipelineState state) { + impl_.producer_tail(state); + } + + CUTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(PipelineState state) { + return impl_.producer_get_barrier(state); + } + + //////////////////// + // Consumer APIs + //////////////////// + CUTLASS_DEVICE + ConsumerToken consumer_try_wait(PipelineState state, uint32_t skip_wait = false) { + return impl_.consumer_try_wait(state, skip_wait); + } + + CUTLASS_DEVICE + void consumer_wait(PipelineState state, ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) { + impl_.consumer_wait(state, barrier_token); + } + + CUTLASS_DEVICE + void consumer_release_from_umma(PipelineState state) { + consumer_release_from_umma(state.index(), false); + } + + CUTLASS_DEVICE + void consumer_release_from_threads(PipelineState state) { + consumer_release_from_threads(state.index()); + } + +private: + Impl impl_; + Params params_; + uint32_t dst_blockid_ = 0; + uint32_t is_signalling_thread_ = 0; + EmptyBarrier *empty_barrier_ptr_; + FullBarrier *full_barrier_ptr_; + uint16_t block_id_mask_ = 0; + static constexpr bool is_2sm_mma = size(AtomThrShape_MNK{}) > 1; + + // Consumer signalling Producer of completion + // Ensures all blocks in the Same Row and Column get notifed. + CUTLASS_DEVICE + void consumer_release_from_umma(uint32_t stage, uint32_t skip) { + uint64_t* smem_ptr = reinterpret_cast(&empty_barrier_ptr_[stage]); + if constexpr (is_2sm_mma) { // Mma cluster shape is 2x1 + if (!skip) { + cutlass::arch::umma_arrive_multicast_2x1SM(smem_ptr, block_id_mask_); + } + } + else { + if (!skip) { + if constexpr (cute::is_static_v and size(ClusterShape{}) == 1) { + cutlass::arch::umma_arrive(smem_ptr); + } + else { + cutlass::arch::umma_arrive_multicast(smem_ptr, block_id_mask_); + } + } + } + } + + CUTLASS_DEVICE + void consumer_release_from_threads(uint32_t stage, uint32_t skip = false) { + empty_barrier_ptr_[stage].arrive(dst_blockid_, is_signalling_thread_ & (!skip)); + #ifndef NDEBUG + if (params_.role == ThreadCategory::Producer || params_.role == ThreadCategory::NonParticipant) { + asm volatile ("brkpt;\n" ::); + } + #endif + } +}; + +} diff --git a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu index 663ba874..f4638545 100644 --- a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu +++ b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu @@ -115,9 +115,29 @@ using namespace cute; ///////////////////////////////////////////////////////////////////////////////////////////////// /// GEMM kernel configurations ///////////////////////////////////////////////////////////////////////////////////////////////// -using MmaType = cutlass::bfloat16_t; -using QuantType = cutlass::int4b_t; -constexpr int TileShapeK = 128 * 8 / sizeof_bits::value; + +// Select MMA type via compile flag +#if defined(CUTLASS_USE_FP16) + using MmaType = cutlass::half_t; // FP16 +#elif defined(CUTLASS_USE_TF32) + using MmaType = cutlass::tfloat32_t; // TF32 (FP32 format with reduced precision) +#else + using MmaType = cutlass::bfloat16_t; // BF16 (default) +#endif + +// Select quantization type via compile flag for this example +#if defined(CUTLASS_MIXED_DTYPE_E2M1) + using QuantType = cutlass::float_e2m1_t; // E2M1 (FP4) +#else + using QuantType = cutlass::int4b_t; // INT4 Two's Complement (default) +#endif + +// TF32 requires K to be multiple of 8; BF16/FP16 can go higher +#if defined(CUTLASS_USE_TF32) + constexpr int TileShapeK = 64; // TF32: K must be multiple of 8, use 64 for good performance +#else + constexpr int TileShapeK = 128 * 8 / sizeof_bits::value; +#endif // A matrix configuration using ElementA = MmaType; // Element type for A matrix operand @@ -139,11 +159,19 @@ using StrideB = cutlass::detail::TagToStrideB_t; // Define the CuTe layout for reoredered quantized tensor B // LayoutAtomQuant places values that will be read by the same thread in contiguous locations in global memory. // It specifies the reordering within a single warp's fragment -//using ValueShuffle = Layout<_1>; // no value reordering +#if defined(CUTLASS_MIXED_DTYPE_E2M1) || defined(CUTLASS_USE_TF32) +// E2M1 & TF32: Use simpler layout without ValueShuffle (like FP8 example) +// ValueShuffle currrently isn't enabled for E2M1 until LayoutAwareConvertImpl specializations support shuffle reordering. +// and for TF32 until we support both the K=8 vs K=16 dimensions for tiles +using LayoutAtomQuant = decltype(cutlass::compute_memory_reordering_atom()); +#else +// INT4: Use ValueShuffle for optimal performance with FP16/BF16 +// using ValueShuffle = Layout<_1>; // no value reordering using ValueShuffle = Layout, Stride<_4,_1>>; // order [0,2,4,6,1,3,5,7] int constexpr NumShuffleAtoms = 1; using MmaAtomShape = Layout>>; using LayoutAtomQuant = decltype(cutlass::compute_memory_reordering_atom()); +#endif using LayoutB_Reordered = decltype(cute::tile_to_shape(LayoutAtomQuant{}, Layout, StrideB>{})); using ElementScale = MmaType; @@ -151,7 +179,7 @@ using ElementZero = ElementScale; using LayoutScale = cutlass::layout::RowMajor; // C/D matrix configuration -using ElementC = cutlass::half_t; // Element type for C and D matrix operands +using ElementC = MmaType; // Element type for C and D matrix operands (matches MMA type) using LayoutC = cutlass::layout::RowMajor; // Layout type for C and D matrix operands constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes) diff --git a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu index 0f1839fd..7c972177 100644 --- a/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu +++ b/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu @@ -120,7 +120,15 @@ using namespace cute; /// GEMM kernel configurations ///////////////////////////////////////////////////////////////////////////////////////////////// using MmaType = cutlass::float_e4m3_t; -using QuantType = cutlass::int4b_t; + +// Select quantization type via compile flag for this example +// templatized throughout code to enable instantiating both int4 and e2m1 versions in the same program +#if defined(CUTLASS_MIXED_DTYPE_E2M1) + using QuantType = cutlass::float_e2m1_t; // E2M1 (FP4) +#else + using QuantType = cutlass::int4b_t; // INT4 Two's Complement (default) +#endif + constexpr int TileShapeK = 128 * 8 / sizeof_bits::value; // A matrix configuration @@ -315,6 +323,7 @@ struct Options : MixedDtypeOptions { ///////////////////////////////////////////////////////////////////////////////////////////////// /// Initialize operands to be used in the GEMM and reference GEMM +template void initialize(Options const& options) { auto shape_B = cute::make_shape(options.n, options.k, options.l); @@ -330,7 +339,7 @@ void initialize(Options const& options) { auto layout_B = make_layout(shape_B, stride_B); - auto a_coord = cutlass::make_Coord(options.m * options.l, options.k); + auto a_coord = cutlass::make_Coord(options.m * options.l, options.k); auto b_coord = cutlass::make_Coord(options.k, options.n * options.l); auto c_coord = cutlass::make_Coord(options.m * options.l, options.n); @@ -346,14 +355,15 @@ void initialize(Options const& options) { block_scale_packed.reset(scale_k * options.l * options.n); block_zero.reset(scale_k * options.l * options.n); + // Initialize all base tensors initialize_tensor(block_A, seed + 2022); initialize_tensor(block_B, seed + 2021); - cutlass::unified_encode_int4b(block_B.get(), block_B_modified.get(), block_B.size()); initialize_tensor(block_C, seed + 2020); initialize_scale(block_scale, options); - cutlass::pack_scale_fp8(block_scale.get(), block_scale_packed.get(), block_scale.size()); + cutlass::pack_scale_fp8(block_scale.get(), block_scale_packed.get(), block_scale.size()); initialize_zero(block_zero, options); + // Compute dequantized reference for validation BEFORE formatting block_B auto shape_scale_zero = cute::make_shape(options.n, scale_k, options.l); stride_S = cutlass::make_cute_packed_stride(StrideS{}, cute::make_shape(options.n, scale_k, options.l)); stride_S_ref = cutlass::make_cute_packed_stride(StrideS_ref{}, cute::make_shape(options.n, scale_k, options.l)); @@ -362,6 +372,28 @@ void initialize(Options const& options) { cudaStream_t stream = cudaStreamDefault; cutlass::dequantize(block_B_dq.get(), block_B.get(), layout_B, block_scale.get(), block_zero.get(), layout_scale_zero, options.g, stream); + // Format B to separate buffer, preserving original block_B in case + // it's needed by the application. + + #ifdef CUTLASS_MIXED_DTYPE_E2M1 + // E2M1: Copy to formatting buffer + cutlass::device_memory::copy_device_to_device(block_B_modified.get(), block_B.get(), block_B.size()); + #else + // INT4: Encode to formatting buffer + cutlass::unified_encode_int4b(block_B.get(), block_B_modified.get(), block_B.size()); + #endif + + + // original code + // if constexpr (cutlass::platform::is_floating_point::value || + // cute::is_same_v) { + // // E2M1: Copy to formatting buffer + // cutlass::device_memory::copy_device_to_device(block_B_modified.get(), block_B.get(), block_B.size()); + // } else { + // // INT4: Encode to formatting buffer + // cutlass::unified_encode_int4b(block_B.get(), block_B_modified.get(), block_B.size()); + // } + if (options.shuffle) { // Repeat the reorder layout atom to tile the whole tensor shape layout_B_reordered = cute::tile_to_shape(LayoutAtomQuant{}, shape_B); @@ -464,7 +496,7 @@ bool verify(Options const& options) { template int run(Options &options) { - initialize(options); + initialize(options); // Instantiate CUTLASS kernel depending on templates Gemm gemm; diff --git a/examples/60_cutlass_import/CMakeLists.txt b/examples/60_cutlass_import/CMakeLists.txt index 6a3acb98..b8d74e34 100644 --- a/examples/60_cutlass_import/CMakeLists.txt +++ b/examples/60_cutlass_import/CMakeLists.txt @@ -54,7 +54,7 @@ target_sources(example PRIVATE main.cpp) target_include_directories( example - PRIVATE + SYSTEM PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} ) diff --git a/examples/77_blackwell_fmha/77_blackwell_fmha_gen.cu b/examples/77_blackwell_fmha/77_blackwell_fmha_gen.cu index ee7eddec..de8fdb2f 100644 --- a/examples/77_blackwell_fmha/77_blackwell_fmha_gen.cu +++ b/examples/77_blackwell_fmha/77_blackwell_fmha_gen.cu @@ -730,8 +730,8 @@ int main_single(int argc, char const **args) { if (props.major != 10 || (props.minor != 0 && props.minor != 3)) { std::cout - << "This example requires a GPU of NVIDIA's Blackwell Architecture " - << "(compute capability 90) and CUDA 12.0 or greater.\n"; + << "This example requires a GPU of NVIDIA's Blackwell Datacenter-class Architecture " + << "(compute capability 100 or 103) and CUDA 12.0 or greater.\n"; return 0; } diff --git a/examples/77_blackwell_fmha/77_blackwell_mla.cu b/examples/77_blackwell_fmha/77_blackwell_mla.cu index e045ee4c..e1436080 100644 --- a/examples/77_blackwell_fmha/77_blackwell_mla.cu +++ b/examples/77_blackwell_fmha/77_blackwell_mla.cu @@ -75,7 +75,11 @@ struct Options { int split_kv = -1; // number of split along k dim. bool is_var_split_kv = false; int max_split_kv = 16; +#ifdef CPASYNC + int page = 1; +#else int page = -1; +#endif float spread = 0.2f; int iterations = 3; bool verify = false; @@ -260,7 +264,7 @@ struct ExampleResult { /////////////////////////////////////////////////////////////////////////////////////////////////// -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)) /////////////////////////////////////////////////////////////////////////////////////////////////// @@ -751,7 +755,7 @@ void run_mla(Options const & options, cutlass::KernelHardwareInfo const& hw_info /////////////////////////////////////////////////////////////////////////////////////////////////// -#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED) /////////////////////////////////////////////////////////////////////////////////////////////////// @@ -796,7 +800,7 @@ int main_single(int argc, char const **args) { return -1; } -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) +#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)) // // Run examples diff --git a/examples/77_blackwell_fmha/CMakeLists.txt b/examples/77_blackwell_fmha/CMakeLists.txt index e3fb2096..799cf4ab 100644 --- a/examples/77_blackwell_fmha/CMakeLists.txt +++ b/examples/77_blackwell_fmha/CMakeLists.txt @@ -112,7 +112,7 @@ set(TEST_MLA_FUSE_REDUCTION --b=1 --k=4096 --split_kv=8 --page=128 --fuse_reduct set(TEST_MLA_LARGE_SPLIT_KV --verify --split_kv=20 --page=128) -if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC_ARCHS MATCHES 100a)) +if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC_ARCHS MATCHES 100a OR CUTLASS_NVCC_ARCHS MATCHES 103a)) foreach(PREC fp8 fp16) string(TOUPPER "${PREC}" PREC_MACRO) diff --git a/examples/77_blackwell_fmha/README.md b/examples/77_blackwell_fmha/README.md index 4db02983..06062b54 100644 --- a/examples/77_blackwell_fmha/README.md +++ b/examples/77_blackwell_fmha/README.md @@ -69,6 +69,8 @@ For detailed information on how to invoke them, check out either the tests in `C * 4.3.0: For variable sequence length, the code requires a batch of valid (but never used) padding memory ahead of the first output batch. No padding is needed for the input tensor, but it requires that the input tensor contain no NaN or Inf values. Note that users should set `total_length` to the `problem_shape`. +* 4.4.0: Added support for Blackwell Ultra (Sm103). + # Copyright Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. diff --git a/examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_convert.hpp b/examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_convert.hpp index 263a23b4..b5c6a7ce 100644 --- a/examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_convert.hpp +++ b/examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_convert.hpp @@ -69,7 +69,7 @@ struct FmhaKernelBwdConvert { static const int MinBlocksPerMultiprocessor = 1; static const int MaxThreadsPerBlock = 128; - using ArchTag = cutlass::arch::Sm90; + using ArchTag = cutlass::arch::Sm100; static const int kBlockSeq = 8; diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp index d916f6e5..d6efb2ec 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp @@ -1508,7 +1508,8 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized { CUTLASS_DEVICE void operator()(Params const& params, char* smem) { -#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) +#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + ! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else int warp_idx = cutlass::canonical_warp_idx_sync(); diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp index be3ec8e7..0e63dce1 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp @@ -1480,7 +1480,8 @@ struct Sm100FmhaBwdMlaKernelTmaWarpSpecialized { CUTLASS_DEVICE void operator()(Params const& params, char* smem) { -#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) +#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + ! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else int warp_idx = cutlass::canonical_warp_idx_sync(); diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp index da20e7dd..b40438d8 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp @@ -251,7 +251,8 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized { } CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { -#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) +#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + ! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp index 5fc99444..bbb1b58d 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp @@ -247,7 +247,8 @@ struct Sm100FmhaGenKernelWarpspecialized { } CUTLASS_DEVICE void operator()(const Params ¶ms, char* smem) { -#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) +#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + ! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else diff --git a/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp index 1f51b5d3..083062f5 100644 --- a/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp +++ b/examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp @@ -507,7 +507,8 @@ struct Sm100FmhaMlaKernelTmaWarpspecialized { CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) { -#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED)) +#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \ + ! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED)) CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n"); #else diff --git a/examples/81_blackwell_gemm_blockwise/README.md b/examples/81_blackwell_gemm_blockwise/README.md index f42faec0..b5af130e 100644 --- a/examples/81_blackwell_gemm_blockwise/README.md +++ b/examples/81_blackwell_gemm_blockwise/README.md @@ -50,7 +50,7 @@ To determine the most performance Blockwise/Groupwise GEMM or Grouped GEMM kerne All Blockwise/Groupwise GEMMs and Group GEMMs with `f32` scaling of `e4m3` or runtime `f8` types can be selected by selecting a subset of kernels when configuring with CMake by passing: -`-DCUTLASS_LIBRARY_KERNELS="cutlass3x*f32xe4m3_*f32xe4m3*,cutlass3x*f32xf8_*f32xf8*"` (you can further reduce the amount of kernels generated by specifying the SFA and SFB scale granularities e.g., `cutlass3x*1x128f32xe4m3_*128x128f32xe4m3*`). +`-DCUTLASS_LIBRARY_KERNELS="cutlass3x*f32xe4m3_*f32xe4m3*,cutlass3x*f32xf8_*f32xf8*"` you can further reduce the amount of kernels generated by specifying the SFA and SFB scale granularities e.g., `cutlass3x*1x128f32xe4m3_*128x128f32xe4m3*`). The simplest way to use the profiler is to pass `m`, `n`, and `k` as well as your `scale_vec_size_m`, `scale_vec_size_n`, and `scale_vec_size_k`. Passing `enable-best-kernel-for-fixed-shape` will do some autotuning diff --git a/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu b/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu index 29062d8f..76ac2e88 100644 --- a/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu +++ b/examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu @@ -32,7 +32,7 @@ /*! \file \brief A GEMM example using CUTLASS for the NVIDIA Blackwell SM103 architecture. - This example demonstrates a simple way to instantiate and run a blockscaled 3xFP4 GEMM on the NVIDIA Blackwell SM103 architecture. + This example demonstrates a simple way to instantiate and run a blockscaled ultra FP4 GEMM on the NVIDIA Blackwell SM103 architecture. Usage: @@ -269,7 +269,7 @@ struct Options { std::ostream & print_usage(std::ostream &out) const { out << "89_sm103_fp4_ultra_gemm\n\n" - << " Sm103 3xFP4 GEMM using a Warp Specialized kernel.\n\n" + << " Sm103 ultra FP4 GEMM using a Warp Specialized kernel.\n\n" << "Options:\n\n" << " --help If specified, displays this usage statement\n\n" << " --m= Sets the M extent of the GEMM\n" diff --git a/examples/90_sm103_fp4_ultra_grouped_gemm/90_sm103_fp4_ultra_grouped_gemm.cu b/examples/90_sm103_fp4_ultra_grouped_gemm/90_sm103_fp4_ultra_grouped_gemm.cu index 89850d4b..2067c45c 100644 --- a/examples/90_sm103_fp4_ultra_grouped_gemm/90_sm103_fp4_ultra_grouped_gemm.cu +++ b/examples/90_sm103_fp4_ultra_grouped_gemm/90_sm103_fp4_ultra_grouped_gemm.cu @@ -441,7 +441,7 @@ struct Options { std::ostream & print_usage(std::ostream &out) const { out << "90_sm103_fp4_ultra_grouped_gemm\n\n" - << " Sm103 3xFP4 Grouped GEMM using a Warp Specialized kernel.\n\n" + << " Sm103 ultra FP4 Grouped GEMM using a Warp Specialized kernel.\n\n" << "Options:\n\n" << " --help If specified, displays this usage statement\n\n" << " --m= Sets the M extent of the GEMM for all groups\n" @@ -963,7 +963,7 @@ int run(Options &options, bool host_problem_shapes_available = true) int main(int argc, char const **args) { std::cout << "\n====================================================" << std::endl; - std::cout << "CUTLASS 3.0 Grouped GEMM Example - 3xfp4 Block Scaled" << std::endl; + std::cout << "CUTLASS 3.0 Grouped GEMM Example - ultra fp4 Block Scaled" << std::endl; std::cout << "====================================================" << std::endl; // CUTLASS must be compiled with CUDA 12.9 Toolkit to run this example diff --git a/examples/93_blackwell_low_latency_gqa/CMakeLists.txt b/examples/93_blackwell_low_latency_gqa/CMakeLists.txt new file mode 100644 index 00000000..8f708db1 --- /dev/null +++ b/examples/93_blackwell_low_latency_gqa/CMakeLists.txt @@ -0,0 +1,36 @@ +# Copyright (c) 2014 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +if (NOT MSVC AND CUTLASS_NVCC_ARCHS MATCHES "100a|100f|103a|103f") + +cutlass_example_add_executable( + 93_blackwell_low_latency_gqa + tgv_gqa.cu +) + +endif() diff --git a/examples/93_blackwell_low_latency_gqa/figures/acc2.png b/examples/93_blackwell_low_latency_gqa/figures/acc2.png new file mode 100644 index 00000000..d9e1490f Binary files /dev/null and b/examples/93_blackwell_low_latency_gqa/figures/acc2.png differ diff --git a/examples/93_blackwell_low_latency_gqa/figures/cta.png b/examples/93_blackwell_low_latency_gqa/figures/cta.png new file mode 100644 index 00000000..17f0281f Binary files /dev/null and b/examples/93_blackwell_low_latency_gqa/figures/cta.png differ diff --git a/examples/93_blackwell_low_latency_gqa/figures/fmax.png b/examples/93_blackwell_low_latency_gqa/figures/fmax.png new file mode 100644 index 00000000..d267829b Binary files /dev/null and b/examples/93_blackwell_low_latency_gqa/figures/fmax.png differ diff --git a/examples/93_blackwell_low_latency_gqa/figures/fsum.png b/examples/93_blackwell_low_latency_gqa/figures/fsum.png new file mode 100644 index 00000000..88648187 Binary files /dev/null and b/examples/93_blackwell_low_latency_gqa/figures/fsum.png differ diff --git a/examples/93_blackwell_low_latency_gqa/figures/tgv_gqa.png b/examples/93_blackwell_low_latency_gqa/figures/tgv_gqa.png new file mode 100644 index 00000000..6307eafe Binary files /dev/null and b/examples/93_blackwell_low_latency_gqa/figures/tgv_gqa.png differ diff --git a/examples/93_blackwell_low_latency_gqa/readme.md b/examples/93_blackwell_low_latency_gqa/readme.md new file mode 100644 index 00000000..09c62508 --- /dev/null +++ b/examples/93_blackwell_low_latency_gqa/readme.md @@ -0,0 +1,73 @@ +# Blackwell Low Latency GQA + +This example introduces TGV GQA, a CuTe C++-based Blackwell kernel optimized for low latency (low batch) generation phase GQA. +To compile and run this example: +```bash +# in cutlass top level directory +mkdir build && cd build +cmake .. -DCUTLASS_NVCC_ARCHS=100a -DCUTLASS_ENABLE_TESTS=OFF -DCUTLASS_ENABLE_EXAMPLES=ON -DCUTLASS_ENABLE_LIBRARY=OFF +cd examples/93_blackwell_low_latency_gqa +make +./93_blackwell_low_latency_gqa --kvL 8192 --kvH 8 --qH 64 --BS 1 +``` + +Supported configs are: +- dH = 64/128 +- bf16/fp8 non block scaling kv cache, fp32/bf16/fp8 output +- QKVO are all dH major +- Arbitrary seq len and batch size +- CUDA graph support +- Flash decoding, configurable number of splits +- Cluster reduction with configurable number of reduction cta +- Attention sink and sliding window + +Unsupported features are: +- Persistent schedule +- MTP +- Paged KV cache + +## Kernel Design + +Each cluster (of size `1x1xMAX_SPLITS`) works on a single kv head in a batch. +It divides the kvL (kv sequence length) evenly into each CTA in the cluster in flash decoding style. +And the final reduction is performed by `NUM_REDUCTION_CTA` number of CTAs in the cluster in parallel. + +![cta](./figures/cta.png) + +The example figure above shows a problem size of 1 batch, 2 KV heads, 4 Q heads. +Each cluster processes 1 batch, 1 KV head, 2 Q heads. +And the KV sequence length is divided into 6 128-length tiles, and we evenly distribute the 6 tiles to 4 CTAs in the cluster. +CTA0 and CTA1 will get 2 tiles while CTA2 and CTA3 will get 1 tile each. +In the reduction phase of flash decoding, in this configuration, only 2 CTAs in the cluster will participate in the final reduction. +So each of CTA in the cluster will send their partial results (`Acc2`) to the 2 reduction CTAs in the cluster. + +We have 7 warps in total, 1 for DMA_Q, 1 for DMA_KV, 1 for MMA, 4 for EPILOG (softmax + cluster reduction). +The imagine below shows the how the data flows across the warps as well as how the control dependencies are established between the warps. + +![tgv_gqa](./figures/tgv_gqa.png) + +## Fmax Reduction Mapping + +We want to get fmax for each column (q token) of the fmax tensor. +Each thread individually get fmax for all q tokens using credux (and inter warp reduction). +Then T0,32,64,96 stores the local fmax to destination cta's dsmem for whole cluster wide fmax reduction. +Each reduction cta will hold the cluster wide fmax values of 2 q tokens (8 q tokens in total, divided by 4 reduction ctas). + +![fmax_mapping](./figures/fmax.png) + +## Fsum Reduction Mapping + +We want to get fsum for each column (q token) of the fsum tensor. +Do a reswizzle of fsum tensor in rmem through smem such that each thread holds a partial column of the fsum tensor. +Then do intra-thread and intra-warp reduction to get the fsum for each column (q token). +Then T0,32,64,96 stores the local fsum to destination cta's dsmem for whole cluster wide fsum reduction. +Each reduction cta will hold the cluster wide fsum values of 2 q tokens (8 q tokens in total, divided by 4 reduction ctas). + +![fsum_mapping](./figures/fsum.png) + +## Acc2 Reduction Mapping + +Each thread in the reduction cta will be responsible for generating 2 q tokens in the final output tensor (8 q tokens in total, divided by 4 reduction ctas). +The reduction is a scaled accumulation (similar to the correction step in attention mainloop). + +![acc2_mapping](./figures/acc2.png) diff --git a/examples/93_blackwell_low_latency_gqa/tgv_gqa.cu b/examples/93_blackwell_low_latency_gqa/tgv_gqa.cu new file mode 100644 index 00000000..90f847bb --- /dev/null +++ b/examples/93_blackwell_low_latency_gqa/tgv_gqa.cu @@ -0,0 +1,750 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Example implementation of low latency GQA for the NVIDIA Blackwell SM100/SM103 + architecture using CUTLASS 3. + + Input tensor shapes: + K has shape (kvL, dH, kvH, BS) + Q has shape ((qHLocal, qL), dH, kvH, BS) + V has shape (dH, kvL, kvH, BS) + O has shape (dH, (qHLocal, qL), kvH, BS) + kvL is max_seq_len, seq_lens[BS] is the actual seq len for each batch + sinks has shape (qHLocal * kvH), i.e. one sink per q head + + Example usage: + $ ./examples/93_blackwell_low_latency_gqa --kvL 8192 --kvH 8 --qH 64 --BS 1 +*/ + +// Standard library includes +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +// Use Thrust to handle host/device allocations +#include +#include + +// Cutlass includes +#include +#include +#include + +// CuTe includes +#include // CuTe tensor implementation + +#include "tgv_gqa.cuh" + +using namespace cute; + +// K (kvL, dH, kvH, BS) +// Q ((qHLocal, qL), dH, kvH, BS) +// V (dH, kvL, kvH, BS) +// O (dH, (qHLocal, qL), kvH, BS) +// seq_lens (BS) +// Sinks ((qHLocal, qL), kvH) +template < + class TypeAcc, + int CTA_kvL, + bool NoSink, + class TensorQ, + class TensorK, + class TensorV, + class TensorO, + class TensorSinks> +void +reference_gqa( + TensorK const& tensor_K, + TensorQ const& tensor_Q, + TensorV const& tensor_V, + TensorO const& tensor_O, + int* seq_lens, + TensorSinks const& tensor_Sinks, + float softmax_scale, + int sliding_window_size) { + + using TypeQKV = typename TensorQ::element_type; + using TypeO = typename TensorO::element_type; + + using namespace cute; + int kvL = size<0>(tensor_K); + int dH = size<1>(tensor_K); + int kvH = size<2>(tensor_K); + int BS = size<3>(tensor_K); + int qHLocal = size<0>(shape<0>(tensor_Q)); + int qL = size<1>(shape<0>(tensor_Q)); + + // reference code don't handle oob either + //assert(kvL % CTA_kvL == 0); + int MaxKVBlocks = cutlass::ceil_div(kvL, CTA_kvL); + + // allocate intermediate tensors + thrust::host_vector host_Acc1(kvL * qHLocal * qL * kvH * BS); + auto tensor_Acc1 = make_tensor(host_Acc1.data(), make_layout(make_shape(kvL, make_shape(qHLocal, qL), kvH, BS))); // (kvL, (qHLocal, qL), kvH, BS) + + // CTA level fmax and fsum + thrust::host_vector host_Fmax(MaxKVBlocks * qHLocal * qL * kvH * BS); + Tensor tensor_Fmax = make_tensor(host_Fmax.data(), make_layout(make_shape(MaxKVBlocks, make_shape(qHLocal, qL), kvH, BS))); // (MaxKVBlocks, (qHLocal, qL), kvH, BS) + fill(tensor_Fmax, -cutlass::platform::numeric_limits::infinity()); + + thrust::host_vector host_Fsum(MaxKVBlocks * qHLocal * qL * kvH * BS); + Tensor tensor_Fsum = make_tensor(host_Fsum.data(), make_layout(make_shape(MaxKVBlocks, make_shape(qHLocal, qL), kvH, BS))); // (MaxKVBlocks, (qHLocal, qL), kvH, BS) + clear(tensor_Fsum); + + thrust::host_vector host_P(CTA_kvL * MaxKVBlocks * qHLocal * qL * kvH * BS); + Tensor tensor_P = make_tensor(host_P.data(), make_layout(make_shape(make_shape(CTA_kvL, MaxKVBlocks), make_shape(qHLocal, qL), kvH, BS))); // ((CTA_kvL, MaxKVBlocks), (qHLocal, qL), kvH, BS) + + thrust::host_vector host_P_converted(CTA_kvL * MaxKVBlocks * qHLocal * qL * kvH * BS); + Tensor tensor_P_converted = make_tensor(host_P_converted.data(), make_layout(make_shape(make_shape(CTA_kvL, MaxKVBlocks), make_shape(qHLocal, qL), kvH, BS))); // ((CTA_kvL, MaxKVBlocks), (qHLocal, qL), kvH, BS) + + thrust::host_vector host_Acc2(dH * qHLocal * qL * kvH * MaxKVBlocks * BS); + auto tensor_Acc2 = make_tensor(host_Acc2.data(), make_layout(make_shape(dH, make_shape(qHLocal, qL), kvH, BS, MaxKVBlocks))); // (dH, (qHLocal, qL), kvH, BS, MaxKVBlocks) + + // cluster level fmax and fsum + thrust::host_vector host_Fmax_cluster(qHLocal * qL * kvH * BS); + Tensor tensor_Fmax_cluster = make_tensor(host_Fmax_cluster.data(), make_layout(make_shape(make_shape(qHLocal, qL), kvH, BS))); // ((qHLocal, qL), kvH, BS) + + thrust::host_vector host_Beta(qHLocal * qL * kvH * MaxKVBlocks * BS); + Tensor tensor_Beta = make_tensor(host_Beta.data(), make_layout(make_shape(make_shape(qHLocal, qL), kvH, BS, MaxKVBlocks))); // ((qHLocal, qL), kvH, BS, MaxKVBlocks) + + thrust::host_vector host_Fsum_cluster(qHLocal * qL * kvH * BS); + Tensor tensor_Fsum_cluster = make_tensor(host_Fsum_cluster.data(), make_layout(make_shape(make_shape(qHLocal, qL), kvH, BS))); // ((qHLocal, qL), kvH, BS) + + float softmax_scale_log2 = softmax_scale * static_cast(M_LOG2E); + + for (int _BS = 0; _BS < BS; ++_BS) { + int seq_len = seq_lens[_BS]; + int NumKVBlocks = cutlass::ceil_div(seq_len, CTA_kvL); + int kvL_start = (sliding_window_size == 0) ? 0 : std::max(0, seq_len - sliding_window_size); + int kvBlock_start = kvL_start / CTA_kvL; + + // bmm1 s = q * k * softmax_scale_log2 + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + int start = _kvBlock * CTA_kvL; + int end = std::min(start + CTA_kvL, seq_len); + assert(start < end); + for (int _kvL = start; _kvL < end; ++_kvL) { + TypeAcc acc = TypeAcc(0.f); + for (int _dH = 0; _dH < dH; ++_dH) { + acc += tensor_K(_kvL, _dH, _kvH, _BS) * tensor_Q(make_coord(_qHLocal, _qL), _dH, _kvH, _BS); + } + tensor_Acc1(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS) = (_kvL < kvL_start) ? -INFINITY : (acc * softmax_scale_log2); + } + } + } + } + } + + // calculate m_ij + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + int start = _kvBlock * CTA_kvL; + int end = std::min(start + CTA_kvL, seq_len); + assert(start < end); + TypeAcc& fmax = tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS); + fmax = tensor_Acc1(start, make_coord(_qHLocal, _qL), _kvH, _BS); + for (int _kvL = start + 1; _kvL < end; ++_kvL) { + fmax = std::max(fmax, tensor_Acc1(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS)); + } + } + } + } + } + + // calculate p = exp2f(s - m_ij) + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + int start = _kvBlock * CTA_kvL; + int end = std::min(start + CTA_kvL, seq_len); + assert(start < end); + for (int _kvL = start; _kvL < end; ++_kvL) { + tensor_P(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS) = std::exp2f(tensor_Acc1(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS) - tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS)); + } + } + } + } + } + + // calculate l_ij = sum(p) + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + int start = _kvBlock * CTA_kvL; + int end = std::min(start + CTA_kvL, seq_len); + assert(start < end); + TypeAcc sum = TypeAcc(0.f); + for (int _kvL = start; _kvL < end; ++_kvL) { + sum += tensor_P(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS); + } + tensor_Fsum(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS) = sum; + } + } + } + } + + // convert P from fp32 to bf16/fp8 + cutlass::NumericConverter converter_p; + for (int i = 0; i < tensor_P(_,_,_,_BS).size(); i++) { + tensor_P_converted(_,_,_,_BS)[i] = converter_p(tensor_P(_,_,_,_BS)[i]); + } + + // bmm2 acc2 = v * p + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _dH = 0; _dH < dH; ++_dH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + int start = _kvBlock * CTA_kvL; + int end = std::min(start + CTA_kvL, seq_len); + assert(start < end); + TypeAcc acc = TypeAcc(0.f); + for (int _kvL = start; _kvL < end; ++_kvL) { + acc += tensor_V(_dH, _kvL, _kvH, _BS) * tensor_P_converted(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS); + } + tensor_Acc2(_dH, make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock) = acc; + } + } + } + } + } + + // calculate cluster level fmax and fsum + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + TypeAcc& fmax = tensor_Fmax_cluster(make_coord(_qHLocal, _qL), _kvH, _BS); + fmax = tensor_Fmax(kvBlock_start, make_coord(_qHLocal, _qL), _kvH, _BS); + for (int _kvBlock = kvBlock_start + 1; _kvBlock < NumKVBlocks; ++_kvBlock) { + fmax = std::max(fmax, tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS)); + } + + // calculate beta + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + tensor_Beta(make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock) = std::exp2f(tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS) - fmax); + } + + // calculate fsum + TypeAcc sum = TypeAcc(0.f); + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + sum += tensor_Fsum(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS) * tensor_Beta(make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock); + } + tensor_Fsum_cluster(make_coord(_qHLocal, _qL), _kvH, _BS) = sum; + if constexpr (!NoSink) { + tensor_Fsum_cluster(make_coord(_qHLocal, _qL), _kvH, _BS) += std::exp2f(tensor_Sinks(make_coord(_qHLocal, _qL), _kvH) * (float)M_LOG2E - fmax); + } + } + } + } + + // convert O from fp32 to bf16/fp8 + cutlass::NumericConverter converter_o; + // final reduction + for (int _kvH = 0; _kvH < kvH; ++_kvH) { + for (int _dH = 0; _dH < dH; ++_dH) { + for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) { + for (int _qL = 0; _qL < qL; ++_qL) { + TypeAcc acc = TypeAcc(0.f); + for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) { + acc += tensor_Acc2(_dH, make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock) * tensor_Beta(make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock); + } + tensor_O(_dH, make_coord(_qHLocal, _qL), _kvH, _BS) = converter_o(acc / tensor_Fsum_cluster(make_coord(_qHLocal, _qL), _kvH, _BS)); + } + } + } + } + } + + /*int example_row = 0; + int example_kvH = 0; + int example_kvBlock = 0; + int example_BS = 3; + print("tensor_Acc1:\t"); print(tensor_Acc1(_, _, example_kvH, example_BS)); print("\n"); + print("tensor_Fmax:\t"); print_tensor(tensor_Fmax(example_kvBlock,_,example_kvH,example_BS)); print("\n"); + print("tensor_P:\t"); print_tensor(tensor_P(make_coord(example_row, example_kvBlock),_,example_kvH,example_BS)); print("\n"); + print("tensor_Fsum:\t"); print_tensor(tensor_Fsum(example_kvBlock,_,example_kvH,example_BS)); print("\n"); + print("tensor_P_converted:\t"); print_tensor(tensor_P_converted(make_coord(example_row, example_kvBlock),_,example_kvH,example_BS)); print("\n"); + print("tensor_Acc2:\t"); print_tensor(tensor_Acc2(example_row,_,example_kvH,example_BS,example_kvBlock)); print("\n"); + print("tensor_Fmax_cluster:\t"); print_tensor(tensor_Fmax_cluster(_,example_kvH,example_BS)); print("\n"); + print("tensor_Beta:\t"); print_tensor(tensor_Beta(_,example_kvH,example_BS,_)); print("\n"); + print("tensor_Fsum_cluster:\t"); print_tensor(tensor_Fsum_cluster(_,example_kvH,example_BS)); print("\n"); + print("tensor_O:\t"); print_tensor(tensor_O(example_row,_,example_kvH,example_BS)); print("\n");*/ +} + +template +void +initialize_tensor(Tensor& tensor, cute::tuple value_range = {-4, 4}) { + using DataType = typename Tensor::element_type; + auto [min, max] = value_range; + for (int i = 0; i < cute::size(tensor); i++) { + tensor(i) = DataType(int((max-min)*(rand() / double(RAND_MAX)) + min)); + } +} + +// Compares two CuTe tensors with torch.allclose semantics +// Returns true if: |input_i - other_i| <= atol + rtol x |other_i| for all elements +template +bool +cute_allclose( + TensorInput const& input, + TensorOther const& other, + float rtol = 1e-05f, + float atol = 1e-08f, + bool equal_nan = false) { + using namespace cute; + + // Tensors must have the same size + if (size(input) != size(other)) { + std::cerr << "Error: Tensor sizes don't match. input size: " << size(input) + << ", other size: " << size(other) << std::endl; + return false; + } + + int mismatches = 0; + const int max_print = 10; // Only print first 10 mismatches + + for (int i = 0; i < size(input); ++i) { + float input_val = float(input(i)); + float other_val = float(other(i)); + + // Handle NaN comparison + bool input_is_nan = std::isnan(input_val); + bool other_is_nan = std::isnan(other_val); + + if (input_is_nan || other_is_nan) { + if (equal_nan && input_is_nan && other_is_nan) { + continue; // Both NaN and equal_nan is true + } + else if (input_is_nan || other_is_nan) { + if (mismatches < max_print) { + std::cerr << "Mismatch at index " << i << ": input=" << input_val + << ", other=" << other_val << " (NaN detected)" << std::endl; + } + mismatches++; + continue; + } + } + + // Check torch.allclose condition: |input - other| <= atol + rtol * |other| + float diff = std::abs(input_val - other_val); + float threshold = atol + rtol * std::abs(other_val); + + if (diff > threshold) { + if (mismatches < max_print) { + std::cerr << "Mismatch at index " << i << ": input=" << input_val + << ", other=" << other_val << ", diff=" << diff + << ", threshold=" << threshold << std::endl; + } + mismatches++; + } + } + + if (mismatches > 0) { + std::cerr << "Total mismatches: " << mismatches << " out of " << size(input) << " elements" << std::endl; + return false; + } + + return true; +} + +struct ProblemStride { + int stride_Q_kvH; + int stride_Q_qHLocal; + int stride_Q_qL; + int stride_Q_dH; + int stride_Q_BS; + + int stride_K_kvH; + int stride_K_kvL; + int stride_K_dH; + int stride_K_BS; + + int stride_V_kvH; + int stride_V_kvL; + int stride_V_dH; + int stride_V_BS; + + int stride_O_kvH; + int stride_O_qHLocal; + int stride_O_qL; + int stride_O_dH; + int stride_O_BS; +}; + +ProblemStride make_gqa_stride(int kvH, int qHLocal, int qL, int kvL, int dH, int BS) { + ProblemStride stride; + + // Q shape ((qHLocal, qL), dH, kvH, BS), where dH is contiguous + // slowest moving dim->fastest moving dim: BS, qL, kvH, qHLocal, dH + stride.stride_Q_kvH = qHLocal * dH; + stride.stride_Q_qHLocal = dH; + stride.stride_Q_qL = kvH * qHLocal * dH; + stride.stride_Q_dH = 1; + stride.stride_Q_BS = kvH * qHLocal * dH * qL; + + // K shape (kvL, dH, kvH, BS), where dH is contiguous + // slowest moving dim->fastest moving dim: BS, kvL, kvH, dH + stride.stride_K_kvH = dH; + stride.stride_K_kvL = kvH * dH; + stride.stride_K_dH = 1; + stride.stride_K_BS = kvL * dH * kvH; + + // V shape (dH, kvL, kvH, BS), where dH is contiguous + // slowest moving dim->fastest moving dim: BS, kvL, kvH, dH + stride.stride_V_kvH = dH; + stride.stride_V_kvL = kvH * dH; + stride.stride_V_dH = 1; + stride.stride_V_BS = kvL * dH * kvH; + + // O shape (dH, (qHLocal, qL), kvH, BS), where dH is contiguous + // slowest moving dim->fastest moving dim: BS, qL, kvH, qHLocal, dH + stride.stride_O_kvH = qHLocal * dH; + stride.stride_O_qHLocal = dH; + stride.stride_O_qL = kvH * qHLocal * dH; + stride.stride_O_dH = 1; + stride.stride_O_BS = kvH * qHLocal * dH * qL; + + return stride; +} + +class GQATester { +public: + // Kernel config constants + using TypeQKV = cutlass::bfloat16_t; + using TypeO = cutlass::bfloat16_t; + using TypeAcc = float; + static constexpr int CTA_qHLocal = 8; + static constexpr int CTA_qL = 1; + static constexpr int CTA_kvL = 128; + static constexpr int CTA_dH = 64; + static constexpr int BMM1_DMA_Stage = 3; + static constexpr int BMM2_DMA_Stage = 3; + static constexpr int MaxSplits = 8; + static constexpr int NumReductionCTA = 8; + static constexpr bool NoSink = true; + +private: + int kvH_, qHLocal_, qL_, kvL_, dH_, BS_; + float softmax_scale_; + ProblemStride stride_; + int sliding_window_size_; + + // Host vectors + thrust::host_vector host_Q_; + thrust::host_vector host_K_; + thrust::host_vector host_V_; + thrust::host_vector host_O_; + thrust::host_vector host_reference_O_; + thrust::host_vector host_seq_lens_; + thrust::host_vector host_sinks_; + + // Device vectors + thrust::device_vector device_Q_; + thrust::device_vector device_K_; + thrust::device_vector device_V_; + thrust::device_vector device_O_; + thrust::device_vector device_seq_lens_; + thrust::device_vector device_sinks_; + +public: + GQATester(int kvH, int qH, int qL, int kvL, int dH, int BS, float softmax_scale, int sliding_window_size) : + kvH_(kvH), qHLocal_(qH / kvH), qL_(qL), kvL_(kvL), dH_(dH), BS_(BS), softmax_scale_(softmax_scale), sliding_window_size_(sliding_window_size) { + assert(sliding_window_size_ >= 0); + // Allocate host memory + host_Q_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_); + host_K_.resize(kvH_ * kvL_ * dH_ * BS_); + host_V_.resize(kvH_ * kvL_ * dH_ * BS_); + host_O_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_); + host_reference_O_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_); + host_seq_lens_.resize(BS_); + host_sinks_.resize(qHLocal_ * kvH_); // one sink per q head + + stride_ = make_gqa_stride(kvH_, qHLocal_, qL_, kvL_, dH_, BS_); + + // Create host CuTe tensors for initialization + auto host_tensor_Q = make_tensor(host_Q_.data(), TGV::gqa::make_layout_Q(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_Q_kvH, stride_.stride_Q_qHLocal, stride_.stride_Q_qL, stride_.stride_Q_dH, stride_.stride_Q_BS)); + auto host_tensor_K = make_tensor(host_K_.data(), TGV::gqa::make_layout_K(kvH_, kvL_, dH_, BS_, stride_.stride_K_kvH, stride_.stride_K_kvL, stride_.stride_K_dH, stride_.stride_K_BS)); + auto host_tensor_V = make_tensor(host_V_.data(), TGV::gqa::make_layout_V(kvH_, kvL_, dH_, BS_, stride_.stride_V_kvH, stride_.stride_V_kvL, stride_.stride_V_dH, stride_.stride_V_BS)); + auto host_tensor_sinks = make_tensor(host_sinks_.data(), TGV::gqa::make_layout_sinks(qHLocal_, qL_, kvH_)); + + // Initialize Q, K, V tensors with random values + initialize_tensor(host_tensor_Q); + initialize_tensor(host_tensor_K); + initialize_tensor(host_tensor_V); + // have batch size matching kvL (i.e. max seq len) for now + bool test_var_seq_lens = false; + for (int i = 0; i < BS_; ++i) { + if (test_var_seq_lens) { + host_seq_lens_[i] = rand() % kvL_ + 1; + } + else { // all the batch have the same seq len + host_seq_lens_[i] = kvL_; + } + } + for (int i = 0; i < qHLocal_ * kvH_; ++i) { + host_sinks_[i] = rand() / (float)RAND_MAX; + } + + // Allocate device memory and copy H2D + device_Q_ = host_Q_; + device_K_ = host_K_; + device_V_ = host_V_; + device_O_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_); + device_seq_lens_ = host_seq_lens_; + device_sinks_ = host_sinks_; + + gpuErrChk(cudaDeviceSynchronize()); + } + + void run_kernel(bool pdl, int pdl_count = -1, cudaStream_t stream = 0) { + TGV::gqa::gqa_host< + TypeQKV, TypeO, TypeAcc, + CTA_qHLocal, CTA_qL, CTA_kvL, CTA_dH, + BMM1_DMA_Stage, BMM2_DMA_Stage, + MaxSplits, + NumReductionCTA>( + device_K_.data().get(), device_Q_.data().get(), device_V_.data().get(), device_O_.data().get(), + device_seq_lens_.data().get(), + NoSink ? nullptr : device_sinks_.data().get(), + kvH_, qHLocal_, qL_, kvL_, dH_, BS_, + stride_.stride_K_kvH, stride_.stride_K_kvL, stride_.stride_K_dH, stride_.stride_K_BS, + stride_.stride_Q_kvH, stride_.stride_Q_qHLocal, stride_.stride_Q_qL, stride_.stride_Q_dH, stride_.stride_Q_BS, + stride_.stride_V_kvH, stride_.stride_V_kvL, stride_.stride_V_dH, stride_.stride_V_BS, + stride_.stride_O_kvH, stride_.stride_O_qHLocal, stride_.stride_O_qL, stride_.stride_O_dH, stride_.stride_O_BS, + softmax_scale_, + sliding_window_size_, + pdl, pdl_count, stream); + } + + bool verify() { + // Run the GPU kernel + run_kernel(false); + gpuErrChk(cudaDeviceSynchronize()); + + // Copy D2H + host_O_ = device_O_; + + // Create tensors for verification using helper methods + auto host_tensor_Q = make_tensor(host_Q_.data(), TGV::gqa::make_layout_Q(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_Q_kvH, stride_.stride_Q_qHLocal, stride_.stride_Q_qL, stride_.stride_Q_dH, stride_.stride_Q_BS)); + auto host_tensor_K = make_tensor(host_K_.data(), TGV::gqa::make_layout_K(kvH_, kvL_, dH_, BS_, stride_.stride_K_kvH, stride_.stride_K_kvL, stride_.stride_K_dH, stride_.stride_K_BS)); + auto host_tensor_V = make_tensor(host_V_.data(), TGV::gqa::make_layout_V(kvH_, kvL_, dH_, BS_, stride_.stride_V_kvH, stride_.stride_V_kvL, stride_.stride_V_dH, stride_.stride_V_BS)); + auto host_tensor_O = make_tensor(host_O_.data(), TGV::gqa::make_layout_O(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_O_kvH, stride_.stride_O_qHLocal, stride_.stride_O_qL, stride_.stride_O_dH, stride_.stride_O_BS)); + auto host_reference_tensor_O = make_tensor(host_reference_O_.data(), TGV::gqa::make_layout_O(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_O_kvH, stride_.stride_O_qHLocal, stride_.stride_O_qL, stride_.stride_O_dH, stride_.stride_O_BS)); + auto host_seq_lens_tensor = make_tensor(host_seq_lens_.data(), make_layout(make_shape(BS_))); + auto host_tensor_sinks = make_tensor(host_sinks_.data(), TGV::gqa::make_layout_sinks(qHLocal_, qL_, kvH_)); + + gpuErrChk(cudaDeviceSynchronize()); + + print("host_seq_lens_tensor:\t"); print_tensor(host_seq_lens_tensor); + print("host_tensor_Q:\t"); print(host_tensor_Q); print("\n"); + print("host_tensor_K:\t"); print(host_tensor_K); print("\n"); + print("host_tensor_V:\t"); print(host_tensor_V); print("\n"); + print("host_tensor_O:\t"); print(host_tensor_O); print("\n"); + print("host_reference_tensor_O:\t"); print(host_reference_tensor_O); print("\n"); + print("host_tensor_sinks:\t"); print(host_tensor_sinks); print("\n"); + + // Execute reference GQA kernel + reference_gqa( + host_tensor_K, host_tensor_Q, host_tensor_V, host_reference_tensor_O, host_seq_lens_.data(), host_tensor_sinks, softmax_scale_, sliding_window_size_); + + // Compare results using torch.allclose semantics + // For bfloat16, use more relaxed tolerances due to reduced precision + bool success = cute_allclose(host_tensor_O, host_reference_tensor_O, + 1e-2f, // rtol (relative tolerance) + 1e-3f); // atol (absolute tolerance) + //print("host_tensor_O:\t"); print_tensor(host_tensor_O(_,_,0)); print("\n"); + //print("host_reference_tensor_O:\t"); print_tensor(host_reference_tensor_O(_,_,0)); print("\n"); + + std::cout << "Execution is " << ((success) ? "successful." : "failed.") << std::endl; + + return success; + } + +}; + +void benchmark_gqa(int kvH, int qH, int qL, int kvL, int dH, int BS, float softmax_scale, int sliding_window_size, bool pdl, int pdl_count, int num_testers = 4, int bench_iters = 100) { + std::cout << "=== GQA Benchmark ===" << std::endl; + std::cout << "Problem size: kvH=" << kvH << ", qH=" << qH << ", qL=" << qL << ", kvL=" << kvL << ", dH=" << dH << ", BS=" << BS << ", sliding_window_size=" << sliding_window_size << std::endl; + std::cout << "Number of testers (L2 thrashing): " << num_testers << std::endl; + std::cout << "Benchmark iterations: " << bench_iters << std::endl; + + // Create multiple tester instances to thrash L2 cache + std::vector> testers; + for (int i = 0; i < num_testers; ++i) { + testers.push_back(std::make_unique(kvH, qH, qL, kvL, dH, BS, softmax_scale, sliding_window_size)); + } + std::cout << "Created " << num_testers << " GQATester instances" << std::endl; + + // Create CUDA stream for graph capture + cudaStream_t stream; + gpuErrChk(cudaStreamCreate(&stream)); + + // Capture CUDA graph + std::cout << "Capturing CUDA graph..." << std::endl; + cudaGraph_t graph; + cudaGraphExec_t graph_exec; + + gpuErrChk(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal)); + // Capture round robin execution pattern + for (int iter = 0; iter < bench_iters; ++iter) { + int tester_idx = iter % num_testers; + // Note: We need to run kernels on the same stream for graph capture + // This requires modifying the kernel launch to accept a stream parameter + // For now, we'll capture a simpler pattern and measure accordingly + testers[tester_idx]->run_kernel(pdl, pdl_count, stream); + } + gpuErrChk(cudaStreamEndCapture(stream, &graph)); + gpuErrChk(cudaGraphInstantiate(&graph_exec, graph, NULL, NULL, 0)); + std::cout << "CUDA graph captured and instantiated" << std::endl; + + // Warmup: run kernels in round robin fashion + std::cout << "Starting warmup..." << std::endl; + gpuErrChk(cudaGraphLaunch(graph_exec, stream)); + gpuErrChk(cudaDeviceSynchronize()); + std::cout << "Warmup completed" << std::endl; + + // Benchmark: replay graph and measure time + std::cout << "Starting benchmark..." << std::endl; + + // Create CUDA events for timing + cudaEvent_t start_event, stop_event; + gpuErrChk(cudaEventCreate(&start_event)); + gpuErrChk(cudaEventCreate(&stop_event)); + + gpuErrChk(cudaProfilerStart()); + gpuErrChk(cudaEventRecord(start_event, stream)); + gpuErrChk(cudaGraphLaunch(graph_exec, stream)); + gpuErrChk(cudaEventRecord(stop_event, stream)); + gpuErrChk(cudaProfilerStop()); + + gpuErrChk(cudaDeviceSynchronize()); + + // Calculate timing + float total_time_ms; + gpuErrChk(cudaEventElapsedTime(&total_time_ms, start_event, stop_event)); + + float avg_time_ms = total_time_ms / bench_iters; + float avg_time_us = avg_time_ms * 1000.0f; + + // Calculate FLOPS + long long ops = 2LL * qH * qL * kvL * dH * 2LL * BS; // 2 ops per multiply-add + double gflops = (ops * bench_iters) / (total_time_ms * 1e6); + double gflops_per_iter = ops / (avg_time_ms * 1e6); + + // Calculate DRAM bandwidth + long long bytes_per_iter = ((long long)kvH * kvL * dH * 2 * sizeof(GQATester::TypeQKV) + (long long)qH * qL * dH * (sizeof(GQATester::TypeQKV) + sizeof(GQATester::TypeO))) * BS; // Q(bf16) + K(bf16) + V(bf16) + O(bf16) + float avg_dram_bw_gbps = (bytes_per_iter / (avg_time_ms / 1000.0f)) / (1024.0f * 1024.0f * 1024.0f); + + // Report results + std::cout << "\n=== Benchmark Results ===" << std::endl; + std::cout << "Average time per iteration: " << avg_time_ms << " ms (" << avg_time_us << " μs)" << std::endl; + std::cout << "GFLOPS per iteration: " << gflops_per_iter << std::endl; + std::cout << "Average DRAM bandwidth: " << avg_dram_bw_gbps << " GB/s" << std::endl; + + // Cleanup + gpuErrChk(cudaGraphExecDestroy(graph_exec)); + gpuErrChk(cudaGraphDestroy(graph)); + gpuErrChk(cudaEventDestroy(start_event)); + gpuErrChk(cudaEventDestroy(stop_event)); + gpuErrChk(cudaStreamDestroy(stream)); + + std::cout << "=== Benchmark Complete ===" << std::endl; +} + + +int main(int argc, char* argv[]) { + srand(time(NULL)); + + int kvH = 8; // num KV head + int qH = 64; // num Q head + int qL = 1; // Q sequence length + int kvL = 2048; // KV sequence length + int dH = 64; // hidden dimension + int BS = 1; // batch size + float softmax_scale = 1.0f / (float)sqrt(dH); + int sliding_window_size = 0; // when sliding_window_size = 0, it's disabled + bool pdl = false; + // don't support it yet + int pdl_count = -1; + + // arg parsing + while (1) { + static struct option long_options[] = { + {"kvL", required_argument, 0, 0}, + {"kvH", required_argument, 0, 0}, + {"qH", required_argument, 0, 0}, + {"qL", required_argument, 0, 0}, + {"BS", required_argument, 0, 0}, + {"sliding_window_size", required_argument, 0, 0}, + {0, 0, 0, 0} // denote end of array + }; + + int option_index = 0; + // M no argument + // M: required argument + // M:: optional argument + int c = getopt_long(argc, argv, "", long_options, &option_index); + + if (c==-1) break; + switch (c) { + case 0: + // Long option + if (option_index == 0) kvL = atoi(optarg); + else if (option_index == 1) kvH = atoi(optarg); + else if (option_index == 2) qH = atoi(optarg); + else if (option_index == 3) qL = atoi(optarg); + else if (option_index == 4) BS = atoi(optarg); + else if (option_index == 5) sliding_window_size = atoi(optarg); + break; + default: assert(false); + } + } + + GQATester tester(kvH, qH, qL, kvL, dH, BS, softmax_scale, sliding_window_size); + bool success = tester.verify(); + std::cout << "Correctness test " << (success ? "PASSED" : "FAILED") << std::endl; + + benchmark_gqa(kvH, qH, qL, kvL, dH, BS, softmax_scale, sliding_window_size, pdl, pdl_count, 100, 1000); +} diff --git a/examples/93_blackwell_low_latency_gqa/tgv_gqa.cuh b/examples/93_blackwell_low_latency_gqa/tgv_gqa.cuh new file mode 100644 index 00000000..303aa876 --- /dev/null +++ b/examples/93_blackwell_low_latency_gqa/tgv_gqa.cuh @@ -0,0 +1,2215 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 +#include +#include +#include +#include + +// Cutlass includes +#include +#include +#include +#include // mma/smem selector, umma::major +#include +#include + +// CuTe includes +#include // CuTe tensor implementation +#include // TMEM allocator for SM100 +#include + +#define gpuErrChk(ans) { gpuAssert2((ans), __FILE__, __LINE__); } +inline void gpuAssert2(cudaError_t code, const char *file, int line, bool abort=true) { + if (code != cudaSuccess) { + fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line); + if (abort) exit(code); + } +} + +namespace TGV { +namespace gqa { + +using namespace cute; + +float constexpr Log2_E = static_cast(M_LOG2E); + +// K has shape (kvL, dH, kvH, BS) +// Q has shape ((qHLocal, qL), dH, kvH, BS) +// V has shape (dH, kvL, kvH, BS) +// O has shape (dH, (qHLocal, qL), kvH, BS) + +// flash decode triton code +/* +# bmm1 +q = tl.load(q_ptrs) # [BLOCK_qL, BLOCK_dH] +k = tl.load(k_ptrs) # [BLOCK_kvL, BLOCK_dH] +qk = tl.dot(q, k.T) * sm_scale # [BLOCK_qL, BLOCK_kvL] + +# softmax(qk), all the variable names are from FA paper +m_ij = tl.max(qk, axis=1) # [BLOCK_qL,] +p = tl.exp(qk - m_ij[:, None]) # [BLOCK_qL, BLOCK_kvL] +# l_ij = exp(qk_n0+1 - m_ij) + ... + exp(qk_n1 - m_ij) +l_ij = tl.sum(p, axis=1) # [BLOCK_qL,] + +# store m_ij and l_ij +tl.store(mij_ptrs, m_ij) +tl.store(lij_ptrs, l_ij) + +# bmm2 +v = tl.load(v_ptrs) # [BLOCK_kvL, BLOCK_dH] +acc = tl.dot(p.to(q.dtype), v) # [BLOCK_qL, BLOCK_dH] +tl.store(acc_ptrs, acc) +*/ + +// Store value to remote shared memory in the cluster +CUTE_DEVICE void +store_shared_remote_f32(float value, uint32_t dsmem_addr, uint32_t remote_barrier_addr) { + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [%0], %1, [%2];" + : : "r"(dsmem_addr), "f"(value), "r"(remote_barrier_addr)); +} + +// given a smem tensor, return the dsmem tensor for the given rank, the tensor addr is in smem addr space (not generic addr space) +template +CUTE_DEVICE auto +get_dsmem_tensor(Tensor tensor, int rank) { + using T = typename decltype(tensor)::value_type; + // tensor.data().get() is the smem addr in the generic addr space, in the generic addr space a region is reserved for smem + // doing ld/st to this region of the generic addr space will be converted into ld.shared/st.shared to the smem addr space by the compiler + // the mapa (and many inline ptx) instruction's input and output addr are in the smem/dsmem addr space, so we need to explicitly convert from generic to shared addr space + uint32_t smem_addr = __cvta_generic_to_shared(tensor.data().get()); // smem addr space + // mapa to get the dsmem addr of this tensor in another CTA + uint32_t dsmem_addr = set_block_rank(smem_addr, rank); // smem addr space + return make_tensor(make_smem_ptr((T*)dsmem_addr), tensor.layout()); +} + +// Helper methods to create layouts +// K always has the shape (kvL, dH, kvH, BS) +// kvH has to be the last dim because we do mma partitioning to the first two dims (M, K) in gemm terminology +// we only index into kvH after cta level partitioning +// it can has arbitrary stride as long as dH (K in gemm terminology) is contiguous +auto make_layout_K( + int kvH, int kvL, int dH, int BS, + int stride_kvH, int stride_kvL, int stride_dH, int stride_BS +) { + assert(stride_dH == 1); + return make_layout(make_shape(kvL, dH, kvH, BS), make_stride(stride_kvL, Int<1>{}, stride_kvH, stride_BS)); +} + +// Q always has the shape ((qHLocal, qL), dH, kvH, BS), qH = kvH * qHLocal +// kvH has to be the last dim because we do mma partitioning to the first two dims (N, K) in gemm terminology +// we only index into kvH after cta level partitioning +// it can has arbitrary stride as long as dH (K in gemm terminology) is contiguous +auto make_layout_Q( + int kvH, int qHLocal, int qL, int dH, int BS, + int stride_kvH, int stride_qHLocal, int stride_qL, int stride_dH, int stride_BS +) { + assert(stride_dH == 1); + return make_layout(make_shape(make_shape(qHLocal, qL), dH, kvH, BS), make_stride(make_stride(stride_qHLocal, stride_qL), Int<1>{}, stride_kvH, stride_BS)); +} + +// V always has the shape (dH, kvL, kvH, BS) +// kvH has to be the last dim because we do mma partitioning to the first two dims (M, K) in gemm terminology +// we only index into kvH after cta level partitioning +// it can has arbitrary stride as long as dH (M in gemm terminology) is contiguous +auto make_layout_V( + int kvH, int kvL, int dH, int BS, + int stride_kvH, int stride_kvL, int stride_dH, int stride_BS +) { + assert(stride_dH == 1); + return make_layout(make_shape(dH, kvL, kvH, BS), make_stride(stride_dH, stride_kvL, stride_kvH, stride_BS)); +} + +// O always has the shape (dH, (qHLocal, qL), kvH, BS), qH = kvH * qHLocal +// kvH has to be the last dim because we do mma partitioning to the first two dims (M, N) in gemm terminology +// we only index into kvH after cta level partitioning +// it can has arbitrary stride as long as dH (M in gemm terminology) is contiguous +auto make_layout_O( + int kvH, int qHLocal, int qL, int dH, int BS, + int stride_kvH, int stride_qHLocal, int stride_qL, int stride_dH, int stride_BS +) { + assert(stride_dH == 1); + return make_layout(make_shape(dH, make_shape(qHLocal, qL), kvH, BS), make_stride(stride_dH, make_stride(stride_qHLocal, stride_qL), stride_kvH, stride_BS)); +} + +// sinks always has the shape ((qHLocal, qL), kvH) +// the real shape is (qHLocal * kvH), but we give qL stride 0 to make CTA indexing cleaner +// so different qL maps to the same sink value +auto make_layout_sinks( + int qHLocal, int qL, int kvH +) { + return make_layout(make_shape(make_shape(qHLocal, qL), kvH), make_stride(make_stride(Int<1>{}, Int<0>{}), qHLocal)); +} + +// simplified from cutlass/include/cutlass/gemm/kernel/static_tile_scheduler.hpp +// allow for easy extention to other types of tile scheduler +struct WorkTileInfo { + int BS_idx = 0; + int kvH_idx = 0; + // each cta process kv block range [kvL_idx_start, kvL_idx_end) + int kvL_idx_start = 0; + int kvL_idx_end = 0; + int dH_idx = 0; // for bmm2 tiling + int qHLocal_idx = 0; + int qL_idx = 0; + bool is_valid_tile = false; + + CUTLASS_HOST_DEVICE bool + is_valid() const { + return is_valid_tile; + } + + CUTLASS_HOST_DEVICE + static WorkTileInfo + invalid_work_tile() { + return {-1, -1, -1, -1, -1, -1, -1, false}; + } +}; + +// The shared memory buffers for Q, K, V matrices. +template < + class TypeQKV, // Tensor Q/K/V data type + class TypeAcc, // Tensor Acc data type + class KSmemLayout, // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + class QSmemLayout, // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + class VSmemLayout, // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + class SSmemLayout, // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) aka C matrix (M, N, 1) for bmm1 + class PSmemLayout, // ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) aka B matrix (N, K, 1) for bmm2 + class WRSmemLayout, // (NumEpilogWarps, (CTA_qHLocal, CTA_qL)), WR stands for warp reduce + class MSMailboxSmemLayout,// (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA), MS stands max and sum + class Acc1SmemLayout, // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + class Acc2MailboxSmemLayout, // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + class SinksSmemLayout, // (CTA_qHLocal * CTA_qL / NumReductionCTA) + int BMM1_DMA_Stage, + int BMM2_DMA_Stage> +struct SharedStorage { + // for bmm1 + alignas(128) cute::ArrayEngine> K; // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + alignas(128) cute::ArrayEngine> Q; // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + // for bmm2 + alignas(128) cute::ArrayEngine> V; // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + alignas(128) cute::ArrayEngine> P; // ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) + // for softmax to get CTA wide max and sum + alignas(128) cute::ArrayEngine> FmaxPartial; // (NumEpilogWarps, (CTA_qHLocal, CTA_qL)) + // for softmax to get cluster wide max and sum + alignas(128) cute::ArrayEngine> FmaxMailbox; // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + alignas(128) cute::ArrayEngine> FsumMailbox; // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + // staging bufferfor making acc1 rmem layout easier for softmax reduction + alignas(128) cute::ArrayEngine> Acc1; // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + // for final flash decode reduction + alignas(128) cute::ArrayEngine> Acc2Mailbox; // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + // for attention sinks + alignas(128) cute::ArrayEngine> Sink; // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + // BMM1 iterate over dH (i.e. K dimension) + alignas(16) cute::uint64_t tma_bmm1_full_barrier[BMM1_DMA_Stage]; // Barrier between TMA/softmax and BMM1, TMA/softmax tells BMM1 the tile is ready/full, BMM1 can start consuming it + alignas(16) cute::uint64_t tma_bmm1_empty_barrier[BMM1_DMA_Stage]; // Barrier between BMM1 and TMA, BMM1 tells TMA the tile is empty, TMA can start loading the next tile + + // BMM2 iterate over kvL (i.e. K dimension) + alignas(16) cute::uint64_t tmasoftmax_bmm2_full_barrier[BMM2_DMA_Stage]; // Barrier between TMA/softmax and BMM2, TMA/softmax tells BMM2 the tile is ready/full, BMM2 can start consuming it + alignas(16) cute::uint64_t tma_bmm2_empty_barrier[BMM2_DMA_Stage]; // Barrier between BMM2 and TMA, BMM2 tells TMA the tile is empty, TMA can start loading the next tile + + alignas(16) cute::uint64_t bmm1_softmax_full_barrier; // Barrier between BMM1 and softmax, BMM1 tells softmax the tile is ready/full, softmax can start consuming it + alignas(16) cute::uint64_t bmm2_epilog_full_barrier; // Barrier between BMM2 and epilog, BMM2 tells epilog the tile is ready/full, epilog can start consuming it + + // for cluster reduction + alignas(16) cute::uint64_t maxsum_mailbox_full_barrier; // barrier indicating the st.async of fmax and fsum are done + alignas(16) cute::uint64_t acc2_mailbox_full_barrier; // barrier indicating the st.async of acc2 are done + + alignas(16) cute::uint32_t bmm1_tmem_base_ptr; // Base pointer for TMEM allocation, BMM1 MMA will allocate TMEM here + alignas(16) cute::uint32_t bmm2_tmem_base_ptr; // Base pointer for TMEM allocation, BMM2 MMA will allocate TMEM here + + // for cp.async to write to + int seq_len; + + CUTE_DEVICE constexpr auto tensor_sK() { return make_tensor(make_smem_ptr(K.begin()), KSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sQ() { return make_tensor(make_smem_ptr(Q.begin()), QSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sV() { return make_tensor(make_smem_ptr(V.begin()), VSmemLayout{}); } + // it's the same tensor P, just different views (layouts) + CUTE_DEVICE constexpr auto tensor_sS() { return make_tensor(make_smem_ptr(P.begin()), SSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sP() { return make_tensor(make_smem_ptr(P.begin()), PSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sFmaxPartial() { return make_tensor(make_smem_ptr(FmaxPartial.begin()), WRSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sAcc1() { return make_tensor(make_smem_ptr(Acc1.begin()), Acc1SmemLayout{}); } + + // mailbox + CUTE_DEVICE constexpr auto tensor_sFmaxMailbox() { return make_tensor(make_smem_ptr(FmaxMailbox.begin()), MSMailboxSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sFsumMailbox() { return make_tensor(make_smem_ptr(FsumMailbox.begin()), MSMailboxSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sAcc2Mailbox() { return make_tensor(make_smem_ptr(Acc2Mailbox.begin()), Acc2MailboxSmemLayout{}); } + CUTE_DEVICE constexpr auto tensor_sSink() { return make_tensor(make_smem_ptr(Sink.begin()), SinksSmemLayout{}); } +}; + +// generic TMA copy device function that works on any tensor for 1 stage +template < + class GTensor, + class STensor, + class TmaAtom, + char Name, + bool DummyIter, // if true, don't issue any real tma load, only set barrier tx count to 0, keep the full/empty barrier pipeline running + bool Print, + int DMA_Stage> +CUTLASS_DEVICE void +TMA_copy( + GTensor gTensor, // ((TMA, NumTma_K), Tiles_K) + STensor sTensor, // ((TMA, NumTma_K), DMA_Stage or 1) + int k_tile, + int k_tile_base, // base index of the kblock that this CTA is going to work on + int& tma_mma_empty_barrier_phase_bit, + int tma_transaction_bytes, + TmaAtom const* tma_atom, + cute::uint64_t* tma_mma_full_barrier, + cute::uint64_t* tma_mma_empty_barrier +) { + + // a single kblock + // wait_barrier's input argument is the old phase bit + // it waits for the smem slot to be empty to start loading the next tile + wait_barrier(tma_mma_empty_barrier[k_tile % DMA_Stage], tma_mma_empty_barrier_phase_bit); + + if constexpr (Print) { + if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("[DMA_%c] barrier empty, kblock %d\n", Name, k_tile); + } + } + + if (elect_one_sync()) { + if constexpr (DummyIter) { + // set 0 transaction bytes for dummy iteration + set_barrier_transaction_bytes(tma_mma_full_barrier[k_tile % DMA_Stage], 0); + } + else { + // set the barrier transaction bytes to the number of bytes to load the tile, has to be done by a single thread + set_barrier_transaction_bytes(tma_mma_full_barrier[k_tile % DMA_Stage], tma_transaction_bytes); + // load Q tile into smem ((CTA_qHLocal, CTA_qL), CTA_dH) + copy(tma_atom->with(tma_mma_full_barrier[k_tile % DMA_Stage]), gTensor(_, k_tile_base + k_tile), sTensor(_,k_tile % DMA_Stage)); + } + } + + if ((k_tile % DMA_Stage) == (DMA_Stage - 1)) { + tma_mma_empty_barrier_phase_bit ^= 1; + } +} + +// generic MMA device function that works on any tensor for 1 stage +template < + class ATensor, + class BTensor, + class CTensor, + class TiledMMA, + char Name, + bool Print> +CUTLASS_DEVICE void +MMA_gemm( + ATensor tCrA, // ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + BTensor tCrB, // ((Mma_N, Mma_K), NumMma_N, NumMma_K, DMA_Stage) + CTensor tCtAcc, // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + TiledMMA tiled_mma, + int& stage_idx, + int& tma_mma_full_barrier_phase_bit, + UMMA::ScaleOut initial_accumulate, + cute::uint64_t* tma_mma_full_barrier, + cute::uint64_t* tma_mma_empty_barrier, + cute::uint64_t& mma_epilog_full_barrier +) { + int constexpr DMA_Stage = size<3>(tCrA); + + // setting it to 0 will clear the TMEM accumulator, setting it to one will reuse existing accumulator + tiled_mma.accumulate_ = initial_accumulate; + + // a single kblock + // wait for tma_mma_full_barrier to be full + wait_barrier(tma_mma_full_barrier[stage_idx], tma_mma_full_barrier_phase_bit); + + if constexpr (Print) { + if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("[MMA_%c] barrier full, kblock stage %d\n", Name, stage_idx); + } + } + + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // execute a Mma_M x Mma_N x Mma_K GEMM + // for K * Q and V * P, both the B matrix has a single stage, so we can just use stage_idx = 0 for B + cute::gemm(tiled_mma, tCrA(_,_,k_block, stage_idx), tCrB(_,_,k_block, 0), tCtAcc); + // after the first mma instruction, we need to start accumulating the result + // i.e. basically do C += A * B + tiled_mma.accumulate_ = UMMA::ScaleOut::One; + } + // notify the DMA warp that the all the MMA issued prior to this tcgen05.commit is done, the smem slot is now empty and can be reused + cutlass::arch::umma_arrive(&tma_mma_empty_barrier[stage_idx]); + + stage_idx++; + // for every DMA_Stage number of iterations, we flip the phase bit such that we reuse the full barrier for another round of mma + if (stage_idx == DMA_Stage) { + tma_mma_full_barrier_phase_bit ^= 1; + stage_idx = 0; + } + + // notify the epilog warp that MMA is done, the tmem slot is now full and can be consumed + cutlass::arch::umma_arrive(&mma_epilog_full_barrier); +} + +enum class ReduceOp { + Sum, + Max, + Min +}; + +// Helper for dependent static_assert in templates +template inline constexpr bool dependent_false = false; + +// https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__SINGLE.html +template +CUTLASS_DEVICE float +reduce_op(float a, float b) { + if constexpr (Op == ReduceOp::Sum) { + return a + b; + } + else if constexpr (Op == ReduceOp::Max) { + return fmaxf(a, b); + } + else if constexpr (Op == ReduceOp::Min) { + return fminf(a, b); + } + else { + // gcc 11.2 evaluate static_assert first before instantiating the template + // so we make the condition template dependent to ensure it is evaluated after if-else branch condition is evaluated + // gcc 13.2 fixes this issue + static_assert(dependent_false, "Invalid reduce op"); + return 0; + } +} + +// generic warp reduce function that works on a single reg value +// according to semantics section of the ptx doc https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-shfl-sync +// shfl.sync waits for all participating threads to reach this point, and then execute the shuffle instruction +// so we don't need to explicitly sync warp beforehand +// only supports full warp now, i.e. all 32 threads are active +template +CUTLASS_DEVICE float +warp_reduce(float val) { + // if it's doing float max or min, using the faster redux.sync instead of warp shuffle + if constexpr (Op == ReduceOp::Max) { + float result; + asm volatile("redux.sync.max.NaN.f32 %0, %1, 0xffffffff;\n" : "=f"(result) : "f"(val)); + return result; + } + else if constexpr (Op == ReduceOp::Min) { + float result; + asm volatile("redux.sync.min.NaN.f32 %0, %1, 0xffffffff;\n" : "=f"(result) : "f"(val)); + return result; + } + else { + float acc = val; + for (int i = 16; i >= 1; i /= 2) { + float new_val = __shfl_xor_sync(0xffffffff, acc, i, 32); + acc = reduce_op(acc, new_val); + } + return acc; + } +} + +// generic cta reduce function that works on a vector of reg values +// this is better than calling the same pattern on each reg value repeatedly because we coalesce the bar.sync before inter warp reduction into a single bar.sync +// the compiler can generate better pipelined code this way +// 1. do a warp reduce (redux.sync or warp shuffle) +// 2. first thread in each warp writes the warp reduce result to smem +// 3. sync across all participating warps/threads using named barrier to ensure smem write visibility +// 4. first thread in each warp reads all the warp reduce result from smem and do a reduction (separately for each warp) +// 5. first thread broadcast the cta reduce result to all threads in the warp +// this should work well when NumWarps is small, if it's big one should use another warp shuffle to do step 3 reduction +template < + ReduceOp Op, int NumWarps, + class RmemTensor, + class OTensor, + class SmemTensor> +CUTLASS_DEVICE void +cta_reduce( + RmemTensor val, // the value tensor to be reduced, we reduce each element across all participating threads + OTensor& out, // the output tensor, same shape as val, each threads hold the entire reduction result for each element + int warp_idx, // relative warp id that participate in the reduction, from 0 to NumWarps-1 + SmemTensor smem, // a tensor of shape (NumWarps, size(val)) to store the partial warp reduce result + cutlass::arch::NamedBarrier& warp_reduce_barrier +) { + + int lane_id = cutlass::canonical_lane_idx(); + for (int i = 0; i < size(val); i++) { + // step 1: do a warp reduce + float warp_reduce_result = warp_reduce(val(i)); + // step 2: first thread in each warp writes the warp reduce result to smem + if (lane_id == 0) { + smem(warp_idx, i) = warp_reduce_result; + } + } + // step 3: sync across all participating warps/threads using named barrier + // according to the ptx page, bar.sync has an implicit fence.cta to ensure smem write visibility + warp_reduce_barrier.arrive_and_wait(); + + // tried smem atomics and it's slower + for (int i = 0; i < size(val); i++) { + // step 4: first thread in each warp reads all the warp reduce result from smem and do a reduction (separately for each warp) + float acc = smem(0, i); + if (lane_id == 0) { + for (int j = 1; j < NumWarps; j++) { + acc = reduce_op(acc, smem(j, i)); + } + } + // step 5: first thread broadcast the cta reduce result to all threads in the warp + acc = __shfl_sync(0xffffffff, acc, 0, 32); + out(i) = acc; + } +} + +// this function acts on the transposed tensor of the above cta_reduce function +// what it effectively does is reducing values between all elements in val tensor and across consecutive NumThreads in the warp +template < + ReduceOp Op, + int NumThreads, // num consecutive threads to reduce + class RmemTensor> +CUTLASS_DEVICE float +cta_reduce_transposed( + RmemTensor val // per thread value tensor +) { + // 1. reduce all elements in val tensor (thread local reduction) + float acc = val(0); + for (int i = 1; i < size(val); i++) { + acc = reduce_op(acc, val(i)); + } + // 2. cross NumThreads thread reduce + for (int i = (NumThreads / 2); i >= 1; i /= 2) { + float new_val = __shfl_xor_sync(0xffffffff, acc, i, 32); + acc = reduce_op(acc, new_val); + } + return acc; +} + +// copied from SM100::TMEM::LOAD::copy_unpack cutlass/include/cute/atom/copy_traits_sm100.hpp +// what it does is given a tmem address, load the data into rmem tensor with the given tcgen05.ld copy op +template < + class CopyOp, + class TD, class DLayout> +CUTLASS_DEVICE void +tmem_load( + uint32_t tmem_addr, + Tensor& dst +) { + static_assert(is_rmem::value, "Expected RMEM dst."); + + using RegTypeDst = typename remove_extent::type; + Tensor rD = recast(dst); + + constexpr int RegNumDst = extent::value; + CUTE_STATIC_ASSERT_V(size(rD) == Int{}, + "The tcgen05.ld CopyOp's size does not match the destination tensor size."); + + detail::explode(CopyOp::copy, + &tmem_addr, seq<0>{}, + rD, make_seq{}); +} + +// copied from SM100::TMEM::STORE::copy_unpack cutlass/include/cute/atom/copy_traits_sm100.hpp +// what it does is given a tmem address, store the data in rmem tensor to the tmem address with the given tcgen05.st copy op +template < + class CopyOp, + class TS, class SLayout> +CUTLASS_DEVICE void +tmem_store( + Tensor& src, + uint32_t tmem_addr +) { + static_assert(is_rmem::value, "Expected RMEM src."); + + using RegTypeSrc = typename remove_extent::type; + Tensor rS = recast(src); + + constexpr int RegNumSrc = extent::value; + CUTE_STATIC_ASSERT_V(size(rS) == Int{}, + "The tcgen05.st CopyOp's size does not match the source tensor size."); + + detail::explode(CopyOp::copy, + rS, make_seq{}, + &tmem_addr, seq<0>{}); +} + +// issue cp.async +CUTLASS_DEVICE void +cp_async( + int* gmem_addr, + int* smem_addr +) { + uint32_t smem_int_ptr = cute::cast_smem_ptr_to_uint(smem_addr); + asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n" + :: "r"(smem_int_ptr), + "l"(gmem_addr), + "n"(sizeof(int))); +} + +// mapping between thread id (T) -> dH (row index of Acc2) +template +CUTLASS_DEVICE auto +make_bmm2_tmem_load_t_dH_layout() { + if constexpr (CTA_dH == 64) { + return make_layout(make_shape(Int<16>{}, Int<2>{}, Int<4>{}), make_stride(Int<1>{}, Int<0>{}, Int<16>{})); + } + else if constexpr (CTA_dH == 128) { + return make_layout(Int<128>{}); + } + else { + static_assert(dependent_false, "CTA_dH must be 64 or 128"); + return make_layout(Int<128>{}); + } +} + +// TMA load Q, guarded by griddepcontrol.wait +template < + class SharedStorage, + class WorkTileInfo, + class QTensor, + class TmaAtomQ, + class TiledBMM1, + int CTA_qHLocal, int CTA_qL, int CTA_dH> +CUTLASS_DEVICE void +DMA_Q_warp( + SharedStorage& shared_storage, + WorkTileInfo work_tile_info, + QTensor mQ, + // when passing tma descriptor as function argument, it has to be pass by pointer/reference, if pass by value, it will live on local memory (i.e. the stack) + // and the tma unit cannot access the local memory, (even if it can, the local memory is strided by thread id, the content for each thread is strided) + TmaAtomQ const* tma_atom_Q, + TiledBMM1 tiled_bmm1) { + + if (!work_tile_info.is_valid()) { + return; + } + + // tCsQ means the mma partitioned smem tensor sQ + Tensor tCsQ = shared_storage.tensor_sQ(); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + // now let's tile the gmem tensor first into the tile this CTA needs (i.e. CTA-level tiling), which is ((CTA_qHLocal, CTA_qL), CTA_dH, Num_CTA_dH) + // mQ has shape ((qHLocal, qL), dH, kvH, BS), after local_tile (without indexing), it becomes ((CTA_qHLocal, CTA_qL), CTA_dH, (Num_CTA_qHLocal, Num_CTA_qL), Num_CTA_dH, kvH, BS) + // there is only one tile on the dH dimension + Tensor gQ = local_tile(mQ, make_shape(make_shape(Int{}, Int{}), Int{}), + make_coord(make_coord(work_tile_info.qHLocal_idx, work_tile_info.qL_idx), _, work_tile_info.kvH_idx, work_tile_info.BS_idx)); // ((CTA_qHLocal, CTA_qL), CTA_dH, Num_CTA_dH) + + // now we want to partition (give the same tensor a new view) the gmem tensor into the mma shape, to match the smem layout + ThrMMA cta_bmm1 = tiled_bmm1.get_slice(0); // 1 sm mma only has 1 thread in tiled_bmm1 + // tCgQ means the mma partitioned gmem tensor gQ + Tensor tCgQ = cta_bmm1.partition_B(gQ); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, Tiles_K) + + // now we want to partition (give the same tensor a new view) the gmem tensor into the tma shape in preparation for tma copy + auto [tBgQ, tBsQ] = tma_partition(*tma_atom_Q, + Int<0>{}, // cta_coord: 1x1 cluster + Layout<_1>{}, // cta_layout: CTA coord -> logical multicast id, no multicast, just identity layout + group_modes<0,3>(tCsQ), group_modes<0,3>(tCgQ)); + // tBgQ: ((TMA, NumTma_N, NumTma_K), Tiles_K) get coalesced to ((TMA, NumTma_K), Tiles_K) since NumTma_N = 1 + // tBsQ: ((TMA, NumTma_N, NumTma_K), 1) get coalesced to ((TMA, NumTma_K), 1) since NumTma_N = 1 + // the shape of the TMA box is (CTA_qHLocal*CTA_qL, CTA_dH) + + /*if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("cta_bmm1:\t"); print(cta_bmm1); print("\n"); + printf("tCsQ:\t"); print(tCsQ); print("\n"); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + printf("gQ:\t"); print(gQ); print("\n"); // ((CTA_qHLocal, CTA_qL), CTA_dH, Num_CTA_dH) + printf("tCgQ:\t"); print(tCgQ); print("\n"); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, Tiles_K) + printf("tBgQ:\t"); print(tBgQ); print("\n"); // ((TMA, NumTma_K), Tiles_K) + printf("tBsQ:\t"); print(tBsQ); print("\n"); // ((TMA, NumTma_K), 1) + }*/ + + // only do griddepcontrol.wait on Q loading, Q is dependent on previous kernel + cutlass::arch::wait_on_dependent_grids(); + + int tma_transaction_bytes = sizeof(make_tensor_like(tBsQ(_,0))); + int k_tile_count = work_tile_info.kvL_idx_end - work_tile_info.kvL_idx_start; + + Tensor tCsK = shared_storage.tensor_sK(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + int constexpr DMA_Stage = size<3>(tCsK); + + /*if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("tBgQ:\t"); print(tBgQ); print("\n"); // ((TMA, NumTma_K), Tiles_K) + printf("tBsQ:\t"); print(tBsQ); print("\n"); // ((TMA, NumTma_K), 1) + printf("TmaBytes: %d\n", tma_transaction_bytes); + printf("k_tile_count: %d\n", k_tile_count); + printf("DMA_Stage: %d\n", DMA_Stage); + }*/ + + // initial phase bit for tma_mma_empty_barrier is 0, and it denotes smem slot is empty + // we wait on the "old" phase bit of 1, when it is "flipped" to 0 (the initial phase bit), the smem slot is empty + // so tma_mma_empty_barrier_phase_bit denotes the phase bit before the flip we are waiting on + // example for DMA_Stage = 2, i.e. 2 smem slots, 2 empty barriers a and b + // kblock 0: old phase bit of barrier a is 1, wait for barrier a's phase bit to change to 0 (the initial phase bit), 0 denotes slot a is empty, 1 denotes slot a is used by mma + // kblock 1: old phase bit of barrier b is 1, wait for barrier b's phase bit to change to 0 (the initial phase bit), 0 denotes slot b is empty, 1 denotes slot b is used by mma + // kblock 2: old phase bit of barrier a is 0, wait for barrier a's phase bit to change to 1 (flipped once), 1 denotes slot a is empty, 0 denotes slot a is used by mma + // kblock 3: old phase bit of barrier b is 0, wait for barrier b's phase bit to change to 1 (flipped once), 1 denotes slot b is empty, 0 denotes slot b is used by mma + // kblock 4: old phase bit of barrier a is 1, wait for barrier a's phase bit to change to 0 (flipped twice), 0 denotes slot a is empty, 1 denotes slot a is used by mma + // kblock 5: old phase bit of barrier b is 1, wait for barrier b's phase bit to change to 0 (flipped twice), 0 denotes slot b is empty, 1 denotes slot b is used by mma + // ... + int tma_mma_empty_barrier_phase_bit = 1; + + // Q0, _, _, ... + int k_tile = 0; + // fetch Q0, there is only 1 Q tile + TMA_copy(tBgQ, tBsQ, k_tile, 0, tma_mma_empty_barrier_phase_bit, tma_transaction_bytes, tma_atom_Q, shared_storage.tma_bmm1_full_barrier, shared_storage.tma_bmm1_empty_barrier); + k_tile++; + + // so the rest of the K tiles all share the same Q tile, so we just do dummy iterations + for (; k_tile < k_tile_count; k_tile++) { + TMA_copy(tBgQ, tBsQ, k_tile, 0, tma_mma_empty_barrier_phase_bit, tma_transaction_bytes, tma_atom_Q, shared_storage.tma_bmm1_full_barrier, shared_storage.tma_bmm1_empty_barrier); + } + + /*wait_barrier(shared_storage.tma_bmm1_full_barrier[0], 0); + if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("tCsQ:\t"); print_tensor(tCsQ(_,0,_,0)); print("\n"); + }*/ +} + +// TMA load K and then V, not guarded by griddepcontrol.wait +template < + class SharedStorage, + class WorkTileInfo, + class KTensor, + class VTensor, + class TmaAtomK, + class TmaAtomV, + class TiledBMM1, + class TiledBMM2, + int CTA_kvL, int CTA_dH> +CUTLASS_DEVICE void +DMA_KV_warp( + SharedStorage& shared_storage, + WorkTileInfo work_tile_info, + KTensor mK, + VTensor mV, + // when passing tma descriptor as function argument, it has to be pass by pointer/reference, if pass by value, it will live on local memory (i.e. the stack) + // and the tma unit cannot access the local memory, (even if it can, the local memory is strided by thread id, the content for each thread is strided) + TmaAtomK const* tma_atom_K, + TmaAtomV const* tma_atom_V, + TiledBMM1 tiled_bmm1, + TiledBMM2 tiled_bmm2) { + + if (!work_tile_info.is_valid()) { + return; + } + + // setup code for K tensor + // tCsK means the mma partitioned smem tensor sK + Tensor tCsK = shared_storage.tensor_sK(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + // now let's tile the gmem tensor first into the tile this CTA needs (i.e. CTA-level tiling), which is (CTA_kvL, CTA_dH, Num_CTA_kvL) + // mK has shape (kvL, dH, kvH, BS), after local_tile (without indexing), it becomes (CTA_kvL, CTA_dH, Num_CTA_kvL, Num_CTA_dH, kvH, BS) + // and this CTA works on gK(_,_,kvL_idx_start:kvL_idx_end) + Tensor gK = local_tile(mK, make_shape(Int{}, Int{}), + make_coord(_, 0, work_tile_info.kvH_idx, work_tile_info.BS_idx)); // (CTA_kvL, CTA_dH, Num_CTA_kvL) + + // now we want to partition (give the same tensor a new view) the gmem tensor into the mma shape, to match the smem layout + ThrMMA cta_bmm1 = tiled_bmm1.get_slice(0); // 1 sm mma only has 1 thread in tiled_bmm1 + // tCgK means the mma partitioned gmem tensor gK + Tensor tCgK = cta_bmm1.partition_A(gK); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, Tiles_M) + + // now we want to partition (give the same tensor a new view) the gmem tensor into the tma shape in preparation for tma copy + auto [tAgK, tAsK] = tma_partition(*tma_atom_K, + Int<0>{}, // cta_coord: 1x1 cluster + Layout<_1>{}, // cta_layout: CTA coord -> logical multicast id, no multicast, just identity layout + group_modes<0,3>(tCsK), group_modes<0,3>(tCgK)); + // tAgK: ((TMA, NumTma_M, NumTma_K), Tiles_M) get coalesced to ((TMA, NumTma_K), Tiles_M) since NumTma_M = 1 + // tAsK: ((TMA, NumTma_M, NumTma_K), BMM1_DMA_Stage) get coalesced to ((TMA, NumTma_K), BMM1_DMA_Stage) since NumTma_M = 1 + // the shape of the TMA box is (CTA_kvL, CTA_dH) + + /*if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("cta_bmm1:\t"); print(cta_bmm1); print("\n"); + printf("tCsK:\t"); print(tCsK); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + printf("gK:\t"); print(gK); print("\n"); // (CTA_kvL, CTA_dH, Num_CTA_kvL) + printf("tCgK:\t"); print(tCgK); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, Tiles_M) + printf("tAgK:\t"); print(tAgK); print("\n"); // ((TMA, NumTma_K), Tiles_M) + printf("tAsK:\t"); print(tAsK); print("\n"); // ((TMA, NumTma_K), BMM1_DMA_Stage) + }*/ + + // setup code for V tensor + // tCsV means the mma partitioned smem tensor sV + Tensor tCsV = shared_storage.tensor_sV(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + // now let's tile the gmem tensor first into the tile this CTA needs (i.e. CTA-level tiling), which is (CTA_dH, CTA_kvL) + // mV has shape (dH, kvL, kvH, BS), after local_tile (without indexing), it becomes (CTA_dH, CTA_kvL, Num_CTA_dH, Num_CTA_kvL, kvH, BS) + // and this CTA works on gV(_,_,kvL_idx_start:kvL_idx_end) + Tensor gV = local_tile(mV, make_shape(Int{}, Int{}), + make_coord(work_tile_info.dH_idx, _, work_tile_info.kvH_idx, work_tile_info.BS_idx)); // (CTA_dH, CTA_kvL, Num_CTA_kvL) + + // now we want to partition (give the same tensor a new view) the gmem tensor into the mma shape, to match the smem layout + ThrMMA cta_bmm2 = tiled_bmm2.get_slice(0); // 1 sm mma only has 1 thread in tiled_bmm2 + // tCgV means the mma partitioned gmem tensor gV + Tensor tCgV = cta_bmm2.partition_A(gV); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, Tiles_K) + + // now we want to partition (give the same tensor a new view) the gmem tensor into the tma shape in preparation for tma copy + auto [tAgV, tAsV] = tma_partition(*tma_atom_V, + Int<0>{}, // cta_coord: 1x1 cluster + Layout<_1>{}, // cta_layout: CTA coord -> logical multicast id, no multicast, just identity layout + group_modes<0,3>(tCsV), group_modes<0,3>(tCgV)); + // tAgV: ((TMA, NumTma_M, NumTma_K), Tiles_K) get coalesced to ((TMA, NumTma_K), Tiles_K) since NumTma_M = 1 + // tAsV: ((TMA, NumTma_M, NumTma_K), BMM2_DMA_Stage) get coalesced to ((TMA, NumTma_K), BMM2_DMA_Stage) since NumTma_M = 1 + // the shape of the TMA box is (CTA_dH, CTA_kvL) + + /*if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("cta_bmm2:\t"); print(cta_bmm2); print("\n"); + printf("tCsV:\t"); print(tCsV); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + printf("gV:\t"); print(gV); print("\n"); // (CTA_dH, CTA_kvL, Num_CTA_kvL) + printf("tCgV:\t"); print(tCgV); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, Tiles_K) + printf("tAgV:\t"); print(tAgV); print("\n"); // ((TMA, NumTma_K), Tiles_K) + printf("tAsV:\t"); print(tAsV); print("\n"); // ((TMA, NumTma_K), BMM2_DMA_Stage) + }*/ + + int tma_K_transaction_bytes = sizeof(make_tensor_like(tAsK(_,0))); + int tma_V_transaction_bytes = sizeof(make_tensor_like(tAsV(_,0))); + int k_tile_count = work_tile_info.kvL_idx_end - work_tile_info.kvL_idx_start; + int constexpr BMM1_DMA_Stage = size<3>(tCsK); + int constexpr BMM2_DMA_Stage = size<3>(tCsV); + + /*if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("TmaKBytes: %d\n", tma_K_transaction_bytes); + printf("TmaVBytes: %d\n", tma_V_transaction_bytes); + printf("k_tile_count: %d\n", k_tile_count); + printf("BMM1_DMA_Stage: %d\n", BMM1_DMA_Stage); + printf("BMM2_DMA_Stage: %d\n", BMM2_DMA_Stage); + }*/ + + // details of how the phase bit works is in DMA_Q_warp + int tma_bmm1_empty_barrier_phase_bit = 1; + int tma_bmm2_empty_barrier_phase_bit = 1; + + int bmm1_k_tile = 0; + int bmm2_k_tile = 0; + + bool constexpr Print = false; + // iterate over kblock + // SOL order is K0, (K1, V0), (K2, V1), (K3, V2), ... + + // sw pipeline prolog + TMA_copy(tAgK, tAsK, bmm1_k_tile, work_tile_info.kvL_idx_start, tma_bmm1_empty_barrier_phase_bit, tma_K_transaction_bytes, tma_atom_K, shared_storage.tma_bmm1_full_barrier, shared_storage.tma_bmm1_empty_barrier); + bmm1_k_tile++; + + for (; bmm1_k_tile < k_tile_count; bmm1_k_tile++, bmm2_k_tile++) { + + TMA_copy(tAgK, tAsK, bmm1_k_tile, work_tile_info.kvL_idx_start, tma_bmm1_empty_barrier_phase_bit, tma_K_transaction_bytes, tma_atom_K, shared_storage.tma_bmm1_full_barrier, shared_storage.tma_bmm1_empty_barrier); + + TMA_copy(tAgV, tAsV, bmm2_k_tile, work_tile_info.kvL_idx_start, tma_bmm2_empty_barrier_phase_bit, tma_V_transaction_bytes, tma_atom_V, shared_storage.tmasoftmax_bmm2_full_barrier, shared_storage.tma_bmm2_empty_barrier); + } + + // sw pipeline epilog + TMA_copy(tAgV, tAsV, bmm2_k_tile, work_tile_info.kvL_idx_start, tma_bmm2_empty_barrier_phase_bit, tma_V_transaction_bytes, tma_atom_V, shared_storage.tmasoftmax_bmm2_full_barrier, shared_storage.tma_bmm2_empty_barrier); + + cutlass::arch::launch_dependent_grids(); +} + +// MMA warp that includes BMM1 and BMM2 +template < + class SharedStorage, + class WorkTileInfo, + class OTensor, + class TiledBMM1, + class TiledBMM2, + int CTA_qHLocal, int CTA_qL, int CTA_kvL, int CTA_dH> +CUTLASS_DEVICE void +MMA_warp( + SharedStorage& shared_storage, + WorkTileInfo work_tile_info, + OTensor mO, + TiledBMM1 tiled_bmm1, + TiledBMM2 tiled_bmm2, + cutlass::arch::NamedBarrier& tmem_allocation_barrier +) { + if (!work_tile_info.is_valid()) { + // we don't allocate tmem for invalid tiles but we still need to relinquish the allocation lock + using TmemAllocator = cute::TMEM::Allocator1Sm; + TmemAllocator tmem_allocator{}; + tmem_allocator.release_allocation_lock(); + return; + } + + // compute BMM1, S = K x Q + // smem tensor prepared by DMA_K and DMA_Q + Tensor tCsK = shared_storage.tensor_sK(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + Tensor tCsQ = shared_storage.tensor_sQ(); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + Tensor sS = shared_storage.tensor_sS(); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + + // BMM1 Fragment Allocation + ThrMMA cta_bmm1 = tiled_bmm1.get_slice(0); // 1 sm mma only has 1 thread in tiled_mma + Tensor tCrK = cta_bmm1.make_fragment_A(tCsK); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + Tensor tCrQ = cta_bmm1.make_fragment_B(tCsQ); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + // then partition it to mma shape + Tensor tCsS = cta_bmm1.partition_C(sS(_,_,0)); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + Tensor tCtAcc1 = cta_bmm1.make_fragment_C(tCsS); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + + // compute BMM2, O = V x P + // smem tensor prepared by DMA_V and softmax + Tensor tCsV = shared_storage.tensor_sV(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + Tensor sP = shared_storage.tensor_sP(); // ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) + // cluster level tiling for mO, gO is what will be produced by the entire cluster + Tensor gO = local_tile(mO, make_shape(Int{}, make_shape(Int{}, Int{})), + make_coord(work_tile_info.dH_idx, make_coord(work_tile_info.qHLocal_idx, work_tile_info.qL_idx), work_tile_info.kvH_idx, work_tile_info.BS_idx)); // (CTA_dH, (CTA_qHLocal, CTA_qL)) + + // BMM2 Fragment Allocation + ThrMMA cta_bmm2 = tiled_bmm2.get_slice(0); // 1 sm mma only has 1 thread in tiled_mma + Tensor tCrV = cta_bmm2.make_fragment_A(tCsV); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + Tensor tCsP = cta_bmm2.partition_B(sP); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + Tensor tCrP = cta_bmm2.make_fragment_B(tCsP); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + // then partition it to mma shape + Tensor tCgO = cta_bmm2.partition_C(gO); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + Tensor tCtAcc2 = cta_bmm2.make_fragment_C(tCgO); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + + // TMEM Allocation + // tmem has 128 lane, 512 column, 1 word/column, each word is 4B, 256KB in total + using TmemAllocator = cute::TMEM::Allocator1Sm; + TmemAllocator tmem_allocator{}; + + int constexpr Acc1_col = CTA_qHLocal * CTA_qL; // num tmem column for tCtAcc1 + int constexpr Acc1_col_max = cute::max(Acc1_col, TmemAllocator::ColumnsPerAllocationSlice); // min allocation granularity is 32 columns + static_assert((Acc1_col_max & (Acc1_col_max - 1)) == 0, "Acc1 column size is not a power of 2"); + + int constexpr Acc2_col = CTA_qHLocal * CTA_qL; // num tmem column for tCtAcc2 + int constexpr Acc2_col_max = cute::max(Acc2_col, TmemAllocator::ColumnsPerAllocationSlice); // min allocation granularity is 32 columns + static_assert((Acc2_col_max & (Acc2_col_max - 1)) == 0, "Acc2 column size is not a power of 2"); + + // allow half of tmem to be used for overlapping with next kernel + static_assert((Acc1_col_max + Acc2_col_max) < TmemAllocator::Sm100TmemCapacityColumns / 2, "Accumulator is too large to fit in half of tmem"); + + // do two separate allocations for bmm1 and bmm2 accumulators, we waste some capacity and cycles + // if needed, in the future we can merge the two allocations into one and do some offsetting + tmem_allocator.allocate(Acc1_col_max, &shared_storage.bmm1_tmem_base_ptr); + tmem_allocator.allocate(Acc2_col_max, &shared_storage.bmm2_tmem_base_ptr); + // notify epilog warp that tmem allocation is complete + tmem_allocation_barrier.arrive(); + + // relinquish early so that prefetch cta can be launched + tmem_allocator.release_allocation_lock(); + + // update the tmem base ptr of the accumulator tensor + tCtAcc1.data() = shared_storage.bmm1_tmem_base_ptr; + tCtAcc2.data() = shared_storage.bmm2_tmem_base_ptr; + + /*if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("tCsK:\t"); print(tCsK); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + printf("tCsQ:\t"); print(tCsQ); print("\n"); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + printf("sS:\t"); print(sS); print("\n"); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + printf("tCsS:\t"); print(tCsS); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("Acc1_col_max:\t%d\n", Acc1_col_max); + printf("tCtAcc1:\t"); print(tCtAcc1); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("tCrK:\t"); print(tCrK); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + printf("tCrQ:\t"); print(tCrQ); print("\n"); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + + printf("tCsV:\t"); print(tCsV); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + printf("sP:\t"); print(sP); print("\n"); // ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) + printf("tCsP:\t"); print(tCsP); print("\n"); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + printf("tCrV:\t"); print(tCrV); print("\n"); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + printf("tCrP:\t"); print(tCrP); print("\n"); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + printf("gO:\t"); print(gO); print("\n"); // (CTA_dH, (CTA_qHLocal, CTA_qL)) + printf("tCgO:\t"); print(tCgO); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("Acc2_col_max:\t%d\n", Acc2_col_max); + printf("tCtAcc2:\t"); print(tCtAcc2); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + }*/ + + bool constexpr Print = false; + int k_tile_count = work_tile_info.kvL_idx_end - work_tile_info.kvL_idx_start; + // iterate over kblock + // SOL order is: K0Q0, (K1Q0, V0P0), (K2Q0, V1P1), (K3Q0, V2P2), ... + + // details of how the phase bit works is in DMA_Q_warp + int tma_bmm1_full_barrier_phase_bit = 0; + int tma_bmm2_full_barrier_phase_bit = 0; + + int bmm1_stage_idx = 0; + int bmm2_stage_idx = 0; + + int k_tile = 0; + + // prolog bmm1 + // bmm1 always overwrites the accumulator, so we set it to UMMA::ScaleOut::Zero + MMA_gemm(tCrK, tCrQ, tCtAcc1, tiled_bmm1, bmm1_stage_idx, tma_bmm1_full_barrier_phase_bit, UMMA::ScaleOut::Zero, shared_storage.tma_bmm1_full_barrier, shared_storage.tma_bmm1_empty_barrier, shared_storage.bmm1_softmax_full_barrier); + + for (; k_tile < (k_tile_count - 1); k_tile++) { + // bmm1 of next iter + MMA_gemm(tCrK, tCrQ, tCtAcc1, tiled_bmm1, bmm1_stage_idx, tma_bmm1_full_barrier_phase_bit, UMMA::ScaleOut::Zero, shared_storage.tma_bmm1_full_barrier, shared_storage.tma_bmm1_empty_barrier, shared_storage.bmm1_softmax_full_barrier); + + // bmm2 of current iter + // for bmm2, only clear the accumulator for the first iter, for subsequent iters, we accumulate the VP result on top of corrected Acc2 + UMMA::ScaleOut bmm2_accumulate = (k_tile == 0) ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + MMA_gemm(tCrV, tCrP, tCtAcc2, tiled_bmm2, bmm2_stage_idx, tma_bmm2_full_barrier_phase_bit, bmm2_accumulate, shared_storage.tmasoftmax_bmm2_full_barrier, shared_storage.tma_bmm2_empty_barrier, shared_storage.bmm2_epilog_full_barrier); + } + + // epilog bmm2 + UMMA::ScaleOut bmm2_accumulate = (k_tile == 0) ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + MMA_gemm(tCrV, tCrP, tCtAcc2, tiled_bmm2, bmm2_stage_idx, tma_bmm2_full_barrier_phase_bit, bmm2_accumulate, shared_storage.tmasoftmax_bmm2_full_barrier, shared_storage.tma_bmm2_empty_barrier, shared_storage.bmm2_epilog_full_barrier); + + // wait for tmem deallocation signal from epilog warp + tmem_allocation_barrier.arrive_and_wait(); + + // deallocate TMEM + tmem_allocator.free(shared_storage.bmm1_tmem_base_ptr, Acc1_col_max); + tmem_allocator.free(shared_storage.bmm2_tmem_base_ptr, Acc2_col_max); +} + +// first do softmax on S, then do flash decode reduction +template < + class SharedStorage, + class WorkTileInfo, + class KTensor, + class OTensor, + class SinkTensor, + class TiledBMM1, + class TiledBMM2, + int CTA_qHLocal, int CTA_qL, int CTA_kvL, int CTA_dH, + int NumReductionCTA, + bool NoSink> +CUTLASS_DEVICE void +EPILOG_warp( + SharedStorage& shared_storage, + WorkTileInfo work_tile_info, + KTensor mK, + OTensor mO, + SinkTensor mSink, + int seq_len, + TiledBMM1 tiled_bmm1, + TiledBMM2 tiled_bmm2, + float softmax_scale_log2, + int sliding_window_size, + cutlass::arch::NamedBarrier& tmem_allocation_barrier, + cutlass::arch::NamedBarrier& epilog_barrier, + int NumSplits, + int tid, // tid local to epilog warp + int warp_idx, // warp id local to epilog warp group + int rank +) { + // bmm1 + Tensor sS = shared_storage.tensor_sS(); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + // BMM1 Fragment Allocation + ThrMMA cta_bmm1 = tiled_bmm1.get_slice(0); // 1 sm mma only has 1 thread in tiled_bmm1 + // then partition it to mma shape + // we only use 1 S buffer, it gives sufficient overlap between MMA and softmax + Tensor tCsS = cta_bmm1.partition_C(sS(_,_,0)); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + // create the tmem tensor view + Tensor tCtAcc1 = cta_bmm1.make_fragment_C(tCsS); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + + // bmm2 + // cluster level tiling for mO, gO is what will be produced by the entire cluster + Tensor gO = local_tile(mO, make_shape(Int{}, make_shape(Int{}, Int{})), + make_coord(work_tile_info.dH_idx, make_coord(work_tile_info.qHLocal_idx, work_tile_info.qL_idx), work_tile_info.kvH_idx, work_tile_info.BS_idx)); // (CTA_dH, (CTA_qHLocal, CTA_qL)) + // BMM2 Fragment Allocation + ThrMMA cta_bmm2 = tiled_bmm2.get_slice(0); // 1 sm mma only has 1 thread in tiled_bmm2 + // then partition it to mma shape + Tensor tCgO = cta_bmm2.partition_C(gO); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + // create the tmem tensor view + Tensor tCtAcc2 = cta_bmm2.make_fragment_C(tCgO); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + + // wait for tmem allocation in mma warp to complete, only do the wait for valid tiles + if (work_tile_info.is_valid()) { + tmem_allocation_barrier.arrive_and_wait(); + } + + // update tmem base ptr of the accumulator tensor + tCtAcc1.data() = shared_storage.bmm1_tmem_base_ptr; + tCtAcc2.data() = shared_storage.bmm2_tmem_base_ptr; + + // BMM1 tcgen05.ld + using TypeAcc = typename decltype(tCtAcc1)::value_type; + constexpr int acc1_col_bits = size<1>(sS) * sizeof_bits_v; + // use tcgen05.ld.32dp32bit simplifies softmax implementation, every thread just hold a row of the accumulator tensor + // softmax is applied to each column, so we do both the sum and max across all 128 threads using warp shuffle + auto bmm1_tmem_load = TMEM::op_repeater(); + TiledCopy tiled_t2r_copy_bmm1 = make_tmem_copy(bmm1_tmem_load, tCtAcc1); + // epilog tid is from 128 to 255, need to offset by -128 when getting the per thread slice + ThrCopy thr_t2r_copy_bmm1 = tiled_t2r_copy_bmm1.get_slice(tid); + + // tD means the partitioning pattern of tcgen05.ld + // now we get the per thread slice of the Acc1 for tcgen05.ld src and dst + Tensor tDtAcc1 = thr_t2r_copy_bmm1.partition_S(tCtAcc1); // (CpyS, NumCpy_M, NumCpy_N), per subpartition view of the Acc1/S tensor + Tensor tDsS = thr_t2r_copy_bmm1.partition_D(tCsS); // (CpyD, NumCpy_M, NumCpy_N), per thread slice of the Acc1/S tensor in smem, i.e. a row of the Acc1 tensor + // allocate per thread rmem space for the Acc1, the shape is the same as the post partition shape of the output tensor + Tensor tDrS = make_tensor(shape(tDsS)); // (CpyD, NumCpy_M, NumCpy_N), per thread slice of the Acc1/S tensor in rmem, i.e. a row of the Acc1 tensor + + // restructure Acc1 in rmem for faster fsum reduction + // because tcgen05.ld.32dp32bit put each row of the Acc1 tensor into 1 thread, for softmax we need cross thread reduction (reduce each column) for fsum/fmax, fmax is fine because we can use redux.sync + // so we restructure Acc1 in rmem such that each thread holds a partial column of the Acc1 tensor, so we do thread local reduction as much as possible, then do warp/cta/cluster reduction + Tensor sAcc1 = shared_storage.tensor_sAcc1(); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + Tensor tCsAcc1 = cta_bmm1.partition_C(sAcc1(_,_,0)); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + Tensor tDsAcc1 = thr_t2r_copy_bmm1.partition_D(tCsAcc1); // (CpyD, NumCpy_M, NumCpy_N), per thread slice of the output tensor + // zipped_divide partition it into (CTA_qHLocal * CTA_qL, (CTA_kvL / CTA_qHLocal / CTA_qL, (CTA_qHLocal, CTA_qL))) + // then the per thread indexing gets the shape to (CTA_qHLocal * CTA_qL) + // tFsum means the thread partitioning pattern for calculating fsum, i.e. each thread holds a partial column of the Acc1 tensor + Tensor tFsumsAcc1 = zipped_divide(sAcc1(_,_,0), Int{})(_, tid); // (CTA_qHLocal * CTA_qL) + Tensor tFsumrAcc1 = make_tensor(shape(tFsumsAcc1)); // (CTA_qHLocal * CTA_qL) + // so what we do is tmem(tDtAcc1)->rmem(tDrS)->smem(tDsAcc1/tFsumsAcc1)->rmem(tFsumrAcc1) + static_assert(CTA_kvL % (CTA_qHLocal * CTA_qL) == 0, "CTA_kvL must be divisible by CTA_qHLocal * CTA_qL"); // such that each column is held by integer number of threads + static_assert((CTA_kvL / (CTA_qHLocal * CTA_qL)) <= 32, "CTA_kvL / (CTA_qHLocal * CTA_qL) must be less than or equal to 32"); // such that fsum reduction won't cross warp boundary + + // predicate tensor for predicating S/P based on kvL + Tensor coordK = make_identity_tensor(shape(mK)); // (kvL, dH, kvH, BS) -> (kvL, dH, kvH, BS) + Tensor cK = local_tile(coordK, make_shape(Int{}, Int{}), + make_coord(_, work_tile_info.dH_idx, work_tile_info.kvH_idx, work_tile_info.BS_idx)); // (CTA_kvL, CTA_dH, Num_CTA_kvL) -> (kvL, dH, kvH, BS) + Tensor cur_kvL = cK(tid, 0, _); // (Num_CTA_kvL) -> (kvL, dH, kvH, BS) + + // BMM2 tcgen05.ld + // now we get the per thread slice of the Acc2 for tcgen05.ld src and dst + constexpr int acc2_col_bits = size<1>(gO) * sizeof_bits_v; + // use tcgen05.ld.32dp32bit simplifies correction implementation, every thread just hold a row of the accumulator tensor + // but cute will be unhappy since we use MMA_M=64 mma but load tmem with 32dp32bit, so we need to bypass cute copy completely by directly using ptx + auto bmm2_tmem_load = TMEM::op_repeater(); + // similarly, use tcgen05.st.32dp32bit to store the corrected Acc2 back to tmem + auto bmm2_tmem_store = TMEM::op_repeater(); + // allocate per thread rmem space for the reshaped Acc2 tensor, it just hold a row of the Acc2 tensor + // tAcc2 means it's the partitioning pattern where each thread holds a row of the Acc2 tensor + // for dH=64, after tcgen05.ld, thread 0-15, 32-47, 64-79, 96-111 will hold the data of Acc2(0:16,:), Acc2(32:48,:), Acc2(64:80,:), Acc2(96:112,:) + // for dH=128, after tcgen05.ld, thread 0-31, 32-63, 64-95, 96-127 will hold the data of Acc2(0:32,:), Acc2(32:64,:), Acc2(64:96,:), Acc2(96:128,:) + Tensor tAcc2rAcc2 = make_tensor(shape<1>(gO)); // (CTA_qHLocal, CTA_qL), per thread slice of the Acc2 tensor in rmem + uint32_t acc2_tmem_addr = raw_pointer_cast(tCtAcc2.data()); + // mapping between thread id (T) -> dH (row index of Acc2), different for dH=64 and dH=128 + Layout bmm2_tmem_load_t_dH_layout = make_bmm2_tmem_load_t_dH_layout(); + + // now we define the warp reduce partial results tensor for fmax + Tensor sFmaxPartial = shared_storage.tensor_sFmaxPartial(); // (NumEpilogWarps, (CTA_qHLocal, CTA_qL)) + // each element in tDrS (CTA_qHLocal, CTA_qL) is from a different column (q token) of Acc1 tensor, softmax takes column sum/max and scale each element + // each column would need a separate fmax, so the meta data max tensor has the same shape as tDrS, and the max value is the same for the same column for all threads + // this is the CTA wide fmax tensor for all kv tiles + Tensor tDrFmax = make_tensor(shape(tDsS)); // (CpyD, NumCpy_M, NumCpy_N) or (CTA_qHLocal * CTA_qL), per thread slice of the meta data max tensor + + // we do things a bit differently for fsum, each warp holds its portion of the fsum tensor, there are CTA_qHLocal * CTA_qL fsum in total for this cta + // and we evenly distribute the q tokens between all epilog warps, so each warp computes CTA_qHLocal * CTA_qL / NumEpilogWarps q tokens + int constexpr NumEpilogWarps = size<0>(sFmaxPartial); + static_assert(CTA_dH <= (NumEpilogWarps * 32), "have enough threads to hold bmm2 tcgen05.ld output"); + int constexpr NumQPerWarp = CTA_qHLocal * CTA_qL / NumEpilogWarps; + // alpha = exp2f(m_i_old - m_i) + Tensor rAlpha = make_tensor(shape(tDrS)); // (CpyD, NumCpy_M, NumCpy_N) or (CTA_qHLocal * CTA_qL), per thread slice + // this is the thread wide fsum tensor for all kv tiles + Tensor l_i = make_tensor(shape(tDrS)); // (CpyD, NumCpy_M, NumCpy_N) or (CTA_qHLocal * CTA_qL), per thread slice + // this is the CTA wide fsum tensor for all kv tiles + Tensor tDrFsum = make_tensor(Int{}); // (NumQPerWarp), per thread slice of the meta data sum tensor + // then we need a layout to represent the mapping between (warp id (T), V) -> q_idx + // T0V0 -> q_idx 0 + // T0V1 -> q_idx 1 + // T1V0 -> q_idx 2 + // T1V1 -> q_idx 3 + // ... + Layout cta_reduce_tv_q_layout = make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); + + // now we do the data type conversion of S/P from fp32 to bf16/fp8 + using TypeQKV = typename decltype(sS)::value_type; + // allocate per thread rmem space for the converted S/P tensor, the shape is the same as the post partition shape of the output tensor + // the only difference with tDrAcc is here per element is TypeQKV instead of TypeAcc + Tensor tDrS_converted = make_tensor(shape(tDsS)); // (CpyD, NumCpy_M, NumCpy_N), per thread slice of the Acc1/S tensor in rmem, i.e. a row of the Acc1 tensor + // create a converter, convert 1 value at a time, from TypeAcc to TypeQKV + cutlass::NumericConverter converter_s; + + // setup for fmax and fsum cluster reduction + // each reduction participating cta handles CTA_qHLocal * CTA_qL / NumReductionCTA number of q tokens, there are NumReductionCTA participating in cluster reduction + // so we want to have a layout/mapping (reduction cta id (T), local q_idx (V)) -> q_idx + // assume we have 8 q_idx, and 2 reduction ctas, then we have the following mapping: + // (reduction cta 0, local q_idx 0) -> q_idx 0 + // (reduction cta 0, local q_idx 1) -> q_idx 1 + // (reduction cta 0, local q_idx 2) -> q_idx 2 + // (reduction cta 0, local q_idx 3) -> q_idx 3 + // (reduction cta 1, local q_idx 0) -> q_idx 4 + // (reduction cta 1, local q_idx 1) -> q_idx 5 + // (reduction cta 1, local q_idx 2) -> q_idx 6 + // (reduction cta 1, local q_idx 3) -> q_idx 7 + // we create the tv layout for cluster reduction (reduction cta id (T), local q_idx (V)) -> q_idx + Layout cluster_reduce_tv_q_layout = make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + + // for fmax/fsum cluster reduction, ultimately we want to have a layout of (warp_id (T), V) -> (reduction cta id (T'), V') for scattering cta scope fmax/fsum values to reduction ctas + // so we do this by (warp_id (T), V) -> q_idx -> (reduction cta id (T'), V') + // we obtain q_idx -> (reduction cta id (T'), V') by taking the right inverse of cluster_reduce_tv_q_layout + // + // regarding whether we choose right-inverse or left-inverse: + // You can think of right-inverse as a truncated version of left-inverse -- truncated to the compact range of a layout. + // if there's a "hole" in the codomain/range, then the right_inverse stops there. E.g. + // (4,5):(1,5) # There's a "hole" at idx=4 + // right_inverse = 4:1 + // left_inverse = (5,5):(1,4) + // since here cluster_reduce_tv_q_layout is bijective, there is no hole, so right-inverse and left-inverse are the same + Layout cluster_reduce_q_tv_layout = right_inverse(cluster_reduce_tv_q_layout); + // (warp_id (T), V) -> (reduction cta id (T'), V') is obtained through composition + Layout cluster_reduce_tv_tv_layout = composition(cluster_reduce_q_tv_layout, cta_reduce_tv_q_layout); + + // setup for fmax cluster reduction + // thread 0 in each epilog warp will send its local portion of fmax (CTA_qHLocal * CTA_qL / NumEpilogWarps q tokens) to corresponding reduction ctas, + // i.e. the entire epilog warp group will send all fmax to all reduction ctas + // for instance we have 8 q tokens and 4 epilog warps, and 8 reduction ctas, fmax tensor is in rmem: + // warp 0: send fmax(0) (q_idx 0) to reduction cta 0 + // warp 0: send fmax(1) (q_idx 1) to reduction cta 1 + // warp 1: send fmax(0) (q_idx 2) to reduction cta 2 + // warp 1: send fmax(1) (q_idx 3) to reduction cta 3 + // warp 2: send fmax(0) (q_idx 4) to reduction cta 4 + // warp 2: send fmax(1) (q_idx 5) to reduction cta 5 + // warp 3: send fmax(0) (q_idx 6) to reduction cta 6 + // warp 3: send fmax(1) (q_idx 7) to reduction cta 7 + Tensor sFmaxMailbox = shared_storage.tensor_sFmaxMailbox(); // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + int constexpr MaxSplits = size<0>(sFmaxMailbox); + // tSRC means cluster reduction src partitioning pattern, i.e. each CTA holds a particular row of the mailbox tensor of all reduction ctas + Tensor tSRCsFmaxMailbox = sFmaxMailbox(rank,_); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + // destination dsmem tensor for each q token + decltype(tSRCsFmaxMailbox) tSRCdsFmaxMailbox[NumQPerWarp]; // send different q tokens to different reduction cta + + // setup for fsum cluster reduction + Tensor sFsumMailbox = shared_storage.tensor_sFsumMailbox(); // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + // tSRC means cluster reduction src partitioning pattern, i.e. each CTA holds a particular row of the mailbox tensor of all reduction ctas + Tensor tSRCsFsumMailbox = sFsumMailbox(rank,_); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + // destination dsmem tensor for each q token + decltype(tSRCsFsumMailbox) tSRCdsFsumMailbox[NumQPerWarp]; // send different q tokens to different reduction cta + + uint32_t dsmem_maxsum_mailbox_full_barrier_addr[NumQPerWarp]; // dsmem addr for mailbox barrier + + // do mapa and get the dsmem addr of each q token + CUTE_UNROLL + for(int i = 0; i < NumQPerWarp; i++) { + auto cluster_tv = cluster_reduce_tv_tv_layout(warp_idx, i); // it's the 1d coordinate of (reduction cta id (T'), V') + // idx2crd convertes 1d coordinate to natural coordinate, i.e. (reduction cta id (T'), V') + int dst_cta_id = get<0>(idx2crd(cluster_tv, shape(cluster_reduce_tv_q_layout))); + tSRCdsFmaxMailbox[i] = get_dsmem_tensor(tSRCsFmaxMailbox, dst_cta_id); + tSRCdsFsumMailbox[i] = get_dsmem_tensor(tSRCsFsumMailbox, dst_cta_id); + uint32_t smem_maxsum_mailbox_full_barrier_addr = __cvta_generic_to_shared(&shared_storage.maxsum_mailbox_full_barrier); + dsmem_maxsum_mailbox_full_barrier_addr[i] = set_block_rank(smem_maxsum_mailbox_full_barrier_addr, dst_cta_id); + } + + // setup for Acc2 cluster reduction + int row_id = bmm2_tmem_load_t_dH_layout(tid); + Tensor sAcc2Mailbox = shared_storage.tensor_sAcc2Mailbox(); // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + // per thread slice of the Acc2 mailbox tensor, tSRC means cluster reduction src partitioning pattern + // each thread sends each row of the Acc2 tensor to many dst CTAs + Tensor tSRCsAcc2Mailbox = sAcc2Mailbox(row_id,_,rank); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + decltype(tSRCsAcc2Mailbox) tSRCdsAcc2Mailbox[NumReductionCTA]; // send different q tokens to different reduction cta + uint32_t dsmem_acc2_mailbox_full_barrier_addr[NumReductionCTA]; // dsmem addr for mailbox barrier + // do mapa and get the dsmem addr of q token + CUTE_UNROLL + for(int i = 0; i < NumReductionCTA; i++) { + tSRCdsAcc2Mailbox[i] = get_dsmem_tensor(tSRCsAcc2Mailbox, i); + uint32_t smem_acc2_mailbox_full_barrier_addr = __cvta_generic_to_shared(&shared_storage.acc2_mailbox_full_barrier); + dsmem_acc2_mailbox_full_barrier_addr[i] = set_block_rank(smem_acc2_mailbox_full_barrier_addr, i); + } + + Tensor tCsK = shared_storage.tensor_sK(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + Tensor tCsV = shared_storage.tensor_sV(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + int constexpr BMM1_DMA_Stage = size<3>(tCsK); + int constexpr BMM2_DMA_Stage = size<3>(tCsV); + int NumKVTiles = work_tile_info.kvL_idx_end - work_tile_info.kvL_idx_start; + + // setup for cluster reduction + // because each CTA is going to produce (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA), i.e. each CTA handles CTA_qHLocal * CTA_qL / NumReductionCTA q tokens + // we further partition gO into per CTA level + static_assert(CTA_qHLocal * CTA_qL % NumReductionCTA == 0, "CTA_qHLocal * CTA_qL must be divisible by NumReductionCTA"); + Tensor tCTAgO = local_tile(gO, make_shape(Int{}, Int{}), make_coord(0, rank)); // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA) + + // now we do the data type conversion of Acc2 from fp32 to bf16/fp8 + using TypeO = typename decltype(mO)::value_type; + // create a converter, convert 1 value at a time, from TypeAcc to TypeO + cutlass::NumericConverter converter_o; + + // setup for fmax/fsum cluster reduction + int maxsum_st_async_transaction_bytes = NumSplits * CTA_qHLocal * CTA_qL / NumReductionCTA * sizeof(TypeAcc) * 2; + + // participating CTA set the barrier transaction bytes + if ((rank < NumReductionCTA) && (tid == 0)) { + set_barrier_transaction_bytes(shared_storage.maxsum_mailbox_full_barrier, maxsum_st_async_transaction_bytes); + } + + // allocate cluster wide fmax and fsum tensor + Tensor rFmaxCluster = make_tensor(size<1>(sFmaxMailbox)); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + // beta = exp2f(m_ij - m_i) + Tensor rBeta = make_tensor(make_shape(size<1>(sFmaxMailbox), size<0>(sFmaxMailbox))); // (CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + Tensor rFsumCluster = make_tensor(size<1>(sFsumMailbox)); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + // setup for Acc2 cluster reduction + int acc2_st_async_transaction_bytes = CTA_dH * CTA_qHLocal * CTA_qL / NumReductionCTA * NumSplits * sizeof(TypeAcc); + + if ((rank < NumReductionCTA) && (tid == 0)) { + set_barrier_transaction_bytes(shared_storage.acc2_mailbox_full_barrier, acc2_st_async_transaction_bytes); + } + + // final per thread (one row of the output tensor) cluster reduction result + Tensor tAcc2rO = make_tensor(size<1>(sFmaxMailbox)); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + Tensor tAcc2rO_converted = make_tensor(size<1>(sFmaxMailbox)); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + // attention sink needed by this cluster is 1 kv head + Tensor gSink = local_tile(mSink, make_shape(make_shape(Int{}, Int{}), Int<1>{}), + make_coord(make_coord(work_tile_info.qHLocal_idx, work_tile_info.qL_idx), work_tile_info.kvH_idx)); // (CTA_qHLocal, CTA_qL) + Tensor sSink = shared_storage.tensor_sSink(); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + // get the sink tensor needed for this cta, i.e. the sink value needed by each reduction cta + // we achieve this by doing a composition of the gSink tensor and the cluster_reduce_tv_q_layout (aka the tv layout for cluster reduction) + // composing with a tv layout is equivalent to reorganizing the tensor layout into shape (T, V) + // after the composition, the tensor shape is (NumReductionCTA, CTA_qHLocal * CTA_qL / NumReductionCTA) + // then we index into the rank-th row to get the sink value needed by this cta + Tensor tClustergSink = composition(gSink, cluster_reduce_tv_q_layout)(rank, _); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + static_assert(size(sSink) < NumEpilogWarps * 32, "one epilog thread loads at most 1 sink value"); + + /*if ((tid == 0) && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + printf("sS:\t"); print(sS); print("\n"); // (CTA_dH, (CTA_qHLocal, CTA_qL), 1) + printf("tCsS:\t"); print(tCsS); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("tCtAcc1:\t"); print(tCtAcc1); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("sAcc1:\t"); print(sAcc1); print("\n"); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + printf("tCsAcc1:\t"); print(tCsAcc1); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + + printf("tiled_t2r_copy_bmm1:\t"); print(tiled_t2r_copy_bmm1); print("\n"); + printf("thr_t2r_copy_bmm1:\t"); print(thr_t2r_copy_bmm1); print("\n"); + printf("tDtAcc1:\t"); print(tDtAcc1); print("\n"); // (CpyS, NumCpy_M, NumCpy_N) + printf("tDsS:\t"); print(tDsS); print("\n"); // (CpyD, NumCpy_M, NumCpy_N) + printf("tDrS:\t"); print(tDrS); print("\n"); // (CpyD, NumCpy_M, NumCpy_N) + + printf("sFmaxPartial:\t"); print(sFmaxPartial); print("\n"); + printf("tDrFmax:\t"); print(tDrFmax); print("\n"); + printf("tDrFsum:\t"); print(tDrFsum); print("\n"); + printf("cta_reduce_tv_q_layout: "); print_layout(cta_reduce_tv_q_layout); print("\n"); + + printf("tDrS_converted:\t"); print(tDrS_converted); print("\n"); // (CpyD, NumCpy_M, NumCpy_N) + + printf("tDsAcc1:\t"); print(tDsAcc1); print("\n"); // (CpyD, NumCpy_M, NumCpy_N) + printf("tFsumsAcc1:\t"); print(tFsumsAcc1); print("\n"); // ((CTA_qHLocal, CTA_qL), (CTA_kvL / CTA_qHLocal / CTA_qL, (CTA_qHLocal, CTA_qL))) + printf("tFsumrAcc1:\t"); print(tFsumrAcc1); print("\n"); // (CTA_qHLocal * CTA_qL) + + printf("coordK: "); print(coordK); print("\n"); // (kvL, dH, kvH, BS) + printf("cK: "); print(cK); print("\n"); // (CTA_kvL, CTA_dH, Num_CTA_kvL) -> (kvL, dH, kvH, BS) + printf("cur_kvL: "); print(cur_kvL); print("\n"); // (Num_CTA_kvL) -> (kvL, dH, kvH, BS) + printf("seq_len: %d\n", seq_len); + + printf("TypeAcc: %lluB\n", sizeof(TypeAcc)); + + printf("gO:\t"); print(gO); print("\n"); // (CTA_dH, (CTA_qHLocal, CTA_qL)) + printf("tCgO:\t"); print(tCgO); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("tCtAcc2:\t"); print(tCtAcc2); print("\n"); // ((Mma_M, Mma_N), NumMma_M, NumMma_N) + printf("tAcc2rAcc2:\t"); print(tAcc2rAcc2); print("\n"); // (CTA_qHLocal, CTA_qL) + + printf("cluster_reduce_tv_q_layout: "); print_layout(cluster_reduce_tv_q_layout); print("\n"); + printf("cluster_reduce_q_tv_layout: "); print_layout(cluster_reduce_q_tv_layout); print("\n"); + printf("cluster_reduce_tv_tv_layout: "); print_layout(cluster_reduce_tv_tv_layout); print("\n"); + + printf("sFmaxMailbox: "); print(sFmaxMailbox); print("\n"); // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tSRCsFmaxMailbox: "); print(tSRCsFmaxMailbox); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tSRCdsFmaxMailbox[1]: "); print(tSRCdsFmaxMailbox[1]); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("sFsumMailbox: "); print(sFsumMailbox); print("\n"); // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tSRCsFsumMailbox: "); print(tSRCsFsumMailbox); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tSRCdsFsumMailbox[1]: "); print(tSRCdsFsumMailbox[1]); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("dsmem_maxsum_mailbox_full_barrier_addr[1]: %x\n", dsmem_maxsum_mailbox_full_barrier_addr[1]); + + printf("sAcc2Mailbox: "); print(sAcc2Mailbox); print("\n"); // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + printf("tSRCsAcc2Mailbox: "); print(tSRCsAcc2Mailbox); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tSRCdsAcc2Mailbox[2]: "); print(tSRCdsAcc2Mailbox[2]); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("dsmem_acc2_mailbox_full_barrier_addr[2]: %x\n", dsmem_acc2_mailbox_full_barrier_addr[2]); + + printf("tCTAgO:\t"); print(tCTAgO); print("\n"); // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA) + + printf("NumSplits: %d\n", NumSplits); + printf("maxsum_st_async_transaction_bytes: %d\n", maxsum_st_async_transaction_bytes); + + printf("rFmaxCluster: "); print(rFmaxCluster); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("rBeta: "); print(rBeta); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + printf("rFsumCluster: "); print(rFsumCluster); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + printf("acc2_st_async_transaction_bytes: %d\n", acc2_st_async_transaction_bytes); + + printf("tAcc2rO: "); print(tAcc2rO); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tAcc2rO_converted: "); print(tAcc2rO_converted); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + printf("gSink: "); print(gSink); print("\n"); // (CTA_qHLocal, CTA_qL) + printf("sSink: "); print(sSink); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + printf("tClustergSink: "); print(tClustergSink); print("\n"); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + // test inf properties + // -inf - 2 = -inf + // exp2f(-inf) = 0 + // refer to cuda math api for full behavior ragarding inf https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__SINGLE.html#group__cuda__math__single_1ga3e2984de99de67ca680c9bb4f4427f81 + printf("-inf: %f\n", -cutlass::platform::numeric_limits::infinity()); + printf("-inf-2: %f\n", -cutlass::platform::numeric_limits::infinity() - 2.0f); + printf("exp2f(-inf-2): %f\n", exp2f(-cutlass::platform::numeric_limits::infinity() - 2.0f)); + }*/ + + static_assert(MaxSplits <= 32, "we can use 1 warp to initialize mailbox"); + // initialize mailbox tensor for fmax and fsum, when NumSplits < MaxSplits, we need to init those value to -inf and 0 + // because we do reduction on the full tensor (of size MaxSplits) not just the valid splits + // there is no need to init sAcc2 because it will be scaled with beta which will be 0 for invalid splits + if (tid < MaxSplits) { + fill(sFmaxMailbox(tid, _), -cutlass::platform::numeric_limits::infinity()); + clear(sFsumMailbox(tid, _)); + } + + // ensure initialized smem is visible to the entire cluster + cutlass::arch::fence_view_async_shared(); + epilog_barrier.arrive_and_wait(); + + // issue cp.async for sink tensor + if constexpr (!NoSink) { + if (rank < NumReductionCTA) { + // each thread loads one sink value + if (tid < size(sSink)) { + cp_async((int*)(&tClustergSink(tid)), (int*)(&sSink(tid))); + } + } + } + + // only do softmax on valid kv tiles + if (work_tile_info.is_valid()) { + // init rmem tensor to make first iteration correct + fill(tDrFmax, -cutlass::platform::numeric_limits::infinity()); + clear(l_i); + // prolog do arrive on tma_bmm1_full_barrier because first BMM1 doesn't need to wait for preceeding tcgen05.ld Acc1 + if (elect_one_sync()) { + arrive_barrier(shared_storage.tma_bmm1_full_barrier[0]); + } + + // initial phase bit for bmm1_softmax_full_barrier is 0 + // it is flipped to 1 when the BMM1 warp is done, and now we can start the softmax + int bmm1_softmax_full_barrier_phase_bit = 0; + + // initial phase bit for bmm2_epilog_full_barrier is 0 + // it is flipped to 1 when the BMM2 warp is done, and now we can start the epilog + int bmm2_epilog_full_barrier_phase_bit = 0; + + // If sliding window feature is enabled, we need to adjust the seq_len_start. + int seq_len_start = (sliding_window_size == 0) ? 0 : cute::max(0, seq_len - sliding_window_size); + + // loop over each kv tile + for (int kv_tile = 0; kv_tile < NumKVTiles; kv_tile++) { + + // wait for the BMM1 warp to finish + wait_barrier(shared_storage.bmm1_softmax_full_barrier, bmm1_softmax_full_barrier_phase_bit); + bmm1_softmax_full_barrier_phase_bit ^= 1; // flip the phase bit in preparation for next BMM1 + + //if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + // printf("[Softmax] tile %d fmax start\n", kv_tile); + //} + + // softmax on S + // load tmem -> rmem + copy(tiled_t2r_copy_bmm1, tDtAcc1, tDrS); + // there isn't a register dependency on tDrS to honor here to signal tcgen05.ld is done, so we need explicit tcgen05.wait::ld to make sure tcgen05.ld is done + // tcgen05.wait::ld has an implicit __syncwarp, so we don't need a bar.sync across all epilog warp group + cutlass::arch::fence_view_async_tmem_load(); + // only 1 thread per warp signals bmm1 of next iter to start, so in aggregate it's 4 arrives, i.e. each warp sync independently with bmm1, no need for all warp group to be in sync + // tcgen05.ld of Acc1 is done, the BMM1 of next iter can start to overwrite tmem + if (elect_one_sync()) { + arrive_barrier(shared_storage.tma_bmm1_full_barrier[(kv_tile + 1) % BMM1_DMA_Stage]); + } + + // fill the row with -inf if the data is not valid (i.e. over the kv seq len), such that after softmax they will be 0 + // the data per thread is a row of the S/Acc1 tensor of shape (CTA_kvL, (CTA_qHLocal, CTA_qL)) + // if the data is not valid, we fill it with -inf (mask off) such that it will be nop in fmax and fsum reduction + auto kv_seq_len = get<0>(cur_kvL(work_tile_info.kvL_idx_start + kv_tile)); + bool row_valid = (kv_seq_len < seq_len) && (kv_seq_len >= seq_len_start); // only kvL within [seq_len_start, seq_len) is valid + if (!row_valid) { + fill(tDrS, -cutlass::platform::numeric_limits::infinity()); + } + // scale it by softmax_scale_log2, s = s * softmax_scale_log2 = s * softmax_scale * log2(e) + // here we leverage the fact that expf(s) = exp2f(s * log2(e)), and doing the exp2f path generates less instructions than expf (and presumably slightly less accurate) + // so we fuse the log2(e) with the softmax scale, hence all later expf (e.g. alpha/beta) should be replaced by exp2f + // with cuda graph, softmax_scale_log2 is already precalculated on the cpu and hardened as kernel parameter/constant, there is no overhead of this fusion + CUTE_UNROLL + for (int i = 0; i < tDrS.size(); i++) { + tDrS[i] *= softmax_scale_log2; + } + + // calculate fmax across all threads for tile (CTA_kvL, (CTA_qHLocal, CTA_qL)) + // m_ij = fmax(s) + Tensor m_ij = make_tensor(shape(tDrFmax)); // (CTA_qHLocal * CTA_qL), per thread slice of the meta data max tensor for the current kv tile + cta_reduce(tDrS, m_ij, warp_idx, sFmaxPartial, epilog_barrier); + + // m_i_old = m_i + Tensor m_i_old = make_tensor(shape(tDrFmax)); + copy(tDrFmax, m_i_old); + + // m_i = max(m_i_old, m_ij) + CUTE_UNROLL + for (int i = 0; i < tDrFmax.size(); i++) { + tDrFmax[i] = reduce_op(tDrFmax(i), m_ij(i)); + } + + // alpha = exp2f(m_i_old - m_i) + CUTE_UNROLL + for (int i = 0; i < tDrFmax.size(); i++) { + rAlpha[i] = exp2f(m_i_old[i] - tDrFmax[i]); + } + + // p = exp2f(s - m_i) + CUTE_UNROLL + for (int i = 0; i < tDrS.size(); i++) { + tDrS[i] = exp2f(tDrS[i] - tDrFmax[i]); + } + + // l_i = alpha * l_i + p + CUTE_UNROLL + for (int i = 0; i < tDrS.size(); i++) { + l_i[i] = rAlpha[i] * l_i[i] + tDrS[i]; + } + + // convert S/P from fp32 to bf16 + CUTE_UNROLL + for (int i = 0; i < tDrS.size(); i++) { + tDrS_converted[i] = converter_s(tDrS[i]); + } + + // wait for the BMM2 warp to finish, such that we can write a new S tensor to smem for the next iteration of BMM2 + if (kv_tile > 0) { // we can safely overwrite S in the first iter since there is no preceeding BMM2 that consumes S + wait_barrier(shared_storage.bmm2_epilog_full_barrier, bmm2_epilog_full_barrier_phase_bit); + bmm2_epilog_full_barrier_phase_bit ^= 1; // flip the phase bit in preparation for next BMM2 + } + + //if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + // printf("[Softmax] tile %d S store start\n", kv_tile); + //} + + // store bf16 S to smem + copy(tDrS_converted, tDsS); + // fence.proxy.async.shared::cta ensures the smem write from generic proxy (i.e. cuda core) is visible to all async proxy (e.g. tma, tcgen05.mma) + // we want the S write to smem to be visible by tcgen05.mma in the mma warp + // so the fence+synchronization pattern is the following: + // 1. every thread in the epilog warp issues a fence.proxy.async.shared::cta such that the smem write by this thread is visible to async proxy + // 2. use named barrier to sync all threads in the epilog warp, so when all threads reach this point, every thread has executed the fence.proxy.async.shared::cta + // at this point the smem write is visible to tcgen05.mma unit + // 3. we use mbarrier to notify mma warp that the S tensor is ready and visible to tcgen05.mma unit and bmm2 is ready to start (if V load is also done) + // + // fence.proxy.async.shared::cta + cutlass::arch::fence_view_async_shared(); + // only need to sync within the warp, baceuse later each warp independently signals bmm2 to start, no need for all warp group to be in sync + __syncwarp(); + + // do the correction here, skip for first iter + // because we already checked BMM2 is done before the S store, so we can tcgen05.ld Acc2 + // also first iter value of alpha=0, if there is perf issue, we don't need the if + if (kv_tile > 0) { + // bmm2 load tmem -> rmem + // only the first 16 threads in each warp contains the Acc2 tensor data, each thread holds a row of the Acc2 tensor + // the rest 16 threads hold bogus data + tmem_load(acc2_tmem_addr, tAcc2rAcc2); + // acc2 = alpha * acc2 + CUTE_UNROLL + for (int i = 0; i < tAcc2rAcc2.size(); i++) { + tAcc2rAcc2[i] *= rAlpha[i]; + } + // now use tcgen05.st to store corrected Acc2 back to tmem + tmem_store(tAcc2rAcc2, acc2_tmem_addr); + // we need a tcgen05.wait::st to make sure tcgen05.st result is visible to async proxy (i.e. tcgen05.mma) + // tcgen05.wait::st has an implicit __syncwarp, so we don't need a bar.sync across all epilog warp group + cutlass::arch::fence_view_async_tmem_store(); + } + + // the warp is in sync, just do arrive per warp to signal bmm2 to start + // only 1 thread per warp signals bmm2 to start, so in aggregate it's 4 arrives + if (elect_one_sync()) { + arrive_barrier(shared_storage.tmasoftmax_bmm2_full_barrier[kv_tile % BMM2_DMA_Stage]); + } + } + + // overlapped with the last BMM2 + // do transpose of l_i in rmem for faster fsum reduction + copy(l_i, tDsAcc1); + // make the smem write visible to all threads in the epilog warp + epilog_barrier.arrive_and_wait(); + // now each thread holds a partial column of the Acc1 tensor + copy(tFsumsAcc1, tFsumrAcc1); + + // l_ij = sum(p) + int constexpr Interval = CTA_kvL / CTA_qHLocal / CTA_qL; + // do fsum reduction among Interval number of threads + tDrFsum(0) = cta_reduce_transposed(tFsumrAcc1); + // get the rest of the fsum values (different q tokens) to thread 0 + // now thread 0, Interval, 2 * Interval, ... have the fsum values for each q token + CUTE_UNROLL + for (int i = 1; i < tDrFsum.size(); i++) { + tDrFsum(i) = __shfl_sync(0xffffffff, tDrFsum(0), i * Interval, 32); + } + + // up until here, the content of tDrFmax and tDrFsum is: + // tDrFmax: warp 0-3: [fmax(q0), fmax(q1), fmax(q2), fmax(q3), fmax(q4), fmax(q5), fmax(q6), fmax(q7)] + // tDrFsum: warp 0 [fsum(q0), fsum(q1)] + // warp 1 [fsum(q2), fsum(q3)] + // warp 2 [fsum(q4), fsum(q5)] + // warp 3 [fsum(q6), fsum(q7)] + // st.async expects the layout of tDrFsum, so we rearrange tDrFmax to the same layout as tDrFsum + static_assert(NumEpilogWarps == 4, "NumEpilogWarps must be 4"); + if (tid == 32) { + int constexpr WarpId = 1; + CUTE_UNROLL + for (int i = 0; i < NumQPerWarp; i++) { + tDrFmax(i) = tDrFmax(WarpId * NumQPerWarp + i); + } + } + else if (tid == 64) { + int constexpr WarpId = 2; + CUTE_UNROLL + for (int i = 0; i < NumQPerWarp; i++) { + tDrFmax(i) = tDrFmax(WarpId * NumQPerWarp + i); + } + } + else if (tid == 96) { + int constexpr WarpId = 3; + CUTE_UNROLL + for (int i = 0; i < NumQPerWarp; i++) { + tDrFmax(i) = tDrFmax(WarpId * NumQPerWarp + i); + } + } + + // we can start sending cta level fmax and fsum to remote using st.async, + // tid 0 in each epilog warp will send its local portion of fmax and fsum to corresponding reduction cta + // when we have 8 q tokens in total and 4 epilog warps, each thread sends 2 q tokens to corresponding reduction cta + if ((tid % 32) == 0) { + CUTE_UNROLL + for (int i = 0; i < NumQPerWarp; i++) { + auto cluster_tv = cluster_reduce_tv_tv_layout(warp_idx, i); // it's the 1d coordinate of (reduction cta id (T'), V') + // idx2crd convertes 1d coordinate to natural coordinate, i.e. (reduction cta id (T'), V') + int local_q_idx = get<1>(idx2crd(cluster_tv, shape(cluster_reduce_tv_q_layout))); + store_shared_remote_f32(tDrFmax(i), (uint32_t)&tSRCdsFmaxMailbox[i](local_q_idx), dsmem_maxsum_mailbox_full_barrier_addr[i]); + store_shared_remote_f32(tDrFsum(i), (uint32_t)&tSRCdsFsumMailbox[i](local_q_idx), dsmem_maxsum_mailbox_full_barrier_addr[i]); + } + } + + // wait for the BMM2 warp to finish + wait_barrier(shared_storage.bmm2_epilog_full_barrier, bmm2_epilog_full_barrier_phase_bit); + + //if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + // printf("[Epilog] barrier full\n"); + //} + + // bmm2 load tmem -> rmem + // for dH=64, only the first 16 threads in each warp contains the Acc2 tensor data, each thread holds a row of the Acc2 tensor + // the rest 16 threads hold bogus data + // for dH=128, all 32 threads hold valid data + tmem_load(acc2_tmem_addr, tAcc2rAcc2); + int lane_id = cutlass::canonical_lane_idx(); + if (lane_id < (CTA_dH / NumEpilogWarps)) { + // do the st.async of acc2 to remote, rmem -> dsmem + CUTE_UNROLL + for (int i = 0; i < tAcc2rAcc2.size(); i++) { + // use the exact same partition layout as fmax and fsum cluster reduction + // i.e. each reduction cta handles CTA_qHLocal * CTA_qL / NumReductionCTA number of q tokens + // here we do q_idx -> (reduction cta id (T'), V'), similarly, we use idx2crd to convert 1d coordinate to natural coordinate + int dst_cta_id = get<0>(idx2crd(cluster_reduce_q_tv_layout(i), shape(cluster_reduce_tv_q_layout))); + int local_q_idx = get<1>(idx2crd(cluster_reduce_q_tv_layout(i), shape(cluster_reduce_tv_q_layout))); + store_shared_remote_f32(tAcc2rAcc2(i), (uint32_t)&tSRCdsAcc2Mailbox[dst_cta_id](local_q_idx), dsmem_acc2_mailbox_full_barrier_addr[dst_cta_id]); + } + } + + // signal the mma warp tcgen05.ld of bmm2 is done, can start deallocate all tmem + tmem_allocation_barrier.arrive(); + } + + // only NumReductionCTA number of reduction ctas will do the reduction + if (rank < NumReductionCTA) { + // handle oob when NumSplits < MaxSplits + // doing computation on MaxSplits is faster than doing NumSplits with if statement + clear(rBeta); + + if constexpr (!NoSink) { + // wait for cp.async of sink tensor to finish + if (tid < size(sSink)) { + cp_async_fence(); + cp_async_wait<0>(); + } + // make the sink tensor visible to all threads in the epilog warp + epilog_barrier.arrive_and_wait(); + } + + // the reduction cta wait for st.async of fmax and fsum to finish + int maxsum_mailbox_full_barrier_phase_bit = 0; + wait_barrier(shared_storage.maxsum_mailbox_full_barrier, maxsum_mailbox_full_barrier_phase_bit); + + //if (elect_one_sync() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + // printf("[Reduction] barrier full\n"); + //} + + // every thread redundantly calculates cluster wide fmax for (CTA_qHLocal * CTA_qL / NumReductionCTA) q tokens + CUTE_UNROLL + for (int i = 0; i < rFmaxCluster.size(); i++) { + // cache the mailbox tensor in rmem to hide ld.shared latency, if we demand load ld.shared, its latency will hurt pipelining/unrolling amount + // so we manually pipeline the ld.shared first and do all the operation on rmem, + // this is similar to sw pipelining but we do it manually because compiler is not smart enough to figure it out + // instruction sequencing is important because we are an in-order issue machine, meaning stall of current instruction issue will block all subsequent instruction issue + // many times compiler is not smart enough to do it, so you need to write the code in the right order to help it + Tensor rFmaxMailbox = make_tensor(Int{}); + copy(sFmaxMailbox(_,i), rFmaxMailbox); + + Tensor rFsumMailbox = make_tensor(Int{}); + copy(sFsumMailbox(_,i), rFsumMailbox); + + rFmaxCluster(i) = rFmaxMailbox(0); + + CUTE_UNROLL + for (int j = 1; j < MaxSplits; j++) { + rFmaxCluster(i) = reduce_op(rFmaxCluster(i), rFmaxMailbox(j)); + } + + // after fmax is done, we can do beta and fsum calculation + // 2 optimizations we are doing: + // 1. if we put NumSplits in the trip count, this makes it into a dynamic loop, and the index j into rBeta becomes a dynamic index + // but rBeta is in RF which doesn't support dynamic indexing, so the compiler will put it onto the stack (local memory) instead and generate ld.local/st.local instructions + // we want to keep rBeta in RF to avoid the local memory access, so we put MaxSplits in the trip count + // 2. we don't use predication of if statement to avoid the access when j >= NumSplits, instead we set the remainder sum/max value to 0 and inf and loop over MaxSplits + // when we use predicate, compiler will generate branch instruction which prevents the loop from doing more software pipelining, i.e. overlapping independent instructions + // from different iterations + TypeAcc acc = 0.0f; + CUTE_UNROLL + for (int j = 0; j < MaxSplits; j++) { + rBeta(i, j) = exp2f(rFmaxMailbox(j) - rFmaxCluster(i)); + acc += rBeta(i, j) * rFsumMailbox(j); + } + rFsumCluster(i) = acc; + if constexpr (!NoSink) { + rFsumCluster(i) += exp2f(sSink(i) * Log2_E - rFmaxCluster(i)); + } + } + + if (tid < CTA_dH) { + // initial phase bit for acc2_mailbox_full_barrier is 0 + // it is flipped to 1 when the acc2 st.async are done, and now we can start the reduction + int acc2_mailbox_full_barrier_phase_bit = 0; + // wait for the acc2 st.async to finish + wait_barrier(shared_storage.acc2_mailbox_full_barrier, acc2_mailbox_full_barrier_phase_bit); + + // now we can do the reduction, reduce across Acc2 of all splits + CUTE_UNROLL + for (int i = 0; i < tAcc2rO.size(); i++) { + TypeAcc acc = 0.0f; + CUTE_UNROLL + for (int j = 0; j < MaxSplits; j++) { + acc += rBeta(i, j) * sAcc2Mailbox(tid, i, j); + } + // o = o / l_i + tAcc2rO(i) = acc / rFsumCluster(i); + } + + // convert from fp32 to bf16 + CUTE_UNROLL + for (int i = 0; i < tAcc2rO.size(); i++) { + tAcc2rO_converted(i) = converter_o(tAcc2rO(i)); + } + + // rmem -> gmem + copy(tAcc2rO_converted, tCTAgO(tid, _)); + } + } + + // debug + /*epilog_barrier.arrive_and_wait(); + if ((tid == 0) && (blockIdx.x == 15) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + int example_row = 0; + + //print("sAcc1:\t"); print(sAcc1); print("\n"); // s = q * k * softmax_scale_log2 + print("tDrFmax:\t"); print_tensor(tDrFmax); print("\n"); // m_ij = fmax(s) + print("tDrS:\t"); print_tensor(tDrS); print("\n"); // p = exp2f(s - m_ij) + print("tDrFsum:\t"); print_tensor(tDrFsum); print("\n"); // l_ij = sum(p) + print("tDrS_converted:\t"); print_tensor(tDrS_converted); print("\n"); + print("sS:\t"); print_tensor(sS(example_row,_,0)); print("\n"); + + Tensor sP = shared_storage.tensor_sP(); // ((CTA_qHLocal, CTA_qL), CTA_kvL, BMM2_DMA_Stage) + print("sP:\t"); print_tensor(sP(_,example_row,0)); print("\n"); + + print("tAcc2rAcc2:\t"); print_tensor(tAcc2rAcc2); print("\n"); + + print("sFmaxMailbox: "); print_tensor(sFmaxMailbox); print("\n"); + print("sFsumMailbox: "); print_tensor(sFsumMailbox); print("\n"); + print("rFmaxCluster: "); print_tensor(rFmaxCluster); print("\n"); + print("rBeta: "); print_tensor(rBeta); print("\n"); + print("rFsumCluster: "); print_tensor(rFsumCluster); print("\n"); + + print("sAcc2Mailbox: "); print_tensor(sAcc2Mailbox(example_row,_,_)); print("\n"); + print("tAcc2rO: "); print_tensor(tAcc2rO); print("\n"); + print("tAcc2rO_converted: "); print_tensor(tAcc2rO_converted); print("\n"); + print("gO: "); print_tensor(gO(example_row,_)); print("\n"); + }*/ +} + +// K has shape (kvL, dH, kvH, BS) +// Q has shape ((qHLocal, qL), dH, kvH, BS) +// V has shape (dH, kvL, kvH, BS) +// O has shape (dH, (qHLocal, qL), kvH, BS) +// sinks has shape ((qHLocal, qL), kvH) +// seq_len has shape (BS) +template < + class SharedStorage, + class KTensor, class QTensor, class VTensor, class OTensor, class SinkTensor, + class TmaAtomK, class TmaAtomQ, class TmaAtomV, + class TiledBMM1, class TiledBMM2, + class TypeAcc, + int CTA_qHLocal, int CTA_qL, int CTA_kvL, int CTA_dH, + int BMM1_DMA_Stage, int BMM2_DMA_Stage, + int MaxSplits, int NumReductionCTA, + bool NoSink> +__maxnreg__(128) +__global__ void +gqa_device( + KTensor mK, + QTensor mQ, + VTensor mV, + OTensor mO, + SinkTensor mSink, + int* seq_lens, + CUTE_GRID_CONSTANT TmaAtomK const tma_atom_K, + CUTE_GRID_CONSTANT TmaAtomQ const tma_atom_Q, + CUTE_GRID_CONSTANT TmaAtomV const tma_atom_V, + TiledBMM1 tiled_bmm1, + TiledBMM2 tiled_bmm2, + float softmax_scale_log2, + int sliding_window_size, + int pdl_count +) { + //if (threadIdx.x == 0) { + // printf("[%d, %d, %d] gqa_device\n", blockIdx.x, blockIdx.y, blockIdx.z); + //} + + // Allocate SMEM + extern __shared__ char shared_memory[]; + SharedStorage& shared_storage = *reinterpret_cast(shared_memory); + + // WorkTileInfo, for non persistent static scheduler, cta id is the work tile info + // since loading the seq_lens is on the critical path of the prolog, we want to start it as soon as possible + int kvH = shape<2>(mK); + int BS_idx = blockIdx.x / kvH; + // only thread 0 issues cp.async to load the seq_len + if (threadIdx.x == 0) { + cp_async(&seq_lens[BS_idx], &shared_storage.seq_len); + } + + //if (threadIdx.x == 0) { + // printf("[%d, %d, %d] work_tile_info: %d, %d, %d, %d, %d\n", blockIdx.x, blockIdx.y, blockIdx.z, work_tile_info.BS_idx, work_tile_info.kvH_idx, work_tile_info.kvL_idx_start, work_tile_info.kvL_idx_end, work_tile_info.dH_idx, work_tile_info.qHLocal_idx, work_tile_info.qL_idx); + //} + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + // barrier initialization, warp 0 does initialization + if (warp_idx == 0) { + // transaction barrier because tma arrive on it, 6 thread arrive: one for DMA_Q warp, one for DMA_KV (K fetch) warp, and 4 for softmax warp (tcgen05.ld Acc1 is done) + cutlass::arch::detail::initialize_barrier_array_aligned(shared_storage.tma_bmm1_full_barrier, /* arrival count */ 6); + // 1 thread (BMM1) arrive to signal DMA_Q and DMA_KV (K fetch) warp + cutlass::arch::detail::initialize_barrier_array_aligned(shared_storage.tma_bmm1_empty_barrier, /* arrival count */ 1); + // transaction barrier because tma arrive on it, 5 thread arrive: one for DMA_KV (V fetch) warp and 4 for softmax warp (S/P store) + cutlass::arch::detail::initialize_barrier_array_aligned(shared_storage.tmasoftmax_bmm2_full_barrier, /* arrival count */ 5); + // 1 thread (BMM2) arrive to signal DMA_KV (V fetch) and softmax warp (P store) + cutlass::arch::detail::initialize_barrier_array_aligned(shared_storage.tma_bmm2_empty_barrier, /* arrival count */ 1); + // 1 thread (BMM1) arrive to signal softmax + cutlass::arch::detail::initialize_barrier_array_aligned(&shared_storage.bmm1_softmax_full_barrier, /* arrival count */ 1); + // 1 thread (BMM2) arrive to signal epilog + cutlass::arch::detail::initialize_barrier_array_aligned(&shared_storage.bmm2_epilog_full_barrier, /* arrival count */ 1); + // 1 thread (epilog) arrive to signal maxsum + cutlass::arch::detail::initialize_barrier_array_aligned(&shared_storage.maxsum_mailbox_full_barrier, /* arrival count */ 1); + // 1 thread (epilog) arrive to signal acc2 + cutlass::arch::detail::initialize_barrier_array_aligned(&shared_storage.acc2_mailbox_full_barrier, /* arrival count */ 1); + } + // Sync tmem allocation status between MMA and softmax/epilogue warps within CTA + // 32 threads (mma) + 128 threads (epilog) to sync + // also used for tmem deallocation between epilog warps and mma warps within CTA + cutlass::arch::NamedBarrier tmem_allocation_barrier(32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + // syncing all threads (128) within 4 epilog warps + cutlass::arch::NamedBarrier epilog_barrier(128, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + + /*if (thread0() && (blockIdx.x == 0) && (blockIdx.y == 0) && (blockIdx.z == 0)) { + Tensor tCsK = shared_storage.tensor_sK(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + Tensor tCsQ = shared_storage.tensor_sQ(); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + Tensor tCsV = shared_storage.tensor_sV(); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + Tensor sS = shared_storage.tensor_sS(); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + Tensor sP = shared_storage.tensor_sP(); // ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) + + printf("tCsK:\t"); print(tCsK); print("\n"); + printf("tCsQ:\t"); print(tCsQ); print("\n"); + printf("tCsV:\t"); print(tCsV); print("\n"); + printf("sS:\t"); print(sS); print("\n"); + printf("sP:\t"); print(sP); print("\n"); + }*/ + + // barrier initialization needs to be visible to all warps + // defer it as late as possible to allow some thread divergence in prolog + cutlass::arch::fence_barrier_init(); +#if 0 + // this will have a membar.gpu to ensure dsmem write visibility within the entire cluster, because there isn't a membar.cluster + // membar.gpu is 0.2us + cluster_sync(); +#else + // the alternative is to use proper fences + // at the ptx level, fence.mbarrier_init.release.cluster act as a release fence (in cluster scope) for mbarrier init op + cutlass::arch::fence_barrier_init(); + // thread 0 waits for its previously issued cp.async to complete + // here we overlap the cp.async with the barrier initialization as much as possible + if (threadIdx.x == 0) { + cp_async_fence(); + cp_async_wait<0>(); + } + + // the cluster sync serves two purposes: + // 1. it waits for all threads in the cluster to see the barrier initialization + // 2. it waits for all threads in the cta to see the cp.async result in smem + cluster_arrive_relaxed(); + cluster_wait(); +#endif + + // bid.y includes rasterization of qHLocal and qL + // so if we have (2, 4) shape (qHLocal, qL) cta, we do the following mapping from linear cta id to 2d id + // 0 -> (0, 0) + // 1 -> (1, 0) + // 2 -> (0, 1) + // 3 -> (1, 1) + // 4 -> (0, 2) + // 5 -> (1, 2) + // 6 -> (0, 3) + // 7 -> (1, 3) + int qHLocal = shape<0>(shape<0>(mQ)); + int num_qHLocal = cutlass::ceil_div(qHLocal, CTA_qHLocal); + // bid.z is the split id for kvL split, try to evenly distribute the kvL blocks to every CTA + int rank = blockIdx.z; + int seq_len = shared_storage.seq_len; + // Sliding window optimization: when enabled, we only process tokens in range + // [seq_len - sliding_window_size, seq_len). To simplify tile distribution, + // we align the skip boundary to CTA_kvL tiles. + // + // Example with seq_len=100, sliding_window_size=30, CTA_kvL=16: + // - Tokens we care about: [70, 100) + // - Skip offset aligned to tiles: floor(70 / 16) * 16 = 64 + // - Workload range: [64, 100), length = 36 + // - Tiles to process: ceil(36 / 16) = 3 tiles + int workload_seq_len = seq_len; + int seq_len_skip_offset = 0; // Number of sequence positions to skip (tile-aligned) + if (sliding_window_size > 0 && sliding_window_size < seq_len) { + int unaligned_skip = seq_len - sliding_window_size; + seq_len_skip_offset = (unaligned_skip / CTA_kvL) * CTA_kvL; // Align to tile boundary + workload_seq_len = seq_len - seq_len_skip_offset; + } + int NumKVTiles = cutlass::ceil_div(workload_seq_len, CTA_kvL); + int NumSplits = cute::min(MaxSplits, NumKVTiles); + int kvL_tile_count_per_cta = NumKVTiles / MaxSplits; + int kvL_tile_count_remainder = NumKVTiles % MaxSplits; + int kvL_tile_count = kvL_tile_count_per_cta + (rank < kvL_tile_count_remainder ? 1 : 0); + int kvL_tile_count_skip_offset = seq_len_skip_offset / CTA_kvL; + int kvL_tile_count_start = rank * kvL_tile_count_per_cta + (rank < kvL_tile_count_remainder ? rank : kvL_tile_count_remainder) + kvL_tile_count_skip_offset; + int kvL_tile_count_end = kvL_tile_count_start + kvL_tile_count; + WorkTileInfo work_tile_info { + .BS_idx = (int32_t)BS_idx, + .kvH_idx = (int32_t)(blockIdx.x % kvH), + .kvL_idx_start = (int32_t)kvL_tile_count_start, + .kvL_idx_end = (int32_t)kvL_tile_count_end, + .dH_idx = 0, // no dH tiling for bmm2 + .qHLocal_idx = (int32_t)(blockIdx.y % num_qHLocal), + .qL_idx = (int32_t)(blockIdx.y / num_qHLocal), + .is_valid_tile = (kvL_tile_count > 0) + }; + + if (warp_idx == 0) { + DMA_Q_warp(shared_storage, work_tile_info, mQ, &tma_atom_Q, tiled_bmm1); + } + else if (warp_idx == 1) { + DMA_KV_warp(shared_storage, work_tile_info, mK, mV, &tma_atom_K, &tma_atom_V, tiled_bmm1, tiled_bmm2); + } + else if (warp_idx == 2) { + MMA_warp(shared_storage, work_tile_info, mO, tiled_bmm1, tiled_bmm2, tmem_allocation_barrier); + } + else if (warp_idx >= 4) { + // epilog tid is from 128 to 255, need to offset by -128 when getting the per thread slice + int tid = threadIdx.x - 128; + // warp_idx - 4 because epilog warp group starts from warp 4 + EPILOG_warp(shared_storage, work_tile_info, mK, mO, mSink, seq_len, tiled_bmm1, tiled_bmm2, softmax_scale_log2, sliding_window_size, tmem_allocation_barrier, epilog_barrier, NumSplits, tid, warp_idx - 4, rank); + } + + __syncthreads(); +} + +// K has shape (kvL, dH, kvH, BS) +// Q has shape ((qHLocal, qL), dH, kvH, BS) +// V has shape (dH, kvL, kvH, BS) +// O has shape (dH, (qHLocal, qL), kvH, BS) +// kvL is max_seq_len, seq_lens[BS] is the actual seq len for each batch +// sinks has shape (qHLocal * kvH), i.e. one sink per q head, when device_ptr_sinks is nullptr, it's disabled +// sliding_window_size is the size of the sliding window, when it's 0, it's disabled +template< + class TypeQKV, class TypeO, class TypeAcc, + int CTA_qHLocal, int CTA_qL, int CTA_kvL, int CTA_dH, + int BMM1_DMA_Stage, int BMM2_DMA_Stage, + int MaxSplits, int NumReductionCTA> +void gqa_host( + TypeQKV* device_ptr_K, + TypeQKV* device_ptr_Q, + TypeQKV* device_ptr_V, + TypeO* device_ptr_O, + int* seq_lens, + TypeAcc* device_ptr_sinks, + int kvH, int qHLocal, int qL, int kvL, int dH, int BS, + int stride_K_kvH, int stride_K_kvL, int stride_K_dH, int stride_K_BS, + int stride_Q_kvH, int stride_Q_qHLocal, int stride_Q_qL, int stride_Q_dH, int stride_Q_BS, + int stride_V_kvH, int stride_V_kvL, int stride_V_dH, int stride_V_BS, + int stride_O_kvH, int stride_O_qHLocal, int stride_O_qL, int stride_O_dH, int stride_O_BS, + float softmax_scale, + int sliding_window_size, + bool pdl, int pdl_count = -1, + cudaStream_t stream = 0 +) { + Layout layout_K = make_layout_K(kvH, kvL, dH, BS, stride_K_kvH, stride_K_kvL, stride_K_dH, stride_K_BS); + Layout layout_Q = make_layout_Q(kvH, qHLocal, qL, dH, BS, stride_Q_kvH, stride_Q_qHLocal, stride_Q_qL, stride_Q_dH, stride_Q_BS); + Layout layout_V = make_layout_V(kvH, kvL, dH, BS, stride_V_kvH, stride_V_kvL, stride_V_dH, stride_V_BS); + Layout layout_O = make_layout_O(kvH, qHLocal, qL, dH, BS, stride_O_kvH, stride_O_qHLocal, stride_O_qL, stride_O_dH, stride_O_BS); + Layout layout_sinks = make_layout_sinks(qHLocal, qL, kvH); + + // how we handle oob: + // oob for K, Q, V are handled by TMA + // oob for O is explicitly handled by predicate in the epilog since it uses simple st.global epilog + // we partition kvL with tile size of CTA_kvL, and we evenly distribute the kvL blocks to MaxSplits number of cta in the cluster + assert(NumReductionCTA <= MaxSplits); + static_assert(((CTA_qHLocal * CTA_qL) % NumReductionCTA) == 0, "each reduction cta must have even number of q tokens"); + + //printf("Running for problem shape (kvH x qHLocal x qL x kvL x dH x BS): %d x %d x %d x %d x %d x %d\n", kvH, qHLocal, qL, kvL, dH, BS); + //printf("with tile size CTA_qHLocal: %d, CTA_qL: %d, CTA_kvL: %d, CTA_dH: %d\n", CTA_qHLocal, CTA_qL, CTA_kvL, CTA_dH); + //printf("layout_K: "); print(layout_K); printf("\n"); + //printf("layout_Q: "); print(layout_Q); printf("\n"); + //printf("layout_V: "); print(layout_V); printf("\n"); + //printf("layout_O: "); print(layout_O); printf("\n"); + + // create the cute tensor + Tensor mK = make_tensor(make_gmem_ptr(device_ptr_K), layout_K); // (kvL, dH, kvH, BS) + Tensor mQ = make_tensor(make_gmem_ptr(device_ptr_Q), layout_Q); // ((qHLocal, qL), dH, kvH, BS) + Tensor mV = make_tensor(make_gmem_ptr(device_ptr_V), layout_V); // (dH, kvL, kvH, BS) + Tensor mO = make_tensor(make_gmem_ptr(device_ptr_O), layout_O); // (dH, (qHLocal, qL), kvH, BS) + Tensor mSink = make_tensor(make_gmem_ptr(device_ptr_sinks), layout_sinks); // ((qHLocal, qL), kvH) + + //printf("mK: "); print(mK); printf("\n"); + //printf("mQ: "); print(mQ); printf("\n"); + //printf("mV: "); print(mV); printf("\n"); + //printf("mO: "); print(mO); printf("\n"); + //printf("mSink: "); print(mSink); printf("\n"); + + static_assert(CTA_kvL == 128, "BMM1's MMA_M needs to be 128 for tcgen05.ld->softmax"); + static_assert(((CTA_qHLocal * CTA_qL) % 8) == 0, "BMM1's MMA_N needs to be divisible by 8 for tcgen05.mma"); + assert(dH == CTA_dH); // bmm1 only has 1 kblock (i.e. 1 Q tile), bmm2 deal with all dH for now, in the foreseable future this is the hardest constraint to lift + static_assert((CTA_dH == 128) || (CTA_dH == 64), "BMM2's MMA_M needs to be at either 128 or 64 for tcgen05.ld->correction"); + // we swap AB so bmm1 is K (CTA_kvL, CTA_dH) x Q (CTA_dH, CTA_qHLocal * CTA_qL) + // both Q and K are dH (K in gemm terminology) major + // M = CTA_kvL, N = CTA_qHLocal * CTA_qL, K = CTA_dH + TiledMMA tiled_bmm1 = cutlass::gemm::collective::detail::sm100_make_1sm_trivial_tiled_mma< + TypeQKV, TypeQKV, TypeAcc, // Mma's A, B, and Accumulator types + Shape, Int, Int>, // TileShape_MNK + Shape<_1, _1, _1>, // ClusterShape_MNK + cute::UMMA::Major::K, cute::UMMA::Major::K>(); + + // we swap AB for bmm2 as well, V (dH, CTA_kvL) x P (CTA_kvL, CTA_qHLocal * CTA_qL) + // V is dH (M in gemm terminology) major, P is CTA_kvL (K in gemm terminology) major in smem after each thread writes P from rmem to smem + // M = CTA_dH, N = CTA_qHLocal * CTA_qL, K = CTA_kvL + TiledMMA tiled_bmm2 = cutlass::gemm::collective::detail::sm100_make_1sm_trivial_tiled_mma< + TypeQKV, TypeQKV, TypeAcc, // Mma's A, B, and Accumulator types + Shape, Int, Int>, // TileShape_MNK + Shape<_1, _1, _1>, // ClusterShape_MNK + cute::UMMA::Major::MN, cute::UMMA::Major::K>(); + + // We can also print and inspect the tiled_mma + //print(tiled_bmm1); + //print(tiled_bmm2); + + // Pre-partitioned smem Tile Shape to post-partitioned smem tile shape ((Mma_M, Mma_K), NumMma_M, NumMma_K, DMA_Stage) + // bmm1 has BMM1_DMA_Stage stage + // tile K shape is (CTA_kvL, CTA_dH, BMM1_DMA_Stage) and is the A operand + // tile Q shape is ((CTA_qHLocal, CTA_qL), CTA_dH, 1) and is the B operand, because all K tiles share the same Q tiles, we only have 1 stage of Q tiles + // tile S shape is (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) and is the C operand + auto shape_K = make_shape(Int{}, Int{}, Int{}); + auto shape_Q = make_shape(make_shape(Int{}, Int{}), Int{}, Int<1>{}); + auto shape_S = make_shape(Int{}, make_shape(Int{}, Int{}), Int<1>{}); + auto mma_shape_K = partition_shape_A(tiled_bmm1, shape_K); + auto mma_shape_Q = partition_shape_B(tiled_bmm1, shape_Q); + + // bmm2 has 1 stage because we explicitly only has 1 kvL tile + // tile V shape is (CTA_dH, CTA_kvL, BMM2_DMA_Stage) and is the A operand + // tile P shape is ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) and is the B operand, since 1 P tile is enough to overlap MMA with softmax + auto shape_V = make_shape(Int{}, Int{}, Int{}); + auto shape_P = select<1, 0, 2>(shape_S); // just a permutation of shape_S + auto mma_shape_V = partition_shape_A(tiled_bmm2, shape_V); + + // Print and inspect shapes for this example. note that this is only the shape of the smem tile, not the actual smem layout + //print("shape_K: "); print(shape_K); print("\n"); + //print("shape_Q: "); print(shape_Q); print("\n"); + //print("shape_S: "); print(shape_S); print("\n"); + //print("shape_V: "); print(shape_V); print("\n"); + //print("shape_P: "); print(shape_P); print("\n"); + //print("mma_shape_K: "); print(mma_shape_K); print("\n"); + //print("mma_shape_Q: "); print(mma_shape_Q); print("\n"); + //print("mma_shape_V: "); print(mma_shape_V); print("\n"); + + // choose the swizzle atom for K, Q, S, V and P + // for bmm1, both K and Q are dH (K in gemm terminology) major + // and S is M major, none of the swizzle layout gives bank conflict free store to smem for S, + // swizzle 64B/128B should give at most 2-way bank conflict which is good enough since the output tile is small + auto SmemLayoutAtomK = cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, // gmem layout of K + TypeQKV, // data type of K + decltype(shape<0>(shape_K)), decltype(shape<1>(shape_K))>(); // tile size of K + auto SmemLayoutAtomQ = cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, // gmem layout of Q + TypeQKV, // data type of Q + decltype(shape<0>(shape_Q)), decltype(shape<1>(shape_Q))>(); // tile size of Q + // we use TypeQKV because we don't actually want to store the bf32 output to smem, we do the softmax and quantize it to bf16 and then store it as P + // hence the smem layout atom should reflect the bf16 precision + // for bmm1 tcgen05.ld, each register is holding a row of S (CTA_kvL is mapped to the thread dimension), if we do st.shared from rmem to smem, to avoid bank conflict, we need to put T0V0, T1V0, T2V0, ... T31V0 contiguously in smem + // then the smem layout of S is M (CTA_kvL) major, so we choose MN major swizzle atom + auto SmemLayoutAtomS = UMMA::Layout_MN_SW128_Atom{}; + + // for bmm2, V is dH (M in gemm terminology) major, P is CTA_kvL (K in gemm terminology) major + auto SmemLayoutAtomV = cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, // gmem layout of V + TypeQKV, // data type of V + decltype(shape<0>(shape_V)), decltype(shape<1>(shape_V))>(); // tile size of V + // swizzle atom for P should be the transpose of the swizzle atom for S, because they literally represent the same tensor just with different dimension order (aka a pytorch transpose) + auto SmemLayoutAtomP = UMMA::Layout_K_SW128_Atom{}; + + // Print and inspect SmemLayoutAtomK, SmemLayoutAtomQ, SmemLayoutAtomS, SmemLayoutAtomV, and SmemLayoutAtomP for this example. + //print("SmemLayoutAtomK: "); print(SmemLayoutAtomK); print("\n"); + //print("SmemLayoutAtomQ: "); print(SmemLayoutAtomQ); print("\n"); + //print("SmemLayoutAtomS: "); print(SmemLayoutAtomS); print("\n"); + //print("SmemLayoutAtomV: "); print(SmemLayoutAtomV); print("\n"); + //print("SmemLayoutAtomP: "); print(SmemLayoutAtomP); print("\n"); + + // finally construct the smem layout for tile K, Q, V, S, and P + auto sK_layout = UMMA::tile_to_mma_shape(SmemLayoutAtomK, mma_shape_K); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM1_DMA_Stage) + auto sQ_layout = UMMA::tile_to_mma_shape(SmemLayoutAtomQ, mma_shape_Q); // ((Mma_N, Mma_K), NumMma_N, NumMma_K, 1) + auto sV_layout = UMMA::tile_to_mma_shape(SmemLayoutAtomV, mma_shape_V); // ((Mma_M, Mma_K), NumMma_M, NumMma_K, BMM2_DMA_Stage) + // S and P use tile_to_shape as we do the mma partition in the kernel later + // here we use Step to enforce same swizzle atom stacking order for S and P, + // with Step<_1, _2, _3>, basically we say the swizzle atom is first stacked along the first dimension (CTA_kvL), then the second dimension ((CTA_qHLocal, CTA_qL)), then the third dimension (1) + auto sS_layout = tile_to_shape(SmemLayoutAtomS, shape_S, Step<_1, _2, _3>{}); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + // with Step<_2, _1, _3>, basically we say the swizzle atom is first stacked along the second dimension (CTA_kvL), then the first dimension ((CTA_qHLocal, CTA_qL)), then the third dimension (1) + auto sP_layout = tile_to_shape(SmemLayoutAtomP, shape_P, Step<_2, _1, _3>{}); // ((CTA_qHLocal, CTA_qL), CTA_kvL, 1) + auto sAcc1_layout = make_layout(shape_S); // (CTA_kvL, (CTA_qHLocal, CTA_qL), 1) + // for storing fmax and fsum warp reduce partial results + int constexpr NumEpilogWarps = 4; + // NumEpilogWarps contiguous because we often ld.shared all NumEpilogWarps from 1/32 threads, this has best vectorization + auto sWarpReduce_layout = make_layout(make_shape(Int{}, make_shape(Int{}, Int{}))); // (NumEpilogWarps, (CTA_qHLocal, CTA_qL)) + // MaxSplits contiguous because we often ld.shared all MaxSplits from 1/32 threads, this has best vectorization + auto sMSMailbox_Layout = make_layout(make_shape(Int{}, Int{})); // (MaxSplits, CTA_qHLocal * CTA_qL / NumReductionCTA) + // default layout is CTA_dH contiguous to maximize st.async/ld.shared bw + auto sAcc2Mailbox_layout = make_layout(make_shape(Int{}, Int{}, Int{})); // (CTA_dH, CTA_qHLocal * CTA_qL / NumReductionCTA, MaxSplits) + auto sSinks_layout = make_layout(Int{}); // (CTA_qHLocal * CTA_qL / NumReductionCTA) + + // Print and inspect sK_layout, sQ_layout, sV_layout, and sP_layout for this example. + //print("sK_layout: "); print(sK_layout); print("\n"); + //print("sQ_layout: "); print(sQ_layout); print("\n"); + //print("sV_layout: "); print(sV_layout); print("\n"); + //print("sS_layout: "); print(sS_layout); print("\n"); + //print("sP_layout: "); print(sP_layout); print("\n"); + //print("sAcc1_layout: "); print(sAcc1_layout); print("\n"); + //print("sWarpReduce_layout: "); print(sWarpReduce_layout); print("\n"); + //print("sMSMailbox_Layout: "); print(sMSMailbox_Layout); print("\n"); + //print("sAcc2Mailbox_layout: "); print(sAcc2Mailbox_layout); print("\n"); + + // Now we can find the SMEM allocation size + using SMEMStorage = SharedStorage; + + static_assert(BMM1_DMA_Stage >= BMM2_DMA_Stage, "otherwise you are wasting BMM2 stage because BMM1 TMA issue will block BMM2 TMA due to insufficient BMM1 stages"); + + // create TMA descriptors for K, Q, V, and P matrices + Copy_Atom tma_atom_K = make_tma_atom( + SM90_TMA_LOAD{}, // TMA Load Op, sm100 reuses sm90 tma atom + mK, // Source GMEM tensor + take<0, 3>(sK_layout), // Destination SMEM layout for 1 DMA_Stage, ((Mma_M, Mma_K), NumMma_M, NumMma_K) + make_shape(get<0>(shape_K), get<1>(shape_K)) // TMA box shape, it's cosize must match the cosize of the destination smem layout + ); + // this is an arithmetic tuple, denoting the coordinate of the top-left corner of the TMA box + Tensor mK_tma = tma_atom_K.get_tma_tensor(shape(mK)); // (kvL, dH, kvH, BS) + //print("tma_atom_K:\t"); print(tma_atom_K); print("\n"); + //print("mK_tma:\t"); print(mK_tma); print("\n"); + + Copy_Atom tma_atom_Q = make_tma_atom( + SM90_TMA_LOAD{}, // TMA Load Op, sm100 reuses sm90 tma atom + mQ, // Source GMEM tensor + // sQ_layout(_,_,_,Int<0>{}) doesn't work + // basically under some corner cases, composedlayout indexing is incorrectly implemented, so we use the take method which is also correct + take<0, 3>(sQ_layout), // Destination SMEM layout for 1 DMA_Stage, ((Mma_N, Mma_K), NumMma_N, NumMma_K) + make_shape(get<0>(shape_Q), get<1>(shape_Q)) // TMA box shape, it's cosize must match the cosize of the destination smem layout + ); + // this is an arithmetic tuple, denoting the coordinate of the top-left corner of the TMA box + Tensor mQ_tma = tma_atom_Q.get_tma_tensor(shape(mQ)); // ((qHLocal, qL), dH, kvH, BS) + //print("tma_atom_Q:\t"); print(tma_atom_Q); print("\n"); + //print("mQ_tma:\t"); print(mQ_tma); print("\n"); + + Copy_Atom tma_atom_V = make_tma_atom( + SM90_TMA_LOAD{}, // TMA Load Op, sm100 reuses sm90 tma atom + mV, // Source GMEM tensor + take<0, 3>(sV_layout), // Destination SMEM layout for 1 DMA_Stage, ((Mma_M, Mma_K), NumMma_M, NumMma_K) + make_shape(get<0>(shape_V), get<1>(shape_V)) // TMA box shape, it's cosize must match the cosize of the destination smem layout + ); + // this is an arithmetic tuple, denoting the coordinate of the top-left corner of the TMA box + Tensor mV_tma = tma_atom_V.get_tma_tensor(shape(mV)); // (dH, kvL, kvH, BS) + //print("tma_atom_V:\t"); print(tma_atom_V); print("\n"); + //print("mV_tma:\t"); print(mV_tma); print("\n"); + + int smemBytes = sizeof(SMEMStorage); + + // print smemBytes + //printf("smemBytes: %d\n", smemBytes); + + // invoke the kernel + cudaLaunchConfig_t config; + cudaLaunchAttribute attrs[2]; + // bid.x: kvH * BS, bid.y: qHLocal * qL, bid.z: kvL + uint32_t Cluster_Size = cute::max(MaxSplits, NumReductionCTA); + config.gridDim = dim3{ + (uint32_t)kvH * BS, + (uint32_t)cutlass::ceil_div(qHLocal, CTA_qHLocal) * cutlass::ceil_div(qL, CTA_qL), + Cluster_Size}; + config.blockDim = 256; // 8 warps + config.dynamicSmemBytes = smemBytes; + config.stream = stream; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim = {1, 1, Cluster_Size}; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = pdl ? 2 : 1; + + if (device_ptr_sinks != nullptr) { + auto *kernel_instance = + &gqa_device; + gpuErrChk(cudaFuncSetAttribute(*kernel_instance, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes)); + // portable max cluster size is 8, but sm100a supports 16, need explicit opt in + gpuErrChk(cudaFuncSetAttribute(*kernel_instance, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); + gpuErrChk(cudaLaunchKernelEx(&config, kernel_instance, mK_tma, mQ_tma, mV_tma, mO, mSink, + seq_lens, + tma_atom_K, tma_atom_Q, tma_atom_V, + tiled_bmm1, tiled_bmm2, + softmax_scale * Log2_E, sliding_window_size, pdl_count)); + } + else { + auto *kernel_instance = + &gqa_device; + gpuErrChk(cudaFuncSetAttribute(*kernel_instance, cudaFuncAttributeMaxDynamicSharedMemorySize, smemBytes)); + // portable max cluster size is 8, but sm100a supports 16, need explicit opt in + gpuErrChk(cudaFuncSetAttribute(*kernel_instance, cudaFuncAttributeNonPortableClusterSizeAllowed, 1)); + gpuErrChk(cudaLaunchKernelEx(&config, kernel_instance, mK_tma, mQ_tma, mV_tma, mO, mSink, + seq_lens, + tma_atom_K, tma_atom_Q, tma_atom_V, + tiled_bmm1, tiled_bmm2, + softmax_scale * Log2_E, sliding_window_size, pdl_count)); + } +} + +} // namespace gqa +} // namespace TGV diff --git a/examples/94_ada_fp8_blockwise/CMakeLists.txt b/examples/94_ada_fp8_blockwise/CMakeLists.txt new file mode 100644 index 00000000..f5c9d30f --- /dev/null +++ b/examples/94_ada_fp8_blockwise/CMakeLists.txt @@ -0,0 +1,32 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +cutlass_example_add_executable( + 94_ada_fp8_blockwise + ada_fp8_blockwise.cu + ) diff --git a/examples/94_ada_fp8_blockwise/ada_fp8_blockwise.cu b/examples/94_ada_fp8_blockwise/ada_fp8_blockwise.cu new file mode 100644 index 00000000..f70cf159 --- /dev/null +++ b/examples/94_ada_fp8_blockwise/ada_fp8_blockwise.cu @@ -0,0 +1,489 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 An FP8 blockwise scaled GEMM example for the NVIDIA Ada SM89 architecture using CUTLASS. +*/ + + +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/device/gemm_blockwise.h" +#include "cutlass/gemm/threadblock/mma_multistage_blockwise.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/gemm.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "helper.h" +#include +#include +#include +#include +#include +#include +#include +#include "cutlass/util/command_line.h" + +using cutlass::ceil_div; + +using ElementA = cutlass::float_e4m3_t; +using ElementB = cutlass::float_e4m3_t; +using ElementOutput = cutlass::bfloat16_t; +using ElementAccumulator = float; +using ElementScale = float; + +using LayoutA = cutlass::layout::RowMajor; +using LayoutB = cutlass::layout::ColumnMajor; +using LayoutC = cutlass::layout::RowMajor; // Currently only RowMajor is supported +using LayoutScale = cutlass::layout::RowMajor; // Currently only RowMajor is supported + +static int const AlignmentA = 16; +static int const AlignmentB = 16; +static int const Stages = 3; +static int const BlockSize = 128; + +constexpr float epsilon = 0.51f; +constexpr float floor_val = 1.0f; + +// ----------------------------------------------------------------------------- +// Command line options structure +// ----------------------------------------------------------------------------- +struct Options { + bool help = false; + int m = 1024; + int n = 1024; + int k = 1024; + float alpha = 1.f; + float beta = 0.f; + bool verify = true; + int iterations = 1000; + int warmup = 1000; + + // Parse command line arguments using CUTLASS helper + void parse(int argc, char const **argv) { + cutlass::CommandLine cmd(argc, argv); + + if (cmd.check_cmd_line_flag("help")) { + help = true; + return; + } + + cmd.get_cmd_line_argument("m", m); + cmd.get_cmd_line_argument("n", n); + cmd.get_cmd_line_argument("k", k); + cmd.get_cmd_line_argument("alpha", alpha, alpha); + cmd.get_cmd_line_argument("beta", beta, beta); + cmd.get_cmd_line_argument("verify", verify); + cmd.get_cmd_line_argument("iterations", iterations); + cmd.get_cmd_line_argument("warmup", warmup); + } + + std::ostream &print_usage(std::ostream &out) const { + out << "94_ada_fp8_blockwise\n\n" + << " FP8 GEMM with blockwise scaling (Ada, Sm89).\n\n" + << "Options:\n\n" + << " --help Display this help string\n" + << " --m= GEMM M dimension (default 1024)\n" + << " --n= GEMM N dimension (default 1024)\n" + << " --k= GEMM K dimension (default 1024)\n" + << " --alpha= Epilogue alpha (default 1.0)\n" + << " --beta= Epilogue beta (default 0.0)\n" + << " --verify= Verify the results (default true)\n" + << " --iterations= Number of timing iterations (default 1000)\n" + << " --warmup= Number of warmup iterations (default 1000)\n"; + return out; + } + + double gflops(double runtime_s) const { + uint64_t flop = uint64_t(2) * m * n * k; + double gflop = double(flop) / double(1.0e9); + return gflop / runtime_s; + } +}; + +using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination< + ElementOutput, 8, ElementAccumulator, ElementAccumulator>; + +using Gemm = cutlass::gemm::device::GemmBlockwise< + ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC, + ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89, + cutlass::gemm::GemmShape<64, 128, 128>, + cutlass::gemm::GemmShape<64, 64, 128>, cutlass::gemm::GemmShape<16, 8, 32>, + ElementScale, LayoutScale, BlockSize, EpilogueOutputOp, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, Stages, + AlignmentA, AlignmentB, false, cutlass::arch::OpMultiplyAdd>; + +// Host-side verification +static bool verify_gemm(int M, int N, int K, ElementA const *A, + ElementB const *B, ElementOutput const *C, + ElementOutput const *D, ElementScale const *scale_A, + ElementScale const *scale_B, float alpha, float beta) { + std::vector A_fp8_host(M * K); + std::vector B_fp8_host(K * N); + std::vector C_bf16_host(M * N); + std::vector D_bf16_host(M * N); + + int kBlocks = ceil_div(K, BlockSize); + int mBlocks = ceil_div(M, BlockSize); + int nBlocks = ceil_div(N, BlockSize); + + std::vector ScaleA(mBlocks * kBlocks); + std::vector ScaleB(nBlocks * kBlocks); + + CUDA_CHECK( + cudaMemcpy(A_fp8_host.data(), A, + sizeof(ElementA) * A_fp8_host.size(), + cudaMemcpyDeviceToHost)); + + CUDA_CHECK( + cudaMemcpy(B_fp8_host.data(), B, + sizeof(ElementB) * B_fp8_host.size(), + cudaMemcpyDeviceToHost)); + + CUDA_CHECK( + cudaMemcpy(C_bf16_host.data(), C, + sizeof(ElementOutput) * C_bf16_host.size(), + cudaMemcpyDeviceToHost)); + + CUDA_CHECK( + cudaMemcpy(D_bf16_host.data(), D, + sizeof(ElementOutput) * D_bf16_host.size(), + cudaMemcpyDeviceToHost)); + + CUDA_CHECK( + cudaMemcpy(ScaleA.data(), scale_A, + sizeof(ElementScale) * ScaleA.size(), + cudaMemcpyDeviceToHost)); + + CUDA_CHECK( + cudaMemcpy(ScaleB.data(), scale_B, + sizeof(ElementScale) * ScaleB.size(), + cudaMemcpyDeviceToHost)); + + // -------------------------------------------------------------------- + // De-quantize A and B into FP32 using blockwise scales + // -------------------------------------------------------------------- + std::vector A_f32(M * K); + for (int m = 0; m < M; ++m) { + int blk_m = m / BlockSize; + for (int k = 0; k < K; ++k) { + int blk_k = k / BlockSize; + float sA = ScaleA[blk_m * kBlocks + blk_k]; + A_f32[m * K + k] = static_cast(A_fp8_host[m * K + k]) * sA; + } + } + + std::vector B_f32(K * N); + for (int n = 0; n < N; ++n) { + int blk_n = n / BlockSize; + for (int k = 0; k < K; ++k) { + int blk_k = k / BlockSize; + float sB = ScaleB[blk_n * kBlocks + blk_k]; + B_f32[k + n * K] = static_cast(B_fp8_host[k + n * K]) * sB; + } + } + + // -------------------------------------------------------------------- + // Prepare tensor refs and run reference GEMM + // -------------------------------------------------------------------- + cutlass::TensorRef A_ref(A_f32.data(), K); + cutlass::TensorRef B_ref(B_f32.data(), K); + + // Convert C (BF16) to FP32 once + std::vector C_f32(M * N); + cutlass::NumericConverter bf16_to_f32; + std::transform(C_bf16_host.begin(), C_bf16_host.end(), C_f32.begin(), + [&](ElementOutput x) { return bf16_to_f32(x); }); + cutlass::TensorRef C_ref(C_f32.data(), N); + + // Output buffer in FP32 (will later hold ReLU result) + std::vector D_ref_f32(M * N, 0.0f); + cutlass::TensorRef D_ref_tensor(D_ref_f32.data(), N); + + cutlass::gemm::GemmCoord problem{M, N, K}; + + cutlass::reference::host::Gemm + gemm_ref; + + gemm_ref(problem, alpha, A_ref, B_ref, + beta, C_ref, D_ref_tensor, 0.0f); + + // Build tensor views for relative comparison + cutlass::TensorView ref_view(D_ref_f32.data(), LayoutC(N), {M, N}); + + // Convert device output (BF16) to float for comparison + std::vector kernel_output(M * N); + cutlass::TensorView kernel_output_view(kernel_output.data(), LayoutC(N), {M, N}); + cutlass::TensorView d_bf16_view(D_bf16_host.data(), LayoutC(N), {M, N}); + cutlass::reference::host::TensorCopy(kernel_output_view, d_bf16_view); + + bool result = cutlass::reference::host::TensorRelativelyEquals( + ref_view, kernel_output_view, epsilon, floor_val); + +// std::cout << "ref:\n" << ref_view << "\n\n\n"; +// std::cout << "compute:\n" << kernel_output_view << "\n\n\n"; + + // Compute error metrics + double mse = cutlass::reference::host::TensorMSE(kernel_output_view, ref_view); + double mre = cutlass::reference::host::TensorMRE(kernel_output_view, ref_view); + double max_error = cutlass::reference::host::TensorGreatestError(kernel_output_view, ref_view); + + std::cout << " Result MSE: " << mse << ", MRE: " << mre << ", greatest error: " << max_error << std::endl; + std::cout << "GEMM Verification result is " << (result ? "PASS" : "FAIL") + << std::endl; + + return result; +} + +static void initialize_tensors( + int M, int N, int K, cutlass::HostTensor &tensor_A, + cutlass::HostTensor &tensor_B, + cutlass::HostTensor &tensor_C, + cutlass::HostTensor &tensor_D, + cutlass::HostTensor &tensor_ScaleA, + cutlass::HostTensor &tensor_ScaleB) { + int mBlocks = ceil_div(M, BlockSize); + int nBlocks = ceil_div(N, BlockSize); + int kBlocks = ceil_div(K, BlockSize); + + uint64_t seed = 2024; + + // Resize tensors ------------------------------------------------------ + tensor_A.resize({M, K}); + tensor_B.resize({K, N}); + tensor_C.resize({M, N}); + tensor_D.resize({M, N}); + + tensor_ScaleA.resize({mBlocks, kBlocks}); + tensor_ScaleB.resize({nBlocks, kBlocks}); + + // Fill A and B with random uniform values in [-2, 2] + cutlass::reference::host::TensorFillRandomUniform( + tensor_A.host_view(), seed + 1, 2.0, -2.0); + cutlass::reference::host::TensorFillRandomUniform( + tensor_B.host_view(), seed + 2, 2.0, -2.0); + + // Fill C and D with random uniform values in [-2, 2] + cutlass::reference::host::TensorFillRandomUniform( + tensor_C.host_view(), seed + 3, 2.0, -2.0); + cutlass::reference::host::TensorFillRandomUniform( + tensor_D.host_view(), seed + 4, 2.0, -2.0); + + // Fill scale tensors with random uniform values in [-1, 1] + cutlass::reference::host::TensorFillRandomUniform( + tensor_ScaleA.host_view(), seed + 5, 1.0, -1.0); + cutlass::reference::host::TensorFillRandomUniform( + tensor_ScaleB.host_view(), seed + 6, 1.0, -1.0); + + tensor_A.sync_device(); + tensor_B.sync_device(); + tensor_C.sync_device(); + tensor_D.sync_device(); + tensor_ScaleA.sync_device(); + tensor_ScaleB.sync_device(); +} + +// This example requires CUDA 12.4 or greater and sm89 or higher +bool sufficient() { + + if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 4)) { + std::cerr << "This example requires CUDA 12.4 or greater." << std::endl; + return false; + } + + size_t smem_size = sizeof(typename Gemm::GemmKernel::SharedStorage); + + cudaDeviceProp properties; + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + std::cerr << "cudaGetDevice() failed with error: " << cudaGetErrorString(result) << std::endl; + return false; + } + + result = cudaGetDeviceProperties(&properties, device_idx); + + if (result != cudaSuccess) { + std::cerr << "cudaGetDeviceProperties() failed with error: " << cudaGetErrorString(result) << std::endl; + return false; + } + + if (properties.major < 8 || (properties.major == 8 && properties.minor < 9)) { + std::cerr << "CUTLASS's Ada FP8 BlockwiseGEMM example requires a device of compute capability 89 or higher.\n" << std::endl; + return false; + } + + if (properties.sharedMemPerBlockOptin < smem_size) { + std::cerr << "Insufficient shared memory. Need " << smem_size + << ", but device only has " << properties.sharedMemPerBlockOptin << std::endl; + return false; + } + + return true; +} + + +int main(int argc, char const **args) { + // --------------------------------------------------------------------------- + // Parse command-line options + // --------------------------------------------------------------------------- + + Options options; + options.parse(argc, args); + + if (options.help) { + options.print_usage(std::cout) << std::endl; + return 0; + } + + // Problem dimensions + const int M = options.m; + const int N = options.n; + const int K = options.k; + + // Waive test if insufficient CUDA device + if (!sufficient()) { + std::cerr << "Insufficient resources to run the kernel." << std::endl; + return 0; + } + + cutlass::HostTensor tensor_A; + cutlass::HostTensor tensor_B; + cutlass::HostTensor tensor_C; + cutlass::HostTensor tensor_D; + cutlass::HostTensor tensor_ScaleA; + cutlass::HostTensor tensor_ScaleB; + + + initialize_tensors(M, N, K, tensor_A, tensor_B, tensor_C, tensor_D, tensor_ScaleA, tensor_ScaleB); + + ElementA const *ptr_A = tensor_A.device_data(); + ElementB const *ptr_B = tensor_B.device_data(); + ElementOutput const *ptr_C = tensor_C.device_data(); + ElementOutput *ptr_D = tensor_D.device_data(); + + float *a_ptr = tensor_ScaleA.device_data(); + float *b_ptr = tensor_ScaleB.device_data(); + + typename Gemm::EpilogueOutputOp::Params epilogue_params(options.alpha, options.beta); + + int kBlocks = ceil_div(K, BlockSize); + int ldA = kBlocks; + int ldB = kBlocks; + + using LayoutScale = Gemm::LayoutScale; + using TensorRefScale = Gemm::TensorRefScale; + + LayoutScale layout_A(ldA); + LayoutScale layout_B(ldB); + TensorRefScale ref_scale_A(a_ptr, layout_A); + TensorRefScale ref_scale_B(b_ptr, layout_B); + + // Construct the argument list ---------------------------------------- + typename Gemm::Arguments arguments( + /* problem_size */ {M, N, K}, + /* A */ {ptr_A, K}, + /* B */ {ptr_B, K}, + /* C */ {ptr_C, N}, + /* D */ {ptr_D, N}, + /* scale_A tensor ref */ ref_scale_A, + /* scale_B tensor ref */ ref_scale_B, + /* epilogue params */ epilogue_params, + /* split-k slices */ 1, + /* gather A indices */ nullptr, + /* gather B indices */ nullptr, + /* scatter D indices */ nullptr); + + Gemm gemm_op; + + cutlass::Status status = gemm_op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "GEMM cannot implement the given problem." << std::endl; + return -1; + } + + size_t workspace_bytes = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_bytes); + + status = gemm_op.initialize(arguments, workspace.get()); + if (status != cutlass::Status::kSuccess) { + std::cerr << "GEMM initialization failed." << std::endl; + return -1; + } + + status = gemm_op(); + if (status != cutlass::Status::kSuccess) { + std::cerr << "GEMM execution failed." << std::endl; + return -1; + } + + // Verification + if (options.verify) { + std::cout << "Verifying GEMM" << std::endl; + bool ok = verify_gemm(M, N, K, ptr_A, ptr_B, ptr_C, ptr_D, a_ptr, b_ptr, + options.alpha, options.beta); + if (!ok) { + std::cerr << "Verification failed." << std::endl; + return -1; + } + } + + // Profiling loop + if (options.iterations > 0) { + GpuTimer timer; + for (int iter = 0; iter < options.warmup + options.iterations; ++iter) { + if (iter == options.warmup) + timer.start(); + CUTLASS_CHECK(gemm_op.run()); + } + timer.stop(); + + float elapsed_ms = timer.elapsed_millis(); + double avg_runtime_ms = double(elapsed_ms) / double(options.iterations); + double gflops = options.gflops(avg_runtime_ms / 1000.0); + std::cout << "Avg runtime: " << avg_runtime_ms << " ms" << std::endl; + std::cout << "GFLOPS: " << gflops << std::endl; + fflush(stdout); + } + + return 0; +} diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 60baa3b2..f27834d9 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1,4 +1,3 @@ - # Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause # @@ -171,6 +170,10 @@ foreach(EXAMPLE 90_sm103_fp4_ultra_grouped_gemm 91_fp4_gemv 92_blackwell_moe_gemm + 93_blackwell_low_latency_gqa + 94_ada_fp8_blockwise + 111_hopper_ssd + 112_blackwell_ssd ) add_subdirectory(${EXAMPLE}) diff --git a/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py new file mode 100644 index 00000000..d68d9fc9 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py @@ -0,0 +1,1832 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +import argparse +from typing import Optional, Tuple, Type, Union +from functools import lru_cache +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.testing as testing +from cutlass.cute.runtime import from_dlpack +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cute.nvgpu import cpasync, tcgen05 + +""" +A high-performance cluster launch control(CLC) dynamic persistent batched dense GEMM example +for the NVIDIA Blackwell SM100 architecture using CUTE DSL. + +The CLC dynamic persistent scheduling technique performs dynamic loading balancing. +It has the ability to adapt available SMs rather than a statically selected number. To support this, +a new instruction is introduced to query for a new tile to compute. This new instruction is similar +to programmatic multicast in context of clusters in that the same starting tile ID for a given cluster +is broadcasted to all threadblocks in the cluster. +See `PTX documentation `. + +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") + +This GEMM kernel supports the following features: + - Utilizes Tensor Memory Access (TMA) for efficient memory operations + - Utilizes Blackwell's tcgen05.mma for matrix multiply-accumulate (MMA) operations (including 2cta mma instructions) + - Implements TMA multicast with cluster to reduce L2 memory traffic + - Support CLC dynamic persistent tile scheduling to have near perfect load balancing + - Support warp specialization to avoid explicit pipelining between mainloop load and mma + +This GEMM works as follows: +1. DMA warp: Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using TMA operations. +2. MMA warp: Perform matrix multiply-accumulate (MMA) operations using tcgen05.mma instruction. +3. EPILOGUE warp: + - Load completed accumulator from tensor memory (TMEM) to registers (RMEM) using tcgen05.ld. + - Type convert C matrix to output type. + - Optionally store C matrix from registers (RMEM) to shared memory (SMEM) to global memory (GMEM) with TMA operations, + or directly store C matrix from registers (RMEM) to global memory (GMEM) without TMA operations. + - Optionally accept an elementwise lambda function epilogue_op to apply to the output tensor: + e.g., relu can set epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + +SM100 tcgen05.mma instructions operate as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Write accumulator to TMEM +The accumulator in TMEM must then be loaded to registers before writing back to GMEM. + +Input arguments to this example is same as dense_gemm.py. + +.. code-block:: bash + + python examples/blackwell/dense_gemm_persistent_dynamic.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/blackwell/dense_gemm_persistent_dynamic.py \ + --ab_dtype Float16 --c_dtype Float16 --acc_dtype Float32 \ + --mma_tiler_mn 256,128 --cluster_shape_mn 2,1 \ + --mnkl 8192,8192,8192,1 \ + --use_tma_store --use_2cta_instrs \ + --warmup_iterations 1 --iterations 10 --skip_ref_check + + +Constraints are same as dense_gemm.py: +* Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2), + see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation +* A/B tensor must have the same data type +* Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) +* Mma tiler N must be 32-256, step 32 +* Cluster shape M/N must be positive and power of 2, total cluster size <= 16 +* Cluster shape M must be multiple of 2 if use_2cta_instrs=True +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 4, 8, and 16 for TFloat32, + Float16/BFloat16, and Int8/Uint8/Float8, respectively. +* OOB tiles are not allowed when TMA store is disabled +""" + + +def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + smem_capacity: int, + occupancy: int, + use_tma_store: bool, + c_smem_layout: Union[cute.Layout, None], +) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param smem_capacity: Total available shared memory capacity in bytes. + :type smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + :param use_tma_store: Whether TMA store is enabled. + :type use_tma_store: bool + :param c_smem_layout: Layout of C operand in shared memory, or None if not using TMA store. + :type c_smem_layout: Union[cute.Layout, None] + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 2 + + # Default C stages + num_c_stage = 2 if use_tma_store else 0 + + # Calculate smem layout and size for one stage of A, B, and C with 1-stage + a_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = utils.sm100.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + + ab_bytes_per_stage = cute.size_in_bytes( + a_dtype, a_smem_layout_stage_one + ) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + if use_tma_store: + num_c_stage += ( + smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage + + +class PersistentDenseGemmKernel: + """This class implements batched matrix multiplication (C = A x B) with support for various data types + and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization. + + :param acc_dtype: Data type for accumulation during computation + :type acc_dtype: type[cutlass.Numeric] + :param use_2cta_instrs: Whether to use CTA group 2 for advanced thread cooperation + :type use_2cta_instrs: bool + :param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N) + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Whether to use Tensor Memory Access (TMA) for storing results + :type use_tma_store: bool + + :note: In current version, A and B tensor must have the same data type + - i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported + + :note: Supported A/B data types: + - TFloat32 + - Float16/BFloat16 + - Int8/Uint8 + - Float8E4M3FN/Float8E5M2 + + :note: Supported accumulator data types: + - Float32 (for all floating point A/B data types) + - Float16 (only for fp16 and fp8 A/B data types) + - Int32 (only for uint8/int8 A/B data types) + + :note: Supported C data types: + - Float32 (for float32 and int32 accumulator data types) + - Int32 (for float32 and int32 accumulator data types) + - Float16/BFloat16 (for fp16 and fp8 accumulator data types) + - Int8/Uint8 (for uint8/int8 accumulator data types) + - Float8E4M3FN/Float8E5M2 (for float32 accumulator data types) + + :note: Constraints: + - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) + - MMA tiler N must be 32-256, step 32 + - Cluster shape M must be multiple of 2 if use_2cta_instrs=True + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + + **Example:** + gemm = PersistentDenseGemmKernel( + acc_dtype=cutlass.Float32, + use_2cta_instrs=True, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(2, 2) + ) + gemm(a, b, c, max_active_clusters, stream) + """ + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + use_tma_store: bool, + ): + """Initializes the configuration for a Blackwell dense GEMM kernel. + + This configuration includes several key aspects: + + 1. MMA Instruction Settings (tcgen05): + - acc_dtype: Data types for MMA accumulator. + - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler. + - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant + with cta_group=2 should be used. + + 2. Cluster Shape: + - cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster. + + 3. Output C tensor store mode: + - use_tma_store: Boolean indicating whether to use Tensor Memory Access (TMA) for storing results. + + :param acc_dtype: Data type of the accumulator. + :type acc_dtype: type[cutlass.Numeric] + :param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction. + :type mma_tiler_mn: Tuple[int, int] + :param use_2cta_instrs: Boolean, True to use cta_group=2 MMA variant. + :type use_2cta_instrs: bool + :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. + :type cluster_shape_mn: Tuple[int, int] + :param use_tma_store: Use Tensor Memory Access (TMA) or normal store for output C tensor. + :type use_tma_store: bool + """ + + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler_mn = mma_tiler_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.use_tma_store = use_tma_store + self.arch = "sm_100" + + self.cta_group = ( + tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + ) + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + *self.epilogue_warp_id, + ) + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + + def _create_tiled_mma(self): + return utils.sm100.make_trivial_tiled_mma( + self.a_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + # Configure tiled mma + tiled_mma = self._create_tiled_mma() + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Compute epilogue subtile + if cutlass.const_expr(self.use_tma_store): + self.epi_tile = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + else: + self.epi_tile = self.cta_tile_shape_mnk[:2] + + c_smem_layout = None + if cutlass.const_expr(self.use_tma_store): + c_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, 1 + ) + + self.smem_capacity = utils.get_smem_capacity_in_bytes() + + # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = _compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.smem_capacity, + self.occupancy, + self.use_tma_store, + c_smem_layout, + ) + + # Setup clc stage by default + self.num_clc_stage = 1 + + # Compute A/B/C shared memory layout + self.a_smem_layout_staged = utils.sm100.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = utils.sm100.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + + self.c_smem_layout_staged = None + if self.use_tma_store: + self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage + ) + + # Compute the number of tensor memory allocation columns + self.num_tmem_alloc_cols = self._compute_num_tmem_alloc_cols( + tiled_mma, self.mma_tiler, self.num_acc_stage, self.arch + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM operation in steps: + - Setup static attributes before smem/grid/tma computation + - Setup TMA load/store atoms and tensors + - Compute grid size with regard to hardware constraints + - Define shared storage for kernel + - Launch the kernel synchronously + + :param a: Input tensor A + :type a: cute.Tensor + :param b: Input tensor B + :type b: cute.Tensor + :param c: Output tensor C + :type c: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + :raises AssertionError: If OOB (Out-Of-Bounds) tiles are present when TMA store is disabled. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + tiled_mma = self._create_tiled_mma() + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, tiled_mma.thr_id + ) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if a.element_type is cutlass.Float32 else None + ), + ) + + # Setup TMA load for B + b_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, tiled_mma.thr_id + ) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=( + cutlass.TFloat32 if b.element_type is cutlass.Float32 else None + ), + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + # Response size is 4B * 4 elements + self.num_clc_response_bytes = 16 + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + if cutlass.const_expr(self.use_tma_store): + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem_layout, self.epi_tile + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, self.cta_tile_shape_mnk, self.cluster_shape_mn + ) + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c if self.use_tma_store else c, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + return + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.use_tma_store): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + # Define shared storage for kernel + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.num_acc_stage * 2 + ] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + clc_ptr: cute.struct.MemRange[cutlass.Int64, self.num_clc_stage * 2] + clc_response_ptr: cute.struct.MemRange[cutlass.Int32, 1] + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * ( + 2 if use_2cta_instrs else 1 + ) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize clc_pipeline (barrier) and states + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + cluster_size = cute.size(self.cluster_shape_mn) + num_clc_consumer_threads = 32 * ( + 1 + cluster_size * (1 + len(self.epilogue_warp_id) + 1) + ) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_pipeline = pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem_dealloc_barrier = None + if cutlass.const_expr(not self.use_tma_store): + tmem_dealloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_dealloc_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + # Initial clc response pointer + clc_response_ptr = storage.clc_response_ptr.data_ptr() + + clc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + # (MMA, MMA_N, MMA_K, STAGE) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + # (bN, bK, RestN, RestK, RestL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gA_mkl, mode=[3]) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgA = thr_mma.partition_A(gA_mkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, self.num_acc_stage) + ) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # + # Construct the scheduler + # + tile_sched = utils.ClcDynamicPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + # + # Persistent tile scheduling loop + # + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestK) + tAgA_slice = tAgA[ + (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2]) + ] + # ((atom_v, rest_v), RestK) + tBgB_slice = tBgB[ + (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2]) + ] + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer empty + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + + # TMA load A/B + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + mcast_mask=b_full_mcast_mask, + ) + + # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + # + # Wait A/B buffer empty + # + ab_producer.tail() + + # + # Sched warp + # + if warp_idx == self.sched_warp_id and is_first_cta_in_cluster: + # + # Persistent tile scheduling loop + # + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, self.num_clc_stage + ) + + while work_tile.is_valid_tile: + # + # Advance to next tile + # + clc_pipeline.producer_acquire(clc_producer_state) + mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) + tile_sched.advance_to_next_work(mbarrier_addr) + clc_producer_state.advance() + + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + clc_pipeline.producer_tail(clc_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop + # + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + # + # Mma mainloop + # + for k_tile in range(k_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + kblk_crd = (None, None, kblk_idx, handle.index) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblk_crd], + tCrB[kblk_crd], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + handle.release() + + # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + # + # Async arrive accumulator buffer full + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + sC = None + if cutlass.const_expr(self.use_tma_store): + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Persistent tile scheduling loop for epilogue + # + if cutlass.const_expr(self.use_tma_store): + assert tma_atom_c is not None and sC is not None + utils.gemm.sm100.epilogue_tma_store( + self, + tidx, + warp_idx, + acc_pipeline, + tiled_mma, + tma_atom_c, + tCtAcc_base, + sC, + tCgC, + epi_tile, + tile_sched, + epilogue_op, + clc_pipeline, + clc_consumer_state, + ) + else: + utils.gemm.sm100.epilogue( + self, + tidx, + acc_pipeline, + tiled_mma, + tCtAcc_base, + tCgC, + epi_tile, + tile_sched, + epilogue_op, + tmem_dealloc_barrier, + None, + None, + clc_pipeline, + clc_consumer_state, + ) + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + ) -> Tuple[utils.ClcDynamicPersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.ClcDynamicPersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params) + + return tile_sched_params, grid + + @staticmethod + def _compute_num_tmem_alloc_cols( + tiled_mma: cute.TiledMma, + mma_tiler: Tuple[int, int, int], + num_acc_stage: int, + arch: str, + ) -> int: + """ + Compute the number of tensor memory allocation columns. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler: The shape (M, N, K) of the MMA tile. + :type mma_tiler: tuple[int, int, int] + :param num_acc_stage: The stage of the accumulator tensor. + :type num_acc_stage: int + + :return: The number of tensor memory allocation columns. + :rtype: int + """ + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, num_acc_stage)) + num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch=arch) + + return num_tmem_alloc_cols + + def check_supported_dtypes( + self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric] + ) -> bool: + """ + Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator + :type acc_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :raises testing.CantImplementError: If the dtypes are invalid + """ + valid_ab_dtypes = { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Uint8, + cutlass.Int8, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + } + if ab_dtype not in valid_ab_dtypes: + raise testing.CantImplementError(f"Unsupported AB dtype: {ab_dtype}") + + if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}: + raise testing.CantImplementError( + f"Unsupported accumulator dtype: {self.acc_dtype}" + ) + + # Define compatibility mapping between accumulator type and AB type + acc_ab_compatibility = { + cutlass.Float32: { + cutlass.Float16, + cutlass.BFloat16, + cutlass.TFloat32, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, # Float32 accumulator supports floating point AB types only + cutlass.Float16: { + cutlass.Float16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }, + cutlass.Int32: {cutlass.Uint8, cutlass.Int8}, + } + # Check compatibility between accumulator type and AB type + if ab_dtype not in acc_ab_compatibility[self.acc_dtype]: + return False + + # Define compatibility mapping between accumulator type and C type + acc_c_compatibility = { + cutlass.Float32: { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + cutlass.Float16: { + cutlass.BFloat16, + cutlass.Float16, + }, + cutlass.Int32: { + cutlass.BFloat16, + cutlass.Float16, + cutlass.Float32, + cutlass.Int32, + cutlass.Int8, + cutlass.Uint8, + }, + } + # Check compatibility between accumulator type and C type + if c_dtype not in acc_c_compatibility[self.acc_dtype]: + return False + + return True + + def check_mma_tiler_and_cluster_shape(self) -> bool: + """Check if the mma tiler and cluster shape are valid. + + :raises testing.CantImplementError: If the mma tiler and cluster shape are invalid + """ + # Skip invalid mma tile shape + if not ( + (not self.use_2cta_instrs and self.mma_tiler_mn[0] in [64, 128]) + or (self.use_2cta_instrs and self.mma_tiler_mn[0] in [128, 256]) + ): + raise testing.CantImplementError( + f"Invalid mma tiler & use_2cta_instrs: {self.mma_tiler_mn}, {self.use_2cta_instrs}" + ) + if self.mma_tiler_mn[1] not in range(32, 257, 32): + raise testing.CantImplementError( + f"Invalid mma tiler N: {self.mma_tiler_mn[1]}" + ) + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + raise testing.CantImplementError( + f"Invalid cluster shape M: {self.cluster_shape_mn[0]}" + ) + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + raise testing.CantImplementError( + f"Invalid cluster shape: {self.cluster_shape_mn}" + ) + + def check_tensor_alignment( + self, + m: int, + n: int, + k: int, + l: int, + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + :param k: The number of columns in the A tensor + :type k: int + :param l: The number of columns in the C tensor + :type l: int + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + + # TODO: move to utils + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + raise testing.CantImplementError( + f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {ab_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}" + ) + + def check_epilog_store_option(self, m: int, n: int) -> bool: + """ + Check if the epilogue store option is valid + + :param m: The number of rows in the A tensor + :type m: int + :param n: The number of columns in the B tensor + :type n: int + + :raises testing.CantImplementError: If the epilogue store option is invalid + """ + # None TMA store version does not have predication, can not support OOB tiles + cta_tile_shape_mn = ( + self.mma_tiler_mn[0] // (2 if self.use_2cta_instrs else 1), + self.mma_tiler_mn[1], + ) + if not self.use_tma_store: + if not (m % cta_tile_shape_mn[0] == 0 and n % cta_tile_shape_mn[1] == 0): + raise testing.CantImplementError( + f"Invalid epilog store option: {m}, {n}" + ) + + def can_implement( + self, + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Determine if the given tensor configuration can be implemented by this kernel. + + :param mnkl: Problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B. + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C. + :type c_dtype: Type[cutlass.Numeric] + :param a_major: Major dimension of the A tensor layout ("m" or "k"). + :type a_major: str + :param b_major: Major dimension of the B tensor layout ("n" or "k"). + :type b_major: str + :param c_major: Major dimension of the C tensor layout ("m" or "n"). + :type c_major: str + :return: True if the kernel supports the given configuration, False otherwise. + :rtype: bool + """ + + try: + # Skip unsupported types + self.check_supported_dtypes(ab_dtype, c_dtype) + + # Skip invalid mma tile shape and cluster shape + self.check_mma_tiler_and_cluster_shape() + + m, n, k, l = mnkl + self.check_tensor_alignment( + m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major + ) + self.check_epilog_store_option(m, n) + except testing.CantImplementError: + return False + return True + + +@cute.jit +def bmm( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, # (l, m, k) + b: cute.Tensor, # (l, k, n) + c: cute.Tensor, # (l, m, n) + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + """ + Wrapper API for persistent GEMM kernel to follow the convention of PyTorch's batch matrix-multiply (bmm). + + Internally, the tensors are permuted to match CuTe's convention: + - a: (m, k, l) + - b: (n, k, l) + - c: (m, n, l) + + :param gemm_op: Kernel operation, expects (a, b, c, max_active_clusters, stream, epilogue_op) + :type gemm_op: cutlass.Constexpr + :param a: Input tensor of shape (l, m, k) + :type a: cute.Tensor + :param b: Input tensor of shape (l, k, n) + :type b: cute.Tensor + :param c: Output tensor of shape (l, m, n) + :type c: cute.Tensor + :param max_active_clusters: Maximum number of hardware clusters to launch + :type max_active_clusters: cutlass.Constexpr + :param epilogue_op: Optional elementwise lambda function to apply per output element, defaults to identity + :type epilogue_op: cutlass.Constexpr, optional + """ + # (l,m,k) -> (m,k,l) + a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])) + # (l,k,n) -> (n,k,l) + b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + # (l,m,n) -> (m,n,l) + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + + gemm_op(a, b, c, max_active_clusters, stream, epilogue_op) + + +@lru_cache(maxsize=1) +def prepare_tensors( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + init_random: bool = True, +): + import torch + from cutlass.torch import dtype as torch_dtype + + m, n, k, l = mnkl + + if a_major == "k": + a = torch.empty((l, m, k), dtype=torch.float32, device="cuda") + elif a_major == "m": + a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + + if b_major == "n": + b = torch.empty((l, k, n), dtype=torch.float32, device="cuda") + elif b_major == "k": + b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1) + + if c_major == "n": + c = torch.empty((l, m, n), dtype=torch.float32, device="cuda") + elif c_major == "m": + c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1) + + if init_random: + a.random_(-2, 3) + b.random_(-2, 3) + c.random_(-2, 3) + + return ( + a.to(dtype=torch_dtype(ab_dtype)), + b.to(dtype=torch_dtype(ab_dtype)), + c.to(dtype=torch_dtype(c_dtype)), + ) + + +@lru_cache(maxsize=1) +def compile_bmm( + mnkl: Tuple[int, int, int, int], + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + max_active_clusters: cutlass.Constexpr = None, + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + epilogue_op: cutlass.Constexpr = lambda x: x, +): + from cutlass.cute.runtime import make_fake_stream + + gemm = PersistentDenseGemmKernel( + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + use_tma_store, + ) + # Check if configuration can be implemented + can_implement = gemm.can_implement( + mnkl, a.element_type, c.element_type, a_major, b_major, c_major + ) + if not can_implement: + raise testing.CantImplementError( + f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, " + f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, " + f"use_tma_store = {use_tma_store}" + ) + + stream = make_fake_stream() + return cute.compile(bmm, gemm, a, b, c, max_active_clusters, stream, epilogue_op) + + +def run( + mnkl: Tuple[int, int, int, int], + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Tuple[int, int] = (2, 1), + use_2cta_instrs: bool = True, + use_tma_store: bool = True, + tolerance: float = 1e-01, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + benchmark: bool = False, + **kwargs, +): + """ + Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking. + + Prepares input tensors, configures and launches the persistent GEMM kernel, + optionally performs reference validation, and benchmarks execution. + + :param mnkl: Problem size as a tuple (M, N, K, L). + :type mnkl: Tuple[int, int, int, int] + :param ab_dtype: Data type for input tensors A and B. + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: Data type for output tensor C. + :type c_dtype: Type[cutlass.Numeric] + :param acc_dtype: Accumulator data type for the matrix multiplication. + :type acc_dtype: Type[cutlass.Numeric] + :param a_major: Memory layout of tensor A. + :type a_major: str + :param b_major: Memory layout of tensor B. + :type b_major: str + :param c_major: Memory layout of tensor C. + :type c_major: str + :param mma_tiler_mn: MMA tiling size (M, N), defaults to (256, 256). + :type mma_tiler_mn: Tuple[int, int], optional + :param cluster_shape_mn: Cluster shape (M, N), defaults to (2, 1). + :type cluster_shape_mn: Tuple[int, int], optional + :param use_2cta_instrs: Whether to use 2CTA MMA instructions, defaults to True. + :type use_2cta_instrs: bool, optional + :param use_tma_store: Whether to use TMA store, defaults to True. + :type use_tma_store: bool, optional + :param tolerance: Tolerance for reference validation, defaults to 1e-01. + :type tolerance: float, optional + :param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0. + :type warmup_iterations: int, optional + :param iterations: Number of benchmark iterations to run, defaults to 1. + :type iterations: int, optional + :param skip_ref_check: Whether to skip reference result validation, defaults to False. + :type skip_ref_check: bool, optional + :param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False. + :type use_cold_l2: bool, optional + :param benchmark: Whether to only benchmark the kernel, defaults to False. + :type benchmark: bool, optional + :raises RuntimeError: If CUDA GPU is not available. + :raises ValueError: If the configuration is invalid or unsupported by the kernel. + :return: Execution time of the GEMM kernel. + :rtype: float + """ + import torch + from cutlass.torch import dtype as torch_dtype + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # Get current CUDA stream from PyTorch + torch_stream = torch.cuda.current_stream() + # Get the raw stream pointer as a CUstream + current_stream = cuda.CUstream(torch_stream.cuda_stream) + + # Check if configuration can be implemented + max_active_clusters = utils.HardwareInfo().get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + # Run and verify BMM with torch + a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major) + + leading_dim_a = 2 if a_major == "k" else 1 + leading_dim_b = 1 if b_major == "k" else 2 + leading_dim_c = 2 if c_major == "n" else 1 + + a_ = from_dlpack( + a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 + ).mark_layout_dynamic(leading_dim=leading_dim_a) + b_ = from_dlpack( + b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 + ).mark_layout_dynamic(leading_dim=leading_dim_b) + c_ = from_dlpack( + c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32 + ).mark_layout_dynamic(leading_dim=leading_dim_c) + + compiled_fn = compile_bmm( + mnkl, + a_, + b_, + c_, + acc_dtype, + a_major, + b_major, + c_major, + mma_tiler_mn, + cluster_shape_mn, + max_active_clusters, + use_2cta_instrs, + use_tma_store, + epilogue_op=lambda x: x, + ) + + print("Running Blackwell Persistent Dense GEMM test with:") + print(f"mnkl: {mnkl}") + print(f"Tolerance: {tolerance}") + print(f"Warmup iterations: {warmup_iterations}") + print(f"Iterations: {iterations}") + print(f"Skip reference checking: {skip_ref_check}") + print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}") + + if not skip_ref_check: + # Use small random number for deterministic result for reference check + compiled_fn(a_, b_, c_, current_stream) + + # Manually quantize to be comparable + ref = ( + torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32)) + .to(dtype=torch_dtype(c_dtype)) + .to(dtype=torch.float32) + ) + torch.testing.assert_close( + c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03 + ) + + if not benchmark: + return 0 + + def generate_tensors(): + init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8] + a, b, c = prepare_tensors( + mnkl, + ab_dtype, + c_dtype, + a_major, + b_major, + c_major, + init_random=not init_normal, + ) + a_ = from_dlpack( + a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 + ).mark_layout_dynamic(leading_dim=leading_dim_a) + b_ = from_dlpack( + b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32 + ).mark_layout_dynamic(leading_dim=leading_dim_b) + c_ = from_dlpack( + c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32 + ).mark_layout_dynamic(leading_dim=leading_dim_c) + return testing.JitArguments(a_, b_, c_, current_stream) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a.numel() * a.element_size() + + b.numel() * b.element_size() + + c.numel() * c.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + # Return execution time in microseconds + return testing.benchmark( + compiled_fn, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + + +def compute_tflops(time_ns, m, n, k): + return 2.0 * m * n * k / time_ns / 1000.0 + + + +def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + +def prepare_parser(): + parser = argparse.ArgumentParser( + description="Example of Dense Persistent GEMM on Blackwell." + ) + + parser.add_argument( + "--mnkl", + type=_parse_comma_separated_ints, + default=(256, 256, 512, 1), + help="mnkl dimensions (comma-separated)", + ) + parser.add_argument( + "--cluster_shape_mn", + type=_parse_comma_separated_ints, + default=(1, 1), + help="Cluster shape (comma-separated)", + ) + parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.TFloat32) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32) + parser.add_argument( + "--use_2cta_instrs", + action="store_true", + help="Enable 2CTA MMA instructions feature", + ) + parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k") + parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument( + "--use_tma_store", action="store_true", help="Use tma store or not" + ) + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + parser.add_argument( + "--benchmark", + type=str, + default="default", + choices=[ + "default", + "none", + ], + help="Benchmark the kernel with nsight or default (cute.testing.benchmark) or none", + ) + parser.add_argument( + "--warmup_iterations", type=int, default=0, help="Warmup iterations" + ) + parser.add_argument( + "--iterations", + type=int, + default=1, + help="Number of iterations to run the kernel", + ) + parser.add_argument( + "--skip_ref_check", action="store_true", help="Skip reference checking" + ) + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + return parser + + +if __name__ == "__main__": + parser = prepare_parser() + parser.add_argument( + "--mma_tiler_mn", + type=_parse_comma_separated_ints, + default=(128, 128), + help="Mma tile shape (comma-separated)", + ) + + args = parser.parse_args() + + if len(args.mnkl) != 4: + parser.error("--mnkl must contain exactly 4 values") + + if len(args.mma_tiler_mn) != 2: + parser.error("--mma_tiler_mn must contain exactly 2 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + print(f"[DSL INFO] Compiling Blackwell Persistent Dense GEMM with:") + print( + f"[DSL INFO] A dtype: {args.ab_dtype}, B dtype: {args.c_dtype}, C dtype: {args.acc_dtype}, Acc dtype: {args.acc_dtype}" + ) + print( + f"[DSL INFO] Matrix majors - A: {args.a_major}, B: {args.b_major}, C: {args.c_major}" + ) + print(f"[DSL INFO] Mma Tiler (M, N): {args.mma_tiler_mn}") + print(f"[DSL INFO] Cluster Shape (M, N): {args.cluster_shape_mn}") + print( + f"[DSL INFO] 2CTA MMA instructions: {'True' if args.use_2cta_instrs else 'False'}" + ) + print(f"[DSL INFO] Use TMA Store: {'True' if args.use_tma_store else 'False'}") + + run( + args.mnkl, + args.ab_dtype, + args.c_dtype, + args.acc_dtype, + args.a_major, + args.b_major, + args.c_major, + args.mma_tiler_mn, + args.cluster_shape_mn, + args.use_2cta_instrs, + args.use_tma_store, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + args.benchmark == "default", + ) + print("PASS") diff --git a/examples/python/CuTeDSL/cute/export/export_to_c.py b/examples/python/CuTeDSL/cute/export/export_to_c.py new file mode 100644 index 00000000..61ec79b9 --- /dev/null +++ b/examples/python/CuTeDSL/cute/export/export_to_c.py @@ -0,0 +1,107 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +import cutlass +import cutlass.cute as cute + +import os +import subprocess +import cuda.bindings.driver as cuda + +"""Example demonstrating how to export a CuTe function. + +This example shows how to: +1. Compile a CuTe function +2. Export the compiled function to a object file and a C header file +3. Compile the object file to a shared library + +To run this example: + +.. code-block:: bash + + # export the compiled function to a object file and a C header file + python cutlass_ir/compiler/python/examples/cute/export/export_to_c.py +""" + + +@cute.kernel +def print_tensor_kernel(a: cute.Tensor): + cute.printf("a: {}", a) + + +@cute.jit +def print_tensor(a: cute.Tensor, stream: cuda.CUstream): + print_tensor_kernel(a).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream) + + +@cute.kernel +def add_one_kernel(a: cute.Tensor, b: cute.Tensor): + a[0] = b[0] + 1 + + +@cute.jit +def add_one(a: cute.Tensor, b: cute.Tensor, stream: cuda.CUstream): + add_one_kernel(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream) + + +def run(): + from cutlass.cute.runtime import make_fake_compact_tensor + + shape = (cute.SymInt(divisibility=16), cute.SymInt()) + a = make_fake_compact_tensor(cutlass.Float32, shape, stride_order=(1, 0)) + b = make_fake_compact_tensor(cutlass.Float32, shape, stride_order=(1, 0)) + stream = cuda.CUstream(0) + compiled_print_tensor = cute.compile(print_tensor, a, stream=stream) + compiled_add_one = cute.compile(add_one, a, b, stream=stream) + + os.makedirs("./build", exist_ok=True) + compiled_print_tensor.export_to_c( + file_path="./build", + file_name="print_tensor_example", + function_prefix="print_tensor", + ) + compiled_add_one.export_to_c( + file_path="./build", file_name="add_one_example", function_prefix="add_one" + ) + + cc = os.environ.get("CC", "gcc") + + # compile the object file to a shared library + cmd = [ + cc, + "-shared", + "-o", + "./build/libexport_example.so", + "./build/print_tensor_example.o", + "./build/add_one_example.o", + ] + subprocess.run(cmd, check=True) + + +if __name__ == "__main__": + run() diff --git a/examples/python/CuTeDSL/cute/export/load_in_python.py b/examples/python/CuTeDSL/cute/export/load_in_python.py new file mode 100644 index 00000000..758fe228 --- /dev/null +++ b/examples/python/CuTeDSL/cute/export/load_in_python.py @@ -0,0 +1,79 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +import cutlass +import cutlass.cute as cute +import cuda.bindings.driver as cuda + + +"""Example demonstrating how to load a CuTe module/function in Python. + +This example shows how to: +1. Load a CuTe module from a object file or a shared library +2. Extract the function from the module and call it in Python + +To run this example: + +.. code-block:: bash + + # prerequesites: export the compiled functions to object files and compile them into a shared library + python examples/cute/export/export_to_c.py + # load the module from a object file or a shared library + python examples/cute/export/load_in_python.py +""" + + +def run(): + import torch + from cutlass.cute.runtime import from_dlpack + + a = ( + torch.arange(16 * 10, dtype=torch.float32, device="cuda") + .reshape(16, 10) + .permute(1, 0) + ) + b = torch.ones(16, 10, dtype=torch.float32, device="cuda").permute(1, 0) + a_cute = from_dlpack(a).mark_layout_dynamic() + b_cute = from_dlpack(b).mark_layout_dynamic() + stream = cuda.CUstream(0) + + print_tensor_mod = cute.runtime.load_module("./build/print_tensor_example.o") + print_tensor_mod.print_tensor(a_cute, stream=stream) + + add_one_mod = cute.runtime.load_module("./build/add_one_example.o") + add_one_mod.add_one(a_cute, b_cute, stream=stream) + assert a[0, 0] == b[0, 0] + 1 + + shared_mod = cute.runtime.load_module("./build/libexport_example.so") + shared_mod.print_tensor(a_cute, stream=stream) + shared_mod.add_one(a_cute, b_cute, stream=stream) + assert a[0, 0] == b[0, 0] + 1 + + +if __name__ == "__main__": + run() diff --git a/examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.cpp b/examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.cpp new file mode 100644 index 00000000..7c58d407 --- /dev/null +++ b/examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.cpp @@ -0,0 +1,254 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary + * + * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual + * property and proprietary rights in and to this material, related + * documentation and any modifications thereto. Any use, reproduction, + * disclosure or distribution of this material and related documentation + * without an express license agreement from NVIDIA CORPORATION or + * its affiliates is strictly prohibited. + */ +/** + * This example demonstrates how to load a CuTe module/function from a object + * file or a shared library and run it in a CUDA kernel. + * + * To run this example: + * + * .. code-block:: bash + * + * # prerequesites: export the compiled functions to object files and compile + * them into a shared library python examples/cute/export/export_to_c.py # run + * the example bash ./examples/cute/export/run_with_dynamic_loading.sh + */ + +#include "CuteDSLRuntime.h" +#include +#include +#include + +std::vector read_file(const std::string &filename) { + std::ifstream file(filename, std::ios::binary); + std::vector content((std::istreambuf_iterator(file)), + std::istreambuf_iterator()); + return content; +} + +void initialize_cuda_context() { + // Initialize cuda context + cudaSetDevice(0); +} + +void check_error(CuteDSLRT_Error_t error) { + if (error != CuteDSLRT_Error_Success) { + printf("Got runtime error: %d\n", error); + exit(1); + } +} + +// Copy the definition of the tensor from `print_tensor_example.h` to here +typedef struct { + void *data; + int32_t dynamic_shapes[2]; + int64_t dynamic_strides[1]; +} print_tensor_Tensor_a_t; + +// Copy the definition of the tensor from `add_one_example.h` to here +typedef struct { + void *data; + int32_t dynamic_shapes[2]; + int64_t dynamic_strides[1]; +} add_one_Tensor_a_t; +typedef struct { + void *data; + int32_t dynamic_shapes[2]; + int64_t dynamic_strides[1]; +} add_one_Tensor_b_t; + +print_tensor_Tensor_a_t prepare_print_tensor_tensor() { + print_tensor_Tensor_a_t tensor; + tensor.data = NULL; + tensor.dynamic_shapes[0] = 32; + tensor.dynamic_shapes[1] = 16; + tensor.dynamic_strides[0] = 16; + return tensor; +} + +add_one_Tensor_a_t prepare_add_one_tensor_a() { + float a_host_ptr[32 * 16]; + for (int i = 0; i < 32 * 16; i++) { + a_host_ptr[i] = i; + } + void *a_device_ptr; + cudaMalloc(&a_device_ptr, sizeof(float) * 32 * 16); + cudaMemcpy(a_device_ptr, a_host_ptr, sizeof(float) * 32 * 16, + cudaMemcpyHostToDevice); + add_one_Tensor_a_t tensor; + tensor.data = (void *)a_device_ptr; + tensor.dynamic_shapes[0] = 32; + tensor.dynamic_shapes[1] = 16; + tensor.dynamic_strides[0] = 16; + return tensor; +} + +add_one_Tensor_b_t prepare_add_one_tensor_b() { + float b_host_ptr[32 * 16]; + for (int i = 0; i < 32 * 16; i++) { + b_host_ptr[i] = 32 * 16 - i; + } + void *b_device_ptr; + cudaMalloc(&b_device_ptr, sizeof(float) * 32 * 16); + cudaMemcpy(b_device_ptr, b_host_ptr, sizeof(float) * 32 * 16, + cudaMemcpyHostToDevice); + add_one_Tensor_b_t tensor; + tensor.data = (void *)b_device_ptr; + tensor.dynamic_shapes[0] = 32; + tensor.dynamic_shapes[1] = 16; + tensor.dynamic_strides[0] = 16; + return tensor; +} + +void run_print_tensor_with_object_file() { + + // prepare tensor + print_tensor_Tensor_a_t tensor = prepare_print_tensor_tensor(); + + // prepare stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // load module and function + CuteDSLRT_Module_t *module = nullptr; + std::vector print_tensor_example_o_bytes = + read_file("build/print_tensor_example.o"); + size_t print_tensor_example_o_size = print_tensor_example_o_bytes.size(); + CuteDSLRT_Error_t error = CuteDSLRT_Module_Create_From_Bytes( + &module, print_tensor_example_o_bytes.data(), print_tensor_example_o_size, + nullptr, 0); + check_error(error); + + CuteDSLRT_Function_t *function = nullptr; + error = CuteDSLRT_Module_Get_Function(&function, module, "print_tensor"); + check_error(error); + + // run kernel, refer to the wrapper function in `print_tensor_example.h` + int32_t cuda_result; + void *args[3] = {&tensor, &stream, &cuda_result}; + error = CuteDSLRT_Function_Run(function, args, 3); + check_error(error); + + // synchronize stream + cudaStreamSynchronize(stream); + + // unload module + error = CuteDSLRT_Module_Destroy(module); + check_error(error); + cudaStreamDestroy(stream); +} + +void run_add_one_with_object_file() { + // prepare a tensor + add_one_Tensor_a_t tensor_a = prepare_add_one_tensor_a(); + + // prepare b tensor + add_one_Tensor_b_t tensor_b = prepare_add_one_tensor_b(); + + // prepare stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // load module + CuteDSLRT_Module_t *module = nullptr; + std::vector add_one_example_o_bytes = + read_file("build/add_one_example.o"); + size_t add_one_example_o_size = add_one_example_o_bytes.size(); + CuteDSLRT_Error_t error = CuteDSLRT_Module_Create_From_Bytes( + &module, add_one_example_o_bytes.data(), add_one_example_o_size, nullptr, + 0); + check_error(error); + + CuteDSLRT_Function_t *function = nullptr; + error = CuteDSLRT_Module_Get_Function(&function, module, "add_one"); + check_error(error); + + // run kernel, refer to the wrapper function in `add_one_example.h` + int32_t cuda_result; + void *args[4] = {&tensor_a, &tensor_b, &stream, &cuda_result}; + error = CuteDSLRT_Function_Run(function, args, 4); + check_error(error); + cudaStreamSynchronize(stream); + + // unload module + error = CuteDSLRT_Module_Destroy(module); + check_error(error); + cudaStreamDestroy(stream); + + // check result + float a_host_ptr[32 * 16]; + cudaMemcpy(a_host_ptr, tensor_a.data, sizeof(float) * 32 * 16, + cudaMemcpyDeviceToHost); + + if (a_host_ptr[0] != 32 * 16 + 1) { + printf("Error: a_host_ptr[0] = %f, expected %d\n", a_host_ptr[0], 32 * 16); + exit(1); + } +} + +void run_example_with_shared_library() { + // prepare tensor + print_tensor_Tensor_a_t tensor = prepare_print_tensor_tensor(); + + // prepare a tensor + add_one_Tensor_a_t tensor_a = prepare_add_one_tensor_a(); + + // prepare b tensor + add_one_Tensor_b_t tensor_b = prepare_add_one_tensor_b(); + + // prepare stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // load module + CuteDSLRT_Module_t *module = nullptr; + const char *shared_libs[] = {"build/libexport_example.so"}; + CuteDSLRT_Error_t error = + CuteDSLRT_Module_Create_From_Bytes(&module, nullptr, 0, shared_libs, 1); + check_error(error); + + // get print tensor function + CuteDSLRT_Function_t *print_tensor_function = nullptr; + error = CuteDSLRT_Module_Get_Function(&print_tensor_function, module, + "print_tensor"); + check_error(error); + + // get add one function + CuteDSLRT_Function_t *add_one_function = nullptr; + error = CuteDSLRT_Module_Get_Function(&add_one_function, module, "add_one"); + check_error(error); + + // run print tensor kernel + int32_t cuda_result; + void *print_tensor_args[3] = {&tensor, &stream, &cuda_result}; + error = CuteDSLRT_Function_Run(print_tensor_function, print_tensor_args, 3); + check_error(error); + cudaStreamSynchronize(stream); + + // run add one kernel + void *add_one_args[4] = {&tensor_a, &tensor_b, &stream, &cuda_result}; + error = CuteDSLRT_Function_Run(add_one_function, add_one_args, 4); + check_error(error); + cudaStreamSynchronize(stream); + + // unload module + error = CuteDSLRT_Module_Destroy(module); + check_error(error); + cudaStreamDestroy(stream); +} + +int main() { + initialize_cuda_context(); + run_print_tensor_with_object_file(); + run_add_one_with_object_file(); + run_example_with_shared_library(); + return 0; +} diff --git a/examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.sh b/examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.sh new file mode 100755 index 00000000..571d06b3 --- /dev/null +++ b/examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.sh @@ -0,0 +1,80 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +#!/bin/bash +set -eu + +# Try to find the wheel path of nvidia-cutlass-dsl +WHEEL_PATH=$(python3 -c "import cutlass, os; print(os.path.dirname(cutlass.__file__))" 2>/dev/null)/../.. + +if [[ -z "$WHEEL_PATH" ]]; then + echo "nvidia-cutlass-dsl wheel not found in the current Python environment." + exit 1 +else + echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH" +fi + +CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/" +export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH}:${CUDA_HOME}/lib64:./build + +if [ -z "$CUDA_HOME" ]; then + CUDA_HOME=/usr/local/cuda + echo "CUDA_HOME not set, using default: $CUDA_HOME" +else + echo "CUDA_HOME found: $CUDA_HOME" +fi +SOURCE_FILE="$(dirname "$0")/run_with_dynamic_loading.cpp" + +echo "Compiling the executable..." +# Search for a common C++ compiler: g++, clang++, or c++ +if [ -n "${CXX-}" ] && command -v "$CXX" &> /dev/null; then + CXX="$CXX" +elif command -v g++ &> /dev/null; then + CXX="g++" +elif command -v clang++ &> /dev/null; then + CXX="clang++" +elif command -v c++ &> /dev/null; then + CXX="c++" +else + echo "Error: No common C++ compiler found (g++, clang++, or c++). Please install a C++ compiler to continue." + exit 1 +fi + +$CXX -o build/run_with_dynamic_loading \ + -I${CUDA_HOME}/include \ + -I${WHEEL_PATH}/include \ + ${SOURCE_FILE} \ + -L${CUTE_DSL_LIB_PATH} \ + -L${CUDA_HOME}/lib64 \ + -L./build \ + -lcudart \ + -lcute_dsl_runtime \ + -lexport_example + +echo "Running the executable..." +./build/run_with_dynamic_loading diff --git a/examples/python/CuTeDSL/cute/export/run_with_static_linking.cpp b/examples/python/CuTeDSL/cute/export/run_with_static_linking.cpp new file mode 100644 index 00000000..96acbc2e --- /dev/null +++ b/examples/python/CuTeDSL/cute/export/run_with_static_linking.cpp @@ -0,0 +1,124 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary + * + * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual + * property and proprietary rights in and to this material, related + * documentation and any modifications thereto. Any use, reproduction, + * disclosure or distribution of this material and related documentation + * without an express license agreement from NVIDIA CORPORATION or + * its affiliates is strictly prohibited. + */ +/** + * This example demonstrates how to load a CuTe module/function from compilation + * and static linking and run it in a CUDA kernel. + * + * To run this example: + * + * .. code-block:: bash + * + * # prerequesites: export the compiled functions to object files and compile + * them into a shared library python examples/cute/export/export_to_c.py # run + * the example bash ./examples/cute/export/run_with_static_linking.sh + */ + +#include "add_one_example.h" +#include "print_tensor_example.h" +#include + +void initialize_cuda_context() { + // Initialize cuda context + int device_id = 0; + cudaSetDevice(device_id); +} + +void run_print_tensor() { + // prepare tensor + print_tensor_Tensor_a_t tensor; + tensor.data = NULL; + tensor.dynamic_shapes[0] = 32; + tensor.dynamic_shapes[1] = 16; + tensor.dynamic_strides[0] = 16; + + // prepare stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // load module + print_tensor_Kernel_Module_t module; + print_tensor_Kernel_Module_Load(&module); + + // run kernel + cute_dsl_print_tensor_wrapper(&module, &tensor, stream); + + // synchronize stream + cudaStreamSynchronize(stream); + + // unload module + print_tensor_Kernel_Module_Unload(&module); + cudaStreamDestroy(stream); +} + +void run_add_one() { + // prepare a tensor + add_one_Tensor_a_t a_tensor; + float a_host_ptr[32 * 16]; + for (int i = 0; i < 32 * 16; i++) { + a_host_ptr[i] = i; + } + void *a_device_ptr; + cudaMalloc(&a_device_ptr, sizeof(float) * 32 * 16); + cudaMemcpy(a_device_ptr, a_host_ptr, sizeof(float) * 32 * 16, + cudaMemcpyHostToDevice); + a_tensor.data = (void *)a_device_ptr; + a_tensor.dynamic_shapes[0] = 32; + a_tensor.dynamic_shapes[1] = 16; + a_tensor.dynamic_strides[0] = 16; + + // prepare b tensor + add_one_Tensor_b_t b_tensor; + float b_host_ptr[32 * 16]; + for (int i = 0; i < 32 * 16; i++) { + b_host_ptr[i] = 32 * 16 - i; + } + void *b_device_ptr; + cudaMalloc(&b_device_ptr, sizeof(float) * 32 * 16); + cudaMemcpy(b_device_ptr, b_host_ptr, sizeof(float) * 32 * 16, + cudaMemcpyHostToDevice); + + b_tensor.data = (void *)b_device_ptr; + b_tensor.dynamic_shapes[0] = 32; + b_tensor.dynamic_shapes[1] = 16; + b_tensor.dynamic_strides[0] = 16; + + // prepare stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // load module + add_one_Kernel_Module_t module; + add_one_Kernel_Module_Load(&module); + + // run kernel + cute_dsl_add_one_wrapper(&module, &a_tensor, &b_tensor, stream); + cudaStreamSynchronize(stream); + + // unload module + add_one_Kernel_Module_Unload(&module); + cudaStreamDestroy(stream); + + // check result + cudaMemcpy(a_host_ptr, a_device_ptr, sizeof(float) * 32 * 16, + cudaMemcpyDeviceToHost); + + if (a_host_ptr[0] != 32 * 16 + 1) { + printf("Error: a_host_ptr[0] = %f, expected %d\n", a_host_ptr[0], 32 * 16); + } +} + +int main() { + initialize_cuda_context(); + run_print_tensor(); + run_add_one(); + return 0; +} diff --git a/examples/python/CuTeDSL/cute/export/run_with_static_linking.sh b/examples/python/CuTeDSL/cute/export/run_with_static_linking.sh new file mode 100755 index 00000000..c3540d73 --- /dev/null +++ b/examples/python/CuTeDSL/cute/export/run_with_static_linking.sh @@ -0,0 +1,76 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +#!/bin/bash +set -eu + +# Try to find the wheel path of nvidia-cutlass-dsl +WHEEL_PATH=$(python3 -c "import cutlass, os; print(os.path.dirname(cutlass.__file__))" 2>/dev/null)/../.. + +if [[ -z "$WHEEL_PATH" ]]; then + echo "nvidia-cutlass-dsl wheel not found in the current Python environment." + exit 1 +else + echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH" +fi +CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/" +export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH} + +if [ -z "$CUDA_HOME" ]; then + CUDA_HOME=/usr/local/cuda + echo "CUDA_HOME not set, using default: $CUDA_HOME" +else + echo "CUDA_HOME found: $CUDA_HOME" +fi +SOURCE_FILE="$(dirname "$0")/run_with_static_linking.cpp" + +echo "Compiling the executable..." +# Search for a common C++ compiler: g++, clang++, or c++ +if [ -n "${CXX-}" ] && command -v "$CXX" &> /dev/null; then + CXX="$CXX" +elif command -v g++ &> /dev/null; then + CXX="g++" +elif command -v clang++ &> /dev/null; then + CXX="clang++" +elif command -v c++ &> /dev/null; then + CXX="c++" +else + echo "Error: No common C++ compiler found (g++, clang++, or c++). Please install a C++ compiler to continue." + exit 1 +fi + +# Note: The options -ldl, -lrt, and -lpthread are optional to satisfy symbol dependencies required by `libcudart_static.a`. +$CXX -o build/run_with_static_linking \ + -I${CUDA_HOME}/include \ + -I./build \ + ${SOURCE_FILE} build/add_one_example.o build/print_tensor_example.o \ + ${CUTE_DSL_LIB_PATH}/libcuda_dialect_runtime_static.a \ + ${CUDA_HOME}/lib64/libcudart_static.a -ldl -lrt -lpthread + +echo "Running the executable..." +./build/run_with_static_linking diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py index a553ada5..57fe9a43 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py @@ -37,21 +37,19 @@ To run this example: .. code-block:: bash - python examples/cute/tvm_ffi/aot_export.py + python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_export.py # run example to use in torch - python examples/cute/tvm_ffi/aot_use_in_torch.py + python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_torch.py # run example to use in jax - python examples/cute/tvm_ffi/aot_use_in_jax.py + python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_jax.py # run example to use in c++ bundle - bash examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh + bash cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh """ from pathlib import Path -import torch import os import subprocess import tvm_ffi -import torch import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack @@ -69,6 +67,8 @@ def add_one(a: cute.Tensor, b: cute.Tensor): def main(): + import torch + # compile the kernel with "--enable-tvm-ffi" option a_torch = torch.arange(10, dtype=torch.float32, device="cuda") b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp index 27883585..e1f0645b 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp @@ -15,32 +15,31 @@ // This example shows how to interface with an AOT compiled function in a C++ // bundle. to build and run the example, run the following command in project // root bash -// examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh +// cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh #include -#include #include #include #include +#include #include + namespace ffi = tvm::ffi; struct CUDANDAlloc { - void AllocData(DLTensor *tensor) { + void AllocData(DLTensor* tensor) { size_t data_size = ffi::GetDataSize(*tensor); - void *ptr = nullptr; + void* ptr = nullptr; cudaError_t err = cudaMalloc(&ptr, data_size); - TVM_FFI_ICHECK_EQ(err, cudaSuccess) - << "cudaMalloc failed: " << cudaGetErrorString(err); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMalloc failed: " << cudaGetErrorString(err); tensor->data = ptr; } - void FreeData(DLTensor *tensor) { + void FreeData(DLTensor* tensor) { if (tensor->data != nullptr) { cudaError_t err = cudaFree(tensor->data); - TVM_FFI_ICHECK_EQ(err, cudaSuccess) - << "cudaFree failed: " << cudaGetErrorString(err); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaFree failed: " << cudaGetErrorString(err); tensor->data = nullptr; } } @@ -51,47 +50,45 @@ inline ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) { } // symbol from the shared library -extern "C" int __tvm_ffi_add_one(void *, const TVMFFIAny *, int32_t, - TVMFFIAny *); +extern "C" int __tvm_ffi_add_one(void*, const TVMFFIAny*, int32_t, TVMFFIAny*); // Redirects into the exported function in object void CallAddOne(ffi::TensorView x, ffi::TensorView y) { - tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y); + tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y); } int main() { - DLDataType f32_dtype{kDLFloat, 32, 1}; - DLDevice cuda_device{kDLCUDA, 0}; + DLDataType f32_dtype{kDLFloat, 32, 1}; + DLDevice cuda_device{kDLCUDA, 0}; - constexpr int ARRAY_SIZE = 10; + constexpr int ARRAY_SIZE = 10; - ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); - ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); + ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); + ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device); - std::vector host_x(ARRAY_SIZE); - for (int i = 0; i < ARRAY_SIZE; ++i) { - host_x[i] = static_cast(i); - } + std::vector host_x(ARRAY_SIZE); + for (int i = 0; i < ARRAY_SIZE; ++i) { + host_x[i] = static_cast(i); + } - size_t nbytes = host_x.size() * sizeof(float); - cudaError_t err = - cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice); - TVM_FFI_ICHECK_EQ(err, cudaSuccess) - << "cudaMemcpy host to device failed: " << cudaGetErrorString(err); + size_t nbytes = host_x.size() * sizeof(float); + cudaError_t err = cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) + << "cudaMemcpy host to device failed: " << cudaGetErrorString(err); - // Call into the FFI function; tensors remain on device because they carry a - // kDLCUDA device tag. - CallAddOne(x, y); + // Call into the FFI function; tensors remain on device because they carry a + // kDLCUDA device tag. + CallAddOne(x, y); - std::vector host_y(host_x.size()); - err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost); - TVM_FFI_ICHECK_EQ(err, cudaSuccess) - << "cudaMemcpy device to host failed: " << cudaGetErrorString(err); + std::vector host_y(host_x.size()); + err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost); + TVM_FFI_ICHECK_EQ(err, cudaSuccess) + << "cudaMemcpy device to host failed: " << cudaGetErrorString(err); - std::cout << "y after add_one_cuda(x, y)" << std::endl; - for (float value : host_y) { - std::cout << value << " "; - } - std::cout << std::endl; - return 0; + std::cout << "y after add_one_cuda(x, y)" << std::endl; + for (float value : host_y) { + std::cout << value << " "; + } + std::cout << std::endl; + return 0; } diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh index dea07c95..3f110a4a 100755 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh @@ -27,8 +27,8 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #!/bin/bash -CUDA_DIALECT_PATH="build/lib/" -export LD_LIBRARY_PATH=${CUDA_DIALECT_PATH}:`tvm-ffi-config --libdir` +# Set up library paths for runtime +export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir) CUDA_HOME=/usr/local/cuda SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp" @@ -38,11 +38,11 @@ g++ -o build/aot_use_in_cpp_bundle \ -I${CUDA_HOME}/include \ `tvm-ffi-config --cxxflags` \ ${SOURCE_FILE} build/add_one.o \ - -L${CUDA_DIALECT_PATH} \ + $(python3 -m cutlass.cute.export.aot_config --ldflags) \ -L${CUDA_HOME}/lib64 \ - -lcuda_dialect_runtime -lcuda -lcudart \ - `tvm-ffi-config --ldflags` \ - `tvm-ffi-config --libs` + $(python3 -m cutlass.cute.export.aot_config --libs) -lcuda -lcudart \ + $(tvm-ffi-config --ldflags) \ + $(tvm-ffi-config --libs) echo "Running the executable..." ./build/aot_use_in_cpp_bundle diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py index 0aaebc1a..ce94567b 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py @@ -37,7 +37,7 @@ def main(): a_jax = jnp.arange(10, dtype=jnp.float32) b_jax = jnp.zeros(10, dtype=jnp.float32) lib_path = "./build/add_one.so" - aot_mod = cute.runtime.load_module(lib_path) + aot_mod = cute.runtime.load_module(lib_path, enable_tvm_ffi=True) jax_tvm_ffi.register_ffi_target("add_one_cute", aot_mod.add_one, platform="gpu") b_jax = jax.ffi.ffi_call( "add_one_cute", diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py index 20a224a9..62ed14b7 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py @@ -27,14 +27,16 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import cutlass.cute as cute -import torch + # now load it back def main(): + import torch + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") lib_path = "./build/add_one.so" - aot_mod = cute.runtime.load_module(lib_path) + aot_mod = cute.runtime.load_module(lib_path, enable_tvm_ffi=True) aot_mod.add_one(a_torch, b_torch) print("result of b after aot_mod.add_one(a, b)") print(b_torch) diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py b/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py index 7a50a232..7b9bb275 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py @@ -36,8 +36,9 @@ To run this example: .. code-block:: bash - python examples/cute/tvm_ffi/error_reporting.py + python cutlass_ir/compiler/python/examples/cute/tvm_ffi/error_reporting.py """ + import torch import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py index 36b3e823..baf52f0b 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py @@ -39,7 +39,7 @@ To run this example: pip install jax-tvm-ffi pip install jax[cuda13] - python examples/cute/tvm_ffi/jit_and_use_in_jax.py + python cutlass_ir/compiler/python/examples/cute/tvm_ffi/jit_and_use_in_jax.py """ import jax diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py index be9867be..97404c15 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py @@ -36,9 +36,9 @@ To run this example: .. code-block:: bash - python examples/cute/tvm_ffi/jit_and_use_in_torch.py + python cutlass_ir/compiler/python/examples/cute/tvm_ffi/jit_and_use_in_torch.py """ -import torch + import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack @@ -56,6 +56,8 @@ def add_one(a: cute.Tensor, b: cute.Tensor): def main(): + import torch + # compile the kernel with "--enable-tvm-ffi" option a_torch = torch.arange(10, dtype=torch.float32, device="cuda") b_torch = torch.zeros(10, dtype=torch.float32, device="cuda") diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/requirements.txt b/examples/python/CuTeDSL/cute/tvm_ffi/requirements.txt index 8a3f9dcd..f3a147e5 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/requirements.txt +++ b/examples/python/CuTeDSL/cute/tvm_ffi/requirements.txt @@ -1,2 +1,2 @@ apache-tvm-ffi -torch-c-dlpack-ext \ No newline at end of file +torch-c-dlpack-ext diff --git a/examples/python/CuTeDSL/jax/cutlass_call_basic.py b/examples/python/CuTeDSL/jax/cutlass_call_basic.py new file mode 100644 index 00000000..e1860403 --- /dev/null +++ b/examples/python/CuTeDSL/jax/cutlass_call_basic.py @@ -0,0 +1,260 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +from functools import partial +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.jax as cjax +import cuda.bindings.driver as cuda + +""" +Examples of calling CuTe DSL from jax.jit function using cutlass_call. + +cutlass_call is a Jax primitive the enables calling of CuTe DSL kernels within a +a jit-compiled Jax function. During the lowering process cutlass_call will +trigger compilation of the kernel and embed it into the HLO computation. It can +then be efficiently launched by XLA without callback to Python. + +This example assumes familiarity with CuTe DSL concepts such as layouts and +dynamic shapes. + +To run this example: + +.. code-block:: bash + + # Run with addition operation + python examples/jax/cutlass_call_basic.py +""" + + +# This is a typical CuTe DSL kernel function that accepts both tensor and scalar values. +@cute.jit +def launch( + A: cute.Tensor, + B: cute.Tensor, + x: cute.Int32, + y: cute.Int32, + C: cute.Tensor, + D: cute.Tensor, + stream: cuda.CUstream, +): + # Print layouts + print("A layout: ", A.layout) + print("B layout: ", B.layout) + print("C layout: ", C.layout) + print("D layout: ", D.layout) + cute.printf("A layout: {}", A.layout) + cute.printf("B layout: {}", B.layout) + cute.printf("C layout: {}", C.layout) + cute.printf("D layout: {}", D.layout) + cute.printf("") + + # Print non-tensor values + print("X is: ", x) + print("Y is: ", y) + cute.printf("X is: {}", x) + cute.printf("Y is: {}", y) + print() + + +# cutlass_call uses a fixed function signature to pass arguments between Jax and CuTeDSL kernel. +# +# Function Signature Requirement: +# stream, inputs, outputs, *, kwargs... +# +# The first argument must be the CUstream that the kernel is run. This stream is managed by the XLA runtime +# and is necessary to schedule and synchronize launches with the rest of your computation. +# +# The second set of arguments are the Jax arrays for inputs and outputs. Inputs must be passed before +# outputs. +# +# Lastly static arguments (i.e. static_argnums or static_argnames) values are passed as keyword only arguments +# by name. +# +# The the kernel does not match this signature a wrapper functions like the one shown below can be written +# or an inline lambda function can be used to rebind the arguments into the appropriate order. +@cute.jit +def launch_jax_wrapper( + stream: cuda.CUstream, + A: cute.Tensor, + B: cute.Tensor, + C: cute.Tensor, + D: cute.Tensor, + *, + x: cute.Int32, + y: cute.Int32, +): + launch(A, B, x, y, C, D, stream) + + +@cute.jit +def launch_aliased( + A: cute.Tensor, B: cute.Tensor, x: cute.Int32, y: cute.Int32, stream: cuda.CUstream +): + # Print layouts + print("A layout: ", A.layout) + print("B layout: ", B.layout) + cute.printf("A layout: {}", A.layout) + cute.printf("B layout: {}", B.layout) + cute.printf("") + + # Print non-tensor values + print("X is: ", x) + print("Y is: ", y) + cute.printf("X is: {}", x) + cute.printf("Y is: {}", y) + print() + + +if __name__ == "__main__": + + @partial(jax.jit, static_argnums=[2, 3]) + def run_cutlass_kernel(a, b, x, y): + call = cjax.cutlass_call( + launch_jax_wrapper, + # Jax requires output shapes/dtype information for each output + output_shape_dtype=( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, a.dtype), + ), + # Static jit arguments are passed via additional keyword arguments + x=x, + y=y, + ) + + # Returned value is a callable to invoke the kernel passing only jax arrays. + return call(a, b) + + print("\nExample: example_basic_call_from_jit") + A = jnp.zeros((512, 32, 64)) + B = jnp.zeros((1, 256, 64, 128)) + C, D = run_cutlass_kernel(A, B, 0, 1) + + @partial(jax.jit, static_argnums=[2, 3]) + def run_cutlass_kernel_lambda(a, b, x, y): + call = cjax.cutlass_call( + # A lambda function may be used to wrap and bind arguments passed by jax + # to the kernel. Alternatively you can wrap using another separate cute.jit + # function. + lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream), + # Jax requires output shapes/dtype information for each output + output_shape_dtype=( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, a.dtype), + ), + # Static jit arguments are passed via additional keyword arguments + x=x, + y=y, + ) + + # Returned value is a callable to invoke the kernel passing only jax arrays. + return call(a, b) + + print("\nExample: run_cutlass_kernel_lambda") + A = jnp.zeros((512, 32, 64)) + B = jnp.zeros((1, 256, 64, 128)) + C, D = run_cutlass_kernel_lambda(A, B, 1, 2) + + @partial(jax.jit, static_argnums=[2, 3]) + def run_cutlass_kernel_static_shapes(a, b, x, y): + call = cjax.cutlass_call( + lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream), + output_shape_dtype=( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, a.dtype), + ), + # By default cutlass_call will treat all tensors as dynamic shape. + # Dynamic shapes are often expected for kernels so this default ensures + # the broadest support. If you know that a kernel can accept fully static + # tensors then you can enable this flag to pass all tensors shapes and + # layouts known at compile time. + use_static_tensors=True, + x=x, + y=y, + ) + return call(a, b) + + print("\nExample: run_cutlass_kernel_static_shapes") + A = jnp.zeros((512, 32, 64)) + B = jnp.zeros((1, 256, 64, 128)) + C, D = run_cutlass_kernel_static_shapes(A, B, 3, 4) + + @partial(jax.jit, static_argnums=[2, 3]) + def run_cutlass_kernel_with_modes(a, b, x, y): + call = cjax.cutlass_call( + lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream), + output_shape_dtype=( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, a.dtype), + ), + # The modes of the layout for each tensor can be specified using the + # TensorSpec. By default modes will align with the physical layout + # but can be mapped to specific index position. If None is passed + # then the default mode is assumed for that tensor. + # + # Individual static/dynamic settings may also be applied. For example + # a specific tensor can be marked to have static shape. + input_spec=( + cjax.TensorSpec(mode=(1, 0, 2), static=True), + cjax.TensorSpec(mode=(3, 1, 2, 0)), + ), + output_spec=(None, cjax.TensorSpec(mode=(0, 1, 3, 2))), + x=x, + y=y, + ) + return call(a, b) + + print("\nExample: run_cutlass_kernel_with_modes") + A = jnp.zeros((512, 32, 64)) + B = jnp.zeros((1, 256, 64, 128)) + C, D = run_cutlass_kernel_with_modes(A, B, 5, 6) + + @partial(jax.jit, static_argnums=[2, 3], donate_argnums=[0, 1]) + def run_cutlass_kernel_aliased_outputs(a, b, x, y): + call = cjax.cutlass_call( + lambda stream, a, b, *, x, y: launch_aliased(a, b, x, y, stream), + output_shape_dtype=( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, b.dtype), + ), + # Can specify the input tensors that are aliasing outputs of this call. + # To avoid allocating separate output buffers. This is useful for kernels + # that update a tensor. + input_output_aliases={0: 0, 1: 1}, + x=x, + y=y, + ) + return call(a, b) + + print("\nExample: run_cutlass_kernel_aliased_outputs") + A = jnp.zeros((512, 32, 64)) + B = jnp.zeros((1, 256, 64, 128)) + A, B = run_cutlass_kernel_aliased_outputs(A, B, 7, 8) diff --git a/examples/python/CuTeDSL/jax/cutlass_call_export.py b/examples/python/CuTeDSL/jax/cutlass_call_export.py new file mode 100644 index 00000000..39a8baf1 --- /dev/null +++ b/examples/python/CuTeDSL/jax/cutlass_call_export.py @@ -0,0 +1,168 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +import pytest +from functools import partial +import argparse + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute + +import jax +import jax.numpy as jnp +from jax import export + +from cutlass.jax import cutlass_call, get_export_disabled_safety_checks +from cutlass.jax.testing import create_tensor + +""" +Examples of using jax.export APIs with functions using cutlass_call. + +This example demonstrates the use of jax.export with CuTe DSL kernel. It assumes +familiarity with CuTe DSL concepts such as layouts and dynamic shapes as well as +Jax's exporting and serialization features: +https://docs.jax.dev/en/latest/export/index.html#export + +To run this example: + +.. code-block:: bash + + # Run with defaults + python examples/jax/cutlass_call_export.py + + # Run with shape (1024, 512) + python examples/jax/cutlass_call_export.py --M 1024 --N 512 + + # Export with symbolic shapes. + python examples/jax/cutlass_call_export.py --export_symbolic +""" + + +@cute.kernel +def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + bdim, _, _ = cute.arch.block_dim() + + thread_idx = bidx * bdim + tidx + + m, n = gA.shape + ni = thread_idx % n + mi = thread_idx // n + + a_val = gA[mi, ni] + b_val = gB[mi, ni] + gC[mi, ni] = a_val + b_val + + +@cute.jit +def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor): + print("mA: ", mA.layout) + print("mB: ", mB.layout) + print("mC: ", mC.layout) + num_threads_per_block = 256 + m, n = mA.shape + kernel(mA, mB, mC).launch( + grid=((m * n) // num_threads_per_block, 1, 1), + block=(num_threads_per_block, 1, 1), + stream=stream, + ) + + +def run_example(M, N, export_symbolic_shapes): + @jax.jit + def f(a, b): + call = cutlass_call(launch, output_shape_dtype=a) + return jax.nn.sigmoid(call(a, b)) + + @jax.jit + def ref_f(a, b): + return jax.nn.sigmoid(a + b) + + # Symbolic or partially shapes are supported by cutlass_call and cute.Tensor + # This allows export of functions calling Cut eDSL kernels w/o having to re-compile + # the kernel for each new shape. + if export_symbolic_shapes: + a, b = export.symbolic_shape("a, b") + export_shape_dtype = jax.ShapeDtypeStruct((a, b), jnp.float32) + else: + export_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32) + + print("Exporting with input signature: ") + print(f"({export_shape_dtype}, {export_shape_dtype})") + + # jax.export can be used to export a jit function containing cutlass_call. + # The function get_export_disabled_safety_checks() returns a list of custom + # call targets that are used by cutlass_call not part of Jax's built-in + # list of stable custom calls. + exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks()) + traced = exported(export_shape_dtype, export_shape_dtype) + + # Serialize the computation to a byte blob. + blob = traced.serialize() + print(f"Serialized computation is {len(blob)} bytes.") + + # Deserialize and run + rehydrated = export.deserialize(blob) + + key = jax.random.key(1123) + a_key, b_key = jax.random.split(key, 2) + + a = create_tensor((M, N), dtype=jnp.float32, key=a_key) + b = create_tensor((M, N), dtype=jnp.float32, key=b_key) + c = rehydrated.call(a, b) + assert jnp.allclose(c, ref_f(a, b)) + + # If the computation was exported with dynamic shapes then we can also + # call it with different shapes. The kernel will not be re-compiled + # even though the shapes are changing. + if export_symbolic_shapes: + a = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=a_key) + b = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=b_key) + c = rehydrated.call(a, b) + assert jnp.allclose(c, ref_f(a, b)) + + a = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=a_key) + b = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=b_key) + c = rehydrated.call(a, b) + assert jnp.allclose(c, ref_f(a, b)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Demonstration of using jax.export with functions with cutlass_call" + ) + parser.add_argument("--M", default=512, type=int) + parser.add_argument("--N", default=256, type=int) + parser.add_argument("--export_symbolic", action="store_true") + + args = parser.parse_args() + run_example(args.M, args.N, args.export_symbolic) + print("PASS") diff --git a/examples/python/CuTeDSL/jax/cutlass_call_sharding.py b/examples/python/CuTeDSL/jax/cutlass_call_sharding.py new file mode 100644 index 00000000..a2d03c34 --- /dev/null +++ b/examples/python/CuTeDSL/jax/cutlass_call_sharding.py @@ -0,0 +1,140 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + +from functools import partial +import argparse + +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from jax.experimental.custom_partitioning import custom_partitioning + +import cutlass +import cutlass.cute as cute +import cutlass.jax as cjax +from cutlass.jax.testing import create_tensor +import cuda.bindings.driver as cuda + + +""" +Examples of combining jax.jit and jax.shard_map for sharding and executing kernels +across multiple GPU devices. + +To run this example: + +.. code-block:: bash + + # Run with addition operation + python examples/jax/cutlass_call_sharding.py +""" + + +@cute.kernel +def kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type) + frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type) + frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type) + + cute.autovec_copy(a[None, tidx, bidx], frgA) + cute.autovec_copy(b[None, tidx, bidx], frgB) + frgC.store(frgA.load() + frgB.load()) + cute.autovec_copy(frgC, c[None, tidx, bidx]) + + +@cute.jit +def launch( + stream: cuda.CUstream, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, +): + cute.printf("a: {}", a.layout) + cute.printf("b: {}", b.layout) + cute.printf("c: {}", c.layout) + kernel(a, b, c).launch( + grid=[a.shape[-1], 1, 1], block=[a.shape[-2], 1, 1], stream=stream + ) + + +def run_example(): + # Create a device mesh with one axis b + ngpu = jax.device_count() + mesh = jax.make_mesh((ngpu,), "b") + + if ngpu == 1: + print("Note: only 1 GPU was detected.") + + # We will shard our 3D tensors over b + sharding = P("b", None, None) + + @partial(jax.jit, static_argnums=[0, 1]) + def allocate_sharded_tensors(shape, dtype): + key = jax.random.key(1123) + a_key, b_keys = jax.random.split(key, 2) + a = create_tensor(shape, dtype, a_key) + b = create_tensor(shape, dtype, b_keys) + a = jax.lax.with_sharding_constraint(a, NamedSharding(mesh, sharding)) + b = jax.lax.with_sharding_constraint(b, NamedSharding(mesh, sharding)) + return a, b + + @jax.jit + def compute(a, b): + # This jax.shard_map partitions the cutlass_call over the mesh. + @partial( + jax.shard_map, + mesh=mesh, + in_specs=(sharding, sharding), + out_specs=(sharding, sharding), + ) + def sharded_call(a_block, b_block): + call = cjax.cutlass_call( + launch, + use_static_tensors=True, + output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype), + ) + ref_result = a_block + b_block + return call(a_block, b_block), ref_result + + return sharded_call(a, b) + + # Allocate (32, 16, 64) on each GPU + shape = (32 * ngpu, 16, 64) + dtype = jnp.float32 + + a, b = allocate_sharded_tensors(shape, dtype) + c, c_ref = compute(a, b) + + assert jnp.allclose(c, c_ref) + + +if __name__ == "__main__": + run_example() + print("PASS") diff --git a/examples/python/CuTeDSL/jax/elementwise_apply_example.py b/examples/python/CuTeDSL/jax/elementwise_apply_example.py new file mode 100644 index 00000000..c958d09c --- /dev/null +++ b/examples/python/CuTeDSL/jax/elementwise_apply_example.py @@ -0,0 +1,329 @@ +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. + + +import argparse +import operator +from functools import partial +from typing import List, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute + +""" +An Elementwise Apply Example using CuTe DSL with cutlass.jax.cutlass_call + +This example is similar to examples/ampere/elementwise_apply.py but demonstrates +how to run the code in a jax specific way using the cutlass_call primitive. It assumes +familiarity with basic CuTe DSL concepts as well as the cutlass_call primitive. + +To run this example: + +.. code-block:: bash + + # Run with addition operation + python examples/jax/elementwise_apply_example.py --M 1024 --N 512 --op add + + # Run with multiplication operation + python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op mul + + # Run with subtraction operation + python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op sub +""" + + +@cute.kernel +def elementwise_apply_kernel( + op: cutlass.Constexpr, + mInputs: List[cute.Tensor], + mC: cute.Tensor, + cC: cute.Tensor, # coordinate tensor + shape: cute.Shape, + tv_layout: cute.Layout, # (tid, vid) -> logic coord +): + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + + ############################################################################### + # Slice to local tile of thread block + ############################################################################### + blk_crd = ((None, None), (bidx, bidy)) + + # Leverage the meta-programming capability of the DSL to slice the tensors for each input + # All for loops below on input tensors would be fully unrolled automatically at compile time + # logical coord -> memory address + gInputs = [t[blk_crd] for t in mInputs] # (TileM, TileN) + gC = mC[blk_crd] # (TileM, TileN) + gCrd = cC[blk_crd] # (TileM, TileN) + + print("[DSL INFO] Sliced Tensors per thread block:") + for i in cutlass.range_constexpr(len(gInputs)): + print(f"[DSL INFO] ctaInputs{i} = {gInputs[i].type}") + print(f"[DSL INFO] gC = {gC.type}") + print(f"[DSL INFO] gCrd = {gCrd.type}") + + ############################################################################### + # Compose with thread block TV layout to map thread & value indices to memory address + ############################################################################### + # (tid, vid) -> memory address + tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs] + tidfrgC = cute.composition(gC, tv_layout) + tidfrgCrd = cute.composition(gCrd, tv_layout) + + # repeat None like vid to remove hierarchy of layout + thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1])) + + ############################################################################### + # Slice to local tile of thread + ############################################################################### + # vid -> address + thrInputs = [t[thr_crd] for t in tidfrgInputs] # (V) + thrC = tidfrgC[thr_crd] # (V) + thrCrd = tidfrgCrd[thr_crd] + + print("[DSL INFO] Sliced Tensors per thread:") + for i in cutlass.range_constexpr(len(thrInputs)): + print(f"[DSL INFO] thrInputs{i} = {thrInputs[i].type}") + print(f"[DSL INFO] thrC = {thrC.type}") + print(f"[DSL INFO] thrCrd = {thrCrd.type}") + + ############################################################################### + # Compute predicate for out of boundary checks + ############################################################################### + frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean) + print(f"[DSL INFO] frgPred = {frgPred.type}") + + for i in cutlass.range_constexpr(cute.size(frgPred)): + frgPred[i] = cute.elem_less(thrCrd[i], shape) + + # if tidx == 0 and bidx == 0: + # cute.print_tensor(frgPred) + + ########################################################## + # Load data and compute result + ########################################################## + + # Load data before use. The compiler will optimize the copy and load + # operations to convert some memory ld/st into register uses. + result = op(*[thrInput.load() for thrInput in thrInputs]) + thrC.store(result) + + +@cute.jit +def elementwise_apply( + op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream +): + """CUDA kernel applying binary operator on each element of two n-D input tensors in + CuTe Python and store to result tensor. + + :param op: Binary operator or lambda function to apply element-wise + :type op: cutlass.Constexpr + :param a: First input tensor + :type a: cute.Tensor + :param b: Second input tensor + :type b: cute.Tensor + :param result: Output tensor to store the results of op(a, b) + :type result: cute.Tensor + :return: None + :rtype: None + """ + + # Baseline: naive TV layout + # * mA layout: (4096, 4096):(4096, 1) + # * TV layout map to (512, 4) tile + # * tidx maps to mode-0 but input layout is contiguous on mode-1, performance will be bad + # tv_layout = cute.make_layout((128, (4, 4)), stride=(4, (512, 1))) + # cta_tiler = (512, 4) + + # Opt-1: better TV layout with better 1D thread layout (SOL with 1D thread layout) + # * mA layout: (4096, 4096):(4096, 1) + # * TV layout map to (4, 512) tile + # * tidx maps to mode-1 which is leading mode of input tensor for coalesced load + # tv_layout = cute.make_layout((128, (4, 4)), stride=(16, (4, 1))) + # cta_tiler = (4, 512) + + # Opt-2: 2D tile but worse + # * mA layout: (4096, 4096):(4096, 1) + # * TV layout map to (128, 16) logical tile + # * V layout is bad as contiguous mode is not on right-most + # * `cute.copy` only supports vectorize when stride-1 of v-layout on right-most ) + # tv_layout = cute.make_layout(((32, 4), (4, 4)), stride=((4, 512), (1, 128))) + # cta_tiler = (128, 16) + + # Opt-3: SOL with 2D thread tile + # * mA layout: (4096, 4096):(4096, 1) + # * TV layout map to (64, 256) logical tile + # * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store + + # Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type + coalesced_ldst_bytes = 16 + + # Compile time validation: expect same element type for all input tensors + assert all(t.element_type == inputs[0].element_type for t in inputs) + dtype = inputs[0].element_type + + thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0)) + val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0)) + val_layout = cute.recast_layout(dtype.width, 8, val_layout) + tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout) + + print("[DSL INFO] Input Tensors:") + for i, t in enumerate(inputs): + print(f"[DSL INFO] inputs{i} = {t}") + print(f"[DSL INFO] result = {result}") + + print("[DSL INFO] Tiling Parameters:") + print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block") + print(f"[DSL INFO] tv_layout = {tv_layout}") + + print("[DSL INFO] Tiled Tensors:") + mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs] + # ((TileM, TileN), (RestM, RestN)) + mC = cute.zipped_divide(result, tiler_mn) + + # (RestM, RestN) -> (RestN, RestM) + remap_block = cute.make_ordered_layout( + cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0) + ) + for i, t in enumerate(mInputs): + print(f"[DSL INFO] gInputs{i} = {mInputs[i]}") + mInputs[i] = cute.composition(t, (None, remap_block)) + print(f"[DSL INFO] gInputs{i} (remapped) = {mInputs[i]}") + + mC = cute.composition(mC, (None, remap_block)) + print(f"[DSL INFO] gC = {mC}") + + idC = cute.make_identity_tensor(result.shape) + cC = cute.zipped_divide(idC, tiler=tiler_mn) + print(f"[DSL INFO] coord tensor = {cC}") + + # Launch the kernel asynchronously + # Group input tensors into a list as a single argument + elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch( + # Compute production at each mode of mC.shape[1] to get multi-dimensional grid size + grid=cute.product_each(mC.shape[1]), + block=[cute.size(tv_layout, mode=[0]), 1, 1], + stream=stream, + ) + + +@cutlass.dsl_user_op +def leaky_relu(x, alpha, *, loc=None, ip=None): + return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip) + + +def leaky_relu_ref(x, alpha): + import jax.numpy as jnp + + return jnp.where(x > 0, x, alpha * x) + + +def run_and_verify(op, M, N, dtype, skip_ref_check=False): + import jax + import jax.numpy as jnp + import cutlass.jax as cjax + import cutlass.jax.testing as testing + + if op == "leaky_relu": + op = partial(leaky_relu, alpha=0.01) + ref_op = partial(leaky_relu_ref, alpha=0.01) + num_inputs = 1 + else: + op = getattr(operator, op) + ref_op = op + num_inputs = 2 + + # This jax function is transformed using jax.jit to compile its contents + # into an efficient HLO executable. + @partial(jax.jit, static_argnums=[1]) + def jax_function(inputs, op): + call = cjax.cutlass_call( + # Bind jax arguments to kernel signature + lambda stream, inputs, output, *, op: elementwise_apply( + op, inputs, output, stream + ), + # Specify output shape/dtype of result + output_shape_dtype=jax.ShapeDtypeStruct(inputs[0].shape, inputs[0].dtype), + # Pass static/constexpr values as kwargs + op=op, + ) + + # Call the kernel! + return call(inputs) + + @partial(jax.jit, static_argnums=[1]) + def jax_ref_function(inputs, op): + return op(*inputs) + + print("\nRunning Elementwise Apply test with:") + print(f"Tensor dimensions: [{M}, {N}]") + print(f"Input and Output Data type: {dtype}") + + jax_dtype = cjax.cutlass_to_jax_dtype(dtype) + keys = jax.random.split(jax.random.key(1435), num_inputs) + inputs = [testing.create_tensor((M, N), jax_dtype, key) for key in keys] + + print("Input tensor shapes:") + for i in range(num_inputs): + print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}") + + epsilon = 1.2 + if op in (operator.truediv, operator.floordiv): + inputs[1] = jnp.where(inputs[1] == 0, epsilon, inputs[1]) + + # Call the jax.jit function which will compile the kernel + c = jax_function(inputs, op) + + if not skip_ref_check: + print("Executing elementwise apply kernel...") + c = jax_function(inputs, op) + print("Verifying results...") + assert jnp.allclose(ref_op(*inputs), c) + print("Results verified successfully!") + print(f"First few elements of result: \n{c[:3, :3]}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Demonstration of calling a kernel with cutlass_call" + ) + parser.add_argument("--M", default=4096, type=int) + parser.add_argument("--N", default=4096, type=int) + parser.add_argument("--op", default="add", type=str) + parser.add_argument("--skip_ref_check", action="store_true") + + args = parser.parse_args() + run_and_verify( + args.op, + args.M, + args.N, + dtype=cutlass.Float32, + skip_ref_check=args.skip_ref_check, + ) + print("\nPASS") diff --git a/include/cute/arch/config.hpp b/include/cute/arch/config.hpp index 8d03794c..d9cecf9f 100644 --- a/include/cute/arch/config.hpp +++ b/include/cute/arch/config.hpp @@ -164,6 +164,9 @@ # define CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED # define CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED # endif +# if (__CUDACC_VER_MAJOR__ == 13 && __CUDACC_VER_MINOR__ >= 1) +# define CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED +# endif #endif #if defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) || defined(CUTLASS_ARCH_MMA_SM103F_ENABLED) diff --git a/include/cute/arch/mma_sm100_desc.hpp b/include/cute/arch/mma_sm100_desc.hpp index 6abc527c..887b5896 100644 --- a/include/cute/arch/mma_sm100_desc.hpp +++ b/include/cute/arch/mma_sm100_desc.hpp @@ -449,7 +449,7 @@ union InstrDescriptorBlockScaled : 1, // b_sf_id_ : 2, // bit [ 4, 6) : Matrix B Scale Factor ID : 1, // - a_format_ : 3, // bit [ 7,10) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean + a_format_ : 3, // bit [ 7, 10): MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean b_format_ : 3, // bit [10,13) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean a_negate_ : 1, // bit [13,14) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format b_negate_ : 1, // bit [14,15) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format diff --git a/include/cute/arch/mma_sm100_umma.hpp b/include/cute/arch/mma_sm100_umma.hpp index 5e611182..d7dfb715 100644 --- a/include/cute/arch/mma_sm100_umma.hpp +++ b/include/cute/arch/mma_sm100_umma.hpp @@ -208,6 +208,91 @@ struct SM100_MMA_F16BF16_TS } }; +template +struct SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN +{ + static_assert(M == 64 || M == 128, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN M-mode size should be 64 or 128 for 1 CTA cluster MMA."); + static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) || + (M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)), + "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN N-mode size should be a multiple of 8 between 8 and 256 for M=64,\ + or a multiple of 16 between 16 and 256 for M=128."); + static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed"); + static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose"); + + using DRegisters = void; + using ARegisters = uint32_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& tmem_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& scaleC, + uint64_t const& idescE) + { +#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED) + if (cute::elect_one_sync()) { + uint32_t mask[4] = {0, 0, 0, 0}; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p; \n\t" + "}\n" + : + : "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC), + "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3])); + } +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN without CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED"); +#endif + } +}; + +template +struct SM100_MMA_TF32_SS_SCALED +{ + static_assert(M == 64 || M == 128, "SM100_MMA_TF32_SS_SCALED M-mode size should be 64 or 128 for 1 CTA cluster MMA."); + static_assert((N % 8 == 0) && (8 <= N) && (N <= 256), + "SM100_MMA_TF32_SS_SCALED N-mode size should be a multiple of 8 between 8 and 256."); + + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& accumulate, + uint64_t const& idescE) + { +#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED) + if (cute::elect_one_sync()) { + // ScaleC input should be a literal or compile time constant + uint32_t mask[4] = {0, 0, 0, 0}; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::tf32 [%0], %1, %2, %3, {%5, %6, %7, %8}, p, %9; \n\t" + "}\n" + : + : "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate), + "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), "n"(ScaleC)); + } +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_SS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED"); +#endif + } +}; + template @@ -249,6 +334,51 @@ struct SM100_MMA_F16BF16_SS_SCALED } }; +template +struct SM100_MMA_TF32_TS_SCALED +{ + static_assert(M == 64 || M == 128, "SM100_MMA_TF32_TS_SCALED M-mode size should be 64 or 128 for 1 CTA cluster MMA."); + static_assert((M == 64 && (N % 8 == 0) && (8 <= N) && (N <= 256)) || + (M == 128 && (N % 16 == 0) && (16 <= N) && (N <= 256)), + "SM100_MMA_TF32_TS_SCALED N-mode size should be a multiple of 8 between 8 and 256 for M=64,\ + or a multiple of 16 between 16 and 256 for M=128."); + static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_TS_SCALED A from TMEM can't be transposed"); + + using DRegisters = void; + using ARegisters = uint32_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& tmem_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& accumulate, + uint64_t const& idescE) + { +#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED) + if (cute::elect_one_sync()) { + // ScaleC input should be a literal or compile time constant + uint32_t mask[4] = {0, 0, 0, 0}; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8}, p, %9; \n\t" + "}\n" + : + : "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate), + "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), "n"(ScaleC)); + } +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_TS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED"); +#endif + } +}; + template +struct SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN +{ + static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN M-mode size should be 128 or 256 for 2 CTA cluster MMA."); + static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN N-mode size should be a multiple of 32 between 32 and 256."); + static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed"); + static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose"); + + using DRegisters = void; + using ARegisters = uint32_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& tmem_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& scaleC, + uint64_t const& idescE) + { +#if defined(CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED) + if (cute::elect_one_sync()) { + uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p; \n\t" + "}\n" + : + : "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(scaleC), + "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), + "r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7])); + } +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN without CUTE_ARCH_TCGEN05_TF32_MMA_ENABLED"); +#endif + } +}; + +template +struct SM100_MMA_TF32_2x1SM_SS_SCALED +{ + static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_SS_SCALED M-mode size should be 128 or 256 for 2 CTA cluster MMA."); + static_assert((N % 16 == 0) && (16 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_SS_SCALED N-mode size should be a multiple of 16 between 16 and 256."); + + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& accumulate, + uint64_t const& idescE) + { +#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED) + if (cute::elect_one_sync()) { + // ScaleC input should be a literal or compile time constant + uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::tf32 [%0], %1, %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p, %13; \n\t" + "}\n" + : + : "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate), + "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), + "r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]), "n"(ScaleC)); + } +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_SS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED"); +#endif + } +}; + template @@ -581,6 +794,49 @@ struct SM100_MMA_F16BF16_2x1SM_SS_SCALED } }; +template +struct SM100_MMA_TF32_2x1SM_TS_SCALED +{ + static_assert(M == 128 || M == 256, "SM100_MMA_TF32_2x1SM_TS_SCALED M-mode size should be 128 or 256 for 2 CTA cluster MMA."); + static_assert((N % 32 == 0) && (32 <= N) && (N <= 256), "SM100_MMA_TF32_2x1SM_TS_SCALED N-mode size should be a multiple of 32 between 32 and 256."); + static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_SCALED A from TMEM can't be transposed"); + + using DRegisters = void; + using ARegisters = uint32_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& tmem_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& accumulate, + uint64_t idescE) + { +#if defined(CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED) + if (cute::elect_one_sync()) { + // ScaleC input should be a literal or compile time constant + uint32_t mask[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::2.kind::tf32 [%0], [%1], %2, %3, {%5, %6, %7, %8, %9, %10, %11, %12}, p, %13; \n\t" + "}\n" + : + : "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(uint32_t(idescE>>32)), "r"(accumulate), + "r"(mask[0]), "r"(mask[1]), "r"(mask[2]), "r"(mask[3]), + "r"(mask[4]), "r"(mask[5]), "r"(mask[6]), "r"(mask[7]), "n"(ScaleC)); + } +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM100_MMA_TF32_2x1SM_TS_SCALED without CUTE_ARCH_TCGEN05_UTFMMA_SCALED_ENABLED"); +#endif + } +}; + template static constexpr uint16_t bidB = 0; CUTE_STATIC_ASSERT(VS == 16 || VS == 32, "Scaling factor vector size has to be 16 or 32 for MXF4NVF4 MMA."); + if constexpr ( VS == 16 ) { +#if defined(CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue8m0 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13}," + "{%14}," + "{%15, %16}," + "{%17}," + "{%18, %19};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "f"(c0), "f"(c1), "f"(c2), "f"(c3), + "r"(uint32_t(sfa0)) , "h"(bidA), "h"(tidA), + "r"(uint32_t(sfb0)) , "h"(bidB), "h"(tidB)); +#else + CUTE_INVALID_CONTROL_PATH("Attempting to use SM120::BLOCKSCALED::SM120_16x8x64_TN_VS without CUTE_ARCH_MXF4NVF4_4X_UE8M0_MMA_ENABLED"); +#endif + } else if constexpr ( VS == 32 ) { + #if defined(CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED) asm volatile( "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::2X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 " @@ -3151,6 +3174,7 @@ struct SM120_16x8x64_TN_VS #else CUTE_INVALID_CONTROL_PATH("Attempting to use SM120::BLOCKSCALED::SM120_16x8x64_TN_VS without CUTE_ARCH_MXF4NVF4_2X_UE8M0_MMA_ENABLED"); #endif + } } }; diff --git a/include/cute/arch/mma_sm120_sparse.hpp b/include/cute/arch/mma_sm120_sparse.hpp index 297ea3b5..a95af101 100644 --- a/include/cute/arch/mma_sm120_sparse.hpp +++ b/include/cute/arch/mma_sm120_sparse.hpp @@ -3324,6 +3324,28 @@ struct SM120_SPARSE_16x8x128_TN_VS struct AuxTmaParams { using GmemStrides = GmemTmaBasisStrides_; // Strides for Gmem mode -> Tma coord mode, may be dynamic GmemStrides g_stride_; - using TmaGmemBasis = TmaGmemBasis_; // Layout for Tma box shape -> Gmem mode(s), always static - static_assert(is_static::value); + using TmaGmemBasis = TmaGmemBasis_; // Layout for Tma box shape -> Gmem mode(s) + // By default, TmaGmemBasis produced by construct_tma_gbasis is fully static. + // The user may construct a dynamic gbasis manually (e.g. to represent smem box with dynamic shape). + // In that case they will need to pass it around via other means. + // We avoid passing it as a data member to avoid ABI impact. + // static_assert(is_static::value); using TmaSwizzle = TmaSwizzle_; // Tma swizzle, always Swizzle static_assert(is_static::value); }; @@ -70,7 +74,7 @@ struct TMA_LOAD_Unpack { static_assert(is_smem::value, "SM90_TMA_LOAD requires the destination be shared memory."); - auto src_coord = src.data().coord_; + auto src_coord = src(Int<0>{}); void* dst_ptr = cute::raw_pointer_cast(dst.data()); #if 0 auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0); @@ -237,7 +241,7 @@ struct Copy_Traits Tensor const& src, Tensor & dst) { - auto src_coord = src.data().coord_; + auto src_coord = src(Int<0>{}); return detail::explode_tuple(detail::CallCOPY{}, traits.opargs_, tuple_seq{}, src_coord, tuple_seq{}); @@ -405,7 +409,7 @@ struct Copy_Traits void const* const desc_ptr = &(traits.tma_desc_); void const* const src_ptr = cute::raw_pointer_cast(src.data()); - auto dst_coord = dst.data().coord_; + auto dst_coord = dst(Int<0>{}); #if 0 auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0); printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n", @@ -446,7 +450,7 @@ struct Copy_Traits void const* const desc_ptr = traits.tma_desc_; void const* const src_ptr = cute::raw_pointer_cast(src.data()); - auto dst_coord = dst.data().coord_; + auto dst_coord = dst(Int<0>{}); #if 0 auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0); printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n", @@ -531,7 +535,8 @@ struct Copy_Traits static_assert(is_smem::value, "Expected smem src for SM90_TMA_REDUCE_ADD"); //static_assert(is_gmem::value, "Expected gmem dst for SM90_TMA_REDUCE_ADD"); // TMA spoofed src tensor - traits.copy_unpack_(cute::raw_pointer_cast(src.data()), dst.data().coord_, tuple_seq{}); + auto dst_coord = dst(Int<0>{}); + traits.copy_unpack_(cute::raw_pointer_cast(src.data()), dst_coord, tuple_seq{}); } }; diff --git a/include/cute/atom/mma_traits_sm100.hpp b/include/cute/atom/mma_traits_sm100.hpp index e9fcd725..c9b42a8f 100644 --- a/include/cute/atom/mma_traits_sm100.hpp +++ b/include/cute/atom/mma_traits_sm100.hpp @@ -1353,6 +1353,186 @@ struct MMA_Traits scaleC) const { return {accumulate, idesc_}; } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } +}; + +// Special instantiation for interleaved complex (emulated) +template +struct MMA_Traits, cutlass::complex, float, + M, N, a_major, b_major, + ScaleC, a_neg, b_neg>> +{ + static_assert(cute::sizeof_bits_v == 16, "Only supports 16bit base types"); + using a_type = complex; + using b_type = complex; + using c_type = float; + + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + + using FrgTypeA = UMMA::smem_desc; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_1sm; + + // Logical shape-K is always 256bits, transform to units of elements + static constexpr int K = 256 / cute::sizeof_bits::value; + + static constexpr uint32_t ScalingFactor = ScaleC; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_1>; + using ALayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint64_t desc_a = A[0]; + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_F16BF16_SS_SCALED::fma(desc_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(UMMA::ScaleOut accumulate, cute::integral_constant scaleC) const { + return {accumulate, idesc_}; + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } +}; + +template +struct MMA_Traits> +{ + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + static_assert(cute::sizeof_bits_v == cute::sizeof_bits_v && cute::sizeof_bits_v == 32, "SM100_MMA_TF32_TS_SCALED supports 32bit types"); + + using FrgTypeA = UMMA::tmem_frg_1sm; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_1sm; + + // Logical shape-K is always 256 bits; transform to units of elements + static constexpr int K = 256 / cute::sizeof_bits::value; + + static constexpr uint32_t ScalingFactor = ScaleC; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_1>; + using ALayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + a_type, b_type, c_type, M, N, a_major, b_major, a_neg, b_neg, c_sat>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint32_t tmem_a = raw_pointer_cast(A.data()); + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_TF32_TS_SCALED::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(UMMA::ScaleOut accumulate, cute::integral_constant scaleC) const { + return {accumulate, idesc_}; + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } }; template scaleC) const { return {accumulate, idesc_}; } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } }; +// Special instantiation for interleaved complex (emulated) +template +struct MMA_Traits, cutlass::complex, float, + M, N, a_major, b_major, + ScaleC, a_neg, b_neg, c_sat>> +{ + static_assert(cute::sizeof_bits_v == 16, "Only supports 16bit base types"); + using a_type = complex; + using b_type = complex; + using c_type = float; + + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + + using FrgTypeA = UMMA::tmem_frg_1sm; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_1sm; + + // Logical shape-K is always 256 bits; transform to units of elements + static constexpr int K = 256 / cute::sizeof_bits::value; + + static constexpr uint32_t ScalingFactor = ScaleC; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_1>; + using ALayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint32_t tmem_a = raw_pointer_cast(A.data()); + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_F16BF16_TS_SCALED::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(UMMA::ScaleOut accumulate, cute::integral_constant scaleC) const { + return {accumulate, idesc_}; + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } +}; + + template +struct MMA_Traits> +{ + using a_type = complex; + using b_type = complex; + using c_type = float; + + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + + using FrgTypeA = UMMA::tmem_frg_1sm; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_1sm; + + // Logical shape-K is always 256 bits; transform to units of elements + static constexpr int K = 256 / cute::sizeof_bits::value; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_1>; + using ALayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride<_0,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + // MMA based interleaved complex GEMM calculates realAcc and imagAcc separately. + // This MMA_traits is used to calculate 1 of the GEMMs below : + // 1. realAcc = realA * realB + (-imagA) * imagB + // 2. imagAcc = imagA * realB + realA * imagB + // So it requires complex type operand A&B and float type operand Acc. + static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed"); + static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose"); + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + tfloat32_t, tfloat32_t, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint32_t tmem_a = raw_pointer_cast(A.data()); + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN< + M, N, + a_major, b_major, + a_neg, b_neg, c_sat>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } +}; + template scaleC) const { return {accumulate, idesc_}; } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } }; + +// Special instantiation for interleaved complex (emulated) +template +struct MMA_Traits, cutlass::complex, float, + M, N, a_major, b_major, + ScaleC, a_neg, b_neg>> +{ + static_assert(cute::sizeof_bits_v == 16, "Only supports 16bit base types"); + using a_type = complex; + using b_type = complex; + using c_type = float; + + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + + using FrgTypeA = UMMA::smem_desc; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_2sm; + + // Size of instructions's K extent is always 256bits, convert to units of element + constexpr static int K = 256 / cute::sizeof_bits::value; + constexpr static uint32_t ScalingFactor = ScaleC; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_2>; + using ALayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint64_t desc_a = A[0]; + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_F16BF16_2x1SM_SS_SCALED::fma(desc_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(UMMA::ScaleOut accumulate, cute::integral_constant scaleC) const { + return {accumulate, idesc_}; + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } +}; + +template +struct MMA_Traits> +{ + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + static_assert(cute::sizeof_bits_v == cute::sizeof_bits_v && cute::sizeof_bits_v == 32, "SM100_MMA_TF32_2x1SM_TS_SCALED supports 32bit types"); + + using FrgTypeA = UMMA::tmem_frg_2sm; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_2sm; + + // Size of instructions' K extent is always 256 bits; convert to units of element + constexpr static int K = 256 / cute::sizeof_bits::value; + constexpr static uint32_t ScalingFactor = ScaleC; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_2>; + using ALayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + a_type, b_type, c_type, M, N, a_major, b_major, a_neg, b_neg, c_sat>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint64_t tmem_a = raw_pointer_cast(A.data()); + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_TF32_2x1SM_TS_SCALED::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(UMMA::ScaleOut accumulate, cute::integral_constant scaleC) const { + return {accumulate, idesc_}; + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } +}; + + template @@ -2011,8 +2543,105 @@ struct MMA_Traits scaleC) const { return {accumulate, idesc_}; } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } }; + +// Special instantiation for interleaved complex (emulated) +template +struct MMA_Traits, cutlass::complex, float, + M, N, a_major, b_major, + ScaleC, a_neg, b_neg, c_sat>> +{ + static_assert(cute::sizeof_bits_v == 16, "Only supports 16bit base types"); + using a_type = complex; + using b_type = complex; + using c_type = float; + + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + + using FrgTypeA = UMMA::tmem_frg_2sm; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_2sm; + + // Size of instructions' K extent is always 256 bits; convert to units of element + constexpr static int K = 256 / cute::sizeof_bits::value; + constexpr static uint32_t ScalingFactor = ScaleC; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_2>; + using ALayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + ab_vtype, ab_vtype, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint64_t tmem_a = raw_pointer_cast(A.data()); + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_F16BF16_2x1SM_TS_SCALED::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(UMMA::ScaleOut accumulate, cute::integral_constant scaleC) const { + return {accumulate, idesc_}; + } + + template + CUTE_HOST_DEVICE constexpr + MMA_Traits> + with(cute::integral_constant) const { + return {accumulate_, UMMA::make_instr_desc()}; + } +}; + + template +struct MMA_Traits> +{ + using a_type = complex; + using b_type = complex; + using c_type = float; + + using ValTypeD = c_type; + using ValTypeA = a_type; + using ValTypeB = b_type; + using ValTypeC = c_type; + + using FrgTypeA = UMMA::tmem_frg_2sm; + using FrgTypeB = UMMA::smem_desc; + using FrgTypeC = UMMA::tmem_frg_2sm; + + // Size of instructions' K extent is always 256 bits; convert to units of element + constexpr static int K = 256 / cute::sizeof_bits::value; + + using Shape_MNK = Shape,Int,Int>; + using ThrID = Layout<_2>; + using ALayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using BLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + using CLayout = Layout,Int>>, + Stride,Stride< _1,Int>>>; + + // Accumulate or overwrite C. 1: read C, 0: ignore C [clear accumulators] + UMMA::ScaleOut accumulate_ = UMMA::ScaleOut::One; + + // Interleaved complex GEMM calculates realAcc and imagAcc separately. + // This MMA_traits is used to calculate 1 of the GEMMs below : + // 1. realAcc = realA * realB + (-imagA) * imagB + // 2. imagAcc = imagA * realB + realA * imagB + // So it requires complex type operand A&B and float type operand Acc. + static_assert(a_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN A from TMEM can't be transposed"); + static_assert(b_major == UMMA::Major::K, "SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN B from SMEM requires non-transpose"); + + UMMA::InstrDescriptor idesc_ = UMMA::make_instr_desc< + tfloat32_t, tfloat32_t, float, M, N, a_major, b_major, a_neg, b_neg, c_sat>(); + + template + CUTE_HOST_DEVICE constexpr friend + void + mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) + { + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected desc registers in MMA_Atom::call"); + static_assert(is_tmem::value, "Expected tmem in MMA_Atom::call"); + + uint64_t tmem_a = raw_pointer_cast(A.data()); + uint64_t desc_b = B[0]; + uint32_t tmem_c = raw_pointer_cast(D.data()); + uint64_t idesc = UMMA::make_runtime_instr_desc<>(traits.idesc_); + + SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN< + M, N, + a_major, b_major, + a_neg, b_neg, c_sat>::fma(tmem_a, desc_b, tmem_c, uint32_t(traits.accumulate_), idesc); + } +}; + template @@ -2640,7 +3345,7 @@ struct MMA_Traits <= 8 && cute::sizeof_bits_v <= 8, "SM100_MMA_F8F6F4_SS supports types with leq 8bit types"); static_assert(M == 64 || M == 128, "SM100_MMA_F8F6F4_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA."); static_assert(((b_major == UMMA::Major::K) && ((N % 8 == 0) && (8 <= N) && (N <= 256))) || - ((b_major == UMMA::Major::MN) && ((N % 16 == 0) && (16 <= N) && (N <= 256))), + ((b_major == UMMA::Major::MN) && ((N % 16 == 0) && (16 <= N) && (N <= 256))), "SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 8 between 8 and 256 when B is K major. \ SM100_MMA_F8F6F4_SS N-mode size should be a multiple of 16 between 16 and 256 when B is MN major."); using FrgTypeA = UMMA::smem_desc; @@ -3058,7 +3763,7 @@ struct MMA_Traits <= 8 && cute::sizeof_bits_v <= 8, "SM100_MMA_F8F6F4_2x1SM_SS supports types with leq 8bit types"); static_assert(M == 128 || M == 256, "SM100_MMA_F8F6F4_2x1SM_SS M-mode size should be 64 or 128 for 1 CTA cluster MMA."); static_assert(((b_major == UMMA::Major::K) && ((N % 16 == 0) && (16 <= N) && (N <= 256))) || - ((b_major == UMMA::Major::MN) && ((N % 32 == 0) && (32 <= N) && (N <= 256))), + ((b_major == UMMA::Major::MN) && ((N % 32 == 0) && (32 <= N) && (N <= 256))), "SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 16 between 16 and 256 when B is K major. \ SM100_MMA_F8F6F4_2x1SM_SS N-mode size should be a multiple of 32 between 32 and 256 when B is MN major."); diff --git a/include/cute/container/array.hpp b/include/cute/container/array.hpp index ae1a10ac..0e57859f 100644 --- a/include/cute/container/array.hpp +++ b/include/cute/container/array.hpp @@ -447,8 +447,20 @@ struct tuple_element> namespace std { +#if (__CUDACC_VER_MAJOR__ >= 13) + #include +#else +#if defined(__CUDACC_RTC__) + template + struct tuple_size; + + template + struct tuple_element; +#endif +#endif + template struct tuple_size> : CUTE_STL_NAMESPACE::integral_constant diff --git a/include/cute/container/array_subbyte.hpp b/include/cute/container/array_subbyte.hpp index 2c6a97be..2e634714 100644 --- a/include/cute/container/array_subbyte.hpp +++ b/include/cute/container/array_subbyte.hpp @@ -617,8 +617,20 @@ struct tuple_element> namespace std { +#if (__CUDACC_VER_MAJOR__ >= 13) + #include +#else +#if defined(__CUDACC_RTC__) + template + struct tuple_size; + + template + struct tuple_element; +#endif +#endif + template struct tuple_size> : CUTE_STL_NAMESPACE::integral_constant diff --git a/include/cute/container/tuple.hpp b/include/cute/container/tuple.hpp index 7a1650ed..f7cfca94 100644 --- a/include/cute/container/tuple.hpp +++ b/include/cute/container/tuple.hpp @@ -701,8 +701,20 @@ struct tuple_element> namespace std { +#if (__CUDACC_VER_MAJOR__ >= 13) + #include +#else +#if defined(__CUDACC_RTC__) + template + struct tuple_size; + + template + struct tuple_element; +#endif +#endif + template struct tuple_size> : CUTE_STL_NAMESPACE::integral_constant diff --git a/include/cute/container/type_list.hpp b/include/cute/container/type_list.hpp index 8c1373cd..62ebf9df 100644 --- a/include/cute/container/type_list.hpp +++ b/include/cute/container/type_list.hpp @@ -103,8 +103,20 @@ struct tuple_element> namespace std { +#if (__CUDACC_VER_MAJOR__ >= 13) + #include +#else +#if defined(__CUDACC_RTC__) + template + struct tuple_size; + + template + struct tuple_element; +#endif +#endif + template struct tuple_size> : CUTE_STL_NAMESPACE::integral_constant diff --git a/include/cute/numeric/arithmetic_tuple.hpp b/include/cute/numeric/arithmetic_tuple.hpp index 4ef9f790..01dde5f1 100644 --- a/include/cute/numeric/arithmetic_tuple.hpp +++ b/include/cute/numeric/arithmetic_tuple.hpp @@ -513,8 +513,20 @@ struct tuple_element> namespace std { +#if (__CUDACC_VER_MAJOR__ >= 13) + #include +#else +#if defined(__CUDACC_RTC__) + template + struct tuple_size; + + template + struct tuple_element; +#endif +#endif + template struct tuple_size> : CUTE_STL_NAMESPACE::integral_constant diff --git a/include/cutlass/arch/barrier.h b/include/cutlass/arch/barrier.h index 903ff0ac..015de3eb 100644 --- a/include/cutlass/arch/barrier.h +++ b/include/cutlass/arch/barrier.h @@ -33,6 +33,7 @@ */ #pragma once +#include "cutlass/cutlass.h" #include #include @@ -285,8 +286,8 @@ class NamedBarrier { #if CUDA_BARRIER_ENABLED asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); cutlass::arch::synclog_emit_named_barrier_arrive_and_wait(__LINE__, num_threads, barrier_id); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -295,8 +296,8 @@ class NamedBarrier { #if CUDA_BARRIER_ENABLED asm volatile("barrier.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); cutlass::arch::synclog_emit_named_barrier_arrive_and_wait(__LINE__, num_threads, barrier_id); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -305,8 +306,8 @@ class NamedBarrier { #if CUDA_BARRIER_ENABLED cutlass::arch::synclog_emit_named_barrier_arrive(__LINE__, num_threads, barrier_id); asm volatile("bar.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -315,8 +316,8 @@ class NamedBarrier { #if CUDA_BARRIER_ENABLED cutlass::arch::synclog_emit_named_barrier_arrive(__LINE__, num_threads, barrier_id); asm volatile("barrier.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads) : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -399,8 +400,8 @@ public: : "r"(arrive_count), "r"(smem_addr) : "memory"); cutlass::arch::synclog_emit_cluster_barrier_init(__LINE__, smem_addr, arrive_count); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -425,8 +426,8 @@ public: : "r"(smem_addr), "r"(phase), "r"(ticks) : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -450,8 +451,8 @@ public: : "memory"); return static_cast(waitComplete); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif return 0; } @@ -474,8 +475,8 @@ public: : "memory"); return static_cast(waitComplete); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif return 0; } @@ -498,8 +499,8 @@ public: } cutlass::arch::synclog_emit_cluster_barrier_arrive_cluster(__LINE__, smem_addr, cta_id, pred); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -516,8 +517,8 @@ public: : "r"(smem_addr) : "memory"); cutlass::arch::synclog_emit_cluster_barrier_arrive(__LINE__, smem_addr); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -532,8 +533,8 @@ public: : : "r"(smem_addr) : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } }; @@ -595,8 +596,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { : "r"(transaction_bytes), "r"(smem_addr) : "memory"); cutlass::arch::synclog_emit_cluster_transaction_barrier_arrive_and_expect_tx(__LINE__, smem_addr, transaction_bytes); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -617,8 +618,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { : : "r"(smem_addr), "r"(cta_id), "r"(pred), "r"(transaction_bytes) : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -635,8 +636,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { : "r"(transaction_bytes), "r"(smem_addr) : "memory"); cutlass::arch::synclog_emit_cluster_transaction_barrier_expect_transaction(__LINE__, smem_addr, transaction_bytes); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -657,8 +658,8 @@ struct ClusterTransactionBarrier : public ClusterBarrier { : "r"(transaction_bytes), "r"(smem_addr), "r"(pred) : "memory"); cutlass::arch::synclog_emit_cluster_transaction_barrier_complete_transaction(__LINE__, smem_addr, dst_cta_id, transaction_bytes, pred); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -717,8 +718,8 @@ void fence_barrier_init() { "}" :: : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -733,8 +734,23 @@ void fence_view_async_shared() { "}" :: : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); +#endif +} + +CUTLASS_DEVICE +void fence_view_shared() { +#if CUDA_BARRIER_ENABLED + cutlass::arch::synclog_emit_fence_view_shared(__LINE__); + asm volatile ( + "{\n\t" + "fence.release.sync_restrict::shared::cta.cluster; \n" + "}" + :: + : "memory"); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -751,8 +767,8 @@ void cpasync_barrier_arrive(uint64_t const* smem_ptr) { : "r"(smem_addr) : "memory"); cutlass::arch::synclog_emit_cpasync_barrier_arrive(__LINE__, smem_addr); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -769,8 +785,8 @@ void cpasync_barrier_arrive_noinc(uint64_t const* smem_ptr) { : "r"(smem_addr) : "memory"); cutlass::arch::synclog_emit_cpasync_barrier_arrive(__LINE__, smem_addr); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -787,8 +803,8 @@ void umma_arrive(uint64_t const* smem_ptr) { :"r"(bar_intptr) : "memory"); } -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -803,8 +819,8 @@ void umma_arrive_2x1SM(uint64_t const* smem_ptr) { :"r"(bar_intptr) : "memory"); } -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -822,8 +838,8 @@ void umma_arrive_multicast(uint64_t const* smem_ptr, uint16_t cta_mask) { :"r"(bar_intptr), "h"(cta_mask) : "memory"); } -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -841,8 +857,8 @@ void umma_arrive_multicast_2x1SM(uint64_t const* smem_ptr, uint16_t cta_mask) { :"r"(bar_intptr), "h"(cta_mask) : "memory"); } -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -861,8 +877,8 @@ void umma_arrive_multicast_no_elect(uint64_t const* smem_ptr, uint16_t cta_mask) : :"r"(bar_intptr), "r"(uint32_t(cta_mask)) : "memory"); -#elif defined(__CUDA_ARCH__) - CUTLASS_NOT_IMPLEMENTED(); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -899,8 +915,8 @@ void umma_arrive_2x1SM_sm0(uint64_t const* smem_ptr) { : "r"(bar_intptr) : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -912,8 +928,8 @@ CUTE_DEVICE static void fence_view_async_tmem_load() { "}" :: : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } @@ -925,8 +941,8 @@ CUTE_DEVICE static void fence_view_async_tmem_store() { "}" :: : "memory"); -#elif defined(__CUDA_ARCH__) - asm volatile ("brkpt;\n" ::); +#else + CUTLASS_NOT_IMPLEMENTED(); #endif } diff --git a/include/cutlass/arch/mma_sm100.h b/include/cutlass/arch/mma_sm100.h index 357a690b..92af9078 100644 --- a/include/cutlass/arch/mma_sm100.h +++ b/include/cutlass/arch/mma_sm100.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/arch/mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/mma_sm70.h b/include/cutlass/arch/mma_sm70.h index 749e71e4..16806493 100644 --- a/include/cutlass/arch/mma_sm70.h +++ b/include/cutlass/arch/mma_sm70.h @@ -33,7 +33,9 @@ */ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/mma_sm75.h b/include/cutlass/arch/mma_sm75.h index 0086d13c..3b5e51df 100644 --- a/include/cutlass/arch/mma_sm75.h +++ b/include/cutlass/arch/mma_sm75.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/arch/wmma.h" diff --git a/include/cutlass/arch/mma_sm80.h b/include/cutlass/arch/mma_sm80.h index 0c56dfc3..48536b96 100644 --- a/include/cutlass/arch/mma_sm80.h +++ b/include/cutlass/arch/mma_sm80.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/mma_sm89.h b/include/cutlass/arch/mma_sm89.h index 6b3a617b..493442ad 100644 --- a/include/cutlass/arch/mma_sm89.h +++ b/include/cutlass/arch/mma_sm89.h @@ -35,7 +35,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/mma_sm90.h b/include/cutlass/arch/mma_sm90.h index b5f018a0..5bcbe270 100644 --- a/include/cutlass/arch/mma_sm90.h +++ b/include/cutlass/arch/mma_sm90.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/mma_sparse_sm80.h b/include/cutlass/arch/mma_sparse_sm80.h index e0433bcd..c39a5695 100644 --- a/include/cutlass/arch/mma_sparse_sm80.h +++ b/include/cutlass/arch/mma_sparse_sm80.h @@ -35,7 +35,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/mma_sparse_sm89.h b/include/cutlass/arch/mma_sparse_sm89.h index d6d5d241..00a0a342 100644 --- a/include/cutlass/arch/mma_sparse_sm89.h +++ b/include/cutlass/arch/mma_sparse_sm89.h @@ -35,7 +35,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "mma.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/arch/synclog.hpp b/include/cutlass/arch/synclog.hpp index 9ca8337e..ffa7baee 100644 --- a/include/cutlass/arch/synclog.hpp +++ b/include/cutlass/arch/synclog.hpp @@ -156,6 +156,10 @@ constexpr bool synclog_enable_fence_view_async_shared = true; constexpr uint32_t synclog_header_fence_view_async_shared = 17; constexpr uint32_t synclog_length_fence_view_async_shared = synclog_length_prefix + 0; +constexpr bool synclog_enable_fence_view_shared = true; +constexpr uint32_t synclog_header_fence_view_shared = 39; +constexpr uint32_t synclog_length_fence_view_shared = synclog_length_prefix + 0; + constexpr bool synclog_enable_cp_async_wait = true; constexpr uint32_t synclog_header_cp_async_wait = 18; constexpr uint32_t synclog_length_cp_async_wait = synclog_length_prefix + 1; @@ -637,6 +641,19 @@ void synclog_emit_fence_view_async_shared(uint32_t line) { #endif // defined(CUTLASS_ENABLE_SYNCLOG) } +CUTLASS_DEVICE +void synclog_emit_fence_view_shared(uint32_t line) { + #if defined(CUTLASS_ENABLE_SYNCLOG) + if constexpr (!synclog_enable_fence_view_shared) return; + if (!synclog_condition_emit()) return; + uint32_t* to = synclog_alloc(synclog_length_fence_view_shared); + if (to == nullptr) return; + synclog_emit_prefix(to, synclog_header_fence_view_shared, line); + #else + CUTLASS_UNUSED(line); + #endif // defined(CUTLASS_ENABLE_SYNCLOG) +} + CUTLASS_DEVICE void synclog_emit_cp_async_wait( uint32_t line, @@ -1091,6 +1108,14 @@ void synclog_print() { continue; } } + if constexpr (synclog_enable_fence_view_shared) { + if (header == synclog_header_fence_view_shared) { + synclog_print_prefix("fence_view_shared", at); + at += synclog_length_fence_view_shared; + printf("\n"); + continue; + } + } if constexpr (synclog_enable_cp_async_wait) { if (header == synclog_header_cp_async_wait) { synclog_print_prefix("cp_async_wait", at); diff --git a/include/cutlass/arch/wmma_sm70.h b/include/cutlass/arch/wmma_sm70.h index 81a54b45..05efaf63 100644 --- a/include/cutlass/arch/wmma_sm70.h +++ b/include/cutlass/arch/wmma_sm70.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/layout/matrix.h" //////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/arch/wmma_sm72.h b/include/cutlass/arch/wmma_sm72.h index 85fa70bb..995f8250 100644 --- a/include/cutlass/arch/wmma_sm72.h +++ b/include/cutlass/arch/wmma_sm72.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/layout/matrix.h" //////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/arch/wmma_sm75.h b/include/cutlass/arch/wmma_sm75.h index 6bee3c64..f67a519b 100644 --- a/include/cutlass/arch/wmma_sm75.h +++ b/include/cutlass/arch/wmma_sm75.h @@ -34,7 +34,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/layout/matrix.h" //////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/array.h b/include/cutlass/array.h index 0e96486f..0ee2591f 100644 --- a/include/cutlass/array.h +++ b/include/cutlass/array.h @@ -718,6 +718,24 @@ struct maximum_absolute_value_reduction, PropogateNaN> { } }; +template +struct maximum_absolute_value_zero_mantissa_reduction, PropagateNaN> { + + CUTLASS_HOST_DEVICE + T operator() (T const& scalar, cutlass::Array const& rhs) const { + + T result = scalar; + maximum_absolute_value_zero_mantissa_reduction scalar_op; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result = scalar_op(result, rhs[i]); + } + + return result; + } +}; + template struct scale> { T const scaling_factor_; diff --git a/include/cutlass/complex.h b/include/cutlass/complex.h index 4449df67..adfe103d 100644 --- a/include/cutlass/complex.h +++ b/include/cutlass/complex.h @@ -813,6 +813,37 @@ struct atomic_add> { } }; +// Maximal exponent reduction for zero-mantissa scaling factors: complex number uses its largest cartesian norm not abs +template +struct maximum_cartesian_norm_zero_mantissa_reduction { + using T = typename TC::value_type; + + CUTLASS_HOST_DEVICE + T operator()(T const &lhs, cutlass::complex const &rhs) const { + maximum_absolute_value_zero_mantissa_reduction red_op; + + return red_op(red_op(lhs, rhs.real()), rhs.imag()); + } +}; + +template +struct maximum_cartesian_norm_zero_mantissa_reduction, PropagateNaN> { + using T = typename TC::value_type; + + CUTLASS_HOST_DEVICE + T operator() (T const& scalar, cutlass::Array const& rhs) const { + + T result = scalar; + maximum_cartesian_norm_zero_mantissa_reduction scalar_op; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result = scalar_op(result, rhs[i]); + } + + return result; + } +}; ////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/cutlass.h b/include/cutlass/cutlass.h index 2442f310..0c7cb201 100644 --- a/include/cutlass/cutlass.h +++ b/include/cutlass/cutlass.h @@ -37,11 +37,7 @@ #include "cutlass/detail/helper_macros.hpp" -#if (__CUDACC_VER_MAJOR__ >= 13) - #define CUDA_STD_HEADER(header) -#else - #define CUDA_STD_HEADER(header) -#endif +#define CUDA_STD_HEADER(header) //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/detail/collective/mixed_input_utils.hpp b/include/cutlass/detail/collective/mixed_input_utils.hpp index 3ea85cdb..7c5ee0b3 100644 --- a/include/cutlass/detail/collective/mixed_input_utils.hpp +++ b/include/cutlass/detail/collective/mixed_input_utils.hpp @@ -862,7 +862,8 @@ public: } else if constexpr (UseScaleLookupTable) { constexpr int num_elements = decltype(size(src))::value; - static_assert(is_same_v, "Lookup table only supports int4 being the quant type now."); + static_assert(is_same_v || is_same_v, + "Lookup table supports int4b_t (Two's Complement) and float_e2m1_t (E2M1/FP4) quant types."); static_assert(sizeof_bits_v == 64, "Lookup table only supports 8 8bit scale values now."); static_assert(num_elements % 4 == 0 && num_elements >= 4, "Lookup table requires a vector size of 4x when converting."); @@ -885,15 +886,31 @@ public: { auto&& scale_neg_ = reinterpret_cast const&>(scales_neg_vm_(i)); auto&& scale_pos_ = reinterpret_cast &>(scales_pos_vm_(i)); - constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa; - asm volatile( - "{\n" - " lop3 .b32 %0, %2, %4, %5, %6;\n" \ - " xor .b32 %1, %3, %5; \n" \ - "}\n" - : "=r"(scale_pos_[0]), "=r"(scale_pos_[1]) - : "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut) - ); + + // Accept CUTLASS pseudo-FP as well + if constexpr (cutlass::platform::is_floating_point::value || + cute::is_same_v) { + // E2M1 (FP4): Sign-magnitude encoding - simple sign flip with two XORs + asm volatile( + "{\n" + " xor .b32 %0, %2, %4;\n" \ + " xor .b32 %1, %3, %4;\n" \ + "}\n" + : "=r"(scale_pos_[0]), "=r"(scale_pos_[1]) + : "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0x80808080) + ); + } else { + // INT4: Two's complement encoding - reorder and sign flip with lop3 + constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa; + asm volatile( + "{\n" + " lop3 .b32 %0, %2, %4, %5, %6;\n" \ + " xor .b32 %1, %3, %5; \n" \ + "}\n" + : "=r"(scale_pos_[0]), "=r"(scale_pos_[1]) + : "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut) + ); + } } } CUTLASS_PRAGMA_UNROLL diff --git a/include/cutlass/detail/sm100_blockscaled_layout.hpp b/include/cutlass/detail/sm100_blockscaled_layout.hpp index bdbd8bfe..e44dd887 100644 --- a/include/cutlass/detail/sm100_blockscaled_layout.hpp +++ b/include/cutlass/detail/sm100_blockscaled_layout.hpp @@ -88,9 +88,14 @@ struct Sm1xxBlockScaledConfig { CUTE_HOST_DEVICE static constexpr auto tile_atom_to_shape_SFA(ProblemShape problem_shape, LayoutSFA layout_sfa = LayoutSFA{}) { - auto problem_shape_MNKL = append<4>(problem_shape, 1); - auto [M, N, K, L] = problem_shape_MNKL; - return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{}); + if constexpr (rank(ProblemShape{}) == 3) { + auto [M, N, K] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(M,K), Step<_2,_1>{}); + } + else { + auto [M, N, K, L] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{}); + } } // The following function is provided for user fill dynamic problem size to the layout_SFB. @@ -98,9 +103,14 @@ struct Sm1xxBlockScaledConfig { CUTE_HOST_DEVICE static constexpr auto tile_atom_to_shape_SFB(ProblemShape problem_shape, LayoutSFB layout_sfb = LayoutSFB{}) { - auto problem_shape_MNKL = append<4>(problem_shape, 1); - auto [M, N, K, L] = problem_shape_MNKL; - return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{}); + if constexpr (rank(ProblemShape{}) == 3) { + auto [M, N, K] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(N,K), Step<_2,_1>{}); + } + else { + auto [M, N, K, L] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{}); + } } template diff --git a/include/cutlass/detail/sm103_blockscaled_layout.hpp b/include/cutlass/detail/sm103_blockscaled_layout.hpp index db89f3a1..800d0191 100644 --- a/include/cutlass/detail/sm103_blockscaled_layout.hpp +++ b/include/cutlass/detail/sm103_blockscaled_layout.hpp @@ -87,9 +87,14 @@ struct Sm103BlockScaledConfig { CUTE_HOST_DEVICE static constexpr auto tile_atom_to_shape_SFA(ProblemShape problem_shape, LayoutSFA layout_sfa = LayoutSFA{}) { - auto problem_shape_MNKL = append<4>(problem_shape, 1); - auto [M, N, K, L] = problem_shape_MNKL; - return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{}); + if constexpr (rank(ProblemShape{}) == 3) { + auto [M, N, K] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(M,K), Step<_2,_1>{}); + } + else { + auto [M, N, K, L] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(M,K,L), Step<_2,_1,_3>{}); + } } // The following function is provided for user fill dynamic problem size to the layout_SFB. @@ -97,9 +102,14 @@ struct Sm103BlockScaledConfig { CUTE_HOST_DEVICE static constexpr auto tile_atom_to_shape_SFB(ProblemShape problem_shape, LayoutSFB layout_sfb = LayoutSFB{}) { - auto problem_shape_MNKL = append<4>(problem_shape, 1); - auto [M, N, K, L] = problem_shape_MNKL; - return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{}); + if constexpr (rank(ProblemShape{}) == 3) { + auto [M, N, K] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(N,K), Step<_2,_1>{}); + } + else { + auto [M, N, K, L] = problem_shape; + return tile_to_shape(SfAtom{}, make_shape(N,K,L), Step<_2,_1,_3>{}); + } } }; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/collective/builders/sm100_builder.inl b/include/cutlass/epilogue/collective/builders/sm100_builder.inl index 211bdfe2..e5fafd52 100644 --- a/include/cutlass/epilogue/collective/builders/sm100_builder.inl +++ b/include/cutlass/epilogue/collective/builders/sm100_builder.inl @@ -1121,6 +1121,23 @@ sm100_dense_dispatch_policy() { else if constexpr (is_base_of_v || is_base_of_v) { return Sm100NoSmemWarpSpecialized{}; } + else if constexpr (is_same_v || + is_same_v) { + return Sm100PtrArrayPlanarComplexNoSmemWarpSpecialized{}; + } + else if constexpr (is_same_v || + is_same_v) { + constexpr bool ReuseSmem_ = (sizeof_bits_v == sizeof_bits_v); // limited smem reuse support for planar complex for now + constexpr int StagesC_ = ReuseSmem_ ? cute::max(cute::min(EpiTiles, 4), StagesD+1) : cute::min(EpiTiles, 4); + constexpr bool DelayTmaStore_ = false; // TMA store delay complicates tensormap updates for Ptr-Array GEMMs + return Sm100PtrArrayPlanarComplexTmaWarpSpecialized{}; + } + else if constexpr (is_same_v || + is_same_v) { + constexpr bool ReuseSmem_ = (sizeof_bits_v == sizeof_bits_v); // limited smem reuse support for planar complex for now + constexpr int StagesC_ = ReuseSmem_ ? cute::max(cute::min(EpiTiles, 4), StagesD+1) : cute::min(EpiTiles, 4); + return Sm100PlanarComplexTmaWarpSpecialized{}; + } else if constexpr (is_same_v || is_same_v) { constexpr bool DelayTmaStore_ = false; // TMA store delay complicates tensormap updates for Ptr-Array GEMMs @@ -1235,6 +1252,16 @@ private: static constexpr auto fusion_callbacks() { + if constexpr (is_same_v || + is_same_v || + is_same_v || + is_same_v) { + static_assert(IsDefaultFusionOp::value, "unsupported schedule + fusion"); + constexpr thread::ScaleType::Kind ScaleType = DisableSource ? thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default; + return thread::LinearCombinationPlanarComplex< + ElementD, FragmentSize, ElementAccumulator, ElementCompute, FusionOp::RoundStyle, ScaleType>({}); + } + else { return typename CallbacksBuilder< decltype(dispatch_policy()), @@ -1514,6 +1541,12 @@ private: return thread::LinearCombination< ElementD, 1, ElementAccumulator, ElementCompute, ScaleType, FusionOp::RoundStyle, ElementC>({}); } + else if constexpr (is_same_v || + is_same_v) { + static_assert(IsDefaultFusionOp::value, "unsupported schedule + fusion"); + return thread::LinearCombinationPlanarComplex< + ElementD, FragmentSize, ElementAccumulator, ElementCompute, FusionOp::RoundStyle, ScaleType>({}); + } else { return typename detail::CallbacksBuilder< DispatchPolicy, @@ -1780,6 +1813,7 @@ struct CollectiveBuilder< CopyAtomR2G, Schedule>; }; + /////////////////////////////////////////////////////////////////////////////// } // namespace cutlass::epilogue::collective diff --git a/include/cutlass/epilogue/collective/detail.hpp b/include/cutlass/epilogue/collective/detail.hpp index 75c14c4b..05078bb3 100644 --- a/include/cutlass/epilogue/collective/detail.hpp +++ b/include/cutlass/epilogue/collective/detail.hpp @@ -736,7 +736,9 @@ public: // Wait for mma warp to fill tmem buffer with accumulator results acc_pipeline.consumer_wait(acc_pipe_consumer_state); - auto [acc_state_next] = (*this).template operator()( + auto [acc_state_next, load_state_next] = (*this).template operator()( + load_pipeline, + load_pipe_consumer_state, acc_pipeline, acc_pipe_consumer_state, problem_shape_mnkl, @@ -746,10 +748,9 @@ public: shared_tensors); // Let mma warp know tmem buffer is consumed and empty - ++load_pipe_consumer_state; ++store_pipe_producer_state; - return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_state_next); + return cute::make_tuple(load_state_next, store_pipe_producer_state, acc_state_next); } // FastF32 API @@ -857,7 +858,9 @@ public: TensorMap tensormap ) { - auto [acc_state_next] = (*this).template operator()( + auto [acc_state_next, load_state_next] = (*this).template operator()( + load_pipeline, + load_pipe_consumer_state, acc_pipeline, acc_pipe_consumer_state, problem_shape_mnkl, @@ -867,10 +870,9 @@ public: shared_tensors); // Let mma warp know tmem buffer is consumed and empty - ++load_pipe_consumer_state; ++store_pipe_producer_state; - return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_state_next); + return cute::make_tuple(load_state_next, store_pipe_producer_state, acc_state_next); } template diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp index fe680480..176e7ad9 100644 --- a/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_nosmem.hpp @@ -169,6 +169,8 @@ public: template< bool ReuseTmem = false, + class LoadPipeline, + class LoadPipelineState, class AccumulatorPipeline, class AccumulatorPipelineState, class ProblemShapeMNKL, @@ -178,6 +180,8 @@ public: > CUTLASS_DEVICE auto operator()( + [[maybe_unused]]LoadPipeline load_pipeline, + [[maybe_unused]]LoadPipelineState load_pipe_consumer_state, AccumulatorPipeline acc_pipeline, AccumulatorPipelineState acc_pipe_consumer_state, ProblemShapeMNKL problem_shape_mnkl, @@ -357,7 +361,7 @@ public: copy_if(tDpD, tTR_rD_src, tR2G_rD_dst); } - return cute::make_tuple(acc_pipe_consumer_state); + return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state); } // API with Global Accumulator in registers for FastFP32 (emulated MMA) kernels. @@ -609,6 +613,8 @@ public: template< bool ReuseTmem = false, + class LoadPipeline, + class LoadPipelineState, class AccumulatorPipeline, class AccumulatorPipelineState, class ProblemShapeMNKL, @@ -618,6 +624,8 @@ public: > CUTLASS_DEVICE auto operator()( + [[maybe_unused]]LoadPipeline load_pipeline, + [[maybe_unused]]LoadPipelineState load_pipe_consumer_state, AccumulatorPipeline acc_pipeline, AccumulatorPipelineState acc_pipe_consumer_state, ProblemShapeMNKL problem_shape_mnkl, @@ -904,7 +912,7 @@ public: // auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks(cst_args); epi_loop_fn(cst_callbacks, is_accumulator_needed); - return cute::make_tuple(acc_pipe_consumer_state); + return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state); } }; diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_nosmem.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_nosmem.hpp new file mode 100644 index 00000000..d3cb7157 --- /dev/null +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_nosmem.hpp @@ -0,0 +1,345 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Functor performing elementwise operations used by Ptr-Array Planar Complex Gemm epilogue. +*/ + + + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/detail.hpp" +#include "cutlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "cute/tensor.hpp" +#include "cute/numeric/numeric_types.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace epilogue { +namespace collective { + +/// Applies an element wise operation to all elements within the fragment +/// and writes it out to destination storage. +template < + class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class ThreadEpilogueOp_, + class CopyOpT2R_ +> +class CollectiveEpilogue< + Sm100PtrArrayPlanarComplexNoSmem, + EpilogueTile_, + ElementC_, + StrideC_, + ElementD_, + StrideD_, + ThreadEpilogueOp_, + CopyOpT2R_ +> { +public: + // + // Type Aliases + // + using DispatchPolicy = Sm100PtrArrayPlanarComplexNoSmem; + using EpilogueTile = EpilogueTile_; + // derived types of output thread level operator + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementOutput = typename ThreadEpilogueOp::ElementOutput; + using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + using ElementScalar = typename ThreadEpilogueOp::ElementScalar; + using ElementC = ElementC_; + using StrideC = StrideC_; + using InternalStrideC = cute::remove_pointer_t; + using ElementD = ElementD_; + using StrideD = StrideD_; + using InternalStrideD = cute::remove_pointer_t; + using CopyOpT2R = CopyOpT2R_; + + using GmemTiledCopyC = void; + using GmemTiledCopyD = void; + + constexpr static int ThreadCount = 128; + constexpr static uint32_t TmaTransactionBytes = 0; + constexpr static int FragmentSize = ThreadEpilogueOp::kCount; + + struct SharedStorage { + struct TensorStorage { }; + struct TensorMapStorage { }; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using TensorMapStorage = typename SharedStorage::TensorMapStorage; + + // Planar complex kernels have two accumulator copies for the real and imaginary tensors. + constexpr static int NumAccumulatorMtxs = 2; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread{}; + ElementC const** ptr_C_real = nullptr; + StrideC dC_real{}; + ElementC const** ptr_C_imag = nullptr; + StrideC dC_imag{}; + ElementD** ptr_D_real = nullptr; + StrideD dD_real{}; + ElementD** ptr_D_imag = nullptr; + StrideD dD_imag{}; + }; + + // Device side epilogue params + using Params = Arguments; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + [[maybe_unused]] ProblemShape const& problem_shape, + Arguments const& args, + [[maybe_unused]] void* workspace) { + return args; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) { + return 0; + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { + return cutlass::Status::kSuccess; + } + + template + static bool + can_implement( + [[maybe_unused]] ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + return true; + } + + CUTLASS_HOST_DEVICE + CollectiveEpilogue(Params const& params, SharedStorage&) : params(params) { }; + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class AccEngine, class AccLayout + > + CUTLASS_DEVICE void + operator()( + ProblemShapeMNKL problem_shape_mnkl, + TileShapeMNK cta_tile_shape_mnk, + TileCoordMNKL cta_coord_mnkl, + cute::Tensor const& accumulators, // (MMA,MMA_M,MMA_N) + [[maybe_unused]] SharedStorage&) { + + using namespace cute; + using X = Underscore; + + static_assert(is_tmem::value, "Accumulator must be TMEM resident."); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4"); + + // Separate out problem shape for convenience + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + // Batches are managed by using appropriate pointers to C and D matrices + const int32_t mock_L = 1; + // Slice to get the tile this CTA is responsible for + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + const int32_t mock_l_coord = 0; + + auto problem_shape_mnl = make_shape(M,N,mock_L); + auto cta_coord_mnl = make_shape(m_coord, n_coord, mock_l_coord); + auto cta_tiler = take<0,2>(cta_tile_shape_mnk); + + // Represent the full output tensor, slice to get the tile this CTA is responsible for + Tensor mC_real = make_tensor(make_gmem_ptr(params.ptr_C_real[l_coord]), problem_shape_mnl, append<3>(params.dC_real,_0{})); // (M,N,L) + Tensor mC_imag = make_tensor(make_gmem_ptr(params.ptr_C_imag[l_coord]), problem_shape_mnl, append<3>(params.dC_imag,_0{})); // (M,N,L) + + Tensor mD_real = make_tensor(make_gmem_ptr(params.ptr_D_real[l_coord]), problem_shape_mnl, append<3>(params.dD_real,_0{})); // (M,N,L) + Tensor mD_imag = make_tensor(make_gmem_ptr(params.ptr_D_imag[l_coord]), problem_shape_mnl, append<3>(params.dD_imag,_0{})); // (M,N,L) + + Tensor gC_real = local_tile(mC_real, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) + Tensor gC_imag = local_tile(mC_imag, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) + + Tensor gD_real = local_tile(mD_real, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) + Tensor gD_imag = local_tile(mD_imag, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) + + // Partition source and destination tiles according to tmem copy T2R partitioning (tTR_) + auto tiled_t2r = make_tmem_copy(CopyOpT2R{}, tensor<0>(accumulators)); + auto thread_t2r = tiled_t2r.get_slice(threadIdx.x % size(tiled_t2r)); + Tensor tTR_gC_real = thread_t2r.partition_D(gC_real); // (T2R,T2R_M,T2R_N) + Tensor tTR_gC_imag = thread_t2r.partition_D(gC_imag); // (T2R,T2R_M,T2R_N) + + Tensor tTR_gD_real = thread_t2r.partition_D(gD_real); // (T2R,T2R_M,T2R_N) + Tensor tTR_gD_imag = thread_t2r.partition_D(gD_imag); // (T2R,T2R_M,T2R_N) + + Tensor tTR_rAcc = make_tensor(append(shape(tTR_gD_real), Int{})); // (T2R,T2R_M,T2R_N,2) + Tensor tTR_rD = make_tensor(append(shape(tTR_gD_real), Int{})); // (T2R,T2R_M,T2R_N,2) + + Tensor tTR_rAcc_frg = recast>(coalesce(tTR_rAcc)); // (EPI_V) + Tensor tTR_rD_frg = recast>(coalesce(tTR_rD)); // (EPI_V) + + Tensor coordD = make_identity_tensor(problem_shape_mnl); // (M,N,L) -> (m,n,l) + Tensor cD = local_tile(coordD, cta_tiler, cta_coord_mnl); // (CTA_M,CTA_N) -> (m,n,l) + Tensor tTR_cD = thread_t2r.partition_D(cD); // (T2R,T2R_M,T2R_N) -> (m,n,l) + + // 1. Load accumulators into register from tmem + auto accumulators_real = accumulators(_,_,_,0); + auto accumulators_imag = accumulators(_,_,_,1); + Tensor tAcc_real = accumulators_real(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + Tensor tAcc_imag = accumulators_imag(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + Tensor tTR_tAcc_real = thread_t2r.partition_S(tAcc_real); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + Tensor tTR_tAcc_imag = thread_t2r.partition_S(tAcc_imag); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + + // tmem -> rmem + copy(tiled_t2r, tTR_tAcc_real, tTR_rAcc(_,_,_,0)); + copy(tiled_t2r, tTR_tAcc_imag, tTR_rAcc(_,_,_,1)); + + // 2. Apply element-wise operation and store to gmem + ThreadEpilogueOp epilogue_op{params.thread}; + // source is needed + if (epilogue_op.is_source_needed()) { + Tensor tTR_rC = make_tensor(append(shape(tTR_gC_real), Int{})); // (T2R,T2R_M,T2R_N,2) + Tensor tTR_rC_frg = recast>(coalesce(tTR_rC)); // (EPI_V) + + auto tTR_rC_real = tTR_rC(_,_,_,0); + auto tTR_rC_imag = tTR_rC(_,_,_,1); + + for( int i = 0; i < size(tTR_gC_real); ++i) { + if (elem_less(tTR_cD(i), problem_shape_mnl)) { + tTR_rC_real(i) = tTR_gC_real(i); + tTR_rC_imag(i) = tTR_gC_imag(i); + } + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc_frg); ++i) { + tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i), tTR_rC_frg(i)); + } + } + // source is not needed, avoid load + else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc_frg); ++i) { + tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i)); + } + } + + auto tTR_rD_real = tTR_rD(_,_,_,0); + auto tTR_rD_imag = tTR_rD(_,_,_,1); + + for( int i = 0; i < size(tTR_gD_real); ++i) { + if (elem_less(tTR_cD(i), problem_shape_mnl)) { + tTR_gD_real(i) = tTR_rD_real(i); + tTR_gD_imag(i) = tTR_rD_imag(i); + } + } + } + +protected: + Params const& params; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// For sm100 kernels requiring warp specialized epilogues +template < + class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class ThreadEpilogueOp_, + class CopyOpT2R_, + class AlignmentC, + class AlignmentD +> +class CollectiveEpilogue< + Sm100PtrArrayPlanarComplexNoSmemWarpSpecialized, + EpilogueTile_, + ElementC_, + StrideC_, + ElementD_, + StrideD_, + ThreadEpilogueOp_, + CopyOpT2R_, + AlignmentC, + AlignmentD +> : public detail::Sm100TmaWarpSpecializedAdapter> +{ +public: + // ctor inheritance + using detail::Sm100TmaWarpSpecializedAdapter>::Sm100TmaWarpSpecializedAdapter; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + + + +} // namespace collective +} // namespace epilogue +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_tma_warpspecialized.hpp new file mode 100644 index 00000000..8c0626a4 --- /dev/null +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_planar_complex_tma_warpspecialized.hpp @@ -0,0 +1,1161 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Functor performing elementwise operations used by Ptr-Array Planar Complex Gemm epilogues. +*/ + + + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/arch/barrier.h" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/detail.hpp" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/epilogue/thread/linear_combination_planar_complex.h" +#include "cutlass/detail/layout.hpp" +#include "cutlass/trace.h" + +#include "cute/tensor.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::epilogue::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int StagesC_, + int StagesD_, + int FragmentSize_, + bool ReuseSmemC_, + bool DelayTmaStore_, + class CtaTileShape_, // (CTA_M,CTA_N,CTA_K, optional: Tile_L) + class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class ThreadEpilogueOp_, + class CopyOpT2R_, + class CopyOpG2S_, + class SmemLayoutAtomC_, + class CopyOpS2R_, + class CopyOpS2G_, + class SmemLayoutAtomD_, + class CopyOpR2S_, + class CopyOpR2R_ +> +class CollectiveEpilogue< + Sm100PtrArrayPlanarComplexTmaWarpSpecialized, + CtaTileShape_, + EpilogueTile_, + ElementC_, + StrideC_, + ElementD_, + StrideD_, + ThreadEpilogueOp_, + CopyOpT2R_, + CopyOpG2S_, + SmemLayoutAtomC_, + CopyOpS2R_, + CopyOpS2G_, + SmemLayoutAtomD_, + CopyOpR2S_, + CopyOpR2R_ +> { +public: + using DispatchPolicy = Sm100PtrArrayPlanarComplexTmaWarpSpecialized; + using CtaTileShape = CtaTileShape_; + using EpilogueTile = EpilogueTile_; + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementC = ElementC_; + using StrideC = StrideC_; + using InternalStrideC = cute::remove_pointer_t; + using ElementD = ElementD_; + using StrideD = StrideD_; + using InternalStrideD = cute::remove_pointer_t; + using CopyOpT2R = CopyOpT2R_; + using CopyOpG2S = CopyOpG2S_; + using SmemLayoutAtomC = SmemLayoutAtomC_; + using CopyOpS2R = CopyOpS2R_; + using CopyOpS2G = CopyOpS2G_; + using SmemLayoutAtomD = SmemLayoutAtomD_; + using CopyOpR2S = CopyOpR2S_; + using CopyOpR2R = CopyOpR2R_; + + using GmemTiledCopyC = CopyOpG2S; + using GmemTiledCopyD = CopyOpS2G; + + constexpr static int ThreadCount = 128; + + static_assert(!is_layout::value && is_tuple::value, "EpilogueTile must be a cute::Tile or cute::Shape"); + static_assert(rank(EpilogueTile{}) == 2, "EpilogueTile must be rank-2: [EPI_TILE_M, EPI_TILE_N]"); + + // Epilog assumes a max scheduler pipe count to calculate the number of asynchronous tma update buffer they need. + constexpr static uint32_t NumMaxSchedulerPipelineStageCount = 8; + +private: + using SmemElementD = typename cutlass::detail::get_unpacked_element_type::type; + using SmemElementC = typename cutlass::detail::get_unpacked_element_type,ElementD,ElementC>>::type; // prevents void ref breakages + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; + constexpr static bool ReuseSmemC = ReuseSmemC_; + constexpr static bool DelayTmaStore = DelayTmaStore_; + constexpr static bool is_source_supported = ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default; + + constexpr static bool is_m_major_C = detail::is_m_major(); + constexpr static bool is_m_major_D = detail::is_m_major(); + + // Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously. + // This should be larger than the total number of TMA requests inflight (from update to issued to returned). + // This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies). + constexpr static uint32_t NumTmaDescriptorsPerSm = NumMaxSchedulerPipelineStageCount + std::max(StagesC, (ReuseSmemC ? StagesC : StagesD)) + 2; + + using SmemLayoutC = decltype(tile_to_shape( + SmemLayoutAtomC{}, + make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int{}), + cute::conditional_t, Step<_1,_2,_3>>{} )); + using SmemLayoutD = decltype(tile_to_shape( + SmemLayoutAtomD{}, + make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int{}), + cute::conditional_t, Step<_1,_2,_3>>{} )); + + constexpr static bool support_smem_reuse = is_source_supported && StagesD <= StagesC + && cosize(take<0,2>(SmemLayoutC{})) == cosize(take<0,2>(SmemLayoutD{})); + static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met"); + + constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{}); + constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{}); + +public : + struct TensorStorageWithC { + alignas(SmemAlignmentC) cute::ArrayEngine> smem_C_real; + alignas(SmemAlignmentC) cute::ArrayEngine> smem_C_imag; + + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_real; + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_imag; + }; + + struct TensorStorageWithoutC { + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_real; + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_imag; + }; + +public: + // TMA pipeline for loading C + using LoadPipeline = cutlass::PipelineTransactionAsync; + using LoadPipelineState = cutlass::PipelineState; + constexpr static uint32_t TmaTransactionBytes = + 2 * ((size(take<0,2>(SmemLayoutC{})) * static_cast(sizeof_bits::value)) / 8); + + // TMA pipeline for storing D + using StorePipeline = cute::conditional_t, + cutlass::PipelineTmaStore>; + using StorePipelineState = cutlass::PipelineState; + + struct SharedStorage { + using TensorStorage = + cute::conditional_t; + TensorStorage tensors; + + struct TensorMapStorage : cute::aligned_struct<128, _0> { + cute::TmaDescriptor smem_tensormap_C_real; + cute::TmaDescriptor smem_tensormap_C_imag; + + cute::TmaDescriptor smem_tensormap_D_real; + cute::TmaDescriptor smem_tensormap_D_imag; + } tensormaps; + + using PipelineStorage = typename LoadPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using TensorMapStorage = typename SharedStorage::TensorMapStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Planar complex kernels have two accumulator copies for the real and imaginary tensors. + constexpr static int NumAccumulatorMtxs = 2; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread{}; + ElementC const** ptr_C_real = nullptr; + StrideC dC_real{}; + ElementC const** ptr_C_imag = nullptr; + StrideC dC_imag{}; + ElementD** ptr_D_real = nullptr; + StrideD dD_real{}; + ElementD** ptr_D_imag = nullptr; + StrideD dD_imag{}; + }; + + // Device side epilogue params + struct Params { + using TensorShapeC = decltype(repeat_like(append<3>(StrideC{}, _1{}), int32_t(0))); + using TensorShapeD = decltype(repeat_like(append<3>(StrideD{}, _1{}), int32_t(0))); + using TMA_C = decltype(make_tma_copy( + CopyOpG2S{}, + make_tensor( + make_gmem_ptr(static_cast,ElementD,ElementC> const*>(nullptr)), + TensorShapeC{}, + append<3>(StrideC{}, _0{})), + take<0,2>(SmemLayoutC{}), + EpilogueTile{}, + _1{})); + using TMA_D = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor( + make_gmem_ptr(static_cast(nullptr)), + TensorShapeD{}, + append<3>(StrideD{}, _0{})), + take<0,2>(SmemLayoutC{}), + EpilogueTile{}, + _1{})); + + typename ThreadEpilogueOp::Params thread{}; + TMA_C tma_load_c_real; + TMA_C tma_load_c_imag; + TMA_D tma_store_d_real; + TMA_D tma_store_d_imag; + cute::TmaDescriptor* tensormaps; + ElementC const** ptr_C_real; + ElementC const** ptr_C_imag; + ElementD** ptr_D_real; + ElementD** ptr_D_imag; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + void* workspace) { + // Tensor shapes for Ptr-Array are initialized correctly here. + auto [M,N,K,mock_L] = problem_shape.get_host_problem_shape(0); + // Batches/Groups are managed by using appropriate pointers to input matrices + mock_L = 1; + + typename Params::TMA_C tma_load_c_real{}; + typename Params::TMA_C tma_load_c_imag{}; + if constexpr (not cute::is_void_v) { + // Tensor pointers will be fixed before the first access + ElementC const* ptr_C_real_first_batch = nullptr; + ElementC const* ptr_C_imag_first_batch = nullptr; + Tensor tensor_c_real = make_tensor(ptr_C_real_first_batch, make_layout(make_shape(M,N,mock_L), append<3>(args.dC_real, _0{}))); + Tensor tensor_c_imag = make_tensor(ptr_C_imag_first_batch, make_layout(make_shape(M,N,mock_L), append<3>(args.dC_imag, _0{}))); + + tma_load_c_real = make_tma_copy(CopyOpG2S{}, tensor_c_real, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{}); + tma_load_c_imag = make_tma_copy(CopyOpG2S{}, tensor_c_imag, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{}); + } + + // Tensor pointers will be fixed before the first access + ElementD* ptr_D_real_first_batch = nullptr; + ElementD* ptr_D_imag_first_batch = nullptr; + Tensor tensor_d_real = make_tensor(ptr_D_real_first_batch, make_layout(make_shape(M,N,mock_L), append<3>(args.dD_real, _0{}))); + Tensor tensor_d_imag = make_tensor(ptr_D_imag_first_batch, make_layout(make_shape(M,N,mock_L), append<3>(args.dD_imag, _0{}))); + + typename Params::TMA_D tma_store_d_real = + make_tma_copy(CopyOpS2G{}, tensor_d_real, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{}); + typename Params::TMA_D tma_store_d_imag = + make_tma_copy(CopyOpS2G{}, tensor_d_imag, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{}); + + return { + args.thread, + tma_load_c_real, + tma_load_c_imag, + tma_store_d_real, + tma_store_d_imag, + static_cast(workspace), + args.ptr_C_real, + args.ptr_C_imag, + args.ptr_D_real, + args.ptr_D_imag + }; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) { + constexpr uint32_t NumTensors = cute::is_void_v ? 2 : 4; + constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor); + + return (NumTensors * SizeOfCuTensorMap * sm_count * NumTmaDescriptorsPerSm); + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { + return cutlass::Status::kSuccess; + } + + template + static bool + can_implement( + ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits_d = cutlass::detail::get_output_alignment_bits(); + auto problem_shape_MNKL = append<4>(problem_shape.get_host_problem_shape(0), 1); + auto [M,N,K,L] = problem_shape_MNKL; + + constexpr int min_tma_aligned_elements_D = tma_alignment_bits_d / cutlass::sizeof_bits::value; + bool implementable = cutlass::detail::check_alignment(cute::make_shape(M,N,L), StrideD{}); + + if constexpr (not cute::is_void_v) { + constexpr int tma_alignment_bits_c = cutlass::detail::get_output_alignment_bits(); + constexpr int min_tma_aligned_elements_C = tma_alignment_bits_c / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,N,L), StrideC{}); + } + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + + bool beta_implementable = true; + + if constexpr (cute::is_void_v) { + if constexpr (detail::has_beta::value) { + beta_implementable = args.thread.beta == 0.0; + } + if constexpr (detail::has_beta_ptr::value) { + beta_implementable = beta_implementable && args.thread.beta_ptr == nullptr; + } + } + + if (!beta_implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Beta/beta pointer was set, but epilogue is sourceless (void-C).\n"); + } + + return implementable && beta_implementable; + } + + template + CUTLASS_HOST_DEVICE + static constexpr int + get_load_pipe_increment(CtaTileMNK cta_tile_mnk) { + // Compute number of epilogue subtiles + constexpr int epi_m = size<0>(cta_tile_mnk) / size<0>(EpilogueTile{}); + constexpr int epi_n = size<1>(cta_tile_mnk) / size<1>(EpilogueTile{}); + + return epi_m * epi_n; + } + + template + CUTLASS_HOST_DEVICE + static constexpr int + get_store_pipe_increment(CtaTileMNK cta_tile_mnk) { + return get_load_pipe_increment(cta_tile_mnk); + } + + CUTLASS_HOST_DEVICE + CollectiveEpilogue(Params const& params_, TensorStorage&) + : params(params_), epilogue_op(params_.thread) {} + + CUTLASS_DEVICE bool + is_producer_load_needed() const { + return epilogue_op.is_source_needed(); + } + + template + CUTLASS_DEVICE auto + load_init( + Params const& params, + TensorMapStorage& shared_tensormap, + int32_t const sm_count, + int32_t const sm_idx) const { + if constexpr (IsTmaAsyncUpdate) { + // Async update kernels will fetch the tensormap directly from tensormaps_init. + return cute::make_tuple(); + } else { + // Fetch a copy of tensormaps for the CTA from Params + constexpr bool IsEpiLoad = true; + auto load_tensormaps = tensormaps_init(params, shared_tensormap, sm_count, sm_idx); + return cute::make_tuple(load_tensormaps); + } + } + + template< + bool ReuseTmem = false, + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma, + class TensorMapC + > + CUTLASS_DEVICE auto + load( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + TensorStorage& shared_tensors, + cute::tuple, bool> load_tensormaps_info, + bool reverse_epi_n = false) { + using namespace cute; + + // Check to see if tensormaps have been replaced in gmem + if (get<1>(load_tensormaps_info) /* did_batch_change */) { + tensormaps_fence_acquire(get<0>(load_tensormaps_info)); + } + + int lane_idx = canonical_lane_idx(); + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + + auto coord_shape = append<3>(make_shape(m_coord, n_coord),Int<0>{}); + // Tile residue + auto m_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + return get<0,i>(problem_shape_mnkl) - get<0,i>(cta_tile_mnk) * get<0,i>(cta_coord_mnkl); + })); + auto n_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + return get<1,i>(problem_shape_mnkl) - get<1,i>(cta_tile_mnk) * get<1,i>(cta_coord_mnkl); + })); + auto residue_mn = make_coord(m_max_coord, n_max_coord); + + // Represent the full source tensor, slice to get the tile this CTA is currently responsible for + Tensor mC_real_mn = params.tma_load_c_real.get_tma_tensor(append<3>(make_shape(M,N),Int<1>{})); // (M,N,L) + Tensor mC_imag_mn = params.tma_load_c_imag.get_tma_tensor(append<3>(make_shape(M,N),Int<1>{})); // (M,N,L) + + Tensor mC_real = coalesce(mC_real_mn, take<0,2>(cta_tile_mnk)); + Tensor mC_imag = coalesce(mC_imag_mn, take<0,2>(cta_tile_mnk)); + + Tensor gC_real = local_tile(mC_real, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + Tensor gC_imag = local_tile(mC_imag, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + + // Apply epilogue subtile, get matching smem tensor + auto ptr_sC_real = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_real.begin(); + } + else { + return shared_tensors.smem_D_real.begin(); + } + }(); + auto ptr_sC_imag = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_imag.begin(); + } + else { + return shared_tensors.smem_D_imag.begin(); + } + }(); + + Tensor gC_real_epi = flat_divide(gC_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor gC_imag_epi = flat_divide(gC_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + Tensor sC_real_epi = make_tensor(make_smem_ptr(ptr_sC_real), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + Tensor sC_imag_epi = make_tensor(make_smem_ptr(ptr_sC_imag), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + + // Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_) + ThrCopy thrblk_g2s_real = params.tma_load_c_real.get_slice(Int<0>{}); + ThrCopy thrblk_g2s_imag = params.tma_load_c_imag.get_slice(Int<0>{}); + + Tensor bGS_gC_real = thrblk_g2s_real.partition_S(gC_real_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N) + Tensor bGS_gC_imag = thrblk_g2s_imag.partition_S(gC_imag_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N) + + Tensor bGS_sC_real = thrblk_g2s_real.partition_D(sC_real_epi); // (TMA,TMA_M,TMA_N,PIPE_C) + Tensor bGS_sC_imag = thrblk_g2s_imag.partition_D(sC_imag_epi); // (TMA,TMA_M,TMA_N,PIPE_C) + + // Predication for TMA load (one thread issues TMA load) + bool issue_tma_load = cute::elect_one_sync(); + + // Acquire the lock for the first stage + uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + load_pipeline.producer_acquire(load_pipe_producer_state); + + CUTLASS_PRAGMA_UNROLL + for (int iter_n = 0; iter_n < size<3>(gC_real_epi); ++iter_n) { + CUTLASS_PRAGMA_UNROLL + for (int iter_m = 0; iter_m < size<2>(gC_real_epi); ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + if constexpr (ReuseTmem) { + if (reverse_epi_n) { + epi_n = size<3>(gC_real_epi) - 1 - iter_n; + } + } + // Acquire the lock for this stage + constexpr uint16_t mcast_mask = 0; + uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + load_pipeline.producer_acquire(load_pipe_producer_state); + + // Execute the TMA load for C + if (issue_tma_load) { + copy(params.tma_load_c_real.with(get<0>(get<0>(load_tensormaps_info)), *tma_barrier, mcast_mask), + bGS_gC_real(_,_,_,epi_m,epi_n), bGS_sC_real(_,_,_,load_pipe_producer_state.index())); + copy(params.tma_load_c_imag.with(get<1>(get<0>(load_tensormaps_info)), *tma_barrier, mcast_mask), + bGS_gC_imag(_,_,_,epi_m,epi_n), bGS_sC_imag(_,_,_,load_pipe_producer_state.index())); + load_pipeline.producer_expect_transaction(load_pipe_producer_state); + } + + // Commit TMA loads for this stage and release the lock + load_pipeline.producer_commit(load_pipe_producer_state); + ++load_pipe_producer_state; + } + } + + return load_pipe_producer_state; + } + + CUTLASS_DEVICE void + load_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + [[maybe_unused]] StorePipeline store_pipeline, + [[maybe_unused]] StorePipelineState store_pipe_producer_state) { + load_pipeline.producer_tail(load_pipe_producer_state); + } + + template + CUTLASS_DEVICE auto + store_init( + Params const& params, + TensorMapStorage& shared_tensormap, + int32_t const sm_count, + int32_t const sm_idx) const { + if constexpr (IsTmaAsyncUpdate) { + return cute::make_tuple(); + } else { + // Fetch a copy of tensormaps for the CTA from Params + constexpr bool IsEpiLoad = false; + cute::tuple store_tensormaps = {nullptr, nullptr}; + int thread_idx = threadIdx.x % ThreadCount; + int warp_idx = thread_idx / NumThreadsPerWarp; + // Only the first epilogue warp needs to perform TMA related operations + if (warp_idx == 0) { + store_tensormaps = tensormaps_init(params, shared_tensormap, sm_count, sm_idx); + } + return cute::make_tuple(store_tensormaps); + } + } + + template< + bool ReuseTmem = false, + class AccumulatorPipeline, + class AccumulatorPipelineState, + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma, + class AccEngine, + class AccLayout, + class TensorMapD + > + CUTLASS_DEVICE auto + store( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + AccumulatorPipeline acc_pipeline, + AccumulatorPipelineState acc_pipe_consumer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + cute::Tensor accumulators, + TensorStorage& shared_tensors, + cute::tuple, bool> store_tensormap_info + ) { + using namespace cute; + using ElementAccumulator = typename AccEngine::value_type; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + + static_assert(is_tmem::value, "Accumulator must be TMEM resident."); + //static_assert(rank(accumulators) == 4, "Accumulators must be MMA-partitioned: [MMA, MMA_M, MMA_N]"); + static_assert(size<1>(accumulators) == 1 && size<2>(accumulators) == 1, "TiledMMA must match partitioned ShapeMN"); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(rank(CtaCoordMNKL{}) == 4, "CoordMNKL must be rank 4"); + + // Indexing variables + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + int thread_idx = threadIdx.x % ThreadCount; + int warp_idx = thread_idx / NumThreadsPerWarp; + [[maybe_unused]] int lane_idx = thread_idx % NumThreadsPerWarp; + + // Check to see if tensormaps have been replaced in gmem + // Only the first epilogue warp needs to perform TMA related operations + if (get<1>(store_tensormap_info) /* did_batch_change */ && warp_idx == 0) { + tensormaps_fence_acquire(get<0>(store_tensormap_info)); + } + + auto accumulators_real = accumulators(_,_,_,0); + auto accumulators_imag = accumulators(_,_,_,1); + + auto coord_shape = append<3>(make_shape(m_coord, n_coord),Int<0>{}); + + // Represent the full output tensor, slice to get the tile this CTA is responsible for + Tensor mD_real_mn = params.tma_store_d_real.get_tma_tensor(append<3>(make_shape(M,N),Int<1>{})); // (M,N,L) + Tensor mD_imag_mn = params.tma_store_d_imag.get_tma_tensor(append<3>(make_shape(M,N),Int<1>{})); // (M,N,L) + + Tensor mD_real = coalesce(mD_real_mn, take<0,2>(cta_tile_mnk)); + Tensor mD_imag = coalesce(mD_imag_mn, take<0,2>(cta_tile_mnk)); + + Tensor gD_real = local_tile(mD_real, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + Tensor gD_imag = local_tile(mD_imag, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + Tensor tAcc_real = accumulators_real(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + Tensor tAcc_imag = accumulators_imag(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + + // Apply epilogue subtiling + Tensor tAcc_real_epi = flat_divide(tAcc_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor tAcc_imag_epi = flat_divide(tAcc_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + Tensor gD_real_epi = flat_divide(gD_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor gD_imag_epi = flat_divide(gD_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Construct the corresponding pipelined smem tensors + auto ptr_sC_real = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_real.begin(); + } + else { + return shared_tensors.smem_D_real.begin(); + } + }(); + auto ptr_sC_imag = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_imag.begin(); + } + else { + return shared_tensors.smem_D_imag.begin(); + } + }(); + + auto ptr_sD_real = shared_tensors.smem_D_real.begin(); + auto ptr_sD_imag = shared_tensors.smem_D_imag.begin(); + + Tensor sC_real_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sC_real), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + Tensor sC_imag_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sC_imag), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + + Tensor sD_real_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sD_real), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + Tensor sD_imag_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sD_imag), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + // (t)hread-partition for (t)mem to (r)egister copy (tTR_) + TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tAcc_real_epi(_,_,_0{},_0{})); + ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_tAcc_real = thread_t2r.partition_S(tAcc_real_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + Tensor tTR_sD_real = thread_t2r.partition_D(sD_real_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N) + Tensor tTR_tAcc_imag = thread_t2r.partition_S(tAcc_imag_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + Tensor tTR_sD_imag = thread_t2r.partition_D(sD_imag_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N) + + // Allocate D and accumulator registers + Tensor tTR_rAcc = make_tensor(append(shape(tTR_sD_real), Int{})); // (T2R,T2R_M,T2R_N,2) + Tensor tTR_rD = make_tensor(append(shape(tTR_sD_real), Int{})); // (T2R,T2R_M,T2R_N,2) + + // Vectorized fragment view + constexpr int FragmentSize = DispatchPolicy::FragmentSize; + Tensor tTR_rAcc_frg = recast>(coalesce(tTR_rAcc)); // (EPI_V) + Tensor tTR_rD_frg = recast>(coalesce(tTR_rD)); // (EPI_V) + + CUTE_STATIC_ASSERT(size(tTR_rAcc) % DispatchPolicy::FragmentSize == 0, "Fragment size does not vectorize properly"); + + // (t)hread-partition for (s)mem to (r)egister copy (tSR_) + TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sC_real = thread_s2r.partition_S(sC_real_epi); // (S2R,S2R_M,S2R_N,PIPE_C) + Tensor tSR_sC_imag = thread_s2r.partition_S(sC_imag_epi); // (S2R,S2R_M,S2R_N,PIPE_C) + + Layout tSR_rC_layout = thread_s2r.retile_D(tTR_rD(_,_,_,_0{})).layout(); // (S2R,S2R_M,S2R_N) + + // Allocate C registers + // If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type + // to eliminate some redundant pack+unpack instruction sequences for sub-word types + constexpr bool IsDirectS2R = cute::is_same_v + && decltype(max_common_vector(tSR_rC_layout, tSR_sC_real.layout()))::value <= 1; + using RegisterElementC = cute::conditional_t; + Tensor tTR_rC = make_tensor(append(shape(tTR_sD_real), _2{})); // (T2R,T2R_M,T2R_N) + Tensor tSR_rC = thread_s2r.retile_D(tTR_rC); // (S2R,S2R_M,S2R_N) + Tensor tTR_rC_frg = recast>(tTR_rC); // (EPI_V) + + // (t)hread-partition for (r)egister to (s)mem copy (tRS_) + TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rD = thread_r2s.retile_S(tTR_rD); // (R2S,R2S_M,R2S_N) + Tensor tRS_sD_real = thread_r2s.partition_D(sD_real_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + Tensor tRS_sD_imag = thread_r2s.partition_D(sD_imag_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + ThrCopy thrblk_s2g = params.tma_store_d_real.get_slice(Int<0>{}); + Tensor bSG_sD_real = thrblk_s2g.partition_S(sD_real_epi); // (S2G,S2G_M,S2G_N,PIPE_D) + Tensor bSG_gD_real = thrblk_s2g.partition_D(gD_real_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + Tensor bSG_sD_imag = thrblk_s2g.partition_S(sD_imag_epi); // (S2G,S2G_M,S2G_N,PIPE_D) + Tensor bSG_gD_imag = thrblk_s2g.partition_D(gD_imag_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + // Coordinate tensors and residue for tile quantization + auto m_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + auto c_m = get<0,i>(problem_shape_mnkl) - get<0,i>(cta_tile_mnk) * get<0,i>(cta_coord_mnkl); + return cute::max(0, c_m); + })); + auto n_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + auto c_n = get<1,i>(problem_shape_mnkl) - get<1,i>(cta_tile_mnk) * get<1,i>(cta_coord_mnkl); + return cute::max(0, c_n); + })); + auto residue_mn = make_coord(m_max_coord, n_max_coord); + Tensor cD = make_identity_tensor(take<0,2>(cta_tile_mnk)); + Tensor tTR_cD = thread_t2r.partition_D(flat_divide(cD, EpilogueTile{})); + + bool is_source_needed = epilogue_op.is_source_needed(); + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; + + // Predication for sub-128 thread T2R tiled copy + Layout tmem_warp_layout = typename decltype(make_tmem_warp_partitioner(tAcc_real_epi(_,_,0,0)))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = size(tmem_warp_layout) != cosize(tmem_warp_layout); + bool issue_tmem_load = true; + + // If tmem doesn't have enough capacity to support double buffering, a portion of tmem (a column of epilogue tiles) + // is overlapped between 2 pseudo-buffers. The shared tmem portion corresponds to the last epilogue tile column of + // tmem accumulator buffer 0, and the first epilogue tile column of tmem accumulator 1. + // Thus, whenever we are processing tmem accumulator buffer 0, we process the epilogue tiles with reversed column order. + // Once the last epilogue tile column is loaded from tmem, the acc_pipeline is released. + // Then, the next accumulation stage for buffer 1 can start. + [[maybe_unused]] bool reverse_epi_n = ReuseTmem && acc_pipe_consumer_state.phase() == 0; + static_assert(not (ReuseTmem && AccumulatorPipeline::Stages != 1), "Tmem reuse requires 1 accumulator stage"); + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = warp_idx == 0; + + // In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight. + // The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in flight, so we can + // only guarantee store completion after StagesD iterations, then we can begin issuing releases on the smem buffer locks. + // store_pipe_producer_state tracks the acquire and load_pipe_consumer_state tracks the release, in circular buffer fashion. + // If TMA store supported async transaction mbarriers we would not need this synchronous release behavior. + LoadPipelineState load_wait_state = load_pipe_consumer_state; + if constexpr (ReuseSmemC) { + load_wait_state = store_pipe_producer_state; + load_wait_state.phase_ ^= 1; + } + + // We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA instructions + // Sync requirements of smem reuse may preclude this optimization + [[maybe_unused]] int epi_m_prev = 0; + [[maybe_unused]] int epi_n_prev = 0; + static_assert(not (DelayTmaStore and ReuseSmemC and StagesC <= StagesD), "This TMA epilogue configuration will deadlock"); + + // The TMA store sequence for one subtile iteration + auto tma_store_fn = [&] (int epi_m, int epi_n) { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d_real.with(get<0>(get<0>(store_tensormap_info))), bSG_sD_real(_,_,_,store_pipe_producer_state.index()), bSG_gD_real(_,_,_,epi_m,epi_n)); + copy(params.tma_store_d_imag.with(get<1>(get<0>(store_tensormap_info))), bSG_sD_imag(_,_,_,store_pipe_producer_state.index()), bSG_gD_imag(_,_,_,epi_m,epi_n)); + } + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + if constexpr (ReuseSmemC) { + // producer_acquire returns when at most StagesD-1 committed stores are pending + bool store_finished = store_pipe_producer_state.count() > StorePipeline::UnacquiredStages; + // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits + if (store_finished) { + if (is_source_needed) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + ++load_pipe_consumer_state; + } + } + }; + + // + // BEGIN EPILOGUE + // + + // Begin the wait for the producer load results + ConsumerToken load_wait_token{BarrierStatus::WaitDone}; + if (is_source_needed) { + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state); + } + // Begin the wait for the accumulator results + ConsumerToken acc_wait_token = acc_pipeline.consumer_try_wait(acc_pipe_consumer_state); + + // For each epilogue subtile within the CTA tile + CUTLASS_PRAGMA_UNROLL + for (int iter_n = 0; iter_n < size<3>(gD_real_epi); ++iter_n) { + CUTLASS_PRAGMA_UNROLL + for (int iter_m = 0; iter_m < size<2>(gD_real_epi); ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + bool is_first_iteration = iter_m == 0 && iter_n == 0; + bool is_last_iteration = iter_m == size<2>(gD_real_epi)-1 && iter_n == size<3>(gD_real_epi)-1; + bool do_acc_release = is_last_iteration; + + // Reverse subtile order for tmem reuse if necessary + if constexpr (ReuseTmem) { + if (reverse_epi_n) { + epi_n = size<3>(gD_real_epi) - 1 - iter_n; + } + do_acc_release = iter_m == size<2>(gD_real_epi)-1 && iter_n == 0; + } + + if (is_source_needed) { + // Wait for the producer load to fill smem + load_pipeline.consumer_wait(load_wait_state, load_wait_token); + + // Copy source tile from smem to register // residual smem -> reg + copy(tiled_s2r, tSR_sC_real(_,_,_,load_wait_state.index()), tSR_rC(_,_,_,0)); + copy(tiled_s2r, tSR_sC_imag(_,_,_,load_wait_state.index()), tSR_rC(_,_,_,1)); + } + + if (is_source_needed) { + // Let producer load warp know smem buffers are consumed and empty + if constexpr (not ReuseSmemC) { + cutlass::arch::fence_view_async_shared(); + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + ++load_wait_state; + } + + if (is_first_iteration) { + // Wait for mma warp to fill tmem buffer with accumulator results + acc_pipeline.consumer_wait(acc_pipe_consumer_state, acc_wait_token); + } + + // The current tile in tmem + Tensor tTR_tAcc_real_mn = tTR_tAcc_real(_,_,_,epi_m,epi_n); + Tensor tTR_tAcc_imag_mn = tTR_tAcc_imag(_,_,_,epi_m,epi_n); + + // Compute tmem load predication if necessary + if constexpr (predicate_tmem_load) { + // Issue tmem load if this tile's tmem subpartition is accessible by this warp + int subpart_idx = (tTR_tAcc_real_mn.data().dp_ / 32) % 4; + issue_tmem_load = warp_idx == subpart_idx; + } + + // Copy accumulator tile from tmem to register + if (issue_tmem_load) { // acc tmem -> reg + copy(tiled_t2r, tTR_tAcc_real_mn, tTR_rAcc(_,_,_,0)); + copy(tiled_t2r, tTR_tAcc_imag_mn, tTR_rAcc(_,_,_,1)); + } + + // After the last tmem load, signal that tmem buffer is consumed and empty + if (do_acc_release) { + cutlass::arch::fence_view_async_tmem_load(); + acc_pipeline.consumer_release(acc_pipe_consumer_state); + ++acc_pipe_consumer_state; + } + + // Vectorized fragment loop with visitor callback entry point + if (epilogue_op.is_source_needed()) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rD_frg); ++i) { + tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i), tTR_rC_frg(i)); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rD_frg); ++i) { + tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i)); + } + } + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the previous subtile + if (not is_first_iteration) { + tma_store_fn(epi_m_prev, epi_n_prev); + } + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + + // Copy output tile from register to smem + bool issue_smem_store = issue_tmem_load; + if (issue_smem_store) { // after scale, reg -> smem + copy(tiled_r2s, tRS_rD(_,_,_,0), tRS_sD_real(_,_,_,store_pipe_producer_state.index())); + copy(tiled_r2s, tRS_rD(_,_,_,1), tRS_sD_imag(_,_,_,store_pipe_producer_state.index())); + } + + if constexpr (not DelayTmaStore) { + // Issue TMA stores for this subtile + tma_store_fn(epi_m, epi_n); + } + + if (is_source_needed) { + // Begin the wait for the next subtile producer load + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state, is_last_iteration); + } + } // for epi_m + } // for epi_n + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the last subtile + tma_store_fn(epi_m_prev, epi_n_prev); + } + + return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_pipe_consumer_state); + } + + template + CUTLASS_DEVICE void + store_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + CtaTileMNK cta_tile_mnk) { + if constexpr (ReuseSmemC) { + if (epilogue_op.is_source_needed()) { + // wait for all TMA stores to complete + store_pipeline.producer_tail(store_pipe_producer_state); + + // Issue releases on up to StagesD-1 previously issued TMA stores + constexpr int release_stages = cute::min(StorePipeline::UnacquiredStages, get_load_pipe_increment(cta_tile_mnk)); + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < release_stages; ++stage) { + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + } + } + } + + // + // Methods to perform different parts of TMA/Tensormap modifications + // + + template + CUTLASS_DEVICE auto + tensormaps_init(Params const& params, + TensorMapStorage& shared_tensormap, + int32_t const sm_count, + int32_t const sm_idx, + bool const is_leader_warp = true) const { + // Define a local struct that provides simple array indexing for TMA descriptors + struct TensorMapArray { + cute::TmaDescriptor* tma_desc_real; + cute::TmaDescriptor* tma_desc_imag; + + TensorMapArray() = default; + + CUTLASS_DEVICE + TensorMapArray(cute::TmaDescriptor* desc_real, cute::TmaDescriptor* desc_imag) + : tma_desc_real(desc_real), tma_desc_imag(desc_imag) {} + + CUTLASS_DEVICE + cute::tuple + operator[](int32_t idx) const { + idx = idx % NumTmaDescriptorsPerSm; + return cute::make_tuple(tma_desc_real + idx, tma_desc_imag + idx); + } + }; + + cute::TmaDescriptor* tma_desc_real = nullptr; + cute::TmaDescriptor* tma_desc_imag = nullptr; + cute::TmaDescriptor* gmem_tensormap = params.tensormaps; + + if (!is_leader_warp) { + if constexpr (IsTmaAsyncUpdate) { + return TensorMapArray{tma_desc_real, tma_desc_imag}; + } else { + return cute::make_tuple(tma_desc_real, tma_desc_imag); + } + } + if constexpr (IsLoad) { + if constexpr (not cute::is_void_v) { + tma_desc_real = &gmem_tensormap[sm_idx * NumTmaDescriptorsPerSm]; + tma_desc_imag = &gmem_tensormap[(sm_idx + sm_count) * NumTmaDescriptorsPerSm]; + + if (cute::elect_one_sync()) { + // Bringing tensormaps from params to smem for modification later + Tensor pC_real_tensormap = make_tensor(params.tma_load_c_real.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sC_real_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_C_real), Int<1>{}, Int<1>{}); + Tensor pC_imag_tensormap = make_tensor(params.tma_load_c_imag.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sC_imag_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_C_imag), Int<1>{}, Int<1>{}); + + copy(recast(pC_real_tensormap), recast(sC_real_tensormap)); + copy(recast(pC_imag_tensormap), recast(sC_imag_tensormap)); + } + __syncwarp(); + } + } else { + int const offset_Ddesc = cute::is_void_v ? 0 : (2 * sm_count); + tma_desc_real = &gmem_tensormap[(sm_idx + offset_Ddesc) * NumTmaDescriptorsPerSm]; + tma_desc_imag = &gmem_tensormap[(sm_idx + offset_Ddesc + sm_count) * NumTmaDescriptorsPerSm]; + + if (cute::elect_one_sync()) { + // Bringing tensormaps from params to gmem for modification later + Tensor pD_real_tensormap = make_tensor(params.tma_store_d_real.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sD_real_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_D_real), Int<1>{}, Int<1>{}); + Tensor pD_imag_tensormap = make_tensor(params.tma_store_d_imag.get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sD_imag_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_D_imag), Int<1>{}, Int<1>{}); + + copy(recast(pD_real_tensormap), recast(sD_real_tensormap)); + copy(recast(pD_imag_tensormap), recast(sD_imag_tensormap)); + } + __syncwarp(); + } + + if constexpr (IsTmaAsyncUpdate) { + return TensorMapArray{tma_desc_real, tma_desc_imag}; + } else { + return cute::make_tuple(tma_desc_real, tma_desc_imag); + } + } + + // Replace address for the global tensor (to be done by single thread) + template + CUTLASS_DEVICE + void + tensormaps_replace_global_address( + TensorMapStorage& shared_tensormap, + Params const& params, + int32_t next_batch) { + // Replacing global_address for the next batch + if constexpr (IsLoad) { + if (not cute::is_void_v) { + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_C_real, + params.ptr_C_real[next_batch]); + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_C_imag, + params.ptr_C_imag[next_batch]); + } + } else { + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_D_real, + params.ptr_D_real[next_batch]); + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_D_imag, + params.ptr_D_imag[next_batch]); + } + } + + template + CUTLASS_DEVICE + void + tensormaps_perform_update( + TensorMapStorage& shared_tensormap, + Params const& params, + cute::tuple const& tensormaps, + [[maybe_unused]] ProblemShape problem_shape, + int32_t next_batch + ) { + if (cute::elect_one_sync()) { + // Replacing global_address for the next batch + tensormaps_replace_global_address(shared_tensormap, params, next_batch); + } + // Ensure warp is converged before issuing tensormap fence release + __syncwarp(); + // Entire warp must do this (ie its aligned) + tensormaps_cp_fence_release( + shared_tensormap, + tensormaps + ); + } + + template + CUTLASS_DEVICE + void + tensormaps_cp_fence_release( + TensorMapStorage& shared_tensormap, + cute::tuple const& tensormaps + ) { + // Commit and wait for all TMA load/store instructions before updating the tensormap in gmem if we're not using async update. + // This operation only happens when the group/batch changes between consecutive tiles. + // If there are no uncommitted instructions then tma_desc_commit_group results in an empty bulk async-group. + auto tma_desc_wait_all_fn = [] () CUTLASS_LAMBDA_FUNC_INLINE { + if (cute::elect_one_sync()) { + cute::tma_desc_commit_group(); + cute::tma_desc_wait_group(); + } + }; + + // Entire warp must do this (ie its aligned) + if constexpr (IsLoad) { + if constexpr (not cute::is_void_v) { + if constexpr (WaitForInflightTmaRequests) { + tma_desc_wait_all_fn(); + } + tma_descriptor_cp_fence_release(get<0>(tensormaps), shared_tensormap.smem_tensormap_C_real); + tma_descriptor_cp_fence_release(get<1>(tensormaps), shared_tensormap.smem_tensormap_C_imag); + } + } else { + if constexpr (WaitForInflightTmaRequests) { + tma_desc_wait_all_fn(); + } + tma_descriptor_cp_fence_release(get<0>(tensormaps), shared_tensormap.smem_tensormap_D_real); + tma_descriptor_cp_fence_release(get<1>(tensormaps), shared_tensormap.smem_tensormap_D_imag); + } + } + + template + CUTLASS_DEVICE + void + tensormaps_fence_acquire(cute::tuple const& tensormaps) { + if constexpr (IsLoad) { + if constexpr (not cute::is_void_v) { + cute::tma_descriptor_fence_acquire(get<0>(tensormaps)); + cute::tma_descriptor_fence_acquire(get<1>(tensormaps)); + } + } else { + cute::tma_descriptor_fence_acquire(get<0>(tensormaps)); + cute::tma_descriptor_fence_acquire(get<1>(tensormaps)); + } + } + +private: + Params const& params; + ThreadEpilogueOp epilogue_op; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::epilogue::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp index 04f1eec8..9ad99b13 100644 --- a/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp +++ b/include/cutlass/epilogue/collective/sm100_epilogue_array_tma_warpspecialized.hpp @@ -1532,6 +1532,7 @@ public: ProblemShape problem_shape, int32_t next_batch ) { + __syncwarp(); if (cute::elect_one_sync()) { // Replacing global_address for the next batch tensormaps_replace_global_address(shared_tensormap, params, next_batch); diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp index bf7e1b6a..c8944f15 100644 --- a/include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp +++ b/include/cutlass/epilogue/collective/sm100_epilogue_nosmem.hpp @@ -214,6 +214,8 @@ protected: public: template< bool ReuseTmem = false, + class LoadPipeline, + class LoadPipelineState, class AccumulatorPipeline, class AccumulatorPipelineState, class ProblemShapeMNKL, @@ -223,6 +225,8 @@ public: > CUTLASS_DEVICE auto operator()( + [[maybe_unused]]LoadPipeline load_pipeline, + [[maybe_unused]]LoadPipelineState load_pipe_consumer_state, AccumulatorPipeline acc_pipeline, AccumulatorPipelineState acc_pipe_consumer_state, ProblemShapeMNKL problem_shape_mnkl, @@ -352,7 +356,7 @@ public: copy_if(tDpD, tTR_rD_src, tR2G_rD_dst); } - return cute::make_tuple(acc_pipe_consumer_state); + return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state); } @@ -571,6 +575,8 @@ public: template< bool ReuseTmem = false, + class LoadPipeline, + class LoadPipelineState, class AccumulatorPipeline, class AccumulatorPipelineState, class ProblemShapeMNKL, @@ -580,6 +586,8 @@ public: > CUTLASS_DEVICE auto operator()( + [[maybe_unused]]LoadPipeline load_pipeline, + [[maybe_unused]]LoadPipelineState load_pipe_consumer_state, AccumulatorPipeline acc_pipeline, AccumulatorPipelineState acc_pipe_consumer_state, ProblemShapeMNKL problem_shape_mnkl, @@ -788,7 +796,7 @@ public: // auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks(cst_args); epi_loop_fn(cst_callbacks); - return cute::make_tuple(acc_pipe_consumer_state); + return cute::make_tuple(acc_pipe_consumer_state, load_pipe_consumer_state); } }; diff --git a/include/cutlass/epilogue/collective/sm100_epilogue_planar_complex_tma_warpspecialized.hpp b/include/cutlass/epilogue/collective/sm100_epilogue_planar_complex_tma_warpspecialized.hpp new file mode 100644 index 00000000..e9b8d18e --- /dev/null +++ b/include/cutlass/epilogue/collective/sm100_epilogue_planar_complex_tma_warpspecialized.hpp @@ -0,0 +1,897 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Functor performing elementwise operations used by Planar Complex Gemm epilogues. +*/ + + + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/arch/barrier.h" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/detail.hpp" +#include "cutlass/epilogue/thread/scale_type.h" +#include "cutlass/epilogue/thread/linear_combination_planar_complex.h" +#include "cutlass/detail/layout.hpp" +#include "cutlass/trace.h" + +#include "cute/tensor.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::epilogue::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int StagesC_, + int StagesD_, + int FragmentSize_, + bool ReuseSmemC_, + bool DelayTmaStore_, + class CtaTileShape_, // (CTA_M,CTA_N,CTA_K, optional: Tile_L) + class EpilogueTile_, // (EPI_TILE_M, EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class ThreadEpilogueOp_, + class CopyOpT2R_, + class CopyOpG2S_, + class SmemLayoutAtomC_, + class CopyOpS2R_, + class CopyOpS2G_, + class SmemLayoutAtomD_, + class CopyOpR2S_, + class CopyOpR2R_ +> +class CollectiveEpilogue< + Sm100PlanarComplexTmaWarpSpecialized, + CtaTileShape_, + EpilogueTile_, + ElementC_, + StrideC_, + ElementD_, + StrideD_, + ThreadEpilogueOp_, + CopyOpT2R_, + CopyOpG2S_, + SmemLayoutAtomC_, + CopyOpS2R_, + CopyOpS2G_, + SmemLayoutAtomD_, + CopyOpR2S_, + CopyOpR2R_ +> { +public: + // + // Type Aliases + // + using DispatchPolicy = Sm100PlanarComplexTmaWarpSpecialized; + using CtaTileShape = CtaTileShape_; + using EpilogueTile = EpilogueTile_; + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementC = ElementC_; + using StrideC = StrideC_; + using ElementD = ElementD_; + using StrideD = StrideD_; + using CopyOpT2R = CopyOpT2R_; + using CopyOpG2S = CopyOpG2S_; + using SmemLayoutAtomC = SmemLayoutAtomC_; + using CopyOpS2R = CopyOpS2R_; + using CopyOpS2G = CopyOpS2G_; + using SmemLayoutAtomD = SmemLayoutAtomD_; + using CopyOpR2S = CopyOpR2S_; + using CopyOpR2R = CopyOpR2R_; + + using GmemTiledCopyC = CopyOpG2S; + using GmemTiledCopyD = CopyOpS2G; + + constexpr static int ThreadCount = 128; + + static_assert(!is_layout::value && is_tuple::value, "EpilogueTile must be a cute::Tile or cute::Shape"); + static_assert(rank(EpilogueTile{}) == 2, "EpilogueTile must be rank-2: [EPI_TILE_M, EPI_TILE_N]"); + +private: + using SmemElementD = typename cutlass::detail::get_unpacked_element_type::type; + using SmemElementC = typename cutlass::detail::get_unpacked_element_type,ElementD,ElementC>>::type; // prevents void ref breakages + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; + constexpr static bool ReuseSmemC = ReuseSmemC_; + constexpr static bool DelayTmaStore = DelayTmaStore_; + constexpr static bool is_source_supported = ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default; + + constexpr static bool is_m_major_C = detail::is_m_major(); + constexpr static bool is_m_major_D = detail::is_m_major(); + + using SmemLayoutC = decltype(tile_to_shape( + SmemLayoutAtomC{}, + make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int{}), + cute::conditional_t, Step<_1,_2,_3>>{} )); + using SmemLayoutD = decltype(tile_to_shape( + SmemLayoutAtomD{}, + make_shape(size<0>(shape(EpilogueTile{})), size<1>(shape(EpilogueTile{})), Int{}), + cute::conditional_t, Step<_1,_2,_3>>{} )); + + constexpr static bool support_smem_reuse = is_source_supported && StagesD <= StagesC + && cosize(take<0,2>(SmemLayoutC{})) == cosize(take<0,2>(SmemLayoutD{})); + static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met"); + + constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{}); + constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{}); + +public : + struct TensorStorageWithC { + alignas(SmemAlignmentC) cute::ArrayEngine> smem_C_real; + alignas(SmemAlignmentC) cute::ArrayEngine> smem_C_imag; + + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_real; + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_imag; + }; + + struct TensorStorageWithoutC { + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_real; + alignas(SmemAlignmentD) cute::ArrayEngine> smem_D_imag; + }; + +public: + // TMA pipeline for loading C + using LoadPipeline = cutlass::PipelineTransactionAsync; + using LoadPipelineState = cutlass::PipelineState; + constexpr static uint32_t TmaTransactionBytes = + 2 * ((size(take<0,2>(SmemLayoutC{})) * static_cast(sizeof_bits::value)) / 8); + + // TMA pipeline for storing D + using StorePipeline = cute::conditional_t, + cutlass::PipelineTmaStore>; + using StorePipelineState = cutlass::PipelineState; + + struct SharedStorage { + using TensorStorage = + cute::conditional_t; + TensorStorage tensors; + + using PipelineStorage = typename LoadPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Planar complex kernels have two accumulator copies for the real and imaginary tensors. + constexpr static int NumAccumulatorMtxs = 2; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread{}; + ElementC const* ptr_C_real = nullptr; + StrideC dC_real{}; + ElementC const* ptr_C_imag = nullptr; + StrideC dC_imag{}; + ElementD* ptr_D_real = nullptr; + StrideD dD_real{}; + ElementD* ptr_D_imag = nullptr; + StrideD dD_imag{}; + }; + + // Device side epilogue params + struct Params { + using TMA_C = decltype(make_tma_copy( + CopyOpG2S{}, + make_tensor( + make_gmem_ptr(static_cast,ElementD,ElementC> const*>(nullptr)), + repeat_like(append<3>(StrideC{}, _1{}), int32_t(0)), + append<3>(StrideC{}, _0{})), + take<0,2>(SmemLayoutC{}), + EpilogueTile{}, + _1{})); + using TMA_D = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor( + make_gmem_ptr(static_cast(nullptr)), + repeat_like(append<3>(StrideD{}, _1{}), int32_t(0)), + append<3>(StrideD{}, _0{})), + take<0,2>(SmemLayoutC{}), + EpilogueTile{}, + _1{})); + + typename ThreadEpilogueOp::Params thread{}; + TMA_C tma_load_c_real; + TMA_C tma_load_c_imag; + TMA_D tma_store_d_real; + TMA_D tma_store_d_imag; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + [[maybe_unused]] void* workspace) { + // Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_mnkl = append<4>(problem_shape, 1); + auto [M, N, K, L] = problem_shape_mnkl; + + typename Params::TMA_C tma_load_c_real{}; + typename Params::TMA_C tma_load_c_imag{}; + if constexpr (not cute::is_void_v) { + Tensor tensor_c_real = make_tensor(make_gmem_ptr(args.ptr_C_real), make_layout(make_shape(M,N,L), append<3>(args.dC_real, _0{}))); + Tensor tensor_c_imag = make_tensor(make_gmem_ptr(args.ptr_C_imag), make_layout(make_shape(M,N,L), append<3>(args.dC_imag, _0{}))); + + tma_load_c_real = make_tma_copy(CopyOpG2S{}, tensor_c_real, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{}); + tma_load_c_imag = make_tma_copy(CopyOpG2S{}, tensor_c_imag, take<0,2>(SmemLayoutC{}), EpilogueTile{}, _1{}); + } + + Tensor tensor_d_real = make_tensor(make_gmem_ptr(args.ptr_D_real), make_layout(make_shape(M,N,L), append<3>(args.dD_real, _0{}))); + Tensor tensor_d_imag = make_tensor(make_gmem_ptr(args.ptr_D_imag), make_layout(make_shape(M,N,L), append<3>(args.dD_imag, _0{}))); + + typename Params::TMA_D tma_store_d_real = + make_tma_copy(CopyOpS2G{}, tensor_d_real, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{}); + typename Params::TMA_D tma_store_d_imag = + make_tma_copy(CopyOpS2G{}, tensor_d_imag, take<0,2>(SmemLayoutD{}), EpilogueTile{}, _1{}); + + return { + args.thread, + tma_load_c_real, + tma_load_c_imag, + tma_store_d_real, + tma_store_d_imag + }; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) { + return 0; + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, + CudaHostAdapter* cuda_adapter = nullptr) { + return cutlass::Status::kSuccess; + } + + template + static bool + can_implement( + ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits_d = cutlass::detail::get_output_alignment_bits(); + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + constexpr int min_tma_aligned_elements_D = tma_alignment_bits_d / cutlass::sizeof_bits::value; + bool implementable = cutlass::detail::check_alignment(cute::make_shape(M,N,L), StrideD{}); + + if constexpr (not cute::is_void_v) { + constexpr int tma_alignment_bits_c = cutlass::detail::get_output_alignment_bits(); + constexpr int min_tma_aligned_elements_C = tma_alignment_bits_c / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,N,L), StrideC{}); + } + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + + bool beta_implementable = true; + + if constexpr (cute::is_void_v) { + if constexpr (detail::has_beta::value) { + beta_implementable = args.thread.beta == 0.0; + } + if constexpr (detail::has_beta_ptr::value) { + beta_implementable = beta_implementable && args.thread.beta_ptr == nullptr; + } + } + + if (!beta_implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Beta/beta pointer was set, but epilogue is sourceless (void-C).\n"); + } + + return implementable && beta_implementable; + } + + template + CUTLASS_HOST_DEVICE + static constexpr int + get_load_pipe_increment(CtaTileMNK cta_tile_mnk) { + // Compute number of epilogue subtiles + constexpr int epi_m = size<0>(cta_tile_mnk) / size<0>(EpilogueTile{}); + constexpr int epi_n = size<1>(cta_tile_mnk) / size<1>(EpilogueTile{}); + + return epi_m * epi_n; + } + + template + CUTLASS_HOST_DEVICE + static constexpr int + get_store_pipe_increment(CtaTileMNK cta_tile_mnk) { + return get_load_pipe_increment(cta_tile_mnk); + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE static void + prefetch_tma_descriptors(Params const& epilogue_params) { + cute::prefetch_tma_descriptor(epilogue_params.tma_load_c_real.get_tma_descriptor()); + cute::prefetch_tma_descriptor(epilogue_params.tma_load_c_imag.get_tma_descriptor()); + cute::prefetch_tma_descriptor(epilogue_params.tma_store_d_real.get_tma_descriptor()); + cute::prefetch_tma_descriptor(epilogue_params.tma_store_d_imag.get_tma_descriptor()); + } + + CUTLASS_HOST_DEVICE + CollectiveEpilogue(Params const& params_, TensorStorage&) + : params(params_), epilogue_op(params_.thread) {} + + CUTLASS_DEVICE bool + is_producer_load_needed() const { + return epilogue_op.is_source_needed(); + } + + template< + bool ReuseTmem = false, + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma + > + CUTLASS_DEVICE auto + load( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + TensorStorage& shared_tensors, + bool reverse_epi_n = false) { + using namespace cute; + + int lane_idx = canonical_lane_idx(); + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + + auto coord_shape = make_coord(m_coord, n_coord, l_coord); + + // Tile residue + auto m_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + return get<0,i>(problem_shape_mnkl) - get<0,i>(cta_tile_mnk) * get<0,i>(cta_coord_mnkl); + })); + auto n_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + return get<1,i>(problem_shape_mnkl) - get<1,i>(cta_tile_mnk) * get<1,i>(cta_coord_mnkl); + })); + auto residue_mn = make_coord(m_max_coord, n_max_coord); + + // Represent the full source tensor, slice to get the tile this CTA is currently responsible for + Tensor mC_real_mn = params.tma_load_c_real.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mC_imag_mn = params.tma_load_c_imag.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + + Tensor mC_real = coalesce(mC_real_mn, take<0,2>(cta_tile_mnk)); + Tensor mC_imag = coalesce(mC_imag_mn, take<0,2>(cta_tile_mnk)); + + Tensor gC_real = local_tile(mC_real, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + Tensor gC_imag = local_tile(mC_imag, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + + // Apply epilogue subtile, get matching smem tensor + auto ptr_sC_real = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_real.begin(); + } + else { + return shared_tensors.smem_D_real.begin(); + } + }(); + auto ptr_sC_imag = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_imag.begin(); + } + else { + return shared_tensors.smem_D_imag.begin(); + } + }(); + + Tensor gC_real_epi = flat_divide(gC_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor gC_imag_epi = flat_divide(gC_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + Tensor sC_real_epi = make_tensor(make_smem_ptr(ptr_sC_real), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + Tensor sC_imag_epi = make_tensor(make_smem_ptr(ptr_sC_imag), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + + // Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_) + ThrCopy thrblk_g2s_real = params.tma_load_c_real.get_slice(Int<0>{}); + ThrCopy thrblk_g2s_imag = params.tma_load_c_imag.get_slice(Int<0>{}); + + Tensor bGS_gC_real = thrblk_g2s_real.partition_S(gC_real_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N) + Tensor bGS_gC_imag = thrblk_g2s_imag.partition_S(gC_imag_epi); // (TMA,TMA_M,TMA_N,EPI_M,EPI_N) + + Tensor bGS_sC_real = thrblk_g2s_real.partition_D(sC_real_epi); // (TMA,TMA_M,TMA_N,PIPE_C) + Tensor bGS_sC_imag = thrblk_g2s_imag.partition_D(sC_imag_epi); // (TMA,TMA_M,TMA_N,PIPE_C) + + // Predication for TMA load (one thread issues TMA load) + bool issue_tma_load = cute::elect_one_sync(); + + // Acquire the lock for the first stage + uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + load_pipeline.producer_acquire(load_pipe_producer_state); + + CUTLASS_PRAGMA_UNROLL + for (int iter_n = 0; iter_n < size<3>(gC_real_epi); ++iter_n) { + CUTLASS_PRAGMA_UNROLL + for (int iter_m = 0; iter_m < size<2>(gC_real_epi); ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + if constexpr (ReuseTmem) { + if (reverse_epi_n) { + epi_n = size<3>(gC_real_epi) - 1 - iter_n; + } + } + // Acquire the lock for this stage + constexpr uint16_t mcast_mask = 0; + uint64_t* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + load_pipeline.producer_acquire(load_pipe_producer_state); + + // Execute the TMA load for C + if (issue_tma_load) { + copy(params.tma_load_c_real.with(*tma_barrier, mcast_mask), + bGS_gC_real(_,_,_,epi_m,epi_n), bGS_sC_real(_,_,_,load_pipe_producer_state.index())); + copy(params.tma_load_c_imag.with(*tma_barrier, mcast_mask), + bGS_gC_imag(_,_,_,epi_m,epi_n), bGS_sC_imag(_,_,_,load_pipe_producer_state.index())); + load_pipeline.producer_expect_transaction(load_pipe_producer_state); + } + + // Commit TMA loads for this stage and release the lock + load_pipeline.producer_commit(load_pipe_producer_state); + ++load_pipe_producer_state; + } + } + + return load_pipe_producer_state; + } + + CUTLASS_DEVICE void + load_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + [[maybe_unused]] StorePipeline store_pipeline, + [[maybe_unused]] StorePipelineState store_pipe_producer_state) { + load_pipeline.producer_tail(load_pipe_producer_state); + } + + template< + bool ReuseTmem = false, + class AccumulatorPipeline, + class AccumulatorPipelineState, + class ProblemShapeMNKL, + class CtaTileMNK, + class CtaCoordMNKL, + class MmaTileMNK, + class TiledMma, + class AccEngine, + class AccLayout + > + CUTLASS_DEVICE auto + store( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + AccumulatorPipeline acc_pipeline, + AccumulatorPipelineState acc_pipe_consumer_state, + ProblemShapeMNKL problem_shape_mnkl, + CtaTileMNK cta_tile_mnk, + CtaCoordMNKL cta_coord_mnkl, + MmaTileMNK mma_tile_mnk, + TiledMma tiled_mma, + cute::Tensor accumulators, + TensorStorage& shared_tensors + ) { + using namespace cute; + using ElementAccumulator = typename AccEngine::value_type; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + + static_assert(is_tmem::value, "Accumulator must be TMEM resident."); + //static_assert(rank(accumulators) == 4, "Accumulators must be MMA-partitioned: [MMA, MMA_M, MMA_N]"); + static_assert(size<1>(accumulators) == 1 && size<2>(accumulators) == 1, "TiledMMA must match partitioned ShapeMN"); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(rank(CtaCoordMNKL{}) == 4, "CoordMNKL must be rank 4"); + + // Indexing variables + auto [M, N, K, L] = problem_shape_mnkl; + auto [m_coord, n_coord, k_coord, l_coord] = cta_coord_mnkl; + int thread_idx = threadIdx.x % ThreadCount; + int warp_idx = thread_idx / NumThreadsPerWarp; + [[maybe_unused]] int lane_idx = thread_idx % NumThreadsPerWarp; + + auto accumulators_real = accumulators(_,_,_,0); + auto accumulators_imag = accumulators(_,_,_,1); + + auto coord_shape = make_coord(m_coord, n_coord, l_coord); + + // Represent the full output tensor, slice to get the tile this CTA is responsible for + Tensor mD_real_mn = params.tma_store_d_real.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + Tensor mD_imag_mn = params.tma_store_d_imag.get_tma_tensor(make_shape(M,N,L)); // (M,N,L) + + Tensor mD_real = coalesce(mD_real_mn, take<0,2>(cta_tile_mnk)); + Tensor mD_imag = coalesce(mD_imag_mn, take<0,2>(cta_tile_mnk)); + + Tensor gD_real = local_tile(mD_real, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + Tensor gD_imag = local_tile(mD_imag, take<0,2>(cta_tile_mnk), coord_shape); // (CTA_M,CTA_N) + + Tensor tAcc_real = accumulators_real(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + Tensor tAcc_imag = accumulators_imag(make_coord(_,_),_0{},_0{}); // (CTA_M,CTA_N) + + // Apply epilogue subtiling + Tensor tAcc_real_epi = flat_divide(tAcc_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor tAcc_imag_epi = flat_divide(tAcc_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + Tensor gD_real_epi = flat_divide(gD_real, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor gD_imag_epi = flat_divide(gD_imag, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Construct the corresponding pipelined smem tensors + auto ptr_sC_real = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_real.begin(); + } + else { + return shared_tensors.smem_D_real.begin(); + } + }(); + auto ptr_sC_imag = [&]() { + if constexpr (not ReuseSmemC and is_source_supported) { + return shared_tensors.smem_C_imag.begin(); + } + else { + return shared_tensors.smem_D_imag.begin(); + } + }(); + + auto ptr_sD_real = shared_tensors.smem_D_real.begin(); + auto ptr_sD_imag = shared_tensors.smem_D_imag.begin(); + + Tensor sC_real_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sC_real), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + Tensor sC_imag_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sC_imag), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C) + + Tensor sD_real_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sD_real), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + Tensor sD_imag_epi = cute::as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(ptr_sD_imag), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D) + + // (t)hread-partition for (t)mem to (r)egister copy (tTR_) + TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tAcc_real_epi(_,_,_0{},_0{})); + ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_tAcc_real = thread_t2r.partition_S(tAcc_real_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + Tensor tTR_sD_real = thread_t2r.partition_D(sD_real_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N) + Tensor tTR_tAcc_imag = thread_t2r.partition_S(tAcc_imag_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N) + Tensor tTR_sD_imag = thread_t2r.partition_D(sD_imag_epi(_,_,_0{})); // (T2R,T2R_M,T2R_N) + + // Allocate D and accumulator registers + Tensor tTR_rAcc = make_tensor(append(shape(tTR_sD_real), Int{})); // (T2R,T2R_M,T2R_N,2) + Tensor tTR_rD = make_tensor(append(shape(tTR_sD_real), Int{})); // (T2R,T2R_M,T2R_N,2) + + // Vectorized fragment view + constexpr int FragmentSize = DispatchPolicy::FragmentSize; + Tensor tTR_rAcc_frg = recast>(coalesce(tTR_rAcc)); // (EPI_V) + Tensor tTR_rD_frg = recast>(coalesce(tTR_rD)); // (EPI_V) + + CUTE_STATIC_ASSERT(size(tTR_rAcc) % DispatchPolicy::FragmentSize == 0, "Fragment size does not vectorize properly"); + + // (t)hread-partition for (s)mem to (r)egister copy (tSR_) + TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sC_real = thread_s2r.partition_S(sC_real_epi); // (S2R,S2R_M,S2R_N,PIPE_C) + Tensor tSR_sC_imag = thread_s2r.partition_S(sC_imag_epi); // (S2R,S2R_M,S2R_N,PIPE_C) + + Layout tSR_rC_layout = thread_s2r.retile_D(tTR_rD(_,_,_,_0{})).layout(); // (S2R,S2R_M,S2R_N) + + // Allocate C registers + // If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type + // to eliminate some redundant pack+unpack instruction sequences for sub-word types + constexpr bool IsDirectS2R = cute::is_same_v + && decltype(max_common_vector(tSR_rC_layout, tSR_sC_real.layout()))::value <= 1; + using RegisterElementC = cute::conditional_t; + Tensor tTR_rC = make_tensor(append(shape(tTR_sD_real), _2{})); // (T2R,T2R_M,T2R_N) + Tensor tSR_rC = thread_s2r.retile_D(tTR_rC); // (S2R,S2R_M,S2R_N) + Tensor tTR_rC_frg = recast>(tTR_rC); // (EPI_V) + + // (t)hread-partition for (r)egister to (s)mem copy (tRS_) + TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rD = thread_r2s.retile_S(tTR_rD); // (R2S,R2S_M,R2S_N) + Tensor tRS_sD_real = thread_r2s.partition_D(sD_real_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + Tensor tRS_sD_imag = thread_r2s.partition_D(sD_imag_epi); // (R2S,R2S_M,R2S_N,PIPE_D) + + // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_) + ThrCopy thrblk_s2g = params.tma_store_d_real.get_slice(Int<0>{}); + Tensor bSG_sD_real = thrblk_s2g.partition_S(sD_real_epi); // (S2G,S2G_M,S2G_N,PIPE_D) + Tensor bSG_gD_real = thrblk_s2g.partition_D(gD_real_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + Tensor bSG_sD_imag = thrblk_s2g.partition_S(sD_imag_epi); // (S2G,S2G_M,S2G_N,PIPE_D) + Tensor bSG_gD_imag = thrblk_s2g.partition_D(gD_imag_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + // Coordinate tensors and residue for tile quantization + auto m_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + auto c_m = get<0,i>(problem_shape_mnkl) - get<0,i>(cta_tile_mnk) * get<0,i>(cta_coord_mnkl); + return cute::max(0, c_m); + })); + auto n_max_coord = unwrap(cute::transform(make_seq(cta_tile_mnk)>{}, [&](auto i) { + auto c_n = get<1,i>(problem_shape_mnkl) - get<1,i>(cta_tile_mnk) * get<1,i>(cta_coord_mnkl); + return cute::max(0, c_n); + })); + auto residue_mn = make_coord(m_max_coord, n_max_coord); + Tensor cD = make_identity_tensor(take<0,2>(cta_tile_mnk)); + Tensor tTR_cD = thread_t2r.partition_D(flat_divide(cD, EpilogueTile{})); + + bool is_source_needed = epilogue_op.is_source_needed(); + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); }; + + // Predication for sub-128 thread T2R tiled copy + Layout tmem_warp_layout = typename decltype(make_tmem_warp_partitioner(tAcc_real_epi(_,_,0,0)))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = size(tmem_warp_layout) != cosize(tmem_warp_layout); + bool issue_tmem_load = true; + + // If tmem doesn't have enough capacity to support double buffering, a portion of tmem (a column of epilogue tiles) + // is overlapped between 2 pseudo-buffers. The shared tmem portion corresponds to the last epilogue tile column of + // tmem accumulator buffer 0, and the first epilogue tile column of tmem accumulator 1. + // Thus, whenever we are processing tmem accumulator buffer 0, we process the epilogue tiles with reversed column order. + // Once the last epilogue tile column is loaded from tmem, the acc_pipeline is released. + // Then, the next accumulation stage for buffer 1 can start. + [[maybe_unused]] bool reverse_epi_n = ReuseTmem && acc_pipe_consumer_state.phase() == 0; + static_assert(not (ReuseTmem && AccumulatorPipeline::Stages != 1), "Tmem reuse requires 1 accumulator stage"); + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = warp_idx == 0; + + // In the reuse smem configuration we have StagesC smem buffers and at most StagesD committed TMA stores in flight. + // The TMA store pipeline producer acquire returns when at most StagesD-1 committed stores are in-flight, so we can + // only guarantee store completion after StagesD iterations, then we can begin issuing releases on the smem buffer locks. + // store_pipe_producer_state tracks the acquire and load_pipe_consumer_state tracks the release, in circular buffer fashion. + // If TMA store supported async transaction mbarriers we would not need this synchronous release behavior. + LoadPipelineState load_wait_state = load_pipe_consumer_state; + if constexpr (ReuseSmemC) { + load_wait_state = store_pipe_producer_state; + load_wait_state.phase_ ^= 1; + } + + // We can delay issue of TMA store by one iteration to achieve better interleaving of non-TMA instructions + // Sync requirements of smem reuse may preclude this optimization + [[maybe_unused]] int epi_m_prev = 0; + [[maybe_unused]] int epi_n_prev = 0; + static_assert(not (DelayTmaStore and ReuseSmemC and StagesC <= StagesD), "This TMA epilogue configuration will deadlock"); + + // The TMA store sequence for one subtile iteration + auto tma_store_fn = [&] (int epi_m, int epi_n) { + // Write the tile from smem to gmem with TMA + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d_real, bSG_sD_real(_,_,_,store_pipe_producer_state.index()), bSG_gD_real(_,_,_,epi_m,epi_n)); + copy(params.tma_store_d_imag, bSG_sD_imag(_,_,_,store_pipe_producer_state.index()), bSG_gD_imag(_,_,_,epi_m,epi_n)); + } + + // Commit the TMA stores for this stage + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for the next smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + if constexpr (ReuseSmemC) { + // producer_acquire returns when at most StagesD-1 committed stores are pending + bool store_finished = store_pipe_producer_state.count() > StorePipeline::UnacquiredStages; + // Let dma warp know earliest smem buffer is consumed and empty after StagesD producer commits + if (store_finished) { + if (is_source_needed) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + ++load_pipe_consumer_state; + } + } + }; + + // + // BEGIN EPILOGUE + // + + // Begin the wait for the producer load results + ConsumerToken load_wait_token{BarrierStatus::WaitDone}; + if (is_source_needed) { + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state); + } + // Begin the wait for the accumulator results + ConsumerToken acc_wait_token = acc_pipeline.consumer_try_wait(acc_pipe_consumer_state); + + // For each epilogue subtile within the CTA tile + CUTLASS_PRAGMA_UNROLL + for (int iter_n = 0; iter_n < size<3>(gD_real_epi); ++iter_n) { + CUTLASS_PRAGMA_UNROLL + for (int iter_m = 0; iter_m < size<2>(gD_real_epi); ++iter_m) { + int epi_m = iter_m, epi_n = iter_n; + bool is_first_iteration = iter_m == 0 && iter_n == 0; + bool is_last_iteration = iter_m == size<2>(gD_real_epi)-1 && iter_n == size<3>(gD_real_epi)-1; + bool do_acc_release = is_last_iteration; + + // Reverse subtile order for tmem reuse if necessary + if constexpr (ReuseTmem) { + if (reverse_epi_n) { + epi_n = size<3>(gD_real_epi) - 1 - iter_n; + } + do_acc_release = iter_m == size<2>(gD_real_epi)-1 && iter_n == 0; + } + + if (is_source_needed) { + // Wait for the producer load to fill smem + load_pipeline.consumer_wait(load_wait_state, load_wait_token); + + // Copy source tile from smem to register // residual smem -> reg + copy(tiled_s2r, tSR_sC_real(_,_,_,load_wait_state.index()), tSR_rC(_,_,_,0)); + copy(tiled_s2r, tSR_sC_imag(_,_,_,load_wait_state.index()), tSR_rC(_,_,_,1)); + } + + if (is_source_needed) { + // Let producer load warp know smem buffers are consumed and empty + if constexpr (not ReuseSmemC) { + cutlass::arch::fence_view_async_shared(); + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + ++load_wait_state; + } + + if (is_first_iteration) { + // Wait for mma warp to fill tmem buffer with accumulator results + acc_pipeline.consumer_wait(acc_pipe_consumer_state, acc_wait_token); + } + + // The current tile in tmem + Tensor tTR_tAcc_real_mn = tTR_tAcc_real(_,_,_,epi_m,epi_n); + Tensor tTR_tAcc_imag_mn = tTR_tAcc_imag(_,_,_,epi_m,epi_n); + + // Compute tmem load predication if necessary + if constexpr (predicate_tmem_load) { + // Issue tmem load if this tile's tmem subpartition is accessible by this warp + int subpart_idx = (tTR_tAcc_real_mn.data().dp_ / 32) % 4; + issue_tmem_load = warp_idx == subpart_idx; + } + + // Copy accumulator tile from tmem to register + if (issue_tmem_load) { // acc tmem -> reg + copy(tiled_t2r, tTR_tAcc_real_mn, tTR_rAcc(_,_,_,0)); + copy(tiled_t2r, tTR_tAcc_imag_mn, tTR_rAcc(_,_,_,1)); + } + + // After the last tmem load, signal that tmem buffer is consumed and empty + if (do_acc_release) { + cutlass::arch::fence_view_async_tmem_load(); + acc_pipeline.consumer_release(acc_pipe_consumer_state); + ++acc_pipe_consumer_state; + } + + // Vectorized fragment loop with visitor callback entry point + if (epilogue_op.is_source_needed()) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rD_frg); ++i) { + tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i), tTR_rC_frg(i)); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rD_frg); ++i) { + tTR_rD_frg(i) = epilogue_op(tTR_rAcc_frg(i)); + } + } + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the previous subtile + if (not is_first_iteration) { + tma_store_fn(epi_m_prev, epi_n_prev); + } + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + + // Copy output tile from register to smem + bool issue_smem_store = issue_tmem_load; + if (issue_smem_store) { // after scale, reg -> smem + copy(tiled_r2s, tRS_rD(_,_,_,0), tRS_sD_real(_,_,_,store_pipe_producer_state.index())); + copy(tiled_r2s, tRS_rD(_,_,_,1), tRS_sD_imag(_,_,_,store_pipe_producer_state.index())); + } + + if constexpr (not DelayTmaStore) { + // Issue TMA stores for this subtile + tma_store_fn(epi_m, epi_n); + } + + if (is_source_needed) { + // Begin the wait for the next subtile producer load + load_wait_token = load_pipeline.consumer_try_wait(load_wait_state, is_last_iteration); + } + } // for epi_m + } // for epi_n + + if constexpr (DelayTmaStore) { + // Issue TMA stores for the last subtile + tma_store_fn(epi_m_prev, epi_n_prev); + } + + return cute::make_tuple(load_pipe_consumer_state, store_pipe_producer_state, acc_pipe_consumer_state); + } + + template + CUTLASS_DEVICE void + store_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + CtaTileMNK cta_tile_mnk) { + if constexpr (ReuseSmemC) { + if (epilogue_op.is_source_needed()) { + // wait for all TMA stores to complete + store_pipeline.producer_tail(store_pipe_producer_state); + + // Issue releases on up to StagesD-1 previously issued TMA stores + constexpr int release_stages = cute::min(StorePipeline::UnacquiredStages, get_load_pipe_increment(cta_tile_mnk)); + CUTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < release_stages; ++stage) { + load_pipeline.consumer_release(load_pipe_consumer_state); + ++load_pipe_consumer_state; + } + } + } + } + +private: + Params const& params; + ThreadEpilogueOp epilogue_op; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::epilogue::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/epilogue/dispatch_policy.hpp b/include/cutlass/epilogue/dispatch_policy.hpp index dc8ba912..f1a53e2a 100644 --- a/include/cutlass/epilogue/dispatch_policy.hpp +++ b/include/cutlass/epilogue/dispatch_policy.hpp @@ -71,11 +71,18 @@ struct PtrArrayFastF32NoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1S struct PtrArrayFastF32NoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {}; struct PtrArrayBlockwiseNoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1Sm {}; struct PtrArrayBlockwiseNoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {}; +struct PtrArrayPlanarComplexNoSmemWarpSpecialized1Sm : PtrArrayNoSmemWarpSpecialized1Sm {}; +struct PtrArrayPlanarComplexNoSmemWarpSpecialized2Sm : PtrArrayNoSmemWarpSpecialized2Sm {}; // Blackwell TMA schedules struct TmaWarpSpecialized1Sm {}; struct TmaWarpSpecialized2Sm {}; struct PtrArrayTmaWarpSpecialized1Sm : TmaWarpSpecialized1Sm {}; struct PtrArrayTmaWarpSpecialized2Sm : TmaWarpSpecialized2Sm {}; + +struct PlanarComplexTmaWarpSpecialized1Sm : TmaWarpSpecialized1Sm {}; +struct PlanarComplexTmaWarpSpecialized2Sm : TmaWarpSpecialized2Sm {}; +struct PtrArrayPlanarComplexTmaWarpSpecialized1Sm : PlanarComplexTmaWarpSpecialized1Sm {}; +struct PtrArrayPlanarComplexTmaWarpSpecialized2Sm : PlanarComplexTmaWarpSpecialized2Sm {}; struct TmaWarpSpecialized1SmNvf4 final : TmaWarpSpecialized1Sm {}; struct TmaWarpSpecialized2SmNvf4 final : TmaWarpSpecialized2Sm {}; struct TmaWarpSpecialized1SmMxf4 final : TmaWarpSpecialized1Sm {}; @@ -253,7 +260,6 @@ struct Sm100NoSmemWarpSpecialized { constexpr static int StagesD = 1; constexpr static int FragmentSize = 1; }; - struct Sm100PtrArrayNoSmem { constexpr static int StagesC = 1; constexpr static int StagesD = 1; @@ -265,6 +271,42 @@ struct Sm100PtrArrayNoSmemWarpSpecialized { constexpr static int StagesD = 1; constexpr static int FragmentSize = 1; }; +struct Sm100PtrArrayPlanarComplexNoSmem {}; +struct Sm100PtrArrayPlanarComplexNoSmemWarpSpecialized {}; + +template< + int StagesC_, + int StagesD_, + int FragmentSize_, + bool ReuseSmemC_, + bool DelayTmaStore_ +> +struct Sm100PlanarComplexTmaWarpSpecialized + : public Sm100TmaWarpSpecialized +{ +}; + + +template< + int StagesC_, + int StagesD_, + int FragmentSize_, + bool ReuseSmemC_, + bool DelayTmaStore_ +> +struct Sm100PtrArrayPlanarComplexTmaWarpSpecialized + : public Sm100TmaWarpSpecialized +{ +}; + template< int StagesC_, int StagesD_, diff --git a/include/cutlass/epilogue/threadblock/epilogue.h b/include/cutlass/epilogue/threadblock/epilogue.h index 5968d6cb..423996a9 100644 --- a/include/cutlass/epilogue/threadblock/epilogue.h +++ b/include/cutlass/epilogue/threadblock/epilogue.h @@ -39,7 +39,9 @@ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/numeric_types.h" #include "cutlass/array.h" diff --git a/include/cutlass/epilogue/threadblock/epilogue_base.h b/include/cutlass/epilogue/threadblock/epilogue_base.h index 94c0a4d1..b6b63608 100644 --- a/include/cutlass/epilogue/threadblock/epilogue_base.h +++ b/include/cutlass/epilogue/threadblock/epilogue_base.h @@ -42,7 +42,9 @@ #include #include #endif +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/matrix_shape.h" #include "cutlass/numeric_types.h" diff --git a/include/cutlass/functional.h b/include/cutlass/functional.h index 6b23b205..2ac6e653 100644 --- a/include/cutlass/functional.h +++ b/include/cutlass/functional.h @@ -38,6 +38,8 @@ #include "cutlass/cutlass.h" #include "cutlass/numeric_types.h" #include "cutlass/platform/platform.h" +#include "cutlass/detail/dependent_false.hpp" + #if defined(__CUDACC_RTC__) #include "cutlass/floating_point_nvrtc.h" #endif @@ -535,6 +537,49 @@ struct maximum_absolute_value_reduction { } }; +// Maximal exponent reduction for zero-mantissa scaling factors +template +struct maximum_absolute_value_zero_mantissa_reduction { + + // Discard mantissa and sign bits for the input. Needs to specify the number of mantissa / exponent bits + template + static CUTLASS_HOST_DEVICE T_ discard_sign_mantissa_impl(T_ x) { + static constexpr UI one = 1; + static constexpr UI n_mantissa = N_Mantissa; + static constexpr UI pos_sign = sizeof(T_) * 8 - 1; // Position of sign bit: bit width - 1. + static constexpr UI mask = ~((one << n_mantissa) - one) & ~(one << pos_sign); + static constexpr UI subnormal_cap = one << n_mantissa; + + UI out = *reinterpret_cast(&x) & mask; + // Subnormals + if (out == 0) { + out = subnormal_cap; + } + return *reinterpret_cast(&out); + } + + // Discard mantissa and sign bits s.t. multipling with this scaling factor only results in an exponent shift + template + static CUTLASS_HOST_DEVICE T_ discard_sign_mantissa(T_ x) { + if constexpr (cute::is_same_v) { + return discard_sign_mantissa_impl(x); + } + else if constexpr (cute::is_same_v) { + return discard_sign_mantissa_impl(x); + } + else { + static_assert(cutlass::detail::dependent_false, "Can't discard mantissa & sign bits for unknown data type"); + } + } + + CUTLASS_HOST_DEVICE + T operator()(T const &lhs, T const &rhs) const { + cutlass::maximum max_op; + + return max_op(lhs, discard_sign_mantissa(rhs)); + } +}; + /// Fused multiply-add template struct multiply_add { diff --git a/include/cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl b/include/cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl new file mode 100644 index 00000000..5cdbdd7e --- /dev/null +++ b/include/cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl @@ -0,0 +1,298 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/gemm/collective/builders/sm100_common.inl" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// FastFP (9xBF16) TCGEN05 kernels builder +// Interleaved complex kernels that provides support for complex data types +template < + class ArchTag, + class GmemLayoutATag, + int AlignmentA, + class TransformA, + class GmemLayoutBTag, + int AlignmentB, + class TransformB, + class ElementAccumulator, + class TileShape_MNK, // The Cluster-level TileShape + class ClusterShape_MNK, + class StageCountType, + class BuilderScheduleTag +> +struct CollectiveBuilder< + ArchTag, + arch::OpClassTensorOp, + cute::tuple, TransformA>, // ElementA + ConjA + GmemLayoutATag, // LayoutA + AlignmentA, + cute::tuple, TransformB>, // ElementB + ConjB + GmemLayoutBTag, // LayoutB + AlignmentB, + ElementAccumulator, + TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + ClusterShape_MNK, // Static cluster shape or dynamic (int, int, int) + StageCountType, + BuilderScheduleTag, + cute::enable_if_t< + (cute::is_same_v + ) && + (not cute::is_tuple::value && not cute::is_tuple::value) && + (cute::is_base_of_v + ) && + ((sizeof(cutlass::complex) * AlignmentA) % detail::tma_alignment_bytes == 0) && + ((sizeof(cutlass::complex) * AlignmentB) % detail::tma_alignment_bytes == 0)>> +{ + static_assert(cute::is_static_v, "TileShape_MNK has to be static"); + static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A(); + static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B(); + static constexpr cute::UMMA::Major UmmaMajorACompute = cute::UMMA::Major::K; + static constexpr cute::UMMA::Major UmmaMajorBCompute = cute::UMMA::Major::K; + static constexpr bool BuilderTagIsSmem = + cute::is_base_of_v + ; + + using ElementA = complex; + using ElementB = complex; + using ElementAMma = complex< + cutlass::bfloat16_t + >; + using ElementBMma = complex< + cutlass::bfloat16_t + >; + static constexpr int ScalingFactor = + 8; + + using TiledMma = decltype(detail::sm100_make_trivial_fastFP32_tiled_mma()); + using AtomThrID = typename TiledMma::AtomThrID; + using AtomThrShapeMNK = Shape(typename TiledMma::ThrLayoutVMNK{})), _1, _1>; + using CtaTileShape_MNK = decltype(shape_div(TileShape_MNK{}, AtomThrShapeMNK{})); + + // ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K) + using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}), + cute::size<2>(TileShape_MNK{})))); + // ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K) + using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}), + cute::size<2>(TileShape_MNK{})))); + + using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{})); + using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{})); + + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + // Take 3 compute buffers into account for swizzle selection + using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + + // Input transform kernel can not use TMA 2SM instructions. + using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<1>(ClusterShape_MNK{}))); + using SmemLayoutAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType< + SmemLayoutAtomA, SmemLayoutAtomACompute>; + + static constexpr int MMA_M = cute::size<0,0>(MmaShapeA_MK{}); + using CopyAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType< + Copy_Atom, ElementA>, + cute::conditional_t<(UmmaMajorACompute == cute::UMMA::Major::K && !BuilderTagIsSmem), + cute::conditional_t<(MMA_M == 64 && size(AtomThrID{}) == 1), SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b8x>, // TS Implementation + Copy_Atom, ElementAMma>> // SS Implementation + >; + + using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{})); + using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{})); + + // Input transform kernel can not use TMA 2SM instructions. + using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<0>(ClusterShape_MNK{}))); + + using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + // Take 3 compute buffers into account for swizzle selection + using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + + using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType< + SmemLayoutAtomB, SmemLayoutAtomBCompute>; + using CopyAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType< + Copy_Atom, ElementB>, + Copy_Atom, ElementBMma> + >; + + // SmemCarveout + static constexpr int NumComplexComponents = 2; + static constexpr int NumComputeMtxs = + 3; + static constexpr int NumBandsToCompute = + 5; + static constexpr int AccPromotionInterval = + 1; + static constexpr int SchedulerPipelineStageCount = 3; + static constexpr bool IsArrayOfPointersGemm = + (cute::is_base_of_v + ); + // CLCPipeline = PipelineCLCFetchAsync + static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage); + // CLC (scheduler) response + static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize; + // CLC Throttle pipeline storage + static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync::SharedStorage); + // Tmem dealloc + static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier); + // Tmem ptr storage + static constexpr auto TmemBasePtrsStorage = sizeof(uint32_t); + // Tensormap Storage + static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0; + + // Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage + static constexpr auto KernelSmemCarveout = static_cast( CLCPipelineStorage + + CLCResponseStorage + + CLCThrottlePipelineStorage + + TmemDeallocStorage + + TmemBasePtrsStorage + + TensorMapStorage); + + // Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations + + static constexpr int ReducedSmemCapacityBytes = + cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout; + static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32< + ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute, + NumComplexComponents, NumComputeMtxs + >(StageCountType{}); + + // Complex 9xBF16 allows TileShape_N = 64, while SmemLayoutAtomB contains Swizzle<3,4,3>. + static constexpr int Load2TransformPipelineStageCount = size<1>(TileShape_MNK{}) == 64 ? get<0>(stage_info) / 2 * 2 : get<0>(stage_info); + static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info); + static constexpr int AccumulatorPipelineStageCount = get<2>(stage_info); + + using AccumulatorCopyAtom = cute::SM100_TMEM_LOAD_32dp32b32x; + + using DispatchPolicy = cute::conditional_t, + cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedFastF32< + Load2TransformPipelineStageCount, + Transform2MmaPipelineStageCount, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + NumBandsToCompute, + ScalingFactor, + AccPromotionInterval, + ClusterShape_MNK, + AccumulatorCopyAtom> + >; + using CollectiveOp = cutlass::gemm::collective::CollectiveMma< + DispatchPolicy, + TileShape_MNK, + ElementA, + cutlass::gemm::TagToStrideA_t, + ElementB, + cutlass::gemm::TagToStrideB_t, + TiledMma, + GmemTiledCopyA, + SmemLayoutAtomPairA, + CopyAtomPairA, + TransformA, + GmemTiledCopyB, + SmemLayoutAtomPairB, + CopyAtomPairB, + TransformB + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// FastFP (9xBF16) TCGEN05 kernels builder +// CUTLASS library compatibility builder without conjugate +template < + class ArchTag, + class GmemLayoutATag, + int AlignmentA, + class GmemLayoutBTag, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, // The Cluster-level TileShape + class ClusterShape_MNK, + class StageCountType, + class BuilderScheduleTag +> +struct CollectiveBuilder< + ArchTag, + arch::OpClassTensorOp, + cutlass::complex, // ElementA + GmemLayoutATag, // LayoutA + AlignmentA, + cutlass::complex, // ElementB + GmemLayoutBTag, // LayoutB + AlignmentB, + ElementAccumulator, + TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + ClusterShape_MNK, // Static cluster shape or dynamic (int, int, int) + StageCountType, + BuilderScheduleTag, + cute::enable_if_t< + (cute::is_same_v + ) && + (not cute::is_tuple::value && not cute::is_tuple::value) && + (cute::is_base_of_v + ) && + ((sizeof(cutlass::complex) * AlignmentA) % detail::tma_alignment_bytes == 0) && + ((sizeof(cutlass::complex) * AlignmentB) % detail::tma_alignment_bytes == 0)>> +{ + using CollectiveOp = typename CollectiveBuilder< + ArchTag, + arch::OpClassTensorOp, + cute::tuple, cute::identity>, + GmemLayoutATag, + AlignmentA, + cute::tuple, cute::identity>, + GmemLayoutBTag, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + BuilderScheduleTag + >::CollectiveOp; +}; + +} // cutlass::gemm::collective diff --git a/include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl b/include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl index c4bcf591..4b345457 100644 --- a/include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl +++ b/include/cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl @@ -61,15 +61,19 @@ sm100_compute_stage_count_or_override_fast_fp32(StageCountAutoCarveout(CtaTileShape_MNK{}); using AtomThrID = typename TiledMma::AtomThrID; constexpr int TmemColumns = 512; + constexpr bool BuilderTagIsSmem = ( + cute::is_base_of_v + ); // Detect 2x2 TMEM layout constexpr int TmemAccWordsPerDP = (CtaM == 64 && size(AtomThrID{}) == 2) ? CtaN/2 : CtaN; constexpr int TmemAWordsPerDP = ComplexComponent * NumComputeMtxs * CtaK / 2; - constexpr bool IsAComputeinTmem = UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v; + constexpr bool IsAComputeinTmem = UmmaMajorA == cute::UMMA::Major::K && !BuilderTagIsSmem; constexpr bool IsAComputeinSmem = !IsAComputeinTmem; constexpr int AccumulatorStageCount = (IsAComputeinTmem) ? (((TmemAccWordsPerDP * ComplexComponent == 128) ? 2 : 3) * ComplexComponent) : (TmemColumns / TmemAccWordsPerDP); - - constexpr int SmemCapacityAfterMma2AccumCarveout = CapacityBytes - (carveout_bytes + AccumulatorStageCount * 32); + + constexpr int SmemCapacityAfterMma2AccumCarveout = CapacityBytes - (carveout_bytes + AccumulatorStageCount * (32 + )); constexpr int TmemInAStageCount_Potential = (IsAComputeinTmem) ? (TmemColumns - AccumulatorStageCount * TmemAccWordsPerDP) / TmemAWordsPerDP : 10000; @@ -87,7 +91,8 @@ sm100_compute_stage_count_or_override_fast_fp32(StageCountAutoCarveout(CtaTileShape_MNK{}) * size<2>(CtaTileShape_MNK{})) + // If ACompute is in TMEM, Acompute buffer has 0 bytes. cutlass::bits_to_bytes(NumComputeMtxs * b_compute_bits * size<1>(CtaTileShape_MNK{}) / size(AtomThrID{}) * size<2>(CtaTileShape_MNK{})) + - static_cast(transform2mma_pipeline_bytes); + static_cast(transform2mma_pipeline_bytes) + ; constexpr int ABComputeStageCount_Potential = SmemCapacityAfterMma2AccumCarveout / (ab_stage_bytes + ab_compute_stage_bytes); // The number of SMEM buffers for A, B. ACompute (if in SMEM), BCompute should be at least Transform2MmaStageCount @@ -135,20 +140,33 @@ struct CollectiveBuilder< (cute::is_same_v ) && (not cute::is_tuple::value && not cute::is_tuple::value) && - (cute::is_base_of_v) && + (cute::is_base_of_v + ) && ((sizeof(float) * AlignmentA) % detail::tma_alignment_bytes == 0) && ((sizeof(float) * AlignmentB) % detail::tma_alignment_bytes == 0)>> { static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A(); static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B(); + static constexpr cute::UMMA::Major UmmaMajorACompute = + UmmaMajorA; + static constexpr cute::UMMA::Major UmmaMajorBCompute = + UmmaMajorB; + static constexpr bool BuilderTagIsSmem = ( + cute::is_base_of_v + ); using ElementA = float; using ElementB = float; - using ElementAMma = cutlass::bfloat16_t; - using ElementBMma = cutlass::bfloat16_t; - static constexpr int ScalingFactor = 8; + using ElementAMma = + cutlass::bfloat16_t + ; + using ElementBMma = + cutlass::bfloat16_t + ; + static constexpr int ScalingFactor = + 8; - using TiledMma = decltype(detail::sm100_make_trivial_fastFP32_tiled_mma()); + using TiledMma = decltype(detail::sm100_make_trivial_fastFP32_tiled_mma()); using AtomThrID = typename TiledMma::AtomThrID; using AtomThrShapeMNK = Shape(typename TiledMma::ThrLayoutVMNK{})), _1, _1>; using CtaTileShape_MNK = decltype(shape_div(TileShape_MNK{}, AtomThrShapeMNK{})); @@ -166,7 +184,7 @@ struct CollectiveBuilder< using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); // Take 3 compute buffers into account for swizzle selection - using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); // Input transform kernel can not use TMA 2SM instructions. @@ -177,7 +195,7 @@ struct CollectiveBuilder< static constexpr int MMA_M = cute::size<0,0>(MmaShapeA_MK{}); using CopyAtomPairA = cutlass::gemm::collective::detail::CollectiveMmaEmulatedCopyType< Copy_Atom, ElementA>, - cute::conditional_t<(UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v), + cute::conditional_t<(UmmaMajorACompute == cute::UMMA::Major::K && !BuilderTagIsSmem), cute::conditional_t<(MMA_M == 64 && size(AtomThrID{}) == 1), SM100_TMEM_STORE_16dp256b1x, SM100_TMEM_STORE_32dp32b8x>, // TS Implementation Copy_Atom, ElementA>> // SS Implementation >; @@ -191,7 +209,7 @@ struct CollectiveBuilder< using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); // Take 3 compute buffers into account for swizzle selection - using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::CollectiveMmaEmulatedLayoutAtomType< @@ -202,11 +220,16 @@ struct CollectiveBuilder< >; // SmemCarveout - static constexpr int NumBandsToCompute = 5; - static constexpr int AccPromotionInterval = 1; + static constexpr int NumComputeMtxs = + 3; + static constexpr int NumBandsToCompute = + 5; + static constexpr int AccPromotionInterval = + 1; static constexpr int SchedulerPipelineStageCount = 3; - static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v); - + static constexpr bool IsArrayOfPointersGemm = + (cute::is_base_of_v + ); // CLCPipeline = PipelineCLCFetchAsync static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage); // CLC (scheduler) response @@ -233,7 +256,9 @@ struct CollectiveBuilder< static constexpr int ReducedSmemCapacityBytes = cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout; static constexpr auto stage_info = cutlass::gemm::collective::detail::sm100_compute_stage_count_or_override_fast_fp32< - ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorA>(StageCountType{}); + ReducedSmemCapacityBytes, CtaTileShape_MNK, TiledMma, BuilderScheduleTag, UmmaMajorACompute, + /*Cmplx=*/ 1, /*Mtxs=*/ NumComputeMtxs + >(StageCountType{}); static constexpr int Load2TransformPipelineStageCount = get<0>(stage_info); static constexpr int Transform2MmaPipelineStageCount = get<1>(stage_info); diff --git a/include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl b/include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl index 78b66196..923ce9cf 100644 --- a/include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl +++ b/include/cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl @@ -73,7 +73,7 @@ sm100_compute_stage_count_or_override_blockscaled(StageCountAutoCarveout::SharedStorage); + constexpr auto mainloop_pipeline_bytes = cutlass::round_up(sizeof(typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage), 128); constexpr auto a_bits = cute::sizeof_bits_v; constexpr auto b_bits = cute::sizeof_bits_v; constexpr auto stage_sfa_bytes = size(filter_zeros(TileShapeSFA{})); diff --git a/include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl b/include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl index e3418444..fc9911d9 100644 --- a/include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl +++ b/include/cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl @@ -102,8 +102,10 @@ sm100_compute_stage_count_or_override_blockwise(StageCountAutoCarveout(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + - cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{}) + ) + + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{}) + ) + cutlass::bits_to_bytes(scale_bits * size<0>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})) + cutlass::bits_to_bytes(scale_bits * size<1>(ScaleShapeMNK{}) * size<2>(ScaleShapeMNK{})), 128) + @@ -441,7 +443,6 @@ struct CollectiveBuilder< >; }; - } // namespace cutlass::gemm::collective ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/builders/sm100_common.inl b/include/cutlass/gemm/collective/builders/sm100_common.inl index 73f43797..284def11 100644 --- a/include/cutlass/gemm/collective/builders/sm100_common.inl +++ b/include/cutlass/gemm/collective/builders/sm100_common.inl @@ -220,7 +220,7 @@ sm100_cluster_shape_to_tma_atom_A(ClusterShapeMNK cluster_shape_mnk, AtomThrId a } else { // In the case of dynamic cluster, multicast decision is not known at compile time. - // A multicast instruction is forced by passing a cute::Int<2>{} to this helper. + // A multicast instruction is forced by passing a cute::Int<2>{} to this helper. return detail::sm90_cluster_shape_to_tma_atom(cute::Int<2>{}); } } @@ -255,7 +255,7 @@ sm100_cluster_shape_to_tma_atom_B(ClusterShapeMNK cluster_shape_mnk, AtomThrId a } else { // In the case of dynamic cluster, multicast decision is not known at compile time. - // A multicast instruction is forced by passing a cute::Int<2>{} to this helper. + // A multicast instruction is forced by passing a cute::Int<2>{} to this helper. return detail::sm90_cluster_shape_to_tma_atom(cute::Int<2>{}); } } @@ -281,7 +281,7 @@ sm100_cluster_shape_to_tma_atom_SFB(ClusterShapeMNK cluster_shape_mnk, AtomThrId } else { // In the case of dynamic cluster, multicast decision is not known at compile time. - // A multicast instruction is forced by passing a cute::Int<2>{} to this helper. + // A multicast instruction is forced by passing a cute::Int<2>{} to this helper. return detail::sm90_cluster_shape_to_tma_atom(cute::Int<2>{}); } } @@ -328,24 +328,24 @@ sm100_make_1sm_trivial_tiled_mma() { return make_tiled_mma(cute::SM100_MMA_S8_SS{}); } - else if constexpr (cute::is_same_v - || cute::is_same_v - || cute::is_same_v + else if constexpr (cute::is_same_v + || cute::is_same_v + || cute::is_same_v || cute::is_same_v || cute::is_same_v || cute::is_same_v || cute::is_same_v || cute::is_same_v ) { - + return make_tiled_mma( cute::MMA_Traits< cute::SM100_MMA_F8F6F4_SS, ElementAMma, - ElementBMma, - ElementAMmaccumulator, - cute::C, - cute::C, + ElementBMma, + ElementAMmaccumulator, + cute::C, + cute::C, cute::integral_constant, cute::integral_constant, cute::integral_constant, @@ -396,9 +396,9 @@ sm100_make_2sm_trivial_tiled_mma() { return make_tiled_mma(cute::SM100_MMA_S8_2x1SM_SS{}); } - else if constexpr (cute::is_same_v - || cute::is_same_v - || cute::is_same_v + else if constexpr (cute::is_same_v + || cute::is_same_v + || cute::is_same_v || cute::is_same_v || cute::is_same_v || cute::is_same_v @@ -408,12 +408,12 @@ sm100_make_2sm_trivial_tiled_mma() { return make_tiled_mma( cute::MMA_Traits< - cute::SM100_MMA_F8F6F4_2x1SM_SS, + cute::SM100_MMA_F8F6F4_2x1SM_SS, ElementAMma, ElementBMma, - ElementAMmaccumulator, - cute::C, - cute::C, + ElementAMmaccumulator, + cute::C, + cute::C, cute::integral_constant, cute::integral_constant, cute::integral_constant, @@ -471,7 +471,7 @@ sm100_make_trivial_tiled_mma() { return sm100_make_1sm_trivial_tiled_mma(); } - // Dynamic cluster shape means we cannot assume we can use 2SM MMA + // Dynamic cluster shape means we cannot assume we can use 2SM MMA } else { return sm100_make_1sm_trivial_tiled_mma +constexpr auto +sm100_make_1sm_ts_trivial_tiled_mma() { + + constexpr int M = cute::size<0>(ClusterTileShape_MNK{}); + static_assert(M == 64 || M == 128, "Invalid TileShape_M."); + + // Do not allow a tiled MMA N mode > 1, as that is not reasonable. + constexpr int N = cute::size<1>(ClusterTileShape_MNK{}); + static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N."); + + if constexpr (cute::is_same_v) { + static_assert(cute::is_same_v, "ElementA and ElementB must match."); + return make_tiled_mma(cute::SM100_MMA_TF32_TS{}); + } + else if constexpr (cute::is_same_v || + cute::is_same_v) { + static_assert(cute::is_same_v, "ElementA and ElementB must match."); + return make_tiled_mma(cute::SM100_MMA_F16BF16_TS{}); + } + else if constexpr (cute::is_same_v || + cute::is_same_v) { + return make_tiled_mma(cute::SM100_MMA_S8_TS{}); + } + else { + static_assert(cutlass::detail::dependent_false, + "Unsupported configuration for SM100 collective builder."); + } +} + +template< + class ElementAMma, + class ElementBMma, + class ElementAccumulator, + class ClusterTileShape_MNK, + class ClusterShape_MNK, + UMMA::Major UmmaMajorA, + UMMA::Major UmmaMajorB, + UMMA::ScaleIn ANeg = UMMA::ScaleIn::One, + UMMA::ScaleIn BNeg = UMMA::ScaleIn::One +> +constexpr auto +sm100_make_2sm_ts_trivial_tiled_mma() { + + constexpr int M = cute::size<0>(ClusterTileShape_MNK{}); + static_assert(M == 128 || M == 256, "Invalid TileShape_M."); + + // Do not allow a tiled MMA N mode > 1, as that is not reasonable. + constexpr int N = cute::size<1>(ClusterTileShape_MNK{}); + static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N."); + + if constexpr (cute::is_same_v) { + static_assert(cute::is_same_v, "For SM100 TF32 MMA, ElementA and ElementB must match."); + return make_tiled_mma(cute::SM100_MMA_TF32_2x1SM_TS{}); + } + else if constexpr (cute::is_same_v || + cute::is_same_v) { + static_assert(cute::is_same_v, "For SM100 F16F32 MMA, ElementA and ElementB must match."); + return make_tiled_mma(cute::SM100_MMA_F16BF16_2x1SM_TS{}); + } + else if constexpr (cute::is_same_v || + cute::is_same_v) { + return make_tiled_mma(cute::SM100_MMA_S8_2x1SM_TS{}); + } + else { + static_assert(cutlass::detail::dependent_false, + "Unsupported configuration for SM100 collective builder."); + } +} + template< class ElementAMma, class ElementBMma, @@ -493,12 +579,16 @@ template< > constexpr auto sm100_make_trivial_fastFP32_tiled_mma() { + constexpr bool TagHasUmmaSs = ( + cute::is_base_of_v + ); + // MMA_2SM requested if constexpr (cute::is_base_of_v ) { using AtomLayout_MNK = decltype(make_layout(shape_div(ClusterShape_MNK{}, Shape<_2,_1,_1>{}))); constexpr int M = cute::size<0>(TileShape_MNK{}); constexpr int N = cute::size<1>(TileShape_MNK{}); - if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v) { + if constexpr (UmmaMajorA == cute::UMMA::Major::K && !TagHasUmmaSs) { return make_tiled_mma(cute::SM100_MMA_F16BF16_2x1SM_TS_SCALED{}); } @@ -512,7 +602,7 @@ sm100_make_trivial_fastFP32_tiled_mma() { // using AtomLayout_MNK = Layout; constexpr int M = cute::size<0>(TileShape_MNK{}); constexpr int N = cute::size<1>(TileShape_MNK{}); - if constexpr (UmmaMajorA == cute::UMMA::Major::K && !cute::is_base_of_v) { + if constexpr (UmmaMajorA == cute::UMMA::Major::K && !TagHasUmmaSs) { return make_tiled_mma(cute::SM100_MMA_F16BF16_TS_SCALED{}); } @@ -531,7 +621,7 @@ sm100_make_trivial_fastFP32_tiled_mma() { // and only M=128 and M=256 are supported, otherwise, fall back to MMA_1SM if constexpr (cute::get<0>(ClusterShape_MNK{}) % 2 == 0 && (cute::get<0>(TileShape_MNK{}) / cute::get<0>(ClusterShape_MNK{})) % 64 == 0) { - if constexpr (!cute::is_base_of_v) { + if constexpr (!TagHasUmmaSs) { return sm100_make_trivial_fastFP32_tiled_mma(); } @@ -541,7 +631,7 @@ sm100_make_trivial_fastFP32_tiled_mma() { } } else { - if constexpr (!cute::is_base_of_v) { + if constexpr (!TagHasUmmaSs) { return sm100_make_trivial_fastFP32_tiled_mma(); } @@ -551,9 +641,9 @@ sm100_make_trivial_fastFP32_tiled_mma() { } } } - // Dynamic cluster shape means we cannot assume we can use 2SM MMA + // Dynamic cluster shape means we cannot assume we can use 2SM MMA else { - if constexpr (!cute::is_base_of_v) { + if constexpr (!TagHasUmmaSs) { return sm100_make_trivial_fastFP32_tiled_mma(); } @@ -569,6 +659,52 @@ sm100_make_trivial_fastFP32_tiled_mma() { } } +template< + class TileShape_MNK, + class ClusterShape_MNK, + class BuilderScheduleTag +> +constexpr auto +sm100_make_trivial_interleaved_complex_tf32_tiled_mma() { + // MMA_2SM requested + if constexpr (cute::is_base_of_v ) { + constexpr int M = cute::size<0>(TileShape_MNK{}); + static_assert(M == 128 || M == 256, "Invalid TileShape_M."); + + // Do not allow a tiled MMA N mode > 1, as that is not reasonable. + constexpr int N = cute::size<1>(TileShape_MNK{}); + static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N."); + return make_tiled_mma(cute::SM100_MMA_TF32_2x1SM_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN{}); + } + // MMA_1SM requested + else if constexpr (cute::is_base_of_v ) { + constexpr int M = cute::size<0>(TileShape_MNK{}); + static_assert(M == 64 || M == 128, "Invalid TileShape_M."); + // Do not allow a tiled MMA N mode > 1, as that is not reasonable. + constexpr int N = cute::size<1>(TileShape_MNK{}); + static_assert(N % 8 == 0 && N <= 256, "Invalid TileShape_N."); + return make_tiled_mma(cute::SM100_MMA_TF32_TS_INTERLEAVED_CF32CTF32CTF32CF32_TN{}); + } + else if constexpr (cute::is_same_v) { + // Static cluster + if constexpr (cute::is_static_v) { + // For MMA_2SM we need a cluster shape that is multiple of 2x1 + // and only M=128 and M=256 are supported, otherwise, fall back to MMA_1SM + if constexpr (cute::get<0>(ClusterShape_MNK{}) % 2 == 0 && + cute::size<0>(TileShape_MNK{}) % 128 == 0) { + return sm100_make_trivial_interleaved_complex_tf32_tiled_mma(); + } + else { + return sm100_make_trivial_interleaved_complex_tf32_tiled_mma(); + } + } + // Dynamic cluster shape means we cannot assume we can use 2SM MMA + else { + return sm100_make_trivial_interleaved_complex_tf32_tiled_mma(); + } + } +} + //Setting mma for Mixed input gemm. Here, ElementAMma should be TACompute template< class ElementAMma, @@ -606,10 +742,10 @@ sm100_make_trivial_mixed_input_tiled_mma() { cute::MMA_Traits< cute::SM100_MMA_F8F6F4_SS, ElementAMma, - ElementBMma, - ElementAccumulator, - cute::C, - cute::C, + ElementBMma, + ElementAccumulator, + cute::C, + cute::C, cute::integral_constant, cute::integral_constant, cute::integral_constant, diff --git a/include/cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl b/include/cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl new file mode 100644 index 00000000..e6801432 --- /dev/null +++ b/include/cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl @@ -0,0 +1,264 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/gemm/collective/builders/sm100_common.inl" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { + +namespace detail { + +// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. +template< + int CapacityBytes, + class TileShapeMNK, + int stages +> +constexpr int +sm100_compute_stage_count_or_override_interleaved_complex_tf32(StageCount stage_count) { + return stages; +} + +// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. +template< + int CapacityBytes, + class TileShapeMNK, + int carveout_bytes +> +constexpr int +sm100_compute_stage_count_or_override_interleaved_complex_tf32(StageCountAutoCarveout stage_count) { + // Each stage include (CollectiveMma::SharedStorage) + // 1. smem for A and smem for B (CollectiveMma::SharedStorage::TensorStorage) + // 2. one Load2TransformPipeline = PipelineTmaTransformAsync + constexpr auto load2transform_pipeline_bytes = sizeof(typename cutlass::PipelineTmaTransformAsync<1>::SharedStorage); + constexpr auto a_bits = cute::sizeof_bits_v>; + constexpr auto b_bits = cute::sizeof_bits_v>; + constexpr int stage_bytes = + cutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + cutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + static_cast(load2transform_pipeline_bytes); + + return (CapacityBytes - carveout_bytes) / stage_bytes; +} + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Interleaved complex tf32 TCGEN05 kernels builder +template < + class ArchTag, + class GmemLayoutATag, + class TransformA, + class GmemLayoutBTag, + class TransformB, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class BuilderScheduleTag +> +struct CollectiveBuilder< + ArchTag, + arch::OpClassTensorOp, + cute::tuple, TransformA>, + GmemLayoutATag, + 2, + cute::tuple, TransformB>, + GmemLayoutBTag, + 2, + cutlass::complex, + TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1) + StageCountType, + BuilderScheduleTag, + cute::enable_if_t< + (cute::is_same_v + ) && + (cute::is_base_of_v || + cute::is_same_v)>> +{ + static_assert(cute::is_static_v, "TileShape_MNK has to be static"); + // ElementA and ElementB are cutlass::complex, which are used as GMEM input and output data type. + // ElementAMma and ElementBMma are cutlass::complex, which are used as SMEM and RF data type. + using ElementA = complex; + using ElementB = complex; + using ElementAccumulator = cutlass::complex; + using ElementAMma = complex; + using ElementBMma = complex; + + static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A(); + static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B(); + + using TiledMma = decltype(detail::sm100_make_trivial_interleaved_complex_tf32_tiled_mma< + TileShape_MNK,ClusterShape_MNK,BuilderScheduleTag>()); + + using AtomThrID = typename TiledMma::AtomThrID; + + // Define A and B block shapes for reduced size TMA_LOADs + // ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K) + using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}), + cute::size<2>(TileShape_MNK{})))); + // ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K) + using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}), + cute::size<2>(TileShape_MNK{})))); + // Input transform kernel can not use TMA 2SM instructions. + using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<1>(ClusterShape_MNK{}))); + + using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{})); + using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{})); + + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + using SmemLayoutAtomACompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + using SmemLayoutAtomPairA = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexLayoutAtomType; + + static constexpr int MMA_M = cute::size<0>(TileShape_MNK{}); + + using CopyAtomPairA = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexCopyType< + Copy_Atom, ElementAMma>, + conditional_t + >; + + // Input transform kernel can not use TMA 2SM instructions. + using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(cute::size<0>(ClusterShape_MNK{}))); + + using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{})); + using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{})); + + using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + using SmemLayoutAtomBCompute = decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + using SmemLayoutAtomPairB = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexLayoutAtomType; + + using CopyAtomPairB = cutlass::gemm::collective::detail::Sm100CollectiveMmaComplexCopyType< + Copy_Atom, ElementBMma>, + Copy_Atom, ElementBMma> + >; + + // Calculate SMEM matrix A and B buffers' pipeline stages + static constexpr int MMA_N = cute::size<1>(TileShape_MNK{}); + static constexpr uint32_t TransformationStageCount = 4; + static constexpr uint32_t AccumulatorPipelineStageCount = (MMA_N >= 128) ? 1 : 2; + static constexpr uint32_t SchedulerPipelineStageCount = 3; + static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v); + + // SmemCarveout + // B needs extra smem for smem tranpose (CollectiveMma::TensorStorageTransformed) + static constexpr auto TensorStorageTransformedSmemBStorage = TransformationStageCount * + static_cast(sizeof(ElementBMma)) * size<0>(BlockTileB_N{}) * size<0>(BlockTileA_K{}); + // CLCPipeline = PipelineCLCFetchAsync + static constexpr auto CLCPipelineStorage = sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage); + // Transform2MmaPipeline = PipelineUmmaConsumerAsync (CollectiveMma) + static constexpr auto Transform2MmaPipelineStorage = sizeof(typename cutlass::PipelineUmmaConsumerAsync::SharedStorage); + // Mma2AccumPipeline = PipelineUmmaAsync (CollectiveMma) + static constexpr auto Mma2AccumPipelineStorage = sizeof(typename cutlass::PipelineUmmaAsync::SharedStorage); + // CLC (scheduler) response + static constexpr auto CLCResponseStorage = SchedulerPipelineStageCount * detail::CLCResponseSize; + // CLC Throttle pipeline storage + static constexpr auto CLCThrottlePipelineStorage = sizeof(typename cutlass::PipelineAsync::SharedStorage); + // Tmem dealloc + static constexpr auto TmemDeallocStorage = sizeof(cutlass::arch::ClusterBarrier); + // Tmem ptr storage + static constexpr auto TmemBasePtrsStorage = sizeof(uint32_t); + // Tensormap Storage + + static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0; + + // Smem usage that's not part of CollectiveEpilogue::SharedStorage & CollectiveMainloop::SharedStorage + static constexpr auto KernelSmemCarveout = static_cast( CLCPipelineStorage + + CLCResponseStorage + + CLCThrottlePipelineStorage + + Transform2MmaPipelineStorage + + Mma2AccumPipelineStorage + + TensorStorageTransformedSmemBStorage + + TmemDeallocStorage + + TmemBasePtrsStorage + + TensorMapStorage); + // Reduce SMEM capacity available for buffers considering extra B smem and barrier smem allocations + + static constexpr int ReducedSmemCapacityBytes = + cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout; + + using SmemTileShape = cute::Shape; + + static constexpr int PipelineStages_ = detail::sm100_compute_stage_count_or_override_interleaved_complex_tf32< + ReducedSmemCapacityBytes, SmemTileShape>(StageCountType{}); + // Complex kernels allow TileShape_N = 64, while SmemLayoutAtomB contains Swizzle<3,4,3>. + static constexpr int PipelineStages = size<1>(TileShape_MNK{}) == 64 ? PipelineStages_ / 2 * 2 : PipelineStages_; + static_assert(PipelineStages >= 2, "Pipeline Stages has to be at least 2"); + + using AccumulatorCopyAtom = cute::SM100_TMEM_LOAD_16dp256b1x; + + using DispatchPolicy = cute::conditional_t, + cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32< + PipelineStages, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + TransformationStageCount, + ClusterShape_MNK, + AccumulatorCopyAtom + > + >; + + using CollectiveOp = cutlass::gemm::collective::CollectiveMma< + DispatchPolicy, + TileShape_MNK, + ElementA, + cutlass::gemm::TagToStrideA_t, + ElementB, + cutlass::gemm::TagToStrideB_t, + TiledMma, + GmemTiledCopyA, + SmemLayoutAtomPairA, + CopyAtomPairA, + TransformA, + GmemTiledCopyB, + SmemLayoutAtomPairB, + CopyAtomPairB, + TransformB + >; +}; + +} // cutlass::gemm::collective diff --git a/include/cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl b/include/cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl new file mode 100644 index 00000000..dee762f1 --- /dev/null +++ b/include/cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl @@ -0,0 +1,182 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/gemm/collective/builders/sm100_common.inl" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +///////////////////////////////////////////////////////////////////////////////////////////////// +// Planar Complex f16/bf16 TCGEN05 kernels builder +template < + class ArchTag, + class ElementA, + class GmemLayoutATag, + class TransformA, + class ElementB, + class GmemLayoutBTag, + class TransformB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class BuilderScheduleTag +> +struct CollectiveBuilder< + ArchTag, + arch::OpClassTensorOp, + cute::tuple, + GmemLayoutATag, + 8, + cute::tuple, + GmemLayoutBTag, + 8, + ElementAccumulator, + TileShape_MNK, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + ClusterShape_MNK, // Static cluster shape or dynamic (int, int, _1) + StageCountType, + BuilderScheduleTag, + cute::enable_if_t< + (cute::is_same_v + ) && + // Element Types AB should be set as real type in Planar Complex f16/bf16 TCGEN05 kernels builder. + (cute::is_same_v || cute::is_same_v) && + (cute::is_same_v || cute::is_same_v) && + // Planar Complex f16/bf16 kernels don't support auto-scheduling for mainloop builder. + cute::is_base_of_v>> +{ + static_assert(cute::is_static_v, "TileShape has to be static"); + + static constexpr cute::UMMA::Major UmmaMajorA = cutlass::gemm::collective::detail::tag_to_umma_major_A(); + static constexpr cute::UMMA::Major UmmaMajorB = cutlass::gemm::collective::detail::tag_to_umma_major_B(); + + using TiledMma = decltype(detail::sm100_make_trivial_tiled_mma< + ElementA, ElementB, ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + UmmaMajorA, UmmaMajorB, BuilderScheduleTag, UMMA::ScaleIn::One>()); + using TiledMmaANeg = decltype(detail::sm100_make_trivial_tiled_mma< + ElementA, ElementB, ElementAccumulator, + TileShape_MNK, ClusterShape_MNK, + UmmaMajorA, UmmaMajorB, BuilderScheduleTag, UMMA::ScaleIn::Neg>()); + using TiledMmaPair = cutlass::gemm::collective::detail::Sm100CollectiveMmaPlanarComplexTiledMmaType; + + using AtomThrID = typename TiledMma::AtomThrID; + + // Define A and B block shapes for reduced size TMA_LOADs + // ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K) + using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(cute::size<0>(TileShape_MNK{}), + cute::size<2>(TileShape_MNK{})))); + // ((MMA_TILE_N,MMA_TILE_K), MMA_N, MMA_K) + using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(cute::size<1>(TileShape_MNK{}), + cute::size<2>(TileShape_MNK{})))); + + using GmemTiledCopyA = decltype(detail::sm100_cluster_shape_to_tma_atom_A(ClusterShape_MNK{}, AtomThrID{})); + + using BlockTileA_M = decltype(cute::size<0,0>(MmaShapeA_MK{}) * cute::size<1>(MmaShapeA_MK{})); + using BlockTileA_K = decltype(cute::size<0,1>(MmaShapeA_MK{}) * cute::size<2>(MmaShapeA_MK{})); + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + UmmaMajorA, ElementA, BlockTileA_M, BlockTileA_K>()); + + using GmemTiledCopyB = decltype(detail::sm100_cluster_shape_to_tma_atom_B(ClusterShape_MNK{}, AtomThrID{})); + + using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{})); + using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{})); + using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + UmmaMajorB, ElementB, BlockTileB_N, BlockTileB_K>()); + + // Calculate SMEM matrix A and B buffers' pipeline stages + static constexpr uint32_t AccumulatorPipelineStageCount = 2; + // Ptr-arry gemm requires extra TensorMap storage + static constexpr bool IsArrayOfPointersGemm = (cute::is_base_of_v); + // Calculate scheduler pipeline stages. Having one more stage than the accumulator allows more latency hiding. + static constexpr uint32_t SchedulerPipelineStageCount = IsArrayOfPointersGemm ? AccumulatorPipelineStageCount + 1: 1; + + static constexpr uint32_t KernelSmemCarveout = detail::Sm100DenseGemmTmaUmmaCarveout< + ClusterShape_MNK, + AccumulatorPipelineStageCount, + SchedulerPipelineStageCount, + detail::CLCResponseSize, + IsArrayOfPointersGemm, + 4 // 4 Tensor maps for A_{imag|real} and B_{imag|real} + >::KernelSmemCarveout; + + // Reduce SMEM capacity available for buffers considering barrier allocations. + + static constexpr int ReducedSmemCapacityBytes = + cutlass::gemm::collective::detail::sm100_smem_capacity_bytes - KernelSmemCarveout; + + using SmemTileShape = cute::Shape; + + // Use complex type to calculate SMEM stage count + using ComplexElementA = cutlass::complex; + using ComplexElementB = cutlass::complex; + + using MainloopPipelineStorage = typename cutlass::PipelineTmaUmmaAsync<1>::SharedStorage; + static constexpr int PipelineStages = detail::sm100_compute_stage_count_or_override< + ReducedSmemCapacityBytes, ComplexElementA, ComplexElementB, SmemTileShape, MainloopPipelineStorage>(StageCountType{}); + + using DispatchPolicy = cute::conditional_t, + cutlass::gemm::MainloopSm100TmaUmmaWarpSpecializedPlanarComplex< + PipelineStages, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + ClusterShape_MNK + > + >; + + using CollectiveOp = cutlass::gemm::collective::CollectiveMma< + DispatchPolicy, + TileShape_MNK, + ElementA, + cutlass::gemm::TagToStrideA_t, + ElementB, + cutlass::gemm::TagToStrideB_t, + TiledMmaPair, + GmemTiledCopyA, + SmemLayoutAtomA, + void, + TransformA, + GmemTiledCopyB, + SmemLayoutAtomB, + void, + TransformB + >; +}; + +} // cutlass::gemm::collective diff --git a/include/cutlass/gemm/collective/collective_builder.hpp b/include/cutlass/gemm/collective/collective_builder.hpp index b16f176a..e5f51df9 100644 --- a/include/cutlass/gemm/collective/collective_builder.hpp +++ b/include/cutlass/gemm/collective/collective_builder.hpp @@ -39,15 +39,15 @@ #include "cutlass/gemm/collective/collective_builder_decl.hpp" #include "cutlass/gemm/collective/builders/sm90_gmma_builder.inl" #include "cutlass/gemm/collective/builders/sm90_sparse_gmma_builder.inl" -#if !defined(__CUDACC_RTC__) -#include "cutlass/gemm/collective/builders/sm100_umma_builder.inl" -#include "cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl" +#if !defined(__CUDACC_RTC__) +#include "cutlass/gemm/collective/builders/sm100_umma_builder.inl" +#include "cutlass/gemm/collective/builders/sm100_9xBF16_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_sparse_umma_builder.inl" -#include "cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl" +#include "cutlass/gemm/collective/builders/sm100_blockscaled_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_blockwise_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_blockscaled_sparse_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_simt_builder.inl" -#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl" +#include "cutlass/gemm/collective/builders/sm100_mixed_input_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_cpasync_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_mixed_tma_cpasync_umma_builder.inl" #include "cutlass/gemm/collective/builders/sm100_blockscaled_mixed_tma_cpasync_umma_builder.inl" @@ -57,6 +57,9 @@ #include "cutlass/gemm/collective/builders/sm120_sparse_mma_builder.inl" #include "cutlass/gemm/collective/builders/sm120_blockscaled_sparse_mma_builder.inl" #include "cutlass/gemm/collective/builders/sm120_blockwise_mma_builder.inl" +#include "cutlass/gemm/collective/builders/sm100_interleaved_complex_umma_builder.inl" +#include "cutlass/gemm/collective/builders/sm100_9xBF16_interleaved_complex_umma_builder.inl" +#include "cutlass/gemm/collective/builders/sm100_planar_complex_umma_builder.inl" #endif ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/collective_mma.hpp b/include/cutlass/gemm/collective/collective_mma.hpp index 8353c50e..2cc5697a 100644 --- a/include/cutlass/gemm/collective/collective_mma.hpp +++ b/include/cutlass/gemm/collective/collective_mma.hpp @@ -42,7 +42,7 @@ #include "cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp" #include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp" #include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp" -#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp" +#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp" #include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp" #include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp" #include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized_fp8.hpp" @@ -61,8 +61,8 @@ #include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_emulated.hpp" #include "cutlass/gemm/collective/sm100_sparse_mma_warpspecialized.hpp" #include "cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp" -#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" -#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp" +#include "cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" +#include "cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized.hpp" #include "cutlass/gemm/collective/sm100_mma_warpspecialized_blockwise_scaling.hpp" #include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_blockwise_scaling.hpp" #include "cutlass/gemm/collective/sm100_mma_warpspecialized_mixed_input.hpp" @@ -74,11 +74,17 @@ #include "cutlass/gemm/collective/sm120_mma_tma.hpp" #include "cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp" #include "cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp" -#include "cutlass/gemm/collective/sm120_sparse_mma_tma.hpp" -#include "cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp" #include "cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp" #include "cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp" -#endif // !defined(__CUDACC_RTC__) +#include "cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp" +#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp" +#include "cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp" +#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp" +#include "cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp" +#include "cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp" +#include "cutlass/gemm/collective/sm120_sparse_mma_tma.hpp" +#include "cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp" +#endif // !defined(__CUDACC_RTC__) diff --git a/include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp b/include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp index 082a13fc..e949f31d 100644 --- a/include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp @@ -604,14 +604,15 @@ struct CollectiveMma< implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); // Check for SFA SFB layout requirement - const auto layout_sfa_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL); - const auto layout_sfb_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL); - implementable = implementable && (layout_sfa_ref == args.layout_SFA); + const auto layout_sfa_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL)); + const auto layout_sfb_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL)); + + implementable = implementable && (layout_sfa_ref == take<0,2>(args.layout_SFA)); if (!implementable) { CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFA mismatch, layout_SFA needs to be K-major\n"); } - implementable = implementable && (layout_sfb_ref == args.layout_SFB); + implementable = implementable && (layout_sfb_ref == take<0,2>(args.layout_SFB)); if (!implementable) { CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFB mismatch, layout_SFB needs to be K-major\n"); } diff --git a/include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp b/include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp index b962296b..e20371d3 100644 --- a/include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp +++ b/include/cutlass/gemm/collective/sm100_blockscaled_sparse_mma_warpspecialized.hpp @@ -772,14 +772,14 @@ struct CollectiveMma< } // Check for SFA SFB layout requirement - const auto layout_sfa_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL); - const auto layout_sfb_ref = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL); - implementable = implementable && (layout_sfa_ref == args.layout_SFA); + const auto layout_sfa_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL)); + const auto layout_sfb_ref = take<0,2>(Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL)); + implementable = implementable && (layout_sfa_ref == take<0,2>(args.layout_SFA)); if (!implementable) { CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFA mismatch, layout_SFA needs to be K-major\n"); } - implementable = implementable && (layout_sfb_ref == args.layout_SFB); + implementable = implementable && (layout_sfb_ref == take<0,2>(args.layout_SFB)); if (!implementable) { CUTLASS_TRACE_HOST(" CAN IMPLEMENT: layout_SFB mismatch, layout_SFB needs to be K-major\n"); } diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp new file mode 100644 index 00000000..d9275cf0 --- /dev/null +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_emulated.hpp @@ -0,0 +1,1197 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/numeric_conversion.h" +#include "cutlass/detail/sm100_tmem_helper.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/arch/mma_sm100.hpp" +#include "cutlass/trace.h" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop for FastF32 Kernels +template < + int Load2TransformPipelineStageCount_, + int Transform2MmaPipelineStageCount_, + int SchedulerPipelineStageCount_, + int AccumulatorPipelineStageCount_, + int NumBandsToCompute_, + int ScalingFactor_, + int AccPromotionInterval_, + class AccumulatorCopyAtom_, + class ClusterShape, + class TileShape_, + class StrideA_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomsA_, + class CopyAtomsA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomsB_, + class CopyAtomsB_, + class TransformB_> +struct CollectiveMma< + MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32< + Load2TransformPipelineStageCount_, + Transform2MmaPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + NumBandsToCompute_, + ScalingFactor_, + AccPromotionInterval_, + ClusterShape, + AccumulatorCopyAtom_>, + TileShape_, + complex, + StrideA_, + complex, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomsA_, + CopyAtomsA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomsB_, + CopyAtomsB_, + TransformB_> +{ + // + // Type Aliases + // + using TileShape = TileShape_; + using TiledMma = TiledMma_; + + // ElementA and ElementB are cutlass::complex, which are used as GMEM input and output data type. + using ElementA = complex; + using StrideA = StrideA_; + using InternalStrideA = cute::remove_pointer_t; + using ElementB = complex; + using StrideB = StrideB_; + using InternalStrideB = cute::remove_pointer_t; + + // For a complex kernel, the MMA output type is real valued, but ElementAccumulator is a complex type for the GETT reference kernel + using ElementAccumulator = complex; + using ElementAccumulatorRaw = typename TiledMma::ValTypeC; + +private: + // ElementAMma and ElementBMma are cutlass::complex, which are used as SMEM and RF data type. + // ElementAMmaRaw and ElementBMmaRaw are cutlass::bfloat16_t, which is the real internal data type set in TMA descriptor and used in TCGEN05 calculation. + using ElementAMma = typename TiledMma::ValTypeA; // complex + using ElementAMmaRaw = typename ElementAMma::value_type; // bfloat16_t + using ElementBMma = typename TiledMma::ValTypeB; // complex + using ElementBMmaRaw = typename ElementBMma::value_type; // bfloat16_t + +public: + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomsA = SmemLayoutAtomsA_; + using SmemLayoutAtomsB = SmemLayoutAtomsB_; + using CopyAtomsA = CopyAtomsA_; + using CopyAtomsB = CopyAtomsB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + + static_assert(cute::is_same_v, "Underlying input type for A should be float"); + static_assert(cute::is_same_v, "Underlying input type for B should be float"); + static_assert(cute::is_same_v, "Underlying compute type for A should be bfloat16_t"); + static_assert(cute::is_same_v, "Underlying compute type for A should be bfloat16_t"); + + // Determine MMA type: MMA_1SM vs MMA_2SM + using AtomThrShapeMNK = Shape(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>; + using DispatchPolicy = MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32< + Load2TransformPipelineStageCount_, + Transform2MmaPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + NumBandsToCompute_, + ScalingFactor_, + AccPromotionInterval_, + ClusterShape, + AccumulatorCopyAtom_>; + static constexpr bool IsDynamicCluster = not cute::is_static_v; + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + + // Define A and B block shapes for reduced size TMA_LOADs + using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + using ArchTag = typename DispatchPolicy::ArchTag; + + using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync< + DispatchPolicy::Load2TransformPipelineStageCount, + AtomThrShapeMNK>; + using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState; + + using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync< + DispatchPolicy::Transform2MmaPipelineStageCount, + AtomThrShapeMNK>; + using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState; + + using Mma2AccumPipeline = cutlass::PipelineUmmaAsync< + DispatchPolicy::Schedule::AccumulatorPipelineStageCount, + AtomThrShapeMNK>; + using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState; + + // Thread Counts + static constexpr uint32_t NumTransformationThreads = 128; + static constexpr uint32_t NumAccumThreads = 128; + + // Get the Algorithm parameters + constexpr static int NumComputeMtxs = 3; + constexpr static int ConjSwapMode = 2; + constexpr static int NumBandsToCompute = DispatchPolicy::NumBandsToCompute; + constexpr static int ScalingFactor = DispatchPolicy::ScalingFactor; + constexpr static int AccPromotionInterval = DispatchPolicy::AccPromotionInterval; + constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; + constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{}) / DispatchPolicy::AccPromotionInterval; + constexpr static int NumBandsMax = 5; + static_assert(NumBandsToCompute <= NumBandsMax && NumBandsToCompute >= 3, "NumBandsToCompute should be less than maximum number of bands"); + static_assert(StagesPerTile * AccPromotionInterval == size<2>(CtaShapeA_MK{}), "PromotionInterval*InstructionK doesn't evenly divide CTA shape"); + + // Copy atom for Accumulator + using AccumulatorCopyAtom = typename DispatchPolicy::AccumulatorCopyAtom; + + static_assert((NumBandsToCompute == 5 || NumBandsToCompute == 4 || NumBandsToCompute == 3), + "9xBF16 with 5/4/3 Bands are supported"); + + using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom; + using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom; + using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom; + using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom; + + using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom; + using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom; + using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom; + using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(CtaShapeA_MK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomACompute{}, + tuple_cat(CtaShapeA_MK{}, tuple, Int>{}))); + + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(CtaShapeB_NK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutBCompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomBCompute{}, + tuple_cat(CtaShapeB_NK{}, tuple, Int, Int>{}))); + + static_assert(DispatchPolicy::Load2TransformPipelineStageCount >= 2 && DispatchPolicy::Load2TransformPipelineStageCount >= 2, + "Specialization requires Stages set to value 2 or more."); + static_assert((cute::is_base_of::value || + cute::is_base_of::value ) && + cute::is_base_of::value, + "MMA atom must A operand from SMEM or TMEM and B operand from SMEM for this mainloop."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyA - invalid TMA copy atom specified."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyB - invalid TMA copy atom specified."); + + struct PipelineStorage { + using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage; + alignas(16) Load2TransformPipelineStorage load2transform_pipeline; + using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage; + alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline; + using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage; + alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline; + }; + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + struct TensorStorageUntransformed { + cute::ArrayEngine> smem_A; + cute::ArrayEngine> smem_B; + }; + + struct TensorStorageTransformedAinSmem { + alignas(1024) cute::ArrayEngine> smem_ACompute; + alignas(1024) cute::ArrayEngine> smem_BCompute; + }; + + union TensorStorageTransformedAinTmem { + alignas(1024) cute::ArrayEngine smem_ACompute; // No smem_ACompute + alignas(1024) cute::ArrayEngine> smem_BCompute; + }; + + using TensorStorageTransformed = cute::conditional_t< + cute::is_base_of::value, + TensorStorageTransformedAinSmem, + TensorStorageTransformedAinTmem>; + + TensorStorageUntransformed input; + TensorStorageTransformed compute; + } tensors; + + struct TensorMapStorage : cute::aligned_struct<128, _0> { + cute::TmaDescriptor smem_tensormap_A; + cute::TmaDescriptor smem_tensormap_B; + } tensormaps; + + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using TensorMapStorage = typename SharedStorage::TensorMapStorage; + + // Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on + // loaded input A and B matrices to convert the data type + static constexpr uint32_t TmaTransactionBytes = + cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast(sizeof_bits::value))+ + cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast(sizeof_bits::value)); + + // Host side kernel arguments + struct Arguments { + ElementA const** ptr_A{nullptr}; + StrideA dA{}; + ElementB const** ptr_B{nullptr}; + StrideB dB{}; + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), + make_tile(typename TiledMma::AtomThrID{}))); + + using TMA_A = decltype(make_tma_atom_A_sm100( + GmemTiledCopyA{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + TMA_A tma_load_a; + TMA_B tma_load_b; + TMA_A tma_load_a_fallback; + TMA_B tma_load_b_fallback; + dim3 cluster_shape_fallback; + cute::TmaDescriptor* tensormaps; + ElementA const** ptr_A; + ElementB const** ptr_B; + }; + + CUTLASS_DEVICE + CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster) + : cluster_shape_(cluster_shape) + , block_rank_in_cluster_(block_rank_in_cluster) { + if constexpr (IsDynamicCluster) { + const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x && + cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y); + observed_tma_load_a_ = is_fallback_cluster ? ¶ms.tma_load_a_fallback : ¶ms.tma_load_a; + observed_tma_load_b_ = is_fallback_cluster ? ¶ms.tma_load_b_fallback : ¶ms.tma_load_b; + } + else { + observed_tma_load_a_ = ¶ms.tma_load_a; + observed_tma_load_b_ = ¶ms.tma_load_b; + } + } + + template + static constexpr Params + to_underlying_arguments(ProblemShape problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) { + (void) workspace; + + // Tensor shapes for Ptr-Array are initialized correctly here. + auto [M,N,K,mock_L] = problem_shape.get_host_problem_shape(0); + // Batches/Groups are managed by using appropriate pointers to input matrices + mock_L = 1; + + // Tensor pointers will be fixed before the first access + ElementA const* ptr_A_first_batch = nullptr; + ElementB const* ptr_B_first_batch = nullptr; + + Tensor tensor_a = make_tensor(ptr_A_first_batch, make_layout(make_shape(M,K,mock_L), args.dA)); + Tensor tensor_b = make_tensor(ptr_B_first_batch, make_layout(make_shape(N,K,mock_L), args.dB)); + + auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape); + // Cluster layout for TMA construction + auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{})); + + auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback); + // Cluster layout for TMA construction + auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + return { + tma_load_a, + tma_load_b, + tma_load_a_fallback, + tma_load_b_fallback, + hw_info.cluster_shape_fallback, + reinterpret_cast(workspace), + reinterpret_cast(args.ptr_A), + reinterpret_cast(args.ptr_B) + }; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) { + constexpr uint32_t NumInputTensors = 2; + constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor); + // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies + return (NumInputTensors * SizeOfCuTensorMap * sm_count); + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) { + return cutlass::Status::kSuccess; + } + + template + static bool + can_implement( + ProblemShape problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits = 128; + auto [M,N,K,L] = problem_shape.get_host_problem_shape(0); + + bool implementable = true; + constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,K,L), StrideA{}); + constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + return implementable; + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE auto + partition_accumulator_shape() { + auto acc_shape = partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N) + + return acc_shape; + } + + /// Produce the inputs to the transform threads by loading inputs from gmem -> smem + template < + class GTensorA, class GTensorB, + class GTensorPartitionedA, class GTensorPartitionedB, + class STensorA, class STensorB, + class TensorMapA, class TensorMapB, + class TileCoordMNKL, + class KTileIterator + > + CUTLASS_DEVICE auto + load( + Params const& params, + Load2TransformPipeline pipeline, + Load2TransformPipelineState load2xform_pipeline_state, + cute::tuple> const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, + KTileIterator k_tile_iter, int k_tile_count) { + + auto [unused_gA, unused_gB, + tAgA_mkl, tBgB_nkl, tAsA, tBsB, + mcast_mask_a, mcast_mask_b, + input_tensormaps] = load_inputs; + + // slice out the work coord from tiled tensors + Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + uint32_t skip_wait = (k_tile_count <= 0); + auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + // Issue the Mainloop loads + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // LOCK mainloop_load2xform_pipeline_state for _writing_ + pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag); + int write_stage = load2xform_pipeline_state.index(); + + using BarrierType = typename Load2TransformPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state); + + // Advance mainloop_pipe + ++load2xform_pipeline_state; + skip_wait = (k_tile_count <= 1); + pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + copy(observed_tma_load_a_->with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage)); + copy(observed_tma_load_b_->with(get<1>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage)); + ++k_tile_iter; + } + return cute::make_tuple(load2xform_pipeline_state, k_tile_iter); + } + + /// Set up the data needed by this collective for load. + /// Returned tuple must contain at least two elements, with the first two elements being: + /// gA_mkl - The tiled tensor for input A + /// gB_nkl - The tiled tensor for input B + // Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks + template + CUTLASS_DEVICE auto + load_init( + ProblemShape_MNKL const& problem_shape_MNKL, + Params const& params, + TensorStorage& shared_storage, + int32_t const sm_count, int32_t const sm_idx) const { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + + Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + + // Define the CTA-in-cluster Layout and Coord + Layout cta_layout_mnk = make_layout(cluster_shape_); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + // Project the cta_layout for tma_a along the n-modes + auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl)); + + // Project the cta_layout for tma_b along the m-modes + auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl)); + + // TMA Multicast Masks + uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + // Fetch a copy of tensormaps for the CTA from Params + auto input_tensormaps = tensormaps_init(params, sm_count, sm_idx); + + return cute::make_tuple( + gA_mkl, gB_nkl, // for scheduler + tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values + mcast_mask_a, mcast_mask_b, // multicast masks + input_tensormaps); // for tma descriptor modification (per-CTA tensormap copy) + } + + template< + class KTileIterator, class Accumulator, + class GTensorA, class DstCopyA, class SrcTensorA, class DstTensorA, + class GTensorB, class SrcTensorB, class DstTensorB + > + CUTLASS_DEVICE auto + transform( + Load2TransformPipeline load2transform_pipeline, + Load2TransformPipelineState load2transform_pipeline_consumer_state, + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_producer_state, + Accumulator accumulators, + cute::tuple input_operands, + KTileIterator k_tile_iter, int k_tile_count) { + + static_assert(cute::is_same_v, "ElementA and ElementB types should be the same."); + static_assert(cute::is_same_v, "ElementAMma and ElementBMma types should be the same."); + + cutlass::arch::NamedBarrier transform_bar(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier); + + // tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tAdA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs(Emulation), SmemStages (In SMEM or TMEM) + // tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, NumComputeMtxs(Complex,Emulation), SmemStages (In SMEM) + auto [unused_tAgA, dst_copy_A, tAsA, tAdACompute, + unused_tBgB, tBsB, tBsBCompute] = input_operands; + + // Create the tensors in registers + auto tArA = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArA_temp = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArACompute = make_tensor(tAsA(_,_,_,_,0).shape()); + + auto tBrB = make_tensor(tBsB(_,_,_,_,0).shape()); + auto tBrB_temp = make_tensor(tBsB(_,_,_,_,0).shape()); + auto tBrBCompute = make_tensor(tBsB(_,_,_,_,0).shape()); + + // For compute, cast to 4 raw elements instead of 2 complex elements. + auto tArA_x4 = recast>(tArA); + auto tArA_temp_x4 = recast>(tArA_temp); + auto tArACompute_x4 = recast>(tArACompute); + + auto tBrB_x4 = recast>(tBrB); + auto tBrB_temp_x4 = recast>(tBrB_temp); + auto tBrBCompute_x4 = recast>(tBrBCompute); + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag); + transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag); + + int load2transform_consumer_index = load2transform_pipeline_consumer_state.index(); + int transform2mma_producer_index = transform2mma_pipeline_producer_state.index(); + + auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state; + auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state; + + // Copy the input B matrix from SMEM + copy(AutoVectorizingCopy{}, tBsB(_,_,_,_,load2transform_consumer_index), tBrB); + // Copy the input A matrix from SMEM + copy(AutoVectorizingCopy{}, tAsA(_,_,_,_,load2transform_consumer_index), tArA); + + /// NOTE: sm100_mma_warpspecialized_interleaved_complex_tf32.hpp introduced about expanding: + /// re(a_complex * b_complex) -> (a_re, a_im) . (b_re,-b_im) = a . b_conj + /// im(a_complex * b_complex) -> (a_re, a_im) . (b_im, b_re) = a . b_swap + /// However, 16b types need to be packed for swapping and negation. + /// Hence, (re | im | re | im) is reordered into (re_x2 | im_x2). + cute::transform(tBrB_x4, tBrB_x4, [&] (auto& f4) -> Array {return {f4[0], f4[2], f4[1], f4[3]};}); + // Conversion b -> b_conj goes first, hence TransformB has a not preceding it. + if constexpr (not cute::is_same_v) { + cute::transform(tBrB_x4, tBrB_x4, [&] (auto& f4) { + auto f2_x2 = *reinterpret_cast,2>*>(&f4); + f2_x2[1] = cutlass::negate>{}(f2_x2[1]); + return *reinterpret_cast*>(&f2_x2); + }); + } + CUTE_UNROLL + for (int comp_mtx_index = 0; comp_mtx_index < NumComputeMtxs; ++comp_mtx_index) { + // Convert from fp32 -> bf16 + cute::transform(tBrB_x4, tBrBCompute_x4, + cutlass::NumericArrayConverter::convert); + // Store as B_conj (for producing C_re) + copy(AutoVectorizingCopy{}, tBrBCompute, tBsBCompute(_,_,_,_,0,comp_mtx_index,transform2mma_producer_index)); + // if it is not the last compute matrix, scale and substract + if (comp_mtx_index < NumComputeMtxs - 1) { + // Convert from bf16 -> fp32 to substract + cute::transform(tBrBCompute_x4, tBrB_temp_x4, + cutlass::NumericArrayConverter::convert); + cute::transform(tBrB_x4, tBrB_temp_x4, tBrB_x4, cutlass::minus>{}); + if constexpr (DispatchPolicy::ScalingFactor != 0) { + cute::transform(tBrB_x4, tBrB_x4, cutlass::scale>{(1 << DispatchPolicy::ScalingFactor)}); + } + } + + // Convert B_conj to B_swap + cute::transform(tBrBCompute_x4, tBrBCompute_x4, [&] (auto& h4) { + // Reinterpret as packed types + auto h2_x2_conj = *reinterpret_cast,2>*>(&h4); + cutlass::negate> neg; + Array,2> h2_x2_swap{ neg(h2_x2_conj[1]), h2_x2_conj[0] }; + return *reinterpret_cast*>(&h2_x2_swap); + }); + // Store as B_swap (for producing C_im) + copy(AutoVectorizingCopy{}, tBrBCompute, tBsBCompute(_,_,_,_,1,comp_mtx_index,transform2mma_producer_index)); + } + + // Loads from SMEM are done. Signal the mainloop load as early as possible + transform_bar.sync(); + load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state); + + // ( re | im | re | im ) -> ( re_x2 | im_x2 ) + cute::transform(tArA_x4, tArA_x4, [&] (auto& f4) -> Array{return {f4[0], f4[2], f4[1], f4[3]};}); + if constexpr (cute::is_same_v) { + cute::transform(tArA_x4, tArA_x4, [&] (auto& f4) { + auto f2_x2 = *reinterpret_cast,2>*>(&f4); + f2_x2[1] = cutlass::negate>{}(f2_x2[1]); + return *reinterpret_cast*>(&f2_x2); + }); + } + CUTE_UNROLL + for (int comp_mtx_index = 0; comp_mtx_index < NumComputeMtxs; ++comp_mtx_index) { + // Convert from fp32 -> bf16 + cute::transform(tArA_x4, tArACompute_x4, + cutlass::NumericArrayConverter::convert); + copy(dst_copy_A, tArACompute, tAdACompute(_,_,_,_,comp_mtx_index,transform2mma_producer_index)); + // if it is not the last compute matrix, scale and substract + if (comp_mtx_index < NumComputeMtxs - 1) { + // Convert from bf16 -> fp32 to substract + cute::transform(tArACompute_x4, tArA_temp_x4, + cutlass::NumericArrayConverter::convert); + cute::transform(tArA_x4, tArA_temp_x4, tArA_x4, cutlass::minus>{}); + if constexpr (DispatchPolicy::ScalingFactor != 0) { + cute::transform(tArA_x4, tArA_x4, cutlass::scale>{(1 << DispatchPolicy::ScalingFactor)}); + } + } + } + + // fence for SMEM writes + cutlass::arch::fence_view_async_shared(); + if constexpr (is_tmem::value) { + // fence for TMEM writes if A operand is coming from TMEM + cutlass::arch::fence_view_async_tmem_store(); + } + + // Let the MMA know we are done transforming + transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state); + // Next pipeline stage + ++load2transform_pipeline_consumer_state; + ++transform2mma_pipeline_producer_state; + + skip_wait = (k_tile_count <= 1); + // Peek the next pipeline stage's barriers + load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + } + return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + transform_init( + Params const& params, + ProblemShape_MNKL const& problem_shape_MNKL, + Accumulator accumulators, + TensorStorage& shared_storage) { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); + Tensor sA = as_position_independent_swizzle_tensor(sA_orig); + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + + Tensor sB_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); + Tensor sB = as_position_independent_swizzle_tensor(sB_orig); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + + // Map input, compute, and fragment tensors to + // Copy strategies and partitioned tensors. These will become the input + // operands of the transform function. Depending on MMA atom type, the + // operands can reside in SMEM or TMEM + auto setup_copy_ops = [&] ( + auto tensor_input, + auto input_copy_atom, + auto tensor_compute, + auto make_fragment, + auto compute_copy_atom) constexpr { + auto fragment_compute = make_fragment(tensor_compute); + if constexpr (cute::is_tmem>::value) { + // For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation. + // Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation + // See: TmemAllocMode::Duplicated. + Tensor tensor_input2x = [&] () constexpr { + if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) { + return make_tensor(tensor_input.data(), + logical_product(tensor_input.layout(), + make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE) + } + else { + return tensor_input; + } + }(); + + fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + // If operand comes from TMEM, create the TMEM_STORE based copy + auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,_,0,0)); + auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x); + auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute); + return cute::make_tuple(reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute); + } + else { + // If the operand comes from SMEM, create SMEM copy. + auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute); + auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout, Stride< _8,_1>>{}, + take<0,3>(tensor_compute.layout())); + + // Source copy is based on the source operand of copy. + auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input); + auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw); + + return cute::make_tuple(AutoVectorizingCopy{}, partitioned_tensor_input, partitioned_tensor_compute); + } + }; + + auto [dst_copy_A, tAsA, tAsACompute] = + setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{}); + + auto [dst_copy_B, tBsB, tBsBCompute] = + setup_copy_ops(sB, InputCopyAtomB{}, sBCompute, [&](auto &arg) {return TiledMma::make_fragment_B(arg);}, ComputeCopyAtomB{}); + + return cute::make_tuple(gA_mkl, dst_copy_A, tAsA, tAsACompute, + gB_nkl, tBsB, tBsBCompute); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class FrgEngine, class FrgLayout, + class TensorA, class TensorB + > + CUTLASS_DEVICE auto + mma( + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_consumer_state, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_producer_state, + cute::Tensor const& accumulators, + cute::tuple const& input_operands, + int k_tile_count + ) { + TiledMma tiled_mma; + + auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + uint32_t skip_wait = (k_tile_count <= 0); + auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + ++next_transform2mma_pipeline_consumer_state; + + // tCrA : (MMA), MMA_M, MMA_K, NumComputeMtxs, SmemStage (In SMEM or TMEM) + // We use SMEM stages to match #buffers in Load <-> Convert + // tCrB : (MMA), MMA_N, MMA_K, NumComputeMtxs, SmemStages (In SMEM) + auto const [tCrA, tCrB] = input_operands; + + using ZeroScaler = cute::integral_constant; + using Scaler = cute::integral_constant; + + int remaining_accum_promotions = k_tile_count * StagesPerTile; + uint32_t mma2accum_skip_wait = (remaining_accum_promotions <= 0); + auto mma2accum_flag = mma2accum_pipeline.producer_try_acquire(mma2accum_pipeline_producer_state, mma2accum_skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag); + + int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index(); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); k_block += DispatchPolicy::AccPromotionInterval, --remaining_accum_promotions) { + // Accum stages are organized as (C_real | C_imag | C_real | C_imag | ...) + CUTLASS_PRAGMA_UNROLL + for (int re_im = 0; re_im < 2; ++re_im) { + mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state, mma2accum_flag); + + int mma2accum_pipeline_producer_state_index = mma2accum_pipeline_producer_state.index(); + auto tCtC = accumulators(_,_,_,mma2accum_pipeline_producer_state_index); + auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state; + + ++mma2accum_pipeline_producer_state; + mma2accum_skip_wait = (remaining_accum_promotions <= 1) && (re_im == 1); + mma2accum_flag = mma2accum_pipeline.producer_try_acquire(mma2accum_pipeline_producer_state, mma2accum_skip_wait); + + auto tCrA0 = tCrA(_,_,_,0,transform2mma_pipeline_consumer_state_index); + auto tCrA1 = tCrA(_,_,_,1,transform2mma_pipeline_consumer_state_index); + auto tCrA2 = tCrA(_,_,_,2,transform2mma_pipeline_consumer_state_index); + + auto tCrB0 = tCrB(_,_,_,re_im,0,transform2mma_pipeline_consumer_state_index); + auto tCrB1 = tCrB(_,_,_,re_im,1,transform2mma_pipeline_consumer_state_index); + auto tCrB2 = tCrB(_,_,_,re_im,2,transform2mma_pipeline_consumer_state_index); + + // MMA instructions Emulation + auto accumulate = UMMA::ScaleOut::Zero; + + // First set of GEMMs that we need to perform for each band are unrolled to set compile-time constant + // scaling parameter. Scaled GEMM operations are only needed for the first MMA operation of each band. + + // Band 5 + if constexpr (NumBandsToCompute == 5) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block), tCrB2(_,_,k_block), tCtC); // A[2]*B[2] + accumulate = UMMA::ScaleOut::One; + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block+s), tCrB2(_,_,k_block+s), tCtC); // A[2]*B[2] + } + } + // Band 4 + if constexpr (NumBandsToCompute >= 4) { + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA1(_,_,k_block), tCrB2(_,_,k_block), tCtC); // A[1]*B[2] + accumulate = UMMA::ScaleOut::One; + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block), tCrB1(_,_,k_block), tCtC); // A[2]*B[1] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block+s), tCrB2(_,_,k_block+s), tCtC); // A[1]*B[2] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block+s), tCrB1(_,_,k_block+s), tCtC); // A[2]*B[1] + } + } + // Band 3 + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA0(_,_,k_block), tCrB2(_,_,k_block), tCtC); // A[2]*B[0] + accumulate = UMMA::ScaleOut::One; + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block), tCrB1(_,_,k_block), tCtC); // A[1]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[0]*B[2] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA0(_,_,k_block+s), tCrB2(_,_,k_block+s), tCtC); // A[2]*B[0] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block+s), tCrB1(_,_,k_block+s), tCtC); // A[1]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block+s), tCrB0(_,_,k_block+s), tCtC); // A[0]*B[2] + } + // Band 2 + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA0(_,_,k_block), tCrB1(_,_,k_block), tCtC); // A[0]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[1]*B[0] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA0(_,_,k_block+s), tCrB1(_,_,k_block+s), tCtC); // A[0]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block+s), tCrB0(_,_,k_block+s), tCtC); // A[1]*B[0] + } + // Band 1 + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA0(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[0]*B[0] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA0(_,_,k_block+s), tCrB0(_,_,k_block+s), tCtC); // A[0]*B[0] + } + mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state); + } + } + + transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state); + + skip_wait = (k_tile_count <= 1); + transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + + curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state; + ++next_transform2mma_pipeline_consumer_state; + } + return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + mma_init(cute::Tensor const& accumulators, TensorStorage& shared_storage) const { + TiledMma tiled_mma; + + Tensor tCrA = [&] () constexpr { + if constexpr (cute::is_base_of::value) { + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + return tiled_mma.make_fragment_A(sACompute); + } + else { + auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{})); + tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + return tCrA; + } + } (); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + Tensor tCrB = tiled_mma.make_fragment_B(sBCompute); + return cute::make_tuple(tCrA, tCrB); + } + + template + CUTLASS_DEVICE auto + accum_init(cute::Tensor const& accumulators, TmemCopyAtom tmem_cp_atom, EpilogueTile epilogue_tile) { + // Obtain a single accumulator + Tensor tAcc = tensor<0>(accumulators(_,_,_,_0{})); + // Apply epilogue subtiling + Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + // Create the TMEM copy for single EpilogueTile. + // Note that EpilogueTile = CtaTile for NoSmem epilogue + auto tiled_t2r = make_tmem_copy(tmem_cp_atom, tAcc_epi(_,_,_0{},_0{})); + auto thread_t2r = tiled_t2r.get_slice(threadIdx.x % size(tiled_t2r)); + Tensor tTR_gC = thread_t2r.partition_D(tAcc_epi); + Tensor tTR_rAcc = make_tensor(shape(tTR_gC)); // (T2R,T2R_M,T2R_N) + Tensor tTR_rGlobAcc = make_tensor(shape(tTR_gC)); // (T2R,T2R_M,T2R_N) + + // Apply epilogue subtiling to bulk accumulator + // We need to tile the whole bulk_tmem allocation with EpilogueTile. + // The accumulation should be aware of the AccumulatorPipelineStages + Tensor tBulkAcc_epi = flat_divide(accumulators(make_coord(_,_),_0{},_0{},_), EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N,PIPE) + Tensor tTR_tBulkAcc = thread_t2r.partition_S(tBulkAcc_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N,PIPE) + return cute::make_tuple(tiled_t2r, thread_t2r, tTR_tBulkAcc, tTR_rAcc, tTR_rGlobAcc); + } + + template + CUTLASS_DEVICE auto + accum(cute::tuple accum_inputs, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_consumer_state, + int k_tile_count) { + auto [tiled_t2r, thread_t2r, tTR_tBulkAcc, + tTR_rAcc, tTR_rGlobAcc] = accum_inputs; + + Tensor tTR_rAcc_float2 = recast>(tTR_rAcc); // (T2R/2,T2R_M,T2R_N) + Tensor tTR_rGlobAcc_float2_x2 = recast,2>>(tTR_rGlobAcc);// (T2R/2,T2R_M,T2R_N) + + // Clear the global accumulator + CUTE_UNROLL + for (int i = 0; i 0; --k_tile_count) { + // The stage is limited to a CTA tile + CUTLASS_PRAGMA_NO_UNROLL + for (int k_block = 0; k_block cute::remove_cvref_t {return {cutlass::plus>{}(f2_x2[0], r2), f2_x2[1]};}); + } + else { + cute::transform(tTR_rGlobAcc_float2_x2, tTR_rAcc_float2, tTR_rGlobAcc_float2_x2, + [&] (auto& f2_x2, auto& i2) -> cute::remove_cvref_t {return {f2_x2[0], cutlass::plus>{}(f2_x2[1], i2)};}); + } + + cutlass::arch::fence_view_async_tmem_load(); // Need a fence bw TMEM_LOAD and arrive + mma2accum_pipeline.consumer_release(mma2accum_pipeline_consumer_state); + + ++mma2accum_pipeline_consumer_state; + skip_wait = ((k_tile_count <= 1) && (k_block >= (StagesPerTile-1))) && (re_im == 1); + mma2accum_flag = mma2accum_pipeline.consumer_try_wait(mma2accum_pipeline_consumer_state, skip_wait); + } + } + } + + // Interleave back (real_x2 | imag_x2) to (real | imag | real | imag) + cute::transform(tTR_rGlobAcc_float2_x2, tTR_rGlobAcc_float2_x2, [&] (auto& f2_x2) -> cute::remove_cvref_t { + Array c0{f2_x2[0][0], f2_x2[1][0]}; + Array c1{f2_x2[0][1], f2_x2[1][1]}; + return {c0, c1}; + }); + + return cute::make_tuple(mma2accum_pipeline_consumer_state, tTR_rGlobAcc); + } + + // + // Methods to perform different parts of TMA/Tensormap modifications + // + + CUTLASS_DEVICE auto + tensormaps_init(Params const& mainloop_params, int32_t const sm_count, int32_t const sm_idx) const { + cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps; + + cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx]; + cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx + sm_count]; + + if (cute::elect_one_sync()) { + // Bringing tensormaps from params to gmem for modification later + Tensor pA_tensormap = make_tensor(observed_tma_load_a_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor gA_tensormap = make_tensor(tma_desc_a, Int<1>{}, Int<1>{}); + Tensor pB_tensormap = make_tensor(observed_tma_load_b_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor gB_tensormap = make_tensor(tma_desc_b, Int<1>{}, Int<1>{}); + + copy(recast(pA_tensormap), recast(gA_tensormap)); + copy(recast(pB_tensormap), recast(gB_tensormap)); + } + + return cute::make_tuple(tma_desc_a, tma_desc_b); + } + + // Bringing tensormaps to smem (to be done by single thread) + template + CUTLASS_DEVICE + void + tensormaps_fetch_to_smem( + TensorMapStorage& shared_tensormap, + cute::tuple const& input_tensormaps) const { + Tensor gA_tensormap = make_tensor(make_gmem_ptr(get<0>(input_tensormaps)), Int<1>{}, Int<1>{}); + Tensor sA_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_A), Int<1>{}, Int<1>{}); + Tensor gB_tensormap = make_tensor(make_gmem_ptr(get<1>(input_tensormaps)), Int<1>{}, Int<1>{}); + Tensor sB_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_B), Int<1>{}, Int<1>{}); + + copy(recast(gA_tensormap), recast(sA_tensormap)); + copy(recast(gB_tensormap), recast(sB_tensormap)); + + cp_async_fence(); + cp_async_wait<0>(); + } + + // Replace address for the global tensor (to be done by single thread) + CUTLASS_DEVICE + void + tensormaps_replace_global_address( + TensorMapStorage& shared_tensormap, + Params const& mainloop_params, + int32_t next_batch) { + // Replacing global_address for the next batch + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A, + mainloop_params.ptr_A[next_batch]); + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B, + mainloop_params.ptr_B[next_batch]); + } + + template + CUTLASS_DEVICE + void + tensormaps_perform_update( + TensorMapStorage& shared_tensormap, + Params const& mainloop_params, + cute::tuple const& input_tensormaps, + int32_t next_batch, + uint32_t lane_predicate) { + if (lane_predicate) { + // Bringing tensormaps to smem + tensormaps_fetch_to_smem(shared_tensormap, input_tensormaps); + + // Replacing global_address for the next batch + tensormaps_replace_global_address(shared_tensormap, mainloop_params, next_batch); + } + } + + template + CUTLASS_DEVICE + void + tensormaps_cp_fence_release ( + TensorMapStorage& shared_tensormap, + cute::tuple const& input_tensormaps) { + if (cute::elect_one_sync()) { + // Perform using same thread as the one that issued TMA store, separate these out as far as possible to hide latency + cute::tma_desc_commit_group(); + cute::tma_desc_wait_group(); + } + // Entire warp must do this (i.e. it's aligned) + tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormap.smem_tensormap_A); + tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormap.smem_tensormap_B); + } + + // The entire warp must call this function collectively (that is, the instructions are aligned) + template + CUTLASS_DEVICE + void + tensormaps_fence_acquire(cute::tuple const& input_tensormaps) { + cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps)); + cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps)); + } + +protected: + + template + CUTLASS_DEVICE + constexpr auto + tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const { + using X = cute::Underscore; + // Separate out problem shape for convenience + auto [M,N,K,L] = problem_shape_MNKL; + + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,L)); + Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,L)); + + // Tile the tensors and defer the slice + Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); + Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); + + return cute::make_tuple(gA_mkl, gB_nkl); + } + + typename Params::TMA_A const* observed_tma_load_a_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_ = nullptr; + + ClusterShape cluster_shape_; + uint32_t block_rank_in_cluster_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp new file mode 100644 index 00000000..9e0a6c5e --- /dev/null +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_interleaved_complex_tf32.hpp @@ -0,0 +1,987 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + +#include "cutlass/cutlass.h" +#include "cutlass/detail/cluster.hpp" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/numeric_conversion.h" +#include "cutlass/detail/sm100_tmem_helper.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/arch/mma_sm100.hpp" +#include "cutlass/trace.h" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/cuda_host_adapter.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop for complex kernels +template < + int ComputationPipelineStageCount_, + int SchedulerPipelineStageCount_, + int AccumulatorPipelineStageCount_, + int TransformationPipelineStageCount_, + class AccumulatorCopyAtom_, + class ClusterShape, // Static cluster shape or dynamic (int, int, _1) + class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + class StrideA_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomsA_, + class CopyAtomsA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomsB_, + class CopyAtomsB_, + class TransformB_> +struct CollectiveMma< + MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32< + ComputationPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + TransformationPipelineStageCount_, + ClusterShape, + AccumulatorCopyAtom_>, + TileShape_, + complex, + StrideA_, + complex, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomsA_, + CopyAtomsA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomsB_, + CopyAtomsB_, + TransformB_> +{ + // + // Type Aliases + // + using TileShape = TileShape_; + using TiledMma = TiledMma_; + + // ElementA and ElementB are cutlass::complex, which are used as GMEM input and output data type. + using ElementA = complex; + using StrideA = StrideA_; + using InternalStrideA = cute::remove_pointer_t; + using ElementB = complex; + using StrideB = StrideB_; + using InternalStrideB = cute::remove_pointer_t; + +private: + // ElementAMma and ElementBMma are cutlass::complex, which are used as SMEM and RF data type. + // ElementAMmaRaw and ElementBMmaRaw are cutlass::tfloat32_t, which is the real internal data type set in TMA descriptor and used in TCGEN05 calculation. + using ElementAMma = typename TiledMma::ValTypeA; // complex + using ElementAMmaRaw = typename ElementAMma::value_type; // tfloat32_t + using ElementBMma = typename TiledMma::ValTypeB; // complex + using ElementBMmaRaw = typename ElementBMma::value_type; // tfloat32_t + +public: + // For a complex kernel, the MMA output type is real valued, but ElementAccumulator is a complex type for the GETT reference kernel + using ElementAccumulator = cutlass::complex; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomsA = SmemLayoutAtomsA_; + using SmemLayoutAtomsB = SmemLayoutAtomsB_; + using CopyAtomsA = CopyAtomsA_; + using CopyAtomsB = CopyAtomsB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + + // Determine MMA type: MMA_1SM vs MMA_2SM + using AtomThrShapeMNK = Shape(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>; + using DispatchPolicy = MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32< + ComputationPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + TransformationPipelineStageCount_, + ClusterShape, + AccumulatorCopyAtom_>; + static constexpr bool IsDynamicCluster = not cute::is_static_v; + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + + // Define A and B block shapes for reduced size TMA_LOADs + using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + using ArchTag = typename DispatchPolicy::ArchTag; + + using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync< + DispatchPolicy::ComputationPipelineStageCount, + AtomThrShapeMNK>; + using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState; + + using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync< + DispatchPolicy::TransformationPipelineStageCount, + AtomThrShapeMNK>; + using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState; + + using Mma2AccumPipeline = cutlass::PipelineUmmaAsync< + DispatchPolicy::Schedule::AccumulatorPipelineStageCount, + AtomThrShapeMNK>; + using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState; + + // Thread Counts + static constexpr uint32_t NumTransformationThreads = 128; + static constexpr uint32_t NumAccumThreads = 128; + + // Get the Algorithm parameters + constexpr static int NumComputeMtxs = 2; + constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; + constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{}); + + // Copy atom for Accumulator + using AccumulatorCopyAtom = typename DispatchPolicy::AccumulatorCopyAtom; + + using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom; + using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom; + using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom; + using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom; + + using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom; + using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom; + using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom; + using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom; + + static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(CtaShapeA_MK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomACompute{}, + append(append(CtaShapeA_MK{}, Int{}), Int{}))); + + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(CtaShapeB_NK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutBCompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomBCompute{}, + append(CtaShapeB_NK{}, Int{}))); + + static_assert(DispatchPolicy::ComputationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more."); + static_assert(DispatchPolicy::TransformationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must have A operand from TMEM and B operand from SMEM for this mainloop."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyA - invalid TMA copy atom specified."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyB - invalid TMA copy atom specified."); + + struct PipelineStorage { + using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage; + alignas(16) Load2TransformPipelineStorage load2transform_pipeline; + using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage; + alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline; + using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage; + alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline; + }; + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + struct TensorStorageUntransformed { + cute::ArrayEngine> smem_A; + cute::ArrayEngine> smem_B; + } input; + + union TensorStorageTransformed { + alignas(1024) cute::ArrayEngine smem_ACompute; // smem_ACompute is actually in tmem + alignas(1024) cute::ArrayEngine> smem_BCompute; + } compute; + } tensors; + + struct TensorMapStorage : cute::aligned_struct<128, _0> { + cute::TmaDescriptor smem_tensormap_A; + cute::TmaDescriptor smem_tensormap_B; + } tensormaps; + + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using TensorMapStorage = typename SharedStorage::TensorMapStorage; + + // Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on + // loaded input A and B matrices to convert the data type + static constexpr uint32_t TmaTransactionBytes = + (size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast(sizeof(ElementAMma))) + + (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast(sizeof(ElementBMma))); + + // Host side kernel arguments + struct Arguments { + ElementA const** ptr_A{nullptr}; + StrideA dA{}; + ElementB const** ptr_B{nullptr}; + StrideB dB{}; + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), + make_tile(typename TiledMma::AtomThrID{}))); + + using TMA_A = decltype(make_tma_atom_A_sm100( + GmemTiledCopyA{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + TMA_A tma_load_a; + TMA_B tma_load_b; + TMA_A tma_load_a_fallback; + TMA_B tma_load_b_fallback; + dim3 cluster_shape_fallback; + cute::TmaDescriptor* tensormaps; + ElementA const** ptr_A; + ElementB const** ptr_B; + }; + + CUTLASS_DEVICE + CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster) + : cluster_shape_(cluster_shape) + , block_rank_in_cluster_(block_rank_in_cluster) { + if constexpr (IsDynamicCluster) { + const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x && + cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y); + observed_tma_load_a_ = is_fallback_cluster ? ¶ms.tma_load_a_fallback : ¶ms.tma_load_a; + observed_tma_load_b_ = is_fallback_cluster ? ¶ms.tma_load_b_fallback : ¶ms.tma_load_b; + } + else { + observed_tma_load_a_ = ¶ms.tma_load_a; + observed_tma_load_b_ = ¶ms.tma_load_b; + } + } + + template + static constexpr Params + to_underlying_arguments(ProblemShape problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) { + (void) workspace; + + // Tensor shapes for Ptr-Array are initialized correctly here. + auto [M,N,K,mock_L] = problem_shape.get_host_problem_shape(0); + // Batches/Groups are managed by using appropriate pointers to input matrices + mock_L = 1; + + // Tensor pointers will be fixed before the first access + ElementA const* ptr_A_first_batch = nullptr; + ElementB const* ptr_B_first_batch = nullptr; + + Tensor tensor_a = make_tensor(ptr_A_first_batch, make_layout(make_shape(M,K,mock_L), args.dA)); + Tensor tensor_b = make_tensor(ptr_B_first_batch, make_layout(make_shape(N,K,mock_L), args.dB)); + + auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape); + // Cluster layout for TMA construction + auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{})); + + auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback); + // Cluster layout for TMA construction + auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + return { + tma_load_a, + tma_load_b, + tma_load_a_fallback, + tma_load_b_fallback, + hw_info.cluster_shape_fallback, + reinterpret_cast(workspace), + reinterpret_cast(args.ptr_A), + reinterpret_cast(args.ptr_B) + }; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) { + constexpr uint32_t NumInputTensors = 2; + constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor); + // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies + return (NumInputTensors * SizeOfCuTensorMap * sm_count); + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) { + return cutlass::Status::kSuccess; + } + + template + static bool + can_implement( + ProblemShape problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits = 128; + auto [M,N,K,L] = problem_shape.get_host_problem_shape(0); + + bool implementable = true; + constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,K,L), StrideA{}); + constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + return implementable; + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE auto + partition_accumulator_shape() { + return append( + partition_shape_C(TiledMma{}, take<0,2>(TileShape{})), + Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2) + } + + /// Produce the inputs to the transform threads by loading inputs from gmem -> smem + template < + class GTensorA, class GTensorB, + class GTensorPartitionedA, class GTensorPartitionedB, + class STensorA, class STensorB, + class TensorMapA, class TensorMapB, + class TileCoordMNKL, + class KTileIterator + > + CUTLASS_DEVICE cute::tuple + load( + Params const& params, + Load2TransformPipeline pipeline, + Load2TransformPipelineState load2xform_pipeline_state, + cute::tuple> const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, + KTileIterator k_tile_iter, int k_tile_count) { + + auto [unused_gA, unused_gB, + tAgA_mkl, tBgB_nkl, tAsA, tBsB, + mcast_mask_a, mcast_mask_b, + input_tensormaps] = load_inputs; + + // slice out the work coord from tiled tensors + Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + uint32_t skip_wait = (k_tile_count <= 0); + auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + // Issue the Mainloop loads + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // LOCK mainloop_load2xform_pipeline_state for _writing_ + pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag); + int write_stage = load2xform_pipeline_state.index(); + + using BarrierType = typename Load2TransformPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state); + + // Advance mainloop_pipe + ++load2xform_pipeline_state; + skip_wait = (k_tile_count <= 1); + pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + copy(observed_tma_load_a_->with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage)); + copy(observed_tma_load_b_->with(get<1>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage)); + ++k_tile_iter; + } + return cute::make_tuple(load2xform_pipeline_state, k_tile_iter); + } + + + /// Set up the data needed by this collective for load. + /// Returned tuple must contain at least two elements, with the first two elements being: + /// gA_mkl - The tiled tensor for input A + /// gB_nkl - The tiled tensor for input B + // Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks + template + CUTLASS_DEVICE auto + load_init( + ProblemShape_MNKL const& problem_shape_MNKL, + Params const& params, + TensorStorage& shared_tensors, + int32_t const sm_count, int32_t const sm_idx) const { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + + Tensor sA = make_tensor(make_smem_ptr(shared_tensors.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + + // Define the CTA-in-cluster Layout and Coord + Layout cta_layout_mnk = make_layout(cluster_shape_); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + // Project the cta_layout for tma_a along the n-modes + auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl)); + + // Project the cta_layout for tma_b along the m-modes + auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl)); + + // TMA Multicast Masks + uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + // Fetch a copy of tensormaps for the CTA from Params + auto input_tensormaps = tensormaps_init(params, sm_count, sm_idx); + + return cute::make_tuple( + gA_mkl, gB_nkl, // for scheduler + tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values + mcast_mask_a, mcast_mask_b, // multicast masks + input_tensormaps); // for tma descriptor modification (per-CTA tensormap copy) + } + + template< + class KTileIterator, class Accumulator, + class GTensorA, class SrcCopyA, class DstCopyA, class SrcTensorA, class DstTensorA, + class GTensorB, class SrcCopyB, class DstCopyB, class SrcTensorB, class DstTensorB + > + CUTLASS_DEVICE auto + transform( + Load2TransformPipeline load2transform_pipeline, + Load2TransformPipelineState load2transform_pipeline_consumer_state, + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_producer_state, + Accumulator accumulators, + cute::tuple input_operands, + KTileIterator k_tile_iter, int k_tile_count) { + cutlass::arch::NamedBarrier transform_barrier(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier); + + // tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tAtACompute : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs, SmemStages (In TMEM) + // tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tBsBCompute : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM) + auto [unused_tAgA, src_copy_A, dst_copy_A, tAsA, tAtACompute, + unused_tBgB, src_copy_B, dst_copy_B, tBsB, tBsBCompute] = input_operands; + + // Create the tensors in registers + auto tArA = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArA_conj = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArA_swap = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tBrB = make_tensor(tBsB(_,_,_,_,0).shape()); + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag); + transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag); + + int load2transform_consumer_index = load2transform_pipeline_consumer_state.index(); + int transform2mma_producer_index = transform2mma_pipeline_producer_state.index(); + + auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state; + auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state; + + // Copy the input A matrix from SMEM + copy(src_copy_A, tAsA(_,_,_,_,load2transform_consumer_index), tArA); + // Copy the input B matrix from SMEM + copy(src_copy_B, tBsB(_,_,_,_,load2transform_consumer_index), tBrB); + + // First MMA, A.real * B.real - A.imag * B.imag + // Compose [real, -imag] copy for A TMEM + // Reflect the conjugation of B through A + if constexpr (cute::is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tArA); i++) { + tArA_conj(i) = {tArA(i).real(), -tArA(i).imag()}; + } + } + else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tArA); i++) { + tArA_conj(i) = tArA(i); + } + } + // Write to TMEM + copy(dst_copy_A, tArA_conj, tAtACompute(_,_,_,_,0,transform2mma_producer_index)); + + // Second MMA, A.imag * B.real + A.real * B.imag + // Compose [imag, real] copy for A TMEM + // Reflect the conjugation of B through A + auto transform_element = [] (ElementAMma const& tArA_i) -> ElementAMma { + if constexpr (cute::is_same_v && cute::is_same_v) { // CC + return {-tArA_i.imag(), -tArA_i.real()}; + } + else if constexpr (cute::is_same_v && not cute::is_same_v) { // CN/CT + return {-tArA_i.imag(), tArA_i.real()}; + } + else if constexpr (not cute::is_same_v && cute::is_same_v) { // NC/TC + return {tArA_i.imag(), -tArA_i.real()}; + } + else { // TN/NT/NN/TT + return {tArA_i.imag(), tArA_i.real()}; + } + }; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tArA); i++) { + tArA_swap(i) = transform_element(tArA(i)); + } + + // Write to TMEM + copy(dst_copy_A, tArA_swap, tAtACompute(_,_,_,_,1,transform2mma_producer_index)); + + // Write the B matrix to SMEM without any changes + copy(dst_copy_B, tBrB, tBsBCompute(_,_,_,_,transform2mma_producer_index)); + + // Loads from SMEM are done. Signal the mainloop load as early as possible + transform_barrier.sync(); + load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state); + + // fence for SMEM writes + cutlass::arch::fence_view_async_shared(); + if constexpr (is_tmem::value) { + // fence for TMEM writes if A operand is coming from TMEM + cutlass::arch::fence_view_async_tmem_store(); + } + + // Let the MMA know we are done transforming + transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state); + // Next pipeline stage + ++load2transform_pipeline_consumer_state; + ++transform2mma_pipeline_producer_state; + + skip_wait = (k_tile_count <= 1); + // Peek the next pipeline stage's barriers + load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + } + return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + transform_init( + Params const& params, + ProblemShape_MNKL const& problem_shape_MNKL, + Accumulator accumulators, + TensorStorage& shared_storage) { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); + Tensor sA = as_position_independent_swizzle_tensor(sA_orig); + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + + Tensor sB_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); + Tensor sB = as_position_independent_swizzle_tensor(sB_orig); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + + // Map input, compute, and fragment tensors to + // Copy strategies and partitioned tensors. These will become the input + // operands of the transform function. Depending on MMA atom type, the + // operands can reside in SMEM or TMEM + auto setup_copy_ops = [&] (auto tensor_input, auto input_copy_atom, + auto tensor_compute, auto make_fragment, auto compute_copy_atom) constexpr { + auto fragment_compute = make_fragment(tensor_compute); + if constexpr (cute::is_tmem>::value) { + // For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation. + // Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation + // See: TmemAllocMode::Duplicated. + Tensor tensor_input2x = [&] () constexpr { + if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) { + return make_tensor(tensor_input.data(), + logical_product(tensor_input.layout(), + make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE) + } + else { + return tensor_input; + } + }(); + + fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + // If operand comes from TMEM, create the TMEM_STORE based copy + auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,_,0,0)); + auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x); + auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute); + // Source copy is based on the source operand of TMEM_STORE copy. + auto smem2reg_tiled_copy = make_tiled_copy_S(Copy_Atom, ElementAMma>{}, reg2tmem_tiled_copy); + return cute::make_tuple(smem2reg_tiled_copy, reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute); + } + else { + // If the operand comes from SMEM, create SMEM copy. + auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute); + auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout, Stride< _8,_1>>{}, + tensor_compute(_,_,_,0).layout()); + + // Source copy is based on the source operand of copy. + auto smem2reg_tiled_copy = make_tiled_copy_S(input_copy_atom, reg2smem_tiled_copy); + auto thr_smem2reg_tiled_copy = smem2reg_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input); + auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw); + + return cute::make_tuple(smem2reg_tiled_copy, reg2smem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute); + } + }; + + auto [src_copy_A, dst_copy_A, tAsA, tAtACompute] = + setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{}); + + auto [src_copy_B, dst_copy_B, tBsB, tBsBCompute] = + setup_copy_ops(sB, InputCopyAtomB{}, sBCompute, [&](auto &arg) {return TiledMma::make_fragment_B(arg);}, ComputeCopyAtomB{}); + + return cute::make_tuple(gA_mkl, src_copy_A, dst_copy_A, tAsA, tAtACompute, + gB_nkl, src_copy_B, dst_copy_B, tBsB, tBsBCompute); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class FrgEngine, class FrgLayout, + class TensorA, class TensorB + > + CUTLASS_DEVICE auto + mma( + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_consumer_state, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_producer_state, + cute::Tensor const& accumulators, + cute::tuple const& input_operands, + int k_tile_count + ) { + TiledMma tiled_mma; + + // tCrA : (MMA), MMA_M, MMA_K, NumComputeMtxs, SmemStage (In TMEM) + // We use SMEM stages to match #buffers in Load <-> Convert + // tCrB : (MMA), MMA_N, MMA_K, SmemStages (In SMEM) + auto const [tCrA, tCrB] = input_operands; + + auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + uint32_t skip_wait = (k_tile_count <= 0); + auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + ++next_transform2mma_pipeline_consumer_state; + + mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state); + + constexpr int RealAccumIndex = 0; + constexpr int ImagAccumIndex = 1; + + int mma2accum_pipeline_producer_state_index = mma2accum_pipeline_producer_state.index(); + auto tCtC_real = accumulators(_,_,_,RealAccumIndex,mma2accum_pipeline_producer_state_index); + auto tCtC_imag = accumulators(_,_,_,ImagAccumIndex,mma2accum_pipeline_producer_state_index); + auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state; + ++mma2accum_pipeline_producer_state; + + // + // PIPELINED MAIN LOOP + // + + tiled_mma.accumulate_ = UMMA::ScaleOut::Zero; + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag); + + int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index(); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < StagesPerTile; ++k_block) { + + auto tCrA_conj = tCrA(_,_,_,Int<0>{},transform2mma_pipeline_consumer_state_index); + auto tCrA_swap = tCrA(_,_,_,Int<1>{},transform2mma_pipeline_consumer_state_index); + + auto tCrB0 = tCrB(_,_,_,transform2mma_pipeline_consumer_state_index); + + // A conjugate * B + cute::gemm(tiled_mma, tCrA_conj(_,_,k_block), tCrB0(_,_,k_block), tCtC_real); // A[0]*B[0] + // A swapped * B + cute::gemm(tiled_mma, tCrA_swap(_,_,k_block), tCrB0(_,_,k_block), tCtC_imag); // A[0]*B[0] + tiled_mma.accumulate_ = UMMA::ScaleOut::One; + } + + transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state); + + skip_wait = (k_tile_count <= 1); + transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + + curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state; + ++next_transform2mma_pipeline_consumer_state; + } + + mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state); + + return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + mma_init(cute::Tensor const& accumulators, TensorStorage& shared_storage) const { + TiledMma tiled_mma; + + Tensor tCrA = [&] () constexpr { + if constexpr (cute::is_base_of::value) { + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + return tiled_mma.make_fragment_A(sACompute); + } + else { + auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{})); + tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + return tCrA; + } + } (); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + Tensor tCrB = tiled_mma.make_fragment_B(sBCompute); + return cute::make_tuple(tCrA, tCrB); + } + + template + CUTLASS_DEVICE auto + accum_init(cute::Tensor const& accumulators, TmemCopyAtom, EpilogueTile) { + return accumulators; + } + + // + // Methods to perform different parts of TMA/Tensormap modifications + // + + CUTLASS_DEVICE auto + tensormaps_init(Params const& mainloop_params, int32_t const sm_count, int32_t const sm_idx) const { + cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps; + + cute::TmaDescriptor* tma_desc_a = &gmem_tensormap[sm_idx]; + cute::TmaDescriptor* tma_desc_b = &gmem_tensormap[sm_idx + sm_count]; + + if (cute::elect_one_sync()) { + // Bringing tensormaps from params to gmem for modification later + Tensor pA_tensormap = make_tensor(observed_tma_load_a_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor gA_tensormap = make_tensor(tma_desc_a, Int<1>{}, Int<1>{}); + Tensor pB_tensormap = make_tensor(observed_tma_load_b_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor gB_tensormap = make_tensor(tma_desc_b, Int<1>{}, Int<1>{}); + + copy(recast(pA_tensormap), recast(gA_tensormap)); + copy(recast(pB_tensormap), recast(gB_tensormap)); + } + + return cute::make_tuple(tma_desc_a, tma_desc_b); + } + + // Bringing tensormaps to smem (to be done by single thread) + template + CUTLASS_DEVICE + void + tensormaps_fetch_to_smem( + TensorMapStorage& shared_tensormap, + cute::tuple const& input_tensormaps) const { + Tensor gA_tensormap = make_tensor(make_gmem_ptr(get<0>(input_tensormaps)), Int<1>{}, Int<1>{}); + Tensor sA_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_A), Int<1>{}, Int<1>{}); + Tensor gB_tensormap = make_tensor(make_gmem_ptr(get<1>(input_tensormaps)), Int<1>{}, Int<1>{}); + Tensor sB_tensormap = make_tensor(make_smem_ptr(&shared_tensormap.smem_tensormap_B), Int<1>{}, Int<1>{}); + + copy(recast(gA_tensormap), recast(sA_tensormap)); + copy(recast(gB_tensormap), recast(sB_tensormap)); + + cp_async_fence(); + cp_async_wait<0>(); + } + + // Replace address for the global tensor (to be done by single thread) + CUTLASS_DEVICE + void + tensormaps_replace_global_address( + TensorMapStorage& shared_tensormap, + Params const& mainloop_params, + int32_t next_batch) { + // Replacing global_address for the next batch + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A, + mainloop_params.ptr_A[next_batch]); + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B, + mainloop_params.ptr_B[next_batch]); + } + + template + CUTLASS_DEVICE + void + tensormaps_perform_update( + TensorMapStorage& shared_tensormap, + Params const& mainloop_params, + cute::tuple const& input_tensormaps, + int32_t next_batch, + uint32_t lane_predicate) { + if (lane_predicate) { + // Bringing tensormaps to smem + tensormaps_fetch_to_smem(shared_tensormap, input_tensormaps); + + // Replacing global_address for the next batch + tensormaps_replace_global_address(shared_tensormap, mainloop_params, next_batch); + } + } + + template + CUTLASS_DEVICE + void + tensormaps_cp_fence_release ( + TensorMapStorage& shared_tensormap, + cute::tuple const& input_tensormaps) { + if (cute::elect_one_sync()) { + // Perform using same thread as the one that issued TMA store, separate these out as far as possible to hide latency + cute::tma_desc_commit_group(); + cute::tma_desc_wait_group(); + } + // Entire warp must do this (i.e. it's aligned) + tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormap.smem_tensormap_A); + tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormap.smem_tensormap_B); + } + + // The entire warp must call this function collectively (that is, the instructions are aligned) + template + CUTLASS_DEVICE + void + tensormaps_fence_acquire(cute::tuple const& input_tensormaps) { + cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps)); + cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps)); + } + +protected: + + template + CUTLASS_DEVICE + constexpr auto + tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const { + using X = cute::Underscore; + // Separate out problem shape for convenience + auto [M,N,K,L] = problem_shape_MNKL; + // Problem Shape and therefore strides that we construct are [M,N,K,L], but since here for the TMA loads + // we are managing TMA descriptors to change batches, we need to neglect the L mode + const int32_t mock_L = 1; + + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,mock_L)); + Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,mock_L)); + + // Tile the tensors and defer the slice + Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l) + Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l) + + return cute::make_tuple(gA_mkl, gB_nkl); + } + + typename Params::TMA_A const* observed_tma_load_a_{nullptr}; + typename Params::TMA_B const* observed_tma_load_b_{nullptr}; + + ClusterShape cluster_shape_; + uint32_t block_rank_in_cluster_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp new file mode 100644 index 00000000..44bdc8cf --- /dev/null +++ b/include/cutlass/gemm/collective/sm100_mma_array_warpspecialized_planar_complex.hpp @@ -0,0 +1,963 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/detail/collective.hpp" +#include "cutlass/detail/cluster.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/gemm/gemm.h" +#include "cutlass/trace.h" +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/cuda_host_adapter.hpp" +#include "cutlass/detail/sm100_tmem_helper.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop +// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one +template < + int Stages, + int SchedulerPipelineStageCount, + int AccumulatorPipelineStageCount, + class ClusterShape, + class TileShape_, // Static cluster shape or dynamic (int, int, _1) + class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMmaPair_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex< + Stages, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + ClusterShape>, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMmaPair_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + + // Determine MMA type: MMA_1SM vs MMA_2SM + using TiledMmaPair = TiledMmaPair_; + using TiledMma = typename TiledMmaPair::TiledMmaAPosAtom; + using TiledMmaANeg = typename TiledMmaPair::TiledMmaANegAtom; + using AtomThrShapeMNK = Shape(typename TiledMma::ThrLayoutVMNK{})), _1, _1>; + + static constexpr bool IsDynamicCluster = not cute::is_static_v; + + using DispatchPolicy = MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex< + Stages, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + ClusterShape>; + using TileShape = TileShape_; + + CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})), + "Static cluster shape used: TileShape should be evenly divided by TiledMma"); + + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + + // Define A and B block shapes for reduced size TMA_LOADs + using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + // Multiple buffer the TMA descriptors for each SM so that we can update them asynchronously. + // This should be larger than the total number of TMA requests inflight (from update to issued to returned). + // This can be calculated by SchedulerStages + max(TmaStages) + 2 (for consumer and producer in-flight accessies). + constexpr static uint32_t NumTmaDescriptorsPerSm = SchedulerPipelineStageCount + Stages + 2; + + using ElementA = ElementA_; + using ElementAMma = typename TiledMma::ValTypeA; + using StrideA = StrideA_; + using InternalStrideA = cute::remove_pointer_t; + using ElementB = ElementB_; + using ElementBMma = typename TiledMma::ValTypeB; + using StrideB = StrideB_; + using InternalStrideB = cute::remove_pointer_t; + 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; + + using MainloopPipeline = cutlass::PipelineTmaUmmaAsync< + DispatchPolicy::Stages, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M, K)"); + static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(cute::is_void_v, + "SM100 UMMA cannot have a non-void copy atom for smem sourced instructions."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N, K)"); + static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(cute::is_void_v, + "SM100 UMMA cannot have a non-void copy atom for smem sourced instructions."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(MmaShapeA_MK{}, Int{}), + cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(MmaShapeB_NK{}, Int{}), + cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source both A and B operand from smem_desc for this mainloop."); + static_assert( + (size(AtomThrShapeMNK{}) == 1 && + (cute::is_same_v || cute::is_same_v)) || + (size(AtomThrShapeMNK{}) == 2 && + (cute::is_same_v || cute::is_same_v)), + "GmemTiledCopy - invalid TMA copy atom specified."); + static_assert( + (size(AtomThrShapeMNK{}) == 1 && + (cute::is_same_v || cute::is_same_v)) || + (size(AtomThrShapeMNK{}) == 2 && + (cute::is_same_v || cute::is_same_v)), + "GmemTiledCopy - invalid TMA copy atom specified."); + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + cute::ArrayEngine> smem_A_real; + cute::ArrayEngine> smem_A_imag; + cute::ArrayEngine> smem_B_real; + cute::ArrayEngine> smem_B_imag; + } tensors; + + struct TensorMapStorage : cute::aligned_struct<128, _0> { + cute::TmaDescriptor smem_tensormap_A_real; + cute::TmaDescriptor smem_tensormap_A_imag; + cute::TmaDescriptor smem_tensormap_B_real; + cute::TmaDescriptor smem_tensormap_B_imag; + } tensormaps; + + using PipelineStorage = typename MainloopPipeline::SharedStorage; + PipelineStorage pipeline; + }; + + // Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them. + using TensorStorage = typename SharedStorage::TensorStorage; + using TensorMapStorage = typename SharedStorage::TensorMapStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly + static constexpr uint32_t TmaTransactionBytes = 2 * ( + cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutA{})) * static_cast(cute::sizeof_bits::value))) + + cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutB{})) * static_cast(cute::sizeof_bits::value)))); + + template + struct TmemStorage { + AccTensor accumulators; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const** ptr_A_real{nullptr}; + StrideA dA_real{}; + ElementA const** ptr_A_imag{nullptr}; + StrideA dA_imag{}; + ElementB const** ptr_B_real{nullptr}; + StrideB dB_real{}; + ElementB const** ptr_B_imag{nullptr}; + StrideB dB_imag{}; + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), make_tile(typename TiledMma::AtomThrID{}))); + + using TMA_A = decltype(make_tma_atom_A_sm100( + GmemTiledCopyA{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + + TMA_A tma_load_a_real; + TMA_A tma_load_a_imag; + TMA_B tma_load_b_real; + TMA_B tma_load_b_imag; + TMA_A tma_load_a_real_fallback; + TMA_A tma_load_a_imag_fallback; + TMA_B tma_load_b_real_fallback; + TMA_B tma_load_b_imag_fallback; + dim3 cluster_shape_fallback; + cute::TmaDescriptor* tensormaps; + ElementA const** ptr_A_real; + ElementA const** ptr_A_imag; + ElementB const** ptr_B_real; + ElementB const** ptr_B_imag; + }; + + CUTLASS_DEVICE + CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster) + : cluster_shape_(cluster_shape) + , block_rank_in_cluster_(block_rank_in_cluster) { + if constexpr (IsDynamicCluster) { + const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x && + cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y); + observed_tma_load_a_real_ = is_fallback_cluster ? ¶ms.tma_load_a_real_fallback : ¶ms.tma_load_a_real; + observed_tma_load_a_imag_ = is_fallback_cluster ? ¶ms.tma_load_a_imag_fallback : ¶ms.tma_load_a_imag; + observed_tma_load_b_real_ = is_fallback_cluster ? ¶ms.tma_load_b_real_fallback : ¶ms.tma_load_b_real; + observed_tma_load_b_imag_ = is_fallback_cluster ? ¶ms.tma_load_b_imag_fallback : ¶ms.tma_load_b_imag; + } + else { + observed_tma_load_a_real_ = ¶ms.tma_load_a_real; + observed_tma_load_a_imag_ = ¶ms.tma_load_a_imag; + observed_tma_load_b_real_ = ¶ms.tma_load_b_real; + observed_tma_load_b_imag_ = ¶ms.tma_load_b_imag; + } + } + + template + static constexpr Params + to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + void* workspace, + cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) { + // Tensor shapes for Ptr-Array are initialized correctly here. + auto [M,N,K,mock_L] = problem_shape.get_host_problem_shape(0); + + // Batches/Groups are managed by using appropriate pointers to input matrices + mock_L = 1; + + // Tensor pointers will be fixed before the first access + ElementA const* ptr_A_real_first_batch = nullptr; + ElementA const* ptr_A_imag_first_batch = nullptr; + + ElementB const* ptr_B_real_first_batch = nullptr; + ElementB const* ptr_B_imag_first_batch = nullptr; + + Tensor tensor_a_real = make_tensor(ptr_A_real_first_batch, make_layout(make_shape(M,K,mock_L), args.dA_real)); + Tensor tensor_a_imag = make_tensor(ptr_A_imag_first_batch, make_layout(make_shape(M,K,mock_L), args.dA_imag)); + + Tensor tensor_b_real = make_tensor(ptr_B_real_first_batch, make_layout(make_shape(N,K,mock_L), args.dB_real)); + Tensor tensor_b_imag = make_tensor(ptr_B_imag_first_batch, make_layout(make_shape(N,K,mock_L), args.dB_imag)); + + auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape); + // Cluster layout for TMA construction + auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{})); + + auto cluster_shape_fallback = conditional_return(make_shape(hw_info.cluster_shape_fallback.x, hw_info.cluster_shape_fallback.y, Int<1>{}), ClusterShape{}); + // Cluster layout for TMA construction + auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_A tma_load_a_real = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_real, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_imag = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_imag, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b_real = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_real, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b_imag = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_imag, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_real_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_real, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_A tma_load_a_imag_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_imag, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_real_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_real, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_imag_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_imag, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + return { + tma_load_a_real, + tma_load_a_imag, + tma_load_b_real, + tma_load_b_imag, + tma_load_a_real_fallback, + tma_load_a_imag_fallback, + tma_load_b_real_fallback, + tma_load_b_imag_fallback, + hw_info.cluster_shape_fallback, + reinterpret_cast(workspace), + reinterpret_cast(args.ptr_A_real), + reinterpret_cast(args.ptr_A_imag), + reinterpret_cast(args.ptr_B_real), + reinterpret_cast(args.ptr_B_imag) + }; + } + + template + static size_t + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args, int sm_count) { + constexpr uint32_t NumInputTensors = 4; + constexpr size_t SizeOfCuTensorMap = sizeof(cute::TmaDescriptor); + // Allocate gmem space for input tensormaps per each SM, A tensormap copies followed by B tensormap copies + return (NumInputTensors * SizeOfCuTensorMap * sm_count * NumTmaDescriptorsPerSm); + } + + template + static cutlass::Status + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream, CudaHostAdapter* cuda_adapter = nullptr) { + return cutlass::Status::kSuccess; + } + + template + static bool + can_implement( + ProblemShape problem_shape, + [[maybe_unused]] Arguments const& args) { + auto [M,N,K,L] = problem_shape.get_host_problem_shape(0); + + bool implementable = true; + constexpr int min_tma_aligned_elements_A = 128 / cute::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,K,L), StrideA{}); + constexpr int min_tma_aligned_elements_B = 128 / cute::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + return implementable; + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE static + auto + partition_accumulator_shape() { + return append(partition_shape_C(TiledMma{}, take<0,2>(TileShape{})), Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2) + } + + template + CUTLASS_DEVICE static + auto + slice_accumulator(TmemStorage tmem_storage, int stage) { + return tmem_storage.accumulators(_,_,_,_,stage); + } + + template + CUTLASS_DEVICE static + auto + init_tmem_tensors(EpilogueTile epi_tile) { + TiledMma tiled_mma; + auto acc_shape = partition_accumulator_shape(); + // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue. + Tensor accumulators = cutlass::detail::make_sm100_accumulator( + tiled_mma, acc_shape, EpilogueTile{}); + TmemStorage tmem_storage; + tmem_storage.accumulators = accumulators; + return tmem_storage; + } + + template + CUTLASS_DEVICE static + void + set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) { + tmem_storage.accumulators.data() = tmem_base_addr; + } + + /// Set up the data needed by this collective for load. + /// Returned tuple must contain at least two elements, with the first two elements being: + /// gA_(real/imag)_mkl - The tiled tma tensor for input A_(real/imag) + /// gB_(real/imag)_nkl - The tiled tma tensor for input B_(real/imag) + /// tAsA_(real/imag) - partitioned smem tensor for A_(real/imag) + /// tBsB_(real/imag) - partitioned smem tensor for B_(real/imag) + /// mcast_mask_a - tma multicast mask for A_(real/imag) + /// mcast_mask_b - tma multicast mask for B_(real/imag) + template + CUTLASS_DEVICE auto + load_init( + ProblemShape_MNKL const& problem_shape_MNKL, + Params const& params, + TensorStorage& shared_tensors, + TensorMapStorage& shared_tensormaps, + int32_t const sm_count, int32_t const sm_idx, + [[maybe_unused]] int32_t num_groups, + [[maybe_unused]] int32_t init_group) const { + using X = Underscore; + + // Separate out problem shape for convenience + auto [M,N,K,L] = problem_shape_MNKL; + // Problem Shape and therefore strides that we construct are [M,N,K,L], but since here for the TMA loads + // we are managing TMA descriptors to change batches, we need to neglect the L mode + const int32_t mock_L = 1; + + // Represent the full tensors -- get these from TMA + Tensor mA_real_mkl = observed_tma_load_a_real_->get_tma_tensor(make_shape(M,K,mock_L)); + Tensor mA_imag_mkl = observed_tma_load_a_imag_->get_tma_tensor(make_shape(M,K,mock_L)); + Tensor mB_real_nkl = observed_tma_load_b_real_->get_tma_tensor(make_shape(N,K,mock_L)); + Tensor mB_imag_nkl = observed_tma_load_b_imag_->get_tma_tensor(make_shape(N,K,mock_L)); + + // Tile the tensors and defer the slice + Tensor gA_real_mkl = local_tile(mA_real_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l) + Tensor gA_imag_mkl = local_tile(mA_imag_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_N, BLK_K, m, k, l) + + Tensor gB_real_nkl = local_tile(mB_real_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l) + Tensor gB_imag_nkl = local_tile(mB_imag_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l) + + // Partition for this CTA + ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + Tensor tCgA_real_mkl = cta_mma.partition_A(gA_real_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + Tensor tCgA_imag_mkl = cta_mma.partition_A(gA_imag_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + + Tensor tCgB_real_nkl = cta_mma.partition_B(gB_real_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + Tensor tCgB_imag_nkl = cta_mma.partition_B(gB_imag_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + + Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + + Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + + // Define the CTA-in-cluster Layout and Coord + Layout cta_layout_mnk = make_layout(cluster_shape_); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + // Project the cta_layout for tma_a along the n-modes + auto [tAgA_real_mkl, tAsA_real] = tma_partition(*observed_tma_load_a_real_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA_real), group_modes<0,3>(tCgA_real_mkl)); + auto [tAgA_imag_mkl, tAsA_imag] = tma_partition(*observed_tma_load_a_imag_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA_imag), group_modes<0,3>(tCgA_imag_mkl)); + + // Project the cta_layout for tma_b along the m-modes + auto [tBgB_real_nkl, tBsB_real] = tma_partition(*observed_tma_load_b_real_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB_real), group_modes<0,3>(tCgB_real_nkl)); + auto [tBgB_imag_nkl, tBsB_imag] = tma_partition(*observed_tma_load_b_imag_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB_imag), group_modes<0,3>(tCgB_imag_nkl)); + // TMA Multicast Masks + uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + auto ret = cute::make_tuple( + gA_real_mkl, gA_imag_mkl, gB_real_nkl, gB_imag_nkl, // for scheduler + tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl, // for input tensor values + tAsA_real, tAsA_imag, tBsB_real, tBsB_imag, // for input tensor values + mcast_mask_a, mcast_mask_b // multicast masks + ); + + if constexpr (IsTensorMapUpdateAsync) { + return ret; + } + else { + // Fetch a copy of tensormaps for the CTA from Params + auto input_tensormaps = tensormaps_init(params, shared_tensormaps, sm_count, sm_idx); + return cute::tuple_cat(ret, cute::make_tuple(input_tensormaps)); + } + } + + /// Set up the data needed by this collective for mma compute. + template + CUTLASS_DEVICE auto + mma_init( + [[maybe_unused]] TmemStorage tmem_storage, + TensorStorage& shared_tensors) const { + Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + + Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // Allocate "fragments/descriptors" for A and B matrices + Tensor tCrA_real = TiledMma::make_fragment_A(sA_real); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrA_imag = TiledMma::make_fragment_A(sA_imag); // (MMA,MMA_M,MMA_K,PIPE) + + Tensor tCrB_real = TiledMma::make_fragment_B(sB_real); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrB_imag = TiledMma::make_fragment_B(sB_imag); // (MMA,MMA_N,MMA_K,PIPE) + + CUTE_STATIC_ASSERT_V(Int{} == size<3>(sA_real)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<3>(sB_real)); // PIPE + + TiledMma tiled_mma_a_pos; + TiledMmaANeg tiled_mma_a_neg; + + return cute::make_tuple(tiled_mma_a_pos, tiled_mma_a_neg, tCrA_real, tCrA_imag, tCrB_real, tCrB_imag); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Producer Perspective + template < + class GTensorA, class GTensorB, + class GTensorPartitionedA, class GTensorPartitionedB, + class STensorA, class STensorB, + class TensorMapA, class TensorMapB, + class TileCoordMNKL, + class KTileIterator + > + CUTLASS_DEVICE auto + load( + Params const& params, + MainloopPipeline mainloop_pipeline, + MainloopPipelineState mainloop_pipe_producer_state, + cute::tuple> const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, + KTileIterator k_tile_iter, int k_tile_count, + bool did_batch_change, + [[maybe_unused]] int curr_batch) { + + auto [unused_gA_real, unused_gA_imag, unused_gB_real, unused_gB_imag, + tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl, + tAsA_real, tAsA_imag, tBsB_real, tBsB_imag, + mcast_mask_a, mcast_mask_b, + input_tensormaps] = load_inputs; + + // Check to see if tensormaps have been replaced in gmem + if (did_batch_change) { + tensormaps_fence_acquire(input_tensormaps); + } + + // slice out the work coord from partitioned tensors + Tensor tAgA_real = tAgA_real_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + Tensor tAgA_imag = tAgA_imag_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + + Tensor tBgB_real = tBgB_real_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + Tensor tBgB_imag = tBgB_imag_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state); + + // Issue the Mainloop loads + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + // LOCK mainloop_pipe_producer_state for _writing_ + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token); + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state); + + if (cute::elect_one_sync()) { + copy(observed_tma_load_a_real_->with(get<0>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA_real(_,*k_tile_iter), tAsA_real(_,write_stage)); + copy(observed_tma_load_a_imag_->with(get<1>(input_tensormaps), *tma_barrier, mcast_mask_a), tAgA_imag(_,*k_tile_iter), tAsA_imag(_,write_stage)); + + copy(observed_tma_load_b_real_->with(get<2>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB_real(_,*k_tile_iter), tBsB_real(_,write_stage)); + copy(observed_tma_load_b_imag_->with(get<3>(input_tensormaps), *tma_barrier, mcast_mask_b), tBgB_imag(_,*k_tile_iter), tBsB_imag(_,write_stage)); + } + --k_tile_count; + ++k_tile_iter; + } + + return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter); + } + + /// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster + CUTLASS_DEVICE void + load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) { + // Issue the epilogue waits + /* This helps avoid early exit of ctas in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class AccumulatorPipeline, + class FrgEngine, class FrgLayout, + class FragmentA, class FragmentB, + class CtaTileCoord + > + CUTLASS_DEVICE auto + mma(cute::tuple pipelines, + cute::tuple pipeline_states, + cute::Tensor const& accumulators, + cute::tuple const& mma_inputs, + CtaTileCoord cta_tile_coord, + int k_tile_count + ) { + static_assert(is_tmem::value, "Accumulator must be tmem resident."); + static_assert(rank(FrgLayout{}) == 4 && size<3>(FrgLayout{}) == _2{}, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N, _2)"); + + auto [tiled_mma_a_pos, tiled_mma_a_neg, tCrA_real, tCrA_imag, tCrB_real, tCrB_imag] = mma_inputs; + + auto [mainloop_pipeline, accumulator_pipeline] = pipelines; + auto [mainloop_pipe_consumer_state, accumulator_pipe_producer_state] = pipeline_states; + + uint32_t skip_wait = k_tile_count <= 0; + auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + + // + // PIPELINED MAIN LOOP + // + tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::Zero; + tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::Zero; + // Wait for tmem accumulator buffer to become empty with a flipped phase + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + + auto accumulators_real = accumulators(_,_,_,0); + auto accumulators_imag = accumulators(_,_,_,1); + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + // WAIT on mainloop_pipe_consumer_state until its data are available + // (phase bit flips from mainloop_pipe_consumer_state.phase() value) + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + + // Compute on k_tile + int read_stage = mainloop_pipe_consumer_state.index(); + // Save current mainlop pipeline read state + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + + // Advance mainloop_pipe + ++mainloop_pipe_consumer_state; + --k_tile_count; + skip_wait = k_tile_count <= 0; + // Peek at next iteration + barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + + // Unroll the K mode manually so we can set scale C to 1 + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA_real); ++k_block) { + // (V,M) x (V,N) => (V,M,N) + + // Calculate real acc, 1st step + // realAcc += realA * realB + cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_real); + + // Calculate imag acc, 1st step + if constexpr (cute::is_same_v) { + // imagAcc += realA * (-imagB) + cute::gemm(tiled_mma_a_neg, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag); + } + else { + // imagAcc += realA * imagB + cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag); + } + + tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::One; + tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::One; + + // Calculate real acc, 2nd step + if constexpr (cute::is_same_v) { + // realAcc -= imagA * imagB + cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real); + } + else { + // realAcc += imagA * imagB + cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real); + } + + // Calculate imag acc, 2nd step + if constexpr (cute::is_same_v) { + // imagAcc += (-imagA) * realB + cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag); + } + else { + // imagAcc += imagA * realB + cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag); + } + } + mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state); + } + + return mainloop_pipe_consumer_state; + } + + // + // Methods to perform different parts of TMA/Tensormap modifications + // + + template + CUTLASS_DEVICE auto + tensormaps_init( + Params const& mainloop_params, + TensorMapStorage& shared_tensormaps, + int32_t const sm_count, + int32_t const sm_idx) const { + cute::TmaDescriptor* gmem_tensormap = mainloop_params.tensormaps; + + cute::TmaDescriptor* tma_desc_a_real = &gmem_tensormap[sm_idx * NumTmaDescriptorsPerSm]; + cute::TmaDescriptor* tma_desc_a_imag = &gmem_tensormap[(sm_idx + sm_count) * NumTmaDescriptorsPerSm]; + + cute::TmaDescriptor* tma_desc_b_real = &gmem_tensormap[(sm_idx + 2 * sm_count) * NumTmaDescriptorsPerSm]; + cute::TmaDescriptor* tma_desc_b_imag = &gmem_tensormap[(sm_idx + 3 * sm_count) * NumTmaDescriptorsPerSm]; + + + if (cute::elect_one_sync()) { + // Bringing tensormaps from params to smem for modification later + Tensor pA_real_tensormap = make_tensor(observed_tma_load_a_real_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sA_real_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_A_real), Int<1>{}, Int<1>{}); + Tensor pA_imag_tensormap = make_tensor(observed_tma_load_a_imag_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sA_imag_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_A_imag), Int<1>{}, Int<1>{}); + + Tensor pB_real_tensormap = make_tensor(observed_tma_load_b_real_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sB_real_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_B_real), Int<1>{}, Int<1>{}); + Tensor pB_imag_tensormap = make_tensor(observed_tma_load_b_imag_->get_tma_descriptor(), Int<1>{}, Int<1>{}); + Tensor sB_imag_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_B_imag), Int<1>{}, Int<1>{}); + + copy(recast(pA_real_tensormap), recast(sA_real_tensormap)); + copy(recast(pA_imag_tensormap), recast(sA_imag_tensormap)); + + copy(recast(pB_real_tensormap), recast(sB_real_tensormap)); + copy(recast(pB_imag_tensormap), recast(sB_imag_tensormap)); + } + __syncwarp(); + + struct TensorMapArray { + cute::TmaDescriptor* tma_desc_a_real; + cute::TmaDescriptor* tma_desc_a_imag; + cute::TmaDescriptor* tma_desc_b_real; + cute::TmaDescriptor* tma_desc_b_imag; + + TensorMapArray() = default; + + CUTLASS_DEVICE + TensorMapArray(cute::TmaDescriptor* tma_desc_a_real, cute::TmaDescriptor* tma_desc_a_imag, + cute::TmaDescriptor* tma_desc_b_real, cute::TmaDescriptor* tma_desc_b_imag) + : tma_desc_a_real(tma_desc_a_real), tma_desc_a_imag(tma_desc_a_imag), + tma_desc_b_real(tma_desc_b_real), tma_desc_b_imag(tma_desc_b_imag) {} + + CUTLASS_DEVICE + cute::tuple + operator[](int32_t idx) { + idx = idx % NumTmaDescriptorsPerSm; + return cute::make_tuple(tma_desc_a_real + idx, tma_desc_a_imag + idx, + tma_desc_b_real + idx, tma_desc_b_imag + idx); + } + }; + + if constexpr (IsTensorMapUpdateAsync) { + return TensorMapArray(tma_desc_a_real, tma_desc_a_imag, tma_desc_b_real, tma_desc_b_imag); + } + else { + return cute::make_tuple(tma_desc_a_real, tma_desc_a_imag, tma_desc_b_real, tma_desc_b_imag); + } + } + + // Replace address for the global tensor (to be done by single thread) + CUTLASS_DEVICE + void + tensormaps_replace_global_address( + TensorMapStorage& shared_tensormap, + Params const& mainloop_params, + int32_t next_batch) { + // Replacing global_address for the next batch + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A_real, + mainloop_params.ptr_A_real[next_batch]); + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_A_imag, + mainloop_params.ptr_A_imag[next_batch]); + + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B_real, + mainloop_params.ptr_B_real[next_batch]); + cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormap.smem_tensormap_B_imag, + mainloop_params.ptr_B_imag[next_batch]); + } + + // The entire warp must call this function collectively (that is, the instructions are aligned) + template + CUTLASS_DEVICE + void + tensormaps_perform_update( + TensorMapStorage& shared_tensormaps, + Params const& mainloop_params, + cute::tuple const& input_tensormaps, + [[maybe_unused]]ProblemShape problem_shape, + int32_t next_batch + ) { + if (cute::elect_one_sync()) { + // Replacing global_address for the next batch + tensormaps_replace_global_address(shared_tensormaps, mainloop_params, next_batch); + } + // Ensure warp is converged before issuing tensormap fence release + __syncwarp(); + // Entire warp must do this (ie its aligned) + tensormaps_cp_fence_release( + shared_tensormaps, + input_tensormaps + ); + } + + template + CUTLASS_DEVICE + void + tensormaps_cp_fence_release ( + TensorMapStorage& shared_tensormap, + cute::tuple const& input_tensormaps + ) { + // Entire warp must do this (i.e., it's aligned) + tma_descriptor_cp_fence_release(get<0>(input_tensormaps), shared_tensormap.smem_tensormap_A_real); + tma_descriptor_cp_fence_release(get<1>(input_tensormaps), shared_tensormap.smem_tensormap_A_imag); + + tma_descriptor_cp_fence_release(get<2>(input_tensormaps), shared_tensormap.smem_tensormap_B_real); + tma_descriptor_cp_fence_release(get<3>(input_tensormaps), shared_tensormap.smem_tensormap_B_imag); + } + + // The entire warp must call this function collectively (that is, the instructions are aligned) + template + CUTLASS_DEVICE + void + tensormaps_fence_acquire(cute::tuple const& input_tensormaps) { + cute::tma_descriptor_fence_acquire(get<0>(input_tensormaps)); + cute::tma_descriptor_fence_acquire(get<1>(input_tensormaps)); + cute::tma_descriptor_fence_acquire(get<2>(input_tensormaps)); + cute::tma_descriptor_fence_acquire(get<3>(input_tensormaps)); + } + +protected: + + typename Params::TMA_A const* observed_tma_load_a_real_ = nullptr; + typename Params::TMA_A const* observed_tma_load_a_imag_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_real_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_imag_ = nullptr; + + ClusterShape cluster_shape_; + uint32_t block_rank_in_cluster_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp new file mode 100644 index 00000000..140b7854 --- /dev/null +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_emulated.hpp @@ -0,0 +1,1072 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/numeric_conversion.h" +#include "cutlass/detail/sm100_tmem_helper.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/arch/mma_sm100.hpp" +#include "cutlass/trace.h" +#include "cutlass/kernel_hardware_info.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop for FastF32 Kernels: Interleaved complex variants +template < + int Load2TransformPipelineStageCount_, + int Transform2MmaPipelineStageCount_, + int SchedulerPipelineStageCount_, + int AccumulatorPipelineStageCount_, + int NumBandsToCompute_, + int ScalingFactor_, + int AccPromotionInterval_, + class AccumulatorCopyAtom_, + class ClusterShape, + class TileShape_, + class StrideA_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomsA_, + class CopyAtomsA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomsB_, + class CopyAtomsB_, + class TransformB_> +struct CollectiveMma< + MainloopSm100TmaUmmaWarpSpecializedFastF32< + Load2TransformPipelineStageCount_, + Transform2MmaPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + NumBandsToCompute_, + ScalingFactor_, + AccPromotionInterval_, + ClusterShape, + AccumulatorCopyAtom_>, + TileShape_, + complex, + StrideA_, + complex, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomsA_, + CopyAtomsA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomsB_, + CopyAtomsB_, + TransformB_> +{ + // + // Type Aliases + // + using TileShape = TileShape_; + using TiledMma = TiledMma_; + + // ElementA and ElementB are cutlass::complex, which are used as GMEM input and output data type. + using ElementA = complex; + using StrideA = StrideA_; + using ElementB = complex; + using StrideB = StrideB_; + + // For a complex kernel, the MMA output type is real valued, but ElementAccumulator is a complex type for the GETT reference kernel + using ElementAccumulator = complex; + using ElementAccumulatorRaw = typename TiledMma::ValTypeC; + +private: + // ElementAMma and ElementBMma are cutlass::complex, which are used as SMEM and RF data type. + // ElementAMmaRaw and ElementBMmaRaw are cutlass::bfloat16_t, which is the real internal data type set in TMA descriptor and used in TCGEN05 calculation. + using ElementAMma = typename TiledMma::ValTypeA; // complex + using ElementAMmaRaw = typename ElementAMma::value_type; // bfloat16_t + using ElementBMma = typename TiledMma::ValTypeB; // complex + using ElementBMmaRaw = typename ElementBMma::value_type; // bfloat16_t + +public: + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomsA = SmemLayoutAtomsA_; + using SmemLayoutAtomsB = SmemLayoutAtomsB_; + using CopyAtomsA = CopyAtomsA_; + using CopyAtomsB = CopyAtomsB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + + static_assert(cute::is_same_v, "Underlying input type for A should be float"); + static_assert(cute::is_same_v, "Underlying input type for B should be float"); + static_assert(cute::is_same_v, "Underlying compute type for A should be bfloat16_t"); + static_assert(cute::is_same_v, "Underlying compute type for A should be bfloat16_t"); + + // Determine MMA type: MMA_1SM vs MMA_2SM + using AtomThrShapeMNK = Shape(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>; + using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedFastF32< + Load2TransformPipelineStageCount_, + Transform2MmaPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + NumBandsToCompute_, + ScalingFactor_, + AccPromotionInterval_, + ClusterShape, + AccumulatorCopyAtom_>; + static constexpr bool IsDynamicCluster = not cute::is_static_v; + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + + // Define A and B block shapes for reduced size TMA_LOADs + using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + using ArchTag = typename DispatchPolicy::ArchTag; + + using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync< + DispatchPolicy::Load2TransformPipelineStageCount, + AtomThrShapeMNK>; + using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState; + + using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync< + DispatchPolicy::Transform2MmaPipelineStageCount, + AtomThrShapeMNK>; + using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState; + + using Mma2AccumPipeline = cutlass::PipelineUmmaAsync< + DispatchPolicy::Schedule::AccumulatorPipelineStageCount, + AtomThrShapeMNK>; + using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState; + + // Thread Counts + static constexpr uint32_t NumTransformationThreads = 128; + static constexpr uint32_t NumAccumThreads = 128; + + // Get the Algorithm parameters + constexpr static int NumComputeMtxs = 3; + constexpr static int ConjSwapMode = 2; + constexpr static int NumBandsToCompute = DispatchPolicy::NumBandsToCompute; + constexpr static int ScalingFactor = DispatchPolicy::ScalingFactor; + constexpr static int AccPromotionInterval = DispatchPolicy::AccPromotionInterval; + constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; + constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{}) / DispatchPolicy::AccPromotionInterval; + constexpr static int NumBandsMax = 5; + static_assert(NumBandsToCompute <= NumBandsMax && NumBandsToCompute >= 3, "NumBandsToCompute should be less than maximum number of bands"); + static_assert(StagesPerTile * AccPromotionInterval == size<2>(CtaShapeA_MK{}), "PromotionInterval*InstructionK doesn't evenly divide CTA shape"); + + // Copy atom for Accumulator + using AccumulatorCopyAtom = typename DispatchPolicy::AccumulatorCopyAtom; + + static_assert((NumBandsToCompute == 5 || NumBandsToCompute == 4 || NumBandsToCompute == 3), + "9xBF16 with 5/4/3 Bands are supported"); + + using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom; + using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom; + using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom; + using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom; + + using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom; + using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom; + using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom; + using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtomCompute must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(CtaShapeA_MK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomACompute{}, + tuple_cat(CtaShapeA_MK{}, tuple, Int>{}))); + + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(CtaShapeB_NK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutBCompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomBCompute{}, + tuple_cat(CtaShapeB_NK{}, tuple, Int, Int>{}))); + + static_assert(DispatchPolicy::Load2TransformPipelineStageCount >= 2 && DispatchPolicy::Load2TransformPipelineStageCount >= 2, + "Specialization requires Stages set to value 2 or more."); + static_assert((cute::is_base_of::value || + cute::is_base_of::value ) && + cute::is_base_of::value, + "MMA atom must A operand from SMEM or TMEM and B operand from SMEM for this mainloop."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyA - invalid TMA copy atom specified."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyB - invalid TMA copy atom specified."); + + struct PipelineStorage { + using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage; + alignas(16) Load2TransformPipelineStorage load2transform_pipeline; + using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage; + alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline; + using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage; + alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline; + }; + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + struct TensorStorageUntransformed { + cute::ArrayEngine> smem_A; + cute::ArrayEngine> smem_B; + }; + + struct TensorStorageTransformedAinSmem { + alignas(1024) cute::ArrayEngine> smem_ACompute; + alignas(1024) cute::ArrayEngine> smem_BCompute; + }; + + union TensorStorageTransformedAinTmem { + alignas(1024) cute::ArrayEngine smem_ACompute; // No smem_ACompute + alignas(1024) cute::ArrayEngine> smem_BCompute; + }; + + using TensorStorageTransformed = cute::conditional_t< + cute::is_base_of::value, + TensorStorageTransformedAinSmem, + TensorStorageTransformedAinTmem>; + + TensorStorageUntransformed input; + TensorStorageTransformed compute; + } tensors; + + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + + // Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on + // loaded input A and B matrices to convert the data type + static constexpr uint32_t TmaTransactionBytes = + cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast(sizeof_bits::value))+ + cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast(sizeof_bits::value)); + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A{nullptr}; + StrideA dA{}; + ElementB const* ptr_B{nullptr}; + StrideB dB{}; + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), + make_tile(typename TiledMma::AtomThrID{}))); + + using TMA_A = decltype(make_tma_atom_A_sm100( + GmemTiledCopyA{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + TMA_A tma_load_a; + TMA_B tma_load_b; + TMA_A tma_load_a_fallback; + TMA_B tma_load_b_fallback; + dim3 cluster_shape_fallback; + }; + + CUTLASS_DEVICE + CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster) + : cluster_shape_(cluster_shape) + , block_rank_in_cluster_(block_rank_in_cluster) { + if constexpr (IsDynamicCluster) { + const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x && + cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y); + observed_tma_load_a_ = is_fallback_cluster ? ¶ms.tma_load_a_fallback : ¶ms.tma_load_a; + observed_tma_load_b_ = is_fallback_cluster ? ¶ms.tma_load_b_fallback : ¶ms.tma_load_b; + } + else { + observed_tma_load_a_ = ¶ms.tma_load_a; + observed_tma_load_b_ = ¶ms.tma_load_b; + } + } + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) { + (void) workspace; + + // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA)); + Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB)); + + auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape); + // Cluster layout for TMA construction + auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{})); + + auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback); + // Cluster layout for TMA construction + auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + return { + tma_load_a, + tma_load_b, + tma_load_a_fallback, + tma_load_b_fallback, + hw_info.cluster_shape_fallback + }; + } + + template + static bool + can_implement( + ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits = 128; + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + bool implementable = true; + constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,K,L), StrideA{}); + constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + return implementable; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE static void + prefetch_tma_descriptors(Params const& params) { + if constexpr (IsDynamicCluster) { + dim3 cs = cute::cluster_shape(); + const bool is_fallback_cluster = (cs.x == params.cluster_shape_fallback.x && cs.y == params.cluster_shape_fallback.y); + if (is_fallback_cluster) { + cute::prefetch_tma_descriptor(params.tma_load_a_fallback.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b_fallback.get_tma_descriptor()); + } + else { + cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + } + } + else { + cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + } + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE auto + partition_accumulator_shape() { + auto acc_shape = partition_shape_C(TiledMma{}, take<0,2>(TileShape{})); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N) + + return acc_shape; + } + + /// Produce the inputs to the transform threads by loading inputs from gmem -> smem + template < + class GTensorA, class GTensorB, + class GTensorPartitionedA, class GTensorPartitionedB, + class STensorA, class STensorB, + class TileCoordMNKL, + class KTileIterator + > + CUTLASS_DEVICE auto + load( + Params const& params, + Load2TransformPipeline pipeline, + Load2TransformPipelineState load2xform_pipeline_state, + cute::tuple const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, + KTileIterator k_tile_iter, int k_tile_count) { + + auto [unused_gA, unused_gB, + tAgA_mkl, tBgB_nkl, tAsA, tBsB, + mcast_mask_a, mcast_mask_b] = load_inputs; + + // slice out the work coord from tiled tensors + Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + uint32_t skip_wait = (k_tile_count <= 0); + auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + // Issue the Mainloop loads + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // LOCK mainloop_load2xform_pipeline_state for _writing_ + pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag); + int write_stage = load2xform_pipeline_state.index(); + + using BarrierType = typename Load2TransformPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state); + + // Advance mainloop_pipe + ++load2xform_pipeline_state; + skip_wait = (k_tile_count <= 1); + pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + copy(observed_tma_load_a_->with(*tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage)); + copy(observed_tma_load_b_->with(*tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage)); + ++k_tile_iter; + } + return cute::make_tuple(load2xform_pipeline_state, k_tile_iter); + } + + /// Set up the data needed by this collective for load. + /// Returned tuple must contain at least two elements, with the first two elements being: + /// gA_mkl - The tiled tensor for input A + /// gB_nkl - The tiled tensor for input B + // Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks + template + CUTLASS_DEVICE auto + load_init( + ProblemShape_MNKL const& problem_shape_MNKL, + Params const& params, + TensorStorage& shared_storage) const { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + + Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + + // Define the CTA-in-cluster Layout and Coord + Layout cta_layout_mnk = make_layout(cluster_shape_); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + // Project the cta_layout for tma_a along the n-modes + auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl)); + + // Project the cta_layout for tma_b along the m-modes + auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl)); + + // TMA Multicast Masks + uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + return cute::make_tuple( + gA_mkl, gB_nkl, // for scheduler + tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values + mcast_mask_a, mcast_mask_b); // multicast masks + } + + template< + class KTileIterator, class Accumulator, + class GTensorA, class DstCopyA, class SrcTensorA, class DstTensorA, + class GTensorB, class SrcTensorB, class DstTensorB + > + CUTLASS_DEVICE auto + transform( + Load2TransformPipeline load2transform_pipeline, + Load2TransformPipelineState load2transform_pipeline_consumer_state, + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_producer_state, + Accumulator accumulators, + cute::tuple input_operands, + KTileIterator k_tile_iter, int k_tile_count) { + + static_assert(cute::is_same_v, "ElementA and ElementB types should be the same."); + static_assert(cute::is_same_v, "ElementAMma and ElementBMma types should be the same."); + + cutlass::arch::NamedBarrier transform_bar(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier); + + // tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tAdA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs(Emulation), SmemStages (In SMEM or TMEM) + // tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, NumComputeMtxs(Complex,Emulation), SmemStages (In SMEM) + auto [unused_tAgA, dst_copy_A, tAsA, tAdACompute, + unused_tBgB, tBsB, tBsBCompute] = input_operands; + + // Create the tensors in registers + auto tArA = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArA_temp = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArACompute = make_tensor(tAsA(_,_,_,_,0).shape()); + + auto tBrB = make_tensor(tBsB(_,_,_,_,0).shape()); + auto tBrB_temp = make_tensor(tBsB(_,_,_,_,0).shape()); + auto tBrBCompute = make_tensor(tBsB(_,_,_,_,0).shape()); + + // For compute, cast to 4 raw elements instead of 2 complex elements. + auto tArA_x4 = recast>(tArA); + auto tArA_temp_x4 = recast>(tArA_temp); + auto tArACompute_x4 = recast>(tArACompute); + + auto tBrB_x4 = recast>(tBrB); + auto tBrB_temp_x4 = recast>(tBrB_temp); + auto tBrBCompute_x4 = recast>(tBrBCompute); + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag); + transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag); + + int load2transform_consumer_index = load2transform_pipeline_consumer_state.index(); + int transform2mma_producer_index = transform2mma_pipeline_producer_state.index(); + + auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state; + auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state; + + // Copy the input B matrix from SMEM + copy(AutoVectorizingCopy{}, tBsB(_,_,_,_,load2transform_consumer_index), tBrB); + // Copy the input A matrix from SMEM + copy(AutoVectorizingCopy{}, tAsA(_,_,_,_,load2transform_consumer_index), tArA); + + /// NOTE: sm100_mma_warpspecialized_interleaved_complex_tf32.hpp introduced about expanding: + /// re(a_complex * b_complex) -> (a_re, a_im) . (b_re,-b_im) = a . b_conj + /// im(a_complex * b_complex) -> (a_re, a_im) . (b_im, b_re) = a . b_swap + /// However, 16b types need to be packed for swapping and negation. + /// Hence, (re | im | re | im) is reordered into (re_x2 | im_x2). + cute::transform(tBrB_x4, tBrB_x4, [&] (auto& f4) -> Array {return {f4[0], f4[2], f4[1], f4[3]};}); + // Conversion b -> b_conj goes first, hence TransformB has a not preceding it. + if constexpr (not cute::is_same_v) { + cute::transform(tBrB_x4, tBrB_x4, [&] (auto& f4) { + auto f2_x2 = *reinterpret_cast,2>*>(&f4); + f2_x2[1] = cutlass::negate>{}(f2_x2[1]); + return *reinterpret_cast*>(&f2_x2); + }); + } + CUTE_UNROLL + for (int comp_mtx_index = 0; comp_mtx_index < NumComputeMtxs; ++comp_mtx_index) { + // Convert from fp32 -> bf16 + cute::transform(tBrB_x4, tBrBCompute_x4, + cutlass::NumericArrayConverter::convert); + // Store as B_conj (for producing C_re) + copy(AutoVectorizingCopy{}, tBrBCompute, tBsBCompute(_,_,_,_,0,comp_mtx_index,transform2mma_producer_index)); + // if it is not the last compute matrix, scale and substract + if (comp_mtx_index < NumComputeMtxs - 1) { + // Convert from bf16 -> fp32 to substract + cute::transform(tBrBCompute_x4, tBrB_temp_x4, + cutlass::NumericArrayConverter::convert); + cute::transform(tBrB_x4, tBrB_temp_x4, tBrB_x4, cutlass::minus>{}); + if constexpr (DispatchPolicy::ScalingFactor != 0) { + cute::transform(tBrB_x4, tBrB_x4, cutlass::scale>{(1 << DispatchPolicy::ScalingFactor)}); + } + } + + // Convert B_conj to B_swap + cute::transform(tBrBCompute_x4, tBrBCompute_x4, [&] (auto& h4) { + // Reinterpret as packed types + auto h2_x2_conj = *reinterpret_cast,2>*>(&h4); + cutlass::negate> neg; + Array,2> h2_x2_swap{ neg(h2_x2_conj[1]), h2_x2_conj[0] }; + return *reinterpret_cast*>(&h2_x2_swap); + }); + // Store as B_swap (for producing C_im) + copy(AutoVectorizingCopy{}, tBrBCompute, tBsBCompute(_,_,_,_,1,comp_mtx_index,transform2mma_producer_index)); + } + + // Loads from SMEM are done. Signal the mainloop load as early as possible + transform_bar.sync(); + load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state); + + // ( re | im | re | im ) -> ( re_x2 | im_x2 ) + cute::transform(tArA_x4, tArA_x4, [&] (auto& f4) -> Array{return {f4[0], f4[2], f4[1], f4[3]};}); + if constexpr (cute::is_same_v) { + cute::transform(tArA_x4, tArA_x4, [&] (auto& f4) { + auto f2_x2 = *reinterpret_cast,2>*>(&f4); + f2_x2[1] = cutlass::negate>{}(f2_x2[1]); + return *reinterpret_cast*>(&f2_x2); + }); + } + CUTE_UNROLL + for (int comp_mtx_index = 0; comp_mtx_index < NumComputeMtxs; ++comp_mtx_index) { + // Convert from fp32 -> bf16 + cute::transform(tArA_x4, tArACompute_x4, + cutlass::NumericArrayConverter::convert); + copy(dst_copy_A, tArACompute, tAdACompute(_,_,_,_,comp_mtx_index,transform2mma_producer_index)); + // if it is not the last compute matrix, scale and substract + if (comp_mtx_index < NumComputeMtxs - 1) { + // Convert from bf16 -> fp32 to substract + cute::transform(tArACompute_x4, tArA_temp_x4, + cutlass::NumericArrayConverter::convert); + cute::transform(tArA_x4, tArA_temp_x4, tArA_x4, cutlass::minus>{}); + if constexpr (DispatchPolicy::ScalingFactor != 0) { + cute::transform(tArA_x4, tArA_x4, cutlass::scale>{(1 << DispatchPolicy::ScalingFactor)}); + } + } + } + + // fence for SMEM writes + cutlass::arch::fence_view_async_shared(); + if constexpr (is_tmem::value) { + // fence for TMEM writes if A operand is coming from TMEM + cutlass::arch::fence_view_async_tmem_store(); + } + + // Let the MMA know we are done transforming + transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state); + // Next pipeline stage + ++load2transform_pipeline_consumer_state; + ++transform2mma_pipeline_producer_state; + + skip_wait = (k_tile_count <= 1); + // Peek the next pipeline stage's barriers + load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + } + return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + transform_init( + Params const& params, + ProblemShape_MNKL const& problem_shape_MNKL, + Accumulator accumulators, + TensorStorage& shared_storage) { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); + Tensor sA = as_position_independent_swizzle_tensor(sA_orig); + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + + Tensor sB_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); + Tensor sB = as_position_independent_swizzle_tensor(sB_orig); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + + // Map input, compute, and fragment tensors to + // Copy strategies and partitioned tensors. These will become the input + // operands of the transform function. Depending on MMA atom type, the + // operands can reside in SMEM or TMEM + auto setup_copy_ops = [&] ( + auto tensor_input, + auto input_copy_atom, + auto tensor_compute, + auto make_fragment, + auto compute_copy_atom) constexpr { + auto fragment_compute = make_fragment(tensor_compute); + if constexpr (cute::is_tmem>::value) { + // For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation. + // Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation + // See: TmemAllocMode::Duplicated. + Tensor tensor_input2x = [&] () constexpr { + if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) { + return make_tensor(tensor_input.data(), + logical_product(tensor_input.layout(), + make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE) + } + else { + return tensor_input; + } + }(); + + fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + // If operand comes from TMEM, create the TMEM_STORE based copy + auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,_,0,0)); + auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x); + auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute); + return cute::make_tuple(reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute); + } + else { + // If the operand comes from SMEM, create SMEM copy. + auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute); + auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout, Stride< _4,_1>>{}, + take<0,3>(tensor_compute.layout())); + + // Source copy is based on the source operand of copy. + auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input); + auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw); + + return cute::make_tuple(AutoVectorizingCopy{}, partitioned_tensor_input, partitioned_tensor_compute); + } + }; + + auto [dst_copy_A, tAsA, tAsACompute] = + setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{}); + + auto [dst_copy_B, tBsB, tBsBCompute] = + setup_copy_ops(sB, InputCopyAtomB{}, sBCompute, [&](auto &arg) {return TiledMma::make_fragment_B(arg);}, ComputeCopyAtomB{}); + return cute::make_tuple(gA_mkl, dst_copy_A, tAsA, tAsACompute, + gB_nkl, tBsB, tBsBCompute); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class FrgEngine, class FrgLayout, + class TensorA, class TensorB + > + CUTLASS_DEVICE auto + mma( + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_consumer_state, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_producer_state, + cute::Tensor const& accumulators, + cute::tuple const& input_operands, + int k_tile_count + ) { + TiledMma tiled_mma; + + auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + uint32_t skip_wait = (k_tile_count <= 0); + auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + ++next_transform2mma_pipeline_consumer_state; + + // tCrA : (MMA), MMA_M, MMA_K, NumComputeMtxs, SmemStage (In SMEM or TMEM) + // We use SMEM stages to match #buffers in Load <-> Convert + // tCrB : (MMA), MMA_N, MMA_K, NumComputeMtxs, SmemStages (In SMEM) + auto const [tCrA, tCrB] = input_operands; + + using ZeroScaler = cute::integral_constant; + using Scaler = cute::integral_constant; + + int remaining_accum_promotions = k_tile_count * StagesPerTile; + uint32_t mma2accum_skip_wait = (remaining_accum_promotions <= 0); + auto mma2accum_flag = mma2accum_pipeline.producer_try_acquire(mma2accum_pipeline_producer_state, mma2accum_skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag); + + int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index(); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); k_block += DispatchPolicy::AccPromotionInterval, --remaining_accum_promotions) { + // Accum stages are organized as (C_real | C_imag | C_real | C_imag | ...) + CUTLASS_PRAGMA_UNROLL + for (int re_im = 0; re_im < 2; ++re_im) { + mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state, mma2accum_flag); + + int mma2accum_pipeline_producer_state_index = mma2accum_pipeline_producer_state.index(); + auto tCtC = accumulators(_,_,_,mma2accum_pipeline_producer_state_index); + auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state; + + ++mma2accum_pipeline_producer_state; + mma2accum_skip_wait = (remaining_accum_promotions <= 1) && (re_im == 1); + mma2accum_flag = mma2accum_pipeline.producer_try_acquire(mma2accum_pipeline_producer_state, mma2accum_skip_wait); + + auto tCrA0 = tCrA(_,_,_,0,transform2mma_pipeline_consumer_state_index); + auto tCrA1 = tCrA(_,_,_,1,transform2mma_pipeline_consumer_state_index); + auto tCrA2 = tCrA(_,_,_,2,transform2mma_pipeline_consumer_state_index); + + auto tCrB0 = tCrB(_,_,_,re_im,0,transform2mma_pipeline_consumer_state_index); + auto tCrB1 = tCrB(_,_,_,re_im,1,transform2mma_pipeline_consumer_state_index); + auto tCrB2 = tCrB(_,_,_,re_im,2,transform2mma_pipeline_consumer_state_index); + + // MMA instructions Emulation + auto accumulate = UMMA::ScaleOut::Zero; + // First set of GEMMs that we need to perform for each band are unrolled to set compile-time constant + // scaling parameter. Scaled GEMM operations are only needed for the first MMA operation of each band. + + // Band 5 + if constexpr (NumBandsToCompute == 5) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block), tCrB2(_,_,k_block), tCtC); // A[2]*B[2] + accumulate = UMMA::ScaleOut::One; + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block+s), tCrB2(_,_,k_block+s), tCtC); // A[2]*B[2] + } + } + // Band 4 + if constexpr (NumBandsToCompute >= 4) { + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA1(_,_,k_block), tCrB2(_,_,k_block), tCtC); // A[1]*B[2] + accumulate = UMMA::ScaleOut::One; + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block), tCrB1(_,_,k_block), tCtC); // A[2]*B[1] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block+s), tCrB2(_,_,k_block+s), tCtC); // A[1]*B[2] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block+s), tCrB1(_,_,k_block+s), tCtC); // A[2]*B[1] + } + } + // Band 3 + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA0(_,_,k_block), tCrB2(_,_,k_block), tCtC); // A[2]*B[0] + accumulate = UMMA::ScaleOut::One; + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block), tCrB1(_,_,k_block), tCtC); // A[1]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[0]*B[2] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA0(_,_,k_block+s), tCrB2(_,_,k_block+s), tCtC); // A[2]*B[0] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block+s), tCrB1(_,_,k_block+s), tCtC); // A[1]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA2(_,_,k_block+s), tCrB0(_,_,k_block+s), tCtC); // A[0]*B[2] + } + // Band 2 + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA0(_,_,k_block), tCrB1(_,_,k_block), tCtC); // A[0]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[1]*B[0] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA0(_,_,k_block+s), tCrB1(_,_,k_block+s), tCtC); // A[0]*B[1] + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA1(_,_,k_block+s), tCrB0(_,_,k_block+s), tCtC); // A[1]*B[0] + } + // Band 1 + cute::gemm(tiled_mma.with(accumulate, Scaler{}), tCrA0(_,_,k_block), tCrB0(_,_,k_block), tCtC); // A[0]*B[0] + CUTLASS_PRAGMA_UNROLL + for (int s = 1; s < DispatchPolicy::AccPromotionInterval; s++) { + cute::gemm(tiled_mma.with(accumulate, ZeroScaler{}), tCrA0(_,_,k_block+s), tCrB0(_,_,k_block+s), tCtC); // A[0]*B[0] + } + mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state); + } + } + + transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state); + + skip_wait = (k_tile_count <= 1); + transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + + curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state; + ++next_transform2mma_pipeline_consumer_state; + } + return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + mma_init(cute::Tensor const& accumulators, TensorStorage& shared_storage) const { + TiledMma tiled_mma; + + Tensor tCrA = [&] () constexpr { + if constexpr (cute::is_base_of::value) { + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + return tiled_mma.make_fragment_A(sACompute); + } + else { + auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{})); + tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + return tCrA; + } + } (); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + Tensor tCrB = tiled_mma.make_fragment_B(sBCompute); + return cute::make_tuple(tCrA, tCrB); + } + + template + CUTLASS_DEVICE auto + accum_init(cute::Tensor const& accumulators, TmemCopyAtom tmem_cp_atom, EpilogueTile epilogue_tile) { + // Obtain a single accumulator + Tensor tAcc = tensor<0>(accumulators(_,_,_,_0{})); + // Apply epilogue subtiling + Tensor tAcc_epi = flat_divide(tAcc, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + // Create the TMEM copy for single EpilogueTile. + // Note that EpilogueTile = CtaTile for NoSmem epilogue + auto tiled_t2r = make_tmem_copy(tmem_cp_atom, tAcc_epi(_,_,_0{},_0{})); + auto thread_t2r = tiled_t2r.get_slice(threadIdx.x % size(tiled_t2r)); + Tensor tTR_gC = thread_t2r.partition_D(tAcc_epi); + Tensor tTR_rAcc = make_tensor(shape(tTR_gC)); // (T2R,T2R_M,T2R_N) + Tensor tTR_rGlobAcc = make_tensor(shape(tTR_gC)); // (T2R,T2R_M,T2R_N) + + // Apply epilogue subtiling to bulk accumulator + // We need to tile the whole bulk_tmem allocation with EpilogueTile. + // The accumulation should be aware of the AccumulatorPipelineStages + Tensor tBulkAcc_epi = flat_divide(accumulators(make_coord(_,_),_0{},_0{},_), EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N,PIPE) + Tensor tTR_tBulkAcc = thread_t2r.partition_S(tBulkAcc_epi); // (T2R,T2R_M,T2R_N,EPI_M,EPI_N,PIPE) + return cute::make_tuple(tiled_t2r, thread_t2r, tTR_tBulkAcc, tTR_rAcc, tTR_rGlobAcc); + } + + template + CUTLASS_DEVICE auto + accum(cute::tuple accum_inputs, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_consumer_state, + int k_tile_count) { + auto [tiled_t2r, thread_t2r, tTR_tBulkAcc, + tTR_rAcc, tTR_rGlobAcc] = accum_inputs; + + Tensor tTR_rAcc_float2 = recast>(tTR_rAcc); // (T2R/2,T2R_M,T2R_N) + Tensor tTR_rGlobAcc_float2_x2 = recast,2>>(tTR_rGlobAcc);// (T2R/2,T2R_M,T2R_N) + + // Clear the global accumulator + CUTE_UNROLL + for (int i = 0; i 0; --k_tile_count) { + // The stage is limited to a CTA tile + CUTLASS_PRAGMA_NO_UNROLL + for (int k_block = 0; k_block cute::remove_cvref_t {return {cutlass::plus>{}(f2_x2[0], r2), f2_x2[1]};}); + } + else { + cute::transform(tTR_rGlobAcc_float2_x2, tTR_rAcc_float2, tTR_rGlobAcc_float2_x2, + [&] (auto& f2_x2, auto& i2) -> cute::remove_cvref_t {return {f2_x2[0], cutlass::plus>{}(f2_x2[1], i2)};}); + } + + cutlass::arch::fence_view_async_tmem_load(); // Need a fence bw TMEM_LOAD and arrive + mma2accum_pipeline.consumer_release(mma2accum_pipeline_consumer_state); + + ++mma2accum_pipeline_consumer_state; + skip_wait = ((k_tile_count <= 1) && (k_block >= (StagesPerTile-1))) && (re_im == 1); + mma2accum_flag = mma2accum_pipeline.consumer_try_wait(mma2accum_pipeline_consumer_state, skip_wait); + } + } + } + + // Interleave back (real_x2 | imag_x2) to (real | imag | real | imag) + cute::transform(tTR_rGlobAcc_float2_x2, tTR_rGlobAcc_float2_x2, [&] (auto& f2_x2) -> cute::remove_cvref_t { + Array c0{f2_x2[0][0], f2_x2[1][0]}; + Array c1{f2_x2[0][1], f2_x2[1][1]}; + return {c0, c1}; + }); + + return cute::make_tuple(mma2accum_pipeline_consumer_state, tTR_rGlobAcc); + } + +protected: + + template + CUTLASS_DEVICE + constexpr auto + tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const { + using X = cute::Underscore; + // Separate out problem shape for convenience + auto [M,N,K,L] = problem_shape_MNKL; + + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,L)); + Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,L)); + + // Tile the tensors and defer the slice + Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); + Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); + + return cute::make_tuple(gA_mkl, gB_nkl); + } + + typename Params::TMA_A const* observed_tma_load_a_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_ = nullptr; + + ClusterShape cluster_shape_; + uint32_t block_rank_in_cluster_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp new file mode 100644 index 00000000..1dc7b02c --- /dev/null +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_interleaved_complex_tf32.hpp @@ -0,0 +1,875 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + +#include "cutlass/cutlass.h" +#include "cutlass/detail/cluster.hpp" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/numeric_conversion.h" +#include "cutlass/detail/sm100_tmem_helper.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/copy_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/arch/mma_sm100.hpp" +#include "cutlass/trace.h" +#include "cutlass/kernel_hardware_info.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +namespace detail { +template +struct Sm100CollectiveMmaComplexLayoutAtomType { + using InputLayoutAtom = InputLayoutAtom_; + using ComputeLayoutAtom = ComputeLayoutAtom_; +}; + +template +struct Sm100CollectiveMmaComplexCopyType { + using InputCopyAtom = InputCopyAtom_; + using ComputeCopyAtom = ComputeCopyAtom_; +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop for complex kernels +template < + int ComputationPipelineStageCount_, + int SchedulerPipelineStageCount_, + int AccumulatorPipelineStageCount_, + int TransformationPipelineStageCount_, + class AccumulatorCopyAtom_, + class ClusterShape, // Static cluster shape or dynamic (int, int, _1) + class TileShape_, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + class StrideA_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomsA_, + class CopyAtomsA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomsB_, + class CopyAtomsB_, + class TransformB_> +struct CollectiveMma< + MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32< + ComputationPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + TransformationPipelineStageCount_, + ClusterShape, + AccumulatorCopyAtom_>, + TileShape_, + complex, + StrideA_, + complex, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomsA_, + CopyAtomsA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomsB_, + CopyAtomsB_, + TransformB_> +{ + // + // Type Aliases + // + using TileShape = TileShape_; + using TiledMma = TiledMma_; + + // ElementA and ElementB are cutlass::complex, which are used as GMEM input and output data type. + using ElementA = complex; + using StrideA = StrideA_; + using ElementB = complex; + using StrideB = StrideB_; + +private: + // ElementAMma and ElementBMma are cutlass::complex, which are used as SMEM and RF data type. + // ElementAMmaRaw and ElementBMmaRaw are cutlass::tfloat32_t, which is the real internal data type set in TMA descriptor and used in TCGEN05 calculation. + using ElementAMma = typename TiledMma::ValTypeA; // complex + using ElementAMmaRaw = typename ElementAMma::value_type; // tfloat32_t + using ElementBMma = typename TiledMma::ValTypeB; // complex + using ElementBMmaRaw = typename ElementBMma::value_type; // tfloat32_t + +public: + // For a complex kernel, the MMA output type is real valued, but ElementAccumulator is a complex type for the GETT reference kernel + using ElementAccumulator = cutlass::complex; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomsA = SmemLayoutAtomsA_; + using SmemLayoutAtomsB = SmemLayoutAtomsB_; + using CopyAtomsA = CopyAtomsA_; + using CopyAtomsB = CopyAtomsB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + + // Determine MMA type: MMA_1SM vs MMA_2SM + using AtomThrShapeMNK = Shape(typename TiledMma_::ThrLayoutVMNK{})), _1, _1>; + using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32< + ComputationPipelineStageCount_, + SchedulerPipelineStageCount_, + AccumulatorPipelineStageCount_, + TransformationPipelineStageCount_, + ClusterShape, + AccumulatorCopyAtom_>; + static constexpr bool IsDynamicCluster = not cute::is_static_v; + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + + // Define A and B block shapes for reduced size TMA_LOADs + using CtaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using CtaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + using ArchTag = typename DispatchPolicy::ArchTag; + + using Load2TransformPipeline = cutlass::PipelineTmaTransformAsync< + DispatchPolicy::ComputationPipelineStageCount, + AtomThrShapeMNK>; + using Load2TransformPipelineState = typename Load2TransformPipeline::PipelineState; + + using Transform2MmaPipeline = cutlass::PipelineUmmaConsumerAsync< + DispatchPolicy::TransformationPipelineStageCount, + AtomThrShapeMNK>; + using Transform2MmaPipelineState = typename Transform2MmaPipeline::PipelineState; + + using Mma2AccumPipeline = cutlass::PipelineUmmaAsync< + DispatchPolicy::Schedule::AccumulatorPipelineStageCount, + AtomThrShapeMNK>; + using Mma2AccumPipelineState = typename Mma2AccumPipeline::PipelineState; + + // Thread Counts + static constexpr uint32_t NumTransformationThreads = 128; + static constexpr uint32_t NumAccumThreads = 128; + + // Get the Algorithm parameters + constexpr static int NumComputeMtxs = 2; + constexpr static int AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount; + constexpr static int StagesPerTile = size<2>(CtaShapeA_MK{}); + + // Copy atom for Accumulator + using AccumulatorCopyAtom = typename DispatchPolicy::AccumulatorCopyAtom; + + using SmemLayoutAtomA = typename SmemLayoutAtomsA::InputLayoutAtom; + using SmemLayoutAtomACompute = typename SmemLayoutAtomsA::ComputeLayoutAtom; + using SmemLayoutAtomB = typename SmemLayoutAtomsB::InputLayoutAtom; + using SmemLayoutAtomBCompute = typename SmemLayoutAtomsB::ComputeLayoutAtom; + + using InputCopyAtomA = typename CopyAtomsA::InputCopyAtom; + using ComputeCopyAtomA = typename CopyAtomsA::ComputeCopyAtom; + using InputCopyAtomB = typename CopyAtomsB::InputCopyAtom; + using ComputeCopyAtomB = typename CopyAtomsB::ComputeCopyAtom; + + static_assert(((size<0,0>(CtaShapeA_MK{}) * size<1>(CtaShapeA_MK{})) % size<0>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeA_MK{}) * size<2>(CtaShapeA_MK{})) % size<1>(SmemLayoutAtomACompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(((size<0,0>(CtaShapeB_NK{}) * size<1>(CtaShapeB_NK{})) % size<0>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(CtaShapeB_NK{}) * size<2>(CtaShapeB_NK{})) % size<1>(SmemLayoutAtomBCompute{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(CtaShapeA_MK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutACompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomACompute{}, + append(append(CtaShapeA_MK{}, Int{}), Int{}))); + + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(CtaShapeB_NK{}, Int{}), + (cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}))); + + using SmemLayoutBCompute = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomBCompute{}, + append(CtaShapeB_NK{}, Int{}))); + + static_assert(DispatchPolicy::ComputationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more."); + static_assert(DispatchPolicy::TransformationPipelineStageCount >= 2, "Specialization requires Stages set to value 2 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must have A operand from TMEM and B operand from SMEM for this mainloop."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyA - invalid TMA copy atom specified."); + static_assert((cute::is_same_v || cute::is_same_v), + "GmemTiledCopyB - invalid TMA copy atom specified."); + + struct PipelineStorage { + using Load2TransformPipelineStorage = typename Load2TransformPipeline::SharedStorage; + alignas(16) Load2TransformPipelineStorage load2transform_pipeline; + using Transform2MmaPipelineStorage = typename Transform2MmaPipeline::SharedStorage; + alignas(16) Transform2MmaPipelineStorage transform2mma_pipeline; + using Mma2AccumPipelineStorage = typename Mma2AccumPipeline::SharedStorage; + alignas(16) Mma2AccumPipelineStorage mma2accum_pipeline; + }; + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + struct TensorStorageUntransformed { + cute::ArrayEngine> smem_A; + cute::ArrayEngine> smem_B; + } input; + + union TensorStorageTransformed { + alignas(1024) cute::ArrayEngine smem_ACompute; // smem_ACompute is actually in tmem + alignas(1024) cute::ArrayEngine> smem_BCompute; + } compute; + } tensors; + + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + + // Different from other GEMM kernels, both CTAs should be aware of loads. Both CTAs will work on + // loaded input A and B matrices to convert the data type + static constexpr uint32_t TmaTransactionBytes = + (size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * size<2>(SmemLayoutA{}) * static_cast(sizeof(ElementAMma))) + + (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * size<2>(SmemLayoutB{}) * static_cast(sizeof(ElementBMma))); + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A{nullptr}; + StrideA dA{}; + ElementB const* ptr_B{nullptr}; + StrideB dB{}; + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), + make_tile(typename TiledMma::AtomThrID{}))); + + using TMA_A = decltype(make_tma_atom_A_sm100( + GmemTiledCopyA{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + TMA_A tma_load_a; + TMA_B tma_load_b; + TMA_A tma_load_a_fallback; + TMA_B tma_load_b_fallback; + dim3 cluster_shape_fallback; + }; + + CUTLASS_DEVICE + CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster) + : cluster_shape_(cluster_shape) + , block_rank_in_cluster_(block_rank_in_cluster) { + if constexpr (IsDynamicCluster) { + const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x && + cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y); + observed_tma_load_a_ = is_fallback_cluster ? ¶ms.tma_load_a_fallback : ¶ms.tma_load_a; + observed_tma_load_b_ = is_fallback_cluster ? ¶ms.tma_load_b_fallback : ¶ms.tma_load_b; + } + else { + observed_tma_load_a_ = ¶ms.tma_load_a; + observed_tma_load_b_ = ¶ms.tma_load_b; + } + } + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) { + (void) workspace; + + // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA)); + Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB)); + + auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape); + // Cluster layout for TMA construction + auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{})); + + auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape_fallback); + // Cluster layout for TMA construction + auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_A tma_load_a = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + return { + tma_load_a, + tma_load_b, + tma_load_a_fallback, + tma_load_b_fallback, + hw_info.cluster_shape_fallback + }; + } + + template + static bool + can_implement( + ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + constexpr int tma_alignment_bits = 128; + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + bool implementable = true; + constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,K,L), StrideA{}); + constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + return implementable; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE static void + prefetch_tma_descriptors(Params const& params) { + if constexpr (IsDynamicCluster) { + dim3 cs = cute::cluster_shape(); + const bool is_fallback_cluster = (cs.x == params.cluster_shape_fallback.x && cs.y == params.cluster_shape_fallback.y); + if (is_fallback_cluster) { + cute::prefetch_tma_descriptor(params.tma_load_a_fallback.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b_fallback.get_tma_descriptor()); + } + else { + cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + } + } + else { + cute::prefetch_tma_descriptor(params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_b.get_tma_descriptor()); + } + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE auto + partition_accumulator_shape() { + return append( + partition_shape_C(TiledMma{}, take<0,2>(TileShape{})), + Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2) + } + + /// Produce the inputs to the transform threads by loading inputs from gmem -> smem + template < + class GTensorA, class GTensorB, + class GTensorPartitionedA, class GTensorPartitionedB, + class STensorA, class STensorB, + class TileCoordMNKL, + class KTileIterator + > + CUTLASS_DEVICE cute::tuple + load( + Params const& params, + Load2TransformPipeline pipeline, + Load2TransformPipelineState load2xform_pipeline_state, + cute::tuple const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, + KTileIterator k_tile_iter, int k_tile_count) { + + auto [unused_gA, unused_gB, + tAgA_mkl, tBgB_nkl, tAsA, tBsB, + mcast_mask_a, mcast_mask_b] = load_inputs; + + // slice out the work coord from tiled tensors + Tensor tAgA = tAgA_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + Tensor tBgB = tBgB_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + uint32_t skip_wait = (k_tile_count <= 0); + auto pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + // Issue the Mainloop loads + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // LOCK mainloop_load2xform_pipeline_state for _writing_ + pipeline.producer_acquire(load2xform_pipeline_state, pipeline_flag); + int write_stage = load2xform_pipeline_state.index(); + + using BarrierType = typename Load2TransformPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(load2xform_pipeline_state); + + // Advance mainloop_pipe + ++load2xform_pipeline_state; + skip_wait = (k_tile_count <= 1); + pipeline_flag = pipeline.producer_try_acquire(load2xform_pipeline_state, skip_wait); + + copy(observed_tma_load_a_->with(*tma_barrier, mcast_mask_a), tAgA(_,*k_tile_iter), tAsA(_,write_stage)); + copy(observed_tma_load_b_->with(*tma_barrier, mcast_mask_b), tBgB(_,*k_tile_iter), tBsB(_,write_stage)); + ++k_tile_iter; + } + return cute::make_tuple(load2xform_pipeline_state, k_tile_iter); + } + + + /// Set up the data needed by this collective for load. + /// Returned tuple must contain at least two elements, with the first two elements being: + /// gA_mkl - The tiled tensor for input A + /// gB_nkl - The tiled tensor for input B + // Other inputs needed for load(): partitioned AB tensors for gmem and smem, and mcast masks + template + CUTLASS_DEVICE auto + load_init( + ProblemShape_MNKL const& problem_shape_MNKL, + Params const& params, + TensorStorage& shared_storage) const { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + Tensor tCgA_mkl = cta_mma.partition_A(gA_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + Tensor tCgB_nkl = cta_mma.partition_B(gB_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + + Tensor sA = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + + // Define the CTA-in-cluster Layout and Coord + Layout cta_layout_mnk = make_layout(cluster_shape_); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + // Project the cta_layout for tma_a along the n-modes + auto [tAgA_mkl, tAsA] = tma_partition(*observed_tma_load_a_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA), group_modes<0,3>(tCgA_mkl)); + + // Project the cta_layout for tma_b along the m-modes + auto [tBgB_nkl, tBsB] = tma_partition(*observed_tma_load_b_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB), group_modes<0,3>(tCgB_nkl)); + + // TMA Multicast Masks + uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + return cute::make_tuple( + gA_mkl, gB_nkl, // for scheduler + tAgA_mkl, tBgB_nkl, tAsA, tBsB, // for input tensor values + mcast_mask_a, mcast_mask_b); // multicast masks + } + + template< + class KTileIterator, class Accumulator, + class GTensorA, class SrcCopyA, class DstCopyA, class SrcTensorA, class DstTensorA, + class GTensorB, class SrcCopyB, class DstCopyB, class SrcTensorB, class DstTensorB + > + CUTLASS_DEVICE auto + transform( + Load2TransformPipeline load2transform_pipeline, + Load2TransformPipelineState load2transform_pipeline_consumer_state, + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_producer_state, + Accumulator accumulators, + cute::tuple input_operands, + KTileIterator k_tile_iter, int k_tile_count) { + cutlass::arch::NamedBarrier transform_barrier(NumTransformationThreads, cutlass::arch::ReservedNamedBarriers::TransformBarrier); + + // tAsA : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tAtACompute : (Copy,#Copy),MMA_Rest,MMA_M_Rest,MMA_K_Rest, NumComputeMtxs, SmemStages (In TMEM) + // tBsB : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM) + // tBsBCompute : (Copy,#Copy),MMA_Rest,MMA_N_Rest,MMA_K_Rest, SmemStages (In SMEM) + auto [unused_tAgA, src_copy_A, dst_copy_A, tAsA, tAtACompute, + unused_tBgB, src_copy_B, dst_copy_B, tBsB, tBsBCompute] = input_operands; + + // Create the tensors in registers + auto tArA = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArA_conj = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tArA_swap = make_tensor(tAsA(_,_,_,_,0).shape()); + auto tBrB = make_tensor(tBsB(_,_,_,_,0).shape()); + + uint32_t skip_wait = (k_tile_count <= 0); + auto load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + auto transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + load2transform_pipeline.consumer_wait(load2transform_pipeline_consumer_state, load2transform_flag); + transform2mma_pipeline.producer_acquire(transform2mma_pipeline_producer_state, transform2mma_flag); + + int load2transform_consumer_index = load2transform_pipeline_consumer_state.index(); + int transform2mma_producer_index = transform2mma_pipeline_producer_state.index(); + + auto curr_load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state; + auto curr_transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state; + + // Copy the input A matrix from SMEM + copy(src_copy_A, tAsA(_,_,_,_,load2transform_consumer_index), tArA); + // Copy the input B matrix from SMEM + copy(src_copy_B, tBsB(_,_,_,_,load2transform_consumer_index), tBrB); + + // First MMA, A.real * B.real - A.imag * B.imag + // Compose [real, -imag] copy for A TMEM + // Reflect the conjugation of B through A + if constexpr (cute::is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tArA); i++) { + tArA_conj(i) = {tArA(i).real(), -tArA(i).imag()}; + } + } + else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tArA); i++) { + tArA_conj(i) = tArA(i); + } + } + // Write to TMEM + copy(dst_copy_A, tArA_conj, tAtACompute(_,_,_,_,0,transform2mma_producer_index)); + + // Second MMA, A.imag * B.real + A.real * B.imag + // Compose [imag, real] copy for A TMEM + // Reflect the conjugation of B through A + auto transform_element = [] (ElementAMma const& tArA_i) -> ElementAMma { + if constexpr (cute::is_same_v && cute::is_same_v) { // CC + return {-tArA_i.imag(), -tArA_i.real()}; + } + else if constexpr (cute::is_same_v && not cute::is_same_v) { // CN/CT + return {-tArA_i.imag(), tArA_i.real()}; + } + else if constexpr (not cute::is_same_v && cute::is_same_v) { // NC/TC + return {tArA_i.imag(), -tArA_i.real()}; + } + else { // TN/NT/NN/TT + return {tArA_i.imag(), tArA_i.real()}; + } + }; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tArA); i++) { + tArA_swap(i) = transform_element(tArA(i)); + } + + // Write to TMEM + copy(dst_copy_A, tArA_swap, tAtACompute(_,_,_,_,1,transform2mma_producer_index)); + + // Write the B matrix to SMEM without any changes + copy(dst_copy_B, tBrB, tBsBCompute(_,_,_,_,transform2mma_producer_index)); + + // Loads from SMEM are done. Signal the mainloop load as early as possible + transform_barrier.sync(); + load2transform_pipeline.consumer_release(curr_load2transform_pipeline_consumer_state); + + // fence for SMEM writes + cutlass::arch::fence_view_async_shared(); + if constexpr (is_tmem::value) { + // fence for TMEM writes if A operand is coming from TMEM + cutlass::arch::fence_view_async_tmem_store(); + } + + // Let the MMA know we are done transforming + transform2mma_pipeline.producer_commit(curr_transform2mma_pipeline_producer_state); + // Next pipeline stage + ++load2transform_pipeline_consumer_state; + ++transform2mma_pipeline_producer_state; + + skip_wait = (k_tile_count <= 1); + // Peek the next pipeline stage's barriers + load2transform_flag = load2transform_pipeline.consumer_try_wait(load2transform_pipeline_consumer_state, skip_wait); + transform2mma_flag = transform2mma_pipeline.producer_try_acquire(transform2mma_pipeline_producer_state, skip_wait); + } + return cute::make_tuple(load2transform_pipeline_consumer_state, transform2mma_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + transform_init( + Params const& params, + ProblemShape_MNKL const& problem_shape_MNKL, + Accumulator accumulators, + TensorStorage& shared_storage) { + auto [gA_mkl, gB_nkl] = tile_input_tensors(params, problem_shape_MNKL); + + Tensor sA_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_A.begin()), SmemLayoutA{}); + Tensor sA = as_position_independent_swizzle_tensor(sA_orig); + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + + Tensor sB_orig = make_tensor(make_smem_ptr(shared_storage.input.smem_B.begin()), SmemLayoutB{}); + Tensor sB = as_position_independent_swizzle_tensor(sB_orig); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + + // Map input, compute, and fragment tensors to + // Copy strategies and partitioned tensors. These will become the input + // operands of the transform function. Depending on MMA atom type, the + // operands can reside in SMEM or TMEM + auto setup_copy_ops = [&] (auto tensor_input, auto input_copy_atom, + auto tensor_compute, auto make_fragment, auto compute_copy_atom) constexpr { + auto fragment_compute = make_fragment(tensor_compute); + if constexpr (cute::is_tmem>::value) { + // For M=128 with 2CTA MMA atoms, the TMEM tensor for A has a duplicated allocation. + // Instead of allocation a 64x16 TMEM tensor, we have a 128x16 allocation + // See: TmemAllocMode::Duplicated. + Tensor tensor_input2x = [&] () constexpr { + if constexpr (decltype(size<0,0>(fragment_compute) == Int<128>{} && size<0,0>(tensor_input) == Int<64>{})::value) { + return make_tensor(tensor_input.data(), + logical_product(tensor_input.layout(), + make_tile(make_tile(Layout<_2,_0>{},_),_,_,_))); // ((128,16),m,k,PIPE) + } + else { + return tensor_input; + } + }(); + + fragment_compute.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + // If operand comes from TMEM, create the TMEM_STORE based copy + auto reg2tmem_tiled_copy = make_tmem_copy(compute_copy_atom, fragment_compute(_,_,_,0,0)); + auto thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input2x); + auto partitioned_tensor_compute = thr_reg2tmem_tiled_copy.partition_D(fragment_compute); + // Source copy is based on the source operand of TMEM_STORE copy. + auto smem2reg_tiled_copy = make_tiled_copy_S(Copy_Atom, ElementAMma>{}, reg2tmem_tiled_copy); + return cute::make_tuple(smem2reg_tiled_copy, reg2tmem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute); + } + else { + // If the operand comes from SMEM, create SMEM copy. + auto tensor_compute_ind_sw = as_position_independent_swizzle_tensor(tensor_compute); + auto reg2smem_tiled_copy = make_cotiled_copy(compute_copy_atom, Layout, Stride< _8,_1>>{}, + tensor_compute(_,_,_,0).layout()); + + // Source copy is based on the source operand of copy. + auto smem2reg_tiled_copy = make_tiled_copy_S(input_copy_atom, reg2smem_tiled_copy); + auto thr_smem2reg_tiled_copy = smem2reg_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(threadIdx.x % NumTransformationThreads); + auto partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(tensor_input); + auto partitioned_tensor_compute = thr_reg2smem_tiled_copy.partition_D(tensor_compute_ind_sw); + + return cute::make_tuple(smem2reg_tiled_copy, reg2smem_tiled_copy, partitioned_tensor_input, partitioned_tensor_compute); + } + }; + + auto [src_copy_A, dst_copy_A, tAsA, tAtACompute] = + setup_copy_ops(sA, InputCopyAtomA{}, sACompute, [&](auto &arg) {return TiledMma::make_fragment_A(arg);}, ComputeCopyAtomA{}); + + auto [src_copy_B, dst_copy_B, tBsB, tBsBCompute] = + setup_copy_ops(sB, InputCopyAtomB{}, sBCompute, [&](auto &arg) {return TiledMma::make_fragment_B(arg);}, ComputeCopyAtomB{}); + + return cute::make_tuple(gA_mkl, src_copy_A, dst_copy_A, tAsA, tAtACompute, + gB_nkl, src_copy_B, dst_copy_B, tBsB, tBsBCompute); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class FrgEngine, class FrgLayout, + class TensorA, class TensorB + > + CUTLASS_DEVICE auto + mma( + Transform2MmaPipeline transform2mma_pipeline, + Transform2MmaPipelineState transform2mma_pipeline_consumer_state, + Mma2AccumPipeline mma2accum_pipeline, + Mma2AccumPipelineState mma2accum_pipeline_producer_state, + cute::Tensor const& accumulators, + cute::tuple const& input_operands, + int k_tile_count + ) { + TiledMma tiled_mma; + + // tCrA : (MMA), MMA_M, MMA_K, NumComputeMtxs, SmemStage (In TMEM) + // We use SMEM stages to match #buffers in Load <-> Convert + // tCrB : (MMA), MMA_N, MMA_K, SmemStages (In SMEM) + auto const [tCrA, tCrB] = input_operands; + + auto curr_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + auto next_transform2mma_pipeline_consumer_state = transform2mma_pipeline_consumer_state; + uint32_t skip_wait = (k_tile_count <= 0); + auto transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + ++next_transform2mma_pipeline_consumer_state; + + mma2accum_pipeline.producer_acquire(mma2accum_pipeline_producer_state); + + constexpr int RealAccumIndex = 0; + constexpr int ImagAccumIndex = 1; + + int mma2accum_pipeline_producer_state_index = mma2accum_pipeline_producer_state.index(); + auto tCtC_real = accumulators(_,_,_,RealAccumIndex,mma2accum_pipeline_producer_state_index); + auto tCtC_imag = accumulators(_,_,_,ImagAccumIndex,mma2accum_pipeline_producer_state_index); + auto curr_mma2accum_pipeline_producer_state = mma2accum_pipeline_producer_state; + ++mma2accum_pipeline_producer_state; + + // + // PIPELINED MAIN LOOP + // + + tiled_mma.accumulate_ = UMMA::ScaleOut::Zero; + CUTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + + transform2mma_pipeline.consumer_wait(curr_transform2mma_pipeline_consumer_state, transform2mma_flag); + + int transform2mma_pipeline_consumer_state_index = curr_transform2mma_pipeline_consumer_state.index(); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < StagesPerTile; ++k_block) { + + auto tCrA_conj = tCrA(_,_,_,Int<0>{},transform2mma_pipeline_consumer_state_index); + auto tCrA_swap = tCrA(_,_,_,Int<1>{},transform2mma_pipeline_consumer_state_index); + + auto tCrB0 = tCrB(_,_,_,transform2mma_pipeline_consumer_state_index); + + // A conjugate * B + cute::gemm(tiled_mma, tCrA_conj(_,_,k_block), tCrB0(_,_,k_block), tCtC_real); // A[0]*B[0] + // A swapped * B + cute::gemm(tiled_mma, tCrA_swap(_,_,k_block), tCrB0(_,_,k_block), tCtC_imag); // A[0]*B[0] + tiled_mma.accumulate_ = UMMA::ScaleOut::One; + } + + transform2mma_pipeline.consumer_release(curr_transform2mma_pipeline_consumer_state); + + skip_wait = (k_tile_count <= 1); + transform2mma_flag = transform2mma_pipeline.consumer_try_wait(next_transform2mma_pipeline_consumer_state, skip_wait); + + curr_transform2mma_pipeline_consumer_state = next_transform2mma_pipeline_consumer_state; + ++next_transform2mma_pipeline_consumer_state; + } + + mma2accum_pipeline.producer_commit(curr_mma2accum_pipeline_producer_state); + + return cute::make_tuple(curr_transform2mma_pipeline_consumer_state, mma2accum_pipeline_producer_state); + } + + template + CUTLASS_DEVICE auto + mma_init(cute::Tensor const& accumulators, TensorStorage& shared_storage) const { + TiledMma tiled_mma; + + Tensor tCrA = [&] () constexpr { + if constexpr (cute::is_base_of::value) { + Tensor sACompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_ACompute.begin()), SmemLayoutACompute{}); + return tiled_mma.make_fragment_A(sACompute); + } + else { + auto tCrA = tiled_mma.make_fragment_A(shape(SmemLayoutACompute{})); + tCrA.data() = accumulators.data().get() + cutlass::detail::find_tmem_tensor_col_offset(accumulators); + return tCrA; + } + } (); + Tensor sBCompute = make_tensor(make_smem_ptr(shared_storage.compute.smem_BCompute.begin()), SmemLayoutBCompute{}); + Tensor tCrB = tiled_mma.make_fragment_B(sBCompute); + return cute::make_tuple(tCrA, tCrB); + } + + template + CUTLASS_DEVICE auto + accum_init(cute::Tensor const& accumulators, TmemCopyAtom, EpilogueTile) { + return accumulators; + } + +protected: + + template + CUTLASS_DEVICE + constexpr auto + tile_input_tensors(Params const& params, ProblemShape_MNKL const& problem_shape_MNKL) const { + using X = cute::Underscore; + // Separate out problem shape for convenience + auto [M,N,K,L] = problem_shape_MNKL; + + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = observed_tma_load_a_->get_tma_tensor(make_shape(M,K,L)); + Tensor mB_nkl = observed_tma_load_b_->get_tma_tensor(make_shape(N,K,L)); + + // Tile the tensors and defer the slice + Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l) + Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l) + + return cute::make_tuple(gA_mkl, gB_nkl); + } + + typename Params::TMA_A const* observed_tma_load_a_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_ = nullptr; + + ClusterShape cluster_shape_; + uint32_t block_rank_in_cluster_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp new file mode 100644 index 00000000..3d9c11e4 --- /dev/null +++ b/include/cutlass/gemm/collective/sm100_mma_warpspecialized_planar_complex.hpp @@ -0,0 +1,829 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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/detail/cluster.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/gemm/gemm.h" +#include "cutlass/trace.h" + +#include "cute/algorithm/functional.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" +#include "cutlass/kernel_hardware_info.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::gemm::collective { +using namespace cute; + +namespace detail { +template +struct Sm100CollectiveMmaPlanarComplexTiledMmaType { + using TiledMmaAPosAtom = TiledMmaAPos_; + using TiledMmaANegAtom = TiledMmaANeg_; +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop +// Both DMA Load and MMA methods of this class must be run by a single thread that's picked by elect_one +template < + int Stages, + int SchedulerPipelineStageCount, + int AccumulatorPipelineStageCount, + class ClusterShape, + class TileShape_, // Static cluster shape or dynamic (int, int, _1) + class ElementA_, // (MmaAtomShapeM, MmaAtomShapeN, TileK) + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMmaPair_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm100TmaUmmaWarpSpecializedPlanarComplex< + Stages, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + ClusterShape>, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMmaPair_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + + // Determine MMA type: MMA_1SM vs MMA_2SM + using TiledMmaPair = TiledMmaPair_; + using TiledMma = typename TiledMmaPair::TiledMmaAPosAtom; + using TiledMmaANeg = typename TiledMmaPair::TiledMmaANegAtom; + using AtomThrShapeMNK = Shape(typename TiledMma::ThrLayoutVMNK{})), _1, _1>; + + static constexpr bool IsDynamicCluster = not cute::is_static_v; + + using DispatchPolicy = MainloopSm100TmaUmmaWarpSpecializedPlanarComplex< + Stages, + SchedulerPipelineStageCount, + AccumulatorPipelineStageCount, + ClusterShape>; + using TileShape = TileShape_; + + CUTE_STATIC_ASSERT_V(evenly_divides(TileShape{}, tile_shape(TiledMma{})), + "Static cluster shape used: TileShape should be evenly divided by TiledMma"); + + using CtaShape_MNK = decltype(shape_div(TileShape{}, AtomThrShapeMNK{})); + + // Define A and B block shapes for reduced size TMA_LOADs + using MmaShapeA_MK = decltype(partition_shape_A(TiledMma{}, make_shape(size<0>(TileShape{}), size<2>(TileShape{})))); + using MmaShapeB_NK = decltype(partition_shape_B(TiledMma{}, make_shape(size<1>(TileShape{}), size<2>(TileShape{})))); + + using ElementA = ElementA_; + using ElementAMma = typename TiledMma::ValTypeA; + using StrideA = StrideA_; + using ElementB = ElementB_; + using ElementBMma = typename TiledMma::ValTypeB; + using StrideB = StrideB_; + 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; + + using MainloopPipeline = cutlass::PipelineTmaUmmaAsync< + DispatchPolicy::Stages, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtomA must be rank 2 (M, K)"); + static_assert(((size<0,0>(MmaShapeA_MK{}) * size<1>(MmaShapeA_MK{})) % size<0>(SmemLayoutAtomA{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(MmaShapeA_MK{}) * size<2>(MmaShapeA_MK{})) % size<1>(SmemLayoutAtomA{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(cute::is_void_v, + "SM100 UMMA cannot have a non-void copy atom for smem sourced instructions."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtomB must be rank 2 (N, K)"); + static_assert(((size<0,0>(MmaShapeB_NK{}) * size<1>(MmaShapeB_NK{})) % size<0>(SmemLayoutAtomB{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(((size<0,1>(MmaShapeB_NK{}) * size<2>(MmaShapeB_NK{})) % size<1>(SmemLayoutAtomB{})) == 0, + "SmemLayoutAtom must evenly divide tile shape."); + static_assert(cute::is_void_v, + "SM100 UMMA cannot have a non-void copy atom for smem sourced instructions."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(MmaShapeA_MK{}, Int{}), + cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + using SmemLayoutB = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(MmaShapeB_NK{}, Int{}), + cute::conditional_t(), Step<_2,_1,_3>, Step<_1,_2,_3>>{})); + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source both A and B operand from smem_desc for this mainloop."); + static_assert( + (size(AtomThrShapeMNK{}) == 1 && + (cute::is_same_v || cute::is_same_v)) || + (size(AtomThrShapeMNK{}) == 2 && + (cute::is_same_v || cute::is_same_v)), + "GmemTiledCopy - invalid TMA copy atom specified."); + static_assert( + (size(AtomThrShapeMNK{}) == 1 && + (cute::is_same_v || cute::is_same_v)) || + (size(AtomThrShapeMNK{}) == 2 && + (cute::is_same_v || cute::is_same_v)), + "GmemTiledCopy - invalid TMA copy atom specified."); + + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128, _0> { + cute::ArrayEngine> smem_A_real; + cute::ArrayEngine> smem_A_imag; + cute::ArrayEngine> smem_B_real; + cute::ArrayEngine> smem_B_imag; + } tensors; + + using PipelineStorage = typename MainloopPipeline::SharedStorage; + PipelineStorage pipeline; + }; + + // Expose shared storage for tensors/pipelines separately to allow kernel layer to reorder them. + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Only one thread issues the TMA and updates the barriers in a 2SM MMA, adjust bytes accordingly + static constexpr uint32_t TmaTransactionBytes = 2 * ( + cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutA{})) * static_cast(cute::sizeof_bits::value))) + + cutlass::bits_to_bytes(size(AtomThrShapeMNK{}) * (cosize(take<0,3>(SmemLayoutB{})) * static_cast(cute::sizeof_bits::value)))); + + template + struct TmemStorage { + AccTensor accumulators; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A_real{nullptr}; + StrideA dA_real{}; + ElementA const* ptr_A_imag{nullptr}; + StrideA dA_imag{}; + ElementB const* ptr_B_real{nullptr}; + StrideB dB_real{}; + ElementB const* ptr_B_imag{nullptr}; + StrideB dB_imag{}; + }; + + template< + class KTileCount, + class GTensorPartitionedA, class GTensorPartitionedB, + class STensorA, class STensorB + > + struct LoadParams { + // For scheduler + KTileCount k_tiles; + // for input tensor values + GTensorPartitionedA tAgA_real_mkl; + GTensorPartitionedA tAgA_imag_mkl; + GTensorPartitionedB tBgB_real_nkl; + GTensorPartitionedB tBgB_imag_nkl; + STensorA tAsA_real; + STensorA tAsA_imag; + STensorB tBsB_real; + STensorB tBsB_imag; + // for input tensor values + uint16_t mcast_mask_a; + uint16_t mcast_mask_b; + + CUTLASS_DEVICE + LoadParams ( + KTileCount k_tiles_, + GTensorPartitionedA tAgA_real_mkl_, GTensorPartitionedA tAgA_imag_mkl_, + GTensorPartitionedB tBgB_real_nkl_, GTensorPartitionedB tBgB_imag_nkl_, + STensorA tAsA_real_, STensorA tAsA_imag_, + STensorB tBsB_real_, STensorB tBsB_imag_, + uint16_t mcast_mask_a_, uint16_t mcast_mask_b_) + : k_tiles(k_tiles_) + , tAgA_real_mkl(tAgA_real_mkl_), tAgA_imag_mkl(tAgA_imag_mkl_) + , tBgB_real_nkl(tBgB_real_nkl_), tBgB_imag_nkl(tBgB_imag_nkl_) + , tAsA_real(tAsA_real_), tAsA_imag(tAsA_imag_) + , tBsB_real(tBsB_real_), tBsB_imag(tBsB_imag_) + , mcast_mask_a(mcast_mask_a_), mcast_mask_b(mcast_mask_b_) {} + }; + + template + struct MmaParams { + TiledMma tiled_mma_a_pos; + TiledMmaANeg tiled_mma_a_neg; + FragmentA tCrA_real; + FragmentA tCrA_imag; + FragmentB tCrB_real; + FragmentB tCrB_imag; + + CUTLASS_DEVICE + MmaParams ( + TiledMma tiled_mma_a_pos_, TiledMmaANeg tiled_mma_a_neg_, + FragmentA tCrA_real_, FragmentA tCrA_imag_, + FragmentB tCrB_real_, FragmentB tCrB_imag_) + : tiled_mma_a_pos(tiled_mma_a_pos_), tiled_mma_a_neg(tiled_mma_a_neg_) + , tCrA_real(tCrA_real_), tCrA_imag(tCrA_imag_) + , tCrB_real(tCrB_real_), tCrB_imag(tCrB_imag_) {} + }; + + // Device side kernel params + struct Params { + using ClusterLayout_VMNK = decltype(tiled_divide(make_layout(conditional_return(make_shape(uint32_t(0), uint32_t(0), Int<1>{}), ClusterShape{})), make_tile(typename TiledMma::AtomThrID{}))); + + using TMA_A = decltype(make_tma_atom_A_sm100( + GmemTiledCopyA{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + + using TMA_B = decltype(make_tma_atom_B_sm100( + GmemTiledCopyB{}, + make_tensor(recast_ptr(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + ClusterLayout_VMNK{}) + ); + + TMA_A tma_load_a_real; + TMA_A tma_load_a_imag; + TMA_B tma_load_b_real; + TMA_B tma_load_b_imag; + TMA_A tma_load_a_real_fallback; + TMA_A tma_load_a_imag_fallback; + TMA_B tma_load_b_real_fallback; + TMA_B tma_load_b_imag_fallback; + dim3 cluster_shape_fallback; + }; + + CUTLASS_DEVICE + CollectiveMma(Params const& params, ClusterShape cluster_shape, uint32_t block_rank_in_cluster) + : cluster_shape_(cluster_shape) + , block_rank_in_cluster_(block_rank_in_cluster) { + if constexpr (IsDynamicCluster) { + const bool is_fallback_cluster = (cute::size<0>(cluster_shape_) == params.cluster_shape_fallback.x && + cute::size<1>(cluster_shape_) == params.cluster_shape_fallback.y); + observed_tma_load_a_real_ = is_fallback_cluster ? ¶ms.tma_load_a_real_fallback : ¶ms.tma_load_a_real; + observed_tma_load_a_imag_ = is_fallback_cluster ? ¶ms.tma_load_a_imag_fallback : ¶ms.tma_load_a_imag; + observed_tma_load_b_real_ = is_fallback_cluster ? ¶ms.tma_load_b_real_fallback : ¶ms.tma_load_b_real; + observed_tma_load_b_imag_ = is_fallback_cluster ? ¶ms.tma_load_b_imag_fallback : ¶ms.tma_load_b_imag; + } + else { + observed_tma_load_a_real_ = ¶ms.tma_load_a_real; + observed_tma_load_a_imag_ = ¶ms.tma_load_a_imag; + observed_tma_load_b_real_ = ¶ms.tma_load_b_real; + observed_tma_load_b_imag_ = ¶ms.tma_load_b_imag; + } + } + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cutlass::KernelHardwareInfo const& hw_info = cutlass::KernelHardwareInfo{}) { + (void) workspace; + + // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + auto ptr_A_real = recast_ptr(args.ptr_A_real); + auto ptr_A_imag = recast_ptr(args.ptr_A_imag); + + auto ptr_B_real = recast_ptr(args.ptr_B_real); + auto ptr_B_imag = recast_ptr(args.ptr_B_imag); + + Tensor tensor_a_real = make_tensor(ptr_A_real, make_layout(make_shape(M,K,L), args.dA_real)); + Tensor tensor_a_imag = make_tensor(ptr_A_imag, make_layout(make_shape(M,K,L), args.dA_imag)); + + Tensor tensor_b_real = make_tensor(ptr_B_real, make_layout(make_shape(N,K,L), args.dB_real)); + Tensor tensor_b_imag = make_tensor(ptr_B_imag, make_layout(make_shape(N,K,L), args.dB_imag)); + + auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, hw_info.cluster_shape); + // Cluster layout for TMA construction + auto cluster_layout_vmnk = tiled_divide(make_layout(cluster_shape), make_tile(typename TiledMma::AtomThrID{})); + + auto cluster_shape_fallback = conditional_return(make_shape(hw_info.cluster_shape_fallback.x, hw_info.cluster_shape_fallback.y, Int<1>{}), ClusterShape{}); + // Cluster layout for TMA construction + auto cluster_layout_vmnk_fallback = tiled_divide(make_layout(cluster_shape_fallback), make_tile(typename TiledMma::AtomThrID{})); + + typename Params::TMA_A tma_load_a_real = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_real, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_imag = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_imag, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b_real = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_real, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_B tma_load_b_imag = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_imag, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk); + + typename Params::TMA_A tma_load_a_real_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_real, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_A tma_load_a_imag_fallback = make_tma_atom_A_sm100( + GmemTiledCopyA{}, + tensor_a_imag, + SmemLayoutA{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_real_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_real, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + typename Params::TMA_B tma_load_b_imag_fallback = make_tma_atom_B_sm100( + GmemTiledCopyB{}, + tensor_b_imag, + SmemLayoutB{}(_,_,_,cute::Int<0>{}), + TileShape{}, + TiledMma{}, + cluster_layout_vmnk_fallback); + + return { + tma_load_a_real, + tma_load_a_imag, + tma_load_b_real, + tma_load_b_imag, + tma_load_a_real_fallback, + tma_load_a_imag_fallback, + tma_load_b_real_fallback, + tma_load_b_imag_fallback, + hw_info.cluster_shape_fallback + }; + } + + template + static bool + can_implement( + ProblemShape const& problem_shape, + [[maybe_unused]] Arguments const& args) { + auto problem_shape_MNKL = append<4>(problem_shape, 1); + auto [M,N,K,L] = problem_shape_MNKL; + + bool implementable = true; + constexpr int min_tma_aligned_elements_A = 128 / cute::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(M,K,L), StrideA{}); + constexpr int min_tma_aligned_elements_B = 128 / cute::sizeof_bits::value; + implementable = implementable && cutlass::detail::check_alignment(cute::make_shape(N,K,L), StrideB{}); + + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + } + return implementable; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE void + prefetch_tma_descriptors() { + cute::prefetch_tma_descriptor(observed_tma_load_a_real_->get_tma_descriptor()); + cute::prefetch_tma_descriptor(observed_tma_load_a_imag_->get_tma_descriptor()); + cute::prefetch_tma_descriptor(observed_tma_load_b_real_->get_tma_descriptor()); + cute::prefetch_tma_descriptor(observed_tma_load_b_imag_->get_tma_descriptor()); + } + + /// Construct A Single Stage's Accumulator Shape + CUTLASS_DEVICE static + auto + partition_accumulator_shape() { + auto acc_shape = append( + partition_shape_C(TiledMma{}, take<0,2>(TileShape{})), + Int<2>{}); // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,TMEM_PIPE,2) + + return acc_shape; + } + + template + CUTLASS_DEVICE static + auto + slice_accumulator(TmemStorage tmem_storage, int stage) { + return cute::make_tuple(tmem_storage.accumulators(_,_,_,_,stage)); + } + + template + CUTLASS_DEVICE static + auto + init_tmem_tensors(EpilogueTile epi_tile) { + TiledMma tiled_mma; + auto acc_shape = partition_accumulator_shape(); + // ((MMA_TILE_M,MMA_TILE_N),MMA_M,MMA_N,ACC_PIPE) where ACC_PIPE=2 so we can double buffer our accumulators for mainloop and epilogue. + Tensor accumulators = cutlass::detail::make_sm100_accumulator( + tiled_mma, acc_shape, EpilogueTile{}); + + TmemStorage tmem_storage; + tmem_storage.accumulators = accumulators; + + return tmem_storage; + } + + template + CUTLASS_DEVICE static + void + set_tmem_offsets(TmemStorage& tmem_storage, uint32_t tmem_base_addr) { + tmem_storage.accumulators.data() = tmem_base_addr; + } + + /// Set up the data needed by this collective for load. + /// Returned tuple must contain at least two elements, with the first two elements being: + /// gA_(real/imag)_mkl - The tiled tma tensor for input A_(real/imag) + /// gB_(real/imag)_nkl - The tiled tma tensor for input B_(real/imag) + /// tAsA_(real/imag) - partitioned smem tensor for A_(real/imag) + /// tBsB_(real/imag) - partitioned smem tensor for B_(real/imag) + /// mcast_mask_a - tma multicast mask for A_(real/imag) + /// mcast_mask_b - tma multicast mask for B_(real/imag) + template + CUTLASS_DEVICE auto + load_init( + ProblemShape_MNKL const& problem_shape_MNKL, + TensorStorage& shared_tensors) const { + using X = Underscore; + + // Separate out problem shape for convenience + auto [M,N,K,L] = problem_shape_MNKL; + + // Represent the full tensors -- get these from TMA + Tensor mA_real_mkl = observed_tma_load_a_real_->get_tma_tensor(make_shape(M,K,L)); + Tensor mA_imag_mkl = observed_tma_load_a_imag_->get_tma_tensor(make_shape(M,K,L)); + Tensor mB_real_nkl = observed_tma_load_b_real_->get_tma_tensor(make_shape(N,K,L)); + Tensor mB_imag_nkl = observed_tma_load_b_imag_->get_tma_tensor(make_shape(N,K,L)); + + // Tile the tensors and defer the slice + Tensor gA_real_mkl = local_tile(mA_real_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M, BLK_K, m, k, l) + Tensor gA_imag_mkl = local_tile(mA_imag_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_N, BLK_K, m, k, l) + + Tensor gB_real_nkl = local_tile(mB_real_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l) + Tensor gB_imag_nkl = local_tile(mB_imag_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N, BLK_K, n, k, l) + + // Partition for this CTA + ThrMMA cta_mma = TiledMma{}.get_slice(blockIdx.x % size(typename TiledMma::AtomThrID{})); + + Tensor tCgA_real_mkl = cta_mma.partition_A(gA_real_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + Tensor tCgA_imag_mkl = cta_mma.partition_A(gA_imag_mkl); // (MMA, MMA_M, MMA_K, m, k, l) + + Tensor tCgB_real_nkl = cta_mma.partition_B(gB_real_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + Tensor tCgB_imag_nkl = cta_mma.partition_B(gB_imag_nkl); // (MMA, MMA_N, MMA_K, n, k, l) + + Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (MMA,MMA_M,MMA_K,PIPE) + + Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (MMA,MMA_N,MMA_K,PIPE) + + // Define the CTA-in-cluster Layout and Coord + Layout cta_layout_mnk = make_layout(cluster_shape_); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMma::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster_); + + // Project the cta_layout for tma_a along the n-modes + auto [tAgA_real_mkl, tAsA_real] = tma_partition(*observed_tma_load_a_real_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA_real), group_modes<0,3>(tCgA_real_mkl)); + auto [tAgA_imag_mkl, tAsA_imag] = tma_partition(*observed_tma_load_a_imag_, + get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(sA_imag), group_modes<0,3>(tCgA_imag_mkl)); + + // Project the cta_layout for tma_b along the m-modes + auto [tBgB_real_nkl, tBsB_real] = tma_partition(*observed_tma_load_b_real_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB_real), group_modes<0,3>(tCgB_real_nkl)); + auto [tBgB_imag_nkl, tBsB_imag] = tma_partition(*observed_tma_load_b_imag_, + get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(sB_imag), group_modes<0,3>(tCgB_imag_nkl)); + // TMA Multicast Masks + uint16_t mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + LoadParams load_params { + shape<3>(gA_real_mkl), // for scheduler + tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl, // for input tensor values + tAsA_real, tAsA_imag, tBsB_real, tBsB_imag, // for input tensor values + mcast_mask_a, mcast_mask_b + }; + return load_params; + } + + /// Set up the data needed by this collective for mma compute. + template + CUTLASS_DEVICE auto + mma_init( + [[maybe_unused]] TmemStorage tmem_storage, + TensorStorage& shared_tensors) const { + Tensor sA_real = make_tensor(make_smem_ptr(shared_tensors.smem_A_real.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sA_imag = make_tensor(make_smem_ptr(shared_tensors.smem_A_imag.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + + Tensor sB_real = make_tensor(make_smem_ptr(shared_tensors.smem_B_real.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + Tensor sB_imag = make_tensor(make_smem_ptr(shared_tensors.smem_B_imag.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // Allocate "fragments/descriptors" for A and B matrices + Tensor tCrA_real = TiledMma::make_fragment_A(sA_real); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrA_imag = TiledMma::make_fragment_A(sA_imag); // (MMA,MMA_M,MMA_K,PIPE) + + Tensor tCrB_real = TiledMma::make_fragment_B(sB_real); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrB_imag = TiledMma::make_fragment_B(sB_imag); // (MMA,MMA_N,MMA_K,PIPE) + + CUTE_STATIC_ASSERT_V(Int{} == size<3>(sA_real)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<3>(sB_real)); // PIPE + + TiledMma tiled_mma_a_pos; + TiledMmaANeg tiled_mma_a_neg; + MmaParams mma_params { + tiled_mma_a_pos, tiled_mma_a_neg, + tCrA_real, tCrA_imag, + tCrB_real, tCrB_imag + }; + + return mma_params; + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Producer Perspective + template < + class LoadParams, + class TileCoordMNKL, + class KTileIterator + > + CUTLASS_DEVICE auto + load( + MainloopPipeline mainloop_pipeline, + MainloopPipelineState mainloop_pipe_producer_state, + LoadParams const& load_inputs, + TileCoordMNKL const& cta_coord_mnkl, + KTileIterator k_tile_iter, int k_tile_count) { + + auto [unused_k_tiles, + tAgA_real_mkl, tAgA_imag_mkl, tBgB_real_nkl, tBgB_imag_nkl, + tAsA_real, tAsA_imag, tBsB_real, tBsB_imag, + mcast_mask_a, mcast_mask_b] = load_inputs; + + // slice out the work coord from partitioned tensors + Tensor tAgA_real = tAgA_real_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + Tensor tAgA_imag = tAgA_imag_mkl(_, get<0>(cta_coord_mnkl) / size(typename TiledMma::AtomThrID{}), _, get<3>(cta_coord_mnkl)); + + Tensor tBgB_real = tBgB_real_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + Tensor tBgB_imag = tBgB_imag_nkl(_, get<1>(cta_coord_mnkl), _, get<3>(cta_coord_mnkl)); + + auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state); + + // Issue the Mainloop loads + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + // LOCK mainloop_pipe_producer_state for _writing_ + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token); + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state); + + if (cute::elect_one_sync()) { + copy(observed_tma_load_a_real_->with(*tma_barrier, mcast_mask_a), tAgA_real(_,*k_tile_iter), tAsA_real(_,write_stage)); + copy(observed_tma_load_a_imag_->with(*tma_barrier, mcast_mask_a), tAgA_imag(_,*k_tile_iter), tAsA_imag(_,write_stage)); + + copy(observed_tma_load_b_real_->with(*tma_barrier, mcast_mask_b), tBgB_real(_,*k_tile_iter), tBsB_real(_,write_stage)); + copy(observed_tma_load_b_imag_->with(*tma_barrier, mcast_mask_b), tBgB_imag(_,*k_tile_iter), tBsB_imag(_,write_stage)); + } + + --k_tile_count; + ++k_tile_iter; + } + + return cute::make_tuple(mainloop_pipe_producer_state, k_tile_iter); + } + + /// Perform a Producer Epilogue to prevent early exit of ctas in a Cluster + CUTLASS_DEVICE void + load_tail(MainloopPipeline mainloop_pipeline, MainloopPipelineState mainloop_pipe_producer_state) { + // Issue the epilogue waits + /* This helps avoid early exit of ctas in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class AccumulatorPipeline, + class FrgEngine, class FrgLayout, + class MmaParams, + class CtaTileCoord + > + CUTLASS_DEVICE auto + mma(cute::tuple pipelines, + cute::tuple pipeline_states, + cute::tuple> const& accumulators_pair, + MmaParams const& mma_inputs, + CtaTileCoord cta_tile_coord, + int k_tile_count + ) { + static_assert(is_tmem::value, "Accumulator must be tmem resident."); + static_assert(rank(FrgLayout{}) == 4 && size<3>(FrgLayout{}) == _2{}, "Accumulator must be MMA-partitioned: (MMA, MMA_M, MMA_N, _2)"); + + auto [tiled_mma_a_pos, tiled_mma_a_neg, tCrA_real, tCrA_imag, tCrB_real, tCrB_imag] = mma_inputs; + + auto [mainloop_pipeline, accumulator_pipeline] = pipelines; + auto [mainloop_pipe_consumer_state, accumulator_pipe_producer_state] = pipeline_states; + + uint32_t skip_wait = k_tile_count <= 0; + auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + + // + // PIPELINED MAIN LOOP + // + tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::Zero; + tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::Zero; + + auto accumulators = get<0>(accumulators_pair); + auto accumulators_real = accumulators(_,_,_,0); + auto accumulators_imag = accumulators(_,_,_,1); + + // Wait for tmem accumulator buffer to become empty with a flipped phase + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + // WAIT on mainloop_pipe_consumer_state until its data are available + // (phase bit flips from mainloop_pipe_consumer_state.phase() value) + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + + // Compute on k_tile + int read_stage = mainloop_pipe_consumer_state.index(); + // Save current mainlop pipeline read state + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + + // Advance mainloop_pipe + ++mainloop_pipe_consumer_state; + --k_tile_count; + skip_wait = k_tile_count <= 0; + // Peek at next iteration + barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + + // Unroll the K mode manually so we can set scale C to 1 + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA_real); ++k_block) { + // (V,M) x (V,N) => (V,M,N) + + // Calculate real acc, 1st step + // realAcc += realA * realB + cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_real); + + // Calculate imag acc, 1st step + if constexpr (cute::is_same_v) { + // imagAcc += realA * (-imagB) + cute::gemm(tiled_mma_a_neg, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag); + } else { + // imagAcc += realA * imagB + cute::gemm(tiled_mma_a_pos, tCrA_real(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_imag); + } + + tiled_mma_a_pos.accumulate_ = UMMA::ScaleOut::One; + tiled_mma_a_neg.accumulate_ = UMMA::ScaleOut::One; + + // Calculate real acc, 2nd step + if constexpr (cute::is_same_v) { + // realAcc -= imagA * imagB + cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real); + } else { + // realAcc += imagA * imagB + cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_imag(_,_,k_block,read_stage), accumulators_real); + } + + // Calculate imag acc, 2nd step + if constexpr (cute::is_same_v) { + // imagAcc += (-imagA) * realB + cute::gemm(tiled_mma_a_neg, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag); + } else { + // imagAcc += imagA * realB + cute::gemm(tiled_mma_a_pos, tCrA_imag(_,_,k_block,read_stage), tCrB_real(_,_,k_block,read_stage), accumulators_imag); + } + } + mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state); + } + + return mainloop_pipe_consumer_state; + } + +protected: + + typename Params::TMA_A const* observed_tma_load_a_real_ = nullptr; + typename Params::TMA_A const* observed_tma_load_a_imag_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_real_ = nullptr; + typename Params::TMA_B const* observed_tma_load_b_imag_ = nullptr; + + ClusterShape cluster_shape_; + uint32_t block_rank_in_cluster_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp b/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp index 34b3aebc..458ee1af 100755 --- a/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp +++ b/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp @@ -259,8 +259,8 @@ struct CollectiveMma< struct TensorStorage : cute::aligned_struct<128, _0> { alignas(1024) cute::ArrayEngine> smem_A; alignas(1024) cute::ArrayEngine> smem_B; - cute::ArrayEngine> smem_SFA; - cute::ArrayEngine> smem_SFB; + alignas(16) cute::ArrayEngine> smem_SFA; + alignas(16) cute::ArrayEngine> smem_SFB; } tensors; struct TensorMapStorage : cute::aligned_struct<128, _0> { diff --git a/include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp b/include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp index edb945b9..9cb80518 100755 --- a/include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp +++ b/include/cutlass/gemm/collective/sm120_blockscaled_mma_tma.hpp @@ -256,8 +256,8 @@ struct CollectiveMma< struct TensorStorage : cute::aligned_struct<128, _0> { alignas(1024) cute::ArrayEngine> smem_A; alignas(1024) cute::ArrayEngine> smem_B; - cute::ArrayEngine> smem_SFA; - cute::ArrayEngine> smem_SFB; + alignas(16) cute::ArrayEngine> smem_SFA; + alignas(16) cute::ArrayEngine> smem_SFB; } tensors; using PipelineStorage = typename MainloopPipeline::SharedStorage; alignas(16) PipelineStorage pipeline_storage; diff --git a/include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp b/include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp index 46ecaa0d..f65629c5 100755 --- a/include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp +++ b/include/cutlass/gemm/collective/sm120_blockscaled_sparse_mma_tma.hpp @@ -296,9 +296,9 @@ struct CollectiveMma< struct TensorStorage : cute::aligned_struct<128> { alignas(1024) cute::ArrayEngine> smem_A; alignas(1024) cute::ArrayEngine> smem_B; - cute::ArrayEngine> smem_SFA; - cute::ArrayEngine> smem_SFB; - cute::ArrayEngine{}> smem_E; + alignas(16) cute::ArrayEngine> smem_SFA; + alignas(16) cute::ArrayEngine> smem_SFB; + alignas(16) cute::ArrayEngine{}> smem_E; } tensors; using PipelineStorageMK = typename MainloopPipelineMK::SharedStorage; diff --git a/include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp b/include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp index 52089b70..30587e46 100644 --- a/include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp +++ b/include/cutlass/gemm/collective/sm120_sparse_mma_tma.hpp @@ -253,7 +253,7 @@ struct CollectiveMma< struct TensorStorage : cute::aligned_struct<128, _0> { alignas(1024) cute::ArrayEngine> smem_A; alignas(1024) cute::ArrayEngine> smem_B; - cute::ArrayEngine{}> smem_E; + alignas(16) cute::ArrayEngine{}> smem_E; } tensors; using PipelineStorageMK = typename MainloopPipelineMK::SharedStorage; diff --git a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp index 20c40956..f27731bc 100644 --- a/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp +++ b/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp @@ -212,8 +212,7 @@ struct CollectiveMma< static_assert(cute::is_same_v, "ElementAccumulator and ElementBlockScale should be same datatype"); - // For TileShapeM < 128, NumSplitsM should be 1 - using NumSplitsM = cute::conditional_t(TileShape_{}) < _128{}, _1, cute::C(TileShape_{}) / 128>>; + using NumSplitsM = cute::C(TileShape_{}) / 128>; static_assert(NumSplitsM{} == 1 || NumSplitsM{} == 2); struct SharedStorage { diff --git a/include/cutlass/gemm/device/gemm_blockwise.h b/include/cutlass/gemm/device/gemm_blockwise.h new file mode 100644 index 00000000..13e22941 --- /dev/null +++ b/include/cutlass/gemm/device/gemm_blockwise.h @@ -0,0 +1,761 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Template for a pipelined GEMM kernel. Does not compute batching or + support split-K. +*/ + +#pragma once + +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" +#include "cutlass/gemm/device/default_gemm_configuration.h" +#include "cutlass/gemm/kernel/default_gemm.h" +#include "cutlass/gemm/kernel/gemm.h" +#include "cutlass/gemm/kernel/gemm_blockwise.h" +#include "cutlass/gemm/threadblock/threadblock_swizzle.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/permute.h" +#include "cutlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient CUTLASS GEMM + kernels that may be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters + onto specific CUTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel + parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most + plausible GEMM configurations for each supported architecture. Consequently, + not all parameters are exposed to the top-level interface. Rather, sensible + defaults at each level of the CUTLASS hierarchy are selected to tradeoff + simplicity of the interface with flexibility. We expect most configurations to + be specified at this level. Applications with more exotic requirements may + construct their kernels of interest using CUTLASS components at the + threadblock, warp, and thread levels of abstraction. + + CUTLASS exposes computations using the functor design pattern in which objects + compose some internal state with an overloaded function call operator. This + enables decoupling of initialization from execution, possibly reducing + overhead during steady state phases of application execution. + + CUTLASS device-level operators expose an Arguments structure encompassing each + logical input to the computation. This is distinct from the kernel-level + Params structure pattern which contains application-specific precomputed state + needed by the device code. + + Example of a CUTLASS GEMM operator implementing the functionality of cuBLAS's + SGEMM NN is as follows: + + // + // Instantiate the CUTLASS GEMM operator. + // + + cutlass::gemm::device::Gemm< + float, + cutlass::layout::ColumnMajor, + float, + cutlass::layout::ColumnMajor, + float, + cutlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + cutlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, {B, ldb}, // + TensorRef ref_B, {C, ldc}, // TensorRef ref_C, {D, ldd}, // + TensorRef ref_D, {alpha, beta} // + EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal accumulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm89, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Element Type for for the scalesl + typename ElementScale_ = float, + /// Layout for the scales. + typename LayoutScale_ = cutlass::layout::RowMajor, + /// Scale Block Size. + int ScaleBlockSize_ = 128, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout = layout::NoPermute> +class GemmBlockwise { +public: + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + using ElementScale = ElementScale_; + using LayoutScale = LayoutScale_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static int const kScaleBlockSize = ScaleBlockSize_; + + + static_assert(kScaleBlockSize == 128, "Scale block size has to be 128 for now."); + // Ensure the threadblock K-dimension is 128 + static_assert(ThreadblockShape::kK == kScaleBlockSize, + "GemmBlockwise requires ThreadblockShape::kK equale to Scale Block Size"); + + static_assert(cutlass::platform::is_same::value, + "Scales have to be row major for now."); + + static_assert(cutlass::platform::is_same::value, + "Scales have to be float."); + + // Tensor reference type for the FP8 scale tensors + using TensorRefScale = cutlass::TensorRef; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Define the kernel + using GemmKernel = typename kernel::GemmBlockwise< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementC, + LayoutC, ElementAccumulator, ElementScale, LayoutScale, + OperatorClass, ArchTag, ThreadblockShape, + WarpShape, InstructionShape, EpilogueOutputOp, ThreadblockSwizzle, + kStages, kSplitKSerial, Operator, SharedMemoryClearOption::kNone, GatherA, + GatherB, ScatterD, PermuteDLayout>::GemmKernel; + + /// Argument structure + struct Arguments { + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + // For gather+scatter operations + int const *gather_A_indices; + int const *gather_B_indices; + int const *scatter_D_indices; + + // Dequantization scale tensors (row-major) + TensorRefScale scale_A; + TensorRefScale scale_B; + + // + // Methods + // + + /// Default ctor + CUTLASS_HOST_DEVICE + Arguments() : problem_size(0, 0, 0), split_k_slices(1) {} + + /// Constructs an Arguments structure + CUTLASS_HOST_DEVICE + Arguments(GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + TensorRefScale scale_A_, + TensorRefScale scale_B_, + typename EpilogueOutputOp::Params epilogue_ = typename EpilogueOutputOp::Params(), + int split_k_slices = 1, int const *gather_A_indices_ = nullptr, + int const *gather_B_indices_ = nullptr, + int const *scatter_D_indices_ = nullptr) + : problem_size(problem_size_), ref_A(ref_A_), ref_B(ref_B_), + ref_C(ref_C_), ref_D(ref_D_), scale_A(scale_A_), + scale_B(scale_B_), epilogue(epilogue_), + split_k_slices(split_k_slices), gather_A_indices(gather_A_indices_), + gather_B_indices(gather_B_indices_), + scatter_D_indices(scatter_D_indices_) {} + }; + +private: + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + /// Constructs the GEMM. + GemmBlockwise() {} + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + // Require the problem K dimension to be an exact multiple of the Threadblock K tile. + if (args.problem_size.k() % ThreadblockShape::kK != 0) { + return Status::kErrorInvalidProblem; + } + + if (!kSplitKSerial && args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + + // ------------------------------------------------------------------ + // Validate scale tensor leading dimensions. + // Row-major layout implies stride(0) equals number of columns (kBlocks). + // Both scale_A (mBlocks × kBlocks) and scale_B (nBlocks × kBlocks) must + // therefore have stride(0) == kBlocks where kBlocks = ceil_div(K, 128). + // ------------------------------------------------------------------ + int const kBlocks = (args.problem_size.k() + ThreadblockShape::kK - 1) / ThreadblockShape::kK; + + if (args.scale_A.stride(0) != kBlocks || args.scale_B.stride(0) != kBlocks) { + return Status::kErrorInvalidProblem; + } + + Status status = GemmKernel::can_implement( + args.problem_size, args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), args.ref_C.non_const_ref(), args.ref_D); + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + cutlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial && args.split_k_slices > 1) { + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, + cudaStream_t stream = nullptr) { + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + cudaError_t result = cudaMemsetAsync(workspace, 0, bytes, stream); + + if (result != cudaSuccess) { + return Status::kErrorInternal; + } + } + } else { + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + // Initialize the Params structure + params_ = typename GemmKernel::Params{args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.scale_A, + args.scale_B, + args.epilogue, + static_cast(workspace), + args.gather_A_indices, + args.gather_B_indices, + args.scatter_D_indices}; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + if (kSplitKSerial && args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + params_.scale_A.reset(args.scale_A.data()); + params_.scale_B.reset(args.scale_B.data()); + params_.output_op = args.epilogue; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(cudaStream_t stream = nullptr) { + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + cudaError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) { + result = cudaFuncSetAttribute(Kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != cudaSuccess) { + return Status::kErrorInternal; + } + } + + cutlass::arch::synclog_setup(); + cutlass::Kernel<<>>(params_); + + result = cudaGetLastError(); + + return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(cudaStream_t stream = nullptr) { return run(stream); } + + /// Runs the kernel using initialized state. + Status operator()(Arguments const &args, void *workspace = nullptr, + cudaStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and +/// operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Element Type for for the scalesl + typename ElementScale_, + /// Layout for the scales. + typename LayoutScale_, + /// Scale Block Size. + int ScaleBlockSize_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// If true, kernel supports split-K as a serial reduction + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator_, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout> +class GemmBlockwise { +public: + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + using ElementScale = ElementScale_; + using LayoutScale = LayoutScale_; + + static int const kScaleBlockShape = ScaleBlockSize_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + + // Alias for per-tile FP8 dequantization scale tensors + using TensorRefScale = cutlass::TensorRef; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + static bool const kSplitKSerial = SplitKSerial; + + using UnderlyingOperator = + GemmBlockwise::type, + ElementA, typename layout::LayoutTranspose::type, + ElementC, layout::RowMajor, ElementAccumulator, + OperatorClass, ArchTag, ThreadblockShape, WarpShape, + InstructionShape, ElementScale, LayoutScale, kScaleBlockShape, + EpilogueOutputOp, ThreadblockSwizzle, + Stages, kAlignmentB, kAlignmentA, SplitKSerial, Operator, + GatherB, GatherA, ScatterD, PermuteDLayout>; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = UnderlyingOperator::kAlignmentC; + + /// Argument structure + struct Arguments { + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + // For gather+scatter operations + int const *gather_A_indices; + int const *gather_B_indices; + int const *scatter_D_indices; + + // Dequantization scale tensors (row-major) + TensorRefScale scale_A; + TensorRefScale scale_B; + + // + // Methods + // + + /// Default ctor + CUTLASS_HOST_DEVICE + Arguments() {} + + /// Constructs an Arguments structure + CUTLASS_HOST_DEVICE + Arguments(GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + TensorRefScale scale_A_, + TensorRefScale scale_B_, + typename EpilogueOutputOp::Params epilogue_ = typename EpilogueOutputOp::Params(), + int split_k_slices = 1, int const *gather_A_indices_ = nullptr, + int const *gather_B_indices_ = nullptr, + int const *scatter_D_indices_ = nullptr) + : problem_size(problem_size_), ref_A(ref_A_), ref_B(ref_B_), + ref_C(ref_C_), ref_D(ref_D_), scale_A(scale_A_), + scale_B(scale_B_), epilogue(epilogue_), + split_k_slices(split_k_slices), gather_A_indices(gather_A_indices_), + gather_B_indices(gather_B_indices_), + scatter_D_indices(scatter_D_indices_) {} + }; + +private: + UnderlyingOperator underlying_operator_; + +public: + /// Constructs the GEMM. + GemmBlockwise() {} + + /// Helper to construct a transposed equivalent for the underying GEMM + /// operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + return UnderlyingArguments( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {args.ref_B.data(), args.ref_B.stride(0)}, + {args.ref_A.data(), args.ref_A.stride(0)}, + {args.ref_C.data(), args.ref_C.stride(0)}, + {args.ref_D.data(), args.ref_D.stride(0)}, + {args.scale_B.data(), args.scale_B.stride(0)}, + {args.scale_A.data(), args.scale_A.stride(0)}, + args.epilogue, + args.split_k_slices, args.gather_B_indices, args.gather_A_indices, + args.scatter_D_indices); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + return UnderlyingOperator::get_workspace_size( + to_underlying_arguments(args)); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) { + return underlying_operator_.initialize(to_underlying_arguments(args), + workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + return underlying_operator_.update(to_underlying_arguments(args), + workspace); + } + + /// Runs the kernel using initialized state. + Status run(cudaStream_t stream = nullptr) { + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(cudaStream_t stream = nullptr) { return run(stream); } + + /// Runs the kernel using initialized state. + Status operator()(Arguments const &args, void *workspace = nullptr, + cudaStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace cutlass +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/device/gemm_universal_base.h b/include/cutlass/gemm/device/gemm_universal_base.h index 1aac3f98..0998735f 100644 --- a/include/cutlass/gemm/device/gemm_universal_base.h +++ b/include/cutlass/gemm/device/gemm_universal_base.h @@ -166,6 +166,7 @@ protected: } } +#ifndef __QNX__ // Update SM occupancy member cudart_result = cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags( &sm_occupancy_, @@ -177,6 +178,7 @@ protected: CUTLASS_TRACE_HOST(" cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned error " << cudaGetErrorString(cudart_result)); return Status::kErrorInternal; } +#endif // Update device ordinal member on success device_ordinal_ = current_ordinal; diff --git a/include/cutlass/gemm/dispatch_policy.hpp b/include/cutlass/gemm/dispatch_policy.hpp index deefe903..3ea71399 100644 --- a/include/cutlass/gemm/dispatch_policy.hpp +++ b/include/cutlass/gemm/dispatch_policy.hpp @@ -787,6 +787,24 @@ struct KernelPtrArrayTmaWarpSpecialized2SmFastFP32Sm100 final : KernelSchedu struct KernelPtrArrayTmaWarpSpecialized1SmFastFP32SmemSm100 final : KernelSchedule1Sm, KernelTmaWarpSpecializedPtrArrayFastFP32SmemSm100 { }; struct KernelPtrArrayTmaWarpSpecialized2SmFastFP32SmemSm100 final : KernelSchedule2Sm, KernelTmaWarpSpecializedPtrArrayFastFP32SmemSm100 { }; +/////////////////////////////////////////////////////////////////////////////////////////////////////// +// SM100 Interleaved Complex GEMM Dispatch Policies +/////////////////////////////////////////////////////////////////////////////////////////////////////// +struct KernelScheduleSm100InterleavedComplexTF32Gemm : KernelScheduleSm100 {}; +// Transform GEMM: Specialize for Interleaved Complex GEMMs +struct KernelTmaWarpSpecialized1SmInterleavedComplexTF32Sm100 final : KernelSchedule1Sm, KernelScheduleSm100InterleavedComplexTF32Gemm { }; +struct KernelTmaWarpSpecialized2SmInterleavedComplexTF32Sm100 final : KernelSchedule2Sm, KernelScheduleSm100InterleavedComplexTF32Gemm { }; + +/////////////////////////////////////////////////////////////////////////////////////////////////////// +// SM100 Ptr-Array Interleaved Complex GEMM Dispatch Policies +/////////////////////////////////////////////////////////////////////////////////////////////////////// +// Interleaved Complex GEMM + (Ptr array or Group GEMM) +struct KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm : KernelScheduleSm100InterleavedComplexTF32Gemm {}; +// Ptr-Array Transform GEMM: Specialize for 1SM vs 2SM Complex GEMM +// Transform GEMM: Specialize for Interleaved Complex GEMMs +struct KernelPtrArrayTmaWarpSpecialized1SmInterleavedComplexTF32Sm100 final : KernelSchedule1Sm, KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm { }; +struct KernelPtrArrayTmaWarpSpecialized2SmInterleavedComplexTF32Sm100 final : KernelSchedule2Sm, KernelScheduleSm100PtrArrayInterleavedComplexTF32Gemm { }; + /////////////////////////////////////////////////////////////////////////////////////////////////////// // SM100 Sparse GEMM Dispatch Policies /////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1133,6 +1151,33 @@ struct MainloopSm100TmaUmmaWarpSpecializedFastF32 { }; +template< + // Number of Pipeline stages for + // MainloopLoad <-> Transformation + int ComputationPipelineStageCount_, + // TileScheduler pipeline depth + int SchedulerPipelineStageCount_, + // Accmulator pipeline depth + int AccumulatorPipelineStageCount_, + // Number of Pipeline stages for + // Transformation <-> MMA + int TransformationPipelineStageCount_, + class ClusterShape_ = Shape<_1,_1,_1>, + class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x +> +struct MainloopSm100TmaUmmaWarpSpecializedInterleavedComplexTF32 { + constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_; + constexpr static int TransformationPipelineStageCount = TransformationPipelineStageCount_; + constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32; + using ClusterShape = ClusterShape_; + using AccumulatorCopyAtom = AccumulatorCopyAtom_; + using ArchTag = arch::Sm100; + using Schedule = KernelTmaWarpSpecializedInputTransformSm100; + + // For backwards compatibility with GemmUniversalAdapter. + constexpr static int Stages = ComputationPipelineStageCount; +}; + // n-buffer in smem, pipelined with Blackwell Mixed Input kernel with UMMA (HwScaled) and TMA, template< // Number of Pipeline stages for @@ -1162,6 +1207,23 @@ struct MainloopSm100TmaUmmaWarpSpecializedMixedInput { }; +// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule +template< + int Stages_, + // TileScheduler pipeline depth + int SchedulerPipelineStageCount_, + // Accmulator pipeline depth + int AccumulatorPipelineStageCount_, + class ClusterShape_ = Shape<_1,_1,_1> +> +struct MainloopSm100TmaUmmaWarpSpecializedPlanarComplex { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + using ArchTag = arch::Sm100; + using Schedule = KernelTmaWarpSpecializedSm100; + constexpr static bool IsOverlappingAccum = false; +}; + // n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule template< int Stages_, @@ -1224,6 +1286,22 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedBlockScaled { +// n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule +template< + int Stages_, + int SchedulerPipelineStageCount_, + int AccumulatorPipelineStageCount_, + class ClusterShape_ = Shape<_1,_1,_1> +> +struct MainloopSm100ArrayTmaUmmaWarpSpecializedPlanarComplex { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + using ArchTag = arch::Sm100; + constexpr static bool IsOverlappingAccum = false; + using Schedule = KernelPtrArrayTmaWarpSpecializedSm100; +}; + + // n-buffer in smem, pipelined with Blackwell Fast FP32 kernel with UMMA (HwScaled) and TMA, // Warp specialized dynamic schedule template< @@ -1271,6 +1349,35 @@ struct MainloopSm100ArrayTmaUmmaWarpSpecializedFastF32 { constexpr static int Stages = Load2TransformPipelineStageCount; }; + +template< + // Number of Pipeline stages for + // MainloopLoad <-> Transformation + int ComputationPipelineStageCount_, + // TileScheduler pipeline depth + int SchedulerPipelineStageCount_, + // Accmulator pipeline depth + int AccumulatorPipelineStageCount_, + // Number of Pipeline stages for + // Transformation <-> MMA + int TransformationPipelineStageCount_, + class ClusterShape_ = Shape<_1,_1,_1>, + class AccumulatorCopyAtom_ = cute::SM100_TMEM_LOAD_16dp256b1x +> +struct MainloopSm100ArrayTmaUmmaWarpSpecializedInterleavedComplexTF32 { + constexpr static int ComputationPipelineStageCount = ComputationPipelineStageCount_; + constexpr static int TransformationPipelineStageCount = TransformationPipelineStageCount_; + constexpr static detail::KernelInputTransformType InputTransformType = detail::KernelInputTransformType::InterleavedComplexTF32; + using ClusterShape = ClusterShape_; + using AccumulatorCopyAtom = AccumulatorCopyAtom_; + using ArchTag = arch::Sm100; + using Schedule = KernelPtrArrayTmaWarpSpecializedInputTransformSm100; + + // For backwards compatibility with GemmUniversalAdapter. + constexpr static int Stages = ComputationPipelineStageCount; +}; + + // n-buffer in smem, pipelined with Blackwell UMMA and TMA, Warp specialized dynamic schedule template< int LoadABPipelineStageCount_, diff --git a/include/cutlass/gemm/kernel/gemm_blockwise.h b/include/cutlass/gemm/kernel/gemm_blockwise.h new file mode 100644 index 00000000..f8ee9d0d --- /dev/null +++ b/include/cutlass/gemm/kernel/gemm_blockwise.h @@ -0,0 +1,223 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 + Default kernel-level GEMM definitions combine threadblock-scoped matrix + multiply-add with the appropriate threadblock-scoped epilogue. + + Note, CUTLASS epilogues universally target row-major outputs. Column-major + outputs are accommodated by exchanging A and B operands and assuming + transposed layouts. Partial specializations here choose + 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "cutlass/arch/wmma.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/epilogue/threadblock/default_epilogue_simt.h" +#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "cutlass/epilogue/threadblock/epilogue.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/kernel/gemm_universal_blockwise.h" +#include "cutlass/gemm/kernel/gemm_pipelined.h" +#include "cutlass/gemm/threadblock/default_mma.h" +#include "cutlass/gemm/threadblock/default_mma_core_simt.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" +#include "cutlass/gemm/threadblock/default_mma_multistage_blockwise.h" +#include "cutlass/gemm/threadblock/threadblock_swizzle.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/permute.h" +#include "cutlass/numeric_types.h" +#include "cutlass/transform/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Element Type for for the scalesl + typename ElementScale, + /// Layout for the scales. + typename LayoutScale, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout = layout::NoPermute, + /// Permute operand A + typename PermuteALayout = layout::NoPermute, + /// Permute operand B + typename PermuteBLayout = layout::NoPermute, + /// + typename Enable = void> +struct GemmBlockwise; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ada Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout> +struct GemmBlockwise { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename cutlass::gemm::threadblock::DefaultMmaBlockwise< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, float, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm89, + ThreadblockShape, WarpShape, InstructionShape, Stages, Operator, false, + SharedMemoryClear, GatherA, GatherB, PermuteALayout, + PermuteBLayout>::ThreadblockMma; + + static_assert(ThreadblockShape::kK % WarpShape::kK == 0, "ThreadblockShape::kK must be divisible by WarpShape::kK."); + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, + EpilogueOutputOp, EpilogueOutputOp::kCount, ScatterD, + PermuteDLayout>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = + kernel::GemmUniversalBlockwise; +}; +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace cutlass diff --git a/include/cutlass/gemm/kernel/gemm_universal_blockwise.h b/include/cutlass/gemm/kernel/gemm_universal_blockwise.h new file mode 100644 index 00000000..72a589f8 --- /dev/null +++ b/include/cutlass/gemm/kernel/gemm_universal_blockwise.h @@ -0,0 +1,359 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Template for a pipelined GEMM kernel. Does not compute batching or + support split-K. +*/ + +#pragma once + +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/semaphore.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct GemmUniversalBlockwise { + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kSplitKSerial = SplitKSerial; + + // Added aliases for per-tile FP8 dequantisation scale tensors + using LayoutScale = cutlass::layout::RowMajor; + using TensorRefScale = cutlass::TensorRef; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + cutlass::gemm::GemmCoord problem_size; + cutlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename OutputOp::Params output_op; + int *semaphore; + int gemm_k_size; + // For gather+scatter operations + int const *gather_A_indices; + int const *gather_B_indices; + int const *scatter_D_indices; + TensorRefScale scale_A; + TensorRefScale scale_B; + + // + // Methods + // + + CUTLASS_HOST_DEVICE + Params() : swizzle_log_tile(0), semaphore(0), gemm_k_size(0) {} + + CUTLASS_HOST_DEVICE + Params(cutlass::gemm::GemmCoord const &problem_size, + cutlass::gemm::GemmCoord const &grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + TensorRefScale scale_A, + TensorRefScale scale_B, + typename OutputOp::Params output_op = typename OutputOp::Params(), + int *workspace = nullptr, int const *gather_A_indices = nullptr, + int const *gather_B_indices = nullptr, + int const *scatter_D_indices = nullptr) + : problem_size(problem_size), grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), ref_A(ref_A), params_B(ref_B.layout()), + ref_B(ref_B), params_C(ref_C.layout()), ref_C(ref_C), + params_D(ref_D.layout()), ref_D(ref_D), scale_A(scale_A), + scale_B(scale_B), output_op(output_op), + gather_A_indices(gather_A_indices), + gather_B_indices(gather_B_indices), + scatter_D_indices(scatter_D_indices) { + int total_gemm_k_iterations = + (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + int gemm_k_iterations = + (total_gemm_k_iterations + grid_tiled_shape.k() - 1) / + grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + + semaphore = workspace; + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + CUTLASS_HOST_DEVICE + GemmUniversalBlockwise() {} + + /// Determines whether kernel satisfies alignment + CUTLASS_HOST_DEVICE + static Status + can_implement(cutlass::gemm::GemmCoord const &problem_size, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D) { + static int const kAlignmentA = + (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = + (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = + (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Epilogue::OutputTileIterator::kElementsPerAccess; + + if (!TensorRef_aligned(ref_A, kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_B, kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_C, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_D, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Executes one GEMM + CUTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + cutlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + return; + } + + // Compute initial location in logical coordinates + cutlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size, + }; + + cutlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + threadblock_tile_offset.n() * Mma::Shape::kN}; + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k = + min(params.problem_size.k(), + (threadblock_tile_offset.k() + 1) * params.gemm_k_size); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = + (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / + Mma::Shape::kK; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, params.ref_A.data(), + {params.problem_size.m(), problem_size_k}, thread_idx, tb_offset_A, + params.gather_A_indices); + + typename Mma::IteratorB iterator_B( + params.params_B, params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, thread_idx, tb_offset_B, + params.gather_B_indices); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx_sync(); + int lane_idx = threadIdx.x % 32; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + if (!kSplitKSerial || gemm_k_iterations > 0) { + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, + params.scale_A, params.scale_B); + } + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // assume identity swizzle + MatrixCoord threadblock_offset(threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * + Mma::Shape::kN); + + int block_idx = threadblock_tile_offset.m() + + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // If performing a reduction via split-K, fetch the initial synchronization + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is + // currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), + params.grid_tiled_shape.k()); + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, params.ref_C.data(), params.problem_size.mn(), + thread_idx, threadblock_offset, params.scatter_D_indices); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, params.ref_D.data(), params.problem_size.mn(), + thread_idx, threadblock_offset, params.scatter_D_indices); + + Epilogue epilogue(shared_storage.epilogue, thread_idx, warp_idx, lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator + // construction + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + // For subsequent threadblocks, the source matrix is held in the 'D' + // tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace cutlass diff --git a/include/cutlass/gemm/kernel/gemv.h b/include/cutlass/gemm/kernel/gemv.h index 4264080c..4d6b42f4 100644 --- a/include/cutlass/gemm/kernel/gemv.h +++ b/include/cutlass/gemm/kernel/gemv.h @@ -247,7 +247,7 @@ public: Status update(Arguments const &args) { output_op = args.output_op; - ref_A = ref_A; + ref_A = args.ref_A; ptr_B = args.ptr_B; ptr_C = args.ptr_C; ptr_D = args.ptr_D; @@ -480,7 +480,7 @@ public: problem_size = args.problem_size; batch_count = args.batch_count; output_op = args.output_op; - ref_A = ref_A; + ref_A = args.ref_A; ptr_B = args.ptr_B; ptr_C = args.ptr_C; ptr_D = args.ptr_D; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp index 47b1ce2c..48b1f738 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp @@ -1146,7 +1146,9 @@ public: // support fixup operations needed by split-/stream-K. These operations are pushed // to the collective layer so that they can reuse the TMEM -> RF copy performed // at the collective layer. - auto [mma2accum_pipeline_state_next] = collective_epilogue( + auto [mma2accum_pipeline_state_next, epi_load_pipe_consumer_state_next] = collective_epilogue( + epi_load_pipeline, + epi_load_pipe_consumer_state, mma2accum_pipeline, mma2accum_pipeline_consumer_state, problem_shape_MNKL, @@ -1157,6 +1159,7 @@ public: ); // Advance the mm2accum pipe mma2accum_pipeline_consumer_state = mma2accum_pipeline_state_next; + epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next; } work_tile_info = next_work_tile_info; diff --git a/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp b/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp index 7e920b2f..67853638 100644 --- a/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp +++ b/include/cutlass/gemm/kernel/sm100_gemm_tma_warpspecialized_input_transform.hpp @@ -144,8 +144,9 @@ public: static constexpr bool IsSchedDynamicPersistent = TileScheduler::IsDynamicPersistent; // Transfer registers from regular warps to Accum warps - static constexpr uint32_t GenericRegisterRequirement = 152; - static constexpr uint32_t AccumRegisterRequirement = 200; + static constexpr uint32_t GenericRegisterRequirement = 64; + static constexpr uint32_t TransformRegisterRequirement = 184; + static constexpr uint32_t AccumRegisterRequirement = 256; // Pipeline and pipeline state types using Load2TransformPipeline = typename CollectiveMainloop::Load2TransformPipeline; @@ -769,7 +770,7 @@ public: else if (is_participant.transformation) { // Register reconfiguration - arch::warpgroup_reg_dealloc(); + arch::warpgroup_reg_alloc(); // Signal the epilogue warps to proceed once the prologue is complete epilogue_throttle_barrier.arrive(); @@ -781,18 +782,19 @@ public: auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{}); auto k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info); auto k_tile_iter = cute::make_coord_iterator(idx2crd(k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl)); - auto [load2transform_pipeline_consumer_state_next, transform2mma_pipeline_producer_state_next] = collective_mainloop.transform( - load2transform_pipeline, - load2transform_pipeline_consumer_state, - transform2mma_pipeline, - transform2mma_pipeline_producer_state, - bulk_tmem, - transform_inputs, - k_tile_iter, k_tile_count - ); - transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state_next; - load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state_next; - + { + auto [load2transform_pipeline_consumer_state_next, transform2mma_pipeline_producer_state_next] = collective_mainloop.transform( + load2transform_pipeline, + load2transform_pipeline_consumer_state, + transform2mma_pipeline, + transform2mma_pipeline_producer_state, + bulk_tmem, + transform_inputs, + k_tile_iter, k_tile_count + ); + transform2mma_pipeline_producer_state = transform2mma_pipeline_producer_state_next; + load2transform_pipeline_consumer_state = load2transform_pipeline_consumer_state_next; + } // Fetch next work tile auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work( work_tile_info, @@ -950,7 +952,9 @@ public: // Wait for tmem allocation tmem_allocation_result_barrier.arrive_and_wait_unaligned(); - auto accum_inputs = collective_mainloop.accum_init(bulk_tmem, typename CollectiveEpilogue::CopyOpT2R{}, typename CollectiveEpilogue::EpilogueTile{}); + auto accum_inputs = [&]() { + return collective_mainloop.accum_init(bulk_tmem, typename CollectiveEpilogue::CopyOpT2R{}, typename CollectiveEpilogue::EpilogueTile{}); + }(); bool do_tail_store = false; do { // Fetch next work tile @@ -967,12 +971,12 @@ public: auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, CtaShape_MNK{}); if constexpr (InputTransformType == cutlass::gemm::detail::KernelInputTransformType::FastF32) { - auto [mma2accum_pipeline_consumer_state_next,tTR_rGlobAcc] = collective_mainloop.accum( - accum_inputs, - mma2accum_pipeline, - mma2accum_pipeline_consumer_state, - k_tile_count); - + auto [mma2accum_pipeline_consumer_state_next,tTR_rGlobAcc] = + collective_mainloop.accum( + accum_inputs, + mma2accum_pipeline, + mma2accum_pipeline_consumer_state, + k_tile_count); mma2accum_pipeline_consumer_state_next = scheduler.template fixup( TiledMma{}, work_tile_info, @@ -1028,7 +1032,9 @@ public: // Epilogue and write to gD // if (scheduler.compute_epilogue(work_tile_info)) { - auto [mma2accum_pipeline_state_next] = collective_epilogue( + auto [mma2accum_pipeline_state_next, epi_load_pipe_consumer_state_next] = collective_epilogue( + epi_load_pipeline, + epi_load_pipe_consumer_state, mma2accum_pipeline, mma2accum_pipeline_consumer_state, problem_shape_MNKL, @@ -1039,6 +1045,7 @@ public: ); // Advance the mm2accum pipe mma2accum_pipeline_consumer_state = mma2accum_pipeline_state_next; + epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next; } } diff --git a/include/cutlass/gemm/kernel/sm100_tile_scheduler.hpp b/include/cutlass/gemm/kernel/sm100_tile_scheduler.hpp index da63fe53..739010c3 100755 --- a/include/cutlass/gemm/kernel/sm100_tile_scheduler.hpp +++ b/include/cutlass/gemm/kernel/sm100_tile_scheduler.hpp @@ -621,7 +621,7 @@ public: , "r"(clc_response.data[1]) , "r"(clc_response.data[2]) , "r"(clc_response.data[3])); - cutlass::arch::fence_view_async_shared(); + cutlass::arch::fence_view_shared(); #endif } diff --git a/include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp b/include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp index 94dcfa65..54d81cab 100644 --- a/include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp +++ b/include/cutlass/gemm/kernel/sm100_tile_scheduler_stream_k.hpp @@ -226,7 +226,7 @@ public: PipelineState advance_to_next_work(Pipeline& clc_pipeline, PipelineState clc_pipe_producer_state) const { return sm100_scheduler_.advance_to_next_work(clc_pipeline, clc_pipe_producer_state); - } + } // Given the inputs, computes the total number of output blocks this problem will compute over template diff --git a/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp index d9127d92..c90b3e3d 100644 --- a/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp @@ -381,11 +381,19 @@ public: // more than 4 CTAs implementable &= (args.hw_info.cluster_shape.x <= 4 && args.hw_info.cluster_shape.y <= 4 && args.hw_info.cluster_shape_fallback.x <= 4 && args.hw_info.cluster_shape_fallback.y <= 4); + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n"); + return implementable; + } } else { // Special cluster check for scale factor multicasts. Due to limited size of scale factors, we can't multicast among // more than 4 CTAs implementable &= ((size<0>(ClusterShape{}) <= 4) && (size<1>(ClusterShape{}) <= 4)); + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n"); + return implementable; + } } } diff --git a/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp b/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp index ae393783..455464e6 100644 --- a/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp +++ b/include/cutlass/gemm/kernel/sm103_blockscaled_gemm_tma_warpspecialized.hpp @@ -322,11 +322,19 @@ public: // more than 4 CTAs implementable &= (args.hw_info.cluster_shape.x <= 4 && args.hw_info.cluster_shape.y <= 4 && args.hw_info.cluster_shape_fallback.x <= 4 && args.hw_info.cluster_shape_fallback.y <= 4); + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n"); + return implementable; + } } else { // Special cluster check for scale factor multicasts. Due to limited size of scale factors, we can't multicast among // more than 4 CTAs implementable &= ((size<0>(ClusterShape{}) <= 4) && (size<1>(ClusterShape{}) <= 4)); + if (!implementable) { + CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Cluster Shapes cannot be greater than 4.\n"); + return implementable; + } } return implementable; diff --git a/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp b/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp index 6d134322..1e206aa6 100644 --- a/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp +++ b/include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp @@ -328,7 +328,7 @@ public: CUTLASS_DEVICE bool is_last_tile(WorkTileInfo work_tile_info, uint32_t advance_count = 1) const { - // Never pass this by reference; it needs a copy, + // Never pass this by reference; it needs a copy, // because continue_current_work will modify it. if (continue_current_work(work_tile_info)) { return false; diff --git a/include/cutlass/gemm/threadblock/default_mma_core_sm80.h b/include/cutlass/gemm/threadblock/default_mma_core_sm80.h index d9817dca..578e5cc7 100644 --- a/include/cutlass/gemm/threadblock/default_mma_core_sm80.h +++ b/include/cutlass/gemm/threadblock/default_mma_core_sm80.h @@ -57,6 +57,7 @@ #include "cutlass/gemm/threadblock/default_mma_core.h" #include "cutlass/gemm/threadblock/default_multistage_mma_complex_core.h" #include "cutlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h" +#include "cutlass/gemm/threadblock/mma_multistage_blockwise.h" #include "cutlass/matrix_shape.h" #include "cutlass/numeric_types.h" diff --git a/include/cutlass/gemm/threadblock/default_mma_multistage_blockwise.h b/include/cutlass/gemm/threadblock/default_mma_multistage_blockwise.h new file mode 100644 index 00000000..7850ed14 --- /dev/null +++ b/include/cutlass/gemm/threadblock/default_mma_multistage_blockwise.h @@ -0,0 +1,212 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Template for a pipelined GEMM kernel for row-major output (OperatorClass TensorOp) and calls MmaMultistageBlockwise threadblock-scoped multistage matrix multiply + that supports additional scaling operands. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "cutlass/arch/arch.h" +#include "cutlass/arch/wmma.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/threadblock/default_mma_core_simt.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm70.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm75.h" +#include "cutlass/gemm/threadblock/default_mma_core_sm80.h" +#include "cutlass/layout/matrix.h" +#include "cutlass/layout/permute.h" +#include "cutlass/numeric_types.h" +#include "cutlass/transform/threadblock/predicated_tile_iterator.h" +#include "cutlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h" + +#if defined(CUTLASS_ARCH_WMMA_ENABLED) +#include "cutlass/gemm/threadblock/default_mma_core_wmma.h" +#endif // CUTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element Type for scales. + typename ElementScale, + /// Layout ytpe for scales. + typename LayoutSclae, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Permute operand A + typename PermuteALayout = layout::NoPermute, + /// Permute operand B + typename PermuteBLayout = layout::NoPermute> +struct DefaultMmaBlockwise; + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Element Type for scales. + typename ElementScale, + /// Layout ytpe for scales. + typename LayoutScale, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout> +struct DefaultMmaBlockwise< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, ElementScale, LayoutScale, arch::OpClassTensorOp, ArchTag, + ThreadblockShape, WarpShape, InstructionShape, Stages, Operator, false, + SharedMemoryClear, GatherA, GatherB, PermuteALayout, PermuteBLayout> { + static_assert(platform::is_same::value || + platform::is_same>::value, + "simt epilogue must be row major"); + + static cutlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? cutlass::arch::CacheOperation::Global + : cutlass::arch::CacheOperation::Always; + + static cutlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? cutlass::arch::CacheOperation::Global + : cutlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = cutlass::Array; + using IteratorA = + cutlass::transform::threadblock::PredicatedTileAccessIterator< + cutlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA, GatherA, + PermuteALayout>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = cutlass::Array; + using IteratorB = + cutlass::transform::threadblock::PredicatedTileAccessIterator< + cutlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB, GatherB, + PermuteBLayout>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = cutlass::gemm::threadblock::MmaMultistageBlockwise< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, LayoutC, + ElementScale, LayoutScale, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClear>; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/threadblock/mma_multistage.h b/include/cutlass/gemm/threadblock/mma_multistage.h index 5b7427b8..e9045c18 100644 --- a/include/cutlass/gemm/threadblock/mma_multistage.h +++ b/include/cutlass/gemm/threadblock/mma_multistage.h @@ -729,6 +729,9 @@ public: // Perform the MAC-iterations gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B); } + + // Expose pipeline state via alias without changing its original access level + using PublicPipeState = PipeState; }; ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/gemm/threadblock/mma_multistage_blockwise.h b/include/cutlass/gemm/threadblock/mma_multistage_blockwise.h new file mode 100644 index 00000000..9903b8b9 --- /dev/null +++ b/include/cutlass/gemm/threadblock/mma_multistage_blockwise.h @@ -0,0 +1,449 @@ +/*************************************************************************************************** + * Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Template for a double-buffered threadblock-scoped GEMM kernel that performs blockwise + scaling dequantization in the MMA for input matrices A and B. +*/ + +#pragma once + +#include "cutlass/aligned_buffer.h" +#include "cutlass/arch/memory.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/gemm/threadblock/mma_base.h" +#include "cutlass/matrix_shape.h" +#include "cutlass/numeric_types.h" +#include "cutlass/gemm/threadblock/mma_multistage.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cutlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + cutlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + cutlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Element Type for for the scalesl + typename ElementScale_, + /// Layout for the scales. + typename LayoutScale_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Used for partial specialization + typename Enable = bool> +class MmaMultistageBlockwise : public MmaMultistage { +public: + ///< Base class + using Base = MmaMultistage; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + /// Data type of scales + using ElementScale = ElementScale_; + /// Layout Type of Scales. + using LayoutScale = LayoutScale_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // Reference to the canonical MmaMultistage specialization with identical + // template arguments. This enables us to reuse its helper structures + // (Detail and PipeState) without redefining them here. + using BaseMma = MmaMultistage; + + /// Internal structure exposed for introspection (aliased from BaseMma). + using Detail = typename BaseMma::Detail; + + // Bring selected base-class helpers into scope so that calls like + // advance_smem_read_stage() resolve correctly in a dependent-name + // context where two-phase lookup would otherwise ignore the base + // class. + using Base::advance_smem_read_stage; + using Base::advance_smem_write_stage; + using Base::copy_tiles_and_advance; + using Base::prologue; + using Base::gmem_wait; + using Base::wind_down; + +private: + // Pipeline state structure reused from the canonical multistage kernel. + using PipeState = typename BaseMma::PublicPipeState; + +private: + // + // Data members + // + + /// Warp-level MMA operator + Operator warp_mma_; + +public: + /// Construct from tensor references + CUTLASS_DEVICE + MmaMultistageBlockwise( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx) + { + // All per-warp iterator adjustments are handled by the base-class + // constructor, so no additional work is required here. + } + + /// Perform a threadblock mainloop iteration of matrix multiply-accumulate + CUTLASS_DEVICE + void mac_loop_iter( + PipeState &pipe_state, ///< [in|out] loop-carried pipeline state + FragmentC &accum, ///< [in|out] destination accumulator tile + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B, ///< [in|out] iterator over B operand in global memory + int &gemm_k_iterations, + cutlass::TensorRef scale_A, // blockwise scale tensor for A + cutlass::TensorRef scale_B, // blockwise scale tensor for B + int k_iter_idx, ///< current K-block index processed by this iteration + int block_m_idx, ///< threadblock index along M dimension (row) + int block_n_idx) ///< threadblock index along N dimension (col) + { + // Unroll the warp-level MMA tiles of a threadblock's mainloop iteration + CUTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load the next warp-tile's A fragment from shared memory + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.load(pipe_state.warp_loaded_frag_A_[(warp_mma_k + 1) % 2]); + ++this->warp_tile_iterator_A_; + + // Load the next warp-tile's B fragment from shared memory + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_[(warp_mma_k + 1) % 2]); + ++this->warp_tile_iterator_B_; + + // Except for the first warp-tile, all warp-tiles convert their incoming shared memory fragments as necessary + if (warp_mma_k > 0) { + warp_mma_.transform(pipe_state.warp_transformed_frag_A_[warp_mma_k % 2], + pipe_state.warp_transformed_frag_B_[warp_mma_k % 2], + pipe_state.warp_loaded_frag_A_[warp_mma_k % 2], + pipe_state.warp_loaded_frag_B_[warp_mma_k % 2]); + } + + // Compute scale factor for the current (M-tile, N-tile, K-tile) triple. + // The K-tile index used for scaling must not exceed the allocated range + // of the scale tensors. This situation can arise in the prologue / + // epilogue iterations of the multistage pipeline when the software + // pipeline executes Stages-1 extra iterations with gemm_k_iterations < 0. + + int ldA = int(scale_A.layout().stride(0)); + int k_block_idx = k_iter_idx; + if (k_block_idx >= ldA) { + k_block_idx = ldA - 1; + } + + float scale_factor = scale_A.at({block_m_idx, k_block_idx}) * + scale_B.at({block_n_idx, k_block_idx}); + + // Perform MMA into a temporary fragment (unscaled) + FragmentC delta; + FragmentC zero_frag; + zero_frag.clear(); + + warp_mma_(delta, pipe_state.warp_transformed_frag_A_[warp_mma_k % 2], + pipe_state.warp_transformed_frag_B_[warp_mma_k % 2], zero_frag); + + // Apply dequantization scaling + CUTLASS_PRAGMA_UNROLL + for (int el = 0; el < FragmentC::kElements; ++el) { + delta[el] *= scale_factor; + } + + // Accumulate the scaled contribution + plus plus_accum; + + if (Detail::kStagedAccumulation) { + pipe_state.tmp_accum_ = plus_accum(pipe_state.tmp_accum_, delta); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, pipe_state.tmp_accum_); + pipe_state.tmp_accum_.clear(); + } + } else { + accum = plus_accum(accum, delta); + } + + // Except for the last warp-tile, all warp-tiles issue their share of + // global->shared fragment copies + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + } + + // The second-to-last warp-tile also: + // - performs the last warp-tile's share of global->shared fragment + // copies + // - moves to the next global fetch stage + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Performs the last warp-tile's share of global->shared fragment copies + int group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + int group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + + // Inserts a memory fence between stages of cp.async instructions. + cutlass::arch::cp_async_fence(); + + // Wait until we have at least one completed global fetch stage + gmem_wait(); + + // Move to the next global fetch stage + advance_smem_write_stage(iterator_A, iterator_B); + advance_smem_read_stage(); + + // Disable global fetching when done with global fetch iterations + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // The last warp-tile also converts the shared memory fragments used by + // the first warp-tile of the next iteration, if necessary (so we can + // immediately start issuing MMA instructions at the top of the loop ) + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma_.transform( + pipe_state.warp_transformed_frag_A_[(warp_mma_k + 1) % 2], + pipe_state.warp_transformed_frag_B_[(warp_mma_k + 1) % 2], + pipe_state.warp_loaded_frag_A_[(warp_mma_k + 1) % 2], + pipe_state.warp_loaded_frag_B_[(warp_mma_k + 1) % 2]); + } + } + } + + /// Perform the specified number of threadblock mainloop iterations of matrix + /// multiply-accumulate. Assumes prologue has been initiated. + CUTLASS_DEVICE + void gemm_iters( + int gemm_k_iterations, ///< number of threadblock mainloop iterations + FragmentC &accum, ///< [in|out] accumulator tile + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B, + cutlass::TensorRef scale_A, // blockwise scale tensor for A + cutlass::TensorRef scale_B, // blockwise scale tensor for B + int block_m_idx, + int block_n_idx) ///< [in|out] iterator over B operand in global memory + { + PipeState pipe_state; + + // Disable global fetching if done with global fetch iterations + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + // Load first warp-tile's A fragment from shared memory + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_A_.load(pipe_state.warp_loaded_frag_A_[0]); + ++this->warp_tile_iterator_A_; + + // Load first warp-tile's B fragment from shared memory + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_[0]); + ++this->warp_tile_iterator_B_; + + // Transform, if necessary, the first warp-tile's shared memory fragments + warp_mma_.transform(pipe_state.warp_transformed_frag_A_[0], + pipe_state.warp_transformed_frag_B_[0], + pipe_state.warp_loaded_frag_A_[0], + pipe_state.warp_loaded_frag_B_[0]); + + if (Detail::kStagedAccumulation) { + pipe_state.tmp_accum_.clear(); + } + + // Mainloop + int k_iter_idx = 0; + + CUTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1); ++k_iter_idx) { + mac_loop_iter(pipe_state, accum, iterator_A, iterator_B, + gemm_k_iterations, scale_A, scale_B, k_iter_idx, + block_m_idx, block_n_idx); + } + + if (Detail::kStagedAccumulation) { + plus plus_accum; + accum = plus_accum(accum, pipe_state.tmp_accum_); + } + + // Commit and drain all pending and predicated cp.async pnz from the GEMM + // mainloop + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + __syncthreads(); + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + CUTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const &src_accum, + cutlass::TensorRef scaleA, + cutlass::TensorRef scaleB) { + // Each scale element corresponds to a 128x128 tile along (M, K) for A and + // (N, K) for B. Grid dimension X enumerates threadblock tiles along M and + // grid dimension Y along N when GemmIdentityThreadblockSwizzle is used with + // the default N = 1 (tile = 1). Therefore, + // blockIdx.x -> tile index along the M dimension + // blockIdx.y -> tile index along the N dimension. + + constexpr int kScaleBlock = 128; + // Row-wise block index for A (and output C/D) – one per 128 rows. + int block_m_idx = (blockIdx.x * Shape::kM) / kScaleBlock; + + // Column-wise block index for B – one per 128 columns. Note that each + // threadblock processes Shape::kN columns, which may be < 128 (64 in this + // kernel). We therefore map two consecutive threadblock tiles onto the + // same 128-wide scale block when Shape::kN < kScaleBlock. + int block_n_idx = (blockIdx.y * Shape::kN) / kScaleBlock; + + // Prologue (start fetching iterations of global fragments into shared + // memory) + prologue(iterator_A, iterator_B, gemm_k_iterations); + + // Wait until we have at least one completed global fetch stage + gmem_wait(); + + // Initialize destination accumulators with source accumulators + accum = src_accum; + + // Perform the MAC-iterations with blockwise dequantization + gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B, scaleA, scaleB, + block_m_idx, block_n_idx); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace cutlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/include/cutlass/layout/permute.h b/include/cutlass/layout/permute.h index 3464d7f5..b311010e 100644 --- a/include/cutlass/layout/permute.h +++ b/include/cutlass/layout/permute.h @@ -39,7 +39,9 @@ */ #pragma once #include "cutlass/cutlass.h" +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/fast_math.h" #include "cutlass/layout/pitch_linear.h" #include "cutlass/layout/matrix.h" diff --git a/include/cutlass/numeric_conversion.h b/include/cutlass/numeric_conversion.h index 42f3c912..e9463b20 100644 --- a/include/cutlass/numeric_conversion.h +++ b/include/cutlass/numeric_conversion.h @@ -3678,30 +3678,537 @@ template < struct NumericArrayConverter : public PackedNumericArrayConverter {}; +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Helpers for PTX conversion with inline assembly (used in the following partial specializations) +// Specifically, we use lookup tables for converting e2m1 to bf16, half and float_e4m3 and float_e5m2 +// +///////////////////////////////////////////////////////////////////////////////////////////////// + + +#if defined(CUDA_PTX_FP4FP6_CVT_ENABLED) && __CUDA_ARCH__ >= 1100 + #define USE_PTX_CONVERT 1 +#endif + + + +namespace detail { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// E2M1 Conversion Helper Functions to BF16 +// +///////////////////////////////////////////////////////////////////////////////////////////////// + + /**************************************************************************** + E2M1 to BF16 Conversion Table: + +-----------------+-------------------------------+---------------+ + | E2M1 Pattern | BF16 Pattern | Numeric Value | + | (s E E M) | (s EEEEEEEE MMMMMMM) | | + +-----------------+-------------------------------+---------------+ + | 0 00 0 | 0 00000000 0000000 | 0.0 | + | 0 00 1 | 0 01111110 0000000 | 0.5 | + | 0 01 0 | 0 01111111 0000000 | 1.0 | + | 0 01 1 | 0 01111111 1000000 | 1.5 | + | 0 10 0 | 0 10000000 0000000 | 2.0 | + | 0 10 1 | 0 10000000 1000000 | 3.0 | + | 0 11 0 | 0 10000001 0000000 | 4.0 | + | 0 11 1 | 0 10000001 1000000 | 6.0 | + +-----------------+-------------------------------+---------------+ + + bits 2-9 go into a LUT, the top 2 bits are inserted from E2M1 value (s E E M) + + +-----------------+-------------------+-----------+ + | E2M1 Pattern | LUT Value (8 bits)| Hex Value | + | (s E E M) | | | + +-----------------+-------------------+-----------+ + | 0 00 0 | 00000000 | 0x00 | + | 0 00 1 | 11111100 | 0xFC | + | 0 01 0 | 11111110 | 0xFE | + | 0 01 1 | 11111111 | 0xFF | + | 0 10 0 | 00000000 | 0x00 | + | 0 10 1 | 00000001 | 0x01 | + | 0 11 0 | 00000010 | 0x02 | + | 0 11 1 | 00000011 | 0x03 | + +-----------------+-------------------+-----------+ + + constexpr unsigned long long E2M1_to_BF16_LUT = 0x03020100FFFEFC00ULL; + constexpr unsigned int E2M1_to_BF16_UPPER_LUT = 0xc0804000U; + + ****************************************************************************/ + + // LUT _e2m1_to_bf16_x2: Direct E2M1->BF16 (converts 2 E2M1 to 2 BF16) + CUTLASS_DEVICE + void _e2m1_to_bf16_x2(unsigned int src, unsigned int& out0) { + constexpr unsigned long long lut = 0x03020100FFFEFC00ULL; + constexpr unsigned int upper_lut = 0xc0804000U; + constexpr unsigned int lut_lo = (unsigned int)lut; + constexpr unsigned int lut_hi = (unsigned int)(lut >> 32); + + asm( + "{\n\t" + ".reg .u32 prmt_ctrl01, target01, upper01;\n\t" + ".reg .u32 target0_1_, upper_0_1;\n\t" + ".reg .u32 upper_prmt_ctrl_01;\n\t" + "and.b32 prmt_ctrl01, %1, 0x0077;\n\t" + "shr.b32 upper_prmt_ctrl_01, %1, 2;\n\t" + "and.b32 upper_prmt_ctrl_01, upper_prmt_ctrl_01, 0x0033;\n\t" + "prmt.b32 target01, %2, %3, prmt_ctrl01;\n\t" + "prmt.b32 target0_1_, target01, target01, 0x3120;\n\t" // both 2 and 3 are known to be zero + "prmt.b32 upper01, %4, %4, upper_prmt_ctrl_01;\n\t" + "prmt.b32 upper_0_1, upper01, upper01, 0x1302;\n\t" // both 2 and 3 are known to be zero + "shl.b32 target0_1_, target0_1_, 6;\n\t" + "or.b32 %0, target0_1_, upper_0_1;\n\t" + "}" + : "=r"(out0) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) // %1, %2, %3, %4 + ); + } + + // LUT x4 _e2m1_to_bf16_x4: Direct E2M1->BF16 (converts 4 E2M1 to 4 BF16 in one call) + CUTLASS_DEVICE + void _e2m1_to_bf16_x4(unsigned int src, unsigned int& out0, unsigned int& out1, unsigned int shift_count=0) { + constexpr unsigned long long lut = 0x03020100FFFEFC00ULL; // bit2 - 9 for e2m1 -> bf16 conversion + constexpr unsigned int upper_lut = 0xc0804000U; // bit0, bot1 for e2m1 -> bf16 conversion + constexpr unsigned int lut_lo = (unsigned int)lut; + constexpr unsigned int lut_hi = (unsigned int)(lut >> 32); + + + if (shift_count == 0) { + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 target0_1_, target2_3_, upper_0_1, upper_2_3;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + "" + "and.b32 prmt_ctrl0123, %2, 0x7777;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %2, 2;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x3333;\n\t" + "prmt.b32 target0123, %3, %4, prmt_ctrl0123;\n\t" + "prmt.b32 target0_1_, target0123, %4, 0x4140;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 target2_3_, target0123, %4, 0x4342;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper0123, %5, %5, upper_prmt_ctrl_0123;\n\t" + "prmt.b32 upper_0_1, upper0123, %4, 0x1404;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper_2_3, upper0123, %4, 0x3424;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "shl.b32 target0_1_, target0_1_, 6;\n\t" + "shl.b32 target2_3_, target2_3_, 6;\n\t" + "or.b32 %0, target0_1_, upper_0_1;\n\t" + "or.b32 %1, target2_3_, upper_2_3;\n\t" + "}" + : "=r"(out0), "=r"(out1) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) // %2, %3, %4, %5 + ); + } else if (shift_count == 16) { + // protect against future changes to the shift_count + assert(shift_count==16 && "shift_count should be 0 or 16"); + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 target0_1_, target2_3_, upper_0_1, upper_2_3;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + "" + "shr.b32 prmt_ctrl0123, %2, 16;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %2, 18;\n\t" + "and.b32 prmt_ctrl0123, prmt_ctrl0123, 0x7777;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x3333;\n\t" + "prmt.b32 target0123, %3, %4, prmt_ctrl0123;\n\t" + "prmt.b32 target0_1_, target0123, %4, 0x4140;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 target2_3_, target0123, %4, 0x4342;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper0123, %5, %5, upper_prmt_ctrl_0123;\n\t" + "prmt.b32 upper_0_1, upper0123, %4, 0x1404;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper_2_3, upper0123, %4, 0x3424;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "shl.b32 target0_1_, target0_1_, 6;\n\t" + "shl.b32 target2_3_, target2_3_, 6;\n\t" + "or.b32 %0, target0_1_, upper_0_1;\n\t" + "or.b32 %1, target2_3_, upper_2_3;\n\t" + "}" + : "=r"(out0), "=r"(out1) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) // %2, %3, %4, %5 + ); + } else { + assert((shift_count==0 || shift_count==16) && "shift_count should be 0 or 16"); + } + } + + // LUT x8 _e2m1_to_bf16_x8: Direct E2M1->BF16 (converts 8 E2M1 to 8 BF16 in one call) + CUTLASS_DEVICE + void _e2m1_to_bf16_x8(unsigned int src, unsigned int& out0, unsigned int& out1, unsigned int& out2, unsigned int& out3) { + _e2m1_to_bf16_x4(src, out0, out1, 0); + _e2m1_to_bf16_x4(src, out2, out3, 16); + } + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// E2M1 Conversion Helper Functions to FP16 +// +///////////////////////////////////////////////////////////////////////////////////////////////// + + + /**************************************************************************** + E2M1 to FP16 Conversion Table: + +-----------------+-----------------------------------+---------------+---------------+ + | E2M1 Pattern | FP16 Pattern | MS Byte (Hex) | Numeric Value | + | (s E E M) | (s EEEEE MMMMMMMMMM) | (s EEEEE MM) | | + +-----------------+-----------------------------------+---------------+---------------+ + | 0 00 0 | 0 00000 0000000000 | 0x00 | 0.0 | + | 0 00 1 | 0 01110 0000000000 | 0x38 | 0.5 | + | 0 01 0 | 0 01111 0000000000 | 0x3C | 1.0 | + | 0 01 1 | 0 01111 1000000000 | 0x3E | 1.5 | + | 0 10 0 | 0 10000 0000000000 | 0x40 | 2.0 | + | 0 10 1 | 0 10000 1000000000 | 0x42 | 3.0 | + | 0 11 0 | 0 10001 0000000000 | 0x44 | 4.0 | + | 0 11 1 | 0 10001 1000000000 | 0x46 | 6.0 | + +-----------------+-----------------------------------+---------------+---------------+ + + constexpr unsigned long long E2M1_to_FP16_LUT = 0x464442403E3C3800ULL; + constexpr unsigned int E2M1_to_FP16_UPPER_LUT = 0xc0804000U; + + ****************************************************************************/ + + // LUT _e2m1_to_half_x2: Direct E2M1->FP16 (converts 2 E2M1 to 2 FP16) + CUTLASS_DEVICE + void _e2m1_to_half_x2(unsigned int src, unsigned int& out0) { + constexpr unsigned long long lut = 0x464442403E3C3800ULL; + constexpr unsigned int upper_lut = 0xc0804000U; + constexpr unsigned int lut_lo = (unsigned int)lut; + constexpr unsigned int lut_hi = (unsigned int)(lut >> 32); + + asm( + "{\n\t" + ".reg .u32 prmt_ctrl01, target01, upper01;\n\t" + ".reg .u32 target0_1_, upper_0_1;\n\t" + ".reg .u32 upper_prmt_ctrl_01;\n\t" + "and.b32 prmt_ctrl01, %1, 0x0077;\n\t" + "shr.b32 upper_prmt_ctrl_01, %1, 2;\n\t" + "and.b32 upper_prmt_ctrl_01, upper_prmt_ctrl_01, 0x0033;\n\t" + "prmt.b32 target01, %2, %3, prmt_ctrl01;\n\t" + "prmt.b32 target0_1_, target01, target01, 0x1302;\n\t" // both 2 and 3 are known to be zero + "prmt.b32 upper01, %4, %4, upper_prmt_ctrl_01;\n\t" + "prmt.b32 upper_0_1, upper01, upper01, 0x1302;\n\t" // both 2 and 3 are known to be zero + "or.b32 %0, target0_1_, upper_0_1;\n\t" + "}" + : "=r"(out0) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) // %2, %3, %4, %5 + ); + } + + // LUT x4 _e2m1_to_half_x4: Direct E2M1->FP16 (converts 4 E2M1 to 4 FP16 in one call) + CUTLASS_DEVICE + void _e2m1_to_half_x4(unsigned int src, unsigned int& out0, unsigned int& out1, unsigned int shift_count=0) { + constexpr unsigned long long lut = 0x464442403E3C3800ULL; // bit0 - 7 for e2m1 -> half conversion + constexpr unsigned int upper_lut = 0xc0804000U; // bit0, bit1 for e2m1 -> FP16 conversion + constexpr unsigned int lut_lo = (unsigned int)lut; + constexpr unsigned int lut_hi = (unsigned int)(lut >> 32); + + assert((shift_count==0 || shift_count==16) && "shift_count should be 0 or 16"); + + if (shift_count == 0) { + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 target0_1_, target2_3_, upper_0_1, upper_2_3;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + "" + "and.b32 prmt_ctrl0123, %2, 0x7777;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %2, 2;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x3333;\n\t" + "prmt.b32 target0123, %3, %4, prmt_ctrl0123;\n\t" + "prmt.b32 target0_1_, target0123, %5, 0x1404;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 target2_3_, target0123, %5, 0x3424;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper0123, %5, %5, upper_prmt_ctrl_0123;\n\t" + "prmt.b32 upper_0_1, upper0123, %5, 0x1404;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper_2_3, upper0123, %5, 0x3424;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "or.b32 %0, target0_1_, upper_0_1;\n\t" + "or.b32 %1, target2_3_, upper_2_3;\n\t" + "}" + : "=r"(out0), "=r"(out1) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) // %2, %3, %4, %5 + ); + } else /* shift_count == 16 */ { + // protect against future changes to the shift_count + assert(shift_count==16 && "shift_count should be 0 or 16"); + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 target0_1_, target2_3_, upper_0_1, upper_2_3;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + "" + "shr.b32 prmt_ctrl0123, %2, 16;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %2, 18;\n\t" + "and.b32 prmt_ctrl0123, prmt_ctrl0123, 0x7777;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x3333;\n\t" + "prmt.b32 target0123, %3, %4, prmt_ctrl0123;\n\t" + "prmt.b32 target0_1_, target0123, %5, 0x1404;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 target2_3_, target0123, %5, 0x3424;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper0123, %5, %5, upper_prmt_ctrl_0123;\n\t" + "prmt.b32 upper_0_1, upper0123, %5, 0x1404;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "prmt.b32 upper_2_3, upper0123, %5, 0x3424;\n\t" // 4 is the low order byte of upper_lut which is 0x00 + "or.b32 %0, target0_1_, upper_0_1;\n\t" + "or.b32 %1, target2_3_, upper_2_3;\n\t" + "}" + : "=r"(out0), "=r"(out1) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) // %2, %3, %4, %5 + ); + } + } + + // LUT x8 _e2m1_to_half_x8: Direct E2M1->FP16 (converts 8 E2M1 to 8 FP16 in one call) + CUTLASS_DEVICE + void _e2m1_to_half_x8(unsigned int src, unsigned int& out0, unsigned int& out1, unsigned int& out2, unsigned int& out3) { + _e2m1_to_half_x4(src, out0, out1, 0); + _e2m1_to_half_x4(src, out2, out3, 16); + } + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// E2M1 -> FP8 (E4M3/E5M2) Shared Helper Functions +// +///////////////////////////////////////////////////////////////////////////////////////////////// + + /**************************************************************************** + E2M1 to E4M3 (FP8) Conversion Table + constexpr unsigned long long E2M1_to_E4M3_LUT_0_7 = 0x4C4844403C383000ULL; + + +-----------------+--------------------------+---------------+-------------+ + | E2M1 Pattern | E4M3 Pattern | Numeric Value | E4M3 Hex | + | (s E E M) | (s EEEE MMM) | | | + +-----------------+--------------------------+---------------+-------------+ + | 0 00 0 | 0 0000 000 | 0.0 | 0x00 | + | 0 00 1 | 0 0110 000 | 0.5 | 0x30 | + | 0 01 0 | 0 0111 000 | 1.0 | 0x38 | + | 0 01 1 | 0 0111 100 | 1.5 | 0x3C | + | 0 10 0 | 0 1000 000 | 2.0 | 0x40 | + | 0 10 1 | 0 1000 100 | 3.0 | 0x44 | + | 0 11 0 | 0 1001 000 | 4.0 | 0x48 | + | 0 11 1 | 0 1001 100 | 6.0 | 0x4C | + +-----------------+--------------------------+---------------+-------------+ + + E2M1 to E5M2 (FP8) Conversion Table + constexpr unsigned long long E2M1_to_E5M2_LUT_0_7 = 0x464442403E3C3800ULL; + + +-----------------+--------------------------+---------------+-------------+ + | E2M1 Pattern | E5M2 Pattern | Numeric Value | E5M2 Hex | + | (s E E M) | (s EEEEE MM) | | | + +-----------------+--------------------------+---------------+-------------+ + | 0 00 0 | 0 00000 00 | 0.0 | 0x00 | + | 0 00 1 | 0 01110 00 | 0.5 | 0x38 | + | 0 01 0 | 0 01111 00 | 1.0 | 0x3C | + | 0 01 1 | 0 01111 10 | 1.5 | 0x3E | + | 0 10 0 | 0 10000 00 | 2.0 | 0x40 | + | 0 10 1 | 0 10000 10 | 3.0 | 0x42 | + | 0 11 0 | 0 10001 00 | 4.0 | 0x44 | + | 0 11 1 | 0 10001 10 | 6.0 | 0x46 | + +-----------------+--------------------------+---------------+-------------+ + ****************************************************************************/ + + // LUT _e2m1_to_fp8_x4: Direct E2M1->FP8 (converts 4 E2M1 to 4 FP8) + template + CUTLASS_DEVICE + static void _e2m1_to_fp8_x4(unsigned int src, unsigned int& out0) { + constexpr unsigned int upper_lut = 0xc0804000U; + constexpr unsigned int lut_lo = (unsigned int)fp8_lut; + constexpr unsigned int lut_hi = (unsigned int)(fp8_lut >> 32); + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + "and.b32 prmt_ctrl0123, %1, 0x7777;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %1, 2;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x3333;\n\t" + "prmt.b32 target0123, %2, %3, prmt_ctrl0123;\n\t" + "prmt.b32 upper0123, %4, %4, upper_prmt_ctrl_0123;\n\t" + "or.b32 %0, target0123, upper0123;\n\t" + "}" + : "=r"(out0) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) + ); + } + + // LUT _e2m1_to_fp8_x2: Direct E2M1->FP8 (converts 2 E2M1 to 2 FP8) + template + CUTLASS_DEVICE + static void _e2m1_to_fp8_x2(unsigned int src, unsigned int& out0) { + constexpr unsigned int upper_lut = 0xc0804000U; + constexpr unsigned int lut_lo = (unsigned int)fp8_lut; + constexpr unsigned int lut_hi = (unsigned int)(fp8_lut >> 32); + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + "and.b32 prmt_ctrl0123, %1, 0x0077;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %1, 2;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x0033;\n\t" + "prmt.b32 target0123, %2, %3, prmt_ctrl0123;\n\t" + "prmt.b32 upper0123, %4, %4, upper_prmt_ctrl_0123;\n\t" + "or.b32 %0, target0123, upper0123;\n\t" // we wiped out the upper half when computing controls + "}" + : "=r"(out0) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) + ); + } + + // LUT _e2m1_to_fp8_x8: Direct E2M1->FP8 (converts 8 E2M1 to 8 FP8) + template + CUTLASS_DEVICE + static void _e2m1_to_fp8_x8(unsigned int src, unsigned int& out0, unsigned int& out1) { + constexpr unsigned int upper_lut = 0xc0804000U; + constexpr unsigned int lut_lo = (unsigned int)fp8_lut; + constexpr unsigned int lut_hi = (unsigned int)(fp8_lut >> 32); + asm( + "{\n\t" + ".reg .u32 prmt_ctrl0123, target0123, upper0123;\n\t" + ".reg .u32 upper_prmt_ctrl_0123;\n\t" + ".reg .u32 prmt_ctrl4567, target4567, upper4567;\n\t" + ".reg .u32 upper_prmt_ctrl_4567;\n\t" + "and.b32 prmt_ctrl0123, %2, 0x7777;\n\t" + "shr.b32 upper_prmt_ctrl_0123, %2, 2;\n\t" + "shr.b32 prmt_ctrl4567, %2, 16;\n\t" + "shr.b32 upper_prmt_ctrl_4567, %2, 18;\n\t" + "and.b32 prmt_ctrl4567, prmt_ctrl4567, 0x7777;\n\t" + "and.b32 upper_prmt_ctrl_0123, upper_prmt_ctrl_0123, 0x3333;\n\t" + "and.b32 upper_prmt_ctrl_4567, upper_prmt_ctrl_4567, 0x3333;\n\t" + "prmt.b32 target0123, %3, %4, prmt_ctrl0123;\n\t" + "prmt.b32 target4567, %3, %4, prmt_ctrl4567;\n\t" + "prmt.b32 upper0123, %5, %5, upper_prmt_ctrl_0123;\n\t" + "prmt.b32 upper4567, %5, %5, upper_prmt_ctrl_4567;\n\t" + "or.b32 %0, target0123, upper0123;\n\t" + "or.b32 %1, target4567, upper4567;\n\t" + "}" + : "=r"(out0), "=r"(out1) + : "r"(src), "r"(lut_lo), "r"(lut_hi), "r"(upper_lut) + ); + } + +} // namespace detail + +namespace detail { + + /* + A helper class that can vectorize a numeric converter with implementation for several vector widths. + + The vector widths must be giving in decreasing order or width, and must be a power of 2. + + The vector converters must produce identical results to the scalar converters for consistency. + */ + class VectorizedConverter { + private: + // Base case to handle remainder elements as scalars. + template + CUTLASS_DEVICE + static void convert_helper( + typename ArrayConverter::result_type& result, + typename ArrayConverter::source_type const& source) { + + using ElementRes = typename ArrayConverter::result_type::Element; + using ElementSrc = typename ArrayConverter::source_type::Element; + // If no more converters, handle the remaining elements as scalars. + constexpr int total_elements = ArrayConverter::result_type::kElements; + constexpr int remainder = total_elements - Offset; + static_assert(remainder == (total_elements % ParentWidth), "Unexpected remainder."); + + typename ArrayConverter::ScalarConverter scalar_converter; + CUTLASS_PRAGMA_UNROLL + for (int i = Offset; i < ArrayConverter::result_type::kElements; ++i) { + result[i] = scalar_converter(ElementSrc(source[i])); + } + } + + template + CUTLASS_DEVICE + static void convert_helper(typename ArrayConverter::result_type& result, typename ArrayConverter::source_type const& source) { + static_assert(sizeof...(OtherVectorArrays) % 2 == 0, "Vector converters must come in {dst, src} pairs"); + static_assert(ResultVectorArray::kElements == SourceVectorArray::kElements, "Vector converters must have the same vector width"); + static_assert(cutlass::platform::is_same::value, + "ResultVectorArray must have the same type ArrayConverter::result_type"); + static_assert(cutlass::platform::is_same::value, + "SourceVectorArray must have the same type ArrayConverter::result_type"); + static_assert(Offset >= 0 && Offset <= ArrayConverter::result_type::kElements, "Offset must be between 0 and N"); + + static_assert(ParentWidth == 0 || ParentWidth > ResultVectorArray::kElements, "Vector arrays must be given in decreasing order of width"); + + constexpr int vector_width = ResultVectorArray::kElements; + static_assert(ispow2(vector_width), "Vector width must be a power of 2"); + + using ElementRes = typename ArrayConverter::result_type::Element; + using ElementSrc = typename ArrayConverter::source_type::Element; + + constexpr int vector_bits_res = vector_width * cutlass::sizeof_bits::value; + constexpr int vector_bits_src = vector_width * cutlass::sizeof_bits::value; + + static_assert(vector_bits_res % 8 == 0, "Result vector type must be byte addressed."); + static_assert(vector_bits_src % 8 == 0, "Source vector type must be byte addressed."); + + constexpr int vector_offset = Offset / vector_width; + ResultVectorArray* packed_result_vec = reinterpret_cast(&result) + vector_offset; + SourceVectorArray const* packed_source_vec = reinterpret_cast(&source) + vector_offset; + + // Convert the remaining elements as vectors. + constexpr int total_elements = ArrayConverter::result_type::kElements; + constexpr int groups_of_vec = (total_elements - Offset) / vector_width; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < groups_of_vec; ++i) { + packed_result_vec[i] = ArrayConverter::template packed_convert(packed_source_vec[i]); + } + + constexpr int new_offset = Offset + vector_width * groups_of_vec; + // Recurse to handle other vector converters, or the scalar base case. + convert_helper(result, source); + } + + public: + /* + A method to convert vectors of elements using the packed_convert method of the converter. + + Converters using this class must implement packed convert and support 1 or more vector conversions. + */ + template + CUTLASS_DEVICE + static void convert(typename ArrayConverter::result_type& result, typename ArrayConverter::source_type const& source) { + convert_helper<0, 0, ArrayConverter, ResultVectorArray, SourceVectorArray, OtherVectorArrays...>(result, source); + } + }; +} + + ///////////////////////////////////////////////////////////////////////////////////////////////// // // Partial specializations for Array <=> Array // ///////////////////////////////////////////////////////////////////////////////////////////////// -/// Partial specialization for Array <= Array +/// Partial specialization for Array <= Array template < + int N, FloatRoundStyle Round > -struct NumericArrayConverter { +struct NumericArrayConverter { using result_element = float; using source_element = cutlass::float_e2m1_t; - - using result_type = Array; - using source_type = Array; + using result_type = Array; + using source_type = Array; static FloatRoundStyle const round_style = Round; - CUTLASS_DEVICE - static result_type convert(source_type const & source) { +private: + using result_type_packed_8 = Array; + using result_type_packed_4 = Array; + using result_type_packed_2 = Array; + using source_type_packed_8 = Array; + using source_type_packed_4 = Array; + using source_type_packed_2 = Array; - #if defined(CUDA_PTX_FP4FP6_CVT_ENABLED) - uint32_t out_fp16[4]; - uint32_t const& src_packed = reinterpret_cast(source); + using ScalarConverter = NumericConverter; + + // Convert 8 elements: E2M1 -> BF16 -> FP32 + CUTLASS_DEVICE + static result_type_packed_8 lut_convert(source_type_packed_8 const &source) { + result_type_packed_8 result; + uint32_t src_packed = *reinterpret_cast(source.data()); + + #if defined(USE_PTX_CONVERT) + + float* out_float = reinterpret_cast(result.data()); + uint32_t halfx2_01, halfx2_23, halfx2_45, halfx2_67; asm volatile( \ "{\n" \ @@ -3711,29 +4218,157 @@ struct NumericArrayConverter { "cvt.rn.f16x2.e2m1x2 %1, byte1;\n" \ "cvt.rn.f16x2.e2m1x2 %2, byte2;\n" \ "cvt.rn.f16x2.e2m1x2 %3, byte3;\n" \ - "}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) , "=r"(out_fp16[2]), "=r"(out_fp16[3]): "r"(src_packed)); + "}\n" + : "=r"(halfx2_01), "=r"(halfx2_23) , "=r"(halfx2_45), "=r"(halfx2_67): "r"(src_packed) + ); - float2 res0 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[0])); - float2 res1 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[1])); - float2 res2 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[2])); - float2 res3 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[3])); + float2 floatx2_01 = __half22float2(reinterpret_cast<__half2 &>(halfx2_01)); + float2 floatx2_23 = __half22float2(reinterpret_cast<__half2 &>(halfx2_23)); + float2 floatx2_45 = __half22float2(reinterpret_cast<__half2 &>(halfx2_45)); + float2 floatx2_67 = __half22float2(reinterpret_cast<__half2 &>(halfx2_67)); - result_type out; - out[0] = res0.x; - out[1] = res0.y; - out[2] = res1.x; - out[3] = res1.y; - out[4] = res2.x; - out[5] = res2.y; - out[6] = res3.x; - out[7] = res3.y; - return out; + out_float[0] = floatx2_01.x; + out_float[1] = floatx2_01.y; + out_float[2] = floatx2_23.x; + out_float[3] = floatx2_23.y; + out_float[4] = floatx2_45.x; + out_float[5] = floatx2_45.y; + out_float[6] = floatx2_67.x; + out_float[7] = floatx2_67.y; + + #else + + // Convert all 8 E2M1 values to BF16, then extend to FP32 + uint4* out_floatx4 = reinterpret_cast(result.data()); + unsigned int bfloatx2_01, bfloatx2_23, bfloatx2_45, bfloatx2_67; + detail::_e2m1_to_bf16_x8(src_packed, bfloatx2_01, bfloatx2_23, bfloatx2_45, bfloatx2_67); + + out_floatx4[0] = make_uint4(bfloatx2_01 << 16, bfloatx2_01 & 0xFFFF0000u, + bfloatx2_23 << 16, bfloatx2_23 & 0xFFFF0000u); + out_floatx4[1] = make_uint4(bfloatx2_45 << 16, bfloatx2_45 & 0xFFFF0000u, + bfloatx2_67 << 16, bfloatx2_67 & 0xFFFF0000u); + + #endif + + return result; + } + + // Convert 4 elements: E2M1 -> BF16 -> FP32 + CUTLASS_DEVICE + static result_type_packed_4 lut_convert(source_type_packed_4 const &source) { + result_type_packed_4 result; + uint16_t src_packed = *reinterpret_cast(source.data()); + + #if defined(USE_PTX_CONVERT) + + float* out_float = reinterpret_cast(result.data()); + uint32_t halfx2_01, halfx2_23; + + asm volatile( \ + "{\n" \ + ".reg .b8 byte0, byte1;\n" \ + "mov.b16 {byte0, byte1}, %2;\n" \ + "cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \ + "cvt.rn.f16x2.e2m1x2 %1, byte1;\n" \ + "}\n" : "=r"(halfx2_01), "=r"(halfx2_23): "h"(src_packed) + ); + + float2 floatx2_01 = __half22float2(reinterpret_cast<__half2 &>(halfx2_01)); + float2 floatx2_23 = __half22float2(reinterpret_cast<__half2 &>(halfx2_23)); + + out_float[0] = floatx2_01.x; + out_float[1] = floatx2_01.y; + out_float[2] = floatx2_23.x; + out_float[3] = floatx2_23.y; + + #else + + uint4* out_floatx4 = reinterpret_cast(result.data()); + unsigned int bfloatx2_01, bfloatx2_23; + detail::_e2m1_to_bf16_x4(src_packed, bfloatx2_01, bfloatx2_23); + + out_floatx4[0] = make_uint4(bfloatx2_01 << 16, bfloatx2_01 & 0xFFFF0000u, + bfloatx2_23 << 16, bfloatx2_23 & 0xFFFF0000u); + + #endif + + return result; + } + + // Convert 2 elements: E2M1 -> BF16 -> FP32 + CUTLASS_DEVICE + static result_type_packed_2 lut_convert(source_type_packed_2 const &source) { + result_type_packed_2 result; + uint8_t src_packed = *reinterpret_cast(source.data()); + + #if defined(USE_PTX_CONVERT) + + float* out_float = reinterpret_cast(result.data()); + uint32_t halfx2_01; + + // extend input to 16 bit since there is no 8 bit register asm constraint available for ptx inline assembly + uint16_t src_packed_16 = src_packed; + + asm volatile( \ + "{\n" \ + ".reg .b8 byte0, byte1;\n" \ + "mov.b16 {byte0, byte1}, %1;\n" \ + "cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \ + "}\n" : "=r"(halfx2_01): "h"(src_packed_16) + ); + + float2 floatx2_01 = __half22float2(reinterpret_cast<__half2 &>(halfx2_01)); + + out_float[0] = floatx2_01.x; + out_float[1] = floatx2_01.y; + + #else + + uint2* out_floatx2 = reinterpret_cast(result.data()); + unsigned int bfloatx2_01; + detail::_e2m1_to_bf16_x2(src_packed, bfloatx2_01); + + out_floatx2[0] = make_uint2(bfloatx2_01 << 16, bfloatx2_01 & 0xFFFF0000u); + + #endif + + return result; + } + + template + CUTLASS_DEVICE + static PackedResultType packed_convert(PackedSrcType const &source) { + static_assert((platform::is_same::value && + platform::is_same::value) || + (platform::is_same::value && + platform::is_same::value) || + (platform::is_same::value && + platform::is_same::value), + "Invalid PackedSrcType/PackedResultType must be 2, 4 or 8."); + + return lut_convert(source); + } + + friend class detail::VectorizedConverter; + +public: + CUTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) + result_type result; + using ConverterType = NumericArrayConverter; + detail::VectorizedConverter::convert(result, source); + + return result; #else result_type result; - NumericConverter converter; + ScalarConverter converter; CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < 8; ++i) { + for (int i = 0; i < N; ++i) { result[i] = converter(source[i]); } @@ -3747,34 +4382,28 @@ struct NumericArrayConverter { } }; -/// Partial specialization for Array <= Array +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for Array <=> Array +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array template < int N, FloatRoundStyle Round > -struct NumericArrayConverter { - static_assert(!(N % 8), "N must be multiple of 8."); - - using result_type = Array; - using source_type = Array; +struct NumericArrayConverter { + using result_type = Array; + using source_type = Array; static FloatRoundStyle const round_style = Round; CUTLASS_HOST_DEVICE static result_type convert(source_type const & source) { - - NumericArrayConverter convert_vector_; - - result_type result; - - Array *result_ptr = reinterpret_cast *>(&result); - Array const *source_ptr = reinterpret_cast const *>(&source); - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < N / 8; ++i) { - result_ptr[i] = convert_vector_(source_ptr[i]); - } - - return result; + // Delegate to float converter, then reinterpret bits + NumericArrayConverter e2m1_to_fp32; + Array fp32_result = e2m1_to_fp32(source); + return reinterpret_cast(fp32_result); } CUTLASS_HOST_DEVICE @@ -4287,93 +4916,7 @@ struct NumericArrayConverter { #endif // Conditional guards to enable partial specialization for packed integers -namespace detail { - /* - A helper class that can vectorize a numeric converter with implementation for several vector widths. - - The vector widths must be giving in decreasing order or width, and must be a power of 2. - - The vector converters must produce identical results to the scalar converters for consistency. - */ - class VectorizedConverter { - private: - // Base case to handle remainder elements as scalars. - template - CUTLASS_DEVICE - static void convert_helper( - typename ArrayConverter::result_type& result, - typename ArrayConverter::source_type const& source) { - - using ElementRes = typename ArrayConverter::result_type::Element; - using ElementSrc = typename ArrayConverter::source_type::Element; - // If no more converters, handle the remaining elements as scalars. - constexpr int total_elements = ArrayConverter::result_type::kElements; - constexpr int remainder = total_elements - Offset; - static_assert(remainder == (total_elements % ParentWidth), "Unexpected remainder."); - - typename ArrayConverter::ScalarConverter scalar_converter; - CUTLASS_PRAGMA_UNROLL - for (int i = Offset; i < ArrayConverter::result_type::kElements; ++i) { - result[i] = scalar_converter(ElementSrc(source[i])); - } - } - - template - CUTLASS_DEVICE - static void convert_helper(typename ArrayConverter::result_type& result, typename ArrayConverter::source_type const& source) { - static_assert(sizeof...(OtherVectorArrays) % 2 == 0, "Vector converters must come in {dst, src} pairs"); - static_assert(ResultVectorArray::kElements == SourceVectorArray::kElements, "Vector converters must have the same vector width"); - static_assert(cutlass::platform::is_same::value, - "ResultVectorArray must have the same type ArrayConverter::result_type"); - static_assert(cutlass::platform::is_same::value, - "SourceVectorArray must have the same type ArrayConverter::result_type"); - static_assert(Offset >= 0 && Offset <= ArrayConverter::result_type::kElements, "Offset must be between 0 and N"); - - static_assert(ParentWidth == 0 || ParentWidth > ResultVectorArray::kElements, "Vector arrays must be given in decreasing order of width"); - - constexpr int vector_width = ResultVectorArray::kElements; - static_assert(ispow2(vector_width), "Vector width must be a power of 2"); - - using ElementRes = typename ArrayConverter::result_type::Element; - using ElementSrc = typename ArrayConverter::source_type::Element; - - constexpr int vector_bits_res = vector_width * cutlass::sizeof_bits::value; - constexpr int vector_bits_src = vector_width * cutlass::sizeof_bits::value; - - static_assert(vector_bits_res % 8 == 0, "Result vector type must be byte addressed."); - static_assert(vector_bits_src % 8 == 0, "Source vector type must be byte addressed."); - - constexpr int vector_offset = Offset / vector_width; - ResultVectorArray* packed_result_vec = reinterpret_cast(&result) + vector_offset; - SourceVectorArray const* packed_source_vec = reinterpret_cast(&source) + vector_offset; - - // Convert the remaining elements as vectors. - constexpr int total_elements = ArrayConverter::result_type::kElements; - constexpr int groups_of_vec = (total_elements - Offset) / vector_width; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < groups_of_vec; ++i) { - packed_result_vec[i] = ArrayConverter::template packed_convert(packed_source_vec[i]); - } - - constexpr int new_offset = Offset + vector_width * groups_of_vec; - // Recurse to handle other vector converters, or the scalar base case. - convert_helper(result, source); - } - - public: - /* - A method to convert vectors of elements using the packed_convert method of the converter. - - Converters using this class must implement packed convert and support 1 or more vector conversions. - */ - template - CUTLASS_DEVICE - static void convert(typename ArrayConverter::result_type& result, typename ArrayConverter::source_type const& source) { - convert_helper<0, 0, ArrayConverter, ResultVectorArray, SourceVectorArray, OtherVectorArrays...>(result, source); - } - }; -} ///////////////////////////////////////////////////////////////////////////////////////////////// /// Partial specialization for Array <= Array @@ -4398,12 +4941,14 @@ private: using ScalarConverter = NumericConverter; - #if defined(CUDA_PTX_FP8_CVT_ENABLED) + + #if defined(USE_PTX_CONVERT) // defined(CUDA_PTX_FP4FP6_CVT_ENABLED) CUTLASS_DEVICE static result_type_packed_8 ptx_convert(source_type_packed_8 const &source) { result_type_packed_8 out; - uint32_t* out_fp16 = reinterpret_cast(&out); + uint32_t* out_halfx2 = reinterpret_cast(&out); uint32_t const& src_packed = reinterpret_cast(source); + asm volatile( \ "{\n" \ ".reg .b8 byte0, byte1, byte2, byte3;\n" \ @@ -4412,39 +4957,87 @@ private: "cvt.rn.f16x2.e2m1x2 %1, byte1;\n" \ "cvt.rn.f16x2.e2m1x2 %2, byte2;\n" \ "cvt.rn.f16x2.e2m1x2 %3, byte3;\n" \ - "}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) , "=r"(out_fp16[2]), "=r"(out_fp16[3]): "r"(src_packed)); + "}\n" : "=r"(out_halfx2[0]), "=r"(out_halfx2[1]) , "=r"(out_halfx2[2]), "=r"(out_halfx2[3]): "r"(src_packed)); + return out; } CUTLASS_DEVICE static result_type_packed_4 ptx_convert(source_type_packed_4 const &source) { result_type_packed_4 out; - uint32_t* out_fp16 = reinterpret_cast(&out); + uint32_t* out_halfx2 = reinterpret_cast(&out); uint16_t const& src_packed = reinterpret_cast(source); + asm volatile( \ "{\n" \ ".reg .b8 byte0, byte1;\n" \ "mov.b16 {byte0, byte1}, %2;\n" \ "cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \ "cvt.rn.f16x2.e2m1x2 %1, byte1;\n" \ - "}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) : "h"(src_packed)); + "}\n" : "=r"(out_halfx2[0]), "=r"(out_halfx2[1]) : "h"(src_packed)); + return out; } CUTLASS_DEVICE static result_type_packed_2 ptx_convert(source_type_packed_2 const &source) { result_type_packed_2 out; - uint32_t* out_fp16 = reinterpret_cast(&out); + uint32_t* out_halfx2 = reinterpret_cast(&out); uint16_t const& src_packed = static_cast(reinterpret_cast(source)); + asm volatile( \ "{\n" \ ".reg .b8 byte0, byte1;\n" \ "mov.b16 {byte0, byte1}, %1;\n" \ "cvt.rn.f16x2.e2m1x2 %0, byte0;\n" \ - "}\n" : "=r"(out_fp16[0]) : "h"(src_packed)); + "}\n" : "=r"(out_halfx2[0]) : "h"(src_packed)); + + return out; + } + +#endif + + CUTLASS_DEVICE + static result_type_packed_8 lut_convert(source_type_packed_8 const &source) { + result_type_packed_8 out; + uint32_t* out_halfx2 = reinterpret_cast(&out); + uint32_t const& src_packed = reinterpret_cast(source); + + // LUT x4 e2m1 to fp16 (SM90) + // CUTLASS_PRAGMA_UNROLL + // for (int i = 0; i < 2; i++) { + // unsigned int lane0123 = src_packed >> (16 * i); + // detail::_e2m1_to_half_x4(lane0123, out_halfx2[2*i], out_halfx2[2*i+1]); + //} + + // LUT x8 e2m1 to fp16 (SM90) + detail::_e2m1_to_half_x8(src_packed, out_halfx2[0], out_halfx2[1], out_halfx2[2], out_halfx2[3]); + return out; } - #endif + + CUTLASS_DEVICE + static result_type_packed_4 lut_convert(source_type_packed_4 const &source) { + result_type_packed_4 out; + uint32_t* out_halfx2 = reinterpret_cast(&out); + uint16_t const& src_packed = reinterpret_cast(source); + + // LUT x4 e2m1 to half + detail::_e2m1_to_half_x4(src_packed, out_halfx2[0], out_halfx2[1]); + return out; + } + + CUTLASS_DEVICE + static result_type_packed_2 lut_convert(source_type_packed_2 const &source) { + result_type_packed_2 out; + uint32_t* out_halfx2 = reinterpret_cast(&out); + uint16_t const& src_packed = static_cast(reinterpret_cast(source)); + + // LUT x2 e2m1 to half + detail::_e2m1_to_half_x2(src_packed, out_halfx2[0]); + return out; + } + template CUTLASS_DEVICE @@ -4457,27 +5050,20 @@ private: platform::is_same::value), "Invalid PackedSrcType/PackedResultType must be 2, 4 or 8 to use private convert dispatch."); - #if defined(CUDA_PTX_FP4FP6_CVT_ENABLED) + // either call lookup table or PTX instruction implementation + #if defined(USE_PTX_CONVERT) // defined(CUDA_PTX_FP4FP6_CVT_ENABLED) return ptx_convert(source); #else - PackedResultType result; - NumericConverter converter; - - const int k_packed = PackedResultType::kElements; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < k_packed; ++i) { - result[i] = converter(source[i]); - } - - return result; + return lut_convert(source); #endif } friend class detail::VectorizedConverter; public: - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) result_type result; using ConverterType = NumericArrayConverter; detail::VectorizedConverter::convert(result, source); return result; + #else + result_type result; + ScalarConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + CUTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round, + int N +> +struct NumericArrayConverter { + using result_element = cutlass::bfloat16_t; + using source_element = cutlass::float_e2m1_t; + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + +private: + using result_type_packed_8 = Array; + using result_type_packed_4 = Array; + using result_type_packed_2 = Array; + using source_type_packed_8 = Array; + using source_type_packed_4 = Array; + using source_type_packed_2 = Array; + + using ScalarConverter = NumericConverter; + + CUTLASS_DEVICE + static result_type_packed_8 lut_convert(source_type_packed_8 const &source) { + + result_type_packed_8 out; + uint32_t* out_bfloatx2 = reinterpret_cast(&out); + uint32_t const& src_packed = reinterpret_cast(source); + + // LUT x4 e2m1 to bf16 (SM90) + // CUTLASS_PRAGMA_UNROLL + // for (int i = 0; i < 2; i++) { + // unsigned int lane0123 = src_packed >> (16 * i); + // detail::_e2m1_to_bf16_x4(lane0123, out_bfloatx2[2*i], out_bfloatx2[2*i+1]); + // } + + // LUT x8 e2m1 to bf16 (SM90) + detail::_e2m1_to_bf16_x8(src_packed, out_bfloatx2[0], out_bfloatx2[1], out_bfloatx2[2], out_bfloatx2[3]); + + return out; + } + + CUTLASS_DEVICE + static result_type_packed_4 lut_convert(source_type_packed_4 const &source) { + result_type_packed_4 out; + uint32_t* out_bfloatx2 = reinterpret_cast(&out); + uint16_t const& src_packed = reinterpret_cast(source); + + // LUT x4 e2m1 to bf16 + detail::_e2m1_to_bf16_x4(src_packed, out_bfloatx2[0], out_bfloatx2[1]); + return out; + } + + CUTLASS_DEVICE + static result_type_packed_2 lut_convert(source_type_packed_2 const &source) { + result_type_packed_2 out; + uint32_t* out_bfloatx2 = reinterpret_cast(&out); + uint16_t const& src_packed = static_cast(reinterpret_cast(source)); + + // LUT x2 e2m1 to bf16 (SM90) + detail::_e2m1_to_bf16_x2(src_packed, out_bfloatx2[0]); + return out; + } + + template + CUTLASS_DEVICE + static PackedResultType packed_convert(PackedSrcType const &source) { + static_assert((platform::is_same::value && + platform::is_same::value) || + (platform::is_same::value && + platform::is_same::value) || + (platform::is_same::value && + platform::is_same::value), + "Invalid PackedSrcType/PackedResultType must be 2, 4 or 8 to use private convert dispatch."); + + // Option to add an optimized PTX path + return lut_convert(source); + } + + friend class detail::VectorizedConverter; + +public: + CUTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) + result_type result; + using ConverterType = NumericArrayConverter; + detail::VectorizedConverter::convert(result, source); + + return result; + #else + result_type result; + ScalarConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif } CUTLASS_HOST_DEVICE @@ -4545,7 +5256,7 @@ private: // [0, 1, -2, -1] encoded as FP8 static constexpr uint32_t E4M3_LUT = 0xB8C03800; - const int iters = PackedSrcType::kElements / 4; + constexpr int iters = PackedSrcType::kElements / 4; #pragma unroll for (int ii = 0; ii < iters; ii += 2, src_reg >>= 16, src_reg_shifted >>= 16) { // This uses a look up table to convert packed int2s to packed fp8s, using the int4 value @@ -4949,6 +5660,202 @@ struct NumericArrayConverter { } }; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template +struct NumericArrayConverter { + using result_type = Array; + using source_type = Array; + + static FloatRoundStyle const round_style = Round; + +private: + using result_type_packed_8 = Array; + using result_type_packed_4 = Array; + using source_type_packed_8 = Array; + using source_type_packed_4 = Array; + + using ScalarConverter = NumericConverter; + + CUTLASS_DEVICE + static uint32_t to_reg(source_type_packed_4 const& source) { + return static_cast( + reinterpret_cast(source)); + } + + CUTLASS_DEVICE + static uint32_t to_reg(source_type_packed_8 const& source) { + return reinterpret_cast(source); + } + + // The core converter uses a lookup table to converts e2m1 -> e4m3. + template + CUTLASS_DEVICE + static PackedResultType packed_convert(PackedSrcType const &source) { + + static_assert((platform::is_same::value && + platform::is_same::value) || + (platform::is_same::value && + platform::is_same::value), + "Invalid PackedSrcType/PackedResultType must be 4 or 8 to use private convert dispatch."); + + // Hold FP8 outputs in reg. We need 1 reg for every 4 outputs. + cutlass::AlignedArray out_fp8; + + // View the input as reg + uint32_t reg = to_reg(source); + + constexpr int iters_x8 = PackedSrcType::kElements / 8; + constexpr int iters_x4 = (PackedSrcType::kElements / 4) & 1; + constexpr unsigned long long E2M1_to_E4M3_LUT_0_7 = 0x4C4844403C383000ULL; + + // we really only get called with 4 or 8 elements, but allow for more arguments for future use + // practically it's one of either x8 or x4, since we unroll everything else will be optimized away + #pragma unroll + for (int ii = 0; ii < iters_x8; ++ii) { + detail::_e2m1_to_fp8_x8(reg, out_fp8[ii*2], out_fp8[ii*2+1]); + } + + if (iters_x4) { + detail::_e2m1_to_fp8_x4(reg, out_fp8[iters_x8*2]); + } + + return reinterpret_cast(out_fp8); + } + + friend class detail::VectorizedConverter; + +public: + CUTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) + result_type result; + using ConverterType = NumericArrayConverter; + detail::VectorizedConverter::convert(result, source); + + return result; + #else + result_type result; + ScalarConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + + CUTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + + +/// Partial specialization for Array <= Array +template +struct NumericArrayConverter { + using result_type = Array; + using source_type = Array; + + static FloatRoundStyle const round_style = Round; + +private: + using result_type_packed_8 = Array; + using result_type_packed_4 = Array; + using source_type_packed_8 = Array; + using source_type_packed_4 = Array; + + using ScalarConverter = NumericConverter; + + CUTLASS_DEVICE + static uint32_t to_reg(source_type_packed_4 const& source) { + return static_cast( + reinterpret_cast(source)); + } + + CUTLASS_DEVICE + static uint32_t to_reg(source_type_packed_8 const& source) { + return reinterpret_cast(source); + } + + // The core converter uses a lookup table to convert e2m1 -> e5m2. + // Uses the existing templated _e2m1_to_fp8_x* helper functions defined above. + template + CUTLASS_DEVICE + static PackedResultType packed_convert(PackedSrcType const &source) { + + static_assert((platform::is_same::value && + platform::is_same::value) || + (platform::is_same::value && + platform::is_same::value), + "Invalid PackedSrcType/PackedResultType must be 4 or 8 to use private convert dispatch."); + + // Hold FP8 outputs in reg. We need 1 reg for every 4 outputs. + cutlass::AlignedArray out_fp8; + + // View the input as reg + uint32_t reg = to_reg(source); + + constexpr int iters_x8 = PackedSrcType::kElements / 8; + constexpr int iters_x4 = (PackedSrcType::kElements / 4) & 1; + // E2M1→E5M2 LUT from table at lines 6109-6124 + constexpr unsigned long long E2M1_to_E5M2_LUT_0_7 = 0x464442403E3C3800ULL; + + // we really only get called with 4 or 8 elements, but allow for more arguments for future use + #pragma unroll + for (int ii = 0; ii < iters_x8; ++ii) { + detail::_e2m1_to_fp8_x8(reg, out_fp8[ii*2], out_fp8[ii*2+1]); + } + + if (iters_x4) { + detail::_e2m1_to_fp8_x4(reg, out_fp8[iters_x8*2]); + } + + return reinterpret_cast(out_fp8); + } + + friend class detail::VectorizedConverter; + +public: + CUTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) + result_type result; + using ConverterType = NumericArrayConverter; + detail::VectorizedConverter::convert(result, source); + + return result; + #else + result_type result; + ScalarConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + + CUTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + + /// Partial specialization for Array <= Array template struct NumericArrayConverter { @@ -5012,7 +5919,7 @@ private: static constexpr uint32_t NEG_E4M3s_REG2 = 0xB8C0C4C8; - const int iters = PackedSrcType::kElements / 4; + constexpr int iters = PackedSrcType::kElements / 4; #pragma unroll for (int ii = 0; ii < iters; ++ii, lut_idx >>=16, sign >>=16) { uint32_t final_prmt_idx = final_prmt_base | sign; @@ -5038,8 +5945,9 @@ private: friend class detail::VectorizedConverter; public: - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) result_type result; using ConverterType = NumericArrayConverter; detail::VectorizedConverter::convert(result, source); return result; + #else + result_type result; + ScalarConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE result_type operator()(source_type const &s) const { return convert(s); } @@ -5119,7 +6038,7 @@ private: static constexpr uint32_t NEG_E5M2s_REG2 = 0xBCC0C2C4; - const int iters = PackedSrcType::kElements / 4; + constexpr int iters = PackedSrcType::kElements / 4; #pragma unroll for (int ii = 0; ii < iters; ++ii, lut_idx >>=16, sign >>=16) { uint32_t final_prmt_idx = final_prmt_base | sign; @@ -5145,8 +6064,9 @@ private: friend class detail::VectorizedConverter; public: - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE static result_type convert(source_type const &source) { + #if defined(__CUDA_ARCH__) result_type result; using ConverterType = NumericArrayConverter; detail::VectorizedConverter::convert(result, source); return result; + #else + result_type result; + ScalarConverter converter; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif } - CUTLASS_DEVICE + CUTLASS_HOST_DEVICE result_type operator()(source_type const &s) const { return convert(s); } @@ -5226,7 +6157,7 @@ private: static constexpr uint32_t E4M3s_REG4 = 0x57565554; - const int iters = PackedSrcType::kElements / 4; + constexpr int iters = PackedSrcType::kElements / 4; #pragma unroll for (int ii = 0; ii < iters; ++ii, lut_idx >>=16, sign >>=16) { uint32_t final_prmt_idx = final_prmt_base | sign; diff --git a/include/cutlass/platform/platform.h b/include/cutlass/platform/platform.h index 3bf82c6a..ee58b054 100644 --- a/include/cutlass/platform/platform.h +++ b/include/cutlass/platform/platform.h @@ -31,6 +31,8 @@ #pragma once +#include "cutlass/tfloat32.h" + /** * \file * \brief C++ features that may be otherwise unimplemented for CUDA device functions. @@ -580,7 +582,7 @@ template struct alignment_of : std::alignment_of {}; #endif - +#if CUDA_VERSION >= 11080 /* 16B specializations where 32-bit Win32 host compiler disagrees with device compiler */ template <> struct alignment_of { @@ -676,7 +678,7 @@ struct alignment_of { }; #endif - +#endif // CUDA_VERSION >= 11080 // Specializations for volatile/const qualified types template @@ -915,6 +917,18 @@ struct numeric_limits { static constexpr bool has_infinity = true; }; +template <> +struct numeric_limits { + CUTLASS_HOST_DEVICE + static tfloat32_t infinity() noexcept { return tfloat32_t::bitcast(0x7f800000);} + CUTLASS_HOST_DEVICE + static tfloat32_t max() noexcept { return tfloat32_t::bitcast(0x7f7fffff);} + CUTLASS_HOST_DEVICE + static tfloat32_t lowest() noexcept { return tfloat32_t::bitcast(0xff7fffff);} + static constexpr bool is_integer = false; + static constexpr bool has_infinity = true; +}; + /// Returns a value that curries the `std::maximum()` function into the identity /// function. No value will compare < than this value. template diff --git a/include/cutlass/predicate_vector.h b/include/cutlass/predicate_vector.h index fe82fbb9..76c02e2d 100644 --- a/include/cutlass/predicate_vector.h +++ b/include/cutlass/predicate_vector.h @@ -40,7 +40,9 @@ #include #endif +#ifndef __QNX__ #include CUDA_STD_HEADER(cassert) +#endif #include "cutlass/platform/platform.h" diff --git a/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp b/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp index 83a891b5..e1b485c9 100644 --- a/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp +++ b/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp @@ -494,8 +494,7 @@ private: // * Output Cta Tensor S to G if (GemmM_within_Cta > 0 && GemmK_within_Cta > 0) { - constexpr int MaxVecBits = 128; // STG.128 - cute::cooperative_copy(threadIdx_X, cEsE, cEgE); + cute::cooperative_copy(threadIdx_X, cEsE, cEgE); } if (GemmMAlignedAC_within_Cta == TensorEAtomM{} && GemmKAlignedAC_within_Cta == TensorEAtomK{}) { @@ -577,9 +576,9 @@ private: // Row Major if constexpr (IsRowMajor) { CUTE_UNROLL - for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(shape<0>(dSrc), ThreadShapeRows * ValueShapeRows); ++iter_row_blk) { + for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(valid_rows, ThreadShapeRows * ValueShapeRows); ++iter_row_blk) { CUTE_UNROLL - for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(shape<1>(dSrc) , ThreadShapeCols * ValueShapeCols); ++col_chunk_i) { + for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(valid_cols, ThreadShapeCols * ValueShapeCols); ++col_chunk_i) { CUTE_UNROLL for (int iter_row_thr = 0; iter_row_thr < ValueShapeRows; ++iter_row_thr) { CUTE_UNROLL @@ -587,7 +586,9 @@ private: const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr; const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr; if constexpr ( (not pred) and (not IsQmmaF6) ) { - dDst(row_i, col_i) = dSrc(row_i, col_i); + if (row_i < valid_rows && col_i < valid_cols) { + dDst(row_i, col_i) = dSrc(row_i, col_i); + } } else { if (row_i < valid_rows && col_i < valid_cols) { @@ -602,9 +603,9 @@ private: // Col Major else { CUTE_UNROLL - for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(shape<1>(dSrc) , ThreadShapeCols * ValueShapeCols); ++col_chunk_i) { + for (int col_chunk_i = 0; col_chunk_i < cutlass::ceil_div(valid_cols, ThreadShapeCols * ValueShapeCols); ++col_chunk_i) { CUTE_UNROLL - for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(shape<0>(dSrc), ThreadShapeRows * ValueShapeRows); ++iter_row_blk) { + for (int iter_row_blk = 0; iter_row_blk < cutlass::ceil_div(valid_rows, ThreadShapeRows * ValueShapeRows); ++iter_row_blk) { CUTE_UNROLL for (int iter_col_thr = 0; iter_col_thr < ValueShapeCols; ++iter_col_thr) { CUTE_UNROLL @@ -612,7 +613,9 @@ private: const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr; const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr; if constexpr ( (not pred) and (not IsQmmaF6) ) { - dDst(row_i, col_i) = dSrc(row_i, col_i); + if (row_i < valid_rows && col_i < valid_cols) { + dDst(row_i, col_i) = dSrc(row_i, col_i); + } } else { if (row_i < valid_rows && col_i < valid_cols) { diff --git a/include/cutlass/workspace.h b/include/cutlass/workspace.h index bb318aa3..610ca8c3 100644 --- a/include/cutlass/workspace.h +++ b/include/cutlass/workspace.h @@ -121,6 +121,8 @@ fill_workspace(void* workspace, T fill_value, size_t fill_count, cudaStream_t st #else CUdeviceptr d_workspace = reinterpret_cast(workspace); CUresult result = CUDA_SUCCESS; + +#ifndef __QNX__ if (sizeof(T) == 4) { result = cuMemsetD32Async(d_workspace, reinterpret_cast(fill_value), fill_count, stream); } @@ -130,6 +132,7 @@ fill_workspace(void* workspace, T fill_value, size_t fill_count, cudaStream_t st else if (sizeof(T) == 1) { result = cuMemsetD8Async(d_workspace, reinterpret_cast(fill_value), fill_count, stream); } +#endif if (CUDA_SUCCESS != result) { const char** error_string_ptr = nullptr; diff --git a/media/docs/cpp/blackwell_functionality.md b/media/docs/cpp/blackwell_functionality.md index 04e4139d..8b410e1c 100644 --- a/media/docs/cpp/blackwell_functionality.md +++ b/media/docs/cpp/blackwell_functionality.md @@ -131,7 +131,7 @@ instructions supported by CUTLASS. |[8](#bs_rows_4_5_7_8_10) | Dense | mx_float6_t | mx_float6_t | TN, NN, NT, TT | 128 | 128 | mxf8f6f4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf6_mxf6_void_bf16_tn_layout.cu)
[NT unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf6_mxf6_void_bf16_nt_layout.cu)| |[9](#bs_rows_6_9_11) | Dense | mx_float6_t | mx_float8_t | TN, NN, NT, TT | 128 | 16 | mxf8f6f4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf6_mxf8_void_f32_tn_layout.cu)
[NT unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf6_mxf8_void_f32_nt_layout.cu)| |[10](#bs_rows_4_5_7_8_10) | Dense | mx_float8_t | mx_float6_t | TN, NN, NT, TT | 16 | 128 | mxf8f6f4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf8_mxf6_f16_f8_tn_layout.cu)
[NT unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf8_mxf6_f16_f8_nt_layout.cu)| -|[11](#bs_rows_6_9_11) | Dense | mx_float8_t | mx_float8_t | TN, NN, NT, TT | 16 | 16 | mxf8f6f4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf8_mxf8_void_f8_tn_layout.cu.cu)
[NT unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf8_mxf8_void_f8_nt_layout.cu)| +|[11](#bs_rows_6_9_11) | Dense | mx_float8_t | mx_float8_t | TN, NN, NT, TT | 16 | 16 | mxf8f6f4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf8_mxf8_void_f8_tn_layout.cu)
[NT unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_tensorop_gemm/mxf8_mxf8_void_f8_nt_layout.cu)| |[12](#bs_rows_1) | Sparse | nv_float4_t | nv_float4_t | TN | 32 (N) / 64 (T) | 32 | mxf4nvf4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_sparse_tensorop_gemm/sm100_bssp_gemm_nvf4_nvf4_f32_void_f16_o_tnn.cu) | |[13](#bs_rows_2) | Sparse | mx_float4_t | mx_float4_t | TN | 32 (N) / 64 (T) | 32 | mxf4, mxf4nvf4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_sparse_tensorop_gemm/sm100_bssp_gemm_mxf4_mxf4_f32_f16_f16_o_tnn.cu) | |[14](#bs_rows_3) | Sparse | mx_float4_t | mx_float4_t | TN, NN, NT, TT | 128 (N) / 256 (T) | 128 | mxf8f6f4 |[TN unit tests](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/sm100_blockscaled_sparse_tensorop_gemm/sm100_bssp_gemm_mxf4_mxf4_f32_f16_f16_q_tnt.cu) | @@ -652,7 +652,7 @@ auto tensor_sfb = make_tensor(bptr, layout_sfb); auto val_a_mk = tensor_sfa(make_coord(m,k,0)); ``` # Blackwell SM120 GEMMs -The NVIDIA RTX 5000 Series GPUs introduce support for new narrow precision (4bit and 6bit) block-scaled and non-block-scaled tensor cores. The PTX ISA has extended the `mma` instructions to support these data formats which are 1x to 4x faster than Ada architecture's fp8 tensor cores. For more detailed information see [`mma` PTX documentation](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#multiply-and-accumulate-instruction-mma). +The NVIDIA RTX 5000 Series GPUs introduce support for new narrow precision (4bit and 6bit) block-scaled and non-block-scaled tensor cores. The PTX ISA has extended the `mma` instructions to support these data formats which are 1x to 4x faster than Ada architecture's fp8 tensor cores. For more detailed information see [`mma` PTX documentation](https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-for-mma). CUTLASS 4.0 has added support for these newly introduced narrow precision GEMMs. Similar to the Blackwell SM100 GEMMs, the SM120 GEMMs can be built using the collective builder interface. See examples in [examples/79_blackwell_geforce_gemm/](../../examples/79_blackwell_geforce_gemm/) and unit tests listed below. diff --git a/media/docs/cpp/functionality.md b/media/docs/cpp/functionality.md index 88fb1678..a5e446f2 100644 --- a/media/docs/cpp/functionality.md +++ b/media/docs/cpp/functionality.md @@ -47,19 +47,19 @@ Hyperlinks to relevant unit tests demonstrate how specific template instances ma | **WmmaTensorOp** | 70+ | 11.4+ | `f16 * f16 + f16 => f16` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16t_f16t_f16n_wmma_tensor_op_f16_sm70.cu) | | **WmmaTensorOp** | 70+ | 11.4+ | `f16 * f16 + f32 => {f16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16t_f16t_f16n_wmma_tensor_op_f32_sm70.cu) | | **WmmaTensorOp** | 75+ | 11.4+ | `s8 * s8 + s32 => {s32, s8}` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_s8t_s8n_s8t_wmma_tensor_op_s32_sm72.cu) | -| **WmmaTensorOp** | 75+ | 11.4+ | `s4 * s4 + s32 => {s32, s4}` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_s4t_s4n_s4t_wmma_tensor_op_s32_sm75.cu) | -| **WmmaTensorOp** | 75+ | 11.4+ | `b1 ^ b1 + s32 => {s32, b1}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_b1t_b1n_b1t_wmma_tensor_op_s32_sm75.cu) | +| **WmmaTensorOp** | 75+ | 11.4+ | `s4 * s4 + s32 => {s32, s4}` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/gemm_s4t_s4n_s32t_wmma_tensor_op_s32_sm75.cu) | +| **WmmaTensorOp** | 75+ | 11.4+ | `b1 ^ b1 + s32 => {s32, b1}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/gemm_b1t_b1n_s32n_wmma_tensor_op_s32_sm75.cu) | | **TensorOp** | 70+ | 11.4+ | `f16 * f16 + f16 => f16` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f16t_volta_tensor_op_f16_sm70.cu) | -| **TensorOp** | 70+ | 11.4+ | `f16 * f16 + f32 => {f16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f16t_volta_tensor_op_f32_sm70.cu) | +| **TensorOp** | 70+ | 11.4+ | `f16 * f16 + f32 => {f16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f32t_volta_tensor_op_f32_sm70.cu) | | **TensorOp** | 75+ | 11.4+ | `f16 * f16 + f16 => f16` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f16_sm75.cu) | -| **TensorOp** | 75+ | 11.4+ | `f16 * f16 + f32 => {f16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm75.cu) | +| **TensorOp** | 75+ | 11.4+ | `f16 * f16 + f32 => {f16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/gemm_f16n_f16t_f32t_tensor_op_f32_sm75.cu) | | **TensorOp** | 75+ | 11.4+ | `s8 * s8 + s32 => {s32, s8}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_s8t_s8n_s32n_tensor_op_s32_sm75.cu) | -| **TensorOp** | 75+ | 11.4+ | `s4 * s4 + s32 => {s32, s4}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_s4t_s4n_s32n_tensor_op_s32_sm75.cu) | +| **TensorOp** | 75+ | 11.4+ | `s4 * s4 + s32 => {s32, s4}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/gemm_s4t_s4n_s32n_tensor_op_s32_sm75.cu) | | **TensorOp** | 75+ | 11.4+ | `b1 ^ b1 + s32 => {s32, b1}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_b1t_b1n_s32n_tensor_op_s32_sm75.cu) | | **TensorOp** | 80+ | 11.4+ | `f16 * f16 + f16 => f16` | {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f16_sm80.cu) | | **TensorOp** | 80+ | 11.4+ | `f16 * f16 + f32 => {f16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f16n_f16t_f16t_tensor_op_f32_sm80.cu) | -| **TensorOp** | 80+ | 11.4+ | `bf16 * bf16 + f32 => {bf16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_bf16n_bf16t_bf16t_tensor_op_f32_sm80.cu) | -| **TensorOp** | 80+ | 11.4+ | `tf32 * tf32 + f32 => f32`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_f32n_f32t_f32t_tensor_op_f32_sm80.cu) | +| **TensorOp** | 80+ | 11.4+ | `bf16 * bf16 + f32 => {bf16, f32}`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/gemm_bf16t_bf16t_bf16t_tensor_op_f32_sm80.cu) | +| **TensorOp** | 80+ | 11.4+ | `tf32 * tf32 + f32 => f32`| {N,T} x {N,T} => {N,T} | [example](https://github.com/NVIDIA/cutlass/blob/main/test/unit/gemm/device/gemm_tf32n_tf32t_f32t_tensor_op_f32_sm80.cu) | | **TensorOp** | 80+ | 11.4+ | `s8 * s8 + s32 => {s32, s8}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_s8t_s8n_s32n_tensor_op_s32_sm80.cu) | | **TensorOp** | 80+ | 11.4+ | `s4 * s4 + s32 => {s32, s4}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_s4t_s4n_s32n_tensor_op_s32_sm80.cu) | | **TensorOp** | 80+ | 11.4+ | `b1 ^ b1 + s32 => {s32, b1}` | { T } x { N } => {N,T} | [example](https://github.com/NVIDIA/cutlass/tree/main/test/unit/gemm/device/gemm_b1t_b1n_s32n_tensor_op_s32_sm80.cu) | diff --git a/media/docs/cpp/pipeline.md b/media/docs/cpp/pipeline.md index bcd5e332..8455bc42 100644 --- a/media/docs/cpp/pipeline.md +++ b/media/docs/cpp/pipeline.md @@ -78,7 +78,7 @@ These classes help developers orchestrate a pipeline of asynchronous producer and consumer threads, without needing to worry about lower-level hardware details. These classes serve a similar function as the various -[pipeline abstractions](https://nvidia.github.io/libcudacxx/extended_api/synchronization_primitives/pipeline.html) +[pipeline abstractions](https://nvidia.github.io/cccl/libcudacxx/extended_api/synchronization_primitives/pipeline.html) in libcu++. #### Pipeline methods diff --git a/media/docs/pythonDSL/cute_dsl.rst b/media/docs/pythonDSL/cute_dsl.rst index 1c129a77..a05c434e 100644 --- a/media/docs/pythonDSL/cute_dsl.rst +++ b/media/docs/pythonDSL/cute_dsl.rst @@ -17,4 +17,6 @@ CuTe DSL Debugging with the DSL Autotuning with the DSL Educational Notebooks + Deprecation Policy Compile with TVM FFI + Ahead-of-Time (AOT) Compilation diff --git a/media/docs/pythonDSL/cute_dsl_api/changelog.rst b/media/docs/pythonDSL/cute_dsl_api/changelog.rst index 49a66caa..e4200e07 100644 --- a/media/docs/pythonDSL/cute_dsl_api/changelog.rst +++ b/media/docs/pythonDSL/cute_dsl_api/changelog.rst @@ -74,7 +74,7 @@ Changelog for CuTe DSL API changes - Introduce BlockScaled layout utilities in `blockscaled_layout.py `_ for creating the required scale factor layouts in global memory, shared memory and tensor memory. * ``cutlass.cute.compile`` now supports compilation options. Refer to `JIT compilation options `_ for more details. -* ``cutlass.cute.testing.assert_`` now works for device JIT function. Specify ``--enable-device-assertions`` as compilation option to enable. +* ``cutlass.cute.testing.assert_`` now works for device JIT function. Specify ``--enable-assertions`` as compilation option to enable. * ``cutlass.cute.make_tiled_copy`` is now deprecated. Please use ``cutlass.cute.make_tiled_copy_tv`` instead. * Shared memory capacity query diff --git a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst index 8459cbae..c82c74d8 100644 --- a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst +++ b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst @@ -4,7 +4,8 @@ Compile with TVM FFI ==================== -Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in the `official documentation `_. +Apache TVM FFI is an open ABI and FFI for machine learning systems. More information can be found in +the `official documentation `_. To install TVM FFI, you can run the following command: @@ -14,7 +15,9 @@ To install TVM FFI, you can run the following command: # optional package for improved torch tensor calling performance pip install torch-c-dlpack-ext -In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster JIT function invocation and provides better interoperability with machine learning frameworks (e.g., directly take ``torch.Tensor`` as arguments). +In |DSL|, TVM FFI can be enabled as an option for JIT-compiled functions. Using TVM FFI can lead to faster +JIT function invocation and provides better interoperability with machine learning frameworks +(e.g., directly take ``torch.Tensor`` as arguments). Enable Apache TVM FFI in |DSL| @@ -129,7 +132,8 @@ stride via the ``stride`` argument in the ``make_fake_tensor`` API. ``cute.Tensor`` adapter for TVM FFI ----------------------------------- -To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the ``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example: +To adapt the ``cute.Tensor`` to the TVM FFI function, you can use the ``cute.runtime.from_dlpack`` function with the +``enable_tvm_ffi=True`` option or the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. For example: .. code-block:: python @@ -570,7 +574,7 @@ Then you can load back the exported module and use it in different ways: from cutlass import cute def example_load_module_add_one(): - mod = cute.runtime.load_module("./add_one.so") + mod = cute.runtime.load_module("./add_one.so", enable_tvm_ffi=True) a_torch = torch.arange(10, dtype=torch.float32, device="cuda") b_torch = torch.empty(10, dtype=torch.float32, device="cuda") mod.add_one(a_torch, b_torch) @@ -587,9 +591,11 @@ in the official documentation. When you build your own libraries, make sure you link against the necessary runtime libraries. You can use ``cute.runtime.find_runtime_libraries(enable_tvm_ffi=True)`` to get the path to these libraries. -``cute.runtime.load_module`` will load these libraries automatically before loading +``cute.runtime.load_module(path, enable_tvm_ffi=True)`` will load these libraries automatically before loading an exported module. You can also manually load these libraries in advanced use cases. +For low-level cute ABI AOT compilation support without TVM FFI, you can refer to :doc:`dsl_ahead_of_time_compilation`. + Keyword Arguments and Defaults ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -683,3 +689,32 @@ The code block below shows how to do this: wrapped_func(a_torch, b_torch, offset=4) print("result of b_torch after wrapped_func(a_torch, b_torch, offset=4)") print(b_torch) + + +Limitations +----------- + +The Fake Tensor flow is ONLY compatible with TVM FFI because TVM FFI supports more flexible constraints on Tensor arguments. +For instance, fake tensor can specify per-mode static shape or constraints on shape and strides which are not supported by +existing ``from_dlpack`` flow. It's expected that JIT function compiled with fake tensor will have different ABI compared to +tensor converted by ``from_dlpack``. + +.. code-block:: python + + import cutlass.cute as cute + import torch + + n = cute.sym_int() + # Dynamic Shape + fake_a = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + + # Compile without tvm-ffi + compiled_fn = cute.compile(foo, fake_a) + + # Wrong, in compatible ABI + compiled_fn(from_dlpack(a)) + + +In order to avoid such issue, it's recommended to use fake tensor only with TVM FFI backend. Practically speaking, +as we only want to call ``from_dlpack`` once and reuse for both compilation and runtime, the benefit of +using fake tensor is limited in this case. diff --git a/media/docs/pythonDSL/cute_dsl_general/debugging.rst b/media/docs/pythonDSL/cute_dsl_general/debugging.rst index 93c75112..759086e3 100644 --- a/media/docs/pythonDSL/cute_dsl_general/debugging.rst +++ b/media/docs/pythonDSL/cute_dsl_general/debugging.rst @@ -172,6 +172,25 @@ For detecting memory errors and race conditions: Please refer to the `compute-sanitizer documentation `_ for more details. +Set function name prefix +~~~~~~~~~~~~~~~~~~~~~~~~~ + +By default, the function name (host function or kernel function) is automatically generated based on the function name and its parameters. +Sometimes you may want to attach some runtime information to the function name to make performance profiling and debugging easier, +e.g., the kernel configs or the rank ids. You can assign a name prefix to the name by calling the ``set_name_prefix`` +method on the host function or kernel function. + +.. code:: python + + @cute.kernel + def kernel(arg1, arg2, ...): + ... + @cute.jit + def launch_kernel(): + kernel.set_name_prefix("your_custom_name_prefix") + kernel(arg1, arg2, ...).launch(grid=[1, 1, 1], block=[1, 1, 1], ...) + +For above example, the generated kernel name will be "your_custom_name_prefix_xxx". Conclusion ---------- diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst new file mode 100644 index 00000000..67b4354f --- /dev/null +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_ahead_of_time_compilation.rst @@ -0,0 +1,239 @@ +.. _dsl_ahead_of_time_compilation: +.. |DSL| replace:: CuTe DSL + +Ahead-of-Time (AOT) Compilation +=============================== + +This guide demonstrates how to use |DSL|'s Ahead-of-Time (AOT) compilation features to export compiled kernels for use in production environments. + +Overview +-------- + +|DSL| Ahead-of-Time (hereinafter referred to as AOT) compilation allows you to: + +* **Compile once, enable cross-compilation**: Write kernels in Python and cross-compile them for multiple GPU architectures. +* **Remove JIT overhead**: Eliminate compilation delays in production by pre-compiling kernels. +* **Flexible integration**: Easily integrate compiled kernels into both Python and C/C++ codebases using flexible deployment options. + +We provide 2 levels of AOT ABI: + +1. **Low-Level CuTe ABI**: This ABI is expressed using CuTe DSL types and tensors, mirroring the original Python function. +2. **High-Level Apache TVM FFI ABI**: For interop with various frameworks (e.g., PyTorch, JAX), and offer high-level stable ABI access. + +This guide will focus on the CuTe ABI AOT. For the Apache TVM FFI AOT, please refer to the section "Exporting Compiled Module" in :doc:`compile_with_tvm_ffi`. + +CuTe ABI AOT Workflow +--------------------- + +Export Interface +~~~~~~~~~~~~~~~~ + +The ``export_to_c`` interface is provided by the ``JitCompiledFunction`` class. It accepts the following parameters: + +* ``file_path``: The path to the directory where the header and object files will be saved. +* ``file_name``: The base name for the header and object files. The same file name will always overwrite existing files. +* ``function_prefix``: The prefix of the function symbol in the generated object file. This should be a unique identifier to avoid symbol conflicts. Users should ensure the function prefix is unique for each exported function. Defaults to the ``file_name``. + +It generates the following files: + +* ``{file_path}/{file_name}.h``: A C header file containing API function declarations. This header specifies the runtime function signatures in C, mirroring the original Python function interfaces. +* ``{file_path}/{file_name}.o``: A standard object file containing the compiled kernel code. You can link this object file into either a static or shared library. It includes the host entry function, fatbin data, and helper functions such as ``cuda_init`` and ``cuda_load_to_device``. Additionally, it embeds metadata for runtime loading and version verification. + +Example: + +.. code-block:: python + + import cutlass.cute as cute + import cutlass.cute.cuda as cuda + + @cute.kernel + def print_tensor_kernel(a: cute.Tensor): + cute.printf("a: {}", a) + + @cute.jit + def print_tensor(a: cute.Tensor, stream: cuda.CUstream): + print_tensor_kernel(a).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream) + + compiled_func = cute.compile(print_tensor) + # Export compiled functions to object files and headers + compiled_func.export_to_c(file_path="./artifacts", file_name="print_tensor_example", function_prefix="print_tensor") + +Loading in Python +~~~~~~~~~~~~~~~~~ + +Load pre-compiled object files or shared libraries into Python for execution. + +.. code-block:: python + + import cutlass.cute as cute + import torch + from cutlass.cute import from_dlpack + import cutlass.cute.cuda as cuda + + # Load module from object file + module = cute.runtime.load_module("./artifacts/print_tensor_example.o") + # or + module = cute.runtime.load_module("./artifacts/libprint_tensor_example.so") + + # Prepare data + a = torch.arange(160, dtype=torch.float32, device="cuda").reshape(16, 10) + a_cute = from_dlpack(a).mark_layout_dynamic() + stream = cuda.CUstream(0) + + # Call the function (no JIT compilation needed!) + module.print_tensor(a_cute, stream=stream) + + # This will fail because 'non_existing_api' was not exported: + # module.non_existing_api() + +C++ Integration with Static Linking +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Integrate compiled kernels directly into your C++ executable during the build process. The generated header file supplies the necessary API for loading the module and invoking the function. + +Example: + +.. code-block:: cpp + + #include "print_tensor_example.h" + #include + + void run_print_tensor() { + // Prepare tensor, the tensor declaration is in the header file + print_tensor_Tensor_a_t tensor_a; + tensor_a.data = nullptr; // GPU memory is set to nullptr. + // Set dynamic shapes and strides + tensor_a.dynamic_shapes[0] = 32; + tensor_a.dynamic_shapes[1] = 16; + tensor_a.dynamic_strides[0] = 16; + + // Create stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // Load module before calling the kernel + print_tensor_Kernel_Module_t module; + print_tensor_Kernel_Module_Load(&module); + + // Call the kernel; the kernel wrapper function is defined in the header file + cute_dsl_print_tensor_wrapper(&module, &tensor_a, stream); + + // Cleanup + print_tensor_Kernel_Module_Unload(&module); + cudaStreamDestroy(stream); + } + +The ``print_tensor_example.h`` header file is generated by the ``export_to_c`` interface. It includes: + +* The ``print_tensor_Kernel_Module_t`` type: Represents the kernel module. +* The ``print_tensor_Tensor_a_t`` type: A tensor-specific type that defines the ABI for a particular CuTe tensor. +* The ``cute_dsl_print_tensor_wrapper`` function: The user-facing entry point to invoke the kernel. + +The compilation of the C++ executable requires the ``libcuda_dialect_runtime.so`` or ``libcuda_dialect_runtime_static.a`` library which is involved in ``/lib``, along with the CUDA driver and runtime libraries, to function properly. + + +C++ Integration with Dynamic Loading +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Dynamically load pre-compiled object files or shared libraries at runtime. By including the ``CuteDSLRuntime.h`` header, you can load the module, look up exported functions, and invoke them. + +.. code-block:: cpp + + #include "CuteDSLRuntime.h" + #include + + void run_print_tensor() { + // Load module from shared library + CuteDSLRT_Module_t *module = nullptr; + CuteDSLRT_Error_t err = CuteDSLRT_Module_Load( + &module, + "./artifacts/libprint_tensor_example.so" + ); + // or + CuteDSLRT_Error_t err = CuteDSLRT_Module_Load( + &module, + "./artifacts/print_tensor_example.o" + ); + check_error(err); + + // Lookup function + CuteDSLRT_Function_t *func = nullptr; + err = CuteDSLRT_Module_Get_Function(&func, module, "print_tensor"); + check_error(err); + + // Prepare arguments, matching the argument type defined in the header file + typedef struct { + void *data; + int32_t dynamic_shapes[2]; + int64_t dynamic_strides[1]; + } print_tensor_Tensor_a_t; + + print_tensor_Tensor_a_t tensor_a; + tensor_a.data = nullptr; + tensor_a.dynamic_shapes[0] = 32; + tensor_a.dynamic_shapes[1] = 16; + tensor_a.dynamic_strides[0] = 16; + + // Create stream + cudaStream_t stream; + cudaStreamCreate(&stream); + + // Call the function; the runtime function accepts packed arguments, refer to the wrapper in the header file + int ret; + void* args[] = {&tensor_a, &stream, &ret}; + err = CuteDSLRT_Function_Run(func, args, 3); + check_error(err); + cudaStreamSynchronize(stream); + + // Cleanup + CuteDSLRT_Module_Destroy(module); + cudaStreamDestroy(stream); + } + +The ``CuteDSLRuntime.h`` header file can be found in ``/include``. It includes: + +* The ``CuteDSLRT_Error_t`` type: Indicates error status. +* The ``CuteDSLRT_Module_Load`` function: Loads the module. +* The ``CuteDSLRT_Module_Get_Function`` function: Gets a function from the loaded module. The runtime API will load the CUDA module for kernel execution. +* The ``CuteDSLRT_Function_Run`` function: Runs the function. +* The ``CuteDSLRT_Module_Destroy`` function: Destroys the module. + +The compilation of the C++ executable requires the ``libcute_dsl_runtime.so`` library which is involved in ``/lib``, along with the CUDA driver and runtime libraries, to function properly. + +Supported Argument Types +------------------------ + +|DSL| supports the following argument types: + +* ``cute.Tensor`` +* ``cute.Shape`` / ``cute.Coord`` / ``cute.Tile`` / ``cute.IntTuple`` / ``cute.Stride`` +* ``cuda.CUstream`` +* ``cutlass.Int8`` / ``cutlass.Int16`` / ``cutlass.Int32`` / ``cutlass.Int64`` / ``cutlass.Boolean`` +* ``cutlass.Uint8`` / ``cutlass.Uint16`` / ``cutlass.Uint32`` / ``cutlass.Uint64`` +* ``cutlass.Float32`` / ``cutlass.TFloat32`` / ``cutlass.Float64`` / ``cutlass.Float16`` + +Note that: + +1. ``cute.Tensor`` is a dynamic tensor type that only contains dynamic shapes and strides in its ABI representation. As a result, different compilations may produce different tensor ABIs. This is why declarations for each tensor type are included in the generated header file. +2. ``strides`` in ``cute.Tensor`` are determined by the ``use_32bit_strides`` compile argument. When ``use_32bit_strides`` is set to ``True``, the strides are 32-bit; when set to ``False``, they are 64-bit. +3. Currently, custom types are not supported for AOT compilation. + +Object File Compatibility Issues +-------------------------------- + +The object file generated by |DSL| depends on the CUDA runtime library. Therefore, ensure that the version of the CUDA runtime/toolkit library matches the version used by |DSL|. Otherwise, ABI compatibility with the CUDA runtime cannot be guaranteed. + +When using C++ static linking integration, compatibility is assured because the header and object files are generated together and guaranteed to match. + +For C++ dynamic loading integration and Python loading, the binary file is loaded at runtime. To ensure compatibility, version information is embedded in the metadata of the generated binary file. At runtime, this version information is checked, and if it does not match the expected version, the binary file will be rejected. + +Relation to Apache TVM FFI AOT +------------------------------ + +Apache TVM FFI AOT offers a comparable capability, enabling TVM functions to be compiled into binary files that can be loaded and executed at runtime. +For more information, see the section "Exporting Compiled Module" in :doc:`compile_with_tvm_ffi`. + +The primary distinction is that, when TVM FFI is enabled, |DSL| generates a dedicated wrapper function on top of the underlying CuTe ABI. This wrapper adheres to the calling conventions defined by TVM FFI. +In contrast, the CuTe ABI entry function is specified directly in the generated header file, which affects how arguments must be provided. + +For instance, with the TVM FFI wrapper function, users are able to pass in arguments such as ``torch.Tensor`` directly. However, when calling the CuTe ABI entry function, arguments should be provided as ``cute.Tensor`` types. diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_control_flow.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_control_flow.rst index f5c2ea35..02f83463 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_control_flow.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_control_flow.rst @@ -184,6 +184,46 @@ Standard Python ``while`` is supported. n += 1 +Summary of Control Flow behavior +--------------------------------- + +.. list-table:: + :header-rows: 1 + :widths: 30 30 30 + + * - **Control Flow** + - **Run time evaluation** + - **Compile time evaluation** + + * - if cutlass.const_expr() + - ❌ + - ✅ + + * - if pred + - ✅ + - ❌ + + * - while cutlass.const_expr() + - ❌ + - ✅ + + * - while pred + - ✅ + - ❌ + + * - for i in cutlass.range_constexpr() + - ❌ + - ✅ + + * - for i in range() + - ✅ + - ❌ + + * - for i in cutlass.range() (support advanced unrolling and pipelining) + - ✅ + - ❌ + + Compile-Time Metaprogramming ---------------------------- diff --git a/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst b/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst index 5f7e7598..8da0c643 100644 --- a/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst +++ b/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.rst @@ -73,7 +73,7 @@ You can use the following code to specify compilation options: jit_executor_with_opt_level_2 = cute.compile(add, 1, 2, options="--opt-level 2") jit_executor_with_opt_level_1 = cute.compile(add, 1, 2, options="--opt-level 1") - jit_executor_with_enable_device_assertions = cute.compile(add, 1, 2, options="--enable-assertions") + jit_executor_with_enable_assertions = cute.compile(add, 1, 2, options="--enable-assertions") jit_executor_with_keep_cubin = cute.compile(add, 1, 2, options="--keep-cubin") jit_executor_with_keep_ptx = cute.compile(add, 1, 2, options="--keep-ptx") jit_executor_with_ptxas_options = cute.compile(add, 1, 2, options="--ptxas-options '--opt-level=2'") @@ -100,7 +100,7 @@ Notebly, boolean options are automatically converted to True instances of the op jit_executor_with_opt_level_2 = cute.compile[OptLevel(2)](add, 1, 2) jit_executor_with_opt_level_1 = cute.compile[OptLevel(1)](add, 1, 2) - jit_executor_with_enable_device_assertions = cute.compile[EnableAssertions](add, 1, 2) + jit_executor_with_enable_assertions = cute.compile[EnableAssertions](add, 1, 2) jit_executor_with_keep_cubin = cute.compile[KeepCUBIN](add, 1, 2) jit_executor_with_keep_ptx = cute.compile[KeepPTX](add, 1, 2) - jit_executor_with_ptxas_options = cute.compile[PtxasOptions("--opt-level=2")](add, 1, 2) \ No newline at end of file + jit_executor_with_ptxas_options = cute.compile[PtxasOptions("--opt-level=2")](add, 1, 2) diff --git a/media/docs/pythonDSL/deprecation.rst b/media/docs/pythonDSL/deprecation.rst new file mode 100644 index 00000000..dd6ad881 --- /dev/null +++ b/media/docs/pythonDSL/deprecation.rst @@ -0,0 +1,56 @@ +.. _deprecation: + +Deprecation Policy +================== + +Purpose +------- + +The goal of this policy is to evolve the DSL and its APIs while keeping user +programs stable. Features or APIs are deprecated only when they are redundant, +unsafe, or block better designs. + +Deprecation Process +------------------- + +**Step 1 — Soft Deprecation** + +When a feature is considered for removal, it is first annotated with the +``@deprecated`` decorator or ``DeprecationWarning`` and documented with a +suggested alternative. At this stage, the feature continues to work normally. + +Users are encouraged to provide feedback and describe their use cases. +If there is strong justification, we may keep or redesign the feature. + +**Step 2 — Removal (the subsequent release)** + +If no valid use cases remain, the deprecated feature will be removed in the +following **minor** release. + +.. note:: + + The release version follows the format ``..``. + +Communication +------------- + +All deprecations are announced through: + +* This page +* In-code warning messages + +Soft Deprecations +----------------- + +**Version 4.2.1** + +* ``cute.arch.warpgroup_reg_alloc`` and ``cute.arch.warpgroup_reg_dealloc`` + → Scheduled for deprecation. Use ``cute.arch.setmaxregister_increase`` and ``cute.arch.setmaxregister_decrease`` instead. + +* ``alignment`` argument in ``CooperativeGroup`` constructor + → Scheduled for deprecation. It was unused; no replacement is suggested. + +Deprecated Features +------------------- + +*(None currently.)* diff --git a/media/docs/pythonDSL/limitations.rst b/media/docs/pythonDSL/limitations.rst index 23673195..f396e4f5 100644 --- a/media/docs/pythonDSL/limitations.rst +++ b/media/docs/pythonDSL/limitations.rst @@ -17,12 +17,8 @@ the DSL. Notable unsupported features ---------------------------- -- Programmatic Dependent Launch (PDL) - convolutions -- full support for ahead of time compilation - preferred clusters -- CLC-based tile schedulers -- EVT support - Windows support Programming Model @@ -82,16 +78,30 @@ Programming Model xs.append(Float32(1.0)) **Python Function** - The DSL currently does not implement support for return values from Python functions, - although this capability is planned for future releases. + The DSL currently has **limited support for return values** from Python functions. + At the moment, only ``constexpr`` values can be returned, while returning **dynamic values** is **not yet supported**. + This capability is planned for a future release. Example: - .. code:: python + .. code-block:: python @cute.jit - def foo(): - return 1 # Currently unsupported in CuTe DSL + def baz(a: cutlass.Constexpr): + return a + 1 + + @cute.jit + def foo(a: cutlass.Int32): + return a + 1 + + @cute.jit + def bar(a: cutlass.Int32): + val = foo(a) # works + + val = baz(10) # works + val = bar(10) # works + foo(10) # currently unsupported in CuTe DSL + **Expression or Statement with Dependent Type** CuTe DSL implements static typing and does not support dependent types. diff --git a/media/docs/pythonDSL/overview.rst b/media/docs/pythonDSL/overview.rst index fbd3abd8..c98bea59 100644 --- a/media/docs/pythonDSL/overview.rst +++ b/media/docs/pythonDSL/overview.rst @@ -105,4 +105,4 @@ You can: - Propose support for additional data types or kernel variants - Help prioritize roadmap features by upvoting GitHub issues -Thank you for helping shape the future of CUTLASS DSLs! \ No newline at end of file +Thank you for helping shape the future of CUTLASS DSLs! diff --git a/media/docs/pythonDSL/quick_start.rst b/media/docs/pythonDSL/quick_start.rst index 495ec8ff..cd6dfc2f 100644 --- a/media/docs/pythonDSL/quick_start.rst +++ b/media/docs/pythonDSL/quick_start.rst @@ -3,28 +3,39 @@ Quick Start Guide ======================= -The CUTLASS DSL 4.0 release currently supports **Linux** and **Python 3.12** only. To install CUTLASS DSLs (limited to CuTe DSL for now), use the following command +The CUTLASS DSL 4.4 release currently supports **Linux** and **Python 3.10 - 3.13** only. To install CUTLASS DSLs (limited to CuTe DSL for now), use the following command Installation ----------------------- To ensure compatibility with the examples and code on `GitHub `_, -use the `requirements.txt `_ file from the corresponding commit in the repository. +use the `setup.sh `_ file from the corresponding commit in the repository. .. code-block:: bash - git clone https://github.com/NVIDIA/cutlass.git - pip install -r cutlass/python/CuTeDSL/requirements.txt - -If you just want to try out the last known stable release of the CUTLASS DSL (may not compatible with the latest examples and code), run: + git clone https://github.com/NVIDIA/cutlass.git + + # For CUDA Toolkit 12.9: + ./cutlass/python/CuTeDSL/setup.sh --cu12 + + # For CUDA Toolkit 13.1: + ./cutlass/python/CuTeDSL/setup.sh --cu13 + +If you just want to try out the last known stable release of the CUTLASS DSL (may not be compatible with the latest examples and code), run: + .. code-block:: bash + # For CUDA Toolkit 12.9: pip install nvidia-cutlass-dsl + # For CUDA Toolkit 13.1: + pip install nvidia-cutlass-dsl[cu13] + + The ``nvidia-cutlass-dsl`` wheel includes everything needed to generate GPU kernels. It requires -the same NVIDIA driver version as the -`CUDA Toolkit 12.9 `_. +the same NVIDIA driver version as the corresponding `CUDA Toolkit `_ +(CUDA Toolkit 12.9 or CUDA Toolkit 13.1). Recommended Dependencies --------------------------------- @@ -35,6 +46,8 @@ To run examples and begin development, we recommend installing: pip install torch jupyter +We recommend installing JAX with CUDA support at version 0.8.1 to run JAX examples. + Recommended Python environment variables for jupyter notebooks -------------------------------------------------------------- diff --git a/python/CuTeDSL/cutlass/__init__.py b/python/CuTeDSL/cutlass/__init__.py index c09463fc..3df1c5c2 100644 --- a/python/CuTeDSL/cutlass/__init__.py +++ b/python/CuTeDSL/cutlass/__init__.py @@ -13,6 +13,8 @@ from ._mlir._mlir_libs import _cutlass_ir _cutlass_ir.populate(_cutlass_ir) +__version__ = "@CUTLASS_IR_WHEEL_RELEASE_VERSION@" + from .cutlass_dsl import ( Constexpr, dsl_user_op, @@ -44,6 +46,7 @@ from .cutlass_dsl import ( # Construction utilities for user-defined classes extract_mlir_values, new_from_mlir_values, + DSLCudaVersion, ) from .cute.typing import * @@ -51,6 +54,7 @@ from .cute.typing import * # Utilities not belonging to CuTe from . import utils as utils from . import pipeline as pipeline +from .utils.version_info import CUDA_VERSION # Used as internal symbol from . import cutlass_dsl as _dsl @@ -61,4 +65,7 @@ register_jit_arg_adapter = _dsl.JitArgAdapterRegistry.register_jit_arg_adapter gpu = _dsl.cutlass_gpu cuda = _dsl.cuda_helpers +# Jax Framework support +from . import jax as jax + CACHE_FILE = "compiled_cache.db" diff --git a/python/CuTeDSL/cutlass/base_dsl/__init__.py b/python/CuTeDSL/cutlass/base_dsl/__init__.py index e3fcbc77..2b4927fa 100644 --- a/python/CuTeDSL/cutlass/base_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/base_dsl/__init__.py @@ -15,3 +15,13 @@ from .runtime import * from ._mlir_helpers import lru_cache_ir, dsl_user_op from .env_manager import get_str_env_var, detect_gpu_arch +from .utils.tree_utils import ( + is_constexpr_field, + tree_flatten, + tree_unflatten, + PyTreeDef, + is_frozen_dataclass, + DSLTreeFlattenError, +) + +from .common import DSLCudaVersion diff --git a/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py b/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py index ef06f8eb..01070140 100644 --- a/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py +++ b/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py @@ -294,7 +294,8 @@ def const(value, ty=None, *, loc=None, ip=None): Generates dynamic expression for constant values. """ from ..typing import Numeric, NumericMeta - from ..dsl import is_dynamic_expression, _numpy_type_to_mlir_type + from ..dsl import is_dynamic_expression + from ..utils.numpy import _numpy_type_to_mlir_type if isinstance(value, Numeric): value = value.value @@ -693,9 +694,9 @@ def _min(lhs, rhs, *, loc=None, ip=None): if isinstance(lhs.type, T.VectorType): elem_type = lhs.type.element_type if arith._is_integer_like_type(elem_type): - assert hasattr( - lhs, "signed" - ), "Should have attribute `signed`, must be a bug" + assert hasattr(lhs, "signed"), ( + "Should have attribute `signed`, must be a bug" + ) if lhs.signed != False: return arith.minsi(lhs, rhs, loc=loc, ip=ip) else: @@ -719,6 +720,7 @@ def _max(lhs, rhs, *, loc=None, ip=None): Assuming the operands have the same type """ from ..dsl import is_dynamic_expression + if not is_dynamic_expression(lhs): if not is_dynamic_expression(rhs): return max(lhs, rhs) @@ -731,9 +733,9 @@ def _max(lhs, rhs, *, loc=None, ip=None): if isinstance(lhs.type, T.VectorType): elem_type = lhs.type.element_type if isinstance(elem_type, ir.IntegerType): - assert hasattr( - lhs, "signed" - ), "Should have attribute `signed`, must be a bug" + assert hasattr(lhs, "signed"), ( + "Should have attribute `signed`, must be a bug" + ) if lhs.signed != False: return arith.maxsi(lhs, rhs, loc=loc, ip=ip) else: diff --git a/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/op.py b/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/op.py index 3ba565ea..ad7893ee 100644 --- a/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/op.py +++ b/python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/op.py @@ -17,6 +17,7 @@ import inspect from functools import wraps from ..._mlir import ir +from ..common import DSLRuntimeError def dsl_user_op(opFunc): @@ -57,7 +58,26 @@ def dsl_user_op(opFunc): ), childLoc=file_loc, ) - res_or_list = opFunc(*args, **kwargs, loc=loc) + + try: + res_or_list = opFunc(*args, **kwargs, loc=loc) + except TypeError as e: + # Provide a helpful error message when function doesn't accept 'loc' + func_name = getattr(opFunc, "__name__", str(opFunc)) + if "unexpected keyword argument 'loc'" in str(e): + raise DSLRuntimeError( + f"Function '{func_name}' decorated with @dsl_user_op does not accept the required 'loc' parameter.", + suggestion=[ + f"1. Add 'loc=None' as a keyword-only parameter to {func_name}:", + f" def {func_name}(..., *, loc=None):", + "", + f"2. Remove the @dsl_user_op decorator if location tracking is not needed", + ], + cause=e, + ) from e + else: + # Re-raise other TypeErrors as-is + raise return res_or_list return wrapper diff --git a/python/CuTeDSL/cutlass/base_dsl/arch.py b/python/CuTeDSL/cutlass/base_dsl/arch.py index 2a69946e..76b75d19 100644 --- a/python/CuTeDSL/cutlass/base_dsl/arch.py +++ b/python/CuTeDSL/cutlass/base_dsl/arch.py @@ -44,7 +44,6 @@ class Arch(Enum): sm_121 = (12, 1, "") sm_121a = (12, 1, "a") sm_121f = (12, 1, "f") - def __init__(self, major, minor, suffix): self.major = major self.minor = minor @@ -59,7 +58,6 @@ class Arch(Enum): return cls((major, minor, suffix)) else: raise ValueError(f"invalid arguments for Arch: {value}") - # attributes to get arch list of specific families @classmethod diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py index 9161127d..c434c982 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py @@ -14,7 +14,8 @@ This module provides helper functions that are generated by the preprocessor. The preprocessor read through python's ast and changes the input code. """ -from typing import Callable, Iterator, Optional, overload +from dataclasses import dataclass +from typing import Callable, Iterator, Optional, overload, List from typing_extensions import deprecated import warnings import inspect @@ -111,6 +112,7 @@ class Executor: unroll=-1, unroll_full=False, prefetch_stages=None, + vectorize=None, ): assert self._loop_execute_range_dynamic, ( "Functions must be set before execution." @@ -128,6 +130,7 @@ class Executor: unroll, unroll_full, prefetch_stages, + vectorize, ) def if_execute( @@ -153,7 +156,6 @@ class Executor: def while_execute( self, - pred, while_before_block: Callable, while_after_block: Callable, write_args=[], @@ -190,9 +192,10 @@ def loop_selector( unroll=-1, unroll_full=False, prefetch_stages=None, + vectorize=None, ): log().debug( - "start [%s] stop [%s] step [%s] write_args [%s] full_write_args_count [%s] write_args_names [%s] unroll [%s] unroll_full [%s] prefetch_stages [%s]", + "start [%s] stop [%s] step [%s] write_args [%s] full_write_args_count [%s] write_args_names [%s] unroll [%s] unroll_full [%s] prefetch_stages [%s] vectorize [%s]", start, stop, step, @@ -202,6 +205,7 @@ def loop_selector( unroll, unroll_full, prefetch_stages, + vectorize, ) from .typing import Integer, Numeric @@ -227,6 +231,7 @@ def loop_selector( unroll, unroll_full, prefetch_stages, + vectorize, ) return ir_loop @@ -247,15 +252,14 @@ def if_selector(pred, write_args=[]): return ir_loop -def while_selector(pred, write_args=[]): +def while_selector(*, write_args=[]): def ir_while_loop(func): - return func(pred, *write_args) + return func(*write_args) return ir_while_loop def while_executor( - pred, while_before_block: Callable, while_after_block: Callable, write_args=[], @@ -263,7 +267,6 @@ def while_executor( write_args_names=[], ): return executor.while_execute( - pred, while_before_block, while_after_block, write_args, @@ -310,15 +313,25 @@ class range: - unroll: Number of iterations to unroll (0 or 1 = no unrolling) - unroll_full: Whether to fully unroll the loop - prefetch_stages: Number of prefetch stages to generate + - vectorize: Whether to vectorize the loop (default: None) """ @overload - def __new__(cls, stop, unroll=0, unroll_full=False, prefetch_stages=None): + def __new__( + cls, stop, unroll=0, unroll_full=False, prefetch_stages=None, vectorize=None + ): pass @overload def __new__( - cls, start, stop, step, unroll=0, unroll_full=False, prefetch_stages=None + cls, + start, + stop, + step, + unroll=0, + unroll_full=False, + prefetch_stages=None, + vectorize=None, ): pass @@ -581,7 +594,11 @@ def copy_members(dest, src): or not hasattr(src, name) ): continue - setattr(dest, name, getattr(src, name)) + try: + setattr(dest, name, getattr(src, name)) + except AttributeError: + # AttributeError means the member is no settable, like property without setter, just ignore + pass def get_locals_or_none(locals, symbols): @@ -609,3 +626,109 @@ def closure_check(closures): f"Function `{closure.__name__}` is a closure that captures variables and is not supported in dynamic control flow", suggestion="Please implicitly pass in captured variables as arguments", ) + + +@dataclass +class FormattedValue: + """ + Represents a value (from an f-string interpolation) along with its conversion and format specification. + + This is a runtime representation of ast.FormattedValue. + See https://docs.python.org/3/library/ast.html#ast.FormattedValue + + This class is used to encapsulate a value that is to be embedded into an f-string formatted string, + along with any conversion flags (such as '!s', '!r', or '!a') and a list of format specification strings + (e.g., '.2f', '05d', etc). + + Attributes: + value (Any): The value to be formatted. + conversion (int): The ASCII code of the conversion character, defaulting to -1 (no conversion). + format_spec (Optional[List[str]]): List of format specifier strings. + """ + + value: Any + conversion: int = -1 + format_spec: Optional[List[str]] = None + + def to_str(self): + """ + Converts the FormattedValue into a format string component and (optional) dynamic argument. + + If the value is a dynamic expression (requires runtime evaluation), returns a printf-style placeholder and + the corresponding dynamic value for later substitution. For static (non-dynamic) values, returns the fully + rendered literal string and None for the dynamic value. + + Returns: + Tuple[str, Optional[Any]]: A tuple where the first element is the format string component and the second + element is the dynamic value (or None if not required). + + Raises: + DSLRuntimeError: If the format spec is invalid or not supported for dynamic expressions. + """ + if executor._is_dynamic_expression(self.value): + if self.conversion != -1: + warnings.warn( + "Conversion may not be honored for dynamic expressions", + stacklevel=3, + ) + if self.format_spec is not None and len(self.format_spec) > 1: + raise DSLRuntimeError( + "At most one format specifier is supported for dynamic expressions", + ) + if self.format_spec is not None: + if self.format_spec[0].startswith(("<", ">", "^")): + raise DSLRuntimeError( + "Alignment specifier is not supported for dynamic expressions", + ) + return f"%{self.format_spec[0]}", self.value + return "{}", self.value + else: + conversion_str = "" + if self.conversion == 97: + conversion_str = "!a" # ASCII + elif self.conversion == 114: + conversion_str = "!r" # repr + elif self.conversion == 115: + conversion_str = "!s" # str + if self.format_spec is not None: + f_str = f"{{{conversion_str}:{self.format_spec[0]}}}" + else: + f_str = f"{{{conversion_str}}}" + return f_str.format(self.value), None + + +def fstring_decompose(joinedStrComponent): + """ + Decomposes a joined f-string component list into a format string and dynamic arguments. + + This function takes a sequence of components representing parts of an f-string, + which may include both literal strings and FormatValue instances. + It assembles a format string with placeholders for dynamic values + and collects the actual dynamic argument values in order. + + Args: + joinedStrComponent (List[Union[str, FormattedValue]]): The components of the f-string. + + Returns: + Tuple[str, ...]: A tuple containing the format string as the first element, + followed by all dynamic argument values. + + Raises: + DSLAstPreprocessorError: If an unsupported component type is encountered. + """ + format_string = "" + dynamic_args = [] + for component in joinedStrComponent: + if isinstance(component, FormattedValue): + s, v = component.to_str() + format_string += s + if v is not None: + # This is a dynamic value, add to dynamic_args + dynamic_args.append(v) + elif isinstance(component, str): + format_string += component + else: + raise DSLAstPreprocessorError( + f"Unsupported component type in f-string: {type(component)}", + ) + return (format_string, *dynamic_args) diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py index 36c2ecb0..4afc61b6 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py @@ -371,7 +371,11 @@ class DSLPreprocessor(ast.NodeTransformer): package_name = module.__package__.rsplit( ".", child_node.level - 1 )[0] - module_name = f"{package_name}.{module_name}" + # For `from . import x`, module name is None, just use package name + if module_name: + module_name = f"{package_name}.{module_name}" + else: + module_name = package_name else: # Handle typically some local import like: # from .common_dense_gemm import DenseGemmKernel @@ -423,12 +427,15 @@ class DSLPreprocessor(ast.NodeTransformer): # Get the module containing the decorated function if module := inspect.getmodule(decorated_func): + if module in self.module_cache: + return self.module_cache[module] try: # Get the module source code source = inspect.getsource(module) module_ast = ast.parse(source) imports = self._get_imports_from_ast(module_ast, module) + self.module_cache[module] = imports except (IOError, TypeError): pass @@ -768,7 +775,7 @@ class DSLPreprocessor(ast.NodeTransformer): return unified_tree def analyze_region_variables( - self, node: Union[ast.For, ast.If], active_symbols: List[Set[str]] + self, node: Union[ast.For, ast.If, ast.While], active_symbols: List[Set[str]] ): """ Analyze variables in different code regions to identify read-only, write-only, @@ -857,7 +864,14 @@ class DSLPreprocessor(ast.NodeTransformer): self.generic_visit(node) analyzer = RegionAnalyzer() - analyzer.visit(ast.Module(body=node)) + analyzer.visit(ast.Module(body=node.body)) + if node.orelse: + analyzer.visit(ast.Module(body=node.orelse)) + + # While's loop condition is executed n times, as loop body + # So collect the variables used in the loop condition + if isinstance(node, ast.While): + analyzer.visit(ast.Module(body=node.test)) # If arg is both write and invoke, remove from invoked_args invoked_args = invoked_args - write_args @@ -917,6 +931,10 @@ class DSLPreprocessor(ast.NodeTransformer): return keywords.get("pipelining", ast.Constant(value=None)) return keywords.get("prefetch_stages", ast.Constant(value=None)) + def extract_vectorize_args(self, iter_node): + keywords = {kw.arg: kw.value for kw in iter_node.keywords} + return keywords.get("vectorize", ast.Constant(value=None)) + def create_loop_function( self, func_name, @@ -927,6 +945,7 @@ class DSLPreprocessor(ast.NodeTransformer): unroll, unroll_full, prefetch_stages, + vectorize, write_args, full_write_args_count, ): @@ -972,6 +991,7 @@ class DSLPreprocessor(ast.NodeTransformer): ast.keyword(arg="unroll", value=unroll), ast.keyword(arg="unroll_full", value=unroll_full), ast.keyword(arg="prefetch_stages", value=prefetch_stages), + ast.keyword(arg="vectorize", value=vectorize), ast.keyword( arg="write_args", value=self.generate_get_locals_or_none_call(write_args), @@ -1387,6 +1407,7 @@ class DSLPreprocessor(ast.NodeTransformer): start_expr, stop_expr, step_expr, has_step = self.extract_range_args(node.iter) unroll, unroll_full = self.extract_unroll_args(node.iter) prefetch_stages = self.extract_prefetch_stages_args(node.iter) + vectorize = self.extract_vectorize_args(node.iter) write_args, full_write_args_count, called_closures = ( self.analyze_region_variables(node, active_symbols) ) @@ -1416,6 +1437,7 @@ class DSLPreprocessor(ast.NodeTransformer): unroll, unroll_full, prefetch_stages, + vectorize, write_args, full_write_args_count, ) @@ -1462,6 +1484,67 @@ class DSLPreprocessor(ast.NodeTransformer): ast.copy_location(new_node, node) return new_node + def processFormattedValue(self, node): + """ + Converts an ast.FormattedValue node into a runtime representation of an ast.FormattedValue. + + This function takes an ast.FormattedValue node and converts it into a runtime representation of ast.FormattedValue. + """ + keywords = [] + if node.conversion != -1: + keywords.append( + ast.keyword(arg="conversion", value=ast.Constant(value=node.conversion)) + ) + if node.format_spec: + keywords.append( + ast.keyword( + arg="format_spec", + value=ast.List(elts=node.format_spec.values, ctx=ast.Load()), + ) + ) + call = ast.Call( + func=_create_module_attribute( + "FormattedValue", + lineno=node.lineno, + col_offset=node.col_offset, + ), + args=[node.value], + keywords=keywords, + ) + return ast.copy_location(call, node) + + def processFString(self, node): + """ + Converts an f-string node into a runtime representation of an f-string. + + This function takes an ast.JoinedStr node and converts it into a list of elements, + where each element is either a literal string or a FormattedValue. + The FormattedValue is converted into a runtime representation of ast.FormattedValue. + """ + elements = [] + joinedStr = node.args[0] + for component in joinedStr.values: + if isinstance(component, ast.Constant): + elements.append(component) + elif isinstance(component, ast.FormattedValue): + elements.append(self.processFormattedValue(component)) + else: + raise DSLAstPreprocessorError( + f"Unsupported component type in f-string: {type(component)}", + filename=self.session_data.file_name, + snippet=ast.unparse(component), + ) + call = ast.Call( + func=_create_module_attribute( + "fstring_decompose", + lineno=node.lineno, + col_offset=node.col_offset, + ), + args=[ast.copy_location(ast.List(elts=elements, ctx=ast.Load()), node)], + keywords=[], + ) + return ast.copy_location(call, node) + def visit_Call(self, node): func = node.func # Visit args and kwargs @@ -1533,45 +1616,62 @@ class DSLPreprocessor(ast.NodeTransformer): ), node, ) - elif isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): - - def create_downcast_call(arg): - return ast.copy_location( - ast.Call( - func=_create_module_attribute( - self.IMPLICIT_DOWNCAST_NUMERIC_TYPE, - submodule_name="typing", - lineno=node.lineno, - col_offset=node.col_offset, - ), - args=[arg], - keywords=[], - ), - arg, - ) - - module = self.session_data.function_globals.get(func.value.id) - if isinstance(module, ModuleType) and module.__package__.endswith( - "._mlir.dialects" + elif ( + func.id == "printf" + and len(node.args) > 0 + and isinstance(node.args[0], ast.JoinedStr) ): - # Check if argument is Numeric, if so, call ir_value() - args = [] - for arg in node.args: - args.append(create_downcast_call(arg)) - kwargs = [] - for kwarg in node.keywords: - kwargs.append( - ast.copy_location( - ast.keyword( - arg=kwarg.arg, - value=create_downcast_call(kwarg.value), + node.args = [ + ast.Starred(value=self.processFString(node), ctx=ast.Load()) + ] + elif isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): + if ( + func.attr == "printf" + and len(node.args) > 0 + and isinstance(node.args[0], ast.JoinedStr) + ): + node.args = [ + ast.Starred(value=self.processFString(node), ctx=ast.Load()) + ] + else: + + def create_downcast_call(arg): + return ast.copy_location( + ast.Call( + func=_create_module_attribute( + self.IMPLICIT_DOWNCAST_NUMERIC_TYPE, + submodule_name="typing", + lineno=node.lineno, + col_offset=node.col_offset, ), - kwarg, - ) + args=[arg], + keywords=[], + ), + arg, + ) + + module = self.session_data.function_globals.get(func.value.id) + if isinstance(module, ModuleType) and module.__package__.endswith( + "._mlir.dialects" + ): + # Check if argument is Numeric, if so, call ir_value() + args = [] + for arg in node.args: + args.append(create_downcast_call(arg)) + kwargs = [] + for kwarg in node.keywords: + kwargs.append( + ast.copy_location( + ast.keyword( + arg=kwarg.arg, + value=create_downcast_call(kwarg.value), + ), + kwarg, + ) + ) + return ast.copy_location( + ast.Call(func=func, args=args, keywords=kwargs), node ) - return ast.copy_location( - ast.Call(func=func, args=args, keywords=kwargs), node - ) else: node.func = self.visit(node.func) @@ -1602,6 +1702,11 @@ class DSLPreprocessor(ast.NodeTransformer): self.generic_visit(node) return node + def visit_AnnAssign(self, node): + self._visit_target(node.target) + self.generic_visit(node) + return node + def visit_Name(self, node): isLoad = isinstance(node.ctx, ast.Load) if node.id in ["max", "min", "any", "all"] and isLoad: @@ -1623,22 +1728,13 @@ class DSLPreprocessor(ast.NodeTransformer): self.generic_visit(node) return node - def check_decorator(self, node: ast.AST) -> bool: - """ - Check if the function has the correct decorator for preprocessing. - """ - if not isinstance(node, ast.FunctionDef): - return False - decorator_list = node.decorator_list - if len(decorator_list) == 0: - return False - - for d in decorator_list: + def get_dsl_decorator_index(self, decorator_list): + for i, d in enumerate(decorator_list): if isinstance(d, ast.Call): if isinstance(d.func, ast.Attribute): if d.func.attr in ["jit", "kernel"]: if d.keywords == []: - return True + return i for keyword in d.keywords: if keyword.arg == "preprocess": try: @@ -1651,9 +1747,32 @@ class DSLPreprocessor(ast.NodeTransformer): elif isinstance(d, ast.Attribute): if d.attr in ["jit", "kernel"]: - return True + return i + return None - return False + def check_decorator(self, node: ast.AST) -> bool: + """ + Check if the function has the correct decorator for preprocessing. + """ + if not isinstance(node, ast.FunctionDef): + return False + decorator_list = node.decorator_list + if len(decorator_list) == 0: + return False + + dsl_decorator_index = self.get_dsl_decorator_index(decorator_list) + + if ( + dsl_decorator_index is not None + and dsl_decorator_index < len(decorator_list) - 1 + ): + decorator = ast.unparse(decorator_list[dsl_decorator_index]) + raise DSLAstPreprocessorError( + f"`{decorator}` decorator must be the inner most decorator", + suggestion=f"Please move the `{decorator}` decorator to the inner most position", + ) + + return dsl_decorator_index is not None def remove_dsl_decorator(self, decorator_list): """ @@ -1707,6 +1826,9 @@ class DSLPreprocessor(ast.NodeTransformer): # Remove .jit and .kernel decorators node.decorator_list = self.remove_dsl_decorator(node.decorator_list) + + # Remove return annotation from processed AST to avoid symbol requirement + node.returns = None return node def visit_With(self, node): @@ -2181,11 +2303,9 @@ class DSLPreprocessor(ast.NodeTransformer): return write_args """ test_expr = self.visit(node.test) - pred_name = self.make_func_param_name("pred", write_args) # Section: decorator construction decorator_keywords = [ - ast.keyword(arg="pred", value=test_expr), ast.keyword( arg="write_args", value=self.generate_get_locals_or_none_call(write_args), @@ -2260,7 +2380,6 @@ class DSLPreprocessor(ast.NodeTransformer): # Section: Execute via executor execute_keywords = [ - ast.keyword(arg="pred", value=ast.Name(id=pred_name, ctx=ast.Load())), ast.keyword( arg="write_args", value=ast.List( @@ -2298,8 +2417,7 @@ class DSLPreprocessor(ast.NodeTransformer): ) # Putting everything together, FunctionDef for while_region - func_args_args = [ast.arg(arg=pred_name, annotation=None)] - func_args_args += [ast.arg(arg=var, annotation=None) for var in write_args] + func_args_args = [ast.arg(arg=var, annotation=None) for var in write_args] func_args = ast.arguments( posonlyargs=[], args=func_args_args, diff --git a/python/CuTeDSL/cutlass/base_dsl/cache_helpers.py b/python/CuTeDSL/cutlass/base_dsl/cache_helpers.py index 1d327f45..ce920311 100644 --- a/python/CuTeDSL/cutlass/base_dsl/cache_helpers.py +++ b/python/CuTeDSL/cutlass/base_dsl/cache_helpers.py @@ -13,21 +13,24 @@ This module provides jit cache load/dump helper functions """ +from collections import OrderedDict import os import io import sys import uuid import random import tempfile -import pwd import time +from typing import Any, Optional from pathlib import Path import hashlib from functools import lru_cache +import weakref import zlib from .utils.logger import log from .jit_executor import JitCompiledFunction +from .common import DSLRuntimeError from .._mlir import ir @@ -36,16 +39,28 @@ from .._mlir import ir # ============================================================================= + def get_current_user(): """ Get the current user. This is used to determine the path to the cache directory. """ - # Try to get the user from the environment variable first user = os.getenv("USER") or os.getenv("USERNAME") - if not user: - # Fallback for Unix-like systems - user = pwd.getpwuid(os.getuid()).pw_name - return user + if user: + return user + # Try Unix-like systems + try: + import pwd + + return pwd.getpwuid(os.getuid()).pw_name + except (ImportError, KeyError, AttributeError, OSError): + raise + + +def normalize_path(path): + """ + Normalize a path to its full long form. + """ + return Path(path).resolve() # default_generated_ir_path is the path to the cache directory. @@ -53,7 +68,6 @@ def get_current_user(): # Otherwise, it is set to a directory controled by TMPDIR defaulting # to /tmp/${USER}/cutlass_python_cache. if not (default_generated_ir_path := os.getenv("CUTE_DSL_CACHE_DIR", None)): - tmp_dir = Path(os.environ.get("TMPDIR", tempfile.gettempdir())) def get_reusable_temp_dir(name): @@ -67,6 +81,7 @@ if not (default_generated_ir_path := os.getenv("CUTE_DSL_CACHE_DIR", None)): default_generated_ir_path = str(tmp_dir / "cutlass_python_cache") print(f"Could not determine user, using default path. Error: {e}") + @lru_cache(maxsize=1) def get_default_file_dump_root(): """ @@ -187,7 +202,7 @@ def save_ir( :rtype: str """ initial_name = f"{dsl_name.lower()}_{fname}.mlir" - save_path = Path(output_dir if output_dir else tempfile.gettempdir()) + save_path = normalize_path(output_dir if output_dir else tempfile.gettempdir()) save_fname = save_path / initial_name # Random ID to avoid any collisions rnd_id = str(uuid.uuid4()) @@ -206,7 +221,9 @@ def save_ir( module.operation.write_bytecode(f) else: with open(temp_fname, "w") as f: - print(module, file=f) + # Always save with the locations in the MLIR assembly textual + # representation. + print(module.operation.get_asm(enable_debug_info=True), file=f) # os.replace is guaranteed to be atomic on POSIX systems if it succeeds # so filepath cannot see a partial write os.replace(temp_fname, save_fname) @@ -287,3 +304,174 @@ def dump_cache_to_path( log().warning( f"{dsl_name} failed with dumping generated IR cache for {file}: {e}" ) + + +class JitCacheDict: + def __init__(self, max_elems: int | None = None): + """ + This is a dictionary-like object that stores JitCompiledFunction objects + and will garbage collect them when the associated function is garbage + collected. This is done to prevent memory leaks of compiled functions. + + If max_elems is not None, the cache will use an LRU eviction policy to + evict the least recently used item when the cache is full. + + :param max_elems: The maximum number of elements in the cache. + If None, the cache is unlimited. Default is None. + :type max_elems: int | None + """ + self._dict = OrderedDict() if max_elems is not None else dict() + self.max_elems = max_elems + + def get(self, key: Any) -> Any | None: + """ + Try to get the JitCompiledFunction object for the given key. If the key + is not in the cache, None will be returned. This has the same semantics + as `dict.get` in that it will not raise a KeyError if the key is not in + the cache. + + This implementation will use the dictionary as a LRU cache. + First it will get the underlying weak reference to the value (compiled + artifact). Then it will convert the weak reference to an object. + + If the object found by the weak reference is None, it will remove the + key from the cache and return None to indicate that the key is not in + the cache. + + Otherwise, it will return the object and move the key to the end of the + cache to indicate that it is recently used. + + :param key: The key to get the value for. + :type key: Any + :return: The value for the key or None if the key is not in the cache. + :rtype: Any | None + """ + if self.max_elems == 0: + return None + + # returns a object, finalizer pair and we just want the object + value = self._dict.get(key) + if value is not None: + # unpack the value into the object and finalizer + obj, _ = value + + # Move the key to the end of the cache to indicate that it is recently used + if self.max_elems is not None: + self._dict.move_to_end(key, last=True) + return obj + else: + # Key is truly missing, return None + return None + + def set(self, key: Any, value: Any, funcBody: Any = None) -> None: + """ + Set the JitCompiledFunction object for the given key. After calling this + method, the key will be in the cache if maxsize is not 0. + + If the cache is disabled, the value will not be set. + + If the cache is full, the least recently used item will be evicted. + + The implementation will use the dictionary as a LRU cache. First we will + remove any existing finalizer for the key to prevent the old finalizer + from interacting with the cache. Then we will add the new value to the + cache with a new finalizer. Finally, we will move the key to the end of + the cache to indicate that it is recently used. + + :param key: The key to set the value for. + :type key: Any + :param value: The value to set for the given key. + :type value: Any + :param funcBody: The function body that is associated with the value. + :type funcBody: Any + """ + if self.max_elems == 0: + # Cache disabled: ignore writes. + return + + # Detach any existing finalizer for this key so that collection of the + # old value cannot accidentally remove or interfere with the new entry. + old = self._dict.get(key) + if old is not None: + _, old_finalize = old + # Finalizer now will not be called anymore for this key + if old_finalize is not None: + old_finalize.detach() + + def _remove_entry(k: Any, self_ref=weakref.ref(self)) -> None: + # Called from GC/finalizer; be defensive and avoid raising. + self_obj = self_ref() + if self_obj is not None: + # Provide None to avoid pop throwing a key error + self_obj.delete(k) + + assert value != funcBody, ( + "Value and funcBody cannot be the same object to avoid circular references" + ) + self._dict[key] = ( + value, + None + if funcBody is None + else weakref.finalize(funcBody, _remove_entry, key), + ) + # Move the key to the end of the cache to indicate that it is recently used + if self.max_elems is not None: + self._dict.move_to_end(key, last=True) + # If the cache is full, the least recently used item will be evicted + while len(self._dict) > self.max_elems: + # pop from the front + evicted_key, evicted_value = self._dict.popitem(last=False) + # Finalizer now will not be called anymore for this key + _, finalize = evicted_value + if finalize is not None: + finalize.detach() + + def __contains__(self, key: str) -> bool: + """ + Check if the given key is in the cache. + + :param key: The key to check if it is in the cache. + :type key: Any + :return: True if the key is in the cache, False otherwise. + :rtype: bool + """ + return key in self._dict + + def __len__(self) -> int: + """ + Get the number of items in the cache. + + :return: The number of items in the cache. + :rtype: int + """ + return len(self._dict) + + def delete(self, key: Any) -> None: + """ + Try to delete the value for the given key. + + If the key is not in the cache, this will do nothing. + + This will detach the finalizer for the key to prevent it from + being called anymore in the case that the value is garbage collected. + + :param key: The key to delete the value for. + :type key: Any + """ + try: + _, finalize = self._dict.pop(key) + if finalize is not None: + finalize.detach() + except KeyError: + # Key is truly missing, do nothing + pass + + def clear(self) -> None: + """ + Clear the cache. + """ + for key, value in self._dict.items(): + _, finalize = value + if finalize is not None: + finalize.detach() + self._dict.clear() diff --git a/python/CuTeDSL/cutlass/base_dsl/common.py b/python/CuTeDSL/cutlass/base_dsl/common.py index 9d33c403..202288ad 100644 --- a/python/CuTeDSL/cutlass/base_dsl/common.py +++ b/python/CuTeDSL/cutlass/base_dsl/common.py @@ -201,8 +201,10 @@ def _get_friendly_cuda_error_message(error_code, error_name): ), } - message = f"{error_name} (error code: {error_code}) \n" \ - f"{additional_info.get(error_name, '')} \n\n{Colors.RESET}" + message = ( + f"{error_name} (error code: {error_code}) \n" + f"{additional_info.get(error_name, '')} \n\n{Colors.RESET}" + ) # Add debug information debug_info = f"\n- {Colors.BOLD}Error name: {error_name}\n" @@ -321,3 +323,31 @@ This error typically occurs when: " • CUTLASS DSL requirements: nvidia-cutlass-dsl documentation", ], ) + + +class DSLCudaVersion: + """ + Class to represent the CUDA version used to build the DSL. + """ + + def __init__(self, version: str): + self.version_tuple = tuple(int(part) for part in version.split(".")) + + def __str__(self): + return f"{self.major}.{self.minor}" + + def __eq__(self, other): + if isinstance(other, DSLCudaVersion): + return self.version_tuple == other.version_tuple + elif isinstance(other, str): + return self == DSLCudaVersion(other) + else: + return False + + @property + def major(self): + return self.version_tuple[0] + + @property + def minor(self): + return self.version_tuple[1] diff --git a/python/CuTeDSL/cutlass/base_dsl/compiler.py b/python/CuTeDSL/cutlass/base_dsl/compiler.py index 560359e2..339f695e 100644 --- a/python/CuTeDSL/cutlass/base_dsl/compiler.py +++ b/python/CuTeDSL/cutlass/base_dsl/compiler.py @@ -332,6 +332,7 @@ class KeepPTX(BooleanBasedFileDumpOption): option_name = "dump-ptx-path" + class LinkLibraries(StringCompileOption): option_name = "link-libraries" @@ -427,9 +428,7 @@ class CompileOptions: else self.options[DumpDir].value ) if self.options[KeepPTX].value: - self.options[KeepPTX].dump_path = os.path.join( - dump_dir, f"{function_name}" - ) + self.options[KeepPTX].dump_path = os.path.join(dump_dir, f"{function_name}") self.options[KeepPTX].full_ptx_path = os.path.join( dump_dir, f"{function_name}.{arch}.ptx" ) @@ -440,7 +439,6 @@ class CompileOptions: self.options[KeepCUBIN].full_cubin_path = os.path.join( dump_dir, f"{function_name}.{arch}.cubin" ) - @property def generate_line_info(self) -> bool: return self.options[GenerateLineInfo].value @@ -501,6 +499,7 @@ class CompileOptions: def _parse_compile_options_from_str(options: str) -> CompileOptions: """ Parse the compile options from a string. + Deprecated and will be removed in the future. """ def _get_compile_option_from_str(option_str: str): @@ -622,6 +621,10 @@ class CompileCallable: func, ) + func_name_prefix = getattr(func, "_name_prefix", None) + if func_name_prefix: + kwargs["_name_prefix"] = func_name_prefix + # If it's a wrapped function created by decorators, get the original function while hasattr(func, "__wrapped__"): func = func.__wrapped__ diff --git a/python/CuTeDSL/cutlass/base_dsl/dsl.py b/python/CuTeDSL/cutlass/base_dsl/dsl.py index 3d381883..118e406b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/dsl.py +++ b/python/CuTeDSL/cutlass/base_dsl/dsl.py @@ -62,7 +62,6 @@ from .arch import Arch # ============================================================================= from .._mlir import ir -from .._mlir.extras import types as T from .._mlir.dialects import func # ============================================================================= @@ -71,78 +70,6 @@ from .._mlir.dialects import func MLIR_DYNAMIC = -9223372036854775808 -# ============================================================================= -# Codegen Utils -# ============================================================================= - - -def _numpy_type_to_mlir_type(dtype): - if dtype == np.float64: - return T.f64() - if dtype == np.float16: - return T.f16() - if dtype == np.float32: - return T.f32() - if dtype == np.int64: - return T.i64() - if dtype == np.int32: - return T.i32() - if dtype == np.int16: - return T.i16() - if dtype == np.int8: - return T.i8() - if dtype == np.uint64: - return T.ui64() - if dtype == np.uint32: - return T.ui32() - if dtype == np.uint16: - return T.ui16() - if dtype == np.uint8: - return T.ui8() - if dtype == np.bool_: - return T.bool() - if dtype == f8E5M2: - return T.f8E5M2() - if dtype == f8E4M3FN: - return T.f8E4M3FN() - if dtype == f8E8M0FNU: - return T.f8E8M0FNU() - if dtype == f6E3M2FN: - return T.f6E3M2FN() - if dtype == f6E2M3FN: - return T.f6E2M3FN() - if dtype == f4E2M1FN: - return T.f4E2M1FN() - assert False, f"Unknown type {type}" - - -def _mlir_type_to_numpy_type(type): - if type == T.f64(): - return np.float64 - if type == T.f16(): - return np.float16 - if type == T.f32(): - return np.float32 - if type == T.i64(): - return np.int64 - if type == T.i32(): - return np.int32 - if type == T.i16(): - return np.int16 - if type == T.i8(): - return np.int8 - if type == T.ui64(): - return np.uint64 - if type == T.ui32(): - return np.uint32 - if type == T.ui16(): - return np.uint16 - if type == T.ui8(): - return np.uint8 - if type == T.bool(): - return np.bool_ - assert False, f"Unknown type {type}" - # ============================================================================= # Main DSL Class @@ -192,6 +119,38 @@ def extract_mlir_values(obj): return res +def extract_mlir_attributes(obj): + """ + Given the `obj`, recursively go through it to extract all contained IR attributes as list of MLIR attributes. + This is used for generating kernel function argument attributes. + """ + res = [] + if hasattr(obj, "__extract_mlir_attributes__"): + res = obj.__extract_mlir_attributes__() + elif isinstance(obj, (tuple, list)): + res = sum((extract_mlir_attributes(x) for x in obj), []) + elif isinstance(obj, SimpleNamespace): + res = [] + for k, v in obj.__dict__.items(): + res.extend(extract_mlir_attributes(v)) + # Can't call is_dynamic_expression as _is_dynamic_expression depends on extract_mlir_values + elif isinstance(obj, set): + raise DSLRuntimeError( + "Sets are not supported in extract_mlir_values to ensure order preservation", + context="The DSL attempted to generate JIT function argument(s) for an argument of type set but failed.", + suggestion="Consider using a list or tuple instead", + ) + elif isinstance(obj, ir.Value): + res = [ir.DictAttr.get({})] + elif isinstance(obj, ir.BlockArgumentList): + res = [ir.DictAttr.get({})] * len(obj) + else: + # Unlike extract_mlir_values we expand in the default case that we do not have an __extract_mlir_attributes__ + res = [ir.DictAttr.get({})] * len(get_mlir_types(obj)) + + return res + + def new_from_mlir_values(obj, values): """ Create a new python object by populating containing MLIR values with list of new values @@ -351,7 +310,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): self.envar = self._env_class(self.name) self.enable_preprocessor = preprocess # This cache uses hash of original ir and env as key, allows dump/load to/from file. Enabled by default - self.jit_cache = dict() + self.jit_cache = JitCacheDict(max_elems=self.envar.jit_cache_max_elems) self.host_jit_decorator_name = f"@{BaseDSL.jit.__name__}" self.device_jit_decorator_name = f"@{BaseDSL.kernel.__name__}" @@ -380,6 +339,11 @@ class BaseDSL(metaclass=DSLSingletonMeta): log().info(f"Initializing {name} DSL") log().debug(f"Logger initialized for {self.name}") + if self.envar.jit_time_profiling: + self.profiler = timer(enable=True) + self.cache_hits = 0 + self.cache_misses = 0 + # Hook excepthook if self.envar.filter_stacktrace: origin_excepthook = sys.excepthook @@ -727,7 +691,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): """ jit_arg_type, jit_arg_attr, jit_exec_arg = [], [], [] - default_attr = ir.DictAttr.get({}) if is_argument_constexpr(arg, arg_spec, arg_name, arg_index, func): jit_exec_arg = jit_arg_type = jit_arg_attr = None @@ -799,10 +762,12 @@ class BaseDSL(metaclass=DSLSingletonMeta): else: jit_exec_arg.extend(get_c_pointers(arg)) jit_arg_type.extend(get_mlir_types(arg)) + jit_arg_attr.extend([default_attr] * len(get_mlir_types(arg))) else: dyn_vals = extract_mlir_values(arg) jit_exec_arg.extend(dyn_vals) jit_arg_type.extend([v.type for v in dyn_vals]) + jit_arg_attr.extend(extract_mlir_attributes(arg)) if not jit_arg_type or not jit_exec_arg: # when it is compile only, we don't have to prepare the executable arguments. @@ -829,8 +794,6 @@ class BaseDSL(metaclass=DSLSingletonMeta): "enable dynamic value conversion at runtime.", ) - jit_arg_attr.extend([default_attr] * len(jit_arg_type)) - if jit_arg_type is not None: jit_exec_args.extend(jit_exec_arg) jit_arg_types.extend(jit_arg_type) @@ -871,16 +834,18 @@ class BaseDSL(metaclass=DSLSingletonMeta): @dataclass class LaunchConfig: cluster: list = None + fallback_cluster: list = None grid: list = field(default_factory=lambda: [1, 1, 1]) block: list = field(default_factory=lambda: [1, 1, 1]) max_number_threads: list = field(default_factory=lambda: [0, 0, 0]) smem: int = None async_deps: list = field(default_factory=list) has_cluster: bool = False + has_fallback_cluster: bool = False min_blocks_per_mp: int = 0 use_pdl: bool = False auto_smem: bool = False - + cooperative: bool = False @staticmethod def _check_and_canonicalize_dim(dim, name): if not isinstance(dim, (list, tuple)): @@ -890,11 +855,11 @@ class BaseDSL(metaclass=DSLSingletonMeta): raise DSLRuntimeError( f"Expected {name} dimension to be less than or equal to 3, but got {len(dim)}" ) - - if any(not isinstance(e, (Integer, int)) for e in dim): - raise DSLRuntimeError( - f"Expected integer for {name} dimension, but got {type(e)}" - ) + for idx, e in enumerate(dim): + if not isinstance(e, (Integer, int)): + raise DSLRuntimeError( + f"Expected integer for {name} dimension at index {idx}, but got {type(e)}" + ) # Pad with 1s to 3-dim vector for grid or block dimensions return list(dim) + [1] * (3 - len(dim)) @@ -913,6 +878,12 @@ class BaseDSL(metaclass=DSLSingletonMeta): elif len(self.cluster) != 3: raise DSLRuntimeError(f"Expect 3d cluster!") + self.has_fallback_cluster = self.fallback_cluster is not None + if self.fallback_cluster is None: + self.fallback_cluster = [None, None, None] + elif len(self.fallback_cluster) != 3: + raise DSLRuntimeError(f"Expect 3d fallback_cluster!") + def has_max_number_threads(self): """Check if max_number_threads is given by user""" return all( @@ -1027,13 +998,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): pass def preprocess_pipeline(self, pipeline, arch) -> str: - if self.envar.cuda_toolkit is None: - self.print_warning( - "CUDA_TOOLKIT_PATH environment variable is not set. Cannot set toolkitPath." - ) - options = { - "toolkitPath": self.envar.cuda_toolkit if self.envar.cuda_toolkit else None, self.pass_sm_arch_name: arch, } @@ -1118,7 +1083,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): try: module.operation.verify() except Exception as e: - raise DSLRuntimeError(f"🧊🧊🧊 ICE IR Verification Failed 🧊🧊🧊", cause=e) + raise DSLRuntimeError("🧊🧊🧊 ICE IR Verification Failed 🧊🧊🧊", cause=e) return module @@ -1158,9 +1123,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): self._build_gpu_module(gpu_module_attrs, loc=loc) ret_types = self.get_return_types() - fop = func.FuncOp( - function_name, (func_types, ret_types), loc=loc - ) + fop = func.FuncOp(function_name, (func_types, ret_types), loc=loc) fop.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get() log().debug("Generated Function OP [%s]", fop) # Attach per-argument source locations if supported by the FuncOp binding. @@ -1189,8 +1152,10 @@ class BaseDSL(metaclass=DSLSingletonMeta): return module, result # Build IR module - profiler = timer(enable=self.envar.jit_time_profiling) - module, result = profiler(build_ir_module)() + if self.envar.jit_time_profiling: + module, result = self.profiler(build_ir_module)() + else: + module, result = build_ir_module() module_hash = self.get_module_hash(module, function_name) module = self.build_module(module, function_name) @@ -1205,6 +1170,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): pipeline, args_spec, no_cache, + no_jit_engine, func_type=JitCompiledFunction, *, full_args=None, @@ -1212,6 +1178,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): dynamic_args=None, dynamic_kwargs=None, original_function_name=None, + funcBody=None, ): # If `gpu-arch` is set by compile_options, use it. Otherwise, use the arch from the environment variable. compile_gpu_arch = ( @@ -1221,13 +1188,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) # If no gpu kernels or compile_gpu_arch is same as the arch from the environment variable, generate a JIT engine. Otherwise, only do the compilation. gen_jit_engine = self.num_kernels == 0 or compile_gpu_arch == self.envar.arch + if no_jit_engine: + gen_jit_engine = False # Preprocess the pipeline. pipeline = self.preprocess_pipeline( self._get_pipeline(pipeline), compile_gpu_arch ) - log().debug(f"Using pipeline = {pipeline}") shared_libs = self.get_shared_libs() - profiler = timer(enable=self.envar.jit_time_profiling) # try load the file cache load_from_file_cache = False if not no_cache: @@ -1236,13 +1203,17 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) if fn is not None: load_from_file_cache = True - self.jit_cache[module_hash] = fn + self.jit_cache.set(module_hash, fn, funcBody=funcBody) - if ( - no_cache - or module_hash not in self.jit_cache - or self.jit_cache[module_hash].ir_module is None - ): + cached_jit_func = None if no_cache else self.jit_cache.get(module_hash) + + if no_cache or cached_jit_func is None or cached_jit_func.ir_module is None: + if self.envar.jit_time_profiling: + self.cache_misses += 1 + log().info( + "Jit cache hit rate=[%f%%]", + self.cache_hits / (self.cache_hits + self.cache_misses) * 100, + ) log().info( "JIT cache miss function=[%s] module_hash=[%s]", function_name, @@ -1250,11 +1221,22 @@ class BaseDSL(metaclass=DSLSingletonMeta): ) # Compile and JIT MLIR module if gen_jit_engine: - engine = profiler(self.compile_and_jit)( - module, pipeline, shared_libs, function_name=function_name - ) + if self.envar.jit_time_profiling: + engine = self.profiler(self.compile_and_jit)( + module, pipeline, shared_libs, function_name=function_name + ) + else: + engine = self.compile_and_jit( + module, pipeline, shared_libs, function_name=function_name + ) else: - profiler(self.compiler_provider.compile)(module, pipeline) + if self.envar.jit_time_profiling: + self.profiler(self.compiler_provider.compile)( + module, + pipeline, + ) + else: + self.compiler_provider.compile(module, pipeline) engine = None else: log().info( @@ -1262,13 +1244,22 @@ class BaseDSL(metaclass=DSLSingletonMeta): function_name, module_hash, ) - module = self.jit_cache[module_hash].ir_module + if self.envar.jit_time_profiling: + self.cache_hits += 1 + log().info( + "JIT cache hit rate=[%f%%]", + self.cache_hits / (self.cache_hits + self.cache_misses) * 100, + ) + module = cached_jit_func.ir_module engine = ( self.compiler_provider.jit(module, shared_libs=shared_libs) if gen_jit_engine else None ) - capi_func = profiler(engine.lookup)(function_name) if engine else None + if self.envar.jit_time_profiling: + capi_func = self.profiler(engine.lookup)(function_name) if engine else None + else: + capi_func = engine.lookup(function_name) if engine else None fn = func_type( module, @@ -1283,14 +1274,14 @@ class BaseDSL(metaclass=DSLSingletonMeta): CUBIN=self.compile_options.full_cubin_path, MLIR=(self.dump_mlir_path if self.envar.keep_ir else None), ), + # set dynamic arguments if the jit_function is a JitCompiledFunction for AOT generation. + dynamic_args=dynamic_args, + dynamic_kwargs=dynamic_kwargs, ) - # set dynamic arguments if the jit_function is a JitCompiledFunction for AOT generation. - fn.set_dynamic_args(dynamic_args, dynamic_kwargs) - if not no_cache: # module stored in cache is compiled. - self.jit_cache[module_hash] = fn + self.jit_cache.set(module_hash, fn, funcBody=funcBody) # write through the file cache if enabled. if not self.envar.disable_file_caching and not load_from_file_cache: dump_cache_to_path( @@ -1330,18 +1321,10 @@ class BaseDSL(metaclass=DSLSingletonMeta): i, funcBody, ): - try: - dynamic_args.append(weakref.proxy(arg)) - except TypeError: - # If arg cannot be weakly referenced (e.g., int, float) - dynamic_args.append(arg) + dynamic_args.append(arg) for i, (k, v) in enumerate(kwargs.items()): if not is_argument_constexpr(v, args_spec.kwonlyargs[i], k, i, funcBody): - try: - dynamic_kwargs[k] = weakref.proxy(v) - except TypeError: - # If v cannot be weakly referenced (e.g., int, float) - dynamic_kwargs[k] = v + dynamic_kwargs[k] = v return dynamic_args, dynamic_kwargs def generate_mlir( @@ -1354,6 +1337,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): args_spec, pipeline, no_cache, + no_jit_engine, compile_only, location=None, ): @@ -1387,10 +1371,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): if self.envar.dryrun: return result + # Get a single reference to the cache since garbage collection + cached_jit_func = None if no_cache else self.jit_cache.get(module_hash) + if ( no_cache - or module_hash not in self.jit_cache - or self.jit_cache[module_hash].capi_func is None + or cached_jit_func is None + or cached_jit_func.capi_func is None ): # no cache or cache miss, do ir generation/compilation/jit engine jit_function = self.compile_and_cache( @@ -1400,11 +1387,13 @@ class BaseDSL(metaclass=DSLSingletonMeta): pipeline, args_spec, no_cache, + no_jit_engine, full_args=args, full_kwargs=kwargs, dynamic_args=dynamic_args, dynamic_kwargs=dynamic_kwargs, original_function_name=original_function_name, + funcBody=funcBody, ) else: # cache hit @@ -1413,7 +1402,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): function_name, module_hash, ) - jit_function = self.jit_cache[module_hash] + jit_function = cached_jit_func finally: self.post_compilation_cleanup() @@ -1537,9 +1526,14 @@ class BaseDSL(metaclass=DSLSingletonMeta): # Disable cache no_cache = kwargs.pop("no_cache", False) + # Disable JIT execution engine + no_jit_engine = kwargs.pop("no_jit_engine", False) + # Always compile(disable cache) and return the result jit_executor compile_only = kwargs.pop("compile_only", False) + func_name_prefix = kwargs.pop("_name_prefix", None) + if not no_cache and ( self.envar.keep_ptx or self.envar.keep_cubin @@ -1560,9 +1554,11 @@ class BaseDSL(metaclass=DSLSingletonMeta): canonicalized_args, canonicalized_kwargs = self._canonicalize_args( sig, *args, **kwargs ) - # Simple name mangling function_name = self.mangle_name(function_name, canonicalized_args, args_spec) + if func_name_prefix: + function_name = f"{func_name_prefix}_{function_name}" + self.compile_options.apply_envar_settings(self.envar, function_name) if not self.compile_options.generate_line_info: self.decorator_location = None @@ -1578,6 +1574,7 @@ class BaseDSL(metaclass=DSLSingletonMeta): args_spec, pipeline, no_cache, + no_jit_engine, compile_only, location=self.decorator_location, ) @@ -1778,6 +1775,8 @@ class BaseDSL(metaclass=DSLSingletonMeta): # The mangled name of Python function is part of the name to # improve readability. kernel_name = f"kernel_{self.mangle_name(kernel_name, args, args_spec)}_{self.num_kernels}" + if hasattr(self, "_name_prefix") and self._name_prefix: + kernel_name = f"{self._name_prefix}_{kernel_name}" self.num_kernels += 1 # Step 0. Preprocess the arguments diff --git a/python/CuTeDSL/cutlass/base_dsl/env_manager.py b/python/CuTeDSL/cutlass/base_dsl/env_manager.py index bd08b158..f135258b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/env_manager.py +++ b/python/CuTeDSL/cutlass/base_dsl/env_manager.py @@ -73,6 +73,27 @@ def get_int_env_var(var_name, default_value=0): return int(value) if value and value.isdigit() else default_value +@lru_cache(maxsize=None) +def get_int_or_none_env_var(var_name, default_value=None): + """ + Get the value of an integer or None union environment variable. + If the value is not a valid integer, the default value 0 is returned. + Note that the value is cached after the first call. + """ + raw = get_str_env_var(var_name) + if raw is None: + return default_value + + value = raw.strip().lower() + if value == "none": + return None + + try: + return int(value) + except ValueError: + return default_value + + @lru_cache(maxsize=None) def has_env_var(var_name): """ @@ -235,17 +256,19 @@ def get_prefix_dsl_libs(prefix: str): return prefix_libs_existing def get_libs_cand(start): - target_cuda_dialect_libs = { - "cuda_dialect_runtime", + target_dsl_runtime_libs = { + "cute_dsl_runtime", } lib_folder_guesses = [ "lib", ] for target_libs in [ - target_cuda_dialect_libs, + target_dsl_runtime_libs, ]: - libs_cand = find_libs_in_ancestors(start, target_libs, lib_folder_guesses) + libs_cand = find_libs_in_ancestors( + start, target_libs, lib_folder_guesses + ) if libs_cand: dsl_libs = ":".join(libs_cand) return dsl_libs @@ -310,6 +333,8 @@ class EnvironmentVarManager(LogEnvironmentManager): - [DSL_NAME]_WARNINGS_IGNORE: Ignore warnings (default: False) - [DSL_NAME]_ENABLE_OPTIMIZATION_WARNINGS: Enable warnings of optimization warnings (default: False) - [DSL_NAME]_JIT_TIME_PROFILING: Whether or not to profile the IR generation/compilation/execution time (default: False) + - [DSL_NAME]_JIT_CACHE_MAX_ELEMS: Maximum number of JIT compiled functions to cache in memory (default: None). If None, the cache is unbounded. + - [DSL_NAME]_NO_CACHE: Disable JIT cache (default: False) - [DSL_NAME]_DISABLE_FILE_CACHING: Disable file caching (default: False) - [DSL_NAME]_LIBS: Path to dependent shared libraries (default: None) - [DSL_NAME]_ENABLE_TVM_FFI: Enable TVM-FFI or not (default: False) @@ -325,12 +350,17 @@ class EnvironmentVarManager(LogEnvironmentManager): self.print_ir = get_bool_env_var(f"{prefix}_PRINT_IR", False) self.filter_stacktrace = get_bool_env_var(f"{prefix}_FILTER_STACKTRACE", True) self.lineinfo = get_bool_env_var(f"{prefix}_LINEINFO", False) + self.no_cache = get_bool_env_var(f"{prefix}_NO_CACHE", False) + self.jit_cache_max_elems = get_int_or_none_env_var( + f"{prefix}_JIT_CACHE_MAX_ELEMS", None + ) + if self.no_cache: + self.jit_cache_max_elems = 0 self.dump_dir = get_str_env_var( f"{prefix}_DUMP_DIR", get_default_file_dump_root() ) self.keep_ptx = get_bool_env_var(f"{prefix}_KEEP_PTX", False) self.keep_cubin = get_bool_env_var(f"{prefix}_KEEP_CUBIN", False) - # File options self.keep_ir = get_bool_env_var(f"{prefix}_KEEP_IR", False) self.cache_dir = get_str_env_var(f"{prefix}_CACHE_DIR", None) diff --git a/python/CuTeDSL/cutlass/base_dsl/export/__init__.py b/python/CuTeDSL/cutlass/base_dsl/export/__init__.py index 5d1c59a3..66eaf403 100644 --- a/python/CuTeDSL/cutlass/base_dsl/export/__init__.py +++ b/python/CuTeDSL/cutlass/base_dsl/export/__init__.py @@ -9,12 +9,23 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from .c_header_generator import CHeaderGenerator -from .export import get_export_module, dump_to_object, export_to_c +from .c_header_generator import CHeaderGenerator, CHeaderArguments +from .export import ( + get_export_module, + encode_metadata_into_ir_module, + decode_metadata_from_execution_engine, +) + +from .export import ArgsSpecProcessor +from .external_binary_module import ExternalBinaryModule, LoadProvider __all__ = [ "CHeaderGenerator", + "CHeaderArguments", "get_export_module", - "dump_to_object", - "export_to_c", + "encode_metadata_into_ir_module", + "decode_metadata_from_execution_engine", + "ArgsSpecProcessor", + "ExternalBinaryModule", + "LoadProvider", ] diff --git a/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py b/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py index 3284280f..f99a3e36 100644 --- a/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py +++ b/python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py @@ -41,6 +41,36 @@ import cuda.bindings.driver as cuda cubin_suffix = "cubin" +class CHeaderArguments: + """This class is used to store the arguments generation of the wrapper function. + The arguments are generated when the JitCompiledFunction is created and used to avoid + the long-term reference of the arguments by the JitCompiledFunction. + """ + + def __init__( + self, + dummy_prefix_name: str, + arguments: List[str], + packed_args: List[str], + declarations: str, + error_msg: str = None, + ): + self.dummy_prefix_name = dummy_prefix_name + self.arguments = arguments + self.packed_args = packed_args + self.declarations = declarations + self.error_msg = error_msg + + def __bool__(self): + return self.error_msg is None + + def __str__(self): + return self.error_msg + + def __repr__(self): + return self.error_msg + + class CHeaderGenerator: """This class provides a Export C Header Generator for c/cpp AOT support.""" @@ -120,11 +150,11 @@ class CHeaderGenerator: return check_cuda - def _generate_kernel_metadata( + def _generate_kernel_module( self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str ): """ - Generate the kernel metadata for the compiled function. + Generate the kernel module for the compiled function. """ function_declarations = [] function_loads = [] @@ -132,31 +162,31 @@ class CHeaderGenerator: function_non_portable_cluster_size_allowed = [] for sym, attrs in kernel_info.items(): function_declaration = f"CUfunction {sym};" - function_load = f'{dsl_name}_CUDA_ERROR_CHECK(cuModuleGetFunction(&metadata->{sym}, metadata->module, "{sym}"));' + function_load = f'{dsl_name}_CUDA_ERROR_CHECK(cuModuleGetFunction(&module->{sym}, module->module, "{sym}"));' function_declarations.append(function_declaration) function_loads.append(function_load) for attr in attrs: function_set_attributes.append( - f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->{sym}, {attr.name}, {attr.value}));" + f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(module->{sym}, {attr.name}, {attr.value}));" ) function_non_portable_cluster_size_allowed.append( - f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->{sym}, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));" + f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(module->{sym}, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));" ) function_declarations_str = "\n ".join(function_declarations) - kernel_metadata_struct = f""" + kernel_module_struct = f""" typedef struct {{ CUmodule module; {function_declarations_str} -}} {symbol_prefix}_Kernel_Metadata_t; +}} {symbol_prefix}_Kernel_Module_t; """ function_loads_str = "\n ".join(function_loads) function_set_attributes_str = "\n ".join(function_set_attributes) function_non_portable_cluster_size_allowed_str = "\n ".join( function_non_portable_cluster_size_allowed ) - kernel_metadata_load = f""" -static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{ - {dsl_name}_CUDA_ERROR_CHECK(cuModuleLoadData(&metadata->module, {symbol_prefix}_{cubin_suffix})); + kernel_module_load = f""" +static inline void {symbol_prefix}_Kernel_Module_Load({symbol_prefix}_Kernel_Module_t *module) {{ + {dsl_name}_CUDA_ERROR_CHECK(cuModuleLoadData(&module->module, {symbol_prefix}_{cubin_suffix})); {function_loads_str} {function_set_attributes_str} int driver_version; @@ -166,13 +196,13 @@ static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_M }} }} """ - kernel_metadata_unload = f""" -static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{ - {dsl_name}_CUDA_ERROR_CHECK(cuModuleUnload(metadata->module)); + kernel_module_unload = f""" +static inline void {symbol_prefix}_Kernel_Module_Unload({symbol_prefix}_Kernel_Module_t *module) {{ + {dsl_name}_CUDA_ERROR_CHECK(cuModuleUnload(module->module)); }} """ - return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload + return kernel_module_struct + kernel_module_load + kernel_module_unload def _generate_arguments( self, @@ -213,29 +243,42 @@ static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel def _generate_wrapper_function( self, + dsl_name: str, symbol_prefix: str, args_spec: ExecutionArgs, function_name: str, kernel_info: Dict[str, List], - dynamic_args: list, - dynamic_kwargs: dict, + c_header_arguments: CHeaderArguments, ): """ Generate the wrapper function for the compiled function which is provided to users as the entry point. It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode. """ # 1. Get the name of the function wrapper - wrapper_function_name = f"{symbol_prefix}_wrapper" + wrapper_function_name = f"{dsl_name.lower()}_{symbol_prefix}_wrapper" capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}" # 2. Generate the signature of the wrapper function - arguments, packed_args, declarations = self._generate_arguments( - symbol_prefix, args_spec, dynamic_args, dynamic_kwargs - ) + if c_header_arguments.error_msg is not None: + raise DSLRuntimeError( + f"Error generating c header arguments: {c_header_arguments.error_msg}" + ) + arguments = [ + arg.replace(c_header_arguments.dummy_prefix_name, symbol_prefix) + for arg in c_header_arguments.arguments + ] + packed_args = [ + arg.replace(c_header_arguments.dummy_prefix_name, symbol_prefix) + for arg in c_header_arguments.packed_args + ] + declarations = [ + declaration.replace(c_header_arguments.dummy_prefix_name, symbol_prefix) + for declaration in c_header_arguments.declarations + ] declarations = "\n".join(declarations) # 3. Generate the wrapper function kernel_symbols = tuple(kernel_info.keys()) - kernel_symbols_str = ", ".join([f"&metadata->{sym}" for sym in kernel_symbols]) + kernel_symbols_str = ", ".join([f"&module->{sym}" for sym in kernel_symbols]) function = ( declarations + f""" @@ -244,7 +287,7 @@ extern "C" #endif void {capi_function_name}(void **args, int32_t num_args); -static inline void {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{ +static inline void {wrapper_function_name}({symbol_prefix}_Kernel_Module_t *module, {", ".join(arguments)}) {{ void *args[{len(packed_args) + len(kernel_symbols)}] = {{ {", ".join(packed_args)}, {kernel_symbols_str} }}; @@ -258,7 +301,6 @@ static inline void {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *me """ Generate the binary of the compiled function. """ - varname = symbol_prefix + "_" + cubin_suffix binary = f""" extern const unsigned char {varname}[]; @@ -272,29 +314,28 @@ extern const unsigned char {varname}[]; args_spec: ExecutionArgs, function_name: str, kernel_info: Dict[str, List], - dynamic_args: list, - dynamic_kwargs: dict, + c_header_arguments: CHeaderArguments, dsl_name: str, ) -> str: if len(kernel_info) > 0: check_cuda = self._generate_check_cuda(dsl_name) - kernel_metadata = self._generate_kernel_metadata( + kernel_module = self._generate_kernel_module( symbol_prefix, kernel_info, dsl_name ) binary = self._generate_binary_declaration(symbol_prefix) else: check_cuda = "" - kernel_metadata = "" + kernel_module = "" binary = "" function = self._generate_wrapper_function( + dsl_name, symbol_prefix, args_spec, function_name, kernel_info, - dynamic_args, - dynamic_kwargs, + c_header_arguments, ) - header = self.includes + check_cuda + binary + kernel_metadata + function + header = self.includes + check_cuda + binary + kernel_module + function return header @@ -333,21 +374,21 @@ extern const unsigned char {varname}[]; # typedef struct { # CUmodule module; # CUfunction kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0; -# } gemm_Kernel_Metadata_t; +# } gemm_Kernel_Module_t; -# static inline void gemm_Kernel_Metadata_Load(gemm_Kernel_Metadata_t *metadata) { -# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleLoadData(&metadata->module, gemm_cubin)); -# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleGetFunction(&metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, metadata->module, "kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0")); +# static inline void gemm_Kernel_Module_Load(gemm_Kernel_Module_t *module) { +# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleLoadData(&module->module, gemm_cubin)); +# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleGetFunction(&module->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, module->module, "kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0")); # int driver_version; # CUTE_DSL_CUDA_ERROR_CHECK(cuDriverGetVersion(&driver_version)); # if (driver_version >= 11080) { -# CUTE_DSL_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)); +# CUTE_DSL_CUDA_ERROR_CHECK(cuFuncSetAttribute(module->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)); # } # } -# static inline void gemm_Kernel_Metadata_Unload(gemm_Kernel_Metadata_t *metadata) { -# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleUnload(metadata->module)); +# static inline void gemm_Kernel_Module_Unload(gemm_Kernel_Module_t *module) { +# CUTE_DSL_CUDA_ERROR_CHECK(cuModuleUnload(module->module)); # } # typedef struct { @@ -375,9 +416,9 @@ extern const unsigned char {varname}[]; # #endif # void _mlir_gemm__mlir_ciface_cutlass_host_func_Tensorgenericoi641_Tensorgenericoi641_Tensorgenericoi641(void **args, int32_t num_args); -# static inline void gemm_wrapper(gemm_Kernel_Metadata_t *metadata, gemm_Tensor_a_t *a, gemm_Tensor_b_t *b, gemm_Tensor_c_t *c) { +# static inline void cute_dsl_gemm_wrapper(gemm_Kernel_Module_t *module, gemm_Tensor_a_t *a, gemm_Tensor_b_t *b, gemm_Tensor_c_t *c) { # void *args[4] = { -# a, b, c, &metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0 +# a, b, c, &module->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0 # }; # _mlir_gemm__mlir_ciface_cutlass_host_func_Tensorgenericoi641_Tensorgenericoi641_Tensorgenericoi641(args, 4); # } diff --git a/python/CuTeDSL/cutlass/base_dsl/export/export.py b/python/CuTeDSL/cutlass/base_dsl/export/export.py index 5342ad15..7e91ab00 100644 --- a/python/CuTeDSL/cutlass/base_dsl/export/export.py +++ b/python/CuTeDSL/cutlass/base_dsl/export/export.py @@ -11,23 +11,26 @@ import io import os -import array -import tempfile from ..common import DSLRuntimeError -from ..jit_executor import JitCompiledFunction, get_escaped_cubin_bytes -from ...base_dsl.dsl import BaseDSL -from ...base_dsl.typing import Int32, Int64, Float32, Float64 from ..._mlir import ir from ..._mlir.dialects import llvm -from .c_header_generator import CHeaderGenerator +import json +import base64 +import ctypes +from inspect import FullArgSpec -from typing import Union +args_spec_suffix = "args_spec" +function_name_suffix = "function_name" +kernel_info_suffix = "kernel_info" +version_suffix = "version" +c_string_suffix = "\0" -cubin_suffix = "cubin" -def get_export_module(ir_module: ir.Module, symbol_prefix: str, *, preserve_symbols = None): +def get_export_module( + ir_module: ir.Module, symbol_prefix: str, *, preserve_symbols=None +): """Get the export module which is cloned from the original compiled ir module, and add the prefix to avoid the symbol conflict. @@ -66,7 +69,10 @@ def get_export_module(ir_module: ir.Module, symbol_prefix: str, *, preserve_symb symbol_prefix + "_" + op.attributes["callee"].value ) # Rename addressof references - elif op.name == "llvm.mlir.addressof" and op.attributes["global_name"].value in defined_symbols: + elif ( + op.name == "llvm.mlir.addressof" + and op.attributes["global_name"].value in defined_symbols + ): op.attributes["global_name"] = ir.FlatSymbolRefAttr.get( symbol_prefix + "_" + op.attributes["global_name"].value ) @@ -76,9 +82,9 @@ def get_export_module(ir_module: ir.Module, symbol_prefix: str, *, preserve_symb renamed_ctors = [] for ctor in ctors: if ctor.value in defined_symbols: - renamed_ctors.append(ir.FlatSymbolRefAttr.get( - symbol_prefix + "_" + ctor.value - )) + renamed_ctors.append( + ir.FlatSymbolRefAttr.get(symbol_prefix + "_" + ctor.value) + ) else: renamed_ctors.append(ctor) if renamed_ctors: @@ -89,9 +95,9 @@ def get_export_module(ir_module: ir.Module, symbol_prefix: str, *, preserve_symb renamed_dtors = [] for dtor in dtors: if dtor.value in defined_symbols: - renamed_dtors.append(ir.FlatSymbolRefAttr.get( - symbol_prefix + "_" + dtor.value - )) + renamed_dtors.append( + ir.FlatSymbolRefAttr.get(symbol_prefix + "_" + dtor.value) + ) else: renamed_dtors.append(dtor) if renamed_dtors: @@ -107,135 +113,127 @@ def get_export_module(ir_module: ir.Module, symbol_prefix: str, *, preserve_symb return export_module +class ArgsSpecProcessor: + """The args spec processor. The args_spec may contain the dsl specific types. The base processor + class is used to define an interface for dumping and loading the args_spec.""" -def dump_to_object( - prefix_name: str, - export_module: ir.Module, - jit_function: Union[JitCompiledFunction, "CudaDialectJitCompiledFunction"], - dsl: BaseDSL, - use_gpu_dialect: bool, -) -> bytes: - """Dump the compiled ir function to a bytes object with ELF format. The bytes object contains the host - launch entry function and cubin inside. + def dumps(self, args_spec: FullArgSpec) -> bytes: + raise NotImplementedError("ArgsSpecProcessor does not support dumps") - @param prefix_name: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. - @param export_module: The export module of the function. This is the module that contains the function to be exported. - @param jit_function: The jit-compiled function. To provided other metadata for the object file. - @param dsl: The dsl object. This is the dsl object to get the compiler provider and shared libs. - @return: The bytes object of the function. - """ - if use_gpu_dialect: - cubin_data = None - - def strip_gpu_binary_op(op): - if op.name == "gpu.binary": - s = io.BytesIO() - op.operation.write_bytecode(s) - nonlocal cubin_data - cubin_data = s.getvalue() - cubin_data = cubin_data.split(b'bin = "')[1].split(b'">')[0] - cubin_data = get_escaped_cubin_bytes(cubin_data) - op.erase() - return ir.WalkResult.ADVANCE - return ir.WalkResult.ADVANCE - - # Strip gpu related to avoid the object file generating builtin module load/unload functions - export_module.operation.walk(strip_gpu_binary_op) - - cubin_array = array.array("b", cubin_data) - with ( - export_module.context, - ir.Location.unknown(), - ir.InsertionPoint(export_module.body), - ): - new_binary_global_op = llvm.GlobalOp( - sym_name="_".join([prefix_name, cubin_suffix]), - global_type=ir.Type.parse(f"!llvm.array<{len(cubin_array)} x i8>"), - linkage=ir.Attribute.parse("#llvm.linkage"), - value=ir.DenseIntElementsAttr.get(cubin_array), - constant=True, - ) - if "gpu.container_module" in export_module.operation.attributes: - del export_module.operation.attributes["gpu.container_module"] - # Generate the object file - export_engine = dsl.compiler_provider.jit( - export_module, shared_libs=dsl.get_shared_libs() - ) - # This lookup is necessary to make sure the compilation is done. - entry_func = export_engine.raw_lookup( - "_".join([prefix_name, jit_function.function_name]) - ) - if not entry_func: - raise DSLRuntimeError( - f"Execution engine cannot find the entry function {prefix_name}_{jit_function.function_name}" - ) - try: - with tempfile.NamedTemporaryFile() as tmp_object_file: - export_engine.dump_to_object_file(tmp_object_file.name) - with open(tmp_object_file.name, "rb") as f: - ret = f.read() - return ret - except Exception as e: - raise DSLRuntimeError(f"Error dumping object file: {e}") from e + def loads(self, args_spec_bytes: bytes): + raise NotImplementedError("ArgsSpecProcessor does not support loads") -def export_to_c( - jit_function: Union[JitCompiledFunction, "CudaDialectJitCompiledFunction"], - file_path: str, - file_name: str, - dsl: BaseDSL, - c_header_generator: CHeaderGenerator, - use_gpu_dialect: bool, +def encode_metadata_into_ir_module( + prefix: str, + ir_module: ir.Module, + args_spec: FullArgSpec, + function_name: str, + kernel_info: dict, + args_spec_processor: ArgsSpecProcessor, + object_file_version: str, ): - """Exports the jit-compiled function to a C compatible files(header/library). - This is used for c/cpp AOT support. - The `file_name` will be used as the symbol prefix of the generated functions, it is guaranteed by - the caller that the generated functions are unique. And the function will always overwrite the existing file. + """Encode the executor metadata into the ir module. The metadata includes: + 1. args_spec: The args_spec of the python function. + 2. function_name: The name mangling function_name of the python host function. + 3. kernel_info: The kernel_info of the jit-compiled function including the kernel name and attributes. + 4. version: The version of the object file. - - The c header file is generated with following components: - 1. The host launch entry function. And the structure definitions of the arguments. - 2. The device metadata load/unload functions. - 3. The cubin data array and len. - - The library contains the binary of the underlying host launch entry function. - - @param jit_function: The jit-compiled function from `cute.compile`. - @param file_path: The path to the directory where the header and object files will be saved. - @param file_name: The name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. - @param dsl: The dsl object. This is the dsl object to get the compiler provider and shared libs. - @param c_header_generator: The c header generator. This is the c header generator to generate the c header file. + @param prefix: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. + @param ir_module: The ir module to encode the metadata into. + @param args_spec: The args_spec of the python function. + @param function_name: The name mangling function_name of the python host function. + @param kernel_info: The kernel_info of the jit-compiled function including the kernel name and attributes. + @param args_spec_processor: The args spec processor. The args_spec may contain the dsl specific types. The processor will be used to dump and load the args_spec. + @param object_file_version: The version of the object file. """ - export_module = get_export_module(jit_function.ir_module, file_name) - # Generate the c header file - header_file_content = c_header_generator( - file_name, - export_module, - jit_function.args_spec, - jit_function.function_name, - jit_function.kernel_info, - jit_function.dynamic_args, - jit_function.dynamic_kwargs, - dsl.name, + if not args_spec: + raise DSLRuntimeError( + "args_spec is empty, please set the args_spec for the python jit function." + ) + version = object_file_version + c_string_suffix + + args_spec_bytes = args_spec_processor.dumps(args_spec) + args_spec_str = base64.b64encode(args_spec_bytes).decode("utf-8") + c_string_suffix + packed_function_name = ( + "_mlir_" + prefix + "__mlir_ciface_" + function_name + c_string_suffix ) - try: - with open(os.path.join(file_path, file_name + ".h"), "w") as f: - f.write(header_file_content) - except Exception as e: - raise DSLRuntimeError(f"Error writing header file: {e}") from e + with ir_module.context, ir.Location.unknown(): + with ir.InsertionPoint(ir_module.body): + args_spec_op = llvm.GlobalOp( + sym_name="_".join([prefix, args_spec_suffix]), + global_type=ir.Type.parse(f"!llvm.array<{len(args_spec_str)} x i8>"), + linkage=ir.Attribute.parse("#llvm.linkage"), + value=ir.StringAttr.get(args_spec_str), + ) + function_name_op = llvm.GlobalOp( + sym_name="_".join([prefix, function_name_suffix]), + global_type=ir.Type.parse( + f"!llvm.array<{len(packed_function_name)} x i8>" + ), + linkage=ir.Attribute.parse("#llvm.linkage"), + value=ir.StringAttr.get(packed_function_name), + ) + # pack the kernel_info from a dict to a global op. + kernel_info = json.dumps(kernel_info) + c_string_suffix + kernel_info_op = llvm.GlobalOp( + sym_name="_".join([prefix, kernel_info_suffix]), + global_type=ir.Type.parse(f"!llvm.array<{len(kernel_info)} x i8>"), + linkage=ir.Attribute.parse("#llvm.linkage"), + value=ir.StringAttr.get(kernel_info), + ) + version_op = llvm.GlobalOp( + sym_name="_".join([prefix, version_suffix]), + global_type=ir.Type.parse(f"!llvm.array<{len(version)} x i8>"), + linkage=ir.Attribute.parse("#llvm.linkage"), + value=ir.StringAttr.get(version), + ) - # Generate the object file - object_file_content = dump_to_object( - file_name, - export_module, - jit_function, - dsl, - use_gpu_dialect, + return ir_module + + +def decode_metadata_from_execution_engine( + prefix: str, + execution_engine: "BinaryExecutionEngine", + args_spec_processor: ArgsSpecProcessor, +): + """Decode the executor metadata from the execution engine. The metadata includes: + 1. args_spec: The args_spec of the python function. + 2. function_name: The name mangling function_name of the python host function. + 3. kernel_info: The kernel_info of the jit-compiled function including the kernel name and attributes. + 4. version: The version of the object file. + + @param prefix: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. + @param execution_engine: The binary execution engine. This is the execution engine to load the cuda module. + @param args_spec_processor: The args spec processor. The args_spec may contain the dsl specific types. The processor will be used to dump and load the args_spec. + @return: The args_spec, function_name, and kernel_info. + """ + args_spec_str_p = execution_engine.lookup("_".join([prefix, args_spec_suffix])) + function_name_str_p = execution_engine.lookup( + "_".join([prefix, function_name_suffix]) ) - try: - with open(os.path.join(file_path, file_name + ".o"), "wb") as f: - f.write(object_file_content) - except Exception as e: - raise DSLRuntimeError(f"Error writing object file: {e}") from e - + kernel_info_str_p = execution_engine.lookup("_".join([prefix, kernel_info_suffix])) + version_str_p = execution_engine.lookup("_".join([prefix, version_suffix])) + if args_spec_str_p: + args_spec_str = ctypes.c_char_p(args_spec_str_p).value.decode("utf-8") + else: + args_spec_str = None + # The StringAttr encodes the string as utf-8 format. + if function_name_str_p: + function_name_str = ctypes.c_char_p(function_name_str_p).value.decode("utf-8") + else: + function_name_str = None + if kernel_info_str_p: + kernel_info_str = ctypes.c_char_p(kernel_info_str_p).value.decode("utf-8") + else: + kernel_info_str = None + if version_str_p: + version_str = ctypes.c_char_p(version_str_p).value.decode("utf-8") + else: + version_str = None + args_spec_bytes = base64.b64decode(args_spec_str) + args_spec = args_spec_processor.loads(args_spec_bytes) + function_name = function_name_str + kernel_info = json.loads(kernel_info_str) + return args_spec, function_name, kernel_info, version_str diff --git a/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py b/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py new file mode 100644 index 00000000..c42c6fa2 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/export/external_binary_module.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import io +import os + +import ctypes +from typing import Callable +from inspect import FullArgSpec + +from ..common import DSLRuntimeError +from ...base_dsl.dsl import BaseDSL +from ...base_dsl.typing import Int32, Int64, Float32, Float64 +from .export import decode_metadata_from_execution_engine + + +def _get_ctypes_return_type(args_spec: FullArgSpec): + """Get the ctypes return type from the args_spec.""" + return_type = args_spec.annotations.get("return", None) + if return_type is None: + raise DSLRuntimeError("Return type is not specified for AOT compiled function.") + type_to_ctype = { + Int32: ctypes.c_int32, + Int64: ctypes.c_int64, + Float32: ctypes.c_float, + Float64: ctypes.c_double, + # Add other types if needed + } + + ctype = type_to_ctype.get(return_type) + if ctype is None: + raise DSLRuntimeError(f"Unsupported return type for AOT loading: {return_type}") + + return ctype + + +class LoadProvider: + """The load provider is a class that stores the necessary information to construct a ExternalBinaryModule, + it could be set later by specific DSL.""" + + def __init__( + self, + dsl: "Type[BaseDSL]", + args_spec_processor: "ArgsSpecProcessor", + version_checker: Callable, + execution_engine_constructor: Callable, + jit_function_constructor: Callable, + ): + self.dsl = dsl + self.args_spec_processor = args_spec_processor + self.version_checker = version_checker + self.execution_engine_constructor = execution_engine_constructor + self.jit_function_constructor = jit_function_constructor + + +class ExternalBinaryModule: + """The exported binary module is a wrapper of the previous exported object files. It is used to load a object file + or a library in memory, allow function lookup and return the corresponding `JitCompiledFunction`. + """ + + load_provider: LoadProvider = None + + def __init__(self, file_path: str): + assert self.load_provider is not None, ( + "Load provider is not set for ExternalBinaryModule." + ) + shared_libs = self.load_provider.dsl._get_dsl().get_shared_libs() + object_file_content = bytes() + if file_path.endswith(".so"): + shared_libs.append(file_path) + else: + try: + with open(file_path, "rb") as f: + object_file_content = f.read() + except Exception as e: + raise DSLRuntimeError(f"Failed to read object file {file_path}: {e}") + # Lifetime of the engine is same as the ExternalBinaryModule. + self.engine = self.load_provider.execution_engine_constructor( + object_file_content, shared_libs + ) + + def __getattr__(self, function_prefix: str) -> "JitCompiledFunction": + """Get the jit_function from the `function_prefix`. The `function_prefix` is specified when users dump the object file. When there is no function_prefix found in the module, the function will raise an error.""" + try: + args_spec, function_name, kernel_info, version_str = ( + decode_metadata_from_execution_engine( + function_prefix, self.engine, self.load_provider.args_spec_processor + ) + ) + except Exception as e: + raise DSLRuntimeError( + f"Function prefix {function_prefix} not found in the module.", cause=e + ) + self.load_provider.version_checker(version_str) + capi_func_p = self.engine.lookup(function_name) + if not capi_func_p: + raise DSLRuntimeError( + "Unknown function: " + + "_mlir_" + + function_prefix + + "__mlir_ciface_" + + function_name + ) + return_type = _get_ctypes_return_type(args_spec) + capi_func = ctypes.CFUNCTYPE(return_type, ctypes.c_void_p)(capi_func_p) + jit_function = self.load_provider.jit_function_constructor( + ir_module=None, + engine=self.engine, + capi_func=capi_func, + args_spec=args_spec, + function_name=function_name, + kernel_info=kernel_info, + jit_time_profiling=self.load_provider.dsl._get_dsl().envar.jit_time_profiling, + jit_function_artifacts=None, + prefix=function_prefix, + load_from_binary=True, + ) + return jit_function diff --git a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py index 5835473d..236049a6 100644 --- a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py +++ b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py @@ -13,6 +13,7 @@ This module provides jit executor related classes """ +import array import ctypes import inspect import io @@ -21,10 +22,12 @@ import weakref import threading import collections import os +import tempfile from dataclasses import dataclass # MLIR modules imports from .._mlir import ir +from .._mlir.dialects import llvm # Local modules imports from . import typing as t @@ -35,6 +38,7 @@ from .typing import get_c_pointers from .utils.logger import log from .utils.timer import timer + class CudaModuleAndKernel: """A loaded CUDA kernel and its metadata.""" @@ -132,17 +136,32 @@ def load_kernels_from_ir_module(module, kernel_info) -> list[CudaModuleAndKernel return list(kernel_modules.values()) - class KwargsWrapperSpec(NamedTuple): """A specification for keyword arguments wrapper.""" + arg_names: list[str] arg_defaults: tuple[Any, ...] kwonly_names: list[str] kwonly_defaults: dict[str, Any] +@dataclass +class ArgMeta: + """Metadata for function arguments.""" + + pos_names: list[str] + kwonly_names: list[str] + all_names: list[str] + annotated_types: list[object] + numeric_flags: list[bool] + name_to_index: dict[str, int] + pos_defaults: list[object] + kwonly_defaults: list[object] + arg_count: int + + class ExecutionArgs: - """Helper that wraps the function signature spec to filter exeuction and compile time arguments.""" + """Helper that wraps the function signature spec to filter execution and compile time arguments.""" def __init__(self, spec, function_name): self.function_name = function_name @@ -150,82 +169,235 @@ class ExecutionArgs: if spec is not None: self.args_spec = self.filter_runtime_arg_spec(spec) self.original_args_spec = spec + self._missing = object() + self._meta = self._build_meta() + self._tls = threading.local() + + def _build_meta(self): + """ + Precompute metadata for the fast-path execution. + This metadata is static per function signature. + """ + spec = self.args_spec + if spec is None: + return ArgMeta( + pos_names=[], + kwonly_names=[], + all_names=[], + annotated_types=[], + numeric_flags=[], + name_to_index={}, + pos_defaults=[], + kwonly_defaults=[], + arg_count=0, + ) + + pos_names = list(spec.args) + kwonly_names = list(spec.kwonlyargs) + all_names = pos_names + kwonly_names + + annotated_types = [spec.annotations.get(n) for n in all_names] + numeric_flags = [isinstance(typ, t.NumericMeta) for typ in annotated_types] + + name_to_index = {n: i for i, n in enumerate(all_names)} + + pos_defaults = [self._missing] * len(pos_names) + if spec.defaults: + start = len(pos_names) - len(spec.defaults) + for i, d in enumerate(spec.defaults): + pos_defaults[start + i] = d + + kwonly_defaults = [self._missing] * len(kwonly_names) + if spec.kwonlydefaults: + for i, n in enumerate(kwonly_names): + if n in spec.kwonlydefaults: + kwonly_defaults[i] = spec.kwonlydefaults[n] + + return ArgMeta( + pos_names=pos_names, + kwonly_names=kwonly_names, + all_names=all_names, + annotated_types=annotated_types, + numeric_flags=numeric_flags, + name_to_index=name_to_index, + pos_defaults=pos_defaults, + kwonly_defaults=kwonly_defaults, + arg_count=len(all_names), + ) + + def generate_execution_args_positional(self, *args): + """Fast execution for positional-only arguments with exact count match. + + This method is optimized for the common case where: + - All arguments are positional (no kwargs) + - Number of arguments matches the function signature exactly + + :param args: The positional arguments tuple + :return: (exe_args, adapted_args) tuple + """ + exe_arg_chunks = [None] * self._meta.arg_count + adapted_args = [] + + tls = self._tls + adapter_caches = getattr(tls, "adapter_caches", None) + if adapter_caches is None: + adapter_caches = [dict() for _ in range(self._meta.arg_count)] + tls.adapter_caches = adapter_caches + + annotated_types = self._meta.annotated_types + numeric_flags = self._meta.numeric_flags + + for index in range(self._meta.arg_count): + arg = args[index] + + cptr_method = getattr(arg, "__c_pointers__", None) + if cptr_method is not None: + exe_arg_chunks[index] = cptr_method() + continue + + if numeric_flags[index]: + arg = t.cast(arg, annotated_types[index]) + exe_arg_chunks[index] = get_c_pointers(arg) + else: + arg_type = type(arg) + cache = adapter_caches[index] + adapter = cache.get(arg_type) + if adapter is None: + adapter = JitArgAdapterRegistry.get_registered_adapter(arg_type) + if adapter is not None: + cache[arg_type] = adapter + + if adapter is not None: + arg = adapter(arg) + adapted_args.append(arg) + + exe_arg_chunks[index] = get_c_pointers(arg) + + exe_args = [p for chunk in exe_arg_chunks for p in chunk] + return exe_args, adapted_args def get_rectified_args(self, args, kwargs): """ This function is used to rectify the args and kwargs to a final runtime argument list according to the args_spec. """ - args_spec = self.args_spec - # Process positional arguments with defaults - rectified_args = list(args) - if args_spec.defaults and len(args) < len(args_spec.args): - rectified_args.extend(args_spec.defaults[len(args) - len(args_spec.args) :]) - for k, v in kwargs.items(): - if k in args_spec.args: - idx = args_spec.args.index(k) - if idx < len(rectified_args): - rectified_args[idx] = v - else: - rectified_args.append(v) - - # Process keyword arguments - rectified_kwargs = collections.OrderedDict( - (k, v) for k, v in kwargs.items() if k not in args_spec.args - ) - if args_spec.kwonlydefaults and len(rectified_kwargs) < len( - args_spec.kwonlyargs - ): - rectified_kwargs.update(args_spec.kwonlydefaults) - - # args/kwargs must match arg_specs - if len(rectified_args) != len(args_spec.args) or len(rectified_kwargs) != len( - args_spec.kwonlyargs - ): + pos_count = len(self._meta.pos_names) + if len(args) > pos_count: raise DSLRuntimeError( "input args/kwargs length does not match runtime function signature!", context={ - "input args length": len(rectified_args), - "input kwargs length": len(rectified_kwargs), - "function signature args length": len(args_spec.args), - "function signature kwonlyargs length": len(args_spec.kwonlyargs), + "input args length": len(args), + "input kwargs length": len(kwargs), + "function signature args length": len(self._meta.pos_names), + "function signature kwonlyargs length": len( + self._meta.kwonly_names + ), }, ) - return rectified_args + list(rectified_kwargs.values()) + + # Start with every slot marked missing, we overwrite as values/defaults bind + rectified = [self._missing] * self._meta.arg_count + pos_len = len(args) + + # Fill positional slots with the values from the caller + rectified[:pos_len] = args + + # Fill positional slots the caller skipped with the defaults + for i in range(pos_len, pos_count): + default = self._meta.pos_defaults[i] + if default is not self._missing: + rectified[i] = default + + # Fill keyword-only slots with the defaults before user kwargs + for j, default in enumerate(self._meta.kwonly_defaults): + idx = pos_count + j + if default is not self._missing: + rectified[idx] = default + + # Fill keyword slots with the values from the caller + for name, value in kwargs.items(): + idx = self._meta.name_to_index.get(name) + if idx is None: + raise DSLRuntimeError( + "unexpected keyword argument", + context={"argument_name": name}, + ) + if idx < pos_len: + raise DSLRuntimeError( + "multiple values for argument", + context={"argument_name": name}, + ) + rectified[idx] = value + + if self._missing in rectified: + missing_args = [ + name + for i, name in enumerate(self._meta.all_names) + if rectified[i] is self._missing + ] + raise DSLRuntimeError( + "input args/kwargs length does not match runtime function signature!", + context={ + "missing": missing_args, + "function signature args length": len(self._meta.pos_names), + "function signature kwonlyargs length": len( + self._meta.kwonly_names + ), + }, + ) + + return rectified def generate_execution_args(self, args, kwargs): """ This function is the prune version of `generate_mlir_function_types` which only generates execution args to get rid of mlir context. """ - args_spec = self.args_spec + if not kwargs and len(args) == self._meta.arg_count: + return self.generate_execution_args_positional(*args) - exe_args = [] + exe_arg_chunks = [None] * self._meta.arg_count adapted_args = [] + tls = self._tls + adapter_caches = getattr(tls, "adapter_caches", None) + if adapter_caches is None: + adapter_caches = [dict() for _ in range(self._meta.arg_count)] + tls.adapter_caches = adapter_caches + input_args = self.get_rectified_args(args, kwargs) - input_arg_names = args_spec.args + args_spec.kwonlyargs - for arg, arg_name in zip(input_args, input_arg_names): - # short-cut for args already converted - if hasattr(arg, "__c_pointers__"): - exe_args.extend(arg.__c_pointers__()) + + for index, arg in enumerate(input_args): + cptr_method = getattr(arg, "__c_pointers__", None) + if cptr_method is not None: + exe_arg_chunks[index] = cptr_method() continue - arg_type = args_spec.annotations.get(arg_name, None) + arg_type_anno = self._meta.annotated_types[index] - # Implicit cast to NumericMeta - if isinstance(arg_type, t.NumericMeta): - arg = t.cast(arg, arg_type) + if self._meta.numeric_flags[index]: + arg = t.cast(arg, arg_type_anno) + exe_arg_chunks[index] = get_c_pointers(arg) else: - # If not any known type, try registered adapter to do the conversion - adapter = JitArgAdapterRegistry.get_registered_adapter(type(arg)) - if adapter: + arg_type = type(arg) + cache = adapter_caches[index] + adapter = cache.get(arg_type) + if adapter is None: + adapter = JitArgAdapterRegistry.get_registered_adapter(arg_type) + if adapter is not None: + cache[arg_type] = adapter + + if adapter is not None: arg = adapter(arg) adapted_args.append(arg) - exe_args.extend(get_c_pointers(arg)) + exe_arg_chunks[index] = get_c_pointers(arg) + + exe_args = [p for chunk in exe_arg_chunks for p in chunk] return exe_args, adapted_args - def get_kwargs_wrapper_spec(self, exclude_arg_names: Sequence[str] = ()) -> KwargsWrapperSpec: + def get_kwargs_wrapper_spec( + self, exclude_arg_names: Sequence[str] = () + ) -> KwargsWrapperSpec: """ This function is used to get the kwargs wrapper spec from the original args_spec. """ @@ -242,7 +414,7 @@ class ExecutionArgs: kwonly_names = [] kwonly_defaults = {} - # Filter arguments and maintain their properties + # Filter arguments and maintain their properties for i, arg_name in enumerate(arg_spec.args): arg_type = arg_spec.annotations.get(arg_name, None) @@ -320,7 +492,7 @@ class ExecutionArgs: "function_name": self.function_name, "expected_args": len(arg_spec.args), "provided_args": len(full_args), - } + }, ) # Filter keyword-only arguments @@ -339,21 +511,23 @@ class ExecutionArgs: elif arg_spec.kwonlydefaults and kwarg in arg_spec.kwonlydefaults: runtime_kwargs[kwarg] = arg_spec.kwonlydefaults[kwarg] - if (len(runtime_args) != len(self.args_spec.args) or - len(runtime_kwargs) != len(self.args_spec.kwonlyargs)): + if len(runtime_args) != len(self.args_spec.args) or len(runtime_kwargs) != len( + self.args_spec.kwonlyargs + ): raise DSLRuntimeError( "input args/kwargs length does not match runtime function signature!", context={ "input args length": len(runtime_args), "input kwargs length": len(runtime_kwargs), "function signature args length": len(self.args_spec.args), - "function signature kwonlyargs length": len(self.args_spec.kwonlyargs), + "function signature kwonlyargs length": len( + self.args_spec.kwonlyargs + ), }, ) return runtime_args + list(runtime_kwargs.values()) - def filter_runtime_arg_spec(self, arg_spec: inspect.FullArgSpec): runtime_args = [] runtime_annotations = {} @@ -444,9 +618,11 @@ class JitExecuteContext: def __init__( self, module: "JitModule", - kernel_fns=[], + kernel_fns=None, context: Optional[cuda_helpers.DevicePrimaryContext] = None, ): + if kernel_fns is None: + kernel_fns = [] self.module = module self.kernel_functions = kernel_fns self.kernel_functions_ptrs = [ctypes.c_void_p(k.getPtr()) for k in kernel_fns] @@ -538,27 +714,71 @@ class JitExecutor: # are garbage collected. self.jit_module = jit_module self.exec_context = exec_context - self.profiler = timer(enable=jit_time_profiling) + self.profiler = timer(enable=jit_time_profiling) if jit_time_profiling else None # Get the cuda result type from the capi function. # This is only set to i32 if CudaDialectJitModule is used. cuda_result_type = self.jit_module.capi_func.restype self.cuda_result = cuda_result_type() if cuda_result_type is not None else None + # Pre-compute flags + self._has_cuda_result = self.cuda_result is not None + self._has_profiler = self.profiler is not None + self._cuda_result_addr = ( + ctypes.addressof(self.cuda_result) if self._has_cuda_result else None + ) + + if jit_time_profiling: + self._get_invoke_packed_args_func = self.profiler( + self._get_invoke_packed_args + ) + self.capi_func = self.profiler(self.jit_module.capi_func) + else: + self._get_invoke_packed_args_func = self._get_invoke_packed_args + self.capi_func = self.jit_module.capi_func + + # Pre-compute packed args metadata + self._num_extra_args = 0 + if self._has_cuda_result: + self._num_extra_args += 1 + if self.exec_context is not None: + self._kernel_ptrs = self.exec_context.kernel_functions_ptrs + self._num_extra_args += len(self._kernel_ptrs) + else: + self._kernel_ptrs = None + self._tls = threading.local() + # Assume each execution args has type `c_void_p` to reduce the overhead of `ctypes.cast`. def _get_invoke_packed_args(self, exe_args): - # If expecting a cuda result, add a pointer to exe_args - if self.cuda_result is not None: - exe_args.append(ctypes.addressof(self.cuda_result)) - if self.exec_context is not None: - exe_args += self.exec_context.kernel_functions_ptrs - packed_args = (ctypes.c_void_p * len(exe_args))() - for argNum in range(len(exe_args)): - arg = exe_args[argNum] - if isinstance(arg, ctypes.c_void_p): - packed_args[argNum] = arg - else: - packed_args[argNum] = ctypes.c_void_p(arg).value + # Pre-calculate sizes once during init and cache + num_base_args = len(exe_args) + total_args = num_base_args + self._num_extra_args + + # Re-use packed args buffer if possible + tls = self._tls + packed_args = getattr(tls, "packed_args", None) + capacity = getattr(tls, "capacity", 0) + + if packed_args is None or capacity < total_args: + packed_args = (ctypes.c_void_p * total_args)() + tls.packed_args = packed_args + tls.capacity = total_args + + idx = num_base_args + if self._has_cuda_result: + packed_args[idx] = self._cuda_result_addr + idx += 1 + if self._kernel_ptrs is not None: + for ptr in self._kernel_ptrs: + packed_args[idx] = ptr + idx += 1 + + for i in range(num_base_args): + arg = exe_args[i] + packed_args[i] = ( + arg if type(arg) is ctypes.c_void_p else ctypes.c_void_p(arg).value + ) + return packed_args def generate_execution_args(self, *args, **kwargs): @@ -566,17 +786,17 @@ class JitExecutor: def run_compiled_program(self, exe_args): try: - packed_args = self.profiler(self._get_invoke_packed_args)(exe_args) - self.profiler(self.jit_module.capi_func)(packed_args) - if self.cuda_result is not None: - if self.cuda_result.value != 0: - error_code = self.cuda_result.value - error_name = cuda_helpers._cudaGetErrorEnum( - cuda_helpers.cuda.CUresult(error_code) - ) - raise DSLCudaRuntimeError(error_code, error_name) - return self.cuda_result.value - return None + packed_args = self._get_invoke_packed_args_func(exe_args) + self.capi_func(packed_args) + if not self._has_cuda_result: + return None + error_code = self.cuda_result.value + if error_code == 0: + return error_code + error_name = cuda_helpers._cudaGetErrorEnum( + cuda_helpers.cuda.CUresult(error_code) + ) + raise DSLCudaRuntimeError(error_code, error_name) except DSLCudaRuntimeError as e: raise e except Exception as e: @@ -584,7 +804,7 @@ class JitExecutor: def __call__(self, *args, **kwargs): exe_args, adapted_args = self.generate_execution_args(*args, **kwargs) - self.run_compiled_program(exe_args) + return self.run_compiled_program(exe_args) @dataclass @@ -616,9 +836,34 @@ class JitFunctionArtifacts: raise DSLRuntimeError(f"Failed to read MLIR file '{self.MLIR}': {e}") +class ExportProvider: + """Holds the dsl specific settings for the export of the jit-compiled function.""" + + dsl: "Type[BaseDSL]" = None + arg_spec_processor: "ArgsSpecProcessor" = None + c_header_generator: "CHeaderGenerator" = None + object_file_version: str = None + + def __init__( + self, + dsl: "Type[BaseDSL]", + arg_spec_processor: "ArgsSpecProcessor", + c_header_generator: "CHeaderGenerator", + object_file_version: str, + mlirExecutionEngine: "MlirExecutionEngine", + ): + self.dsl = dsl + self.arg_spec_processor = arg_spec_processor + self.c_header_generator = c_header_generator + self.object_file_version = object_file_version + self.mlirExecutionEngine = mlirExecutionEngine + + class JitCompiledFunction: """Holds a compiled function.""" + export_provider: ExportProvider = None + def __init__( self, ir_module, @@ -631,6 +876,8 @@ class JitCompiledFunction: jit_function_artifacts, prefix=None, load_from_binary=False, + dynamic_args=None, + dynamic_kwargs=None, ): self.ir_module = ir_module self.engine = engine @@ -656,6 +903,9 @@ class JitCompiledFunction: self._executor_lock = threading.RLock() self._default_executor = None + # This is used to do early generation of the c header arguments to release the reference to the dynamic arguments. + self._generate_c_header_arguments(dynamic_args, dynamic_kwargs) + @property def __ptx__(self): """Returns the PTX code of the JIT-compiled function.""" @@ -704,7 +954,7 @@ class JitCompiledFunction: if self.engine is None: raise DSLRuntimeError( "The compiled function does not have a valid execution engine.", - suggestion="For cross-compilation, please use `cute.export.export_to_c` to serialize the compiled function and load/execute it on target device.", + suggestion="For cross-compilation, please use `JitCompiledFunction.export_to_c` to serialize the compiled function and load/execute it on target device.", ) def to(self, device=None) -> JitExecutor: @@ -735,11 +985,6 @@ class JitCompiledFunction: context = self.jit_module.get_device_execute_context(device) return JitExecutor(self.jit_module, context, self.jit_time_profiling) - def set_dynamic_args(self, dynamic_args, dynamic_kwargs): - """Sets the dynamic argument information required for export to c code generation.""" - self.dynamic_args = dynamic_args - self.dynamic_kwargs = dynamic_kwargs - def generate_execution_args(self, *args, **kwargs): return self.args_spec.generate_execution_args(args, kwargs) @@ -750,7 +995,10 @@ class JitCompiledFunction: CUDA errors. If you need to call the kernel on multiple devices use `to` to return a per-device function. """ - exe_args, adapted_args = self.generate_execution_args(*args, **kwargs) + exe_args, adapted_args = self.args_spec.generate_execution_args(args, kwargs) + executor = self._default_executor + if executor is not None: # Only lock on first call + return executor.run_compiled_program(exe_args) return self.run_compiled_program(exe_args) def run_compiled_program(self, exe_args): @@ -768,3 +1016,168 @@ class JitCompiledFunction: proxy_self = weakref.proxy(self) self._default_executor = proxy_self.to(None) return self._default_executor.run_compiled_program(exe_args) + + def _generate_c_header_arguments(self, dynamic_args, dynamic_kwargs): + """Generates the c header arguments for the AOT C header generation.""" + self.c_header_arguments = None + from .export import CHeaderArguments + + if dynamic_args is not None or dynamic_kwargs is not None: + self.dummy_prefix_name = "dummy_prefix_name" + try: + # This arguments may be generated failure due to not all the arguments (e.g. custom types) are supported by the AOT C header generator. + c_header_arguments, packed_args, declarations = ( + self.export_provider.c_header_generator._generate_arguments( + self.dummy_prefix_name, + self.args_spec, + dynamic_args, + dynamic_kwargs, + ) + ) + self.c_header_arguments = CHeaderArguments( + self.dummy_prefix_name, + c_header_arguments, + packed_args, + declarations, + None, + ) + except Exception as e: + self.c_header_arguments = CHeaderArguments( + self.dummy_prefix_name, [], [], [], str(e) + ) + + def dump_to_object( + self, + function_prefix: str, + ) -> bytes: + """Dump the compiled ir function to a bytes object with ELF format. The bytes object contains the host + launch entry function and cubin inside. + + @param function_prefix: The prefix name of the function. This is the user provided unique identifier name of the function to avoid symbol conflict in the generated object file. + @return: The bytes object of the function. + """ + # lazy import to avoid circular dependency + from .export import get_export_module, encode_metadata_into_ir_module + + assert self.export_provider is not None, ( + "Export provider is not set for JitCompiledFunction." + ) + export_module = get_export_module(self.ir_module, function_prefix) + export_module = encode_metadata_into_ir_module( + function_prefix, + export_module, + self.args_spec.args_spec, + self.function_name, + self.kernel_info, + self.export_provider.arg_spec_processor, + self.export_provider.object_file_version, + ) + + cubin_data = None + + def strip_gpu_binary_op(op): + if op.name == "gpu.binary": + s = io.BytesIO() + op.operation.write_bytecode(s) + nonlocal cubin_data + cubin_data = s.getvalue() + cubin_data = cubin_data.split(b'bin = "')[1].split(b'">')[0] + cubin_data = get_escaped_cubin_bytes(cubin_data) + op.erase() + return ir.WalkResult.ADVANCE + return ir.WalkResult.ADVANCE + + # Strip gpu related to avoid the object file generating builtin module load/unload functions + export_module.operation.walk(strip_gpu_binary_op) + cubin_suffix = "cubin" + if cubin_data is not None: + cubin_array = array.array("b", cubin_data) + with ( + export_module.context, + ir.Location.unknown(), + ir.InsertionPoint(export_module.body), + ): + new_binary_global_op = llvm.GlobalOp( + sym_name="_".join([function_prefix, cubin_suffix]), + global_type=ir.Type.parse(f"!llvm.array<{len(cubin_array)} x i8>"), + linkage=ir.Attribute.parse("#llvm.linkage"), + value=ir.DenseIntElementsAttr.get(cubin_array), + constant=True, + ) + + if "gpu.container_module" in export_module.operation.attributes: + del export_module.operation.attributes["gpu.container_module"] + # Generate the object file + + try: + with tempfile.NamedTemporaryFile() as tmp_object_file: + self.export_provider.mlirExecutionEngine.dump_object_file_pic( + export_module, + tmp_object_file.name, + "_".join([function_prefix, self.function_name]), + ) + with open(tmp_object_file.name, "rb") as f: + ret = f.read() + return ret + except Exception as e: + raise DSLRuntimeError(f"Error dumping object file: {e}") from e + + def export_to_c( + self, + file_path: str, + file_name: str, + function_prefix: str = "", + ): + """Exports the jit-compiled function to a C compatible files(header/library). + This is used for c/cpp AOT support. + The `file_path` will be used as the directory to save the header and object files. + The `file_name` will be used as the name of the header and object files. The same file name will always overwrite the existing file. + The `function_prefix` will be used as the symbol prefix of the generated functions, it is guaranteed by + the caller that the generated functions are unique. + + + The c header file is generated with following components: + 1. The host launch entry function. And the structure definitions of the arguments. + 2. The device metadata load/unload functions. + 3. The cubin data array and len. + + The library contains the binary of the underlying host launch entry function. + + @param jit_function: The jit-compiled function from `cute.compile`. + @param file_path: The path to the directory where the header and object files will be saved. + @param file_name: The name of the header and object files. + @param function_prefix: The prefix of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file. Default to the `file_name`. + """ + if function_prefix is None or function_prefix == "": + function_prefix = file_name + + assert self.export_provider is not None, ( + "Export provider is not set for JitCompiledFunction." + ) + # lazy import to avoid circular dependency + from .export import get_export_module + + export_module = get_export_module(self.ir_module, function_prefix) + # Generate the c header file + header_file_content = self.export_provider.c_header_generator( + function_prefix, + export_module, + self.args_spec, + self.function_name, + self.kernel_info, + self.c_header_arguments, + self.export_provider.dsl._get_dsl().name, + ) + try: + with open(os.path.join(file_path, file_name + ".h"), "w") as f: + f.write(header_file_content) + except Exception as e: + raise DSLRuntimeError(f"Error writing header file: {e}") from e + + # Generate the object file + object_file_content = self.dump_to_object(function_prefix) + try: + with open(os.path.join(file_path, file_name + ".o"), "wb") as f: + f.write(object_file_content) + except Exception as e: + raise DSLRuntimeError(f"Error writing object file: {e}") from e diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py index 940e27d5..e7a60a3d 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py @@ -347,6 +347,16 @@ def initialize_cuda_context(device_id: int = 0, flags: int = 0): # Initialize CUDA Driver API _log().info(f"cuInit {flags}") checkCudaErrors(cuda.cuInit(flags)) + + driver_version = get_driver_version() + + # Check the CUDA driver version works for the installed cuda-python package + if driver_version < 13000 and cuda.CUDA_VERSION >= 13000: + raise DSLRuntimeError( + f"CUDA driver version {driver_version} is below the minimum required version for the installed cuda-python package {cuda.CUDA_VERSION}.", + suggestion=f"Consider updating your NVIDIA driver to version 580 or above. Or install cuda-python package with version 12.9 or below.", + ) + # Retrieve handle for device cuDevice = get_device(device_id) # Create context diff --git a/python/CuTeDSL/cutlass/base_dsl/runtime/tensor_descriptor.py b/python/CuTeDSL/cutlass/base_dsl/runtime/tensor_descriptor.py index 864e6675..5b946950 100644 --- a/python/CuTeDSL/cutlass/base_dsl/runtime/tensor_descriptor.py +++ b/python/CuTeDSL/cutlass/base_dsl/runtime/tensor_descriptor.py @@ -94,6 +94,14 @@ class TensorDescriptor: return get_tensor_desc_device_id(self._capsule) return -1 + @property + def pointer(self): + """ + Returns the pointer to the tensor data. This is either the device pointer or the data pointer if the data is not + in a device. + """ + return self.device_pointer if self.device_pointer is not None else self.data_ptr + @property def element_type(self): """Return the corresponding Python type based on DLPack dtype metadata.""" diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py index 1bc21a7e..0ce03475 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py @@ -113,15 +113,20 @@ class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): if isinstance(stride, int) and stride == 1: return i raise ValueError("stride=1 index not found, needed for f4 tensor") + stride_one_index = find_stride_one_index() - def map_shape_for_tensor_dtype_f4x2_to_f4(index: int, value: ir.Value) -> ir.Value: + def map_shape_for_tensor_dtype_f4x2_to_f4( + index: int, value: ir.Value + ) -> ir.Value: if index == stride_one_index: with ir.InsertionPoint(current_block): return self.mul(value, self.integer_constant(value.type, 2)) return value - def map_stride_for_tensor_dtype_f4x2_to_f4(index, value: ir.Value) -> ir.Value: + def map_stride_for_tensor_dtype_f4x2_to_f4( + index, value: ir.Value + ) -> ir.Value: if index != stride_one_index: with ir.InsertionPoint(current_block): return self.mul(value, self.integer_constant(value.type, 2)) @@ -141,7 +146,9 @@ class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): if param.strides is not None: for index, dim in enumerate(param.strides): if isinstance(dim, spec.Var): - strides.append(map_stride_value(index, context.matched_var_binding[dim])) + strides.append( + map_stride_value(index, context.matched_var_binding[dim]) + ) flatten_struct, alloca = self.pack_values_to_alloca( current_block, context.entry_block, [data, *shape, *strides] ) diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py index 2d6c9d11..2d5ab527 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py @@ -233,7 +233,9 @@ class MLIRBuilder(MLIRTypeBuilder): remainder = llvm.urem(value, align_val) return self.equal(remainder, self.i64(0)) - def br(self, target_block: ir.Block, *, args: Optional[list[ir.Value]] = None) -> None: + def br( + self, target_block: ir.Block, *, args: Optional[list[ir.Value]] = None + ) -> None: """Create an unconditional branch. Parameters @@ -338,7 +340,7 @@ class MLIRBuilder(MLIRTypeBuilder): false_dest_operands=false_dest_operands, true_dest=true_block, false_dest=false_block, - branch_weights=ir.DenseI32ArrayAttr.get(list(branch_weights)) + branch_weights=ir.DenseI32ArrayAttr.get(list(branch_weights)), ) else: llvm.cond_br( @@ -436,8 +438,7 @@ class MLIRBuilder(MLIRTypeBuilder): internal: bool = False, llvm_func_attrs: Sequence[str] = (), ) -> tuple[list[ir.Value], ir.Block]: - """Create a function with the given signature. - """ + """Create a function with the given signature.""" func_op = llvm.func( name, function_type=self.as_attr( @@ -476,7 +477,9 @@ class MLIRBuilder(MLIRTypeBuilder): ) func_op.attributes["llvm.linkage"] = ir.StringAttr.get("external") - def create_alloca(self, entry_block: ir.Block, alloca_type: ir.Type, array_size: int) -> ir.Value: + def create_alloca( + self, entry_block: ir.Block, alloca_type: ir.Type, array_size: int + ) -> ir.Value: """Create an alloca operation.""" with ir.InsertionPoint(entry_block.operations[0]): # declare the struct type diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py index d9d76904..13c24835 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py @@ -10,6 +10,7 @@ # is strictly prohibited. """Kernel specification classes for TVM-FFI function parameters.""" + from abc import ABC from collections.abc import Sequence @@ -20,6 +21,7 @@ try: except ModuleNotFoundError: pass + class DefaultConfig: """Default configuration with context manager support.""" @@ -105,7 +107,7 @@ class Var(Param): dtype: Union[str, "tvm_ffi.dtype"], *, divisibility: Optional[int] = None, - ) -> None: + ) -> None: """Initialize a Var parameter. Parameters @@ -138,7 +140,8 @@ class Shape(Param): def __init__( self, - name: str, shape: list[Union[int, Var]], + name: str, + shape: list[Union[int, Var]], ) -> None: """Initialize a Shape parameter. @@ -220,7 +223,9 @@ class Tensor(Param): self.data = Var(name + ".data", tvm_ffi.dtype("handle")) self.shape: list[Union[int, Var]] = list(shape) self.dtype = tvm_ffi.dtype(dtype) - self.strides: Optional[list[Var]] = list(strides) if strides is not None else None + self.strides: Optional[list[Var]] = ( + list(strides) if strides is not None else None + ) self.data_alignment = data_alignment # Use default device type if none specified @@ -319,6 +324,7 @@ class ConstNone(Param): name : str The parameter name. """ + name: str def __init__(self, name: str) -> None: @@ -340,6 +346,7 @@ class TupleParam(Param): name : str The parameter name. """ + name: str params: list[Param] @@ -471,6 +478,7 @@ def create_map_tensor_dtype_f4x2_to_f4_spec(f4_tensor_spec: Tensor) -> Tensor: if isinstance(stride, int) and stride == 1: return i raise ValueError("Cannot find dimension with stride=1") + stride_one_index = find_stride_one_index() def divisibility_divide_by_2(value: Var) -> Optional[int]: @@ -487,7 +495,9 @@ def create_map_tensor_dtype_f4x2_to_f4_spec(f4_tensor_spec: Tensor) -> Tensor: raise ValueError(f"Dimension {index} with stride=1 must be even") return value // 2 # create a new var with the same name and dtype - return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value)) + return Var( + value.name, value.dtype, divisibility=divisibility_divide_by_2(value) + ) return value def map_stride(index: int, value: Union[int, Var]) -> Union[int, Var]: @@ -497,7 +507,9 @@ def create_map_tensor_dtype_f4x2_to_f4_spec(f4_tensor_spec: Tensor) -> Tensor: raise ValueError(f"Dimension {index} with stride != 1 must be even") return value // 2 # create a new var with the same name and dtype - return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value)) + return Var( + value.name, value.dtype, divisibility=divisibility_divide_by_2(value) + ) return value new_shape = [map_shape(i, x) for i, x in enumerate(f4_tensor_spec.shape)] diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py index 4b976d0c..9948d500 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py @@ -38,6 +38,7 @@ class ArgContext: :ivar tuple_indices: List of tuple indices for nested tuple access (e.g., [0, 1] for tuple[0][1]). :vartype tuple_indices: list[int] """ + param_name: str arg_index: int tuple_indices: list[int] @@ -54,7 +55,10 @@ class ArgContext: else: # Nested tuple element: "on my_tuple[0][1] in argument #0" indices_str = "".join(f"[{i}]" for i in self.tuple_indices) - return [f"on {self.param_name}{indices_str} in argument ", f"#{self.arg_index}"] + return [ + f"on {self.param_name}{indices_str} in argument ", + f"#{self.arg_index}", + ] def get_field_name(self, field_suffix: str) -> str: """Get the field name with tuple indices for shape/stride access. @@ -289,7 +293,9 @@ class TVMFFIBuilder(MLIRBuilder): elem_type=self.tvm_ffi_object_type, ) - def load_ffi_any_array_item_type_index(self, args: ir.Value, index: int) -> ir.Value: + def load_ffi_any_array_item_type_index( + self, args: ir.Value, index: int + ) -> ir.Value: """Get the type index from the index-th field of tvm_ffi_any_type struct. Semantics as follows: @@ -321,11 +327,7 @@ class TVMFFIBuilder(MLIRBuilder): ) return llvm.load(self.i32_type, type_index_ptr) - def load_ffi_any_array_item_v_int64( - self, - args: ir.Value, - index: int - ) -> ir.Value: + def load_ffi_any_array_item_v_int64(self, args: ir.Value, index: int) -> ir.Value: """Get the v_int64 from the index-th field of tvm_ffi_any_type struct. Semantics as follows: @@ -362,10 +364,7 @@ class TVMFFIBuilder(MLIRBuilder): return llvm.load(self.f64_type, v_float64_ptr) def load_ffi_any_array_item_v_ptr( - self, - args: ir.Value, - index: int, - address_space: Optional[int] = None + self, args: ir.Value, index: int, address_space: Optional[int] = None ) -> ir.Value: """Get the v_ptr from the index-th field of tvm_ffi_any_type struct. @@ -546,8 +545,10 @@ class TVMFFIBuilder(MLIRBuilder): ) -> ir.Value: """Downcast i64 to lower bits.""" overflow_flags = llvm.IntegerOverflowFlags.none - if (hasattr(tvm_ffi._dtype.DataTypeCode, "BOOL") and - target_dtype.type_code == tvm_ffi._dtype.DataTypeCode.BOOL): + if ( + hasattr(tvm_ffi._dtype.DataTypeCode, "BOOL") + and target_dtype.type_code == tvm_ffi._dtype.DataTypeCode.BOOL + ): # LLVM use i1 (boolean) for boolean return llvm.icmp(llvm.ICmpPredicate.ne, v_int64, self.i64(0)) if target_dtype.bits == 64: @@ -805,7 +806,7 @@ class TVMFFIBuilder(MLIRBuilder): true_block=subsequent_block, false_block=error_block, # likely to be true - branch_weights=self.BRANCH_WEIGHTS_LIKELY + branch_weights=self.BRANCH_WEIGHTS_LIKELY, ) with ir.InsertionPoint(error_block): self.raise_error_and_return(error_kind, error_message_parts) @@ -869,7 +870,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode the integer parameter at the given index.""" # read the type index with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) # Check if type is int or bool (both use v_int64, bool can be converted to int) is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) is_bool = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIBool)) @@ -915,7 +918,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode the float parameter at the given index.""" # read the type index with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) # Check if type is float, int, or bool (int and bool can be converted to float) is_float = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIFloat)) is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) @@ -954,7 +959,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # Handle float type with ir.InsertionPoint(float_block): - v_float64: ir.Value = self.load_ffi_any_array_item_v_float64(args, arg_index) + v_float64: ir.Value = self.load_ffi_any_array_item_v_float64( + args, arg_index + ) if param.dtype.bits == 64: float_result = v_float64 elif param.dtype.bits == 32: @@ -1023,7 +1030,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode the opaque handle parameter at the given index.""" # read the type index with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) # Check if type is opaque pointer is_opaque_ptr = self.equal( type_index, self.i32(TVMFFITypeIndex.kTVMFFIOpaquePtr) @@ -1075,7 +1084,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode the opaque handle parameter at the given index.""" # read the type index with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) # Check if type is a nullptr is_nullptr = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFINone)) @@ -1108,6 +1119,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): return current_block from tvm_ffi._dtype import DataTypeCode + is_uint = dtype.type_code == DataTypeCode.UINT # compute out the upper and lower bounds @@ -1121,7 +1133,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): error_msg = [ "Out of bound ", *error_msg_context, - f", expected to be in {str(dtype)} range [{lower_bound}, {upper_bound}]" + f", expected to be in {str(dtype)} range [{lower_bound}, {upper_bound}]", ] # Check bounds using appropriate comparison predicates @@ -1129,19 +1141,24 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lower_bound_value = self.i64(lower_bound) upper_bound_value = self.i64(upper_bound) if is_uint: - is_above_lower = llvm.icmp(llvm.ICmpPredicate.uge, value, lower_bound_value) - is_below_upper = llvm.icmp(llvm.ICmpPredicate.ule, value, upper_bound_value) + is_above_lower = llvm.icmp( + llvm.ICmpPredicate.uge, value, lower_bound_value + ) + is_below_upper = llvm.icmp( + llvm.ICmpPredicate.ule, value, upper_bound_value + ) else: - is_above_lower = llvm.icmp(llvm.ICmpPredicate.sge, value, lower_bound_value) - is_below_upper = llvm.icmp(llvm.ICmpPredicate.sle, value, upper_bound_value) + is_above_lower = llvm.icmp( + llvm.ICmpPredicate.sge, value, lower_bound_value + ) + is_below_upper = llvm.icmp( + llvm.ICmpPredicate.sle, value, upper_bound_value + ) in_bounds = self.and_(is_above_lower, is_below_upper) return self.check_condition( - current_block, - lambda: in_bounds, - "ValueError", - error_msg + current_block, lambda: in_bounds, "ValueError", error_msg ) def check_int_value_divisibility( @@ -1173,6 +1190,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ir.Block The updated block after checks. """ + def check_divisibility() -> ir.Value: cond = self.i64_divisible_const(value, divisibility) if skip_check_predicate is not None: @@ -1182,14 +1200,11 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): error_msg = [ "Invalid ", *error_msg_context, - f", expected to be divisible by {divisibility}" + f", expected to be divisible by {divisibility}", ] return self.check_condition( - current_block, - check_divisibility, - "ValueError", - error_msg + current_block, check_divisibility, "ValueError", error_msg ) def set_or_check_matched_var_binding( @@ -1218,14 +1233,15 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # check divisibility if specified if var.divisibility is not None: current_block = self.check_int_value_divisibility( - current_block, value, var.divisibility, error_msg_context, + current_block, + value, + var.divisibility, + error_msg_context, skip_check_predicate=skip_check_predicate, ) # store the source value with parameter info with ir.InsertionPoint(current_block): - target_value = self.downcast_i64_to_lower_bits( - value, var.dtype - ) + target_value = self.downcast_i64_to_lower_bits(value, var.dtype) else: target_value = value # store the source value @@ -1251,7 +1267,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): error_msg_mismatch = [ error_prefix_mismatch, *error_msg_context, - f", expected to be {var}" + f", expected to be {var}", ] def check_value_mismatch() -> ir.Value: @@ -1261,10 +1277,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): return cond return self.check_condition( - current_block, - check_value_mismatch, - error_kind, - error_msg_mismatch + current_block, check_value_mismatch, error_kind, error_msg_mismatch ) def set_or_check_matched_var_binding_from_shape( @@ -1288,8 +1301,12 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self._fn_call_context, ] return self.set_or_check_matched_var_binding( - current_block, var, value, error_msg, arg_field_name, - skip_check_predicate=skip_check_predicate + current_block, + var, + value, + error_msg, + arg_field_name, + skip_check_predicate=skip_check_predicate, ) def decode_param_shape_from_ffi_array( @@ -1325,9 +1342,13 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ) -> tuple[ir.Block, ir.Value]: """Validate and load a single shape element from the array.""" with ir.InsertionPoint(block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(array_data, index) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + array_data, index + ) # Check if type is int or bool (both use v_int64, bool can be converted to int) - is_int_val = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) + is_int_val = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt) + ) # Check that the element is an integer field_name = arg_context.get_field_name("") @@ -1346,7 +1367,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ) with ir.InsertionPoint(block): - v_int64: ir.Value = self.load_ffi_any_array_item_v_int64(array_data, index) + v_int64: ir.Value = self.load_ffi_any_array_item_v_int64( + array_data, index + ) return block, v_int64 @@ -1382,7 +1405,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ) with ir.InsertionPoint(current_block): - load_shapes = [self.load_i64_array_item(shape_data, i) for i in range(len(param.shape))] + load_shapes = [ + self.load_i64_array_item(shape_data, i) for i in range(len(param.shape)) + ] return (current_block, load_shapes) @@ -1396,10 +1421,16 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ) -> ir.Block: """Decode the shape parameter at the given index.""" with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) # Check if type is ffi.Shape or ffi.Array - is_ffi_shape = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIShape)) - is_ffi_array = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIArray)) + is_ffi_shape = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFIShape) + ) + is_ffi_array = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFIArray) + ) # Create error block and subsequent blocks error_block = current_block.create_after() @@ -1407,7 +1438,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ffi_array_check_block = ffi_shape_block.create_after() ffi_array_block = ffi_array_check_block.create_after() # Create subsequent_block with i64 arguments for each shape dimension - subsequent_block = ffi_array_block.create_after(*[self.i64_type] * len(param.shape)) + subsequent_block = ffi_array_block.create_after( + *[self.i64_type] * len(param.shape) + ) # Branch from current_block: check FFI shape first with ir.InsertionPoint(current_block): @@ -1427,15 +1460,17 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # ffi array check block: verify it's actually an Array with ir.InsertionPoint(ffi_array_check_block): self.cond_br( - is_ffi_array, ffi_array_block, error_block, - branch_weights=self.BRANCH_WEIGHTS_LIKELY + is_ffi_array, + ffi_array_block, + error_block, + branch_weights=self.BRANCH_WEIGHTS_LIKELY, ) # ffi array block with ir.InsertionPoint(ffi_array_block): array_cell_ptr = self.get_object_cell_ptr( self.load_ffi_any_array_item_v_ptr(args, arg_index) - ) + ) ffi_array_block, load_shapes = self.decode_param_shape_from_ffi_array( ffi_array_block, param, arg_context, array_cell_ptr ) @@ -1476,9 +1511,15 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode tensor step0: check index and find out DLTensor*.""" # read the type index with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) - is_ffi_tensor = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFITensor)) - is_dl_tensor = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIDLTensorPtr)) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) + is_ffi_tensor = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFITensor) + ) + is_dl_tensor = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFIDLTensorPtr) + ) # Create error block and subsequent block error_block = current_block.create_after() @@ -1493,15 +1534,19 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # ffi tensor block with ir.InsertionPoint(ffi_tensor_block): - ffi_tensor_ptr: ir.Value = self.load_ffi_any_array_item_v_ptr(args, arg_index) + ffi_tensor_ptr: ir.Value = self.load_ffi_any_array_item_v_ptr( + args, arg_index + ) dl_tensor_ptr = self.get_object_cell_ptr(ffi_tensor_ptr) self.br(subsequent_block, args=[dl_tensor_ptr]) # dltensor check block: verify it's actually a DLTensor with ir.InsertionPoint(dl_tensor_check_block): self.cond_br( - is_dl_tensor, dl_tensor_block, error_block, - branch_weights=self.BRANCH_WEIGHTS_LIKELY + is_dl_tensor, + dl_tensor_block, + error_block, + branch_weights=self.BRANCH_WEIGHTS_LIKELY, ) # dltensor block @@ -1552,6 +1597,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # check data alignment if specified if param.data_alignment is not None: + def check_alignment() -> ir.Value: # Convert pointer to integer data_as_int = llvm.ptrtoint(self.i64_type, data) @@ -1574,7 +1620,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # store the matched values, these do not need constraint checks self.matched_var_binding[param.data] = data self.matched_var_source[param.data] = param.data - self.matched_var_arg_field_name[param.data] = arg_context.get_field_name(".data") + self.matched_var_arg_field_name[param.data] = arg_context.get_field_name( + ".data" + ) # check device_id constraint if user specifies a device_id variable current_block = self.set_or_check_matched_var_binding( @@ -1718,8 +1766,12 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode the stream parameter at the given index.""" # stream is decoded as opaque handle return self.decode_param_opaque_handle( - current_block, param.var, args, arg_index, arg_context, - allow_int_as_ptr=True + current_block, + param.var, + args, + arg_index, + arg_context, + allow_int_as_ptr=True, ) def decode_param_data_pointer( @@ -1790,8 +1842,12 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """Decode the tuple parameter at the given index.""" # Check if type is kTVMFFIArray with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) - is_ffi_array = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIArray)) + type_index: ir.Value = self.load_ffi_any_array_item_type_index( + args, arg_index + ) + is_ffi_array = self.equal( + type_index, self.i32(TVMFFITypeIndex.kTVMFFIArray) + ) # Check that the type is an array current_block = self.check_condition( @@ -1831,7 +1887,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): for i, tuple_param in enumerate(param.params): # Create nested context for the tuple element nested_context = arg_context.get_element_context(i) - current_block = self.decode_param(current_block, tuple_param, array_data, i, nested_context) + current_block = self.decode_param( + current_block, tuple_param, array_data, i, nested_context + ) return current_block @@ -1860,15 +1918,25 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): """ if isinstance(param, spec.Var): if param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.INT: - return self.decode_param_int(current_block, param, args, arg_index, arg_context) + return self.decode_param_int( + current_block, param, args, arg_index, arg_context + ) elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.UINT: # UINT uses the same logic as INT since both are stored in v_int64 - return self.decode_param_int(current_block, param, args, arg_index, arg_context) - elif (hasattr(tvm_ffi._dtype.DataTypeCode, "BOOL") and - param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.BOOL): - return self.decode_param_int(current_block, param, args, arg_index, arg_context) + return self.decode_param_int( + current_block, param, args, arg_index, arg_context + ) + elif ( + hasattr(tvm_ffi._dtype.DataTypeCode, "BOOL") + and param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.BOOL + ): + return self.decode_param_int( + current_block, param, args, arg_index, arg_context + ) elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.FLOAT: - return self.decode_param_float(current_block, param, args, arg_index, arg_context) + return self.decode_param_float( + current_block, param, args, arg_index, arg_context + ) elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.HANDLE: return self.decode_param_opaque_handle( current_block, param, args, arg_index, arg_context @@ -1876,20 +1944,32 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): else: raise ValueError(f"Unsupported parameter type: {param.dtype.type_code}") elif isinstance(param, spec.Shape): - return self.decode_param_shape(current_block, param, args, arg_index, arg_context) + return self.decode_param_shape( + current_block, param, args, arg_index, arg_context + ) elif isinstance(param, spec.Tensor): - return self.decode_param_tensor(current_block, param, args, arg_index, arg_context) + return self.decode_param_tensor( + current_block, param, args, arg_index, arg_context + ) elif isinstance(param, spec.Stream): - return self.decode_param_stream(current_block, param, args, arg_index, arg_context) + return self.decode_param_stream( + current_block, param, args, arg_index, arg_context + ) elif isinstance(param, spec.EnvStream): # decode of env stream is deferred after we go through all parameters return current_block elif isinstance(param, spec.DataPointer): - return self.decode_param_data_pointer(current_block, param, args, arg_index, arg_context) + return self.decode_param_data_pointer( + current_block, param, args, arg_index, arg_context + ) elif isinstance(param, spec.ConstNone): - return self.decode_param_const_none(current_block, param, args, arg_index, arg_context) + return self.decode_param_const_none( + current_block, param, args, arg_index, arg_context + ) elif isinstance(param, spec.TupleParam): - return self.decode_param_tuple(current_block, param, args, arg_index, arg_context) + return self.decode_param_tuple( + current_block, param, args, arg_index, arg_context + ) else: raise ValueError(f"Unsupported parameter type: {type(param)}") @@ -1993,7 +2073,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): arg_index=ffi_arg_index, tuple_indices=[], ) - current_block = self.decode_param(current_block, param, args, ffi_arg_index, arg_context) + current_block = self.decode_param( + current_block, param, args, ffi_arg_index, arg_context + ) ffi_arg_index += 1 with ir.InsertionPoint(current_block): @@ -2053,7 +2135,9 @@ def attach_ffi_func( def rename_tvm_ffi_function( - module: ir.Module, old_name: str, new_name: str, + module: ir.Module, + old_name: str, + new_name: str, ) -> None: """Rename the TVM FFI function in the module. diff --git a/python/CuTeDSL/cutlass/base_dsl/typing.py b/python/CuTeDSL/cutlass/base_dsl/typing.py index 507b67ff..89fc6fb0 100644 --- a/python/CuTeDSL/cutlass/base_dsl/typing.py +++ b/python/CuTeDSL/cutlass/base_dsl/typing.py @@ -10,6 +10,7 @@ # is strictly prohibited. import ctypes +from itertools import chain import numpy as np import operator from typing_extensions import deprecated @@ -229,7 +230,7 @@ def get_c_pointers(obj): if hasattr(obj, "__c_pointers__"): return obj.__c_pointers__() elif isinstance(obj, (tuple, list)): - return sum((get_c_pointers(x) for x in obj), []) + return list(chain.from_iterable(get_c_pointers(x) for x in obj)) elif isinstance(obj, set): raise DSLRuntimeError( "Sets are not supported in get_c_pointers to ensure order preservation", diff --git a/python/CuTeDSL/cutlass/base_dsl/utils/numpy.py b/python/CuTeDSL/cutlass/base_dsl/utils/numpy.py new file mode 100644 index 00000000..61a181e3 --- /dev/null +++ b/python/CuTeDSL/cutlass/base_dsl/utils/numpy.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +This module provides a NumPy related utility functions that are needed for +the DSL. +""" + +import numpy as np +from ..._mlir.extras import types as T + +# ============================================================================= +# Codegen Utils +# ============================================================================= + + +def _numpy_type_to_mlir_type(dtype): + if dtype == np.float64: + return T.f64() + if dtype == np.float16: + return T.f16() + if dtype == np.float32: + return T.f32() + if dtype == np.int64: + return T.i64() + if dtype == np.int32: + return T.i32() + if dtype == np.int16: + return T.i16() + if dtype == np.int8: + return T.i8() + if dtype == np.uint64: + return T.ui64() + if dtype == np.uint32: + return T.ui32() + if dtype == np.uint16: + return T.ui16() + if dtype == np.uint8: + return T.ui8() + if dtype == np.bool_: + return T.bool() + if dtype == f8E5M2: + return T.f8E5M2() + if dtype == f8E4M3FN: + return T.f8E4M3FN() + if dtype == f8E8M0FNU: + return T.f8E8M0FNU() + if dtype == f6E3M2FN: + return T.f6E3M2FN() + if dtype == f6E2M3FN: + return T.f6E2M3FN() + if dtype == f4E2M1FN: + return T.f4E2M1FN() + raise TypeError(f"Unknown NumPy dtype for MLIR conversion: {dtype!r}") + + +def _mlir_type_to_numpy_type(mlir_type): + if mlir_type == T.f64(): + return np.float64 + if mlir_type == T.f16(): + return np.float16 + if mlir_type == T.f32(): + return np.float32 + if mlir_type == T.i64(): + return np.int64 + if mlir_type == T.i32(): + return np.int32 + if mlir_type == T.i16(): + return np.int16 + if mlir_type == T.i8(): + return np.int8 + if mlir_type == T.ui64(): + return np.uint64 + if mlir_type == T.ui32(): + return np.uint32 + if mlir_type == T.ui16(): + return np.uint16 + if mlir_type == T.ui8(): + return np.uint8 + if mlir_type == T.bool(): + return np.bool_ + raise TypeError(f"Unknown MLIR type for NumPy conversion: {mlir_type!r}") diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tree_utils.py b/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py similarity index 86% rename from python/CuTeDSL/cutlass/cutlass_dsl/tree_utils.py rename to python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py index da1a6b5b..e447f353 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/tree_utils.py +++ b/python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py @@ -14,15 +14,15 @@ import dataclasses import itertools as it from types import SimpleNamespace -from ..base_dsl.typing import as_numeric, Numeric, Constexpr -from ..base_dsl._mlir_helpers.arith import ArithValue -from ..base_dsl.common import DSLBaseError -from .._mlir import ir +from ..typing import as_numeric, Numeric, Constexpr +from .._mlir_helpers.arith import ArithValue +from ..common import DSLBaseError +from ..._mlir import ir NoneType = type(None) # ============================================================================= -# Tree Utils +# Tree Utilsß # ============================================================================= @@ -43,6 +43,18 @@ def unzip2(pairs: Iterable[tuple[Any, Any]]) -> tuple[list[Any], list[Any]]: return lst1, lst2 +def unzip3( + triples: Iterable[tuple[Any, Any, Any]], +) -> tuple[list[Any], list[Any], list[Any]]: + """Unzip a sequence of triples into three lists.""" + lst1, lst2, lst3 = [], [], [] + for x1, x2, x3 in triples: + lst1.append(x1) + lst2.append(x2) + lst3.append(x3) + return lst1, lst2, lst3 + + def get_fully_qualified_class_name(x: Any) -> str: """ Get the fully qualified class name of an object. @@ -251,6 +263,7 @@ def set_dataclass_attributes( instance: Any, fields: list[str], values: Iterable[Any], + is_frozen: bool, constexpr_fields: list[str], ) -> Any: """ @@ -268,10 +281,16 @@ def set_dataclass_attributes( if not fields: return instance - kwargs = dict(zip(fields, values)) - for field in constexpr_fields: - kwargs[field] = getattr(instance, field) - return dataclasses.replace(instance, **kwargs) + if is_frozen: + kwargs = dict(zip(fields, values)) + for field in constexpr_fields: + kwargs[field] = getattr(instance, field) + return dataclasses.replace(instance, **kwargs) + else: + for field, value in zip(fields, values): + setattr(instance, field, value) + + return instance def default_dataclass_from_iterable( @@ -290,12 +309,11 @@ def default_dataclass_from_iterable( The reconstructed dataclass instance """ instance = metadata.original_obj + is_frozen = is_frozen_dataclass(instance) - new_instance = set_dataclass_attributes( - instance, metadata.fields, children, metadata.constexpr_fields + return set_dataclass_attributes( + instance, metadata.fields, children, is_frozen, metadata.constexpr_fields ) - metadata.original_obj = new_instance - return new_instance def dynamic_expression_to_iterable(x: Any) -> tuple[SimpleNamespace, list[Any]]: @@ -490,7 +508,7 @@ unflattened_a should be structurally identical to a, and unflattened_b should be """ -def tree_flatten(x: Any) -> tuple[list[Any], PyTreeDef]: +def tree_flatten(x: Any) -> tuple[list[Any], list[ir.Attribute], PyTreeDef]: """ Flatten a nested structure into a flat list of values and a tree definition. @@ -512,8 +530,8 @@ def tree_flatten(x: Any) -> tuple[list[Any], PyTreeDef]: >>> tree_flatten([1, [2, 3], 4]) ([1, 2, 3, 4], PyTreeDef(...)) """ - children_iter, treedef = _tree_flatten(x) - return list(children_iter), treedef + children_iter, child_attrs_iter, treedef = _tree_flatten(x) + return list(children_iter), list[ir.Attribute](child_attrs_iter), treedef def get_registered_node_types_or_insert(x: Any) -> Union[NodeType, None]: @@ -576,7 +594,9 @@ def create_leaf_for_value( ) -def _tree_flatten(x: Any) -> tuple[Iterable[Any], Union[PyTreeDef, Leaf]]: +def _tree_flatten( + x: Any, +) -> tuple[Iterable[Any], Iterable[ir.Attribute], Union[PyTreeDef, Leaf]]: """ Internal function to flatten a tree structure. @@ -595,28 +615,46 @@ def _tree_flatten(x: Any) -> tuple[Iterable[Any], Union[PyTreeDef, Leaf]]: DSLTreeFlattenError: If the object type is not supported """ if x is None: - return [], create_leaf_for_value(x, is_none=True) + return [], [], create_leaf_for_value(x, is_none=True) elif isinstance(x, ArithValue) and is_dynamic_expression(x): v = x.__extract_mlir_values__() - return v, create_leaf_for_value( - x, - node_metadata=SimpleNamespace(is_dynamic_expression=1, original_obj=x), - ir_type_str=str(v[0].type), + a = ( + [ir.DictAttr.get({})] + if not hasattr(x, "__extract_mlir_attributes__") + else x.__extract_mlir_attributes__() + ) + return ( + v, + a, + create_leaf_for_value( + x, + node_metadata=SimpleNamespace(is_dynamic_expression=1, original_obj=x), + ir_type_str=str(v[0].type), + ), ) elif isinstance(x, ArithValue): - return [x], create_leaf_for_value(x, is_numeric=True) + return [x], [ir.DictAttr.get({})], create_leaf_for_value(x, is_numeric=True) elif isinstance(x, ir.Value): - return [x], create_leaf_for_value(x) + return [x], [ir.DictAttr.get({})], create_leaf_for_value(x) elif isinstance(x, Numeric): v = x.__extract_mlir_values__() - return v, create_leaf_for_value( - x, - node_metadata=SimpleNamespace(is_dynamic_expression=1, original_obj=x), - ir_type_str=str(v[0].type), + a = ( + [ir.DictAttr.get({})] + if not hasattr(x, "__extract_mlir_attributes__") + else x.__extract_mlir_attributes__() + ) + return ( + v, + a, + create_leaf_for_value( + x, + node_metadata=SimpleNamespace(is_dynamic_expression=1, original_obj=x), + ir_type_str=str(v[0].type), + ), ) else: @@ -628,14 +666,31 @@ def _tree_flatten(x: Any) -> tuple[Iterable[Any], Union[PyTreeDef, Leaf]]: raise DSLTreeFlattenError( "Flatten Error: children is None", get_fully_qualified_class_name(x) ) - children_flat, child_trees = unzip2(map(_tree_flatten, children)) + children_flat, child_attrs_flat, child_trees = unzip3( + map(_tree_flatten, children) + ) flattened = it.chain.from_iterable(children_flat) - return flattened, PyTreeDef(node_type, node_metadata, tuple(child_trees)) + + if hasattr(x, "__extract_mlir_attributes__"): + # If x has extract mlir attributes it overrides the child's default attributes + child_attrs_flat = [x.__extract_mlir_attributes__()] * len( + child_attrs_flat + ) + child_attrs_flattened = it.chain.from_iterable(child_attrs_flat) + return ( + flattened, + child_attrs_flattened, + PyTreeDef(node_type, node_metadata, tuple(child_trees)), + ) # Try to convert to numeric try: nval = as_numeric(x).ir_value() - return [nval], create_leaf_for_value(nval, is_numeric=True) + return ( + [nval], + [ir.DictAttr.get({})], + create_leaf_for_value(nval, is_numeric=True), + ) except Exception: raise DSLTreeFlattenError( "Flatten Error", get_fully_qualified_class_name(x) @@ -658,7 +713,7 @@ def tree_unflatten(treedef: PyTreeDef, xs: list[Any]) -> Any: The reconstructed nested structure Example: - >>> flat_values, treedef = tree_flatten([1, [2, 3], 4]) + >>> flat_values, _, treedef = tree_flatten([1, [2, 3], 4]) >>> tree_unflatten(treedef, flat_values) [1, [2, 3], 4] """ @@ -709,6 +764,7 @@ def _check_tree_equal(lhs: Union[PyTreeDef, Leaf], rhs: Union[PyTreeDef, Leaf]) """ if isinstance(lhs, Leaf) and isinstance(rhs, Leaf): return lhs.is_none == rhs.is_none and lhs.ir_type_str == rhs.ir_type_str + elif isinstance(lhs, PyTreeDef) and isinstance(rhs, PyTreeDef): lhs_metadata = lhs.node_metadata rhs_metadata = rhs.node_metadata @@ -725,6 +781,7 @@ def _check_tree_equal(lhs: Union[PyTreeDef, Leaf], rhs: Union[PyTreeDef, Leaf]) and len(lhs.child_treedefs) == len(rhs.child_treedefs) and all(map(_check_tree_equal, lhs.child_treedefs, rhs.child_treedefs)) ) + else: return False @@ -744,8 +801,8 @@ def check_tree_equal(lhs: PyTreeDef, rhs: PyTreeDef) -> int: int: Index of the first differing child, or -1 if trees are equal Example: - >>> treedef1 = tree_flatten([1, [2, 3]])[1] - >>> treedef2 = tree_flatten([1, [2, 4]])[1] + >>> treedef1 = tree_flatten([1, [2, 3]])[2] + >>> treedef2 = tree_flatten([1, [2, 4]])[2] >>> check_tree_equal(treedef1, treedef2) 1 # The second child differs """ diff --git a/python/CuTeDSL/cutlass/cute/__init__.py b/python/CuTeDSL/cutlass/cute/__init__.py index 2be8c7f5..5522bd67 100644 --- a/python/CuTeDSL/cutlass/cute/__init__.py +++ b/python/CuTeDSL/cutlass/cute/__init__.py @@ -57,7 +57,6 @@ from .core import ( make_composed_layout, make_layout_tv, make_swizzle, - make_sparse_elem, recast_ptr, get, select, @@ -99,6 +98,8 @@ from .core import ( local_partition, local_tile, printf, + get_nonswizzle_portion, + get_swizzle_portion, # Wrapper classes Swizzle, E, @@ -118,6 +119,8 @@ from .core import ( # FastDivmod operations FastDivmodDivisor, fast_divmod_create_divisor, + basis_value, + basis_get, ) from .tuple import ( @@ -130,6 +133,9 @@ from .tuple import ( product_like, product_each, elem_less, + tuple_cat, + transform_apply, + filter_tuple, ) from .tensor import ( TensorSSA, @@ -178,6 +184,8 @@ from .atom import ( ) from .algorithm import gemm, copy, basic_copy, basic_copy_if, autovec_copy, prefetch +from . import typing as typing_module +from . import core from . import arch from . import export @@ -215,6 +223,7 @@ _tvm_ffi_args_spec_converter.attach_args_spec_converter(_dsl.CuTeDSL._get_dsl()) # Explicitly export all symbols for documentation generation __all__ = [ # Core types + *core.__all__, "AddressSpace", "CacheEvictionPriority", "Tensor", @@ -252,6 +261,8 @@ __all__ = [ "make_composed_layout", "make_layout_tv", "make_layout_image_mask", + "get_nonswizzle_portion", + "get_swizzle_portion", # Tensor functions "make_ptr", "make_tensor", @@ -271,6 +282,8 @@ __all__ = [ "find", "find_if", "transform_leaf", + "basis_value", + "basis_get", "coalesce", "group_modes", "cosize", @@ -287,6 +300,9 @@ __all__ = [ "prepend_ones", "append_ones", "elem_less", + "tuple_cat", + "transform_apply", + "filter_tuple", # Math operations "ceil_div", "round_up", diff --git a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py index e4c51781..8a5ebc9e 100644 --- a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py +++ b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py @@ -353,6 +353,8 @@ def _convert_single_arg( elem_param = _convert_single_arg(elem, elem_name, None, ctx) tuple_params.append(elem_param) return spec.TupleParam(arg_name, tuple_params) + elif isinstance(arg, bool): + return spec.Var(arg_name, NumericToTVMFFIDtype[Boolean]) elif isinstance(arg, int): # in cute.compile, unannotated const int is converted to int32 return spec.Var(arg_name, NumericToTVMFFIDtype[Int32]) diff --git a/python/CuTeDSL/cutlass/cute/arch/__init__.py b/python/CuTeDSL/cutlass/cute/arch/__init__.py index a69ad0bc..6d10c66e 100644 --- a/python/CuTeDSL/cutlass/cute/arch/__init__.py +++ b/python/CuTeDSL/cutlass/cute/arch/__init__.py @@ -16,6 +16,7 @@ from .nvvm_wrappers import * from .smem import * from .tmem import * from .numeric_conversion import * +from .clc import * # __all__ is required here for documentation generation __all__ = [ @@ -73,12 +74,24 @@ __all__ = [ "vote_any_sync", "vote_all_sync", "vote_uni_sync", + "atomic_add", + "atomic_and", + "atomic_or", + "atomic_xor", + "atomic_max", + "atomic_min", + "atomic_exch", + "atomic_cas", + "store", + "load", "popc", "fence_proxy", "fence_view_async_tmem_load", "fence_view_async_tmem_store", "warpgroup_reg_alloc", "warpgroup_reg_dealloc", + "setmaxregister_increase", + "setmaxregister_decrease", "fma_packed_f32x2", "mul_packed_f32x2", "add_packed_f32x2", @@ -100,6 +113,8 @@ __all__ = [ # # tmem.py # + "get_max_tmem_alloc_cols", + "get_min_tmem_alloc_cols", "retrieve_tmem_ptr", "alloc_tmem", "relinquish_tmem_alloc_permit", @@ -115,4 +130,9 @@ __all__ = [ "cvt_i8x2_to_f32x2", "cvt_i8_bf16", "cvt_f32x2_bf16x2", + # + # clc.py + # + "issue_clc_query", + "clc_response", ] diff --git a/python/CuTeDSL/cutlass/cute/arch/clc.py b/python/CuTeDSL/cutlass/cute/arch/clc.py new file mode 100644 index 00000000..29af7348 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/arch/clc.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Tuple +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import nvvm, vector + +from ..typing import Int32, Pointer, Int128 + + +@dsl_user_op +def issue_clc_query( + mbar_ptr: Pointer, + clc_response_ptr: Pointer, + loc=None, + ip=None, +) -> None: + """ + The clusterlaunchcontrol.try_cancel instruction requests atomically cancelling the launch + of a cluster that has not started running yet. It asynchronously writes an opaque response + to shared memory indicating whether the operation succeeded or failed. On success, the + opaque response contains the ctaid of the first CTA of the canceled cluster. + + :param mbar_ptr: A pointer to the mbarrier address in SMEM + :type mbar_ptr: Pointer + :param clc_response_ptr: A pointer to the cluster launch control response address in SMEM + :type clc_response_ptr: Pointer + """ + mbar_llvm_ptr = mbar_ptr.llvm_ptr + clc_response_llvm_ptr = clc_response_ptr.llvm_ptr + nvvm.clusterlaunchcontrol_try_cancel_multicast( + clc_response_llvm_ptr, + mbar_llvm_ptr, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def clc_response( + result_addr: Pointer, loc=None, ip=None +) -> Tuple[Int32, Int32, Int32, Int32]: + """ + After loading response from clusterlaunchcontrol.try_cancel instruction into 16-byte + register, it can be further queried using clusterlaunchcontrol.query_cancel instruction. + If the cluster is canceled successfully, predicate p is set to true; otherwise, it is + set to false. If the request succeeded, clusterlaunchcontrol.query_cancel.get_first_ctaid + extracts the CTA id of the first CTA in the canceled cluster. By default, the instruction + returns a .v4 vector whose first three elements are the x, y and z coordinate of first CTA + in canceled cluster. + + :param result_addr: A pointer to the cluster launch control response address in SMEM + :type result_addr: Pointer + """ + from cutlass.cute import recast_ptr, make_tensor, make_layout + + clc_ptr_i128 = recast_ptr(result_addr, dtype=Int128, loc=loc, ip=ip) + clc_tensor = make_tensor( + clc_ptr_i128, make_layout(1, loc=loc, ip=ip), loc=loc, ip=ip + ) + + # Load the 128-bit value from shared memory + clc_result_vec = clc_tensor.load(loc=loc, ip=ip) + + # Extract the i128 scalar from the vector<1xi128> + clc_result_i128 = vector.extract( + clc_result_vec.ir_value(loc=loc, ip=ip), + [], + [0], + ) + # Query if the cluster was canceled + pred = nvvm.clusterlaunchcontrol_query_cancel_is_canceled( + T.bool(), + clc_result_i128, + loc=loc, + ip=ip, + ) + is_valid = Int32(pred) + + # Get first CTA ID x component + m_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_x( + T.i32(), + clc_result_i128, + loc=loc, + ip=ip, + ) + + # Get first CTA ID y component + n_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_y( + T.i32(), + clc_result_i128, + loc=loc, + ip=ip, + ) + + # Get first CTA ID z component + l_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_z( + T.i32(), + clc_result_i128, + loc=loc, + ip=ip, + ) + + m_idx = Int32(m_idx_i32) + n_idx = Int32(n_idx_i32) + l_idx = Int32(l_idx_i32) + + return m_idx, n_idx, l_idx, is_valid diff --git a/python/CuTeDSL/cutlass/cute/arch/elect.py b/python/CuTeDSL/cutlass/cute/arch/elect.py index ebd6d840..9f51484e 100644 --- a/python/CuTeDSL/cutlass/cute/arch/elect.py +++ b/python/CuTeDSL/cutlass/cute/arch/elect.py @@ -69,6 +69,8 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion: # Only one thread in the warp executes the code in this context pass """ + from cutlass.base_dsl.arch import Arch + BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) is_thread_leader = nvvm.elect_sync(T.bool()) if_op = scf.IfOp(is_thread_leader, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/cute/arch/mbar.py b/python/CuTeDSL/cutlass/cute/arch/mbar.py index 469ef3d5..3f9f2ec5 100644 --- a/python/CuTeDSL/cutlass/cute/arch/mbar.py +++ b/python/CuTeDSL/cutlass/cute/arch/mbar.py @@ -35,7 +35,10 @@ def mbarrier_init(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None: :type cnt: Int """ nvvm.mbarrier_init_shared( - mbar_ptr.llvm_ptr, Int32(cnt).ir_value(loc=loc, ip=ip), loc=loc, ip=ip + mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), + Int32(cnt).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, ) @@ -65,7 +68,7 @@ def mbarrier_arrive_and_expect_tx( """ BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = mbar_ptr.llvm_ptr + mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) if peer_cta_rank_in_cluster is not None: mbar_llvm_ptr = nvvm.mapa_shared_cluster( mbar_llvm_ptr.type, @@ -105,7 +108,7 @@ def mbarrier_expect_tx( """ BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = mbar_ptr.llvm_ptr + mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) if peer_cta_rank_in_cluster is not None: mbar_llvm_ptr = nvvm.mapa( mbar_llvm_ptr.type, @@ -144,7 +147,7 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None: # This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX # The timeout in ns only applies to the latter and this call is truly blocking nvvm.mbarrier_try_wait_parity_shared( - mbar_ptr.llvm_ptr, + mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), Int32(phase).ir_value(loc=loc, ip=ip), Int32(timeout_ns).ir_value(loc=loc, ip=ip), loc=loc, @@ -169,7 +172,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo return Boolean( nvvm.mbarrier_wait_parity( T.bool(), - mbar_ptr.llvm_ptr, + mbar_ptr.to_llvm_ptr(loc=loc, ip=ip), Int32(phase).ir_value(loc=loc, ip=ip), nvvm.MBarrierWaitKind.TRY, loc=loc, @@ -223,7 +226,7 @@ def mbarrier_arrive( the mbarrier is converted to a remote address in the peer CTA's SMEM. """ - mbar_llvm_ptr = mbar_ptr.llvm_ptr + mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) if peer_cta_rank_in_cluster is not None: BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) @@ -261,10 +264,5 @@ def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> N """ BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90) - mbar_llvm_ptr = mbar_ptr.llvm_ptr - nvvm.cp_async_mbarrier_arrive_shared( - mbar_llvm_ptr, - noinc=True, - loc=loc, - ip=ip, - ) + mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip) + nvvm.cp_async_mbarrier_arrive_shared(mbar_llvm_ptr, noinc=True, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py index 073cf6d9..6484136f 100644 --- a/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py +++ b/python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py @@ -30,6 +30,7 @@ from .nvvm_wrappers import ( cvt_f4e2m1x4_to_f16x4, cvt_f4e2m1x2_to_f16x2, cvt_f4e2m1_f16, + sext_unpacked_i4x4_to_i8x4, ) from ..typing import ( @@ -272,6 +273,38 @@ def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None): return vec_dst +@dsl_user_op +def sext_unpacked_i4_i8_intrinsic(vec_unpacked_i4, length, *, loc=None, ip=None): + """ + Sign extend vector of int4 unpacked in 8b containers to packed int8 + + :param vec_unpacked_i4: The input vector of unpacked int4. + :type vec_unpacked_i4: 1D vector of unpacked int4 + :param length: The length of the input vector. + :type length: int + :return: The output 1D vector of int8 with the same length as the input vector. + :rtype: 1D vector of int8 + """ + assert length % 4 == 0, "unsupported length" + + vec_i8x4_type = ir.VectorType.get([4], Int8.mlir_type, loc=loc) + vec_i8_type = ir.VectorType.get([length], Int8.mlir_type, loc=loc) + vec_i8 = llvm.mlir_zero(vec_i8_type, loc=loc, ip=ip) + + for pos in range(0, length, 4): + vec_unpacked_i4x4 = vector.extract_strided_slice( + vec_i8x4_type, vec_unpacked_i4, [pos], [4], [1], loc=loc, ip=ip + ) + vec_i8x4 = sext_unpacked_i4x4_to_i8x4( + vec_unpacked_i4x4, loc=loc, ip=ip + ) + vec_i8 = vector.insert_strided_slice( + vec_i8x4, vec_i8, [pos], [1], loc=loc, ip=ip + ) + + return vec_i8 + + # Expose supported architectures via the intrinsic symbol cvt_i8_bf16_intrinsic.supported_archs = ( *Arch.AmpereArchs(), diff --git a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py index bb7e006b..78f68639 100644 --- a/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py +++ b/python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py @@ -10,19 +10,16 @@ # is strictly prohibited. from functools import partial -from typing import Optional, Tuple, Union, Callable, TYPE_CHECKING +from typing import Optional, Tuple, Union, Callable, Literal from typing_extensions import deprecated -from cutlass.cutlass_dsl import T, dsl_user_op, cutlass_arith +from cutlass.cutlass_dsl import T, dsl_user_op import cutlass.cutlass_dsl as cutlass_dsl from cutlass._mlir import ir from cutlass._mlir.dialects import arith, llvm, nvvm, vector -if TYPE_CHECKING: - from cutlass.tensor import TensorSSA - # Forward nvvm enums from cutlass._mlir.dialects.nvvm import ( ProxyKind, @@ -54,6 +51,81 @@ WARP_SIZE = 32 FULL_MASK = 0xFFFFFFFF +# ============================================================================ +# Enum String Mapping Helper +# ============================================================================ +# This section provides a helper to convert string literals to NVVM enum types +# by introspecting the enum's __str__() method. Each function imports and +# enhances only the enums it needs, avoiding namespace pollution. +# +# Usage within functions: +# MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind) +# sem = MemOrderKind.from_str("relaxed") +## ============================================================================ + + +def _enhance_enum_with_str_mapping(enum_class): + """ + Enhance an IntEnum class with automatic string-to-enum conversion. + + Builds a reverse mapping from __str__() output to enum members and adds + a from_str() class method for conversion. Safe to call multiple times + (idempotent - won't re-enhance if already enhanced). + + :param enum_class: The enum class to enhance + :return: The enhanced enum class (for chaining) + """ + # Skip if already enhanced + if hasattr(enum_class, "from_str"): + return enum_class + + # Build reverse mapping from string representation to enum member + str_to_enum_map = {} + for member in enum_class: + str_repr = str(member) + if str_repr in str_to_enum_map: + raise ValueError( + f"Duplicate string representation '{str_repr}' in {enum_class.__name__}" + ) + str_to_enum_map[str_repr] = member + + # Add from_str class method + @classmethod + def from_str(cls, s): + """ + Convert a string literal to the corresponding enum member. + + :param s: String representation of the enum member, or an enum member itself (deprecated) + :return: The enum member (or None if s is None) + :raises ValueError: If the string is not a valid enum member + """ + import warnings + + if s is None: + return None + + # Check if s is already an enum member of the correct type + if isinstance(s, cls): + warnings.warn( + f"Passing enum member directly to {cls.__name__}.from_str() is deprecated. " + f"Please use string literals instead (e.g., '{str(s)}' instead of {cls.__name__}.{s.name}).", + DeprecationWarning, + stacklevel=2, + ) + return s + + if s not in str_to_enum_map: + valid_options = sorted(str_to_enum_map.keys()) + raise ValueError( + f"Invalid {cls.__name__} string: '{s}'. " + f"Valid options are: {valid_options}" + ) + return str_to_enum_map[s] + + enum_class.from_str = from_str + return enum_class + + @dsl_user_op def lane_idx(*, loc=None, ip=None) -> Int32: """ @@ -374,7 +446,6 @@ def warp_reduction( offset = offset // 2 return val - warp_reduction_max = partial( warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else cutlass_dsl.max(x, y), @@ -389,13 +460,34 @@ def barrier(*, barrier_id=None, number_of_threads=None, loc=None, ip=None) -> No """ if barrier_id is not None: barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip) + else: + barrier_id = Int32(0).ir_value(loc=loc, ip=ip) if number_of_threads is not None: number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip) - - nvvm.barrier( - barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip - ) + llvm.inline_asm( + None, + [barrier_id, number_of_threads], + "bar.sync $0, $1;", + "r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + else: + llvm.inline_asm( + None, + [barrier_id], + "bar.sync $0;", + "r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) @dsl_user_op @@ -404,6 +496,8 @@ def barrier_arrive( ) -> None: if barrier_id is not None: barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip) + else: + barrier_id = Int32(0).ir_value(loc=loc, ip=ip) if number_of_threads is None: raise ValueError( @@ -411,8 +505,14 @@ def barrier_arrive( ) number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip) - nvvm.barrier_arrive( - barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip + llvm.inline_asm( + None, + [barrier_id, number_of_threads], + "bar.arrive $0, $1;", + "r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, ) @@ -682,7 +782,7 @@ def popc(value: Numeric, *, loc=None, ip=None) -> Numeric: @dsl_user_op def fence_view_async_tmem_op( - kind: Tcgen05WaitKind, + kind: Literal["load", "store"], *, loc=None, ip=None, @@ -715,18 +815,20 @@ def fence_view_async_tmem_op( ``` - :param kind: The kind of fence operation to perform including LOAD and STORE. - :type kind: Tcgen05WaitKind + :param kind: The kind of fence operation to perform ("load", "store"). + :type kind: Literal["load", "store"] """ - nvvm.tcgen05_wait(kind, loc=loc, ip=ip) + from cutlass._mlir.dialects.nvvm import Tcgen05WaitKind + + # Enhance enum and convert string literal to enum type + Tcgen05WaitKind_enhanced = _enhance_enum_with_str_mapping(Tcgen05WaitKind) + kind = Tcgen05WaitKind_enhanced.from_str(kind) + + nvvm.tcgen05_wait(kind=kind, loc=loc, ip=ip) -fence_view_async_tmem_load = partial( - fence_view_async_tmem_op, kind=Tcgen05WaitKind.LOAD -) -fence_view_async_tmem_store = partial( - fence_view_async_tmem_op, kind=Tcgen05WaitKind.STORE -) +fence_view_async_tmem_load = partial(fence_view_async_tmem_op, kind="load") +fence_view_async_tmem_store = partial(fence_view_async_tmem_op, kind="store") @dsl_user_op @@ -751,23 +853,45 @@ def fence_view_async_shared( @dsl_user_op -def warpgroup_reg_realloc_op( +def setmaxregister_increase( + reg_count: int, + *, + loc=None, + ip=None, +): + return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip) + + +@dsl_user_op +def setmaxregister_decrease( + reg_count: int, + *, + loc=None, + ip=None, +): + return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip) + + +@dsl_user_op +@deprecated("API is deprecated, use setmaxregister_increase instead") +def warpgroup_reg_alloc( reg_count: int, - kind: SetMaxRegisterAction, *, loc=None, ip=None, ) -> None: - nvvm.setmaxregister(reg_count, kind, loc=loc, ip=ip) + nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip) -warpgroup_reg_alloc = partial( - warpgroup_reg_realloc_op, kind=SetMaxRegisterAction.increase -) -warpgroup_reg_dealloc = partial( - warpgroup_reg_realloc_op, kind=SetMaxRegisterAction.decrease -) - +@dsl_user_op +@deprecated("API is deprecated, use setmaxregister_decrease instead") +def warpgroup_reg_dealloc( + reg_count: int, + *, + loc=None, + ip=None, +) -> None: + nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip) @dsl_user_op def calc_packed_f32x2_op( @@ -776,11 +900,17 @@ def calc_packed_f32x2_op( src_c: Optional[Tuple[Float32, Float32]], calc_func: Callable, *, - rnd=RoundingModeKind.RZ, - ftz=True, + rnd: Optional[Literal["rn", "rz", "rm", "rp", "none"]] = "rn", + ftz=None, loc=None, ip=None, ) -> Tuple[Float32, Float32]: + from cutlass._mlir.dialects.nvvm import RoundingModeKind + + # Enhance enum and convert string literal to enum type + RoundingModeKind_enhanced = _enhance_enum_with_str_mapping(RoundingModeKind) + rnd = RoundingModeKind_enhanced.from_str(rnd) + vec_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc) vec_src_a = vector.from_elements( vec_type, @@ -1259,6 +1389,15 @@ def cvt_i4x8_to_bf16x8(src_vec8, *, loc=None, ip=None): vec_bf16x8 = llvm.bitcast(vec_bf16x8_type, rst_i32, loc=loc, ip=ip) return vec_bf16x8 +# Sign extend 4 int4 unpacked in 8b containers +@dsl_user_op +def sext_unpacked_i4x4_to_i8x4(src_vec4, *, loc=None, ip=None): + imm_u32 = arith.constant(Uint32.mlir_type, 0x78787878, loc=loc, ip=ip) + src_u32 = llvm.bitcast(Uint32.mlir_type, src_vec4, loc=loc, ip=ip) + dst_u32 = arith.addi(src_u32, imm_u32, loc=loc, ip=ip) + dst_u32 = arith.xori(dst_u32, imm_u32, loc=loc, ip=ip) + return llvm.bitcast(src_vec4.type, dst_u32, loc=loc, ip=ip) + @dsl_user_op def log2_of_pow2_int(a: Int32, *, loc=None, ip=None) -> Int32: @@ -1346,6 +1485,621 @@ def griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None: +def _normalize_ptr(addr, *, loc=None, ip=None) -> ir.Value: + """ + Helper function to normalize pointer types to MLIR ir.Value. + + Supports: + - ir.Value (LLVM pointer): returned as-is + - cute.ptr (_Pointer instance): converted via to_llvm_ptr() + + :param addr: Address in various pointer formats + :return: Normalized MLIR pointer value + :rtype: ir.Value + """ + # If it's already an MLIR ir.Value, return as-is + if isinstance(addr, ir.Value): + return addr + + # If it has to_llvm_ptr method (cute._Pointer instances) + if hasattr(addr, "to_llvm_ptr") and callable(addr.to_llvm_ptr): + return addr.to_llvm_ptr(loc=loc, ip=ip) + + # If none of the above, return as-is and let NVVM handle it + # This allows for future pointer types without breaking existing code + return addr + + +def _atomic( + ptr, + val: Union[Numeric, ir.Value], + *, + op: Literal[ + "add", + "fadd", + "max", + "min", + "and", + "or", + "xor", + "exch", + ], + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Union[Numeric, ir.Value]: + """ + General atomic operation function. + + Atomically adds `val` to the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location. Supports: + - ir.Value (LLVM pointer) + - cute.ptr (_Pointer instance) + :param val: Value to add (scalar Numeric or vector ir.Value) + :type val: Union[Numeric, ir.Value] + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :param op: Atomic operation ("add", "fadd", "max", "min", "and", "or", "xor", "exch") + :type op: Literal["add", "fadd", "max", "min", "and", "or", "xor", "exch"] + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Union[Numeric, ir.Value] + """ + from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind + from cutlass.utils.version_info import CUDA_VERSION + + # Enhance enums and convert string literals to enum types + AtomicOpKind = _enhance_enum_with_str_mapping(AtomicOpKind) + MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind) + MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind) + + op = AtomicOpKind.from_str(op) + sem = MemOrderKind.from_str(sem) + scope = MemScopeKind.from_str(scope) + + # Normalize pointer type to MLIR ir.Value + ptr = _normalize_ptr(ptr, loc=loc, ip=ip) + + # * Handle `val` Type - scalar Numeric or vector ir.Value + is_vector = isinstance(val, ir.Value) and isinstance(val.type, ir.VectorType) + + if is_vector: + # Vector type atomic - val is already an ir.Value + val_ir = val + val_type = val.type + # Check if it's a floating-point vector type + elem_type = val.type.element_type + is_float_vector = ( + elem_type == Float16.mlir_type + or elem_type == BFloat16.mlir_type + or elem_type == Float32.mlir_type + ) + + # Vector atomics for f16/bf16/f32 only support ADD (FADD) + if is_float_vector and op == AtomicOpKind.ADD: + op = AtomicOpKind.FADD + else: + # Scalar type atomic - convert to Numeric + if not isinstance(val, Numeric): + val = as_numeric(val) + val_type = type(val) + val_ir = val.ir_value(loc=loc, ip=ip) + + # * Float + # For .f32, .f64, .f16, .bf16, .f16x2, .bf16x2, only .add (FADD) is supported + # For .u32 .u64, .s32, .s64, .add .and .or .xor .cas .exch .min .max are supported + if val_type.is_float: + # For floating-point types, only ADD is supported + if op == AtomicOpKind.ADD: + # Convert ADD to FADD for floating-point types + op = AtomicOpKind.FADD + + # * NVVM call based on nvvm version + if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: + # Old API: requires explicit result type as first positional argument + # For vectors: pass val_type (ir.VectorType), for scalars: pass val_type.mlir_type + result_type = val_type if is_vector else val_type.mlir_type + result = nvvm.atomicrmw( + result_type, + op=op, + ptr=ptr, + a=val_ir, + mem_order=sem, + syncscope=scope, + loc=loc, + ip=ip, + ) + else: + # New API: infers result type automatically + result = nvvm.atomicrmw( + op=op, + ptr=ptr, + a=val_ir, + mem_order=sem, + syncscope=scope, + loc=loc, + ip=ip, + ) + # Return raw result for vectors, wrapped for scalars + return result if is_vector else val_type(result) + + +def atomic_add( + ptr, + val: Union[Numeric, ir.Value], + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Union[Numeric, ir.Value]: + """ + Performs an atomic addition operation. + + Atomically adds `val` to the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value to add (scalar Numeric or vector ir.Value) + :type val: Union[Numeric, ir.Value] + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Union[Numeric, ir.Value] + """ + return _atomic(ptr, val, op="add", sem=sem, scope=scope, loc=loc, ip=ip) + + +def atomic_and( + ptr, + val: Numeric, + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic bitwise AND operation. + + Atomically computes bitwise AND of `val` with the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value for AND operation + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + return _atomic(ptr, val, op="and", sem=sem, scope=scope, loc=loc, ip=ip) + + +def atomic_or( + ptr, + val: Numeric, + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic bitwise OR operation. + + Atomically computes bitwise OR of `val` with the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value for OR operation + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + return _atomic(ptr, val, op="or", sem=sem, scope=scope, loc=loc, ip=ip) + + +def atomic_xor( + ptr, + val: Numeric, + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic bitwise XOR operation. + + Atomically computes bitwise XOR of `val` with the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value for XOR operation + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + return _atomic(ptr, val, op="xor", sem=sem, scope=scope, loc=loc, ip=ip) + + +def atomic_max( + ptr, + val: Numeric, + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic maximum operation. + + Atomically computes maximum of `val` and the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value for MAX operation + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + return _atomic(ptr, val, op="max", sem=sem, scope=scope, loc=loc, ip=ip) + + +def atomic_min( + ptr, + val: Numeric, + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic minimum operation. + + Atomically computes minimum of `val` and the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value for MIN operation + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + return _atomic(ptr, val, op="min", sem=sem, scope=scope, loc=loc, ip=ip) + + +def atomic_exch( + ptr, + val: Numeric, + *, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic exchange operation. + + Atomically exchanges `val` with the value at memory location `ptr` and returns the old value. + + :param ptr: Pointer to memory location + :param val: Value to exchange + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + return _atomic(ptr, val, op="exch", sem=sem, scope=scope, loc=loc, ip=ip) + + +@dsl_user_op +def atomic_cas( + ptr, + *, + cmp: Numeric, + val: Numeric, + sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> Numeric: + """ + Performs an atomic compare-and-swap (CAS) operation. + + Atomically compares the value at the memory location with `cmp`. If they are equal, + stores `val` at the memory location and returns the old value. + + :param ptr: Pointer to memory location. Supports: + - ir.Value (LLVM pointer) + - cute.ptr (_Pointer instance) + :param cmp: Value to compare against current memory value + :type cmp: Numeric + :param val: Value to store if comparison succeeds + :type val: Numeric + :param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel") + :type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] + :param scope: Memory scope ("gpu", "cta", "cluster", "sys") + :type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] + :return: Old value at memory location + :rtype: Numeric + """ + from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind + from cutlass.utils.version_info import CUDA_VERSION + + # Enhance enums and convert string literals to enum types + MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind) + MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind) + + sem = MemOrderKind.from_str(sem) + scope = MemScopeKind.from_str(scope) + + # Normalize pointer type to MLIR ir.Value + ptr = _normalize_ptr(ptr, loc=loc, ip=ip) + + # * Hanldle `val`, `cmp` Numeric Type + if not isinstance(cmp, Numeric): + cmp = as_numeric(cmp) + if not isinstance(val, Numeric): + val = as_numeric(val) + cmp_type = type(cmp) + cmp_ir = cmp.ir_value(loc=loc, ip=ip) + val_ir = val.ir_value(loc=loc, ip=ip) + + # * NVVM call based on nvvm version + if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9: + result = nvvm.atomicrmw( + cmp_type.mlir_type, + op=AtomicOpKind.CAS, + ptr=ptr, + a=val_ir, + b=cmp_ir, + mem_order=sem, + syncscope=scope, + loc=loc, + ip=ip, + ) + else: + result = nvvm.atomicrmw( + op=AtomicOpKind.CAS, + ptr=ptr, + a=cmp_ir, + b=val_ir, + mem_order=sem, + syncscope=scope, + loc=loc, + ip=ip, + ) + return cmp_type(result) + + +@dsl_user_op +def store( + ptr, + val: Union[Numeric, ir.Value], + *, + level1_eviction_priority: Optional[ + Literal[ + "evict_normal", + "evict_first", + "evict_last", + "evict_no_allocate", + "evict_unchanged", + ] + ] = None, + cop: Optional[Literal["wb", "cg", "cs", "wt"]] = None, + ss: Optional[Literal["cta", "cluster"]] = None, + sem: Optional[Literal["relaxed", "release"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + loc=None, + ip=None, +) -> None: + """ + Store a value to a memory location. + + :param ptr: Pointer to store to. Supports: + - ir.Value (LLVM pointer) + - cute.ptr (_Pointer instance) + :param val: Value to store (scalar Numeric or vector ir.Value) + :type val: Union[Numeric, ir.Value] + :param level1_eviction_priority: L1 cache eviction policy string literal: + "evict_normal" : .level1::eviction_priority = .L1::evict_normal + "evict_first" : .level1::eviction_priority = .L1::evict_first + "evict_last" : .level1::eviction_priority = .L1::evict_last + "evict_no_allocate" : .level1::eviction_priority = .L1::no_allocate + "evict_unchanged" : .level1::eviction_priority = .L1::evict_unchanged + :param cop: Store cache modifier string literal: + :param ss: Shared memory space string literal: + "cta" : .ss = .shared::cta + "cluster" : .ss = .shared::cluster + None : .ss = .global + :param sem: Memory semantic string literal: + :param scope: Memory scope string literal: + + """ + from cutlass._mlir.dialects.nvvm import ( + MemOrderKind, + MemScopeKind, + StoreCacheModifierKind, + EvictKind, + SharedSpace, + ) + + # Enhance enums and convert string literals to enum types + MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind) + MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind) + StoreCacheModifierKind = _enhance_enum_with_str_mapping(StoreCacheModifierKind) + EvictKind = _enhance_enum_with_str_mapping(EvictKind) + SharedSpace = _enhance_enum_with_str_mapping(SharedSpace) + + sem = MemOrderKind.from_str(sem) + scope = MemScopeKind.from_str(scope) + cop = StoreCacheModifierKind.from_str(cop) + level1_eviction_priority = EvictKind.from_str(level1_eviction_priority) + ss = SharedSpace.from_str(ss) + + # Normalize pointer type to MLIR ir.Value + ptr = _normalize_ptr(ptr, loc=loc, ip=ip) + + # Handle both scalar Numeric and vector ir.Value + is_vector = isinstance(val, ir.Value) and isinstance(val.type, ir.VectorType) + + if is_vector: + # Vector type store - val is already an ir.Value + val_ir = val + else: + # Scalar type store - ensure val is a Numeric and convert to MLIR Value + if not isinstance(val, Numeric): + val = as_numeric(val) + val_ir = val.ir_value(loc=loc, ip=ip) + + nvvm.store_ext( + val_ir, + ptr, + order=sem, + scope=scope, + evict=level1_eviction_priority, + cache_modifier=cop, + shared_space=ss, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def load( + ptr, + dtype: Union[type[Numeric], ir.VectorType], + *, + sem: Optional[Literal["relaxed", "acquire"]] = None, + scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None, + level1_eviction_priority: Optional[ + Literal[ + "evict_normal", + "evict_first", + "evict_last", + "evict_no_allocate", + "evict_unchanged", + ] + ] = None, + cop: Optional[Literal["ca", "cg", "cs", "lu", "cv"]] = None, + ss: Optional[Literal["cta", "cluster"]] = None, + level_prefetch_size: Optional[Literal["size_64b", "size_128b", "size_256b"]] = None, + loc=None, + ip=None, +) -> Union[Numeric, ir.Value]: + """ + Load a value from a memory location. + + :param ptr: Pointer to load from. Supports: + - ir.Value (LLVM pointer) + - cute.ptr (_Pointer instance) + :param dtype: Data type to load. Can be: + - Scalar: Numeric type class (Int8, Uint8, Int32, Float32, etc.) + - Vector: ir.VectorType for vectorized load (e.g., ir.VectorType.get([4], Int64.mlir_type)) + :type dtype: Union[type[Numeric], ir.VectorType] + :param sem: Memory semantic string literal: + :param scope: Memory scope string literal: + :param level1_eviction_priority: L1 cache eviction policy string literal: + "evict_normal" : .level1::eviction_priority = .L1::evict_normal + "evict_first" : .level1::eviction_priority = .L1::evict_first + "evict_last" : .level1::eviction_priority = .L1::evict_last + "evict_no_allocate" : .level1::eviction_priority = .L1::no_allocate + "evict_unchanged" : .level1::eviction_priority = .L1::evict_unchanged + :param cop: Load cache modifier string literal: + :param ss: Shared memory space string literal: + "cta" : .ss = .shared::cta + "cluster" : .ss = .shared::cluster + None : .ss = .global + :param level_prefetch_size: L2 cache prefetch size hint string literal: + "size_64b" : .level::prefetch_size = .L2::64B + "size_128b" : .level::prefetch_size = .L2::128B + "size_256b" : .level::prefetch_size = .L2::256B + :return: Loaded value (scalar Numeric or vector ir.Value) + :rtype: Union[Numeric, ir.Value] + """ + from cutlass._mlir.dialects.nvvm import ( + MemOrderKind, + MemScopeKind, + LoadCacheModifierKind, + EvictKind, + SharedSpace, + L2PrefetchSize, + ) + + # Enhance enums and convert string literals to enum types + MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind) + MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind) + LoadCacheModifierKind = _enhance_enum_with_str_mapping(LoadCacheModifierKind) + EvictKind = _enhance_enum_with_str_mapping(EvictKind) + SharedSpace = _enhance_enum_with_str_mapping(SharedSpace) + L2PrefetchSize = _enhance_enum_with_str_mapping(L2PrefetchSize) + + sem = MemOrderKind.from_str(sem) + scope = MemScopeKind.from_str(scope) + cop = LoadCacheModifierKind.from_str(cop) + level1_eviction_priority = EvictKind.from_str(level1_eviction_priority) + ss = SharedSpace.from_str(ss) + level_prefetch_size = L2PrefetchSize.from_str(level_prefetch_size) + + # Normalize pointer type to MLIR ir.Value + ptr = _normalize_ptr(ptr, loc=loc, ip=ip) + + # Determine if dtype is a vector type or scalar type + is_vector = isinstance(dtype, ir.VectorType) and isinstance(dtype, ir.VectorType) + + if is_vector: + # Vector load: dtype is already an ir.VectorType + mlir_type = dtype + scalar_dtype = None # We don't need to wrap the result + else: + # Scalar load: dtype is a Numeric type class + mlir_type = dtype.mlir_type + scalar_dtype = dtype + + result = nvvm.load_ext( + res=mlir_type, + addr=ptr, + order=sem, + scope=scope, + evict=level1_eviction_priority, + cache_modifier=cop, + shared_space=ss, + prefetch=level_prefetch_size, + loc=loc, + ip=ip, + ) + + # Return raw ir.Value for vectors, wrapped Numeric for scalars + if is_vector: + return result + else: + return scalar_dtype(result) + + @dsl_user_op def cvt_f4e2m1_f16(src, *, loc=None, ip=None): # 0 padding for upper 4 bits diff --git a/python/CuTeDSL/cutlass/cute/arch/tmem.py b/python/CuTeDSL/cutlass/cute/arch/tmem.py index f87fbc37..2e393672 100644 --- a/python/CuTeDSL/cutlass/cute/arch/tmem.py +++ b/python/CuTeDSL/cutlass/cute/arch/tmem.py @@ -16,11 +16,61 @@ from cutlass.cutlass_dsl import dsl_user_op import cutlass._mlir.dialects.cute as _cute_ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir -from ..typing import Pointer, Int, Int32, Numeric, NumericMeta +from ..typing import Pointer, Int, Int32, Numeric, NumericMeta, Tensor + +SM100_TMEM_CAPACITY_COLUMNS = ( + 512 # deprecated; use get_max_tmem_alloc_cols(arch="sm_100") instead +) +SM100_TMEM_MIN_ALLOC_COLUMNS = ( + 32 # deprecated; use get_min_tmem_alloc_cols(arch="sm_100") instead +) + +TMEM_MAX_ALLOC_COLUMNS_MAP = { + "sm_120": 512, + "sm_103": 512, + "sm_100": 512, +} + +TMEM_MIN_ALLOC_COLUMNS_MAP = { + "sm_120": 32, + "sm_103": 32, + "sm_100": 32, +} -SM100_TMEM_CAPACITY_COLUMNS = 512 -SM100_TMEM_MIN_ALLOC_COLUMNS = 32 +def get_max_tmem_alloc_cols(compute_capability: str) -> int: + """Get the tensor memory capacity in columns for a given compute capability. + + Returns the maximum TMEM capacity in columns available for the specified + GPU compute capability. + + :param compute_capability: The compute capability string (e.g. "sm_100", "sm_103") + :type compute_capability: str + :return: The TMEM capacity in columns + :rtype: int + :raises ValueError: If the compute capability is not supported + """ + if compute_capability not in TMEM_MAX_ALLOC_COLUMNS_MAP: + raise ValueError(f"Unsupported compute capability: {compute_capability}") + return TMEM_MAX_ALLOC_COLUMNS_MAP[compute_capability] + + + +def get_min_tmem_alloc_cols(compute_capability: str) -> int: + """Get the minimum TMEM allocation columns for a given compute capability. + + Returns the minimum TMEM allocation columns available for the specified + GPU compute capability. + + :param compute_capability: The compute capability string (e.g. "sm_100", "sm_103") + :type compute_capability: str + :return: The minimum TMEM allocation columns + :rtype: int + :raises ValueError: If the compute capability is not supported + """ + if compute_capability not in TMEM_MIN_ALLOC_COLUMNS_MAP: + raise ValueError(f"Unsupported compute capability: {compute_capability}") + return TMEM_MIN_ALLOC_COLUMNS_MAP[compute_capability] @dsl_user_op @@ -64,6 +114,7 @@ def alloc_tmem( smem_ptr_to_write_address: Pointer, is_two_cta=None, *, + arch: str = "sm_100", loc=None, ip=None, ) -> None: @@ -76,15 +127,19 @@ def alloc_tmem( to :type smem_ptr_to_write_address: Pointer :param is_two_cta: Optional boolean parameter for 2-CTA MMAs + :param arch: The architecture of the GPU. + :type arch: str """ + tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch) + tmem_min_alloc_cols = get_min_tmem_alloc_cols(arch) if isinstance(num_columns, int): if ( - num_columns < SM100_TMEM_MIN_ALLOC_COLUMNS - or num_columns > SM100_TMEM_CAPACITY_COLUMNS + num_columns < tmem_min_alloc_cols + or num_columns > tmem_max_alloc_cols or not (num_columns & (num_columns - 1) == 0) ): raise ValueError( - f"num_columns must be between 32 and 512, and must be pow of 2, but got {num_columns}" + f"num_columns must be between {tmem_min_alloc_cols} and {tmem_max_alloc_cols}, and must be pow of 2, but got {num_columns}" ) _cute_nvgpu_ir.arch_sm100_alloc_tmem( Int32(num_columns).ir_value(loc=loc, ip=ip), @@ -112,6 +167,7 @@ def dealloc_tmem( num_columns: Int, is_two_cta=None, *, + arch: str = "sm_100", loc=None, ip=None, ) -> None: @@ -123,15 +179,19 @@ def dealloc_tmem( :param num_columns: The number of columns in the TMEM allocation :type num_columns: Int :param is_two_cta: Optional boolean parameter for 2-CTA MMAs + :param arch: The architecture of the GPU. + :type arch: str """ + tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch) + tmem_min_alloc_cols = get_min_tmem_alloc_cols(arch) if isinstance(num_columns, int): if ( - num_columns < SM100_TMEM_MIN_ALLOC_COLUMNS - or num_columns > SM100_TMEM_CAPACITY_COLUMNS + num_columns < tmem_min_alloc_cols + or num_columns > tmem_max_alloc_cols or not (num_columns & (num_columns - 1) == 0) ): raise ValueError( - f"num_columns must be between 32 and 512, and must be pow of 2, but got {num_columns}" + f"num_columns must be between {tmem_min_alloc_cols} and {tmem_max_alloc_cols}, and must be pow of 2, but got {num_columns}" ) _cute_nvgpu_ir.arch_sm100_dealloc_tmem( tmem_ptr.value, diff --git a/python/CuTeDSL/cutlass/cute/atom.py b/python/CuTeDSL/cutlass/cute/atom.py index 4881b05e..ba453317 100644 --- a/python/CuTeDSL/cutlass/cute/atom.py +++ b/python/CuTeDSL/cutlass/cute/atom.py @@ -136,10 +136,15 @@ class Atom(ABC): self._trait = trait def __extract_mlir_values__(self): - return extract_mlir_values(self._trait) + return extract_mlir_values(self._trait) + extract_mlir_values(self._op) def __new_from_mlir_values__(self, values): - return self.__class__(self.op, new_from_mlir_values(self._trait, values)) + traits_value = values[: len(extract_mlir_values(self._trait))] + op_value = values[len(extract_mlir_values(self._trait)) :] + + new_trait = new_from_mlir_values(self._trait, traits_value) + new_op = new_from_mlir_values(self._op, op_value) + return self.__class__(new_op, new_trait) @property def op(self) -> Op: @@ -318,24 +323,29 @@ class TiledMma(MmaAtom): # @property - def tv_layout_A_tiled(self) -> Layout: - return static(self._trait.value.type.layout_a_tv_tiled) + @dsl_user_op + def tv_layout_A_tiled(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.layout_a_tv_tiled, loc=loc, ip=ip) @property - def tv_layout_B_tiled(self) -> Layout: - return static(self._trait.value.type.layout_b_tv_tiled) + @dsl_user_op + def tv_layout_B_tiled(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.layout_b_tv_tiled, loc=loc, ip=ip) @property - def tv_layout_C_tiled(self) -> Layout: - return static(self._trait.value.type.layout_c_tv_tiled) + @dsl_user_op + def tv_layout_C_tiled(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.layout_c_tv_tiled, loc=loc, ip=ip) @property - def permutation_mnk(self) -> Tile: - return _unpack_x_tuple(self._trait.value.type.permutation_mnk) + @dsl_user_op + def permutation_mnk(self, *, loc=None, ip=None) -> Tile: + return _unpack_x_tuple(self._trait.value.type.permutation_mnk, loc=loc, ip=ip) @property - def thr_layout_vmnk(self) -> Layout: - return static(self._trait.value.type.thr_layout_vmnk) + @dsl_user_op + def thr_layout_vmnk(self, *, loc=None, ip=None) -> Layout: + return static(self._trait.value.type.thr_layout_vmnk, loc=loc, ip=ip) @property def size(self) -> int: @@ -598,6 +608,33 @@ class CopyAtom(Atom): def layout_dst_tv(self) -> Layout: return static(self._trait.value.type.layout_dst_tv) + @property + def smem_layout(self): + """ + Convenience property to access the SMEM layout for TMA copy atoms. + + This is a shortcut for ``atom.op.smem_layout`` that checks if the operation + is a TMA operation and provides a clearer error message if not. + + :return: The SMEM layout + :rtype: Layout or ComposedLayout + :raises TypeError: If the operation is not a TMA operation + :raises ValueError: If the SMEM layout is not set + + Example: + >>> layout = tma_atom.smem_layout # Instead of tma_atom.op.smem_layout + """ + # Import here to avoid circular dependency + from .nvgpu.cpasync.copy import TmaCopyOp + + if not isinstance(self.op, TmaCopyOp): + raise TypeError( + f"smem_layout is only available for TMA copy operations, " + f"but this atom uses {type(self.op).__name__}" + ) + + return self.op.smem_layout + class TiledCopy(CopyAtom): """ diff --git a/python/CuTeDSL/cutlass/cute/core.py b/python/CuTeDSL/cutlass/cute/core.py index 4ad53b5e..632a3b77 100644 --- a/python/CuTeDSL/cutlass/cute/core.py +++ b/python/CuTeDSL/cutlass/cute/core.py @@ -17,7 +17,7 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload from typing_extensions import deprecated from cutlass._mlir import ir -from cutlass._mlir.dialects import builtin, llvm +from cutlass._mlir.dialects import builtin, llvm, vector from cutlass._mlir.dialects import cute as _cute_ir from cutlass._mlir.dialects.cute import ( Ratio as _Ratio, @@ -66,6 +66,91 @@ from .typing import ( is_integer, ) +__all__ = [ + # Classes + "IntValue", + "Swizzle", + "struct", + # Utility functions + "E", + "get_divisibility", + "is_valid_leaf", + "is_static", + "has_underscore", + "has_scaled_basis", + "pretty_str", + "printf", + # Layout operations + "front", + "is_major", + "assume", + "make_swizzle", + "static", + "get_leaves", + "depth", + "rank", + "is_congruent", + "is_weakly_congruent", + "get", + "select", + "group_modes", + "slice_", + "dice", + "prepend", + "append", + "prepend_ones", + "append_ones", + "repeat_as_tuple", + "repeat", + "repeat_like", + "flatten", + "filter_zeros", + "filter", + "size", + "shape_div", + "ceil_div", + "round_up", + "make_layout", + "make_identity_layout", + "make_ordered_layout", + "make_layout_like", + "make_composed_layout", + "cosize", + "size_in_bytes", + "coalesce", + "crd2idx", + "idx2crd", + "recast_layout", + "slice_and_offset", + "shape", + "recast_ptr", + "make_ptr", + "composition", + "complement", + "right_inverse", + "left_inverse", + "logical_product", + "zipped_product", + "tiled_product", + "flat_product", + "raked_product", + "blocked_product", + "logical_divide", + "zipped_divide", + "tiled_divide", + "flat_divide", + "max_common_layout", + "max_common_vector", + "tile_to_shape", + "local_partition", + "local_tile", + "make_layout_image_mask", + "leading_dim", + "make_layout_tv", + "get_nonswizzle_portion", + "get_swizzle_portion", +] + #################################################################################################### # # Internal IntTuple helpers @@ -131,8 +216,8 @@ def _pack_tile(tile: Tile, *, loc=None, ip=None) -> ir.Value: leaves = [] for e in tile: if isinstance(e, _Layout): - leaves.extend(list(flatten_to_tuple(e.shape))) - leaves.extend(list(flatten_to_tuple(e.stride))) + leaves.extend(list(flatten_to_tuple(e.shape_method(loc=loc, ip=ip)))) + leaves.extend(list(flatten_to_tuple(e.stride_method(loc=loc, ip=ip)))) else: leaves.append(e) return leaves @@ -143,6 +228,7 @@ def _pack_tile(tile: Tile, *, loc=None, ip=None) -> ir.Value: _get_typed_value(x) for x in dyn_elems if isinstance(x, (Integer, ir.Value)) ] + tile = transform_leaf(_get_typed_value, tile) res_ty = _cute_ir.pack_tile(tile) return _cute_ir.make_tile(res_ty, dyn_elems, loc=loc, ip=ip) @@ -332,16 +418,18 @@ class IntValue(cutlass_arith.ArithValue): # Dispatch to `__rmul__` of `other` return NotImplemented - return IntValue(op(self, other_val, **kwargs)) + return IntValue( + op(self, other_val, **kwargs), + loc=kwargs.get("loc"), + ip=kwargs.get("ip"), + ) return wrapper @dsl_user_op @_binary_op def __add__(self, other, *, loc=None, ip=None): - return _cute_ir.add_offset( - self.get_typed_value(loc=loc, ip=ip), other, loc=loc, ip=ip - ) + return _cute_ir.tuple_add(self.get_typed_value(), other, loc=loc, ip=ip) @dsl_user_op @_binary_op @@ -373,10 +461,8 @@ class IntValue(cutlass_arith.ArithValue): @dsl_user_op @_binary_op - def __radd__(self, other, *, loc=None, ip=None) -> "IntValue": - return _cute_ir.add_offset( - other, self.get_typed_value(loc=loc, ip=ip), loc=loc, ip=ip - ) + def __radd__(self, other, *, loc=None, ip=None): + return _cute_ir.tuple_add(other, self.get_typed_value(), loc=loc, ip=ip) @dsl_user_op @_binary_op @@ -579,6 +665,9 @@ class ScaledBasis: if isinstance(self._value, Integer): scale = self._value.ir_value(loc=loc, ip=ip) return _ScaledBasis(scale, self._mode, get_divisibility(scale)) + elif isinstance(self._value, cutlass_arith.ArithValue): + scale = self._value + return _ScaledBasis(scale, self._mode, get_divisibility(scale)) else: scale = self._value return _ScaledBasis(scale, self._mode) @@ -658,6 +747,27 @@ class ScaledBasis: return ScaledBasis(scale * value, self.mode) # type: ignore + def __mul__( + self, scale: Union[Int, ir.Value, Ratio], *, loc=None, ip=None + ) -> "ScaledBasis": + """Multiplication by a scale factor. + This operation is used in layout algebra to scale basis elements, + which is essential for operations like composition and partitioning. + + :param scale: The scale factor + :type scale: Union[Int, ir.Value, Ratio] + :param loc: The source location for the operation, defaults to None + :type loc: Location, optional + :param ip: The insertion point for the operation, defaults to None + :type ip: InsertionPoint, optional + :return: A new scaled basis element + :rtype: ScaledBasis + :raises TypeError: If scale is not of a supported type + :raises NotImplementedError: If scaling a basis element with a ratio value + """ + + return self.__rmul__(scale, loc=loc, ip=ip) + def __extract_mlir_values__(self): if isinstance(self.value, Ratio): # Ratio is always static @@ -716,6 +826,79 @@ def get_divisibility(x: Union[int, Integer]) -> int: return 1 +def basis_value(e: Union[ScaledBasis, Any]) -> Union[Int, ir.Value, Ratio]: + """Extract the value from a ScaledBasis or return the input as-is. + + If the input is a ScaledBasis, returns its value component. + Otherwise, returns the input unchanged. + + :param e: The input element (ScaledBasis or any other type) + :type e: Any + :return: The value of the ScaledBasis or the input itself + :rtype: Any + + **Examples:** + + .. code-block:: python + + >>> basis_value(ScaledBasis(5, 0)) + 5 + >>> basis_value(42) + 42 + """ + if isinstance(e, ScaledBasis): + return e.value + else: + return e + + +@dsl_user_op +def basis_get( + basis: Union[ScaledBasis, Numeric, int], + t: Union[XTuple, Layout, ComposedLayout], + *, + loc=None, + ip=None, +) -> Union[XTuple, Layout, ComposedLayout]: + """Apply the mode indices from a ScaledBasis to get an element from a tuple, layout, or composed layout. + + If the basis is a ScaledBasis or Numeric with mode indices, this function uses those + indices to extract the corresponding element from the tuple using hierarchical + indexing. If the basis is not a ScaledBasis or has no modes, returns the tuple, layout, or composed layout as-is. + + :param basis: The basis element (ScaledBasis) + :type basis: ScaledBasis + :param t: The tuple, layout, or composed layout to index into + :type t: Union[XTuple, Layout, ComposedLayout] + :return: The element at the position specified by the basis modes, or t itself + :rtype: Union[XTuple, Layout, ComposedLayout] + + **Examples:** + + .. code-block:: python + + >>> basis_get(ScaledBasis(2, 1), (10, 20, 30)) + 20 + >>> basis_get(ScaledBasis(2, [0, 1]), ((10, 20), (30, 40))) + 20 + >>> basis_get(5, (10, 20, 30)) # Non-basis returns tuple as-is + (10, 20, 30) + """ + if isinstance(basis, ScaledBasis): + modes = basis.mode + if len(modes) == 0: + return t + else: + # Use hierarchical indexing with the mode list + return get(t, modes, loc=loc, ip=ip) + elif isinstance(basis, (Numeric, int)): + return t + else: + raise TypeError( + f"basis must be a ScaledBasis or Numeric, but got {type(basis)}" + ) + + @ir.register_value_caster(_cute_ir.SwizzleType.get_static_typeid(), replace=True) class Swizzle(ir.Value): """ @@ -751,6 +934,41 @@ class Swizzle(ir.Value): # Cut off the MLIR type's string for making pretty_str more concise return self.type.__str__()[15 : 15 + 8] + def __eq__(self, other) -> Union[bool, Boolean]: + """Check if this Swizzle is equal to another Swizzle. Since num_bits, num_base, and num_shift are static, + this is a constant expression. + + Two Swizzles are equal if they have the same num_bits, num_base, and num_shift. + + :param other: The Swizzle to compare with. + :return: True if Swizzles are equal, False otherwise. + """ + if isinstance(other, Swizzle): + return self.type == other.type + else: + return False + + @property + def num_bits(self) -> int: + """ + Returns the number of bits in the mask (B in Sw). + """ + return self.type.num_bits + + @property + def num_base(self) -> int: + """ + Returns the number of least-significant bits to keep constant (M in Sw). + """ + return self.type.num_base + + @property + def num_shift(self) -> int: + """ + Returns the distance to shift the mask (S in Sw). + """ + return self.type.num_shift + @ir.register_value_caster(_cute_ir.LayoutType.get_static_typeid(), replace=True) class _Layout(Layout): @@ -794,12 +1012,26 @@ class _Layout(Layout): """ super().__init__(op_result) - def __str__(self) -> str: + def __repr__(self, *, loc=None, ip=None) -> str: + return self.__str__(loc=loc, ip=ip) + + def __str__(self, *, loc=None, ip=None) -> str: """Return a string representation of the layout. :return: A string in the format "shape:stride". """ - return f"{pretty_str(self.shape)}:{pretty_str(self.stride)}" + type_str = self.type.__str__() + return type_str[type_str.find("<") + 2 : type_str.rfind(">") - 1] + + @lru_cache_ir() + def shape_method(self, *, loc=None, ip=None) -> Shape: + return _unpack_x_tuple(_cute_ir.get_shape(self, loc=loc, ip=ip), loc=loc, ip=ip) + + @lru_cache_ir() + def stride_method(self, *, loc=None, ip=None) -> Stride: + return _unpack_x_tuple( + _cute_ir.get_stride(self, loc=loc, ip=ip), loc=loc, ip=ip + ) @property @dsl_user_op @@ -812,7 +1044,7 @@ class _Layout(Layout): :return: The hierarchical shape of the layout. """ - return _unpack_x_tuple(_cute_ir.get_shape(self, loc=loc, ip=ip), loc=loc, ip=ip) + return self.shape_method(loc=loc, ip=ip) @property @dsl_user_op @@ -824,9 +1056,7 @@ class _Layout(Layout): :return: The hierarchical stride of the layout. """ - return _unpack_x_tuple( - _cute_ir.get_stride(self, loc=loc, ip=ip), loc=loc, ip=ip - ) + return self.stride_method(loc=loc, ip=ip) @property def max_alignment(self) -> int: @@ -977,6 +1207,10 @@ class _ComposedLayout(ComposedLayout): @property @dsl_user_op def shape(self, *, loc=None, ip=None) -> Shape: + return self.shape_method(loc=loc, ip=ip) + + @dsl_user_op + def shape_method(self, *, loc=None, ip=None) -> Shape: return _unpack_x_tuple( _cute_ir.get_shape(self.value, loc=loc, ip=ip), loc=loc, ip=ip ) @@ -1088,11 +1322,15 @@ class _Pointer(Pointer): @property @lru_cache_ir() - def dtype(self) -> Union[Type[Numeric], _cute_ir.SparseElemType]: - if isinstance(self.value.type.value_type, _cute_ir.SparseElemType): - return self.value.type.value_type - else: - return Numeric.from_mlir_type(self.value.type.value_type) + def dtype( + self, + ) -> Union[ + Type[Numeric], + ]: + ret_type = None + if ret_type is None: + ret_type = Numeric.from_mlir_type(self.value.type.value_type) + return ret_type @property def alignment(self) -> int: @@ -1121,6 +1359,25 @@ class _Pointer(Pointer): """ Get the LLVM pointer representation of this pointer. + :param loc: Source location for MLIR, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for MLIR, defaults to None + :type ip: Optional[InsertionPoint] + :return: The LLVM pointer representation + :rtype: ir.Value + """ + return self.to_llvm_ptr(loc=loc, ip=ip) + + @dsl_user_op + @lru_cache_ir() + def to_llvm_ptr(self, *, loc=None, ip=None) -> ir.Value: + """ + Get the LLVM pointer representation of this pointer. (Used by internal API to propagate loc and ip) + + :param loc: Source location for MLIR, defaults to None + :type loc: Optional[Location] + :param ip: Insertion point for MLIR, defaults to None + :type ip: Optional[InsertionPoint] :return: The LLVM pointer representation :rtype: ir.Value """ @@ -1212,16 +1469,16 @@ class _Pointer(Pointer): #################################################################################################### -def _op_wrapper(op_fn, input): +def _op_wrapper(op_fn, input, *, loc=None, ip=None): from .tensor import _Tensor if isinstance(input, Tensor): - res = op_fn(input.value) - return _Tensor(res, dtype=input.element_type) + res = op_fn(input.value, loc=loc, ip=ip) + return _Tensor(res, dtype=input.element_type, loc=loc, ip=ip) elif isinstance(input, _ComposedLayout): - return op_fn(input.value) + return op_fn(input.value, loc=loc, ip=ip) else: - return op_fn(input) + return op_fn(input, loc=loc, ip=ip) # @@ -1486,17 +1743,14 @@ def assume(src, divby=None, *, loc=None, ip=None): @dsl_user_op def make_swizzle(b, m, s, *, loc=None, ip=None): # canonicalize to <0, 4, 3> for identity swizzle (as compiler assumes <0, 4, 3>) + if not isinstance(b, int) or not isinstance(m, int) or not isinstance(s, int): + raise ValueError("b, m, and s must be int") if b == 0: m, s = 4, 3 ty = ir.Type.parse(f'!cute.swizzle<"S<{b},{m},{s}>">') return Swizzle(static(ty, loc=loc, ip=ip)) -@dsl_user_op -def make_sparse_elem(num_logical, num_phys, elem_type, *, loc=None, ip=None): - return _cute_ir.SparseElemType.get(num_logical, num_phys, elem_type.mlir_type) - - @dsl_user_op def static(value, *, loc=None, ip=None): return _cute_ir.static(value, loc=loc, ip=ip) @@ -1846,7 +2100,7 @@ def group_modes(input, begin: int, end: Optional[int] = None, *, loc=None, ip=No return (*input[:begin], (input[begin:end]), *input[end:]) return _op_wrapper( - partial(_cute_ir.group_modes, begin=begin, end=end, loc=loc, ip=ip), input + partial(_cute_ir.group_modes, begin=begin, end=end), input, loc=loc, ip=ip ) @@ -1939,7 +2193,7 @@ def slice_(src, coord: Coord, *, loc=None, ip=None): return () coord_val = _pack_coord(coord, loc=loc, ip=ip) - return _op_wrapper(partial(_cute_ir.slice, coord=coord_val, loc=loc, ip=ip), src) + return _op_wrapper(partial(_cute_ir.slice, coord=coord_val), src, loc=loc, ip=ip) @overload @@ -2015,7 +2269,7 @@ def dice(src, dicer, *, loc=None, ip=None): dicer_val = _pack_coord(dicer, loc=loc, ip=ip) return _op_wrapper( - partial(_cute_ir.dice, coord=dicer_val.type.attribute, loc=loc, ip=ip), src + partial(_cute_ir.dice, coord=dicer_val.type.attribute), src, loc=loc, ip=ip ) @@ -2030,7 +2284,7 @@ def _extend(func, input, elem, up_to_rank, loc, ip): raise TypeError(f"Input type of elem ({type(elem)}) is not accepted!") N = rank(input) + 1 if up_to_rank is None else up_to_rank - return _op_wrapper(partial(func, N, element=elem, loc=loc, ip=ip), input) + return _op_wrapper(partial(func, N, element=elem), input, loc=loc, ip=ip) if is_valid_leaf(input) or isinstance(input, tuple): if elem is None: @@ -2717,9 +2971,7 @@ class _ComposedLayoutWithInnerFunc(ComposedLayout): delta = self._outer(coord) delta_val = _pack_int_tuple(delta, loc=loc, ip=ip) - offset_val_new = _cute_ir.add_offset( - self._offset_val, delta_val, loc=loc, ip=ip - ) + offset_val_new = _cute_ir.tuple_add(self._offset_val, delta_val, loc=loc, ip=ip) offset_new = _unpack_x_tuple(offset_val_new, loc=loc, ip=ip) return self._inner(offset_new) @@ -2906,7 +3158,7 @@ def coalesce(input, *, target_profile: Coord = None, loc=None, ip=None): profile_val = None return _op_wrapper( - partial(_cute_ir.coalesce, target_profile=profile_val, loc=loc, ip=ip), input + partial(_cute_ir.coalesce, target_profile=profile_val), input, loc=loc, ip=ip ) @@ -2989,15 +3241,13 @@ def idx2crd(idx, shape, *, loc=None, ip=None): import cutlass.cute as cute @cute.jit def foo(): - coord = cute.idx2crd(11, (5,4)) + coord = cute.idx2crd(11, (5, 4)) # idx2crd is always col-major # For shape (m, n, l, ...), coord = (idx % m, idx // m % n, idx // m // n % l, ... # Computed as: (11 % 5, 11 // 5 % 4) = (1, 2) print(coord) - foo() # Expected output: (1, 2) - **Note:** - Python DSL is aligned with C++ DSL. + foo() # Expected output: (1, 2) """ if is_integer(idx) and is_integer(shape): return idx @@ -3008,7 +3258,56 @@ def idx2crd(idx, shape, *, loc=None, ip=None): @dsl_user_op -def recast_layout(new_type_bits, old_type_bits, src_layout, *, loc=None, ip=None): +def recast_layout( + new_type_bits: int, + old_type_bits: int, + src_layout: Union[Layout, ComposedLayout], + *, + loc=None, + ip=None, +): + """ + Recast a layout from one data type to another. + + :param new_type_bits: The new data type bits + :type new_type_bits: int + :param old_type_bits: The old data type bits + :type old_type_bits: int + :param src_layout: The layout to recast + :type src_layout: Union[Layout, ComposedLayout] + :param loc: Optional location information for IR diagnostics. + :type loc: optional + :param ip: Optional instruction pointer or context for underlying IR functions. + :type ip: optional + :return: The recast layout + :rtype: Layout or ComposedLayout + + **Example:** + + .. code-block:: python + + import cutlass.cute as cute + @cute.jit + def foo(): + # Create a layout + L = cute.make_layout((2, 3, 4)) + # Recast the layout to a different data type + L_recast = cute.recast_layout(16, 8, L) + print(L_recast) + foo() # Expected output: (2, 3, 4) + """ + if not isinstance(new_type_bits, int): + raise TypeError( + f"new_type_bits must be an integer instead got {type(new_type_bits)}" + ) + if not isinstance(old_type_bits, int): + raise TypeError( + f"old_type_bits must be an integer instead got {type(old_type_bits)}" + ) + if not isinstance(src_layout, (Layout, ComposedLayout)): + raise TypeError( + f"src_layout must be a layout or composed layout instead got {type(src_layout)}" + ) if isinstance(src_layout, _ComposedLayout): src_layout = src_layout.value return _cute_ir.recast_layout( @@ -3086,15 +3385,14 @@ def recast_ptr( loc=None, ip=None, ) -> Pointer: + cvt_type = None if dtype is not None: - if isinstance(dtype, _cute_ir.SparseElemType): - # use SparseElemType as dtype - pass - else: + if cvt_type is None: if not isclass(dtype) or not issubclass(dtype, Numeric): raise TypeError(f"dtype must be a type of Numeric, but got {dtype}") - dtype = dtype.mlir_type + cvt_type = dtype.mlir_type + dtype = cvt_type value_type = ptr.type.value_type if dtype is None else dtype swizzle = swizzle_.type.attribute if swizzle_ is not None else None res_ty = _cute_ir.PtrType.get(value_type, ptr.memspace, ptr.alignment, swizzle) @@ -3405,7 +3703,7 @@ def logical_divide(target, tiler: Tiler, *, loc=None, ip=None): if isinstance(tiler, tuple): tiler = _pack_tile(tiler, loc=loc, ip=ip) # type: ignore return _op_wrapper( - partial(_cute_ir.logical_divide, tiler=tiler, loc=loc, ip=ip), target + partial(_cute_ir.logical_divide, tiler=tiler), target, loc=loc, ip=ip ) @@ -3420,7 +3718,7 @@ def zipped_divide(target, tiler: Tiler, *, loc=None, ip=None): if isinstance(tiler, tuple): tiler = _pack_tile(tiler, loc=loc, ip=ip) # type: ignore return _op_wrapper( - partial(_cute_ir.zipped_divide, tiler=tiler, loc=loc, ip=ip), target + partial(_cute_ir.zipped_divide, tiler=tiler), target, loc=loc, ip=ip ) @@ -3435,7 +3733,7 @@ def tiled_divide(target, tiler: Tiler, *, loc=None, ip=None): if isinstance(tiler, tuple): tiler = _pack_tile(tiler, loc=loc, ip=ip) return _op_wrapper( - partial(_cute_ir.tiled_divide, tiler=tiler, loc=loc, ip=ip), target + partial(_cute_ir.tiled_divide, tiler=tiler), target, loc=loc, ip=ip ) @@ -3450,7 +3748,7 @@ def flat_divide(target, tiler: Tile, *, loc=None, ip=None): if isinstance(tiler, tuple): tiler = _pack_tile(tiler, loc=loc, ip=ip) return _op_wrapper( - partial(_cute_ir.flat_divide, tiler=tiler, loc=loc, ip=ip), target + partial(_cute_ir.flat_divide, tiler=tiler), target, loc=loc, ip=ip ) @@ -3654,7 +3952,6 @@ def leading_dim(shape: Shape, stride: Stride) -> Union[int, Tuple[int, ...], Non return find_if(stride, pred_fn=pred_fn) - @dsl_user_op def make_layout_tv( thr_layout: Layout, val_layout: Layout, *, loc=None, ip=None @@ -3732,11 +4029,76 @@ def make_layout_tv( right_inverse(layout_mn, loc=loc, ip=ip), tmp, loc=loc, ip=ip ) - tiler_mn = product_each(layout_mn.shape, loc=loc, ip=ip) + tiler_mn = product_each(layout_mn.shape_method(loc=loc, ip=ip), loc=loc, ip=ip) return (tiler_mn, layout_tv) +@dsl_user_op +def get_nonswizzle_portion( + layout: Union[Layout, ComposedLayout], *, loc=None, ip=None +) -> Union[Layout, ComposedLayout]: + """ + Extract the non-swizzle portion from a layout. + + For a simple Layout, the entire layout is considered non-swizzled and is returned as-is. + For a ComposedLayout, the inner layout (non-swizzled portion) is extracted and returned, + effectively separating the base layout from any swizzle transformation that may be applied. + + :param layout: A Layout or ComposedLayout from which to extract the non-swizzle portion. + :type layout: Union[Layout, ComposedLayout] + :param loc: Optional location information for IR diagnostics. + :type loc: optional + :param ip: Optional + :type ip: optional + :returns: The non-swizzle portion of the input layout. For Layout objects, returns the layout itself. + For ComposedLayout objects, returns the outer layout component. + :rtype: Layout + :raises TypeError: If the layout is neither a Layout nor a ComposedLayout. + """ + if isinstance(layout, Layout): + return layout + elif isinstance(layout, ComposedLayout): + return layout.outer + else: + raise TypeError(f"expects a Layout or ComposedLayout, but got {type(layout)}") + + +@dsl_user_op +def get_swizzle_portion( + layout: Union[Layout, ComposedLayout], *, loc=None, ip=None +) -> Swizzle: + """ + Extract or create the swizzle portion from a layout. + + For a simple Layout (which has no explicit swizzle), a default identity swizzle is created. + For a ComposedLayout, the outer layout is checked and returned if it is a Swizzle object. + Otherwise, a default identity swizzle is created. The default identity swizzle has parameters + (0, 4, 3), which represents a no-op swizzle transformation. + + :param layout: A Layout or ComposedLayout from which to extract the swizzle portion. + :type layout: Union[Layout, ComposedLayout] + :param loc: Optional location information for IR diagnostics. + :type loc: optional + :param ip: Optional + :type ip: optional + :returns: The swizzle portion of the layout. For Layout objects or ComposedLayout objects without + a Swizzle outer component, returns a default identity swizzle (0, 4, 3). For ComposedLayout + objects with a Swizzle outer component, returns that swizzle. + :rtype: Swizzle + :raises TypeError: If the layout is neither a Layout nor a ComposedLayout. + """ + if isinstance(layout, Layout): + return make_swizzle(0, 4, 3, loc=loc, ip=ip) + elif isinstance(layout, ComposedLayout): + if isinstance(layout.inner, Swizzle): + return layout.inner + else: + return make_swizzle(0, 4, 3, loc=loc, ip=ip) + else: + raise TypeError(f"expects a Layout or ComposedLayout, but got {type(layout)}") + + ############################################################################## # User defined struct ############################################################################## @@ -3870,7 +4232,8 @@ class struct: self._size = size self._base = base - def data_ptr(self): + @dsl_user_op + def data_ptr(self, *, loc=None, ip=None): """ Returns start pointer to the data in this memory range. @@ -3878,9 +4241,10 @@ class struct: :raises AssertionError: If the size of the memory range is negative. """ assert self._size >= 0 - return recast_ptr(self._base, dtype=self._dtype) + return recast_ptr(self._base, dtype=self._dtype, loc=loc, ip=ip) - def get_tensor(self, layout, swizzle=None, dtype=None): + @dsl_user_op + def get_tensor(self, layout, swizzle=None, dtype=None, *, loc=None, ip=None): """ Creates a tensor from the memory range. @@ -3898,8 +4262,8 @@ class struct: if isinstance(layout, ComposedLayout) and (swizzle is not None): raise TypeError("incompatible layout with swizzle") elem_type = self._dtype if dtype is None else dtype - ptr = recast_ptr(self._base, swizzle, dtype=elem_type) - res = make_tensor(ptr, layout) + ptr = recast_ptr(self._base, swizzle, dtype=elem_type, loc=loc, ip=ip) + res = make_tensor(ptr, layout, loc=loc, ip=ip) return res def __getitem__(self, index: int) -> Any: @@ -4046,7 +4410,8 @@ class struct: self._size_of = self.align_offset(offset, alignment) # create the __init__ method for decorated struct - def __call__(self, base: Any) -> None: + @dsl_user_op + def __call__(self, base: Any, *, loc=None, ip=None) -> None: """ Creates a new instance of the decorated struct. @@ -4065,7 +4430,7 @@ class struct: if isinstance(obj, struct._AlignMeta): obj = obj.dtype if struct._is_scalar_type(obj): - new_obj = recast_ptr(base + off, dtype=obj) + new_obj = recast_ptr(base + off, dtype=obj, loc=loc, ip=ip) setattr(cls, name, new_obj) elif isinstance(obj, struct._MemRangeMeta): new_obj = struct._MemRangeData(obj._dtype, obj._size, base + off) diff --git a/python/CuTeDSL/cutlass/cute/experimental/__init__.py b/python/CuTeDSL/cutlass/cute/experimental/__init__.py new file mode 100755 index 00000000..d913a607 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/experimental/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +raise NotImplementedError( + "CuTe Experimental module is only supported on Cuda toolkit 13.1 and above!" +) diff --git a/python/CuTeDSL/cutlass/cute/export/__init__.py b/python/CuTeDSL/cutlass/cute/export/__init__.py index 9eb79612..1b9546ba 100644 --- a/python/CuTeDSL/cutlass/cute/export/__init__.py +++ b/python/CuTeDSL/cutlass/cute/export/__init__.py @@ -11,28 +11,39 @@ from .c_header_generator import CuteCHeaderGenerator -from ...base_dsl.export import ( - get_export_module, - dump_to_object as _dump_to_object, - export_to_c as _export_to_c, -) -from ...cutlass_dsl import CuTeDSL -from functools import partial as _partial -from ...cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction +from .export import object_file_version as _object_file_version +from .export import CuteArgsSpecProcessor as _CuteArgsSpecProcessor -dump_to_object = _partial( - _dump_to_object, - dsl=CuTeDSL._get_dsl(), +from ...base_dsl.jit_executor import ExportProvider as _ExportProvider +from ...cutlass_dsl import CuTeDSL as _CuTeDSL +from ...cutlass_dsl.cuda_jit_executor import ( + CudaDialectJitCompiledFunction as _CudaDialectJitCompiledFunction, ) -export_to_c = _partial( - _export_to_c, - dsl=CuTeDSL._get_dsl(), +from ..._mlir._mlir_libs._cutlass_ir import _mlirExecutionEngine + +_CudaDialectJitCompiledFunction.export_provider = _ExportProvider( + dsl=_CuTeDSL, + arg_spec_processor=_CuteArgsSpecProcessor(), c_header_generator=CuteCHeaderGenerator(), - use_gpu_dialect=False, + object_file_version=_object_file_version, + mlirExecutionEngine=_mlirExecutionEngine, ) + +from ...base_dsl.export import ExternalBinaryModule as _ExternalBinaryModule +from ...base_dsl.export import LoadProvider as _LoadProvider +from .load import version_checker as _version_checker +from ..._mlir._mlir_libs._cutlass_ir._execution_engine import ( + BinaryExecutionEngine as _BinaryExecutionEngine, +) + +_ExternalBinaryModule.load_provider = _LoadProvider( + dsl=_CuTeDSL, + args_spec_processor=_CuteArgsSpecProcessor(), + version_checker=_version_checker, + execution_engine_constructor=_BinaryExecutionEngine, + jit_function_constructor=_CudaDialectJitCompiledFunction, +) + __all__ = [ "CuteCHeaderGenerator", - "get_export_module", - "dump_to_object", - "export_to_c", ] diff --git a/python/CuTeDSL/cutlass/cute/export/aot_config.py b/python/CuTeDSL/cutlass/cute/export/aot_config.py new file mode 100644 index 00000000..c143fac9 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/export/aot_config.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +""" +CLI tool to help with AOT compilation configuration. + +Similar to tvm-ffi-config or llvm-config, this tool provides compiler flags +for linking against CuTe DSL runtime libraries. + +Usage: + python -m cutlass.cute.export.aot_config --libdir # Returns the library directory path + python -m cutlass.cute.export.aot_config --ldflags # Returns -L flags for linking + python -m cutlass.cute.export.aot_config --libs # Returns -l flags for linking + +Examples: + # Compile and link a shared library using shell substitution + g++ -shared -o kernel.so kernel.o \\ + $(python -m cutlass.cute.export.aot_config --ldflags) \\ + $(python -m cutlass.cute.export.aot_config --libs) + + # Or using backticks + g++ -shared -o kernel.so kernel.o `python -m cutlass.cute.export.aot_config --ldflags` `python -m cutlass.cute.export.aot_config --libs` +""" + +import argparse +import sys +from pathlib import Path + + +def get_libdir() -> str: + """ + Get the library directory path containing libcuda_dialect_runtime.so. + + :return: Path to the library directory + :rtype: str + """ + from ..runtime import find_runtime_libraries + + libs = find_runtime_libraries(enable_tvm_ffi=False) + if libs: + # Return the directory containing the first library found + return str(Path(libs[0]).parent) + + return "" + + +def get_libs(enable_tvm_ffi: bool = False) -> str: + """ + Get the -l flags needed for AOT compilation linking. + + Similar to `tvm-ffi-config --libs` which returns `-ltvm_ffi`, + this returns `-lcuda_dialect_runtime` (and `-ltvm_ffi` if TVM-FFI is enabled). + + :param enable_tvm_ffi: Whether to include TVM-FFI library + :return: Space-separated -l flags (e.g., "-lcuda_dialect_runtime -ltvm_ffi") + :rtype: str + """ + from ..runtime import find_runtime_libraries + + libs = find_runtime_libraries(enable_tvm_ffi=enable_tvm_ffi) + # Convert full paths to -l flags + # e.g., /path/to/libcuda_dialect_runtime.so -> -lcuda_dialect_runtime + flags = [] + for lib in libs: + lib_path = Path(lib) + lib_name = lib_path.stem # e.g., "libcuda_dialect_runtime" + if lib_name.startswith("lib"): + lib_name = lib_name[3:] + flags.append(f"-l{lib_name}") + return " ".join(flags) + + +def get_lib_paths(enable_tvm_ffi: bool = False) -> list[str]: + """ + Get the full paths to runtime libraries. + + :param enable_tvm_ffi: Whether to include TVM-FFI library + :return: List of full library paths + :rtype: list[str] + """ + from ..runtime import find_runtime_libraries + + return find_runtime_libraries(enable_tvm_ffi=enable_tvm_ffi) + + +def get_ldflags() -> str: + """ + Get the -L flags for the linker. + + Similar to `tvm-ffi-config --ldflags` which returns `-L`. + + :return: -L flag with library directory path + :rtype: str + """ + libdir = get_libdir() + if libdir: + return f"-L{libdir}" + return "" + + +def main(): + parser = argparse.ArgumentParser( + description="AOT configuration helper for CuTe DSL (similar to tvm-ffi-config)", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Get library directory path + python -m cutlass.cute.export.aot_config --libdir + + # Get -L flags for linking + python -m cutlass.cute.export.aot_config --ldflags + + # Get -l flags for linking + python -m cutlass.cute.export.aot_config --libs + + # Compile a shared library + g++ -shared -o kernel.so kernel.o \\ + $(python -m cutlass.cute.export.aot_config --ldflags) \\ + $(python -m cutlass.cute.export.aot_config --libs) + """, + ) + + parser.add_argument( + "--libdir", + action="store_true", + help="Print the library directory path containing runtime libraries", + ) + parser.add_argument( + "--ldflags", + action="store_true", + help="Print -L flags for linking (e.g., -L/path/to/lib)", + ) + parser.add_argument( + "--libs", + action="store_true", + help="Print -l flags for linking (e.g., -lcuda_dialect_runtime)", + ) + parser.add_argument( + "--with-tvm-ffi", + action="store_true", + help="Include TVM-FFI library in --libs output (disabled by default)", + ) + + args = parser.parse_args() + + if not args.libdir and not args.ldflags and not args.libs: + parser.print_help() + sys.exit(1) + + enable_tvm_ffi = args.with_tvm_ffi + + if args.libdir: + print(get_libdir()) + + if args.ldflags: + print(get_ldflags()) + + if args.libs: + print(get_libs(enable_tvm_ffi=enable_tvm_ffi)) + + +if __name__ == "__main__": + main() + diff --git a/python/CuTeDSL/cutlass/cute/export/c_header_generator.py b/python/CuTeDSL/cutlass/cute/export/c_header_generator.py index 65a2cb13..7d5ff0e5 100644 --- a/python/CuTeDSL/cutlass/cute/export/c_header_generator.py +++ b/python/CuTeDSL/cutlass/cute/export/c_header_generator.py @@ -10,7 +10,7 @@ # is strictly prohibited. from cutlass.cute.typing import NumericMeta, Integer -from cutlass.base_dsl.export import CHeaderGenerator +from cutlass.base_dsl.export import CHeaderGenerator, CHeaderArguments from cutlass.base_dsl.dsl import is_dynamic_expression from cutlass.base_dsl.common import DSLRuntimeError from cutlass.base_dsl.jit_executor import ExecutionArgs @@ -29,6 +29,22 @@ import cuda.bindings.driver as cuda class CuteCHeaderGenerator(CHeaderGenerator): """This class provides a Export C Header Generator for cute c/cpp AOT support.""" + includes = """ +#pragma once + +#include +#include +#include +#include + +""" + cuda_error_check = r"""_CUDA_ERROR_CHECK(err) { \ + if ((err) != cudaSuccess) { \ + printf("Got Cuda Error %s: %s\n", cudaGetErrorName(err), cudaGetErrorString(err)); \ + } \ +} +""" + def _get_cute_algebra_type(self, arg_type: Any, arg: Any) -> str: """Judge if the dynamic elements of the cute algebra type are same(Int32 or Int64). If so, generate the corresponding C type. Otherwise, refuse to generate the argument @@ -65,55 +81,50 @@ class CuteCHeaderGenerator(CHeaderGenerator): """ return "" - def _generate_kernel_metadata( + def _generate_kernel_module( self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str ): """ - Generate the kernel metadata for the compiled function. + Generate the kernel module for the compiled function. """ - kernel_metadata_struct = f""" + kernel_module_struct = f""" typedef struct {{ - CUlibrary module; -}} {symbol_prefix}_Kernel_Metadata_t; + cudaLibrary_t module; +}} {symbol_prefix}_Kernel_Module_t; """ - kernel_metadata_load = f""" + kernel_module_load = f""" #ifdef __cplusplus extern "C" {{ #endif void _mlir_{symbol_prefix}_cuda_init(void **); -void _mlir_{symbol_prefix}_cuda_load(void **); -static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{ - CUlibrary *libraryPtr = &(metadata->module); - int32_t ret; +void _mlir_{symbol_prefix}_cuda_load_to_device(void **); +static inline void {symbol_prefix}_Kernel_Module_Load({symbol_prefix}_Kernel_Module_t *module) {{ + cudaLibrary_t *libraryPtr = &(module->module); + cudaError_t ret; struct {{ - CUlibrary **libraryPtr; - int32_t *ret; + cudaLibrary_t **libraryPtr; + cudaError_t *ret; }} initArgs = {{&libraryPtr, &ret}}; _mlir_{symbol_prefix}_cuda_init((void **)(&initArgs)); - {dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret)); + {dsl_name}_CUDA_ERROR_CHECK(ret); + int32_t device_id = 0; struct {{ - CUlibrary *library; - int32_t *ret; - }} loadArgs = {{libraryPtr, &ret}}; - _mlir_{symbol_prefix}_cuda_load((void **)(&loadArgs)); - {dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret)); - CUdevice device; - {dsl_name}_CUDA_ERROR_CHECK(cuCtxGetDevice(&device)); - int max_shared_memory_per_block_optin; - {dsl_name}_CUDA_ERROR_CHECK(cuDeviceGetAttribute(&max_shared_memory_per_block_optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device)); - unsigned int num_kernels; - {dsl_name}_CUDA_ERROR_CHECK(cuLibraryGetKernelCount(&num_kernels, metadata->module)); - CUkernel *kernels = (CUkernel *)malloc(num_kernels * sizeof(CUkernel)); - {dsl_name}_CUDA_ERROR_CHECK(cuLibraryEnumerateKernels(kernels, num_kernels, metadata->module)); - for (unsigned int i = 0; i < num_kernels; i++) {{ - {dsl_name}_CUDA_ERROR_CHECK(cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, max_shared_memory_per_block_optin, kernels[i], device)); + cudaLibrary_t **library; + int32_t *device_id; + cudaError_t *ret; + }} loadArgs = {{&libraryPtr, &device_id, &ret}}; + int32_t device_count; + {dsl_name}_CUDA_ERROR_CHECK(cudaGetDeviceCount(&device_count)); + for (int32_t i = 0; i < device_count; i++) {{ + device_id = i; + _mlir_{symbol_prefix}_cuda_load_to_device((void **)(&loadArgs)); + {dsl_name}_CUDA_ERROR_CHECK(ret); }} - free(kernels); }} """ - kernel_metadata_unload = f""" -static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{ - {dsl_name}_CUDA_ERROR_CHECK(cuLibraryUnload(metadata->module)); + kernel_module_unload = f""" +static inline void {symbol_prefix}_Kernel_Module_Unload({symbol_prefix}_Kernel_Module_t *module) {{ + {dsl_name}_CUDA_ERROR_CHECK(cudaLibraryUnload(module->module)); }} #ifdef __cplusplus @@ -121,7 +132,7 @@ static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel #endif """ - return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload + return kernel_module_struct + kernel_module_load + kernel_module_unload def _generate_arguments( self, @@ -174,13 +185,13 @@ typedef struct {{ elif isinstance(arg_type, NumericMeta): arguments.append(self._generate_numeric_argument(arg_name, arg_type)) packed_args.append("&" + arg_name) - elif is_cute_algebra_type(arg_type): + elif is_cute_algebra_type(arg_type) or isinstance(arg, (tuple, list)): c_type = self._get_cute_algebra_type(arg_type, arg) arguments.append(f"{c_type}*{arg_name}") for i in range(self._count_dynamic_expression(arg)): packed_args.append("&" + arg_name + "[" + str(i) + "]") elif isclass(arg_type) and issubclass(arg_type, cuda.CUstream): - arguments.append("CUstream " + arg_name) + arguments.append("cudaStream_t " + arg_name) packed_args.append("&" + arg_name) else: raise DSLRuntimeError( @@ -191,24 +202,38 @@ typedef struct {{ def _generate_wrapper_function( self, + dsl_name: str, symbol_prefix: str, args_spec: ExecutionArgs, function_name: str, kernel_info: Dict[str, List], - dynamic_args: list, - dynamic_kwargs: dict, + c_header_arguments: CHeaderArguments, ): """ Generate the wrapper function for the compiled function which is provided to users as the entry point. It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode. """ # 1. Get the name of the function wrapper - wrapper_function_name = f"{symbol_prefix}_wrapper" + wrapper_function_name = f"{dsl_name.lower()}_{symbol_prefix}_wrapper" capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}" # 2. Generate the signature of the wrapper function - arguments, packed_args, declarations = self._generate_arguments( - symbol_prefix, args_spec, dynamic_args, dynamic_kwargs - ) + if c_header_arguments.error_msg is not None: + raise DSLRuntimeError( + f"Error generating c header arguments: {c_header_arguments.error_msg}" + ) + arguments = [ + arg.replace(c_header_arguments.dummy_prefix_name, symbol_prefix) + for arg in c_header_arguments.arguments + ] + packed_args = [ + arg.replace(c_header_arguments.dummy_prefix_name, symbol_prefix) + for arg in c_header_arguments.packed_args + ] + declarations = [ + declaration.replace(c_header_arguments.dummy_prefix_name, symbol_prefix) + for declaration in c_header_arguments.declarations + ] + # 3. Get the return type of the wrapper function. # Note that this requires the return type to be properly annotated in python. return_type = args_spec.args_spec.annotations.get("return", None) @@ -227,7 +252,7 @@ extern "C" #endif void {capi_function_name}(void **args, int32_t num_args); -static inline {return_type} {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{ +static inline {return_type} {wrapper_function_name}({symbol_prefix}_Kernel_Module_t *module, {", ".join(arguments)}) {{ {return_type} ret; void *args[{len(packed_args) + 1}] = {{ {", ".join(packed_args)}, diff --git a/python/CuTeDSL/cutlass/cute/export/export.py b/python/CuTeDSL/cutlass/cute/export/export.py new file mode 100644 index 00000000..fa27ce8b --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/export/export.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import os +import pickle +import copy +from ..typing import IntTuple, Shape, Stride, Coord, Tile +from inspect import FullArgSpec + +from cutlass.base_dsl.export import ( + ArgsSpecProcessor, +) + +cute_algebra_types_dump = { + IntTuple: "IntTuple", + Shape: "Shape", + Stride: "Stride", + Coord: "Coord", + Tile: "Tile", +} +cute_algebra_types_load = { + "IntTuple": IntTuple, + "Shape": Shape, + "Stride": Stride, + "Coord": Coord, + "Tile": Tile, +} + + +class CuteArgsSpecProcessor(ArgsSpecProcessor): + def dumps(self, args_spec: FullArgSpec) -> bytes: + new_args_spec = copy.deepcopy(args_spec) + for arg, arg_type in new_args_spec.annotations.items(): + if arg_type in cute_algebra_types_dump.keys(): + new_args_spec.annotations[arg] = cute_algebra_types_dump[arg_type] + return pickle.dumps(new_args_spec) + + def loads(self, args_spec_bytes: bytes) -> FullArgSpec: + args_spec = pickle.loads(args_spec_bytes) + for arg, arg_type in args_spec.annotations.items(): + if arg_type in cute_algebra_types_load.keys(): + args_spec.annotations[arg] = cute_algebra_types_load[arg_type] + return args_spec + + +# This is the version of the object file. It is used to check the version of the object file is compatible with the current dsl version or not. +object_file_version = "1.1" diff --git a/python/CuTeDSL/cutlass/cute/export/load.py b/python/CuTeDSL/cutlass/cute/export/load.py new file mode 100644 index 00000000..00e1a9b2 --- /dev/null +++ b/python/CuTeDSL/cutlass/cute/export/load.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from cutlass.base_dsl.common import DSLRuntimeError + + +def version_checker(version: str) -> bool: + """Check the version of the object file is compatible with the current dsl version or not.""" + if version not in ["1.0", "1.1"]: + raise DSLRuntimeError("Incompatible version: " + version) + return True diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/__init__.py b/python/CuTeDSL/cutlass/cute/nvgpu/__init__.py index ca08ecd4..c3e70442 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/__init__.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/__init__.py @@ -17,10 +17,17 @@ from . import tcgen05 from .common import * from .helpers import * +from . import common +from . import helpers + # __all__ is required here for documentation generation __all__ = [ - "OpError", - "MmaUniversalOp", - "CopyUniversalOp", + *common.__all__, + *helpers.__all__, + # submodules With namespace + "warp", + "cpasync", + "warpgroup", + "tcgen05", ] diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/common.py b/python/CuTeDSL/cutlass/cute/nvgpu/common.py index e2c5e830..3fa1a321 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/common.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/common.py @@ -20,9 +20,19 @@ from cutlass._mlir import ir from .. import atom from ..typing import Float16, Float32, Float64, Numeric -from cutlass import cute +__all__ = [ + "OpError", + "MmaUniversalOp", + "MmaUniversalTrait", + "CopyUniversalOp", + "CopyUniversalTrait", + "MemoryOrder", + "MemoryScope", + "CacheEvictionPriority", +] + class OpError(DSLBaseError): """ An exception class for Op construction errors. @@ -83,7 +93,7 @@ class MmaUniversalOp(atom.MmaOp): self.abacc_dtype.mlir_type, self.abacc_dtype.mlir_type, ) - return MmaUniversalTrait(cute.make_atom(atom_ty, loc=loc, ip=ip)) + return MmaUniversalTrait(atom.make_atom(atom_ty, loc=loc, ip=ip)) def _verify_fragment_A(self, input, *, loc=None, ip=None): pass @@ -140,6 +150,23 @@ class MemoryScope(enum.Enum): return self.value +class CacheEvictionPriority(enum.Enum): + EVICT_NORMAL = _cute_ir.CacheEvictionPriority.EVICT_NORMAL + EVICT_FIRST = _cute_ir.CacheEvictionPriority.EVICT_FIRST + EVICT_LAST = _cute_ir.CacheEvictionPriority.EVICT_LAST + EVICT_UNCHANGED = _cute_ir.CacheEvictionPriority.EVICT_UNCHANGED + NO_ALLOCATE = _cute_ir.CacheEvictionPriority.NO_ALLOCATE + + def __str__(self) -> str: + return f"{self.__class__.__name__}.{self.name}" + + def __repr__(self) -> str: + return f"<{self.__class__.__name__}.{self.name}>" + + def _to_ir(self) -> _cute_ir.CacheEvictionPriority: + return self.value + + @dataclass(frozen=True) class CopyUniversalOp(atom.CopyOp): """ @@ -150,7 +177,12 @@ class CopyUniversalOp(atom.CopyOp): .. code-block:: python op = cute.nvgpu.CopyUniversalOp() - atom = cute.make_copy_atom(op, tensor_dtype, num_bits_per_copy=64) + atom = cute.make_copy_atom( + op, + tensor_dtype, + num_bits_per_copy=64, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_NORMAL + ) - ``tensor_dtype`` is the data type used to build the reference TV Layout (either the source \ or the destination TV Layout) in unit of tensor elements and is used for partitioning by \ @@ -158,6 +190,12 @@ class CopyUniversalOp(atom.CopyOp): - ``num_bits_per_copy`` is a kw argument specifying the number of bits to copy per Atom \ execution. This can be larger than the width of the above data type. When not provided, \ the compiler will do a best effort at auto-vectorizing. + - ``l1c_evict_priority`` is a kw argument specifying the L1 cache eviction priority hint for \ + the copy operation. Defaults to ``EVICT_NORMAL`` if not provided. + - ``invariant`` is a kw argument specifying whether the load is invariant (read-only data \ + that never changes). This enables compiler optimizations like instruction reordering. \ + Defaults to ``False`` if not provided. + """ def __str__(self) -> str: @@ -167,25 +205,27 @@ class CopyUniversalOp(atom.CopyOp): self, copy_internal_type: Type[Numeric], *, + num_bits_per_copy: int = 0, + memory_order: MemoryOrder = MemoryOrder.WEAK, + memory_scope: MemoryScope = MemoryScope.CTA, + l1c_evict_priority: CacheEvictionPriority = CacheEvictionPriority.EVICT_NORMAL, + invariant: bool = False, loc=None, ip=None, - **kwargs, ) -> "CopyUniversalTrait": - num_bits_per_copy = kwargs.get("num_bits_per_copy", 0) - memory_order = kwargs.get("memory_order", MemoryOrder.WEAK) - memory_scope = kwargs.get("memory_scope", MemoryScope.CTA) - if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0): + if not isinstance(num_bits_per_copy, int) or num_bits_per_copy < 0: raise ValueError( - "expects a 'num_bits_per_copy' kw argument of type int that is non-negative " - f"when creating a copy Atom for {self.__class__.__name__}" + f"'num_bits_per_copy' must be a non-negative int when creating a copy Atom for {self.__class__.__name__!r}" ) - ty = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get( + atom_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get( copy_internal_type.mlir_type, num_bits_per_copy, memory_order._to_ir(), memory_scope._to_ir(), + l1c_evict_priority._to_ir(), + invariant, ) - return CopyUniversalTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return CopyUniversalTrait(atom.make_atom(atom_type, loc=loc, ip=ip)) class CopyUniversalTrait(atom.Trait): diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py index 09192497..fbf9d863 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py @@ -13,15 +13,14 @@ import enum from dataclasses import dataclass from typing import Optional, Type -from cutlass import cute from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir +from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp from cutlass._mlir import ir -from ...atom import CopyOp, Trait -from ...tensor import ReductionOp +from ...atom import CopyOp, Trait, make_atom from ...typing import Int16, Int64, Pointer, Integer, Numeric from ..common import OpError from ..tcgen05.mma import CtaGroup @@ -29,7 +28,7 @@ from ..tcgen05.mma import CtaGroup #################################################################################################### # -# Aynchronous copies +# Asynchronous copies # #################################################################################################### @@ -96,7 +95,7 @@ class CopyG2SOp(CopyOp): ty = _cute_nvgpu_ir.CopyAtomSIMTAsyncCopyType.get( copy_internal_type.mlir_type, self.cache_mode._to_ir(), num_bits_per_copy ) - return CopyG2STrait(cute.make_atom(ty, loc=loc, ip=ip)) + return CopyG2STrait(make_atom(ty, loc=loc, ip=ip)) class CopyG2STrait(Trait): @@ -121,7 +120,16 @@ class TmaCopyOp(CopyOp): Base class for all TMA copy operations. """ - pass + def __init__(self, smem_layout: Optional[ir.Value] = None) -> None: + self.smem_layout = smem_layout + + def __extract_mlir_values__(self): + return [self.smem_layout] + + def __new_from_mlir_values__(self, values): + res = self.__class__() + res.smem_layout = values[0] + return res # @@ -129,7 +137,7 @@ class TmaCopyOp(CopyOp): # -@dataclass(frozen=True) +@dataclass class CopyBulkTensorTileG2SOp(TmaCopyOp): """ Bulk tensor asynchrnous GMEM to SMEM Copy Operation using the TMA unit. @@ -246,7 +254,7 @@ class CopyBulkTensorTileG2STrait(Trait): # -@dataclass(frozen=True) +@dataclass class CopyBulkTensorTileG2SMulticastOp(TmaCopyOp): """ Bulk tensor asynchrnous multicast GMEM to SMEM Copy Operation using the TMA unit. @@ -375,7 +383,7 @@ class CopyBulkTensorTileG2SMulticastTrait(Trait): # -@dataclass(frozen=True) +@dataclass class CopyBulkTensorTileS2GOp(TmaCopyOp): """ Bulk tensor asynchronous SMEM to GMEM Copy Operation using the TMA unit. @@ -449,7 +457,11 @@ class CopyBulkTensorTileS2GTrait(Trait): pass -@dataclass(frozen=True) +class CopyBulkTensorTileS2GTrait(Trait): + pass + + +@dataclass class CopyReduceBulkTensorTileS2GOp(TmaCopyOp): """ Bulk tensor asynchronous SMEM to GMEM Reduction Operation using the TMA unit. @@ -585,7 +597,7 @@ class CopyBulkG2SOp(CopyOp): ty = _cute_nvgpu_ir.CopyAtomBulkCopyG2SType.get( copy_internal_type.mlir_type, num_bits_per_copy, False ) - return CopyBulkG2STrait(cute.make_atom(ty, loc=loc, ip=ip)) + return CopyBulkG2STrait(make_atom(ty, loc=loc, ip=ip)) class CopyBulkG2STrait(Trait): @@ -670,7 +682,7 @@ class CopyBulkG2SMulticastOp(CopyOp): ty = _cute_nvgpu_ir.CopyAtomBulkCopyG2SType.get( copy_internal_type.mlir_type, num_bits_per_copy, True ) - return CopyBulkG2SMulticastTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return CopyBulkG2SMulticastTrait(make_atom(ty, loc=loc, ip=ip)) class CopyBulkG2SMulticastTrait(Trait): @@ -764,8 +776,248 @@ class CopyBulkS2GOp(CopyOp): ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2GType.get( copy_internal_type.mlir_type, num_bits_per_copy, False ) - return CopyBulkS2GTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return CopyBulkS2GTrait(make_atom(ty, loc=loc, ip=ip)) class CopyBulkS2GTrait(Trait): pass + + +# +# Bulk SMEM -> GMEM mask copies +# + + +@dataclass(frozen=True) +class CopyBulkS2GByteMaskOp(CopyOp): + """ + Bulk copy asynchrnous SMEM to GMEM Copy Operation with mask. + The i-th bit in the 16-bit wide byteMask operand specifies whether + the i-th byte of each 16-byte wide chunk of source data is copied to the destination. + + See the `PTX documentation `__. + """ + + def __post_init__(self) -> None: + # Arch verification + arch: Arch = CuTeDSL._get_dsl().get_arch_enum() + if not arch >= Arch.sm_100: + raise OpError( + self, + f"expects arch to be at least {Arch.sm_100.name}, but got {arch.name}", + suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + ) + + def __str__(self) -> str: + res = "cp.async SMEM -> GMEM bulk copy Operation" + return res + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "CopyBulkS2GByteMaskTrait": + num_bits_per_copy = kwargs.get("num_bits_per_copy", 0) + if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0): + raise ValueError( + "expects a 'num_bits_per_copy' kw argument of type int that is positive " + f"when creating a copy Atom for {self.__class__.__name__}" + ) + ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2GType.get( + copy_internal_type.mlir_type, num_bits_per_copy, True + ) + return CopyBulkS2GByteMaskTrait(make_atom(ty, loc=loc, ip=ip)) + + +class CopyBulkS2GByteMaskTrait(Trait): + def unpack( + self, + *, + loc=None, + ip=None, + byte_mask=None, + **kwargs, + ): + """ + Custom implementation of unpack for bulk copy store with mask. + + The bulk store with mask requires `byte_mask` keyword argument to be provided when + using `copy`. Any other kw arguments will be ignored instead of triggering an error. + """ + if not isinstance(byte_mask, Integer): + raise ValueError( + "expects a byte mask to be provided via the byte_mask kw argument" + ) + # Support for .cp_mask qualifier requires sm_100 or higher. + attr_str = f"#cute_nvgpu.atom_copy_field_bulks2g<{TMA_BYTE_MASK_FIELD_NAME}>" + attr = ir.Attribute.parse(attr_str) + val = _cute_nvgpu_ir.atom_set_value( + self.value, + attr, + Int16(byte_mask).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + return val + + +# +# Bulk SMEM CTA to Cluster copies +# + + +@dataclass(frozen=True) +class CopyBulkS2SOp(CopyOp): + """ + Bulk copy asynchrnous SMEM CTA to Cluster Copy Operation. + + See the `PTX documentation `__. + """ + + def __post_init__(self) -> None: + # Arch verification + arch: Arch = CuTeDSL._get_dsl().get_arch_enum() + if not arch >= Arch.sm_90: + raise OpError( + self, + f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}", + suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + ) + + def __str__(self) -> str: + res = "cp.async CTA -> Cluster bulk copy Operation" + return res + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "CopyBulkS2STrait": + num_bits_per_copy = kwargs.get("num_bits_per_copy", 0) + if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0): + raise ValueError( + "expects a 'num_bits_per_copy' kw argument of type int that is positive " + f"when creating a copy Atom for {self.__class__.__name__}" + ) + ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2SType.get( + copy_internal_type.mlir_type, num_bits_per_copy + ) + return CopyBulkS2STrait(make_atom(ty, loc=loc, ip=ip)) + + +class CopyBulkS2STrait(Trait): + def unpack( + self, + *, + loc=None, + ip=None, + mbar_ptr: Optional[Pointer] = None, + cta_rank: Optional[Integer] = None, + **kwargs, + ): + """ + Custom implementation of unpack for bulk copy cta to cluster. + + The bulk cta to cluster copy requires a `mbar_ptr` and `cta_rank` keyword argument to be provided + when using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error. + """ + + if not isinstance(mbar_ptr, Pointer): + raise ValueError( + "expects a pointer to an mbarrier to be provided via the mbar_ptr kw argument" + ) + if not isinstance(cta_rank, Integer): + raise ValueError( + "expects a cta rank of int32 to be provided via the cta_rank kw argument" + ) + attr_str = f"#cute_nvgpu.atom_copy_field_bulks2s<{TMA_MBAR_PTR_FIELD_NAME}>" + attr = ir.Attribute.parse(attr_str) + val = _cute_nvgpu_ir.atom_set_value( + self.value, attr, mbar_ptr.value, loc=loc, ip=ip + ) + attr_str = f"#cute_nvgpu.atom_copy_field_bulks2s<{TMA_CTA_RANK_FIELD_NAME}>" + attr = ir.Attribute.parse(attr_str) + val = _cute_nvgpu_ir.atom_set_value( + val, attr, Int32(cta_rank).ir_value(loc=loc, ip=ip), loc=loc, ip=ip + ) + return val + + +#################################################################################################### +# +# Aynchronous distributed shared memory stores +# +#################################################################################################### + +MBAR_PTR_FIELD_NAME = "mbar_ptr" + + +@dataclass(frozen=True) +class CopyDsmemStoreOp(CopyOp): + """ + Asynchronous Store operation to DSMEM with explicit synchronization. + + See the `PTX documentation `__. + """ + + def __post_init__(self) -> None: + # Arch verification + arch: Arch = CuTeDSL._get_dsl().get_arch_enum() + if not arch >= Arch.sm_90: + raise OpError( + self, + f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}", + suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + ) + + def __str__(self) -> str: + res = "st.async RMEM -> DSMEM copy Operation" + return res + + def _make_trait( + self, + copy_internal_type: Type[Numeric], + *, + loc=None, + ip=None, + **kwargs, + ) -> "CopyDsmemStoreTrait": + num_bits_per_copy = kwargs.get("num_bits_per_copy", 0) + if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0): + raise ValueError( + "expects a 'num_bits_per_copy' kw argument of type int that is non-negative " + f"when creating a copy Atom for {self.__class__.__name__}" + ) + ty = _cute_nvgpu_ir.CopyAtomDsmemStoreType.get( + copy_internal_type.mlir_type, num_bits_per_copy + ) + return CopyDsmemStoreTrait(make_atom(ty, loc=loc, ip=ip)) + + +class CopyDsmemStoreTrait(Trait): + def unpack( + self, + *, + loc=None, + ip=None, + mbar_ptr: Optional[Pointer] = None, + **kwargs, + ): + """ + Custom implementation of unpack for dsmem async copy. + + The dsmem async copy requires `mbar_ptr` keyword argument to be provided when using `cute.copy`. + Any other kw arguments will be ignored instead of triggering an error. + """ + if not isinstance(mbar_ptr, Pointer): + raise ValueError( + "expects a pointer to an mbarrier to be provided via the mbar_ptr kw argument", + ) + attr_str = f"#cute_nvgpu.atom_copy_field_dsmem_store<{MBAR_PTR_FIELD_NAME}>" + attr = ir.Attribute.parse(attr_str) + val = _cute_nvgpu_ir.atom_set_value( + self.value, + attr, + mbar_ptr.value, + loc=loc, + ip=ip, + ) + return val + + diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py index de51c1e7..b724c2cd 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py @@ -10,6 +10,7 @@ # is strictly prohibited. from typing import Optional, Tuple, Type, Union +from typing_extensions import deprecated from cutlass.cutlass_dsl import dsl_user_op @@ -39,15 +40,16 @@ from .copy import ( CopyReduceBulkTensorTileS2GNonExecTrait, ) +TMAOp = Union[ + CopyBulkTensorTileG2SOp, + CopyBulkTensorTileG2SMulticastOp, + CopyBulkTensorTileS2GOp, + CopyReduceBulkTensorTileS2GOp, +] @dsl_user_op def make_tiled_tma_atom( - op: Union[ - CopyBulkTensorTileG2SOp, - CopyBulkTensorTileG2SMulticastOp, - CopyBulkTensorTileS2GOp, - CopyReduceBulkTensorTileS2GOp, - ], + op: TMAOp, gmem_tensor: Tensor, smem_layout: Union[Layout, ComposedLayout], cta_tiler: Tiler, @@ -69,37 +71,30 @@ def make_tiled_tma_atom( this function figures out the bulk tensor asynchronous copy instruction to use with the maximum "TMA vector length" to copy tiles of the GMEM tensor to/from an SMEM buffer with the provided - layout and consistent with the provided Tiler. + layout while maintaining consistency with the provided Tiler. This function returns two results: 1. the Copy Atom - 2. the so-called TMA tensor used to map logical coordinates of the GMEM tensor to coordinates \ - that the TMA unit can consume. TMA tensors have so-called basis stride elements so that the \ - associated layout can output coordinates. Otherwise, TMA tensors can be partitioned \ - similarly to any other CuTe tensors using the algebra. + 2. a TMA tensor that maps logical coordinates of the GMEM tensor to coordinates consumed by the \ + TMA unit. TMA tensors contain basis stride elements that enable their associated layout to \ + compute coordinates. Like other CuTe tensors, TMA tensors can be partitioned. - :param op: The Copy Operation to construct an Atom for - :type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileS2GOp, CopyReduceBulkTensorTileS2GOp] + :param op: The TMA Copy Operation to construct an Atom + :type op: TMAOp :param gmem_tensor: The GMEM tensor involved in the Copy :type gmem_tensor: Tensor - :param smem_layout: The SMEM layout to construct the Copy Atom for + :param smem_layout: The SMEM layout to construct the Copy Atom :type smem_layout: Union[Layout, ComposedLayout] :param cta_tiler: The CTA Tiler to use :type cta_tiler: Tiler :param num_multicast: The multicast factor :type num_multicast: int - :param internal_type: An optional parameter for the internal data type to use when the actual data type is not supported by the TMA unit + :param internal_type: Optional internal data type to use when the tensor data type is not supported by the TMA unit :type internal_type: Type[Numeric] - :return: A Copy Atom for this Operation and the associated TMA tensor + :return: A TMA Copy Atom associated with the TMA tensor :rtype: Tuple[atom.CopyAtom, Tensor] """ - - if internal_type is not None: - if not isinstance(internal_type, NumericMeta): - raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") - internal_type = internal_type.mlir_type - cta_v_map = core.composition( core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip), cta_tiler, @@ -110,6 +105,26 @@ def make_tiled_tma_atom( if isinstance(smem_layout, core._ComposedLayout): smem_layout = smem_layout.value + # Set the smem_layout on the operation for later retrieval + op.smem_layout = ( + smem_layout.value + if isinstance(smem_layout, core._ComposedLayout) + else smem_layout + ) + + tma_format = None + if internal_type is not None: + if not isinstance(internal_type, NumericMeta): + raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") + + use_unpack = (internal_type.width == 8 and + isinstance(gmem_tensor.element_type, NumericMeta) and + gmem_tensor.element_type.width < 8) + internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type + tma_format = _cute_nvgpu_ir.TmaDataFormat( + _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) + ) + if isinstance(op, CopyBulkTensorTileG2SOp): if num_multicast != 1: raise ValueError( @@ -122,7 +137,7 @@ def make_tiled_tma_atom( cta_v_map, op._to_ir(), num_multicast=num_multicast, - internal_type=internal_type, + tma_format=tma_format, loc=loc, ip=ip, ) @@ -139,7 +154,7 @@ def make_tiled_tma_atom( cta_v_map, op._to_ir(), num_multicast=num_multicast, - internal_type=internal_type, + tma_format=tma_format, loc=loc, ip=ip, ) @@ -152,7 +167,7 @@ def make_tiled_tma_atom( gmem_tensor.value, smem_layout, cta_v_map, - internal_type=internal_type, + tma_format=tma_format, loc=loc, ip=ip, ) @@ -163,7 +178,7 @@ def make_tiled_tma_atom( smem_layout, cta_v_map, op._to_ir(), - internal_type=internal_type, + tma_format=tma_format, loc=loc, ip=ip, ) @@ -314,6 +329,8 @@ def fence_tma_desc_acquire( has_side_effects=True, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, ) @@ -342,6 +359,8 @@ def cp_fence_tma_desc_release( has_side_effects=True, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, ) @@ -358,10 +377,13 @@ def fence_tma_desc_release(*, loc=None, ip=None) -> None: has_side_effects=True, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, ) @dsl_user_op +@deprecated("`group_bulk_copy_modes` is deprecated, use `group_modes` instead") def group_bulk_copy_modes(src: Tensor, dst: Tensor, loc=None, ip=None) -> Tuple: """ Copy async bulk need group mode 0, acquiring whole tensor for bulk copy diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py index cff45047..6c786be5 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/helpers.py @@ -26,6 +26,11 @@ from .cpasync.copy import ( ) +__all__ = [ + "make_tiled_tma_atom_A", + "make_tiled_tma_atom_B", +] + #################################################################################################### # # TMA creation helpers for tcgen05 MMAs @@ -91,10 +96,6 @@ def make_tiled_tma_atom_A( """ - if internal_type is not None: - if not isinstance(internal_type, NumericMeta): - raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") - internal_type = internal_type.mlir_type check_type_in( op, [CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp], @@ -102,6 +103,13 @@ def make_tiled_tma_atom_A( "make_tiled_tma_atom_A", ) + # Set the smem_layout on the operation for later retrieval + op.smem_layout = ( + smem_layout.value + if isinstance(smem_layout, core._ComposedLayout) + else smem_layout + ) + ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip) mma_tiler_mk = (mma_tiler_mnk[0], *mma_tiler_mnk[2:]) g_tile = core.composition(ident, mma_tiler_mk, loc=loc, ip=ip) @@ -123,6 +131,19 @@ def make_tiled_tma_atom_A( if isinstance(smem_layout, core._ComposedLayout): smem_layout = smem_layout.value + tma_format = None + if internal_type is not None: + if not isinstance(internal_type, NumericMeta): + raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") + + use_unpack = (internal_type.width == 8 and + isinstance(gmem_tensor.element_type, NumericMeta) and + gmem_tensor.element_type.width < 8) + internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type + tma_format = _cute_nvgpu_ir.TmaDataFormat( + _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) + ) + # res[0] = the IR Value for the non-executable atom instance # res[1] = the IR Value for the associated TMA tensor res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( @@ -131,7 +152,7 @@ def make_tiled_tma_atom_A( cta_v_map, op._to_ir(), num_multicast=num_multicast, - internal_type=internal_type, + tma_format=tma_format, loc=loc, ip=ip, ) @@ -203,10 +224,6 @@ def make_tiled_tma_atom_B( """ - if internal_type is not None: - if not isinstance(internal_type, NumericMeta): - raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") - internal_type = internal_type.mlir_type check_type_in( op, [CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp], @@ -214,6 +231,13 @@ def make_tiled_tma_atom_B( "make_tiled_tma_atom_B", ) + # Set the smem_layout on the operation for later retrieval + op.smem_layout = ( + smem_layout.value + if isinstance(smem_layout, core._ComposedLayout) + else smem_layout + ) + ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip) mma_tiler_nk = (mma_tiler_mnk[1], *mma_tiler_mnk[2:]) g_tile = core.composition(ident, mma_tiler_nk, loc=loc, ip=ip) @@ -235,6 +259,19 @@ def make_tiled_tma_atom_B( if isinstance(smem_layout, core._ComposedLayout): smem_layout = smem_layout.value + tma_format = None + if internal_type is not None: + if not isinstance(internal_type, NumericMeta): + raise TypeError(f"internal_type must be a Numeric, but got {internal_type}") + + use_unpack = (internal_type.width == 8 and + isinstance(gmem_tensor.element_type, NumericMeta) and + gmem_tensor.element_type.width < 8) + internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type + tma_format = _cute_nvgpu_ir.TmaDataFormat( + _cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack) + ) + # res[0] = the IR Value for the non-executable atom instance # res[1] = the IR Value for the associated TMA tensor res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load( @@ -243,7 +280,7 @@ def make_tiled_tma_atom_B( cta_v_map, op._to_ir(), num_multicast=num_multicast, - internal_type=internal_type, + tma_format=tma_format, loc=loc, ip=ip, ) @@ -255,9 +292,3 @@ def make_tiled_tma_atom_B( atom.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])), res[1], ) - - -__all__ = [ - "make_tiled_tma_atom_A", - "make_tiled_tma_atom_B", -] diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py index 6a3dcaff..76b2198b 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py @@ -13,7 +13,6 @@ import enum from dataclasses import dataclass from typing import Type -from cutlass import cute from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL @@ -21,7 +20,7 @@ import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir from ..common import OpError -from ...atom import CopyOp, Trait +from ...atom import CopyOp, Trait, make_atom from ...typing import Numeric from .mma import CtaGroup @@ -188,7 +187,7 @@ class Ld16x64bOp(_LdBase): self.repeat.value, ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None, ) - return Ld16x64bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Ld16x64bTrait(make_atom(ty, loc=loc, ip=ip)) class Ld16x64bTrait(Trait): @@ -246,7 +245,7 @@ class Ld16x128bOp(_LdBase): self.repeat.value, ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None, ) - return Ld16x128bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Ld16x128bTrait(make_atom(ty, loc=loc, ip=ip)) class Ld16x128bTrait(Trait): @@ -304,7 +303,7 @@ class Ld16x256bOp(_LdBase): self.repeat.value, ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None, ) - return Ld16x256bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Ld16x256bTrait(make_atom(ty, loc=loc, ip=ip)) class Ld16x256bTrait(Trait): @@ -344,7 +343,7 @@ class Ld16x32bx2Op(_LdBase): self.repeat.value, ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None, ) - return Ld16x32bx2Trait(cute.make_atom(ty, loc=loc, ip=ip)) + return Ld16x32bx2Trait(make_atom(ty, loc=loc, ip=ip)) class Ld16x32bx2Trait(Trait): @@ -384,7 +383,7 @@ class Ld32x32bOp(_LdBase): self.repeat.value, ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None, ) - return Ld32x32bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Ld32x32bTrait(make_atom(ty, loc=loc, ip=ip)) class Ld32x32bTrait(Trait): @@ -478,7 +477,7 @@ class St16x64bOp(_StBase): self.repeat.value, ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None, ) - return St16x64bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return St16x64bTrait(make_atom(ty, loc=loc, ip=ip)) class St16x64bTrait(Trait): @@ -513,7 +512,7 @@ class St16x128bOp(_StBase): self.repeat.value, ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None, ) - return St16x128bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return St16x128bTrait(make_atom(ty, loc=loc, ip=ip)) class St16x128bTrait(Trait): @@ -548,7 +547,7 @@ class St16x256bOp(_StBase): self.repeat.value, ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None, ) - return St16x256bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return St16x256bTrait(make_atom(ty, loc=loc, ip=ip)) class St16x256bTrait(Trait): @@ -574,7 +573,7 @@ class St16x32bx2Op(_StBase): self.repeat.value, ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None, ) - return St16x32bx2Trait(cute.make_atom(ty, loc=loc, ip=ip)) + return St16x32bx2Trait(make_atom(ty, loc=loc, ip=ip)) class St16x32bx2Trait(Trait): @@ -600,7 +599,7 @@ class St32x32bOp(_StBase): self.repeat.value, ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None, ) - return St32x32bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return St32x32bTrait(make_atom(ty, loc=loc, ip=ip)) class St32x32bTrait(Trait): @@ -681,7 +680,7 @@ class Cp128x256bOp(_S2TCopyBase): self.cta_group.value, _cute_nvgpu_ir.CopyS2TBroadcast.none, ) - return Cp128x256bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Cp128x256bTrait(make_atom(ty, loc=loc, ip=ip)) class Cp128x256bTrait(Trait): @@ -707,7 +706,7 @@ class Cp128x128bOp(_S2TCopyBase): self.cta_group.value, _cute_nvgpu_ir.CopyS2TBroadcast.none, ) - return Cp128x128bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Cp128x128bTrait(make_atom(ty, loc=loc, ip=ip)) class Cp128x128bTrait(Trait): @@ -733,7 +732,7 @@ class Cp4x256bOp(_S2TCopyBase): self.cta_group.value, _cute_nvgpu_ir.CopyS2TBroadcast.none, ) - return Cp4x256bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Cp4x256bTrait(make_atom(ty, loc=loc, ip=ip)) class Cp4x256bTrait(Trait): @@ -759,7 +758,7 @@ class Cp4x32x128bOp(_S2TCopyBase): self.cta_group.value, _cute_nvgpu_ir.CopyS2TBroadcast.x4, ) - return Cp4x32x128bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return Cp4x32x128bTrait(make_atom(ty, loc=loc, ip=ip)) class Cp4x32x128bTrait(Trait): @@ -785,7 +784,7 @@ class Cp2x64x128b0213Op(_S2TCopyBase): self.cta_group.value, _cute_nvgpu_ir.CopyS2TBroadcast.lw_0213, ) - return Cp2x64x128b0213Trait(cute.make_atom(ty, loc=loc, ip=ip)) + return Cp2x64x128b0213Trait(make_atom(ty, loc=loc, ip=ip)) class Cp2x64x128b0213Trait(Trait): @@ -811,7 +810,7 @@ class Cp2x64x128b0123Op(_S2TCopyBase): self.cta_group.value, _cute_nvgpu_ir.CopyS2TBroadcast.lw_0123, ) - return Cp2x64x128b0123Trait(cute.make_atom(ty, loc=loc, ip=ip)) + return Cp2x64x128b0123Trait(make_atom(ty, loc=loc, ip=ip)) class Cp2x64x128b0123Trait(Trait): diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py index 1439571c..b70a6ea0 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py @@ -102,27 +102,17 @@ def make_smem_layout_atom( SmemLayoutAtomKind.MN_SW128_32B, ): # M/N-major layout - return core.make_composed_layout( - sw, - 0, - core.make_layout( - (num_contiguous_elems, 8), stride=(1, num_contiguous_elems) - ), - loc=loc, - ip=ip, + outer = core.make_layout( + (num_contiguous_elems, 8), stride=(1, num_contiguous_elems), loc=loc, ip=ip ) else: # K-major layout - return core.make_composed_layout( - sw, - 0, - core.make_layout( - (8, num_contiguous_elems), stride=(num_contiguous_elems, 1) - ), - loc=loc, - ip=ip, + outer = core.make_layout( + (8, num_contiguous_elems), stride=(num_contiguous_elems, 1), loc=loc, ip=ip ) + return core.make_composed_layout(sw, 0, outer, loc=loc, ip=ip) + @overload def tile_to_mma_shape( diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py index 1fd4ca52..705f366f 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py @@ -13,7 +13,6 @@ import enum from dataclasses import dataclass from typing import Type, Any -from cutlass import cute from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL, T @@ -24,9 +23,9 @@ from cutlass._mlir import ir from ..common import OpError from ... import core, atom from ...core import _pack_shape, rank, depth -from ...tensor import _Tensor from ...typing import ( Shape, + Tensor, Float4E2M1FN, Float8E8M0FNU, Float8E5M2, @@ -43,9 +42,7 @@ from ...typing import ( AddressSpace, Pointer, ) -from ...atom import Trait - -from ..warp.mma import SparseMetadataFormat +from ...atom import Trait, make_atom #################################################################################################### @@ -242,7 +239,7 @@ class MmaOp(Tcgen05MmaOp): + f"\n Instruction shape MNK = {self.shape_mnk}" ) - def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None): if input.memspace == AddressSpace.smem and isinstance( input.layout.type, _cute_ir.ComposedLayoutType ): @@ -254,7 +251,7 @@ class MmaOp(Tcgen05MmaOp): ) return True - def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None): if input.memspace == AddressSpace.smem and isinstance( input.layout.type, _cute_ir.ComposedLayoutType ): @@ -388,7 +385,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp): + f"\n Instruction shape MNK = {self.shape_mnk}" ) - def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None): if input.memspace == AddressSpace.smem and isinstance( input.layout.type, _cute_ir.ComposedLayoutType ): @@ -400,7 +397,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp): ) return True - def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None): if input.memspace == AddressSpace.smem and isinstance( input.layout.type, _cute_ir.ComposedLayoutType ): @@ -452,161 +449,6 @@ class BlockScaledMmaTraits(Trait): ) -# Base class for all tcgen05 Sparse MMA Ops with syntax `tcgen05.mma.cta_group.kind.sparse` used to factor out some internal code -@dataclass(frozen=True) -class SparseMmaOp(Tcgen05MmaOp): - a_dtype: Type[Numeric] - b_dtype: Type[Numeric] - acc_dtype: Type[Numeric] - shape_mnk: Shape - cta_group: CtaGroup - a_src: OperandSource - a_major_mode: OperandMajorMode - b_major_mode: OperandMajorMode - sparse_metadata_format: SparseMetadataFormat - - admissible_archs = Arch.filter( - lambda arch: arch.is_family_of(Arch.sm_100f) or arch.is_family_of(Arch.sm_110f) - ) - - def __post_init__(self) -> None: - # Verify arch - arch = BaseDSL._get_dsl().get_arch_enum() - if arch not in self.admissible_archs: - raise OpError( - self, - f"expects arch to be one of {self.admissible_archs}, but got {arch}", - suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", - ) - # Verify that the user provided enum values - if not isinstance(self.cta_group, CtaGroup): - raise OpError( - self, - "expects the 'cta_group' Op parameter to be a tcgen05.CtaGroup instance", - ) - if not isinstance(self.a_src, OperandSource): - raise OpError( - self, - "expects the 'a_src' Op parameter to be a tcgen05.OperandSource instance", - ) - if not isinstance(self.a_major_mode, OperandMajorMode): - raise OpError( - self, - "expects the 'a_major_mode' Op parameter to be a tcgen05.OperandMajorMode instance", - ) - if not isinstance(self.b_major_mode, OperandMajorMode): - raise OpError( - self, - "expects the 'b_major_mode' Op parameter to be a tcgen05.OperandMajorMode instance", - ) - if not isinstance(self.sparse_metadata_format, SparseMetadataFormat): - raise OpError( - self, - "expects the 'sparse_metadata_format' Op parameter to be a tcgen05.SparseMetadataFormat instance", - ) - # Verify the instruction shape - if (rank(self.shape_mnk) not in [2, 3]) or (depth(self.shape_mnk) != 1): - raise OpError( - self, - f"expected a flat rank 2 or 3 tuple for the 'shape_mnk' Op parameter, " - f"but got {self.shape_mnk}", - ) - m, n = self.shape_mnk[0], self.shape_mnk[1] - # For sparse MMA, the shape validation follows the same rules as dense MMA - # but the K dimension is typically doubled in the derived classes - if self.cta_group == CtaGroup.ONE: - if m not in [64, 128]: - raise OpError(self, f"expects the M-mode to be 64 or 128, but got {m}") - if m == 64: - if (n < 8) or (n > 256) or (n % 8 != 0): - raise OpError( - self, - f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}", - ) - elif m == 128: - if (n < 16) or (n > 256) or (n % 16 != 0): - raise OpError( - self, - f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}", - ) - else: - if m not in [128, 256]: - raise OpError(self, f"expects the M-mode to be 128 or 256, but got {m}") - if (n < 32) or (n > 256) or (n % 32 != 0): - raise OpError( - self, - f"expects the N-mode to satisfy 32 <= N <= 256 and N % 32 == 0, but got {n}", - ) - - def __str__(self) -> str: - return ( - self.__class__.descriptive_name # type: ignore - + f"\n A data type = {self.a_dtype}" - + f"\n B data type = {self.b_dtype}" - + f"\n Accumulator data type = {self.acc_dtype}" - + f"\n CTA group = {self.cta_group}" - + f"\n A source location = {self.a_src}" - + f"\n A major mode = {self.a_major_mode}" - + f"\n B major mode = {self.b_major_mode}" - + f"\n Instruction shape MNK = {self.shape_mnk}" - + f"\n Sparse metadata format = {self.sparse_metadata_format}" - ) - - def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None): - if input.memspace == AddressSpace.smem and isinstance( - input.layout.type, _cute_ir.ComposedLayoutType - ): - raise OpError( - self, - f"Expected affine layout for {self._make_trait()}'s operand A, " - f"but got composed layout instead: {input.layout}" - f"\nPlease use recast_ptr(ptr, {input.layout.inner}, element_type) operation to move swizzle to the ptr", - ) - return True - - def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None): - if input.memspace == AddressSpace.smem and isinstance( - input.layout.type, _cute_ir.ComposedLayoutType - ): - raise OpError( - self, - f"Expected affine layout for {self._make_trait()}'s operand B, " - f"but got composed layout instead: {input.layout}" - f"\nPlease use recast_ptr(ptr, {input.layout.inner}, element_type) operation to move swizzle to the ptr", - ) - return True - - -class SparseMmaTraits(Trait): - admissible_fields = [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B] - - def set(self, field, value, *, loc=None, ip=None) -> None: - if field not in self.admissible_fields: - raise ValueError( - f"expects field to be one of {self.admissible_fields}, but got {field}" - ) - field_name = ( - f"#cute_nvgpu.atom_mma_field_sm100_sparse<{field._to_ir_field_name()}>" - ) - attr = ir.Attribute.parse(field_name) - self.value = _cute_nvgpu_ir.atom_set_value( - self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip - ) - - def get(self, field, *, loc=None, ip=None) -> Any: - if field not in self.admissible_fields: - raise ValueError( - f"expects field to be one of {self.admissible_fields}, but got {field}" - ) - field_name = ( - f"#cute_nvgpu.atom_mma_field_sm100_sparse<{field._to_ir_field_name()}>" - ) - attr = ir.Attribute.parse(field_name) - return _cute_nvgpu_ir.atom_get_value( - Boolean.mlir_type, self.value, attr, loc=loc, ip=ip - ) - - # # TF32 MMA # @@ -669,7 +511,7 @@ class MmaTF32Op(MmaOp): 0, ) return MmaTF32Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), @@ -763,7 +605,7 @@ class MmaF16BF16Op(MmaOp): 0, ) return MmaF16BF16Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), @@ -780,109 +622,6 @@ class MmaF16BF16Trait(MmaTraits): pass -@dataclass(frozen=True) -class MmaF16BF16SparseOp(SparseMmaOp): - """ - F16/BF16 tcgen05 Sparse MMA Operation. - - See the `PTX documentation `__. - This Operation corresponds to the ``.kind::f16`` qualifier with sparse support. - """ - - descriptive_name = "tcgen05 F16/BF16 Sparse MMA Operation" - - def __init__( - self, - ab_dtype: Type[Numeric], - acc_dtype: Type[Numeric], - instruction_shape: Shape, - cta_group: CtaGroup, - a_src: OperandSource, - a_major_mode: OperandMajorMode, - b_major_mode: OperandMajorMode, - sparse_metadata_format: SparseMetadataFormat, - ) -> None: - super().__init__( - ab_dtype, - ab_dtype, - acc_dtype, - instruction_shape, - cta_group, - a_src, - a_major_mode, - b_major_mode, - sparse_metadata_format, - ) - self._verify() - - def _verify(self) -> None: - # Input data type verification - if self.a_dtype not in [Float16, BFloat16]: - raise OpError( - self, - "expects the 'ab_dtype' Op parameter to be one of Float16 or BFloat16", - ) - assert self.b_dtype == self.a_dtype, "a_dtype and b_dtype must be the same" - # Accumulator data type verification - if self.acc_dtype not in [Float16, Float32]: - raise OpError( - self, - "expects the 'acc_dtype' Op parameter to be one of Float16 or Float32", - ) - # Instruction shape verification - instruction_k = 32 # For sparse, K is doubled compared to dense F16/BF16 - if rank(self.shape_mnk) == 2: - object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k)) - if self.shape_mnk[2] != instruction_k: - raise OpError( - self, - f"expects the instruction extent in the K-mode to be {instruction_k}, " - f"but got {self.shape_mnk[2]}", - ) - - def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF16BF16SparseTrait": - shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) - ty = _cute_nvgpu_ir.MmaAtomSM100UMMASparseType.get( - shape_mnk.type.attribute, - self.cta_group.value, - self.a_major_mode._to_ir(), - self.b_major_mode._to_ir(), - self.a_dtype.mlir_type, - self.b_dtype.mlir_type, - self.acc_dtype.mlir_type, - T.ui8(), - self.sparse_metadata_format._to_ir(), - self.a_src._to_ir(), - 0, # cScaleExp - ) - - def get_e_ptr(): - ptr_type = _cute_ir.PtrType.get(T.ui8(), _cute_ir.AddressSpace.tmem, 8) - address_value = Int32(0).ir_value(loc=loc, ip=ip) - aligned_ty = _cute_ir.ConstrainedIntType.get(8, 32) - aligned_intptr = _cute_ir.assume(aligned_ty, address_value, loc=loc, ip=ip) - ui8_tmem_ptr = _cute_ir.inttoptr(ptr_type, aligned_intptr, loc=loc, ip=ip) - return ui8_tmem_ptr - - return MmaF16BF16SparseTrait( - cute.make_atom( - ty, - ( - Boolean(False).ir_value(loc=loc, ip=ip), - Boolean(False).ir_value(loc=loc, ip=ip), - Boolean(False).ir_value(loc=loc, ip=ip), - get_e_ptr().value, - ), - loc=loc, - ip=ip, - ) - ) - - -class MmaF16BF16SparseTrait(SparseMmaTraits): - pass - - # # I8 MMA # @@ -953,7 +692,7 @@ class MmaI8Op(MmaOp): 0, ) return MmaI8Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), @@ -1046,7 +785,7 @@ class MmaFP8Op(MmaOp): 0, ) return MmaFP8Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), @@ -1136,7 +875,7 @@ class MmaMXF8Op(BlockScaledMmaOp): self.sf_vec_size, ) return MmaMXF8Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), @@ -1222,7 +961,7 @@ class MmaMXF4Op(BlockScaledMmaOp): self.sf_vec_size, ) return MmaMXF4Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), @@ -1315,7 +1054,7 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp): self.sf_vec_size, ) return MmaMXF4NVF4Trait( - cute.make_atom( + make_atom( ty, ( Boolean(False).ir_value(loc=loc, ip=ip), diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py index 36e5a3fe..6d1e3034 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py @@ -16,9 +16,13 @@ from .mma import * # __all__ is required here for documentation generation __all__ = [ # mma.py + "Field", "MmaF16BF16Op", + "MmaMXF4Op", + "MmaMXF4NVF4Op", # copy.py "LdMatrix8x8x16bOp", + "LdMatrix16x8x8bOp", "LdMatrix16x16x8bOp", "StMatrix8x8x16bOp", "StMatrix16x8x8bOp", diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py index 26750b2b..09cde138 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py @@ -12,20 +12,20 @@ from dataclasses import dataclass from typing import Type -from cutlass import cute import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir import ir from ..common import OpError from ...core import _pack_shape -from ...typing import Numeric -from ...atom import CopyOp, Trait +from ...typing import Numeric, Optional +from ...atom import CopyOp, Trait, make_atom @dataclass(frozen=True) class BaseOp(CopyOp): transpose: bool = False num_matrices: int = 1 + unpack_bits: Optional[int] = None def __post_init__(self) -> None: if not isinstance(self.transpose, bool): @@ -41,6 +41,8 @@ class BaseOp(CopyOp): ) if self.transpose: res += "\n transposed" + if self.unpack_bits is not None: + res += f"\n unpack {self.unpack_bits}b to 8b" return res @@ -60,6 +62,8 @@ class LdMatrix8x8x16bOp(BaseOp): self, "expects the 'num_matrices' Op parameter to be one of [1,2,4]", ) + if self.unpack_bits is not None: + raise OpError(self, "Op doesn't support unpacking") def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs @@ -72,46 +76,140 @@ class LdMatrix8x8x16bOp(BaseOp): self.num_matrices, ir.UnitAttr.get() if self.transpose else None, ) - return LdMatrix8x8x16bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return LdMatrix8x8x16bTrait(make_atom(ty, loc=loc, ip=ip)) class LdMatrix8x8x16bTrait(Trait): pass +@dataclass(frozen=True) +class LdMatrix8x16x8bOp(BaseOp): + """ + 8x16 ``ldmatrix`` Operation with unpacking to 8b container. + Packed source container is 16x4b elements with 64b padding + or 16x6b elements with 32b padding (total 128b per 16 elements) + + See the `PTX documentation `__. + This operation corresponds to the ``.m8n16`` and the ``.b4x16_p64``, ``.b6x16_p32`` qualifiers. + """ + + def __post_init__(self) -> None: + super().__post_init__() + if self.transpose: + raise OpError(self, "Op doesn't support transpose") + if self.num_matrices not in [1, 2, 4]: + raise OpError( + self, + "expects the 'num_matrices' Op parameter to be one of [1,2,4]", + ) + if self.unpack_bits not in [4, 6]: + raise OpError(self, "Op unpack bits must be 4 or 6") + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "LdMatrix8x16x8bTrait": + mode = _pack_shape((8, 16), loc=loc, ip=ip) + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 + if self.unpack_bits == 6: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8 + ty = _cute_nvgpu_ir.CopyAtomLdsmType.get( + copy_internal_type.mlir_type, + mode.type.attribute, + sz_pattern, + self.num_matrices, + None, + ) + return LdMatrix8x16x8bTrait(make_atom(ty, loc=loc, ip=ip)) + + +class LdMatrix8x16x8bTrait(Trait): + pass + +@dataclass(frozen=True) +class LdMatrix16x8x8bOp(BaseOp): + """ + 16x8 8b ``ldmatrix`` Operation with transpose + + There is no direct PTX correspondance to this Op. + This actually lowers to ldmatrix with the ``.m16n16`` qualifier and + additional address and value permutations to match stmatrix.m16n8.trans. + Useful for vectorizing with Ampere-style 8x8 matrix thread-value layouts + """ + + def __post_init__(self) -> None: + super().__post_init__() + if not self.transpose: + raise OpError(self, "Op only supports transpose") + if self.num_matrices not in [2, 4]: + raise OpError( + self, + "expects the 'num_matrices' Op parameter to be one of [2,4]", + ) + if self.unpack_bits not in [None, 4, 6]: + raise OpError(self, "Op unpack bits must be 4 or 6 or None") + + def _make_trait( + self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs + ) -> "LdMatrix16x8x8bTrait": + mode = _pack_shape((16, 8), loc=loc, ip=ip) + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u8 + if self.unpack_bits == 4: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 + elif self.unpack_bits == 6: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8 + ty = _cute_nvgpu_ir.CopyAtomLdsmType.get( + copy_internal_type.mlir_type, + mode.type.attribute, + sz_pattern, + self.num_matrices, + ir.UnitAttr.get(), + ) + return LdMatrix16x8x8bTrait(make_atom(ty, loc=loc, ip=ip)) + +class LdMatrix16x8x8bTrait(Trait): + pass + @dataclass(frozen=True) class LdMatrix16x16x8bOp(BaseOp): """ - 16x16 8-bit ``ldmatrix`` Operation. - + 16x16 ``ldmatrix`` Operation with transpose and optional unpacking to 8b container. + Packed source container is 16x4b elements with 64b padding + or 16x6b elements with 32b padding (total 128b per 16 elements) + See the `PTX documentation `__. - This operation corresponds to the ``.m16n16`` and the ``.b16`` qualifiers. + This operation corresponds to the ``.m16n16`` and the ``.b4x16_p64``,``.b6x16_p32``,``.b8`` qualifiers. """ - def __init__(self, num_matrices: int) -> None: - super().__init__(transpose=True, num_matrices=num_matrices) - self._verify() - - def _verify(self): - assert self.transpose, "transpose must be True" + def __post_init__(self) -> None: + super().__post_init__() + if not self.transpose: + raise OpError(self, "Op only supports transpose") if self.num_matrices not in [1, 2]: raise OpError( self, "expects the 'num_matrices' Op parameter to be one of [1,2]", ) + if self.unpack_bits not in [None, 4, 6]: + raise OpError(self, "Op unpack bits must be 4 or 6 or None") def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs ) -> "LdMatrix16x16x8bTrait": mode = _pack_shape((16, 16), loc=loc, ip=ip) + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u8 + if self.unpack_bits == 4: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8 + elif self.unpack_bits == 6: + sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8 ty = _cute_nvgpu_ir.CopyAtomLdsmType.get( copy_internal_type.mlir_type, mode.type.attribute, - _cute_nvgpu_ir.LdsmSzPattern.u8, + sz_pattern, self.num_matrices, ir.UnitAttr.get(), ) - return LdMatrix16x16x8bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return LdMatrix16x16x8bTrait(make_atom(ty, loc=loc, ip=ip)) class LdMatrix16x16x8bTrait(Trait): @@ -134,6 +232,8 @@ class StMatrix8x8x16bOp(BaseOp): self, "expects the 'num_matrices' Op parameter to be one of [1,2,4]", ) + if self.unpack_bits is not None: + raise OpError(self, "Op doesn't support unpacking") def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs @@ -145,7 +245,7 @@ class StMatrix8x8x16bOp(BaseOp): self.num_matrices, ir.UnitAttr.get() if self.transpose else None, ) - return StMatrix8x8x16bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return StMatrix8x8x16bTrait(make_atom(ty, loc=loc, ip=ip)) class StMatrix8x8x16bTrait(Trait): @@ -161,17 +261,17 @@ class StMatrix16x8x8bOp(BaseOp): This operation corresponds to the ``m16n8`` qualifier. """ - def __init__(self, num_matrices: int) -> None: - super().__init__(transpose=True, num_matrices=num_matrices) - self._verify() - - def _verify(self): + def __post_init__(self) -> None: + super().__post_init__() + if not self.transpose: + raise OpError(self, "Op only supports transpose") if self.num_matrices not in [1, 2, 4]: - assert self.transpose, "transpose must be True" raise OpError( self, "expects the 'num_matrices' Op parameter to be one of [1,2,4]", ) + if self.unpack_bits is not None: + raise OpError(self, "Op doesn't support unpacking") def _make_trait( self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs @@ -183,7 +283,7 @@ class StMatrix16x8x8bOp(BaseOp): self.num_matrices, ir.UnitAttr.get(), ) - return StMatrix16x8x8bTrait(cute.make_atom(ty, loc=loc, ip=ip)) + return StMatrix16x8x8bTrait(make_atom(ty, loc=loc, ip=ip)) class StMatrix16x8x8bTrait(Trait): diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py index a23dc862..145d9a88 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py @@ -10,19 +10,33 @@ # is strictly prohibited. from dataclasses import dataclass -from typing import Type +from typing import Type, Any import enum from cutlass import cute +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import CuTeDSL + from ..common import OpError -from ...typing import Shape, Float16, BFloat16, Float32, Numeric +from ...typing import ( + Shape, + Float4E2M1FN, + Float8E8M0FNU, + Float8E4M3FN, + Float16, + BFloat16, + Float32, + Boolean, + Numeric, + Pointer, +) from ...core import _pack_shape -from ...tensor import _Tensor -from ...atom import MmaOp, Trait +from ...typing import Tensor +from ...atom import MmaOp, Trait, make_atom +from cutlass._mlir import ir import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir -from cutlass._mlir.dialects.cute_nvgpu import SparseMetadataFormat #################################################################################################### @@ -43,7 +57,7 @@ class WarpMmaOp(MmaOp): @dataclass(frozen=True) class MmaF16BF16Op(WarpMmaOp): """ - F16/BF16 tcgen05 MMA Operation. + F16/BF16 warp-level MMA Operation. See the `PTX documentation `__. This Operation covers the instructions using the ``.f16`` or ``.bf16`` qualifiers for the input operands. @@ -83,7 +97,7 @@ class MmaF16BF16Op(WarpMmaOp): self.ab_dtype.mlir_type, self.acc_dtype.mlir_type, ) - return MmaF16BF16Trait(cute.make_atom(ty, loc=loc, ip=ip)) + return MmaF16BF16Trait(make_atom(ty, loc=loc, ip=ip)) def __str__(self) -> str: return ( @@ -93,10 +107,10 @@ class MmaF16BF16Op(WarpMmaOp): + f"\n Instruction shape MNK = {self.shape_mnk}" ) - def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None): pass - def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None): pass @@ -104,12 +118,86 @@ class MmaF16BF16Trait(Trait): pass -class SparseMetadataFormat(enum.Enum): +# Base class for SM120 Blockscaled MMA Ops +@dataclass(frozen=True) +class MmaSM120BlockScaledOp(MmaOp): + ab_dtype: Type[Numeric] + acc_dtype: Type[Numeric] + shape_mnk: Shape + sf_type: Type[Numeric] + sf_vec_size: int + use_sf_layout_TV: bool = False + + admissible_archs = [ + "sm_120a", + ] + + def __post_init__(self) -> None: + # Verify arch + arch = CuTeDSL._get_dsl().get_arch_enum() + if not arch == Arch.sm_120a: + raise OpError( + self, + f"expects arch to be one of {self.admissible_archs}, but got {arch}", + suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture", + ) + if self.ab_dtype != Float4E2M1FN: + raise OpError( + self, + "expects the 'ab_dtype' Op parameter to be Float4E2M1FN", + ) + if self.acc_dtype != Float32: + raise OpError( + self, + "expects the 'acc_dtype' Op parameter to be Float32", + ) + if self.shape_mnk != (16, 8, 64): + raise OpError( + self, + "expects the 'shape_mnk' Op parameter to be (16,8,64)", + ) + + if self.sf_vec_size == 16: + if self.sf_type != Float8E4M3FN: + raise OpError( + self, + "expects the 'sf_type' Op parameter to be Float8E4M3FN", + ) + elif self.sf_vec_size == 32: + if self.sf_type != Float8E8M0FNU: + raise OpError( + self, + "expects the 'sf_type' Op parameter to be Float8E8M0FNU", + ) + else: + raise OpError( + self, + "expects the 'sf_vec_size' Op parameter to be 16 or 32", + ) + def __str__(self) -> str: + return ( + "warp-level MXF4/MXF4NVF4 MMA Operation" + + f"\n A/B data type = {self.ab_dtype}" + + f"\n Accumulator data type = {self.acc_dtype}" + + f"\n Instruction shape MNK = {self.shape_mnk}" + + f"\n Vector size = {self.sf_vec_size}" + + f"\n SF data type = {self.sf_type}" + ) + + def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None): + pass + + def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None): + pass + +class Field(enum.Enum): """ - An enumeration for the sparse metadata format of the MMA. + An enumeration for the fields of the MMA Atom that can be modified at runtime. """ - TID = SparseMetadataFormat.tid + ACCUMULATE = "accum_c" + SFA = "sf_a" + SFB = "sf_b" def __str__(self) -> str: return f"{self.__class__.__name__}.{self.name}" @@ -117,66 +205,144 @@ class SparseMetadataFormat(enum.Enum): def __repr__(self) -> str: return f"<{self.__class__.__name__}.{self.name}>" - def _to_ir(self) -> _cute_nvgpu_ir.SparseMetadataFormat: + def _to_ir_field_name(self) -> str: return self.value +class MmaBlockScaledTrait(Trait): + admissible_fields = [ + Field.ACCUMULATE, + Field.SFA, + Field.SFB, + ] + + def set(self, field, value, *, loc=None, ip=None) -> None: + if field not in self.admissible_fields: + raise ValueError( + f"expects field to be one of {self.admissible_fields}, but got {field}" + ) + if field == Field.ACCUMULATE: + value = Boolean(value).ir_value(loc=loc, ip=ip) + elif field in [Field.SFA, Field.SFB]: + if not isinstance(value, Pointer): + raise ValueError( + f"expects value to be a pointer for {field}, but got {type(value).__name__}" + ) + value = value.value + + field_name = f"#cute_nvgpu.atom_mma_field_sm120_block_scaled<{field._to_ir_field_name()}>" + attr = ir.Attribute.parse(field_name) + self.value = _cute_nvgpu_ir.atom_set_value( + self.value, attr, value, loc=loc, ip=ip + ) + + def get(self, field, *, loc=None, ip=None) -> Any: + if field not in [Field.ACCUMULATE]: + raise ValueError(f"the get method for {field} is not supported") + field_name = f"#cute_nvgpu.atom_mma_field_sm120_block_scaled<{field._to_ir_field_name()}>" + attr = ir.Attribute.parse(field_name) + return _cute_nvgpu_ir.atom_get_value( + Boolean.mlir_type, self.value, attr, loc=loc, ip=ip + ) + + +# +# MXF4 MMA +# + + @dataclass(frozen=True) -class MmaF16BF16SparseOp(WarpMmaOp): - ab_dtype: Type[Numeric] - acc_dtype: Type[Numeric] - shape_mnk: Shape - sparse_metadata_format: SparseMetadataFormat +class MmaMXF4Op(MmaSM120BlockScaledOp): + """ + MXF4 warp-level MMA Operation. - def __post_init__(self) -> None: - # verify field after initialization - if not isinstance(self.sparse_metadata_format, SparseMetadataFormat): - raise OpError( - self, - "expects the 'sparse_metadata_format' Op parameter to be a SparseMetadataFormat instance", - ) - # verify the instruction shape - if self.ab_dtype not in [Float16, BFloat16]: - raise OpError( - self, - "expects the 'ab_dtype' Op parameter to be one of Float16 or BFloat16", - ) - if self.acc_dtype not in [Float16, Float32]: - raise OpError( - self, - "expects the 'acc_dtype' Op parameter to be one of Float16 or Float32", - ) - if (self.ab_dtype == BFloat16) and (self.acc_dtype != Float32): - raise OpError( - self, - "expects the 'acc_dtype' Op parameter to be Float32 when 'ab_dtype' is BFloat16", - ) - if self.shape_mnk not in [(16, 8, 16), (16, 8, 32)]: - raise OpError( - self, - "expects the 'shape_mnk' Op parameter to be one of (16,8,16) or (16,8,32)", - ) + See the `PTX documentation `__. + This Operation covers the instructions using the ``.e2m1`` qualifiers for the input operands. + .kind = {.kind::mxf4}; + .scale_vec_size = {.scale_vec::2X}; + .stype = {.ue8m0}; + """ - def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF16BF16SparseTrait": + descriptive_name = "warp-level MXF4 MMA Operation" + + def __init__( + self, + ab_dtype: Type[Numeric], + acc_dtype: Type[Numeric], + sf_type: Type[Numeric], + ) -> None: + super().__init__( + ab_dtype, + acc_dtype, + (16, 8, 64), + sf_type, + 32, + ) + + def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4Trait": shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) - ty = _cute_nvgpu_ir.MmaAtomSM80SparseType.get( + ty = _cute_nvgpu_ir.MmaAtomSM120BlockScaledType.get( shape_mnk.type.attribute, + 32, + False, self.ab_dtype.mlir_type, self.ab_dtype.mlir_type, self.acc_dtype.mlir_type, - self.sparse_metadata_format._to_ir(), - ) - return MmaF16BF16SparseTrait(cute.make_atom(ty, loc=loc, ip=ip)) - - def __str__(self) -> str: - return ( - "warp-level F16/BF16 Sparse MMA Operation" - + f"\n A/B data type = {self.ab_dtype}" - + f"\n Accumulator data type = {self.acc_dtype}" - + f"\n Instruction shape MNK = {self.shape_mnk}" - + f"\n Sparse metadata format = {self.sparse_metadata_format}" + self.sf_type.mlir_type, ) + return MmaMXF4Trait(make_atom(ty, loc=loc, ip=ip)) -class MmaF16BF16SparseTrait(Trait): +class MmaMXF4Trait(MmaBlockScaledTrait): + pass + + +# +# MXF4NVF4 MMA +# + + +@dataclass(frozen=True) +class MmaMXF4NVF4Op(MmaSM120BlockScaledOp): + """ + MXF4NVF4 warp-level MMA Operation. + + See the `PTX documentation `__. + This Operation covers the instructions using the ``.e2m1`` qualifiers for the input operands. + .kind = {.kind::mxf4nvf4}; + .scale_vec_size = {.scale_vec::2X, .scale_vec::4X}; + .stype = {.ue8m0, .ue4m3}; + """ + + descriptive_name = "warp-level MXF4NVF4 MMA Operation" + + def __init__( + self, + ab_dtype: Type[Numeric], + acc_dtype: Type[Numeric], + sf_type: Type[Numeric], + ) -> None: + super().__init__( + ab_dtype, + acc_dtype, + (16, 8, 64), + sf_type, + 16, + ) + + def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4NVF4Trait": + shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip) + ty = _cute_nvgpu_ir.MmaAtomSM120BlockScaledType.get( + shape_mnk.type.attribute, + 16, + False, + self.ab_dtype.mlir_type, + self.ab_dtype.mlir_type, + self.acc_dtype.mlir_type, + self.sf_type.mlir_type, + ) + return MmaMXF4NVF4Trait(make_atom(ty, loc=loc, ip=ip)) + + +class MmaMXF4NVF4Trait(MmaBlockScaledTrait): pass diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py index a35a978f..3555790e 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py @@ -13,7 +13,6 @@ import enum from dataclasses import dataclass from typing import Type, Any -from cutlass import cute from cutlass.base_dsl.arch import Arch from cutlass.cutlass_dsl import BaseDSL, T @@ -23,9 +22,9 @@ from cutlass._mlir import ir from ..common import OpError from ...core import _pack_shape, rank, depth -from ...tensor import _Tensor from ...typing import ( Shape, + Tensor, Float16, BFloat16, Float32, @@ -38,7 +37,7 @@ from ...typing import ( Numeric, AddressSpace, ) -from ...atom import MmaOp, Trait +from ...atom import MmaOp, Trait, make_atom #################################################################################################### @@ -181,7 +180,7 @@ class MmaOp(WarpGroupMmaOp): + f"\n Instruction shape MNK = {self.shape_mnk}" ) - def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None): if input.memspace == AddressSpace.smem and isinstance( input.layout.type, _cute_ir.ComposedLayoutType ): @@ -193,7 +192,7 @@ class MmaOp(WarpGroupMmaOp): ) return True - def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None): + def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None): if input.memspace == AddressSpace.smem and isinstance( input.layout.type, _cute_ir.ComposedLayoutType ): @@ -305,12 +304,7 @@ class MmaF16BF16Op(MmaOp): self.a_src._to_ir(), ) return MmaF16BF16Trait( - cute.make_atom( - ty, - (Boolean(False).ir_value(loc=loc, ip=ip),), - loc=loc, - ip=ip, - ) + make_atom(ty, (Boolean(False).ir_value(loc=loc, ip=ip),), loc=loc, ip=ip) ) @@ -391,12 +385,7 @@ class MmaF8Op(MmaOp): self.a_src._to_ir(), ) return MmaF8Trait( - cute.make_atom( - ty, - (Boolean(False).ir_value(loc=loc, ip=ip),), - loc=loc, - ip=ip, - ) + make_atom(ty, (Boolean(False).ir_value(loc=loc, ip=ip),), loc=loc, ip=ip) ) @@ -486,12 +475,7 @@ class MmaI8Op(MmaOp): self.a_src._to_ir(), ) return MmaI8Trait( - cute.make_atom( - ty, - (Boolean(False).ir_value(loc=loc, ip=ip),), - loc=loc, - ip=ip, - ) + make_atom(ty, (Boolean(False).ir_value(loc=loc, ip=ip),), loc=loc, ip=ip) ) diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index 4b95c16a..87082656 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -25,9 +25,19 @@ import cutlass._mlir.dialects.cuda as _cuda_dialect from cutlass.cutlass_dsl import JitArgAdapterRegistry, CuTeDSL as _CuTeDSL from cutlass.base_dsl.common import DSLRuntimeError +from cutlass.base_dsl.export import ExternalBinaryModule # Local modules imports -from .typing import AddressSpace, Layout, Tensor, Pointer, Numeric, SymInt +from .typing import ( + AddressSpace, + Layout, + Tensor, + Pointer, + Numeric, + SymInt, + Float32, + TFloat32, +) from . import core from .tensor import _Tensor as CoreTensor @@ -69,7 +79,10 @@ class _Pointer(Pointer): else: self._assumed_align = assumed_align - self._c_pointer = None + self._desc = ctypes.c_void_p(int(self._pointer)) + self._c_pointer = ctypes.addressof(self._desc) + self._c_pointers_cache = [self._c_pointer] + assert int(self._pointer) % self._assumed_align == 0, ( f"pointer must be {self._assumed_align} bytes aligned" ) @@ -85,10 +98,7 @@ class _Pointer(Pointer): return self._pointer def __c_pointers__(self): - if self._c_pointer is None: - self._desc = ctypes.c_void_p(int(self._pointer)) - self._c_pointer = ctypes.addressof(self._desc) - return [self._c_pointer] + return self._c_pointers_cache def __new_from_mlir_values__(self, values): assert len(values) == 1 @@ -151,29 +161,24 @@ class _Tensor(Tensor): self._memref_desc = None self._dtype = None self._use_32bit_stride = use_32bit_stride + self._c_pointers_cache = None @property def __class__(self) -> Type[Tensor]: # Cheat to let `type(_Tensor())` to return cute.Tensor return Tensor - def lazily_load_dltensor(func): - """Decorator to lazily load the DLTensorWrapper. + def load_dltensor(self): + """Lazily load the DLTensorWrapper. - This decorator loads the DLTensorWrapper when needed, + This function loads the DLTensorWrapper when needed, avoiding overhead in the critical path of calling JIT functions. """ + if self._dltensor_wrapper is None: + self._dltensor_wrapper = _cute_ir.DLTensorWrapper( + self._dlpack_data, self._use_32bit_stride + ) - def wrapper(self, *args, **kwargs): - if self._dltensor_wrapper is None: - self._dltensor_wrapper = _cute_ir.DLTensorWrapper( - self._dlpack_data, self._use_32bit_stride - ) - return func(self, *args, **kwargs) - - return wrapper - - @lazily_load_dltensor def mark_layout_dynamic(self, leading_dim: Optional[int] = None): """Marks the tensor layout as dynamic based on the leading dimension. @@ -193,10 +198,10 @@ class _Tensor(Tensor): :return: The tensor with dynamic layout :rtype: _Tensor """ + self.load_dltensor() self._dltensor_wrapper.mark_layout_dynamic(leading_dim) return self - @lazily_load_dltensor def mark_compact_shape_dynamic( self, mode: int, @@ -208,6 +213,7 @@ class _Tensor(Tensor): :param mode: The mode of the compact shape, defaults to 0 :type mode: int :param stride_order: Consistent with `torch.Tensor.dim_order`. Defaults to None. + Indicates the order of the modes (dimensions) if the current layout were converted to row-major order. It starts from the outermost to the innermost dimension. :type stride_order: tuple[int, ...], optional @@ -228,17 +234,18 @@ class _Tensor(Tensor): Using `torch.Tensor.dim_order()` to get the stride order of the torch tensor. .. code-block:: python - a = torch.empty(3, 4) - t = cute.runtime.from_dlpack(a) - t = t.mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()) + a = torch.empty(3, 4) + t = cute.runtime.from_dlpack(a) + t = t.mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order()) """ + self.load_dltensor() self._dltensor_wrapper.mark_compact_shape_dynamic( mode, stride_order, divisibility ) return self @property - @lazily_load_dltensor def element_type(self) -> Type[Numeric]: + self.load_dltensor() if self._dtype is None: self._dtype = self._dltensor_wrapper.dtype return self._dtype @@ -274,24 +281,24 @@ class _Tensor(Tensor): self._dtype = new_type @property - @lazily_load_dltensor def memspace(self): + self.load_dltensor() return self._dltensor_wrapper.address_space @property - @lazily_load_dltensor def size_in_bytes(self) -> int: + self.load_dltensor() return self._dltensor_wrapper.size_in_bytes() @property - @lazily_load_dltensor def mlir_type(self) -> ir.Type: + self.load_dltensor() return self._dltensor_wrapper.get_type( self.element_type.mlir_type, self._assumed_align ) - @lazily_load_dltensor def __str__(self) -> str: + self.load_dltensor() return f"Tensor<0x{self._dltensor_wrapper.str}>" def __repr__(self): @@ -304,8 +311,8 @@ class _Tensor(Tensor): raise TypeError("runtime._Tensor is not indexable") @property - @lazily_load_dltensor def iterator(self): + self.load_dltensor() return _Pointer( self._dltensor_wrapper.data_ptr, self.element_type, @@ -320,13 +327,13 @@ class _Tensor(Tensor): ) @property - @lazily_load_dltensor def shape(self): + self.load_dltensor() return self._dltensor_wrapper.shape @property - @lazily_load_dltensor def stride(self): + self.load_dltensor() strides = self._dltensor_wrapper.stride if strides is None: strides = itertools.accumulate( @@ -356,28 +363,30 @@ class _Tensor(Tensor): raise TypeError("fill function is not supported in runtime") @property - @lazily_load_dltensor def data_ptr(self): + self.load_dltensor() return self._dltensor_wrapper.data_ptr @property - @lazily_load_dltensor def dynamic_shapes_mask(self): """Get the mask of dynamic shapes in the tensor.""" + self.load_dltensor() return self._dltensor_wrapper.get_dynamic_shapes_mask() @property - @lazily_load_dltensor def dynamic_strides_mask(self): """Get the mask of dynamic strides in the tensor.""" + self.load_dltensor() return self._dltensor_wrapper.get_dynamic_strides_mask() - @lazily_load_dltensor def __c_pointers__(self): - self._memref_desc = self._dltensor_wrapper.build_memref_desc( - self._assumed_align - ) - return [_cute_ir.pycapsule_get_pointer(self._memref_desc)] + if self._c_pointers_cache is None: + self.load_dltensor() + self._memref_desc = self._dltensor_wrapper.build_memref_desc( + self._assumed_align + ) + self._c_pointers_cache = [_cute_ir.pycapsule_get_pointer(self._memref_desc)] + return self._c_pointers_cache def __get_mlir_types__(self): return [self.mlir_type] @@ -480,6 +489,12 @@ class _FakeCompactTensor(Tensor): def stride(self): return self._stride + @property + def leading_dim(self): + for dim, order in enumerate(self._stride_order): + if order == 0: + return dim + @property def dynamic_shapes_mask(self): return tuple(1 if isinstance(e, SymInt) else 0 for e in self._shape) @@ -618,7 +633,8 @@ def make_fake_compact_tensor( :param shape: Shape of the tensor. :type shape: tuple[int, ...] :param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions. - If None, the default layout is col-major. Otherwise, it should be a permutation of the dimension indices. + If None, the default layout is left-to-right order (known as column-major order for flatten layout). + Otherwise, it should be a permutation order of the dimension indices. :type stride_order: tuple[int, ...], optional :param memspace: Memory space where the fake tensor resides. Optional. :type memspace: str, optional @@ -646,6 +662,9 @@ def make_fake_compact_tensor( # Compiled function will take a tensor with the type: # tensor o (100,?{div=8}):(?{i32 div=8},1)> compiled_foo = cute.compile(foo, x) + + # Default stride order is left-to-right order: (1, 8) + y = make_fake_compact_tensor(cutlass.Float32, (8, 3)) """ return _FakeCompactTensor( @@ -730,6 +749,7 @@ def from_dlpack( use_32bit_stride=False, *, enable_tvm_ffi=False, + force_tf32=False, ) -> Tensor: """Convert from tensor object supporting __dlpack__() to a CuTe Tensor. @@ -745,6 +765,8 @@ def from_dlpack( :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to False. When True, the tensor will be converted to a TVM-FFI function compatible tensor. :type enable_tvm_ffi: bool, optional + :param force_tf32: Whether to force the element type to TFloat32 if the element type is Float32. + :type force_tf32: bool, optional :return: A CuTe Tensor object :rtype: Tensor @@ -764,12 +786,15 @@ def from_dlpack( # If the environment variable `CUTE_DSL_ENABLE_TVM_FFI` is set to True, the tensor will be converted to # a TVM-FFI function compatible tensor. enable_tvm_ffi = enable_tvm_ffi or _CuTeDSL._get_dsl().envar.enable_tvm_ffi - return _Tensor( + res = _Tensor( tensor_dlpack, assumed_align=assumed_align, use_32bit_stride=use_32bit_stride, enable_tvm_ffi=enable_tvm_ffi, ) + if force_tf32 and res.element_type == Float32: + res.element_type = TFloat32 + return res def make_ptr( @@ -851,15 +876,21 @@ class TensorAdapter: def __init__(self, arg): self._arg = from_dlpack(arg).mark_layout_dynamic() + self._c_pointers_cache = None + self._mlir_types_cache = None def __new_from_mlir_values__(self, values): return self._arg.__new_from_mlir_values__(values) def __c_pointers__(self): - return self._arg.__c_pointers__() + if self._c_pointers_cache is None: + self._c_pointers_cache = self._arg.__c_pointers__() + return self._c_pointers_cache def __get_mlir_types__(self): - return self._arg.__get_mlir_types__() + if self._mlir_types_cache is None: + self._mlir_types_cache = self._arg.__get_mlir_types__() + return self._mlir_types_cache def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: @@ -872,7 +903,7 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: :rtype: list """ - def _get_cuda_dialect_runtime_path(): + def _get_cute_dsl_runtime_path(): libs = get_prefix_dsl_libs("CUTE_DSL") if libs is None: return None @@ -884,15 +915,15 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: libs = libs.split(":") for path in libs: - if path.endswith("libcuda_dialect_runtime.so"): + if path.endswith("libcute_dsl_runtime.so"): return path return None libs = [] - cuda_dialect_runtime_path = _get_cuda_dialect_runtime_path() - if cuda_dialect_runtime_path: - libs.append(cuda_dialect_runtime_path) + cute_dsl_runtime_path = _get_cute_dsl_runtime_path() + if cute_dsl_runtime_path: + libs.append(cute_dsl_runtime_path) if enable_tvm_ffi: import tvm_ffi @@ -905,8 +936,8 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: _LOAD_MODULE_LIBS_CACHE = [] -def load_module(file_path: str, *, enable_tvm_ffi: bool = True): - """Load a module from a file path. Today only support TVM-FFI module. +def load_module(file_path: str, *, enable_tvm_ffi: bool = False): + """Load a module from a file path. :param file_path: The path to the module file :type file_path: str @@ -937,9 +968,7 @@ def load_module(file_path: str, *, enable_tvm_ffi: bool = True): # compatible with tvm-ffi < 0.1.6 return tvm_ffi.load_module(file_path) else: - raise DSLRuntimeError( - "Unimplemented, please load the module with enable_tvm_ffi=True." - ) + return ExternalBinaryModule(file_path) # ------------------------------------------------------------------------- # Try to register_jit_arg_adapter for TensorAdapter diff --git a/python/CuTeDSL/cutlass/cute/tensor.py b/python/CuTeDSL/cutlass/cute/tensor.py index d24295a8..a8f0ff7b 100644 --- a/python/CuTeDSL/cutlass/cute/tensor.py +++ b/python/CuTeDSL/cutlass/cute/tensor.py @@ -28,6 +28,30 @@ from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir from cutlass._mlir.dialects import vector, arith +from .typing import ( + Numeric, + Integer, + Boolean, + Int4, + Uint8, + Int8, + Int32, + Int64, + BFloat16, + IntTuple, + Coord, + Shape, + Stride, + Pointer, + Layout, + ComposedLayout, + Tensor, + AddressSpace, + is_integer, + is_int_tuple, + as_numeric, +) + from .core import ( _unpack_x_tuple, _pack_int_tuple, @@ -82,6 +106,29 @@ from .tuple import transform_leaf, product, product_like, flatten_to_tuple from .arch import cvt_i8_bf16_intrinsic, cvt_i4_bf16_intrinsic, cvt_f4e2m1_f16_intrinsic +__all__ = [ + "TensorSSA", + "ReductionOp", + "make_tensor", + "make_identity_tensor", + "make_fragment", + "make_fragment_like", + "make_rmem_tensor_like", + "make_rmem_tensor", + "recast_tensor", + "domain_offset", + "print_tensor", + "full", + "full_like", + "empty_like", + "ones_like", + "zeros_like", + "where", + "any_", + "all_", +] + + @ir.register_value_caster(_cute_ir.MemRefType.get_static_typeid(), replace=True) @ir.register_value_caster(_cute_ir.CoordTensorType.get_static_typeid(), replace=True) @ir.register_value_caster( @@ -132,8 +179,9 @@ class _Tensor(Tensor): iter_val = _cute_ir.get_iter(self.value, loc=loc, ip=ip) if isinstance(iter_val, Pointer): self._iterator = iter_val - elif isinstance(iter_val.type, _cute_ir.IntTupleType): - self._iterator = _unpack_x_tuple(iter_val) + elif isinstance(iter_val.type, _cute_ir.ArithTupleIteratorType): + itup_val = _cute_ir.deref_arith_tuple_iter(iter_val) + self._iterator = _unpack_x_tuple(itup_val) elif isinstance(iter_val, ir.Value): # Example: SMEM descriptor iterator, not well supported today self._iterator = iter_val @@ -152,6 +200,9 @@ class _Tensor(Tensor): else: raise TypeError(f"unsupported iterator type, got {type(self.iterator)}") + def __repr__(self): + return self.__str__() + def __str__(self): from .core import pretty_str @@ -257,7 +308,8 @@ class _Tensor(Tensor): res = _cute_ir.get_iter( slice_(self, crd, loc=loc, ip=ip).value, loc=loc, ip=ip ) - return _unpack_x_tuple(res, loc=loc, ip=ip) + itup_val = _cute_ir.deref_arith_tuple_iter(res) + return _unpack_x_tuple(itup_val) else: self._check_can_load_store() self._check_can_dereference() @@ -387,9 +439,10 @@ class _Tensor(Tensor): return _cute_ir.get_layout(self.value, loc=loc, ip=ip) @property + @dsl_user_op @lru_cache_ir() - def shape(self) -> Shape: - return self.layout.shape + def shape(self, *, loc=None, ip=None) -> Shape: + return self.layout.shape_method(loc=loc, ip=ip) @property @lru_cache_ir() @@ -633,7 +686,7 @@ def make_tensor( """ if isinstance(layout, _ComposedLayoutWithInnerFunc): raise ValueError( - "CuTe DSL tensor does not support composed layouts with inner functions: {layout}" + f"CuTe DSL tensor does not support composed layouts with inner functions: {layout}" ) if not isinstance(layout, (Layout, ComposedLayout)): @@ -644,8 +697,12 @@ def make_tensor( res_ty = None if is_integer(iterator) or isinstance(iterator, tuple): - iterator = _pack_int_tuple(iterator, loc=loc, ip=ip) - res_ty = _cute_ir.CoordTensorType.get(iterator.type, layout.type) + itup_val = _pack_int_tuple(iterator, loc=loc, ip=ip) + iter_ty = _cute_ir.ArithTupleIteratorType.get(itup_val.type) + iterator = _cute_ir.make_arith_tuple_iter( + iter=iter_ty, value=itup_val, loc=loc, ip=ip + ) + res_ty = _cute_ir.CoordTensorType.get(itup_val.type, layout.type) elif isinstance(iterator, Pointer): iterator = iterator.value res_ty = _cute_ir.MemRefType.get(iterator.type, layout.type) @@ -773,7 +830,7 @@ def make_fragment( @dsl_user_op def make_rmem_tensor_like( - src: Union[Layout, ComposedLayout, Tensor], + src: Union[Layout, ComposedLayout, Tensor, "TensorSSA"], dtype: Optional[Type[Numeric]] = None, *, loc=None, @@ -826,7 +883,7 @@ def make_rmem_tensor_like( create register storage for intermediate results. """ - if not isinstance(src, (Layout, ComposedLayout, Tensor)): + if not isinstance(src, (Layout, ComposedLayout, Tensor, TensorSSA)): raise TypeError( f"src must be a Layout or ComposedLayout or Tensor, got {type(src)}" ) @@ -844,6 +901,9 @@ def make_rmem_tensor_like( else: res_dtype = dtype or src.element_type src_layout = src.layout + elif isinstance(src, TensorSSA): + res_dtype = dtype or src.element_type + src_layout = make_layout(src.shape, loc=loc, ip=ip) else: if dtype is None: raise ValueError("dtype must be provided when src is a layout") @@ -918,7 +978,7 @@ def domain_offset(coord: Coord, tensor: Tensor, *, loc=None, ip=None) -> Tensor: ip=ip, ) elif is_integer(tensor.iterator) or isinstance(tensor.iterator, tuple): - new_iter = _cute_ir.add_offset( + new_iter = _cute_ir.tuple_add( _pack_int_tuple(tensor.iterator, loc=loc, ip=ip), _pack_int_tuple(offset, loc=loc, ip=ip), loc=loc, @@ -1136,10 +1196,12 @@ class TensorSSA(cutlass_arith.ArithValue): def _apply_op( self, op, other: "TensorSSA", flip=False, *, loc, ip ) -> "TensorSSA": ... + @overload def _apply_op( self, op, other: cutlass_arith.ArithValue, flip=False, *, loc, ip ) -> "TensorSSA": ... + @overload def _apply_op( self, op, other: Union[int, float, bool], flip=False, *, loc, ip diff --git a/python/CuTeDSL/cutlass/cute/testing.py b/python/CuTeDSL/cutlass/cute/testing.py index 8cbc339a..9b99209f 100644 --- a/python/CuTeDSL/cutlass/cute/testing.py +++ b/python/CuTeDSL/cutlass/cute/testing.py @@ -22,12 +22,24 @@ import cuda.bindings.runtime as cuda_runtime import cutlass import cutlass.base_dsl.jit_executor -from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op +import cutlass.cutlass_dsl.cuda_jit_executor +from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op, const_expr -from .typing import Numeric, Int8, Boolean +from .typing import Numeric, Int8, Boolean, Tensor, Layout, Shape -import cutlass.cute as cute -from cutlass.cute import nvgpu +from . import nvgpu +from .core import recast_layout, make_layout, composition, get, rank, size, zipped_divide +from .tuple import elem_less +from .tensor import ( + make_rmem_tensor, + recast_tensor, + make_identity_tensor, + TensorSSA, + _Tensor, +) +from .atom import make_copy_atom +from .algorithm import copy +from .runtime import from_dlpack from cutlass._mlir.dialects import builtin, cf, nvvm, vector @@ -37,14 +49,252 @@ def assert_(cond, msg=None, *, loc=None, ip=None): cf.assert_(Boolean(cond).ir_value(), msg if msg else "", loc=loc, ip=ip) -def _maybe_recast_tensor_from_f4(src: cute.Tensor, tv_layout: cute.Layout): +################################################ +# Runtime Assertion Helper Utilities For Testing +################################################ + + +class AssertionError(RuntimeError): + """Custom assertion error for runtime assertions.""" + + pass + + +class Assertion: + """Base class for runtime assertion.""" + + pass + + +class _CompileTimeAssertion(Assertion): + """Compile-time assertion helper that tracks assertion results during execution. + + This assertion is used internally when RuntimeAssertion is passed through + JIT compilation. It stores assertion results in a tensor and provides compile-time + tracking of assertion results. + """ + + def __init__( + self, + tensor: _Tensor, + num_assertions: int = 1, + msgs=None, + device=None, + disable: bool = False, + init_value: bool = False, + used_indices: set = None, + ): + """Initialize _CompileTimeAssertion. + + :param tensor: Tensor to store assertion results + :param num_assertions: Number of assertions to support + :param msgs: List of assertion messages + :param device: Device to run assertions on + :param disable: If True, assertions are disabled + :param init_value: Initial value for assertion tensor + :param used_indices: Set of used assertion indices + """ + if msgs is None: + msgs = [] + self._tensor = tensor + self._num_assertions = num_assertions + self._device = device + self._disable = disable + self._msgs = msgs + self._init_value = init_value + self._used_indices = used_indices + + def __new_from_mlir_values__(self, values): + if self._disable: + return _CompileTimeAssertion( + None, + self._num_assertions, + self._msgs, + self._device, + self._disable, + self._init_value, + self._used_indices, + ) + return _CompileTimeAssertion( + _Tensor(values[0], dtype=Boolean), + self._num_assertions, + self._msgs, + self._device, + self._disable, + self._init_value, + self._used_indices, + ) + + def __extract_mlir_values__(self): + if self._disable: + return [] + return self._tensor.__extract_mlir_values__() + + @dsl_user_op + @CuTeDSL.jit + def store(self, idx: Constexpr, pred: Boolean, msg: str = "", *, loc=None, ip=None): + """Assert a predicate condition. + + :param idx: Assertion index + :type idx: int + :param pred: Predicate condition to assert + :type pred: Boolean + :param msg: Assertion message + :type msg: str, optional + :param loc: MLIR location information for debugging, defaults to None + :type loc: optional + :param ip: MLIR insertion point for code generation, defaults to None + :type ip: optional + """ + if const_expr(self._disable): + return + if const_expr(not isinstance(idx, int)): + raise ValueError(f"expects idx to be 'int', but got {type(idx)}") + if const_expr(idx >= self._num_assertions): + raise ValueError(f"please increase the number of assertions!!!") + if const_expr(self._init_value is True): + self._tensor[idx] = pred and self._tensor[idx] + else: + self._tensor[idx] = pred + self._msgs[idx] = f"{msg}\nAt {loc}" + self._used_indices.add(idx) + + def __enter__(self): + """Enter context manager.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Exit context manager and verify assertions if no exception occurred.""" + # Only verify if there was no exception in the with block + if exc_type is None and not self._disable: + # _CompileTimeAssertion doesn't have verify method as it's checked at compile time + pass + return False # Don't suppress exceptions + + +class RuntimeAssertion(Assertion): + """Runtime assertion helper that verifies conditions at runtime. + ```python + There are two modes to use RuntimeAssertion: + 1. Manual mode - explicitly call verify(): + ```python + @cute.jit + def jit_func(assertions: Assertion): + assertions.store(0, pred, "assertion failed") + assertions = cute.testing.RuntimeAssertion(num_assertions=1) + jit_func(assertions) + assertions.verify() + ``` + + 2. Context manager mode - automatically verifies on exit: + ```python + with cute.testing.RuntimeAssertion(num_assertions=1) as assertions: + jit_func(assertions) + # verify() is called automatically after the with block + ``` + """ + + def __init__( + self, + num_assertions: int = 1, + device=None, + disable: bool = False, + init_value: bool = False, + ): + """Initialize _RuntimeAssertion. + + :param num_assertions: Number of assertions to support + :param device: Device to run assertions on (None for CPU, "cuda" for GPU) + :param disable: If True, assertions are disabled + :param init_value: Initial value for assertion tensor + """ + self._num_assertions = num_assertions + self._device = device + self._disable = disable + self._msgs = [""] * num_assertions + self._init_value = init_value + self._used_indices = set() + if self._disable: + return + import torch + + self._torch_tensor = torch.full( + (self._num_assertions,), + device=self._device, + dtype=torch.bool, + fill_value=init_value, + ) + self._tensor = from_dlpack(self._torch_tensor) + + def __c_pointers__(self): + """Get C pointers for passing to JIT functions.""" + if self._disable: + return [] + return self._tensor.__c_pointers__() + + def __get_mlir_types__(self): + """Get MLIR types for code generation.""" + if self._disable: + return [] + return self._tensor.__get_mlir_types__() + + def __new_from_mlir_values__(self, values): + """Create new instance from MLIR values (for JIT compilation).""" + if self._disable: + return _CompileTimeAssertion( + None, + self._num_assertions, + self._msgs, + self._device, + self._disable, + self._init_value, + self._used_indices, + ) + return _CompileTimeAssertion( + _Tensor(values[0], dtype=Boolean), + self._num_assertions, + self._msgs, + self._device, + self._disable, + self._init_value, + self._used_indices, + ) + + def verify(self): + """Verify all assertions have passed.""" + if self._disable: + return + import torch + + if self._device is not None: + torch.cuda.synchronize() + false_indices = torch.where(self._torch_tensor == False)[0].tolist() + valid_indices = [idx for idx in false_indices if idx in self._used_indices] + if len(valid_indices) > 0: + # emit the first assertion error. + raise AssertionError(self._msgs[valid_indices[0]]) + + def __enter__(self): + """Enter the context manager, returns self for use in 'with' statement.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Exit the context manager, automatically calls verify().""" + if exc_type is None: + # Only verify if no exception occurred in the with block + self.verify() + # Return False to propagate any exception that occurred + return False + + +def _maybe_recast_tensor_from_f4(src: Tensor, tv_layout: Layout): if src.element_type.width == 4: - tv_layout = cute.recast_layout(8, 4, tv_layout) - src = cute.recast_tensor(src, dtype=Int8) + tv_layout = recast_layout(8, 4, tv_layout) + src = recast_tensor(src, dtype=Int8) return src, tv_layout -def _maybe_recast_to_f4(input: cute.TensorSSA, dtype: Type[Numeric]): +def _maybe_recast_to_f4(input: TensorSSA, dtype: Type[Numeric]): """Conditionally recasts the tensor to 4-bit type if the destination type is 4-bit. :param input: The input tensor to recast. @@ -56,18 +306,18 @@ def _maybe_recast_to_f4(input: cute.TensorSSA, dtype: Type[Numeric]): raise TypeError(f"dst_ty must be a type of Numeric, but got {dtype}") if dtype.width == 4: - recast_shape = cute.recast_layout(4, 8, cute.make_layout(input.shape)).shape + recast_shape = recast_layout(4, 8, make_layout(input.shape)).shape i4_vec = vector.bitcast( T.vector(input.type.shape[0] * 2, T.i(4)), input.maybe_downcast() ) res_vect = builtin.unrealized_conversion_cast( [T.vector(i4_vec.type.shape[0], dtype.mlir_type)], [i4_vec] ) - return cute.TensorSSA(res_vect, recast_shape, dtype) + return TensorSSA(res_vect, recast_shape, dtype) return input -def _maybe_recast_from_f4(input: cute.TensorSSA, src_dtype: Type[Numeric]): +def _maybe_recast_from_f4(input: TensorSSA, src_dtype: Type[Numeric]): """Conditionally recasts the tensor from 4-bit type if the source type is 4-bit. :param input: The input tensor to recast. @@ -79,23 +329,23 @@ def _maybe_recast_from_f4(input: cute.TensorSSA, src_dtype: Type[Numeric]): raise TypeError(f"src_ty must be a type of Numeric, but got {src_dtype}") if src_dtype.width == 4: - recast_shape = cute.recast_layout(8, 4, cute.make_layout(input.shape)).shape + recast_shape = recast_layout(8, 4, make_layout(input.shape)).shape i4_vec = builtin.unrealized_conversion_cast( [T.vector(input.type.shape[0], T.i(4))], [input.maybe_downcast()] ) res_vect = vector.bitcast(T.vector(i4_vec.type.shape[0] // 2, T.i8()), i4_vec) - return cute.TensorSSA(res_vect, recast_shape, Int8) + return TensorSSA(res_vect, recast_shape, Int8) return input @CuTeDSL.kernel def _convert_kernel( - gSrc: cute.Tensor, - gDst: cute.Tensor, - cSrc: cute.Tensor, - src_tv_layout: cute.Layout, - dst_tv_layout: cute.Layout, - src_shape: cute.Shape, + gSrc: Tensor, + gDst: Tensor, + cSrc: Tensor, + src_tv_layout: Layout, + dst_tv_layout: Layout, + src_shape: Shape, src_ty, dst_ty, ): @@ -111,9 +361,9 @@ def _convert_kernel( # compose with CTA TV layout # tid, vid -> address - tidfrgSrc = cute.composition(ctaSrc, src_tv_layout) # (T,V) - tidfrgDst = cute.composition(ctaDst, dst_tv_layout) # (T,V) - tidfrgCSrc = cute.composition(ctaCSrc, src_tv_layout) # (T,V) + tidfrgSrc = composition(ctaSrc, src_tv_layout) # (T,V) + tidfrgDst = composition(ctaDst, dst_tv_layout) # (T,V) + tidfrgCSrc = composition(ctaCSrc, src_tv_layout) # (T,V) # print(f"tidfrgSrc = {tidfrgSrc.type}") # slice for threads @@ -124,19 +374,19 @@ def _convert_kernel( # print(f"thrSrc = {thrSrc.type}") # predicate - if cute.elem_less(thrCSrc[0], src_shape): + if elem_less(thrCSrc[0], src_shape): # allocate fragments for gmem->rmem - frgSrc = cute.make_rmem_tensor( - cute.get(src_tv_layout, mode=[1]), gSrc.element_type + frgSrc = make_rmem_tensor( + get(src_tv_layout, mode=[1]), gSrc.element_type ) # (V) - frgDst = cute.make_rmem_tensor( - cute.get(dst_tv_layout, mode=[1]), gDst.element_type + frgDst = make_rmem_tensor( + get(dst_tv_layout, mode=[1]), gDst.element_type ) # (V) # print(f"frgSrc = {frgSrc.type}") # Move data to reg address space - copy_atom_load = cute.make_copy_atom(nvgpu.CopyUniversalOp(), gSrc.element_type) - cute.copy(copy_atom_load, thrSrc, frgSrc) + copy_atom_load = make_copy_atom(nvgpu.CopyUniversalOp(), gSrc.element_type) + copy(copy_atom_load, thrSrc, frgSrc) vec_src = frgSrc.load() vec_src = _maybe_recast_to_f4(vec_src, src_ty) @@ -145,14 +395,14 @@ def _convert_kernel( frgDst.store(vec_dst) # Copy the results back to c - copy_atom_stg = cute.make_copy_atom(nvgpu.CopyUniversalOp(), gDst.element_type) - cute.copy(copy_atom_stg, frgDst, thrDst) + copy_atom_stg = make_copy_atom(nvgpu.CopyUniversalOp(), gDst.element_type) + copy(copy_atom_stg, frgDst, thrDst) @CuTeDSL.jit(preprocess=False) def _convert( - src: cute.Tensor, - dst: cute.Tensor, + src: Tensor, + dst: Tensor, leading_mode: Constexpr, elem_per_copy: Constexpr, ): @@ -160,35 +410,29 @@ def _convert( src_ty = src.element_type dst_ty = dst.element_type - tv_layout = cute.make_layout((128, elem_per_copy), stride=(elem_per_copy, 1)) + tv_layout = make_layout((128, elem_per_copy), stride=(elem_per_copy, 1)) # Step 2. maybe recast from f4 tensor src, src_tv_layout = _maybe_recast_tensor_from_f4(src, tv_layout) dst, dst_tv_layout = _maybe_recast_tensor_from_f4(dst, tv_layout) src_shape = src.shape # predicate tensor - idA = cute.make_identity_tensor(src.shape) + idA = make_identity_tensor(src.shape) # Step 3. select a proper tiling pattern as (...,TileV, ...) src_cta_tiler = [ 1, - ] * cute.rank(src.layout) - src_cta_tiler[leading_mode] = cute.size(src_tv_layout) # (...,TileV,...) + ] * rank(src.layout) + src_cta_tiler[leading_mode] = size(src_tv_layout) # (...,TileV,...) dst_cta_tiler = [ 1, - ] * cute.rank(dst.layout) - dst_cta_tiler[leading_mode] = cute.size(dst_tv_layout) # (...,TileV,...) + ] * rank(dst.layout) + dst_cta_tiler[leading_mode] = size(dst_tv_layout) # (...,TileV,...) # Step 4. partition input and output tensor by cta tiler. - gS = cute.zipped_divide( - src, tuple(src_cta_tiler) - ) # ((...,TileV,...),(...,RestV,...)) - cS = cute.zipped_divide( - idA, tuple(src_cta_tiler) - ) # ((...,TileV,...),(...,RestV,...)) - gD = cute.zipped_divide( - dst, tuple(dst_cta_tiler) - ) # ((...,TileV,...),(...,RestV,...)) + gS = zipped_divide(src, tuple(src_cta_tiler)) # ((...,TileV,...),(...,RestV,...)) + cS = zipped_divide(idA, tuple(src_cta_tiler)) # ((...,TileV,...),(...,RestV,...)) + gD = zipped_divide(dst, tuple(dst_cta_tiler)) # ((...,TileV,...),(...,RestV,...)) # print(f"{gS.type=}") _convert_kernel( @@ -201,8 +445,8 @@ def _convert( src_ty, dst_ty, ).launch( - grid=[cute.size(gS, mode=[1]), 1, 1], - block=[cute.size(src_tv_layout, mode=[0]), 1, 1], + grid=[size(gS, mode=[1]), 1, 1], + block=[size(src_tv_layout, mode=[0]), 1, 1], ) @@ -210,7 +454,7 @@ def _convert( # And when src or dst dtype is narrow precision(Float4E2M1FN/Float8E8M0FNU/Float8E4M3FN), the shape of # their leading dimension should be 4(fp8)/8(fp4) element align. (nvgpu.cvt_fptrunc/cvt_fpext # needs 32-bits aligned input/output) -def convert(src: cute.Tensor, dst: cute.Tensor): +def convert(src: Tensor, dst: Tensor): assert len(src.shape) == len(dst.shape), ( "Shape of src and dst tensors should be the same rank." ) @@ -292,6 +536,14 @@ class JitArguments: def __init__(self, *args, **kwargs): self.args = args self.kwargs = kwargs + self.references = list() + + def add_to_scope(self, references: Any) -> None: + """ + Keeps references to external variables (e.g., Torch tensors when taking a view) + in the scope of the lifetime of the JitArguments object. + """ + self.references.extend(references) def _cuda_success( @@ -428,6 +680,9 @@ def benchmark( :rtype: float """ + import cutlass.base_dsl.jit_executor as jit_executor + import cutlass.cutlass_dsl.cuda_jit_executor as cuda_jit_executor + if stream is None: stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT) @@ -697,7 +952,7 @@ def _benchmark_for_autotune( _cuda_success(err, "Error on querying event") execution_time_ms.append(elapsed_time) # unit: us - time_us = sum(execution_time_ms) / len(execution_time_ms) + time_us = sum(execution_time_ms) * 1e3 / len(execution_time_ms) except Exception as e: print(f"This config execution error: {e}") time_us = float("inf") @@ -775,6 +1030,7 @@ class autotune_jit: Returns: Decorated wrapper function """ + from cutlass.cute import compile # Initialize autotune parameters if not hasattr(func, "_autotune_params"): @@ -825,7 +1081,7 @@ class autotune_jit: # For example, if current_config contains "cluster_shape_mn": (2, 1) # It will override func's default parameter value merged_kwargs = {**kwargs, **current_config} - compiled_func = cute.compile( + compiled_func = compile( func._original_func, *args, **merged_kwargs ) diff --git a/python/CuTeDSL/cutlass/cute/tuple.py b/python/CuTeDSL/cutlass/cute/tuple.py index 52ab318c..15c4114e 100644 --- a/python/CuTeDSL/cutlass/cute/tuple.py +++ b/python/CuTeDSL/cutlass/cute/tuple.py @@ -17,7 +17,17 @@ from cutlass.cutlass_dsl import is_dynamic_expression, dsl_user_op from cutlass._mlir import ir import cutlass._mlir.dialects.cute as _cute_ir -from .typing import XTuple, IntTuple, Shape, Coord, Boolean, is_integer +from .typing import ( + ComposedLayout, + Layout, + Stride, + XTuple, + IntTuple, + Shape, + Coord, + Boolean, + is_integer, +) def wrap(x) -> Tuple[Any, ...]: @@ -184,7 +194,7 @@ def product_each(a: IntTuple, *, loc=None, ip=None) -> IntTuple: def find_if( t: Union[tuple, ir.Value, int], - pred_fn: Callable[[int, Tuple[int, ...]], bool], + pred_fn: Callable[[Union[tuple, ir.Value, int], int], bool], *, loc=None, ip=None, @@ -197,7 +207,8 @@ def find_if( :type t: Union[tuple, ir.Value, int] :param pred_fn: A callable object (lambda, function, etc.) that predicates the value and position in t. It takes the current leaf value and position, returns True if the value or position is satisfied. - :type pred_fn: Callable[[int, Tuple[int, ...]], bool] + The type must be compatible with rank(t). + :type pred_fn: Callable[[Union[tuple, ir.Value, int], int], bool] :return: Index if found at top level, tuple of indices showing nested position, or None if not found :rtype: Union[int, Tuple[int, ...], None] @@ -329,3 +340,154 @@ def elem_less( lhs_val = _pack_coord(lhs, loc=loc, ip=ip) rhs_val = _pack_coord(rhs, loc=loc, ip=ip) return Boolean(_cute_ir.elem_less(lhs_val, rhs_val, loc=loc, ip=ip)) + + +def tuple_cat(*tuples): + """Concatenate multiple tuples into a single tuple. + + This function takes any number of tuples and concatenates them into a single tuple. + Non-tuple arguments are treated as single-element tuples. + + :param tuples: Variable number of tuples to concatenate + :type tuples: tuple or any + :return: A single concatenated tuple + :rtype: tuple + + **Examples:** + + .. code-block:: python + + >>> tuple_cat((1, 2), (3, 4)) + (1, 2, 3, 4) + >>> tuple_cat((1,), (2, 3), (4,)) + (1, 2, 3, 4) + >>> tuple_cat(1, (2, 3)) + (1, 2, 3) + """ + result = () + for t in tuples: + if isinstance(t, tuple): + result += t + else: + result += (t,) + return result + + +def transform_apply(*args, f: Callable, g: Callable): + """Transform elements of tuple(s) with f, then apply g to all results. + + This function applies f to corresponding elements across input tuple(s), + then applies g to all transformed results. It mimics the C++ CuTe implementation. + + Supports multiple signatures: + - transform_apply(t, f, g): For single tuple, computes g(f(t[0]), f(t[1]), ...) + - transform_apply(t0, t1, f, g): For two tuples, computes g(f(t0[0], t1[0]), f(t0[1], t1[1]), ...) + - transform_apply(t0, t1, t2, ..., f, g): For multiple tuples of same length + + For non-tuple inputs, f is applied to the input(s) and g is applied to that single result. + + :param args: One or more tuples (or non-tuples) to transform + :param f: The function to apply to each element (or corresponding elements across tuples) + :type f: Callable + :param g: The function to apply to all transformed elements + :type g: Callable + :param loc: Source location for MLIR, defaults to None + :type loc: optional + :param ip: Insertion point, defaults to None + :type ip: optional + :return: The result of applying g to all transformed elements + :rtype: any + + **Examples:** + + .. code-block:: python + + >>> transform_apply((1, 2, 3), f=lambda x: x * 2, g=lambda *args: sum(args)) + 12 # (1*2 + 2*2 + 3*2) = 12 + >>> transform_apply((1, 2), f=lambda x: (x, x+1), g=tuple_cat) + (1, 2, 2, 3) + >>> transform_apply((1, 2), (3, 4), f=lambda x, y: x + y, g=lambda *args: args) + (4, 6) + """ + if not isinstance(f, Callable): + raise TypeError(f"f must be callable, but got {type(f)}") + if not isinstance(g, Callable): + raise TypeError(f"g must be callable, but got {type(g)}") + + if not args: + raise ValueError("transform_apply requires at least one argument") + + # Check if first argument is a tuple to determine behavior + if isinstance(args[0], tuple): + # Verify all args are tuples of the same length + if not all(isinstance(arg, tuple) for arg in args): + raise TypeError("All arguments must be tuples or all must be non-tuples") + + tuple_length = len(args[0]) + for i, arg in enumerate(args[1:], 1): + if len(arg) != tuple_length: + raise ValueError( + f"All tuple arguments must have the same length. " + f"arg[0] has length {tuple_length}, but arg[{i}] has length {len(arg)}" + ) + + # Apply f to corresponding elements across all tuples: g(f(args[0][i], args[1][i], ...), ...) + transformed_results = tuple( + f(*(arg[i] for arg in args)) for i in range(tuple_length) + ) + return g(*transformed_results) + else: + # Non-tuple case: apply f to all args, then g to that single result + result = f(*args) + return g(result) + + +def filter_tuple(*args, f: Callable): + """Filter and flatten tuple elements by applying a function. + + The function f should return tuples, which are then concatenated together + to produce the final result. This is useful for filtering and transforming + tuple structures in a single pass. + + :param t: The tuple to filter + :type t: Union[tuple, ir.Value, int] + :param f: The function to apply to each element of t + :type f: Callable + :param loc: Source location for MLIR, defaults to None + :type loc: optional + :param ip: Insertion point, defaults to None + :type ip: optional + :return: A concatenated tuple of all results + :rtype: tuple + + **Examples:** + + .. code-block:: python + + >>> # Keep only even numbers, wrapped in tuples + >>> filter_tuple((1, 2, 3, 4), lambda x: (x,) if x % 2 == 0 else ()) + (2, 4) + >>> # Duplicate each element + >>> filter_tuple((1, 2, 3), lambda x: (x, x)) + (1, 1, 2, 2, 3, 3) + """ + if not isinstance(f, Callable): + raise TypeError(f"f must be callable, but got {type(f)}") + + return transform_apply(*args, f=f, g=lambda *args: tuple_cat(*args)) + + +__all__ = [ + "transform_leaf", + "find_if", + "find", + "flatten_to_tuple", + "unflatten", + "product", + "product_like", + "product_each", + "elem_less", + "tuple_cat", + "transform_apply", + "filter_tuple", +] diff --git a/python/CuTeDSL/cutlass/cute/typing.py b/python/CuTeDSL/cutlass/cute/typing.py index f8b16864..80a084da 100644 --- a/python/CuTeDSL/cutlass/cute/typing.py +++ b/python/CuTeDSL/cutlass/cute/typing.py @@ -52,6 +52,15 @@ class SymInt: [self._width == other._width, self._divisibility == other._divisibility] ) + def __mod__(self, other: int) -> Union["SymInt", int]: + if self._divisibility % other != 0: + from math import gcd + + div = gcd(self._divisibility, other) + return SymInt(self._width, divisibility=div) + else: + return 0 + def __c_pointers__(self): return [ctypes.c_void_p(0).value] @@ -381,10 +390,14 @@ __all__ = [ "Float6E2M3FN", "Float6E3M2FN", "IntTuple", - "Layout", + "ScaledBasis", "Coord", "Shape", "Stride", + "Layout", + "ComposedLayout", + "Pointer", + "Tensor", "Tile", "Tiler", "XTuple", diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py index 022cb79d..690546c6 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py @@ -33,6 +33,8 @@ from ..base_dsl.ast_helpers import ( copy_members, get_locals_or_none, closure_check, + fstring_decompose, + FormattedValue, ) from ..base_dsl import * diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py index 644465b6..2c5e8a6e 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py @@ -12,6 +12,7 @@ """ This module provides jit executor related classes for CUTLASS. """ + import ctypes import functools import weakref @@ -32,6 +33,7 @@ from ..base_dsl.common import DSLRuntimeError from ..base_dsl.typing import Int32 from ..base_dsl.runtime.cuda import checkCudaErrors + class CudaDialectJitModule: """Holds the execution engine and cuda libraries.""" @@ -80,22 +82,23 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): jit_function_artifacts, prefix=None, load_from_binary=False, + dynamic_args=None, + dynamic_kwargs=None, ): - self.ir_module = ir_module - self.engine = engine - self.capi_func = capi_func - self.function_name = function_name - self.kernel_info = kernel_info - if args_spec is not None: - self.args_spec = ExecutionArgs(args_spec, self.function_name) - self.jit_time_profiling = jit_time_profiling - self.prefix = prefix - assert ( - isinstance(jit_function_artifacts, JitFunctionArtifacts) - or jit_function_artifacts is None + super().__init__( + ir_module, + engine, + capi_func, + args_spec, + function_name, + kernel_info, + jit_time_profiling, + jit_function_artifacts, + prefix, + load_from_binary, + dynamic_args, + dynamic_kwargs, ) - self.artifacts = jit_function_artifacts - self.load_from_binary = load_from_binary # Set cuda result return type. # When execution engine/capi function is None, do not set the return type. @@ -104,13 +107,6 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): if self.args_spec: self.args_spec.args_spec.annotations["return"] = Int32 - # This runtime state is stored here so that we can preserve the module - # in the compiler cache. Callers can extend the lifetime of the module - # by creating and retaining the executor. - self.jit_module = None - self._executor_lock = threading.RLock() - self._default_executor = None - @functools.cached_property def num_devices(self): """Returns the number of CUDA devices available.""" @@ -222,7 +218,11 @@ class CudaDialectJitCompiledFunction(JitCompiledFunction): device_id = ctypes.c_int32(0) pointer_to_device_id = ctypes.pointer(device_id) - cuda_load_args = [pointer_to_library, pointer_to_device_id, pointer_to_err] + cuda_load_args = [ + pointer_to_pointer_to_library, + pointer_to_device_id, + pointer_to_err, + ] packed_args = (ctypes.c_void_p * len(cuda_load_args))() for i, arg in enumerate(cuda_load_args): packed_args[i] = ctypes.cast(arg, ctypes.c_void_p) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py index ca27d3b7..83196af2 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py @@ -29,6 +29,7 @@ from typing import ( ) import functools import pkgutil +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import is_dataclass, fields from math import ceil from itertools import chain @@ -36,6 +37,8 @@ from pathlib import Path from collections.abc import Sequence import builtins import ctypes +import hashlib +import os from ..base_dsl import * from ..base_dsl import compiler @@ -93,17 +96,8 @@ from .cutlass_ast_decorators import ( _while_execute_dynamic, ) -from .tree_utils import ( - is_constexpr_field, - tree_flatten, - tree_unflatten, - PyTreeDef, - is_frozen_dataclass, - DSLTreeFlattenError, -) from ..base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry - # ============================================================================= # Cutlass DSL Base Abstract Class # ============================================================================= @@ -283,6 +277,18 @@ class CutlassBaseDSL(BaseDSL): log().info(f"GPU module: {self.gpu_module}") return ir.InsertionPoint(self.gpu_module.bodyRegion.blocks[0]) + @staticmethod + def generate_func_ret_op(loc=None, ip=None): + raise NotImplementedError( + "generate_func_ret_op() must be implemented by subclasses." + ) + + @staticmethod + def generate_func_op(arg_types, arg_attrs, kernel_name, loc=None): + raise NotImplementedError( + "generate_func_op() must be implemented by subclasses." + ) + def _generate_kernel_attrs(self, config: BaseDSL.LaunchConfig) -> dict: assert isinstance(config, BaseDSL.LaunchConfig), ( f"Expect LaunchConfig for @kernel, but got {type(config)}" @@ -326,43 +332,105 @@ class CutlassBaseDSL(BaseDSL): Get the version of cutlass dsl, used for computing the hash key of the cache. Including source python files and the shared library. """ + + def _hash_chunk( + key: str, path: str, idx: int, start: int, size: int + ) -> tuple[str, int, bytes]: + """Hash one chunk of a file with SHA-256.""" + h = hashlib.sha256() + if size > 0: + try: + with open(path, "rb") as f: + f.seek(start) + h.update(f.read(size)) + except Exception as e: + raise DSLRuntimeError( + f"Failed to read module file {key}." + "The file may not exist or may not be readable." + "Please re-install the package." + ) from e + return key, idx, h.digest() + + def _iter_jobs(): + """Chunk jobs generator to hash files in parallel""" + for key, path, size in files: + # empty files still get a deterministic hash from SHA-256 of zero bytes + for i in range(max(1, -(-size // chunk_size))): # ceil division + start = i * chunk_size + computed_size = min(chunk_size, max(size - start, 0)) + yield (key, path, i, start, computed_size) + dsl_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - # get the version hash of the cutlass shared library - version_hash = hashlib.sha256() - # update the version hash of the source python files - for lib in pkgutil.walk_packages([dsl_path], prefix="cutlass."): - try: - with open(lib.module_finder.find_spec(lib.name).origin, "rb") as f: - version_hash.update(f.read()) - except Exception: - raise DSLRuntimeError( - f"Failed to read module file {lib.name}. The file may not exist or may not be readable." - "Please re-install the package." - ) + files = [] + + # Keep large dso file first in the list to reduce tail effect + giant_dso_name = str( + next( + (Path(dsl_path) / "_mlir" / "_mlir_libs").glob("_cutlass_ir.cpython*") + ).name + ) + so_path = os.path.join(dsl_path, "_mlir", "_mlir_libs", giant_dso_name) try: # update the version hash of the cutlass shared library - giant_dso_name = str( - next( - (Path(dsl_path) / "_mlir" / "_mlir_libs").glob( - "_cutlass_ir.cpython*" - ) - ).name - ) - with open( - os.path.join(dsl_path, f"_mlir/_mlir_libs/{giant_dso_name}"), - "rb", - ) as f: - while True: - chunk = f.read(1024**2) - if not chunk: - break - version_hash.update(chunk) - except Exception: + so_size = os.path.getsize(so_path) + except Exception as e: raise DSLRuntimeError( f"Failed to read the shared library file {giant_dso_name}." "The file may not exist or may not be readable." "Please re-install the package." + ) from e + files.append((giant_dso_name, so_path, so_size)) + + def handle_import_error(exc): + """Handle errors during package walking, ignoring ImportError and NotImplementedError.""" + if isinstance(exc, (ImportError, NotImplementedError)): + log().info(f"Skipping module due to {type(exc).__name__}: {exc}") + else: + log().warning(f"Unexpected error during package walk: {exc}") + + for lib in pkgutil.walk_packages( + [dsl_path], prefix="cutlass.", onerror=handle_import_error + ): + spec = lib.module_finder.find_spec(lib.name) + if not spec or not spec.origin: + continue + path = spec.origin + try: + size = os.path.getsize(path) + except Exception as e: + raise DSLRuntimeError( + f"Failed to read module file {lib.name}. The file may not exist or may not be readable." + "Please re-install the package." + ) from e + files.append((lib.name, path, size)) + + # Submit chunks to a job queue + chunk_size = 1 << 24 # 16 MB (tuned) + per_file_chunks = {} + # 16 threads max to avoid context switching overhead + # To avoid oversubscription, we use half of cpu_count() + max_workers = min(16, (os.cpu_count() or 8) // 2) + with ThreadPoolExecutor(max_workers=max_workers) as ex: + futures = [ex.submit(_hash_chunk, *job) for job in _iter_jobs()] + for fut in as_completed(futures): + key, idx, digest = fut.result() + per_file_chunks.setdefault(key, []).append((idx, digest)) + + # update the version hash of the cutlass shared library using tree-hash + version_hash = hashlib.sha256() + # Since files list is in arbitrary order, we sort by key to get deterministic order + for key, path, size in sorted(files, key=lambda t: t[0]): + chunks = per_file_chunks.get(key) + file_hash = hashlib.sha256( + b"".join( + digest + for _, digest in sorted( + chunks + ) # Similarily, sort chunks by index to get deterministic order + ) ) + file_hash.update(key.encode("utf-8")) + version_hash.update(file_hash.digest()) return version_hash @@ -399,12 +467,14 @@ class CutlassBaseDSL(BaseDSL): pipeline, args_spec, no_cache, + no_jit_engine, *, full_args=None, full_kwargs=None, dynamic_args=None, dynamic_kwargs=None, original_function_name=None, + funcBody=None, ): """ Compile the module and cache the result. @@ -415,6 +485,7 @@ class CutlassBaseDSL(BaseDSL): :param pipeline: The pipeline to use for compilation. :param args_spec: The args spec to use for compilation. :param no_cache: Whether to cache the result. + :param no_jit_engine: Whether to create JIT execution engine. :param full_args: The full arguments to use for compilation. :param full_kwargs: The full keyword arguments to use for compilation. :param dynamic_args: The dynamic arguments to use for compilation. @@ -433,8 +504,10 @@ class CutlassBaseDSL(BaseDSL): from cutlass.base_dsl.tvm_ffi_builder import attach_ffi_func assert self._tvm_ffi_args_spec_converter is not None - tvm_ffi_spec_params, kwargs_wrapper_spec = self._tvm_ffi_args_spec_converter( - function_name, args_spec, full_args, full_kwargs + tvm_ffi_spec_params, kwargs_wrapper_spec = ( + self._tvm_ffi_args_spec_converter( + function_name, args_spec, full_args, full_kwargs + ) ) tvm_ffi_provider = TVMFFICuteCallProvider(function_name) @@ -454,8 +527,7 @@ class CutlassBaseDSL(BaseDSL): def _make_compiled_func(*args, **kwargs): if kwargs_wrapper_spec.kwonly_names or kwargs_wrapper_spec.arg_defaults: return TVMFFIJitCompiledFunctionWithKwargs( - *args, **kwargs, - kwargs_wrapper_spec=kwargs_wrapper_spec + *args, **kwargs, kwargs_wrapper_spec=kwargs_wrapper_spec ) else: return TVMFFIJitCompiledFunction(*args, **kwargs) @@ -472,11 +544,13 @@ class CutlassBaseDSL(BaseDSL): pipeline, args_spec, no_cache, + no_jit_engine, _make_compiled_func, full_args=full_args, full_kwargs=full_kwargs, dynamic_args=dynamic_args, dynamic_kwargs=dynamic_kwargs, + funcBody=funcBody, ) return super().compile_and_cache( @@ -486,12 +560,14 @@ class CutlassBaseDSL(BaseDSL): pipeline, args_spec, no_cache, + no_jit_engine, CudaDialectJitCompiledFunction, full_args=full_args, full_kwargs=full_kwargs, dynamic_args=dynamic_args, dynamic_kwargs=dynamic_kwargs, original_function_name=original_function_name, + funcBody=funcBody, ) @staticmethod @@ -539,14 +615,21 @@ class CutlassBaseDSL(BaseDSL): cluster_size_x=None, cluster_size_y=None, cluster_size_z=None, + preferred_cluster_size_x=None, + preferred_cluster_size_y=None, + preferred_cluster_size_z=None, dynamic_shared_memory_size=None, use_pdl=False, + cooperative=False, loc=None, ip=None, ): + # set to 3 for PDL, cluster size, and cooperative + max_num_attributes = 3 + + if preferred_cluster_size_x is not None: + max_num_attributes += 1 - # Max number of attributes in the launch config is set to 2 for PDL and cluster size - max_num_attributes = 2 launch_config_type = cuda_dialect.LaunchConfigType.get(max_num_attributes) if len(stream) == 0: @@ -571,7 +654,7 @@ class CutlassBaseDSL(BaseDSL): # Launch config type launch_config_type, # Max num of attributes the launch config can hold - # set to 2 for PDL and cluster size + # set to 3 for PDL, cluster size, and cooperative ir.IntegerAttr.get(ir.IntegerType.get_signless(32), max_num_attributes), block_size_x, block_size_y, @@ -585,10 +668,9 @@ class CutlassBaseDSL(BaseDSL): ip=ip, ) - if use_pdl: - cuda_dialect.launch_cfg_programmatic_stream_serialization_allowed( - cfg, Int32(use_pdl).ir_value(), loc=loc, ip=ip - ) + cuda_dialect.launch_cfg_programmatic_stream_serialization_allowed( + cfg, Int32(use_pdl).ir_value(), loc=loc, ip=ip + ) if cluster_size_x is not None: if cluster_size_y is None: @@ -602,11 +684,27 @@ class CutlassBaseDSL(BaseDSL): cfg, cluster_size_x, cluster_size_y, cluster_size_z, loc=loc, ip=ip ) + if preferred_cluster_size_x is not None: + preferred_cluster_size_y = preferred_cluster_size_y or 1 + preferred_cluster_size_z = preferred_cluster_size_z or 1 + preferred_x = Int32(preferred_cluster_size_x).ir_value(loc=loc, ip=ip) + preferred_y = Int32(preferred_cluster_size_y).ir_value(loc=loc, ip=ip) + preferred_z = Int32(preferred_cluster_size_z).ir_value(loc=loc, ip=ip) + cuda_dialect.launch_cfg_preferred_cluster_dim( + cfg, preferred_x, preferred_y, preferred_z, loc=loc, ip=ip + ) + + cuda_dialect.launch_cfg_cooperative( + cfg, Int32(cooperative).ir_value(loc=loc, ip=ip), loc=loc, ip=ip + ) + op = cuda_dialect.launch_ex( cuda_dialect.ResultType.get(), kernel, cfg, kernel_operands, + # This is true for any DSL generated kernel + assume_kernel_attr=ir.Attribute.parse("#cuda.assume_kernel_attr"), loc=loc, ip=ip, ) @@ -662,7 +760,7 @@ class CutlassBaseDSL(BaseDSL): loc=loc, ip=ip, ) - op.attributes["use_pdl"] = ir.BoolAttr.get(use_pdl) + op.attributes["use_pdl"] = use_pdl.ir_value() return _get_op_result_or_op_results(op) def _kernel_helper(self, funcBody, *args, **kwargs): @@ -674,27 +772,14 @@ class CutlassBaseDSL(BaseDSL): def generate_func_op(self, arg_types, arg_attrs, kernel_name, loc=None): super().generate_func_op(arg_types, arg_attrs, kernel_name) - self.func_op = cuda_dialect.KernelOp( - kernel_name, ir.FunctionType.get(arg_types, []), loc=loc + self.func_op = self.dsl.generate_func_op( + arg_types, arg_attrs, kernel_name, loc ) self.arg_types = arg_types - self.func_op.attributes["cu_attrs"] = ir.DictAttr.get( - { - str( - cuda_dialect.CUFunctionAttribute.non_portable_cluster_size_allowed - ): ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1), - str( - cuda_dialect.CUFunctionAttribute.max_dynamic_shared_size_bytes - ): cuda_dialect.DevMaxSharedMemoryOptinAttr.get(), - } - ) - if arg_attrs is not None: - log().debug(arg_attrs) - self.func_op.arg_attrs = arg_attrs return self.func_op def generate_func_ret_op(self, loc=None, ip=None): - return cuda_dialect.ReturnOp([], loc=loc, ip=ip) + self.dsl.generate_func_ret_op(loc, ip) def get_func_body_start(self): assert self.func_op is not None, "Invalid func_op is not expected!" @@ -736,31 +821,72 @@ class CutlassBaseDSL(BaseDSL): async_deps = cfg.async_deps if not isinstance(cfg.async_deps, (list, tuple)): async_deps = [cfg.async_deps] + + launch_kwargs = {} + + if cfg.has_fallback_cluster: + launch_kwargs.update( + dict( + zip( + ("cluster_size_x", "cluster_size_y", "cluster_size_z"), + tuple(cfg.fallback_cluster), + ) + ) + ) + launch_kwargs.update( + dict( + zip( + ( + "preferred_cluster_size_x", + "preferred_cluster_size_y", + "preferred_cluster_size_z", + ), + tuple(cfg.cluster), + ) + ) + ) + elif cfg.has_cluster: + launch_kwargs.update( + dict( + zip( + ("cluster_size_x", "cluster_size_y", "cluster_size_z"), + tuple(cfg.cluster), + ) + ) + ) + CutlassBaseDSL.cuda_launch_func( async_deps, kernelSym, *cfg.grid, *cfg.block, kernelOperands, - **dict( - zip( - ("cluster_size_x", "cluster_size_y", "cluster_size_z"), - tuple(cfg.cluster), - ) - ), + **launch_kwargs, dynamic_shared_memory_size=cfg.smem, use_pdl=cfg.use_pdl, + cooperative=cfg.cooperative, loc=loc, ) return None - return KernelLauncher( - self, - lambda: _CutlassIrKernelGenHelper(self), - funcBody, - *args, - **kwargs, - ) + custom_name = kwargs.pop("_name_prefix", None) + if custom_name: + return KernelLauncher( + self, + lambda: _CutlassIrKernelGenHelper(self), + funcBody, + *args, + **kwargs, + _name_prefix=custom_name, + ) + else: + return KernelLauncher( + self, + lambda: _CutlassIrKernelGenHelper(self), + funcBody, + *args, + **kwargs, + ) def _preprocess_launch_config_args(self, args, kwargs): """Helper to preprocess args and kwargs for LaunchConfig""" @@ -809,7 +935,10 @@ class CutlassBaseDSL(BaseDSL): # skip generic types such as List[int], Tuple[int, int], etc. for performance consideration? pass - elif getattr(arg_annotation, "__name__", None) == "CUstream" and arg.__class__.__name__ == "_FakeStream": + elif ( + getattr(arg_annotation, "__name__", None) == "CUstream" + and arg.__class__.__name__ == "_FakeStream" + ): # allow FakeStream to be passed as CUDA stream pass @@ -875,14 +1004,14 @@ class CutlassBaseDSL(BaseDSL): ): # Try tree_flatten try: - dyn_vals, _ = tree_flatten(arg) + dyn_vals, attr_vals, _ = tree_flatten(arg) except DSLTreeFlattenError: # If fails, just return the original arg return jit_exec_arg, jit_arg_type, jit_arg_attr if dyn_vals: jit_arg_type.extend([v.type for v in dyn_vals]) - jit_arg_attr.extend([default_attr] * len(dyn_vals)) + jit_arg_attr.extend(attr_vals) jit_exec_arg.extend( _get_c_pointers_cutlass(arg) if is_host else dyn_vals ) @@ -910,7 +1039,7 @@ class CutlassBaseDSL(BaseDSL): ): # Try tree_unflatten try: - dyn_vals, tree_def = tree_flatten(arg) + dyn_vals, _, tree_def = tree_flatten(arg) block_args = fop_args[iv_block_args : iv_block_args + len(dyn_vals)] ir_arg.append(tree_unflatten(tree_def, block_args)) iv_block_args += len(dyn_vals) @@ -937,6 +1066,30 @@ class CuTeDSL(CutlassBaseDSL): super().__init__(name, compiler_provider, pass_sm_arch_name, preprocess=True) + @staticmethod + def generate_func_op(arg_types, arg_attrs, kernel_name, loc=None): + func_op = cuda_dialect.KernelOp( + kernel_name, ir.FunctionType.get(arg_types, []), loc=loc + ) + func_op.attributes["cu_attrs"] = ir.DictAttr.get( + { + str( + cuda_dialect.CUFunctionAttribute.non_portable_cluster_size_allowed + ): ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1), + str( + cuda_dialect.CUFunctionAttribute.max_dynamic_shared_size_bytes + ): cuda_dialect.DevMaxSharedMemoryOptinAttr.get(), + } + ) + if arg_attrs is not None: + log().debug(arg_attrs) + func_op.arg_attrs = arg_attrs + return func_op + + @staticmethod + def generate_func_ret_op(loc=None, ip=None): + return cuda_dialect.ReturnOp([], loc=loc, ip=ip) + # ============================================================================= # KernelLauncher @@ -974,6 +1127,8 @@ class KernelLauncher: self.func_args = func_args self.func_kwargs = func_kwargs + self._name_prefix = func_kwargs.pop("_name_prefix", None) + self._check_func_args(funcBody, *func_args, **func_kwargs) def _check_func_args(self, funcBody, *func_args, **func_kwargs): @@ -1002,6 +1157,9 @@ class KernelLauncher: config = self.dsl.LaunchConfig(*args, **kwargs) kernel_attrs = _build_kernel_attrs(config) + if hasattr(self, "_name_prefix") and self._name_prefix: + self.dsl._name_prefix = self._name_prefix + kernel_generator = self.dsl.kernel_launcher( requiredArgs=["config"], unitAttrNames=["gpu.kernel", "cute.kernel"], @@ -1119,7 +1277,7 @@ def unpack_to_irvalue( log().debug("[%d]: will-unpacked: [type:%s] %s", idx, type(packed), packed) try: - unpacked_values, treedef = tree_flatten( + unpacked_values, _, treedef = tree_flatten( remove_read_only_frozen_dataclass(mixed_values, full_write_args_count) ) except DSLTreeFlattenError as e: @@ -1182,6 +1340,7 @@ def to_index(value): return res + def _validate_iter_args_structure(iter_args, ir_values): """ Validates that iter_args structure contains the same number of atomic values @@ -2087,3 +2246,15 @@ executor.set_functions( all_executor=all_, builtin_redirector=_builtin_redirector, ) + + +class DSLCudaVerNotImplemented(DSLNotImplemented): + """ + Exception raised when a feature is not supported for DSL with CUDA version less than the required version. + """ + + def __init__(self, feature: str, required_version: str): + super().__init__( + message=f"{feature} is not supported for DSL with CUDA version less than {required_version}", + suggestion=f"Consider upgrading to CUDA version {required_version}+ based DSL", + ) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py index 84d0c3a8..5c032d1c 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py @@ -18,10 +18,10 @@ from collections.abc import Sequence from ..base_dsl.dsl import is_dynamic_expression from ..base_dsl.ast_helpers import * from ..base_dsl.utils.logger import log -from ..base_dsl import typing as t +from ..base_dsl import typing as t, Arch from ..base_dsl.typing import Boolean, Numeric, as_numeric +from ..base_dsl.utils.tree_utils import PyTreeDef, check_tree_equal from . import cutlass as cutlass_dsl -from .tree_utils import PyTreeDef, check_tree_equal # ============================================================================= # AST Helpers @@ -289,6 +289,7 @@ def _loop_execute_range_dynamic( unroll: int = -1, unroll_full: bool = False, prefetch_stages: int = None, + vectorize: bool = None, ): """ Example: build an scf.for with optional unroll, using our universal helper. @@ -341,6 +342,17 @@ def _loop_execute_range_dynamic( ) log().debug("prefetch_stages attribute: %s", prefetch_stages_attr) + vectorize_attr = None + if vectorize: + arch = cutlass_dsl.CuTeDSL._get_dsl().get_arch_enum() + if arch < Arch.sm_100: + raise DSLRuntimeError( + f"vectorize is supported for sm_100 and above, got {arch}." + ) + _attr_const_check(vectorize, bool, "vectorize") + vectorize_attr = ir.BoolAttr.get(True) + log().debug("vectorize attribute: %s", vectorize_attr) + log().debug( "Creating scf.ForOp \n\t\tstart=%s: type : %s\n\t\tstop=%s: type : %s\n\t\tstep=%s: type : %s", start_, @@ -373,6 +385,9 @@ def _loop_execute_range_dynamic( if prefetch_stages_attr is not None: for_op.attributes["cutlass.pipelining"] = prefetch_stages_attr + if vectorize_attr is not None: + for_op.attributes["cutlass.vectorize"] = vectorize_attr + return for_op def for_body_builder( diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py index a2060cab..a7c1e7e2 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py @@ -255,8 +255,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): current_device: Optional[ir.Value], target_device: Optional[ir.Value], ) -> ir.Block: - """Set the CUDA device index if it differs from the target device. - """ + """Set the CUDA device index if it differs from the target device.""" # If either device is None, no switching needed if current_device is None: assert target_device is None @@ -274,7 +273,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): self.cond_br( cond=devices_differ, true_block=switch_device_block, - false_block=continuation_block + false_block=continuation_block, ) # Switch device block: call cudaSetDevice @@ -288,7 +287,9 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): ) # Check for errors and branch to continuation - switch_device_block = self.check_cuda_error(result, switch_device_block, context) + switch_device_block = self.check_cuda_error( + result, switch_device_block, context + ) with ir.InsertionPoint(switch_device_block): self.br(continuation_block) @@ -319,7 +320,9 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): op_bundle_sizes=[], op_bundle_operands=[], ) - current_block = self.check_cuda_error(get_device_result, current_block, context) + current_block = self.check_cuda_error( + get_device_result, current_block, context + ) # Load the current device index from the alloca with ir.InsertionPoint(current_block): @@ -351,7 +354,6 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): return current_block - def find_cuda_device_index_from_params(self, context: CallContext): """Find the CUDA device index from tensor parameters.""" for param in context.params: @@ -363,12 +365,9 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): return None def create_shared_cuda_error_block( - self, - current_block: ir.Block, - context: CallContext + self, current_block: ir.Block, context: CallContext ) -> ir.Block: - """Create a shared error handling block for all CUDA errors. - """ + """Create a shared error handling block for all CUDA errors.""" # Create the shared error block after the current block (setup phase) # This block will be branched to from multiple error checking sites # It accepts the error code as a block argument @@ -398,7 +397,9 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): current_block = self.append_unload_to_global_dtors(current_block, context) # Create shared CUDA error handling block after the setup blocks # This reduces code duplication - all CUDA errors branch to this single block - self.cuda_error_handle_block = self.create_shared_cuda_error_block(current_block, context) + self.cuda_error_handle_block = self.create_shared_cuda_error_block( + current_block, context + ) # setup device index, will be set around the call to the target function self.cuda_device_index = self.find_cuda_device_index_from_params(context) current_block = super().__call__(current_block, context) @@ -416,6 +417,7 @@ def _inplace_hide_symbols(ir_module: ir.Module, hide_check: Callable[[str], bool @return: The ir module with the symbols hidden. """ defined_symbols = set() + def walk_llvm_func_op(op): # not a declaration if ( @@ -452,6 +454,7 @@ def _get_format_from_object_file_path(object_file_path: str) -> str: class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): """Base class for TVM FFI compiled function.""" + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -464,10 +467,12 @@ class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): return self.__call__(*exe_args) def export_to_c( - self, object_file_path: str, function_name: str = None, + self, + object_file_path: str, + function_name: str = None, *, enable_pic: bool = True, - export_only_tvm_ffi_symbols: bool = False + export_only_tvm_ffi_symbols: bool = False, ): """Export the TVM FFI function to an object file. @@ -481,8 +486,9 @@ class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): internal_symbol_prefix = "__cute_internal_" + function_name mod = self.ir_module mod = get_export_module( - self.ir_module, internal_symbol_prefix, - preserve_symbols=[f"__tvm_ffi_{self.function_name}"] + self.ir_module, + internal_symbol_prefix, + preserve_symbols=[f"__tvm_ffi_{self.function_name}"], ) rename_tvm_ffi_function(mod, self.function_name, function_name) @@ -498,8 +504,7 @@ class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): f.write(out_bytes) def _create_tvm_ffi_function(self): - """Create the tvm_ffi.Function from the current execution engine. - """ + """Create the tvm_ffi.Function from the current execution engine.""" if self.engine is not None: # trigger eager compile of init callbacks cuda_init = self.engine.raw_lookup("cuda_init") @@ -512,14 +517,14 @@ class TVMFFIJitCompiledFunctionBase(CudaDialectJitCompiledFunction): "__tvm_ffi_" + self.function_name ) tvm_ffi_function = tvm_ffi.Function.__from_mlir_packed_safe_call__( - tvm_ffi_function_ptr, keep_alive_object=self.engine) + tvm_ffi_function_ptr, keep_alive_object=self.engine + ) return tvm_ffi_function return None class TVMFFIJitCompiledFunction(tvm_ffi.Function, TVMFFIJitCompiledFunctionBase): - """TVM FFI Function that directly subclasses the tvm_ffi.Function for pos only arguments. - """ + """TVM FFI Function that directly subclasses the tvm_ffi.Function for pos only arguments.""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -537,8 +542,7 @@ class TVMFFIJitCompiledFunction(tvm_ffi.Function, TVMFFIJitCompiledFunctionBase) class TVMFFIJitCompiledFunctionWithKwargs(TVMFFIJitCompiledFunctionBase): - """TVM FFI Function with kwargs wrapper support - """ + """TVM FFI Function with kwargs wrapper support""" def __init__(self, *args, **kwargs): assert "kwargs_wrapper_spec" in kwargs, "kwargs_wrapper_spec is required" @@ -549,6 +553,7 @@ class TVMFFIJitCompiledFunctionWithKwargs(TVMFFIJitCompiledFunctionBase): if kwargs_wrapper_spec.kwonly_names or kwargs_wrapper_spec.arg_defaults: try: from tvm_ffi.utils import kwargs_wrapper # type: ignore + self._kwargs_wrapper = kwargs_wrapper.make_kwargs_wrapper( self._tvm_ffi_function, arg_names=kwargs_wrapper_spec.arg_names, @@ -557,14 +562,15 @@ class TVMFFIJitCompiledFunctionWithKwargs(TVMFFIJitCompiledFunctionBase): kwonly_defaults=kwargs_wrapper_spec.kwonly_defaults, ) except ImportError: - raise DSLRuntimeError("install apache-tvm-ffi>=0.1.5 to enable kwargs/defaults") + raise DSLRuntimeError( + "install apache-tvm-ffi>=0.1.5 to enable kwargs/defaults" + ) else: # positional only is probably fine self._kwargs_wrapper = self._tvm_ffi_function def __call__(self, *args, **kwargs): - """Call the TVM FFI function with kwargs wrapper. - """ + """Call the TVM FFI function with kwargs wrapper.""" return self._kwargs_wrapper(*args, **kwargs) def __tvm_ffi_object__(self): @@ -574,7 +580,8 @@ class TVMFFIJitCompiledFunctionWithKwargs(TVMFFIJitCompiledFunctionBase): def supports_kwargs_wrapper() -> bool: """Check if the kwargs wrapper is supported.""" try: - from tvm_ffi.utils import kwargs_wrapper # type: ignore + from tvm_ffi.utils import kwargs_wrapper # type: ignore + return True except ImportError: return False diff --git a/python/CuTeDSL/cutlass/jax/__init__.py b/python/CuTeDSL/cutlass/jax/__init__.py new file mode 100644 index 00000000..e9fbd808 --- /dev/null +++ b/python/CuTeDSL/cutlass/jax/__init__.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from functools import cache + +@cache +def is_available(): + """Returns true of Jax support is enabled.""" + try: + import jax + + _HAVE_JAX = True + except ImportError as e: + _HAVE_JAX = False + + return _HAVE_JAX + + +if is_available(): + from .primitive import cutlass_call + from .types import ( + jax_to_cutlass_dtype, + cutlass_to_jax_dtype, + from_dlpack, + JaxArray, + TensorSpec, + ) + from .compile import ( + release_compile_cache, + initialize_cutlass_dsl, + ) + from .ffi import ( + get_export_disabled_safety_checks, + find_cute_dsl_runtime_library, + register_ffi, + is_ffi_registered, + ) + from . import testing + + # This is a legacy name for TensorSpec. It will be removed eventually. + TensorMode = TensorSpec + + # This explicit init method ensures that we avoid initialization at + # unexpected times in jax tracing. + initialize_cutlass_dsl() + + __all__ = [ + "cutlass_call", + "jax_to_cutlass_dtype", + "cutlass_to_jax_dtype", + "from_dlpack", + "JaxArray", + "TensorSpec", + "TensorMode", + "release_compile_cache", + "get_export_disabled_safety_checks", + "is_available", + "testing", + ] +else: + # export is_available check for callers or tests. + __all__ = ["is_available"] diff --git a/python/CuTeDSL/cutlass/jax/compile.py b/python/CuTeDSL/cutlass/jax/compile.py new file mode 100644 index 00000000..f57e9dfe --- /dev/null +++ b/python/CuTeDSL/cutlass/jax/compile.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +import os +import gc +import ctypes +import inspect +from typing import Any, Callable, Optional, Sequence +from dataclasses import dataclass +from functools import partial +from pathlib import Path + +import time +import logging +import threading +import hashlib + +import cuda.bindings.driver as cuda + +import jax +import jax.numpy as jnp +import jaxlib + +from .types import ( + jax_to_cutlass_dtype, + from_dlpack, + JaxArray, + JaxArrayList, + TensorSpec, + JaxTracedArray, + DEFAULT_CUTLASS_DEVICE_MEMSPACE, + DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT, +) + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl.cutlass import CuTeDSL + +logger = logging.getLogger(__name__) + +_CUTLASS_COMPILE_CACHE = {} +_EXPORT_PREFIX = "cutlass_call" + + +@dataclass(frozen=True) +class Arg: + idx: int # position in pytree + shape: tuple[Any, ...] + dtype: jnp.dtype + spec: TensorSpec + + def get_static_flag(self, use_static_tensors: bool): + if self.spec.static is None: + return use_static_tensors + else: + return self.spec.static + + +@dataclass(frozen=True) +class FunctionSpec: + """Contains a specification of the inputs and outputs to the kernel.""" + + in_args: tuple[Arg, ...] + input_tree: Any + out_args: tuple[Arg, ...] + output_tree: Any + input_output_aliases: tuple[tuple[int, int], ...] + input_spec: tuple[TensorSpec, ...] + output_spec: tuple[TensorSpec, ...] + compile_options: str + use_static_tensors: bool + kwargs: tuple[tuple[str, Any]] + + def get_compile_args(self): + """Returns the arguments to provide to cute.compile.""" + compiler_ins = [ + JaxTracedArray( + jax_to_cutlass_dtype(leaf.dtype), + leaf.shape, + DEFAULT_CUTLASS_DEVICE_MEMSPACE, + leaf.spec.ptr_assumed_align, + leaf.spec.layout, + leaf.spec.mode, + leaf.get_static_flag(self.use_static_tensors), + ) + for leaf in self.in_args + ] + compiler_outs = [ + JaxTracedArray( + jax_to_cutlass_dtype(leaf.dtype), + leaf.shape, + DEFAULT_CUTLASS_DEVICE_MEMSPACE, + leaf.spec.ptr_assumed_align, + leaf.spec.layout, + leaf.spec.mode, + leaf.get_static_flag(self.use_static_tensors), + ) + for leaf in self.out_args + ] + return JaxArrayList(tuple(compiler_ins + compiler_outs)) + + +@cute.jit +def jit_wrapper( + stream: cuda.CUstream, + args: JaxArrayList, + *, + wrapped_fn: cutlass.Constexpr, + spec: cutlass.Constexpr, +): + # split buffer argument into inputs and outputs and return to tree + ins, outs = args[: len(spec.in_args)], args[(len(spec.in_args)) :] + ins = [x.get_tensor() for x in ins] + outs = [x.get_tensor() for x in outs] + ins = jax.tree.unflatten(spec.input_tree, ins) + outs = jax.tree.unflatten(spec.output_tree, outs) + wrapped_fn(stream, *ins, *outs, **dict(spec.kwargs)) + + +@dataclass +class CompileResult: + """Holds reference to the compiled kernel and argument spec.""" + + module: bytes + fingerprint: bytes + spec: FunctionSpec + + +def _check_is_valid_type(x, is_input): + if not is_input: + if not isinstance(x, jax.ShapeDtypeStruct): + raise TypeError("Invalid output value passed.", x) + else: + if not isinstance(x, jax.Array): + raise TypeError("Invalid type passed.", x) + + +def build_function_spec( + ins, + in_tree, + outs, + out_tree, + input_spec, + output_spec, + input_output_aliases, + compile_options, + use_static_tensors, + kwargs, +): + in_args = [] + for idx, (arg, spec) in enumerate(zip(ins, input_spec)): + _check_is_valid_type(arg, is_input=True) + in_args.append(Arg(idx, arg.shape, arg.dtype, spec)) + + out_args = [] + for idx, (arg, spec) in enumerate(zip(outs, output_spec)): + _check_is_valid_type(arg, is_input=False) + out_args.append(Arg(idx, arg.shape, arg.dtype, spec)) + + # Return the argument specs to the original pytree structure + # We need this structure to sanely match index positions of the + # arguments to the kernel. + ins_args_structured = jax.tree.unflatten(in_tree, in_args) + out_args_structured = jax.tree.unflatten(out_tree, out_args) + + # Assign per-leaf aliases + input_output_aliases_per_leaf = {} + for input_arg_alias_idx in input_output_aliases: + flat_in, _ = jax.tree.flatten(ins_args_structured[input_arg_alias_idx]) + flat_out, _ = jax.tree.flatten( + out_args_structured[input_output_aliases[input_arg_alias_idx]] + ) + for i, o in zip(flat_in, flat_out): + input_output_aliases_per_leaf[i.idx] = o.idx + + # Remove aliased arguments from output set since they are also provided + # as inputs. This is done at the very top level of the tree to simplify + # how we handle aliasing. The assumption is that the entire pytree is + # aliased. + out_args_structured = list(out_args_structured) + for out_idx in sorted(tuple(set(input_output_aliases.values())), reverse=True): + try: + out_args_structured.pop(out_idx) + except IndexError: + raise ValueError(f"Invalid output alias {out_idx} in input_output_aliases.") + out_args_structured = tuple(out_args_structured) + + in_args_flat, _ = jax.tree.flatten(ins_args_structured) + out_args_flat, out_tree = jax.tree.flatten(out_args_structured) + + spec = FunctionSpec( + tuple(in_args_flat), + in_tree, + tuple(out_args_flat), + out_tree, + tuple(input_output_aliases_per_leaf.items()), + tuple(input_spec), + tuple(output_spec), + compile_options, + use_static_tensors, + tuple((k, kwargs[k]) for k in kwargs), + ) + + return spec + + +_compile_lock = threading.Lock() + + +def get_or_compile_kernel(fn, spec): + """Gets or compiles fn and returns a CutlassCompileResult. + + The function and its specification is used as a key to determine if a new + function must be compiled. + """ + cache_key = (fn, spec) + if cache_key in _CUTLASS_COMPILE_CACHE: + return _CUTLASS_COMPILE_CACHE[cache_key] + + # Don't allow more than 1 thead to compile at any time. + # We assume that the cache key is per thread so we don't need to lock + # the above check in compile cache, + compiled_fn = None + with _compile_lock: + start = time.time() + try: + cute_compile = cutlass.cute.compile + if spec.compile_options: + cute_compile = partial(cute_compile, options=spec.compile_options) + + compiled_fn = cute_compile( + jit_wrapper, + cuda.CUstream(0), + spec.get_compile_args(), + wrapped_fn=fn, + spec=spec, + ) + + module = compiled_fn.dump_to_object(_EXPORT_PREFIX) + fingerprint = bytes.fromhex(hashlib.sha256(module).hexdigest()) + + except Exception as e: + # Log here because Jax can obscure the exception details. + logger.exception("Compilation failure for kernel.") + raise e + end = time.time() + logger.debug(f"Took {end - start} to compile cute kernel.") + + result = CompileResult(module=module, spec=spec, fingerprint=fingerprint) + _CUTLASS_COMPILE_CACHE[cache_key] = result + return result + + +def release_compile_cache(): + """Releases entries from the compile cache. + + Note that is may prevent cute dsl from saving its persistent compilation cache entries. + """ + _CUTLASS_COMPILE_CACHE.clear() + dsl = CuTeDSL._get_dsl() + dsl.jit_cache.clear() + # TODO: This is needed to release frames being held in the DSL + # We should avoid holding such references as they unexpectedly + # extend object lifetime. + dsl.frame = None + gc.collect() + + +class _DummyInitKernel: + @cute.kernel + def kernel(self): + pass + + @cute.jit + def init(self): + pass + + +_CUTLASS_DSL_INITIALIZED = False + + +def initialize_cutlass_dsl(): + """Initializes cutlass DSL.""" + global _CUTLASS_DSL_INITIALIZED + if _CUTLASS_DSL_INITIALIZED: + return + + # Call compiler to ensure we've pre-processed any kernels inside cutedsl. + kernel = _DummyInitKernel() + with _compile_lock: + logger.debug("Initializing cutlass dsl...") + _ = cutlass.cute.compile(kernel.init) + + _CUTLASS_DSL_INITIALIZED = True diff --git a/python/CuTeDSL/cutlass/jax/ffi.py b/python/CuTeDSL/cutlass/jax/ffi.py new file mode 100644 index 00000000..55964cbf --- /dev/null +++ b/python/CuTeDSL/cutlass/jax/ffi.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Sequence +from pathlib import Path +from functools import cache +import os +import logging +import ctypes + +import cutlass +from cutlass.base_dsl.env_manager import find_libs_in_ancestors +from cutlass.cutlass_dsl.cutlass import CuTeDSL + +import jax +import jax.export + +logger = logging.getLogger(__name__) + +_CUTE_DSL_RUNTIME_LIBRARY_NAME = "cute_dsl_runtime" + +_CUTLASS_CALL_TARGETS = { + "CuteDSLRT_NvJaxCutlassCall": {"execute": "CuteDSLRT_NvJaxCutlassCallExecute"}, + "CuteDSLRT_NvJaxCutlassCallNoCudaGraph": { + "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph" + }, +} + + +def get_cutlass_call_ffi_name(allow_cuda_graph): + """Returns the FFI target to call when running cutlass_call functions.""" + disable_cuda_graph = not allow_cuda_graph + if not disable_cuda_graph: + return "CuteDSLRT_NvJaxCutlassCall" + else: + return "CuteDSLRT_NvJaxCutlassCallNoCudaGraph" + + +def get_export_disabled_safety_checks() -> Sequence[jax.export.DisabledSafetyCheck]: + """Returns jax.export.DisabledSafetyCheck to allow cutlass_call kernels.""" + checks = [] + for target in _CUTLASS_CALL_TARGETS: + checks.append(jax.export.DisabledSafetyCheck.custom_call(target)) + return tuple(checks) + + +@cache +def find_cute_dsl_runtime_library(): + """Searches for the CuTeDSL runtime library.""" + dsl = CuTeDSL._get_dsl() + candidate_libs = [] + + try: + # Prefer the environment variable if we find it there. + if dsl.envar.shared_libs: + for lib in dsl.envar.shared_libs.split(":"): + if lib.endswith(f"{_CUTE_DSL_RUNTIME_LIBRARY_NAME}.so"): + return lib + + # Otherwise try to search for the library inside the wheel + def get_libs_cand(start): + libs_cand = find_libs_in_ancestors( + start, [_CUTE_DSL_RUNTIME_LIBRARY_NAME], ["lib"] + ) + if libs_cand: + return libs_cand + return [] + + dsl_libs = get_libs_cand(cutlass.__file__) + if not dsl_libs: + dsl_libs = get_libs_cand(Path(__file__).parent.parent.resolve()) + + candidate_libs.extend(dsl_libs) + + except Exception as e: + logger.debug(f"Failed to locate libraries due to an exception:", e) + + for lib in candidate_libs: + if lib.endswith(f"{_CUTE_DSL_RUNTIME_LIBRARY_NAME}.so"): + return lib + + return None + + +_FFI_CALLS_REGISTERED = False + + +def register_ffi(): + """Registers custom calls with Jax/XLA runtime.""" + global _FFI_CALLS_REGISTERED + if _FFI_CALLS_REGISTERED: + return + + runtime_library = find_cute_dsl_runtime_library() + if not runtime_library: + logger.debug( + "No CuTeDSL runtime library found - skipping python ffi registration." + ) + return + + lib = ctypes.CDLL(runtime_library) + + def _capsule(funcptr): + destructor = ctypes.CFUNCTYPE(None, ctypes.py_object) + builder = ctypes.pythonapi.PyCapsule_New + builder.restype = ctypes.py_object + builder.argtypes = (ctypes.c_void_p, ctypes.c_char_p, destructor) + return builder(funcptr, None, destructor(0)) + + def _register_ffi_targets(lib, targets): + for target_name, target in targets.items(): + handler = {} + for stage, fn_name in target.items(): + fn = getattr(lib, fn_name) + fn.restype = ctypes.c_void_p + handler[stage] = _capsule(fn) + logger.debug(f"Registering ffi handler: {target_name}, {handler}") + jax.ffi.register_ffi_target( + target_name, handler["execute"], platform="CUDA" + ) + + # Register the custom FFI targets + _register_ffi_targets(lib, _CUTLASS_CALL_TARGETS) + + _FFI_CALLS_REGISTERED = True + + +def is_ffi_registered(): + """Returns true if the FFI calls have been registered with Jax/XLA.""" + global _FFI_CALLS_REGISTERED + return _FFI_CALLS_REGISTERED diff --git a/python/CuTeDSL/cutlass/jax/primitive.py b/python/CuTeDSL/cutlass/jax/primitive.py new file mode 100644 index 00000000..f9d838ce --- /dev/null +++ b/python/CuTeDSL/cutlass/jax/primitive.py @@ -0,0 +1,309 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Any, Union, Sequence, Callable +from functools import partial +import logging +import os + +import cuda.bindings.driver as cuda + +import jax, jax.numpy as jnp +import jax.extend +from jax.interpreters import mlir +from jax._src.interpreters import ad +from jax._src.interpreters import batching +from jax._src import ffi +from jax.tree import flatten, unflatten + +import cutlass + +from .compile import get_or_compile_kernel, build_function_spec +from .types import row_major_layout, default_tensor_spec, TensorSpec +from .ffi import get_cutlass_call_ffi_name, is_ffi_registered, register_ffi + +logger = logging.getLogger(__name__) + +cutlass_call_inner_p = jax.extend.core.Primitive("cutlass_call_inner") +cutlass_call_inner_p.multiple_results = True + + +def cutlass_call( + fn: Callable[..., None], + *, + output_shape_dtype: Any, + input_spec: Any = None, + output_spec: Any = None, + input_mode: Any = None, + output_mode: Any = None, + input_output_aliases=None, + allow_cuda_graph=True, + compile_options=None, + use_static_tensors=False, + **kwargs, +): + """Creates a callable that invokes a @cute.jit function. + + Args: + fn: A @cute.jit decorated function that launches a cutlass kernel. + output_shape_dtype: A pytree representing the shape and dtype of the output buffers. + input_output_aliases: Optional mapping of input to output aliases. Positions are specified assuming + a flattened input and output pytree. + input_spec: Specifies a cute.Tensor dimension order for input tensors. If None then the order + will assume the corresponding layout order. + output_spec: Specifies a cute.Tensor dimension order for output tensors. If None then the order + will assume the corresponding layout order. + input_mode: Legacy alias for input_spec. This parameter may be removed in future versions. + output_spec: Legacy alias for output_spec. This parameter may be removed in future versions. + allow_cuda_graph: If false will prevent XLA from building a cuda graph of for this call. + compile_options: Optional compiler arguments to pass into cute.compile. + use_static_tensors: If True, tensor shapes and strides are treated as constexpr values by + default. This can improve performance through compiler specialization but may not work + properly with all kernels. Specific tensors may be marked static or dynamic using the mode + and override this flag. + kwargs: Optional constexpr parameters to pass into the kernel fn. + + Note: This API is experimental and subject to change! + """ + output_shape_dtype = jax.tree.map( + lambda leaf: jax.ShapeDtypeStruct(leaf.shape, leaf.dtype), output_shape_dtype + ) + + if input_output_aliases is None: + input_output_aliases = {} + + if input_spec and input_mode: + raise ValueError( + "input_spec and input_mode can not both be set. Use input_spec only as input_mode is deprecated." + ) + if input_mode: + input_spec = input_mode + + if output_spec and output_mode: + raise ValueError( + "input_spec and input_mode can not both be set. Use input_spec only as input_mode is deprecated." + ) + if output_mode: + output_spec = output_mode + + return _cutlass_call_impl( + fn, + output_shape_dtype=output_shape_dtype, + input_spec=input_spec, + output_spec=output_spec, + input_output_aliases=input_output_aliases, + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + **kwargs, + ) + + +def _normalize_tensor_spec(value: Any): + if value is None: + return [None] + elif isinstance(value, (tuple, list)): + if isinstance(value[0], int): # single tuple of modes + return TensorSpec(mode=tuple(value)) + else: + flat, _ = jax.tree.flatten( + [_normalize_tensor_spec(x) for x in value], + is_leaf=lambda x: x is None or isinstance(x, TensorSpec), + ) + return flat + elif isinstance(value, TensorSpec): + return [value] + else: + raise TypeError(f"Unexpected value for TensorMode {value} {type(value)}") + + +def _cutlass_call_impl( + fn, + *, + output_shape_dtype: Any, + input_spec: Any, + output_spec: Any, + input_output_aliases, + allow_cuda_graph, + compile_options, + use_static_tensors, + **kwargs, +): + multiple_results = isinstance(output_shape_dtype, Sequence) + if not multiple_results: + output_shape_dtype = (output_shape_dtype,) + output_shape_dtype_flat, output_tree = jax.tree.flatten(output_shape_dtype) + + @partial(jax.jit, inline=True) + def call_wrapper(*args): + args_flat, args_tree = jax.tree.flatten(args) + + if input_spec is None: + input_spec_flat = tuple(default_tensor_spec(x) for x in args_flat) + else: + input_spec_flat = _normalize_tensor_spec(input_spec) + for idx, (spec, arg) in enumerate(zip(input_spec_flat, args_flat)): + if spec is None: + input_spec_flat[idx] = default_tensor_spec(arg) + input_spec_flat = tuple(input_spec_flat) + + if output_spec is None: + output_spec_flat = tuple( + default_tensor_spec(x) for x in output_shape_dtype_flat + ) + else: + output_spec_flat = _normalize_tensor_spec(output_spec) + for idx, (spec, arg) in enumerate( + zip(output_spec_flat, output_shape_dtype_flat) + ): + if spec is None: + output_spec_flat[idx] = default_tensor_spec(arg) + output_spec_flat = tuple(output_spec_flat) + if len(input_spec_flat) != len(args_flat): + raise ValueError( + f"Must has same number of input modes ({len(input_spec_flat)}) as input arrays ({len(args_flat)})." + ) + + if len(output_spec_flat) != len(output_shape_dtype_flat): + raise ValueError( + f"Must has same number of output modes ({len(output_spec_flat)}) as output arrays ({len(output_shape_dtype_flat)})." + ) + + # Validate dynamic mode settings match whatever static shape + # information we got as input. + for idx, (arg, spec) in enumerate(zip(args_flat, input_spec_flat)): + if spec.layout is not None and len(spec.layout) != len(arg.shape): + raise ValueError( + f"Input #{idx} has invalid layout {spec.layout} for shape {arg.shape}." + ) + if spec.mode is not None and len(spec.mode) != len(arg.shape): + raise ValueError( + f"Input #{idx} has invalid mode {spec.mode} for shape {arg.shape}." + ) + + for idx, (arg, spec) in enumerate( + zip(output_shape_dtype_flat, output_spec_flat) + ): + if spec.layout is not None and len(spec.layout) != len(arg.shape): + raise ValueError( + f"Output #{idx} has invalid layout {spec.layout} for shape {arg.shape}." + ) + + if spec.mode is not None and len(spec.mode) != len(arg.shape): + raise ValueError( + f"Output #{idx} has invalid mode {spec.mode} for shape {arg.shape}." + ) + + output_flat = cutlass_call_inner_p.bind( + *args_flat, + fn=fn, + args_tree=args_tree, + output_shape_dtype_flat=tuple(output_shape_dtype_flat), + output_tree=output_tree, + input_spec_flat=tuple(input_spec_flat), + output_spec_flat=tuple(output_spec_flat), + input_output_aliases=tuple(input_output_aliases.items()), + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + **kwargs, + ) + + output = jax.tree.unflatten(output_tree, output_flat) + return output if multiple_results else output[0] + + return call_wrapper + + +@cutlass_call_inner_p.def_abstract_eval +def cutlass_call_inner_p_abstract(*_, output_shape_dtype_flat, **__): + return [jax.core.ShapedArray(x.shape, x.dtype) for x in output_shape_dtype_flat] + + +def cutlass_call_inner_p_impl( + *args_flat, + fn, + args_tree: Any, + output_shape_dtype_flat: Any, + output_tree: Any, + input_spec_flat: Any, + output_spec_flat: Any, + input_output_aliases, + allow_cuda_graph, + compile_options, + use_static_tensors, + **kwargs, +): + input_output_aliases = dict(input_output_aliases) + spec = build_function_spec( + args_flat, + args_tree, + output_shape_dtype_flat, + output_tree, + input_spec_flat, + output_spec_flat, + input_output_aliases, + compile_options, + use_static_tensors, + kwargs, + ) + + kernel = get_or_compile_kernel(fn, spec) + + # Ensure our FFI target is registered. We do this lazily here + # so that we only load the dependant library if needed. + if not is_ffi_registered(): + register_ffi() + + call_name = get_cutlass_call_ffi_name(allow_cuda_graph) + fun = jax.ffi.ffi_call( + call_name, + result_shape_dtypes=output_shape_dtype_flat, + input_output_aliases=dict(spec.input_output_aliases), + ) + + return fun(*args_flat, module=kernel.module, key=kernel.fingerprint) + + +def _cutlass_call_jvp_rule(*args, **kwargs): + del args, kwargs + raise NotImplementedError( + "cutlass_call does not support VJP. Please use `jax.custom_jvp` for taking gradients." + ) + + +ad.primitive_jvps[cutlass_call_inner_p] = _cutlass_call_jvp_rule + + +def _cutlass_call_transpose_rule(*args, **kwargs): + del args, kwargs + raise NotImplementedError( + "cutlass_call does not support transpose. Please use `jax.custom_vjp` for taking gradients." + ) + + +ad.primitive_transposes[cutlass_call_inner_p] = _cutlass_call_transpose_rule + + +def _cutlass_call_vmap_rule(*args, **kwargs): + del args, kwargs + raise NotImplementedError( + "cutlass_call does not support batching with jax.vmap. Please " + "use jax.custom_batching.custom_vmap for applying vmap. " + ) + + +batching.primitive_batchers[cutlass_call_inner_p] = _cutlass_call_vmap_rule + +jax._src.dispatch.simple_impl(cutlass_call_inner_p) +jax._src.dispatch.prim_requires_devices_during_lowering.add(cutlass_call_inner_p) +lowering = mlir.lower_fun(cutlass_call_inner_p_impl, multiple_results=True) +mlir.register_lowering(cutlass_call_inner_p, lowering, platform="cuda") diff --git a/python/CuTeDSL/cutlass/jax/testing.py b/python/CuTeDSL/cutlass/jax/testing.py new file mode 100644 index 00000000..781671a2 --- /dev/null +++ b/python/CuTeDSL/cutlass/jax/testing.py @@ -0,0 +1,258 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from functools import partial + +import jax +import jax.numpy as jnp + +import cutlass.cute as cute +from cutlass.cutlass_dsl import dsl_user_op + +def reorder_modes(src: str, target: str) -> tuple[int, ...]: + """Computes the mode given a source and target order.""" + src = tuple(src) + target = tuple(target) + src_map = {} + for idx, s in enumerate(src): + src_map[s] = idx + return tuple([src_map[d] for d in target]) + + +def gemm_a_major(d: str): + """Returns order for A tensor major mode.""" + return {"k": "lmk", "m": "lkm"}[d] + + +def gemm_a_mode(d: str) -> tuple[int, ...]: + """Returns mode for A tensor major mode.""" + return reorder_modes(gemm_a_major(d), "mkl") + + +def gemm_b_major(d: str): + """Returns order for B tensor major mode.""" + return {"k": "lnk", "n": "lkn"}[d] + + +def gemm_b_mode(d: str) -> tuple[int, ...]: + """Returns mode for B tensor major mode.""" + return reorder_modes(gemm_b_major(d), "nkl") + + +def gemm_c_major(d: str): + """Returns order for C tensor major mode.""" + return {"n": "lmn", "m": "lnm"}[d] + + +def gemm_c_mode(d: str) -> tuple[int, ...]: + """Returns mode for C tensor major mode.""" + return reorder_modes(gemm_c_major(d), "mnl") + + +def gemm_a_shape(l, m, k, major) -> tuple[int, ...]: + """Returns shape for A tensor given major mode.""" + assert major in ("k", "m") + shape = (l, m, k) if major == "k" else (l, k, m) + return shape + + +def gemm_b_shape(l, n, k, major) -> tuple[int, ...]: + """Returns shape for B tensor given major mode.""" + assert major in ("k", "n") + shape = (l, n, k) if major == "k" else (l, k, n) + return shape + + +def gemm_c_shape(l, m, n, major) -> tuple[int, ...]: + """Returns shape for C tensor given major mode.""" + assert major in ("m", "n") + shape = (l, m, n) if major == "n" else (l, n, m) + return shape + + +@dsl_user_op +def get_gemm_shape_from_tensors( + a: cute.Tensor, b: cute.Tensor, *, loc=None, ip=None +) -> tuple[int, int, int, int]: + """Returns a tuple of (M, N, K, L) from A/B gemm tensors.""" + # mkl, nkl + m, k, l = a.shape[:] + n = b.shape[0] + return (m, n, k, l) + +def create_tensor( + shape, dtype, key, *, minval=-2.0, maxval=2.0, fill_value=None, fill_arange=False +): + if fill_arange: + tensor = jnp.ones(shape, dtype=dtype) + tensor = tensor * jnp.arange(tensor.size, dtype=tensor.dtype).reshape( + tensor.shape + ) + elif fill_value is not None: + tensor = jnp.full(shape, fill_value, dtype=dtype) + else: + tensor = jax.random.uniform( + key, shape, dtype=jnp.float32, minval=minval, maxval=maxval + ) + tensor = tensor.astype(dtype) + return tensor + + +def create_a_tensor( + l, + m, + k, + major, + dtype, + key, + minval=-2.0, + maxval=2.0, + fill_value=None, + fill_arange=False, +): + shape = gemm_a_shape(l, m, k, major) + tensor = create_tensor( + shape, + dtype, + key, + minval=minval, + maxval=maxval, + fill_value=fill_value, + fill_arange=fill_arange, + ) + return tensor + + +def create_b_tensor( + l, + n, + k, + major, + dtype, + key, + minval=-2.0, + maxval=2.0, + fill_value=None, + fill_arange=False, +): + shape = gemm_b_shape(l, n, k, major) + tensor = create_tensor( + shape, + dtype, + key, + minval=minval, + maxval=maxval, + fill_value=fill_value, + fill_arange=fill_arange, + ) + return tensor + + +def create_cd_tensor( + l, + m, + n, + major, + dtype, + key, + *, + minval=-2.0, + maxval=2.0, + fill_value=None, + fill_arange=False, +): + shape = gemm_c_shape(l, m, n, major) + tensor = create_tensor( + shape, + dtype, + key, + minval=minval, + maxval=maxval, + fill_value=fill_value, + fill_arange=fill_arange, + ) + return tensor + + +def gemm_reference_einsum( + a, + b, + acc_dtype, + c_dtype, + a_major, + b_major, + c_major, + sf_a=None, + sf_b=None, + precision="highest", +): + a_idx = gemm_a_major(a_major) + b_idx = gemm_b_major(b_major) + c_idx = gemm_c_major(c_major) + spec = f"{a_idx},{b_idx}->{c_idx}" + + # If block scaled pre-scale input at higher precision + # Assumes we only use it for fp8 and smaller. + if sf_a is not None: + sf_vec_size = int(a.shape[-1] // sf_a.shape[-1]) + sf_a = jnp.repeat(sf_a, sf_vec_size, axis=-1) + a = a.astype(jnp.float16) * sf_a.astype(jnp.float16) + + if sf_b is not None: + sf_vec_size = int(b.shape[-1] // sf_b.shape[-1]) + sf_b = jnp.repeat(sf_b, sf_vec_size, axis=-1) + b = b.astype(jnp.float16) * sf_b.astype(jnp.float16) + + return jax.jit( + lambda a, b: jnp.einsum( + spec, a, b, preferred_element_type=acc_dtype, precision=precision + ).astype(c_dtype) + )(a, b) + + +def create_attn_tensors( + b, s, hq, hkv, d, dtype, key, *, minval=-2.0, maxval=2.0, fill_value=None +): + qkey, kkey, vkey = jax.random.split(key, 3) + return ( + create_tensor( + (b, s, hq, d), + dtype, + qkey, + minval=minval, + maxval=maxval, + fill_value=fill_value, + ), + create_tensor( + (b, s, hkv, d), + dtype, + kkey, + minval=minval, + maxval=maxval, + fill_value=fill_value, + ), + create_tensor( + (b, s, hkv, d), + dtype, + vkey, + minval=minval, + maxval=maxval, + fill_value=fill_value, + ), + ) + + +def attn_ref(q, k, v, is_causal: bool): + return jax.jit( + lambda q, k, v: jax.nn.dot_product_attention( + q, k, v, is_causal=is_causal, implementation="cudnn" + ) + )(q, k, v) diff --git a/python/CuTeDSL/cutlass/jax/types.py b/python/CuTeDSL/cutlass/jax/types.py new file mode 100644 index 00000000..e6f84ac2 --- /dev/null +++ b/python/CuTeDSL/cutlass/jax/types.py @@ -0,0 +1,425 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Type, Optional, Sequence, Union, Callable, Any, TypeVar +import sys +import ctypes +import math +import inspect +from dataclasses import dataclass, field +from functools import partial, reduce +from operator import mul +from itertools import chain +from typing import Annotated + +import cuda.bindings.driver as cuda + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack as _from_dlpack +from cutlass.cute import AddressSpace, Numeric, IntTuple +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, arith +import cutlass._mlir.dialects.cute as _cute_ir + +JAX_DTYPE_TO_CUTLASS_DTYPE = { + jnp.bool.dtype: cutlass.Boolean, + jnp.int8.dtype: cutlass.Int8, + jnp.int16.dtype: cutlass.Int16, + jnp.int32.dtype: cutlass.Int32, + jnp.int64.dtype: cutlass.Int64, + jnp.uint8.dtype: cutlass.Uint8, + jnp.uint16.dtype: cutlass.Uint16, + jnp.uint32.dtype: cutlass.Uint32, + jnp.uint64.dtype: cutlass.Uint64, + jnp.bfloat16.dtype: cutlass.BFloat16, + jnp.float16.dtype: cutlass.Float16, + jnp.float32.dtype: cutlass.Float32, + jnp.float64.dtype: cutlass.Float64, + jnp.float8_e8m0fnu.dtype: cutlass.Float8E8M0FNU, + jnp.float8_e5m2.dtype: cutlass.Float8E5M2, + jnp.float8_e4m3.dtype: cutlass.Float8E4M3, + jnp.float8_e4m3fn.dtype: cutlass.Float8E4M3FN, + jnp.float8_e4m3b11fnuz.dtype: cutlass.Float8E4M3B11FNUZ, + jnp.float4_e2m1fn.dtype: cutlass.Float4E2M1FN, +} +CUTLASS_DTYPE_TO_JAX_DTYPE = { + value: key for key, value in JAX_DTYPE_TO_CUTLASS_DTYPE.items() +} + +DEFAULT_CUTLASS_DEVICE_MEMSPACE = AddressSpace.gmem +DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT = 256 + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class TensorSpec: + """Provides a specification of cute.Tensor modes and additional metadata about + dynamic/static shapes for compilation. + + Arguments: + layout : Specifies the position of stries as they relate to the framework + tensor (S0, S1, ... SN) + mode : Specifies the position of each mode in the tensor (M0, M1, ... MN) + static : Specifies the tensor shape is represented as static constexpr. + ptr_assumed_align: Specifies the pointer alignment. + """ + + # Specifies the layout of the Jax array. If not it will be assumed that the layout + # is row major. + layout: tuple[int, ...] | None = field(metadata=dict(static=True), default=None) + # Indicates the order of modes. If unspecified the modes will match exactly with + # the layout of the Jax tensor (e.g. row-major). Typically used to map from the + # input layout to kernel layouts (e.g. MKL/NKL/MNL). + mode: tuple[int, ...] | None = field(metadata=dict(static=True), default=None) + # Indicates the shape and strides will be defined statically. Setting True ay enable + # additional optimization. Kernels that do not support static shapes will generate + # compile errors if this is enabled so we leave it off by default. + static: bool = field(metadata=dict(static=True), default=None) + # Overrides the default pointer alignment. Generally this should not be changed + # but is left here to provide a hook. + ptr_assumed_align: int = field( + metadata=dict(static=True), default=DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT + ) + +def row_major_layout(shaped): + """Returns a row major layout given a shaped value. + + Row major layout is (N-1, N-2, ... 1, 0) for an N-dimensional tensor. + """ + if hasattr(shaped, "shape"): + shaped = shaped.shape + return tuple(reversed(range(len(shaped)))) + + +def default_tensor_mode(shaped): + """Returns a default tensor mode given a shaped value. + + Default mode is (0, 1, ... N-2, N-1) for an N_dimensional tensor. + """ + if hasattr(shaped, "shape"): + shaped = shaped.shape + return tuple(range(len(shaped))) + + +def default_tensor_spec(shaped) -> TensorSpec: + """Returns a default tensor spec given a shaped value. + + Default layout is (N-1, N-2, ... 1, 0) for an N-dimensional tensor. + Default mode is (0, 1, ... N-2, N-1) for an N_dimensional tensor. + """ + if hasattr(shaped, "shape"): + shaped = shaped.shape + return TensorSpec(layout=row_major_layout(shaped), mode=default_tensor_mode(shaped)) + + +def jax_to_cutlass_dtype(dtype): + """Gets the corresponding cutlass dtype given a jax dtype.""" + dtype = jnp.dtype(dtype) + if dtype not in JAX_DTYPE_TO_CUTLASS_DTYPE: + raise ValueError(f"Jax dtype [{dtype}] has no equivalent cutlass dtype.") + return JAX_DTYPE_TO_CUTLASS_DTYPE[dtype] + + +def cutlass_to_jax_dtype(dtype): + """Gets the corresponding cutlass dtype given a jax dtype.""" + if dtype not in CUTLASS_DTYPE_TO_JAX_DTYPE: + raise ValueError(f"Cutlass dtype [{dtype}] has no equivalent jax dtype.") + return CUTLASS_DTYPE_TO_JAX_DTYPE[dtype] + + +def from_dlpack(array, assumed_align: int = DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT): + """Convert jax.Array to a DL pack tensor.""" + return _from_dlpack(array, assumed_align=assumed_align) + + +class JaxArray: + """Base class for JaxArray argument type. + + JaxArray provides glue between XLA/JAX FFI tensors and cute.Tensor. + + The following fields/properties provide control over the conversion + to cute.Tensor as part of jax.jit lowering. These properties are + constexpr and compiled into the kernel. + + 1. dtype: The tensor data type defined by the jax array. + 2. shape: The tensor shape defined at jit tracing time. This shape + can be concrete or symbolic in the case of jax.export. + 3. mem_space: The memory space of the tensor. Defaults to gmem. + 4. assumed_align: The alignment of the tensor. Defaults to XLA alignment. + 5. order: Specifies the order of the shape to determine strides. + 6. mode: Specifies how to map ordered elements to the modes od a cute.Layout. + 7. static: If True, tensor shapes and strides are compiled statically. + """ + + def __init__( + self, + dtype, + shape, + mem_space, + assumed_align, + order=None, + mode=None, + static=False, + ): + self.dtype = dtype + self.shape = tuple(shape) + self.ndim = len(self.shape) + self.mem_space = mem_space + self.assumed_align = assumed_align + if order is None: + order = row_major_layout(shape) + if mode is None: + mode = default_tensor_mode(shape) + + if len(order) != len(shape): + raise ValueError(f"layout must be same length as shape", order, shape) + for s in order: + if s < 0 or s >= len(shape): + raise ValueError(f"Invalid index {s} in stride order", order, shape) + if len(tuple(set(order))) != len(order): + raise ValueError(f"layout has duplicate indices", order) + + if len(mode) != len(shape): + raise ValueError(f"mode must be same length as shape", mode, shape) + for s in mode: + if s < 0 or s >= len(shape): + raise ValueError(f"Invalid index {s} in stride order", mode, shape) + if len(tuple(set(mode))) != len(mode): + raise ValueError(f"mode has duplicate indices", mode) + + self.order = tuple(order) + self.mode = tuple(mode) + + if any([jax.export.is_symbolic_dim(s) for s in self.shape]) and static: + raise ValueError( + f"{self.shape} contains one or more symbolic dimensions requires static=False" + ) + self.static = static + + +class JaxArrayValue(JaxArray): + """The IR representation of the JaxArray.""" + + def __init__( + self, + ir_value, + dtype, + shape, + mem_space, + assumed_align, + order, + mode, + static, + ): + super().__init__(dtype, shape, mem_space, assumed_align, order, mode, static) + self.value = ir_value + + def __str__(self): + return f"JaxArrayValue<{self.value}:{self.dtype}:{self.shape}:{self.order}:{self.mode}:{self.static}>" + + def __repr__(self): + return str(self) + + def _make_ordered_layout_dynamic_strides( + self, shape, order: tuple[int, ...], *, loc=None, ip=None + ): + i32 = ir.IntegerType.get_signless(32) + i64 = ir.IntegerType.get_signless(64) + one = arith.constant(i64, 1) + zero = arith.constant(i64, 0) + pairs = sorted(zip(shape, order), key=lambda x: x[1]) + + # Compute strides for each element in order. + strides = [1] # static 1 for leading + if len(shape) > 1: + strides.append(pairs[0][0]) + for i, idx in enumerate(range(len(pairs[:-2]))): + strides.append(arith.muli(pairs[i][0], strides[-1])) + + # Apply the order to strides + strides_ordered = [] + for i in range(len(shape)): + strides_ordered.append(strides[order[i]]) + + # zero out any stride for a shape of size 1 to align with make_ordered_layout + # We ignore the leading dimension of 1 + final_stride = [] + for i in range(len(shape)): + x = arith.cmpi(0, one, shape[i]) + s = strides_ordered[i] + if isinstance(s, int) and s == 1: + final_stride.append(s) + else: + final_stride.append(arith.select(x, zero, s)) + + # Shapes are expected to be int32 so truncate to that before creating layout + shape = tuple([arith.trunci(i32, s) for s in shape]) + + return cute.make_layout(shape, stride=tuple(final_stride)) + + def _load_dynamic_shapes(self, ffi_buffer, *, loc=None, ip=None): + i64 = ir.IntegerType.get_signless(64) + shape_array = llvm.extractvalue( + llvm.PointerType.get(), + ffi_buffer, + [1], + loc=loc, + ip=ip, + ) + + shape_i64 = [] + for i in range(len(self.shape)): + r = llvm.getelementptr( + llvm.PointerType.get(), + shape_array, + [], + raw_constant_indices=ir.DenseI32ArrayAttr.get([i]), + elem_type=i64, + loc=loc, + ip=ip, + ) + shape_i64.append(llvm.load(i64, r, loc=loc, ip=ip)) + + return tuple(shape_i64) + + def _load_pointer(self, ffi_buffer, *, loc=None, ip=None): + raw_ptr = llvm.extractvalue( + llvm.PointerType.get(), + ffi_buffer, + [0], + loc=loc, + ip=ip, + ) + return cute.make_ptr( + self.dtype, + raw_ptr, + self.mem_space, + assumed_align=self.assumed_align, + loc=loc, + ip=ip, + ) + + def get_tensor(self, *, loc=None, ip=None): + ffi_buffer_type = llvm.StructType.get_literal( + [llvm.PointerType.get(), llvm.PointerType.get()] + ) + + ffi_buffer = llvm.load(ffi_buffer_type, self.value, loc=loc, ip=ip) + pointer = self._load_pointer(ffi_buffer) + + if self.static: + shape = tuple(self.shape) + layout = cute.make_ordered_layout(shape, order=self.order, loc=loc, ip=ip) + else: + shape = self._load_dynamic_shapes(ffi_buffer) + layout = self._make_ordered_layout_dynamic_strides(shape, self.order) + + # Apply mode order + if self.mode is not None: + layout = cute.select(layout, self.mode, loc=loc, ip=ip) + + return cute.make_tensor(pointer, layout, loc=loc, ip=ip) + + def __extract_mlir_values__(self): + return [self.value] + + def __new_from_mlir_values__(self, values): + return JaxArrayValue( + values[0], + self.dtype, + self.shape, + self.mem_space, + self.assumed_align, + self.order, + self.mode, + self.static, + ) + + +class JaxTracedArray(JaxArray): + """Represents a traced array value that is used for cute.compile. + + Traced values are not real tensors or allocated on the device. + """ + + def __init__( + self, + dtype, + shape, + mem_space, + assumed_align, + order, + mode, + static, + ): + super().__init__(dtype, shape, mem_space, assumed_align, order, mode, static) + + def __str__(self): + return f"JaxTracedArray<{self.dtype}:{self.shape}:{self.order}:{self.mode}:{self.static}>" + + def __repr__(self): + return str(self) + + def __get_mlir_types__(self): + # Struct passed as opaque object. + return [llvm.PointerType.get()] + + def __new_from_mlir_values__(self, values): + return JaxArrayValue( + values, + self.dtype, + self.shape, + self.mem_space, + self.assumed_align, + self.order, + self.mode, + self.static, + ) + + def __c_pointers__(self): + return [0] + + +class JaxArrayList: + """Holds list of JaxArray or JaxTracedArray. + This class facilitates conversion of JaxTracedArray to JaxArray when crossing + the jit boundary. + """ + + def __init__(self, arrays: Sequence[JaxArray]): + self.arrays = tuple(arrays) + + def __getitem__(self, idx): + return self.arrays[idx] + + def __len__(self): + return len(self.arrays) + + def __iter__(self): + return iter(self.arrays) + + def __c_pointers__(self): + return [x.__c_pointers__()[0] for x in self.arrays] + + def __get_mlir_types__(self): + return [x.__get_mlir_types__()[0] for x in self.arrays] + + def __extract_mlir_values__(self): + return [x.__extract_mlir_values__()[0] for x in self.arrays] + + def __new_from_mlir_values__(self, values): + return JaxArrayList( + [x.__new_from_mlir_values__(v) for x, v in zip(self.arrays, values)] + ) diff --git a/python/CuTeDSL/cutlass/pipeline/__init__.py b/python/CuTeDSL/cutlass/pipeline/__init__.py index 461f01dd..1263f3fc 100644 --- a/python/CuTeDSL/cutlass/pipeline/__init__.py +++ b/python/CuTeDSL/cutlass/pipeline/__init__.py @@ -35,7 +35,6 @@ from .sm90 import ( PipelineAsync, PipelineCpAsync, PipelineTmaAsync, - PipelineTmaMultiConsumersAsync, PipelineTmaStore, PipelineOrder, PipelineProducer, @@ -46,6 +45,8 @@ from .sm100 import ( PipelineTmaUmma, PipelineAsyncUmma, PipelineUmmaAsync, + PipelineClcFetchAsync, + PipelineTmaMultiConsumersAsync, ) __all__ = [ @@ -63,9 +64,10 @@ __all__ = [ "PipelineCpAsync", "PipelineTmaAsync", "PipelineTmaUmma", - "PipelineTmaMultiConsumersAsync", "PipelineAsyncUmma", "PipelineUmmaAsync", + "PipelineClcFetchAsync", + "PipelineTmaMultiConsumersAsync", "PipelineTmaStore", "PipelineProducer", "PipelineConsumer", diff --git a/python/CuTeDSL/cutlass/pipeline/helpers.py b/python/CuTeDSL/cutlass/pipeline/helpers.py index a422037b..20479c15 100644 --- a/python/CuTeDSL/cutlass/pipeline/helpers.py +++ b/python/CuTeDSL/cutlass/pipeline/helpers.py @@ -102,6 +102,8 @@ class PipelineOp(enum.Enum): TCGen05Mma = enum.auto() # Tensor Memory Accelerator load TmaLoad = enum.auto() + # Cluster launch cancel response load + ClcLoad = enum.auto() # TMA Store consuming smem produced by AsyncThread TmaStore = enum.auto() # Composite of multiple PipelineOps @@ -161,12 +163,16 @@ class MbarrierArray(SyncObject): MbarrierArray implements an abstraction for an array of smem barriers. """ + @dsl_user_op def __init__( self, barrier_storage: cute.Pointer, num_stages: int, agent: tuple[PipelineOp, CooperativeGroup], tx_count: int = 0, + *, + loc=None, + ip=None, ) -> None: self.barrier_storage = barrier_storage self.tx_count = tx_count @@ -187,7 +193,7 @@ class MbarrierArray(SyncObject): self.mbarrier_base = self.barrier_storage # Mbarrier initialization in constructor - self.mbarrier_init() + self.mbarrier_init(loc=loc, ip=ip) def recast_to_new_op_type(self, new_op_type: PipelineOp) -> "MbarrierArray": """ @@ -207,6 +213,7 @@ class MbarrierArray(SyncObject): return new_mbarrier_array # Mbarrier initialization + @dsl_user_op def mbarrier_init(self, *, loc=None, ip=None) -> None: """ Initializes an array of mbarriers using warp 0. @@ -226,6 +233,7 @@ class MbarrierArray(SyncObject): if_generate(warp_idx == 0, then_body, loc=loc, ip=ip) + @dsl_user_op def arrive( self, index: int, @@ -257,7 +265,13 @@ class MbarrierArray(SyncObject): ) self.arrive_tcgen05mma(index, dst, cta_group, loc=loc, ip=ip) elif self.op_type in [PipelineOp.TmaLoad]: - self.arrive_and_expect_tx(index, self.tx_count) + # TMA operation signals local mbarrier only + self.arrive_and_expect_tx(index, self.tx_count, loc=loc, ip=ip) + elif self.op_type in [PipelineOp.ClcLoad]: + # Multiple threads in CTA 0 each signal a different remote CTA in cluster's mbarrier + self.arrive_and_expect_tx_with_dst( + index, self.tx_count, dst, loc=loc, ip=ip + ) elif self.op_type is PipelineOp.AsyncLoad: self.arrive_cp_async_mbarrier(index, loc=loc, ip=ip) else: @@ -265,6 +279,7 @@ class MbarrierArray(SyncObject): f"Error: MbarrierArray is not supported for PipelineOp: {_get_pipeline_op(self.op_type)}." ) + @dsl_user_op def arrive_mbarrier( self, index: int, dst_rank: Optional[int] = None, *, loc=None, ip=None ) -> None: @@ -277,11 +292,13 @@ class MbarrierArray(SyncObject): self.get_barrier(index, loc=loc, ip=ip), dst_rank, loc=loc, ip=ip ) + @dsl_user_op def arrive_cp_async_mbarrier(self, index: int, *, loc=None, ip=None): cute.arch.cp_async_mbarrier_arrive_noinc( self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip ) + @dsl_user_op def arrive_tcgen05mma( self, index: int, @@ -306,20 +323,34 @@ class MbarrierArray(SyncObject): ip=ip, ) - def arrive_and_expect_tx(self, index: int, tx_count: int) -> None: - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx(self.get_barrier(index), tx_count) + @dsl_user_op + def arrive_and_expect_tx( + self, index: int, tx_count: int, *, loc=None, ip=None + ) -> None: + with cute.arch.elect_one(loc=loc, ip=ip): + cute.arch.mbarrier_arrive_and_expect_tx( + self.get_barrier(index, loc=loc, ip=ip), tx_count, loc=loc, ip=ip + ) + @dsl_user_op + def arrive_and_expect_tx_with_dst( + self, index: int, tx_count: int, dst: Optional[int] = None, *, loc=None, ip=None + ) -> None: + cute.arch.mbarrier_arrive_and_expect_tx(self.get_barrier(index), tx_count, dst, loc=loc, ip=ip) + + @dsl_user_op def try_wait(self, index: int, phase: int, *, loc=None, ip=None) -> Boolean: return cute.arch.mbarrier_try_wait( self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip ) + @dsl_user_op def wait(self, index: int, phase: int, *, loc=None, ip=None) -> None: cute.arch.mbarrier_wait( self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip ) + @dsl_user_op def arrive_and_wait( self, index: int, @@ -333,9 +364,11 @@ class MbarrierArray(SyncObject): arrive(index, dst, cta_group, loc=loc, ip=ip) wait(index, phase, loc=loc, ip=ip) + @dsl_user_op def arrive_and_drop(self, *, loc=None, ip=None) -> None: raise NotImplementedError("Error: Not yet supported.") + @dsl_user_op def get_barrier(self, index: int, *, loc=None, ip=None) -> cute.Pointer: return self.mbarrier_base + index @@ -440,12 +473,15 @@ class NamedBarrier(SyncObject): ip=ip, ) + @dsl_user_op def arrive_and_drop(self, *, loc=None, ip=None) -> None: raise NotImplementedError("Error: Not supported.") + @dsl_user_op def sync(self, *, loc=None, ip=None) -> None: self.arrive_and_wait() + @dsl_user_op def get_barrier(self, *, loc=None, ip=None) -> int: return self.barrier_id @@ -490,6 +526,7 @@ class TmaStoreFence(SyncObject): raise NotImplementedError("Error: Not supported.") # TmaStoreFence doesn't have mbarriers + @dsl_user_op def get_barrier(self, *, loc=None, ip=None) -> None: assert False, ( "Error: TmaStoreFence doesn't use mbarriers and cannot return a barrier." @@ -511,6 +548,7 @@ class TmaStoreFence(SyncObject): class PipelineUserType(enum.Enum): Producer = enum.auto() Consumer = enum.auto() + ProducerConsumer = enum.auto() ############################################################################## @@ -548,11 +586,12 @@ class PipelineState: def phase(self) -> Int32: return self._phase - def reset_count(self): - self._count = Int32(0) + @dsl_user_op + def reset_count(self, *, loc=None, ip=None): + self._count = Int32(0, loc=loc, ip=ip) @dsl_user_op - def advance(self, *, loc=None, ip=None): + def advance(self, *, loc=None, ip=None) -> None: self._index += 1 self._count += 1 @@ -613,18 +652,19 @@ class PipelineState: ) +@dsl_user_op def make_pipeline_state(type: PipelineUserType, stages: int, *, loc=None, ip=None): """ Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1. """ - if type is PipelineUserType.Producer: + if type in (PipelineUserType.Producer, PipelineUserType.ProducerConsumer): return PipelineState( stages, Int32(0, loc=loc, ip=ip), Int32(0, loc=loc, ip=ip), Int32(1, loc=loc, ip=ip), ) - elif type is PipelineUserType.Consumer: + elif type in (PipelineUserType.Consumer, PipelineUserType.ProducerConsumer): return PipelineState( stages, Int32(0, loc=loc, ip=ip), @@ -642,6 +682,7 @@ def make_pipeline_state(type: PipelineUserType, stages: int, *, loc=None, ip=Non ############################################################################## +@dsl_user_op def pipeline_init_arrive( cluster_shape_mn: Optional[cute.Layout] = None, is_relaxed: bool = False, @@ -667,6 +708,7 @@ def pipeline_init_arrive( cute.arch.cluster_arrive(loc=loc, ip=ip) +@dsl_user_op def pipeline_init_wait( cluster_shape_mn: Optional[cute.Layout] = None, *, loc=None, ip=None ): @@ -682,11 +724,13 @@ def pipeline_init_wait( cute.arch.cluster_wait(loc=loc, ip=ip) +@dsl_user_op def _sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None): warnings.warn("_sync is deprecated. Please use agent_sync instead.") agent_sync(group, is_relaxed, loc=loc, ip=ip) +@dsl_user_op def agent_sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None): """ Syncs all threads within an agent. @@ -720,6 +764,7 @@ def _mbarrier_i64_to_ptr(val: Int64) -> cute.Pointer: # NamedBarrier free functions +@dsl_user_op def arrive(barrier_id: int, num_threads: int, *, loc=None, ip=None): """ The aligned flavor of arrive is used when all threads in the CTA will execute the @@ -730,6 +775,7 @@ def arrive(barrier_id: int, num_threads: int, *, loc=None, ip=None): ) +@dsl_user_op def arrive_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None): """ The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA. @@ -745,6 +791,7 @@ def arrive_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None): ) +@dsl_user_op def wait(barrier_id: int, num_threads: int): """ NamedBarriers do not have a standalone wait like mbarriers, only an arrive_and_wait. @@ -759,6 +806,7 @@ def wait(barrier_id: int, num_threads: int): arrive_and_wait(loc=loc, ip=ip) +@dsl_user_op def wait_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None): warnings.warn( "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()." @@ -774,11 +822,13 @@ def wait_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None): ) +@dsl_user_op def arrive_and_wait(barrier_id: int, num_threads: int, *, loc=None, ip=None): cute.arch.barrier( barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip ) +@dsl_user_op def sync(barrier_id: int = 0, *, loc=None, ip=None): cute.arch.barrier(barrier_id=barrier_id, loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/pipeline/sm100.py b/python/CuTeDSL/cutlass/pipeline/sm100.py index a962bdde..76a17b24 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm100.py +++ b/python/CuTeDSL/cutlass/pipeline/sm100.py @@ -14,12 +14,15 @@ from typing import Optional import cutlass import cutlass.cute as cute -from cutlass.cutlass_dsl import Boolean, if_generate +from cutlass.cutlass_dsl import Boolean, Int32, if_generate, dsl_user_op from cutlass.pipeline import ( Agent, CooperativeGroup, PipelineOp, + SyncObject, + MbarrierArray, + TmaStoreFence, PipelineState, PipelineAsync, agent_sync, @@ -39,23 +42,66 @@ class PipelineTmaUmma(PipelineAsync): is_leader_cta: bool cta_group: cute.nvgpu.tcgen05.CtaGroup + @dsl_user_op + @staticmethod + def _make_sync_object( + barrier_storage: cute.Pointer, + num_stages: int, + agent: tuple[PipelineOp, CooperativeGroup], + tx_count: int = 0, + *, + loc=None, + ip=None, + ) -> SyncObject: + """ + Returns a SyncObject corresponding to an agent's PipelineOp. + """ + if agent[0] in [ + PipelineOp.AsyncThread, + PipelineOp.TmaLoad, + PipelineOp.TCGen05Mma, + PipelineOp.Composite, + PipelineOp.AsyncLoad, + PipelineOp.ClcLoad, + ]: + return MbarrierArray( + barrier_storage=barrier_storage, + num_stages=num_stages, + agent=agent, + tx_count=tx_count, + loc=loc, + ip=ip, + ) + elif agent[0] is PipelineOp.TmaStore: + # Path taken for AsyncTmaStore + return TmaStoreFence(num_stages=num_stages) + else: + assert False, "Error: Invalid PipelineOp specified." + + @dsl_user_op @staticmethod def _compute_mcast_arrival_mask( - cta_layout_vmnk: cute.Layout, mcast_mode_mn: tuple[int, int] + cta_layout_vmnk: cute.Layout, + mcast_mode_mn: tuple[int, int], + *, + loc=None, + ip=None, ): """ Computes a mask for signaling arrivals to multicasting threadblocks. """ cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + cute.arch.block_idx_in_cluster(loc=loc, ip=ip), loc=loc, ip=ip + ) + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord( + cta_rank_in_cluster, loc=loc, ip=ip ) - cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster) tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask( - cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2 + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2, loc=loc, ip=ip ) tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask( - cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1 + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1, loc=loc, ip=ip ) block_in_cluster_coord_vmnk_peer = ( @@ -63,10 +109,18 @@ class PipelineTmaUmma(PipelineAsync): *cta_in_cluster_coord_vmnk[1:], ) tma_mcast_mask_a_peer = cute.nvgpu.cpasync.create_tma_multicast_mask( - cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=2 + cta_layout_vmnk, + block_in_cluster_coord_vmnk_peer, + mcast_mode=2, + loc=loc, + ip=ip, ) tma_mcast_mask_b_peer = cute.nvgpu.cpasync.create_tma_multicast_mask( - cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=1 + cta_layout_vmnk, + block_in_cluster_coord_vmnk_peer, + mcast_mode=1, + loc=loc, + ip=ip, ) assert not (mcast_mode_mn[0] == 0 and mcast_mode_mn[1] == 0) @@ -82,21 +136,23 @@ class PipelineTmaUmma(PipelineAsync): assert mcast_mode_mn[0] == 1 return tma_mcast_mask_a | tma_mcast_mask_a_peer + @dsl_user_op @staticmethod - def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout): + def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout, *, loc=None, ip=None): """ Computes leader threadblocks for 2CTA kernels. For 1CTA, all threadblocks are leaders. """ - bidx, bidy, _ = cute.arch.block_idx() + bidx, bidy, _ = cute.arch.block_idx(loc=loc, ip=ip) mma_coord_vmnk = ( - bidx % cute.size(cta_layout_vmnk, mode=[0]), - bidx // cute.size(cta_layout_vmnk, mode=[0]), + bidx % cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip), + bidx // cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip), bidy, None, ) return mma_coord_vmnk[0] == 0 + @dsl_user_op @staticmethod def create( *, @@ -108,6 +164,8 @@ class PipelineTmaUmma(PipelineAsync): cta_layout_vmnk: Optional[cute.Layout] = None, mcast_mode_mn: tuple[int, int] = (1, 1), defer_sync: bool = False, + loc=None, + ip=None, ): """Creates and initializes a new PipelineTmaUmma instance. @@ -130,7 +188,7 @@ class PipelineTmaUmma(PipelineAsync): :rtype: PipelineTmaUmma """ if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -140,27 +198,39 @@ class PipelineTmaUmma(PipelineAsync): producer = (producer_type, producer_group) consumer = (consumer_type, consumer_group) - sync_object_full = PipelineAsync._make_sync_object( - barrier_storage.align(min_align=8), num_stages, producer, tx_count + sync_object_full = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8), + num_stages, + producer, + tx_count, + loc=loc, + ip=ip, ) - sync_object_empty = PipelineAsync._make_sync_object( - barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + sync_object_empty = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, + num_stages, + consumer, + loc=loc, + ip=ip, ) - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: # No mcast mask if not using clusters producer_mask = None # All threadblocks are leaders if not using clusters is_leader_cta = True else: producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask( - cta_layout_vmnk, mcast_mode_mn + cta_layout_vmnk, mcast_mode_mn, loc=loc, ip=ip + ) + is_leader_cta = PipelineTmaUmma._compute_is_leader_cta( + cta_layout_vmnk, loc=loc, ip=ip ) - is_leader_cta = PipelineTmaUmma._compute_is_leader_cta(cta_layout_vmnk) cta_group = ( cute.nvgpu.tcgen05.CtaGroup.ONE - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1 + if cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) == 1 else cute.nvgpu.tcgen05.CtaGroup.TWO ) @@ -168,7 +238,7 @@ class PipelineTmaUmma(PipelineAsync): if not defer_sync: cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: agent_sync(Agent.ThreadBlock) else: agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) @@ -183,25 +253,41 @@ class PipelineTmaUmma(PipelineAsync): cta_group, ) - def consumer_release(self, state: PipelineState): + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): """ UMMA consumer release buffer empty, cta_group needs to be provided. """ - self.sync_object_empty.arrive(state.index, self.consumer_mask, self.cta_group) + self.sync_object_empty.arrive( + state.index, self.consumer_mask, self.cta_group, loc=loc, ip=ip + ) def producer_acquire( - self, state: PipelineState, try_acquire_token: Optional[Boolean] = None + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, ): """ TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks. """ if_generate( try_acquire_token is None or try_acquire_token == 0, - lambda: self.sync_object_empty.wait(state.index, state.phase), + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) if_generate( self.is_leader_cta, - lambda: self.sync_object_full.arrive(state.index, self.producer_mask), + lambda: self.sync_object_full.arrive( + state.index, self.producer_mask, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) def producer_commit(self, state: PipelineState): @@ -219,51 +305,65 @@ class PipelineAsyncUmma(PipelineAsync): cta_group: cute.nvgpu.tcgen05.CtaGroup + @dsl_user_op @staticmethod - def _compute_leading_cta_rank(cta_v_size): + def _compute_leading_cta_rank(cta_v_size, *, loc=None, ip=None): """ Computes the leading CTA rank. """ cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + cute.arch.block_idx_in_cluster(loc=loc, ip=ip), + loc=loc, + ip=ip, ) return cta_rank_in_cluster // cta_v_size * cta_v_size + @dsl_user_op @staticmethod - def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout): + def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout, *, loc=None, ip=None): """ Computes leader threadblocks for 2CTA kernels. For 1CTA, all threadblocks are leaders. """ - bidx, bidy, _ = cute.arch.block_idx() + bidx, bidy, _ = cute.arch.block_idx(loc=loc, ip=ip) mma_coord_vmnk = ( - bidx % cute.size(cta_layout_vmnk, mode=[0]), - bidx // cute.size(cta_layout_vmnk, mode=[0]), + bidx % cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip), + bidx // cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip), bidy, None, ) return mma_coord_vmnk[0] == 0 + @dsl_user_op @staticmethod - def _compute_peer_cta_mask(cta_layout_vmnk: cute.Layout): + def _compute_peer_cta_mask(cta_layout_vmnk: cute.Layout, *, loc=None, ip=None): """ Computes a mask for signaling arrivals to multicasting threadblocks. """ cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + cute.arch.block_idx_in_cluster(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord( + cta_rank_in_cluster, loc=loc, ip=ip ) - cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster) mask_self = cute.nvgpu.cpasync.create_tma_multicast_mask( - cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=0 + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=0, loc=loc, ip=ip ) block_in_cluster_coord_vmnk_peer = ( cta_in_cluster_coord_vmnk[0] ^ 1, *cta_in_cluster_coord_vmnk[1:], ) mask_peer = cute.nvgpu.cpasync.create_tma_multicast_mask( - cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=0 + cta_layout_vmnk, + block_in_cluster_coord_vmnk_peer, + mcast_mode=0, + loc=loc, + ip=ip, ) return mask_self | mask_peer + @dsl_user_op @staticmethod def create( *, @@ -273,6 +373,8 @@ class PipelineAsyncUmma(PipelineAsync): barrier_storage: cute.Pointer = None, cta_layout_vmnk: Optional[cute.Layout] = None, defer_sync: bool = False, + loc=None, + ip=None, ): """Creates and initializes a new PipelineAsyncUmma instance. @@ -291,7 +393,7 @@ class PipelineAsyncUmma(PipelineAsync): :rtype: PipelineAsyncUmma """ if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -301,37 +403,53 @@ class PipelineAsyncUmma(PipelineAsync): producer = (producer_type, producer_group) consumer = (consumer_type, consumer_group) - sync_object_full = PipelineAsync._make_sync_object( + sync_object_full = PipelineTmaUmma._make_sync_object( barrier_storage.align(min_align=8), num_stages, producer, + loc=loc, + ip=ip, ) - sync_object_empty = PipelineAsync._make_sync_object( - barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + sync_object_empty = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, + num_stages, + consumer, + loc=loc, + ip=ip, ) cta_v_size = ( - cute.size(cta_layout_vmnk, mode=[0]) if cta_layout_vmnk is not None else 1 + cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) + if cta_layout_vmnk is not None + else 1 ) cta_group = ( cute.nvgpu.tcgen05.CtaGroup.ONE - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1 + if cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) == 1 else cute.nvgpu.tcgen05.CtaGroup.TWO ) - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1: + if ( + cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) == 1 + ): # No mcast mask if we're not using 2CTA tcgen05 MMA producer_mask = None consumer_mask = None else: # If we're using 2CTA UMMAs, producer will arrive the mbar on leading CTA # We need to get the target cta_rank - producer_mask = PipelineAsyncUmma._compute_leading_cta_rank(cta_v_size) + producer_mask = PipelineAsyncUmma._compute_leading_cta_rank( + cta_v_size, loc=loc, ip=ip + ) # consumer needs to get the mask to signal - consumer_mask = PipelineAsyncUmma._compute_peer_cta_mask(cta_layout_vmnk) + consumer_mask = PipelineAsyncUmma._compute_peer_cta_mask( + cta_layout_vmnk, loc=loc, ip=ip + ) if not defer_sync: cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: agent_sync(Agent.ThreadBlock) else: agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) @@ -345,7 +463,8 @@ class PipelineAsyncUmma(PipelineAsync): cta_group, ) - def consumer_release(self, state: PipelineState): + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): """ UMMA consumer release buffer empty, cta_group needs to be provided. """ @@ -360,29 +479,38 @@ class PipelineUmmaAsync(PipelineAsync): cta_group: cute.nvgpu.tcgen05.CtaGroup + @dsl_user_op @staticmethod - def _compute_tmem_sync_mask(cta_layout_vmnk: cute.Layout): + def _compute_tmem_sync_mask(cta_layout_vmnk: cute.Layout, *, loc=None, ip=None): """ Computes a mask to signal completion of tmem buffers for 2CTA kernels. """ cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + cute.arch.block_idx_in_cluster(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord( + cta_rank_in_cluster, loc=loc, ip=ip ) - cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster) return cute.make_layout_image_mask( - cta_layout_vmnk, cta_in_cluster_coord_vmnk, mode=0 + cta_layout_vmnk, cta_in_cluster_coord_vmnk, mode=0, loc=loc, ip=ip ) + @dsl_user_op @staticmethod - def _compute_peer_cta_rank(): + def _compute_peer_cta_rank(*, loc=None, ip=None): """ Computes a mask to signal release of tmem buffers for 2CTA kernels. """ cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + cute.arch.block_idx_in_cluster(loc=loc, ip=ip), + loc=loc, + ip=ip, ) return cta_rank_in_cluster // 2 * 2 + @dsl_user_op @staticmethod def create( *, @@ -392,6 +520,8 @@ class PipelineUmmaAsync(PipelineAsync): barrier_storage: cute.Pointer = None, cta_layout_vmnk: Optional[cute.Layout] = None, defer_sync: bool = False, + loc=None, + ip=None, ): """Creates an instance of PipelineUmmaAsync with computed attributes. @@ -410,7 +540,7 @@ class PipelineUmmaAsync(PipelineAsync): :rtype: PipelineUmmaAsync """ if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -420,34 +550,44 @@ class PipelineUmmaAsync(PipelineAsync): producer = (producer_type, producer_group) consumer = (consumer_type, consumer_group) - sync_object_full = PipelineAsync._make_sync_object( - barrier_storage.align(min_align=8), num_stages, producer + sync_object_full = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8), num_stages, producer, loc=loc, ip=ip ) - sync_object_empty = PipelineAsync._make_sync_object( - barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + sync_object_empty = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, + num_stages, + consumer, + loc=loc, + ip=ip, ) - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: # Set mask to None if not using clusters (i.e. 1CTA kernels) producer_mask = None else: - producer_mask = PipelineUmmaAsync._compute_tmem_sync_mask(cta_layout_vmnk) + producer_mask = PipelineUmmaAsync._compute_tmem_sync_mask( + cta_layout_vmnk, loc=loc, ip=ip + ) - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1: + if ( + cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) == 1 + ): # Set mask to None if not using 2CTA instructions consumer_mask = None else: - consumer_mask = PipelineUmmaAsync._compute_peer_cta_rank() + consumer_mask = PipelineUmmaAsync._compute_peer_cta_rank(loc=loc, ip=ip) cta_group = ( cute.nvgpu.tcgen05.CtaGroup.ONE - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1 + if cta_layout_vmnk is None + or cute.size(cta_layout_vmnk, mode=[0], loc=loc, ip=ip) == 1 else cute.nvgpu.tcgen05.CtaGroup.TWO ) if not defer_sync: cute.arch.mbarrier_init_fence() - if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, loc=loc, ip=ip) == 1: agent_sync(Agent.ThreadBlock) else: agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) @@ -461,14 +601,18 @@ class PipelineUmmaAsync(PipelineAsync): cta_group, ) - def producer_commit(self, state: PipelineState): + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): """ UMMA producer commit buffer full, cta_group needs to be provided. """ - self.sync_object_full.arrive(state.index, self.producer_mask, self.cta_group) + self.sync_object_full.arrive( + state.index, self.producer_mask, self.cta_group, loc=loc, ip=ip + ) + @dsl_user_op @cute.jit - def producer_tail(self, state: PipelineState): + def producer_tail(self, state: PipelineState, *, loc=None, ip=None): """ Make sure the last used buffer empty signal is visible to producer. Producer tail is usually executed by producer before exit, to avoid dangling @@ -477,8 +621,9 @@ class PipelineUmmaAsync(PipelineAsync): :param state: The pipeline state that points to next useful buffer :type state: PipelineState """ + bidx_in_cluster = cute.arch.block_idx_in_cluster(loc=loc, ip=ip) cta_rank_in_cluster = cute.arch.make_warp_uniform( - cute.arch.block_idx_in_cluster() + bidx_in_cluster, loc=loc, ip=ip ) is_leader_cta = cta_rank_in_cluster % 2 == 0 @@ -486,5 +631,375 @@ class PipelineUmmaAsync(PipelineAsync): # Assume state contains that next useful buffer # So we only need to advance to num_stages - 1 times to last used buffer for i in cutlass.range_constexpr(self.num_stages - 1): - state.advance() - self.producer_acquire(state) + state.advance(loc=loc, ip=ip) + self.producer_acquire(state, loc=loc, ip=ip) + + +@dataclass(frozen=True) +class PipelineClcFetchAsync: + """ + PipelineClcFetchAsync implements a producer-consumer pipeline for Cluster Launch + Control based dynamic scheduling. Both producer and consumer operate asynchronously + using barrier synchronization to coordinate across pipeline stages and cluster CTAs. + + - Producer: waits for empty buffer, signals full barrier with transection bytes + across all CTAs in cluster, hardware autosignals each CTA's mbarrier when + transaction bytes are written, then the satte advance to next buffer slot. + - Consumer: waits for full barrier, then load respinse from local SMEM, then + sigals CTA 0's empty barrier to allow buffer reuse. + """ + + sync_object_full: SyncObject + sync_object_empty: SyncObject + num_stages: int + producer_mask: Optional[Int32] + consumer_mask: Optional[Int32] + is_signalling_thread: Boolean + + @staticmethod + @cute.jit + def _init_full_barrier_arrive_signal(cta_layout_vmnk: cute.Layout, tidx: Int32): + """ + Computes producer barrier signaling parameters, returns destination CTA rank + (0 to cluster_size-1) based on thread ID, and a boolean flag indicating if + this thread participates in signaling. + + :param cta_layout_vmnk: Cluster layout defining CTA count + :param tidx: Thread ID within the CTA + """ + dst_rank = tidx % 32 + is_signalling_thread = dst_rank < cute.size(cta_layout_vmnk) + return dst_rank, is_signalling_thread + + @staticmethod + def create( + *, + num_stages: int, + producer_group: CooperativeGroup, + consumer_group: CooperativeGroup, + tx_count: int, + barrier_storage: cute.Pointer = None, + producer_mask: Int32 = None, + consumer_mask: Int32 = None, + cta_layout_vmnk: Optional[cute.Layout] = None, + defer_sync: bool = False, + ): + """ + This helper function computes any necessary attributes and returns an instance of PipelineClcFetchAsync. + :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers + :type barrier_storage: cute.Pointer + :param num_stages: Number of buffer stages for this pipeline + :type num_stages: int + :param producer_group: `CooperativeGroup` for the producer agent + :type producer_group: CooperativeGroup + :param consumer_group: `CooperativeGroup` for the consumer agent + :type consumer_group: CooperativeGroup + :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage + :type tx_count: int + :param producer_mask: Mask for signaling arrives for the producer agent, defaults to ``None`` + :type producer_mask: Int32, optional + :param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to ``None`` + :type consumer_mask: Int32, optional + """ + if not isinstance(barrier_storage, cute.Pointer): + raise TypeError( + f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" + ) + + producer_type = PipelineOp.ClcLoad + consumer_type = PipelineOp.AsyncThread + + producer = (producer_type, producer_group) + consumer = (consumer_type, consumer_group) + sync_object_full = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8), num_stages, producer, tx_count + ) + sync_object_empty = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + ) + + if cta_layout_vmnk is None: + cta_layout_vmnk = cute.make_layout((1, 1, 1, 1)) + + tidx, _, _ = cute.arch.thread_idx() + # All signalling happens from CTA 0's threads, each thread + # in CTA 0 signals a different remote CTA's mbarrier. + (producer_mask, is_signalling_thread) = ( + PipelineClcFetchAsync._init_full_barrier_arrive_signal( + cta_layout_vmnk, tidx + ) + ) + + # The producer (sched warp) runs ONLY in CTA 0, all consumers + # across the cluster must arrive at CTA 0's empty barrier + consumer_mask = 0 + + if not defer_sync: + cute.arch.mbarrier_init_fence() + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + agent_sync(Agent.ThreadBlock) + else: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) + + return PipelineClcFetchAsync( + sync_object_full, + sync_object_empty, + num_stages, + producer_mask, + consumer_mask, + is_signalling_thread, + ) + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + """ + Producer acquire waits for empty buffer and sets transaction expectation on full barrier. + + :param state: Pipeline state pointing to the current buffer stage + :param try_acquire_token: Optional token to skip the empty barrier wait + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + if_generate( + self.is_signalling_thread, + lambda: self.sync_object_full.arrive( + state.index, self.producer_mask, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + + @dsl_user_op + def consumer_wait( + self, + state: PipelineState, + try_wait_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + """ + Consumer waits for full barrier to be signaled by hardware multicast. + + :param state: Pipeline state pointing to the current buffer stage + :param try_wait_token: Optional token to skip the full barrier wait + """ + if_generate( + try_wait_token is None or try_wait_token == 0, + lambda: self.sync_object_full.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + self.sync_object_empty.arrive(state.index, self.consumer_mask, loc=loc, ip=ip) + + @dsl_user_op + def producer_get_barrier( + self, state: PipelineState, *, loc=None, ip=None + ) -> cute.Pointer: + return self.sync_object_full.get_barrier(state.index, loc=loc, ip=ip) + + @dsl_user_op + def producer_tail( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + """ + Ensures all in-flight buffers are released before producer exits. + + :param state: Pipeline state with current position in the buffer + :param try_acquire_token: Optional token to skip the empty barrier waits + + """ + for i in range(self.num_stages): + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + state.advance(loc=loc, ip=ip) + + +@dataclass(frozen=True) +class PipelineTmaMultiConsumersAsync(PipelineAsync): + """ + PipelineTmaMultiConsumersAsync is used for TMA producers and UMMA+Async consumers. + """ + + is_leader_cta: bool + sync_object_empty_umma: SyncObject + sync_object_empty_async: SyncObject + cta_group: cute.nvgpu.tcgen05.CtaGroup + + @staticmethod + def create( + *, + num_stages: int, + producer_group: CooperativeGroup, + consumer_group_umma: CooperativeGroup, + consumer_group_async: CooperativeGroup, + tx_count: int, + barrier_storage: cute.Pointer = None, + cta_layout_vmnk: Optional[cute.Layout] = None, + defer_sync: bool = False, + ): + """ + This helper function computes any necessary attributes and returns an instance of PipelineTmaMultiConsumersAsync. + :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers + :type barrier_storage: cute.Pointer + :param num_stages: Number of buffer stages for this pipeline + :type num_stages: Int32 + :param producer_group: `CooperativeGroup` for the producer agent + :type producer_group: CooperativeGroup + :param consumer_group_umma: `CooperativeGroup` for the UMMA consumer agent + :type consumer_group_umma: CooperativeGroup + :param consumer_group_async: `CooperativeGroup` for the AsyncThread consumer agent + :type consumer_group_async: CooperativeGroup + :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage + :type tx_count: int + :param cta_layout_vmnk: Layout of the cluster shape + :type cta_layout_vmnk: cute.Layout | None + """ + if not isinstance(barrier_storage, cute.Pointer): + raise TypeError( + f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" + ) + + producer_type = PipelineOp.TmaLoad + consumer_type = PipelineOp.Composite + consumer_type_umma = PipelineOp.TCGen05Mma + consumer_type_async = PipelineOp.AsyncThread + + if consumer_group_umma.agent != consumer_group_async.agent: + raise ValueError( + "UMMA and AsyncThread consumer groups must be the same agent" + ) + + if cta_layout_vmnk is not None and cute.size(cta_layout_vmnk) != 1: + raise ValueError( + f"PipelineTmaMultiConsumersAsync is not verified for cta_layout_vmnk != 1, cta_layout_vmnk:{cta_layout_vmnk}" + ) + + consumer_group = CooperativeGroup( + consumer_group_umma.agent, + consumer_group_umma.size + consumer_group_async.size, + ) + + producer = (producer_type, producer_group) + consumer = (consumer_type, consumer_group) + + sync_object_full = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8), num_stages, producer, tx_count + ) + sync_object_empty = PipelineTmaUmma._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + ) + sync_object_empty_umma = sync_object_empty.recast_to_new_op_type( + consumer_type_umma + ) + sync_object_empty_async = sync_object_empty.recast_to_new_op_type( + consumer_type_async + ) + + # No mcast mask if not using clusters + producer_mask = None + consumer_mask = None + # All threadblocks are leaders if not using clusters + is_leader_cta = True + cta_group = ( + cute.nvgpu.tcgen05.CtaGroup.ONE + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1 + else cute.nvgpu.tcgen05.CtaGroup.TWO + ) + + if not defer_sync: + cute.arch.mbarrier_init_fence() + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + agent_sync(Agent.ThreadBlock) + else: + agent_sync(Agent.ThreadBlockCluster, is_relaxed=True) + + return PipelineTmaMultiConsumersAsync( + sync_object_full, + sync_object_empty, + num_stages, + producer_mask, + consumer_mask, + is_leader_cta, + sync_object_empty_umma, + sync_object_empty_async, + cta_group, + ) + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + """ + TMA producer acquire waits on buffer empty and sets the transaction barrier for leader threadblocks. + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + if_generate( + self.is_leader_cta, + lambda: self.sync_object_full.arrive(state.index, self.producer_mask), + loc=loc, + ip=ip, + ) + + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): + """ + TMA producer commit is a noop since TMA instruction itself updates the transaction count. + """ + pass + + @dsl_user_op + def consumer_release( + self, state: PipelineState, op_type: PipelineOp, *, loc=None, ip=None + ): + if op_type == PipelineOp.TCGen05Mma: + self.sync_object_empty_umma.arrive( + state.index, self.consumer_mask, self.cta_group, loc=loc, ip=ip + ) + elif op_type == PipelineOp.AsyncThread: + self.sync_object_empty_async.arrive( + state.index, self.consumer_mask, loc=loc, ip=ip + ) + else: + raise ValueError(f"Invalid PipelineOp specified. op_type:{op_type}") diff --git a/python/CuTeDSL/cutlass/pipeline/sm90.py b/python/CuTeDSL/cutlass/pipeline/sm90.py index 2393409f..385f8fa1 100644 --- a/python/CuTeDSL/cutlass/pipeline/sm90.py +++ b/python/CuTeDSL/cutlass/pipeline/sm90.py @@ -134,7 +134,6 @@ class PipelineAsync: if agent[0] in [ PipelineOp.AsyncThread, PipelineOp.TmaLoad, - PipelineOp.TCGen05Mma, PipelineOp.Composite, PipelineOp.AsyncLoad, ]: @@ -183,7 +182,7 @@ class PipelineAsync: :rtype: PipelineAsync """ if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -212,38 +211,66 @@ class PipelineAsync: consumer_mask, ) + @dsl_user_op def producer_acquire( - self, state: PipelineState, try_acquire_token: Optional[Boolean] = None + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, ): if_generate( try_acquire_token is None or try_acquire_token == 0, - lambda: self.sync_object_empty.wait(state.index, state.phase), + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) - def producer_try_acquire(self, state: PipelineState): - return self.sync_object_empty.try_wait(state.index, state.phase) + @dsl_user_op + def producer_try_acquire(self, state: PipelineState, *, loc=None, ip=None): + return self.sync_object_empty.try_wait(state.index, state.phase, loc=loc, ip=ip) - def producer_commit(self, state: PipelineState): - self.sync_object_full.arrive(state.index, self.producer_mask) + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): + self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip) + @dsl_user_op def consumer_wait( - self, state: PipelineState, try_wait_token: Optional[Boolean] = None + self, + state: PipelineState, + try_wait_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, ): if_generate( try_wait_token is None or try_wait_token == 0, - lambda: self.sync_object_full.wait(state.index, state.phase), + lambda: self.sync_object_full.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) - def consumer_try_wait(self, state: PipelineState): - return self.sync_object_full.try_wait(state.index, state.phase) + @dsl_user_op + def consumer_try_wait(self, state: PipelineState, *, loc=None, ip=None): + return self.sync_object_full.try_wait(state.index, state.phase, loc=loc, ip=ip) - def consumer_release(self, state: PipelineState): - self.sync_object_empty.arrive(state.index, self.consumer_mask) + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + self.sync_object_empty.arrive(state.index, self.consumer_mask, loc=loc, ip=ip) - def producer_get_barrier(self, state: PipelineState) -> cute.Pointer: - return self.sync_object_full.get_barrier(state.index) + @dsl_user_op + def producer_get_barrier( + self, state: PipelineState, *, loc=None, ip=None + ) -> cute.Pointer: + return self.sync_object_full.get_barrier(state.index, loc=loc, ip=ip) - def producer_tail(self, state: PipelineState): + @dsl_user_op + def producer_tail(self, state: PipelineState, *, loc=None, ip=None): """ Make sure the last used buffer empty signal is visible to producer. Producer tail is usually executed by producer before exit, to avoid dangling @@ -255,20 +282,27 @@ class PipelineAsync: # Assume state contains that next useful buffer # So we only need to advance to num_stages - 1 times to last used buffer for i in range(self.num_stages - 1): - state.advance() - self.producer_acquire(state) + state.advance(loc=loc, ip=ip) + self.producer_acquire(state, loc=loc, ip=ip) - # Util methods to manage produer and consumer - def make_producer(self): - state = make_pipeline_state(PipelineUserType.Producer, self.num_stages) + # Util methods to manage producer and consumer + @dsl_user_op + def make_producer(self, *, loc=None, ip=None): + state = make_pipeline_state( + PipelineUserType.Producer, self.num_stages, loc=loc, ip=ip + ) return PipelineProducer(self, state, self.sync_object_full.cg) - def make_consumer(self): - state = make_pipeline_state(PipelineUserType.Consumer, self.num_stages) + @dsl_user_op + def make_consumer(self, *, loc=None, ip=None): + state = make_pipeline_state( + PipelineUserType.Consumer, self.num_stages, loc=loc, ip=ip + ) return PipelineConsumer(self, state, self.sync_object_empty.cg) - def make_participants(self): - return self.make_producer(), self.make_consumer() + @dsl_user_op + def make_participants(self, *, loc=None, ip=None): + return self.make_producer(loc=loc, ip=ip), self.make_consumer(loc=loc, ip=ip) @dataclass(frozen=True) @@ -341,7 +375,9 @@ class PipelineTmaAsync(PipelineAsync): @staticmethod @cute.jit def init_empty_barrier_arrive_signal( - cta_layout_vmnk: cute.Layout, tidx: Int32, mcast_mode_mn: tuple[int, int] = (1, 1) + cta_layout_vmnk: cute.Layout, + tidx: Int32, + mcast_mode_mn: tuple[int, int] = (1, 1), ): """Initialize the empty barrier arrive signal. @@ -430,7 +466,7 @@ class PipelineTmaAsync(PipelineAsync): :rtype: PipelineTmaAsync """ if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -496,10 +532,13 @@ class PipelineTmaAsync(PipelineAsync): lambda: self.sync_object_empty.wait( state.index, state.phase, loc=loc, ip=ip ), + loc=loc, + ip=ip, ) self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip) - def producer_commit(self, state: PipelineState): + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): """ TMA producer commit is a noop since TMA instruction itself updates the transaction count. """ @@ -564,7 +603,7 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): :rtype: PipelineTmaMultiConsumersAsync """ if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -580,7 +619,8 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): if cta_layout_vmnk is not None and cute.size(cta_layout_vmnk) != 1: raise ValueError( - f"PipelineTmaMultiConsumersAsync is not verified for cta_layout_vmnk != 1, cta_layout_vmnk:{cta_layout_vmnk}" + "PipelineTmaMultiConsumersAsync is not verified for cta_layout_vmnk != 1, " + f"cta_layout_vmnk:{cta_layout_vmnk}" ) consumer_group = CooperativeGroup( @@ -633,34 +673,54 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync): cta_group, ) + @dsl_user_op def producer_acquire( - self, state: PipelineState, try_acquire_token: Optional[Boolean] = None + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, ): """ TMA producer acquire waits on buffer empty and sets the transaction barrier for leader threadblocks. """ if_generate( try_acquire_token is None or try_acquire_token == 0, - lambda: self.sync_object_empty.wait(state.index, state.phase), + lambda: self.sync_object_empty.wait( + state.index, state.phase, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) if_generate( self.is_leader_cta, - lambda: self.sync_object_full.arrive(state.index, self.producer_mask), + lambda: self.sync_object_full.arrive( + state.index, self.producer_mask, loc=loc, ip=ip + ), + loc=loc, + ip=ip, ) - def producer_commit(self, state: PipelineState): + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): """ TMA producer commit is a noop since TMA instruction itself updates the transaction count. """ pass - def consumer_release(self, state: PipelineState, op_type: PipelineOp): + @dsl_user_op + def consumer_release( + self, state: PipelineState, op_type: PipelineOp, *, loc=None, ip=None + ): if op_type == PipelineOp.TCGen05Mma: self.sync_object_empty_umma.arrive( - state.index, self.consumer_mask, self.cta_group + state.index, self.consumer_mask, self.cta_group, loc=loc, ip=ip ) elif op_type == PipelineOp.AsyncThread: - self.sync_object_empty_async.arrive(state.index, self.consumer_mask) + self.sync_object_empty_async.arrive( + state.index, self.consumer_mask, loc=loc, ip=ip + ) else: raise ValueError(f"Invalid PipelineOp specified. op_type:{op_type}") @@ -766,7 +826,7 @@ class PipelineOrder: defer_sync: bool = False, ): if not isinstance(barrier_storage, cute.Pointer): - raise ValueError( + raise TypeError( f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" ) @@ -805,7 +865,7 @@ class PipelineOrder: signalling_id = (self.group_id + 1) % self.length idx = self.get_barrier_for_current_stage_idx(signalling_id) cute.arch.mbarrier_arrive( - self.sync_object_full.get_barrier(idx), loc=loc, ip=ip + self.sync_object_full.get_barrier(idx, loc=loc, ip=ip), loc=loc, ip=ip ) self.state.advance(loc=loc, ip=ip) @@ -813,7 +873,10 @@ class PipelineOrder: def wait(self, *, loc=None, ip=None): idx = self.get_barrier_for_current_stage_idx(self.group_id) cute.arch.mbarrier_wait( - self.sync_object_full.get_barrier(idx), self.state.phase, loc=loc, ip=ip + self.sync_object_full.get_barrier(idx, loc=loc, ip=ip), + self.state.phase, + loc=loc, + ip=ip, ) @@ -922,13 +985,14 @@ class PipelineProducer: self._ImmutableResourceHandle__immutable_state ) - def commit(self): + @dsl_user_op + def commit(self, *, loc=None, ip=None): """Signal that data production is complete for the current stage. This allows consumers to start processing the data. """ self.get_origin().producer_commit( - self._ImmutableResourceHandle__immutable_state + self._ImmutableResourceHandle__immutable_state, loc=loc, ip=ip ) def __init__(self, pipeline, state, group: CooperativeGroup): @@ -949,13 +1013,18 @@ class PipelineProducer: """Create a new Producer instance with the same state.""" return PipelineProducer(self.__pipeline, self.__state.clone(), self.__group) - def reset(self): + @dsl_user_op + def reset(self, *, loc=None, ip=None): """Reset the count of how many handles this producer has committed.""" - self.__state.reset_count() + self.__state.reset_count(loc=loc, ip=ip) + @dsl_user_op def acquire( self, try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, ) -> ImmutableResourceHandle: """Wait for the current buffer to be empty before producing data. This is a blocking operation. @@ -965,18 +1034,22 @@ class PipelineProducer: :return: A handle to the producer for committing the data :rtype: ImmutableResourceHandle """ - self.__pipeline.producer_acquire(self.__state, try_acquire_token) + self.__pipeline.producer_acquire( + self.__state, try_acquire_token, loc=loc, ip=ip + ) handle = PipelineProducer.ImmutableResourceHandle( self.__pipeline, self.__state.clone() ) return handle - def advance(self): + @dsl_user_op + def advance(self, *, loc=None, ip=None): """Move to the next pipeline stage.""" - self.__state.advance() + self.__state.advance(loc=loc, ip=ip) + @dsl_user_op def acquire_and_advance( - self, try_acquire_token: Optional[Boolean] = None + self, try_acquire_token: Optional[Boolean] = None, *, loc=None, ip=None ) -> ImmutableResourceHandle: """Acquire the current buffer and advance to the next pipeline stage. @@ -992,11 +1065,12 @@ class PipelineProducer: acquired buffer stage :rtype: ImmutableResourceHandle """ - handle = self.acquire(try_acquire_token) - self.advance() + handle = self.acquire(try_acquire_token, loc=loc, ip=ip) + self.advance(loc=loc, ip=ip) return handle - def try_acquire(self) -> Boolean: + @dsl_user_op + def try_acquire(self, *, loc=None, ip=None) -> Boolean: """Attempt to acquire the current buffer without blocking. This method tries to acquire the current buffer stage for producing data @@ -1006,9 +1080,12 @@ class PipelineProducer: :return: A boolean token indicating whether the buffer was successfully acquired :rtype: Boolean """ - return self.__pipeline.producer_try_acquire(self.__state) + return self.__pipeline.producer_try_acquire(self.__state, loc=loc, ip=ip) - def commit(self, handle: Optional[ImmutableResourceHandle] = None): + @dsl_user_op + def commit( + self, handle: Optional[ImmutableResourceHandle] = None, *, loc=None, ip=None + ): """Signal that data production is complete for the current stage. This allows consumers to start processing the data. @@ -1018,19 +1095,20 @@ class PipelineProducer: :raises AssertionError: If provided handle does not belong to this producer """ if handle is not None: - assert handle.get_origin() is self, ( + assert handle.get_origin() is self.__pipeline, ( "ResourceHandle does not belong to this PipelineProducer instance" ) - handle.commit() + handle.commit(loc=loc, ip=ip) else: - self.__pipeline.producer_commit(self.__state) + self.__pipeline.producer_commit(self.__state, loc=loc, ip=ip) - def tail(self): + @dsl_user_op + def tail(self, *, loc=None, ip=None): """Ensure all used buffers are properly synchronized before producer exit. This should be called before the producer finishes to avoid dangling signals. """ - self.__pipeline.producer_tail(self.__state) + self.__pipeline.producer_tail(self.__state, loc=loc, ip=ip) def __extract_mlir_values__(self): """Extract MLIR values from the current state. @@ -1099,12 +1177,13 @@ class PipelineConsumer: @dataclass(frozen=True) class ImmutableResourceHandle(ImmutableResourceHandle): - def release(self): + @dsl_user_op + def release(self, *, loc=None, ip=None): """Signal that data production is complete for the current stage. This allows consumers to start processing the data. """ self.get_origin().consumer_release( - self._ImmutableResourceHandle__immutable_state + self._ImmutableResourceHandle__immutable_state, loc=loc, ip=ip ) def __init__(self, pipeline, state: PipelineState, group: CooperativeGroup): @@ -1125,11 +1204,15 @@ class PipelineConsumer: """Create a new Consumer instance with the same state.""" return PipelineConsumer(self.__pipeline, self.__state.clone(), self.__group) - def reset(self): + @dsl_user_op + def reset(self, *, loc=None, ip=None): """Reset the count of how many handles this consumer has consumed.""" - self.__state.reset_count() + self.__state.reset_count(loc=loc, ip=ip) - def wait(self, try_wait_token: Optional[Boolean] = None) -> ImmutableResourceHandle: + @dsl_user_op + def wait( + self, try_wait_token: Optional[Boolean] = None, *, loc=None, ip=None + ) -> ImmutableResourceHandle: """Wait for data to be ready in the current buffer. This is a blocking operation that will not return until data is available. @@ -1140,22 +1223,24 @@ class PipelineConsumer: once data consumption is complete :rtype: ImmutableResourceHandle """ - self.__pipeline.consumer_wait(self.__state, try_wait_token) + self.__pipeline.consumer_wait(self.__state, try_wait_token, loc=loc, ip=ip) handle = PipelineConsumer.ImmutableResourceHandle( self.__pipeline, self.__state.clone() ) return handle - def advance(self): + @dsl_user_op + def advance(self, *, loc=None, ip=None): """Advance the consumer to the next pipeline stage. This updates the internal state to point to the next buffer in the pipeline. Should be called after consuming data from the current buffer. """ - self.__state.advance() + self.__state.advance(loc=loc, ip=ip) + @dsl_user_op def wait_and_advance( - self, try_wait_token: Optional[Boolean] = None + self, try_wait_token: Optional[Boolean] = None, *, loc=None, ip=None ) -> ImmutableResourceHandle: """Atomically wait for data and advance to next pipeline stage. @@ -1170,11 +1255,12 @@ class PipelineConsumer: once data consumption is complete :rtype: ImmutableResourceHandle """ - handle = self.wait(try_wait_token) - self.advance() + handle = self.wait(try_wait_token, loc=loc, ip=ip) + self.advance(loc=loc, ip=ip) return handle - def try_wait(self) -> Boolean: + @dsl_user_op + def try_wait(self, *, loc=None, ip=None) -> Boolean: """Non-blocking check if data is ready in the current buffer. This method provides a way to test if data is available without blocking. @@ -1183,19 +1269,22 @@ class PipelineConsumer: :return: True if data is ready to be consumed, False if the buffer is not yet ready :rtype: Boolean """ - return self.__pipeline.consumer_try_wait(self.__state) + return self.__pipeline.consumer_try_wait(self.__state, loc=loc, ip=ip) - def release(self, handle: Optional[ImmutableResourceHandle] = None): + @dsl_user_op + def release( + self, handle: Optional[ImmutableResourceHandle] = None, *, loc=None, ip=None + ): """Signal that data consumption is complete for the current stage. This allows producers to start producing new data. """ if handle is not None: - assert handle.get_origin() is self, ( + assert handle.get_origin() is self.__pipeline, ( "ResourceHandle does not belong to this PipelineConsumer instance" ) - handle.release() + handle.release(loc=loc, ip=ip) else: - self.__pipeline.consumer_release(self.__state) + self.__pipeline.consumer_release(self.__state, loc=loc, ip=ip) def __extract_mlir_values__(self): """Extract MLIR values from the current state. diff --git a/python/CuTeDSL/cutlass/torch.py b/python/CuTeDSL/cutlass/torch.py index a795cbf6..754beed3 100644 --- a/python/CuTeDSL/cutlass/torch.py +++ b/python/CuTeDSL/cutlass/torch.py @@ -13,7 +13,7 @@ import ctypes from math import prod from dataclasses import dataclass from enum import Enum -from typing import Optional, Type, Union +from typing import Optional, Type, Union, Tuple, Literal from cutlass.cute.typing import ( Numeric, @@ -183,7 +183,7 @@ def convert_cute_tensor( ) -> Tensor: """ Change the value of the cute tensor to make its value converted from a fp32 torch tensor. - Used for fp8 and int4 types tensor creatation now. + Used for fp8 and int4 types tensor creation now. """ # if torch_tensor is on cpu, create a gpu copy if f32_torch_tensor.device.type == "cpu": @@ -289,7 +289,9 @@ def cute_tensor_like( assumed_align: Optional[int] = None, ) -> tuple[Tensor, torch.Tensor]: """ - Create a cute tensor use a torch tensor as the data source + Create a cute tensor use a torch tensor as the data source. + + The cute tensor is a managed reference to the torch tensor. :param data_ref: torch tensor as the data source :param cutlass_dtype: cutlass dtype of the cute tensor @@ -314,9 +316,11 @@ def cute_tensor_like( leading_dim = get_leading_dim(torch_tensor) cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + is_empty_tensor = torch_tensor.numel() == 0 # initialize the cute tensor data - if (cutlass_dtype.is_float and cutlass_dtype.width <= 8) or ( - cutlass_dtype.is_integer and cutlass_dtype.width == 4 + if not is_empty_tensor and ( + (cutlass_dtype.is_float and cutlass_dtype.width <= 8) + or (cutlass_dtype.is_integer and cutlass_dtype.width == 4) ): cute_tensor = convert_cute_tensor( data_ref.to(dtype=torch.float32), @@ -327,3 +331,32 @@ def cute_tensor_like( else: torch_tensor.copy_(data_ref.to(dtype=torch_dtype)) return cute_tensor, torch_tensor + + +def prepare_tensors_for_gemm( + mnkl: Tuple[int, int, int, int] | Tuple[int, int, int], + a_dtype: Type[Numeric], + b_dtype: Type[Numeric], + c_dtype: Type[Numeric], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Prepare tensors for GEMM + """ + if len(mnkl) == 4: + m, n, k, l = mnkl + a = torch.empty(l, m, k, dtype=dtype(a_dtype), device="cuda").permute(1, 2, 0) + b = torch.empty(l, n, k, dtype=dtype(b_dtype), device="cuda").permute(1, 2, 0) + c = torch.empty(l, m, n, dtype=dtype(c_dtype), device="cuda").permute(1, 2, 0) + elif len(mnkl) == 3: + m, n, k = mnkl + a = torch.empty(m, k, dtype=dtype(a_dtype), device="cuda") + b = torch.empty(n, k, dtype=dtype(b_dtype), device="cuda") + c = torch.empty(m, n, dtype=dtype(c_dtype), device="cuda") + else: + raise ValueError(f"mnkl must be a tuple of length 3 or 4, but got {mnkl}") + + a = a.random_(-2, 2) + b = b.random_(-2, 2) + c = c.random_(-2, 2) + + return a, b, c diff --git a/python/CuTeDSL/cutlass/utils/README.md b/python/CuTeDSL/cutlass/utils/README.md index 3a583ed4..b4a84681 100644 --- a/python/CuTeDSL/cutlass/utils/README.md +++ b/python/CuTeDSL/cutlass/utils/README.md @@ -4,6 +4,7 @@ This folder contains various utilties for kernel authoring. Specifically, the im followings can be considered experimental and subject to breaking changes: - static persistent tile scheduler defined in [`static_persistent_tile_scheduler.py`](./static_persistent_tile_scheduler.py) +- dynamic persistent tile scheduler defined in [`dynamic_persistent_tile_scheduler.py`](./dynamic_persistent_tile_scheduler.py) - pipeline abstractions defined in [`pipeline.py`](./pipeline.py) - grouped GEMM utilties defined [`grouped_gemm_tile_scheduler_helper.py`](./grouped_gemm_tile_scheduler_helper.py) and [`tensormap_manager.py`](./tensormap_manager.py) diff --git a/python/CuTeDSL/cutlass/utils/__init__.py b/python/CuTeDSL/cutlass/utils/__init__.py index 89b1f003..adddfc98 100644 --- a/python/CuTeDSL/cutlass/utils/__init__.py +++ b/python/CuTeDSL/cutlass/utils/__init__.py @@ -16,15 +16,17 @@ from .static_persistent_tile_scheduler import ( StaticPersistentRuntimeTileScheduler, ) -from .hardware_info import ( - HardwareInfo, +from .dynamic_persistent_tile_scheduler import ( + ClcDynamicPersistentTileSchedulerParams, + ClcDynamicPersistentTileScheduler, ) +from .hardware_info import HardwareInfo + from .blackwell_helpers import ( compute_epilogue_tile_shape, get_smem_store_op, get_tmem_load_op, - get_num_tmem_alloc_cols, make_smem_layout_a, make_smem_layout_b, make_smem_layout_epi, @@ -49,11 +51,13 @@ from .blockscaled_layout import ( make_tmem_layout_sfb, ) -from .grouped_gemm_tile_scheduler_helper import ( +from .grouped_gemm_persistent_tile_scheduler import ( GroupSearchResult, GroupedGemmGroupSearchState, - GroupedGemmTileSchedulerHelper, create_initial_search_state, + GroupedWorkTileInfo, + StaticPersistentGroupTileScheduler, + GroupedGemmTileSchedulerHelper, ) from .tensormap_manager import ( @@ -62,10 +66,11 @@ from .tensormap_manager import ( ) from .smem_allocator import SmemAllocator, get_smem_capacity_in_bytes -from .tmem_allocator import TmemAllocator +from .tmem_allocator import TmemAllocator, get_num_tmem_alloc_cols from .layout import LayoutEnum +from . import gemm from . import distributed from .mixed_input_helpers import ( @@ -73,6 +78,9 @@ from .mixed_input_helpers import ( scale_tma_partition, transform_partition, scale_partition, + epilog_gmem_copy_and_partition, + epilog_smem_copy_and_partition, + epilog_tmem_copy_and_partition, get_gmem_layout_scale, get_smem_layout_scale, compute_smem_layout, @@ -80,7 +88,15 @@ from .mixed_input_helpers import ( get_tma_atom_kind, get_copy_atom_a_transform, is_valid_scale_granularity, + is_shuffle_a, + is_valid_tensor_alignment, + is_valid_mma_tiler_and_cluster_shape, get_divisibility, + create_initial_contiguous_group_search_state, + contiguous_group_search, + make_contiguous_group_work_tile_info, + cvt_tensor_a, + store_transformed_a, ) from . import hopper_helpers as sm90 @@ -92,6 +108,7 @@ __all__ = [ "get_smem_capacity_in_bytes", "SmemAllocator", "TmemAllocator", + "get_num_tmem_alloc_cols", "LayoutEnum", "WorkTileInfo", "PersistentTileSchedulerParams", @@ -116,10 +133,11 @@ __all__ = [ "get_copy_atom_a_transform", "is_valid_scale_granularity", "get_divisibility", + "epilogue_tma_store", + "epilogue", "compute_epilogue_tile_shape", "get_smem_store_op", "get_tmem_load_op", - "get_num_tmem_alloc_cols", "make_smem_layout_a", "make_smem_layout_b", "make_smem_layout_epi", @@ -129,5 +147,8 @@ __all__ = [ "sm100", "print_latex", "print_latex_tv", + "gemm", "distributed", + "ClcDynamicPersistentTileSchedulerParams", + "ClcDynamicPersistentTileScheduler", ] diff --git a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py index 6eb08bd8..f6fa0f84 100644 --- a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py +++ b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py @@ -9,8 +9,8 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from math import log2, ceil from typing import List, Type, Union, Tuple +from typing_extensions import deprecated from cutlass.cutlass_dsl import ( Float16, @@ -47,7 +47,6 @@ from cutlass.cute.nvgpu.tcgen05 import ( Ld32x32bOp, Repetition, Pack, - find_tmem_tensor_col_offset, SmemLayoutAtomKind, make_smem_layout_atom, tile_to_mma_shape, @@ -64,6 +63,21 @@ from cutlass.utils.layout import LayoutEnum OperandSource = Tcgen05OperandSource +@dsl_user_op +@deprecated("API is deprecated, use cutlass.utils.get_num_tmem_alloc_cols instead") +def get_num_tmem_alloc_cols( + tmem_tensors: Union[cute.Tensor, List[cute.Tensor]], + rounding=True, + *, + loc=None, + ip=None, +) -> int: + import cutlass.utils as utils + return utils.get_num_tmem_alloc_cols( + tmem_tensors, rounding, arch="sm_100", loc=loc, ip=ip + ) + + @dsl_user_op def compute_epilogue_tile_shape( cta_tile_shape: cute.Shape, @@ -305,13 +319,13 @@ def get_smem_store_op( op = StMatrix8x8x16bOp(is_m_major, 2) return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip) elif use_stmatrix_m16n8_4x: - op = StMatrix16x8x8bOp(4) + op = StMatrix16x8x8bOp(transpose=True, num_matrices=4) return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip) elif use_stmatrix_m16n8_2x: - op = StMatrix16x8x8bOp(2) + op = StMatrix16x8x8bOp(transpose=True, num_matrices=2) return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip) elif use_stmatrix_m16n8_1x: - op = StMatrix16x8x8bOp(1) + op = StMatrix16x8x8bOp(transpose=True, num_matrices=1) return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip) else: op = CopyUniversalOp() @@ -512,50 +526,6 @@ def get_tmem_load_op( raise ValueError() -def get_num_tmem_alloc_cols( - tmem_tensors: Union[cute.Tensor, List[cute.Tensor]], rounding=True -) -> int: - """Get the total number of TMEM allocation columns for the given TMEM tensors. - - :param tmem_tensors: The TMEM tensors to get the number of allocation columns for. - :type tmem_tensors: Union[cute.Tensor, List[cute.Tensor]] - :param rounding: Whether to round up the number of allocation columns to the nearest power of 2. - :type rounding: bool - - :return: The total number of TMEM allocation columns. - :rtype: int - - :raises ValueError: If the number of TMEM allocation columns exceeds the maximum capacity of 512 or is less than 32. - """ - # Turn tmem_tensors into a list - if isinstance(tmem_tensors, cute.Tensor): - tmem_tensors = [tmem_tensors] - - # For each tensor in tmem_tensors, find the tmem_tensor_col_offset - num_tmem_alloc_cols_per_tensor = [ - find_tmem_tensor_col_offset(t) for t in tmem_tensors - ] - - # Sum up the num_tmem_alloc_cols_per_tensor - num_tmem_alloc_cols = sum(num_tmem_alloc_cols_per_tensor) - - # Round up num_tmem_cols_total to the nearest power of 2 and make sure it is at least 32 - if rounding: - num_tmem_alloc_cols = max(1 << ceil(log2(num_tmem_alloc_cols)), 32) - - # Validate the number of TMEM allocation columns - SM100_TMEM_CAPACITY_COLUMNS = 512 - SM100_TMEM_MIN_ALLOC_COLUMNS = 32 - if ( - num_tmem_alloc_cols > SM100_TMEM_CAPACITY_COLUMNS - or num_tmem_alloc_cols < SM100_TMEM_MIN_ALLOC_COLUMNS - ): - raise ValueError( - f"TMEM allocation columns {num_tmem_alloc_cols} exceeds the maximum capacity of {SM100_TMEM_CAPACITY_COLUMNS} or less than {SM100_TMEM_MIN_ALLOC_COLUMNS}" - ) - return num_tmem_alloc_cols - - def get_smem_layout_atom_ab( major_mode: OperandMajorMode, element_type: Type[Numeric], @@ -1156,11 +1126,54 @@ def cluster_shape_to_tma_atom_SFB( ) +@dsl_user_op +def get_permutation_mnk( + tile_shape_mnk: cute.Shape, + sf_vec_size: int, + use_mxf8f6f4: bool, + *, + loc=None, + ip=None, +) -> Tuple[int, int, int]: + """ + Get the permutation of M, N, K for the tiled MMA. + + :param tile_shape_mnk: The shape of the tile + :type tile_shape_mnk: cute.Shape + :param sf_vec_size: The vector size of the Scale Factor. + :type sf_vec_size: int + :param use_mxf8f6f4: Whether to use MXF8F6F4 or MXF4NVF4. + :type use_mxf8f6f4: bool + + :return: The permutation of M, N, K + :rtype: Tuple[int, int, int] + + :raise ValueError: If the tile shape is not divisible by the sf_vec_size + """ + perm_m = min(tile_shape_mnk[0], 128) + # refer to C++ code: + # /include/cutlass/gemm/collective/builders/sm120_common.inl?ref_type=heads#L158 + if sf_vec_size == 32 or sf_vec_size == 16: + perm_n_shape = (8, 2, 2) + perm_n_stride = (1, 16, 8) + else: + raise ValueError(f"Unsupported sf_vec_size, got {sf_vec_size}") + + perm_n_layout = cute.make_layout(perm_n_shape, stride=perm_n_stride) + perm_k = 32 if use_mxf8f6f4 else 64 + permutation_mnk = ( + perm_m, + perm_n_layout, + perm_k, + ) + + return permutation_mnk + + __all__ = [ "compute_epilogue_tile_shape", "get_smem_store_op", "get_tmem_load_op", - "get_num_tmem_alloc_cols", "make_smem_layout_a", "make_smem_layout_b", "make_smem_layout_epi", @@ -1169,4 +1182,6 @@ __all__ = [ "cluster_shape_to_tma_atom_A", "cluster_shape_to_tma_atom_B", "cluster_shape_to_tma_atom_SFB", + "get_permutation_mnk", + "get_num_tmem_alloc_cols", # deprecated; use cutlass.utils.get_num_tmem_alloc_cols instead ] diff --git a/python/CuTeDSL/cutlass/utils/blockscaled_layout.py b/python/CuTeDSL/cutlass/utils/blockscaled_layout.py index 40072d68..49078db5 100644 --- a/python/CuTeDSL/cutlass/utils/blockscaled_layout.py +++ b/python/CuTeDSL/cutlass/utils/blockscaled_layout.py @@ -214,6 +214,7 @@ def make_smem_layout_sfb( return sfb_smem_layout_staged + @dsl_user_op def make_tmem_layout_sfa( tiled_mma: cute.TiledMma, diff --git a/python/CuTeDSL/cutlass/utils/distributed.py b/python/CuTeDSL/cutlass/utils/distributed.py index 6741259c..36c458f1 100644 --- a/python/CuTeDSL/cutlass/utils/distributed.py +++ b/python/CuTeDSL/cutlass/utils/distributed.py @@ -16,13 +16,8 @@ import cutlass.cute as cute from cutlass.cute.typing import Pointer, Int32 from cutlass.cutlass_dsl import T, dsl_user_op from cutlass._mlir import ir -from cutlass._mlir.dialects import llvm, nvvm -from cutlass._mlir.dialects.nvvm import ( - MemOrderKind, - MemScopeKind, - AtomicOpKind, -) - +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import Literal __all__ = [ # misc @@ -43,14 +38,12 @@ __all__ = [ @dsl_user_op -def atomicAdd(dst_ptr: Pointer, val: Int32, loc=None, ip=None) -> Int32: - return nvvm.atomicrmw( - T.i32(), - AtomicOpKind.ADD, - dst_ptr.llvm_ptr, - val.ir_value(loc=loc, ip=ip), - mem_order=MemOrderKind.RELAXED, - syncscope=MemScopeKind.SYS, +def atomicAdd(dst_ptr: Pointer, val: Int32, *, loc=None, ip=None) -> Int32: + return cute.arch.atomic_add( + ptr=dst_ptr.llvm_ptr, + val=val, + sem="relaxed", + scope="sys", loc=loc, ip=ip, ) @@ -273,41 +266,24 @@ def spin_lock_atom_cas_relaxed_wait( *, expected_val: Int32, reset_val: Int32, - scope: str, + scope: Literal["gpu", "sys"], loc=None, ip=None, ) -> None: """ wait on a spin lock until the expected count is reached. Reset flag to reset_val if the expected count is reached. """ - if scope == "gpu": - result = 0 - while result != expected_val: - result = nvvm.atomicrmw( - T.i32(), - AtomicOpKind.CAS, - lock_ptr.llvm_ptr, - Int32(reset_val).ir_value(loc=loc, ip=ip), - b=Int32(expected_val).ir_value(loc=loc, ip=ip), - mem_order=MemOrderKind.RELAXED, - syncscope=MemScopeKind.GPU, - loc=loc, - ip=ip, - ) - elif scope == "sys": - result = 0 - while result != expected_val: - result = nvvm.atomicrmw( - T.i32(), - AtomicOpKind.CAS, - lock_ptr.llvm_ptr, - Int32(reset_val).ir_value(loc=loc, ip=ip), - b=Int32(expected_val).ir_value(loc=loc, ip=ip), - mem_order=MemOrderKind.RELAXED, - syncscope=MemScopeKind.SYS, - loc=loc, - ip=ip, - ) + result = 0 + while result != expected_val: + result = cute.arch.atomic_cas( + ptr=lock_ptr.llvm_ptr, + cmp=Int32(expected_val), + val=Int32(reset_val), + sem="relaxed", + scope=scope, + loc=loc, + ip=ip, + ) ######################################################## diff --git a/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py new file mode 100644 index 00000000..dec19c90 --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/dynamic_persistent_tile_scheduler.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import Tuple + +from cutlass.cutlass_dsl import ( + Boolean, + Integer, + Int32, + min, + extract_mlir_values, + new_from_mlir_values, + dsl_user_op, + T, +) +from cutlass._mlir import ir +from cutlass.utils.static_persistent_tile_scheduler import ( + WorkTileInfo, +) +import cutlass.cute as cute + +class ClcDynamicPersistentTileSchedulerParams: + """A class to represent parameters for a dynamic persistent tile scheduler. + + This class is designed to manage and compute the layout of clusters and tiles + in a batched gemm problem. + + :ivar cluster_shape_mn: Shape of the cluster in (m, n) dimensions (K dimension cta count must be 1). + :type cluster_shape_mn: tuple + """ + + def __init__( + self, + problem_shape_ntile_mnl: cute.Shape, + cluster_shape_mnk: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the ClcDynamicPersistentTileSchedulerParams with the given parameters. + + :param problem_shape_ntile_mnl: The shape of the problem in terms of + number of CTA (Cooperative Thread Array) in (m, n, l) dimensions. + :type problem_shape_ntile_mnl: cute.Shape + :param cluster_shape_mnk: The shape of the cluster in (m, n) dimensions. + :type cluster_shape_mnk: cute.Shape + + :raises ValueError: If cluster_shape_k is not 1. + """ + + if cluster_shape_mnk[2] != 1: + raise ValueError(f"unsupported cluster_shape_k {cluster_shape_mnk[2]}") + + self.problem_shape_ntile_mnl = problem_shape_ntile_mnl + # cluster_shape_mnk is kept for reconstruction + self._cluster_shape_mnk = cluster_shape_mnk + self.cluster_shape_mn = cluster_shape_mnk[:2] + self._loc = loc + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.problem_shape_ntile_mnl, self._cluster_shape_mnk]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip( + [self.problem_shape_ntile_mnl, self._cluster_shape_mnk], self._values_pos + ): + obj_list.append(new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return ClcDynamicPersistentTileSchedulerParams( + *(tuple(obj_list)), loc=self._loc + ) + + @dsl_user_op + def get_grid_shape(self, *, loc=None, ip=None) -> Tuple[Integer, Integer, Integer]: + """ + Computes the grid shape based on the problem shape and cluster shape. + + :return: the grid is the CTA numbers that has aligned with cluster shape. + """ + + problem_ceiling_cta_mnl = cute.round_up( + self.problem_shape_ntile_mnl, self._cluster_shape_mnk + ) + return problem_ceiling_cta_mnl + +class ClcDynamicPersistentTileScheduler: + """A scheduler for dynamic persistent tile execution in CUTLASS/CuTe kernels. + + :ivar params: Tile schedule related params, including cluster shape. + :type params: ClcDynamicPersistentTileSchedulerParams + :ivar cta_id_in_cluster: ID of the CTA within its cluster + :type cta_id_in_cluster: cute.Coord + :ivar _num_tiles_executed: Counter for executed tiles + :type _num_tiles_executed: Int32 + """ + + def __init__( + self, + params: ClcDynamicPersistentTileSchedulerParams, + cta_id_in_cluster: cute.Coord, + num_tiles_executed: Int32, + clc_response_ptr: cute.Pointer, + block_idx: Tuple[Integer, Integer, Integer], + ): + """ + Initializes the ClcDynamicPersistentTileScheduler with the given parameters. + + :param params: Tile schedule related params, including cluster shape. + :type params: ClcDynamicPersistentTileSchedulerParams + :param cta_id_in_cluster: ID of the CTA within its cluster. + :type cta_id_in_cluster: cute.Coord + :param num_tiles_executed: Counter for executed tiles. + :type num_tiles_executed: Int32 + :param clc_response_ptr: Pointer of the clc rsponse. + :type clc_response_ptr: Tuple[Integer, Integer, Integer, Integer] + :param block_idx: The block index. + :type block_idx: Tuple[Integer, Integer, Integer] + """ + self.params = params + self.cta_id_in_cluster = cta_id_in_cluster + self._num_tiles_executed = num_tiles_executed + self._clc_response_ptr = clc_response_ptr + self._block_idx = block_idx + + def __extract_mlir_values__(self) -> list[ir.Value]: + values = extract_mlir_values(self.cta_id_in_cluster) + values.extend(extract_mlir_values(self._num_tiles_executed)) + values.extend(extract_mlir_values(self._clc_response_ptr)) + values.extend(extract_mlir_values(self._block_idx)) + return values + + def __new_from_mlir_values__( + self, values: list[ir.Value] + ) -> "ClcDynamicPersistentTileScheduler": + assert len(values) == 8 + new_cta_id_in_cluster = new_from_mlir_values( + self.cta_id_in_cluster, values[0:3] + ) + new_num_tiles_executed = new_from_mlir_values( + self._num_tiles_executed, [values[3]] + ) + new_clc_response_ptr = new_from_mlir_values(self._clc_response_ptr, [values[4]]) + new_block_idx = new_from_mlir_values(self._block_idx, values[5:8]) + return ClcDynamicPersistentTileScheduler( + self.params, + new_cta_id_in_cluster, + new_num_tiles_executed, + new_clc_response_ptr, + new_block_idx, + ) + + @dsl_user_op + @staticmethod + def create( + params: ClcDynamicPersistentTileSchedulerParams, + block_idx: Tuple[Integer, Integer, Integer], + grid_dim: Tuple[Integer, Integer, Integer], + clc_response_ptr: cute.Pointer, + *, + loc=None, + ip=None, + ): + """Initialize the dynamic persistent tile scheduler. + + :param params: Parameters for the persistent + tile scheduler. + :type params: ClcDynamicPersistentTileSchedulerParams + :param block_idx: The 3d block index in the format (bidx, bidy, bidz). + :type block_idx: Tuple[Integer, Integer, Integer] + :param grid_dim: The 3d grid dimensions for kernel launch. + :type grid_dim: Tuple[Integer, Integer, Integer] + + :return: A ClcDynamicPersistentTileScheduler object. + :rtype: ClcDynamicPersistentTileScheduler + """ + params = params + + bidx, bidy, bidz = block_idx + + # CTA id in the cluster + cta_id_in_cluster = ( + Int32(bidx % params.cluster_shape_mn[0]), + Int32(bidy % params.cluster_shape_mn[1]), + Int32(0), + ) + + # Initialize number of tiles executed to zero + num_tiles_executed = Int32(0) + # Initialize clc response pointer + clc_response_ptr = clc_response_ptr + # The block index + block_idx = block_idx + + return ClcDynamicPersistentTileScheduler( + params, + cta_id_in_cluster, + num_tiles_executed, + clc_response_ptr, + block_idx, + ) + + # called by host + @dsl_user_op + def get_grid_shape( + params: ClcDynamicPersistentTileSchedulerParams, + *, + loc=None, + ip=None, + ) -> Tuple[Integer, Integer, Integer]: + """Calculates the grid shape to be launched on GPU using problem shape, + threadblock shape, and active cluster size. + + :param params: Parameters for grid shape calculation. + :type params: ClcDynamicPersistentTileSchedulerParams + + :return: The calculated 3d grid shape. + :rtype: Tuple[Integer, Integer, Integer] + """ + + return params.get_grid_shape(loc=loc, ip=ip) + + @dsl_user_op + def work_tile_info_from_clc_response( + self, result_addr: Int32, *, loc=None, ip=None + ) -> WorkTileInfo: + """ + Simulates parsing CLC response data in Python. + result_addr: 16-byte response data (simulating shared memory access) + """ + m_idx, n_idx, l_idx, vld = cute.arch.clc_response(result_addr, loc=loc, ip=ip) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + cta_idx_in_cluster, cta_idy_in_cluster, _ = self.cta_id_in_cluster + cur_tile_coord = (m_idx + cta_idx_in_cluster, n_idx + cta_idy_in_cluster, l_idx) + return WorkTileInfo(cur_tile_coord, vld) + + @dsl_user_op + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + smem_addr = self._clc_response_ptr + work_tile = self.work_tile_info_from_clc_response(smem_addr) + return work_tile + + @dsl_user_op + def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo: + bidx, bidy, bidz = self._block_idx + return WorkTileInfo((bidx, bidy, bidz), True) + + @dsl_user_op + def advance_to_next_work(self, mbarrier_addr, loc=None, ip=None): + # Query new work tile + with cute.arch.elect_one(): + cute.arch.issue_clc_query( + mbarrier_addr, self._clc_response_ptr, loc=loc, ip=ip + ) + self._num_tiles_executed += Int32(1) + + @property + def num_tiles_executed(self) -> Int32: + return self._num_tiles_executed diff --git a/python/CuTeDSL/cutlass/utils/gemm/__init__.py b/python/CuTeDSL/cutlass/utils/gemm/__init__.py new file mode 100644 index 00000000..d9d89020 --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/gemm/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited + +from . import sm100 + +__all__ = [ + "sm100", +] diff --git a/python/CuTeDSL/cutlass/utils/gemm/sm100.py b/python/CuTeDSL/cutlass/utils/gemm/sm100.py new file mode 100644 index 00000000..1089a91a --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/gemm/sm100.py @@ -0,0 +1,917 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited + +from typing import Tuple, Union +import cutlass.cute as cute +from cutlass.cutlass_dsl import Int32, Boolean, Constexpr, const_expr +import cutlass.pipeline as pipeline +from cutlass.utils.static_persistent_tile_scheduler import StaticPersistentTileScheduler +from cutlass.utils.dynamic_persistent_tile_scheduler import ( + ClcDynamicPersistentTileScheduler, +) +from cutlass.utils.blackwell_helpers import get_tmem_load_op, get_smem_store_op +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.common import CacheEvictionPriority + +__all__ = [ + "epilogue_tma_store", + "epilogue", + "epilogue_tma_store_release_flag", + "epilogue_release_flag", +] + + +def transform_partitioned_tensor_layout(tensor: cute.Tensor) -> cute.Tensor: + """ + Transform MMA layout from ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, ...rest) + to ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), ...rest). + + This groups MMA_ATOM_M with MMA_M and MMA_ATOM_N with MMA_N. + + :param tensor: Input tensor with layout ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, ...rest) + :type tensor: cute.Tensor + :return: Transformed tensor with layout ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), ...rest) + :rtype: cute.Tensor + """ + layout = tensor.layout + shape = layout.shape + stride = layout.stride + + # Build new shape: ((shape[0][0], shape[1]), (shape[0][1], shape[2]), ...rest) + new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:]) + + # Build new stride: ((stride[0][0], stride[1]), (stride[0][1], stride[2]), ...rest) + new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:]) + + new_layout = cute.make_layout(shape=new_shape, stride=new_stride) + return cute.make_tensor(tensor.iterator, new_layout) + + +def epilogue_tmem_copy_and_partition( + gemm_kernel, + tidx: Int32, + tAcc: cute.Tensor, + tCgC: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[Boolean, bool], +) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param gemm_kernel: The kernel instance + :type gemm_kernel: Any + :param tidx: The thread index in epilogue warp groups + :type tidx: Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param tCgC: The global tensor C to be copied and partitioned + :type tCgC: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: Union[Boolean, bool] + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc: The accumulated tensor in register used to hold t2r results + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = get_tmem_load_op( + gemm_kernel.cta_tile_shape_mnk, + gemm_kernel.c_layout, + gemm_kernel.c_dtype, + gemm_kernel.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc, + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(tCgC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, gemm_kernel.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + +def epilogue_smem_copy_and_partition( + gemm_kernel, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: Int32, + sC: cute.Tensor, +) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = get_smem_store_op( + gemm_kernel.c_layout, gemm_kernel.c_dtype, gemm_kernel.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + +@cute.jit +def epilogue_tma_store( + gemm_kernel, + epi_tidx: Int32, + warp_idx: Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tma_atom_c: cute.CopyAtom, + # Input of epilogue + tCtAcc_base: cute.Tensor, + # Staging of epilogue + sC: cute.Tensor, + # Output of epilogue + tCgC_base: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: Union[StaticPersistentTileScheduler, ClcDynamicPersistentTileScheduler], + epilogue_op: Constexpr, + clc_pipeline: Union[pipeline.PipelineClcFetchAsync, None] = None, + clc_consumer_state: Union[pipeline.PipelineState, None] = None, +) -> None: + # Layout transformation for tCgC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) + tCgC = transform_partitioned_tensor_layout(tCgC_base) + + # Layout transformation for tCtAcc_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE) + tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base) + + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = epilogue_tmem_copy_and_partition( + gemm_kernel, epi_tidx, tCtAcc, tCgC, epi_tile, gemm_kernel.use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition( + gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, gemm_kernel.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(gemm_kernel.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=gemm_kernel.num_c_stage, producer_group=c_producer_group + ) + + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=gemm_kernel.epilog_sync_bar_id, + num_threads=32 * len(gemm_kernel.epilogue_warp_id), + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + # Check if tile_sched is StaticPersistentTileScheduler or any subclass inheriting from it + if const_expr(isinstance(tile_sched, StaticPersistentTileScheduler)): + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + elif const_expr(isinstance(tile_sched, ClcDynamicPersistentTileScheduler)): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + else: + # Not match + pass + + # Wait for C store complete + c_pipeline.producer_tail() + + +@cute.jit +def epilogue( + gemm_kernel, + epi_tidx: Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tCtAcc_base: cute.Tensor, + tCgC_base: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: Union[StaticPersistentTileScheduler, ClcDynamicPersistentTileScheduler], + epilogue_op: Constexpr, + tmem_dealloc_barrier: pipeline.NamedBarrier, + tCcC_base: cute.Tensor = None, + mC_mnl: cute.Tensor = None, + clc_pipeline: Union[pipeline.PipelineClcFetchAsync, None] = None, + clc_consumer_state: Union[pipeline.PipelineState, None] = None, +) -> None: + """ + Epilogue function that stores accumulator results directly to global memory. + Used when TMA store is not enabled. + + :param gemm_kernel: The kernel instance + :type gemm_kernel: Any + :param epi_tidx: Thread index in epilogue warp groups + :type epi_tidx: Int32 + :param acc_pipeline: Accumulator pipeline for async operations + :type acc_pipeline: pipeline.PipelineAsync + :param tiled_mma: The tiled MMA configuration + :type tiled_mma: cute.TiledMma + :param tCtAcc_base: Base accumulator tensor in tensor memory + :type tCtAcc_base: cute.Tensor + :param tCgC_base: The global memory tensor C to be copied and partitioned + :type tCgC_base: cute.Tensor + :param epi_tile: Epilogue tile configuration + :type epi_tile: cute.Tile + :param tile_sched: Tile scheduler for persistent scheduling + :type tile_sched: StaticPersistentTileScheduler + :param epilogue_op: Optional elementwise operation to apply + :type epilogue_op: Constexpr + :param tmem_dealloc_barrier: Barrier for tensor memory deallocation + :type tmem_dealloc_barrier: pipeline.NamedBarrier + :param alignment_bytes: Alignment bytes for global memory store + :type alignment_bytes: int + :param tCcC_base: Identity/coordinate tensor C + :type tCcC_base: cute.Tensor + :param mC_mnl: Global memory tensor C (full tensor for predicate computation) + :type mC_mnl: cute.Tensor + :param clc_pipeline: Pipeline for dynamic persistent tile scheduling + :type clc_pipeline: Union[pipeline.PipelineClcFetchAsync, None] + :param clc_consumer_state: Consumer state for dynamic persistent tile scheduling + :type clc_consumer_state: Union[pipeline.PipelineState, None] + """ + + # Layout transformation for tCgC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) + tCgC = transform_partitioned_tensor_layout(tCgC_base) + + # Layout transformation for tCtAcc_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE) + tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base) + + # + # Partition for epilogue + # + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = epilogue_tmem_copy_and_partition( + gemm_kernel, epi_tidx, tCtAcc, tCgC, epi_tile, gemm_kernel.use_2cta_instrs + ) + + gC_epi = cute.flat_divide(tCgC, epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, gemm_kernel.c_dtype + ) + + mclD = cute.max_common_layout( + tTR_rC.layout, tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].layout + ) + num_bits_per_copy = min( + tTR_gC_partitioned.iterator.alignment * 8, + cute.size(mclD) * gemm_kernel.c_dtype.width, + 256, + ) + + simt_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + gemm_kernel.c_dtype, + num_bits_per_copy=num_bits_per_copy, + l1c_evict_priority=CacheEvictionPriority.NO_ALLOCATE, + ) + use_predication = tCcC_base is not None and mC_mnl is not None + + if const_expr(use_predication): + # Layout transformation for tCcC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) + tCcC = transform_partitioned_tensor_layout(tCcC_base) + cC_epi = cute.flat_divide(tCcC, epi_tile) + tTR_cC_partitioned = thr_copy_t2r.partition_D(cC_epi) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, gemm_kernel.num_acc_stage + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # + # Pre-advance to next tile + # + if const_expr(isinstance(tile_sched, StaticPersistentTileScheduler)): + tile_sched.advance_to_next_work() + next_work_tile = tile_sched.get_current_work() + + # Get tile coord from current work tile + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + if const_expr(use_predication): + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_cC = tTR_cC_partitioned[ + ( + None, + None, + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + tTR_cC = cute.group_modes(tTR_cC, 3, cute.rank(tTR_cC)) + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in range(subtile_cnt): + # + # Get the destination and coordinate slices for this subtile + # + tTR_gC_subtile = tTR_gC[(None, None, None, subtile_idx)] + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + # Async arrive accumulator buffer empty + # Release early for perf + if subtile_idx == subtile_cnt - 1: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) + tTR_rC.store(acc_vec) + + if const_expr(use_predication): + # compute predicate + tTR_cC_subtile = tTR_cC[(None, None, None, subtile_idx)] + pred_C_shape = (1, *tTR_cC_subtile.shape[1:]) + pred_C = cute.make_rmem_tensor(pred_C_shape, Boolean) + for m_idx in range(tTR_cC_subtile.shape[1]): + for n_idx in range(tTR_cC_subtile.shape[2]): + vector_first_coord = tTR_cC_subtile[(0, m_idx, n_idx)] + pred_C[(0, m_idx, n_idx)] = cute.elem_less( + vector_first_coord, mC_mnl.shape + ) + # Store C to global memory with predication + cute.copy(simt_atom, tTR_rC, tTR_gC_subtile, pred=pred_C) + else: + # Store C directly to global memory + cute.copy(simt_atom, tTR_rC, tTR_gC_subtile) + + if const_expr(isinstance(tile_sched, StaticPersistentTileScheduler)): + work_tile = next_work_tile + elif const_expr(isinstance(tile_sched, ClcDynamicPersistentTileScheduler)): + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + + +@cute.jit +def epilogue_tma_store_release_flag( + gemm_kernel, + epi_tidx: Int32, + warp_idx: Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tma_atom_c: cute.CopyAtom, + # Input of epilogue + tCtAcc_base: cute.Tensor, + # Staging of epilogue + sC: cute.Tensor, + # Output of epilogue + tCgC_base: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: StaticPersistentTileScheduler, + epilogue_op: Constexpr, + flag_base: cute.Tensor, + flag_mem_scope: str, +) -> None: + # Layout transformation for tCgC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) + tCgC = transform_partitioned_tensor_layout(tCgC_base) + # Layout transformation for tCtAcc_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE) + tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base) + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = epilogue_tmem_copy_and_partition( + gemm_kernel, epi_tidx, tCtAcc, tCgC, epi_tile, gemm_kernel.use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, gemm_kernel.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition( + gemm_kernel, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + tCgC_epi = cute.flat_divide(tCgC, epi_tile) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, gemm_kernel.num_acc_stage + ) + + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(gemm_kernel.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=gemm_kernel.num_c_stage, producer_group=c_producer_group + ) + + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=gemm_kernel.epilog_sync_bar_id, + num_threads=32 * len(gemm_kernel.epilogue_warp_id), + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # + # Slice to per mma tile index + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_prev_subtiles + subtile_idx) % gemm_kernel.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory + # + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Set Per Output Tile Flag with Release + # + import cutlass.utils as utils + from cutlass._mlir.dialects.nvvm import ( + MemOrderKind, + MemScopeKind, + ) + + # 1D linear index of current output tile + tile_id_linear = Int32( + tile_sched._current_work_linear_idx + * cute.size(gemm_kernel.cluster_shape_mn) + + cute.arch.block_idx_in_cluster() + ) + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + # Update flag with release semantic with GPU scope + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = flag_base.iterator + tile_id_linear + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + order="release", + scope=flag_mem_scope, + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + +@cute.jit +def epilogue_release_flag( + gemm_kernel, + epi_tidx: Int32, + acc_pipeline: pipeline.PipelineAsync, + tiled_mma: cute.TiledMma, + tCtAcc_base: cute.Tensor, + tCgC_base: cute.Tensor, + epi_tile: cute.Tile, + tile_sched: StaticPersistentTileScheduler, + epilogue_op: Constexpr, + tmem_dealloc_barrier: pipeline.NamedBarrier, + flag_base: cute.Tensor, + flag_mem_scope: str, +) -> None: + """ + Epilogue function that stores accumulator results directly to global memory. + Used when TMA store is not enabled. + + :param gemm_kernel: The kernel instance + :type gemm_kernel: Any + :param epi_tidx: Thread index in epilogue warp groups + :type epi_tidx: Int32 + :param acc_pipeline: Accumulator pipeline for async operations + :type acc_pipeline: pipeline.PipelineAsync + :param tiled_mma: The tiled MMA configuration + :type tiled_mma: cute.TiledMma + :param tCtAcc_base: Base accumulator tensor in tensor memory + :type tCtAcc_base: cute.Tensor + :param tCgC_base: The global memory tensor C to be copied and partitioned + :type tCgC_base: cute.Tensor + :param epi_tile: Epilogue tile configuration + :type epi_tile: cute.Tile + :param tile_sched: Tile scheduler for persistent scheduling + :type tile_sched: StaticPersistentTileScheduler + :param epilogue_op: Optional elementwise operation to apply + :type epilogue_op: Constexpr + :param tmem_dealloc_barrier: Barrier for tensor memory deallocation + :type tmem_dealloc_barrier: pipeline.NamedBarrier + :param flag_base: Base flag tensor + :type flag_base: cute.Tensor + :param flag_mem_scope: Memory scope for flag + :type flag_mem_scope: str + """ + # Layout transformation for tCgC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, TILE_M, TILE_N, TILE_K) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), TILE_M, TILE_N, TILE_K) + tCgC = transform_partitioned_tensor_layout(tCgC_base) + # Layout transformation for tCtAcc_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE) + tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base) + # + # Partition for epilogue + # + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = epilogue_tmem_copy_and_partition( + gemm_kernel, epi_tidx, tCtAcc, tCgC, epi_tile, gemm_kernel.use_2cta_instrs + ) + + gC_epi = cute.flat_divide(tCgC, epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tTR_gC_partitioned = thr_copy_t2r.partition_D(gC_epi) + # (T2R, T2R_M, T2R_N) + tTR_rC = cute.make_rmem_tensor( + tTR_gC_partitioned[(None, None, None, 0, 0, 0, 0, 0)].shape, gemm_kernel.c_dtype + ) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), gemm_kernel.c_dtype) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, gemm_kernel.num_acc_stage + ) + + work_tile = tile_sched.initial_work_tile_info() + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_mnl = ( + cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape), + cur_tile_coord[1], + cur_tile_coord[2], + ) + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N) + tTR_gC = tTR_gC_partitioned[ + ( + None, + None, + None, + None, + None, + *mma_tile_coord_mnl, + ) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + # Async arrive accumulator buffer empty + # Release early for perf + if subtile_idx == subtile_cnt - 1: + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Convert to C type + # + acc_vec = tTR_rAcc.load() + acc_vec = epilogue_op(acc_vec.to(gemm_kernel.c_dtype)) + tTR_rC.store(acc_vec) + + # + # Store C directly to global memory + # + cute.copy(simt_atom, tTR_rC, tTR_gC[(None, None, None, subtile_idx)]) + + # + # Set Per Output Tile Flag with Release + # + import cutlass.utils as utils + + # 1D linear index of current output tile + tile_id_linear = Int32( + tile_sched._current_work_linear_idx + * cute.size(gemm_kernel.cluster_shape_mn) + + cute.arch.block_idx_in_cluster() + ) + # Wait for C store complete + # Unlike regular epilogue where we only wait C store complete once at end of each kernel. + # Here we need to wait for C store complete for each output tile before we set the release flag. + c_pipeline.producer_tail() + # Update flag with release semantic with GPU scope + if warp_idx == gemm_kernel.epilogue_warp_id[0]: + with cute.arch.elect_one(): + flag_curr_tile = flag_base.iterator + tile_id_linear + utils.distributed.multimem_red_add1( + lock_ptr=flag_curr_tile, + order="release", + scope=flag_mem_scope, + ) + + # + # Advance to next tile + # + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Synchronize before TMEM dealloc (done by the caller) + tmem_dealloc_barrier.arrive_and_wait() + diff --git a/python/CuTeDSL/cutlass/utils/gemm_helper.py b/python/CuTeDSL/cutlass/utils/gemm_helper.py deleted file mode 100644 index e69de29b..00000000 diff --git a/python/CuTeDSL/cutlass/utils/grouped_gemm_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/grouped_gemm_persistent_tile_scheduler.py new file mode 100644 index 00000000..227eafbe --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/grouped_gemm_persistent_tile_scheduler.py @@ -0,0 +1,1035 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from typing import List, Tuple + +import cutlass.cute as cute +from cutlass.cutlass_dsl import ( + Boolean, + Int32, + Integer, + extract_mlir_values, + new_from_mlir_values, + const_expr, + dsl_user_op, + min, +) +from cutlass._mlir import ir +from typing_extensions import deprecated + +from cutlass.utils import ( + PersistentTileSchedulerParams, + StaticPersistentTileScheduler, + WorkTileInfo, +) + + +class GroupSearchResult: + """ + The result of the group search for grouped gemm. + + :param group_idx: The result group index + :type group_idx: Int32 + :param cta_tile_idx_m: CTA tile index along M dimension after rasterization + :type cta_tile_idx_m: Int32 + :param cta_tile_idx_n: CTA tile index along N dimension after rasterization + :type cta_tile_idx_n: Int32 + :param problem_shape_m: The M dimension of the gemm problem + :type problem_shape_m: Int32 + :param problem_shape_n: The N dimension of the gemm problem + :type problem_shape_n: Int32 + :param problem_shape_k: The K dimension of the gemm problem + :type problem_shape_k: Int32 + :param cta_tile_count_k: Number of tiles along K dimension + :type cta_tile_count_k: Int32 + """ + + def __init__( + self, + group_idx: Int32, + cta_tile_idx_m: Int32, + cta_tile_idx_n: Int32, + problem_shape_m: Int32, + problem_shape_n: Int32, + problem_shape_k: Int32, + cta_tile_count_k: Int32, + ) -> None: + self.group_idx = group_idx + self.cta_tile_idx_m = cta_tile_idx_m + self.cta_tile_idx_n = cta_tile_idx_n + self.problem_shape_m = problem_shape_m + self.problem_shape_n = problem_shape_n + self.problem_shape_k = problem_shape_k + self.cta_tile_count_k = cta_tile_count_k + + def __extract_mlir_values__(self) -> List[ir.Value]: + values = extract_mlir_values(self.group_idx) + values.extend(extract_mlir_values(self.cta_tile_idx_m)) + values.extend(extract_mlir_values(self.cta_tile_idx_n)) + values.extend(extract_mlir_values(self.problem_shape_m)) + values.extend(extract_mlir_values(self.problem_shape_n)) + values.extend(extract_mlir_values(self.problem_shape_k)) + values.extend(extract_mlir_values(self.cta_tile_count_k)) + return values + + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GroupSearchResult": + assert len(values) == 7 + return GroupSearchResult(*tuple(values)) + + +class GroupedGemmGroupSearchState: + """ + The state of group index search for grouped gemm. + + The state will be initialized once and updated in every round of group index search. + + :param start_group_idx: The group idx to start the search with + :type start_group_idx: Int32 + :param tile_count_prev_group: Number of tiles before the matched group + :type tile_count_prev_group: Int32 + :param tile_count_searched: Number of tiles we have searched. When the matched group is found, + it records the number of tiles including the matched group + :type tile_count_searched: Int32 + """ + + def __init__( + self, + start_group_idx: Int32, + tile_count_prev_group: Int32, + tile_count_searched: Int32, + found: Boolean, + ) -> None: + self.start_group_idx = start_group_idx + self.tile_count_prev_group = tile_count_prev_group + self.tile_count_searched = tile_count_searched + self.found = Boolean(found) + + def __extract_mlir_values__(self) -> List[ir.Value]: + values = extract_mlir_values(self.start_group_idx) + values.extend(extract_mlir_values(self.tile_count_prev_group)) + values.extend(extract_mlir_values(self.tile_count_searched)) + values.extend(extract_mlir_values(self.found)) + return values + + def __new_from_mlir_values__( + self, values: List[ir.Value] + ) -> "GroupedGemmGroupSearchState": + assert len(values) == 4 + start_group_idx = new_from_mlir_values(self.start_group_idx, [values[0]]) + tile_count_prev_group = new_from_mlir_values( + self.tile_count_prev_group, [values[1]] + ) + tile_count_searched = new_from_mlir_values( + self.tile_count_searched, [values[2]] + ) + found = new_from_mlir_values(self.found, [values[3]]) + return GroupedGemmGroupSearchState( + start_group_idx, tile_count_prev_group, tile_count_searched, found + ) + + +def create_initial_search_state() -> GroupedGemmGroupSearchState: + """ + Create an initial search state for grouped gemm. + + :return: A new search state with initial values + :rtype: GroupedGemmGroupSearchState + """ + return GroupedGemmGroupSearchState( + start_group_idx=Int32(0), + tile_count_prev_group=Int32(0), + tile_count_searched=Int32(0), + found=Boolean(False), + ) + + +# Grouped Work Tile Information +class GroupedWorkTileInfo(WorkTileInfo): + """A class to represent information about a work tile. + + :ivar tile_idx: The index of the tile. + :type tile_idx: cute.Coord + :ivar is_valid_tile: Whether the tile is valid. + :type is_valid_tile: Boolean + :ivar group_search_result: Group work tile information. + :type group_search_result: GroupSearchResult + """ + + def __init__( + self, + tile_idx: cute.Coord, + is_valid_tile: Boolean, + group_search_result: GroupSearchResult, + ): + super().__init__(tile_idx, is_valid_tile) + self.group_search_result = group_search_result + + def __extract_mlir_values__(self) -> list[ir.Value]: + values = extract_mlir_values(self.tile_idx) + values.extend(extract_mlir_values(self.is_valid_tile)) + values.extend(extract_mlir_values(self.group_search_result)) + return values + + def __new_from_mlir_values__(self, values: list[ir.Value]) -> "GroupedWorkTileInfo": + if len(values) != 11: + raise ValueError("Length of mlir values extracted is incorrect.") + new_tile_idx = new_from_mlir_values(self._tile_idx, values[:3]) + new_is_valid_tile = new_from_mlir_values(self._is_valid_tile, [values[3]]) + new_group_search_result = new_from_mlir_values( + self.group_search_result, values[4:11] + ) + return GroupedWorkTileInfo( + new_tile_idx, new_is_valid_tile, new_group_search_result + ) + + +# Static Persistent Grouped GEMM +class StaticPersistentGroupTileScheduler(StaticPersistentTileScheduler): + """A scheduler for static persistent group-based tile execution in CUTLASS/CuTe kernels. + + :ivar params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl + :type params: PersistentTileSchedulerParams + :ivar num_persistent_clusters: Number of persistent clusters that can be launched + :type num_persistent_clusters: Int32 + :ivar _current_work_linear_idx: Current cluster index + :type _current_work_linear_idx: Int32 + :ivar cta_id_in_cluster: ID of the CTA within its cluster + :type cta_id_in_cluster: cute.Coord + :ivar _num_tiles_executed: Counter for executed tiles + :type _num_tiles_executed: Int32 + :ivar cluster_tile_shape_mnk: The shape of cluster tile as (m, n, k) + :type cluster_tile_shape_mnk: tuple[int, int, int] + :ivar search_state: The initial search state + :type search_state: GroupedGemmGroupSearchState + :ivar group_count: Number of groups in current grouped gemm problem + :type group_count: int + :ivar problem_shape_mnkl: Problem shape tensor for groups + :type problem_shape_mnkl: cute.Tensor + """ + + def __init__( + self, + params: PersistentTileSchedulerParams, + num_persistent_clusters: Int32, + current_work_linear_idx: Int32, + cta_id_in_cluster: cute.Coord, + num_tiles_executed: Int32, + cluster_tile_shape_mnk: tuple[int, int, int], + search_state: GroupedGemmGroupSearchState, + group_count: int, + problem_shape_mnkl: cute.Tensor, + ): + StaticPersistentTileScheduler.__init__( + self, + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + ) + self.group_count = group_count + self.lane_idx = cute.arch.lane_idx() + self.cluster_tile_shape_mnk = cluster_tile_shape_mnk + self.search_state = search_state + self.problem_shape_mnkl = problem_shape_mnkl + + def __extract_mlir_values__(self) -> list[ir.Value]: + values = extract_mlir_values(self.num_persistent_clusters) + values.extend(extract_mlir_values(self._current_work_linear_idx)) + values.extend(extract_mlir_values(self.cta_id_in_cluster)) + values.extend(extract_mlir_values(self._num_tiles_executed)) + values.extend(extract_mlir_values(self.search_state)) + values.extend(extract_mlir_values(self.problem_shape_mnkl)) + values.extend(extract_mlir_values(self.params)) + return values + + def __new_from_mlir_values__( + self, values: list[ir.Value] + ) -> "StaticPersistentGroupTileScheduler": + if len(values) < 11: + raise ValueError("Length of mlir values extracted is incorrect.") + new_num_persistent_clusters = new_from_mlir_values( + self.num_persistent_clusters, [values[0]] + ) + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[1]] + ) + new_cta_id_in_cluster = new_from_mlir_values( + self.cta_id_in_cluster, values[2:5] + ) + new_num_tiles_executed = new_from_mlir_values( + self._num_tiles_executed, [values[5]] + ) + search_state = new_from_mlir_values(self.search_state, values[6:10]) + problem_shape_mnkl = new_from_mlir_values(self.problem_shape_mnkl, [values[10]]) + params = new_from_mlir_values(self.params, values[11:]) + + return StaticPersistentGroupTileScheduler( + params, + new_num_persistent_clusters, + new_current_work_linear_idx, + new_cta_id_in_cluster, + new_num_tiles_executed, + self.cluster_tile_shape_mnk, + search_state, + self.group_count, + problem_shape_mnkl, + ) + + @staticmethod + @dsl_user_op + def create( + params: PersistentTileSchedulerParams, + block_idx: Tuple[Integer, Integer, Integer], + grid_dim: Tuple[Integer, Integer, Integer], + cluster_tile_shape_mnk: tuple[int, int, int], + initial_search_state: GroupedGemmGroupSearchState, + group_count: int, + problem_shape_mnkl: cute.Tensor, + *, + loc=None, + ip=None, + ): + """Initialize the static persistent group-based tile scheduler. + + :param params: Parameters for the persistent + tile scheduler. + :type params: PersistentTileSchedulerParams + :param block_idx: The 3d block index in the format (bidx, bidy, bidz). + :type block_idx: Tuple[Integer, Integer, Integer] + :param grid_dim: The 3d grid dimensions for kernel launch. + :type grid_dim: Tuple[Integer, Integer, Integer] + :param cluster_tile_shape_mnk: The shape of cluster tile as (m, n, k) + :type cluster_tile_shape_mnk: tuple[int, int, int] + :param initial_search_state: The initial search state + :type initial_search_state: GroupedGemmGroupSearchState + :param group_count: Number of groups in current grouped gemm problem + :type group_count: int + :param problem_shape_mnkl: Problem shape tensor for groups + :type problem_shape_mnkl: cute.Tensor + + :return: A StaticPersistentGroupTileScheduler object. + :rtype: StaticPersistentGroupTileScheduler + """ + + # Calculate the number of persistent clusters by dividing the total grid size + # by the number of CTAs per cluster + num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size( + params.cluster_shape_mn, loc=loc, ip=ip + ) + + bidx, bidy, bidz = block_idx + + # Initialize workload index equals to the cluster index in the grid + current_work_linear_idx = Int32(bidz) + + # CTA id in the cluster + cta_id_in_cluster = ( + Int32(bidx % params.cluster_shape_mn[0]), + Int32(bidy % params.cluster_shape_mn[1]), + Int32(0), + ) + + # Initialize number of tiles executed to zero + num_tiles_executed = Int32(0) + return StaticPersistentGroupTileScheduler( + params, + num_persistent_clusters, + current_work_linear_idx, + cta_id_in_cluster, + num_tiles_executed, + cluster_tile_shape_mnk, + initial_search_state, + group_count, + problem_shape_mnkl, + ) + + @dsl_user_op + @cute.jit + def _prefix_sum(self, value_per_thread: Int32, *, loc=None, ip=None) -> Int32: + """ + Perform prefix sum within a full warp. + + :param value_per_thread: The value for this thread to contribute to the prefix sum + :type value_per_thread: Int32 + :return: The prefix sum result for this thread + :rtype: Int32 + """ + clamp_value = 0 + idx = 1 + sum_per_thread = value_per_thread + while const_expr(idx < cute.arch.WARP_SIZE): + value = cute.arch.shuffle_sync_up( + sum_per_thread, idx, mask_and_clamp=clamp_value, loc=loc, ip=ip + ) + if self.lane_idx >= idx: + sum_per_thread += value + idx = idx << 1 + return sum_per_thread + + @dsl_user_op + def _get_problem_for_group( + self, problem_shape_mnkl: cute.Tensor, group_idx: Int32, *, loc=None, ip=None + ) -> cute.Tensor: + """ + Load gemm problem (m,n,k,l) for the specified group from global memory to register. + + :param problem_shape_mnkl: Tensor in global memory with layout (group_count, 4):(4, 1) + :type problem_shape_mnkl: cute.Tensor + :param group_idx: The index of the group to load + :type group_idx: Int32 + :return: The problem shape tensor for the specified group + :rtype: cute.Tensor + """ + cur_problem_mnkl = cute.make_rmem_tensor( + cute.make_layout(4), problem_shape_mnkl.element_type, loc=loc, ip=ip + ) + cute.autovec_copy( + problem_shape_mnkl[(group_idx, None)], cur_problem_mnkl, loc=loc, ip=ip + ) + return cur_problem_mnkl + + @dsl_user_op + def _get_cluster_tile_count_mn( + self, problem_shape: cute.Tensor, *, loc=None, ip=None + ) -> Int32: + """ + Compute total cluster count. + + :param problem_shape: Tensor containing problem shape (m, n, k, l) + :type problem_shape: cute.Tensor + :return: The total cluster tile count for M and N dimensions + :rtype: Int32 + """ + cur_ntile_m = ( + problem_shape[0] + self.cluster_tile_shape_mnk[0] - 1 + ) // self.cluster_tile_shape_mnk[0] + cur_ntile_n = ( + problem_shape[1] + self.cluster_tile_shape_mnk[1] - 1 + ) // self.cluster_tile_shape_mnk[1] + cur_ntile_mn = cur_ntile_m * cur_ntile_n + return cur_ntile_mn + + @dsl_user_op + def _compute_cta_tile_coord( + self, + cluster_tile_idx: Int32, + cta_tile_coord_in_cluster: tuple, + cluster_tile_count_m: Int32, + cluster_tile_count_n: Int32, + *, + loc=None, + ip=None, + ) -> tuple: + """ + Compute CTA tile indices along M and N dimensions based on the linear index within a group. + + It uses the AlongM mode to decompose the linear index onto M and N dimensions. + + :param cluster_tile_idx: The linear index within a group + :type cluster_tile_idx: Int32 + :param cta_tile_coord_in_cluster: CTA indices along M and N dimensions within a cluster + :type cta_tile_coord_in_cluster: tuple of Int32 + :param cluster_tile_count_m: The number of clusters along M dimension of the matched group + :type cluster_tile_count_m: Int32 + :param cluster_tile_count_n: The number of clusters along N dimension of the matched group + :type cluster_tile_count_n: Int32 + :return: A tuple containing CTA tile indices along M and N dimensions + :rtype: tuple of (Int32, Int32) + """ + cluster_layout_mn = cute.make_layout( + (cluster_tile_count_m, cluster_tile_count_n), loc=loc, ip=ip + ) + (mi, ni) = cluster_layout_mn.get_hier_coord(cluster_tile_idx) + cta_tile_idx_m = ( + mi * self.params.cluster_shape_mn[0] + cta_tile_coord_in_cluster[0] + ) + cta_tile_idx_n = ( + ni * self.params.cluster_shape_mn[1] + cta_tile_coord_in_cluster[1] + ) + return (cta_tile_idx_m, cta_tile_idx_n) + + @dsl_user_op + @cute.jit + def _group_search( + self, + linear_idx: Int32, + problem_shape_mnkl: cute.Tensor, + init_group_idx: Int32, + init_tile_count_searched: Int32, + *, + loc=None, + ip=None, + ) -> GroupedGemmGroupSearchState: + """ + Search which group the linear index belongs to. + + :param linear_idx: The linear index to be decomposed + :type linear_idx: Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups + :type problem_shape_mnkl: cute.Tensor + :param init_group_idx: The group idx to start the search with + :type init_group_idx: Int32 + :param init_tile_count_searched: The number of tiles we have searched + :type init_tile_count_searched: Int32 + :return: The updated search state + :rtype: GroupedGemmGroupSearchState + """ + c_0 = Int32(0).ir_value() + last_lane_idx = cute.arch.WARP_SIZE - 1 + + tile_count_searched = init_tile_count_searched + start_group_idx = init_group_idx + not_found = linear_idx >= tile_count_searched + start_not_found = not_found + tile_count_prev_group = self.search_state.tile_count_prev_group + + while not_found and start_group_idx < self.group_count: + # get group to search for current lane + cur_group_idx = start_group_idx + self.lane_idx + # check if the group to be checked is out of range + inside_group_bound = cur_group_idx < self.group_count + + cur_ntile_mn = c_0 + if inside_group_bound: + # get problem size of current group + cur_problem_mnkl = self._get_problem_for_group( + problem_shape_mnkl, cur_group_idx, loc=loc, ip=ip + ) + cur_ntile_mn = self._get_cluster_tile_count_mn( + cur_problem_mnkl, loc=loc, ip=ip + ) + + # compute tile count from beginning to current group(included) + total_cluster_tile_count_ps_per_thread = self._prefix_sum( + cur_ntile_mn, loc=loc, ip=ip + ) + cluster_tile_count_end_per_thread = ( + total_cluster_tile_count_ps_per_thread + tile_count_searched + ) + + group_not_in_window = linear_idx >= cluster_tile_count_end_per_thread + hitted_group_idx_in_search_window = cute.arch.popc( + cute.arch.vote_ballot_sync(group_not_in_window, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + not_found = hitted_group_idx_in_search_window == cute.arch.WARP_SIZE + start_group_idx = hitted_group_idx_in_search_window + start_group_idx + + hit_the_1st_problem_in_search_window = ( + hitted_group_idx_in_search_window == c_0 + ) + tile_count_prev_group = tile_count_searched + if hit_the_1st_problem_in_search_window == False: + tile_count_prev_group = cute.arch.shuffle_sync( + cluster_tile_count_end_per_thread, + hitted_group_idx_in_search_window - 1, + ) + + # If no matched group, then get new_cluster_tile_count_end from last lane + # Otherwise, get new_cluster_tile_count_end from the hitted group + lane_idx_for_cluster_tile_count_end = hitted_group_idx_in_search_window + if not_found: + lane_idx_for_cluster_tile_count_end = last_lane_idx + tile_count_searched = cute.arch.shuffle_sync( + cluster_tile_count_end_per_thread, + lane_idx_for_cluster_tile_count_end, + ) + + # The tile is invalid if not_found doesn't change before and after the while loop. + end_not_found = not_found + is_valid = start_not_found != end_not_found + + return GroupedGemmGroupSearchState( + start_group_idx, + tile_count_prev_group, + tile_count_searched, + is_valid, + ) + + @dsl_user_op + @cute.jit + def _group_search_and_load_problem_shape( + self, + linear_idx: Int32, + problem_shape_mnkl: cute.Tensor, + start_group_idx: Int32, + tile_count_searched: Int32, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, cute.Tensor]: + """ + Perform group search and load problem shape for the matched group. + + :param linear_idx: The linear index to be decomposed + :type linear_idx: Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups + :type problem_shape_mnkl: cute.Tensor + :param start_group_idx: The group idx to start the search with + :type start_group_idx: Int32 + :param tile_count_searched: The number of tiles we have searched + :type tile_count_searched: Int32 + :return: A tuple containing the final group index and the problem shape tensor + :rtype: Tuple[Int32, cute.Tensor] + """ + self.search_state = self._group_search( + linear_idx, + problem_shape_mnkl, + start_group_idx, + tile_count_searched, + loc=loc, + ip=ip, + ) + # get final group search state + found = self.search_state.found + + final_group_idx = -1 + problem_mnkl = cute.make_rmem_tensor( + cute.make_layout(4), problem_shape_mnkl.element_type, loc=loc, ip=ip + ) + if found: + final_group_idx = self.search_state.start_group_idx + # let's revisit if it's better to broadcast problem_shape_mnk in group_search + problem_mnkl = self._get_problem_for_group( + problem_shape_mnkl, final_group_idx, loc=loc, ip=ip + ) + return found, final_group_idx, problem_mnkl + + @dsl_user_op + def delinearize_z( + self, + cta_tile_coord: tuple, + *, + loc=None, + ip=None, + ) -> GroupSearchResult: + """ + Delinearize the linear z index and return GroupSearchResult. + + This function should be used by warps that need to know the CTA tile index on M and N dimensions. + + :param cta_tile_coord: The raw CTA coordinate from tile scheduler + :type cta_tile_coord: tuple of Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for each group + :type problem_shape_mnkl: cute.Tensor + :return: The search result containing group index and tile coordinates + :rtype: GroupSearchResult + """ + # delinear the z coord + + linear_idx = self._current_work_linear_idx + + found, group_idx, problem_mnkl = self._group_search_and_load_problem_shape( + linear_idx, + self.problem_shape_mnkl, + self.search_state.start_group_idx, + self.search_state.tile_count_prev_group, + loc=loc, + ip=ip, + ) + + # The work_tile is valid if its linear index could be mapped to a group in the problem shapes + is_valid = found + + # linear index local to current group + cluster_tile_idx_in_current_group = ( + linear_idx - self.search_state.tile_count_prev_group + ) + + cluster_count_m, cluster_count_n, cluster_count_k = cute.ceil_div( + (problem_mnkl[0], problem_mnkl[1], problem_mnkl[2]), + ( + self.cluster_tile_shape_mnk[0], + self.cluster_tile_shape_mnk[1], + self.cluster_tile_shape_mnk[2], + ), + loc=loc, + ip=ip, + ) + + # decompose to get indices on M and N + cta_tile_idx_m, cta_tile_idx_n = self._compute_cta_tile_coord( + cluster_tile_idx_in_current_group, + cta_tile_coord, + cluster_count_m, + cluster_count_n, + loc=loc, + ip=ip, + ) + + group_search_result = GroupSearchResult( + group_idx, + cta_tile_idx_m, + cta_tile_idx_n, + problem_mnkl[0], + problem_mnkl[1], + problem_mnkl[2], + cluster_count_k, + ) + + return GroupedWorkTileInfo(cta_tile_coord, is_valid, group_search_result) + + @dsl_user_op + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + work_tile = self._get_current_work_for_linear_idx( + self._current_work_linear_idx, loc=loc, ip=ip + ) + grouped_work_tile = self.delinearize_z(work_tile.tile_idx, loc=loc, ip=ip) + return grouped_work_tile + +@deprecated("API is deprecated, use cutlass.utils.StaticPersistentGroupTileScheduler instead") +class GroupedGemmTileSchedulerHelper: + """ + A helper to translate the raw block index (x, y, z) from tile scheduler to real CTA tile index for grouped gemm. + + :param group_count: Number of groups in current grouped gemm problem + :type group_count: int + :param tile_sched_params: Parameter used to create the tile scheduler this helper works with + :type tile_sched_params: PersistentTileSchedulerParams + :param cluster_tile_shape_mnk: The shape of cluster tile as (m, n, k) + :type cluster_tile_shape_mnk: tuple[int, int, int] + :param search_state: The initial search state + :type search_state: GroupedGemmGroupSearchState + """ + + def __init__( + self, + group_count: int, + tile_sched_params: PersistentTileSchedulerParams, + cluster_tile_shape_mnk: tuple[int, int, int], + search_state: GroupedGemmGroupSearchState, + ) -> None: + self.tile_sched_params = tile_sched_params + self.group_count = group_count + self.lane_idx = cute.arch.lane_idx() + self.cluster_tile_shape_mnk = cluster_tile_shape_mnk + self.search_state = search_state + + def __extract_mlir_values__(self) -> List[ir.Value]: + values = extract_mlir_values(self.tile_sched_params) + values.extend(extract_mlir_values(self.search_state)) + return values + + def __new_from_mlir_values__( + self, values: List[ir.Value] + ) -> "GroupedGemmTileSchedulerHelper": + # Reconstruct tile_sched_params and determine how many values it consumed. + # NOTE: tile_sched_params may contain FastDivmod divisors (when swizzle_size == 1), + # which adds extra MLIR values. + params_values = extract_mlir_values(self.tile_sched_params) + n_params_values = len(params_values) + tile_sched_params = new_from_mlir_values( + self.tile_sched_params, values[:n_params_values] + ) + + # Reconstruct search_state from remaining values + search_state = new_from_mlir_values(self.search_state, values[n_params_values:]) + + return GroupedGemmTileSchedulerHelper( + self.group_count, + tile_sched_params, + self.cluster_tile_shape_mnk, + search_state, + ) + + def delinearize_z( + self, + cta_tile_coord: tuple, + problem_shape_mnkl: cute.Tensor, + ) -> GroupSearchResult: + """ + Delinearize the linear z index and return GroupSearchResult. + + This function should be used by warps that need to know the CTA tile index on M and N dimensions. + + :param cta_tile_coord: The raw CTA coordinate from tile scheduler + :type cta_tile_coord: tuple of Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for each group + :type problem_shape_mnkl: cute.Tensor + :return: The search result containing group index and tile coordinates + :rtype: GroupSearchResult + """ + # delinear the z coord + linear_idx = cta_tile_coord[2] + group_idx, problem_mnkl = self._group_search_and_load_problem_shape( + linear_idx, + problem_shape_mnkl, + self.search_state.start_group_idx, + self.search_state.tile_count_prev_group, + ) + # linear index local to current group + cluster_tile_idx_in_current_group = ( + linear_idx - self.search_state.tile_count_prev_group + ) + cluster_count_m, cluster_count_n, cluster_count_k = cute.ceil_div( + (problem_mnkl[0], problem_mnkl[1], problem_mnkl[2]), + ( + self.cluster_tile_shape_mnk[0], + self.cluster_tile_shape_mnk[1], + self.cluster_tile_shape_mnk[2], + ), + ) + # decompose to get indices on M and N + cta_tile_idx_m, cta_tile_idx_n = self._compute_cta_tile_coord( + cluster_tile_idx_in_current_group, + cta_tile_coord, + cluster_count_m, + cluster_count_n, + ) + return GroupSearchResult( + group_idx, + cta_tile_idx_m, + cta_tile_idx_n, + problem_mnkl[0], + problem_mnkl[1], + problem_mnkl[2], + cluster_count_k, + ) + + def search_cluster_tile_count_k( + self, + cta_tile_coord: tuple, + problem_shape_mnkl: cute.Tensor, + ) -> Tuple[Int32, Int32]: + """ + Search the matched group for given linear index and compute the number of tiles along K dimension for the matched group. + + This function should be used by warps that are only interested in the number of tiles along K dimension. + + :param cta_tile_coord: The raw CTA coordinate from tile scheduler + :type cta_tile_coord: tuple of Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups + :type problem_shape_mnkl: cute.Tensor + :return: A tuple containing cluster count along K dimension and the group index + :rtype: Tuple[Int32, Int32] + """ + group_idx, problem_mnk = self._group_search_and_load_problem_shape( + cta_tile_coord[2], + problem_shape_mnkl, + self.search_state.start_group_idx, + self.search_state.tile_count_prev_group, + ) + cluster_count_k = ( + problem_mnk[2] + self.cluster_tile_shape_mnk[2] - 1 + ) // self.cluster_tile_shape_mnk[2] + return cluster_count_k, group_idx + + @cute.jit + def _prefix_sum(self, value_per_thread: Int32) -> Int32: + """ + Perform prefix sum within a full warp. + + :param value_per_thread: The value for this thread to contribute to the prefix sum + :type value_per_thread: Int32 + :return: The prefix sum result for this thread + :rtype: Int32 + """ + clamp_value = 0 + idx = 1 + sum_per_thread = value_per_thread + while const_expr(idx < cute.arch.WARP_SIZE): + value = cute.arch.shuffle_sync_up( + sum_per_thread, idx, mask_and_clamp=clamp_value + ) + if self.lane_idx >= idx: + sum_per_thread += value + idx = idx << 1 + return sum_per_thread + + def _get_problem_for_group( + self, problem_shape_mnkl: cute.Tensor, group_idx: Int32 + ) -> cute.Tensor: + """ + Load gemm problem (m,n,k,l) for the specified group from global memory to register. + + :param problem_shape_mnkl: Tensor in global memory with layout (group_count, 4):(4, 1) + :type problem_shape_mnkl: cute.Tensor + :param group_idx: The index of the group to load + :type group_idx: Int32 + :return: The problem shape tensor for the specified group + :rtype: cute.Tensor + """ + cur_problem_mnkl = cute.make_rmem_tensor( + cute.make_layout(4), problem_shape_mnkl.element_type + ) + cute.autovec_copy(problem_shape_mnkl[(group_idx, None)], cur_problem_mnkl) + return cur_problem_mnkl + + def _get_cluster_tile_count_mn(self, problem_shape: cute.Tensor) -> Int32: + """ + Compute total cluster count. + + :param problem_shape: Tensor containing problem shape (m, n, k, l) + :type problem_shape: cute.Tensor + :return: The total cluster tile count for M and N dimensions + :rtype: Int32 + """ + cur_ntile_m = ( + problem_shape[0] + self.cluster_tile_shape_mnk[0] - 1 + ) // self.cluster_tile_shape_mnk[0] + cur_ntile_n = ( + problem_shape[1] + self.cluster_tile_shape_mnk[1] - 1 + ) // self.cluster_tile_shape_mnk[1] + cur_ntile_mn = cur_ntile_m * cur_ntile_n + return cur_ntile_mn + + def _compute_cta_tile_coord( + self, + cluster_tile_idx: Int32, + cta_tile_coord_in_cluster: tuple, + cluster_tile_count_m: Int32, + cluster_tile_count_n: Int32, + ) -> tuple: + """ + Compute CTA tile indices along M and N dimensions based on the linear index within a group. + + It uses the AlongM mode to decompose the linear index onto M and N dimensions. + + :param cluster_tile_idx: The linear index within a group + :type cluster_tile_idx: Int32 + :param cta_tile_coord_in_cluster: CTA indices along M and N dimensions within a cluster + :type cta_tile_coord_in_cluster: tuple of Int32 + :param cluster_tile_count_m: The number of clusters along M dimension of the matched group + :type cluster_tile_count_m: Int32 + :param cluster_tile_count_n: The number of clusters along N dimension of the matched group + :type cluster_tile_count_n: Int32 + :return: A tuple containing CTA tile indices along M and N dimensions + :rtype: tuple of (Int32, Int32) + """ + cluster_layout_mn = cute.make_layout( + (cluster_tile_count_m, cluster_tile_count_n) + ) + (mi, ni) = cluster_layout_mn.get_hier_coord(cluster_tile_idx) + cta_tile_idx_m = ( + mi * self.tile_sched_params.cluster_shape_mn[0] + + cta_tile_coord_in_cluster[0] + ) + cta_tile_idx_n = ( + ni * self.tile_sched_params.cluster_shape_mn[1] + + cta_tile_coord_in_cluster[1] + ) + return (cta_tile_idx_m, cta_tile_idx_n) + + @cute.jit + def _group_search( + self, + linear_idx: Int32, + problem_shape_mnkl: cute.Tensor, + init_group_idx: Int32, + init_tile_count_searched: Int32, + ) -> GroupedGemmGroupSearchState: + """ + Search which group the linear index belongs to. + + :param linear_idx: The linear index to be decomposed + :type linear_idx: Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups + :type problem_shape_mnkl: cute.Tensor + :param init_group_idx: The group idx to start the search with + :type init_group_idx: Int32 + :param init_tile_count_searched: The number of tiles we have searched + :type init_tile_count_searched: Int32 + :return: The updated search state + :rtype: GroupedGemmGroupSearchState + """ + c_0 = Int32(0).ir_value() + last_lane_idx = cute.arch.WARP_SIZE - 1 + + tile_count_searched = init_tile_count_searched + start_group_idx = init_group_idx + not_found = linear_idx >= tile_count_searched + tile_count_prev_group = self.search_state.tile_count_prev_group + while not_found: + # get group to search for current lane + cur_group_idx = start_group_idx + self.lane_idx + # check if the group to be checked is out of range + inside_group_bound = cur_group_idx < self.group_count + cur_ntile_mn = c_0 + if inside_group_bound: + # get problem size of current group + cur_problem_mnkl = self._get_problem_for_group( + problem_shape_mnkl, cur_group_idx + ) + cur_ntile_mn = self._get_cluster_tile_count_mn(cur_problem_mnkl) + # compute tile count from beginning to current group(included) + total_cluster_tile_count_ps_per_thread = self._prefix_sum(cur_ntile_mn) + cluster_tile_count_end_per_thread = ( + total_cluster_tile_count_ps_per_thread + tile_count_searched + ) + + group_not_in_window = linear_idx >= cluster_tile_count_end_per_thread + hitted_group_idx_in_search_window = cute.arch.popc( + cute.arch.vote_ballot_sync(group_not_in_window) + ) + not_found = hitted_group_idx_in_search_window == cute.arch.WARP_SIZE + start_group_idx = hitted_group_idx_in_search_window + start_group_idx + hit_the_1st_problem_in_search_window = ( + hitted_group_idx_in_search_window == c_0 + ) + tile_count_prev_group = tile_count_searched + if hit_the_1st_problem_in_search_window == False: + tile_count_prev_group = cute.arch.shuffle_sync( + cluster_tile_count_end_per_thread, + hitted_group_idx_in_search_window - 1, + ) + + # If no matched group, then get new_cluster_tile_count_end from last lane + # Otherwise, get new_cluster_tile_count_end from the hitted group + lane_idx_for_cluster_tile_count_end = hitted_group_idx_in_search_window + if not_found: + lane_idx_for_cluster_tile_count_end = last_lane_idx + tile_count_searched = cute.arch.shuffle_sync( + cluster_tile_count_end_per_thread, + lane_idx_for_cluster_tile_count_end, + ) + + return GroupedGemmGroupSearchState( + start_group_idx, + tile_count_prev_group, + tile_count_searched, + 1 # found will always be 1 for old api + ) + + def _group_search_and_load_problem_shape( + self, + linear_idx: Int32, + problem_shape_mnkl: cute.Tensor, + start_group_idx: Int32, + tile_count_searched: Int32, + ) -> Tuple[Int32, cute.Tensor]: + """ + Perform group search and load problem shape for the matched group. + + :param linear_idx: The linear index to be decomposed + :type linear_idx: Int32 + :param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups + :type problem_shape_mnkl: cute.Tensor + :param start_group_idx: The group idx to start the search with + :type start_group_idx: Int32 + :param tile_count_searched: The number of tiles we have searched + :type tile_count_searched: Int32 + :return: A tuple containing the final group index and the problem shape tensor + :rtype: Tuple[Int32, cute.Tensor] + """ + self.search_state = self._group_search( + linear_idx, + problem_shape_mnkl, + start_group_idx, + tile_count_searched, + ) + # get final group search state + final_group_idx = self.search_state.start_group_idx + # let's revisit if it's better to broadcast problem_shape_mnk in group_search + problem_mnkl = self._get_problem_for_group(problem_shape_mnkl, final_group_idx) + return final_group_idx, problem_mnkl + diff --git a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py index 0468f5d9..41ab59bd 100644 --- a/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py +++ b/python/CuTeDSL/cutlass/utils/mixed_input_helpers.py @@ -1,30 +1,13 @@ -# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. 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. +# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. from __future__ import annotations @@ -34,12 +17,16 @@ from typing import Optional, Union import cutlass import cutlass.cute as cute +from cutlass.cutlass_dsl import ( + extract_mlir_values, + new_from_mlir_values, +) +from cutlass.utils.layout import LayoutEnum import cutlass.utils.blackwell_helpers as sm100_utils from cutlass.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.runtime import from_dlpack """ -This file contain common utilities for mixed-input GEMM. +This file contains common utilities for mixed-input GEMM. """ @@ -61,7 +48,7 @@ def scale_tma_partition( ) -> tuple[cute.Tensor, cute.Tensor]: """ Perform TMA partition for scale tensor. - This method partitions the gobal memory and shared memory buffer for scale tensor for TMA load. + This method partitions the global memory and shared memory buffer for the scale tensor for TMA load. :param tCsS: Input scale shared memory tensor :type tCsS: cute.Tensor :param tCgS: Input scale global memory tensor @@ -120,7 +107,9 @@ def transform_partition( sA_input: cute.Tensor, A_transform: cute.Tensor, transform_local_tidx: cutlass.Int32, -) -> tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]: +) -> tuple[ + Optional[cute.TiledCopy], Optional[cute.TiledCopy], cute.Tensor, cute.Tensor +]: """ Partition tensors for transform input and output. This method sets up the copy atoms and partitions the shared/tensor memory @@ -146,7 +135,7 @@ def transform_partition( * tAsA_input: Partitioned input tensor A * tA_transform: Partitioned transformed tensor A - :rtype: tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor] + :rtype: tuple[Optional[cute.TiledCopy], Optional[cute.TiledCopy], cute.Tensor, cute.Tensor] """ if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM): if cutlass.const_expr( @@ -244,6 +233,185 @@ def scale_partition( return smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS +def epilog_gmem_copy_and_partition( + c_dtype: type[cutlass.Numeric], + tidx: cutlass.Int32, + tma_atom_c: Optional[cute.CopyAtom], + tiled_copy_t2r: Optional[cute.TiledCopy], + gC_mnl_tma: Optional[cute.Tensor], + gC_mnl_simt: Optional[cute.Tensor], + epi_tile: cute.Tile, + sC: cute.Tensor, +) -> tuple[ + Optional[cute.CopyAtom], + Optional[cute.Tensor], + Optional[cute.Tensor], + Optional[cute.CopyAtom], + Optional[cute.Tensor], +]: + """ + Partitions source and destination tensors for a global memory store. + This method generates a tiled copy for storing results to global memory + and partitions the source (register or shared memory) and destination + (global memory) tensors accordingly. If tma_atom_c is not None, then + partition for TMA store will be performed. If tiled_copy_t2r is not None, then + partition for SIMT store will be performed. + :param c_dtype: The data type of the tensor C. + :type c_dtype: type[cutlass.Numeric] + :param tidx: The thread index in epilogue warp groups. + :type tidx: cutlass.Int32 + :param tma_atom_c: The TMA copy atom. + :type tma_atom_c: Optional[cute.CopyAtom] + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy. + :type tiled_copy_t2r: Optional[cute.TiledCopy] + :param gC_mnl_tma: The global tensor C for TMA. + :type gC_mnl_tma: Optional[cute.Tensor] + :param gC_mnl_simt: The global tensor C for SIMT Copy. + :type gC_mnl_simt: Optional[cute.Tensor] + :param epi_tile: The epilogue tiler. + :type epi_tile: cute.Tile + :param sC: The shared memory tensor C. + :return: A tuple containing the appropriate copy atom and partitioned + source and destination tensors for the store operation. + :rtype: tuple[Optional[cute.CopyAtom], Optional[cute.Tensor], Optional[cute.Tensor], Optional[cute.CopyAtom], Optional[cute.Tensor]] + """ + bSG_sC = None + bSG_gC = None + simt_atom = None + tTR_gC = None + if tma_atom_c is not None: + gC_epi_tma = cute.flat_divide( + gC_mnl_tma[((None, None), 0, 0, None, None, None)], epi_tile + ) + # TMA store + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi_tma, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + if tiled_copy_t2r is not None: + # SIMT Store + gC_epi_simt = cute.flat_divide( + gC_mnl_simt[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_gC = thr_copy_t2r.partition_D(gC_epi_simt) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype) + return tma_atom_c, bSG_sC, bSG_gC, simt_atom, tTR_gC + + +def epilog_smem_copy_and_partition( + c_layout: LayoutEnum, + c_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, +) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Partitions source and destination tensors for a shared memory store. + This method generates a tiled copy for storing results to shared memory + and partitions the source (register) and destination (shared memory) + tensors accordingly. + :param c_layout: The layout of the tensor C. + :type c_layout: LayoutEnum + :param c_dtype: The data type of the tensor C. + :type c_dtype: type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator tensor. + :type acc_dtype: type[cutlass.Numeric] + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy. + :param tTR_rC: The partitioned accumulator tensor. + :param tidx: The thread index in epilogue warp groups. + :param sC: The shared memory tensor to be copied and partitioned. + :return: A tuple containing the tiled copy for the store operation and + the partitioned source and destination tensors. + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + c_layout, c_dtype, acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + +def epilog_tmem_copy_and_partition( + cta_tile_shape_mnk: tuple[int, int, int], + c_layout: LayoutEnum, + c_dtype: type[cutlass.Numeric], + acc_dtype: type[cutlass.Numeric], + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], +) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Partitions source and destination tensors for a tensor memory load. + This method generates a tiled copy for loading accumulators from tensor + memory and partitions the source (tensor memory) and destination + (register) tensors accordingly. + :param cta_tile_shape_mnk: The shape of the CTA tile (M, N, K). + :type cta_tile_shape_mnk: tuple[int, int, int] + :param c_layout: The layout of the tensor C. + :type c_layout: LayoutEnum + :param c_dtype: The data type of the tensor C. + :type c_dtype: type[cutlass.Numeric] + :param acc_dtype: The data type of the accumulator tensor. + :type acc_dtype: type[cutlass.Numeric] + :param tidx: The thread index in epilogue warp groups. + :param tAcc: The accumulator tensor to be copied and partitioned. + :param gC_mnl: The global tensor C. + :param epi_tile: The epilogue tiler. + :param use_2cta_instrs: Whether use_2cta_instrs is enabled. + :return: A tuple containing the tiled copy for the load operation and + the partitioned source and destination tensors. + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + cta_tile_shape_mnk, + c_layout, + c_dtype, + acc_dtype, + epi_tile, + use_2cta_instrs, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)] + ) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gC_mnl_epi = cute.flat_divide( + gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile + ) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def get_gmem_layout_scale( scale_shape_mkl: tuple[int, int, int], scale_granularity_m: int, @@ -345,10 +513,9 @@ def get_smem_layout_scale( cute.size(mma_tiler[2]) % cute.size(smem_layout_scale_per_stage.outer[1]) == 0 ), "smem_layout_scale_per_stage must evenly divide tile k shape." # Shared memory buffer for scale must be at least 128B to satisfy TMA requirement - - assert cute.size_in_bytes(a_scale_dtype, smem_layout_scale_per_stage) >= 128, ( - "smem size for scale must be at least 128B" - ) + assert ( + cute.size_in_bytes(a_scale_dtype, smem_layout_scale_per_stage) >= 128 + ), "smem size for scale must be at least 128B" # Scale layout in smem with multiple stages smem_layout_scale = cute.append( smem_layout_scale_per_stage, @@ -496,6 +663,98 @@ def is_valid_scale_granularity( return True +def is_shuffle_a( + a_major: str, + k: int, + a_dtype: type[cutlass.Numeric], + mma_dtype: type[cutlass.Numeric], + scale_granularity_k: int, +) -> bool: + # Enable shuffle on the k mode of A tensor when + # 1) tensor a is k-major and, + # 2) k is exactly divisible by 8 and, + # 3) a_dtype is Int4 and mma_dtype is BFloat16 and, + # 4) scale granularity k is greater than 8 + shuffle_a = ( + a_major == "k" + and a_dtype == cutlass.Int4 + and k % 8 == 0 + and mma_dtype == cutlass.BFloat16 + and scale_granularity_k >= 8 + ) + return shuffle_a + + +def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + a_dtype: type[cutlass.Numeric], + b_dtype: type[cutlass.Numeric], + c_dtype: type[cutlass.Numeric], + scale_dtype: type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + mma_tiler_mnk: tuple[int, int, int], + use_2cta_instrs: bool, + cluster_shape_mn: tuple[int, int], + scale_granularity_m: int, + scale_granularity_k: int, +) -> bool: + """ + Check if the tensor alignments are valid for the given problem size and data types. + """ + + def check_contiguous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if not ( + check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k)) + and check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k)) + and check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n)) + and ( + scale_granularity_k == 0 + or check_contiguous_16B_alignment( + b_dtype, True, (m, k // scale_granularity_k) + ) + ) + ): + return False + # Check if scale tensor matches the TMA load 128B alignment requirement + cta_tile_shape_mnk = ( + mma_tiler_mnk[0] // (2 if use_2cta_instrs else 1), + mma_tiler_mnk[1], + mma_tiler_mnk[2], + ) + if ( + scale_granularity_m > 0 + and (cta_tile_shape_mnk[0] // cluster_shape_mn[1] // scale_granularity_m) + * (scale_dtype.width // 8) + < 128 + ): + return False + return True + + +def is_valid_mma_tiler_and_cluster_shape( + mma_tiler: tuple[int, int, int], + cluster_shape_mn: tuple[int, int], + use_2cta_instrs: bool, +) -> bool: + """ + Check if the MMA tiler and cluster shape are valid for the given problem size. + """ + if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: + return False + if (mma_tiler[0] // (2 if use_2cta_instrs else 1)) not in [64, 128]: + return False + return True + + def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int: """ Calculate the largest power of 2 divisibility factor for memory alignment. @@ -505,3 +764,244 @@ def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int: if contiguous_dim_size % (2**i) == 0: return min(2**i, upper_bound) return 1 + + +class ContiguousGGSearchState: + """ + The state of group search for grouped GEMM with contiguous offsets. + + The state records the progress of the group search algorithm on one mode. It will be + initialized once and updated in every round of group index search. + + :param last_tile_count: Number of cluster tiles before the current group + :type last_tile_count: cutlass.Int32 + :param cur_boundary: The boundary of the current group, which is the size along the seach + mode before the next group + :type cur_boundary: cutlass.Int32 + :param cur_tile_count: Number of cluster tiles searched so far + :type cur_tile_count: cutlass.Int32 + :param cur_group_idx: The index of the current group + :type cur_group_idx: cutlass.Int32 + :param cur_offset: The starting offset of the current group along the search mode + :type cur_offset: cutlass.Int32 + :param cur_start: The starting offset of the current cluster tile size along the search mode + when group search is done + :type cur_start: cutlass.Int32 + """ + + def __init__( + self, + last_tile_count: cutlass.Int32, + cur_boundary: cutlass.Int32, + cur_tile_count: cutlass.Int32, + cur_group_idx: cutlass.Int32, + cur_offset: cutlass.Int32, + cur_start: cutlass.Int32, + ): + self.last_tile_count = last_tile_count + self.cur_boundary = cur_boundary + self.cur_tile_count = cur_tile_count + self.cur_group_idx = cur_group_idx + self.cur_offset = cur_offset + self.cur_start = cur_start + + def __extract_mlir_values__(self): + values = extract_mlir_values(self.last_tile_count) + values.extend(extract_mlir_values(self.cur_boundary)) + values.extend(extract_mlir_values(self.cur_tile_count)) + values.extend(extract_mlir_values(self.cur_group_idx)) + values.extend(extract_mlir_values(self.cur_offset)) + values.extend(extract_mlir_values(self.cur_start)) + return values + + def __new_from_mlir_values__(self, values) -> "ContiguousGGSearchState": + last_tile_count = new_from_mlir_values(self.last_tile_count, [values[0]]) + cur_boundary = new_from_mlir_values(self.cur_boundary, [values[1]]) + cur_tile_count = new_from_mlir_values(self.cur_tile_count, [values[2]]) + cur_group_idx = new_from_mlir_values(self.cur_group_idx, [values[3]]) + cur_offset = new_from_mlir_values(self.cur_offset, [values[4]]) + cur_start = new_from_mlir_values(self.cur_start, [values[5]]) + return ContiguousGGSearchState( + last_tile_count, + cur_boundary, + cur_tile_count, + cur_group_idx, + cur_offset, + cur_start, + ) + + +def create_initial_contiguous_group_search_state() -> ContiguousGGSearchState: + """ + Create an initial search state for grouped gemm with contiguous offsets. + """ + return ContiguousGGSearchState( + last_tile_count=cutlass.Int32(0), + cur_boundary=cutlass.Int32(0), + cur_tile_count=cutlass.Int32(0), + cur_group_idx=cutlass.Int32(0), + cur_offset=cutlass.Int32(0), + cur_start=cutlass.Int32(0), + ) + + +class ContiguousGroupWorkTileInfo: + """ + Tile info for grouped GEMM with contiguous offsets. + It's constructed from the search state and contains information needed for different warps. + + :param group_count: The total number of groups + :type group_count: int + :param cta_coord_m: The coordinate of the current CTA tile along the M mode + :type cta_coord_m: cutlass.Int32 + :param coord_n: The starting offset on N mode for the current CTA tile + :type coord_n: cutlass.Int32 + :param group_idx: The index of the current group + :type group_idx: cutlass.Int32 + :param distance_to_boundary: The distance to the boundary of the current group + :type distance_to_boundary: cutlass.Int32 + """ + + def __init__( + self, + group_count: int, + cta_coord_m: cutlass.Int32, + coord_n: cutlass.Int32, + group_idx: cutlass.Int32, + distance_to_boundary: cutlass.Int32, + ): + self.cta_coord_m = cta_coord_m + self.coord_n = coord_n + self.group_idx = group_idx + self.distance_to_boundary = distance_to_boundary + self.group_count = group_count + + def __extract_mlir_values__(self): + values = extract_mlir_values(self.cta_coord_m) + values.extend(extract_mlir_values(self.coord_n)) + values.extend(extract_mlir_values(self.group_idx)) + values.extend(extract_mlir_values(self.distance_to_boundary)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 4 + new_cta_coord_m = new_from_mlir_values(self.cta_coord_m, [values[0]]) + new_coord_n = new_from_mlir_values(self.coord_n, [values[1]]) + new_group_idx = new_from_mlir_values(self.group_idx, [values[2]]) + new_distance_to_boundary = new_from_mlir_values( + self.distance_to_boundary, [values[3]] + ) + return ContiguousGroupWorkTileInfo( + self.group_count, + new_cta_coord_m, + new_coord_n, + new_group_idx, + new_distance_to_boundary, + ) + + @property + def is_valid_tile(self): + return self.group_idx < self.group_count + + +@cute.jit +def contiguous_group_search( + cluster_tile_shape_mnk: tuple[int, int, int], + group_count: cutlass.Int32, + linear_idx: cutlass.Int32, + search_state: ContiguousGGSearchState, + cumsum: cute.Tensor, + search_mode: int, +) -> ContiguousGGSearchState: + """ + Group search for contiguous grouped gemm. + """ + not_found = linear_idx >= search_state.cur_tile_count + next_boundary = cutlass.Int32(0) + cur_group_idx = search_state.cur_group_idx + cur_offset = search_state.cur_offset + last_tile_count = search_state.last_tile_count + cur_boundary = search_state.cur_boundary + cur_tile_count = search_state.cur_tile_count + if not_found: + cur_group_idx = cur_group_idx + 1 + while not_found and cur_group_idx <= group_count: + next_boundary = cumsum[cur_group_idx] + num_m_blocks = cute.ceil_div( + (next_boundary - cur_boundary), + cluster_tile_shape_mnk[search_mode], + ) + next_tile_count = num_m_blocks + cur_tile_count + not_found = linear_idx >= next_tile_count + + last_tile_count = cur_tile_count + cur_offset = cur_boundary + cur_boundary = next_boundary + cur_tile_count = next_tile_count + if not_found: + cur_group_idx = cur_group_idx + 1 + cur_start = cur_offset + cluster_tile_shape_mnk[search_mode] * ( + linear_idx - last_tile_count + ) + return ContiguousGGSearchState( + last_tile_count, + cur_boundary, + cur_tile_count, + cur_group_idx, + cur_offset, + cur_start, + ) + + +def make_contiguous_group_work_tile_info(group_count: int, sTile_info: cute.Tensor): + """ + Generate ContiguousGroupWorkTileInfo from tile_info tensor generated by contiguous_group_search + """ + tile_info = cute.make_rmem_tensor(sTile_info.shape, sTile_info.element_type) + cute.autovec_copy(sTile_info, tile_info) + return ContiguousGroupWorkTileInfo( + group_count, tile_info[0], tile_info[1], tile_info[2], tile_info[3] + ) + + +def cvt_tensor_a( + src: cute.Tensor, dtype: type[cutlass.Numeric], shuffle: bool +) -> cute.TensorSSA: + """ + Convert tensor src to the given data type. If shuffle is True, use shuffle intrinsic + for int4-to-bf16 conversion. + """ + from cutlass import CUDA_VERSION + + # shuffle is supported since CUDA 13.1 + shuffle_supported = True + if CUDA_VERSION.major < 13 or (CUDA_VERSION == 13 and CUDA_VERSION.minor < 1): + shuffle_supported = False + shuffle = shuffle and shuffle_supported + rst = src.load() + if cutlass.const_expr(shuffle): + # conversion with shuffle + rst = cute.TensorSSA( + cute.arch.cvt_i4_bf16_intrinsic( + rst, + cute.size(rst.shape), + with_shuffle=True, + ), + rst.shape, + dtype, + ) + else: + rst = rst.to(dtype) + return rst + + +def store_transformed_a( + src_a: cute.Tensor, dst_a: cute.Tensor, copy_atom_a: Optional[cute.CopyAtom] +) -> None: + """ + Store transformed A tensor to the given destination tensor. If copy_atom_a is not None, use autovec_copy. + """ + if cutlass.const_expr(copy_atom_a is not None): + cute.copy(copy_atom_a, src_a, dst_a) + else: + cute.autovec_copy(src_a, dst_a) diff --git a/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py b/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py index 164224ae..925ff522 100644 --- a/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py +++ b/python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py @@ -87,6 +87,7 @@ class PersistentTileSchedulerParams: :type problem_layout_ncluster_mnl: cute.Layout """ + @dsl_user_op def __init__( self, problem_shape_ntile_mnl: cute.Shape, @@ -124,7 +125,7 @@ class PersistentTileSchedulerParams: self._cluster_shape_mnk = cluster_shape_mnk self.cluster_shape_mn = cluster_shape_mnk[:2] self.swizzle_size = swizzle_size - self._raster_along_m = raster_along_m + self.raster_along_m = raster_along_m self._loc = loc # By default, we follow m major (col-major) raster order, so make a col-major layout @@ -188,20 +189,27 @@ class PersistentTileSchedulerParams: problem_layout_size, loc=loc, ip=ip ) - # cluster_shape_m_fdd: Used to decode work_unit_id to cluster coordinates - self.cluster_shape_m_fdd = cute.fast_divmod_create_divisor( - cluster_count_m, loc=loc, ip=ip + if raster_along_m: + cluster_count_major = cluster_count_m + cluster_count_minor = cluster_count_n + else: + cluster_count_major = cluster_count_n + cluster_count_minor = cluster_count_m + + # cluster_shape_major_fdd: Used to decode work_unit_id to cluster coordinates + self.cluster_shape_major_fdd = cute.fast_divmod_create_divisor( + cluster_count_major, loc=loc, ip=ip ) - # cluster_shape_n_fdd: Used for the second level decomposition - self.cluster_shape_n_fdd = cute.fast_divmod_create_divisor( - cluster_count_n, loc=loc, ip=ip + # cluster_shape_minor_fdd: Used for the second level decomposition + self.cluster_shape_minor_fdd = cute.fast_divmod_create_divisor( + cluster_count_minor, loc=loc, ip=ip ) else: # FastDivmod not applicable with swizzling, set to None self.batch_fdd = None - self.cluster_shape_m_fdd = None - self.cluster_shape_n_fdd = None + self.cluster_shape_major_fdd = None + self.cluster_shape_minor_fdd = None def __extract_mlir_values__(self): values, self._values_pos = [], [] @@ -209,7 +217,7 @@ class PersistentTileSchedulerParams: self.problem_shape_ntile_mnl, self._cluster_shape_mnk, self.swizzle_size, - self._raster_along_m, + self.raster_along_m, ]: obj_values = extract_mlir_values(obj) values += obj_values @@ -223,8 +231,8 @@ class PersistentTileSchedulerParams: for i, (fdd_name, fdd_obj) in enumerate( [ ("batch_fdd", self.batch_fdd), - ("cluster_shape_m_fdd", self.cluster_shape_m_fdd), - ("cluster_shape_n_fdd", self.cluster_shape_n_fdd), + ("cluster_shape_major_fdd", self.cluster_shape_major_fdd), + ("cluster_shape_minor_fdd", self.cluster_shape_minor_fdd), ] ): if fdd_obj is not None: @@ -251,7 +259,7 @@ class PersistentTileSchedulerParams: self.problem_shape_ntile_mnl, self._cluster_shape_mnk, self.swizzle_size, - self._raster_along_m, + self.raster_along_m, ], self._values_pos[:-1], # Exclude FastDivmod count ): @@ -263,7 +271,7 @@ class PersistentTileSchedulerParams: new_params = PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc) # Restore FastDivmod divisors from remaining values - fdd_names = ["batch_fdd", "cluster_shape_m_fdd", "cluster_shape_n_fdd"] + fdd_names = ["batch_fdd", "cluster_shape_major_fdd", "cluster_shape_minor_fdd"] if hasattr(self, "_fastdivmod_indices") and len(self._fastdivmod_indices) > 0: # Override the FastDivmod divisors created by __init__ with reconstructed ones @@ -538,15 +546,24 @@ class StaticPersistentTileScheduler: # Step 2: Decode work_unit_id using FastDivmod objects # The layout structure is: problem_layout_ncluster_mnl has shape (cluster_count_m, cluster_count_n, batch_count) - # work_unit_id needs to be decomposed into (batch_l, cluster_n, cluster_m) in little-endian order + # work_unit_id needs to be decomposed into (batch_l, cluster_minor, cluster_major) in little-endian order - # First, get cluster_m using cluster_shape_m_fdd - cluster_n_batch, cluster_m = divmod( - work_unit_id, self.params.cluster_shape_m_fdd + # First, get cluster_major using cluster_shape_major_fdd + cluster_minor_batch, cluster_major = divmod( + work_unit_id, self.params.cluster_shape_major_fdd ) - # Then decode cluster_n_batch to get cluster_n and batch_l using FastDivmod - batch_l, cluster_n = divmod(cluster_n_batch, self.params.cluster_shape_n_fdd) + # Then decode cluster_minor_batch to get cluster_minor and batch_l using FastDivmod + batch_l, cluster_minor = divmod( + cluster_minor_batch, self.params.cluster_shape_minor_fdd + ) + + if self.params.raster_along_m: + cluster_m = cluster_major + cluster_n = cluster_minor + else: + cluster_m = cluster_minor + cluster_n = cluster_major return (cluster_m, cluster_n, batch_l) diff --git a/python/CuTeDSL/cutlass/utils/tensormap_manager.py b/python/CuTeDSL/cutlass/utils/tensormap_manager.py index 716eb804..f5c7160d 100644 --- a/python/CuTeDSL/cutlass/utils/tensormap_manager.py +++ b/python/CuTeDSL/cutlass/utils/tensormap_manager.py @@ -79,34 +79,50 @@ class TensorMapManager: # init tensormap pointed by dst_ptr with the one inside copy_atom. # dst_ptr should be pointing to a global memory location or a smem location # warp_id specifies which warp to perform the initialization + @dsl_user_op @cute.jit def init_tensormap_from_atom( - self, copy_atom: cute.CopyAtom, dst_ptr: cute.Pointer, warp_id: int + self, + copy_atom: cute.CopyAtom, + dst_ptr: cute.Pointer, + warp_id: int, + *, + loc=None, + ip=None, ) -> None: - warp_idx = cute.arch.warp_idx() - warp_idx = cute.arch.make_warp_uniform(warp_idx) + warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) + warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip) if warp_idx == warp_id: - with cute.arch.elect_one(): - cute.nvgpu.cpasync.copy_tensormap(copy_atom, dst_ptr) - cute.arch.sync_warp() + with cute.arch.elect_one(loc=loc, ip=ip): + cute.nvgpu.cpasync.copy_tensormap(copy_atom, dst_ptr, loc=loc, ip=ip) + cute.arch.sync_warp(loc=loc, ip=ip) return # Perform a fence operation to ensure previous `init_tensormap_from_atom` calls have been completed + @dsl_user_op def fence_tensormap_initialization( self, + *, + loc=None, + ip=None, ) -> None: if self.tensormap_update_mode == TensorMapUpdateMode.GMEM: - cute.arch.fence_acq_rel_cta() + cute.arch.fence_acq_rel_cta(loc=loc, ip=ip) return # Perform a fence operation to ensure previous `update_tensormap` calls have been completed + @dsl_user_op def fence_tensormap_update( self, tensormap_ptr: cute.Pointer, + *, + loc=None, + ip=None, ) -> None: - cute.nvgpu.cpasync.fence_tma_desc_acquire(tensormap_ptr) + cute.nvgpu.cpasync.fence_tma_desc_acquire(tensormap_ptr, loc=loc, ip=ip) return + @dsl_user_op @cute.jit def update_tensormap( self, @@ -115,8 +131,13 @@ class TensorMapManager: tensormap_gmem_ptr: Tuple[cute.Pointer, ...], warp_id: int, tensormap_smem_ptr: Tuple[cute.Pointer, ...], + *, + loc=None, + ip=None, ) -> None: - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + warp_idx = cute.arch.make_warp_uniform( + cute.arch.warp_idx(loc=loc, ip=ip), loc=loc, ip=ip + ) # updates before touching tensormap in global memory if warp_idx == warp_id: if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM): @@ -124,23 +145,25 @@ class TensorMapManager: tma_copy_atom, tensor_gmem, tensormap_smem_ptr ): cute.nvgpu.cpasync.update_tma_descriptor( - copy_atom, tensor, smem_ptr + copy_atom, tensor, smem_ptr, loc=loc, ip=ip ) # wait until it's safe to update tensormap in global memory - with cute.arch.elect_one(): - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.sync_warp() + with cute.arch.elect_one(loc=loc, ip=ip): + cute.arch.cp_async_bulk_commit_group(loc=loc, ip=ip) + cute.arch.cp_async_bulk_wait_group(0, read=True, loc=loc, ip=ip) + cute.arch.sync_warp(loc=loc, ip=ip) # updates to tensormap in global memory if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM): for gmem_ptr, smem_ptr in zip(tensormap_gmem_ptr, tensormap_smem_ptr): - cute.nvgpu.cpasync.cp_fence_tma_desc_release(gmem_ptr, smem_ptr) + cute.nvgpu.cpasync.cp_fence_tma_desc_release( + gmem_ptr, smem_ptr, loc=loc, ip=ip + ) else: for copy_atom, tensor, gmem_ptr in zip( tma_copy_atom, tensor_gmem, tensormap_gmem_ptr ): cute.nvgpu.cpasync.update_tma_descriptor( - copy_atom, tensor, gmem_ptr + copy_atom, tensor, gmem_ptr, loc=loc, ip=ip ) - cute.arch.sync_warp() - cute.nvgpu.cpasync.fence_tma_desc_release() + cute.arch.sync_warp(loc=loc, ip=ip) + cute.nvgpu.cpasync.fence_tma_desc_release(loc=loc, ip=ip) diff --git a/python/CuTeDSL/cutlass/utils/tmem_allocator.py b/python/CuTeDSL/cutlass/utils/tmem_allocator.py index 19818fdb..925f7005 100644 --- a/python/CuTeDSL/cutlass/utils/tmem_allocator.py +++ b/python/CuTeDSL/cutlass/utils/tmem_allocator.py @@ -9,7 +9,8 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. -from typing import Optional, Type +from typing import Optional, Type, Union, List +from math import ceil, log2 import inspect from cutlass import const_expr @@ -23,10 +24,12 @@ from cutlass.cutlass_dsl import ( import cutlass.pipeline as pipeline import cutlass.cute as cute from cutlass._mlir import ir +from cutlass.cute.nvgpu.tcgen05 import find_tmem_tensor_col_offset +from cutlass.cute.arch import get_max_tmem_alloc_cols, get_min_tmem_alloc_cols class TmemAllocator: - """A class for managing tensor memory allocation on Blackwell GPU. + """A class for managing tensor memory allocation. This class manages allocation/deallocation of tensor memory, including the mbarrier synchronization for two cta use case. @@ -43,6 +46,8 @@ class TmemAllocator: :type _num_allocated_columns: int :ivar _two_cta_tmem_dealloc_mbar_ptr: The mbarrier pointer required when deallocating tensor memory for two cta. :type _two_cta_tmem_dealloc_mbar_ptr: cute.Pointer + :ivar _arch: The architecture of the GPU. + :type _arch: str """ @dsl_user_op @@ -75,6 +80,7 @@ class TmemAllocator: num_allocated_columns: int = 0, two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] = None, *, + arch: str = "sm_100", loc=None, ip=None, ): @@ -116,6 +122,8 @@ class TmemAllocator: self._num_allocated_columns = num_allocated_columns self._two_cta_tmem_dealloc_mbar_ptr = two_cta_tmem_dealloc_mbar_ptr self._barrier_for_retrieve = barrier_for_retrieve + self._arch = arch + self._max_tmem_columns = get_max_tmem_alloc_cols(arch) # Init tmem dealloc mbarrier if two cta if const_expr(self._is_two_cta): @@ -150,6 +158,7 @@ class TmemAllocator: self._is_two_cta, self._num_allocated_columns, new_two_cta_tmem_dealloc_mbar_ptr, + arch=self._arch, ) @cute.jit @@ -157,13 +166,13 @@ class TmemAllocator: """Check if the number of columns is valid. This method checks if the number of columns is valid. - It checks if the number of columns is larger than 0, smaller than 512, a multiple of 32, and a power of two. + It checks if the number of columns is larger than 0, smaller than max capacity, a multiple of 32, and a power of two. """ # larger than 0 if const_expr(num_columns < 0): return False - # smaller than 512 - if const_expr(num_columns > 512): + # smaller than max capacity + if const_expr(num_columns > self._max_tmem_columns): return False # multiple of 32 if const_expr(num_columns % 32 != 0): @@ -183,10 +192,10 @@ class TmemAllocator: """ assert self.check_valid_num_columns(num_columns), ( - "num_columns must be multiple of 32 and power of two, and between 0 and 512" + f"num_columns must be multiple of 32 and power of two, and between 0 and {self._max_tmem_columns}" ) - assert self._num_allocated_columns + num_columns <= 512, ( - "total allocated columns must be less than or equal to 512" + assert self._num_allocated_columns + num_columns <= self._max_tmem_columns, ( + f"total allocated columns must be less than or equal to {self._max_tmem_columns}" ) warp_idx = cute.arch.warp_idx(loc=loc, ip=ip) @@ -197,6 +206,7 @@ class TmemAllocator: num_columns, self._alloc_result_dst_smem_ptr, is_two_cta=self._is_two_cta, + arch=self._arch, loc=loc, ip=ip, ) @@ -290,6 +300,7 @@ class TmemAllocator: tmem_ptr, num_deallocate_columns, is_two_cta=self._is_two_cta, + arch=self._arch, loc=loc, ip=ip, ) @@ -335,3 +346,56 @@ TmemAllocator.__init__.__signature__ = inspect.Signature( ), ] ) + + +def get_num_tmem_alloc_cols( + tmem_tensors: Union[cute.Tensor, List[cute.Tensor]], + rounding=True, + *, + arch: str = "sm_100", + loc=None, + ip=None, +) -> int: + """Get the total number of TMEM allocation columns for the given TMEM tensors. + + :param tmem_tensors: The TMEM tensors to get the number of allocation columns for. + :type tmem_tensors: Union[cute.Tensor, List[cute.Tensor]] + :param rounding: Whether to round up the number of allocation columns to the nearest power of 2. + :type rounding: bool + :param arch: The architecture of the GPU. + :type arch: str + :return: The total number of TMEM allocation columns. + :rtype: int + + :raises ValueError: If the number of TMEM allocation columns exceeds the maximum capacity or is less than 32. + """ + # Turn tmem_tensors into a list + if isinstance(tmem_tensors, cute.Tensor): + tmem_tensors = [tmem_tensors] + + tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch) + tmem_min_alloc_cols = get_min_tmem_alloc_cols(arch) + + # For each tensor in tmem_tensors, find the tmem_tensor_col_offset + num_tmem_alloc_cols_per_tensor = [ + find_tmem_tensor_col_offset(t) for t in tmem_tensors + ] + + # Sum up the num_tmem_alloc_cols_per_tensor + num_tmem_alloc_cols = sum(num_tmem_alloc_cols_per_tensor) + + # Round up num_tmem_cols_total to the nearest power of 2 and make sure it is at least 32 + if rounding: + num_tmem_alloc_cols = max( + 1 << ceil(log2(num_tmem_alloc_cols)), tmem_min_alloc_cols + ) + + # Validate the number of TMEM allocation columns + if ( + num_tmem_alloc_cols > tmem_max_alloc_cols + or num_tmem_alloc_cols < tmem_min_alloc_cols + ): + raise ValueError( + f"TMEM allocation columns {num_tmem_alloc_cols} exceeds the maximum capacity of {tmem_max_alloc_cols} or less than {tmem_min_alloc_cols}" + ) + return num_tmem_alloc_cols diff --git a/python/CuTeDSL/cutlass/utils/version_info.py b/python/CuTeDSL/cutlass/utils/version_info.py new file mode 100644 index 00000000..231856ec --- /dev/null +++ b/python/CuTeDSL/cutlass/utils/version_info.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NvidiaProprietary +# +# Use of this software is governed by the terms and conditions of the +# NVIDIA End User License Agreement (EULA), available at: +# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html +# +# Any use, reproduction, disclosure, or distribution of this software +# and related documentation outside the scope permitted by the EULA +# is strictly prohibited. + +from ..cutlass_dsl import DSLCudaVersion, DSLRuntimeError + +try: + from .._mlir._mlir_libs._cutlass_ir._base_dsl import get_cuda_version + CUDA_VERSION = DSLCudaVersion(get_cuda_version()) +except Exception as e: + raise DSLRuntimeError( + "💥💥💥 Failed to get CUDA version 💥💥💥", + cause=e, + suggestion="Consider re-installing the package." + ) from e diff --git a/python/CuTeDSL/requirements.txt b/python/CuTeDSL/requirements.txt index 508ab5d4..4c4fb948 100644 --- a/python/CuTeDSL/requirements.txt +++ b/python/CuTeDSL/requirements.txt @@ -1,3 +1,3 @@ # Use `pip install -r requirements.txt` with the present file to install a # wheel consistent with the present state of the github repository -nvidia-cutlass-dsl==4.3.5 +nvidia-cutlass-dsl==4.4.0.dev0 diff --git a/python/cutlass_library/emit_kernel_listing.py b/python/cutlass_library/emit_kernel_listing.py index a68898e4..4041caee 100755 --- a/python/cutlass_library/emit_kernel_listing.py +++ b/python/cutlass_library/emit_kernel_listing.py @@ -48,6 +48,7 @@ import csv import json import math import os +import re try: import builtins @@ -413,13 +414,49 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode raise Exception(error_message) elif mode == "functional_L1": + sm100_mma_data_type_general = [ + 'gemm_f16_f16_f16_f16_f16', + 'gemm_f16_f16_f16_void_f16', + 'gemm_f16_f16_f32_f32_f32', + 'gemm_bf16_bf16_f32_bf16_bf16', + 'gemm_bf16_bf16_f32_f32_f32', + 'gemm_e2m1_e2m1_f32_f32_f32', + 'gemm_e2m1_e3m2_f32_f32_f32', + 'gemm_e2m1_e4m3_f32_f32_f32', + 'gemm_e3m2_e2m1_f32_f32_f32', + 'gemm_e3m2_e3m2_f32_f32_f32', + 'gemm_e3m2_e4m3_f32_f32_f32', + 'gemm_e4m3_e2m1_f32_f32_f32', + 'gemm_e4m3_e3m2_f32_f32_f32', + 'gemm_e4m3_e4m3_f32_bf16_bf16', + 'gemm_e4m3_e5m2_f32_bf16_e4m3', + 'gemm_e5m2_e4m3_f32_f16_e4m3', + 'gemm_s8_s8_s32_s32_s32', + 'gemm_s8_s8_s32_s8_s8', + 'gemm_f4_f4_f32_f32_f32', + 'gemm_f4_f6_f32_f32_f32', + 'gemm_f4_f8_f32_f32_f32', + 'gemm_f6_f4_f32_f32_f32', + 'gemm_f6_f6_f32_f32_f32', + 'gemm_f6_f8_f32_f32_f32', + 'gemm_f8_f4_f32_f32_f32', + 'gemm_f8_f6_f32_f32_f32', + 'gemm_f8_f8_f32_bf16_bf16', + 'gemm_f8_f8_f32_bf16_e4m3', + 'gemm_f8_f8_f32_bf16_e5m2', + 'gemm_f8_f8_f32_f16_e4m3', + 'gemm_f8_f8_f32_f16_e5m2', + 'gemm_f8_f8_f32_f16_f16', + 'gemm_f8_f8_f32_f32_f32', + 'tf32gemm_*', + ] sm100_mma_cluster_size = [ '0x0x1' # dynamic cluster ] # Restrict to two layouts to reduce L1 build and test time. sm100_mma_layouts = ['tnt', 'ntn'] - sm100_mma_filter_regex_1sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*1sm.*" - sm100_mma_filter_regex_2sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*2sm.*" + sm100_mma_filter_regex_1sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_data_type_general, sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*1sm.*" + sm100_mma_filter_regex_2sm = "cutlass3x_sm100_tensorop.*(" + ").*(".join([ "|".join(x) for x in [sm100_mma_data_type_general, sm100_mma_cluster_size, sm100_mma_layouts]]) + ").*2sm.*" block_scaled_data_type = [ 'ue8m0xe2m1_ue8m0xe2m1_f32_f16_e5m2', 'ue8m0xe2m1_ue8m0xe2m3_f32_f16_e5m2', @@ -446,7 +483,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode filter_regex_sm100_mma = f"({sm100_mma_filter_regex_1sm})|" \ f"({sm100_mma_filter_regex_2sm})|" \ f"({block_scaled_filter_regex_1sm})|" \ - f"({block_scaled_filter_regex_2sm})" \ + f"({block_scaled_filter_regex_2sm})|" \ f"({sm103_block_scaled_filter_regex_1sm})|" \ f"({sm103_block_scaled_filter_regex_2sm})" # CTA tiles for sm120 MMA - only run one tile size to reduce build/test times @@ -483,7 +520,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode filter_regex_sm120_mma = f"({filter_regex_sm120_mma_0})|({filter_regex_sm120_mma_1})|({filter_regex_sm120_mma_2})|({filter_regex_sm120_mma_3})" - problem_waves = [0.5, 1.25, 2.5] + problem_waves = [0.5, 2.5] if arch in ["120a", "120f", "121a", "121f"]: kernel_filter = f"({filter_regex_sm120_mma})" @@ -538,6 +575,7 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode runtime_input_datatypes = [None] if dynamic_datatype: + # Standard runtime datatype kernels encoded as f4_f4 / f6_f6 / f8_f8, etc. if "f4_f4" in kernel_name: runtime_input_datatypes = [['e2m1','e2m1']] elif "f4_f6" in kernel_name: @@ -588,6 +626,23 @@ def emit_gemm_kernel_testlist(manifest, curr_build_dir, arch, mode elif "ue8m0xf8_ue8m0xf8" in kernel_name: runtime_input_datatypes = [['e4m3','e4m3']] + # Blockwise runtime-datatype kernels encode the fp8 selector together with the + # accumulator precision and block tile, e.g.: + # gemm_64x128f32xf8_32x128f32xf8_... + # which does not contain an "f8_f8" substring. As a fallback, detect this + # encoding and map it to the same runtime input datatypes as the symmetric + # f4_f4 / f6_f6 / f8_f8 cases above. + if runtime_input_datatypes == [None]: + m = re.search(r"f32x(f[468])", kernel_name) + if m: + fp_token = m.group(1) + if fp_token == "f4": + runtime_input_datatypes = [['e2m1', 'e2m1']] + elif fp_token == "f6": + runtime_input_datatypes = [['e3m2', 'e3m2']] + elif fp_token == "f8": + runtime_input_datatypes = [['e4m3', 'e4m3']] + if "bstensorop" in kernel_name or is_blockwise(manifest.operations_by_name[kernel_name].gemm_kind): profiler_flags_for_verification = "host" diff --git a/python/cutlass_library/generator.py b/python/cutlass_library/generator.py index 350ebef5..546f193f 100644 --- a/python/cutlass_library/generator.py +++ b/python/cutlass_library/generator.py @@ -7487,7 +7487,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kin tile_schedulers = [ TileSchedulerType.Default ] - # Some SM100 NoSmem epilogue instantiations rely on CUTE's shape_div, which enforces a compile-time # divisibility condition between CTA N and the epilogue N tile. Keep this conservative and scoped: # only apply the divisibility filter for selected common (c_type, d_type) pairs. @@ -7498,7 +7497,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kin (DataType.void, DataType.f16): 64, (DataType.void, DataType.bf16): 64, } - # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] @@ -7618,7 +7616,6 @@ def GenerateSM100_TensorOp_fp8_UMMA_alignx_gemm(manifest, cuda_version, gemm_kin kernel_schedule = KernelScheduleType.WarpSpecialized1SmSm100 epi_schedule = EpilogueScheduleType.NoSmemWarpSpecialized1Sm - # SM100 NoSmem epilogue uses EpilogueTileAuto with N-tile = min(64, cta_n). # CUTE's shape_div then requires a compile-time divisibility condition between cta_n and 64. # Only instantiate kernels where cta_n <= 64 or cta_n is an exact multiple of 64 to avoid @@ -10625,6 +10622,232 @@ def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version): [[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) + +# SM100 Interleaved Complex Tf32 Kernels +def GenerateSM100_TensorOp_32b_UMMA_gemm_complex(manifest, cuda_version): + if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): + return + + # layouts for ABC and their alignments. + layouts = [ + [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]], + [[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2]] + ] + + complex_transforms = [ + (ComplexTransform.none, ComplexTransform.none), + (ComplexTransform.conj, ComplexTransform.none), + (ComplexTransform.none, ComplexTransform.conj), + (ComplexTransform.conj, ComplexTransform.conj) + ] + + data_types = [ + { + "a_type" : DataType.cf32, + "b_type" : DataType.cf32, + "c_type" : DataType.cf32, + "d_type" : DataType.cf32, + "acc_type" : DataType.cf32, + "epi_type" : DataType.cf32, + }, + { + "a_type" : DataType.cf32, + "b_type" : DataType.cf32, + "c_type" : DataType.void, + "d_type" : DataType.cf32, + "acc_type" : DataType.cf32, + "epi_type" : DataType.cf32, + } + ] + + thor_sm = ThorSMRenumbering(cuda_version) + + min_cc = 100 + max_cc = 100 + max_cc = max(max_cc, thor_sm) + + math_instructions_1sm = [ + # tf32 -> f32 + MathInstruction( + [128, 64, 4], + DataType.tf32, DataType.tf32, DataType.f32, + OpcodeClass.TensorOp, + MathOperation.multiply_add_complex) + ] + + cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1], [4,4,1] + , DynamicClusterShape + ] + + if thor_sm in manifest.compute_capabilities_baseline : + cluster_shapes_1sm = [[1,2,1], [1,1,1], [1,4,1] + , DynamicClusterShape + ] + + tile_schedulers = [ + TileSchedulerType.Default, TileSchedulerType.StreamK + ] + + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes_1sm: + multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_1sm[0], + math_inst.instruction_shape[1] * multiplier_1sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, + [[KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized1SmSm100, EpilogueScheduleType.NoSmemWarpSpecialized1Sm]], + complex_transforms, + tile_schedulers=tile_schedulers) + + # 2xSM MMA kernels + math_instructions_2sm = [ + # tf32 -> f32 + MathInstruction( + [256, 64, 4], + DataType.tf32, DataType.tf32, DataType.f32, + OpcodeClass.TensorOp, + MathOperation.multiply_add_complex) + ] + + cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1] + , DynamicClusterShape + ] + + if thor_sm in manifest.compute_capabilities_baseline : + cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1] + , DynamicClusterShape + ] + + for math_inst in math_instructions_2sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes_2sm: + multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_2sm[0], + math_inst.instruction_shape[1] * multiplier_2sm[1], + math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, + [[KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized2SmSm100, EpilogueScheduleType.NoSmemWarpSpecialized2Sm]], + complex_transforms, + tile_schedulers=tile_schedulers) + +def GenerateSM100_TensorOp_FastF32_UMMA_gemm_complex_stream_k(manifest, cuda_version): + if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): + return + + # layouts for ABC and their alignments. + layouts = [ + [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]], + [[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2]], + [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2]], + + ] + + data_types = [ + { + "a_type" : DataType.cf32, + "b_type" : DataType.cf32, + "c_type" : DataType.cf32, + "d_type" : DataType.cf32, + "acc_type" : DataType.cf32, + "epi_type" : DataType.cf32, + } + ] + + # Unsupported yet + complex_transforms = None + # [ + # (ComplexTransform.none, ComplexTransform.none), + # (ComplexTransform.conj, ComplexTransform.none), + # (ComplexTransform.none, ComplexTransform.conj), + # (ComplexTransform.conj, ComplexTransform.conj) + # ] + + min_cc = 100 + max_cc = 100 + + math_instructions_1sm = [ + MathInstruction( + [128, 64, 8], + DataType.cbf16, DataType.cbf16, DataType.cf32, + OpcodeClass.TensorOp, + MathOperation.multiply_add), + ] + + cluster_shapes_1sm = [ + [1,1,1], [4,4,1] + , DynamicClusterShape + ] + + tile_schedulers = [ + TileSchedulerType.Default, TileSchedulerType.StreamK, + ] + + # 1xSM MMA kernels + for math_inst in math_instructions_1sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes_1sm: + multiplier = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier[0], + math_inst.instruction_shape[1] * multiplier[1], + math_inst.instruction_shape[2] * 2], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, + [[KernelScheduleType.TmaWarpSpecialized1SmFastFP32Sm100, EpilogueScheduleType.FastF32NoSmemWarpSpecialized1Sm]], + complex_transforms, + tile_schedulers=tile_schedulers) + + # 2xSM MMA kernels + math_instructions_2sm = [ + MathInstruction( + [256, 64, 8], + DataType.cbf16, DataType.cbf16, DataType.cf32, + OpcodeClass.TensorOp, + MathOperation.multiply_add), + ] + + cluster_shapes_2sm = [ + [2,1,1], [4,4,1] + , DynamicClusterShape + ] + + for math_inst in math_instructions_2sm: + tile_descriptions = [] + for cluster_shape in cluster_shapes_2sm: + multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) + tile_descriptions.append( + TileDescription([ + math_inst.instruction_shape[0] * multiplier_2sm[0], + math_inst.instruction_shape[1] * multiplier_2sm[1], + math_inst.instruction_shape[2] * 2], + 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) + CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, + [[KernelScheduleType.TmaWarpSpecialized2SmFastFP32Sm100, EpilogueScheduleType.FastF32NoSmemWarpSpecialized2Sm]], + complex_transforms, + tile_schedulers=tile_schedulers) + + # Conv Utility functions def make_dims_and_alignments_triple(dim: int, bit_per_element_A: int, bit_per_element_B: int, bit_per_element_C: int): bit_alignment_required_by_tma = 128 @@ -11780,6 +12003,10 @@ def GenerateSM100(manifest, cuda_version): GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version) GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockwiseUniversal3x) + GenerateSM100_TensorOp_32b_UMMA_gemm_complex(manifest, cuda_version) + # CGemm with 9xBF16 + GenerateSM100_TensorOp_FastF32_UMMA_gemm_complex_stream_k(manifest, cuda_version) + # # Sparse Gemm # diff --git a/python/cutlass_library/library.py b/python/cutlass_library/library.py index 8a26474c..e78db65b 100644 --- a/python/cutlass_library/library.py +++ b/python/cutlass_library/library.py @@ -562,6 +562,12 @@ class KernelScheduleType(enum.Enum): SparseNvf4TmaWarpSpecialized2SmSm100 = enum_auto() SparseMxf8f6f4TmaWarpSpecialized1SmSm100 = enum_auto() SparseMxf8f6f4TmaWarpSpecialized2SmSm100 = enum_auto() + + InterleavedComplexTF32TmaWarpSpecialized1SmSm100 = enum_auto() + InterleavedComplexTF32TmaWarpSpecialized2SmSm100 = enum_auto() + TmaWarpSpecialized1SmFastFP32Sm100 = enum_auto() + TmaWarpSpecialized2SmFastFP32Sm100 = enum_auto() + # FP4 Ultra MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103 = enum_auto() MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103 = enum_auto() @@ -679,7 +685,10 @@ KernelScheduleTag = { KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103DisablePrefetch', KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch', KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch: 'cutlass::gemm::KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs32Sm103DisablePrefetch', - + KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized1SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmInterleavedComplexTF32Sm100', + KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized2SmSm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmInterleavedComplexTF32Sm100', + KernelScheduleType.TmaWarpSpecialized1SmFastFP32Sm100: 'cutlass::gemm::KernelTmaWarpSpecialized1SmFastFP32Sm100', + KernelScheduleType.TmaWarpSpecialized2SmFastFP32Sm100: 'cutlass::gemm::KernelTmaWarpSpecialized2SmFastFP32Sm100', KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative', KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum', KernelScheduleType.PtrArrayTmaWarpSpecializedPingpong: 'cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong', @@ -799,7 +808,10 @@ KernelScheduleSuffixes = { KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch: '_o_vs16_ultra_2sm_tmapf', KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_1sm_tmapf', KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch: '_o_vs32_ultra_2sm_tmapf', - + KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized1SmSm100: '_1sm', + KernelScheduleType.InterleavedComplexTF32TmaWarpSpecialized2SmSm100: '_2sm', + KernelScheduleType.TmaWarpSpecialized1SmFastFP32Sm100: '_FastF32_1sm', + KernelScheduleType.TmaWarpSpecialized2SmFastFP32Sm100: '_FastF32_2sm', KernelScheduleType.PtrArrayTmaWarpSpecializedCooperative: '_warpspecialized_cooperative', KernelScheduleType.PtrArrayTmaWarpSpecializedCooperativeFP8FastAccum: '_warpspecialized_cooperative_fp8_fastaccum', KernelScheduleType.PtrArrayTmaWarpSpecializedPingpong: '_warpspecialized_pingpong', diff --git a/test/self_contained_includes/CMakeLists.txt b/test/self_contained_includes/CMakeLists.txt index 094cfeea..97d1ce0c 100644 --- a/test/self_contained_includes/CMakeLists.txt +++ b/test/self_contained_includes/CMakeLists.txt @@ -260,11 +260,11 @@ set(header_files_to_check ) # for each header in _header_files: -# create a .cu file with the same name as the header's path, except with / replaced with % +# create a .cu file with the same name as the header's path, except with / replaced with # # have the .cu file include that header set(_gen_source_files "") foreach(header_file ${header_files_to_check}) - string(REPLACE "/" "%" header_file_esc ${header_file}) + string(REPLACE "/" "#" header_file_esc ${header_file}) file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/${header_file_esc}.cu" "#include <${header_file}>") diff --git a/test/unit/gemm/device/gemm_testbed_3x_planar_complex.hpp b/test/unit/gemm/device/gemm_testbed_3x_planar_complex.hpp new file mode 100644 index 00000000..d723bc2c --- /dev/null +++ b/test/unit/gemm/device/gemm_testbed_3x_planar_complex.hpp @@ -0,0 +1,482 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide Planar Complex GEMM interface +*/ + + + +#pragma once + +#include +#include +#include + +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x.hpp" + +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/host_tensor_planar_complex.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/gemm_planar_complex.h" +#include "cutlass/numeric_types.h" + +#include "testbed_utils.h" + +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/layout/matrix.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/fast_math.h" +#include "cutlass/platform/platform.h" + +#include "cute/int_tuple.hpp" +#include "cute/layout.hpp" + + +namespace test { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Testbed3xPlanarComplex { + // Kernel data types + using ElementA = typename Gemm::GemmKernel::ElementA; + using StrideA = typename Gemm::GemmKernel::StrideA; + using ElementB = typename Gemm::GemmKernel::ElementB; + using StrideB = typename Gemm::GemmKernel::StrideB; + using ElementC = std::conditional_t, + typename Gemm::GemmKernel::ElementD,typename Gemm::GemmKernel::ElementC>; + using StrideC = typename Gemm::GemmKernel::StrideC; + using ElementD = typename Gemm::GemmKernel::ElementD; + using StrideD = typename Gemm::GemmKernel::StrideD; + using ElementAccumulator = typename Gemm::GemmKernel::ElementAccumulator; + using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape; + using EpilogueOutputOp = typename Gemm::EpilogueOutputOp; + using ClusterShapeType = typename Gemm::GemmKernel::CollectiveMainloop::DispatchPolicy::ClusterShape; + /// For custom EVTs + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + using ElementScalar = typename EpilogueOutputOp::ElementScalar; + + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + + static constexpr uint32_t mma_promotion_interval = 4; + + // Looks at Cute Stride to check Row / Column Major + template + static constexpr bool is_row_or_col_major(){ + int stride_0 = int(cute::size<0>(Stride{})); + int stride_1 = int(cute::size<1>(Stride{})); + int depth = cute::depth(Stride{}); + return ((stride_0 == 1) || (stride_1 == 1)) && (depth == 1); + } + + // Note: this limitation comes from testbed / not the library + static_assert(is_row_or_col_major(), + "ERROR : A Layout is neither Row / Column Major)"); + static_assert(is_row_or_col_major(), + "ERROR : B Layout is neither Row / Column Major)"); + static_assert(is_row_or_col_major(), + "ERROR : C Layout is neither Row / Column Major)"); + static_assert(is_row_or_col_major(), + "ERROR : D Layout is neither Row / Column Major)"); + + // Deduce Cutlass Layouts (RowMajor & ColumnMajor) + using LayoutTagA = cutlass::detail::StrideToLayoutTagA_t; + using LayoutTagB = cutlass::detail::StrideToLayoutTagB_t; + using LayoutTagC = cutlass::detail::StrideToLayoutTagA_t; + using LayoutTagD = cutlass::detail::StrideToLayoutTagA_t; + + /// Initialization + StrideA stride_a; + StrideB stride_b; + StrideC stride_c; + StrideD stride_d; + + typename LayoutTagA::Stride stride_factor_A; + typename LayoutTagB::Stride stride_factor_B; + typename LayoutTagC::Stride stride_factor_C; + typename LayoutTagD::Stride stride_factor_D; + + cutlass::Distribution::Kind init_A; + cutlass::Distribution::Kind init_B; + cutlass::Distribution::Kind init_C; + uint64_t seed; + static constexpr uint64_t kDefaultSeed = 4096; + + // Data members + cutlass::HostTensorPlanarComplex tensor_A; + cutlass::HostTensorPlanarComplex tensor_B; + cutlass::HostTensorPlanarComplex tensor_C; + cutlass::HostTensorPlanarComplex tensor_D; + cutlass::HostTensorPlanarComplex reference_D; + + uint32_t sm_count; + + // Used to force multi-wave tests for persistent kernel schedules + constexpr static int MaxSmCount = 16; + using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions; + using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode; + + cutlass::ComplexTransform TransformA = Gemm::kTransformA; + cutlass::ComplexTransform TransformB = Gemm::kTransformB; + + // + // Methods + // + Testbed3xPlanarComplex( + cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + uint64_t seed_ = kDefaultSeed + ): + stride_factor_A(typename LayoutTagA::Stride()), + stride_factor_B(typename LayoutTagB::Stride()), + stride_factor_C(typename LayoutTagC::Stride()), + stride_factor_D(typename LayoutTagD::Stride()), + init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { } + + /// Helper to initialize a tensor view + template + bool initialize_tensor( + cutlass::TensorViewPlanarComplex view, + cutlass::Distribution::Kind dist_kind, + uint64_t seed) { + + if (dist_kind == cutlass::Distribution::Uniform) { + double scope_max, scope_min; + int bits_input = cutlass::sizeof_bits::value; + int bits_output = cutlass::sizeof_bits::value; + + if (bits_input == 1) { + scope_max = 2; + scope_min = 0; + } + else if (bits_input <= 8) { + scope_max = 2; + scope_min = -2; + } + else if (bits_output == 16) { + scope_max = 5; + scope_min = -5; + } + else { + scope_max = 8; + scope_min = -8; + } + cutlass::reference::host::TensorFillRandomUniform( + view, seed, scope_max, scope_min, 0); + } + + else if (dist_kind == cutlass::Distribution::Gaussian) { + cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + } + + else if (dist_kind == cutlass::Distribution::AllOnes) { + cutlass::reference::host::TensorFill(view, {Element(1), Element(0)}); + } + else { + EXPECT_TRUE(false) << "Not implemented"; + return false; + } + + return true; + } + + /// Initializes data structures + void initialize(ProblemShapeType problem_size) { + // + // Allocate the GEMM workspace + // + auto problem_shape_MNKL = cute::append<4>(problem_size, 1); + auto M = cute::size<0>(problem_shape_MNKL); + auto N = cute::size<1>(problem_shape_MNKL); + auto K = cute::size<2>(problem_shape_MNKL); + auto L = cute::size<3>(problem_shape_MNKL); + + 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)); + + // 2.x host tensor does not natively contain a batch stride or coord, so we spoof if by folding it into the outer mode + auto a_coord = cutlass::make_Coord(M * L, K); + auto c_coord = cutlass::make_Coord(M * L, N); + // Cutlass has Row/Col major refers to MxK times KxN matrix product, + // so the HostTensorB should be treated as KxN in "coord"'s view + auto b_coord = cutlass::make_Coord(K, N * L); + + tensor_A.resize(a_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(a_coord, stride_factor_A)); + tensor_B.resize(b_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(b_coord, stride_factor_B)); + tensor_C.resize(c_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(c_coord, stride_factor_C)); + tensor_D.resize(c_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(c_coord, stride_factor_D)); + + reference_D.resize(c_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(c_coord, stride_factor_D), false); + + EXPECT_TRUE(initialize_tensor(tensor_A.host_view(), init_A, seed + 2022)); + EXPECT_TRUE(initialize_tensor(tensor_B.host_view(), init_B, seed + 2021)); + EXPECT_TRUE(initialize_tensor(tensor_C.host_view(), init_C, seed + 2020)); + + cutlass::reference::host::TensorFill(tensor_D.host_view(), cutlass::complex()); + cutlass::reference::host::TensorFill(reference_D.host_view(), cutlass::complex()); + + tensor_A.sync_device(); + tensor_B.sync_device(); + tensor_C.sync_device(); + tensor_D.sync_device(); + } + + /// Verifies the result is a GEMM + bool verify( + ProblemShapeType problem_size, + ElementScalar alpha, + ElementScalar beta) + { + auto problem_shape_MNKL = cute::append<4>(problem_size, 1); + auto M = cute::size<0>(problem_shape_MNKL); + auto N = cute::size<1>(problem_shape_MNKL); + auto K = cute::size<2>(problem_shape_MNKL); + auto L = cute::size<3>(problem_shape_MNKL); + +#if 0 + std::cout << " M : " << M << " N : " << N << " K : " << K << " L : " << L << std::endl; +#endif + // + // Compute reference + // + cutlass::reference::host::GemmPlanarComplex< + ElementA, LayoutTagA, + ElementB, LayoutTagB, + ElementC, LayoutTagC, + ElementAccumulator + >( + cutlass::gemm::GemmCoord(M,N,K), + alpha, + tensor_A.host_ref(), + TransformA, + tensor_B.host_ref(), + TransformB, + beta, + tensor_C.host_ref(), + reference_D.host_ref() + ); + + bool passed = false; + + tensor_D.sync_host(); + passed = cutlass::reference::host::TensorEquals( + tensor_D.host_view(), + reference_D.host_view() + ); + + EXPECT_TRUE(passed); + + if (!passed) { + std::stringstream fname; + fname << "error_Planar_Complex_Gemm_device_" + << M << "x" << N << "x" << K << "x" << L << "_" + << cute::get<0>(typename Gemm::GemmKernel::TileShape{}) << "_" + << cute::get<1>(typename Gemm::GemmKernel::TileShape{}) << "_" + << cute::get<2>(typename Gemm::GemmKernel::TileShape{}) << ".txt"; + + std::ofstream file(fname.str()); + file + << "problem: " << ' ' << M << "x" << N << "x" << K << ", Batch count = " << L + << ", alpha: " << alpha << ", beta: " << beta + << "\n\n"; + + file + << "A =\n" << tensor_A.host_view() + << "\nB =\n" << tensor_B.host_view() + << "\nC =\n" << tensor_C.host_view() + << "\n\nReference =\n" << reference_D.host_view() + << "\n\nComputed =\n" << tensor_D.host_view(); + } + + return passed; + } + + /// Returns true if the CUDA device is sufficient to execute the kernel. + bool sufficient() { + // + // Determine SMEM requirements and waive if not satisfied + // + + int smem_size = Gemm::GemmKernel::SharedStorageSize; + + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + cudaDeviceProp properties; + result = cudaGetDeviceProperties(&properties, device_idx); + this->sm_count = properties.multiProcessorCount; + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDeviceProperties() failed"); + } + + if (properties.sharedMemPerBlockOptin < size_t(smem_size)) { + return false; + } + + return true; + } + + bool run( + ProblemShapeType problem_size, + ElementScalar alpha = ElementScalar(1), + ElementScalar beta = ElementScalar(0), + RasterOrderOptions raster_order = RasterOrderOptions::Heuristic, + detail::MaxSwizzleSize max_swizzle = detail::MaxSwizzleSize{}, + detail::Splits splits = detail::Splits{}, + DecompositionMode decomposition_mode = DecompositionMode::Heuristic, + unsigned int cluster_m = 0, + unsigned int cluster_n = 0, + unsigned int cluster_m_fallback = 0, + unsigned int cluster_n_fallback = 0 + ) { + // Waive test if insufficient CUDA device + if (!sufficient()) { + if (CUTLASS_TEST_UNIT_ENABLE_WARNINGS) { + std::cerr << "Test waived due to insufficient CUDA device." << std::endl; + } + return true; + } + + this->initialize(problem_size); + + // + // Launch device kernel + // + + // + // Initialize the GEMM operator + // + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + + if (cute::is_static_v) { + this->sm_count = cutlass::platform::min(MaxSmCount, cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id)); + hw_info.sm_count = this->sm_count; + } + else { + this->sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + hw_info.sm_count = this->sm_count; + + // Runtime and preferred cluster setting + hw_info.cluster_shape = {cluster_m, cluster_n, 1}; + hw_info.cluster_shape_fallback = {cluster_m_fallback, cluster_n_fallback, 1}; + } + + typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args; + if constexpr (cute::is_same_v) { + scheduler_args = { static_cast(splits), static_cast(max_swizzle), raster_order, decomposition_mode }; + } + else { + scheduler_args = { static_cast(max_swizzle), raster_order }; + } + + auto arguments = typename Gemm::Arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + problem_size, + { + tensor_A.device_data(), stride_a, tensor_A.device_data_imag(), stride_a, + tensor_B.device_data(), stride_b, tensor_B.device_data_imag(), stride_b + }, + { + {alpha, beta}, + tensor_C.device_data(), stride_c, tensor_C.device_data_imag(), stride_c, + tensor_D.device_data(), stride_d, tensor_D.device_data_imag(), stride_d + }, + hw_info, + scheduler_args + }; + + Gemm gemm_op; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + cutlass::Status status = gemm_op.can_implement(arguments); + + if (status != cutlass::Status::kSuccess) { + cudaError_t error = cudaGetLastError(); + std::cerr << "This test is not supported: " << cudaGetErrorString(error) << "\n"; + return true; + } + + // + // Run the GEMM + // + + cudaError_t result; + status = gemm_op.initialize(arguments, workspace.get()); + status = gemm_op.run(); + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + EXPECT_EQ(result, cudaSuccess) << "Error at Kernel Sync."; + return false; + } + + EXPECT_TRUE(status == cutlass::Status::kSuccess) << to_string(status); + // + // Verify + // + bool passed = this->verify(problem_size, alpha, beta); + + if (!passed) { + std::cout << "Error : Failed : with alpha: " << alpha << ", beta: " << beta + << "\n"; + } + + return passed; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace test diff --git a/test/unit/gemm/device/gemm_testbed_3x_ptr_array_planar_complex.hpp b/test/unit/gemm/device/gemm_testbed_3x_ptr_array_planar_complex.hpp new file mode 100644 index 00000000..01488c92 --- /dev/null +++ b/test/unit/gemm/device/gemm_testbed_3x_ptr_array_planar_complex.hpp @@ -0,0 +1,536 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Testbed for Ptr-Array and Grouped Planar Complex GEMM interface +*/ + + + +#pragma once + +#include +#include +#include + +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x.hpp" + +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/host_tensor_planar_complex.h" +#include "cutlass/util/tensor_view_io.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/host/tensor_fill.h" +#include "cutlass/util/reference/host/tensor_copy.h" +#include "cutlass/util/reference/host/tensor_compare.h" +#include "cutlass/util/reference/host/tensor_norm.h" +#include "cutlass/util/reference/host/gemm_planar_complex.h" +#include "cutlass/numeric_types.h" + +#include "testbed_utils.h" + +#include "cutlass/kernel_hardware_info.hpp" +#include "cutlass/layout/matrix.h" +#include "cutlass/matrix_coord.h" +#include "cutlass/gemm/gemm.h" +#include "cutlass/fast_math.h" +#include "cutlass/platform/platform.h" + +#include "cute/int_tuple.hpp" +#include "cute/layout.hpp" + + +namespace test { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Testbed3xPlanarComplex { + // Kernel data types + using ElementA = typename Gemm::GemmKernel::ElementA; + using StrideA = typename Gemm::GemmKernel::StrideA; + using ElementB = typename Gemm::GemmKernel::ElementB; + using StrideB = typename Gemm::GemmKernel::StrideB; + using ElementC = std::conditional_t, + typename Gemm::GemmKernel::ElementD,typename Gemm::GemmKernel::ElementC>; + using StrideC = typename Gemm::GemmKernel::StrideC; + using ElementD = typename Gemm::GemmKernel::ElementD; + using StrideD = typename Gemm::GemmKernel::StrideD; + using ElementAccumulator = typename Gemm::GemmKernel::ElementAccumulator; + using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape; + using EpilogueOutputOp = typename Gemm::EpilogueOutputOp; + using ClusterShapeType = typename Gemm::GemmKernel::CollectiveMainloop::DispatchPolicy::ClusterShape; + /// For custom EVTs + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + using ElementScalar = typename EpilogueOutputOp::ElementScalar; + + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + + static constexpr uint32_t mma_promotion_interval = 4; + + // Looks at Cute Stride to check Row / Column Major + template + static constexpr bool is_row_or_col_major(){ + int stride_0 = int(cute::size<0>(Stride{})); + int stride_1 = int(cute::size<1>(Stride{})); + int depth = cute::depth(Stride{}); + return ((stride_0 == 1) || (stride_1 == 1)) && (depth == 1); + } + + // Note: this limitation comes from testbed / not the library + static_assert(is_row_or_col_major(), + "ERROR : A Layout is neither Row / Column Major)"); + static_assert(is_row_or_col_major(), + "ERROR : B Layout is neither Row / Column Major)"); + static_assert(is_row_or_col_major(), + "ERROR : C Layout is neither Row / Column Major)"); + static_assert(is_row_or_col_major(), + "ERROR : D Layout is neither Row / Column Major)"); + + // Deduce Cutlass Layouts (RowMajor & ColumnMajor) + using LayoutTagA = cutlass::detail::StrideToLayoutTagA_t; + using LayoutTagB = cutlass::detail::StrideToLayoutTagB_t; + using LayoutTagC = cutlass::detail::StrideToLayoutTagA_t; + using LayoutTagD = cutlass::detail::StrideToLayoutTagA_t; + + /// Initialization + StrideA stride_a; + StrideB stride_b; + StrideC stride_c; + StrideD stride_d; + + typename LayoutTagA::Stride stride_factor_A; + typename LayoutTagB::Stride stride_factor_B; + typename LayoutTagC::Stride stride_factor_C; + typename LayoutTagD::Stride stride_factor_D; + + cutlass::Distribution::Kind init_A; + cutlass::Distribution::Kind init_B; + cutlass::Distribution::Kind init_C; + uint64_t seed; + static constexpr uint64_t kDefaultSeed = 4096; + + // Data members + std::vector> tensors_A; + std::vector> tensors_B; + std::vector> tensors_C; + std::vector> tensors_D; + std::vector> references_D; + + cutlass::DeviceAllocation device_tensors_A_real; + cutlass::DeviceAllocation device_tensors_A_imag; + + cutlass::DeviceAllocation device_tensors_B_real; + cutlass::DeviceAllocation device_tensors_B_imag; + + cutlass::DeviceAllocation device_tensors_C_real; + cutlass::DeviceAllocation device_tensors_C_imag; + + cutlass::DeviceAllocation device_tensors_D_real; + cutlass::DeviceAllocation device_tensors_D_imag; + + uint32_t sm_count; + + // Used to force multi-wave tests for persistent kernel schedules + constexpr static int MaxSmCount = 16; + using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions; + using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode; + + cutlass::ComplexTransform TransformA = Gemm::kTransformA; + cutlass::ComplexTransform TransformB = Gemm::kTransformB; + + // + // Methods + // + Testbed3xPlanarComplex( + cutlass::Distribution::Kind init_A_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_B_ = cutlass::Distribution::Uniform, + cutlass::Distribution::Kind init_C_ = cutlass::Distribution::Uniform, + uint64_t seed_ = kDefaultSeed + ): + stride_factor_A(typename LayoutTagA::Stride()), + stride_factor_B(typename LayoutTagB::Stride()), + stride_factor_C(typename LayoutTagC::Stride()), + stride_factor_D(typename LayoutTagD::Stride()), + init_A(init_A_), init_B(init_B_), init_C(init_C_), seed(seed_) { } + + /// Helper to initialize a tensor view + template + bool initialize_tensor( + cutlass::TensorViewPlanarComplex view, + cutlass::Distribution::Kind dist_kind, + uint64_t seed) { + + if (dist_kind == cutlass::Distribution::Uniform) { + double scope_max, scope_min; + int bits_input = cutlass::sizeof_bits::value; + int bits_output = cutlass::sizeof_bits::value; + + if (bits_input == 1) { + scope_max = 2; + scope_min = 0; + } + else if (bits_input <= 8) { + scope_max = 2; + scope_min = -2; + } + else if (bits_output == 16) { + scope_max = 5; + scope_min = -5; + } + else { + scope_max = 8; + scope_min = -8; + } + cutlass::reference::host::TensorFillRandomUniform( + view, seed, scope_max, scope_min, 0); + } + + else if (dist_kind == cutlass::Distribution::Gaussian) { + cutlass::reference::host::TensorFillRandomGaussian(view, seed, 0, 0.5); + } + + else if (dist_kind == cutlass::Distribution::AllOnes) { + cutlass::reference::host::TensorFill(view, {Element(1), Element(0)}); + } + else { + EXPECT_TRUE(false) << "Not implemented"; + return false; + } + + return true; + } + + /// Initializes data structures + void initialize(ProblemShapeType problem_shapes) { + // + // Allocate the GEMM workspace + // + auto [M, N, K, L] = problem_shapes.get_host_problem_shape(); + + stride_a = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1)); + stride_b = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1)); + stride_c = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1)); + stride_d = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1)); + + auto a_coord = cutlass::make_Coord(M, K); + auto c_coord = cutlass::make_Coord(M, N); + // Cutlass has Row/Col major refers to MxK times KxN matrix product, + // so the HostTensorB should be treated as KxN in "coord"'s view + auto b_coord = cutlass::make_Coord(K, N); + + tensors_A.clear(); + tensors_B.clear(); + tensors_C.clear(); + tensors_D.clear(); + references_D.clear(); + + for (int32_t i = 0; i < L; ++i) { + tensors_A.push_back(cutlass::HostTensorPlanarComplex(a_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(a_coord, stride_factor_A))); + tensors_B.push_back(cutlass::HostTensorPlanarComplex(b_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(b_coord, stride_factor_B))); + tensors_C.push_back(cutlass::HostTensorPlanarComplex(c_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(c_coord, stride_factor_C))); + tensors_D.push_back(cutlass::HostTensorPlanarComplex(c_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(c_coord, stride_factor_D))); + references_D.push_back(cutlass::HostTensorPlanarComplex(c_coord, cutlass::layout::Affine2Layout_Factory::layout_factory(c_coord, stride_factor_D), false)); + + EXPECT_TRUE(initialize_tensor(tensors_A[i].host_view(), init_A, seed + 2022 + i)); + EXPECT_TRUE(initialize_tensor(tensors_B[i].host_view(), init_B, seed + 2021 + i)); + EXPECT_TRUE(initialize_tensor(tensors_C[i].host_view(), init_B, seed + 2020 + i)); + + cutlass::reference::host::TensorFill(tensors_D[i].host_view(), cutlass::complex()); + cutlass::reference::host::TensorFill(references_D[i].host_view(), cutlass::complex()); + + tensors_A[i].sync_device(); + tensors_B[i].sync_device(); + tensors_C[i].sync_device(); + tensors_D[i].sync_device(); + } + } + + /// Verifies the result is a GEMM + bool verify( + ProblemShapeType problem_shapes, + ElementScalar alpha, + ElementScalar beta) + { + using namespace cute; + auto [M, N, K, L] = problem_shapes.get_host_problem_shape(); + +#if 0 + std::cout << " M : " << M << " N : " << N << " K : " << K << " L : " << L << std::endl; +#endif + + // + // Compute reference + // + bool passed = false; + for(int batch = 0; batch < L; ++batch) { + cutlass::reference::host::GemmPlanarComplex< + ElementA, LayoutTagA, + ElementB, LayoutTagB, + ElementC, LayoutTagC, + ElementAccumulator + >( + cutlass::gemm::GemmCoord(M,N,K), + alpha, + tensors_A[batch].host_ref(), + TransformA, + tensors_B[batch].host_ref(), + TransformB, + beta, + tensors_C[batch].host_ref(), + references_D[batch].host_ref() + ); + + tensors_D[batch].sync_host(); + passed = cutlass::reference::host::TensorEquals( + tensors_D[batch].host_view(), + references_D[batch].host_view() + ); + + EXPECT_TRUE(passed); + + if (!passed) { + std::stringstream fname; + fname << "error_Planar_Complex_Gemm_device_" + << M << "x" << N << "x" << K << "x" << batch << "_" + << cute::get<0>(typename Gemm::GemmKernel::TileShape{}) << "_" + << cute::get<1>(typename Gemm::GemmKernel::TileShape{}) << "_" + << cute::get<2>(typename Gemm::GemmKernel::TileShape{}) << ".txt"; + + std::ofstream file(fname.str()); + file + << "problem: " << ' ' << M << "x" << N << "x" << K << ", Batch idx = " << batch + << ", alpha: " << alpha << ", beta: " << beta + << "\n\n"; + + file + << "A =\n" << tensors_A[batch].host_view() + << "\nB =\n" << tensors_B[batch].host_view() + << "\nC =\n" << tensors_C[batch].host_view() + << "\n\nReference =\n" << references_D[batch].host_view() + << "\n\nComputed =\n" << tensors_D[batch].host_view(); + } + } + return passed; + } + + /// Returns true if the CUDA device is sufficient to execute the kernel. + bool sufficient() { + // + // Determine SMEM requirements and waive if not satisfied + // + + int smem_size = Gemm::GemmKernel::SharedStorageSize; + + int device_idx; + cudaError_t result = cudaGetDevice(&device_idx); + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDevice() API call failed."); + } + + cudaDeviceProp properties; + result = cudaGetDeviceProperties(&properties, device_idx); + this->sm_count = properties.multiProcessorCount; + + if (result != cudaSuccess) { + throw std::runtime_error("cudaGetDeviceProperties() failed"); + } + + if (properties.sharedMemPerBlockOptin < size_t(smem_size)) { + return false; + } + + return true; + } + + bool run( + ProblemShapeType problem_shapes, + ElementScalar alpha = ElementScalar(1), + ElementScalar beta = ElementScalar(0), + RasterOrderOptions raster_order = RasterOrderOptions::Heuristic, + detail::MaxSwizzleSize max_swizzle = detail::MaxSwizzleSize{}, + detail::Splits splits = detail::Splits{}, + DecompositionMode decomposition_mode = DecompositionMode::Heuristic, + unsigned int cluster_m = 0, + unsigned int cluster_n = 0, + unsigned int cluster_m_fallback = 0, + unsigned int cluster_n_fallback = 0 + ) { + // Waive test if insufficient CUDA device + if (!sufficient()) { + if (CUTLASS_TEST_UNIT_ENABLE_WARNINGS) { + std::cerr << "Test waived due to insufficient CUDA device." << std::endl; + } + return true; + } + + this->initialize(problem_shapes); + + // + // Launch device kernel + // + + // + // Initialize the GEMM operator + // + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + + this->sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + hw_info.sm_count = this->sm_count; + if (not cute::is_static_v) { + // Runtime and preferred cluster setting + hw_info.cluster_shape = {cluster_m, cluster_n, 1}; + hw_info.cluster_shape_fallback = {cluster_m_fallback, cluster_n_fallback, 1}; + } + + typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args; + if constexpr (cute::is_same_v) { + scheduler_args = { static_cast(splits), static_cast(max_swizzle), raster_order, decomposition_mode }; + } + else { + scheduler_args = { static_cast(max_swizzle), raster_order }; + } + + auto [M, N, K, L] = problem_shapes.get_host_problem_shape(); + std::vector ptr_A_real_host(L); + std::vector ptr_A_imag_host(L); + + std::vector ptr_B_real_host(L); + std::vector ptr_B_imag_host(L); + + std::vector ptr_C_real_host(L); + std::vector ptr_C_imag_host(L); + + std::vector ptr_D_real_host(L); + std::vector ptr_D_imag_host(L); + + for (int32_t i = 0; i < L; ++i) { + ptr_A_real_host.at(i) = tensors_A[i].device_data(); + ptr_A_imag_host.at(i) = tensors_A[i].device_data_imag(); + + ptr_B_real_host.at(i) = tensors_B[i].device_data(); + ptr_B_imag_host.at(i) = tensors_B[i].device_data_imag(); + + ptr_C_real_host.at(i) = tensors_C[i].device_data(); + ptr_C_imag_host.at(i) = tensors_C[i].device_data_imag(); + + ptr_D_real_host.at(i) = tensors_D[i].device_data(); + ptr_D_imag_host.at(i) = tensors_D[i].device_data_imag(); + } + + device_tensors_A_real.reset(L); + device_tensors_A_real.copy_from_host(ptr_A_real_host.data()); + device_tensors_A_imag.reset(L); + device_tensors_A_imag.copy_from_host(ptr_A_imag_host.data()); + + device_tensors_B_real.reset(L); + device_tensors_B_real.copy_from_host(ptr_B_real_host.data()); + device_tensors_B_imag.reset(L); + device_tensors_B_imag.copy_from_host(ptr_B_imag_host.data()); + + device_tensors_C_real.reset(L); + device_tensors_C_real.copy_from_host(ptr_C_real_host.data()); + device_tensors_C_imag.reset(L); + device_tensors_C_imag.copy_from_host(ptr_C_imag_host.data()); + + device_tensors_D_real.reset(L); + device_tensors_D_real.copy_from_host(ptr_D_real_host.data()); + device_tensors_D_imag.reset(L); + device_tensors_D_imag.copy_from_host(ptr_D_imag_host.data()); + + auto arguments = typename Gemm::Arguments { + cutlass::gemm::GemmUniversalMode::kArray, + problem_shapes, + { + device_tensors_A_real.get(), stride_a, device_tensors_A_imag.get(), stride_a, + device_tensors_B_real.get(), stride_b, device_tensors_B_imag.get(), stride_b + }, + { + {alpha, beta}, + device_tensors_C_real.get(), stride_c, device_tensors_C_imag.get(), stride_c, + device_tensors_D_real.get(), stride_d, device_tensors_D_imag.get(), stride_d + }, + hw_info, + scheduler_args + }; + + Gemm gemm_op; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + cutlass::Status status = gemm_op.can_implement(arguments); + + if (status != cutlass::Status::kSuccess) { + cudaError_t error = cudaGetLastError(); + std::cerr << "This test is not supported: " << cudaGetErrorString(error) << "\n"; + return true; + } + + // + // Run the GEMM + // + + cudaError_t result; + status = gemm_op.initialize(arguments, workspace.get()); + status = gemm_op.run(); + result = cudaDeviceSynchronize(); + if (result != cudaSuccess) { + EXPECT_EQ(result, cudaSuccess) << "Error at Kernel Sync."; + return false; + } + + EXPECT_TRUE(status == cutlass::Status::kSuccess) << to_string(status); + // + // Verify + // + bool passed = this->verify(problem_shapes, alpha, beta); + + if (!passed) { + std::cout << "Error : Failed : with alpha: " << alpha << ", beta: " << beta + << "\n"; + } + + return passed; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace test diff --git a/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32.cu b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32.cu new file mode 100644 index 00000000..1af4d4d7 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32.cu @@ -0,0 +1,163 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x128x64 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TTT +TEST(SM100_Device_Gemm_Planar_cbf16t_cbf16t_f32t_tensorop_1sm, 128x128x64_1x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TMA MULTICAST /////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TNN +TEST(SM100_Device_Gemm_Planar_cbf16t_cbf16n_f32n_tensorop_1sm, 512x256x64_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_2sm.cu b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_2sm.cu new file mode 100644 index 00000000..02b20856 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_2sm.cu @@ -0,0 +1,163 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 128x64x64 /////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NNN +TEST(SM100_Device_Gemm_Planar_cbf16n_cbf16n_f32n_tensorop_2sm, 128x64x64_2x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_2,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TMA MULTICAST ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NTT +TEST(SM100_Device_Gemm_Planar_cbf16n_cbf16t_f32t_tensorop_2sm, 256x256x64_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_conjugate_layout.cu b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_conjugate_layout.cu new file mode 100644 index 00000000..d281ee52 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_conjugate_layout.cu @@ -0,0 +1,211 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TCGEN05.1CTA /////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// CN +TEST(SM100_Device_Gemm_Planar_cbf16c_cbf16n_f32n_tensorop_1sm, 128x64x64_1x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::conjugate; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +// NC +TEST(SM100_Device_Gemm_Planar_cbf16n_cbf16c_f32n_tensorop_1sm, 128x128x64_1x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::conjugate; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TCGEN05.2CTA //////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TC +TEST(SM100_Device_Gemm_Planar_cbf16t_cbf16c_f32t_tensorop_2sm, 128x64x64_2x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::conjugate; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_2,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif diff --git a/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_preferred_cluster.cu b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_preferred_cluster.cu new file mode 100644 index 00000000..1b0ff81d --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_preferred_cluster.cu @@ -0,0 +1,165 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////// 64x128x64_0x0x1 TCGEN05.1CTA /////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +// TTT +TEST(SM100_Device_Gemm_Planar_cbf16t_cbf16t_f32t_tensorop_1sm_preferred_cluster, 64x128x64_0x0x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_64,_128,_64>; + using ClusterShape = cute::Shape; // Runtime cluster shape + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexPreferredClusterSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////// 128x128x64_0x0x1 TCGEN05.2CTA ////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +// NNN +TEST(SM100_Device_Gemm_Planar_cbf16n_cbf16n_f32n_tensorop_2sm_preferred_cluster, 128x128x64_0x0x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = cute::Shape; // Runtime cluster shape + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + // 2Sm preferred cluster kernel should directly use 2Sm KernelSchedule for now + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexPreferredClusterSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_ptr_array.cu b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_ptr_array.cu new file mode 100644 index 00000000..341d1652 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cbf16_cbf16_cbf16_tensor_op_cf32_ptr_array.cu @@ -0,0 +1,282 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide Ptr-Array GEMM interface +*/ + + + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../common/cutlass_unit_test.h" + +#include "gemm_testbed_3x_ptr_array_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 128x64x64, CLUSTER_SIZE : 1x1x1, CLUSTER_TILE_SIZE : 128x64x64, TCGEN05.1CTA ////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TT +TEST(SM100_Device_Gemm_Planar_cbf16t_cbf16t_f32n_tensorop_ptr_array_1sm, 128x64x64_1x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized1Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 64x128x64, CLUSTER_SIZE : 4x4x1, CLUSTER_TILE_SIZE : 256x256x64, TCGEN05.1CTA ///////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NN +TEST(SM100_Device_Gemm_Planar_cbf16n_cbf16n_f32n_tensorop_ptr_array_1sm, 256x512x64_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_64,_128,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized1Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 64x128x64, CLUSTER_SIZE : 2x1x1, CLUSTER_TILE_SIZE : 128x128x64, TCGEN05.1CTA ///////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TN +TEST(SM100_Device_Gemm_Planar_cbf16t_cbf16n_f32n_tensorop_ptr_array_2sm, 128x128x64_2x1x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_2,_1,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized2Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 128x64x64, CLUSTER_SIZE : 4x4x1, CLUSTER_TILE_SIZE : 512x256x64, TCGEN05.1CTA ///////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NT +TEST(SM100_Device_Gemm_Planar_cbf16n_cbf16t_f32n_tensorop_ptr_array_2sm, 512x256x64_4x4x1) { + using ElementA = cutlass::bfloat16_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::bfloat16_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_256,_64,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized2Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::bfloat16_t, LayoutC, 8, + cutlass::bfloat16_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32.cu b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32.cu new file mode 100644 index 00000000..099a105b --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32.cu @@ -0,0 +1,163 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////// 64x64x64 /////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NNN +TEST(SM100_Device_Gemm_Planar_cf16n_cf16n_f32n_tensorop_1sm, 64x64x64_1x1x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_64,_64,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TMA MULTICAST /////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NTT +TEST(SM100_Device_Gemm_Planar_cf16n_cf16t_f32t_tensorop_1sm, 512x512x64_4x4x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_2sm.cu b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_2sm.cu new file mode 100644 index 00000000..7f77e456 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_2sm.cu @@ -0,0 +1,163 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// 256x64x64 ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TTT +TEST(SM100_Device_Gemm_Planar_cf16t_cf16t_f32t_tensorop_2sm, 256x64x64_2x1x1) { + using ElementA = cutlass::half_t; + using LayoutA = cutlass::layout::RowMajor; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + + using ElementB = cutlass::half_t; + using LayoutB = cutlass::layout::RowMajor; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_256,_64,_64>; + using ClusterShape = Shape<_2,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TMA MULTICAST ////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TNN +TEST(SM100_Device_Gemm_Planar_cf16t_cf16n_f32n_tensorop_2sm, 256x512x64_4x4x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_conjugate_layout.cu b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_conjugate_layout.cu new file mode 100644 index 00000000..a595cfd7 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_conjugate_layout.cu @@ -0,0 +1,212 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TCGEN05.1CTA //////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NC +TEST(SM100_Device_Gemm_Planar_cf16n_cf16c_f32n_tensorop_1sm, 128x128x64_1x1x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::conjugate; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +// CT +TEST(SM100_Device_Gemm_Planar_cf16c_cf16t_f32n_tensorop_1sm, 64x64x64_1x1x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::conjugate; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_64,_64,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////// TCGEN05.2CTA //////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// CC +TEST(SM100_Device_Gemm_Planar_cf16c_cf16c_f32t_tensorop_2sm, 256x64x64_2x1x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::conjugate; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::conjugate; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_256,_64,_64>; + using ClusterShape = Shape<_2,_1,_1>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + cutlass::gemm::StreamKScheduler + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif diff --git a/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_preferred_cluster.cu b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_preferred_cluster.cu new file mode 100644 index 00000000..6e8542d4 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_preferred_cluster.cu @@ -0,0 +1,165 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide GEMM interface +*/ + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/numeric_types.h" +#include "cutlass/arch/mma_sm100.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cute/atom/mma_traits_sm100.hpp" +#include "../../common/cutlass_unit_test.h" +#include "gemm_testbed_3x_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////// 128x128x64_0x0x1 TCGEN05.1CTA ////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +// TNT +TEST(SM100_Device_Gemm_Planar_cf16t_cf16n_f32t_tensorop_1sm_preferred_cluster, 128x128x64_0x0x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::RowMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = cute::Shape; // Runtime cluster shape + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized1Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecialized1SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexPreferredClusterSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////// 256x128x64_0x0x1 MMA.2SM /////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// + +// NTN +TEST(SM100_Device_Gemm_Planar_cf16n_cf16t_f32n_tensorop_2sm_preferred_cluster, 256x128x64_0x0x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_256,_128,_64>; + using ClusterShape = cute::Shape; // Runtime cluster shape + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + cutlass::epilogue::PlanarComplexTmaWarpSpecialized2Sm + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + // 2Sm preferred cluster kernel should directly use 2Sm KernelSchedule for now + cutlass::gemm::KernelTmaWarpSpecialized2SmPlanarComplexSm100 + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexPreferredClusterSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_ptr_array.cu b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_ptr_array.cu new file mode 100644 index 00000000..10cbaf51 --- /dev/null +++ b/test/unit/gemm/device/sm100_gemm_planar_cf16_cf16_cf16_tensor_op_cf32_ptr_array.cu @@ -0,0 +1,282 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. 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 Tests for device-wide Ptr-Array GEMM interface +*/ + + + +#include + +#include "cutlass/cutlass.h" +#include "cute/tensor.hpp" +#include "cute/atom/mma_atom.hpp" + +#include "cutlass/numeric_types.h" + +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/gemm/kernel/tile_scheduler.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/sm70_epilogue_vectorized.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" + +#include "../../common/cutlass_unit_test.h" + +#include "gemm_testbed_3x_ptr_array_planar_complex.hpp" + +using namespace cute; + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +////////// CTA_TILE_SIZE : 128x128x64, CLUSTER_SIZE : 1x1x1, CLUSTER_TILE_SIZE : 128x128x64, TCGEN05.1CTA ///////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TN +TEST(SM100_Device_Gemm_Planar_cf16t_cf16n_f32n_tensorop_ptr_array_1sm, 128x128x64_1x1x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_128,_64>; + using ClusterShape = Shape<_1,_1,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized1Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 64x64x64, CLUSTER_SIZE : 2x2x1, CLUSTER_TILE_SIZE : 128x128x64, TCGEN05.1CTA ////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NT +TEST(SM100_Device_Gemm_Planar_cf16n_cf16t_f32n_tensorop_ptr_array_1sm, 128x128x64_2x2x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_64,_64,_64>; + using ClusterShape = Shape<_2,_2,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized1Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 64x64x64, CLUSTER_SIZE : 2x1x1, CLUSTER_TILE_SIZE : 128x64x64, TCGEN05.2CTA /////////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// TT +TEST(SM100_Device_Gemm_Planar_cf16t_cf16t_f32n_tensorop_ptr_array_2sm, 128x64x64_2x1x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::RowMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::RowMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_128,_64,_64>; + using ClusterShape = Shape<_2,_1,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized2Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +/////////// CTA_TILE_SIZE : 128x128x64, CLUSTER_SIZE : 4x4x1, CLUSTER_TILE_SIZE : 512x512x64, TCGEN05.2CTA //////////// +/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// +// NN +TEST(SM100_Device_Gemm_Planar_cf16n_cf16n_f32n_tensorop_ptr_array_2sm, 512x512x64_4x4x1) { + using ElementA = cutlass::half_t; + using TransformA = cute::identity; + using ElementPairA = cute::tuple; + using LayoutA = cutlass::layout::ColumnMajor; + + using ElementB = cutlass::half_t; + using TransformB = cute::identity; + using ElementPairB = cute::tuple; + using LayoutB = cutlass::layout::ColumnMajor; + + using ElementAccumulator = float; + using LayoutC = cutlass::layout::ColumnMajor; + + using MmaTileShape = cute::Shape<_256,_128,_64>; + using ClusterShape = Shape<_4,_4,_1>; + + using MainloopSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmPlanarComplexSm100; // Kernel to launch + using EpilogueSchedule = cutlass::epilogue::PtrArrayPlanarComplexTmaWarpSpecialized2Sm; // Epilogue to launch + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + float, float, + cutlass::half_t, LayoutC, 8, + cutlass::half_t, LayoutC, 8, + EpilogueSchedule + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + ElementPairA, LayoutA, 8, + ElementPairB, LayoutB, 8, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout< + static_cast(sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cutlass::gemm::ArrayProblemShape>, + CollectiveMainloop, + CollectiveEpilogue + >; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + EXPECT_TRUE(test::gemm::device::TestPlanarComplexSmall()); +} + +#endif // #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp b/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp index cf3bfdde..0955f6f7 100644 --- a/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp +++ b/test/unit/transform/device/testbed_sparse_gemm_compressor.hpp @@ -306,14 +306,14 @@ public: compressor_utility.structure_sparse_zero_mask_fill(datas.tensor_A.host_data(), seed + 6); - // Check for failed devide + // Check for failed device CUDA_CHECK_FALSE(cudaGetLastError()); datas.tensor_A.sync_device(); datas.tensor_A_Comp.sync_device(); datas.tensor_E.sync_device(); - // Check for failed devide + // Check for failed device CUDA_CHECK_FALSE(cudaGetLastError()); return true; @@ -376,18 +376,6 @@ public: datas.tensor_A_Comp.sync_host(); datas.tensor_E.sync_host(); - #if 0 - { - printf("\n--> DEVICE OUTPUT\n"); - printf("datas.tensor_A\n"); - std::cout << datas.tensor_A.host_view() << std::endl << std::endl; - printf("datas.tensor_A_Comp\n"); - std::cout << datas.tensor_A_Comp.host_view() << std::endl << std::endl; - printf("datas.tensor_E\n"); - std::cout << datas.tensor_E.host_view() << std::endl << std::endl; - } - #endif - return true; } @@ -816,9 +804,11 @@ public: printf("compare_reference() DEVICE <-> LEGACY HOST fail\n"); return false; } - // else { - // printf("DEVICE <-> HOST PASS\n"); - // } + #if 0 + else { + printf("DEVICE <-> HOST PASS\n"); + } + #endif return true; } diff --git a/tools/library/CMakeLists.txt b/tools/library/CMakeLists.txt index 6358f03a..39676ca4 100644 --- a/tools/library/CMakeLists.txt +++ b/tools/library/CMakeLists.txt @@ -123,6 +123,19 @@ function(cutlass_add_cutlass_library) PRIVATE cutlass_library_internal_interface ) + # Use medium code model to handle large libraries (>2GB) on Linux x86_64 + # Conservative approach: only apply on Linux with GCC/Clang/Intel compilers + # This addresses the "relocation truncated to fit" linker error when libraries exceed 2GB + if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND + CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64" AND + CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang|IntelLLVM") + message(STATUS "CUTLASS: Applying -mcmodel=medium for large library support (Linux x86_64, ${CMAKE_CXX_COMPILER_ID})") + target_compile_options(${__NAME}_objs PRIVATE + $<$:-mcmodel=medium> + $<$:-Xcompiler=-mcmodel=medium> + ) + endif() + if (CUTLASS_BUILD_MONO_LIBRARY AND __SUFFIX) # If we're only building a single monolithic library then we diff --git a/tools/library/src/reference/gemm_f6_f8_f32.cu b/tools/library/src/reference/gemm_f6_f8_f32.cu index 08bdc5e3..cefea554 100644 --- a/tools/library/src/reference/gemm_f6_f8_f32.cu +++ b/tools/library/src/reference/gemm_f6_f8_f32.cu @@ -56,6 +56,10 @@ namespace library { // 2. e3m2_e4m3_f32_f16_e5m2 // 3. e3m2_e4m3_f32_f16_f16 // 4. e3m2_e4m3_f32_f32_f32 +// 5. e3m2_e5m2_f32_f16_e4m3 +// 6. e3m2_e5m2_f32_f16_e5m2 +// 7. e3m2_e5m2_f32_f16_f16 +// 8. e3m2_e5m2_f32_f32_f32 void initialize_gemm_reference_operations_f6_f8_f32(Manifest &manifest) { @@ -138,6 +142,86 @@ void initialize_gemm_reference_operations_f6_f8_f32(Manifest &manifest) { float, // ElementAccumulator float // ElementD >(manifest); + + // 5. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + half_t, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 6. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + half_t, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 7. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + half_t, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 8. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + float, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); + + // 5. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e4m3_t // ElementD + >(manifest); + + // 6. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float_e5m2_t // ElementD + >(manifest); + + // 7. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + half_t // ElementD + >(manifest); + + // 8. + make_gemm_real_canonical_layouts< + float_e3m2_t, // ElementA + float_e5m2_t, // ElementB + void, // ElementC + float, // ElementScalar + float, // ElementAccumulator + float // ElementD + >(manifest); } /////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h b/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h index 366133b5..e9582f53 100644 --- a/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h +++ b/tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h @@ -318,8 +318,30 @@ protected: int swizzle_size, bool is_dynamic_cluster_enabled); - /// Verifies CUTLASS against host and device references - bool verify_with_reference_( + /// Verifies CUTLASS against device reference (for regular grouped GEMM) + bool verify_with_device_reference_( + Options const& options, + PerformanceReport& report, + DeviceContext& device_context, + library::Operation const* operation, + ProblemSpace const& problem_space, + ProblemSpace::Problem const& problem, + cutlass::library::NumericTypeID element_A, + cutlass::library::NumericTypeID element_B); + + /// Verifies CUTLASS against host reference (for regular grouped GEMM) + bool verify_regular_with_host_reference_( + Options const& options, + PerformanceReport& report, + DeviceContext& device_context, + library::Operation const* operation, + ProblemSpace const& problem_space, + ProblemSpace::Problem const& problem, + cutlass::library::NumericTypeID element_A, + cutlass::library::NumericTypeID element_B); + + /// Verifies CUTLASS against host reference (for block-scaled/blockwise grouped GEMM) + bool verify_block_with_host_reference_( Options const& options, PerformanceReport& report, DeviceContext& device_context, diff --git a/tools/profiler/src/block_scaled_gemm_operation_profiler.cu b/tools/profiler/src/block_scaled_gemm_operation_profiler.cu index dd3ee072..0fbe84e9 100644 --- a/tools/profiler/src/block_scaled_gemm_operation_profiler.cu +++ b/tools/profiler/src/block_scaled_gemm_operation_profiler.cu @@ -1137,37 +1137,31 @@ bool BlockScaledGemmOperationProfiler::verify_cutlass( } - bool verification_status = verify_with_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B); - - // Update disposition to worst case verification outcome among all - // verification providers which are supported - bool is_any_verification_run_passed = false; - for (auto &m : results_.back().verification_map) { - if (m.second == Disposition::kFailed || m.second == Disposition::kIncorrect) { - results_.back().disposition = m.second; - return true; - } - if (!is_any_verification_run_passed && m.second == Disposition::kPassed) { - is_any_verification_run_passed = true; - } + if (!verify_with_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B)) { + return false; } - if (is_any_verification_run_passed) { - results_.back().disposition = Disposition::kPassed; + // Block-scaled GEMM only supports host verification - check it directly + switch (results_.back().verification_map[library::Provider::kReferenceHost]) { + case Disposition::kFailed: + results_.back().disposition = Disposition::kFailed; + return true; + case Disposition::kIncorrect: + results_.back().disposition = Disposition::kIncorrect; + return true; + case Disposition::kPassed: + results_.back().disposition = Disposition::kPassed; + return true; + default: + break; } } - // if verification.required is set, then return success iff at least one ref-check was run - if (options.verification.required) { - bool did_any_verification_run = false; - for (auto provider : options.verification.providers) { - did_any_verification_run |= (Disposition::kNotRun != results_.back().verification_map[provider]); - } - - if (not did_any_verification_run) { - results_.back().status = Status::kErrorNotSupported; - return false; - } + // if verification.required is set, check if host verification ran (the only supported provider) + if (options.verification.required && + results_.back().verification_map[library::Provider::kReferenceHost] == Disposition::kNotRun) { + results_.back().status = Status::kErrorNotSupported; + return false; } // Return true means continue profiling @@ -1664,3 +1658,4 @@ Status BlockScaledGemmOperationProfiler::profile_cutlass_( } // namespace cutlass ///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/tools/profiler/src/blockwise_gemm_operation_profiler.cu b/tools/profiler/src/blockwise_gemm_operation_profiler.cu index 7b52ad7d..910fe54c 100644 --- a/tools/profiler/src/blockwise_gemm_operation_profiler.cu +++ b/tools/profiler/src/blockwise_gemm_operation_profiler.cu @@ -988,37 +988,31 @@ bool BlockwiseGemmOperationProfiler::verify_cutlass( } - bool verification_status = verify_with_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B); - - // Update disposition to worst case verification outcome among all - // verification providers which are supported - bool is_any_verification_run_passed = false; - for (auto &m : results_.back().verification_map) { - if (m.second == Disposition::kFailed || m.second == Disposition::kIncorrect) { - results_.back().disposition = m.second; - return true; - } - if (!is_any_verification_run_passed && m.second == Disposition::kPassed) { - is_any_verification_run_passed = true; - } + if (!verify_with_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B)) { + return false; } - if (is_any_verification_run_passed) { - results_.back().disposition = Disposition::kPassed; + // Blockwise GEMM only supports host verification - check it directly + switch (results_.back().verification_map[library::Provider::kReferenceHost]) { + case Disposition::kFailed: + results_.back().disposition = Disposition::kFailed; + return true; + case Disposition::kIncorrect: + results_.back().disposition = Disposition::kIncorrect; + return true; + case Disposition::kPassed: + results_.back().disposition = Disposition::kPassed; + return true; + default: + break; } } - // if verification.required is set, then return success iff at least one ref-check was run - if (options.verification.required) { - bool did_any_verification_run = false; - for (auto provider : options.verification.providers) { - did_any_verification_run |= (Disposition::kNotRun != results_.back().verification_map[provider]); - } - - if (not did_any_verification_run) { - results_.back().status = Status::kErrorNotSupported; - return false; - } + // if verification.required is set, check if host verification ran (the only supported provider) + if (options.verification.required && + results_.back().verification_map[library::Provider::kReferenceHost] == Disposition::kNotRun) { + results_.back().status = Status::kErrorNotSupported; + return false; } // Return true means continue profiling @@ -1555,3 +1549,4 @@ Status BlockwiseGemmOperationProfiler::profile_cutlass_( } // namespace cutlass ///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/tools/profiler/src/grouped_gemm_operation_profiler.cu b/tools/profiler/src/grouped_gemm_operation_profiler.cu index f680bf4c..91beed03 100644 --- a/tools/profiler/src/grouped_gemm_operation_profiler.cu +++ b/tools/profiler/src/grouped_gemm_operation_profiler.cu @@ -109,9 +109,9 @@ GroupedGemmOperationProfiler::GroupedGemmOperationProfiler(Options const& option {"beta", "epilogue::beta"}, "Epilogue scalar beta (applied to all GEMMs in group)."}, {ArgumentTypeID::kEnumerated, {"runtime_input_datatype_a", "runtime-input-datatype::a"}, - "Runtime datatype (e4m3, e5m2, e3m2, e2m3, e2m1)"}, + "Runtime datatype (e4m3, e5m2, e3m2, e2m3, e2m1)"}, {ArgumentTypeID::kEnumerated, {"runtime_input_datatype_b", "runtime-input-datatype::b"}, - "Runtime datatype (e4m3, e5m2, e3m2, e2m3, e2m1)"}, + "Runtime datatype (e4m3, e5m2, e3m2, e2m3, e2m1)"}, {ArgumentTypeID::kEnumerated, {"raster_order", "raster-order"}, "Raster order (heuristic, along_n, along_m)"}, {ArgumentTypeID::kInteger, {"swizzle_size", "swizzle-size"}, "Size to swizzle"}, @@ -198,9 +198,9 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( int n0 = 32; int k0 = 64; for (int i = 0; i < num_groups; i++) { - auto m = m0 * (i + 1); + auto m = is_moe ? m0 * num_groups : m0 * (i + 1); auto n = n0 * (i + 1); - auto k = k0 * (i + 1); + auto k = is_moe ? k0 * num_groups : k0 * (i + 1); problem_sizes[i] = {m, n, k}; problem_sizes_3x[i] = {m, n, k}; } @@ -275,7 +275,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( max_problem_size_3x = {max_m, max_n, max_k}; } if (is_moe) { - for(size_t group_idx = 0; group_idx < problem_sizes.size(); group_idx++) { + for(size_t group_idx = 0; group_idx < problem_sizes.size(); group_idx++) { if (problem_sizes[group_idx].m() != max_problem_size_3x[0] || problem_sizes[group_idx].k() != max_problem_size_3x[2]) { std::cerr << "Problem size M:"<< problem_sizes[group_idx].m() << "K:" << problem_sizes[group_idx].k() << " for group " << group_idx << "should be equal to " @@ -336,7 +336,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse( // default value this->use_pdl = false; } - + if (!arg_as_RuntimeDatatype(this->runtime_input_datatype_a, "runtime_input_datatype_a", problem_space, problem)) { // default value this->runtime_input_datatype_a = cutlass::library::RuntimeDatatype::kStatic; @@ -547,7 +547,7 @@ void GroupedGemmOperationProfiler::GroupedGemmProblem::initialize_result( set_argument(result, "raster_order", problem_space, library::to_string(raster_order)); set_argument(result, "swizzle_size", problem_space, swizzle_size); set_argument(result, "use_pdl", problem_space, library::to_string(use_pdl)); - + set_argument(result, "runtime_input_datatype_a", problem_space, library::to_string(runtime_input_datatype_a)); set_argument(result, "runtime_input_datatype_b", problem_space, library::to_string(runtime_input_datatype_b)); @@ -614,8 +614,8 @@ Status GroupedGemmOperationProfiler::initialize_configuration( std::string sf_tuple = "\\d+x\\d+"; std::string datatypes_regex = "\\w?f\\d+|e\\dm\\d"; // bf16 | f16 | f32 | e4m3 | ... - std::string blockwise_regex_string = sf_tuple + "(" + datatypes_regex + ")x(" + - datatypes_regex + ")_" + sf_tuple + "(" + + std::string blockwise_regex_string = sf_tuple + "(" + datatypes_regex + ")x(" + + datatypes_regex + ")_" + sf_tuple + "(" + datatypes_regex + ")x(" + datatypes_regex + ")"; if (std::string(operation_desc.gemm.name).find("bstensor") != std::string::npos) { @@ -650,7 +650,7 @@ Status GroupedGemmOperationProfiler::initialize_configuration( gemm_workspace_.arguments.swizzle_size = problem_.swizzle_size; gemm_workspace_.arguments.raster_order = problem_.raster_order; - + gemm_workspace_.arguments.runtime_input_datatype_a = problem_.runtime_input_datatype_a; gemm_workspace_.arguments.runtime_input_datatype_b = problem_.runtime_input_datatype_b; @@ -782,7 +782,7 @@ Status GroupedGemmOperationProfiler::initialize_workspace( gemm_workspace_.reference_ptr_array_host.resize(num_groups); int seed_shift = 0; - if (not operation_desc.is_moe) { + if (not operation_desc.is_moe) { for (size_t group_idx = 0; group_idx < num_groups; group_idx++) { auto group_str = std::to_string(group_idx); gemm_workspace_.A_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor( @@ -1410,29 +1410,30 @@ bool GroupedGemmOperationProfiler::verify_cutlass( init_arguments(options); library::Operation const* underlying_operation = operation; - results_.back().status = underlying_operation->initialize_with_arguments(&gemm_workspace_.arguments); - if (results_.back().status != Status::kSuccess) { + auto& result = results_.back(); + + result.status = underlying_operation->initialize_with_arguments(&gemm_workspace_.arguments); + if (result.status != Status::kSuccess) { return false; } - results_.back().status = underlying_operation->run( + result.status = underlying_operation->run( &gemm_workspace_.arguments, gemm_workspace_.host_workspace.data(), gemm_workspace_.device_workspace.data()); - if (results_.back().status != Status::kSuccess) { - results_.back().disposition = Disposition::kFailed; + if (result.status != Status::kSuccess) { + result.disposition = Disposition::kFailed; return false; } - cudaError_t result = cudaDeviceSynchronize(); - if (result != cudaSuccess) { - results_.back().disposition = Disposition::kFailed; + if (cudaDeviceSynchronize() != cudaSuccess) { + result.disposition = Disposition::kFailed; return false; } // CUTLASS op ran the but not yet verified against any verification provider - results_.back().disposition = Disposition::kNotVerified; + result.disposition = Disposition::kNotVerified; // // Run verification providers @@ -1443,7 +1444,7 @@ bool GroupedGemmOperationProfiler::verify_cutlass( #if CUTLASS_ENABLE_CUBLAS if (options.verification.provider_enabled(library::Provider::kCUBLAS)) { // set verification map for cublas to not supported - results_.back().verification_map[library::Provider::kCUBLAS] = Disposition::kNotSupported; + result.verification_map[library::Provider::kCUBLAS] = Disposition::kNotSupported; } #endif // #if CUTLASS_ENABLE_CUBLAS @@ -1457,10 +1458,10 @@ bool GroupedGemmOperationProfiler::verify_cutlass( bool is_runtime_datatype_b = runtime_datatype_b != cutlass::library::RuntimeDatatype::kStatic; assert(is_runtime_datatype_a == is_runtime_datatype_b && "runtime datatype should be both dynamic or static."); - + cutlass::library::NumericTypeID element_A = desc.gemm.A.element; cutlass::library::NumericTypeID element_B = desc.gemm.B.element; - + if (is_runtime_datatype_a) { element_A = cutlass::library::dynamic_datatype_to_id(runtime_datatype_a); } @@ -1469,44 +1470,85 @@ bool GroupedGemmOperationProfiler::verify_cutlass( element_B = cutlass::library::dynamic_datatype_to_id(runtime_datatype_b); } - bool verification_status = verify_with_reference_( - options, - report, - device_context, - operation, - problem_space, - problem, - element_A, - element_B); - - // Update disposition to worst case verification outcome among all - // verification providers which are supported - bool is_any_verification_run_passed = false; - for (auto& m : results_.back().verification_map) { - if (m.second == Disposition::kFailed || m.second == Disposition::kIncorrect) { - results_.back().disposition = m.second; - return true; + // Block-scaled and blockwise GEMMs only support host verification - always terminates early + if (is_block_scaled || is_blockwise) { + if (!verify_block_with_host_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B)) { + return false; } - if (!is_any_verification_run_passed && m.second == Disposition::kPassed) { - is_any_verification_run_passed = true; + + switch (result.verification_map[library::Provider::kReferenceHost]) { + case Disposition::kFailed: + result.disposition = Disposition::kFailed; + return true; + case Disposition::kIncorrect: + result.disposition = Disposition::kIncorrect; + return true; + case Disposition::kPassed: + result.disposition = Disposition::kPassed; + return true; + case Disposition::kNotSupported: + result.disposition = Disposition::kNotSupported; + return true; + default: + // verification.required check for block-scaled/blockwise (for kNotRun, kNotVerified, etc.) + if (options.verification.required) { + result.status = Status::kErrorNotSupported; + return false; + } + return true; } } - if (is_any_verification_run_passed) { - results_.back().disposition = Disposition::kPassed; + // Regular grouped GEMM verification flow + + // Step 1: Try device reference if enabled (also run if cuBLAS requested since it's not supported, or as fallback if host not enabled) + if (options.verification.provider_enabled(library::Provider::kReferenceDevice) || + options.verification.provider_enabled(library::Provider::kCUBLAS) || + !options.verification.provider_enabled(library::Provider::kReferenceHost)) { + verify_with_device_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B); + + auto device_it = result.verification_map.find(library::Provider::kReferenceDevice); + if (device_it != result.verification_map.end()) { + if (device_it->second == Disposition::kFailed || device_it->second == Disposition::kIncorrect) { + result.disposition = device_it->second; + return true; + } + if (device_it->second == Disposition::kPassed) { + result.disposition = Disposition::kPassed; + return true; + } } } - // if verification.required is set, then return success iff at least one ref-check was run - if (options.verification.required) { - bool did_any_verification_run = false; - for (auto provider : options.verification.providers) { - did_any_verification_run |= - (Disposition::kNotRun != results_.back().verification_map[provider]); - } + // Step 2: Try host reference if enabled + if (options.verification.provider_enabled(library::Provider::kReferenceHost)) { + verify_regular_with_host_reference_(options, report, device_context, operation, problem_space, problem, element_A, element_B); - if (not did_any_verification_run) { - results_.back().status = Status::kErrorNotSupported; + auto host_it = result.verification_map.find(library::Provider::kReferenceHost); + if (host_it != result.verification_map.end()) { + if (host_it->second == Disposition::kFailed || host_it->second == Disposition::kIncorrect) { + result.disposition = host_it->second; + return true; + } + if (host_it->second == Disposition::kPassed) { + result.disposition = Disposition::kPassed; + return true; + } + } + } + } + + // Step 3: Check verification.required for regular grouped GEMM + if (options.verification.required) { + auto device_it = result.verification_map.find(library::Provider::kReferenceDevice); + auto host_it = result.verification_map.find(library::Provider::kReferenceHost); + + bool did_any_verification_run = + (device_it != result.verification_map.end() && device_it->second != Disposition::kNotRun) || + (host_it != result.verification_map.end() && host_it->second != Disposition::kNotRun); + + if (!did_any_verification_run) { + result.status = Status::kErrorNotSupported; return false; } } @@ -1515,8 +1557,8 @@ bool GroupedGemmOperationProfiler::verify_cutlass( return true; } -/// Verifies CUTLASS against host and device references -bool GroupedGemmOperationProfiler::verify_with_reference_( +/// Verifies CUTLASS against device reference (for regular grouped GEMM) +bool GroupedGemmOperationProfiler::verify_with_device_reference_( Options const& options, PerformanceReport& report, DeviceContext& device_context, @@ -1528,98 +1570,20 @@ bool GroupedGemmOperationProfiler::verify_with_reference_( library::GroupedGemmDescription const& desc = static_cast(operation->description()); - for (auto provider : options.verification.providers) { - - // Skip providers that are not enabled - if (!options.verification.provider_enabled(provider)) { - continue; - } - - // we only have a block scaled reference kernel implemented on the host - if ((is_block_scaled || is_blockwise) && provider != library::Provider::kReferenceHost) { - continue; - } - auto status = Status::kSuccess; auto disposition = Disposition::kFailed; - // we don't have grouped GEMM reference kernels so we loop over the groups and perform - // a regular GEMM for each group + + // Loop over groups and perform a regular GEMM for each group using device reference for (size_t group_idx = 0, num_groups = problem_.problem_sizes.size(); group_idx < num_groups; group_idx++) { + // Use device pointers directly void* ptr_A = gemm_workspace_.A_ptr_array_host[group_idx]->data(); void* ptr_B = gemm_workspace_.B_ptr_array_host[group_idx]->data(); void* ptr_C = gemm_workspace_.C_ptr_array_host[group_idx]->data(); void* ptr_D = gemm_workspace_.reference_ptr_array_host[group_idx]->data(); - // To support the host-side reference, conditionally allocate and - // copy tensors to host memory. - std::vector host_data_A; - std::vector host_data_B; - std::vector host_data_C; - std::vector host_data_D; - std::vector host_data_SFA; - std::vector host_data_SFB; - std::vector host_data_SFC; - std::vector host_data_SFD; - std::vector host_data_norm_constant; - - void* ptr_SFA{nullptr}; - void* ptr_SFB{nullptr}; - void* ptr_SFD{nullptr}; - void* ptr_norm_constant{nullptr}; - - if (provider == library::Provider::kReferenceHost) { - host_data_A.resize(gemm_workspace_.A_ptr_array_host[group_idx]->bytes()); - ptr_A = host_data_A.data(); - gemm_workspace_.A_ptr_array_host[group_idx]->copy_to_host( - ptr_A); // this is copying all the data for L2 busting as well - - host_data_B.resize(gemm_workspace_.B_ptr_array_host[group_idx]->bytes()); - ptr_B = host_data_B.data(); - gemm_workspace_.B_ptr_array_host[group_idx]->copy_to_host(ptr_B); - - host_data_C.resize(gemm_workspace_.C_ptr_array_host[group_idx]->bytes()); - ptr_C = host_data_C.data(); - gemm_workspace_.C_ptr_array_host[group_idx]->copy_to_host(ptr_C); - - host_data_D.resize(gemm_workspace_.reference_ptr_array_host[group_idx]->bytes()); - ptr_D = host_data_D.data(); - - if (is_block_scaled) { - auto const& ws = gemm_workspace_.block_scales.value(); - - host_data_SFA.resize(ws.SFA_ptr_array_host[group_idx]->bytes()); - ptr_SFA = host_data_SFA.data(); - ws.SFA_ptr_array_host[group_idx]->copy_to_host(ptr_SFA); - host_data_SFB.resize(ws.SFB_ptr_array_host[group_idx]->bytes()); - ptr_SFB = host_data_SFB.data(); - ws.SFB_ptr_array_host[group_idx]->copy_to_host(ptr_SFB); - - host_data_SFD.resize(ws.SFD_reference_ptr_array_host[group_idx]->bytes()); - ptr_SFD = host_data_SFD.data(); - - host_data_norm_constant.resize(ws.norm_constant->bytes()); - ptr_norm_constant = host_data_norm_constant.data(); - ws.norm_constant->copy_to_host(ptr_norm_constant); - } - else if (is_blockwise) { - auto const& ws = gemm_workspace_.block_scales.value(); - - host_data_SFA.resize(ws.SFA_ptr_array_host[group_idx]->bytes()); - ptr_SFA = host_data_SFA.data(); - ws.SFA_ptr_array_host[group_idx]->copy_to_host(ptr_SFA); - host_data_SFB.resize(ws.SFB_ptr_array_host[group_idx]->bytes()); - ptr_SFB = host_data_SFB.data(); - ws.SFB_ptr_array_host[group_idx]->copy_to_host(ptr_SFB); - } - } - - const auto &desc = static_cast(operation->description()); - const auto& gemm_desc = desc.gemm; - - if (!is_block_scaled and !is_blockwise) { library::Handle handle; - handle.set_provider(provider); + handle.set_provider(library::Provider::kReferenceDevice); status = handle.gemm_universal( library::GemmUniversalMode::kGemm, @@ -1659,241 +1623,462 @@ bool GroupedGemmOperationProfiler::verify_with_reference_( gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(), gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(), gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride()); - } - else if (is_block_scaled) { - auto const& block_scale_desc = desc.block_scales.value(); - auto& block_scale_ws = gemm_workspace_.block_scales.value(); - library::BlockScaledGemmFunctionalKey blockScaledGemm_key( - library::Provider::kReferenceHost, - library::GemmKind::kUniversal, - library::OperationKind::kBlockScaledGemm, - gemm_desc.tile_description.math_instruction.element_accumulator, - gemm_desc.element_epilogue, - element_A, - gemm_desc.A.layout, - block_scale_desc.SFA.element, - element_B, - gemm_desc.B.layout, - block_scale_desc.SFB.element, - gemm_desc.C.element, - gemm_desc.C.layout, - gemm_desc.D.element, - gemm_desc.D.layout, - block_scale_desc.SFD.element, - block_scale_desc.SFD.layout, - block_scale_desc.SFKVecSize, - block_scale_desc.EpilogueSFVecSize); - - auto operators_it = - library::Singleton::get().operation_table.block_scaled_gemm_operations.find( - blockScaledGemm_key); - if ( - operators_it == - library::Singleton::get().operation_table.block_scaled_gemm_operations.end()) { - disposition = Disposition::kNotSupported; + if (status != Status::kSuccess) { break; } - if (operators_it->second.empty()) { - disposition = Disposition::kNotSupported; + disposition = compare_tensors( + options, + *gemm_workspace_.D_ptr_array_host[group_idx], + *gemm_workspace_.reference_ptr_array_host[group_idx], + gemm_workspace_.D_ptr_array_host[group_idx]->batch_stride()); + + if (disposition != Disposition::kPassed) { break; } + } - auto cc_it = operators_it->second.begin(); - if (cc_it == operators_it->second.end()) { - disposition = Disposition::kNotSupported; - break; - } - - // host reference has only one instances in BlockScaledOperationVectorMap - library::Operation const* reference_op = cc_it->second[0]; - library::BlockScaledGemmArguments arguments{ - {int(problem_.m(group_idx)), int(problem_.n(group_idx)), int(problem_.k(group_idx))}, - {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}, - {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}, - 1, // batch count - ptr_A, - ptr_B, - ptr_SFA, - ptr_SFB, - ptr_C, - ptr_D, - ptr_SFD, - problem_.alpha.data(), - problem_.beta.data(), - library::ScalarPointerMode::kHost, - problem_.lda[group_idx], - problem_.ldb[group_idx], - problem_.ldc[group_idx], - problem_.ldc[group_idx], - gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(), - gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(), - gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(), - gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride(), - ptr_norm_constant}; - - library::GemmUniversalConfiguration configuration{ - library::GemmUniversalMode::kGemm, - problem_.problem_sizes[group_idx], - {problem_.cluster_m, problem_.cluster_n, problem_.cluster_k}, - {problem_.cluster_m_fallback, problem_.cluster_n_fallback, problem_.cluster_k_fallback}, - 1, - problem_.lda[group_idx], - problem_.ldb[group_idx], - problem_.ldc[group_idx], - problem_.ldc[group_idx], - 1, - }; - uint64_t host_workspace_size_needed = reference_op->get_host_workspace_size(&gemm_workspace_.configuration); - std::vector host_workspace(host_workspace_size_needed); - status = reference_op->initialize(&configuration, host_workspace.data()); if (status != Status::kSuccess) { - break; - } - - status = reference_op->run(&arguments, host_workspace.data()); - - block_scale_ws.SFD_reference_ptr_array_host[group_idx]->copy_from_host(ptr_SFD); + results_.back().verification_map[library::Provider::kReferenceDevice] = Disposition::kNotVerified; } else { - // Blockwise - auto const& block_scale_desc = desc.block_scales.value(); - auto& block_scale_ws = gemm_workspace_.block_scales.value(); + results_.back().status = status; + results_.back().verification_map[library::Provider::kReferenceDevice] = disposition; + } - library::BlockwiseGemmFunctionalKey blockwiseGemm_key( - library::Provider::kReferenceHost, - library::GemmKind::kUniversal, - library::OperationKind::kBlockwiseGemm, - gemm_desc.tile_description.math_instruction.element_accumulator, - gemm_desc.element_epilogue, - element_A, - gemm_desc.A.layout, - block_scale_desc.SFA.element, - element_B, - gemm_desc.B.layout, - block_scale_desc.SFB.element, - gemm_desc.C.element, - gemm_desc.C.layout, - gemm_desc.D.element, - gemm_desc.D.layout, - block_scale_desc.SFMVecSize, - block_scale_desc.SFNVecSize, - block_scale_desc.SFKVecSize - ); + if (options.verification.save_workspace == SaveWorkspace::kIncorrect && + results_.back().verification_map[library::Provider::kReferenceDevice] == Disposition::kIncorrect) { + save_workspace(device_context, options, desc, library::Provider::kCUTLASS, library::Provider::kReferenceDevice); + } - auto operators_it = library::Singleton::get().operation_table.blockwise_gemm_operations.find(blockwiseGemm_key); - if ( - operators_it == - library::Singleton::get().operation_table.blockwise_gemm_operations.end()) { - disposition = Disposition::kNotSupported; - break; - } + return true; // continue profiling +} - if (operators_it->second.empty()) { - disposition = Disposition::kNotSupported; - break; - } +///////////////////////////////////////////////////////////////////////////////////////////////// - auto cc_it = operators_it->second.begin(); - if (cc_it == operators_it->second.end()) { - disposition = Disposition::kNotSupported; - break; - } +/// Verifies CUTLASS against host reference (for regular grouped GEMM) +bool GroupedGemmOperationProfiler::verify_regular_with_host_reference_( + Options const& options, + PerformanceReport& report, + DeviceContext& device_context, + library::Operation const* operation, + ProblemSpace const& problem_space, + ProblemSpace::Problem const& problem, + cutlass::library::NumericTypeID element_A, + cutlass::library::NumericTypeID element_B) { + library::GroupedGemmDescription const& desc = + static_cast(operation->description()); - // host reference has only one instances in BlockScaledOperationVectorMap - library::Operation const* reference_op = cc_it->second[0]; + auto status = Status::kSuccess; + auto disposition = Disposition::kFailed; - library::BlockwiseGemmArguments arguments { - {int(problem_.m(group_idx)), int(problem_.n(group_idx)), int(problem_.k(group_idx))}, - {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}, - {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}, - 1, // batch_count + // Loop over groups and perform a regular GEMM for each group using host reference + for (size_t group_idx = 0, num_groups = problem_.problem_sizes.size(); group_idx < num_groups; + group_idx++) { + // Copy data to host memory + std::vector host_data_A(gemm_workspace_.A_ptr_array_host[group_idx]->bytes()); + std::vector host_data_B(gemm_workspace_.B_ptr_array_host[group_idx]->bytes()); + std::vector host_data_C(gemm_workspace_.C_ptr_array_host[group_idx]->bytes()); + std::vector host_data_D(gemm_workspace_.reference_ptr_array_host[group_idx]->bytes()); + + void* ptr_A = host_data_A.data(); + void* ptr_B = host_data_B.data(); + void* ptr_C = host_data_C.data(); + void* ptr_D = host_data_D.data(); + + gemm_workspace_.A_ptr_array_host[group_idx]->copy_to_host(ptr_A); + gemm_workspace_.B_ptr_array_host[group_idx]->copy_to_host(ptr_B); + gemm_workspace_.C_ptr_array_host[group_idx]->copy_to_host(ptr_C); + + library::Handle handle; + handle.set_provider(library::Provider::kReferenceHost); + + status = handle.gemm_universal( + library::GemmUniversalMode::kGemm, + problem_.m(group_idx), + problem_.n(group_idx), + problem_.k(group_idx), + problem_.cluster_m, + problem_.cluster_n, + problem_.cluster_k, + problem_.cluster_m_fallback, + problem_.cluster_n_fallback, + problem_.cluster_k_fallback, + desc.gemm.tile_description.math_instruction.element_accumulator, + desc.gemm.element_epilogue, + problem_.alpha.data(), + element_A, + desc.gemm.A.layout, + desc.gemm.transform_A, ptr_A, + int(problem_.lda[group_idx]), + element_B, + desc.gemm.B.layout, + desc.gemm.transform_B, ptr_B, - ptr_SFA, - ptr_SFB, + int(problem_.ldb[group_idx]), + problem_.beta.data(), + desc.gemm.C.element, + desc.gemm.C.layout, ptr_C, + int(problem_.ldc[group_idx]), + desc.gemm.D.element, + desc.gemm.D.layout, ptr_D, - problem_.alpha.data(), - problem_.beta.data(), - library::ScalarPointerMode::kHost, - problem_.lda[group_idx], - problem_.ldb[group_idx], - problem_.ldc[group_idx], - problem_.ldc[group_idx], + int(problem_.ldc[group_idx]), + 1, gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(), gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(), gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(), - gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride(), - }; + gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride()); - library::GemmUniversalConfiguration configuration{ - library::GemmUniversalMode::kGemm, - problem_.problem_sizes[group_idx], - {problem_.cluster_m, problem_.cluster_n, problem_.cluster_k}, - {problem_.cluster_m_fallback, problem_.cluster_n_fallback, problem_.cluster_k_fallback}, - 1, - problem_.lda[group_idx], - problem_.ldb[group_idx], - problem_.ldc[group_idx], - problem_.ldc[group_idx], - 1, - }; - uint64_t host_workspace_size_needed = reference_op->get_host_workspace_size(&gemm_workspace_.configuration); - std::vector host_workspace(host_workspace_size_needed); - status = reference_op->initialize(&configuration, host_workspace.data()); if (status != Status::kSuccess) { break; } - status = reference_op->run(&arguments, host_workspace.data()); - } - - if (status != Status::kSuccess) { - break; - } - - if (provider == library::Provider::kReferenceHost) { + // Copy result back to device gemm_workspace_.reference_ptr_array_host[group_idx]->copy_from_host(ptr_D); - } disposition = compare_tensors( options, *gemm_workspace_.D_ptr_array_host[group_idx], *gemm_workspace_.reference_ptr_array_host[group_idx], gemm_workspace_.D_ptr_array_host[group_idx]->batch_stride()); - if (disposition != Disposition::kPassed) { - break; - } - if (is_block_scaled) { - auto& ws = gemm_workspace_.block_scales.value(); - auto const& block_scale_desc = desc.block_scales.value(); - if (block_scale_desc.SFD.element != library::NumericTypeID::kVoid) { - disposition = compare_tensors( - options, - *ws.SFD_ptr_array_host[group_idx], - *ws.SFD_reference_ptr_array_host[group_idx], - ws.SFD_ptr_array_host[group_idx]->batch_stride()); if (disposition != Disposition::kPassed) { break; } } + + if (status != Status::kSuccess) { + results_.back().verification_map[library::Provider::kReferenceHost] = Disposition::kNotVerified; + } + else { + results_.back().status = status; + results_.back().verification_map[library::Provider::kReferenceHost] = disposition; + } + + if (options.verification.save_workspace == SaveWorkspace::kIncorrect && + results_.back().verification_map[library::Provider::kReferenceHost] == Disposition::kIncorrect) { + save_workspace(device_context, options, desc, library::Provider::kCUTLASS, library::Provider::kReferenceHost); + } + + return true; // continue profiling +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Verifies CUTLASS against host reference (for block-scaled/blockwise grouped GEMM) +bool GroupedGemmOperationProfiler::verify_block_with_host_reference_( + Options const& options, + PerformanceReport& report, + DeviceContext& device_context, + library::Operation const* operation, + ProblemSpace const& problem_space, + ProblemSpace::Problem const& problem, + cutlass::library::NumericTypeID element_A, + cutlass::library::NumericTypeID element_B) { + library::GroupedGemmDescription const& desc = + static_cast(operation->description()); + + auto status = Status::kSuccess; + auto disposition = Disposition::kFailed; + + // we don't have grouped GEMM reference kernels so we loop over the groups and perform + // a regular GEMM for each group + for (size_t group_idx = 0, num_groups = problem_.problem_sizes.size(); group_idx < num_groups; + group_idx++) { + // To support the host-side reference, allocate and copy tensors to host memory. + std::vector host_data_A; + std::vector host_data_B; + std::vector host_data_C; + std::vector host_data_D; + std::vector host_data_SFA; + std::vector host_data_SFB; + std::vector host_data_SFD; + std::vector host_data_norm_constant; + + host_data_A.resize(gemm_workspace_.A_ptr_array_host[group_idx]->bytes()); + void* ptr_A = host_data_A.data(); + gemm_workspace_.A_ptr_array_host[group_idx]->copy_to_host(ptr_A); + + host_data_B.resize(gemm_workspace_.B_ptr_array_host[group_idx]->bytes()); + void* ptr_B = host_data_B.data(); + gemm_workspace_.B_ptr_array_host[group_idx]->copy_to_host(ptr_B); + + host_data_C.resize(gemm_workspace_.C_ptr_array_host[group_idx]->bytes()); + void* ptr_C = host_data_C.data(); + gemm_workspace_.C_ptr_array_host[group_idx]->copy_to_host(ptr_C); + + host_data_D.resize(gemm_workspace_.reference_ptr_array_host[group_idx]->bytes()); + void* ptr_D = host_data_D.data(); + + void* ptr_SFA{nullptr}; + void* ptr_SFB{nullptr}; + void* ptr_SFD{nullptr}; + void* ptr_norm_constant{nullptr}; + + if (is_block_scaled) { + auto const& ws = gemm_workspace_.block_scales.value(); + + host_data_SFA.resize(ws.SFA_ptr_array_host[group_idx]->bytes()); + ptr_SFA = host_data_SFA.data(); + ws.SFA_ptr_array_host[group_idx]->copy_to_host(ptr_SFA); + host_data_SFB.resize(ws.SFB_ptr_array_host[group_idx]->bytes()); + ptr_SFB = host_data_SFB.data(); + ws.SFB_ptr_array_host[group_idx]->copy_to_host(ptr_SFB); + + host_data_SFD.resize(ws.SFD_reference_ptr_array_host[group_idx]->bytes()); + ptr_SFD = host_data_SFD.data(); + + host_data_norm_constant.resize(ws.norm_constant->bytes()); + ptr_norm_constant = host_data_norm_constant.data(); + ws.norm_constant->copy_to_host(ptr_norm_constant); + } + else if (is_blockwise) { + auto const& ws = gemm_workspace_.block_scales.value(); + + host_data_SFA.resize(ws.SFA_ptr_array_host[group_idx]->bytes()); + ptr_SFA = host_data_SFA.data(); + ws.SFA_ptr_array_host[group_idx]->copy_to_host(ptr_SFA); + host_data_SFB.resize(ws.SFB_ptr_array_host[group_idx]->bytes()); + ptr_SFB = host_data_SFB.data(); + ws.SFB_ptr_array_host[group_idx]->copy_to_host(ptr_SFB); + } + + const auto& gemm_desc = desc.gemm; + + if (is_block_scaled) { + auto const& block_scale_desc = desc.block_scales.value(); + auto& block_scale_ws = gemm_workspace_.block_scales.value(); + + library::BlockScaledGemmFunctionalKey blockScaledGemm_key( + library::Provider::kReferenceHost, + library::GemmKind::kUniversal, + library::OperationKind::kBlockScaledGemm, + gemm_desc.tile_description.math_instruction.element_accumulator, + gemm_desc.element_epilogue, + element_A, + gemm_desc.A.layout, + block_scale_desc.SFA.element, + element_B, + gemm_desc.B.layout, + block_scale_desc.SFB.element, + gemm_desc.C.element, + gemm_desc.C.layout, + gemm_desc.D.element, + gemm_desc.D.layout, + block_scale_desc.SFD.element, + block_scale_desc.SFD.layout, + block_scale_desc.SFKVecSize, + block_scale_desc.EpilogueSFVecSize); + + auto operators_it = + library::Singleton::get().operation_table.block_scaled_gemm_operations.find( + blockScaledGemm_key); + if (operators_it == + library::Singleton::get().operation_table.block_scaled_gemm_operations.end()) { + disposition = Disposition::kNotSupported; + break; + } + + if (operators_it->second.empty()) { + disposition = Disposition::kNotSupported; + break; + } + + auto cc_it = operators_it->second.begin(); + if (cc_it == operators_it->second.end()) { + disposition = Disposition::kNotSupported; + break; + } + + // host reference has only one instances in BlockScaledOperationVectorMap + library::Operation const* reference_op = cc_it->second[0]; + library::BlockScaledGemmArguments arguments{ + {int(problem_.m(group_idx)), int(problem_.n(group_idx)), int(problem_.k(group_idx))}, + {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}, + {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}, + 1, // batch count + ptr_A, + ptr_B, + ptr_SFA, + ptr_SFB, + ptr_C, + ptr_D, + ptr_SFD, + problem_.alpha.data(), + problem_.beta.data(), + library::ScalarPointerMode::kHost, + problem_.lda[group_idx], + problem_.ldb[group_idx], + problem_.ldc[group_idx], + problem_.ldc[group_idx], + gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(), + gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(), + gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(), + gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride(), + ptr_norm_constant}; + + library::GemmUniversalConfiguration configuration{ + library::GemmUniversalMode::kGemm, + problem_.problem_sizes[group_idx], + {problem_.cluster_m, problem_.cluster_n, problem_.cluster_k}, + {problem_.cluster_m_fallback, problem_.cluster_n_fallback, problem_.cluster_k_fallback}, + 1, + problem_.lda[group_idx], + problem_.ldb[group_idx], + problem_.ldc[group_idx], + problem_.ldc[group_idx], + 1, + }; + uint64_t host_workspace_size_needed = reference_op->get_host_workspace_size(&gemm_workspace_.configuration); + std::vector host_workspace(host_workspace_size_needed); + status = reference_op->initialize(&configuration, host_workspace.data()); + if (status != Status::kSuccess) { + break; + } + + status = reference_op->run(&arguments, host_workspace.data()); + + block_scale_ws.SFD_reference_ptr_array_host[group_idx]->copy_from_host(ptr_SFD); + } + else { + // Blockwise + auto const& block_scale_desc = desc.block_scales.value(); + + library::BlockwiseGemmFunctionalKey blockwiseGemm_key( + library::Provider::kReferenceHost, + library::GemmKind::kUniversal, + library::OperationKind::kBlockwiseGemm, + gemm_desc.tile_description.math_instruction.element_accumulator, + gemm_desc.element_epilogue, + element_A, + gemm_desc.A.layout, + block_scale_desc.SFA.element, + element_B, + gemm_desc.B.layout, + block_scale_desc.SFB.element, + gemm_desc.C.element, + gemm_desc.C.layout, + gemm_desc.D.element, + gemm_desc.D.layout, + block_scale_desc.SFMVecSize, + block_scale_desc.SFNVecSize, + block_scale_desc.SFKVecSize + ); + + auto operators_it = library::Singleton::get().operation_table.blockwise_gemm_operations.find(blockwiseGemm_key); + if (operators_it == + library::Singleton::get().operation_table.blockwise_gemm_operations.end()) { + disposition = Disposition::kNotSupported; + break; + } + + if (operators_it->second.empty()) { + disposition = Disposition::kNotSupported; + break; + } + + auto cc_it = operators_it->second.begin(); + if (cc_it == operators_it->second.end()) { + disposition = Disposition::kNotSupported; + break; + } + + // host reference has only one instances in BlockScaledOperationVectorMap + library::Operation const* reference_op = cc_it->second[0]; + + library::BlockwiseGemmArguments arguments { + {int(problem_.m(group_idx)), int(problem_.n(group_idx)), int(problem_.k(group_idx))}, + {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}, + {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}, + 1, // batch_count + ptr_A, + ptr_B, + ptr_SFA, + ptr_SFB, + ptr_C, + ptr_D, + problem_.alpha.data(), + problem_.beta.data(), + library::ScalarPointerMode::kHost, + problem_.lda[group_idx], + problem_.ldb[group_idx], + problem_.ldc[group_idx], + problem_.ldc[group_idx], + gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(), + gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(), + gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(), + gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride(), + }; + + library::GemmUniversalConfiguration configuration{ + library::GemmUniversalMode::kGemm, + problem_.problem_sizes[group_idx], + {problem_.cluster_m, problem_.cluster_n, problem_.cluster_k}, + {problem_.cluster_m_fallback, problem_.cluster_n_fallback, problem_.cluster_k_fallback}, + 1, + problem_.lda[group_idx], + problem_.ldb[group_idx], + problem_.ldc[group_idx], + problem_.ldc[group_idx], + 1, + }; + uint64_t host_workspace_size_needed = reference_op->get_host_workspace_size(&gemm_workspace_.configuration); + std::vector host_workspace(host_workspace_size_needed); + status = reference_op->initialize(&configuration, host_workspace.data()); + if (status != Status::kSuccess) { + break; + } + + status = reference_op->run(&arguments, host_workspace.data()); + } + + if (status != Status::kSuccess) { + break; + } + + gemm_workspace_.reference_ptr_array_host[group_idx]->copy_from_host(ptr_D); + + disposition = compare_tensors( + options, + *gemm_workspace_.D_ptr_array_host[group_idx], + *gemm_workspace_.reference_ptr_array_host[group_idx], + gemm_workspace_.D_ptr_array_host[group_idx]->batch_stride()); + if (disposition != Disposition::kPassed) { + break; + } + + if (is_block_scaled) { + auto& ws = gemm_workspace_.block_scales.value(); + auto const& block_scale_desc = desc.block_scales.value(); + if (block_scale_desc.SFD.element != library::NumericTypeID::kVoid) { + disposition = compare_tensors( + options, + *ws.SFD_ptr_array_host[group_idx], + *ws.SFD_reference_ptr_array_host[group_idx], + ws.SFD_ptr_array_host[group_idx]->batch_stride()); + if (disposition != Disposition::kPassed) { + break; + } } } - if (status != Status::kSuccess) { - results_.back().verification_map[provider] = Disposition::kNotVerified; - continue; - } - results_.back().status = status; - results_.back().verification_map[provider] = disposition; + } - if ( - options.verification.save_workspace == SaveWorkspace::kIncorrect && - results_.back().verification_map[provider] == Disposition::kIncorrect) { - save_workspace(device_context, options, desc, library::Provider::kCUTLASS, provider); - } + if (status != Status::kSuccess) { + results_.back().verification_map[library::Provider::kReferenceHost] = Disposition::kNotVerified; + return true; // continue profiling + } + + results_.back().status = status; + results_.back().verification_map[library::Provider::kReferenceHost] = disposition; + + if ( + options.verification.save_workspace == SaveWorkspace::kIncorrect && + disposition == Disposition::kIncorrect) { + save_workspace(device_context, options, desc, library::Provider::kCUTLASS, library::Provider::kReferenceHost); } return true; // continue profiling @@ -1912,7 +2097,7 @@ bool GroupedGemmOperationProfiler::profile( if (options.profiling.provider_enabled(library::Provider::kCUTLASS)) { if (options.profiling.enable_kernel_performance_search) { - std::cerr << "Exhaustive performance search is not available for Grouped GEMMs. " + std::cerr << "Exhaustive performance search is not available for Grouped GEMMs. " << "Please use --enable-best-kernel-for-fixed-shape to profile a specific problem size " << "with --problem-sizes or --problem-sizes-file.\n"; } @@ -1948,6 +2133,14 @@ Status GroupedGemmOperationProfiler::profile_cutlass_( return result.status; } + //For dynamic clusters, we update arguments, hence need to rerun can_implement() + result.status = underlying_operation->can_implement( + &gemm_workspace_.configuration, + &gemm_workspace_.arguments); + if (result.status != Status::kSuccess) { + return result.status; + } + auto func = [&](cudaStream_t stream, int iteration) { // Iterate over copies of the problem in memory int workspace_idx = options.profiling.warmup_iterations + iteration; diff --git a/tools/util/include/cutlass/util/mixed_dtype_utils.hpp b/tools/util/include/cutlass/util/mixed_dtype_utils.hpp index e92e3d90..0b271ec8 100644 --- a/tools/util/include/cutlass/util/mixed_dtype_utils.hpp +++ b/tools/util/include/cutlass/util/mixed_dtype_utils.hpp @@ -35,6 +35,7 @@ #pragma once #include +#include #include "cute/layout.hpp" #include "cute/tensor.hpp" #include "cute/arch/mma_sm90.hpp" @@ -175,21 +176,79 @@ static void dequantize(DequantizedElement* dq_buffer, CUDA_CHECK(cudaStreamSynchronize(stream)); } -template +template class packed_scale_t { public: - static_assert(cute::sizeof_bits_v == 8, - "only 8 bit arithmetic types are supported."); + static_assert( + cute::sizeof_bits_v == 8, + "ElementScale must be a supported 8-bit type."); + CUTLASS_HOST_DEVICE - explicit packed_scale_t(T val) { - if constexpr (!cute::is_unsigned_v) { - // Only pack negative values. The positive values are generated in flight in the mainloop. - storage[0] = pack4(T(float(val) * -8.f), T(float(val) * -7.f), T(float(val) * -6.f), T(float(val) * -5.f)); - storage[1] = pack4(T(float(val) * -4.f), T(float(val) * -3.f), T(float(val) * -2.f), -val); - } - else { - storage[0] = pack4(T(float(val) * 8.f), T(float(val) * 7.f), T(float(val) * 6.f), T(float(val) * 5.f)); - storage[1] = pack4(T(float(val) * 4.f), T(float(val) * 3.f), T(float(val) * 2.f), val); + explicit packed_scale_t(ElementScale val) { + + if constexpr (cutlass::platform::is_floating_point::value || + cute::is_same_v) { + // floating point QuantType needs to be converted to a signed type + // or a floating point ttype + static_assert( + !std::is_unsigned_v || + cutlass::platform::is_floating_point::value || + cute::is_same_v || + cute::is_same_v, + "E2M1 quantization requires ElementScale to be signed or FP8"); + + // E2M1 quantization: Use E2M1 LUT values + storage[0] = pack4( + ElementScale(float(val) * 0.0f), // E2M1: 0b000 => 0.0 + ElementScale(float(val) * (-0.5f)), // E2M1: 0b001 => -0.5 + ElementScale(float(val) * (-1.0f)), // E2M1: 0b010 => -1.0 + ElementScale(float(val) * (-1.5f)) // E2M1: 0b011 => -1.5 + ); + storage[1] = pack4( + ElementScale(float(val) * (-2.0f)), // E2M1: 0b100 => -2.0 + ElementScale(float(val) * (-3.0f)), // E2M1: 0b101 => -3.0 + ElementScale(float(val) * (-4.0f)), // E2M1: 0b110 => -4.0 + ElementScale(float(val) * (-6.0f)) // E2M1: 0b111 => -6.0 + ); + } else if constexpr (!std::is_unsigned_v) { + static_assert( + (std::is_integral_v && std::is_signed_v) || + cute::is_same_v, + "Two's complement LUT requires signed integer QuantType (int4b_t)"); + + // INT4 Two's Complement quantization: Use TC LUT values + storage[0] = pack4( + ElementScale(float(val) * (-8.0f)), // TC: 0b000 => -8 + ElementScale(float(val) * (-7.0f)), // TC: 0b001 => -7 + ElementScale(float(val) * (-6.0f)), // TC: 0b010 => -6 + ElementScale(float(val) * (-5.0f)) // TC: 0b011 => -5 + ); + storage[1] = pack4( + ElementScale(float(val) * (-4.0f)), // TC: 0b100 => -4 + ElementScale(float(val) * (-3.0f)), // TC: 0b101 => -3 + ElementScale(float(val) * (-2.0f)), // TC: 0b110 => -2 + ElementScale(float(val) * (-1.0f)) // TC: 0b111 => -1 + ); + } else { + // Unsigned ElementScale and QuantType: Pack positive LUT values (for UINT4) + static_assert(((std::is_integral_v && std::is_unsigned_v) || + cute::is_same_v) && + std::is_unsigned_v, + "Uint4 LUT requires unsigned QuantType (uint4b_t) and ElementScale"); + + // UINT4 LUT: pack positive LUT values + storage[0] = pack4( + ElementScale(float(val) * 8.0f), + ElementScale(float(val) * 7.0f), + ElementScale(float(val) * 6.0f), + ElementScale(float(val) * 5.0f) + ); + storage[1] = pack4( + ElementScale(float(val) * 4.0f), + ElementScale(float(val) * 3.0f), + ElementScale(float(val) * 2.0f), + ElementScale(float(val) * 1.0f) + ); } } CUTLASS_HOST_DEVICE @@ -230,7 +289,7 @@ private: Storage storage[2] {}; CUTLASS_HOST_DEVICE - static Storage pack4(T c1, T c2, T c3, T c4) { + static Storage pack4(ElementScale c1, ElementScale c2, ElementScale c3, ElementScale c4) { Storage result = 0; result |= (static_cast(reinterpret_cast(c4)) << 24); result |= (static_cast(reinterpret_cast(c3)) << 16); @@ -239,25 +298,25 @@ private: return result; } CUTLASS_HOST_DEVICE - T get() const { + ElementScale get() const { auto stage = static_cast(storage[0] >> 8); #if defined(__CUDA_ARCH__) - return reinterpret_cast(stage); + return reinterpret_cast(stage); #else - T tmp; + ElementScale tmp; std::memcpy(&tmp, &stage, sizeof(Stage)); return tmp; #endif } CUTLASS_HOST_DEVICE - T get(int idx) const { + ElementScale get(int idx) const { Stage stage; if (idx < 4) stage = static_cast(storage[0] >> (8 * idx)); else stage = static_cast(storage[1] >> (8 * idx - 32)); #if defined(__CUDA_ARCH__) - return reinterpret_cast(stage); + return reinterpret_cast(stage); #else - T tmp; + ElementScale tmp; std::memcpy(&tmp, &stage, sizeof(Stage)); return tmp; #endif @@ -301,7 +360,7 @@ static bool unified_encode_int4b(cutlass::int4b_t const *block_in, cutlass::int4 return true; } -template +template static bool pack_scale_fp8(ElementScale const *block_in, cutlass::Array *block_out, const size_t block_size) { std::vector data_in(block_size); std::vector> data_out(block_size); @@ -315,7 +374,7 @@ static bool pack_scale_fp8(ElementScale const *block_in, cutlass::Array tmp(data_in[i]); + cutlass::packed_scale_t tmp(data_in[i]); data_out[i] = reinterpret_cast const&>(tmp); }