feat(_inductor): Support DVM for the r2.12 branch-pytorch-AtomGit

M.gitmodules+4-0

	@@ -21,6 +21,10 @@
21	path = third_party/torch-mlir	21	path = third_party/torch-mlir
22	url = https://gitcode.com/gh_mirrors/to/torch-mlir.git	22	url = https://gitcode.com/gh_mirrors/to/torch-mlir.git
23	update = none	23	update = none
		24	+[submodule "third_party/dvm/dvm"]
		25	+ path = third_party/dvm/dvm
		26	+ url = https://gitcode.com/mindspore/dvm.git
		27	+ branch = r2.10
24	[submodule "third_party/acl_src/runtime"]	28	[submodule "third_party/acl_src/runtime"]
25	path = third_party/acl_src/runtime	29	path = third_party/acl_src/runtime
26	url = https://gitcode.com/cann/runtime.git	30	url = https://gitcode.com/cann/runtime.git

MCMakeLists.txt+4-0

	@@ -321,6 +321,10 @@ endif()
321		321
322	add_subdirectory(${TORCHNPU_THIRD_PARTY_ROOT}/fmt EXCLUDE_FROM_ALL)	322	add_subdirectory(${TORCHNPU_THIRD_PARTY_ROOT}/fmt EXCLUDE_FROM_ALL)
323		323
		324	+add_subdirectory(${PROJECT_SOURCE_DIR}/third_party/dvm)
		325	+add_dependencies(${PLUGIN_NAME} dvm_build)
		326	+target_link_libraries(${PLUGIN_NAME} PRIVATE ${PROJECT_SOURCE_DIR}/third_party/dvm/dvm/libdvm.a)
		327	+
324	link_directories(${PYTORCH_INSTALL_DIR}/lib)	328	link_directories(${PYTORCH_INSTALL_DIR}/lib)
325	link_directories(${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs)	329	link_directories(${TORCHNPU_THIRD_PARTY_ROOT}/acl/libs)
326		330

Atest/_inductor/test_dvm_elu_native.py+52-0

@@ -0,0 +1,52 @@
	1	+# Owner(s): ["module: tests"]
	2	+"""End-to-end: elu/elu_backward are numerically correct through the DVM
	3	+backend's native aclnn fallback (elu is excluded from decomposition; the policy
	4	+assertions live in test_dvm_decomp.py).
	5	+
	6	+The DVM backend is pinned via torch.compile(options={"npu_backend": "dvm"}).
	7	+We deliberately do NOT import torch_npu._inductor at module scope: importing it
	8	+loads a backend at import time, which would turn the first torch.compile into a
	9	+mid-process backend switch (default -> dvm). Plain ``import torch_npu`` does not
	10	+load _inductor.
	11	+"""
	12	+import unittest
	13	+
	14	+import torch
	15	+import torch_npu
	16	+from torch.testing._internal.common_utils import (
	17	+ instantiate_parametrized_tests,
	18	+ parametrize,
	19	+ run_tests,
	20	+ TestCase,
	21	+)
	22	+
	23	+
	24	+@unittest.skipIf(not torch_npu.npu.is_available(), "requires an NPU device")
	25	+class TestDvmEluNative(TestCase):
	26	+ @parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
	27	+ def test_elu_forward_backward_matches_eager(self, dtype):
	28	+ tol = 1e-3 if dtype == torch.float32 else 4e-3
	29	+ # build in fp32 then cast (npu normal kernel has no bf16 support)
	30	+ ref = torch.randn(256, 4096, device="npu").to(dtype)
	31	+ x_e = ref.detach().clone().requires_grad_(True)
	32	+ x_c = ref.detach().clone().requires_grad_(True)
	33	+
	34	+ def fn(t):
	35	+ return torch.nn.functional.elu(t)
	36	+
	37	+ out_e = fn(x_e)
	38	+ out_e.float().sum().backward()
	39	+
	40	+ compiled = torch.compile(
	41	+ fn, backend="inductor", options={"npu_backend": "dvm"}
	42	+ )
	43	+ out_c = compiled(x_c)
	44	+ out_c.float().sum().backward()
	45	+
	46	+ self.assertEqual(out_e, out_c, atol=tol, rtol=tol)
	47	+ self.assertEqual(x_e.grad, x_c.grad, atol=tol, rtol=tol)
	48	+
	49	+
	50	+instantiate_parametrized_tests(TestDvmEluNative)
	51	+if __name__ == "__main__":
	52	+ run_tests()

Atest/_inductor/test_dvm_graph_fusion.py+175-0

@@ -0,0 +1,175 @@
	1	+from unittest import mock
	2	+
	3	+import torch
	4	+
	5	+from torch.testing._internal.common_utils import TestCase
	6	+from torch.testing._internal.common_utils import (
	7	+ run_tests,
	8	+ parametrize,
	9	+ instantiate_parametrized_tests,
	10	+)
	11	+from torch_npu._inductor.dvm.graph_fusion import (
	12	+ DvmGraphFusionPatch,
	13	+ _dvm_generate_fallback_kernel,
	14	+ _fused_metas,
	15	+)
	16	+from torch_npu._inductor.dvm.graph_build import is_fx_dynamic
	17	+
	18	+
	19	+class TestModule(torch.nn.Module):
	20	+ def __init__(self, args, *kwargs):
	21	+ super().__init__(args, *kwargs)
	22	+
	23	+ def forward(self, a, b, c):
	24	+ add = a + b
	25	+ mul = add * c
	26	+ return torch.sum(mul, dim=(0,), keepdim=True) + 1
	27	+
	28	+
	29	+class MatMulModule(torch.nn.Module):
	30	+ def __init__(self, args, *kwargs):
	31	+ super().__init__(args, *kwargs)
	32	+
	33	+ def forward(self, a, b):
	34	+ mm = torch.mm(a.t(), b)
	35	+ mm = mm.to(torch.float32)
	36	+ return mm + 1
	37	+
	38	+
	39	+class TestDvmByGraphFusion(TestCase):
	40	+ @parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16])
	41	+ @parametrize("is_dynamic", [True, False])
	42	+ def test_basic_partitioning(self, dtype, is_dynamic):
	43	+ a = torch.normal(0, 0.1, size=(512, 1024), dtype=dtype).npu()
	44	+ b = torch.normal(0, 0.1, size=(512, 1024), dtype=torch.float16).npu()
	45	+ c = torch.normal(0, 0.1, size=(1, 1024), dtype=dtype).npu()
	46	+ model = TestModule()
	47	+
	48	+ with DvmGraphFusionPatch():
	49	+ dvm_compiled_model = torch.compile(
	50	+ model, backend="inductor", dynamic=is_dynamic
	51	+ )
	52	+ with torch.no_grad():
	53	+ expect = model(a, b, c)
	54	+ result = dvm_compiled_model(a, b, c)
	55	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	56	+
	57	+ @parametrize("dtype", [torch.float16])
	58	+ @parametrize("is_dynamic", [False])
	59	+ def test_basic_partitioning_npugraph(self, dtype, is_dynamic):
	60	+ a = torch.normal(0, 0.1, size=(512, 1024), dtype=dtype).npu()
	61	+ b = torch.normal(0, 0.1, size=(512, 1024), dtype=torch.float16).npu()
	62	+ c = torch.normal(0, 0.1, size=(1, 1024), dtype=dtype).npu()
	63	+ model = TestModule()
	64	+ with DvmGraphFusionPatch():
	65	+ dvm_compiled_model = torch.compile(
	66	+ model,
	67	+ backend="inductor",
	68	+ dynamic=is_dynamic,
	69	+ options={"triton.cudagraphs": True},
	70	+ )
	71	+ with torch.no_grad():
	72	+ expect = model(a, b, c)
	73	+ result = dvm_compiled_model(a, b, c)
	74	+ result = dvm_compiled_model(a, b, c)
	75	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	76	+
	77	+ @parametrize("k,n,m", [(1280, 2560, 640), (1024, 1280, 2560)])
	78	+ @parametrize("dtype", [torch.float16, torch.bfloat16])
	79	+ @parametrize("is_dynamic", [True, False])
	80	+ def test_matmul(self, k, n, m, dtype, is_dynamic):
	81	+ a = torch.normal(0, 0.02, size=(k, n), dtype=dtype).npu()
	82	+ b = torch.normal(0, 0.02, size=(k, m), dtype=dtype).npu()
	83	+ model = MatMulModule()
	84	+
	85	+ with DvmGraphFusionPatch():
	86	+ dvm_compiled_model = torch.compile(
	87	+ model, backend="inductor", dynamic=is_dynamic
	88	+ )
	89	+ with torch.no_grad():
	90	+ expect = model(a, b)
	91	+ result = dvm_compiled_model(a, b)
	92	+ self.assertEqual(expect, result, atol=2e-3, rtol=2e-3)
	93	+
	94	+
	95	+instantiate_parametrized_tests(TestDvmByGraphFusion)
	96	+
	97	+
	98	+class _AddModule(torch.nn.Module):
	99	+ def forward(self, x):
	100	+ return x + 1
	101	+
	102	+
	103	+class TestDvmFallbackStridePatchGuard(TestCase):
	104	+ """Guard adb97b9: skip stride patch in dynamic fused subgraph codegen."""
	105	+
	106	+ def _make_fallback_kernel(self, fused_id=0):
	107	+ fallback_kernel = mock.MagicMock()
	108	+ fallback_kernel.op_overload._name = "dvm::fused_graph_0"
	109	+ fallback_kernel.codegen_args.return_value = ["buf0", fused_id]
	110	+ fallback_kernel.codegen_kwargs.return_value = []
	111	+ fallback_kernel.get_name.return_value = "buf_out"
	112	+ return fallback_kernel
	113	+
	114	+ def _make_codegen_wrapper(self):
	115	+ wrapper = mock.MagicMock()
	116	+ wrapper.header = mock.MagicMock()
	117	+ return wrapper
	118	+
	119	+ def _make_fused_meta(self):
	120	+ meta = mock.MagicMock()
	121	+ meta.gm = mock.MagicMock()
	122	+ meta.name = "dvm_graph_fused_0"
	123	+ codegen = mock.MagicMock()
	124	+ codegen.cont_flag_input = [True]
	125	+ codegen.need_trans_input = [False]
	126	+ meta.codegen.return_value = (codegen, "# dvm kernel\n")
	127	+ return meta
	128	+
	129	+ @mock.patch(
	130	+ "torch_npu._inductor.dvm.graph_fusion.patch_gm_placeholder_strides_from_codegen_args"
	131	+ )
	132	+ def test_fallback_kernel_stride_patch_guarded_by_is_fx_dynamic(self, mock_patch):
	133	+ gm_static = torch.fx.symbolic_trace(_AddModule())
	134	+ placeholder = next(n for n in gm_static.graph.nodes if n.op == "placeholder")
	135	+ placeholder.meta["val"] = torch.randn(2, 3)
	136	+ self.assertFalse(is_fx_dynamic(gm_static))
	137	+
	138	+ batch = torch.export.Dim("batch", min=1, max=1024)
	139	+ exported = torch.export.export(
	140	+ _AddModule(),
	141	+ (torch.randn(2, 3),),
	142	+ dynamic_shapes={"x": {0: batch}},
	143	+ )
	144	+ self.assertTrue(is_fx_dynamic(exported.graph_module))
	145	+
	146	+ meta = self._make_fused_meta()
	147	+ try:
	148	+ _fused_metas[0] = meta
	149	+ with mock.patch(
	150	+ "torch_npu._inductor.dvm.graph_fusion.is_fx_dynamic", return_value=True
	151	+ ):
	152	+ _dvm_generate_fallback_kernel(
	153	+ self._make_codegen_wrapper(),
	154	+ self._make_fallback_kernel(),
	155	+ )
	156	+ mock_patch.assert_not_called()
	157	+
	158	+ mock_patch.reset_mock()
	159	+ # _dvm_generate_fallback_kernel pops fused_id from _fused_metas; re-seed
	160	+ # before exercising the static-shape branch in the same test.
	161	+ _fused_metas[0] = meta
	162	+ with mock.patch(
	163	+ "torch_npu._inductor.dvm.graph_fusion.is_fx_dynamic", return_value=False
	164	+ ):
	165	+ _dvm_generate_fallback_kernel(
	166	+ self._make_codegen_wrapper(),
	167	+ self._make_fallback_kernel(),
	168	+ )
	169	+ mock_patch.assert_called_once_with(meta.gm, ["buf0"])
	170	+ finally:
	171	+ _fused_metas.pop(0, None)
	172	+
	173	+
	174	+if __name__ == "__main__":
	175	+ run_tests()

Atest/_inductor/test_dvm_kernel.py+39-0

@@ -0,0 +1,39 @@
	1	+import torch
	2	+
	3	+from torch.testing._internal.common_utils import TestCase
	4	+from torch.testing._internal.common_utils import run_tests
	5	+
	6	+from torch_npu._inductor import dvm
	7	+
	8	+
	9	+def fused_add_sum(k: dvm.Kernel):
	10	+ x = k.load([-1, -1, -1], dvm.float32)
	11	+ y = k.load([-1, -1, -1], dvm.float32)
	12	+ scalar = k.scalar(dvm.float32)
	13	+ a = k.add(x, y)
	14	+ b = k.add(a, scalar)
	15	+ c = k.sum(b, [0, 1], True)
	16	+ k.store(c)
	17	+
	18	+
	19	+class TestDvmKernelOp(TestCase):
	20	+ def test_dvm_kernel_op(self):
	21	+ a = torch.normal(0, 0.1, size=(512, 128, 256), dtype=torch.float32).npu()
	22	+ b = torch.normal(0, 0.1, size=(512, 1, 256), dtype=torch.float32).npu()
	23	+ scalar = 1.22
	24	+ expect = torch.sum((a + b + 1.22), dim=[0, 1], keepdim=True)
	25	+ result = torch.empty((1, 1, 256), device="npu")
	26	+ kernel1 = dvm.kernel(ktype="vector", dyn_shape=True)(fused_add_sum)
	27	+ kernel1.run(a, b, scalar, result)
	28	+ kernel1.run(a, b, scalar, result)
	29	+ kernel1.run(a, b, scalar, result)
	30	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	31	+ kernel2 = dvm.kernel(ktype="split", dyn_shape=True)(fused_add_sum)
	32	+ result = kernel2(a, b, scalar)
	33	+ result = kernel2(a, b, scalar)
	34	+ result = kernel2(a, b, scalar)
	35	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	36	+
	37	+
	38	+if __name__ == "__main__":
	39	+ run_tests()

Atest/_inductor/test_dvm_mlir_fusion.py+321-0

@@ -0,0 +1,321 @@
	1	+# Owner(s): ["module: tests"]
	2	+import os
	3	+
	4	+import torch
	5	+from torch._inductor.utils import run_and_get_code
	6	+from torch.testing._internal.common_utils import (
	7	+ instantiate_parametrized_tests,
	8	+ parametrize,
	9	+ run_tests,
	10	+ TestCase,
	11	+)
	12	+
	13	+
	14	+class TestModule(torch.nn.Module):
	15	+ def __init__(self, args, *kwargs):
	16	+ super().__init__(args, *kwargs)
	17	+
	18	+ def forward(self, a, b, c):
	19	+ b = torch.transpose(b, 0, 1)
	20	+ add = a + b
	21	+ sub = c - a
	22	+ mul = add * sub
	23	+ mul = mul + 3
	24	+ return mul, torch.sum(mul, dim=[0, 2], keepdim=True)
	25	+
	26	+
	27	+class ReduceCaseModel(torch.nn.Module):
	28	+ def __init__(self):
	29	+ super().__init__()
	30	+
	31	+ def forward(self, arg0_1, arg1_1, arg2_1):
	32	+ sum_1 = torch.ops.aten.sum.dim_IntList(arg0_1, [0, 2, 3], True)
	33	+ div = torch.ops.aten.div.Scalar(sum_1, 9800.0)
	34	+ view = torch.ops.aten.reshape.default(div, [64])
	35	+ mul = torch.ops.aten.mul.Scalar(view, 0.1)
	36	+ mul_1 = torch.ops.aten.mul.Scalar(arg1_1, 0.9)
	37	+ add = torch.ops.aten.add.Tensor(mul, mul_1)
	38	+ expand = torch.ops.aten.expand.default(div, [8, 64, 35, 35])
	39	+ sub = torch.ops.aten.sub.Tensor(arg0_1, expand)
	40	+ pow_1 = torch.ops.aten.pow.Tensor_Scalar(sub, 2)
	41	+ sum_2 = torch.ops.aten.sum.dim_IntList(pow_1, [0, 2, 3], True)
	42	+ div_1 = torch.ops.aten.div.Scalar(sum_2, 9800.0)
	43	+ add_1 = torch.ops.aten.add.Scalar(div_1, 0.001)
	44	+ rsqrt = torch.ops.aten.rsqrt.default(add_1)
	45	+ view_1 = torch.ops.aten.reshape.default(div_1, [64])
	46	+ mul_2 = torch.ops.aten.mul.Scalar(view_1, 1.0001020512297174)
	47	+ mul_3 = torch.ops.aten.mul.Scalar(mul_2, 0.1)
	48	+ mul_4 = torch.ops.aten.mul.Scalar(arg2_1, 0.9)
	49	+ add_2 = torch.ops.aten.add.Tensor(mul_3, mul_4)
	50	+ return (div, add, rsqrt, add_2)
	51	+
	52	+
	53	+class DeterministicReduceModel(torch.nn.Module):
	54	+ def __init__(self):
	55	+ super().__init__()
	56	+
	57	+ def forward(self, arg0):
	58	+ return torch.ops.aten.sum.default(arg0)
	59	+
	60	+
	61	+class BitwiseBoolModel(torch.nn.Module):
	62	+ def forward(self, arg0, arg1):
	63	+ bitwise_and = torch.ops.aten.bitwise_and.Tensor(arg0, arg1)
	64	+ bitwise_not = torch.ops.aten.bitwise_not.default(arg0)
	65	+ return torch.ops.aten.bitwise_or.Tensor(bitwise_and, bitwise_not)
	66	+
	67	+
	68	+class BitwiseIntModel(torch.nn.Module):
	69	+ def forward(self, arg0):
	70	+ return torch.ops.aten.bitwise_not.default(arg0)
	71	+
	72	+
	73	+class MmTransposeBackwardModel(torch.nn.Module):
	74	+ def forward(self, a, b):
	75	+ loss = torch.mm(a, b.t()).sum()
	76	+ grad_a, grad_b = torch.autograd.grad(loss, (a, b))
	77	+ return loss, grad_a, grad_b
	78	+
	79	+
	80	+class CopyInplaceModel(torch.nn.Module):
	81	+ def forward(self, dst, src):
	82	+ add = torch.ops.aten.add.Tensor(src, 1.0)
	83	+ torch.ops.aten.copy_.default(dst, add)
	84	+ return ()
	85	+
	86	+
	87	+class Int64AddModel(torch.nn.Module):
	88	+ def forward(self, x, y):
	89	+ return x + y
	90	+
	91	+
	92	+class Int64CompareModel(torch.nn.Module):
	93	+ def forward(self, x, y):
	94	+ return x > y
	95	+
	96	+
	97	+class Int64PointwiseFusionModel(torch.nn.Module):
	98	+ def forward(self, x, y, ids):
	99	+ sum_xy = x + y
	100	+ mask = sum_xy > y
	101	+ pointwise = torch.where(mask, sum_xy, x)
	102	+ dense_mask = torch.unsqueeze(
	103	+ torch.where(
	104	+ ids >= 0,
	105	+ torch.ones_like(ids, dtype=torch.float32),
	106	+ torch.zeros_like(ids),
	107	+ ),
	108	+ dim=-1,
	109	+ )
	110	+ ids = torch.where(ids == -1, torch.zeros_like(ids), ids)
	111	+ return pointwise, dense_mask, ids
	112	+
	113	+
	114	+class TestDvmByMlir(TestCase):
	115	+ def _run_and_get_code_with_dvm(self, model, *args):
	116	+ original_backend = os.environ.get("TORCHINDUCTOR_NPU_BACKEND")
	117	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
	118	+ try:
	119	+ compiled_model = torch.compile(model, backend="inductor", dynamic=False)
	120	+ return run_and_get_code(compiled_model, *args)
	121	+ finally:
	122	+ if original_backend is None:
	123	+ os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None)
	124	+ else:
	125	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = original_backend
	126	+
	127	+ def test_int64_add_fuses_into_dvm(self):
	128	+ arg0 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	129	+ arg1 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	130	+ model = Int64AddModel()
	131	+
	132	+ with torch.no_grad():
	133	+ expect = model(arg0, arg1)
	134	+ result, codes = self._run_and_get_code_with_dvm(model, arg0, arg1)
	135	+
	136	+ code = "\n".join(codes)
	137	+ self.assertEqual(expect, result)
	138	+ self.assertIn("k.add", code)
	139	+
	140	+ def test_int64_compare_fuses_into_dvm(self):
	141	+ arg0 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	142	+ arg1 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	143	+ model = Int64CompareModel()
	144	+
	145	+ with torch.no_grad():
	146	+ expect = model(arg0, arg1)
	147	+ result, codes = self._run_and_get_code_with_dvm(model, arg0, arg1)
	148	+
	149	+ code = "\n".join(codes)
	150	+ self.assertEqual(expect, result)
	151	+ self.assertIn("k.greater", code)
	152	+
	153	+ def test_int64_pointwise_chain_fuses_into_dvm(self):
	154	+ arg0 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	155	+ arg1 = torch.randint(-8, 8, (32, 32), dtype=torch.int64, device="npu")
	156	+ ids = torch.randint(-2, 4, (32, 32), dtype=torch.int64, device="npu")
	157	+ model = Int64PointwiseFusionModel()
	158	+
	159	+ with torch.no_grad():
	160	+ expect = model(arg0, arg1, ids)
	161	+ result, codes = self._run_and_get_code_with_dvm(model, arg0, arg1, ids)
	162	+
	163	+ code = "\n".join(codes)
	164	+ self.assertEqual(expect, result)
	165	+ self.assertIn("k.add", code)
	166	+ self.assertIn("k.greater", code)
	167	+ self.assertIn("k.greater_equal", code)
	168	+ self.assertIn("k.equal", code)
	169	+ self.assertEqual(code.count("k.select("), 3)
	170	+
	171	+
	172	+ @parametrize("dtype", [torch.float16, torch.float32, torch.bfloat16])
	173	+ @parametrize("is_dynamic", [True, False])
	174	+ def test_basic_partitioning(self, dtype, is_dynamic):
	175	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
	176	+ a = torch.normal(0, 0.01, size=(512, 1), dtype=dtype).npu()
	177	+ b = torch.normal(0, 0.01, size=(512, 4, 256), dtype=dtype).npu()
	178	+ c = torch.normal(0, 0.01, size=(1, 256), dtype=dtype).npu()
	179	+ model = TestModule()
	180	+ dvm_compiled_model = torch.compile(
	181	+ model, backend="inductor", dynamic=is_dynamic
	182	+ )
	183	+ with torch.no_grad():
	184	+ expect = model(a, b, c)
	185	+ result = dvm_compiled_model(a, b, c)
	186	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	187	+ del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
	188	+
	189	+ @parametrize("dtype", [torch.bfloat16])
	190	+ @parametrize("is_dynamic", [False])
	191	+ def test_basic_partitioning_npugraph(self, dtype, is_dynamic):
	192	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
	193	+ a = torch.normal(0, 0.01, size=(512, 1), dtype=dtype).npu()
	194	+ b = torch.normal(0, 0.01, size=(512, 4, 256), dtype=dtype).npu()
	195	+ c = torch.normal(0, 0.01, size=(1, 256), dtype=dtype).npu()
	196	+ model = TestModule()
	197	+ dvm_compiled_model = torch.compile(
	198	+ model,
	199	+ backend="inductor",
	200	+ dynamic=is_dynamic,
	201	+ options={"triton.cudagraphs": True},
	202	+ )
	203	+ with torch.no_grad():
	204	+ expect = model(a, b, c)
	205	+ result = dvm_compiled_model(a, b, c)
	206	+ result = dvm_compiled_model(a, b, c)
	207	+ result = dvm_compiled_model(a, b, c)
	208	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	209	+ del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
	210	+
	211	+ @parametrize("dtype", [torch.float16, torch.float32])
	212	+ @parametrize("is_dynamic", [True, False])
	213	+ def test_reduce_case(self, dtype, is_dynamic):
	214	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
	215	+ arg0 = torch.empty_strided(
	216	+ torch.Size((8, 64, 35, 35)),
	217	+ (78400, 1225, 35, 1),
	218	+ dtype=dtype,
	219	+ device="npu",
	220	+ ).uniform_(0, 1)
	221	+ arg1 = torch.empty_strided(
	222	+ torch.Size((64,)), (1,), dtype=dtype, device="npu"
	223	+ ).uniform_(0, 1)
	224	+ arg2 = torch.empty_strided(
	225	+ torch.Size((64,)), (1,), dtype=dtype, device="npu"
	226	+ ).uniform_(0, 1)
	227	+ model = ReduceCaseModel()
	228	+ dvm_compiled_model = torch.compile(
	229	+ model, backend="inductor", dynamic=is_dynamic
	230	+ )
	231	+ with torch.no_grad():
	232	+ expect = model(arg0, arg1, arg2)
	233	+ result = dvm_compiled_model(arg0, arg1, arg2)
	234	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	235	+ del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
	236	+
	237	+ def test_deterministic_reduce_case(self):
	238	+ os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"
	239	+ deterministic_state = torch.are_deterministic_algorithms_enabled()
	240	+ deterministic_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
	241	+ arg0 = torch.normal(
	242	+ 0, 0.1, size=(16, 128, 64, 64), dtype=torch.float32, device="npu"
	243	+ )
	244	+ model = DeterministicReduceModel()
	245	+ try:
	246	+ torch.use_deterministic_algorithms(True)
	247	+ dvm_compiled_model = torch.compile(
	248	+ model, backend="inductor", dynamic=False
	249	+ )
	250	+ with torch.no_grad():
	251	+ first_result = dvm_compiled_model(arg0)
	252	+ second_result = dvm_compiled_model(arg0)
	253	+ self.assertEqual(first_result, second_result, atol=0, rtol=0)
	254	+ finally:
	255	+ torch.use_deterministic_algorithms(
	256	+ deterministic_state, warn_only=deterministic_warn_only
	257	+ )
	258	+ del os.environ["TORCHINDUCTOR_NPU_BACKEND"]
	259	+
	260	+ def test_bitwise_bool_ops_codegen(self):
	261	+ arg0 = torch.randint(0, 2, (32, 32), dtype=torch.bool, device="npu")
	262	+ arg1 = torch.randint(0, 2, (32, 32), dtype=torch.bool, device="npu")
	263	+ model = BitwiseBoolModel()
	264	+
	265	+ with torch.no_grad():
	266	+ expect = model(arg0, arg1)
	267	+ result, codes = self._run_and_get_code_with_dvm(model, arg0, arg1)
	268	+
	269	+ code = "\n".join(codes)
	270	+ self.assertEqual(expect, result)
	271	+ self.assertIn("k.logical_and", code)
	272	+ self.assertIn("k.logical_or", code)
	273	+ self.assertIn("k.logical_not", code)
	274	+
	275	+ def test_bitwise_int_rule_fallback(self):
	276	+ arg0 = torch.randint(-8, 8, (32, 32), dtype=torch.int32, device="npu")
	277	+ model = BitwiseIntModel()
	278	+
	279	+ with torch.no_grad():
	280	+ expect = model(arg0)
	281	+ result, codes = self._run_and_get_code_with_dvm(model, arg0)
	282	+
	283	+ code = "\n".join(codes)
	284	+ self.assertEqual(expect, result)
	285	+ self.assertNotIn("k.logical_not", code)
	286	+
	287	+ def test_mm_t_backward_no_dvm_fused_matmul_backward(self):
	288	+ a = torch.randn((4, 8), dtype=torch.float32, device="npu")
	289	+ b = torch.randn((3, 8), dtype=torch.float32, device="npu")
	290	+ a_eager = a.detach().clone().requires_grad_(True)
	291	+ b_eager = b.detach().clone().requires_grad_(True)
	292	+ a_compiled = a.detach().clone().requires_grad_(True)
	293	+ b_compiled = b.detach().clone().requires_grad_(True)
	294	+ model = MmTransposeBackwardModel()
	295	+
	296	+ expect = model(a_eager, b_eager)
	297	+ result, codes = self._run_and_get_code_with_dvm(
	298	+ model, a_compiled, b_compiled
	299	+ )
	300	+
	301	+ code = "\n".join(codes)
	302	+ self.assertEqual(expect, result, atol=1e-3, rtol=1e-3)
	303	+ self.assertNotIn("dvm_fused_matmul_backward", code)
	304	+
	305	+ def test_copy_inplace_codegen(self):
	306	+ src = torch.randn((128,), dtype=torch.float32, device="npu")
	307	+ dst = torch.zeros((128,), dtype=torch.float32, device="npu")
	308	+ expect_dst = dst.clone()
	309	+ actual_dst = dst.clone()
	310	+ model = CopyInplaceModel()
	311	+
	312	+ with torch.no_grad():
	313	+ model(expect_dst, src)
	314	+ self._run_and_get_code_with_dvm(model, actual_dst, src)
	315	+
	316	+ self.assertEqual(expect_dst, actual_dst)
	317	+
	318	+
	319	+instantiate_parametrized_tests(TestDvmByMlir)
	320	+if __name__ == "__main__":
	321	+ run_tests()

Mtest/npu/test_public_bindings.py+8-0

	@@ -666,6 +666,14 @@ class TestPublicBindings(TestCase):
666	"torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.ir",	666	"torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.ir",
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",
		670	+ "torch_npu._inductor.dvm.decomp",
		671	+ "torch_npu._inductor.dvm.fx_pass",
		672	+ "torch_npu._inductor.dvm.fx_test",
		673	+ "torch_npu._inductor.dvm.graph_build",
		674	+ "torch_npu._inductor.dvm.graph_fusion",
		675	+ "torch_npu._inductor.dvm.mlir_fusion",
		676	+ "torch_npu._inductor.dvm.op_emitter",
669	}	677	}
670		678
671	# No new entries should be added to this list.	679	# No new entries should be added to this list.

Dthird_party/acl/inc/experiment/msprof/toolchain/prof_api.h+0-254

	@@ -1,254 +0,0 @@
