AI assist developer for python dt fourth batch for 2.8.0-pytorch-AtomGit

Atest/_inductor/test_npu_device.py+81-0

@@ -0,0 +1,81 @@
	1	+import torch
	2	+from torch.testing._internal.common_utils import TestCase, run_tests
	3	+from torch_npu.npu import device_count
	4	+from torch_npu.utils._dynamo_device import NpuInterface, current_device, set_device
	5	+from torch_npu.utils._inductor import NPUDeviceOpOverrides
	6	+from torch_npu._inductor.config import config as npu_config
	7	+from torch_npu._inductor.npu_device import NewNPUDeviceOpOverrides
	8	+
	9	+
	10	+class TestNpuDevice(TestCase):
	11	+ def test_aoti_get_stream(self):
	12	+ overrides = NewNPUDeviceOpOverrides()
	13	+ result = overrides.aoti_get_stream()
	14	+ excepted = "aoti_torch_get_current_cuda_stream"
	15	+ self.assertEqual(result, excepted)
	16	+
	17	+ def test_cpp_stream_type(self):
	18	+ overrides = NewNPUDeviceOpOverrides()
	19	+ result = overrides.cpp_stream_type()
	20	+ excepted = "aclrtStream"
	21	+ self.assertEqual(result, excepted)
	22	+
	23	+ def test_abi_compatible_header(self):
	24	+ overrides = NewNPUDeviceOpOverrides()
	25	+ result = overrides.abi_compatible_header()
	26	+ self.assertIn("#include <fstream>", result)
	27	+ self.assertIn("#include <vector>", result)
	28	+ self.assertIn("#include <iostream>", result)
	29	+ self.assertIn("#include <string>", result)
	30	+ self.assertIn("#include <tuple>", result)
	31	+ self.assertIn("#include <unordered_map>", result)
	32	+ self.assertIn("#include <memory>", result)
	33	+ self.assertIn("#include <filesystem>", result)
	34	+ self.assertIn("#include <assert.h>", result)
	35	+ self.assertIn("#include <stdbool.h>", result)
	36	+ self.assertIn("#include <sys/syscall.h>", result)
	37	+ self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)
	38	+ self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)
	39	+ self.assertIn("#include \"experiment/runtime/runtime/rt.h\"", result)
	40	+
	41	+ def test_cpp_aoti_stream_guard(self):
	42	+ overrides = NewNPUDeviceOpOverrides()
	43	+ result = overrides.cpp_aoti_stream_guard()
	44	+ excepted = "AOTICudaStreamGuard"
	45	+ self.assertEqual(result, excepted)
	46	+
	47	+ def test_cpp_aoti_device_guard_not_implemented(self):
	48	+ overrides = NewNPUDeviceOpOverrides()
	49	+ with self.assertRaises(NotImplementedError):
	50	+ overrides.cpp_aoti_device_guard()
	51	+
	52	+ def test_device_guard(self):
	53	+ overrides = NewNPUDeviceOpOverrides()
	54	+ result = overrides.device_guard(0)
	55	+ excepted = "torch.npu.utils.device(0)"
	56	+ self.assertEqual(result, excepted)
	57	+
	58	+ def test_synchronize(self):
	59	+ overrides = NewNPUDeviceOpOverrides()
	60	+ result = overrides.synchronize()
	61	+ excepted = """
	62	+ stream = torch.npu.current_stream()
	63	+ stream.synchronize()
	64	+ """
	65	+ self.assertEqual(result, excepted)
	66	+
	67	+ def test_set_device(self):
	68	+ overrides = NewNPUDeviceOpOverrides()
	69	+ result = overrides.set_device(0)
	70	+ excepted = "torch.npu.set_device(0)"
	71	+ self.assertEqual(result, excepted)
	72	+
	73	+ def test_import_get_raw_stream_as(self):
	74	+ overrides = NewNPUDeviceOpOverrides()
	75	+ result = overrides.import_get_raw_stream_as("test_name")
	76	+ excepted = "from torch_npu._inductor import get_current_raw_stream as test_name"
	77	+ self.assertEqual(result, excepted)
	78	+
	79	+
	80	+if __name__ == "__main__":
	81	+ run_tests()

