已合并
【fix】 fix inductor upcast_to_fp32 codegen #33344
kkjocker创建于 4月8日
【fix】 fix inductor upcast_to_fp32 codegen #33344
已合并
共 2 个文件变更+42-0
| @@ -0,0 +1,35 @@ | |||
| 1 | +import unittest | ||
| 2 | +import torch | ||
| 3 | + | ||
| 4 | +from testutils import TestUtils | ||
| 5 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 6 | +from torch._inductor import config | ||
| 7 | +from torch._inductor.utils import run_and_get_code | ||
| 8 | + | ||
| 9 | +import torch_npu | ||
| 10 | +import torch_npu._inductor | ||
| 11 | + | ||
| 12 | +DEVICE = "npu" | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +class TestCodegenUpcastToFP32(TestUtils): | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + def test_codegen_upcast_to_fp32_emits_cast(self, dtype, upcast_flag): | ||
| 19 | + | ||
| 20 | + def func(x): | ||
| 21 | + return torch.abs(x) | ||
| 22 | + | ||
| 23 | + x = torch.randn((1024, 1024), device=DEVICE, dtype=dtype) | ||
| 24 | + | ||
| 25 | + with config.patch("triton.codegen_upcast_to_fp32", upcast_flag): | ||
| 26 | + opt_func = torch._dynamo.optimize("inductor")(func) | ||
| 27 | + out, code = run_and_get_code(opt_func, x) | ||
| 28 | + | ||
| 29 | + self.assertTrue(".to(tl.float32)" in code[0]) | ||
| 30 | + self.assertEqual(func(x), opt_func(x)) | ||
| 31 | + | ||
| 32 | +instantiate_parametrized_tests(TestCodegenUpcastToFP32) | ||
| 33 | + | ||
| 34 | +if __name__ == "__main__": | ||
| 35 | + run_tests() | ||
| @@ -1379,6 +1379,13 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1379 | index_str = indexing.index_str | 1379 | index_str = indexing.index_str |
| 1380 | mask_str = indexing.mask_str | 1380 | mask_str = indexing.mask_str |
| 1381 | line = f"tl.load({var} + ({index_str}), {mask_str}{ep}{other})" | 1381 | line = f"tl.load({var} + ({index_str}), {mask_str}{ep}{other})" |
| 1382 | + | ||
| 1383 | + if ( | ||
| 1384 | + dtype in (torch.float16, torch.bfloat16) | ||
| 1385 | + and config.triton.codegen_upcast_to_fp32 | ||
| 1386 | + ): | ||
| 1387 | + line += ".to(tl.float32)" | ||
| 1388 | + dtype = torch.float32 | ||
| 1382 | 1389 | ||
| 1383 | dtype = V.graph.get_dtype(name) | 1390 | dtype = V.graph.get_dtype(name) |
| 1384 | if dtype in (torch.bfloat16,): | 1391 | if dtype in (torch.bfloat16,): |