1	-/**
2	- * @file prof_api.h
3	- *
4	- * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
5	- *
6	- * This program is distributed in the hope that it will be useful,
7	- * but WITHOUT ANY WARRANTY; without even the implied warranty of
8	- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
9	- *
10	- */
11	-
12	-#ifndef PROF_API_H
13	-#define PROF_API_H
14	-
15	-#include <stddef.h>
16	-#include <stdint.h>
17	-#include <stdbool.h>
18	-#include "prof_common.h"
19	-
20	-#ifdef __cplusplus
21	-extern "C" {
22	-#endif // __cplusplus
23	-
24	-#if (defined(_WIN32) \|\| defined(_WIN64) \|\| defined(_MSC_VER))
25	-#define MSVP_PROF_API __declspec(dllexport)
26	-#else
27	-#define MSVP_PROF_API __attribute__((visibility("default")))
28	-#endif
29	-
30	-/*
31	- * @ingroup libprofapi
32	- * @name profRegReporterCallback
33	- * @brief register report callback interface for atlas
34	- * @param [in] reporter: reporter callback handle
35	- * @return 0:SUCCESS, !0:FAILED
36	- */
37	-MSVP_PROF_API int32_t profRegReporterCallback(MsprofReportHandle reporter);
38	-
39	-/*
40	- * @ingroup libprofapi
41	- * @name profRegCtrlCallback
42	- * @brief register control callback, interface for atlas
43	- * @param [in] handle: control callback handle
44	- * @return 0:SUCCESS, !0:FAILED
45	- */
46	-MSVP_PROF_API int32_t profRegCtrlCallback(MsprofCtrlHandle handle);
47	-
48	-/*
49	- * @ingroup libprofapi
50	- * @name profRegDeviceStateCallback
51	- * @brief register device state notify callback, interface for atlas
52	- * @param [in] handle: handle of ProfNotifySetDevice
53	- * @return 0:SUCCESS, !0:FAILED
54	- */
55	-MSVP_PROF_API int32_t profRegDeviceStateCallback(MsprofSetDeviceHandle handle);
56	-
57	-/*
58	- * @ingroup libprofapi
59	- * @name profGetDeviceIdByGeModelIdx
60	- * @brief get device id by model id, interface for atlas
61	- * @param [in] modelIdx: ge model id
62	- * @param [out] deviceId: device id
63	- * @return 0:SUCCESS, !0:FAILED
64	- */
65	-MSVP_PROF_API int32_t profGetDeviceIdByGeModelIdx(const uint32_t modelIdx, uint32_t *deviceId);
66	-
67	-/*
68	- * @ingroup libprofapi
69	- * @name profSetProfCommand
70	- * @brief register set profiling command, interface for atlas
71	- * @param [in] command: 0 isn't aging, !0 is aging
72	- * @param [in] len: api of timestamp data
73	- * @return 0:SUCCESS, !0:FAILED
74	- */
75	-MSVP_PROF_API int32_t profSetProfCommand(VOID_PTR command, uint32_t len);
76	-
77	-/*
78	- * @ingroup libprofapi
79	- * @name profSetStepInfo
80	- * @brief set step info for torch, interface for atlas
81	- * @param [in] indexId: id of iteration index
82	- * @param [in] tagId: id of tag
83	- * @param [in] stream: api of timestamp data
84	- * @return 0:SUCCESS, !0:FAILED
85	- */
86	-MSVP_PROF_API int32_t profSetStepInfo(const uint64_t indexId, const uint16_t tagId, void* const stream);
87	-
88	-/*
89	- * @ingroup libprofapi
90	- * @name MsprofRegisterProfileCallback
91	- * @brief register profile callback by callback type, interface for atlas
92	- * @param [in] callbackType: type of callback(reporter/ctrl/device state/command)
93	- * @param [in] callback: callback of profile
94	- * @param [in] len: callback length
95	- * @return 0:SUCCESS, !0:FAILED
96	- */
97	-MSVP_PROF_API int32_t MsprofRegisterProfileCallback(int32_t callbackType, VOID_PTR callback, uint32_t len);
98	-
99	-/**
100	- * @ingroup libprofapi
101	- * @name MsprofInit
102	- * @brief Profiling module init
103	- * @param [in] dataType: profiling type: ACL Env/ACL Json/GE Option
104	- * @param [in] data: profiling switch data
105	- * @param [in] dataLen: Length of data
106	- * @return 0:SUCCESS, >0:FAILED
107	- */
108	-MSVP_PROF_API int32_t MsprofInit(uint32_t dataType, VOID_PTR data, uint32_t dataLen);
109	-
110	-/**
111	- * @ingroup libprofapi
112	- * @name MsprofSetConfig
113	- * @brief Set profiling config
114	- * @return 0:SUCCESS, !0:FAILED
115	- */
116	-MSVP_PROF_API int32_t MsprofSetConfig(uint32_t configType, const char *config, size_t configLength);
117	-
118	-/**
119	- * @ingroup libprofapi
120	- * @name MsprofRegisterCallback
121	- * @brief register profiling switch callback for module
122	- * @param [in] agingFlag: 0 isn't aging, !0 is aging
123	- * @param [in] api: api of timestamp data
124	- * @return 0:SUCCESS, !0:FAILED
125	- */
126	-MSVP_PROF_API int32_t MsprofRegisterCallback(uint32_t moduleId, ProfCommandHandle handle);
127	-
128	-/**
129	- * @ingroup libprofapi
130	- * @name MsprofReportData
131	- * @brief report profiling data of module
132	- * @param [in] moduleId: module id
133	- * @param [in] type: report type(init/uninit/max length/hash)
134	- * @param [in] data: profiling data
135	- * @param [in] len: length of profiling data
136	- * @return 0:SUCCESS, !0:FAILED
137	- */
138	-MSVP_PROF_API int32_t MsprofReportData(uint32_t moduleId, uint32_t type, VOID_PTR data, uint32_t len);
139	-
140	-/*
141	- * @ingroup libprofapi
142	- * @name MsprofReportApi
143	- * @brief report api timestamp
144	- * @param [in] agingFlag: 0 isn't aging, !0 is aging
145	- * @param [in] api: api of timestamp data
146	- * @return 0:SUCCESS, !0:FAILED
147	- */
148	-MSVP_PROF_API int32_t MsprofReportApi(uint32_t agingFlag, const struct MsprofApi *api);
149	-
150	-/*
151	- * @ingroup libprofapi
152	- * @name MsprofReportEvent
153	- * @brief report event timestamp
154	- * @param [in] agingFlag: 0 isn't aging, !0 is aging
155	- * @param [in] event: event of timestamp data
156	- * @return 0:SUCCESS, !0:FAILED
157	- */
158	-MSVP_PROF_API int32_t MsprofReportEvent(uint32_t agingFlag, const struct MsprofEvent *event);
159	-
160	-/*
161	- * @ingroup libprofapi
162	- * @name MsprofReportCompactInfo
163	- * @brief report profiling compact infomation
164	- * @param [in] agingFlag: 0 isn't aging, !0 is aging
165	- * @param [in] data: profiling data of compact infomation
166	- * @param [in] length: length of profiling data
167	- * @return 0:SUCCESS, !0:FAILED
168	- */
169	-MSVP_PROF_API int32_t MsprofReportCompactInfo(uint32_t agingFlag, const VOID_PTR data, uint32_t length);
170	-
171	-/*
172	- * @ingroup libprofapi
173	- * @name MsprofReportAdditionalInfo
174	- * @brief report profiling additional infomation
175	- * @param [in] agingFlag: 0 isn't aging, !0 is aging
176	- * @param [in] data: profiling data of additional infomation
177	- * @param [in] length: length of profiling data
178	- * @return 0:SUCCESS, !0:FAILED
179	- */
180	-MSVP_PROF_API int32_t MsprofReportAdditionalInfo(uint32_t agingFlag, const VOID_PTR data, uint32_t length);
181	-
182	-/*
183	- * @ingroup libprofapi
184	- * @name MsprofRegTypeInfo
185	- * @brief reg mapping info of type id and type name
186	- * @param [in] level: level is the report struct's level
187	- * @param [in] typeId: type id is the report struct's type
188	- * @param [in] typeName: label of type id for presenting user
189	- * @return 0:SUCCESS, !0:FAILED
190	- */
191	-MSVP_PROF_API int32_t MsprofRegTypeInfo(uint16_t level, uint32_t typeId, const char *typeName);
192	-
193	-/*
194	- * @ingroup libprofapi
195	- * @name MsprofGetHashId
196	- * @brief return hash id of hash info
197	- * @param [in] hashInfo: infomation to be hashed
198	- * @param [in] length: the length of infomation to be hashed
199	- * @return hash id
200	- */
201	-MSVP_PROF_API uint64_t MsprofGetHashId(const char *hashInfo, size_t length);
202	-
203	-/**
204	- * @ingroup libprofapi
205	- * @name MsprofSetDeviceIdByGeModelIdx
206	- * @brief insert device id by model id
207	- * @param [in] geModelIdx: ge model id
208	- * @param [in] deviceId: device id
209	- * @return 0:SUCCESS, !0:FAILED
210	- */
211	-MSVP_PROF_API int32_t MsprofSetDeviceIdByGeModelIdx(const uint32_t geModelIdx, const uint32_t deviceId);
212	-
213	-/**
214	- * @ingroup libprofapi
215	- * @name MsprofUnsetDeviceIdByGeModelIdx
216	- * @brief delete device id by model id
217	- * @param [in] geModelIdx: ge model id
218	- * @param [in] deviceId: device id
219	- * @return 0:SUCCESS, !0:FAILED
220	- */
221	-MSVP_PROF_API int32_t MsprofUnsetDeviceIdByGeModelIdx(const uint32_t geModelIdx, const uint32_t deviceId);
222	-
223	-/**
224	- * @ingroup libprofapi
225	- * @name register report interface for atlas
226	- * @brief report api timestamp
227	- * @param [in] chipId: multi die's chip
228	- * @param [in] deviceId: device id
229	- * @param [in] isOpen: device is open
230	- * @return 0:SUCCESS, !0:FAILED
231	- */
232	-MSVP_PROF_API int32_t MsprofNotifySetDevice(uint32_t chipId, uint32_t deviceId, bool isOpen);
233	-
234	-/**
235	- * @ingroup libprofapi
236	- * @name MsprofFinalize
237	- * @brief profiling finalize
238	- * @return 0:SUCCESS, !0:FAILED
239	- */
240	-MSVP_PROF_API int32_t MsprofFinalize();
241	-
242	-/*
243	- * @ingroup libprofapi
244	- * @name MsprofSysCycleTime
245	- * @brief get systime cycle time of CPU
246	- * @return system cycle time of CPU
247	- */
248	-MSVP_PROF_API uint64_t MsprofSysCycleTime();
249	-
250	-#ifdef __cplusplus
251	-}
252	-#endif
253	-
254	-#endif

Dthird_party/acl/inc/experiment/msprof/toolchain/prof_common.h+0-1101

	@@ -1,1101 +0,0 @@
1	-/**
2	- * @file prof_common.h
3	- *
4	- * Copyright (c) Huawei Technologies Co., Ltd. 2019-2024. All rights reserved.
5	- *
6	- * This program is distributed in the hope that it will be useful,
7	- * but WITHOUT ANY WARRANTY; without even the implied warranty of
8	- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
9	- *
10	- */
11	-#ifndef MSPROFILER_PROF_COMMON_H
12	-#define MSPROFILER_PROF_COMMON_H
13	-
14	-#include <stdint.h>
15	-#include <stddef.h>
16	-
17	-#ifdef __cplusplus
18	-extern "C" {
19	-#endif // __cplusplus
20	-
21	-#define MSPROF_DATA_HEAD_MAGIC_NUM 0x5A5AU
22	-#define MSPROF_REPORT_DATA_MAGIC_NUM 0x5A5AU
23	-#define MSPROF_TASK_TIME_L0 0x00000800ULL // mean PROF_TASK_TIME
24	-#define MSPROF_EVENT_FLAG 0xFFFFFFFFFFFFFFFFULL
25	-typedef void* VOID_PTR;
26	-typedef const void* ConstVoidPtr;
27	-typedef int32_t (*MsprofReportHandle)(uint32_t moduleId, uint32_t type, VOID_PTR data, uint32_t len);
28	-typedef int32_t (*MsprofCtrlHandle)(uint32_t type, VOID_PTR data, uint32_t len);
29	-typedef int32_t (*MsprofSetDeviceHandle)(VOID_PTR data, uint32_t len);
30	-typedef int32_t (MsprofCtrlCallback)(uint32_t type, void data, uint32_t len);
31	-typedef int32_t (MsprofReporterCallback)(uint32_t moduleId, uint32_t type, void data, uint32_t len);
32	-
33	-/**
34	- * @name ProfCommandHandle
35	- * @brief callback to start/stop profiling
36	- * @param type [IN] enum call back type
37	- * @param data [IN] callback data
38	- * @param len [IN] callback data size
39	- * @return enum MsprofErrorCode
40	- */
41	-typedef int32_t (*ProfCommandHandle)(uint32_t type, VOID_PTR data, uint32_t len);
42	-
43	-/* Msprof report level */
44	-#define MSPROF_REPORT_PYTORCH_LEVEL 30000U
45	-#define MSPROF_REPORT_PTA_LEVEL 25000U
46	-#define MSPROF_REPORT_ACL_LEVEL 20000U
47	-#define MSPROF_REPORT_MODEL_LEVEL 15000U
48	-#define MSPROF_REPORT_NODE_LEVEL 10000U
49	-#define MSPROF_REPORT_AICPU_LEVEL 6000U
50	-#define MSPROF_REPORT_HCCL_NODE_LEVEL 5500U
51	-#define MSPROF_REPORT_RUNTIME_LEVEL 5000U
52	-#define MSPROF_REPORT_PROF_LEVEL 4500U
53	-#define MSPROF_REPORT_DPU_LEVEL 4000U
54	-
55	-/* Msprof report type of acl(20000) level(acl), offset: 0x000000 */
56	-#define MSPROF_REPORT_ACL_OP_BASE_TYPE 0x010000U
57	-#define MSPROF_REPORT_ACL_MODEL_BASE_TYPE 0x020000U
58	-#define MSPROF_REPORT_ACL_RUNTIME_BASE_TYPE 0x030000U
59	-#define MSPROF_REPORT_ACL_OTHERS_BASE_TYPE 0x040000U
60	-
61	-
62	-/* Msprof report type of acl(20000) level(host api), offset: 0x050000 */
63	-#define MSPROF_REPORT_ACL_NN_BASE_TYPE 0x050000U
64	-#define MSPROF_REPORT_ACL_ASCENDC_TYPE 0x060000U
65	-#define MSPROF_REPORT_ACL_HOST_HCCL_BASE_TYPE 0x070000U
66	-#define MSPROF_REPORT_ACL_DVPP_BASE_TYPE 0x090000U
67	-#define MSPROF_REPORT_ACL_GRAPH_BASE_TYPE 0x0A0000U
68	-
69	-/* Msprof report type of model(15000) level, offset: 0x000000 */
70	-#define MSPROF_REPORT_MODEL_GRAPH_ID_MAP_TYPE 0U /* type info: graph_id_map */
71	-#define MSPROF_REPORT_MODEL_EXECUTE_TYPE 1U /* type info: execute */
72	-#define MSPROF_REPORT_MODEL_LOAD_TYPE 2U /* type info: load */
73	-#define MSPROF_REPORT_MODEL_INPUT_COPY_TYPE 3U /* type info: IntputCopy */
74	-#define MSPROF_REPORT_MODEL_OUTPUT_COPY_TYPE 4U /* type info: OutputCopy */
75	-#define MSPROF_REPORT_MODEL_LOGIC_STREAM_TYPE 7U /* type info: logic_stream_info */
76	-#define MSPROF_REPORT_MODEL_EXEOM_TYPE 8U /* type info: exeom */
77	-#define MSPROF_REPORT_MODEL_UDF_BASE_TYPE 0x010000U /* type info: udf_info */
78	-#define MSPROF_REPORT_MODEL_AICPU_BASE_TYPE 0x020000U /* type info: aicpu */
79	-
80	-/* Msprof report type of node(10000) level, offset: 0x000000 */
81	-#define MSPROF_REPORT_NODE_BASIC_INFO_TYPE 0U /* type info: node_basic_info */
82	-#define MSPROF_REPORT_NODE_TENSOR_INFO_TYPE 1U /* type info: tensor_info */
83	-#define MSPROF_REPORT_NODE_FUSION_OP_INFO_TYPE 2U /* type info: funsion_op_info */
84	-#define MSPROF_REPORT_NODE_CONTEXT_ID_INFO_TYPE 4U /* type info: context_id_info */
85	-#define MSPROF_REPORT_NODE_LAUNCH_TYPE 5U /* type info: launch */
86	-#define MSPROF_REPORT_NODE_TASK_MEMORY_TYPE 6U /* type info: task_memory_info */
87	-#define MSPROF_REPORT_NODE_HOST_OP_EXEC_TYPE 8U /* type info: op exec */
88	-#define MSPROF_REPORT_NODE_ATTR_INFO_TYPE 9U /* type info: node_attr_info */
89	-#define MSPROF_REPORT_NODE_HCCL_OP_INFO_TYPE 10U /* type info: hccl_op_info */
90	-#define MSPROF_REPORT_NODE_STATIC_OP_MEM_TYPE 11U /* type info: static_op_mem */
91	-#define MSPROF_REPORT_NODE_MC2_COMMINFO_TYPE 12U /* type info: mc2_comm_info */
92	-
93	-/* Msprof report type of node(10000) level(ge api), offset: 0x010000 */
94	-#define MSPROF_REPORT_NODE_GE_API_BASE_TYPE 0x010000U /* type info: ge api */
95	-#define MSPROF_REPORT_NODE_HCCL_BASE_TYPE 0x020000U /* type info: hccl api */
96	-#define MSPROF_REPORT_NODE_DVPP_API_BASE_TYPE 0x030000U /* type info: dvpp api */
97	-/* Msprof report type of aicpu(6000), offset: 0x000000 */
98	-#define MSPROF_REPORT_AICPU_NODE_TYPE 0U /* type info: DATA_PREPROCESS.AICPU */
99	-#define MSPROF_REPORT_AICPU_DP_TYPE 1U /* type info: DATA_PREPROCESS.DP */
100	-#define MSPROF_REPORT_AICPU_MODEL_TYPE 2U /* type info: DATA_PREPROCESS.AICPU_MODEL */
101	-#define MSPROF_REPORT_AICPU_MI_TYPE 3U /* type info: DATA_PREPROCESS.AICPUMI */
102	-#define MSPROF_REPORT_AICPU_MC2_EXECUTE_COMM_TIME 4U /* 通信时刻信息 */
103	-#define MSPROF_REPORT_AICPU_MC2_EXECUTE_COMP_TIME 5U /* 计算时刻信息 */
104	-#define MSPROF_REPORT_AICPU_MC2_HCCL_INFO 6U /* task信息 */
105	-
106	-/* Msprof report type of hccl(5500) level(op api), offset: 0x010000 */
107	-#define MSPROF_REPORT_HCCL_NODE_BASE_TYPE 0x010000U
108	-#define MSPROF_REPORT_HCCL_MASTER_TYPE 0x010001U
109	-#define MSPROF_REPORT_HCCL_SLAVE_TYPE 0x010002U
110	-
111	-/* Msprof report type of hccl(4000U) level(dpu), offset: 0x000000 */
112	-#define MSPROF_REPORT_DPU_TRACK_TYPE 0U /* type info: dpu_track */
113	-
114	-/* use with AdprofCheckFeatureIsOn */
115	-#define ADPROF_TASK_TIME_L0 0x00000008ULL
116	-#define ADPROF_TASK_TIME_L1 0x00000010ULL
117	-#define ADPROF_TASK_TIME_L2 0x00000020ULL
118	-
119	-/* Msprof report type of profiling(4500) */
120	-#define MSPROF_REPORT_DIAGNOSTIC_INFO_TYPE 0x010000U
121	-
122	-enum ProfileCallbackType {
123	- PROFILE_CTRL_CALLBACK = 0,
124	- PROFILE_DEVICE_STATE_CALLBACK,
125	- PROFILE_REPORT_API_CALLBACK,
126	- PROFILE_REPORT_EVENT_CALLBACK,
127	- PROFILE_REPORT_COMPACT_CALLBACK,
128	- PROFILE_REPORT_ADDITIONAL_CALLBACK,
129	- PROFILE_REPORT_REG_TYPE_INFO_CALLBACK,
130	- PROFILE_REPORT_GET_HASH_ID_CALLBACK,
131	- PROFILE_HOST_FREQ_IS_ENABLE_CALLBACK,
132	- PROFILE_REPORT_API_C_CALLBACK,
133	- PROFILE_REPORT_EVENT_C_CALLBACK,
134	- PROFILE_REPORT_REG_TYPE_INFO_C_CALLBACK,
135	- PROFILE_REPORT_GET_HASH_ID_C_CALLBACK,
136	- PROFILE_HOST_FREQ_IS_ENABLE_C_CALLBACK,
137	-};
138	-
139	-enum MsprofDataTag {
140	- MSPROF_ACL_DATA_TAG = 0, // acl data tag, range: 0~19
141	- MSPROF_GE_DATA_TAG_MODEL_LOAD = 20, // ge data tag, range: 20~39
142	- MSPROF_GE_DATA_TAG_FUSION = 21,
143	- MSPROF_GE_DATA_TAG_INFER = 22,
144	- MSPROF_GE_DATA_TAG_TASK = 23,
145	- MSPROF_GE_DATA_TAG_TENSOR = 24,
146	- MSPROF_GE_DATA_TAG_STEP = 25,
147	- MSPROF_GE_DATA_TAG_ID_MAP = 26,
148	- MSPROF_GE_DATA_TAG_HOST_SCH = 27,
149	- MSPROF_RUNTIME_DATA_TAG_API = 40, // runtime data tag, range: 40~59
150	- MSPROF_RUNTIME_DATA_TAG_TRACK = 41,
151	- MSPROF_AICPU_DATA_TAG = 60, // aicpu data tag, range: 60~79
152	- MSPROF_AICPU_MODEL_TAG = 61,
153	- MSPROF_HCCL_DATA_TAG = 80, // hccl data tag, range: 80~99
154	- MSPROF_DP_DATA_TAG = 100, // dp data tag, range: 100~119
155	- MSPROF_MSPROFTX_DATA_TAG = 120, // hccl data tag, range: 120~139
156	- MSPROF_DATA_TAG_MAX = 65536, // data tag value type is uint16_t
157	-};
158	-
159	-enum MsprofMindsporeNodeTag {
160	- GET_NEXT_DEQUEUE_WAIT = 1,
161	-};
162	-
163	-/**
164	- * @brief struct of mixed data
165	- */
166	-#define MSPROF_MIX_DATA_RESERVE_BYTES 7
167	-#define MSPROF_MIX_DATA_STRING_LEN 120
168	-enum MsprofMixDataType {
169	- MSPROF_MIX_DATA_HASH_ID = 0,
170	- MSPROF_MIX_DATA_STRING,
171	-};
172	-struct MsprofMixData {
173	- uint8_t type; // MsprofMixDataType
174	- uint8_t rsv[MSPROF_MIX_DATA_RESERVE_BYTES];
175	- union {
176	- uint64_t hashId;
177	- char dataStr[MSPROF_MIX_DATA_STRING_LEN];
178	- } data;
179	-};
180	-
181	-#define PATH_LEN_MAX 1023
182	-#define PARAM_LEN_MAX 4095
183	-struct MsprofCommandHandleParams {
184	- uint32_t pathLen;
185	- uint32_t storageLimit; // MB
186	- uint32_t profDataLen;
187	- char path[PATH_LEN_MAX + 1];
188	- char profData[PARAM_LEN_MAX + 1];
189	-};
190	-
191	-/**
192	- * @brief profiling command info
193	- */
194	-#define MSPROF_MAX_DEV_NUM 64
195	-struct MsprofCommandHandle {
196	- uint64_t profSwitch;
197	- uint64_t profSwitchHi;
198	- uint32_t devNums;
199	- uint32_t devIdList[MSPROF_MAX_DEV_NUM];
200	- uint32_t modelId;
201	- uint32_t type;
202	- uint32_t cacheFlag;
203	- struct MsprofCommandHandleParams params;
204	-};
205	-
206	-/**
207	- * @brief struct of data reported by acl
208	- */
209	-#define MSPROF_ACL_DATA_RESERVE_BYTES 32
210	-#define MSPROF_ACL_API_NAME_LEN 64
211	-enum MsprofAclApiType {
212	- MSPROF_ACL_API_TYPE_OP = 1,
213	- MSPROF_ACL_API_TYPE_MODEL,
214	- MSPROF_ACL_API_TYPE_RUNTIME,
215	- MSPROF_ACL_API_TYPE_OTHERS,
216	-};
217	-struct MsprofAclProfData {
218	-#ifdef __cplusplus
219	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
220	- uint16_t dataTag = MSPROF_ACL_DATA_TAG;
221	-#else
222	- uint16_t magicNumber;
223	- uint16_t dataTag;
224	-#endif
225	- uint32_t apiType; // enum MsprofAclApiType
226	- uint64_t beginTime;
227	- uint64_t endTime;
228	- uint32_t processId;
229	- uint32_t threadId;
230	- char apiName[MSPROF_ACL_API_NAME_LEN];
231	- uint8_t reserve[MSPROF_ACL_DATA_RESERVE_BYTES];
232	-};
233	-
234	-/**
235	- * @brief struct of data reported by GE
236	- */
237	-#define MSPROF_GE_MODELLOAD_DATA_RESERVE_BYTES 104
238	-struct MsprofGeProfModelLoadData {
239	-#ifdef __cplusplus
240	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
241	- uint16_t dataTag = MSPROF_GE_DATA_TAG_MODEL_LOAD;
242	-#else
243	- uint16_t magicNumber;
244	- uint16_t dataTag;
245	-#endif
246	- uint32_t modelId;
247	- struct MsprofMixData modelName;
248	- uint64_t startTime;
249	- uint64_t endTime;
250	- uint8_t reserve[MSPROF_GE_MODELLOAD_DATA_RESERVE_BYTES];
251	-};
252	-
253	-#define MSPROF_GE_FUSION_DATA_RESERVE_BYTES 8
254	-#define MSPROF_GE_FUSION_OP_NUM 8
255	-struct MsprofGeProfFusionData {
256	-#ifdef __cplusplus
257	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
258	- uint16_t dataTag = MSPROF_GE_DATA_TAG_FUSION;
259	-#else
260	- uint16_t magicNumber;
261	- uint16_t dataTag;
262	-#endif
263	- uint32_t modelId;
264	- struct MsprofMixData fusionName;
265	- uint64_t inputMemSize;
266	- uint64_t outputMemSize;
267	- uint64_t weightMemSize;
268	- uint64_t workspaceMemSize;
269	- uint64_t totalMemSize;
270	- uint64_t fusionOpNum;
271	- uint64_t fusionOp[MSPROF_GE_FUSION_OP_NUM];
272	- uint8_t reserve[MSPROF_GE_FUSION_DATA_RESERVE_BYTES];
273	-};
274	-
275	-#define MSPROF_GE_INFER_DATA_RESERVE_BYTES 64
276	-struct MsprofGeProfInferData {
277	-#ifdef __cplusplus
278	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
279	- uint16_t dataTag = MSPROF_GE_DATA_TAG_INFER;
280	-#else
281	- uint16_t magicNumber;
282	- uint16_t dataTag;
283	-#endif
284	- uint32_t modelId;
285	- struct MsprofMixData modelName;
286	- uint32_t requestId;
287	- uint32_t threadId;
288	- uint64_t inputDataStartTime;
289	- uint64_t inputDataEndTime;
290	- uint64_t inferStartTime;
291	- uint64_t inferEndTime;
292	- uint64_t outputDataStartTime;
293	- uint64_t outputDataEndTime;
294	- uint8_t reserve[MSPROF_GE_INFER_DATA_RESERVE_BYTES];
295	-};
296	-
297	-#define MSPROF_GE_TASK_DATA_RESERVE_BYTES 12
298	-#define MSPROF_GE_OP_TYPE_LEN 56
299	-enum MsprofGeTaskType {
300	- MSPROF_GE_TASK_TYPE_AI_CORE = 0,
301	- MSPROF_GE_TASK_TYPE_AI_CPU,
302	- MSPROF_GE_TASK_TYPE_AIV,
303	- MSPROF_GE_TASK_TYPE_WRITE_BACK,
304	- MSPROF_GE_TASK_TYPE_MIX_AIC,
305	- MSPROF_GE_TASK_TYPE_MIX_AIV,
306	- MSPROF_GE_TASK_TYPE_FFTS_PLUS,
307	- MSPROF_GE_TASK_TYPE_DSA,
308	- MSPROF_GE_TASK_TYPE_DVPP,
309	- MSPROF_GE_TASK_TYPE_HCCL,
310	- MSPROF_GE_TASK_TYPE_FUSION,
311	- MSPROF_GE_TASK_TYPE_INVALID
312	-};
313	-
314	-enum MsprofGeShapeType {
315	- MSPROF_GE_SHAPE_TYPE_STATIC = 0,
316	- MSPROF_GE_SHAPE_TYPE_DYNAMIC,
317	-};
318	-struct MsprofGeOpType {
319	- uint8_t type; // MsprofMixDataType
320	- uint8_t rsv[MSPROF_MIX_DATA_RESERVE_BYTES];
321	- union {
322	- uint64_t hashId;
323	- char dataStr[MSPROF_GE_OP_TYPE_LEN];
324	- } data;
325	-};
326	-struct MsprofGeProfTaskData {
327	-#ifdef __cplusplus
328	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
329	- uint16_t dataTag = MSPROF_GE_DATA_TAG_TASK;
330	-#else
331	- uint16_t magicNumber;
332	- uint16_t dataTag;
333	-#endif
334	- uint32_t taskType; // MsprofGeTaskType
335	- struct MsprofMixData opName;
336	- struct MsprofGeOpType opType;
337	- uint64_t curIterNum;
338	- uint64_t timeStamp;
339	- uint32_t shapeType; // MsprofGeShapeType
340	- uint32_t blockDims;
341	- uint32_t modelId;
342	- uint32_t streamId;
343	- uint32_t taskId;
344	- uint32_t threadId;
345	- uint32_t contextId;
346	- uint8_t reserve[MSPROF_GE_TASK_DATA_RESERVE_BYTES];
347	-};
348	-
349	-#define MSPROF_GE_TENSOR_DATA_RESERVE_BYTES 8
350	-#define MSPROF_GE_TENSOR_DATA_SHAPE_LEN 8
351	-#define MSPROF_GE_TENSOR_DATA_NUM 5
352	-enum MsprofGeTensorType {
353	- MSPROF_GE_TENSOR_TYPE_INPUT = 0,
354	- MSPROF_GE_TENSOR_TYPE_OUTPUT,
355	-};
356	-struct MsprofGeTensorData {
357	- uint32_t tensorType; // MsprofGeTensorType
358	- uint32_t format;
359	- uint32_t dataType;
360	- uint32_t shape[MSPROF_GE_TENSOR_DATA_SHAPE_LEN];
361	-};
362	-
363	-struct MsprofGeProfTensorData {
364	-#ifdef __cplusplus
365	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
366	- uint16_t dataTag = MSPROF_GE_DATA_TAG_TENSOR;
367	-#else
368	- uint16_t magicNumber;
369	- uint16_t dataTag;
370	-#endif
371	- uint32_t modelId;
372	- uint64_t curIterNum;
373	- uint32_t streamId;
374	- uint32_t taskId;
375	- uint32_t tensorNum;
376	- struct MsprofGeTensorData tensorData[MSPROF_GE_TENSOR_DATA_NUM];
377	- uint8_t reserve[MSPROF_GE_TENSOR_DATA_RESERVE_BYTES];
378	-};
379	-
380	-#define MSPROF_GE_STEP_DATA_RESERVE_BYTES 27
381	-enum MsprofGeStepTag {
382	- MSPROF_GE_STEP_TAG_BEGIN = 0,
383	- MSPROF_GE_STEP_TAG_END,
384	-};
385	-struct MsprofGeProfStepData {
386	-#ifdef __cplusplus
387	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
388	- uint16_t dataTag = MSPROF_GE_DATA_TAG_STEP;
389	-#else
390	- uint16_t magicNumber;
391	- uint16_t dataTag;
392	-#endif
393	- uint32_t modelId;
394	- uint32_t streamId;
395	- uint32_t taskId;
396	- uint64_t timeStamp;
397	- uint64_t curIterNum;
398	- uint32_t threadId;
399	- uint8_t tag; // MsprofGeStepTag
400	- uint8_t reserve[MSPROF_GE_STEP_DATA_RESERVE_BYTES];
401	-};
402	-
403	-#define MSPROF_GE_ID_MAP_DATA_RESERVE_BYTES 6
404	-struct MsprofGeProfIdMapData {
405	-#ifdef __cplusplus
406	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
407	- uint16_t dataTag = MSPROF_GE_DATA_TAG_ID_MAP;
408	-#else
409	- uint16_t magicNumber;
410	- uint16_t dataTag;
411	-#endif
412	- uint32_t graphId;
413	- uint32_t modelId;
414	- uint32_t sessionId;
415	- uint64_t timeStamp;
416	- uint16_t mode;
417	- uint8_t reserve[MSPROF_GE_ID_MAP_DATA_RESERVE_BYTES];
418	-};
419	-
420	-#define MSPROF_GE_HOST_SCH_DATA_RESERVE_BYTES 24
421	-struct MsprofGeProfHostSchData {
422	-#ifdef __cplusplus
423	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
424	- uint16_t dataTag = MSPROF_GE_DATA_TAG_HOST_SCH;
425	-#else
426	- uint16_t magicNumber;
427	- uint16_t dataTag;
428	-#endif
429	- uint32_t threadId; // record in start event
430	- uint64_t element;
431	- uint64_t event;
432	- uint64_t startTime; // record in start event
433	- uint64_t endTime; // record in end event
434	- uint8_t reserve[MSPROF_GE_HOST_SCH_DATA_RESERVE_BYTES];
435	-};
436	-
437	-/**
438	- * @brief struct of data reported by RunTime
439	- */
440	-#define MSPROF_AICPU_DATA_RESERVE_BYTES 9
441	-struct MsprofAicpuProfData {
442	-#ifdef __cplusplus
443	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
444	- uint16_t dataTag = MSPROF_AICPU_DATA_TAG;
445	-#else
446	- uint16_t magicNumber;
447	- uint16_t dataTag;
448	-#endif
449	- uint16_t streamId;
450	- uint16_t taskId;
451	- uint64_t runStartTime;
452	- uint64_t runStartTick;
453	- uint64_t computeStartTime;
454	- uint64_t memcpyStartTime;
455	- uint64_t memcpyEndTime;
456	- uint64_t runEndTime;
457	- uint64_t runEndTick;
458	- uint32_t threadId;
459	- uint32_t deviceId;
460	- uint64_t submitTick;
461	- uint64_t scheduleTick;
462	- uint64_t tickBeforeRun;
463	- uint64_t tickAfterRun;
464	- uint32_t kernelType;
465	- uint32_t dispatchTime;
466	- uint32_t totalTime;
467	- uint16_t fftsThreadId;
468	- uint8_t version;
469	- uint8_t reserve[MSPROF_AICPU_DATA_RESERVE_BYTES];
470	-};
471	-
472	-struct MsprofAicpuModelProfData {
473	-#ifdef __cplusplus
474	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
475	- uint16_t dataTag = MSPROF_AICPU_MODEL_TAG;
476	-#else
477	- uint16_t magicNumber;
478	- uint16_t dataTag;
479	-#endif
480	- uint32_t rsv; // Ensure 8-byte alignment
481	- uint64_t timeStamp;
482	- uint64_t indexId;
483	- uint32_t modelId;
484	- uint16_t tagId;
485	- uint16_t rsv1;
486	- uint64_t eventId;
487	- uint8_t reserve[24];
488	-};
489	-
490	-/**
491	- * @brief struct of data reported by DP
492	- */
493	-#define MSPROF_DP_DATA_RESERVE_BYTES 16
494	-#define MSPROF_DP_DATA_ACTION_LEN 16
495	-#define MSPROF_DP_DATA_SOURCE_LEN 64
496	-#define MSPROF_CTX_ID_MAX_NUM 55
497	-
498	-struct MsprofDpProfData {
499	-#ifdef __cplusplus
500	- uint16_t magicNumber = MSPROF_DATA_HEAD_MAGIC_NUM;
501	- uint16_t dataTag = MSPROF_DP_DATA_TAG;
502	-#else
503	- uint16_t magicNumber;
504	- uint16_t dataTag;
505	-#endif
506	- uint32_t rsv; // Ensure 8-byte alignment
507	- uint64_t timeStamp;
508	- char action[MSPROF_DP_DATA_ACTION_LEN];
509	- char source[MSPROF_DP_DATA_SOURCE_LEN];
510	- uint64_t index;
511	- uint64_t size;
512	- uint8_t reserve[MSPROF_DP_DATA_RESERVE_BYTES];
513	-};
514	-
515	-struct MsprofAicpuNodeAdditionalData {
516	- uint16_t streamId;
517	- uint16_t taskId;
518	- uint64_t runStartTime;
519	- uint64_t runStartTick;
520	- uint64_t computeStartTime;
521	- uint64_t memcpyStartTime;
522	- uint64_t memcpyEndTime;
523	- uint64_t runEndTime;
524	- uint64_t runEndTick;
525	- uint32_t threadId;
526	- uint32_t deviceId;
527	- uint64_t submitTick;
528	- uint64_t scheduleTick;
529	- uint64_t tickBeforeRun;
530	- uint64_t tickAfterRun;
531	- uint32_t kernelType;
532	- uint32_t dispatchTime;
533	- uint32_t totalTime;
534	- uint16_t fftsThreadId;
535	- uint8_t version;
536	- uint8_t reserve[MSPROF_AICPU_DATA_RESERVE_BYTES];
537	-};
538	-
539	-struct MsprofAicpuModelAdditionalData {
540	- uint64_t indexId;
541	- uint32_t modelId;
542	- uint16_t tagId;
543	- uint16_t rsv1;
544	- uint64_t eventId;
545	- uint8_t reserve[24];
546	-};
547	-
548	-struct MsprofAicpuDpAdditionalData {
549	- char action[MSPROF_DP_DATA_ACTION_LEN];
550	- char source[MSPROF_DP_DATA_SOURCE_LEN];
551	- uint64_t index;
552	- uint64_t size;
553	- uint8_t reserve[MSPROF_DP_DATA_RESERVE_BYTES];
554	-};
555	-
556	-struct MsprofAicpuMiAdditionalData {
557	- uint32_t nodeTag; // MsprofMindsporeNodeTag:1
558	- uint32_t reserve;
559	- uint64_t queueSize;
560	- uint64_t runStartTime;
561	- uint64_t runEndTime;
562	-};
563	-
564	-// AICPU kfc算子执行时间
565	-struct AicpuKfcProfCommTurn {
566	- uint64_t waitNotifyStartTime; // 开始等待通信参数
567	- uint64_t kfcAlgExeStartTime; // 开始通信算法执行
568	- uint64_t sendTaskStartTime; // 开始下发task
569	- uint64_t waitActiveStartTime; // 开始等待激活
570	- uint64_t acitveStartTime; // 开始激活处理
571	- uint64_t waitExeEndStartTime; // 开始等待任务执行结束
572	- uint64_t rtsqExeEndTime; // 任务执行结束时间
573	- uint64_t dataLen; // 本轮通信数据长度
574	- uint32_t deviceId;
575	- uint16_t streamId;
576	- uint16_t taskId;
577	- uint8_t version;
578	- uint8_t commTurn; // 总通信轮次
579	- uint8_t currentTurn;
580	- uint8_t reserve[5];
581	-};
582	-
583	-// Aicore算子执行时间
584	-struct AicpuKfcProfComputeTurn {
585	- uint64_t waitComputeStartTime; // 开始等待计算
586	- uint64_t computeStartTime; // 开始计算
587	- uint64_t computeExeEndTime; // 计算执行结束
588	- uint64_t dataLen; // 本轮计算数据长度
589	- uint32_t deviceId;
590	- uint16_t streamId;
591	- uint16_t taskId;
592	- uint8_t version;
593	- uint8_t computeTurn; // 总计算轮次
594	- uint8_t currentTurn;
595	- uint8_t reserve[5];
596	-};
597	-
598	-/**
599	- * @brief struct of data reported by HCCL
600	- */
601	-#pragma pack(4)
602	-struct MsprofHcclProfNotify {
603	- uint32_t taskID;
604	- uint64_t notifyID;
605	- uint32_t stage;
606	- uint32_t remoteRank;
607	- uint32_t transportType;
608	- uint32_t role; // role {0: dst, 1:src}
609	- double durationEstimated;
610	-};
611	-
612	-struct MsprofHcclProfReduce {
613	- uint32_t taskID;
614	- uint64_t src;
615	- uint64_t dst;
616	- uint64_t size;
617	- uint32_t op; // {0: sum, 1: mul, 2: max, 3: min}
618	- uint32_t dataType; // data type {0: INT8, 1: INT16, 2: INT32, 3: FP16, 4:FP32, 5:INT64, 6:UINT64}
619	- uint32_t linkType; // link type {0: 'OnChip', 1: 'HCCS', 2: 'PCIe', 3: 'RoCE'}
620	- uint32_t remoteRank;
621	- uint32_t transportType; // transport type {0: SDMA, 1: RDMA, 2:LOCAL}
622	- uint32_t role; // role {0: dst, 1:src}
623	- double durationEstimated;
624	-};
625	-
626	-struct MsprofHcclProfRDMA {
627	- uint32_t taskID;
628	- uint64_t src;
629	- uint64_t dst;
630	- uint64_t size;
631	- uint64_t notifyID;
632	- uint32_t linkType; // link type {0: 'OnChip', 1: 'HCCS', 2: 'PCIe', 3: 'RoCE'}
633	- uint32_t remoteRank;
634	- uint32_t transportType; // transport type {0: RDMA, 1:SDMA, 2:LOCAL}
635	- uint32_t role; // role {0: dst, 1:src}
636	- uint32_t type; // RDMA type {0: RDMASendNotify, 1:RDMASendPayload}
637	- double durationEstimated;
638	-};
639	-
640	-struct MsprofHcclProfMemcpy {
641	- uint32_t taskID;
642	- uint64_t src;
643	- uint64_t dst;
644	- uint64_t size;
645	- uint64_t notifyID;
646	- uint32_t linkType; // link type {0: 'OnChip', 1: 'HCCS', 2: 'PCIe', 3: 'RoCE'}
647	- uint32_t remoteRank;
648	- uint32_t transportType; // transport type {0: RDMA, 1:SDMA, 2:LOCAL}
649	- uint32_t role; // role {0: dst, 1:src}
650	- double durationEstimated;
651	-};
652	-
653	-struct MsprofHcclProfStageStep {
654	- uint32_t rank;
655	- uint32_t rankSize;
656	-};
657	-
658	-struct MsprofHcclProfFlag {
659	- uint64_t cclTag;
660	- uint64_t groupName;
661	- uint32_t localRank;
662	- uint32_t workFlowMode;
663	-};
664	-
665	-#define MSPROF_HCCL_INVALID_UINT 0xFFFFFFFFU
666	-struct MsprofHcclInfo {
667	- uint64_t itemId;
668	- uint64_t cclTag;
669	- uint64_t groupName;
670	- uint32_t localRank;
671	- uint32_t remoteRank;
672	- uint32_t rankSize;
673	- uint32_t workFlowMode;
674	- uint32_t planeID;
675	- uint32_t ctxId;
676	- uint64_t notifyID;
677	- uint32_t stage;
678	- uint32_t role; // role {0: dst, 1:src}
679	- double durationEstimated;
680	- uint64_t srcAddr;
681	- uint64_t dstAddr;
682	- uint64_t dataSize; // bytes
683	- uint32_t opType; // {0: sum, 1: mul, 2: max, 3: min}
684	- uint32_t dataType; // data type {0: INT8, 1: INT16, 2: INT32, 3: FP16, 4:FP32, 5:INT64, 6:UINT64}
685	- uint32_t linkType; // link type {0: 'OnChip', 1: 'HCCS', 2: 'PCIe', 3: 'RoCE'}
686	- uint32_t transportType; // transport type {0: SDMA, 1: RDMA, 2:LOCAL}
687	- uint32_t rdmaType; // RDMA type {0: RDMASendNotify, 1:RDMASendPayload}
688	- uint32_t reserve2;
689	-#ifdef __cplusplus
690	- MsprofHcclInfo() : role(MSPROF_HCCL_INVALID_UINT), opType(MSPROF_HCCL_INVALID_UINT),
691	- dataType(MSPROF_HCCL_INVALID_UINT), linkType(MSPROF_HCCL_INVALID_UINT),
692	- transportType(MSPROF_HCCL_INVALID_UINT), rdmaType(MSPROF_HCCL_INVALID_UINT)
693	- {
694	- }
695	-#endif
696	-};
697	-
698	-struct MsprofAicpuMC2HcclInfo {
699	- uint64_t itemId;
700	- uint64_t cclTag;
701	- uint64_t groupName;
702	- uint32_t localRank;
703	- uint32_t remoteRank;
704	- uint32_t rankSize;
705	- uint32_t workFlowMode;
706	- uint32_t planeID;
707	- uint32_t ctxId;
708	- uint64_t notifyID;
709	- uint32_t stage;
710	- uint32_t role; // role {0: dst, 1:src}
711	- double durationEstimated;
712	- uint64_t srcAddr;
713	- uint64_t dstAddr;
714	- uint64_t dataSize; // bytes
715	- uint32_t opType; // {0: sum, 1: mul, 2: max, 3: min}
716	- uint32_t dataType; // data type {0: INT8, 1: INT16, 2: INT32, 3: FP16, 4:FP32, 5:INT64, 6:UINT64}
717	- uint32_t linkType; // link type {0: 'OnChip', 1: 'HCCS', 2: 'PCIe', 3: 'RoCE'}
718	- uint32_t transportType; // transport type {0: SDMA, 1: RDMA, 2:LOCAL}
719	- uint32_t rdmaType; // RDMA type {0: RDMASendNotify, 1:RDMASendPayload}
720	- uint32_t taskId;
721	- uint16_t streamId;
722	- uint16_t reserve[3];
723	-};
724	-
725	-struct ProfilingDeviceCommResInfo {
726	- uint64_t groupName; // 通信域
727	- uint32_t rankSize; // 通信域内rank总数
728	- uint32_t rankId; // 当前device rankId，通信域内编号
729	- uint32_t usrRankId; // 当前device rankId，全局编号
730	- uint32_t aicpuKfcStreamId; // MC2中launch aicpu kfc算子的stream
731	- uint32_t commStreamSize; // 当前device侧使用的通信stream数量
732	- uint32_t commStreamIds[8]; // 具体streamId
733	- uint32_t reserve;
734	-};
735	-
736	-#define MSPROF_MULTI_THREAD_MAX_NUM 25
737	-struct MsprofMultiThread {
738	- uint32_t threadNum;
739	- uint32_t threadId[MSPROF_MULTI_THREAD_MAX_NUM];
740	-};
741	-#pragma pack()
742	-
743	-
744	-#pragma pack(1)
745	-struct MsprofNodeBasicInfo {
746	- uint64_t opName;
747	- uint32_t taskType;
748	- uint64_t opType;
749	- uint32_t blockDim;
750	- uint32_t opFlag;
751	-};
752	-
753	-enum AttrType {
754	- OP_ATTR = 0,
755	-};
756	-
757	-struct MsprofAttrInfo {
758	- uint64_t opName;
759	- uint32_t attrType;
760	- uint64_t hashId;
761	-};
762	-
763	-struct MsrofTensorData {
764	- uint32_t tensorType;
765	- uint32_t format;
766	- uint32_t dataType;
767	- uint32_t shape[MSPROF_GE_TENSOR_DATA_SHAPE_LEN];
768	-};
769	-
770	-struct MsprofTensorInfo {
771	- uint64_t opName;
772	- uint32_t tensorNum;
773	- struct MsrofTensorData tensorData[MSPROF_GE_TENSOR_DATA_NUM];
774	-};
775	-
776	-struct MsprofHCCLOPInfo { // for MsprofReportCompactInfo buffer data
777	- uint8_t relay : 1; // 借轨通信
778	- uint8_t retry : 1; // 重传标识
779	- uint8_t dataType; // 跟HcclDataType类型保存一致
780	- uint16_t algType; // algtype 每4位表示一个算法阶段
781	- uint64_t count; // 发送数据个数
782	- uint64_t groupName; // group hash id
783	-};
784	-
785	-struct ProfFusionOpInfo {
786	-uint64_t opName;
787	-uint32_t fusionOpNum;
788	-uint64_t inputMemsize;
789	-uint64_t outputMemsize;
790	-uint64_t weightMemSize;
791	-uint64_t workspaceMemSize;
792	-uint64_t totalMemSize;
793	-uint64_t fusionOpId[MSPROF_GE_FUSION_OP_NUM];
794	-};
795	-
796	-struct MsprofContextIdInfo {
797	- uint64_t opName;
798	- uint32_t ctxIdNum;
799	- uint32_t ctxIds[MSPROF_CTX_ID_MAX_NUM];
800	-};
801	-
802	-struct MsprofGraphIdInfo {
803	- uint64_t modelName;
804	- uint32_t graphId;
805	- uint32_t modelId;
806	-};
807	-
808	-struct MsprofMemoryInfo {
809	- uint64_t addr;
810	- int64_t size;
811	- uint64_t nodeId; // op name hash id
812	- uint64_t totalAllocateMemory;
813	- uint64_t totalReserveMemory;
814	- uint32_t deviceId;
815	- uint32_t deviceType;
816	-};
817	-
818	-/**
819	- * @name MsprofStampInfo
820	- * @brief struct of data reported by msproftx
821	- */
822	-struct MsprofStampInfo {
823	- uint16_t magicNumber;
824	- uint16_t dataTag;
825	- uint32_t processId;
826	- uint32_t threadId;
827	- uint32_t category; // marker category
828	- uint32_t eventType;
829	- int32_t payloadType;
830	- union PayloadValue {
831	- uint64_t ullValue;
832	- int64_t llValue;
833	- double dValue;
834	- uint32_t uiValue[2];
835	- int32_t iValue[2];
836	- float fValue[2];
837	- } payload; // payload info for marker
838	- uint64_t startTime;
839	- uint64_t endTime;
840	- int32_t messageType;
841	- char message[128];
842	-};
843	-
844	-struct MsprofStaticOpMem {
845	- int64_t size; // op memory size
846	- uint64_t opName; // op name hash id
847	- uint64_t lifeStart; // serial number of op memory used
848	- uint64_t lifeEnd; // serial number of op memory used
849	- uint64_t totalAllocateMemory; // static graph total allocate memory
850	- uint64_t dynOpName; // 0: invalid， other： dynamic op name of root
851	- uint32_t graphId; // multipe model
852	-};
853	-
854	-#define MSPROF_PHYSIC_STREAM_ID_MAX_NUM 56
855	-struct MsprofLogicStreamInfo {
856	- uint32_t logicStreamId;
857	- uint32_t physicStreamNum;
858	- uint32_t physicStreamId[MSPROF_PHYSIC_STREAM_ID_MAX_NUM];
859	-};
860	-
861	-struct MsprofExeomLoadInfo {
862	- uint32_t modelId;
863	- uint32_t reserve;
864	- uint64_t modelName; /* name hash */
865	-};
866	-#pragma pack()
867	-
868	-struct MsprofApi { // for MsprofReportApi
869	-#ifdef __cplusplus
870	- uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
871	-#else
872	- uint16_t magicNumber;
873	-#endif
874	- uint16_t level;
875	- uint32_t type;
876	- uint32_t threadId;
877	- uint32_t reserve;
878	- uint64_t beginTime;
879	- uint64_t endTime;
880	- uint64_t itemId;
881	-};
882	-
883	-struct MsprofEvent { // for MsprofReportEvent
884	-#ifdef __cplusplus
885	- uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
886	-#else
887	- uint16_t magicNumber;
888	-#endif
889	- uint16_t level;
890	- uint32_t type;
891	- uint32_t threadId;
892	- uint32_t requestId; // 0xFFFF means single event
893	- uint64_t timeStamp;
894	-#ifdef __cplusplus
895	- uint64_t eventFlag = MSPROF_EVENT_FLAG;
896	-#else
897	- uint64_t eventFlag;
898	-#endif
899	- uint64_t itemId;
900	-};
901	-
902	-struct MsprofRuntimeTrack { // for MsprofReportCompactInfo buffer data
903	- uint16_t deviceId;
904	- uint16_t streamId;
905	- uint32_t taskId;
906	- uint64_t taskType; // task message hash id
907	-};
908	-
909	-struct MsprofDpuTrack { // for MsprofReportCompactInfo buffer data
910	- uint16_t deviceId; // high 4 bits, devType: dpu: 1, low 12 bits device id
911	- uint16_t streamId;
912	- uint32_t taskId;
913	- uint32_t taskType; // task type enum
914	- uint32_t res;
915	- uint64_t startTime; // start time
916	-};
917	-
918	-#define MSPROF_COMPACT_INFO_DATA_LENGTH (40)
919	-struct MsprofCompactInfo { // for MsprofReportCompactInfo buffer data
920	-#ifdef __cplusplus
921	- uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
922	-#else
923	- uint16_t magicNumber;
924	-#endif
925	- uint16_t level;
926	- uint32_t type;
927	- uint32_t threadId;
928	- uint32_t dataLen;
929	- uint64_t timeStamp;
930	- union {
931	- uint8_t info[MSPROF_COMPACT_INFO_DATA_LENGTH];
932	- struct MsprofRuntimeTrack runtimeTrack;
933	- struct MsprofNodeBasicInfo nodeBasicInfo;
934	- struct MsprofHCCLOPInfo hcclopInfo;
935	- struct MsprofDpuTrack dpuTack;
936	- } data;
937	-};
938	-
939	-#define MSPROF_ADDTIONAL_INFO_DATA_LENGTH (232)
940	-struct MsprofAdditionalInfo { // for MsprofReportAdditionalInfo buffer data
941	-#ifdef __cplusplus
942	- uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
943	-#else
944	- uint16_t magicNumber;
945	-#endif
946	- uint16_t level;
947	- uint32_t type;
948	- uint32_t threadId;
949	- uint32_t dataLen;
950	- uint64_t timeStamp;
951	- uint8_t data[MSPROF_ADDTIONAL_INFO_DATA_LENGTH];
952	-};
953	-
954	-/**
955	- * @name MsprofErrorCode
956	- * @brief error code
957	- */
958	-enum MsprofErrorCode {
959	- MSPROF_ERROR_NONE = 0,
960	- MSPROF_ERROR_MEM_NOT_ENOUGH,
961	- MSPROF_ERROR_GET_ENV,
962	- MSPROF_ERROR_CONFIG_INVALID,
963	- MSPROF_ERROR_ACL_JSON_OFF,
964	- MSPROF_ERROR,
965	- MSPROF_ERROR_UNINITIALIZE,
966	-};
967	-
968	-#define MSPROF_ENGINE_MAX_TAG_LEN (63)
969	-
970	-/**
971	- * @name ReporterData
972	- * @brief struct of data to report
973	- */
974	-struct ReporterData {
975	- char tag[MSPROF_ENGINE_MAX_TAG_LEN + 1]; // the sub-type of the module, data with different tag will be writen
976	- int32_t deviceId; // the index of device
977	- size_t dataLen; // the length of send data
978	- uint8_t *data; // the data content
979	-};
980	-
981	-/**
982	- * @name MsprofHashData
983	- * @brief struct of data to hash
984	- */
985	-struct MsprofHashData {
986	- int32_t deviceId; // the index of device
987	- size_t dataLen; // the length of data
988	- uint8_t *data; // the data content
989	- uint64_t hashId; // the id of hashed data
990	-};
991	-
992	-enum MsprofConfigParamType {
993	- DEV_CHANNEL_RESOURCE = 0, // device channel resource
994	- HELPER_HOST_SERVER // helper host server
995	-};
996	-
997	-/**
998	- * @name MsprofConfigParam
999	- * @brief struct of set config
1000	- */
1001	-struct MsprofConfigParam {
1002	- uint32_t deviceId; // the index of device
1003	- uint32_t type; // DEV_CHANNEL_RESOURCE; HELPER_HOST_SERVER
1004	- uint32_t value; // DEV_CHANNEL_RESOURCE: 1 off; HELPER_HOST_SERVER: 1 on
1005	-};
1006	-
1007	-/**
1008	- * @name MsprofReporterModuleId
1009	- * @brief module id of data to report
1010	- */
1011	-enum MsprofReporterModuleId {
1012	- MSPROF_MODULE_DATA_PREPROCESS = 0, // DATA_PREPROCESS
1013	- MSPROF_MODULE_HCCL, // HCCL
1014	- MSPROF_MODULE_ACL, // AclModule
1015	- MSPROF_MODULE_FRAMEWORK, // Framework
1016	- MSPROF_MODULE_RUNTIME, // runtime
1017	- MSPROF_MODULE_MSPROF // msprofTx
1018	-};
1019	-
1020	-/**
1021	- * @name MsprofReporterCallbackType
1022	- * @brief reporter callback request type
1023	- */
1024	-enum MsprofReporterCallbackType {
1025	- MSPROF_REPORTER_REPORT = 0, // report data
1026	- MSPROF_REPORTER_INIT, // init reporter
1027	- MSPROF_REPORTER_UNINIT, // uninit reporter
1028	- MSPROF_REPORTER_DATA_MAX_LEN, // data max length for calling report callback
1029	- MSPROF_REPORTER_HASH // hash data to id
1030	-};
1031	-
1032	-#define MSPROF_OPTIONS_DEF_LEN_MAX (2048U)
1033	-
1034	-/**
1035	- * @name MsprofGeOptions
1036	- * @brief struct of MSPROF_CTRL_INIT_GE_OPTIONS
1037	- */
1038	-struct MsprofGeOptions {
1039	- char jobId[MSPROF_OPTIONS_DEF_LEN_MAX];
1040	- char options[MSPROF_OPTIONS_DEF_LEN_MAX];
1041	-};
1042	-
1043	-/**
1044	- * @name MsprofCtrlCallbackType
1045	- * @brief ctrl callback request type
1046	- */
1047	-enum MsprofCtrlCallbackType {
1048	- MSPROF_CTRL_INIT_ACL_ENV = 0, // start profiling with acl env
1049	- MSPROF_CTRL_INIT_ACL_JSON = 1, // start pro with acl.json
1050	- MSPROF_CTRL_INIT_GE_OPTIONS = 2, // start profiling with ge env and options
1051	- MSPROF_CTRL_FINALIZE = 3, // stop profiling
1052	- MSPROF_CTRL_INIT_HELPER = 4, // start profiling in helper device
1053	- MSPROF_CTRL_INIT_PURE_CPU = 5, // start profiling in pure cpu
1054	- MSPROF_CTRL_INIT_DYNA = 0xFF, // start profiling for dynamic profiling
1055	-};
1056	-
1057	-enum MsprofCommandHandleType {
1058	- PROF_COMMANDHANDLE_TYPE_INIT = 0,
1059	- PROF_COMMANDHANDLE_TYPE_START,
1060	- PROF_COMMANDHANDLE_TYPE_STOP,
1061	- PROF_COMMANDHANDLE_TYPE_FINALIZE,
1062	- PROF_COMMANDHANDLE_TYPE_MODEL_SUBSCRIBE,
1063	- PROF_COMMANDHANDLE_TYPE_MODEL_UNSUBSCRIBE,
1064	- PROF_COMMANDHANDLE_TYPE_MAX
1065	-};
1066	-
1067	-enum MsprofConfigType {
1068	- MSPROF_CONFIG_HELPER_HOST = 0
1069	-};
1070	-
1071	-/**
1072	- * @brief profiling command type
1073	- */
1074	-enum ProfCtrlType {
1075	- PROF_CTRL_INVALID = 0,
1076	- PROF_CTRL_SWITCH,
1077	- PROF_CTRL_REPORTER,
1078	- PROF_CTRL_STEPINFO,
1079	- PROF_CTRL_BUTT
1080	-};
1081	-
1082	-/**
1083	- * @brief Prof Chip ID
1084	- */
1085	-enum Prof_Chip_ID {
1086	- PROF_CHIP_ID0 = 0
1087	-};
1088	-
1089	-/**
1090	- * @brief the struct of profiling set setp info
1091	- */
1092	-typedef struct ProfStepInfoCmd {
1093	- uint64_t index_id;
1094	- uint16_t tag_id;
1095	- void *stream;
1096	-} ProfStepInfoCmd_t;
1097	-
1098	-#ifdef __cplusplus
1099	-}
1100	-#endif
1101	-#endif // MSPROFILER_PROF_COMMON_H_

