Star1738
1648
Code
Code
Issues470
Pull Requests550
Actions
Discussion
Wiki
Analysis
Star1738
1648
Closed
fix(op_graph): 迁移 RaggedBinCount 和 MulNoNan 推导注册 #6548
raoliang_sacAdd on 20 days agoClosed at 15 days ago
fix(op_graph): 迁移 RaggedBinCount 和 MulNoNan 推导注册 #6548
Closed
raoliang_sacAdd on 20 days agoClosed at 15 days ago
Merged fromraoliang_sac/ops-mathfix/native-infer-registration-20261020intocann/ops-mathmaster
Total: 7 files changed+268-349
Amath/mul_no_nan/op_graph/mul_no_nan_graph_infer.cpp+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+#include "log/log.h"
12+#include "register/op_impl_registry.h"
13+ 
14+namespace ops {
15+ge::graphStatus InferDataType4MulNoNan(gert::InferDataTypeContext* context)
16+{
17+ if (context == nullptr) {
18+ return ge::GRAPH_FAILED;
19+ }
20+ const auto dtype = context->GetInputDataType(0);
21+ OP_CHECK_IF(dtype != context->GetInputDataType(1),
22+ OP_LOGE(context->GetNodeName(), "x1 and x2 must have the same dtype."), return ge::GRAPH_FAILED);
23+ return context->SetOutputDataType(0, dtype);
24+}
25+ 
26+IMPL_OP(MulNoNan).InferDataType(InferDataType4MulNoNan);
27+} // namespace ops
Mmath/mul_no_nan/op_host/mul_no_nan_infershape.cpp+1-12
@@ -18,23 +18,12 @@
18 18 
19using namespace ge;19using namespace ge;
20namespace ops {20namespace ops {
21-constexpr size_t MUL_NO_NAN_INPUT_X1_INDEX = 0;
22-constexpr size_t MUL_NO_NAN_OUTPUT_Y_INDEX = 0;
23- 
24static ge::graphStatus InferShape4MulNoNan(gert::InferShapeContext* context)21static ge::graphStatus InferShape4MulNoNan(gert::InferShapeContext* context)
25{22{
26 OP_LOGD(context->GetNodeName(), "Begin to do InferShape4MulNoNan in ops-math");23 OP_LOGD(context->GetNodeName(), "Begin to do InferShape4MulNoNan in ops-math");
27 return Ops::Base::InferShape4Broadcast(context);24 return Ops::Base::InferShape4Broadcast(context);
28}25}
29 26 
30-static ge::graphStatus InferDataType4MulNoNan(gert::InferDataTypeContext* context)27+IMPL_OP_INFERSHAPE(MulNoNan).InferShape(InferShape4MulNoNan);
31-{
32- OP_LOGD(context->GetNodeName(), "Begin to do InferDataType4MulNoNan in ops-math");
33- // Align with op_proto: output dtype is the same as input x1.
34- context->SetOutputDataType(MUL_NO_NAN_OUTPUT_Y_INDEX, context->GetInputDataType(MUL_NO_NAN_INPUT_X1_INDEX));
35- return ge::GRAPH_SUCCESS;
36-}
37- 
38-IMPL_OP_INFERSHAPE(MulNoNan).InferShape(InferShape4MulNoNan).InferDataType(InferDataType4MulNoNan);
39 28 
40} // namespace ops29} // namespace ops
Amath/mul_no_nan/tests/ut/op_graph/test_mul_no_nan_graph_infer_pass.cpp+100-0
@@ -0,0 +1,100 @@
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+#include <gtest/gtest.h>
12+#include "base/context_builder/op_infer_datatype_context_builder.h"
13+ 
14+namespace ops {
15+ge::graphStatus InferDataType4MulNoNan(gert::InferDataTypeContext* context);
16+} // namespace ops
17+ 
18+// InferDataType: output dtype follows x1 (fp16).
19+TEST(MulNoNanGraphInfer, infer_datatype_fp16)
20+{
21+ gert::OpInferDataTypeContextBuilder builder;
22+ builder.OpType("MulNoNan").OpName("MulNoNan");
23+ builder.IONum(2, 1);
24+ builder.InputTensorDesc(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
25+ builder.InputTensorDesc(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
26+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
27+ auto contextHolder = builder.Build();
28+ auto* context = contextHolder.GetContext();
29+ ASSERT_NE(context, nullptr);
30+ 
31+ auto ret = ops::InferDataType4MulNoNan(context);
32+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
33+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT16);
34+}
35+ 
36+// InferDataType: output dtype follows x1 (fp32).
37+TEST(MulNoNanGraphInfer, infer_datatype_fp32)
38+{
39+ gert::OpInferDataTypeContextBuilder builder;
40+ builder.OpType("MulNoNan").OpName("MulNoNan");
41+ builder.IONum(2, 1);
42+ builder.InputTensorDesc(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
43+ builder.InputTensorDesc(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
44+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
45+ auto contextHolder = builder.Build();
46+ auto* context = contextHolder.GetContext();
47+ ASSERT_NE(context, nullptr);
48+ 
49+ auto ret = ops::InferDataType4MulNoNan(context);
50+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
51+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
52+}
53+ 
54+// InferDataType: output dtype follows x1 (bf16).
55+TEST(MulNoNanGraphInfer, infer_datatype_bf16)
56+{
57+ gert::OpInferDataTypeContextBuilder builder;
58+ builder.OpType("MulNoNan").OpName("MulNoNan");
59+ builder.IONum(2, 1);
60+ builder.InputTensorDesc(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
61+ builder.InputTensorDesc(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
62+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
63+ auto contextHolder = builder.Build();
64+ auto* context = contextHolder.GetContext();
65+ ASSERT_NE(context, nullptr);
66+ 
67+ auto ret = ops::InferDataType4MulNoNan(context);
68+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
69+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_BF16);
70+}
71+ 
72+// InferDataType: output dtype follows x1 (int32).
73+TEST(MulNoNanGraphInfer, infer_datatype_int32)
74+{
75+ gert::OpInferDataTypeContextBuilder builder;
76+ builder.OpType("MulNoNan").OpName("MulNoNan");
77+ builder.IONum(2, 1);
78+ builder.InputTensorDesc(0, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND);
79+ builder.InputTensorDesc(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND);
80+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
81+ auto contextHolder = builder.Build();
82+ auto* context = contextHolder.GetContext();
83+ ASSERT_NE(context, nullptr);
84+ 
85+ auto ret = ops::InferDataType4MulNoNan(context);
86+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
87+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_INT32);
88+}
89+ 
90+TEST(MulNoNanGraphInfer, RejectsMismatchedInputsWithoutLegacyVerifier)
91+{
92+ gert::OpInferDataTypeContextBuilder builder;
93+ builder.OpType("MulNoNan").OpName("mul_no_nan_mismatched_dtype").IONum(2, 1);
94+ builder.InputTensorDesc(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
95+ builder.InputTensorDesc(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
96+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
97+ auto holder = builder.Build();
98+ ASSERT_NE(holder.GetContext(), nullptr);
99+ EXPECT_EQ(ops::InferDataType4MulNoNan(holder.GetContext()), ge::GRAPH_FAILED);
100+}
Mmath/mul_no_nan/tests/ut/op_host/test_mul_no_nan_infershape.cpp+0-97
@@ -12,7 +12,6 @@
12#include <iostream>12#include <iostream>
13#include "infershape_context_faker.h"13#include "infershape_context_faker.h"
14#include "infershape_case_executor.h"14#include "infershape_case_executor.h"
15-#include "op_infer_datatype_context_builder.h"
16#include "base/registry/op_impl_space_registry_v2.h"15#include "base/registry/op_impl_space_registry_v2.h"
17 16 
18using namespace ge;17using namespace ge;
@@ -203,99 +202,3 @@ TEST_F(MulNoNanInfershape, cross_rank_x2_bigger_int32)
203 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 4, 5}};202 std::vector<std::vector<int64_t>> expectOutputShape = {{3, 4, 5}};
204 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);203 ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectOutputShape);
205}204}
206- 
207-// InferDataType: output dtype follows x1 (fp16).
208-TEST_F(MulNoNanInfershape, infer_datatype_fp16)
209-{
210- auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
211- ASSERT_NE(spaceRegistry, nullptr);
212- auto opImpl = spaceRegistry->GetOpImpl("MulNoNan");
213- ASSERT_NE(opImpl, nullptr);
214- ASSERT_NE(opImpl->infer_datatype, nullptr);
215- 
216- gert::OpInferDataTypeContextBuilder builder;
217- builder.OpType("MulNoNan").OpName("MulNoNan");
218- builder.IONum(2, 1);
219- builder.InputTensorDesc(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
220- builder.InputTensorDesc(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
221- builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
222- auto contextHolder = builder.Build();
223- auto* context = contextHolder.GetContext();
224- ASSERT_NE(context, nullptr);
225- 
226- auto ret = opImpl->infer_datatype(context);
227- EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
228- EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT16);
229-}
230- 
231-// InferDataType: output dtype follows x1 (fp32).
232-TEST_F(MulNoNanInfershape, infer_datatype_fp32)
233-{
234- auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
235- ASSERT_NE(spaceRegistry, nullptr);
236- auto opImpl = spaceRegistry->GetOpImpl("MulNoNan");
237- ASSERT_NE(opImpl, nullptr);
238- ASSERT_NE(opImpl->infer_datatype, nullptr);
239- 
240- gert::OpInferDataTypeContextBuilder builder;
241- builder.OpType("MulNoNan").OpName("MulNoNan");
242- builder.IONum(2, 1);
243- builder.InputTensorDesc(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
244- builder.InputTensorDesc(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
245- builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
246- auto contextHolder = builder.Build();
247- auto* context = contextHolder.GetContext();
248- ASSERT_NE(context, nullptr);
249- 
250- auto ret = opImpl->infer_datatype(context);
251- EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
252- EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
253-}
254- 
255-// InferDataType: output dtype follows x1 (bf16).
256-TEST_F(MulNoNanInfershape, infer_datatype_bf16)
257-{
258- auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
259- ASSERT_NE(spaceRegistry, nullptr);
260- auto opImpl = spaceRegistry->GetOpImpl("MulNoNan");
261- ASSERT_NE(opImpl, nullptr);
262- ASSERT_NE(opImpl->infer_datatype, nullptr);
263- 
264- gert::OpInferDataTypeContextBuilder builder;
265- builder.OpType("MulNoNan").OpName("MulNoNan");
266- builder.IONum(2, 1);
267- builder.InputTensorDesc(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
268- builder.InputTensorDesc(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
269- builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
270- auto contextHolder = builder.Build();
271- auto* context = contextHolder.GetContext();
272- ASSERT_NE(context, nullptr);
273- 
274- auto ret = opImpl->infer_datatype(context);
275- EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
276- EXPECT_EQ(context->GetOutputDataType(0), ge::DT_BF16);
277-}
278- 
279-// InferDataType: output dtype follows x1 (int32).
280-TEST_F(MulNoNanInfershape, infer_datatype_int32)
281-{
282- auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
283- ASSERT_NE(spaceRegistry, nullptr);
284- auto opImpl = spaceRegistry->GetOpImpl("MulNoNan");
285- ASSERT_NE(opImpl, nullptr);
286- ASSERT_NE(opImpl->infer_datatype, nullptr);
287- 
288- gert::OpInferDataTypeContextBuilder builder;
289- builder.OpType("MulNoNan").OpName("MulNoNan");
290- builder.IONum(2, 1);
291- builder.InputTensorDesc(0, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND);
292- builder.InputTensorDesc(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND);
293- builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
294- auto contextHolder = builder.Build();
295- auto* context = contextHolder.GetContext();
296- ASSERT_NE(context, nullptr);
297- 
298- auto ret = opImpl->infer_datatype(context);
299- EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
300- EXPECT_EQ(context->GetOutputDataType(0), ge::DT_INT32);
301-}
Mmath/ragged_bin_count/op_graph/ragged_bin_count_graph_infer.cpp+0-72
@@ -15,78 +15,6 @@
15 15 
16#include "ragged_bin_count_graph_infer_internal.h"16#include "ragged_bin_count_graph_infer_internal.h"
17 17 
18-#include <cstdint>
19-#include <cstring>
20-#include <utility>
21-#include <vector>
22- 
23-#include "graph/operator_reg.h"
24- 
25-namespace ge {
26-namespace {
27-graphStatus ReadLegacyBinCount(const Operator& op, int64_t& bins)
28-{
29- bins = UNKNOWN_DIM;
30- Tensor sizeTensor;
31- if (op.GetInputConstData("size", sizeTensor) != GRAPH_SUCCESS) {
32- return GRAPH_SUCCESS;
33- }
34- const auto dtype = op.GetInputDescByName("size").GetDataType();
35- if (sizeTensor.GetData() == nullptr || sizeTensor.GetTensorDesc().GetDataType() != dtype) {
36- return GRAPH_FAILED;
37- }
38- if (dtype == DT_INT32 && sizeTensor.GetSize() == sizeof(int32_t)) {
39- int32_t value = 0;
40- std::memcpy(&value, sizeTensor.GetData(), sizeof(value));
41- bins = value;
42- } else if (dtype == DT_INT64 && sizeTensor.GetSize() == sizeof(int64_t)) {
43- std::memcpy(&bins, sizeTensor.GetData(), sizeof(bins));
44- } else {
45- return GRAPH_FAILED;
46- }
47- return bins >= 0 ? GRAPH_SUCCESS : GRAPH_FAILED;
48-}
49-} // namespace
50- 
51-IMPLEMT_COMMON_INFERFUNC(RaggedBinCountLegacyInferShape)
52-{
53- const auto splits = op.GetInputDescByName("splits").GetShape().GetDims();
54- const auto size = op.GetInputDescByName("size").GetShape().GetDims();
55- if (splits != UNKNOWN_RANK && (splits.size() != 1U || (splits[0] != UNKNOWN_DIM && splits[0] < 2))) {
56- return GRAPH_FAILED;
57- }
58- if (size != UNKNOWN_RANK && (size.size() != 1U || (size[0] != UNKNOWN_DIM && size[0] != 1))) {
59- return GRAPH_FAILED;
60- }
61- 
62- const int64_t rows = splits.size() == 1U && splits[0] >= 2 ? splits[0] - 1 : UNKNOWN_DIM;
63- int64_t bins = UNKNOWN_DIM;
64- if (ReadLegacyBinCount(op, bins) != GRAPH_SUCCESS) {
65- return GRAPH_FAILED;
66- }
67- auto output = op.GetOutputDescByName("output");
68- output.SetShape(Shape({rows, bins}));
69- // Keep the legacy IR dtype surface; native 950 dtype checks remain in the
70- // existing InferDataType and tiling callbacks.
71- output.SetDataType(op.GetInputDescByName("weights").GetDataType());
72- if (rows == UNKNOWN_DIM || bins == UNKNOWN_DIM) {
73- const auto rowRange = rows == UNKNOWN_DIM ? std::make_pair(int64_t{0}, UNKNOWN_DIM) :
74- std::make_pair(rows, rows);
75- const auto binRange = bins == UNKNOWN_DIM ? std::make_pair(int64_t{0}, UNKNOWN_DIM) :
76- std::make_pair(bins, bins);
77- if (output.SetShapeRange({rowRange, binRange}) != GRAPH_SUCCESS) {
78- return GRAPH_FAILED;
79- }
80- }
81- return op.UpdateOutputDesc("output", output);
82-}
83- 
84-// GEIR can use the legacy registry before the native InferShape callback.
85-// Its built-in implementation resizes weights to the rank of values, losing
86-// dimensions (and elements). Infer only the output and preserve every input.
87-COMMON_INFER_FUNC_REG(RaggedBinCount, RaggedBinCountLegacyInferShape);
88-} // namespace ge
89- 
90namespace ops {18namespace ops {
91static ge::graphStatus InferDataTypeForRaggedBinCount(gert::InferDataTypeContext* context)19static ge::graphStatus InferDataTypeForRaggedBinCount(gert::InferDataTypeContext* context)
92{20{
Dmath/ragged_bin_count/tests/ut/op_graph/test_ragged_bin_count_legacy_infer_pass.cpp+0-168
@@ -1,168 +0,0 @@
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-#include <cstdint>
12-#include <vector>
13- 
14-#include "gtest/gtest.h"
15-#include "../../../op_graph/ragged_bin_count_proto.h"
16- 
17-namespace {
18-ge::TensorDesc MakeDesc(const std::vector<int64_t>& dims, ge::DataType dtype)
19-{
20- ge::TensorDesc desc(ge::Shape(dims), ge::FORMAT_ND, dtype);
21- desc.SetOriginShape(ge::Shape(dims));
22- desc.SetOriginFormat(ge::FORMAT_ND);
23- return desc;
24-}
25- 
26-class SizeConstant : public ge::Operator {
27-public:
28- explicit SizeConstant(const ge::Tensor& value) : ge::Operator("size_constant", "Const")
29- {
30- OutputRegister("y");
31- UpdateOutputDesc("y", value.GetTensorDesc());
32- SetAttr("value", value);
33- }
34-};
35- 
36-ge::op::RaggedBinCount MakeOp(const std::vector<int64_t>& values, const std::vector<int64_t>& weights,
37- ge::DataType indexType = ge::DT_INT32, const std::vector<int64_t>& splits = {2})
38-{
39- ge::op::RaggedBinCount op("ragged_bin_count_legacy_infer");
40- EXPECT_EQ(op.update_input_desc_splits(MakeDesc(splits, ge::DT_INT64)), ge::GRAPH_SUCCESS);
41- EXPECT_EQ(op.update_input_desc_values(MakeDesc(values, indexType)), ge::GRAPH_SUCCESS);
42- EXPECT_EQ(op.update_input_desc_size(MakeDesc({1}, indexType)), ge::GRAPH_SUCCESS);
43- EXPECT_EQ(op.update_input_desc_weights(MakeDesc(weights, ge::DT_FLOAT)), ge::GRAPH_SUCCESS);
44- op.update_output_desc_output(ge::TensorDesc(ge::Shape(), ge::FORMAT_ND, ge::DT_FLOAT));
45- EXPECT_EQ(op.get_input_desc_splits().GetShape().GetDims(), splits);
46- EXPECT_EQ(op.get_input_desc_values().GetShape().GetDims(), values);
47- EXPECT_EQ(op.get_input_desc_size().GetShape().GetDims(), (std::vector<int64_t>{1}));
48- EXPECT_EQ(op.get_input_desc_weights().GetShape().GetDims(), weights);
49- return op;
50-}
51- 
52-void CheckPreservedShapes(const std::vector<int64_t>& values, const std::vector<int64_t>& weights)
53-{
54- auto op = MakeOp(values, weights);
55- ASSERT_EQ(op.InferShapeAndType(), ge::GRAPH_SUCCESS);
56- EXPECT_EQ(op.get_input_desc_values().GetShape().GetDims(), values);
57- EXPECT_EQ(op.get_input_desc_weights().GetShape().GetDims(), weights);
58- EXPECT_EQ(op.get_output_desc_output().GetShape().GetDims(), (std::vector<int64_t>{1, -1}));
59- EXPECT_EQ(op.get_output_desc_output().GetDataType(), ge::DT_FLOAT);
60-}
61- 
62-template <typename T>
63-void CheckConstantSize(ge::DataType indexType, T bins, ge::graphStatus expected)
64-{
65- auto op = MakeOp({6}, {2, 3}, indexType);
66- const ge::TensorDesc desc = MakeDesc({1}, indexType);
67- const ge::Tensor tensor(desc, reinterpret_cast<const uint8_t*>(&bins), sizeof(bins));
68- SizeConstant constant(tensor);
69- op.set_input_size(constant, "y");
70- ASSERT_EQ(op.InferShapeAndType(), expected);
71- if (expected == ge::GRAPH_SUCCESS) {
72- EXPECT_EQ(op.get_output_desc_output().GetShape().GetDims(),
73- (std::vector<int64_t>{1, static_cast<int64_t>(bins)}));
74- EXPECT_EQ(op.get_input_desc_weights().GetShape().GetDims(), (std::vector<int64_t>{2, 3}));
75- }
76-}
77-} // namespace
78- 
79-TEST(RaggedBinCountLegacyInferTest, PreservesTwoDimensionalWeightsForVectorValues)
80-{
81- CheckPreservedShapes({6}, {2, 3});
82-}
83- 
84-TEST(RaggedBinCountLegacyInferTest, PreservesVectorWeightsForTwoDimensionalValues)
85-{
86- CheckPreservedShapes({2, 3}, {6});
87-}
88- 
89-TEST(RaggedBinCountLegacyInferTest, PreservesDifferentTwoDimensionalShapes) { CheckPreservedShapes({2, 3}, {3, 2}); }
90- 
91-TEST(RaggedBinCountLegacyInferTest, PreservesEmptyWeights)
92-{
93- for (const auto& weights : std::vector<std::vector<int64_t>>{{0}, {0, 3}, {2, 0}}) {
94- CheckPreservedShapes({2, 3}, weights);
95- }
96-}
97- 
98-TEST(RaggedBinCountLegacyInferTest, PreservesScalarWeights) { CheckPreservedShapes({1}, {}); }
99- 
100-TEST(RaggedBinCountLegacyInferTest, ReadsInt32ConstantSize)
101-{
102- CheckConstantSize<int32_t>(ge::DT_INT32, 7, ge::GRAPH_SUCCESS);
103-}
104- 
105-TEST(RaggedBinCountLegacyInferTest, ReadsInt64ConstantSizeWithoutTruncation)
106-{
107- CheckConstantSize<int64_t>(ge::DT_INT64, (int64_t{1} << 32) + 7, ge::GRAPH_SUCCESS);
108-}
109- 
110-TEST(RaggedBinCountLegacyInferTest, AcceptsZeroBins) { CheckConstantSize<int64_t>(ge::DT_INT64, 0, ge::GRAPH_SUCCESS); }
111- 
112-TEST(RaggedBinCountLegacyInferTest, RejectsNegativeBins)
113-{
114- CheckConstantSize<int32_t>(ge::DT_INT32, -1, ge::GRAPH_FAILED);
115- CheckConstantSize<int64_t>(ge::DT_INT64, -1, ge::GRAPH_FAILED);
116-}
117- 
118-TEST(RaggedBinCountLegacyInferTest, RetainsKnownRowsWhenSizeIsUnknown)
119-{
120- auto op = MakeOp({6}, {2, 3}, ge::DT_INT32, {4});
121- ASSERT_EQ(op.InferShapeAndType(), ge::GRAPH_SUCCESS);
122- EXPECT_EQ(op.get_output_desc_output().GetShape().GetDims(), (std::vector<int64_t>{3, -1}));
123- std::vector<std::pair<int64_t, int64_t>> ranges;
124- ASSERT_EQ(op.get_output_desc_output().GetShapeRange(ranges), ge::GRAPH_SUCCESS);
125- EXPECT_EQ(ranges, (std::vector<std::pair<int64_t, int64_t>>{{3, 3}, {0, -1}}));
126-}
127- 
128-TEST(RaggedBinCountLegacyInferTest, PreservesUnknownInputDimensions)
129-{
130- auto op = MakeOp({-1}, {-1, 3}, ge::DT_INT32, {-1});
131- ASSERT_EQ(op.InferShapeAndType(), ge::GRAPH_SUCCESS);
132- EXPECT_EQ(op.get_input_desc_weights().GetShape().GetDims(), (std::vector<int64_t>{-1, 3}));
133- EXPECT_EQ(op.get_output_desc_output().GetShape().GetDims(), (std::vector<int64_t>{-1, -1}));
134-}
135- 
136-TEST(RaggedBinCountLegacyInferTest, PreservesUnknownInputRanks)
137-{
138- auto op = MakeOp({-2}, {-2}, ge::DT_INT32, {-2});
139- ASSERT_EQ(op.InferShapeAndType(), ge::GRAPH_SUCCESS);
140- EXPECT_EQ(op.get_input_desc_weights().GetShape().GetDims(), (std::vector<int64_t>{-2}));
141- EXPECT_EQ(op.get_output_desc_output().GetShape().GetDims(), (std::vector<int64_t>{-1, -1}));
142-}
143- 
144-TEST(RaggedBinCountLegacyInferTest, PreservesLegacyOutputDataTypes)
145-{
146- for (const auto dtype : {ge::DT_INT32, ge::DT_INT64, ge::DT_FLOAT, ge::DT_DOUBLE}) {
147- auto op = MakeOp({6}, {2, 3});
148- ASSERT_EQ(op.update_input_desc_weights(MakeDesc({2, 3}, dtype)), ge::GRAPH_SUCCESS);
149- ASSERT_EQ(op.InferShapeAndType(), ge::GRAPH_SUCCESS);
150- EXPECT_EQ(op.get_output_desc_output().GetDataType(), dtype);
151- }
152-}
153- 
154-TEST(RaggedBinCountLegacyInferTest, RejectsInvalidSplitsAndSizeShapes)
155-{
156- for (const auto& splits : std::vector<std::vector<int64_t>>{{}, {1}, {2, 1}}) {
157- auto op = MakeOp({6}, {2, 3}, ge::DT_INT32, splits);
158- EXPECT_EQ(op.InferShapeAndType(), ge::GRAPH_FAILED);
159- }
160- auto op = MakeOp({6}, {2, 3});
161- op.update_input_desc_size(ge::TensorDesc(ge::Shape(), ge::FORMAT_ND, ge::DT_INT32));
162- EXPECT_EQ(op.InferShapeAndType(), ge::GRAPH_FAILED);
163- for (const int64_t sizeElements : {int64_t{1} << 31, int64_t{1} << 32}) {
164- auto oversizedSizeOp = MakeOp({6}, {2, 3});
165- ASSERT_EQ(oversizedSizeOp.update_input_desc_size(MakeDesc({sizeElements}, ge::DT_INT32)), ge::GRAPH_SUCCESS);
166- EXPECT_EQ(oversizedSizeOp.InferShapeAndType(), ge::GRAPH_FAILED);
167- }
168-}
Amath/ragged_bin_count/tests/ut/op_host/test_ragged_bin_count_registry_infershape.cpp+140-0
@@ -0,0 +1,140 @@
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+#include <gtest/gtest.h>
12+#include "base/context_builder/op_infer_shape_context_builder.h"
13+#include "base/registry/op_impl_space_registry_v2.h"
14+#include "exe_graph/runtime/runtime_tensor.h"
15+ 
16+namespace {
17+gert::StorageShape MakeStorageShape(const gert::Shape& shape)
18+{
19+ gert::StorageShape storage;
20+ storage.MutableOriginShape() = shape;
21+ storage.MutableStorageShape() = shape;
22+ return storage;
23+}
24+ 
25+void CheckRegisteredInfer(const gert::Shape& valuesShape, const gert::Shape& weightsShape, const gert::Shape& expected,
26+ const gert::Shape& splitsShape = {2}, const gert::Shape& sizeShape = {1},
27+ ge::DataType indexType = ge::DT_INT32, void* sizeData = nullptr,
28+ ge::graphStatus status = ge::GRAPH_SUCCESS)
29+{
30+ auto registry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
31+ ASSERT_NE(registry, nullptr);
32+ const auto* impl = registry->GetOpImpl("RaggedBinCount");
33+ ASSERT_NE(impl, nullptr);
34+ ASSERT_NE(impl->infer_shape, nullptr);
35+ EXPECT_TRUE(impl->IsInputDataDependency(2U));
36+ const gert::StorageFormat format(ge::FORMAT_ND, ge::FORMAT_ND, {});
37+ gert::Tensor splits(MakeStorageShape(splitsShape), format, ge::DT_INT64);
38+ gert::Tensor values(MakeStorageShape(valuesShape), format, indexType);
39+ gert::Tensor size(MakeStorageShape(sizeShape), format, gert::TensorPlacement::kOnHost, indexType, sizeData);
40+ gert::Tensor weights(MakeStorageShape(weightsShape), format, ge::DT_FLOAT);
41+ gert::OpInferShapeContextBuilder builder;
42+ builder.OpType("RaggedBinCount").OpName("ragged_bin_count_registered").IONum(4, 1);
43+ builder.InputTensors({&splits, &values, &size, &weights});
44+ builder.OutputTensorDesc(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
45+ auto holder = builder.Build();
46+ ASSERT_NE(holder.GetContext(), nullptr);
47+ auto* context = holder.GetContext();
48+ ASSERT_NE(context->GetOutputShape(0), nullptr);
49+ *context->GetOutputShape(0) = gert::Shape({91, 92, 93});
50+ ASSERT_EQ(impl->infer_shape(context), status);
51+ EXPECT_EQ(*context->GetInputShape(0), splitsShape);
52+ EXPECT_EQ(*context->GetInputShape(1), valuesShape);
53+ EXPECT_EQ(*context->GetInputShape(2), sizeShape);
54+ EXPECT_EQ(*context->GetInputShape(3), weightsShape);
55+ if (status == ge::GRAPH_SUCCESS) {
56+ EXPECT_EQ(*context->GetOutputShape(0), expected);
57+ ASSERT_EQ(impl->infer_shape(context), ge::GRAPH_SUCCESS);
58+ EXPECT_EQ(*context->GetOutputShape(0), expected);
59+ EXPECT_EQ(*context->GetInputShape(1), valuesShape);
60+ EXPECT_EQ(*context->GetInputShape(3), weightsShape);
61+ }
62+}
63+} // namespace
64+ 
65+TEST(RaggedBinCountRegisteredInfer, PreservesMatrixWeightsForVectorValues)
66+{
67+ CheckRegisteredInfer({6}, {2, 3}, {1, -1});
68+}
69+ 
70+TEST(RaggedBinCountRegisteredInfer, PreservesVectorWeightsForMatrixValues)
71+{
72+ CheckRegisteredInfer({2, 3}, {6}, {1, -1});
73+}
74+ 
75+TEST(RaggedBinCountRegisteredInfer, PreservesDifferentMatrixShapes) { CheckRegisteredInfer({2, 3}, {3, 2}, {1, -1}); }
76+ 
77+TEST(RaggedBinCountRegisteredInfer, PreservesEmptyWeights)
78+{
79+ CheckRegisteredInfer({2, 3}, {0}, {1, -1});
80+ CheckRegisteredInfer({2, 3}, {0, 3}, {1, -1});
81+ CheckRegisteredInfer({2, 3}, {2, 0}, {1, -1});
82+}
83+ 
84+TEST(RaggedBinCountRegisteredInfer, PreservesScalarWeights) { CheckRegisteredInfer({1}, {}, {1, -1}); }
85+ 
86+TEST(RaggedBinCountRegisteredInfer, ReadsInt32ConstantSize)
87+{
88+ int32_t bins = 7;
89+ CheckRegisteredInfer({6}, {2, 3}, {1, 7}, {2}, {1}, ge::DT_INT32, &bins);
90+}
91+ 
92+TEST(RaggedBinCountRegisteredInfer, ReadsInt64ConstantSizeWithoutTruncation)
93+{
94+ int64_t bins = (int64_t{1} << 32) + 7;
95+ CheckRegisteredInfer({6}, {2, 3}, {1, bins}, {2}, {1}, ge::DT_INT64, &bins);
96+}
97+ 
98+TEST(RaggedBinCountRegisteredInfer, AcceptsZeroBins)
99+{
100+ int64_t bins = 0;
101+ CheckRegisteredInfer({6}, {2, 3}, {1, 0}, {2}, {1}, ge::DT_INT64, &bins);
102+}
103+ 
104+TEST(RaggedBinCountRegisteredInfer, RejectsNegativeBins)
105+{
106+ int32_t bins32 = -1;
107+ int64_t bins64 = -1;
108+ CheckRegisteredInfer({6}, {2, 3}, {}, {2}, {1}, ge::DT_INT32, &bins32, ge::GRAPH_FAILED);
109+ CheckRegisteredInfer({6}, {2, 3}, {}, {2}, {1}, ge::DT_INT64, &bins64, ge::GRAPH_FAILED);
110+}
111+ 
112+TEST(RaggedBinCountRegisteredInfer, RetainsKnownRowsWhenSizeIsUnknown)
113+{
114+ CheckRegisteredInfer({6}, {2, 3}, {3, -1}, {4});
115+}
116+ 
117+TEST(RaggedBinCountRegisteredInfer, PreservesUnknownInputDimensions)
118+{
119+ CheckRegisteredInfer({-1}, {-1, 3}, {-1, -1}, {-1});
120+}
121+ 
122+TEST(RaggedBinCountRegisteredInfer, PreservesUnknownInputRanks)
123+{
124+ CheckRegisteredInfer({-2}, {-2}, {-1, -1}, {-2}, {-2});
125+}
126+ 
127+TEST(RaggedBinCountRegisteredInfer, RejectsInvalidSplitsAndSizeShapes)
128+{
129+ CheckRegisteredInfer({6}, {2, 3}, {}, {}, {1}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
130+ CheckRegisteredInfer({6}, {2, 3}, {}, {1}, {1}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
131+ CheckRegisteredInfer({6}, {2, 3}, {}, {2, 1}, {1}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
132+ CheckRegisteredInfer({6}, {2, 3}, {}, {2}, {}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
133+ CheckRegisteredInfer({6}, {2, 3}, {}, {2}, {int64_t{1} << 31}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
134+ CheckRegisteredInfer({6}, {2, 3}, {}, {2}, {int64_t{1} << 32}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
135+}
136+ 
137+TEST(RaggedBinCountRegisteredInfer, RejectsMismatchedNonemptyElementCounts)
138+{
139+ CheckRegisteredInfer({6}, {2, 2}, {}, {2}, {1}, ge::DT_INT32, nullptr, ge::GRAPH_FAILED);
140+}