[feat][compilation]按torch版本兼容rms_norm/rope pattern的dtype cast分解-MindIE-SD-AtomGit

Mcsrc/ops/norm_rope_concat/op_kernel/norm_rope_concat.h+62-28

	@@ -79,7 +79,54 @@ private:
79		79
80	#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	80	#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
81	__aicore__ inline void DoMulAddVfCall(const LocalTensor<float> &x, uint32_t heads);	81	__aicore__ inline void DoMulAddVfCall(const LocalTensor<float> &x, uint32_t heads);
82	- __simd_vf__ inline void DoMulAddVf(__ubuf__ float* xBuf, __ubuf__ float* weightBuf, __ubuf__ float* biasBuf, uint32_t heads);	82	+ __simd_vf__ static inline void DoMulAddVf(__ubuf__ float* xBuf, __ubuf__ float* weightBuf,
		83	+ __ubuf__ float* biasBuf, uint32_t heads, uint32_t alignedNormDim)
		84	+ {
		85	+ RegTensor<float> xRegTensor;
		86	+ RegTensor<float> weightRegTensor;
		87	+ RegTensor<float> biasRegTensor;
		88	+
		89	+ uint32_t oneRepeatSize = AscendC::GetVecLen() / sizeof(float);
		90	+ uint32_t repeatTimes = (alignedNormDim + oneRepeatSize - 1) / oneRepeatSize;
		91	+
		92	+ for (uint32_t i = 0; i < heads; ++i) {
		93	+ uint32_t len = alignedNormDim;
		94	+ __ubuf__ float* tempXBuf = xBuf + i * len;
		95	+ __ubuf__ float* tempWeightBuf = weightBuf;
		96	+ __ubuf__ float* tempBiasBuf = biasBuf;
		97	+ for (uint32_t j = 0; j < repeatTimes; ++j) {
		98	+ MaskReg maskReg = UpdateMask<float>(len);
		99	+ LoadAlign(xRegTensor, tempXBuf + j * oneRepeatSize);
		100	+ LoadAlign(weightRegTensor, tempWeightBuf + j * oneRepeatSize);
		101	+ LoadAlign(biasRegTensor, tempBiasBuf + j * oneRepeatSize);
		102	+ MulDstAdd(xRegTensor, weightRegTensor, biasRegTensor, maskReg);
		103	+ StoreAlign(tempXBuf + j * oneRepeatSize, xRegTensor, maskReg);
		104	+ }
		105	+ }
		106	+ }
		107	+ __aicore__ inline void DoMulVfCall(const LocalTensor<float> &x, uint32_t heads);
		108	+ __simd_vf__ static inline void DoMulVf(__ubuf__ float* xBuf, __ubuf__ float* weightBuf, uint32_t heads,
		109	+ uint32_t alignedNormDim)
		110	+ {
		111	+ RegTensor<float> xRegTensor;
		112	+ RegTensor<float> weightRegTensor;
		113	+
		114	+ uint32_t oneRepeatSize = AscendC::GetVecLen() / sizeof(float);
		115	+ uint32_t repeatTimes = (alignedNormDim + oneRepeatSize - 1) / oneRepeatSize;
		116	+
		117	+ for (uint32_t i = 0; i < heads; ++i) {
		118	+ uint32_t len = alignedNormDim;
		119	+ __ubuf__ float* tempXBuf = xBuf + i * len;
		120	+ __ubuf__ float* tempWeightBuf = weightBuf;
		121	+ for (uint32_t j = 0; j < repeatTimes; ++j) {
		122	+ MaskReg maskReg = UpdateMask<float>(len);
		123	+ LoadAlign(xRegTensor, tempXBuf + j * oneRepeatSize);
		124	+ LoadAlign(weightRegTensor, tempWeightBuf + j * oneRepeatSize);
		125	+ Mul(xRegTensor, xRegTensor, weightRegTensor, maskReg);
		126	+ StoreAlign(tempXBuf + j * oneRepeatSize, xRegTensor, maskReg);
		127	+ }
		128	+ }
		129	+ }
