From 9a11381780b8d5b6bf6ed1b31ed96f82ad790c10 Mon Sep 17 00:00:00 2001 From: hbhu_bin Date: Mon, 7 Jul 2025 17:34:56 +0800 Subject: [PATCH] support ASDSIP lib --- .jenkins/check/config/filter_cpplint.txt | 1 + cmake/ascend_variables.cmake | 4 +- cmake/package.cmake | 7 + cmake/package_plugin.cmake | 7 + mindspore/ccsrc/ms_extension/api.h | 3 + .../ccsrc/ms_extension/ascend/CMakeLists.txt | 5 + .../ms_extension/ascend/asdsip/CMakeLists.txt | 9 + .../ascend/asdsip/asdsip_common.cc | 128 ++++++++++++ .../ascend/asdsip/asdsip_common.h | 87 ++++++++ .../ms_extension/ascend/asdsip/asdsip_utils.h | 191 ++++++++++++++++++ .../pynative/pyboost_extension.cc | 7 - .../ms_extension/pynative/pyboost_extension.h | 7 +- .../mindspore/ops/operations/custom_ops.py | 14 ++ .../custom/jit_test_files/asdsip_fftc2c.cpp | 44 ++++ .../custom/test_custom_pyboost_ascend.py | 21 ++ 15 files changed, 526 insertions(+), 9 deletions(-) create mode 100644 mindspore/ccsrc/ms_extension/ascend/asdsip/CMakeLists.txt create mode 100644 mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.cc create mode 100644 mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.h create mode 100644 mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_utils.h create mode 100644 tests/st/graph_kernel/custom/jit_test_files/asdsip_fftc2c.cpp diff --git a/.jenkins/check/config/filter_cpplint.txt b/.jenkins/check/config/filter_cpplint.txt index 33877bc771a..da048616a08 100644 --- a/.jenkins/check/config/filter_cpplint.txt +++ b/.jenkins/check/config/filter_cpplint.txt @@ -85,6 +85,7 @@ "mindspore/mindspore/ops/kernel/ascend/pyboost/customize/conv2d_padding.h" "runtime/references" "mindspore/mindspore/ops/kernel/ascend/pyboost/customize/conv3d_padding.cc" "runtime/references" "mindspore/mindspore/ops/kernel/ascend/pyboost/customize/conv3d_padding.h" "runtime/references" +"mindspore/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_utils.h" "runtime/references" "mindspore/mindspore/ccsrc/plugin/res_manager/ascend/symbol_interface/acl_tdt_symbol.cc" "build/include_subdir" "mindspore/mindspore/ccsrc/plugin/res_manager/ascend/symbol_interface/acl_prof_symbol.cc" "build/include_subdir" "mindspore/mindspore/ccsrc/plugin/res_manager/ascend/symbol_interface/acl_rt_symbol.cc" "build/include_subdir" diff --git a/cmake/ascend_variables.cmake b/cmake/ascend_variables.cmake index ffee8fe7456..eae3699205e 100644 --- a/cmake/ascend_variables.cmake +++ b/cmake/ascend_variables.cmake @@ -26,7 +26,7 @@ set(ASCEND_TOOLKIT_PLUGIN_PATH ${ASCEND_TOOLKIT_RUNTIME_PATH}/plugin/opskernel) set(ASCEND_NNAE_RUNTIME_PATH ${ASCEND_PATH}/nnae/latest/lib64) set(ASCEND_NNAE_OPP_PATH ${ASCEND_PATH}/nnae/latest/opp/built-in/op_impl/ai_core/tbe/op_tiling) -# nnal packages (for ATB kernel) +# nnal packages (for ATB kernel and ASDSIP kernel) if(DEFINED ENV{ASCEND_CUSTOM_PATH}) set(ASCEND_NNAL_RUNTIME_PATH ${ASCEND_PATH}/../nnal/) else() @@ -40,3 +40,5 @@ if(NOT ENABLE_GLIBCXX) add_compile_definitions(ENABLE_ATB) endif() endif() + +set(ASCEND_NNAL_ASDSIP_PATH ${ASCEND_NNAL_RUNTIME_PATH}/asdsip/latest/) diff --git a/cmake/package.cmake b/cmake/package.cmake index 0e72968f5fb..ab75f68b39a 100644 --- a/cmake/package.cmake +++ b/cmake/package.cmake @@ -248,6 +248,13 @@ if(ENABLE_D) COMPONENT mindspore ) endif() + if(EXISTS ${ASCEND_NNAL_ASDSIP_PATH}) + install( + TARGETS mindspore_extension_ascend_asdsip ARCHIVE + DESTINATION ${INSTALL_PLUGIN_DIR}/ascend + COMPONENT mindspore + ) + endif() if(ENABLE_MPI) install( TARGETS ascend_collective d_collective diff --git a/cmake/package_plugin.cmake b/cmake/package_plugin.cmake index d39a86a299d..4bd08d190ef 100644 --- a/cmake/package_plugin.cmake +++ b/cmake/package_plugin.cmake @@ -99,6 +99,13 @@ if(ENABLE_D) COMPONENT mindspore ) endif() + if(EXISTS ${ASCEND_NNAL_ASDSIP_PATH}) + install( + TARGETS mindspore_extension_ascend_asdsip ARCHIVE + DESTINATION ${INSTALL_PLUGIN_DIR}/ascend + COMPONENT mindspore + ) + endif() install( TARGETS ms_atb_boost DESTINATION ${INSTALL_PLUGIN_DIR}/ascend diff --git a/mindspore/ccsrc/ms_extension/api.h b/mindspore/ccsrc/ms_extension/api.h index 73e1a8f149c..d08c112a996 100644 --- a/mindspore/ccsrc/ms_extension/api.h +++ b/mindspore/ccsrc/ms_extension/api.h @@ -29,5 +29,8 @@ #ifdef CUSTOM_ENABLE_ATB #include "ms_extension/ascend/atb/atb_common.h" #endif // CUSTOM_ENABLE_ATB +#ifdef CUSTOM_ENABLE_ASDSIP +#include "ms_extension/ascend/asdsip/asdsip_common.h" +#endif // CUSTOM_ENABLE_ASDSIP #endif // CUSTOM_ASCEND_OP #endif // MINDSPORE_CCSRC_MS_EXTENSION_API_H_ diff --git a/mindspore/ccsrc/ms_extension/ascend/CMakeLists.txt b/mindspore/ccsrc/ms_extension/ascend/CMakeLists.txt index 4046f336941..ad60c1fe639 100644 --- a/mindspore/ccsrc/ms_extension/ascend/CMakeLists.txt +++ b/mindspore/ccsrc/ms_extension/ascend/CMakeLists.txt @@ -4,3 +4,8 @@ if(EXISTS ${ASCEND_NNAL_ATB_PATH}) add_subdirectory(atb) endif() + +if(EXISTS ${ASCEND_NNAL_ASDSIP_PATH}) + add_definitions(-DASDSIP_PATH="${ASCEND_NNAL_ASDSIP_PATH}") + add_subdirectory(asdsip) +endif() diff --git a/mindspore/ccsrc/ms_extension/ascend/asdsip/CMakeLists.txt b/mindspore/ccsrc/ms_extension/ascend/asdsip/CMakeLists.txt new file mode 100644 index 00000000000..36b12cad1ad --- /dev/null +++ b/mindspore/ccsrc/ms_extension/ascend/asdsip/CMakeLists.txt @@ -0,0 +1,9 @@ +file(GLOB_RECURSE EXTENSION_ASCEND_ASDSIP_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} + "*.cc" +) + +set_property(SOURCE ${EXTENSION_ASCEND_ASDSIP_SRC_LIST} PROPERTY COMPILE_DEFINITIONS + SUBMODULE_ID=mindspore::SubModuleId::SM_EXTENSION) +add_library(mindspore_extension_ascend_asdsip STATIC ${EXTENSION_ASCEND_ASDSIP_SRC_LIST}) + +target_link_libraries(mindspore_extension_ascend_asdsip PRIVATE mindspore_extension mindspore_ops_ascend) diff --git a/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.cc b/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.cc new file mode 100644 index 00000000000..d91c1f76cfe --- /dev/null +++ b/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.cc @@ -0,0 +1,128 @@ +/** + * Copyright 2025 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "ascend/asdsip/asdsip_common.h" + +namespace ms::pynative { +constexpr size_t kDefaultCacheSize = 50; +using FFTHandlePair = std::pair; + +class FFTCache { + public: + static FFTCache &GetInstance() { + static FFTCache instance(kDefaultCacheSize); + return instance; + } + asdFftHandle Get(const FFTParam ¶m) { + uint64_t hash = HashFFTParam(param); + std::lock_guard lock(mtx_); + auto iter = fft_map_.find(hash); + if (iter != fft_map_.end()) { + fft_cache_.splice(fft_cache_.begin(), fft_cache_, iter->second); + return fft_cache_.front().second; + } else { + asdFftHandle handle = CreateHandle(param); + fft_cache_.push_front({hash, handle}); + fft_map_[hash] = fft_cache_.begin(); + if (fft_cache_.size() > capaticy_) { + DestrotyHandle(fft_cache_.back().second); + fft_map_.erase(fft_cache_.back().first); + fft_cache_.pop_back(); + } + return handle; + } + } + size_t GetSize() { return fft_cache_.size(); } + void SetSize(size_t capaticy) { + capaticy_ = capaticy; + std::lock_guard lock(mtx_); + while (fft_cache_.size() > capaticy_) { + DestrotyHandle(fft_cache_.back().second); + fft_map_.erase(fft_cache_.back().first); + fft_cache_.pop_back(); + } + } + void Clear() { + for (auto &item : fft_cache_) { + DestrotyHandle(item.second); + } + fft_map_.clear(); + fft_cache_.clear(); + } + + private: + explicit FFTCache(size_t capaticy) : capaticy_(capaticy) {} + size_t capaticy_; + std::list fft_cache_; + std::unordered_map::iterator> fft_map_; + std::mutex mtx_; +}; + +size_t AsdSipFFTOpRunner::cache_capaticy_ = kDefaultCacheSize; +bool AsdSipFFTOpRunner::cache_set_flag_ = false; +void AsdSipFFTOpRunner::SetCacheSize(size_t capaticy) { cache_capaticy_ = capaticy; } + +void AsdSipFFTOpRunner::Init(const FFTParam ¶m) { + if (!cache_set_flag_) { + FFTCache::GetInstance().SetSize(cache_capaticy_); + cache_set_flag_ = true; + } + _device_context_->device_res_manager_->BindDeviceToCurrentThread(false); + asd_fft_handle_ = FFTCache::GetInstance().Get(param); +} + +inner::MemBlockPtr AsdSipFFTOpRunner::_MallocDeviceAddress() { + auto mem_block = PyboostRunner::_MallocDeviceAddress(); + AsdFftSetWorkSpace(asd_fft_handle_, _workspace_ptr_); + AsdFftSetStream(asd_fft_handle_, stream()); + if (_inputs_[0].is_defined() && !_inputs_[0].is_contiguous()) { + input_tensor_ = mindspore::device::ascend::ConvertType(_inputs_[0].contiguous().tensor()); + } else { + input_tensor_ = mindspore::device::ascend::ConvertType(_inputs_[0].tensor()); + } + output_tensor_ = mindspore::device::ascend::ConvertType(_outputs_[0].tensor()); + return mem_block; +} + +size_t AsdSipFFTOpRunner::CalcWorkspace() { + AsdFftGetWorkSpaceSize(asd_fft_handle_, workspace_size_); + return workspace_size_; +} + +void AsdSipFFTOpRunner::_Run() { PyboostRunner::_Run(); } + +void AsdSipFFTOpRunner::LaunchKernel() { + static std::unordered_map asd_fft_exec_map; + auto iter = asd_fft_exec_map.find(_op_name_); + AsdFftExecFunc asd_fft_exec = nullptr; + if (iter == asd_fft_exec_map.end()) { + auto asd_fft_exec_symbol = GetAsdSipApiFuncAddr(_op_name_.c_str()); + asd_fft_exec = reinterpret_cast(asd_fft_exec_symbol); + asd_fft_exec_map[_op_name_] = asd_fft_exec; + } else { + asd_fft_exec = iter->second; + } + if (asd_fft_exec == nullptr) { + MS_LOG(EXCEPTION) << "Get asdFftExec " << _op_name_ << "failed"; + } + auto ret = asd_fft_exec(asd_fft_handle_, input_tensor_, output_tensor_); + if (ret != 0) { + MS_LOG(EXCEPTION) << "Execute " << _op_name_ << " failed, ret: " << ret; + } + mindspore::device::ascend::Release(input_tensor_); + mindspore::device::ascend::Release(output_tensor_); +} +} // namespace ms::pynative diff --git a/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.h b/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.h new file mode 100644 index 00000000000..e1a13d646c3 --- /dev/null +++ b/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_common.h @@ -0,0 +1,87 @@ +/** + * Copyright 2025 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef MINDSPORE_CCSRC_MS_EXTENSION_ASCEND_ASDSIP_COMMON_H_ +#define MINDSPORE_CCSRC_MS_EXTENSION_ASCEND_ASDSIP_COMMON_H_ + +#include +#include +#include +#include +#include +#include +#include "ms_extension/ascend/asdsip/asdsip_utils.h" + +namespace ms::pynative { +/// \brief A runner for executing Ascend SiP Boost (ASDSIP) FFT operations. +/// +/// This class extends the `PyboostRunner` class and provides additional functionality +/// for initializing and running ASDSIP FFT operations on Ascend hardware. It manages input +/// and output tensors, memory allocation, and kernel launches specific to ASDSIP. +class EXTENSION_API AsdSipFFTOpRunner : public PyboostRunner { + public: + explicit AsdSipFFTOpRunner(std::string op_name) : PyboostRunner(op_name) {} + inner::MemBlockPtr _MallocDeviceAddress() override; + size_t CalcWorkspace() override; + void LaunchKernel() override; + /// \brief [API] Initializes the ASDSIP operation with the given parameters. + /// + /// This method retrieves the corresponding ASDSIP operation instance using + /// the provided parameters and operation name. + /// The operation instance is used to configure and execute the operation. + /// + /// \param[in] param The parameters required to initialize the ASDSIP operation. + /// + /// \details + /// - The `param` is passed to the `FFTCache` singleton, which manages + /// caching of ASDSIP operation instances. + void Init(const FFTParam ¶m); + static void SetCacheSize(size_t capaticy); + + protected: + void _Run() override; + + private: + asdFftHandle asd_fft_handle_; + uint64_t workspace_size_{0}; + aclTensor *input_tensor_; + aclTensor *output_tensor_; + static bool cache_set_flag_; + static size_t cache_capaticy_; +}; + +/// \brief [API] Executes an ASDSIP FFT operation using the provided parameters, inputs, and outputs. +/// +/// This function creates an `AsdSipFFTOpRunner` instance, initializes it with the given parameters, +/// and executes the operation on Ascend hardware. +/// +/// \param[in] op_name The name of the ASDSIP FFT operation to execute. +/// \param[in] fft_param The parameters required to initialize the ASDSIP FFT operation. +/// \param[in] input An input tensor for the operation. +/// \param[in] output An output tensor for the operation. +/// +/// \details +/// - The function first creates a shared pointer to an `AsdSipFFTOpRunner` instance. +/// - It calls the `Init` method of the runner to initialize the operation. +/// - Finally, it invokes the `Run` method of the runner to execute the operation. +inline void LaunchAsdSipFFT(const std::string &op_name, const FFTParam &fft_param, const ms::Tensor &input, + const ms::Tensor &output) { + auto runner = std::make_shared(op_name); + runner->Init(fft_param); + runner->Run({input}, {output}); +} +} // namespace ms::pynative +#endif // MINDSPORE_CCSRC_MS_EXTENSION_ASCEND_ASDSIP_COMMON_H_ diff --git a/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_utils.h b/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_utils.h new file mode 100644 index 00000000000..28f6fb4cc98 --- /dev/null +++ b/mindspore/ccsrc/ms_extension/ascend/asdsip/asdsip_utils.h @@ -0,0 +1,191 @@ +/** + * Copyright 2025 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_CCSRC_MS_EXTENSION_ASCEND_ASDSIP_UTILS_H_ +#define MINDSPORE_CCSRC_MS_EXTENSION_ASCEND_ASDSIP_UTILS_H_ + +#include +#include +#include +#include +#include "acl/acl.h" +#include "ms_extension/pynative/pyboost_extension.h" +#include "kernel/ascend/acl_ir/op_api_convert.h" + +namespace ms::pynative { +using mindspore::device::ascend::aclTensor; + +enum asdFftType { + ASCEND_FFT_C2C = 0x10, + ASCEND_FFT_C2R = 0x11, + ASCEND_FFT_R2C = 0x12, + ASCEND_STFT_C2C = 0x20, + ASCEND_STFT_C2R = 0x21, + ASCEND_STFT_R2C = 0x22, +}; + +enum asdFftDirection { + ASCEND_FFT_FORWARD = 0x10, + ASCEND_FFT_BACKWARD = 0x11, +}; + +enum asdFft1dDimType { + ASCEND_FFT_HORIZONTAL = 0x10, + ASCEND_FFT_VERTICAL = 0x11, +}; + +struct FFTParam { + int64_t fftXSize = 0; + int64_t fftYSize = 0; + asdFftType fftType = asdFftType::ASCEND_FFT_C2C; + asdFftDirection direction = asdFftDirection::ASCEND_FFT_FORWARD; + int64_t batchSize = 0; + asdFft1dDimType dimType = asdFft1dDimType::ASCEND_FFT_HORIZONTAL; +}; + +typedef void *asdFftHandle; +typedef int (*_asdFftCreate)(asdFftHandle &handle); +typedef int (*_asdFftDestroy)(asdFftHandle handle); +typedef int (*_asdFftSetStream)(asdFftHandle handle, void *stream); +typedef int (*_asdFftSynchronize)(asdFftHandle handle); +typedef int (*_asdFftGetWorkspaceSize)(asdFftHandle handle, size_t &size); +typedef int (*_asdFftSetWorkspace)(asdFftHandle handle, void *work_space); +typedef int (*_asdFftMakePlan1D)(asdFftHandle handle, int64_t fftXSize, asdFftType fftType, asdFftDirection direction, + int64_t batch_size, asdFft1dDimType dim_type); +typedef int (*_asdFftMakePlan2D)(asdFftHandle handle, int64_t fftXSize, int64_t fftYSize, asdFftType fftType, + asdFftDirection direction, int64_t batch_size); + +using AsdFftExecFunc = int (*)(asdFftHandle handle, aclTensor *input, aclTensor *output); + +#define GET_API_FUNC(func_name) reinterpret_cast<_##func_name>(GetAsdSipApiFuncAddr(#func_name)) + +inline std::string GetAsdSipLibPath() { +#ifdef ASDSIP_PATH + return std::string(ASDSIP_PATH) + "/lib/libasdsip.so"; +#else + MS_LOG(EXCEPTION) << "Find libasdsip.so failed, you should install the nnal package first."; +#endif +} + +inline uint64_t HashFFTParam(const FFTParam ¶m) { + static std::hash hash_func; + std::string param_str = std::to_string(param.fftXSize) + std::to_string(param.fftYSize) + + std::to_string(static_cast(param.fftType)) + + std::to_string(static_cast(param.direction)) + std::to_string(param.batchSize) + + std::to_string(static_cast(param.dimType)); + return hash_func(param_str); +} + +inline void *GetAsdSipApiFuncAddr(const char *api_name) { + static auto handle = dlopen(GetAsdSipLibPath().c_str(), RTLD_LAZY | RTLD_LOCAL); + if (handle == nullptr) { + MS_LOG(EXCEPTION) + << "Load libasdsip.so from " << GetAsdSipLibPath() << " failed, error is: " << dlerror() + << ", you should check the nnal package whether installed correctly and source the set_env.sh in Asdsip."; + } + auto func_addr = dlsym(handle, api_name); + if (func_addr == nullptr) { + MS_LOG(EXCEPTION) << "Get api " << api_name << " from " << GetAsdSipLibPath() << " failed."; + } + return func_addr; +} + +inline int AsdFftCreate(asdFftHandle &handle) { + static auto asd_fft_create = GET_API_FUNC(asdFftCreate); + if (asd_fft_create == nullptr) { + MS_LOG(EXCEPTION) << "asdFftCreate is nullptr."; + } + return asd_fft_create(handle); +} + +inline int AsdFftMakePlan1D(asdFftHandle handle, int64_t fftXSize, asdFftType fftType, asdFftDirection direction, + int64_t batch_size, asdFft1dDimType dim_type) { + static auto asd_sip_make_plan_1d = GET_API_FUNC(asdFftMakePlan1D); + if (asd_sip_make_plan_1d == nullptr) { + MS_LOG(EXCEPTION) << "asdSipFftMakePlan1D is nullptr."; + } + return asd_sip_make_plan_1d(handle, fftXSize, fftType, direction, batch_size, dim_type); +} + +inline int AsdFftMakePlan2D(asdFftHandle handle, int64_t fftXSize, int64_t fftYSize, asdFftType fftType, + asdFftDirection direction, int64_t batch_size) { + static auto asd_sip_make_plan_2d = GET_API_FUNC(asdFftMakePlan2D); + if (asd_sip_make_plan_2d == nullptr) { + MS_LOG(EXCEPTION) << "asdSipFftMakePlan2D is nullptr."; + } + return asd_sip_make_plan_2d(handle, fftXSize, fftYSize, fftType, direction, batch_size); +} + +inline asdFftHandle CreateHandle(const FFTParam ¶m) { + asdFftHandle handle; + auto ret = AsdFftCreate(handle); + if (ret != 0) { + MS_LOG(EXCEPTION) << "AsdFftCreate failed."; + } + if (param.fftYSize == 0) { + AsdFftMakePlan1D(handle, param.fftXSize, param.fftType, param.direction, param.batchSize, param.dimType); + } else { + AsdFftMakePlan2D(handle, param.fftXSize, param.fftYSize, param.fftType, param.direction, param.batchSize); + } + return handle; +} + +inline int AsdFftDestroy(asdFftHandle handle) { + static auto asd_fft_destrory = GET_API_FUNC(asdFftDestroy); + if (asd_fft_destrory == nullptr) { + MS_LOG(EXCEPTION) << "asdFftDestroy is nullptr."; + } + return asd_fft_destrory(handle); +} + +inline int AsdFftSetStream(asdFftHandle handle, void *stream) { + static auto asd_fft_set_stream = GET_API_FUNC(asdFftSetStream); + if (asd_fft_set_stream == nullptr) { + MS_LOG(EXCEPTION) << "asdFftSetStream is nullptr."; + } + return asd_fft_set_stream(handle, stream); +} + +inline int AsdFftSynchronize(asdFftHandle handle) { + static auto asd_fft_synchronize = GET_API_FUNC(asdFftSynchronize); + if (asd_fft_synchronize == nullptr) { + MS_LOG(EXCEPTION) << "asdFftSynchronize is nullptr."; + } + return asd_fft_synchronize(handle); +} + +inline int AsdFftGetWorkSpaceSize(asdFftHandle handle, size_t &size) { + static auto asd_fft_get_workspace = GET_API_FUNC(asdFftGetWorkspaceSize); + if (asd_fft_get_workspace == nullptr) { + MS_LOG(EXCEPTION) << "asdFftGetWorkSpaceSize is nullptr."; + } + return asd_fft_get_workspace(handle, size); +} + +inline int AsdFftSetWorkSpace(asdFftHandle handle, void *work_space) { + static auto asd_fft_set_workspace = GET_API_FUNC(asdFftSetWorkspace); + if (asd_fft_set_workspace == nullptr) { + MS_LOG(EXCEPTION) << "asdFftSetWorkSpace is nullptr."; + } + return asd_fft_set_workspace(handle, work_space); +} + +inline void DestrotyHandle(asdFftHandle handle) { + AsdFftSynchronize(handle); + AsdFftDestroy(handle); +} + +} // namespace ms::pynative +#endif // MINDSPORE_CCSRC_MS_EXTENSION_ASCEND_ASDSIP_UTILS_H_ diff --git a/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.cc b/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.cc index 0a5fc76b211..e856ed74bf3 100644 --- a/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.cc +++ b/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.cc @@ -57,12 +57,6 @@ MemBlock::MemBlock(const mindspore::device::DeviceContext *device_context, size_ } MemBlock::~MemBlock() { device_context_->device_res_manager_->FreeMemory(ptr_); } - -std::string GetDeviceTarget() { - auto msctx = mindspore::MsContext::GetInstance(); - MS_EXCEPTION_IF_NULL(msctx); - return msctx->get_param(mindspore::MsCtxParam::MS_CTX_DEVICE_TARGET); -} } // namespace inner namespace pynative { @@ -73,7 +67,6 @@ void PyboostRunner::Run(const std::vector &inputs, const std::vector_PrepareStream(); this->_PrepareDeviceAddress(); PyBoostUtils::DispatchRun(std::make_shared([runner = shared_from_this()]() { diff --git a/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.h b/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.h index 7ed7926976e..3fb7aab3ab3 100644 --- a/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.h +++ b/mindspore/ccsrc/ms_extension/pynative/pyboost_extension.h @@ -157,7 +157,12 @@ class EXTENSION_EXPORT PyboostRunner : public std::enable_shared_from_thisget_param(mindspore::MsCtxParam::MS_CTX_DEVICE_TARGET); + _device_context_ = mindspore::runtime::OpRunner::GetDeviceContext(target); + } /** * @brief Virtual destructor for PyboostRunner. diff --git a/mindspore/python/mindspore/ops/operations/custom_ops.py b/mindspore/python/mindspore/ops/operations/custom_ops.py index 95f12fec5f8..7cdf26c558e 100644 --- a/mindspore/python/mindspore/ops/operations/custom_ops.py +++ b/mindspore/python/mindspore/ops/operations/custom_ops.py @@ -1217,6 +1217,9 @@ class CustomOpBuilder: - enable_atb (bool, optional): Whether to call ATB (Ascend Transformer Boost) operator. If set to ``True``, the `backend` must be ``Ascend`` or left empty. Default: ``False``. + - enable_asdsip (bool, optional): Whether to call ASDSIP (Ascend SiP Boost) operator. If set to ``True``, + the `backend` must be ``Ascend`` or left empty. Default: ``False``. + .. note:: - If the `backend` argument is provided, additional default flags will be automatically added to the compilation and linking steps to support the operator's target backend. The default options @@ -1247,6 +1250,7 @@ class CustomOpBuilder: self.ldflags = ldflags self.build_dir = kwargs.get("build_dir") self.enable_atb = kwargs.get("enable_atb", False) + self.enable_asdsip = kwargs.get("enable_asdsip", False) self.debug_mode = kwargs.get("debug_mode", False) self._ms_path = os.path.dirname(os.path.abspath(ms.__file__)) if self.enable_atb: @@ -1254,6 +1258,11 @@ class CustomOpBuilder: raise ValueError("For 'CustomOpBuilder', when 'enable_atb' is set to True, the 'backend' must be " f"'Ascend' (or left implicit), but got '{backend}'") self.backend = "Ascend" + if self.enable_asdsip: + if backend is not None and backend != "Ascend": + raise ValueError("For 'CustomOpBuilder', when 'enable_asdsip' is set to True, the 'backend' must be " + f"'Ascend' (or left implicit), but got '{backend}'") + self.backend = "Ascend" if self.backend == "Ascend": ascend_opp_path = os.getenv("ASCEND_OPP_PATH") if not ascend_opp_path: @@ -1324,6 +1333,8 @@ class CustomOpBuilder: flags.append('-DCUSTOM_ASCEND_OP') if self.enable_atb: flags.append('-DCUSTOM_ENABLE_ATB') + if self.enable_asdsip: + flags.append('-DCUSTOM_ENABLE_ASDSIP') if self.cflags is not None: flags.append(self.cflags) return flags @@ -1356,6 +1367,9 @@ class CustomOpBuilder: flags.append('-lmindspore_extension_ascend_atb') flags.append(f"-L{os.path.abspath(os.path.join(self.atb_home_path, 'lib'))}") flags.append('-latb') + if self.enable_asdsip: + flags.append(f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin', 'ascend'))}") + flags.append('-lmindspore_extension_ascend_asdsip') if self.ldflags is not None: flags.append(self.ldflags) return flags diff --git a/tests/st/graph_kernel/custom/jit_test_files/asdsip_fftc2c.cpp b/tests/st/graph_kernel/custom/jit_test_files/asdsip_fftc2c.cpp new file mode 100644 index 00000000000..115ed91fd18 --- /dev/null +++ b/tests/st/graph_kernel/custom/jit_test_files/asdsip_fftc2c.cpp @@ -0,0 +1,44 @@ +/** + * Copyright 2025 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "ms_extension/api.h" + +ms::Tensor InferFFTForward(const ms::Tensor &input) { + ShapeVector out_tensor_shape(input.shape()); + return ms::Tensor(input.data_type(), out_tensor_shape); +} + +ms::Tensor npu_fft(const ms::Tensor &input, int64_t n, int64_t batch_size) { + ms::pynative::FFTParam param; + param.fftXSize = n; + param.fftYSize = 0; + param.fftType = ms::pynative::asdFftType::ASCEND_FFT_C2C; + param.direction = ms::pynative::asdFftDirection::ASCEND_FFT_FORWARD; + param.batchSize = batch_size; + param.dimType = ms::pynative::asdFft1dDimType::ASCEND_FFT_HORIZONTAL; + auto output = InferFFTForward(input); + ms::pynative::LaunchAsdSipFFT("asdFftExecC2C", param, input, output); + return output; +} + +auto pyboost_npu_fft(const ms::Tensor &input, int64_t n, int64_t batch_size) { + return ms::pynative::PyboostRunner::Call<1>(npu_fft, input, n, batch_size); +} + +PYBIND11_MODULE(MS_EXTENSION_NAME, m) { + m.def("fft", &pyboost_npu_fft, "FFT C2C", pybind11::arg("input"), pybind11::arg("n"), + pybind11::arg("batch_size")); +} diff --git a/tests/st/graph_kernel/custom/test_custom_pyboost_ascend.py b/tests/st/graph_kernel/custom/test_custom_pyboost_ascend.py index cb9cfb37dee..deafac0d083 100644 --- a/tests/st/graph_kernel/custom/test_custom_pyboost_ascend.py +++ b/tests/st/graph_kernel/custom/test_custom_pyboost_ascend.py @@ -109,6 +109,27 @@ def test_pyboost_atb_rope(): assert np.allclose(run_key.asnumpy(), benchmark_key.asnumpy(), 1e-3, 1e-3) +@arg_mark(plat_marks=['platform_ascend910b'], level_mark='level1', card_mark='onecard', essential_mark='essential') +def test_pyboost_asdsip_fft(): + """ + Feature: CustomOpBuilder. + Description: Custom asdsip op. + Expectation: success. + """ + ms.set_device("Ascend") + my_ops = CustomOpBuilder("asdsip_fftc2c", "jit_test_files/asdsip_fftc2c.cpp", enable_asdsip=True).load() + input_np = np.random.rand(2, 16) + real_np = input_np.astype(np.float32) + imag_np = input_np.astype(np.float32) + complex_np = real_np + 1j * imag_np + + input_tensor = ms.Tensor(complex_np, dtype=ms.dtype.complex64) + output_tensor = my_ops.fft(input_tensor, 16, 2) + output_np = np.fft.fft(complex_np) + assert np.allclose(output_tensor.asnumpy(), output_np, 1e-3, 1e-3) + + + @arg_mark(plat_marks=['platform_ascend910b'], level_mark='level1', card_mark='onecard', essential_mark='essential') def test_pyboost_aclnn(): """ -- Gitee