已合并
add cdist_grad 950 #2981
chenxingyu18创建于 5月27日
add cdist_grad 950 #2981
已合并
chenxingyu18创建于 5月27日
11 个文件变更+939-23
Mmath/cdist_grad/CMakeLists.txt+12-7
@@ -1,12 +1,17 @@
1# ----------------------------------------------------------------------------1# ----------------------------------------------------------------------------
2-# This program is free software, you can redistribute it and/or modify it.2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# Copyright (c) 2025 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-# This file is a part of the CANN Open Software.4+# CANN Open Software License Agreement Version 2.0 (the "License").
5-# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
7-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8-# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
10# ----------------------------------------------------------------------------9# ----------------------------------------------------------------------------
11 10 
12-add_all_modules_sources()11+# 设置算子定义时支持的芯片类型
12+set(SUPPORT_COMPUTE_UNIT "ascend950")
13+# 设置每种芯片类型对应的tiling文件目录,即采用op_host目录下哪个文件夹下的tiling文件编译
14+set(SUPPORT_TILING_DIR "arch35")
15+add_all_modules_sources(OPTYPE cdist_grad ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR})
16+ 
17+# L0 op source is in op_api/, picked up by GLOB for the default path
Mmath/cdist_grad/op_api/cdist_grad.cpp+38-16
@@ -18,6 +18,7 @@
18#include "opdev/op_executor.h"18#include "opdev/op_executor.h"
19#include "opdev/op_log.h"19#include "opdev/op_log.h"
20#include "opdev/shape_utils.h"20#include "opdev/shape_utils.h"
21+#include "op_api/aclnn_check.h"
21 22 
22using namespace op;23using namespace op;
23 24 
@@ -29,6 +30,9 @@ OP_TYPE_REGISTER(CdistGrad);
29static const std::initializer_list<DataType> AICORE_DTYPE_SUPPORT_LIST = {30static const std::initializer_list<DataType> AICORE_DTYPE_SUPPORT_LIST = {
30 op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16};31 op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16};
31 32 
33+static const std::initializer_list<DataType> AICORE_DTYPE_SUPPORT_LIST_950 = {
34+ op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16};
35+ 
32static inline bool IsAiCoreSupport(36static inline bool IsAiCoreSupport(
33 const aclTensor* grad, const aclTensor* x1, const aclTensor* x2, const aclTensor* cdist)37 const aclTensor* grad, const aclTensor* x1, const aclTensor* x2, const aclTensor* cdist)
34{38{
@@ -36,10 +40,12 @@ static inline bool IsAiCoreSupport(
36 grad->GetDataType() != cdist->GetDataType()) {40 grad->GetDataType() != cdist->GetDataType()) {
37 return false;41 return false;
38 }42 }
39- return op::CheckType(grad->GetDataType(), AICORE_DTYPE_SUPPORT_LIST) &&43+ auto npuArch = op::GetCurrentPlatformInfo().GetCurNpuArch();
40- op::CheckType(x1->GetDataType(), AICORE_DTYPE_SUPPORT_LIST) &&44+ auto dtypeList = IsRegBase(npuArch) ? AICORE_DTYPE_SUPPORT_LIST_950 : AICORE_DTYPE_SUPPORT_LIST;
41- op::CheckType(x2->GetDataType(), AICORE_DTYPE_SUPPORT_LIST) &&45+ return op::CheckType(grad->GetDataType(), dtypeList) &&
42- op::CheckType(cdist->GetDataType(), AICORE_DTYPE_SUPPORT_LIST);46+ op::CheckType(x1->GetDataType(), dtypeList) &&
47+ op::CheckType(x2->GetDataType(), dtypeList) &&
48+ op::CheckType(cdist->GetDataType(), dtypeList);
43}49}
44 50 
45// AICORE算子kernel51// AICORE算子kernel
@@ -61,20 +67,36 @@ const aclTensor* CdistGrad(
61 const aclTensor* grad, const aclTensor* x1, const aclTensor* x2, const aclTensor* cdist, float p,67 const aclTensor* grad, const aclTensor* x1, const aclTensor* x2, const aclTensor* cdist, float p,
62 aclOpExecutor* executor)68 aclOpExecutor* executor)
63{69{
64- op::Shape outputShape;70+ L0_DFX(CdistGrad, grad, x1, x2, cdist);
65- auto dimnum = grad->GetViewShape().GetDimNum();71+ aclTensor* yOut = nullptr;
66- for (size_t i = 0; i < dimnum - NUMBER_TWO; i++) {72+ if (IsRegBase(op::GetCurrentPlatformInfo().GetCurNpuArch())) {
67- outputShape.AppendDim(grad->GetViewShape().GetDim(i));73+ // ascend950 path: reduce axis is removed from output
68- }74+ // e.g. input [P,R,M] → output [P,M], not [P,1,M]
69- outputShape.AppendDim(grad->GetViewShape().GetDim(dimnum - 1));75+ op::Shape outputShape;
70- // 根据输出shape申请输出tensor76+ auto dimnum = grad->GetViewShape().GetDimNum();
71- auto out = executor->AllocTensor(outputShape, grad->GetDataType());77+ for (size_t i = 0; i < dimnum; i++) {
72- if (out == nullptr) {78+ if (static_cast<int64_t>(i) != static_cast<int64_t>(dimnum - NUMBER_TWO)) {
73- OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "alloc out tensor failed.");79+ outputShape.AppendDim(grad->GetViewShape().GetDim(i));
74- return nullptr;80+ }
81+ }
82+ yOut = executor->AllocTensor(outputShape, grad->GetDataType(), grad->GetStorageFormat());
83+ CHECK_RET(yOut != nullptr, nullptr);
84+ } else {
85+ op::Shape outputShape;
86+ auto dimnum = grad->GetViewShape().GetDimNum();
87+ for (size_t i = 0; i < dimnum - NUMBER_TWO; i++) {
88+ outputShape.AppendDim(grad->GetViewShape().GetDim(i));
89+ }
90+ outputShape.AppendDim(grad->GetViewShape().GetDim(dimnum - 1));
91+ yOut = executor->AllocTensor(outputShape, grad->GetDataType());
92+ if (yOut == nullptr) {
93+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "alloc out tensor failed.");
94+ return nullptr;
95+ }
75 }96 }
97+ 
76 if (IsAiCoreSupport(grad, x1, x2, cdist)) {98 if (IsAiCoreSupport(grad, x1, x2, cdist)) {
77- return CdistGradAiCore(grad, x1, x2, cdist, p, out, executor);99+ return CdistGradAiCore(grad, x1, x2, cdist, p, yOut, executor);
78 } else {100 } else {
79 OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "Datatype not supported.");101 OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "Datatype not supported.");
80 return nullptr;102 return nullptr;
Amath/cdist_grad/op_host/arch35/cdist_grad_tiling_arch35.cpp+234-0
@@ -0,0 +1,234 @@
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_tiling_arch35.cpp
13+ * \brief Tiling for CdistGrad on ascend950 (arch35).
14+ *
15+ * Prototype: INPUT(grad, x1, x2, cdist) -> OUTPUT(y), ATTR(p)
16+ * normMode is encoded into TilingKey as a compile-time parameter (not in tilingData).
17+ */
18+ 
19+#include <cmath>
20+#include <vector>
21+#include "log/log.h"
22+#include "register/op_impl_registry.h"
23+#include "atvoss/reduce/reduce_tiling.h"
24+#include "atvoss/reduce/reduce_tiling_data.h"
25+#include "util/platform_util.h"
26+#include "math/cdist_grad/op_kernel/cdist_grad_tiling_data.h"
27+#include "math/cdist_grad/op_kernel/arch35/cdist_grad_dag.h"
28+#include "math/cdist_grad/op_kernel/arch35/cdist_grad_tiling_key.h"
29+ 
30+using namespace Ops::Base;
31+ 
32+namespace optiling {
33+ 
34+struct CdistGradCompileInfo {
35+ ReduceOpCompileInfo opInfo;
36+};
37+ 
38+static constexpr int32_t SIZE4 = 4;
39+static constexpr int32_t SIZE2 = 2;
40+ 
41+static ge::graphStatus DoTilingAscendC(
42+ gert::TilingContext* context, const CdistGradCompileInfo* compileInfo,
43+ ReduceOpInputParam& opInput, ReduceTilingKey& key,
44+ ReduceOpTilingData* reduceTiling, int32_t normMode)
45+{
46+ ge::graphStatus status = ge::GRAPH_FAILED;
47+ 
48+ if (ge::GetSizeByDataType(opInput.inputDtype) == SIZE4) {
49+ // float32
50+ if (normMode == CdistGrad::NORM_MODE_INF) {
51+ status = Tiling4ReduceOp<CdistGrad::CdistGradInfDag<float, float>::OpDag>(
52+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
53+ } else if (normMode == CdistGrad::NORM_MODE_LARGE_P) {
54+ status = Tiling4ReduceOp<CdistGrad::CdistGradLargePDag<float, float>::OpDag>(
55+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
56+ } else if (normMode == CdistGrad::NORM_MODE_P0) {
57+ status = Tiling4ReduceOp<CdistGrad::CdistGradP0Dag<float, float>::OpDag>(
58+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
59+ } else if (normMode == CdistGrad::NORM_MODE_P1) {
60+ status = Tiling4ReduceOp<CdistGrad::CdistGradP1Dag<float, float>::OpDag>(
61+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
62+ } else if (normMode == CdistGrad::NORM_MODE_P2) {
63+ status = Tiling4ReduceOp<CdistGrad::CdistGradP2Dag<float, float>::OpDag>(
64+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
65+ } else {
66+ status = Tiling4ReduceOp<CdistGrad::CdistGradDag<float, float>::OpDag>(
67+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
68+ }
69+ } else if (ge::GetSizeByDataType(opInput.inputDtype) == SIZE2) {
70+ // float16 / bfloat16 — both use half for tiling calc, DTYPE_GRAD handles actual type
71+ if (normMode == CdistGrad::NORM_MODE_INF) {
72+ status = Tiling4ReduceOp<CdistGrad::CdistGradInfDag<half, float>::OpDag>(
73+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
74+ } else if (normMode == CdistGrad::NORM_MODE_LARGE_P) {
75+ status = Tiling4ReduceOp<CdistGrad::CdistGradLargePDag<half, float>::OpDag>(
76+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
77+ } else if (normMode == CdistGrad::NORM_MODE_P0) {
78+ status = Tiling4ReduceOp<CdistGrad::CdistGradP0Dag<half, float>::OpDag>(
79+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
80+ } else if (normMode == CdistGrad::NORM_MODE_P1) {
81+ status = Tiling4ReduceOp<CdistGrad::CdistGradP1Dag<half, float>::OpDag>(
82+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
83+ } else if (normMode == CdistGrad::NORM_MODE_P2) {
84+ status = Tiling4ReduceOp<CdistGrad::CdistGradP2Dag<half, float>::OpDag>(
85+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
86+ } else {
87+ status = Tiling4ReduceOp<CdistGrad::CdistGradDag<half, float>::OpDag>(
88+ context, opInput, key, &compileInfo->opInfo, reduceTiling);
89+ }
90+ }
91+ 
92+ OP_CHECK_IF(
93+ (status == ge::GRAPH_FAILED),
94+ OP_LOGE(context->GetNodeName(),
95+ "ReduceOp Tiling failed, dtype should be in (float16/bfloat16/float)"),
96+ return ge::GRAPH_FAILED);
97+ return status;
98+}
99+ 
100+static int32_t ComputeNormMode(float p)
101+{
102+ if (std::isinf(p) || p == -1.0f) {
103+ return CdistGrad::NORM_MODE_INF;
104+ } else if (p == 0.0f) {
105+ return CdistGrad::NORM_MODE_P0;
106+ } else if (p == 1.0f) {
107+ return CdistGrad::NORM_MODE_P1;
108+ } else if (p == 2.0f) {
109+ return CdistGrad::NORM_MODE_P2;
110+ } else if (p > 2.0f) {
111+ return CdistGrad::NORM_MODE_LARGE_P;
112+ } else {
113+ return CdistGrad::NORM_MODE_GENERAL;
114+ }
115+}
116+ 
117+static ge::graphStatus Tiling4CdistGrad(gert::TilingContext* context)
118+{
119+ auto compileInfo = reinterpret_cast<const CdistGradCompileInfo*>(context->GetCompileInfo());
120+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
121+ 
122+ ReduceOpInputParam opInput;
123+ OP_CHECK_IF(
124+ (ReduceOpTmpl::GetInputParam(context, opInput, 0) == ge::GRAPH_FAILED),
125+ OP_LOGE(context->GetNodeName(), "ReduceOp get grad input param failed"), return ge::GRAPH_FAILED);
126+ 
127+ int64_t dimNum = static_cast<int64_t>(opInput.shape.size());
128+ OP_CHECK_IF(
129+ dimNum < 2,
130+ OP_LOGE(context->GetNodeName(), "CdistGrad requires at least 2D input, got: %ld", dimNum),
131+ return ge::GRAPH_FAILED);
132+ opInput.axes = {dimNum - 2};
133+ 
134+ // Get attr p
135+ auto attrs = context->GetAttrs();
136+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
137+ float p = 2.0f;
138+ if (attrs->GetAttrNum() > 0) {
139+ const float* pAttr = attrs->GetAttrPointer<float>(0);
140+ p = (pAttr == nullptr) ? 2.0f : *pAttr;
141+ }
142+ 
143+ // Compute normMode from p — only used for TilingKey and DAG selection, not stored in tilingData
144+ int32_t normMode = ComputeNormMode(p);
145+ 
146+ auto tilingData = context->GetTilingData<CdistGradTilingData>();
147+ OP_CHECK_NULL_WITH_CONTEXT(context, tilingData);
148+ 
149+ tilingData->powCdist = p - 1.0f;
150+ tilingData->powDiff = p - 2.0f;
151+ 
152+ OP_LOGI(context->GetNodeName(), "CdistGrad attr p = %f, normMode = %d", p, normMode);
153+ 
154+ ReduceTilingKey key;
155+ OP_CHECK_IF(
156+ (DoTilingAscendC(context, compileInfo, opInput, key,
157+ &(tilingData->reduceTiling), normMode) == ge::GRAPH_FAILED),
158+ OP_LOGE(context->GetNodeName(), "DoTiling Failed for CdistGrad"),
159+ return ge::GRAPH_FAILED);
160+ 
161+ uint64_t tilingKey;
162+ GEN_REDUCE_TILING_KEY(tilingKey, key, static_cast<uint32_t>(normMode));
163+ OP_LOGI(
164+ context->GetNodeName(),
165+ "patternID:%u, loopARCount:%u, loopInnerARCount:%u, normMode:%d, Tiling Key is:%lu",
166+ key.patternID, key.loopARCount, key.loopInnerARCount, normMode, tilingKey);
167+ context->SetTilingKey(tilingKey);
168+ return ge::GRAPH_SUCCESS;
169+}
170+ 
171+template <typename ContextT>
172+static ge::graphStatus TilingPrepare4ReduceOp(ContextT* context, ReduceOpCompileInfo* compileInfo)
173+{
174+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
175+ auto platformInfo = context->GetPlatformInfo();
176+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
177+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
178+ compileInfo->vectorCoreNum = ascendcPlatform.GetCoreNumAiv();
179+ OP_CHECK_IF(
180+ (compileInfo->vectorCoreNum == 0UL),
181+ OP_LOGE(context->GetNodeName(), "ReduceOp GetHardwareInfo Failed, vectorCoreNum:%lu",
182+ compileInfo->vectorCoreNum),
183+ return ge::GRAPH_FAILED);
184+ 
185+ uint64_t ubSize = 0;
186+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
187+ OP_CHECK_IF(
188+ ubSize <= CACHE_BUF_SIZE,
189+ OP_LOGE(context->GetNodeName(), "ReduceOp GetHardwareInfo Failed, ubSize:%lu, at least:%lu.",
190+ ubSize, CACHE_BUF_SIZE),
191+ return ge::GRAPH_FAILED);
192+ compileInfo->ubSize = ubSize;
193+ 
194+ compileInfo->cacheLineSize = GetCacheLineSize(context);
195+ OP_CHECK_IF(
196+ compileInfo->cacheLineSize == 0UL,
197+ OP_LOGE(context->GetNodeName(), "ReduceOp GetHardwareInfo Failed, cacheLineSize:%lu.",
198+ compileInfo->cacheLineSize),
199+ return ge::GRAPH_FAILED);
200+ 
201+ compileInfo->ubBlockSize = GetUbBlockSize(context);
202+ OP_CHECK_IF(
203+ compileInfo->ubBlockSize == 0UL,
204+ OP_LOGE(context->GetNodeName(), "ReduceOp GetHardwareInfo Failed, ubBlockSize:%lu.",
205+ compileInfo->ubBlockSize),
206+ return ge::GRAPH_FAILED);
207+ 
208+ compileInfo->vRegSize = GetVRegSize(context);
209+ OP_CHECK_IF(
210+ compileInfo->vRegSize == 0UL,
211+ OP_LOGE(context->GetNodeName(), "ReduceOp GetHardwareInfo Failed, vRegSize:%lu.", compileInfo->vRegSize),
212+ return ge::GRAPH_FAILED);
213+ 
214+ return ge::GRAPH_SUCCESS;
215+}
216+ 
217+template <typename T>
218+inline static T* GetCompileInfoPtr(gert::TilingParseContext* context)
219+{
220+ return context->GetCompiledInfo<T>();
221+}
222+ 
223+static ge::graphStatus TilingParse4CdistGrad(gert::TilingParseContext* context)
224+{
225+ auto compileInfo = GetCompileInfoPtr<CdistGradCompileInfo>(context);
226+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
227+ return TilingPrepare4ReduceOp(context, &compileInfo->opInfo);
228+}
229+ 
230+IMPL_OP_OPTILING(CdistGrad)
231+ .Tiling(Tiling4CdistGrad)
232+ .TilingParse<CdistGradCompileInfo>(TilingParse4CdistGrad);
233+ 
234+} // namespace optiling
Amath/cdist_grad/op_host/cdist_grad_def.cpp+62-0
@@ -0,0 +1,62 @@
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_def.cpp
13+ * \brief AICore info for CdistGrad op
14+ */
15+ 
16+#include "register/op_def_registry.h"
17+ 
18+namespace ops {
19+static const std::vector<ge::DataType> dataType = {ge::DT_FLOAT16, ge::DT_FLOAT};
20+static const std::vector<ge::Format> format = {ge::FORMAT_ND, ge::FORMAT_ND};
21+ 
22+class CdistGrad : public OpDef {
23+public:
24+ explicit CdistGrad(const char* name) : OpDef(name)
25+ {
26+ this->Input("grad")
27+ .ParamType(REQUIRED)
28+ .DataType(dataType)
29+ .Format(format)
30+ .UnknownShapeFormat(format);
31+ this->Input("x1")
32+ .ParamType(REQUIRED)
33+ .DataType(dataType)
34+ .Format(format)
35+ .UnknownShapeFormat(format);
36+ this->Input("x2")
37+ .ParamType(REQUIRED)
38+ .DataType(dataType)
39+ .Format(format)
40+ .UnknownShapeFormat(format);
41+ this->Input("cdist")
42+ .ParamType(REQUIRED)
43+ .DataType(dataType)
44+ .Format(format)
45+ .UnknownShapeFormat(format);
46+ this->Output("y")
47+ .ParamType(REQUIRED)
48+ .DataType(dataType)
49+ .Format(format)
50+ .UnknownShapeFormat(format);
51+ this->Attr("p").AttrType(OPTIONAL).Float(2.0);
52+ 
53+ OpAICoreConfig aicoreConfig;
54+ aicoreConfig.DynamicCompileStaticFlag(true)
55+ .DynamicRankSupportFlag(true)
56+ .DynamicShapeSupportFlag(true)
57+ .ExtendCfgInfo("opFile.value", "cdist_grad_apt");
58+ this->AICore().AddConfig("ascend950", aicoreConfig);
59+ }
60+};
61+OP_ADD(CdistGrad);
62+} // namespace ops
Amath/cdist_grad/op_host/cdist_grad_infershape.cpp+60-0
@@ -0,0 +1,60 @@
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_infershape.cpp
13+ * \brief InferShape for CdistGrad operator
14+ *
15+ * CdistGrad reduces along axis -2 (without keepdims).
16+ * Input shape: [..., P, R, M] (after broadcast)
17+ * Output shape: [..., P, M] (R dimension is reduced away)
18+ */
19+ 
20+#include "log/log.h"
21+#include "register/op_impl_registry.h"
22+ 
23+using namespace ge;
24+ 
25+namespace ops {
26+constexpr int64_t NUMBER_TWO = 2;
27+ 
28+static ge::graphStatus InferShape4CdistGrad(gert::InferShapeContext* context)
29+{
30+ OP_LOGD(context->GetNodeName(), "Begin to do InferShape4CdistGrad");
31+ const gert::Shape* gradShape = context->GetInputShape(0);
32+ OP_CHECK_NULL_WITH_CONTEXT(context, gradShape);
33+ const gert::Shape* x1Shape = context->GetInputShape(1);
34+ OP_CHECK_NULL_WITH_CONTEXT(context, x1Shape);
35+ const gert::Shape* cdistShape = context->GetInputShape(3);
36+ OP_CHECK_NULL_WITH_CONTEXT(context, cdistShape);
37+ 
38+ int64_t dimNum = gradShape->GetDimNum();
39+ OP_CHECK_IF(dimNum < NUMBER_TWO,
40+ OP_LOGE(context->GetNodeName(),
41+ "Only supports at least 2D tensors, got: %ld.", dimNum),
42+ return ge::GRAPH_FAILED);
43+ 
44+ gert::Shape* yShape = context->GetOutputShape(0);
45+ OP_CHECK_NULL_WITH_CONTEXT(context, yShape);
46+ 
47+ // CdistGrad reduces on axis -2: output has dimNum - 1 dimensions
48+ // Output dims: grad_shape[0..dimNum-3] + grad_shape[dimNum-1]
49+ yShape->SetDimNum(dimNum - 1);
50+ for (int64_t i = 0; i < dimNum - NUMBER_TWO; i++) {
51+ yShape->SetDim(i, gradShape->GetDim(i));
52+ }
53+ yShape->SetDim(dimNum - NUMBER_TWO, gradShape->GetDim(dimNum - 1));
54+ 
55+ OP_LOGD(context->GetNodeName(), "End to do InferShape4CdistGrad");
56+ return GRAPH_SUCCESS;
57+}
58+ 
59+IMPL_OP_INFERSHAPE(CdistGrad).InferShape(InferShape4CdistGrad);
60+} // namespace ops
Amath/cdist_grad/op_host/config/ascend950/cdist_grad_binary.json+50-0
@@ -0,0 +1,50 @@
1+{
2+ "op_type": "CdistGrad",
3+ "op_list": [
4+ {
5+ "bin_filename": "CdistGrad_6d3844989f83c1de15cb106d07138083",
6+ "inputs": [
7+ { "name": "grad", "index": 0, "dtype": "float16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
8+ { "name": "x1", "index": 1, "dtype": "float16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
9+ { "name": "x2", "index": 2, "dtype": "float16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
10+ { "name": "cdist", "index": 3, "dtype": "float16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" }
11+ ],
12+ "outputs": [
13+ { "name": "y", "index": 0, "dtype": "float16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" }
14+ ],
15+ "attrs": [
16+ { "name": "p", "dtype": "float", "value": null }
17+ ]
18+ },
19+ {
20+ "bin_filename": "CdistGrad_876161686b0bc0b329614d9b9c3b5adc",
21+ "inputs": [
22+ { "name": "grad", "index": 0, "dtype": "bfloat16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
23+ { "name": "x1", "index": 1, "dtype": "bfloat16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
24+ { "name": "x2", "index": 2, "dtype": "bfloat16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
25+ { "name": "cdist", "index": 3, "dtype": "bfloat16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" }
26+ ],
27+ "outputs": [
28+ { "name": "y", "index": 0, "dtype": "bfloat16", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" }
29+ ],
30+ "attrs": [
31+ { "name": "p", "dtype": "float", "value": null }
32+ ]
33+ },
34+ {
35+ "bin_filename": "CdistGrad_632f457d4787205a3df995bf22221bbe",
36+ "inputs": [
37+ { "name": "grad", "index": 0, "dtype": "float32", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
38+ { "name": "x1", "index": 1, "dtype": "float32", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
39+ { "name": "x2", "index": 2, "dtype": "float32", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" },
40+ { "name": "cdist", "index": 3, "dtype": "float32", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" }
41+ ],
42+ "outputs": [
43+ { "name": "y", "index": 0, "dtype": "float32", "format": "ND", "paramType": "required", "shape": [-2], "format_match_mode": "FormatAgnostic" }
44+ ],
45+ "attrs": [
46+ { "name": "p", "dtype": "float", "value": null }
47+ ]
48+ }
49+ ]
50+}
Amath/cdist_grad/op_host/config/ascend950/cdist_grad_simplified_key.ini+2-0
@@ -0,0 +1,2 @@
1+[CdistGrad]
2+default=0
Amath/cdist_grad/op_kernel/arch35/cdist_grad_dag.h+323-0
@@ -0,0 +1,323 @@
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_dag.h
13+ * \brief cdist_grad dag — six DAGs, all Compare+Select replaced by arithmetic
14+ *
15+ * Arithmetic replacements:
16+ * sign: diff / (|diff| + eps) replaces Compare(NE) + Select
17+ * mask: (d + |d| + eps) / (2|d| + eps) replaces Compare(GE/EQ) + Select
18+ * nz_x: x + eps replaces Compare(NE) + Select for safe divisor
19+ * mask_nz: x / (x + eps) replaces Compare(NE) + Select for zero-out
20+ *
21+ * CdistGradP0Dag: p == 0 → output zeros
22+ * CdistGradP1Dag: p == 1 → grad * sign
23+ * CdistGradP2Dag: p == 2 → grad * diff / cdist
24+ * CdistGradDag: 0<p<2 → sign * |diff|^(p-1) * grad / |cdist|^(p-1)
25+ * CdistGradLargePDag: p>2 → diff * |diff|^(p-2) * grad / |cdist|^(p-1)
26+ * CdistGradInfDag: p==inf → grad * sign * mask(|diff| >= cdist)
27+ */
28+ 
29+#ifndef CDIST_GRAD_DAG_H
30+#define CDIST_GRAD_DAG_H
31+ 
32+#include "atvoss/util/elems.h"
33+#include "atvoss/util/dag.h"
34+#include "atvoss/util/vec.h"
35+#include "atvoss/util/placeholder.h"
36+#include "atvoss/reduce/reduce_operator.h"
37+ 
38+namespace CdistGrad {
39+using namespace Ops::Base;
40+ 
41+constexpr int CAST_MODE_NONE = 0;
42+constexpr int CAST_MODE_RINT = 1;
43+ 
44+constexpr int32_t NORM_MODE_GENERAL = 0; // 0 < p < 2, p != 1
45+constexpr int32_t NORM_MODE_INF = 1; // p == inf
46+constexpr int32_t NORM_MODE_LARGE_P = 2; // p > 2
47+constexpr int32_t NORM_MODE_P0 = 3; // p == 0
48+constexpr int32_t NORM_MODE_P1 = 4; // p == 1
49+constexpr int32_t NORM_MODE_P2 = 5; // p == 2
50+ 
51+// ---------------------------------------------------------------------------
52+// CdistGradP0Dag — p == 0 → output zeros
53+// Var<0>: zero_scalar = 0.0
54+// ---------------------------------------------------------------------------
55+template <typename T, typename PromoteT>
56+struct CdistGradP0Dag {
57+ using OpCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
58+ using CastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInGrad>;
59+ 
60+ using OpZero = Bind<Vec::Muls<PromoteT>, CastGrad, Placeholder::Var<PromoteT, 0>>;
61+ using ReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, OpZero>;
62+ using CastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, ReduceOp0>;
63+ using OpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, CastOut>;
64+ 
65+ using Outputs = Elems<OpCopyOut>;
66+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
67+ using OpDag = DAGSch<Outputs, void, MemCfg>;
68+};
69+ 
70+// ---------------------------------------------------------------------------
71+// CdistGradP1Dag — p == 1
72+// sign = diff / (|diff| + eps) replaces Compare(NE)+Select for nz_diff
73+// result = grad * sign
74+// ---------------------------------------------------------------------------
75+template <typename T, typename PromoteT>
76+struct CdistGradP1Dag {
77+ using Eps = MAKE_CONST(PromoteT, 1e-30);
78+ 
79+ using OpCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
80+ using CastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInGrad>;
81+ 
82+ using OpCopyInX1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
83+ using CastX1 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInX1>;
84+ 
85+ using OpCopyInX2 = Bind<Vec::CopyIn<T>, Placeholder::In2<T>>;
86+ using CastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInX2>;
87+ 
88+ using OpDiff = Bind<Vec::Sub<PromoteT>, CastX1, CastX2>;
89+ using OpDiffAbs = Bind<Vec::Abs<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+ 
93+ using OpRes = Bind<Vec::Mul<PromoteT>, CastGrad, OpSign>;
94+ 
95+ using ReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, OpRes>;
96+ using CastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, ReduceOp0>;
97+ using OpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, CastOut>;
98+ 
99+ using Outputs = Elems<OpCopyOut>;
100+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
101+ using OpDag = DAGSch<Outputs, void, MemCfg>;
102+};
103+ 
104+// ---------------------------------------------------------------------------
105+// CdistGradP2Dag — p == 2
106+// result = grad * diff / (cdist + eps)
107+// When cdist = 0: diff = 0, numerator = 0, result = 0 automatically
108+// ---------------------------------------------------------------------------
109+template <typename T, typename PromoteT>
110+struct CdistGradP2Dag {
111+ using Eps = MAKE_CONST(PromoteT, 1e-30);
112+ 
113+ using OpCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
114+ using CastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInGrad>;
115+ 
116+ using OpCopyInX1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
117+ using CastX1 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInX1>;
118+ 
119+ using OpCopyInX2 = Bind<Vec::CopyIn<T>, Placeholder::In2<T>>;
120+ using CastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInX2>;
121+ 
122+ using OpCopyInCdist = Bind<Vec::CopyIn<T>, Placeholder::In3<T>>;
123+ using CastCdist = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, OpCopyInCdist>;
124+ 
125+ using OpDiff = Bind<Vec::Sub<PromoteT>, CastX1, CastX2>;
126+ using SafeCdist = Bind<Vec::Adds<PromoteT>, CastCdist, Eps>; // cdist + eps
127+ using OpNumerator = Bind<Vec::Mul<PromoteT>, CastGrad, OpDiff>; // grad * diff
128+ using OpResult = Bind<Vec::Div<PromoteT>, OpNumerator, SafeCdist>; // grad * diff / (cdist+eps)
129+ 
130+ using ReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, OpResult>;
131+ using CastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, ReduceOp0>;
132+ using OpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, CastOut>;
133+ 
134+ using Outputs = Elems<OpCopyOut>;
135+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
136+ using OpDag = DAGSch<Outputs, void, MemCfg>;
137+};
138+ 
139+// ---------------------------------------------------------------------------
140+// CdistGradDag — 0 < p < 2, p != 1
141+// sign = diff / (|diff| + eps)
142+// safe_diff = |diff| + eps (prevents log(0))
143+// safe_cdist = cdist + eps (prevents log(0) and div-by-0)
144+// mask_diff = |diff| / (|diff| + eps) (zero-out when |diff|=0)
145+// mask_cdist = cdist / (cdist + eps) (zero-out when cdist=0)
146+// num = sign * exp(log(safe_diff) * (p-1))
147+// numerator = num * grad
148+// denominator = exp(log(safe_cdist) * (p-1))
149+// result = numerator / denominator * mask_cdist * mask_diff
150+// Var<0>: power = p - 1
151+// ---------------------------------------------------------------------------
152+template <typename T, typename PromoteT>
153+struct CdistGradDag {
154+ using DagEps = MAKE_CONST(PromoteT, 1e-30);
155+ 
156+ using DagCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
157+ using DagCastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, DagCopyInGrad>;
158+ 
159+ using DagCopyInX1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
160+ using DagCastX1 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, DagCopyInX1>;
161+ 
162+ using DagCopyInX2 = Bind<Vec::CopyIn<T>, Placeholder::In2<T>>;
163+ using DagCastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, DagCopyInX2>;
164+ 
165+ using DagCopyInCdist = Bind<Vec::CopyIn<T>, Placeholder::In3<T>>;
166+ using DagCastCdist = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, DagCopyInCdist>;
167+ 
168+ // diff & abs
169+ using DagDiff = Bind<Vec::Sub<PromoteT>, DagCastX1, DagCastX2>;
170+ 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+ 
174+ // sign = diff / (|diff| + eps)
175+ using DagSign = Bind<Vec::Div<PromoteT>, DagDiff, DagSafeAbsDiff>;
176+ 
177+ // masks: 0 when input=0, ~1 otherwise
178+ using DagMaskDiff = Bind<Vec::Div<PromoteT>, DagDiffAbs, DagSafeAbsDiff>; // |diff|/(|diff|+eps)
179+ using DagMaskCdist = Bind<Vec::Div<PromoteT>, DagCastCdist, DagSafeCdist>; // cdist/(cdist+eps)
180+ 
181+ // power: (safe_x)^(p-1) via log/exp, Var<0> = p-1
182+ using DagPowDiff = Bind<Vec::Exp<PromoteT>,
183+ Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, DagSafeAbsDiff>,
184+ Placeholder::Var<PromoteT, 0>>>;
185+ using DagPowCdist = Bind<Vec::Exp<PromoteT>,
186+ Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, DagSafeCdist>,
187+ Placeholder::Var<PromoteT, 0>>>;
188+ 
189+ // num = sign * |diff|^(p-1), numerator = num * grad
190+ using DagNum = Bind<Vec::Mul<PromoteT>, DagSign, DagPowDiff>;
191+ using DagNumerator = Bind<Vec::Mul<PromoteT>, DagNum, DagCastGrad>;
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
197+ using DagMaskedCdist = Bind<Vec::Mul<PromoteT>, DagDivResult, DagMaskCdist>;
198+ using DagResult = Bind<Vec::Mul<PromoteT>, DagMaskedCdist, DagMaskDiff>;
199+ 
200+ using DagReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, DagResult>;
201+ using DagCastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, DagReduceOp0>;
202+ using DagCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, DagCastOut>;
203+ 
204+ using Outputs = Elems<DagCopyOut>;
205+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
206+ using OpDag = DAGSch<Outputs, void, MemCfg>;
207+};
208+ 
209+// ---------------------------------------------------------------------------
210+// CdistGradLargePDag — p > 2
211+// safe_cdist = cdist + eps
212+// mask_cdist = cdist / (cdist + eps)
213+// num = diff * |diff|^(p-2) (0 when |diff|=0 for p>2)
214+// numerator = num * grad
215+// denominator = |cdist|^(p-1)
216+// result = numerator / denominator * mask_cdist
217+// Var<0>: power_diff = p - 2
218+// Var<1>: power_cdist = p - 1
219+// ---------------------------------------------------------------------------
220+template <typename T, typename PromoteT>
221+struct CdistGradLargePDag {
222+ using LpEps = MAKE_CONST(PromoteT, 1e-30);
223+ 
224+ using LpCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
225+ using LpCastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, LpCopyInGrad>;
226+ 
227+ using LpCopyInX1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
228+ using LpCastX1 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, LpCopyInX1>;
229+ 
230+ using LpCopyInX2 = Bind<Vec::CopyIn<T>, Placeholder::In2<T>>;
231+ using LpCastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, LpCopyInX2>;
232+ 
233+ using LpCopyInCdist = Bind<Vec::CopyIn<T>, Placeholder::In3<T>>;
234+ using LpCastCdist = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, LpCopyInCdist>;
235+ 
236+ using LpDiff = Bind<Vec::Sub<PromoteT>, LpCastX1, LpCastX2>;
237+ 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
240+ using LpMaskCdist = Bind<Vec::Div<PromoteT>, LpCastCdist, LpSafeCdist>; // cdist/(cdist+eps)
241+ 
242+ // |diff|^(p-2), Var<0> = p-2
243+ using LpPowDiff = Bind<Vec::Exp<PromoteT>,
244+ Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, LpSafeAbsDiff>,
245+ Placeholder::Var<PromoteT, 0>>>;
246+ 
247+ // |cdist|^(p-1), Var<1> = p-1
248+ using LpPowCdist = Bind<Vec::Exp<PromoteT>,
249+ Bind<Vec::Muls<PromoteT>, Bind<Vec::Log<PromoteT>, LpSafeCdist>,
250+ Placeholder::Var<PromoteT, 1>>>;
251+ 
252+ // num = diff * |diff|^(p-2)
253+ using LpNum = Bind<Vec::Mul<PromoteT>, LpDiff, LpPowDiff>;
254+ // numerator = num * grad
255+ using LpNumerator = Bind<Vec::Mul<PromoteT>, LpNum, LpCastGrad>;
256+ // res = numerator / |cdist|^(p-1)
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+ 
261+ using LpReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, LpResult>;
262+ using LpCastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, LpReduceOp0>;
263+ using LpCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, LpCastOut>;
264+ 
265+ using Outputs = Elems<LpCopyOut>;
266+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
267+ using OpDag = DAGSch<Outputs, void, MemCfg>;
268+};
269+ 
270+// ---------------------------------------------------------------------------
271+// CdistGradInfDag — p == inf
272+// sign = diff / (|diff| + eps)
273+// mask = (d + |d| + eps) / (2|d| + eps) where d = |diff| - cdist
274+// result = grad * sign * mask
275+// ---------------------------------------------------------------------------
276+template <typename T, typename PromoteT>
277+struct CdistGradInfDag {
278+ using InfEps = MAKE_CONST(PromoteT, 1e-30);
279+ 
280+ using InfCopyInGrad = Bind<Vec::CopyIn<T>, Placeholder::In0<T>>;
281+ using InfCastGrad = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, InfCopyInGrad>;
282+ 
283+ using InfCopyInX1 = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
284+ using InfCastX1 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, InfCopyInX1>;
285+ 
286+ using InfCopyInX2 = Bind<Vec::CopyIn<T>, Placeholder::In2<T>>;
287+ using InfCastX2 = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, InfCopyInX2>;
288+ 
289+ using InfCopyInCdist = Bind<Vec::CopyIn<T>, Placeholder::In3<T>>;
290+ using InfCastCdist = Bind<Vec::Cast<PromoteT, T, CAST_MODE_NONE>, InfCopyInCdist>;
291+ 
292+ using InfDiff = Bind<Vec::Sub<PromoteT>, InfCastX1, InfCastX2>;
293+ using InfDiffAbs = Bind<Vec::Abs<PromoteT>, InfDiff>;
294+ 
295+ // sign = diff / (|diff| + eps)
296+ using InfSafeAbsDiff = Bind<Vec::Adds<PromoteT>, InfDiffAbs, InfEps>;
297+ using InfSign = Bind<Vec::Div<PromoteT>, InfDiff, InfSafeAbsDiff>;
298+ 
299+ // mask: d = |diff| - cdist, mask = (d + |d| + eps) / (2|d| + eps)
300+ using InfD = Bind<Vec::Sub<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+ 
308+ // result = grad * sign * mask
309+ using InfGradSign = Bind<Vec::Mul<PromoteT>, InfCastGrad, InfSign>;
310+ using InfResult = Bind<Vec::Mul<PromoteT>, InfGradSign, InfMask>;
311+ 
312+ using InfReduceOp0 = Bind<Vec::ReduceSumOp<PromoteT>, InfResult>;
313+ using InfCastOut = Bind<Vec::Cast<T, PromoteT, CAST_MODE_RINT>, InfReduceOp0>;
314+ using InfCopyOut = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, InfCastOut>;
315+ 
316+ using Outputs = Elems<InfCopyOut>;
317+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
318+ using OpDag = DAGSch<Outputs, void, MemCfg>;
319+};
320+ 
321+} // namespace CdistGrad
322+ 
323+#endif // CDIST_GRAD_DAG_H
Amath/cdist_grad/op_kernel/arch35/cdist_grad_tiling_key.h+39-0
@@ -0,0 +1,39 @@
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_tiling_key.h
13+ * \brief cdist_grad tiling key declare — normMode encoded as compile-time param
14+ */
15+ 
16+#ifndef _CDIST_GRAD_TILING_KEY_H_
17+#define _CDIST_GRAD_TILING_KEY_H_
18+ 
19+#include "atvoss/reduce/reduce_tiling_key_decl.h"
20+ 
21+ASCENDC_TPL_ARGS_DECL(CdistGrad, REDUCE_TPL_KEY_DECL(),
22+ ASCENDC_TPL_UINT_DECL(normMode, 8, ASCENDC_TPL_UI_RANGE, 1, 0, 5));
23+ 
24+ASCENDC_TPL_SEL(
25+ ASCENDC_TPL_ARGS_SEL(REDUCE_TPL_KEY_SEL_EMPTY(),
26+ ASCENDC_TPL_UINT_SEL(normMode, ASCENDC_TPL_UI_RANGE, 1, 0, 5)),
27+ ASCENDC_TPL_ARGS_SEL(REDUCE_TPL_KEY_SEL_A(),
28+ ASCENDC_TPL_UINT_SEL(normMode, ASCENDC_TPL_UI_RANGE, 1, 0, 5)),
29+ ASCENDC_TPL_ARGS_SEL(REDUCE_TPL_KEY_SEL_AR_NORMAL(),
30+ ASCENDC_TPL_UINT_SEL(normMode, ASCENDC_TPL_UI_RANGE, 1, 0, 5)),
31+ ASCENDC_TPL_ARGS_SEL(REDUCE_TPL_KEY_SEL_AR_GROUP(),
32+ ASCENDC_TPL_UINT_SEL(normMode, ASCENDC_TPL_UI_RANGE, 1, 0, 5)),
33+ ASCENDC_TPL_ARGS_SEL(REDUCE_TPL_KEY_SEL_ARA_NORMAL(),
34+ ASCENDC_TPL_UINT_SEL(normMode, ASCENDC_TPL_UI_RANGE, 1, 0, 5)),
35+ ASCENDC_TPL_ARGS_SEL(REDUCE_TPL_KEY_SEL_ARA_GROUP(),
36+ ASCENDC_TPL_UINT_SEL(normMode, ASCENDC_TPL_UI_RANGE, 1, 0, 5))
37+);
38+ 
39+#endif // _CDIST_GRAD_TILING_KEY_H_
Amath/cdist_grad/op_kernel/cdist_grad_apt.cpp+92-0
@@ -0,0 +1,92 @@
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_apt.cpp
13+ * \brief CdistGrad kernel entry (arch35)
14+ *
15+ * normMode is a compile-time template parameter — each binary contains only ONE DAG.
16+ * TilingKey encodes normMode so the framework loads the correct binary at runtime.
17+ */
18+ 
19+#include "atvoss/reduce/reduce_sch.h"
20+#include "arch35/cdist_grad_dag.h"
21+#include "arch35/cdist_grad_tiling_key.h"
22+#include "cdist_grad_tiling_data.h"
23+ 
24+using namespace Ops::Base::ReduceOpTmpl;
25+using namespace AscendC;
26+ 
27+#define CDIST_GRAD_LAUNCH(DagT, ...) \
28+ using _CGOp = ReduceSch<REDUCE_TPL_VALUE, DagT::OpDag>; \
29+ _CGOp _cgOp(&tilingData.reduceTiling); \
30+ _cgOp.Init(&pipe, __VA_ARGS__); \
31+ _cgOp.Process(static_cast<DTYPE_GRAD>(0))
32+ 
33+template <REDUCE_TPL_PARAM, int32_t normMode>
34+__global__ __aicore__ void cdist_grad(
35+ GM_ADDR grad, GM_ADDR x1, GM_ADDR x2, GM_ADDR cdist,
36+ GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
37+{
38+ if (g_coreType == AIC) {
39+ return;
40+ }
41+ if (workspace == nullptr) {
42+ return;
43+ }
44+ SetSysWorkspace(workspace);
45+ GM_ADDR userWS = GetUserWorkspace(workspace);
46+ if (userWS == nullptr) {
47+ return;
48+ }
49+ 
50+ REGISTER_TILING_DEFAULT(CdistGradTilingData);
51+ GET_TILING_DATA_WITH_STRUCT(CdistGradTilingData, tilingData, tiling);
52+ 
53+ TPipe pipe;
54+ using PromoteType = __reduceType::GetPromoteType<DTYPE_GRAD>::T;
55+ 
56+ if constexpr (normMode == CdistGrad::NORM_MODE_INF) {
57+ using Dag = CdistGrad::CdistGradInfDag<DTYPE_GRAD, PromoteType>;
58+ CDIST_GRAD_LAUNCH(Dag, grad, x1, x2, cdist, y, userWS);
59+ } else if constexpr (normMode == CdistGrad::NORM_MODE_LARGE_P) {
60+ using OpLp = ReduceSch<REDUCE_TPL_VALUE,
61+ CdistGrad::CdistGradLargePDag<DTYPE_GRAD, PromoteType>::OpDag>;
62+ OpLp opLp(&tilingData.reduceTiling);
63+ opLp.template SetVar<PromoteType, 0>(static_cast<PromoteType>(tilingData.powDiff));
64+ opLp.template SetVar<PromoteType, 1>(static_cast<PromoteType>(tilingData.powCdist));
65+ opLp.Init(&pipe, grad, x1, x2, cdist, y, userWS);
66+ opLp.Process(static_cast<DTYPE_GRAD>(0));
67+ } else if constexpr (normMode == CdistGrad::NORM_MODE_P0) {
68+ using OpP0 = ReduceSch<REDUCE_TPL_VALUE,
69+ CdistGrad::CdistGradP0Dag<DTYPE_GRAD, PromoteType>::OpDag>;
70+ OpP0 opP0(&tilingData.reduceTiling);
71+ opP0.template SetVar<PromoteType, 0>(static_cast<PromoteType>(0.0f));
72+ // P0 DAG uses only In0, Out0 — no In1/In2/In3
73+ opP0.Init(&pipe, grad, y, userWS);
74+ opP0.Process(static_cast<DTYPE_GRAD>(0));
75+ } else if constexpr (normMode == CdistGrad::NORM_MODE_P1) {
76+ using Dag = CdistGrad::CdistGradP1Dag<DTYPE_GRAD, PromoteType>;
77+ CDIST_GRAD_LAUNCH(Dag, grad, x1, x2, y, userWS);
78+ } else if constexpr (normMode == CdistGrad::NORM_MODE_P2) {
79+ using Dag = CdistGrad::CdistGradP2Dag<DTYPE_GRAD, PromoteType>;
80+ CDIST_GRAD_LAUNCH(Dag, grad, x1, x2, cdist, y, userWS);
81+ } else {
82+ // NORM_MODE_GENERAL: 0 < p < 2, p != 1
83+ using Op = ReduceSch<REDUCE_TPL_VALUE,
84+ CdistGrad::CdistGradDag<DTYPE_GRAD, PromoteType>::OpDag>;
85+ Op op(&tilingData.reduceTiling);
86+ op.template SetVar<PromoteType, 0>(static_cast<PromoteType>(tilingData.powCdist));
87+ op.Init(&pipe, grad, x1, x2, cdist, y, userWS);
88+ op.Process(static_cast<DTYPE_GRAD>(0));
89+ }
90+}
91+ 
92+#undef CDIST_GRAD_LAUNCH
Amath/cdist_grad/op_kernel/cdist_grad_tiling_data.h+27-0
@@ -0,0 +1,27 @@
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_tiling_data.h
13+ * \brief Tiling data for CdistGrad operator.
14+ */
15+ 
16+#ifndef CDIST_GRAD_TILING_DATA_H
17+#define CDIST_GRAD_TILING_DATA_H
18+ 
19+#include "atvoss/reduce/reduce_tiling_data.h"
20+ 
21+struct CdistGradTilingData {
22+ Ops::Base::ReduceOpTilingData reduceTiling;
23+ float powCdist; // p - 1
24+ float powDiff; // p - 2 (used by CdistGradLargePDag)
25+};
26+ 
27+#endif // CDIST_GRAD_TILING_DATA_H