[CK_BUILDER] Forward convolution builder improvements (#3179)

Proposed changes
Improve the forward convolution builder implementation and addressed leftover feedback left from PR #3138. Main changes

Refactored tests such that they reflect better the builder pattern. The templates and types for the convolution algorithm concepts are created via factory that facilitates programmatic creation of the device op instances.
Moved tests into anonymous namespace.
The convolution factory had lot of if-else constructs when CK Builder types were converted into CK library types. I had initially trouble in using static_assert in the default branch of switch as the static_assert was evaluated at compile time even for valid types. However, if we change the static_assert to throw "<error message>", it will result in a compile-time error only if the default branch is actually hit. This assumes that the function is consteval. Hence, changed all conversions in the convolution factory to use switch, which is more intuitive.
Removed the explicit device op definition from convolution signature and the corresponding predicate file. The device ops are defined by the corresponding concepts. This allowed to remove lot of boilerplate code from the convolution factory.
Adde inheritance and convolution algorithm specialization to handle device ops that are specialization of a more generic ones. The large tensor support is more naturally expressed by this pattern.
Added support for the FP8 data type.

* WIP: Builder for expected test results.

* Improve ckb fwd conv instance tests.

* clang-format

* Change if-else statements into switch in conv factory.

* Fix clang-formatting.

* Removed unnecessary includes.

* Added missing copyright.

* Remove explicit device op flag from from convolution signature.

* Add missing concept.

* Fix build.

* clang-format

* Add test for building conv fwd FP8 instances.

* Add missing header to instance traits.

* Clean-up recently added instances.

* Introduce inheritance and specialization.

* Use builder to build conv algorithm templates and types.

* clang-format

* Fix conv description tests.

---------

Co-authored-by: John Shumway <[email protected]>

Author: vpietila-amd

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