Add support for 64-bit indexing on transpose kernels

projects/rocfft/CHANGELOG.md

@@ -8,6 +8,7 @@ Documentation for rocFFT is available at
 ### Added
 
 * Added amdgcnspirv architecture to client programs, so that they are functional even on gfx architectures that they have not been explicitly compiled in.
+* Added support for very large FFTs on gfx1250.
 
 ## rocFFT 1.0.39 for ROCm 10.0
 

projects/rocfft/library/src/CMakeLists.txt

@@ -324,6 +324,7 @@ add_library( rocfft-rtc-launch OBJECT
   rtc_chirp_kernel.cpp
   load_store_ops_kernel.cpp
   tree_node_callback.cpp
+  tree_node_rtc.cpp
 )
 target_link_libraries( rocfft-rtc-launch PRIVATE rocfft-rtc-cache )
 

projects/rocfft/library/src/include/rtc_kernel.h

@@ -28,6 +28,7 @@
 #include <future>
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -41,11 +42,46 @@ class TreeNode;
 class LeafNode;
 struct GridParam;
 
+// Width of the integer type used for index/offset arithmetic inside
+// generated kernels. Kernels declare such arguments as "index_type".
+enum class IndexType
+{
+    _32BIT,
+    _64BIT,
+};
+
 // Helper class that handles alignment of kernel arguments
 class RTCKernelArgs
 {
 public:
     RTCKernelArgs() = default;
+    explicit RTCKernelArgs(IndexType itype)
+        : itype(itype)
+    {
+    }
+    // append a value for an argument declared as "index_type"
+    void append_index(size_t value, std::optional<IndexType> itype = std::nullopt)
+    {
+        itype = itype.has_value() ? itype.value() : this->itype;
+
+        switch(itype.value())
+        {
+        case IndexType::_32BIT:
+        {
+            if(value > std::numeric_limits<unsigned int>::max())
+                throw std::runtime_error("index value overflows 32-bit kernel index_type");
+            unsigned int v = static_cast<unsigned int>(value);
+            append(&v, sizeof(v));
+            break;
+        }
+        case IndexType::_64BIT:
+        {
+            unsigned long long v = value;
+            append(&v, sizeof(v));
+            break;
+        }
+        }
+    }
     void append_ptr(const void* ptr)
     {
         append(&ptr, sizeof(void*));
@@ -104,6 +140,7 @@ class RTCKernelArgs
     }
 
     std::vector<char> buf;
+    IndexType         itype = IndexType::_32BIT;
 };
 
 // Base class for a runtime compiled kernel.  Subclassed for
@@ -281,6 +318,45 @@ static const char* rtc_array_type_name(rocfft_array_type type)
     }
 }
 
+static const char* rtc_index_name(IndexType itype)
+{
+    switch(itype)
+    {
+    case IndexType::_32BIT:
+        return "_i32";
+    case IndexType::_64BIT:
+        return "_i64";
+    }
+
+    throw std::runtime_error("Invalid index type");
+}
+
+static const char* rtc_index_type(IndexType itype)
+{
+    switch(itype)
+    {
+    case IndexType::_32BIT:
+        return "unsigned int";
+    case IndexType::_64BIT:
+        return "unsigned long long";
+    }
+
+    throw std::runtime_error("Invalid index type");
+}
+
+static const char* rtc_index_type_decl(IndexType itype)
+{
+    switch(itype)
+    {
+    case IndexType::_32BIT:
+        return "typedef unsigned int index_type;\n";
+    case IndexType::_64BIT:
+        return "typedef unsigned long long index_type;\n";
+    }
+
+    throw std::runtime_error("Invalid index type");
+}
+
 static const char* rtc_precision_name(rocfft_precision precision)
 {
     switch(precision)

projects/rocfft/library/src/include/rtc_transpose_gen.h

@@ -29,6 +29,7 @@
 
 struct TransposeSpecs
 {
+    IndexType               itype;
     unsigned int            tileX;
     unsigned int            tileY;
     size_t                  dim;

projects/rocfft/library/src/include/rtc_transpose_kernel.h

@@ -28,11 +28,15 @@ struct RTCKernelTranspose : public RTCKernel
     RTCKernelTranspose(const std::string&                       kernel_name,
                        std::shared_future<hipModule_wrapper_t>& module,
                        dim3                                     gridDim,
-                       dim3                                     blockDim)
+                       dim3                                     blockDim,
+                       IndexType                                itype)
         : RTCKernel(kernel_name, module, gridDim, blockDim)
+        , itype(itype)
     {
     }
 
+    IndexType itype;
+
     static RTCKernel::RTCGenerator generate_from_node(const LeafNode&    node,
                                                       const std::string& gpu_arch,
                                                       bool               enable_callbacks);

projects/rocfft/library/src/include/tree_node.h

@@ -1039,6 +1039,7 @@ class LeafNode : public InternalNode
     bool         CreateDeviceResources() override;
     void         SetupGridParam(GridParam& gp) override;
     FMKey        GetKernelKey() const override;
+    bool         KernelNeeds64BitIndexing() const;
     virtual void GetKernelFactors();
     virtual void GetKernelPartialPassFactors();
 };

projects/rocfft/library/src/rtc_transpose_gen.cpp

@@ -48,6 +48,7 @@ std::string transpose_rtc_kernel_name(const TransposeSpecs& specs)
         break;
     }
 
+    kernel_name += rtc_index_name(specs.itype);
     kernel_name += rtc_precision_name(specs.precision);
     kernel_name += rtc_array_type_name(specs.inArrayType);
     kernel_name += rtc_array_type_name(specs.outArrayType);
@@ -88,6 +89,11 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs&
     src += device_enum_h;
     src += callback_h;
 
+    // NOTE:
+    // Index variables declared as 32BIT are all bounded by grid limits,
+    // and widening them to 64BIT would cost registers for nothing
+
+    src += rtc_index_type_decl(specs.itype);
     src += rtc_precision_type_decl(specs.precision, array_type_is_complex(specs.inArrayType));
 
     src += rtc_const_cbtype_decl(specs.cbtype);
@@ -99,24 +105,24 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs&
     Variable input_var{"input", "scalar_type", true, true};
     Variable output_var{"output", "scalar_type", true, true};
     Variable twiddles_large_var{"twiddles_large", "const scalar_type", true, true};
-    Variable dim_var{"dim", "unsigned int"};
-    Variable length0_var{"length0", "unsigned int"};
-    Variable length1_var{"length1", "unsigned int"};
-    Variable length2_var{"length2", "unsigned int"};
+    Variable dim_var{"dim", rtc_index_type(IndexType::_32BIT)};
+    Variable length0_var{"length0", rtc_index_type(IndexType::_32BIT)};
+    Variable length1_var{"length1", rtc_index_type(IndexType::_32BIT)};
+    Variable length2_var{"length2", rtc_index_type(IndexType::_32BIT)};
+    Variable gridX{"gridX", "const " + std::string(rtc_index_type(IndexType::_32BIT))};
+    Variable gridY{"gridY", "const " + std::string(rtc_index_type(IndexType::_32BIT))};
+    Variable gridZ{"gridZ", "const " + std::string(rtc_index_type(IndexType::_32BIT))};
     Variable lengths_var{"lengths", "const size_t", true, true};
-    Variable stride_in0_var{"stride_in0", "unsigned int"};
-    Variable stride_in1_var{"stride_in1", "unsigned int"};
-    Variable stride_in2_var{"stride_in2", "unsigned int"};
+    Variable stride_in0_var{"stride_in0", "index_type"};
+    Variable stride_in1_var{"stride_in1", "index_type"};
+    Variable stride_in2_var{"stride_in2", "index_type"};
     Variable stride_in_var{"stride_in", "const size_t", true, true};
-    Variable idist_var{"idist", "unsigned int"};
-    Variable stride_out0_var{"stride_out0", "unsigned int"};
-    Variable stride_out1_var{"stride_out1", "unsigned int"};
-    Variable stride_out2_var{"stride_out2", "unsigned int"};
+    Variable idist_var{"idist", "index_type"};
+    Variable stride_out0_var{"stride_out0", "index_type"};
+    Variable stride_out1_var{"stride_out1", "index_type"};
+    Variable stride_out2_var{"stride_out2", "index_type"};
     Variable stride_out_var{"stride_out", "const size_t", true, true};
