已合并
test: 补充reg_func算子UT覆盖缺失分支(#210) #1953
test: 补充reg_func算子UT覆盖缺失分支(#210) #1953
已合并
Leechi666创建于 9月3日
共 5 个文件变更+356-0
@@ -841,5 +841,51 @@ TEST_F(CalcCompareTmpSizeTest, CalcEQTmpSize_INT64_ShouldReturnCorrectSize_WhenN
841 ASSERT_EQ(result[0]->size, sym::Min(af::Symbol(512) * s0 + af::Symbol(256), MAX_TMP_BUFFER_SIZE));841 ASSERT_EQ(result[0]->size, sym::Min(af::Symbol(512) * s0 + af::Symbol(256), MAX_TMP_BUFFER_SIZE));
842 ASSERT_EQ(result[0]->life_time_axis_id, -1);842 ASSERT_EQ(result[0]->life_time_axis_id, -1);
843}843}
844+ 
845+// Test: two vectorized axes with BF16 dtype hits the default-size fallback branch
846+TEST_F(CalcCompareTmpSizeTest, CalcGeTmpSize_ShouldReturnDefaultSize_WhenTwoAxesAndBf16Dtype) {
847+ af::AscGraph graph("test");
848+ auto s0 = graph.CreateSizeVar("s0");
849+ auto s1 = graph.CreateSizeVar("s1");
850+ 
851+ auto z0 = graph.CreateAxis("z0", s0);
852+ auto z1 = graph.CreateAxis("z1", s1);
853+ 
854+ af::ascir_op::Data x1("x1", graph);
855+ af::ascir_op::Load load1("load1");
856+ af::ascir_op::Ge ge("compare");
857+ af::ascir_op::Store store("store");
858+ af::ascir_op::Output y("y");
859+ 
860+ x1.attr.sched.axis = {z0.id, z1.id};
861+ x1.y.dtype = af::DT_BF16;
862+ *x1.y.axis = {z0.id, z1.id};
863+ *x1.y.repeats = {s0, s1};
864+ *x1.y.strides = {s1, Symbol(1)};
865+ 
866+ load1.x = x1.y;
867+ load1.attr.sched.axis = {z0.id, z1.id};
868+ load1.y.dtype = af::DT_BF16;
869+ *load1.y.axis = {z0.id, z1.id};
870+ *load1.y.repeats = {s0, s1};
871+ *load1.y.strides = {s1, Symbol(1)};
872+ *load1.y.vectorized_axis = {z0.id, z1.id};
873+ 
874+ ge.x1 = load1.y;
875+ ge.x2 = load1.y;
876+ ge.attr.sched.axis = {z0.id, z1.id};
877+ ge.y.dtype = af::DT_BF16;
878+ *ge.y.axis = {z0.id, z1.id};
879+ *ge.y.repeats = {s0, s1};
880+ *ge.y.strides = {s1, Symbol(1)};
881+ 
882+ std::shared_ptr<af::AscNode> node = graph.FindNode("compare");
883+ node->inputs[0].attr.vectorized_strides = {s1, Symbol(1)};
884+ std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcGeTmpSize(*node);
885+ ASSERT_EQ(result.size(), 1);
886+ ASSERT_EQ(result[0]->size, af::Symbol(8192));
887+ ASSERT_EQ(result[0]->life_time_axis_id, -1);
888+}
889+ 
844} // namespace ascir890} // namespace ascir
845} // namespace af891} // namespace af
@@ -179,5 +179,156 @@ TEST_F(CalcReduceTmpSizeTest, CalcReduceTmpSize_test_1) {
179 std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcReduceTmpSize(*node);179 std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcReduceTmpSize(*node);
180 ASSERT_EQ(result.size(), 2);180 ASSERT_EQ(result.size(), 2);
181}181}
182+ 
183+// 复制自 CreateGraphReduce,算子换成 ArgMax(覆盖 CalcReduceTmpSize 的 ArgMax 分支)
184+template <af::DataType T>
185+void CreateGraphArgMax(af::AscGraph &graph, Expression &s1, Expression &s2) {
186+ af::Expression One = af::Symbol(1);
187+ af::Expression Zero = af::Symbol(0);
188+ auto s0 = graph.CreateSizeVar("s0");
189+ s1 = graph.CreateSizeVar("s1");
190+ s2 = graph.CreateSizeVar("s2");
191+ 
192+ auto z0 = graph.CreateAxis("z0", s0);
193+ auto z1 = graph.CreateAxis("z1", s1);
194+ auto z2 = graph.CreateAxis("z2", s2);
195+ 
196+ Data x1("x1", graph);
197+ Load load1("load1");
198+ af::ascir_op::ArgMax argmax0("argmax0");
199+ Store store("store");
200+ Output y("y");
201+ 
202+ x1.attr.sched.axis = {z0.id, z1.id, z2.id};
203+ x1.y.dtype = T;
204+ *x1.y.axis = {z0.id, z1.id, z2.id};
205+ *x1.y.repeats = {s0, s1, s2};
206+ *x1.y.strides = {s1 * s2, s2, One};
207+ 
208+ load1.x = x1.y;
209+ load1.attr.sched.axis = {z0.id, z1.id, z2.id};
210+ load1.y.dtype = T;
211+ *load1.y.axis = {z0.id, z1.id, z2.id};
212+ *load1.y.repeats = {s0, s1, s2};
213+ *load1.y.strides = {s1 * s2, s2, One};
214+ *load1.y.vectorized_axis = {z1.id, z2.id};
215+ 
216+ argmax0.x = load1.y;
217+ argmax0.attr.sched.axis = {z0.id, z1.id, z2.id};
218+ argmax0.attr.sched.loop_axis = {z0.id};
219+ argmax0.y.dtype = T;
220+ *argmax0.y.axis = {z0.id, z1.id, z2.id};
221+ *argmax0.y.repeats = {s0, s1, One};
222+ *argmax0.y.strides = {s2, One, Zero};
223+ *argmax0.y.vectorized_axis = {z1.id, z2.id};
224+ 
225+ store.x = argmax0.y;
226+ store.attr.sched.axis = {z0.id, z1.id, z2.id};
227+ store.y.dtype = T;
228+ *store.y.axis = {z0.id, z1.id, z2.id};
229+ *store.y.repeats = {s0, s1, s2};
230+ *store.y.strides = {s1 * s2, s2, One};
231+ 
232+ y.x = store.y;
233+ y.attr.sched.axis = {z0.id, z1.id, z2.id};
234+ y.y.dtype = T;
235+ *y.y.axis = {z0.id, z1.id, z2.id};
236+ *y.y.repeats = {s0, s1, s2};
237+ *y.y.strides = {s1 * s2, s2, One};
238+}
239+ 
240+// 复制自 CreateGraphReduce,算子换成 ArgMaxMultiRPhase1(覆盖对应分支)
241+template <af::DataType T>
242+void CreateGraphArgMaxPhase1(af::AscGraph &graph, Expression &s1, Expression &s2) {
243+ af::Expression One = af::Symbol(1);
244+ af::Expression Zero = af::Symbol(0);
245+ auto s0 = graph.CreateSizeVar("s0");
246+ s1 = graph.CreateSizeVar("s1");
247+ s2 = graph.CreateSizeVar("s2");
248+ 
249+ auto z0 = graph.CreateAxis("z0", s0);
250+ auto z1 = graph.CreateAxis("z1", s1);
251+ auto z2 = graph.CreateAxis("z2", s2);
252+ 
253+ Data x1("x1", graph);
254+ Load load1("load1");
255+ af::ascir_op::ArgMaxMultiRPhase1 argmaxp1("argmaxp1");
256+ Store store("store");
257+ Output y("y");
258+ 
259+ x1.attr.sched.axis = {z0.id, z1.id, z2.id};
260+ x1.y.dtype = T;
261+ *x1.y.axis = {z0.id, z1.id, z2.id};
262+ *x1.y.repeats = {s0, s1, s2};
263+ *x1.y.strides = {s1 * s2, s2, One};
264+ 
265+ load1.x = x1.y;
266+ load1.attr.sched.axis = {z0.id, z1.id, z2.id};
267+ load1.y.dtype = T;
268+ *load1.y.axis = {z0.id, z1.id, z2.id};
269+ *load1.y.repeats = {s0, s1, s2};
270+ *load1.y.strides = {s1 * s2, s2, One};
271+ *load1.y.vectorized_axis = {z1.id, z2.id};
272+ 
273+ argmaxp1.x = load1.y;
274+ argmaxp1.attr.sched.axis = {z0.id, z1.id, z2.id};
275+ argmaxp1.attr.sched.loop_axis = {z0.id};
276+ argmaxp1.value.dtype = T;
277+ *argmaxp1.value.axis = {z0.id, z1.id, z2.id};
278+ *argmaxp1.value.repeats = {s0, s1, One};
279+ *argmaxp1.value.strides = {s2, One, Zero};
280+ *argmaxp1.value.vectorized_axis = {z1.id, z2.id};
281+ 
282+ store.x = argmaxp1.value;
283+ store.attr.sched.axis = {z0.id, z1.id, z2.id};
284+ store.y.dtype = T;
285+ *store.y.axis = {z0.id, z1.id, z2.id};
286+ *store.y.repeats = {s0, s1, s2};
287+ *store.y.strides = {s1 * s2, s2, One};
288+ 
289+ y.x = store.y;
290+ y.attr.sched.axis = {z0.id, z1.id, z2.id};
291+ y.y.dtype = T;
292+ *y.y.axis = {z0.id, z1.id, z2.id};
293+ *y.y.repeats = {s0, s1, s2};
294+ *y.y.strides = {s1 * s2, s2, One};
295+}
296+ 
297+// ArgMax 型:返回 4 个 tmp buf,生命周期序列 -1,0,1,2 是该分支独有指纹
298+TEST_F(CalcReduceTmpSizeTest, CalcReduceTmpSize_WhenNodeTypeIsArgMax) {
299+ af::AscGraph graph("test_argmax");
300+ af::Expression One = af::Symbol(1);
301+ af::Expression Zero = af::Symbol(0);
302+ Expression s1;
303+ Expression s2;
304+ CreateGraphArgMax<af::DT_FLOAT>(graph, s1, s2);
305+ std::shared_ptr<af::AscNode> node = graph.FindNode("argmax0");
306+ node->inputs[0].attr.vectorized_strides = {s2, One};
307+ node->outputs[0].attr.vectorized_strides = {One, Zero};
308+ std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcReduceTmpSize(*node);
309+ ASSERT_EQ(result.size(), 4U);
310+ ASSERT_EQ(result[0]->life_time_axis_id, -1);
311+ ASSERT_EQ(result[1]->life_time_axis_id, 0);
312+ ASSERT_EQ(result[2]->life_time_axis_id, 1);
313+ ASSERT_EQ(result[3]->life_time_axis_id, 2);
314+}
315+ 
316+// ArgMaxMultiRPhase1 型:返回 3 个 tmp buf,生命周期序列 -1,0,1
317+TEST_F(CalcReduceTmpSizeTest, CalcReduceTmpSize_WhenNodeTypeIsArgMaxMultiRPhase1) {
318+ af::AscGraph graph("test_argmax_p1");
319+ af::Expression One = af::Symbol(1);
320+ af::Expression Zero = af::Symbol(0);
321+ Expression s1;
322+ Expression s2;
323+ CreateGraphArgMaxPhase1<af::DT_FLOAT>(graph, s1, s2);
324+ std::shared_ptr<af::AscNode> node = graph.FindNode("argmaxp1");
325+ node->inputs[0].attr.vectorized_strides = {s2, One};
326+ node->outputs[0].attr.vectorized_strides = {One, Zero};
327+ std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcReduceTmpSize(*node);
328+ ASSERT_EQ(result.size(), 3U);
329+ ASSERT_EQ(result[0]->life_time_axis_id, -1);
330+ ASSERT_EQ(result[1]->life_time_axis_id, 0);
331+ ASSERT_EQ(result[2]->life_time_axis_id, 1);
332+}
182} // namespace ascir333} // namespace ascir
183} // namespace af334} // namespace af
@@ -143,6 +143,68 @@ TEST_F(CalcRemainderTmpSizeTest, CalcRemainderTmpSize_ShouldReturnCorrectSize_Wh
143 ASSERT_EQ(result[0]->size, af::Symbol(8192));143 ASSERT_EQ(result[0]->size, af::Symbol(8192));
144 ASSERT_EQ(result[0]->life_time_axis_id, -1);144 ASSERT_EQ(result[0]->life_time_axis_id, -1);
145}145}
146+// Test: Tensor input (second input repeats is a real tensor, not all 1)
147+TEST_F(CalcRemainderTmpSizeTest, CalcRemainderTmpSize_ShouldReturnCorrectSize_WhenSecondInputIsTensor) {
148+ af::AscGraph graph("test");
149+ auto s0 = graph.CreateSizeVar("s0");
150+ auto s1 = graph.CreateSizeVar("s1");
146 151 
152+ auto z0 = graph.CreateAxis("z0", s0);
153+ auto z1 = graph.CreateAxis("z1", s1);
154+ 
155+ af::ascir_op::Data x1("x1", graph);
156+ af::ascir_op::Data x2("x2", graph);
157+ af::ascir_op::Load load1("load1");
158+ af::ascir_op::Load load2("load2");
159+ af::ascir_op::Remainder remainder("remainder");
160+ af::ascir_op::Store store("store");
161+ af::ascir_op::Output y("y");
162+ 
163+ x1.attr.sched.axis = {z0.id, z1.id};
164+ x1.y.dtype = af::DT_FLOAT;
165+ *x1.y.axis = {z0.id, z1.id};
166+ *x1.y.repeats = {s0, s1};
167+ *x1.y.strides = {s1, Symbol(1)};
168+ 
169+ x2.attr.sched.axis = {z0.id, z1.id};
170+ x2.y.dtype = af::DT_FLOAT;
171+ *x2.y.axis = {z0.id, z1.id};
172+ *x2.y.repeats = {s0, s1};
173+ *x2.y.strides = {s1, Symbol(1)};
174+ 
175+ load1.x = x1.y;
176+ load1.attr.sched.axis = {z0.id, z1.id};
177+ load1.y.dtype = af::DT_FLOAT;
178+ *load1.y.axis = {z0.id, z1.id};
179+ *load1.y.repeats = {s0, s1};
180+ *load1.y.strides = {s1, Symbol(1)};
181+ *load1.y.vectorized_axis = {z0.id, z1.id};
182+ 
183+ load2.x = x2.y;
184+ load2.attr.sched.axis = {z0.id, z1.id};
185+ load2.y.dtype = af::DT_FLOAT;
186+ *load2.y.axis = {z0.id, z1.id};
187+ *load2.y.repeats = {s0, s1};
188+ *load2.y.strides = {s1, Symbol(1)};
189+ *load2.y.vectorized_axis = {z0.id, z1.id};
190+ 
191+ remainder.x1 = load1.y;
192+ remainder.x2 = load2.y;
193+ remainder.attr.sched.axis = {z0.id, z1.id};
194+ remainder.y.dtype = af::DT_FLOAT;
195+ *remainder.y.axis = {z0.id, z1.id};
196+ *remainder.y.repeats = {s0, s1};
197+ *remainder.y.strides = {s1, Symbol(1)};
198+ 
199+ std::shared_ptr<af::AscNode> node = graph.FindNode("remainder");
200+ // vectorized_strides must be set manually: GetInputSize relies on it to locate the vectorized axis
201+ node->inputs[0].attr.vectorized_strides = {s1, Symbol(1)};
202+ 
203+ std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcRemainderTmpSize(*node);
204+ ASSERT_EQ(result.size(), 1);
205+ ASSERT_EQ(result[0]->life_time_axis_id, -1);
206+ // tensor input: 3 intermediate buffers * sizeof(float) * input elements (s0*s1), capped by GetTmpBuffer
207+ ASSERT_EQ(result[0]->size, sym::Min(Symbol(65312), Symbol(12) * s0 * s1));
208+}
147} // namespace ascir209} // namespace ascir
148} // namespace af210} // namespace af
@@ -95,5 +95,41 @@ TEST_F(CalcRsqrtTmpSizeTest, CalcRsqrtTmpSize_ShouldReturnCorrectSize_WhenNodels
95 ASSERT_EQ(result[0]->size, sym::Min(af::Symbol(4) * s1 * s0, MAX_TMP_BUFFER_SIZE));95 ASSERT_EQ(result[0]->size, sym::Min(af::Symbol(4) * s1 * s0, MAX_TMP_BUFFER_SIZE));
96 ASSERT_EQ(result[0]->life_time_axis_id, -1);96 ASSERT_EQ(result[0]->life_time_axis_id, -1);
97}97}
98+ 
99+// Test: all-scalar input -> GetInputDataSizeTmpBuffer falls back to the last input index
100+TEST_F(CalcRsqrtTmpSizeTest, CalcRsqrtTmpSize_ShouldReturnTypeSize_WhenInputIsAllScalar) {
101+ af::AscGraph graph("test");
102+ auto s0 = graph.CreateSizeVar("s0");
103+ auto s1 = graph.CreateSizeVar("s1");
104+ 
105+ auto z0 = graph.CreateAxis("z0", s0);
106+ auto z1 = graph.CreateAxis("z1", s1);
107+ 
108+ af::ascir_op::Scalar x1("x1", graph);
109+ af::ascir_op::Rsqrt rsqrt("rsqrt");
110+ af::ascir_op::Store store("store");
111+ af::ascir_op::Output y("y");
112+ 
113+ x1.attr.sched.axis = {z0.id, z1.id};
114+ x1.y.dtype = af::DT_FLOAT;
115+ *x1.y.axis = {};
116+ *x1.y.repeats = {};
117+ *x1.y.strides = {};
118+ 
119+ rsqrt.x = x1.y;
120+ rsqrt.attr.sched.axis = {z0.id, z1.id};
121+ rsqrt.y.dtype = af::DT_FLOAT;
122+ *rsqrt.y.axis = {z0.id, z1.id};
123+ *rsqrt.y.repeats = {s0, s1};
124+ *rsqrt.y.strides = {s1, Symbol(1)};
125+ 
126+ std::shared_ptr<af::AscNode> node = graph.FindNode("rsqrt");
127+ std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcRsqrtTmpSize(*node);
128+ ASSERT_EQ(result.size(), 1);
129+ // all-scalar input: input_size=1, sizeof(float)=4 -> tmp size 4 (capped by GetTmpBuffer)
130+ ASSERT_EQ(result[0]->size, af::Symbol(4));
131+ ASSERT_EQ(result[0]->life_time_axis_id, -1);
132+}
133+ 
98} // namespace ascir134} // namespace ascir
99} // namespace af135} // namespace af
@@ -175,5 +175,66 @@ TEST_F(CalcTrueDivTmpSizeTest, CalcTrueDivTmpSize_ShouldReturnCorrectSize_WhenIn
175 ASSERT_EQ(result[0]->size, af::Symbol(8192));175 ASSERT_EQ(result[0]->size, af::Symbol(8192));
176 ASSERT_EQ(result[0]->life_time_axis_id, -1);176 ASSERT_EQ(result[0]->life_time_axis_id, -1);
177}177}
178+ 
179+// Test: second input is ub_scalar (repeats all 1) -> HasScalarOrUbScalar is true -> default tmp size
180+TEST_F(CalcTrueDivTmpSizeTest, CalcTrueDivTmpSize_ShouldReturnDefaultSize_WhenSecondInputIsUbScalar) {
181+ af::AscGraph graph("test");
182+ auto s0 = graph.CreateSizeVar("s0");
183+ auto s1 = graph.CreateSizeVar("s1");
184+ 
185+ auto z0 = graph.CreateAxis("z0", s0);
186+ auto z1 = graph.CreateAxis("z1", s1);
187+ 
188+ af::ascir_op::Data x1("x1", graph);
189+ af::ascir_op::Data x2("x2", graph);
190+ af::ascir_op::Load load1("load1");
191+ af::ascir_op::Load load2("load2");
192+ af::ascir_op::TrueDiv true_div("true_div");
193+ af::ascir_op::Store store("store");
194+ af::ascir_op::Output y("y");
195+ 
196+ x1.attr.sched.axis = {z0.id, z1.id};
197+ x1.y.dtype = af::DT_FLOAT;
198+ *x1.y.axis = {z0.id, z1.id};
199+ *x1.y.repeats = {s0, s1};
200+ *x1.y.strides = {s1, Symbol(1)};
201+ 
202+ x2.attr.sched.axis = {z0.id, z1.id};
203+ x2.y.dtype = af::DT_FLOAT;
204+ *x2.y.axis = {z0.id, z1.id};
205+ *x2.y.repeats = {Symbol(1), Symbol(1)};
206+ *x2.y.strides = {Symbol(1), Symbol(1)};
207+ 
208+ load1.x = x1.y;
209+ load1.attr.sched.axis = {z0.id, z1.id};
210+ load1.y.dtype = af::DT_FLOAT;
211+ *load1.y.axis = {z0.id, z1.id};
212+ *load1.y.repeats = {s0, s1};
213+ *load1.y.strides = {s1, Symbol(1)};
214+ *load1.y.vectorized_axis = {z0.id, z1.id};
215+ 
216+ load2.x = x2.y;
217+ load2.attr.sched.axis = {z0.id, z1.id};
218+ load2.y.dtype = af::DT_FLOAT;
219+ *load2.y.axis = {z0.id, z1.id};
220+ *load2.y.repeats = {Symbol(1), Symbol(1)};
221+ *load2.y.strides = {Symbol(1), Symbol(1)};
222+ *load2.y.vectorized_axis = {z0.id, z1.id};
223+ 
224+ true_div.x1 = load1.y;
225+ true_div.x2 = load2.y;
226+ true_div.attr.sched.axis = {z0.id, z1.id};
227+ true_div.y.dtype = af::DT_FLOAT;
228+ *true_div.y.axis = {z0.id, z1.id};
229+ *true_div.y.repeats = {s0, s1};
230+ *true_div.y.strides = {s1, Symbol(1)};
231+ 
232+ std::shared_ptr<af::AscNode> node = graph.FindNode("true_div");
233+ std::vector<std::unique_ptr<af::TmpBufDesc>> result = CalcTrueDivTmpSize(*node);
234+ ASSERT_EQ(result.size(), 1);
235+ ASSERT_EQ(result[0]->size, af::Symbol(8192));
236+ ASSERT_EQ(result[0]->life_time_axis_id, -1);
237+}
238+ 
178} // namespace ascir239} // namespace ascir
179} // namespace af240} // namespace af