83	#endif	130	#endif
84	private:	131	private:
85	TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQue_;	132	TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQue_;
	@@ -128,38 +175,21 @@ __aicore__ inline void NormOperationForward<isTraining>::DoMulAdd(const LocalTen
128	}	175	}
129		176
130	#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	177	#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
131	- template <bool isTraining>
132	- __simd_vf__ inline void NormOperationForward<isTraining>::DoMulAddVf(
133	- __ubuf__ float* xBuf, __ubuf__ float* weightBuf, __ubuf__ float* biasBuf, uint32_t heads) {
134	- RegTensor<float> xRegTensor;
135	- RegTensor<float> weightRegTensor;
136	- RegTensor<float> biasRegTensor;
137	-
138	- uint32_t oneRepeatSize = AscendC::GetVecLen() / sizeof(float);
139	- uint32_t repeatTimes = (this->alignedNormDim_ + oneRepeatSize - 1) / oneRepeatSize;
140	- for (uint32_t i = 0; i < heads; ++i) {
141	- uint32_t len = this->alignedNormDim_;
142	- __ubuf__ float* tempXBuf = xBuf + i * len;
143	- __ubuf__ float* tempWeightBuf = weightBuf;
144	- __ubuf__ float* tempBiasBuf = biasBuf;
145	- for (uint32_t j = 0; j < repeatTimes; ++j) {
146	- MaskReg maskReg = UpdateMask<float>(len);
147	- LoadAlign(xRegTensor, tempXBuf + j * oneRepeatSize);
148	- LoadAlign(weightRegTensor, tempWeightBuf + j * oneRepeatSize);
149	- LoadAlign(biasRegTensor, tempBiasBuf + j * oneRepeatSize);
150	- MulDstAdd(xRegTensor, weightRegTensor, biasRegTensor, maskReg);
151	- StoreAlign(tempXBuf + j * oneRepeatSize, xRegTensor, maskReg);
152	- }
153	- }
154	- }
155	-
156	template <bool isTraining>	178	template <bool isTraining>
157	__aicore__ inline void NormOperationForward<isTraining>::DoMulAddVfCall(const LocalTensor<float> &x, uint32_t heads)	179	__aicore__ inline void NormOperationForward<isTraining>::DoMulAddVfCall(const LocalTensor<float> &x, uint32_t heads)
158	{	180	{
159	__ubuf__ float* xBuf = (__ubuf__ float*)x.GetPhyAddr();	181	__ubuf__ float* xBuf = (__ubuf__ float*)x.GetPhyAddr();
160	__ubuf__ float* weightBuf = (__ubuf__ float*)weight_.GetPhyAddr();	182	__ubuf__ float* weightBuf = (__ubuf__ float*)weight_.GetPhyAddr();
161	__ubuf__ float* biasBuf = (__ubuf__ float*)bias_.GetPhyAddr();	183	__ubuf__ float* biasBuf = (__ubuf__ float*)bias_.GetPhyAddr();
162	- DoMulAddVf(xBuf, weightBuf, biasBuf, heads);	184	+ DoMulAddVf(xBuf, weightBuf, biasBuf, heads, this->alignedNormDim_);
		185	+ }
		186	+
		187	+ template <bool isTraining>
		188	+ __aicore__ inline void NormOperationForward<isTraining>::DoMulVfCall(const LocalTensor<float> &x, uint32_t heads)
		189	+ {
		190	+ __ubuf__ float* xBuf = (__ubuf__ float*)x.GetPhyAddr();
		191	+ __ubuf__ float* weightBuf = (__ubuf__ float*)weight_.GetPhyAddr();
		192	+ DoMulVf(xBuf, weightBuf, heads, this->alignedNormDim_);
163	}	193	}
164	#endif	194	#endif
165		195
	@@ -323,7 +353,11 @@ __aicore__ inline void NormOperationForward<isTraining>::DoRMSNorm(const LocalTe
323	Div(x, x, tmp0, size);	353	Div(x, x, tmp0, size);
324	PipeBarrier<PIPE_V>();	354	PipeBarrier<PIPE_V>();
325	if constexpr (normType == NormType::RMS_NORM_AFFINE) {	355	if constexpr (normType == NormType::RMS_NORM_AFFINE) {
326	- DoMul(x, heads);	356	+ #if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
		357	+ DoMulVfCall(x, heads);
		358	+ #else
		359	+ DoMul(x, heads);
		360	+ #endif
327	}	361	}
328	PipeBarrier<PIPE_V>();	362	PipeBarrier<PIPE_V>();
329	}	363	}

Mmindiesd/compilation/patterns/rms_norm_pattern.py+13-1

	@@ -10,22 +10,34 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
		13	+import re
		14	+
13	import torch	15	import torch
		16	+