-    Variable odist_var{"odist", "unsigned int"};
-    Variable gridX{"gridX", "const unsigned int"};
-    Variable gridY{"gridY", "const unsigned int"};
-    Variable gridZ{"gridZ", "const unsigned int"};
+    Variable odist_var{"odist", "index_type"};
 
     Function func(kernel_name);
     func.launch_bounds = specs.tileX * specs.tileY;
@@ -164,14 +170,12 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs&
     func.body += CommentLines{"since gridDim is passed as {gridX, 1, 1}, use the",
                               "following variables to recover block indices in a 3-D fashion:"};
 
-    Variable old_blockIdx_x{"old_blockIdx_x", "unsigned int"};
-    Variable old_blockIdx_y{"old_blockIdx_y", "unsigned int"};
-    Variable old_blockIdx_z{"old_blockIdx_z", "unsigned int"};
-
-    Variable tileBlockIdx_y{"tileBlockIdx_y", "unsigned int"};
-    Variable tileBlockIdx_x{"tileBlockIdx_x", "unsigned int"};
-
-    Variable remaining{"remaining", "unsigned int"};
+    Variable old_blockIdx_x{"old_blockIdx_x", rtc_index_type(IndexType::_32BIT)};
+    Variable old_blockIdx_y{"old_blockIdx_y", rtc_index_type(IndexType::_32BIT)};
+    Variable old_blockIdx_z{"old_blockIdx_z", rtc_index_type(IndexType::_32BIT)};
+    Variable tileBlockIdx_y{"tileBlockIdx_y", rtc_index_type(IndexType::_32BIT)};
+    Variable tileBlockIdx_x{"tileBlockIdx_x", rtc_index_type(IndexType::_32BIT)};
+    Variable remaining{"remaining", rtc_index_type(IndexType::_32BIT)};
 
     // if a 1-D grid was provided because creating a natural 3-D grid exceeded allowed limits, then remap it to a 3-D grid.
     if(!specs.grid3D)
@@ -221,14 +225,14 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs&
         func.body += Assign{length2_var, 1};
     }
 
-    Variable tile_x_index{"tile_x_index", "unsigned int"};
-    Variable tile_y_index{"tile_y_index", "unsigned int"};
+    Variable tile_x_index{"tile_x_index", rtc_index_type(IndexType::_32BIT)};
+    Variable tile_y_index{"tile_y_index", rtc_index_type(IndexType::_32BIT)};
     func.body += Declaration{tile_x_index, "threadIdx.x"};
     func.body += Declaration{tile_y_index, "threadIdx.y"};
 
     func.body += CommentLines{"work out offset for dimensions after the first 3"};
