Star674
261
代码
代码
Issues51
Pull Requests118
Actions
项目讨论
Wiki
分析
Star674
261
已合并
feat: 支持静态Shape根据ub大小过滤非法template #964
zhang_shengjie创建于 6月21日
feat: 支持静态Shape根据ub大小过滤非法template #964
已合并
zhang_shengjie创建于 6月21日
从zhang_shengjie/graph-autofusionfeat/static-ub-template-filter合入到cann/graph-autofusiondevelop
共 34 个文件变更+3148-797
Mautofuse/att/gen_model_info/gen_model_info.cpp+74-0
@@ -10,7 +10,10 @@
10 10 
11#include "gen_model_info.h"11#include "gen_model_info.h"
12#include <fstream>12#include <fstream>
13+#include <set>
13#include "common/util/mem_utils.h"14#include "common/util/mem_utils.h"
15+#include "common/ub_expr/asc_graph_ub_expr_builder.h"
16+#include "common/ub_expr/ub_expr_utils.h"
14#include "common/autofuse_config/auto_fuse_config_utils.h"17#include "common/autofuse_config/auto_fuse_config_utils.h"
15#include "graph/utils/file_utils.h"18#include "graph/utils/file_utils.h"
16#include "expr_gen/generate_tiling_expr.h"19#include "expr_gen/generate_tiling_expr.h"
@@ -37,6 +40,76 @@ constexpr uint32_t kConstType = 1U;
37constexpr uint32_t kVarType = 2U;40constexpr uint32_t kVarType = 2U;
38constexpr uint32_t kDefaultAlignValue = 1U;41constexpr uint32_t kDefaultAlignValue = 1U;
39const std::string kModelInfoFilePath = "./";42const std::string kModelInfoFilePath = "./";
43+ 
44+std::set<std::string> GetUbContainerNames(const ModelInfo &model_info) {
45+ std::set<std::string> names;
46+ const auto ub_iter = model_info.hardware_cons.find(HardwareDef::UB);
47+ if (ub_iter == model_info.hardware_cons.cend()) {
48+ return names;
49+ }
50+ for (const auto &symbol : ub_iter->second.FreeSymbols()) {
51+ const auto name_iter = model_info.variable_name_map.find(symbol);
52+ const auto name = name_iter == model_info.variable_name_map.cend() ? Str(symbol) : name_iter->second;
53+ if (model_info.container_exprs.find(name) != model_info.container_exprs.cend()) {
54+ names.insert(name);
55+ }
56+ }
57+ return names;
58+}
59+ 
60+void EraseVariableByName(ExprExprMap &variable_expr_map, std::map<Expr, std::string, ExprCmp> &variable_name_map,
61+ const std::string &name) {
62+ for (auto iter = variable_name_map.begin(); iter != variable_name_map.end();) {
63+ if (iter->second == name || Str(iter->first) == name) {
64+ variable_expr_map.erase(iter->first);
65+ iter = variable_name_map.erase(iter);
66+ } else {
67+ ++iter;
68+ }
69+ }
70+ for (auto iter = variable_expr_map.begin(); iter != variable_expr_map.end();) {
71+ if (Str(iter->first) == name) {
72+ iter = variable_expr_map.erase(iter);
73+ } else {
74+ ++iter;
75+ }
76+ }
77+}
78+ 
79+void EraseContainerExprs(const std::set<std::string> &container_names, ModelInfo &model_info) {
80+ for (const auto &name : container_names) {
81+ model_info.container_exprs.erase(name);
82+ EraseVariableByName(model_info.variable_expr_map, model_info.variable_name_map, name);
83+ }
84+}
85+ 
86+void ApplyCommonUbExprContext(const ascir::UbExprContext &context, ModelInfo &model_info) {
87+ const auto old_ub_container_names = GetUbContainerNames(model_info);
88+ const auto ub_expr_result = ascir::UbExprUtils::BuildUbExpr(context);
89+ if (ub_expr_result.has_ub_expr) {
90+ model_info.hardware_cons[HardwareDef::UB] = ub_expr_result.ub_expr;
91+ } else {
92+ model_info.hardware_cons.erase(HardwareDef::UB);
93+ }
94+ EraseContainerExprs(old_ub_container_names, model_info);
95+ for (const auto &container : context.container_names) {
96+ const auto expr_iter = context.container_expr.find(container.first);
97+ if (expr_iter == context.container_expr.cend()) {
98+ continue;
99+ }
100+ model_info.container_exprs[container.second] = expr_iter->second;
101+ model_info.variable_expr_map[container.first] = expr_iter->second;
102+ model_info.variable_name_map[container.first] = container.second;
103+ }
104+}
105+ 
106+af::Status RefreshCommonUbExprContext(const af::AscGraph &graph, ModelInfo &model_info) {
107+ ascir::UbExprContext context;
108+ GE_ASSERT_SUCCESS(ascir::AscGraphUbExprBuilder().Build(graph, context), "Build common UB expr failed, graph:[%s].",
109+ graph.GetName().c_str());
110+ ApplyCommonUbExprContext(context, model_info);
111+ return af::SUCCESS;
112+}
40} // namespace113} // namespace
41 114 
42af::Status GenerateModelInfo(const af::AscGraph &graph, ModelInfo &model_info, TuningSpacePtr &tuning_space,115af::Status GenerateModelInfo(const af::AscGraph &graph, ModelInfo &model_info, TuningSpacePtr &tuning_space,
@@ -52,6 +125,7 @@ af::Status GenerateModelInfo(const af::AscGraph &graph, ModelInfo &model_info, T
52 // step2: get basic expr constraint125 // step2: get basic expr constraint
53 att::GenerateTilingExpr tiling_expr(tuning_space);126 att::GenerateTilingExpr tiling_expr(tuning_space);
54 GE_ASSERT_SUCCESS(tiling_expr.Generate(model_info), "Get basic expr constraint failed.");127 GE_ASSERT_SUCCESS(tiling_expr.Generate(model_info), "Get basic expr constraint failed.");
128+ GE_ASSERT_SUCCESS(RefreshCommonUbExprContext(graph, model_info), "Refresh common UB expr failed.");
55 // step3: call passes to get configs129 // step3: call passes to get configs
56 ATTConfig att_config;130 ATTConfig att_config;
57 std::vector<PassFunc> pass_funcs;131 std::vector<PassFunc> pass_funcs;
Mautofuse/att/generator/solver_pass_gen/axes_reorder_solver/axes_reorder_solver_gen.cpp+10-1
@@ -14,6 +14,7 @@
14#include "graph/symbolizer/symbolic_utils.h"14#include "graph/symbolizer/symbolic_utils.h"
15#include "common_utils.h"15#include "common_utils.h"
16#include "named_origin_buf_expr_generator.h"16#include "named_origin_buf_expr_generator.h"
17+#include "ub_named_expr_builder.h"
17 18 
18using namespace ascgen_utils;19using namespace ascgen_utils;
19 20 
@@ -403,7 +404,15 @@ std::string AxesReorderSolverGen::GenGetUbSizeStaticFunc() {
403 bool ub_exist = false;404 bool ub_exist = false;
404 auto ub_iter = hardware_use_map_.find(HardwareDef::UB);405 auto ub_iter = hardware_use_map_.find(HardwareDef::UB);
405 if (ub_iter != hardware_use_map_.end()) {406 if (ub_iter != hardware_use_map_.end()) {
406- auto tmp_func_pair = GenNamedOriginBufExpr(ub_iter->second, " ");407+ ascir::UbExprContext context;
408+ context.ub_expr = ub_iter->second;
409+ for (const auto &item : container_expr_) {
410+ context.container_expr[item.first] = item.second;
411+ }
412+ for (const auto &item : container_names_) {
413+ context.container_names[item.first] = item.second;
414+ }
415+ auto tmp_func_pair = BuildNamedUbExpr(context, " ");
407 std::string tmp_def = tmp_func_pair.first;416 std::string tmp_def = tmp_func_pair.first;
408 std::string func_return_expr = tmp_func_pair.second;417 std::string func_return_expr = tmp_func_pair.second;
409 codes += tmp_def;418 codes += tmp_def;
Aautofuse/att/generator/solver_pass_gen/axes_reorder_solver/ub_named_expr_builder.cpp+48-0
@@ -0,0 +1,48 @@
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 "ub_named_expr_builder.h"
12+ 
13+#include "common/ub_expr/ub_expr_utils.h"
14+#include "named_origin_buf_expr_generator.h"
15+ 
16+namespace att {
17+namespace {
18+ 
19+ExprExprMap ToAttExprMap(const ascir::UbExprExprMap &expr_map) {
20+ ExprExprMap result;
21+ for (const auto &item : expr_map) {
22+ result[item.first] = item.second;
23+ }
24+ return result;
25+}
26+ 
27+std::map<Expr, std::string, ExprCmp> ToAttExprNameMap(
28+ const std::map<ascir::UbExpr, std::string, ascir::UbExprCmp> &expr_names) {
29+ std::map<Expr, std::string, ExprCmp> result;
30+ for (const auto &item : expr_names) {
31+ result[item.first] = item.second;
32+ }
33+ return result;
34+}
35+ 
36+} // namespace
37+ 
38+std::pair<std::string, std::string> BuildNamedUbExpr(const ascir::UbExprContext &context, const std::string &indent) {
39+ const auto ub_expr_result = ascir::UbExprUtils::BuildUbExpr(context);
40+ if (!ub_expr_result.has_ub_expr) {
41+ return {"", ""};
42+ }
43+ const auto container_expr = ToAttExprMap(context.container_expr);
44+ const auto container_names = ToAttExprNameMap(context.container_names);
45+ return NamedOriginBufExprGenerator(container_expr, container_names).Generate(ub_expr_result.ub_expr, indent);
46+}
47+ 
48+} // namespace att
Aautofuse/att/generator/solver_pass_gen/axes_reorder_solver/ub_named_expr_builder.h+25-0
@@ -0,0 +1,25 @@
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+#ifndef ATT_UB_NAMED_EXPR_BUILDER_H_
12+#define ATT_UB_NAMED_EXPR_BUILDER_H_
13+ 
14+#include <string>
15+#include <utility>
16+ 
17+#include "common/ub_expr/ub_expr_types.h"
18+ 
19+namespace att {
20+ 
21+std::pair<std::string, std::string> BuildNamedUbExpr(const ascir::UbExprContext &context, const std::string &indent);
22+ 
23+} // namespace att
24+ 
25+#endif // ATT_UB_NAMED_EXPR_BUILDER_H_
Mautofuse/common/platform_context.cpp+31-0
@@ -49,6 +49,10 @@ void PlatformContext::SetPlatform(const std::string &platform_name) {
49 49 
50void PlatformContext::SetPlatformInfo(const PlatformInfo &platform_info) {50void PlatformContext::SetPlatformInfo(const PlatformInfo &platform_info) {
51 std::lock_guard<std::mutex> lg(mutex_);51 std::lock_guard<std::mutex> lg(mutex_);
52+ if (platform_info.ub_size > 0) {
53+ ub_size_override_ = platform_info.ub_size;
54+ has_ub_size_override_ = true;
55+ }
52 if (!platform_info.soc_ver.empty()) {56 if (!platform_info.soc_ver.empty()) {
53 platform_info_ = platform_info;57 platform_info_ = platform_info;
54 initialized_ = true;58 initialized_ = true;
@@ -57,6 +61,13 @@ void PlatformContext::SetPlatformInfo(const PlatformInfo &platform_info) {
57 }61 }
58}62}
59 63 
64+void PlatformContext::SetUbSizeOverride(int64_t ub_size) {
65+ std::lock_guard<std::mutex> lg(mutex_);
66+ ub_size_override_ = ub_size;
67+ has_ub_size_override_ = true;
68+ GELOGI("Set UB size override: ub_size=%lld", ub_size_override_);
69+}
70+ 
60af::Status PlatformContext::GetCurrentPlatformString(std::string &platform_name) {71af::Status PlatformContext::GetCurrentPlatformString(std::string &platform_name) {
61 if (!initialized_) {72 if (!initialized_) {
62 GE_ASSERT_SUCCESS(Initialize(), "Failed to init platform info with name %s.", platform_name.c_str());73 GE_ASSERT_SUCCESS(Initialize(), "Failed to init platform info with name %s.", platform_name.c_str());
@@ -108,4 +119,24 @@ af::Status PlatformContext::GetPlatformInfo(PlatformInfo &platform_info) {
108 platform_info_.aiv_num, platform_info_.ub_size);119 platform_info_.aiv_num, platform_info_.ub_size);
109 return af::SUCCESS;120 return af::SUCCESS;
110}121}
122+ 
123+bool PlatformContext::TryGetInitializedPlatformInfo(PlatformInfo &platform_info) {
124+ std::lock_guard<std::mutex> lg(mutex_);
125+ if (!initialized_) {
126+ platform_info = PlatformInfo{};
127+ return false;
128+ }
129+ platform_info = platform_info_;
130+ return true;
131+}
132+ 
133+bool PlatformContext::TryGetUbSizeOverride(int64_t &ub_size) const {
134+ std::lock_guard<std::mutex> lg(mutex_);
135+ if (!has_ub_size_override_) {
136+ ub_size = 0;
137+ return false;
138+ }
139+ ub_size = ub_size_override_;
140+ return true;
141+}
111} // namespace ge142} // namespace ge
Aautofuse/common/ub_expr/asc_graph_ub_expr_builder.cpp+333-0
@@ -0,0 +1,333 @@
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 "common/ub_expr/asc_graph_ub_expr_builder.h"
12+ 
13+#include <algorithm>
14+#include <set>
15+ 
16+#include "ascir_ops.h"
17+#include "common/checker.h"
18+#include "common_utils.h"
19+#include "graph/types.h"
20+ 
21+namespace ascir {
22+namespace {
23+constexpr int64_t kMinTmpBufferSize = 8 * 1024;
24+constexpr int64_t kSimtDcacheSize = 32 * 1024;
25+constexpr int64_t kBlockAlignBytes = 32;
26+ 
27+std::string MakeQueueName(int64_t id) {
28+ return "q" + std::to_string(id) + "_size";
29+}
30+ 
31+std::string MakeBufferName(int64_t id) {
32+ return "b" + std::to_string(id) + "_size";
33+}
34+ 
35+std::string GetQueueName(const af::AscTensorAttr &attr) {
36+ return attr.que.name.empty() ? MakeQueueName(attr.que.id) : attr.que.name;
37+}
38+ 
39+std::string GetBufferName(const af::AscTensorAttr &attr) {
40+ return attr.buf.name.empty() ? MakeBufferName(attr.buf.id) : attr.buf.name;
41+}
42+ 
43+uint32_t GetTypeSize(ge::DataType dtype) {
44+ uint32_t type_size = 0U;
45+ if (!ge::TypeUtils::GetDataTypeLength(dtype, type_size)) {
46+ return 1U;
47+ }
48+ return type_size > 0U ? type_size : 1U;
49+}
50+ 
51+UbExpr AlignToBlock(const UbExpr &size) {
52+ return af::sym::Mul(af::Symbol(kBlockAlignBytes), af::sym::Ceiling(af::sym::Div(size, af::Symbol(kBlockAlignBytes))));
53+}
54+ 
55+bool HasPrintableExpr(const UbExpr &expr) {
56+ if (!expr.IsValid()) {
57+ return false;
58+ }
59+ const auto expr_str = expr.Str();
60+ return expr_str != nullptr && expr_str[0] != '\0';
61+}
62+ 
63+UbExpr TensorVectorizedElementSize(const af::AscGraph &graph, const af::AscTensorAttr &attr) {
64+ if (attr.vectorized_axis.empty()) {
65+ return af::sym::kSymbolOne;
66+ }
67+ for (size_t i = 0UL; i < attr.vectorized_axis.size(); ++i) {
68+ if (i >= attr.vectorized_strides.size()) {
69+ continue;
70+ }
71+ const auto &stride = attr.vectorized_strides[i];
72+ if (stride == af::sym::kSymbolZero) {
73+ continue;
74+ }
75+ if (attr.axis.empty() && i < attr.repeats.size()) {
76+ return stride == af::sym::kSymbolOne ? attr.repeats[i] : attr.repeats[i] * stride;
77+ }
78+ auto axis_iter = std::find(attr.axis.cbegin(), attr.axis.cend(), attr.vectorized_axis[i]);
79+ if (axis_iter == attr.axis.cend()) {
80+ GELOGW("[AscGraphUbExprBuilder] cannot find vectorized axis %ld in tensor attr of graph %s",
81+ attr.vectorized_axis[i], graph.GetName().c_str());
82+ return {};
83+ }
84+ const auto axis_index = static_cast<size_t>(std::distance(attr.axis.cbegin(), axis_iter));
85+ if (axis_index >= attr.repeats.size()) {
86+ GELOGW("[AscGraphUbExprBuilder] vectorized axis %ld repeat index %zu is invalid in graph %s",
87+ attr.vectorized_axis[i], axis_index, graph.GetName().c_str());
88+ return {};
89+ }
90+ return stride == af::sym::kSymbolOne ? attr.repeats[axis_index] : attr.repeats[axis_index] * stride;
91+ }
92+ return af::sym::kSymbolOne;
93+}
94+ 
95+UbExpr TensorBytes(const af::AscGraph &graph, const af::AscTensorAttr &attr) {
96+ return AlignToBlock(TensorVectorizedElementSize(graph, attr) * af::Symbol(GetTypeSize(attr.dtype)));
97+}
98+ 
99+void AppendUniqueVar(const UbExpr &var, std::vector<UbExpr> &vars) {
100+ if (!HasPrintableExpr(var)) {
101+ return;
102+ }
103+ const auto iter = std::find_if(vars.cbegin(), vars.cend(), [&var](const auto &item) { return item == var; });
104+ if (iter == vars.cend()) {
105+ vars.emplace_back(var);
106+ }
107+}
108+ 
109+void CollectVars(const UbExpr &expr, std::vector<UbExpr> &vars) {
110+ if (!HasPrintableExpr(expr)) {
111+ return;
112+ }
113+ for (const auto &var : expr.FreeSymbols()) {
114+ AppendUniqueVar(var, vars);
115+ }
116+}
117+ 
118+void AppendMax(UbExpr &expr, const UbExpr &item) {
119+ if (!HasPrintableExpr(item)) {
120+ return;
121+ }
122+ expr = HasPrintableExpr(expr) ? af::sym::Max(expr, item) : item;
123+}
124+ 
125+void AppendAdd(UbExpr &expr, const UbExpr &item) {
126+ if (!HasPrintableExpr(item)) {
127+ return;
128+ }
129+ expr = HasPrintableExpr(expr) ? expr + item : item;
130+}
131+ 
132+struct ContainerState {
133+ std::string name;
134+ UbExpr normal_max;
135+ uint32_t buf_num = 0U;
136+ std::map<int64_t, UbExpr> share_group_bytes;
137+};
138+ 
139+void AddContainerTensor(const af::AscGraph &graph, const af::AscTensorAttr &attr, ContainerState &state) {
140+ const auto bytes = TensorBytes(graph, attr);
141+ AppendMax(state.normal_max, bytes);
142+ if (attr.mem.reuse_id == af::kIdNone) {
143+ return;
144+ }
145+ AppendAdd(state.share_group_bytes[attr.mem.reuse_id], bytes);
146+}
147+ 
148+void AddQueueTensor(const af::AscGraph &graph, const af::AscTensorAttr &attr, ContainerState &state) {
149+ if (state.name.empty()) {
150+ state.name = GetQueueName(attr);
151+ }
152+ AddContainerTensor(graph, attr, state);
153+ state.buf_num = std::max(state.buf_num, static_cast<uint32_t>(std::max<int64_t>(attr.que.buf_num, 0)));
154+}
155+ 
156+void AddBufferTensor(const af::AscGraph &graph, const af::AscTensorAttr &attr, ContainerState &state) {
157+ if (state.name.empty()) {
158+ state.name = GetBufferName(attr);
159+ }
160+ AddContainerTensor(graph, attr, state);
161+}
162+ 
163+UbExpr ContainerSlotBytes(const ContainerState &state) {
164+ UbExpr size = state.normal_max;
165+ for (const auto &item : state.share_group_bytes) {
166+ AppendMax(size, item.second);
167+ }
168+ return size;
169+}
170+ 
171+UbExpr QueueTotalBytes(const ContainerState &state, const UbExpr &slot_bytes) {
172+ if (!slot_bytes.IsValid()) {
173+ return {};
174+ }
175+ const uint32_t buf_num = state.buf_num == 0U ? 1U : state.buf_num;
176+ return buf_num == 1U ? slot_bytes : slot_bytes * af::Symbol(buf_num);
177+}
178+ 
179+bool IsUbAlloc(const af::AscTensorAttr &attr) {
180+ return attr.mem.hardware == af::MemHardware::kMemHardwareUB &&
181+ (attr.mem.alloc_type == af::AllocType::kAllocTypeQueue ||
182+ attr.mem.alloc_type == af::AllocType::kAllocTypeBuffer);
183+}
184+ 
185+bool IsUbTmpBuffer(const af::TmpBuffer &tmp_buffer) {
186+ return tmp_buffer.id != af::kIdNone;
187+}
188+ 
189+void AddContainer(UbExprContext &context, const std::string &name, const UbExpr &expr, const UbExpr &total_expr) {
190+ if (!HasPrintableExpr(expr) || !HasPrintableExpr(total_expr)) {
191+ return;
192+ }
193+ const UbExpr symbol = af::Symbol(name.c_str());
194+ context.container_expr[symbol] = expr;
195+ context.container_names[symbol] = name;
196+ AppendAdd(context.ub_expr, total_expr);
197+ CollectVars(expr, context.ub_related_vars);
198+}
199+ 
200+void AddContainer(UbExprContext &context, const std::string &name, const UbExpr &expr) {
201+ const UbExpr symbol = af::Symbol(name.c_str());
202+ AddContainer(context, name, expr, symbol);
203+}
204+ 
205+bool IsConstZero(const UbExpr &expr) {
206+ int64_t value = 0;
207+ return HasPrintableExpr(expr) && expr.GetConstValue(value) && value == 0;
208+}
209+ 
210+void AddBuiltinTmpBuffer(const af::AscGraph &graph, UbExprContext &context) {
211+ const auto builtin_tmp_buffer = ascgen_utils::CalcExtraTmpBufForAscGraph(graph);
212+ if (IsConstZero(builtin_tmp_buffer)) {
213+ return;
214+ }
215+ AppendAdd(context.ub_expr, builtin_tmp_buffer);
216+ CollectVars(builtin_tmp_buffer, context.ub_related_vars);
217+}
218+ 
219+UbExpr CalcReservedUbSize(const af::AscGraph &graph) {
220+ UbExpr reserved_ub_size = af::Symbol(ascgen_utils::CalcReservedTmpBufSizeForAscGraph(graph));
221+ for (const auto &node : graph.GetAllNodes()) {
222+ GE_ASSERT_NOTNULL(node);
223+ if (node->GetType() == af::ascir_op::Gather::Type) {
224+ reserved_ub_size = reserved_ub_size + af::Symbol(kSimtDcacheSize);
225+ break;
226+ }
227+ }
228+ return reserved_ub_size;
229+}
230+ 
231+void AddReservedUb(const af::AscGraph &graph, UbExprContext &context) {
232+ const auto reserved_ub_size = CalcReservedUbSize(graph);
233+ if (IsConstZero(reserved_ub_size)) {
234+ return;
235+ }
236+ AppendAdd(context.ub_expr, reserved_ub_size);
237+}
238+ 
239+bool IsDynamicSizeVar(const af::SizeVarPtr &size_var) {
240+ return size_var != nullptr && HasPrintableExpr(size_var->expr) && !size_var->expr.IsConstExpr();
241+}
242+ 
243+void FillSizeVars(const af::AscGraph &graph, UbExprContext &context) {
244+ for (const auto &size_var : graph.GetAllSizeVar()) {
245+ if (IsDynamicSizeVar(size_var)) {
246+ AppendUniqueVar(size_var->expr, context.dynamic_size_vars);
247+ }
248+ }
249+}
250+ 
251+void FillTileVars(const af::AscGraph &graph, UbExprContext &context) {
252+ for (const auto &axis : graph.GetAllAxis()) {
253+ if (axis == nullptr || !HasPrintableExpr(axis->size)) {
254+ continue;
255+ }
256+ if (axis->type == af::Axis::kAxisTypeTileInner) {
257+ context.var_min_values[axis->size] = af::sym::kSymbolOne;
258+ AppendUniqueVar(axis->size, context.ub_related_vars);
259+ continue;
260+ }
261+ if (axis->type == af::Axis::kAxisTypeOriginal && !axis->size.IsConstExpr()) {
262+ AppendUniqueVar(axis->size, context.dynamic_size_vars);
263+ }
264+ }
265+}
266+ 
267+void AddTmpBuffer(const af::TmpBuffer &tmp_buffer, std::map<int64_t, UbExpr> &node_tmp_buffer_bytes) {
268+ if (!IsUbTmpBuffer(tmp_buffer) || !HasPrintableExpr(tmp_buffer.buf_desc.size)) {
269+ return;
270+ }
271+ AppendAdd(node_tmp_buffer_bytes[tmp_buffer.id], tmp_buffer.buf_desc.size);
272+}
273+ 
274+void MergeNodeTmpBuffers(const std::map<int64_t, UbExpr> &node_tmp_buffer_bytes,
275+ std::map<int64_t, ContainerState> &buffer_bytes) {
276+ for (const auto &item : node_tmp_buffer_bytes) {
277+ AppendMax(buffer_bytes[item.first].normal_max, af::sym::Max(item.second, af::Symbol(kMinTmpBufferSize)));
278+ }
279+}
280+ 
281+void AddTensor(const af::AscGraph &graph, const af::AscTensorAttr &attr, std::map<int64_t, ContainerState> &queue_bytes,
282+ std::map<int64_t, ContainerState> &buffer_bytes) {
283+ if (!IsUbAlloc(attr)) {
284+ return;
285+ }
286+ if (attr.mem.alloc_type == af::AllocType::kAllocTypeQueue && attr.que.id != af::kIdNone) {
287+ AddQueueTensor(graph, attr, queue_bytes[attr.que.id]);
288+ return;
289+ }
290+ if (attr.mem.alloc_type == af::AllocType::kAllocTypeBuffer && attr.buf.id != af::kIdNone) {
291+ AddBufferTensor(graph, attr, buffer_bytes[attr.buf.id]);
292+ }
293+}
294+ 
295+} // namespace
296+ 
297+af::Status AscGraphUbExprBuilder::Build(const af::AscGraph &graph, UbExprContext &context) const {
298+ context = UbExprContext{};
299+ context.graph_name = graph.GetName();
300+ context.tiling_case_id = graph.GetTilingKey();
301+ FillSizeVars(graph, context);
302+ FillTileVars(graph, context);
303+ 
304+ std::map<int64_t, ContainerState> queue_bytes;
305+ std::map<int64_t, ContainerState> buffer_bytes;
306+ for (const auto &node : graph.GetAllNodes()) {
307+ GE_ASSERT_NOTNULL(node);
308+ for (const auto &output : node->outputs()) {
309+ AddTensor(graph, output->attr, queue_bytes, buffer_bytes);
310+ }
311+ std::map<int64_t, UbExpr> node_tmp_buffer_bytes;
312+ for (const auto &tmp_buffer : node->attr.tmp_buffers) {
313+ AddTmpBuffer(tmp_buffer, node_tmp_buffer_bytes);
314+ }
315+ MergeNodeTmpBuffers(node_tmp_buffer_bytes, buffer_bytes);
316+ }
317+ 
318+ for (const auto &item : queue_bytes) {
319+ const auto name = item.second.name.empty() ? MakeQueueName(item.first) : item.second.name;
320+ const auto symbol = af::Symbol(name.c_str());
321+ const auto slot_bytes = ContainerSlotBytes(item.second);
322+ AddContainer(context, name, slot_bytes, QueueTotalBytes(item.second, symbol));
323+ }
324+ for (const auto &item : buffer_bytes) {
325+ const auto name = item.second.name.empty() ? MakeBufferName(item.first) : item.second.name;
326+ AddContainer(context, name, ContainerSlotBytes(item.second));
327+ }
328+ AddBuiltinTmpBuffer(graph, context);
329+ AddReservedUb(graph, context);
330+ return af::SUCCESS;
331+}
332+ 
333+} // namespace ascir
Aautofuse/common/ub_expr/asc_graph_ub_expr_builder.h+26-0
@@ -0,0 +1,26 @@
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+#ifndef ASCGEN_DEV_BASE_COMMON_ASC_GRAPH_UB_EXPR_BUILDER_H_
12+#define ASCGEN_DEV_BASE_COMMON_ASC_GRAPH_UB_EXPR_BUILDER_H_
13+ 
14+#include "common/ub_expr/ub_expr_types.h"
15+#include "ascendc_ir/ascendc_ir_core/ascendc_ir.h"
16+ 
17+namespace ascir {
18+ 
19+class AscGraphUbExprBuilder {
20+ public:
21+ af::Status Build(const af::AscGraph &graph, UbExprContext &context) const;
22+};
23+ 
24+} // namespace ascir
25+ 
26+#endif // ASCGEN_DEV_BASE_COMMON_ASC_GRAPH_UB_EXPR_BUILDER_H_
Aautofuse/common/ub_expr/ub_expr_types.h+63-0
@@ -0,0 +1,63 @@
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+#ifndef ASCGEN_DEV_BASE_COMMON_UB_EXPR_TYPES_H_
12+#define ASCGEN_DEV_BASE_COMMON_UB_EXPR_TYPES_H_
13+ 
14+#include <map>
15+#include <string>
16+#include <utility>
17+#include <vector>
18+ 
19+#include "graph/symbolizer/symbolic.h"
20+ 
21+namespace ascir {
22+ 
23+using UbExpr = af::Expression;
24+ 
25+struct UbExprCmp {
26+ bool operator()(const UbExpr &lhs, const UbExpr &rhs) const {
27+ if (!lhs.IsValid()) {
28+ return rhs.IsValid();
29+ }
30+ if (!rhs.IsValid()) {
31+ return false;
32+ }
33+ return lhs.Compare(rhs) < 0;
34+ }
35+};
36+ 
37+using UbExprExprMap = std::map<UbExpr, UbExpr, UbExprCmp>;
38+using UbExprUintMap = std::map<UbExpr, uint32_t, UbExprCmp>;
39+ 
40+struct UbExprContext {
41+ UbExpr ub_expr;
42+ UbExprExprMap container_expr;
43+ std::map<UbExpr, std::string, UbExprCmp> container_names;
44+ std::vector<UbExpr> ub_related_vars;
45+ UbExprExprMap var_min_values;
46+ UbExprUintMap const_vars;
47+ UbExprExprMap static_size_vars;
48+ std::vector<UbExpr> dynamic_size_vars;
49+ std::vector<std::pair<UbExpr, UbExpr>> expr_relations;
50+ std::string graph_name;
51+ std::string template_name;
52+ int64_t tiling_case_id = 0;
53+};
54+ 
55+struct UbExprBuildResult {
56+ bool has_ub_expr = false;
57+ UbExpr ub_expr;
58+ std::string origin_expr;
59+};
60+ 
61+} // namespace ascir
62+ 
63+#endif // ASCGEN_DEV_BASE_COMMON_UB_EXPR_TYPES_H_
Aautofuse/common/ub_expr/ub_expr_utils.cpp+26-0
@@ -0,0 +1,26 @@
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 "common/ub_expr/ub_expr_utils.h"
12+ 
13+namespace ascir {
14+ 
15+UbExprBuildResult UbExprUtils::BuildUbExpr(const UbExprContext &context) {
16+ UbExprBuildResult result;
17+ if (!context.ub_expr.IsValid() || context.ub_expr.Str() == nullptr) {
18+ return result;
19+ }
20+ result.has_ub_expr = true;
21+ result.ub_expr = context.ub_expr;
22+ result.origin_expr = context.ub_expr.Str().get();
23+ return result;
24+}
25+ 
26+} // namespace ascir
Aautofuse/common/ub_expr/ub_expr_utils.h+25-0
@@ -0,0 +1,25 @@
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+#ifndef ASCGEN_DEV_BASE_COMMON_UB_EXPR_UTILS_H_
12+#define ASCGEN_DEV_BASE_COMMON_UB_EXPR_UTILS_H_
13+ 
14+#include "common/ub_expr/ub_expr_types.h"
15+ 
16+namespace ascir {
17+ 
18+class UbExprUtils {
19+ public:
20+ static UbExprBuildResult BuildUbExpr(const UbExprContext &context);
21+};
22+ 
23+} // namespace ascir
24+ 
25+#endif // ASCGEN_DEV_BASE_COMMON_UB_EXPR_UTILS_H_
Mautofuse/inc/common/platform_context.h+10-0
@@ -36,6 +36,8 @@ class PlatformContext {
36 36 
37 void SetPlatformInfo(const PlatformInfo &platform_info);37 void SetPlatformInfo(const PlatformInfo &platform_info);
38 38 
39+ void SetUbSizeOverride(int64_t ub_size);
40+ 
39 // 获取当前 platform 字符串41 // 获取当前 platform 字符串
40 af::Status GetCurrentPlatformString(std::string &platform_name);42 af::Status GetCurrentPlatformString(std::string &platform_name);
41 43 
@@ -44,15 +46,23 @@ class PlatformContext {
44 platform_info_.soc_ver = "";46 platform_info_.soc_ver = "";
45 platform_info_.aiv_num = 0;47 platform_info_.aiv_num = 0;
46 platform_info_.ub_size = 0;48 platform_info_.ub_size = 0;
49+ ub_size_override_ = 0;
50+ has_ub_size_override_ = false;
47 }51 }
48 52 
49 af::Status GetPlatformInfo(PlatformInfo &platform_info);53 af::Status GetPlatformInfo(PlatformInfo &platform_info);
50 54 
55+ bool TryGetInitializedPlatformInfo(PlatformInfo &platform_info);
56+ 
57+ bool TryGetUbSizeOverride(int64_t &ub_size) const;
58+ 
51 private:59 private:
52 af::Status Initialize();60 af::Status Initialize();
53 af::Status InitPlatformInfo();61 af::Status InitPlatformInfo();
54 PlatformContext() = default;62 PlatformContext() = default;
55 PlatformInfo platform_info_;63 PlatformInfo platform_info_;
64+ int64_t ub_size_override_ = 0;
65+ bool has_ub_size_override_ = false;
56 bool initialized_ = false;66 bool initialized_ = false;
57 static std::mutex mutex_;67 static std::mutex mutex_;
58};68};
Mautofuse/optimize/buffer_allocate/buf_que_allocator.cpp+80-7
@@ -160,6 +160,12 @@ Status BufQueAllocator::AllocBufQueForSingleImplGraph(af::AscGraph &impl_graph,
160}160}
161 161 
162Status BufQueAllocator::AllocBufQue(::ascir::FusedScheduledResult &fused_scheduled_result) {162Status BufQueAllocator::AllocBufQue(::ascir::FusedScheduledResult &fused_scheduled_result) {
163+ GE_CHK_STATUS_RET(PrepareImplGraphMemoryPlan(fused_scheduled_result), "PrepareImplGraphMemoryPlan failed");
164+ GE_CHK_STATUS_RET(CollectFusedIoNodes(fused_scheduled_result), "CollectFusedIoNodes failed");
165+ return ge::GRAPH_SUCCESS;
166+}
167+ 
168+Status BufQueAllocator::PrepareImplGraphMemoryPlan(::ascir::FusedScheduledResult &fused_scheduled_result) {
163 const auto &platform = PlatformFactory::GetInstance().GetPlatform();169 const auto &platform = PlatformFactory::GetInstance().GetPlatform();
164 GE_CHECK_NOTNULL(platform, "Platform is not found.");170 GE_CHECK_NOTNULL(platform, "Platform is not found.");
165 171 
@@ -179,6 +185,25 @@ Status BufQueAllocator::AllocBufQue(::ascir::FusedScheduledResult &fused_schedul
179 return ge::GRAPH_SUCCESS;185 return ge::GRAPH_SUCCESS;
180}186}
181 187 
188+Status BufQueAllocator::CollectFusedIoNodes(::ascir::FusedScheduledResult &fused_scheduled_result) {
189+ fused_scheduled_result.input_nodes.clear();
190+ fused_scheduled_result.output_nodes.clear();
191+ fused_scheduled_result.workspace_nodes.clear();
192+ node_type_to_index_to_node_.clear();
193+ for (const auto &scheduled_results : fused_scheduled_result.node_idx_to_scheduled_results) {
194+ for (const auto &result : scheduled_results) {
195+ for (const auto &schedule_group : result.schedule_groups) {
196+ for (const auto &impl_graph : schedule_group.impl_graphs) {
197+ GE_CHK_STATUS_RET(CollectIoNodes(impl_graph), "CollectIoNodes failed, graph = %s",
198+ impl_graph.GetName().c_str());
199+ }
200+ }
201+ }
202+ }
203+ AppendCollectedIoNodes(fused_scheduled_result);
204+ return ge::GRAPH_SUCCESS;
205+}
206+ 
182Status BufQueAllocator::AllocateForIoNodes(const af::AscGraph &impl_graph) {207Status BufQueAllocator::AllocateForIoNodes(const af::AscGraph &impl_graph) {
183 for (const auto &node : impl_graph.GetAllNodes()) {208 for (const auto &node : impl_graph.GetAllNodes()) {
184 GE_ASSERT_NOTNULL(node);209 GE_ASSERT_NOTNULL(node);
@@ -193,7 +218,8 @@ Status BufQueAllocator::AllocateForIoNodes(const af::AscGraph &impl_graph) {
193 tensor_id = it->second;218 tensor_id = it->second;
194 GELOGI("same index, cur_node: %s", node->GetName().c_str());219 GELOGI("same index, cur_node: %s", node->GetName().c_str());
195 auto &index_to_node = node_type_to_index_to_node_[node->GetType()];220 auto &index_to_node = node_type_to_index_to_node_[node->GetType()];
196- if (node->GetName().size() < index_to_node[index]->GetName().size()) {221+ auto node_it = index_to_node.find(index);
222+ if (node_it == index_to_node.end() || ShouldReplaceRepresentative(node_it->second, node)) {
197 index_to_node[index] = node;223 index_to_node[index] = node;
198 }224 }
199 } else {225 } else {
@@ -244,18 +270,65 @@ Status BufQueAllocator::AllocateForIoNodes(::ascir::FusedScheduledResult &fused_
244 }270 }
245 }271 }
246 }272 }
273+ return ge::GRAPH_SUCCESS;
274+}
275+ 
276+bool BufQueAllocator::ShouldReplaceRepresentative(const ascir::NodeView &current, const ascir::NodeView &candidate) {
277+ if (current == nullptr) {
278+ return true;
279+ }
280+ if (candidate == nullptr) {
281+ return false;
282+ }
283+ return candidate->GetName().size() < current->GetName().size();
284+}
285+ 
286+Status BufQueAllocator::CollectIoNodes(const af::AscGraph &impl_graph) {
287+ for (const auto &node : impl_graph.GetAllNodes()) {
288+ GE_ASSERT_NOTNULL(node);
289+ if (ScheduleUtils::IsDataInput(node) || IsOps<Output>(node)) {
290+ int64_t index = -1;
291+ GE_CHK_STATUS_RET(node->attr.ir_attr->GetAttrValue("index", index), "Get attr index failed, node = %s[%s]",
292+ node->GetNamePtr(), node->GetTypePtr());
293+ auto &index_to_node = node_type_to_index_to_node_[node->GetType()];
294+ const auto iter = index_to_node.find(index);
295+ if (iter == index_to_node.end() || ShouldReplaceRepresentative(iter->second, node)) {
296+ index_to_node[index] = node;
297+ }
298+ continue;
299+ }
300+ if (IsOps<Workspace>(node)) {
301+ const auto iter = workspace_name_to_tensor_id_.find(node->GetName());
302+ GE_ASSERT_TRUE(iter != workspace_name_to_tensor_id_.end(), "Workspace tensor id is not allocated, node = %s",
303+ node->GetNamePtr());
304+ node_type_to_index_to_node_[node->GetType()][iter->second] = node;
305+ }
306+ }
307+ return ge::GRAPH_SUCCESS;
308+}
309+ 
310+void BufQueAllocator::AppendCollectedIoNodes(::ascir::FusedScheduledResult &fused_scheduled_result) const {
247 for (const auto &type : {Data::Type, af::ascir_op::ScalarData::Type}) {311 for (const auto &type : {Data::Type, af::ascir_op::ScalarData::Type}) {
248- for (const auto &index_and_node : node_type_to_index_to_node_[type]) {312+ const auto iter = node_type_to_index_to_node_.find(type);
313+ if (iter == node_type_to_index_to_node_.cend()) {
314+ continue;
315+ }
316+ for (const auto &index_and_node : iter->second) {
249 fused_scheduled_result.input_nodes.emplace_back(index_and_node.second);317 fused_scheduled_result.input_nodes.emplace_back(index_and_node.second);
250 }318 }
251 }319 }
252- for (const auto &index_and_node : node_type_to_index_to_node_[Output::Type]) {320+ const auto output_iter = node_type_to_index_to_node_.find(Output::Type);
253- fused_scheduled_result.output_nodes.emplace_back(index_and_node.second);321+ if (output_iter != node_type_to_index_to_node_.cend()) {
322+ for (const auto &index_and_node : output_iter->second) {
323+ fused_scheduled_result.output_nodes.emplace_back(index_and_node.second);
324+ }
254 }325 }
255- for (const auto &index_and_node : node_type_to_index_to_node_[Workspace::Type]) {326+ const auto workspace_iter = node_type_to_index_to_node_.find(Workspace::Type);
256- fused_scheduled_result.workspace_nodes.emplace_back(index_and_node.second);327+ if (workspace_iter != node_type_to_index_to_node_.cend()) {
328+ for (const auto &index_and_node : workspace_iter->second) {
329+ fused_scheduled_result.workspace_nodes.emplace_back(index_and_node.second);
330+ }
257 }331 }
258- return ge::GRAPH_SUCCESS;
259}332}
260 333 
261Status BufQueAllocator::SetOutputTensorAttr(const af::AscGraph &impl_graph) const {334Status BufQueAllocator::SetOutputTensorAttr(const af::AscGraph &impl_graph) const {
Mautofuse/optimize/buffer_allocate/buf_que_allocator.h+5-0
@@ -22,6 +22,8 @@ namespace optimize {
22class BufQueAllocator {22class BufQueAllocator {
23 public:23 public:
24 Status AllocBufQue(::ascir::FusedScheduledResult &fused_scheduled_result);24 Status AllocBufQue(::ascir::FusedScheduledResult &fused_scheduled_result);
25+ Status PrepareImplGraphMemoryPlan(::ascir::FusedScheduledResult &fused_scheduled_result);
26+ Status CollectFusedIoNodes(::ascir::FusedScheduledResult &fused_scheduled_result);
25 27 
26 private:28 private:
27 Status AllocBufQueForSingleImplGraph(af::AscGraph &impl_graph, size_t max_que_num,29 Status AllocBufQueForSingleImplGraph(af::AscGraph &impl_graph, size_t max_que_num,
@@ -30,8 +32,11 @@ class BufQueAllocator {
30 bool is_reduce_mem_reuse);32 bool is_reduce_mem_reuse);
31 Status AllocateForIoNodes(::ascir::FusedScheduledResult &fused_scheduled_result);33 Status AllocateForIoNodes(::ascir::FusedScheduledResult &fused_scheduled_result);
32 Status AllocateForIoNodes(const af::AscGraph &impl_graph);34 Status AllocateForIoNodes(const af::AscGraph &impl_graph);
35+ Status CollectIoNodes(const af::AscGraph &impl_graph);
36+ void AppendCollectedIoNodes(::ascir::FusedScheduledResult &fused_scheduled_result) const;
33 Status SetOutputTensorAttr(const af::AscGraph &impl_graph) const;37 Status SetOutputTensorAttr(const af::AscGraph &impl_graph) const;
34 static void SetGlobalMemInfo(const af::AscTensor &tensor, int64_t tensor_id);38 static void SetGlobalMemInfo(const af::AscTensor &tensor, int64_t tensor_id);
39+ static bool ShouldReplaceRepresentative(const ascir::NodeView &current, const ascir::NodeView &candidate);
35 void InitTensorReuseInfoAndLifeTime(const ascir::NodeView &node, const af::AscTensor *output, TensorInfo &tensor_info,40 void InitTensorReuseInfoAndLifeTime(const ascir::NodeView &node, const af::AscTensor *output, TensorInfo &tensor_info,
36 bool is_reduce_mem_reuse, bool is_cube_none_db) const;41 bool is_reduce_mem_reuse, bool is_cube_none_db) const;
37 void InitTensorReuseInfo(const ascir::NodeView &node, const af::AscTensor *output, TensorInfo &tensor_info,42 void InitTensorReuseInfo(const ascir::NodeView &node, const af::AscTensor *output, TensorInfo &tensor_info,
Mautofuse/optimize/optimize.cpp+9-2
@@ -31,6 +31,7 @@
31#include "graph/symbolizer/symbolic_utils.h"31#include "graph/symbolizer/symbolic_utils.h"
32#include "optimize/graph_completeness/dtype_consistency.h"32#include "optimize/graph_completeness/dtype_consistency.h"
33#include "pre_process/pre_process.h"33#include "pre_process/pre_process.h"
34+#include "static_ub_template_filter.h"
34 35 
35using namespace ascir;36using namespace ascir;
36using namespace optimize;37using namespace optimize;
@@ -638,7 +639,10 @@ Status Optimizer::OptimizeFusedAscBackend(const af::ComputeGraphPtr &fused_graph
638 }639 }
639 }640 }
640 fused_scheduled_result.origin_vars.assign(original_var_set.begin(), original_var_set.end());641 fused_scheduled_result.origin_vars.assign(original_var_set.begin(), original_var_set.end());
641- GE_CHK_STATUS_RET(BufQueAllocator().AllocBufQue(fused_scheduled_result));642+ BufQueAllocator allocator;
643+ GE_CHK_STATUS_RET(allocator.PrepareImplGraphMemoryPlan(fused_scheduled_result));
644+ GE_CHK_STATUS_RET(StaticUbTemplateFilter().Filter(fused_scheduled_result));
645+ GE_CHK_STATUS_RET(allocator.CollectFusedIoNodes(fused_scheduled_result));
642 GELOGI("AllocBufQue end");646 GELOGI("AllocBufQue end");
643 TryEnableGroupParallel(fused_scheduled_result);647 TryEnableGroupParallel(fused_scheduled_result);
644 for (auto &scheduled_results : fused_scheduled_result.node_idx_to_scheduled_results) {648 for (auto &scheduled_results : fused_scheduled_result.node_idx_to_scheduled_results) {
@@ -957,11 +961,14 @@ Status Optimizer::Optimize(af::AscGraph &hint_graph, FusedScheduledResult &fused
957 GE_ASSERT_SUCCESS(OptimizeForHintGraph(hint_graph, fused_scheduled_result.node_idx_to_scheduled_results[0UL]),961 GE_ASSERT_SUCCESS(OptimizeForHintGraph(hint_graph, fused_scheduled_result.node_idx_to_scheduled_results[0UL]),
958 "Failed to optimize for graph:[%s].", hint_graph.GetName().c_str());962 "Failed to optimize for graph:[%s].", hint_graph.GetName().c_str());
959 // 内存分配963 // 内存分配
960- GE_CHK_STATUS_RET(BufQueAllocator().AllocBufQue(fused_scheduled_result));964+ BufQueAllocator allocator;
965+ GE_CHK_STATUS_RET(allocator.PrepareImplGraphMemoryPlan(fused_scheduled_result));
961 if (options_.graph_type == GraphType::kAscGraph) {966 if (options_.graph_type == GraphType::kAscGraph) {
962 fused_scheduled_result.fused_graph_name = hint_graph.GetName().c_str();967 fused_scheduled_result.fused_graph_name = hint_graph.GetName().c_str();
963 fused_scheduled_result.origin_vars.assign(original_var_set.begin(), original_var_set.end());968 fused_scheduled_result.origin_vars.assign(original_var_set.begin(), original_var_set.end());
964 }969 }
970+ GE_CHK_STATUS_RET(StaticUbTemplateFilter().Filter(fused_scheduled_result));
971+ GE_CHK_STATUS_RET(allocator.CollectFusedIoNodes(fused_scheduled_result));
965 GELOGI("AllocBufQue end");972 GELOGI("AllocBufQue end");
966 TryEnableGroupParallel(fused_scheduled_result);973 TryEnableGroupParallel(fused_scheduled_result);
967 ExecSeqAdvancedOfLoad(fused_scheduled_result);974 ExecSeqAdvancedOfLoad(fused_scheduled_result);
Aautofuse/optimize/static_ub_template_filter.cpp+416-0
@@ -0,0 +1,416 @@
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 "static_ub_template_filter.h"
12+ 
13+#include <algorithm>
14+#include <cmath>
15+#include <cstdint>
16+#include <exception>
17+#include <limits>
18+#include <string>
19+#include <utility>
20+#include <vector>
21+ 
22+#include "common/platform_context.h"
23+#include "common/ub_expr/asc_graph_ub_expr_builder.h"
24+#include "common/ub_expr/ub_expr_utils.h"
25+#include "rt_external_base.h"
26+ 
27+namespace optimize {
28+namespace {
29+constexpr const char *kLogPrefix = "[StaticUbTemplateFilter]";
30+constexpr size_t kMaxExprLogLength = 512UL;
31+constexpr uint32_t kMaxRuntimeSpecValueLen = 16U;
32+constexpr const char *kAicoreSpecLabel = "AICoreSpec";
33+constexpr const char *kUbSizeSpecKey = "ub_size";
34+ 
35+enum class EvalStatus {
36+ kKnown,
37+ kUnknown,
38+ kFailed,
39+};
40+ 
41+struct EvalResult {
42+ EvalStatus status = EvalStatus::kUnknown;
43+ int64_t min_ub_usage = 0;
44+ af::Expression origin_expr;
45+ af::Expression min_expr;
46+};
47+ 
48+struct TemplatePosition {
49+ size_t node_idx = 0UL;
50+ size_t result_idx = 0UL;
51+ size_t group_idx = 0UL;
52+ size_t impl_idx = 0UL;
53+};
54+ 
55+struct FilterState {
56+ ascir::FusedScheduledResult &fused_scheduled_result;
57+ int64_t ub_size = 0;
58+};
59+ 
60+struct UbLimitResult {
61+ bool has_ub_limit = false;
62+ int64_t ub_size = 0;
63+};
64+ 
65+std::string ExprToString(const af::Expression &expr) {
66+ if (!expr.IsValid()) {
67+ return "<invalid>";
68+ }
69+ const auto expr_str_ptr = expr.Str();
70+ if (expr_str_ptr == nullptr || expr_str_ptr.get() == nullptr) {
71+ return "<null>";
72+ }
73+ std::string expr_str = expr_str_ptr.get();
74+ if (expr_str.size() <= kMaxExprLogLength) {
75+ return expr_str;
76+ }
77+ return expr_str.substr(0UL, kMaxExprLogLength) + "...";
78+}
79+ 
80+const char *GetGraphName(const ge::AscendString &graph_name) {
81+ const auto *name = graph_name.GetString();
82+ return (name == nullptr || name[0] == '\0') ? "<null>" : name;
83+}
84+ 
85+std::string BuildTemplateName(size_t node_idx, size_t result_idx, size_t group_idx, size_t impl_idx) {
86+ return "node" + std::to_string(node_idx) + "_result" + std::to_string(result_idx) + "_group" +
87+ std::to_string(group_idx) + "_impl" + std::to_string(impl_idx);
88+}
89+ 
90+void AppendReplacement(const af::Expression &src, const af::Expression &dst,
91+ std::vector<std::pair<af::Expression, af::Expression>> &replacements) {
92+ if (!src.IsValid() || !dst.IsValid()) {
93+ return;
94+ }
95+ const auto iter =
96+ std::find_if(replacements.cbegin(), replacements.cend(), [&src](const auto &item) { return item.first == src; });
97+ if (iter == replacements.cend()) {
98+ replacements.emplace_back(src, dst);
99+ }
100+}
101+ 
102+void AppendVarOneReplacements(const std::vector<af::Expression> &vars,
103+ std::vector<std::pair<af::Expression, af::Expression>> &replacements) {
104+ for (const auto &var : vars) {
105+ AppendReplacement(var, af::sym::kSymbolOne, replacements);
106+ }
107+}
108+ 
109+void BuildReplacements(const ascir::UbExprContext &context,
110+ std::vector<std::pair<af::Expression, af::Expression>> &container_replacements,
111+ std::vector<std::pair<af::Expression, af::Expression>> &var_replacements) {
112+ for (const auto &item : context.container_expr) {
113+ AppendReplacement(item.first, item.second, container_replacements);
114+ }
115+ for (const auto &item : context.var_min_values) {
116+ AppendReplacement(item.first, item.second, var_replacements);
117+ }
118+ for (const auto &item : context.const_vars) {
119+ AppendReplacement(item.first, af::Symbol(item.second), var_replacements);
120+ }
121+ for (const auto &item : context.static_size_vars) {
122+ AppendReplacement(item.first, item.second, var_replacements);
123+ }
124+ AppendVarOneReplacements(context.dynamic_size_vars, var_replacements);
125+}
126+ 
127+af::Expression ApplyReplacements(const af::Expression &expr,
128+ const std::vector<std::pair<af::Expression, af::Expression>> &complex_replacements,
129+ const std::vector<std::pair<af::Expression, af::Expression>> &symbol_replacements) {
130+ af::Expression result = expr;
131+ if (!complex_replacements.empty()) {
132+ result = result.Replace(complex_replacements);
133+ }
134+ if (!symbol_replacements.empty()) {
135+ result = result.Subs(symbol_replacements);
136+ }
137+ if (result.IsValid()) {
138+ result = result.Simplify();
139+ }
140+ return result;
141+}
142+ 
143+bool TryGetConstInt64(const af::Expression &expr, int64_t &value) {
144+ if (!expr.IsValid() || !expr.IsConstExpr()) {
145+ return false;
146+ }
147+ if (expr.GetExprType() == af::ExprType::kExprConstantInteger) {
148+ return expr.GetConstValue(value);
149+ }
150+ if (expr.GetExprType() != af::ExprType::kExprConstantRation &&
151+ expr.GetExprType() != af::ExprType::kExprConstantRealDouble) {
152+ return false;
153+ }
154+ double double_value = 0.0;
155+ if (!expr.GetConstValue(double_value) || !std::isfinite(double_value)) {
156+ return false;
157+ }
158+ if (double_value < static_cast<double>(std::numeric_limits<int64_t>::min()) ||
159+ double_value > static_cast<double>(std::numeric_limits<int64_t>::max())) {
160+ return false;
161+ }
162+ double integer_value = 0.0;
163+ const double fraction_value = std::modf(double_value, &integer_value);
164+ if (fraction_value > 0.0 || fraction_value < 0.0) {
165+ return false;
166+ }
167+ value = static_cast<int64_t>(integer_value);
168+ return true;
169+}
170+ 
171+EvalResult EvalMinUbUsage(const ascir::UbExprContext &context) {
172+ EvalResult eval_result;
173+ const auto build_result = ascir::UbExprUtils::BuildUbExpr(context);
174+ if (!build_result.has_ub_expr) {
175+ eval_result.status = EvalStatus::kUnknown;
176+ return eval_result;
177+ }
178+ eval_result.origin_expr = build_result.ub_expr;
179+ 
180+ std::vector<std::pair<af::Expression, af::Expression>> container_replacements;
181+ std::vector<std::pair<af::Expression, af::Expression>> var_replacements;
182+ BuildReplacements(context, container_replacements, var_replacements);
183+ eval_result.min_expr = ApplyReplacements(build_result.ub_expr, container_replacements, var_replacements);
184+ if (!eval_result.min_expr.IsValid()) {
185+ eval_result.status = EvalStatus::kFailed;
186+ return eval_result;
187+ }
188+ if (!eval_result.min_expr.FreeSymbols().empty()) {
189+ eval_result.status = EvalStatus::kUnknown;
190+ return eval_result;
191+ }
192+ if (!TryGetConstInt64(eval_result.min_expr, eval_result.min_ub_usage) || eval_result.min_ub_usage < 0) {
193+ eval_result.status = EvalStatus::kFailed;
194+ return eval_result;
195+ }
196+ eval_result.status = EvalStatus::kKnown;
197+ return eval_result;
198+}
199+ 
200+EvalResult EvalGraphMinUbUsage(const af::AscGraph &graph) {
201+ ascir::UbExprContext context;
202+ try {
203+ const auto status = ascir::AscGraphUbExprBuilder().Build(graph, context);
204+ if (status != af::SUCCESS) {
205+ GELOGD("%s build UB expr failed, graph=%s", kLogPrefix, graph.GetName().c_str());
206+ EvalResult result;
207+ result.status = EvalStatus::kFailed;
208+ return result;
209+ }
210+ return EvalMinUbUsage(context);
211+ } catch (const std::exception &e) {
212+ GELOGD("%s eval UB expr failed, graph=%s, reason=%s", kLogPrefix, graph.GetName().c_str(), e.what());
213+ } catch (...) {
214+ GELOGD("%s eval UB expr failed, graph=%s, reason=unknown exception", kLogPrefix, graph.GetName().c_str());
215+ }
216+ EvalResult result;
217+ result.status = EvalStatus::kFailed;
218+ result.origin_expr = context.ub_expr;
219+ return result;
220+}
221+ 
222+bool TryParseRuntimeUbSize(const char *ub_size_str, int64_t &ub_size) {
223+ try {
224+ size_t parsed_size = 0UL;
225+ const std::string str_value(ub_size_str);
226+ ub_size = static_cast<int64_t>(std::stoll(str_value, &parsed_size));
227+ if (parsed_size == str_value.size()) {
228+ return true;
229+ }
230+ GELOGD("%s parse runtime ub_size failed, value=%s", kLogPrefix, ub_size_str);
231+ } catch (const std::exception &e) {
232+ GELOGD("%s parse runtime ub_size failed, value=%s, reason=%s", kLogPrefix, ub_size_str, e.what());
233+ } catch (...) {
234+ GELOGD("%s parse runtime ub_size failed, value=%s, reason=unknown exception", kLogPrefix, ub_size_str);
235+ }
236+ ub_size = 0;
237+ return false;
238+}
239+ 
240+bool TryGetRuntimeUbSize(int64_t &ub_size) {
241+ char ub_size_str[kMaxRuntimeSpecValueLen] = {};
242+ const auto ret = rtGetSocSpec(kAicoreSpecLabel, kUbSizeSpecKey, ub_size_str, kMaxRuntimeSpecValueLen);
243+ if (ret != RT_ERROR_NONE) {
244+ GELOGD("%s get runtime ub_size failed, ret=%d", kLogPrefix, ret);
245+ ub_size = 0;
246+ return false;
247+ }
248+ return TryParseRuntimeUbSize(ub_size_str, ub_size);
249+}
250+ 
251+bool ShouldDropImplGraph(const af::AscGraph &impl_graph, int64_t ub_size, const TemplatePosition &position) {
252+ const auto eval_result = EvalGraphMinUbUsage(impl_graph);
253+ const auto template_name =
254+ BuildTemplateName(position.node_idx, position.result_idx, position.group_idx, position.impl_idx);
255+ if (eval_result.status == EvalStatus::kUnknown) {
256+ GELOGD(
257+ "%s keep template due to unknown UB expr, graph=%s, template=%s, tiling_case=%ld, ub_expr=%s, "
258+ "platform_ub_size=%ld",
259+ kLogPrefix, impl_graph.GetName().c_str(), template_name.c_str(), impl_graph.GetTilingKey(),
260+ ExprToString(eval_result.origin_expr).c_str(), ub_size);
261+ return false;
262+ }
263+ if (eval_result.status == EvalStatus::kFailed) {
264+ GELOGD(
265+ "%s keep template due to failed UB expr eval, graph=%s, template=%s, tiling_case=%ld, ub_expr=%s, min_expr=%s, "
266+ "platform_ub_size=%ld",
267+ kLogPrefix, impl_graph.GetName().c_str(), template_name.c_str(), impl_graph.GetTilingKey(),
268+ ExprToString(eval_result.origin_expr).c_str(), ExprToString(eval_result.min_expr).c_str(), ub_size);
269+ return false;
270+ }
271+ if (eval_result.min_ub_usage <= ub_size) {
272+ GELOGD(
273+ "%s keep template, graph=%s, template=%s, tiling_case=%ld, ub_expr=%s, min_expr=%s, min_ub_usage=%ld, "
274+ "platform_ub_size=%ld",
275+ kLogPrefix, impl_graph.GetName().c_str(), template_name.c_str(), impl_graph.GetTilingKey(),
276+ ExprToString(eval_result.origin_expr).c_str(), ExprToString(eval_result.min_expr).c_str(),
277+ eval_result.min_ub_usage, ub_size);
278+ return false;
279+ }
280+ GELOGD(
281+ "%s drop template, graph=%s, template=%s, tiling_case=%ld, ub_expr=%s, min_expr=%s, min_ub_usage=%ld, "
282+ "platform_ub_size=%ld, reason=min_ub_usage exceeds platform_ub_size",
283+ kLogPrefix, impl_graph.GetName().c_str(), template_name.c_str(), impl_graph.GetTilingKey(),
284+ ExprToString(eval_result.origin_expr).c_str(), ExprToString(eval_result.min_expr).c_str(),
285+ eval_result.min_ub_usage, ub_size);
286+ return true;
287+}
288+ 
289+size_t CountImplGraphs(const std::vector<ascir::ScheduledResult> &scheduled_results) {
290+ size_t count = 0UL;
291+ for (const auto &scheduled_result : scheduled_results) {
292+ for (const auto &schedule_group : scheduled_result.schedule_groups) {
293+ count += schedule_group.impl_graphs.size();
294+ }
295+ }
296+ return count;
297+}
298+ 
299+bool FilterScheduledResult(FilterState &state, ascir::ScheduledResult &scheduled_result, size_t node_idx,
300+ size_t result_idx) {
301+ bool keep_scheduled_result = true;
302+ for (size_t group_idx = 0UL; group_idx < scheduled_result.schedule_groups.size(); ++group_idx) {
303+ auto &schedule_group = scheduled_result.schedule_groups[group_idx];
304+ const size_t before_group_size = schedule_group.impl_graphs.size();
305+ size_t impl_idx = 0UL;
306+ schedule_group.impl_graphs.erase(
307+ std::remove_if(schedule_group.impl_graphs.begin(), schedule_group.impl_graphs.end(),
308+ [&](const af::AscGraph &impl_graph) {
309+ const size_t current_impl_idx = impl_idx++;
310+ const TemplatePosition position = {node_idx, result_idx, group_idx, current_impl_idx};
311+ const bool drop = ShouldDropImplGraph(impl_graph, state.ub_size, position);
312+ if (drop) {
313+ schedule_group.graph_name_to_score_funcs.erase(impl_graph.GetName());
314+ }
315+ return drop;
316+ }),
317+ schedule_group.impl_graphs.end());
318+ if (before_group_size > 0UL && schedule_group.impl_graphs.empty()) {
319+ GELOGD(
320+ "%s drop scheduled result because all templates in group are filtered, graph=%s, node_idx=%zu, "
321+ "result_idx=%zu, group_idx=%zu",
322+ kLogPrefix, GetGraphName(state.fused_scheduled_result.fused_graph_name), node_idx, result_idx, group_idx);
323+ keep_scheduled_result = false;
324+ }
atomgit-bot
atomgit-botatomgit-bot6月21日

🔵 Low Priority

在 FilterScheduledResult 函数中（第 265-270 行），当某个 schedule_group 的所有 impl_graphs 全部被过滤掉后，函数立即 return false，导致后续的 schedule_group（更高索引）不会被处理。虽然外层 remove_if 会整体移除该 ScheduledResult，功能上正确，但 state.total_dropped 只会计数到当前处理过的 group 中被显式丢弃的模板数量，后续未处理的 group 中被隐式丢弃的模板不会被计数。最终在 Filter::Filter() 第 338 行日志中输出的 total_dropped 会低于实际丢弃的模板总数（低估）。

而 FilterNodeScheduledResults 第 291 行日志中使用的 dropped = before - kept 则能正确反映净减少数量。两者不一致会误导调试。

建议：两种修复方向：1) 如果需要精确计数，可将早期返回改为先遍历完所有 group（仅标记结果为 false），确保所有 group 的丢弃都被计入 total_dropped；2) 如果接受近似计数，可将第 338 行日志中的 total_dropped 改为使用 before - kept 差值，或在注释中说明 total_dropped 仅计数显式丢弃。

likedislike
不准确？
325+ }
326+ return keep_scheduled_result;
327+}
328+ 
329+af::Status FilterNodeScheduledResults(ascir::FusedScheduledResult &fused_scheduled_result,
330+ std::vector<ascir::ScheduledResult> &scheduled_results, int64_t ub_size,
331+ size_t node_idx, size_t &total_dropped) {
332+ const size_t before = CountImplGraphs(scheduled_results);
333+ FilterState state = {fused_scheduled_result, ub_size};
334+ size_t result_idx = 0UL;
335+ scheduled_results.erase(std::remove_if(scheduled_results.begin(), scheduled_results.end(),
336+ [&](ascir::ScheduledResult &scheduled_result) {
337+ const size_t current_result_idx = result_idx++;
338+ return !FilterScheduledResult(state, scheduled_result, node_idx,
339+ current_result_idx);
340+ }),
341+ scheduled_results.end());
342+ const size_t kept = CountImplGraphs(scheduled_results);
343+ const size_t dropped = before - kept;
344+ total_dropped += dropped;
345+ GELOGI("%s graph=%s, node_idx=%zu, templates_before=%zu, dropped=%zu, kept=%zu", kLogPrefix,
346+ GetGraphName(fused_scheduled_result.fused_graph_name), node_idx, before, dropped, kept);
347+ GE_ASSERT_TRUE(before == 0UL || !scheduled_results.empty(), "%s all templates are filtered, graph=%s, node_idx=%zu",
348+ kLogPrefix, GetGraphName(fused_scheduled_result.fused_graph_name), node_idx);
349+ return af::SUCCESS;
350+}
351+ 
352+UbLimitResult GetUbLimit() {
353+ UbLimitResult result;
354+ int64_t ub_size_override = 0;
355+ const bool has_ub_size_override = ge::PlatformContext::GetInstance().TryGetUbSizeOverride(ub_size_override);
356+ if (has_ub_size_override && ub_size_override > 0) {
357+ result.has_ub_limit = true;
358+ result.ub_size = ub_size_override;
359+ return result;
360+ }
361+ 
362+ ge::PlatformInfo platform_info;
363+ const bool has_platform_info = ge::PlatformContext::GetInstance().TryGetInitializedPlatformInfo(platform_info);
364+ if (has_platform_info && platform_info.ub_size > 0) {
365+ result.has_ub_limit = true;
366+ result.ub_size = platform_info.ub_size;
367+ return result;
368+ }
369+ 
370+ int64_t runtime_ub_size = 0;
371+ if (TryGetRuntimeUbSize(runtime_ub_size)) {
372+ result.has_ub_limit = true;
373+ result.ub_size = runtime_ub_size;
374+ return result;
375+ }
376+ 
377+ if (has_ub_size_override) {
378+ result.has_ub_limit = true;
379+ result.ub_size = ub_size_override;
380+ return result;
381+ }
382+ if (has_platform_info) {
383+ result.has_ub_limit = true;
384+ result.ub_size = platform_info.ub_size;
385+ }
386+ return result;
387+}
388+ 
389+} // namespace
390+ 
391+af::Status StaticUbTemplateFilter::Filter(ascir::FusedScheduledResult &fused_scheduled_result) const {
392+ const auto ub_limit = GetUbLimit();
393+ if (!ub_limit.has_ub_limit) {
394+ GELOGD("%s skip UB template filter because platform info is not initialized, graph=%s", kLogPrefix,
395+ GetGraphName(fused_scheduled_result.fused_graph_name));
396+ return af::SUCCESS;
397+ }
398+ const int64_t ub_size = ub_limit.ub_size;
399+ if (ub_size <= 0) {
400+ GELOGD("%s skip UB template filter because ub_size is invalid, graph=%s, ub_size=%ld", kLogPrefix,
401+ GetGraphName(fused_scheduled_result.fused_graph_name), ub_size);
402+ return af::SUCCESS;
403+ }
404+ 
405+ size_t total_dropped = 0UL;
406+ for (size_t node_idx = 0UL; node_idx < fused_scheduled_result.node_idx_to_scheduled_results.size(); ++node_idx) {
407+ auto &scheduled_results = fused_scheduled_result.node_idx_to_scheduled_results[node_idx];
408+ GE_CHK_STATUS_RET(
409+ FilterNodeScheduledResults(fused_scheduled_result, scheduled_results, ub_size, node_idx, total_dropped));
410+ }
411+ GELOGI("%s graph=%s, total_dropped=%zu, platform_ub_size=%ld", kLogPrefix,
412+ GetGraphName(fused_scheduled_result.fused_graph_name), total_dropped, ub_size);
413+ return af::SUCCESS;
414+}
415+ 
416+} // namespace optimize
Aautofuse/optimize/static_ub_template_filter.h+25-0
@@ -0,0 +1,25 @@
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+#ifndef ASCGEN_DEV_BASE_OPTIMIZE_STATIC_UB_TEMPLATE_FILTER_H_
12+#define ASCGEN_DEV_BASE_OPTIMIZE_STATIC_UB_TEMPLATE_FILTER_H_
13+ 
14+#include "schedule_result.h"
15+ 
16+namespace optimize {
17+ 
18+class StaticUbTemplateFilter {
19+ public:
20+ af::Status Filter(ascir::FusedScheduledResult &fused_scheduled_result) const;
21+};
22+ 
23+} // namespace optimize
24+ 
25+#endif // ASCGEN_DEV_BASE_OPTIMIZE_STATIC_UB_TEMPLATE_FILTER_H_
Mautofuse/tests/depends/runtime/src/runtime_stub.cc+12-0
@@ -43,6 +43,18 @@ rtError_t RuntimeStub::rtGetSocVersion(char *version, const uint32_t maxLen) {
43 43 
44rtError_t RuntimeStub::rtGetSocSpec(const char *label, const char *key, char *val, const uint32_t maxLen) {44rtError_t RuntimeStub::rtGetSocSpec(const char *label, const char *key, char *val, const uint32_t maxLen) {
45 (void)label;45 (void)label;
46+ if (strcmp(key, "NpuArch") == 0) {
47+ (void)strcpy_s(val, maxLen, "2201");
48+ return RT_ERROR_NONE;
49+ }
50+ if (strcmp(key, "vector_core_cnt") == 0) {
51+ (void)strcpy_s(val, maxLen, "48");
52+ return RT_ERROR_NONE;
53+ }
54+ if (strcmp(key, "ub_size") == 0) {
55+ (void)strcpy_s(val, maxLen, "245760");
56+ return RT_ERROR_NONE;
57+ }
46 // 返回 padding_size = 32 (兼容旧平台)58 // 返回 padding_size = 32 (兼容旧平台)
47 if (strcmp(key, "padding_size") == 0) {59 if (strcmp(key, "padding_size") == 0) {
48 (void)strcpy_s(val, maxLen, "32");60 (void)strcpy_s(val, maxLen, "32");
Mautofuse/tests/depends/runtime/src/runtime_stub.h+30-4
@@ -12,6 +12,7 @@
12#define __INC_LLT_RUNTIME_STUB_H12#define __INC_LLT_RUNTIME_STUB_H
13 13 
14#include <cstdint>14#include <cstdint>
15+#include <cstring>
15#include <vector>16#include <vector>
16#include <memory>17#include <memory>
17#include <mutex>18#include <mutex>
@@ -21,6 +22,34 @@
21using rtError_t = int32_t;22using rtError_t = int32_t;
22constexpr rtError_t RT_ERROR_NONE = 0;23constexpr rtError_t RT_ERROR_NONE = 0;
23namespace ge {24namespace ge {
25+struct RuntimeSocSpecDefaults {
26+ const char *npu_arch = "3510";
27+ const char *fallback = "0";
28+};
29+ 
30+struct RuntimeSocSpecValue {
31+ const char *key;
32+ const char *value;
33+};
34+ 
35+inline rtError_t CopyRuntimeSocSpecValue(const char *label, const char *key, char *val, const uint32_t maxLen,
36+ const RuntimeSocSpecDefaults &defaults) {
37+ (void)label;
38+ const RuntimeSocSpecValue specs[] = {
39+ {"NpuArch", defaults.npu_arch},
40+ {"vector_core_cnt", "48"},
41+ {"ub_size", "245760"},
42+ };
43+ for (const auto &spec : specs) {
44+ if (std::strcmp(key, spec.key) == 0) {
45+ (void)strcpy_s(val, maxLen, spec.value);
46+ return RT_ERROR_NONE;
47+ }
48+ }
49+ (void)strcpy_s(val, maxLen, defaults.fallback);
50+ return RT_ERROR_NONE;
51+}
52+ 
24class RuntimeStub {53class RuntimeStub {
25 public:54 public:
26 virtual ~RuntimeStub() = default;55 virtual ~RuntimeStub() = default;
@@ -55,10 +84,7 @@ class RuntimeStubV2Common : public RuntimeStub {
55 }84 }
56 85 
57 rtError_t rtGetSocSpec(const char *label, const char *key, char *val, const uint32_t maxLen) override {86 rtError_t rtGetSocSpec(const char *label, const char *key, char *val, const uint32_t maxLen) override {
58- (void)label;87+ return CopyRuntimeSocSpecValue(label, key, val, maxLen, RuntimeSocSpecDefaults{"3510", "0"});
59- (void)key;
60- (void)strcpy_s(val, maxLen, "3510");
61- return RT_ERROR_NONE;
62 }88 }
63};89};
64} // namespace ge90} // namespace ge
Mautofuse/tests/st/common/runtime_stub.h+2-4
@@ -10,6 +10,7 @@
10 10 
11#ifndef AUTOFUSE_RUNTIME_STUB_H11#ifndef AUTOFUSE_RUNTIME_STUB_H
12#define AUTOFUSE_RUNTIME_STUB_H12#define AUTOFUSE_RUNTIME_STUB_H
13+#include <cstring>
13#include_next <runtime_stub.h>14#include_next <runtime_stub.h>
14namespace af {15namespace af {
15using ge::RuntimeStub;16using ge::RuntimeStub;
@@ -21,10 +22,7 @@ class RuntimeStubV2 : public ge::RuntimeStub {
21 }22 }
22 23 
23 rtError_t rtGetSocSpec(const char *label, const char *key, char *val, const uint32_t maxLen) override {24 rtError_t rtGetSocSpec(const char *label, const char *key, char *val, const uint32_t maxLen) override {
24- (void)label;25+ return ge::CopyRuntimeSocSpecValue(label, key, val, maxLen, ge::RuntimeSocSpecDefaults{"3510", "0"});
25- (void)key;
26- (void)strcpy_s(val, maxLen, "3510");
27- return RT_ERROR_NONE;
28 }26 }
29};27};
30} // namespace af28} // namespace af
Mautofuse/tests/st/optimize/optimize_st.cpp+1-1
@@ -4053,7 +4053,7 @@ TEST_F(OptimizerSt, NodeCacheMarkerReduce) {
4053 ::ascir::FusedScheduledResult fused_scheduled_result;4053 ::ascir::FusedScheduledResult fused_scheduled_result;
4054 Status res = optimizer.Optimize(graph, fused_scheduled_result);4054 Status res = optimizer.Optimize(graph, fused_scheduled_result);
4055 EXPECT_EQ(res, af::SUCCESS);4055 EXPECT_EQ(res, af::SUCCESS);
4056- EXPECT_EQ(fused_scheduled_result.node_idx_to_scheduled_results[0].size(), 4UL);4056+ EXPECT_EQ(fused_scheduled_result.node_idx_to_scheduled_results[0].size(), 3UL);
4057 const auto &impl_graphs = fused_scheduled_result.node_idx_to_scheduled_results[0][0].schedule_groups[0].impl_graphs;4057 const auto &impl_graphs = fused_scheduled_result.node_idx_to_scheduled_results[0][0].schedule_groups[0].impl_graphs;
4058 EXPECT_EQ(impl_graphs.size(), 2);4058 EXPECT_EQ(impl_graphs.size(), 2);
4059 4059 
Mautofuse/tests/st/python/test_python_ascir.py+458-378
@@ -1,21 +1,21 @@
1# -*- coding: utf-8 -*-1# -*- coding: utf-8 -*-
2# -----------------------------------------------------------------------------------------------------------2# -----------------------------------------------------------------------------------------------------------
3# Copyright (c) 2025 Huawei Technologies Co., Ltd.3# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5# CANN Open Software License Agreement Version 2.0 (the "License").5# CANN Open Software License Agreement Version 2.0 (the "License").
6# Please refer to the License for details. You may not use this file except in compliance with the License.6# Please refer to the License for details. You may not use this file except in compliance with the License.
7-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.8# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9# See LICENSE in the root of the software repository for the full text of the License.9# See LICENSE in the root of the software repository for the full text of the License.
10# -----------------------------------------------------------------------------------------------------------10# -----------------------------------------------------------------------------------------------------------
11 11 
12-import math
13import pytest12import pytest
14import json13import json
15import time14import time
16import os15import os
17import shutil16import shutil
18from autofuse.pyautofuse import ascir, Autofuser, AutofuserOptions, Schedule, CodeGen17from autofuse.pyautofuse import ascir, Autofuser, AutofuserOptions, Schedule, CodeGen
18+ 
19try:19try:
20 from autofuse import ascir_api20 from autofuse import ascir_api
21except ImportError:21except ImportError:
@@ -24,19 +24,19 @@ from ascir import Max, Min, Mod
24 24 
25PYF_PATH = os.path.dirname(os.path.realpath(__file__))25PYF_PATH = os.path.dirname(os.path.realpath(__file__))
26 26 
27-ascir.utils.set_platform("2201", 1, 1024)27+ascir.utils.set_platform("2201", 1, 245760)
28 28 
29 29 
30# pyascir 构图能力暂不支持30# pyascir 构图能力暂不支持
31-class TestAscir():31+class TestAscir:
32 @staticmethod32 @staticmethod
33 def test_graph_create_size_expr_by_long():33 def test_graph_create_size_expr_by_long():
34 s0 = ascir.SizeExpr()34 s0 = ascir.SizeExpr()
35 assert s0 == 135 assert s0 == 1
36 try:36 try:
37- s0 = ascir.SizeExpr('100')37+ s0 = ascir.SizeExpr("100")
38 except Exception as e:38 except Exception as e:
39- assert e.args[0] == 'Only support type of SizeExpr or long'39+ assert e.args[0] == "Only support type of SizeExpr or long"
40 40 
41 assert s0 == 141 assert s0 == 1
42 s0 = ascir.SizeExpr(0)42 s0 = ascir.SizeExpr(0)
@@ -67,20 +67,22 @@ class TestAscir():
67 @staticmethod67 @staticmethod
68 def test_graph_create_size():68 def test_graph_create_size():
69 graph = ascir.HintGraph("test")69 graph = ascir.HintGraph("test")
70- s0 = graph.create_size("s0")70+ graph.create_size("s0")
71- s1 = graph.create_size("s1")71+ graph.create_size("s1")
72- s2 = graph.create_size("s2")72+ graph.create_size("s2")
73 73 
74 debug_str = ascir.utils.debug_str(graph)74 debug_str = ascir.utils.debug_str(graph)
75- assert debug_str == "".join([75+ assert debug_str == "".join(
76- "Graph: test\n",76+ [
77- "Sizes:\n",77+ "Graph: test\n",
78- " s0: VAR\n",78+ "Sizes:\n",
79- " s1: VAR\n",79+ " s0: VAR\n",
80- " s2: VAR\n",80+ " s1: VAR\n",
81- "Axis:\n",81+ " s2: VAR\n",
82- "Nodes:\n",82+ "Axis:\n",
83- ])83+ "Nodes:\n",
84+ ]
85+ )
84 86 
85 @staticmethod87 @staticmethod
86 def test_graph_create_axis():88 def test_graph_create_axis():
@@ -89,33 +91,35 @@ class TestAscir():
89 s1 = graph.create_size("s1")91 s1 = graph.create_size("s1")
90 s2 = graph.create_size("s2")92 s2 = graph.create_size("s2")
91 93 
92- z0 = graph.create_axis("z0", s0)94+ graph.create_axis("z0", s0)
93- z1 = graph.create_axis("z1", s1)95+ graph.create_axis("z1", s1)
94- z2 = graph.create_axis("z2", s2)96+ graph.create_axis("z2", s2)
95- z3 = graph.create_axis("z3", 512)97+ graph.create_axis("z3", 512)
96- z4 = graph.create_axis("z4", s1 + s2)98+ graph.create_axis("z4", s1 + s2)
97 99 
98 debug_str = ascir.utils.debug_str(graph)100 debug_str = ascir.utils.debug_str(graph)
99- assert debug_str == "".join([101+ assert debug_str == "".join(
100- "Graph: test\n",102+ [
101- "Sizes:\n",103+ "Graph: test\n",
102- " s0: VAR\n",104+ "Sizes:\n",
103- " s1: VAR\n",105+ " s0: VAR\n",
104- " s2: VAR\n",106+ " s1: VAR\n",
105- "Axis:\n",107+ " s2: VAR\n",
106- " z0(0) : ORIGINAL, size:s0, \n",108+ "Axis:\n",
107- " z1(1) : ORIGINAL, size:s1, \n",109+ " z0(0) : ORIGINAL, size:s0, \n",
108- " z2(2) : ORIGINAL, size:s2, \n",110+ " z1(1) : ORIGINAL, size:s1, \n",
109- " z3(3) : ORIGINAL, size:512, \n",111+ " z2(2) : ORIGINAL, size:s2, \n",
110- " z4(4) : ORIGINAL, size:(s1 + s2), \n",112+ " z3(3) : ORIGINAL, size:512, \n",
111- "Nodes:\n",113+ " z4(4) : ORIGINAL, size:(s1 + s2), \n",
112- ])114+ "Nodes:\n",
115+ ]
116+ )
113 117 
114 @staticmethod118 @staticmethod
115 def test_graph_create_node():119 def test_graph_create_node():
116 graph = ascir.HintGraph("test")120 graph = ascir.HintGraph("test")
117 121 
118- x = ascir.ops.Data("x", graph)122+ ascir.ops.Data("x", graph)
119 debug_str = ascir.utils.debug_str(graph)123 debug_str = ascir.utils.debug_str(graph)
120 assert debug_str124 assert debug_str
121 125 
@@ -134,8 +138,12 @@ class TestAscir():
134 try:138 try:
135 graph.infer_dtypes()139 graph.infer_dtypes()
136 except Exception as e:140 except Exception as e:
137- assert e.args[0] == 'Check dtype failed for cast Cast; input_dtypes: [DT_INT8], output_dytpes: [DT_INT4]'141+ assert (
142+ e.args[0]
143+ == "Check dtype failed for cast Cast; input_dtypes: [DT_INT8], output_dytpes: [DT_INT4]"
144+ )
138 import sys145 import sys
146+ 
139 # 通常为2 （变量 + getrefcount 参数）147 # 通常为2 （变量 + getrefcount 参数）
140 print(f"x.attr ref count is {sys.getrefcount(x.attr)}")148 print(f"x.attr ref count is {sys.getrefcount(x.attr)}")
141 del x.attr149 del x.attr
@@ -145,21 +153,23 @@ class TestAscir():
145 graph = ascir.HintGraph("test")153 graph = ascir.HintGraph("test")
146 154 
147 x = ascir_api.Data(graph, dtype=ascir.dtypes.int8)155 x = ascir_api.Data(graph, dtype=ascir.dtypes.int8)
148- cast = None
149 try:156 try:
150- cast = ascir_api.Cast(graph, x, dtype=ascir.dtypes.int4, axis=[])157+ ascir_api.Cast(graph, x, dtype=ascir.dtypes.int4, axis=[])
151 except Exception as e:158 except Exception as e:
152- assert e.args[0] == 'Check dtype failed for cast_0 Cast; input_dtypes: [DT_INT8], output_dytpes: [DT_INT4]'159+ assert (
160+ e.args[0]
161+ == "Check dtype failed for cast_0 Cast; input_dtypes: [DT_INT8], output_dytpes: [DT_INT4]"
162+ )
153 163 
154 @staticmethod164 @staticmethod
155 def test_graph_create_const_node_with_value_str_attr():165 def test_graph_create_const_node_with_value_str_attr():
156 graph = ascir.HintGraph("test")166 graph = ascir.HintGraph("test")
157 167 
158 x = ascir.ops.Scalar("x", graph)168 x = ascir.ops.Scalar("x", graph)
159- x.attr.ir_attr.value = '11.1'169+ x.attr.ir_attr.value = "11.1"
160 x.y.dtype = ascir.dtypes.float32170 x.y.dtype = ascir.dtypes.float32
161 debug_str = ascir.utils.debug_str(graph)171 debug_str = ascir.utils.debug_str(graph)
162- assert x.attr.ir_attr.value == '11.1'172+ assert x.attr.ir_attr.value == "11.1"
163 assert debug_str173 assert debug_str
164 174 
165 @staticmethod175 @staticmethod
@@ -176,7 +186,7 @@ class TestAscir():
176 x = ascir.ops.Data("x", graph)186 x = ascir.ops.Data("x", graph)
177 x.attr.sched.axis = [z0, z1, z2]187 x.attr.sched.axis = [z0, z1, z2]
178 188 
179- load = ascir.ops.Load("load")189+ ascir.ops.Load("load")
180 x.y.dtype = ascir.dtypes.float16190 x.y.dtype = ascir.dtypes.float16
181 x.y.axis = [z0, z1, z2]191 x.y.axis = [z0, z1, z2]
182 assert x.y.axis == [z0.id, z1.id, z2.id]192 assert x.y.axis == [z0.id, z1.id, z2.id]
@@ -198,7 +208,10 @@ class TestAscir():
198 try:208 try:
199 graph.infer_dtypes()209 graph.infer_dtypes()
200 except Exception as e:210 except Exception as e:
201- assert e.args[0] == 'Infer dtype failed for load Load; input_dtypes: [DT_INT64] is not supportted now'211+ assert (
212+ e.args[0]
213+ == "Infer dtype failed for load Load; input_dtypes: [DT_INT64] is not supportted now"
214+ )
202 graph.infer_dtypes()215 graph.infer_dtypes()
203 debug_str = ascir.utils.debug_str(graph)216 debug_str = ascir.utils.debug_str(graph)
204 assert debug_str217 assert debug_str
@@ -220,25 +233,27 @@ class TestAscir():
220 start = time.time()233 start = time.time()
221 time.sleep(0.1)234 time.sleep(0.1)
222 end = time.time()235 end = time.time()
223- ascir.utils.duration_record(["device", "fused_graph"], int(start * 1e9), int((end - start) * 1e9))236+ ascir.utils.duration_record(
237+ ["device", "fused_graph"], int(start * 1e9), int((end - start) * 1e9)
238+ )
224 ascir.utils.report_durations()239 ascir.utils.report_durations()
225 try:240 try:
226 ascir.utils.duration_record(["device", "fused_graph"], "time")241 ascir.utils.duration_record(["device", "fused_graph"], "time")
227 except TypeError as e:242 except TypeError as e:
228- assert e.args[0] == 'UtilsDurationRecord param parse failed'243+ assert e.args[0] == "UtilsDurationRecord param parse failed"
229 244 
230 try:245 try:
231 ascir.utils.duration_record(["device", "fused_graph"], int(-1), int(-1))246 ascir.utils.duration_record(["device", "fused_graph"], int(-1), int(-1))
232 except TypeError as e:247 except TypeError as e:
233- assert e.args[0] == 'duration param is invalid'248+ assert e.args[0] == "duration param is invalid"
234 249 
235 try:250 try:
236 ascir.utils.duration_record([0, 1], int(-1), int(-1))251 ascir.utils.duration_record([0, 1], int(-1), int(-1))
237 except TypeError as e:252 except TypeError as e:
238- assert e.args[0] == 'target param is invalid'253+ assert e.args[0] == "target param is invalid"
239 254 
240 255 
241-class TestAutofuseLoadAbsStore():256+class TestAutofuseLoadAbsStore:
242 @staticmethod257 @staticmethod
243 def construct_graph():258 def construct_graph():
244 graph = ascir.HintGraph("LoadAbsStore")259 graph = ascir.HintGraph("LoadAbsStore")
@@ -266,7 +281,7 @@ class TestAutofuseLoadAbsStore():
266 try:281 try:
267 load.attr.ir_attr.offset = "3"282 load.attr.ir_attr.offset = "3"
268 except Exception as e:283 except Exception as e:
269- assert e.args[0] == 'Only support type of SizeExpr or long'284+ assert e.args[0] == "Only support type of SizeExpr or long"
270 offset_of_0 = ascir.SizeExpr(0)285 offset_of_0 = ascir.SizeExpr(0)
271 load.attr.ir_attr.offset = offset_of_0286 load.attr.ir_attr.offset = offset_of_0
272 assert load.attr.ir_attr.offset.expression == "0"287 assert load.attr.ir_attr.offset.expression == "0"
@@ -289,7 +304,7 @@ class TestAutofuseLoadAbsStore():
289 try:304 try:
290 store.attr.ir_attr.offset = "4"305 store.attr.ir_attr.offset = "4"
291 except Exception as e:306 except Exception as e:
292- assert e.args[0] == 'Only support type of SizeExpr or long'307+ assert e.args[0] == "Only support type of SizeExpr or long"
293 store.attr.ir_attr.offset = offset_of_0 + 1308 store.attr.ir_attr.offset = offset_of_0 + 1
294 assert store.attr.ir_attr.offset.expression == "1"309 assert store.attr.ir_attr.offset.expression == "1"
295 store.x = abs_op310 store.x = abs_op
@@ -322,9 +337,9 @@ class TestAutofuseLoadAbsStore():
322 337 
323 hint_graph = self.construct_graph()338 hint_graph = self.construct_graph()
324 schedule_results = fuser.schedule(hint_graph)339 schedule_results = fuser.schedule(hint_graph)
325- graph_name = schedule_results.get_name()340+ schedule_results.get_name()
326- input_num = schedule_results.get_input_num()341+ schedule_results.get_input_num()
327- output_num = schedule_results.get_output_num()342+ schedule_results.get_output_num()
328 343 
329 def test_codegen(self):344 def test_codegen(self):
330 options = AutofuserOptions()345 options = AutofuserOptions()
@@ -351,10 +366,12 @@ class TestAutofuseLoadAbsStore():
351 assert isinstance(kernel_dict, dict), "kernel_dict should be a dictionary"366 assert isinstance(kernel_dict, dict), "kernel_dict should be a dictionary"
352 367 
353 # 验证字典中包含预期的键（可能是"ub"、"common"或"default"）368 # 验证字典中包含预期的键（可能是"ub"、"common"或"default"）
354- assert any(key in tiling_dict for key in ["ub", "common", "default"]), \369+ assert any(key in tiling_dict for key in ["ub", "common", "default"]), (
355 "tiling_dict should contain at least one of the expected keys"370 "tiling_dict should contain at least one of the expected keys"
356- assert any(key in kernel_dict for key in ["ub", "common", "default"]), \371+ )
372+ assert any(key in kernel_dict for key in ["ub", "common", "default"]), (
357 "kernel_dict should contain at least one of the expected keys"373 "kernel_dict should contain at least one of the expected keys"
374+ )
358 375 
359 def test_host_code_generator(self):376 def test_host_code_generator(self):
360 # 测试host_code_generator方法377 # 测试host_code_generator方法
@@ -381,9 +398,9 @@ class TestAutofuseLoadAbsStore():
381 assert isinstance(infer_shape, str), "infer_shape should be a string"398 assert isinstance(infer_shape, str), "infer_shape should be a string"
382 399 
383 # 验证py_tilings中包含预期的键（可能是"ub"、"common"或"default"）400 # 验证py_tilings中包含预期的键（可能是"ub"、"common"或"default"）
384- assert any(key in py_tilings for key in ["ub", "common", "default"]), \401+ assert any(key in py_tilings for key in ["ub", "common", "default"]), (
385 "py_tilings should contain at least one of the expected keys"402 "py_tilings should contain at least one of the expected keys"
386- 403+ )
387 404 
388 def test_autofuse_backend(self):405 def test_autofuse_backend(self):
389 options = AutofuserOptions()406 options = AutofuserOptions()
@@ -393,7 +410,7 @@ class TestAutofuseLoadAbsStore():
393 tiling_def, host_tiling, op_kernel = fuser.autofuse_backend(hint_graph)410 tiling_def, host_tiling, op_kernel = fuser.autofuse_backend(hint_graph)
394 411 
395 412 
396-class TestAutofuseLoadMatMulStore():413+class TestAutofuseLoadMatMulStore:
397 @staticmethod414 @staticmethod
398 def construct_graph():415 def construct_graph():
399 graph = ascir.HintGraph("LoadCubeStore")416 graph = ascir.HintGraph("LoadCubeStore")
@@ -421,7 +438,7 @@ class TestAutofuseLoadMatMulStore():
421 try:438 try:
422 load.attr.ir_attr.offset = "3"439 load.attr.ir_attr.offset = "3"
423 except Exception as e:440 except Exception as e:
424- assert e.args[0] == 'Only support type of SizeExpr or long'441+ assert e.args[0] == "Only support type of SizeExpr or long"
425 offset_of_0 = ascir.SizeExpr(0)442 offset_of_0 = ascir.SizeExpr(0)
426 load.attr.ir_attr.offset = offset_of_0443 load.attr.ir_attr.offset = offset_of_0
427 assert load.attr.ir_attr.offset.expression == "0"444 assert load.attr.ir_attr.offset.expression == "0"
@@ -451,7 +468,7 @@ class TestAutofuseLoadMatMulStore():
451 try:468 try:
452 store.attr.ir_attr.offset = "4"469 store.attr.ir_attr.offset = "4"
453 except Exception as e:470 except Exception as e:
454- assert e.args[0] == 'Only support type of SizeExpr or long'471+ assert e.args[0] == "Only support type of SizeExpr or long"
455 store.attr.ir_attr.offset = offset_of_0 + 1472 store.attr.ir_attr.offset = offset_of_0 + 1
456 assert store.attr.ir_attr.offset.expression == "1"473 assert store.attr.ir_attr.offset.expression == "1"
457 store.x = matmul_op474 store.x = matmul_op
@@ -480,8 +497,8 @@ class TestAutofuseLoadMatMulStore():
480 497 
481 hint_graph = self.construct_graph()498 hint_graph = self.construct_graph()
482 schedule_results = fuser.schedule(hint_graph)499 schedule_results = fuser.schedule(hint_graph)
483- attr = schedule_results.is_cube_type()500+ schedule_results.is_cube_type()
484- attr = schedule_results.get_cube_attributes()501+ schedule_results.get_cube_attributes()
485 502 
486 def test_device_code_generator(self):503 def test_device_code_generator(self):
487 # 测试device_code_generator方法504 # 测试device_code_generator方法
@@ -500,10 +517,12 @@ class TestAutofuseLoadMatMulStore():
500 assert isinstance(kernel_dict, dict), "kernel_dict should be a dictionary"517 assert isinstance(kernel_dict, dict), "kernel_dict should be a dictionary"
501 518 
502 # 验证字典中包含预期的键（可能是"ub"、"common"或"default"）519 # 验证字典中包含预期的键（可能是"ub"、"common"或"default"）
503- assert any(key in tiling_dict for key in ["ub", "common", "default"]), \520+ assert any(key in tiling_dict for key in ["ub", "common", "default"]), (
504 "tiling_dict should contain at least one of the expected keys"521 "tiling_dict should contain at least one of the expected keys"
505- assert any(key in kernel_dict for key in ["ub", "common", "default"]), \522+ )
523+ assert any(key in kernel_dict for key in ["ub", "common", "default"]), (
506 "kernel_dict should contain at least one of the expected keys"524 "kernel_dict should contain at least one of the expected keys"
525+ )
507 526 
508 def test_host_code_generator(self):527 def test_host_code_generator(self):
509 # 测试host_code_generator方法528 # 测试host_code_generator方法
@@ -530,15 +549,18 @@ class TestAutofuseLoadMatMulStore():
530 assert isinstance(infer_shape, str), "infer_shape should be a string"549 assert isinstance(infer_shape, str), "infer_shape should be a string"
531 550 
532 # 验证py_tilings中包含预期的键（可能是"ub"、"common"或"default"）551 # 验证py_tilings中包含预期的键（可能是"ub"、"common"或"default"）
533- assert any(key in py_tilings for key in ["ub", "common", "default"]), \552+ assert any(key in py_tilings for key in ["ub", "common", "default"]), (
534 "py_tilings should contain at least one of the expected keys"553 "py_tilings should contain at least one of the expected keys"
554+ )
535 555 
536 # 验证每个值也是字典556 # 验证每个值也是字典
537 for key, value in py_tilings.items():557 for key, value in py_tilings.items():
538- assert isinstance(value, dict), f"Value for key {key} should be a dictionary"558+ assert isinstance(value, dict), (
559+ f"Value for key {key} should be a dictionary"
560+ )
539 561 
540 562 
541-class TestAutofuseLoadBatchMatmulStore():563+class TestAutofuseLoadBatchMatmulStore:
542 @staticmethod564 @staticmethod
543 def construct_graph():565 def construct_graph():
544 graph = ascir.HintGraph("LoadBatchMatMulStore")566 graph = ascir.HintGraph("LoadBatchMatMulStore")
@@ -566,7 +588,7 @@ class TestAutofuseLoadBatchMatmulStore():
566 try:588 try:
567 load.attr.ir_attr.offset = "3"589 load.attr.ir_attr.offset = "3"
568 except Exception as e:590 except Exception as e:
569- assert e.args[0] == 'Only support type of SizeExpr or long'591+ assert e.args[0] == "Only support type of SizeExpr or long"
570 offset_of_0 = ascir.SizeExpr(0)592 offset_of_0 = ascir.SizeExpr(0)
571 load.attr.ir_attr.offset = offset_of_0593 load.attr.ir_attr.offset = offset_of_0
572 assert load.attr.ir_attr.offset.expression == "0"594 assert load.attr.ir_attr.offset.expression == "0"
@@ -596,7 +618,7 @@ class TestAutofuseLoadBatchMatmulStore():
596 try:618 try:
597 store.attr.ir_attr.offset = "4"619 store.attr.ir_attr.offset = "4"
598 except Exception as e:620 except Exception as e:
599- assert e.args[0] == 'Only support type of SizeExpr or long'621+ assert e.args[0] == "Only support type of SizeExpr or long"
600 store.attr.ir_attr.offset = offset_of_0 + 1622 store.attr.ir_attr.offset = offset_of_0 + 1
601 assert store.attr.ir_attr.offset.expression == "1"623 assert store.attr.ir_attr.offset.expression == "1"
602 store.x = matmul_op624 store.x = matmul_op
@@ -625,11 +647,11 @@ class TestAutofuseLoadBatchMatmulStore():
625 647 
626 hint_graph = self.construct_graph()648 hint_graph = self.construct_graph()
627 schedule_results = fuser.schedule(hint_graph)649 schedule_results = fuser.schedule(hint_graph)
628- attr = schedule_results.is_cube_type()650+ schedule_results.is_cube_type()
629- attr = schedule_results.get_cube_attributes()651+ schedule_results.get_cube_attributes()
630 652 
631 653 
632-class TestAutofuseLoadMatMulStoreNew():654+class TestAutofuseLoadMatMulStoreNew:
633 @staticmethod655 @staticmethod
634 def construct_graph():656 def construct_graph():
635 graph = ascir.HintGraph("LoadCubeStoreNew")657 graph = ascir.HintGraph("LoadCubeStoreNew")
@@ -695,8 +717,8 @@ class TestAutofuseLoadMatMulStoreNew():
695 matmul_op.attr.ir_attr.enable_hf32 = 1717 matmul_op.attr.ir_attr.enable_hf32 = 1
696 matmul_op.attr.ir_attr.transpose_x1 = 0 # X1不转置718 matmul_op.attr.ir_attr.transpose_x1 = 0 # X1不转置
697 matmul_op.attr.ir_attr.transpose_x2 = 0 # X2不转置719 matmul_op.attr.ir_attr.transpose_x2 = 0 # X2不转置
698- matmul_op.attr.ir_attr.has_relu = 0 # 不使用ReLU720+ matmul_op.attr.ir_attr.has_relu = 0 # 不使用ReLU
699- matmul_op.attr.ir_attr.offset_x = 0 # 偏移量为0721+ matmul_op.attr.ir_attr.offset_x = 0 # 偏移量为0
700 matmul_op.y.dtype = ascir.dtypes.float16722 matmul_op.y.dtype = ascir.dtypes.float16
701 matmul_op.y.axis = [m, n]723 matmul_op.y.axis = [m, n]
702 matmul_op.y.size = [s0, s2]724 matmul_op.y.size = [s0, s2]
@@ -743,8 +765,8 @@ class TestAutofuseLoadMatMulStoreNew():
743 765 
744 hint_graph = self.construct_graph()766 hint_graph = self.construct_graph()
745 schedule_results = fuser.schedule(hint_graph)767 schedule_results = fuser.schedule(hint_graph)
746- attr = schedule_results.is_cube_type()768+ schedule_results.is_cube_type()
747- attr = schedule_results.get_cube_attributes()769+ schedule_results.get_cube_attributes()
748 770 
749 def test_device_code_generator(self):771 def test_device_code_generator(self):
750 # 测试device_code_generator方法772 # 测试device_code_generator方法
@@ -763,15 +785,19 @@ class TestAutofuseLoadMatMulStoreNew():
763 assert isinstance(kernel_dict, dict), "kernel_dict should be a dictionary"785 assert isinstance(kernel_dict, dict), "kernel_dict should be a dictionary"
764 786 
765 # 验证字典中只包含"ub"或"common"键787 # 验证字典中只包含"ub"或"common"键
766- assert any(key in tiling_dict for key in ["ub", "common", "default"]), \788+ assert any(key in tiling_dict for key in ["ub", "common", "default"]), (
767 "tiling_dict should contain either 'ub', 'common' or 'default' key"789 "tiling_dict should contain either 'ub', 'common' or 'default' key"
768- assert all(key in ["ub", "common", "default"] for key in tiling_dict.keys()), \790+ )
791+ assert all(key in ["ub", "common", "default"] for key in tiling_dict.keys()), (
769 "tiling_dict should only contain 'ub', 'common' or 'default' keys"792 "tiling_dict should only contain 'ub', 'common' or 'default' keys"
793+ )
770 794 
771- assert any(key in kernel_dict for key in ["ub", "common", "default"]), \795+ assert any(key in kernel_dict for key in ["ub", "common", "default"]), (
772 "kernel_dict should contain either 'ub', 'common' or 'default' key"796 "kernel_dict should contain either 'ub', 'common' or 'default' key"
773- assert all(key in ["ub", "common", "default"] for key in kernel_dict.keys()), \797+ )
798+ assert all(key in ["ub", "common", "default"] for key in kernel_dict.keys()), (
774 "kernel_dict should only contain 'ub', 'common' or 'default' keys"799 "kernel_dict should only contain 'ub', 'common' or 'default' keys"
800+ )
775 801 
776 def test_host_code_generator(self):802 def test_host_code_generator(self):
777 # 测试host_code_generator方法803 # 测试host_code_generator方法
@@ -798,14 +824,18 @@ class TestAutofuseLoadMatMulStoreNew():
798 assert isinstance(infer_shape, str), "infer_shape should be a string"824 assert isinstance(infer_shape, str), "infer_shape should be a string"
799 825 
800 # 验证py_tilings中只包含"ub"或"common"键826 # 验证py_tilings中只包含"ub"或"common"键
801- assert any(key in py_tilings for key in ["ub", "common", "default"]), \827+ assert any(key in py_tilings for key in ["ub", "common", "default"]), (
802 "py_tilings should contain either 'ub', 'common' or 'default' key"828 "py_tilings should contain either 'ub', 'common' or 'default' key"
803- assert all(key in ["ub", "common", "default"] for key in py_tilings.keys()), \829+ )
830+ assert all(key in ["ub", "common", "default"] for key in py_tilings.keys()), (
804 "py_tilings should only contain 'ub', 'common' or 'default' keys"831 "py_tilings should only contain 'ub', 'common' or 'default' keys"
832+ )
805 833 
806 # 验证每个值也是字典834 # 验证每个值也是字典
807 for key, value in py_tilings.items():835 for key, value in py_tilings.items():
808- assert isinstance(value, dict), f"Value for key {key} should be a dictionary"836+ assert isinstance(value, dict), (
837+ f"Value for key {key} should be a dictionary"
838+ )
809 839 
810 840 
811class TestCubeAttributes:841class TestCubeAttributes:
@@ -931,7 +961,7 @@ class TestCubeAttributes:
931 # assert attr_dict["input_num"] == 2, "input_num should be 2"961 # assert attr_dict["input_num"] == 2, "input_num should be 2"
932 962 
933 963 
934-class TestAutofuseGatherAbsStore():964+class TestAutofuseGatherAbsStore:
935 @staticmethod965 @staticmethod
936 def construct_graph():966 def construct_graph():
937 graph = ascir.HintGraph("GatherAbsStore")967 graph = ascir.HintGraph("GatherAbsStore")
@@ -1004,7 +1034,7 @@ class TestAutofuseGatherAbsStore():
1004 options = AutofuserOptions()1034 options = AutofuserOptions()
1005 fuser = Autofuser(options)1035 fuser = Autofuser(options)
1006 hint_graph = self.construct_graph()1036 hint_graph = self.construct_graph()
1007- schedule_results = fuser.autofuse(hint_graph)1037+ fuser.autofuse(hint_graph)
1008 1038 
1009 @pytest.mark.skip1039 @pytest.mark.skip
1010 def test_codegen(self):1040 def test_codegen(self):
@@ -1016,25 +1046,25 @@ class TestAutofuseGatherAbsStore():
1016 tiling_def, host_tiling, op_kernel = fuser.codegen(impl_graphs)1046 tiling_def, host_tiling, op_kernel = fuser.codegen(impl_graphs)
1017 1047 
1018 1048 
1019-class TestAutofuseLoadConcatStore():1049+class TestAutofuseLoadConcatStore:
1020 @staticmethod1050 @staticmethod
1021 def construct_graph():1051 def construct_graph():
1022 try:1052 try:
1023- NpuKernel0Graph = ascir.HintGraph(100)1053+ npu_kernel0_graph = ascir.HintGraph(100)
1024 except Exception as e:1054 except Exception as e:
1025- assert e.args[0] == 'argument 1 must be str, not int'1055+ assert e.args[0] == "argument 1 must be str, not int"
1026- NpuKernel0Graph = ascir.HintGraph('LoadConcatStore')1056+ npu_kernel0_graph = ascir.HintGraph("LoadConcatStore")
1027- s0 = NpuKernel0Graph.create_size("s0")1057+ s0 = npu_kernel0_graph.create_size("s0")
1028- s1 = NpuKernel0Graph.create_size("s1")1058+ s1 = npu_kernel0_graph.create_size("s1")
1029- z0 = NpuKernel0Graph.create_axis("z0", s0)1059+ z0 = npu_kernel0_graph.create_axis("z0", s0)
1030- z1 = NpuKernel0Graph.create_axis("z1", s1 * 2)1060+ z1 = npu_kernel0_graph.create_axis("z1", s1 * 2)
1031- arg2_1 = ascir.ops.Data('arg2_1', NpuKernel0Graph)1061+ arg2_1 = ascir.ops.Data("arg2_1", npu_kernel0_graph)
1032 arg2_1.y.dtype = ascir.dtypes.float161062 arg2_1.y.dtype = ascir.dtypes.float16
1033- load = ascir.ops.Load('load')1063+ load = ascir.ops.Load("load")
1034 try:1064 try:
1035 load.infer_dtype()1065 load.infer_dtype()
1036 except Exception as e:1066 except Exception as e:
1037- assert e.args[0] == 'node load Load need set input before call infer dype'1067+ assert e.args[0] == "node load Load need set input before call infer dype"
1038 load.attr.sched.axis = [z0, z1]1068 load.attr.sched.axis = [z0, z1]
1039 load.x = arg2_1.y1069 load.x = arg2_1.y
1040 load.y.axis = [z0, z1]1070 load.y.axis = [z0, z1]
@@ -1042,32 +1072,32 @@ class TestAutofuseLoadConcatStore():
1042 load.y.size = [s0, s1]1072 load.y.size = [s0, s1]
1043 load.infer_dtype()1073 load.infer_dtype()
1044 assert load.y.dtype == ascir.dtypes.float161074 assert load.y.dtype == ascir.dtypes.float16
1045- arg3_1 = ascir.ops.Data('arg3_1', NpuKernel0Graph)1075+ arg3_1 = ascir.ops.Data("arg3_1", npu_kernel0_graph)
1046 arg3_1.y.dtype = ascir.dtypes.float161076 arg3_1.y.dtype = ascir.dtypes.float16
1047- load1 = ascir.ops.Load('load1')1077+ load1 = ascir.ops.Load("load1")
1048 load1.attr.sched.axis = [z0, z1]1078 load1.attr.sched.axis = [z0, z1]
1049 assert load1.attr.sched.axis == [z0.id, z1.id]1079 assert load1.attr.sched.axis == [z0.id, z1.id]
1050 load1.x = arg3_1.y1080 load1.x = arg3_1.y
1051 load1.y.axis = [z0, z1]1081 load1.y.axis = [z0, z1]
1052 load1.y.strides = [s1, ascir.SizeExpr(1)]1082 load1.y.strides = [s1, ascir.SizeExpr(1)]
1053 load1.y.size = [s0, s1]1083 load1.y.size = [s0, s1]
1054- concat = ascir.ops.Concat('concat')1084+ concat = ascir.ops.Concat("concat")
1055 concat.attr.sched.axis = [z0, z1]1085 concat.attr.sched.axis = [z0, z1]
1056 concat.x = [load, load1.y]1086 concat.x = [load, load1.y]
1057 concat.y.axis = [z0, z1]1087 concat.y.axis = [z0, z1]
1058 concat.y.strides = [s1 + s1, ascir.SizeExpr(1)]1088 concat.y.strides = [s1 + s1, ascir.SizeExpr(1)]
1059 concat.y.size = [s0, s1 * 2]1089 concat.y.size = [s0, s1 * 2]
1060- store = ascir.ops.Store('store')1090+ store = ascir.ops.Store("store")
1061 store.attr.sched.axis = [z0, z1]1091 store.attr.sched.axis = [z0, z1]
1062 store.x = concat.y1092 store.x = concat.y
1063 store.y.axis = [z0, z1]1093 store.y.axis = [z0, z1]
1064 store.y.strides = [s1 * 2, ascir.SizeExpr(1)]1094 store.y.strides = [s1 * 2, ascir.SizeExpr(1)]
1065 store.y.size = [s0, s1 * 2]1095 store.y.size = [s0, s1 * 2]
1066- buf0 = ascir.ops.Output('buf0')1096+ buf0 = ascir.ops.Output("buf0")
1067 buf0.x = store.y1097 buf0.x = store.y
1068 buf0.y.dtype = ascir.dtypes.float161098 buf0.y.dtype = ascir.dtypes.float16
1069- NpuKernel0Graph.infer_dtypes()1099+ npu_kernel0_graph.infer_dtypes()
1070- return NpuKernel0Graph1100+ return npu_kernel0_graph
1071 1101 
1072 def test_construct_graph(self):1102 def test_construct_graph(self):
1073 graph = self.construct_graph()1103 graph = self.construct_graph()
@@ -1079,10 +1109,10 @@ class TestAutofuseLoadConcatStore():
1079 fuser = Autofuser(options)1109 fuser = Autofuser(options)
1080 1110 
1081 hint_graph = self.construct_graph()1111 hint_graph = self.construct_graph()
1082- schedule_results = fuser.schedule(hint_graph)1112+ fuser.schedule(hint_graph)
1083 1113 
1084 1114 
1085-class TestAutofuseLoadSplitStore():1115+class TestAutofuseLoadSplitStore:
1086 @staticmethod1116 @staticmethod
1087 def construct_graph():1117 def construct_graph():
1088 graph = ascir.HintGraph("LoadSplitStore")1118 graph = ascir.HintGraph("LoadSplitStore")
@@ -1148,51 +1178,51 @@ class TestAutofuseLoadSplitStore():
1148 debug_graph = ascir.utils.debug_str(graph)1178 debug_graph = ascir.utils.debug_str(graph)
1149 assert debug_graph != ""1179 assert debug_graph != ""
1150 finally:1180 finally:
1151- ascir.utils.set_platform("2201", 1, 1024)1181+ ascir.utils.set_platform("2201", 1, 245760)
1152 1182 
1153 1183 
1154-class TestWorkspaceOptimize():1184+class TestWorkspaceOptimize:
1155 @staticmethod1185 @staticmethod
1156 def construct_graph():1186 def construct_graph():
1157- NpuKernel0Graph = ascir.HintGraph('workspace')1187+ npu_kernel0_graph = ascir.HintGraph("workspace")
1158- s0 = NpuKernel0Graph.create_size("s0")1188+ s0 = npu_kernel0_graph.create_size("s0")
1159- z0 = NpuKernel0Graph.create_axis("z0", s0)1189+ z0 = npu_kernel0_graph.create_axis("z0", s0)
1160- arg2_1 = ascir.ops.Data('arg2_1', NpuKernel0Graph)1190+ arg2_1 = ascir.ops.Data("arg2_1", npu_kernel0_graph)
1161 arg2_1.y.dtype = ascir.dtypes.float161191 arg2_1.y.dtype = ascir.dtypes.float16
1162- load = ascir.ops.Load('load')1192+ load = ascir.ops.Load("load")
1163 load.attr.sched.axis = [z0]1193 load.attr.sched.axis = [z0]
1164 load.x = arg2_1.y1194 load.x = arg2_1.y
1165 load.y.axis = [z0]1195 load.y.axis = [z0]
1166 load.y.strides = [ascir.SizeExpr(1)]1196 load.y.strides = [ascir.SizeExpr(1)]
1167 load.y.size = [s0]1197 load.y.size = [s0]
1168- store = ascir.ops.Store('store')1198+ store = ascir.ops.Store("store")
1169 store.attr.sched.axis = [z0]1199 store.attr.sched.axis = [z0]
1170 store.x = load.y1200 store.x = load.y
1171 store.y.axis = [z0]1201 store.y.axis = [z0]
1172 store.y.strides = [ascir.SizeExpr(1)]1202 store.y.strides = [ascir.SizeExpr(1)]
1173 store.y.size = [s0]1203 store.y.size = [s0]
1174- ws = ascir.ops.Workspace('buf8')1204+ ws = ascir.ops.Workspace("buf8")
1175 ws.attr.sched.axis = [z0]1205 ws.attr.sched.axis = [z0]
1176 ws.x = store.y1206 ws.x = store.y
1177 ws.y.size = [s0]1207 ws.y.size = [s0]
1178 ws.y.dtype = ascir.dtypes.float161208 ws.y.dtype = ascir.dtypes.float16
1179 ws.y.axis = [z0]1209 ws.y.axis = [z0]
1180 ws.y.strides = [ascir.SizeExpr(1)]1210 ws.y.strides = [ascir.SizeExpr(1)]
1181- load1 = ascir.ops.Load('load1')1211+ load1 = ascir.ops.Load("load1")
1182 load1.attr.sched.axis = [z0]1212 load1.attr.sched.axis = [z0]
1183 load1.x = ws.y1213 load1.x = ws.y
1184 load1.y.axis = [z0]1214 load1.y.axis = [z0]
1185 load1.y.strides = [ascir.SizeExpr(1)]1215 load1.y.strides = [ascir.SizeExpr(1)]
1186 load1.y.size = [s0]1216 load1.y.size = [s0]
1187 1217 
1188- store1 = ascir.ops.Store('store1')1218+ store1 = ascir.ops.Store("store1")
1189 store1.attr.sched.axis = [z0]1219 store1.attr.sched.axis = [z0]
1190 store1.x = load1.y1220 store1.x = load1.y
1191 store1.y.axis = [z0]1221 store1.y.axis = [z0]
1192 store1.y.strides = [ascir.SizeExpr(1)]1222 store1.y.strides = [ascir.SizeExpr(1)]
1193 store1.y.size = [s0]1223 store1.y.size = [s0]
1194 1224 
1195- ws1 = ascir.ops.Workspace('buf2')1225+ ws1 = ascir.ops.Workspace("buf2")
1196 ws1.attr.sched.axis = [z0]1226 ws1.attr.sched.axis = [z0]
1197 ws1.x = store1.y1227 ws1.x = store1.y
1198 ws1.y.size = [s0]1228 ws1.y.size = [s0]
@@ -1200,21 +1230,21 @@ class TestWorkspaceOptimize():
1200 ws1.y.axis = [z0]1230 ws1.y.axis = [z0]
1201 ws1.y.strides = [ascir.SizeExpr(1)]1231 ws1.y.strides = [ascir.SizeExpr(1)]
1202 1232 
1203- load2 = ascir.ops.Load('load2')1233+ load2 = ascir.ops.Load("load2")
1204 load2.attr.sched.axis = [z0]1234 load2.attr.sched.axis = [z0]
1205 load2.x = ws1.y1235 load2.x = ws1.y
1206 load2.y.axis = [z0]1236 load2.y.axis = [z0]
1207 load2.y.strides = [ascir.SizeExpr(1)]1237 load2.y.strides = [ascir.SizeExpr(1)]
1208 load2.y.size = [s0]1238 load2.y.size = [s0]
1209 1239 
1210- load3 = ascir.ops.Load('load3')1240+ load3 = ascir.ops.Load("load3")
1211 load3.attr.sched.axis = [z0]1241 load3.attr.sched.axis = [z0]
1212 load3.x = ws1.y1242 load3.x = ws1.y
1213 load3.y.axis = [z0]1243 load3.y.axis = [z0]
1214 load3.y.strides = [ascir.SizeExpr(1)]1244 load3.y.strides = [ascir.SizeExpr(1)]
1215 load3.y.size = [s0]1245 load3.y.size = [s0]
1216- NpuKernel0Graph.infer_dtypes()1246+ npu_kernel0_graph.infer_dtypes()
1217- return NpuKernel0Graph1247+ return npu_kernel0_graph
1218 1248 
1219 def test_construct_graph(self):1249 def test_construct_graph(self):
1220 graph = self.construct_graph()1250 graph = self.construct_graph()
@@ -1226,9 +1256,10 @@ class TestWorkspaceOptimize():
1226 fuser = Autofuser(options)1256 fuser = Autofuser(options)
1227 1257 
1228 hint_graph = self.construct_graph()1258 hint_graph = self.construct_graph()
1229- schedule_results = fuser.schedule(hint_graph)1259+ fuser.schedule(hint_graph)
1230 1260 
1231-class TestCodeGenLoadAbsStore():1261+ 
1262+class TestCodeGenLoadAbsStore:
1232 @staticmethod1263 @staticmethod
1233 def construct_graph():1264 def construct_graph():
1234 graph = ascir.HintGraph("LoadAbsStore")1265 graph = ascir.HintGraph("LoadAbsStore")
@@ -1298,7 +1329,7 @@ class TestCodeGenLoadAbsStore():
1298 scheduler = Schedule(options)1329 scheduler = Schedule(options)
1299 1330 
1300 hint_graph = self.construct_graph()1331 hint_graph = self.construct_graph()
1301- impl_graphs = scheduler.schedule(hint_graph)1332+ scheduler.schedule(hint_graph)
1302 1333 
1303 @pytest.mark.skip1334 @pytest.mark.skip
1304 def test_codegen(self):1335 def test_codegen(self):
@@ -1307,9 +1338,13 @@ class TestCodeGenLoadAbsStore():
1307 1338 
1308 hint_graph = self.construct_graph()1339 hint_graph = self.construct_graph()
1309 impl_graphs = scheduler.schedule(hint_graph)1340 impl_graphs = scheduler.schedule(hint_graph)
1310- shape_info = ascir.ShapeInfo({"s0": "GetDimValueFromGraphInputData(0, 0);",1341+ shape_info = ascir.ShapeInfo(
1311- "s1": "GetDimValueFromGraphInputData(0, 1);",1342+ {
1312- "s2": "GetDimValueFromGraphInputData(1, 0);"})1343+ "s0": "GetDimValueFromGraphInputData(0, 0);",
1344+ "s1": "GetDimValueFromGraphInputData(0, 1);",
1345+ "s2": "GetDimValueFromGraphInputData(1, 0);",
1346+ }
1347+ )
1313 1348 
1314 kernel_path = "./fused_graph_kernel.o"1349 kernel_path = "./fused_graph_kernel.o"
1315 with open(kernel_path, "wb") as o_file:1350 with open(kernel_path, "wb") as o_file:
@@ -1319,73 +1354,79 @@ class TestCodeGenLoadAbsStore():
1319 "name": "Alice",1354 "name": "Alice",
1320 "age": 30,1355 "age": 30,
1321 "is_student": False,1356 "is_student": False,
1322- "courses": ["Math", "Science", "History"]1357+ "courses": ["Math", "Science", "History"],
1323 }1358 }
1324 json_path = "./fused_graph_kernel.json"1359 json_path = "./fused_graph_kernel.json"
1325 with open(json_path, "w") as json_file:1360 with open(json_path, "w") as json_file:
1326 json.dump(data, json_file, indent=4)1361 json.dump(data, json_file, indent=4)
1327 1362 
1328 tiling_data, op_kernel = codegen.device_code_generator(hint_graph, impl_graphs)1363 tiling_data, op_kernel = codegen.device_code_generator(hint_graph, impl_graphs)
1329- tiling, infer_shape = codegen.host_code_generator(hint_graph, impl_graphs, shape_info, "", ["", ""])1364+ tiling, infer_shape = codegen.host_code_generator(
1365+ hint_graph, impl_graphs, shape_info, "", ["", ""]
1366+ )
1330 get_kernel = codegen.get_kernel_and_json_generator(kernel_path, json_path)1367 get_kernel = codegen.get_kernel_and_json_generator(kernel_path, json_path)
1331 os.remove(kernel_path)1368 os.remove(kernel_path)
1332 os.remove(json_path)1369 os.remove(json_path)
1333- assert tiling_data == "".join([1370+ assert tiling_data == "".join(
1334- "#ifndef __Autofuse_Tiling_Data_H__\n"1371+ [
1335- "#define __Autofuse_Tiling_Data_H__\n"1372+ "#ifndef __Autofuse_Tiling_Data_H__\n"
1336- "#include <stdint.h>\n"1373+ "#define __Autofuse_Tiling_Data_H__\n"
1337- "#include \"kernel_tiling/kernel_tiling.h\"\n"1374+ "#include <stdint.h>\n"
1338- "#define BEGIN_TILING_DATA_DEF_T(name) struct name {\n"1375+ '#include "kernel_tiling/kernel_tiling.h"\n'
1339- "#define TILING_DATA_FIELD_DEF_T(type, name) \\\n"1376+ "#define BEGIN_TILING_DATA_DEF_T(name) struct name {\n"
1340- " type name; \\\n"1377+ "#define TILING_DATA_FIELD_DEF_T(type, name) \\\n"
1341- " inline void set_##name(type value) { name = value; } \\\n",1378+ " type name; \\\n"
1342- " inline type get_##name() { return name; } \\\n"1379+ " inline void set_##name(type value) { name = value; } \\\n",
1343- " inline type* get_addr_##name() {return &name;}\n"1380+ " inline type get_##name() { return name; } \\\n"
1344- "#define END_TILING_DATA_DEF_T };\n"1381+ " inline type* get_addr_##name() {return &name;}\n"
1345- "#define TILING_DATA_FIELD_DEF_T_STRUCT(struct_type, filed_name) \\\n"1382+ "#define END_TILING_DATA_DEF_T };\n"
1346- " struct_type filed_name;\n\n"1383+ "#define TILING_DATA_FIELD_DEF_T_STRUCT(struct_type, filed_name) \\\n"
1347- "BEGIN_TILING_DATA_DEF_T(AutofuseTilingData)\n"1384+ " struct_type filed_name;\n\n"
1348- " TILING_DATA_FIELD_DEF_T(uint32_t, block_dim);\n"1385+ "BEGIN_TILING_DATA_DEF_T(AutofuseTilingData)\n"
1349- " TILING_DATA_FIELD_DEF_T(uint32_t, corenum);\n"1386+ " TILING_DATA_FIELD_DEF_T(uint32_t, block_dim);\n"
1350- " TILING_DATA_FIELD_DEF_T(uint32_t, ub_size);\n"1387+ " TILING_DATA_FIELD_DEF_T(uint32_t, corenum);\n"
1351- " TILING_DATA_FIELD_DEF_T(uint32_t, hbm_size);\n"1388+ " TILING_DATA_FIELD_DEF_T(uint32_t, ub_size);\n"
1352- " TILING_DATA_FIELD_DEF_T(uint32_t, tiling_key);\n"1389+ " TILING_DATA_FIELD_DEF_T(uint32_t, hbm_size);\n"
1353- " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2t_size);\n"1390+ " TILING_DATA_FIELD_DEF_T(uint32_t, tiling_key);\n"
1354- " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2Tb_size);\n"1391+ " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2t_size);\n"
1355- "END_TILING_DATA_DEF_T;\n\n"1392+ " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2Tb_size);\n"
1356- "struct AutofuseTilingDataPerf {\n"1393+ "END_TILING_DATA_DEF_T;\n\n"
1357- " AutofuseTilingData tiling_data;\n"1394+ "struct AutofuseTilingDataPerf {\n"
1358- " double best_perf;\n"1395+ " AutofuseTilingData tiling_data;\n"
1359- "};\n"1396+ " double best_perf;\n"
1360- "#endif\n"1397+ "};\n"
1361- ])1398+ "#endif\n",
1399+ ]
1400+ )
1362 1401 
1363- assert infer_shape == "".join([1402+ assert infer_shape == "".join([""])
1364- ""])
1365 1403 
1366- assert get_kernel == "".join([1404+ assert get_kernel == "".join(
1367- "#include <cstdint>\n"1405+ [
1368- "#include <cstring>\n"1406+ "#include <cstdint>\n"
1369- "#include <vector>\n"1407+ "#include <cstring>\n"
1370- "extern \"C\" void GetKernelBin(std::vector<char> &kernel_bin) {\n"1408+ "#include <vector>\n"
1371- " std::vector<uint8_t> temp_kernel = {\n"1409+ 'extern "C" void GetKernelBin(std::vector<char> &kernel_bin) {\n'
1372- " 84, 104, 105, 115, 32, 105, 115, 32, 97, 32, 46, 111, 32, 102, 105, 108, 101, 32, 99, 111, \n"1410+ " std::vector<uint8_t> temp_kernel = {\n"
1373- " 110, 116, 101, 110, 116, 46, };\n"1411+ " 84, 104, 105, 115, 32, 105, 115, 32, 97, 32, 46, 111, 32, 102, 105, 108, 101, 32, 99, 111, \n"
1374- " kernel_bin.resize(temp_kernel.size());\n"1412+ " 110, 116, 101, 110, 116, 46, };\n"
1375- " std::memcpy(kernel_bin.data(), temp_kernel.data(), temp_kernel.size() * sizeof(uint8_t));\n"1413+ " kernel_bin.resize(temp_kernel.size());\n"
1376- "}"])1414+ " std::memcpy(kernel_bin.data(), temp_kernel.data(), temp_kernel.size() * sizeof(uint8_t));\n"
1415+ "}"
1416+ ]
1417+ )
1377 1418 
1378 1419 
1379-class TestComputeGraphInput():1420+class TestComputeGraphInput:
1380 @staticmethod1421 @staticmethod
1381 def construct_compute_graph():1422 def construct_compute_graph():
1382 test_graph = os.path.join(PYF_PATH, "test_graph.txt")1423 test_graph = os.path.join(PYF_PATH, "test_graph.txt")
1383- with open(test_graph, 'r', encoding='utf-8') as file:1424+ with open(test_graph, "r", encoding="utf-8") as file:
1384 content = file.read()1425 content = file.read()
1385 compute_graph = ascir.utils.deserialize("compute_graph", content)1426 compute_graph = ascir.utils.deserialize("compute_graph", content)
1386 print(compute_graph.get_name(), flush=True)1427 print(compute_graph.get_name(), flush=True)
1387 print(compute_graph.get_info(), flush=True)1428 print(compute_graph.get_info(), flush=True)
1388- assert compute_graph != None1429+ assert compute_graph is not None
1389 return compute_graph1430 return compute_graph
1390 1431 
1391 @pytest.mark.skip1432 @pytest.mark.skip
@@ -1394,7 +1435,7 @@ class TestComputeGraphInput():
1394 scheduler = Schedule(options)1435 scheduler = Schedule(options)
1395 1436 
1396 compute_graph = self.construct_compute_graph()1437 compute_graph = self.construct_compute_graph()
1397- schedule_results = scheduler.scheduleV2(compute_graph)1438+ scheduler.scheduleV2(compute_graph)
1398 1439 
1399 def test_scheduleV2_fail(self):1440 def test_scheduleV2_fail(self):
1400 options = AutofuserOptions()1441 options = AutofuserOptions()
@@ -1403,8 +1444,9 @@ class TestComputeGraphInput():
1403 compute_graph = ascir.HintComputeGraph("test")1444 compute_graph = ascir.HintComputeGraph("test")
1404 try:1445 try:
1405 scheduler.scheduleV2(compute_graph)1446 scheduler.scheduleV2(compute_graph)
1406- except RuntimeError as e:1447+ except RuntimeError:
1407 pass1448 pass
1449+ 
1408 @pytest.mark.skip1450 @pytest.mark.skip
1409 def test_computegraph_codegen(self):1451 def test_computegraph_codegen(self):
1410 scheduler = Schedule()1452 scheduler = Schedule()
@@ -1412,9 +1454,13 @@ class TestComputeGraphInput():
1412 1454 
1413 compute_graph = self.construct_compute_graph()1455 compute_graph = self.construct_compute_graph()
1414 schedule_results = scheduler.scheduleV2(compute_graph)1456 schedule_results = scheduler.scheduleV2(compute_graph)
1415- shape_info = ascir.ShapeInfo({"s0": "GetDimValueFromGraphInputData(0, 0);",1457+ shape_info = ascir.ShapeInfo(
1416- "s1": "GetDimValueFromGraphInputData(0, 1);",1458+ {
1417- "s2": "GetDimValueFromGraphInputData(1, 0);"})1459+ "s0": "GetDimValueFromGraphInputData(0, 0);",
1460+ "s1": "GetDimValueFromGraphInputData(0, 1);",
1461+ "s2": "GetDimValueFromGraphInputData(1, 0);",
1462+ }
1463+ )
1418 1464 
1419 kernel_path = "./fused_graph_kernel.o"1465 kernel_path = "./fused_graph_kernel.o"
1420 with open(kernel_path, "wb") as o_file:1466 with open(kernel_path, "wb") as o_file:
@@ -1424,90 +1470,98 @@ class TestComputeGraphInput():
1424 "name": "Alice",1470 "name": "Alice",
1425 "age": 30,1471 "age": 30,
1426 "is_student": False,1472 "is_student": False,
1427- "courses": ["Math", "Science", "History"]1473+ "courses": ["Math", "Science", "History"],
1428 }1474 }
1429 json_path = "./fused_graph_kernel.json"1475 json_path = "./fused_graph_kernel.json"
1430 with open(json_path, "w") as json_file:1476 with open(json_path, "w") as json_file:
1431 json.dump(data, json_file, indent=4)1477 json.dump(data, json_file, indent=4)
1432 1478 
1433 tiling_data, op_kernel = codegen.device_code_generator(schedule_results)1479 tiling_data, op_kernel = codegen.device_code_generator(schedule_results)
1434- assert tiling_data == "".join([1480+ assert tiling_data == "".join(
1435- "#ifndef __Autofuse_Tiling_Data_H__\n"1481+ [
1436- "#define __Autofuse_Tiling_Data_H__\n"1482+ "#ifndef __Autofuse_Tiling_Data_H__\n"
1437- "#include <stdint.h>\n"1483+ "#define __Autofuse_Tiling_Data_H__\n"
1438- "#include \"kernel_tiling/kernel_tiling.h\"\n"1484+ "#include <stdint.h>\n"
1439- "#define BEGIN_TILING_DATA_DEF_T(name) struct name {\n"1485+ '#include "kernel_tiling/kernel_tiling.h"\n'
1440- "#define TILING_DATA_FIELD_DEF_T(type, name) \\\n"1486+ "#define BEGIN_TILING_DATA_DEF_T(name) struct name {\n"
1441- " type name; \\\n"1487+ "#define TILING_DATA_FIELD_DEF_T(type, name) \\\n"
1442- " inline void set_##name(type value) { name = value; } \\\n",1488+ " type name; \\\n"
1443- " inline type get_##name() { return name; } \\\n"1489+ " inline void set_##name(type value) { name = value; } \\\n",
1444- " inline type* get_addr_##name() {return &name;}\n"1490+ " inline type get_##name() { return name; } \\\n"
1445- "#define END_TILING_DATA_DEF_T };\n"1491+ " inline type* get_addr_##name() {return &name;}\n"
1446- "#define TILING_DATA_FIELD_DEF_T_STRUCT(struct_type, filed_name) \\\n"1492+ "#define END_TILING_DATA_DEF_T };\n"
1447- " struct_type filed_name;\n\n"1493+ "#define TILING_DATA_FIELD_DEF_T_STRUCT(struct_type, filed_name) \\\n"
1448- "BEGIN_TILING_DATA_DEF_T(AutofuseTilingData)\n"1494+ " struct_type filed_name;\n\n"
1449- " TILING_DATA_FIELD_DEF_T(uint32_t, block_dim);\n"1495+ "BEGIN_TILING_DATA_DEF_T(AutofuseTilingData)\n"
1450- " TILING_DATA_FIELD_DEF_T(uint32_t, corenum);\n"1496+ " TILING_DATA_FIELD_DEF_T(uint32_t, block_dim);\n"
1451- " TILING_DATA_FIELD_DEF_T(uint32_t, ub_size);\n"1497+ " TILING_DATA_FIELD_DEF_T(uint32_t, corenum);\n"
1452- " TILING_DATA_FIELD_DEF_T(uint32_t, hbm_size);\n"1498+ " TILING_DATA_FIELD_DEF_T(uint32_t, ub_size);\n"
1453- " TILING_DATA_FIELD_DEF_T(uint32_t, tiling_key);\n"1499+ " TILING_DATA_FIELD_DEF_T(uint32_t, hbm_size);\n"
1454- " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2t_size);\n"1500+ " TILING_DATA_FIELD_DEF_T(uint32_t, tiling_key);\n"
1455- " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2Tb_size);\n"1501+ " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2t_size);\n"
1456- " TILING_DATA_FIELD_DEF_T(uint32_t, q0_size);\n"1502+ " TILING_DATA_FIELD_DEF_T(uint32_t, z0z1z2Tb_size);\n"
1457- " TILING_DATA_FIELD_DEF_T(uint32_t, q1_size);\n"1503+ " TILING_DATA_FIELD_DEF_T(uint32_t, q0_size);\n"
1458- " TILING_DATA_FIELD_DEF_T(uint32_t, b0_size);\n"1504+ " TILING_DATA_FIELD_DEF_T(uint32_t, q1_size);\n"
1459- "END_TILING_DATA_DEF_T;\n\n"1505+ " TILING_DATA_FIELD_DEF_T(uint32_t, b0_size);\n"
1460- "struct AutofuseTilingDataPerf {\n"1506+ "END_TILING_DATA_DEF_T;\n\n"
1461- " AutofuseTilingData tiling_data;\n"1507+ "struct AutofuseTilingDataPerf {\n"
1462- " double best_perf;\n"1508+ " AutofuseTilingData tiling_data;\n"
1463- "};\n"1509+ " double best_perf;\n"
1464- "#endif\n"1510+ "};\n"
1465- ])1511+ "#endif\n",
1512+ ]
1513+ )
1466 1514 
1467 output_shape = [["s0", "s1"]]1515 output_shape = [["s0", "s1"]]
1468 vector_core_num = "0"1516 vector_core_num = "0"
1469- tiling, infer_shape = codegen.host_code_generator(schedule_results, shape_info, output_shape, "", vector_core_num)1517+ tiling, infer_shape = codegen.host_code_generator(
1470- pgo_src = codegen.pgo_code_generator(schedule_results, "")1518+ schedule_results, shape_info, output_shape, "", vector_core_num
1519+ )
1520+ codegen.pgo_code_generator(schedule_results, "")
1471 get_kernel = codegen.get_kernel_and_json_generator(kernel_path, json_path)1521 get_kernel = codegen.get_kernel_and_json_generator(kernel_path, json_path)
1472 os.remove(kernel_path)1522 os.remove(kernel_path)
1473 os.remove(json_path)1523 os.remove(json_path)
1474- assert get_kernel == "".join([1524+ assert get_kernel == "".join(
1475- "#include <cstdint>\n"1525+ [
1476- "#include <cstring>\n"1526+ "#include <cstdint>\n"
1477- "#include <vector>\n"1527+ "#include <cstring>\n"
1478- "extern \"C\" void GetKernelBin(std::vector<char> &kernel_bin) {\n"1528+ "#include <vector>\n"
1479- " std::vector<uint8_t> temp_kernel = {\n"1529+ 'extern "C" void GetKernelBin(std::vector<char> &kernel_bin) {\n'
1480- " 84, 104, 105, 115, 32, 105, 115, 32, 97, 32, 46, 111, 32, 102, 105, 108, 101, 32, 99, 111, \n"1530+ " std::vector<uint8_t> temp_kernel = {\n"
1481- " 110, 116, 101, 110, 116, 46, };\n"1531+ " 84, 104, 105, 115, 32, 105, 115, 32, 97, 32, 46, 111, 32, 102, 105, 108, 101, 32, 99, 111, \n"
1482- " kernel_bin.resize(temp_kernel.size());\n"1532+ " 110, 116, 101, 110, 116, 46, };\n"
1483- " std::memcpy(kernel_bin.data(), temp_kernel.data(), temp_kernel.size() * sizeof(uint8_t));\n"1533+ " kernel_bin.resize(temp_kernel.size());\n"
1484- "}"])1534+ " std::memcpy(kernel_bin.data(), temp_kernel.data(), temp_kernel.size() * sizeof(uint8_t));\n"
1535+ "}"
1536+ ]
1537+ )
1485 1538 
1486 try:1539 try:
1487 output_shape = ["s0"]1540 output_shape = ["s0"]
1488 vector_core_num = "0"1541 vector_core_num = "0"
1489- tiling, infer_shape = codegen.host_code_generator(schedule_results, shape_info,1542+ tiling, infer_shape = codegen.host_code_generator(
1490- output_shape, "", vector_core_num)1543+ schedule_results, shape_info, output_shape, "", vector_core_num
1491- except ValueError as e:1544+ )
1545+ except ValueError:
1492 pass1546 pass
1493 1547 
1494 try:1548 try:
1495- pgo_src = codegen.pgo_code_generator(schedule_results)1549+ codegen.pgo_code_generator(schedule_results)
1496- except ValueError as e:1550+ except ValueError:
1497 pass1551 pass
1498 1552 
1499 try:1553 try:
1500 get_kernel = codegen.get_kernel_and_json_generator(kernel_path, json_path)1554 get_kernel = codegen.get_kernel_and_json_generator(kernel_path, json_path)
1501- except ValueError as e:1555+ except ValueError:
1502 pass1556 pass
1503 1557 
1504 try:1558 try:
1505 get_kernel = codegen.get_kernel_and_json_generator(kernel_path)1559 get_kernel = codegen.get_kernel_and_json_generator(kernel_path)
1506- except ValueError as e:1560+ except ValueError:
1507 pass1561 pass
1508 1562 
1509 1563 
1510-class TestHintGraph():1564+class TestHintGraph:
1511 @staticmethod1565 @staticmethod
1512 def construct_graph():1566 def construct_graph():
1513 graph = ascir.HintGraph("LoadAbsStore")1567 graph = ascir.HintGraph("LoadAbsStore")
@@ -1567,35 +1621,35 @@ class TestHintGraph():
1567 asc_graph = self.construct_graph()1621 asc_graph = self.construct_graph()
1568 try:1622 try:
1569 asc_graph.set_name(2)1623 asc_graph.set_name(2)
1570- except TypeError as e:1624+ except TypeError:
1571 assert asc_graph.get_name() == "".join(["LoadAbsStore"])1625 assert asc_graph.get_name() == "".join(["LoadAbsStore"])
1572 1626 
1573 asc_graph.set_name("test_graph")1627 asc_graph.set_name("test_graph")
1574 assert asc_graph.get_name() == "".join(["test_graph"])1628 assert asc_graph.get_name() == "".join(["test_graph"])
1575 1629 
1576 1630 
1577-class TestFusedGraph():1631+class TestFusedGraph:
1578 @staticmethod1632 @staticmethod
1579 def construct_add_ascgraph(name: str) -> ascir.HintGraph:1633 def construct_add_ascgraph(name: str) -> ascir.HintGraph:
1580- NpuKernel0Graph = ascir.HintGraph(name)1634+ npu_kernel0_graph = ascir.HintGraph(name)
1581- s0 = NpuKernel0Graph.create_size("s0")1635+ s0 = npu_kernel0_graph.create_size("s0")
1582- s1 = NpuKernel0Graph.create_size("s1")1636+ s1 = npu_kernel0_graph.create_size("s1")
1583- z0 = NpuKernel0Graph.create_axis("z0", s0)1637+ z0 = npu_kernel0_graph.create_axis("z0", s0)
1584- z1 = NpuKernel0Graph.create_axis("z1", s1)1638+ z1 = npu_kernel0_graph.create_axis("z1", s1)
1585- sub_data0 = ascir.ops.Data('sub_data0', NpuKernel0Graph)1639+ sub_data0 = ascir.ops.Data("sub_data0", npu_kernel0_graph)
1586 sub_data0.y.dtype = ascir.dtypes.float161640 sub_data0.y.dtype = ascir.dtypes.float16
1587 sub_data0.attr.ir_attr.index = 01641 sub_data0.attr.ir_attr.index = 0
1588- load0 = ascir.ops.Load('load')1642+ load0 = ascir.ops.Load("load")
1589 load0.attr.ir_attr.offset = 01643 load0.attr.ir_attr.offset = 0
1590 load0.attr.sched.axis = [z0, z1]1644 load0.attr.sched.axis = [z0, z1]
1591 load0.x = sub_data0.y1645 load0.x = sub_data0.y
1592 load0.y.axis = [z0, z1]1646 load0.y.axis = [z0, z1]
1593 load0.y.strides = [s1, ascir.SizeExpr(1)]1647 load0.y.strides = [s1, ascir.SizeExpr(1)]
1594 load0.y.size = [s0, s1]1648 load0.y.size = [s0, s1]
1595- sub_data1 = ascir.ops.Data('sub_data1', NpuKernel0Graph)1649+ sub_data1 = ascir.ops.Data("sub_data1", npu_kernel0_graph)
1596 sub_data1.y.dtype = ascir.dtypes.float161650 sub_data1.y.dtype = ascir.dtypes.float16
1597 sub_data1.attr.ir_attr.index = 11651 sub_data1.attr.ir_attr.index = 1
1598- load1 = ascir.ops.Load('load')1652+ load1 = ascir.ops.Load("load")
1599 load1.attr.ir_attr.offset = ascir.SizeExpr(0)1653 load1.attr.ir_attr.offset = ascir.SizeExpr(0)
1600 load1.attr.sched.axis = [z0, z1]1654 load1.attr.sched.axis = [z0, z1]
1601 load1.x = sub_data1.y1655 load1.x = sub_data1.y
@@ -1603,7 +1657,7 @@ class TestFusedGraph():
1603 load1.y.strides = [s1, ascir.SizeExpr(1)]1657 load1.y.strides = [s1, ascir.SizeExpr(1)]
1604 load1.y.size = [s0, s1]1658 load1.y.size = [s0, s1]
1605 1659 
1606- add0 = ascir.ops.Add('add')1660+ add0 = ascir.ops.Add("add")
1607 add0.attr.sched.axis = [z0, z1]1661 add0.attr.sched.axis = [z0, z1]
1608 add0.x1 = load0.y1662 add0.x1 = load0.y
1609 add0.x2 = load1.y1663 add0.x2 = load1.y
@@ -1611,52 +1665,52 @@ class TestFusedGraph():
1611 add0.y.strides = [s1 + s1, ascir.SizeExpr(1)]1665 add0.y.strides = [s1 + s1, ascir.SizeExpr(1)]
1612 add0.y.size = [s0, s1 * 2]1666 add0.y.size = [s0, s1 * 2]
1613 1667 
1614- store0 = ascir.ops.Store('store')1668+ store0 = ascir.ops.Store("store")
1615 store0.attr.ir_attr.offset = ascir.SizeExpr(0)1669 store0.attr.ir_attr.offset = ascir.SizeExpr(0)
1616 store0.attr.sched.axis = [z0, z1]1670 store0.attr.sched.axis = [z0, z1]
1617 store0.x = add0.y1671 store0.x = add0.y
1618 store0.y.axis = [z0, z1]1672 store0.y.axis = [z0, z1]
1619- store0.y.strides = [s1 ** 2, ascir.SizeExpr(1)]1673+ store0.y.strides = [s1**2, ascir.SizeExpr(1)]
1620 store0.y.size = [s0, s1 * 2]1674 store0.y.size = [s0, s1 * 2]
1621 1675 
1622- store1 = ascir.ops.Store('store')1676+ store1 = ascir.ops.Store("store")
1623 store1.attr.ir_attr.offset = ascir.SizeExpr(10)1677 store1.attr.ir_attr.offset = ascir.SizeExpr(10)
1624 store1.attr.sched.axis = [z0, z1]1678 store1.attr.sched.axis = [z0, z1]
1625 store1.x = add0.y1679 store1.x = add0.y
1626 store1.y.axis = [z0, z1]1680 store1.y.axis = [z0, z1]
1627 store1.y.strides = [s1 * 2, ascir.SizeExpr(1)]1681 store1.y.strides = [s1 * 2, ascir.SizeExpr(1)]
1628 store1.y.size = [s0, s1 * 2]1682 store1.y.size = [s0, s1 * 2]
1629- buf0 = ascir.ops.Output('buf0')1683+ buf0 = ascir.ops.Output("buf0")
1630 buf0.attr.ir_attr.index = 01684 buf0.attr.ir_attr.index = 0
1631 # store0, strore1 写到同一个output上，偏移不同1685 # store0, strore1 写到同一个output上，偏移不同
1632 buf0.x = [store0.y, store1]1686 buf0.x = [store0.y, store1]
1633 buf0.y.dtype = ascir.dtypes.float161687 buf0.y.dtype = ascir.dtypes.float16
1634- buf1 = ascir.ops.Output('buf1')1688+ buf1 = ascir.ops.Output("buf1")
1635 buf1.attr.ir_attr.index = 11689 buf1.attr.ir_attr.index = 1
1636 buf1.x = store1.y1690 buf1.x = store1.y
1637- NpuKernel0Graph.infer_dtypes()1691+ npu_kernel0_graph.infer_dtypes()
1638- ascir.utils.dump(NpuKernel0Graph)1692+ ascir.utils.dump(npu_kernel0_graph)
1639- return NpuKernel0Graph1693+ return npu_kernel0_graph
1640 1694 
1641 @staticmethod1695 @staticmethod
1642 def construct_add_ascgraph_without_data(name: str) -> ascir.HintGraph:1696 def construct_add_ascgraph_without_data(name: str) -> ascir.HintGraph:
1643- NpuKernel0Graph = ascir.HintGraph(name)1697+ npu_kernel0_graph = ascir.HintGraph(name)
1644- s0 = NpuKernel0Graph.create_size("s0")1698+ s0 = npu_kernel0_graph.create_size("s0")
1645- s1 = NpuKernel0Graph.create_size("s1")1699+ s1 = npu_kernel0_graph.create_size("s1")
1646- z0 = NpuKernel0Graph.create_axis("z0", s0)1700+ z0 = npu_kernel0_graph.create_axis("z0", s0)
1647- z1 = NpuKernel0Graph.create_axis("z1", s1)1701+ z1 = npu_kernel0_graph.create_axis("z1", s1)
1648- sub_data0 = ascir.ops.Scalar('sub_data0', NpuKernel0Graph)1702+ sub_data0 = ascir.ops.Scalar("sub_data0", npu_kernel0_graph)
1649 sub_data0.y.dtype = ascir.dtypes.float161703 sub_data0.y.dtype = ascir.dtypes.float16
1650- load0 = ascir.ops.Load('load')1704+ load0 = ascir.ops.Load("load")
1651 load0.attr.ir_attr.offset = 01705 load0.attr.ir_attr.offset = 0
1652 load0.attr.sched.axis = [z0, z1]1706 load0.attr.sched.axis = [z0, z1]
1653 load0.x = sub_data0.y1707 load0.x = sub_data0.y
1654 load0.y.axis = [z0, z1]1708 load0.y.axis = [z0, z1]
1655 load0.y.strides = [s1, ascir.SizeExpr(1)]1709 load0.y.strides = [s1, ascir.SizeExpr(1)]
1656 load0.y.size = [s0, s1]1710 load0.y.size = [s0, s1]
1657- sub_data1 = ascir.ops.Scalar('sub_data1', NpuKernel0Graph)1711+ sub_data1 = ascir.ops.Scalar("sub_data1", npu_kernel0_graph)
1658 sub_data1.y.dtype = ascir.dtypes.float161712 sub_data1.y.dtype = ascir.dtypes.float16
1659- load1 = ascir.ops.Load('load')1713+ load1 = ascir.ops.Load("load")
1660 load1.attr.ir_attr.offset = ascir.SizeExpr(0)1714 load1.attr.ir_attr.offset = ascir.SizeExpr(0)
1661 load1.attr.sched.axis = [z0, z1]1715 load1.attr.sched.axis = [z0, z1]
1662 load1.x = sub_data1.y1716 load1.x = sub_data1.y
@@ -1664,7 +1718,7 @@ class TestFusedGraph():
1664 load1.y.strides = [s1, ascir.SizeExpr(1)]1718 load1.y.strides = [s1, ascir.SizeExpr(1)]
1665 load1.y.size = [s0, s1]1719 load1.y.size = [s0, s1]
1666 1720 
1667- add0 = ascir.ops.Add('add')1721+ add0 = ascir.ops.Add("add")
1668 add0.attr.sched.axis = [z0, z1]1722 add0.attr.sched.axis = [z0, z1]
1669 add0.x1 = load0.y1723 add0.x1 = load0.y
1670 add0.x2 = load1.y1724 add0.x2 = load1.y
@@ -1672,134 +1726,158 @@ class TestFusedGraph():
1672 add0.y.strides = [s1 + s1, ascir.SizeExpr(1)]1726 add0.y.strides = [s1 + s1, ascir.SizeExpr(1)]
1673 add0.y.size = [s0, s1 * 2]1727 add0.y.size = [s0, s1 * 2]
1674 1728 
1675- store0 = ascir.ops.Store('store')1729+ store0 = ascir.ops.Store("store")
1676 store0.attr.ir_attr.offset = ascir.SizeExpr(0)1730 store0.attr.ir_attr.offset = ascir.SizeExpr(0)
1677 store0.attr.sched.axis = [z0, z1]1731 store0.attr.sched.axis = [z0, z1]
1678 store0.x = add0.y1732 store0.x = add0.y
1679 store0.y.axis = [z0, z1]1733 store0.y.axis = [z0, z1]
1680- store0.y.strides = [s1 ** 2, ascir.SizeExpr(1)]1734+ store0.y.strides = [s1**2, ascir.SizeExpr(1)]
1681 store0.y.size = [s0, s1 * 2]1735 store0.y.size = [s0, s1 * 2]
1682 1736 
1683- store1 = ascir.ops.Store('store')1737+ store1 = ascir.ops.Store("store")
1684 store1.attr.ir_attr.offset = ascir.SizeExpr(10)1738 store1.attr.ir_attr.offset = ascir.SizeExpr(10)
1685 store1.attr.sched.axis = [z0, z1]1739 store1.attr.sched.axis = [z0, z1]
1686 store1.x = add0.y1740 store1.x = add0.y
1687 store1.y.axis = [z0, z1]1741 store1.y.axis = [z0, z1]
1688 store1.y.strides = [s1 * 2, ascir.SizeExpr(1)]1742 store1.y.strides = [s1 * 2, ascir.SizeExpr(1)]
1689 store1.y.size = [s0, s1 * 2]1743 store1.y.size = [s0, s1 * 2]
1690- buf0 = ascir.ops.Output('buf0')1744+ buf0 = ascir.ops.Output("buf0")
1691 buf0.attr.ir_attr.index = 01745 buf0.attr.ir_attr.index = 0
1692 # store0, strore1 写到同一个output上，偏移不同1746 # store0, strore1 写到同一个output上，偏移不同
1693 buf0.x = [store0.y, store1]1747 buf0.x = [store0.y, store1]
1694 buf0.y.dtype = ascir.dtypes.float161748 buf0.y.dtype = ascir.dtypes.float16
1695- buf1 = ascir.ops.Output('buf1')1749+ buf1 = ascir.ops.Output("buf1")
1696 buf1.attr.ir_attr.index = 11750 buf1.attr.ir_attr.index = 1
1697 buf1.x = store1.y1751 buf1.x = store1.y
1698- NpuKernel0Graph.infer_dtypes()1752+ npu_kernel0_graph.infer_dtypes()
1699- ascir.utils.dump(NpuKernel0Graph)1753+ ascir.utils.dump(npu_kernel0_graph)
1700- return NpuKernel0Graph1754+ return npu_kernel0_graph
1701 1755 
1702 def test_fused_graph_construct_and_dump_with_ascbackend_node(self):1756 def test_fused_graph_construct_and_dump_with_ascbackend_node(self):
1703- FusedGraph = ascir.FusedGraph('fused_graph')1757+ fused_graph = ascir.FusedGraph("fused_graph")
1704- data0 = ascir.ops.Data('data0', FusedGraph)1758+ data0 = ascir.ops.Data("data0", fused_graph)
1705 data0.attr.ir_attr.index = 01759 data0.attr.ir_attr.index = 0
1706- data1 = ascir.ops.Data('data1', FusedGraph)1760+ data1 = ascir.ops.Data("data1", fused_graph)
1707 data1.attr.ir_attr.index = 01761 data1.attr.ir_attr.index = 0
1708- ascgraph_node0 = ascir.ops.AscBackend("ascgraph_node0", self.construct_add_ascgraph_without_data("ascgraph0"),1762+ ascgraph_node0 = ascir.ops.AscBackend(
1709- FusedGraph)1763+ "ascgraph_node0",
1710- ascgraph_node1 = ascir.ops.AscBackend("ascgraph_node1", self.construct_add_ascgraph("ascgraph1"), FusedGraph)1764+ self.construct_add_ascgraph_without_data("ascgraph0"),
1765+ fused_graph,
1766+ )
1767+ ascgraph_node1 = ascir.ops.AscBackend(
1768+ "ascgraph_node1", self.construct_add_ascgraph("ascgraph1"), fused_graph
1769+ )
1711 ascgraph_node1.x = [data0.y, data1.y]1770 ascgraph_node1.x = [data0.y, data1.y]
1712- ascgraph_node2 = ascir.ops.AscBackend("ascgraph_node2", self.construct_add_ascgraph("ascgraph2"), FusedGraph)1771+ ascgraph_node2 = ascir.ops.AscBackend(
1772+ "ascgraph_node2", self.construct_add_ascgraph("ascgraph2"), fused_graph
1773+ )
1713 ascgraph_node2.x = [ascgraph_node0.y[0], ascgraph_node1.y[1]]1774 ascgraph_node2.x = [ascgraph_node0.y[0], ascgraph_node1.y[1]]
1714- output = ascir.ops.Output('output', FusedGraph)1775+ output = ascir.ops.Output("output", fused_graph)
1715 output.x = ascgraph_node2.y[1]1776 output.x = ascgraph_node2.y[1]
1716- ascir.utils.dump(FusedGraph)1777+ ascir.utils.dump(fused_graph)
1717 1778 
1718 def test_fused_graph_inductor(self):1779 def test_fused_graph_inductor(self):
1719- FusedGraph = ascir.FusedGraph('fused_graph')1780+ fused_graph = ascir.FusedGraph("fused_graph")
1720 1781 
1721 options = AutofuserOptions()1782 options = AutofuserOptions()
1722- scheduler = Schedule(options)1783+ Schedule(options)
1723 fuser = Autofuser(options)1784 fuser = Autofuser(options)
1724 try:1785 try:
1725- schedule_results = fuser.schedule(FusedGraph)1786+ fuser.schedule(fused_graph)
1726- tiling_def, host_tiling, op_kernel = fuser.autofuse_backend(FusedGraph)1787+ tiling_def, host_tiling, op_kernel = fuser.autofuse_backend(fused_graph)
1727- except RuntimeError as e:1788+ except RuntimeError:
1728 pass1789 pass
1729 1790 
1730 def test_fused_graph_construct_and_dump_with_ascgraph_node(self):1791 def test_fused_graph_construct_and_dump_with_ascgraph_node(self):
1731- FusedGraph = ascir.FusedGraph('fused_graph')1792+ fused_graph = ascir.FusedGraph("fused_graph")
1732- data0 = ascir.ops.Data('data0', FusedGraph)1793+ data0 = ascir.ops.Data("data0", fused_graph)
1733 data0.attr.ir_attr.index = 01794 data0.attr.ir_attr.index = 0
1734- data1 = ascir.ops.Data('data1', FusedGraph)1795+ data1 = ascir.ops.Data("data1", fused_graph)
1735 data1.attr.ir_attr.index = 01796 data1.attr.ir_attr.index = 0
1736- ascgraph_node0 = ascir.ops.AscGraph("ascgraph_node0", self.construct_add_ascgraph("ascgraph0"), FusedGraph)1797+ ascgraph_node0 = ascir.ops.AscGraph(
1798+ "ascgraph_node0", self.construct_add_ascgraph("ascgraph0"), fused_graph
1799+ )
1737 ascgraph_node0.x = [data0.y, data1.y]1800 ascgraph_node0.x = [data0.y, data1.y]
1738- ascgraph_node1 = ascir.ops.AscGraph("ascgraph_node1", self.construct_add_ascgraph("ascgraph1"), FusedGraph)1801+ ascgraph_node1 = ascir.ops.AscGraph(
1802+ "ascgraph_node1", self.construct_add_ascgraph("ascgraph1"), fused_graph
1803+ )
1739 ascgraph_node1.x = [data0.y, data1.y]1804 ascgraph_node1.x = [data0.y, data1.y]
1740- ascgraph_node2 = ascir.ops.AscGraph("ascgraph_node2", self.construct_add_ascgraph("ascgraph2"), FusedGraph)1805+ ascgraph_node2 = ascir.ops.AscGraph(
1806+ "ascgraph_node2", self.construct_add_ascgraph("ascgraph2"), fused_graph
1807+ )
1741 ascgraph_node2.x = [ascgraph_node0.y[0], ascgraph_node1.y[0]]1808 ascgraph_node2.x = [ascgraph_node0.y[0], ascgraph_node1.y[0]]
1742- output = ascir.ops.Output('output', FusedGraph)1809+ output = ascir.ops.Output("output", fused_graph)
1743 output.x = ascgraph_node2.y[0]1810 output.x = ascgraph_node2.y[0]
1744 try:1811 try:
1745 ascgraph_node2.x = [ascgraph_node0.y[0].dtype, ascgraph_node1.y[0]]1812 ascgraph_node2.x = [ascgraph_node0.y[0].dtype, ascgraph_node1.y[0]]
1746 except TypeError as e:1813 except TypeError as e:
1747 assert e.args[0] == "Input Type is invalid."1814 assert e.args[0] == "Input Type is invalid."
1748 1815 
1749- ascir.utils.dump(FusedGraph)1816+ ascir.utils.dump(fused_graph)
1750 try:1817 try:
1751 ascir.utils.dump(data0)1818 ascir.utils.dump(data0)
1752 except TypeError as e:1819 except TypeError as e:
1753- assert e.args[0] == "Argument must be a HintGraph or FusedGraph object, got Data"1820+ assert (
1821+ e.args[0]
1822+ == "Argument must be a HintGraph or FusedGraph object, got Data"
1823+ )
1754 1824 
1755 1825 
1756-class TestFusedGraphByApi():1826+class TestFusedGraphByApi:
1757 @staticmethod1827 @staticmethod
1758 def construct_add_ascgraph(name: str) -> ascir.HintGraph:1828 def construct_add_ascgraph(name: str) -> ascir.HintGraph:
1759- NpuKernel0Graph = ascir.HintGraph(name)1829+ npu_kernel0_graph = ascir.HintGraph(name)
1760- s0 = NpuKernel0Graph.create_size("s0")1830+ s0 = npu_kernel0_graph.create_size("s0")
1761- s1 = NpuKernel0Graph.create_size("s1")1831+ s1 = npu_kernel0_graph.create_size("s1")
1762- z0 = NpuKernel0Graph.create_axis("z0", s0)1832+ z0 = npu_kernel0_graph.create_axis("z0", s0)
1763- z1 = NpuKernel0Graph.create_axis("z1", s1)1833+ z1 = npu_kernel0_graph.create_axis("z1", s1)
1764- sub_data0 = ascir_api.Data(NpuKernel0Graph, dtype=ascir.dtypes.float16)1834+ sub_data0 = ascir_api.Data(npu_kernel0_graph, dtype=ascir.dtypes.float16)
1765- load0 = ascir_api.Load(NpuKernel0Graph, sub_data0, offset=0, axis=[z0, z1])1835+ load0 = ascir_api.Load(npu_kernel0_graph, sub_data0, offset=0, axis=[z0, z1])
1766 assert load0.axis == [z0.id, z1.id]1836 assert load0.axis == [z0.id, z1.id]
1767 assert load0.size == [s0, s1]1837 assert load0.size == [s0, s1]
1768 assert load0.strides == [s1, 1]1838 assert load0.strides == [s1, 1]
1769- sub_data1 = ascir_api.Data(NpuKernel0Graph, dtype=ascir.dtypes.float16)1839+ sub_data1 = ascir_api.Data(npu_kernel0_graph, dtype=ascir.dtypes.float16)
1770- load1 = ascir_api.Load(NpuKernel0Graph, sub_data1, offset=0, axis=[z0, z1])1840+ load1 = ascir_api.Load(npu_kernel0_graph, sub_data1, offset=0, axis=[z0, z1])
1771- add0 = ascir_api.Add(NpuKernel0Graph, load0, load1, axis=[z0, z1])1841+ add0 = ascir_api.Add(npu_kernel0_graph, load0, load1, axis=[z0, z1])
1772 assert add0.axis == [z0.id, z1.id]1842 assert add0.axis == [z0.id, z1.id]
1773 assert add0.size == [s0, s1]1843 assert add0.size == [s0, s1]
1774 assert add0.strides == [s1, 1]1844 assert add0.strides == [s1, 1]
1775- store0 = ascir_api.Store(NpuKernel0Graph, add0, offset=0, axis=[z0, z1])1845+ store0 = ascir_api.Store(npu_kernel0_graph, add0, offset=0, axis=[z0, z1])
1776- store1 = ascir_api.Store(NpuKernel0Graph, add0, offset=10, axis=[z0, z1])1846+ store1 = ascir_api.Store(npu_kernel0_graph, add0, offset=10, axis=[z0, z1])
1777 # store0, strore1 写到同一个output上，偏移不同1847 # store0, strore1 写到同一个output上，偏移不同
1778- buf0 = ascir_api.Output(NpuKernel0Graph, [store0, store1], dtype=ascir.dtypes.float16)1848+ ascir_api.Output(
1779- buf1 = ascir_api.Output(NpuKernel0Graph, store1) # infer1849+ npu_kernel0_graph, [store0, store1], dtype=ascir.dtypes.float16
1850+ )
1851+ buf1 = ascir_api.Output(npu_kernel0_graph, store1) # infer
1780 assert buf1.dtype == ascir.dtypes.float161852 assert buf1.dtype == ascir.dtypes.float16
1781- print(ascir.utils.debug_str(NpuKernel0Graph))1853+ print(ascir.utils.debug_str(npu_kernel0_graph))
1782- return NpuKernel0Graph1854+ return npu_kernel0_graph
1783 1855 
1784 def test_fused_graph_construct_and_dump_with_ascbackend_node(self):1856 def test_fused_graph_construct_and_dump_with_ascbackend_node(self):
1785- FusedGraph = ascir.FusedGraph('fused_graph')1857+ fused_graph = ascir.FusedGraph("fused_graph")
1786- data0 = ascir.ops.Data('data0', FusedGraph)1858+ data0 = ascir.ops.Data("data0", fused_graph)
1787 data0.attr.ir_attr.index = 01859 data0.attr.ir_attr.index = 0
1788- data1 = ascir.ops.Data('data1', FusedGraph)1860+ data1 = ascir.ops.Data("data1", fused_graph)
1789 data1.attr.ir_attr.index = 01861 data1.attr.ir_attr.index = 0
1790- ascgraph_node0 = ascir.ops.AscGraph("ascgraph_node0", self.construct_add_ascgraph("ascgraph0"),1862+ ascgraph_node0 = ascir.ops.AscGraph(
1791- FusedGraph)1863+ "ascgraph_node0", self.construct_add_ascgraph("ascgraph0"), fused_graph
1864+ )
1792 ascgraph_node0.x = [data0.y, data1.y]1865 ascgraph_node0.x = [data0.y, data1.y]
1793- ascgraph_node1 = ascir.ops.AscGraph("ascgraph_node1", self.construct_add_ascgraph("ascgraph1"), FusedGraph)1866+ ascgraph_node1 = ascir.ops.AscGraph(
1867+ "ascgraph_node1", self.construct_add_ascgraph("ascgraph1"), fused_graph
1868+ )
1794 ascgraph_node1.x = [data0.y, data1.y]1869 ascgraph_node1.x = [data0.y, data1.y]
1795- ascgraph_node2 = ascir.ops.AscGraph("ascgraph_node2", self.construct_add_ascgraph("ascgraph2"), FusedGraph)1870+ ascgraph_node2 = ascir.ops.AscGraph(
1871+ "ascgraph_node2", self.construct_add_ascgraph("ascgraph2"), fused_graph
1872+ )
1796 ascgraph_node2.x = [ascgraph_node0.y[0], ascgraph_node1.y[1]]1873 ascgraph_node2.x = [ascgraph_node0.y[0], ascgraph_node1.y[1]]
1797- output = ascir.ops.Output('output', FusedGraph)1874+ output = ascir.ops.Output("output", fused_graph)
1798 output.x = ascgraph_node2.y[1]1875 output.x = ascgraph_node2.y[1]
1799- ascir.utils.dump(FusedGraph)1876+ ascir.utils.dump(fused_graph)
1877+ 
1800 1878 
1801# 测试包含transpose的sched, codegen的流程, 执行不抛异常, 返回结果非空1879# 测试包含transpose的sched, codegen的流程, 执行不抛异常, 返回结果非空
1802-class TestAutofuseLoadTransposeStore():1880+class TestAutofuseLoadTransposeStore:
1803 @staticmethod1881 @staticmethod
1804 def construct_invalid_graph():1882 def construct_invalid_graph():
1805 graph = ascir.HintGraph("LoadTransposeStore")1883 graph = ascir.HintGraph("LoadTransposeStore")
@@ -1835,7 +1913,7 @@ class TestAutofuseLoadTransposeStore():
1835 buf_z2 = graph.create_axis("buf_z2", s2)1913 buf_z2 = graph.create_axis("buf_z2", s2)
1836 1914 
1837 arg3_1 = ascir.ops.Data("arg3_1", graph)1915 arg3_1 = ascir.ops.Data("arg3_1", graph)
1838- arg3_1.attr.ir_attr.index= 01916+ arg3_1.attr.ir_attr.index = 0
1839 arg3_1.attr.sched.axis = [z0, z1, z2]1917 arg3_1.attr.sched.axis = [z0, z1, z2]
1840 arg3_1.y.dtype = ascir.dtypes.float161918 arg3_1.y.dtype = ascir.dtypes.float16
1841 arg3_1.y.axis = [z0, z1, z2]1919 arg3_1.y.axis = [z0, z1, z2]
@@ -1846,7 +1924,7 @@ class TestAutofuseLoadTransposeStore():
1846 try:1924 try:
1847 load.attr.ir_attr.offset = "3"1925 load.attr.ir_attr.offset = "3"
1848 except Exception as e:1926 except Exception as e:
1849- assert e.args[0] == 'Only support type of SizeExpr or long'1927+ assert e.args[0] == "Only support type of SizeExpr or long"
1850 offset_of_0 = ascir.SizeExpr(0)1928 offset_of_0 = ascir.SizeExpr(0)
1851 load.attr.ir_attr.offset = offset_of_01929 load.attr.ir_attr.offset = offset_of_0
1852 assert load.attr.ir_attr.offset.expression == "0"1930 assert load.attr.ir_attr.offset.expression == "0"
@@ -1869,7 +1947,7 @@ class TestAutofuseLoadTransposeStore():
1869 try:1947 try:
1870 store.attr.ir_attr.offset = "4"1948 store.attr.ir_attr.offset = "4"
1871 except Exception as e:1949 except Exception as e:
1872- assert e.args[0] == 'Only support type of SizeExpr or long'1950+ assert e.args[0] == "Only support type of SizeExpr or long"
1873 store.attr.ir_attr.offset = offset_of_0 + 11951 store.attr.ir_attr.offset = offset_of_0 + 1
1874 assert store.attr.ir_attr.offset.expression == "1"1952 assert store.attr.ir_attr.offset.expression == "1"
1875 store.x = transpose0_op1953 store.x = transpose0_op
@@ -1888,7 +1966,7 @@ class TestAutofuseLoadTransposeStore():
1888 buf1.y.axis = [z1, z0, z2]1966 buf1.y.axis = [z1, z0, z2]
1889 buf1.y.size = [s1, s0, s2]1967 buf1.y.size = [s1, s0, s2]
1890 buf1.y.strides = [s0 * s2, s2, ascir.SizeExpr(1)]1968 buf1.y.strides = [s0 * s2, s2, ascir.SizeExpr(1)]
1891- graph.set_axis_map({z0:[buf_z0], z1:[buf_z1], z2:[buf_z2]})1969+ graph.set_axis_map({z0: [buf_z0], z1: [buf_z1], z2: [buf_z2]})
1892 return graph1970 return graph
1893 1971 
1894 def test_construct_graph(self):1972 def test_construct_graph(self):
@@ -1897,33 +1975,35 @@ class TestAutofuseLoadTransposeStore():
1897 assert debug_str1975 assert debug_str
1898 1976 
1899 def test_autofuse_backend(self):1977 def test_autofuse_backend(self):
1900- options = AutofuserOptions()1978+ options = AutofuserOptions()
1901- fuser = Autofuser(options)1979+ fuser = Autofuser(options)
1902- try:1980+ try:
1903 hint_graph = self.construct_graph()1981 hint_graph = self.construct_graph()
1904 sched_result = fuser.schedule(hint_graph)1982 sched_result = fuser.schedule(hint_graph)
1905 tiling_def, host_tiling, op_kernel = fuser.codegen(sched_result)1983 tiling_def, host_tiling, op_kernel = fuser.codegen(sched_result)
1906 assert len(tiling_def) > 01984 assert len(tiling_def) > 0
1907 assert len(host_tiling) > 01985 assert len(host_tiling) > 0
1908 assert len(op_kernel) > 01986 assert len(op_kernel) > 0
1909- except RuntimeError as e:1987+ except RuntimeError:
1910 pass1988 pass
1989+ 
1911 import os1990 import os
1991+ 
1912 def test_autofuse_backend_faild_dump_graph(self):1992 def test_autofuse_backend_faild_dump_graph(self):
1913 options = AutofuserOptions()1993 options = AutofuserOptions()
1914 fuser = Autofuser(options)1994 fuser = Autofuser(options)
1915 hint_graph = self.construct_invalid_graph()1995 hint_graph = self.construct_invalid_graph()
1916- with pytest.raises(RuntimeError, match=r'^Optimize fail$'):1996+ with pytest.raises(RuntimeError, match=r"^Optimize fail$"):
1917- sched_result = fuser.schedule(hint_graph)1997+ fuser.schedule(hint_graph)
1918- target_dir = './'1998+ target_dir = "./"
1919 for item in os.listdir(target_dir):1999 for item in os.listdir(target_dir):
1920 item_path = os.path.join(target_dir, item)2000 item_path = os.path.join(target_dir, item)
1921- if os.path.isdir(item_path) and item.startswith('ascgen_dump_pid'):2001+ if os.path.isdir(item_path) and item.startswith("ascgen_dump_pid"):
1922 print(f"delete dump dir :{item_path}")2002 print(f"delete dump dir :{item_path}")
1923 shutil.rmtree(item_path)2003 shutil.rmtree(item_path)
1924 2004 
1925 2005 
1926-class TestSizeExprMaxMin():2006+class TestSizeExprMaxMin:
1927 """Test Max and Min functions for SizeExpr"""2007 """Test Max and Min functions for SizeExpr"""
1928 2008 
1929 @staticmethod2009 @staticmethod
@@ -2001,7 +2081,7 @@ class TestSizeExprMaxMin():
2001 s1 = ascir.SizeExpr(20)2081 s1 = ascir.SizeExpr(20)
2002 s2 = ascir.SizeExpr(30)2082 s2 = ascir.SizeExpr(30)
2003 expr1 = s0 + s1 # 302083 expr1 = s0 + s1 # 30
2004- expr2 = s2 # 302084+ expr2 = s2 # 30
2005 max_expr = Max(expr1, expr2)2085 max_expr = Max(expr1, expr2)
2006 assert max_expr == 302086 assert max_expr == 30
2007 2087 
@@ -2012,7 +2092,7 @@ class TestSizeExprMaxMin():
2012 s1 = ascir.SizeExpr(20)2092 s1 = ascir.SizeExpr(20)
2013 s2 = ascir.SizeExpr(5)2093 s2 = ascir.SizeExpr(5)
2014 expr1 = s0 + s1 # 302094 expr1 = s0 + s1 # 30
2015- expr2 = s2 # 52095+ expr2 = s2 # 5
2016 min_expr = Min(expr1, expr2)2096 min_expr = Min(expr1, expr2)
2017 assert min_expr == 52097 assert min_expr == 5
2018 2098 
@@ -2045,7 +2125,7 @@ class TestSizeExprMaxMin():
2045 assert min_val == 202125 assert min_val == 20
2046 2126 
2047 2127 
2048-class TestSizeExprMod():2128+class TestSizeExprMod:
2049 """Test Mod function for SizeExpr"""2129 """Test Mod function for SizeExpr"""
2050 2130 
2051 @staticmethod2131 @staticmethod
@@ -2103,7 +2183,7 @@ class TestSizeExprMod():
2103 assert mod_abc == 22183 assert mod_abc == 2
2104 2184 
2105 2185 
2106-class TestSizeExprArithmetic():2186+class TestSizeExprArithmetic:
2107 """Test SizeExpr arithmetic operators in various scenarios"""2187 """Test SizeExpr arithmetic operators in various scenarios"""
2108 2188 
2109 @staticmethod2189 @staticmethod
@@ -2113,8 +2193,8 @@ class TestSizeExprArithmetic():
2113 base_size = graph.create_size("base")2193 base_size = graph.create_size("base")
2114 2194 
2115 # Block size = base^22195 # Block size = base^2
2116- block_size = base_size ** 22196+ block_size = base_size**2
2117- z0 = graph.create_axis("z0", block_size)2197+ graph.create_axis("z0", block_size)
2118 2198 
2119 debug_str = ascir.utils.debug_str(graph)2199 debug_str = ascir.utils.debug_str(graph)
2120 assert debug_str2200 assert debug_str
@@ -2130,7 +2210,7 @@ class TestSizeExprArithmetic():
2130 2210 
2131 # Total elements = batch * seq * hidden2211 # Total elements = batch * seq * hidden
2132 total_elements = batch_size * seq_len * hidden_size2212 total_elements = batch_size * seq_len * hidden_size
2133- z0 = graph.create_axis("z0", total_elements)2213+ graph.create_axis("z0", total_elements)
2134 2214 
2135 debug_str = ascir.utils.debug_str(graph)2215 debug_str = ascir.utils.debug_str(graph)
2136 assert debug_str2216 assert debug_str
@@ -2145,7 +2225,7 @@ class TestSizeExprArithmetic():
2145 2225 
2146 # Split size = total / num_splits2226 # Split size = total / num_splits
2147 split_size = total_size / num_splits2227 split_size = total_size / num_splits
2148- z0 = graph.create_axis("z0", split_size)2228+ graph.create_axis("z0", split_size)
2149 2229 
2150 debug_str = ascir.utils.debug_str(graph)2230 debug_str = ascir.utils.debug_str(graph)
2151 assert debug_str2231 assert debug_str
@@ -2161,12 +2241,11 @@ class TestSizeExprArithmetic():
2161 2241 
2162 # Total size = size1 + size2 + constant2242 # Total size = size1 + size2 + constant
2163 total_size = size1 + size2 + constant2243 total_size = size1 + size2 + constant
2164- z0 = graph.create_axis("z0", total_size)2244+ graph.create_axis("z0", total_size)
2165 2245 
2166 debug_str = ascir.utils.debug_str(graph)2246 debug_str = ascir.utils.debug_str(graph)
2167 assert debug_str2247 assert debug_str
2168 2248 
2169- 
2170 @staticmethod2249 @staticmethod
2171 def test_subtraction_basic():2250 def test_subtraction_basic():
2172 """Test basic subtraction between SizeExpr"""2251 """Test basic subtraction between SizeExpr"""
@@ -2213,7 +2292,7 @@ class TestSizeExprArithmetic():
2213 assert result == 72292 assert result == 7
2214 2293 
2215 2294 
2216-class TestSizeExprEdgeCases():2295+class TestSizeExprEdgeCases:
2217 """Test SizeExpr edge cases and boundary conditions"""2296 """Test SizeExpr edge cases and boundary conditions"""
2218 2297 
2219 @staticmethod2298 @staticmethod
@@ -2277,7 +2356,7 @@ class TestSizeExprEdgeCases():
2277 assert min_val == 02356 assert min_val == 0
2278 2357 
2279 2358 
2280-class TestSizeExprInRealScenarios():2359+class TestSizeExprInRealScenarios:
2281 """Test SizeExpr in real-world scenarios"""2360 """Test SizeExpr in real-world scenarios"""
2282 2361 
2283 @staticmethod2362 @staticmethod
@@ -2292,7 +2371,7 @@ class TestSizeExprInRealScenarios():
2292 # Clamp tile size: at least min_tile, at most max_tile2371 # Clamp tile size: at least min_tile, at most max_tile
2293 clamped_size = Min(Max(requested_size, min_tile), max_tile)2372 clamped_size = Min(Max(requested_size, min_tile), max_tile)
2294 2373 
2295- z0 = graph.create_axis("z0", clamped_size)2374+ graph.create_axis("z0", clamped_size)
2296 2375 
2297 debug_str = ascir.utils.debug_str(graph)2376 debug_str = ascir.utils.debug_str(graph)
2298 assert debug_str2377 assert debug_str
@@ -2311,13 +2390,13 @@ class TestSizeExprInRealScenarios():
2311 # offset = n*C*H*W + c*H*W + h*W + w2390 # offset = n*C*H*W + c*H*W + h*W + w
2312 offset = n * c * h * w + c * h * w + h * w + w2391 offset = n * c * h * w + c * h * w + h * w + w
2313 2392 
2314- z0 = graph.create_axis("z0", offset)2393+ graph.create_axis("z0", offset)
2315 2394 
2316 debug_str = ascir.utils.debug_str(graph)2395 debug_str = ascir.utils.debug_str(graph)
2317 assert debug_str2396 assert debug_str
2318 2397 
2319 2398 
2320-class TestSizeExprOperatorCombination():2399+class TestSizeExprOperatorCombination:
2321 """Test SizeExpr operator combinations"""2400 """Test SizeExpr operator combinations"""
2322 2401 
2323 @staticmethod2402 @staticmethod
@@ -2333,8 +2412,8 @@ class TestSizeExprOperatorCombination():
2333 remainder = total_size % block_size2412 remainder = total_size % block_size
2334 2413 
2335 # Create axes for both2414 # Create axes for both
2336- z_blocks = graph.create_axis("z_blocks", num_blocks)2415+ graph.create_axis("z_blocks", num_blocks)
2337- z_remain = graph.create_axis("z_remain", remainder)2416+ graph.create_axis("z_remain", remainder)
2338 2417 
2339 debug_str = ascir.utils.debug_str(graph)2418 debug_str = ascir.utils.debug_str(graph)
2340 assert debug_str2419 assert debug_str
@@ -2352,7 +2431,7 @@ class TestSizeExprOperatorCombination():
2352 # Clamp value between min and max2431 # Clamp value between min and max
2353 clamped = Min(Max(value, min_val), max_val)2432 clamped = Min(Max(value, min_val), max_val)
2354 2433 
2355- z0 = graph.create_axis("z0", clamped)2434+ graph.create_axis("z0", clamped)
2356 2435 
2357 debug_str = ascir.utils.debug_str(graph)2436 debug_str = ascir.utils.debug_str(graph)
2358 assert debug_str2437 assert debug_str
@@ -2369,7 +2448,7 @@ class TestSizeExprOperatorCombination():
2369 # For concat, output size in non-concat dim is max of inputs2448 # For concat, output size in non-concat dim is max of inputs
2370 max_dim_size = Max(Max(size1, size2), size3)2449 max_dim_size = Max(Max(size1, size2), size3)
2371 2450 
2372- z0 = graph.create_axis("z0", max_dim_size)2451+ graph.create_axis("z0", max_dim_size)
2373 2452 
2374 debug_str = ascir.utils.debug_str(graph)2453 debug_str = ascir.utils.debug_str(graph)
2375 assert debug_str2454 assert debug_str
@@ -2386,13 +2465,13 @@ class TestSizeExprOperatorCombination():
2386 # Min helps check the smaller dimension2465 # Min helps check the smaller dimension
2387 min_size = Min(size1, size2)2466 min_size = Min(size1, size2)
2388 2467 
2389- z0 = graph.create_axis("z0", min_size)2468+ graph.create_axis("z0", min_size)
2390 2469 
2391 debug_str = ascir.utils.debug_str(graph)2470 debug_str = ascir.utils.debug_str(graph)
2392 assert debug_str2471 assert debug_str
2393 2472 
2394 2473 
2395-class TestSizeExprErrorScenarios():2474+class TestSizeExprErrorScenarios:
2396 """Test SizeExpr error handling and edge cases"""2475 """Test SizeExpr error handling and edge cases"""
2397 2476 
2398 @staticmethod2477 @staticmethod
@@ -2402,9 +2481,9 @@ class TestSizeExprErrorScenarios():
2402 graph.create_size("size1")2481 graph.create_size("size1")
2403 # Max with no arguments should raise TypeError2482 # Max with no arguments should raise TypeError
2404 try:2483 try:
2405- result = Max()2484+ Max()
2406 assert False, "Expected TypeError for Max() with no arguments"2485 assert False, "Expected TypeError for Max() with no arguments"
2407- except (TypeError, AttributeError) as e:2486+ except (TypeError, AttributeError):
2408 # Expected - invalid number of arguments2487 # Expected - invalid number of arguments
2409 pass2488 pass
2410 2489 
@@ -2415,7 +2494,7 @@ class TestSizeExprErrorScenarios():
2415 size1 = graph.create_size("size1")2494 size1 = graph.create_size("size1")
2416 # Max with single argument should raise TypeError2495 # Max with single argument should raise TypeError
2417 try:2496 try:
2418- result = Max(size1)2497+ Max(size1)
2419 assert False, "Expected TypeError for Max() with single argument"2498 assert False, "Expected TypeError for Max() with single argument"
2420 except TypeError:2499 except TypeError:
2421 # Expected - invalid number of arguments2500 # Expected - invalid number of arguments
@@ -2430,7 +2509,7 @@ class TestSizeExprErrorScenarios():
2430 size3 = graph.create_size("size3")2509 size3 = graph.create_size("size3")
2431 # Max with three arguments should raise TypeError2510 # Max with three arguments should raise TypeError
2432 try:2511 try:
2433- result = Max(size1, size2, size3)2512+ Max(size1, size2, size3)
2434 assert False, "Expected TypeError for Max() with three arguments"2513 assert False, "Expected TypeError for Max() with three arguments"
2435 except TypeError:2514 except TypeError:
2436 # Expected - invalid number of arguments2515 # Expected - invalid number of arguments
@@ -2443,14 +2522,14 @@ class TestSizeExprErrorScenarios():
2443 size1 = graph.create_size("size1")2522 size1 = graph.create_size("size1")
2444 # Max with string argument2523 # Max with string argument
2445 try:2524 try:
2446- result = Max(size1, "invalid")2525+ Max(size1, "invalid")
2447 # If it doesn't raise, at least verify it handles gracefully2526 # If it doesn't raise, at least verify it handles gracefully
2448 except (TypeError, AttributeError, SystemError):2527 except (TypeError, AttributeError, SystemError):
2449 # Expected - invalid type2528 # Expected - invalid type
2450 pass2529 pass
2451 # Max with None argument2530 # Max with None argument
2452 try:2531 try:
2453- result = Max(size1, None)2532+ Max(size1, None)
2454 # If it doesn't raise, at least verify it handles gracefully2533 # If it doesn't raise, at least verify it handles gracefully
2455 except (TypeError, AttributeError, SystemError):2534 except (TypeError, AttributeError, SystemError):
2456 # Expected - invalid type2535 # Expected - invalid type
@@ -2463,7 +2542,7 @@ class TestSizeExprErrorScenarios():
2463 graph.create_size("size1")2542 graph.create_size("size1")
2464 # Min with no arguments should raise TypeError2543 # Min with no arguments should raise TypeError
2465 try:2544 try:
2466- result = Min()2545+ Min()
2467 assert False, "Expected TypeError for Min() with no arguments"2546 assert False, "Expected TypeError for Min() with no arguments"
2468 except (TypeError, AttributeError):2547 except (TypeError, AttributeError):
2469 # Expected - invalid number of arguments2548 # Expected - invalid number of arguments
@@ -2476,7 +2555,7 @@ class TestSizeExprErrorScenarios():
2476 size1 = graph.create_size("size1")2555 size1 = graph.create_size("size1")
2477 # Min with single argument should raise TypeError2556 # Min with single argument should raise TypeError
2478 try:2557 try:
2479- result = Min(size1)2558+ Min(size1)
2480 assert False, "Expected TypeError for Min() with single argument"2559 assert False, "Expected TypeError for Min() with single argument"
2481 except TypeError:2560 except TypeError:
2482 # Expected - invalid number of arguments2561 # Expected - invalid number of arguments
@@ -2489,14 +2568,14 @@ class TestSizeExprErrorScenarios():
2489 size1 = graph.create_size("size1")2568 size1 = graph.create_size("size1")
2490 # Min with integer argument2569 # Min with integer argument
2491 try:2570 try:
2492- result = Min(size1, 42)2571+ Min(size1, 42)
2493 # If it doesn't raise, at least verify it handles gracefully2572 # If it doesn't raise, at least verify it handles gracefully
2494 except (TypeError, AttributeError, SystemError):2573 except (TypeError, AttributeError, SystemError):
2495 # Expected - invalid type2574 # Expected - invalid type
2496 pass2575 pass
2497 # Min with dict argument2576 # Min with dict argument
2498 try:2577 try:
2499- result = Min(size1, {"key": "value"})2578+ Min(size1, {"key": "value"})
2500 # If it doesn't raise, at least verify it handles gracefully2579 # If it doesn't raise, at least verify it handles gracefully
2501 except (TypeError, AttributeError, SystemError):2580 except (TypeError, AttributeError, SystemError):
2502 # Expected - invalid type2581 # Expected - invalid type
@@ -2509,7 +2588,7 @@ class TestSizeExprErrorScenarios():
2509 graph.create_size("size1")2588 graph.create_size("size1")
2510 # Mod with no arguments should raise TypeError2589 # Mod with no arguments should raise TypeError
2511 try:2590 try:
2512- result = Mod()2591+ Mod()
2513 assert False, "Expected TypeError for Mod() with no arguments"2592 assert False, "Expected TypeError for Mod() with no arguments"
2514 except (TypeError, AttributeError):2593 except (TypeError, AttributeError):
2515 # Expected - invalid number of arguments2594 # Expected - invalid number of arguments
@@ -2522,7 +2601,7 @@ class TestSizeExprErrorScenarios():
2522 size1 = graph.create_size("size1")2601 size1 = graph.create_size("size1")
2523 # Mod with single argument should raise TypeError2602 # Mod with single argument should raise TypeError
2524 try:2603 try:
2525- result = Mod(size1)2604+ Mod(size1)
2526 assert False, "Expected TypeError for Mod() with single argument"2605 assert False, "Expected TypeError for Mod() with single argument"
2527 except TypeError:2606 except TypeError:
2528 # Expected - invalid number of arguments2607 # Expected - invalid number of arguments
@@ -2535,14 +2614,14 @@ class TestSizeExprErrorScenarios():
2535 size1 = graph.create_size("size1")2614 size1 = graph.create_size("size1")
2536 # Mod with list argument2615 # Mod with list argument
2537 try:2616 try:
2538- result = Mod(size1, [1, 2, 3])2617+ Mod(size1, [1, 2, 3])
2539 # If it doesn't raise, at least verify it handles gracefully2618 # If it doesn't raise, at least verify it handles gracefully
2540 except (TypeError, AttributeError, SystemError):2619 except (TypeError, AttributeError, SystemError):
2541 # Expected - invalid type2620 # Expected - invalid type
2542 pass2621 pass
2543 # Mod with tuple argument2622 # Mod with tuple argument
2544 try:2623 try:
2545- result = Mod(size1, (1, 2))2624+ Mod(size1, (1, 2))
2546 # If it doesn't raise, at least verify it handles gracefully2625 # If it doesn't raise, at least verify it handles gracefully
2547 except (TypeError, AttributeError, SystemError):2626 except (TypeError, AttributeError, SystemError):
2548 # Expected - invalid type2627 # Expected - invalid type
@@ -2555,7 +2634,7 @@ class TestSizeExprErrorScenarios():
2555 s1 = ascir.SizeExpr(50)2634 s1 = ascir.SizeExpr(50)
2556 # This tests the operator in reverse - string formatting with SizeExpr2635 # This tests the operator in reverse - string formatting with SizeExpr
2557 try:2636 try:
2558- result = "invalid" % s12637+ "invalid" % s1
2559 # String % with SizeExpr - might work differently2638 # String % with SizeExpr - might work differently
2560 except (TypeError, AttributeError):2639 except (TypeError, AttributeError):
2561 # Expected - string formatting doesn't support SizeExpr2640 # Expected - string formatting doesn't support SizeExpr
@@ -2566,7 +2645,7 @@ class TestSizeExprErrorScenarios():
2566 """Test FloorDiv with negative value edge case"""2645 """Test FloorDiv with negative value edge case"""
2567 s0 = ascir.SizeExpr(100)2646 s0 = ascir.SizeExpr(100)
2568 # Negative divisor - verify it handles without crashing2647 # Negative divisor - verify it handles without crashing
2569- result = s0 // -52648+ s0 // -5
2570 # Just verify it doesn't crash2649 # Just verify it doesn't crash
2571 2650 
2572 @staticmethod2651 @staticmethod
@@ -2574,7 +2653,7 @@ class TestSizeExprErrorScenarios():
2574 """Test Max with None as left argument"""2653 """Test Max with None as left argument"""
2575 # None is not a valid SizeExpr - should raise SystemError2654 # None is not a valid SizeExpr - should raise SystemError
2576 try:2655 try:
2577- result = Max(None, ascir.SizeExpr(10))2656+ Max(None, ascir.SizeExpr(10))
2578 except (SystemError, TypeError):2657 except (SystemError, TypeError):
2579 # Expected - None is not a valid SizeExpr2658 # Expected - None is not a valid SizeExpr
2580 pass2659 pass
@@ -2584,7 +2663,7 @@ class TestSizeExprErrorScenarios():
2584 """Test Max with None as right argument"""2663 """Test Max with None as right argument"""
2585 # None is not a valid SizeExpr - should raise SystemError2664 # None is not a valid SizeExpr - should raise SystemError
2586 try:2665 try:
2587- result = Max(ascir.SizeExpr(10), None)2666+ Max(ascir.SizeExpr(10), None)
2588 except (SystemError, TypeError):2667 except (SystemError, TypeError):
2589 # Expected - None is not a valid SizeExpr2668 # Expected - None is not a valid SizeExpr
2590 pass2669 pass
@@ -2594,7 +2673,7 @@ class TestSizeExprErrorScenarios():
2594 """Test Min with None as both arguments"""2673 """Test Min with None as both arguments"""
2595 # None is not a valid SizeExpr - should raise SystemError2674 # None is not a valid SizeExpr - should raise SystemError
2596 try:2675 try:
2597- result = Min(None, None)2676+ Min(None, None)
2598 except (SystemError, TypeError):2677 except (SystemError, TypeError):
2599 # Expected - None is not a valid SizeExpr2678 # Expected - None is not a valid SizeExpr
2600 pass2679 pass
@@ -2641,3 +2720,4 @@ class TestSizeExprErrorScenarios():
2641 2720 
2642 result = ascir.utils.set_platform("2201", 1, 1024)2721 result = ascir.utils.set_platform("2201", 1, 1024)
2643 assert result is None, "set_platform should return None for valid input"2722 assert result is None, "set_platform should return None for valid input"
2723+ ascir.utils.set_platform("2201", 1, 245760)
Mautofuse/tests/ut/att/testcase/gen_model_info/core/test_gen_model_info.cpp+24-0
@@ -73,6 +73,15 @@ ascir::FusedScheduledResult CreateScheduleResultWithSingleAscGraph(
73 fused_schedule_result.node_idx_to_scheduled_results.emplace_back(schedule_results);73 fused_schedule_result.node_idx_to_scheduled_results.emplace_back(schedule_results);
74 return fused_schedule_result;74 return fused_schedule_result;
75}75}
76+ 
77+bool HasVariableName(const ModelInfo &model_info, const std::string &name) {
78+ for (const auto &item : model_info.variable_name_map) {
79+ if (item.second == name) {
80+ return true;
81+ }
82+ }
83+ return false;
84+}
76} // namespace85} // namespace
77 86 
78namespace af {87namespace af {
@@ -667,6 +676,21 @@ TEST_F(TestGenModelInfo, ModelInfoParser) {
667 EXPECT_EQ(MakeJson(model_info_list, json_info), af::SUCCESS);676 EXPECT_EQ(MakeJson(model_info_list, json_info), af::SUCCESS);
668}677}
669 678 
679+TEST_F(TestGenModelInfo, RefreshCommonUbExprKeepsNonContainerVariables) {
680+ std::vector<af::AscGraph> graphs;
681+ TilingModelInfo model_info_list;
682+ af::AscGraph graph("graph");
683+ ASSERT_EQ(af::ascir::cg::Build2DTransposeAscendGraph(graph, {1, 0}), af::SUCCESS);
684+ graphs.emplace_back(graph);
685+ const auto &tiling_data_name = graph.GetName() + "TilingData";
686+ ASSERT_EQ(GenerateModelInfo(graphs, model_info_list, {{kTilingDataTypeName, tiling_data_name}}), af::SUCCESS);
687+ 
688+ ASSERT_EQ(model_info_list.size(), 1U);
689+ EXPECT_TRUE(HasVariableName(model_info_list[0], "transpose_output_0"));
690+ EXPECT_TRUE(HasVariableName(model_info_list[0], "add_output_0"));
691+ EXPECT_NE(model_info_list[0].container_exprs.find("q0_size"), model_info_list[0].container_exprs.end());
692+}
693+ 
670TEST_F(TestGenModelInfo, ModelInfoParserForTranspose10ApiTiling) {694TEST_F(TestGenModelInfo, ModelInfoParserForTranspose10ApiTiling) {
671 std::string json_info;695 std::string json_info;
672 std::vector<af::AscGraph> graphs;696 std::vector<af::AscGraph> graphs;
Mautofuse/tests/ut/att/testcase/solver_pass_gen/axes_reorder_gen/test_axes_reorder_gen.cpp+17-2
@@ -10,8 +10,10 @@
10 10 
11#include "gtest/gtest.h"11#include "gtest/gtest.h"
12#include "base/base_types.h"12#include "base/base_types.h"
13+#include "common/ub_expr/ub_expr_types.h"
13#define private public14#define private public
14#include "generator/solver_pass_gen/axes_reorder_solver/axes_reorder_solver_gen.h"15#include "generator/solver_pass_gen/axes_reorder_solver/axes_reorder_solver_gen.h"
16+#include "generator/solver_pass_gen/axes_reorder_solver/ub_named_expr_builder.h"
15#include "gen_model_info/api_perf_register/v1/perf_param_v1.h"17#include "gen_model_info/api_perf_register/v1/perf_param_v1.h"
16using namespace att;18using namespace att;
17 19 
@@ -270,8 +272,7 @@ TEST_F(TestAxesReorderSolverGen, GenSolverFuncImplAppliesRuntimeReorderOnceBefor
270 EXPECT_LT(swap_pos, solver_pos);272 EXPECT_LT(swap_pos, solver_pos);
271}273}
272 274 
273-TEST_F(TestAxesReorderSolverGen, TEST_GEN_SOLVER_case2)275+TEST_F(TestAxesReorderSolverGen, TEST_GEN_SOLVER_case2) {
274-{
275 Expr x0 = CreateExpr("x0");276 Expr x0 = CreateExpr("x0");
276 Expr x1 = CreateExpr("block_dim");277 Expr x1 = CreateExpr("block_dim");
277 std::vector<Expr> cut_cons;278 std::vector<Expr> cut_cons;
@@ -906,6 +907,20 @@ TEST_F(TestAxesReorderSolverGen, GenGetUbSizeStaticFunc_UsesNamedExprForSemantic
906 EXPECT_NE(actual.find("ub_size = (32 * Ceiling((Rational(1,32) * tensor_size_0)))"), std::string::npos) << actual;907 EXPECT_NE(actual.find("ub_size = (32 * Ceiling((Rational(1,32) * tensor_size_0)))"), std::string::npos) << actual;
907}908}
908 909 
910+TEST_F(TestAxesReorderSolverGen, BuildNamedUbExprUsesCommonOriginExpr) {
911+ Expr tensor = CreateExpr("tensor_size");
912+ Expr tensor_value = CreateExpr("s0") * CreateExpr("s1");
913+ ascir::UbExprContext context;
914+ context.ub_expr = af::sym::Mul(CreateExpr(32), af::sym::Ceiling(af::sym::Div(tensor, CreateExpr(32))));
915+ context.container_expr[tensor] = tensor_value;
916+ context.container_names[tensor] = "load0";
917+ 
918+ const auto actual = BuildNamedUbExpr(context, " ");
919+ 
920+ EXPECT_NE(actual.first.find("auto tensor_size_0 = (s0 * s1);"), std::string::npos) << actual.first;
921+ EXPECT_NE(actual.second.find("tensor_size_0"), std::string::npos) << actual.second;
922+}
923+ 
909TEST_F(TestAxesReorderSolverGen, GenGetUbSizeStaticFunc_NamesExpandedUbExpr) {924TEST_F(TestAxesReorderSolverGen, GenGetUbSizeStaticFunc_NamesExpandedUbExpr) {
910 AxesReorderSolverGen solver("case_test", "TilingData");925 AxesReorderSolverGen solver("case_test", "TilingData");
911 Expr s2 = CreateExpr("S2");926 Expr s2 = CreateExpr("S2");
Aautofuse/tests/ut/common/test_platform_context.cpp+89-0
@@ -0,0 +1,89 @@
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+ 
13+#include "common/platform_context.h"
14+#include "tests/depends/runtime/src/runtime_stub.h"
15+ 
16+class PlatformContextTest : public testing::Test {
17+ protected:
18+ void TearDown() override {
19+ ge::PlatformContext::GetInstance().Reset();
20+ }
21+};
22+ 
23+TEST_F(PlatformContextTest, TryGetInitializedPlatformInfoReturnsFalseAfterReset) {
24+ ge::PlatformContext::GetInstance().Reset();
25+ 
26+ ge::PlatformInfo info;
27+ EXPECT_FALSE(ge::PlatformContext::GetInstance().TryGetInitializedPlatformInfo(info));
28+ EXPECT_EQ(info.ub_size, 0);
29+}
30+ 
31+TEST_F(PlatformContextTest, TryGetInitializedPlatformInfoReturnsInjectedInfo) {
32+ ge::PlatformContext::GetInstance().Reset();
33+ ge::PlatformInfo injected;
34+ injected.soc_ver = "mock_soc";
35+ injected.aiv_num = 20;
36+ injected.ub_size = 262144;
37+ ge::PlatformContext::GetInstance().SetPlatformInfo(injected);
38+ 
39+ ge::PlatformInfo actual;
40+ EXPECT_TRUE(ge::PlatformContext::GetInstance().TryGetInitializedPlatformInfo(actual));
41+ EXPECT_EQ(actual.soc_ver, "mock_soc");
42+ EXPECT_EQ(actual.aiv_num, 20);
43+ EXPECT_EQ(actual.ub_size, 262144);
44+}
45+ 
46+TEST_F(PlatformContextTest, SetPlatformInfoWithEmptySocVerStoresUbSizeOverride) {
47+ ge::PlatformContext::GetInstance().Reset();
48+ ge::PlatformInfo injected;
49+ injected.ub_size = 262144;
50+ ge::PlatformContext::GetInstance().SetPlatformInfo(injected);
51+ 
52+ ge::PlatformInfo actual;
53+ EXPECT_FALSE(ge::PlatformContext::GetInstance().TryGetInitializedPlatformInfo(actual));
54+ 
55+ int64_t ub_size = 0;
56+ EXPECT_TRUE(ge::PlatformContext::GetInstance().TryGetUbSizeOverride(ub_size));
57+ EXPECT_EQ(ub_size, 262144);
58+}
59+ 
60+TEST_F(PlatformContextTest, SetUbSizeOverrideDoesNotInitializePlatformInfo) {
61+ ge::PlatformContext::GetInstance().Reset();
62+ ge::PlatformContext::GetInstance().SetUbSizeOverride(2048);
63+ 
64+ ge::PlatformInfo actual;
65+ int64_t ub_size = 0;
66+ EXPECT_FALSE(ge::PlatformContext::GetInstance().TryGetInitializedPlatformInfo(actual));
67+ EXPECT_TRUE(ge::PlatformContext::GetInstance().TryGetUbSizeOverride(ub_size));
68+ EXPECT_EQ(ub_size, 2048);
69+}
70+ 
71+TEST_F(PlatformContextTest, ResetClearsUbSizeOverride) {
72+ ge::PlatformContext::GetInstance().Reset();
73+ ge::PlatformContext::GetInstance().SetUbSizeOverride(262144);
74+ ge::PlatformContext::GetInstance().Reset();
75+ 
76+ int64_t ub_size = 0;
77+ EXPECT_FALSE(ge::PlatformContext::GetInstance().TryGetUbSizeOverride(ub_size));
78+ EXPECT_EQ(ub_size, 0);
79+}
80+ 
81+TEST_F(PlatformContextTest, GetPlatformInfoInitializesUbSizeFromRuntime) {
82+ ge::RuntimeStub::Reset();
83+ 
84+ ge::PlatformInfo actual;
85+ EXPECT_EQ(ge::PlatformContext::GetInstance().GetPlatformInfo(actual), af::SUCCESS);
86+ EXPECT_EQ(actual.soc_ver, "2201");
87+ EXPECT_EQ(actual.aiv_num, 48);
88+ EXPECT_EQ(actual.ub_size, 245760);
89+}
Aautofuse/tests/ut/common/test_ub_expr_utils.cpp+341-0
@@ -0,0 +1,341 @@
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+ 
13+#include "ascir_ops.h"
14+#include "common/ub_expr/asc_graph_ub_expr_builder.h"
15+#include "common/ub_expr/ub_expr_utils.h"
16+#include "graph/utils/graph_utils.h"
17+ 
18+namespace {
19+struct UbQueueAttrConfig {
20+ ge::DataType dtype = ge::DT_FLOAT16;
21+ std::vector<af::Expression> repeats;
22+ std::vector<int64_t> vectorized_axis;
23+ std::vector<af::Expression> vectorized_strides;
24+ af::Position position = af::Position::kPositionVecIn;
25+ int64_t que_id = 0;
26+ int64_t buf_num = 1;
27+};
28+ 
29+af::AscNodePtr BuildLoadNode(af::AscGraph &graph, const std::string &name) {
30+ const auto data_name = name + "_data";
31+ af::ascir_op::Data data_op(data_name.c_str(), graph);
32+ af::ascir_op::Load load_op(name.c_str());
33+ graph.AddNode(load_op);
34+ load_op.x = data_op.y;
35+ return graph.FindNode(name.c_str());
36+}
37+ 
38+void SetUbQueueOutputAttr(af::AscTensorAttr &attr, const UbQueueAttrConfig &config) {
39+ attr.dtype = config.dtype;
40+ attr.repeats = config.repeats;
41+ attr.vectorized_axis = config.vectorized_axis;
42+ attr.vectorized_strides = config.vectorized_strides;
43+ attr.mem.alloc_type = af::AllocType::kAllocTypeQueue;
44+ attr.mem.hardware = af::MemHardware::kMemHardwareUB;
45+ attr.mem.position = config.position;
46+ attr.que.id = config.que_id;
47+ attr.que.buf_num = config.buf_num;
48+}
49+ 
50+af::Expression ReplaceContainers(const ascir::UbExprContext &context) {
51+ std::vector<std::pair<af::Expression, af::Expression>> replacements;
52+ for (const auto &container_expr : context.container_expr) {
53+ replacements.emplace_back(container_expr.first, container_expr.second);
54+ }
55+ return context.ub_expr.Replace(replacements).Simplify();
56+}
57+} // namespace
58+ 
59+TEST(UbExprUtilsTest, BuildUbExprReturnsFalseForInvalidExpr) {
60+ ascir::UbExprContext context;
61+ 
62+ const auto result = ascir::UbExprUtils::BuildUbExpr(context);
63+ 
64+ EXPECT_FALSE(result.has_ub_expr);
65+}
66+ 
67+TEST(UbExprUtilsTest, BuildUbExprReturnsOriginExpr) {
68+ ascir::UbExprContext context;
69+ context.ub_expr = af::sym::Add(af::Symbol("s0"), af::Symbol(32));
70+ 
71+ const auto result = ascir::UbExprUtils::BuildUbExpr(context);
72+ 
73+ EXPECT_TRUE(result.has_ub_expr);
74+ EXPECT_EQ(result.ub_expr, context.ub_expr);
75+ EXPECT_EQ(result.origin_expr, context.ub_expr.Str().get());
76+}
77+ 
78+TEST(AscGraphUbExprBuilderTest, BuildReturnsFalseForGraphWithoutUbAllocation) {
79+ af::AscGraph graph("no_ub_alloc");
80+ ascir::UbExprContext context;
81+ 
82+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
83+ 
84+ EXPECT_FALSE(ascir::UbExprUtils::BuildUbExpr(context).has_ub_expr);
85+}
86+ 
87+TEST(AscGraphUbExprBuilderTest, BuildAggregatesQueueBufferAndTmpBuffer) {
88+ af::AscGraph graph("ub_alloc");
89+ auto &axis = graph.CreateAxis("s0", af::Symbol("s0"));
90+ auto &tile_axis = graph.CreateAxis("t0", af::Axis::kAxisTypeTileInner, af::Symbol("t0"), {axis.id}, af::kIdNone);
91+ auto node = BuildLoadNode(graph, "load");
92+ ASSERT_NE(node, nullptr);
93+ 
94+ auto &output_attr = node->outputs[0].attr;
95+ output_attr.dtype = ge::DT_FLOAT16;
96+ output_attr.repeats = {tile_axis.size};
97+ output_attr.vectorized_axis = {tile_axis.id};
98+ output_attr.vectorized_strides = {af::Symbol(1)};
99+ output_attr.mem.alloc_type = af::AllocType::kAllocTypeQueue;
100+ output_attr.mem.hardware = af::MemHardware::kMemHardwareUB;
101+ output_attr.mem.position = af::Position::kPositionVecIn;
102+ output_attr.mem.reuse_id = 0;
103+ output_attr.que.id = 0;
104+ output_attr.que.buf_num = 2;
105+ 
106+ af::TmpBuffer tmp_buffer;
107+ tmp_buffer.id = 1;
108+ tmp_buffer.buf_desc.size = af::Symbol(64);
109+ tmp_buffer.mem.alloc_type = af::AllocType::kAllocTypeBuffer;
110+ tmp_buffer.mem.hardware = af::MemHardware::kMemHardwareUB;
111+ node->attr.tmp_buffers.emplace_back(tmp_buffer);
112+ 
113+ ascir::UbExprContext context;
114+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
115+ const auto result = ascir::UbExprUtils::BuildUbExpr(context);
116+ 
117+ EXPECT_TRUE(result.has_ub_expr);
118+ EXPECT_NE(result.origin_expr.find("q0_size"), std::string::npos);
119+ EXPECT_NE(result.origin_expr.find("b1_size"), std::string::npos);
120+ EXPECT_FALSE(context.container_expr.empty());
121+ EXPECT_FALSE(context.ub_related_vars.empty());
122+ EXPECT_EQ(context.graph_name, "ub_alloc");
123+}
124+ 
125+TEST(AscGraphUbExprBuilderTest, BuildAlignsTensorBytesAfterDtypeSize) {
126+ af::AscGraph graph("fp16_align");
127+ auto &axis = graph.CreateAxis("s0", af::Symbol(33));
128+ auto node = BuildLoadNode(graph, "load");
129+ ASSERT_NE(node, nullptr);
130+ 
131+ auto &output_attr = node->outputs[0].attr;
132+ SetUbQueueOutputAttr(output_attr, {ge::DT_FLOAT16, {axis.size}, {axis.id}, {af::Symbol(1)}});
133+ 
134+ ascir::UbExprContext context;
135+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
136+ ASSERT_EQ(context.container_expr.size(), 1U);
137+ auto min_expr = ReplaceContainers(context);
138+ 
139+ int64_t min_ub_usage = 0;
140+ EXPECT_TRUE(min_expr.GetConstValue(min_ub_usage)) << min_expr.Str().get();
141+ EXPECT_EQ(min_ub_usage, 96);
142+}
143+ 
144+TEST(AscGraphUbExprBuilderTest, BuildUsesRepeatStrideForVectorizedTensorSize) {
145+ af::AscGraph graph("repeat_stride_size");
146+ auto &outer_axis = graph.CreateAxis("z0t", af::Symbol(100));
147+ auto &inner_axis = graph.CreateAxis("z1", af::Symbol(8));
148+ auto node = BuildLoadNode(graph, "load");
149+ ASSERT_NE(node, nullptr);
150+ 
151+ auto &output_attr = node->outputs[0].attr;
152+ SetUbQueueOutputAttr(output_attr, {ge::DT_INT32,
153+ {outer_axis.size, af::Symbol(16)},
154+ {outer_axis.id, inner_axis.id},
155+ {af::Symbol(16), af::Symbol(1)},
156+ af::Position::kPositionVecOut,
157+ 1});
158+ output_attr.axis = {outer_axis.id, inner_axis.id};
159+ 
160+ ascir::UbExprContext context;
161+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
162+ 
163+ const auto q1_iter = context.container_expr.find(af::Symbol("q1_size"));
164+ ASSERT_NE(q1_iter, context.container_expr.end());
165+ int64_t q1_size = 0;
166+ EXPECT_TRUE(q1_iter->second.GetConstValue(q1_size)) << q1_iter->second.Str().get();
167+ EXPECT_EQ(q1_size, 6400);
168+}
169+ 
170+TEST(AscGraphUbExprBuilderTest, BuildKeepsQueueContainerExprAsSingleBufferSlot) {
171+ af::AscGraph graph("queue_buf_num");
172+ auto &axis = graph.CreateAxis("s0", af::Symbol(16));
173+ auto node = BuildLoadNode(graph, "load");
174+ ASSERT_NE(node, nullptr);
175+ 
176+ auto &output_attr = node->outputs[0].attr;
177+ SetUbQueueOutputAttr(output_attr,
178+ {ge::DT_FLOAT16, {axis.size}, {axis.id}, {af::Symbol(1)}, af::Position::kPositionVecIn, 0, 2});
179+ 
180+ ascir::UbExprContext context;
181+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
182+ 
183+ const auto q0_iter = context.container_expr.find(af::Symbol("q0_size"));
184+ ASSERT_NE(q0_iter, context.container_expr.end());
185+ int64_t q0_size = 0;
186+ EXPECT_TRUE(q0_iter->second.GetConstValue(q0_size)) << q0_iter->second.Str().get();
187+ EXPECT_EQ(q0_size, 32);
188+ 
189+ int64_t ub_usage = 0;
190+ auto min_expr = ReplaceContainers(context);
191+ EXPECT_TRUE(min_expr.GetConstValue(ub_usage)) << min_expr.Str().get();
192+ EXPECT_EQ(ub_usage, 64);
193+}
194+ 
195+TEST(AscGraphUbExprBuilderTest, BuildKeepsUserDefinedContainerNames) {
196+ af::AscGraph graph("named_container");
197+ auto &axis = graph.CreateAxis("s0", af::Symbol(16));
198+ auto queue_node = BuildLoadNode(graph, "queue_load");
199+ auto buffer_node = BuildLoadNode(graph, "buffer_load");
200+ ASSERT_NE(queue_node, nullptr);
201+ ASSERT_NE(buffer_node, nullptr);
202+ 
203+ auto &queue_attr = queue_node->outputs[0].attr;
204+ queue_attr.dtype = ge::DT_FLOAT16;
205+ queue_attr.repeats = {axis.size};
206+ queue_attr.vectorized_axis = {axis.id};
207+ queue_attr.vectorized_strides = {af::Symbol(1)};
208+ queue_attr.mem.alloc_type = af::AllocType::kAllocTypeQueue;
209+ queue_attr.mem.hardware = af::MemHardware::kMemHardwareUB;
210+ queue_attr.que.id = 0;
211+ queue_attr.que.name = "custom_queue_size";
212+ queue_attr.que.buf_num = 2;
213+ 
214+ auto &buffer_attr = buffer_node->outputs[0].attr;
215+ buffer_attr.dtype = ge::DT_FLOAT16;
216+ buffer_attr.repeats = {axis.size};
217+ buffer_attr.vectorized_axis = {axis.id};
218+ buffer_attr.vectorized_strides = {af::Symbol(1)};
219+ buffer_attr.mem.alloc_type = af::AllocType::kAllocTypeBuffer;
220+ buffer_attr.mem.hardware = af::MemHardware::kMemHardwareUB;
221+ buffer_attr.buf.id = 1;
222+ buffer_attr.buf.name = "custom_buffer_size";
223+ 
224+ ascir::UbExprContext context;
225+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
226+ 
227+ EXPECT_NE(context.container_expr.find(af::Symbol("custom_queue_size")), context.container_expr.end());
228+ EXPECT_NE(context.container_expr.find(af::Symbol("custom_buffer_size")), context.container_expr.end());
229+ EXPECT_EQ(context.container_expr.find(af::Symbol("q0_size")), context.container_expr.end());
230+ EXPECT_EQ(context.container_expr.find(af::Symbol("b1_size")), context.container_expr.end());
231+}
232+ 
233+TEST(AscGraphUbExprBuilderTest, BuildSumsCoexistBufferTensorsWithSameReuseId) {
234+ af::AscGraph graph("buffer_reuse_group");
235+ auto &axis = graph.CreateAxis("s0", af::Symbol(16));
236+ auto node0 = BuildLoadNode(graph, "load0");
237+ auto node1 = BuildLoadNode(graph, "load1");
238+ ASSERT_NE(node0, nullptr);
239+ ASSERT_NE(node1, nullptr);
240+ 
241+ auto set_buffer_attr = [&axis](af::AscTensorAttr &attr) {
242+ attr.dtype = ge::DT_FLOAT16;
243+ attr.repeats = {axis.size};
244+ attr.vectorized_axis = {axis.id};
245+ attr.vectorized_strides = {af::Symbol(1)};
246+ attr.mem.alloc_type = af::AllocType::kAllocTypeBuffer;
247+ attr.mem.hardware = af::MemHardware::kMemHardwareUB;
248+ attr.mem.position = af::Position::kPositionVecIn;
249+ attr.mem.reuse_id = 3;
250+ attr.buf.id = 2;
251+ };
252+ set_buffer_attr(node0->outputs[0].attr);
253+ set_buffer_attr(node1->outputs[0].attr);
254+ 
255+ ascir::UbExprContext context;
256+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
257+ 
258+ const auto b2_iter = context.container_expr.find(af::Symbol("b2_size"));
259+ ASSERT_NE(b2_iter, context.container_expr.end());
260+ int64_t b2_size = 0;
261+ EXPECT_TRUE(b2_iter->second.GetConstValue(b2_size)) << b2_iter->second.Str().get();
262+ EXPECT_EQ(b2_size, 64);
263+}
264+ 
265+TEST(AscGraphUbExprBuilderTest, BuildIncludesBuiltinTmpBuffer) {
266+ af::AscGraph graph("builtin_tmp_buffer");
267+ auto &axis = graph.CreateAxis("s0", af::Symbol(1));
268+ auto load_node = BuildLoadNode(graph, "load");
269+ ASSERT_NE(load_node, nullptr);
270+ auto &load_attr = load_node->outputs[0].attr;
271+ load_attr.dtype = ge::DT_BOOL;
272+ load_attr.repeats = {axis.size};
273+ load_attr.vectorized_axis = {axis.id};
274+ load_attr.vectorized_strides = {af::Symbol(1)};
275+ 
276+ af::ascir_op::LogicalNot logical_not_op("logical_not");
277+ graph.AddNode(logical_not_op);
278+ auto logical_not_node = graph.FindNode("logical_not");
279+ ASSERT_NE(logical_not_node, nullptr);
280+ ASSERT_EQ(af::GraphUtils::AddEdge(load_node->GetOutDataAnchor(0), logical_not_node->GetInDataAnchor(0)),
281+ ge::GRAPH_SUCCESS);
282+ auto &logical_not_attr = logical_not_node->outputs[0].attr;
283+ logical_not_attr.dtype = ge::DT_BOOL;
284+ logical_not_attr.repeats = {axis.size};
285+ logical_not_attr.vectorized_axis = {axis.id};
286+ logical_not_attr.vectorized_strides = {af::Symbol(1)};
287+ 
288+ ascir::UbExprContext context;
289+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
290+ 
291+ int64_t min_ub_usage = 0;
292+ EXPECT_TRUE(context.ub_expr.Simplify().GetConstValue(min_ub_usage)) << context.ub_expr.Str().get();
293+ EXPECT_EQ(min_ub_usage, 32);
294+}
295+ 
296+TEST(AscGraphUbExprBuilderTest, BuildSumsTmpBuffersWithSameIdInOneNode) {
297+ af::AscGraph graph("same_id_tmp_buffer");
298+ af::Operator op("Compute", "Compute");
299+ auto node = graph.AddNode(op);
300+ ASSERT_NE(node, nullptr);
301+ node->attr.tmp_buffers.emplace_back(af::TmpBuffer{{af::Symbol(4096), -1}, af::MemAttr(), 0});
302+ node->attr.tmp_buffers.emplace_back(af::TmpBuffer{{af::Symbol(8192), -1}, af::MemAttr(), 0});
303+ 
304+ ascir::UbExprContext context;
305+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
306+ 
307+ int64_t min_ub_usage = 0;
308+ auto min_expr = ReplaceContainers(context);
309+ EXPECT_TRUE(min_expr.GetConstValue(min_ub_usage)) << min_expr.Str().get();
310+ EXPECT_EQ(min_ub_usage, 12288);
311+}
312+ 
313+TEST(AscGraphUbExprBuilderTest, BuildIgnoresInvalidTmpBufferSize) {
314+ af::AscGraph graph("invalid_tmp_buffer_size");
315+ af::Operator op("Compute", "Compute");
316+ auto node = graph.AddNode(op);
317+ ASSERT_NE(node, nullptr);
318+ af::TmpBuffer tmp_buffer;
319+ tmp_buffer.id = 99;
320+ node->attr.tmp_buffers.emplace_back(tmp_buffer);
321+ 
322+ ascir::UbExprContext context;
323+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
324+ 
325+ EXPECT_EQ(context.container_expr.find(af::Symbol("b99_size")), context.container_expr.end());
326+ EXPECT_EQ(ascir::UbExprUtils::BuildUbExpr(context).origin_expr.find("Max(,"), std::string::npos);
327+}
328+ 
329+TEST(AscGraphUbExprBuilderTest, BuildIncludesReservedUbForGather) {
330+ af::AscGraph graph("reserved_ub");
331+ af::Operator op("Gather", "Gather");
332+ auto node = graph.AddNode(op);
333+ ASSERT_NE(node, nullptr);
334+ 
335+ ascir::UbExprContext context;
336+ EXPECT_EQ(ascir::AscGraphUbExprBuilder().Build(graph, context), af::SUCCESS);
337+ 
338+ int64_t min_ub_usage = 0;
339+ EXPECT_TRUE(context.ub_expr.Simplify().GetConstValue(min_ub_usage)) << context.ub_expr.Str().get();
340+ EXPECT_EQ(min_ub_usage, 40960);
341+}
Mautofuse/tests/ut/optimize/test_buf_que_allocator.cpp+81-0
@@ -54,6 +54,87 @@ class BufQueAllocatorUT : public ::testing::Test {
54};54};
55} // namespace optimize55} // namespace optimize
56 56 
57+static af::AscGraph MakeStaticLoadStoreGraph(const std::string &name, int64_t size, int64_t data_index = 0) {
58+ af::AscGraph graph(name.c_str());
59+ const af::Expression s0 = graph.CreateSizeVar(size);
60+ auto z0 = graph.CreateAxis("z0", s0);
61+ 
62+ af::ascir_op::Data data(("data" + std::to_string(data_index)).c_str(), graph);
63+ data.ir_attr.SetIndex(data_index);
64+ data.y.dtype = ge::DT_UINT8;
65+ 
66+ af::ascir_op::Load load("load");
67+ load.x = data.y;
68+ load.attr.api.compute_type = af::ComputeType::kComputeLoad;
69+ load.attr.api.unit = af::ComputeUnit::kUnitMTE2;
70+ load.y.dtype = ge::DT_UINT8;
71+ *load.y.axis = {z0.id};
72+ *load.y.repeats = {s0};
73+ *load.y.strides = {af::ops::One};
74+ 
75+ af::ascir_op::Store store("store");
76+ store.x = load.y;
77+ store.attr.api.compute_type = af::ComputeType::kComputeStore;
78+ store.attr.api.unit = af::ComputeUnit::kUnitMTE2;
79+ store.y.dtype = ge::DT_UINT8;
80+ *store.y.axis = {z0.id};
81+ *store.y.repeats = {s0};
82+ *store.y.strides = {af::ops::One};
83+ 
84+ af::ascir_op::Output output(("output" + std::to_string(data_index)).c_str());
85+ output.x = store.y;
86+ output.ir_attr.SetIndex(data_index);
87+ return graph;
88+}
89+ 
90+static ascir::FusedScheduledResult MakeFusedScheduledResultWithGraphs(std::vector<af::AscGraph> &&impl_graphs) {
91+ ascir::FusedScheduledResult fused_result;
92+ fused_result.node_idx_to_scheduled_results.resize(1UL);
93+ auto &scheduled_result = fused_result.node_idx_to_scheduled_results[0].emplace_back();
94+ auto &group = scheduled_result.schedule_groups.emplace_back();
95+ group.impl_graphs = std::move(impl_graphs);
96+ return fused_result;
97+}
98+ 
99+TEST_F(BufQueAllocatorUT, PrepareImplGraphMemoryPlanDoesNotPopulateFusedIoNodes) {
100+ auto fused_result = MakeFusedScheduledResultWithGraphs({MakeStaticLoadStoreGraph("prepared", 16)});
101+ 
102+ BufQueAllocator allocator;
103+ ASSERT_EQ(allocator.PrepareImplGraphMemoryPlan(fused_result), af::SUCCESS);
104+ 
105+ EXPECT_TRUE(fused_result.input_nodes.empty());
106+ EXPECT_TRUE(fused_result.output_nodes.empty());
107+ EXPECT_TRUE(fused_result.workspace_nodes.empty());
108+ auto &impl_graph = fused_result.node_idx_to_scheduled_results[0][0].schedule_groups[0].impl_graphs[0];
109+ auto load = impl_graph.FindNode("load");
110+ ASSERT_NE(load, nullptr);
111+ EXPECT_EQ(load->outputs[0].attr.mem.hardware, af::MemHardware::kMemHardwareUB);
112+}
113+ 
114+TEST_F(BufQueAllocatorUT, CollectFusedIoNodesOnlyUsesRetainedImplGraphs) {
115+ auto fused_result = MakeFusedScheduledResultWithGraphs(
116+ {MakeStaticLoadStoreGraph("dropped", 16, 0), MakeStaticLoadStoreGraph("kept", 16, 1)});
117+ 
118+ BufQueAllocator allocator;
119+ ASSERT_EQ(allocator.PrepareImplGraphMemoryPlan(fused_result), af::SUCCESS);
120+ auto &impl_graphs = fused_result.node_idx_to_scheduled_results[0][0].schedule_groups[0].impl_graphs;
121+ impl_graphs.erase(impl_graphs.begin());
122+ ASSERT_EQ(allocator.CollectFusedIoNodes(fused_result), af::SUCCESS);
123+ 
124+ ASSERT_EQ(fused_result.input_nodes.size(), 1UL);
125+ ASSERT_EQ(fused_result.output_nodes.size(), 1UL);
126+ EXPECT_EQ(fused_result.input_nodes[0]->GetName(), "data1");
127+ EXPECT_EQ(fused_result.output_nodes[0]->GetName(), "output1");
128+}
129+ 
130+TEST_F(BufQueAllocatorUT, AllocBufQueKeepsOldOneShotBehavior) {
131+ auto fused_result = MakeFusedScheduledResultWithGraphs({MakeStaticLoadStoreGraph("one_shot", 16)});
132+ 
133+ ASSERT_EQ(BufQueAllocator().AllocBufQue(fused_result), af::SUCCESS);
134+ ASSERT_EQ(fused_result.input_nodes.size(), 1UL);
135+ ASSERT_EQ(fused_result.output_nodes.size(), 1UL);
136+}
137+ 
57TEST_F(BufQueAllocatorUT, test_reuse_id_vecacc) {138TEST_F(BufQueAllocatorUT, test_reuse_id_vecacc) {
58 af::AscGraph graph("test_reuse_id_vecacc");139 af::AscGraph graph("test_reuse_id_vecacc");
59 const af::Expression s0 = graph.CreateSizeVar("s0");140 const af::Expression s0 = graph.CreateSizeVar("s0");