Dthird_party/acl/inc/experiment/msprof/toolchain/prof_data_config.h+0-99

	@@ -1,99 +0,0 @@
1	-/**
2	- * @file prof_data_config.h
3	- *
4	- * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
5	- *
6	- * This program is distributed in the hope that it will be useful,
7	- * but WITHOUT ANY WARRANTY; without even the implied warranty of
8	- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
9	- *
10	- */
11	-
12	-#ifndef MSPROFILER_API_PROF_DATA_CONFIG_H
13	-#define MSPROFILER_API_PROF_DATA_CONFIG_H
14	-
15	-// DataTypeConfig
16	-#define PROF_ACL_API 0x00000001ULL
17	-#define PROF_TASK_TIME_L1 0x00000002ULL
18	-#define PROF_AICORE_METRICS 0x00000004ULL
19	-#define PROF_AICPU_TRACE 0x00000008ULL
20	-#define PROF_L2CACHE 0x00000010ULL
21	-#define PROF_HCCL_TRACE 0x00000020ULL
22	-#define PROF_TRAINING_TRACE 0x00000040ULL
23	-#define PROF_MSPROFTX 0x00000080ULL
24	-#define PROF_RUNTIME_API 0x00000100ULL
25	-#define PROF_FWK_SCHEDULE_L0 0x00000200ULL
26	-#define PROF_TASK_TSFW 0x00000400ULL
27	-#define PROF_TASK_TIME 0x00000800ULL
28	-#define PROF_TASK_MEMORY 0x00001000ULL
29	-#define PROF_TASK_TIME_L2 0x00002000ULL
30	-#define PROF_OP_ATTR 0x00004000ULL
31	-
32	-// system profilinig switch
33	-#define PROF_CPU 0x00010000ULL
34	-#define PROF_HARDWARE_MEMORY 0x00020000ULL
35	-#define PROF_IO 0x00040000ULL
36	-#define PROF_INTER_CONNECTION 0x00080000ULL
37	-#define PROF_DVPP 0x00100000ULL
38	-#define PROF_SYS_AICORE_SAMPLE 0x00200000ULL
39	-#define PROF_AIVECTORCORE_SAMPLE 0x00400000ULL
40	-#define PROF_INSTR 0x00800000ULL
41	-
42	-#define PROF_FWK_SCHEDULE_L1 0x0000001000000ULL
43	-#define PROF_PURE_CPU 0x0000002000000ULL
44	-#define PROF_RUNTIME_TRACE 0x0000004000000ULL
45	-#define PROF_SCHEDULE_TIMELINE 0x0000008000000ULL
46	-#define PROF_SCHEDULE_TRACE 0x0000010000000ULL
47	-#define PROF_AIVECTORCORE_METRICS 0x0000020000000ULL
48	-#define PROF_SUBTASK_TIME 0x0000040000000ULL
49	-#define PROF_OP_DETAIL 0x0000080000000ULL
50	-
51	-#define PROF_AICPU_MODEL 0x4000000000000000ULL
52	-#define PROF_MODEL_LOAD 0x8000000000000000ULL
53	-
54	-#define PROF_TASK_TRACE PROF_RUNTIME_TRACE \| PROF_TASK_TIME \| PROF_TRAINING_TRACE \| PROF_HCCL_TRACE \| PROF_TASK_TIME_L1
55	-
56	-// DataTypeConfig MASK
57	-#define PROF_ACL_API_MASK 0x00000001ULL
58	-#define PROF_TASK_TIME_L1_MASK 0x00000002ULL
59	-#define PROF_AICORE_METRICS_MASK 0x00000004ULL
60	-#define PROF_AICPU_TRACE_MASK 0x00000008ULL
61	-#define PROF_L2CACHE_MASK 0x00000010ULL
62	-#define PROF_HCCL_TRACE_MASK 0x00000020ULL
63	-#define PROF_TRAINING_TRACE_MASK 0x00000040ULL
64	-#define PROF_MSPROFTX_MASK 0x00000080ULL
65	-#define PROF_RUNTIME_API_MASK 0x00000100ULL
66	-#define PROF_TASK_FRAMEWORK_MASK 0x00000200ULL
67	-#define PROF_FWK_SCHEDULE_L0_MASK 0x00000200ULL
68	-#define PROF_TASK_TSFW_MASK 0x00000400ULL
69	-#define PROF_TASK_TIME_MASK 0x00000800ULL
70	-#define PROF_TASK_MEMORY_MASK 0x00001000ULL
71	-#define PROF_TASK_TIME_L2_MASK 0x00002000ULL
72	-#define PROF_OP_ATTR_MASK 0x00004000ULL
73	-
74	-// system profilinig mask
75	-#define PROF_CPU_MASK 0x00010000ULL
76	-#define PROF_HARDWARE_MEMORY_MASK 0x00020000ULL
77	-#define PROF_IO_MASK 0x00040000ULL
78	-#define PROF_INTER_CONNECTION_MASK 0x00080000ULL
79	-#define PROF_DVPP_MASK 0x00100000ULL
80	-#define PROF_SYS_AICORE_SAMPLE_MASK 0x00200000ULL
81	-#define PROF_AIVECTORCORE_SAMPLE_MASK 0x00400000ULL
82	-#define PROF_INSTR_MASK 0x00800000ULL
83	-
84	-#define PROF_MODEL_EXECUTE_MASK 0x0000001000000ULL
85	-#define PROF_FWK_SCHEDULE_L1_MASK 0x0000001000000ULL
86	-#define PROF_RUNTIME_TRACE_MASK 0x0000004000000ULL
87	-#define PROF_SCHEDULE_TIMELINE_MASK 0x0000008000000ULL
88	-#define PROF_SCHEDULE_TRACE_MASK 0x0000010000000ULL
89	-#define PROF_AIVECTORCORE_METRICS_MASK 0x0000020000000ULL
90	-#define PROF_SUBTASK_TIME_MASK 0x0000040000000ULL
91	-#define PROF_OP_DETAIL_MASK 0x0000080000000ULL
92	-
93	-#define PROF_AICPU_MODEL_MASK 0x4000000000000000ULL
94	-#define PROF_MODEL_LOAD_MASK 0x8000000000000000ULL
95	-
96	-// devprof config
97	-#define PROF_HI_MC2 0x1000000000000ULL // bit0 of the upper 16 bits of profSwitchHi means MC2 switch
98	-
99	-#endif // MSPROFILER_API_PROF_DATA_CONFIG_H_

Athird_party/acl/inc/profiling/aprof_pub.h+180-0

