From fa7dba2cd1c5333bb1ee556a535179a586ea2bae Mon Sep 17 00:00:00 2001
From: Qianfeng Zhang <Qianfeng.Zhang@amd.com>
Date: Fri, 8 Aug 2025 06:20:38 +0000
Subject: [PATCH] [ck_tile] Merge get_partition_index() and
 get_partition_index_v2() to get_partition_index() with bool_constant
 parameter

---
 .../ck_tile/core/tensor/tile_distribution.hpp  | 18 ++----------------
 1 file changed, 2 insertions(+), 16 deletions(-)

diff --git a/include/ck_tile/core/tensor/tile_distribution.hpp b/include/ck_tile/core/tensor/tile_distribution.hpp
index 9392b9be48..8fe16a1693 100644
--- a/include/ck_tile/core/tensor/tile_distribution.hpp
+++ b/include/ck_tile/core/tensor/tile_distribution.hpp
@@ -114,22 +114,8 @@ struct tile_distribution
         }
         else if constexpr(NDimP == 2)
         {
-            return array<index_t, 2>{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<index_t, 1>{get_lane_id()};
-        }
-        else if constexpr(NDimP == 2)
-        {
-            return array<index_t, 2>{get_warp_id_to_vgpr(), get_lane_id()};
+            return array<index_t, 2>{get_warp_id(bool_constant<save_warp_id_in_sgpr>{}),
+                                     get_lane_id()};
         }
     }