cdist_grad 950 use compare_select-ops-math-AtomGit

已合并

cdist_grad 950 use compare_select #3278

chenxingyu18创建于 6月11日从chenxingyu18/ops-math： master 合入到 cann/ops-math： master

cdist_grad 950 use compare_select #3278

已合并

chenxingyu18创建于 6月11日

从chenxingyu18/ops-math master合入到cann/ops-math master

共 3 个文件变更+277-87

Mmath/cdist_grad/op_kernel/arch35/cdist_grad_dag.h+76-85

	@@ -10,19 +10,18 @@
10		10
11	/*!	11	/*!
12	* \file cdist_grad_dag.h	12	* \file cdist_grad_dag.h
13	- * \brief cdist_grad dag — six DAGs, all Compare+Select replaced by arithmetic	13	+ * \brief cdist_grad dag — six DAGs, Compare+Select via custom MicroAPI operators
14	*	14	*
15	- * Arithmetic replacements:	15	+ * Custom operators (MicroAPI CompareScalar + Select, avoids DAG multi-consumer):
16	- * sign: diff / (\|diff\| + eps) replaces Compare(NE) + Select	16	+ * CdistGradSignOp: sign(x) via GT/LT compare → 1.0 / -1.0 / 0.0
17	- * mask: (d + \|d\| + eps) / (2\|d\| + eps) replaces Compare(GE/EQ) + Select	17	+ * CdistGradMaskEQOp: a == b ? 1.0 : 0.0 via Sub + EQ compare
18	- * nz_x: x + eps replaces Compare(NE) + Select for safe divisor	18	+ * CdistGradMaskNEZeroOp: x != 0 ? 1.0 : 0.0 via EQ compare (inverted)
19	- * mask_nz: x / (x + eps) replaces Compare(NE) + Select for zero-out
20	*	19	*
21	* CdistGradP0Dag: p == 0 → output zeros	20	* CdistGradP0Dag: p == 0 → output zeros
22	- * CdistGradP1Dag: p == 1 → grad * sign	21	+ * CdistGradP1Dag: p == 1 → grad * sign(diff)
23	* CdistGradP2Dag: p == 2 → grad * diff / cdist	22	* CdistGradP2Dag: p == 2 → grad * diff / cdist
24	- * CdistGradDag: 0<p<2 → sign * \|diff\|^(p-1) * grad / \|cdist\|^(p-1)	23	+ * CdistGradDag: 0<p<2 → sign * (\|diff\|/cdist)^(p-1) * grad * masks
25	- * CdistGradLargePDag: p>2 → diff * \|diff\|^(p-2) * grad / \|cdist\|^(p-1)	24	+ * CdistGradLargePDag: p>2 → sign * (\|diff\|/cdist)^(p-1) * grad * mask_cdist
26	* CdistGradInfDag: p==inf → grad * sign * mask(\|diff\| >= cdist)	25	* CdistGradInfDag: p==inf → grad * sign * mask(\|diff\| >= cdist)
27	*/	26	*/
28		27
	@@ -34,6 +33,7 @@
34	#include "atvoss/util/vec.h"	33	#include "atvoss/util/vec.h"
35	#include "atvoss/util/placeholder.h"	34	#include "atvoss/util/placeholder.h"
36	#include "atvoss/reduce/reduce_operator.h"	35	#include "atvoss/reduce/reduce_operator.h"
		36	+#include "cdist_grad_operator.h"
