已合并
wrap_triton #26742
cuiduo创建于 2025年11月19日
wrap_triton #26742
已合并
共 6 个文件变更+62-4
| @@ -0,0 +1,53 @@ | |||
| 1 | +import torch | ||
| 2 | +import triton | ||
| 3 | +from torch.library import triton_op, wrap_triton | ||
| 4 | +from triton import language as tl | ||
| 5 | +from torch._inductor.utils import run_and_get_code | ||
| 6 | +from torch.testing._internal.common_utils import ( | ||
| 7 | + run_tests, | ||
| 8 | + parametrize, | ||
| 9 | + instantiate_parametrized_tests, | ||
| 10 | +) | ||
| 11 | +from testutils import TestUtils | ||
| 12 | +import torch_npu | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +class TestWrapTriton(TestUtils): | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + def sin_kernel(in_ptr0, out_ptr, n_elements, BLOCK_SIZE: "tl.constexpr"): | ||
| 20 | + pid = tl.program_id(axis=0) | ||
| 21 | + block_start = pid * BLOCK_SIZE | ||
| 22 | + offsets = block_start + tl.arange(0, BLOCK_SIZE) | ||
| 23 | + mask = offsets < n_elements | ||
| 24 | + x = tl.load(in_ptr0 + offsets, mask=mask) | ||
| 25 | + output = tl.sin(x) | ||
| 26 | + tl.store(out_ptr + offsets, output, mask=mask) | ||
| 27 | + | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + def mysin(x: torch.Tensor) -> torch.Tensor: | ||
| 31 | + out = torch.empty_like(x) | ||
| 32 | + n_elements = x.numel() | ||
| 33 | + wrap_triton(TestWrapTriton.sin_kernel)[(n_elements,)]( | ||
| 34 | + x, out, n_elements, BLOCK_SIZE=4 | ||
| 35 | + ) | ||
| 36 | + return out | ||
| 37 | + | ||
| 38 | + def op_calc(self, x): | ||
| 39 | + return self.mysin(x) + x | ||
| 40 | + | ||
| 41 | + | ||
| 42 | + def test_wrap_triton(self, shape): | ||
| 43 | + x = torch.randn(shape, requires_grad=False, dtype=torch.float32, device="npu") | ||
| 44 | + std_out = self.op_calc(x) | ||
| 45 | + compile_out, codes = run_and_get_code(torch.compile(self.op_calc), x) | ||
| 46 | + self.assertEqual(std_out, compile_out, atol=1e-5, rtol=1e-5) | ||
| 47 | + self.assertTrue('sin_kernel_0.run' in codes[0]) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +instantiate_parametrized_tests(TestWrapTriton) | ||
| 51 | + | ||
| 52 | +if __name__ == "__main__": | ||
| 53 | + run_tests() | ||
| @@ -1,4 +1,5 @@ | |||
| 1 | import os | 1 | import os |
| 2 | +from unittest import skip | ||
| 2 | import random | 3 | import random |
| 3 | import numpy as np | 4 | import numpy as np |
| 4 | 5 | ||
| @@ -31,7 +32,8 @@ class TestModel(TestUtils): | |||
| 31 | calc = forward_calc(primals_4) | 32 | calc = forward_calc(primals_4) |
| 32 | ref = forward(primals_4) | 33 | ref = forward(primals_4) |
| 33 | self.assertEqual(ref, calc, atol=1e-4, rtol=1e-4, equal_nan=True) | 34 | self.assertEqual(ref, calc, atol=1e-4, rtol=1e-4, equal_nan=True) |
| 34 | - | 35 | + |
| 36 | + | ||
| 35 | def test_opensora_cases_model_11_inference(self): | 37 | def test_opensora_cases_model_11_inference(self): |
| 36 | def forward(arg0_1: "f32[1, 1, 9600]", arg1_1: "f32[1, 1, 512]"): | 38 | def forward(arg0_1: "f32[1, 1, 9600]", arg1_1: "f32[1, 1, 512]"): |
| 37 | unsqueeze: "f32[1, 1, 1, 9600]" = torch.ops.aten.unsqueeze.default(arg0_1, 1) | 39 | unsqueeze: "f32[1, 1, 1, 9600]" = torch.ops.aten.unsqueeze.default(arg0_1, 1) |
| @@ -278,4 +278,4 @@ def _register_npu_inductor_fallbacks(): | |||
| 278 | 278 | ||
| 279 | make_fallback(aten._log_softmax) | 279 | make_fallback(aten._log_softmax) |
| 280 | make_fallback(aten.gather) | 280 | make_fallback(aten.gather) |
| 281 | - make_fallback(aten.nll_loss_forward) | 281 | + make_fallback(aten.nll_loss_forward) |
| @@ -2299,4 +2299,4 @@ def _register_npu_inductor_fallbacks(): | |||
| 2299 | 2299 | ||
| 2300 | lowering.make_fallback(aten._log_softmax) | 2300 | lowering.make_fallback(aten._log_softmax) |
| 2301 | lowering.make_fallback(aten.gather) | 2301 | lowering.make_fallback(aten.gather) |
| 2302 | - lowering.make_fallback(aten.nll_loss_forward) | 2302 | + lowering.make_fallback(aten.nll_loss_forward) |
| @@ -1,4 +1,5 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | +from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation | ||
| 2 | from torch_npu import npu_dtype_cast, _npu_dtype_cast | 3 | from torch_npu import npu_dtype_cast, _npu_dtype_cast |
| 3 | 4 | ||
| 4 | aten = torch.ops.aten | 5 | aten = torch.ops.aten |
| @@ -75,6 +76,7 @@ GENERATE_LIST = [ | |||
| 75 | aten.squeeze, | 76 | aten.squeeze, |
| 76 | aten.copy, | 77 | aten.copy, |
| 77 | aten.reciprocal, | 78 | aten.reciprocal, |
| 79 | + triton_kernel_wrapper_mutation, | ||
| 78 | ] | 80 | ] |
| 79 | 81 | ||
| 80 | GENERATE_LIST2 = [ | 82 | GENERATE_LIST2 = [ |