fix(cv-ubfuse): host UB 约束按原始 tmp 需求计算, 修复 matmul+pow UB 越界-graph-autofusion-AtomGit

已合并

fix(cv-ubfuse): host UB 约束按原始 tmp 需求计算, 修复 matmul+pow UB 越界 #2297

tianbaolin创建于 9月22日从tianbaolin/graph-autofusion： tbl/fix/cv_ub_out_of_boundry 合入到 cann/graph-autofusion： develop

fix(cv-ubfuse): host UB 约束按原始 tmp 需求计算, 修复 matmul+pow UB 越界 #2297

已合并

tianbaolin创建于 9月22日

从tianbaolin/graph-autofusion tbl/fix/cv_ub_out_of_boundry合入到cann/graph-autofusion develop

共 5 个文件变更+65-9

Mautofuse/att/gen_model_info/gen_model_info.cpp+4-4

	@@ -175,10 +175,10 @@ void ApplyCommonUbExprContext(const ascir::UbExprContext &context, ModelInfo &mo
175	}	175	}
176	}	176	}
177		177
178	-af::Status RefreshCommonUbExprContext(const af::AscGraph &graph, ModelInfo &model_info) {	178	+af::Status RefreshCommonUbExprContext(const af::AscGraph &graph, ModelInfo &model_info, bool is_cv_ub_fusion) {
179	ascir::UbExprContext context;	179	ascir::UbExprContext context;
180	- GE_ASSERT_SUCCESS(ascir::AscGraphUbExprBuilder().Build(graph, context), "Build common UB expr failed, graph:[%s].",	180	+ GE_ASSERT_SUCCESS(ascir::AscGraphUbExprBuilder().Build(graph, context, is_cv_ub_fusion),
181	- graph.GetName().c_str());	181	+ "Build common UB expr failed, graph:[%s].", graph.GetName().c_str());
182	ApplyCommonUbExprContext(context, model_info);	182	ApplyCommonUbExprContext(context, model_info);
183	return af::SUCCESS;	183	return af::SUCCESS;
184	}	184	}
	@@ -201,7 +201,7 @@ static af::Status GenerateSingleModelInfoWithContext(const af::AscGraph &graph,
201	// step2: get basic expr constraint	201	// step2: get basic expr constraint
202	att::GenerateTilingExpr tiling_expr(tuning_space);	202	att::GenerateTilingExpr tiling_expr(tuning_space);
203	GE_ASSERT_SUCCESS(tiling_expr.Generate(model_info), "Get basic expr constraint failed.");	203	GE_ASSERT_SUCCESS(tiling_expr.Generate(model_info), "Get basic expr constraint failed.");
204	- GE_ASSERT_SUCCESS(RefreshCommonUbExprContext(graph, model_info), "Refresh common UB expr failed.");	204	+ GE_ASSERT_SUCCESS(RefreshCommonUbExprContext(graph, model_info, is_cv_ub_fusion), "Refresh common UB expr failed.");
205	GELOGI("[DFX]End to generate model info for graph %s of tiling case id %u", graph.GetName().c_str(), tiling_case_id);	205	GELOGI("[DFX]End to generate model info for graph %s of tiling case id %u", graph.GetName().c_str(), tiling_case_id);
206	return af::SUCCESS;	206	return af::SUCCESS;
207	}	207	}

