已合并
feat: add aclnnReselectStaticKernelWithPath interface #40386
rich创建于 7月7日
feat: add aclnnReselectStaticKernelWithPath interface #40386
已合并
共 4 个文件变更+80-0
| @@ -17,6 +17,7 @@ TORCH_NPU_REGISTER_LIBRARY(libopapi) | |||
| 17 | TORCH_NPU_LOAD_FUNC(aclnnSilentCheck) | 17 | TORCH_NPU_LOAD_FUNC(aclnnSilentCheck) |
| 18 | TORCH_NPU_LOAD_FUNC(aclnnSilentCheckV2) | 18 | TORCH_NPU_LOAD_FUNC(aclnnSilentCheckV2) |
| 19 | TORCH_NPU_LOAD_FUNC(aclnnReselectStaticKernel) | 19 | TORCH_NPU_LOAD_FUNC(aclnnReselectStaticKernel) |
| 20 | +TORCH_NPU_LOAD_FUNC(aclnnReselectStaticKernelWithPath) | ||
| 20 | 21 | ||
| 21 | bool IsExistAclnnSilentCheck() | 22 | bool IsExistAclnnSilentCheck() |
| 22 | { | 23 | { |
| @@ -39,5 +40,20 @@ aclnnStatus ReselectStaticKernel() | |||
| 39 | return ret; | 40 | return ret; |
| 40 | } | 41 | } |
| 41 | 42 | ||
| 43 | +aclnnStatus ReselectStaticKernelWithPath(const std::string &path) | ||
| 44 | +{ | ||
| 45 | + typedef aclnnStatus (*AclnnApiFunc)(const char *); | ||
| 46 | + static AclnnApiFunc aclnnReselectStaticKernelWithPathFunc = nullptr; | ||
| 47 | + if (aclnnReselectStaticKernelWithPathFunc == nullptr) { | ||
| 48 | + aclnnReselectStaticKernelWithPathFunc = | ||
| 49 | + (AclnnApiFunc)TORCH_NPU_GET_FUNC(aclnnReselectStaticKernelWithPath); | ||
| 50 | + } | ||
| 51 | + TORCH_CHECK(aclnnReselectStaticKernelWithPathFunc, | ||
 | |||
| 52 | + "Failed to find function ", "aclnnReselectStaticKernelWithPath", | ||
| 53 | + PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 54 | + auto ret = aclnnReselectStaticKernelWithPathFunc(path.c_str()); | ||
| 55 | + return ret; | ||
| 56 | +} | ||
| 57 | + | ||
| 42 | } // namespace opapi | 58 | } // namespace opapi |
| 43 | } // namespace c10_npu | 59 | } // namespace c10_npu |
| @@ -1,6 +1,7 @@ | |||
| 1 | 1 | ||
| 2 | 2 | ||
| 3 | 3 | ||
| 4 | + | ||
| 4 | 5 | ||
| 5 | namespace c10_npu { | 6 | namespace c10_npu { |
| 6 | namespace opapi { | 7 | namespace opapi { |
| @@ -16,5 +17,11 @@ bool IsExistAclnnSilentCheck(); | |||
| 16 | */ | 17 | */ |
| 17 | aclnnStatus ReselectStaticKernel(); | 18 | aclnnStatus ReselectStaticKernel(); |
| 18 | 19 | ||
| 20 | +/** | ||
| 21 | + This Api is used to reselect static kernel with a specified path, | ||
| 22 | + it need to be called once at process. | ||
| 23 | + */ | ||
| 24 | +aclnnStatus ReselectStaticKernelWithPath(const std::string &path); | ||
| 25 | + | ||
| 19 | } // namespace opapi | 26 | } // namespace opapi |
| 20 | } // namespace c10_npu | 27 | } // namespace c10_npu |
| @@ -1,6 +1,7 @@ | |||
| 1 | 1 | ||
| 2 | 2 | ||
| 3 | 3 | ||
| 4 | + | ||
| 4 | 5 | ||
| 5 | 6 | ||
| 6 | 7 | ||
| @@ -2076,6 +2077,53 @@ PyObject* THNPModule_aclnn_reselect_static_kernel( | |||
| 2076 | END_HANDLE_TH_ERRORS | 2077 | END_HANDLE_TH_ERRORS |
| 2077 | } | 2078 | } |
| 2078 | 2079 | ||
| 2080 | +PyObject* THNPModule_aclnn_reselect_static_kernel_with_path( | ||
| 2081 | + PyObject* self, | ||
| 2082 | + PyObject* arg) { | ||
| 2083 | + HANDLE_TH_ERRORS | ||
| 2084 | + TORCH_CHECK(THPUtils_checkString(arg), | ||
| 2085 | + "path must be a string", | ||
| 2086 | + PTA_ERROR(ErrCode::PARAM)); | ||
| 2087 | + std::string path = THPUtils_unpackString(arg); | ||
| 2088 | + TORCH_CHECK(path.find('\0') == std::string::npos, | ||
| 2089 | + "path must not contain null byte", | ||
| 2090 | + PTA_ERROR(ErrCode::PARAM)); | ||
| 2091 | + struct stat st; | ||
| 2092 | + TORCH_CHECK(stat(path.c_str(), &st) == 0, | ||
| 2093 | + "path does not exist: ", path, | ||
| 2094 | + PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 2095 | + TORCH_CHECK(S_ISDIR(st.st_mode), | ||
| 2096 | + "path must be a directory: ", path, | ||
| 2097 | + PTA_ERROR(ErrCode::PARAM)); | ||
| 2098 | + | ||
| 2099 | + NPUStatus ret = c10_npu::emptyAllNPUStream(); | ||
| 2100 | + TORCH_CHECK( | ||
| 2101 | + ret == NPU_STATUS_SUCCESS, | ||
| 2102 | + "Failed to empty NPU task queue, ret:", | ||
| 2103 | + ret, | ||
| 2104 | + PTA_ERROR(ErrCode::INTERNAL)); | ||
| 2105 | + | ||
| 2106 | + static const auto task_queue_enable = | ||
| 2107 | + c10_npu::option::OptionsManager::GetTaskQueueEnable(); | ||
| 2108 | + if (task_queue_enable == 2) { | ||
| 2109 | + auto acl_call = [path]() -> int { | ||
| 2110 | + return c10_npu::opapi::ReselectStaticKernelWithPath(path); | ||
| 2111 | + }; | ||
| 2112 | + at_npu::native::OpCommand::RunOpApiV2("reselect_static_kernel_with_path", acl_call); | ||
| 2113 | + NPUStatus ret = c10_npu::emptyAllNPUStream(); | ||
| 2114 | + TORCH_CHECK( | ||
| 2115 | + ret == NPU_STATUS_SUCCESS, | ||
| 2116 | + "Failed to empty NPU task queue, ret:", | ||
| 2117 | + ret, | ||
| 2118 | + PTA_ERROR(ErrCode::INTERNAL)); | ||
| 2119 | + } else { | ||
| 2120 | + NPU_CHECK_ERROR(c10_npu::opapi::ReselectStaticKernelWithPath(path)); | ||
| 2121 | + } | ||
| 2122 | + | ||
| 2123 | + Py_RETURN_NONE; | ||
| 2124 | + END_HANDLE_TH_ERRORS | ||
| 2125 | +} | ||
| 2126 | + | ||
| 2079 | PyObject* THNPModule_npu_set_thread_affinity(PyObject* self, PyObject* args) { | 2127 | PyObject* THNPModule_npu_set_thread_affinity(PyObject* self, PyObject* args) { |
| 2080 | HANDLE_TH_ERRORS | 2128 | HANDLE_TH_ERRORS |
| 2081 | int core_start, core_end; | 2129 | int core_start, core_end; |
| @@ -2702,6 +2750,10 @@ static struct PyMethodDef THNPModule_methods[] = { | |||
| 2702 | (PyCFunction)THNPModule_aclnn_reselect_static_kernel, | 2750 | (PyCFunction)THNPModule_aclnn_reselect_static_kernel, |
| 2703 | METH_NOARGS, | 2751 | METH_NOARGS, |
| 2704 | nullptr}, | 2752 | nullptr}, |
| 2753 | + {"_aclnn_reselect_static_kernel_with_path", | ||
| 2754 | + (PyCFunction)THNPModule_aclnn_reselect_static_kernel_with_path, | ||
| 2755 | + METH_O, | ||
| 2756 | + nullptr}, | ||
| 2705 | {"_npu_set_thread_affinity", | 2757 | {"_npu_set_thread_affinity", |
| 2706 | (PyCFunction)THNPModule_npu_set_thread_affinity, | 2758 | (PyCFunction)THNPModule_npu_set_thread_affinity, |
| 2707 | METH_VARARGS, | 2759 | METH_VARARGS, |
| @@ -618,6 +618,11 @@ def _aclnn_reselect_static_kernel(): | |||
| 618 | torch_npu._C._aclnn_reselect_static_kernel() | 618 | torch_npu._C._aclnn_reselect_static_kernel() |
| 619 | 619 | ||
| 620 | 620 | ||
| 621 | +def _aclnn_reselect_static_kernel_with_path(path): | ||
| 622 | + torch_npu.npu._lazy_init() | ||
| 623 | + torch_npu._C._aclnn_reselect_static_kernel_with_path(path) | ||
| 624 | + | ||
| 625 | + | ||
| 621 | from .random import * # noqa: F403 | 626 | from .random import * # noqa: F403 |
| 622 | from .memory import * # noqa: F403 | 627 | from .memory import * # noqa: F403 |
| 623 | 628 | ||
这个check应该可以放到
if (aclnnReselectStaticKernelWithPathFunc == nullptr)内部