Merge branch 'develop' into ck_tile/fmha_receipt_aiter

.azuredevops/rocm-ci.yml

@@ -14,6 +14,7 @@ trigger:
   branches:
     include:
     - develop
+    - amd-develop
   paths:
     exclude:
     - .github

CMakeLists.txt

@@ -92,6 +92,7 @@ endif()
 add_compile_options(-Wno-bit-int-extension)
 add_compile_options(-Wno-pass-failed)
 add_compile_options(-Wno-switch-default)
+add_compile_options(-Wno-unique-object-duplication)
 
 if(DL_KERNELS)
     add_definitions(-DDL_KERNELS)
@@ -196,17 +197,20 @@ if (SUPPORTED_GPU_TARGETS MATCHES "gfx9")
     add_definitions(-DCK_USE_XDL)
     set(CK_USE_XDL "ON")
 endif()
-if (SUPPORTED_GPU_TARGETS MATCHES "gfx94")
+if (SUPPORTED_GPU_TARGETS MATCHES "gfx94" OR SUPPORTED_GPU_TARGETS MATCHES "gfx95")
     message("Enabling FP8 gemms on native architectures")
     add_definitions(-DCK_USE_GFX94)
     set(CK_USE_GFX94 "ON")
 endif()
+if (SUPPORTED_GPU_TARGETS MATCHES "gfx95")
+	add_definitions(-DCK_USE_AMD_MFMA_GFX950)
+endif()
 if (SUPPORTED_GPU_TARGETS MATCHES "gfx11" OR SUPPORTED_GPU_TARGETS MATCHES "gfx12")
     message("Enabling WMMA instances")
     add_definitions(-DCK_USE_WMMA)
     set(CK_USE_WMMA "ON")
 endif()
-if (SUPPORTED_GPU_TARGETS MATCHES "gfx12")
+if (SUPPORTED_GPU_TARGETS MATCHES "gfx12" OR SUPPORTED_GPU_TARGETS MATCHES "gfx950")
     add_definitions(-DCK_USE_OCP_FP8)
     set(CK_USE_OCP_FP8 "ON")
 endif()
@@ -214,6 +218,10 @@ if (SUPPORTED_GPU_TARGETS MATCHES "gfx90a" OR SUPPORTED_GPU_TARGETS MATCHES "gfx
     add_definitions(-DCK_USE_FNUZ_FP8)
     set(CK_USE_FNUZ_FP8 "ON")
 endif()
+if (SUPPORTED_GPU_TARGETS MATCHES "gfx950")
+    add_definitions(-DCK_USE_NATIVE_MX_SUPPORT)
+    set(CK_USE_NATIVE_MX_SUPPORT "ON")
+endif()
 
 option(CK_USE_FP8_ON_UNSUPPORTED_ARCH "Enable FP8 GEMM instances on older architectures" OFF)
 if(CK_USE_FP8_ON_UNSUPPORTED_ARCH AND (SUPPORTED_GPU_TARGETS MATCHES "gfx90a" OR SUPPORTED_GPU_TARGETS MATCHES "gfx908"))

Jenkinsfile

@@ -117,7 +117,7 @@ def getDockerImage(Map conf=[:]){
     {
         echo "Pulling down image: ${image}"
         retimage = docker.image("${image}")
-        withDockerRegistry([ credentialsId: "docker_test_cred", url: "" ]) {
+        withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
             retimage.pull()
         }
     }
@@ -148,7 +148,7 @@ def buildDocker(install_prefix){
             //force building the new docker if that parameter is true
             echo "Building image: ${image_name}"
             retimage = docker.build("${image_name}", dockerArgs)
-            withDockerRegistry([ credentialsId: "docker_test_cred", url: "" ]) {
+            withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
                 retimage.push()
             }
             sh 'docker images -q -f dangling=true | xargs --no-run-if-empty docker rmi'
@@ -162,7 +162,7 @@ def buildDocker(install_prefix){
     catch(Exception ex){
         echo "Unable to locate image: ${image_name}. Building image now"
         retimage = docker.build("${image_name}", dockerArgs + ' .')
-        withDockerRegistry([ credentialsId: "docker_test_cred", url: "" ]) {
+        withDockerRegistry([ credentialsId: "ck_docker_cred", url: "" ]) {
             retimage.push()
         }
     }
@@ -795,8 +795,8 @@ pipeline {
             description: "Run the ck_tile FMHA tests (default: OFF)")
         booleanParam(
             name: "RUN_CK_TILE_GEMM_TESTS",
-            defaultValue: false,
-            description: "Run the ck_tile GEMM tests (default: OFF)")
+            defaultValue: true,
+            description: "Run the ck_tile GEMM tests (default: ON)")
         booleanParam(
             name: "BUILD_INSTANCES_ONLY",
             defaultValue: false,

client_example/01_gemm/README.md

+[Back to supported operations](../../../include/ck/README.md)
+# Composable Kernel GEMM
+
+## GEMM
+General matrix multiplications operation. In CK GEMM operation is called as `DeviceGemm` and requires following types as template parameters:
+
+* **ALayout** - A matrix layout (RowMajor/ColumnMajor).
+* **BLayout** - B matrix layout (RowMajor/ColumnMajor).
+* **CLayout** - B matrix layout (RowMajor/ColumnMajor).
+* **ADataType** - A matrix data type.
+* **BDataType** - B matrix data type.
+* **CDataType** - B matrix data type.
+* **AElementwiseOperation** - Fused operation on tensor A before GEMM.
+* **BElementwiseOperation** - Fused operation on tensor B before GEMM.
+* **CElementwiseOperation** - Fused operation on tensor C after GEMM.
+
+For matrices with large K dimension `DeviceGemmSplitK` implementation is available. This implementation allows user to split K dimension between work groups. This implementation uses `AtomicAdd` operation on global memory, thus need to zero-out output buffer for correct results.
+
+For fused operations with additional tensor there are `DeviceGemmMultipleABD` or `DeviceGemmMultipleD` operation which require following parameters:
+* **DsLayout** - layouts for additional tensors for fused operations.
+* **DsDataType** - data types for additional tensors for fused operations.
+
+For `DeviceGemmMultipleABD` **ALayout**, **BLayout**, **ADataType** and **BDataType** user should pass a tuple.
+
+List of the device operations in CK:
+
+* **DeviceGemmDl** - Device operation with DL instructions.
+* **DeviceGemmDpp** - Device operation with DL instructions with DPP instructions during data load.
+* **DeviceGemmWmma_CShuffle** - Device operation with WMMA instructions with CShuffle optimization for more optimized data store.
+* **DeviceGemm_Xdl_CShuffle_LdsDirectLoad** - Device operation with XDL instructions and CShuffle optimization for more optimized data store and direct load from global memory to shared memory.
+* **DeviceGemm_Xdl_CShuffle** - Device operation with XDL instructions with CShuffle optimization for more optimized data store.
+* **DeviceGemm_Xdl_CShuffleV2** - Device operation with XDL instructions with CShuffle optimization for more optimized data store. GEMM pipeline has been optimized compared to **DeviceGemm_Xdl_CShuffle**.
+* **DeviceGemmXdlSkipBLds** - Device operation with XDL instructions. Load to shared memory has been skiped for B matrix.
+* **DeviceGemm_Xdl_WaveletModel_CShuffle** - Device operation with XDL instructions with CShuffle optimization for more optimized data store. Producer and consumer scheme cooperation between waves in workgroup.
+* **DeviceGemmXdl** - Device operation with XDL instructions.
+
+Table of supported cases by instance factory with XDL instruction for Row/Row/Row, Row/Column/Row, Column/Row/Row or Column/Column/Row:
+
+|       |Is supported|
+|-------|---|
+|bf16|✓|
+|fp16|✓|
+|fp32|✓|
+|int8|✓|
+|fp8 |✓|
+
+Table of supported cases by instance factory with WMMA instruction for Row/Row/Row, Row/Column/Row, Column/Row/Row or Column/Column/Row:
+
+|       |Is supported|
+|-------|---|
+|bf16|✓|
+|fp16|✓|
+|fp32|✗|
+|int8|✓|
+|fp8 |✗|
+
+Table of supported cases by instance factory with DL instruction for Row/Row/Row, Row/Column/Row, Column/Row/Row or Column/Column/Row:
+
+|       |Is supported|
+|-------|---|
+|bf16|✗|
+|fp16|✓|
+|fp32|✓|
+|int8|✓|
+|fp8 |✗|
+
+Table of supported cases by instance factory with fused output elementwise operation:
+
+* **B Matrix Multiply + Add + Gelu** - bf16 (int8 for B matrix)
+* **B Matrix Multiply + Add** - bf16 (int8 for B matrix)
+* **B Matrix Multiply + Gelu** - bf16 (int8 for B matrix)
+* **B Matrix Multiply** - bf16 (int8 for B matrix)
+
+* **Add + Add + Gelu** - fp16
+* **Add + Gelu** - fp16, bf16 (int8 for B matrix) for Row/Column/Row
+* **Multiply** - fp16
+* **Add + Multiply** - fp16
+* **Add + Relu** - fp16 (int8 for B matrix)  for Row/Column/Row, bf16 (int8 for B matrix) for Row/Column/Row
+* **Add + Silu** - fp16 (int8 for B matrix)  for Row/Column/Row, bf16 (int8 for B matrix) for Row/Column/Row
+* **Add** - fp16 (int8 for B matrix)  for Row/Column/Row, bf16 (int8 for B matrix) for Row/Column/Row
+* **Bilinear** - fp16, int8
+* **Gelu** - fp16
+* **Multiply + Add** - fp16 for Row/Column/Row and Row/Row/Row, fp16 (int8 for B matrix, fp32 for Bias) for Row/Column/Row and Row/Row/Row, 
+* **Quantization** - int8
+
+## GEMM V2 (Universal GEMM)
+General matrix multiplications operation optimized for MI300 series. Operation is called as `DeviceGemmV2` and requires following types as template parameters:
+
+* **ALayout** - A matrix layout (RowMajor/ColumnMajor).
+* **BLayout** - B matrix layout (RowMajor/ColumnMajor).
+* **CLayout** - B matrix layout (RowMajor/ColumnMajor).
+* **ADataType** - A matrix data type.
+* **BDataType** - B matrix data type.
+* **CDataType** - B matrix data type.
+* **AElementwiseOperation** - Fused operation on tensor A before GEMM.
+* **BElementwiseOperation** - Fused operation on tensor B before GEMM.
+* **CElementwiseOperation** - Fused operation on tensor C after GEMM.
+
+This implementation allows user to split K dimension between work groups. This implementation requires AtomicAdd operation on global memory (output buffer must be set to zeroes if splitK parameter is larger than one).
+
+List of the device operations for in CK:
+
+* **DeviceGemm_Xdl_CShuffleV3** - Device operation with XDL instructions with CShuffle optimization for more optimized data store.
+* **DeviceGemm_Xdl_CShuffleV3R1** - Device operation with XDL instructions with CShuffle optimization for more optimized data store. This implementation perform reduction on splitted K dimension after GEMM instead of AtomicAdd instruction. 
+
+Table of supported cases by instance factory with XDL instruction for Row/Row/Row, Row/Column/Row, Column/Row/Row or Column/Column/Row:
+
+|       |Is supported|
+|-------|---|
+|bf16|✓|
+|fp16|✓|
+|fp32|✗|
+|int8|✗|
+|fp8 (C bf16)|✓|
+|fp16 (A fp8)|✓|
+|fp16 (B fp8)|✓|
+
+## Others
+
+* **DeviceGemm_dequantB** - GEMM with dequantization (implemented with WMMA instructions).
+* **DeviceGemmMultipleD_ABScale** - GEMM with scale for A and B matrix.
+* **DeviceGemmMultipleDLayernorm** - GEMM fused with layernorm.
+* **DeviceGemmMultipleDMultipleR** - GEMM fused with reductions and custom global reductions operators.
+* **DeviceGemmReduce** - GEMM fused with reduction.
+* **DeviceGemm_Streamk_V2** - GEMM stream K implementation. Implementation allows to use reduction instead of AtomicAdd.
+* **DeviceGemmStreamK** - GEMM stream K implementation using AtomicAdd.

client_example/07_grouped_convnd_fwd/README.md

+[Back to supported operations](../../../include/ck/README.md)
+# Composable Kernel Grouped Convolution
+
+## Grouped Convolution Forward
+Grouped convolution operation for 1D, 2D or 3D spatial dimensions. Convolution utilizes GEMM kernel after tensor coordinate transform. In CK Grouped Convolution Forward operation is called as `DeviceGroupedConvFwdMultipleABD` and requires following types as template parameters:
+
+* **NumDimSpatial** - number of spatial dimensions (1D, 2D, 3D).
+* **InLayout** - input layout (NHWGC, GNHWC, NGCHW).
+* **WeiLayout** - weight layout (GKYXC).
+* **DsLayout** - layouts for additional tensors for fused operations.
+* **OutLayout** - output layout (NHWGK, GNHWK, NGKHW).
+* **ADataType** - input data type. Pass tuple if there is fused operation with input.
+* **BDataType** - weight data type. Pass tuple if there is fused operation with weight.
+* **DsDataType** - data types for additional tensors for fused operations.
+* **EDataType** - Output data type.
+* **AElementwiseOperation** - fused operation on tensor A (input).
+* **BElementwiseOperation** - fused operation on tensor B (weight).
+* **CDEElementwiseOperation** - fused operation on tensor C (output).
+* **AComputeType** - compute data type of tensor A for mfma instruction (ADataType by default).
+* **BComputeType** - compute data type of tensor B for mfma instruction (AComputeType by default).
+
+Grouped convolution forward support tensors larger than 2GB.
+
+List of the device operations for grouped convolution forward in CK:
+
+* **DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3** - Device operation with XDL instructions. Optimized for AMD Instinct MI300 series.
+* **DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle** - Device operation with XDL instructions and support of fused operations to input, weight and output.
+* **DeviceGroupedConvFwdMultipleD_Wmma_CShuffle** - Device operation with WMMA instructions.
+* **DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK** - Device operation with DL instructions.
+
+Table of supported cases by instance factory with XDL instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|bf16 |2D, 3D|2D|1D, 2D, 3D|
+|fp16 |2D, 3D|2D|1D, 2D, 3D|
+|fp32 |2D, 3D|2D|1D, 2D, 3D|
+|int8 |2D, 3D|2D|1D, 3D|
+|fp8  |3D|✗|✗|
+|bf8  |3D|✗|✗|
+
+Table of supported cases by instance factory with WMMA instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|fp16 |2D, 3D|✗|2D, 3D|
+|int8 |2D, 3D|✗|2D, 3D|
+
+Table of supported cases by instance factory with DL instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|bf16 |✗|✗|2D|
+|fp16 |✗|✗|2D|
+|fp32 |✗|✗|2D|
+|int8 |✗|✗|2D|
+
+Table of supported cases by instance factory with fused elementwise operation:
+
+* **Dynamic elementwise operation** - 2D/3D, NHWGC, bf16/fp16/fp32/int8
+* **Bilinear** - 3D, NHWGC, bf16/fp16/fp32/int8
+* **ConvInvScale** - 3D, NHWGC, fp8
+* **ConvScale** - 3D, NHWGC, fp8/bf8
+* **ConvScale + Add** - 3D, NHWGC, fp8
+* **ConvScale + Relu** - 3D, NHWGC, fp8
+* **Scale** - 3D, NHWGC, bf16/fp16/fp32/int8
+* **Scale + Add (for A and B)** - 3D, NHWGC, bf16/fp16/fp32/int8
+* **Scale + Add + Scale + Add + Relu** - 3D, NHWGC, bf16/fp16/fp32/int8

client_example/10_grouped_convnd_bwd_data/README.md

+[Back to supported operations](../../../include/ck/README.md)
+# Composable Kernel Grouped Convolution
+
+## Grouped Convolution Backward Data
+
+Grouped convolution operation for 1D, 2D or 3D spatial dimensions. Convolution utilizes GEMM kernel after tensor coordinate transform. In CK Grouped Convolution Backward Data operation is called as `DeviceGroupedConvBwdDataMultipleD` and requires following types as template parameters:
+
+* **NumDimSpatial** - number of spatial dimensions (1D, 2D, 3D).
+* **ALayout** - output layout (NHWGK, GNHWK, NGKHW).
+* **BLayout** - weight layout (GKYXC).
+* **DsLayout** - layouts for additional tensors for fused operations.
+* **ELayout** - input layout (NHWGC, GNHWC, NGCHW).
+* **ADataType** - output data type.
+* **BDataType** - weight data type.
+* **DsDataType** - data types for additional tensors for fused operations.
+* **EDataType** - input data type.
+* **AElementwiseOperation** - fused operation on tensor A (output).
+* **BElementwiseOperation** - fused operation on tensor B (weight).
+* **CDEElementwiseOperation** - fused operation on tensor C (input).
+* **AComputeType** - compute data type of tensor A for mfma instruction (ADataType by default).
+* **BComputeType** - compute data type of tensor B for mfma instruction (AComputeType by default).
+
+Grouped convolution backward data supports tensors larger than 2GB (except when image is larger than 2GB).
+
+List of the device operations for grouped convolution backward data in CK:
+
+* **DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1** - Device operation with XDL instructions and support of fused operations to input.
+* **DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle** - Device operation with WMMA instructions.
+
+Table of supported cases by instance factory with XDL instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|bf16|2D, 3D|✗|2D, 3D|
+|fp16 |2D, 3D|✗|2D, 3D|
+|fp32  |2D, 3D|✗|2D, 3D|
+
+Table of supported cases by instance factory with WMMA instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|fp16 |2D, 3D|✗|2D, 3D|
+|int8 |2D, 3D|✗|2D, 3D|
+
+Table of supported cases by instance factory with fused elementwise operation:
+
+* **Bilinear** - 3D, NHWGC, bf16/fp16/fp32
+* **Scale** - 3D, NHWGC, bf16/fp16/fp32

client_example/11_grouped_conv_bwd_weight/README.md

+[Back to supported operations](../../../include/ck/README.md)
+# Composable Kernel Grouped Convolution
+
+## Grouped Convolution Backward Weight
+
+Grouped convolution operation for 1D, 2D or 3D spatial dimensions. Convolution utilizes GEMM kernel after tensor coordinate transform. Backward weight version uses splitK feature (due to large GEMM K dimension). In CK Grouped Convolution Backward Weight operation is called as `DeviceGroupedConvBwdWeight` and requires following types as template parameters:
+
+* **NumDimSpatial** - number of spatial dimensions (1D, 2D, 3D).
+* **InLayout** - input layout (NHWGC, GNHWC, NGCHW).
+* **WeiLayout** - weight layout (GKYXC).
+* **OutLayout** - output layout (NHWGK, GNHWK, NGKHW).
+* **InDataType** - input data type.
+* **WeiDataType** - weight data type.
+* **OutDataType** - output data type.
+* **InElementwiseOperation** - fused operation on tensor input.
+* **WeiElementwiseOperation** - fused operation on tensor weight.
+* **OutElementwiseOperation** - fused operation on tensor output.
+* **ComputeTypeA** - compute data type of tensor A for mfma instruction (ADataType by default).
+* **ComputeTypeB** - compute data type of tensor B for mfma instruction (ComputeTypeA by default).
+
+For fused operations with additional tensor there is `DeviceGroupedConvBwdWeightMultipleD` operation which requires following parameters:
+* **DsLayout** - layouts for additional tensors for fused operations.
+* **DsDataType** - data types for additional tensors for fused operations.
+
+Grouped convolution backward weight doesn't supports tensors larger than 2GB.
+
+List of the device operations for grouped convolution backward weight in CK:
+
+* **DeviceGroupedConvBwdWeight_Xdl_CShuffle** - Device operation with XDL instructions.
+* **DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle** - Device operation with XDL instructions. Optimized for small C or K.
+* **DeviceGroupedConvBwdWeight_Wmma_CShuffle** - Device operation with WMMA instructions.
+* **DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle** - Device operation with XDL instructions and support of fused operations to output.
+* **DeviceGroupedConvBwdWeight_Dl** - Device operation with DL instructions.
+
+Table of supported cases by instance factory with XDL instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|bf16|2D, 3D|✗|✗|
+|bf16(fp32 for weight)|2D, 3D|✗|1D, 2D, 3D|
+|fp16 |2D, 3D|✗|1D, 2D, 3D|
+|fp32  |2D, 3D|✗|1D, 2D, 3D|
+
+Table of supported cases by instance factory with WMMA instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|fp16 |3D|✗|3D|
+|int8 |3D|✗|3D|
+
+Table of supported cases by instance factory with DL instruction:
+
+|       |NHWGC/GKYXC/NHWGK|NGCHW/GKYXC/NGKHW|GNHWC/GKYXC/GNHWK|
+|-------|---|---|---|
+|bf16(fp32 for weight)|1D, 2D, 3D|✗|1D, 2D, 3D|
+|fp16 |1D, 2D, 3D|✗|1D, 2D, 3D|
+|fp32  |1D, 2D, 3D|✗|1D, 2D, 3D|
+
+Table of supported cases by instance factory with fused elementwise operation:
+
+* **Bilinear** - 3D, NHWGC, bf16(fp32 for weight)/fp16/fp32
+* **Scale** - 3D, NHWGC, bf16(fp32 for weight)/fp16/fp32

client_example/CMakeLists.txt

@@ -56,7 +56,7 @@ if (GPU_TARGETS)
         add_definitions(-DCK_USE_WMMA)
         set(CK_USE_WMMA "ON")
     endif()
-    if (GPU_TARGETS MATCHES "gfx12")
+    if (GPU_TARGETS MATCHES "gfx12" OR GPU_TARGETS MATCHES "gfx950")
         add_definitions(-DCK_USE_OCP_FP8)
         set(CK_USE_OCP_FP8 "ON")
     endif()

codegen/driver/main.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <functional>
 #include <iostream>

codegen/include/ck/host/device_batched_gemm_softmax_gemm/operation.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+#include <vector>
+#include <string>
+#include "ck/host/types.hpp"
+#include "ck/host/operation/gemm.hpp"
+#include "ck/host/device_batched_gemm_softmax_gemm/problem.hpp"
+
+namespace ck {
+namespace host {
+namespace device_batched_gemm_softmax_gemm {
+
+// defines all values need for an instance of fwd conv
+struct Operation_Xdl_CShuffle
+{
+    // returns a vector of instances, only given fusion operators: will use default problem spec
+    static std::vector<std::vector<Operation_Xdl_CShuffle>>
+    CreateOperations(const std::string& prologue, const std::string& epilogue);
+    // returns a vector of instances, given a problem spec and fusion operators
+    static std::vector<Operation_Xdl_CShuffle>
+    CreateOperations(const Problem& prob, const std::string& prologue, const std::string& epilogue);
+    TensorDesc A{};
+    TensorDesc B{};
+    TensorDesc B1{};
+    TensorDesc C{};
+    DataType acc                    = DataType::Float;
+    DataType cs_type                = DataType::Half;
+    std::string a_elem_op           = PassThrough;
+    std::string b_elem_op           = PassThrough;
+    std::string b1_elem_op          = PassThrough;
+    std::string c_elem_op           = PassThrough;
+    std::string acc_elem_op         = Scale;
+    std::string prologue            = "";
+    std::string epilogue            = "";
+    std::string gemm_specialization = "ck::tensor_operation::device::GemmSpecialization::Default";
+    // tuning parameters
+    operation::TileDescGemmGemm tile_desc{};
+    operation::BlockTransferDesc a_block_transfer{};
+    operation::BlockTransferDesc b0_block_transfer{};
+    operation::BlockTransferDesc b1_block_transfer{};
+    operation::CShuffleDesc cshuffle{};
+    operation::CBlockTransferDesc c_block_transfer{};
+
+    bool mask_out_upper_triangle = false;
+
+    // functions to update fusion operators if provided
+    void update_prologue(const std::string& prologue);
+    void update_epilogue(const std::string& epilogue);
+    /**constexpr**/ bool
+    IsSupported(std::size_t MRaw_, std::size_t NRaw_, std::size_t KRaw_, std::size_t Gemm1NRaw_);
+    // returns a templated instance
+    Solution ToSolution() const;
+};
+
+} // namespace device_batched_gemm_softmax_gemm
+} // namespace host
+} // namespace ck

codegen/include/ck/host/device_batched_gemm_softmax_gemm/problem.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+#include <vector>
+#include <string>
+#include "ck/host/types.hpp"
+
+namespace ck {
+namespace host {
+namespace device_batched_gemm_softmax_gemm {
+
+// defines the problem specification for a GEMM operation
+struct Problem
+{
+    std::size_t M            = 0;
+    std::size_t N            = 0;
+    std::size_t K            = 0;
+    std::size_t O            = 0;
+    bool TransA              = false;
+    bool TransB              = false;
+    bool TransB1             = false;
+    bool TransC              = false;
+    DataType ADataType       = DataType::Half;
+    DataType BDataType       = DataType::Half;
+    DataType B1DataType      = DataType::Half;
+    DataType CDataType       = DataType::Half;
+    std::string AElementOp   = PassThrough;
+    std::string BElementOp   = PassThrough;
+    std::string B1ElementOp  = PassThrough;
+    std::string CElementOp   = PassThrough;
+    std::string AccElementOp = Scale;
+
+    // returns the correct device op file for the operation
+    std::string GetIncludeHeader() const;
+
+    // returns a list of instances based on the problem spec and provided fusion operations
+    std::vector<Solution> GetSolutions(const std::string& arch,
+                                       const std::string& prologue,
+                                       const std::string& epilogue) const;
+};
+
+} // namespace device_batched_gemm_softmax_gemm
+} // namespace host
+} // namespace ck

codegen/include/ck/host/device_gemm_multiple_d/operation.hpp

@@ -41,6 +41,8 @@ struct Operation_Xdl_CShuffle
     operation::BlockTransferDesc b_block_transfer{};
     operation::CShuffleDesc cshuffle{};
     operation::CBlockTransferDesc c_block_transfer{};
+    LoopScheduler loop_scheduler{};
+    PipelineVersion pipeline_version{};
 
     // functions to update fusion operators if provided
     void update_prologue(const std::string& prologue);

codegen/include/ck/host/operation/gemm.hpp

@@ -23,6 +23,26 @@ struct TileDesc
     int n_Xdl_per_wave           = 0;
     int num_gemmk_prefetch_stage = 0;
 };
+
+struct TileDescGemmGemm
+{
+    int block_size               = 0;
+    int gemm01_m_per_block       = 0;
+    int gemm0_n_per_block        = 0;
+    int gemm0_k_per_block        = 0;
+    int gemm1_n_per_block        = 0;
+    int gemm1_k_per_block        = 0;
+    int ak1                      = 0;
+    int bk1                      = 0;
+    int b1k1                     = 0;
+    int m_per_XDL                = 0;
+    int n_per_XDL                = 0;
+    int gemm0_m_Xdl_per_wave     = 0;
+    int gemm0_n_Xdl_per_wave     = 0;
+    int gemm1_n_Xdl_per_wave     = 0;
+    int num_gemmk_prefetch_stage = 0;
+};
+
 struct BlockTransferDesc
 {
     std::string thread_cluster_length        = "";

codegen/include/ck/host/types.hpp

@@ -66,6 +66,20 @@ enum class GemmType
 };
 std::string ToString(GemmType gt);
 
+enum class LoopScheduler
+{
+    Default,
+    Interwave,
+};
+std::string ToString(LoopScheduler ls);
+
+enum class PipelineVersion
+{
+    v1,
+    v2
+};
+std::string ToString(PipelineVersion pv);
+
 struct TensorDesc
 {
     DataType element;
@@ -84,6 +98,7 @@ const std::string S = SequenceStr({xs...});
 
 constexpr const char* PassThrough = "ck::tensor_operation::element_wise::PassThrough";
 constexpr const char* Bilinear    = "ck::tensor_operation::element_wise::Bilinear";
+constexpr const char* Scale       = "ck::tensor_operation::element_wise::Scale";
 
 } // namespace host
 } // namespace ck

codegen/src/device_batched_gemm_softmax_gemm.cpp

+
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "ck/host/device_batched_gemm_softmax_gemm/problem.hpp"
+#include "ck/host/device_batched_gemm_softmax_gemm/operation.hpp"
+#include "ck/host/utils.hpp"
+#include <algorithm>
+
+namespace ck {
+namespace host {
+namespace device_batched_gemm_softmax_gemm {
+
+// return the relevant device op file based on the operation
+std::string Problem::GetIncludeHeader() const
+{
+    return "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_xdl_cshuffle.hpp";
+}
+
+// returns templated instances when provided with a problem specification
+std::vector<Solution> Problem::GetSolutions(const std::string& arch,
+                                            const std::string& prologue,
+                                            const std::string& epilogue) const
+{
+    if(get_xdlop_archs().count(arch) == 0)
+        return {};
+    auto ops = ck::host::device_batched_gemm_softmax_gemm::Operation_Xdl_CShuffle::CreateOperations(
+        *this, prologue, epilogue); // obtains vector of instances
+    std::vector<Solution> result;
+    std::transform(ops.begin(), ops.end(), std::back_inserter(result), [&](const auto& op) {
+        return op.ToSolution(); // template instance with correct values
+    });
+    return result;
+}
+
+} // namespace device_batched_gemm_softmax_gemm
+} // namespace host
+} // namespace ck

codegen/src/device_gemm_multiple_d_operation_xdl_cshuffle.cpp

@@ -62,6 +62,12 @@ void Operation_Xdl_CShuffle::update_epilogue(const std::string& epi)
 // accounts for all possible combinations of Row/Col major
 static Layout ToLayout(bool Trans) { return Trans ? Layout::Column : Layout::Row; }
 
+// clang-format off
+// DeviceGemmMultipleD_Xdl_CShuffle<    Col,    Row, Row_Row_Tuple,    Row,   F16,   F16,     F32,      F32, F16_F16_Tuple,   F16, PassThrough, PassThrough, AddAddFastGelu, GemmMNKPadding,        1,    64,    16,    16,    32,   8,   8,   16,   16,    1,    1,     S<4, 16, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,         1,     S<4, 16, 1>,     S<0, 2, 1>,     S<0, 2, 1>,             1,              1,              8,         1,           1,           1,               S<1, 16, 1, 4>,               1,
+
+// DeviceGemmMultipleD_Xdl_CShuffle<    Row,    Col, Row_Row_Tuple,    Row,   F16,   F16,     F32,      F32, F16_F16_Tuple,   F16, PassThrough, PassThrough, AddAddFastGelu, GemmMNKPadding,        1,    64,    16,    16,    32,   8,   8,   16,   16,    1,    1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              1,              8,         1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              1,              8,         1,           1,           1,               S<1, 16, 1, 4>,               1, LoopScheduler::Default,        PipelineVersion::v1>
+// clang-format on
+
 // Hard-code tuning parameters in modularized fashion, string them together into a vector of
 // instances
 std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
@@ -83,6 +89,8 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
   {   128,    64,   128,    32,   8,   8,   32,   32,    2,    2,        1},
   {   256,   128,    64,    32,   8,   8,   32,   32,    2,    1,        1},
   {   256,    64,   128,    32,   8,   8,   32,   32,    1,    2,        1},
+//  Irregular tile
+  {    64,    16,    16,    32,   8,   8,   16,   16,    1,    1,        1},
         // clang-format on
     };
 
@@ -100,6 +108,8 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
   {    S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1},
   {    S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1},
   {    S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1},
+//  Irregular tile
+  {    S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              1,              8,         1},
         // clang-format on
     };
 
@@ -109,15 +119,17 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
 //   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|
 // Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          |
 //                |               |               |               |               |               |          |
+  {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              4,              8,         1},
+  {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              2,              8,         1},
+  {    S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              4,              8,         1},
+  {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              2,              8,         1},
+  {    S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              4,              8,         1},
+  {    S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              2,              8,         1},
+  {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              2,              8,         1},
+  {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,         1},
+//  Irregular tile
+  {    S<4, 16, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,         1},
         // clang-format on
-        {S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 8, 1},
-        {S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1},
-        {S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 8, 1},
-        {S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1},
-        {S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 8, 1},
-        {S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1},
-        {S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1},
-        {S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 1},
     };
 
     std::vector<operation::BlockTransferDesc> b_block_descriptions_rowmajor = {
@@ -134,6 +146,8 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
   {    S<4, 32, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              4,              8,         1},
   {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,         1},
   {    S<4, 64, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              2,              8,         1},
+//  Irregular tile
+  {    S<4, 16, 1>,     S<0, 2, 1>,     S<0, 2, 1>,              1,              1,              8,         1},
         // clang-format on
     };
 