-    Variable offset_in{"offset_in", "unsigned int"};
-    Variable offset_out{"offset_out", "unsigned int"};
+    Variable offset_in{"offset_in", "index_type"};
+    Variable offset_out{"offset_out", "index_type"};
     if(specs.grid3D)
     {
         func.body += Declaration{remaining, "blockIdx.z"};
@@ -243,7 +247,7 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs&
     // use specified dim to avoid loops if possible
     if(specs.dim > 3)
     {
-        Variable d{"d", "unsigned int"};
+        Variable d{"d", rtc_index_type(IndexType::_32BIT)};
         For      offset_loop{
             d,
             3,
@@ -266,7 +270,7 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs&
     func.body += CallbackStoreDeclaration("scalar_type", "cbtype");
 
     // loop variables for reading/writing
-    Variable i{"i", "unsigned int"};
+    Variable i{"i", rtc_index_type(IndexType::_32BIT)};
     Variable logical_row{"logical_row", "auto"};
     Variable logical_col{"logical_col", "auto"};
     Variable idx0{"idx0", "auto"};

projects/rocfft/library/src/rtc_transpose_kernel.cpp

@@ -90,7 +90,13 @@ RTCKernel::RTCGenerator RTCKernelTranspose::generate_from_node(const LeafNode&
 
     bool tileAligned = node.length[0] % tileX == 0 && node.length[1] % tileX == 0;
 
-    TransposeSpecs specs{tileX,
+    // Determine index type based on whether the kernel needs 64-bit indexing.
+    // This runs after buffer assignment, fusion and padding, so the node's
+    // lengths and strides are final.
+    IndexType itype = node.KernelNeeds64BitIndexing() ? IndexType::_64BIT : IndexType::_32BIT;
+
+    TransposeSpecs specs{itype,
+                         tileX,
                          tileY,
                          node.length.size(),
                          node.precision,
@@ -115,14 +121,14 @@ RTCKernel::RTCGenerator RTCKernelTranspose::generate_from_node(const LeafNode&
                                         dim3                                     gridDim,
                                         dim3                                     blockDim) {
         return std::unique_ptr<RTCKernel>(
-            new RTCKernelTranspose(kernel_name, module, gridDim, blockDim));
+            new RTCKernelTranspose(kernel_name, module, gridDim, blockDim, itype));
     };
     return generator;
 }
 
 RTCKernelArgs RTCKernelTranspose::get_launch_args(DeviceCallIn& data)
 {
-    RTCKernelArgs kargs;
+    RTCKernelArgs kargs{itype};
     kargs.append_ptr(data.bufIn[0]);
     if(array_type_is_planar(data.node->inArrayType))
         kargs.append_ptr(data.bufIn[1]);
@@ -131,24 +137,28 @@ RTCKernelArgs RTCKernelTranspose::get_launch_args(DeviceCallIn& data)
         kargs.append_ptr(data.bufOut[1]);
     kargs.append_ptr(data.node->twiddles_large);
 
+    // NOTE:
+    // Kargs appended as index type 32BIT are all bounded by grid limits,
+    // and widening them to 64BIT would cost registers for nothing
+
     auto num_lengths = data.node->length.size();
-    kargs.append_unsigned_int(num_lengths);
-    kargs.append_unsigned_int(data.node->length[0]);
-    kargs.append_unsigned_int(data.node->length[1]);
-    kargs.append_unsigned_int(num_lengths > 2 ? data.node->length[2] : 1);
+    kargs.append_index(num_lengths, IndexType::_32BIT);
+    kargs.append_index(data.node->length[0], IndexType::_32BIT);
+    kargs.append_index(data.node->length[1], IndexType::_32BIT);
+    kargs.append_index(num_lengths > 2 ? data.node->length[2] : 1, IndexType::_32BIT);
     kargs.append_ptr(kargs_lengths(data.node->devKernArg));
 
-    kargs.append_unsigned_int(data.node->inStride[0]);
-    kargs.append_unsigned_int(data.node->inStride[1]);
-    kargs.append_unsigned_int(num_lengths > 2 ? data.node->inStride[2] : 0);
+    kargs.append_index(data.node->inStride[0]);
+    kargs.append_index(data.node->inStride[1]);
+    kargs.append_index(num_lengths > 2 ? data.node->inStride[2] : 0);
     kargs.append_ptr(kargs_stride_in(data.node->devKernArg));
-    kargs.append_unsigned_int(data.node->iDist);
+    kargs.append_index(data.node->iDist);
 
-    kargs.append_unsigned_int(data.node->outStride[0]);
-    kargs.append_unsigned_int(data.node->outStride[1]);
-    kargs.append_unsigned_int(num_lengths > 2 ? data.node->outStride[2] : 0);
+    kargs.append_index(data.node->outStride[0]);
+    kargs.append_index(data.node->outStride[1]);
+    kargs.append_index(num_lengths > 2 ? data.node->outStride[2] : 0);
     kargs.append_ptr(kargs_stride_out(data.node->devKernArg));
-    kargs.append_unsigned_int(data.node->oDist);
+    kargs.append_index(data.node->oDist);
 
     // pass gridX, gridY and gridZ to restore a 3-D GPU grid, if needed for large grids
     unsigned int tileX = data.node->precision == rocfft_precision_single ? 64 : 32;
@@ -163,9 +173,9 @@ RTCKernelArgs RTCKernelTranspose::get_launch_args(DeviceCallIn& data)
                                          data.node->batch,
                                          std::multiplies<unsigned int>());
 
-    kargs.append_unsigned_int(gridX);
-    kargs.append_unsigned_int(gridY);
-    kargs.append_unsigned_int(gridZ);
+    kargs.append_index(gridX, IndexType::_32BIT);
+    kargs.append_index(gridY, IndexType::_32BIT);
+    kargs.append_index(gridZ, IndexType::_32BIT);
 
     // callback params
     kargs.append_ptr(data.callbacks.load_cb_fn);