feat(dvm): support A5 and optimizations-pytorch-AtomGit

M.gitmodules+1-1

	@@ -24,4 +24,4 @@
24	[submodule "third_party/dvm/dvm"]	24	[submodule "third_party/dvm/dvm"]
25	path = third_party/dvm/dvm	25	path = third_party/dvm/dvm
26	url = https://gitcode.com/mindspore/dvm.git	26	url = https://gitcode.com/mindspore/dvm.git
27	- branch = r2.9	27	+ branch = r2.10

Mthird_party/dvm/CMakeLists.txt+6-12

	@@ -1,7 +1,11 @@
1	set(DVM_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/dvm)	1	set(DVM_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/dvm)
2	set(DVM_LIB ${DVM_ROOT}/libdvm.a)	2	set(DVM_LIB ${DVM_ROOT}/libdvm.a)
3		3
4	-set(DVM_MAKE_ARGS_STR "PRE_ASCEND=1")	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	+)
5		9
6	file(GLOB DVM_DEPENDS	10	file(GLOB DVM_DEPENDS
7	CONFIGURE_DEPENDS	11	CONFIGURE_DEPENDS
	@@ -12,19 +16,9 @@ file(GLOB DVM_DEPENDS
12	"${DVM_ROOT}/include/*.h"	16	"${DVM_ROOT}/include/*.h"
13	)	17	)
14		18
15	-set(ACL_INC_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../acl/inc)
16	-if (EXISTS "${ACL_INC_ROOT}")
17	- file(GLOB ACL_INC_DIRS LIST_DIRECTORIES true "${ACL_INC_ROOT}/*")
18	- foreach(acl_inc_dir IN LISTS ACL_INC_DIRS)
19	- if (IS_DIRECTORY "${acl_inc_dir}")
20	- file(COPY "${acl_inc_dir}" DESTINATION "${DVM_ROOT}/include")
21	- endif()
22	- endforeach()
23	-endif()
24	-
25	add_custom_command(	19	add_custom_command(
26	OUTPUT ${DVM_LIB}	20	OUTPUT ${DVM_LIB}
27	- COMMAND make -C ${DVM_ROOT} ${DVM_MAKE_ARGS_STR} libdvm.a	21	+ COMMAND make -C ${DVM_ROOT} ${DVM_MAKE_ARGS} libdvm.a
28	WORKING_DIRECTORY ${DVM_ROOT}	22	WORKING_DIRECTORY ${DVM_ROOT}
29	DEPENDS ${DVM_DEPENDS}	23	DEPENDS ${DVM_DEPENDS}
30	COMMENT "Building DVM static library"	24	COMMENT "Building DVM static library"

Mthird_party/dvm/dvm+1-1

	@@ -1 +1 @@
1	-Subproject commit fe91de40ccbb3d9fbb2077423ee0f109760fb2f6	1	+Subproject commit d59888fdc684320445b92addaa4fd37eb1943902

Mtorch_npu/_inductor/dvm/__init__.py+10-3

	@@ -11,6 +11,7 @@ from torch_npu._C.dvm import (
11	NDObject,	11	NDObject,
12	TorchKernel as Kernel,	12	TorchKernel as Kernel,
13	)	13	)
		14	+from torch_npu.npu._backends import get_soc_version
14		15
15	from .fx_test import _accuracy_check_run	16	from .fx_test import _accuracy_check_run
16		17
	@@ -24,6 +25,12 @@ float32 = DataType.float32
24	int32 = DataType.int32	25	int32 = DataType.int32
25	int64 = DataType.int64	26	int64 = DataType.int64
26		27
		28	+Ascend910B1 = 220
		29	+Ascend310B1 = 240
		30	+Ascend910_9391 = 250
		31	+Ascend950 = 260
		32	+is_ascend950 = get_soc_version() >= Ascend950
		33	+
27	KERNEL_FACTORY = {	34	KERNEL_FACTORY = {
28	("mix", True): partial(DynKernel, Kernel.K_MIX, Kernel.F_DYN),	35	("mix", True): partial(DynKernel, Kernel.K_MIX, Kernel.F_DYN),
29	("mix", False): partial(Kernel, Kernel.K_MIX, 0),	36	("mix", False): partial(Kernel, Kernel.K_MIX, 0),
	@@ -39,9 +46,8 @@ KERNEL_FACTORY = {
39	def kernel(	46	def kernel(
40	ktype: str = "split",	47	ktype: str = "split",
41	dyn_shape: bool = False,	48	dyn_shape: bool = False,
42	- mfusion: bool = False,
43	):	49	):
44	- r"""kernel(ktype="split", dyn_shape=False, mfusion=False)	50	+ r"""kernel(ktype="split", dyn_shape=False)
45		51
46	Return a decorator that builds and executes a DVM kernel.	52	Return a decorator that builds and executes a DVM kernel.
47		53
	@@ -188,4 +194,5 @@ def _install_bf16_promote():
188	setattr(Kernel, name, _promote_bf16(op_fn))	194	setattr(Kernel, name, _promote_bf16(op_fn))
189		195
190		196
191	-_install_bf16_promote()	197	+if not is_ascend950:
		198	+ _install_bf16_promote()

Mtorch_npu/_inductor/dvm/graph_build.py+1-3

	@@ -33,7 +33,6 @@ class DvmCodegenInterpreter(torch.fx.Interpreter):
33	):	33	):
34	super().__init__(gm)	34	super().__init__(gm)
35	self.gm = gm	35	self.gm = gm
36	- self.is_mfusion = bool(getattr(gm, "_mfusion", False))
37	self.ktype = ktype	36	self.ktype = ktype
38	self.is_mix_kernel = annotate_mm_transpose_flags(gm)	37	self.is_mix_kernel = annotate_mm_transpose_flags(gm)
39	if is_dynamic is None:	38	if is_dynamic is None:
	@@ -53,8 +52,7 @@ class DvmCodegenInterpreter(torch.fx.Interpreter):
53	self.ktype = "spec"	52	self.ktype = "spec"
54	self.code.splice(f'\n"""\n{self.gm.print_readable(print_output=False)}\n"""')	53	self.code.splice(f'\n"""\n{self.gm.print_readable(print_output=False)}\n"""')
55	decorator = (	54	decorator = (
56	- f"{chr(64)}dvm.kernel(ktype={self.ktype!r}, dyn_shape={self.is_dynamic}, "	55	+ f"{chr(64)}dvm.kernel(ktype={self.ktype!r}, dyn_shape={self.is_dynamic})"
57	- f"mfusion={self.is_mfusion!r})"
58	)	56	)
59	self.code.splice(decorator)	57	self.code.splice(decorator)
60	self.code.splice(f"def {self.KERNEL_NAME_PLACEHOLDER}(k):")	58	self.code.splice(f"def {self.KERNEL_NAME_PLACEHOLDER}(k):")

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

	@@ -76,7 +76,8 @@ anir_config.GENERATE_LIST = [
76	aten.scalar_tensor,	76	aten.scalar_tensor,
77	aten.unsqueeze,	77	aten.unsqueeze,
78	aten.squeeze,	78	aten.squeeze,
79	- aten.clone,	79	+ # aten.reshape,
		80	+ # aten.clone,
80	]	81	]
81		82
82		83
	@@ -262,7 +263,6 @@ def _dvm_can_fuse_horizontal(self, node1, node2):
262	def _patch_lowering_type_checks():	263	def _patch_lowering_type_checks():
263	import torch._inductor.graph as inductor_graph	264	import torch._inductor.graph as inductor_graph
264	import torch._inductor.lowering as inductor_lowering	265	import torch._inductor.lowering as inductor_lowering
265	- import torch._inductor.pattern_matcher as pattern_matcher
266		266
267	import torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering as npu_lowering_mod	267	import torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu.inductor_patch.lowering as npu_lowering_mod
268		268
	@@ -281,7 +281,7 @@ def _patch_lowering_type_checks():
281	if not isinstance(meta, torch._subclasses.FakeTensor):	281	if not isinstance(meta, torch._subclasses.FakeTensor):
282	continue	282	continue
283		283
284	- if meta.is_cpu:	284	+ if meta.is_cpu and config.disable_cpp_codegen:
285	return True	285	return True
286		286
287	if node.target in DVM_OP_REGISTRY:	287	if node.target in DVM_OP_REGISTRY:
	@@ -292,9 +292,6 @@ def _patch_lowering_type_checks():
292	inductor_lowering.fallback_node_due_to_unsupported_type = (	292	inductor_lowering.fallback_node_due_to_unsupported_type = (
293	_fallback_node_due_to_unsupported_type	293	_fallback_node_due_to_unsupported_type
294	)	294	)
295	- pattern_matcher.fallback_node_due_to_unsupported_type = (
296	- _fallback_node_due_to_unsupported_type
297	- )
298	npu_lowering_mod.fallback_node_due_to_unsupported_type = (	295	npu_lowering_mod.fallback_node_due_to_unsupported_type = (
299	_fallback_node_due_to_unsupported_type	296	_fallback_node_due_to_unsupported_type
300	)	297	)

Mtorch_npu/_inductor/dvm/op_emitter.py+12-7

	@@ -2,6 +2,8 @@ import torch
2	import torch.fx	2	import torch.fx
3	import torch.utils._pytree as pytree	3	import torch.utils._pytree as pytree
4		4
		5	+from . import is_ascend950
		6	+
5		7
6	aten = torch.ops.aten	8	aten = torch.ops.aten
7	prims = torch.ops.prims	9	prims = torch.ops.prims
	@@ -43,7 +45,7 @@ def to_dvm_dtype(dtype):
43	return dtype	45	return dtype
44		46
45		47
46	-def _check_dtype(inputs, supported_dtypes):	48	+def _check_dtype(inputs, supported_dtypes, allow_cpu=False):
47	for inp in inputs:	49	for inp in inputs:
48	if not isinstance(inp, torch.fx.Node):	50	if not isinstance(inp, torch.fx.Node):
49	continue	51	continue
	@@ -55,7 +57,7 @@ def _check_dtype(inputs, supported_dtypes):
55	continue	57	continue
56	if meta.dtype not in supported_dtypes:	58	if meta.dtype not in supported_dtypes:
57	return False	59	return False
58	- if meta.is_cpu:	60	+ if meta.is_cpu and not allow_cpu:
59	return False	61	return False
60	return True	62	return True
61		63
	@@ -67,7 +69,7 @@ def common_rule(node: torch.fx.Node):
67		69
68		70
69	def full_rule(node: torch.fx.Node):	71	def full_rule(node: torch.fx.Node):
70	- return _check_dtype([node], supported_dtypes=DVM_SUPPORT_TYPE)	72	+ return _check_dtype([node], supported_dtypes=DVM_SUPPORT_TYPE, allow_cpu=True)
71		73
72		74
73	def cast_rule(node: torch.fx.Node):	75	def cast_rule(node: torch.fx.Node):
	@@ -114,11 +116,14 @@ def mm_rule(node: torch.fx.Node):
114		116
115	def check(input_node):	117	def check(input_node):
116	t = input_node.meta["val"]	118	t = input_node.meta["val"]
		119	+ if t.dim() > 4 or t.dim() < 2:
		120	+ return False
117	inner_axis = inner_axis_length(t)	121	inner_axis = inner_axis_length(t)
118	- if isinstance(inner_axis, torch.SymInt):	122	+ if not is_ascend950:
119	- return False	123	+ if isinstance(inner_axis, torch.SymInt):
120	- if t.dim() > 4 or t.dim() < 2 or inner_axis > MAX_INNER:	124	+ return False
121	- return False	125	+ if inner_axis > MAX_INNER:
		126	+ return False
122	return True	127	return True
123		128
124	def check_output(output_node):	129	def check_output(output_node):

Mtorch_npu/_inductor/mfusion/README.md+2-2

	@@ -105,10 +105,10 @@ with MFusionPatch():
105		105
106	启用 MFusion 后，`torch.compile` 会生成融合后的 kernel 代码。如需查看融合算子的详细信息，可通过设置环境变量 export TORCH_COMPILE_DEBUG=1 开启调试模式，相关信息会输出到torch_compile_debug目录下的output_code.py 文件中。	106	启用 MFusion 后，`torch.compile` 会生成融合后的 kernel 代码。如需查看融合算子的详细信息，可通过设置环境变量 export TORCH_COMPILE_DEBUG=1 开启调试模式，相关信息会输出到torch_compile_debug目录下的output_code.py 文件中。
107		107
108	-融合 kernel 通过 `@dvm.kernel(..., mfusion=True)` 装饰器标识：	108	+kernel 会生成 `mfusion_dvm_*` 命名的 DVM kernel：
109		109
110	```Python	110	```Python
111	-@dvm.kernel(ktype='vector', dyn_shape=False, mfusion=True)	111	+@dvm.kernel(ktype='vector', dyn_shape=False)
112	def mfusion_dvm_0(k):	112	def mfusion_dvm_0(k):
113	arg0 = k.load([32], dvm.float32)	113	arg0 = k.load([32], dvm.float32)
114	arg1 = k.load([32], dvm.float32)	114	arg1 = k.load([32], dvm.float32)

Mtorch_npu/csrc/inductor/dvm/pybind_api.cpp+10-0

	@@ -206,6 +206,16 @@ py::object TorchKernelPy::Store(py::object obj, DataTypePy type)
206	return ObjToPy(stores_.emplace_back(op));	206	return ObjToPy(stores_.emplace_back(op));
207	}	207	}
208		208
		209	+py::object TorchKernelPy::ViewStore(py::object obj, py::object stride, DataTypePy type)
		210	+{
		211	+ auto in_obj = PyToObj(obj);
		212	+ if (type != kDataTypeEnd) {
		213	+ in_obj = kernel_.Cast(in_obj, type);
		214	+ }
		215	+ auto op = kernel_.Store(nullptr, in_obj, GetShapeRef(stride));
		216	+ return ObjToPy(stores_.emplace_back(op));
		217	+}
		218	+
209	void TorchKernelPy::Setup()	219	void TorchKernelPy::Setup()
210	{	220	{
211	SetupRelocs();	221	SetupRelocs();

Mtorch_npu/csrc/inductor/dvm/pybind_api.h+1-0

	@@ -64,6 +64,7 @@ public:
64	py::object Load(py::object shape, DataTypePy type) override;	64	py::object Load(py::object shape, DataTypePy type) override;
65	py::object ViewLoad(py::object shape, py::object stride, DataTypePy type) override;	65	py::object ViewLoad(py::object shape, py::object stride, DataTypePy type) override;
66	py::object Store(py::object obj, DataTypePy type) override;	66	py::object Store(py::object obj, DataTypePy type) override;
		67	+ py::object ViewStore(py::object obj, py::object stride, DataTypePy type) override;
67	IntArrayRef* GetShapeRef(py::object shape) override;	68	IntArrayRef* GetShapeRef(py::object shape) override;
68		69
69	void SetKernelInfo(const std::string& op_name, const std::string& op_fullname,	70	void SetKernelInfo(const std::string& op_name, const std::string& op_fullname,