[fix]user_autotune_npu-pytorch-AtomGit

Atest/_inductor/test_user_autotune_npu.py+48-0

@@ -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	+ @triton.autotune(
	12	+ configs=[
	13	+ triton.Config({"BLOCK_SIZE": 64}),
	14	+ triton.Config({"BLOCK_SIZE": 32}),
	15	+ ],
	16	+ key=["n_elements"],
	17	+ )
	18	+ @triton.jit
	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()

Mtorch_npu/_inductor/codegen/wrapper.py+29-2

	@@ -20,9 +20,10 @@ from torch._inductor.virtualized import V
20	from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode	20	from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
21	from torch.utils._sympy.singleton_int import SingletonInt	21	from torch.utils._sympy.singleton_int import SingletonInt
22	from torch._inductor.ir import GraphPartitionSignature	22	from torch._inductor.ir import GraphPartitionSignature
23	-
24	from torch_npu._inductor import config as npu_config	23	from torch_npu._inductor import config as npu_config
25	import torch_npu.npu.aclnn	24	import torch_npu.npu.aclnn
		25	+from torch_npu._inductor.codegen.triton import gen_triton_ext_imports
		26	+from torch_npu._inductor.npu_triton_heuristics import PrecomputedGridNpu, user_autotune_npu
26		27
27		28
28	class NPUWrapperCodeGen(PythonWrapperCodegen):	29	class NPUWrapperCodeGen(PythonWrapperCodegen):
	@@ -264,4 +265,30 @@ class NPUWrapperCodeGen(PythonWrapperCodegen):
264	with self.wrapper_call.indent():	265	with self.wrapper_call.indent():
265	self.wrapper_call.writeline('exc_info=(None, None, None)')	266	self.wrapper_call.writeline('exc_info=(None, None, None)')
266	self.wrapper_call.writeline('static_kernel_complier.__exit__(*exc_info)')	267	self.wrapper_call.writeline('static_kernel_complier.__exit__(*exc_info)')
		AtlasAccount3月16日变量命名问题: 继续使用了拼写错误的变量名'static_kernel_complier'，与前面几行的问题相同。问题类型: 变量命名问题文件路径: `torch_npu/_inductor/codegen/wrapper.py` 行号: 295 问题代码: self.wrapper_call.writeline('static_kernel_complier.__exit__(exc_info)') 修改建议: 将'static_kernel_complier'修改为'static_kernel_compiler'：self.wrapper_call.writeline('static_kernel_compiler.__exit__(exc_info)') --- 此评论由代码审查工具自动生成
267	- super().generate_return(output_refs)	268	+ super().generate_return(output_refs)
		269	+
		270	+ def define_kernel(
		271	+ self,
		272	+ kernel_name: str,
		273	+ kernel_body: str,
		274	+ metadata: Optional[str] = None,
		275	+ gpu: bool = True,
		276	+ cpp_definition: Optional[str] = None,
		277	+ ):
		278	+ #重写父类逻辑
		279	+ #修改user_autotune为user_autotune_npu进行适配，必满core dump错误
		280	+ if "user_autotune" in kernel_body and "user_autotune_npu" not in kernel_body:
		281	+ kernel_body = kernel_body.replace(
		282	+ "triton_heuristics.user_autotune(",
		283	+ "npu_triton_heuristics.user_autotune_npu("
		284	+ )
		285	+ kernel_body = kernel_body.replace(
		286	+ "PrecomputedGrid",
		287	+ "PrecomputedGridNpu"
		288	+ )
		289	+ #import npu_triton_heuristicsd相关头文件
		290	+ kernel_body = kernel_body.replace(
		291	+ "'''\n",
		292	+ "'''\n" + gen_triton_ext_imports() + "\n"
		293	+ )
		294	+ super().define_kernel(kernel_name, kernel_body, metadata, gpu, cpp_definition)

