tensor_descriptor_utils.hpp Source File

tensor_descriptor_utils.hpp Source File#

Composable Kernel: tensor_descriptor_utils.hpp Source File
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
 
#pragma once
 
#include "ck_tile/core.hpp"
 
namespace ck_tile {
template <ck_tile::index_t NumDimG,
          ck_tile::index_t NumDimM,
          ck_tile::index_t NumDimN,
          ck_tile::index_t NumDimK>
struct TensorDescriptorUtils
{
    CK_TILE_HOST static constexpr auto
    Make_A_GridDescriptor_M_K(const std::vector<ck_tile::index_t>& A_dims    = {},
                              const std::vector<ck_tile::index_t>& A_strides = {})
    {
        const auto to_tuple = [&](auto& vec, auto start, auto end) {
            return generate_tuple([&](auto i) { return vec[start + i]; }, number<end - start>{});
        };
 
        // Remove G Dimensions
        const auto A_dims_M_K =
            to_tuple(A_dims, number<NumDimG>{}, number<NumDimG + NumDimM + NumDimK>{});
        const auto A_strides_M_K =
            to_tuple(A_strides, number<NumDimG>{}, number<NumDimG + NumDimM + NumDimK>{});
 
        // dimension Ids for M and K
        constexpr auto A_dims_M_ids = typename arithmetic_sequence_gen<0, NumDimM, 1>::type{};
        constexpr auto A_dims_K_ids =
            typename arithmetic_sequence_gen<NumDimM, NumDimM + NumDimK, 1>::type{};
 
        // Dimensions for M [M0, M1, ...] and K [K0, K1, ...]
        const auto dims_M = get_container_subset(A_dims_M_K, A_dims_M_ids);
        const auto dims_K = get_container_subset(A_dims_M_K, A_dims_K_ids);
 
        // naive tensor A[M0, M1, M2, ..., K0, K1, K2...] Discriptor
        const auto A_grid_desc_Ms_Ks =
            ck_tile::make_naive_tensor_descriptor(A_dims_M_K, A_strides_M_K);
 
        // transformed tensor to flatten M and K dimensions  [M_total = M0 * M1 * M2 * ... , K_total
        // = K0 * K1 * K2 * ...]
        const auto A_grid_desc_Mflat_Kflat = ck_tile::transform_tensor_descriptor(
            A_grid_desc_Ms_Ks,
            make_tuple(make_merge_transform(dims_M), make_merge_transform(dims_K)),
            make_tuple(A_dims_M_ids, A_dims_K_ids),
            make_tuple(sequence<0>{}, sequence<1>{}));
 
        return A_grid_desc_Mflat_Kflat;
    }
    CK_TILE_HOST static constexpr auto
    Make_B_GridDescriptor_N_K(const std::vector<ck_tile::index_t>& B_dims    = {},
                              const std::vector<ck_tile::index_t>& B_strides = {})
    {
        const auto to_tuple = [&](auto& vec, auto start, auto end) {
            return generate_tuple([&](auto i) { return vec[start + i]; }, number<end - start>{});
        };
 
        // Remove G Dimensions
        const auto B_dims_N_K =
            to_tuple(B_dims, number<NumDimG>{}, number<NumDimG + NumDimN + NumDimK>{});
        const auto B_strides_N_K =
            to_tuple(B_strides, number<NumDimG>{}, number<NumDimG + NumDimN + NumDimK>{});
 
        // dimension Ids for N and K
        constexpr auto B_dims_N_ids = typename arithmetic_sequence_gen<0, NumDimN, 1>::type{};
        constexpr auto B_dims_K_ids =
            typename arithmetic_sequence_gen<NumDimN, NumDimN + NumDimK, 1>::type{};
 
        // Dimensions for N [N0, N1, ...] and K [K0, K1, ...]
        const auto dims_N = get_container_subset(B_dims_N_K, B_dims_N_ids);
        const auto dims_K = get_container_subset(B_dims_N_K, B_dims_K_ids);
 
        // naive tensor B[N0, N1, N2, ..., K0, K1, K2...] Discriptor
        const auto B_grid_desc_Ns_Ks =
            ck_tile::make_naive_tensor_descriptor(B_dims_N_K, B_strides_N_K);
 
        // transformed tensor to flatten N and K dimensions  [N_total = N0 * N1 * N2 * ... , K_total
        // = K0 * K1 * K2 * ...]
        const auto B_grid_desc_Nflat_Kflat = ck_tile::transform_tensor_descriptor(
            B_grid_desc_Ns_Ks,
            make_tuple(make_merge_transform(dims_N), make_merge_transform(dims_K)),
            make_tuple(B_dims_N_ids, B_dims_K_ids),
            make_tuple(sequence<0>{}, sequence<1>{}));
 
        return B_grid_desc_Nflat_Kflat;
    }
    CK_TILE_HOST static constexpr auto
    Make_E_GridDescriptor_M_N(const std::vector<ck_tile::index_t>& E_dims    = {},
                              const std::vector<ck_tile::index_t>& E_strides = {})
    {
        const auto to_tuple = [&](auto& vec, auto start, auto end) {
            return generate_tuple([&](auto i) { return vec[start + i]; }, number<end - start>{});
        };
 
        // Remove G dimensions
        const auto E_dims_M_N =
            to_tuple(E_dims, number<NumDimG>{}, number<NumDimG + NumDimM + NumDimN>{});
        const auto E_strides_M_N =
            to_tuple(E_strides, number<NumDimG>{}, number<NumDimG + NumDimM + NumDimN>{});
 
        // dimension Ids for M and N
        constexpr auto E_dims_M_ids = typename arithmetic_sequence_gen<0, NumDimM, 1>::type{};
        constexpr auto E_dims_N_ids =
            typename arithmetic_sequence_gen<NumDimM, NumDimM + NumDimN, 1>::type{};
 
        // Dimensions for M and N
        const auto dims_M = get_container_subset(E_dims_M_N, E_dims_M_ids);
        const auto dims_N = get_container_subset(E_dims_M_N, E_dims_N_ids);
 
        // naive tensor E[M0, M1, M2, ..., N0, N1, N2...] Discriptor
        const auto E_grid_desc_Ms_Ns =
            ck_tile::make_naive_tensor_descriptor(E_dims_M_N, E_strides_M_N);
 
        // transformed tensor to flatten M and N dimensions   [M_total = M0 * M1 * M2 * ... ,
        // N_total = N0 * N1 * N2 * ...]
        const auto E_grid_desc_Mflat_Nflat = ck_tile::transform_tensor_descriptor(
            E_grid_desc_Ms_Ns,
            make_tuple(make_merge_transform(dims_M), make_merge_transform(dims_N)),
            make_tuple(E_dims_M_ids, E_dims_N_ids),
            make_tuple(sequence<0>{}, sequence<1>{}));
 
        return E_grid_desc_Mflat_Nflat;
    }
};
 
} // namespace ck_tile