@@ -0,0 +1,180 @@
	1	+/**
	2	+ * @file aprof_pub.h
	3	+ *
	4	+ * Minimal profiling declarations required by DVM and MLIR launcher code.
	5	+ */
	6	+#ifndef APROF_PUB_H
	7	+#define APROF_PUB_H
	8	+
	9	+#include <stddef.h>
	10	+#include <stdint.h>
	11	+
	12	+#ifdef __cplusplus
	13	+extern "C" {
	14	+#endif
	15	+
	16	+#if (defined(_WIN32) \|\| defined(_WIN64) \|\| defined(_MSC_VER))
	17	+#define MSVP_PROF_API __declspec(dllexport)
	18	+#else
	19	+#define MSVP_PROF_API __attribute__((visibility("default")))
	20	+#endif
	21	+
	22	+typedef void *VOID_PTR;
	23	+
	24	+#define MSPROF_REPORT_DATA_MAGIC_NUM 0x5A5AU
	25	+#define MSPROF_COMPACT_INFO_DATA_LENGTH 40
	26	+#define PATH_LEN_MAX 1023
	27	+#define PARAM_LEN_MAX 4095
	28	+#define MSPROF_MAX_DEV_NUM 64
	29	+#define MSPROF_GE_TENSOR_DATA_SHAPE_LEN 8
	30	+#define MSPROF_GE_TENSOR_DATA_NUM 5
	31	+#define MSPROF_CTX_ID_MAX_NUM 55
	32	+#define MSPROF_ADDTIONAL_INFO_DATA_LENGTH 232
	33	+
	34	+#define MSPROF_REPORT_NODE_LEVEL 10000U
	35	+#define MSPROF_REPORT_NODE_BASIC_INFO_TYPE 0U
	36	+#define MSPROF_REPORT_NODE_TENSOR_INFO_TYPE 1U
	37	+#define MSPROF_REPORT_NODE_CONTEXT_ID_INFO_TYPE 4U
	38	+#define MSPROF_REPORT_NODE_LAUNCH_TYPE 5U
	39	+
	40	+#define PROF_TASK_TIME_L1_MASK 0x00000002ULL
	41	+#define PROF_TASK_TIME_MASK 0x00000800ULL
	42	+
	43	+enum MsprofErrorCode {
	44	+ MSPROF_ERROR_NONE = 0,
	45	+};
	46	+
	47	+enum MsprofCommandHandleType {
	48	+ PROF_COMMANDHANDLE_TYPE_INIT = 0,
	49	+ PROF_COMMANDHANDLE_TYPE_START,
	50	+ PROF_COMMANDHANDLE_TYPE_STOP,
	51	+ PROF_COMMANDHANDLE_TYPE_FINALIZE,
	52	+ PROF_COMMANDHANDLE_TYPE_MODEL_SUBSCRIBE,
	53	+ PROF_COMMANDHANDLE_TYPE_MODEL_UNSUBSCRIBE,
	54	+ PROF_COMMANDHANDLE_TYPE_MAX
	55	+};
	56	+
	57	+enum MsprofGeTaskType {
	58	+ MSPROF_GE_TASK_TYPE_AI_CORE = 0,
	59	+};
	60	+
	61	+enum MsprofGeTensorType {
	62	+ MSPROF_GE_TENSOR_TYPE_INPUT = 0,
	63	+ MSPROF_GE_TENSOR_TYPE_OUTPUT,
	64	+};
	65	+
	66	+enum ProfCtrlType {
	67	+ PROF_CTRL_INVALID = 0,
	68	+ PROF_CTRL_SWITCH,
	69	+};
	70	+
	71	+struct MsprofCommandHandleParams {
	72	+ uint32_t pathLen;
	73	+ uint32_t storageLimit;
	74	+ uint32_t profDataLen;
	75	+ char path[PATH_LEN_MAX + 1];
	76	+ char profData[PARAM_LEN_MAX + 1];
	77	+};
	78	+
	79	+struct MsprofCommandHandle {
	80	+ uint64_t profSwitch;
	81	+ uint64_t profSwitchHi;
	82	+ uint32_t devNums;
	83	+ uint32_t devIdList[MSPROF_MAX_DEV_NUM];
	84	+ uint32_t modelId;
	85	+ uint32_t type;
	86	+ uint32_t cacheFlag;
	87	+ struct MsprofCommandHandleParams params;
	88	+};
	89	+
	90	+#pragma pack(1)
	91	+
	92	+struct MsprofNodeBasicInfo {
	93	+ uint64_t opName;
	94	+ uint32_t taskType;
	95	+ uint64_t opType;
	96	+ uint32_t blockDim;
	97	+ uint32_t opFlag;
	98	+};
	99	+
	100	+struct MsrofTensorData {
	101	+ uint32_t tensorType;
	102	+ uint32_t format;
	103	+ uint32_t dataType;
	104	+ uint32_t shape[MSPROF_GE_TENSOR_DATA_SHAPE_LEN];
	105	+};
	106	+
	107	+struct MsprofTensorInfo {
	108	+ uint64_t opName;
	109	+ uint32_t tensorNum;
	110	+ struct MsrofTensorData tensorData[MSPROF_GE_TENSOR_DATA_NUM];
	111	+};
	112	+
	113	+struct MsprofContextIdInfo {
	114	+ uint64_t opName;
	115	+ uint32_t ctxIdNum;
	116	+ uint32_t ctxIds[MSPROF_CTX_ID_MAX_NUM];
	117	+};
	118	+
	119	+#pragma pack()
	120	+
	121	+struct MsprofApi {
	122	+#ifdef __cplusplus
	123	+ uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
	124	+#else
	125	+ uint16_t magicNumber;
	126	+#endif
	127	+ uint16_t level;
	128	+ uint32_t type;
	129	+ uint32_t threadId;
	130	+ uint32_t reserve;
	131	+ uint64_t beginTime;
	132	+ uint64_t endTime;
	133	+ uint64_t itemId;
	134	+};
	135	+
	136	+struct MsprofCompactInfo {
	137	+#ifdef __cplusplus
	138	+ uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
	139	+#else
	140	+ uint16_t magicNumber;
	141	+#endif
	142	+ uint16_t level;
	143	+ uint32_t type;
	144	+ uint32_t threadId;
	145	+ uint32_t dataLen;
	146	+ uint64_t timeStamp;
	147	+ union {
	148	+ uint8_t info[MSPROF_COMPACT_INFO_DATA_LENGTH];
	149	+ struct MsprofNodeBasicInfo nodeBasicInfo;
	150	+ } data;
	151	+};
	152	+
	153	+struct MsprofAdditionalInfo {
	154	+#ifdef __cplusplus
	155	+ uint16_t magicNumber = MSPROF_REPORT_DATA_MAGIC_NUM;
	156	+#else
	157	+ uint16_t magicNumber;
	158	+#endif
	159	+ uint16_t level;
	160	+ uint32_t type;
	161	+ uint32_t threadId;
	162	+ uint32_t dataLen;
	163	+ uint64_t timeStamp;
	164	+ uint8_t data[MSPROF_ADDTIONAL_INFO_DATA_LENGTH];
	165	+};
	166	+
	167	+typedef int32_t (ProfCommandHandle)(uint32_t type, void data, uint32_t len);
	168	+
	169	+MSVP_PROF_API int32_t MsprofRegisterCallback(uint32_t moduleId, ProfCommandHandle handle);
	170	+MSVP_PROF_API int32_t MsprofReportApi(uint32_t nonPersistantFlag, const struct MsprofApi *api);
	171	+MSVP_PROF_API int32_t MsprofReportCompactInfo(uint32_t nonPersistantFlag, const VOID_PTR data, uint32_t length);
	172	+MSVP_PROF_API int32_t MsprofReportAdditionalInfo(uint32_t nonPersistantFlag, const VOID_PTR data, uint32_t length);
	173	+MSVP_PROF_API uint64_t MsprofGetHashId(const char *hashInfo, size_t length);
	174	+MSVP_PROF_API uint64_t MsprofSysCycleTime(void);
	175	+
	176	+#ifdef __cplusplus
	177	+}
	178	+#endif
	179	+
	180	+#endif

Athird_party/acl/inc/profiling/prof_api.h+11-0

@@ -0,0 +1,11 @@
	1	+/**
	2	+ * @file prof_api.h
	3	+ *
	4	+ * Minimal profiling API declarations required by DVM and MLIR launcher code.
	5	+ */
	6	+#ifndef PROF_API_H
	7	+#define PROF_API_H
	8	+
	9	+#include "prof_common.h"
	10	+
	11	+#endif

Athird_party/acl/inc/profiling/prof_common.h+24-0

@@ -0,0 +1,24 @@
	1	+/**
	2	+ * @file prof_common.h
	3	+ *
	4	+ * Minimal common profiling declarations required by DVM and MLIR launcher code.
	5	+ */
	6	+#ifndef MSPROFILER_PROF_COMMON_H
	7	+#define MSPROFILER_PROF_COMMON_H
	8	+
	9	+#include "aprof_pub.h"
	10	+
	11	+#ifdef __cplusplus
	12	+extern "C" {
	13	+#endif
	14	+
	15	+#define MSPROF_DATA_HEAD_MAGIC_NUM 0x5A5AU
	16	+#define MSPROF_TASK_TIME_L0 0x00000800ULL
	17	+
	18	+typedef const void *ConstVoidPtr;
	19	+
	20	+#ifdef __cplusplus
	21	+}
	22	+#endif
	23	+
	24	+#endif

Athird_party/dvm/CMakeLists.txt+31-0

@@ -0,0 +1,31 @@
	1	+set(DVM_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/dvm)
	2	+set(DVM_LIB ${DVM_ROOT}/libdvm.a)
	3	+
	4	+set(ACL_INC_ROOT "${TORCHNPU_THIRD_PARTY_ROOT}/acl/inc")
	5	+set(DVM_MAKE_ARGS
	6	+ PRE_ASCEND=1
	7	+ DVM_CUSTOM_FLAGS=-I${ACL_INC_ROOT}
	8	+)
	9	+
	10	+file(GLOB DVM_DEPENDS
	11	+ CONFIGURE_DEPENDS
	12	+ "${DVM_ROOT}/Makefile"
	13	+ "${DVM_ROOT}/env.sh"
	14	+ "${DVM_ROOT}/src/*.cc"
	15	+ "${DVM_ROOT}/src/*.h"
	16	+ "${DVM_ROOT}/include/*.h"
	17	+)
	18	+
	19	+add_custom_command(
	20	+ OUTPUT ${DVM_LIB}
	21	+ COMMAND make -C ${DVM_ROOT} ${DVM_MAKE_ARGS} libdvm.a
	22	+ WORKING_DIRECTORY ${DVM_ROOT}
	23	+ DEPENDS ${DVM_DEPENDS}
	24	+ COMMENT "Building DVM static library"
	25	+ VERBATIM
	26	+)
	27	+
	28	+add_custom_target(dvm_build DEPENDS ${DVM_LIB})
	29	+add_library(dvm STATIC IMPORTED GLOBAL)
	30	+set_target_properties(dvm PROPERTIES IMPORTED_LOCATION ${DVM_LIB})
	31	+add_dependencies(dvm dvm_build)

Athird_party/dvm/dvm+1-0

	@@ -0,0 +1 @@
		1	+Subproject commit edefdc04ec22d14fe578c3e7dc63cd7c54e20632

Mtorch_npu/_inductor/__init__.py+14-0

	@@ -32,6 +32,19 @@ def _load_mlir_backend():
32	apply_mlir_inductor_patch()	32	apply_mlir_inductor_patch()
33	register_mlir_codegen_backend()	33	register_mlir_codegen_backend()
34		34
		35	+
		36	+def _load_dvm_backend():
		37	+ import torch
		38	+ from .ascend_npu_ir.ascend_npu_ir.npu import npu_inductor_plugin
		39	+ from .lowering_patch import apply_mlir_inductor_patch
		40	+ from .ascend_npu_ir.ascend_npu_ir.npu.npu_inductor_plugin import (
		41	+ register_mlir_codegen_backend,
		42	+ )
		43	+ apply_mlir_inductor_patch()
		44	+ register_mlir_codegen_backend()
		45	+ from .dvm import mlir_fusion
		46	+
		47	+
35	def _load_triton_backend():	48	def _load_triton_backend():
36	import os	49	import os
37	import torch	50	import torch
	@@ -158,6 +171,7 @@ def _load_triton_backend():
158		171
159	_BACKEND_LOADERS = {	172	_BACKEND_LOADERS = {
160	"mlir": _load_mlir_backend,	173	"mlir": _load_mlir_backend,
		174	+ "dvm": _load_dvm_backend,
161	"default": _load_triton_backend,	175	"default": _load_triton_backend,
162	}	176	}
163		177

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/codecache.py+0-1

	@@ -218,7 +218,6 @@ class CustomAsyncCompile(AsyncCompile):
218	assert get_compile_threads() > 1	218	assert get_compile_threads() > 1
219	# Wrapper around ProcessPoolExecutor forks in a new process we control	219	# Wrapper around ProcessPoolExecutor forks in a new process we control
220	log.info("Creating subprocess pool with %d workers", get_compile_threads())	220	log.info("Creating subprocess pool with %d workers", get_compile_threads())
221	- os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
222	pool = SubprocPool(get_compile_threads())	221	pool = SubprocPool(get_compile_threads())
223		222
224	# Set an attribute we can check to see if the pool is ready.	223	# Set an attribute we can check to see if the pool is ready.

Mtorch_npu/_inductor/ascend_npu_ir/ascend_npu_ir/npu/mlir_compiler.py+0-2

	@@ -32,8 +32,6 @@ akg_spec = importlib.util.find_spec("akg")
32	if akg_spec is not None:	32	if akg_spec is not None:
33	from akg.kernel import Kernel as MlirKernel	33	from akg.kernel import Kernel as MlirKernel
34		34
35	-os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
36	-
37		35
38	reinterpret_tensor = torch.ops.inductor._reinterpret_tensor	36	reinterpret_tensor = torch.ops.inductor._reinterpret_tensor
39	global_cache = set()	37	global_cache = set()

Atorch_npu/_inductor/dvm/__init__.py+198-0

@@ -0,0 +1,198 @@
	1	+import os
	2	+import sys
	3	+from functools import partial, wraps
	4	+
	5	+import torch
	6	+
	7	+import torch_npu
	8	+from torch_npu._C.dvm import (
	9	+ NDObject,
	10	+ DataType,
	11	+ DynKernel,
	12	+ GraphSplitKernel,
	13	+ DynGraphSplitKernel,
	14	+ TorchKernel as Kernel,
	15	+)
	16	+from torch_npu._inductor.npu_compare import check_accuracy_dvm
	17	+from torch_npu.npu._backends import get_soc_version
	18	+
	19	+
	20	+debug_mode = os.environ.get("INDUCTOR_DVM_DEBUG_MODE", "0") == "1"
	21	+bf16_vector_keep_promoted = False
	22	+
	23	+bool_ = DataType.bool
	24	+float16 = DataType.float16
	25	+bfloat16 = DataType.bfloat16
	26	+float32 = DataType.float32
	27	+int32 = DataType.int32
	28	+int64 = DataType.int64
	29	+
	30	+Ascend910B1 = 220
	31	+Ascend310B1 = 240
	32	+Ascend910_9391 = 250
	33	+Ascend950 = 260
	34	+is_ascend950 = get_soc_version() >= Ascend950
	35	+
	36	+KERNEL_FACTORY = {
	37	+ ("mix", True): partial(DynKernel, Kernel.K_MIX, Kernel.F_DYN),
	38	+ ("mix", False): partial(Kernel, Kernel.K_MIX, 0),
	39	+ ("split", True): DynGraphSplitKernel,
	40	+ ("split", False): GraphSplitKernel,
	41	+ ("spec", True): partial(
	42	+ DynKernel, Kernel.K_VEC, Kernel.F_DYN \| Kernel.F_SPEC
	43	+ ),
	44	+ ("spec", False): partial(Kernel, Kernel.K_VEC, Kernel.F_SPEC),
	45	+ ("vector", True): partial(DynKernel, Kernel.K_VEC, Kernel.F_DYN),
	46	+ ("vector", False): partial(Kernel, Kernel.K_VEC, 0),
	47	+}
	48	+
	49	+
	50	+def kernel(
	51	+ ktype: str = "split",
	52	+ dyn_shape: bool = False,
	53	+):
	54	+ r"""kernel(ktype="split", dyn_shape=False)
	55	+
	56	+ Return a decorator that builds and executes a DVM kernel.
	57	+
	58	+ The decorated function ``builder(kobj)`` constructs the kernel object.
	59	+ The returned callable supports two execution styles:
	60	+
	61	+ - ``fn(args, *kwargs)``: takes input tensors only and returns output tensor(s).
	62	+ - ``fn.run(args, *kwargs)``: takes input and output tensors, writes results
	63	+ into the provided output buffers, and returns ``None``.
	64	+
	65	+ Args:
	66	+ ktype (str, optional): kernel type. Default: ``"split"``.
	67	+ dyn_shape (bool, optional): enable dynamic shapes. Default: ``False``.
	68	+
	69	+ Returns:
	70	+ Callable: a callable kernel wrapper with a ``run`` method and ``kobj`` attribute.
	71	+ """
	72	+
	73	+ def decorate(builder):
	74	+ kobj = KERNEL_FACTORY[(ktype, dyn_shape)]()
	75	+ kernel_name = getattr(builder, "__name__", "<unknown>")
	76	+
	77	+ builder(kobj)
	78	+ kobj.setup()
	79	+
	80	+ def _format_args(args):
	81	+ dump_parts = []
	82	+ arg_summaries = []
	83	+ for a in args:
	84	+ if isinstance(a, torch.Tensor):
	85	+ shape = tuple(a.shape)
	86	+ dump_parts.append(str(a.shape))
	87	+ arg_summaries.append(
	88	+ f"Tensor(shape={shape}, dtype={a.dtype}, device={a.device})"
	89	+ )
	90	+ elif isinstance(a, (torch.SymInt, torch.SymFloat)):
	91	+ sym_name = type(a).__name__
	92	+ dump_parts.append(sym_name)
	93	+ arg_summaries.append(sym_name)
	94	+ else:
	95	+ arg_summaries.append(type(a).__name__)
	96	+ return dump_parts, arg_summaries
	97	+
	98	+ def _post_run(args):
	99	+ try:
	100	+ torch_npu.npu.synchronize()
	101	+ except Exception as exc:
	102	+ dump_text = kobj.dump()
	103	+ das_text = kobj.das()
	104	+ dump_parts, arg_summaries = _format_args(args)
	105	+ dump_id = ",".join(dump_parts)
	106	+ msg = [
	107	+ "DVM debug sync failed.",
	108	+ f"kernel_name={kernel_name}",
	109	+ f"dump_id={dump_id}",
	110	+ f"args={arg_summaries}",
	111	+ f"dump={dump_text}",
	112	+ f"das={das_text}",
	113	+ f"error={type(exc).__name__}: {exc}",
	114	+ ]
	115	+ print("\n".join(msg), file=sys.stderr)
	116	+ raise
	117	+
	118	+ @wraps(builder)
	119	+ def fn(args, *kwargs):
	120	+ outputs = kobj(*args)
	121	+ if debug_mode:
	122	+ _post_run(args)
	123	+ return outputs
	124	+
	125	+ def run(args, *kwargs):
	126	+ kobj.run(*args)
	127	+ if debug_mode:
	128	+ _post_run(args)
	129	+
	130	+ fn.run = run
	131	+ fn.kobj = kobj
	132	+
	133	+ return fn
	134	+
	135	+ return decorate
	136	+
	137	+
	138	+def _install_bf16_promote():
	139	+ unsupported_bf16_ops = (
	140	+ "sqrt",
	141	+ "abs",
	142	+ "log",
	143	+ "exp",
	144	+ "reciprocal",
	145	+ "logical_not",
	146	+ "round",
	147	+ "floor",
	148	+ "ceil",
	149	+ "trunc",
	150	+ "equal",
	151	+ "not_equal",
	152	+ "greater",
	153	+ "greater_equal",
	154	+ "less",
	155	+ "less_equal",
	156	+ "add",
	157	+ "sub",
	158	+ "mul",
	159	+ "div",
	160	+ "pow",
	161	+ "maximum",
	162	+ "minimum",
	163	+ "logical_and",
	164	+ "logical_or",
	165	+ "select",
	166	+ "sum",
	167	+ "max",
	168	+ "min",
	169	+ )
	170	+
	171	+ def _promote_bf16(op_fn):
	172	+ @wraps(op_fn)
	173	+ def wrapper(self, *args):
	174	+ need_cast_back = False
	175	+
	176	+ def maybe_cast_arg(x):
	177	+ nonlocal need_cast_back
	178	+ if isinstance(x, NDObject) and x.dtype() == bfloat16:
	179	+ need_cast_back = True
	180	+ return self.cast(x, float32)
	181	+ return x
	182	+
	183	+ new_args = tuple(maybe_cast_arg(a) for a in args)
	184	+ out = op_fn(self, *new_args)
	185	+ if bf16_vector_keep_promoted and need_cast_back :
	186	+ out = self.cast(out, bfloat16)
	187	+ return out
	188	+
	189	+ return wrapper
	190	+
	191	+ for name in unsupported_bf16_ops:
	192	+ op_fn = getattr(Kernel, name, None)
	193	+ if op_fn is not None:
	194	+ setattr(Kernel, name, _promote_bf16(op_fn))
	195	+
	196	+
	197	+if not is_ascend950:
	198	+ _install_bf16_promote()

Atorch_npu/_inductor/dvm/decomp.py+171-0

@@ -0,0 +1,171 @@
	1	+import math
	2	+
	3	+import torch
	4	+from torch._decomp import remove_decompositions
	5	+from torch._inductor import decomposition as inductor_decomp
	6	+
	7	+
	8	+aten = torch.ops.aten
	9	+prims = torch.ops.prims
	10	+quantized = torch.ops.quantized
	11	+
	12	+decomps_to_exclude_npu = [
	13	+ aten._batch_norm_no_update,
	14	+ aten._batch_norm_with_update,
	15	+ aten._batch_norm_with_update_functional,
	16	+ aten._log_softmax,
	17	+ aten._log_softmax_backward_data,
	18	+ aten.batch_norm_backward,
	19	+ aten.convolution_backward,
	20	+ aten.embedding,
	21	+ aten.embedding_backward,
	22	+ aten.embedding_dense_backward,
	23	+ aten.gelu.default,
	24	+ aten.gelu_backward.default,
	25	+ aten.grid_sampler_2d,
	26	+ aten.grid_sampler_2d_backward,
	27	+ aten.linalg_vector_norm,
	28	+ aten.max_pool2d_with_indices,
	29	+ aten.max_pool2d_with_indices_backward,
	30	+ aten.native_batch_norm,
	31	+ aten.native_group_norm,
	32	+ aten.nll_loss2d_backward,
	33	+ aten.nll_loss2d_forward,
	34	+ aten.nll_loss_backward,
	35	+ aten.nll_loss_forward,
	36	+ aten.reflection_pad2d,
	37	+ aten.reflection_pad2d_backward,
	38	+ aten.slice.Tensor,
	39	+ aten.triu,
	40	+ aten.upsample_bilinear2d,
	41	+ aten.upsample_bilinear2d_backward,
	42	+ aten.upsample_nearest1d,
	43	+ aten.upsample_nearest1d_backward,
	44	+ aten.upsample_nearest2d,
	45	+ aten.upsample_nearest2d_backward,
	46	+ aten.upsample_nearest3d,
	47	+ aten.upsample_nearest3d_backward,
	48	+ torch.ops.npu.npu_rotary_mul,
	49	+ torch.ops.npu.npu_rotary_mul_backward,
	50	+]
	51	+
	52	+FP32_MIN_V2 = -8.8
	53	+FP32_MAX_V2 = 8.8
	54	+DOUBLE_X = 2.0
	55	+
	56	+
	57	+def tanh(a):
	58	+ """
	59	+ y = (exp(2x) - 1) / (exp(2x) + 1)
	60	+ with x clipped to [-8.8, 8.8] in float32 before multiply-by-2.
	61	+ """
	62	+ orig_dtype = a.dtype
	63	+ if orig_dtype != torch.float32:
	64	+ a = a.to(torch.float32)
	65	+ x = torch.clamp(a, min=FP32_MIN_V2, max=FP32_MAX_V2)
	66	+ x2 = x * DOUBLE_X
	67	+ e2x = torch.exp(x2)
	68	+ out = (e2x - 1.0) / (e2x + 1.0)
	69	+
	70	+ if orig_dtype != torch.float32:
	71	+ out = out.to(orig_dtype)
	72	+ return out
	73	+
	74	+
	75	+def gelu(a: torch.Tensor, approximate: str = "none"):
	76	+ """
	77	+ y = -sqrt(8/pi) * (x + 0.044715 * x^3)
	78	+ out = x / (1 + exp(y))
	79	+ """
	80	+ orig_dtype = a.dtype
	81	+ if orig_dtype != torch.float32:
	82	+ a = a.to(torch.float32)
	83	+
	84	+ M_SQRT2 = math.sqrt(2)
	85	+ M_2_SQRTPI = 2.0 / math.sqrt(math.pi)
	86	+ kBeta = M_SQRT2 * M_2_SQRTPI
	87	+ kKappa = 0.044715
	88	+
	89	+ a_cube = a * a * a
	90	+ inner = a + kKappa * a_cube
	91	+ y = -kBeta * inner
	92	+ out = a / (1.0 + torch.exp(y))
	93	+
	94	+ if orig_dtype != torch.float32:
	95	+ out = out.to(orig_dtype)
	96	+ return out
	97	+
	98	+
	99	+def gelu_backward(grad, self, approximate: str = "none"):
	100	+ orig_dtype = grad.dtype
	101	+ if orig_dtype != torch.float32:
	102	+ grad = grad.to(torch.float32)
	103	+ self = self.to(torch.float32)
	104	+ M_SQRT2 = math.sqrt(2)
	105	+ M_2_SQRTPI = 2.0 / math.sqrt(math.pi)
	106	+ kBeta = M_SQRT2 * M_2_SQRTPI * 0.5
	107	+ kKappa = 0.044715
	108	+ x_sq = self * self
	109	+ x_cube = x_sq * self
	110	+ inner = kBeta * (self + kKappa * x_cube)
	111	+ tanh_inner = torch.tanh(inner)
	112	+
	113	+ left = 0.5 * self
	114	+ right = 1.0 + tanh_inner
	115	+
	116	+ left_derivative = 0.5 * right
	117	+
	118	+ tanh_derivative = (tanh_inner * tanh_inner) * -1.0 + 1.0
	119	+ inner_derivative = kBeta * (1.0 + 3.0 * kKappa * x_sq)
	120	+ right_derivative = left * tanh_derivative * inner_derivative
	121	+ out = grad * (left_derivative + right_derivative)
	122	+
	123	+ if orig_dtype != torch.float32:
	124	+ out = out.to(orig_dtype)
	125	+ return out
	126	+
	127	+
	128	+def sigmoid(a: torch.Tensor) -> torch.Tensor:
	129	+ orig_dtype = a.dtype
	130	+ if orig_dtype != torch.float32:
	131	+ a = a.to(torch.float32)
	132	+ out = 1 / (1.0 + torch.exp(torch.neg(a)))
	133	+ if orig_dtype != torch.float32:
	134	+ out = out.to(orig_dtype)
	135	+ return out
	136	+
	137	+
	138	+_dvm_inductor_decomp_patched = False
	139	+
	140	+
	141	+def _register_inductor_decomposition_safe(overloads, fn):
	142	+ """Register a custom Inductor decomposition; ignore duplicate registration."""
	143	+ try:
	144	+ inductor_decomp.register_decomposition(overloads)(fn)
	145	+ except (RuntimeError, ValueError) as e:
	146	+ msg = str(e).lower()
	147	+ if any(
	148	+ s in msg
	149	+ for s in (
	150	+ "duplicate",
	151	+ "already",
	152	+ "exists",
	153	+ "re-register",
	154	+ "re_register",
	155	+ )
	156	+ ):
	157	+ return
	158	+ raise
	159	+
	160	+
	161	+def patch_decomp():
	162	+ """Patch Inductor decomposition for DVM paths (idempotent)."""
	163	+ global _dvm_inductor_decomp_patched
	164	+ if _dvm_inductor_decomp_patched:
	165	+ return
	166	+ remove_decompositions(inductor_decomp.decompositions, decomps_to_exclude_npu)
	167	+ _register_inductor_decomposition_safe([aten.sigmoid.default], sigmoid)
	168	+ _register_inductor_decomposition_safe([aten.gelu_backward.default], gelu_backward)
	169	+ _register_inductor_decomposition_safe([aten.gelu.default], gelu)
	170	+ _register_inductor_decomposition_safe([aten.tanh.default], tanh)
	171	+ _dvm_inductor_decomp_patched = True

Atorch_npu/_inductor/dvm/fx_pass.py+251-0

@@ -0,0 +1,251 @@
	1	+from dataclasses import dataclass
	2	+
	3	+import torch
	4	+from torch.fx import GraphModule, Node
	5	+from torch._subclasses import FakeTensor
	6	+from torch._prims_common import elementwise_dtypes, ELEMENTWISE_TYPE_PROMOTION_KIND
	7	+from .op_emitter import _is_last2_transpose_tensor
	8	+
	9	+
	10	+aten = torch.ops.aten
	11	+prims = torch.ops.prims
	12	+
	13	+
	14	+def need_fallback_gm(gm: torch.fx.GraphModule) -> bool:
	15	+ for node in gm.graph.nodes:
	16	+ if node.op != "call_function":
	17	+ continue
	18	+ if node.target not in (
	19	+ aten.reshape.default,
	20	+ aten.expand.default,
	21	+ ):
	22	+ return False
	23	+ return True
	24	+
	25	+
	26	+def annotate_mm_transpose_flags(gm: torch.fx.GraphModule):
	27	+ flag = False
	28	+ for node in gm.graph.nodes:
	29	+ if node.op != "call_function":
	30	+ continue
	31	+ if node.target in [aten.mm.default, aten.bmm.default]:
	32	+ lhs = node.args[0]
	33	+ rhs = node.args[1]
	34	+ flag = True
	35	+ elif node.target is aten.addmm.default:
	36	+ add = node.args[0]
	37	+ lhs = node.args[1]
	38	+ rhs = node.args[2]
	39	+ if (
	40	+ add.meta["val"].dim() == 1
	41	+ and node.kwargs.get("beta", 1) == 1
	42	+ and node.kwargs.get("alpha", 1) == 1
	43	+ ):
	44	+ node.meta["use_bias"] = True
	45	+ flag = True
	46	+ else:
	47	+ continue
	48	+ node.meta["trans_a"] = False
	49	+ node.meta["trans_b"] = False
	50	+ if lhs.op == "placeholder" and _is_last2_transpose_tensor(lhs.meta["val"]):
	51	+ node.meta["trans_a"] = True
	52	+ lhs.meta["trans"] = True
	53	+ if rhs.op == "placeholder" and _is_last2_transpose_tensor(rhs.meta["val"]):
	54	+ node.meta["trans_b"] = True
	55	+ rhs.meta["trans"] = True
	56	+
	57	+ return flag
	58	+
	59	+
	60	+def make_cast_node(g, src: Node, target_dtype: torch.dtype) -> Node:
	61	+ cast = g.call_function(
	62	+ prims.convert_element_type.default,
	63	+ args=(src, target_dtype),
	64	+ )
	65	+
	66	+ cast.meta["val"] = src.meta["val"].to(dtype=target_dtype)
	67	+ return cast
	68	+
	69	+
	70	+def decompose_k1_matmul_to_mul(gm: GraphModule) -> GraphModule:
	71	+ g = gm.graph
	72	+ changed = False
	73	+
	74	+ for node in list(g.nodes):
	75	+ if node.op != "call_function":
	76	+ continue
	77	+ if node.target not in (aten.mm.default, aten.bmm.default):
	78	+ continue
	79	+
	80	+ lhs, rhs = node.args[:2]
	81	+
	82	+ if not isinstance(lhs, Node) or not isinstance(rhs, Node):
	83	+ continue
	84	+
	85	+ lhs_val = lhs.meta.get("val", None)
	86	+ rhs_val = rhs.meta.get("val", None)
	87	+ if not isinstance(lhs_val, FakeTensor) or not isinstance(rhs_val, FakeTensor):
	88	+ continue
	89	+
	90	+ if _is_last2_transpose_tensor(lhs_val):
	91	+ continue
	92	+ if _is_last2_transpose_tensor(rhs_val):
	93	+ continue
	94	+
	95	+ lhs_k = lhs_val.shape[-1]
	96	+ rhs_k = rhs_val.shape[-2]
	97	+
	98	+ if isinstance(lhs_k, torch.SymInt) or isinstance(rhs_k, torch.SymInt):
	99	+ continue
	100	+ if lhs_k != 1 or rhs_k != 1:
	101	+ continue
	102	+
	103	+ with g.inserting_before(node):
	104	+ mul_node = g.call_function(aten.mul.Tensor, args=(lhs, rhs))
	105	+ mul_node.meta["val"] = node.meta["val"]
	106	+
	107	+ node.replace_all_uses_with(mul_node)
	108	+ g.erase_node(node)
	109	+ changed = True
	110	+
	111	+ if changed:
	112	+ g.lint()
	113	+ gm.recompile()
	114	+ return gm
	115	+
	116	+
	117	+def insert_sum_fp32_prepost_cast_prims(gm: GraphModule):
	118	+ g = gm.graph
	119	+ for node in g.nodes:
	120	+ if node.op != "call_function":
	121	+ continue
	122	+ if node.target not in [aten.sum.default, aten.sum.dim_IntList]:
	123	+ continue
	124	+
	125	+ out_val = node.meta.get("val", None)
	126	+ if not isinstance(out_val, FakeTensor):
	127	+ continue
	128	+ orig_out_dtype = out_val.dtype
	129	+
	130	+ if not node.args:
	131	+ continue
	132	+ x = node.args[0]
	133	+ if not isinstance(x, Node):
	134	+ continue
	135	+
	136	+ in_val = x.meta.get("val", None)
	137	+ if not isinstance(in_val, FakeTensor):
	138	+ continue
	139	+ in_dtype = in_val.dtype
	140	+
	141	+ if in_dtype == torch.float32:
	142	+ continue
	143	+
	144	+ with g.inserting_before(node):
	145	+ x_fp32 = make_cast_node(g, x, torch.float32)
	146	+
	147	+ new_args = list(node.args)
	148	+ new_args[0] = x_fp32
	149	+ node.args = tuple(new_args)
	150	+
	151	+ if orig_out_dtype != torch.float32:
	152	+ with g.inserting_after(node):
	153	+ y = make_cast_node(g, node, orig_out_dtype)
	154	+ node.replace_all_uses_with(y)
	155	+ y.args = (node, orig_out_dtype)
	156	+
	157	+ g.lint()
	158	+ gm.recompile()
	159	+ return gm
	160	+
	161	+
	162	+@dataclass(frozen=True)
	163	+class PromoteRule:
	164	+ pos: tuple[int, ...]
	165	+ kind: object = ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
	166	+
	167	+
	168	+PROMOTE_TYPE_OP = {
	169	+ # ========= Elementwise =========
	170	+ aten.add.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	171	+ aten.sub.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	172	+ aten.mul.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	173	+ aten.div.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	174	+ aten.pow.Tensor_Tensor: PromoteRule(
	175	+ (0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
	176	+ ),
	177	+ aten.lt.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL),
	178	+ aten.le.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL),
	179	+ aten.gt.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL),
	180	+ aten.ge.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL),
	181	+ aten.eq.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL),
	182	+ aten.ne.Tensor: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL),
	183	+ aten.maximum.default: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	184	+ aten.minimum.default: PromoteRule((0, 1), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	185	+ # ========= Where =========
	186	+ # where(cond, x, y) → promote x/y
	187	+ # aten.where.default: PromoteRule((1, 2), ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT),
	188	+}
	189	+
	190	+
	191	+def insert_promote_cast_by_pos_prims(gm: GraphModule) -> GraphModule:
	192	+ g = gm.graph
	193	+
	194	+ for node in g.nodes:
	195	+ if node.op != "call_function":
	196	+ continue
	197	+
	198	+ rule = PROMOTE_TYPE_OP.get(node.target, None)
	199	+ if rule is None:
	200	+ continue
	201	+
	202	+ arg_vals = []
	203	+ arg_nodes = {}
	204	+
	205	+ for idx in rule.pos:
	206	+ if idx >= len(node.args):
	207	+ continue
	208	+ arg = node.args[idx]
	209	+ if not isinstance(arg, Node):
	210	+ continue
	211	+ val = arg.meta.get("val", None)
	212	+ if not isinstance(val, FakeTensor):
	213	+ continue
	214	+ arg_vals.append(val)
	215	+ arg_nodes[idx] = arg
	216	+
	217	+ if len(arg_vals) <= 1:
	218	+ continue
	219	+
	220	+ dtypes = [v.dtype for v in arg_vals]
	221	+ if all(dt == dtypes[0] for dt in dtypes[1:]):
	222	+ continue
	223	+
	224	+ compute_dtype, _ = elementwise_dtypes(
	225	+ *arg_vals,
	226	+ type_promotion_kind=rule.kind,
	227	+ )
	228	+
	229	+ new_args = list(node.args)
	230	+ for idx, arg in arg_nodes.items():
	231	+ in_val = arg.meta.get("val", None)
	232	+ if in_val.dtype == compute_dtype:
	233	+ continue
	234	+ with g.inserting_before(node):
	235	+ cast = make_cast_node(g, arg, compute_dtype)
	236	+ new_args[idx] = cast
	237	+
	238	+ node.args = tuple(new_args)
	239	+
	240	+ g.lint()
	241	+ gm.recompile()
	242	+ return gm
	243	+
	244	+
	245	+def expand_to_reshape(gm: GraphModule) -> GraphModule:
	246	+ for node in gm.graph.find_nodes(op="call_function", target=aten.expand.default):
	247	+ x = node.args[0]
	248	+ in_val = x.meta.get("val")
	249	+ out_val = node.meta.get("val")
	250	+ if tuple(in_val.shape) == tuple(out_val.shape):
	251	+ node.target = aten.reshape.default

Atorch_npu/_inductor/dvm/fx_test.py+145-0

@@ -0,0 +1,145 @@
	1	+import os
	2	+import hashlib
	3	+import torch
	4	+
	5	+
	6	+DEFAULT_OUTPUT_DIR = "./dvm_fx_regression_cases"
	7	+
	8	+
	9	+def generate_dvm_fx_case(
	10	+ gm: torch.fx.GraphModule,
	11	+ output_dir: str = DEFAULT_OUTPUT_DIR,
	12	+ fusion_type: str = "graph",
	13	+):
	14	+ if fusion_type not in ("graph", "mlir"):
	15	+ raise ValueError(f"unsupported fusion_type: {fusion_type}")
	16	+
	17	+ def _indent(code: str, spaces: int) -> str:
	18	+ pad = " " * spaces
	19	+ return "\n".join(
	20	+ pad + line if line.strip() else line
	21	+ for line in code.split("\n")
	22	+ )
	23	+
	24	+ os.makedirs(output_dir, exist_ok=True)
	25	+
	26	+ readable = gm.print_readable(print_output=False)
	27	+
	28	+ sig = []
	29	+ inputs = [n for n in gm.graph.nodes if n.op == "placeholder"]
	30	+ for i, n in enumerate(inputs):
	31	+ v = n.meta["val"]
	32	+ sig.append(f"arg{i}:{tuple(v.shape)},{tuple(v.stride())},{v.dtype}")
	33	+
	34	+ h = hashlib.sha256(
	35	+ (fusion_type + "\n" + readable + "\n" + "\n".join(sig)).encode("utf-8")
	36	+ ).hexdigest()[:16]
	37	+ case_name = f"test_{fusion_type}_{h}"
	38	+
	39	+ file_path = os.path.join(output_dir, f"{case_name}.py")
	40	+ if os.path.exists(file_path):
	41	+ print(f"[skip] {file_path}")
	42	+ return None
	43	+ class_name = "TestModel"
	44	+
	45	+ input_lines = []
	46	+ for i, n in enumerate(inputs):
	47	+ v = n.meta["val"]
	48	+ fill = "random_()" if v.dtype == torch.bool else "uniform_(0, 1)"
	49	+ input_lines.append(
	50	+ f"arg{i} = torch.empty_strided("
	51	+ f"torch.Size({tuple(v.shape)}), "
	52	+ f"{tuple(v.stride())}, "
	53	+ f"dtype={v.dtype}, device='npu').{fill}"
	54	+ )
	55	+
	56	+ input_code = "\n ".join(input_lines)
	57	+ fwd_args = ", ".join(f"arg{i}" for i in range(len(inputs)))
	58	+
	59	+ if fusion_type == "graph":
	60	+ fusion_env = ""
	61	+ fusion_imports = (
	62	+ "from torch_npu._inductor.dvm.graph_fusion "
	63	+ "import DvmGraphFusionPatch"
	64	+ )
	65	+ compile_lines = [
	66	+ "with DvmGraphFusionPatch():",
	67	+ " compiled = torch.compile(model, backend=\"inductor\", dynamic=False)",
	68	+ f" out = compiled({fwd_args})",
	69	+ " deterministic_state = torch.are_deterministic_algorithms_enabled()",
	70	+ " deterministic_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()",
	71	+ " try:",
	72	+ " torch.use_deterministic_algorithms(True)",
	73	+ " deterministic_compiled = torch.compile(model, backend=\"inductor\", dynamic=False)",
	74	+ f" deterministic_out = deterministic_compiled({fwd_args})",
	75	+ " finally:",
	76	+ " torch.use_deterministic_algorithms(deterministic_state, warn_only=deterministic_warn_only)",
	77	+ ]
	78	+ else:
	79	+ fusion_env = 'os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "dvm"'
	80	+ fusion_imports = "from torch_npu._inductor.dvm import mlir_fusion"
	81	+ compile_lines = [
	82	+ "compiled = torch.compile(model, backend=\"inductor\", dynamic=False)",
	83	+ f"out = compiled({fwd_args})",
	84	+ "deterministic_state = torch.are_deterministic_algorithms_enabled()",
	85	+ "deterministic_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()",
	86	+ "try:",
	87	+ " torch.use_deterministic_algorithms(True)",
	88	+ " deterministic_compiled = torch.compile(model, backend=\"inductor\", dynamic=False)",
	89	+ f" deterministic_out = deterministic_compiled({fwd_args})",
	90	+ "finally:",
	91	+ " torch.use_deterministic_algorithms(deterministic_state, warn_only=deterministic_warn_only)",
	92	+ ]
	93	+ compile_code = "\n ".join(compile_lines)
	94	+ env_lines = fusion_env
	95	+
	96	+ test_code = f"""import os
	97	+
	98	+os.environ["INDUCTOR_DVM_DUMP_FX_TEST"] = "0"
	99	+{env_lines}
	100	+
	101	+import torch
	102	+import torch_npu
	103	+from torch import device, tensor
	104	+from math import inf, nan
	105	+from torch.utils._pytree import tree_flatten
	106	+
	107	+
	108	+class {class_name}(torch.nn.Module):
	109	+ def __init__(self):
	110	+ super().__init__()
	111	+{_indent(gm.code, 4)}
	112	+
	113	+
	114	+def _assert_close(ref, out, atol=2e-3, rtol=2e-3):
	115	+ rf, rs = tree_flatten(ref)
	116	+ of, os = tree_flatten(out)
	117	+ assert rs == os, f"pytree mismatch\\nref={{rs}}\\nout={{os}}"
	118	+ for r, o in zip(rf, of):
	119	+ torch.testing.assert_close(r, o, atol=atol, rtol=rtol, equal_nan=True)
	120	+
	121	+
	122	+def test_case():
	123	+ {input_code}
	124	+
	125	+ model = {class_name}().npu()
	126	+ ref = model({fwd_args})
	127	+
	128	+ {fusion_imports}
	129	+
	130	+ {compile_code}
	131	+
	132	+ _assert_close(ref, out)
	133	+ _assert_close(ref, deterministic_out)
	134	+
	135	+
	136	+if __name__ == "__main__":
	137	+ test_case()
	138	+ print("PASS")
	139	+"""
	140	+
	141	+ with open(file_path, "w", encoding="utf-8") as f:
	142	+ f.write(test_code)
	143	+
	144	+ print(f"[ok] generated: {file_path}")
	145	+ return file_path

Atorch_npu/_inductor/dvm/graph_build.py+165-0

@@ -0,0 +1,165 @@
	1	+import torch
	2	+import torch.utils._pytree as pytree
	3	+from torch._inductor.utils import IndentedBuffer
	4	+from torch.fx.node import Argument, Target
	5	+
	6	+from .fx_pass import annotate_mm_transpose_flags
	7	+from .op_emitter import DVM_OP_REGISTRY, load, store
	8	+from .util import codegen_maybe_view_load
	9	+
	10	+
	11	+aten = torch.ops.aten
	12	+
	13	+
	14	+def is_fx_dynamic(graph):
	15	+ for node in graph.graph.nodes:
	16	+ if node.op == "placeholder" or node.op == "call_function":
	17	+ val = node.meta.get("val")
	18	+ if val is None:
	19	+ continue
	20	+ if isinstance(val, torch.Tensor):
	21	+ if any(isinstance(dim, torch.SymInt) for dim in val.shape):
	22	+ return True
	23	+ elif isinstance(val, (torch.SymInt, torch.SymFloat)):
	24	+ return True
	25	+ return False
	26	+
	27	+
	28	+class DvmCodegenInterpreter(torch.fx.Interpreter):
	29	+ KERNEL_NAME_PLACEHOLDER = "__DVM_KERNEL_NAME__"
	30	+
	31	+ def __init__(
	32	+ self,
	33	+ gm: torch.fx.GraphModule,
	34	+ ktype: str,
	35	+ view_fusion_level=1,
	36	+ is_dynamic: bool \| None = None,
	37	+ ):
	38	+ super().__init__(gm)
	39	+ self.gm = gm
	40	+ self.ktype = ktype
	41	+ self.is_mix_kernel = annotate_mm_transpose_flags(gm)
	42	+ if is_dynamic is None:
	43	+ self.is_dynamic = is_fx_dynamic(gm)
	44	+ else:
	45	+ if not isinstance(is_dynamic, bool):
	46	+ raise TypeError("is_dynamic must be bool when provided")
	47	+ self.is_dynamic = is_dynamic
	48	+ self.current_node = None
	49	+ self.cont_flag_input = []
	50	+ self.need_trans_input = []
	51	+ self.view_fusion_level = view_fusion_level
	52	+ self.code = IndentedBuffer()
	53	+
	54	+ self.spec_nodes = set()
	55	+ if self.ktype == "vector" and self.need_spec():
	56	+ self.ktype = "spec"
	57	+ self.code.splice(f'\n"""\n{self.gm.print_readable(print_output=False)}\n"""')
	58	+ decorator = (
	59	+ f"{chr(64)}dvm.kernel(ktype={self.ktype!r}, dyn_shape={self.is_dynamic})"
	60	+ )
	61	+ self.code.splice(decorator)
	62	+ self.code.splice(f"def {self.KERNEL_NAME_PLACEHOLDER}(k):")
	63	+ self.code.do_indent()
	64	+
	65	+ def need_spec(self) -> bool:
	66	+ self.spec_nodes.clear()
	67	+ for node in self.gm.graph.nodes:
	68	+ if node.op != "call_function":
	69	+ continue
	70	+ for input_node in node.all_input_nodes:
	71	+ if input_node.op == "call_function" and input_node.target in [
	72	+ aten.sum.default,
	73	+ aten.sum.dim_IntList,
	74	+ aten.amax.default,
	75	+ aten.amin.default,
	76	+ ]:
	77	+ self.spec_nodes.add(input_node)
	78	+ return len(self.spec_nodes) > 0
	79	+
	80	+ def run_node(self, n: torch.fx.Node) -> Argument:
	81	+ self.current_node = n
	82	+ expr = super().run_node(n)
	83	+ if n.op == "output":
	84	+ for _expr in pytree.tree_leaves(expr):
	85	+ self.code.splice(f"{_expr}")
	86	+ else:
	87	+ self.code.splice(f"{n} = {expr}")
	88	+ if n in self.spec_nodes:
	89	+ self.code.splice("k.spec_next()")
	90	+ return f"{n}"
	91	+
	92	+ def placeholder(
	93	+ self, target: "Target", args: tuple[Argument], kwargs: dict[str, Argument]
	94	+ ) -> Argument:
	95	+ meta = self.current_node.meta
	96	+ val = meta["val"]
	97	+ if isinstance(val, torch.SymInt):
	98	+ self.cont_flag_input.append(True)
	99	+ return "k.scalar(dvm.int64)"
	100	+ if isinstance(val, torch.SymFloat):
	101	+ self.cont_flag_input.append(True)
	102	+ return "k.scalar(dvm.float32)"
	103	+
	104	+ shape, stride, dtype = val.shape, val.stride(), val.dtype
	105	+ is_symbolic = any(
	106	+ isinstance(s, torch.SymInt) and s.node.is_symbolic() for s in shape
	107	+ )
	108	+ is_contiguous = val.is_contiguous()
	109	+
	110	+ if self.is_mix_kernel:
	111	+ trans = meta.get("trans", False)
	112	+ self.need_trans_input.append(trans)
	113	+ self.cont_flag_input.append(trans or is_contiguous)
	114	+ if trans:
	115	+ shape = val.mT.shape
	116	+ shape = [-1 if isinstance(s, torch.SymInt) else s for s in shape]
	117	+ return load(shape, dtype)
	118	+
	119	+ shape = [-1 if isinstance(s, torch.SymInt) else s for s in shape]
	120	+ stride = [-1 if isinstance(s, torch.SymInt) else s for s in stride]
	121	+ if is_contiguous:
	122	+ expr, skip_cont = load(shape, dtype), True
	123	+ else:
	124	+ expr, skip_cont = codegen_maybe_view_load(
	125	+ shape,
	126	+ stride,
	127	+ dtype,
	128	+ view_fusion_level=self.view_fusion_level,
	129	+ is_symbolic=is_symbolic,
	130	+ )
	131	+ self.cont_flag_input.append(skip_cont)
	132	+ return expr
	133	+
	134	+ def call_function(
	135	+ self, target: "Target", args: tuple[Argument, ...], kwargs: dict[str, Argument]
	136	+ ) -> Argument:
	137	+ if target not in DVM_OP_REGISTRY:
	138	+ raise NotImplementedError(f"{target} not implemented in DVM")
	139	+ func, _ = DVM_OP_REGISTRY.get(target)
	140	+ meta = self.current_node.meta
	141	+
	142	+ if target in (aten.mm.default, aten.bmm.default):
	143	+ args = (*args, meta.get("trans_a", False), meta.get("trans_b", False))
	144	+
	145	+ elif target is aten.addmm.default:
	146	+ args = (
	147	+ *args,
	148	+ meta.get("trans_a", False),
	149	+ meta.get("trans_b", False),
	150	+ meta.get("use_bias", False),
	151	+ )
	152	+
	153	+ return func(args, *kwargs)
	154	+
	155	+ def output(
	156	+ self, target: "Target", args: tuple[Argument, ...], kwargs: dict[str, Argument]
	157	+ ) -> Argument:
	158	+ outs = super().output(target, args, kwargs)
	159	+
	160	+ def codegen(out, node):
	161	+ if isinstance(node, torch.fx.Node):
	162	+ return store(out, node.meta["val"].dtype)
	163	+ return ""
	164	+
	165	+ return pytree.tree_map(codegen, outs, self.current_node.args[0])

Atorch_npu/_inductor/dvm/graph_fusion.py+437-0

@@ -0,0 +1,437 @@
	1	+import os
	2	+from collections import defaultdict
	3	+from types import SimpleNamespace
	4	+
	5	+import torch
	6	+from torch.fx import Graph, GraphModule, Node
	7	+from torch.library import Library
	8	+from torch.utils._pytree import tree_map
	9	+from torch._inductor import config as inductor_config
	10	+from torch._inductor.codegen.wrapper import PythonWrapperCodegen
	11	+from torch._inductor.virtualized import V
	12	+from torch._subclasses import FakeTensor
	13	+
	14	+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
	15	+from torch.fx.passes.operator_support import OperatorSupportBase
	16	+from torch.fx.passes.tools_common import legalize_graph
	17	+from torch.fx.passes.utils.fuser_utils import (
	18	+ topo_sort,
	19	+ fuse_as_graphmodule,
	20	+ erase_nodes,
	21	+)
	22	+
	23	+from .graph_build import DvmCodegenInterpreter, is_fx_dynamic
	24	+from .util import patch_gm_placeholder_strides_from_codegen_args
	25	+from .fx_test import generate_dvm_fx_case
	26	+from .op_emitter import DVM_OP_REGISTRY
	27	+from .fx_pass import (
	28	+ decompose_k1_matmul_to_mul,
	29	+ insert_promote_cast_by_pos_prims,
	30	+ insert_sum_fp32_prepost_cast_prims,
	31	+ expand_to_reshape,
	32	+ need_fallback_gm,
	33	+)
	34	+
	35	+dump_fx_test = os.environ.get("INDUCTOR_DVM_DUMP_FX_TEST", "0") == "1"
	36	+
	37	+view_fusion_level = int(os.environ.get("INDUCTOR_DVM_VIEW_FUSION_LEVEL", "1"))
	38	+
	39	+aten = torch.ops.aten
	40	+prims = torch.ops.prims
	41	+
	42	+GRAPH_FUSION_SUPPORT_OP = [
	43	+ aten.add.Tensor,
	44	+ aten.add.Scalar,
	45	+ aten.sub.Tensor,
	46	+ aten.sub.Scalar,
	47	+ aten.mul.Tensor,
	48	+ aten.mul.Scalar,
	49	+ aten.div.Tensor,
	50	+ aten.div.Scalar,
	51	+ aten.pow.Tensor_Tensor,
	52	+ aten.pow.Tensor_Scalar,
	53	+ aten.pow.Scalar,
	54	+ aten.lt.Tensor,
	55	+ aten.lt.Scalar,
	56	+ aten.le.Tensor,
	57	+ aten.le.Scalar,
	58	+ aten.gt.Tensor,
	59	+ aten.gt.Scalar,
	60	+ aten.ge.Tensor,
	61	+ aten.ge.Scalar,
	62	+ aten.eq.Tensor,
	63	+ aten.eq.Scalar,
	64	+ aten.ne.Tensor,
	65	+ aten.ne.Scalar,
	66	+ aten.maximum.default,
	67	+ aten.minimum.default,
	68	+ aten.sqrt.default,
	69	+ aten.rsqrt.default,
	70	+ aten.abs.default,
	71	+ aten.log.default,
	72	+ aten.exp.default,
	73	+ aten.reciprocal.default,
	74	+ aten.isfinite.default,
	75	+ prims.convert_element_type.default,
	76	+ torch.ops.npu.npu_dtype_cast.default,
	77	+ torch.ops.npu.npu_dtype_cast_backward.default,
	78	+ torch.ops.npu._npu_dtype_cast.default,
	79	+ torch.ops.npu._npu_dtype_cast_backward.default,
	80	+ aten.sum.dim_IntList,
	81	+ aten.sum.default,
	82	+ aten.neg.default,
	83	+ aten.relu.default,
	84	+ aten.mm.default,
	85	+ aten.bmm.default,
	86	+ aten.addmm.default,
	87	+ aten.where.default,
	88	+ aten.where.self,
	89	+ # aten.expand.default,
	90	+ aten.full.default,
	91	+ # aten.reshape.default,
	92	+]
	93	+
	94	+class UnionFind:
	95	+ def __init__(self) -> None:
	96	+ self.parent: dict[Node, Node] = {}
	97	+ self.rank: dict[Node, int] = {}
	98	+
	99	+ def find(self, x: Node) -> Node:
	100	+ p = self.parent.get(x, x)
	101	+ if p != x:
	102	+ self.parent[x] = self.find(p)
	103	+ else:
	104	+ self.parent[x] = x
	105	+ return self.parent[x]
	106	+
	107	+ def union(self, a: Node, b: Node) -> None:
	108	+ ra, rb = self.find(a), self.find(b)
	109	+ if ra == rb:
	110	+ return
	111	+
	112	+ rka = self.rank.get(ra, 0)
	113	+ rkb = self.rank.get(rb, 0)
	114	+ if rka < rkb:
	115	+ ra, rb = rb, ra
	116	+ rka, rkb = rkb, rka
	117	+
	118	+ self.parent[rb] = ra
	119	+ if rka == rkb:
	120	+ self.rank[ra] = rka + 1
	121	+
	122	+
	123	+class DvmOpSupport(OperatorSupportBase):
	124	+ def is_node_supported(self, submodules, node):
	125	+ if node.op == "call_function" and node.target in GRAPH_FUSION_SUPPORT_OP:
	126	+ _, rule = DVM_OP_REGISTRY.get(node.target)
	127	+ return rule(node)
	128	+ return False
	129	+
	130	+
	131	+def split_partition_with_union_find(
	132	+ partition_nodes: dict[Node, None],
	133	+) -> list[dict[Node, None]]:
	134	+ """
	135	+ Split a partition into connected components based on data-dependency edges
	136	+ within the partition: if u is an input of v and both are in partition,
	137	+ they belong to the same component.
	138	+ """
	139	+ nodes = list(partition_nodes.keys())
	140	+ node_set = set(nodes)
	141	+
	142	+ uf = UnionFind()
	143	+ for n in nodes:
	144	+ uf.find(n)
	145	+
	146	+ for v in nodes:
	147	+ for u in v.all_input_nodes:
	148	+ if u in node_set:
	149	+ uf.union(u, v)
	150	+
	151	+ groups: dict[Node, dict[Node, None]] = defaultdict(dict)
	152	+ for n in nodes:
	153	+ root = uf.find(n)
	154	+ groups[root][n] = None
	155	+
	156	+ return list(groups.values())
	157	+
	158	+
	159	+class _FusedMeta:
	160	+ def __init__(self, name: str, gm: GraphModule, input_nodes: list[Node]):
	161	+ self.name = name
	162	+ self.gm = gm
	163	+
	164	+ def codegen(self):
	165	+ """
	166	+ Return (codegen_interpreter, python_source_string).
	167	+ """
	168	+ if dump_fx_test:
	169	+ generate_dvm_fx_case(self.gm, fusion_type="graph")
	170	+ cg = DvmCodegenInterpreter(
	171	+ self.gm, ktype="split", view_fusion_level=view_fusion_level
	172	+ )
	173	+ cg.run()
	174	+ code = cg.code.getvalue().replace(cg.KERNEL_NAME_PLACEHOLDER, self.name)
	175	+ return cg, code
	176	+
	177	+
	178	+# fused_id -> meta
	179	+_fused_metas: dict[int, _FusedMeta] = {}
	180	+_fused_libs: dict[str, Library] = {}
	181	+
	182	+
	183	+def _fused_run_stub(args, *kwargs):
	184	+ raise AssertionError("This op should never run at eager runtime.")
	185	+
	186	+
	187	+def _fused_run_fake(args, *kwargs):
	188	+ """
	189	+ Fake implementation for custom_op. The last arg is fused_id.
	190	+ We execute sub_gm.forward on meta tensors and return contiguous meta outputs.
	191	+ """
	192	+ fused_id = int(args[-1])
	193	+ meta = _fused_metas[fused_id]
	194	+ out = meta.gm.forward(*args[:-1])
	195	+ return tree_map(lambda t: t if t.is_contiguous() else t.contiguous(), out)
	196	+
	197	+
	198	+class GraphFusionPartitioner(CapabilityBasedPartitioner):
	199	+ fused_op_map_: dict[str, object] = {}
	200	+ fused_id_: int = 0
	201	+
	202	+ def _input_schema_types(self, input_nodes: list[Node]) -> list[str]:
	203	+ input_types: list[str] = []
	204	+ for node in input_nodes:
	205	+ val = node.meta.get("val", None)
	206	+ if isinstance(val, torch.SymInt):
	207	+ input_types.append("SymInt")
	208	+ elif isinstance(val, torch.SymFloat):
	209	+ input_types.append("SymFloat")
	210	+ else:
	211	+ input_types.append("Tensor")
	212	+ return input_types
	213	+
	214	+ def _build_schema(self, input_types: list[str], output_len: int) -> str:
	215	+ extra_inputs = "int fused_id"
	216	+ input_schema = ", ".join(
	217	+ f"{input_type} x{i}" for i, input_type in enumerate(input_types)
	218	+ )
	219	+ input_schema = (
	220	+ f"{input_schema}, {extra_inputs}" if input_schema else extra_inputs
	221	+ )
	222	+
	223	+ if output_len == 1:
	224	+ output_schema = "Tensor"
	225	+ else:
	226	+ output_schema = f'({", ".join(["Tensor"] * output_len)})'
	227	+ return f"({input_schema}) -> {output_schema}"
	228	+
	229	+ def _get_or_create_custom_op(self, input_nodes: list[Node], output_len: int):
	230	+ input_types = self._input_schema_types(input_nodes)
	231	+ type_suffix = "_".join(
	232	+ "t" if t == "Tensor" else "si" if t == "SymInt" else "sf"
	233	+ for t in input_types
	234	+ )
	235	+ input_len = len(input_types)
	236	+ op_def = f"{input_len}_{output_len}"
	237	+ if any(t != "Tensor" for t in input_types):
	238	+ op_def = f"{op_def}_{type_suffix}"
	239	+ custom = self.fused_op_map_.get(op_def, None)
	240	+ if custom is not None:
	241	+ return custom
	242	+ schema = self._build_schema(input_types, output_len)
	243	+ op_name = "fused_graph_" + op_def
	244	+ qualname = "dvm::" + op_name
	245	+
	246	+ lib = Library("dvm", "FRAGMENT")
	247	+ # Use flexible_layout so shared producers can stay shared, instead of
	248	+ # each fused_graph consumer materializing its own fixed-stride copy.
	249	+ lib.define(op_name + schema, tags=[torch._C.Tag.flexible_layout])
	250	+ lib._register_fake(op_name, _fused_run_fake, _stacklevel=1)
	251	+ lib.impl(op_name, _fused_run_stub, "CompositeExplicitAutograd")
	252	+
	253	+ opoverload = getattr(torch.ops.dvm, op_name).default
	254	+ custom = SimpleNamespace(_opoverload=opoverload, _lib=lib, _qualname=qualname)
	255	+ _fused_libs[qualname] = lib
	256	+ self.fused_op_map_[op_def] = custom
	257	+ return custom
	258	+
	259	+ def _should_fuse(
	260	+ self,
	261	+ sub_gm: GraphModule,
	262	+ orig_outputs: list[Node],
	263	+ orig_inputs: list[Node],
	264	+ ) -> bool:
	265	+ if len(orig_outputs) == 0:
	266	+ return False
	267	+ if need_fallback_gm(sub_gm):
	268	+ return False
	269	+ return any(
	270	+ isinstance(node.meta.get("val", None), FakeTensor)
	271	+ for node in [orig_inputs, orig_outputs]
	272	+ )
	273	+
	274	+ def partition_and_fuse(self) -> GraphModule:
	275	+ partitions = self.propose_partitions()
	276	+
	277	+ # further split each proposed partition by union-find connectivity
	278	+ partition_nodes_list: list[dict[Node, None]] = [
	279	+ sp_nodes
	280	+ for partition in partitions
	281	+ for sp_nodes in split_partition_with_union_find(partition.nodes)
	282	+ ]
	283	+
	284	+ for partition_nodes in partition_nodes_list:
	285	+ fused_id = self.fused_id_
	286	+ self.fused_id_ += 1
	287	+
	288	+ fused_name = f"dvm_graph_fused_{fused_id}"
	289	+ sorted_nodes = topo_sort(list(partition_nodes))
	290	+
	291	+ sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(
	292	+ self.graph_module,
	293	+ sorted_nodes,
	294	+ fused_name,
	295	+ partition_nodes,
	296	+ )
	297	+
	298	+ # post-processing inside sub graph
	299	+ decompose_k1_matmul_to_mul(sub_gm)
	300	+ insert_promote_cast_by_pos_prims(sub_gm)
	301	+ insert_sum_fp32_prepost_cast_prims(sub_gm)
	302	+ expand_to_reshape(sub_gm)
	303	+
	304	+ if not self._should_fuse(sub_gm, orig_outputs, orig_inputs):
	305	+ continue
	306	+
	307	+ _fused_metas[fused_id] = _FusedMeta(fused_name, sub_gm, orig_inputs)
	308	+
	309	+ output_len = len(orig_outputs)
	310	+ custom = self._get_or_create_custom_op(orig_inputs, output_len)
	311	+
	312	+ # create fused call node in original graph
	313	+ args = (*orig_inputs, fused_id)
	314	+ new_node = self.graph_module.graph.call_function(
	315	+ custom._opoverload, tuple(args), None
	316	+ )
	317	+
	318	+ new_meta_vals = []
	319	+ with V.fake_mode:
	320	+ for orig_output in orig_outputs:
	321	+ meta_val = orig_output.meta["val"]
	322	+ new_meta_vals.append(
	323	+ torch.empty(
	324	+ meta_val.size(),
	325	+ dtype=meta_val.dtype,
	326	+ device=meta_val.device,
	327	+ requires_grad=meta_val.requires_grad,
	328	+ )
	329	+ )
	330	+
	331	+ if output_len == 1:
	332	+ orig_outputs[0].replace_all_uses_with(new_node)
	333	+ new_node.meta["val"] = new_meta_vals[0]
	334	+ else:
	335	+ for i, (orig_output, meta_val) in enumerate(
	336	+ zip(orig_outputs, new_meta_vals)
	337	+ ):
	338	+ proxy_out = torch.fx.Proxy(new_node)[i].node
	339	+ proxy_out.meta["val"] = meta_val
	340	+ orig_output.replace_all_uses_with(proxy_out)
	341	+ new_node.meta["val"] = tuple(new_meta_vals)
	342	+
	343	+ # erase old nodes
	344	+ erase_nodes(self.graph_module, sorted_nodes)
	345	+
	346	+ legalize_graph(self.graph_module)
	347	+ output_node = self.graph_module.graph.find_nodes(op="output")[0]
	348	+ next(iter(reversed(self.graph_module.graph.nodes))).append(output_node)
	349	+ return self.graph_module
	350	+
	351	+
	352	+def dvm_graph_fusion(graph: Graph):
	353	+ gm: GraphModule = graph.owning_module
	354	+
	355	+ dvm_support = DvmOpSupport()
	356	+ fusion_part = GraphFusionPartitioner(
	357	+ gm,
	358	+ dvm_support,
	359	+ allows_single_node_partition=True,
	360	+ )
	361	+ fusion_part.partition_and_fuse()
	362	+
	363	+
	364	+def _dvm_generate_fallback_kernel(self, fallback_kernel):
	365	+ """
	366	+ Patch point: PythonWrapperCodegen.generate_fallback_kernel
	367	+ If it's our custom op, pop meta and emit dvm kernel code.
	368	+ """
	369	+ args = [fallback_kernel.codegen_args(), fallback_kernel.codegen_kwargs()]
	370	+
	371	+ if not fallback_kernel.op_overload._name.startswith("dvm::fused_graph_"):
	372	+ DvmGraphFusionPatch._orig_generate_fallback_kernel(self, fallback_kernel)
	373	+ return
	374	+
	375	+ fused_id = int(args[-1])
	376	+ meta = _fused_metas.pop(fused_id)
	377	+
	378	+ args_list = list(args[:-1])
	379	+ if not is_fx_dynamic(meta.gm):
	380	+ patch_gm_placeholder_strides_from_codegen_args(meta.gm, args_list)
	381	+ cg, code = meta.codegen()
	382	+ self.header.splice(code)
	383	+
	384	+ buf_name = fallback_kernel.get_name()
	385	+
	386	+ args_list = list(args[:-1])
	387	+ # cont/trans handling based on codegen interpreter
	388	+ for i, skip_cont in enumerate(cg.cont_flag_input):
	389	+ if not skip_cont:
	390	+ args_list[i] += ".contiguous()"
	391	+ for i, trans in enumerate(cg.need_trans_input):
	392	+ if trans:
	393	+ args_list[i] += ".mT"
	394	+
	395	+ self.writeline(f"{buf_name} = {meta.name}({', '.join(args_list)})")
	396	+ self.add_import_once("from torch_npu._inductor import dvm")
	397	+
	398	+
	399	+class DvmGraphFusionPatch:
	400	+ _enabled = False
	401	+ _orig_generate_fallback_kernel = None
	402	+ _orig_post_grad_custom_post_pass = None
	403	+
	404	+ @staticmethod
	405	+ def enable() -> None:
	406	+ if not DvmGraphFusionPatch._enabled:
	407	+ DvmGraphFusionPatch._orig_generate_fallback_kernel = (
	408	+ PythonWrapperCodegen.generate_fallback_kernel
	409	+ )
	410	+ DvmGraphFusionPatch._orig_post_grad_custom_post_pass = (
	411	+ inductor_config.post_grad_custom_post_pass
	412	+ )
	413	+ PythonWrapperCodegen.generate_fallback_kernel = (
	414	+ _dvm_generate_fallback_kernel
	415	+ )
	416	+ inductor_config.post_grad_custom_post_pass = dvm_graph_fusion
	417	+ DvmGraphFusionPatch._enabled = True
	418	+
	419	+ @staticmethod
	420	+ def disable() -> None:
	421	+ if not DvmGraphFusionPatch._enabled:
	422	+ return
	423	+ PythonWrapperCodegen.generate_fallback_kernel = (
	424	+ DvmGraphFusionPatch._orig_generate_fallback_kernel
	425	+ )
	426	+ inductor_config.post_grad_custom_post_pass = (
	427	+ DvmGraphFusionPatch._orig_post_grad_custom_post_pass
	428	+ )
	429	+ DvmGraphFusionPatch._enabled = False
	430	+
	431	+ def __enter__(self) -> "DvmGraphFusionPatch":
	432	+ DvmGraphFusionPatch.enable()
	433	+ return self
	434	+
	435	+ def __exit__(self, exc_type, exc, tb) -> bool:
	436	+ DvmGraphFusionPatch.disable()
	437	+ return False

