Merged
feat(cdist_grad): 补充 InferDataType RT2.0 推导并新增 UT 看护 #7246
feat(cdist_grad): 补充 InferDataType RT2.0 推导并新增 UT 看护 #7246
Merged
fitZepHYrAdd on 6 days ago
Total: 2 files changed+123-3
@@ -10,11 +10,13 @@
10 10 
11/*!11/*!
12 * \file cdist_grad_infershape.cpp12 * \file cdist_grad_infershape.cpp
13- * \brief InferShape for CdistGrad operator13+ * \brief InferShape and InferDataType for CdistGrad operator
14 *14 *
15 * CdistGrad reduces along axis -2 (without keepdims).15 * CdistGrad reduces along axis -2 (without keepdims).
16 * Input shape: [..., P, R, M] (after broadcast)16 * Input shape: [..., P, R, M] (after broadcast)
17 * Output shape: [..., P, M] (R dimension is reduced away)17 * Output shape: [..., P, M] (R dimension is reduced away)
18+ *
19+ * Output dtype is the same as input grad (all inputs share the same dtype).
18 */20 */
19 21 
20#include "log/log.h"22#include "log/log.h"
@@ -24,6 +26,8 @@ using namespace ge;
24 26 
25namespace ops {27namespace ops {
26constexpr int64_t NUMBER_TWO = 2;28constexpr int64_t NUMBER_TWO = 2;
29+constexpr size_t CDIST_GRAD_INPUT_GRAD_INDEX = 0;
30+constexpr size_t CDIST_GRAD_OUTPUT_Y_INDEX = 0;
27 31 
28static ge::graphStatus InferShape4CdistGrad(gert::InferShapeContext* context)32static ge::graphStatus InferShape4CdistGrad(gert::InferShapeContext* context)
29{33{
@@ -68,5 +72,15 @@ static ge::graphStatus InferShape4CdistGrad(gert::InferShapeContext* context)
68 return GRAPH_SUCCESS;72 return GRAPH_SUCCESS;
69}73}
70 74 
71-IMPL_OP_INFERSHAPE(CdistGrad).InferShape(InferShape4CdistGrad);75+static ge::graphStatus InferDataType4CdistGrad(gert::InferDataTypeContext* context)
76+{
77+ OP_LOGD(context->GetNodeName(), "Begin to do InferDataType4CdistGrad");
78+ // grad、x1、x2、cdist 要求同 dtype,输出 y 的 dtype 与 grad 一致
79+ const auto gradDtype = context->GetInputDataType(CDIST_GRAD_INPUT_GRAD_INDEX);
80+ context->SetOutputDataType(CDIST_GRAD_OUTPUT_Y_INDEX, gradDtype);
81+ OP_LOGD(context->GetNodeName(), "End to do InferDataType4CdistGrad");
82+ return GRAPH_SUCCESS;
83+}
84+ 
85+IMPL_OP_INFERSHAPE(CdistGrad).InferShape(InferShape4CdistGrad).InferDataType(InferDataType4CdistGrad);
72} // namespace ops86} // namespace ops
@@ -10,13 +10,15 @@
10 10 
11/*!11/*!
12 * \file test_cdist_grad_infershape.cpp12 * \file test_cdist_grad_infershape.cpp
13- * \brief13+ * \brief CdistGrad InferShape and InferDataType UT.
14 */14 */
15 15 
16#include <gtest/gtest.h>16#include <gtest/gtest.h>
17#include <iostream>17#include <iostream>
18#include "infershape_context_faker.h"18#include "infershape_context_faker.h"
19#include "infershape_case_executor.h"19#include "infershape_case_executor.h"
20+#include "op_infer_datatype_context_builder.h"
21+#include "base/registry/op_impl_space_registry_v2.h"
20 22 
21class CdistGradInfershape : public testing::Test {23class CdistGradInfershape : public testing::Test {
22protected:24protected:
@@ -158,3 +160,107 @@ TEST_F(CdistGradInfershape, cdist_grad_infershape_1d_failed_test)
158 std::vector<std::vector<int64_t>> expectOutputShape = {};160 std::vector<std::vector<int64_t>> expectOutputShape = {};
159 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);161 ExecuteTestCase(infershapeContextPara, ge::GRAPH_FAILED, expectOutputShape);
160}162}
163+ 
164+// InferDataType:输出 dtype 与 grad 一致(float32)
165+TEST_F(CdistGradInfershape, cdist_grad_infer_datatype_float)
166+{
167+ auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
168+ ASSERT_NE(spaceRegistry, nullptr);
169+ auto opImpl = spaceRegistry->GetOpImpl("CdistGrad");
170+ ASSERT_NE(opImpl, nullptr);
171+ ASSERT_NE(opImpl->infer_datatype, nullptr);
172+ 
173+ gert::OpInferDataTypeContextBuilder builder;
174+ builder.OpType("CdistGrad").OpName("CdistGrad");
175+ builder.IONum(4, 1);
176+ builder.InputTensorDesc(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
177+ builder.InputTensorDesc(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
178+ builder.InputTensorDesc(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
179+ builder.InputTensorDesc(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
180+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
181+ auto contextHolder = builder.Build();
182+ auto* context = contextHolder.GetContext();
183+ ASSERT_NE(context, nullptr);
184+ 
185+ auto ret = opImpl->infer_datatype(context);
186+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
187+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
188+}
189+ 
190+// InferDataType:输出 dtype 与 grad 一致(float16)
191+TEST_F(CdistGradInfershape, cdist_grad_infer_datatype_float16)
192+{
193+ auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
194+ ASSERT_NE(spaceRegistry, nullptr);
195+ auto opImpl = spaceRegistry->GetOpImpl("CdistGrad");
196+ ASSERT_NE(opImpl, nullptr);
197+ ASSERT_NE(opImpl->infer_datatype, nullptr);
198+ 
199+ gert::OpInferDataTypeContextBuilder builder;
200+ builder.OpType("CdistGrad").OpName("CdistGrad");
201+ builder.IONum(4, 1);
202+ builder.InputTensorDesc(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
203+ builder.InputTensorDesc(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
204+ builder.InputTensorDesc(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
205+ builder.InputTensorDesc(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
206+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
207+ auto contextHolder = builder.Build();
208+ auto* context = contextHolder.GetContext();
209+ ASSERT_NE(context, nullptr);
210+ 
211+ auto ret = opImpl->infer_datatype(context);
212+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
213+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT16);
214+}
215+ 
216+// InferDataType:输出 dtype 与 grad 一致(bfloat16)
217+TEST_F(CdistGradInfershape, cdist_grad_infer_datatype_bf16)
218+{
219+ auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
220+ ASSERT_NE(spaceRegistry, nullptr);
221+ auto opImpl = spaceRegistry->GetOpImpl("CdistGrad");
222+ ASSERT_NE(opImpl, nullptr);
223+ ASSERT_NE(opImpl->infer_datatype, nullptr);
224+ 
225+ gert::OpInferDataTypeContextBuilder builder;
226+ builder.OpType("CdistGrad").OpName("CdistGrad");
227+ builder.IONum(4, 1);
228+ builder.InputTensorDesc(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
229+ builder.InputTensorDesc(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
230+ builder.InputTensorDesc(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
231+ builder.InputTensorDesc(3, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND);
232+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
233+ auto contextHolder = builder.Build();
234+ auto* context = contextHolder.GetContext();
235+ ASSERT_NE(context, nullptr);
236+ 
237+ auto ret = opImpl->infer_datatype(context);
238+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
239+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_BF16);
240+}
241+ 
242+// InferDataType:输入 dtype 不一致时,输出 dtype 仍取 grad(输入 0)
243+TEST_F(CdistGradInfershape, cdist_grad_infer_datatype_follows_grad)
244+{
245+ auto spaceRegistry = gert::DefaultOpImplSpaceRegistryV2::GetInstance().GetSpaceRegistry();
246+ ASSERT_NE(spaceRegistry, nullptr);
247+ auto opImpl = spaceRegistry->GetOpImpl("CdistGrad");
248+ ASSERT_NE(opImpl, nullptr);
249+ ASSERT_NE(opImpl->infer_datatype, nullptr);
250+ 
251+ gert::OpInferDataTypeContextBuilder builder;
252+ builder.OpType("CdistGrad").OpName("CdistGrad");
253+ builder.IONum(4, 1);
254+ builder.InputTensorDesc(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND);
255+ builder.InputTensorDesc(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
256+ builder.InputTensorDesc(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
257+ builder.InputTensorDesc(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND);
258+ builder.OutputTensorDesc(0, ge::FORMAT_ND, ge::FORMAT_ND);
259+ auto contextHolder = builder.Build();
260+ auto* context = contextHolder.GetContext();
261+ ASSERT_NE(context, nullptr);
262+ 
263+ auto ret = opImpl->infer_datatype(context);
264+ EXPECT_EQ(ret, ge::GRAPH_SUCCESS);
265+ EXPECT_EQ(context->GetOutputDataType(0), ge::DT_FLOAT);
266+}