Atest/_inductor/test_npu_fusion_attention_graph.py+97-0

@@ -0,0 +1,97 @@
	1	+import functools
	2	+import sympy
	3	+import torch
	4	+import torch.nn.functional as F
	5	+from torch.autograd import Function
	6	+from torch.library import Library, impl
	7	+from torch.testing._internal.common_utils import TestCase, run_tests
	8	+import torch_npu
	9	+from torch_npu._inductor.npu_fusion_attention_graph import NpuGraphAttentionFunction
	10	+
	11	+
	12	+class TestNpuFusionAttentionGraph(TestCase):
	13	+ def test_npu_graph_attention_function(self):
	14	+ query = torch.randn(2, 4, 8, 16, device='npu', requires_grad=True)
	15	+ key = torch.randn(2, 4, 8, 16, device='npu')
	16	+ value = torch.randn(2, 4, 8, 16, device='npu')
	17	+ head_num = 4
	18	+ input_layout = "BNSD"
	19	+
	20	+ output = NpuGraphAttentionFunction.apply(
	21	+ query, key, value, head_num, input_layout
	22	+ )
	23	+
	24	+ self.assertEqual(output[0].shape, query.shape)
	25	+ self.assertEqual(output[1].shape, (2, 4, 8, 8))
	26	+ self.assertEqual(output[2].shape, (2, 4, 8, 8))
	27	+ self.assertEqual(output[3].shape, (0,))
	28	+
	29	+ grad_outputs = (
	30	+ torch.randn_like(output[0]),
	31	+ torch.randn_like(output[1]),
	32	+ torch.randn_like(output[2]),
	33	+ torch.randn_like(output[3]),
	34	+ torch.randn(1, device='npu'),
	35	+ torch.randn(1, device='npu'),
	36	+ torch.randn(1, device='npu')
	37	+ )
	38	+
	39	+ output[0].backward(grad_outputs[0])
	40	+
	41	+ def test_npu_fa_forward_scale_handling(self):
	42	+ query = torch.randn(2, 4, 8, 16, device='npu')
	43	+ key = torch.randn(2, 4, 8, 16, device='npu')
	44	+ value = torch.randn(2, 4, 8, 16, device='npu')
	45	+ head_num = 4
	46	+ input_layout = "BNSD"
	47	+
	48	+ result = torch.ops.npu_graph.npu_fa(
	49	+ query, key, value, head_num, input_layout, scale=2.0
	50	+ )
	51	+
	52	+ self.assertEqual(result[0].shape, query.shape)
	53	+ self.assertEqual(result[1].shape, (2, 4, 8, 8))
	54	+ self.assertEqual(result[2].shape, (2, 4, 8, 8))
	55	+ self.assertEqual(result[3].shape, (0,))
	56	+
	57	+ def test_npu_fa_backward_meta_impl(self):
	58	+ query = torch.randn(2, 4, 8, 16, device='meta')
	59	+ key = torch.randn(2, 4, 8, 16, device='meta')
	60	+ value = torch.randn(2, 4, 8, 16, device='meta')
	61	+ dy = torch.randn(2, 4, 8, 16, device='meta')
	62	+ head_num = 4
	63	+ input_layout = "BSH"
	64	+
	65	+ result = torch.ops.npu_graph.npu_fa_backward(
	66	+ query, key, value, dy, head_num, input_layout
	67	+ )
	68	+
	69	+ self.assertEqual(result[0].shape, query.shape)
	70	+ self.assertEqual(result[1].shape, key.shape)
	71	+ self.assertEqual(result[2].shape, value.shape)
	72	+ self.assertIsNone(result[3])
	73	+
	74	+ def test_npu_fa_backward_scale_value_handling(self):
	75	+ query = torch.randn(1, 2, 4, 8, device='npu', requires_grad=True)
	76	+ key = torch.randn(1, 2, 4, 8, device='npu')
	77	+ value = torch.randn(1, 2, 4, 8, device='npu')
	78	+ dy = torch.randn(1, 2, 4, 8, device='npu')
	79	+ head_num = 2
	80	+ input_layout = "BNSD"
	81	+
	82	+ try:
	83	+ result = torch.ops.npu_graph.npu_fa_backward(
	84	+ query, key, value, dy, head_num, input_layout, scale_value=2.0
	85	+ )
	86	+ self.assertEqual(result[0].shape, query.shape)
	87	+ self.assertEqual(result[1].shape, key.shape)
	88	+ self.assertEqual(result[2].shape, value.shape)
	89	+ except RuntimeError as e:
	90	+ if "aclnnFlashAttentionScoreGrad" in str(e):
	91	+ pass
	92	+ else:
	93	+ raise e
	94	+
	95	+
	96	+if __name__ == "__main__":
	97	+ run_tests()

