Merge remote-tracking branch 'origin/develop' into vpietila/ggemm-profiling

example/README.md

+[Back to the main page](../README.md)
+# Composable Kernel examples
\ No newline at end of file

example/ck_tile/03_gemm/gemm_basic.cpp

@@ -92,6 +92,11 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s)
     const dim3 grids      = Kernel::GridSize(args.M, args.N, args.kbatch);
     constexpr dim3 blocks = Kernel::BlockSize();
 
+    if(!Kernel::IsSupportedArgument(kargs))
+    {
+        throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n");
+    }
+
     if(s.log_level_ > 0)
     {
         std::cout << "Launching kernel with args:"

example/ck_tile/03_gemm/universal_gemm.cpp

@@ -119,6 +119,11 @@ float gemm_calc(const gemm_basic_args& args, const ck_tile::stream_config& s)
         const dim3 grids      = Kernel::GridSize(args.M, args.N, args.kbatch);
         constexpr dim3 blocks = Kernel::BlockSize();
 
+        if(!Kernel::IsSupportedArgument(kargs))
+        {
+            throw std::runtime_error("Wrong! Arguments not supported! Skipping gemm!\n");
+        }
+
         if(s.log_level_ > 0)
         {
             std::cout << "Launching kernel with args:"

example/ck_tile/17_grouped_gemm/CMakeLists.txt

+add_executable(tile_example_grouped_gemm EXCLUDE_FROM_ALL grouped_gemm.cpp)
+

example/ck_tile/17_grouped_gemm/README.md

+# Grouped CShuffle GEMM
+
+This folder contains example for Grouped GEMM using ck_tile tile-programming implementation. Currently, it only supports the basic feature of the CK Tile GEMM, but creates the placeholders for the future support on different GEMM pipeline and different GEMM modules. In the near future, we will gradually migrate all the GEMM features from old CK to CK Tile.
+
+## build
+```
+# in the root of ck_tile
+mkdir build && cd build
+# you can replace <arch> with the appropriate architecture (for example gfx90a or gfx942) or leave it blank
+sh ../script/cmake-ck-dev.sh  ../ <arch>
+# The basic pipeline method on the gemm calculation
+make tile_example_grouped_gemm -j
+```
+This will result in an executable `build/bin/tile_example_grouped_gemm`
+
+## example
+```
+args:
+   -a_layout    Tensor A layout (default:R)
+   -b_layout    Tensor B layout (default:R)
+   -c_layout    Tensor C layout (default:R)
+          -v    0. No validation, 1. Validation on CPU
+     -warmup    number of iterations before benchmark the kernel (default:10)
+     -repeat    number of iterations to benchmark the kernel (default:100)
+```

example/ck_tile/17_grouped_gemm/grouped_gemm.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <hip/hip_runtime.h>
+
+#include <cstring>
+#include <iostream>
+#include <ostream>
+#include <string>
+#include <tuple>
+#include <memory>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/epilogue.hpp"
+#include "ck_tile/ops/gemm.hpp"
+#include "ck_tile/host.hpp"
+#include "grouped_gemm.hpp"
+#include "utils.hpp"
+
+namespace {
+
+struct GroupedGemmKernelParam
+{
+    static const bool kPadM        = false;
+    static const bool kPadN        = false;
+    static const bool kPadK        = false;
+    static const bool kTilePermute = false;
+
+    static const ck_tile::index_t kOutputRank = 2;
+
+    static const int kBlockPerCu         = 1;
+    static const ck_tile::index_t M_Tile = 128;
+    static const ck_tile::index_t N_Tile = 128;
+    static const ck_tile::index_t K_Tile = 32;
+
+    static const ck_tile::index_t M_Warp = 2;
+    static const ck_tile::index_t N_Warp = 2;
+    static const ck_tile::index_t K_Warp = 1;
+
+    static const ck_tile::index_t M_Warp_Tile = 32;
+    static const ck_tile::index_t N_Warp_Tile = 32;
+    static const ck_tile::index_t K_Warp_Tile = 8;
+};
+
+using CodegenGemmShape =
+    ck_tile::TileGemmShape<ck_tile::sequence<GroupedGemmKernelParam::M_Tile,
+                                             GroupedGemmKernelParam::N_Tile,
+                                             GroupedGemmKernelParam::K_Tile>,
+                           ck_tile::sequence<GroupedGemmKernelParam::M_Warp,
+                                             GroupedGemmKernelParam::N_Warp,
+                                             GroupedGemmKernelParam::K_Warp>,
+                           ck_tile::sequence<GroupedGemmKernelParam::M_Warp_Tile,
+                                             GroupedGemmKernelParam::N_Warp_Tile,
+                                             GroupedGemmKernelParam::K_Warp_Tile>>;
+
+using TilePartitioner = ck_tile::GemmTile1DPartitioner<CodegenGemmShape>;
+
+template <typename CLayout>
+using GemmEpilogue = std::conditional_t<
+    std::is_same_v<CLayout, ck_tile::tensor_layout::gemm::ColumnMajor>,
+    ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<AccDataType,
+                                                               CDataType,
+                                                               GroupedGemmKernelParam::kPadM,
+                                                               GroupedGemmKernelParam::kPadN,
+                                                               GroupedGemmKernelParam::kTilePermute,
+                                                               GroupedGemmKernelParam::kOutputRank,
+                                                               1,
+                                                               0,
+                                                               TilePartitioner::MPerBlock,
+                                                               TilePartitioner::NPerBlock>>,
+    ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<AccDataType,
+                                                                 CDataType,
+                                                                 GroupedGemmKernelParam::kPadM,
+                                                                 GroupedGemmKernelParam::kPadN>>>;
+
+template <typename ALayout, typename BLayout, typename CLayout>
+using CodegenGemmTraits = ck_tile::TileGemmTraits<GroupedGemmKernelParam::kPadM,
+                                                  GroupedGemmKernelParam::kPadN,
+                                                  GroupedGemmKernelParam::kPadK,
+                                                  ALayout,
+                                                  BLayout,
+                                                  CLayout>;
+
+template <typename ALayout, typename BLayout, typename CLayout>
+using CodegenPipelineProblem =
+    ck_tile::GemmPipelineProblem<ADataType,
+                                 BDataType,
+                                 AccDataType,
+                                 CodegenGemmShape,
+                                 CodegenGemmTraits<ALayout, BLayout, CLayout>>;
+
+using CodegenGemmPolicy = ck_tile::UniversalGemmPipelineAgBgCrPolicy;
+
+template <typename ALayout, typename BLayout, typename CLayout>
+using CodegenGemmPipeline =
+    ck_tile::GemmPipelineAGmemBGmemCRegV1<CodegenPipelineProblem<ALayout, BLayout, CLayout>,
+                                          CodegenGemmPolicy>;
+
+template <typename ALayout, typename BLayout, typename CLayout>
+using Kernel = ck_tile::GroupedGemmKernel<TilePartitioner,
+                                          CodegenGemmPipeline<ALayout, BLayout, CLayout>,
+                                          GemmEpilogue<CLayout>>;
+}; // namespace
+
+std::size_t GetWorkspaceSize(const std::vector<grouped_gemm_kargs>& gemm_descs)
+{
+    return ::Kernel<std::nullptr_t, std::nullptr_t, std::nullptr_t>::GetWorkSpaceSize(gemm_descs);
+}
+
+template <typename ALayout, typename BLayout, typename CLayout>
+float grouped_gemm(const std::vector<grouped_gemm_kargs>& gemm_descs,
+                   const ck_tile::stream_config& s,
+                   void* p_workspace_)
+{
+    using GroupedGemmKernel = ::Kernel<ALayout, BLayout, CLayout>;
+
+    auto arguments = GroupedGemmKernel::MakeKargs(gemm_descs);
+
+    const dim3 grids      = GroupedGemmKernel::GridSize(gemm_descs);
+    constexpr dim3 blocks = GroupedGemmKernel::BlockSize();
+
+    ck_tile::hip_check_error(hipMemcpyWithStream(
+        p_workspace_,
+        arguments.data(),
+        arguments.size() * sizeof(typename GroupedGemmKernel::GemmTransKernelArg),
+        hipMemcpyHostToDevice,
+        s.stream_id_));
+
+    if(s.log_level_ > 0)
+    {
+        std::cout << "Launching kernel with args:"
+                  << " grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                  << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
+                  << std::endl;
+    }
+
+    float ave_time =
+        ck_tile::launch_kernel(s,
+                               ck_tile::make_kernel<blocks.x, GroupedGemmKernelParam::kBlockPerCu>(
+                                   GroupedGemmKernel{},
+                                   grids,
+                                   blocks,
+                                   0,
+                                   ck_tile::cast_pointer_to_constant_address_space(p_workspace_),
+                                   gemm_descs.size()));
+    return ave_time;
+}
+
+#include "run_grouped_gemm_example.inc"
+
+int main(int argc, char* argv[]) { return !run_grouped_gemm_example(argc, argv); }

example/ck_tile/17_grouped_gemm/grouped_gemm.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <string>
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp"
+
+template <typename DataType>
+struct GemmBasicTypeConfig;
+
+template <>
+struct GemmBasicTypeConfig<ck_tile::half_t>
+{
+    using ADataType   = ck_tile::half_t;
+    using BDataType   = ck_tile::half_t;
+    using CDataType   = ck_tile::half_t;
+    using AccDataType = float;
+};
+
+using Types = GemmBasicTypeConfig<ck_tile::half_t>;
+
+// Specific type aliases for easy access
+using ADataType   = Types::ADataType;
+using BDataType   = Types::BDataType;
+using AccDataType = Types::AccDataType;
+using CDataType   = Types::CDataType;
+
+using grouped_gemm_kargs = ck_tile::GroupedGemmHostArgs;
+
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("a_layout", "R", "A tensor data layout - Row by default")
+        .insert("b_layout", "R", "B tensor data layout - Row by default")
+        .insert("c_layout", "R", "C tensor data layout - Row by default")
+        .insert("validate", "1", "0. No validation, 1. Validation on CPU")
+        .insert("warmup", "10", "number of iterations before benchmark the kernel")
+        .insert("repeat", "100", "number of iterations to benchmark the kernel")
+        .insert("group_count", "16", "group count");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}
+
+std::size_t GetWorkspaceSize(const std::vector<grouped_gemm_kargs>& gemm_descs);
+
+float grouped_gemm_calc(const std::vector<grouped_gemm_kargs>& gemm_descs,
+                        const ck_tile::stream_config& s,
+                        void* p_workspace_);

example/ck_tile/17_grouped_gemm/run_grouped_gemm_example.inc

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+template <typename ALayout, typename BLayout, typename CLayout>
+float invoke_gemm(int n_warmup,
+                  int n_repeat,
+                  int group_count,
+                  const std::vector<grouped_gemm_kargs>& args)
+{
+
+    ck_tile::DeviceMem gemm_workspace;
+    gemm_workspace.Realloc(GetWorkspaceSize(args));
+
+    float ave_time = grouped_gemm<ALayout, BLayout, CLayout>(
+        args,
+        ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat},
+        gemm_workspace.GetDeviceBuffer());
+
+    std::string op_name{"Grouped Gemm"};
+
+    std::size_t flop = 0, num_btype = 0;
+    for(int j = 0; j < group_count; ++j)
+    {
+        flop += std::size_t(2) * args[j].M * args[j].N * args[j].K;
+
+        num_btype += sizeof(ADataType) * args[j].M * args[j].K +
+                     sizeof(BDataType) * args[j].K * args[j].N +
+                     sizeof(CDataType) * args[j].M * args[j].N;
+    }
+
+    float tflops     = static_cast<float>(flop) / 1.E9 / ave_time;
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, "
+              << gb_per_sec << " GB/s, " << op_name << std::endl;
+
+    return ave_time;
+}
+
+template <typename ALayout, typename BLayout, typename CLayout>
+int run_grouped_gemm_example_with_layouts(int argc,
+                                          char* argv[],
+                                          const ALayout a_layout                  = ALayout{},
+                                          const BLayout b_layout                  = BLayout{},
+                                          [[maybe_unused]] const CLayout c_layout = CLayout{})
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+
+    if(!result)
+    {
+        return -1;
+    };
+
+    const int group_count = arg_parser.get_int("group_count");
+    const int repeat      = arg_parser.get_int("repeat");
+    const int warmup      = arg_parser.get_int("warmup");
+
+    std::vector<ck_tile::index_t> Ms;
+    std::vector<ck_tile::index_t> Ns;
+    std::vector<ck_tile::index_t> Ks;
+    std::vector<ck_tile::index_t> stride_As;
+    std::vector<ck_tile::index_t> stride_Bs;
+    std::vector<ck_tile::index_t> stride_Cs;
+
+    for(int i = 0; i < group_count; i++)
+    {
+        Ms.push_back(256 + 256 * i);
+        Ns.push_back(128 + 128 * i);
+        Ks.push_back(128 + 64 * i);
+
+        stride_As.push_back(Ks[i]);
+        stride_Bs.push_back(Ks[i]);
+        stride_Cs.push_back(Ns[i]);
+    }
+
+    std::vector<ck_tile::HostTensor<ADataType>> a_m_k_tensors;
+    std::vector<ck_tile::HostTensor<BDataType>> b_k_n_tensors;
+    std::vector<ck_tile::HostTensor<CDataType>> c_m_n_tensors;
+
+    a_m_k_tensors.reserve(group_count);
+    b_k_n_tensors.reserve(group_count);
+    c_m_n_tensors.reserve(group_count);
+
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> a_m_k_dev_buf;
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> b_k_n_dev_buf;
+    std::vector<std::unique_ptr<ck_tile::DeviceMem>> c_m_n_dev_buf;
+
+    a_m_k_dev_buf.reserve(group_count);
+    b_k_n_dev_buf.reserve(group_count);
+    c_m_n_dev_buf.reserve(group_count);
+
+    std::vector<grouped_gemm_kargs> gemm_descs;
+    gemm_descs.reserve(group_count);
+
+    for(int i = 0; i < group_count; ++i)
+    {
+        const ck_tile::index_t M = Ms[i];
+        const ck_tile::index_t N = Ns[i];
+        const ck_tile::index_t K = Ks[i];
+
+        stride_As[i] = f_get_default_stride(M, N, stride_As[i], a_layout);
+        stride_Bs[i] = f_get_default_stride(K, N, stride_Bs[i], b_layout);
+        stride_Cs[i] = f_get_default_stride(M, N, stride_Cs[i], CLayout{});
+
+        a_m_k_tensors.push_back(
+            ck_tile::HostTensor<ADataType>(f_host_tensor_descriptor(M, K, stride_As[i], a_layout)));
+        b_k_n_tensors.push_back(
+            ck_tile::HostTensor<BDataType>(f_host_tensor_descriptor(K, N, stride_Bs[i], b_layout)));
+        c_m_n_tensors.push_back(ck_tile::HostTensor<CDataType>(
+            f_host_tensor_descriptor(M, N, stride_Cs[i], CLayout{})));
+
+        std::cout << "gemm[" << i << "]"
+                  << " a_m_k: " << a_m_k_tensors[i].mDesc << " b_k_n: " << b_k_n_tensors[i].mDesc
+                  << " c_m_n: " << c_m_n_tensors[i].mDesc << std::endl;
+
+        ck_tile::FillUniformDistribution<ADataType>{-5.f, 5.f}(a_m_k_tensors[i]);
+        ck_tile::FillUniformDistribution<BDataType>{-5.f, 5.f}(b_k_n_tensors[i]);
+
+        a_m_k_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(
+            a_m_k_tensors[i].get_element_space_size_in_bytes()));
+        b_k_n_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(
+            b_k_n_tensors[i].get_element_space_size_in_bytes()));
+        c_m_n_dev_buf.push_back(std::make_unique<ck_tile::DeviceMem>(
+            c_m_n_tensors[i].get_element_space_size_in_bytes()));
+
+        a_m_k_dev_buf[i]->ToDevice(a_m_k_tensors[i].data());
+        b_k_n_dev_buf[i]->ToDevice(b_k_n_tensors[i].data());
+        c_m_n_dev_buf[i]->SetZero();
+        c_m_n_tensors[i].SetZero();
+
+        const void* p_a = a_m_k_dev_buf[i]->GetDeviceBuffer();
+        const void* p_b = b_k_n_dev_buf[i]->GetDeviceBuffer();
+        void* p_c       = c_m_n_dev_buf[i]->GetDeviceBuffer();
+
+        gemm_descs.push_back({p_a, p_b, p_c, M, N, K, stride_As[i], stride_Bs[i], stride_Cs[i]});
+    }
+
+    invoke_gemm<ALayout, BLayout, CLayout>(warmup, repeat, group_count, gemm_descs);
+
+    for(int i = 0; i < group_count; i++)
+    {
+        c_m_n_dev_buf[i]->FromDevice(c_m_n_tensors[i].data());
+    }
+
+    bool pass{true};
+    if(arg_parser.get_int("validate"))
+    {
+        for(int i = 0; i < group_count; ++i)
+        {
+            ck_tile::HostTensor<CDataType> c_m_n_host_ref(
+                f_host_tensor_descriptor(Ms[i], Ns[i], stride_Cs[i], CLayout{}));
+            c_m_n_host_ref.SetZero();
+            ck_tile::reference_gemm<ADataType, BDataType, AccDataType, CDataType>(
+                a_m_k_tensors[i], b_k_n_tensors[i], c_m_n_host_ref);
+            pass &= ck_tile::check_err(c_m_n_tensors[i], c_m_n_host_ref);
+        }
+        std::cout << "The CPU veification result is:" << (pass ? "correct" : "fail") << std::endl;
+    }
+
+    return pass;
+}
+
+int run_grouped_gemm_example(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+    {
+        return -1;
+    }
+
+    const std::string a_layout = arg_parser.get_str("a_layout");
+    const std::string b_layout = arg_parser.get_str("b_layout");
+
+    using Row = ck_tile::tensor_layout::gemm::RowMajor;
+    using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
+
+    if(a_layout == "R" && b_layout == "C")
+    {
+        return run_grouped_gemm_example_with_layouts(argc, argv, Row{}, Col{}, Row{});
+    }
+    else if(a_layout == "R" && b_layout == "R")
+    {
+        return run_grouped_gemm_example_with_layouts(argc, argv, Row{}, Row{}, Row{});
+    }
+    else
+    {
+        throw std::runtime_error("Unsupported data layout configuration for A,B and C tensors!");
+    }
+}

example/ck_tile/17_grouped_gemm/utils.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+template <typename TLayout>
+constexpr auto
+f_host_tensor_descriptor(std::size_t row, std::size_t col, std::size_t stride, TLayout layout)
+{
+    using namespace ck_tile::literals;
+
+    if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)
+    {
+        return ck_tile::HostTensorDescriptor({row, col}, {stride, 1_uz});
+    }
+    else
+    {
+        return ck_tile::HostTensorDescriptor({row, col}, {1_uz, stride});
+    }
+}
+template <typename TLayout>
+constexpr auto
+f_get_default_stride(std::size_t row, std::size_t col, std::size_t stride, TLayout layout)
+{
+    if(stride == 0)
+    {
+        if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)
+        {
+            return col;
+        }
+        else
+        {
+            return row;
+        }
+    }
+    else
+        return stride;
+}

example/ck_tile/CMakeLists.txt

@@ -16,3 +16,4 @@ add_subdirectory(13_moe_sorting)
 add_subdirectory(14_moe_smoothquant)
 add_subdirectory(15_fused_moe)
 add_subdirectory(16_batched_gemm)
+add_subdirectory(17_grouped_gemm)

include/ck/README.md

+[Back to the main page](../../README.md)
+# Composable Kernel supported operations
+## Supported device operations
+* [Average pooling]()
+* [Batched contraction]()
+* [Batched gemm]()
+* [Batchnorm]()
+* [CGEMM]()
+* [Contraction]()
+* [Convolution]()
+* [Image to Column and Column to Image]()
+* [Elementwise]()
+* [GEMM]()
+* [Max pooling]()
+* [Reduce]()
+* [Normalization]()
+* [Permute]()
+* [Put]()
+* [Softmax]()

include/ck/library/utility/host_tensor.hpp

@@ -326,7 +326,7 @@ struct Tensor
 
     std::size_t GetElementSpaceSizeInBytes() const { return sizeof(T) * GetElementSpaceSize(); }
 
-    void SetZero() { ck::ranges::fill<T>(mData, 0); }
+    void SetZero() { ck::ranges::fill<T>(mData, T{0}); }
 
     template <typename F>
     void ForEach_impl(F&& f, std::vector<size_t>& idx, size_t rank)

include/ck/library/utility/host_tensor_generator.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -37,7 +37,7 @@ struct GeneratorTensor_1<ck::half_t>
     float value = 1.0;
 
     template <typename... Is>
-    ck::bhalf_t operator()(Is...)
+    ck::half_t operator()(Is...)
     {
         return ck::type_convert<ck::half_t>(value);
     }
@@ -62,7 +62,7 @@ struct GeneratorTensor_1<ck::f8_t>
     float value = 1.0;
 
     template <typename... Is>
-    ck::bhalf_t operator()(Is...)
+    ck::f8_t operator()(Is...)
     {
         return ck::type_convert<ck::f8_t>(value);
     }
@@ -256,14 +256,33 @@ struct GeneratorTensor_Checkboard
     }
 };
 
-template <ck::index_t Dim>
+/**
+ * @brief Is used to generate sequential values based on the specified dimension.
+ *
+ * @tparam T The type of the tensor values.
+ * @tparam Dim The specific dimension used for generation.
+ *
+ * GeneratorTensor_Sequential<1>{} will generate the following values for a 3x3 tensor:
+ *
+ * 0 1 2
+ * 0 1 2
+ * 0 1 2
+ *
+ * Essentially, the values generated are logical coordinates of the generated element that
+ * correspond to dimension Dim. E.g. for 2-dimensional tensor and Dim=1, the values are the column
+ * indices.
+ *
+ */
+template <typename T, ck::index_t Dim>
 struct GeneratorTensor_Sequential
 {
     template <typename... Ts>
-    float operator()(Ts... Xs) const
+    T operator()(Ts... Xs) const
     {
         std::array<ck::index_t, sizeof...(Ts)> dims = {{static_cast<ck::index_t>(Xs)...}};
-        return dims[Dim];
+
+        float tmp = dims[Dim];
+        return ck::type_convert<T>(tmp);
     }
 };
 

include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2023-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -106,89 +106,35 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
     static constexpr auto I3           = Number<3>{};
     static constexpr index_t KPerBlock = K0PerBlock * K1;
 
