feat(rocfft): Add support for 32/64 bit index arithmetic on Stockham Kernels

projects/rocfft/library/src/chirp.cpp

@@ -57,9 +57,9 @@ void launch_chirp_kernel(const size_t           N,
     auto blockSize = CHIRP_THREADS;
     auto numBlocks = DivRoundingUp<size_t>(N, blockSize);
 
-    auto          kernel = RTCKernelChirp::generate(deviceProp.gcnArchName, precision);
-    RTCKernelArgs kargs;
-    kargs.append_size_t(N);
+    auto          kernel = RTCKernelChirp::generate(deviceProp.gcnArchName, N, precision);
+    RTCKernelArgs kargs(RTCKernelChirp::itype(deviceProp.gcnArchName, N));
+    kargs.append_kint(N);
     kargs.append_ptr(output);
     kernel.get()->launch(kargs, dim3(numBlocks), dim3(blockSize), 0, deviceProp, stream);
 }

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

@@ -20,6 +20,7 @@
 
 #pragma once
 
+#include "../../include/rtc_kernel.h"
 #include "generator.h"
 
 enum BluesteinOperationType
@@ -47,14 +48,15 @@ struct BluesteinData
     // templates
     //
     Variable scalar_type{"scalar_type", "typename"};
+    Variable integer_type{"integer_type", "typename"};
 
     //
     // internal variables
     //
     Variable chirp{"chirp", "const scalar_type", true, true};
-    Variable data_idx{"data_idx", "size_t"};
-    Variable data_voffset{"data_voffset", "size_t"};
-    Variable data_soffset{"data_soffset", "size_t"};
+    Variable data_idx{"data_idx", "integer_type"};
+    Variable data_voffset{"data_voffset", "integer_type"};
+    Variable data_soffset{"data_soffset", "integer_type"};
     Variable data_rw_flag{"data_rw_flag", "bool"};
     Variable buf_in{"buf_in", "scalar_type", true};
     Variable buf_inre{"buf_inre", "real_type_t<scalar_type>", true, true};
@@ -66,21 +68,21 @@ struct BluesteinData
     Variable data_bufre{"data_bufre", "real_type_t<scalar_type>", true, true};
     Variable data_bufim{"data_bufim", "real_type_t<scalar_type>", true, true};
     Variable data_elem{"data_elem", "scalar_type"};
-    Variable length_N_blue{"length_N_blue", "const size_t"};
-    Variable length_M_blue{"length_M_blue", "const size_t"};
-    Variable global_stride_in_0{"global_stride_in_0", "const size_t"};
-    Variable global_stride_in_1{"global_stride_in_1", "const size_t"};
-    Variable global_idist{"global_idist", "const size_t"};
-    Variable global_stride_out_0{"global_stride_out_0", "const size_t"};
-    Variable global_stride_out_1{"global_stride_out_1", "const size_t"};
-    Variable global_odist{"global_odist", "const size_t"};
-    Variable transform_idx{"transform_idx", "const size_t"};
+    Variable length_N_blue{"length_N_blue", "const integer_type"};
+    Variable length_M_blue{"length_M_blue", "const integer_type"};
+    Variable global_stride_in_0{"global_stride_in_0", "const integer_type"};
+    Variable global_stride_in_1{"global_stride_in_1", "const integer_type"};
+    Variable global_idist{"global_idist", "const integer_type"};
+    Variable global_stride_out_0{"global_stride_out_0", "const integer_type"};
+    Variable global_stride_out_1{"global_stride_out_1", "const integer_type"};
+    Variable global_odist{"global_odist", "const integer_type"};
+    Variable transform_idx{"transform_idx", "const integer_type"};
 
     //
     // variables borrowed from stockham_gen_base.h
     //
-    Variable global_data_id{"global_data_id", "size_t"};
-    Variable global_transf_id{"global_transf_id", "size_t"};
+    Variable global_data_id{"global_data_id", "integer_type"};
+    Variable global_transf_id{"global_transf_id", "integer_type"};
     Variable load_cb_data{"load_cb_data", "void*"};
     Variable load_cb_fn{"load_cb_fn", "void", true, true};
     Variable store_cb_data{"store_cb_data", "void*"};