14	from ..passes.register_pattern_to_pass import PatternBase	17	from ..passes.register_pattern_to_pass import PatternBase
15		18
16	if hasattr(torch.npu, "is_available"):	19	if hasattr(torch.npu, "is_available"):
17	npu_available = torch.npu.is_available()	20	npu_available = torch.npu.is_available()
18	if npu_available:	21	if npu_available:
19	import torch_npu	22	import torch_npu
		23	+
20	import mindiesd	24	import mindiesd
21		25
		26	+_torch_version_tuple = tuple(int(x) for x in re.match(r"(\d+)\.(\d+)", torch.__version__).groups())
		lijinxi8月19日抽取公共逻辑，在后续整改https://gitcode.com/Ascend/MindIE-SD/issues/315
		27	+# torch_npu >= 2.9 移除 Tensor.to 的 NPU 过适配(MR 30358)后,
		28	+# rms_norm 分解中的 dtype cast 由 npu._npu_dtype_cast 变为 aten._to_copy。
		29	+IS_TORCH_GE_29 = _torch_version_tuple >= (2, 9)
		30	+
22		31
23	def create(dtype, epsilon=1e-6):	32	def create(dtype, epsilon=1e-6):
24	if "2.6.0" in torch.__version__:	33	if "2.6.0" in torch.__version__:
25	_dtype_cast_func = torch.ops.npu.npu_dtype_cast.default	34	_dtype_cast_func = torch.ops.npu.npu_dtype_cast.default
26	_eps_in_bf16 = torch.tensor(epsilon, dtype=torch.bfloat16, device="cpu").item()	35	_eps_in_bf16 = torch.tensor(epsilon, dtype=torch.bfloat16, device="cpu").item()
27	else:	36	else:
28	- _dtype_cast_func = torch.ops.npu._npu_dtype_cast.default	37	+ if IS_TORCH_GE_29:
		38	+ _dtype_cast_func = lambda x, dtype: torch.ops.aten._to_copy.default(x, dtype=dtype)
		39	+ else:
		40	+ _dtype_cast_func = torch.ops.npu._npu_dtype_cast.default
29	_eps_in_fp32 = torch.tensor(epsilon, dtype=torch.float32, device="cpu").item()	41	_eps_in_fp32 = torch.tensor(epsilon, dtype=torch.float32, device="cpu").item()
30		42
31		43

Mmindiesd/compilation/patterns/rope_pattern.py+14-3

	@@ -10,21 +10,33 @@
10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	10	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
11	# See the Mulan PSL v2 for more details.	11	# See the Mulan PSL v2 for more details.
12		12
		13	+import re
		14	+
13	import torch	15	import torch
		16	+
14	from ..passes.register_pattern_to_pass import PatternBase	17	from ..passes.register_pattern_to_pass import PatternBase
15		18
16	if hasattr(torch.npu, "is_available"):	19	if hasattr(torch.npu, "is_available"):
17	npu_available = torch.npu.is_available()	20	npu_available = torch.npu.is_available()
18	if npu_available:	21	if npu_available:
19	import torch_npu	22	import torch_npu
		23	+
20	import mindiesd	24	import mindiesd
21		25
		26	+_torch_version_tuple = tuple(int(x) for x in re.match(r"(\d+)\.(\d+)", torch.__version__).groups())
		27	+# torch_npu >= 2.9 移除 Tensor.to 的 NPU 过适配(MR 30358)后,
		28	+# dtype cast 由 npu._npu_dtype_cast 变为 aten._to_copy。
		29	+IS_TORCH_GE_29 = _torch_version_tuple >= (2, 9)
		30	+
22		31
23	def create(dtype):	32	def create(dtype):
24	if "2.6.0" in torch.__version__:	33	if "2.6.0" in torch.__version__:
25	_dtype_cast_func = torch.ops.npu.npu_dtype_cast.default	34	_dtype_cast_func = torch.ops.npu.npu_dtype_cast.default
26	else:	35	else:
27	- _dtype_cast_func = torch.ops.npu._npu_dtype_cast.default	36	+ if IS_TORCH_GE_29:
		37	+ _dtype_cast_func = lambda x, dtype: torch.ops.aten._to_copy.default(x, dtype=dtype)
		38	+ else:
		39	+ _dtype_cast_func = torch.ops.npu._npu_dtype_cast.default
28		40
29		41
30	class RopePattern(PatternBase):	42	class RopePattern(PatternBase):
	@@ -51,8 +63,7 @@ def create(dtype):
