已合并
fix: 修复动态Shape下Reduce Tile 运行时轴重排的问题 #1537
fix: 修复动态Shape下Reduce Tile 运行时轴重排的问题 #1537
已合并
zhang_shengjie创建于 7月28日
共 16 个文件变更+753-379
@@ -231,12 +231,11 @@ struct CacheLineConfig {
231};231};
232 232 
233struct RuntimeReorderRule {233struct RuntimeReorderRule {
234- Expr preferred_axis;
235- Expr fallback_axis;
236 Expr condition_axis;234 Expr condition_axis;
237 Expr compare_axis;235 Expr compare_axis;
238 uint32_t condition_threshold{0U};236 uint32_t condition_threshold{0U};
239 uint32_t compare_threshold{0U};237 uint32_t compare_threshold{0U};
238+ std::vector<Expr> preferred_order;
240};239};
241 240 
242class TilingScheduleConfigTable {241class TilingScheduleConfigTable {
@@ -309,8 +308,8 @@ struct ModelInfo {
309 std::vector<Expr> sizes{}; // 图原始Sizes信息308 std::vector<Expr> sizes{}; // 图原始Sizes信息
310 vector<CacheLineConfig> cache_line_config; // ub->gm/gm->ub节点的cache配置信息309 vector<CacheLineConfig> cache_line_config; // ub->gm/gm->ub节点的cache配置信息
311 const TilingScheduleConfigTable *tiling_schedule_config_table{nullptr};310 const TilingScheduleConfigTable *tiling_schedule_config_table{nullptr};
312- TilingScheduleConfig tiling_schedule_config; // Model 级别的 Tiling 调度配置311+ TilingScheduleConfig tiling_schedule_config; // Model 级别的 Tiling 调度配置
313- bool is_enable_equal_order_tiling{false}; // 使能等order tiling算法312+ bool is_enable_equal_order_tiling{false}; // 使能等order tiling算法
314 std::vector<RuntimeReorderRule> runtime_reorder_rules; // 运行时按shape调整单模板内轴优先级313 std::vector<RuntimeReorderRule> runtime_reorder_rules; // 运行时按shape调整单模板内轴优先级
315};314};
316 315 
@@ -263,8 +263,6 @@ bool ArgListReorder::TryBuildReduceTileRuntimeReorderRule(const NodeInfo &node,
263 if (!GetReduceAxisDataTypeSize(node, axis.get(), reduce_axis_ori_axes_set, data_type_size)) {263 if (!GetReduceAxisDataTypeSize(node, axis.get(), reduce_axis_ori_axes_set, data_type_size)) {
264 continue;264 continue;
265 }265 }
266- rule.preferred_axis = axis->repeat;
267- rule.fallback_axis = axis->repeat;
268 rule.compare_axis = GetOriginalAxisRepeat(axis.get());266 rule.compare_axis = GetOriginalAxisRepeat(axis.get());
269 rule.compare_threshold = GetVectorLenSize() / data_type_size;267 rule.compare_threshold = GetVectorLenSize() / data_type_size;
270 return true;268 return true;
@@ -292,7 +290,6 @@ bool ArgListReorder::HasSmallTailLargeReduceTile(const NodeInfo &node,
292 if (tensor->data_type_size == 0U) {290 if (tensor->data_type_size == 0U) {
293 continue;291 continue;
294 }292 }
295- rule.preferred_axis = tensor->repeat[i];
296 rule.condition_axis = GetOriginalAxisRepeat(dim);293 rule.condition_axis = GetOriginalAxisRepeat(dim);
297 rule.condition_threshold = CeilDiv(GetCacheLineSize(), tensor->data_type_size);294 rule.condition_threshold = CeilDiv(GetCacheLineSize(), tensor->data_type_size);
298 uint64_t tail_bytes = 0UL;295 uint64_t tail_bytes = 0UL;
@@ -561,6 +558,62 @@ std::vector<AttAxisPtr> ArgListReorder::GetNewArgList(const std::vector<size_t>
561 return new_arg_list;558 return new_arg_list;
562}559}
563 560 
561+std::vector<Expr> ArgListReorder::GetTileSplitOrder(const std::vector<size_t> &topo_order,
562+ const std::vector<AttAxisPtr> &arg_list) const {
563+ std::vector<Expr> order;
564+ for (const size_t order_id : topo_order) {
565+ if ((order_id < kOrderIdStart) || (order_id > arg_list.size() + kOrderIdStart - 1U)) {
566+ continue;
567+ }
568+ const auto &axis = arg_list[order_id - kOrderIdStart];
569+ if (AttUtils::IsTileSplitAxis(axis)) {
570+ order.emplace_back(axis->size->symbol_expr);
571+ }
572+ }
573+ return order;
574+}
575+ 
576+bool ArgListReorder::SetRuntimePreferredOrder(const std::vector<AttAxisPtr> &canonical_arg_list,
577+ const std::vector<AttAxisPtr> &source_arg_list,
578+ const std::vector<size_t> &preferred_topo_order,
579+ RuntimeReorderRule &rule) const {
580+ std::vector<Expr> canonical_order;
581+ std::map<size_t, std::vector<Expr>> equal_order_groups;
582+ for (const auto &axis : canonical_arg_list) {
583+ if (!AttUtils::IsTileSplitAxis(axis)) {
584+ continue;
585+ }
586+ canonical_order.emplace_back(axis->size->symbol_expr);
587+ equal_order_groups[axis->order].emplace_back(axis->size->symbol_expr);
588+ }
589+ auto preferred_order = GetTileSplitOrder(preferred_topo_order, source_arg_list);
590+ for (const auto &group : equal_order_groups) {
591+ if (group.second.size() < 2U) {
592+ continue;
593+ }
594+ std::vector<size_t> positions;
595+ for (size_t i = 0U; i < preferred_order.size(); ++i) {
596+ if (std::find(group.second.begin(), group.second.end(), preferred_order[i]) != group.second.end()) {
597+ positions.emplace_back(i);
598+ }
599+ }
600+ if (positions.size() != group.second.size()) {
601+ return false;
602+ }
603+ for (size_t i = 0U; i < positions.size(); ++i) {
604+ preferred_order[positions[i]] = group.second[i];
605+ }
606+ }
607+ std::set<Expr, ExprCmp> canonical_set(canonical_order.begin(), canonical_order.end());
608+ std::set<Expr, ExprCmp> preferred_set(preferred_order.begin(), preferred_order.end());
609+ if ((canonical_order.size() != preferred_order.size()) || (canonical_set.size() != canonical_order.size()) ||
610+ (canonical_set != preferred_set)) {
611+ return false;
612+ }
613+ rule.preferred_order = std::move(preferred_order);
614+ return true;
615+}
616+ 
564// 排序的入口函数617// 排序的入口函数
565af::Status ArgListReorder::SortArgList(vector<AttAxisPtr> &arg_list, vector<AttAxisPtr> &tiling_R_arg_list,618af::Status ArgListReorder::SortArgList(vector<AttAxisPtr> &arg_list, vector<AttAxisPtr> &tiling_R_arg_list,
566 std::vector<RuntimeReorderRule> *runtime_reorder_rules) {619 std::vector<RuntimeReorderRule> *runtime_reorder_rules) {
@@ -579,7 +632,15 @@ af::Status ArgListReorder::SortArgList(vector<AttAxisPtr> &arg_list, vector<AttA
579 GE_ASSERT_SUCCESS(BuildArgListPriorityGraph(arg_list, prefer_reduce_tile_), "build arg list graph failed");632 GE_ASSERT_SUCCESS(BuildArgListPriorityGraph(arg_list, prefer_reduce_tile_), "build arg list graph failed");
580 std::vector<AttAxisPtr> new_arg_list = GetNewArgList(graph_->TopologicalSort(), arg_list);633 std::vector<AttAxisPtr> new_arg_list = GetNewArgList(graph_->TopologicalSort(), arg_list);
581 if ((runtime_reorder_rules != nullptr) && has_dynamic_reduce_tile_reorder_) {634 if ((runtime_reorder_rules != nullptr) && has_dynamic_reduce_tile_reorder_) {
582- runtime_reorder_rules->emplace_back(dynamic_reduce_tile_reorder_rule_);635+ graph_ = af::MakeShared<ArgPriorityGraph>(arg_list.size());
636+ GE_ASSERT_NOTNULL(graph_, "Create preferred graph failed");
637+ RuntimeReorderRule runtime_rule = dynamic_reduce_tile_reorder_rule_;
638+ if ((BuildArgListPriorityGraph(arg_list, true) == af::SUCCESS) &&
639+ SetRuntimePreferredOrder(new_arg_list, arg_list, graph_->TopologicalSort(), runtime_rule)) {
640+ runtime_reorder_rules->emplace_back(std::move(runtime_rule));
641+ } else {
642+ GELOGW("[ATT][ReduceTileReorder] Invalid preferred order, keep canonical order.");
643+ }
583 }644 }
584 645 
585 if (tiling_R_) {646 if (tiling_R_) {
@@ -233,6 +233,11 @@ class ArgListReorder {
233 af::Status ApplyPriorityRules(bool tiling_R, const AxisCategories &categories);233 af::Status ApplyPriorityRules(bool tiling_R, const AxisCategories &categories);
234 std::vector<AttAxisPtr> GetNewArgList(const std::vector<size_t> &topo_order,234 std::vector<AttAxisPtr> GetNewArgList(const std::vector<size_t> &topo_order,
235 const std::vector<AttAxisPtr> &arg_list) const;235 const std::vector<AttAxisPtr> &arg_list) const;
236+ std::vector<Expr> GetTileSplitOrder(const std::vector<size_t> &topo_order,
237+ const std::vector<AttAxisPtr> &arg_list) const;
238+ bool SetRuntimePreferredOrder(const std::vector<AttAxisPtr> &canonical_arg_list,
239+ const std::vector<AttAxisPtr> &source_arg_list,
240+ const std::vector<size_t> &preferred_topo_order, RuntimeReorderRule &rule) const;
236 void MakeSureLoadStoreInnerestSameOrder(const std::vector<AttAxisPtr> &arg_list) const;241 void MakeSureLoadStoreInnerestSameOrder(const std::vector<AttAxisPtr> &arg_list) const;
237 bool HandleProperty(const SubAxis *dim, att::ArgListReorder::AxisProperty property, bool is_reduce,242 bool HandleProperty(const SubAxis *dim, att::ArgListReorder::AxisProperty property, bool is_reduce,
238 bool is_broadcast);243 bool is_broadcast);
@@ -721,6 +721,46 @@ ExprUintMap ArgsManager::GetAxesPriority() const {
721 return axes_pirority;721 return axes_pirority;
722}722}
723 723 
724+ExprUintMap ArgsManager::GetAxesOrder() const {
725+ ExprUintMap original_orders;
726+ std::set<Expr, ExprCmp> ambiguous_originals;
727+ uint32_t next_order = 1U;
728+ for (const auto &arg_axis : model_info_.arg_list) {
729+ GE_ASSERT_NOTNULL(arg_axis);
730+ GE_ASSERT_NOTNULL(arg_axis->size);
731+ const uint32_t order = static_cast<uint32_t>(arg_axis->order);
732+ next_order = std::max(next_order, order + 1U);
733+ const auto result = original_orders.emplace(arg_axis->size->symbol_expr, order);
734+ if (!result.second) {
735+ ambiguous_originals.emplace(arg_axis->size->symbol_expr);
736+ }
737+ }
738+ const ExprExprMap var_relations = GetVarsRelations();
739+ const std::vector<Expr> searchable_vars = GetSearchableVars();
740+ std::map<Expr, size_t, ExprCmp> original_counts;
741+ for (const auto &var : searchable_vars) {
742+ const auto relation_iter = var_relations.find(var);
743+ const Expr &original_var = relation_iter == var_relations.end() ? var : relation_iter->second;
744+ ++original_counts[original_var];
745+ }
746+ ExprUintMap axes_order;
747+ for (const auto &var : searchable_vars) {
748+ Expr original_var = var;
749+ const auto relation_iter = var_relations.find(var);
750+ if (relation_iter != var_relations.end()) {
751+ original_var = relation_iter->second;
752+ }
753+ const auto order_iter = original_orders.find(original_var);
754+ if ((order_iter != original_orders.end()) && (ambiguous_originals.count(original_var) == 0U) &&
755+ (original_counts.at(original_var) == 1U)) {
756+ axes_order[var] = order_iter->second;
757+ } else {
758+ axes_order[var] = next_order++;
759+ }
760+ }
761+ return axes_order;
762+}
763+ 
724void ArgsManager::Reset() {764void ArgsManager::Reset() {
725 vars_infos_.clear();765 vars_infos_.clear();
726 hardware_cons_.clear();766 hardware_cons_.clear();
@@ -193,6 +193,7 @@ class ArgsManager {
193 * @brief 获取轴优先级信息193 * @brief 获取轴优先级信息
194 */194 */
195 ExprUintMap GetAxesPriority() const;195 ExprUintMap GetAxesPriority() const;
196+ ExprUintMap GetAxesOrder() const;
196 /**197 /**
197 * @brief 获取Model Info198 * @brief 获取Model Info
198 */199 */
@@ -52,7 +52,7 @@ std::string GenPgoSolverGenerateAllTilingDataBody() {
52 tilingDataVar = input_.pure_mc_vars[index - input_.local_buffer_vars_size];52 tilingDataVar = input_.pure_mc_vars[index - input_.local_buffer_vars_size];
53 from_local_buffer_vars = false;53 from_local_buffer_vars = false;
54 } else {54 } else {
55- tilingDataVar = input_.local_buffer_vars[index];55+ tilingDataVar = input_.ordered_local_buffer_vars[index];
56 from_local_buffer_vars = true;56 from_local_buffer_vars = true;
57 }57 }
58 auto min_ = tilingDataVar->align;58 auto min_ = tilingDataVar->align;
@@ -91,7 +91,7 @@ std::string GenPgoSolverGenerateAllTilingDataTail() {
91 auto tilingDataVarTmp = input_.pure_mc_vars[tmp - input_.local_buffer_vars_size];91 auto tilingDataVarTmp = input_.pure_mc_vars[tmp - input_.local_buffer_vars_size];
92 tilingDataVarTmp->value = tilingDataVarTmp->align;92 tilingDataVarTmp->value = tilingDataVarTmp->align;
93 } else {93 } else {
94- auto tilingDataVarTmp = input_.local_buffer_vars[tmp];94+ auto tilingDataVarTmp = input_.ordered_local_buffer_vars[tmp];
95 tilingDataVarTmp->value = tilingDataVarTmp->align;95 tilingDataVarTmp->value = tilingDataVarTmp->align;
96 }96 }
97 tmp += 1;97 tmp += 1;
@@ -106,7 +106,11 @@ std::string GenPgoSolverGenerateAllTilingDataTail() {
106 continue;106 continue;
107 }107 }
108 }108 }
109- ans_item[index] = tilingDataVar->value;109+ if (from_local_buffer_vars) {
110+ ans_item[input_.ordered_to_canonical[index]] = tilingDataVar->value;
111+ } else {
112+ ans_item[index] = tilingDataVar->value;
113+ }
110 PgoSolverGenerateAllTilingDataInner(index + 1, ans_item, ans, step_max);114 PgoSolverGenerateAllTilingDataInner(index + 1, ans_item, ans, step_max);
111 }115 }
112}116}
@@ -157,6 +157,8 @@ std::string GenAxesReorderSolverInput() {
157 codes += " TilingVariable **tiling_vars = nullptr;\n";157 codes += " TilingVariable **tiling_vars = nullptr;\n";
158 codes += " TilingVariable **pure_mc_vars = nullptr;\n";158 codes += " TilingVariable **pure_mc_vars = nullptr;\n";
159 codes += " TilingVariable **local_buffer_vars = nullptr;\n";159 codes += " TilingVariable **local_buffer_vars = nullptr;\n";
160+ codes += " TilingVariable **ordered_local_buffer_vars = nullptr;\n";
161+ codes += " uint32_t *ordered_to_canonical = nullptr;\n";
160 codes += " Constraint **all_cons = nullptr;\n";162 codes += " Constraint **all_cons = nullptr;\n";
161 codes += GenReorderSolverInputDebugString();163 codes += GenReorderSolverInputDebugString();
162 codes += "};\n";164 codes += "};\n";
@@ -114,6 +114,9 @@ bool AxesReorderSolver::IdentifyEqualPriorityAxes(const uint32_t axes_num, uint3
114 uint32_t index = 0;114 uint32_t index = 0;
115 constexpr uint32_t kSupportMaxEqualPriorityAxes = 2;115 constexpr uint32_t kSupportMaxEqualPriorityAxes = 2;
116 for (const auto &pair : order_to_axis_indices) {116 for (const auto &pair : order_to_axis_indices) {
117+ if (pair.second.size() < kSupportMaxEqualPriorityAxes) {
118+ continue;
119+ }
117 for (const auto &id : pair.second) {120 for (const auto &id : pair.second) {
118 if (index >= axes_num) {121 if (index >= axes_num) {
119 OP_LOGI(OP_NAME, "Axes num %u is not enough to identify equal priority axes, index=%u", axes_num, index);122 OP_LOGI(OP_NAME, "Axes num %u is not enough to identify equal priority axes, index=%u", axes_num, index);
@@ -41,25 +41,26 @@ std::string GenMcRelatedNaiveTiling() {
41 int64_t upper_bound_satisfied_ub_threshold = var->value;41 int64_t upper_bound_satisfied_ub_threshold = var->value;
42 int64_t lower_bound_satisfied_ub_threshold = var->align;42 int64_t lower_bound_satisfied_ub_threshold = var->align;
43 // 若未找到满足UB利用率的解,则不需要进一步处理43 // 若未找到满足UB利用率的解,则不需要进一步处理
44- if (BinaryFindLowerBoundSatisfiedUBThresholdCond(var, i, var->align, lower_bound_satisfied_ub_threshold)) {44+ if (BinaryFindLowerBoundSatisfiedUBThresholdCond(var, canonical_idx, var->align,
45+ lower_bound_satisfied_ub_threshold)) {
45 OP_LOGD(OP_NAME, "Found lower_bound_satisfied_ub_threshold:%ld, upper:%ld, lower:%ld, i:%u, input: %s",46 OP_LOGD(OP_NAME, "Found lower_bound_satisfied_ub_threshold:%ld, upper:%ld, lower:%ld, i:%u, input: %s",
46- lower_bound_satisfied_ub_threshold, upper_bound_satisfied_ub_threshold, var->align, i,47+ lower_bound_satisfied_ub_threshold, upper_bound_satisfied_ub_threshold, var->align, canonical_idx,
47 input_.DebugString().c_str());48 input_.DebugString().c_str());
48 var->SetValue(upper_bound_satisfied_ub_threshold);49 var->SetValue(upper_bound_satisfied_ub_threshold);
49 // 3.1) 在[var->align,upper_bound_satisfied_ub_threshold]范围内,二分查找满足UB利用率的边界50 // 3.1) 在[var->align,upper_bound_satisfied_ub_threshold]范围内,二分查找满足UB利用率的边界
50- auto satisfied_core_threshold = BinaryFindLowerBoundSatisfiedCoreNum(var, i,51+ auto satisfied_core_threshold = BinaryFindLowerBoundSatisfiedCoreNum(var, canonical_idx,
51 lower_bound_satisfied_ub_threshold);52 lower_bound_satisfied_ub_threshold);
52 auto satisfied_core_threshold_left = satisfied_core_threshold.first;53 auto satisfied_core_threshold_left = satisfied_core_threshold.first;
53 auto satisfied_core_threshold_right = satisfied_core_threshold.second;54 auto satisfied_core_threshold_right = satisfied_core_threshold.second;
54 OP_LOGD(OP_NAME, "Found lower bound satisfied core num:%ld, %ld, var upper:%ld, lower:%ld, i:%u, input: %s",55 OP_LOGD(OP_NAME, "Found lower bound satisfied core num:%ld, %ld, var upper:%ld, lower:%ld, i:%u, input: %s",
55 satisfied_core_threshold_left, satisfied_core_threshold_right, upper_bound_satisfied_ub_threshold,56 satisfied_core_threshold_left, satisfied_core_threshold_right, upper_bound_satisfied_ub_threshold,
56- var->align, i, input_.DebugString().c_str());57+ var->align, canonical_idx, input_.DebugString().c_str());
57 // 3.2)先尝试Tile块更大的值,若有解,则更新var58 // 3.2)先尝试Tile块更大的值,若有解,则更新var
58 int64_t available_core_num_right = 0L;59 int64_t available_core_num_right = 0L;
59 var->SetValue(satisfied_core_threshold_right);60 var->SetValue(satisfied_core_threshold_right);
60 if (InitMulticoreVars() && MulticoreTilingCore(false) && CalRealUsedCoreNum(available_core_num_right)) {61 if (InitMulticoreVars() && MulticoreTilingCore(false) && CalRealUsedCoreNum(available_core_num_right)) {
61 OP_LOGD(OP_NAME, "Found larger tile size:%ld, available_core_num:%ld, i:%u",62 OP_LOGD(OP_NAME, "Found larger tile size:%ld, available_core_num:%ld, i:%u",
62- satisfied_core_threshold_right, available_core_num_right, i);63+ satisfied_core_threshold_right, available_core_num_right, canonical_idx);
63 }64 }
64 int64_t available_core_num_left = 0L;65 int64_t available_core_num_left = 0L;
65 if ((satisfied_core_threshold_left != satisfied_core_threshold_right) && (satisfied_core_threshold_left > 0L)) {66 if ((satisfied_core_threshold_left != satisfied_core_threshold_right) && (satisfied_core_threshold_left > 0L)) {
@@ -67,7 +68,7 @@ std::string GenMcRelatedNaiveTiling() {
67 if (InitMulticoreVars() && MulticoreTilingCore(false) && CalRealUsedCoreNum(available_core_num_left) &&68 if (InitMulticoreVars() && MulticoreTilingCore(false) && CalRealUsedCoreNum(available_core_num_left) &&
68 (available_core_num_left > available_core_num_right)) {69 (available_core_num_left > available_core_num_right)) {
69 OP_LOGD(OP_NAME, "Found smaller tile size:%ld, available_core_num:%ld, i:%u",70 OP_LOGD(OP_NAME, "Found smaller tile size:%ld, available_core_num:%ld, i:%u",
70- satisfied_core_threshold_left, available_core_num_left, i);71+ satisfied_core_threshold_left, available_core_num_left, canonical_idx);
71 } else {72 } else {
72 var->SetValue(satisfied_core_threshold_right);73 var->SetValue(satisfied_core_threshold_right);
73 }74 }
@@ -79,24 +80,25 @@ std::string GenMcRelatedNaiveTiling() {
79 80 
80std::string GenNaiveLocalBufTilingImpl() {81std::string GenNaiveLocalBufTilingImpl() {
81 std::string kNaiveLocalBufTilingImpl = R"(82 std::string kNaiveLocalBufTilingImpl = R"(
82- for (uint32_t i = 0u; i < num_vars; ++i) {83+ for (uint32_t ordered_idx = 0U; ordered_idx < num_vars; ++ordered_idx) {
83- auto &var = vars[i];84+ auto &var = vars[ordered_idx];
85+ const uint32_t canonical_idx = input_.ordered_to_canonical[ordered_idx];
84 auto upper_bound = var->upper_bound(var->upper_bound_vars);86 auto upper_bound = var->upper_bound(var->upper_bound_vars);
85 int64_t boundary = (upper_bound / var->align) * var->align;87 int64_t boundary = (upper_bound / var->align) * var->align;
86 if (boundary < var->align) {88 if (boundary < var->align) {
87 OP_LOGW(OP_NAME, "Invalid aligned upper bound:%ld, raw upper:%ld, align:%ld, i:%u, input: %s.",89 OP_LOGW(OP_NAME, "Invalid aligned upper bound:%ld, raw upper:%ld, align:%ld, i:%u, input: %s.",
88- boundary, upper_bound, var->align, i, input_.DebugString().c_str());90+ boundary, upper_bound, var->align, canonical_idx, input_.DebugString().c_str());
89 return false;91 return false;
90 }92 }
91 var->SetValue(boundary);93 var->SetValue(boundary);
92 int64_t upper_bound_satisfied_ub = -1L;94 int64_t upper_bound_satisfied_ub = -1L;
93 if (!BinaryFindUpperBoundSatisfiedUBLimit(var, var->align, upper_bound_satisfied_ub)) {95 if (!BinaryFindUpperBoundSatisfiedUBLimit(var, var->align, upper_bound_satisfied_ub)) {
94 OP_LOGW(OP_NAME, "BinaryFindUpperBoundSatisfiedUBLimit failed, upper:%ld, lower:%ld, i:%u, input: %s.",96 OP_LOGW(OP_NAME, "BinaryFindUpperBoundSatisfiedUBLimit failed, upper:%ld, lower:%ld, i:%u, input: %s.",
95- upper_bound, var->align, i, input_.DebugString().c_str());97+ upper_bound, var->align, canonical_idx, input_.DebugString().c_str());
96 return false;98 return false;
97 }99 }
98 OP_LOGD(OP_NAME, "Found upper_bound_satisfied_ub:%ld, upper_bound:%ld, lower_bound:%ld, i:%u, input: %s",100 OP_LOGD(OP_NAME, "Found upper_bound_satisfied_ub:%ld, upper_bound:%ld, lower_bound:%ld, i:%u, input: %s",
99- upper_bound_satisfied_ub, boundary, var->align, i, input_.DebugString().c_str());101+ upper_bound_satisfied_ub, boundary, var->align, canonical_idx, input_.DebugString().c_str());
100 var->SetValue(upper_bound_satisfied_ub);102 var->SetValue(upper_bound_satisfied_ub);
101)";103)";
102 std::string kNaiveLocalBufTilingImplPostProcess = R"(104 std::string kNaiveLocalBufTilingImplPostProcess = R"(
@@ -154,11 +156,13 @@ std::string GenNaiveLocalBufTilingWithEqualOrderImpl() {
154 }156 }
155 // 处理剩余的轴157 // 处理剩余的轴
156 int64_t max_core_num = static_cast<int64_t>(input_.corenum_threshold * input_.core_num);158 int64_t max_core_num = static_cast<int64_t>(input_.corenum_threshold * input_.core_num);
157- for (uint32_t i = 0; i < input_.local_buffer_vars_size; ++i) {159+ for (uint32_t ordered_idx = 0U; ordered_idx < input_.local_buffer_vars_size; ++ordered_idx) {
158- if (solved_axes[i]) {160+ const uint32_t canonical_idx = input_.ordered_to_canonical[ordered_idx];
161+ if (solved_axes[canonical_idx]) {
159 continue;162 continue;
160 }163 }
161- if (!ProcessSingleAxisNaive(input_.local_buffer_vars[i], i, max_core_num)) {164+ auto *var = input_.ordered_local_buffer_vars[ordered_idx];
165+ if (!ProcessSingleAxisNaive(var, canonical_idx, max_core_num)) {
162 return false;166 return false;
163 }167 }
164 }168 }
@@ -181,7 +185,7 @@ std::string GenNaiveLocalBufTiling(bool enable_equal_order_tiling) {
181 codes.append(GenInitLocalMCVars());185 codes.append(GenInitLocalMCVars());
182 codes.append(R"(186 codes.append(R"(
183 uint32_t num_vars = input_.local_buffer_vars_size;187 uint32_t num_vars = input_.local_buffer_vars_size;
184- auto *vars = input_.local_buffer_vars;188+ auto *vars = input_.ordered_local_buffer_vars;
185)");189)");
186 codes.append(GenNaiveLocalBufTilingImpl());190 codes.append(GenNaiveLocalBufTilingImpl());
187 }191 }
@@ -191,11 +195,12 @@ std::string GenNaiveLocalBufTiling(bool enable_equal_order_tiling) {
191std::string GenBinaryLocalBufTilingCore() {195std::string GenBinaryLocalBufTilingCore() {
192 return R"(196 return R"(
193bool AxesReorderSolver::BinaryLocalBufTilingCore(const std::vector<bool> &solved_axes) {197bool AxesReorderSolver::BinaryLocalBufTilingCore(const std::vector<bool> &solved_axes) {
194- for (uint32_t i = 0u; i < input_.local_buffer_vars_size; ++i) {198+ for (uint32_t ordered_idx = 0U; ordered_idx < input_.local_buffer_vars_size; ++ordered_idx) {
195- if (solved_axes[i]) {199+ const uint32_t canonical_idx = input_.ordered_to_canonical[ordered_idx];
200+ if (solved_axes[canonical_idx]) {
196 continue;201 continue;
197 }202 }
198- auto &var = input_.local_buffer_vars[i];203+ auto &var = input_.ordered_local_buffer_vars[ordered_idx];
199 auto upper_bound = var->upper_bound(var->upper_bound_vars);204 auto upper_bound = var->upper_bound(var->upper_bound_vars);
200 int64_t boundary = (upper_bound / var->align) * var->align;205 int64_t boundary = (upper_bound / var->align) * var->align;
201 int64_t init_val = var->value;206 int64_t init_val = var->value;
@@ -30,7 +30,7 @@ inline bool AxesReorderSolver::WorkloadBalance() {
30 constexpr double EPS = 1e-6;30 constexpr double EPS = 1e-6;
31 bool related = false;31 bool related = false;
32 uint32_t num_vars = input_.local_buffer_vars_size;32 uint32_t num_vars = input_.local_buffer_vars_size;
33- auto *vars = input_.local_buffer_vars;33+ auto *vars = input_.ordered_local_buffer_vars;
34 uint32_t index = num_vars - 1;34 uint32_t index = num_vars - 1;
35 for (uint32_t i = 0; i < num_vars; i++) {35 for (uint32_t i = 0; i < num_vars; i++) {
36 if (vars[i]->mc_related) {36 if (vars[i]->mc_related) {
@@ -908,6 +908,10 @@ std::string AxesReorderSolverGen::InitiateArgs() {
908 }908 }
909 for (const auto &local_arg : local_buffer_tiling_vars_) {909 for (const auto &local_arg : local_buffer_tiling_vars_) {
910 strs += " TilingVariable " + Str(local_arg) + ";\n";910 strs += " TilingVariable " + Str(local_arg) + ";\n";
911+ const auto order_iter = axes_order_.find(local_arg);
912+ if (order_iter != axes_order_.end()) {
913+ strs += " " + Str(local_arg) + ".order = " + std::to_string(order_iter->second) + "UL;\n";
914+ }
911 }915 }
912 return strs;916 return strs;
913}917}
@@ -1117,6 +1121,20 @@ std::string AxesReorderSolverGen::SetTilingVars(VarsType var_type) {
1117 strs += "};\n";1121 strs += "};\n";
1118 strs += " input." + var_name + "_size = " + std::to_string(vars.size()) + "u;\n";1122 strs += " input." + var_name + "_size = " + std::to_string(vars.size()) + "u;\n";
1119 strs += " input." + var_name + " = " + var_name + ";\n";1123 strs += " input." + var_name + " = " + var_name + ";\n";
1124+ if (var_type == VarsType::LOCALBUFFER) {
1125+ strs += " TilingVariable* ordered_local_buffer_vars[" + std::to_string(vars.size()) + "] = {";
1126+ for (uint32_t i = 0U; i < vars.size(); ++i) {
1127+ strs += "&" + Str(vars[i]) + ", ";
1128+ }
1129+ strs += "};\n";
1130+ strs += " uint32_t ordered_to_canonical[" + std::to_string(vars.size()) + "] = {";
1131+ for (uint32_t i = 0U; i < vars.size(); ++i) {
1132+ strs += std::to_string(i) + "u, ";
1133+ }
1134+ strs += "};\n";
1135+ strs += " input.ordered_local_buffer_vars = ordered_local_buffer_vars;\n";
1136+ strs += " input.ordered_to_canonical = ordered_to_canonical;\n";
1137+ }
1120 }1138 }
1121 return strs;1139 return strs;
1122}1140}
@@ -1130,6 +1148,43 @@ int32_t AxesReorderSolverGen::GetLocalBufferVarIndex(const Expr &expr) const {
1130 return -1;1148 return -1;
1131}1149}
1132 1150 
1151+int32_t AxesReorderSolverGen::ResolveLocalBufferVarIndex(const Expr &expr) const {
1152+ int32_t resolved_idx = -1;
1153+ for (size_t i = 0U; i < local_buffer_tiling_vars_.size(); ++i) {
1154+ const Expr &search_var = local_buffer_tiling_vars_[i];
1155+ bool matched = search_var == expr;
1156+ const auto relation_iter = vars_relations_.find(search_var);
1157+ matched = matched || ((relation_iter != vars_relations_.end()) && (relation_iter->second == expr));
1158+ if (!matched) {
1159+ continue;
1160+ }
1161+ if (resolved_idx >= 0) {
1162+ return -1;
1163+ }
1164+ resolved_idx = static_cast<int32_t>(i);
1165+ }
1166+ return resolved_idx;
1167+}
1168+ 
1169+std::vector<int32_t> AxesReorderSolverGen::GetRuntimePreferredIndices(const RuntimeReorderRule &rule) const {
1170+ std::vector<int32_t> indices;
1171+ std::set<int32_t> unique_indices;
1172+ for (const Expr &expr : rule.preferred_order) {
1173+ const int32_t index = ResolveLocalBufferVarIndex(expr);
1174+ if (index < 0) {
1175+ continue;
1176+ }
1177+ if (!unique_indices.emplace(index).second) {
1178+ return {};
1179+ }
1180+ indices.emplace_back(index);
1181+ }
1182+ if (indices.size() != local_buffer_tiling_vars_.size()) {
1183+ return {};
1184+ }
1185+ return indices;
1186+}
1187+ 
1133std::string AxesReorderSolverGen::GenRuntimeExprValue(const Expr &expr) const {1188std::string AxesReorderSolverGen::GenRuntimeExprValue(const Expr &expr) const {
1134 if (expr.IsConstExpr()) {1189 if (expr.IsConstExpr()) {
1135 return Str(expr);1190 return Str(expr);
@@ -1167,31 +1222,19 @@ std::string AxesReorderSolverGen::GenRuntimeCompoundExprValue(const Expr &expr)
1167}1222}
1168 1223 
1169std::string AxesReorderSolverGen::GenRuntimeReorderRule(const RuntimeReorderRule &rule) {1224std::string AxesReorderSolverGen::GenRuntimeReorderRule(const RuntimeReorderRule &rule) {
1170- const int32_t preferred_idx = GetLocalBufferVarIndex(rule.preferred_axis);1225+ const std::vector<int32_t> preferred_indices = GetRuntimePreferredIndices(rule);
1171- const int32_t fallback_idx = GetLocalBufferVarIndex(rule.fallback_axis);1226+ if (preferred_indices.empty()) {
1172- if ((preferred_idx < 0) || (fallback_idx < 0) || (preferred_idx == fallback_idx)) {
1173 return "";1227 return "";
1174 }1228 }
1175 std::string code;1229 std::string code;
1176 code += " if ((" + GenRuntimeExprValue(rule.condition_axis) + " < " + std::to_string(rule.condition_threshold) +1230 code += " if ((" + GenRuntimeExprValue(rule.condition_axis) + " < " + std::to_string(rule.condition_threshold) +
1177 ") && (" + GenRuntimeExprValue(rule.compare_axis) + " > " + std::to_string(rule.compare_threshold) + ")) {\n";1231 ") && (" + GenRuntimeExprValue(rule.compare_axis) + " > " + std::to_string(rule.compare_threshold) + ")) {\n";
1178- code += " OP_LOGI(OP_NAME, \"[ATT][ReduceTileReorder] Runtime reduce tile reorder chooses preferred axis ";1232+ for (size_t i = 0U; i < preferred_indices.size(); ++i) {
1179- code += Str(rule.preferred_axis);1233+ code += " input.ordered_local_buffer_vars[" + std::to_string(i) + "] = input.local_buffer_vars[" +
1180- code += " before fallback axis ";1234+ std::to_string(preferred_indices[i]) + "];\n";
1181- code += Str(rule.fallback_axis);1235+ code += " input.ordered_to_canonical[" + std::to_string(i) + "] = " + std::to_string(preferred_indices[i]) +
1182- code += ", prefer splitting reduce axis.\");\n";1236+ "u;\n";
1183- if (preferred_idx > fallback_idx) {
1184- code += " auto *runtime_preferred_var = input.local_buffer_vars[" + std::to_string(preferred_idx) + "];\n";
1185- code += " input.local_buffer_vars[" + std::to_string(preferred_idx) + "] = input.local_buffer_vars[" +
1186- std::to_string(fallback_idx) + "];\n";
1187- code += " input.local_buffer_vars[" + std::to_string(fallback_idx) + "] = runtime_preferred_var;\n";
1188 }1237 }
1189- code += " } else {\n";
1190- code += " OP_LOGI(OP_NAME, \"[ATT][ReduceTileReorder] Runtime reduce tile reorder keeps fallback axis ";
1191- code += Str(rule.fallback_axis);
1192- code += " before preferred axis ";
1193- code += Str(rule.preferred_axis);
1194- code += ", prefer splitting tail axis.\");\n";
1195 code += " }\n";1238 code += " }\n";
1196 return code;1239 return code;
1197}1240}
@@ -1613,6 +1656,7 @@ std::string AxesReorderSolverGen::GenPGOSolverFuncImpl() {
1613 codes += InitiateArgs();1656 codes += InitiateArgs();
1614 codes += GenInputInfo(all_cons, local_buffer_cons, mc_mixed_cons);1657 codes += GenInputInfo(all_cons, local_buffer_cons, mc_mixed_cons);
1615 codes += GenInput(TradeOffConfig(), all_cons);1658 codes += GenInput(TradeOffConfig(), all_cons);
1659+ codes += GenRuntimeReorderRules();
1616 // topn search_cfg override: directly apply ub_threshold/corenum_threshold to solver input1660 // topn search_cfg override: directly apply ub_threshold/corenum_threshold to solver input
1617 codes += " if (search_cfg != nullptr) {\n";1661 codes += " if (search_cfg != nullptr) {\n";
1618 codes += " if (search_cfg->ub_threshold_enabled) {\n";1662 codes += " if (search_cfg->ub_threshold_enabled) {\n";
@@ -24,227 +24,254 @@
24#include "generator/solver_pass_gen/pgo_config_setters_mixin.h"24#include "generator/solver_pass_gen/pgo_config_setters_mixin.h"
25 25 
26namespace att {26namespace att {
27- enum class ConsType {27+enum class ConsType {
28- BUFFER = 0,28+ BUFFER = 0,
29- CUT = 1,29+ CUT = 1,
30- MCMIXED = 2,30+ MCMIXED = 2,
31- ALL = 3,31+ ALL = 3,
32+};
33+ 
34+enum class InputType {
35+ INPUT = 0,
36+ TILING = 1,
37+};
38+ 
39+enum class VarsType {
40+ PUREMC = 0,
41+ LOCALBUFFER = 1,
42+};
43+ 
44+class AxesReorderSolverGen : public SolverGen,
45+ public InputOutputSettersMixin<AxesReorderSolverGen>,
46+ public PgoConfigSettersMixin<AxesReorderSolverGen> {
47+ public:
48+ explicit AxesReorderSolverGen(const std::string &tiling_case_id, const std::string &type_name)
49+ : SolverGen(tiling_case_id, type_name) {}
50+ ~AxesReorderSolverGen() override = default;
51+ std::string GenSolverClassImpl() override;
52+ std::string GenSolverFuncImpl() override;
53+ std::string GenPGOSolverFilter();
54+ std::string GenSolverFuncInvoke() override;
55+ std::string GenPGOSolverClassImpl();
56+ std::string GenPGOSolverFuncImpl();
57+ 
58+ void SetInputArgs(const std::vector<Expr> &input_args) {
59+ input_args_ = input_args;
60+ }
61+ void SetConstArgs(const ExprUintMap &const_vars) {
62+ std::vector<Expr> const_args;
63+ const_vars_map_ = const_vars;
64+ for (const auto &pair : const_vars_map_) {
65+ const_args.push_back(pair.first);
66+ }
67+ const_args_ = const_args;
68+ }
69+ void SetBufferUseAlg(const std::map<HardwareDef, Expr> &hardware_use_map) {
70+ hardware_use_map_ = hardware_use_map;
71+ }
72+ void SetArgAlignMap(const ExprExprMap &arg_align_map) {
73+ arg_align_map_ = arg_align_map;
74+ }
75+ void SetArgPromptAlignMap(const ExprUintMap &arg_prompt_align_map) {
76+ arg_prompt_align_map_ = arg_prompt_align_map;
77+ }
78+ void SetArgDataTypeSizeMap(const ExprUintMap &data_type_size_map) {
79+ data_type_size_map_ = data_type_size_map;
80+ }
81+ void SetInputAlign(const ExprExprMap &input_align);
82+ void SetTotalCutCons(const std::vector<Expr> &total_cut_cons) {
83+ total_cut_cons_ = total_cut_cons;
84+ }
85+ void SetFromAxesMap(const std::map<Expr, std::vector<Expr>, ExprCmp> &from_axes_map) {
86+ from_axes_map_ = from_axes_map;
87+ }
88+ void SetVarPriority(const ExprUintMap &priority) {
89+ priority_map_ = priority;
90+ }
91+ void SetAxesOrder(const ExprUintMap &axes_order) {
92+ axes_order_ = axes_order;
93+ }
94+ void SetVarsRelations(const ExprExprMap &vars_relations) {
95+ vars_relations_ = vars_relations;
96+ }
97+ void SetContainerExpr(const ExprExprMap &container_expr) {
98+ container_expr_ = container_expr;
99+ }
100+ void SetContainerNames(const std::map<Expr, std::string, ExprCmp> &container_names) {
101+ container_names_ = container_names;
102+ }
103+ void SetReplaceVars(const std::vector<std::pair<Expr, Expr>> &replace_vars) {
104+ for (const auto &var : replace_vars) {
105+ replace_vars_.emplace_back(var);
106+ }
107+ }
108+ void SetTernaryOps(const std::map<Expr, TernaryOp, ExprCmp> &ternary_ops) {
109+ ternary_ops_ = ternary_ops;
110+ }
111+ 
112+ void SetExeTimeMap(const std::map<Expr, std::vector<Expr>, ExprCmp> &exe_time_map) {
113+ for (const auto &pair : exe_time_map) {
114+ exe_time_map_[pair.first] = pair.second;
115+ }
116+ }
117+ void Arrange();
118+ void SetObjFunc(const Expr &head_cost, const std::map<PipeType, Expr> pipe_2_obj_map) {
119+ head_cost_ = head_cost;
120+ pipe_2_obj_map_ = pipe_2_obj_map;
121+ }
122+ void SetIsConcatOuterMap(const ExprUintMap &is_concat_outer_map) {
123+ is_concat_outer_map_ = is_concat_outer_map;
124+ }
125+ void SetConcatInnerDims(const std::vector<Expr> &concat_inner_dims) {
126+ concat_inner_dims_ = concat_inner_dims;
127+ }
128+ void SetUBThreshold(const double &ub_threshold) {
129+ ub_threshold_ = ub_threshold;
130+ }
131+ void SetCoreNumThreshold(const double &corenum_threshold) {
132+ corenum_threshold_ = corenum_threshold;
32 };133 };
33- 134+ void SetReservedUbSize(const Expr &reserved_ub_size) {
34- enum class InputType {135+ reserved_ub_size_ = reserved_ub_size;
35- INPUT = 0,
36- TILING = 1,
37 };136 };
137+ void SetHighPerfTiling(const bool enable_high_perf) {
138+ enable_high_perf_ = enable_high_perf;
139+ }
140+ void SetEnableEqualOrder(const bool enable_equal_order) {
141+ enable_equal_order_ = enable_equal_order;
142+ }
143+ void SetSearchArgs(const std::vector<Expr> &search_args) {
144+ search_args_ = search_args;
145+ }
146+ void SetArrangeCode(const std::string &arrange_code) {
147+ arrange_code_ = arrange_code;
148+ }
149+ void SetTilingScheduleConfigTable(const TilingScheduleConfigTable *tiling_schedule_config_table) {
150+ tiling_schedule_config_table_ = tiling_schedule_config_table;
151+ }
152+ void SetTilingScheduleConfig(const TilingScheduleConfig &tiling_schedule_config) {
153+ tiling_schedule_config_ = tiling_schedule_config;
154+ }
155+ void SetRuntimeReorderRules(const std::vector<RuntimeReorderRule> &runtime_reorder_rules) {
156+ runtime_reorder_rules_ = runtime_reorder_rules;
157+ }
158+ void SetCacheLineConfig(const vector<CacheLineConfig> *cache_line_config) {
159+ cache_line_config_ = cache_line_config;
160+ }
161+ void SetEnableParallel(bool enable_parallel) {
162+ enable_group_parallel_ = enable_parallel;
163+ }
164+ void SetTilingCaseIdent(TilingCaseIdent tiling_case_ident) {
165+ tiling_case_ident_ = tiling_case_ident;
166+ }
167+ void SetGroupNum(size_t group_num) {
168+ group_num_ = group_num;
169+ }
38 170 
39- enum class VarsType {171+ private:
40- PUREMC = 0,172+ static bool VarCmp(Expr &a, Expr &b);
41- LOCALBUFFER = 1,173+ void ReorderVars();
42- };174+ void GetMCArgs();
43- 175+ void GetLocalBufferTilingVars();
44- class AxesReorderSolverGen : public SolverGen, public InputOutputSettersMixin<AxesReorderSolverGen>,176+ void GetRelatedArgs(const Expr &expr, std::vector<Expr> &related_args) const;
45- public PgoConfigSettersMixin<AxesReorderSolverGen> {177+ bool NeedUBMultiCoreBalance();
46- public:178+ std::string GenGetStaticInputParam(const HardwareDef &hardware_type, bool no_type = false) const;
47- explicit AxesReorderSolverGen(const std::string &tiling_case_id, const std::string &type_name)179+ std::string GenGetObjStaticInputParam(bool no_type = false);
48- : SolverGen(tiling_case_id, type_name) {}180+ std::string GenGetObjStaticFunc();
49- ~AxesReorderSolverGen() override = default;181+ void CollectInitialWorkList(std::vector<Expr> &work_list) const;
50- std::string GenSolverClassImpl() override;182+ void CollectNeededTenaryVarsClosure(std::vector<Expr> &work_list, std::set<std::string> &needed_vars) const;
51- std::string GenSolverFuncImpl() override;183+ std::string GenTenaryVarDecls(const std::set<std::string> &needed_vars);
52- std::string GenPGOSolverFilter();184+ std::string GenSingleTenaryVar(const TernaryOp &op, const std::string &var_name, std::set<std::string> &declared_vars,
53- std::string GenSolverFuncInvoke() override;185+ std::map<std::string, std::string> &content_to_first_var);
54- std::string GenPGOSolverClassImpl();186+ std::string GenGetTilingDataObjStaticFunc();
55- std::string GenPGOSolverFuncImpl();187+ std::string GenObjFunc();
56- 188+ std::string GenGetUbSizeStaticFunc();
57- void SetInputArgs(const std::vector<Expr> &input_args) { input_args_ = input_args; }189+ std::string GenGetTilingDataUbSizeStaticFunc();
58- void SetConstArgs(const ExprUintMap &const_vars) {190+ std::string GenGetBlockDimStatic(Expr &corenum_cons);
59- std::vector<Expr> const_args;191+ std::string GenGetTilingDataBlockDimStatic(Expr &corenum_cons);
60- const_vars_map_ = const_vars;192+ std::string GenUBThresholdFunc();
61- for (const auto &pair : const_vars_map_) {193+ std::string GenUBSizeCacheLineFunc();
62- const_args.push_back(pair.first);194+ std::string GenCoreNumFunc();
63- }195+ std::pair<std::vector<Expr>, std::vector<Expr>> SortConsArgs(const Expr &expr, bool &is_mc_mixed);
64- const_args_ = const_args;196+ std::string ObtainRelatedVars(Expr &expr);
65- }197+ std::string InitiateArgs();
66- void SetBufferUseAlg(const std::map<HardwareDef, Expr> &hardware_use_map) {198+ std::string InitiateBufferConsArgs(uint32_t cons_idx, HardwareDef hardware, const Expr &cons);
67- hardware_use_map_ = hardware_use_map;199+ std::string InitiateCutConsArgs(uint32_t cons_idx, const Expr &cons, bool &is_mc_mixed);
68- }200+ std::string GenConsUbFunc(uint32_t cons_idx, const std::vector<Expr> &rel_tiling_vars,
69- void SetArgAlignMap(const ExprExprMap &arg_align_map) {201+ const std::vector<Expr> &rel_cons_vars) const;
70- arg_align_map_ = arg_align_map;202+ std::string GenConsFunc(uint32_t cons_idx, ConsType cons_type, const Expr &cons,
71- }203+ const std::vector<Expr> &rel_tiling_vars, const std::vector<Expr> &rel_cons_vars) const;
72- void SetArgPromptAlignMap(const ExprUintMap &arg_prompt_align_map) {204+ std::string SetVarCons(const Expr &arg, const std::vector<Expr> &all_cons) const;
73- arg_prompt_align_map_ = arg_prompt_align_map;205+ std::string GenUpperBoundFunc(const Expr &var);
74- }206+ std::string GenUpperBoundInfo(const Expr &var);
75- void SetArgDataTypeSizeMap(const ExprUintMap &data_type_size_map) {207+ std::string SetInputVars(InputType input_type);
76- data_type_size_map_ = data_type_size_map;208+ std::string SetInputCons(std::vector<Expr> cons) const;
77- }209+ std::string SetTilingVars(VarsType var_type);
78- void SetInputAlign(const ExprExprMap &input_align);210+ std::string GenRuntimeReorderRules();
79- void SetTotalCutCons(const std::vector<Expr> &total_cut_cons) { total_cut_cons_ = total_cut_cons; }211+ std::string GenRuntimeReorderRule(const RuntimeReorderRule &rule);
80- void SetFromAxesMap(const std::map<Expr, std::vector<Expr>, ExprCmp> &from_axes_map) { from_axes_map_ = from_axes_map; }212+ std::string GenRuntimeExprValue(const Expr &expr) const;
81- void SetVarPriority(const ExprUintMap &priority) { priority_map_ = priority; }213+ std::string GenRuntimeCompoundExprValue(const Expr &expr) const;
82- void SetContainerExpr(const ExprExprMap &container_expr) { container_expr_ = container_expr; }214+ int32_t GetLocalBufferVarIndex(const Expr &expr) const;
83- void SetContainerNames(const std::map<Expr, std::string, ExprCmp> &container_names) { container_names_ = container_names; }215+ int32_t ResolveLocalBufferVarIndex(const Expr &expr) const;
84- void SetReplaceVars(const std::vector<std::pair<Expr, Expr>> &replace_vars) {216+ std::vector<int32_t> GetRuntimePreferredIndices(const RuntimeReorderRule &rule) const;
85- for (const auto &var : replace_vars) {217+ void InitConcatPromptAlign(const Expr &local_var, const uint32_t prompt_align, std::string &strs);
86- replace_vars_.emplace_back(var);218+ std::string GenInputInfo(std::vector<Expr> &all_cons, std::vector<Expr> &local_buffer_cons,
87- }219+ std::vector<Expr> &mc_mixed_cons);
88- }220+ std::string GenInput(const TradeOffConfig &trade_off_config, std::vector<Expr> &all_cons);
89- void SetTernaryOps(const std::map<Expr, TernaryOp, ExprCmp> &ternary_ops) {221+ std::string GenSetTiling();
90- ternary_ops_ = ternary_ops;222+ std::string GenSolverRunInvoke(const std::string &class_name);
91- }223+ std::string GenEmptyTensorCheckInSolver();
92- 224+ std::string GenOriginExpr(const std::vector<Expr> &exprs, const std::string &indent) const;
93- void SetExeTimeMap(const std::map<Expr, std::vector<Expr>, ExprCmp> &exe_time_map) {225+ std::pair<std::string, std::string> GenOriginBufExpr(const Expr &expr, const std::string &indent) const;
94- for (const auto &pair : exe_time_map) {226+ std::pair<std::string, std::string> GenNamedOriginBufExpr(const Expr &expr, const std::string &indent) const;
95- exe_time_map_[pair.first] = pair.second;227+ std::string GenPgoSetTiling();
96- }228+ std::string GenPgoSetMaxBlockDim() const;
97- }229+ std::vector<uint32_t> GetArgRelateCons(const Expr &arg, const std::vector<Expr> &all_cons) const;
98- void Arrange();230+ std::string IsEnableBlockLoopTradeOffByPerf() const;
99- void SetObjFunc(const Expr &head_cost, const std::map<PipeType, Expr> pipe_2_obj_map) {231+ std::string GenPendingSearchConfigOverride();
100- head_cost_ = head_cost;232+ std::vector<Expr> mc_args_;
101- pipe_2_obj_map_ = pipe_2_obj_map;233+ std::vector<Expr> input_args_;
102- }234+ std::vector<Expr> const_args_;
103- void SetIsConcatOuterMap(const ExprUintMap &is_concat_outer_map) { is_concat_outer_map_ = is_concat_outer_map; }235+ std::vector<Expr> total_cut_cons_;
104- void SetConcatInnerDims(const std::vector<Expr> &concat_inner_dims) { concat_inner_dims_ = concat_inner_dims; }236+ std::vector<Expr> local_buffer_tiling_vars_;
105- void SetUBThreshold(const double &ub_threshold) {237+ ExprExprMap input_align_;
106- ub_threshold_ = ub_threshold;238+ ExprUintMap const_vars_map_;
107- }239+ ExprExprMap arg_align_map_;
108- void SetCoreNumThreshold(const double &corenum_threshold) {240+ ExprUintMap arg_prompt_align_map_;
109- corenum_threshold_ = corenum_threshold;241+ ExprUintMap data_type_size_map_;
110- };242+ ExprExprMap container_expr_;
111- void SetReservedUbSize(const Expr &reserved_ub_size) {243+ std::vector<std::pair<Expr, Expr>> replace_vars_;
112- reserved_ub_size_ = reserved_ub_size;244+ std::map<Expr, TernaryOp, ExprCmp> ternary_ops_;
113- };245+ std::map<Expr, std::vector<Expr>, ExprCmp> exe_time_map_;
114- void SetHighPerfTiling(const bool enable_high_perf) {246+ std::map<Expr, std::string, ExprCmp> container_names_;
115- enable_high_perf_ = enable_high_perf;247+ std::map<HardwareDef, Expr> hardware_use_map_;
116- }248+ std::map<Expr, std::vector<Expr>, ExprCmp> from_axes_map_;
117- void SetEnableEqualOrder(const bool enable_equal_order) {249+ static ExprUintMap priority_map_;
118- enable_equal_order_ = enable_equal_order;250+ ExprUintMap axes_order_;
119- }251+ ExprExprMap vars_relations_;
120- void SetSearchArgs(const std::vector<Expr> &search_args) {252+ std::map<PipeType, Expr> pipe_2_obj_map_;
121- search_args_ = search_args;253+ Expr head_cost_;
122- }254+ ExprUintMap is_concat_outer_map_;
123- void SetArrangeCode(const std::string &arrange_code) {255+ std::vector<Expr> concat_inner_dims_;
124- arrange_code_ = arrange_code;256+ ExprUintMap mc_related_ub_args_map_;
125- }257+ std::vector<Expr> search_args_;
126- void SetTilingScheduleConfigTable(const TilingScheduleConfigTable *tiling_schedule_config_table) {258+ double ub_threshold_{0.2};
127- tiling_schedule_config_table_ = tiling_schedule_config_table;259+ Expr reserved_ub_size_{CreateExpr(0)};
128- }260+ double corenum_threshold_{0.4};
129- void SetTilingScheduleConfig(const TilingScheduleConfig &tiling_schedule_config) {261+ bool enable_high_perf_{false};
130- tiling_schedule_config_ = tiling_schedule_config;262+ bool enable_equal_order_{false};
131- }263+ std::string arrange_code_;
132- void SetRuntimeReorderRules(const std::vector<RuntimeReorderRule> &runtime_reorder_rules) {264+ const TilingScheduleConfigTable *tiling_schedule_config_table_{nullptr};
133- runtime_reorder_rules_ = runtime_reorder_rules;265+ TilingScheduleConfig tiling_schedule_config_; // Model 级别的 Tiling 调度配置
134- }266+ std::vector<RuntimeReorderRule> runtime_reorder_rules_;
135- void SetCacheLineConfig(const vector<CacheLineConfig> *cache_line_config) {267+ const vector<CacheLineConfig> *cache_line_config_{nullptr};
136- cache_line_config_ = cache_line_config;268+ bool enable_group_parallel_{false};
137- }269+ size_t group_num_{1UL};
138- void SetEnableParallel(bool enable_parallel) {270+ TilingCaseIdent tiling_case_ident_{ScheduleGroupIdent{}, 0U, ""};
139- enable_group_parallel_ = enable_parallel;271+};
140- }272+bool CheckExist(const std::vector<Expr> &args, const Expr &check_arg);
141- void SetTilingCaseIdent(TilingCaseIdent tiling_case_ident) {273+std::string SetRelatedVars(const std::vector<Expr> &rel_tiling_vars, const std::vector<Expr> &rel_cons_vars);
142- tiling_case_ident_ = tiling_case_ident;274+std::string GenRelatedVars(uint32_t cons_idx, const std::vector<Expr> &rel_tiling_vars,
143- }275+ const std::vector<Expr> &rel_cons_vars);
144- void SetGroupNum(size_t group_num) {276+} // namespace att
145- group_num_ = group_num;
146- }
147- 
148- private:
149- static bool VarCmp(Expr &a, Expr &b);
150- void ReorderVars();
151- void GetMCArgs();
152- void GetLocalBufferTilingVars();
153- void GetRelatedArgs(const Expr &expr, std::vector<Expr> &related_args) const;
154- bool NeedUBMultiCoreBalance();
155- std::string GenGetStaticInputParam(const HardwareDef &hardware_type, bool no_type = false) const;
156- std::string GenGetObjStaticInputParam(bool no_type = false);
157- std::string GenGetObjStaticFunc();
158- void CollectInitialWorkList(std::vector<Expr> &work_list) const;
159- void CollectNeededTenaryVarsClosure(std::vector<Expr> &work_list, std::set<std::string> &needed_vars) const;
160- std::string GenTenaryVarDecls(const std::set<std::string> &needed_vars);
161- std::string GenSingleTenaryVar(const TernaryOp &op, const std::string &var_name,
162- std::set<std::string> &declared_vars,
163- std::map<std::string, std::string> &content_to_first_var);
164- std::string GenGetTilingDataObjStaticFunc();
165- std::string GenObjFunc();
166- std::string GenGetUbSizeStaticFunc();
167- std::string GenGetTilingDataUbSizeStaticFunc();
168- std::string GenGetBlockDimStatic(Expr &corenum_cons);
169- std::string GenGetTilingDataBlockDimStatic(Expr &corenum_cons);
170- std::string GenUBThresholdFunc();
171- std::string GenUBSizeCacheLineFunc();
172- std::string GenCoreNumFunc();
173- std::pair<std::vector<Expr>, std::vector<Expr>> SortConsArgs(const Expr &expr, bool &is_mc_mixed);
174- std::string ObtainRelatedVars(Expr &expr);
175- std::string InitiateArgs();
176- std::string InitiateBufferConsArgs(uint32_t cons_idx, HardwareDef hardware, const Expr &cons);
177- std::string InitiateCutConsArgs(uint32_t cons_idx, const Expr &cons, bool &is_mc_mixed);
178- std::string GenConsUbFunc(uint32_t cons_idx, const std::vector<Expr> &rel_tiling_vars,
179- const std::vector<Expr> &rel_cons_vars) const;
180- std::string GenConsFunc(uint32_t cons_idx, ConsType cons_type, const Expr &cons,
181- const std::vector<Expr> &rel_tiling_vars, const std::vector<Expr> &rel_cons_vars) const;
182- std::string SetVarCons(const Expr &arg, const std::vector<Expr> &all_cons) const;
183- std::string GenUpperBoundFunc(const Expr &var);
184- std::string GenUpperBoundInfo(const Expr &var);
185- std::string SetInputVars(InputType input_type);
186- std::string SetInputCons(std::vector<Expr> cons) const;
187- std::string SetTilingVars(VarsType var_type);
188- std::string GenRuntimeReorderRules();
189- std::string GenRuntimeReorderRule(const RuntimeReorderRule &rule);
190- std::string GenRuntimeExprValue(const Expr &expr) const;
191- std::string GenRuntimeCompoundExprValue(const Expr &expr) const;
192- int32_t GetLocalBufferVarIndex(const Expr &expr) const;
193- void InitConcatPromptAlign(const Expr &local_var, const uint32_t prompt_align, std::string &strs);
194- std::string GenInputInfo(std::vector<Expr> &all_cons, std::vector<Expr> &local_buffer_cons,
195- std::vector<Expr> &mc_mixed_cons);
196- std::string GenInput(const TradeOffConfig &trade_off_config, std::vector<Expr> &all_cons);
197- std::string GenSetTiling();
198- std::string GenSolverRunInvoke(const std::string &class_name);
199- std::string GenEmptyTensorCheckInSolver();
200- std::string GenOriginExpr(const std::vector<Expr> &exprs, const std::string &indent) const;
201- std::pair<std::string, std::string> GenOriginBufExpr(const Expr &expr, const std::string &indent) const;
202- std::pair<std::string, std::string> GenNamedOriginBufExpr(const Expr &expr, const std::string &indent) const;
203- std::string GenPgoSetTiling();
204- std::string GenPgoSetMaxBlockDim() const;
205- std::vector<uint32_t> GetArgRelateCons(const Expr &arg, const std::vector<Expr> &all_cons) const;
206- std::string IsEnableBlockLoopTradeOffByPerf() const;
207- std::string GenPendingSearchConfigOverride();
208- std::vector<Expr> mc_args_;
209- std::vector<Expr> input_args_;
210- std::vector<Expr> const_args_;
211- std::vector<Expr> total_cut_cons_;
212- std::vector<Expr> local_buffer_tiling_vars_;
213- ExprExprMap input_align_;
214- ExprUintMap const_vars_map_;
215- ExprExprMap arg_align_map_;
216- ExprUintMap arg_prompt_align_map_;
217- ExprUintMap data_type_size_map_;
218- ExprExprMap container_expr_;
219- std::vector<std::pair<Expr, Expr>> replace_vars_;
220- std::map<Expr, TernaryOp, ExprCmp> ternary_ops_;
221- std::map<Expr, std::vector<Expr>, ExprCmp> exe_time_map_;
222- std::map<Expr, std::string, ExprCmp> container_names_;
223- std::map<HardwareDef, Expr> hardware_use_map_;
224- std::map<Expr, std::vector<Expr>, ExprCmp> from_axes_map_;
225- static ExprUintMap priority_map_;
226- std::map<PipeType, Expr> pipe_2_obj_map_;
227- Expr head_cost_;
228- ExprUintMap is_concat_outer_map_;
229- std::vector<Expr> concat_inner_dims_;
230- ExprUintMap mc_related_ub_args_map_;
231- std::vector<Expr> search_args_;
232- double ub_threshold_{0.2};
233- Expr reserved_ub_size_{CreateExpr(0)};
234- double corenum_threshold_{0.4};
235- bool enable_high_perf_{false};
236- bool enable_equal_order_{false};
237- std::string arrange_code_;
238- const TilingScheduleConfigTable *tiling_schedule_config_table_{nullptr};
239- TilingScheduleConfig tiling_schedule_config_; // Model 级别的 Tiling 调度配置
240- std::vector<RuntimeReorderRule> runtime_reorder_rules_;
241- const vector<CacheLineConfig> *cache_line_config_ {nullptr};
242- bool enable_group_parallel_{false};
243- size_t group_num_{1UL};
244- TilingCaseIdent tiling_case_ident_{ScheduleGroupIdent{}, 0U, ""};
245- };
246- bool CheckExist(const std::vector<Expr> &args, const Expr &check_arg);
247- std::string SetRelatedVars(const std::vector<Expr> &rel_tiling_vars, const std::vector<Expr> &rel_cons_vars);
248- std::string GenRelatedVars(uint32_t cons_idx, const std::vector<Expr> &rel_tiling_vars, const std::vector<Expr> &rel_cons_vars);
249-}
250#endif277#endif
@@ -470,6 +470,8 @@ void SolverPassManager::InitSolverGen(AxesReorderSolverGen &solver_gen) {
470 solver_gen.SetExeTimeMap(args_manager_.GetTernaryOpRelatedVars());470 solver_gen.SetExeTimeMap(args_manager_.GetTernaryOpRelatedVars());
471 solver_gen.SetInputAlign(GetOriginalInputAlign());471 solver_gen.SetInputAlign(GetOriginalInputAlign());
472 solver_gen.SetVarPriority(args_manager_.GetAxesPriority());472 solver_gen.SetVarPriority(args_manager_.GetAxesPriority());
473+ solver_gen.SetAxesOrder(args_manager_.GetAxesOrder());
474+ solver_gen.SetVarsRelations(args_manager_.GetVarsRelations());
473 solver_gen.SetObjFunc(args_manager_.GetHeadCost(), args_manager_.GetObjectFunc());475 solver_gen.SetObjFunc(args_manager_.GetHeadCost(), args_manager_.GetObjectFunc());
474 solver_gen.SetUBThreshold(ub_threshold_);476 solver_gen.SetUBThreshold(ub_threshold_);
475 solver_gen.SetReservedUbSize(reserved_ub_size_);477 solver_gen.SetReservedUbSize(reserved_ub_size_);
@@ -1,9 +1,9 @@
1/**1/**
2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of 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").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.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, 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.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
@@ -15,26 +15,23 @@
15#define private public15#define private public
16#include "expr_gen/arg_list_reorder.h"16#include "expr_gen/arg_list_reorder.h"
17 17 
18- 18+namespace att {
19-namespace att{
20class TestArgListReorder : public ::testing::Test {19class TestArgListReorder : public ::testing::Test {
21 public:20 public:
22- static void TearDownTestCase()21+ static void TearDownTestCase() {
23- {
24 std::cout << "Test end." << std::endl;22 std::cout << "Test end." << std::endl;
25 }23 }
26- static void SetUpTestCase()24+ static void SetUpTestCase() {
27- {
28 std::cout << "Test begin." << std::endl;25 std::cout << "Test begin." << std::endl;
29 }26 }
30 void SetUp() override {27 void SetUp() override {
31- // Code here will be called immediately after the constructor (right28+ // Code here will be called immediately after the constructor (right
32- // before each test).29+ // before each test).
33 }30 }
34 31 
35 void TearDown() override {32 void TearDown() override {
36- // Code here will be called immediately after each test (right33+ // Code here will be called immediately after each test (right
37- // before the destructor).34+ // before the destructor).
38 }35 }
39};36};
40 37 
@@ -47,20 +44,33 @@ class TestTilingScheduleConfigTable : public TilingScheduleConfigTable {
47 TestTilingScheduleConfigTable(uint32_t cache_line_size, uint32_t vector_len_size)44 TestTilingScheduleConfigTable(uint32_t cache_line_size, uint32_t vector_len_size)
48 : cache_line_size_(cache_line_size), vector_len_size_(vector_len_size) {}45 : cache_line_size_(cache_line_size), vector_len_size_(vector_len_size) {}
49 46 
50- [[nodiscard]] bool IsEnableBlockLoopAutoTune() const override { return false; }47+ [[nodiscard]] bool IsEnableBlockLoopAutoTune() const override {
51- [[nodiscard]] bool IsEnableCacheLineCheck() const override { return true; }48+ return false;
52- [[nodiscard]] TradeOffConfig GetTradeOffConfig() const override { return {}; }49+ }
53- [[nodiscard]] double GetUbThresholdPerfValEffect() const override { return 0.0; }50+ [[nodiscard]] bool IsEnableCacheLineCheck() const override {
54- [[nodiscard]] TilingScheduleConfig GetModelTilingScheduleConfig() const override51+ return true;
55- {52+ }
53+ [[nodiscard]] TradeOffConfig GetTradeOffConfig() const override {
54+ return {};
55+ }
56+ [[nodiscard]] double GetUbThresholdPerfValEffect() const override {
57+ return 0.0;
58+ }
59+ [[nodiscard]] TilingScheduleConfig GetModelTilingScheduleConfig() const override {
56 TilingScheduleConfig config;60 TilingScheduleConfig config;
57 config.cache_line_size = cache_line_size_;61 config.cache_line_size = cache_line_size_;
58 config.vector_len_size = vector_len_size_;62 config.vector_len_size = vector_len_size_;
59 return config;63 return config;
60 }64 }
61- [[nodiscard]] uint32_t GetCacheLineSize() const override { return cache_line_size_; }65+ [[nodiscard]] uint32_t GetCacheLineSize() const override {
62- [[nodiscard]] uint32_t GetVectorLenSize() const override { return vector_len_size_; }66+ return cache_line_size_;
63- [[nodiscard]] bool IsCoreNumThresholdPenaltyEnable() const override { return false; }67+ }
68+ [[nodiscard]] uint32_t GetVectorLenSize() const override {
69+ return vector_len_size_;
70+ }
71+ [[nodiscard]] bool IsCoreNumThresholdPenaltyEnable() const override {
72+ return false;
73+ }
64 74 
65 private:75 private:
66 uint32_t cache_line_size_;76 uint32_t cache_line_size_;
@@ -79,15 +89,13 @@ struct ReduceTailSubAxes {
79 std::unique_ptr<SubAxis> tail;89 std::unique_ptr<SubAxis> tail;
80};90};
81 91 
82-std::shared_ptr<AttAxis> MakeAttAxis(const std::string &name)92+std::shared_ptr<AttAxis> MakeAttAxis(const std::string &name) {
83-{
84 auto axis = std::make_shared<AttAxis>();93 auto axis = std::make_shared<AttAxis>();
85 axis->name = name;94 axis->name = name;
86 return axis;95 return axis;
87}96}
88 97 
89-size_t GetArgIndex(const std::vector<AttAxisPtr> &arg_list, const std::string &name)98+size_t GetArgIndex(const std::vector<AttAxisPtr> &arg_list, const std::string &name) {
90-{
91 for (size_t i = 0U; i < arg_list.size(); ++i) {99 for (size_t i = 0U; i < arg_list.size(); ++i) {
92 if (arg_list[i]->name == name) {100 if (arg_list[i]->name == name) {
93 return i;101 return i;
@@ -198,16 +206,15 @@ ReduceTailSortCase BuildReduceTailSortCaseWithOriginalAxes(const Expr &reduce_si
198}206}
199} // namespace207} // namespace
200 208 
201-TEST_F(TestArgListReorder, case0)209+TEST_F(TestArgListReorder, case0) {
202-{210+ // Define TuningSpace
203- //Define TuningSpace211+ // Create node: MatMul
204- //Create node: MatMul212+ // input : [m, k][k, n]
205- //input : [m, k][k, n]213+ // repeat : [M, K][K, N]
206- //repeat : [M, K][K, N]214+ // stride : [MM, KK][KK, NN]
207- //stride : [MM, KK][KK, NN]215+ // output : [m, n, k]
208- //output : [m, n, k]216+ // repeat : [M, N, ONE]
209- //repeat : [M, N, ONE]217+ // stride : [MM, NN, ZERO]
210- //stride : [MM, NN, ZERO]
211 NodeInfo MatMul;218 NodeInfo MatMul;
212 Tensor Tensor00;219 Tensor Tensor00;
213 std::shared_ptr<Tensor> tensor00 = std::make_shared<Tensor>(Tensor00);220 std::shared_ptr<Tensor> tensor00 = std::make_shared<Tensor>(Tensor00);
@@ -250,13 +257,13 @@ TEST_F(TestArgListReorder, case0)
250 MatMul.inputs = {tensor00, tensor01};257 MatMul.inputs = {tensor00, tensor01};
251 MatMul.outputs = {tensor02};258 MatMul.outputs = {tensor02};
252 259 
253- //Create node: Load260+ // Create node: Load
254- //input : [a, b]261+ // input : [a, b]
255- //repeat : [A, B]262+ // repeat : [A, B]
256- //stride : [AA, BB]263+ // stride : [AA, BB]
257- //output : [a, b]264+ // output : [a, b]
258- //repeat : [A, B]265+ // repeat : [A, B]
259- //stride : [AA, BB]266+ // stride : [AA, BB]
260 NodeInfo Load;267 NodeInfo Load;
261 Tensor Tensor10;268 Tensor Tensor10;
262 std::shared_ptr<Tensor> tensor10 = std::make_shared<Tensor>(Tensor10);269 std::shared_ptr<Tensor> tensor10 = std::make_shared<Tensor>(Tensor10);
@@ -292,9 +299,9 @@ TEST_F(TestArgListReorder, case0)
292 tuning_space_->sub_axes.emplace_back(std::move(n));299 tuning_space_->sub_axes.emplace_back(std::move(n));
293 tuning_space_->sub_axes.emplace_back(std::move(a));300 tuning_space_->sub_axes.emplace_back(std::move(a));
294 tuning_space_->sub_axes.emplace_back(std::move(b));301 tuning_space_->sub_axes.emplace_back(std::move(b));
295- //End Define302+ // End Define
296 303 
297- //Define Modelinfo304+ // Define Modelinfo
298 ModelInfo model_info;305 ModelInfo model_info;
299 AttAxis att_m;306 AttAxis att_m;
300 std::shared_ptr<AttAxis> attaxis_m = std::make_shared<AttAxis>(att_m);307 std::shared_ptr<AttAxis> attaxis_m = std::make_shared<AttAxis>(att_m);
@@ -312,13 +319,13 @@ TEST_F(TestArgListReorder, case0)
312 attaxis_a->name = "a";319 attaxis_a->name = "a";
313 attaxis_b->name = "b";320 attaxis_b->name = "b";
314 model_info.arg_list = {attaxis_m, attaxis_k, attaxis_n, attaxis_a, attaxis_b};321 model_info.arg_list = {attaxis_m, attaxis_k, attaxis_n, attaxis_a, attaxis_b};
315- //End Define322+ // End Define
316 323 
317 ArgListReorder arg_list_reorder(tuning_space_);324 ArgListReorder arg_list_reorder(tuning_space_);
318 std::vector<AttAxisPtr> tiling_R_arg_list;325 std::vector<AttAxisPtr> tiling_R_arg_list;
319 EXPECT_EQ(arg_list_reorder.SortArgList(model_info.arg_list, tiling_R_arg_list), af::SUCCESS);326 EXPECT_EQ(arg_list_reorder.SortArgList(model_info.arg_list, tiling_R_arg_list), af::SUCCESS);
320 std::map<std::string, size_t> arg_id_map;327 std::map<std::string, size_t> arg_id_map;
321- for (size_t i=0; i < model_info.arg_list.size(); i++) {328+ for (size_t i = 0; i < model_info.arg_list.size(); i++) {
322 auto arg = model_info.arg_list[i];329 auto arg = model_info.arg_list[i];
323 arg_id_map[arg->name] = i;330 arg_id_map[arg->name] = i;
324 }331 }
@@ -326,16 +333,15 @@ TEST_F(TestArgListReorder, case0)
326 EXPECT_EQ(arg_id_map["n"], 1);333 EXPECT_EQ(arg_id_map["n"], 1);
327}334}
328 335 
329-TEST_F(TestArgListReorder, case1)336+TEST_F(TestArgListReorder, case1) {
330-{337+ // Define TuningSpace
331- //Define TuningSpace338+ // Create node: MatMul
332- //Create node: MatMul339+ // input : [m, k][k, n]
333- //input : [m, k][k, n]340+ // repeat : [M, K][K, N]
334- //repeat : [M, K][K, N]341+ // stride : [MM, KK][KK, NN]
335- //stride : [MM, KK][KK, NN]342+ // output : [m, n, k]
336- //output : [m, n, k]343+ // repeat : [M, N, ONE]
337- //repeat : [M, N, ONE]344+ // stride : [MM, NN, ZERO]
338- //stride : [MM, NN, ZERO]
339 NodeInfo node1;345 NodeInfo node1;
340 std::shared_ptr<Tensor> tensor0 = std::make_shared<Tensor>();346 std::shared_ptr<Tensor> tensor0 = std::make_shared<Tensor>();
341 std::shared_ptr<Tensor> tensor1 = std::make_shared<Tensor>();347 std::shared_ptr<Tensor> tensor1 = std::make_shared<Tensor>();
@@ -346,7 +352,7 @@ TEST_F(TestArgListReorder, case1)
346 auto z0z1 = std::make_unique<SubAxis>();352 auto z0z1 = std::make_unique<SubAxis>();
347 auto z0z1t = std::make_unique<SubAxis>();353 auto z0z1t = std::make_unique<SubAxis>();
348 auto z2t = std::make_unique<SubAxis>();354 auto z2t = std::make_unique<SubAxis>();
349- 355+ 
350 z0->name = "z0";356 z0->name = "z0";
351 z1->name = "z1";357 z1->name = "z1";
352 z2->name = "z2";358 z2->name = "z2";
@@ -368,8 +374,6 @@ TEST_F(TestArgListReorder, case1)
368 node1.inputs = {tensor0};374 node1.inputs = {tensor0};
369 node1.outputs = {tensor1};375 node1.outputs = {tensor1};
370 376 
371- 
372- 
373 auto tuning_space_ = std::make_shared<TuningSpace>();377 auto tuning_space_ = std::make_shared<TuningSpace>();
374 tuning_space_->node_infos = {node1};378 tuning_space_->node_infos = {node1};
375 tuning_space_->sub_axes.emplace_back(std::move(z0));379 tuning_space_->sub_axes.emplace_back(std::move(z0));
@@ -378,9 +382,9 @@ TEST_F(TestArgListReorder, case1)
378 tuning_space_->sub_axes.emplace_back(std::move(z0z1));382 tuning_space_->sub_axes.emplace_back(std::move(z0z1));
379 tuning_space_->sub_axes.emplace_back(std::move(z0z1t));383 tuning_space_->sub_axes.emplace_back(std::move(z0z1t));
380 tuning_space_->sub_axes.emplace_back(std::move(z2t));384 tuning_space_->sub_axes.emplace_back(std::move(z2t));
381- //End Define385+ // End Define
382 386 
383- //Define Modelinfo387+ // Define Modelinfo
384 ModelInfo model_info;388 ModelInfo model_info;
385 std::shared_ptr<AttAxis> att_z0 = std::make_shared<AttAxis>();389 std::shared_ptr<AttAxis> att_z0 = std::make_shared<AttAxis>();
386 std::shared_ptr<AttAxis> att_z1 = std::make_shared<AttAxis>();390 std::shared_ptr<AttAxis> att_z1 = std::make_shared<AttAxis>();
@@ -388,7 +392,7 @@ TEST_F(TestArgListReorder, case1)
388 std::shared_ptr<AttAxis> att_z0z1 = std::make_shared<AttAxis>();392 std::shared_ptr<AttAxis> att_z0z1 = std::make_shared<AttAxis>();
389 std::shared_ptr<AttAxis> att_z0z1t = std::make_shared<AttAxis>();393 std::shared_ptr<AttAxis> att_z0z1t = std::make_shared<AttAxis>();
390 std::shared_ptr<AttAxis> att_z2t = std::make_shared<AttAxis>();394 std::shared_ptr<AttAxis> att_z2t = std::make_shared<AttAxis>();
391- 395+ 
392 att_z0->name = "z0";396 att_z0->name = "z0";
393 att_z1->name = "z1";397 att_z1->name = "z1";
394 att_z2->name = "z2";398 att_z2->name = "z2";
@@ -402,21 +406,20 @@ TEST_F(TestArgListReorder, case1)
402 att_z2t->from_axis.emplace_back(att_z2.get());406 att_z2t->from_axis.emplace_back(att_z2.get());
403 407 
404 model_info.arg_list = {att_z0, att_z1, att_z2, att_z0z1, att_z0z1t, att_z2t};408 model_info.arg_list = {att_z0, att_z1, att_z2, att_z0z1, att_z0z1t, att_z2t};
405- //End Define409+ // End Define
406 410 
407 ArgListReorder arg_list_reorder(tuning_space_);411 ArgListReorder arg_list_reorder(tuning_space_);
408 std::vector<AttAxisPtr> tiling_R_arg_list;412 std::vector<AttAxisPtr> tiling_R_arg_list;
409 EXPECT_EQ(arg_list_reorder.SortArgList(model_info.arg_list, tiling_R_arg_list), af::SUCCESS);413 EXPECT_EQ(arg_list_reorder.SortArgList(model_info.arg_list, tiling_R_arg_list), af::SUCCESS);
410 std::map<std::string, size_t> arg_id_map;414 std::map<std::string, size_t> arg_id_map;
411- for (size_t i=0; i < model_info.arg_list.size(); i++) {415+ for (size_t i = 0; i < model_info.arg_list.size(); i++) {
412 auto arg = model_info.arg_list[i];416 auto arg = model_info.arg_list[i];
413 arg_id_map[arg->name] = i;417 arg_id_map[arg->name] = i;
414 }418 }
415 EXPECT_EQ(arg_id_map["z0z1t"] < arg_id_map["z2t"], true);419 EXPECT_EQ(arg_id_map["z0z1t"] < arg_id_map["z2t"], true);
416}420}
417 421 
418-TEST_F(TestArgListReorder, case2)422+TEST_F(TestArgListReorder, case2) {
419-{
420 NodeInfo node1;423 NodeInfo node1;
421 std::shared_ptr<Tensor> tensor0 = std::make_shared<Tensor>();424 std::shared_ptr<Tensor> tensor0 = std::make_shared<Tensor>();
422 std::shared_ptr<Tensor> tensor1 = std::make_shared<Tensor>();425 std::shared_ptr<Tensor> tensor1 = std::make_shared<Tensor>();
@@ -427,7 +430,7 @@ TEST_F(TestArgListReorder, case2)
427 auto z0z1 = std::make_unique<SubAxis>();430 auto z0z1 = std::make_unique<SubAxis>();
428 auto z0z1t = std::make_unique<SubAxis>();431 auto z0z1t = std::make_unique<SubAxis>();
429 auto z2t = std::make_unique<SubAxis>();432 auto z2t = std::make_unique<SubAxis>();
430- 433+ 
431 z0->name = "z0";434 z0->name = "z0";
432 z1->name = "z1";435 z1->name = "z1";
433 z2->name = "z2";436 z2->name = "z2";
@@ -449,8 +452,6 @@ TEST_F(TestArgListReorder, case2)
449 node1.inputs = {tensor1};452 node1.inputs = {tensor1};
450 node1.outputs = {tensor0};453 node1.outputs = {tensor0};
451 454 
452- 
453- 
454 auto tuning_space_ = std::make_shared<TuningSpace>();455 auto tuning_space_ = std::make_shared<TuningSpace>();
455 tuning_space_->node_infos = {node1};456 tuning_space_->node_infos = {node1};
456 tuning_space_->sub_axes.emplace_back(std::move(z0));457 tuning_space_->sub_axes.emplace_back(std::move(z0));
@@ -459,9 +460,9 @@ TEST_F(TestArgListReorder, case2)
459 tuning_space_->sub_axes.emplace_back(std::move(z0z1));460 tuning_space_->sub_axes.emplace_back(std::move(z0z1));
460 tuning_space_->sub_axes.emplace_back(std::move(z0z1t));461 tuning_space_->sub_axes.emplace_back(std::move(z0z1t));
461 tuning_space_->sub_axes.emplace_back(std::move(z2t));462 tuning_space_->sub_axes.emplace_back(std::move(z2t));
462- //End Define463+ // End Define
463 464 
464- //Define Modelinfo465+ // Define Modelinfo
465 ModelInfo model_info;466 ModelInfo model_info;
466 std::shared_ptr<AttAxis> att_z0 = std::make_shared<AttAxis>();467 std::shared_ptr<AttAxis> att_z0 = std::make_shared<AttAxis>();
467 std::shared_ptr<AttAxis> att_z1 = std::make_shared<AttAxis>();468 std::shared_ptr<AttAxis> att_z1 = std::make_shared<AttAxis>();
@@ -469,7 +470,7 @@ TEST_F(TestArgListReorder, case2)
469 std::shared_ptr<AttAxis> att_z0z1 = std::make_shared<AttAxis>();470 std::shared_ptr<AttAxis> att_z0z1 = std::make_shared<AttAxis>();
470 std::shared_ptr<AttAxis> att_z0z1t = std::make_shared<AttAxis>();471 std::shared_ptr<AttAxis> att_z0z1t = std::make_shared<AttAxis>();
471 std::shared_ptr<AttAxis> att_z2t = std::make_shared<AttAxis>();472 std::shared_ptr<AttAxis> att_z2t = std::make_shared<AttAxis>();
472- 473+ 
473 att_z0->name = "z0";474 att_z0->name = "z0";
474 att_z1->name = "z1";475 att_z1->name = "z1";
475 att_z2->name = "z2";476 att_z2->name = "z2";
@@ -483,13 +484,13 @@ TEST_F(TestArgListReorder, case2)
483 att_z2t->from_axis.emplace_back(att_z2.get());484 att_z2t->from_axis.emplace_back(att_z2.get());
484 485 
485 model_info.arg_list = {att_z0, att_z1, att_z2, att_z0z1, att_z0z1t, att_z2t};486 model_info.arg_list = {att_z0, att_z1, att_z2, att_z0z1, att_z0z1t, att_z2t};
486- //End Define487+ // End Define
487 488 
488 ArgListReorder arg_list_reorder(tuning_space_);489 ArgListReorder arg_list_reorder(tuning_space_);
489 std::vector<AttAxisPtr> tiling_R_arg_list;490 std::vector<AttAxisPtr> tiling_R_arg_list;
490 EXPECT_EQ(arg_list_reorder.SortArgList(model_info.arg_list, tiling_R_arg_list), af::SUCCESS);491 EXPECT_EQ(arg_list_reorder.SortArgList(model_info.arg_list, tiling_R_arg_list), af::SUCCESS);
491 std::map<std::string, size_t> arg_id_map;492 std::map<std::string, size_t> arg_id_map;
492- for (size_t i=0; i < model_info.arg_list.size(); i++) {493+ for (size_t i = 0; i < model_info.arg_list.size(); i++) {
493 auto arg = model_info.arg_list[i];494 auto arg = model_info.arg_list[i];
494 arg_id_map[arg->name] = i;495 arg_id_map[arg->name] = i;
495 }496 }
@@ -512,15 +513,17 @@ TEST_F(TestArgListReorder, keep_tiling_r_arg_list_when_reduce_block_split_withou
512 auto tail_not_small_case = BuildReduceTailSortCase(af::Symbol(256), af::Symbol(64));513 auto tail_not_small_case = BuildReduceTailSortCase(af::Symbol(256), af::Symbol(64));
513 ArgListReorder tail_not_small_reorder(tail_not_small_case.tuning_space);514 ArgListReorder tail_not_small_reorder(tail_not_small_case.tuning_space);
514 std::vector<AttAxisPtr> tail_not_small_tiling_R_arg_list;515 std::vector<AttAxisPtr> tail_not_small_tiling_R_arg_list;
515- EXPECT_EQ(tail_not_small_reorder.SortArgList(tail_not_small_case.model_info.arg_list,516+ EXPECT_EQ(
516- tail_not_small_tiling_R_arg_list), af::SUCCESS);517+ tail_not_small_reorder.SortArgList(tail_not_small_case.model_info.arg_list, tail_not_small_tiling_R_arg_list),
518+ af::SUCCESS);
517 EXPECT_FALSE(tail_not_small_tiling_R_arg_list.empty());519 EXPECT_FALSE(tail_not_small_tiling_R_arg_list.empty());
518 520 
519 auto reduce_not_large_case = BuildReduceTailSortCase(af::Symbol(128), af::Symbol(32));521 auto reduce_not_large_case = BuildReduceTailSortCase(af::Symbol(128), af::Symbol(32));
520 ArgListReorder reduce_not_large_reorder(reduce_not_large_case.tuning_space);522 ArgListReorder reduce_not_large_reorder(reduce_not_large_case.tuning_space);
521 std::vector<AttAxisPtr> reduce_not_large_tiling_R_arg_list;523 std::vector<AttAxisPtr> reduce_not_large_tiling_R_arg_list;
522 EXPECT_EQ(reduce_not_large_reorder.SortArgList(reduce_not_large_case.model_info.arg_list,524 EXPECT_EQ(reduce_not_large_reorder.SortArgList(reduce_not_large_case.model_info.arg_list,
523- reduce_not_large_tiling_R_arg_list), af::SUCCESS);525+ reduce_not_large_tiling_R_arg_list),
526+ af::SUCCESS);
524 EXPECT_FALSE(reduce_not_large_tiling_R_arg_list.empty());527 EXPECT_FALSE(reduce_not_large_tiling_R_arg_list.empty());
525}528}
526 529 
@@ -537,8 +540,9 @@ TEST_F(TestArgListReorder, keep_default_single_template_for_reduce_tile_without_
537 auto tail_not_small_case = BuildReduceTailSortCase(af::Symbol(512), af::Symbol(48), 64U, 512U, false);540 auto tail_not_small_case = BuildReduceTailSortCase(af::Symbol(512), af::Symbol(48), 64U, 512U, false);
538 ArgListReorder tail_not_small_reorder(tail_not_small_case.tuning_space);541 ArgListReorder tail_not_small_reorder(tail_not_small_case.tuning_space);
539 std::vector<AttAxisPtr> tail_not_small_tiling_R_arg_list;542 std::vector<AttAxisPtr> tail_not_small_tiling_R_arg_list;
540- EXPECT_EQ(tail_not_small_reorder.SortArgList(tail_not_small_case.model_info.arg_list,543+ EXPECT_EQ(
541- tail_not_small_tiling_R_arg_list), af::SUCCESS);544+ tail_not_small_reorder.SortArgList(tail_not_small_case.model_info.arg_list, tail_not_small_tiling_R_arg_list),
545+ af::SUCCESS);
542 EXPECT_TRUE(tail_not_small_tiling_R_arg_list.empty());546 EXPECT_TRUE(tail_not_small_tiling_R_arg_list.empty());
543 EXPECT_LT(GetArgIndex(tail_not_small_case.model_info.arg_list, "reduce"),547 EXPECT_LT(GetArgIndex(tail_not_small_case.model_info.arg_list, "reduce"),
544 GetArgIndex(tail_not_small_case.model_info.arg_list, "tail"));548 GetArgIndex(tail_not_small_case.model_info.arg_list, "tail"));
@@ -547,7 +551,8 @@ TEST_F(TestArgListReorder, keep_default_single_template_for_reduce_tile_without_
547 ArgListReorder reduce_not_large_reorder(reduce_not_large_case.tuning_space);551 ArgListReorder reduce_not_large_reorder(reduce_not_large_case.tuning_space);
548 std::vector<AttAxisPtr> reduce_not_large_tiling_R_arg_list;552 std::vector<AttAxisPtr> reduce_not_large_tiling_R_arg_list;
549 EXPECT_EQ(reduce_not_large_reorder.SortArgList(reduce_not_large_case.model_info.arg_list,553 EXPECT_EQ(reduce_not_large_reorder.SortArgList(reduce_not_large_case.model_info.arg_list,
550- reduce_not_large_tiling_R_arg_list), af::SUCCESS);554+ reduce_not_large_tiling_R_arg_list),
555+ af::SUCCESS);
551 EXPECT_TRUE(reduce_not_large_tiling_R_arg_list.empty());556 EXPECT_TRUE(reduce_not_large_tiling_R_arg_list.empty());
552 EXPECT_LT(GetArgIndex(reduce_not_large_case.model_info.arg_list, "reduce"),557 EXPECT_LT(GetArgIndex(reduce_not_large_case.model_info.arg_list, "reduce"),
553 GetArgIndex(reduce_not_large_case.model_info.arg_list, "tail"));558 GetArgIndex(reduce_not_large_case.model_info.arg_list, "tail"));
@@ -587,16 +592,39 @@ TEST_F(TestArgListReorder, record_runtime_reorder_for_dynamic_reduce_tile) {
587 EXPECT_TRUE(tiling_R_arg_list.empty());592 EXPECT_TRUE(tiling_R_arg_list.empty());
588 EXPECT_EQ(test_case.model_info.runtime_reorder_rules.size(), 1U);593 EXPECT_EQ(test_case.model_info.runtime_reorder_rules.size(), 1U);
589 const auto &rule = test_case.model_info.runtime_reorder_rules[0];594 const auto &rule = test_case.model_info.runtime_reorder_rules[0];
590- EXPECT_EQ(Str(rule.preferred_axis), "tail_size");595+ ASSERT_EQ(rule.preferred_order.size(), 2U);
591- EXPECT_EQ(Str(rule.fallback_axis), "reduce_size");596+ EXPECT_EQ(Str(rule.preferred_order[0]), "tail_size");
597+ EXPECT_EQ(Str(rule.preferred_order[1]), "reduce_size");
592 EXPECT_EQ(Str(rule.condition_axis), "origin_tail_size");598 EXPECT_EQ(Str(rule.condition_axis), "origin_tail_size");
593 EXPECT_EQ(Str(rule.compare_axis), "origin_reduce_size");599 EXPECT_EQ(Str(rule.compare_axis), "origin_reduce_size");
594 EXPECT_EQ(rule.condition_threshold, 64U);600 EXPECT_EQ(rule.condition_threshold, 64U);
595 EXPECT_EQ(rule.compare_threshold, 128U);601 EXPECT_EQ(rule.compare_threshold, 128U);
596}602}
597 603 
598-TEST_F(TestArgListReorder, v2_micro_api_len_equals_schedule_vector_len)604+TEST_F(TestArgListReorder, keep_canonical_relative_order_inside_equal_order_group) {
599-{605+ auto axis0 = MakeAttAxis("axis0");
606+ auto axis1 = MakeAttAxis("axis1");
607+ auto axis2 = MakeAttAxis("axis2");
608+ InitAttAxis(axis0, "axis0", CreateExpr("axis0"));
609+ InitAttAxis(axis1, "axis1", CreateExpr("axis1"));
610+ InitAttAxis(axis2, "axis2", CreateExpr("axis2"));
611+ axis0->order = 3U;
612+ axis1->order = 3U;
613+ axis2->order = 4U;
614+ const std::vector<AttAxisPtr> canonical_axes{axis0, axis1, axis2};
615+ 
616+ ArgListReorder arg_list_reorder(std::make_shared<TuningSpace>());
617+ RuntimeReorderRule rule;
618+ EXPECT_TRUE(arg_list_reorder.SetRuntimePreferredOrder(canonical_axes, canonical_axes, {3U, 2U, 1U}, rule));
619+ ASSERT_EQ(rule.preferred_order.size(), 3U);
620+ EXPECT_EQ(Str(rule.preferred_order[0]), "axis2");
621+ EXPECT_EQ(Str(rule.preferred_order[1]), "axis0");
622+ EXPECT_EQ(Str(rule.preferred_order[2]), "axis1");
623+ EXPECT_EQ(axis0->order, 3U);
624+ EXPECT_EQ(axis1->order, 3U);
625+}
626+ 
627+TEST_F(TestArgListReorder, v2_micro_api_len_equals_schedule_vector_len) {
600 PerfParamTableV2 perf_param_table;628 PerfParamTableV2 perf_param_table;
601 TilingScheduleConfigTableV2 tiling_schedule_config_table;629 TilingScheduleConfigTableV2 tiling_schedule_config_table;
602 EXPECT_EQ(perf_param_table.GetMicroApiLen(), tiling_schedule_config_table.GetVectorLenSize());630 EXPECT_EQ(perf_param_table.GetMicroApiLen(), tiling_schedule_config_table.GetVectorLenSize());
@@ -604,12 +632,11 @@ TEST_F(TestArgListReorder, v2_micro_api_len_equals_schedule_vector_len)
604 tiling_schedule_config_table.GetVectorLenSize());632 tiling_schedule_config_table.GetVectorLenSize());
605}633}
606 634 
607-TEST_F(TestArgListReorder, v1_micro_api_len_equals_schedule_vector_len)635+TEST_F(TestArgListReorder, v1_micro_api_len_equals_schedule_vector_len) {
608-{
609 PerfParamTableV1 perf_param_table;636 PerfParamTableV1 perf_param_table;
610 TilingScheduleConfigTableV1 tiling_schedule_config_table;637 TilingScheduleConfigTableV1 tiling_schedule_config_table;
611 EXPECT_EQ(perf_param_table.GetMicroApiLen(), tiling_schedule_config_table.GetVectorLenSize());638 EXPECT_EQ(perf_param_table.GetMicroApiLen(), tiling_schedule_config_table.GetVectorLenSize());
612 EXPECT_EQ(tiling_schedule_config_table.GetModelTilingScheduleConfig().vector_len_size,639 EXPECT_EQ(tiling_schedule_config_table.GetModelTilingScheduleConfig().vector_len_size,
613 tiling_schedule_config_table.GetVectorLenSize());640 tiling_schedule_config_table.GetVectorLenSize());
614}641}
615-} // namespace att642+} // namespace att
@@ -378,6 +378,44 @@ TEST_F(ArgsManagerUtest, test_get_searchable_vars2) {
378 EXPECT_EQ(exprs.size(), 0);378 EXPECT_EQ(exprs.size(), 0);
379}379}
380 380 
381+TEST_F(ArgsManagerUtest, test_get_axes_order_uses_unique_fallback_for_missing_or_ambiguous_axes) {
382+ const Expr valid = CreateExpr("valid");
383+ const Expr ambiguous = CreateExpr("ambiguous");
384+ const Expr missing = CreateExpr("missing");
385+ const Expr shared = CreateExpr("shared");
386+ const Expr shared_search0 = CreateExpr("shared_search0");
387+ const Expr shared_search1 = CreateExpr("shared_search1");
388+ ModelInfo model_info;
389+ auto valid_axis = MakeAxis("valid", AxisPosition::INNER, false, false, false, MakeSymVar(valid));
390+ auto ambiguous_axis0 = MakeAxis("ambiguous0", AxisPosition::INNER, false, false, false, MakeSymVar(ambiguous));
391+ auto ambiguous_axis1 = MakeAxis("ambiguous1", AxisPosition::INNER, false, false, false, MakeSymVar(ambiguous));
392+ auto shared_axis = MakeAxis("shared", AxisPosition::INNER, false, false, false, MakeSymVar(shared));
393+ valid_axis->order = 3U;
394+ ambiguous_axis0->order = 1U;
395+ ambiguous_axis1->order = 1U;
396+ shared_axis->order = 4U;
397+ model_info.arg_list = {valid_axis, ambiguous_axis0, ambiguous_axis1, shared_axis};
398+ 
399+ ArgsManager args_manager(model_info);
400+ VarInfo searchable_info;
401+ searchable_info.do_search = true;
402+ VarInfo shared_info = searchable_info;
403+ shared_info.replacement.orig_expr = shared;
404+ args_manager.vars_infos_ = {{valid, searchable_info},
405+ {ambiguous, searchable_info},
406+ {missing, searchable_info},
407+ {shared_search0, shared_info},
408+ {shared_search1, shared_info}};
409+ 
410+ const ExprUintMap axes_order = args_manager.GetAxesOrder();
411+ ASSERT_EQ(axes_order.size(), 5U);
412+ EXPECT_EQ(axes_order.at(valid), 3U);
413+ const std::set<uint32_t> fallback_orders = {axes_order.at(ambiguous), axes_order.at(missing),
414+ axes_order.at(shared_search0), axes_order.at(shared_search1)};
415+ EXPECT_EQ(fallback_orders.size(), 4U);
416+ EXPECT_GE(*fallback_orders.begin(), 5U);
417+}
418+ 
381TEST_F(ArgsManagerUtest, test_get_replaced_vars) {419TEST_F(ArgsManagerUtest, test_get_replaced_vars) {
382 ModelInfo info;420 ModelInfo info;
383 ArgsManager args_manager(info);421 ArgsManager args_manager(info);
@@ -13,6 +13,7 @@
13#include "common/ub_expr/ub_expr_types.h"13#include "common/ub_expr/ub_expr_types.h"
14#define private public14#define private public
15#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/axes_reorder_solver_code.h"
16#include "generator/solver_pass_gen/axes_reorder_solver/ub_named_expr_builder.h"17#include "generator/solver_pass_gen/axes_reorder_solver/ub_named_expr_builder.h"
17#include "gen_model_info/api_perf_register/v1/perf_param_v1.h"18#include "gen_model_info/api_perf_register/v1/perf_param_v1.h"
18using namespace att;19using namespace att;
@@ -194,8 +195,7 @@ TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleForDynamicReduceTile) {
194 Expr reduce = CreateExpr("reduce");195 Expr reduce = CreateExpr("reduce");
195 Expr tail = CreateExpr("tail");196 Expr tail = CreateExpr("tail");
196 RuntimeReorderRule rule;197 RuntimeReorderRule rule;
197- rule.preferred_axis = tail;198+ rule.preferred_order = {tail, reduce};
198- rule.fallback_axis = reduce;
199 rule.condition_axis = tail;199 rule.condition_axis = tail;
200 rule.compare_axis = reduce;200 rule.compare_axis = reduce;
201 rule.condition_threshold = 64U;201 rule.condition_threshold = 64U;
@@ -210,13 +210,11 @@ TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleForDynamicReduceTile) {
210 std::string code = solver_gen.GenRuntimeReorderRules();210 std::string code = solver_gen.GenRuntimeReorderRules();
211 EXPECT_TRUE(code.find("(tail.upper_bound(tail.upper_bound_vars) < 64)") != std::string::npos);211 EXPECT_TRUE(code.find("(tail.upper_bound(tail.upper_bound_vars) < 64)") != std::string::npos);
212 EXPECT_TRUE(code.find("(reduce.upper_bound(reduce.upper_bound_vars) > 128)") != std::string::npos);212 EXPECT_TRUE(code.find("(reduce.upper_bound(reduce.upper_bound_vars) > 128)") != std::string::npos);
213- EXPECT_TRUE(code.find("Runtime reduce tile reorder chooses preferred axis tail before fallback axis reduce") !=213+ EXPECT_TRUE(code.find("input.ordered_local_buffer_vars[0] = input.local_buffer_vars[1]") != std::string::npos);
214- std::string::npos);214+ EXPECT_TRUE(code.find("input.ordered_to_canonical[0] = 1u") != std::string::npos);
215- EXPECT_TRUE(code.find("Runtime reduce tile reorder keeps fallback axis reduce before preferred axis tail") !=215+ EXPECT_TRUE(code.find("input.ordered_local_buffer_vars[1] = input.local_buffer_vars[0]") != std::string::npos);
216- std::string::npos);216+ EXPECT_TRUE(code.find("input.ordered_to_canonical[1] = 0u") != std::string::npos);
217- EXPECT_TRUE(code.find("auto *runtime_preferred_var = input.local_buffer_vars[1]") != std::string::npos);217+ EXPECT_TRUE(code.find("input.local_buffer_vars[0] =") == std::string::npos);
218- EXPECT_TRUE(code.find("input.local_buffer_vars[1] = input.local_buffer_vars[0]") != std::string::npos);
219- EXPECT_TRUE(code.find("input.local_buffer_vars[0] = runtime_preferred_var") != std::string::npos);
220}218}
221 219 
222TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxesForDynamicReduceTile) {220TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxesForDynamicReduceTile) {
@@ -225,8 +223,7 @@ TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxesForDynamic
225 Expr origin_reduce = CreateExpr("origin_reduce");223 Expr origin_reduce = CreateExpr("origin_reduce");
226 Expr origin_tail = CreateExpr("origin_tail");224 Expr origin_tail = CreateExpr("origin_tail");
227 RuntimeReorderRule rule;225 RuntimeReorderRule rule;
228- rule.preferred_axis = tail;226+ rule.preferred_order = {tail, reduce};
229- rule.fallback_axis = reduce;
230 rule.condition_axis = origin_tail;227 rule.condition_axis = origin_tail;
231 rule.compare_axis = origin_reduce;228 rule.compare_axis = origin_reduce;
232 rule.condition_threshold = 64U;229 rule.condition_threshold = 64U;
@@ -242,9 +239,8 @@ TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxesForDynamic
242 std::string code = solver_gen.GenRuntimeReorderRules();239 std::string code = solver_gen.GenRuntimeReorderRules();
243 EXPECT_TRUE(code.find("(origin_tail.value < 64)") != std::string::npos);240 EXPECT_TRUE(code.find("(origin_tail.value < 64)") != std::string::npos);
244 EXPECT_TRUE(code.find("(origin_reduce.value > 128)") != std::string::npos);241 EXPECT_TRUE(code.find("(origin_reduce.value > 128)") != std::string::npos);
245- EXPECT_TRUE(code.find("auto *runtime_preferred_var = input.local_buffer_vars[1]") != std::string::npos);242+ EXPECT_TRUE(code.find("input.ordered_local_buffer_vars[0] = input.local_buffer_vars[1]") != std::string::npos);
246- EXPECT_TRUE(code.find("input.local_buffer_vars[1] = input.local_buffer_vars[0]") != std::string::npos);243+ EXPECT_TRUE(code.find("input.ordered_local_buffer_vars[1] = input.local_buffer_vars[0]") != std::string::npos);
247- EXPECT_TRUE(code.find("input.local_buffer_vars[0] = runtime_preferred_var") != std::string::npos);
248}244}
249 245 
250TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxisProductForDynamicReduceTile) {246TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxisProductForDynamicReduceTile) {
@@ -255,8 +251,7 @@ TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderRuleUsesOriginalAxisProductFor
255 Expr origin_tail0 = CreateExpr("origin_tail0");251 Expr origin_tail0 = CreateExpr("origin_tail0");
256 Expr origin_tail1 = CreateExpr("origin_tail1");252 Expr origin_tail1 = CreateExpr("origin_tail1");
257 RuntimeReorderRule rule;253 RuntimeReorderRule rule;
258- rule.preferred_axis = tail;254+ rule.preferred_order = {tail, reduce};
259- rule.fallback_axis = reduce;
260 rule.condition_axis = origin_tail0 * origin_tail1;255 rule.condition_axis = origin_tail0 * origin_tail1;
261 rule.compare_axis = origin_reduce0 * origin_reduce1;256 rule.compare_axis = origin_reduce0 * origin_reduce1;
262 rule.condition_threshold = 64U;257 rule.condition_threshold = 64U;
@@ -280,8 +275,7 @@ TEST_F(TestAxesReorderSolverGen, GenSolverFuncImplAppliesRuntimeReorderOnceBefor
280 Expr reduce = CreateExpr("reduce");275 Expr reduce = CreateExpr("reduce");
281 Expr tail = CreateExpr("tail");276 Expr tail = CreateExpr("tail");
282 RuntimeReorderRule rule;277 RuntimeReorderRule rule;
283- rule.preferred_axis = tail;278+ rule.preferred_order = {tail, reduce};
284- rule.fallback_axis = reduce;
285 rule.condition_axis = tail;279 rule.condition_axis = tail;
286 rule.compare_axis = reduce;280 rule.compare_axis = reduce;
287 rule.condition_threshold = 64U;281 rule.condition_threshold = 64U;
@@ -293,13 +287,135 @@ TEST_F(TestAxesReorderSolverGen, GenSolverFuncImplAppliesRuntimeReorderOnceBefor
293 solver_gen.SetRuntimeReorderRules({rule});287 solver_gen.SetRuntimeReorderRules({rule});
294 288 
295 const std::string code = solver_gen.GenSolverFuncImpl();289 const std::string code = solver_gen.GenSolverFuncImpl();
296- const std::string swap_code = "auto *runtime_preferred_var = input.local_buffer_vars[1]";290+ const std::string reorder_code = "input.ordered_local_buffer_vars[0] = input.local_buffer_vars[1]";
297- const size_t swap_pos = code.find(swap_code);291+ const size_t reorder_pos = code.find(reorder_code);
298 const size_t solver_pos = code.find("AxesReorderSolvercase_test solver(input);");292 const size_t solver_pos = code.find("AxesReorderSolvercase_test solver(input);");
299- EXPECT_EQ(CountSubstr(code, swap_code), 1U);293+ EXPECT_EQ(CountSubstr(code, reorder_code), 1U);
300- EXPECT_NE(swap_pos, std::string::npos);294+ EXPECT_NE(reorder_pos, std::string::npos);
301 EXPECT_NE(solver_pos, std::string::npos);295 EXPECT_NE(solver_pos, std::string::npos);
302- EXPECT_LT(swap_pos, solver_pos);296+ EXPECT_LT(reorder_pos, solver_pos);
297+ EXPECT_NE(code.find("TilingVariable* ordered_local_buffer_vars[2] = {&reduce, &tail"), std::string::npos);
298+ EXPECT_NE(code.find("uint32_t ordered_to_canonical[2] = {0u, 1u"), std::string::npos);
299+ EXPECT_NE(code.find("tiling_data.set_reduce(input.local_buffer_vars[0]->value)"), std::string::npos);
300+ EXPECT_NE(code.find("tiling_data.set_tail(input.local_buffer_vars[1]->value)"), std::string::npos);
301+}
302+ 
303+TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderUsesCompletePermutationAndRejectsInvalidOrder) {
304+ Expr axis0 = CreateExpr("axis0");
305+ Expr axis1 = CreateExpr("axis1");
306+ Expr axis2 = CreateExpr("axis2");
307+ RuntimeReorderRule rule;
308+ rule.preferred_order = {axis2, axis0, axis1};
309+ rule.condition_axis = axis0;
310+ rule.compare_axis = axis1;
311+ rule.condition_threshold = 64U;
312+ rule.compare_threshold = 128U;
313+ 
314+ AxesReorderSolverGen solver_gen("case_test", "TilingData");
315+ solver_gen.local_buffer_tiling_vars_ = {axis0, axis1, axis2};
316+ EXPECT_NE(
317+ solver_gen.GenRuntimeReorderRule(rule).find("input.ordered_local_buffer_vars[2] = input.local_buffer_vars[1]"),
318+ std::string::npos);
319+ 
320+ Expr ignored_axis = CreateExpr("ignored_axis");
321+ rule.preferred_order = {axis2, ignored_axis, axis0, axis1};
322+ EXPECT_NE(
323+ solver_gen.GenRuntimeReorderRule(rule).find("input.ordered_local_buffer_vars[2] = input.local_buffer_vars[1]"),
324+ std::string::npos);
325+ 
326+ rule.preferred_order = {axis2, axis2, axis1};
327+ EXPECT_TRUE(solver_gen.GenRuntimeReorderRule(rule).empty());
328+}
329+ 
330+TEST_F(TestAxesReorderSolverGen, GenRuntimeReorderResolvesReplacedVarsOneToOne) {
331+ Expr origin_reduce = CreateExpr("origin_reduce");
332+ Expr origin_tail = CreateExpr("origin_tail");
333+ Expr search_reduce = CreateExpr("search_reduce");
334+ Expr search_tail = CreateExpr("search_tail");
335+ RuntimeReorderRule rule;
336+ rule.preferred_order = {origin_tail, origin_reduce};
337+ rule.condition_axis = origin_tail;
338+ rule.compare_axis = origin_reduce;
339+ rule.condition_threshold = 64U;
340+ rule.compare_threshold = 128U;
341+ 
342+ AxesReorderSolverGen solver_gen("case_test", "TilingData");
343+ solver_gen.local_buffer_tiling_vars_ = {search_reduce, search_tail};
344+ solver_gen.vars_relations_ = {{search_reduce, origin_reduce}, {search_tail, origin_tail}};
345+ const std::string code = solver_gen.GenRuntimeReorderRule(rule);
346+ EXPECT_NE(code.find("input.ordered_local_buffer_vars[0] = input.local_buffer_vars[1]"), std::string::npos);
347+ 
348+ Expr duplicate_search = CreateExpr("duplicate_search");
349+ solver_gen.local_buffer_tiling_vars_ = {search_reduce, duplicate_search, search_tail};
350+ solver_gen.vars_relations_[duplicate_search] = origin_reduce;
351+ rule.preferred_order = {origin_reduce, origin_tail, duplicate_search};
352+ EXPECT_TRUE(solver_gen.GenRuntimeReorderRule(rule).empty());
353+}
354+ 
355+TEST_F(TestAxesReorderSolverGen, GenPGOSolverAppliesRuntimeOrderBeforeSolverConstruction) {
356+ Expr reduce = CreateExpr("reduce");
357+ Expr tail = CreateExpr("tail");
358+ RuntimeReorderRule rule;
359+ rule.preferred_order = {tail, reduce};
360+ rule.condition_axis = tail;
361+ rule.compare_axis = reduce;
362+ rule.condition_threshold = 64U;
363+ rule.compare_threshold = 128U;
364+ 
365+ AxesReorderSolverGen solver_gen("case_test", "TilingData");
366+ solver_gen.local_buffer_tiling_vars_ = {reduce, tail};
367+ solver_gen.hardware_use_map_[HardwareDef::UB] = reduce + tail;
368+ solver_gen.SetRuntimeReorderRules({rule});
369+ 
370+ const std::string code = solver_gen.GenPGOSolverFuncImpl();
371+ const size_t reorder_pos = code.find("input.ordered_local_buffer_vars[0] = input.local_buffer_vars[1]");
372+ const size_t solver_pos = code.find("PGOSolvercase_test solver(input);");
373+ ASSERT_NE(reorder_pos, std::string::npos);
374+ ASSERT_NE(solver_pos, std::string::npos);
375+ EXPECT_LT(reorder_pos, solver_pos);
376+}
377+ 
378+TEST_F(TestAxesReorderSolverGen, GenLocalBufferSolversUseOrderedAxesWithCanonicalIndices) {
379+ const std::string naive_code = GenNaiveLocalBufTiling(true);
380+ EXPECT_NE(naive_code.find("input_.ordered_local_buffer_vars[ordered_idx]"), std::string::npos);
381+ EXPECT_NE(naive_code.find("input_.ordered_to_canonical[ordered_idx]"), std::string::npos);
382+ EXPECT_NE(naive_code.find("solved_axes[canonical_idx]"), std::string::npos);
383+ EXPECT_NE(naive_code.find("ProcessSingleAxisNaive(var, canonical_idx"), std::string::npos);
384+ 
385+ const std::string binary_code = GenBinaryLocalBufTilingCore();
386+ EXPECT_NE(binary_code.find("input_.ordered_local_buffer_vars[ordered_idx]"), std::string::npos);
387+ EXPECT_NE(binary_code.find("input_.ordered_to_canonical[ordered_idx]"), std::string::npos);
388+ EXPECT_NE(binary_code.find("solved_axes[canonical_idx]"), std::string::npos);
389+}
390+ 
391+TEST_F(TestAxesReorderSolverGen, GenWorkloadBalanceUsesOrderedAxes) {
392+ const std::string code = GenWorkloadBalancePrepare();
393+ EXPECT_NE(code.find("auto *vars = input_.ordered_local_buffer_vars"), std::string::npos);
394+}
395+ 
396+TEST_F(TestAxesReorderSolverGen, GenEqualOrderAndAxisOrderRemainCanonical) {
397+ const std::string equal_order_code = GenNaiveLocalBufTiling(true);
398+ EXPECT_NE(equal_order_code.find("IdentifyEqualPriorityAxes"), std::string::npos);
399+ EXPECT_NE(equal_order_code.find("solved_axes[canonical_idx]"), std::string::npos);
400+ const std::string identify_code = GenIdentifyEqualPriorityAxes();
401+ EXPECT_NE(identify_code.find("if (pair.second.size() < kSupportMaxEqualPriorityAxes)"), std::string::npos);
402+ 
403+ Expr axis0 = CreateExpr("axis0");
404+ Expr axis1 = CreateExpr("axis1");
405+ AxesReorderSolverGen solver_gen("case_test", "TilingData");
406+ solver_gen.local_buffer_tiling_vars_ = {axis0, axis1};
407+ solver_gen.axes_order_ = {{axis0, 3U}, {axis1, 3U}};
408+ const std::string init_code = solver_gen.InitiateArgs();
409+ EXPECT_NE(init_code.find("axis0.order = 3UL"), std::string::npos);
410+ EXPECT_NE(init_code.find("axis1.order = 3UL"), std::string::npos);
411+}
412+ 
413+TEST_F(TestAxesReorderSolverGen, GenPGOEnumeratesOrderedAxesAndStoresCanonicalCandidates) {
414+ const std::string code = GenPgoSolverGenerateAllTilingData();
415+ EXPECT_NE(code.find("tilingDataVar = input_.ordered_local_buffer_vars[index]"), std::string::npos);
416+ EXPECT_NE(code.find("input_.ordered_local_buffer_vars[tmp]"), std::string::npos);
417+ EXPECT_NE(code.find("ans_item[input_.ordered_to_canonical[index]] = tilingDataVar->value"), std::string::npos);
418+ EXPECT_NE(code.find("ans_item[index] = tilingDataVar->value"), std::string::npos);
303}419}
304 420 
305TEST_F(TestAxesReorderSolverGen, TEST_GEN_SOLVER_case2) {421TEST_F(TestAxesReorderSolverGen, TEST_GEN_SOLVER_case2) {