From 2b0f3791a6e8edf47d199600a61deb05ca1c4d42 Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Wed, 6 Aug 2025 14:40:38 +0000 Subject: [PATCH] [ck_tile] Add get_partition_index_v2 which uses warp_id in vgpr and to be used by tile_windows on lds-based tensor_view --- .../ck_tile/core/tensor/tile_distribution.hpp | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/include/ck_tile/core/tensor/tile_distribution.hpp b/include/ck_tile/core/tensor/tile_distribution.hpp index 606b346e19..9392b9be48 100644 --- a/include/ck_tile/core/tensor/tile_distribution.hpp +++ b/include/ck_tile/core/tensor/tile_distribution.hpp @@ -114,7 +114,22 @@ struct tile_distribution } else if constexpr(NDimP == 2) { - return array{get_warp_id(), get_lane_id()}; + return array{get_warp_id_to_sgpr(), get_lane_id()}; + } + } + + CK_TILE_HOST_DEVICE static auto _get_partition_index_v2() + { + // only support warp-tile and block-tile + static_assert(NDimP == 1 or NDimP == 2, "wrong!"); + + if constexpr(NDimP == 1) + { + return array{get_lane_id()}; + } + else if constexpr(NDimP == 2) + { + return array{get_warp_id_to_vgpr(), get_lane_id()}; } }