fix(npu): preserve native spatial ops for matting-pytorch-AtomGit

Mtest/_inductor/test_triton_experimental_regressions.py+79-1

	@@ -5,8 +5,11 @@ from unittest import mock
5		5
6	import sympy	6	import sympy
7	import torch	7	import torch
		8	+import torch.nn.functional as F
		9	+from torch._inductor import config
8	from torch._inductor.codegen.triton import IndexingOptions, TritonKernel	10	from torch._inductor.codegen.triton import IndexingOptions, TritonKernel
9	from torch._inductor.fx_passes.control_dependencies import control_deps	11	from torch._inductor.fx_passes.control_dependencies import control_deps
		12	+from torch._inductor.graph import GraphLowering
10	from torch.testing._internal.common_utils import (	13	from torch.testing._internal.common_utils import (
11	TestCase,	14	TestCase,
12	instantiate_parametrized_tests,	15	instantiate_parametrized_tests,
	@@ -22,6 +25,81 @@ from torch_npu._inductor.triton_experimental.codegen import triton as npu_triton
22		25
23	@instantiate_parametrized_tests	26	@instantiate_parametrized_tests
24	class TestTritonExperimentalRegressions(TestCase):	27	class TestTritonExperimentalRegressions(TestCase):
		28	+ def test_spatial_ops_use_explicit_fallback_list(self):
		29	+ expected = (
		30	+ torch.ops.aten.reflection_pad2d.default,
		31	+ torch.ops.aten.reflection_pad2d_backward.default,
		32	+ torch.ops.aten.upsample_bilinear2d.default,
		33	+ torch.ops.aten.upsample_bilinear2d_backward.default,
		34	+ )
		35	+ for op in expected:
		36	+ self.assertIn(op, lowering_override_list.EXPLICIT_FALLBACK_LIST)
		37	+
		38	+ def _check_native_forward_backward(self, fn, input, native_ops):
		39	+ eager_input = input.detach().clone().requires_grad_()
		40	+ compiled_input = input.detach().clone().requires_grad_()
		41	+
		42	+ eager_output = fn(eager_input)
		43	+ grad_output = torch.randn_like(eager_output)
		44	+ eager_output.backward(grad_output)
		45	+
		46	+ compiled = torch.compile(
		47	+ fn,
		48	+ backend="inductor",
		49	+ fullgraph=True,
		50	+ options={"npu_backend": "triton_experimental"},
		51	+ )
		52	+ source_codes = []
		53	+
		54	+ def save_output_code(code):
		55	+ source_codes.append(code)
		56	+
		57	+ try:
		58	+ with (
		59	+ config.patch("force_disable_caches", True),
		60	+ config.patch("implicit_fallbacks", False),
		61	+ mock.patch.object(GraphLowering, "save_output_code", save_output_code),
		62	+ ):
		63	+ torch._dynamo.reset()
		64	+ compiled_output = compiled(compiled_input)
		65	+ compiled_output.backward(grad_output)
		66	+ finally:
		67	+ torch._dynamo.reset()
		68	+
		69	+ self.assertEqual(eager_output, compiled_output)
		70	+ self.assertEqual(eager_input.grad, compiled_input.grad)
		71	+ code = "\n".join(source_codes)
		72	+ for native_op in native_ops:
		73	+ self.assertIn(f"torch.ops.aten.{native_op}.default(", code)
		74	+
		75	+ def test_reflection_pad2d_native_forward_backward(self):
		76	+ def fn(input):
		77	+ return F.pad(input, (1, 2, 1, 2), mode="reflect")
		78	+
		79	+ self._check_native_forward_backward(
		80	+ fn,
		81	+ torch.randn(2, 3, 5, 6, device="npu"),
		82	+ ("reflection_pad2d", "reflection_pad2d_backward"),
		83	+ )
		84	+
		85	+ def _check_bilinear_upsample(self, **kwargs):
		86	+ def fn(input):
		87	+ return F.interpolate(
		88	+ input, mode="bilinear", align_corners=False, **kwargs
		89	+ )
		90	+
		91	+ self._check_native_forward_backward(
		92	+ fn,
		93	+ torch.randn(2, 3, 5, 7, device="npu"),
		94	+ ("upsample_bilinear2d", "upsample_bilinear2d_backward"),
		95	+ )
		96	+
		97	+ def test_bilinear_upsample_size_native_forward_backward(self):
		98	+ self._check_bilinear_upsample(size=(9, 11))
		99	+
		100	+ def test_bilinear_upsample_scale_factor_native_forward_backward(self):
		101	+ self._check_bilinear_upsample(scale_factor=(1.5, 2.0))
		102	+
25	def test_control_deps_is_not_replaced_with_fallback(self):	103	def test_control_deps_is_not_replaced_with_fallback(self):
26	self.assertIn(control_deps, lowering_override_list.KEEP_UPSTREAM_LOWERING)	104	self.assertIn(control_deps, lowering_override_list.KEEP_UPSTREAM_LOWERING)
27		105
	@@ -36,7 +114,7 @@ class TestTritonExperimentalRegressions(TestCase):
36	):	114	):
37	experimental_lowering._register_npu_inductor_fallbacks()	115	experimental_lowering._register_npu_inductor_fallbacks()
38		116
39	- make_fallback.assert_not_called()	117	+ self.assertNotIn(mock.call(control_deps), make_fallback.call_args_list)
40		118
41	def test_constant_index_normalizes_emitted_and_cse_shapes(self):	119	def test_constant_index_normalizes_emitted_and_cse_shapes(self):
42	upstream_result = IndexingOptions(	120	upstream_result = IndexingOptions(

Mtorch_npu/_inductor/decomposition.py+26-0

	@@ -886,6 +886,27 @@ def _override_native_dropout_decomp():
886	extra_random_decomps[aten.native_dropout.default] = native_dropout	886	extra_random_decomps[aten.native_dropout.default] = native_dropout
887		887
888		888
		889	+def _register_upsample_bilinear2d_vec_dispatcher():
		890	+ """Route the vec overload to the native NPU-compatible default overload."""
		891	+ from torch._decomp.decompositions import upsample_compute_output_size
		892	+
		893	+ def upsample_bilinear2d_vec(
		894	+ input, output_size, align_corners, scale_factors
		895	+ ):
		896	+ osize = upsample_compute_output_size(
		897	+ input.size(), output_size, scale_factors
		898	+ )
		899	+ scales = scale_factors if scale_factors else (None, None)
		900	+ return aten.upsample_bilinear2d.default(
		901	+ input, osize, align_corners, scales[0], scales[1]
		902	+ )
		903	+
		904	+ for dispatch_key in (DispatchKey.Autograd, DispatchKey.CompositeImplicitAutograd):
		905	+ aten.upsample_bilinear2d.vec.py_impl(dispatch_key)(
		906	+ upsample_bilinear2d_vec
		907	+ )
		908	+
		909	+
889	def _register_triton_experimental_decompositions():	910	def _register_triton_experimental_decompositions():
890	"""Install all triton_experimental decomposition-table overrides and prune the	911	"""Install all triton_experimental decomposition-table overrides and prune the
891	exclusion list. Called directly by ``_load_triton_experimental_backend``.	912	exclusion list. Called directly by ``_load_triton_experimental_backend``.
	@@ -909,6 +930,10 @@ def _register_triton_experimental_decompositions():
909	aten.embedding,	930	aten.embedding,
910	aten.embedding_dense_backward,	931	aten.embedding_dense_backward,
911	aten.expm1,	932	aten.expm1,
		933	+ aten.reflection_pad2d,
		934	+ aten.reflection_pad2d_backward,
		935	+ aten.upsample_bilinear2d,
		936	+ aten.upsample_bilinear2d_backward,
912	]	937	]
913	# On A5 (910_95), let these ops go through decomposition instead of	938	# On A5 (910_95), let these ops go through decomposition instead of
914	# falling back, so drop them from the exclusion list.	939	# falling back, so drop them from the exclusion list.
	@@ -922,6 +947,7 @@ def _register_triton_experimental_decompositions():
922	npu_decomps_to_exclude.remove(op)	947	npu_decomps_to_exclude.remove(op)
923		948
924	disable_implicit_decomposition()	949	disable_implicit_decomposition()
		950	+ _register_upsample_bilinear2d_vec_dispatcher()
925	remove_decompositions(decompositions, npu_decomps_to_exclude)	951	remove_decompositions(decompositions, npu_decomps_to_exclude)
926		952
927	# Also remove from fast_random_decomps cache (used by select_decomp_table	953	# Also remove from fast_random_decomps cache (used by select_decomp_table

Mtorch_npu/_inductor/triton_experimental/lowering.py+9-1

	@@ -18,7 +18,11 @@ import logging
18	import sympy	18	import sympy
19	from . import config as ncfg	19	from . import config as ncfg
20	from . import device_props as _device_props	20	from . import device_props as _device_props
21	-from .lowering_override_list import GENERATE_LIST, KEEP_UPSTREAM_LOWERING	21	+from .lowering_override_list import (
		22	+ EXPLICIT_FALLBACK_LIST,
		23	+ GENERATE_LIST,
		24	+ KEEP_UPSTREAM_LOWERING,
		25	+)
22	import torch	26	import torch
23	from torch._inductor.lowering import (	27	from torch._inductor.lowering import (
24	lowerings,	28	lowerings,
	@@ -101,6 +105,10 @@ def _register_npu_inductor_fallbacks():
101	make_fallback(op)	105	make_fallback(op)
102	FALLBACK_LIST.append(op)	106	FALLBACK_LIST.append(op)
103		107
		108	+ for op in EXPLICIT_FALLBACK_LIST:
		109	+ make_fallback(op)
		110	+ FALLBACK_LIST.append(op)
		111	+
104	# Bernoulli's NPU op-api still consumes host-side seed/offset values via	112	# Bernoulli's NPU op-api still consumes host-side seed/offset values via
105	# NPUGeneratorImpl::philox_engine_inputs(). CANN does not currently expose	113	# NPUGeneratorImpl::philox_engine_inputs(). CANN does not currently expose
106	# a Bernoulli overload accepting the tensor seed/offset carried by	114	# a Bernoulli overload accepting the tensor seed/offset carried by

Mtorch_npu/_inductor/triton_experimental/lowering_override_list.py+14-2

	@@ -4,8 +4,10 @@
4		4
5	``GENERATE_LIST``: ops the backend keeps an inductor lowering for (a kernel is	5	``GENERATE_LIST``: ops the backend keeps an inductor lowering for (a kernel is
6	generated). Every other op that is not a decomposition is turned into a fallback by	6	generated). Every other op that is not a decomposition is turned into a fallback by
7	-``lowering._register_npu_inductor_fallbacks``. ``KEEP_UPSTREAM_LOWERING``: ops whose	7	+``lowering._register_npu_inductor_fallbacks``. ``EXPLICIT_FALLBACK_LIST``: ops that
8	-upstream lowering must be preserved verbatim rather than clobbered into a fallback.	8	+need a native fallback but have no existing lowering for that bulk registration to
		9	+discover. ``KEEP_UPSTREAM_LOWERING``: ops whose upstream lowering must be preserved
		10	+verbatim rather than clobbered into a fallback.
9	"""	11	"""
10	import torch	12	import torch
11	from torch._inductor.fx_passes.control_dependencies import control_deps	13	from torch._inductor.fx_passes.control_dependencies import control_deps
	@@ -109,6 +111,16 @@ GENERATE_LIST = [
109	if _device_props.is_a5():	111	if _device_props.is_a5():
110	GENERATE_LIST += []	112	GENERATE_LIST += []
111		113
		114	+# These decompositions are disabled for triton_experimental so the native NPU
		115	+# kernels are retained. They have no upstream lowering, therefore they must be
		116	+# registered explicitly for strict mode (implicit_fallbacks=False).
		117	+EXPLICIT_FALLBACK_LIST = [
		118	+ aten.reflection_pad2d.default,
		119	+ aten.reflection_pad2d_backward.default,
		120	+ aten.upsample_bilinear2d.default,
		121	+ aten.upsample_bilinear2d_backward.default,
		122	+]
		123	+
112	# Higher-order and runtime-assertion ops whose intentional upstream lowerings must	124	# Higher-order and runtime-assertion ops whose intentional upstream lowerings must
113	# not be replaced with generic fallbacks. control_deps carries a Subgraph argument	125	# not be replaced with generic fallbacks. control_deps carries a Subgraph argument
114	# that only its dedicated lowering understands; FallbackKernel treats it as a tensor	126	# that only its dedicated lowering understands; FallbackKernel treats it as a tensor