refine code

example/ck_tile/39_gemm_softmax_grouped_topk/CMakeLists.txt

@@ -1,9 +1,9 @@
-set(EXAMPLE_REDUCE "tile_example_basic_gemm_softmax")
+set(EXAMPLE_REDUCE "tile_example_basic_gemm_softmax_grouped_topk")
 # not using add_example_executable() to add this target, since we don't want this to have
 # to be included in "make all/install/check"
 message("adding example ${EXAMPLE_REDUCE}")
 
-add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL gemm_softmax.cpp)
+add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL gemm_softmax_grouped_topk.cpp)
 target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
 set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
 

example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp

@@ -18,7 +18,7 @@ namespace ck_tile {
 //  B Tile Window: global memory
 //  C Distributed tensor: register
 template <typename Problem, typename Policy = ck_tile::BlockGemmPipelineAGmemBGmemCRegDefaultPolicy>
-struct BlockGemmSoftmaxPipelineAGmemBGmemCReg
+struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
 {
     using ADataType      = remove_cvref_t<typename Problem::ADataType>;
     using BDataType      = remove_cvref_t<typename Problem::BDataType>;

example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp

@@ -5,8 +5,8 @@
 
 #include "config.h"
 #include "ck_tile/host.hpp"
-#include "gemm.hpp"
-#include "reference_gemm.hpp"
+#include "gemm_softmax_grouped_topk.hpp"
+#include "reference_gemm_softmax_grouped_topk.hpp"
 
 /*
  * Toy code of GEMM
@@ -227,21 +227,13 @@ int main(int argc, char* argv[])
                                                 Ldb,
                                                 Ldout,
                                                 CElementFunction{}));
-    // auto pass      = true;
+
     bool rtn = true;
     if(verification)
     {
-        // reference gemm
-        // ck_tile::HostTensor<CDataType> c_host_ref(c_lengths, c_strides);
-        // reference_basic_gemm_softmax<ADataType, ADataType, AccDataType, CDataType>(
-        //     a_host, b_host, c_host_ref);
-        // c_buf.FromDevice(c_host_dev.mData.data());
-        // pass &= ck_tile::check_err(c_host_dev, c_host_ref);
-        // std::cout << "valid:" << (pass ? "y" : "n") << std::endl;
-
         ck_tile::HostTensor<WeightType> value_ref(out_lengths, out_strides);
         ck_tile::HostTensor<IndexType> index_ref(out_lengths, out_strides);
-        reference_basic_gemm_softmax_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
+        reference_basic_gemm_softmax_grouped_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
             a_host, b_host, value_ref, index_ref, topk);
         value_buf.FromDevice(value_host_dev.mData.data());
         index_buf.FromDevice(index_host_dev.mData.data());
@@ -285,6 +277,5 @@ int main(int argc, char* argv[])
     std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
               << std::endl;
 
-    // return !pass;
     return rtn;
 }

example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp

@@ -8,9 +8,9 @@
 #include "ck_tile/ops/common.hpp"
 #include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
 
-#include "block_gemm_pipeline_agmem_bgmem_creg.hpp"
+#include "block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp"
 #include "config.h"
-#include "grid_gemm.hpp"
+#include "grid_gemm_softmax_grouped_topk.hpp"
 
 namespace ck_tile {
 
@@ -171,7 +171,7 @@ struct Gemm
                                          IndexType,
                                          kBlockSize,
                                          TileGemmShape<kMPerBlock, kNPerBlock, kKPerBlock, kTopKPerBlock>>;
-            return BlockGemmSoftmaxPipelineAGmemBGmemCReg<BlockGemmPipelineProblem_>{};
+            return BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg<BlockGemmPipelineProblem_>{};
         }
     };
 

example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp

@@ -7,7 +7,7 @@
 #include "ck_tile/host/host_tensor.hpp"
 
 template <typename ADataType, typename BDataType, typename AccDataType, typename WeightType, typename IndexType>
-void reference_basic_gemm_softmax_topk(const ck_tile::HostTensor<ADataType>& a_m_k,
+void reference_basic_gemm_softmax_grouped_topk(const ck_tile::HostTensor<ADataType>& a_m_k,
                           const ck_tile::HostTensor<BDataType>& b_n_k,
                           ck_tile::HostTensor<WeightType>& y_values,
                           ck_tile::HostTensor<IndexType>& y_indices,