feat: Unified multi-branch to master branch, to build package for several pytorch version from master branch.-pytorch-AtomGit

Mci/build.sh+56-0

	@@ -4,8 +4,12 @@ set -e
4		4
5	CUR_DIR=$(dirname $(readlink -f $0))	5	CUR_DIR=$(dirname $(readlink -f $0))
6	SUPPORTED_PY_VERSION=(3.9 3.10 3.11 3.12 3.13)	6	SUPPORTED_PY_VERSION=(3.9 3.10 3.11 3.12 3.13)
		7	+SUPPORTED_TORCH_VERSION=(2.10.0 2.11.0 2.12.0)
7	# Default supported python version is 3.9	8	# Default supported python version is 3.9
8	PY_VERSION="3.9"	9	PY_VERSION="3.9"
		10	+# Torch version to validate against installed PyTorch (empty = skip check)
		11	+# Also written to version.txt before building
		12	+TORCH_VERSION=""
9		13
10	# Parse arguments inside script	14	# Parse arguments inside script
11	function parse_script_args() {	15	function parse_script_args() {
	@@ -24,6 +28,11 @@ function parse_script_args() {
24	args_num=$((args_num-1))	28	args_num=$((args_num-1))
25	shift	29	shift
26	;;	30	;;
		31	+ --torch=*)
		32	+ TORCH_VERSION=$(echo "${1}"\|cut -d"=" -f2)
		33	+ args_num=$((args_num-1))
		34	+ shift
		35	+ ;;
27	--disable_torchair)	36	--disable_torchair)
28	export DISABLE_INSTALL_TORCHAIR=TRUE	37	export DISABLE_INSTALL_TORCHAIR=TRUE
29	args_num=$((args_num-1))	38	args_num=$((args_num-1))
	@@ -91,6 +100,45 @@ function check_python_version() {
91	fi	100	fi
92	}	101	}
93		102
		103	+function check_torch_version() {
		104	+ if [ -z "${TORCH_VERSION}" ]; then
		105	+ return 0
		106	+ fi
		107	+ local matched="false"
		108	+ for ver in ${SUPPORTED_TORCH_VERSION[*]}; do
		109	+ if [ "${TORCH_VERSION}" = "${ver}" ]; then
		110	+ matched="true"
		111	+ break
		112	+ fi
		113	+ done
		114	+ if [ "${matched}" = "false" ]; then
		115	+ echo "${TORCH_VERSION} is an unsupported torch version, we suggest ${SUPPORTED_TORCH_VERSION[*]}"
		116	+ exit 1
		117	+ fi
		118	+}
		119	+
		120	+function check_torch_installed() {
		121	+ local installed
		122	+ installed=$(python"${PY_VERSION}" -c "import torch; print(torch.__version__)" 2>/dev/null)
		123	+ if [ -z "${installed}" ]; then
		124	+ echo "PyTorch is not installed for python${PY_VERSION}. Please install it before building."
		125	+ exit 1
		126	+ fi
		127	+ # Strip local tag (e.g. 2.11.0+cpu -> 2.11.0)
		128	+ local installed_base="${installed%%+*}"
		129	+ if [ -n "${TORCH_VERSION}" ]; then
		130	+ # Compare major.minor only (ignore patch and local tag)
		131	+ local requested_mm installed_mm
		132	+ requested_mm=$(echo "${TORCH_VERSION}" \| cut -d. -f1,2)
		133	+ installed_mm=$(echo "${installed_base}" \| cut -d. -f1,2)
		134	+ if [ "${installed_mm}" != "${requested_mm}" ]; then
		135	+ echo "PyTorch version mismatch: requested ${TORCH_VERSION}, but ${installed} is installed."
		136	+ exit 1
		137	+ fi
		138	+ fi
		139	+ echo "Using PyTorch ${installed}"
		140	+}
		141	+
