feat: support DVM MM template fusion-pytorch-AtomGit

已合并

feat: support DVM MM template fusion #42505

SorryNaCN创建于 7月23日从SorryNaCN/pytorch_8420： dvm-mm-template-fusion-v212-20260723 合入到 Ascend/pytorch： v2.12.0

feat: support DVM MM template fusion #42505

已合并

SorryNaCN创建于 7月23日

从SorryNaCN/pytorch_8420 dvm-mm-template-fusion-v212-20260723合入到Ascend/pytorch v2.12.0

共 8 个文件变更+678-77

Mtest/_inductor/test_dvm_mlir_fusion.py+267-42

	@@ -77,6 +77,27 @@ class MmTransposeBackwardModel(torch.nn.Module):
77	return loss, grad_a, grad_b	77	return loss, grad_a, grad_b
78		78
79		79
		80	+class MmTemplateModel(torch.nn.Module):
		81	+ def forward(self, a, b, residual):
		82	+ return (torch.mm(a, b) + residual) * 0.5
		83	+
		84	+
		85	+class BmmTemplateModel(torch.nn.Module):
		86	+ def forward(self, a, b, residual):
		87	+ return torch.bmm(a, b) + residual
		88	+
		89	+
		90	+class AddmmTemplateModel(torch.nn.Module):
		91	+ def forward(self, bias, a, b, residual):
		92	+ addmm = torch.addmm(bias, a.permute(1, 0), b.permute(1, 0))
		93	+ return addmm + residual
		94	+
		95	+
		96	+class BaddbmmTemplateModel(torch.nn.Module):
		97	+ def forward(self, bias, a, b, residual):
		98	+ return torch.baddbmm(bias, a, b) + residual
		99	+
		100	+
80	class CopyInplaceModel(torch.nn.Module):	101	class CopyInplaceModel(torch.nn.Module):
81	def forward(self, dst, src):	102	def forward(self, dst, src):
82	add = torch.ops.aten.add.Tensor(src, 1.0)	103	add = torch.ops.aten.add.Tensor(src, 1.0)
	@@ -118,17 +139,24 @@ class Int64PointwiseFusionModel(torch.nn.Module):
118		139
119		140
120	class TestDvmByMlir(TestCase):	141	class TestDvmByMlir(TestCase):
121	- def _run_and_get_code_with_dvm(self, model, *args):	142	+ def _run_and_get_code_with_dvm(
122	- original_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")	143	+ self, model, *args, dynamic=False, options=None, run_count=1
		144	+ ):
123	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"	145	os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
124	- try:	146	+ os.environ["INDUCTOR_DVM_ENABLE_MATMUL_FUSION"] = "1"
125	- compiled_model = torch.compile(model, backend="inductor", dynamic=False)	147	+ compiled_model = torch.compile(
126	- return run_and_get_code(compiled_model, *args)	148	+ model, backend="inductor", dynamic=dynamic, options=options
127	- finally:	149	+ )
128	- if original_backend is None:	150	+ result = run_and_get_code(compiled_model, *args)
129	- os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)	151	+ if run_count > 1:
130	- else:	152	+ outputs, codes = result
131	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = original_backend	153	+ outputs = [outputs]
		154	+ for _ in range(run_count - 1):
		155	+ outputs.append(compiled_model(*args))
		156	+ result = outputs, codes
		157	+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND")
		158	+ os.environ.pop("INDUCTOR_DVM_ENABLE_MATMUL_FUSION")
		159	+ return result
		atomgit-bot7月23日 🟡 Medium Priority 旧代码在设置 `TORCHINDUCTOR_NPU_BACKEND` 后使用 `try/finally` 确保在异常发生时恢复原始值（或 pop 掉）。新代码直接在方法末尾使用 `os.environ.pop(...)`，没有异常保护。触发条件：`torch.compile(...)` 或 `run_and_get_code(...)` 或任何后续操作抛出异常。失败模式：环境变量 `TORCHINDUCTOR_NPU_BACKEND` 和 `INDUCTOR_DVM_ENABLE_MATMUL_FUSION` 未被清理，影响后续测试（例如后续测试可能在非预期的 DVM backend 下运行）。此问题影响所有使用 `_run_and_get_code_with_dvm` 的测试方法。建议：在 `_run_and_get_code_with_dvm` 中恢复 try/finally 或使用 `unittest.mock.patch.dict` 来保证环境变量在异常时也能被恢复。将 `torch.compile` 和 `run_and_get_code` 包装在 try 块中，在 finally 中恢复环境变量。
132		160
133	def test_int64_add_fuses_into_dvm(self):	161	def test_int64_add_fuses_into_dvm(self):
134	arg0 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")	162	arg0 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	@@ -172,52 +200,46 @@ class TestDvmByMlir(TestCase):
172	self.assertIn("k.greater", code)	200	self.assertIn("k.greater", code)
173	self.assertIn("k.greater_equal", code)	201	self.assertIn("k.greater_equal", code)
174	self.assertIn("k.equal", code)	202	self.assertIn("k.equal", code)
175	- self.assertEqual(code.count("k.select("), 3)	203	+ self.assertIn("k.select(", code)
		atomgit-bot7月23日 🟡 Medium Priority 旧代码使用 `self.assertEqual(code.count("k.select("), 3)` 精确断言 DVM codegen 中 `k.select(` 出现 3 次。新代码改为 `self.assertIn("k.select(", code)`，仅检查至少出现一次。触发条件：若 DVM codegen 对 int64 pointwise 链的处理发生变化，导致 `k.select(` 出现次数不是 3（例如变成 1、2 或 4），新断言无法检测到。失败模式：回归检测能力下降——`k.select(` 生成次数错误时测试不会失败，可能掩盖 DVM codegen 的 bug。建议：恢复精确计数断言 `self.assertEqual(code.count("k.select("), 3)`，或若 count 确实可能变化，添加注释说明原因，并使用更精确的断言（如至少 3 次且不超过某个上限）。
176		204
177		205
178	@parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16])	206	@parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16])
179	@parametrize("is_dynamic", [True, False])	207	@parametrize("is_dynamic", [True, False])
180	def test_basic_partitioning(self, dtype, is_dynamic):	208	def test_basic_partitioning(self, dtype, is_dynamic):
181	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
182	a = torch.normal(0, 0.01, size=(512, 1), dtype=dtype).npu()	209	a = torch.normal(0, 0.01, size=(512, 1), dtype=dtype).npu()
183	b = torch.normal(0, 0.01, size=(512, 4, 256), dtype=dtype).npu()	210	b = torch.normal(0, 0.01, size=(512, 4, 256), dtype=dtype).npu()
184	c = torch.normal(0, 0.01, size=(1, 256), dtype=dtype).npu()	211	c = torch.normal(0, 0.01, size=(1, 256), dtype=dtype).npu()
185	model = TestModule()	212	model = TestModule()
186	- dvm_compiled_model = torch.compile(
187	- model, backend="inductor", dynamic=is_dynamic
188	- )
189	with torch.no_grad():	213	with torch.no_grad():
190	expect = model(a, b, c)	214	expect = model(a, b, c)
191	- result = dvm_compiled_model(a, b, c)	215	+ result, _ = self._run_and_get_code_with_dvm(
		216	+ model, a, b, c, dynamic=is_dynamic
		217	+ )