@@ -602,7 +604,7 @@ class BluesteinFunction
 
     std::string render_template()
     {
-        return "<" + data.scalar_type.render() + ">";
+        return "<" + data.scalar_type.render() + "," + data.integer_type.render() + ">";
     }
 
     const std::vector<std::string> function_name = {"bluestein_load_cc_fwd_chirp_device",
@@ -716,6 +718,7 @@ class BluesteinKernel
     {
         TemplateList tpls;
         tpls.append(blueData.scalar_type);
+        tpls.append(blueData.integer_type);
 
         return tpls;
     }

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

@@ -30,6 +30,7 @@
 
 #pragma once
 #include "../../../../shared/arithmetic.h"
+#include "../../include/rtc_kernel.h"
 #include "generator.h"
 
 #include <cmath>
@@ -244,10 +245,14 @@ struct FFTComputeOffsets
 
     std::shared_ptr<Context> context;
 
-    Variable transform{"transform", "size_t"};
-    Variable remaining{"remaining", "size_t"};
-    Variable index_along_d{"index_along_d", "size_t"};
-    Variable d{"d", "int"};
+    // transform is blockIdx.x scaled by transforms_per_block, so it
+    // outgrows 32 bits while the grid itself is still legal; remaining
+    // and index_along_d are decomposed from it against lengths.  d only
+    // counts dimensions, so it stays narrow.
+    Variable transform{"transform", "integer_type"};
+    Variable remaining{"remaining", "integer_type"};
+    Variable index_along_d{"index_along_d", "integer_type"};
+    Variable d{"d", rtc_kint_type(KIntType::U32)};
 
     FFTComputeOffsets() = delete;
     FFTComputeOffsets(unsigned int             length0,
@@ -1059,18 +1064,20 @@ struct StockhamTransform
     std::shared_ptr<Context> context;
 
     FFTBuffer R{"R", Literal{0}, Literal{1}, 0};
-    FFTBuffer lds{"lds", Variable{"offset_lds", "int"}, Variable{"stride_lds", "int"}};
-    FFTBuffer X{"X", Variable{"offset_lds", "size_t"}, Literal{1}};
+    FFTBuffer lds{"lds",
+                  Variable{"offset_lds", rtc_kint_type(KIntType::U32)},
+                  Variable{"stride_lds", rtc_kint_type(KIntType::U32)}};
+    FFTBuffer X{"X", Variable{"offset_lds", rtc_kint_type(KIntType::U32)}, Literal{1}};
 
-    Variable dim{"dim", "unsigned int"};
-    Variable nbatch{"nbatch", "size_t"};
-    Variable lengths{"lengths", "size_t", true, true};
-    Variable stride{"stride", "size_t", true, true};
-    Variable offset{"offset", "size_t"};
+    Variable dim{"dim", rtc_kint_type(KIntType::U32)};
+    Variable nbatch{"nbatch", "integer_type"};
+    Variable lengths{"lengths", "integer_type", true, true};
+    Variable stride{"stride", "integer_type", true, true};
+    Variable offset{"offset", "integer_type"};
 
     Variable write{"write", "bool"};
-    Variable thread{"thread", "size_t"};
-    Variable batch{"batch", "size_t"};
+    Variable thread{"thread", rtc_kint_type(KIntType::U32)};
+    Variable batch{"batch", "integer_type"};
 
     Variable twiddles{"twiddles", "const scalar_type", true, true};
     Variable load_cb_fn{"load_cb_fn", "void*"};
@@ -1152,23 +1159,25 @@ struct BluesteinTransform
     // FFT registers
     FFTBuffer R{"R", Literal{0}, Literal{1}, 0};
     // LDS buffer
-    FFTBuffer A{"A", Variable{"offset_lds", "int"}, Variable{"stride_lds", "int"}};
+    FFTBuffer A{"A",
+                Variable{"offset_lds", rtc_kint_type(KIntType::U32)},
+                Variable{"stride_lds", rtc_kint_type(KIntType::U32)}};
     // FFTed chirp signal (second half of chirp buffer)
     FFTBuffer B{"B", Literal{0}, Literal{1}};
     // chirp signal (first half of chirp buffer)
     FFTBuffer a{"a", Literal{0}, Literal{1}};
     // user data
-    FFTBuffer X{"X", Variable{"offset", "size_t"}, Variable{"stride0", "size_t"}};
+    FFTBuffer X{"X", Variable{"offset", "integer_type"}, Variable{"stride0", "integer_type"}};
 
-    Variable dim{"dim", "unsigned int"};
-    Variable nbatch{"nbatch", "size_t"};
-    Variable lengths{"lengths", "size_t", true, true};
-    Variable stride{"stride", "size_t", true, true};
-    Variable offset{"offset", "size_t"};
+    Variable dim{"dim", rtc_kint_type(KIntType::U32)};
+    Variable nbatch{"nbatch", "integer_type"};
+    Variable lengths{"lengths", "integer_type", true, true};
+    Variable stride{"stride", "integer_type", true, true};
+    Variable offset{"offset", "integer_type"};
 
     Variable write{"write", "bool"};
-    Variable thread{"thread", "size_t"};
-    Variable batch{"batch", "size_t"};
+    Variable thread{"thread", rtc_kint_type(KIntType::U32)};
+    Variable batch{"batch", "integer_type"};
     Variable val{"val", "scalar_type"};
 
     Variable buf_temp{"buf_temp", "scalar_type", true, true};

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

@@ -854,14 +854,26 @@ 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, {}, precision, gcn_arch_name, workgroup_size[0], scheme, transform_type);
+            StockhamGeneratorSpecs specs1(KIntType::U32,
+                                          factors1,
+                                          {},
+                                          precision,
+                                          gcn_arch_name,
+                                          workgroup_size[0],
+                                          scheme,
+                                          transform_type);
             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, {}, precision, gcn_arch_name, workgroup_size[1], scheme, transform_type);
+            StockhamGeneratorSpecs specs2(KIntType::U32,
+                                          factors2,
+                                          {},
+                                          precision,
+                                          gcn_arch_name,
+                                          workgroup_size[1],
+                                          scheme,
+                                          transform_type);
             specs2.direct_to_from_reg    = direct_to_from_reg[1];
             specs2.threads_per_transform = threads_per_transform[1];
             specs2.wgs_is_derived        = true;
@@ -900,7 +912,8 @@ int main()
             ++arg;
             factors = parse_uints_csv(*arg);
 
-            StockhamGeneratorSpecs specs(factors,
+            StockhamGeneratorSpecs specs(KIntType::U32,
+                                         factors,
                                          factors2d,
                                          precision,
                                          gcn_arch_name,
@@ -915,7 +928,8 @@ int main()
             specs.threads_per_transform = threads_per_transform.front();
 
             // second dimension for 2D_SINGLE
-            StockhamGeneratorSpecs specs2d(factors2d,
+            StockhamGeneratorSpecs specs2d(KIntType::U32,
+                                           factors2d,
                                            factors,
                                            precision,
                                            gcn_arch_name,

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

@@ -22,6 +22,7 @@
 
 #pragma once
 #include "../../../../shared/arithmetic.h"
+#include "../../include/rtc_kernel.h"
 #include "../kernels/device_enum.h"
 #include "rocfft/rocfft.h"
 #include <optional>
@@ -31,14 +32,16 @@
 
 struct StockhamGeneratorSpecs
 {
-    StockhamGeneratorSpecs(const std::vector<unsigned int>&   factors,
+    StockhamGeneratorSpecs(const KIntType&                    itype,
+                           const std::vector<unsigned int>&   factors,
                            const std::vector<unsigned int>&   factors2d,
                            unsigned int                       precision,
                            const std::string&                 gcn_arch_name,
                            unsigned int                       workgroup_size,
                            const std::string&                 scheme,
                            const std::optional<unsigned int>& transform_type = std::nullopt)
-        : factors(factors)
+        : itype(itype)
+        , factors(factors)
         , factors2d(factors2d)
         , precision(precision)
         , gcn_arch_name(gcn_arch_name)
@@ -50,6 +53,7 @@ struct StockhamGeneratorSpecs
     {
     }
 
+    KIntType                  itype;
     std::vector<unsigned int> factors;
     std::vector<unsigned int> factors2d;
     std::vector<unsigned int> factors_pp;

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

@@ -238,12 +238,12 @@ struct StockhamKernelFused2D : public StockhamKernelRR
                                  + " threads per 2d transform",
                              "  (1 2d slab per thread block)"};
 
-        Variable d{"d", "int"};
-        Variable index_along_d{"index_along_d", "size_t"};
-        Variable remaining{"remaining", "size_t"};
-        Variable plength{"plength", "size_t"};
+        Variable d{"d", rtc_kint_type(KIntType::U32)};
+        Variable index_along_d{"index_along_d", "integer_type"};
+        Variable remaining{"remaining", "integer_type"};
+        Variable plength{"plength", "integer_type"};
 
-        Variable batch0{"batch0", "size_t"};
+        Variable batch0{"batch0", "integer_type"};
 
         Variable SB_1ST{"SB_1ST", "const StrideBin"};
         Variable SB_2ND{"SB_2ND", "const StrideBin"};