94	function main()	142	function main()
95	{	143	{
96	if ! parse_script_args "$@"; then	144	if ! parse_script_args "$@"; then
	@@ -98,8 +146,16 @@ function main()
98	exit 1	146	exit 1
99	fi	147	fi
100	check_python_version	148	check_python_version
		149	+ check_torch_version
		150	+ check_torch_installed
101		151
102	cd ${CUR_DIR}/..	152	cd ${CUR_DIR}/..
		153	+
		154	+ if [ -n "${TORCH_VERSION}" ]; then
		155	+ export TORCH_VERSION
		156	+ echo "${TORCH_VERSION}" > version.txt
		157	+ echo "Set package version to ${TORCH_VERSION}"
		158	+ fi
103	# if you add or delete file/files in the project, you need to remove the following comment	159	# if you add or delete file/files in the project, you need to remove the following comment
104	# make clean	160	# make clean
105		161

Msetup.py+14-0

	@@ -130,6 +130,17 @@ def generate_torch_npu_version():
130	generate_torch_npu_version()	130	generate_torch_npu_version()
131		131
132		132
		133	+def _get_torch_requires():
		134	+ torch_version = os.environ.get("TORCH_VERSION", "")
		135	+ if not torch_version:
		136	+ try:
		137	+ import torch
		138	+ torch_version = torch.__version__.split("+")[0]
		139	+ except ImportError:
		140	+ pass
		141	+ return ["torch==" + torch_version] if torch_version else []
		142	+
		143	+
133	def which(thefile):	144	def which(thefile):
134	path = os.environ.get("PATH", os.defpath).split(os.pathsep)	145	path = os.environ.get("PATH", os.defpath).split(os.pathsep)
135	for d in path:	146	for d in path:
	@@ -731,6 +742,9 @@ setup(
731	define_macros=[('_GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0'), ('GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0')]	742	define_macros=[('_GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0'), ('GLIBCXX_USE_CXX11_ABI', '1' if USE_CXX11_ABI else '0')]
732	),	743	),
733	],	744	],
		745	+ install_requires=[
		746	+ *_get_torch_requires(),
		747	+ ],
734	extras_require={	748	extras_require={
735	},	749	},
736	package_data={	750	package_data={

Mtest/_inductor/test_npu_device.py+1-1

	@@ -36,7 +36,7 @@ class TestNpuDevice(TestCase):
36	self.assertIn("#include <sys/syscall.h>", result)	36	self.assertIn("#include <sys/syscall.h>", result)
37	self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)	37	self.assertIn("#include <torch_npu/csrc/framework/OpCommand.h>", result)
38	self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)	38	self.assertIn("#include <torch_npu/csrc/core/npu/NPUStream.h>", result)
39	- self.assertIn("#include \"experiment/runtime/runtime/rt.h\"", result)	39	+ self.assertIn("#include \"runtime/runtime/rt.h\"", result)
40		40
41	def test_cpp_aoti_stream_guard(self):	41	def test_cpp_aoti_stream_guard(self):
42	overrides = NewNPUDeviceOpOverrides()	42	overrides = NewNPUDeviceOpOverrides()

Atorch_npu/_compat/__init__.py+0-0

The file is empty

Atorch_npu/_compat/distributed.py+9-0

@@ -0,0 +1,9 @@
	1	+from torch_npu._compat.version import CURRENT_VERSION
	2	+
	3	+# COMPAT(>= 2.11): register_op_strategy / register_prop_rule moved from
	4	+# _ops.registration to _ops.utils in PyTorch 2.11.
	5	+# CAN REMOVE else branch when MIN_SUPPORTED >= (2, 11)
	6	+if CURRENT_VERSION >= (2, 11):
	7	+ from torch.distributed.tensor._ops.utils import register_op_strategy, register_prop_rule
	8	+else:
	9	+ from torch.distributed.tensor._ops.registration import register_op_strategy, register_prop_rule

Atorch_npu/_compat/inductor.py+31-0