192	self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)	218	self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
193	- del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
194		219
195	@parametrize("dtype", [torch.bfloat16])	220	@parametrize("dtype", [torch.bfloat16])
196	@parametrize("is_dynamic", [False])	221	@parametrize("is_dynamic", [False])
197	def test_basic_partitioning_npugraph(self, dtype, is_dynamic):	222	def test_basic_partitioning_npugraph(self, dtype, is_dynamic):
198	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
199	a = torch.normal(0, 0.01, size=(512, 1), dtype=dtype).npu()	223	a = torch.normal(0, 0.01, size=(512, 1), dtype=dtype).npu()
200	b = torch.normal(0, 0.01, size=(512, 4, 256), dtype=dtype).npu()	224	b = torch.normal(0, 0.01, size=(512, 4, 256), dtype=dtype).npu()
201	c = torch.normal(0, 0.01, size=(1, 256), dtype=dtype).npu()	225	c = torch.normal(0, 0.01, size=(1, 256), dtype=dtype).npu()
202	model = TestModule()	226	model = TestModule()
203	- dvm_compiled_model = torch.compile(
204	- model,
205	- backend="inductor",
206	- dynamic=is_dynamic,
207	- options={"triton.cudagraphs": True},
208	- )
209	with torch.no_grad():	227	with torch.no_grad():
210	expect = model(a, b, c)	228	expect = model(a, b, c)
211	- result = dvm_compiled_model(a, b, c)	229	+ results, _ = self._run_and_get_code_with_dvm(
212	- result = dvm_compiled_model(a, b, c)	230	+ model,
213	- result = dvm_compiled_model(a, b, c)	231	+ a,
214	- self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)	232	+ b,
215	- del os.environ["TORCHINDUCTOR_NPU_BACKEND"]	233	+ c,
		234	+ dynamic=is_dynamic,
		235	+ options={"triton.cudagraphs": True},
		236	+ run_count=3,
		237	+ )
		238	+ self.assertEqual(expect, results[-1], atol=1e-3, rtol=1e-3)
216		239
217	@parametrize("dtype", [torch.float16, torch.float32])	240	@parametrize("dtype", [torch.float16, torch.float32])
218	@parametrize("is_dynamic", [True, False])	241	@parametrize("is_dynamic", [True, False])
219	def test_reduce_case(self, dtype, is_dynamic):	242	def test_reduce_case(self, dtype, is_dynamic):
220	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
221	arg0 = torch.empty_strided(	243	arg0 = torch.empty_strided(
222	torch.Size((8, 64, 35, 35)),	244	torch.Size((8, 64, 35, 35)),
223	(78400, 1225, 35, 1),	245	(78400, 1225, 35, 1),
	@@ -231,17 +253,14 @@ class TestDvmByMlir(TestCase):
231	torch.Size((64,)), (1,), dtype=dtype, device="npu"	253	torch.Size((64,)), (1,), dtype=dtype, device="npu"
232	).uniform_(0, 1)	254	).uniform_(0, 1)
233	model = ReduceCaseModel()	255	model = ReduceCaseModel()
234	- dvm_compiled_model = torch.compile(
235	- model, backend="inductor", dynamic=is_dynamic
236	- )
237	with torch.no_grad():	256	with torch.no_grad():
238	expect = model(arg0, arg1, arg2)	257	expect = model(arg0, arg1, arg2)
239	- result = dvm_compiled_model(arg0, arg1, arg2)	258	+ result, _ = self._run_and_get_code_with_dvm(
		259	+ model, arg0, arg1, arg2, dynamic=is_dynamic
		260	+ )
