已合并
[fix]user_autotune_npu #32036
cuiduo创建于 3月19日
[fix]user_autotune_npu #32036
已合并
共 4 个文件变更+141-3
| @@ -0,0 +1,48 @@ | |||
| 1 | +import torch | ||
| 2 | +import triton | ||
| 3 | +import triton.language as tl | ||
| 4 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 5 | +import torch_npu | ||
| 6 | +import torch_npu._inductor | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +class TestUserAutotuneNpu(TestCase): | ||
| 10 | + def test_user_autotune_npu(self): | ||
| 11 | + | ||
| 12 | + configs=[ | ||
| 13 | + triton.Config({"BLOCK_SIZE": 64}), | ||
| 14 | + triton.Config({"BLOCK_SIZE": 32}), | ||
| 15 | + ], | ||
| 16 | + key=["n_elements"], | ||
| 17 | + ) | ||
| 18 | + | ||
| 19 | + def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: "tl.constexpr"): | ||
| 20 | + pid = tl.program_id(axis=0) | ||
| 21 | + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) | ||
| 22 | + mask = offsets < n_elements | ||
| 23 | + x = tl.load(x_ptr + offsets, mask=mask) | ||
| 24 | + y = tl.load(y_ptr + offsets, mask=mask) | ||
| 25 | + tl.store(output_ptr + offsets, x + y, mask=mask) | ||
| 26 | + | ||
| 27 | + def add(x, y): | ||
| 28 | + output = torch.empty_like(x) | ||
| 29 | + n_elements = output.numel() | ||
| 30 | + | ||
| 31 | + def grid(meta): | ||
| 32 | + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) | ||
| 33 | + | ||
| 34 | + add_kernel[grid](x, y, output, n_elements) | ||
| 35 | + return output | ||
| 36 | + | ||
| 37 | + x = torch.randn(64, device="npu") | ||
| 38 | + y = torch.randn(64, device="npu") | ||
| 39 | + expected = x + y | ||
| 40 | + | ||
| 41 | + compiled = torch.compile(add, backend="inductor") | ||
| 42 | + result = compiled(x, y) | ||
| 43 | + | ||
| 44 | + self.assertTrue(torch.allclose(result, expected, atol=1e-3)) | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +if __name__ == "__main__": | ||
| 48 | + run_tests() | ||
| @@ -467,7 +467,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 467 | def _get_grid_type(self) -> type[triton_heuristics.GridExpr]: | 467 | def _get_grid_type(self) -> type[triton_heuristics.GridExpr]: |
| 468 | return npu_triton_heuristics.GridNpu | 468 | return npu_triton_heuristics.GridNpu |
| 469 | 469 | ||
| 470 | - def gen_triton_ext_imports(self): | 470 | + @staticmethod |
| 471 | + def gen_triton_ext_imports(): | ||
| 471 | imports = IndentedBuffer() | 472 | imports = IndentedBuffer() |
| 472 | imports.splice( | 473 | imports.splice( |
| 473 | """ | 474 | """ |
| @@ -18,6 +18,8 @@ from torch._inductor.ir import GraphPartitionSignature | |||
| 18 | 18 | ||
| 19 | from torch_npu._inductor import config as npu_config | 19 | from torch_npu._inductor import config as npu_config |
| 20 | import torch_npu.npu.aclnn | 20 | import torch_npu.npu.aclnn |
| 21 | +from torch_npu._inductor.npu_triton_heuristics import PrecomputedGridNpu, user_autotune_npu | ||
| 22 | +from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel | ||
| 21 | 23 | ||
| 22 | 24 | ||
| 23 | class NPUWrapperCodeGen(PythonWrapperCodegen): | 25 | class NPUWrapperCodeGen(PythonWrapperCodegen): |
| @@ -270,4 +272,34 @@ class NPUWrapperCodeGen(PythonWrapperCodegen): | |||
| 270 | with self.wrapper_call.indent(): | 272 | with self.wrapper_call.indent(): |
| 271 | self.wrapper_call.writeline('exc_info=(None, None, None)') | 273 | self.wrapper_call.writeline('exc_info=(None, None, None)') |
| 272 | self.wrapper_call.writeline('static_kernel_compiler.__exit__(*exc_info)') | 274 | self.wrapper_call.writeline('static_kernel_compiler.__exit__(*exc_info)') |
| 273 | - super().generate_return(output_refs) | 275 | + super().generate_return(output_refs) |
| 276 | + | ||
| 277 | + def define_kernel( | ||
| 278 | + self, | ||
| 279 | + kernel_name: str, | ||
| 280 | + kernel_body: str, | ||
| 281 | + metadata: Optional[str] = None, | ||
| 282 | + gpu: bool = True, | ||
| 283 | + cpp_definition: Optional[str] = None, | ||
| 284 | + ): | ||
| 285 | + # 重写父类逻辑，将triton_heuristics.user_autotune替换为npu_triton_heuristics.user_autotune_npu， | ||
| 286 | + # 将PrecomputedGrid替换为PrecomputedGridNpu，以适配NPU设备，避免core dump错误。 | ||
| 287 | + if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body: | ||
| 288 | + kernel_body = kernel_body.replace( | ||
| 289 | + "triton_heuristics.user_autotune(", | ||
| 290 | + "npu_triton_heuristics.user_autotune_npu(" | ||
| 291 | + ) | ||
| 292 | + kernel_body = kernel_body.replace( | ||
| 293 | + "PrecomputedGrid", | ||
| 294 | + "PrecomputedGridNpu" | ||
| 295 | + ) | ||
| 296 | + kernel_body = kernel_body.replace( | ||
| 297 | + "FixedGrid", | ||
| 298 | + "FixedGridNpu" | ||
| 299 | + ) | ||
| 300 | + #import npu_triton_heuristicsd相关头文件 | ||
| 301 | + kernel_body = kernel_body.replace( | ||
| 302 | + "'''\n", | ||
| 303 | + "'''\n" + NPUIndexTritonKernel.gen_triton_ext_imports() + "\n" | ||
| 304 | + ) | ||
| 305 | + super().define_kernel(kernel_name, kernel_body, metadata, gpu, cpp_definition) | ||
| @@ -48,7 +48,8 @@ from torch._inductor.runtime.triton_heuristics import ( | |||
| 48 | NoTritonConfigsError, | 48 | NoTritonConfigsError, |
| 49 | TritonCompileResult, | 49 | TritonCompileResult, |
| 50 | GridExpr, | 50 | GridExpr, |
| 51 | - config_to_dict | 51 | + config_to_dict, |
| 52 | + config_from_dict | ||
| 52 | ) | 53 | ) |
| 53 | from torch._inductor.runtime.runtime_utils import triton_hash_to_path_key | 54 | from torch._inductor.runtime.runtime_utils import triton_hash_to_path_key |
| 54 | from triton.compiler import CompiledKernel | 55 | from triton.compiler import CompiledKernel |
| @@ -206,6 +207,39 @@ class GridExprNpu(GridExpr): | |||
| 206 | return grid | 207 | return grid |
| 207 | 208 | ||
| 208 | 209 | ||
| 210 | + | ||
| 211 | +class PrecomputedGridNpu(GridNpu): | ||
| 212 | + def __init__(self, *, inductor_meta, mode="python", **kwargs): | ||
| 213 | + super().__init__(inductor_meta=inductor_meta, mode=mode, numels=kwargs.get("numels")) | ||
| 214 | + | ||
| 215 | + def generate(self, meta: dict[str, int]) -> None: | ||
| 216 | + for candidate in self.inductor_meta["precomputed_grids"]: | ||
| 217 | + if all(meta.get(k) == v for k, v in candidate["config"].items()): | ||
| 218 | + self.x_grid, self.y_grid, self.z_grid = candidate[self.mode] | ||
| 219 | + return | ||
| 220 | + raise AssertionError( | ||
| 221 | + f"Precomputed grid not found for {meta} in {self.inductor_meta['precomputed_grids']}" | ||
| 222 | + ) | ||
| 223 | + | ||
| 224 | + | ||
| 225 | + | ||
| 226 | +class FixedGridNpu(GridNpu): | ||
| 227 | + def __init__(self, *, inductor_meta, mode="python", **kwargs): | ||
| 228 | + super().__init__(inductor_meta=inductor_meta, mode=mode, numels=kwargs.get("numels")) | ||
| 229 | + | ||
| 230 | + | ||
| 231 | + def setup_grid_as_args() -> dict[str, Any]: | ||
| 232 | + """Inductor meta so the launcher takes three extra grid arguments""" | ||
| 233 | + return { | ||
| 234 | + "grid_type": FixedGridNpu.__name__, | ||
| 235 | + "fixed_grid": ["_grid_0", "_grid_1", "_grid_2"], | ||
| 236 | + "extra_launcher_args": ["_grid_0", "_grid_1", "_grid_2"], | ||
| 237 | + } | ||
| 238 | + | ||
| 239 | + def generate(self, meta: dict[str, int]) -> None: | ||
| 240 | + self.x_grid, self.y_grid, self.z_grid = self.inductor_meta["fixed_grid"] | ||
| 241 | + | ||
| 242 | + | ||
| 209 | class TritonCompileResultNpu(TritonCompileResult): | 243 | class TritonCompileResultNpu(TritonCompileResult): |
| 210 | def make_launcher(self): | 244 | def make_launcher(self): |
| 211 | cfg = self.config | 245 | cfg = self.config |
| @@ -1338,3 +1372,26 @@ def benchmark_all_configs(self, *args, input_grid, **kwargs): | |||
| 1338 | k.shared, | 1372 | k.shared, |
| 1339 | ) | 1373 | ) |
| 1340 | return timings | 1374 | return timings |
| 1375 | + | ||
| 1376 | + | ||
| 1377 | +def user_autotune_npu( | ||
| 1378 | + configs, | ||
| 1379 | + triton_meta, | ||
| 1380 | + filename=None, | ||
| 1381 | + inductor_meta=None, | ||
| 1382 | + custom_kernel=False, | ||
| 1383 | +): | ||
| 1384 | + | ||
| 1385 | + if len(configs) == 0: | ||
| 1386 | + configs = [triton.Config({})] | ||
| 1387 | + else: | ||
| 1388 | + configs = [*map(config_from_dict, configs)] | ||
| 1389 | + return cached_autotune( | ||
| 1390 | + None, | ||
| 1391 | + configs, | ||
| 1392 | + triton_meta=triton_meta, | ||
| 1393 | + heuristic_type=HeuristicType.USER_AUTOTUNE, | ||
| 1394 | + filename=filename, | ||
| 1395 | + inductor_meta=inductor_meta, | ||
| 1396 | + custom_kernel=custom_kernel, | ||
| 1397 | + ) | ||