[rocFFT] Add ability to configure kernel per architecture

projects/rocfft/CHANGELOG.md

@@ -5,6 +5,10 @@ Documentation for rocFFT is available at
 
 ## rocFFT 1.0.36 (unreleased)
 
+### Added
+
+* Added support for per precision and architecture kernel configuration entries in the library.
+
 ### Optimized
 
 * Removed a potential unnecessary global transpose operation from MPI 3D multi-GPU pencil decompositions.

projects/rocfft/library/src/CMakeLists.txt

@@ -151,7 +151,7 @@ set( kgen_logic_files
      # python code that decides kernel parameters
      ${CMAKE_SOURCE_DIR}/library/src/device/kernel-generator.py
      ${CMAKE_SOURCE_DIR}/library/src/device/generator.py
-     ${CMAKE_SOURCE_DIR}/library/src/device/kernels/configs/config_lds.py
+     ${CMAKE_SOURCE_DIR}/library/src/device/kernels/configs/config_arch.py
      ${CMAKE_SOURCE_DIR}/library/src/device/kernels/configs/config_sbrr.py
      ${CMAKE_SOURCE_DIR}/library/src/device/kernels/configs/config_sbcc.py
      ${CMAKE_SOURCE_DIR}/library/src/device/kernels/configs/config_sbcr.py

projects/rocfft/library/src/device/generator/stockham_gen.cpp

@@ -22,6 +22,7 @@
 using namespace std::placeholders;
 
 #include "../../../../shared/arithmetic.h"
+#include "../../../../shared/device_properties.h"
 #include "../../../../shared/precision_type.h"
 #include "generator.h"
 #include "stockham_gen.h"
@@ -52,7 +53,8 @@ struct GeneratedLauncher
                       const std::vector<unsigned int>& pp_factors_other,
                       const unsigned int&              pp_current_dim,
                       const unsigned int&              pp_off_dim,
-                      bool                             double_precision,
+                      const unsigned int&              precision_type,
+                      const std::string&               gcn_arch_name,
                       const std::string&               sbrc_type,
                       const std::string&               sbrc_transpose_type)
         : scheme(scheme)
@@ -69,7 +71,9 @@ struct GeneratedLauncher
         , direct_to_from_reg(kernel.direct_to_from_reg)
         , sbrc_type(sbrc_type)
         , sbrc_transpose_type(sbrc_transpose_type)
-        , double_precision(double_precision)
+        , precision_type(precision_type)
+        , gcn_arch_name(gcn_arch_name)
+
     {
     }
 
@@ -90,7 +94,10 @@ struct GeneratedLauncher
     // SBRC transpose type
     std::string sbrc_type;
     std::string sbrc_transpose_type;
-    bool        double_precision;
+
+    unsigned int precision_type;
+
+    std::string gcn_arch_name;
 
     // output a json object that the python generator can parse to know
     // how to build the function pool
@@ -129,7 +136,8 @@ struct GeneratedLauncher
         add_member("direct_to_from_reg", direct_to_from_reg ? "true" : "false");
         add_member("sbrc_type", quote_str(sbrc_type));
         add_member("sbrc_transpose_type", quote_str(sbrc_transpose_type));
-        add_member("double_precision", double_precision ? "true" : "false");
+        add_member("precision_type", std::to_string(precision_type));
+        add_member("gcn_arch_name", quote_str(gcn_arch_name));
         add_member("pp_child_scheme", quote_str(pp_child_scheme));
         add_member("pp_factors_curr", vec_to_list(pp_factors_curr));
         add_member("pp_factors_other", vec_to_list(pp_factors_other));
@@ -149,31 +157,31 @@ struct LaunchSuffix
     std::string sbrc_transpose_type;
 };
 