-    static constexpr auto transform_conv_to_gemm =
-        TransformConvBwdDataToGemm_v1<NDimSpatial,
-                                      ConvBackwardDataSpecialization,
-                                      K1,
-                                      K1,
-                                      MPerBlock,
-                                      NPerBlock,
-                                      KPerBlock,
-                                      true /* DoPadGemmM */,
-                                      true /* DoPadGemmN */>{};
-
-    static auto GetDummyABDsEGridDescriptor()
-    {
-        const std::array<index_t, NDimSpatial + 3> dummy_tensor_lengths = {1};
-        const std::array<index_t, NDimSpatial + 3> dummy_tensor_strides = {1};
-        const std::array<index_t, NDimSpatial> dummy_spatial_lengths    = {1};
-
-        const auto a_grid_desc_ak0_m_ak1 =
-            transform_conv_to_gemm.template MakeADescriptor_AK0_M_AK1<ALayout>(
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths);
-
-        const auto b_grid_desc_bk0_n_bk1 =
-            transform_conv_to_gemm.template MakeBDescriptor_BK0_N_BK1<BLayout>(
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths);
-
-        const auto ds_grid_desc_m_n = generate_tuple(
-            [&](auto i) {
-                using DLayout = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
-
-                return transform_conv_to_gemm.template MakeCDescriptor_M_N<DLayout>(
-                    dummy_tensor_lengths,
-                    dummy_tensor_strides,
-                    dummy_tensor_lengths,
-                    dummy_tensor_strides,
-                    dummy_tensor_lengths,
-                    dummy_tensor_strides,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths);
-            },
-            Number<NumDTensor>{});
-
-        const auto e_grid_desc_m_n =
-            transform_conv_to_gemm.template MakeCDescriptor_M_N<ELayout>(dummy_tensor_lengths,
-                                                                         dummy_tensor_strides,
-                                                                         dummy_tensor_lengths,
-                                                                         dummy_tensor_strides,
-                                                                         dummy_tensor_lengths,
-                                                                         dummy_tensor_strides,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths);
+    using ConvToGemmBwdDataTransform = TransformConvBwdDataToGemm_v1<NDimSpatial,
+                                                                     ConvBackwardDataSpecialization,
+                                                                     K1,
+                                                                     K1,
+                                                                     MPerBlock,
+                                                                     NPerBlock,
+                                                                     KPerBlock,
+                                                                     true /* DoPadGemmM */,
+                                                                     true /* DoPadGemmN */,
+                                                                     ALayout,
+                                                                     BLayout,
+                                                                     ELayout>;
 
+    static auto
+    GetDummyABDsEGridDescriptor(const ConvToGemmBwdDataTransform& conv_to_gemm_transform)
+    {
+        const auto a_grid_desc_ak0_m_ak1 = conv_to_gemm_transform.MakeADescriptor_AK0_M_AK1();
+        const auto b_grid_desc_bk0_n_bk1 = conv_to_gemm_transform.MakeBDescriptor_BK0_N_BK1();
+        const auto ds_grid_desc_m_n =
+            generate_tuple([&](auto) { return conv_to_gemm_transform.MakeCDescriptor_M_N(); },
+                           Number<NumDTensor>{});
+        const auto e_grid_desc_m_n = conv_to_gemm_transform.MakeCDescriptor_M_N();
         return make_tuple(
             a_grid_desc_ak0_m_ak1, b_grid_desc_bk0_n_bk1, ds_grid_desc_m_n, e_grid_desc_m_n);
     }
 
     // desc
-    using ABDsEGridDesc = decltype(GetDummyABDsEGridDescriptor());
+    constexpr static ConvToGemmBwdDataTransform dummy_conv_to_gemm_transform;
+    using ABDsEGridDesc = decltype(GetDummyABDsEGridDescriptor(dummy_conv_to_gemm_transform));
 
     using AGridDesc_AK0_M_AK1 = remove_cvref_t<tuple_element_t<0, ABDsEGridDesc>>;
     using BGridDesc_BK0_N_BK1 = remove_cvref_t<tuple_element_t<1, ABDsEGridDesc>>;
@@ -270,7 +216,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
                  const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_lengths,
                  const std::array<index_t, NDimSpatial + 3>& b_g_k_c_xs_strides,
                  const std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor>&
-                     ds_g_n_c_wis_lengths,
+                 /*ds_g_n_c_wis_lengths*/,
                  const std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor>&
                      ds_g_n_c_wis_strides,
                  const std::array<index_t, NDimSpatial + 3>& e_g_n_c_wis_lengths,
@@ -291,15 +237,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
               b_element_op_{b_element_op},
               cde_element_op_{cde_element_op},
               a_g_n_k_wos_lengths_{a_g_n_k_wos_lengths},
-              a_g_n_k_wos_strides_{a_g_n_k_wos_strides},
               b_g_k_c_xs_lengths_{b_g_k_c_xs_lengths},
-              b_g_k_c_xs_strides_{b_g_k_c_xs_strides},
-              ds_g_n_c_wis_lengths_{ds_g_n_c_wis_lengths},
-              ds_g_n_c_wis_strides_{ds_g_n_c_wis_strides},
-              e_g_n_c_wis_lengths_{e_g_n_c_wis_lengths},
-              e_g_n_c_wis_strides_{e_g_n_c_wis_strides},
               conv_filter_strides_{conv_filter_strides},
-              conv_filter_dilations_{conv_filter_dilations},
               input_left_pads_{input_left_pads},
               input_right_pads_{input_right_pads}
         {
@@ -382,68 +321,47 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
                             tildes = {i_ztilde, i_ytilde, i_xtilde};
                         }
 
+                        ConvToGemmBwdDataTransform conv_to_gemm_transform_{a_g_n_k_wos_lengths,
+                                                                           a_g_n_k_wos_strides,
+                                                                           b_g_k_c_xs_lengths,
+                                                                           b_g_k_c_xs_strides,
+                                                                           e_g_n_c_wis_lengths,
+                                                                           e_g_n_c_wis_strides,
+                                                                           conv_filter_strides,
+                                                                           conv_filter_dilations,
+                                                                           input_left_pads,
+                                                                           input_right_pads,
+                                                                           tildes};
+
                         const auto a_grid_desc_ak0_m_ak1 =
-                            transform_conv_to_gemm.template MakeADescriptor_AK0_M_AK1<ALayout>(
-                                a_g_n_k_wos_lengths,
-                                a_g_n_k_wos_strides,
-                                b_g_k_c_xs_lengths,
-                                b_g_k_c_xs_strides,
-                                e_g_n_c_wis_lengths,
-                                e_g_n_c_wis_strides,
-                                conv_filter_strides,
-                                conv_filter_dilations,
-                                input_left_pads,
-                                input_right_pads,
-                                tildes);
+                            conv_to_gemm_transform_.MakeADescriptor_AK0_M_AK1();
 
                         const auto b_grid_desc_bk0_n_bk1 =
-                            transform_conv_to_gemm.template MakeBDescriptor_BK0_N_BK1<BLayout>(
-                                a_g_n_k_wos_lengths,
-                                a_g_n_k_wos_strides,
-                                b_g_k_c_xs_lengths,
-                                b_g_k_c_xs_strides,
-                                e_g_n_c_wis_lengths,
-                                e_g_n_c_wis_strides,
-                                conv_filter_strides,
-                                conv_filter_dilations,
-                                input_left_pads,
-                                input_right_pads,
-                                tildes);
+                            conv_to_gemm_transform_.MakeBDescriptor_BK0_N_BK1();
 
                         DsGridDesc_M_N ds_grid_desc_m_n;
 
                         // populate Ds desc
                         static_for<0, NumDTensor, 1>{}([&](auto i) {
                             using DLayout = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
-
-                            ds_grid_desc_m_n(i) =
-                                transform_conv_to_gemm.template MakeCDescriptor_M_N<DLayout>(
-                                    a_g_n_k_wos_lengths,
-                                    a_g_n_k_wos_strides,
-                                    b_g_k_c_xs_lengths,
-                                    b_g_k_c_xs_strides,
-                                    ds_g_n_c_wis_lengths[i],
-                                    ds_g_n_c_wis_strides[i],
-                                    conv_filter_strides,
-                                    conv_filter_dilations,
-                                    input_left_pads,
-                                    input_right_pads,
-                                    tildes);
-                        });
-
-                        const auto e_grid_desc_m_n =
-                            transform_conv_to_gemm.template MakeCDescriptor_M_N<ELayout>(
+                            static_assert(is_same_v<DLayout, ELayout>);
+                            ConvToGemmBwdDataTransform conv_to_gemm_transform_d{
                                 a_g_n_k_wos_lengths,
                                 a_g_n_k_wos_strides,
                                 b_g_k_c_xs_lengths,
                                 b_g_k_c_xs_strides,
                                 e_g_n_c_wis_lengths,
-                                e_g_n_c_wis_strides,
+                                ds_g_n_c_wis_strides[i],
                                 conv_filter_strides,
                                 conv_filter_dilations,
                                 input_left_pads,
                                 input_right_pads,
-                                tildes);
+                                tildes};
+
+                            ds_grid_desc_m_n(i) = conv_to_gemm_transform_d.MakeCDescriptor_M_N();
+                        });
+
+                        const auto e_grid_desc_m_n = conv_to_gemm_transform_.MakeCDescriptor_M_N();
 
                         // for check validity
                         ds_grid_desc_m_n_container_.push_back(ds_grid_desc_m_n);
@@ -522,17 +440,9 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
         BElementwiseOp b_element_op_;
         CDEElementwiseOp cde_element_op_;
 
-        // for checking IsSupportedArgument()
         std::array<index_t, NDimSpatial + 3> a_g_n_k_wos_lengths_;
-        std::array<index_t, NDimSpatial + 3> a_g_n_k_wos_strides_;
         std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_lengths_;
-        std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_strides_;
-        std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_c_wis_lengths_;
-        std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_c_wis_strides_;
-        std::array<index_t, NDimSpatial + 3> e_g_n_c_wis_lengths_;
-        std::array<index_t, NDimSpatial + 3> e_g_n_c_wis_strides_;
         std::array<index_t, NDimSpatial> conv_filter_strides_;
-        std::array<index_t, NDimSpatial> conv_filter_dilations_;
         std::array<index_t, NDimSpatial> input_left_pads_;
         std::array<index_t, NDimSpatial> input_right_pads_;
     };

include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp

@@ -54,15 +54,16 @@ template <typename GridwiseGemm,
           typename ABDataType,
           typename DsPointer,
           typename EDataType,
-          typename AElementwiseOperation,
-          typename BElementwiseOperation,
-          typename CDEElementwiseOperation,
+          typename AElementwiseOp,
+          typename BElementwiseOp,
+          typename CDEElementwiseOp,
           typename AGridDesc_AK0_M_AK1,
           typename BGridDesc_BK0_N_BK1,
           typename DsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
           typename EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,
           typename Block2ETileMap,
           typename ComputePtrOffsetOfBatch,
+          typename ComputePtrOffsetOfN,
           bool HasMainKBlockLoop>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
@@ -73,10 +74,9 @@ __global__ void
             const ABDataType* __restrict__ p_b_grid,
             DsPointer p_ds_grid,
             EDataType* __restrict__ p_e_grid,
-            const AElementwiseOperation a_element_op,
-            const BElementwiseOperation b_element_op,
-            const CDEElementwiseOperation cde_element_op,
-            const index_t batch_count,
+            const AElementwiseOp a_element_op,
+            const BElementwiseOp b_element_op,
+            const CDEElementwiseOp cde_element_op,
             const AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1,
             const BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1,
             const DsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock
@@ -84,24 +84,29 @@ __global__ void
             const EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock
                 e_grid_desc_mblock_mperblock_nblock_nperblock_,
             const Block2ETileMap block_2_ctile_map,
-            const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch)
+            const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch,
+            const ComputePtrOffsetOfN compute_ptr_offset_of_n)
 {
 #if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
     defined(__gfx94__))
     // offset base pointer for each work-group
-    const index_t num_blocks_per_batch =
-        __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
-    const index_t g_idx = __builtin_amdgcn_readfirstlane(get_block_1d_id() / num_blocks_per_batch);
+    const index_t n_idx = __builtin_amdgcn_readfirstlane(blockIdx.z);
+    const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.y);
 
-    const long_index_t a_batch_offset = amd_wave_read_first_lane(
-        static_cast<long_index_t>(compute_ptr_offset_of_batch.GetAPtrOffset(g_idx)));
-    const long_index_t b_batch_offset = amd_wave_read_first_lane(
-        static_cast<long_index_t>(compute_ptr_offset_of_batch.GetBPtrOffset(g_idx)));
-    const long_index_t e_batch_offset = amd_wave_read_first_lane(
-        static_cast<long_index_t>(compute_ptr_offset_of_batch.GetEPtrOffset(g_idx)));
+    const long_index_t a_batch_offset =
+        amd_wave_read_first_lane(compute_ptr_offset_of_batch.GetAPtrOffset(g_idx));
+    const long_index_t b_batch_offset =
+        amd_wave_read_first_lane(compute_ptr_offset_of_batch.GetBPtrOffset(g_idx));
+    const long_index_t e_batch_offset =
+        amd_wave_read_first_lane(compute_ptr_offset_of_batch.GetEPtrOffset(g_idx));
 
     const auto ds_batch_offset = compute_ptr_offset_of_batch.GetDsPtrOffset(g_idx);
 
+    const long_index_t a_n_offset =
+        amd_wave_read_first_lane(compute_ptr_offset_of_n.GetAPtrOffset(n_idx));
+    const long_index_t e_n_offset =
+        amd_wave_read_first_lane(compute_ptr_offset_of_n.GetEPtrOffset(n_idx));
+
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
 
     DsPointer p_ds_grid_grp;
@@ -112,10 +117,10 @@ __global__ void
     static_for<0, NumDTensor, 1>{}(
         [&](auto i) { p_ds_grid_grp(i) = p_ds_grid[i] + ds_batch_offset[i]; });
 
-    GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid + a_batch_offset,
+    GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid + a_batch_offset + a_n_offset,
                                                   p_b_grid + b_batch_offset,
                                                   p_ds_grid_grp,
-                                                  p_e_grid + e_batch_offset,
+                                                  p_e_grid + e_batch_offset + e_n_offset,
                                                   p_shared,
                                                   a_element_op,
                                                   b_element_op,
@@ -130,7 +135,6 @@ __global__ void
     ignore = p_b_grid;
     ignore = p_ds_grid;
     ignore = p_e_grid;
-    ignore = batch_count;
     ignore = a_grid_desc_ak0_m_ak1;
     ignore = b_grid_desc_bk0_n_bk1;
     ignore = ds_grid_desc_mblock_mperblock_nblock_nperblock;
@@ -139,6 +143,7 @@ __global__ void
     ignore = b_element_op;
     ignore = cde_element_op;
     ignore = compute_ptr_offset_of_batch;
+    ignore = compute_ptr_offset_of_n;
     ignore = block_2_ctile_map;
 #endif
 }
@@ -233,82 +238,54 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
     static constexpr auto I2 = Number<2>{};
     static constexpr auto I3 = Number<3>{};
 
-    static constexpr auto transform_conv_to_gemm =
-        TransformConvBwdDataToGemm_v1<NDimSpatial,
-                                      ConvBackwardDataSpecialization,
-                                      AK1,
-                                      BK1,
-                                      MPerBlock,
-                                      NPerBlock,
-                                      KPerBlock,
-                                      DoPadGemmM,
-                                      DoPadGemmN>{};
-
-    static auto GetDummyABDsEGridDescriptor()
+    using ConvToGemmBwdDataTransform = TransformConvBwdDataToGemm_v1<NDimSpatial,
+                                                                     ConvBackwardDataSpecialization,
+                                                                     AK1,
+                                                                     BK1,
+                                                                     MPerBlock,
+                                                                     NPerBlock,
+                                                                     KPerBlock,
+                                                                     DoPadGemmM,
+                                                                     DoPadGemmN,
+                                                                     ALayout,
+                                                                     BLayout,
+                                                                     ELayout,
+                                                                     true, /*SplitConvN*/
+                                                                     ABDataType,
+                                                                     EDataType>;
+
+    static auto
+    GetDummyABDsEGridDescriptor(const ConvToGemmBwdDataTransform& conv_to_gemm_transform)
     {
-        const std::array<index_t, NDimSpatial + 3> dummy_tensor_lengths = {1};
-        const std::array<index_t, NDimSpatial + 3> dummy_tensor_strides = {1};
-        const std::array<index_t, NDimSpatial> dummy_spatial_lengths    = {1};
-
-        const auto a_grid_desc_ak0_m_ak1 =
-            transform_conv_to_gemm.template MakeADescriptor_AK0_M_AK1<ALayout>(
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths);
-
-        const auto b_grid_desc_bk0_n_bk1 =
-            transform_conv_to_gemm.template MakeBDescriptor_BK0_N_BK1<BLayout>(
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_tensor_lengths,
-                dummy_tensor_strides,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths,
-                dummy_spatial_lengths);
+        const auto a_grid_desc_ak0_m_ak1 = conv_to_gemm_transform.MakeADescriptor_AK0_M_AK1();
+
+        const auto b_grid_desc_bk0_n_bk1 = conv_to_gemm_transform.MakeBDescriptor_BK0_N_BK1();
 
         const auto ds_grid_desc_m_n = generate_tuple(
             [&](auto i) {
-                using DLayout = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
-
-                return transform_conv_to_gemm.template MakeCDescriptor_M_N<DLayout>(
-                    dummy_tensor_lengths,
-                    dummy_tensor_strides,
-                    dummy_tensor_lengths,
-                    dummy_tensor_strides,
-                    dummy_tensor_lengths,
-                    dummy_tensor_strides,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths,
-                    dummy_spatial_lengths);
+                using DLayout   = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
+                using DDataType = remove_cvref_t<tuple_element_t<i.value, DsDataType>>;
+                using ConvToGemmBwdDataTransformD =
+                    TransformConvBwdDataToGemm_v1<NDimSpatial,
+                                                  ConvBackwardDataSpecialization,
+                                                  AK1,
+                                                  BK1,
+                                                  MPerBlock,
+                                                  NPerBlock,
+                                                  KPerBlock,
+                                                  DoPadGemmM,
+                                                  DoPadGemmN,
+                                                  ALayout,
+                                                  BLayout,
+                                                  DLayout,
+                                                  true, /*SplitConvN*/
+                                                  ABDataType,
+                                                  DDataType>;
+                return ConvToGemmBwdDataTransformD{}.MakeCDescriptor_M_N();
             },
             Number<NumDTensor>{});
 
-        const auto e_grid_desc_m_n =
-            transform_conv_to_gemm.template MakeCDescriptor_M_N<ELayout>(dummy_tensor_lengths,
-                                                                         dummy_tensor_strides,
-                                                                         dummy_tensor_lengths,
-                                                                         dummy_tensor_strides,
-                                                                         dummy_tensor_lengths,
-                                                                         dummy_tensor_strides,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths,
-                                                                         dummy_spatial_lengths);
+        const auto e_grid_desc_m_n = conv_to_gemm_transform.MakeCDescriptor_M_N();
 
         return make_tuple(
             a_grid_desc_ak0_m_ak1, b_grid_desc_bk0_n_bk1, ds_grid_desc_m_n, e_grid_desc_m_n);
@@ -377,7 +354,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
     }
 
     // desc
-    using ABDsEGridDesc = decltype(GetDummyABDsEGridDescriptor());
+    constexpr static ConvToGemmBwdDataTransform dummy_conv_to_gemm_transform;
+    using ABDsEGridDesc = decltype(GetDummyABDsEGridDescriptor(dummy_conv_to_gemm_transform));
 
     using AGridDesc_AK0_M_AK1 = remove_cvref_t<tuple_element_t<0, ABDsEGridDesc>>;
     using BGridDesc_BK0_N_BK1 = remove_cvref_t<tuple_element_t<1, ABDsEGridDesc>>;
@@ -431,15 +409,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
               b_element_op_{b_element_op},
               cde_element_op_{cde_element_op},
               a_g_n_k_wos_lengths_{a_g_n_k_wos_lengths},
-              a_g_n_k_wos_strides_{a_g_n_k_wos_strides},
               b_g_k_c_xs_lengths_{b_g_k_c_xs_lengths},
-              b_g_k_c_xs_strides_{b_g_k_c_xs_strides},
-              ds_g_n_c_wis_lengths_{ds_g_n_c_wis_lengths},
-              ds_g_n_c_wis_strides_{ds_g_n_c_wis_strides},
-              e_g_n_c_wis_lengths_{e_g_n_c_wis_lengths},
-              e_g_n_c_wis_strides_{e_g_n_c_wis_strides},
               conv_filter_strides_{conv_filter_strides},
-              conv_filter_dilations_{conv_filter_dilations},
               input_left_pads_{input_left_pads},
               input_right_pads_{input_right_pads}
         {
@@ -450,11 +421,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                 p_ds_grid_(i) = static_cast<const DDataType*>(p_ds[i]);
             });
 
-            // A/B/Ds/E Batch Stride
-            compute_ptr_offset_of_batch_.BatchStrideA_ = a_g_n_k_wos_strides[0];
-            compute_ptr_offset_of_batch_.BatchStrideB_ = b_g_k_c_xs_strides[0];
-            compute_ptr_offset_of_batch_.BatchStrideE_ = e_g_n_c_wis_strides[0];
-
             static_for<0, NumDTensor, 1>{}([&](auto i) {
                 compute_ptr_offset_of_batch_.BatchStrideDs_(i) = ds_g_n_c_wis_strides[i][0];
             });
@@ -526,68 +492,65 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                             throw std::runtime_error("wrong! only implemented for 2D and 3D now");
                         }
 
+                        ConvToGemmBwdDataTransform conv_to_gemm_transform_{a_g_n_k_wos_lengths,
+                                                                           a_g_n_k_wos_strides,
+                                                                           b_g_k_c_xs_lengths,
+                                                                           b_g_k_c_xs_strides,
+                                                                           e_g_n_c_wis_lengths,
+                                                                           e_g_n_c_wis_strides,
+                                                                           conv_filter_strides,
+                                                                           conv_filter_dilations,
+                                                                           input_left_pads,
+                                                                           input_right_pads,
+                                                                           tildes};
+
+                        conv_N_per_block_ = conv_to_gemm_transform_.N_;
+
                         const auto a_grid_desc_ak0_m_ak1 =
-                            transform_conv_to_gemm.template MakeADescriptor_AK0_M_AK1<ALayout>(
-                                a_g_n_k_wos_lengths,
-                                a_g_n_k_wos_strides,
-                                b_g_k_c_xs_lengths,
-                                b_g_k_c_xs_strides,
-                                e_g_n_c_wis_lengths,
-                                e_g_n_c_wis_strides,
-                                conv_filter_strides,
-                                conv_filter_dilations,
-                                input_left_pads,
-                                input_right_pads,
-                                tildes);
+                            conv_to_gemm_transform_.MakeADescriptor_AK0_M_AK1();
 
                         const auto b_grid_desc_bk0_n_bk1 =
