/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/permute/kernel/generic_permute_kernel.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/permute/kernel/generic_permute_kernel.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/permute/kernel/generic_permute_kernel.hpp Source File

Go to the documentation of this file.

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/common.hpp"
 // #include "ck_tile/ops/permute/pipeline/generic_petmute_problem.hpp"
  
 namespace ck_tile {
  
 /* independent host side argument, no template
  */
 struct GenericPermuteHostArgs
 {
     static constexpr index_t kMaxRanks = 8; // TODO: hardcoded
  
     const void* p_src;
     void* p_dst;
     index_t rank;
     index_t shape[kMaxRanks]; // input shape
     index_t perm[kMaxRanks];  // permute index
 };
  
 /*
 simulate torch.permute:
 x_ = x_.view(x.shape[0],
                     x.shape[1]//16, 16,
                     x.shape[2]//32, 4, 8)
 x_ = x_.permute(0,1,3,4,2,5)
 x_ = x_.contiguous()
 x_ = x_.view(x.shape[0], x.shape[1], x.shape[2]);//
  
 this kernel is supposed not to be performant(just OK), with functional support up to kMaxRanks
 dim of permutation, with a single kernel
  
 */
 template <typename Problem_>
 struct GenericPermute
 {
     using Problem = ck_tile::remove_cvref_t<Problem_>;
  
     using DataType                      = remove_cvref_t<typename Problem::DataType>;
     static constexpr index_t kBlockSize = Problem::kBlockSize;
     static constexpr index_t kMaxRanks  = Problem::kMaxRanks;
     static constexpr bool KeepLastDim   = Problem::KeepLastDim;
  
     struct __attribute__((packed)) Kargs
     {
         const void* p_src;
         void* p_dst;
         // index_t rank;
         index_t num_elements;
         index_t perm_length[kMaxRanks]; // tensor length after permutation
         index_t perm_stride[kMaxRanks]; // tensor stride after permutation
     };
  
     CK_TILE_HOST static constexpr index_t TotalElements(const GenericPermuteHostArgs& h)
     {
         index_t n = 1;
         for(auto i = 0; i < h.rank; i++)
         {
             n *= h.shape[i];
         }
         return n;
     }
  
     CK_TILE_HOST static constexpr Kargs MakeKargs(const GenericPermuteHostArgs& h)
     {
         Kargs a;
         a.p_src = h.p_src;
         a.p_dst = h.p_dst;
  
         // assert rank <= kMaxRanks
         index_t i = 0;
  
         index_t perm[kMaxRanks];
         index_t x_shape[kMaxRanks];
         index_t x_stride[kMaxRanks];
         // index_t perm_length[kMaxRanks];
  
         for(; i < h.rank; i++)
         {
             x_shape[i] = h.shape[i];
             perm[i]    = h.perm[i];
         }
         for(; i < kMaxRanks; i++)
         {
             x_shape[i] = 1;
             perm[i]    = i; // will index to len = 1
         }
  
         index_t stride = 1;
         for(index_t j = kMaxRanks - 1; j >= 0; j--)
         {
             x_stride[j] = stride;
             stride *= x_shape[j];
         }
  
         for(index_t j = 0; j < kMaxRanks; j++)
         {
             a.perm_length[j] = x_shape[perm[j]];
             a.perm_stride[j] = x_stride[perm[j]];
         }
  
         a.num_elements = TotalElements(h);
         return a;
     }
  
     CK_TILE_HOST static constexpr auto GridSize(GenericPermuteHostArgs h)
     {
         auto total = TotalElements(h);
         auto grids = dim3((total + BlockSize() - 1) / BlockSize());
         //  printf("### total:%d, grids:%dx%dx%d\n", total, );
         return grids;
     }
  
     CK_TILE_HOST_DEVICE static constexpr auto BlockSize() { return Problem::kBlockSize; }
  
     CK_TILE_DEVICE void operator()(Kargs kargs) const
     {
         index_t id = blockIdx.x * BlockSize() + threadIdx.x;
  
         if(id >= kargs.num_elements)
             return;
  
         const auto perm_length =
             generate_tuple([&](auto I) { return kargs.perm_length[I]; }, number<kMaxRanks>{});
         const auto perm_stride =
             generate_tuple([&](auto I) { return kargs.perm_stride[I]; }, number<kMaxRanks>{});
  
         const DataType* p_src = reinterpret_cast<const DataType*>(kargs.p_src);
         DataType* p_dst       = reinterpret_cast<DataType*>(kargs.p_dst);
  
         const auto src_view_0 = make_naive_tensor_view<address_space_enum::global>(
             p_src, perm_length, perm_stride, number<1>{}, number<1>{});
  
         const auto src_view = transform_tensor_view(
             src_view_0,
             make_tuple(make_merge_transform(perm_length)),
             make_tuple(typename arithmetic_sequence_gen<0, kMaxRanks, 1>::type{}),
             make_tuple(sequence<0>{}));
  
         auto dst_view_0 = make_naive_tensor_view_packed<address_space_enum::global>(
             p_dst, perm_length, number<1>{});
  
         auto dst_view = transform_tensor_view(
             dst_view_0,
             make_tuple(make_merge_transform(perm_length)),
             make_tuple(typename arithmetic_sequence_gen<0, kMaxRanks, 1>::type{}),
             make_tuple(sequence<0>{}));
  
         // TODO: hard code to vector 1
         using vector_t = thread_buffer<DataType, 1>;
  
         const auto src_coord =
             make_tensor_coordinate(src_view.get_tensor_descriptor(), array<index_t, 1>{id});
         const auto dst_coord =
             make_tensor_coordinate(dst_view.get_tensor_descriptor(), array<index_t, 1>{id});
  
         // printf("src id:%d, os:%d\n", id, src_coord.get_offset());
         // printf("dst id:%d, os:%d\n", id, dst_coord.get_offset());
  
         const vector_t x = src_view.template get_vectorized_elements<vector_t>(src_coord, 0);
         dst_view.template set_vectorized_elements<vector_t>(dst_coord, 0, x);
     }
 };
  
 } // namespace ck_tile