-void make_launcher(const std::vector<unsigned int>& precision_types,
+void make_launcher(const unsigned int&              precision_type,
                    const std::vector<LaunchSuffix>& launcher_suffixes,
                    StockhamKernel&                  kernel,
+                   const std::string&               gcn_arch_name,
                    const std::string&               pp_child_scheme,
                    const std::vector<unsigned int>& pp_factors_curr,
                    const std::vector<unsigned int>& pp_factors_other,
                    const unsigned int&              pp_current_dim,
                    const unsigned int&              pp_off_dim,
                    std::vector<GeneratedLauncher>&  generated_launchers)
 {
-    for(auto precision_type : precision_types)
+
+    for(auto&& launcher : launcher_suffixes)
     {
-        for(auto&& launcher : launcher_suffixes)
-        {
-            generated_launchers.emplace_back(kernel,
-                                             launcher.scheme,
-                                             pp_child_scheme,
-                                             pp_factors_curr,
-                                             pp_factors_other,
-                                             pp_current_dim,
-                                             pp_off_dim,
-                                             precision_type == rocfft_precision_double,
-                                             launcher.sbrc_type,
-                                             launcher.sbrc_transpose_type);
-        }
+        generated_launchers.emplace_back(kernel,
+                                         launcher.scheme,
+                                         pp_child_scheme,
+                                         pp_factors_curr,
+                                         pp_factors_other,
+                                         pp_current_dim,
+                                         pp_off_dim,
+                                         precision_type,
+                                         gcn_arch_name,
+                                         launcher.sbrc_type,
+                                         launcher.sbrc_transpose_type);
     }
 }
 
@@ -255,9 +263,10 @@ void stockham_partial_pass_variants(const std::string&               kernel_name
         if(params_1.current_dim == 0 && params_2.current_dim == 2)
         {
             StockhamPartialPassKernelRR kernelRR(specs1, params_1);
-            make_launcher(specs1.precisions,
+            make_launcher(specs1.precision,
                           {{"pp_stoc", specs1.scheme, "", ""}},
                           kernelRR,
+                          specs1.gcn_arch_name,
                           "CS_KERNEL_STOCKHAM_PP",
                           params_1.pp_factors_curr,
                           params_1.pp_factors_other,
@@ -266,9 +275,10 @@ void stockham_partial_pass_variants(const std::string&               kernel_name
                           launchers);
 
             StockhamPartialPassKernelCC kernelCC(specs2, params_2, false);
-            make_launcher(specs2.precisions,
+            make_launcher(specs2.precision,
                           {{"pp_sbcc", specs2.scheme, "", ""}},
                           kernelCC,
+                          specs2.gcn_arch_name,
                           "CS_KERNEL_STOCKHAM_PP_BLOCK_CC",
                           params_2.pp_factors_curr,
                           params_2.pp_factors_other,
@@ -279,9 +289,10 @@ void stockham_partial_pass_variants(const std::string&               kernel_name
         else if(params_1.current_dim == 2 && params_2.current_dim == 0)
         {
             StockhamPartialPassKernelRR kernelCC(specs1, params_1);
-            make_launcher(specs1.precisions,
+            make_launcher(specs1.precision,
                           {{"pp_sbcc", specs1.scheme, "", ""}},
                           kernelCC,
+                          specs1.gcn_arch_name,
                           "CS_KERNEL_STOCKHAM_PP_BLOCK_CC",
                           params_1.pp_factors_curr,
                           params_1.pp_factors_other,
@@ -290,9 +301,10 @@ void stockham_partial_pass_variants(const std::string&               kernel_name
                           launchers);
 
             StockhamPartialPassKernelCC kernelRR(specs2, params_2, false);
-            make_launcher(specs2.precisions,
+            make_launcher(specs2.precision,
                           {{"pp_stoc", specs2.scheme, "", ""}},
                           kernelRR,
+                          specs2.gcn_arch_name,
                           "CS_KERNEL_STOCKHAM_PP",
                           params_2.pp_factors_curr,
                           params_2.pp_factors_other,
@@ -341,9 +353,10 @@ void stockham_variants(const std::string&            kernel_name,
     if(specs.scheme == "CS_KERNEL_STOCKHAM")
     {
         StockhamKernelRR kernel(specs);
-        make_launcher(specs.precisions,
+        make_launcher(specs.precision,
                       {{"stoc", specs.scheme, "", ""}},
                       kernel,
+                      specs.gcn_arch_name,
                       "CS_NONE",
                       std::vector<unsigned int>(),
                       std::vector<unsigned int>(),
@@ -354,9 +367,10 @@ void stockham_variants(const std::string&            kernel_name,
     else if(specs.scheme == "CS_KERNEL_STOCKHAM_BLOCK_CC")
     {
         StockhamKernelCC kernel(specs, false, false);
-        make_launcher(specs.precisions,
+        make_launcher(specs.precision,
                       {{"sbcc", specs.scheme, "", ""}},
                       kernel,
+                      specs.gcn_arch_name,
                       "CS_NONE",
                       std::vector<unsigned int>(),
                       std::vector<unsigned int>(),
@@ -405,9 +419,10 @@ void stockham_variants(const std::string&            kernel_name,
                             "SBRC_3D_FFT_ERC_TRANS_Z_XY",
                             "TILE_UNALIGNED"});
 
-        make_launcher(specs.precisions,
+        make_launcher(specs.precision,
                       suffixes,
                       kernel,
+                      specs.gcn_arch_name,
                       "CS_NONE",
                       std::vector<unsigned int>(),
                       std::vector<unsigned int>(),
@@ -419,9 +434,10 @@ void stockham_variants(const std::string&            kernel_name,
     {
         StockhamKernelCR kernel(specs);
 
-        make_launcher(specs.precisions,
+        make_launcher(specs.precision,
                       {{"sbcr", specs.scheme, "", ""}},
                       kernel,
+                      specs.gcn_arch_name,
                       "CS_NONE",
                       std::vector<unsigned int>(),
                       std::vector<unsigned int>(),
@@ -433,39 +449,24 @@ void stockham_variants(const std::string&            kernel_name,
     {
         StockhamKernelFused2D fused2d(specs, specs2d);
 
-        // output 2D launchers
-        for(auto prec_type : specs.precisions)
-        {
-            launchers.emplace_back(fused2d,
-                                   specs.scheme,
-                                   "CS_NONE",
-                                   std::vector<unsigned int>(),
-                                   std::vector<unsigned int>(),
-                                   0,
-                                   0,
-                                   (prec_type == rocfft_precision_double),
-                                   "",
-                                   "");
-        }
+        launchers.emplace_back(fused2d,
+                               specs.scheme,
+                               "CS_NONE",
+                               std::vector<unsigned int>(),
+                               std::vector<unsigned int>(),
+                               0,
+                               0,
+                               specs.precision,
+                               specs.gcn_arch_name,
+                               "",
+                               "");
     }
     else
         throw std::runtime_error("unhandled scheme");
 
     output_json(launchers, kernel_name, output);
 }
 
-static size_t max_bytes_per_element(const std::vector<unsigned int>& precisions)
-{
-    // generate for the maximum element size in the available
-    // precisions
-    size_t element_size = 0;
-    for(auto p : precisions)
-    {
-        element_size = std::max(element_size, complex_type_size(static_cast<rocfft_precision>(p)));
-    }
-    return element_size;
-}
-
 // =========================================================
 // Partial pass parameters row-major ordering helpers.
 // Kernel configuration parameters for CS_3D_PP in
@@ -636,9 +637,11 @@ int main()
     std::string line;
 
     std::string  kernel_name;
+    std::string  gcn_arch_name;
     std::string  scheme;
     bool         half_lds;
     unsigned int lds_size_bytes;
+    unsigned int bytes_per_element;
 
     const char* DELIM = "";
     std::cout << "{";
@@ -661,6 +664,9 @@ int main()
         lds_size_bytes = std::stoul(*arg);
         ++arg;
 
+        bytes_per_element = std::stoul(*arg);
+        ++arg;
+
         std::string kernel_name = *arg;
 
         ++arg;
@@ -698,8 +704,11 @@ int main()
         threads_per_transform = parse_uints_csv(*arg);
 
         ++arg;
-        std::vector<unsigned int> precisions;
-        precisions = parse_uints_csv(*arg);
+        gcn_arch_name = *arg;
+
+        ++arg;
+        unsigned int precision;
+        precision = std::stoul(*arg);
 
         // create spec and pass to stockham_variants, writes partial output to stdout
         std::cout << DELIM;
@@ -733,12 +742,14 @@ int main()
             if(direct_to_from_reg.size() != 2)
                 throw std::runtime_error("CS_3D_PP requires two direct_to_from_reg configuration");
 
-            StockhamGeneratorSpecs specs1(factors1, {}, precisions, workgroup_size[0], scheme);
+            StockhamGeneratorSpecs specs1(
+                factors1, {}, precision, gcn_arch_name, workgroup_size[0], scheme);
             specs1.direct_to_from_reg    = direct_to_from_reg[0];
             specs1.threads_per_transform = threads_per_transform[0];
             specs1.wgs_is_derived        = true;
 
-            StockhamGeneratorSpecs specs2(factors2, {}, precisions, workgroup_size[1], scheme);
+            StockhamGeneratorSpecs specs2(
+                factors2, {}, precision, gcn_arch_name, workgroup_size[1], scheme);
             specs2.direct_to_from_reg    = direct_to_from_reg[1];
             specs2.threads_per_transform = threads_per_transform[1];
             specs2.wgs_is_derived        = true;
@@ -769,24 +780,28 @@ int main()
             ++arg;
             factors = parse_uints_csv(*arg);
 
-            StockhamGeneratorSpecs specs(factors, factors2d, precisions, workgroup_size[0], scheme);
+            StockhamGeneratorSpecs specs(
+                factors, factors2d, precision, gcn_arch_name, workgroup_size[0], scheme);
             specs.half_lds           = half_lds;
             specs.direct_to_from_reg = direct_to_from_reg[0];
 
-            specs.bytes_per_element = max_bytes_per_element(precisions);
+            specs.bytes_per_element = bytes_per_element;
 
             specs.threads_per_transform = threads_per_transform.front();
 
             // second dimension for 2D_SINGLE
             StockhamGeneratorSpecs specs2d(
-                factors2d, factors, precisions, workgroup_size[0], scheme);
+                factors2d, factors, precision, gcn_arch_name, workgroup_size[0], scheme);
 
             if(!threads_per_transform.empty())
                 specs2d.threads_per_transform = threads_per_transform.back();
 
             // 2D_SINGLE kernels use the specified workgroup size
-            // directly
-            if(scheme == "CS_KERNEL_2D_SINGLE")
+            // directly.
+            // Kernels with an architecture other than generic will
+            // also use the workgroup size directly, as the calculations
+            // for wgs_is_derived=false assume an LDS size of 64KiB.
+            if(scheme == "CS_KERNEL_2D_SINGLE" || gcn_arch_name != generic_gcn_arch_name)
             {
                 specs.wgs_is_derived   = true;
                 specs2d.wgs_is_derived = true;

projects/rocfft/library/src/device/generator/stockham_gen.h

@@ -32,12 +32,14 @@ struct StockhamGeneratorSpecs
 {
     StockhamGeneratorSpecs(const std::vector<unsigned int>& factors,
                            const std::vector<unsigned int>& factors2d,
-                           const std::vector<unsigned int>& precisions,
+                           unsigned int                     precision,
+                           const std::string&               gcn_arch_name,
                            unsigned int                     workgroup_size,
                            const std::string&               scheme)
         : factors(factors)
         , factors2d(factors2d)
-        , precisions(precisions)
+        , precision(precision)
+        , gcn_arch_name(gcn_arch_name)
         , length(product(factors.begin(), factors.end()))
         , length2d(product(factors2d.begin(), factors2d.end()))
         , workgroup_size(workgroup_size)
@@ -47,7 +49,8 @@ struct StockhamGeneratorSpecs
 
     std::vector<unsigned int> factors;
     std::vector<unsigned int> factors2d;
-    std::vector<unsigned int> precisions; // mapped from rocfft_precision
+    unsigned int              precision; // mapped from rocfft_precision
+    std::string               gcn_arch_name;
     unsigned int              length;
     unsigned int              length2d = 0;