-                            transform_conv_to_gemm.template MakeBDescriptor_BK0_N_BK1<BLayout>(
-                                a_g_n_k_wos_lengths,
-                                a_g_n_k_wos_strides,
-                                b_g_k_c_xs_lengths,
-                                b_g_k_c_xs_strides,
-                                e_g_n_c_wis_lengths,
-                                e_g_n_c_wis_strides,
-                                conv_filter_strides,
-                                conv_filter_dilations,
-                                input_left_pads,
-                                input_right_pads,
-                                tildes);
+                            conv_to_gemm_transform_.MakeBDescriptor_BK0_N_BK1();
 
                         DsGridDesc_M_N ds_grid_desc_m_n;
 
                         // populate Ds desc
                         static_for<0, NumDTensor, 1>{}([&](auto i) {
-                            using DLayout = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
-
-                            ds_grid_desc_m_n(i) =
-                                transform_conv_to_gemm.template MakeCDescriptor_M_N<DLayout>(
-                                    a_g_n_k_wos_lengths,
-                                    a_g_n_k_wos_strides,
-                                    b_g_k_c_xs_lengths,
-                                    b_g_k_c_xs_strides,
-                                    ds_g_n_c_wis_lengths[i],
-                                    ds_g_n_c_wis_strides[i],
-                                    conv_filter_strides,
-                                    conv_filter_dilations,
-                                    input_left_pads,
-                                    input_right_pads,
-                                    tildes);
-                        });
-
-                        const auto e_grid_desc_m_n =
-                            transform_conv_to_gemm.template MakeCDescriptor_M_N<ELayout>(
+                            using DLayout   = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
+                            using DDataType = remove_cvref_t<tuple_element_t<i.value, DsDataType>>;
+                            using ConvToGemmBwdDataTransformD =
+                                TransformConvBwdDataToGemm_v1<NDimSpatial,
+                                                              ConvBackwardDataSpecialization,
+                                                              AK1,
+                                                              BK1,
+                                                              MPerBlock,
+                                                              NPerBlock,
+                                                              KPerBlock,
+                                                              DoPadGemmM,
+                                                              DoPadGemmN,
+                                                              ALayout,
+                                                              BLayout,
+                                                              DLayout,
+                                                              true, /*SplitConvN*/
+                                                              ABDataType,
+                                                              DDataType>;
+                            ConvToGemmBwdDataTransformD conv_to_gemm_transform_d{
                                 a_g_n_k_wos_lengths,
                                 a_g_n_k_wos_strides,
                                 b_g_k_c_xs_lengths,
                                 b_g_k_c_xs_strides,
-                                e_g_n_c_wis_lengths,
-                                e_g_n_c_wis_strides,
+                                ds_g_n_c_wis_lengths[i],
+                                ds_g_n_c_wis_strides[i],
                                 conv_filter_strides,
                                 conv_filter_dilations,
                                 input_left_pads,
                                 input_right_pads,
-                                tildes);
+                                tildes};
+
+                            ds_grid_desc_m_n(i) = conv_to_gemm_transform_d.MakeCDescriptor_M_N();
+                        });
+
+                        const auto e_grid_desc_m_n = conv_to_gemm_transform_.MakeCDescriptor_M_N();
 
                         // desc for problem definition
                         const auto a_grid_desc_m_k =
@@ -628,6 +591,13 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                     }
                 }
             }
+            // A/B/Ds/E Batch Stride
+            compute_ptr_offset_of_batch_.BatchStrideA_ = a_g_n_k_wos_strides[0];
+            compute_ptr_offset_of_batch_.BatchStrideB_ = b_g_k_c_xs_strides[0];
+            compute_ptr_offset_of_batch_.BatchStrideE_ = e_g_n_c_wis_strides[0];
+
+            compute_ptr_offset_of_n_.BatchStrideA_ = a_g_n_k_wos_strides[1] * conv_N_per_block_;
+            compute_ptr_offset_of_n_.BatchStrideE_ = e_g_n_c_wis_strides[1] * conv_N_per_block_;
         }
 
         void Print() const
@@ -660,6 +630,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
 
         // tensor descriptor for problem definition
         index_t num_group_;
+        index_t conv_N_per_block_;
         std::vector<AGridDesc_M_K> a_grid_desc_m_k_container_;
         std::vector<BGridDesc_N_K> b_grid_desc_n_k_container_;
         std::vector<DsGridDesc_M_N> ds_grid_desc_m_n_container_;
@@ -678,23 +649,16 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
 
         // for computing batch offset
         ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor> compute_ptr_offset_of_batch_;
+        ComputePtrOffsetOfStridedBatch<I1, I1, I0> compute_ptr_offset_of_n_;
 
         // element-wise op
         AElementwiseOp a_element_op_;
         BElementwiseOp b_element_op_;
         CDEElementwiseOp cde_element_op_;
 
-        // for checking IsSupportedArgument()
         std::array<index_t, NDimSpatial + 3> a_g_n_k_wos_lengths_;
-        std::array<index_t, NDimSpatial + 3> a_g_n_k_wos_strides_;
         std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_lengths_;
-        std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_strides_;
-        std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_c_wis_lengths_;
-        std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_c_wis_strides_;
-        std::array<index_t, NDimSpatial + 3> e_g_n_c_wis_lengths_;
-        std::array<index_t, NDimSpatial + 3> e_g_n_c_wis_strides_;
         std::array<index_t, NDimSpatial> conv_filter_strides_;
-        std::array<index_t, NDimSpatial> conv_filter_dilations_;
         std::array<index_t, NDimSpatial> input_left_pads_;
         std::array<index_t, NDimSpatial> input_right_pads_;
     };
@@ -711,8 +675,12 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                 arg.Print();
             }
 
-            float ave_time = 0;
+            const index_t gdy = arg.num_group_;
+            const index_t num_workgroups_per_Conv_N =
+                arg.a_g_n_k_wos_lengths_[I1] / arg.conv_N_per_block_;
+            const index_t gdz = num_workgroups_per_Conv_N;
 
+            float ave_time = 0;
             for(std::size_t i = 0; i < arg.a_grid_desc_ak0_m_ak1_container_.size(); i++)
             {
                 if(!GridwiseGemm::CheckValidity(arg.a_grid_desc_m_k_container_[i],
@@ -724,9 +692,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                     throw std::runtime_error("wrong! device_op has invalid setting");
                 }
 
-                const index_t grid_size = arg.block_2_etile_map_container_[i].CalculateGridSize(
-                                              arg.e_grid_desc_m_n_container_[i]) *
-                                          arg.num_group_;
+                const index_t gdx = arg.block_2_etile_map_container_[i].CalculateGridSize(
+                    arg.e_grid_desc_m_n_container_[i]);
 
                 const auto GemmK = arg.a_grid_desc_m_k_container_[i].GetLength(I1);
 
@@ -747,12 +714,13 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                         DeviceOp::EGridDesc_MBlock_MPerBlock_NBlock_NPerBlock,
                         Block2ETileMap,
                         ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor>,
+                        ComputePtrOffsetOfStridedBatch<I1, I1, I0>,
                         has_main_loop>;
 
                     return launch_and_time_kernel(
                         stream_config,
                         kernel,
-                        dim3(grid_size),
+                        dim3(gdx, gdy, gdz),
                         dim3(BlockSize),
                         0,
                         arg.p_a_grid_,
@@ -762,13 +730,13 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
                         arg.a_element_op_,
                         arg.b_element_op_,
                         arg.cde_element_op_,
-                        arg.a_g_n_k_wos_lengths_[0], // Group count
                         arg.a_grid_desc_ak0_m_ak1_container_[i],
                         arg.b_grid_desc_bk0_n_bk1_container_[i],
                         arg.ds_grid_desc_mblock_mperblock_nblock_nperblock_container_[i],
                         arg.e_grid_desc_mblock_mperblock_nblock_nperblock_container_[i],
                         arg.block_2_etile_map_container_[i],
-                        arg.compute_ptr_offset_of_batch_);
+                        arg.compute_ptr_offset_of_batch_,
+                        arg.compute_ptr_offset_of_n_);
                 };
 
                 if(GridwiseGemm::CalculateHasMainKBlockLoop(GemmK))

include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp

@@ -111,8 +111,7 @@ __global__ void
             [[maybe_unused]] const ComputePtrOffsetOfBatch compute_ptr_offset_of_batch,
             [[maybe_unused]] const index_t num_k_per_block)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     // offset base pointer for each work-group
     const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.z * NumGroupsToMerge);
     const index_t k_idx = __builtin_amdgcn_readfirstlane(blockIdx.y * num_k_per_block);

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp

@@ -38,8 +38,7 @@ __global__ void
     // __attribute__((amdgpu_waves_per_eu(1, 1)))
     kernel_gemm_xdl_cshuffle_v3(typename GridwiseGemm::Argument karg)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
 
     GridwiseGemm::template Run<HasMainKBlockLoop, CGlobalMemoryDataOperation, TailNum>(

include/ck/tensor_operation/operator_transform/transform_conv_bwd_data_to_gemm_v1.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -13,245 +13,614 @@
 namespace ck {
 namespace tensor_operation {
 
-namespace {
 template <
     index_t NDimSpatial,
+    ck::tensor_operation::device::ConvolutionBackwardDataSpecialization ConvBwdDataSpecialization,
+    index_t AK1,
+    index_t BK1,
+    index_t GemmMPerBlock,
+    index_t GemmNPerBlock,
+    index_t GemmKPerBlock,
+    bool DoPadGemmM,
+    bool DoPadGemmN,
     typename ALayout,
-    ck::tensor_operation::device::ConvolutionBackwardDataSpecialization ConvBwdDataSpecialization>
-constexpr auto make_out_grid_desc(const index_t N,
-                                  const index_t Do,
-                                  const index_t Ho,
-                                  const index_t Wo,
-                                  const index_t K,
-                                  const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_strides)
+    typename BLayout,
+    typename CLayout,
+    bool SplitN              = false,
+    typename ADataType       = float,
+    typename CDataType       = float,
+    index_t NumGroupsToMerge = 1,
+    typename IndexType       = index_t>
+struct TransformConvBwdDataToGemm_v1
 {
-    const auto KStride = Number<1>{};
+    private:
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+    static constexpr auto I3 = Number<3>{};
 
-    if constexpr(is_same_v<ALayout, tensor_layout::convolution::NHWGK>)
-    {
-        const index_t NStride  = out_g_n_k_wos_strides[1];
-        const index_t HiStride = out_g_n_k_wos_strides[3];
-        const index_t WiStride = out_g_n_k_wos_strides[4];
-        if constexpr(ConvBwdDataSpecialization ==
-                     ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
-                         Filter1x1Stride1Pad0)
-        {
+    static constexpr auto NonSpatialDimsNum = Number<3>{};
 
-            return make_naive_tensor_descriptor(make_tuple(N * Ho * Wo, K),
-                                                make_tuple(WiStride, KStride));
-        }
-        else
+    static constexpr auto DIdx = NonSpatialDimsNum;
+    static constexpr auto HIdx =
+        NDimSpatial == 2 ? NonSpatialDimsNum : Number<NonSpatialDimsNum + 1>{};
+    static constexpr auto WIdx =
+        NDimSpatial == 2 ? Number<NonSpatialDimsNum + 1>{} : Number<NonSpatialDimsNum + 2>{};
+
+    static constexpr auto ZIdx = NonSpatialDimsNum;
+    static constexpr auto YIdx =
+        NDimSpatial == 2 ? NonSpatialDimsNum : Number<NonSpatialDimsNum + 1>{};
+    static constexpr auto XIdx =
+        NDimSpatial == 2 ? Number<NonSpatialDimsNum + 1>{} : Number<NonSpatialDimsNum + 2>{};
+
+    template <typename ConvDimsType>
+    static long_index_t calculate_element_space_size_impl(const ConvDimsType& lengths,
+                                                          const ConvDimsType& strides,
+                                                          index_t i)
+    {
+        long_index_t acc = 1;
+        for(; i < (NDimSpatial + 3); i++)
         {
-            return make_naive_tensor_descriptor(make_tuple(N, Ho, Wo, K),
-                                                make_tuple(NStride, HiStride, WiStride, KStride));
+            acc +=
+                static_cast<long_index_t>(lengths[i] - I1) * static_cast<long_index_t>(strides[i]);
         }
+
+        return acc;
     }
-    else if constexpr(is_same_v<ALayout, tensor_layout::convolution::NDHWGK>)
+
+    template <typename ConvDimsType>
+    static IndexType GetSplitedNSize(const ConvDimsType& a_g_n_k_wos_lengths,
+                                     const ConvDimsType& a_g_n_k_wos_strides,
+                                     const ConvDimsType& c_g_n_c_wis_lengths,
+                                     const ConvDimsType& c_g_n_c_wis_strides)
     {
-        const index_t NStride  = out_g_n_k_wos_strides[1];
-        const index_t DoStride = out_g_n_k_wos_strides[3];
-        const index_t HoStride = out_g_n_k_wos_strides[4];
-        const index_t WoStride = out_g_n_k_wos_strides[5];
-        if constexpr(ConvBwdDataSpecialization ==
-                     ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
-                         Filter1x1Stride1Pad0)
+        const long_index_t a_element_space_size =
+            calculate_element_space_size_impl(a_g_n_k_wos_lengths, a_g_n_k_wos_strides, I1);
+        const long_index_t c_element_space_size =
+            calculate_element_space_size_impl(c_g_n_c_wis_lengths, c_g_n_c_wis_strides, I1);
+        const long_index_t element_space_size = math::max(a_element_space_size * sizeof(ADataType),
+                                                          c_element_space_size * sizeof(CDataType));
+        constexpr long_index_t TwoGB          = (long_index_t{1} << 31);
+
+        const IndexType N = a_g_n_k_wos_lengths[I1];
+
+        if(element_space_size > TwoGB)
         {
+            // Minimum divisor of N to not exceed 2GB
+            const auto divisor = math::integer_divide_ceil(element_space_size, TwoGB);
 
-            return make_naive_tensor_descriptor(make_tuple(N * Do * Ho * Wo, K),
-                                                make_tuple(WoStride, KStride));
+            if(divisor <= static_cast<double>(N))
+            {
+                // Find least divisor of N larger than element_space_size / TwoGB
+                // Iterate up to sqrt(N). There are no divisors above this value.
+                for(IndexType least_divisor = divisor; least_divisor * least_divisor <= N;
+                    least_divisor++)
+                {
+                    if(N % least_divisor == 0)
+                    {
+                        return N / least_divisor;
+                    }
+                }
+                // Not found, process one Convolution N per block
+                return 1;
+            }
+            else
+            {
+                // Not possible to support even after split N.
+                // Too large tensor.
+                return N;
+            }
         }
         else
         {
-            return make_naive_tensor_descriptor(
-                make_tuple(N, Do, Ho, Wo, K),
-                make_tuple(NStride, DoStride, HoStride, WoStride, KStride));
+            // Split N is not needed.
+            return N;
         }
     }
-    else if constexpr(is_same_v<ALayout, tensor_layout::convolution::GNHWK>)
+
+    public:
+    __host__ __device__ constexpr TransformConvBwdDataToGemm_v1() {}
+
+    template <typename TransformConvBwdDataToGemm_v1Base>
+    __host__ __device__ TransformConvBwdDataToGemm_v1(
+        const TransformConvBwdDataToGemm_v1Base& transform_conv_bwd_data_to_gemm_base)
+        : N_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.N_)},
+          Di_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Di_)},
+          Hi_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Hi_)},
+          Wi_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Wi_)},
+          Do_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Do_)},
+          Ho_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Ho_)},
+          Wo_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Wo_)},
+          Z_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Z_)},
+          Y_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.Y_)},
+          X_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.X_)},
+          K_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.K_)},
+          C_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.C_)},
+          DiStride_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.DiStride_)},
+          HiStride_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.HiStride_)},
+          WiStride_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.WiStride_)},
+          DoStride_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.DoStride_)},
+          HoStride_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.HoStride_)},
+          WoStride_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.WoStride_)},
+          CStrideTensorB_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.CStrideTensorB_)},
+          CStrideTensorC_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.CStrideTensorC_)},
+          KStrideTensorA_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.KStrideTensorA_)},
+          KStrideTensorB_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.KStrideTensorB_)},
+          NStrideTensorA_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.NStrideTensorA_)},
+          NStrideTensorC_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.NStrideTensorC_)},
+          ConvStrideD_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ConvStrideD_)},
+          ConvStrideH_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ConvStrideH_)},
+          ConvStrideW_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ConvStrideW_)},
+          ConvDilationD_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ConvDilationD_)},
+          ConvDilationH_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ConvDilationH_)},
+          ConvDilationW_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ConvDilationW_)},
+          InLeftPadD_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.InLeftPadD_)},
+          InLeftPadH_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.InLeftPadH_)},
+          InLeftPadW_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.InLeftPadW_)},
+          InRightPadD_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.InRightPadD_)},
+          InRightPadH_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.InRightPadH_)},
+          InRightPadW_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.InRightPadW_)},
+          IdxZTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.IdxZTilde_)},
+          IdxYTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.IdxYTilde_)},
+          IdxXTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.IdxXTilde_)},
+          GcdStrideDilationD_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.GcdStrideDilationD_)},
+          GcdStrideDilationH_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.GcdStrideDilationH_)},
+          GcdStrideDilationW_{
+              static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.GcdStrideDilationW_)},
+          ZTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ZTilde_)},
+          YTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.YTilde_)},
+          XTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.XTilde_)},
+          DTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.DTilde_)},
+          HTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.HTilde_)},
+          WTilde_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.WTilde_)},
+          ZDot_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.ZDot_)},
+          YDot_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.YDot_)},
+          XDot_{static_cast<IndexType>(transform_conv_bwd_data_to_gemm_base.XDot_)}
     {
-        // assume packed
-        if constexpr(ConvBwdDataSpecialization ==
-                     ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
-                         Filter1x1Stride1Pad0)
+    }
+
+    template <typename ConvDimsType, typename ConvSpatialDimsType>
+    __host__ __device__
+    TransformConvBwdDataToGemm_v1(const ConvDimsType& a_g_n_k_wos_lengths,
+                                  const ConvDimsType& a_g_n_k_wos_strides,
+                                  const ConvDimsType& b_g_k_c_xs_lengths,
+                                  const ConvDimsType& b_g_k_c_xs_strides,
+                                  const ConvDimsType& c_g_n_c_wis_lengths,
+                                  const ConvDimsType& c_g_n_c_wis_strides,
+                                  const ConvSpatialDimsType& conv_filter_strides,
+                                  const ConvSpatialDimsType& conv_filter_dilations,
+                                  const ConvSpatialDimsType& input_left_pads,
+                                  const ConvSpatialDimsType& input_right_pads,
+                                  const ConvSpatialDimsType& tildes)
+        : Hi_{c_g_n_c_wis_lengths[HIdx]},
+          Wi_{c_g_n_c_wis_lengths[WIdx]},
+          Ho_{a_g_n_k_wos_lengths[HIdx]},
+          Wo_{a_g_n_k_wos_lengths[WIdx]},
+          Y_{b_g_k_c_xs_lengths[YIdx]},
+          X_{b_g_k_c_xs_lengths[XIdx]},
+          K_{a_g_n_k_wos_lengths[I2]},
+          C_{b_g_k_c_xs_lengths[I2]},
+          HiStride_{c_g_n_c_wis_strides[HIdx]},
+          WiStride_{c_g_n_c_wis_strides[WIdx]},
+          HoStride_{a_g_n_k_wos_strides[HIdx]},
+          WoStride_{a_g_n_k_wos_strides[WIdx]},
+          CStrideTensorB_{b_g_k_c_xs_strides[I2]},
+          CStrideTensorC_{c_g_n_c_wis_strides[I2]},
+          KStrideTensorA_{a_g_n_k_wos_strides[I2]},
+          KStrideTensorB_{b_g_k_c_xs_strides[I1]},
+          NStrideTensorA_{a_g_n_k_wos_strides[I1]},
+          NStrideTensorC_{c_g_n_c_wis_strides[I1]},
+          ConvStrideH_{conv_filter_strides[HIdx - NonSpatialDimsNum]},
+          ConvStrideW_{conv_filter_strides[WIdx - NonSpatialDimsNum]},
+          ConvDilationH_{conv_filter_dilations[HIdx - NonSpatialDimsNum]},
+          ConvDilationW_{conv_filter_dilations[WIdx - NonSpatialDimsNum]},
+          InLeftPadH_{input_left_pads[HIdx - NonSpatialDimsNum]},
+          InLeftPadW_{input_left_pads[WIdx - NonSpatialDimsNum]},
+          InRightPadH_{input_right_pads[HIdx - NonSpatialDimsNum]},
+          InRightPadW_{input_right_pads[WIdx - NonSpatialDimsNum]},
+          IdxYTilde_{tildes[YIdx - NonSpatialDimsNum]},
+          IdxXTilde_{tildes[XIdx - NonSpatialDimsNum]}
+    {
+        static_assert(is_same_v<ConvSpatialDimsType, std::array<IndexType, NDimSpatial>> ||
+                      is_same_v<ConvSpatialDimsType, ck::Array<IndexType, NDimSpatial>>);
+        static_assert(is_same_v<ConvDimsType, std::array<IndexType, NDimSpatial + I3>> ||
+                      is_same_v<ConvDimsType, ck::Array<IndexType, NDimSpatial + I3>>);
+
+        if constexpr(SplitN)
         {
-            return make_naive_tensor_descriptor_packed(make_tuple(N * Ho * Wo, K));
+            N_ = GetSplitedNSize(
+                a_g_n_k_wos_lengths, a_g_n_k_wos_strides, c_g_n_c_wis_lengths, c_g_n_c_wis_strides);
         }
         else
         {
-            return make_naive_tensor_descriptor_packed(make_tuple(N, Ho, Wo, K));
+            N_ = c_g_n_c_wis_lengths[I1];
         }
-    }
-    else if constexpr(is_same_v<ALayout, tensor_layout::convolution::GNDHWK>)
-    {
-        // assume packed
-        if constexpr(ConvBwdDataSpecialization ==
-                     ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
-                         Filter1x1Stride1Pad0)
+        if constexpr(NDimSpatial == 3)
         {
-            return make_naive_tensor_descriptor_packed(make_tuple(N * Do * Ho * Wo, K));
+            Di_                 = c_g_n_c_wis_lengths[DIdx];
+            Do_                 = a_g_n_k_wos_lengths[DIdx];
+            Z_                  = b_g_k_c_xs_lengths[ZIdx];
+            DiStride_           = c_g_n_c_wis_strides[DIdx];
+            DoStride_           = a_g_n_k_wos_strides[DIdx];
+            ConvStrideD_        = conv_filter_strides[DIdx - NonSpatialDimsNum];
+            ConvDilationD_      = conv_filter_dilations[DIdx - NonSpatialDimsNum];
+            InLeftPadD_         = input_left_pads[DIdx - NonSpatialDimsNum];
+            InRightPadD_        = input_right_pads[DIdx - NonSpatialDimsNum];
+            IdxZTilde_          = tildes[ZIdx - NonSpatialDimsNum];
+            GcdStrideDilationD_ = math::gcd(ConvStrideD_, ConvDilationD_);
+            ZTilde_             = ConvStrideD_ / GcdStrideDilationD_;
+            DTilde_ = Do_ + math::integer_divide_ceil(ConvDilationD_ * (Z_ - I1), ConvStrideD_);
+            ZDot_   = math::integer_divide_ceil(Z_, ZTilde_);
         }
         else
         {
-            return make_naive_tensor_descriptor_packed(make_tuple(N, Do, Ho, Wo, K));
+            Di_ = Do_ = Z_ = ZTilde_ = ConvStrideD_ = DTilde_ = ZDot_ = 1;
+            InLeftPadD_ = InRightPadD_ = DiStride_ = DoStride_ = IdxZTilde_ = 0;
         }
-    }
-    else
-    {
-        throw std::runtime_error("wrong! unsupported layout: " + ALayout::name());
-    }
-}
 
