已合并
add_random_op #30653
Ambi创建于 2月9日
add_random_op #30653
已合并
从已删除 :v2.8.0合入到Ascend/pytorchv2.8.0
共 6 个文件变更+380-0
| @@ -0,0 +1,56 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + parametrize, | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | +) | ||
| 7 | +from testutils import TestUtils | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestRngprimsPhiloxRand(TestUtils): | ||
| 12 | + def op_calc(self, x): | ||
| 13 | + size = list(x.shape) | ||
| 14 | + | ||
| 15 | + seed = torch.tensor(1234, device=x.device, dtype=torch.int64) | ||
| 16 | + offset = torch.tensor(0, device=x.device, dtype=torch.int64) | ||
| 17 | + | ||
| 18 | + rand, new_offset = torch.ops.rngprims.philox_rand( | ||
| 19 | + size, # SymInt[] | ||
| 20 | + seed, # Tensor | ||
| 21 | + offset, # Tensor | ||
| 22 | + None, # stride | ||
| 23 | + x.device, # device | ||
| 24 | + x.dtype, # dtype | ||
| 25 | + ) | ||
| 26 | + | ||
| 27 | + return rand * 2.0 + x | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + def test_philox_rand_eager_vs_inductor(self, shape, dtype): | ||
| 32 | + device = "npu" | ||
| 33 | + | ||
| 34 | + x = torch.ones(shape, device=device, dtype=dtype) | ||
| 35 | + torch.manual_seed(0) | ||
| 36 | + y1_eager = self.op_calc(x) | ||
| 37 | + y2_eager = self.op_calc(x) | ||
| 38 | + | ||
| 39 | + compiled_op = torch.compile( | ||
| 40 | + self.op_calc, | ||
| 41 | + backend="inductor", | ||
| 42 | + fullgraph=True, | ||
| 43 | + dynamic=False, | ||
| 44 | + ) | ||
| 45 | + | ||
| 46 | + y1_ind = compiled_op(x) | ||
| 47 | + y2_ind = compiled_op(x) | ||
| 48 | + | ||
| 49 | + self.assertEqual(y1_eager, y2_eager) | ||
| 50 | + self.assertEqual(y1_ind, y2_ind) | ||
| 51 | + self.assertEqual(y1_eager, y1_ind) | ||
| 52 | + | ||
| 53 | +instantiate_parametrized_tests(TestRngprimsPhiloxRand) | ||
| 54 | + | ||
| 55 | +if __name__ == "__main__": | ||
| 56 | + run_tests() | ||
| @@ -0,0 +1,56 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + parametrize, | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | +) | ||
| 7 | +from testutils import TestUtils | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestRngprimsPhiloxRandDefault(TestUtils): | ||
| 12 | + def op_calc(self, x): | ||
| 13 | + size = list(x.shape) | ||
| 14 | + | ||
| 15 | + seed = torch.tensor(1234, device=x.device, dtype=torch.int64) | ||
| 16 | + offset = torch.tensor(0, device=x.device, dtype=torch.int64) | ||
| 17 | + | ||
| 18 | + rand, new_offset = torch.ops.rngprims.philox_rand.default( | ||
| 19 | + size, # SymInt[] | ||
| 20 | + seed, # Tensor | ||
| 21 | + offset, # Tensor | ||
| 22 | + None, # stride | ||
| 23 | + x.device, # device | ||
| 24 | + x.dtype, # dtype | ||
| 25 | + ) | ||
| 26 | + | ||
| 27 | + return rand * 2.0 + x | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + def test_philox_rand_eager_vs_inductor(self, shape, dtype): | ||
| 32 | + device = "npu" | ||
| 33 | + | ||
| 34 | + x = torch.ones(shape, device=device, dtype=dtype) | ||
| 35 | + torch.manual_seed(0) | ||
| 36 | + y1_eager = self.op_calc(x) | ||
| 37 | + y2_eager = self.op_calc(x) | ||
| 38 | + | ||
| 39 | + compiled_op = torch.compile( | ||
| 40 | + self.op_calc, | ||
| 41 | + backend="inductor", | ||
| 42 | + fullgraph=True, | ||
| 43 | + dynamic=False, | ||
| 44 | + ) | ||
| 45 | + | ||
| 46 | + y1_ind = compiled_op(x) | ||
| 47 | + y2_ind = compiled_op(x) | ||
| 48 | + | ||
| 49 | + self.assertEqual(y1_eager, y2_eager) | ||
| 50 | + self.assertEqual(y1_ind, y2_ind) | ||
| 51 | + self.assertEqual(y1_eager, y1_ind) | ||
| 52 | + | ||
| 53 | +instantiate_parametrized_tests(TestRngprimsPhiloxRandDefault) | ||
| 54 | + | ||
| 55 | +if __name__ == "__main__": | ||
| 56 | + run_tests() | ||
| @@ -0,0 +1,48 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + parametrize, | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | +) | ||
| 7 | +from testutils import TestUtils | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestRunAndSaveRngState(TestUtils): | ||
| 12 | + def op_calc(self, like, device, dtype): | ||
| 13 | + rng_state1, res1 = torch._prims.rng_prims.run_and_save_rng_state( | ||
| 14 | + torch.ops.aten.rand_like.default, | ||
| 15 | + like, | ||
| 16 | + device=device, | ||
| 17 | + dtype=dtype, | ||
| 18 | + ) | ||
| 19 | + | ||
| 20 | + torch_npu.npu.set_rng_state(rng_state1) | ||
| 21 | + | ||
| 22 | + rng_state2, res2 = torch._prims.rng_prims.run_and_save_rng_state( | ||
| 23 | + torch.ops.aten.rand_like.default, | ||
| 24 | + like, | ||
| 25 | + device=device, | ||
| 26 | + dtype=dtype, | ||
| 27 | + ) | ||
| 28 | + | ||
| 29 | + return rng_state1, res1, rng_state2, res2 | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + def test_rng_state_with_compile(self, shape, dtype): | ||
| 34 | + device = "npu" | ||
| 35 | + | ||
| 36 | + like = torch.empty(shape, device=device, dtype=dtype) | ||
| 37 | + | ||
| 38 | + # eager | ||
| 39 | + rng_state1_eager, res1_eager, rng_state2_eager, res2_eager = \ | ||
| 40 | + self.op_calc(like, device, dtype) | ||
| 41 | + | ||
| 42 | + self.assertEqual(res1_eager, res2_eager) | ||
| 43 | + self.assertTrue(torch.equal(rng_state1_eager, rng_state2_eager)) | ||
| 44 | + | ||
| 45 | +instantiate_parametrized_tests(TestRunAndSaveRngState) | ||
| 46 | + | ||
| 47 | +if __name__ == "__main__": | ||
| 48 | + run_tests() | ||
| @@ -0,0 +1,55 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + parametrize, | ||
| 5 | + instantiate_parametrized_tests, | ||
| 6 | +) | ||
| 7 | +from testutils import TestUtils | ||
| 8 | +import torch_npu | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestRunWithRngState(TestUtils): | ||
| 12 | + def op_calc(self, current_state, like, device, dtype): | ||
| 13 | + res1 = torch._prims.rng_prims.run_with_rng_state( | ||
| 14 | + current_state, | ||
| 15 | + torch.ops.aten.rand_like.default, | ||
| 16 | + like, | ||
| 17 | + device=device, | ||
| 18 | + dtype=dtype, | ||
| 19 | + ) | ||
| 20 | + | ||
| 21 | + res2 = torch._prims.rng_prims.run_with_rng_state( | ||
| 22 | + current_state, | ||
| 23 | + torch.ops.aten.rand_like.default, | ||
| 24 | + like, | ||
| 25 | + device=device, | ||
| 26 | + dtype=dtype, | ||
| 27 | + ) | ||
| 28 | + | ||
| 29 | + return res1, res2 | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + def test_rng_state_with_compile(self, shape, dtype): | ||
| 34 | + device = "npu" | ||
| 35 | + torch.manual_seed(0) | ||
| 36 | + current_state = torch_npu.npu.get_rng_state() | ||
| 37 | + like = torch.empty(shape, device=device, dtype=dtype) | ||
| 38 | + | ||
| 39 | + # eager | ||
| 40 | + res1_eager, res2_eager = \ | ||
| 41 | + self.op_calc(current_state, like, device, dtype) | ||
| 42 | + | ||
| 43 | + # compiled | ||
| 44 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||
| 45 | + res1_ind, res2_ind = \ | ||
| 46 | + compiled_op_calc(current_state, like, device, dtype) | ||
| 47 | + | ||
| 48 | + self.assertEqual(res1_eager, res2_eager) | ||
| 49 | + self.assertEqual(res1_ind, res2_ind) | ||
| 50 | + self.assertEqual(res1_ind, res1_eager) | ||
| 51 | + | ||
| 52 | +instantiate_parametrized_tests(TestRunWithRngState) | ||
| 53 | + | ||
| 54 | +if __name__ == "__main__": | ||
| 55 | + run_tests() | ||
| @@ -6,6 +6,7 @@ import torch | |||
| 6 | import torch.nn.functional as F | 6 | import torch.nn.functional as F |
| 7 | from torch.autograd import Function | 7 | from torch.autograd import Function |
| 8 | from torch.library import Library, impl | 8 | from torch.library import Library, impl |
| 9 | +from torch._inductor.pattern_matcher import init_once_fakemode | ||
| 9 | import torch_npu | 10 | import torch_npu |
| 10 | 11 | ||
| 11 | npu_def = Library("npu_graph", "DEF") | 12 | npu_def = Library("npu_graph", "DEF") |
| @@ -156,6 +157,7 @@ def npu_fusion_attention_graph(query, key, value, head_num, input_layout, pse=No | |||
| 156 | torch_npu.npu_fusion_attention_graph = npu_fusion_attention_graph | 157 | torch_npu.npu_fusion_attention_graph = npu_fusion_attention_graph |
| 157 | 158 | ||
| 158 | 159 | ||
| 160 | + | ||
| 159 | def register_fa_pass(): | 161 | def register_fa_pass(): |
| 160 | TOKEN_MAX = 2147483647 | 162 | TOKEN_MAX = 2147483647 |
| 161 | from torch._inductor.pattern_matcher import register_replacement, fwd_only, joint_fwd_bwd | 163 | from torch._inductor.pattern_matcher import register_replacement, fwd_only, joint_fwd_bwd |
| @@ -1,9 +1,11 @@ | |||
| 1 | +from typing import Optional | ||
| 1 | import operator | 2 | import operator |
| 2 | from functools import reduce | 3 | from functools import reduce |
| 3 | 4 | ||
| 4 | import torch | 5 | import torch |
| 5 | from torch._prims_common import TensorLike | 6 | from torch._prims_common import TensorLike |
| 6 | from torch._inductor.codegen.common import DeviceOpOverrides, register_device_op_overrides | 7 | from torch._inductor.codegen.common import DeviceOpOverrides, register_device_op_overrides |
| 8 | +from torch._prims.rng_prims import register_rng_prim | ||
| 7 | 9 | ||
| 8 | 10 | ||
| 9 | class NPUDeviceOpOverrides(DeviceOpOverrides): | 11 | class NPUDeviceOpOverrides(DeviceOpOverrides): |
| @@ -44,3 +46,164 @@ def _max_unpoolnd_patch( | |||
| 44 | ).view(output.shape) | 46 | ).view(output.shape) |
| 45 | 47 | ||
| 46 | torch._decomp.decompositions._max_unpoolnd = _max_unpoolnd_patch | 48 | torch._decomp.decompositions._max_unpoolnd = _max_unpoolnd_patch |
| 49 | + | ||
| 50 | + | ||
| 51 | +def patch_philox_rand_offset(): | ||
| 52 | + def get_philox_rand_offset_patch(shape): | ||
| 53 | + numel_scalar = 1 | ||
| 54 | + for dim_size in shape: | ||
| 55 | + numel_scalar *= dim_size | ||
| 56 | + numel = torch.scalar_tensor(numel_scalar, dtype=torch.int64) | ||
| 57 | + | ||
| 58 | + return numel | ||
| 59 | + torch._prims.rng_prims.philox_rand_offset = get_philox_rand_offset_patch | ||
| 60 | + | ||
| 61 | + | ||
| 62 | +def patch_register_philox_rand(): | ||
| 63 | + rng_prims = torch._prims.rng_prims | ||
| 64 | + philox_rand_offset_meta = rng_prims.philox_rand_offset_meta | ||
| 65 | + philox_rand_offset = rng_prims.philox_rand_offset | ||
| 66 | + make_contiguous_strides_for = rng_prims.make_contiguous_strides_for | ||
| 67 | + _prims = torch._prims | ||
| 68 | + _device = rng_prims._device | ||
| 69 | + _dtype = rng_prims._dtype | ||
| 70 | + CUDARngStateHelper = rng_prims.CUDARngStateHelper | ||
| 71 | + | ||
| 72 | + | ||
| 73 | + def get_register_philox_rand_patch(): | ||
| 74 | + name = "philox_rand" | ||
| 75 | + schema = "(SymInt[] size, Tensor seed, Tensor offset, int[]? stride, Device? device=None, ScalarType? dtype=None) -> (Tensor, Tensor)" # noqa: B950 | ||
| 76 | + | ||
| 77 | + | ||
| 78 | + def _philox_rand_meta( | ||
| 79 | + shape: torch.Size, | ||
| 80 | + seed: torch.Tensor, | ||
| 81 | + offset: torch.Tensor, | ||
| 82 | + stride: Optional[tuple[int, ...]], | ||
| 83 | + device: _device, | ||
| 84 | + dtype: _dtype, | ||
| 85 | + ): | ||
| 86 | + stride = make_contiguous_strides_for(shape) | ||
| 87 | + random_values = _prims.TensorMeta( | ||
| 88 | + shape=shape, strides=stride, dtype=dtype, device=device | ||
| 89 | + ) | ||
| 90 | + offset = philox_rand_offset_meta(shape) | ||
| 91 | + return (random_values, offset) | ||
| 92 | + | ||
| 93 | + | ||
| 94 | + def _philox_rand( | ||
| 95 | + shape: torch.Size, | ||
| 96 | + seed: torch.Tensor, | ||
| 97 | + offset: torch.Tensor, | ||
| 98 | + stride: Optional[tuple[int, ...]], | ||
| 99 | + device: _device, | ||
| 100 | + dtype: _dtype, | ||
| 101 | + ): | ||
| 102 | + if device.type == "cpu": | ||
| 103 | + devices = [] | ||
| 104 | + else: | ||
| 105 | + devices = [device] | ||
| 106 | + | ||
| 107 | + with torch.random.fork_rng(devices, device_type="npu"): | ||
| 108 | + CUDARngStateHelper.set_torch_state_tensor(seed, offset) | ||
| 109 | + random_values = torch.rand(shape, device=device, dtype=dtype) | ||
| 110 | + | ||
| 111 | + return random_values, philox_rand_offset(shape) | ||
| 112 | + | ||
| 113 | + | ||
| 114 | + register_rng_prim( | ||
| 115 | + name=name, | ||
| 116 | + schema=schema, | ||
| 117 | + impl_aten=_philox_rand, | ||
| 118 | + impl_meta=_philox_rand_meta, | ||
| 119 | + doc="Philox based stateless rand operator", | ||
| 120 | + tags=(torch.Tag.nondeterministic_seeded,), | ||
| 121 | + ) | ||
| 122 | + | ||
| 123 | + torch._prims.rng_prims.register_philox_rand = get_register_philox_rand_patch | ||
| 124 | + torch._prims.rng_prims.register_philox_rand() | ||
| 125 | + | ||
| 126 | + | ||
| 127 | +def patch_register_run_and_save_rng_state_op(): | ||
| 128 | + from torch._prims import rng_prims | ||
| 129 | + from torch._C import DispatchKey | ||
| 130 | + | ||
| 131 | + run_and_save_rng_state = getattr( | ||
| 132 | + rng_prims, "run_and_save_rng_state", None | ||
| 133 | + ) | ||
| 134 | + | ||
| 135 | + | ||
| 136 | + | ||
| 137 | + def impl_npu(op, *args, **kwargs): | ||
| 138 | + import torch_npu | ||
| 139 | + return torch_npu.npu.get_rng_state(), op(*args, **kwargs) | ||
| 140 | + | ||
| 141 | + | ||
| 142 | + backend_select_impl = run_and_save_rng_state.py_kernels.get( | ||
| 143 | + DispatchKey.BackendSelect, None | ||
| 144 | + ) | ||
| 145 | + | ||
| 146 | + | ||
| 147 | + def backend_select_with_npu(op, *args, **kwargs): | ||
| 148 | + from torch._prims.rng_prims import get_device | ||
| 149 | + | ||
| 150 | + device = get_device(args, kwargs) | ||
| 151 | + | ||
| 152 | + if device == "npu": | ||
| 153 | + return impl_npu(op, *args, **kwargs) | ||
| 154 | + | ||
| 155 | + return backend_select_impl(op, *args, **kwargs) | ||
| 156 | + | ||
| 157 | + run_and_save_rng_state.py_kernels[ | ||
| 158 | + DispatchKey.BackendSelect | ||
| 159 | + ] = backend_select_with_npu | ||
| 160 | + | ||
| 161 | + | ||
| 162 | +def patch_register_run_with_rng_state_op(): | ||
| 163 | + from torch._prims import rng_prims | ||
| 164 | + from torch._C import DispatchKey | ||
| 165 | + | ||
| 166 | + run_with_rng_state = getattr( | ||
| 167 | + rng_prims, "run_with_rng_state", None | ||
| 168 | + ) | ||
| 169 | + | ||
| 170 | + if getattr(run_with_rng_state, "_npu_patched", False): | ||
| 171 | + return | ||
| 172 | + | ||
| 173 | + | ||
| 174 | + | ||
| 175 | + def impl_npu(rng_state, op, *args, **kwargs): | ||
| 176 | + import torch_npu | ||
| 177 | + current_state = torch_npu.npu.get_rng_state() | ||
| 178 | + torch_npu.npu.set_rng_state(rng_state) | ||
| 179 | + try: | ||
| 180 | + out = op(*args, **kwargs) | ||
| 181 | + finally: | ||
| 182 | + torch_npu.npu.set_rng_state(current_state) | ||
| 183 | + return out | ||
| 184 | + | ||
| 185 | + | ||
| 186 | + backend_select_impl = run_with_rng_state.py_kernels.get( | ||
| 187 | + DispatchKey.BackendSelect, None | ||
| 188 | + ) | ||
| 189 | + | ||
| 190 | + | ||
| 191 | + def backend_select_with_npu(rng_state, op, *args, **kwargs): | ||
| 192 | + from torch._prims.rng_prims import get_device | ||
| 193 | + | ||
| 194 | + device = get_device(args, kwargs) | ||
| 195 | + | ||
| 196 | + if device == "npu": | ||
| 197 | + return impl_npu(rng_state, op, *args, **kwargs) | ||
| 198 | + | ||
| 199 | + return backend_select_impl(rng_state, op, *args, **kwargs) | ||
| 200 | + | ||
| 201 | + run_with_rng_state.py_kernels[ | ||
| 202 | + DispatchKey.BackendSelect | ||
| 203 | + ] = backend_select_with_npu | ||
| 204 | + | ||
| 205 | + | ||
| 206 | +patch_register_run_and_save_rng_state_op() | ||
| 207 | +patch_register_run_with_rng_state_op() | ||
| 208 | +patch_philox_rand_offset() | ||
| 209 | +patch_register_philox_rand() | ||