37		37
38	namespace CdistGrad {	38	namespace CdistGrad {
39	using namespace Ops::Base;	39	using namespace Ops::Base;
	@@ -69,13 +69,11 @@ struct CdistGradP0Dag {
69		69
70	// ---------------------------------------------------------------------------	70	// ---------------------------------------------------------------------------
71	// CdistGradP1Dag — p == 1	71	// CdistGradP1Dag — p == 1
72	-// sign = diff / (\|diff\| + eps) replaces Compare(NE)+Select for nz_diff	72	+// sign = CdistGradSignOp(diff) via MicroAPI CompareScalar GT/LT + Select
73	// result = grad * sign	73	// result = grad * sign
74	// ---------------------------------------------------------------------------	74	// ---------------------------------------------------------------------------
75	template <typename T, typename PromoteT>	75	template <typename T, typename PromoteT>
76	struct CdistGradP1Dag {	76	struct CdistGradP1Dag {
77	- using Eps = MAKE_CONST(PromoteT, 1e-30);
78	-
79	using OpCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;	77	using OpCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
80	using CastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInGrad>;	78	using CastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInGrad>;
81		79
	@@ -86,9 +84,7 @@ struct CdistGradP1Dag {
86	using CastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInX2>;	84	using CastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInX2>;
87		85
88	using OpDiff = Bind<Vec::Sub<PromoteT>, CastX1, CastX2>;	86	using OpDiff = Bind<Vec::Sub<PromoteT>, CastX1, CastX2>;
89	- using OpDiffAbs = Bind<Vec::Abs<PromoteT>, OpDiff>;	87	+ using OpSign = Bind<CdistGradSignOp<PromoteT>, OpDiff>;
90	- using SafeAbsDiff = Bind<Vec::Adds<PromoteT>, OpDiffAbs, Eps>; // \|diff\| + eps
91	- using OpSign = Bind<Vec::Div<PromoteT>, OpDiff, SafeAbsDiff>; // diff / (\|diff\| + eps)
92		88
93	using OpRes = Bind<Vec::Mul<PromoteT>, CastGrad, OpSign>;	89	using OpRes = Bind<Vec::Mul<PromoteT>, CastGrad, OpSign>;
94		90
	@@ -103,8 +99,9 @@ struct CdistGradP1Dag {
103		99
104	// ---------------------------------------------------------------------------	100	// ---------------------------------------------------------------------------
105	// CdistGradP2Dag — p == 2	101	// CdistGradP2Dag — p == 2
106	-// result = grad * diff / (cdist + eps)	102	+// mask = CdistGradMaskNEZeroOp(cdist) via MicroAPI
107	-// When cdist = 0: diff = 0, numerator = 0, result = 0 automatically	103	+// result = grad * diff / (cdist + eps) * mask
		104	+// Matches PyTorch: dist == 0 ? 0 : grad * diff / dist
108	// ---------------------------------------------------------------------------	105	// ---------------------------------------------------------------------------
109	template <typename T, typename PromoteT>	106	template <typename T, typename PromoteT>
110	struct CdistGradP2Dag {	107	struct CdistGradP2Dag {
	@@ -125,7 +122,11 @@ struct CdistGradP2Dag {
125	using OpDiff = Bind<Vec::Sub<PromoteT>, CastX1, CastX2>;	122	using OpDiff = Bind<Vec::Sub<PromoteT>, CastX1, CastX2>;
126	using SafeCdist = Bind<Vec::Adds<PromoteT>, CastCdist, Eps>; // cdist + eps	123	using SafeCdist = Bind<Vec::Adds<PromoteT>, CastCdist, Eps>; // cdist + eps
127	using OpNumerator = Bind<Vec::Mul<PromoteT>, CastGrad, OpDiff>; // grad * diff	124	using OpNumerator = Bind<Vec::Mul<PromoteT>, CastGrad, OpDiff>; // grad * diff
128	- using OpResult = Bind<Vec::Div<PromoteT>, OpNumerator, SafeCdist>; // grad * diff / (cdist+eps)	125	+ using OpDivResult = Bind<Vec::Div<PromoteT>, OpNumerator, SafeCdist>; // grad * diff / (cdist+eps)
		126	+
		127	+ // mask: cdist != 0 ? 1.0 : 0.0 — zeros out when cdist=0
		128	+ using OpMask = Bind<CdistGradMaskNEZeroOp<PromoteT>, CastCdist>;
		129	+ using OpResult = Bind<Vec::Mul<PromoteT>, OpDivResult, OpMask>;
129		130
130	using ReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, OpResult>;	131	using ReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, OpResult>;
131	using CastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, ReduceOp0>;	132	using CastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, ReduceOp0>;
	@@ -138,15 +139,16 @@ struct CdistGradP2Dag {
138		139
139	// ---------------------------------------------------------------------------	140	// ---------------------------------------------------------------------------
140	// CdistGradDag — 0 < p < 2, p != 1	141	// CdistGradDag — 0 < p < 2, p != 1
141	-// sign = diff / (\|diff\| + eps)	142	+// sign = CdistGradSignOp(diff) via MicroAPI
142	-// safe_diff = \|diff\| + eps (prevents log(0))	143	+// safe_cdist = cdist + eps (prevents div-by-0)
143	-// safe_cdist = cdist + eps (prevents log(0) and div-by-0)	144	+// ratio = \|diff\| / safe_cdist
144	-// mask_diff = \|diff\| / (\|diff\| + eps) (zero-out when \|diff\|=0)	145	+// pow_ratio = exp(log(ratio + eps) * (p-1)) single log+exp chain
145	-// mask_cdist = cdist / (cdist + eps) (zero-out when cdist=0)	146	+// mask_diff = CdistGradMaskNEZeroOp(\|diff\|) via MicroAPI
146	-// num = sign * exp(log(safe_diff) * (p-1))	147	+// mask_cdist = CdistGradMaskNEZeroOp(cdist) via MicroAPI
147	-// numerator = num * grad	148	+// result = sign * pow_ratio * grad * mask_cdist * mask_diff
148	-// denominator = exp(log(safe_cdist) * (p-1))	149	+//
149	-// result = numerator / denominator * mask_cdist * mask_diff	150	+// Optimization: (\|diff\|/cdist)^(p-1) replaces separate \|diff\|^(p-1)/cdist^(p-1),
		151	+// cutting transcendental ops from 2 log + 2 exp + 1 div → 1 log + 1 exp + 1 div.
150	// Var<0>: power = p - 1	152	// Var<0>: power = p - 1
151	// ---------------------------------------------------------------------------	153	// ---------------------------------------------------------------------------
152	template <typename T, typename PromoteT>	154	template <typename T, typename PromoteT>
	@@ -168,33 +170,28 @@ struct CdistGradDag {
168	// diff & abs	170	// diff & abs
169	using DagDiff = Bind<Vec::Sub<PromoteT>, DagCastX1, DagCastX2>;	171	using DagDiff = Bind<Vec::Sub<PromoteT>, DagCastX1, DagCastX2>;
170	using DagDiffAbs = Bind<Vec::Abs<PromoteT>, DagDiff>;	172	using DagDiffAbs = Bind<Vec::Abs<PromoteT>, DagDiff>;
171	- using DagSafeAbsDiff = Bind<Vec::Adds<PromoteT>, DagDiffAbs, DagEps>; // \|diff\| + eps
172	using DagSafeCdist = Bind<Vec::Adds<PromoteT>, DagCastCdist, DagEps>; // cdist + eps	173	using DagSafeCdist = Bind<Vec::Adds<PromoteT>, DagCastCdist, DagEps>; // cdist + eps
173		174
174	- // sign = diff / (\|diff\| + eps)	175	+ // sign via MicroAPI
175	- using DagSign = Bind<Vec::Div<PromoteT>, DagDiff, DagSafeAbsDiff>;	176	+ using DagSign = Bind<CdistGradSignOp<PromoteT>, DagDiff>;
176		177
177	- // masks: 0 when input=0, ~1 otherwise	178	+ // masks via MicroAPI: 0 when input=0, 1 otherwise
178	- using DagMaskDiff = Bind<Vec::Div<PromoteT>, DagDiffAbs, DagSafeAbsDiff>; // \|diff\|/(\|diff\|+eps)	179	+ using DagMaskDiff = Bind<CdistGradMaskNEZeroOp<PromoteT>, DagDiffAbs>;
179	- using DagMaskCdist = Bind<Vec::Div<PromoteT>, DagCastCdist, DagSafeCdist>; // cdist/(cdist+eps)	180	+ using DagMaskCdist = Bind<CdistGradMaskNEZeroOp<PromoteT>, DagCastCdist>;
180		181
181	- // power: (safe_x)^(p-1) via log/exp, Var<0> = p-1	182	+ // ratio = \|diff\| / (cdist + eps), then (ratio)^(p-1) via single log+exp chain
182	- using DagPowDiff = Bind<Vec::Exp<PromoteT>,	183	+ using DagRatio = Bind<Vec::Div<PromoteT>, DagDiffAbs, DagSafeCdist>;
183	- Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, DagSafeAbsDiff>,	184	+ using DagSafeRatio = Bind<Vec::Adds<PromoteT>, DagRatio, DagEps>;
184	- Placeholder::Var<PromoteT, 0>>>;	185	+ using DagPowRatio = Bind<Vec::Exp<PromoteT>,
185	- using DagPowCdist = Bind<Vec::Exp<PromoteT>,	186	+ Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, DagSafeRatio>,
186	- Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, DagSafeCdist>,	187	+ Placeholder::Var<PromoteT, 0>>>;
187	- Placeholder::Var<PromoteT, 0>>>;
188		188
189	- // num = sign * \|diff\|^(p-1), numerator = num * grad	189	+ // num = sign * (\|diff\|/cdist)^(p-1), result = num * grad * masks
190	- using DagNum = Bind<Vec::Mul<PromoteT>, DagSign, DagPowDiff>;	190	+ using DagNum = Bind<Vec::Mul<PromoteT>, DagSign, DagPowRatio>;
191	using DagNumerator = Bind<Vec::Mul<PromoteT>, DagNum, DagCastGrad>;	191	using DagNumerator = Bind<Vec::Mul<PromoteT>, DagNum, DagCastGrad>;
192		192
193	- // res = numerator / denominator
194	- using DagDivResult = Bind<Vec::Div<PromoteT>, DagNumerator, DagPowCdist>;
195	-
196	// apply masks: zero-out when cdist=0 or \|diff\|=0	193	// apply masks: zero-out when cdist=0 or \|diff\|=0
197	- using DagMaskedCdist = Bind<Vec::Mul<PromoteT>, DagDivResult, DagMaskCdist>;	194	+ using DagMaskedCdist = Bind<Vec::Mul<PromoteT>, DagNumerator, DagMaskCdist>;
198	using DagResult = Bind<Vec::Mul<PromoteT>, DagMaskedCdist, DagMaskDiff>;	195	using DagResult = Bind<Vec::Mul<PromoteT>, DagMaskedCdist, DagMaskDiff>;
199		196
200	using DagReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, DagResult>;	197	using DagReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, DagResult>;
	@@ -208,14 +205,20 @@ struct CdistGradDag {
208		205
209	// ---------------------------------------------------------------------------	206	// ---------------------------------------------------------------------------
210	// CdistGradLargePDag — p > 2	207	// CdistGradLargePDag — p > 2
		208	+// Equivalent to: sign(diff) * (\|diff\|/cdist)^(p-1) * grad
		209	+// since diff * \|diff\|^(p-2) / cdist^(p-1) = sign * \|diff\|^(p-1) / cdist^(p-1)
		210	+// = sign * (\|diff\|/cdist)^(p-1)
		211	+//
		212	+// sign = CdistGradSignOp(diff) via MicroAPI
211	// safe_cdist = cdist + eps	213	// safe_cdist = cdist + eps
212	-// mask_cdist = cdist / (cdist + eps)	214	+// ratio = \|diff\| / safe_cdist
213	-// num = diff * \|diff\|^(p-2) (0 when \|diff\|=0 for p>2)	215	+// pow_ratio = exp(log(ratio + eps) * (p-1)) single log+exp chain
214	-// numerator = num * grad	216	+// mask_cdist = CdistGradMaskNEZeroOp(cdist) via MicroAPI
215	-// denominator = \|cdist\|^(p-1)	217	+// result = sign * pow_ratio * grad * mask_cdist
216	-// result = numerator / denominator * mask_cdist	218	+//
217	-// Var<0>: power_diff = p - 2	219	+// Optimization: same ratio approach as CdistGradDag, reducing from
218	-// Var<1>: power_cdist = p - 1	220	+// 2 log + 2 exp + 1 div → 1 log + 1 exp + 1 div.
		221	+// Var<0>: power = p - 1
219	// ---------------------------------------------------------------------------	222	// ---------------------------------------------------------------------------
220	template <typename T, typename PromoteT>	223	template <typename T, typename PromoteT>
221	struct CdistGradLargePDag {	224	struct CdistGradLargePDag {
	@@ -235,28 +238,25 @@ struct CdistGradLargePDag {
235		238
236	using LpDiff = Bind<Vec::Sub<PromoteT>, LpCastX1, LpCastX2>;	239	using LpDiff = Bind<Vec::Sub<PromoteT>, LpCastX1, LpCastX2>;
237	using LpDiffAbs = Bind<Vec::Abs<PromoteT>, LpDiff>;	240	using LpDiffAbs = Bind<Vec::Abs<PromoteT>, LpDiff>;
238	- using LpSafeAbsDiff = Bind<Vec::Adds<PromoteT>, LpDiffAbs, LpEps>; // \|diff\| + eps (prevents log(0))
239	using LpSafeCdist = Bind<Vec::Adds<PromoteT>, LpCastCdist, LpEps>; // cdist + eps	241	using LpSafeCdist = Bind<Vec::Adds<PromoteT>, LpCastCdist, LpEps>; // cdist + eps
240	- using LpMaskCdist = Bind<Vec::Div<PromoteT>, LpCastCdist, LpSafeCdist>; // cdist/(cdist+eps)
241		242
242	- // \|diff\|^(p-2), Var<0> = p-2	243	+ // sign via MicroAPI
243	- using LpPowDiff = Bind<Vec::Exp<PromoteT>,	244	+ using LpSign = Bind<CdistGradSignOp<PromoteT>, LpDiff>;
244	- Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, LpSafeAbsDiff>,
245	- Placeholder::Var<PromoteT, 0>>>;
246		245
247	- // \|cdist\|^(p-1), Var<1> = p-1	246	+ // mask via MicroAPI: cdist != 0 ? 1.0 : 0.0
248	- using LpPowCdist = Bind<Vec::Exp<PromoteT>,	247	+ using LpMaskCdist = Bind<CdistGradMaskNEZeroOp<PromoteT>, LpCastCdist>;
249	- Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, LpSafeCdist>,
250	- Placeholder::Var<PromoteT, 1>>>;
251		248
252	- // num = diff * \|diff\|^(p-2)	249	+ // ratio = \|diff\| / (cdist + eps), then (ratio)^(p-1) via single log+exp chain
253	- using LpNum = Bind<Vec::Mul<PromoteT>, LpDiff, LpPowDiff>;	250	+ using LpRatio = Bind<Vec::Div<PromoteT>, LpDiffAbs, LpSafeCdist>;
254	- // numerator = num * grad	251	+ using LpSafeRatio = Bind<Vec::Adds<PromoteT>, LpRatio, LpEps>;
		252	+ using LpPowRatio = Bind<Vec::Exp<PromoteT>,
		253	+ Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, LpSafeRatio>,
		254	+ Placeholder::Var<PromoteT, 0>>>;
		255	+
		256	+ // num = sign * (\|diff\|/cdist)^(p-1), result = num * grad * mask
		257	+ using LpNum = Bind<Vec::Mul<PromoteT>, LpSign, LpPowRatio>;
255	using LpNumerator = Bind<Vec::Mul<PromoteT>, LpNum, LpCastGrad>;	258	using LpNumerator = Bind<Vec::Mul<PromoteT>, LpNum, LpCastGrad>;
256	- // res = numerator / \|cdist\|^(p-1)	259	+ using LpResult = Bind<Vec::Mul<PromoteT>, LpNumerator, LpMaskCdist>;
257	- using LpRawResult = Bind<Vec::Div<PromoteT>, LpNumerator, LpPowCdist>;
258	- // zero-out when cdist = 0
259	- using LpResult = Bind<Vec::Mul<PromoteT>, LpRawResult, LpMaskCdist>;
260		260
261	using LpReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, LpResult>;	261	using LpReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, LpResult>;
262	using LpCastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, LpReduceOp0>;	262	using LpCastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, LpReduceOp0>;
	@@ -269,14 +269,12 @@ struct CdistGradLargePDag {
269		269
270	// ---------------------------------------------------------------------------	270	// ---------------------------------------------------------------------------
271	// CdistGradInfDag — p == inf	271	// CdistGradInfDag — p == inf
272	-// sign = diff / (\|diff\| + eps)	272	+// sign = CdistGradSignOp(diff) via MicroAPI CompareScalar GT/LT + Select
273	-// mask = (d + \|d\| + eps) / (2\|d\| + eps) where d = \|diff\| - cdist	273	+// mask = CdistGradMaskEQOp(\|diff\|, cdist) via MicroAPI Sub + CompareScalar EQ + Select
274	// result = grad * sign * mask	274	// result = grad * sign * mask
275	// ---------------------------------------------------------------------------	275	// ---------------------------------------------------------------------------
276	template <typename T, typename PromoteT>	276	template <typename T, typename PromoteT>
277	struct CdistGradInfDag {	277	struct CdistGradInfDag {
278	- using InfEps = MAKE_CONST(PromoteT, 1e-30);
279	-
280	using InfCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;	278	using InfCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
281	using InfCastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, InfCopyInGrad>;	279	using InfCastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, InfCopyInGrad>;
282		280
	@@ -292,18 +290,11 @@ struct CdistGradInfDag {
292	using InfDiff = Bind<Vec::Sub<PromoteT>, InfCastX1, InfCastX2>;	290	using InfDiff = Bind<Vec::Sub<PromoteT>, InfCastX1, InfCastX2>;
293	using InfDiffAbs = Bind<Vec::Abs<PromoteT>, InfDiff>;	291	using InfDiffAbs = Bind<Vec::Abs<PromoteT>, InfDiff>;
294		292
295	- // sign = diff / (\|diff\| + eps)	293	+ // sign = MicroAPI sign(diff)
296	- using InfSafeAbsDiff = Bind<Vec::Adds<PromoteT>, InfDiffAbs, InfEps>;	294	+ using InfSign = Bind<CdistGradSignOp<PromoteT>, InfDiff>;
297	- using InfSign = Bind<Vec::Div<PromoteT>, InfDiff, InfSafeAbsDiff>;
298		295
299	- // mask: d = \|diff\| - cdist, mask = (d + \|d\| + eps) / (2\|d\| + eps)	296	+ // mask = MicroAPI (\|diff\| == cdist) ? 1.0 : 0.0
300	- using InfD = Bind<Vec::Sub<PromoteT>, InfDiffAbs, InfCastCdist>;	297	+ using InfMask = Bind<CdistGradMaskEQOp<PromoteT>, InfDiffAbs, InfCastCdist>;
301	- using InfDAbs = Bind<Vec::Abs<PromoteT>, InfD>;
302	- using InfTwoAbsD = Bind<Vec::Add<PromoteT>, InfDAbs, InfDAbs>;
303	- using InfNumer = Bind<Vec::Add<PromoteT>, InfD, InfDAbs>;
304	- using InfNumerEps = Bind<Vec::Adds<PromoteT>, InfNumer, InfEps>;
305	- using InfDenom = Bind<Vec::Adds<PromoteT>, InfTwoAbsD, InfEps>;
306	- using InfMask = Bind<Vec::Div<PromoteT>, InfNumerEps, InfDenom>;
307		298
308	// result = grad * sign * mask	299	// result = grad * sign * mask
309	using InfGradSign = Bind<Vec::Mul<PromoteT>, InfCastGrad, InfSign>;	300	using InfGradSign = Bind<Vec::Mul<PromoteT>, InfCastGrad, InfSign>;

Amath/cdist_grad/op_kernel/arch35/cdist_grad_operator.h+200-0

@@ -0,0 +1,200 @@
	1	+/**
	2	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+ * CANN Open Software License Agreement Version 2.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+
	11	+/*!
	12	+ * \file cdist_grad_operator.h
	13	+ * \brief Custom Vec operators for cdist_grad — Compare+Select via MicroAPI
	14	+ *
	15	+ * By encapsulating CompareScalar + Select inside a single Vec operator node,
	16	+ * we avoid the multi-consumer buffer conflicts that occur when Compare and Select
	17	+ * are separate DAG nodes.
	18	+ *
	19	+ * Operators:
	20	+ * CdistGradSignOp — sign(x): 1.0 / -1.0 / 0.0 (ElemwiseUnaryOP)
	21	+ * CdistGradMaskGEOp — a >= b ? 1.0 : 0.0 (ElemwiseBinaryOP)
	22	+ * CdistGradMaskNEZeroOp — x != 0 ? 1.0 : 0.0 (ElemwiseUnaryOP)
	23	+ */
	24	+
	25	+#ifndef CDIST_GRAD_OPERATOR_H
	26	+#define CDIST_GRAD_OPERATOR_H
	27	+
	28	+#include "atvoss/util/vec.h"
	29	+
	30	+namespace CdistGrad {
	31	+using namespace Ops::Base;
	32	+
	33	+/**
	34	+ * \brief sign(x): 1.0 if x > 0, -1.0 if x < 0, 0.0 if x == 0
	35	+ *
	36	+ * Uses CompareScalar GT/LT to detect positive/negative, Select to choose values.
	37	+ * Replaces arithmetic approximation: diff / (\|diff\| + eps)
	38	+ */
	39	+template <typename PromoteT>
	40	+struct CdistGradSignOp : public Vec::ElemwiseUnaryOP<PromoteT, PromoteT> {
	41	+ __aicore__ inline CdistGradSignOp(LocalTensor<PromoteT>& dst, LocalTensor<PromoteT>& src, uint32_t count)
	42	+ {
	43	+#ifdef __CCE_AICORE__
	44	+ uint32_t dtypeSize = sizeof(PromoteT);
	45	+ uint32_t VL = AscendC::VECTOR_REG_WIDTH / dtypeSize;
	46	+ uint16_t loopNum = CeilDivision(count, VL);
	47	+ uint32_t vlSize = VL;
	48	+ __VEC_SCOPE__
	49	+ {
	50	+ __ubuf__ PromoteT* srcAddr = (__ubuf__ PromoteT*)src.GetPhyAddr();
	51	+ __ubuf__ PromoteT* dstAddr = (__ubuf__ PromoteT*)dst.GetPhyAddr();
	52	+
	53	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregSrc;
	54	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregDst;
	55	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregOne;
	56	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregNegOne;
	57	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregZero;
	58	+ AscendC::MicroAPI::MaskReg opMask;
	59	+ AscendC::MicroAPI::MaskReg posMask;
	60	+ AscendC::MicroAPI::MaskReg negMask;
	61	+
	62	+ AscendC::MicroAPI::Duplicate(vregOne, static_cast<PromoteT>(1.0));
	63	+ AscendC::MicroAPI::Duplicate(vregNegOne, static_cast<PromoteT>(-1.0));
	64	+ AscendC::MicroAPI::Duplicate(vregZero, static_cast<PromoteT>(0.0));
	65	+
	66	+ for (uint16_t loopIdx = 0; loopIdx < loopNum; loopIdx++) {
	67	+ opMask = AscendC::MicroAPI::UpdateMask<PromoteT, AscendC::MicroAPI::RegTraitNumOne>(count);
	68	+
	69	+ AscendC::MicroAPI::DataCopy(vregSrc, srcAddr + loopIdx * vlSize);
	70	+
	71	+ // posMask = (x > 0)
	72	+ AscendC::MicroAPI::CompareScalar<PromoteT, CMPMODE::GT>(
	73	+ posMask, vregSrc, static_cast<PromoteT>(0), opMask);
	74	+ // negMask = (x < 0)
	75	+ AscendC::MicroAPI::CompareScalar<PromoteT, CMPMODE::LT>(
	76	+ negMask, vregSrc, static_cast<PromoteT>(0), opMask);
	77	+
	78	+ // Select(dst, true_val, false_val, mask): mask=1 → true_val, mask=0 → false_val
	79	+ // Step 1: posMask=1 → 1.0, posMask=0 → 0.0
	80	+ AscendC::MicroAPI::Select(vregDst, vregOne, vregZero, posMask);
	81	+ // Step 2: negMask=1 → -1.0, negMask=0 → keep step1 result
	82	+ AscendC::MicroAPI::Select(vregDst, vregNegOne, vregDst, negMask);
	83	+
	84	+ AscendC::MicroAPI::DataCopy(dstAddr + loopIdx * vlSize, vregDst, opMask);
	85	+ }
	86	+ }
	87	+#endif
	88	+ }
	89	+};
	90	+
	91	+/**
	92	+ * \brief a == b ? 1.0 : 0.0 (binary mask: exact equality)
	93	+ *
	94	+ * Computes diff = a - b, then CompareScalar EQ(diff, 0) → mask.
	95	+ * Select(1.0, 0.0, mask) gives 1.0 where a == b, 0.0 otherwise.
	96	+ * Matches PyTorch's p=inf backward: mask = 1 - min(1, ceil(\| \|diff\| - dist \|))
	97	+ */
	98	+template <typename PromoteT>
	99	+struct CdistGradMaskEQOp : public Vec::ElemwiseBinaryOP<PromoteT, PromoteT, PromoteT> {
	100	+ __aicore__ inline CdistGradMaskEQOp(LocalTensor<PromoteT>& dst, LocalTensor<PromoteT>& src1,
	101	+ LocalTensor<PromoteT>& src2, uint32_t count)
	102	+ {
	103	+#ifdef __CCE_AICORE__
	104	+ uint32_t dtypeSize = sizeof(PromoteT);
	105	+ uint32_t VL = AscendC::VECTOR_REG_WIDTH / dtypeSize;
	106	+ uint16_t loopNum = CeilDivision(count, VL);
	107	+ uint32_t vlSize = VL;
	108	+ __VEC_SCOPE__
	109	+ {
	110	+ __ubuf__ PromoteT* src1Addr = (__ubuf__ PromoteT*)src1.GetPhyAddr();
	111	+ __ubuf__ PromoteT* src2Addr = (__ubuf__ PromoteT*)src2.GetPhyAddr();
	112	+ __ubuf__ PromoteT* dstAddr = (__ubuf__ PromoteT*)dst.GetPhyAddr();
	113	+
	114	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregA;
	115	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregB;
	116	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregDiff;
	117	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregOne;
	118	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregZero;
	119	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregResult;
	120	+ AscendC::MicroAPI::MaskReg opMask;
	121	+ AscendC::MicroAPI::MaskReg eqMask;
	122	+
	123	+ AscendC::MicroAPI::Duplicate(vregOne, static_cast<PromoteT>(1.0));
	124	+ AscendC::MicroAPI::Duplicate(vregZero, static_cast<PromoteT>(0.0));
	125	+
	126	+ for (uint16_t loopIdx = 0; loopIdx < loopNum; loopIdx++) {
	127	+ opMask = AscendC::MicroAPI::UpdateMask<PromoteT, AscendC::MicroAPI::RegTraitNumOne>(count);
	128	+
	129	+ AscendC::MicroAPI::DataCopy(vregA, src1Addr + loopIdx * vlSize);
	130	+ AscendC::MicroAPI::DataCopy(vregB, src2Addr + loopIdx * vlSize);
	131	+
	132	+ // diff = a - b
	133	+ AscendC::MicroAPI::Sub(vregDiff, vregA, vregB, opMask);
	134	+
	135	+ // eqMask = (diff == 0) i.e. a == b
	136	+ AscendC::MicroAPI::CompareScalar<PromoteT, CMPMODE::EQ>(
	137	+ eqMask, vregDiff, static_cast<PromoteT>(0), opMask);
	138	+
	139	+ // vregResult = 1.0 where eqMask=1(a==b), 0.0 where eqMask=0(a!=b)
	140	+ AscendC::MicroAPI::Select(vregResult, vregOne, vregZero, eqMask);
	141	+ AscendC::MicroAPI::DataCopy(dstAddr + loopIdx * vlSize, vregResult, opMask);
	142	+ }
	143	+ }
	144	+#endif
	145	+ }
	146	+};
	147	+
	148	+/**
	149	+ * \brief x != 0 ? 1.0 : 0.0 (non-zero mask)
	150	+ *
	151	+ * CompareScalar EQ(x, 0) → eqMask, then Select(1.0, 0.0, eqMask).
	152	+ * Where x==0: eqMask=1 → Select picks 0.0 (true branch = first arg).
	153	+ * Where x!=0: eqMask=0 → Select picks 1.0 (false branch = second arg).
	154	+ * Replaces arithmetic approximation: x / (x + eps)
	155	+ */
	156	+template <typename PromoteT>
	157	+struct CdistGradMaskNEZeroOp : public Vec::ElemwiseUnaryOP<PromoteT, PromoteT> {
	158	+ __aicore__ inline CdistGradMaskNEZeroOp(LocalTensor<PromoteT>& dst, LocalTensor<PromoteT>& src, uint32_t count)
	159	+ {
	160	+#ifdef __CCE_AICORE__
	161	+ uint32_t elemSize = sizeof(PromoteT);
	162	+ uint32_t vecLen = AscendC::VECTOR_REG_WIDTH / elemSize;
	163	+ uint16_t loopCnt = CeilDivision(count, vecLen);
	164	+ uint32_t vlLen = vecLen;
	165	+ __VEC_SCOPE__
	166	+ {
	167	+ __ubuf__ PromoteT* inputAddr = (__ubuf__ PromoteT*)src.GetPhyAddr();
	168	+ __ubuf__ PromoteT* outputAddr = (__ubuf__ PromoteT*)dst.GetPhyAddr();
	169	+
	170	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregIn;
	171	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregValOne;
	172	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregValZero;
	173	+ AscendC::MicroAPI::RegTensor<PromoteT, AscendC::MicroAPI::RegTraitNumOne> vregOut;
	174	+ AscendC::MicroAPI::MaskReg cmpMask;
	175	+ AscendC::MicroAPI::MaskReg zeroMask;
	176	+
	177	+ AscendC::MicroAPI::Duplicate(vregValOne, static_cast<PromoteT>(1.0));
	178	+ AscendC::MicroAPI::Duplicate(vregValZero, static_cast<PromoteT>(0.0));
	179	+
	180	+ for (uint16_t idx = 0; idx < loopCnt; idx++) {
	181	+ cmpMask = AscendC::MicroAPI::UpdateMask<PromoteT, AscendC::MicroAPI::RegTraitNumOne>(count);
	182	+
	183	+ AscendC::MicroAPI::DataCopy(vregIn, inputAddr + idx * vlLen);
	184	+
	185	+ // zeroMask = (x == 0)
	186	+ AscendC::MicroAPI::CompareScalar<PromoteT, CMPMODE::EQ>(
	187	+ zeroMask, vregIn, static_cast<PromoteT>(0), cmpMask);
	188	+
	189	+ // vregOut = 0.0 where zeroMask=1(x==0), 1.0 where zeroMask=0(x!=0)
	190	+ AscendC::MicroAPI::Select(vregOut, vregValZero, vregValOne, zeroMask);
	191	+ AscendC::MicroAPI::DataCopy(outputAddr + idx * vlLen, vregOut, cmpMask);
	192	+ }
	193	+ }
	194	+#endif
	195	+ }
	196	+};
	197	+
	198	+} // namespace CdistGrad
	199	+
	200	+#endif // CDIST_GRAD_OPERATOR_H

Mmath/cdist_grad/op_kernel/cdist_grad_apt.cpp+1-2

	@@ -60,8 +60,7 @@ __global__ __aicore__ void cdist_grad(
60	using OpLp = ReduceSch<REDUCE_TPL_VALUE,	60	using OpLp = ReduceSch<REDUCE_TPL_VALUE,
61	CdistGrad::CdistGradLargePDag<DTYPE_GRAD, PromoteType>::OpDag>;	61	CdistGrad::CdistGradLargePDag<DTYPE_GRAD, PromoteType>::OpDag>;
62	OpLp opLp(&tilingData.reduceTiling);	62	OpLp opLp(&tilingData.reduceTiling);
63	- opLp.template SetVar<PromoteType, 0>(static_cast<PromoteType>(tilingData.powDiff));	63	+ opLp.template SetVar<PromoteType, 0>(static_cast<PromoteType>(tilingData.powCdist));
64	- opLp.template SetVar<PromoteType, 1>(static_cast<PromoteType>(tilingData.powCdist));
65	opLp.Init(&pipe, grad, x1, x2, cdist, y, userWS);	64	opLp.Init(&pipe, grad, x1, x2, cdist, y, userWS);
66	opLp.Process(static_cast<DTYPE_GRAD>(0));	65	opLp.Process(static_cast<DTYPE_GRAD>(0));
67	} else if constexpr (normMode == CdistGrad::NORM_MODE_P0) {	66	} else if constexpr (normMode == CdistGrad::NORM_MODE_P0) {