-template <typename BLayout>
-constexpr auto make_wei_grid_desc(
-    const index_t K, const index_t Z, const index_t Y, const index_t X, const index_t C)
-{
+        GcdStrideDilationH_ = math::gcd(ConvStrideH_, ConvDilationH_);
+        GcdStrideDilationW_ = math::gcd(ConvStrideW_, ConvDilationW_);
 
-    if constexpr(is_same_v<BLayout, tensor_layout::convolution::GKYXC>)
-    {
-        return make_naive_tensor_descriptor_packed(make_tuple(K, Y, X, C));
-    }
-    else if constexpr(is_same_v<BLayout, tensor_layout::convolution::GKZYXC>)
-    {
-        return make_naive_tensor_descriptor_packed(make_tuple(K, Z, Y, X, C));
-    }
-    else
-    {
-        throw std::runtime_error("wrong! unsupported layout: " + BLayout::name());
-    }
-}
-
-template <index_t NDimSpatial, typename CLayout>
-constexpr auto make_in_grid_desc(const index_t N,
-                                 const index_t Di,
-                                 const index_t Hi,
-                                 const index_t Wi,
-                                 const index_t C,
-                                 const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_strides)
-{
+        YTilde_ = ConvStrideH_ / GcdStrideDilationH_;
+        XTilde_ = ConvStrideW_ / GcdStrideDilationW_;
 
-    if constexpr(is_same_v<CLayout, tensor_layout::convolution::GNHWC> ||
-                 is_same_v<CLayout, tensor_layout::convolution::NHWGC> ||
-                 is_same_v<CLayout, tensor_layout::convolution::G_NHW_C>)
-    {
-        return make_naive_tensor_descriptor(make_tuple(N, Hi, Wi, C),
-                                            make_tuple(in_g_n_c_wis_strides[1],
-                                                       in_g_n_c_wis_strides[3],
-                                                       in_g_n_c_wis_strides[4],
-                                                       in_g_n_c_wis_strides[2]));
+        HTilde_ = Ho_ + math::integer_divide_ceil(ConvDilationH_ * (Y_ - I1), ConvStrideH_);
+        WTilde_ = Wo_ + math::integer_divide_ceil(ConvDilationW_ * (X_ - I1), ConvStrideW_);
+
+        YDot_ = math::integer_divide_ceil(Y_, YTilde_);
+        XDot_ = math::integer_divide_ceil(X_, XTilde_);
     }
-    else if constexpr(is_same_v<CLayout, tensor_layout::convolution::GNDHWC> ||
-                      is_same_v<CLayout, tensor_layout::convolution::NDHWGC>)
+
+#if 0 // At now not supported to split tensor
+    __host__ bool AreDescriptorsSmallerThan2GB() const
     {
-        return make_naive_tensor_descriptor(make_tuple(N, Di, Hi, Wi, C),
-                                            make_tuple(in_g_n_c_wis_strides[1],
-                                                       in_g_n_c_wis_strides[3],
-                                                       in_g_n_c_wis_strides[4],
-                                                       in_g_n_c_wis_strides[5],
-                                                       in_g_n_c_wis_strides[2]));
+        constexpr long_index_t TwoGB = (long_index_t{1} << 31);
+
+        const long_index_t in_desc_space_size =
+            I1 + (N_ - I1) * NStrideTensorC_ + (Di_ - I1) * DiStride_ + (Hi_ - I1) * HiStride_ +
+            (Wi_ - I1) * WiStride_ + (C_ - I1) * CStrideTensorC_;
+        const long_index_t out_desc_space_size =
+            I1 + (N_ - I1) * NStrideTensorA_ + (Do_ - I1) * DoStride_ + (Ho_ - I1) * HoStride_ +
+            (Wo_ - I1) * WoStride_ + (K_ - I1) * KStrideTensorA_;
+
+        bool is_a_descriptor_smaller_than_2GB = (out_desc_space_size * sizeof(ADataType)) <= TwoGB;
+        bool is_c_descriptor_smaller_than_2GB = (in_desc_space_size * sizeof(CDataType)) <= TwoGB;
+
+        return is_a_descriptor_smaller_than_2GB && is_c_descriptor_smaller_than_2GB;
     }
-    else
+
+    __host__ auto SplitConvProblem(const ADataType* a_grid_ptr_base,
+                                   CDataType* c_grid_ptr_base) const
     {
-        throw std::runtime_error("wrong! unsupported layout: " + CLayout::name());
-    }
-}
+        // Create copies
+        auto conv_to_gemm_transformer_left  = *this;
+        auto conv_to_gemm_transformer_right = *this;
+        IndexType a_right_offset            = 0;
+        IndexType c_right_offset            = 0;
+        // Calculate real filter size
+        const IndexType z_eff = (Z_ - 1) * ConvDilationD_ + 1;
+        const IndexType y_eff = (Y_ - 1) * ConvDilationH_ + 1;
+        const IndexType x_eff = (X_ - 1) * ConvDilationW_ + 1;
+        // Calculate start position in input for right tensor
+        const IndexType di_right_transformer_start_idx = (Do_ / 2) * ConvStrideD_;
+        const IndexType hi_right_transformer_start_idx = (Ho_ / 2) * ConvStrideH_;
+        const IndexType wi_right_transformer_start_idx = (Wo_ / 2) * ConvStrideW_;
+        // Calculate last position in input for left tensor
+        const IndexType di_left_transformer_end_idx = (Do_ / 2 - 1) * ConvStrideD_ + z_eff;
+        const IndexType hi_left_transformer_end_idx = (Ho_ / 2 - 1) * ConvStrideH_ + y_eff;
+        const IndexType wi_left_transformer_end_idx = (Wo_ / 2 - 1) * ConvStrideW_ + x_eff;
+        // Allow to split if whole left padding will be in left tensor and right padding in right
+        // tensor
+        const bool is_possible_to_split_d = Do_ != 1 &&
+                                            di_right_transformer_start_idx > InLeftPadD_ &&
+                                            di_left_transformer_end_idx <= (InLeftPadD_ + Di_);
+        const bool is_possible_to_split_h = Ho_ != 1 &&
+                                            hi_right_transformer_start_idx > InLeftPadH_ &&
+                                            hi_left_transformer_end_idx <= (InLeftPadH_ + Hi_);
+        const bool is_possible_to_split_w = Wo_ != 1 &&
+                                            wi_right_transformer_start_idx > InLeftPadW_ &&
+                                            wi_left_transformer_end_idx <= (InLeftPadW_ + Wi_);
+
+        if(is_possible_to_split_d)
+        {
+            // Apply new sizes
+            // Split output on half
+            conv_to_gemm_transformer_left.Do_  = Do_ / 2;
+            conv_to_gemm_transformer_right.Do_ = Do_ - Do_ / 2;
+            // Assign left padding to left convolution
+            conv_to_gemm_transformer_left.InLeftPadD_  = InLeftPadD_;
+            conv_to_gemm_transformer_right.InLeftPadD_ = 0;
+            // Assign right padding to right convolution
+            conv_to_gemm_transformer_left.InRightPadD_  = 0;
+            conv_to_gemm_transformer_right.InRightPadD_ = InRightPadD_;
+            // Calculate new input size
+            conv_to_gemm_transformer_left.Di_ = di_left_transformer_end_idx - InLeftPadD_;
+            conv_to_gemm_transformer_right.Di_ =
+                math::min(Di_ - (di_right_transformer_start_idx - InLeftPadD_),
+                          (conv_to_gemm_transformer_right.Do_ - 1) * ConvStrideD_ + z_eff);
+            ;
+            // Calcualte offsets
+            a_right_offset = (Do_ / 2) * DoStride_;
+            c_right_offset = ((Do_ / 2) * ConvStrideD_ - InLeftPadD_) * DiStride_;
+        }
+        else if(is_possible_to_split_h)
+        {
+            conv_to_gemm_transformer_left.Ho_  = Ho_ / 2;
+            conv_to_gemm_transformer_right.Ho_ = Ho_ - Ho_ / 2;
 
-} // namespace
+            conv_to_gemm_transformer_left.InLeftPadH_  = InLeftPadH_;
+            conv_to_gemm_transformer_right.InLeftPadH_ = 0;
 
-template <
-    index_t NDimSpatial,
-    ck::tensor_operation::device::ConvolutionBackwardDataSpecialization ConvBwdDataSpecialization,
-    index_t AK1,
-    index_t BK1,
-    index_t GemmMPerBlock,
-    index_t GemmNPerBlock,
-    index_t GemmKPerBlock,
-    bool DoPadGemmM,
-    bool DoPadGemmN>
-struct TransformConvBwdDataToGemm_v1
-{
-    static constexpr auto I0 = Number<0>{};
-    static constexpr auto I1 = Number<1>{};
+            conv_to_gemm_transformer_left.InRightPadH_  = 0;
+            conv_to_gemm_transformer_right.InRightPadH_ = InRightPadH_;
 
-    static constexpr auto NonSpatialDimsNum = Number<3>{};
+            conv_to_gemm_transformer_left.Hi_ = hi_left_transformer_end_idx - InLeftPadH_;
+            conv_to_gemm_transformer_right.Hi_ =
+                math::min(Hi_ - (hi_right_transformer_start_idx - InLeftPadH_),
+                          (conv_to_gemm_transformer_right.Ho_ - 1) * ConvStrideH_ + y_eff);
+            a_right_offset = (Ho_ / 2) * HoStride_;
+            c_right_offset = ((Ho_ / 2) * ConvStrideH_ - InLeftPadH_) * HiStride_;
+        }
+        else if(is_possible_to_split_w)
+        {
+            conv_to_gemm_transformer_left.Wo_  = Wo_ / 2;
+            conv_to_gemm_transformer_right.Wo_ = Wo_ - Wo_ / 2;
 
-    static constexpr auto DIdx = Number<NonSpatialDimsNum>{};
-    static constexpr auto HIdx =
-        NDimSpatial == 2 ? Number<NonSpatialDimsNum>{} : Number<NonSpatialDimsNum + 1>{};
-    static constexpr auto WIdx =
-        NDimSpatial == 2 ? Number<NonSpatialDimsNum + 1>{} : Number<NonSpatialDimsNum + 2>{};
+            conv_to_gemm_transformer_left.InLeftPadW_  = InLeftPadW_;
+            conv_to_gemm_transformer_right.InLeftPadW_ = 0;
 
-    static constexpr auto ZIdx = Number<NonSpatialDimsNum>{};
-    static constexpr auto YIdx =
-        NDimSpatial == 2 ? Number<NonSpatialDimsNum>{} : Number<NonSpatialDimsNum + 1>{};
-    static constexpr auto XIdx =
-        NDimSpatial == 2 ? Number<NonSpatialDimsNum + 1>{} : Number<NonSpatialDimsNum + 2>{};
+            conv_to_gemm_transformer_left.InRightPadW_  = 0;
+            conv_to_gemm_transformer_right.InRightPadW_ = InRightPadW_;
 
-    template <typename ALayout,
-              typename std::enable_if<(NDimSpatial == 2 || NDimSpatial == 3) &&
-                                          (is_same_v<ALayout, tensor_layout::convolution::GNHWK> ||
-                                           is_same_v<ALayout, tensor_layout::convolution::GNDHWK> ||
-                                           is_same_v<ALayout, tensor_layout::convolution::NHWGK> ||
-                                           is_same_v<ALayout, tensor_layout::convolution::NDHWGK>),
-                                      bool>::type = false>
-    static auto MakeADescriptor_AK0_M_AK1(
-        const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_lengths,
-        const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_strides,
-        const std::array<index_t, NDimSpatial + 3>& wei_g_k_c_xs_lengths,
-        const std::array<index_t, NDimSpatial + 3>& /* wei_g_k_c_xs_strides */,
-        const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_lengths,
-        const std::array<index_t, NDimSpatial + 3>& /* in_g_n_c_wis_strides */,
-        const std::array<index_t, NDimSpatial>& conv_filter_strides,
-        const std::array<index_t, NDimSpatial>& conv_filter_dilations,
-        const std::array<index_t, NDimSpatial>& input_left_pads,
-        const std::array<index_t, NDimSpatial>& /* input_right_pads */,
-        const std::array<index_t, NDimSpatial>& tildes)
+            conv_to_gemm_transformer_left.Wi_ = wi_left_transformer_end_idx - InLeftPadW_;
+            conv_to_gemm_transformer_right.Wi_ =
+                math::min(Wi_ - (wi_right_transformer_start_idx - InLeftPadW_),
+                          (conv_to_gemm_transformer_right.Wo_ - 1) * ConvStrideW_ + x_eff);
+
+            a_right_offset = (Wo_ / 2) * WoStride_;
+            c_right_offset = ((Wo_ / 2) * ConvStrideW_ - InLeftPadW_) * WiStride_;
+        }
+        // Return left transform, right transformer, right offset to Input and right offset to
+        // Output
+        return ck::make_tuple(conv_to_gemm_transformer_left,
+                              conv_to_gemm_transformer_right,
+                              a_grid_ptr_base + a_right_offset,
+                              c_grid_ptr_base + c_right_offset);
+    }
+
+    __host__ auto SplitConvProblem(const ADataType* a_grid_ptr_base,
+                                   CDataType* c_grid_ptr_base) const
     {
-        index_t i_ztilde = tildes[ZIdx - NonSpatialDimsNum];
-        index_t i_ytilde = tildes[YIdx - NonSpatialDimsNum];
-        index_t i_xtilde = tildes[XIdx - NonSpatialDimsNum];
+        // Create copies
+        auto conv_to_gemm_transformer_left  = *this;
+        auto conv_to_gemm_transformer_right = *this;
+        IndexType a_right_offset            = 0;
+        IndexType c_right_offset            = 0;
+
+        // Calculate start position in input for right tensor
+        const IndexType do_right_transformer_start_idx = math::integer_divide_ceil((Di_ / 2) + InLeftPadD_ - ((Z_ - 1) * ConvDilationD_), ConvStrideD_);
+        const IndexType ho_right_transformer_start_idx = math::integer_divide_ceil((Hi_ / 2) + InLeftPadH_ - ((Y_ - 1) * ConvDilationH_), ConvStrideH_);
+        const IndexType wo_right_transformer_start_idx = math::integer_divide_ceil((Wi_ / 2) + InLeftPadW_ - ((X_ - 1) * ConvDilationW_), ConvStrideW_);
+        // Calculate last position in input for left tensor
+        const IndexType do_left_transformer_end_idx = math::integer_divide_ceil((Di_ / 2 - 1) + InLeftPadD_, ConvStrideD_);
+        const IndexType ho_left_transformer_end_idx = math::integer_divide_ceil((Hi_ / 2 - 1) + InLeftPadH_, ConvStrideH_);
+        const IndexType wo_left_transformer_end_idx = math::integer_divide_ceil((Wi_ / 2 - 1) + InLeftPadW_, ConvStrideW_);
+
+
+        if(Di_!=1)
+        {
+            // Apply new sizes
+            // Split output on half
+            conv_to_gemm_transformer_left.Di_  = Di_ / 2;
+            conv_to_gemm_transformer_right.Di_ = Di_ - Di_ / 2;
+            // Assign left padding to left convolution
+            conv_to_gemm_transformer_left.InLeftPadD_  = InLeftPadD_;
+            conv_to_gemm_transformer_right.InLeftPadD_ = 0;
+            // // Assign right padding to right convolution
+            conv_to_gemm_transformer_left.InRightPadD_  = 0;
+            conv_to_gemm_transformer_right.InRightPadD_ = InRightPadD_;
+            // Calculate new input size
+            conv_to_gemm_transformer_left.Do_ = do_left_transformer_end_idx;
+            conv_to_gemm_transformer_right.Do_ = Do_ - do_right_transformer_start_idx;
+            ;
+            // Calcualte offsets
+            a_right_offset = do_right_transformer_start_idx * DoStride_;
+            c_right_offset = (Di_ / 2) * DiStride_;
+        }
+        else if(Hi_!=1)
+        {
+            // Apply new sizes
+            // Split output on half
+            conv_to_gemm_transformer_left.Hi_  = Hi_ / 2;
+            conv_to_gemm_transformer_right.Hi_ = Hi_ - Hi_ / 2;
+            // Assign left padding to left convolution
+            conv_to_gemm_transformer_left.InLeftPadH_  = InLeftPadH_;
+            conv_to_gemm_transformer_right.InLeftPadH_ = 0;
+            // // Assign right padding to right convolution
+            conv_to_gemm_transformer_left.InRightPadH_  = 0;
+            conv_to_gemm_transformer_right.InRightPadH_ = InRightPadH_;
+            // Calculate new input size
+            conv_to_gemm_transformer_left.Ho_ = ho_left_transformer_end_idx ;
+            conv_to_gemm_transformer_right.Ho_ = Ho_ - ho_right_transformer_start_idx ;
+            ;
+            // Calcualte offsets
+            a_right_offset = ho_right_transformer_start_idx * HoStride_;
+            c_right_offset = (Hi_ / 2) * HiStride_;
+        }
+        else if(Wi_!=1)
+        {
+            // Apply new sizes
+            // Split output on half
+            conv_to_gemm_transformer_left.Wi_  = Wi_ / 2;
+            conv_to_gemm_transformer_right.Wi_ = Wi_ - Wi_ / 2;
+            // Assign left padding to left convolution
+            conv_to_gemm_transformer_left.InLeftPadW_  = InLeftPadW_;
+            conv_to_gemm_transformer_right.InLeftPadW_ = 0;
+            // Assign right padding to right convolution
+            conv_to_gemm_transformer_left.InRightPadW_  = 0;
+            conv_to_gemm_transformer_right.InRightPadW_ = InRightPadW_;
+            // Calculate new input size
+            conv_to_gemm_transformer_left.Wo_ = wo_left_transformer_end_idx;
+            conv_to_gemm_transformer_right.Wo_ = Wo_ - wo_right_transformer_start_idx;
+            ;
+            // Calcualte offsets
+            a_right_offset = wo_right_transformer_start_idx * WoStride_;
+            c_right_offset = (Wi_ / 2) * WiStride_;
+        }
+        // Return left transform, right transformer, right offset to Input and right offset to
+        // Output
+        return ck::make_tuple(conv_to_gemm_transformer_left,
+                              conv_to_gemm_transformer_right,
+                              a_grid_ptr_base + a_right_offset,
+                              c_grid_ptr_base + c_right_offset);
+    }
+#endif
 
-        const index_t N = in_g_n_c_wis_lengths[1];
-        const index_t K = wei_g_k_c_xs_lengths[1];
+    __host__ __device__ auto MakeOutGridDesc() const
+    {
+        if constexpr(is_same_v<ALayout, tensor_layout::convolution::NHWGK>)
+        {
+            if constexpr(ConvBwdDataSpecialization ==
+                         ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
+                             Filter1x1Stride1Pad0)
+            {
 
-        const index_t Di = NDimSpatial == 3 ? in_g_n_c_wis_lengths[DIdx] : 1;
-        const index_t Hi = in_g_n_c_wis_lengths[HIdx];
-        const index_t Wi = in_g_n_c_wis_lengths[WIdx];
+                return make_naive_tensor_descriptor(make_tuple(N_ * Ho_ * Wo_, K_),
+                                                    make_tuple(WoStride_, KStrideTensorA_));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(N_, Ho_, Wo_, K_),
+                    make_tuple(NStrideTensorA_, HoStride_, WoStride_, KStrideTensorA_));
+            }
+        }
+        else if constexpr(is_same_v<ALayout, tensor_layout::convolution::NDHWGK>)
+        {
+            if constexpr(ConvBwdDataSpecialization ==
+                         ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
+                             Filter1x1Stride1Pad0)
+            {
 
-        const index_t Do = NDimSpatial == 3 ? out_g_n_k_wos_lengths[DIdx] : 1;
-        const index_t Ho = out_g_n_k_wos_lengths[HIdx];
-        const index_t Wo = out_g_n_k_wos_lengths[WIdx];
+                return make_naive_tensor_descriptor(make_tuple(N_ * Do_ * Ho_ * Wo_, K_),
+                                                    make_tuple(WoStride_, KStrideTensorA_));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(N_, Do_, Ho_, Wo_, K_),
+                    make_tuple(NStrideTensorA_, DoStride_, HoStride_, WoStride_, KStrideTensorA_));
+            }
+        }
+        else if constexpr(is_same_v<ALayout, tensor_layout::convolution::GNHWK>)
+        {
+            // assume packed
+            if constexpr(ConvBwdDataSpecialization ==
+                         ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
+                             Filter1x1Stride1Pad0)
+            {
+                return make_naive_tensor_descriptor_packed(make_tuple(N_ * Ho_ * Wo_, K_));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_packed(make_tuple(N_, Ho_, Wo_, K_));
+            }
+        }
+        else if constexpr(is_same_v<ALayout, tensor_layout::convolution::GNDHWK>)
+        {
+            // assume packed
+            if constexpr(ConvBwdDataSpecialization ==
+                         ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
+                             Filter1x1Stride1Pad0)
+            {
+                return make_naive_tensor_descriptor_packed(make_tuple(N_ * Do_ * Ho_ * Wo_, K_));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_packed(make_tuple(N_, Do_, Ho_, Wo_, K_));
+            }
+        }
+        else
+        {
+            throw std::runtime_error("wrong! unsupported layout: " + ALayout::name());
+        }
+    }
 
-        const index_t Z = NDimSpatial == 3 ? wei_g_k_c_xs_lengths[ZIdx] : 1;
-        const index_t Y = wei_g_k_c_xs_lengths[YIdx];
-        const index_t X = wei_g_k_c_xs_lengths[XIdx];
+    __host__ __device__ auto MakeWeiGridDesc() const
+    {
 
-        const index_t InLeftPadD = input_left_pads[DIdx - NonSpatialDimsNum];
-        const index_t InLeftPadH = input_left_pads[HIdx - NonSpatialDimsNum];
-        const index_t InLeftPadW = input_left_pads[WIdx - NonSpatialDimsNum];
+        if constexpr(is_same_v<BLayout, tensor_layout::convolution::GKYXC>)
+        {
+            return make_naive_tensor_descriptor_packed(make_tuple(K_, Y_, X_, C_));
+        }
+        else if constexpr(is_same_v<BLayout, tensor_layout::convolution::GKZYXC>)
+        {
+            return make_naive_tensor_descriptor_packed(make_tuple(K_, Z_, Y_, X_, C_));
+        }
+        else
+        {
+            throw std::runtime_error("wrong! unsupported layout: " + BLayout::name());
+        }
+    }
 