Atorch_npu/_inductor/dvm/mlir_fusion.py+395-0

@@ -0,0 +1,395 @@
	1	+import os
	2	+
	3	+import torch
	4	+from torch._inductor import config
	5	+from torch._inductor.fx_passes.control_dependencies import control_deps
	6	+from torch._inductor.codegen.common import IndentedBuffer
	7	+from torch._inductor.codegen.simd import code_hash, SIMDKernel
	8	+from torch._inductor.scheduler import WhyNoFuse
	9	+from torch._inductor.utils import get_fused_kernel_name
	10	+from torch._inductor.virtualized import V
	11	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir import config as anir_config
	12	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.codegen.mlir import (
	13	+ NpuMlirKernel,
	14	+ NpuMlirScheduling,
	15	+)
	16	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch import (
	17	+ lowering as npu_lowering,
	18	+)
	19	+from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.utils import (
	20	+ get_num_call_functions,
	21	+ to_folder,
	22	+)
	23	+
	24	+from .decomp import patch_decomp
	25	+from .fx_test import generate_dvm_fx_case
	26	+from .graph_build import DvmCodegenInterpreter
	27	+from .op_emitter import common_rule, DVM_OP_REGISTRY, DVM_SUPPORT_TYPE, _extra_int_types
	28	+
	29	+
	30	+dump_fx_test = os.environ.get("INDUCTOR_DVM_DUMP_FX_TEST", "0") == "1"
	31	+view_fusion_level = int(os.environ.get("INDUCTOR_DVM_VIEW_FUSION_LEVEL", "1"))
	32	+disable_post_reduce_fusion = (
	33	+ os.environ.get("INDUCTOR_DVM_DISABLE_POST_REDUCE_FUSION", "0") == "1"
	34	+)
	35	+aten = torch.ops.aten
	36	+prims = torch.ops.prims
	37	+quantized = torch.ops.quantized
	38	+_quantized = torch.ops._quantized
	39	+
	40	+
	41	+anir_config.GENERATE_LIST = [
	42	+ control_deps,
	43	+ aten._assert_scalar,
	44	+ aten.mul,
	45	+ aten.add,
	46	+ aten.sub,
	47	+ aten.div,
	48	+ aten.clamp_min,
	49	+ aten.clamp_max,
	50	+ aten.maximum,
	51	+ aten.minimum,
	52	+ aten.abs,
	53	+ aten.reciprocal,
	54	+ aten.log,
	55	+ aten.exp,
	56	+ aten.pow,
	57	+ aten.sqrt,
	58	+ aten.rsqrt,
	59	+ aten.neg,
	60	+ aten.lt,
	61	+ aten.le,
	62	+ aten.gt,
	63	+ aten.ge,
	64	+ aten.eq,
	65	+ aten.ne,
	66	+ aten.bitwise_and,
	67	+ aten.bitwise_or,
	68	+ aten.bitwise_not,
	69	+ aten.where,
	70	+ prims.convert_element_type,
	71	+ torch.ops.npu.npu_dtype_cast,
	72	+ torch.ops.npu.npu_dtype_cast_backward,
	73	+ torch.ops.npu._npu_dtype_cast,
	74	+ torch.ops.npu._npu_dtype_cast_backward,
	75	+ aten.expand,
	76	+ aten.var_mean,
	77	+ aten.sum,
	78	+ aten.mean,
	79	+ aten.amax,
	80	+ aten.amin,
	81	+ aten.full,
	82	+ aten.relu,
	83	+ aten.where,
	84	+ aten.scalar_tensor,
	85	+ aten.unsqueeze,
	86	+ aten.squeeze,
	87	+ aten.reshape,
	88	+ # aten.copy,
	89	+ # aten.copy_,
	90	+ # aten.clone,
	91	+]
	92	+
	93	+def _is_node_supported_by_dvm_rule(node, allow_common_rule=False):
	94	+ if node.target in DVM_OP_REGISTRY:
	95	+ _, rule = DVM_OP_REGISTRY.get(node.target)
	96	+ return rule(node)
	97	+ return allow_common_rule and common_rule(node)
	98	+
	99	+
	100	+def _codegen_dvm_kernel(self, Name=None):
	101	+ def is_node_dvm_supported(node):
	102	+ if node.op == "placeholder":
	103	+ meta = node.meta["val"]
	104	+ if isinstance(meta, torch._subclasses.FakeTensor):
	105	+ return meta.dtype in [DVM_SUPPORT_TYPE, _extra_int_types]
	106	+ if node.op == "call_function":
	107	+ return _is_node_supported_by_dvm_rule(node)
	108	+ return True
	109	+
	110	+ if all(is_node_dvm_supported(node) for node in self._gm.graph.nodes):
	111	+ self.dvm_codegen = DvmCodegenInterpreter(
	112	+ self._gm, ktype="vector", view_fusion_level=view_fusion_level
	113	+ )
	114	+ self.dvm_codegen.run()
	115	+ return self.dvm_codegen.code.getvalue()
	116	+ else:
	117	+ self.dvm_codegen = None
	118	+ return self._gm.print_readable(print_output=False)
	119	+
	120	+
	121	+def _kernel_layout_key(mlir_kernel):
	122	+ non_contiguous_key = tuple(
	123	+ (name, tuple(indices))
	124	+ for name, indices in sorted(mlir_kernel.non_contiguous_indices.items())
	125	+ )
	126	+ dvm_codegen = getattr(mlir_kernel, "dvm_codegen", None)
	127	+ cont_flag_input = tuple(getattr(dvm_codegen, "cont_flag_input", ()))
	128	+ return non_contiguous_key, cont_flag_input
	129	+
	130	+
	131	+def _define_dvm_kernel(self, src_code, mlir_kernel, traced_graph, mode=None):
	132	+ kernel_key = (src_code, _kernel_layout_key(mlir_kernel))
	133	+ wrapper = V.graph.wrapper_code
	134	+
	135	+ if kernel_key in wrapper.src_to_kernel:
	136	+ kernel_name = wrapper.src_to_kernel[kernel_key]
	137	+ else:
	138	+ fused_kernel_name = "dvm_" + get_fused_kernel_name(
	139	+ mlir_kernel._snodes, config.triton.descriptive_names
	140	+ )
	141	+ kernel_suffix = V.graph.wrapper_code.next_kernel_suffix()
	142	+ kernel_name = "_".join([fused_kernel_name, kernel_suffix])
	143	+
	144	+ traced_graph_hash = code_hash(
	145	+ traced_graph.print_readable(print_output=False) + kernel_name
	146	+ )
	147	+
	148	+ kernel_info = {}
	149	+
	150	+ wrapper.src_to_kernel[kernel_key] = kernel_name
	151	+ current_device = V.graph.get_current_device_or_throw()
	152	+
	153	+ compile_wrapper = IndentedBuffer()
	154	+ if (
	155	+ mlir_kernel.dvm_codegen is None
	156	+ or kernel_name in anir_config.force_fallback_kernel_names
	157	+ ):
	158	+ num_call_functions = get_num_call_functions(mlir_kernel._gm)
	159	+ kernel_meta = {
	160	+ "device_str": current_device.type,
	161	+ "device_index": current_device.index,
	162	+ "num_outputs": mlir_kernel.num_outputs,
	163	+ "non_contiguous_indices": mlir_kernel.non_contiguous_indices,
	164	+ "dynamic": mlir_kernel._is_dynamic,
	165	+ "mutated_indices": mlir_kernel.mutated_indices,
	166	+ "traced_graph_cache": anir_config.traced_graph_cache,
	167	+ "traced_graph_hash": traced_graph_hash,
	168	+ "num_call_functions": num_call_functions,
	169	+ **kernel_info,
	170	+ }
	171	+ compile_wrapper.writeline(
	172	+ f"async_compile.import_fx({kernel_name!r}, kernel_meta={kernel_meta})"
	173	+ )
	174	+ metadata_comment = (
	175	+ f'"""\n{mlir_kernel._gm.print_readable(print_output=False)}\n"""'
	176	+ )
	177	+ wrapper.define_kernel(
	178	+ kernel_name, compile_wrapper.getvalue(), metadata_comment
	179	+ )
	180	+ dump_path = os.path.join(
	181	+ os.getenv("TORCHINDUCTOR_CACHE_DIR"),
	182	+ anir_config.traced_graph_cache,
	183	+ str(current_device.index),
	184	+ traced_graph_hash,
	185	+ )
	186	+ if not os.path.exists(dump_path):
	187	+ os.makedirs(dump_path, exist_ok=True)
	188	+ to_folder(
	189	+ mlir_kernel._gm,
	190	+ dump_path,
	191	+ graph_hash=traced_graph_hash,
	192	+ module_name=traced_graph_hash,
	193	+ )
	194	+ else:
	195	+ wrapper.add_import_once("from torch_npu._inductor import dvm")
	196	+ if dump_fx_test:
	197	+ generate_dvm_fx_case(mlir_kernel._gm, fusion_type="mlir")
	198	+ kernel_meta = {
	199	+ "kernel_name": fused_kernel_name,
	200	+ "kernel_fullname": kernel_name,
	201	+ }
	202	+ code = mlir_kernel.dvm_codegen.code
	203	+ code.splice(
	204	+ f"""
	205	+ k.set_kernel_info(
	206	+ {kernel_meta.get("kernel_name")!r}, # kernel_name
	207	+ {kernel_meta.get("kernel_fullname")!r}, # kernel_fullname
	208	+ )
	209	+ """,
	210	+ strip=True,
	211	+ )
	212	+ func_name = kernel_name + "_build"
	213	+ func_code = code.getvalue().replace(
	214	+ mlir_kernel.dvm_codegen.KERNEL_NAME_PLACEHOLDER, func_name
	215	+ )
	216	+ compile_wrapper.writeline(func_name)
	217	+
	218	+ if anir_config.online_acc_comp:
	219	+ dump_path = os.path.join(
	220	+ os.getenv("TORCHINDUCTOR_CACHE_DIR"),
	221	+ anir_config.traced_graph_cache,
	222	+ str(current_device.index),
	223	+ traced_graph_hash,
	224	+ )
	225	+ if not os.path.exists(dump_path):
	226	+ os.makedirs(dump_path, exist_ok=True)
	227	+ to_folder(
	228	+ mlir_kernel._gm,
	229	+ dump_path,
	230	+ graph_hash=traced_graph_hash,
	231	+ module_name=traced_graph_hash,
	232	+ )
	233	+ compile_wrapper.writeline(
	234	+ f"{kernel_name}._acc_meta = {{"
	235	+ f"'traced_graph_hash': {traced_graph_hash!r}, "
	236	+ f"'traced_graph_cache': {anir_config.traced_graph_cache!r}, "
	237	+ f"'device_index': {current_device.index}, "
	238	+ f"'num_outputs': {mlir_kernel.num_outputs}}}"
	239	+ )
	240	+
	241	+ wrapper.define_kernel(kernel_name, compile_wrapper.getvalue(), func_code)
	242	+
	243	+ return kernel_name
	244	+
	245	+
	246	+def _dvm_can_fuse_vertical(self, node1, node2):
	247	+ _, (numel1, rnumel1) = node1.group
	248	+ _, (numel2, rnumel2) = node2.group
	249	+ why = WhyNoFuse(node1, node2)
	250	+
	251	+ if node1.is_reduction():
	252	+ return False
	253	+
	254	+ if not node2.is_reduction():
	255	+ return numel1 == numel2 and rnumel1 == rnumel2
	256	+ else:
	257	+ if numel1 == numel2 * rnumel2:
	258	+ if not all(
	259	+ SIMDKernel.is_compatible((numel2, rnumel2), n.get_ranges())
	260	+ for n in node1.get_nodes()
	261	+ ):
	262	+ why("nodes numel/rnumel incompatibility")
	263	+ return False
	264	+
	265	+ return True
	266	+ if numel1 != numel2:
	267	+ why("nodes numel incompatibility")
	268	+ return numel1 == numel2
	269	+
	270	+
	271	+def _dvm_can_fuse_horizontal(self, node1, node2):
	272	+ return False
	273	+
	274	+
	275	+def _patch_lowering_type_checks():
	276	+ import torch._inductor.graph as inductor_graph
	277	+ import torch._inductor.lowering as inductor_lowering
	278	+ import torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering as npu_lowering_mod
	279	+ import torch_npu._inductor.graph as npu_graph_mod
	280	+
	281	+ fallback_node_due_to_unsupported_type = (
	282	+ inductor_lowering.fallback_node_due_to_unsupported_type
	283	+ )
	284	+
	285	+ def _fallback_node_due_to_unsupported_type(
	286	+ node: torch.fx.Node, allow_cpu_inputs=True
	287	+ ):
	288	+ if fallback_node_due_to_unsupported_type(node, allow_cpu_inputs):
	289	+ return True
	290	+
	291	+ if node.target is torch.ops.higher_order.triton_kernel_wrapper_functional:
	292	+ return False
	293	+ if node.target is torch.ops.higher_order.triton_kernel_wrapper_mutation:
	294	+ return False
	295	+ if node.target is aten.lift_fresh_copy.default:
	296	+ return False
	297	+
	298	+ return not _is_node_supported_by_dvm_rule(node, allow_common_rule=True)
	299	+
	300	+ inductor_lowering.fallback_node_due_to_unsupported_type = (
	301	+ _fallback_node_due_to_unsupported_type
	302	+ )
	303	+ npu_lowering_mod.fallback_node_due_to_unsupported_type = (
	304	+ _fallback_node_due_to_unsupported_type
	305	+ )
	306	+ inductor_graph.fallback_node_due_to_unsupported_type = (
	307	+ _fallback_node_due_to_unsupported_type
	308	+ )
	309	+ npu_graph_mod.fallback_node_due_to_unsupported_type = (
	310	+ _fallback_node_due_to_unsupported_type
	311	+ )
	312	+
	313	+
	314	+def _patch_lowering():
	315	+ from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering import (
	316	+ is_boolean_dtype,
	317	+ is_integer_dtype,
	318	+ make_reduction,
	319	+ to_dtype,
	320	+ )
	321	+
	322	+ def get_overloads(aten_fn):
	323	+ if not isinstance(aten_fn, (list, tuple)):
	324	+ aten_fn = [aten_fn]
	325	+ else:
	326	+ aten_fn = list(aten_fn)
	327	+
	328	+ for fn in list(aten_fn):
	329	+ if isinstance(fn, torch._ops.OpOverloadPacket):
	330	+ for overload in fn.overloads():
	331	+ other_fn = getattr(fn, overload)
	332	+ aten_fn.append(other_fn)
	333	+
	334	+ return aten_fn
	335	+
	336	+ def sum_(x, axis=None, keepdims=False, *, dtype=None):
	337	+ if axis and any(ax < 0 for ax in axis):
	338	+ offset = len(x.get_size())
	339	+ axis = [ax + offset if ax < 0 else ax for ax in axis]
	340	+ if (
	341	+ is_integer_dtype(x.get_dtype()) or is_boolean_dtype(x.get_dtype())
	342	+ ) and dtype is None:
	343	+ dtype = torch.int64
	344	+
	345	+ out_dtype = x.get_dtype() if dtype is None else dtype
	346	+
	347	+ fn = make_reduction("sum", override_return_dtype=torch.float32)
	348	+ r = fn(x, axis, keepdims, dtype=torch.float32)
	349	+
	350	+ if out_dtype != torch.float32:
	351	+ r = to_dtype(r, out_dtype)
	352	+
	353	+ return r
	354	+
	355	+ anir_config.disable_any_pbr = False
	356	+ ops = get_overloads([aten.sum, prims.sum])
	357	+ npu_lowering.register_lowering(ops)(sum_)
	358	+ npu_lowering.make_fallback(
	359	+ aten.matmul_backward.default,
	360	+ layout_constraint=None,
	361	+ warn=False,
	362	+ override_decomp=True,
	363	+ )
	364	+ npu_lowering.add_layout_constraint(aten.matmul_backward.default, None)
	365	+
	366	+
	367	+class DvmMlirFusionPatch:
	368	+ _enabled = False
	369	+
	370	+ @staticmethod
	371	+ def enable() -> None:
	372	+ if DvmMlirFusionPatch._enabled:
	373	+ return
	374	+ from torch._dynamo import config as dynamo_config
	375	+ from torch._inductor import config as inductor_config
	376	+
	377	+ dynamo_config.specialize_float = True # enable float specialization until launch with scalar supported
	378	+ inductor_config.unroll_reductions_threshold = 1 # disable unroll reductions
	379	+ inductor_config.size_asserts = (
	380	+ False # npu ops always return contiguous tensors which maybe different from meta outputs
	381	+ )
	382	+ inductor_config.allow_buffer_reuse = False
	383	+ inductor_config.comprehensive_padding = False
	384	+ patch_decomp()
	385	+ _patch_lowering_type_checks()
	386	+ _patch_lowering()
	387	+ NpuMlirKernel.codegen_kernel = _codegen_dvm_kernel
	388	+ NpuMlirScheduling.define_kernel = _define_dvm_kernel
	389	+ if disable_post_reduce_fusion:
	390	+ NpuMlirScheduling.can_fuse_horizontal = _dvm_can_fuse_horizontal
	391	+ NpuMlirScheduling.can_fuse_vertical = _dvm_can_fuse_vertical
	392	+ DvmMlirFusionPatch._enabled = True
	393	+
	394	+
	395	+DvmMlirFusionPatch.enable()

Atorch_npu/_inductor/dvm/op_emitter.py+577-0

@@ -0,0 +1,577 @@
	1	+import torch
	2	+import torch.fx
	3	+import torch.utils._pytree as pytree
	4	+
	5	+from . import is_ascend950
	6	+
	7	+
	8	+aten = torch.ops.aten
	9	+prims = torch.ops.prims
	10	+
	11	+DVM_OP_REGISTRY = {}
	12	+
	13	+DVM_SUPPORT_TYPE = [
	14	+ torch.bfloat16,
	15	+ torch.float16,
	16	+ torch.float32,
	17	+ torch.int32,
	18	+ torch.bool,
	19	+]
	20	+
	21	+
	22	+DVM_SUPPORT_FLOAT_TYPE = [
	23	+ torch.bfloat16,
	24	+ torch.float16,
	25	+ torch.float32,
	26	+]
	27	+
	28	+_extra_int_types = [torch.int32] if is_ascend950 else [torch.int32, torch.int64]
	29	+DVM_SUPPORT_FLOAT_INT_TYPE = [
	30	+ *DVM_SUPPORT_FLOAT_TYPE,
	31	+ *_extra_int_types,
	32	+]
	33	+
	34	+DVM_SUPPORT_OPTIONAL_INT64_TYPE = DVM_SUPPORT_TYPE if is_ascend950 else [*DVM_SUPPORT_TYPE, torch.int64]
	35	+
	36	+DVM_DTYPE_MAP = {
	37	+ torch.bfloat16: "dvm.bfloat16",
	38	+ torch.float16: "dvm.float16",
	39	+ torch.float32: "dvm.float32",
	40	+ torch.int32: "dvm.int32",
	41	+ torch.int64: "dvm.int64",
	42	+ torch.bool: "dvm.bool_",
	43	+}
	44	+
	45	+
	46	+def to_dvm_dtype(dtype):
	47	+ if dtype is None:
	48	+ return None
	49	+ if isinstance(dtype, torch.dtype):
	50	+ if dtype not in DVM_DTYPE_MAP:
	51	+ raise NotImplementedError(f"Unsupported dtype for DVM: {dtype}")
	52	+ return DVM_DTYPE_MAP[dtype]
	53	+ return dtype
	54	+
	55	+
	56	+def _check_dtype(inputs, supported_dtypes):
	57	+ for inp in inputs:
	58	+ if not isinstance(inp, torch.fx.Node):
	59	+ continue
	60	+ if "val" not in inp.meta:
	61	+ continue
	62	+
	63	+ for meta in pytree.tree_leaves(inp.meta["val"]):
	64	+ if not isinstance(meta, torch._subclasses.FakeTensor):
	65	+ continue
	66	+ if meta.dtype not in supported_dtypes:
	67	+ return False
	68	+ return True
	69	+
	70	+
	71	+def where_rule(node: torch.fx.Node):
	72	+ return _check_dtype(node.args[1:], DVM_SUPPORT_FLOAT_INT_TYPE)
	73	+
	74	+
	75	+def _is_last2_transpose_tensor(t: torch._subclasses.FakeTensor) -> bool:
	76	+ if t.dim() < 2:
	77	+ return False
	78	+
	79	+ if t.is_contiguous():
	80	+ return False
	81	+
	82	+ if not (t.stride(-2) == 1 and t.stride(-1) == t.size(-2)):
	83	+ return False
	84	+
	85	+ batch = t.size(-1) * t.size(-2)
	86	+ for d in range(t.dim() - 3, -1, -1):
	87	+ if t.stride(d) != batch:
	88	+ return False
	89	+ batch *= t.size(d)
	90	+
	91	+ return True
	92	+
	93	+
	94	+def mm_rule(node: torch.fx.Node):
	95	+ UINT16_MAX = (1 << 16) - 1
	96	+ UINT8_MAX = (1 << 8) - 1
	97	+ MAX_INNER = UINT16_MAX - UINT8_MAX
	98	+ SMALL_OUTPUT_MAX = 256
	99	+
	100	+ def inner_axis_length(t: torch._subclasses.FakeTensor):
	101	+ if _is_last2_transpose_tensor(t):
	102	+ return t.mT.size(-1)
	103	+ return t.size(-1)
	104	+
	105	+ def check(input_node):
	106	+ t = input_node.meta["val"]
	107	+ if t.dim() > 4 or t.dim() < 2:
	108	+ return False
	109	+ inner_axis = inner_axis_length(t)
	110	+ if not is_ascend950:
	111	+ if isinstance(inner_axis, torch.SymInt):
	112	+ return False
	113	+ if inner_axis > MAX_INNER:
	114	+ return False
	115	+ return True
	116	+
	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):
	127	+ lhs_t = lhs_node.meta["val"]
	128	+ rhs_t = rhs_node.meta["val"]
	129	+ lhs_k = lhs_t.shape[-1]
	130	+ rhs_k = rhs_t.shape[-2]
	131	+ if isinstance(lhs_k, torch.SymInt) or isinstance(rhs_k, torch.SymInt):
	132	+ return True
	133	+ if lhs_k == 1 and rhs_k == 1:
	134	+ return (not _is_last2_transpose_tensor(lhs_t)) and (
	135	+ not _is_last2_transpose_tensor(rhs_t)
	136	+ )
	137	+ return True
	138	+
	139	+ if node.target in [aten.mm.default, aten.bmm.default]:
	140	+ lhs = node.args[0]
	141	+ rhs = node.args[1]
	142	+ elif node.target is aten.addmm.default:
	143	+ lhs = node.args[1]
	144	+ rhs = node.args[2]
	145	+ else:
	146	+ return False
	147	+ if node.meta["val"].dtype not in (torch.float16, torch.bfloat16):
	148	+ return False
	149	+
	150	+ return (
	151	+ check(lhs) and check(rhs) and check_output(node) and check_k1_fusion(lhs, rhs)
	152	+ )
	153	+
	154	+
	155	+class DvmOpInfo:
	156	+ def __init__(
	157	+ self,
	158	+ func,
	159	+ input_dtypes=DVM_SUPPORT_FLOAT_TYPE,
	160	+ output_dtypes=DVM_SUPPORT_FLOAT_TYPE,
	161	+ rule=None,
	162	+ ):
	163	+ self.func = func
	164	+ self.input_dtypes = input_dtypes
	165	+ self.output_dtypes = output_dtypes
	166	+ self.rule = rule
	167	+
	168	+ def is_supported(self, node: torch.fx.Node):
	169	+ inputs = pytree.arg_tree_leaves(node.args, *node.kwargs)
	170	+ return (
	171	+ (
	172	+ self.input_dtypes is None
	173	+ or _check_dtype(inputs, self.input_dtypes)
	174	+ )
	175	+ and (
	176	+ self.output_dtypes is None
	177	+ or _check_dtype([node], self.output_dtypes)
	178	+ )
	179	+ and (self.rule is None or self.rule(node))
	180	+ )
	181	+
	182	+ def __iter__(self):
	183	+ yield self.func
	184	+ yield self.is_supported
	185	+
	186	+
	187	+def register_dvm_op(
	188	+ *ops,
	189	+ input_dtypes=DVM_SUPPORT_FLOAT_TYPE,
	190	+ output_dtypes=DVM_SUPPORT_FLOAT_TYPE,
	191	+ rule=None,
	192	+):
	193	+ def decorator(func):
	194	+ info = DvmOpInfo(
	195	+ func,
	196	+ input_dtypes=input_dtypes,
	197	+ output_dtypes=output_dtypes,
	198	+ rule=rule,
	199	+ )
	200	+ for op in ops:
	201	+ DVM_OP_REGISTRY[op] = info
	202	+ return func
	203	+
	204	+ return decorator
	205	+
	206	+
	207	+_DEFAULT_OP_INFO = DvmOpInfo(None)
	208	+
	209	+
	210	+def common_rule(node: torch.fx.Node):
	211	+ return _DEFAULT_OP_INFO.is_supported(node)
	212	+
	213	+
	214	+def format_shape(shape):
	215	+ if isinstance(shape, (int, torch.SymInt)):
	216	+ shape = [shape]
	217	+ return "[" + ", ".join(map(str, shape)) + "]"
	218	+
	219	+
	220	+@register_dvm_op(
	221	+ aten.add.Tensor,
	222	+ aten.add.Scalar,
	223	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	224	+ output_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	225	+)
	226	+def add(x, y, alpha=1):
	227	+ if alpha != 1:
	228	+ y = mul(y, alpha)
	229	+ return f"k.add({x}, {y})"
	230	+
	231	+
	232	+@register_dvm_op(
	233	+ aten.sub.Tensor,
	234	+ aten.sub.Scalar,
	235	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	236	+ output_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	237	+)
	238	+def sub(x, y):
	239	+ return f"k.sub({x}, {y})"
	240	+
	241	+
	242	+@register_dvm_op(
	243	+ aten.mul.Tensor,
	244	+ aten.mul.Scalar,
	245	+ input_dtypes=[*DVM_SUPPORT_FLOAT_TYPE, torch.int32],
	246	+ output_dtypes=[*DVM_SUPPORT_FLOAT_TYPE, torch.int32],
	247	+)
	248	+def mul(x, y):
	249	+ return f"k.mul({x}, {y})"
	250	+
	251	+
	252	+@register_dvm_op(aten.div.Tensor, aten.div.Scalar)
	253	+def div(x, y):
	254	+ return f"k.div({x}, {y})"
	255	+
	256	+
	257	+@register_dvm_op(aten.pow.Tensor_Tensor, aten.pow.Tensor_Scalar, aten.pow.Scalar)
	258	+def pow_op(x, y):
	259	+ return f"k.pow({x}, {y})"
	260	+
	261	+
	262	+@register_dvm_op(
	263	+ aten.lt.Tensor,
	264	+ aten.lt.Scalar,
	265	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	266	+ output_dtypes=[torch.bool],
	267	+)
	268	+def less(x, y):
	269	+ return f"k.less({x}, {y})"
	270	+
	271	+
	272	+@register_dvm_op(
	273	+ aten.le.Tensor,
	274	+ aten.le.Scalar,
	275	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	276	+ output_dtypes=[torch.bool],
	277	+)
	278	+def less_equal(x, y):
	279	+ return f"k.less_equal({x}, {y})"
	280	+
	281	+
	282	+@register_dvm_op(
	283	+ aten.gt.Tensor,
	284	+ aten.gt.Scalar,
	285	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	286	+ output_dtypes=[torch.bool],
	287	+)
	288	+def greater(x, y):
	289	+ return f"k.greater({x}, {y})"
	290	+
	291	+
	292	+@register_dvm_op(
	293	+ aten.ge.Tensor,
	294	+ aten.ge.Scalar,
	295	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	296	+ output_dtypes=[torch.bool],
	297	+)
	298	+def greater_equal(x, y):
	299	+ return f"k.greater_equal({x}, {y})"
	300	+
	301	+
	302	+@register_dvm_op(aten.maximum.default)
	303	+def maximum(x, y):
	304	+ return f"k.maximum({x}, {y})"
	305	+
	306	+
	307	+@register_dvm_op(aten.minimum.default)
	308	+def minimum(x, y):
	309	+ return f"k.minimum({x}, {y})"
	310	+
	311	+
	312	+@register_dvm_op(aten.clamp_min.default)
	313	+def clamp_min(x, min_value):
	314	+ return maximum(x, min_value)
	315	+
	316	+
	317	+@register_dvm_op(aten.clamp_max.default)
	318	+def clamp_max(x, max_value):
	319	+ return minimum(x, max_value)
	320	+
	321	+
	322	+@register_dvm_op(
	323	+ aten.logical_and.default,
	324	+ aten.bitwise_and.Tensor,
	325	+ input_dtypes=[torch.bool],
	326	+ output_dtypes=[torch.bool],
	327	+)
	328	+def logical_and(x, y):
	329	+ return f"k.logical_and({x}, {y})"
	330	+
	331	+
	332	+@register_dvm_op(
	333	+ aten.logical_or.default,
	334	+ aten.bitwise_or.Tensor,
	335	+ input_dtypes=[torch.bool],
	336	+ output_dtypes=[torch.bool],
	337	+)
	338	+def logical_or(x, y):
	339	+ return f"k.logical_or({x}, {y})"
	340	+
	341	+
	342	+@register_dvm_op(
	343	+ aten.eq.Tensor,
	344	+ aten.eq.Scalar,
	345	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	346	+ output_dtypes=[torch.bool],
	347	+)
	348	+def equal(x, y):
	349	+ return f"k.equal({x}, {y})"
	350	+
	351	+
	352	+@register_dvm_op(
	353	+ aten.ne.Tensor,
	354	+ aten.ne.Scalar,
	355	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	356	+ output_dtypes=[torch.bool],
	357	+)
	358	+def not_equal(x, y):
	359	+ return f"k.not_equal({x}, {y})"
	360	+
	361	+
	362	+@register_dvm_op(aten.sqrt.default)
	363	+def sqrt(x):
	364	+ return f"k.sqrt({x})"
	365	+
	366	+
	367	+@register_dvm_op(aten.rsqrt.default)
	368	+def rsqrt(x):
	369	+ return div(1, sqrt(x))
	370	+
	371	+
	372	+@register_dvm_op(aten.abs.default)
	373	+def abs_op(x):
	374	+ return f"k.abs({x})"
	375	+
	376	+
	377	+@register_dvm_op(aten.log.default)
	378	+def log(x):
	379	+ return f"k.log({x})"
	380	+
	381	+
	382	+@register_dvm_op(aten.exp.default)
	383	+def exp(x):
	384	+ return f"k.exp({x})"
	385	+
	386	+
	387	+@register_dvm_op(aten.reciprocal.default)
	388	+def reciprocal(x):
	389	+ return f"k.reciprocal({x})"
	390	+
	391	+
	392	+@register_dvm_op(aten.isfinite.default, output_dtypes=[torch.bool])
	393	+def is_finite(x):
	394	+ return f"k.is_finite({x})"
	395	+
	396	+
	397	+@register_dvm_op(
	398	+ aten.logical_not.default,
	399	+ aten.bitwise_not.default,
	400	+ input_dtypes=[torch.bool],
	401	+ output_dtypes=[torch.bool],
	402	+)
	403	+def logical_not(x):
	404	+ return f"k.logical_not({x})"
	405	+
	406	+
	407	+@register_dvm_op(aten.round.default, aten.round.decimals)
	408	+def round_op(x):
	409	+ return f"k.round({x})"
	410	+
	411	+
	412	+@register_dvm_op(aten.floor.default)
	413	+def floor(x):
	414	+ return f"k.floor({x})"
	415	+
	416	+
	417	+@register_dvm_op(aten.ceil.default)
	418	+def ceil(x):
	419	+ return f"k.ceil({x})"
	420	+
	421	+
	422	+@register_dvm_op(aten.trunc.default)
	423	+def trunc(x):
	424	+ return f"k.trunc({x})"
	425	+
	426	+
	427	+@register_dvm_op(
	428	+ aten._to_copy.default,
	429	+ prims.convert_element_type.default,
	430	+ torch.ops.npu.npu_dtype_cast.default,
	431	+ torch.ops.npu.npu_dtype_cast_backward.default,
	432	+ torch.ops.npu._npu_dtype_cast.default,
	433	+ torch.ops.npu._npu_dtype_cast_backward.default,
	434	+ input_dtypes=DVM_SUPPORT_OPTIONAL_INT64_TYPE,
	435	+ output_dtypes=DVM_SUPPORT_OPTIONAL_INT64_TYPE,
	436	+)
	437	+def cast(x, dtype):
	438	+ dtype = to_dvm_dtype(dtype)
	439	+ return f"k.cast({x}, {dtype})"
	440	+
	441	+
	442	+@register_dvm_op(
	443	+ aten.expand.default,
	444	+ input_dtypes=DVM_SUPPORT_TYPE,
	445	+ output_dtypes=DVM_SUPPORT_TYPE,
	446	+)
	447	+def broadcast(x, shape):
	448	+ shape = format_shape(shape)
	449	+ return f"k.broadcast({x}, {shape})"
	450	+
	451	+
	452	+@register_dvm_op(
	453	+ aten.where.default,
	454	+ aten.where.self,
	455	+ input_dtypes=None,
	456	+ output_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	457	+ rule=where_rule,
	458	+)
	459	+def select(x, y, z):
	460	+ return f"k.select({x}, {y}, {z})"
	461	+
	462	+
	463	+@register_dvm_op(aten.sum.dim_IntList, aten.sum.default)
	464	+def reduce_sum(x, dim=None, keepdim=False, dtype=None):
	465	+ if dim is None:
	466	+ dim = []
	467	+ dim = format_shape(dim)
	468	+ return f"k.sum({x}, {dim}, {keepdim})"
	469	+
	470	+
	471	+@register_dvm_op(aten.amax.default)
	472	+def reduce_max(x, dim=None, keepdim=False):
	473	+ if dim is None:
	474	+ dim = []
	475	+ dim = format_shape(dim)
	476	+ return f"k.max({x}, {dim}, {keepdim})"
	477	+
	478	+
	479	+@register_dvm_op(aten.amin.default)
	480	+def reduce_min(x, dim=None, keepdim=False):
	481	+ if dim is None:
	482	+ dim = []
	483	+ dim = format_shape(dim)
	484	+ return f"k.min({x}, {dim}, {keepdim})"
	485	+
	486	+
	487	+@register_dvm_op(
	488	+ aten.view.default,
	489	+ aten.reshape.default,
	490	+ aten._unsafe_view.default,
	491	+ input_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	492	+ output_dtypes=DVM_SUPPORT_FLOAT_INT_TYPE,
	493	+)
	494	+def reshape(x, shape):
	495	+ shape = format_shape(shape)
	496	+ return f"k.reshape({x}, {shape})"
	497	+
	498	+
	499	+@register_dvm_op(aten.neg.default)
	500	+def neg(x):
	501	+ return mul(x, -1)
	502	+
	503	+
	504	+@register_dvm_op(aten.relu.default)
	505	+def relu(x):
	506	+ return maximum(x, 0)
	507	+
	508	+
	509	+def copy(x):
	510	+ return f"k.copy({x})"
	511	+
	512	+
	513	+@register_dvm_op(
	514	+ aten.copy.default,
	515	+ aten.copy_.default,
	516	+ input_dtypes=DVM_SUPPORT_TYPE,
	517	+ output_dtypes=DVM_SUPPORT_TYPE,
	518	+)
	519	+def copy_(dst, src, non_blocking=False):
	520	+ return src
	521	+
	522	+
	523	+@register_dvm_op(aten.clone.default)
	524	+def clone(x, memory_format=None):
	525	+ return copy(x)
	526	+
	527	+
	528	+@register_dvm_op(
	529	+ aten.full.default,
	530	+ output_dtypes=DVM_SUPPORT_TYPE,
	531	+)
	532	+def full(
	533	+ size,
	534	+ fill_value,
	535	+ **kwargs,
	536	+):
	537	+ size = format_shape(size)
	538	+ dtype = to_dvm_dtype(kwargs.get("dtype"))
	539	+ return f"k.full({fill_value}, {size}, {dtype})"
	540	+
	541	+
	542	+@register_dvm_op(aten.mm.default, aten.bmm.default, rule=mm_rule)
	543	+def matmul(x, y, trans_a, trans_b):
	544	+ return f"k.matmul({x}, {y}, {trans_a}, {trans_b})"
	545	+
	546	+
	547	+def matmul_bias(bias, x, y, trans_a, trans_b, beta=1, alpha=1):
	548	+ return f"k.matmul({x}, {y}, {trans_a}, {trans_b},{bias})"
	549	+
	550	+
	551	+@register_dvm_op(aten.addmm.default, rule=mm_rule)
	552	+def addmm(z, x, y, trans_a, trans_b, use_bias, beta=1, alpha=1):
	553	+ if use_bias:
	554	+ return matmul_bias(z, x, y, trans_a, trans_b)
	555	+ if beta != 1:
	556	+ z = mul(z, beta)
	557	+ mm = matmul(x, y, trans_a, trans_b)
	558	+ if alpha != 1:
	559	+ mm = mul(mm, alpha)
	560	+ return add(mm, z)
	561	+
	562	+
	563	+def load(shape, dtype):
	564	+ dtype = to_dvm_dtype(dtype)
	565	+ return f"k.load({shape}, {dtype})"
	566	+
	567	+
	568	+def view_load(shape, stride, dtype):
	569	+ dtype = to_dvm_dtype(dtype)
	570	+ return f"k.view_load({shape}, {stride}, {dtype})"
	571	+
	572	+
	573	+def store(x, dtype=None):
	574	+ if dtype is None:
	575	+ return f"k.store({x})"
	576	+ dtype = to_dvm_dtype(dtype)
	577	+ return f"k.store({x}, {dtype})"