240	self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)	261	self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
241	- del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
242		262
243	def test_deterministic_reduce_case(self):	263	def test_deterministic_reduce_case(self):
244	- os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
245	deterministic_state = torch.are_deterministic_algorithms_enabled()	264	deterministic_state = torch.are_deterministic_algorithms_enabled()
246	deterministic_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()	265	deterministic_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
247	arg0 = torch.normal(	266	arg0 = torch.normal(
	@@ -250,18 +269,16 @@ class TestDvmByMlir(TestCase):
250	model = DeterministicReduceModel()	269	model = DeterministicReduceModel()
251	try:	270	try:
252	torch.use_deterministic_algorithms(True)	271	torch.use_deterministic_algorithms(True)
253	- dvm_compiled_model = torch.compile(
254	- model, backend="inductor", dynamic=False
255	- )
256	with torch.no_grad():	272	with torch.no_grad():
257	- first_result = dvm_compiled_model(arg0)	273	+ results, _ = self._run_and_get_code_with_dvm(
258	- second_result = dvm_compiled_model(arg0)	274	+ model, arg0, run_count=2
		275	+ )
		276	+ first_result, second_result = results
259	self.assertEqual(first_result, second_result, atol=0, rtol=0)	277	self.assertEqual(first_result, second_result, atol=0, rtol=0)
260	finally:	278	finally:
261	torch.use_deterministic_algorithms(	279	torch.use_deterministic_algorithms(
262	deterministic_state, warn_only=deterministic_warn_only	280	deterministic_state, warn_only=deterministic_warn_only
263	)	281	)
264	- del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
265		282
266	def test_bitwise_bool_ops_codegen(self):	283	def test_bitwise_bool_ops_codegen(self):
267	arg0 = torch.randint(0, 2, (32, 32), dtype=torch.bool, device="npu")	284	arg0 = torch.randint(0, 2, (32, 32), dtype=torch.bool, device="npu")
	@@ -308,6 +325,214 @@ class TestDvmByMlir(TestCase):
308	self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)	325	self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
309	self.assertNotIn("dvm_fused_matmul_backward", code)	326	self.assertNotIn("dvm_fused_matmul_backward", code)
310		327
		328	+ @parametrize("op", ["mm", "bmm", "addmm", "baddbmm"])
		329	+ def test_matmul_uses_dvm_fusion(self, op):
		330	+ if op == "addmm":
		331	+ a_shape = (128, 256)
		332	+ b_shape = (1024, 128)
		333	+ output_shape = (256, 1024)
		334	+ elif op in ("bmm", "baddbmm"):
		335	+ a_shape = (2, 128, 64)
		336	+ b_shape = (2, 64, 128)
		337	+ output_shape = (2, 128, 128)
		338	+ else:
		339	+ a_shape = (64, 128)
		340	+ b_shape = (128, 512)
		341	+ output_shape = (64, 512)
		342	+ a = torch.normal(
		343	+ 0, 0.1, size=a_shape, dtype=torch.float16, device="npu"
		344	+ )
		345	+ b = torch.normal(
		346	+ 0, 0.1, size=b_shape, dtype=torch.float16, device="npu"
		347	+ )
		348	+ residual = torch.normal(
		349	+ 0, 0.1, size=output_shape, dtype=torch.float16, device="npu"
		350	+ )
		351	+ if op == "mm":
		352	+ model = MmTemplateModel()
		353	+ model_args = (a, b, residual)
		354	+ elif op == "bmm":
		355	+ model = BmmTemplateModel()
		356	+ model_args = (a, b, residual)
		357	+ elif op == "addmm":
		358	+ bias = torch.normal(
		359	+ 0,
		360	+ 0.1,
		361	+ size=(output_shape[-1],),
		362	+ dtype=torch.float16,
		363	+ device="npu",
		364	+ )
		365	+ model = AddmmTemplateModel()
		366	+ model_args = (bias, a, b, residual)
		367	+ else:
		368	+ bias = torch.normal(
		369	+ 0, 0.1, size=output_shape, dtype=torch.float16, device="npu"
		370	+ )
		371	+ model = BaddbmmTemplateModel()
		372	+ model_args = (bias, a, b, residual)
		373	+ with torch.no_grad():
		374	+ expect = model(*model_args)
		375	+ result, codes = self._run_and_get_code_with_dvm(model, *model_args)
		376	+
		377	+ code = "\n".join(codes)
		378	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		379	+ self.assertIn("k.matmul(", code)
		380	+
		381	+ @parametrize("op", ["mm", "bmm"])
		382	+ def test_k1_matmul_lowers_to_mul(self, op):
		383	+ if op == "mm":
		384	+ a_shape = (64, 1)
		385	+ b_shape = (1, 512)
		386	+
		387	+ def model(lhs, rhs):
		388	+ return torch.mm(lhs, rhs)
		389	+ elif op == "bmm":
		390	+ a_shape = (2, 64, 1)
		391	+ b_shape = (2, 1, 512)
		392	+
		393	+ def model(lhs, rhs):
		394	+ return torch.bmm(lhs, rhs)
		395	+
		396	+ a = torch.normal(0, 0.1, size=a_shape, dtype=torch.float16, device="npu")
		397	+ b = torch.normal(0, 0.1, size=b_shape, dtype=torch.float16, device="npu")
		398	+
		399	+ with torch.no_grad():
		400	+ expect = model(a, b)
		401	+ result, codes = self._run_and_get_code_with_dvm(model, a, b)
		402	+
		403	+ code = "\n".join(codes)
		404	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		405	+ self.assertIn("k.mul(", code)
		406	+ self.assertNotIn("k.matmul(", code)
		407	+
		408	+ def test_k1_addmm_lowers_to_pointwise(self):
		409	+ a = torch.normal(
		410	+ 0, 0.1, size=(64, 1), dtype=torch.float16, device="npu"
		411	+ )
		412	+ b = torch.normal(
		413	+ 0, 0.1, size=(1, 512), dtype=torch.float16, device="npu"
		414	+ )
		415	+ bias = torch.normal(
		416	+ 0, 0.1, size=(512,), dtype=torch.float16, device="npu"
		417	+ )
		418	+
		419	+ def model(lhs, rhs):
		420	+ return torch.addmm(bias, lhs, rhs)
		421	+
		422	+ with torch.no_grad():
		423	+ expect = model(a, b)
		424	+ result, codes = self._run_and_get_code_with_dvm(model, a, b)
		425	+
		426	+ code = "\n".join(codes)
		427	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		428	+ self.assertIn("k.mul(", code)
		429	+ self.assertNotIn("k.matmul(", code)
		430	+
		431	+ def test_matmul_fusion_output_with_multiple_users(self):
		432	+ def model(a, b, denom, scale):
		433	+ mm = torch.mm(a, b).reshape(4, 8, 16)
		434	+ reduced = mm.float().sum(dim=(0, 1), keepdim=True)
		435	+ scaled_reduced = ((mm / denom) * scale).float().sum(
		436	+ dim=(0, 1), keepdim=True
		437	+ )
		438	+ return reduced, scaled_reduced
		439	+
		440	+ a = torch.normal(
		441	+ 0, 0.01, size=(32, 64), dtype=torch.float16, device="npu"
		442	+ )
		443	+ b = torch.normal(
		444	+ 0, 0.01, size=(64, 16), dtype=torch.float16, device="npu"
		445	+ )
		446	+ denom = torch.rand((4, 8, 1), dtype=torch.float16, device="npu") + 0.5
		447	+ scale = torch.normal(
		448	+ 0, 0.01, size=(4, 8, 16), dtype=torch.float16, device="npu"
		449	+ )
		450	+ with torch.no_grad():
		451	+ expect = model(a, b, denom, scale)
		452	+ result, codes = self._run_and_get_code_with_dvm(
		453	+ model, a, b, denom, scale
		454	+ )
		455	+
		456	+ code = "\n".join(codes)
		457	+ self.assertEqual(expect, result, atol=1e-2, rtol=1e-2)
		458	+ self.assertIn("k.matmul(", code)
		459	+
		460	+ def test_matmul_does_not_fuse_view_only_epilogue(self):
		461	+ def model(a, b):
		462	+ return torch.mm(a, b).reshape(8, 8, 512)
		463	+
		464	+ a = torch.normal(
		465	+ 0, 0.01, size=(64, 128), dtype=torch.float16, device="npu"
		466	+ )
		467	+ b = torch.normal(
		468	+ 0, 0.01, size=(128, 512), dtype=torch.float16, device="npu"
		469	+ )
		470	+ with torch.no_grad():
		471	+ expect = model(a, b)
		472	+ result, codes = self._run_and_get_code_with_dvm(model, a, b)
		473	+
		474	+ code = "\n".join(codes)
		475	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		476	+ self.assertIn("k.matmul(", code)
		477	+
		478	+ def test_matmul_fuses_view_with_pointwise_epilogue(self):
		479	+ def model(a, b, residual):
		480	+ view = torch.mm(a, b).reshape(8, 8, 512)
		481	+ return view + residual
		482	+
		483	+ a = torch.normal(
		484	+ 0, 0.01, size=(64, 128), dtype=torch.float16, device="npu"
		485	+ )
		486	+ b = torch.normal(
		487	+ 0, 0.01, size=(128, 512), dtype=torch.float16, device="npu"
		488	+ )
		489	+ residual = torch.normal(
		490	+ 0, 0.01, size=(8, 8, 512), dtype=torch.float16, device="npu"
		491	+ )
		492	+ with torch.no_grad():
		493	+ expect = model(a, b, residual)
		494	+ result, codes = self._run_and_get_code_with_dvm(model, a, b, residual)
		495	+
		496	+ code = "\n".join(codes)
		497	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		498	+ self.assertIn("k.matmul(", code)
		499	+
		500	+ def test_bmm_with_view_input_uses_dvm_fusion(self):
		501	+ def model(a, b):
		502	+ softmax = torch.softmax(a, dim=-1)
		503	+ return torch.bmm(softmax.reshape(4, 128, 128), b)
		504	+
		505	+ a = torch.normal(
		506	+ 0, 0.01, size=(2, 2, 128, 128), dtype=torch.float16, device="npu"
		507	+ )
		508	+ b = torch.normal(
		509	+ 0, 0.01, size=(4, 128, 64), dtype=torch.float16, device="npu"
		510	+ )
		511	+ with torch.no_grad():
		512	+ expect = model(a, b)
		513	+ result, codes = self._run_and_get_code_with_dvm(model, a, b)
		514	+
		515	+ code = "\n".join(codes)
		516	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		517	+ self.assertIn("k.matmul(", code)
		518	+
		519	+ def test_bmm_same_buffer_views_keep_distinct_input_meta(self):
		520	+ def model(x):
		521	+ lhs = x.reshape(4, 64, 128)
		522	+ rhs = x.reshape(4, 128, 64)
		523	+ return torch.bmm(lhs, rhs)
		524	+
		525	+ x = torch.normal(
		526	+ 0, 0.01, size=(4, 8192), dtype=torch.float16, device="npu"
		527	+ )
		528	+ with torch.no_grad():
		529	+ expect = model(x)
		530	+ result, codes = self._run_and_get_code_with_dvm(model, x)
		531	+
		532	+ code = "\n".join(codes)
		533	+ self.assertEqual(expect, result, atol=5e-3, rtol=5e-3)
		534	+ self.assertIn("k.matmul(", code)
		535	+
311	def test_copy_inplace_codegen(self):	536	def test_copy_inplace_codegen(self):
312	src = torch.randn((128,), dtype=torch.float32, device="npu")	537	src = torch.randn((128,), dtype=torch.float32, device="npu")
313	dst = torch.zeros((128,), dtype=torch.float32, device="npu")	538	dst = torch.zeros((128,), dtype=torch.float32, device="npu")

Mtest/npu/test_public_bindings.py+2-0

	@@ -667,6 +667,8 @@ class TestPublicBindings(TestCase):
667	"torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering",	667	"torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering",
668	"torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.scheduler",	668	"torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.scheduler",
669	"torch_npu._inductor.dvm",	669	"torch_npu._inductor.dvm",
		670	+ "torch_npu._inductor.dvm.config",
		671	+ "torch_npu._inductor.dvm.template",
670	"torch_npu._inductor.dvm.decomp",	672	"torch_npu._inductor.dvm.decomp",
671	"torch_npu._inductor.dvm.fx_pass",	673	"torch_npu._inductor.dvm.fx_pass",
672	"torch_npu._inductor.dvm.fx_test",	674	"torch_npu._inductor.dvm.fx_test",

Mtorch_npu/_inductor/dvm/config.py+6-0

	@@ -1,5 +1,7 @@
1	"""Shared configuration for DVM Inductor integration."""	1	"""Shared configuration for DVM Inductor integration."""
2		2
		3	+import os
		4	+
3	# Run post-launch DVM debug checks.	5	# Run post-launch DVM debug checks.
4	debug_mode = False	6	debug_mode = False
5	# Emit standalone FX regression cases for DVM-fused graphs.	7	# Emit standalone FX regression cases for DVM-fused graphs.
	@@ -8,5 +10,9 @@ dump_fx_test = False
8	view_fusion_level = 1	10	view_fusion_level = 1
9	# Use DVM-specific fusion rules that prevent post-reduction fusion.	11	# Use DVM-specific fusion rules that prevent post-reduction fusion.
10	disable_post_reduce_fusion = False	12	disable_post_reduce_fusion = False
		13	+# Enable DVM matmul fusion for mm, bmm, addmm, and baddbmm.
		14	+enable_matmul_fusion = (
		15	+ os.environ.get("INDUCTOR_DVM_ENABLE_MATMUL_FUSION", "0") == "1"
		16	+)
11	# Cast promoted BF16 vector-operation results back to BF16.	17	# Cast promoted BF16 vector-operation results back to BF16.
12	bf16_vector_keep_promoted = False	18	bf16_vector_keep_promoted = False

Mtorch_npu/_inductor/dvm/fx_pass.py+3-3

	@@ -32,12 +32,12 @@ def annotate_mm_transpose_flags(gm: torch.fx.GraphModule):
32	lhs = node.args[0]	32	lhs = node.args[0]
33	rhs = node.args[1]	33	rhs = node.args[1]
34	flag = True	34	flag = True
35	- elif node.target is aten.addmm.default:	35	+ elif node.target in (aten.addmm.default, aten.baddbmm.default):
36	- add = node.args[0]
37	lhs = node.args[1]	36	lhs = node.args[1]
38	rhs = node.args[2]	37	rhs = node.args[2]
		38	+ bias = node.args[0]
39	if (	39	if (
40	- add.meta["val"].dim() == 1	40	+ bias.meta["val"].dim() == 1
41	and node.kwargs.get("beta", 1) == 1	41	and node.kwargs.get("beta", 1) == 1
42	and node.kwargs.get("alpha", 1) == 1	42	and node.kwargs.get("alpha", 1) == 1
43	):	43	):

Mtorch_npu/_inductor/dvm/graph_build.py+9-4

	@@ -36,7 +36,6 @@ class DvmCodegenInterpreter(torch.fx.Interpreter):
36	):	36	):
37	super().__init__(gm)	37	super().__init__(gm)
38	self.gm = gm	38	self.gm = gm
39	- self.ktype = ktype
40	self.is_mix_kernel = annotate_mm_transpose_flags(gm)	39	self.is_mix_kernel = annotate_mm_transpose_flags(gm)
41	if is_dynamic is None:	40	if is_dynamic is None:
42	self.is_dynamic = is_fx_dynamic(gm)	41	self.is_dynamic = is_fx_dynamic(gm)
	@@ -50,8 +49,7 @@ class DvmCodegenInterpreter(torch.fx.Interpreter):