51	cos_part = _dtype_cast_func(x, dtype) * cos	63	cos_part = _dtype_cast_func(x, dtype) * cos
52	sin_part = _dtype_cast_func(x_rotated, dtype) * sin	64	sin_part = _dtype_cast_func(x_rotated, dtype) * sin
53	x_out = cos_part + sin_part	65	x_out = cos_part + sin_part
54	- x_out.type_as(x)	66	+ return torch.ops.aten._to_copy.default(x_out, dtype=x.dtype)
55	- return x_out
56	return func(x, cos, sin)	67	return func(x, cos, sin)
57		68
58	@staticmethod	69	@staticmethod

Atests/compilation/patterns/test_dtype_cast_compat.py+60-0

@@ -0,0 +1,60 @@
	1	+import os
	2	+import unittest
	3	+
	4	+import torch
	5	+from torch.fx.experimental.proxy_tensor import make_fx
	6	+
	7	+from mindiesd.compilation.patterns.rms_norm_pattern import (
	8	+ IS_TORCH_GE_29,
	9	+)
	10	+from mindiesd.compilation.patterns.rms_norm_pattern import (
	11	+ create as create_rmsnorm_pattern,
	12	+)
	13	+from mindiesd.compilation.patterns.rope_pattern import create as create_rope_pattern
	14	+
	15	+
	16	+def _trace(pattern, args):
	17	+ return make_fx(pattern)(*args)
	18	+
	19	+
	20	+@unittest.skipIf(
	21	+ os.environ.get("MINDIE_TEST_MODE", "ALL") == "CPU",
	22	+ "Skip NPU-dependent tests when MINDIE_TEST_MODE is CPU.",
	23	+)
	24	+@unittest.skipUnless(IS_TORCH_GE_29, "dtype cast 适配分支仅对 torch >= 2.9 生效")
	C C changetheway8月19日这里参考其他测试用例补充一下测试设备分类 Z zqxu 8月19日评论： 8月19日评论：
	25	+class TestPatternDtypeCastCompat(unittest.TestCase):
	26	+ """回归测试：torch_npu >= 2.9 移除 Tensor.to 的 NPU 过适配(MR 30358)后,
	27	+ AOT 分解中的 dtype cast 由 npu._npu_dtype_cast 变为 aten._to_copy。
	28	+ 本测试验证 pattern 期望图已适配为 _to_copy, 否则 RMSNorm/RoPE
	29	+ 融合在 torch 2.9+ 环境会全部失配。
	30	+ """
	31	+
	32	+ def test_rmsnorm_pattern_first_cast_is_to_copy(self):
	33	+ pattern_cls = create_rmsnorm_pattern(dtype=torch.bfloat16, epsilon=1e-6)
	34	+ hidden_states = torch.empty(2, 2, 2, 2, dtype=torch.bfloat16, device="meta")
	35	+ weight = torch.empty(2, dtype=torch.bfloat16, device="meta")
	36	+
	37	+ gm = _trace(pattern_cls.pattern, (hidden_states, weight))
	38	+ call_targets = [n.target for n in gm.graph.nodes if n.op == "call_function"]
	39	+
	40	+ # rms_norm 分解序列的第一个 op 必须是 _to_copy（fp32 提升）
	41	+ self.assertEqual(call_targets[0], torch.ops.aten._to_copy.default)
	42	+
	43	+ def test_rope_pattern_casts_are_to_copy(self):
	44	+ pattern_cls = create_rope_pattern(torch.float32)
	45	+ x = torch.empty(2, 2, 2, 2, dtype=torch.bfloat16, device="meta")
	46	+ cos = torch.empty(1, 2, 1, 2, dtype=torch.float32, device="meta")
	47	+ sin = torch.empty(1, 2, 1, 2, dtype=torch.float32, device="meta")
	48	+
	49	+ gm = _trace(pattern_cls.pattern, (x, cos, sin))
	50	+ call_targets = [n.target for n in gm.graph.nodes if n.op == "call_function"]
	51	+
	52	+ to_copy_count = sum(t == torch.ops.aten._to_copy.default for t in call_targets)
	53	+ # x 与 x_rotated 的 fp32 cast + 末尾 .to(x.dtype) cast
	54	+ self.assertGreaterEqual(to_copy_count, 3)
	55	+ # 最后一个 op 必须是 _to_copy（对应真实图 apply_rotary_emb 的 .to(x.dtype)）
	56	+ self.assertEqual(call_targets[-1], torch.ops.aten._to_copy.default)
	57	+
	58	+
	59	+if __name__ == "__main__":
	60	+ unittest.main()