From 33d547afc8d5e7640f5691c8aaa9e47f0affb8ba Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Mon, 17 Aug 2026 10:24:17 -0600 Subject: [PATCH 01/19] feat(hipfft): implement JIT callback API, deprecate function pointer API JIRA ID: AIFFT-5 New APIs to allow specifying JIT callbacks: - hipFFT: * hipfftXtSetJITCallback In rocFFT-backed hipFFT, the callback function is specified as SPIR-V bitcode. cuFFT-backed hipFFT accepts LTO-IR bitcode. Enabled JIT callback tests and removed the running of function pointer callback tests by default. Function pointer callbacks may still be explicitly requested for manual tests with --callback=funcptr. --- projects/hipfft/CHANGELOG.md | 14 ++ projects/hipfft/clients/hipfft_params.h | 3 - .../clients/tests/accuracy_test_callback.cpp | 2 +- .../clients/tests/multi_device_test.cpp | 3 +- .../docs/reference/hipfft-api-usage.rst | 1 + .../hipfft/library/include/hipfft/hipfftXt.h | 80 ++++++++- .../hipfft/library/src/amd_detail/hipfft.cpp | 162 ++++++++++++++---- .../library/src/nvidia_detail/hipfft.cpp | 21 +++ 8 files changed, 243 insertions(+), 43 deletions(-) diff --git a/projects/hipfft/CHANGELOG.md b/projects/hipfft/CHANGELOG.md index 29817e55afd1..81028ef45638 100644 --- a/projects/hipfft/CHANGELOG.md +++ b/projects/hipfft/CHANGELOG.md @@ -9,6 +9,14 @@ Documentation for hipFFT is available at * Added amdgcnspirv architecture to client programs, so that they are functional even on gfx architectures that they have not been explicitly compiled in. +* Implemented `hipfftXtSetJITCallback` API, to allow for user-defined device functions to be called when loading + input or storing output of a transform. These callback functions are specified after a plan is allocated with + `hipfftCreate` but before the plan is initialized with one of the `MakePlan` functions. The backend FFT library + will Just-In-Time (JIT) compile the code into its own kernels. + + On AMD platforms, the device function is provided as SPIR-V. On CUDA platforms, the device function is provided as + LTO-IR fatbin. + ### Changed * Modified the rocFFT backend's implementation details of hipFFT so that cuFFT backend's @@ -19,6 +27,12 @@ Documentation for hipFFT is available at Multi-device, unbatched one-dimensional transforms remain unimplemented pending further analyses of the exact behavior(s) to be matched. +### Deprecations + +* The `hipfftXtSetCallback` and `hipfftXtClearCallback` APIs are now deprecated and will be removed in a future + release. They allow for specifying callbacks as device function pointers at plan execution time, but rocFFT cannot + optimize the combined code. Instead, users should specify JIT callbacks on plan descriptions. + ## hipFFT 1.0.25 for ROCm 10.0 ### Changed diff --git a/projects/hipfft/clients/hipfft_params.h b/projects/hipfft/clients/hipfft_params.h index 602a72dd20d6..f909fe99ee5d 100644 --- a/projects/hipfft/clients/hipfft_params.h +++ b/projects/hipfft/clients/hipfft_params.h @@ -1256,8 +1256,6 @@ class hipfft_params : public fft_params { return HIPFFT_SUCCESS; } - throw unimplemented_exception("jit callbacks not implemented"); -#if 0 hipfftResult_t ret{HIPFFT_INVALID_PLAN}; hipfftXtCallbackType cbtype = HIPFFT_CB_UNDEFINED; switch(itype) @@ -1362,7 +1360,6 @@ class hipfft_params : public fft_params ? nullptr : store_jit_cb_state->get_raw_data_ptrs().data()); return ret; -#endif } // call hipfftCreate + hipfftMake* functions, inserting calls to diff --git a/projects/hipfft/clients/tests/accuracy_test_callback.cpp b/projects/hipfft/clients/tests/accuracy_test_callback.cpp index 1ed44ab5b50a..1907b32c3d6e 100644 --- a/projects/hipfft/clients/tests/accuracy_test_callback.cpp +++ b/projects/hipfft/clients/tests/accuracy_test_callback.cpp @@ -104,7 +104,7 @@ const static std::vector> ooffset_range = {{0, 0}, {1, 1}}; auto transform_types = {fft_transform_type_complex_forward, fft_transform_type_real_forward}; #ifdef __HIP__ -INSTANTIATE_TEST_SUITE_P(DISABLED_callback_no_offset, +INSTANTIATE_TEST_SUITE_P(callback_no_offset, accuracy_test, ::testing::ValuesIn(param_generator_base(test_prob, transform_types, diff --git a/projects/hipfft/clients/tests/multi_device_test.cpp b/projects/hipfft/clients/tests/multi_device_test.cpp index 613a96239262..41b21a36f473 100644 --- a/projects/hipfft/clients/tests/multi_device_test.cpp +++ b/projects/hipfft/clients/tests/multi_device_test.cpp @@ -82,7 +82,8 @@ std::vector param_generator_multi_gpu(const std::optional // function pointer callbacks need -fgpu-rdc, but that causes build // nondeterminism in kpack - auto multi_device_callbacks = {fft_callback_type_none, /*fft_callback_type_funcptr, */}; + auto multi_device_callbacks + = {fft_callback_type_none, /*fft_callback_type_funcptr, */ fft_callback_type_jit}; // gather cases to test as single-device params, then distribute // to multiple GPUs diff --git a/projects/hipfft/docs/reference/hipfft-api-usage.rst b/projects/hipfft/docs/reference/hipfft-api-usage.rst index 747bf9d737bb..42d71fd95b5c 100644 --- a/projects/hipfft/docs/reference/hipfft-api-usage.rst +++ b/projects/hipfft/docs/reference/hipfft-api-usage.rst @@ -187,6 +187,7 @@ Callbacks .. doxygenfunction:: hipfftXtSetCallback .. doxygenfunction:: hipfftXtClearCallback .. doxygenfunction:: hipfftXtSetCallbackSharedSize +.. doxygenfunction:: hipfftXtSetJITCallback Single-process multi-GPU transforms diff --git a/projects/hipfft/library/include/hipfft/hipfftXt.h b/projects/hipfft/library/include/hipfft/hipfftXt.h index 5642a87e57c2..d5355b8ee6ef 100644 --- a/projects/hipfft/library/include/hipfft/hipfftXt.h +++ b/projects/hipfft/library/include/hipfft/hipfftXt.h @@ -1,5 +1,5 @@ /****************************************************************************** - * Copyright (C) 2021 - 2023 Advanced Micro Devices, Inc. All rights + * Copyright (C) 2021 - 2026 Advanced Micro Devices, Inc. All rights * reserved. * * Permission is hereby granted, free of charge, to any person obtaining a copy @@ -85,12 +85,13 @@ typedef void (*hipfftCallbackStoreR)( typedef void (*hipfftCallbackStoreD)( void* dataOut, size_t offset, hipfftDoubleReal element, void* callerInfo, void* sharedPointer); -/*! @brief Set a callback on a plan +/*! @brief Set a function pointer callback on a plan (deprecated) * - * @details Set either a load or store callback to run with a plan. - * The type of callback is specified with the 'cbtype' parameter. - * An array ofcallback and callback data pointers must be given - - * one per device executing the plan. + * @details Set either a load or store callback to run with a plan, + * specified as a device function pointer. The type of callback is + * specified with the 'cbtype' parameter. An array of callback and + * callback data pointers must be given - one per device executing + * the plan. * * @param[in] plan The FFT plan. * @param[in] callbacks Array of callback function pointers. @@ -102,9 +103,9 @@ HIPFFT_EXPORT hipfftResult hipfftXtSetCallback(hipfftHandle plan, hipfftXtCallbackType cbtype, void** callbackData); -/*! @brief Remove a callback from a plan +/*! @brief Remove a function pointer callback from a plan (deprecated) * - * @details Remove a previously-set callback from a plan. + * @details Remove a previously-set function pointer callback from a plan. * * @param[in] plan The FFT plan. * @param[in] cbtype Type of callback being removed. @@ -125,6 +126,69 @@ HIPFFT_EXPORT hipfftResult hipfftXtSetCallbackSharedSize(hipfftHandle pl hipfftXtCallbackType cbtype, size_t sharedSize); +typedef hipfftComplex (*hipfftJITCallbackLoadC)(void* data, + unsigned long long offset, + void* cbdata, + void* sharedMem); +typedef hipfftDoubleComplex (*hipfftJITCallbackLoadZ)(void* data, + unsigned long long offset, + void* cbdata, + void* sharedMem); +typedef hipfftReal (*hipfftJITCallbackLoadR)(void* data, + unsigned long long offset, + void* cbdata, + void* sharedMem); +typedef hipfftDoubleReal (*hipfftJITCallbackLoadD)(void* data, + unsigned long long offset, + void* cbdata, + void* sharedMem); + +typedef void (*hipfftJITCallbackStoreC)( + void* data, unsigned long long offset, hipfftComplex element, void* cbdata, void* sharedMem); +typedef void (*hipfftJITCallbackStoreZ)(void* data, + unsigned long long offset, + hipfftDoubleComplex element, + void* cbdata, + void* sharedMem); +typedef void (*hipfftJITCallbackStoreR)( + void* data, unsigned long long offset, hipfftReal element, void* cbdata, void* sharedMem); +typedef void (*hipfftJITCallbackStoreD)( + void* data, unsigned long long offset, hipfftDoubleReal element, void* cbdata, void* sharedMem); + +/*! @brief Set a JIT callback on a plan + * + * @details Set either a load or store callback to run with a plan. + * The type of callback is specified with the 'cbtype' parameter. + * The callback is provided as SPIR-V on AMD platforms and as + * LTO-IR fatbin on CUDA platforms. + * + * 'cbdata' is an optional array of pointers of data that is passed + * to the callback function, one per visible HIP device. hipFFT may + * execute the callback function on the current HIP device as well + * as any device used by a multi-device transform. The + * corresponding pointer for a device in the array will be passed to + * the callback function's `cbdata` parameter when it is executed on + * that device. + * + * This function must be called after the plan is allocated using + * ::hipfftCreate, but before the plan is initialized by any of the + * "MakePlan" functions. Therefore, API functions that combine + * creation and initialization (::hipfftPlan1d, ::hipfftPlan2d, + * ::hipfftPlan3d, and ::hipfftPlanMany) cannot set a JIT callback. + * + * @param[in] plan The FFT plan. + * @param[in] symbol_name Name of the symbol in the compiled bitcode. + * @param[in] bitcode_data Pointer to bitcode data. + * @param[in] bitcode_len_bytes Number of bytes in the bitcode data. + * @param[in] callbackData Array of callback function data pointers. + */ +HIPFFT_EXPORT hipfftResult hipfftXtSetJITCallback(hipfftHandle plan, + const char* symbol_name, + const void* bitcode_data, + size_t bitcode_len_bytes, + hipfftXtCallbackType cbtype, + void** cbdata); + /*! @brief Initialize a batched rank-dimensional FFT plan with advanced data layout and specified input, output, execution data types. diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index dd4bffe19c2b..65edac46af5f 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -829,12 +829,19 @@ struct hipfftHandle_t rocfft_execution_info_wrapper_t info; std::vector device_contexts; - void** load_callback_ptrs = nullptr; - void** load_callback_data = nullptr; - size_t load_callback_lds_bytes = 0; - void** store_callback_ptrs = nullptr; - void** store_callback_data = nullptr; - size_t store_callback_lds_bytes = 0; + void** load_callback_ptrs = nullptr; + const char* load_callback_symbol = nullptr; + const void* load_callback_bitcode = nullptr; + size_t load_callback_bitcode_len = 0; + std::vector load_callback_data; + size_t load_callback_lds_bytes = 0; + + void** store_callback_ptrs = nullptr; + const char* store_callback_symbol = nullptr; + const void* store_callback_bitcode = nullptr; + size_t store_callback_bitcode_len = 0; + std::vector store_callback_data; + size_t store_callback_lds_bytes = 0; // Multi-processing communicator rocfft_comm_type comm_type = rocfft_comm_none; @@ -992,6 +999,28 @@ struct hipfftHandle_t } return ret; } + + // Helpers to set/get callback data, since we can be passed a + // void** but need to copy the right number of pointers. When + // passing callback data to library APIs, we also want to pass + // nullptr if we were never given any callback data. + static std::vector callback_data_from_ptrs(void** ptrs) + { + std::vector cbdata_vec; + if(ptrs) + { + // Copy one pointer per device + std::copy( + ptrs, ptrs + rocfft_scoped_device::device_count(), std::back_inserter(cbdata_vec)); + } + return cbdata_vec; + } + static void** callback_data_to_ptrs(std::vector& callback_data) + { + if(callback_data.empty()) + return nullptr; + return callback_data.data(); + } }; static inline hipfftResult handle_exception() noexcept @@ -1216,6 +1245,30 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_comm(desc, plan->comm_type, plan->comm_handle)); + // set JIT callbacks if specified + if(plan->load_callback_symbol && plan->load_callback_bitcode + && plan->load_callback_bitcode_len) + { + ROCFFT_EXPECT_SUCCESS(rocfft_plan_description_set_load_callback( + desc, + plan->load_callback_symbol, + plan->load_callback_bitcode, + plan->load_callback_bitcode_len, + hipfftHandle_t::callback_data_to_ptrs(plan->load_callback_data), + plan->load_callback_lds_bytes)); + } + if(plan->store_callback_symbol && plan->store_callback_bitcode + && plan->store_callback_bitcode_len) + { + ROCFFT_EXPECT_SUCCESS(rocfft_plan_description_set_store_callback( + desc, + plan->store_callback_symbol, + plan->store_callback_bitcode, + plan->store_callback_bitcode_len, + hipfftHandle_t::callback_data_to_ptrs(plan->store_callback_data), + plan->store_callback_lds_bytes)); + } + if(plan->device_contexts.size() > 1) { // Determine the sub-formats of this rocfft plan's input and output fields. Unbatched @@ -2024,7 +2077,7 @@ try || plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; + plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_LD_COMPLEX_DOUBLE: @@ -2032,7 +2085,7 @@ try || plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; + plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_LD_REAL: @@ -2040,7 +2093,7 @@ try || !plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; + plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_LD_REAL_DOUBLE: @@ -2048,7 +2101,7 @@ try || !plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; + plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_COMPLEX: @@ -2056,7 +2109,7 @@ try || plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; + plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_COMPLEX_DOUBLE: @@ -2064,7 +2117,7 @@ try || plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; + plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_REAL: @@ -2072,7 +2125,7 @@ try || !plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; + plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_REAL_DOUBLE: @@ -2080,7 +2133,7 @@ try || !plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; + plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_UNDEFINED: @@ -2088,16 +2141,18 @@ try } rocfft_status res; - res = rocfft_execution_info_set_load_callback(plan->info, - plan->load_callback_ptrs, - plan->load_callback_data, - plan->load_callback_lds_bytes); + res = rocfft_execution_info_set_load_callback( + plan->info, + plan->load_callback_ptrs, + hipfftHandle_t::callback_data_to_ptrs(plan->load_callback_data), + plan->load_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; - res = rocfft_execution_info_set_store_callback(plan->info, - plan->store_callback_ptrs, - plan->store_callback_data, - plan->store_callback_lds_bytes); + res = rocfft_execution_info_set_store_callback( + plan->info, + plan->store_callback_ptrs, + hipfftHandle_t::callback_data_to_ptrs(plan->store_callback_data), + plan->store_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; return HIPFFT_SUCCESS; @@ -2143,16 +2198,18 @@ try } rocfft_status res; - res = rocfft_execution_info_set_load_callback(plan->info, - plan->load_callback_ptrs, - plan->load_callback_data, - plan->load_callback_lds_bytes); + res = rocfft_execution_info_set_load_callback( + plan->info, + plan->load_callback_ptrs, + hipfftHandle_t::callback_data_to_ptrs(plan->load_callback_data), + plan->load_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; - res = rocfft_execution_info_set_store_callback(plan->info, - plan->store_callback_ptrs, - plan->store_callback_data, - plan->store_callback_lds_bytes); + res = rocfft_execution_info_set_store_callback( + plan->info, + plan->store_callback_ptrs, + hipfftHandle_t::callback_data_to_ptrs(plan->store_callback_data), + plan->store_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; return HIPFFT_SUCCESS; @@ -2162,6 +2219,51 @@ catch(...) return handle_exception(); } +hipfftResult hipfftXtSetJITCallback(hipfftHandle plan, + const char* symbol_name, + const void* bitcode_data, + size_t bitcode_len_bytes, + hipfftXtCallbackType cbtype, + void** cbdata) +try +{ + if(!plan) + return HIPFFT_INVALID_PLAN; + + switch(cbtype) + { + case HIPFFT_CB_LD_COMPLEX: + case HIPFFT_CB_LD_COMPLEX_DOUBLE: + case HIPFFT_CB_LD_REAL: + case HIPFFT_CB_LD_REAL_DOUBLE: + { + plan->load_callback_symbol = symbol_name; + plan->load_callback_bitcode = bitcode_data; + plan->load_callback_bitcode_len = bitcode_len_bytes; + plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(cbdata); + break; + } + case HIPFFT_CB_ST_COMPLEX: + case HIPFFT_CB_ST_COMPLEX_DOUBLE: + case HIPFFT_CB_ST_REAL: + case HIPFFT_CB_ST_REAL_DOUBLE: + { + plan->store_callback_symbol = symbol_name; + plan->store_callback_bitcode = bitcode_data; + plan->store_callback_bitcode_len = bitcode_len_bytes; + plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(cbdata); + break; + } + case HIPFFT_CB_UNDEFINED: + return HIPFFT_INVALID_VALUE; + } + return HIPFFT_SUCCESS; +} +catch(...) +{ + return handle_exception(); +} + hipfftResult hipfftXtMakePlanMany(hipfftHandle plan, int rank, long long int* n, diff --git a/projects/hipfft/library/src/nvidia_detail/hipfft.cpp b/projects/hipfft/library/src/nvidia_detail/hipfft.cpp index fafb4c8b6455..d65483c9d5e1 100644 --- a/projects/hipfft/library/src/nvidia_detail/hipfft.cpp +++ b/projects/hipfft/library/src/nvidia_detail/hipfft.cpp @@ -787,6 +787,27 @@ catch(...) return handle_exception(); } +hipfftResult hipfftXtSetJITCallback(hipfftHandle plan, + const char* symbol_name, + const void* bitcode_data, + size_t bitcode_len_bytes, + hipfftXtCallbackType cbtype, + void** cbdata) +try +{ + return cufftResultToHipResult( + cufftXtSetJITCallback(plan, + symbol_name, + bitcode_data, + bitcode_len_bytes, + hipfftCallbackTypeToCufftCallbackType(cbtype), + cbdata)); +} +catch(...) +{ + return handle_exception(); +} + hipfftResult hipfftXtMakePlanMany(hipfftHandle plan, int rank, long long int* n, From e47c82e420ce8f79951129747812fb037d62fecb Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Wed, 19 Aug 2026 14:15:00 -0600 Subject: [PATCH 02/19] use execution_info jit cb data APIs --- projects/hipfft/clients/hipfft_params.h | 23 ++- .../hipfft/library/include/hipfft/hipfftXt.h | 30 ++-- .../hipfft/library/src/amd_detail/hipfft.cpp | 143 ++++++++---------- 3 files changed, 99 insertions(+), 97 deletions(-) diff --git a/projects/hipfft/clients/hipfft_params.h b/projects/hipfft/clients/hipfft_params.h index f909fe99ee5d..7a1326f2616a 100644 --- a/projects/hipfft/clients/hipfft_params.h +++ b/projects/hipfft/clients/hipfft_params.h @@ -163,6 +163,11 @@ class hipfft_params : public fft_params // will be provided with externally-managed work area(s): static std::vector externally_managed_workareas; + // JIT callback data pointers, stored in a params-level vector so + // that they live as long as the plan + std::vector load_jit_cb_data_ptrs; + std::vector store_jit_cb_data_ptrs; + static std::vector externally_managed_extra_vram_footprint() { std::vector footprint; @@ -1300,14 +1305,17 @@ class hipfft_params : public fft_params } check_jit_callback_state(); + load_jit_cb_data_ptrs.resize(load_jit_cb_state->data.size()); + std::transform(load_jit_cb_state->data.begin(), + load_jit_cb_state->data.end(), + load_jit_cb_data_ptrs.begin(), + [](gpubuf& buf) { return buf.data(); }); ret = hipfftXtSetJITCallback(plan, load_jit_cb_state->symbol, load_jit_cb_state->func.data(), load_jit_cb_state->func.size(), cbtype, - load_jit_cb_state->data.empty() - ? nullptr - : load_jit_cb_state->get_raw_data_ptrs().data()); + load_jit_cb_data_ptrs.data()); if(ret != HIPFFT_SUCCESS) return ret; @@ -1351,14 +1359,17 @@ class hipfft_params : public fft_params throw std::runtime_error("unsupported data type for store callback"); } } + store_jit_cb_data_ptrs.resize(store_jit_cb_state->data.size()); + std::transform(store_jit_cb_state->data.begin(), + store_jit_cb_state->data.end(), + store_jit_cb_data_ptrs.begin(), + [](gpubuf& buf) { return buf.data(); }); ret = hipfftXtSetJITCallback(plan, store_jit_cb_state->symbol, store_jit_cb_state->func.data(), store_jit_cb_state->func.size(), cbtype, - store_jit_cb_state->data.empty() - ? nullptr - : store_jit_cb_state->get_raw_data_ptrs().data()); + store_jit_cb_data_ptrs.data()); return ret; } diff --git a/projects/hipfft/library/include/hipfft/hipfftXt.h b/projects/hipfft/library/include/hipfft/hipfftXt.h index d5355b8ee6ef..4227a79a5669 100644 --- a/projects/hipfft/library/include/hipfft/hipfftXt.h +++ b/projects/hipfft/library/include/hipfft/hipfftXt.h @@ -93,6 +93,9 @@ typedef void (*hipfftCallbackStoreD)( * callback data pointers must be given - one per device executing * the plan. * + * The specified arrays must remain valid until either the plan is + * destroyed or callbacks are cleared from the plan. + * * @param[in] plan The FFT plan. * @param[in] callbacks Array of callback function pointers. * @param[in] cbtype Type of callback being set. @@ -162,19 +165,22 @@ typedef void (*hipfftJITCallbackStoreD)( * The callback is provided as SPIR-V on AMD platforms and as * LTO-IR fatbin on CUDA platforms. * - * 'cbdata' is an optional array of pointers of data that is passed - * to the callback function, one per visible HIP device. hipFFT may - * execute the callback function on the current HIP device as well - * as any device used by a multi-device transform. The - * corresponding pointer for a device in the array will be passed to - * the callback function's `cbdata` parameter when it is executed on - * that device. + * 'cbdata' is an optional array of pointers of data that is passed + * to the callback function, one per visible HIP device. hipFFT may + * execute the callback function on the current HIP device as well + * as any device used by a multi-device transform. The + * corresponding pointer for a device in the array will be passed to + * the callback function's `cbdata` parameter when it is executed on + * that device. + * + * If specified, the 'cbdata' array must remain valid until either + * the plan is destroyed or callbacks are cleared from the plan. * - * This function must be called after the plan is allocated using - * ::hipfftCreate, but before the plan is initialized by any of the - * "MakePlan" functions. Therefore, API functions that combine - * creation and initialization (::hipfftPlan1d, ::hipfftPlan2d, - * ::hipfftPlan3d, and ::hipfftPlanMany) cannot set a JIT callback. + * This function must be called after the plan is allocated using + * ::hipfftCreate, but before the plan is initialized by any of the + * "MakePlan" functions. Therefore, API functions that combine + * creation and initialization (::hipfftPlan1d, ::hipfftPlan2d, + * ::hipfftPlan3d, and ::hipfftPlanMany) cannot set a JIT callback. * * @param[in] plan The FFT plan. * @param[in] symbol_name Name of the symbol in the compiled bitcode. diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index 65edac46af5f..88afa42db795 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -829,19 +829,19 @@ struct hipfftHandle_t rocfft_execution_info_wrapper_t info; std::vector device_contexts; - void** load_callback_ptrs = nullptr; - const char* load_callback_symbol = nullptr; - const void* load_callback_bitcode = nullptr; - size_t load_callback_bitcode_len = 0; - std::vector load_callback_data; - size_t load_callback_lds_bytes = 0; - - void** store_callback_ptrs = nullptr; - const char* store_callback_symbol = nullptr; - const void* store_callback_bitcode = nullptr; - size_t store_callback_bitcode_len = 0; - std::vector store_callback_data; - size_t store_callback_lds_bytes = 0; + void** load_callback_ptrs = nullptr; + const char* load_callback_symbol = nullptr; + const void* load_callback_bitcode = nullptr; + size_t load_callback_bitcode_len = 0; + void** load_callback_data = nullptr; + size_t load_callback_lds_bytes = 0; + + void** store_callback_ptrs = nullptr; + const char* store_callback_symbol = nullptr; + const void* store_callback_bitcode = nullptr; + size_t store_callback_bitcode_len = 0; + void** store_callback_data = nullptr; + size_t store_callback_lds_bytes = 0; // Multi-processing communicator rocfft_comm_type comm_type = rocfft_comm_none; @@ -999,28 +999,6 @@ struct hipfftHandle_t } return ret; } - - // Helpers to set/get callback data, since we can be passed a - // void** but need to copy the right number of pointers. When - // passing callback data to library APIs, we also want to pass - // nullptr if we were never given any callback data. - static std::vector callback_data_from_ptrs(void** ptrs) - { - std::vector cbdata_vec; - if(ptrs) - { - // Copy one pointer per device - std::copy( - ptrs, ptrs + rocfft_scoped_device::device_count(), std::back_inserter(cbdata_vec)); - } - return cbdata_vec; - } - static void** callback_data_to_ptrs(std::vector& callback_data) - { - if(callback_data.empty()) - return nullptr; - return callback_data.data(); - } }; static inline hipfftResult handle_exception() noexcept @@ -1249,24 +1227,22 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, if(plan->load_callback_symbol && plan->load_callback_bitcode && plan->load_callback_bitcode_len) { - ROCFFT_EXPECT_SUCCESS(rocfft_plan_description_set_load_callback( - desc, - plan->load_callback_symbol, - plan->load_callback_bitcode, - plan->load_callback_bitcode_len, - hipfftHandle_t::callback_data_to_ptrs(plan->load_callback_data), - plan->load_callback_lds_bytes)); + ROCFFT_EXPECT_SUCCESS( + rocfft_plan_description_set_load_callback(desc, + plan->load_callback_symbol, + plan->load_callback_bitcode, + plan->load_callback_bitcode_len, + plan->load_callback_lds_bytes)); } if(plan->store_callback_symbol && plan->store_callback_bitcode && plan->store_callback_bitcode_len) { - ROCFFT_EXPECT_SUCCESS(rocfft_plan_description_set_store_callback( - desc, - plan->store_callback_symbol, - plan->store_callback_bitcode, - plan->store_callback_bitcode_len, - hipfftHandle_t::callback_data_to_ptrs(plan->store_callback_data), - plan->store_callback_lds_bytes)); + ROCFFT_EXPECT_SUCCESS( + rocfft_plan_description_set_store_callback(desc, + plan->store_callback_symbol, + plan->store_callback_bitcode, + plan->store_callback_bitcode_len, + plan->store_callback_lds_bytes)); } if(plan->device_contexts.size() > 1) @@ -1365,6 +1341,19 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, } } + // if JIT callbacks are used, pass the cbdata to the execution info + if(plan->load_callback_symbol && plan->load_callback_bitcode && plan->load_callback_bitcode_len) + { + ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_load_callback_data( + plan->info, plan->load_callback_data, plan->device_contexts.size())); + } + if(plan->store_callback_symbol && plan->store_callback_bitcode + && plan->store_callback_bitcode_len) + { + ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_store_callback_data( + plan->info, plan->store_callback_data, plan->device_contexts.size())); + } + return HIPFFT_SUCCESS; } @@ -2077,7 +2066,7 @@ try || plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->load_callback_data = callbackData; plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_LD_COMPLEX_DOUBLE: @@ -2085,7 +2074,7 @@ try || plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->load_callback_data = callbackData; plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_LD_REAL: @@ -2093,7 +2082,7 @@ try || !plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->load_callback_data = callbackData; plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_LD_REAL_DOUBLE: @@ -2101,7 +2090,7 @@ try || !plan->io_type.is_real_to_complex()) return HIPFFT_INVALID_VALUE; plan->load_callback_ptrs = callbacks; - plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->load_callback_data = callbackData; plan->load_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_COMPLEX: @@ -2109,7 +2098,7 @@ try || plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->store_callback_data = callbackData; plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_COMPLEX_DOUBLE: @@ -2117,7 +2106,7 @@ try || plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->store_callback_data = callbackData; plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_REAL: @@ -2125,7 +2114,7 @@ try || !plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->store_callback_data = callbackData; plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_ST_REAL_DOUBLE: @@ -2133,7 +2122,7 @@ try || !plan->io_type.is_complex_to_real()) return HIPFFT_INVALID_VALUE; plan->store_callback_ptrs = callbacks; - plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(callbackData); + plan->store_callback_data = callbackData; plan->store_callback_lds_bytes = 0; break; case HIPFFT_CB_UNDEFINED: @@ -2141,18 +2130,16 @@ try } rocfft_status res; - res = rocfft_execution_info_set_load_callback( - plan->info, - plan->load_callback_ptrs, - hipfftHandle_t::callback_data_to_ptrs(plan->load_callback_data), - plan->load_callback_lds_bytes); + res = rocfft_execution_info_set_load_callback(plan->info, + plan->load_callback_ptrs, + plan->load_callback_data, + plan->load_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; - res = rocfft_execution_info_set_store_callback( - plan->info, - plan->store_callback_ptrs, - hipfftHandle_t::callback_data_to_ptrs(plan->store_callback_data), - plan->store_callback_lds_bytes); + res = rocfft_execution_info_set_store_callback(plan->info, + plan->store_callback_ptrs, + plan->store_callback_data, + plan->store_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; return HIPFFT_SUCCESS; @@ -2198,18 +2185,16 @@ try } rocfft_status res; - res = rocfft_execution_info_set_load_callback( - plan->info, - plan->load_callback_ptrs, - hipfftHandle_t::callback_data_to_ptrs(plan->load_callback_data), - plan->load_callback_lds_bytes); + res = rocfft_execution_info_set_load_callback(plan->info, + plan->load_callback_ptrs, + plan->load_callback_data, + plan->load_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; - res = rocfft_execution_info_set_store_callback( - plan->info, - plan->store_callback_ptrs, - hipfftHandle_t::callback_data_to_ptrs(plan->store_callback_data), - plan->store_callback_lds_bytes); + res = rocfft_execution_info_set_store_callback(plan->info, + plan->store_callback_ptrs, + plan->store_callback_data, + plan->store_callback_lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; return HIPFFT_SUCCESS; @@ -2240,7 +2225,7 @@ try plan->load_callback_symbol = symbol_name; plan->load_callback_bitcode = bitcode_data; plan->load_callback_bitcode_len = bitcode_len_bytes; - plan->load_callback_data = hipfftHandle_t::callback_data_from_ptrs(cbdata); + plan->load_callback_data = cbdata; break; } case HIPFFT_CB_ST_COMPLEX: @@ -2251,7 +2236,7 @@ try plan->store_callback_symbol = symbol_name; plan->store_callback_bitcode = bitcode_data; plan->store_callback_bitcode_len = bitcode_len_bytes; - plan->store_callback_data = hipfftHandle_t::callback_data_from_ptrs(cbdata); + plan->store_callback_data = cbdata; break; } case HIPFFT_CB_UNDEFINED: From 3cb6cda14feb85dcea4af1a649fb8e5ba8d8cea0 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 13:52:55 -0600 Subject: [PATCH 03/19] callback_prob -> 0.1 --- projects/hipfft/clients/tests/gtest_main.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/hipfft/clients/tests/gtest_main.cpp b/projects/hipfft/clients/tests/gtest_main.cpp index a2e60332c29c..584fec3e3780 100644 --- a/projects/hipfft/clients/tests/gtest_main.cpp +++ b/projects/hipfft/clients/tests/gtest_main.cpp @@ -387,7 +387,7 @@ int main(int argc, char* argv[]) app.add_option("--callback_prob", callback_prob_factor, "Probability multiplier for running individual callback transforms") - ->default_val(0.0) + ->default_val(0.1) ->check(CLI::NonNegativeNumber); constexpr auto emulation_quick = "quick"; constexpr auto emulation_smoke = "smoke"; From ad1a881d14fa25acbdc29d0edca6d65dada216a1 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 13:54:17 -0600 Subject: [PATCH 04/19] deprecate legacy callback APIs --- projects/hipfft/library/include/hipfft/hipfft.h | 12 +++++++++++- projects/hipfft/library/include/hipfft/hipfftXt.h | 2 ++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/projects/hipfft/library/include/hipfft/hipfft.h b/projects/hipfft/library/include/hipfft/hipfft.h index 8b68d121928b..9747868a7490 100644 --- a/projects/hipfft/library/include/hipfft/hipfft.h +++ b/projects/hipfft/library/include/hipfft/hipfft.h @@ -1,5 +1,5 @@ /****************************************************************************** - * Copyright (C) 2016 - 2022 Advanced Micro Devices, Inc. All rights + * Copyright (C) 2016 - 2026 Advanced Micro Devices, Inc. All rights * reserved. * * Permission is hereby granted, free of charge, to any person obtaining a copy @@ -59,6 +59,16 @@ DISABLE_WARNING_IGNORED_ATTRIBUTES #include DISABLE_WARNING_POP +#if defined(__cplusplus) && __cplusplus >= 201402L +#define HIPFFT_DEPRECATED_MSG(msg) [[deprecated(msg)]] +#elif defined(__GNUC__) +#define HIPFFT_DEPRECATED_MSG(msg) __attribute__((deprecated(msg))) +#elif defined(_MSC_VER) +#define HIPFFT_DEPRECATED_MSG(msg) __declspec(deprecated(msg)) +#else // no-op +#define HIPFFT_DEPRECATED_MSG(msg) +#endif + #ifdef __cplusplus #include extern "C" { diff --git a/projects/hipfft/library/include/hipfft/hipfftXt.h b/projects/hipfft/library/include/hipfft/hipfftXt.h index 4227a79a5669..2c77d915fe67 100644 --- a/projects/hipfft/library/include/hipfft/hipfftXt.h +++ b/projects/hipfft/library/include/hipfft/hipfftXt.h @@ -101,6 +101,7 @@ typedef void (*hipfftCallbackStoreD)( * @param[in] cbtype Type of callback being set. * @param[in] callbackData Array of callback function data pointers */ +HIPFFT_DEPRECATED_MSG("Use hipfftXtSetJITCallback instead.") HIPFFT_EXPORT hipfftResult hipfftXtSetCallback(hipfftHandle plan, void** callbacks, hipfftXtCallbackType cbtype, @@ -113,6 +114,7 @@ HIPFFT_EXPORT hipfftResult hipfftXtSetCallback(hipfftHandle plan, * @param[in] plan The FFT plan. * @param[in] cbtype Type of callback being removed. */ +HIPFFT_DEPRECATED_MSG("Use hipfftXtSetJITCallback instead.") HIPFFT_EXPORT hipfftResult hipfftXtClearCallback(hipfftHandle plan, hipfftXtCallbackType cbtype); /*! @brief Set shared memory size for callback. From f9ee401c8087a3262f4bdcc0d21a93fa373e81d3 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 14:11:31 -0600 Subject: [PATCH 05/19] make copies of jit symbol name and bitcode --- .../hipfft/library/src/amd_detail/hipfft.cpp | 65 +++++++++---------- 1 file changed, 31 insertions(+), 34 deletions(-) diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index 88afa42db795..2c7dbc362815 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -829,19 +829,17 @@ struct hipfftHandle_t rocfft_execution_info_wrapper_t info; std::vector device_contexts; - void** load_callback_ptrs = nullptr; - const char* load_callback_symbol = nullptr; - const void* load_callback_bitcode = nullptr; - size_t load_callback_bitcode_len = 0; - void** load_callback_data = nullptr; - size_t load_callback_lds_bytes = 0; - - void** store_callback_ptrs = nullptr; - const char* store_callback_symbol = nullptr; - const void* store_callback_bitcode = nullptr; - size_t store_callback_bitcode_len = 0; - void** store_callback_data = nullptr; - size_t store_callback_lds_bytes = 0; + void** load_callback_ptrs = nullptr; + std::string load_callback_symbol; + std::vector load_callback_bitcode; + void** load_callback_data = nullptr; + size_t load_callback_lds_bytes = 0; + + void** store_callback_ptrs = nullptr; + std::string store_callback_symbol; + std::vector store_callback_bitcode; + void** store_callback_data = nullptr; + size_t store_callback_lds_bytes = 0; // Multi-processing communicator rocfft_comm_type comm_type = rocfft_comm_none; @@ -1224,24 +1222,22 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, rocfft_plan_description_set_comm(desc, plan->comm_type, plan->comm_handle)); // set JIT callbacks if specified - if(plan->load_callback_symbol && plan->load_callback_bitcode - && plan->load_callback_bitcode_len) + if(!plan->load_callback_symbol.empty() && !plan->load_callback_bitcode.empty()) { ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_load_callback(desc, - plan->load_callback_symbol, - plan->load_callback_bitcode, - plan->load_callback_bitcode_len, + plan->load_callback_symbol.c_str(), + plan->load_callback_bitcode.data(), + plan->load_callback_bitcode.size(), plan->load_callback_lds_bytes)); } - if(plan->store_callback_symbol && plan->store_callback_bitcode - && plan->store_callback_bitcode_len) + if(!plan->store_callback_symbol.empty() && !plan->store_callback_bitcode.empty()) { ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_store_callback(desc, - plan->store_callback_symbol, - plan->store_callback_bitcode, - plan->store_callback_bitcode_len, + plan->store_callback_symbol.c_str(), + plan->store_callback_bitcode.data(), + plan->store_callback_bitcode.size(), plan->store_callback_lds_bytes)); } @@ -1342,13 +1338,12 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, } // if JIT callbacks are used, pass the cbdata to the execution info - if(plan->load_callback_symbol && plan->load_callback_bitcode && plan->load_callback_bitcode_len) + if(!plan->load_callback_symbol.empty() && !plan->load_callback_bitcode.empty()) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_load_callback_data( plan->info, plan->load_callback_data, plan->device_contexts.size())); } - if(plan->store_callback_symbol && plan->store_callback_bitcode - && plan->store_callback_bitcode_len) + if(!plan->store_callback_symbol.empty() && plan->store_callback_bitcode.empty()) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_store_callback_data( plan->info, plan->store_callback_data, plan->device_contexts.size())); @@ -2222,10 +2217,11 @@ try case HIPFFT_CB_LD_REAL: case HIPFFT_CB_LD_REAL_DOUBLE: { - plan->load_callback_symbol = symbol_name; - plan->load_callback_bitcode = bitcode_data; - plan->load_callback_bitcode_len = bitcode_len_bytes; - plan->load_callback_data = cbdata; + plan->load_callback_symbol = symbol_name; + plan->load_callback_bitcode.assign(static_cast(bitcode_data), + static_cast(bitcode_data) + + bitcode_len_bytes); + plan->load_callback_data = cbdata; break; } case HIPFFT_CB_ST_COMPLEX: @@ -2233,10 +2229,11 @@ try case HIPFFT_CB_ST_REAL: case HIPFFT_CB_ST_REAL_DOUBLE: { - plan->store_callback_symbol = symbol_name; - plan->store_callback_bitcode = bitcode_data; - plan->store_callback_bitcode_len = bitcode_len_bytes; - plan->store_callback_data = cbdata; + plan->store_callback_symbol = symbol_name; + plan->store_callback_bitcode.assign(static_cast(bitcode_data), + static_cast(bitcode_data) + + bitcode_len_bytes); + plan->store_callback_data = cbdata; break; } case HIPFFT_CB_UNDEFINED: From 24dd73176da2f09bae1a72ee66b6fb6b5fdc13ab Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 14:18:24 -0600 Subject: [PATCH 06/19] reject JIT callbacks + multi-GPU --- .../hipfft/library/src/amd_detail/hipfft.cpp | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index 2c7dbc362815..f42d57749a81 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -845,6 +845,16 @@ struct hipfftHandle_t rocfft_comm_type comm_type = rocfft_comm_none; void* comm_handle = nullptr; + bool jit_load_callback_enabled() const + { + return !load_callback_symbol.empty() && !load_callback_bitcode.empty(); + } + + bool jit_store_callback_enabled() const + { + return !store_callback_symbol.empty() && !store_callback_bitcode.empty(); + } + bool is_valid_for(const hipLibXtDesc& lib_desc, fft_io desc_io_label) const { // Plan must be initialized for multi-device usage @@ -1222,7 +1232,14 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, rocfft_plan_description_set_comm(desc, plan->comm_type, plan->comm_handle)); // set JIT callbacks if specified - if(!plan->load_callback_symbol.empty() && !plan->load_callback_bitcode.empty()) + + // JIT callbacks cannot currently be combined with multi-GPU + // transforms. + if((plan->jit_load_callback_enabled() || plan->jit_store_callback_enabled()) + && plan->device_contexts.size() > 1) + return HIPFFT_NOT_IMPLEMENTED; + + if(plan->jit_load_callback_enabled()) { ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_load_callback(desc, @@ -1231,7 +1248,7 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, plan->load_callback_bitcode.size(), plan->load_callback_lds_bytes)); } - if(!plan->store_callback_symbol.empty() && !plan->store_callback_bitcode.empty()) + if(plan->jit_store_callback_enabled()) { ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_store_callback(desc, @@ -1338,12 +1355,12 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, } // if JIT callbacks are used, pass the cbdata to the execution info - if(!plan->load_callback_symbol.empty() && !plan->load_callback_bitcode.empty()) + if(plan->jit_load_callback_enabled()) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_load_callback_data( plan->info, plan->load_callback_data, plan->device_contexts.size())); } - if(!plan->store_callback_symbol.empty() && plan->store_callback_bitcode.empty()) + if(plan->jit_store_callback_enabled()) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_store_callback_data( plan->info, plan->store_callback_data, plan->device_contexts.size())); From 51254385919cdf639661abc32799b0abb122793c Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 14:19:27 -0600 Subject: [PATCH 07/19] document the inability of JIT callbacks to work with multi-GPU --- projects/hipfft/CHANGELOG.md | 3 +++ .../hipfft/library/include/hipfft/hipfftXt.h | 20 ++++++++++--------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/projects/hipfft/CHANGELOG.md b/projects/hipfft/CHANGELOG.md index 81028ef45638..bf38d2c388a9 100644 --- a/projects/hipfft/CHANGELOG.md +++ b/projects/hipfft/CHANGELOG.md @@ -17,6 +17,9 @@ Documentation for hipFFT is available at On AMD platforms, the device function is provided as SPIR-V. On CUDA platforms, the device function is provided as LTO-IR fatbin. + These APIs are not currently compatible with multi-GPU transforms. This + support will be added in a future release of hipFFT. + ### Changed * Modified the rocFFT backend's implementation details of hipFFT so that cuFFT backend's diff --git a/projects/hipfft/library/include/hipfft/hipfftXt.h b/projects/hipfft/library/include/hipfft/hipfftXt.h index 2c77d915fe67..4d533b674211 100644 --- a/projects/hipfft/library/include/hipfft/hipfftXt.h +++ b/projects/hipfft/library/include/hipfft/hipfftXt.h @@ -163,19 +163,21 @@ typedef void (*hipfftJITCallbackStoreD)( /*! @brief Set a JIT callback on a plan * * @details Set either a load or store callback to run with a plan. - * The type of callback is specified with the 'cbtype' parameter. + * The type of callback is specified with the `cbtype` parameter. * The callback is provided as SPIR-V on AMD platforms and as * LTO-IR fatbin on CUDA platforms. * - * 'cbdata' is an optional array of pointers of data that is passed - * to the callback function, one per visible HIP device. hipFFT may - * execute the callback function on the current HIP device as well - * as any device used by a multi-device transform. The - * corresponding pointer for a device in the array will be passed to - * the callback function's `cbdata` parameter when it is executed on - * that device. + * Note: JIT callbacks cannot currently be used on multi-GPU transforms. + * This support will be added in a future release of hipFFT. * - * If specified, the 'cbdata' array must remain valid until either + * `cbdata` is an optional array of pointers of data that is passed + * to the callback function. As JIT callbacks are currently only + * supported for single-device transforms, this array must have one + * pointer element in it if it is specified. The pointer will be + * passed to the callback function's `cbdata` parameter when it is + * executed on that device. + * + * If specified, the `cbdata` array must remain valid until either * the plan is destroyed or callbacks are cleared from the plan. * * This function must be called after the plan is allocated using From d779180162530221c9321f741bb77b25b8d7803f Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 14:19:46 -0600 Subject: [PATCH 08/19] disable multi-GPU JIT callback tests --- projects/hipfft/clients/tests/multi_device_test.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/projects/hipfft/clients/tests/multi_device_test.cpp b/projects/hipfft/clients/tests/multi_device_test.cpp index 41b21a36f473..9f47ea539b1d 100644 --- a/projects/hipfft/clients/tests/multi_device_test.cpp +++ b/projects/hipfft/clients/tests/multi_device_test.cpp @@ -81,9 +81,10 @@ std::vector param_generator_multi_gpu(const std::optional static const std::vector> stride_range = {{1}}; // function pointer callbacks need -fgpu-rdc, but that causes build - // nondeterminism in kpack + // nondeterminism in kpack. + // JIT callbacks are not yet supported on multi-GPU transforms auto multi_device_callbacks - = {fft_callback_type_none, /*fft_callback_type_funcptr, */ fft_callback_type_jit}; + = {fft_callback_type_none, /*fft_callback_type_funcptr, fft_callback_type_jit*/}; // gather cases to test as single-device params, then distribute // to multiple GPUs From 09ea8e50e25f69afcbd5b1ff68ec6a54b32d2331 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 14:49:52 -0600 Subject: [PATCH 09/19] don't try to precompile JITCB tokens --- projects/hipfft/clients/tests/gtest_main.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/projects/hipfft/clients/tests/gtest_main.cpp b/projects/hipfft/clients/tests/gtest_main.cpp index 584fec3e3780..2a96c87f65ab 100644 --- a/projects/hipfft/clients/tests/gtest_main.cpp +++ b/projects/hipfft/clients/tests/gtest_main.cpp @@ -258,6 +258,12 @@ void precompile_test_kernels(const std::string& precompile_file) { hipfft_params params; params.from_token(token); + + // JIT callbacks would require JIT state to be + // specified which we're not doing here + if(params.run_callbacks == fft_callback_type_jit) + continue; + params.validate(); params.create_plan(); if(params.is_forward()) From 6eceb92d61f76f1ba1fbaf16e6f5db0a30a43b16 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 15:09:01 -0600 Subject: [PATCH 10/19] document cbtype, fix name of doc string for cbdata --- projects/hipfft/library/include/hipfft/hipfftXt.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/projects/hipfft/library/include/hipfft/hipfftXt.h b/projects/hipfft/library/include/hipfft/hipfftXt.h index 4d533b674211..9bdbfd760cba 100644 --- a/projects/hipfft/library/include/hipfft/hipfftXt.h +++ b/projects/hipfft/library/include/hipfft/hipfftXt.h @@ -190,7 +190,8 @@ typedef void (*hipfftJITCallbackStoreD)( * @param[in] symbol_name Name of the symbol in the compiled bitcode. * @param[in] bitcode_data Pointer to bitcode data. * @param[in] bitcode_len_bytes Number of bytes in the bitcode data. - * @param[in] callbackData Array of callback function data pointers. + * @param[in] cbtype Type of callback function. + * @param[in] cbdata Array of callback function data pointers. */ HIPFFT_EXPORT hipfftResult hipfftXtSetJITCallback(hipfftHandle plan, const char* symbol_name, From 23cbaf90ee45b2e3cba7d67f5b29b868df9c8b1a Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 15:09:26 -0600 Subject: [PATCH 11/19] use get_raw_data_ptrs --- projects/hipfft/clients/hipfft_params.h | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/projects/hipfft/clients/hipfft_params.h b/projects/hipfft/clients/hipfft_params.h index 7a1326f2616a..2794c6796c7f 100644 --- a/projects/hipfft/clients/hipfft_params.h +++ b/projects/hipfft/clients/hipfft_params.h @@ -1305,12 +1305,8 @@ class hipfft_params : public fft_params } check_jit_callback_state(); - load_jit_cb_data_ptrs.resize(load_jit_cb_state->data.size()); - std::transform(load_jit_cb_state->data.begin(), - load_jit_cb_state->data.end(), - load_jit_cb_data_ptrs.begin(), - [](gpubuf& buf) { return buf.data(); }); - ret = hipfftXtSetJITCallback(plan, + load_jit_cb_data_ptrs = load_jit_cb_state->get_raw_data_ptrs(); + ret = hipfftXtSetJITCallback(plan, load_jit_cb_state->symbol, load_jit_cb_state->func.data(), load_jit_cb_state->func.size(), @@ -1359,12 +1355,8 @@ class hipfft_params : public fft_params throw std::runtime_error("unsupported data type for store callback"); } } - store_jit_cb_data_ptrs.resize(store_jit_cb_state->data.size()); - std::transform(store_jit_cb_state->data.begin(), - store_jit_cb_state->data.end(), - store_jit_cb_data_ptrs.begin(), - [](gpubuf& buf) { return buf.data(); }); - ret = hipfftXtSetJITCallback(plan, + store_jit_cb_data_ptrs = store_jit_cb_state->get_raw_data_ptrs(); + ret = hipfftXtSetJITCallback(plan, store_jit_cb_state->symbol, store_jit_cb_state->func.data(), store_jit_cb_state->func.size(), From 59bca62831c3629d3740ac6d45737dfc2d5fd02d Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Tue, 25 Aug 2026 15:40:20 -0600 Subject: [PATCH 12/19] introduce callback_info_t, centralize validation of cbtype vs iotype --- .../hipfft/library/src/amd_detail/hipfft.cpp | 237 ++++++++++-------- 1 file changed, 132 insertions(+), 105 deletions(-) diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index f42d57749a81..bcd780a0da45 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -829,32 +829,69 @@ struct hipfftHandle_t rocfft_execution_info_wrapper_t info; std::vector device_contexts; - void** load_callback_ptrs = nullptr; - std::string load_callback_symbol; - std::vector load_callback_bitcode; - void** load_callback_data = nullptr; - size_t load_callback_lds_bytes = 0; - - void** store_callback_ptrs = nullptr; - std::string store_callback_symbol; - std::vector store_callback_bitcode; - void** store_callback_data = nullptr; - size_t store_callback_lds_bytes = 0; + struct callback_info_t + { + hipfftXtCallbackType type = HIPFFT_CB_UNDEFINED; + void** funcptrs = nullptr; + std::string symbol; + std::vector bitcode; + void** data = nullptr; + size_t lds_bytes = 0; + + bool jit_enabled() const + { + return !symbol.empty() && !bitcode.empty(); + } + + // Throws HIPFFT_INVALID_VALUE if the callback type is not + // valid for the precision and I/O + static void validate_type(hipfftXtCallbackType cbtype, hipfftIOType io_type) + { + switch(cbtype) + { + case HIPFFT_CB_LD_COMPLEX: + if(io_type.precision() != rocfft_precision_single || io_type.is_real_to_complex()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_LD_COMPLEX_DOUBLE: + if(io_type.precision() != rocfft_precision_double || io_type.is_real_to_complex()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_LD_REAL: + if(io_type.precision() != rocfft_precision_single || !io_type.is_real_to_complex()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_LD_REAL_DOUBLE: + if(io_type.precision() != rocfft_precision_double || !io_type.is_real_to_complex()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_ST_COMPLEX: + if(io_type.precision() != rocfft_precision_single || io_type.is_complex_to_real()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_ST_COMPLEX_DOUBLE: + if(io_type.precision() != rocfft_precision_double || io_type.is_complex_to_real()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_ST_REAL: + if(io_type.precision() != rocfft_precision_single || !io_type.is_complex_to_real()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_ST_REAL_DOUBLE: + if(io_type.precision() != rocfft_precision_double || !io_type.is_complex_to_real()) + throw HIPFFT_INVALID_VALUE; + break; + case HIPFFT_CB_UNDEFINED: + throw HIPFFT_INVALID_VALUE; + } + } + }; + callback_info_t load_callback, store_callback; // Multi-processing communicator rocfft_comm_type comm_type = rocfft_comm_none; void* comm_handle = nullptr; - bool jit_load_callback_enabled() const - { - return !load_callback_symbol.empty() && !load_callback_bitcode.empty(); - } - - bool jit_store_callback_enabled() const - { - return !store_callback_symbol.empty() && !store_callback_bitcode.empty(); - } - bool is_valid_for(const hipLibXtDesc& lib_desc, fft_io desc_io_label) const { // Plan must be initialized for multi-device usage @@ -1235,27 +1272,29 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, // JIT callbacks cannot currently be combined with multi-GPU // transforms. - if((plan->jit_load_callback_enabled() || plan->jit_store_callback_enabled()) + if((plan->load_callback.jit_enabled() || plan->store_callback.jit_enabled()) && plan->device_contexts.size() > 1) return HIPFFT_NOT_IMPLEMENTED; - - if(plan->jit_load_callback_enabled()) + if(plan->load_callback.jit_enabled()) { + hipfftHandle_t::callback_info_t::validate_type(plan->load_callback.type, plan->io_type); ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_load_callback(desc, - plan->load_callback_symbol.c_str(), - plan->load_callback_bitcode.data(), - plan->load_callback_bitcode.size(), - plan->load_callback_lds_bytes)); + plan->load_callback.symbol.c_str(), + plan->load_callback.bitcode.data(), + plan->load_callback.bitcode.size(), + plan->load_callback.lds_bytes)); } - if(plan->jit_store_callback_enabled()) + if(plan->store_callback.jit_enabled()) { + hipfftHandle_t::callback_info_t::validate_type(plan->store_callback.type, + plan->io_type); ROCFFT_EXPECT_SUCCESS( rocfft_plan_description_set_store_callback(desc, - plan->store_callback_symbol.c_str(), - plan->store_callback_bitcode.data(), - plan->store_callback_bitcode.size(), - plan->store_callback_lds_bytes)); + plan->store_callback.symbol.c_str(), + plan->store_callback.bitcode.data(), + plan->store_callback.bitcode.size(), + plan->store_callback.lds_bytes)); } if(plan->device_contexts.size() > 1) @@ -1355,15 +1394,15 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, } // if JIT callbacks are used, pass the cbdata to the execution info - if(plan->jit_load_callback_enabled()) + if(plan->load_callback.jit_enabled()) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_load_callback_data( - plan->info, plan->load_callback_data, plan->device_contexts.size())); + plan->info, plan->load_callback.data, plan->device_contexts.size())); } - if(plan->jit_store_callback_enabled()) + if(plan->store_callback.jit_enabled()) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_store_callback_data( - plan->info, plan->store_callback_data, plan->device_contexts.size())); + plan->info, plan->store_callback.data, plan->device_contexts.size())); } return HIPFFT_SUCCESS; @@ -2066,6 +2105,8 @@ try if(!plan) return HIPFFT_INVALID_PLAN; + hipfftHandle_t::callback_info_t::validate_type(cbtype, plan->io_type); + // check that the input/output type matches what's being requested // // NOTE: cufft explicitly does not save shared memory bytes when @@ -2074,68 +2115,52 @@ try switch(cbtype) { case HIPFFT_CB_LD_COMPLEX: - if(plan->io_type.precision() != rocfft_precision_single - || plan->io_type.is_real_to_complex()) - return HIPFFT_INVALID_VALUE; - plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; - plan->load_callback_lds_bytes = 0; + plan->load_callback.type = cbtype; + plan->load_callback.funcptrs = callbacks; + plan->load_callback.data = callbackData; + plan->load_callback.lds_bytes = 0; break; case HIPFFT_CB_LD_COMPLEX_DOUBLE: - if(plan->io_type.precision() != rocfft_precision_double - || plan->io_type.is_real_to_complex()) - return HIPFFT_INVALID_VALUE; - plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; - plan->load_callback_lds_bytes = 0; + plan->load_callback.type = cbtype; + plan->load_callback.funcptrs = callbacks; + plan->load_callback.data = callbackData; + plan->load_callback.lds_bytes = 0; break; case HIPFFT_CB_LD_REAL: - if(plan->io_type.precision() != rocfft_precision_single - || !plan->io_type.is_real_to_complex()) - return HIPFFT_INVALID_VALUE; - plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; - plan->load_callback_lds_bytes = 0; + plan->load_callback.type = cbtype; + plan->load_callback.funcptrs = callbacks; + plan->load_callback.data = callbackData; + plan->load_callback.lds_bytes = 0; break; case HIPFFT_CB_LD_REAL_DOUBLE: - if(plan->io_type.precision() != rocfft_precision_double - || !plan->io_type.is_real_to_complex()) - return HIPFFT_INVALID_VALUE; - plan->load_callback_ptrs = callbacks; - plan->load_callback_data = callbackData; - plan->load_callback_lds_bytes = 0; + plan->load_callback.type = cbtype; + plan->load_callback.funcptrs = callbacks; + plan->load_callback.data = callbackData; + plan->load_callback.lds_bytes = 0; break; case HIPFFT_CB_ST_COMPLEX: - if(plan->io_type.precision() != rocfft_precision_single - || plan->io_type.is_complex_to_real()) - return HIPFFT_INVALID_VALUE; - plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; - plan->store_callback_lds_bytes = 0; + plan->store_callback.type = cbtype; + plan->store_callback.funcptrs = callbacks; + plan->store_callback.data = callbackData; + plan->store_callback.lds_bytes = 0; break; case HIPFFT_CB_ST_COMPLEX_DOUBLE: - if(plan->io_type.precision() != rocfft_precision_double - || plan->io_type.is_complex_to_real()) - return HIPFFT_INVALID_VALUE; - plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; - plan->store_callback_lds_bytes = 0; + plan->store_callback.type = cbtype; + plan->store_callback.funcptrs = callbacks; + plan->store_callback.data = callbackData; + plan->store_callback.lds_bytes = 0; break; case HIPFFT_CB_ST_REAL: - if(plan->io_type.precision() != rocfft_precision_single - || !plan->io_type.is_complex_to_real()) - return HIPFFT_INVALID_VALUE; - plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; - plan->store_callback_lds_bytes = 0; + plan->store_callback.type = cbtype; + plan->store_callback.funcptrs = callbacks; + plan->store_callback.data = callbackData; + plan->store_callback.lds_bytes = 0; break; case HIPFFT_CB_ST_REAL_DOUBLE: - if(plan->io_type.precision() != rocfft_precision_double - || !plan->io_type.is_complex_to_real()) - return HIPFFT_INVALID_VALUE; - plan->store_callback_ptrs = callbacks; - plan->store_callback_data = callbackData; - plan->store_callback_lds_bytes = 0; + plan->store_callback.type = cbtype; + plan->store_callback.funcptrs = callbacks; + plan->store_callback.data = callbackData; + plan->store_callback.lds_bytes = 0; break; case HIPFFT_CB_UNDEFINED: return HIPFFT_INVALID_VALUE; @@ -2143,15 +2168,15 @@ try rocfft_status res; res = rocfft_execution_info_set_load_callback(plan->info, - plan->load_callback_ptrs, - plan->load_callback_data, - plan->load_callback_lds_bytes); + plan->load_callback.funcptrs, + plan->load_callback.data, + plan->load_callback.lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; res = rocfft_execution_info_set_store_callback(plan->info, - plan->store_callback_ptrs, - plan->store_callback_data, - plan->store_callback_lds_bytes); + plan->store_callback.funcptrs, + plan->store_callback.data, + plan->store_callback.lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; return HIPFFT_SUCCESS; @@ -2184,13 +2209,13 @@ try case HIPFFT_CB_LD_COMPLEX_DOUBLE: case HIPFFT_CB_LD_REAL: case HIPFFT_CB_LD_REAL_DOUBLE: - plan->load_callback_lds_bytes = sharedSize; + plan->load_callback.lds_bytes = sharedSize; break; case HIPFFT_CB_ST_COMPLEX: case HIPFFT_CB_ST_COMPLEX_DOUBLE: case HIPFFT_CB_ST_REAL: case HIPFFT_CB_ST_REAL_DOUBLE: - plan->store_callback_lds_bytes = sharedSize; + plan->store_callback.lds_bytes = sharedSize; break; case HIPFFT_CB_UNDEFINED: return HIPFFT_INVALID_VALUE; @@ -2198,15 +2223,15 @@ try rocfft_status res; res = rocfft_execution_info_set_load_callback(plan->info, - plan->load_callback_ptrs, - plan->load_callback_data, - plan->load_callback_lds_bytes); + plan->load_callback.funcptrs, + plan->load_callback.data, + plan->load_callback.lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; res = rocfft_execution_info_set_store_callback(plan->info, - plan->store_callback_ptrs, - plan->store_callback_data, - plan->store_callback_lds_bytes); + plan->store_callback.funcptrs, + plan->store_callback.data, + plan->store_callback.lds_bytes); if(res != rocfft_status_success) return HIPFFT_INVALID_VALUE; return HIPFFT_SUCCESS; @@ -2234,11 +2259,12 @@ try case HIPFFT_CB_LD_REAL: case HIPFFT_CB_LD_REAL_DOUBLE: { - plan->load_callback_symbol = symbol_name; - plan->load_callback_bitcode.assign(static_cast(bitcode_data), + plan->load_callback.type = cbtype; + plan->load_callback.symbol = symbol_name; + plan->load_callback.bitcode.assign(static_cast(bitcode_data), static_cast(bitcode_data) + bitcode_len_bytes); - plan->load_callback_data = cbdata; + plan->load_callback.data = cbdata; break; } case HIPFFT_CB_ST_COMPLEX: @@ -2246,11 +2272,12 @@ try case HIPFFT_CB_ST_REAL: case HIPFFT_CB_ST_REAL_DOUBLE: { - plan->store_callback_symbol = symbol_name; - plan->store_callback_bitcode.assign(static_cast(bitcode_data), + plan->store_callback.type = cbtype; + plan->store_callback.symbol = symbol_name; + plan->store_callback.bitcode.assign(static_cast(bitcode_data), static_cast(bitcode_data) + bitcode_len_bytes); - plan->store_callback_data = cbdata; + plan->store_callback.data = cbdata; break; } case HIPFFT_CB_UNDEFINED: From c7dfe1f8fa46d128ac5ef2eb5b9da7e0c976cca3 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Wed, 26 Aug 2026 15:59:20 -0600 Subject: [PATCH 13/19] check plan initialized in XtSetCallback --- projects/hipfft/library/src/amd_detail/hipfft.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index bcd780a0da45..bee49ef2e438 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -2102,7 +2102,7 @@ hipfftResult hipfftXtSetCallback(hipfftHandle plan, void** callbackData) try { - if(!plan) + if(!plan || !plan->initialized()) return HIPFFT_INVALID_PLAN; hipfftHandle_t::callback_info_t::validate_type(cbtype, plan->io_type); From d57691e13fc393fa5ee430b5fb85c61cb80b37ed Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Wed, 26 Aug 2026 15:59:30 -0600 Subject: [PATCH 14/19] disallow JIT callbacks with MPI --- projects/hipfft/library/src/amd_detail/hipfft.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index bee49ef2e438..2298c3a63e44 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -1272,8 +1272,13 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, // JIT callbacks cannot currently be combined with multi-GPU // transforms. + const bool is_multi_gpu_transform = plan->device_contexts.size() > 1 +#ifdef HIPFFT_MPI_ENABLE + || plan->mp_input_brick || plan->mp_output_brick +#endif + ; if((plan->load_callback.jit_enabled() || plan->store_callback.jit_enabled()) - && plan->device_contexts.size() > 1) + && is_multi_gpu_transform) return HIPFFT_NOT_IMPLEMENTED; if(plan->load_callback.jit_enabled()) { From f2347cf801de6da2a1272fba906f6dc0d136ba2b Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Thu, 27 Aug 2026 13:45:51 -0600 Subject: [PATCH 15/19] doc fixes: - consistently say SPIR-V is bitcode - clang++ -> amdclang++ - add section with description and example for passing data to callbacks - align load/store API docs --- .../docs/how-to/load-store-callbacks.rst | 76 +++++++++++++++++-- .../rocfft/library/include/rocfft/rocfft.h | 7 +- 2 files changed, 72 insertions(+), 11 deletions(-) diff --git a/projects/rocfft/docs/how-to/load-store-callbacks.rst b/projects/rocfft/docs/how-to/load-store-callbacks.rst index 6e6f2d669222..515a6238e265 100644 --- a/projects/rocfft/docs/how-to/load-store-callbacks.rst +++ b/projects/rocfft/docs/how-to/load-store-callbacks.rst @@ -97,15 +97,15 @@ Compiling functions to SPIR-V for JIT callbacks :cpp:func:`rocfft_plan_description_set_load_callback` and :cpp:func:`rocfft_plan_description_set_store_callback` accept -callback functions as a named symbol in compiled SPIR-V code. +callback functions as a named symbol in compiled SPIR-V bitcode. Symbol names can only contain digits (0-9), letters (a-z, A-Z), and underscores, and cannot begin with a digit. A callback function written as HIP code must first be compiled to -SPIR-V before it can be added to a plan description. The following -example demonstrates how to compile such code using the ``clang++`` -compiler. +SPIR-V bitcode before it can be added to a plan description. The +following example demonstrates how to compile such code using the +``amdclang++`` compiler. An example load callback function for a single-precision real-complex forward transform might look like: @@ -141,19 +141,19 @@ the file and the file's length are then passed to rocFFT: rocfft_plan_description create_plan_desc_with_callback() { // Read the compiled callback into a vector - std::vector code; + std::vector bitcode; std::ifstream infile("load_callback.spv", std::ios::binary | std::ios::ate); auto size = infile.tellg(); - code.resize(size); + bitcode.resize(size); infile.seekg(0); - infile.read(code.data(), size); + infile.read(bitcode.data(), size); // Create a plan description and set the load callback rocfft_plan_description desc = nullptr; if(rocfft_plan_description_create(&desc) != rocfft_status_success) return nullptr; if(rocfft_plan_description_set_load_callback(desc, "load_callback", - code.data(), size, 0) != rocfft_status_success) + bitcode.data(), size, 0) != rocfft_status_success) { rocfft_plan_description_destroy(desc); return nullptr; @@ -161,6 +161,66 @@ the file and the file's length are then passed to rocFFT: return desc; } +Passing data to callback functions +---------------------------------- + +rocFFT can optionally pass a user-specified pointer value to callback +functions. This is useful in cases where the callback function +requires extra data on top of the input/output pointer and offset +that are already provided. + +Callback data is specified on a rocFFT execution info object using +:cpp:func:`rocfft_execution_info_set_load_callback_data` and +:cpp:func:`rocfft_execution_info_set_store_callback_data` for load +and store callbacks, respectively. + +These functions accept an array of callback data pointers, one per +brick in the input fields (for load callbacks) or output fields (for +store callbacks) of the transform. A transform which does not +specify a field and brick layout for input (or output) is +considered to have a single brick for input (or output). + +Here is an example showing how to pass filtering data to a load +callback. + +.. code-block:: c++ + + // Define a structure to hold arbitrary amounts of data to pass to + // the callback function. This example has just one data member + // but it could be extended with additional data members. + struct load_callback_data + { + hipDoubleComplex* filter = nullptr; + }; + + // Initialize the structure on the host + load_callback_data cbdata_host; + + // Set the filter pointer in the host structure. Code to allocate and + // initialize this filter on the device has been omitted but would + // depend on the details of the filtering operation. + cbdata_host.filter = device_filter; + + // Copy the structure to the device + load_callback_data* cbdata_device = nullptr; + hipMalloc(&cbdata_device, sizeof(load_callback_data)); + hipMemcpy(cbdata_device, &cbdata_host, sizeof(load_callback_data), hipMemcpyHostToDevice); + + // Initialize an array of device pointers on the host. This example + // creates an array of length 1 as the input has only one brick. + void* cbdata_ptrs[1]; + cbdata_ptrs[0] = cbdata_device; + + // Create an execution info object and set the device pointer array on it. + rocfft_execution_info info = nullptr; + rocfft_execution_info_create(&info); + rocfft_execution_info_set_load_callback_data(info, cbdata_ptrs, 1); + + // When the execution info is passed to rocfft_execute, the load + // callback receives the cbdata_device pointer as its callback_data + // parameter. The callback can then cast that pointer from 'void*' to + // 'load_callback_data*' and access the filter. + Legacy function pointer callbacks (deprecated) ---------------------------------------------- diff --git a/projects/rocfft/library/include/rocfft/rocfft.h b/projects/rocfft/library/include/rocfft/rocfft.h index dd231c0e2191..f055d90c64c6 100644 --- a/projects/rocfft/library/include/rocfft/rocfft.h +++ b/projects/rocfft/library/include/rocfft/rocfft.h @@ -443,6 +443,7 @@ ROCFFT_EXPORT rocfft_status * * ::rocfft_execution_info_set_store_callback_data can optionally be * used to set the `cbdata` value received by the callback function. + * rocFFT will pass nullptr for `cbdata` by default. * * Currently, `shared_mem_bytes` must be 0. Callbacks are not * supported on transforms that use planar formats for either input @@ -803,9 +804,9 @@ ROCFFT_EXPORT rocfft_status rocfft_execution_info_set_store_callback(rocfft_exec * used at plan creation. * * The provided data pointers replace any previously-specified - * store callback data for this execution info handle. `cb_data` may be - * nullptr, to indicate that all store callbacks should receive - * nullptr for their `cbdata` parameters. + * store callback data for this execution info handle. `cb_data` may + * be nullptr along with a count of zero, to indicate that all store + * callbacks should receive nullptr for their `cbdata` parameters. * * @param[in] info execution info handle * @param[in] cb_data callback function data, passed to the JIT-compiled store callback when it is called From 5ba45d6512a986632dcfad5944bce0e7d02847e8 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Thu, 27 Aug 2026 15:30:55 -0600 Subject: [PATCH 16/19] review fixes: - don't talk about plan descriptions in changelog - clarify deprecation message - allow clearing JIT callbacks --- projects/hipfft/CHANGELOG.md | 3 ++- .../hipfft/library/include/hipfft/hipfftXt.h | 10 ++++++-- .../hipfft/library/src/amd_detail/hipfft.cpp | 24 ++++++++++++++----- 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/projects/hipfft/CHANGELOG.md b/projects/hipfft/CHANGELOG.md index bf38d2c388a9..4ab74d26eed9 100644 --- a/projects/hipfft/CHANGELOG.md +++ b/projects/hipfft/CHANGELOG.md @@ -34,7 +34,8 @@ Documentation for hipFFT is available at * The `hipfftXtSetCallback` and `hipfftXtClearCallback` APIs are now deprecated and will be removed in a future release. They allow for specifying callbacks as device function pointers at plan execution time, but rocFFT cannot - optimize the combined code. Instead, users should specify JIT callbacks on plan descriptions. + optimize the combined code. Instead, users should specify JIT callbacks with `hipfftXtSetJITCallback` before + initializing the plan. ## hipFFT 1.0.25 for ROCm 10.0 diff --git a/projects/hipfft/library/include/hipfft/hipfftXt.h b/projects/hipfft/library/include/hipfft/hipfftXt.h index 9bdbfd760cba..f31d7a57aad2 100644 --- a/projects/hipfft/library/include/hipfft/hipfftXt.h +++ b/projects/hipfft/library/include/hipfft/hipfftXt.h @@ -101,7 +101,8 @@ typedef void (*hipfftCallbackStoreD)( * @param[in] cbtype Type of callback being set. * @param[in] callbackData Array of callback function data pointers */ -HIPFFT_DEPRECATED_MSG("Use hipfftXtSetJITCallback instead.") +HIPFFT_DEPRECATED_MSG( + "Function pointer callbacks are deprecated; use hipfftXtSetJITCallback for JIT callbacks.") HIPFFT_EXPORT hipfftResult hipfftXtSetCallback(hipfftHandle plan, void** callbacks, hipfftXtCallbackType cbtype, @@ -114,7 +115,8 @@ HIPFFT_EXPORT hipfftResult hipfftXtSetCallback(hipfftHandle plan, * @param[in] plan The FFT plan. * @param[in] cbtype Type of callback being removed. */ -HIPFFT_DEPRECATED_MSG("Use hipfftXtSetJITCallback instead.") +HIPFFT_DEPRECATED_MSG( + "Function pointer callbacks are deprecated; use hipfftXtSetJITCallback for JIT callbacks.") HIPFFT_EXPORT hipfftResult hipfftXtClearCallback(hipfftHandle plan, hipfftXtCallbackType cbtype); /*! @brief Set shared memory size for callback. @@ -170,6 +172,10 @@ typedef void (*hipfftJITCallbackStoreD)( * Note: JIT callbacks cannot currently be used on multi-GPU transforms. * This support will be added in a future release of hipFFT. * + * JIT callbacks can be cleared from the plan by passing nullptr + * for `symbol_name` and `bitcode_data`, and 0 for + * `bitcode_len_bytes`. + * * `cbdata` is an optional array of pointers of data that is passed * to the callback function. As JIT callbacks are currently only * supported for single-device transforms, this array must have one diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index 2298c3a63e44..d2ce9d935562 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -1399,12 +1399,12 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, } // if JIT callbacks are used, pass the cbdata to the execution info - if(plan->load_callback.jit_enabled()) + if(plan->load_callback.jit_enabled() && plan->load_callback.data) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_load_callback_data( plan->info, plan->load_callback.data, plan->device_contexts.size())); } - if(plan->store_callback.jit_enabled()) + if(plan->store_callback.jit_enabled() && plan->store_callback.data) { ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_store_callback_data( plan->info, plan->store_callback.data, plan->device_contexts.size())); @@ -2257,6 +2257,12 @@ try if(!plan) return HIPFFT_INVALID_PLAN; + // Return an error if we're neither clearing callbacks nor setting + // symbol + bitcode together. + if(!(!symbol_name && !bitcode_data && !bitcode_len_bytes) + && !(symbol_name && bitcode_data && bitcode_len_bytes)) + return HIPFFT_INVALID_VALUE; + switch(cbtype) { case HIPFFT_CB_LD_COMPLEX: @@ -2264,8 +2270,11 @@ try case HIPFFT_CB_LD_REAL: case HIPFFT_CB_LD_REAL_DOUBLE: { - plan->load_callback.type = cbtype; - plan->load_callback.symbol = symbol_name; + plan->load_callback.type = cbtype; + if(symbol_name) + plan->load_callback.symbol = symbol_name; + else + plan->load_callback.symbol.clear(); plan->load_callback.bitcode.assign(static_cast(bitcode_data), static_cast(bitcode_data) + bitcode_len_bytes); @@ -2277,8 +2286,11 @@ try case HIPFFT_CB_ST_REAL: case HIPFFT_CB_ST_REAL_DOUBLE: { - plan->store_callback.type = cbtype; - plan->store_callback.symbol = symbol_name; + plan->store_callback.type = cbtype; + if(symbol_name) + plan->store_callback.symbol = symbol_name; + else + plan->store_callback.symbol.clear(); plan->store_callback.bitcode.assign(static_cast(bitcode_data), static_cast(bitcode_data) + bitcode_len_bytes); From 79945c3ed9630d3e64bb684e7114bd874e0ef9f7 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Thu, 27 Aug 2026 16:02:41 -0600 Subject: [PATCH 17/19] ensure plan is not initialized when setting JIT callbacks --- projects/hipfft/library/src/amd_detail/hipfft.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/hipfft/library/src/amd_detail/hipfft.cpp b/projects/hipfft/library/src/amd_detail/hipfft.cpp index d2ce9d935562..72349796267f 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -2254,7 +2254,7 @@ hipfftResult hipfftXtSetJITCallback(hipfftHandle plan, void** cbdata) try { - if(!plan) + if(!plan || plan->initialized()) return HIPFFT_INVALID_PLAN; // Return an error if we're neither clearing callbacks nor setting From 0d5500de3cba530bbd53f38f0bbc4ba72d8637cb Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Fri, 28 Aug 2026 22:33:04 -0600 Subject: [PATCH 18/19] add doc note about cbdata array only being length-1 for now --- projects/rocfft/docs/how-to/load-store-callbacks.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/projects/rocfft/docs/how-to/load-store-callbacks.rst b/projects/rocfft/docs/how-to/load-store-callbacks.rst index 515a6238e265..1689673d8112 100644 --- a/projects/rocfft/docs/how-to/load-store-callbacks.rst +++ b/projects/rocfft/docs/how-to/load-store-callbacks.rst @@ -180,6 +180,13 @@ store callbacks) of the transform. A transform which does not specify a field and brick layout for input (or output) is considered to have a single brick for input (or output). +.. note:: + As JIT callbacks cannot currently be used on transforms that have + fields or bricks specified on the plan description, the length of the + array of callback data pointers will always be 1 if callback data is + specified. + + Here is an example showing how to pass filtering data to a load callback. From b3ffb6b326488562758ca521cd9e43bcfaced3a7 Mon Sep 17 00:00:00 2001 From: Steve Leung Date: Mon, 31 Aug 2026 08:55:05 -0600 Subject: [PATCH 19/19] disable callback tests Also put the multi-GPU + callback kernels in the same place as the other coverage kernels, since the db is designed to be appended to. Might save some compilation time. --- projects/hipfft/clients/tests/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/projects/hipfft/clients/tests/CMakeLists.txt b/projects/hipfft/clients/tests/CMakeLists.txt index cc607da3316a..9b0a702cf657 100644 --- a/projects/hipfft/clients/tests/CMakeLists.txt +++ b/projects/hipfft/clients/tests/CMakeLists.txt @@ -423,7 +423,7 @@ if (WIN32) endif() option(BUILD_CODE_COVERAGE "Build with code coverage flags (clang only)" OFF) -set(COVERAGE_TEST_OPTIONS "--smoketest;--gtest_filter=-*call*" CACHE STRING "Command line arguments for hipfft-test when generating a code coverage report (Note: an additional run of hipfft-test targeting multi_gpu* and callback* tests is always executed and coverage results are aggregated)") +set(COVERAGE_TEST_OPTIONS "--smoketest;--gtest_filter=-*call*" CACHE STRING "Command line arguments for hipfft-test when generating a code coverage report (Note: an additional run of hipfft-test targeting multi_gpu* is always executed and coverage results are aggregated)") if (BUILD_CODE_COVERAGE) add_custom_target( code_cov_tests @@ -431,7 +431,7 @@ if (BUILD_CODE_COVERAGE) COMMAND ${CMAKE_COMMAND} -E rm -rf ./coverage-report COMMAND ${CMAKE_COMMAND} -E make_directory ./coverage-report/profraw COMMAND ${CMAKE_COMMAND} -E env LLVM_PROFILE_FILE="./coverage-report/profraw/hipfft-coverage_%p.profraw" GTEST_LISTENER=NO_PASS_LINE_IN_LOG $ --precompile=./clients/staging/hipfft-test-precompile.db ${COVERAGE_TEST_OPTIONS} - COMMAND ${CMAKE_COMMAND} -E env LLVM_PROFILE_FILE="./coverage-report/profraw/hipfft-coverage_%p.profraw" GTEST_LISTENER=NO_PASS_LINE_IN_LOG $ --precompile=./clients/staging/hipfft-test-precompile-multi_gpu-plus-callback.db --gtest_filter=multi_gpu*:callback* + COMMAND ${CMAKE_COMMAND} -E env LLVM_PROFILE_FILE="./coverage-report/profraw/hipfft-coverage_%p.profraw" GTEST_LISTENER=NO_PASS_LINE_IN_LOG $ --precompile=./clients/staging/hipfft-test-precompile.db --gtest_filter=multi_gpu* WORKING_DIRECTORY ${CMAKE_BINARY_DIR} )