diff --git a/projects/hipfft/CHANGELOG.md b/projects/hipfft/CHANGELOG.md index 7b6737ea876a..f9a21d33d6cb 100644 --- a/projects/hipfft/CHANGELOG.md +++ b/projects/hipfft/CHANGELOG.md @@ -9,6 +9,17 @@ 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. + + 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 @@ -19,6 +30,13 @@ 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 with `hipfftXtSetJITCallback` before + initializing the plan. + ### Resolved issues * Fixed a hang when creating a plan with a zero length or zero batch. These now return `HIPFFT_INVALID_SIZE`. diff --git a/projects/hipfft/clients/hipfft_params.h b/projects/hipfft/clients/hipfft_params.h index 602a72dd20d6..2794c6796c7f 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; @@ -1256,8 +1261,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) @@ -1302,14 +1305,13 @@ class hipfft_params : public fft_params } check_jit_callback_state(); - 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(), 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; @@ -1353,16 +1355,14 @@ class hipfft_params : public fft_params throw std::runtime_error("unsupported data type for store callback"); } } - 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(), 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; -#endif } // call hipfftCreate + hipfftMake* functions, inserting calls to 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} ) 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/gtest_main.cpp b/projects/hipfft/clients/tests/gtest_main.cpp index a2e60332c29c..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()) @@ -387,7 +393,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"; diff --git a/projects/hipfft/clients/tests/multi_device_test.cpp b/projects/hipfft/clients/tests/multi_device_test.cpp index 613a96239262..9f47ea539b1d 100644 --- a/projects/hipfft/clients/tests/multi_device_test.cpp +++ b/projects/hipfft/clients/tests/multi_device_test.cpp @@ -81,8 +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 - auto multi_device_callbacks = {fft_callback_type_none, /*fft_callback_type_funcptr, */}; + // 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*/}; // 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/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 5642a87e57c2..f31d7a57aad2 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,30 +85,38 @@ 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. + * + * 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. * @param[in] callbackData Array of callback function data pointers */ +HIPFFT_DEPRECATED_MSG( + "Function pointer callbacks are deprecated; use hipfftXtSetJITCallback for JIT callbacks.") HIPFFT_EXPORT hipfftResult hipfftXtSetCallback(hipfftHandle plan, void** callbacks, 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. */ +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. @@ -125,6 +133,79 @@ 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. + * + * 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 + * 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 + * ::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] cbtype Type of callback function. + * @param[in] cbdata 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 5be7c7137b84..1415c94605c5 100644 --- a/projects/hipfft/library/src/amd_detail/hipfft.cpp +++ b/projects/hipfft/library/src/amd_detail/hipfft.cpp @@ -829,12 +829,64 @@ 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; + 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; @@ -1223,6 +1275,40 @@ 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 + + // 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()) + && is_multi_gpu_transform) + return HIPFFT_NOT_IMPLEMENTED; + 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)); + } + 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)); + } + if(plan->device_contexts.size() > 1) { // Determine the sub-formats of this rocfft plan's input and output fields. Unbatched @@ -1319,6 +1405,18 @@ static hipfftResult hipfftMakePlan_internal(hipfftHandle plan, } } + // if JIT callbacks are used, pass the cbdata to the execution info + 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() && plan->store_callback.data) + { + ROCFFT_EXPECT_SUCCESS(rocfft_execution_info_set_store_callback_data( + plan->info, plan->store_callback.data, plan->device_contexts.size())); + } + return HIPFFT_SUCCESS; } @@ -2016,9 +2114,11 @@ 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); + // check that the input/output type matches what's being requested // // NOTE: cufft explicitly does not save shared memory bytes when @@ -2027,68 +2127,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; @@ -2096,15 +2180,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; @@ -2137,13 +2221,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; @@ -2151,15 +2235,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; @@ -2169,6 +2253,67 @@ 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 || plan->initialized()) + 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: + case HIPFFT_CB_LD_COMPLEX_DOUBLE: + case HIPFFT_CB_LD_REAL: + case HIPFFT_CB_LD_REAL_DOUBLE: + { + 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); + plan->load_callback.data = 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.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); + plan->store_callback.data = 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, diff --git a/projects/rocfft/docs/how-to/load-store-callbacks.rst b/projects/rocfft/docs/how-to/load-store-callbacks.rst index 6e6f2d669222..1689673d8112 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,73 @@ 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). + +.. 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. + +.. 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