Atest/_inductor/test_npu_static_kernel.py+60-0

@@ -0,0 +1,60 @@
	1	+import os
	2	+import subprocess
	3	+import datetime
	4	+from pathlib import Path
	5	+import stat
	6	+from torch.testing._internal.common_utils import TestCase, run_tests
	7	+import torch_npu
	8	+from torch_npu._inductor.config import log
	9	+from torch_npu._inductor.npu_static_kernel import StaticKernelCompiler
	10	+from torch_npu._inductor.npu_static_kernel import safe_resolve_output_dir
	11	+
	12	+
	13	+class TestNpuStaticKernel(TestCase):
	14	+ def test_safe_resolve_output_dir_dot_dot(self):
	15	+ with self.assertRaises(ValueError):
	16	+ safe_resolve_output_dir("test/../dir")
	17	+
	18	+ def test_safe_resolve_output_dir_null_byte(self):
	19	+ with self.assertRaises(ValueError):
	20	+ safe_resolve_output_dir("test/../x00dir")
	21	+
	22	+ def test_safe_resolve_output_dir_permission_error(self):
	23	+ import tempfile
	24	+ with tempfile.TemporaryDirectory() as tmpdir:
	25	+ readonly_dir = Path(tmpdir) / "readonly"
	26	+ readonly_dir.mkdir()
	27	+ readonly_dir.chmod(stat.S_IRUSR \| stat.S_IXUSR)
	28	+ build_path = str(readonly_dir / 'subdir')
	29	+ with self.assertRaises(RuntimeError):
	30	+ safe_resolve_output_dir(build_path)
	31	+
	32	+ def test_safe_resolve_output_dir_symlink(self):
	33	+ import tempfile
	34	+ with tempfile.TemporaryDirectory() as tmpdir:
	35	+ real_dir = Path(tmpdir) / "real"
	36	+ real_dir.mkdir()
	37	+ symlink_dir = Path(tmpdir) / "symlink"
	38	+ symlink_dir.symlink_to(real_dir)
	39	+ build_path = str(symlink_dir / "subdir")
	40	+ with self.assertRaises(ValueError):
	41	+ safe_resolve_output_dir(build_path)
	42	+
	43	+ def test_uninstall_static_kernel_no_path(self):
	44	+ from torch_npu._inductor.npu_static_kernel import _uninstall_path, uninstall_static_kernel
	45	+ _uninstall_path = None
	46	+ uninstall_static_kernel()
	47	+
	48	+ def test_safe_resolve_output_dir_absolute_path(self):
	49	+ import tempfile
	50	+ import shutil
	51	+ with tempfile.TemporaryDirectory() as tmpdir:
	52	+ abs_dir = Path(tmpdir) / "test_build"
	53	+ abs_dir.mkdir()
	54	+ result = safe_resolve_output_dir(str(abs_dir))
	55	+ self.assertTrue(result.exists())
	56	+ self.assertIn("kernel_aot_optimization_build_outputs", str(result))
	57	+
	58	+
	59	+if __name__ == "__main__":
	60	+ run_tests()