Mtorch_npu/_inductor/npu_triton_heuristics.py+43-2

	@@ -49,7 +49,8 @@ from torch._inductor.runtime.triton_heuristics import (
49	NoTritonConfigsError,	49	NoTritonConfigsError,
50	TritonCompileResult,	50	TritonCompileResult,
51	GridExpr,	51	GridExpr,
52	- config_to_dict	52	+ config_to_dict,
		53	+ config_from_dict
53	)	54	)
54	from torch._inductor.runtime.runtime_utils import triton_hash_to_path_key	55	from torch._inductor.runtime.runtime_utils import triton_hash_to_path_key
55	from triton.compiler import CompiledKernel	56	from triton.compiler import CompiledKernel
	@@ -227,6 +228,21 @@ class FixedGridNpu(GridExpr):
227	}	228	}
228		229
229		230
		231	+@dataclasses.dataclass
		232	+class PrecomputedGridNpu(GridNpu):
		233	+ def __init__(self, , inductor_meta, mode="python", *kwargs):
		234	+ super().__init__(inductor_meta=inductor_meta, mode=mode, numels=kwargs.get("numels"))
		235	+
		236	+ def generate(self, meta: dict[str, int]) -> None:
		237	+ for candidate in self.inductor_meta["precomputed_grids"]:
		238	+ if all(meta.get(k) == v for k, v in candidate["config"].items()):
		239	+ self.x_grid, self.y_grid, self.z_grid = candidate[self.mode]
		240	+ return
		241	+ raise AssertionError(
		242	+ f"Precomputed grid not found for {meta} in {self.inductor_meta['precomputed_grids']}"
		243	+ )
		244	+
		245	+
230	def is_namedtuple_isinstance(obj):	246	def is_namedtuple_isinstance(obj):
231	return (	247	return (
232	isinstance(obj, tuple) and	248	isinstance(obj, tuple) and
	@@ -792,6 +808,7 @@ class NPUCachingAutotuner(CachingAutotuner):
792		808
793	def kernel_call():	809	def kernel_call():
794	cloned_args, cloned_kwargs = self.clone_args(args, *kwargs)	810	cloned_args, cloned_kwargs = self.clone_args(args, *kwargs)
		811	+ self.reset_to_zero_args(args, *kwargs)
795	launcher(	812	launcher(
796	*cloned_args,	813	*cloned_args,
797	**cloned_kwargs,	814	**cloned_kwargs,
	@@ -1485,6 +1502,7 @@ def _benchmark_all_configs(self, args, *kwargs):
1485	{dict(zip(self.arg_names, args)), launcher.config.kwargs}	1502	{dict(zip(self.arg_names, args)), launcher.config.kwargs}
1486	)	1503	)
1487	cloned_args, cloned_kwargs = self.clone_args(args, *kwargs)	1504	cloned_args, cloned_kwargs = self.clone_args(args, *kwargs)
		1505	+ self.reset_to_zero_args(args, *kwargs)
1488	launcher(	1506	launcher(
1489	*cloned_args,	1507	*cloned_args,
1490	**cloned_kwargs,	1508	**cloned_kwargs,
	@@ -1623,4 +1641,27 @@ def precompile_parallel(
1623		1641
1624	self._precompile_worker_parallel()	1642	self._precompile_worker_parallel()
1625	self._make_launchers()	1643	self._make_launchers()
1626	- log.info(f"kernel: {self.get_fn_name()} precompile elapsed time: {time.perf_counter() - start_time}s")	1644	+ log.info(f"kernel: {self.get_fn_name()} precompile elapsed time: {time.perf_counter() - start_time}s")
		1645	+
		1646	+
		1647	+def user_autotune_npu(
		1648	+ configs,
		1649	+ triton_meta,
		1650	+ filename=None,
		1651	+ inductor_meta=None,
		1652	+ custom_kernel=False,
		1653	+):
		1654	+
		1655	+ if len(configs) == 0:
		1656	+ configs = [triton.Config({})]
		1657	+ else:
		1658	+ configs = [*map(config_from_dict, configs)]
		1659	+ return cached_autotune(
		1660	+ None,
		1661	+ configs,
		1662	+ triton_meta=triton_meta,
		1663	+ heuristic_type=HeuristicType.USER_AUTOTUNE,
		1664	+ filename=filename,
		1665	+ inductor_meta=inductor_meta,
		1666	+ custom_kernel=custom_kernel,
		1667	+ )