@@ -0,0 +1,31 @@
	1	+from torch_npu._compat.version import CURRENT_VERSION
	2	+
	3	+
	4	+# COMPAT(>= 2.11): CachingAutotuner moved from runtime.triton_heuristics to triton_heuristics.
	5	+# Lazy function (not module-level import): torch._inductor.triton_heuristics does not exist
	6	+# in CPU-only builds, so importing at module load time would fail.
	7	+# CAN REMOVE when MIN_SUPPORTED >= (2, 11): import from triton_heuristics directly
	8	+def get_CachingAutotuner():
	9	+ if CURRENT_VERSION >= (2, 11):
	10	+ from torch._inductor.triton_heuristics import CachingAutotuner
	11	+ else:
	12	+ from torch._inductor.runtime.triton_heuristics import CachingAutotuner
	13	+ return CachingAutotuner
	14	+
	15	+
	16	+# COMPAT(>= 2.11): sizevars.var_to_val renamed to sizevars.backed_var_to_val.
	17	+# CAN REMOVE when MIN_SUPPORTED >= (2, 11): use sizevars.backed_var_to_val directly
	18	+def get_sizevars_backed_var_to_val(sizevars):
	19	+ if hasattr(sizevars, "backed_var_to_val"):
	20	+ return sizevars.backed_var_to_val
	21	+ return sizevars.var_to_val
	22	+
	23	+
	24	+# COMPAT(>= 2.11): gen_common_triton_imports changed from module-level function
	25	+# to instance method on the kernel object.
	26	+# CAN REMOVE when MIN_SUPPORTED >= (2, 11): call kernel.gen_common_triton_imports() directly
	27	+def gen_common_triton_imports(kernel):
	28	+ if CURRENT_VERSION >= (2, 11):
	29	+ return kernel.gen_common_triton_imports()
	30	+ from torch._inductor.codegen.triton import gen_common_triton_imports as _fn
	31	+ return _fn()

Atorch_npu/_compat/version.py+12-0

@@ -0,0 +1,12 @@
	1	+import torch
	2	+
	3	+
	4	+def _parse(version_str: str) -> tuple:
	5	+ parts = version_str.split("+")[0].split(".")
	6	+ return (int(parts[0]), int(parts[1]))
	7	+
	8	+
	9	+CURRENT_VERSION: tuple = _parse(torch.__version__)
	10	+
	11	+# Bump this when dropping old version support; run tools/check_compat.py to find stale COMPAT blocks.
	12	+MIN_SUPPORTED_VERSION: tuple = (2, 10)

Mtorch_npu/_inductor/__init__.py+5-1

	@@ -72,6 +72,9 @@ else:
72	patch_constant_fold_uniform_value()	72	patch_constant_fold_uniform_value()
73	patch_fallback_kernel_codegen()	73	patch_fallback_kernel_codegen()
74		74
		75	+ from .ir import patch_extern_kernel_codegen_size_asserts
		76	+ patch_extern_kernel_codegen_size_asserts()
		77	+
75	patch_aot_code_compiler_compile()	78	patch_aot_code_compiler_compile()
76		79
77		80
	@@ -92,7 +95,8 @@ else:
92		95
93	# register fx_pass should be put behind of _register_npu_inductor_decompositons	96	# register fx_pass should be put behind of _register_npu_inductor_decompositons
94	def _replace_benchmark_all_configs():	97	def _replace_benchmark_all_configs():
95	- from torch._inductor.triton_heuristics import CachingAutotuner	98	+ from torch_npu._compat.inductor import get_CachingAutotuner
		99	+ CachingAutotuner = get_CachingAutotuner()
96	from .npu_triton_heuristics import benchmark_all_configs	100	from .npu_triton_heuristics import benchmark_all_configs
97	CachingAutotuner.benchmark_all_configs = benchmark_all_configs	101	CachingAutotuner.benchmark_all_configs = benchmark_all_configs
98		102

