[CK_BUILDER] Forward convolution builder improvements

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

@@ -95,7 +95,8 @@ concept AccessOrderDescriptor = requires(T t) {
     { t.order } -> std::convertible_to<std::array<size_t, 3>>;
 };
 
-// No requirements yet for a ConvAlgorithm concept.
+// Base requirement for all ConvAlgorithm concepts, i.e., all conv algorithm concepts must meet this
+// concept.
 template <typename T>
 concept ConvAlgorithmDescriptor = std::is_class_v<T>;
 
@@ -183,6 +184,12 @@ concept SpecifiesLoopScheduler = requires {
     { T::loop_scheduler } -> std::convertible_to<PipelineScheduler>;
 };
 
+template <typename T>
+concept SpecifiesLargeTensorSupport = requires {
+    { T::specialization } -> std::convertible_to<ConvAlgorithmSpecialization>;
+    requires T::specialization == ConvAlgorithmSpecialization::LARGE_TENSOR;
+};
+
 /******************************************** */
 /* DL-specific descriptors and requirements   */
 /******************************************** */
@@ -204,21 +211,9 @@ concept DlThreadClusterDescriptor = requires(T t) {
     { t.n1_xs } -> std::convertible_to<std::array<size_t, 2>>;
 };
 
-// Concept for DL block transfer K0_M0_M1_K1 format
-template <typename T>
-concept DlBlockTransferK0M0M1K1Descriptor = requires(T t) {
-    { t.thread_slice_lengths } -> std::convertible_to<std::array<size_t, 4>>;
-    { t.thread_cluster_lengths } -> std::convertible_to<std::array<size_t, 4>>;
-    { t.thread_cluster_arrange_order } -> std::convertible_to<std::array<size_t, 4>>;
-    { t.src_access_order } -> std::convertible_to<std::array<size_t, 4>>;
-    { t.src_vector_tensor_lengths } -> std::convertible_to<std::array<size_t, 4>>;
-    { t.src_vector_tensor_contiguous_dim_order } -> std::convertible_to<std::array<size_t, 4>>;
-    { t.dst_vector_tensor_lengths } -> std::convertible_to<std::array<size_t, 4>>;
-};
-
-// Concept for DL block transfer K0_N0_N1_K1 format
+// Concept for DL block transfer
 template <typename T>
-concept DlBlockTransferK0N0N1K1Descriptor = requires(T t) {
+concept DlBlockTransferDescriptor = requires(T t) {
     { t.thread_slice_lengths } -> std::convertible_to<std::array<size_t, 4>>;
     { t.thread_cluster_lengths } -> std::convertible_to<std::array<size_t, 4>>;
     { t.thread_cluster_arrange_order } -> std::convertible_to<std::array<size_t, 4>>;
@@ -228,9 +223,9 @@ concept DlBlockTransferK0N0N1K1Descriptor = requires(T t) {
     { t.dst_vector_tensor_lengths } -> std::convertible_to<std::array<size_t, 4>>;
 };
 
-// Concept for DL C thread transfer
+// Concept for DL epilogue
 template <typename T>
-concept DlCThreadTransferDescriptor = requires(T t) {
+concept DlEpilogueDescriptor = requires(T t) {
     { t.src_dst_access_order } -> std::convertible_to<std::array<size_t, 6>>;
     { t.src_dst_vector_dim } -> std::convertible_to<size_t>;
     { t.dst_scalar_per_vector } -> std::convertible_to<size_t>;
@@ -239,31 +234,63 @@ concept DlCThreadTransferDescriptor = requires(T t) {
 // Concept to check if algorithm specifies DL thread config
 template <typename T>
 concept SpecifiesDlThreadConfig = requires {
-    { T::dl_thread_config } -> DlThreadConfigDescriptor;
+    { T::thread_config } -> DlThreadConfigDescriptor;
 };
 
 // Concept to check if algorithm specifies DL thread cluster
 template <typename T>
 concept SpecifiesDlThreadCluster = requires {
-    { T::dl_thread_cluster } -> DlThreadClusterDescriptor;
+    { T::thread_cluster } -> DlThreadClusterDescriptor;
 };
 
-// Concept to check if algorithm specifies DL A block transfer
+// Concept to check if algorithm specifies DL block transfer
 template <typename T>
-concept SpecifiesDlBlockTransferA = requires {
-    { T::dl_block_transfer_a } -> DlBlockTransferK0M0M1K1Descriptor;
+concept SpecifiesDlBlockTransfer = requires {
+    { T::block_transfer_a } -> DlBlockTransferDescriptor;
+    { T::block_transfer_b } -> DlBlockTransferDescriptor;
 };
 
-// Concept to check if algorithm specifies DL B block transfer
+// Concept to check if algorithm specifies DL C thread transfer
 template <typename T>
-concept SpecifiesDlBlockTransferB = requires {
-    { T::dl_block_transfer_b } -> DlBlockTransferK0N0N1K1Descriptor;
+concept SpecifiesDlEpilogue = requires {
+    { T::epilogue_c } -> DlEpilogueDescriptor;
 };
 
-// Concept to check if algorithm specifies DL C thread transfer
+/******************************************** */
+/* Concepts for the different device ops */
+/******************************************** */
+
 template <typename T>
-concept SpecifiesDlCThreadTransfer = requires {
-    { T::dl_c_thread_transfer } -> DlCThreadTransferDescriptor;
-};
+concept DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3 =
+    ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesGridwiseXdlGemm<T> &&
+    SpecifiesBlockTransfer<T> && SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
+    SpecifiesSourceAccessOrder<T> && SpecifiesFwdConcSpecialization<T> &&
+    SpecifiesGemmSpecialization<T> && SpecifiesBlockGemm<T>;
+
+template <typename T>
+concept DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle =
+    ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesGridwiseXdlGemm<T> &&
+    SpecifiesBlockTransfer<T> && SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
+    SpecifiesSourceAccessOrder<T> && SpecifiesFwdConcSpecialization<T> &&
+    SpecifiesGemmSpecialization<T> && SpecifiesNumPrefetchStages<T> &&
+    SpecifiesNumGroupsToMerge<T> && SpecifiesLoopScheduler<T>;
+
+template <typename T>
+concept DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle =
+    ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesGridwiseWmmaGemm<T> &&
+    SpecifiesBlockTransfer<T> && SpecifiesLdsTransfer<T> && SpecifiesThreadClusterAccessOrder<T> &&
+    SpecifiesSourceAccessOrder<T> && SpecifiesFwdConcSpecialization<T> &&
+    SpecifiesGemmSpecialization<T> && SpecifiesNumPrefetchStages<T> && SpecifiesLoopScheduler<T>;
+
+template <typename T>
+concept DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK =
+    ConvAlgorithmDescriptor<T> && SpecifiesThreadBlock<T> && SpecifiesFwdConcSpecialization<T> &&
+    SpecifiesGemmSpecialization<T> && SpecifiesDlThreadConfig<T> && SpecifiesDlThreadCluster<T> &&
+    SpecifiesDlBlockTransfer<T> && SpecifiesDlEpilogue<T>;
+
+template <typename T>
+concept DeviceGroupedConvFwdMultipleD_Xdl_CShuffle_Large_Tensor =
+    DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<decltype(T::base_algorithm)> &&
+    SpecifiesLargeTensorSupport<T>;
 
 } // namespace ck_tile::builder