50	self.code = IndentedBuffer()	49	self.code = IndentedBuffer()
51		50
52	self.spec_nodes = set()	51	self.spec_nodes = set()
53	- if self.ktype == "vector" and self.need_spec():	52	+ self.set_kernel_ktype(ktype)
54	- self.ktype = "spec"
55	self.code.splice(f'\n"""\n{self.gm.print_readable(print_output=False)}\n"""')	53	self.code.splice(f'\n"""\n{self.gm.print_readable(print_output=False)}\n"""')
56	decorator = (	54	decorator = (
57	f"{chr(64)}dvm.kernel(ktype={self.ktype!r}, dyn_shape={self.is_dynamic})"	55	f"{chr(64)}dvm.kernel(ktype={self.ktype!r}, dyn_shape={self.is_dynamic})"
	@@ -60,6 +58,13 @@ class DvmCodegenInterpreter(torch.fx.Interpreter):
60	self.code.splice(f"def {self.KERNEL_NAME_PLACEHOLDER}(k):")	58	self.code.splice(f"def {self.KERNEL_NAME_PLACEHOLDER}(k):")
61	self.code.do_indent()	59	self.code.do_indent()
62		60
		61	+ def set_kernel_ktype(self, ktype: str) -> None:
		62	+ self.ktype = ktype
		63	+ if self.ktype != "split" and self.is_mix_kernel:
		64	+ self.ktype = "mix"
		65	+ elif self.ktype == "vector" and self.need_spec():
		66	+ self.ktype = "spec"
		67	+
63	def need_spec(self) -> bool:	68	def need_spec(self) -> bool:
64	self.spec_nodes.clear()	69	self.spec_nodes.clear()
65	for node in self.gm.graph.nodes:	70	for node in self.gm.graph.nodes:
	@@ -139,7 +144,7 @@ class DvmCodegenInterpreter(torch.fx.Interpreter):
139	if target in (aten.mm.default, aten.bmm.default):	144	if target in (aten.mm.default, aten.bmm.default):
140	args = (*args, meta.get("trans_a", False), meta.get("trans_b", False))	145	args = (*args, meta.get("trans_a", False), meta.get("trans_b", False))
141		146
142	- elif target is aten.addmm.default:	147	+ elif target in (aten.addmm.default, aten.baddbmm.default):
143	args = (	148	args = (
144	*args,	149	*args,
145	meta.get("trans_a", False),	150	meta.get("trans_a", False),

Mtorch_npu/_inductor/dvm/mlir_fusion.py+86-3

	@@ -1,10 +1,11 @@
		1	+from typing import List
1	import torch	2	import torch
2	from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation	3	from torch._higher_order_ops.triton_kernel_wrap import triton_kernel_wrapper_mutation
3	from torch._inductor import config	4	from torch._inductor import config
4	from torch._inductor.fx_passes.control_dependencies import control_deps	5	from torch._inductor.fx_passes.control_dependencies import control_deps
5	from torch._inductor.codegen.common import IndentedBuffer, register_backend_for_device	6	from torch._inductor.codegen.common import IndentedBuffer, register_backend_for_device
6	from torch._inductor.codegen.simd import code_hash, SIMDKernel	7	from torch._inductor.codegen.simd import code_hash, SIMDKernel
7	-from torch._inductor.scheduler import WhyNoFuse	8	+from torch._inductor.scheduler import SchedulerNode, WhyNoFuse
8	from torch._inductor.utils import get_fused_kernel_name	9	from torch._inductor.utils import get_fused_kernel_name
9	from torch._inductor.virtualized import V	10	from torch._inductor.virtualized import V
10	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config	11	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config
	@@ -13,6 +14,7 @@ from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.mlir import (
13	)	14	)
14	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.meta_kernel import (	15	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.meta_kernel import (
15	NpuMetaScheduling,	16	NpuMetaScheduling,
		17	+ create_fx_from_snodes_by_traced_graph,
16	)	18	)
17	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.wrapper import (	19	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.wrapper import (
18	NpuMlirWrapperCodeGen,	20	NpuMlirWrapperCodeGen,
	@@ -24,7 +26,11 @@ from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.utils import (
24	get_num_call_functions,	26	get_num_call_functions,
25	)	27	)
26		28
27	-from .config import disable_post_reduce_fusion, dump_fx_test	29	+from .config import (
		30	+ disable_post_reduce_fusion,
		31	+ dump_fx_test,
		32	+ enable_matmul_fusion,
		33	+)
28	from .decomp import patch_decomp	34	from .decomp import patch_decomp
29	from .fx_test import generate_dvm_fx_case	35	from .fx_test import generate_dvm_fx_case
30	from .graph_build import DvmCodegenInterpreter	36	from .graph_build import DvmCodegenInterpreter
	@@ -34,6 +40,11 @@ from .op_emitter import (
34	DVM_SUPPORT_TYPE,	40	DVM_SUPPORT_TYPE,
35	_extra_int_types,	41	_extra_int_types,
36	)	42	)
		43	+from .template import (
		44	+ can_fuse_dvm_epilogue,
		45	+ DvmTemplateBuffer,
		46	+ patch_dvm_matmul_template_fusion,
		47	+)
37		48
38	aten = torch.ops.aten	49	aten = torch.ops.aten
39	prims = torch.ops.prims	50	prims = torch.ops.prims
	@@ -226,6 +237,12 @@ class NpuDvmScheduling(NpuMetaScheduling):
226	return kernel_name	237	return kernel_name
227		238
228	def can_fuse_vertical(self, node1, node2):	239	def can_fuse_vertical(self, node1, node2):
		240	+ template1 = node1.get_template_node()
		241	+ template2 = node2.get_template_node()
		242	+ if isinstance(template1, DvmTemplateBuffer):
		243	+ return can_fuse_dvm_epilogue(node1, node2)
		244	+ if isinstance(template2, DvmTemplateBuffer):
		245	+ return False
229	if not disable_post_reduce_fusion:	246	if not disable_post_reduce_fusion:
230	return super().can_fuse_vertical(node1, node2)	247	return super().can_fuse_vertical(node1, node2)
231		248
	@@ -253,10 +270,75 @@ class NpuDvmScheduling(NpuMetaScheduling):
253	return numel1 == numel2	270	return numel1 == numel2
254		271
255	def can_fuse_horizontal(self, node1, node2):	272	def can_fuse_horizontal(self, node1, node2):
		273	+ template1 = node1.get_template_node()
		274	+ template2 = node2.get_template_node()
		275	+ if isinstance(template1, DvmTemplateBuffer) or isinstance(
		276	+ template2, DvmTemplateBuffer
		277	+ ):
		278	+ return False
256	if not disable_post_reduce_fusion:	279	if not disable_post_reduce_fusion:
257	return super().can_fuse_horizontal(node1, node2)	280	return super().can_fuse_horizontal(node1, node2)
258	return False	281	return False
259		282
		283	+ def codegen_template(
		284	+ self,
		285	+ template_node: SchedulerNode,
		286	+ epilogue_nodes: List[SchedulerNode],
		287	+ prologue_nodes: List[SchedulerNode] = (),
		288	+ ):
		289	+ template_buffer = template_node.get_template_node()
		290	+ if not isinstance(template_buffer, DvmTemplateBuffer):
		291	+ return super().codegen_template(
		292	+ template_node, epilogue_nodes, prologue_nodes
		293	+ )
		294	+ if prologue_nodes:
		295	+ raise RuntimeError(
		296	+ "DVM matmul template only supports pointwise epilogue fusion"
		297	+ )
		298	+
		299	+ snodes = [template_node, *epilogue_nodes]
		300	+ fused_node_names = set()
		301	+ for snode in snodes:
		302	+ fused_node_names.update(snode.get_operation_names())
		303	+ removed_buffers = {
		304	+ name
		305	+ for snode in snodes
		306	+ for name in snode.get_buffer_names()
		307	+ if self.scheduler.can_buffer_be_removed_through_fusion(
		308	+ name, fused_node_names
		309	+ )
		310	+ }
		311	+ V.graph.removed_buffers \|= removed_buffers
		312	+
		313	+ traced_graph, call_args, compile_kwargs = (
		314	+ create_fx_from_snodes_by_traced_graph(snodes, None)
		315	+ )
		316	+ mlir_kernel = self.meta_kernel_type(
		317	+ traced_graph, snodes, call_args, **compile_kwargs
		318	+ )
		319	+ with V.set_kernel_handler(mlir_kernel):
		320	+ src_code = mlir_kernel.codegen_kernel()
		321	+
		322	+ need_trans_input = getattr(
		323	+ mlir_kernel.dvm_codegen, "need_trans_input", ()
		324	+ )
		325	+ for index, arg in enumerate(call_args):
		326	+ if arg in template_buffer.input_bindings:
		327	+ arg = V.graph.wrapper_code.val_to_arg_str(
		328	+ template_buffer.input_bindings[arg]
		329	+ )
		330	+ if index < len(need_trans_input) and need_trans_input[index]:
		331	+ arg += ".mT"
		332	+ call_args[index] = arg
		333	+
		334	+ kernel_name = self.define_kernel(src_code, mlir_kernel, traced_graph)
		335	+ with V.set_kernel_handler(mlir_kernel):
		336	+ for node in snodes:
		337	+ node.mark_run()
		338	+ self.codegen_comment(snodes)
		339	+ mlir_kernel.call_kernel(kernel_name, template_node.node)
		340	+ self.free_buffers_in_scheduler()
		341	+
260		342
261	def _patch_lowering_type_checks():	343	def _patch_lowering_type_checks():
262	import torch._inductor.graph as inductor_graph	344	import torch._inductor.graph as inductor_graph
	@@ -280,7 +362,6 @@ def _patch_lowering_type_checks():
280	return False	362	return False
281	if node.target is aten.lift_fresh_copy.default:	363	if node.target is aten.lift_fresh_copy.default:
282	return False	364	return False
283	-
284	return not _is_node_supported_by_dvm_rule(node, allow_common_rule=True)	365	return not _is_node_supported_by_dvm_rule(node, allow_common_rule=True)
285		366
286	inductor_lowering.fallback_node_due_to_unsupported_type = (	367	inductor_lowering.fallback_node_due_to_unsupported_type = (
	@@ -371,6 +452,8 @@ class DvmMlirFusionPatch:
371	patch_decomp()	452	patch_decomp()
372	_patch_lowering_type_checks()	453	_patch_lowering_type_checks()
373	_patch_lowering()	454	_patch_lowering()
		455	+ if enable_matmul_fusion:
		456	+ patch_dvm_matmul_template_fusion()
374	register_backend_for_device(	457	register_backend_for_device(
375	"npu", NpuDvmScheduling, NpuMlirWrapperCodeGen	458	"npu", NpuDvmScheduling, NpuMlirWrapperCodeGen
376	)	459	)

Mtorch_npu/_inductor/dvm/op_emitter.py+13-25

	@@ -95,7 +95,6 @@ def mm_rule(node: torch.fx.Node):
95	UINT16_MAX = (1 << 16) - 1	95	UINT16_MAX = (1 << 16) - 1
96	UINT8_MAX = (1 << 8) - 1	96	UINT8_MAX = (1 << 8) - 1
97	MAX_INNER = UINT16_MAX - UINT8_MAX	97	MAX_INNER = UINT16_MAX - UINT8_MAX
98	- SMALL_OUTPUT_MAX = 256
99		98
100	def inner_axis_length(t: torch._subclasses.FakeTensor):	99	def inner_axis_length(t: torch._subclasses.FakeTensor):
101	if _is_last2_transpose_tensor(t):	100	if _is_last2_transpose_tensor(t):
	@@ -114,15 +113,6 @@ def mm_rule(node: torch.fx.Node):
114	return False	113	return False
115	return True	114	return True
116		115
117	- def check_output(output_node):
118	- t = output_node.meta["val"]
119	- last_two_dims = t.shape[-2:]
120	- if all(not isinstance(dim, torch.SymInt) for dim in last_two_dims) and all(
121	- dim <= SMALL_OUTPUT_MAX for dim in last_two_dims
122	- ):
123	- return False
124	- return True
125	-
126	def check_k1_fusion(lhs_node, rhs_node):	116	def check_k1_fusion(lhs_node, rhs_node):
127	lhs_t = lhs_node.meta["val"]	117	lhs_t = lhs_node.meta["val"]
128	rhs_t = rhs_node.meta["val"]	118	rhs_t = rhs_node.meta["val"]
	@@ -136,10 +126,10 @@ def mm_rule(node: torch.fx.Node):
136	)	126	)
137	return True	127	return True
138		128
139	- if node.target in [aten.mm.default, aten.bmm.default]:	129	+ if node.target in (aten.mm.default, aten.bmm.default):
140	lhs = node.args[0]	130	lhs = node.args[0]
141	rhs = node.args[1]	131	rhs = node.args[1]
142	- elif node.target is aten.addmm.default:	132	+ elif node.target in (aten.addmm.default, aten.baddbmm.default):
143	lhs = node.args[1]	133	lhs = node.args[1]
144	rhs = node.args[2]	134	rhs = node.args[2]
145	else:	135	else:
	@@ -147,9 +137,7 @@ def mm_rule(node: torch.fx.Node):
147	if node.meta["val"].dtype not in (torch.float16, torch.bfloat16):	137	if node.meta["val"].dtype not in (torch.float16, torch.bfloat16):
148	return False	138	return False
149		139
150	- return (	140	+ return check(lhs) and check(rhs) and check_k1_fusion(lhs, rhs)
151	- check(lhs) and check(rhs) and check_output(node) and check_k1_fusion(lhs, rhs)
152	- )
153		141
154		142
155	class DvmOpInfo:	143	class DvmOpInfo:
	@@ -544,20 +532,20 @@ def matmul(x, y, trans_a, trans_b):
544	return f"k.matmul({x}, {y}, {trans_a}, {trans_b})"	532	return f"k.matmul({x}, {y}, {trans_a}, {trans_b})"
545		533
546		534
547	-def matmul_bias(bias, x, y, trans_a, trans_b, beta=1, alpha=1):	535	+def matmul_bias(bias, mat1, mat2, trans_a, trans_b):
548	- return f"k.matmul({x}, {y}, {trans_a}, {trans_b},{bias})"	536	+ return f"k.matmul({mat1}, {mat2}, {trans_a}, {trans_b},{bias})"
549		537
550		538
551	-@register_dvm_op(aten.addmm.default, rule=mm_rule)	539	+@register_dvm_op(aten.addmm.default, aten.baddbmm.default, rule=mm_rule)
552	-def addmm(z, x, y, trans_a, trans_b, use_bias, beta=1, alpha=1):	540	+def addmm(inp, mat1, mat2, trans_a, trans_b, use_bias, beta=1, alpha=1):
553	if use_bias:	541	if use_bias:
554	- return matmul_bias(z, x, y, trans_a, trans_b)	542	+ return matmul_bias(inp, mat1, mat2, trans_a, trans_b)
555	- if beta != 1:	543	+ acc = matmul(mat1, mat2, trans_a, trans_b)
556	- z = mul(z, beta)
557	- mm = matmul(x, y, trans_a, trans_b)
558	if alpha != 1:	544	if alpha != 1:
559	- mm = mul(mm, alpha)	545	+ acc = mul(acc, alpha)
560	- return add(mm, z)	546	+ if beta != 1:
		547	+ inp = mul(inp, beta)
		548	+ return add(acc, inp)
561		549
562		550
563	def load(shape, dtype):	551	def load(shape, dtype):

Atorch_npu/_inductor/dvm/template.py+292-0

@@ -0,0 +1,292 @@
	1	+import torch
	2	+from torch._inductor import config, ir, scheduler
	3	+from torch._inductor.dependencies import MemoryDep
	4	+from torch._inductor.kernel.mm_common import mm_args
	5	+from torch._inductor.utils import sympy_product
	6	+from torch._inductor.virtualized import V
	7	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config
	8	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch import (
	9	+ ir as npu_ir,
	10	+ lowering as npu_lowering,
	11	+)
	12	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering import (
	13	+ fetch_graphs,
	14	+ merge_traced_graphs,
	15	+)
	16	+from torch_npu._inductor.lowering_common import TracedGraph, create_fake_input
	17	+
	18	+from .op_emitter import mm_rule
	19	+
	20	+aten = torch.ops.aten
	21	+_orig_npu_subtract_graph = npu_ir.subtract_graph
	22	+
	23	+
	24	+def _make_template_input_graph(inp, name=None):
	25	+ traced_graph = TracedGraph()
	26	+ placeholder = traced_graph.graph.placeholder(name or inp.get_name())
	27	+ placeholder.meta["val"] = create_fake_input(
	28	+ inp.get_size(), inp.get_stride(), inp.get_device(), inp.get_dtype()
	29	+ )
	30	+ traced_graph.last_node = placeholder
	31	+ return traced_graph
	32	+
	33	+
	34	+def _make_matmul_input_graphs(inputs, node_name):
	35	+ graphs = []
	36	+ bindings = {}
	37	+ for index, inp in enumerate(inputs):
	38	+ logical_name = f"_dvm_{node_name}_mat_input_{index}"
	39	+ graphs.append(_make_template_input_graph(inp, logical_name))
	40	+ bindings[logical_name] = inp
	41	+ return graphs, bindings
	42	+
	43	+
	44	+class _DvmTemplateGraph(TracedGraph):
	45	+ def __init__(self, traced_graph, input_bindings):
	46	+ self.graph = traced_graph.graph
	47	+ self.last_node = traced_graph.last_node
	48	+ self.sym_nodes = traced_graph.sym_nodes
	49	+ self.input_bindings = input_bindings
	50	+
	51	+ def get_placeholder_names(self):
	52	+ return {
	53	+ self.input_bindings[name].get_name()
	54	+ if name in self.input_bindings
	55	+ else name
	56	+ for name in super().get_placeholder_names()
	57	+ }
	58	+
	59	+
	60	+class DvmTemplateBuffer(ir.TemplateBuffer):
	61	+ def __init__(self, layout, inputs, traced_graph, input_bindings):
	62	+ self.traced_graph = _DvmTemplateGraph(traced_graph, input_bindings)
	63	+ self.input_bindings = input_bindings
	64	+ super().__init__(layout, inputs, make_kernel_render=None)
	65	+ # Reuse NPU meta_kernel's snode.node.data.traced_graph rebuild path.
	66	+ self.data = self
	67	+
	68	+ def get_traced_graph(self):
	69	+ return _make_template_input_graph(self)
	70	+
	71	+
	72	+def _subtract_dvm_template_graph(graph1, graph2, node_name=None):
	73	+ if (
	74	+ node_name is not None
	75	+ and isinstance(V.graph.try_get_buffer(node_name), DvmTemplateBuffer)
	76	+ and graph2.last_node.op == "placeholder"
	77	+ and graph2.last_node.name == node_name
	78	+ ):
	79	+ node_name = None
	80	+ return _orig_npu_subtract_graph(graph1, graph2, node_name)
	81	+
	82	+
	83	+def _register_dvm_mm_template_lowerings():
	84	+ for op in (aten.mm, aten.bmm, aten.addmm, aten.baddbmm):
	85	+ if op in anir_config.FALLBACK_LIST:
	86	+ anir_config.FALLBACK_LIST.remove(op)
	87	+ if op not in anir_config.GENERATE_LIST:
	88	+ anir_config.GENERATE_LIST.append(op)
	89	+
	90	+ def make_mm_template(op, mat1, mat2, *, layout=None):
	91	+ if V.graph.cpp_wrapper:
	92	+ return npu_lowering.fallback_handler(op)(mat1, mat2)
	93	+
	94	+ _, _, k, layout, mat1, mat2 = mm_args(mat1, mat2, layout=layout)
	95	+ current_node = V.graph.current_node
	96	+ if current_node is None or "val" not in current_node.meta:
	97	+ return npu_lowering.fallback_handler(op)(mat1, mat2)
	98	+ if k == 1:
	99	+ mat1 = npu_lowering.expand(ir.TensorBox.create(mat1), layout.size)
	100	+ mat2 = npu_lowering.expand(ir.TensorBox.create(mat2), layout.size)
	101	+ return npu_lowering.lowerings[aten.mul.Tensor](mat1, mat2)
	102	+
	103	+ inputs = [mat1, mat2]
	104	+ input_graphs, input_bindings = _make_matmul_input_graphs(
	105	+ inputs, current_node.name
	106	+ )
	107	+ traced_graph = merge_traced_graphs(input_graphs, op, current_node.name)
	108	+ traced_graph.last_node.meta["val"] = current_node.meta["val"]
	109	+ if not mm_rule(traced_graph.last_node):
	110	+ return npu_lowering.fallback_handler(op)(mat1, mat2)
	111	+
	112	+ return ir.TensorBox.create(
	113	+ DvmTemplateBuffer(
	114	+ layout,
	115	+ inputs,
	116	+ traced_graph,
	117	+ input_bindings=input_bindings,
	118	+ )
	119	+ )
	120	+
	121	+ @npu_lowering.register_lowering([aten.mm.default], type_promotion_kind=None)
	122	+ def dvm_mm(mat1, mat2, *, layout=None):
	123	+ return make_mm_template(aten.mm.default, mat1, mat2, layout=layout)
	124	+
	125	+ @npu_lowering.register_lowering([aten.bmm.default], type_promotion_kind=None)
	126	+ def dvm_bmm(mat1, mat2, *, layout=None):
	127	+ return make_mm_template(aten.bmm.default, mat1, mat2, layout=layout)
	128	+
	129	+ def make_addmm_template(
	130	+ op,
	131	+ inp,
	132	+ mat1,
	133	+ mat2,
	134	+ *,
	135	+ alpha=1,
	136	+ beta=1,
	137	+ layout=None,
	138	+ ):
	139	+ _, _, k, layout, mat1, mat2, expanded_inp = mm_args(
	140	+ mat1, mat2, inp, layout=layout
	141	+ )
	142	+ if op is aten.addmm.default and k == 1:
	143	+ mul = npu_lowering.lowerings[aten.mul.Tensor]
	144	+ add = npu_lowering.lowerings[aten.add.Tensor]
	145	+ inp = ir.TensorBox.create(expanded_inp)
	146	+ product = mul(mat1, mat2)
	147	+ if alpha != 1:
	148	+ product = mul(product, alpha)
	149	+ if beta != 1:
	150	+ inp = mul(inp, beta)
	151	+ return add(product, inp)
	152	+
	153	+ current_node = V.graph.current_node
	154	+ matrix_inputs = [mat1, mat2]
	155	+ matrix_graphs, input_bindings = _make_matmul_input_graphs(
	156	+ matrix_inputs, current_node.name
	157	+ )
	158	+ input_graphs = [fetch_graphs([inp]), matrix_graphs]
	159	+ traced_graph = merge_traced_graphs(
	160	+ input_graphs, op, current_node.name, alpha=alpha, beta=beta
	161	+ )
	162	+ traced_graph.last_node.meta["val"] = current_node.meta["val"]
	163	+ if not mm_rule(traced_graph.last_node):
	164	+ return npu_lowering.fallback_handler(op)(
	165	+ inp, mat1, mat2, alpha=alpha, beta=beta
	166	+ )
	167	+
	168	+ return ir.TensorBox.create(
	169	+ DvmTemplateBuffer(
	170	+ layout,
	171	+ [inp, *matrix_inputs],
	172	+ traced_graph,
	173	+ input_bindings=input_bindings,
	174	+ )
	175	+ )
	176	+
	177	+ @npu_lowering.register_lowering([aten.addmm.default], type_promotion_kind=None)
	178	+ def dvm_addmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
	179	+ return make_addmm_template(
	180	+ aten.addmm.default,
	181	+ inp,
	182	+ mat1,
	183	+ mat2,
	184	+ alpha=alpha,
	185	+ beta=beta,
	186	+ layout=layout,
	187	+ )
	188	+
	189	+ @npu_lowering.register_lowering([aten.baddbmm.default], type_promotion_kind=None)
	190	+ def dvm_baddbmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None):
	191	+ return make_addmm_template(
	192	+ aten.baddbmm.default,
	193	+ inp,
	194	+ mat1,
	195	+ mat2,
	196	+ alpha=alpha,
	197	+ beta=beta,
	198	+ layout=layout,
	199	+ )
	200	+
	201	+
	202	+def _keep_addmm_for_dvm_template():
	203	+ from torch._inductor.fx_passes import post_grad
	204	+
	205	+ for entries in post_grad.pass_patterns[2].patterns.values():
	206	+ for entry in entries:
	207	+ if entry.extra_check is post_grad.should_prefer_unfused_addmm:
	208	+ entry.extra_check = lambda _match: False
	209	+ return
	210	+
	211	+
	212	+def patch_dvm_matmul_template_fusion() -> None:
	213	+ _keep_addmm_for_dvm_template()
	214	+ npu_ir.subtract_graph = _subtract_dvm_template_graph
	215	+ _register_dvm_mm_template_lowerings()
	216	+
	217	+
	218	+def _buffer_numel(buffer):
	219	+ buffer = getattr(buffer, "node", buffer)
	220	+ return sympy_product(buffer.get_size())
	221	+
	222	+
	223	+def _has_unsupported_epilogue_broadcast(node: scheduler.SchedulerNode) -> bool:
	224	+ output_shape = tuple(node.node.get_size())
	225	+ if len(output_shape) <= 2:
	226	+ return False
	227	+
	228	+ for dep in node.read_writes.reads:
	229	+ if not isinstance(dep, MemoryDep):
	230	+ continue
	231	+
	232	+ input_buffer = V.graph.try_get_buffer(dep.name)
	233	+ input_buffer = getattr(input_buffer, "node", input_buffer)
	234	+ if input_buffer is None or not hasattr(input_buffer, "get_size"):
	235	+ return True
	236	+
	237	+ input_shape = tuple(input_buffer.get_size())
	238	+ if len(input_shape) > len(output_shape):
	239	+ return True
	240	+ input_shape = (1,) * (len(output_shape) - len(input_shape)) + input_shape
	241	+ outer_shape_mismatch = any(
	242	+ not V.graph.sizevars.statically_known_equals(input_dim, output_dim)
	243	+ for input_dim, output_dim in zip(input_shape[:-2], output_shape[:-2])
	244	+ )
	245	+ is_broadcast_access = any(
	246	+ var not in dep.index.free_symbols for var in dep.var_names
	247	+ )
	248	+ if outer_shape_mismatch and is_broadcast_access:
	249	+ return True
	250	+ return False
	251	+
	252	+
	253	+def can_fuse_dvm_epilogue(
	254	+ template_node: scheduler.BaseSchedulerNode,
	255	+ epilogue_node: scheduler.BaseSchedulerNode,
	256	+) -> bool:
	257	+ if not config.epilogue_fusion:
	258	+ return False
	259	+ template_buffer = template_node.get_template_node()
	260	+ if not isinstance(template_buffer, DvmTemplateBuffer):
	261	+ return False
	262	+ if isinstance(epilogue_node.get_template_node(), DvmTemplateBuffer):
	263	+ return False
	264	+ if epilogue_node.is_reduction():
	265	+ return False
	266	+ _, (numel1, rnumel1) = template_node.group
	267	+ _, (numel2, rnumel2) = epilogue_node.group
	268	+ if numel1 != numel2 or rnumel1 != rnumel2:
	269	+ return False
	270	+
	271	+ matmul_numel = _buffer_numel(template_buffer)
	272	+ if not all(
	273	+ V.graph.sizevars.statically_known_equals(
	274	+ matmul_numel, _buffer_numel(output)
	275	+ )
	276	+ for output in epilogue_node.get_outputs()
	277	+ ):
	278	+ return False
	279	+ if not (template_node.get_buffer_names() & epilogue_node.used_buffer_names()):
	280	+ return False
	281	+
	282	+ for node in epilogue_node.get_nodes():
	283	+ if (
	284	+ not isinstance(node, scheduler.SchedulerNode)
	285	+ or not isinstance(node.node, ir.ComputedBuffer)
	286	+ or not isinstance(node.node.data, ir.Pointwise)
	287	+ ):
	288	+ return False
	289	+ if _has_unsupported_epilogue_broadcast(node):
	290	+ return False
	291	+
	292	+ return True