Mautofuse/common/ub_expr/asc_graph_ub_expr_builder.cpp+29-4

	@@ -24,6 +24,27 @@ namespace ascir {
24	namespace {	24	namespace {
25	constexpr int64_t kMinTmpBufferSize = 8 * 1024;	25	constexpr int64_t kMinTmpBufferSize = 8 * 1024;
26	constexpr int64_t kBlockAlignBytes = 32;	26	constexpr int64_t kBlockAlignBytes = 32;
		27	+// GetTmpBuffer 会为 tmp 公式套上 v1 repeat cap 壳: Min(size, 255 * 256 + 32)。
		28	+// CV UBFuse 的 epilogue 内联公式按原始需求分配 tmp, host UB 约束须与之一致, 否则
		29	+// 求解器会选出使 tmp 越界的 stage 尺寸。数值须与 default_reg_func.cpp 的 cap 保持一致。
		30	+constexpr int64_t kMaxV1RepeatTmpBytes = 255LL * 256LL + 32LL;
		31	+
		32	+UbExpr StripV1RepeatCap(UbExpr size) {
		33	+ if (!size.IsValid()) {
		34	+ return size;
		35	+ }
		36	+ auto args = size.GetArgs();
		37	+ if (args.size() != 2UL) {
		38	+ return size;
		39	+ }
		40	+ for (size_t i = 0UL; i < args.size(); i++) {
		41	+ int64_t cap_value = 0;
		42	+ if (args[i].IsConstExpr() && args[i].GetConstValue(cap_value) && cap_value == kMaxV1RepeatTmpBytes) {
		43	+ return args[1UL - i];
		44	+ }
		45	+ }
		46	+ return size;
		47	+}
27		48
28	std::string MakeQueueName(int64_t id) {	49	std::string MakeQueueName(int64_t id) {
29	return "q" + std::to_string(id) + "_size";	50	return "q" + std::to_string(id) + "_size";
	@@ -266,11 +287,13 @@ void FillTileVars(const af::AscGraph &graph, UbExprContext &context) {
266	}	287	}
267	}	288	}
268		289
269	-void AddTmpBuffer(const af::TmpBuffer &tmp_buffer, std::map<int64_t, UbExpr> &node_tmp_buffer_bytes) {	290	+void AddTmpBuffer(const af::TmpBuffer &tmp_buffer, bool use_raw_tmp_size,
		291	+ std::map<int64_t, UbExpr> &node_tmp_buffer_bytes) {
270	if (!IsUbTmpBuffer(tmp_buffer) \|\| !HasPrintableExpr(tmp_buffer.buf_desc.size)) {	292	if (!IsUbTmpBuffer(tmp_buffer) \|\| !HasPrintableExpr(tmp_buffer.buf_desc.size)) {
271	return;	293	return;
272	}	294	}
273	- AppendAdd(node_tmp_buffer_bytes[tmp_buffer.id], tmp_buffer.buf_desc.size);	295	+ const UbExpr size = use_raw_tmp_size ? StripV1RepeatCap(tmp_buffer.buf_desc.size) : tmp_buffer.buf_desc.size;
		296	+ AppendAdd(node_tmp_buffer_bytes[tmp_buffer.id], size);
274	}	297	}
275		298
276	void MergeNodeTmpBuffers(const std::map<int64_t, UbExpr> &node_tmp_buffer_bytes,	299	void MergeNodeTmpBuffers(const std::map<int64_t, UbExpr> &node_tmp_buffer_bytes,
	@@ -296,10 +319,12 @@ void AddTensor(const af::AscGraph &graph, const af::AscTensorAttr &attr, std::ma
296		319
297	} // namespace	320	} // namespace
298		321
299	-af::Status AscGraphUbExprBuilder::Build(const af::AscGraph &graph, UbExprContext &context) const {	322	+af::Status AscGraphUbExprBuilder::Build(const af::AscGraph &graph, UbExprContext &context,
		323	+ bool use_raw_tmp_size) const {
300	context = UbExprContext{};	324	context = UbExprContext{};
301	context.graph_name = graph.GetName();	325	context.graph_name = graph.GetName();
302	context.tiling_case_id = graph.GetTilingKey();	326	context.tiling_case_id = graph.GetTilingKey();
		327	+ context.use_raw_tmp_size = use_raw_tmp_size;
303	FillSizeVars(graph, context);	328	FillSizeVars(graph, context);
304	FillTileVars(graph, context);	329	FillTileVars(graph, context);
305		330
	@@ -315,7 +340,7 @@ af::Status AscGraphUbExprBuilder::Build(const af::AscGraph &graph, UbExprContext
315	}	340	}
316	std::map<int64_t, UbExpr> node_tmp_buffer_bytes;	341	std::map<int64_t, UbExpr> node_tmp_buffer_bytes;
317	for (const auto &tmp_buffer : node->attr.tmp_buffers) {	342	for (const auto &tmp_buffer : node->attr.tmp_buffers) {
318	- AddTmpBuffer(tmp_buffer, node_tmp_buffer_bytes);	343	+ AddTmpBuffer(tmp_buffer, use_raw_tmp_size, node_tmp_buffer_bytes);
319	}	344	}
320	MergeNodeTmpBuffers(node_tmp_buffer_bytes, buffer_bytes);	345	MergeNodeTmpBuffers(node_tmp_buffer_bytes, buffer_bytes);
321	}	346	}

Mautofuse/common/ub_expr/asc_graph_ub_expr_builder.h+1-1

	@@ -18,7 +18,7 @@ namespace ascir {
18		18
19	class AscGraphUbExprBuilder {	19	class AscGraphUbExprBuilder {
20	public:	20	public:
21	- af::Status Build(const af::AscGraph &graph, UbExprContext &context) const;	21	+ af::Status Build(const af::AscGraph &graph, UbExprContext &context, bool use_raw_tmp_size = false) const;
22	};	22	};
23		23
24	} // namespace ascir	24	} // namespace ascir

Mautofuse/common/ub_expr/ub_expr_types.h+3-0

	@@ -50,6 +50,9 @@ struct UbExprContext {
50	std::string graph_name;	50	std::string graph_name;
51	std::string template_name;	51	std::string template_name;
52	int64_t tiling_case_id = 0;	52	int64_t tiling_case_id = 0;
		53	+ // CV UBFuse 的 epilogue 内联公式不走 tiling 求解, tmp 需求按原始公式分配(无 v1 repeat cap),
		54	+ // UB 约束须与 device 侧分配保持一致, 否则求解器会选出使 tmp 越界的 stage 尺寸。
		55	+ bool use_raw_tmp_size = false;
53	};	56	};
54		57
55	struct UbExprBuildResult {	58	struct UbExprBuildResult {

Mautofuse/tests/ut/common/test_ub_expr_utils.cpp+28-0

	@@ -123,6 +123,34 @@ TEST(AscGraphUbExprBuilderTest, BuildAggregatesQueueBufferAndTmpBuffer) {
123	EXPECT_EQ(context.graph_name, "ub_alloc");	123	EXPECT_EQ(context.graph_name, "ub_alloc");
124	}	124	}
125		125
		126	+TEST(AscGraphUbExprBuilderTest, BuildStripsV1RepeatCapForRawTmpSize) {
		127	+ af::AscGraph graph("raw_tmp");
		128	+ auto node = BuildLoadNode(graph, "load");
		129	+ ASSERT_NE(node, nullptr);
		130	+
		131	+ const auto raw_tmp_size = af::Symbol("raw_tmp_size");
		132	+ af::TmpBuffer tmp_buffer;
		133	+ tmp_buffer.id = 1;
		134	+ tmp_buffer.buf_desc.size = af::sym::Min(raw_tmp_size, af::Symbol(255 * 256 + 32));
		135	+ tmp_buffer.mem.alloc_type = af::AllocType::kAllocTypeBuffer;
		136	+ tmp_buffer.mem.hardware = af::MemHardware::kMemHardwareUB;
		137	+ node->attr.tmp_buffers.emplace_back(tmp_buffer);
		138	+
		139	+ ascir::UbExprContext capped_context;
		140	+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, capped_context), af::SUCCESS);
		141	+ auto capped_expr = ReplaceContainers(capped_context).Replace({{raw_tmp_size, af::Symbol(80000)}}).Simplify();
		142	+ int64_t capped_usage = 0;
		143	+ EXPECT_TRUE(capped_expr.GetConstValue(capped_usage)) << capped_expr.Str().get();
		144	+ EXPECT_EQ(capped_usage, 65312);
		145	+
		146	+ ascir::UbExprContext raw_context;
		147	+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, raw_context, true), af::SUCCESS);
		148	+ auto raw_expr = ReplaceContainers(raw_context).Replace({{raw_tmp_size, af::Symbol(80000)}}).Simplify();
		149	+ int64_t raw_usage = 0;
		150	+ EXPECT_TRUE(raw_expr.GetConstValue(raw_usage)) << raw_expr.Str().get();
		151	+ EXPECT_EQ(raw_usage, 80000);
		152	+}
		153	+
126	TEST(AscGraphUbExprBuilderTest, BuildAlignsTensorBytesAfterDtypeSize) {	154	TEST(AscGraphUbExprBuilderTest, BuildAlignsTensorBytesAfterDtypeSize) {
127	af::AscGraph graph("fp16_align");	155	af::AscGraph graph("fp16_align");
128	auto &axis = graph.CreateAxis("s0", af::Symbol(33));	156	auto &axis = graph.CreateAxis("s0", af::Symbol(33));