Atorch_npu/_inductor/dvm/util.py+64-0

@@ -0,0 +1,64 @@
	1	+"""DVM load/view_load selection based on real tensor strides."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+from typing import Sequence
	6	+
	7	+import torch
	8	+from torch._inductor.virtualized import V
	9	+
	10	+from .op_emitter import load, view_load
	11	+
	12	+
	13	+def patch_gm_placeholder_strides_from_codegen_args(
	14	+ gm: torch.fx.GraphModule,
	15	+ arg_names: Sequence[str],
	16	+) -> None:
	17	+ """Patch placeholder meta['val'] with Inductor buffer layout strides at codegen time."""
	18	+ placeholders = [n for n in gm.graph.nodes if n.op == "placeholder"]
	19	+ for node, name in zip(placeholders, arg_names):
	20	+ val = node.meta.get("val")
	21	+ if isinstance(val, (torch.SymInt, torch.SymFloat)):
	22	+ continue
	23	+ if not isinstance(val, torch.Tensor):
	24	+ continue
	25	+ buf = V.graph.try_get_buffer(name)
	26	+ if buf is None:
	27	+ continue
	28	+ layout_stride = tuple(buf.get_stride())
	29	+ if layout_stride == tuple(val.stride()):
	30	+ continue
	31	+ node.meta["val"] = torch.empty_strided(
	32	+ val.shape,
	33	+ layout_stride,
	34	+ dtype=val.dtype,
	35	+ device=val.device,
	36	+ )
	37	+
	38	+
	39	+def codegen_maybe_view_load(
	40	+ shape: Sequence,
	41	+ stride: Sequence,
	42	+ dtype: torch.dtype,
	43	+ *,
	44	+ view_fusion_level: int,
	45	+ is_symbolic: bool,
	46	+) -> tuple[str, bool]:
	47	+ """Return (expr, skip_cont).
	48	+
	49	+ skip_cont=False means the caller must manually materialize a non-contiguous
	50	+ input with .contiguous() before launching the DVM kernel.
	51	+ """
	52	+ if view_fusion_level == 0:
	53	+ return load(shape, dtype), False
	54	+
	55	+ if view_fusion_level == 2:
	56	+ return view_load(shape, stride, dtype), True
	57	+
	58	+ if is_symbolic:
	59	+ return load(shape, dtype), False
	60	+
	61	+ if stride[-1] == 1 and shape[-1] != 1:
	62	+ return view_load(shape, stride, dtype), True
	63	+
	64	+ return load(shape, dtype), False

Mtorch_npu/csrc/InitNpuBindings.cpp+2-0

	@@ -157,6 +157,7 @@ void THNPEvent_init(PyObject *module);
157	void THNPGraph_init(PyObject *module);	157	void THNPGraph_init(PyObject *module);
158	void THNPMemPool_init(PyObject* module);	158	void THNPMemPool_init(PyObject* module);
159	void THNPMLIR_init(PyObject* module);	159	void THNPMLIR_init(PyObject* module);
		160	+void THDVM_init(PyObject* module);
160	PyMethodDef* THNPModule_get_methods();	161	PyMethodDef* THNPModule_get_methods();
161		162
162	static std::vector<PyMethodDef> methods;	163	static std::vector<PyMethodDef> methods;
	@@ -199,6 +200,7 @@ PyObject* initModule()
199	THNPGraph_init(module);	200	THNPGraph_init(module);
200	THNPMemPool_init(module);	201	THNPMemPool_init(module);
201	THNPMLIR_init(module);	202	THNPMLIR_init(module);
		203	+ THDVM_init(module);
202		204
203	RegisterNPUDeviceProperties(module);	205	RegisterNPUDeviceProperties(module);
204	BindGetDeviceProperties(module);	206	BindGetDeviceProperties(module);

Mtorch_npu/csrc/inductor/CMakeLists.txt+2-1

	@@ -3,9 +3,10 @@ FILE(GLOB _INDUCTOR_SRCS
3	aoti_runner/*.cpp	3	aoti_runner/*.cpp
4	aoti_torch/*.cpp	4	aoti_torch/*.cpp
5	aoti_torch/generated/*.cpp	5	aoti_torch/generated/*.cpp
		6	+ dvm/*.cpp
6	mlir/*.cpp)	7	mlir/*.cpp)
7		8
8	LIST(APPEND INDUCTOR_SRCS ${_INDUCTOR_SRCS})	9	LIST(APPEND INDUCTOR_SRCS ${_INDUCTOR_SRCS})
9		10
10	# Pass to parent	11	# Pass to parent
11	-set(INDUCTOR_SRCS ${INDUCTOR_SRCS} PARENT_SCOPE)	12	+set(INDUCTOR_SRCS ${INDUCTOR_SRCS} PARENT_SCOPE)

Atorch_npu/csrc/inductor/dvm/pybind_api.cpp+811-0

@@ -0,0 +1,811 @@
	1	+// Copyright (c) 2025 Huawei Technologies Co., Ltd
	2	+// All rights reserved.
	3	+//
	4	+// Licensed under the BSD 3-Clause License (the "License");
	5	+// you may not use this file except in compliance with the License.
	6	+// You may obtain a copy of the License at
	7	+//
	8	+// Unless required by applicable law or agreed to in writing, software
	9	+// distributed under the License is distributed on an "AS IS" BASIS,
	10	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	11	+// See the License for the specific language governing permissions and
	12	+// limitations under the License.
	13	+#ifndef BUILD_LIBTORCH
	14	+#include "torch_npu/csrc/inductor/dvm/pybind_api.h"
	15	+
	16	+#include <algorithm>
	17	+#include <ATen/Context.h>
	18	+#include <c10/core/SymFloat.h>
	19	+#include <torch/csrc/THP.h>
	20	+#include <fstream>
	21	+#include <stdexcept>
	22	+#include <utility>
	23	+#include "torch_npu/csrc/core/npu/NPUGraphsUtils.h"
	24	+#include "torch_npu/csrc/core/npu/NPUStream.h"
	25	+#include "torch_npu/csrc/framework/OpCommand.h"
	26	+
	27	+template <typename T> using shared_ptr_class_ = py::class_<T, std::shared_ptr<T> >;
	28	+
	29	+void TORCH_NPU_API THDVM_init(PyObject* module)
	30	+{
	31	+ using namespace dvm;
	32	+ auto torch_C_m = py::handle(module).cast<py::module>();
	33	+ auto dvm_m = torch_C_m.def_submodule("dvm", "DVM bindings");
	34	+ RegDvmPy(dvm_m);
	35	+
	36	+ pybind11::class_<TorchKernelPy, KernelPy, std::shared_ptr<TorchKernelPy> >(dvm_m, "TorchKernel")
	37	+ .def(py::init<int, uint32_t>())
	38	+ .def("set_kernel_info", &TorchKernelPy::SetKernelInfo, "set_kernel_info")
	39	+ .def("setup", &TorchKernelPy::Setup, "setup")
	40	+ .def("run", &TorchKernelPy::Run, "run kernel")
	41	+ .def("__call__", &TorchKernelPy::Call, "call kernel");
	42	+
	43	+ pybind11::class_<DynKernelPy, TorchKernelPy, std::shared_ptr<DynKernelPy> >(dvm_m, "DynKernel")
	44	+ .def(py::init<int, uint32_t>())
	45	+ .def("scalar", &DynKernelPy::MakeScalar, "setup", py::arg("dtype") = DataTypePy(kDataTypeEnd));
	46	+
	47	+ pybind11::class_<GraphSplitKernelPy, TorchKernelPy, std::shared_ptr<GraphSplitKernelPy> >(dvm_m, "GraphSplitKernel")
	48	+ .def(py::init<>());
	49	+
	50	+ pybind11::class_<DynGraphSplitKernelPy, DynKernelPy, std::shared_ptr<DynGraphSplitKernelPy> >(dvm_m,
	51	+ "DynGraphSplitKernel")
	52	+ .def(py::init<>());
	53	+}
	54	+
	55	+namespace dvm {
	56	+namespace {
	57	+at::ScalarType DvmDType2TorchDtype(DType dtype)
	58	+{
	59	+ switch (dtype) {
	60	+ case DType::kFloat32:
	61	+ return at::ScalarType::Float;
	62	+ case DType::kFloat16:
	63	+ return at::ScalarType::Half;
	64	+ case DType::kBFloat16:
	65	+ return at::ScalarType::BFloat16;
	66	+ case DType::kInt32:
	67	+ return at::ScalarType::Int;
	68	+ case DType::kInt64:
	69	+ return at::ScalarType::Long;
	70	+ case DType::kBool:
	71	+ return at::ScalarType::Bool;
	72	+ default:
	73	+ throw std::runtime_error("Unsupported Dtype conversion");
	74	+ }
	75	+}
	76	+
	77	+bool IsCurrentStreamCapturing()
	78	+{
	79	+ return static_cast<int>(c10_npu::currentStreamCaptureStatusMayInitCtx()) != 0;
	80	+}
	81	+
	82	+class FakeWsAllocator : public WsAllocator {
	83	+public:
	84	+ explicit FakeWsAllocator(TorchKernelPy* kernel_py) : kernel_py_(kernel_py) {}
	85	+
	86	+ void* Alloc(size_t size) override
	87	+ {
	88	+ kernel_py_->SetWorkspaceSize(size);
	89	+ return nullptr;
	90	+ }
	91	+
	92	+private:
	93	+ TorchKernelPy* kernel_py_ = nullptr;
	94	+};
	95	+
	96	+class ExternalWsAllocator : public WsAllocator {
	97	+public:
	98	+ explicit ExternalWsAllocator(void* workspace_ptr) : workspace_ptr_(workspace_ptr) {}
	99	+
	100	+ void* Alloc(size_t) override { return workspace_ptr_; }
	101	+
	102	+private:
	103	+ void* workspace_ptr_ = nullptr;
	104	+};
	105	+
	106	+inline std::pair<at::Tensor, void*> AllocWorkspaceV1(size_t size)
	107	+{
	108	+ if (size == 0) {
	109	+ return {at::Tensor(), nullptr};
	110	+ }
	111	+ at::TensorOptions options = at::TensorOptions(c10::DeviceType::PrivateUse1);
	112	+ auto workspace_tensor = at::empty({static_cast<int64_t>(size)}, options.dtype(at::kByte));
	113	+ auto workspace_ptr = const_cast<void*>(workspace_tensor.storage().data());
	114	+ return {workspace_tensor, workspace_ptr};
	115	+}
	116	+
	117	+inline std::pair<at::Tensor, void*> AllocWorkspaceV2(size_t size, aclrtStream stream)
	118	+{
	119	+ if (size == 0) {
	120	+ return {at::Tensor(), nullptr};
	121	+ }
	122	+ auto workspace_tensor = at_npu::native::allocate_workspace(size, stream);
	123	+ auto workspace_ptr = const_cast<void*>(workspace_tensor.storage().data());
	124	+ return {workspace_tensor, workspace_ptr};
	125	+}
	126	+
	127	+} // namespace
	128	+
	129	+TorchKernelPy::TorchKernelPy(int kernel_type, uint32_t flags)
	130	+ : ws_size_(0), kernel_type_(kernel_type), kernel_flags_(flags)
	131	+{
	132	+ SetDeterm(at::globalContext().deterministicAlgorithms());
	133	+ kernel_.Reset(static_cast<KernelType>(kernel_type), flags);
	134	+}
	135	+
	136	+TorchKernelPy::~TorchKernelPy()
	137	+{
	138	+ for (auto ref : shapes_) {
	139	+ delete ref;
	140	+ }
	141	+}
	142	+
	143	+void TorchKernelPy::SetDeterm(bool enable)
	144	+{
	145	+ auto& conf = Config::Instance();
	146	+ if (enable) {
	147	+ conf.SetDeterm();
	148	+ } else {
	149	+ conf.UnsetDeterm();
	150	+ }
	151	+}
	152	+
	153	+void TorchKernelPy::SetTuning(bool enable)
	154	+{
	155	+ auto& conf = Config::Instance();
	156	+ if (enable) {
	157	+ conf.SetOnlineTuner().SetLazyTuner();
	158	+ } else {
	159	+ conf.UnsetOnlineTuner().UnsetLazyTuner();
	160	+ }
	161	+}
	162	+
	163	+IntArrayRef* TorchKernelPy::GetShapeRef(py::object shape) { return SymIntArraytoShapeRef(shape); }
	164	+
	165	+ShapeRef* TorchKernelPy::SymIntArraytoShapeRef(py::object shape)
	166	+{
	167	+ auto shape_array = shape.cast<py::sequence>();
	168	+ auto& ref = shapes_.emplace_back(new ShapeWithRef(shape_array.size()));
	169	+ for (size_t i = 0; i < ref->size; ++i) {
	170	+ ref->shape_data[i] = shape_array[i].cast<int64_t>();
	171	+ }
	172	+ return ref;
	173	+}
	174	+
	175	+ShapeRef* TorchKernelPy::SymIntArraytoShapeRef(at::IntArrayRef shape_array)
	176	+{
	177	+ auto& ref = shapes_.emplace_back(new ShapeWithRef(shape_array.size()));
	178	+ for (size_t i = 0; i < ref->size; ++i) {
	179	+ ref->shape_data[i] = shape_array[i];
	180	+ }
	181	+ return ref;
	182	+}
	183	+
	184	+py::object TorchKernelPy::Load(py::object shape, DataTypePy type)
	185	+{
	186	+ ShapeRef* shape_ref = SymIntArraytoShapeRef(shape);
	187	+ auto op = kernel_.Load(nullptr, shape_ref, type);
	188	+ loads_.emplace_back(op);
	189	+ load_types_.emplace_back(LoadType::Load);
	190	+ return ObjToPy(op);
	191	+}
	192	+
	193	+py::object TorchKernelPy::GlobalAccess(py::object shape, DataTypePy type)
	194	+{
	195	+ ShapeRef* shape_ref = SymIntArraytoShapeRef(shape);
	196	+ auto op = kernel_.GlobalAccess(nullptr, shape_ref, type);
	197	+ loads_.emplace_back(op);
	198	+ return ObjToPy(op);
	199	+}
	200	+
	201	+py::object TorchKernelPy::ViewLoad(py::object shape, py::object stride, DataTypePy type)
	202	+{
	203	+ ShapeRef* shape_ref = SymIntArraytoShapeRef(shape);
	204	+ ShapeRef* stride_ref = SymIntArraytoShapeRef(stride);
	205	+ auto op = kernel_.Load(nullptr, shape_ref, stride_ref, type);
	206	+ loads_.emplace_back(op);
	207	+ load_types_.emplace_back(LoadType::View);
	208	+ return ObjToPy(op);
	209	+}
	210	+
	211	+py::object TorchKernelPy::GatherLoad(py::object shape, py::object index, DataTypePy type, int axis)
	212	+{
	213	+ ShapeRef* shape_ref = SymIntArraytoShapeRef(shape);
	214	+ auto op = kernel_.GatherLoad(nullptr, shape_ref, PyToObj(index), axis, type);
	215	+ loads_.emplace_back(op);
	216	+ return ObjToPy(op);
	217	+}
	218	+
	219	+py::object TorchKernelPy::Store(py::object obj, DataTypePy type)
	220	+{
	221	+ auto in_obj = PyToObj(obj);
	222	+ if (type != kDataTypeEnd) {
	223	+ in_obj = kernel_.Cast(in_obj, type);
	224	+ }
	225	+ auto op = kernel_.Store(nullptr, in_obj);
	226	+ return ObjToPy(stores_.emplace_back(op));
	227	+}
	228	+
	229	+py::object TorchKernelPy::ViewStore(py::object obj, py::object stride, DataTypePy type)
	230	+{
	231	+ auto in_obj = PyToObj(obj);
	232	+ if (type != kDataTypeEnd) {
	233	+ in_obj = kernel_.Cast(in_obj, type);
	234	+ }
	235	+ auto op = kernel_.Store(nullptr, in_obj, GetShapeRef(stride));
	236	+ return ObjToPy(stores_.emplace_back(op));
	237	+}
	238	+
	239	+void TorchKernelPy::Setup()
	240	+{
	241	+ SetupRelocs();
	242	+ ws_size_ = kernel_.CodeGen();
	243	+}
	244	+
	245	+void TorchKernelPy::SetupRelocs()
	246	+{
	247	+ relocs_.clear();
	248	+ relocs_.reserve(loads_.size() + stores_.size());
	249	+ for (auto op : loads_) {
	250	+ relocs_.emplace_back(op, nullptr);
	251	+ }
	252	+ for (auto op : stores_) {
	253	+ relocs_.emplace_back(op, nullptr);
	254	+ }
	255	+}
	256	+
	257	+TorchKernelPy::ParsedCallInputs TorchKernelPy::ParseTensorCallInputs(py::args inputs,
	258	+ std::vector<at::Tensor>& tensor_refs) const
	259	+{
	260	+ TORCH_CHECK(inputs.size() == loads_.size(), "Call expects ", loads_.size(), " input tensors, got ", inputs.size());
	261	+
	262	+ auto addr = std::make_shared<std::vector<void*> >();
	263	+ addr->resize(relocs_.size());
	264	+ tensor_refs.reserve(loads_.size());
	265	+ at::TensorOptions options = at::TensorOptions(c10::DeviceType::PrivateUse1);
	266	+
	267	+ for (size_t i = 0; i < loads_.size(); ++i) {
	268	+ auto tensor = inputs[i].cast<at::Tensor>();
	269	+ if (tensor.device().type() == c10::DeviceType::CPU) {
	270	+ tensor = at_npu::native::OpPreparation::copy_tensor_host_to_device(tensor);
	271	+ }
	272	+ (*addr)[i] = tensor.data_ptr();
	273	+ options = tensor.options();
	274	+ tensor_refs.emplace_back(tensor);
	275	+ }
	276	+ return {addr, options};
	277	+}
	278	+
	279	+void TorchKernelPy::Run(py::args args)
	280	+{
	281	+ const auto num_inputs = loads_.size();
	282	+ const auto num_outputs = stores_.size();
	283	+ TORCH_CHECK(args.size() == num_inputs + num_outputs, "DVM kernel run expects ", num_inputs + num_outputs,
	284	+ " tensors, got ", args.size());
	285	+ auto addr = std::make_shared<std::vector<void*> >();
	286	+ tensor_list_.reserve(args.size());
	287	+ out_refs_.reserve(num_outputs);
	288	+ addr->resize(num_inputs + num_outputs);
	289	+ for (size_t i = 0; i < args.size(); i++) {
	290	+ auto tensor = args[i].cast<at::Tensor>();
	291	+ if (i < num_inputs && tensor.device().type() == c10::DeviceType::CPU) {
	292	+ tensor = at_npu::native::OpPreparation::copy_tensor_host_to_device(tensor);
	293	+ }
	294	+ if (i < num_inputs) {
	295	+ if (!tensor.is_contiguous() && load_types_[i] == LoadType::Load) {
	296	+ tensor = tensor.contiguous();
	297	+ }
	298	+ } else if (!tensor.is_contiguous()) {
	299	+ tensor = out_refs_.emplace_back(tensor, at::empty(tensor.sizes(), tensor.options())).second;
	300	+ }
	301	+ tensor_list_.emplace_back(tensor);
	302	+ (*addr)[i] = tensor.data_ptr();
	303	+ }
	304	+
	305	+ aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
	306	+ if (!IsCurrentStreamCapturing()) {
	307	+ auto dvm_call = [this, addr, stream]() { return LaunchV2(addr->data(), stream); };
	308	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, dvm_call);
	309	+ } else {
	310	+ auto [workspace_tensor, workspace_ptr] = AllocWorkspaceV1(ws_size_);
	311	+ auto dvm_call = [this, addr, stream, workspace_ptr]() {
	312	+ return LaunchV1(addr->data(), stream, workspace_ptr);
	313	+ };
	314	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, dvm_call);
	315	+ }
	316	+
	317	+ for (auto& [ori_tensor, cur_tensor] : out_refs_) {
	318	+ ori_tensor.copy_(cur_tensor);
	319	+ }
	320	+ tensor_list_.clear();
	321	+ out_refs_.clear();
	322	+}
	323	+
	324	+py::object TorchKernelPy::Call(py::args args)
	325	+{
	326	+ const auto num_inputs = loads_.size();
	327	+ const auto num_outputs = stores_.size();
	328	+ std::vector<at::Tensor> tensor_list;
	329	+ auto parsed = ParseTensorCallInputs(args, tensor_list);
	330	+ auto ret = CreateOutputs(parsed.options, parsed.addr->data() + num_inputs);
	331	+ aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
	332	+ if (!IsCurrentStreamCapturing()) {
	333	+ auto dvm_call = [this, addr = parsed.addr, stream]() { return LaunchV2(addr->data(), stream); };
	334	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, dvm_call);
	335	+ } else {
	336	+ auto [workspace_tensor, workspace_ptr] = AllocWorkspaceV1(ws_size_);
	337	+ auto dvm_call = [this, addr = parsed.addr, stream, workspace_ptr]() {
	338	+ return LaunchV1(addr->data(), stream, workspace_ptr);
	339	+ };
	340	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, dvm_call);
	341	+ }
	342	+ return ret;
	343	+}
	344	+
	345	+py::object TorchKernelPy::CreateOutputs(const at::TensorOptions& options, void** addr)
	346	+{
	347	+ auto create_output = [this, options](NDObject* store) -> at::Tensor {
	348	+ auto shape_ref = kernel_.GetShape(store);
	349	+ c10::IntArrayRef shape = c10::IntArrayRef(shape_ref->data, shape_ref->size);
	350	+ at::ScalarType dtype = DvmDType2TorchDtype(kernel_.GetDType(store));
	351	+ return at_npu::native::OpPreparation::apply_tensor_without_format(shape, options.dtype(dtype));
	352	+ };
	353	+ if (stores_.size() == 1) {
	354	+ at::Tensor tensor = create_output(stores_.front());
	355	+ *addr = tensor.data_ptr();
	356	+ return py::cast(tensor);
	357	+ }
	358	+ py::tuple tuple_ret(stores_.size());
	359	+ for (size_t i = 0; i < stores_.size(); ++i) {
	360	+ at::Tensor tensor = create_output(stores_[i]);
	361	+ *addr = tensor.data_ptr();
	362	+ addr++;
	363	+ tuple_ret[i] = py::cast(tensor);
	364	+ }
	365	+ return tuple_ret;
	366	+}
	367	+
	368	+int TorchKernelPy::LaunchV2(void** addr, aclrtStream stream)
	369	+{
	370	+ auto workspace_pair = AllocWorkspaceV2(ws_size_, stream);
	371	+ auto workspace_ptr = workspace_pair.second;
	372	+ return LaunchV1(addr, stream, workspace_ptr);
	373	+}
	374	+
	375	+int TorchKernelPy::LaunchV1(void** addr, aclrtStream stream, void* workspace_ptr)
	376	+{
	377	+ for (size_t i = 0; i < relocs_.size(); ++i) {
	378	+ relocs_[i].addr = addr[i];
	379	+ }
	380	+ return kernel_.Launch(relocs_.data(), relocs_.size(), workspace_ptr, stream);
	381	+}
	382	+
	383	+void* GraphSplitBase::Alloc(size_t size)
	384	+{
	385	+ auto [workspace_tensor, workspace_ptr] = AllocWorkspaceV2(size, stream_);
	386	+ ws_ = std::move(workspace_tensor);
	387	+ return workspace_ptr;
	388	+}
	389	+
	390	+int GraphSplitBase::LaunchV1(Kernel& kernel, void** addr, aclrtStream stream, std::vector<RelocEntry>& relocs,
	391	+ void* workspace_ptr)
	392	+{
	393	+ for (size_t i = 0; i < relocs.size(); ++i) {
	394	+ relocs[i].addr = addr[i];
	395	+ }
	396	+ ExternalWsAllocator allocator(workspace_ptr);
	397	+ kernel.CodeGen(relocs.data(), relocs.size(), &allocator);
	398	+ int ret = kernel.Launch(stream);
	399	+ return ret;
	400	+}
	401	+
	402	+int GraphSplitBase::LaunchV2(Kernel& kernel, void** addr, aclrtStream stream, std::vector<RelocEntry>& relocs)
	403	+{
	404	+ stream_ = stream;
	405	+ for (size_t i = 0; i < relocs.size(); ++i) {
	406	+ relocs[i].addr = addr[i];
	407	+ }
	408	+ kernel.CodeGen(relocs.data(), relocs.size(), this);
	409	+ int ret = kernel.Launch(stream);
	410	+ ws_.reset();
	411	+ return ret;
	412	+}
	413	+
	414	+void GraphSplitKernelPy::Setup()
	415	+{
	416	+ SetupRelocs();
	417	+ kernel_.Infer();
	418	+ SetWorkspaceSize(0);
	419	+ FakeWsAllocator fake_alloc(this);
	420	+ kernel_.CodeGen(nullptr, 0, &fake_alloc);
	421	+}
	422	+
	423	+void GraphSplitKernelPy::Run(py::args)
	424	+{
	425	+ TORCH_CHECK(false, "GraphSplitKernel::Run is unsupported. Use Call/__call__ to create outputs internally.");
	426	+}
	427	+
	428	+py::object GraphSplitKernelPy::Call(py::args inputs)
	429	+{
	430	+ std::vector<at::Tensor> tensor_list;
	431	+ auto parsed = ParseTensorCallInputs(inputs, tensor_list);
	432	+ auto ret = CreateOutputs(parsed.options, parsed.addr->data() + loads_.size());
	433	+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
	434	+ if (!IsCurrentStreamCapturing()) {
	435	+ auto launch_call = [this, stream, addr = parsed.addr]() -> int {
	436	+ return GraphSplitBase::LaunchV2(kernel_, addr->data(), stream, relocs_);
	437	+ };
	438	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, launch_call);
	439	+ } else {
	440	+ auto [workspace_tensor, workspace_ptr] = AllocWorkspaceV1(ws_size_);
	441	+ auto launch_call = [this, stream, addr = parsed.addr, workspace_ptr]() -> int {
	442	+ return GraphSplitBase::LaunchV1(kernel_, addr->data(), stream, relocs_, workspace_ptr);
	443	+ };
	444	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, launch_call);
	445	+ }
	446	+ return ret;
	447	+}
	448	+
	449	+std::unique_ptr<DynKernelPy> DynGraphSplitKernelPy::CloneExecutor() const
	450	+{
	451	+ auto executor = std::make_unique<DynGraphSplitKernelPy>();
	452	+ CloneExecutorStateTo(*executor);
	453	+ return executor;
	454	+}
	455	+
	456	+void DynGraphSplitKernelPy::Setup() { DynKernelPy::Setup(); }
	457	+
	458	+void DynGraphSplitKernelPy::Run(py::args)
	459	+{
	460	+ TORCH_CHECK(false, "DynGraphSplitKernel::Run is unsupported. Use Call/__call__ to create outputs internally.");
	461	+}
	462	+
	463	+py::object DynGraphSplitKernelPy::Call(py::args inputs)
	464	+{
	465	+ TORCH_CHECK(!IsCurrentStreamCapturing(),
	466	+ "DynGraphSplitKernel does not support stream capture: dynamic shape requires runtime Infer.");
	467	+ auto* executor = static_cast<DynGraphSplitKernelPy*>(AcquireExecutor());
	468	+ std::vector<at::Tensor> tensor_list;
	469	+ auto parsed = executor->ParseDynCallInputs(inputs, tensor_list);
	470	+
	471	+ executor->UpdateSymShapeData();
	472	+ executor->kernel_.Infer();
	473	+ auto ret = executor->CreateOutputs(parsed.options, parsed.addr->data() + executor->loads_.size());
	474	+ auto stream = c10_npu::getCurrentNPUStream().stream(false);
	475	+ auto launch_call = [this, executor, stream, addr = parsed.addr]() -> int {
	476	+ int ret = executor->GraphSplitBase::LaunchV2(executor->kernel_, addr->data(), stream, executor->relocs_);
	477	+ ReleaseExecutor(executor);
	478	+ return ret;
	479	+ };
	480	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, launch_call);
	481	+ return ret;
	482	+}
	483	+
	484	+DynKernelPy::~DynKernelPy()
	485	+{
	486	+ for (auto ref : dyn_load_shapes_) {
	487	+ delete ref;
	488	+ }
	489	+}
	490	+
	491	+DynKernelPy::LoadShapeRef* DynKernelPy::GetDynLoadShapeRef(size_t dim_size)
	492	+{
	493	+ static int64_t init_dyn_data[ShapeWithRef::MAX_SIZE] = {-1};
	494	+ auto ref = new LoadShapeRef();
	495	+ ref->shape.data = init_dyn_data;
	496	+ ref->shape.size = dim_size;
	497	+ ref->stride.data = nullptr;
	498	+ ref->stride.size = 0;
	499	+ dyn_load_shapes_.push_back(ref);
	500	+ return ref;
	501	+}
	502	+
	503	+py::object DynKernelPy::Load(py::object shape, DataTypePy type)
	504	+{
	505	+ auto shape_seq = shape.cast<py::sequence>();
	506	+ auto ref = GetDynLoadShapeRef(shape_seq.size());
	507	+ ShapeRef* shape_ref = &ref->shape;
	508	+ auto op = kernel_.Load(nullptr, shape_ref, type);
	509	+ loads_.emplace_back(op);
	510	+ load_types_.emplace_back(LoadType::Load);
	511	+ return ObjToPy(op);
	512	+}
	513	+
	514	+py::object DynKernelPy::GlobalAccess(py::object shape, DataTypePy type)
	515	+{
	516	+ auto shape_seq = shape.cast<py::sequence>();
	517	+ auto ref = GetDynLoadShapeRef(shape_seq.size());
	518	+ ShapeRef* shape_ref = &ref->shape;
	519	+ auto op = kernel_.GlobalAccess(nullptr, shape_ref, type);
	520	+ loads_.emplace_back(op);
	521	+ return ObjToPy(op);
	522	+}
	523	+
	524	+py::object DynKernelPy::ViewLoad(py::object shape, py::object stride, DataTypePy type)
	525	+{
	526	+ auto shape_seq = shape.cast<py::sequence>();
	527	+ auto ref = GetDynLoadShapeRef(shape_seq.size());
	528	+ ref->stride = ref->shape;
	529	+ ShapeRef* shape_ref = &ref->shape;
	530	+ ShapeRef* stride_ref = &ref->stride;
	531	+ auto op = kernel_.Load(nullptr, shape_ref, stride_ref, type);
	532	+ loads_.emplace_back(op);
	533	+ load_types_.emplace_back(LoadType::View);
	534	+ return ObjToPy(op);
	535	+}
	536	+
	537	+py::object DynKernelPy::GatherLoad(py::object shape, py::object index, DataTypePy type, int axis)
	538	+{
	539	+ auto shape_seq = shape.cast<py::sequence>();
	540	+ auto ref = GetDynLoadShapeRef(shape_seq.size());
	541	+ ShapeRef* shape_ref = &ref->shape;
	542	+ auto op = kernel_.GatherLoad(nullptr, shape_ref, PyToObj(index), axis, type);
	543	+ loads_.emplace_back(op);
	544	+ return ObjToPy(op);
	545	+}
	546	+
	547	+ShapeRef* DynKernelPy::SymIntArraytoShapeRef(py::object shape)
	548	+{
	549	+ auto shape_array = shape.cast<py::sequence>();
	550	+ auto& ref = shapes_.emplace_back(new ShapeWithRef(shape_array.size()));
	551	+ auto& sym_shape = sym_shape_.emplace_back();
	552	+ for (size_t i = 0; i < ref->size; ++i) {
	553	+ ref->shape_data[i] = -1;
	554	+ if (py::isinstance<py::int_>(shape_array[i])) {
	555	+ auto sym_ptr = const_input_.emplace_back(std::make_shared<ScalarRefPy>());
	556	+ sym_ptr->data_ = shape_array[i].cast<int64_t>();
	557	+ sym_shape.emplace_back(sym_ptr);
	558	+ } else {
	559	+ sym_shape.emplace_back(shape_array[i].cast<ScalarRefPyPtr>());
	560	+ }
	561	+ }
	562	+ return ref;
	563	+}
	564	+
	565	+void DynKernelPy::UpdateSymShapeData()
	566	+{
	567	+ for (size_t i = 0; i < sym_shape_.size(); i++) {
	568	+ for (size_t j = 0; j < shapes_[i]->size; j++) {
	569	+ shapes_[i]->shape_data[j] = sym_shape_[i][j]->data_.i64;
	570	+ }
	571	+ }
	572	+}
	573	+
	574	+std::unique_ptr<DynKernelPy> DynKernelPy::CloneExecutor() const
	575	+{
	576	+ auto executor = std::make_unique<DynKernelPy>(kernel_type_, kernel_flags_);
	577	+ CloneExecutorStateTo(*executor);
	578	+ return executor;
	579	+}
	580	+
	581	+void DynKernelPy::CloneExecutorStateTo(DynKernelPy& executor) const
	582	+{
	583	+ SplitCloneHelper helper;
	584	+ std::unordered_map<ScalarRef*, ScalarRefPyPtr> scalar_to_owner;
	585	+
	586	+ executor.ws_size_ = ws_size_;
	587	+ executor.op_name_ = op_name_;
	588	+ executor.op_fullname_ = op_fullname_;
	589	+ executor.load_types_ = load_types_;
	590	+ executor.kernel_.SetNameHint(executor.op_name_.c_str(), executor.op_fullname_.c_str());
	591	+
	592	+ executor.shapes_.reserve(shapes_.size());
	593	+ for (auto ref : shapes_) {
	594	+ auto clone_ref = new ShapeWithRef(ref->size);
	595	+ for (size_t i = 0; i < ref->size; ++i) {
	596	+ clone_ref->shape_data[i] = ref->shape_data[i];
	597	+ }
	598	+ executor.shapes_.push_back(clone_ref);
	599	+ helper.ref_map_[ref] = clone_ref;
	600	+ }
	601	+
	602	+ executor.dyn_load_shapes_.reserve(dyn_load_shapes_.size());
	603	+ for (auto ref : dyn_load_shapes_) {
	604	+ auto clone_ref = new LoadShapeRef();
	605	+ clone_ref->shape = ref->shape;
	606	+ clone_ref->stride = ref->stride;
	607	+ executor.dyn_load_shapes_.push_back(clone_ref);
	608	+ helper.ref_map_[&ref->shape] = &clone_ref->shape;
	609	+ helper.ref_map_[&ref->stride] = &clone_ref->stride;
	610	+ }
	611	+
	612	+ auto clone_scalar = [&helper, &scalar_to_owner](const ScalarRefPyPtr& src, std::vector<ScalarRefPyPtr>& dst) {
	613	+ auto clone = std::make_shared<ScalarRefPy>();
	614	+ clone->data_ = src->data_;
	615	+ helper.ref_map_[&src->data_] = &clone->data_;
	616	+ scalar_to_owner[&clone->data_] = clone;
	617	+ dst.push_back(clone);
	618	+ };
	619	+ for (const auto& scalar : const_input_) {
	620	+ clone_scalar(scalar, executor.const_input_);
	621	+ }
	622	+ for (const auto& scalar : sym_scalar_input_) {
	623	+ clone_scalar(scalar, executor.sym_scalar_input_);
	624	+ }
	625	+
	626	+ executor.kernel_.Clone(kernel_, helper);
	627	+ executor.loads_.reserve(loads_.size());
	628	+ for (auto op : loads_) {
	629	+ executor.loads_.push_back(helper.GetClone(op));
	630	+ }
	631	+ executor.stores_.reserve(stores_.size());
	632	+ for (auto op : stores_) {
	633	+ executor.stores_.push_back(helper.GetClone(op));
	634	+ }
	635	+
	636	+ executor.sym_shape_.reserve(sym_shape_.size());
	637	+ for (const auto& shape_scalars : sym_shape_) {
	638	+ auto& clone_shape_scalars = executor.sym_shape_.emplace_back();
	639	+ clone_shape_scalars.reserve(shape_scalars.size());
	640	+ for (const auto& scalar : shape_scalars) {
	641	+ auto it = helper.ref_map_.find(&scalar->data_);
	642	+ TORCH_CHECK(it != helper.ref_map_.end(), "Failed to clone symbolic shape scalar reference.");
	643	+ auto clone_scalar_ref = static_cast<ScalarRef*>(it->second);
	644	+ auto owner_it = scalar_to_owner.find(clone_scalar_ref);
	645	+ TORCH_CHECK(owner_it != scalar_to_owner.end(), "Failed to remap cloned symbolic shape scalar.");
	646	+ clone_shape_scalars.push_back(owner_it->second);
	647	+ }
	648	+ }
	649	+
	650	+ executor.SetupRelocs();
	651	+}
	652	+
	653	+TorchKernelPy::ParsedCallInputs DynKernelPy::ParseDynCallInputs(py::args inputs,
	654	+ std::vector<at::Tensor>& tensor_refs) const
	655	+{
	656	+ TORCH_CHECK(inputs.size() == dyn_load_shapes_.size() + sym_scalar_input_.size(), "Dynamic call expects ",
	657	+ dyn_load_shapes_.size() + sym_scalar_input_.size(), " args, got ", inputs.size());
	658	+
	659	+ auto addr = std::make_shared<std::vector<void*> >();
	660	+ addr->resize(relocs_.size());
	661	+ tensor_refs.reserve(dyn_load_shapes_.size());
	662	+ at::TensorOptions options = at::TensorOptions(c10::DeviceType::PrivateUse1);
	663	+
	664	+ size_t input_index = 0;
	665	+ size_t sym_scalar_index = 0;
	666	+ for (size_t i = 0; i < inputs.size(); ++i) {
	667	+ if (THPVariable_Check(inputs[i].ptr())) {
	668	+ auto tensor = inputs[i].cast<at::Tensor>();
	669	+ if (tensor.device().type() == c10::DeviceType::CPU) {
	670	+ tensor = at_npu::native::OpPreparation::copy_tensor_host_to_device(tensor);
	671	+ }
	672	+ (*addr)[input_index] = tensor.data_ptr();
	673	+ auto ref = dyn_load_shapes_[input_index++];
	674	+ if (ref->stride.data) {
	675	+ ref->stride.data = tensor.strides().data();
	676	+ }
	677	+ ref->shape.data = tensor.sizes().data();
	678	+ options = tensor.options();
	679	+ tensor_refs.emplace_back(tensor);
	680	+ } else if (py::isinstance<c10::SymFloat>(inputs[i])) {
	681	+ sym_scalar_input_[sym_scalar_index++]->data_ =
	682	+ static_cast<float>(inputs[i].cast<c10::SymFloat>().expect_float());
	683	+ } else if (py::isinstance<c10::SymInt>(inputs[i])) {
	684	+ sym_scalar_input_[sym_scalar_index++]->data_ = inputs[i].cast<c10::SymInt>().expect_int();
	685	+ } else if (py::isinstance<py::float_>(inputs[i])) {
	686	+ sym_scalar_input_[sym_scalar_index++]->data_ = inputs[i].cast<float>();
	687	+ } else if (py::isinstance<py::int_>(inputs[i])) {
	688	+ sym_scalar_input_[sym_scalar_index++]->data_ = inputs[i].cast<int64_t>();
	689	+ } else {
	690	+ const char* py_type = inputs[i].ptr() ? Py_TYPE(inputs[i].ptr())->tp_name : "<null>";
	691	+ TORCH_CHECK(
	692	+ false, "Unsupported dynamic input type at arg[", i,
	693	+ "]. Expected one of: Tensor, c10::SymFloat, c10::SymInt, float, int. Got Python type: ", py_type);
	694	+ }
	695	+ }
	696	+
	697	+ TORCH_CHECK(input_index == dyn_load_shapes_.size(), "Dynamic call expects ", dyn_load_shapes_.size(),
	698	+ " tensor inputs, got ", input_index);
	699	+ TORCH_CHECK(sym_scalar_index == sym_scalar_input_.size(), "Dynamic call expects ", sym_scalar_input_.size(),
	700	+ " scalar inputs, got ", sym_scalar_index);
	701	+ return {addr, options};
	702	+}
	703	+
	704	+void DynKernelPy::Setup()
	705	+{
	706	+ SetupRelocs();
	707	+ dyn_executors_.push_back(this);
	708	+}
	709	+
	710	+void DynKernelPy::Run(py::args args)
	711	+{
	712	+ TORCH_CHECK(!IsCurrentStreamCapturing(),
	713	+ "DynKernel does not support stream capture: dynamic shape requires runtime CodeGen.");
	714	+ const auto num_inputs = loads_.size();
	715	+ const auto num_outputs = stores_.size();
	716	+ const auto num_sym = sym_scalar_input_.size();
	717	+ TORCH_CHECK(args.size() == num_inputs + num_outputs + num_sym, "DynKernel expects ",
	718	+ num_inputs + num_outputs + num_sym, " args, got ", args.size());
	719	+ auto info = std::make_shared<DynamicInfo>();
	720	+ info->shape.reserve(num_inputs);
	721	+ info->strides.reserve(num_inputs);
	722	+ info->scalars.reserve(num_sym);
	723	+ info->addr.reserve(num_inputs + num_outputs);
	724	+ tensor_list_.reserve(num_inputs + num_outputs);
	725	+ out_refs_.reserve(num_outputs);
	726	+ for (size_t i = 0; i < args.size(); i++) {
	727	+ if (THPVariable_Check(args[i].ptr())) {
	728	+ auto tensor = args[i].cast<at::Tensor>();
	729	+ if (i < num_inputs + num_sym && tensor.device().type() == c10::DeviceType::CPU) {
	730	+ tensor = at_npu::native::OpPreparation::copy_tensor_host_to_device(tensor);
	731	+ }
	732	+ if (i < num_inputs + num_sym) {
	733	+ if (!tensor.is_contiguous() && load_types_[i] == LoadType::Load) {
	734	+ tensor = tensor.contiguous();
	735	+ }
	736	+ } else if (!tensor.is_contiguous()) {
	737	+ tensor = out_refs_.emplace_back(tensor, at::empty(tensor.sizes(), tensor.options())).second;
	738	+ }
	739	+ if (i < num_inputs + num_sym) {
	740	+ info->shape.emplace_back(tensor.sizes().vec());
	741	+ info->strides.emplace_back(tensor.strides().vec());
	742	+ }
	743	+ tensor_list_.emplace_back(tensor);
	744	+ info->addr.emplace_back(tensor.data_ptr());
	745	+ } else if (py::isinstance<c10::SymFloat>(args[i])) {
	746	+ info->scalars.emplace_back(static_cast<float>(args[i].cast<c10::SymFloat>().expect_float()));
	747	+ } else if (py::isinstance<c10::SymInt>(args[i])) {
	748	+ info->scalars.emplace_back(args[i].cast<c10::SymInt>().expect_int());
	749	+ } else if (py::isinstance<py::float_>(args[i])) {
	750	+ info->scalars.emplace_back(args[i].cast<float>());
	751	+ } else if (py::isinstance<py::int_>(args[i])) {
	752	+ info->scalars.emplace_back(args[i].cast<int64_t>());
	753	+ } else {
	754	+ TORCH_CHECK(false, "Unsupported dynamic input type for DynKernel.");
	755	+ }
	756	+ }
	757	+
	758	+ aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
	759	+
	760	+ auto dvm_call = [this, info, stream]() {
	761	+ for (size_t i = 0; i < loads_.size(); i++) {
	762	+ auto& ref = dyn_load_shapes_[i];
	763	+ if (ref->stride.data) {
	764	+ ref->stride.data = info->strides[i].data();
	765	+ }
	766	+ ref->shape.data = info->shape[i].data();
	767	+ }
	768	+ for (size_t i = 0; i < sym_scalar_input_.size(); i++) {
	769	+ if (info->scalars[i].type == kInt64) {
	770	+ sym_scalar_input_[i]->data_ = info->scalars[i].i64;
	771	+ } else {
	772	+ sym_scalar_input_[i]->data_ = info->scalars[i].f32;
	773	+ }
	774	+ }
	775	+ UpdateSymShapeData();
	776	+ ws_size_ = kernel_.CodeGen();
	777	+ return LaunchV2(info->addr.data(), stream);
	778	+ };
	779	+
	780	+ at_npu::native::OpCommand::RunOpApiV2(op_name_, dvm_call);
	781	+
	782	+ for (auto& [ori_tensor, cur_tensor] : out_refs_) {
	783	+ ori_tensor.copy_(cur_tensor);
	784	+ }
	785	+ tensor_list_.clear();
	786	+ out_refs_.clear();
	787	+}
	788	+
	789	+py::object DynKernelPy::Call(py::args args)
	790	+{
	791	+ TORCH_CHECK(!IsCurrentStreamCapturing(),
	792	+ "DynKernel does not support stream capture: dynamic shape requires runtime CodeGen.");
	793	+ auto* executor = AcquireExecutor();
	794	+ const auto num_inputs = executor->loads_.size();
	795	+ std::vector<at::Tensor> tensor_list;
	796	+ auto parsed = executor->ParseDynCallInputs(args, tensor_list);
	797	+
	798	+ executor->UpdateSymShapeData();
	799	+ executor->ws_size_ = executor->kernel_.CodeGen();
	800	+ auto ret = executor->CreateOutputs(parsed.options, parsed.addr->data() + num_inputs);
	801	+ aclrtStream stream = c10_npu::getCurrentNPUStream().stream(false);
	802	+ auto dvm_call = [this, executor, addr = parsed.addr, stream]() {
	803	+ int ret = executor->LaunchV2(addr->data(), stream);
	804	+ ReleaseExecutor(executor);
	805	+ return ret;
	806	+ };
	807	+ at_npu::native::OpCommand::RunOpApiV2(executor->op_name_, dvm_call);
	808	+ return ret;
	809	+}
	810	+} // namespace dvm
	811	+#endif // BUILD_LIBTORCH

Atorch_npu/csrc/inductor/dvm/pybind_api.h+236-0

@@ -0,0 +1,236 @@
	1	+// Copyright (c) 2025 Huawei Technologies Co., Ltd
	2	+// All rights reserved.
	3	+//
	4	+// Licensed under the BSD 3-Clause License (the "License");
	5	+// you may not use this file except in compliance with the License.
	6	+// You may obtain a copy of the License at
	7	+//
	8	+// Unless required by applicable law or agreed to in writing, software
	9	+// distributed under the License is distributed on an "AS IS" BASIS,
	10	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	11	+// See the License for the specific language governing permissions and
	12	+// limitations under the License.
	13	+
	14	+#ifndef _DVM_PYBIND_API_H_
	15	+#define _DVM_PYBIND_API_H_
	16	+#ifndef BUILD_LIBTORCH
	17	+#include <c10/core/SymFloat.h>
	18	+#include <c10/core/SymInt.h>
	19	+#include <cstdint>
	20	+#include <memory>
	21	+#include <mutex>
	22	+#include <utility>
	23	+#include <unordered_map>
	24	+#include <vector>
	25	+#include <torch/csrc/THP.h>
	26	+#include <torch/csrc/utils/pybind.h>
	27	+#include <torch/csrc/python_headers.h>
	28	+#include "torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h"
	29	+#include "torch_npu/csrc/framework/utils/OpPreparation.h"
	30	+#include "third_party/dvm/dvm/include/dvm.h"
	31	+#include "third_party/dvm/dvm/include/dvm_py.h"
	32	+
	33	+
	34	+namespace dvm {
	35	+struct SplitCloneHelper : public dvm::CloneHelper {
	36	+ IntArrayRef* GetClone(IntArrayRef* shape) override
	37	+ {
	38	+ auto it = ref_map_.find(shape);
	39	+ return it != ref_map_.end() ? static_cast<IntArrayRef*>(it->second) : shape;
	40	+ }
	41	+
	42	+ ScalarRef* GetClone(ScalarRef* scalar) override
	43	+ {
	44	+ auto it = ref_map_.find(scalar);
	45	+ return it != ref_map_.end() ? static_cast<ScalarRef*>(it->second) : scalar;
	46	+ }
	47	+
	48	+ NDObject* GetClone(NDObject* op) override
	49	+ {
	50	+ auto it = op_map_.find(op);
	51	+ return it != op_map_.end() ? it->second : nullptr;
	52	+ }
	53	+
	54	+ void SetClone(NDObject* op, NDObject* clone) override { op_map_[op] = clone; }
	55	+
	56	+ std::unordered_map<NDObject, NDObject> op_map_;
	57	+ std::unordered_map<void, void> ref_map_;
	58	+};
	59	+class TorchKernelPy : public KernelPy {
	60	+public:
	61	+ TorchKernelPy(int kernel_type, uint32_t flags);
	62	+ ~TorchKernelPy();
	63	+
	64	+ py::object Load(py::object shape, DataTypePy type) override;
	65	+ py::object GlobalAccess(py::object shape, DataTypePy type) override;
	66	+ py::object ViewLoad(py::object shape, py::object stride, DataTypePy type) override;
	67	+ py::object GatherLoad(py::object shape, py::object index, DataTypePy type, int axis) override;
	68	+ py::object Store(py::object obj, DataTypePy type) override;
	69	+ py::object ViewStore(py::object obj, py::object stride, DataTypePy type) override;
	70	+ IntArrayRef* GetShapeRef(py::object shape) override;
	71	+
	72	+ void SetKernelInfo(const std::string& op_name, const std::string& op_fullname)
	73	+ {
	74	+ op_name_ = op_name;
	75	+ op_fullname_ = op_fullname;
	76	+ kernel_.SetNameHint(op_name_.c_str(), op_fullname_.c_str());
	77	+ }
	78	+
	79	+ int LaunchV1(void** addr, aclrtStream stream, void* workspace_ptr);
	80	+ int LaunchV2(void** addr, aclrtStream stream);
	81	+ virtual void Setup();
	82	+ virtual void Run(py::args args);
	83	+ virtual py::object Call(py::args args);
	84	+
	85	+ static void SetDeterm(bool enable);
	86	+ static void SetTuning(bool enable);
	87	+ virtual ShapeRef* SymIntArraytoShapeRef(py::object shape);
	88	+ ShapeRef* SymIntArraytoShapeRef(at::IntArrayRef shape_array);
	89	+ py::object CreateOutputs(const at::TensorOptions& options, void** addr);
	90	+ void SetupRelocs();
	91	+ void SetWorkspaceSize(size_t size) { ws_size_ = size; }
	92	+
	93	+ struct ShapeWithRef : public ShapeRef {
	94	+ enum { MAX_SIZE = 8 };
	95	+ ShapeWithRef(size_t sz)
	96	+ {
	97	+ data = shape_data;
	98	+ size = sz;
	99	+ }
	100	+ int64_t shape_data[MAX_SIZE];
	101	+ };
	102	+
	103	+protected:
	104	+ enum class LoadType {
	105	+ Load,
	106	+ View,
	107	+ SymBol,
	108	+ };
	109	+
	110	+ struct ParsedCallInputs {
	111	+ std::shared_ptr<std::vector<void*> > addr;
	112	+ at::TensorOptions options;
	113	+ };
	114	+
	115	+ ParsedCallInputs ParseTensorCallInputs(py::args inputs, std::vector<at::Tensor>& tensor_refs) const;
	116	+
	117	+ std::vector<RelocEntry> relocs_;
	118	+
	119	+ std::vector<ShapeWithRef*> shapes_;
	120	+ std::vector<NDObject*> loads_;
	121	+ std::vector<NDObject*> stores_;
	122	+ std::vector<LoadType> load_types_;
	123	+ std::vector<at::Tensor> tensor_list_;
	124	+ std::vector<std::pair<at::Tensor, at::Tensor> > out_refs_;
	125	+ size_t ws_size_;
	126	+ int kernel_type_;
	127	+ uint32_t kernel_flags_;
	128	+ std::string op_name_;
	129	+ std::string op_fullname_;
	130	+};
	131	+
	132	+class GraphSplitBase : public WsAllocator {
	133	+public:
	134	+ virtual ~GraphSplitBase() = default;
	135	+ int LaunchV1(Kernel& kernel, void** addr, aclrtStream stream, std::vector<RelocEntry>& relocs,
	136	+ void* workspace_ptr);
	137	+ int LaunchV2(Kernel& kernel, void** addr, aclrtStream stream, std::vector<RelocEntry>& relocs);
	138	+
	139	+ void* Alloc(size_t size) override;
	140	+
	141	+protected:
	142	+ at::Tensor ws_;
	143	+ aclrtStream stream_;
	144	+};
	145	+class DynKernelPy : public TorchKernelPy {
	146	+public:
	147	+ DynKernelPy(int kernel_type, uint32_t flags) : TorchKernelPy(kernel_type, flags) {}
	148	+ ~DynKernelPy();
	149	+ py::object Load(py::object shape, DataTypePy type) override;
	150	+ py::object GlobalAccess(py::object shape, DataTypePy type) override;
	151	+ py::object ViewLoad(py::object shape, py::object stride, DataTypePy type) override;
	152	+ py::object GatherLoad(py::object shape, py::object index, DataTypePy type, int axis) override;
	153	+ struct LoadShapeRef {
	154	+ ShapeRef shape;
	155	+ ShapeRef stride;
	156	+ };
	157	+ LoadShapeRef* GetDynLoadShapeRef(size_t dim_size);
	158	+ void Setup() override;
	159	+ void Run(py::args args) override;
	160	+ py::object Call(py::args args) override;
	161	+ ShapeRef* SymIntArraytoShapeRef(py::object shape) override;
	162	+ void UpdateSymShapeData();
	163	+
	164	+ py::object MakeScalar(DataTypePy type = DataTypePy(kDataTypeEnd))
	165	+ {
	166	+ auto scalar = std::make_shared<ScalarRefPy>(type);
	167	+ sym_scalar_input_.emplace_back(scalar);
	168	+ load_types_.emplace_back(LoadType::SymBol);
	169	+ return py::cast(scalar);
	170	+ }
	171	+
	172	+ struct DynamicInfo {
	173	+ std::vector<void*> addr;
	174	+ std::vector<ScalarRef> scalars;
	175	+ std::vector<std::vector<int64_t> > shape;
	176	+ std::vector<std::vector<int64_t> > strides;
	177	+ };
	178	+
	179	+protected:
	180	+ ParsedCallInputs ParseDynCallInputs(py::args inputs, std::vector<at::Tensor>& tensor_refs) const;
	181	+ DynKernelPy* AcquireExecutor()
	182	+ {
	183	+ std::lock_guard<std::mutex> lock(dyn_executor_mutex_);
	184	+ if (!dyn_executors_.empty()) {
	185	+ auto* executor = dyn_executors_.back();
	186	+ dyn_executors_.pop_back();
	187	+ return executor;
	188	+ }
	189	+ auto new_executor = CloneExecutor();
	190	+ auto* executor = new_executor.get();
	191	+ dyn_owned_executors_.push_back(std::move(new_executor));
	192	+ return executor;
	193	+ }
	194	+ void ReleaseExecutor(DynKernelPy* executor)
	195	+ {
	196	+ std::lock_guard<std::mutex> lock(dyn_executor_mutex_);
	197	+ dyn_executors_.push_back(executor);
	198	+ }
	199	+ virtual std::unique_ptr<DynKernelPy> CloneExecutor() const;
	200	+ void CloneExecutorStateTo(DynKernelPy& executor) const;
	201	+
	202	+ std::vector<LoadShapeRef*> dyn_load_shapes_;
	203	+ std::vector<ScalarRefPyPtr> sym_scalar_input_;
	204	+ std::vector<ScalarRefPyPtr> const_input_;
	205	+ std::vector<std::vector<ScalarRefPyPtr> > sym_shape_;
	206	+
	207	+private:
	208	+ std::mutex dyn_executor_mutex_;
	209	+ std::vector<DynKernelPy*> dyn_executors_;
	210	+ std::vector<std::unique_ptr<DynKernelPy>> dyn_owned_executors_;
	211	+};
	212	+
	213	+class GraphSplitKernelPy : public TorchKernelPy, public GraphSplitBase {
	214	+public:
	215	+ GraphSplitKernelPy() : TorchKernelPy(KernelPy::K_SPLIT, KernelPy::F_UWS) {}
	216	+ void Setup() override;
	217	+ void Run(py::args inputs) override;
	218	+ py::object Call(py::args inputs) override;
	219	+};
	220	+
	221	+class DynGraphSplitKernelPy : public DynKernelPy, public GraphSplitBase {
	222	+public:
	223	+ DynGraphSplitKernelPy()
	224	+ : DynKernelPy(KernelPy::K_SPLIT, KernelPy::F_UWS \| KernelPy::F_DYN)
	225	+ {
	226	+ }
	227	+ void Setup() override;
	228	+ void Run(py::args inputs) override;
	229	+ py::object Call(py::args inputs) override;
	230	+
	231	+private:
	232	+ std::unique_ptr<DynKernelPy> CloneExecutor() const override;
	233	+};
	234	+} // namespace dvm
	235	+#endif // BUILD_LIBTORCH
	236	+#endif // _DVM_PYBIND_API_H_

Mtorch_npu/csrc/inductor/mlir/cpp_common.cpp+3-3

	@@ -8,8 +8,8 @@
8	#include <torch_npu/csrc/framework/OpCommand.h>	8	#include <torch_npu/csrc/framework/OpCommand.h>
9	#include <torch_npu/csrc/profiler/profiler_mgr.h>	9	#include <torch_npu/csrc/profiler/profiler_mgr.h>
10		10
11	-#include "third_party/acl/inc/experiment/msprof/toolchain/prof_api.h"	11	+#include "third_party/acl/inc/profiling/prof_api.h"
12	-#include "third_party/acl/inc/experiment/msprof/toolchain/prof_common.h"	12	+#include "third_party/acl/inc/profiling/prof_common.h"
13		13
14	struct TilingMem {	14	struct TilingMem {
15	std::unique_ptr<void, decltype(&aclrtFreeHost)> arg_tiling_host;	15	std::unique_ptr<void, decltype(&aclrtFreeHost)> arg_tiling_host;
	@@ -167,4 +167,4 @@ void TORCH_NPU_API opcommand_call(const char* name, std::function<int()> launch_
167	at_npu::native::OpCommand cmd;	167	at_npu::native::OpCommand cmd;
168	cmd.Name(name).SetCustomHandler(launch_call).Run();	168	cmd.Name(name).SetCustomHandler(launch_call).Run();
169	}	169	}
170	-#endif // BUILD_LIBTORCH	170	+#endif // BUILD_LIBTORCH

Mtorch_npu/utils/_dynamo.py+2-2

	@@ -228,12 +228,12 @@ def patch_inductor_wrapper():
228	else:	228	else:
229	src_init(self, mode, options, dynamic)	229	src_init(self, mode, options, dynamic)
230	backend = _resolve_npu_backend_from_wrapper(self)	230	backend = _resolve_npu_backend_from_wrapper(self)
231	- if backend=="mlir":	231	+ if backend == "mlir":
232	with _NpuBackendScope(backend):	232	with _NpuBackendScope(backend):
233	log.info("Running MLIR backend")	233	log.info("Running MLIR backend")
234	device_id = torch_npu.npu.current_device()	234	device_id = torch_npu.npu.current_device()
235	torch_npu._C._recovery_all_npu_stream(device_id)	235	torch_npu._C._recovery_all_npu_stream(device_id)
236	- if backend=="dvm":	236	+ if backend == "dvm":
237	with _NpuBackendScope(backend):	237	with _NpuBackendScope(backend):
238	log.info("Running dvm backend")	238	log.info("Running dvm backend")
239		239