已合并
support torch.npu.npurt #39457
support torch.npu.npurt #39457
已合并
bellatan创建于 6月27日
8 个文件变更+268-8
Atest/npu/test_npurt.py+106-0
@@ -0,0 +1,106 @@
1+#!/usr/bin/env python3
2+import ctypes
3+ 
4+import torch
5+import torch_npu
6+from torch.testing._internal.common_utils import TestCase, run_tests
7+ 
8+ 
9+ACL_HOST_REG_MAPPED = 0x2
10+ACL_HOST_REG_PINNED = 0x10000000
11+ 
12+ 
13+class TestNPURT(TestCase):
14+ def _check_host_register_unregister(self, flag):
15+ rt = torch_npu.npu.npurt()
16+ 
17+ tensor = torch.empty(4096, dtype=torch.uint8).share_memory_()
18+ ptr = tensor.data_ptr()
19+ size = tensor.numel() * tensor.element_size()
20+ 
21+ ret = rt.npuHostRegister(ptr, size, flag)
22+ self.assertEqual(int(ret), 0)
23+ 
24+ try:
25+ self.assertGreater(ptr, 0)
26+ self.assertGreater(size, 0)
27+ self.assertEqual(ptr % 4096, 0)
28+ self.assertEqual(size % 4096, 0)
29+ finally:
30+ ret = rt.npuHostUnregister(ptr)
31+ self.assertEqual(int(ret), 0)
32+ 
33+ def test_npurt_api_exists(self):
34+ rt = torch_npu.npu.npurt()
35+ self.assertIs(rt, torch_npu.npu.npurt())
36+ self.assertTrue(hasattr(torch_npu._C, "_npurt"))
37+ 
38+ for name in (
39+ "npuHostRegister",
40+ "npuHostUnregister",
41+ "npuStreamCreate",
42+ "npuStreamDestroy",
43+ ):
44+ self.assertTrue(hasattr(rt, name))
45+ 
46+ def test_npurt_stream_create_destroy(self):
47+ rt = torch_npu.npu.npurt()
48+ 
49+ stream = ctypes.c_void_p()
50+ stream_p_int = ctypes.addressof(stream)
51+ 
52+ ret = rt.npuStreamCreate(stream_p_int)
53+ self.assertEqual(int(ret), 0)
54+ 
55+ try:
56+ self.assertIsNotNone(stream.value)
57+ self.assertNotEqual(stream.value, 0)
58+ finally:
59+ if stream.value:
60+ ret = rt.npuStreamDestroy(stream.value)
61+ self.assertEqual(int(ret), 0)
62+ 
63+ def test_npurt_host_register_supported_flags(self):
64+ # Host memory satisfies the ACL registration requirements (shared memory /
65+ # page-aligned allocation). In this scenario ACL_HOST_REG_MAPPED and
66+ # ACL_HOST_REG_PINNED (or their combination) are all expected to succeed.
67+ for flag in (
68+ ACL_HOST_REG_PINNED,
69+ ACL_HOST_REG_MAPPED,
70+ ACL_HOST_REG_PINNED | ACL_HOST_REG_MAPPED,
71+ ):
72+ self._check_host_register_unregister(flag)
73+ 
74+ def test_npurt_register_only_pinned(self):
75+ # Ordinary CPU tensor allocated by PyTorch (non page-aligned / non-shared
76+ # memory). Only ACL_HOST_REG_PINNED is expected to work, matching the
77+ # current PyTorch pin_memory registration path.
78+ ACL_HOST_REG_PINNED = 0x10000000
79+ t = torch.ones(20)
80+ npurt = torch_npu.npu.npurt()
81+ r = npurt.npuHostRegister(t.data_ptr(), t.numel() * t.element_size(), ACL_HOST_REG_PINNED)
82+ self.assertEqual(r, 0)
83+ try:
84+ self.assertGreater(t.data_ptr(), 0)
85+ finally:
86+ r = npurt.npuHostUnregister(t.data_ptr())
87+ self.assertEqual(r, 0)
88+ 
89+ def test_npurt_host_register_invalid_size(self):
90+ rt = torch_npu.npu.npurt()
91+ 
92+ tensor = torch.empty(4096, dtype=torch.uint8).share_memory_()
93+ ret = rt.npuHostRegister(tensor.data_ptr(), 0, ACL_HOST_REG_PINNED)
94+ self.assertNotEqual(int(ret), 0)
95+ 
96+ def test_npurt_null_ptr_returns_error(self):
97+ rt = torch_npu.npu.npurt()
98+ 
99+ self.assertNotEqual(int(rt.npuHostRegister(0, 4096, ACL_HOST_REG_PINNED)), 0)
100+ self.assertNotEqual(int(rt.npuHostUnregister(0)), 0)
101+ self.assertNotEqual(int(rt.npuStreamCreate(0)), 0)
102+ self.assertNotEqual(int(rt.npuStreamDestroy(0)), 0)
103+ 
104+ 
105+if __name__ == "__main__":
106+ run_tests()
Mtest/npu/test_torch_npu.py+3-2
@@ -1,3 +1,4 @@
1+#!/usr/bin/env python3
1import unittest2import unittest
2import contextlib3import contextlib
3import collections4import collections
@@ -27,9 +28,9 @@ class TorchNPUDeviceTestCase(TestCase):
27 def test_npu_can_device_access_peer(self):28 def test_npu_can_device_access_peer(self):
28 res = torch_npu.npu.can_device_access_peer(0, 0)29 res = torch_npu.npu.can_device_access_peer(0, 0)
29 self.assertEqual(res, False)30 self.assertEqual(res, False)
30- with self.assertRaisesRegex(AssertionError, "Invalid devide id"):31+ with self.assertRaisesRegex(AssertionError, "Invalid device id"):
31 torch_npu.npu.can_device_access_peer(-1, 0)32 torch_npu.npu.can_device_access_peer(-1, 0)
32- with self.assertRaisesRegex(AssertionError, "Invalid peer devide id"):33+ with self.assertRaisesRegex(AssertionError, "Invalid peer device id"):
33 torch_npu.npu.can_device_access_peer(0, -1)34 torch_npu.npu.can_device_access_peer(0, -1)
34 35 
35 def test_npu_device(self):36 def test_npu_device(self):
Mtest/torch_npu_schema.json+3-0
@@ -2915,5 +2915,8 @@
2915 "torch_c_func: batch_isend_irecv": {2915 "torch_c_func: batch_isend_irecv": {
2916 "signature": "(std::vector<std::string>& op_type, std::vector<at::Tensor>& tensors, std::vector<int64_t> remote_rank_list) -> c10::intrusive_ptr<c10d::Work>",2916 "signature": "(std::vector<std::string>& op_type, std::vector<at::Tensor>& tensors, std::vector<int64_t> remote_rank_list) -> c10::intrusive_ptr<c10d::Work>",
2917 "file": "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp"2917 "file": "torch_npu/csrc/distributed/ProcessGroupHCCL.hpp"
2918+ },
2919+ "torch_npu.npu.npurt": {
2920+ "signature": "()"
2918 }2921 }
2919}2922}
Mtorch_npu/_init/core/module_loader.py+3-0
@@ -1,3 +1,4 @@
1+#!/usr/bin/env python3
1import importlib2import importlib
2import inspect3import inspect
3import os4import os
@@ -13,6 +14,7 @@ _REQUIRED_C_EXTENSION_CHILDREN = [
13 "_flops_count",14 "_flops_count",
14 "_profiler",15 "_profiler",
15 "_distributed_c10d",16 "_distributed_c10d",
17+ "_npurt",
16]18]
17 19 
18 20 
@@ -59,6 +61,7 @@ def _initialize_c_extension_children(required_children):
59 _create_child_once(_C, "_cd", "_cd_init")61 _create_child_once(_C, "_cd", "_cd_init")
60 _create_child_once(_C, "_logging", "_logging_init")62 _create_child_once(_C, "_logging", "_logging_init")
61 _create_child_once(_C, "_flops_count", "_flops_count_init")63 _create_child_once(_C, "_flops_count", "_flops_count_init")
64+ _create_child_once(_C, "_npurt", "_npurt_init")
62 65 
63 _register_c_extension_submodules(_C)66 _register_c_extension_submodules(_C)
64 missing = [name for name in required_children if not hasattr(_C, name)]67 missing = [name for name in required_children if not hasattr(_C, name)]
Mtorch_npu/csrc/InitNpuBindings.cpp+2-0
@@ -19,6 +19,7 @@
19#include "torch_npu/csrc/logging/Init.h"19#include "torch_npu/csrc/logging/Init.h"
20#include "torch_npu/csrc/ipc/StorageSharing.h"20#include "torch_npu/csrc/ipc/StorageSharing.h"
21#include "torch_npu/csrc/npu/Module.h"21#include "torch_npu/csrc/npu/Module.h"
22+#include "torch_npu/csrc/npu/npurt.h"
22#include "torch_npu/csrc/custom_dtype/Init.h"23#include "torch_npu/csrc/custom_dtype/Init.h"
23#include "torch_npu/csrc/npu/Stress_detect.h"24#include "torch_npu/csrc/npu/Stress_detect.h"
24#include "torch_npu/csrc/utils/TensorType.h"25#include "torch_npu/csrc/utils/TensorType.h"
@@ -182,6 +183,7 @@ PyObject* initModule()
182 AddPyMethodDefs(methods, torch_npu::flopcount::flops_count_functions());183 AddPyMethodDefs(methods, torch_npu::flopcount::flops_count_functions());
183 AddPyMethodDefs(methods, torch_npu::logging::logging_functions());184 AddPyMethodDefs(methods, torch_npu::logging::logging_functions());
184 AddPyMethodDefs(methods, torch_npu::reductions::reductions_functions());185 AddPyMethodDefs(methods, torch_npu::reductions::reductions_functions());
186+ AddPyMethodDefs(methods, torch_npu::npurt::npurt_functions());
185 AddPyMethodDefs(methods, c10_npu::custom_dtype_functions());187 AddPyMethodDefs(methods, c10_npu::custom_dtype_functions());
186 AddPyMethodDefs(methods, torch_npu::afd::python_functions());188 AddPyMethodDefs(methods, torch_npu::afd::python_functions());
187 static struct PyModuleDef torchnpu_module = {189 static struct PyModuleDef torchnpu_module = {
Atorch_npu/csrc/npu/npurt.cpp+91-0
@@ -0,0 +1,91 @@
1+#include "torch_npu/csrc/npu/npurt.h"
2+ 
3+#include <cstdint>
4+ 
5+#include <pybind11/pybind11.h>
6+#include <pybind11/stl.h>
7+#include <torch/csrc/python_headers.h>
8+#include <torch/csrc/utils/object_ptr.h>
9+#include <torch/csrc/utils/pybind.h>
10+ 
11+#include "torch_npu/csrc/core/npu/NPUGuard.h"
12+#include "torch_npu/csrc/core/npu/NPUException.h"
13+#include "torch_npu/csrc/core/npu/interface/AclInterface.h"
14+ 
15+namespace py = pybind11;
16+ 
17+namespace torch_npu {
18+namespace npurt {
19+ 
20+namespace {
21+int npuHostRegister(uintptr_t ptr, size_t size, uint32_t flags)
22+{
23+ if (ptr == 0) {
24+ return static_cast<int>(ACL_ERROR_INVALID_PARAM);
25+ }
26+ py::gil_scoped_release no_gil;
27+ return static_cast<int>(
28+ c10_npu::acl::AclrtHostRegisterV2(reinterpret_cast<void*>(ptr), size, flags));
29+}
30+ 
31+int npuHostUnregister(uintptr_t ptr)
32+{
33+ if (ptr == 0) {
34+ return static_cast<int>(ACL_ERROR_INVALID_PARAM);
35+ }
36+ py::gil_scoped_release no_gil;
37+ return static_cast<int>(
38+ c10_npu::acl::AclrtHostUnregister(reinterpret_cast<void*>(ptr)));
39+}
40+ 
41+int npuStreamCreate(uintptr_t ptr)
42+{
43+ if (ptr == 0) {
44+ return static_cast<int>(ACL_ERROR_INVALID_PARAM);
45+ }
46+ py::gil_scoped_release no_gil;
47+ return static_cast<int>(
48+ aclrtCreateStream(reinterpret_cast<aclrtStream*>(ptr)));
49+}
50+ 
51+int npuStreamDestroy(uintptr_t ptr)
52+{
53+ if (ptr == 0) {
54+ return static_cast<int>(ACL_ERROR_INVALID_PARAM);
55+ }
56+ py::gil_scoped_release no_gil;
57+ return static_cast<int>(
58+ aclrtDestroyStream(reinterpret_cast<aclrtStream>(ptr)));
59+}
60+ 
61+PyObject* npurt_init(PyObject* /* unused */, PyObject* /* noargs */)
62+{
63+ auto torch_npu_C_module = THPObjectPtr(PyImport_ImportModule("torch_npu._C"));
64+ if (!torch_npu_C_module) {
65+ return nullptr;
66+ }
67+ 
68+ auto torch_npu_C_m = py::handle(torch_npu_C_module).cast<py::module>();
69+ auto npurt = torch_npu_C_m.def_submodule("_npurt", "NPU runtime API bindings");
70+ npurt.def("npuHostRegister", &npuHostRegister);
71+ npurt.def("npuHostUnregister", &npuHostUnregister);
72+ npurt.def("npuStreamCreate", &npuStreamCreate);
73+ npurt.def("npuStreamDestroy", &npuStreamDestroy);
74+ 
75+ Py_RETURN_TRUE;
76+}
77+ 
78+static PyMethodDef NPURTMethods[] = {
79+ {"_npurt_init", npurt_init, METH_NOARGS, nullptr},
80+ {nullptr, nullptr, 0, nullptr}
81+};
82+ 
83+} // namespace
84+ 
85+PyMethodDef* npurt_functions()
86+{
87+ return NPURTMethods;
88+}
89+ 
90+} // namespace npurt
91+} // namespace torch_npu
Atorch_npu/csrc/npu/npurt.h+13-0
@@ -0,0 +1,13 @@
1+#pragma once
2+ 
3+#include <Python.h>
4+ 
5+#include "torch_npu/csrc/core/npu/NPUMacros.h"
6+ 
7+namespace torch_npu {
8+namespace npurt {
9+ 
10+TORCH_NPU_API PyMethodDef* npurt_functions();
11+ 
12+} // namespace npurt
13+} // namespace torch_npu
Mtorch_npu/npu/__init__.py+47-6
@@ -1,3 +1,4 @@
1+#!/usr/bin/env python3
1__all__ = [2__all__ = [
2 "is_initialized",3 "is_initialized",
3 "init",4 "init",
@@ -140,6 +141,7 @@ __all__ = [
140 "register_npu_graph_handler",141 "register_npu_graph_handler",
141 "super_kernel_scope_begin",142 "super_kernel_scope_begin",
142 "super_kernel_scope_end",143 "super_kernel_scope_end",
144+ "npurt",
143]145]
144 146 
145from typing import Tuple, Union, List, cast, Optional147from typing import Tuple, Union, List, cast, Optional
@@ -168,7 +170,7 @@ from .npu_config import * # noqa: F403
168from .autocast_utils import * # noqa: F403170from .autocast_utils import * # noqa: F403
169from .backends import * # noqa: F403171from .backends import * # noqa: F403
170from ._backends import * # noqa: F403172from ._backends import * # noqa: F403
171-from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403173+from .deterministic import enable_deterministic_with_backward, disable_deterministic_with_backward # noqa: F403
172from . import npugraph_ex174from . import npugraph_ex
ascend-robotascend-robot
ascend-robotascend-robot7月1日

【openlibing.ci】识别到代码检查告警抑制注释,匹配工具:flake8,ruff,请Committer检视其合理性。

likedislike
bellatan
bellatan
7月1日 评论:
ascend-robotascend-robot7月1日

此条代码评论区间+172+174

【openlibing.ci】识别到代码检查告警抑制注释,匹配工具:flake8,ruff,请Committer检视其合理性。

likedislike
173 175 
174from .graphs import (176from .graphs import (
@@ -241,6 +243,40 @@ def init():
241 torch_npu.npu._lazy_init()243 torch_npu.npu._lazy_init()
242 244 
243 245 
246+def npurt():
247+ r"""Retrieves the NPU runtime API module.
248+ 
249+ This function initializes the NPU runtime environment if it is not already
250+ initialized and returns the NPU runtime API module (_npurt). The module
251+ provides access to a subset of NPU runtime functions.
252+ 
253+ Available APIs include:
254+ - npuHostRegister: Register host memory for device access.
255+ The `flags` parameter must follow CANN ACL Host Register
256+ definitions.
257+ 
258+ Refer to:
259+ https://www.hiascend.com/document/detail/zh/canncommercial/900/API/runtimeapi/aclcppdevg_03_2128.html
260+ 
261+ - npuHostUnregister: Unregister previously registered host memory.
262+ 
263+ - npuStreamCreate: Create a raw runtime stream.
264+ 
265+ - npuStreamDestroy: Destroy a raw runtime stream.
266+ 
267+ Returns:
268+ The NPU runtime API module (_npurt).
269+ 
270+ Raises:
271+ RuntimeError: If the NPU runtime cannot be initialized or the runtime
272+ API module is unavailable.
273+ """
274+ torch_npu.npu._lazy_init()
275+ if not hasattr(torch_npu._C, "_npurt"):
276+ raise RuntimeError("torch_npu._C._npurt is unavailable in this build.")
277+ return torch_npu._C._npurt
278+ 
279+ 
244def _lazy_init():280def _lazy_init():
245 def _queue_call(queued_calls):281 def _queue_call(queued_calls):
246 for queued_call, orig_traceback in queued_calls:282 for queued_call, orig_traceback in queued_calls:
@@ -471,9 +507,9 @@ def can_device_access_peer(device_id, peer_device_id):
471 device_id = _get_device_index(device_id, optional=True)507 device_id = _get_device_index(device_id, optional=True)
472 peer_device_id = _get_device_index(peer_device_id, optional=True)508 peer_device_id = _get_device_index(peer_device_id, optional=True)
473 if device_id < 0 or device_id >= device_count():509 if device_id < 0 or device_id >= device_count():
474- raise AssertionError("Invalid devide id" + pta_error(ErrCode.VALUE))510+ raise AssertionError("Invalid device id" + pta_error(ErrCode.VALUE))
475 if peer_device_id < 0 or peer_device_id >= device_count():511 if peer_device_id < 0 or peer_device_id >= device_count():
476- raise AssertionError("Invalid peer devide id" + pta_error(ErrCode.VALUE))512+ raise AssertionError("Invalid peer device id" + pta_error(ErrCode.VALUE))
477 return torch_npu._C._npu_canDeviceAccessPeer(device_id, peer_device_id)513 return torch_npu._C._npu_canDeviceAccessPeer(device_id, peer_device_id)
478 514 
479 515 
@@ -514,7 +550,8 @@ def get_device_capability(device=None):
514 The format should be "major.minor", e.g., "9.0" or "8.0".550 The format should be "major.minor", e.g., "9.0" or "8.0".
515 551 
516 .. note::552 .. note::
517- The return value of get_device_capability is only for compatibility with PyTorch and does not represent the actual capability of the NPU device.553+ The return value of get_device_capability is only for compatibility with PyTorch
554+ and does not represent the actual capability of the NPU device.
518 555 
519 Args:556 Args:
520 device (torch.device or int, optional): The device parameter has no practical meaning.557 device (torch.device or int, optional): The device parameter has no practical meaning.
@@ -535,7 +572,10 @@ def get_device_capability(device=None):
535 global _cached_device_capability, _cached_device_capability_env572 global _cached_device_capability, _cached_device_capability_env
536 573 
537 capability_env = os.getenv("TORCH_NPU_DEVICE_CAPABILITY")574 capability_env = os.getenv("TORCH_NPU_DEVICE_CAPABILITY")
538- warning_str = "The return value of get_device_capability is only for compatibility with PyTorch and does not represent the actual capability of the NPU device."575+ warning_str = (
576+ "The return value of get_device_capability is only for compatibility with PyTorch "
577+ "and does not represent the actual capability of the NPU device."
578+ )
539 if not capability_env:579 if not capability_env:
540 warnings.warn(f"You can set the device capability via the environment variable TORCH_NPU_DEVICE_CAPABILITY. {warning_str}")580 warnings.warn(f"You can set the device capability via the environment variable TORCH_NPU_DEVICE_CAPABILITY. {warning_str}")
541 return None581 return None
@@ -544,7 +584,8 @@ def get_device_capability(device=None):
544 if _cached_device_capability_env == capability_env and _cached_device_capability is not None:584 if _cached_device_capability_env == capability_env and _cached_device_capability is not None:
545 return _cached_device_capability585 return _cached_device_capability
546 586 
547- # Validate the format of the environment variable, expected format is 'major.minor' where major and minor are non-negative integers (e.g., '8.0')587+ # Validate the format of the environment variable, expected format is 'major.minor' where major
588+ # and minor are non-negative integers (e.g., '8.0')
548 pattern = r'^(\d+)\.(\d+)$'589 pattern = r'^(\d+)\.(\d+)$'
549 match = re.match(pattern, capability_env)590 match = re.match(pattern, capability_env)
550 if not match:591 if not match: