Star850
1331
代码
代码
Issues1514
Pull Requests2239
流水线
项目讨论
Wiki
分析
Star850
1331
已合并
[feat] add nestedtensor backend #30979
culechan创建于 2月25日
[feat] add nestedtensor backend #30979
已合并
culechan创建于 2月25日
从culechan/pytorchcherry-pick-mr-30508-1771989684490-auto合入到Ascend/pytorchv2.8.0
共 2 个文件变更+159-62
Mtest/test_nestedtensor.py+23-11
@@ -23,7 +23,7 @@ class TestNestedTensor(TestCase):
23 length = np.random.randint(1, max_seq_len)23 length = np.random.randint(1, max_seq_len)
24 row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))24 row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))
25 data.append(row)25 data.append(row)
26- nested_tensor_ref_list.append(torch.tensor(row))26+ nested_tensor_ref_list.append(torch.tensor(row).npu())
27 nested_tensor = torch.nested.nested_tensor(data)27 nested_tensor = torch.nested.nested_tensor(data)
28 nested_tensor_list = nested_tensor.unbind()28 nested_tensor_list = nested_tensor.unbind()
29 for i in range(batch_size):29 for i in range(batch_size):
@@ -43,7 +43,7 @@ class TestNestedTensor(TestCase):
43 row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))43 row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))
44 row = [list(item * np.arange(max_seq_len)) for item in row]44 row = [list(item * np.arange(max_seq_len)) for item in row]
45 data.append(row)45 data.append(row)
46- nested_tensor_ref_list.append(torch.tensor(row))46+ nested_tensor_ref_list.append(torch.tensor(row).npu())
47 nested_tensor = torch.nested.nested_tensor(data)47 nested_tensor = torch.nested.nested_tensor(data)
48 nested_tensor_list = nested_tensor.unbind()48 nested_tensor_list = nested_tensor.unbind()
49 for i in range(batch_size):49 for i in range(batch_size):
@@ -63,45 +63,57 @@ class TestNestedTensor(TestCase):
63 row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))63 row = list(np.random.randint(low=0, high=vocab_size, size=(length,)))
64 row = [list(item * np.arange(max_seq_len)) for item in row]64 row = [list(item * np.arange(max_seq_len)) for item in row]
65 data.append(row)65 data.append(row)
66- nested_tensor_ref_list.append(torch.tensor(row))66+ nested_tensor_ref_list.append(torch.tensor(row).npu())
67 nested_tensor = torch.nested.nested_tensor(data)67 nested_tensor = torch.nested.nested_tensor(data)
68 nested_tensor_list = nested_tensor.unbind()68 nested_tensor_list = nested_tensor.unbind()
69 for i in range(batch_size):69 for i in range(batch_size):
70 self.assertEqual(nested_tensor_list[i], nested_tensor_ref_list[i].type(torch.float32))70 self.assertEqual(nested_tensor_list[i], nested_tensor_ref_list[i].type(torch.float32))
71 71 
72 def _test_unbind_case(self, a, b):72 def _test_unbind_case(self, a, b):
73- nt = torch.nested.nested_tensor([a, b], dtype=a.dtype)73+ nt = torch.nested.nested_tensor([a.npu(), b.npu()], dtype=a.dtype)
74 nt_list = nt.unbind()74 nt_list = nt.unbind()
75 self.assertEqual(len(nt_list), 2)75 self.assertEqual(len(nt_list), 2)
76 self.assertEqual(nt_list[0], a)76 self.assertEqual(nt_list[0], a)
77 self.assertEqual(nt_list[1], b)77 self.assertEqual(nt_list[1], b)
78+ 
79+ def _test_asnested_case(self, a, b):
80+ nt = torch.nested.nested_tensor([a.npu(), b.npu()], dtype=a.dtype)
81+ nt_as = torch.nested.as_nested_tensor([a.npu(), b.npu()], dtype=a.dtype)
82+ nt_list = nt.unbind()
83+ nt_as_list = nt_as.unbind()
84+ self.assertEqual(len(nt_as_list), len(nt_list))
85+ self.assertEqual(nt_as_list[0], nt_list[0])
86+ self.assertEqual(nt_as_list[1], nt_list[1])
78 87
79- def test_unbind_case1(self):88+ def test_unbind_and_asnested_int64(self):
80 a = torch.tensor([[1, 2, 3], [4, 5, 6]])89 a = torch.tensor([[1, 2, 3], [4, 5, 6]])
81 b = torch.tensor([[7, 8], [10, 11]])90 b = torch.tensor([[7, 8], [10, 11]])
82 self._test_unbind_case(a, b)91 self._test_unbind_case(a, b)
92+ self._test_asnested_case(a, b)
83 93
84- def test_unbind_case2(self):94+ def test_unbind_and_asnested_float32(self):
85 a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)95 a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
86 b = torch.tensor([[7, 8], [10, 11]], dtype=torch.float32)96 b = torch.tensor([[7, 8], [10, 11]], dtype=torch.float32)
87 self._test_unbind_case(a, b)97 self._test_unbind_case(a, b)
98+ self._test_asnested_case(a, b)
88 99 
89- def test_unbind_case3(self):100+ def test_unbind_and_asnested_empty(self):
90 a = torch.tensor([[], []])101 a = torch.tensor([[], []])
91 b = torch.tensor([[], [], []])102 b = torch.tensor([[], [], []])
92 self._test_unbind_case(a, b)103 self._test_unbind_case(a, b)
104+ self._test_asnested_case(a, b)
93 105 
94 def test_default_options_nested_tensor(self):106 def test_default_options_nested_tensor(self):
95- default_nested_tensor = torch.nested.nested_tensor([])107+ default_nested_tensor = torch.nested.nested_tensor([], device="npu:0")
96- default_tensor = torch.tensor([])108+ default_tensor = torch.tensor([]).npu()
97 self.assertEqual(default_nested_tensor.dtype, default_tensor.dtype)109 self.assertEqual(default_nested_tensor.dtype, default_tensor.dtype)
98 self.assertEqual(default_nested_tensor.device, default_tensor.device)110 self.assertEqual(default_nested_tensor.device, default_tensor.device)
99 self.assertEqual(default_nested_tensor.layout, default_tensor.layout)111 self.assertEqual(default_nested_tensor.layout, default_tensor.layout)
100 self.assertEqual(default_nested_tensor.dim(), default_tensor.dim())112 self.assertEqual(default_nested_tensor.dim(), default_tensor.dim())
101 self.assertEqual(default_nested_tensor.requires_grad, default_tensor.requires_grad)113 self.assertEqual(default_nested_tensor.requires_grad, default_tensor.requires_grad)
102 114
103- def test_nested_tensor_size(self):115+ def test_nested_tensor_errsize(self):
104- nt = torch.nested.nested_tensor([torch.tensor([[1, 2, 3], [4, 5, 6]]), torch.tensor([[7, 8], [10, 11], [12, 13]])])116+ nt = torch.nested.nested_tensor([torch.tensor([[1, 2, 3], [4, 5, 6]]).npu(), torch.tensor([[7, 8], [10, 11], [12, 13]]).npu()])
105 self.assertEqual(nt.size(0), 2)117 self.assertEqual(nt.size(0), 2)
106 self.assertRaisesRegex(RuntimeError,118 self.assertRaisesRegex(RuntimeError,
107 "Given dimension 1 is irregular and does not have a size", 119 "Given dimension 1 is irregular and does not have a size",
Mtorchnpugen/gen_backend_stubs.py+136-51
@@ -14,6 +14,7 @@
14# See the License for the specific language governing permissions and14# See the License for the specific language governing permissions and
15# limitations under the License.15# limitations under the License.
16 16 
17+from dataclasses import dataclass
17import pathlib18import pathlib
18import argparse19import argparse
19import os20import os
@@ -209,7 +210,8 @@ def parse_backend_yaml(
209 if not isinstance(yaml_values, dict):210 if not isinstance(yaml_values, dict):
210 raise TypeError("yaml_values is not dict")211 raise TypeError("yaml_values is not dict")
211 212 
212- valid_keys = ['backend', 'cpp_namespace', 'supported', 'autograd', 'custom', 'custom_autograd', 'symint', 'quant']213+ # NestedTensor is a stub for now. (Not yet implemented.)
214+ valid_keys = ['backend', 'cpp_namespace', 'supported', 'autograd', 'custom', 'custom_autograd', 'symint', 'quant', 'nestedtensor']
213 215 
214 yaml_backend = yaml_values.pop('backend', None)216 yaml_backend = yaml_values.pop('backend', None)
215 true_backend = 'PrivateUse1' if yaml_backend == 'NPU' else yaml_backend217 true_backend = 'PrivateUse1' if yaml_backend == 'NPU' else yaml_backend
@@ -273,6 +275,11 @@ def parse_backend_yaml(
273 raise TypeError(f'expected "quant" to be a list, but got: {quant}')275 raise TypeError(f'expected "quant" to be a list, but got: {quant}')
274 quant = [op['func'].split("(")[0] if isinstance(op, Dict) else op for op in quant]276 quant = [op['func'].split("(")[0] if isinstance(op, Dict) else op for op in quant]
275 277 
278+ nestedtensor = yaml_values.pop('nestedtensor', [])
279+ if not isinstance(nestedtensor, list):
280+ raise TypeError(f'expected "nestedtensor" to be a list, but got: {nestedtensor}')
281+ nestedtensor = [op['func'].split("(")[0] if isinstance(op, Dict) else op for op in nestedtensor]
282+ 
276 # custom_supported is only supported for filt expose api, and is not useful here.283 # custom_supported is only supported for filt expose api, and is not useful here.
277 yaml_values.pop('custom_supported', [])284 yaml_values.pop('custom_supported', [])
278 if (len(yaml_values.keys()) > 0):285 if (len(yaml_values.keys()) > 0):
@@ -314,6 +321,13 @@ the behavior of autograd for some operators on your backend. However "Autograd{b
314 raise KeyError("quant_key should not be in backend_indices.")321 raise KeyError("quant_key should not be in backend_indices.")
315 backend_indices[str(backend_key) + quant_key] = quant_idx322 backend_indices[str(backend_key) + quant_key] = quant_idx
316 323 
324+ nestedtensor_key = "Nestedtensor"
325+ if len(nestedtensor) > 0:
326+ nestedtensor_idx = create_backend_index(nestedtensor, symint_set, backend_key, native_functions_map, cpp_namespace)
327+ if nestedtensor_key in backend_indices:
328+ raise KeyError("nestedtensor_key should not be in backend_indices.")
329+ backend_indices[str(backend_key) + nestedtensor_key] = nestedtensor_idx
330+ 
317 # check_grouped_native_functions(backend_key, autograd_key, backend_indices, grouped_native_functions)331 # check_grouped_native_functions(backend_key, autograd_key, backend_indices, grouped_native_functions)
318 return ParsedExternalYaml(true_backend, backend_key, autograd_key, cpp_namespace, backend_indices)332 return ParsedExternalYaml(true_backend, backend_key, autograd_key, cpp_namespace, backend_indices)
319 333 
@@ -535,61 +549,131 @@ m.impl("${schema}", TORCH_FN(at::native::${kernel}));"""
535 })549 })
536 550 
537 551 
552+# 定义配置数据类
553+@dataclass
554+class SpecialRegisterConfig:
555+ dispatch_key: str
556+ filename: str
557+ header: str
558+ extra_impls: List[str]
559+ 
560+KERNEL_TEMPLATE = CodeTemplate("""\
561+m.impl("${schema}", TORCH_FN(op_plugin::${kernel}));""")
562+ 
563+ 
564+def _gen_special_registration_body(
565+ backend_indices: BackendIndex,
566+ config: SpecialRegisterConfig,
567+) -> str:
568+ """生成特殊注册的主体内容"""
569+ kernel_regs = [
570+ KERNEL_TEMPLATE.substitute(schema=op_name, kernel=metadata.kernel)
571+ for op_name, metadata in backend_indices.index.items()
572+ ]
573+
574+ template = CodeTemplate("""\
575+TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) {
576+$kernel_registrations
577+$extra_impls
578+};""")
579+
580+ return template.substitute(
581+ dispatch_key=config.dispatch_key,
582+ kernel_registrations=kernel_regs,
583+ extra_impls="\n".join(config.extra_impls),
584+ )
585+ 
586+ 
587+def _write_special_register(
588+ fm: FileManager,
589+ config: SpecialRegisterConfig,
590+ static_init_dispatch_registrations: str,
591+) -> None:
592+ """写入特殊注册文件"""
593+ ns_helper = NamespaceHelper(namespace_str="at")
594+
595+ dispatch_definitions = fm.substitute_with_template(
596+ 'RegisterDispatchDefinitions.ini',
597+ lambda: {
598+ 'ns_prologue': ns_helper.prologue,
599+ 'ns_epilogue': ns_helper.epilogue,
600+ 'static_init_dispatch_registrations': static_init_dispatch_registrations,
601+ 'deferred_dispatch_registrations': '',
602+ 'dispatch_namespace': '',
603+ 'dispatch_namespaced_definitions': '',
604+ 'dispatch_anonymous_definitions': '',
605+ },
606+ ).split('\n')
607+
608+ fm.write_with_template(
609+ f'{config.filename}.cpp',
610+ 'RegisterDispatchKey.cpp',
611+ lambda: {
612+ 'extra_cuda_headers': '',
613+ 'external_backend_headers': config.header,
614+ 'namespaced_headers': '',
615+ 'DispatchKey': 'NPU',
616+ 'dispatch_headers': '',
617+ 'ops_headers': '',
618+ 'dispatch_helpers': '',
619+ 'dispatch_definitions': dispatch_definitions,
620+ }
621+ )
622+ 
623+SPECIAL_REGISTERS = {
624+ 'quantize': SpecialRegisterConfig(
625+ dispatch_key="QuantizedPrivateUse1",
626+ filename="QuantizedRegister",
627+ header='''\
628+#include <ATen/ops/quantize_per_tensor.h>
629+#include "op_plugin/OpInterface.h"
630+''',
631+ extra_impls=[
632+ 'm.impl("q_scale", TORCH_FN(at::native::q_scale_quant));',
633+ 'm.impl("q_per_channel_scales", TORCH_FN(at::native::q_per_channel_scales));',
634+ 'm.impl("q_zero_point", TORCH_FN(at::native::q_zero_point_quant));',
635+ 'm.impl("q_per_channel_zero_points", TORCH_FN(at::native::q_per_channel_zero_points));',
636+ 'm.impl("q_per_channel_axis", TORCH_FN(at::native::q_per_channel_axis));',
637+ 'm.impl("qscheme", TORCH_FN(at::native::qscheme_quant));',
638+ ],
639+ ),
640+ 'nestedtensor': SpecialRegisterConfig(
641+ dispatch_key="NestedTensorPrivateUse1",
642+ filename="NestedTensorRegister",
643+ header="",
644+ extra_impls=[
645+ 'm.impl("unbind.int", TORCH_FN(at::native::NestedTensor_unbind));',
646+ 'm.impl("values", TORCH_FN(at::native::values_nested));',
647+ 'm.impl("_nested_tensor_size", TORCH_FN(at::native::_nested_tensor_size));',
648+ ],
649+ ),
650+}
651+ 
652+ 
538def gen_quantize_register(653def gen_quantize_register(
539 fm: FileManager,654 fm: FileManager,
540 backend_indices: BackendIndex,655 backend_indices: BackendIndex,
541-):656+) -> None:
542- ns_helper = NamespaceHelper(namespace_str="at")657+ """生成量化注册"""
658+ config = SPECIAL_REGISTERS['quantize']
659+ static_init = _gen_special_registration_body(
660+ backend_indices["NPUQuantize"],
661+ config,
662+ )
663+ _write_special_register(fm, config, static_init)
543 664 
544- quantize_dict: Dict[str, str] = {}
545- for op_name, metadata in backend_indices.index.items():
546- quantize_dict[op_name] = metadata.kernel
547 665 
548- native_func_header = """\666+def gen_nestedtensor_register(
549-#include <ATen/ops/quantize_per_tensor.h>667+ fm: FileManager,
550-#include "op_plugin/OpInterface.h"668+ backend_indices: BackendIndex,
551-"""669+) -> None:
552- static_template = CodeTemplate(670+ """生成嵌套张量注册"""
553- """\671+ config = SPECIAL_REGISTERS['nestedtensor']
554-TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) {672+ static_init = _gen_special_registration_body(
555-$dispatch_registrations_body673+ backend_indices["NPUNestedtensor"],
556-m.impl("q_scale", TORCH_FN(at::native::q_scale_quant));674+ config,
557-m.impl("q_per_channel_scales", TORCH_FN(at::native::q_per_channel_scales));
558-m.impl("q_zero_point", TORCH_FN(at::native::q_zero_point_quant));
559-m.impl("q_per_channel_zero_points", TORCH_FN(at::native::q_per_channel_zero_points));
560-m.impl("q_per_channel_axis", TORCH_FN(at::native::q_per_channel_axis));
561-m.impl("qscheme", TORCH_FN(at::native::qscheme_quant));
562-};"""
563 )675 )
564- kernel_template = CodeTemplate(676+ _write_special_register(fm, config, static_init)
565- """\
566-m.impl("${schema}", TORCH_FN(op_plugin::${kernel}));"""
567- )
568- static_init_dispatch_registrations = static_template.substitute(
569- dispatch_key="QuantizedPrivateUse1",
570- dispatch_registrations_body=[kernel_template.substitute(schema=kv[0], kernel=kv[1]) for kv in quantize_dict.items()]
571- )
572- fm.write_with_template(f'QuantizedRegister.cpp', 'RegisterDispatchKey.cpp', lambda: {
573- 'extra_cuda_headers': '',
574- 'external_backend_headers': native_func_header,
575- 'namespaced_headers': '',
576- 'DispatchKey': 'NPU',
577- 'dispatch_headers': '',
578- 'ops_headers': '',
579- 'dispatch_helpers': '',
580- 'dispatch_definitions': fm.substitute_with_template(
581- 'RegisterDispatchDefinitions.ini',
582- lambda: {
583- 'ns_prologue': ns_helper.prologue,
584- 'ns_epilogue': ns_helper.epilogue,
585- 'static_init_dispatch_registrations': static_init_dispatch_registrations,
586- 'deferred_dispatch_registrations': '',
587- 'dispatch_namespace': '',
588- 'dispatch_namespaced_definitions': '',
589- 'dispatch_anonymous_definitions': '',
590- },
591- ).split('\n'),
592- })
593 677 
594 678 
595def gen_functionalization(fm: FileManager,679def gen_functionalization(fm: FileManager,
@@ -734,7 +818,8 @@ def run(source_yaml: str, output_dir: str, dry_run: bool,
734 register_dispatch_key_func=dest.RegisterDispatchKey,818 register_dispatch_key_func=dest.RegisterDispatchKey,
735 )819 )
736 820 
737- gen_quantize_register(fm, backend_indices=backend_indices["NPUQuantize"])821+ gen_quantize_register(fm, backend_indices)
822+ gen_nestedtensor_register(fm, backend_indices)
738 823 
739 pta_template_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "templates")824 pta_template_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "templates")
740 fm = FileManager(install_dir=output_dir, template_dir=pta_template_dir, dry_run=dry_run)825 fm = FileManager(install_dir=output_dir, template_dir=pta_template_dir, dry_run=dry_run)