Merge branch 'amd-develop' into amd-master

CHANGELOG.md

@@ -25,8 +25,8 @@ Full documentation for Composable Kernel is not yet available.
 - Added multi-embeddings support (#542).
 - Added Navi3x blockwise GEMM and real GEMM support (#541).
 - Added Navi grouped ConvBwdWeight support (#505).
-- Added pool3d forward (#697).
-- Added maxpool backward (#750).
+- Added MaxPool, AvgPool forward (#815).
+- Added MaxPool backward (#750).
 
 ### Changed
 - Changed ...

CMakeLists.txt

@@ -93,7 +93,7 @@ find_package(hip)
 # SWDEV-413293 and https://reviews.llvm.org/D155213
 math(EXPR hip_VERSION_FLAT "(${hip_VERSION_MAJOR} * 1000 + ${hip_VERSION_MINOR}) * 100000 + ${hip_VERSION_PATCH}")
 message("hip_version_flat=${hip_VERSION_FLAT}")
-if(${hip_VERSION_FLAT} GREATER 500700000)
+if(${hip_VERSION_FLAT} GREATER 500723302)
    message("Adding the fno-offload-uniform-block compiler flag")
    add_compile_options(-fno-offload-uniform-block)
 endif()

CONTRIBUTORS.md

@@ -6,9 +6,11 @@ This is the list of developers and contributors to Composable Kernel library
 ## Developers
 [Chao Liu](https://github.com/asroy), [Jing Zhang](https://github.com/zjing14), 2018-2023
 
-[Letao Qin](https://github.com/ltqin), [Qianfeng Zhang](https://github.com/qianfengz), [Liang Huang](https://github.com/carlushuang), [Shaojie Wang](https://github.com/shaojiewang), 2019-2022
+[Letao Qin](https://github.com/ltqin), [Qianfeng Zhang](https://github.com/qianfengz), [Liang Huang](https://github.com/carlushuang), [Shaojie Wang](https://github.com/shaojiewang), 2019-2023
 
-[Anthony Chang](https://github.com/rosenrodt), [Chunyu Lai](https://github.com/rocking5566), [Illia Silin](https://github.com/illsilin), [Adam Osewski](https://github.com/aosewski), [Poyen Chen](https://github.com/poyenc), [Rosty Geyyer](https://github.com/geyyer), 2022
+[Anthony Chang](https://github.com/rosenrodt), [Chunyu Lai](https://github.com/rocking5566), [Illia Silin](https://github.com/illsilin), [Adam Osewski](https://github.com/aosewski), [Poyen Chen](https://github.com/poyenc), [Rosty Geyyer](https://github.com/geyyer), [Astha Rai](https://github.com/arai713), [Shi YanXing](https://github.com/Yanxing-Shi), 2022-2023
+
+[Hari Sadasivan](https://github.com/hsadasiv), [Bartlomiej Kocot](https://github.com/bartekxk), [Bartlomiej Wroblewski](https://github.com/bwroblew), 2023
 
 Hanwen Chang, 2019-2021,
 

Dockerfile

@@ -29,6 +29,11 @@ RUN if [ "$ROCMVERSION" != "5.7" ]; then \
         sh -c "wget http://artifactory-cdn.amd.com/artifactory/list/amdgpu-deb/amdgpu-install-internal_5.7-20.04-1_all.deb" && \
         apt update && apt-get install -y ./amdgpu-install-internal_5.7-20.04-1_all.deb && \
         amdgpu-repo --amdgpu-build=1609671 --rocm-build=compute-rocm-npi-mi300/1354; \
+    elif [ "$ROCMVERSION" = "5.7" ] && [ "$compiler_version" = "rc1" ]; then \
+        sh -c "wget http://artifactory-cdn.amd.com/artifactory/list/amdgpu-deb/amdgpu-install-internal_5.7-20.04-1_all.deb" && \
+        apt update && apt-get install -y ./amdgpu-install-internal_5.7-20.04-1_all.deb && \
+        sh -c 'echo deb [arch=amd64 trusted=yes] http://compute-artifactory.amd.com/artifactory/list/rocm-release-archive-20.04-deb/ 5.7 rel-19 > /etc/apt/sources.list.d/rocm-build.list' && \
+        amdgpu-repo --amdgpu-build=1637781; \
     fi
 
 RUN sh -c "echo deb http://mirrors.kernel.org/ubuntu focal main universe | tee -a /etc/apt/sources.list"

Jenkinsfile

@@ -612,7 +612,7 @@ def process_results(Map conf=[:]){
 }
 
 //launch develop branch daily at 23:00 UT in FULL_QA mode and at 19:00 UT with latest staging compiler version
-CRON_SETTINGS = BRANCH_NAME == "develop" ? '''0 23 * * * % RUN_FULL_QA=true
+CRON_SETTINGS = BRANCH_NAME == "develop" ? '''0 23 * * * % RUN_FULL_QA=true;ROCMVERSION=5.7;COMPILER_VERSION=rc1
                                               0 21 * * * % ROCMVERSION=5.6;COMPILER_VERSION=;COMPILER_COMMIT=
                                               0 19 * * * % BUILD_DOCKER=true;COMPILER_VERSION=amd-stg-open;COMPILER_COMMIT=''' : ""
 
@@ -710,8 +710,8 @@ pipeline {
                     }
                     agent{ label rocmnode("gfx908 || gfx90a") }
                     environment{
-                        setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx908;gfx90a;gfx940" """
-                        execute_args = """ cd ../client_example && rm -rf build && mkdir build && cd build && cmake -D CMAKE_PREFIX_PATH="${env.WORKSPACE}/install;/opt/rocm" -DGPU_TARGETS="gfx908;gfx90a;gfx940" -D CMAKE_CXX_COMPILER="${build_compiler()}" .. && make -j """ 
+                        setup_args = """ -DCMAKE_INSTALL_PREFIX=../install -DGPU_TARGETS="gfx908;gfx90a;gfx940;gfx941" """
+                        execute_args = """ cd ../client_example && rm -rf build && mkdir build && cd build && cmake -D CMAKE_PREFIX_PATH="${env.WORKSPACE}/install;/opt/rocm" -DGPU_TARGETS="gfx908;gfx90a;gfx940;gfx941" -D CMAKE_CXX_COMPILER="${build_compiler()}" .. && make -j """ 
                     }
                     steps{
                         Build_CK_and_Reboot(setup_args: setup_args, config_targets: "install", no_reboot:true, build_type: 'Release', execute_cmd: execute_args, prefixpath: '/usr/local')

client_example/19_pool_fwd/avg_pool3d_fwd.cpp

@@ -16,6 +16,9 @@ using InDataType    = ck::half_t;
 using OutDataType   = ck::half_t;
 using IndexDataType = int32_t;
 
+using InLayout  = ck::tensor_layout::convolution::NDHWC;
+using OutLayout = ck::tensor_layout::convolution::NDHWC;
+
 constexpr ck::index_t InOutRank  = 5;
 constexpr ck::index_t WindowRank = 3;
 #if 0
@@ -55,6 +58,9 @@ int main(int argc, char* argv[])
     ck::index_t window_stride_d   = 2;
     ck::index_t window_stride_h   = 2;
     ck::index_t window_stride_w   = 2;
+    ck::index_t window_dilation_d = 1;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_d     = 1;
     ck::index_t in_left_pad_h     = 1;
     ck::index_t in_left_pad_w     = 1;
@@ -62,15 +68,20 @@ int main(int argc, char* argv[])
     ck::index_t in_right_pad_h    = 1;
     ck::index_t in_right_pad_w    = 1;
 
-    ck::index_t Do = (Di + in_left_pad_d + in_right_pad_d - Z) / window_stride_d + 1;
-    ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
-    ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;
+    const ck::index_t Zs = (Z - 1) * window_dilation_d + 1;
+    const ck::index_t Ys = (Y - 1) * window_dilation_h + 1;
+    const ck::index_t Xs = (X - 1) * window_dilation_w + 1;
+    ck::index_t Do       = (Di + in_left_pad_d + in_right_pad_d - Zs) / window_stride_d + 1;
+    ck::index_t Ho       = (Hi + in_left_pad_h + in_right_pad_h - Ys) / window_stride_h + 1;
+    ck::index_t Wo       = (Wi + in_left_pad_w + in_right_pad_w - Xs) / window_stride_w + 1;
 
     // Pool API only support the order of NCDHW
     std::vector<ck::index_t> in_length              = {N, C, Di, Hi, Wi};
     std::vector<ck::index_t> out_length             = {N, C, Do, Ho, Wo};
     std::vector<ck::index_t> window_spatial_lengths = {Z, Y, X};
     std::vector<ck::index_t> window_strides = {window_stride_d, window_stride_h, window_stride_w};
+    std::vector<ck::index_t> window_dilations{
+        window_dilation_d, window_dilation_h, window_dilation_w};
     std::vector<ck::index_t> input_left_pads  = {in_left_pad_d, in_left_pad_h, in_left_pad_w};
     std::vector<ck::index_t> input_right_pads = {in_right_pad_d, in_right_pad_h, in_right_pad_w};
 
@@ -90,6 +101,8 @@ int main(int argc, char* argv[])
                                                                  InDataType,
                                                                  OutDataType,
                                                                  IndexDataType,
+                                                                 InLayout,
+                                                                 OutLayout,
                                                                  ReduceOpId,
                                                                  OutputIndex>;
 
@@ -122,6 +135,7 @@ int main(int argc, char* argv[])
             out_tensor_stride,
             out_tensor_stride,
             window_strides,
+            window_dilations,
             input_left_pads,
             input_right_pads,
             {2, 3, 4});
@@ -181,6 +195,7 @@ int main(int argc, char* argv[])
             out_tensor_stride,
             out_tensor_stride,
             window_strides,
+            window_dilations,
             input_left_pads,
             input_right_pads,
             {2, 3, 4});

client_example/19_pool_fwd/max_pool2d_fwd.cpp

@@ -10,14 +10,18 @@
 #include "ck/tensor_operation/gpu/device/device_pool_fwd.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
-#include "ck/library/tensor_operation_instance/gpu/pool2d_fwd.hpp"
+#include "ck/library/tensor_operation_instance/gpu/pool3d_fwd.hpp"
 
 using InDataType    = ck::half_t;
 using OutDataType   = ck::half_t;
 using IndexDataType = int32_t;
 
-constexpr ck::index_t InOutRank  = 4;
-constexpr ck::index_t WindowRank = 2;
+// We use pool3d to implement pool2d in this example
+using InLayout  = ck::tensor_layout::convolution::NDHWC;
+using OutLayout = ck::tensor_layout::convolution::NDHWC;
+
+constexpr ck::index_t InOutRank  = 5;
+constexpr ck::index_t WindowRank = 3;
 #if 1
 constexpr auto ReduceOpId  = ck::ReduceTensorOp::MAX;
 constexpr bool OutputIndex = true;
@@ -42,6 +46,35 @@ struct SimpleDeviceMem
     void* p_mem_;
 };
 
+void TransformPool2dparamToPool3d(std::vector<ck::index_t>& input_lengths,
+                                  std::vector<ck::index_t>& window_lengths,
+                                  std::vector<ck::index_t>& output_lengths,
+                                  std::vector<ck::index_t>& input_stride,
+                                  std::vector<ck::index_t>& output_stride,
+                                  std::vector<ck::index_t>& indices_stride,
+                                  std::vector<ck::index_t>& window_strides,
+                                  std::vector<ck::index_t>& window_dilations,
+                                  std::vector<ck::index_t>& input_left_pads,
+                                  std::vector<ck::index_t>& input_right_pads,
+                                  std::vector<ck::index_t>& pooling_dims)
+{
+    // NCHW to NCDHW
+    input_lengths.insert(input_lengths.begin() + 2, 1);
+    output_lengths.insert(output_lengths.begin() + 2, 1);
+    input_stride.insert(input_stride.begin() + 2, 0);
+    output_stride.insert(output_stride.begin() + 2, 0);
+    indices_stride.insert(indices_stride.begin() + 2, 0);
+
+    // YX to ZYX
+    window_lengths.insert(window_lengths.begin(), 1);
+    window_strides.insert(window_strides.begin(), 0);
+    window_dilations.insert(window_dilations.begin(), 0);
+    input_left_pads.insert(input_left_pads.begin(), 0);
+    input_right_pads.insert(input_right_pads.begin(), 0);
+
+    pooling_dims = {2, 3, 4};
+}
+
 int main(int argc, char* argv[])
 {
     ck::index_t N                 = 2;
@@ -52,21 +85,27 @@ int main(int argc, char* argv[])
     ck::index_t Wi                = 30;
     ck::index_t window_stride_h   = 2;
     ck::index_t window_stride_w   = 2;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_h     = 1;
     ck::index_t in_left_pad_w     = 1;
     ck::index_t in_right_pad_h    = 1;
     ck::index_t in_right_pad_w    = 1;
 
-    ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
-    ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;
+    const ck::index_t Ys = (Y - 1) * window_dilation_h + 1;
+    const ck::index_t Xs = (X - 1) * window_dilation_w + 1;
+    ck::index_t Ho       = (Hi + in_left_pad_h + in_right_pad_h - Ys) / window_stride_h + 1;
+    ck::index_t Wo       = (Wi + in_left_pad_w + in_right_pad_w - Xs) / window_stride_w + 1;
 
     // Pool API only support the order of NCHW
     std::vector<ck::index_t> in_length              = {N, C, Hi, Wi};
     std::vector<ck::index_t> out_length             = {N, C, Ho, Wo};
     std::vector<ck::index_t> window_spatial_lengths = {Y, X};
     std::vector<ck::index_t> window_strides         = {window_stride_h, window_stride_w};
+    std::vector<ck::index_t> window_dilations       = {window_dilation_h, window_dilation_w};
     std::vector<ck::index_t> input_left_pads        = {in_left_pad_h, in_left_pad_w};
     std::vector<ck::index_t> input_right_pads       = {in_right_pad_h, in_right_pad_w};
+    std::vector<ck::index_t> pooling_dims           = {2, 3};
 
     std::size_t in_tensor_size  = N * C * Hi * Wi;
     std::size_t out_tensor_size = N * C * Ho * Wo;
@@ -75,6 +114,18 @@ int main(int argc, char* argv[])
     std::vector<ck::index_t> in_tensor_stride  = {C * Hi * Wi, 1, Wi * C, C};
     std::vector<ck::index_t> out_tensor_stride = {C * Ho * Wo, 1, Wo * C, C};
 
+    TransformPool2dparamToPool3d(in_length,
+                                 window_spatial_lengths,
+                                 out_length,
+                                 in_tensor_stride,
+                                 out_tensor_stride,
+                                 out_tensor_stride,
+                                 window_strides,
+                                 window_dilations,
+                                 input_left_pads,
+                                 input_right_pads,
+                                 pooling_dims);
+
     SimpleDeviceMem in_device_buf(sizeof(InDataType) * in_tensor_size);
     SimpleDeviceMem out_device_buf(sizeof(OutDataType) * out_tensor_size);
     SimpleDeviceMem out_indices_device_buf(sizeof(IndexDataType) * out_tensor_size);
@@ -84,6 +135,8 @@ int main(int argc, char* argv[])
                                                                  InDataType,
                                                                  OutDataType,
                                                                  IndexDataType,
+                                                                 InLayout,
+                                                                 OutLayout,
                                                                  ReduceOpId,
                                                                  OutputIndex>;
 
@@ -116,9 +169,10 @@ int main(int argc, char* argv[])
             out_tensor_stride,
             out_tensor_stride,
             window_strides,
+            window_dilations,
             input_left_pads,
             input_right_pads,
-            {2, 3});
+            pooling_dims);
 
         auto invoker_ptr = op_ptr->MakeInvokerPointer();
 
@@ -175,9 +229,10 @@ int main(int argc, char* argv[])
             out_tensor_stride,
             out_tensor_stride,
             window_strides,
+            window_dilations,
             input_left_pads,
             input_right_pads,
-            {2, 3});
+            pooling_dims);
 
         auto invoker_ptr = op_ptr->MakeInvokerPointer();
 

docs/sphinx/requirements.in

-rocm-docs-core==0.10.3
+rocm-docs-core>=0.20.0
 sphinxcontrib-bibtex==2.5.0

docs/sphinx/requirements.txt

@@ -38,6 +38,8 @@ docutils==0.16
     #   pydata-sphinx-theme
     #   sphinx
     #   sphinxcontrib-bibtex
+fastjsonschema==2.18.0
+    # via rocm-docs-core
 gitdb==4.0.10
     # via gitpython
 gitpython==3.1.31
@@ -46,20 +48,12 @@ idna==3.4
     # via requests
 imagesize==1.4.1
     # via sphinx
-importlib-metadata==6.0.0
-    # via
-    #   sphinx
-    #   sphinxcontrib-bibtex
-importlib-resources==5.12.0
-    # via rocm-docs-core
 jinja2==3.1.2
     # via
     #   myst-parser
     #   sphinx
 latexcodec==2.0.1
     # via pybtex
-linkify-it-py==1.0.3
-    # via myst-parser
 markdown-it-py==2.2.0
     # via
     #   mdit-py-plugins
@@ -70,7 +64,7 @@ mdit-py-plugins==0.3.5
     # via myst-parser
 mdurl==0.1.2
     # via markdown-it-py
-myst-parser[linkify]==1.0.0
+myst-parser==1.0.0
     # via rocm-docs-core
 packaging==23.0
     # via
@@ -99,18 +93,17 @@ pyjwt[crypto]==2.6.0
     # via pygithub
 pynacl==1.5.0
     # via pygithub
-pytz==2023.3
-    # via babel
 pyyaml==6.0
     # via
     #   myst-parser
     #   pybtex
+    #   rocm-docs-core
     #   sphinx-external-toc
 requests==2.28.2
     # via
     #   pygithub
     #   sphinx
-rocm-docs-core==0.10.3
+rocm-docs-core>=0.20.0
     # via -r requirements.in
 six==1.16.0
     # via
@@ -160,13 +153,7 @@ sphinxcontrib-serializinghtml==1.1.5
     # via sphinx
 typing-extensions==4.5.0
     # via pydata-sphinx-theme
-uc-micro-py==1.0.1
-    # via linkify-it-py
 urllib3==1.26.15
     # via requests
 wrapt==1.15.0
     # via deprecated
-zipp==3.15.0
-    # via
-    #   importlib-metadata
-    #   importlib-resources

example/01_gemm/CMakeLists.txt

@@ -6,6 +6,8 @@ if(DL_KERNELS)
   if(DTYPES MATCHES "fp16" OR NOT DEFINED DTYPES)
     add_example_executable(example_gemm_dl_fp16 gemm_dl_fp16.cpp)
     add_dependencies(example_gemm_dl example_gemm_dl_fp16)
+    add_example_executable(example_gemm_dl_dpp8_fp16 gemm_dl_dpp8_fp16.cpp)
+    add_dependencies(example_gemm_dl example_gemm_dl_dpp8_fp16)
   endif()
   if(DTYPES MATCHES "int8" OR NOT DEFINED DTYPES)
     add_example_executable(example_gemm_dl_int8 gemm_dl_int8.cpp)
@@ -64,3 +66,6 @@ if(DTYPES MATCHES "fp8" OR NOT DEFINED DTYPES)
     add_dependencies(example_gemm_xdl example_gemm_xdl_f8)
   endif()
 endif()
+
+add_example_executable(example_gemm_xdl_fp16_f8 gemm_xdl_fp16_f8.cpp)
+add_dependencies(example_gemm_xdl example_gemm_xdl_fp16_f8)

example/01_gemm/gemm_dl_dpp8_fp16.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "common.hpp"
+
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_dl_dpp8.hpp"
+
+using ADataType   = ck::half_t;
+using BDataType   = ck::half_t;
+using CDataType   = ck::half_t;
+using AccDataType = float;
+
+using ALayout = Col;
+using BLayout = Row;
+using CLayout = Row;
+
+using AElementOp = PassThrough;
+using BElementOp = PassThrough;
+using CElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+
+// clang-format off
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemmDlDpp8
+// ######|     AData|     BData|     CData|     AccData| ALayout| BLayout| CLayout|           A|           B|           C|           GEMM| Block|  MPer|  NPer| K0Per| K1|      M1Per|      N1Per|   KPer|  M11N11Thread|  M11N11Thread|     ABlockTransfer|       ABlockTransfer| ABlockTransfer| ABlockTransfer|      ABlockTransfer|     ABlockTransfer|      ABlockTransfer|     BBlockTransfer|       BBlockTransfer| BBlockTransfer| BBlockTransfer|      BBlockTransfer|     BBlockTransfer|      BBlockTransfer|     CThreadTransfer| CThreadTransfer|    CThreadTransfer|
+// ######|      Type|      Type|      Type|        Type|        |        |        | Elementwise| Elementwise| Elementwise| Spacialization|  Size| Block| Block| Block|   | ThreadM111| ThreadN111| Thread| ClusterM110Xs| ClusterN110Xs| ThreadSliceLengths| ThreadClusterLengths|  ThreadCluster|      SrcAccess|     SrcVectorTensor|    SrcVectorTensor|     DstVectorTensor| ThreadSliceLengths| ThreadClusterLengths|  ThreadCluster|      SrcAccess|     SrcVectorTensor|    SrcVectorTensor|     DstVectorTensor|        SrcDstAccess| SrcDstVectorDim| DstScalarPerVector|
+// ######|          |          |          |            |        |        |        |   Operation|   Operation|   Operation|               |      |      |      |      |   |           |           |       |              |              |        K0_M0_M1_K1|          K0_M0_M1_K1|   ArrangeOrder|          Order| Lengths_K0_M0_M1_K1| ContiguousDimOrder| Lengths_K0_M0_M1_K1|        K0_N0_N1_K1|          K0_N0_N1_K1|   ArrangeOrder|          Order| Lengths_K0_N0_N1_K1| ContiguousDimOrder| Lengths_K0_N0_N1_K1|               Order|                |                   |
+// ######|          |          |          |            |        |        |        |            |            |            |               |      |      |      |      |   |           |           |       |              |              |                   |                     |               |               |                    |                   |                    |                   |                     |               |               |                    |                   |                    |                    |                |                   |
+         < ADataType, BDataType, CDataType, AccDataType, ALayout, BLayout, CLayout,  AElementOp,  BElementOp,  CElementOp,    GemmDefault,   256,   128,   128,    16,  2,          1,          8,      8,       S<8, 8>,       S<4, 1>,      S<2, 1, 4, 2>,      S<8, 1,  32, 1>,  S<0, 3, 1, 2>,  S<0, 3, 1, 2>,       S<1, 1, 4, 1>,      S<0, 3, 1, 2>,       S<1, 1, 4, 2>,      S<2, 1, 4, 2>,       S<8, 1, 32, 1>,  S<0, 3, 1, 2>,  S<0, 3, 1, 2>,       S<1, 1, 4, 1>,      S<0, 3, 1, 2>,       S<1, 1, 4, 2>, S<0, 1, 2, 3, 4, 5>,               5,                  4>;
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::
+    ReferenceGemm<ADataType, BDataType, CDataType, AccDataType, AElementOp, BElementOp, CElementOp>;
+
+#include "run_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_gemm_example(argc, argv); }

example/01_gemm/gemm_xdl_fp16_f8.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "common.hpp"
+
+#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle.hpp"
+
+using ADataType        = ck::f8_t;
+using BDataType        = ck::half_t;
+using CDataType        = ck::half_t;
+using AccDataType      = float;
+using CShuffleDataType = ck::half_t;
+
+using ALayout = Row;
+using BLayout = Col;
+using CLayout = Row;
+
+using AElementOp = PassThrough;
+using BElementOp = PassThrough;
+using CElementOp = PassThrough;
+
+static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
+static constexpr auto LoopSched   = ck::make_default_loop_scheduler();
+static constexpr auto PipelineVer = ck::PipelineVersion::v1;
+using ComputeType                 = ck::half_t;
+
+// clang-format off
+using DeviceGemmInstance = ck::tensor_operation::device::DeviceGemm_Xdl_CShuffle
+// ######| ALayout| BLayout| CLayout|     AData|     BData|     CData|     AccData|         CShuffle|           A|           B|           C|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer|      Loop|    Pipeline|  ComputeType|
+// ######|        |        |        |      Type|      Type|      Type|        Type|         DataType| Elementwise| Elementwise| Elementwise| Spacialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|         _MBlock_MWaveMPerXdl| ScalarPerVector| Scheduler|     Version|             |
+// ######|        |        |        |          |          |          |            |                 |   Operation|   Operation|   Operation|               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|         _NBlock_NWaveNPerXdl|   _NWaveNPerXdl|          |            |             |
+// ######|        |        |        |          |          |          |            |                 |            |            |            |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |          |            |             |
+        < ALayout, BLayout, CLayout, ADataType, BDataType, CDataType, AccDataType, CShuffleDataType,  AElementOp,  BElementOp,  CElementOp,    GemmDefault,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     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,           1,           1,               S<1, 32, 1, 8>,               8,  LoopSched, PipelineVer, ComputeType>;
+// clang-format on
+
+using ReferenceGemmInstance = ck::tensor_operation::host::
+    ReferenceGemm<ADataType, BDataType, CDataType, AccDataType, AElementOp, BElementOp, CElementOp>;
+
+#include "run_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_gemm_example(argc, argv); }

example/13_pool2d_fwd/pool2d_fwd_common.hpp

@@ -39,17 +39,18 @@ bool pool_test(bool do_verification,
                ck::index_t Wi,
                ck::index_t window_stride_h,
                ck::index_t window_stride_w,
+               ck::index_t window_dilation_h,
+               ck::index_t window_dilation_w,
                ck::index_t in_left_pad_h,
                ck::index_t in_left_pad_w,
                ck::index_t in_right_pad_h,
                ck::index_t in_right_pad_w)
 {
     using DevicePoolFwdInstance =
-        ck::tensor_operation::device::DevicePool2dFwd_Input_N_Hi_Wi_C_Output_N_Ho_Wo_C<
-            InDataType,      // InDataType
-            OutDataType,     // OutDataType
-            IndexDataType,   // IndexDataType
-            ComputeDataType, // ComputeDataType
+        ck::tensor_operation::device::DevicePool2dFwd_NHWC_NHWC<InDataType,
+                                                                OutDataType,
+                                                                IndexDataType,
+                                                                ComputeDataType,
                                                                 ReduceOpId,
                                                                 OutputIndex,
                                                                 64, // BlockSize
@@ -57,13 +58,16 @@ bool pool_test(bool do_verification,
                                                                 1,  // ReduceKThreadClusterSize
                                                                 4,  // ReduceMThreadSliceSize
                                                                 1,  // ReduceKThreadSliceSize
-            4>; // InSrcOutDstVectorSize
+                                                                1>; // InSrcOutDstVectorSize
 
-    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
-    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;
+    const ck::index_t Ys = (Y - 1) * window_dilation_h + 1;
+    const ck::index_t Xs = (X - 1) * window_dilation_w + 1;
+    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Ys) / window_stride_h + 1;
+    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - Xs) / window_stride_w + 1;
 
     const std::vector<ck::index_t> window_spatial_lengths{Y, X};
     const std::vector<ck::index_t> window_strides{window_stride_h, window_stride_w};
+    const std::vector<ck::index_t> window_dilations{window_dilation_h, window_dilation_w};
     const std::vector<ck::index_t> input_left_pads{in_left_pad_h, in_left_pad_w};
     const std::vector<ck::index_t> input_right_pads{in_right_pad_h, in_right_pad_w};
 
@@ -123,6 +127,7 @@ bool pool_test(bool do_verification,
         {C * Ho * Wo, 1, Wo * C, C},
         {C * Ho * Wo, 1, Wo * C, C},
         window_strides,
+        window_dilations,
         input_left_pads,
         input_right_pads,
         {2, 3});
@@ -144,8 +149,8 @@ bool pool_test(bool do_verification,
 
     float gb_per_sec = num_btype / 1.E6 / ave_time;
 
-    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
-              << std::endl;
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec
+              << " GB / s " << std::endl;
 
     bool pass = true;
 
@@ -169,6 +174,7 @@ bool pool_test(bool do_verification,
                                                              out_indices_n_c_ho_wo_host,
                                                              window_spatial_lengths,
                                                              window_strides,
+                                                             window_dilations,
                                                              input_left_pads,
                                                              input_right_pads);
 

example/13_pool2d_fwd/pool2d_fwd_fp16.cpp

@@ -42,6 +42,8 @@ int main(int argc, char* argv[])
     ck::index_t Wi                = 71;
     ck::index_t window_stride_h   = 2;
     ck::index_t window_stride_w   = 2;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_h     = 1;
     ck::index_t in_left_pad_w     = 1;
     ck::index_t in_right_pad_h    = 1;
@@ -59,7 +61,7 @@ int main(int argc, char* argv[])
         init_method     = std::stoi(argv[2]);
         time_kernel     = static_cast<bool>(std::stoi(argv[3]));
     }
-    else if(argc == 16)
+    else if(argc == 18)
     {
         do_verification = std::stoi(argv[1]);
         init_method     = std::stoi(argv[2]);
@@ -73,17 +75,19 @@ int main(int argc, char* argv[])
         Wi                = std::stoi(argv[9]);
         window_stride_h   = std::stoi(argv[10]);
         window_stride_w   = std::stoi(argv[11]);
-        in_left_pad_h   = std::stoi(argv[12]);
-        in_left_pad_w   = std::stoi(argv[13]);
-        in_right_pad_h  = std::stoi(argv[14]);
-        in_right_pad_w  = std::stoi(argv[15]);
+        window_dilation_h = std::stoi(argv[12]);
+        window_dilation_w = std::stoi(argv[13]);
+        in_left_pad_h     = std::stoi(argv[14]);
+        in_left_pad_w     = std::stoi(argv[15]);
+        in_right_pad_h    = std::stoi(argv[16]);
+        in_right_pad_w    = std::stoi(argv[17]);
     }
     else
     {
         printf("arg1: verification (0=no, 1=yes)\n");
         printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
         printf("arg3: time kernel (0=no, 1=yes)\n");
-        printf("arg4 to 15: N, C, Y, X, Hi, Wi, Sy, Sx, LeftPy, LeftPx, RightPy, "
+        printf("arg4 to 15: N, C, Y, X, Hi, Wi, Sy, Sx, Dy, Dx, LeftPy, LeftPx, RightPy, "
                "RightPx\n");
         exit(0);
     }
@@ -107,6 +111,8 @@ int main(int argc, char* argv[])
                                        Wi,
                                        window_stride_h,
                                        window_stride_w,
+                                       window_dilation_h,
+                                       window_dilation_w,
                                        in_left_pad_h,
                                        in_left_pad_w,
                                        in_right_pad_h,

example/13_pool2d_fwd/pool2d_fwd_fp32.cpp

@@ -42,6 +42,8 @@ int main(int argc, char* argv[])
     ck::index_t Wi                = 71;
     ck::index_t window_stride_h   = 2;
     ck::index_t window_stride_w   = 2;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_h     = 1;
     ck::index_t in_left_pad_w     = 1;
     ck::index_t in_right_pad_h    = 1;
@@ -59,7 +61,7 @@ int main(int argc, char* argv[])
         init_method     = std::stoi(argv[2]);
         time_kernel     = static_cast<bool>(std::stoi(argv[3]));
     }
-    else if(argc == 16)
+    else if(argc == 18)
     {
         do_verification = std::stoi(argv[1]);
         init_method     = std::stoi(argv[2]);
@@ -73,17 +75,19 @@ int main(int argc, char* argv[])
         Wi                = std::stoi(argv[9]);
         window_stride_h   = std::stoi(argv[10]);
         window_stride_w   = std::stoi(argv[11]);
-        in_left_pad_h   = std::stoi(argv[12]);
-        in_left_pad_w   = std::stoi(argv[13]);
-        in_right_pad_h  = std::stoi(argv[14]);
-        in_right_pad_w  = std::stoi(argv[15]);
+        window_dilation_h = std::stoi(argv[12]);
+        window_dilation_w = std::stoi(argv[13]);
+        in_left_pad_h     = std::stoi(argv[14]);
+        in_left_pad_w     = std::stoi(argv[15]);
+        in_right_pad_h    = std::stoi(argv[16]);
+        in_right_pad_w    = std::stoi(argv[17]);
     }
     else
     {
         printf("arg1: verification (0=no, 1=yes)\n");
         printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
         printf("arg3: time kernel (0=no, 1=yes)\n");
-        printf("arg4 to 15: N, C, Y, X, Hi, Wi, Sy, Sx, LeftPy, LeftPx, RightPy, "
+        printf("arg4 to 15: N, C, Y, X, Hi, Wi, Sy, Sx, Dy, Dx, LeftPy, LeftPx, RightPy, "
                "RightPx\n");
         exit(0);
     }
@@ -107,6 +111,8 @@ int main(int argc, char* argv[])
                                        Wi,
                                        window_stride_h,
                                        window_stride_w,
+                                       window_dilation_h,
+                                       window_dilation_w,
                                        in_left_pad_h,
                                        in_left_pad_w,
                                        in_right_pad_h,

example/48_pool3d_fwd/pool3d_fwd_common.hpp

@@ -18,7 +18,45 @@
 #include "ck/library/utility/literals.hpp"
 #include "ck/library/reference_tensor_operation/cpu/reference_pool_fwd.hpp"
 
-template <typename InDataType,
+template <typename TensorLayout>
+std::vector<ck::index_t> f_tensor_strides_ncdhw(ck::index_t N_,
+                                                ck::index_t C_,
+                                                ck::index_t D,
+                                                ck::index_t H,
+                                                ck::index_t W,
+                                                TensorLayout layout)
+{
+    using namespace ck::literals;
+    (void)N_;
+    if constexpr(ck::is_same<decltype(layout), ck::tensor_layout::convolution::NCDHW>::value)
+        return {C_ * D * H * W, D * H * W, H * W, W, 1_uz};
+    else if constexpr(ck::is_same<decltype(layout), ck::tensor_layout::convolution::NDHWC>::value)
+        return {D * C_ * H * W, 1_uz, C_ * H * W, W * C_, C_};
+};
+
+template <typename TensorLayout>
+HostTensorDescriptor f_host_tensor_descriptor(std::size_t N_,
+                                              std::size_t C_,
+                                              std::size_t D,
+                                              std::size_t H,
+                                              std::size_t W,
+                                              TensorLayout layout)
+{
+    using namespace ck::literals;
+
+    if constexpr(ck::is_same<decltype(layout), ck::tensor_layout::convolution::NCDHW>::value)
+    {
+        return HostTensorDescriptor({N_, C_, D, H, W}, {C_ * D * H * W, D * H * W, H * W, W, 1_uz});
+    }
+    else if constexpr(ck::is_same<decltype(layout), ck::tensor_layout::convolution::NDHWC>::value)
+    {
+        return HostTensorDescriptor({N_, C_, D, H, W},
+                                    {D * C_ * H * W, 1_uz, C_ * H * W, W * C_, C_});
+    }
+};
+
+template <typename DevicePoolFwdInstance,
+          typename InDataType,
           typename OutDataType,
           typename ComputeDataType,
           typename IndexDataType,
@@ -40,6 +78,9 @@ bool pool3d_test(bool do_verification,
                  ck::index_t window_stride_d,
                  ck::index_t window_stride_h,
                  ck::index_t window_stride_w,
+                 ck::index_t window_dilation_d,
+                 ck::index_t window_dilation_h,
+                 ck::index_t window_dilation_w,
                  ck::index_t in_left_pad_d,
                  ck::index_t in_left_pad_h,
                  ck::index_t in_left_pad_w,
@@ -47,53 +88,21 @@ bool pool3d_test(bool do_verification,
                  ck::index_t in_right_pad_h,
                  ck::index_t in_right_pad_w)
 {
-    using DevicePoolFwdInstance =
-        ck::tensor_operation::device::DevicePool3dFwd_Input_N_Di_Hi_Wi_C_Output_N_Do_Ho_Wo_C<
-            InDataType,      // InDataType
-            OutDataType,     // OutDataType
-            IndexDataType,   // IndexDataType
-            ComputeDataType, // ComputeDataType
-            ReduceOpId,
-            OutputIndex,
-            64, // BlockSize
-            64, // ReduceMThreadClusterSize
-            1,  // ReduceKThreadClusterSize
-            4,  // ReduceMThreadSliceSize
-            1,  // ReduceKThreadSliceSize
-            4>; // InSrcOutDstVectorSize
-
-    const ck::index_t Do = (Di + in_left_pad_d + in_right_pad_d - Z) / window_stride_d + 1;
-    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
-    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;
+    const ck::index_t Zs = (Z - 1) * window_dilation_d + 1;
+    const ck::index_t Ys = (Y - 1) * window_dilation_h + 1;
+    const ck::index_t Xs = (X - 1) * window_dilation_w + 1;
+    const ck::index_t Do = (Di + in_left_pad_d + in_right_pad_d - Zs) / window_stride_d + 1;
+    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Ys) / window_stride_h + 1;
+    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - Xs) / window_stride_w + 1;
 
     const std::vector<ck::index_t> window_spatial_lengths{Z, Y, X};
     const std::vector<ck::index_t> window_strides{
         window_stride_d, window_stride_h, window_stride_w};
+    const std::vector<ck::index_t> window_dilations{
+        window_dilation_d, window_dilation_h, window_dilation_w};
     const std::vector<ck::index_t> input_left_pads{in_left_pad_d, in_left_pad_h, in_left_pad_w};
     const std::vector<ck::index_t> input_right_pads{in_right_pad_d, in_right_pad_h, in_right_pad_w};
 
-    // tensor layout
-    auto f_host_tensor_descriptor = [](std::size_t N_,
-                                       std::size_t C_,
-                                       std::size_t D,
-                                       std::size_t H,
-                                       std::size_t W,
-                                       auto layout) {
-        using namespace ck::literals;
-
-        if constexpr(ck::is_same<decltype(layout), ck::tensor_layout::convolution::NCDHW>::value)
-        {
-            return HostTensorDescriptor({N_, C_, D, H, W},
-                                        {C_ * D * H * W, D * H * W, H * W, W, 1_uz});
-        }
-        else if constexpr(ck::is_same<decltype(layout),
-                                      ck::tensor_layout::convolution::NDHWC>::value)
-        {
-            return HostTensorDescriptor({N_, C_, D, H, W},
-                                        {D * C_ * H * W, 1_uz, C_ * H * W, W * C_, C_});
-        }
-    };
-
     Tensor<InDataType> in_n_c_di_hi_wi(f_host_tensor_descriptor(N, C, Di, Hi, Wi, InLayout{}));
     Tensor<OutDataType> out_n_c_do_ho_wo_host(
         f_host_tensor_descriptor(N, C, Do, Ho, Wo, OutLayout{}));
@@ -126,10 +135,11 @@ bool pool3d_test(bool do_verification,
         {N, C, Di, Hi, Wi},
         {Z, Y, X},
         {N, C, Do, Ho, Wo},
-        {Di * C * Hi * Wi, 1, C * Hi * Wi, Wi * C, C},
-        {Do * C * Ho * Wo, 1, C * Ho * Wo, Wo * C, C},
-        {Do * C * Ho * Wo, 1, C * Ho * Wo, Wo * C, C},
+        f_tensor_strides_ncdhw(N, C, Di, Hi, Wi, InLayout{}),
+        f_tensor_strides_ncdhw(N, C, Do, Ho, Wo, OutLayout{}),
+        f_tensor_strides_ncdhw(N, C, Do, Ho, Wo, OutLayout{}),
         window_strides,
+        window_dilations,
         input_left_pads,
         input_right_pads,
         {2, 3, 4});
@@ -165,6 +175,7 @@ bool pool3d_test(bool do_verification,
                                                              out_indices_n_c_do_ho_wo_host,
                                                              window_spatial_lengths,
                                                              window_strides,
+                                                             window_dilations,
                                                              input_left_pads,
                                                              input_right_pads);
 

example/48_pool3d_fwd/pool3d_fwd_fp16.cpp

@@ -27,6 +27,20 @@ static constexpr auto ReduceOpId = ck::ReduceTensorOp::AVG;
 static constexpr bool OutputIndex  = false;
 static constexpr bool PropagateNan = false;
 
+using DevicePoolFwdInstance =
+    ck::tensor_operation::device::DevicePool3dFwd_NDHWC_NDHWC<InDataType,
+                                                              OutDataType,
+                                                              IndexDataType,
+                                                              ComputeDataType,
+                                                              ReduceOpId,
+                                                              OutputIndex,
+                                                              64, // BlockSize
+                                                              64, // ReduceMThreadClusterSize
+                                                              1,  // ReduceKThreadClusterSize
+                                                              1,  // ReduceMThreadSliceSize
+                                                              1,  // ReduceKThreadSliceSize
+                                                              1>; // InSrcOutDstVectorSize
+
 int main()
 {
     bool do_verification = true;
@@ -44,6 +58,9 @@ int main()
     ck::index_t window_stride_d   = 2;
     ck::index_t window_stride_h   = 2;
     ck::index_t window_stride_w   = 2;
+    ck::index_t window_dilation_d = 1;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_d     = 1;
     ck::index_t in_left_pad_h     = 1;
     ck::index_t in_left_pad_w     = 1;
@@ -51,7 +68,8 @@ int main()
     ck::index_t in_right_pad_h    = 1;
     ck::index_t in_right_pad_w    = 1;
 
-    bool pass = pool3d_test<InDataType,
+    bool pass = pool3d_test<DevicePoolFwdInstance,
+                            InDataType,
                             OutDataType,
                             ComputeDataType,
                             IndexDataType,
@@ -72,6 +90,9 @@ int main()
                                          window_stride_d,
                                          window_stride_h,
                                          window_stride_w,
+                                         window_dilation_d,
+                                         window_dilation_h,
+                                         window_dilation_w,
                                          in_left_pad_d,
                                          in_left_pad_h,
                                          in_left_pad_w,

example/49_maxpool2d_bwd/maxpool2d_bwd_bf16.cpp

@@ -32,6 +32,8 @@ int main()
     ck::index_t Wi                = 32;
     ck::index_t window_stride_h   = 1;
     ck::index_t window_stride_w   = 1;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_h     = 0;
     ck::index_t in_left_pad_w     = 0;
     ck::index_t in_right_pad_h    = 0;
@@ -53,6 +55,8 @@ int main()
                                                Wi,
                                                window_stride_h,
                                                window_stride_w,
+                                               window_dilation_h,
+                                               window_dilation_w,
                                                in_left_pad_h,
                                                in_left_pad_w,
                                                in_right_pad_h,

example/49_maxpool2d_bwd/maxpool2d_bwd_common.hpp

@@ -36,6 +36,8 @@ bool maxpool_bwd_test(bool do_verification,
                       ck::index_t Wi,
                       ck::index_t window_stride_h,
                       ck::index_t window_stride_w,
+                      ck::index_t window_dilation_h,
+                      ck::index_t window_dilation_w,
                       ck::index_t in_left_pad_h,
                       ck::index_t in_left_pad_w,
                       ck::index_t in_right_pad_h,
@@ -44,13 +46,12 @@ bool maxpool_bwd_test(bool do_verification,
     using PassThrough = ck::tensor_operation::element_wise::PassThrough;
 
     using DevicePoolFwdInstance =
-        ck::tensor_operation::device::DevicePool2dFwd_Input_N_Hi_Wi_C_Output_N_Ho_Wo_C<
-            InDataType,      // InDataType
+        ck::tensor_operation::device::DevicePool2dFwd_NHWC_NHWC<InDataType,      // InDataType
                                                                 OutDataType,     // OutDataType
                                                                 IndexDataType,   // IndexDataType
                                                                 ComputeDataType, // ComputeDataType
                                                                 ck::ReduceTensorOp::MAX,
-            true, // OutputIndex
+                                                                true,
                                                                 64, // BlockSize
                                                                 64, // ReduceMThreadClusterSize
                                                                 1,  // ReduceKThreadClusterSize
@@ -61,11 +62,14 @@ bool maxpool_bwd_test(bool do_verification,
     using DeviceMaxPoolBwdInstance = ck::tensor_operation::device::
         DeviceIndexPoolBwdImpl<DOutDataType, IndexDataType, DInDataType, 4>;
 
-    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
-    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;
+    const ck::index_t Ys = (Y - 1) * window_dilation_h + 1;
+    const ck::index_t Xs = (X - 1) * window_dilation_w + 1;
+    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Ys) / window_stride_h + 1;
+    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - Xs) / window_stride_w + 1;
 
     const std::vector<ck::index_t> window_spatial_lengths{Y, X};
     const std::vector<ck::index_t> window_strides{window_stride_h, window_stride_w};
+    const std::vector<ck::index_t> window_dilations{window_dilation_h, window_dilation_w};
     const std::vector<ck::index_t> input_left_pads{in_left_pad_h, in_left_pad_w};
     const std::vector<ck::index_t> input_right_pads{in_right_pad_h, in_right_pad_w};
 
@@ -128,6 +132,7 @@ bool maxpool_bwd_test(bool do_verification,
         {C * Ho * Wo, 1, Wo * C, C},
         {C * Ho * Wo, 1, Wo * C, C},
         window_strides,
+        window_dilations,
         input_left_pads,
         input_right_pads,
         {2, 3});
@@ -191,6 +196,7 @@ bool maxpool_bwd_test(bool do_verification,
                                                                      indices_n_c_ho_wo_host,
                                                                      window_spatial_lengths,
                                                                      window_strides,
+                                                                     window_dilations,
                                                                      input_left_pads,
                                                                      input_right_pads);
         ref_pooling_fwd_invoker.Run(ref_pooling_fwd_argument);

example/49_maxpool2d_bwd/maxpool2d_bwd_fp16.cpp

@@ -32,6 +32,8 @@ int main()
     ck::index_t Wi                = 32;
     ck::index_t window_stride_h   = 1;
     ck::index_t window_stride_w   = 1;
+    ck::index_t window_dilation_h = 1;
+    ck::index_t window_dilation_w = 1;
     ck::index_t in_left_pad_h     = 0;
     ck::index_t in_left_pad_w     = 0;
     ck::index_t in_right_pad_h    = 0;
@@ -53,6 +55,8 @@ int main()
                                                Wi,
                                                window_stride_h,
                                                window_stride_w,
+                                               window_dilation_h,
+                                               window_dilation_w,
                                                in_left_pad_h,
                                                in_left_pad_w,
                                                in_right_pad_h,