Few fixes to B2C map and new functionality.

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

@@ -1225,6 +1225,12 @@ struct BlockToCTileMap_LinearKSplit
         return make_tuple(M0_idx_, N0_idx_, K0_idx_);
     }
 
+    __host__ __device__ index_t GetOutputTileIdx() const
+    {
+        const auto N0 = math::integer_divide_ceil(N_, NPerBlock);
+        return M0_idx_ * N0 + N0_idx_;
+    }
+
     template <typename CTileIdx, typename CTileDim>
     __host__ __device__ bool ValidCTileIndex(const CTileIdx& /* c_tile_idx */,
                                              const CTileDim& /* c_tile_dim */) const
@@ -1234,8 +1240,13 @@ struct BlockToCTileMap_LinearKSplit
 
     __host__ __device__ bool GetNextKTileIdx()
     {
-        K0_idx_++;
-        return K0_idx_ < KSplit_;
+        if(K0_idx_ + 1 < KSplit_)
+        {
+            K0_idx_++;
+            return true;
+        }
+        else
+            return false;
     }
 
     ///
@@ -1247,7 +1258,7 @@ struct BlockToCTileMap_LinearKSplit
     ///
     __host__ __device__ bool IsFirstKSplitBlock(index_t tiles_per_block) const
     {
-        return (K0_idx_ - tiles_per_block) <= 0;
+        return (K0_idx_ + 1 - tiles_per_block) <= 0;
     }
 
     __host__ __device__ index_t GetTileMIdx() const { return M0_idx_; }

test/block_to_ctile_map/test_block_to_ctile_map.cpp

@@ -364,11 +364,12 @@ TEST(BlockToCTileMap, BlockToCTileMap_LinearKSplit_BottomIndex)
 
 TEST(BlockToCTileMap, BlockToCTileMap_LinearKSplit_NextKTile)
 {
-    const index_t M         = 768;
-    const index_t N         = 384;
-    const index_t MPerBlock = 128;
-    const index_t NPerBlock = 64;
-    const index_t KSplit    = 3;
+    const index_t M               = 768;
+    const index_t N               = 384;
+    const index_t MPerBlock       = 128;
+    const index_t NPerBlock       = 64;
+    const index_t KSplit          = 3;
+    const index_t tiles_per_block = 1;
 
     auto c_grid_desc_m_n = make_naive_tensor_descriptor_packed(make_tuple(M, N));
 
@@ -377,6 +378,7 @@ TEST(BlockToCTileMap, BlockToCTileMap_LinearKSplit_NextKTile)
     auto m0n0k0_idx = tile_map.CalculateBottomIndex(3);
     EXPECT_EQ((std::vector<int>{m0n0k0_idx[I0], m0n0k0_idx[I1], m0n0k0_idx[I2]}),
               (std::vector<int>{0, 1, 0}));
+    EXPECT_TRUE(tile_map.IsFirstKSplitBlock(tiles_per_block));
 
     for(index_t i = 0; i < KSplit - 1; i++)
     {
@@ -384,9 +386,11 @@ TEST(BlockToCTileMap, BlockToCTileMap_LinearKSplit_NextKTile)
         m0n0k0_idx = tile_map.GetBottomIndex();
         EXPECT_EQ((std::vector<int>{m0n0k0_idx[I0], m0n0k0_idx[I1], m0n0k0_idx[I2]}),
                   (std::vector<int>{0, 1, i + 1}));
+        EXPECT_FALSE(tile_map.IsFirstKSplitBlock(tiles_per_block));
     }
     EXPECT_FALSE(tile_map.GetNextKTileIdx());
     m0n0k0_idx = tile_map.GetBottomIndex();
     EXPECT_EQ((std::vector<int>{m0n0k0_idx[I0], m0n0k0_idx[I1], m0n0k0_idx[I2]}),
-              (std::vector<int>{0, 1, 3}));
+              (std::vector<int>{0, 1, 2}));
+    EXPECT_FALSE(tile_map.IsFirstKSplitBlock(tiles_per_block));
 }