Mtorch_npu/_inductor/codegen/ir.py+5-4

	@@ -166,12 +166,13 @@ def substituted_dims_in_indexing(self, indexing, kernel, range_tree_nodes_substi
166	return substituted	166	return substituted
167		167
168		168
169	-def generate_body_indexing(body, indices):	169	+def generate_body_indexing(body, indices, allow_same_symbol_in_index=False):
170	index = list(itertools.chain.from_iterable(indices))	170	index = list(itertools.chain.from_iterable(indices))
171	if not (len(index) == len(body.var_ranges)):	171	if not (len(index) == len(body.var_ranges)):
172	raise RuntimeError("assert len(index) == len(body.var_ranges), (index, body.var_ranges)")	172	raise RuntimeError("assert len(index) == len(body.var_ranges), (index, body.var_ranges)")
173	- if not (all(v not in body.var_ranges for v in index)):	173	+ if not allow_same_symbol_in_index:
174	- raise RuntimeError("assert all(v not in body.var_ranges for v in index)")	174	+ if not (all(v not in body.var_ranges for v in index)):
		175	+ raise RuntimeError("assert all(v not in body.var_ranges for v in index)")
175		176
176	replacements = dict(zip(body.var_ranges.keys(), index))	177	replacements = dict(zip(body.var_ranges.keys(), index))
177	indexing_map = dict(zip(index, body.var_ranges.keys()))	178	indexing_map = dict(zip(index, body.var_ranges.keys()))
	@@ -193,7 +194,7 @@ def transform_dims_in_indexing(self, indices):
193	# select tiling axis, recover missing dimensions,	194	# select tiling axis, recover missing dimensions,
194	def loopbody__call__(self, *indices, allow_same_symbol_in_index=False):	195	def loopbody__call__(self, *indices, allow_same_symbol_in_index=False):
195	if self.indexing is None:	196	if self.indexing is None:
196	- generate_body_indexing(self, indices)	197	+ generate_body_indexing(self, indices, allow_same_symbol_in_index)
197	result = self.root_block()	198	result = self.root_block()
198	self.indexing = None	199	self.indexing = None
199	return result	200	return result

Mtorch_npu/_inductor/codegen/kernel_analysis.py+2-1

	@@ -4,11 +4,12 @@ from torch._inductor import ir
4	from torch._inductor.scheduler import SchedulerNode	4	from torch._inductor.scheduler import SchedulerNode
5	from torch._inductor.utils import sympy_index_symbol	5	from torch._inductor.utils import sympy_index_symbol
6	from torch._inductor.virtualized import V	6	from torch._inductor.virtualized import V
		7	+from torch_npu._compat.inductor import get_sizevars_backed_var_to_val
7		8
8		9
9	class IndexAnalysis:	10	class IndexAnalysis:
10	def __init__(self, kernel, raw_index, is_store_index=False, is_index_expr=False):	11	def __init__(self, kernel, raw_index, is_store_index=False, is_index_expr=False):
11	- self.index = raw_index.subs(V.graph.sizevars.backed_var_to_val)	12	+ self.index = raw_index.subs(get_sizevars_backed_var_to_val(V.graph.sizevars))
12	self.kernel = kernel	13	self.kernel = kernel
13	self.tiling_axis = [x.symbol() for x in self.kernel.tiling_axis]	14	self.tiling_axis = [x.symbol() for x in self.kernel.tiling_axis]
14	self.stride_list = None # stride list [1,2,4,24]	15	self.stride_list = None # stride list [1,2,4,24]

Mtorch_npu/_inductor/codegen/split_tiling.py+3-2

	@@ -6,6 +6,7 @@ from torch._inductor.loop_body import MemoryUsageType
6	from torch._inductor.runtime.runtime_utils import next_power_of_2	6	from torch._inductor.runtime.runtime_utils import next_power_of_2
7	from torch._inductor.utils import ModularIndexing, sympy_subs	7	from torch._inductor.utils import ModularIndexing, sympy_subs
8	from torch._inductor.virtualized import V	8	from torch._inductor.virtualized import V
		9	+from torch_npu._compat.inductor import get_sizevars_backed_var_to_val
9		10
10	from .kernel_analysis import IndexAnalysis	11	from .kernel_analysis import IndexAnalysis
11	from .triton_utils import get_aligned_numel	12	from .triton_utils import get_aligned_numel
	@@ -254,10 +255,10 @@ class SplitTiling:
254	xnumel = x	255	xnumel = x
255	ynumel = y	256	ynumel = y
256	if isinstance(xnumel, (sympy.Symbol, sympy.Expr)) and not isinstance(xnumel, sympy.Integer):	257	if isinstance(xnumel, (sympy.Symbol, sympy.Expr)) and not isinstance(xnumel, sympy.Integer):
257	- xnumel = xnumel.subs(V.graph.sizevars.backed_var_to_val)	258	+ xnumel = xnumel.subs(get_sizevars_backed_var_to_val(V.graph.sizevars))
258		259
259	if isinstance(ynumel, (sympy.Symbol, sympy.Expr)) and not isinstance(ynumel, sympy.Integer):	260	if isinstance(ynumel, (sympy.Symbol, sympy.Expr)) and not isinstance(ynumel, sympy.Integer):
260	- ynumel = ynumel.subs(V.graph.sizevars.backed_var_to_val)	261	+ ynumel = ynumel.subs(get_sizevars_backed_var_to_val(V.graph.sizevars))
261		262
262	if isinstance(xnumel, sympy.Integer) and isinstance(ynumel, int):	263	if isinstance(xnumel, sympy.Integer) and isinstance(ynumel, int):
263	ynumel = sympy.Integer(ynumel)	264	ynumel = sympy.Integer(ynumel)

Mtorch_npu/_inductor/codegen/triton.py+5-3

	@@ -82,6 +82,8 @@ from .kernel_analysis import IndexAnalysis, ReductionAnalysis
82	from .npu_kernel_features import NumelList	82	from .npu_kernel_features import NumelList
83	from ..runtime import NPUDeviceProperties	83	from ..runtime import NPUDeviceProperties
84	from .. import npu_triton_heuristics	84	from .. import npu_triton_heuristics
		85	+from torch_npu._compat.inductor import get_sizevars_backed_var_to_val
		86	+from torch_npu._compat.inductor import gen_common_triton_imports as _compat_gen_common_triton_imports
85		87
86		88
87	def flatten(nums):	89	def flatten(nums):
	@@ -516,7 +518,7 @@ class NPUIndexTritonKernel(TritonKernel):
516	def initialize_range_tree(self, pid_cache):	518	def initialize_range_tree(self, pid_cache):
517	for k, x in self.numels.items():	519	for k, x in self.numels.items():
518	if not isinstance(x, sympy.Integer):	520	if not isinstance(x, sympy.Integer):
519	- x = x.subs(V.graph.sizevars.backed_var_to_val)	521	+ x = x.subs(get_sizevars_backed_var_to_val(V.graph.sizevars))
520	self.numels[k] = x	522	self.numels[k] = x
521		523
522	no_r_dim = not self.inside_reduction or self.numels["r"] == 1	524	no_r_dim = not self.inside_reduction or self.numels["r"] == 1
	@@ -730,7 +732,7 @@ class NPUIndexTritonKernel(TritonKernel):
730	size_hints = self.get_size_hints()	732	size_hints = self.get_size_hints()
731	heuristics = self._get_heuristic()	733	heuristics = self._get_heuristic()
732	if name is None:	734	if name is None:
733	- code.splice(self.gen_common_triton_imports())	735	+ code.splice(_compat_gen_common_triton_imports(self))
734	# Note: add extra imports for extensions	736	# Note: add extra imports for extensions
735	code.splice(self.gen_triton_ext_imports())	737	code.splice(self.gen_triton_ext_imports())
736		738
	@@ -1120,7 +1122,7 @@ class NPUIndexTritonKernel(TritonKernel):
1120		1122
1121	# all are load indexings, select the longest as gold	1123	# all are load indexings, select the longest as gold
1122	for index in self.load_store_indexing:	1124	for index in self.load_store_indexing:
1123	- index = index.subs(V.graph.sizevars.backed_var_to_val)	1125	+ index = index.subs(get_sizevars_backed_var_to_val(V.graph.sizevars))
1124	analyze = IndexAnalysis(self, index)	1126	analyze = IndexAnalysis(self, index)
1125	if len(analyze.var_list) > maximum_length and all_tiling_in_var_list(analyze.var_list):	1127	if len(analyze.var_list) > maximum_length and all_tiling_in_var_list(analyze.var_list):
1126	longest = analyze.var_list	1128	longest = analyze.var_list

Mtorch_npu/_inductor/config.py+8-0

	@@ -86,6 +86,14 @@ dump_fx_graph = os.environ.get("INDUCTOR_ASCEND_DUMP_FX_GRAPH", False) \
86	# (2) [1, 2, 10] means try to fallback kernel like triton_xxx_1, triton_xxx_2 and triton_xxx_10	86	# (2) [1, 2, 10] means try to fallback kernel like triton_xxx_1, triton_xxx_2 and triton_xxx_10
87	force_fallback_kernel_id = []	87	force_fallback_kernel_id = []
88		88
		89	+# Control whether to skip stride assertions for ops that may change stride
		90	+# at runtime (like _to_copy on NPU forcing Contiguous memory format).
		91	+#
		92	+# Usage:
		93	+# - Skip specific ops: skip_specific_stride_asserts = [torch.ops.aten._to_copy.default, ...]
		94	+# - Disable skip: skip_specific_stride_asserts = [] (default)
		95	+skip_specific_stride_asserts = []
		96	+
89	acc_comp_tol = {	97	acc_comp_tol = {
90	torch.float32: {'rtol': 1.3e-6, 'atol': 1e-5},	98	torch.float32: {'rtol': 1.3e-6, 'atol': 1e-5},
91	torch.float16: {'rtol': 1e-3, 'atol': 1e-5},	99	torch.float16: {'rtol': 1e-3, 'atol': 1e-5},

Mtorch_npu/_inductor/ir.py+27-1

	@@ -85,4 +85,30 @@ def patch_fallback_kernel_codegen():
85	self.codegen_unbacked_symbol_defs(wrapper)	85	self.codegen_unbacked_symbol_defs(wrapper)
86		86
87	from torch._inductor.ir import FallbackKernel	87	from torch._inductor.ir import FallbackKernel
88	- FallbackKernel.codegen = codegen_npu	88	+ FallbackKernel.codegen = codegen_npu
		89	+
		90	+
		91	+def patch_extern_kernel_codegen_size_asserts():
		92	+ from torch._inductor.ir import ExternKernel
		93	+ from . import config as npu_config
		94	+ original_codegen_size_asserts = ExternKernel.codegen_size_asserts
		95	+
		96	+ def npu_codegen_size_asserts(self, wrapper):
		97	+ fx_node = getattr(self, 'fx_node', None)
		98	+ should_skip = False
		99	+ if fx_node and fx_node.target:
		100	+ skip_config = getattr(npu_config, 'skip_specific_stride_asserts', [])
		101	+ if isinstance(skip_config, (list, tuple)):
		102	+ should_skip = fx_node.target in skip_config
		103	+ if should_skip:
		104	+ if config.size_asserts and not V.graph.cpp_wrapper:
		105	+ from torch._inductor.utils import sympy_product
		106	+ if sympy_product(self.get_size()) == 0:
		107	+ return
		108	+ wrapper.writeline(
		109	+ f"# NPU: Skipping stride assertion for {fx_node.target}"
		110	+ )
		111	+ else:
		112	+ original_codegen_size_asserts(self, wrapper)
		113	+
		114	+ ExternKernel.codegen_size_asserts = npu_codegen_size_asserts

Mtorch_npu/distributed/tensor/_math_ops.py+1-1

	@@ -16,9 +16,9 @@ from torch.distributed.tensor._op_schema import (
16	from torch.distributed.tensor._ops.utils import (	16	from torch.distributed.tensor._ops.utils import (
17	generate_redistribute_costs,	17	generate_redistribute_costs,
18	normalize_dim,	18	normalize_dim,
19	- register_op_strategy,
20	expand_to_full_mesh_op_strategy,	19	expand_to_full_mesh_op_strategy,
21	)	20	)
		21	+from torch_npu._compat.distributed import register_op_strategy
22	from torch.distributed.tensor._ops._math_ops import (	22	from torch.distributed.tensor._ops._math_ops import (
23	_replicate_dims_start_at,	23	_replicate_dims_start_at,
24	_infer_reduce_dims_map,	24	_infer_reduce_dims_map,

Mtorch_npu/distributed/tensor/_matrix_ops.py+2-1

	@@ -4,7 +4,8 @@ import os
4	import torch	4	import torch
5	from torch.distributed._tensor.experimental import register_sharding	5	from torch.distributed._tensor.experimental import register_sharding
6	from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta	6	from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
7	-from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy, register_op_strategy	7	+from torch.distributed.tensor._ops.utils import expand_to_full_mesh_op_strategy
		8	+from torch_npu._compat.distributed import register_op_strategy
8	from torch.distributed.tensor import DTensor, Partial, Replicate, Shard	9	from torch.distributed.tensor import DTensor, Partial, Replicate, Shard
9	from torch.distributed.tensor._op_schema import (	10	from torch.distributed.tensor._op_schema import (
10	OpInfo,	11	OpInfo,

Mtorch_npu/distributed/tensor/_pointwise_ops.py+1-1

	@@ -1,7 +1,7 @@
1		1
2	import torch	2	import torch
3	from torch.distributed.tensor._op_schema import OpSchema, RuntimeSchemaInfo	3	from torch.distributed.tensor._op_schema import OpSchema, RuntimeSchemaInfo
4	-from torch.distributed.tensor._ops.utils import register_op_strategy	4	+from torch_npu._compat.distributed import register_op_strategy
5	from torch.distributed.tensor._ops._pointwise_ops import pointwise_strategy	5	from torch.distributed.tensor._ops._pointwise_ops import pointwise_strategy
6		6
7		7

Mtorch_npu/dynamo/trace_rule.py+1-0

	@@ -45,6 +45,7 @@ torch_non_c_binding_in_graph_functions_npu = dict.fromkeys(
45	"torch.npu._memory_viz._frame_fmt",	45	"torch.npu._memory_viz._frame_fmt",
46	"torch.npu.amp.autocast_mode.custom_bwd",	46	"torch.npu.amp.autocast_mode.custom_bwd",
47	"torch.npu.amp.autocast_mode.custom_fwd",	47	"torch.npu.amp.autocast_mode.custom_fwd",
		48	+ "torch.npu.is_initialized",
48	"torch.npu._get_current_allocator",	49	"torch.npu._get_current_allocator",
49	"torch.npu.is_bf16_supported",	50	"torch.npu.is_bf16_supported",
50	"torch.npu.memory._get_current_allocator",	51	"torch.npu.memory._get_current_allocator",

Mtorch_npu/utils/_dynamo.py+3-1

	@@ -185,8 +185,10 @@ def patch_inductor_wrapper():
185	try:	185	try:
186	import torch_mlir	186	import torch_mlir
187	from torch_mlir import ir	187	from torch_mlir import ir
188	- except ImportError as e:	188	+ except ImportError as e:
189	raise ImportError("torch_mlir is not installed, install it first.") from e	189	raise ImportError("torch_mlir is not installed, install it first.") from e
		190	+ from torch._inductor.lowering import make_fallback as ori_make_fallback
		191	+ torch._inductor.lowering.make_fallback = ori_make_fallback
190	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu import (	192	from torch_npu._inductor.ascend_npu_ir.ascend_npu_ir.npu import (
191	npu_inductor_plugin,	193	npu_inductor_plugin,
192	)	194	)

Mtorch_npu/utils/dtensor.py+2-1

	@@ -1,6 +1,7 @@
1	import torch	1	import torch
2	from torch.distributed.tensor._ops._common_rules import pointwise_rule	2	from torch.distributed.tensor._ops._common_rules import pointwise_rule
3	-from torch.distributed.tensor._ops.utils import normalize_dims, register_prop_rule	3	+from torch.distributed.tensor._ops.utils import normalize_dims
		4	+from torch_npu._compat.distributed import register_prop_rule
4	from torch.distributed.tensor._ops._matrix_ops import bmm_strategy	5	from torch.distributed.tensor._ops._matrix_ops import bmm_strategy
5	from torch.distributed.tensor._ops._view_ops import (	6	from torch.distributed.tensor._ops._view_ops import (
6	register_op_strategy_map,	7	register_op_strategy_map,

Mtorch_npu/utils/storage.py+15-2

	@@ -16,11 +16,24 @@ def _rebuild_npu_tensor(storage, storage_offset, size, stride, requires_grad, ba
16	"please use 2.1 and newer torch to re-store the weight file."	16	"please use 2.1 and newer torch to re-store the weight file."
17	)	17	)
18	se._warn_legacy_serialization(warn_massages, "oldfile")	18	se._warn_legacy_serialization(warn_massages, "oldfile")
19	- tensor = torch.tensor([], dtype=storage.dtype, device=storage.device)	19	+ tensor = torch.empty(
		20	+ (0,),
		21	+ dtype=storage.dtype,
		22	+ device=storage._untyped_storage.device,
		23	+ requires_grad=requires_grad,
		24	+ )
20	tensor.set_(storage, storage_offset, size, stride)	25	tensor.set_(storage, storage_offset, size, stride)
21	tensor.requires_grad = requires_grad	26	tensor.requires_grad = requires_grad
22	tensor._backward_hooks = backward_hooks	27	tensor._backward_hooks = backward_hooks
23	- if not se.RE_MAP_CPU:	28	+ target_device = torch.device("cpu") if se.RE_MAP_CPU else torch.device("npu")
		29	+ is_fake_mode = (
		30	+ hasattr(torch, "_guards")
		31	+ and torch._guards.detect_fake_mode(None) is not None
		32	+ )
		33	+
		34	+ if is_fake_mode:
		35	+ tensor.fake_device = target_device
		36	+ elif not se.RE_MAP_CPU:
24	if isinstance(npu_storage_info, bool):	37	if isinstance(npu_storage_info, bool):
25	tensor = tensor.npu()	38	tensor = tensor.npu()
26	else:	39	else: