Star945
1861
代码
代码
Issues223
Pull Requests333
Actions
项目讨论
Wiki
分析
Star945
1861
已合并
bn_training_reduce算子支持NCDHW #9089
yulianjie创建于 21 天前
bn_training_reduce算子支持NCDHW #9089
已合并
yulianjie创建于 21 天前
从lianjieyu/ops-nnbn_training_reduce合入到cann/ops-nnmaster
共 13 个文件变更+286-94
Mnorm/bn_training_reduce/README.md+8-7
@@ -13,7 +13,8 @@
13 13 
14## 功能说明14## 功能说明
15 15 
16-- 算子功能：对四维NCHW或NHWC输入的N、H、W轴执行归约，输出每个通道的元素和与平方和。16+- 算子功能：保留输入的C轴并归约其余所有轴，输出每个通道的元素和与平方和。Ascend 950 支持
17+ NCHW 2~4维、NHWC 4维和NCDHW 5维。
17- 计算公式：18- 计算公式：
18 19 
19 $$20 $$
@@ -45,30 +46,30 @@
45 <tr>46 <tr>
46 <td>x</td>47 <td>x</td>
47 <td>输入</td>48 <td>输入</td>
48- <td>待统计的四维训练激活。</td>49+ <td>待统计的训练激活。</td>
49 <td>FLOAT16、BFLOAT16、FLOAT</td>50 <td>FLOAT16、BFLOAT16、FLOAT</td>
50- <td>NCHW、NHWC</td>51+ <td>NCHW、NHWC、NCDHW</td>
51 </tr>52 </tr>
52 <tr>53 <tr>
53 <td>sum</td>54 <td>sum</td>
54 <td>输出</td>55 <td>输出</td>
55 <td>沿N、H、W轴归约得到的每通道元素和。</td>56 <td>沿N、H、W轴归约得到的每通道元素和。</td>
56 <td>FLOAT</td>57 <td>FLOAT</td>
57- <td>ND</td>58+ <td>与x相同</td>
58 </tr>59 </tr>
59 <tr>60 <tr>
60 <td>square_sum</td>61 <td>square_sum</td>
61 <td>输出</td>62 <td>输出</td>
62 <td>沿N、H、W轴归约得到的每通道平方和。</td>63 <td>沿N、H、W轴归约得到的每通道平方和。</td>
63 <td>FLOAT</td>64 <td>FLOAT</td>
64- <td>ND</td>65+ <td>与x相同</td>
65 </tr>66 </tr>
66</tbody></table>67</tbody></table>
67 68 
68## 约束说明69## 约束说明
69 70 
70-- x必须为四维NCHW或NHWC张量，通道轴分别为第1维或第3维。71+- Ascend 950 支持NCHW 2~4维、NHWC 4维和NCDHW 5维，不支持NDC1HWC0；输出格式与x相同。
71-- sum和square_sum必须为一维ND张量，长度等于x的C维，数据类型固定为FLOAT。72+- sum和square_sum的逻辑shape均为一维[C]，长度等于x的C维，数据类型固定为FLOAT。
72- FLOAT16和BFLOAT16输入按FLOAT精度执行平方与累加。73- FLOAT16和BFLOAT16输入按FLOAT精度执行平方与累加。
73- 算子无属性，支持空Tensor；归约集合为空时输出为零或空向量。74- 算子无属性，支持空Tensor；归约集合为空时输出为零或空向量。
74 75 
Mnorm/bn_training_reduce/examples/arch35/test_geir_bn_training_reduce.cpp+14-2
@@ -214,7 +214,16 @@ const char* DtypeName(ge::DataType dtype)
214 }214 }
215}215}
216 216 
217-ge::Format ParseFormat(const std::string& format) { return format == "NHWC" ? ge::FORMAT_NHWC : ge::FORMAT_NCHW; }217+ge::Format ParseFormat(const std::string& format)
218+{
219+ if (format == "NHWC") {
220+ return ge::FORMAT_NHWC;
221+ }
222+ if (format == "NCDHW") {
223+ return ge::FORMAT_NCDHW;
224+ }
225+ return ge::FORMAT_NCHW;
226+}
218 227 
219const char* FormatName(ge::Format format)228const char* FormatName(ge::Format format)
220{229{
@@ -227,6 +236,9 @@ const char* FormatName(ge::Format format)
227 if (format == ge::FORMAT_NHWC) {236 if (format == ge::FORMAT_NHWC) {
228 return "NHWC";237 return "NHWC";
229 }238 }
239+ if (format == ge::FORMAT_NCDHW) {
240+ return "NCDHW";
241+ }
230 return "OTHER";242 return "OTHER";
231}243}
232 244 
@@ -570,7 +582,7 @@ GraphBundle BuildGraph(const CaseDef& test)
570 bundle.node = test.id == "route-target" ? "bn_training_reduce_route" : "bn_training_reduce_" + test.id;582 bundle.node = test.id == "route-target" ? "bn_training_reduce_route" : "bn_training_reduce_" + test.id;
571 bundle.graph = ge::Graph((bundle.node + "_graph").c_str());583 bundle.graph = ge::Graph((bundle.node + "_graph").c_str());
572 auto reduce = ge::op::BNTrainingReduce(bundle.node.c_str());584 auto reduce = ge::op::BNTrainingReduce(bundle.node.c_str());
573- const ge::TensorDesc outDesc(ge::Shape(ge::UNKNOWN_RANK), ge::FORMAT_ND, ge::DT_FLOAT);585+ const ge::TensorDesc outDesc(ge::Shape(ge::UNKNOWN_RANK), ParseFormat(test.inputFormat), ge::DT_FLOAT);
574 reduce.update_output_desc_sum(outDesc);586 reduce.update_output_desc_sum(outDesc);
575 reduce.update_output_desc_square_sum(outDesc);587 reduce.update_output_desc_square_sum(outDesc);
576 588 
Mnorm/bn_training_reduce/examples/test_geir_bn_training_reduce.cpp+14-2
@@ -214,7 +214,16 @@ const char* DtypeName(ge::DataType dtype)
214 }214 }
215}215}
216 216 
217-ge::Format ParseFormat(const std::string& format) { return format == "NHWC" ? ge::FORMAT_NHWC : ge::FORMAT_NCHW; }217+ge::Format ParseFormat(const std::string& format)
218+{
219+ if (format == "NHWC") {
220+ return ge::FORMAT_NHWC;
221+ }
222+ if (format == "NCDHW") {
223+ return ge::FORMAT_NCDHW;
224+ }
225+ return ge::FORMAT_NCHW;
226+}
218 227 
219const char* FormatName(ge::Format format)228const char* FormatName(ge::Format format)
220{229{
@@ -227,6 +236,9 @@ const char* FormatName(ge::Format format)
227 if (format == ge::FORMAT_NHWC) {236 if (format == ge::FORMAT_NHWC) {
228 return "NHWC";237 return "NHWC";
229 }238 }
239+ if (format == ge::FORMAT_NCDHW) {
240+ return "NCDHW";
241+ }
230 return "OTHER";242 return "OTHER";
231}243}
232 244 
@@ -570,7 +582,7 @@ GraphBundle BuildGraph(const CaseDef& test)
570 bundle.node = test.id == "route-target" ? "bn_training_reduce_route" : "bn_training_reduce_" + test.id;582 bundle.node = test.id == "route-target" ? "bn_training_reduce_route" : "bn_training_reduce_" + test.id;
571 bundle.graph = ge::Graph((bundle.node + "_graph").c_str());583 bundle.graph = ge::Graph((bundle.node + "_graph").c_str());
572 auto reduce = ge::op::BNTrainingReduce(bundle.node.c_str());584 auto reduce = ge::op::BNTrainingReduce(bundle.node.c_str());
573- const ge::TensorDesc outDesc(ge::Shape(ge::UNKNOWN_RANK), ge::FORMAT_ND, ge::DT_FLOAT);585+ const ge::TensorDesc outDesc(ge::Shape(ge::UNKNOWN_RANK), ParseFormat(test.inputFormat), ge::DT_FLOAT);
574 reduce.update_output_desc_sum(outDesc);586 reduce.update_output_desc_sum(outDesc);
575 reduce.update_output_desc_square_sum(outDesc);587 reduce.update_output_desc_square_sum(outDesc);
576 588 
Mnorm/bn_training_reduce/op_api/aclnn_batch_norm_reduce.cpp+18-27
@@ -34,8 +34,6 @@ static const int64_t IN_SIZE = 4;
34static constexpr size_t OUT_SIZE = 1;34static constexpr size_t OUT_SIZE = 1;
35static constexpr size_t DIM_ZERO = 0;35static constexpr size_t DIM_ZERO = 0;
36static constexpr size_t DIM_ONE = 1;36static constexpr size_t DIM_ONE = 1;
37-static constexpr size_t DIM_TWO = 2;
38-static constexpr size_t DIM_THREE = 3;
39static constexpr size_t TENSOR_NUM = 2;37static constexpr size_t TENSOR_NUM = 2;
40 38 
41// 根据API定义，需要列出所能支持的所有dtype39// 根据API定义，需要列出所能支持的所有dtype
@@ -64,8 +62,8 @@ static bool CheckFormat(const aclTensor* x, const aclTensor* sum, const aclTenso
64 auto xFormat = x->GetStorageFormat();62 auto xFormat = x->GetStorageFormat();
65 auto sumFormat = sum->GetStorageFormat();63 auto sumFormat = sum->GetStorageFormat();
66 auto squareSumFormat = squareSum->GetStorageFormat();64 auto squareSumFormat = squareSum->GetStorageFormat();
67- if (xFormat != op::Format::FORMAT_NCHW && xFormat != op::Format::FORMAT_NHWC) {65+ if (xFormat != op::Format::FORMAT_NCHW) {
68- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of x only supports [NCHW, NHWC], but format is [%s].",66+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of x only supports [NCHW], but format is [%s].",
69 op::ToString(xFormat).GetString());67 op::ToString(xFormat).GetString());
70 return false;68 return false;
71 }69 }
@@ -99,10 +97,7 @@ static bool CheckShape(const aclTensor* x, const aclTensor* sum, const aclTensor
99 sumShape.GetDimNum()),97 sumShape.GetDimNum()),
100 return false);98 return false);
101 99 
102- int64_t inputC = xShape.GetDim(DIM_ONE);100+ const int64_t inputC = xShape.GetDim(DIM_ONE);
103- if (x->GetStorageFormat() == op::Format::FORMAT_NHWC) {
104- inputC = xShape.GetDim(DIM_THREE);
105- }
106 OP_CHECK(sumShape.GetDim(DIM_ZERO) == inputC,101 OP_CHECK(sumShape.GetDim(DIM_ZERO) == inputC,
107 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "It is expected size of sum equals to input_C(%ld), but size is %zu",102 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "It is expected size of sum equals to input_C(%ld), but size is %zu",
108 inputC, sumShape.GetDim(DIM_ZERO)),103 inputC, sumShape.GetDim(DIM_ZERO)),
@@ -129,12 +124,10 @@ static aclnnStatus CheckParams(const aclTensor* x, aclTensor* sum, aclTensor* sq
129 return ACLNN_SUCCESS;124 return ACLNN_SUCCESS;
130}125}
131 126 
132-const aclTensor* ResizeTo4D(const aclTensor* input, op::Format format, aclOpExecutor* executor)127+const aclTensor* ResizeTo4D(const aclTensor* input, aclOpExecutor* executor)
133{128{
134- const int64_t appendDimNchw[] = {0, 2, 3};129+ const int64_t appendDim[] = {0, 2, 3};
135- const int64_t appendDimNhwc[] = {0, 1, 2};130+ constexpr size_t appendDimNum = sizeof(appendDim) / sizeof(appendDim[0]);
136- const int64_t* appendDim = format == op::Format::FORMAT_NHWC ? appendDimNhwc : appendDimNchw;
137- constexpr size_t appendDimNum = sizeof(appendDimNchw) / sizeof(appendDimNchw[0]);
138 aclIntArray* newShape = executor->AllocIntArray(appendDim, appendDimNum);131 aclIntArray* newShape = executor->AllocIntArray(appendDim, appendDimNum);
139 132 
140 auto inputUnsqueeze = l0op::UnsqueezeNd(input, newShape, executor);133 auto inputUnsqueeze = l0op::UnsqueezeNd(input, newShape, executor);
@@ -144,18 +137,16 @@ const aclTensor* ResizeTo4D(const aclTensor* input, op::Format format, aclOpExec
144 auto formatTensor = executor == nullptr ? const_cast<aclTensor*>(inputUnsqueeze) :137 auto formatTensor = executor == nullptr ? const_cast<aclTensor*>(inputUnsqueeze) :
145 executor->CreateView(inputUnsqueeze, inputUnsqueeze->GetViewShape(),138 executor->CreateView(inputUnsqueeze, inputUnsqueeze->GetViewShape(),
146 inputUnsqueeze->GetViewOffset());139 inputUnsqueeze->GetViewOffset());
147- formatTensor->SetViewFormat(format);140+ formatTensor->SetViewFormat(Format::FORMAT_NCHW);
148- formatTensor->SetOriginalFormat(format);141+ formatTensor->SetOriginalFormat(Format::FORMAT_NCHW);
149- formatTensor->SetStorageFormat(format);142+ formatTensor->SetStorageFormat(Format::FORMAT_NCHW);
150 return formatTensor;143 return formatTensor;
151}144}
152 145 
153const aclTensor* ResizeTo1D(const aclTensor* input, aclOpExecutor* executor)146const aclTensor* ResizeTo1D(const aclTensor* input, aclOpExecutor* executor)
154{147{
155- const int64_t removeDimNchw[] = {0, 2, 3};148+ const int64_t removeDim[] = {0, 2, 3};
156- const int64_t removeDimNhwc[] = {0, 1, 2};149+ constexpr size_t removeDimNum = sizeof(removeDim) / sizeof(removeDim[0]);
157- const int64_t* removeDim = input->GetStorageFormat() == op::Format::FORMAT_NHWC ? removeDimNhwc : removeDimNchw;
158- constexpr size_t removeDimNum = sizeof(removeDimNchw) / sizeof(removeDimNchw[0]);
159 aclIntArray* newShape = executor->AllocIntArray(removeDim, removeDimNum);150 aclIntArray* newShape = executor->AllocIntArray(removeDim, removeDimNum);
160 151 
161 auto inputSqueeze = l0op::SqueezeNd(input, newShape, executor);152 auto inputSqueeze = l0op::SqueezeNd(input, newShape, executor);
@@ -183,8 +174,11 @@ aclnnStatus aclnnBatchNormReduceGetWorkspaceSize(const aclTensor* x, aclTensor*
183 // 检查必选输入/输出是否为空指针174 // 检查必选输入/输出是否为空指针
184 CHECK_RET(CheckNotNull(x, sum, squareSum), ACLNN_ERR_PARAM_NULLPTR);175 CHECK_RET(CheckNotNull(x, sum, squareSum), ACLNN_ERR_PARAM_NULLPTR);
185 176 
177+ auto ret = CheckParams(x, sum, squareSum);
178+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
179+ 
186 if (x->IsEmpty() || sum->IsEmpty() || squareSum->IsEmpty()) {180 if (x->IsEmpty() || sum->IsEmpty() || squareSum->IsEmpty()) {
187- auto ret = op::ProcessEmptyTensorWithValue(sum, 0, uniqueExecutor.get());181+ ret = op::ProcessEmptyTensorWithValue(sum, 0, uniqueExecutor.get());
188 CHECK_RET(ret == ACLNN_SUCCESS, ret);182 CHECK_RET(ret == ACLNN_SUCCESS, ret);
189 ret = op::ProcessEmptyTensorWithValue(squareSum, 0, uniqueExecutor.get());183 ret = op::ProcessEmptyTensorWithValue(squareSum, 0, uniqueExecutor.get());
190 CHECK_RET(ret == ACLNN_SUCCESS, ret);184 CHECK_RET(ret == ACLNN_SUCCESS, ret);
@@ -193,16 +187,13 @@ aclnnStatus aclnnBatchNormReduceGetWorkspaceSize(const aclTensor* x, aclTensor*
193 return ACLNN_SUCCESS;187 return ACLNN_SUCCESS;
194 }188 }
195 189 
196- auto ret = CheckParams(x, sum, squareSum);
197- CHECK_RET(ret == ACLNN_SUCCESS, ret);
198- 
199 auto xContiguous = l0op::Contiguous(x, uniqueExecutor.get());190 auto xContiguous = l0op::Contiguous(x, uniqueExecutor.get());
200 CHECK_RET(xContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);191 CHECK_RET(xContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
201 192 
202- auto sum4D = ResizeTo4D(sum, xContiguous->GetStorageFormat(), uniqueExecutor.get());193+ auto sumNCHW = ResizeTo4D(sum, uniqueExecutor.get());
203- CHECK_RET(sum4D != nullptr, ACLNN_ERR_INNER_NULLPTR);194+ CHECK_RET(sumNCHW != nullptr, ACLNN_ERR_INNER_NULLPTR);
204 195 
205- std::array<aclTensor*, TENSOR_NUM> sumTensor = l0op::BNTrainingReduce(xContiguous, sum4D->GetViewShape(),196+ std::array<aclTensor*, TENSOR_NUM> sumTensor = l0op::BNTrainingReduce(xContiguous, sumNCHW->GetViewShape(),
206 uniqueExecutor.get());197 uniqueExecutor.get());
207 198 
208 auto sumND = ResizeTo1D(sumTensor[0], uniqueExecutor.get());199 auto sumND = ResizeTo1D(sumTensor[0], uniqueExecutor.get());
Mnorm/bn_training_reduce/op_host/arch35/bn_training_reduce_tiling_arch35.cpp+35-17文件内容审核中，请稍后刷新重试
Mnorm/bn_training_reduce/op_host/arch35/bn_training_reduce_tiling_public.cpp+28-9
@@ -68,7 +68,7 @@ int64_t DTypeSize(BNTrainingReducePublicDType dtype)
68 68 
69size_t ChannelIndex(BNTrainingReducePublicFormat format)69size_t ChannelIndex(BNTrainingReducePublicFormat format)
70{70{
71- return format == BNTrainingReducePublicFormat::NCHW ? 1U : 3U;71+ return format == BNTrainingReducePublicFormat::NHWC ? 3U : 1U;
72}72}
73 73 
74int64_t ChannelSize(const BNTrainingReducePublicInputs& inputs) { return inputs.shape[ChannelIndex(inputs.format)]; }74int64_t ChannelSize(const BNTrainingReducePublicInputs& inputs) { return inputs.shape[ChannelIndex(inputs.format)]; }
@@ -135,16 +135,30 @@ struct TilingContext {
135 bool arithmeticOverflow = false;135 bool arithmeticOverflow = false;
136};136};
137 137 
138-BNTrainingReducePublicStatus ValidateInputs(const BNTrainingReducePublicInputs& inputs)138+bool IsSupportedFormatAndRank(const BNTrainingReducePublicInputs& inputs)
139+{
140+ switch (inputs.format) {
141+ case BNTrainingReducePublicFormat::NCHW:
142+ return inputs.rank >= 2 && inputs.rank <= 4;
143+ case BNTrainingReducePublicFormat::NHWC:
144+ return inputs.rank == 4;
145+ case BNTrainingReducePublicFormat::NCDHW:
146+ return inputs.rank == 5;
147+ default:
148+ return false;
149+ }
150+}
151+ 
152+BNTrainingReducePublicStatus ValidateInputsImpl(const BNTrainingReducePublicInputs& inputs)
139{153{
140 if (!inputs.inputPresent) {154 if (!inputs.inputPresent) {
141 return BNTrainingReducePublicStatus::NULL_INPUT;155 return BNTrainingReducePublicStatus::NULL_INPUT;
142 }156 }
143- if (inputs.rank != 4 ||157+ if (!IsSupportedFormatAndRank(inputs)) {
144- (inputs.format != BNTrainingReducePublicFormat::NCHW && inputs.format != BNTrainingReducePublicFormat::NHWC)) {
145 return BNTrainingReducePublicStatus::SHAPE_MISMATCH;158 return BNTrainingReducePublicStatus::SHAPE_MISMATCH;
146 }159 }
147- if (std::any_of(inputs.shape.begin(), inputs.shape.end(), [](int64_t dim) { return dim < 0; })) {160+ const auto shapeEnd = inputs.shape.begin() + inputs.rank;
161+ if (std::any_of(inputs.shape.begin(), shapeEnd, [](int64_t dim) { return dim < 0; })) {
148 return BNTrainingReducePublicStatus::SHAPE_MISMATCH;162 return BNTrainingReducePublicStatus::SHAPE_MISMATCH;
149 }163 }
150 if (DTypeSize(inputs.inputDtype) == 0) {164 if (DTypeSize(inputs.inputDtype) == 0) {
@@ -173,7 +187,7 @@ BNTrainingReduceEmptyKind ClassifyEmpty(const BNTrainingReducePublicInputs& inpu
173 if (inputs.shape[channelIndex] == 0) {187 if (inputs.shape[channelIndex] == 0) {
174 return BNTrainingReduceEmptyKind::EMPTY_A;188 return BNTrainingReduceEmptyKind::EMPTY_A;
175 }189 }
176- for (size_t i = 0; i < inputs.shape.size(); ++i) {190+ for (size_t i = 0; i < static_cast<size_t>(inputs.rank); ++i) {
177 if (i != channelIndex && inputs.shape[i] == 0) {191 if (i != channelIndex && inputs.shape[i] == 0) {
178 return BNTrainingReduceEmptyKind::EMPTY_R;192 return BNTrainingReduceEmptyKind::EMPTY_R;
179 }193 }
@@ -189,9 +203,9 @@ bool NormalizePattern(const BNTrainingReducePublicInputs& inputs, TilingContext&
189 ctx.cacheLineSize = inputs.cacheLineSize;203 ctx.cacheLineSize = inputs.cacheLineSize;
190 ctx.dtypeSize = DTypeSize(inputs.inputDtype);204 ctx.dtypeSize = DTypeSize(inputs.inputDtype);
191 205 
192- std::array<bool, 4> initialTypes = {true, true, true, true};206+ std::array<bool, 5> initialTypes = {true, true, true, true, true};
193 initialTypes[ChannelIndex(inputs.format)] = false;207 initialTypes[ChannelIndex(inputs.format)] = false;
194- for (size_t i = 0; i < inputs.shape.size(); ++i) {208+ for (size_t i = 0; i < static_cast<size_t>(inputs.rank); ++i) {
195 if (inputs.shape[i] != 1) {209 if (inputs.shape[i] != 1) {
196 ctx.axisShape.push_back(inputs.shape[i]);210 ctx.axisShape.push_back(inputs.shape[i]);
197 ctx.isReduce.push_back(initialTypes[i]);211 ctx.isReduce.push_back(initialTypes[i]);
@@ -775,7 +789,7 @@ BNTrainingReducePublicResult ComputeEmptyTiling(const BNTrainingReducePublicInpu
775BNTrainingReducePublicResult ComputeAllRoutes(const BNTrainingReducePublicInputs& inputs)789BNTrainingReducePublicResult ComputeAllRoutes(const BNTrainingReducePublicInputs& inputs)
776{790{
777 BNTrainingReducePublicResult result;791 BNTrainingReducePublicResult result;
778- result.status = ValidateInputs(inputs);792+ result.status = ValidateInputsImpl(inputs);
779 if (result.status != BNTrainingReducePublicStatus::SUCCESS) {793 if (result.status != BNTrainingReducePublicStatus::SUCCESS) {
780 return result;794 return result;
781 }795 }
@@ -834,6 +848,11 @@ BNTrainingReducePublicResult ComputeAllRoutes(const BNTrainingReducePublicInputs
834 848 
835} // namespace849} // namespace
836 850 
851+BNTrainingReducePublicStatus ValidateBNTrainingReducePublicInputs(const BNTrainingReducePublicInputs& inputs)
852+{
853+ return ValidateInputsImpl(inputs);
854+}
855+ 
837BNTrainingReducePublicResult ComputeBNTrainingReducePublicTiling(const BNTrainingReducePublicInputs& inputs)856BNTrainingReducePublicResult ComputeBNTrainingReducePublicTiling(const BNTrainingReducePublicInputs& inputs)
838{857{
839 return ComputeAllRoutes(inputs);858 return ComputeAllRoutes(inputs);
Mnorm/bn_training_reduce/op_host/arch35/bn_training_reduce_tiling_public.h+4-1
@@ -32,6 +32,7 @@ enum class BNTrainingReducePublicStatus : int32_t {
32enum class BNTrainingReducePublicFormat : int32_t {32enum class BNTrainingReducePublicFormat : int32_t {
33 NCHW = 0,33 NCHW = 0,
34 NHWC,34 NHWC,
35+ NCDHW,
35};36};
36 37 
37enum class BNTrainingReducePublicDType : int32_t {38enum class BNTrainingReducePublicDType : int32_t {
@@ -61,7 +62,7 @@ enum class BNTrainingReduceTilingKey : int64_t {
61struct BNTrainingReducePublicInputs {62struct BNTrainingReducePublicInputs {
62 bool inputPresent = true;63 bool inputPresent = true;
63 int32_t rank = 4;64 int32_t rank = 4;
64- std::array<int64_t, 4> shape = {1, 1, 1, 1};65+ std::array<int64_t, 5> shape = {1, 1, 1, 1, 1};
65 BNTrainingReducePublicFormat format = BNTrainingReducePublicFormat::NCHW;66 BNTrainingReducePublicFormat format = BNTrainingReducePublicFormat::NCHW;
66 BNTrainingReducePublicDType inputDtype = BNTrainingReducePublicDType::FLOAT32;67 BNTrainingReducePublicDType inputDtype = BNTrainingReducePublicDType::FLOAT32;
67 68 
@@ -90,6 +91,8 @@ struct BNTrainingReducePublicResult {
90 BNTrainingReduceTilingData tilingData = {};91 BNTrainingReduceTilingData tilingData = {};
91};92};
92 93 
94+BNTrainingReducePublicStatus ValidateBNTrainingReducePublicInputs(const BNTrainingReducePublicInputs& inputs);
95+ 
93// Side-effect-free Host Tiling implementation used by the runtime glue.96// Side-effect-free Host Tiling implementation used by the runtime glue.
94BNTrainingReducePublicResult ComputeBNTrainingReducePublicTiling(const BNTrainingReducePublicInputs& inputs);97BNTrainingReducePublicResult ComputeBNTrainingReducePublicTiling(const BNTrainingReducePublicInputs& inputs);
95 98 
Mnorm/bn_training_reduce/op_host/bn_training_reduce_def.cpp+18-13
@@ -17,23 +17,28 @@ public:
17 {17 {
18 this->Input("x")18 this->Input("x")
19 .ParamType(REQUIRED)19 .ParamType(REQUIRED)
20- .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT})20+ .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT,
21- .Format(21+ ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT})
22- {ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC, ge::FORMAT_NHWC})22+ .Format({ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC,
23- .UnknownShapeFormat(23+ ge::FORMAT_NHWC, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW})
24- {ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC, ge::FORMAT_NHWC});24+ .UnknownShapeFormat({ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC,
25+ ge::FORMAT_NHWC, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW});
25 this->Output("sum")26 this->Output("sum")
26 .ParamType(REQUIRED)27 .ParamType(REQUIRED)
27- .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})28+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
28- .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})29+ ge::DT_FLOAT, ge::DT_FLOAT})
29- .UnknownShapeFormat(30+ .Format({ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC,
30- {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});31+ ge::FORMAT_NHWC, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW})
32+ .UnknownShapeFormat({ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC,
33+ ge::FORMAT_NHWC, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW});
31 this->Output("square_sum")34 this->Output("square_sum")
32 .ParamType(REQUIRED)35 .ParamType(REQUIRED)
33- .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})36+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
34- .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})37+ ge::DT_FLOAT, ge::DT_FLOAT})
35- .UnknownShapeFormat(38+ .Format({ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC,
36- {ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});39+ ge::FORMAT_NHWC, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW})
40+ .UnknownShapeFormat({ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NCHW, ge::FORMAT_NHWC, ge::FORMAT_NHWC,
41+ ge::FORMAT_NHWC, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW, ge::FORMAT_NCDHW});
37 this->FormatMatchMode(FormatCheckOption::STRICT);42 this->FormatMatchMode(FormatCheckOption::STRICT);
38 OpAICoreConfig aiCoreConfig;43 OpAICoreConfig aiCoreConfig;
39 aiCoreConfig.DynamicCompileStaticFlag(true)44 aiCoreConfig.DynamicCompileStaticFlag(true)
Mnorm/bn_training_reduce/op_host/bn_training_reduce_infershape.cpp+14-8
@@ -35,18 +35,24 @@ static ge::graphStatus InferShape4BNTrainingReduce(gert::InferShapeContext* cont
35 *squareSumShape = *xShape;35 *squareSumShape = *xShape;
36 return GRAPH_SUCCESS;36 return GRAPH_SUCCESS;
37 }37 }
38- if (xShape->GetDimNum() != 4) {38+ 
39- OP_LOGE(context, "BNTrainingReduce input x rank must be 4, but got %zu.", xShape->GetDimNum());39+ const size_t rank = xShape->GetDimNum();
40- return GRAPH_FAILED;
41- }
42 const ge::Format format = xDesc->GetOriginFormat();40 const ge::Format format = xDesc->GetOriginFormat();
43- if (format != FORMAT_NCHW && format != FORMAT_NHWC) {41+ size_t channelIndex = 0;
44- OP_LOGE(context, "BNTrainingReduce input x format must be NCHW or NHWC, but got %d.",42+ if (format == FORMAT_NCHW && rank >= 2U && rank <= 4U) {
45- static_cast<int32_t>(format));43+ channelIndex = 1U;
44+ } else if (format == FORMAT_NHWC && rank == 4U) {
45+ channelIndex = 3U;
46+ } else if (format == FORMAT_NCDHW && rank == 5U) {
47+ channelIndex = 1U;
48+ } else {
49+ OP_LOGE(context,
50+ "BNTrainingReduce on Ascend 950 only supports NCHW rank 2-4, NHWC rank 4 and NCDHW rank 5, "
51+ "but got format %d and rank %zu.",
52+ static_cast<int32_t>(format), rank);
46 return GRAPH_FAILED;53 return GRAPH_FAILED;
47 }54 }
48 55 
49- const size_t channelIndex = format == FORMAT_NCHW ? 1U : 3U;
50 const int64_t channel = xShape->GetDim(channelIndex);56 const int64_t channel = xShape->GetDim(channelIndex);
51 sumShape->SetDimNum(1);57 sumShape->SetDimNum(1);
52 sumShape->SetDim(0, channel);58 sumShape->SetDim(0, channel);
Mnorm/bn_training_reduce/op_kernel/bn_training_reduce.cpp+5-1
@@ -44,7 +44,11 @@ __global__ __aicore__ void bn_training_reduce(GM_ADDR x, GM_ADDR sum, GM_ADDR sq
44 }44 }
45 } else if constexpr (templateType) {45 } else if constexpr (templateType) {
46 NsBNTrainingReduce::BNTrainingReduceKernel<DTYPE_X, isTailR, isDeterministic> op;46 NsBNTrainingReduce::BNTrainingReduceKernel<DTYPE_X, isTailR, isDeterministic> op;
47- op.InitGroup(x, sum, squareSum, workspace, &tilingData);47+ GM_ADDR userWorkspace = workspace;
48+ if constexpr (isDeterministic) {
49+ userWorkspace = AscendC::GetUserWorkspace(workspace);
50+ }
51+ op.InitGroup(x, sum, squareSum, userWorkspace, &tilingData);
48 for (int32_t outputIdx = 0; outputIdx < 2; ++outputIdx) {52 for (int32_t outputIdx = 0; outputIdx < 2; ++outputIdx) {
49 op.ProcessGroup(outputIdx);53 op.ProcessGroup(outputIdx);
50 AscendC::SyncAll();54 AscendC::SyncAll();
Mnorm/bn_training_reduce/tests/ut/op_host/arch35/test_bn_training_reduce_tiling.cpp+62-5
@@ -8,6 +8,8 @@
8 * 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.
9 */9 */
10 10 
11+#include <initializer_list>
12+ 
11#include <gtest/gtest.h>13#include <gtest/gtest.h>
12 14 
13#include "../../../../op_host/arch35/bn_training_reduce_tiling_public.h"15#include "../../../../op_host/arch35/bn_training_reduce_tiling_public.h"
@@ -17,15 +19,19 @@ namespace {
17 19 
18constexpr int64_t kTestUbSizeBytes = 256 * 1024;20constexpr int64_t kTestUbSizeBytes = 256 * 1024;
19 21 
20-BNTrainingReducePublicInputs MakeInputs(const std::array<int64_t, 4>& shape, BNTrainingReducePublicFormat format)22+BNTrainingReducePublicInputs MakeInputs(std::initializer_list<int64_t> shape, BNTrainingReducePublicFormat format)
21{23{
22 BNTrainingReducePublicInputs inputs;24 BNTrainingReducePublicInputs inputs;
23- inputs.shape = shape;25+ inputs.rank = static_cast<int32_t>(shape.size());
26+ size_t index = 0;
27+ for (const int64_t dim : shape) {
28+ inputs.shape[index++] = dim;
29+ }
24 inputs.format = format;30 inputs.format = format;
25 inputs.ubSize = kTestUbSizeBytes;31 inputs.ubSize = kTestUbSizeBytes;
26- const size_t channelIndex = format == BNTrainingReducePublicFormat::NCHW ? 1 : 3;32+ const size_t channelIndex = format == BNTrainingReducePublicFormat::NHWC ? 3U : 1U;
27- inputs.sumDim0 = shape[channelIndex];33+ inputs.sumDim0 = inputs.shape[channelIndex];
28- inputs.squareSumDim0 = shape[channelIndex];34+ inputs.squareSumDim0 = inputs.shape[channelIndex];
29 return inputs;35 return inputs;
30}36}
31 37 
@@ -79,5 +85,56 @@ TEST(BNTrainingReduceTilingTest, SupportsNhwcEmptyReduceAxis)
79 EXPECT_GT(result.tilingData.usedCoreNum, 0);85 EXPECT_GT(result.tilingData.usedCoreNum, 0);
80}86}
81 87 
88+TEST(BNTrainingReduceTilingTest, SupportsNchwRank2)
89+{
90+ const auto inputs = MakeInputs({32, 8}, BNTrainingReducePublicFormat::NCHW);
91+ const auto result = ComputeBNTrainingReducePublicTiling(inputs);
92+ 
93+ EXPECT_EQ(result.status, BNTrainingReducePublicStatus::SUCCESS);
94+ EXPECT_GT(result.blockDim, 0U);
95+}
96+ 
97+TEST(BNTrainingReduceTilingTest, SupportsNcdhw)
98+{
99+ const auto inputs = MakeInputs({2, 3, 4, 5, 6}, BNTrainingReducePublicFormat::NCDHW);
100+ const auto result = ComputeBNTrainingReducePublicTiling(inputs);
101+ 
102+ EXPECT_EQ(result.status, BNTrainingReducePublicStatus::SUCCESS);
103+ EXPECT_GT(result.blockDim, 0U);
104+}
105+ 
106+TEST(BNTrainingReduceTilingTest, RejectsUnsupportedFormat)
107+{
108+ auto inputs = MakeInputs({2, 3, 4, 5}, BNTrainingReducePublicFormat::NCHW);
109+ inputs.format = static_cast<BNTrainingReducePublicFormat>(99);
110+ const auto result = ComputeBNTrainingReducePublicTiling(inputs);
111+ 
112+ EXPECT_EQ(result.status, BNTrainingReducePublicStatus::SHAPE_MISMATCH);
113+}
114+ 
115+TEST(BNTrainingReduceTilingTest, RejectsWrongRankForNhwc)
116+{
117+ const auto inputs = MakeInputs({2, 4, 3}, BNTrainingReducePublicFormat::NHWC);
118+ const auto result = ComputeBNTrainingReducePublicTiling(inputs);
119+ 
120+ EXPECT_EQ(result.status, BNTrainingReducePublicStatus::SHAPE_MISMATCH);
121+}
122+ 
123+TEST(BNTrainingReduceTilingTest, RejectsWrongRankForNcdhw)
124+{
125+ const auto inputs = MakeInputs({2, 3, 4, 5}, BNTrainingReducePublicFormat::NCDHW);
126+ const auto result = ComputeBNTrainingReducePublicTiling(inputs);
127+ 
128+ EXPECT_EQ(result.status, BNTrainingReducePublicStatus::SHAPE_MISMATCH);
129+}
130+ 
131+TEST(BNTrainingReduceTilingTest, SharedValidationRejectsInvalidSmallROutputs)
132+{
133+ auto inputs = MakeInputs({1, 64, 1, 3}, BNTrainingReducePublicFormat::NCHW);
134+ inputs.sumRank = 2;
135+ 
136+ EXPECT_EQ(ValidateBNTrainingReducePublicInputs(inputs), BNTrainingReducePublicStatus::SHAPE_MISMATCH);
137+}
138+ 
82} // namespace139} // namespace
83} // namespace optiling140} // namespace optiling
Mnorm/bn_training_reduce/tests/ut/op_host/op_api/test_aclnn_batch_norm_reduce.cpp+15-2
@@ -34,7 +34,7 @@ TEST_F(l2BatchNormReduceTest, l2_batch_norm_reduce_bfloat16)
34 EXPECT_EQ(aclRet, ACL_SUCCESS);34 EXPECT_EQ(aclRet, ACL_SUCCESS);
35}35}
36 36 
37-TEST_F(l2BatchNormReduceTest, l2_batch_norm_reduce_nhwc)37+TEST_F(l2BatchNormReduceTest, l2_batch_norm_reduce_rejects_nhwc)
38{38{
39 auto x = TensorDesc({3, 3, 8, 5}, ACL_FLOAT, ACL_FORMAT_NHWC);39 auto x = TensorDesc({3, 3, 8, 5}, ACL_FLOAT, ACL_FORMAT_NHWC);
40 auto sum = TensorDesc({5}, ACL_FLOAT, ACL_FORMAT_ND);40 auto sum = TensorDesc({5}, ACL_FLOAT, ACL_FORMAT_ND);
@@ -44,5 +44,18 @@ TEST_F(l2BatchNormReduceTest, l2_batch_norm_reduce_nhwc)
44 44 
45 uint64_t workspaceSize = 0;45 uint64_t workspaceSize = 0;
46 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);46 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
47- EXPECT_EQ(aclRet, ACL_SUCCESS);47+ EXPECT_NE(aclRet, ACL_SUCCESS);
48+}
49+ 
50+TEST_F(l2BatchNormReduceTest, l2_batch_norm_reduce_rejects_empty_nhwc)
51+{
52+ auto x = TensorDesc({0, 3, 8, 5}, ACL_FLOAT, ACL_FORMAT_NHWC);
53+ auto sum = TensorDesc({5}, ACL_FLOAT, ACL_FORMAT_ND);
54+ auto squareSum = TensorDesc({5}, ACL_FLOAT, ACL_FORMAT_ND);
55+ 
56+ auto ut = OP_API_UT(aclnnBatchNormReduce, INPUT(x), OUTPUT(sum, squareSum));
57+ 
58+ uint64_t workspaceSize = 0;
59+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
60+ EXPECT_NE(aclRet, ACL_SUCCESS);
48}61}
Mnorm/bn_training_reduce/tests/ut/op_host/test_bn_training_reduce_infershape.cpp+51-0
@@ -53,6 +53,18 @@ TEST(BNTrainingReduceInferShapeTest, SupportsNchw)
53 EXPECT_EQ(squareSumShape.GetDim(0), 3);53 EXPECT_EQ(squareSumShape.GetDim(0), 3);
54}54}
55 55 
56+TEST(BNTrainingReduceInferShapeTest, SupportsNchwRank2)
57+{
58+ gert::Shape inputShape = {4, 3};
59+ gert::Shape sumShape;
60+ gert::Shape squareSumShape;
61+ 
62+ ASSERT_EQ(RunInferShape(inputShape, ge::FORMAT_NCHW, sumShape, squareSumShape), ge::GRAPH_SUCCESS);
63+ ASSERT_EQ(sumShape.GetDimNum(), 1U);
64+ EXPECT_EQ(sumShape.GetDim(0), 3);
65+ EXPECT_EQ(squareSumShape.GetDim(0), 3);
66+}
67+ 
56TEST(BNTrainingReduceInferShapeTest, SupportsNhwc)68TEST(BNTrainingReduceInferShapeTest, SupportsNhwc)
57{69{
58 gert::Shape inputShape = {2, 4, 5, 3};70 gert::Shape inputShape = {2, 4, 5, 3};
@@ -65,6 +77,18 @@ TEST(BNTrainingReduceInferShapeTest, SupportsNhwc)
65 EXPECT_EQ(squareSumShape.GetDim(0), 3);77 EXPECT_EQ(squareSumShape.GetDim(0), 3);
66}78}
67 79 
80+TEST(BNTrainingReduceInferShapeTest, SupportsNcdhw)
81+{
82+ gert::Shape inputShape = {2, 3, 4, 5, 6};
83+ gert::Shape sumShape;
84+ gert::Shape squareSumShape;
85+ 
86+ ASSERT_EQ(RunInferShape(inputShape, ge::FORMAT_NCDHW, sumShape, squareSumShape), ge::GRAPH_SUCCESS);
87+ ASSERT_EQ(sumShape.GetDimNum(), 1U);
88+ EXPECT_EQ(sumShape.GetDim(0), 3);
89+ EXPECT_EQ(squareSumShape.GetDim(0), 3);
90+}
91+ 
68TEST(BNTrainingReduceInferShapeTest, RejectsNd)92TEST(BNTrainingReduceInferShapeTest, RejectsNd)
69{93{
70 gert::Shape inputShape = {2, 3, 4, 5};94 gert::Shape inputShape = {2, 3, 4, 5};
@@ -74,5 +98,32 @@ TEST(BNTrainingReduceInferShapeTest, RejectsNd)
74 EXPECT_EQ(RunInferShape(inputShape, ge::FORMAT_ND, sumShape, squareSumShape), ge::GRAPH_FAILED);98 EXPECT_EQ(RunInferShape(inputShape, ge::FORMAT_ND, sumShape, squareSumShape), ge::GRAPH_FAILED);
75}99}
76 100 
101+TEST(BNTrainingReduceInferShapeTest, RejectsNdc1hwc0OnAscend950)
102+{
103+ gert::Shape inputShape = {2, 1, 3, 4, 5, 16};
104+ gert::Shape sumShape;
105+ gert::Shape squareSumShape;
106+ 
107+ EXPECT_EQ(RunInferShape(inputShape, ge::FORMAT_NDC1HWC0, sumShape, squareSumShape), ge::GRAPH_FAILED);
108+}
109+ 
110+TEST(BNTrainingReduceInferShapeTest, RejectsWrongRankForNhwc)
111+{
112+ gert::Shape inputShape = {2, 4, 3};
113+ gert::Shape sumShape;
114+ gert::Shape squareSumShape;
115+ 
116+ EXPECT_EQ(RunInferShape(inputShape, ge::FORMAT_NHWC, sumShape, squareSumShape), ge::GRAPH_FAILED);
117+}
118+ 
119+TEST(BNTrainingReduceInferShapeTest, RejectsWrongRankForNcdhw)
120+{
121+ gert::Shape inputShape = {2, 3, 4, 5};
122+ gert::Shape sumShape;
123+ gert::Shape squareSumShape;
124+ 
125+ EXPECT_EQ(RunInferShape(inputShape, ge::FORMAT_NCDHW, sumShape, squareSumShape), ge::GRAPH_FAILED);
126+}
127+ 
77} // namespace128} // namespace
78} // namespace ops129} // namespace ops