Addition of the derived structs for the new Stream-K TilePartitioner

There are 2 derived structs based on whether Stream-K is persistent or not.
If it's persistent that means that both the data parallel and Stream-K sections
are data parallel. If it's non-persistent that means that only the
Stream-K section is persistent, while the data parallel section will have
separate workgroups allocated for it. Both structs will have a template
argument for Persistent.

The 2 derived classes will inherit common variables and functions from the
Stream-K TilePartitioner base class. There are additional variables for the
differing data parallel sections that will be added to each derived class,
that are in charge of the indexing/bookkeeping for the data parallel sections.
The only additional function that will differ between the 2 structs is GridSize(),
as the non-persistent will allocate extra workgroups for data parallel.

Unit tests for the derived structs are included.

include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp

@@ -10,6 +10,8 @@
 
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/common.hpp"
+#include <format>
+#include <iostream>
 
 namespace ck_tile {
 
@@ -810,4 +812,5 @@ struct StreamKTilePartitioner
     uint32_t M_, N_, K_;
     uint32_t num_tile_m_, num_tile_n_, num_tile_k_;
 };
+
 } // namespace ck_tile

include/ck_tile/ops/gemm/kernel/streamk_gemm_tile_partitioner.hpp

@@ -202,6 +202,117 @@ struct StreamKTilePartitionerBase
     index_t n_;
 };
 
+/**
+ * @brief Template for the Stream-K tile partitioner derived struct.
+ *
+ * This partitioner is responsible for mapping workgroups to tiles in the C tensor
+ * for the Stream-K algorithm. This struct is derived from
+ * StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategy>. Behavior of the
+ * StreamKTilePartitioner based on persistency will be in the template specializations.
+ *
+ *  @tparam BlockGemmShapeType  A class providing basic GEMM parameters.
+ *  @tparam ReductionStrategy   An enum that defines the reduction strategy for the results in the C
+ * Tensor.
+ *  @tparam Persistent          A bool that indicates whether to use a Persistent approach
+ */
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy, bool Persistent>
+struct StreamKTilePartitioner_v2;
+
+/**
+ * @brief Persistent Stream-K tile partitioner derived struct.
+ *
+ * This partitioner is responsible for mapping workgroups to tiles in the C tensor
+ * for the Stream-K algorithm when using a Persistent approach where no extra workgroups
+ * are allocated for data parallel.
+ *
+ *  @tparam BlockGemmShapeType  A class providing basic GEMM parameters.
+ *  @tparam ReductionStrategy   An enum that defines the reduction strategy for the results in the C
+ * Tensor.
+ */
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+struct StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, true>
+    : StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategy>
+{
+    StreamKTilePartitioner_v2(ck_tile::index_t m,
+                              ck_tile::index_t n,
+                              ck_tile::index_t k,
+                              ck_tile::index_t grid);
+
+    public:
+    /**
+     * @brief Calculates the launching grid size for the Stream-K kernel. In the Persistent
+     * case, no extra workgroups are allocated for the data parallel section, making the grid
+     * size num_cu * occupancy.
+     *
+     * @return dim_3           The launching grid size for the kernel.
+     */
+    CK_TILE_HOST auto grid_size() const noexcept -> dim3;
+
+    CK_TILE_HOST_DEVICE index_t get_dp_tiles_per_cta() const noexcept;
+    CK_TILE_HOST_DEVICE index_t get_extra_dp_tiles() const noexcept;
+
+    protected:
+    /**
+     * @brief The total number of DP tiles per workgroup.
+     */
+    int dp_tiles_per_cta_;
+
+    /**
+     * @brief The total number of DP tiles left over when dp_tiles is not evenly divisible by grid.
+     */
+    int extra_dp_tiles_;
+};
+
+/**
+ * @brief Non-Persistent Stream-K tile partitioner derived struct.
+ *
+ * This partitioner is responsible for mapping workgroups to tiles in the C tensor
+ * for the Stream-K algorithm when using a Non-Persistent approach where extra workgroups
+ * are allocated for the data parallel section.
+ *
+ *  @tparam BlockGemmShapeType  A class providing basic GEMM parameters.
+ *  @tparam ReductionStrategy   An enum that defines the reduction strategy for the results in the C
+ * Tensor.
+ */
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+struct StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, false>
+    : StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategy>
+{
+    StreamKTilePartitioner_v2(ck_tile::index_t m,
+                              ck_tile::index_t n,
+                              ck_tile::index_t k,
+                              ck_tile::index_t grid);
+
+    public:
+    /**
+     * @brief Calculates the launching grid size for the Stream-K kernel. In the Non-Persistent
+     * case, extra workgroups are allocated for the data parallel section, making the grid
+     * size the total number of Stream-K and data parallel workgroups.
+     *
+     * @return dim_3           The launching grid size for the kernel.
+     */
+    CK_TILE_HOST auto grid_size() const noexcept -> dim3;
+    CK_TILE_HOST_DEVICE index_t get_dp_ctas() const noexcept;
+    CK_TILE_HOST_DEVICE index_t get_dp_start_block_idx() const noexcept;
+    CK_TILE_HOST_DEVICE index_t get_sk_start_block_idx() const noexcept;
+
+    protected:
+    /**
+     * @brief The total number of DP workgroups.
+     */
+    int dp_ctas_;
+
+    /**
+     * @brief The index that starts the DP workgroups, always 0 in our implementation.
+     */
+    int dp_start_block_idx_;
+
+    /**
+     * @brief The index that starts the Stream-K workgroups, set to the number of dp_tiles.
+     */
+    int sk_start_block_idx_;
+};
+
 } // namespace ck_tile
 
 #include "streamk_gemm_tile_partitioner_impl.hpp"

