diff --git a/projects/rocfft/CHANGELOG.md b/projects/rocfft/CHANGELOG.md index 067a40d94484..b52eae0d7ffe 100644 --- a/projects/rocfft/CHANGELOG.md +++ b/projects/rocfft/CHANGELOG.md @@ -13,6 +13,10 @@ Documentation for rocFFT is available at * Fixed a memory leak when `rocfft_plan_create` fails. * Fixed a potential write-after-free if plans are destroyed during process teardown. +### Added + +* Added further support for very large FFTs (length greater than 2^32). + ## rocFFT 1.0.40 for ROCm 10.1 ### Added diff --git a/projects/rocfft/clients/tests/accuracy_test_adhoc.cpp b/projects/rocfft/clients/tests/accuracy_test_adhoc.cpp index 36acefc0f9f8..b86f6bf7b2dc 100644 --- a/projects/rocfft/clients/tests/accuracy_test_adhoc.cpp +++ b/projects/rocfft/clients/tests/accuracy_test_adhoc.cpp @@ -666,6 +666,58 @@ INSTANTIATE_TEST_SUITE_P(adhoc_kint_index_boundary, adhoc_kint_index_tokens)), accuracy_test::TestName); +// Test cases for 64-bit index kernels that require large twiddle arrays. +const auto adhoc_64_bit_idx_twiddle_large_tokens = { + // clang-format off + // CS_L1D_TRTRT twl 5 + "complex_forward_len_8589934592_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_8589934592_odist_8589934592_ioffset_0_0_ooffset_0_0", // ~128 GiB VRAM + // Bluestein PAD_MUL / FFT_MUL count boundary. count = lengthBlue * batch, + // lengthBlue = 524288, so batch 8192 => count = 2^32 (last U32 case) and + // batch 8193 => count = 4295491584 (first U64 case). + "complex_forward_len_196597_single_ip_batch_8192_istride_1_CI_ostride_1_CI_idist_196597_odist_196597_ioffset_0_0_ooffset_0_0", // ~76 GiB VRAM + "complex_forward_len_196597_single_ip_batch_8193_istride_1_CI_ostride_1_CI_idist_196597_odist_196597_ioffset_0_0_ooffset_0_0", // ~76 GiB VRAM + // Bluestein chirp large-twiddle step count. chirp large1D = 2 * length, + // so the twl tier boundaries land at length 128, 32768, 8388608, 2^31. + // twl 1 + "complex_forward_len_19_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_19_odist_19_ioffset_0_0_ooffset_0_0", + "complex_inverse_len_19_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_19_odist_19_ioffset_0_0_ooffset_0_0", + // twl 2 + "complex_forward_len_2053_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_2053_odist_2053_ioffset_0_0_ooffset_0_0", + // twl 1 + twl 2, both forced non-fused by the 2D parent + "complex_forward_len_19_2053_single_ip_batch_1_istride_2053_1_CI_ostride_2053_1_CI_idist_39007_odist_39007_ioffset_0_0_ooffset_0_0", + // twl 3 + "complex_forward_len_196597_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_196597_odist_196597_ioffset_0_0_ooffset_0_0", + // twl 3 / twl 4 boundary: large1D 16777186 vs 16777234, either side of 2^24 + "complex_forward_len_8388593_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_8388593_odist_8388593_ioffset_0_0_ooffset_0_0", + "complex_forward_len_8388617_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_8388617_odist_8388617_ioffset_0_0_ooffset_0_0", + // twl 4 + "complex_forward_len_25165813_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_25165813_odist_25165813_ioffset_0_0_ooffset_0_0", + // twl 4 / twl 5 boundary: large1D 4294967294 left side of 2^32 + "complex_forward_len_2147483647_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_2147483647_odist_2147483647_ioffset_0_0_ooffset_0_0", // ~176GiB VRAM + // clang-format on +}; +INSTANTIATE_TEST_SUITE_P( + adhoc_64_bit_idx_twiddle_large, + accuracy_test, + ::testing::ValuesIn(param_generator_token(test_prob, adhoc_64_bit_idx_twiddle_large_tokens)), + accuracy_test::TestName); + +// ~320 GiB. The host-side reference is set up before the VRAM check can +// skip it, so this costs real time even on machines that can't run it. +// Sample it instead of running it every time. +const auto adhoc_64_bit_idx_twiddle_large_slow_tokens = { + // clang-format off + // twl 5 boundary: large1D 4294967294 right side of 2^32 + "complex_forward_len_2147483659_single_ip_batch_1_istride_1_CI_ostride_1_CI_idist_2147483659_odist_2147483659_ioffset_0_0_ooffset_0_0", + // clang-format on +}; +INSTANTIATE_TEST_SUITE_P( + adhoc_64_bit_idx_twiddle_large_slow, + accuracy_test, + ::testing::ValuesIn(param_generator_token(test_prob* very_large_prob_factor, + adhoc_64_bit_idx_twiddle_large_slow_tokens)), + accuracy_test::TestName); + inline auto param_even_real_odd_base_index() { std::vector params; diff --git a/projects/rocfft/clients/tests/gtest_main.cpp b/projects/rocfft/clients/tests/gtest_main.cpp index e5e2bd116bcf..0a150bc29ec5 100644 --- a/projects/rocfft/clients/tests/gtest_main.cpp +++ b/projects/rocfft/clients/tests/gtest_main.cpp @@ -71,7 +71,8 @@ double real_prob_factor; double complex_planar_prob_factor; // Modifier for probability of running tests with callbacks double callback_prob_factor; - +// Modifier for probability of running very large tests +double very_large_prob_factor; // Number of random tests per suite size_t n_random_tests = 0; @@ -384,6 +385,11 @@ int main(int argc, char* argv[]) "Probability multiplier for running individual callback transforms") ->default_val(0.2) ->check(CLI::Range(0.0, 1.0)); + app.add_option("--very_large_prob", + very_large_prob_factor, + "Probability multiplier for running very large transforms") + ->default_val(0.2) + ->check(CLI::Range(0.0, 1.0)); constexpr auto emulation_quick = "quick"; constexpr auto emulation_smoke = "smoke"; @@ -399,6 +405,7 @@ int main(int argc, char* argv[]) "--unittest_prob", "--nrand", "--callback_prob", + "--very_large_prob", "--R") ->each([&](const std::string& emulationtype) { // Emulation test suites focus on well-established software paths; we are looking for @@ -412,6 +419,9 @@ int main(int argc, char* argv[]) // Callbacks are not an emulation test target. callback_prob_factor = 0; + // Very large transforms are not an emulation test target. + very_large_prob_factor = 0; + if(emulationtype == emulation_quick) { // Configuration specific for "quick simulation test" category, the whole test run diff --git a/projects/rocfft/library/src/device/kernels/common.h b/projects/rocfft/library/src/device/kernels/common.h index 8a83d9f481e7..98abdde35b24 100644 --- a/projects/rocfft/library/src/device/kernels/common.h +++ b/projects/rocfft/library/src/device/kernels/common.h @@ -155,59 +155,52 @@ using complex_type_t = typename complex_type::type; // complex_type_t double_complex_val; template -__device__ T TWLstep1(const T* twiddles, size_t u) +__device__ T TWLstep(const T* twiddles, size_t u, unsigned int step) { size_t j = u & 255; T result = twiddles[j]; + +#pragma unroll + for(unsigned int i = 0; i < step - 1; i++) + { + u >>= 8; + j = u & 255; + result = T( + (result.x * twiddles[256 * (i + 1) + j].x - result.y * twiddles[256 * (i + 1) + j].y), + (result.y * twiddles[256 * (i + 1) + j].x + result.x * twiddles[256 * (i + 1) + j].y)); + } + return result; } +template +__device__ T TWLstep1(const T* twiddles, size_t u) +{ + return TWLstep(twiddles, u, 1); +} + template __device__ T TWLstep2(const T* twiddles, size_t u) { - size_t j = u & 255; - T result = twiddles[j]; - u >>= 8; - j = u & 255; - result = T((result.x * twiddles[256 + j].x - result.y * twiddles[256 + j].y), - (result.y * twiddles[256 + j].x + result.x * twiddles[256 + j].y)); - return result; + return TWLstep(twiddles, u, 2); } template __device__ T TWLstep3(const T* twiddles, size_t u) { - size_t j = u & 255; - T result = twiddles[j]; - u >>= 8; - j = u & 255; - result = T((result.x * twiddles[256 + j].x - result.y * twiddles[256 + j].y), - (result.y * twiddles[256 + j].x + result.x * twiddles[256 + j].y)); - u >>= 8; - j = u & 255; - result = T((result.x * twiddles[512 + j].x - result.y * twiddles[512 + j].y), - (result.y * twiddles[512 + j].x + result.x * twiddles[512 + j].y)); - return result; + return TWLstep(twiddles, u, 3); } template __device__ T TWLstep4(const T* twiddles, size_t u) { - size_t j = u & 255; - T result = twiddles[j]; - u >>= 8; - j = u & 255; - result = T((result.x * twiddles[256 + j].x - result.y * twiddles[256 + j].y), - (result.y * twiddles[256 + j].x + result.x * twiddles[256 + j].y)); - u >>= 8; - j = u & 255; - result = T((result.x * twiddles[512 + j].x - result.y * twiddles[512 + j].y), - (result.y * twiddles[512 + j].x + result.x * twiddles[512 + j].y)); - u >>= 8; - j = u & 255; - result = T((result.x * twiddles[768 + j].x - result.y * twiddles[768 + j].y), - (result.y * twiddles[768 + j].x + result.x * twiddles[768 + j].y)); - return result; + return TWLstep(twiddles, u, 4); +} + +template +__device__ T TWLstep5(const T* twiddles, size_t u) +{ + return TWLstep(twiddles, u, 5); } #define TWIDDLE_STEP_MUL_FWD(TWFUNC, TWIDDLES, INDEX, REG) \ diff --git a/projects/rocfft/library/src/include/rtc_bluestein_gen.h b/projects/rocfft/library/src/include/rtc_bluestein_gen.h index 99b7ba728774..3ca030c2224b 100644 --- a/projects/rocfft/library/src/include/rtc_bluestein_gen.h +++ b/projects/rocfft/library/src/include/rtc_bluestein_gen.h @@ -62,6 +62,10 @@ struct BluesteinMultiSpecs CallbackType cbtype; std::optional loadOps; std::optional storeOps; + // Whether the launch grid spills into Y, which decides how the + // kernel derives its thread index. Part of the kernel name, since + // the two forms compile to different code. + bool splitGrid; }; std::string bluestein_multi_rtc_kernel_name(const BluesteinMultiSpecs& specs); diff --git a/projects/rocfft/library/src/include/rtc_kernel.h b/projects/rocfft/library/src/include/rtc_kernel.h index ec66001bc1a9..28f5a7e82678 100644 --- a/projects/rocfft/library/src/include/rtc_kernel.h +++ b/projects/rocfft/library/src/include/rtc_kernel.h @@ -42,7 +42,6 @@ struct DeviceCallIn; class TreeNode; class LeafNode; struct GridParam; - // Helper class that handles alignment of kernel arguments class RTCKernelArgs { diff --git a/projects/rocfft/library/src/rtc_bluestein_gen.cpp b/projects/rocfft/library/src/rtc_bluestein_gen.cpp index 2860c554d9d9..ec15da7b4316 100644 --- a/projects/rocfft/library/src/rtc_bluestein_gen.cpp +++ b/projects/rocfft/library/src/rtc_bluestein_gen.cpp @@ -62,6 +62,35 @@ std::string bluestein_single_rtc_kernel_name(const BluesteinSingleSpecs& specs) return kernel_name; } +// Type wide enough to hold the full tx * tx product for any tx that +// integer_type can express. +static std::string chirp_rtc_wide_type_decl(const KIntType& itype) +{ + return itype == KIntType::U64 ? "typedef __uint128_t wide_type;\n" + : "typedef unsigned long long wide_type;\n"; +} + +// Thread index for the one-thread-per-element grid laid out by +// flat_grid_dim. +// +// A grid that fits in X alone indexes in integer_type, which keeps the +// index provably 32-bit for the common case. That range matters well +// past the index itself: tx feeds the division and modulus by numof, +// which cost noticeably more once the dividend can be any 64-bit value. +// +// Once the grid spills into Y the index no longer fits, so it is +// computed in size_t. integer_type would not do even when it is 64-bit +// wide: it is derived from the data extents, which say nothing about the +// Bluestein work item count that sizes the grid. +static const char* flat_grid_thread_index(bool splitGrid) +{ + if(!splitGrid) + return "threadIdx.x + static_cast(blockIdx.x) * blockDim.x"; + + return "threadIdx.x + static_cast(blockIdx.x) * blockDim.x" + " + static_cast(blockIdx.y) * gridDim.x * blockDim.x"; +} + std::string bluestein_single_rtc(const std::string& kernel_name, const BluesteinSingleSpecs& specs) { auto length = specs.length; @@ -164,6 +193,9 @@ std::string bluestein_multi_rtc_kernel_name(const BluesteinMultiSpecs& specs) kernel_name += rtc_precision_name(specs.precision); kernel_name += rtc_array_type_name(specs.inArrayType); kernel_name += rtc_array_type_name(specs.outArrayType); + // only suffixed when set, to leave the usual kernel names alone + if(specs.splitGrid) + kernel_name += "_gridxy"; kernel_name += load_store_name_suffix(specs.loadOps, specs.storeOps); kernel_name += rtc_cbtype_name(specs.cbtype); @@ -191,10 +223,10 @@ static std::string bluestein_multi_chirp_rtc(const std::string& kernel_n func.arguments.append(twl); func.arguments.append(dir); - Variable tx{"tx", "size_t"}; + Variable tx{"tx", "wide_type"}; Variable val{"val", "scalar_type"}; - func.body += Declaration{tx, "threadIdx.x + blockIdx.x * blockDim.x"}; + func.body += Declaration{tx, Literal{flat_grid_thread_index(specs.splitGrid)}}; func.body += Declaration{val, CallExpr{"scalar_type", {Literal{"0.0"}, Literal{"0.0"}}}}; func.body @@ -205,6 +237,8 @@ static std::string bluestein_multi_chirp_rtc(const std::string& kernel_n {Assign{val, CallExpr{"TWLstep3", {twiddles_large, (tx * tx) % (2 * N)}}}}}; func.body += ElseIf{twl == 4, {Assign{val, CallExpr{"TWLstep4", {twiddles_large, (tx * tx) % (2 * N)}}}}}; + func.body += ElseIf{twl == 5, + {Assign{val, CallExpr{"TWLstep5", {twiddles_large, (tx * tx) % (2 * N)}}}}}; func.body += MultiplyAssign(val.y(), CallExpr{"real_type_t", {dir}}); @@ -239,6 +273,7 @@ std::string bluestein_multi_rtc(const std::string& kernel_name, const BluesteinM src += common_h; src += device_enum_h; src += rtc_precision_type_decl(specs.precision); + src += chirp_rtc_wide_type_decl(specs.itype); src += rtc_kint_type_decl(specs.itype); src += load_store_decls(specs.loadOps, specs.storeOps, specs.cbtype); src += callback_h; @@ -293,7 +328,7 @@ std::string bluestein_multi_rtc(const std::string& kernel_name, const BluesteinM Variable chirp{"chirp", "scalar_type", true}; Variable out_elem{"out_elem", "scalar_type"}; - func.body += Declaration{tx, "threadIdx.x + blockIdx.x * blockDim.x"}; + func.body += Declaration{tx, Literal{flat_grid_thread_index(specs.splitGrid)}}; func.body += If{tx >= totalWI, {Return{}}}; diff --git a/projects/rocfft/library/src/rtc_bluestein_kernel.cpp b/projects/rocfft/library/src/rtc_bluestein_kernel.cpp index d5f0133da819..46e9e7c46970 100644 --- a/projects/rocfft/library/src/rtc_bluestein_kernel.cpp +++ b/projects/rocfft/library/src/rtc_bluestein_kernel.cpp @@ -157,16 +157,46 @@ RTCKernel::RTCGenerator RTCKernelBluesteinMulti::generate_from_node(const LeafNo count *= node.length[i]; count *= numof; + // Lay out a flat work item count as a grid of blocks, for kernels that + // assign one thread per element. + // + // Total work items along a grid dimension is counted in uint32_t in the + // dispatch packet, so one dimension cannot dispatch more than that many + // work items. The block count stays well inside the grid size limits + // long before that point, so exceeding it truncates the dispatch instead + // of failing the launch. Spill into Y once X alone can no longer cover + // the range. + // + // Kernels using this must derive their thread index as + // threadIdx.x + blockIdx.x * blockDim.x + blockIdx.y * gridDim.x * blockDim.x + // computed in 64 bits, and must tolerate a grid that rounds up past the + // work item count. + auto flat_grid_dim = [](size_t work_items, unsigned int block_size) -> dim3 { + // whole blocks a single dimension can dispatch + const size_t max_blocks = std::numeric_limits::max() / block_size; + + auto num_blocks = DivRoundingUp(work_items, block_size); + + if(num_blocks <= max_blocks) + return {static_cast(num_blocks), 1, 1}; + + // fewest Y rows that leave X able to cover the rest + auto grid_y = DivRoundingUp(num_blocks, max_blocks); + if(grid_y > std::numeric_limits::max()) + throw std::runtime_error("work item count too large to dispatch as a 2D grid"); + auto grid_x = DivRoundingUp(num_blocks, grid_y); + + return {static_cast(grid_x), static_cast(grid_y), 1}; + }; + if(scheme == CS_KERNEL_CHIRP) { - generator.gridDim - = {static_cast((M - N) / LAUNCH_BOUNDS_BLUESTEIN_MULTI_KERNEL + 1)}; + generator.gridDim = flat_grid_dim(M - N + 1, LAUNCH_BOUNDS_BLUESTEIN_MULTI_KERNEL); generator.blockDim = {LAUNCH_BOUNDS_BLUESTEIN_MULTI_KERNEL}; } else { - generator.gridDim - = {(static_cast(count) - 1) / LAUNCH_BOUNDS_BLUESTEIN_MULTI_KERNEL + 1}; + generator.gridDim = flat_grid_dim(count, LAUNCH_BOUNDS_BLUESTEIN_MULTI_KERNEL); generator.blockDim = {LAUNCH_BOUNDS_BLUESTEIN_MULTI_KERNEL}; } @@ -179,7 +209,8 @@ RTCKernel::RTCGenerator RTCKernelBluesteinMulti::generate_from_node(const LeafNo node.outArrayType, cbtype, node.loadOps, - node.storeOps}; + node.storeOps, + generator.gridDim.y > 1}; generator.generate_name = [=]() { return bluestein_multi_rtc_kernel_name(specs); }; @@ -205,8 +236,10 @@ RTCKernelArgs RTCKernelBluesteinMulti::get_launch_args(DeviceCallIn& data) { int twl = 0; - if(data.node->large1D > (size_t)256 * 256 * 256 * 256) + if(data.node->large1D > (size_t)256 * 256 * 256 * 256 * 256) throw std::runtime_error("large1D twiddle size too large error"); + else if(data.node->large1D > (size_t)256 * 256 * 256 * 256) + twl = 5; else if(data.node->large1D > (size_t)256 * 256 * 256) twl = 4; else if(data.node->large1D > (size_t)256 * 256) diff --git a/projects/rocfft/library/src/rtc_transpose_gen.cpp b/projects/rocfft/library/src/rtc_transpose_gen.cpp index 72fe7c3c7f1d..d88d2350ab2a 100644 --- a/projects/rocfft/library/src/rtc_transpose_gen.cpp +++ b/projects/rocfft/library/src/rtc_transpose_gen.cpp @@ -270,9 +270,9 @@ std::string transpose_rtc(const std::string& kernel_name, const TransposeSpecs& Variable i{"i", rtc_kint_type(KIntType::U32)}; Variable logical_row{"logical_row", "auto"}; Variable logical_col{"logical_col", "auto"}; - Variable idx0{"idx0", "auto"}; - Variable idx1{"idx1", "auto"}; - Variable idx2{"idx2", "auto"}; + Variable idx0{"idx0", "integer_type"}; + Variable idx1{"idx1", "integer_type"}; + Variable idx2{"idx2", "integer_type"}; Variable global_read_idx{"global_read_idx", "auto"}; Variable global_write_idx{"global_write_idx", "auto"}; Variable elem{"elem", "scalar_type"}; diff --git a/projects/rocfft/library/src/rtc_transpose_kernel.cpp b/projects/rocfft/library/src/rtc_transpose_kernel.cpp index 8700f61e5575..d9d319d485ad 100644 --- a/projects/rocfft/library/src/rtc_transpose_kernel.cpp +++ b/projects/rocfft/library/src/rtc_transpose_kernel.cpp @@ -69,8 +69,10 @@ RTCKernel::RTCGenerator RTCKernelTranspose::generate_from_node(const LeafNode& generator.blockDim = {tileX, tileY}; size_t largeTwdSteps = 0; - if(node.large1D > (size_t)256 * 256 * 256 * 256) + if(node.large1D > (size_t)256 * 256 * 256 * 256 * 256) throw std::runtime_error("large1D twiddle size too large error"); + if(node.large1D > (size_t)256 * 256 * 256 * 256) + largeTwdSteps = 5; else if(node.large1D > (size_t)256 * 256 * 256) largeTwdSteps = 4; else if(node.large1D > (size_t)256 * 256) diff --git a/projects/rocfft/shared/device_properties.h b/projects/rocfft/shared/device_properties.h index 522d8c48b73e..9a0a5368361c 100644 --- a/projects/rocfft/shared/device_properties.h +++ b/projects/rocfft/shared/device_properties.h @@ -1,4 +1,4 @@ -// Copyright (C) 2023 Advanced Micro Devices, Inc. All rights reserved. +// Copyright (C) 2023 - 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 @@ -23,6 +23,7 @@ #include #include +#include #include #include @@ -106,6 +107,15 @@ static void launch_limits_check(const std::string& kernel_name, || gridDim.y > static_cast(deviceProp.maxGridSize[1]) || gridDim.z > static_cast(deviceProp.maxGridSize[2])) throw std::runtime_error("max grid size exceeded: " + kernel_name); + + // Total work items along each dimension is counted in uint32_t in the + // dispatch packet, so the grid/block product must fit even when each + // factor on its own is well within the grid and block limits above. + constexpr uint64_t max_work_items_per_dim = std::numeric_limits::max(); + if(static_cast(gridDim.x) * blockDim.x > max_work_items_per_dim + || static_cast(gridDim.y) * blockDim.y > max_work_items_per_dim + || static_cast(gridDim.z) * blockDim.z > max_work_items_per_dim) + throw std::runtime_error("max work items per dim exceeded: " + kernel_name); } #endif diff --git a/projects/rocfft/shared/test_params.h b/projects/rocfft/shared/test_params.h index b1f7a6e50c9a..308a831b01f8 100644 --- a/projects/rocfft/shared/test_params.h +++ b/projects/rocfft/shared/test_params.h @@ -37,6 +37,7 @@ extern double complex_interleaved_prob_factor; extern double real_prob_factor; extern double complex_planar_prob_factor; extern double callback_prob_factor; +extern double very_large_prob_factor; extern double half_epsilon; extern double single_epsilon;