diff --git a/projects/rocfft/CHANGELOG.md b/projects/rocfft/CHANGELOG.md index 1528dd427b0d..83e96527f79d 100644 --- a/projects/rocfft/CHANGELOG.md +++ b/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 diff --git a/projects/rocfft/library/src/CMakeLists.txt b/projects/rocfft/library/src/CMakeLists.txt index 66267e9ebd77..a9bc13234b09 100644 --- a/projects/rocfft/library/src/CMakeLists.txt +++ b/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 ) diff --git a/projects/rocfft/library/src/include/rtc_kernel.h b/projects/rocfft/library/src/include/rtc_kernel.h index b28924bbda8e..b8cc4176d67a 100644 --- a/projects/rocfft/library/src/include/rtc_kernel.h +++ b/projects/rocfft/library/src/include/rtc_kernel.h @@ -28,6 +28,7 @@ #include #include #include +#include #include #include @@ -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 itype = std::nullopt) + { + itype = itype.has_value() ? itype.value() : this->itype; + + switch(itype.value()) + { + case IndexType::_32BIT: + { + if(value > std::numeric_limits::max()) + throw std::runtime_error("index value overflows 32-bit kernel index_type"); + unsigned int v = static_cast(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 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) diff --git a/projects/rocfft/library/src/include/rtc_transpose_gen.h b/projects/rocfft/library/src/include/rtc_transpose_gen.h index 3f0a4be522e6..e3d8da75b881 100644 --- a/projects/rocfft/library/src/include/rtc_transpose_gen.h +++ b/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; diff --git a/projects/rocfft/library/src/include/rtc_transpose_kernel.h b/projects/rocfft/library/src/include/rtc_transpose_kernel.h index 1d36d5ef63db..b7f61e1bf844 100644 --- a/projects/rocfft/library/src/include/rtc_transpose_kernel.h +++ b/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& 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); diff --git a/projects/rocfft/library/src/include/tree_node.h b/projects/rocfft/library/src/include/tree_node.h index a01f1b11eb3d..458c7c92610a 100644 --- a/projects/rocfft/library/src/include/tree_node.h +++ b/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(); }; diff --git a/projects/rocfft/library/src/rtc_transpose_gen.cpp b/projects/rocfft/library/src/rtc_transpose_gen.cpp index fa23dde7dfc5..db70bc950a39 100644 --- a/projects/rocfft/library/src/rtc_transpose_gen.cpp +++ b/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"}; diff --git a/projects/rocfft/library/src/rtc_transpose_kernel.cpp b/projects/rocfft/library/src/rtc_transpose_kernel.cpp index c50477f257db..09b8f520b771 100644 --- a/projects/rocfft/library/src/rtc_transpose_kernel.cpp +++ b/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( - 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()); - 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); diff --git a/projects/rocfft/library/src/tree_node_rtc.cpp b/projects/rocfft/library/src/tree_node_rtc.cpp new file mode 100644 index 000000000000..5e49ee83a210 --- /dev/null +++ b/projects/rocfft/library/src/tree_node_rtc.cpp @@ -0,0 +1,51 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +// THE SOFTWARE. + +// TreeNode methods that kernel generation needs. These live apart +// from tree_node.cpp because the RTC helper executables link +// rocfft-rtc-launch without the rest of the rocFFT library. + +#include "tree_node.h" + +#include "../../shared/ptrdiff.h" + +#include + +bool LeafNode::KernelNeeds64BitIndexing() const +{ + for(auto io : {io_data_label::INPUT, io_data_label::OUTPUT}) + { + const auto& io_stride = io == io_data_label::INPUT ? inStride : outStride; + const auto& io_dist = io == io_data_label::INPUT ? iDist : oDist; + const auto& io_array_type = io == io_data_label::INPUT ? inArrayType : outArrayType; + const auto& io_offset = io == io_data_label::INPUT ? iOffset : oOffset; + const auto io_length = io == io_data_label::INPUT ? length : GetOutputLength(); + + // Hermitian interleaved data may be re-interpreted as real data internally. + if((io_offset + compute_ptrdiff(io_length, io_stride, batch, io_dist)) + * (io_array_type == rocfft_array_type_hermitian_interleaved ? 2 : 1) + > static_cast(INT32_MAX) + 1) + { + return true; + } + } + + return false; +} diff --git a/projects/rocfft/shared/accuracy_test.h b/projects/rocfft/shared/accuracy_test.h index 467cc9d33d94..fb29c0822d5e 100644 --- a/projects/rocfft/shared/accuracy_test.h +++ b/projects/rocfft/shared/accuracy_test.h @@ -524,11 +524,6 @@ inline void fft_vs_reference_impl(Tparams& params, bool round_trip) // Make sure that the parameters make sense: ASSERT_TRUE(params.valid(verbose)); - // TODO: temporary workaround awaiting robust support for - // 64-bit indexing in rocfft kernels. - if(params.may_need_64bit_indexing()) - throw ROCFFT_SKIP{"This test may require kernel support for 64-bit integer arithmetic"}; - // Create reference results as early as possible so that system memory is reserved for (an // estimation of) the possible FFTW plan's workspace (if needed), providing some guard against // OOM kills thereafter. diff --git a/projects/rocfft/shared/fft_params.h b/projects/rocfft/shared/fft_params.h index ed46b39790c5..5414acf2f03d 100644 --- a/projects/rocfft/shared/fft_params.h +++ b/projects/rocfft/shared/fft_params.h @@ -1889,30 +1889,6 @@ class fft_params return true; } - // TODO: temporary workaround awaiting robust support for - // 64-bit indexing in rocfft kernels. - bool may_need_64bit_indexing() const - { - for(auto io : {fft_io::fft_io_in, fft_io::fft_io_out}) - { - const auto& io_stride = io == fft_io::fft_io_in ? istride : ostride; - const auto& io_dist = io == fft_io::fft_io_in ? idist : odist; - const auto& io_array_type = io == fft_io::fft_io_in ? itype : otype; - const auto& io_offset = io == fft_io::fft_io_in ? ioffset : ooffset; - const auto io_length = io == fft_io::fft_io_in ? ilength() : olength(); - const auto max_offset - = io_offset.empty() ? 0 : *std::max_element(io_offset.begin(), io_offset.end()); - // Hermitian interleaved data may be re-interpreted as real data internally. - if((max_offset + compute_ptrdiff(io_length, io_stride, nbatch, io_dist)) - * (io_array_type == fft_array_type_hermitian_interleaved ? 2 : 1) - > static_cast(UINT32_MAX) + 1) - { - return true; - } - } - return false; - } - // Fill in any missing parameters. void validate() {