include/ck_tile/ops/gemm/kernel/streamk_gemm_tile_partitioner_impl.hpp

@@ -211,4 +211,91 @@ StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategy>::get_n() const
     return n_;
 }
 
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy, bool Persistent>
+struct StreamKTilePartitioner_v2;
+
+// child class for Persistent Tile Partitioner
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, true>::StreamKTilePartitioner_v2(
+    ck_tile::index_t m, ck_tile::index_t n, ck_tile::index_t k, ck_tile::index_t grid)
+    : StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategy>(m, n, k, grid)
+{ // inherit from base constructor
+    dp_tiles_per_cta_ = this->dp_tiles_ / this->grid_;
+    extra_dp_tiles_   = this->dp_tiles_ % this->grid_;
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST auto
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, true>::grid_size() const noexcept
+    -> dim3
+{
+    if(extra_dp_tiles_ == 0)
+    {
+        return dim3(this->grid_, 1, 1);
+    }
+    else
+    {
+        return dim3(this->num_tiles_, 1, 1);
+    }
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, true>::get_dp_tiles_per_cta()
+    const noexcept
+{
+    return dp_tiles_per_cta_;
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, true>::get_extra_dp_tiles()
+    const noexcept
+{
+    return extra_dp_tiles_;
+}
+
+// child class for Non-Persistent Tile Partitioner
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, false>::StreamKTilePartitioner_v2(
+    ck_tile::index_t m, ck_tile::index_t n, ck_tile::index_t k, ck_tile::index_t grid)
+    : StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategy>(m, n, k, grid)
+{ // inherit from base constructor
+    dp_ctas_            = this->dp_tiles_;
+    dp_start_block_idx_ = 0;
+    sk_start_block_idx_ = this->dp_tiles_;
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST auto
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, false>::grid_size() const noexcept
+    -> dim3
+{
+    return dim3(dp_ctas_ + this->get_sk_ctas(), 1, 1);
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, false>::get_dp_ctas()
+    const noexcept
+{
+    return dp_ctas_;
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, false>::get_dp_start_block_idx()
+    const noexcept
+{
+    return dp_start_block_idx_;
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategy>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitioner_v2<BlockGemmShapeType, ReductionStrategy, false>::get_sk_start_block_idx()
+    const noexcept
+{
+    return sk_start_block_idx_;
+}
+
 } // namespace ck_tile