@@ -151,6 +165,8 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
   {    S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1},
   {    S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1},
   {    S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1},
+//  Irregular tile
+  {    S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              1,              8,         1},
         // clang-format on
     };
 
@@ -167,6 +183,7 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
   {          1,           1},
   {          1,           1},
   {          1,           1},
+  {          1,           1},
   {          1,           1},
         // clang-format on
     };
@@ -185,6 +202,8 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
   {              S<1, 16, 1, 8>,               8},
   {              S<1, 32, 1, 8>,               8},
   {              S<1, 32, 1, 8>,               8},
+//  Irregular tile
+  {              S<1, 16, 1, 4>,               1},
         // clang-format on
     };
 
@@ -199,6 +218,14 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
     assert(tile_descriptions.size() == cshuffle_descriptions.size());
     assert(tile_descriptions.size() == c_block_descriptions.size());
 
+    const std::vector<std::tuple<LoopScheduler, PipelineVersion>> scheduler_pipeline_descriptions =
+        {
+            {LoopScheduler::Default, PipelineVersion::v1},
+            {LoopScheduler::Interwave, PipelineVersion::v1},
+            {LoopScheduler::Default, PipelineVersion::v2},
+        };
+    for(auto [loop_scheduler, pipeline_version] : scheduler_pipeline_descriptions)
+    {
         // Put all values together into a single operation > store into the result vector
         for(std::size_t i = 0; i < tile_descriptions.size(); i++)
         {
@@ -223,10 +250,13 @@ std::vector<Operation_Xdl_CShuffle> Operation_Xdl_CShuffle::CreateOperations(
                                                 x.tile_desc.m_per_block,
                                                 x.tile_desc.n_per_block,
                                                 x.tile_desc.k_per_block);
+            x.loop_scheduler      = loop_scheduler;
+            x.pipeline_version    = pipeline_version;
             x.update_prologue(prologue);
             x.update_epilogue(epilogue);
             result.push_back(x);
         }
+    }
     return result;
 }
 
