[CK_BUILDER] Refactor warp GEMM and conv algorithm concepts

experimental/builder/include/ck_tile/builder/conv_algorithm_concepts.hpp

@@ -30,62 +30,57 @@ concept ThreadBlockDescriptor = requires(T t) {
 
 // Concept for parameters that describe a gridwise XDL GEMM problem.
 template <typename T>
-concept GridwiseXdlGemmDescriptor = requires(T t) {
-    { t.m_per_xdl } -> SizeType;
-    { t.n_per_xdl } -> SizeType;
-    { t.m_xdl_per_wave } -> SizeType;
-    { t.n_xdl_per_wave } -> SizeType;
+concept WarpGemmDescriptor = requires(T t) {
+    { t.matrix_instruction } -> std::convertible_to<MatrixInstructionType>;
+    { t.gemm_m_per_instruction } -> SizeType;
+    { t.gemm_n_per_instruction } -> SizeType;
+    { t.gemm_m_iters_per_wave } -> SizeType;
+    { t.gemm_n_iters_per_wave } -> SizeType;
 };
 
-// Concept for parameter that describe block GEMM problem.
+// Concept for parameters that describe the GEMM pipeline.
 template <typename T>
-concept BlockGemmPipelineDescriptor = requires(T t) {
+concept GemmPipelineDescriptor = requires(T t) {
+    { t.num_conv_groups_to_merge } -> SizeType;
+    { t.num_gemm_k_prefetch_stages } -> SizeType;
     { t.pipeline_version } -> std::convertible_to<PipelineVersion>;
     { t.scheduler } -> std::convertible_to<PipelineScheduler>;
 };
 
-// Concept for parameters that describe a gridwise WMMA GEMM problem.
-template <typename T>
-concept GridwiseWmmaGemmDescriptor = requires(T t) {
-    { t.k1 } -> SizeType;
-    { t.m_per_wmma } -> SizeType;
-    { t.n_per_wmma } -> SizeType;
-    { t.m_wmma_per_wave } -> SizeType;
-    { t.n_wmma_per_wave } -> SizeType;
-};
-
 // Concept for vectorized data transfer for convolution input tensors.
 template <typename T>
-concept BlockTransferDescriptor3D = requires(T t) {
+concept InputTileThreadClusterDescriptor3D = requires(T t) {
     { t.k0 } -> SizeType;
     { t.m_n } -> SizeType;
     { t.k1 } -> SizeType;
 };
 
 template <typename T>
-concept BlockTransferDescriptor4D = requires(T t) {
+concept InputTileThreadClusterDescriptor4D = requires(T t) {
     { t.k0 } -> SizeType;
     { t.m_n } -> SizeType;
     { t.k1 } -> SizeType;
     { t.k_batch_size } -> SizeType;
 };
 
 template <typename T, size_t ThreadClusterRank>
-concept BlockTransferDescriptor = (ThreadClusterRank == 3 && BlockTransferDescriptor3D<T>) ||
-                                  (ThreadClusterRank == 4 && BlockTransferDescriptor4D<T>);
+concept InputTileThreadClusterDescriptor =
+    (ThreadClusterRank == 3 && InputTileThreadClusterDescriptor3D<T>) ||
+    (ThreadClusterRank == 4 && InputTileThreadClusterDescriptor4D<T>);
 
 // Concept for thread cluster dimensions for GEMM output tensor.
 template <typename T>
-concept ThreadClusterDescriptor = requires(T t) {
-    { t.m_block } -> SizeType;
-    { t.m_wave_per_xdl } -> SizeType;
-    { t.n_block } -> SizeType;
-    { t.n_wave_per_xdl } -> SizeType;
+concept OutputTileThreadClusterDescriptor = requires(T t) {
+    { t.gemm_m_block_size } -> SizeType;
+    { t.gemm_m_per_block } -> SizeType;
+    { t.gemm_n_block_size } -> SizeType;
+    { t.gemm_n_per_block } -> SizeType;
 };
 
 // Concept for the LDS transfer for the convolution input tensors.
 template <typename T>
 concept LdsTransferDescriptor = requires(T t) {
+    { t.global_memory_vector_load_size } -> SizeType;
     { t.src_vector_dim } -> SizeType;
     { t.src_scalar_per_vector } -> SizeType;
     { t.lds_dst_scalar_per_vector } -> SizeType;
@@ -172,45 +167,18 @@ concept SpecifiesTileThreadBlock = requires {
     { T::thread_block } -> TileThreadBlockDescriptor;
 };
 
-// Concept to check if a struct specifies gridwise XDL GEMM info.
-template <typename T>
-concept GridwiseFwdXdlGemmDescriptor = requires(T t) {
-    { t.ak1 } -> SizeType;
-    { t.bk1 } -> SizeType;
-    { t.xdl_params } -> GridwiseXdlGemmDescriptor;
-};
-
-// Concept to check if a struct specifies gridwise XDL GEMM info.
-template <typename T>
-concept GridwiseBwdXdlGemmDescriptor = requires(T t) {
-    { t.k1 } -> SizeType;
-    { t.xdl_params } -> GridwiseXdlGemmDescriptor;
-};
-
-// Concept to check if a struct specifies gridwise XDL GEMM info.
+// Concept to check if a struct specifies warp GEMM info.
 template <typename T>
-concept SpecifiesGridwiseFwdXdlGemm = requires(T t) {
-    { t.gridwise_gemm } -> GridwiseFwdXdlGemmDescriptor;
-};
-
-// Concept to check if a struct specifies gridwise XDL GEMM info.
-template <typename T>
-concept SpecifiesGridwiseBwdXdlGemm = requires(T t) {
-    { t.gridwise_gemm } -> GridwiseBwdXdlGemmDescriptor;
-};
-
-// Concept to check if a struct specifies gridwise WMMA GEMM info.
-template <typename T>
-concept SpecifiesGridwiseWmmaGemm = requires(T t) {
-    { t.gridwise_gemm } -> GridwiseWmmaGemmDescriptor;
+concept SpecifiesWarpGemm = requires {
+    { T::warp_gemm } -> WarpGemmDescriptor;
 };
 
 // Concept to check if a struct specifies convolution input and output block transfer info.
 template <typename T, size_t ThreadClusterRank = 3>
-concept SpecifiesBlockTransfer = requires(T t) {
-    { T::transfer.a.block_transfer } -> BlockTransferDescriptor<ThreadClusterRank>;
-    { T::transfer.b.block_transfer } -> BlockTransferDescriptor<ThreadClusterRank>;
-    { T::transfer.c.thread_cluster_dims } -> ThreadClusterDescriptor;
+concept SpecifiesThreadClusters = requires(T t) {
+    { T::transfer.a.thread_cluster } -> InputTileThreadClusterDescriptor<ThreadClusterRank>;
+    { T::transfer.b.thread_cluster } -> InputTileThreadClusterDescriptor<ThreadClusterRank>;
+    { T::transfer.c.thread_cluster } -> OutputTileThreadClusterDescriptor;
 };
 
 // Concept to check if a struct specifies convolution scalar per vector infor for A, B and C.
@@ -232,8 +200,8 @@ concept SpecifiesLdsTransfer = requires(T t) {
 // Concept to check if a struct specifies thread cluster access order info.
 template <typename T>
 concept SpecifiesThreadClusterAccessOrder = requires(T t) {
-    { T::transfer.a.block_transfer_access_order } -> AccessOrderDescriptor;
-    { T::transfer.b.block_transfer_access_order } -> AccessOrderDescriptor;
+    { T::transfer.a.thread_cluster_access_order } -> AccessOrderDescriptor;
+    { T::transfer.b.thread_cluster_access_order } -> AccessOrderDescriptor;
 };
 
 // Concept to check if a struct specifies source access order info.
@@ -245,13 +213,8 @@ concept SpecifiesSourceAccessOrder = requires(T t) {
 
 // Concept to check if struct specifies block GEMM.
 template <typename T>
-concept SpecifiesBlockGemm = requires {
-    { T::block_gemm_pipeline } -> BlockGemmPipelineDescriptor;
-};
-
-template <typename T>
-concept SpecifiesGridwiseGemmPipeline = requires {
-    { T::pipeline_version } -> std::convertible_to<PipelineVersion>;
+concept SpecifiesGemmPipeline = requires {
+    { T::gemm_pipeline } -> GemmPipelineDescriptor;
 };
 
 // Concept to check if struct specifies block GEMM (CK Tile).
@@ -307,16 +270,6 @@ concept SpecifiesNumGroupsToMerge = requires {
     { T::num_conv_groups_to_merge } -> SizeType;
 };
 
-template <typename T>
-concept SpecifiesLoopScheduler = requires {
-    { T::loop_scheduler } -> std::convertible_to<PipelineScheduler>;
-};
-
-template <typename T>
-concept SpecifiesGenericInstance = !requires {
-    { T::specialization };
-};
-
 template <typename T>
 concept SpecifiesTransposeTransfer = requires {
     { T::max_transpose_transfer_src_scalar_per_vector } -> SizeType;
@@ -333,38 +286,58 @@ template <typename T>
 concept TransposeTransferWellDefinedIfProvided =
     !HasTransposeTransfer<T> || SpecifiesTransposeTransfer<T>;
 
-template <typename T>
-concept SpecifiesGemmBatchOptions = requires {
-    { T::num_conv_groups_to_merge } -> SizeType;
-};
-
 /******************************************** */
 /* Algorithm specialization concepts          */
 /******************************************** */
 template <typename T>
 concept SpecifiesLargeTensorSupport = requires {
     { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
-    requires T::specialization == ConvAlgorithmSpecialization::LARGE_TENSOR;
+    requires !!(T::specialization & ConvAlgorithmSpecialization::LARGE_TENSOR);
 };
 
 template <typename T>
 concept SpecifiesReferenceAlgorithm = requires {
     { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
-    requires T::specialization == ConvAlgorithmSpecialization::REFERENCE;
+    requires !!(T::specialization & ConvAlgorithmSpecialization::REFERENCE);
 };
 
 template <typename T>
 concept SpecifiesTwoStageSupport = requires {
     { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
-    requires T::specialization == ConvAlgorithmSpecialization::TWO_STAGE;
+    requires !!(T::specialization & ConvAlgorithmSpecialization::TWO_STAGE);
 };
 
 template <typename T>
 concept SpecifiesMultipleDSupport = requires {
     { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
-    requires T::specialization == ConvAlgorithmSpecialization::MULTIPLE_D;
+    requires !!(T::specialization & ConvAlgorithmSpecialization::MULTIPLE_D);
+};
+
+template <typename T>
+concept SpecifiesPipelineV3 = requires {
+    { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
+    requires !!(T::specialization & ConvAlgorithmSpecialization::PIPELINE_V3);
 };
 
+template <typename T>
+concept SpecifiesGenericInstance = !requires {
+    { T::specialization };
+} || requires {
+    { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
+    requires !!(T::specialization == ConvAlgorithmSpecialization::NONE);
+};
+
+template <auto Algorithm>
+concept SpecifiesXdl =
+    requires { requires Algorithm.warp_gemm.matrix_instruction == MatrixInstructionType::XDL; };
+
+template <auto Algorithm>
+concept SpecifiesWmma =
+    requires { requires Algorithm.warp_gemm.matrix_instruction == MatrixInstructionType::WMMA; };
+
+template <auto Algorithm>
+concept SpecifiesValidWarpGemm = SpecifiesXdl<Algorithm> || SpecifiesWmma<Algorithm>;
+
 /******************************************** */
 /* DL-specific descriptors and requirements   */
 /******************************************** */