已关闭
test(dynamo): cover assume_constant_result on NPU #42034
2501_93637465创建于 7月18日关闭于 8月6日
test(dynamo): cover assume_constant_result on NPU #42034
已关闭
共 1 个文件变更+64-0
| @@ -0,0 +1,64 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +""" | ||
| 17 | +Add validation cases for torch.compiler.assume_constant_result on NPU: | ||
| 18 | +1. PyTorch community functionally validates torch._dynamo.assume_constant_result | ||
| 19 | + and only checks the public API signature, without direct NPU coverage. | ||
| 20 | +2. This file directly validates torch.compiler.assume_constant_result with NPU | ||
| 21 | + tensors and checks eager/compiled result consistency. | ||
| 22 | +""" | ||
| 23 | + | ||
| 24 | +import torch | ||
| 25 | +import torch_npu | ||
| 26 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 27 | + | ||
| 28 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 29 | + | ||
| 30 | + | ||
| 31 | +class TestAssumeConstantResult(TestCase): | ||
| 32 | + def test_assume_constant_result(self): | ||
| 33 | + torch._dynamo.reset() | ||
| 34 | + call_count = 0 | ||
| 35 | + | ||
| 36 | + def constant_scale(): | ||
| 37 | + nonlocal call_count | ||
| 38 | + call_count += 1 | ||
| 39 | + return 2.0 | ||
| 40 | + | ||
| 41 | + marked_constant_scale = torch.compiler.assume_constant_result( | ||
| 42 | + constant_scale | ||
| 43 | + ) | ||
| 44 | + self.assertIs(marked_constant_scale, constant_scale) | ||
| 45 | + | ||
| 46 | + def fn(x): | ||
| 47 | + return x * marked_constant_scale() | ||
| 48 | + | ||
| 49 | + compiled_fn = torch.compile(fn, backend="eager", fullgraph=True) | ||
| 50 | + x = torch.arange(4, dtype=torch.float32).to(device_type) | ||
| 51 | + | ||
| 52 | + actual = compiled_fn(x) | ||
| 53 | + next_actual = compiled_fn(x + 1) | ||
| 54 | + torch_npu.npu.synchronize() | ||
| 55 | + | ||
| 56 | + self.assertTrue(torch.equal(actual.cpu(), (x * 2.0).cpu())) | ||
| 57 | + self.assertTrue( | ||
| 58 | + torch.equal(next_actual.cpu(), ((x + 1) * 2.0).cpu()) | ||
| 59 | + ) | ||
| 60 | + self.assertEqual(call_count, 1) | ||
| 61 | + | ||
| 62 | + | ||
| 63 | +if __name__ == "__main__": | ||
| 64 | + run_tests() | ||