-        const index_t ConvStrideD = conv_filter_strides[DIdx - NonSpatialDimsNum];
-        const index_t ConvStrideH = conv_filter_strides[HIdx - NonSpatialDimsNum];
-        const index_t ConvStrideW = conv_filter_strides[WIdx - NonSpatialDimsNum];
+    __host__ __device__ auto MakeInGridDesc() const
+    {
 
-        const index_t ConvDilationD = conv_filter_dilations[DIdx - NonSpatialDimsNum];
-        const index_t ConvDilationH = conv_filter_dilations[HIdx - NonSpatialDimsNum];
-        const index_t ConvDilationW = conv_filter_dilations[WIdx - NonSpatialDimsNum];
+        if constexpr(is_same_v<CLayout, tensor_layout::convolution::GNHWC> ||
+                     is_same_v<CLayout, tensor_layout::convolution::NHWGC> ||
+                     is_same_v<CLayout, tensor_layout::convolution::G_NHW_C>)
+        {
+            return make_naive_tensor_descriptor(
+                make_tuple(N_, Hi_, Wi_, C_),
+                make_tuple(NStrideTensorC_, HiStride_, WiStride_, CStrideTensorC_));
+        }
+        else if constexpr(is_same_v<CLayout, tensor_layout::convolution::GNDHWC> ||
+                          is_same_v<CLayout, tensor_layout::convolution::NDHWGC>)
+        {
+            return make_naive_tensor_descriptor(
+                make_tuple(N_, Di_, Hi_, Wi_, C_),
+                make_tuple(NStrideTensorC_, DiStride_, HiStride_, WiStride_, CStrideTensorC_));
+        }
+        else
+        {
+            throw std::runtime_error("wrong! unsupported layout: " + CLayout::name());
+        }
+    }
 
+    template <
+        typename ALayout_                   = ALayout,
+        typename std::enable_if<(NDimSpatial == 2 || NDimSpatial == 3) &&
+                                    (is_same_v<ALayout_, tensor_layout::convolution::GNHWK> ||
+                                     is_same_v<ALayout_, tensor_layout::convolution::GNDHWK> ||
+                                     is_same_v<ALayout_, tensor_layout::convolution::NHWGK> ||
+                                     is_same_v<ALayout_, tensor_layout::convolution::NDHWGK>),
+                                bool>::type = false>
+    __host__ __device__ auto MakeADescriptor_AK0_M_AK1() const
+    {
         // n_do_ho_wo_k for 3d or n_ho_wo_k for 2d
-        const auto out_grid_desc =
-            make_out_grid_desc<NDimSpatial, ALayout, ConvBwdDataSpecialization>(
-                N, Do, Ho, Wo, K, out_g_n_k_wos_strides);
+        const auto out_grid_desc = MakeOutGridDesc();
 
         if constexpr(ConvBwdDataSpecialization ==
                      ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
                          Filter1x1Stride1Pad0)
         {
-            const index_t AK0 = math::integer_divide_ceil(K, AK1);
+            const index_t AK0 = math::integer_divide_ceil(K_, AK1);
 
             // A: output tensor
             const auto out_gemmak0_gemmmraw_gemmak1_grid_desc = transform_tensor_descriptor(
                 out_grid_desc,
-                make_tuple(make_pass_through_transform(N * Do * Ho * Wo),
+                make_tuple(make_pass_through_transform(N_ * Do_ * Ho_ * Wo_),
                            make_unmerge_transform(make_tuple(AK0, AK1))),
                 make_tuple(Sequence<0>{}, Sequence<1>{}),
                 make_tuple(Sequence<1>{}, Sequence<0, 2>{}));
@@ -266,82 +635,63 @@ struct TransformConvBwdDataToGemm_v1
         }
         else
         {
-            const auto GcdStrideDilationD = math::gcd(ConvStrideD, ConvDilationD);
-            const auto GcdStrideDilationH = math::gcd(ConvStrideH, ConvDilationH);
-            const auto GcdStrideDilationW = math::gcd(ConvStrideW, ConvDilationW);
-
-            const auto ZTilde = ConvStrideD / GcdStrideDilationD;
-            const auto YTilde = ConvStrideH / GcdStrideDilationH;
-            const auto XTilde = ConvStrideW / GcdStrideDilationW;
-
-            const auto ZDot = math::integer_divide_ceil(Z, ZTilde);
-            const auto YDot = math::integer_divide_ceil(Y, YTilde);
-            const auto XDot = math::integer_divide_ceil(X, XTilde);
-
-            const auto DTilde =
-                Do + math::integer_divide_ceil(ConvDilationD * (Z - I1), ConvStrideD);
-            const auto HTilde =
-                Ho + math::integer_divide_ceil(ConvDilationH * (Y - I1), ConvStrideH);
-            const auto WTilde =
-                Wo + math::integer_divide_ceil(ConvDilationW * (X - I1), ConvStrideW);
-
             // only work on HTilde and WTilde that contribute to non-padding area of input tensor
             const auto IDTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadD - ConvDilationD * (ZTilde - I1)), ConvStrideD);
+                math::max(I0, InLeftPadD_ - ConvDilationD_ * (ZTilde_ - I1)), ConvStrideD_);
             const auto IHTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadH - ConvDilationH * (YTilde - I1)), ConvStrideH);
+                math::max(I0, InLeftPadH_ - ConvDilationH_ * (YTilde_ - I1)), ConvStrideH_);
             const auto IWTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadW - ConvDilationW * (XTilde - I1)), ConvStrideW);
+                math::max(I0, InLeftPadW_ - ConvDilationW_ * (XTilde_ - I1)), ConvStrideW_);
 
             const auto IDTildeSliceEnd = math::min(
-                DTilde, math::integer_divide_ceil(InLeftPadD + Di - I1, ConvStrideD) + I1);
+                DTilde_, math::integer_divide_ceil(InLeftPadD_ + Di_ - I1, ConvStrideD_) + I1);
             const auto IHTildeSliceEnd = math::min(
-                HTilde, math::integer_divide_ceil(InLeftPadH + Hi - I1, ConvStrideH) + I1);
+                HTilde_, math::integer_divide_ceil(InLeftPadH_ + Hi_ - I1, ConvStrideH_) + I1);
             const auto IWTildeSliceEnd = math::min(
-                WTilde, math::integer_divide_ceil(InLeftPadW + Wi - I1, ConvStrideW) + I1);
+                WTilde_, math::integer_divide_ceil(InLeftPadW_ + Wi_ - I1, ConvStrideW_) + I1);
 
             const auto DTildeSlice = IDTildeSliceEnd - IDTildeSliceBegin;
             const auto HTildeSlice = IHTildeSliceEnd - IHTildeSliceBegin;
             const auto WTildeSlice = IWTildeSliceEnd - IWTildeSliceBegin;
 
             // GemmK is different for each GEMM
-            const auto ZDotSlice = math::integer_divide_ceil(Z - i_ztilde, ZTilde);
-            const auto YDotSlice = math::integer_divide_ceil(Y - i_ytilde, YTilde);
-            const auto XDotSlice = math::integer_divide_ceil(X - i_xtilde, XTilde);
+            const auto ZDotSlice = math::integer_divide_ceil(Z_ - IdxZTilde_, ZTilde_);
+            const auto YDotSlice = math::integer_divide_ceil(Y_ - IdxYTilde_, YTilde_);
+            const auto XDotSlice = math::integer_divide_ceil(X_ - IdxXTilde_, XTilde_);
 
             if constexpr(NDimSpatial == 2)
             {
                 // A: output tensor
                 const auto out_n_hop_wop_k_grid_desc = transform_tensor_descriptor(
                     out_grid_desc,
-                    make_tuple(make_pass_through_transform(N),
-                               make_pad_transform(Ho, I0, I0),
-                               make_pad_transform(Wo, I0, I0),
-                               make_pass_through_transform(K)),
+                    make_tuple(make_pass_through_transform(N_),
+                               make_pad_transform(Ho_, I0, I0),
+                               make_pad_transform(Wo_, I0, I0),
+                               make_pass_through_transform(K_)),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
 
                 const auto out_n_ydot_htilde_xdot_wtilde_k_grid_desc = transform_tensor_descriptor(
                     out_n_hop_wop_k_grid_desc,
                     make_tuple(
-                        make_pass_through_transform(N),
-                        make_embed_transform(make_tuple(YDot, HTilde),
-                                             make_tuple(-ConvDilationH / GcdStrideDilationH, I1)),
-                        make_embed_transform(make_tuple(XDot, WTilde),
-                                             make_tuple(-ConvDilationW / GcdStrideDilationW, I1)),
-                        make_pass_through_transform(K)),
+                        make_pass_through_transform(N_),
+                        make_embed_transform(make_tuple(YDot_, HTilde_),
+                                             make_tuple(-ConvDilationH_ / GcdStrideDilationH_, I1)),
+                        make_embed_transform(make_tuple(XDot_, WTilde_),
+                                             make_tuple(-ConvDilationW_ / GcdStrideDilationW_, I1)),
+                        make_pass_through_transform(K_)),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
 
                 const auto out_n_ydotslice_htildeslice_xdotslice_wtildeslice_k_grid_desc =
                     transform_tensor_descriptor(
                         out_n_ydot_htilde_xdot_wtilde_k_grid_desc,
-                        make_tuple(make_pass_through_transform(N),
-                                   make_slice_transform(YDot, I0, YDotSlice),
-                                   make_slice_transform(HTilde, IHTildeSliceBegin, HTildeSlice),
-                                   make_slice_transform(XDot, I0, XDotSlice),
-                                   make_slice_transform(WTilde, IWTildeSliceBegin, WTildeSlice),
-                                   make_pass_through_transform(K)),
+                        make_tuple(make_pass_through_transform(N_),
+                                   make_slice_transform(YDot_, I0, YDotSlice),
+                                   make_slice_transform(HTilde_, IHTildeSliceBegin, HTildeSlice),
+                                   make_slice_transform(XDot_, I0, XDotSlice),
+                                   make_slice_transform(WTilde_, IWTildeSliceBegin, WTildeSlice),
+                                   make_pass_through_transform(K_)),
                         make_tuple(Sequence<0>{},
                                    Sequence<1>{},
                                    Sequence<2>{},
@@ -357,8 +707,8 @@ struct TransformConvBwdDataToGemm_v1
 
                 const auto out_gemmk_gemmmraw_grid_desc = transform_tensor_descriptor(
                     out_n_ydotslice_htildeslice_xdotslice_wtildeslice_k_grid_desc,
-                    make_tuple(make_merge_transform(make_tuple(YDotSlice, XDotSlice, K)),
-                               make_merge_transform(make_tuple(N, HTildeSlice, WTildeSlice))),
+                    make_tuple(make_merge_transform(make_tuple(YDotSlice, XDotSlice, K_)),
+                               make_merge_transform(make_tuple(N_, HTildeSlice, WTildeSlice))),
                     make_tuple(Sequence<1, 3, 5>{}, Sequence<0, 2, 4>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}));
 
@@ -385,11 +735,11 @@ struct TransformConvBwdDataToGemm_v1
                 // A: output tensor
                 const auto out_n_hop_wop_k_grid_desc = transform_tensor_descriptor(
                     out_grid_desc,
-                    make_tuple(make_pass_through_transform(N),
-                               make_pad_transform(Do, I0, I0),
-                               make_pad_transform(Ho, I0, I0),
-                               make_pad_transform(Wo, I0, I0),
-                               make_pass_through_transform(K)),
+                    make_tuple(make_pass_through_transform(N_),
+                               make_pad_transform(Do_, I0, I0),
+                               make_pad_transform(Ho_, I0, I0),
+                               make_pad_transform(Wo_, I0, I0),
+                               make_pass_through_transform(K_)),
                     make_tuple(
                         Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
                     make_tuple(
@@ -398,17 +748,17 @@ struct TransformConvBwdDataToGemm_v1
                 const auto out_n_zdot_dtilde_ydot_htilde_xdot_wtilde_k_grid_desc =
                     transform_tensor_descriptor(
                         out_n_hop_wop_k_grid_desc,
-                        make_tuple(make_pass_through_transform(N),
+                        make_tuple(make_pass_through_transform(N_),
                                    make_embed_transform(
-                                       make_tuple(ZDot, DTilde),
-                                       make_tuple(-ConvDilationD / GcdStrideDilationD, I1)),
+                                       make_tuple(ZDot_, DTilde_),
+                                       make_tuple(-ConvDilationD_ / GcdStrideDilationD_, I1)),
                                    make_embed_transform(
-                                       make_tuple(YDot, HTilde),
-                                       make_tuple(-ConvDilationH / GcdStrideDilationH, I1)),
+                                       make_tuple(YDot_, HTilde_),
+                                       make_tuple(-ConvDilationH_ / GcdStrideDilationH_, I1)),
                                    make_embed_transform(
-                                       make_tuple(XDot, WTilde),
-                                       make_tuple(-ConvDilationW / GcdStrideDilationW, I1)),
-                                   make_pass_through_transform(K)),
+                                       make_tuple(XDot_, WTilde_),
+                                       make_tuple(-ConvDilationW_ / GcdStrideDilationW_, I1)),
+                                   make_pass_through_transform(K_)),
                         make_tuple(Sequence<0>{},
                                    Sequence<1>{},
                                    Sequence<2>{},
@@ -424,14 +774,15 @@ struct TransformConvBwdDataToGemm_v1
                     out_n_zdotslice_dtildeslice_ydotslice_htildeslice_xdotslice_wtildeslice_k_grid_desc =
                         transform_tensor_descriptor(
                             out_n_zdot_dtilde_ydot_htilde_xdot_wtilde_k_grid_desc,
-                            make_tuple(make_pass_through_transform(N),
-                                       make_slice_transform(ZDot, I0, ZDotSlice),
-                                       make_slice_transform(DTilde, IDTildeSliceBegin, DTildeSlice),
-                                       make_slice_transform(YDot, I0, YDotSlice),
-                                       make_slice_transform(HTilde, IHTildeSliceBegin, HTildeSlice),
-                                       make_slice_transform(XDot, I0, XDotSlice),
-                                       make_slice_transform(WTilde, IWTildeSliceBegin, WTildeSlice),
-                                       make_pass_through_transform(K)),
+                            make_tuple(
+                                make_pass_through_transform(N_),
+                                make_slice_transform(ZDot_, I0, ZDotSlice),
+                                make_slice_transform(DTilde_, IDTildeSliceBegin, DTildeSlice),
+                                make_slice_transform(YDot_, I0, YDotSlice),
+                                make_slice_transform(HTilde_, IHTildeSliceBegin, HTildeSlice),
+                                make_slice_transform(XDot_, I0, XDotSlice),
+                                make_slice_transform(WTilde_, IWTildeSliceBegin, WTildeSlice),
+                                make_pass_through_transform(K_)),
                             make_tuple(Sequence<0>{},
                                        Sequence<1>{},
                                        Sequence<2>{},
@@ -452,8 +803,9 @@ struct TransformConvBwdDataToGemm_v1
                 const auto out_gemmk_gemmmraw_grid_desc = transform_tensor_descriptor(
                     out_n_zdotslice_dtildeslice_ydotslice_htildeslice_xdotslice_wtildeslice_k_grid_desc,
                     make_tuple(
-                        make_merge_transform(make_tuple(ZDotSlice, YDotSlice, XDotSlice, K)),
-                        make_merge_transform(make_tuple(N, DTildeSlice, HTildeSlice, WTildeSlice))),
+                        make_merge_transform(make_tuple(ZDotSlice, YDotSlice, XDotSlice, K_)),
+                        make_merge_transform(
+                            make_tuple(N_, DTildeSlice, HTildeSlice, WTildeSlice))),
                     make_tuple(Sequence<1, 3, 5, 7>{}, Sequence<0, 2, 4, 6>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}));
 
@@ -482,66 +834,31 @@ struct TransformConvBwdDataToGemm_v1
         }
     }
 
-    template <typename BLayout,
+    template <typename BLayout_                   = BLayout,
               typename std::enable_if<(NDimSpatial == 2 || NDimSpatial == 3) &&
-                                          (is_same_v<BLayout, tensor_layout::convolution::GKYXC> ||
-                                           is_same_v<BLayout, tensor_layout::convolution::GKZYXC>),
+                                          (is_same_v<BLayout_, tensor_layout::convolution::GKYXC> ||
+                                           is_same_v<BLayout_, tensor_layout::convolution::GKZYXC>),
                                       bool>::type = false>
-    static auto MakeBDescriptor_BK0_N_BK1(
-        const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_lengths,
-        const std::array<index_t, NDimSpatial + 3>& /* out_g_n_k_wos_strides */,
-        const std::array<index_t, NDimSpatial + 3>& wei_g_k_c_xs_lengths,
-        const std::array<index_t, NDimSpatial + 3>& /* wei_g_k_c_xs_strides */,
-        const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_lengths,
-        const std::array<index_t, NDimSpatial + 3>& /* in_g_n_c_wis_strides */,
-        const std::array<index_t, NDimSpatial>& conv_filter_strides,
-        const std::array<index_t, NDimSpatial>& conv_filter_dilations,
-        const std::array<index_t, NDimSpatial>& /* input_left_pads */,
-        const std::array<index_t, NDimSpatial>& /* input_right_pads */,
-        const std::array<index_t, NDimSpatial>& tildes)
+    __host__ __device__ auto MakeBDescriptor_BK0_N_BK1() const
     {
-        index_t i_ztilde = tildes[ZIdx - NonSpatialDimsNum];
-        index_t i_ytilde = tildes[YIdx - NonSpatialDimsNum];
-        index_t i_xtilde = tildes[XIdx - NonSpatialDimsNum];
-
-        const index_t N = in_g_n_c_wis_lengths[1];
-        const index_t K = wei_g_k_c_xs_lengths[1];
-        const index_t C = wei_g_k_c_xs_lengths[2];
-
-        const index_t Do = NDimSpatial == 3 ? out_g_n_k_wos_lengths[DIdx] : 1;
-        const index_t Ho = out_g_n_k_wos_lengths[HIdx];
-        const index_t Wo = out_g_n_k_wos_lengths[WIdx];
-
-        const index_t Z = NDimSpatial == 3 ? wei_g_k_c_xs_lengths[ZIdx] : 1;
-        const index_t Y = wei_g_k_c_xs_lengths[YIdx];
-        const index_t X = wei_g_k_c_xs_lengths[XIdx];
-
-        const index_t ConvStrideD = conv_filter_strides[DIdx - NonSpatialDimsNum];
-        const index_t ConvStrideH = conv_filter_strides[HIdx - NonSpatialDimsNum];
-        const index_t ConvStrideW = conv_filter_strides[WIdx - NonSpatialDimsNum];
-
-        const index_t ConvDilationD = conv_filter_dilations[DIdx - NonSpatialDimsNum];
-        const index_t ConvDilationH = conv_filter_dilations[HIdx - NonSpatialDimsNum];
-        const index_t ConvDilationW = conv_filter_dilations[WIdx - NonSpatialDimsNum];
-
         // assume packed
         // k_y_x_c for 2d or k_z_y_x_c for 3d
-        const auto wei_grid_desc = make_wei_grid_desc<BLayout>(K, Z, Y, X, C);
+        const auto wei_grid_desc = MakeWeiGridDesc();
 
         if constexpr(ConvBwdDataSpecialization ==
                      ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
                          Filter1x1Stride1Pad0)
         {
-            const index_t BK0 = math::integer_divide_ceil(K, BK1);
+            const index_t BK0 = math::integer_divide_ceil(K_, BK1);
 
             // B: weight tensor
             const auto wei_gemmbk0_gemmnraw_gemmbk1_grid_desc =
-                transform_tensor_descriptor(make_naive_tensor_descriptor_packed(make_tuple(K, C)),
+                transform_tensor_descriptor(make_naive_tensor_descriptor_packed(make_tuple(K_, C_)),
                                             make_tuple(make_unmerge_transform(make_tuple(BK0, BK1)),
-                                                       make_pass_through_transform(C)),
+                                                       make_pass_through_transform(C_)),
                                             make_tuple(Sequence<0>{}, Sequence<1>{}),
                                             make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
-            make_naive_tensor_descriptor(make_tuple(N * Do * Ho * Wo, C), make_tuple(I0, I1));
+            make_naive_tensor_descriptor(make_tuple(N_ * Do_ * Ho_ * Wo_, C_), make_tuple(I0, I1));
 
             const auto wei_gemmbk0_gemmn_gemmbk1_grid_desc =
                 ck::tensor_operation::device::PadTensorDescriptor(
@@ -553,22 +870,10 @@ struct TransformConvBwdDataToGemm_v1
         }
         else
         {
-            const auto GcdStrideDilationD = math::gcd(ConvStrideD, ConvDilationD);
-            const auto GcdStrideDilationH = math::gcd(ConvStrideH, ConvDilationH);
-            const auto GcdStrideDilationW = math::gcd(ConvStrideW, ConvDilationW);
-
-            const auto ZTilde = ConvStrideD / GcdStrideDilationD;
-            const auto YTilde = ConvStrideH / GcdStrideDilationH;
-            const auto XTilde = ConvStrideW / GcdStrideDilationW;
-
-            const auto ZDot = math::integer_divide_ceil(Z, ZTilde);
-            const auto YDot = math::integer_divide_ceil(Y, YTilde);
-            const auto XDot = math::integer_divide_ceil(X, XTilde);
-
             // GemmK is different for each GEMM
-            const auto ZDotSlice = math::integer_divide_ceil(Z - i_ztilde, ZTilde);
-            const auto YDotSlice = math::integer_divide_ceil(Y - i_ytilde, YTilde);
-            const auto XDotSlice = math::integer_divide_ceil(X - i_xtilde, XTilde);
+            const auto ZDotSlice = math::integer_divide_ceil(Z_ - IdxZTilde_, ZTilde_);
+            const auto YDotSlice = math::integer_divide_ceil(Y_ - IdxYTilde_, YTilde_);
+            const auto XDotSlice = math::integer_divide_ceil(X_ - IdxXTilde_, XTilde_);
 
             // B weight tensor
             if constexpr(NDimSpatial == 2)
@@ -576,23 +881,23 @@ struct TransformConvBwdDataToGemm_v1
                 const auto wei_k_ydot_ytilde_xdot_xtilde_c_grid_desc = transform_tensor_descriptor(
                     wei_grid_desc,
                     make_tuple(
-                        make_pass_through_transform(K),
-                        make_embed_transform(make_tuple(YDot, YTilde),
-                                             make_tuple(ConvStrideH / GcdStrideDilationH, I1)),
-                        make_embed_transform(make_tuple(XDot, XTilde),
-                                             make_tuple(ConvStrideW / GcdStrideDilationW, I1)),
-                        make_pass_through_transform(C)),
+                        make_pass_through_transform(K_),
+                        make_embed_transform(make_tuple(YDot_, YTilde_),
+                                             make_tuple(ConvStrideH_ / GcdStrideDilationH_, I1)),
+                        make_embed_transform(make_tuple(XDot_, XTilde_),
+                                             make_tuple(ConvStrideW_ / GcdStrideDilationW_, I1)),
+                        make_pass_through_transform(C_)),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
 
                 const auto wei_k_ydotslice_xdotslice_c_grid_desc = transform_tensor_descriptor(
                     wei_k_ydot_ytilde_xdot_xtilde_c_grid_desc,
-                    make_tuple(make_pass_through_transform(K),
-                               make_slice_transform(YDot, I0, YDotSlice),
-                               make_slice_transform(XDot, I0, XDotSlice),
-                               make_freeze_transform(i_ytilde),
-                               make_freeze_transform(i_xtilde),
-                               make_pass_through_transform(C)),
+                    make_tuple(make_pass_through_transform(K_),
+                               make_slice_transform(YDot_, I0, YDotSlice),
+                               make_slice_transform(XDot_, I0, XDotSlice),
+                               make_freeze_transform(IdxYTilde_),
+                               make_freeze_transform(IdxXTilde_),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<0>{},
                                Sequence<1>{},
                                Sequence<3>{},
@@ -608,8 +913,8 @@ struct TransformConvBwdDataToGemm_v1
 
                 const auto wei_gemmk_gemmnraw_grid_desc = transform_tensor_descriptor(
                     wei_k_ydotslice_xdotslice_c_grid_desc,
-                    make_tuple(make_merge_transform(make_tuple(YDotSlice, XDotSlice, K)),
-                               make_pass_through_transform(C)),
+                    make_tuple(make_merge_transform(make_tuple(YDotSlice, XDotSlice, K_)),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<1, 2, 0>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}));
 
@@ -636,15 +941,17 @@ struct TransformConvBwdDataToGemm_v1
                 const auto wei_k_zdot_ztilde_ydot_ytilde_xdot_xtilde_c_grid_desc =
                     transform_tensor_descriptor(
                         wei_grid_desc,
-                        make_tuple(
-                            make_pass_through_transform(K),
-                            make_embed_transform(make_tuple(ZDot, ZTilde),
-                                                 make_tuple(ConvStrideD / GcdStrideDilationD, I1)),
-                            make_embed_transform(make_tuple(YDot, YTilde),
-                                                 make_tuple(ConvStrideH / GcdStrideDilationH, I1)),
-                            make_embed_transform(make_tuple(XDot, XTilde),
-                                                 make_tuple(ConvStrideW / GcdStrideDilationW, I1)),
-                            make_pass_through_transform(C)),
+                        make_tuple(make_pass_through_transform(K_),
+                                   make_embed_transform(
+                                       make_tuple(ZDot_, ZTilde_),
+                                       make_tuple(ConvStrideD_ / GcdStrideDilationD_, I1)),
+                                   make_embed_transform(
+                                       make_tuple(YDot_, YTilde_),
+                                       make_tuple(ConvStrideH_ / GcdStrideDilationH_, I1)),
+                                   make_embed_transform(
+                                       make_tuple(XDot_, XTilde_),
+                                       make_tuple(ConvStrideW_ / GcdStrideDilationW_, I1)),
+                                   make_pass_through_transform(C_)),
                         make_tuple(Sequence<0>{},
                                    Sequence<1>{},
                                    Sequence<2>{},
@@ -659,14 +966,14 @@ struct TransformConvBwdDataToGemm_v1
                 const auto wei_gemmk_zdotslice_ydotslice_xdotslice_c_grid_desc =
                     transform_tensor_descriptor(
                         wei_k_zdot_ztilde_ydot_ytilde_xdot_xtilde_c_grid_desc,
-                        make_tuple(make_pass_through_transform(K),
-                                   make_slice_transform(ZDot, I0, ZDotSlice),
-                                   make_slice_transform(YDot, I0, YDotSlice),
-                                   make_slice_transform(XDot, I0, XDotSlice),
-                                   make_freeze_transform(i_ztilde),
-                                   make_freeze_transform(i_ytilde),
-                                   make_freeze_transform(i_xtilde),
-                                   make_pass_through_transform(C)),
+                        make_tuple(make_pass_through_transform(K_),
+                                   make_slice_transform(ZDot_, I0, ZDotSlice),
+                                   make_slice_transform(YDot_, I0, YDotSlice),
+                                   make_slice_transform(XDot_, I0, XDotSlice),
+                                   make_freeze_transform(IdxZTilde_),
+                                   make_freeze_transform(IdxYTilde_),
+                                   make_freeze_transform(IdxXTilde_),
+                                   make_pass_through_transform(C_)),
                         make_tuple(Sequence<0>{},
                                    Sequence<1>{},
                                    Sequence<3>{},
@@ -686,8 +993,9 @@ struct TransformConvBwdDataToGemm_v1
 
                 const auto wei_gemmk_gemmnraw_grid_desc = transform_tensor_descriptor(
                     wei_gemmk_zdotslice_ydotslice_xdotslice_c_grid_desc,
-                    make_tuple(make_merge_transform(make_tuple(ZDotSlice, YDotSlice, XDotSlice, K)),
-                               make_pass_through_transform(C)),
+                    make_tuple(
+                        make_merge_transform(make_tuple(ZDotSlice, YDotSlice, XDotSlice, K_)),
+                        make_pass_through_transform(C_)),
                     make_tuple(Sequence<1, 2, 3, 0>{}, Sequence<4>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}));
 
@@ -716,66 +1024,20 @@ struct TransformConvBwdDataToGemm_v1
         }
     }
 
-    template <typename CLayout,
-              typename std::enable_if<(NDimSpatial == 2 || NDimSpatial == 3) &&
-                                          (is_same_v<CLayout, tensor_layout::convolution::GNHWC> ||
-                                           is_same_v<CLayout, tensor_layout::convolution::GNDHWC> ||
-                                           is_same_v<CLayout, tensor_layout::convolution::NHWGC> ||
-                                           is_same_v<CLayout, tensor_layout::convolution::NDHWGC> ||
-                                           is_same_v<CLayout, tensor_layout::convolution::G_NHW_C>),
-                                      bool>::type = false>
-    static auto
-    MakeCDescriptor_M_N(const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_lengths,
-                        const std::array<index_t, NDimSpatial + 3>& /* out_g_n_k_wos_strides */,
-                        const std::array<index_t, NDimSpatial + 3>& wei_g_k_c_xs_lengths,
-                        const std::array<index_t, NDimSpatial + 3>& /* wei_g_k_c_xs_strides */,
-                        const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_lengths,
-                        const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_strides,
-                        const std::array<index_t, NDimSpatial>& conv_filter_strides,
-                        const std::array<index_t, NDimSpatial>& conv_filter_dilations,
-                        const std::array<index_t, NDimSpatial>& input_left_pads,
-                        const std::array<index_t, NDimSpatial>& input_right_pads,
-                        const std::array<index_t, NDimSpatial>& tildes)
+    template <
+        typename CLayout_                   = CLayout,
+        typename std::enable_if<(NDimSpatial == 2 || NDimSpatial == 3) &&
+                                    (is_same_v<CLayout_, tensor_layout::convolution::GNHWC> ||
+                                     is_same_v<CLayout_, tensor_layout::convolution::GNDHWC> ||
+                                     is_same_v<CLayout_, tensor_layout::convolution::NHWGC> ||
+                                     is_same_v<CLayout_, tensor_layout::convolution::NDHWGC> ||
+                                     is_same_v<CLayout_, tensor_layout::convolution::G_NHW_C>),
+                                bool>::type = false>
+    __host__ __device__ auto MakeCDescriptor_M_N() const
     {
-        index_t i_ztilde = tildes[ZIdx - NonSpatialDimsNum];
-        index_t i_ytilde = tildes[YIdx - NonSpatialDimsNum];
-        index_t i_xtilde = tildes[XIdx - NonSpatialDimsNum];
-
-        const index_t N = in_g_n_c_wis_lengths[1];
-        const index_t C = wei_g_k_c_xs_lengths[2];
-
-        const index_t Di = NDimSpatial == 3 ? in_g_n_c_wis_lengths[DIdx] : 1;
-        const index_t Hi = in_g_n_c_wis_lengths[HIdx];
-        const index_t Wi = in_g_n_c_wis_lengths[WIdx];
-
-        const index_t Do = NDimSpatial == 3 ? out_g_n_k_wos_lengths[DIdx] : 1;
-        const index_t Ho = out_g_n_k_wos_lengths[HIdx];
-        const index_t Wo = out_g_n_k_wos_lengths[WIdx];
-
-        const index_t Z = NDimSpatial == 3 ? wei_g_k_c_xs_lengths[ZIdx] : 1;
-        const index_t Y = wei_g_k_c_xs_lengths[YIdx];
-        const index_t X = wei_g_k_c_xs_lengths[XIdx];
-
-        const index_t InLeftPadD = input_left_pads[DIdx - NonSpatialDimsNum];
-        const index_t InLeftPadH = input_left_pads[HIdx - NonSpatialDimsNum];
-        const index_t InLeftPadW = input_left_pads[WIdx - NonSpatialDimsNum];
-
-        const index_t InRightPadD = input_right_pads[DIdx - NonSpatialDimsNum];
-        const index_t InRightPadH = input_right_pads[HIdx - NonSpatialDimsNum];
-        const index_t InRightPadW = input_right_pads[WIdx - NonSpatialDimsNum];
-
-        const index_t ConvStrideD = conv_filter_strides[DIdx - NonSpatialDimsNum];
-        const index_t ConvStrideH = conv_filter_strides[HIdx - NonSpatialDimsNum];
-        const index_t ConvStrideW = conv_filter_strides[WIdx - NonSpatialDimsNum];
-
-        const index_t ConvDilationD = conv_filter_dilations[DIdx - NonSpatialDimsNum];
-        const index_t ConvDilationH = conv_filter_dilations[HIdx - NonSpatialDimsNum];
-        const index_t ConvDilationW = conv_filter_dilations[WIdx - NonSpatialDimsNum];
-
         // assume strided
         // n_hi_wi_c for 2d n_di_hi_wi_c for 3d
-        const auto in_grid_desc =
-            make_in_grid_desc<NDimSpatial, CLayout>(N, Di, Hi, Wi, C, in_g_n_c_wis_strides);
+        const auto in_grid_desc = MakeInGridDesc();
 
         if constexpr(ConvBwdDataSpecialization ==
                      ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
@@ -787,10 +1049,10 @@ struct TransformConvBwdDataToGemm_v1
                 const auto in_n_y_ho_x_wo_c_grid_desc = transform_tensor_descriptor(
                     in_grid_desc,
                     make_tuple(
-                        make_pass_through_transform(N),
-                        make_embed_transform(make_tuple(I1, Ho), make_tuple(I1, ConvStrideH)),
-                        make_embed_transform(make_tuple(I1, Wo), make_tuple(I1, ConvStrideW)),
-                        make_pass_through_transform(C)),
+                        make_pass_through_transform(N_),
+                        make_embed_transform(make_tuple(I1, Ho_), make_tuple(I1, ConvStrideH_)),
+                        make_embed_transform(make_tuple(I1, Wo_), make_tuple(I1, ConvStrideW_)),
+                        make_pass_through_transform(C_)),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
 
@@ -798,8 +1060,8 @@ struct TransformConvBwdDataToGemm_v1
                     in_n_y_ho_x_wo_c_grid_desc,
                     make_tuple(make_freeze_transform(I0),
                                make_freeze_transform(I0),
-                               make_merge_transform(make_tuple(N, Ho, Wo)),
-                               make_pass_through_transform(C)),
+                               make_merge_transform(make_tuple(N_, Ho_, Wo_)),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<1>{}, Sequence<3>{}, Sequence<0, 2, 4>{}, Sequence<5>{}),
                     make_tuple(Sequence<>{}, Sequence<>{}, Sequence<0>{}, Sequence<1>{}));
 
@@ -818,11 +1080,11 @@ struct TransformConvBwdDataToGemm_v1
                 const auto in_n_x_do_y_ho_x_wo_c_grid_desc = transform_tensor_descriptor(
                     in_grid_desc,
                     make_tuple(
-                        make_pass_through_transform(N),
-                        make_embed_transform(make_tuple(I1, Do), make_tuple(I1, ConvStrideD)),
-                        make_embed_transform(make_tuple(I1, Ho), make_tuple(I1, ConvStrideH)),
-                        make_embed_transform(make_tuple(I1, Wo), make_tuple(I1, ConvStrideW)),
-                        make_pass_through_transform(C)),
+                        make_pass_through_transform(N_),
+                        make_embed_transform(make_tuple(I1, Do_), make_tuple(I1, ConvStrideD_)),
+                        make_embed_transform(make_tuple(I1, Ho_), make_tuple(I1, ConvStrideH_)),
+                        make_embed_transform(make_tuple(I1, Wo_), make_tuple(I1, ConvStrideW_)),
+                        make_pass_through_transform(C_)),
                     make_tuple(
                         Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
                     make_tuple(Sequence<0>{},
@@ -836,8 +1098,8 @@ struct TransformConvBwdDataToGemm_v1
                     make_tuple(make_freeze_transform(I0),
                                make_freeze_transform(I0),
                                make_freeze_transform(I0),
-                               make_merge_transform(make_tuple(N, Do, Ho, Wo)),
-                               make_pass_through_transform(C)),
+                               make_merge_transform(make_tuple(N_, Do_, Ho_, Wo_)),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<1>{},
                                Sequence<3>{},
                                Sequence<5>{},
@@ -861,36 +1123,21 @@ struct TransformConvBwdDataToGemm_v1
         }
         else
         {
-            const auto GcdStrideDilationD = math::gcd(ConvStrideD, ConvDilationD);
-            const auto GcdStrideDilationH = math::gcd(ConvStrideH, ConvDilationH);
-            const auto GcdStrideDilationW = math::gcd(ConvStrideW, ConvDilationW);
-
-            const auto ZTilde = ConvStrideD / GcdStrideDilationD;
-            const auto YTilde = ConvStrideH / GcdStrideDilationH;
-            const auto XTilde = ConvStrideW / GcdStrideDilationW;
-
-            const auto DTilde =
-                Do + math::integer_divide_ceil(ConvDilationD * (Z - I1), ConvStrideD);
-            const auto HTilde =
-                Ho + math::integer_divide_ceil(ConvDilationH * (Y - I1), ConvStrideH);
-            const auto WTilde =
-                Wo + math::integer_divide_ceil(ConvDilationW * (X - I1), ConvStrideW);
-
             // only work on DTilde, HTilde and WTilde that contribute to
             // non-padding area of input tensor
             const auto IDTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadD - ConvDilationD * (ZTilde - I1)), ConvStrideD);
+                math::max(I0, InLeftPadD_ - ConvDilationD_ * (ZTilde_ - I1)), ConvStrideD_);
             const auto IHTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadH - ConvDilationH * (YTilde - I1)), ConvStrideH);
+                math::max(I0, InLeftPadH_ - ConvDilationH_ * (YTilde_ - I1)), ConvStrideH_);
             const auto IWTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadW - ConvDilationW * (XTilde - I1)), ConvStrideW);
+                math::max(I0, InLeftPadW_ - ConvDilationW_ * (XTilde_ - I1)), ConvStrideW_);
 
             const auto IDTildeSliceEnd = math::min(
-                DTilde, math::integer_divide_ceil(InLeftPadD + Di - I1, ConvStrideD) + I1);
+                DTilde_, math::integer_divide_ceil(InLeftPadD_ + Di_ - I1, ConvStrideD_) + I1);
             const auto IHTildeSliceEnd = math::min(
-                HTilde, math::integer_divide_ceil(InLeftPadH + Hi - I1, ConvStrideH) + I1);
+                HTilde_, math::integer_divide_ceil(InLeftPadH_ + Hi_ - I1, ConvStrideH_) + I1);
             const auto IWTildeSliceEnd = math::min(
-                WTilde, math::integer_divide_ceil(InLeftPadW + Wi - I1, ConvStrideW) + I1);
+                WTilde_, math::integer_divide_ceil(InLeftPadW_ + Wi_ - I1, ConvStrideW_) + I1);
 
             const auto DTildeSlice = IDTildeSliceEnd - IDTildeSliceBegin;
             const auto HTildeSlice = IHTildeSliceEnd - IHTildeSliceBegin;
@@ -901,34 +1148,34 @@ struct TransformConvBwdDataToGemm_v1
             {
                 const auto in_n_hip_wip_c_grid_desc = transform_tensor_descriptor(
                     in_grid_desc,
-                    make_tuple(make_pass_through_transform(N),
-                               make_pad_transform(Hi, InLeftPadH, InRightPadH),
-                               make_pad_transform(Wi, InLeftPadW, InRightPadW),
-                               make_pass_through_transform(C)),
+                    make_tuple(make_pass_through_transform(N_),
+                               make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
+                               make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
 
                 const auto in_n_ytilde_htilde_xtilde_wtilde_c_grid_desc =
                     transform_tensor_descriptor(
                         in_n_hip_wip_c_grid_desc,
-                        make_tuple(make_pass_through_transform(N),
-                                   make_embed_transform(make_tuple(YTilde, HTilde),
-                                                        make_tuple(ConvDilationH, ConvStrideH)),
-                                   make_embed_transform(make_tuple(XTilde, WTilde),
-                                                        make_tuple(ConvDilationW, ConvStrideW)),
-                                   make_pass_through_transform(C)),
+                        make_tuple(make_pass_through_transform(N_),
+                                   make_embed_transform(make_tuple(YTilde_, HTilde_),
+                                                        make_tuple(ConvDilationH_, ConvStrideH_)),
+                                   make_embed_transform(make_tuple(XTilde_, WTilde_),
+                                                        make_tuple(ConvDilationW_, ConvStrideW_)),
+                                   make_pass_through_transform(C_)),
                         make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
                         make_tuple(
                             Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
 
                 const auto in_n_htildeslice_wtildeslice_c_grid_desc = transform_tensor_descriptor(
                     in_n_ytilde_htilde_xtilde_wtilde_c_grid_desc,
-                    make_tuple(make_pass_through_transform(N),
-                               make_freeze_transform(i_ytilde),
-                               make_slice_transform(HTilde, IHTildeSliceBegin, HTildeSlice),
-                               make_freeze_transform(i_xtilde),
-                               make_slice_transform(WTilde, IWTildeSliceBegin, WTildeSlice),
-                               make_pass_through_transform(C)),
+                    make_tuple(make_pass_through_transform(N_),
+                               make_freeze_transform(IdxYTilde_),
+                               make_slice_transform(HTilde_, IHTildeSliceBegin, HTildeSlice),
+                               make_freeze_transform(IdxXTilde_),
+                               make_slice_transform(WTilde_, IWTildeSliceBegin, WTildeSlice),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<0>{},
                                Sequence<1>{},
                                Sequence<2>{},
@@ -944,8 +1191,8 @@ struct TransformConvBwdDataToGemm_v1
 
                 const auto in_gemmmraw_gemmnraw_grid_desc = transform_tensor_descriptor(
                     in_n_htildeslice_wtildeslice_c_grid_desc,
-                    make_tuple(make_merge_transform(make_tuple(N, HTildeSlice, WTildeSlice)),
-                               make_pass_through_transform(C)),
+                    make_tuple(make_merge_transform(make_tuple(N_, HTildeSlice, WTildeSlice)),
+                               make_pass_through_transform(C_)),
                     make_tuple(Sequence<0, 1, 2>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}));
 
@@ -961,11 +1208,11 @@ struct TransformConvBwdDataToGemm_v1
             {
                 const auto in_n_dip_hip_wip_c_grid_desc = transform_tensor_descriptor(
                     in_grid_desc,
-                    make_tuple(make_pass_through_transform(N),
-                               make_pad_transform(Di, InLeftPadD, InRightPadD),
-                               make_pad_transform(Hi, InLeftPadH, InRightPadH),
-                               make_pad_transform(Wi, InLeftPadW, InRightPadW),
-                               make_pass_through_transform(C)),
+                    make_tuple(make_pass_through_transform(N_),
+                               make_pad_transform(Di_, InLeftPadD_, InRightPadD_),
+                               make_pad_transform(Hi_, InLeftPadH_, InRightPadH_),
+                               make_pad_transform(Wi_, InLeftPadW_, InRightPadW_),
+                               make_pass_through_transform(C_)),
                     make_tuple(
                         Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}, Sequence<4>{}),
                     make_tuple(
@@ -974,14 +1221,14 @@ struct TransformConvBwdDataToGemm_v1
                 const auto in_n_ztilde_dtilde_ytilde_htilde_xtilde_wtilde_c_grid_desc =
                     transform_tensor_descriptor(
                         in_n_dip_hip_wip_c_grid_desc,
-                        make_tuple(make_pass_through_transform(N),
-                                   make_embed_transform(make_tuple(ZTilde, DTilde),
-                                                        make_tuple(ConvDilationD, ConvStrideD)),
-                                   make_embed_transform(make_tuple(YTilde, HTilde),
-                                                        make_tuple(ConvDilationH, ConvStrideH)),
-                                   make_embed_transform(make_tuple(XTilde, WTilde),
-                                                        make_tuple(ConvDilationW, ConvStrideW)),
-                                   make_pass_through_transform(C)),
+                        make_tuple(make_pass_through_transform(N_),
+                                   make_embed_transform(make_tuple(ZTilde_, DTilde_),
+                                                        make_tuple(ConvDilationD_, ConvStrideD_)),
+                                   make_embed_transform(make_tuple(YTilde_, HTilde_),
+                                                        make_tuple(ConvDilationH_, ConvStrideH_)),
+                                   make_embed_transform(make_tuple(XTilde_, WTilde_),
+                                                        make_tuple(ConvDilationW_, ConvStrideW_)),
+                                   make_pass_through_transform(C_)),
                         make_tuple(Sequence<0>{},
                                    Sequence<1>{},
                                    Sequence<2>{},
@@ -996,14 +1243,14 @@ struct TransformConvBwdDataToGemm_v1
                 const auto in_n_dtildeslice_htildeslice_wtildeslice_c_grid_desc =
                     transform_tensor_descriptor(
                         in_n_ztilde_dtilde_ytilde_htilde_xtilde_wtilde_c_grid_desc,
-                        make_tuple(make_pass_through_transform(N),
-                                   make_freeze_transform(i_ztilde),
-                                   make_slice_transform(DTilde, IDTildeSliceBegin, DTildeSlice),
-                                   make_freeze_transform(i_ytilde),
-                                   make_slice_transform(HTilde, IHTildeSliceBegin, HTildeSlice),
-                                   make_freeze_transform(i_xtilde),
-                                   make_slice_transform(WTilde, IWTildeSliceBegin, WTildeSlice),
-                                   make_pass_through_transform(C)),
+                        make_tuple(make_pass_through_transform(N_),
+                                   make_freeze_transform(IdxZTilde_),
+                                   make_slice_transform(DTilde_, IDTildeSliceBegin, DTildeSlice),
+                                   make_freeze_transform(IdxYTilde_),
+                                   make_slice_transform(HTilde_, IHTildeSliceBegin, HTildeSlice),
+                                   make_freeze_transform(IdxXTilde_),
+                                   make_slice_transform(WTilde_, IWTildeSliceBegin, WTildeSlice),
+                                   make_pass_through_transform(C_)),
                         make_tuple(Sequence<0>{},
                                    Sequence<1>{},
                                    Sequence<2>{},
@@ -1024,8 +1271,8 @@ struct TransformConvBwdDataToGemm_v1
                 const auto in_gemmmraw_gemmnraw_grid_desc = transform_tensor_descriptor(
                     in_n_dtildeslice_htildeslice_wtildeslice_c_grid_desc,
                     make_tuple(
-                        make_merge_transform(make_tuple(N, DTildeSlice, HTildeSlice, WTildeSlice)),
-                        make_pass_through_transform(C)),
+                        make_merge_transform(make_tuple(N_, DTildeSlice, HTildeSlice, WTildeSlice)),
+                        make_pass_through_transform(C_)),
                     make_tuple(Sequence<0, 1, 2, 3>{}, Sequence<4>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}));
 
@@ -1044,84 +1291,41 @@ struct TransformConvBwdDataToGemm_v1
     }
 
     // for input bias
-    template <typename CLayout,
+    template <typename CLayout_                   = CLayout,
               typename std::enable_if<NDimSpatial == 2 &&
-                                          (is_same_v<CLayout, tensor_layout::convolution::GC> ||
-                                           is_same_v<CLayout, tensor_layout::convolution::G_C>),
+                                          (is_same_v<CLayout_, tensor_layout::convolution::GC> ||
+                                           is_same_v<CLayout_, tensor_layout::convolution::G_C>),
                                       bool>::type = false>
-    static auto
-    MakeCDescriptor_M_N(const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_lengths,
-                        const std::array<index_t, NDimSpatial + 3>& /* out_g_n_k_wos_strides */,
-                        const std::array<index_t, NDimSpatial + 3>& wei_g_k_c_xs_lengths,
-                        const std::array<index_t, NDimSpatial + 3>& /* wei_g_k_c_xs_strides */,
-                        const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_lengths,
-                        const std::array<index_t, NDimSpatial + 3>& /* in_g_n_c_wis_strides */,
-                        const std::array<index_t, NDimSpatial>& conv_filter_strides,
-                        const std::array<index_t, NDimSpatial>& conv_filter_dilations,
-                        const std::array<index_t, NDimSpatial>& input_left_pads,
-                        const std::array<index_t, NDimSpatial>& /* input_right_pads */,
-                        const std::array<index_t, NDimSpatial>& /* tildes */)
+    __host__ __device__ auto MakeCDescriptor_M_N() const
     {
-        const index_t N = in_g_n_c_wis_lengths[1];
-        const index_t C = wei_g_k_c_xs_lengths[2];
-
-        const index_t Hi = in_g_n_c_wis_lengths[3];
-        const index_t Wi = in_g_n_c_wis_lengths[4];
-
-        const index_t Ho = out_g_n_k_wos_lengths[3];
-        const index_t Wo = out_g_n_k_wos_lengths[4];
-
-        const index_t Y = wei_g_k_c_xs_lengths[3];
-        const index_t X = wei_g_k_c_xs_lengths[4];
-
-        const index_t InLeftPadH = input_left_pads[0];
-        const index_t InLeftPadW = input_left_pads[1];
-
-        const index_t ConvStrideH = conv_filter_strides[0];
-        const index_t ConvStrideW = conv_filter_strides[1];
-
-        const index_t ConvDilationH = conv_filter_dilations[0];
-        const index_t ConvDilationW = conv_filter_dilations[1];
-
         if constexpr(ConvBwdDataSpecialization ==
                      ck::tensor_operation::device::ConvolutionBackwardDataSpecialization::
                          Filter1x1Stride1Pad0)
         {
             const auto in_gemmm_gemmn_grid_desc =
-                make_naive_tensor_descriptor(make_tuple(N * Ho * Wo, C), make_tuple(I0, I1));
+                make_naive_tensor_descriptor(make_tuple(N_ * Ho_ * Wo_, C_), make_tuple(I0, I1));
 
             return in_gemmm_gemmn_grid_desc;
         }
         else
         {
-            const auto GcdStrideDilationH = math::gcd(ConvStrideH, ConvDilationH);
-            const auto GcdStrideDilationW = math::gcd(ConvStrideW, ConvDilationW);
-
-            const auto YTilde = ConvStrideH / GcdStrideDilationH;
-            const auto XTilde = ConvStrideW / GcdStrideDilationW;
-
-            const auto HTilde =
-                Ho + math::integer_divide_ceil(ConvDilationH * (Y - I1), ConvStrideH);
-            const auto WTilde =
-                Wo + math::integer_divide_ceil(ConvDilationW * (X - I1), ConvStrideW);
-
             // only work on HTilde and WTilde that contribute to non-padding area of input tensor
             const auto IHTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadH - ConvDilationH * (YTilde - I1)), ConvStrideH);
+                math::max(I0, InLeftPadH_ - ConvDilationH_ * (YTilde_ - I1)), ConvStrideH_);
             const auto IWTildeSliceBegin = math::integer_divide_floor(
-                math::max(I0, InLeftPadW - ConvDilationW * (XTilde - I1)), ConvStrideW);
+                math::max(I0, InLeftPadW_ - ConvDilationW_ * (XTilde_ - I1)), ConvStrideW_);
 
             const auto IHTildeSliceEnd = math::min(
-                HTilde, math::integer_divide_ceil(InLeftPadH + Hi - I1, ConvStrideH) + I1);
+                HTilde_, math::integer_divide_ceil(InLeftPadH_ + Hi_ - I1, ConvStrideH_) + I1);
             const auto IWTildeSliceEnd = math::min(
-                WTilde, math::integer_divide_ceil(InLeftPadW + Wi - I1, ConvStrideW) + I1);
+                WTilde_, math::integer_divide_ceil(InLeftPadW_ + Wi_ - I1, ConvStrideW_) + I1);
 
             const auto HTildeSlice = IHTildeSliceEnd - IHTildeSliceBegin;
             const auto WTildeSlice = IWTildeSliceEnd - IWTildeSliceBegin;
 
             // bias tensor
             const auto in_gemmmraw_gemmnraw_grid_desc = make_naive_tensor_descriptor(
-                make_tuple(N * HTildeSlice * WTildeSlice, C), make_tuple(I0, I1));
+                make_tuple(N_ * HTildeSlice * WTildeSlice, C_), make_tuple(I0, I1));
 
             const auto in_gemmm_gemmn_grid_desc = ck::tensor_operation::device::PadTensorDescriptor(
                 in_gemmmraw_gemmnraw_grid_desc,
@@ -1131,6 +1335,25 @@ struct TransformConvBwdDataToGemm_v1
             return in_gemmm_gemmn_grid_desc;
         }
     }
+
+    IndexType N_;
+    IndexType Di_, Hi_, Wi_;
+    IndexType Do_, Ho_, Wo_;
+    IndexType Z_, Y_, X_;
+    IndexType K_, C_;
+    IndexType DiStride_, HiStride_, WiStride_;
+    IndexType DoStride_, HoStride_, WoStride_;
+    IndexType CStrideTensorB_, CStrideTensorC_, KStrideTensorA_, KStrideTensorB_;
+    IndexType NStrideTensorA_, NStrideTensorC_;
+    IndexType ConvStrideD_, ConvStrideH_, ConvStrideW_;
+    IndexType ConvDilationD_, ConvDilationH_, ConvDilationW_;
+    IndexType InLeftPadD_, InLeftPadH_, InLeftPadW_;
+    IndexType InRightPadD_, InRightPadH_, InRightPadW_;
+    IndexType IdxZTilde_, IdxYTilde_, IdxXTilde_;
+    IndexType GcdStrideDilationD_, GcdStrideDilationH_, GcdStrideDilationW_;
+    IndexType ZTilde_, YTilde_, XTilde_;
+    IndexType DTilde_, HTilde_, WTilde_;
+    IndexType ZDot_, YDot_, XDot_;
 };
 
 } // namespace tensor_operation

include/ck/utility/amd_buffer_addressing.hpp

@@ -549,8 +549,10 @@ __device__ void amd_buffer_store_impl(const typename vector_type<T, N>::type src
             (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
             (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
             (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
-            (is_same<T, f8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
-            (is_same<T, bf8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, f8_fnuz_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, bf8_fnuz_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
+            (is_same<T, fp8_storage_t>::value &&
+             (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)) ||
             (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
         "wrong! not implemented");
 
@@ -843,8 +845,8 @@ amd_buffer_load_invalid_element_return_zero(const T* p_src_wave,
 
 #else
 
-    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size, coherence>(
-        src_wave_buffer_resource, src_thread_addr_offset, 0);
+    vector_t tmp{amd_buffer_load_impl<scalar_t, vector_size, coherence>(
+        src_wave_buffer_resource, src_thread_addr_offset, 0)};
     return src_thread_element_valid ? tmp : vector_t(0);
 #endif
 }
@@ -873,8 +875,8 @@ amd_buffer_load_invalid_element_return_customized_value(const T* p_src_wave,
 
     constexpr index_t vector_size = scalar_type<vector_t>::vector_size;
 
-    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size, coherence>(
-        src_wave_buffer_resource, src_thread_addr_offset, 0);
+    vector_t tmp{amd_buffer_load_impl<scalar_t, vector_size, coherence>(
+        src_wave_buffer_resource, src_thread_addr_offset, 0)};
 
     return src_thread_element_valid ? tmp : vector_t(customized_value);
 }

include/ck/utility/amd_ck_fp8.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck/utility/random_gen.hpp"
+#include "ck/utility/type.hpp"
+
+#ifdef CK_USE_FNUZ_FP8
+#define CK_USE_FNUZ_FP8 1
+#else
+#define CK_USE_FNUZ_FP8 0
+#endif
+
+#ifdef CK_USE_OCP_FP8
+#define CK_USE_OCP_FP8 1
+#else
+#define CK_USE_OCP_FP8 0
+#endif
+
+namespace ck {
+
+using f8_fnuz_t  = _BitInt(8);
+using bf8_fnuz_t = unsigned _BitInt(8);
+
+#if(defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) || defined(__gfx1200__) || \
+    defined(__gfx1201__)) &&                                                                     \
+    __HIP_DEVICE_COMPILE__
+#define CK_FP8_CVT_FAST_PATH 1
+#else
+#define CK_FP8_CVT_FAST_PATH 0
+#endif
+
+#if(defined(__gfx1200__) || defined(__gfx1201__)) && __HIP_DEVICE_COMPILE__
+#define CK_OCP_FP8_CVT_FAST_PATH 1
+#else
+#define CK_OCP_FP8_CVT_FAST_PATH 0
+#endif
+
+typedef unsigned char fp8_storage_t;
+
+/**
+ * \brief Describes FP8 interpretation
+ */
+enum class ck_fp8_interpretation_t
+{
+    CK_E4M3_OCP  = 0, // OCP E4M3
+    CK_E5M2_OCP  = 1, // OCP E5M2
+    CK_E4M3_FNUZ = 2, // FP8
+    CK_E5M2_FNUZ = 3, // BF8
+};
+
+/**
+ * \brief Describes saturation behavior
+ */
+enum class ck_saturation_t
+{
+    CK_NOSAT     = 0, // No saturation - replace with NaN or Inf
+    CK_SATFINITE = 1, // Saturate to finite
+};
+
+namespace fp8_impl {
+
+typedef fp8_storage_t fp8x2_storage_t __attribute__((ext_vector_type(2)));
+typedef float float2_t __attribute__((ext_vector_type(2)));
+
+__host__ __device__ static inline constexpr bool fnuz_f8_is_nan(f8_fnuz_t a)
+{
+    return static_cast<unsigned char>(a) == 0x80;
+}
+__host__ __device__ static inline constexpr bool fnuz_bf8_is_nan(bf8_fnuz_t a)
+{
+    return static_cast<unsigned char>(a) == 0x80;
+}
+
+__host__ __device__ static inline constexpr bool ocp_f8_is_nan(fp8_storage_t a)
+{
+    return (a & 0x7f) == 0x7f;
+}
+__host__ __device__ static inline constexpr bool ocp_bf8_is_nan(fp8_storage_t a)
+{
+    return (a & 0x7f) > 0x7c;
+}
+
+// The conversion function is from rocblas
+// https://github.com/ROCm/rocBLAS/blob/9b7f692abe3c54b88d1e77e045a7db7f1f188b69/library/include/internal/rocblas_hip_f8_impl.h#L220
+// This has been modified to handle double types as well
+template <typename T, int wm, int we, bool is_fnuz, bool clip = false>
+__host__ __device__ static inline T cast_from_f8(fp8_storage_t x)
+{
+    constexpr bool is_half   = __hip_internal::is_same<T, _Float16>::value;
+    constexpr bool is_float  = __hip_internal::is_same<T, float>::value;
+    constexpr bool is_double = __hip_internal::is_same<T, double>::value;
+    static_assert(is_half || is_float || is_double, "only half, float and double are supported");
+
+    constexpr int weo = is_half ? 5 : (is_float ? 8 : 11);
+    constexpr int wmo = is_half ? 10 : (is_float ? 23 : 52);
+
+    T fInf, fNegInf, fNaN, fNeg0, fmax, fmin;
+    if constexpr(is_half)
+    {
+        const unsigned short int ihInf    = 0x7C00;
+        const unsigned short int ihNegInf = 0xFC00;
+        const unsigned short int ihNaN    = 0x7C01;
+        const unsigned short int ihNeg0   = 0x8000;
+        /* Max number in e5m2 57344*/
+        const unsigned short int ifmax = 0x7B00;
+        const unsigned short int ifmin = 0xFB00;
+
+        fInf    = bit_cast<_Float16>(ihInf);
+        fNegInf = bit_cast<_Float16>(ihNegInf);
+        fNaN    = bit_cast<_Float16>(ihNaN);
+        fNeg0   = bit_cast<_Float16>(ihNeg0);
+        fmax    = bit_cast<_Float16>(ifmax);
+        fmin    = bit_cast<_Float16>(ifmin);
+    }
+    else if constexpr(is_float)
+    {
+        const unsigned int ifInf    = 0x7F800000;
+        const unsigned int ifNegInf = 0xFF800000;
+        const unsigned int ifNaN    = 0x7F800001;
+        const unsigned int ifNeg0   = 0x80000000;
+        /* Max number in e5m2 57344*/
+        const unsigned int ifmax = 0x47600000;
+        const unsigned int ifmin = 0xC7600000;
+
+        fInf    = bit_cast<float>(ifInf);
+        fNegInf = bit_cast<float>(ifNegInf);
+        fNaN    = bit_cast<float>(ifNaN);
+        fNeg0   = bit_cast<float>(ifNeg0);
+        fmax    = bit_cast<float>(ifmax);
+        fmin    = bit_cast<float>(ifmin);
+    }
+    else if constexpr(is_double)
+    {
+        const unsigned long long ifInf    = 0x7FF0000000000000ull;
+        const unsigned long long ifNegInf = 0xFFF0000000000000ull;
+        const unsigned long long ifNaN    = 0x7FF0000000000001ull;
+        const unsigned long long ifNeg0   = 0x8000000000000000ull;
+        /* Max number in e5m2 57344*/
+        const unsigned long long ifmax = 0x40EC000000000000ull;
+        const unsigned long long ifmin = 0xC0EC000000000000ull;
+
+        fInf    = bit_cast<double>(ifInf);
+        fNegInf = bit_cast<double>(ifNegInf);
+        fNaN    = bit_cast<double>(ifNaN);
+        fNeg0   = bit_cast<double>(ifNeg0);
+        fmax    = bit_cast<double>(ifmax);
+        fmin    = bit_cast<double>(ifmin);
+    }
+
+    if(x == 0)
+    {
+        return 0;
+    }
+
+    unsigned long long sign     = x >> 7;
+    unsigned long long mantissa = x & ((1 << wm) - 1);
+    int exponent                = (x & 0x7F) >> wm;
+    if constexpr(is_fnuz)
+    {
+        if(x == 0x80)
+        {
+            return fNaN;
+        }
+    }
+    else
+    {
+        if(x == 0x80)
+        {
+            return fNeg0;
+        }
+        if constexpr(we == 4)
+        { // e4m3
+            if((x & 0x7F) == 0x7F)
+            {
+                return fNaN;
+            }
+        }
+        else if((x & 0x7C) == 0x7C)
+        { // e5m2
+            if((x & 0x3) == 0)
+            {
+                if constexpr(clip)
+                {
+                    return sign ? fmin : fmax;
+                }
+                return sign ? fNegInf : fInf;
+            }
+            return fNaN;
+        }
+    }
+
+    typename __hip_internal::conditional<
+        sizeof(T) == 2,
+        unsigned short int,
+        typename __hip_internal::conditional<sizeof(T) == 4, unsigned int, unsigned long long>::
+            type>::type retval;
+
+    if constexpr(we == 5 && is_half && !is_fnuz)
+    {
+        retval = x << 8;
+        return bit_cast<T>(retval);
+    }
+
+    const int exp_low_cutoff = (1 << (weo - 1)) - (1 << (we - 1)) + 1 - (is_fnuz ? 1 : 0);
+
+    // subnormal input
+    if(exponent == 0)
+    {
+#if defined(__HIP_DEVICE_COMPILE__) && __HIP_DEVICE_COMPILE__
+        // guaranteed mantissa!=0 since cases 0x0 and 0x80 are handled above
+        int sh = 1 + __clz(mantissa) - (32 - wm);
+#else
+        int sh = 1 + __builtin_clz(mantissa) - (32 - wm);
+#endif
+        mantissa <<= sh;
+        exponent += 1 - sh;
+        mantissa &= ((1ull << wm) - 1);
+    }
+    exponent += exp_low_cutoff - 1;
+    mantissa <<= wmo - wm;
+
+    // subnormal output (occurs when T=half, we=5, negative_zero_nan=true)
+    if(exponent <= 0)
+    {
+        mantissa |= 1 << wmo;
+        mantissa >>= 1 - exponent;
+        exponent = 0;
+    }
+
+    if constexpr(sizeof(T) == 2)
+        retval = (sign << 15) | (exponent << 10) | mantissa;
+    else if constexpr(sizeof(T) == 4)
+        retval = (sign << 31) | (exponent << 23) | mantissa;
+    else
+        retval = (sign << 63) | (static_cast<unsigned long long>(exponent) << 52) | mantissa;
+
+    return bit_cast<T>(retval);
+}
+
+#if CK_FP8_CVT_FAST_PATH
+template <ck_fp8_interpretation_t interpret>
+static __device__ float cast_to_f32_from_f8(fp8_storage_t v)
+{
+    union
+    {
+        unsigned int i32val;
+        unsigned char i8val[4];
+    } val;
+    val.i8val[0] = v;
+
+    static_assert(interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ ||
+                      interpret == ck_fp8_interpretation_t::CK_E4M3_OCP ||
+                      interpret == ck_fp8_interpretation_t::CK_E5M2_FNUZ ||
+                      interpret == ck_fp8_interpretation_t::CK_E5M2_OCP,
+                  "Only FNUZ and OCP interpretations are supported");
+
+    if constexpr((interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ) ||
+                 (interpret == ck_fp8_interpretation_t::CK_E4M3_OCP))
+    {
+        return __builtin_amdgcn_cvt_f32_fp8(val.i32val, 0);
+    }
+    else
+    {
+        return __builtin_amdgcn_cvt_f32_bf8(val.i32val, 0);
+    }
+}
+
+template <ck_fp8_interpretation_t interpret>
+static __device__ float2_t cast_to_f32x2_from_f8x2(fp8x2_storage_t v)
+{
+    const auto i16val = bit_cast<uint16_t>(v);
+
+    static_assert(interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ ||
+                      interpret == ck_fp8_interpretation_t::CK_E4M3_OCP ||
+                      interpret == ck_fp8_interpretation_t::CK_E5M2_FNUZ ||
+                      interpret == ck_fp8_interpretation_t::CK_E5M2_OCP,
+                  "Only FNUZ and OCP interpretations are supported");
+
+    if constexpr((interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ) ||
+                 (interpret == ck_fp8_interpretation_t::CK_E4M3_OCP))
+    {
+        return __builtin_amdgcn_cvt_pk_f32_fp8(i16val, false);
+    }
+    else
+    {
+        return __builtin_amdgcn_cvt_pk_f32_bf8(i16val, false);
+    }
+}
+
+#endif
+
+} // namespace fp8_impl
+
+struct f8_ocp_t
+{
+    using data_type = fp8_storage_t;
+    data_type data;
+
+    static constexpr ck_saturation_t default_saturation = ck_saturation_t::CK_SATFINITE;
+    static constexpr ck_fp8_interpretation_t default_interpret =
+        ck_fp8_interpretation_t::CK_E4M3_OCP;
+
+    static constexpr unsigned int we = 4; // exponent width
+    static constexpr unsigned int wm = 3; // mantissa width
+
+    __host__ __device__ constexpr bool operator==(const f8_ocp_t& other) const
+    {
+        return (data == other.data) && (fp8_impl::ocp_f8_is_nan(data) == false); // NaN != NaN
+    }
+
+#if CK_USE_OCP_FP8
+    __host__ __device__ explicit operator float() const
+#else
+    __host__ explicit operator float() const
+#endif
+    {
+#if CK_OCP_FP8_CVT_FAST_PATH
+        return fp8_impl::cast_to_f32_from_f8<default_interpret>(this->data);
+#else
+        return fp8_impl::cast_from_f8<float, wm, we, false>(
+            this->data); // XXX: clip==false must be consistent with operator _Float16
+#endif
+    }
+
+#if CK_USE_OCP_FP8
+    __host__ __device__ explicit operator _Float16() const
+#else
+    __host__ explicit operator _Float16() const
+#endif
+    {
+#if CK_OCP_FP8_CVT_FAST_PATH
+        return static_cast<_Float16>(fp8_impl::cast_to_f32_from_f8<default_interpret>(this->data));
+#else
+        return fp8_impl::cast_from_f8<_Float16, wm, we, false>(
+            this->data); // XXX: clip==false must be consistent with operator float
+#endif
+    }
+};
+
+struct bf8_ocp_t
+{
+    using data_type = fp8_storage_t;
+    data_type data;
+
+    static constexpr ck_saturation_t default_saturation = ck_saturation_t::CK_SATFINITE;
+    static constexpr ck_fp8_interpretation_t default_interpret =
+        ck_fp8_interpretation_t::CK_E5M2_OCP;
+
+    static constexpr unsigned int we = 5; // exponent width
+    static constexpr unsigned int wm = 2; // mantissa width
+
+    __host__ __device__ constexpr bool operator==(const bf8_ocp_t& other) const
+    {
+        return (data == other.data) && (fp8_impl::ocp_bf8_is_nan(data) == false); // NaN != NaN
+    }
+
+#if CK_USE_OCP_FP8
+    __host__ __device__ explicit operator float() const
+
+#else
+    __host__ explicit operator float() const
+#endif
+    {
+#if defined(__gfx1200__) || defined(__gfx1201__)
+        return fp8_impl::cast_to_f32_from_f8<default_interpret>(this->data);
+#else
+        return fp8_impl::cast_from_f8<float, wm, we, false>(
+            this->data); // XXX: clip==false must be consistent with operator _Float16
+#endif
+    }
+
+#if CK_USE_OCP_FP8
+    __host__ __device__ explicit operator _Float16() const
+#else
+    __host__ explicit operator _Float16() const
+#endif
+    {
+#if defined(__gfx1200__) || defined(__gfx1201__)
+        return static_cast<_Float16>(fp8_impl::cast_to_f32_from_f8<default_interpret>(this->data));
+#else
+        return fp8_impl::cast_from_f8<_Float16, wm, we, false>(
+            this->data); // XXX: clip==false must be consistent with operator float
+#endif
+    }
+};
+
+template <typename T>
+__host__ __device__ static inline constexpr bool fp8_is_nan(T);
+
+template <>
+__host__ __device__ inline constexpr bool fp8_is_nan(f8_ocp_t a)
+{
+    return fp8_impl::ocp_f8_is_nan(a.data);
+}
+template <>
+__host__ __device__ inline constexpr bool fp8_is_nan(bf8_ocp_t a)
+{
+    return fp8_impl::ocp_bf8_is_nan(a.data);
+}
+template <>
+__host__ __device__ inline constexpr bool fp8_is_nan(f8_fnuz_t a)
+{
+    return fp8_impl::fnuz_f8_is_nan(a);
+}
+template <>
+__host__ __device__ inline constexpr bool fp8_is_nan(bf8_fnuz_t a)
+{
+    return fp8_impl::fnuz_bf8_is_nan(a);
+}
+
+template <typename T,
+          std::enable_if_t<std::is_same_v<T, bf8_ocp_t> || std::is_same_v<T, f8_ocp_t> ||
+                               std::is_same_v<T, bf8_fnuz_t> || std::is_same_v<T, f8_fnuz_t>,
+                           bool> = true>
+__host__ __device__ static inline constexpr bool fp8_is_inf(T)
+{
+    return false;
+}
+template <>
+__host__ __device__ inline constexpr bool fp8_is_inf(bf8_ocp_t a)
+{
+    return (a.data & 0x7f) == 0x7c;
+}
+
+namespace fp8_impl {
+
+// Assertions to check for supported conversion types
+#define __assert_ocp_support(interp)                                               \
+    {                                                                              \
+        if(interp != ck_fp8_interpretation_t::CK_E4M3_OCP &&                       \
+           interp != ck_fp8_interpretation_t::CK_E5M2_OCP)                         \
+        {                                                                          \
+            __hip_assert(false && "type is unsupported by current target device"); \
+        }                                                                          \
+    }
+#define __assert_fnuz_support(interp)                                              \
+    {                                                                              \
+        if(interp != ck_fp8_interpretation_t::CK_E4M3_FNUZ &&                      \
+           interp != ck_fp8_interpretation_t::CK_E5M2_FNUZ)                        \
+        {                                                                          \
+            __hip_assert(false && "type is unsupported by current target device"); \
+        }                                                                          \
+    }
+
+__host__ __device__ static inline void
+__is_interpret_supported([[maybe_unused]] ck_fp8_interpretation_t interp)
+{
+#if defined(__HIP_DEVICE_COMPILE__) && __HIP_DEVICE_COMPILE__
+#if CK_USE_OCP_FP8
+    __assert_ocp_support(interp);
+#endif
+#if CK_USE_FNUZ_FP8
+    __assert_fnuz_support(interp);
+#endif
+#endif
+}
+
+#if CK_FP8_CVT_FAST_PATH
+// The conversion function is from rocblas
+// https://github.com/ROCm/rocBLAS/blob/9b7f692abe3c54b88d1e77e045a7db7f1f188b69/library/include/internal/rocblas_float8.h#L79
+template <ck_fp8_interpretation_t interpret, bool saturate, bool stochastic_rounding = false>
+static __device__ fp8_storage_t cast_to_f8_from_f32(float v, unsigned int rng = 0)
+{
+    fp8_storage_t i8data;
+    union
+    {
+        float fval;
+        unsigned int i32val;
+        unsigned char i8val[4]; // NOTE: not endian independent
+    } val;
+
+    unsigned int ival = 0;
+    val.fval          = v;
+
+    if constexpr(saturate)
+    {
+        if constexpr(interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ)
+        {
+            if((val.i32val & 0x7F800000) != 0x7F800000)
+            { /// propagate NAN/INF, no clipping
+                val.fval = __builtin_amdgcn_fmed3f(val.fval, 240.0, -240.0);
+            }
+        }
+        else if constexpr(interpret == ck_fp8_interpretation_t::CK_E4M3_OCP)
+        { // OCP type
+            if((val.i32val & 0x7F800000) != 0x7F800000)
+            { /// propagate NAN/INF, no clipping
+                val.fval = __builtin_amdgcn_fmed3f(val.fval, 448.0, -448.0);
+            }
+        }
+        else
+        {
+            if((val.i32val & 0x7F800000) != 0x7F800000)
+            { /// propagate NAN/INF, no clipping
+                val.fval = __builtin_amdgcn_fmed3f(val.fval, 57344.0, -57344.0);
+            }
+        }
+    }
+
+    if constexpr(stochastic_rounding)
+    {
+        ival       = (interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ) ||
+                       (interpret == ck_fp8_interpretation_t::CK_E4M3_OCP)
+                         ? __builtin_amdgcn_cvt_sr_fp8_f32(val.fval, rng, ival, 0)
+                         : __builtin_amdgcn_cvt_sr_bf8_f32(val.fval, rng, ival, 0); // 0 pos
+        val.i32val = ival;
+        i8data     = val.i8val[0]; // little endian
+    }
+    else
+    { // RNE CVT
+        ival       = (interpret == ck_fp8_interpretation_t::CK_E4M3_FNUZ) ||
+                       (interpret == ck_fp8_interpretation_t::CK_E4M3_OCP)
+                         ? __builtin_amdgcn_cvt_pk_fp8_f32(val.fval, val.fval, ival, false)
+                         : __builtin_amdgcn_cvt_pk_bf8_f32(val.fval,
+                                                     val.fval,
+                                                     ival,
+                                                     false); // false -> WORD0
+        val.i32val = ival;
+        i8data     = val.i8val[0];
+    }
+    return i8data;
+}
+#endif // CK_FP8_CVT_FAST_PATH
+
+// The conversion function is from rocblas
+// https://github.com/ROCm/rocBLAS/blob/9b7f692abe3c54b88d1e77e045a7db7f1f188b69/library/include/internal/rocblas_hip_f8_impl.h#L39
+// This has been modified to add double types conversion as well
+template <typename T, int wm, int we, bool is_fnuz, bool clip = false, bool stoch = false>
+__host__ __device__ static inline fp8_storage_t cast_to_f8(T _x, unsigned int rng = 0)
+{
+    constexpr bool is_half   = __hip_internal::is_same<T, _Float16>::value;
+    constexpr bool is_float  = __hip_internal::is_same<T, float>::value;
+    constexpr bool is_double = __hip_internal::is_same<T, double>::value;
+    static_assert(is_half || is_float || is_double,
+                  "Only half, float and double can be cast to f8");
+
+    constexpr int mfmt = (sizeof(T) == 8) ? 52 : ((sizeof(T) == 4) ? 23 : 10);
+
+    using T_bitwise = typename __hip_internal::conditional<
+        sizeof(T) == 2,
+        unsigned short int,
+        typename __hip_internal::conditional<sizeof(T) == 4, unsigned int, unsigned long long>::
+            type>::type;
+    T_bitwise x_bitwise = bit_cast<T_bitwise>(_x);
+
+    unsigned long long x{x_bitwise};
+
+    unsigned long long head, mantissa;
+    int exponent, bias;
+    unsigned int sign;
+    unsigned long long fInf, mask;
+
+    if constexpr(sizeof(T) == 8)
+    {
+        head     = x & 0xFFF0000000000000ull;
+        mantissa = x & 0xFFFFFFFFFFFFFull;
+        exponent = (head >> 52) & 0x7FF;
+        sign     = head >> 63;
+        bias     = 1023;
+        fInf     = 0x7FF0000000000000ull;
+        mask     = 0x7FFFFFFFFFFFFFFFull;
+    }
+    else if constexpr(sizeof(T) == 4)
+    {
+        head     = x & 0xFF800000;
+        mantissa = x & 0x7FFFFF;
+        exponent = (head >> 23) & 0xFF;
+        sign     = head >> 31;
+        bias     = 127;
+        fInf     = 0x7F800000;
+        mask     = 0x7FFFFFFF;
+    }
+    else
+    {
+        head     = x & 0xFC00;
+        mantissa = x & 0x3FF;
+        exponent = (head >> 10) & 0x1F;
+        sign     = head >> 15;
+        bias     = 15;
+        fInf     = 0x7C00;
+        mask     = 0x7FFF;
+    }
+    unsigned int signed_inf = 0;
+    unsigned int nan        = 0;
+    if constexpr(is_fnuz)
+    {
+        signed_inf = clip ? ((sign << 7) + 0x7f) : 0x80;
+        nan        = 0x80;
+    }
+    else
+    {
+        if constexpr(we == 4)
+        { // e4m3
+            signed_inf = (sign << 7) + (clip ? 0x7e : 0x7f);
+        }
+        else
+        { // e5m2
+            signed_inf = (sign << 7) + (clip ? 0x7b : 0x7c);
+        }
+        nan = (sign << 7) + 0x7f;
+    }
+    // Max values
+    unsigned long long ifmax = 0;
+    if constexpr(sizeof(T) == 8)
+    {
+        if constexpr(we == 5)
+        { // 57344
+            ifmax = 0x40EC000000000000ull;
+        }
+        else
+        {
+            if constexpr(is_fnuz)
+            { // 240
+                ifmax = 0x406E000000000000ull;
+            }
+            else
+            { // 448
+                ifmax = 0x407C000000000000ull;
+            }
+        }
+    }
+    else if(sizeof(T) == 4)
+    {
+        if constexpr(we == 5)
+        {
+            ifmax = 0x47600000;
+        }
+        else
+        {
+            if constexpr(is_fnuz)
+            {
+                ifmax = 0x43700000;
+            }
+            else
+            {
+                ifmax = 0x43E00000;
+            }
+        }
+    }
+    else
+    {
+        if constexpr(we == 5)
+        {
+            ifmax = 0x7B00;
+        }
+        else
+        {
+            if constexpr(is_fnuz)
+            {
+                ifmax = 0x5B80;
+            }
+            else
+            {
+                ifmax = 0x5F00;
+            }
+        }
+    }
+    // Deal with inf and NaNs
+    if((x & fInf) == fInf)
+    {
+        if constexpr(is_fnuz)
+            return signed_inf;
+
+        return mantissa != 0 ? nan : signed_inf;
+    }
+
+    if((x & mask) > ifmax)
+    {
+        return signed_inf;
+    }
+
+    if(x == 0)
+    {
+        return 0;
+    }
+
+    // First need to check if it is normal or denorm as there is a difference of
+    // implicit 1 Then need to adjust the exponent to align with the F8 exponent,
+    // in the meanwhile, shift The mantissa. Then for stochastic rounding, add rng
+    // to mantissa and truncate. And for RNE, no need to add rng. Then probably
+    // need to check whether there is carry and adjust exponent and mantissa again
+
+    // For IEEE bias mode, the bias is 2^(k-1) -1 where k is the width of exponent
+    // bits
+    const int f8_bias                  = (1 << (we - 1)) - 1 + (is_fnuz ? 1 : 0);
+    const int f8_denormal_act_exponent = 1 - f8_bias; // actual exponent of f8 denormal
+    // act_exponent is the actual exponent of fp32/fp16 (after subtracting bias)
+    // f8_exponent is the converted f8 exponent with bias encoding
+    // exponent_diff is the diff between fp32/fp16 exponent and f8 exponent,
+    // the difference needs to be adjusted and mantissa shifted
+    int act_exponent, f8_exponent, exponent_diff;
+
+    if(exponent == 0)
+    { // fp32/fp16 is in denormal.
+        /* fp32 denormal is below 2^-127 so it is usually not a concern here, we
+    mostly concern fp16 here. In this case, f8 is usually in denormal. But there
+    could be exceptions. fp16 denormal has exponent bias 15 while bf8 with NANOO has
+    exponent bias 16. It means that there are some numbers in fp16 denormal but they
+    are bf8 (NANOO) normals - smallest bf8 (NANOO) normal is 2^-15. fp16 numbers
+    where exponent==0 (actual exponent -14) and highest bit of mantissa is 1 are bf8
+    (NANOO) normal. In this case, the fp16 mantissa should be shift left by 1  */
+        act_exponent  = exponent - bias + 1;
+        exponent_diff = f8_denormal_act_exponent -
+                        act_exponent; // actual exponent is exponent-bias+1 as it is denormal
+    }
+    else
+    { // fp32/fp16 is normal with implicit 1
+        act_exponent = exponent - bias;
+        if(act_exponent <= f8_denormal_act_exponent)
+        {
+            /* This is the case where fp32/fp16 is normal but it is in f8 denormal
+      range. For example fp8 nanoo mode, denormal exponent is -7, but if the fp32/fp16
+      actual exponent is -7, it is actually larger due to the implicit 1,
+      Therefore it needs to be adjust to -6 and mantissa shift right by 1.
+      So for fp32/fp16, exponent -8 is the cut point to convert to fp8 nanoo */
+            exponent_diff = f8_denormal_act_exponent - act_exponent;
+        }
+        else
+        {                      // both fp32/fp16 and f8 are in normal range
+            exponent_diff = 0; // exponent_diff=0 does not mean there is no difference
+                               // for this case, act_exponent could be larger. Just
+                               // that it does not need shift mantissa
+        }
+        mantissa += (1ull << mfmt); // Add the implicit 1 into mantissa
+    }
+
+    bool midpoint = (mantissa & ((1ull << (mfmt - wm + exponent_diff)) - 1)) ==
+                    (1ull << (mfmt - wm + exponent_diff - 1));
+    /* This part is a bit tricky. The judgment of whether it is a tie needs to be
+  done before we shift right as shift right could rip off some residual part and
+  make something not midpoint look like midpoint. For example, the fp16 number
+  0x1002 (0 00100 0000000010), it is larger than midpoint, but after shift right
+  by 4 bits, it would look like midpoint.
+  */
+
+    if(exponent_diff > 0)
+        mantissa >>= exponent_diff;
+    else if(exponent_diff == -1)
+        mantissa <<= -exponent_diff;
+    bool implicit_one = mantissa & (1ull << mfmt);
+    // if there is no implicit 1, it  means the f8 is denormal and need to adjust
+    // to denorm exponent
+    f8_exponent =
+        (act_exponent + exponent_diff) /*actual f8 exponent*/ + f8_bias - (implicit_one ? 0 : 1);
+
+    // Now we have the exponent and mantissa adjusted
+    unsigned long long drop_mask = (1ull << (mfmt - wm)) - 1;
+    bool odd =
+        mantissa & (1ull << (mfmt - wm)); // if the least significant bit that is not truncated is 1
+    mantissa +=
+        (stoch ? rng : (midpoint ? (odd ? mantissa : mantissa - 1ull) : mantissa)) & drop_mask;
+
+    // Now we deal with overflow
+    if(f8_exponent == 0)
+    {
+        if((1ull << mfmt) & mantissa)
+        {
+            f8_exponent = 1; // denormal overflow to become normal, promote exponent
+        }
+    }
+    else
+    {
+        if((1ull << (mfmt + 1)) & mantissa)
+        {
+            mantissa >>= 1;
+            f8_exponent++;
+        }
+    }
+
+    mantissa >>= (mfmt - wm);
+
+    // above range: quantize to maximum possible float of the same sign
+    const int max_exp = (1 << we) - 1;
+    if(f8_exponent > max_exp)
+    {
+        if constexpr(clip)
+        {
+            mantissa    = (1 << wm) - 1;
+            f8_exponent = max_exp;
+        }
+        else
+        {
+            return signed_inf;
+        }
+    }
+
+    if(f8_exponent == 0 && mantissa == 0)
+        return is_fnuz ? 0 : (sign << 7);
+    mantissa &= (1 << wm) - 1;
+    return (sign << 7) | (f8_exponent << wm) | mantissa;
+}
+
+/**
+ * \brief convert float to @p fp8_storage_t
+ *
+ * \tparam interp interpretation of fp8
+ * \tparam sat saturation of fp8
+ * \param f float number
+ * \return fp8_storage_t
+ */
+template <ck_fp8_interpretation_t interp,
+          ck_saturation_t sat      = ck_saturation_t::CK_SATFINITE,
+          bool stochastic_rounding = false>
+#if CK_FP8_CVT_FAST_PATH
+__host__ __device__ static inline fp8_storage_t cvt_float_to_fp8(const float f)
+{
+    __is_interpret_supported(interp);
+    uint32_t rng = 0;
+    if constexpr(stochastic_rounding)
+    {
+        constexpr int seed = 1254739;
+        rng                = prand_generator<float, seed>(reinterpret_cast<uintptr_t>(&f), f);
+    }
+    return cast_to_f8_from_f32<interp, sat == ck_saturation_t::CK_SATFINITE, stochastic_rounding>(
+        f, rng);
+#else
+#if CK_USE_OCP_FP8
+__host__ __device__ static inline fp8_storage_t cvt_float_to_fp8(const float f)
+{
+#else
+__host__ static inline fp8_storage_t cvt_float_to_fp8(const float f)
+{
+#endif
+    uint32_t rng = 0;
+    if constexpr(stochastic_rounding)
+    {
+        constexpr int seed = 1254739;
+        rng = prand_generator<float, seed>(reinterpret_cast<uintptr_t>(&f), f);
+    }
+
+    if constexpr(interp == ck_fp8_interpretation_t::CK_E4M3_FNUZ)
+    {
+        return cast_to_f8<float,
+                          3,
+                          4,
+                          true,
+                          sat == ck_saturation_t::CK_SATFINITE,
+                          stochastic_rounding>(f, rng);
+    }
+    else if constexpr(interp == ck_fp8_interpretation_t::CK_E5M2_FNUZ)
+    {
+        return cast_to_f8<float,
+                          2,
+                          5,
+                          true,
+                          sat == ck_saturation_t::CK_SATFINITE,
+                          stochastic_rounding>(f, rng);
+    }
+    else if constexpr(interp == ck_fp8_interpretation_t::CK_E4M3_OCP)
+    {
+        return cast_to_f8<float,
+                          3,
+                          4,
+                          false,
+                          sat == ck_saturation_t::CK_SATFINITE,
+                          stochastic_rounding>(f, rng);
+    }
+    else if constexpr(interp == ck_fp8_interpretation_t::CK_E5M2_OCP)
+    {
+        return cast_to_f8<float,
+                          2,
+                          5,
+                          false,
+                          sat == ck_saturation_t::CK_SATFINITE,
+                          stochastic_rounding>(f, rng);
+    }
+    else
+    {
+        __hip_assert(false && "FP8 type is not supported by current target device");
+        return 0;
+    }
+#endif // CK_FP8_CVT_FAST_PATH
+}
+
+/**
+ * \brief convert _Float16 to @p fp8_storage_t
+ *
+ * \tparam sat saturation of fp8
+ * \tparam interp interpretation of fp8
+ * \tparam stochastic_rounding switch between RNE and SR
+ * \param x _Float16 value
+ * \return fp8_storage_t
+ */
+template <ck_fp8_interpretation_t interp,
+          ck_saturation_t sat      = ck_saturation_t::CK_SATFINITE,
+          bool stochastic_rounding = false>
+#if CK_FP8_CVT_FAST_PATH || CK_USE_OCP_FP8
+__host__ __device__ static inline fp8_storage_t cvt_half_t_to_fp8(const _Float16 x)
+#else
+__host__ static inline fp8_storage_t cvt_half_t_to_fp8(const _Float16 x)
+#endif
+{
+    return cvt_float_to_fp8<interp, sat, stochastic_rounding>(static_cast<float>(x));
+}
+
+} // namespace fp8_impl
+
+// Declare a template function for fp8 conversion using RNE
+template <typename Y, typename X>
+__host__ __device__ constexpr Y f8_convert_rne(X x);
+
+// convert fp32 to fp8 with rounding to nearest even
+template <>
+inline __host__ __device__ f8_ocp_t f8_convert_rne<f8_ocp_t, float>(float x)
+{
+    return f8_ocp_t{
+        fp8_impl::cvt_float_to_fp8<f8_ocp_t::default_interpret, f8_ocp_t::default_saturation>(x)};
+}
+
+// convert fp32 to bf8 with rounding to nearest even
+template <>
+inline __host__ __device__ bf8_ocp_t f8_convert_rne<bf8_ocp_t, float>(float x)
+{
+    return bf8_ocp_t{
+        fp8_impl::cvt_float_to_fp8<bf8_ocp_t::default_interpret, bf8_ocp_t::default_saturation>(x)};
+}
+
+// convert _Float16 to fp8 with rounding to nearest even
+template <>
+inline __host__ __device__ f8_ocp_t f8_convert_rne<f8_ocp_t, _Float16>(_Float16 x)
+{
+    return f8_ocp_t{
+        fp8_impl::cvt_half_t_to_fp8<f8_ocp_t::default_interpret, f8_ocp_t::default_saturation>(x)};
+}
+
+template <>
+inline __host__ __device__ bf8_ocp_t f8_convert_rne<bf8_ocp_t, _Float16>(_Float16 x)
+{
+    return bf8_ocp_t{
+        fp8_impl::cvt_half_t_to_fp8<bf8_ocp_t::default_interpret, bf8_ocp_t::default_saturation>(
+            x)};
+}
+
+// Declare a template function for fp8 conversion using RNE
+template <typename Y, typename X>
+__host__ __device__ constexpr Y f8_convert_sr(X x);
+
+// convert fp32 to fp8 with stochastic rounding
+template <>
+inline __host__ __device__ f8_ocp_t f8_convert_sr<f8_ocp_t, float>(float x)
+{
+    return f8_ocp_t{
+        fp8_impl::cvt_float_to_fp8<f8_ocp_t::default_interpret, f8_ocp_t::default_saturation, true>(
+            x)};
+}
+
+// convert fp32 to bf8 with stochastic rounding
+template <>
+inline __host__ __device__ bf8_ocp_t f8_convert_sr<bf8_ocp_t, float>(float x)
+{
+    return bf8_ocp_t{fp8_impl::cvt_float_to_fp8<bf8_ocp_t::default_interpret,
+                                                bf8_ocp_t::default_saturation,
+                                                true>(x)};
+}
+
+// convert _Float16 to fp8 with stochastic rounding
+template <>
+inline __host__ __device__ f8_ocp_t f8_convert_sr<f8_ocp_t, _Float16>(_Float16 x)
+{
+    return f8_ocp_t{fp8_impl::cvt_half_t_to_fp8<f8_ocp_t::default_interpret,
+                                                f8_ocp_t::default_saturation,
+                                                true>(x)};
+}
+
+// convert _Float16 to bf8 with stochastic rounding
+template <>
+inline __host__ __device__ bf8_ocp_t f8_convert_sr<bf8_ocp_t, _Float16>(_Float16 x)
+{
+    return bf8_ocp_t{fp8_impl::cvt_half_t_to_fp8<bf8_ocp_t::default_interpret,
+                                                 bf8_ocp_t::default_saturation,
+                                                 true>(x)};
+}
+
+#if CK_USE_OCP_FP8
+using f8_t  = f8_ocp_t;
+using bf8_t = bf8_ocp_t;
+#define CK_FP8_TYPE_FNUZ 0
+#define CK_FP8_TYPE_OCP 1
+#else
+using f8_t = f8_fnuz_t;
+using bf8_t = bf8_fnuz_t;
+#define CK_FP8_TYPE_FNUZ 1
+#define CK_FP8_TYPE_OCP 0
+#endif
+
+} // namespace ck