@@ -263,7 +293,7 @@ static const char* const DeviceGemmMultipleD_Xdl_CShuffleTemplate =
     "${BBlockTransferSrcScalarPerVector}, ${BBlockTransferDstScalarPerVector_BK1}, "
     "${BBlockLdsExtraN}, ${CShuffleMXdlPerWavePerShuffle}, ${CShuffleNXdlPerWavePerShuffle}, "
     "${CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock}, "
-    "${CDEBlockTransferScalarPerVector_NPerBlock}>";
+    "${CDEBlockTransferScalarPerVector_NPerBlock}, ${LoopScheduler}, ${PipelineVersion}>";
 
 // use hardcoded instances from vector of operations to substitute values into instance template
 Solution Operation_Xdl_CShuffle::ToSolution() const
@@ -336,6 +366,8 @@ Solution Operation_Xdl_CShuffle::ToSolution() const
          this->c_block_transfer.cluster_lengths_m_block_m_wave_m_per_Xdl_n_block_n_wave_n_per_Xdl},
         {"CDEBlockTransferScalarPerVector_NPerBlock",
          std::to_string(this->c_block_transfer.scalar_per_vector_n_wave_n_per_Xdl)},
+        {"LoopScheduler", ToString(this->loop_scheduler)},
+        {"PipelineVersion", ToString(this->pipeline_version)},
     };
 
     return Solution{InterpolateString(DeviceGemmMultipleD_Xdl_CShuffleTemplate, values),

codegen/src/headers.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
 #include "ck/host/headers.hpp"
 #include "ck_headers.hpp"
 

codegen/src/types.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
 #include "ck/host/types.hpp"
 #include "ck/host/stringutils.hpp"
 #include <algorithm>
@@ -56,6 +59,26 @@ std::string ToString(GemmType gt)
     throw std::runtime_error("Incorrect gemm type");
 }
 
+std::string ToString(LoopScheduler ls)
+{
+    switch(ls)
+    {
+    case LoopScheduler::Default: return "ck::LoopScheduler::Default";
+    case LoopScheduler::Interwave: return "ck::LoopScheduler::Interwave";
+    }
+    throw std::runtime_error("Incorrect LoopScheduler type");
+}
+
+std::string ToString(PipelineVersion pv)
+{
+    switch(pv)
+    {
+    case PipelineVersion::v1: return "ck::PipelineVersion::v1";
+    case PipelineVersion::v2: return "ck::PipelineVersion::v2";
+    }
+    throw std::runtime_error("Incorrect PipelineVersion type");
+}
+
 std::string SequenceStr(const std::vector<int>& v)
 {
     return "ck::Sequence<" +

codegen/test/gemm_multiple_d.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+
 #include "ck/host/device_gemm_multiple_d/problem.hpp"
 #include "ck/host/device_gemm_multiple_d/operation.hpp"
 #include "ck/host/headers.hpp"