Star837
282
代码
代码
Issues40
Pull Requests182
Actions
项目讨论
Wiki
分析
Star837
282
已合并
feat: 完善 IndirectLoad SIMD 和 SIMT 支持 #1542
xiebangrui2025创建于 7月28日
feat: 完善 IndirectLoad SIMD 和 SIMT 支持 #1542
已合并
xiebangrui2025创建于 7月28日
从xiebangrui2025/graph-autofusionIL合入到cann/graph-autofusiondevelop
共 31 个文件变更+1293-593
Mautofuse/att/gen_model_info/parser/ascend_graph_parser.cpp+1-1
@@ -878,7 +878,7 @@ af::Status AscendGraphParser::ConvertToTuningSpace(const af::AscGraph &graph) {
878 continue;878 continue;
879 }879 }
880 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);880 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);
881- if (indirect_load_behavior.uses_direct_gm_pipeline || indirect_load_behavior.skips_api_emit) {881+ if (indirect_load_behavior.uses_direct_gm_pipeline) {
882 GELOGD("[IndirectLoad] Skip tuning-space node info for node[%s], direct_gm[%d], skips_emit[%d].",882 GELOGD("[IndirectLoad] Skip tuning-space node info for node[%s], direct_gm[%d], skips_emit[%d].",
883 node->GetNamePtr(), static_cast<int32_t>(indirect_load_behavior.uses_direct_gm_pipeline),883 node->GetNamePtr(), static_cast<int32_t>(indirect_load_behavior.uses_direct_gm_pipeline),
884 static_cast<int32_t>(indirect_load_behavior.skips_api_emit));884 static_cast<int32_t>(indirect_load_behavior.skips_api_emit));
Mautofuse/codegen/codegen_kernel.cpp+2-7
@@ -28,7 +28,6 @@
28#include "indirect_load_utils.h"28#include "indirect_load_utils.h"
29#include "optimize/platform/platform_factory.h"29#include "optimize/platform/platform_factory.h"
30#include "optimize/schedule_utils.h"30#include "optimize/schedule_utils.h"
31-#include "v35/codegen/simt_scalar_call/simt_scalar_emitter.h"
32#include "common/platform_context.h"31#include "common/platform_context.h"
33#include "codegen_graph_check.h"32#include "codegen_graph_check.h"
34 33 
@@ -67,10 +66,7 @@ Status AllocateQues(TPipe &tpipe, const QueCollection &collection) {
67 const auto new_que = tpipe.ques.emplace(iter.first, TQue{iter.first, iter.second, position});66 const auto new_que = tpipe.ques.emplace(iter.first, TQue{iter.first, iter.second, position});
68 GE_CHK_BOOL_RET_STATUS(new_que.second, af::FAILED, "Codegen emplace que [%ld] failed", iter.first);67 GE_CHK_BOOL_RET_STATUS(new_que.second, af::FAILED, "Codegen emplace que [%ld] failed", iter.first);
69 }68 }
70- const auto skip_iter = collection.direct_gm_ques.find(iter.first);69+ tpipe.ques.at(iter.first).skip_init_for_simt_direct_gm = collection.direct_gm_ques.at(iter.first);
71- if (skip_iter != collection.direct_gm_ques.end()) {
72- tpipe.ques.at(iter.first).skip_init_for_simt_direct_gm = skip_iter->second;
73- }
74 }70 }
75 for (auto &[id, que] : tpipe.ques) {71 for (auto &[id, que] : tpipe.ques) {
76 if (id != tpipe.cube_output_que_id) {72 if (id != tpipe.cube_output_que_id) {
@@ -2215,8 +2211,7 @@ Status Kernel::ParseGraph(const ascir::ImplGraph &graph, const ascir::FusedSched
2215 continue;2211 continue;
2216 }2212 }
2217 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);2213 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);
2218- if (indirect_load_behavior.skips_api_emit && indirect_load_behavior.uses_direct_gm_pipeline &&2214+ if (indirect_load_behavior.skips_api_emit && !IsOps<Store>(node)) {
2219- !IsOps<Store>(node)) {
2220 continue;2215 continue;
2221 }2216 }
2222 2217 
Mautofuse/codegen/codegen_kernel_loop.cpp+2-5
@@ -19,7 +19,6 @@
19#include "api_call/utils/api_call_factory.h"19#include "api_call/utils/api_call_factory.h"
20#include "graph/ascendc_ir/utils/asc_tensor_utils.h"20#include "graph/ascendc_ir/utils/asc_tensor_utils.h"
21#include "indirect_load_utils.h"21#include "indirect_load_utils.h"
22-#include "v35/codegen/simt_scalar_call/simt_scalar_emitter.h"
23 22 
24using namespace std;23using namespace std;
25using namespace af::ops;24using namespace af::ops;
@@ -305,7 +304,7 @@ Status Loop::ConstructFromNodes(ascir::NodeViewVisitorConst nodes, const Tiler &
305 TraverseGraphForReduceNodes(nodes, current_loop->is_graph_has_reduce_node, current_loop->is_ar);304 TraverseGraphForReduceNodes(nodes, current_loop->is_graph_has_reduce_node, current_loop->is_ar);
306 auto lifecycle_edge = GetLifecycleEdge(nodes, tpipe);305 auto lifecycle_edge = GetLifecycleEdge(nodes, tpipe);
307 for (auto node : nodes) {306 for (auto node : nodes) {
308- if (IsSkippedApiEmitProcessNode(node)) {307+ if (ascgen_utils::indirect_load::GetTemplateBehavior(std::dynamic_pointer_cast<af::AscNode>(node)).skips_api_emit) {
309 GE_CHK_STATUS_RET(AddSkippedApiEmitProcessCall(node, current_loop, current_axis, tensor_calls));308 GE_CHK_STATUS_RET(AddSkippedApiEmitProcessCall(node, current_loop, current_axis, tensor_calls));
310 continue;309 continue;
311 }310 }
@@ -716,9 +715,7 @@ Status Loop::GenerateLoop(const Tiler &tiler, const TPipe &tpipe, std::vector<as
716 }715 }
717 ss << "for (" << axis.AsArg() << " = 0; " << axis << " < " << axis.loop_size.Str() << "; " << axis << "++) "716 ss << "for (" << axis.AsArg() << " = 0; " << axis << " < " << axis.loop_size.Str() << "; " << axis << "++) "
718 << "{" << std::endl;717 << "{" << std::endl;
719- const bool skip_calc_from_axis_for_indirect_load_simt =718+ if (tpipe.cv_fusion_type != ascir::CubeTemplateType::kUBFuse) {
720- axis.type == Axis::Type::kAxisTypeTileInner && FindIndirectLoadOpCall(this, true) != nullptr;
721- if (tpipe.cv_fusion_type != ascir::CubeTemplateType::kUBFuse && !skip_calc_from_axis_for_indirect_load_simt) {
722 ss << tiler.CalcFromAxis(axis.id);719 ss << tiler.CalcFromAxis(axis.id);
723 }720 }
724 GenerateEnCacheCondition(tiler, tpipe, axis, ss);721 GenerateEnCacheCondition(tiler, tpipe, axis, ss);
Mautofuse/common/indirect_load_utils.cpp+3-13
@@ -14,7 +14,6 @@
14 14 
15#include "ascir_ops.h"15#include "ascir_ops.h"
16#include "ascir_ops_utils.h"16#include "ascir_ops_utils.h"
17-#include "common_utils.h"
18#include "schedule_result.h"17#include "schedule_result.h"
19 18 
20namespace ascgen_utils::indirect_load {19namespace ascgen_utils::indirect_load {
@@ -39,30 +38,24 @@ TemplateBehavior GetBehavior(TemplateRole role) {
39 TemplateBehavior behavior;38 TemplateBehavior behavior;
40 switch (role) {39 switch (role) {
41 case TemplateRole::kSimdInputPre:40 case TemplateRole::kSimdInputPre:
42- behavior.skips_main_schedule_tiling = true;41+ behavior.excludes_tiling_group = true;
43- behavior.skips_api_emit = true;
44 behavior.preserves_vectorized_axis = true;42 behavior.preserves_vectorized_axis = true;
45 break;43 break;
46- case TemplateRole::kSimdOp:
47- break;
48 case TemplateRole::kSimtInputBoundary:44 case TemplateRole::kSimtInputBoundary:
49 behavior.skips_main_schedule_tiling = true;45 behavior.skips_main_schedule_tiling = true;
50 behavior.skips_api_emit = true;46 behavior.skips_api_emit = true;
51 behavior.uses_direct_gm_pipeline = true;47 behavior.uses_direct_gm_pipeline = true;
52- behavior.skips_ub_expr = true;
53 behavior.preserves_vectorized_axis = true;48 behavior.preserves_vectorized_axis = true;
54 break;49 break;
55 case TemplateRole::kSimtDirectGmBoundary:50 case TemplateRole::kSimtDirectGmBoundary:
56 case TemplateRole::kSimtInlineTransform:51 case TemplateRole::kSimtInlineTransform:
57 behavior.skips_api_emit = true;52 behavior.skips_api_emit = true;
58 behavior.uses_direct_gm_pipeline = true;53 behavior.uses_direct_gm_pipeline = true;
59- behavior.skips_ub_expr = true;
60 behavior.preserves_vectorized_axis = true;54 behavior.preserves_vectorized_axis = true;
61 break;55 break;
62 case TemplateRole::kSimtOp:56 case TemplateRole::kSimtOp:
63 behavior.uses_direct_gm_pipeline = true;57 behavior.uses_direct_gm_pipeline = true;
64 behavior.skips_ub_lifecycle = true;58 behavior.skips_ub_lifecycle = true;
65- behavior.skips_ub_expr = true;
66 behavior.preserves_vectorized_axis = true;59 behavior.preserves_vectorized_axis = true;
67 break;60 break;
68 case TemplateRole::kNone:61 case TemplateRole::kNone:
@@ -82,8 +75,7 @@ TemplateBehavior GetTemplateBehavior(const af::AscNodePtr &node) {
82 75 
83af::Status InheritTemplateRoleIfIL(af::AscGraph &graph, const std::string &vf_node_name, const af::AscNodePtr &src) {76af::Status InheritTemplateRoleIfIL(af::AscGraph &graph, const std::string &vf_node_name, const af::AscNodePtr &src) {
84 GE_ASSERT_NOTNULL(src);77 GE_ASSERT_NOTNULL(src);
85- const auto &behavior = GetTemplateBehavior(src);78+ if (GetTemplateRole(src) == TemplateRole::kNone) {
86- if (!behavior.skips_api_emit && !behavior.uses_direct_gm_pipeline) {
87 return af::SUCCESS;79 return af::SUCCESS;
88 }80 }
89 auto vf_node = graph.FindNode(vf_node_name.c_str());81 auto vf_node = graph.FindNode(vf_node_name.c_str());
@@ -157,9 +149,7 @@ bool ShouldApplyInputInnerVectorization(const af::AscNodePtr &node) {
157}149}
158 150 
159bool ShouldDisableRegularVectorFunc(const af::AscNodePtr &node) {151bool ShouldDisableRegularVectorFunc(const af::AscNodePtr &node) {
160- const TemplateRole role = GetTemplateRole(node);152+ return GetTemplateBehavior(node).uses_direct_gm_pipeline;
161- return role == TemplateRole::kSimtInputBoundary || role == TemplateRole::kSimtDirectGmBoundary ||
162- role == TemplateRole::kSimtInlineTransform || role == TemplateRole::kSimtOp;
163}153}
164 154 
165af::AscNodePtr GetInputProducer(const af::AscNodePtr &node, size_t input_index) {155af::AscNodePtr GetInputProducer(const af::AscNodePtr &node, size_t input_index) {
Mautofuse/common/indirect_load_utils.h+1-2
@@ -21,7 +21,6 @@ namespace ascgen_utils::indirect_load {
21enum class TemplateRole : int64_t {21enum class TemplateRole : int64_t {
22 kNone,22 kNone,
23 kSimdInputPre,23 kSimdInputPre,
24- kSimdOp,
25 kSimtInputBoundary,24 kSimtInputBoundary,
26 kSimtDirectGmBoundary,25 kSimtDirectGmBoundary,
27 kSimtInlineTransform,26 kSimtInlineTransform,
@@ -29,11 +28,11 @@ enum class TemplateRole : int64_t {
29};28};
30 29 
31struct TemplateBehavior {30struct TemplateBehavior {
31+ bool excludes_tiling_group = false;
32 bool skips_main_schedule_tiling = false;32 bool skips_main_schedule_tiling = false;
33 bool skips_api_emit = false;33 bool skips_api_emit = false;
34 bool uses_direct_gm_pipeline = false;34 bool uses_direct_gm_pipeline = false;
35 bool skips_ub_lifecycle = false;35 bool skips_ub_lifecycle = false;
36- bool skips_ub_expr = false;
37 bool preserves_vectorized_axis = false;36 bool preserves_vectorized_axis = false;
38};37};
39 38 
Mautofuse/common/ub_expr/asc_graph_ub_expr_builder.cpp+1-1
@@ -307,7 +307,7 @@ af::Status AscGraphUbExprBuilder::Build(const af::AscGraph &graph, UbExprContext
307 std::map<int64_t, ContainerState> buffer_bytes;307 std::map<int64_t, ContainerState> buffer_bytes;
308 for (const auto &node : graph.GetAllNodes()) {308 for (const auto &node : graph.GetAllNodes()) {
309 GE_ASSERT_NOTNULL(node);309 GE_ASSERT_NOTNULL(node);
310- if (ascgen_utils::indirect_load::GetTemplateBehavior(node).skips_ub_expr) {310+ if (ascgen_utils::indirect_load::GetTemplateBehavior(node).uses_direct_gm_pipeline) {
311 continue;311 continue;
312 }312 }
313 for (const auto &output : node->outputs()) {313 for (const auto &output : node->outputs()) {
Mautofuse/inc/graph_metadef/graph/ascendc_ir/ascir_registry.h+0-13
@@ -188,19 +188,6 @@ class AscIrCodegen {
188 (void)node;188 (void)node;
189 return true;189 return true;
190 }190 }
191- 
192- virtual bool IsSimtScalarSupported(const AscNode &node) const {
193- (void)node;
194- return false;
195- }
196- 
197- virtual ge::graphStatus GenerateSimtScalarExpr(const AscNode &node, const std::vector<std::string> &inputs,
198- std::string &expr) const {
199- (void)node;
200- (void)inputs;
201- (void)expr;
202- return af::FAILED;
203- }
204};191};
205 192 
206class AscIrAtt {193class AscIrAtt {
Mautofuse/optimize/autoschedule/schedule.cpp+13-6
@@ -520,7 +520,8 @@ Status AddIndirectLoadSyntheticOuterAxis(const af::AscNodePtr &node, ascir::Axis
520 return af::SUCCESS;520 return af::SUCCESS;
521}521}
522 522 
523-Status ApplyIndirectLoadTemplateMerge(ascir::ImplGraph &graph, const af::AscNodePtr &node, ascir::AxisId axis_id) {523+Status ApplyIndirectLoadTemplateMerge(ascir::ImplGraph &graph, const af::AscNodePtr &node, ascir::AxisId axis_id,
524+ bool merge_tensor_axis = true) {
524 if (axis_id == af::kIdNone) {525 if (axis_id == af::kIdNone) {
525 return af::SUCCESS;526 return af::SUCCESS;
526 }527 }
@@ -531,15 +532,19 @@ Status ApplyIndirectLoadTemplateMerge(ascir::ImplGraph &graph, const af::AscNode
531 node->GetNamePtr());532 node->GetNamePtr());
532 GE_ASSERT_TRUE(graph.ApplySchedAxisMerge(node, axis_id), "Failed to merge schedule axis[%ld] for node[%s].",533 GE_ASSERT_TRUE(graph.ApplySchedAxisMerge(node, axis_id), "Failed to merge schedule axis[%ld] for node[%s].",
533 axis_id, node->GetNamePtr());534 axis_id, node->GetNamePtr());
534- GE_ASSERT_TRUE(graph.ApplyTensorAxisMerge(node, axis_id), "Failed to merge tensor axis[%ld] for node[%s].", axis_id,535+ if (merge_tensor_axis) {
535- node->GetNamePtr());536+ GE_ASSERT_TRUE(graph.ApplyTensorAxisMerge(node, axis_id), "Failed to merge tensor axis[%ld] for node[%s].",
537+ axis_id, node->GetNamePtr());
538+ }
536 } else {539 } else {
537 GELOGD("[IndirectLoad] Graph[%s] apply single-axis template merge axis[%ld] for node[%s].", graph.GetName().c_str(),540 GELOGD("[IndirectLoad] Graph[%s] apply single-axis template merge axis[%ld] for node[%s].", graph.GetName().c_str(),
538 axis_id, node->GetNamePtr());541 axis_id, node->GetNamePtr());
539 GE_ASSERT_TRUE(graph.ApplySchedAxisMerge(node, axis_id, {axis_id}),542 GE_ASSERT_TRUE(graph.ApplySchedAxisMerge(node, axis_id, {axis_id}),
540 "Failed to merge schedule axis[%ld] for node[%s].", axis_id, node->GetNamePtr());543 "Failed to merge schedule axis[%ld] for node[%s].", axis_id, node->GetNamePtr());
541- GE_ASSERT_TRUE(graph.ApplyTensorAxisMerge(node, axis_id, {axis_id}),544+ if (merge_tensor_axis) {
542- "Failed to merge tensor axis[%ld] for node[%s].", axis_id, node->GetNamePtr());545+ GE_ASSERT_TRUE(graph.ApplyTensorAxisMerge(node, axis_id, {axis_id}),
546+ "Failed to merge tensor axis[%ld] for node[%s].", axis_id, node->GetNamePtr());
547+ }
543 }548 }
544 return af::SUCCESS;549 return af::SUCCESS;
545}550}
@@ -598,8 +603,10 @@ Status Scheduler::ApplyIndirectLoadNodeAxes(const af::AscNodePtr &node, bool &sk
598 return af::SUCCESS;603 return af::SUCCESS;
599 }604 }
600 if (ascgen_utils::indirect_load::ShouldApplyInputInnerVectorization(node)) {605 if (ascgen_utils::indirect_load::ShouldApplyInputInnerVectorization(node)) {
601- skip_main_tiling = true;
602 GE_ASSERT_SUCCESS(ApplyInputInnerVectorizedAxis(graph_, node, indirect_load_axes_.input_inner_axis));606 GE_ASSERT_SUCCESS(ApplyInputInnerVectorizedAxis(graph_, node, indirect_load_axes_.input_inner_axis));
607+ GE_ASSERT_SUCCESS(AddIndirectLoadSyntheticOuterAxis(node, indirect_load_axes_.outer_axis,
608+ has_indirect_load_synthetic_outer_axis_));
609+ GE_ASSERT_SUCCESS(ApplyIndirectLoadTemplateMerge(graph_, node, indirect_load_axes_.outer_axis, false));
603 return af::SUCCESS;610 return af::SUCCESS;
604 }611 }
605 if (ascgen_utils::indirect_load::ShouldSkipMainScheduleTiling(node)) {612 if (ascgen_utils::indirect_load::ShouldSkipMainScheduleTiling(node)) {
Mautofuse/optimize/autoschedule/tiling_group.cpp+2-2
@@ -484,8 +484,8 @@ Status TilingGroup::GenTilingGroup(const ascir::ImplGraph &impl_graph, AxisGroup
484 std::vector<std::pair<std::string, AxisGroup>> node_name_to_tiling_group;484 std::vector<std::pair<std::string, AxisGroup>> node_name_to_tiling_group;
485 std::set<af::AxisId> n_groupset;485 std::set<af::AxisId> n_groupset;
486 for (const auto &node : impl_graph.GetAllNodes()) {486 for (const auto &node : impl_graph.GetAllNodes()) {
487- if (ScheduleUtils::IsBuffer(node) ||487+ const auto behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);
488- ascgen_utils::indirect_load::GetTemplateBehavior(node).skips_main_schedule_tiling) {488+ if (ScheduleUtils::IsBuffer(node) || behavior.skips_main_schedule_tiling || behavior.excludes_tiling_group) {
489 continue;489 continue;
490 }490 }
491 AxisGroup single_node_axes_group;491 AxisGroup single_node_axes_group;
Mautofuse/optimize/platform/common/base_alignment_strategy.cpp+1-2
@@ -386,8 +386,7 @@ af::Status BaseAlignmentStrategy::AlignVectorizedStrides(ascir::ImplGraph &impl_
386 for (const auto &node : impl_graph.GetAllNodes()) {386 for (const auto &node : impl_graph.GetAllNodes()) {
387 GE_ASSERT_NOTNULL(node);387 GE_ASSERT_NOTNULL(node);
388 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);388 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node);
389- if (ScheduleUtils::IsBuffer(node) || indirect_load_behavior.uses_direct_gm_pipeline ||389+ if (ScheduleUtils::IsBuffer(node) || indirect_load_behavior.uses_direct_gm_pipeline) {
390- indirect_load_behavior.skips_api_emit) {
391 continue;390 continue;
392 }391 }
393 GE_ASSERT_SUCCESS(SetVectorizedStridesForOneNode(node));392 GE_ASSERT_SUCCESS(SetVectorizedStridesForOneNode(node));
Mautofuse/optimize/task_generator/indirect_load_schedule_case_generator.cpp+95-53
@@ -12,12 +12,12 @@
12 12 
13#include "ascir_ops.h"13#include "ascir_ops.h"
14#include "ascir_ops_utils.h"14#include "ascir_ops_utils.h"
15+#include "v35/ascir/ascir_codegen_v2.h"
15#include "common_utils.h"16#include "common_utils.h"
16#include "graph_utils.h"17#include "graph_utils.h"
17#include "graph/symbolizer/symbolic.h"18#include "graph/symbolizer/symbolic.h"
18#include "indirect_load_utils.h"19#include "indirect_load_utils.h"
19#include "schedule_result.h"20#include "schedule_result.h"
20-#include "schedule_utils.h"
21 21 
22#include <algorithm>22#include <algorithm>
23#include <sstream>23#include <sstream>
@@ -76,15 +76,25 @@ bool HasExp2(const af::AscGraph &graph) {
76 return false;76 return false;
77}77}
78 78 
79-bool IsTemplateCandidateLegal(ascir::TemplateId template_id, const IndirectLoadGraphPaths &paths) {79+af::Status ValidateTemplateCandidate(ascir::TemplateId template_id, const IndirectLoadGraphPaths &paths,
80+ bool &is_candidate_legal) {
81+ is_candidate_legal = true;
80 if (template_id != ascir::TemplateId::kIndirectLoadSimt) {82 if (template_id != ascir::TemplateId::kIndirectLoadSimt) {
81- return true;83+ return af::SUCCESS;
82 }84 }
85+ is_candidate_legal = false;
83 if (paths.data_input.empty()) {86 if (paths.data_input.empty()) {
84- return false;87+ return af::SUCCESS;
85 }88 }
86 const auto &input_node = paths.data_input.front();89 const auto &input_node = paths.data_input.front();
87- return af::ops::IsOps<af::ascir_op::Data>(input_node) || af::ops::IsOps<af::ascir_op::Load>(input_node);90+ if (af::ops::IsOps<af::ascir_op::Data>(input_node)) {
91+ is_candidate_legal = true;
92+ } else if (af::ops::IsOps<af::ascir_op::Load>(input_node)) {
93+ GE_ASSERT_SUCCESS(ascgen_utils::indirect_load::SetTemplateRole(
94+ input_node, ascgen_utils::indirect_load::TemplateRole::kSimtInputBoundary));
95+ is_candidate_legal = true;
96+ }
97+ return af::SUCCESS;
88}98}
89 99 
90af::Status GetIndirectLoadAxis(const af::AscNodePtr &node, int64_t &axis) {100af::Status GetIndirectLoadAxis(const af::AscNodePtr &node, int64_t &axis) {
@@ -142,12 +152,12 @@ af::Status MergeAxesForTemplate(af::AscGraph &graph, const std::vector<af::AxisI
142 return af::SUCCESS;152 return af::SUCCESS;
143}153}
144 154 
145-af::Status CreateFixedTileSplit(af::AscGraph &graph, af::AxisId axis_id, af::AxisId &outer_id, af::AxisId &inner_id) {155+af::Status CreateFixedTileSplit(af::AscGraph &graph, af::AxisId axis_id) {
146 const auto *axis = graph.FindAxis(axis_id);156 const auto *axis = graph.FindAxis(axis_id);
147 GE_ASSERT_NOTNULL(axis, "IndirectLoad fixed tile axis %ld is not found.", axis_id);157 GE_ASSERT_NOTNULL(axis, "IndirectLoad fixed tile axis %ld is not found.", axis_id);
148- outer_id =158+ const af::AxisId outer_id =
149 graph.CreateAxis(axis->name + "T", ascir::Axis::Type::kAxisTypeTileOuter, axis->size, {axis_id}, af::kIdNone).id;159 graph.CreateAxis(axis->name + "T", ascir::Axis::Type::kAxisTypeTileOuter, axis->size, {axis_id}, af::kIdNone).id;
150- inner_id =160+ const af::AxisId inner_id =
151 graph161 graph
152 .CreateAxis(axis->name + "t", ascir::Axis::Type::kAxisTypeTileInner, af::sym::kSymbolOne, {axis_id}, outer_id)162 .CreateAxis(axis->name + "t", ascir::Axis::Type::kAxisTypeTileInner, af::sym::kSymbolOne, {axis_id}, outer_id)
153 .id;163 .id;
@@ -168,12 +178,11 @@ af::Status BuildSimdInputInnerAxis(af::AscGraph &graph, const af::AscNodePtr &in
168 return af::SUCCESS;178 return af::SUCCESS;
169}179}
170af::Status BuildAxisViewByBoundary(af::AscGraph &graph, const std::vector<af::AxisId> &axes, size_t boundary,180af::Status BuildAxisViewByBoundary(af::AscGraph &graph, const std::vector<af::AxisId> &axes, size_t boundary,
171- af::AxisId &outer_axis, af::AxisId &inner_axis, std::vector<af::AxisId> &outer_axes,181+ af::AxisId &outer_axis, af::AxisId &inner_axis) {
172- std::vector<af::AxisId> &inner_axes) {
173 GE_ASSERT_TRUE(!axes.empty(), "IndirectLoad output axis is empty.");182 GE_ASSERT_TRUE(!axes.empty(), "IndirectLoad output axis is empty.");
174 const size_t split = std::min(boundary, axes.size());183 const size_t split = std::min(boundary, axes.size());
175- outer_axes.assign(axes.begin(), axes.begin() + static_cast<int64_t>(split));184+ const std::vector<af::AxisId> outer_axes(axes.begin(), axes.begin() + static_cast<int64_t>(split));
176- inner_axes.assign(axes.begin() + static_cast<int64_t>(split), axes.end());185+ const std::vector<af::AxisId> inner_axes(axes.begin() + static_cast<int64_t>(split), axes.end());
177 if (outer_axes.empty()) {186 if (outer_axes.empty()) {
178 outer_axis = graph.CreateAxis("indirect_load_single_outer", af::sym::kSymbolOne).id;187 outer_axis = graph.CreateAxis("indirect_load_single_outer", af::sym::kSymbolOne).id;
179 } else {188 } else {
@@ -192,13 +201,8 @@ af::Status NormalizeAxesForTemplate(af::AscGraph &graph, const af::AscNodePtr &i
192 const size_t boundary = static_cast<size_t>(axis < 0L ? axis + rank : axis);201 const size_t boundary = static_cast<size_t>(axis < 0L ? axis + rank : axis);
193 ascir::AxisId outer_axis = af::kIdNone;202 ascir::AxisId outer_axis = af::kIdNone;
194 ascir::AxisId inner_axis = af::kIdNone;203 ascir::AxisId inner_axis = af::kIdNone;
195- std::vector<ascir::AxisId> outer_axes;204+ GE_ASSERT_SUCCESS(BuildAxisViewByBoundary(graph, output_axes, boundary, outer_axis, inner_axis));
196- std::vector<ascir::AxisId> inner_axes;205+ GE_ASSERT_SUCCESS(CreateFixedTileSplit(graph, outer_axis));
197- GE_ASSERT_SUCCESS(
198- BuildAxisViewByBoundary(graph, output_axes, boundary, outer_axis, inner_axis, outer_axes, inner_axes));
199- ascir::AxisId fixed_tile_outer_axis = af::kIdNone;
200- ascir::AxisId fixed_tile_inner_axis = af::kIdNone;
201- GE_ASSERT_SUCCESS(CreateFixedTileSplit(graph, outer_axis, fixed_tile_outer_axis, fixed_tile_inner_axis));
202 GE_ASSERT_SUCCESS(206 GE_ASSERT_SUCCESS(
203 ascgen_utils::indirect_load::SetTemplateAxes(indirect_load, {outer_axis, inner_axis, input_inner_axis}));207 ascgen_utils::indirect_load::SetTemplateAxes(indirect_load, {outer_axis, inner_axis, input_inner_axis}));
204 return af::SUCCESS;208 return af::SUCCESS;
@@ -217,6 +221,22 @@ af::Status NormalizeSimdAxesForTemplate(af::AscGraph &graph, const af::AscNodePt
217 if (input_producer != nullptr && !af::ops::IsOps<af::ascir_op::Data>(input_producer)) {221 if (input_producer != nullptr && !af::ops::IsOps<af::ascir_op::Data>(input_producer)) {
218 GE_ASSERT_SUCCESS(BuildSimdInputInnerAxis(graph, input_producer, axis_index, input_inner_axis));222 GE_ASSERT_SUCCESS(BuildSimdInputInnerAxis(graph, input_producer, axis_index, input_inner_axis));
219 }223 }
224+ for (const auto &node : paths.data_input) {
225+ if (af::ops::IsOps<af::ascir_op::Data>(node)) {
226+ for (const auto &output : node->outputs()) {
227+ std::copy_n(output_axes.begin(), axis_index, output->attr.axis.begin());
228+ }
229+ break;
230+ }
231+ if (ascgen_utils::indirect_load::GetTemplateRole(node) !=
232+ ascgen_utils::indirect_load::TemplateRole::kSimdInputPre) {
233+ break;
234+ }
235+ std::copy_n(output_axes.begin(), axis_index, node->attr.sched.axis.begin());
236+ for (const auto &output : node->outputs()) {
237+ std::copy_n(output_axes.begin(), axis_index, output->attr.axis.begin());
238+ }
239+ }
220 return NormalizeAxesForTemplate(graph, indirect_load, axis, rank, input_inner_axis);240 return NormalizeAxesForTemplate(graph, indirect_load, axis, rank, input_inner_axis);
221}241}
222 242 
@@ -232,7 +252,8 @@ bool CanEmitSimtScalar(const af::AscNodePtr &node) {
232 return false;252 return false;
233 }253 }
234 const auto impl = ascgen_utils::GetAscIrCodegenImpl(node->GetType());254 const auto impl = ascgen_utils::GetAscIrCodegenImpl(node->GetType());
235- return impl != nullptr && impl->IsSimtScalarSupported(*node);255+ const auto *v2_impl = impl == nullptr ? nullptr : dynamic_cast<af::ascir::AscIrCodegenV2 *>(impl.get());
256+ return v2_impl != nullptr && v2_impl->IsSimtScalarSupported(*node);
236}257}
237 258 
238bool CollectMovableSimtInputPreNodes(const af::AscNodePtr &indirect_load, const NodePath &data_input, NodePath &nodes) {259bool CollectMovableSimtInputPreNodes(const af::AscNodePtr &indirect_load, const NodePath &data_input, NodePath &nodes) {
@@ -252,7 +273,34 @@ bool CollectMovableSimtInputPreNodes(const af::AscNodePtr &indirect_load, const
252 return false;273 return false;
253}274}
254 275 
255-af::Status MoveSimtInputPreNode(const af::AscNodePtr &node, const af::AscNodePtr &indirect_load) {276+bool IsSimdInputShapeEnlarged(const af::AscNodePtr &indirect_load) {
277+ af::Expression input_numel = af::sym::kSymbolOne;
278+ for (const af::Expression &repeat : indirect_load->inputs()[0]->attr.repeats) {
279+ input_numel = input_numel * repeat;
280+ }
281+ af::Expression output_numel = af::sym::kSymbolOne;
282+ for (const af::Expression &repeat : indirect_load->outputs()[0]->attr.repeats) {
283+ output_numel = output_numel * repeat;
284+ }
285+ return af::SymbolicUtils::StaticCheckGt(output_numel, input_numel) == af::TriBool::kTrue;
286+}
287+ 
288+void CollectMovableSimdInputPreNodes(const af::AscNodePtr &indirect_load, const NodePath &data_input, NodePath &nodes) {
289+ af::AscNodePtr consumer = indirect_load;
290+ for (const auto &current : data_input) {
291+ if (af::ops::IsOps<af::ascir_op::Data>(current) || af::ops::IsOps<af::ascir_op::Load>(current) ||
292+ current->inputs.Size() != 1UL || current->outputs().size() != 1UL ||
293+ current->attr.api.compute_type != af::ComputeType::kComputeElewise || current->GetInControlNodesSize() != 0UL ||
294+ current->GetOutControlNodesSize() != 0UL ||
295+ ascgen_utils::indirect_load::GetOnlyOutputConsumer(current) != consumer) {
296+ return;
297+ }
298+ nodes.emplace_back(current);
299+ consumer = current;
300+ }
301+}
302+ 
303+af::Status MoveInputPreNode(const af::AscNodePtr &node, const af::AscNodePtr &indirect_load) {
256 const auto producer = ascgen_utils::indirect_load::GetInputProducer(node, 0UL);304 const auto producer = ascgen_utils::indirect_load::GetInputProducer(node, 0UL);
257 GE_ASSERT_NOTNULL(producer);305 GE_ASSERT_NOTNULL(producer);
258 const auto producer_out = producer->GetOutDataAnchor(0UL);306 const auto producer_out = producer->GetOutDataAnchor(0UL);
@@ -281,7 +329,7 @@ af::Status MoveSimtInputPreNode(const af::AscNodePtr &node, const af::AscNodePtr
281 node_output->attr.repeats = indirect_output->attr.repeats;329 node_output->attr.repeats = indirect_output->attr.repeats;
282 node_output->attr.strides = indirect_output->attr.strides;330 node_output->attr.strides = indirect_output->attr.strides;
283 indirect_output->attr.dtype = producer->outputs()[0]->attr.dtype;331 indirect_output->attr.dtype = producer->outputs()[0]->attr.dtype;
284- GELOGD("[IndirectLoad] Move SIMT input pre node[%s] after IndirectLoad[%s].", node->GetNamePtr(),332+ GELOGD("[IndirectLoad] Move input pre node[%s] after IndirectLoad[%s].", node->GetNamePtr(),
285 indirect_load->GetNamePtr());333 indirect_load->GetNamePtr());
286 return af::SUCCESS;334 return af::SUCCESS;
287}335}
@@ -295,7 +343,23 @@ af::Status MoveSimtInputPreNodes(const af::AscNodePtr &indirect_load, IndirectLo
295 return af::SUCCESS;343 return af::SUCCESS;
296 }344 }
297 for (const auto &node : nodes) {345 for (const auto &node : nodes) {
298- GE_ASSERT_SUCCESS(MoveSimtInputPreNode(node, indirect_load));346+ GE_ASSERT_SUCCESS(MoveInputPreNode(node, indirect_load));
347+ }
348+ paths.data_input.erase(paths.data_input.begin(), paths.data_input.begin() + static_cast<int64_t>(nodes.size()));
349+ paths.output.insert(paths.output.begin(), nodes.rbegin(), nodes.rend());
350+ return af::SUCCESS;
351+}
352+ 
353+af::Status MoveSimdInputPreNodes(const af::AscNodePtr &indirect_load, IndirectLoadGraphPaths &paths) {
354+ if (IsSimdInputShapeEnlarged(indirect_load) || indirect_load->GetInControlNodesSize() != 0UL ||
355+ indirect_load->GetOutControlNodesSize() != 0UL ||
356+ ascgen_utils::indirect_load::GetOnlyOutputConsumer(indirect_load) == nullptr) {
357+ return af::SUCCESS;
358+ }
359+ NodePath nodes;
360+ CollectMovableSimdInputPreNodes(indirect_load, paths.data_input, nodes);
361+ for (const auto &node : nodes) {
362+ GE_ASSERT_SUCCESS(MoveInputPreNode(node, indirect_load));
299 }363 }
300 paths.data_input.erase(paths.data_input.begin(), paths.data_input.begin() + static_cast<int64_t>(nodes.size()));364 paths.data_input.erase(paths.data_input.begin(), paths.data_input.begin() + static_cast<int64_t>(nodes.size()));
301 paths.output.insert(paths.output.begin(), nodes.rbegin(), nodes.rend());365 paths.output.insert(paths.output.begin(), nodes.rbegin(), nodes.rend());
@@ -346,24 +410,22 @@ af::Status RecordTemplateLogicalView(const af::AscNodePtr &indirect_load) {
346af::Status PrepareCandidateGraph(const af::AscNodePtr &indirect_load, ascir::TemplateId template_id,410af::Status PrepareCandidateGraph(const af::AscNodePtr &indirect_load, ascir::TemplateId template_id,
347 IndirectLoadGraphPaths &paths, bool &is_candidate_legal) {411 IndirectLoadGraphPaths &paths, bool &is_candidate_legal) {
348 paths = CollectGraphPaths(indirect_load);412 paths = CollectGraphPaths(indirect_load);
349- if (template_id == ascir::TemplateId::kIndirectLoadSimt) {413+ if (template_id == ascir::TemplateId::kIndirectLoadSimd) {
414+ GE_ASSERT_SUCCESS(MoveSimdInputPreNodes(indirect_load, paths));
415+ } else if (template_id == ascir::TemplateId::kIndirectLoadSimt) {
350 GE_ASSERT_SUCCESS(MoveSimtInputPreNodes(indirect_load, paths));416 GE_ASSERT_SUCCESS(MoveSimtInputPreNodes(indirect_load, paths));
351 }417 }
352- is_candidate_legal = IsTemplateCandidateLegal(template_id, paths);418+ GE_ASSERT_SUCCESS(ValidateTemplateCandidate(template_id, paths, is_candidate_legal));
353 if (!is_candidate_legal) {419 if (!is_candidate_legal) {
354 return af::SUCCESS;420 return af::SUCCESS;
355 }421 }
356 GE_ASSERT_SUCCESS(PropagateInputTensorAttrs(paths));422 GE_ASSERT_SUCCESS(PropagateInputTensorAttrs(paths));
357 GE_ASSERT_SUCCESS(RecordTemplateLogicalView(indirect_load));423 GE_ASSERT_SUCCESS(RecordTemplateLogicalView(indirect_load));
358- GE_ASSERT_SUCCESS(::ascir::SetTemplateId(indirect_load, template_id));
359 return af::SUCCESS;424 return af::SUCCESS;
360}425}
361 426 
362af::Status AnnotateSimdTemplateRoles(const af::AscNodePtr &indirect_load, const IndirectLoadGraphPaths &paths) {427af::Status AnnotateSimdTemplateRoles(const af::AscNodePtr &indirect_load, const IndirectLoadGraphPaths &paths) {
363 GE_ASSERT_NOTNULL(indirect_load);428 GE_ASSERT_NOTNULL(indirect_load);
364- GE_ASSERT_SUCCESS(
365- ascgen_utils::indirect_load::SetTemplateRole(indirect_load, ascgen_utils::indirect_load::TemplateRole::kSimdOp));
366- 
367 for (const auto &node : paths.data_input) {429 for (const auto &node : paths.data_input) {
368 if (af::ops::IsOps<af::ascir_op::Data>(node)) {430 if (af::ops::IsOps<af::ascir_op::Data>(node)) {
369 break;431 break;
@@ -375,7 +437,7 @@ af::Status AnnotateSimdTemplateRoles(const af::AscNodePtr &indirect_load, const
375 }437 }
376 break;438 break;
377 }439 }
378- if (!ascgen_utils::IsNodeSupportsVectorFunction(node) && !af::ops::IsOps<af::ascir_op::Load>(node)) {440+ if (node->inputs.Size() != 1UL || node->outputs().size() != 1UL) {
379 break;441 break;
380 }442 }
381 GE_ASSERT_SUCCESS(443 GE_ASSERT_SUCCESS(
@@ -384,24 +446,6 @@ af::Status AnnotateSimdTemplateRoles(const af::AscNodePtr &indirect_load, const
384 return af::SUCCESS;446 return af::SUCCESS;
385}447}
386 448 
387-af::Status AnnotateSimtInputRoles(const NodePath &path, bool &is_candidate_legal) {
388- is_candidate_legal = false;
389- for (const auto &node : path) {
390- if (af::ops::IsOps<af::ascir_op::Data>(node)) {
391- is_candidate_legal = true;
392- return af::SUCCESS;
393- }
394- if (af::ops::IsOps<af::ascir_op::Load>(node)) {
395- GE_ASSERT_SUCCESS(ascgen_utils::indirect_load::SetTemplateRole(
396- node, ascgen_utils::indirect_load::TemplateRole::kSimtInputBoundary));
397- is_candidate_legal = true;
398- return af::SUCCESS;
399- }
400- return af::SUCCESS;
401- }
402- return af::SUCCESS;
403-}
404- 
405af::Status AnnotateSimtTransformRoles(const NodePath &path, const std::string &terminal_type,449af::Status AnnotateSimtTransformRoles(const NodePath &path, const std::string &terminal_type,
406 const std::string &direct_gm_boundary_type, bool &is_candidate_legal) {450 const std::string &direct_gm_boundary_type, bool &is_candidate_legal) {
407 is_candidate_legal = false;451 is_candidate_legal = false;
@@ -431,10 +475,6 @@ af::Status AnnotateSimtTemplateRoles(const af::AscNodePtr &indirect_load, const
431 is_candidate_legal = false;475 is_candidate_legal = false;
432 GE_ASSERT_SUCCESS(476 GE_ASSERT_SUCCESS(
433 ascgen_utils::indirect_load::SetTemplateRole(indirect_load, ascgen_utils::indirect_load::TemplateRole::kSimtOp));477 ascgen_utils::indirect_load::SetTemplateRole(indirect_load, ascgen_utils::indirect_load::TemplateRole::kSimtOp));
434- GE_ASSERT_SUCCESS(AnnotateSimtInputRoles(paths.data_input, is_candidate_legal));
435- if (!is_candidate_legal) {
436- return af::SUCCESS;
437- }
438 GE_ASSERT_SUCCESS(AnnotateSimtTransformRoles(paths.index_input, af::ascir_op::Data::Type, af::ascir_op::Load::Type,478 GE_ASSERT_SUCCESS(AnnotateSimtTransformRoles(paths.index_input, af::ascir_op::Data::Type, af::ascir_op::Load::Type,
439 is_candidate_legal));479 is_candidate_legal));
440 if (!is_candidate_legal) {480 if (!is_candidate_legal) {
@@ -452,8 +492,9 @@ af::Status ApplySimdGraphPass(af::AscGraph &graph, const af::AscNodePtr &indirec
452 if (!is_candidate_legal) {492 if (!is_candidate_legal) {
453 return af::SUCCESS;493 return af::SUCCESS;
454 }494 }
455- GE_ASSERT_SUCCESS(NormalizeSimdAxesForTemplate(graph, indirect_load, paths));
456 GE_ASSERT_SUCCESS(AnnotateSimdTemplateRoles(indirect_load, paths));495 GE_ASSERT_SUCCESS(AnnotateSimdTemplateRoles(indirect_load, paths));
496+ GE_ASSERT_SUCCESS(NormalizeSimdAxesForTemplate(graph, indirect_load, paths));
497+ GE_ASSERT_SUCCESS(::ascir::SetTemplateId(indirect_load, ascir::TemplateId::kIndirectLoadSimd));
457 return af::SUCCESS;498 return af::SUCCESS;
458}499}
459 500 
@@ -472,6 +513,7 @@ af::Status ApplySimtGraphPass(af::AscGraph &graph, const af::AscNodePtr &indirec
472 }513 }
473 GE_ASSERT_SUCCESS(::ascir::SetDcacheSize(indirect_load, kIndirectLoadSimtDcacheSize));514 GE_ASSERT_SUCCESS(::ascir::SetDcacheSize(indirect_load, kIndirectLoadSimtDcacheSize));
474 GE_ASSERT_SUCCESS(NormalizeSimtAxesForTemplate(graph, indirect_load));515 GE_ASSERT_SUCCESS(NormalizeSimtAxesForTemplate(graph, indirect_load));
516+ GE_ASSERT_SUCCESS(::ascir::SetTemplateId(indirect_load, ascir::TemplateId::kIndirectLoadSimt));
475 return af::SUCCESS;517 return af::SUCCESS;
476}518}
477 519 
Mautofuse/tests/framework/share_graph/include/share_graph.h+24-1
@@ -10,10 +10,28 @@
10 10 
11#ifndef INC_SHARE_GRAPH_H11#ifndef INC_SHARE_GRAPH_H
12#define INC_SHARE_GRAPH_H12#define INC_SHARE_GRAPH_H
13+#include <cstdint>
14+ 
13#include "graph/compute_graph.h"15#include "graph/compute_graph.h"
14#include "ascendc_ir.h"16#include "ascendc_ir.h"
15 17 
16namespace ascir {18namespace ascir {
19+enum class IndirectLoadInputPreType : uint32_t {
20+ kNone = 0U,
21+ kRelu,
22+ kExp2,
23+ kReluExp2,
24+ kReluRelu,
25+ kPrecisionCast,
26+};
27+ 
28+enum class IndirectLoadOutputPostType : uint32_t {
29+ kDefault = 0U,
30+ kModifiedBesselK0,
31+ kDirectStore,
32+ kCastExp2CastBack,
33+};
34+ 
17struct ShareGraph {35struct ShareGraph {
18 static af::ComputeGraphPtr LoadLog2StoreFusedGraph(size_t dims_size);36 static af::ComputeGraphPtr LoadLog2StoreFusedGraph(size_t dims_size);
19 static af::ComputeGraphPtr ModFusedGraph(size_t dims_size);37 static af::ComputeGraphPtr ModFusedGraph(size_t dims_size);
@@ -113,7 +131,12 @@ struct ShareGraph {
113 static af::ComputeGraphPtr LoadGatherTailAbsStore(int64_t gather_axis, af::DataType data_type);131 static af::ComputeGraphPtr LoadGatherTailAbsStore(int64_t gather_axis, af::DataType data_type);
114 static af::ComputeGraphPtr LoadGatherOneAxisAbsStore(int64_t gather_axis, af::DataType data_type);132 static af::ComputeGraphPtr LoadGatherOneAxisAbsStore(int64_t gather_axis, af::DataType data_type);
115 static af::ComputeGraphPtr IndirectLoadStoreFusedGraph(size_t rank, int64_t axis, af::DataType data_type,133 static af::ComputeGraphPtr IndirectLoadStoreFusedGraph(size_t rank, int64_t axis, af::DataType data_type,
116- bool has_input_pre, bool use_exp2);134+ af::DataType index_type,
135+ IndirectLoadInputPreType input_pre_type, bool use_exp2,
136+ IndirectLoadOutputPostType output_post_type,
137+ const std::vector<int64_t> &input_shape = {},
138+ const std::vector<int64_t> &output_shape = {},
139+ bool mixed_index_pre = false);
117 static af::ComputeGraphPtr MatMulFusedGraph(size_t dims_size);140 static af::ComputeGraphPtr MatMulFusedGraph(size_t dims_size);
118 static af::ComputeGraphPtr GatherReduceStore(int64_t gather_axis, af::DataType data_type);141 static af::ComputeGraphPtr GatherReduceStore(int64_t gather_axis, af::DataType data_type);
119 static af::ComputeGraphPtr LoadWhereReduceStoreFusedGraph(size_t dims_size, bool x2_scalar, bool x3_scalar);142 static af::ComputeGraphPtr LoadWhereReduceStoreFusedGraph(size_t dims_size, bool x2_scalar, bool x3_scalar);
Mautofuse/tests/framework/share_graph/src/share_graph.cc+204-25
@@ -8299,10 +8299,12 @@ void SetIndirectLoadNodeView(Op &op, const std::vector<int64_t> &axes, const std
8299}8299}
8300 8300 
8301static void CreateIndirectLoadAxes(af::AscGraph &graph, size_t rank, size_t offset, std::vector<int64_t> &axes,8301static void CreateIndirectLoadAxes(af::AscGraph &graph, size_t rank, size_t offset, std::vector<int64_t> &axes,
8302- std::vector<af::Expression> &repeats, std::vector<af::Expression> &strides) {8302+ std::vector<af::Expression> &repeats, std::vector<af::Expression> &strides,
8303+ const std::vector<int64_t> &shape) {
8303 for (size_t i = 0UL; i < rank; ++i) {8304 for (size_t i = 0UL; i < rank; ++i) {
8304 const size_t dim = offset + i;8305 const size_t dim = offset + i;
8305- repeats.emplace_back(graph.CreateSizeVar("s" + std::to_string(dim)));8306+ repeats.emplace_back(shape.empty() ? graph.CreateSizeVar("s" + std::to_string(dim))
8307+ : graph.CreateSizeVar(shape[i]));
8306 axes.emplace_back(graph.CreateAxis("z" + std::to_string(dim), repeats.back()).id);8308 axes.emplace_back(graph.CreateAxis("z" + std::to_string(dim), repeats.back()).id);
8307 }8309 }
8308 strides.assign(rank, af::ops::One);8310 strides.assign(rank, af::ops::One);
@@ -8311,10 +8313,30 @@ static void CreateIndirectLoadAxes(af::AscGraph &graph, size_t rank, size_t offs
8311 }8313 }
8312}8314}
8313 8315 
8316+static void AddIndirectLoadPrecisionCastDataChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8317+ const std::vector<af::Expression> &repeats,
8318+ const std::vector<af::Expression> &strides,
8319+ af::ascir_op::IndirectLoad &indirect_load) {
8320+ af::ascir_op::Data x("x");
8321+ graph.AddNode(x);
8322+ x.y.dtype = af::DT_FLOAT16;
8323+ x.ir_attr.SetIndex(0);
8324+ af::ascir_op::Load input_load("input_load");
8325+ graph.AddNode(input_load);
8326+ input_load.x = x.y;
8327+ SetIndirectLoadNodeView(input_load, axes, repeats, strides, af::DT_FLOAT16);
8328+ af::ascir_op::Cast input_cast("input_cast");
8329+ graph.AddNode(input_cast);
8330+ input_cast.x = input_load.y;
8331+ SetIndirectLoadNodeView(input_cast, axes, repeats, strides, af::DT_FLOAT);
8332+ indirect_load.x1 = input_cast.y;
8333+}
8334+ 
8314static void AddIndirectLoadDataChain(af::AscGraph &graph, const std::vector<int64_t> &axes,8335static void AddIndirectLoadDataChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8315 const std::vector<af::Expression> &repeats,8336 const std::vector<af::Expression> &repeats,
8316 const std::vector<af::Expression> &strides, af::DataType data_type,8337 const std::vector<af::Expression> &strides, af::DataType data_type,
8317- bool has_input_pre, af::ascir_op::IndirectLoad &indirect_load) {8338+ IndirectLoadInputPreType input_pre_type,
8339+ af::ascir_op::IndirectLoad &indirect_load) {
8318 af::ascir_op::Data x("x");8340 af::ascir_op::Data x("x");
8319 graph.AddNode(x);8341 graph.AddNode(x);
8320 x.y.dtype = data_type;8342 x.y.dtype = data_type;
@@ -8323,43 +8345,185 @@ static void AddIndirectLoadDataChain(af::AscGraph &graph, const std::vector<int6
8323 graph.AddNode(input_load);8345 graph.AddNode(input_load);
8324 input_load.x = x.y;8346 input_load.x = x.y;
8325 SetIndirectLoadNodeView(input_load, axes, repeats, strides, data_type);8347 SetIndirectLoadNodeView(input_load, axes, repeats, strides, data_type);
8326- if (!has_input_pre) {8348+ if (input_pre_type == IndirectLoadInputPreType::kNone) {
8327 indirect_load.x1 = input_load.y;8349 indirect_load.x1 = input_load.y;
8328 return;8350 return;
8329 }8351 }
8352+ if (input_pre_type == IndirectLoadInputPreType::kExp2) {
8353+ af::ascir_op::Exp2 input_exp2("input_exp2");
8354+ graph.AddNode(input_exp2);
8355+ input_exp2.x = input_load.y;
8356+ SetIndirectLoadNodeView(input_exp2, axes, repeats, strides, data_type);
8357+ indirect_load.x1 = input_exp2.y;
8358+ return;
8359+ }
8330 af::ascir_op::Relu input_relu("input_relu");8360 af::ascir_op::Relu input_relu("input_relu");
8331 graph.AddNode(input_relu);8361 graph.AddNode(input_relu);
8332 input_relu.x = input_load.y;8362 input_relu.x = input_load.y;
8333 SetIndirectLoadNodeView(input_relu, axes, repeats, strides, data_type);8363 SetIndirectLoadNodeView(input_relu, axes, repeats, strides, data_type);
8334- indirect_load.x1 = input_relu.y;8364+ if (input_pre_type == IndirectLoadInputPreType::kRelu) {
8365+ indirect_load.x1 = input_relu.y;
8366+ return;
8367+ }
8368+ if (input_pre_type == IndirectLoadInputPreType::kReluRelu) {
8369+ af::ascir_op::Relu input_relu2("input_relu2");
8370+ graph.AddNode(input_relu2);
8371+ input_relu2.x = input_relu.y;
8372+ SetIndirectLoadNodeView(input_relu2, axes, repeats, strides, data_type);
8373+ indirect_load.x1 = input_relu2.y;
8374+ return;
8375+ }
8376+ af::ascir_op::Exp2 input_exp2("input_exp2");
8377+ graph.AddNode(input_exp2);
8378+ input_exp2.x = input_relu.y;
8379+ SetIndirectLoadNodeView(input_exp2, axes, repeats, strides, data_type);
8380+ indirect_load.x1 = input_exp2.y;
8335}8381}
8336 8382 
8337static void AddIndirectLoadIndexChain(af::AscGraph &graph, const std::vector<int64_t> &axes,8383static void AddIndirectLoadIndexChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8338 const std::vector<af::Expression> &repeats,8384 const std::vector<af::Expression> &repeats,
8339- const std::vector<af::Expression> &strides,8385+ const std::vector<af::Expression> &strides, af::DataType index_type,
8340- af::ascir_op::IndirectLoad &indirect_load) {8386+ bool mixed_index_pre, af::ascir_op::IndirectLoad &indirect_load) {
8341 af::ascir_op::Data index("index");8387 af::ascir_op::Data index("index");
8342 graph.AddNode(index);8388 graph.AddNode(index);
8343- index.y.dtype = af::DT_INT32;8389+ index.y.dtype = index_type;
8344 index.ir_attr.SetIndex(1);8390 index.ir_attr.SetIndex(1);
8345 af::ascir_op::Load index_load("index_load");8391 af::ascir_op::Load index_load("index_load");
8346 graph.AddNode(index_load);8392 graph.AddNode(index_load);
8347 index_load.x = index.y;8393 index_load.x = index.y;
8348- SetIndirectLoadNodeView(index_load, axes, repeats, strides, af::DT_INT32);8394+ SetIndirectLoadNodeView(index_load, axes, repeats, strides, index_type);
8349 af::ascir_op::Abs index_abs("index_abs");8395 af::ascir_op::Abs index_abs("index_abs");
8350 graph.AddNode(index_abs);8396 graph.AddNode(index_abs);
8351 index_abs.x = index_load.y;8397 index_abs.x = index_load.y;
8352- SetIndirectLoadNodeView(index_abs, axes, repeats, strides, af::DT_INT32);8398+ SetIndirectLoadNodeView(index_abs, axes, repeats, strides, index_type);
8353- indirect_load.x2 = index_abs.y;8399+ if (!mixed_index_pre) {
8400+ indirect_load.x2 = index_abs.y;
8401+ return;
8402+ }
8403+ af::ascir_op::Abs index_abs2("index_abs2");
8404+ graph.AddNode(index_abs2);
8405+ index_abs2.x = index_abs.y;
8406+ SetIndirectLoadNodeView(index_abs2, axes, repeats, strides, index_type);
8407+ af::ascir_op::Cast index_cast("index_cast_float");
8408+ graph.AddNode(index_cast);
8409+ index_cast.x = index_abs2.y;
8410+ SetIndirectLoadNodeView(index_cast, axes, repeats, strides, af::DT_FLOAT);
8411+ af::ascir_op::Exp2 index_exp2("index_exp2");
8412+ graph.AddNode(index_exp2);
8413+ index_exp2.x = index_cast.y;
8414+ SetIndirectLoadNodeView(index_exp2, axes, repeats, strides, af::DT_FLOAT);
8415+ af::ascir_op::Log2 index_log2("index_log2");
8416+ graph.AddNode(index_log2);
8417+ index_log2.x = index_exp2.y;
8418+ SetIndirectLoadNodeView(index_log2, axes, repeats, strides, af::DT_FLOAT);
8419+ af::ascir_op::FloorToInt index_floor_to_int("index_floor_to_int");
8420+ graph.AddNode(index_floor_to_int);
8421+ index_floor_to_int.x = index_log2.y;
8422+ SetIndirectLoadNodeView(index_floor_to_int, axes, repeats, strides, index_type);
8423+ indirect_load.x2 = index_floor_to_int.y;
8424+}
8425+ 
8426+template <typename Op>
8427+static void AddIndirectLoadCustomOutputChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8428+ const std::vector<af::Expression> &repeats,
8429+ const std::vector<af::Expression> &strides, af::DataType data_type,
8430+ IndirectLoadOutputPostType output_post_type, Op &op) {
8431+ af::AscOpOutput *output = &op.y;
8432+ if (output_post_type == IndirectLoadOutputPostType::kModifiedBesselK0) {
8433+ af::ascir_op::ModifiedBesselK0 modified_bessel_k0("output_modified_bessel_k0");
8434+ graph.AddNode(modified_bessel_k0);
8435+ modified_bessel_k0.x = *output;
8436+ SetIndirectLoadNodeView(modified_bessel_k0, axes, repeats, strides, data_type);
8437+ output = &modified_bessel_k0.y;
8438+ }
8439+ af::ascir_op::Store store("store");
8440+ graph.AddNode(store);
8441+ store.x = *output;
8442+ SetIndirectLoadNodeView(store, axes, repeats, strides, data_type);
8443+ af::ascir_op::Output y("y");
8444+ graph.AddNode(y);
8445+ y.x = store.y;
8446+ SetIndirectLoadNodeView(y, axes, repeats, strides, data_type);
8447+ y.ir_attr.SetIndex(0);
8448+}
8449+ 
8450+static void AddIndirectLoadPrecisionOutputChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8451+ const std::vector<af::Expression> &repeats,
8452+ const std::vector<af::Expression> &strides, bool use_exp2,
8453+ af::ascir_op::IndirectLoad &indirect_load) {
8454+ af::ascir_op::Exp exp("output_exp");
8455+ af::ascir_op::Exp2 exp2("output_exp");
8456+ af::AscOpOutput *math_output = nullptr;
8457+ if (use_exp2) {
8458+ graph.AddNode(exp2);
8459+ exp2.x = indirect_load.y;
8460+ SetIndirectLoadNodeView(exp2, axes, repeats, strides, af::DT_FLOAT);
8461+ math_output = &exp2.y;
8462+ } else {
8463+ graph.AddNode(exp);
8464+ exp.x = indirect_load.y;
8465+ SetIndirectLoadNodeView(exp, axes, repeats, strides, af::DT_FLOAT);
8466+ math_output = &exp.y;
8467+ }
8468+ af::ascir_op::Cast output_cast("output_cast");
8469+ graph.AddNode(output_cast);
8470+ output_cast.x = *math_output;
8471+ SetIndirectLoadNodeView(output_cast, axes, repeats, strides, af::DT_FLOAT16);
8472+ AddIndirectLoadCustomOutputChain(graph, axes, repeats, strides, af::DT_FLOAT16,
8473+ IndirectLoadOutputPostType::kDirectStore, output_cast);
8474+}
8475+ 
8476+static void AddIndirectLoadUint32OutputChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8477+ const std::vector<af::Expression> &repeats,
8478+ const std::vector<af::Expression> &strides,
8479+ af::ascir_op::IndirectLoad &indirect_load) {
8480+ af::ascir_op::Cast cast_int32("output_cast_int32");
8481+ graph.AddNode(cast_int32);
8482+ cast_int32.x = indirect_load.y;
8483+ SetIndirectLoadNodeView(cast_int32, axes, repeats, strides, af::DT_INT32);
8484+ af::ascir_op::Cast cast_float("output_cast_float");
8485+ graph.AddNode(cast_float);
8486+ cast_float.x = cast_int32.y;
8487+ SetIndirectLoadNodeView(cast_float, axes, repeats, strides, af::DT_FLOAT);
8488+ af::ascir_op::Exp2 exp2("output_exp");
8489+ graph.AddNode(exp2);
8490+ exp2.x = cast_float.y;
8491+ SetIndirectLoadNodeView(exp2, axes, repeats, strides, af::DT_FLOAT);
8492+ af::ascir_op::Cast cast_back_int32("output_cast_back_int32");
8493+ graph.AddNode(cast_back_int32);
8494+ cast_back_int32.x = exp2.y;
8495+ SetIndirectLoadNodeView(cast_back_int32, axes, repeats, strides, af::DT_INT32);
8496+ af::ascir_op::Cast output_cast("output_cast");
8497+ graph.AddNode(output_cast);
8498+ output_cast.x = cast_back_int32.y;
8499+ SetIndirectLoadNodeView(output_cast, axes, repeats, strides, af::DT_UINT32);
8500+ AddIndirectLoadCustomOutputChain(graph, axes, repeats, strides, af::DT_UINT32,
8501+ IndirectLoadOutputPostType::kDirectStore, output_cast);
8354}8502}
8355 8503 
8356static void AddIndirectLoadOutputChain(af::AscGraph &graph, const std::vector<int64_t> &axes,8504static void AddIndirectLoadOutputChain(af::AscGraph &graph, const std::vector<int64_t> &axes,
8357 const std::vector<af::Expression> &repeats,8505 const std::vector<af::Expression> &repeats,
8358 const std::vector<af::Expression> &strides, af::DataType data_type,8506 const std::vector<af::Expression> &strides, af::DataType data_type,
8359- bool use_exp2, af::ascir_op::IndirectLoad &indirect_load) {8507+ IndirectLoadInputPreType input_pre_type, bool use_exp2,
8508+ IndirectLoadOutputPostType output_post_type,
8509+ af::ascir_op::IndirectLoad &indirect_load) {
8510+ if (input_pre_type == IndirectLoadInputPreType::kPrecisionCast) {
8511+ AddIndirectLoadPrecisionOutputChain(graph, axes, repeats, strides, use_exp2, indirect_load);
8512+ return;
8513+ }
8514+ if (output_post_type == IndirectLoadOutputPostType::kCastExp2CastBack) {
8515+ AddIndirectLoadUint32OutputChain(graph, axes, repeats, strides, indirect_load);
8516+ return;
8517+ }
8518+ if (output_post_type != IndirectLoadOutputPostType::kDefault) {
8519+ AddIndirectLoadCustomOutputChain(graph, axes, repeats, strides, data_type, output_post_type, indirect_load);
8520+ return;
8521+ }
8360 af::ascir_op::Neg output_neg("output_neg");8522 af::ascir_op::Neg output_neg("output_neg");
8361 graph.AddNode(output_neg);8523 graph.AddNode(output_neg);
8362- if (use_exp2) {8524+ if (input_pre_type == IndirectLoadInputPreType::kExp2 || input_pre_type == IndirectLoadInputPreType::kReluExp2) {
8525+ output_neg.x = indirect_load.y;
8526+ } else if (use_exp2) {
8363 af::ascir_op::Exp2 exp2("output_exp");8527 af::ascir_op::Exp2 exp2("output_exp");
8364 graph.AddNode(exp2);8528 graph.AddNode(exp2);
8365 exp2.x = indirect_load.y;8529 exp2.x = indirect_load.y;
@@ -8387,27 +8551,41 @@ static void AddIndirectLoadOutputChain(af::AscGraph &graph, const std::vector<in
8387}8551}
8388 8552 
8389static void IndirectLoadStore_BeforeAutofuse(af::AscGraph &graph, size_t rank, int64_t axis, af::DataType data_type,8553static void IndirectLoadStore_BeforeAutofuse(af::AscGraph &graph, size_t rank, int64_t axis, af::DataType data_type,
8390- bool has_input_pre, bool use_exp2) {8554+ af::DataType index_type, IndirectLoadInputPreType input_pre_type,
8555+ bool use_exp2, IndirectLoadOutputPostType output_post_type,
8556+ const std::vector<int64_t> &input_shape,
8557+ const std::vector<int64_t> &output_shape, bool mixed_index_pre) {
8391 std::vector<int64_t> input_axes;8558 std::vector<int64_t> input_axes;
8392 std::vector<af::Expression> input_repeats;8559 std::vector<af::Expression> input_repeats;
8393 std::vector<af::Expression> input_strides;8560 std::vector<af::Expression> input_strides;
8394- CreateIndirectLoadAxes(graph, rank, 0UL, input_axes, input_repeats, input_strides);8561+ CreateIndirectLoadAxes(graph, rank, 0UL, input_axes, input_repeats, input_strides, input_shape);
8395 std::vector<int64_t> output_axes;8562 std::vector<int64_t> output_axes;
8396 std::vector<af::Expression> output_repeats;8563 std::vector<af::Expression> output_repeats;
8397 std::vector<af::Expression> output_strides;8564 std::vector<af::Expression> output_strides;
8398- CreateIndirectLoadAxes(graph, rank, rank, output_axes, output_repeats, output_strides);8565+ CreateIndirectLoadAxes(graph, rank, rank, output_axes, output_repeats, output_strides, output_shape);
8399 af::ascir_op::IndirectLoad indirect_load("indirect_load");8566 af::ascir_op::IndirectLoad indirect_load("indirect_load");
8400 graph.AddNode(indirect_load);8567 graph.AddNode(indirect_load);
8401- AddIndirectLoadDataChain(graph, input_axes, input_repeats, input_strides, data_type, has_input_pre, indirect_load);8568+ if (input_pre_type == IndirectLoadInputPreType::kPrecisionCast) {
8402- AddIndirectLoadIndexChain(graph, output_axes, output_repeats, output_strides, indirect_load);8569+ AddIndirectLoadPrecisionCastDataChain(graph, input_axes, input_repeats, input_strides, indirect_load);
8570+ } else {
8571+ AddIndirectLoadDataChain(graph, input_axes, input_repeats, input_strides, data_type, input_pre_type, indirect_load);
8572+ }
8573+ AddIndirectLoadIndexChain(graph, output_axes, output_repeats, output_strides, index_type, mixed_index_pre,
8574+ indirect_load);
8403 indirect_load.ir_attr.SetAxis(axis);8575 indirect_load.ir_attr.SetAxis(axis);
8404- SetIndirectLoadNodeView(indirect_load, output_axes, output_repeats, output_strides, data_type);8576+ const af::DataType indirect_load_type =
8405- AddIndirectLoadOutputChain(graph, output_axes, output_repeats, output_strides, data_type, use_exp2, indirect_load);8577+ input_pre_type == IndirectLoadInputPreType::kPrecisionCast ? af::DT_FLOAT : data_type;
8578+ SetIndirectLoadNodeView(indirect_load, output_axes, output_repeats, output_strides, indirect_load_type);
8579+ AddIndirectLoadOutputChain(graph, output_axes, output_repeats, output_strides, data_type, input_pre_type, use_exp2,
8580+ output_post_type, indirect_load);
8406}8581}
8407 8582 
8408-af::ComputeGraphPtr ShareGraph::IndirectLoadStoreFusedGraph(size_t rank, int64_t axis, af::DataType data_type,8583+af::ComputeGraphPtr ShareGraph::IndirectLoadStoreFusedGraph(
8409- bool has_input_pre, bool use_exp2) {8584+ size_t rank, int64_t axis, af::DataType data_type, af::DataType index_type, IndirectLoadInputPreType input_pre_type,
8410- if (rank < 2UL || rank > 4UL || axis < -static_cast<int64_t>(rank) || axis >= static_cast<int64_t>(rank)) {8585+ bool use_exp2, IndirectLoadOutputPostType output_post_type, const std::vector<int64_t> &input_shape,
8586+ const std::vector<int64_t> &output_shape, bool mixed_index_pre) {
8587+ if (rank < 2UL || rank > 4UL || axis < -static_cast<int64_t>(rank) || axis >= static_cast<int64_t>(rank) ||
8588+ (!input_shape.empty() && input_shape.size() != rank) || (!output_shape.empty() && output_shape.size() != rank)) {
8411 return nullptr;8589 return nullptr;
8412 }8590 }
8413 af::AscGraph fused_graph("indirect_load_store_test");8591 af::AscGraph fused_graph("indirect_load_store_test");
@@ -8423,7 +8601,7 @@ af::ComputeGraphPtr ShareGraph::IndirectLoadStoreFusedGraph(size_t rank, int64_t
8423 auto data_desc = std::make_shared<GeTensorDesc>();8601 auto data_desc = std::make_shared<GeTensorDesc>();
8424 data_desc->SetDataType(data_type);8602 data_desc->SetDataType(data_type);
8425 auto index_desc = std::make_shared<GeTensorDesc>();8603 auto index_desc = std::make_shared<GeTensorDesc>();
8426- index_desc->SetDataType(af::DT_INT32);8604+ index_desc->SetDataType(index_type);
8427 auto backend_desc = std::make_shared<OpDesc>("asc_backend", "AscBackend");8605 auto backend_desc = std::make_shared<OpDesc>("asc_backend", "AscBackend");
8428 backend_desc->AddInputDesc(data_desc->Clone());8606 backend_desc->AddInputDesc(data_desc->Clone());
8429 backend_desc->AddInputDesc(index_desc->Clone());8607 backend_desc->AddInputDesc(index_desc->Clone());
@@ -8434,7 +8612,8 @@ af::ComputeGraphPtr ShareGraph::IndirectLoadStoreFusedGraph(size_t rank, int64_t
8434 }8612 }
8435 8613 
8436 auto sub_graph = std::make_shared<af::AscGraph>("indirect_load_store_test");8614 auto sub_graph = std::make_shared<af::AscGraph>("indirect_load_store_test");
8437- IndirectLoadStore_BeforeAutofuse(*sub_graph, rank, axis, data_type, has_input_pre, use_exp2);8615+ IndirectLoadStore_BeforeAutofuse(*sub_graph, rank, axis, data_type, index_type, input_pre_type, use_exp2,
8616+ output_post_type, input_shape, output_shape, mixed_index_pre);
8438 auto fuse_attrs = backend->GetOpDesc()->GetOrCreateAttrsGroup<af::AutoFuseAttrs>();8617 auto fuse_attrs = backend->GetOpDesc()->GetOrCreateAttrsGroup<af::AutoFuseAttrs>();
8439 if (fuse_attrs == nullptr) {8618 if (fuse_attrs == nullptr) {
8440 return nullptr;8619 return nullptr;
Mautofuse/tests/ut/optimize/task_generator/test_indirect_load_schedule_case_generator.cpp+27-24
@@ -297,26 +297,6 @@ TEST(IndirectLoadScheduleCaseGeneratorTest, DoesNotStoreFixedTileAxesAsTemplateM
297 af::kIdNone);297 af::kIdNone);
298}298}
299 299 
300-TEST(IndirectLoadScheduleCaseGeneratorTest, StoresInputVectorizedAxisAsTemplateMetadata) {
301- auto graph = BuildIndirectLoadGraph(2, true);
302- optimize::IndirectLoadScheduleCaseGenerator generator;
303- std::vector<af::AscGraph> graphs;
304- std::vector<std::string> score_functions;
305- ASSERT_EQ(generator.Generate(graph, graphs, score_functions), af::SUCCESS);
306- auto &simd_graph = FindGeneratedGraphByTemplate(graphs, ascir::TemplateId::kIndirectLoadSimd);
307- const auto indirect_load = simd_graph.FindNode("indirect_load");
308- const auto pre_abs = simd_graph.FindNode("pre_abs");
309- ASSERT_NE(indirect_load, nullptr);
310- ASSERT_NE(pre_abs, nullptr);
311- 
312- ascgen_utils::indirect_load::TemplateAxes axes;
313- ASSERT_EQ(ascgen_utils::indirect_load::GetTemplateAxes(indirect_load, axes), af::SUCCESS);
314- ExpectAxisNames(simd_graph, pre_abs->attr.sched.axis, {"z0", "z1", "z2", "z3"});
315- ExpectAxisNames(simd_graph, pre_abs->outputs()[0]->attr.axis, {"z0", "z1", "z2", "z3"});
316- EXPECT_TRUE(pre_abs->outputs()[0]->attr.vectorized_axis.empty());
317- ExpectAxisNames(simd_graph, {axes.input_inner_axis}, {"indirect_load_input_inner"});
318-}
319- 
320TEST(IndirectLoadScheduleCaseGeneratorTest, GeneratedSimdCandidateKeepsPublicBehavior) {300TEST(IndirectLoadScheduleCaseGeneratorTest, GeneratedSimdCandidateKeepsPublicBehavior) {
321 auto graph = BuildIndirectLoadGraph(2, true);301 auto graph = BuildIndirectLoadGraph(2, true);
322 optimize::IndirectLoadScheduleCaseGenerator generator;302 optimize::IndirectLoadScheduleCaseGenerator generator;
@@ -332,7 +312,6 @@ TEST(IndirectLoadScheduleCaseGeneratorTest, GeneratedSimdCandidateKeepsPublicBeh
332 EXPECT_FALSE(il_behavior.skips_api_emit);312 EXPECT_FALSE(il_behavior.skips_api_emit);
333 EXPECT_FALSE(il_behavior.uses_direct_gm_pipeline);313 EXPECT_FALSE(il_behavior.uses_direct_gm_pipeline);
334 EXPECT_FALSE(il_behavior.skips_ub_lifecycle);314 EXPECT_FALSE(il_behavior.skips_ub_lifecycle);
335- EXPECT_FALSE(il_behavior.skips_ub_expr);
336 EXPECT_FALSE(il_behavior.preserves_vectorized_axis);315 EXPECT_FALSE(il_behavior.preserves_vectorized_axis);
337 EXPECT_FALSE(ascgen_utils::indirect_load::ShouldDisableRegularVectorFunc(indirect_load));316 EXPECT_FALSE(ascgen_utils::indirect_load::ShouldDisableRegularVectorFunc(indirect_load));
338}317}
@@ -348,7 +327,6 @@ TEST(IndirectLoadScheduleCaseGeneratorTest, GeneratedSimtCandidateKeepsPublicBeh
348 EXPECT_FALSE(behavior.skips_api_emit);327 EXPECT_FALSE(behavior.skips_api_emit);
349 EXPECT_TRUE(behavior.uses_direct_gm_pipeline);328 EXPECT_TRUE(behavior.uses_direct_gm_pipeline);
350 EXPECT_TRUE(behavior.skips_ub_lifecycle);329 EXPECT_TRUE(behavior.skips_ub_lifecycle);
351- EXPECT_TRUE(behavior.skips_ub_expr);
352 EXPECT_TRUE(behavior.preserves_vectorized_axis);330 EXPECT_TRUE(behavior.preserves_vectorized_axis);
353 EXPECT_FALSE(ascgen_utils::indirect_load::ShouldApplyInputInnerVectorization(indirect_load));331 EXPECT_FALSE(ascgen_utils::indirect_load::ShouldApplyInputInnerVectorization(indirect_load));
354 EXPECT_FALSE(ascgen_utils::indirect_load::ShouldSkipMainScheduleTiling(indirect_load));332 EXPECT_FALSE(ascgen_utils::indirect_load::ShouldSkipMainScheduleTiling(indirect_load));
@@ -356,7 +334,26 @@ TEST(IndirectLoadScheduleCaseGeneratorTest, GeneratedSimtCandidateKeepsPublicBeh
356 EXPECT_TRUE(ascgen_utils::indirect_load::ShouldDisableRegularVectorFunc(indirect_load));334 EXPECT_TRUE(ascgen_utils::indirect_load::ShouldDisableRegularVectorFunc(indirect_load));
357}335}
358 336 
359-TEST(IndirectLoadScheduleCaseGeneratorTest, SimtMovesInputPrecisionCastAfterIndirectLoad) {337+TEST(IndirectLoadScheduleCaseGeneratorTest, SimdMovesInputPreAfterIndirectLoad) {
338+ auto graph = BuildIndirectLoadGraph(2, true);
339+ optimize::IndirectLoadScheduleCaseGenerator generator;
340+ std::vector<af::AscGraph> graphs;
341+ std::vector<std::string> score_functions;
342+ ASSERT_EQ(generator.Generate(graph, graphs, score_functions), af::SUCCESS);
343+ auto &simd_graph = FindGeneratedGraphByTemplate(graphs, ascir::TemplateId::kIndirectLoadSimd);
344+ const auto indirect_load = simd_graph.FindNode("indirect_load");
345+ const auto pre_abs = simd_graph.FindNode("pre_abs");
346+ ASSERT_NE(indirect_load, nullptr);
347+ ASSERT_NE(pre_abs, nullptr);
348+ 
349+ EXPECT_EQ(ascgen_utils::indirect_load::GetInputProducer(indirect_load, 0UL)->GetName(), "x");
350+ EXPECT_EQ(ascgen_utils::indirect_load::GetInputProducer(pre_abs, 0UL), indirect_load);
351+ EXPECT_EQ(ascgen_utils::indirect_load::GetOnlyOutputConsumer(indirect_load), pre_abs);
352+ ExpectAxisNames(simd_graph, pre_abs->attr.sched.axis, {"z4", "z5", "z6", "z7"});
353+ ExpectAxisNames(simd_graph, pre_abs->outputs()[0]->attr.axis, {"z4", "z5", "z6", "z7"});
354+}
355+ 
356+TEST(IndirectLoadScheduleCaseGeneratorTest, MovesInputPrecisionCastAfterIndirectLoad) {
360 auto graph = BuildIndirectLoadPrecisionCastGraph();357 auto graph = BuildIndirectLoadPrecisionCastGraph();
361 optimize::IndirectLoadScheduleCaseGenerator generator;358 optimize::IndirectLoadScheduleCaseGenerator generator;
362 std::vector<af::AscGraph> graphs;359 std::vector<af::AscGraph> graphs;
@@ -372,8 +369,14 @@ TEST(IndirectLoadScheduleCaseGeneratorTest, SimtMovesInputPrecisionCastAfterIndi
372 EXPECT_NE(simd_graph.FindNode("output_cast"), nullptr);369 EXPECT_NE(simd_graph.FindNode("output_cast"), nullptr);
373 ASSERT_NE(simd_graph.FindNode("indirect_load"), nullptr);370 ASSERT_NE(simd_graph.FindNode("indirect_load"), nullptr);
374 ASSERT_NE(simd_graph.FindNode("output_exp"), nullptr);371 ASSERT_NE(simd_graph.FindNode("output_exp"), nullptr);
375- EXPECT_EQ(simd_graph.FindNode("indirect_load")->outputs()[0]->attr.dtype, af::DT_FLOAT);372+ EXPECT_EQ(simd_graph.FindNode("indirect_load")->outputs()[0]->attr.dtype, af::DT_FLOAT16);
376 EXPECT_EQ(simd_graph.FindNode("output_exp")->outputs()[0]->attr.dtype, af::DT_FLOAT);373 EXPECT_EQ(simd_graph.FindNode("output_exp")->outputs()[0]->attr.dtype, af::DT_FLOAT);
374+ EXPECT_EQ(ascgen_utils::indirect_load::GetInputProducer(simd_graph.FindNode("indirect_load"), 0UL)->GetName(),
375+ "input_load");
376+ EXPECT_EQ(ascgen_utils::indirect_load::GetInputProducer(simd_graph.FindNode("input_cast"), 0UL),
377+ simd_graph.FindNode("indirect_load"));
378+ EXPECT_EQ(ascgen_utils::indirect_load::GetOnlyOutputConsumer(simd_graph.FindNode("input_cast"))->GetName(),
379+ "output_exp");
377 380 
378 auto &simt_graph = FindGeneratedGraphByTemplate(graphs, ascir::TemplateId::kIndirectLoadSimt);381 auto &simt_graph = FindGeneratedGraphByTemplate(graphs, ascir::TemplateId::kIndirectLoadSimt);
379 const auto input_cast = simt_graph.FindNode("input_cast");382 const auto input_cast = simt_graph.FindNode("input_cast");
Mautofuse/tests/v35/st/backend_e2e_v2/indirect_load_store_test/CMakeLists.txt+88-12
@@ -1,4 +1,4 @@
1-function(add_indirect_load_e2e_case test_name rank axis use_exp2 tiling_key x0 x1 x2 x3 index0 index1 index2 index3)1+function(add_indirect_load_e2e_case test_name rank axis input_pre_type use_exp2 tiling_key x0 x1 x2 x3 index0 index1 index2 index3)
2 set(case_workdir ${CMAKE_CURRENT_BINARY_DIR}/${test_name})2 set(case_workdir ${CMAKE_CURRENT_BINARY_DIR}/${test_name})
3 file(MAKE_DIRECTORY ${case_workdir})3 file(MAKE_DIRECTORY ${case_workdir})
4 do_backend_e2e_st_test(${test_name}4 do_backend_e2e_st_test(${test_name}
@@ -14,7 +14,7 @@ function(add_indirect_load_e2e_case test_name rank axis use_exp2 tiling_key x0 x
14 set(case_definitions14 set(case_definitions
15 IL_RANK=${rank}15 IL_RANK=${rank}
16 IL_AXIS=${axis}16 IL_AXIS=${axis}
17- IL_HAS_INPUT_PRE=117+ IL_INPUT_PRE_TYPE=${input_pre_type}
18 IL_USE_EXP2=${use_exp2}18 IL_USE_EXP2=${use_exp2}
19 IL_TILING_KEY=${tiling_key}19 IL_TILING_KEY=${tiling_key}
20 IL_X_S0=${x0}20 IL_X_S0=${x0}
@@ -29,13 +29,89 @@ function(add_indirect_load_e2e_case test_name rank axis use_exp2 tiling_key x0 x
29 target_compile_definitions(${test_name}_e2e_v2 PRIVATE ${case_definitions})29 target_compile_definitions(${test_name}_e2e_v2 PRIVATE ${case_definitions})
30endfunction()30endfunction()
31 31 
32-add_indirect_load_e2e_case(indirect_load_rank2_axis1_simt 2 1 0 1 5 8 1 1 4 6 1 1)32+function(mark_indirect_load_static_shape test_name)
33-add_indirect_load_e2e_case(indirect_load_rank2_axis1_simd 2 1 1 0 5 8 1 1 4 6 1 1)33+ target_compile_definitions(${test_name}_codegen_v2 PRIVATE IL_STATIC_SHAPE)
34-add_indirect_load_e2e_case(indirect_load_rank2_axis0_simt 2 0 0 1 4 9 1 1 6 8 1 1)34+ target_compile_definitions(${test_name}_e2e_v2 PRIVATE IL_STATIC_SHAPE)
35-add_indirect_load_e2e_case(indirect_load_rank2_axis0_simd 2 0 1 0 4 9 1 1 6 8 1 1)35+endfunction()
36-add_indirect_load_e2e_case(indirect_load_rank2_axisneg2_simt 2 -2 0 1 4 9 1 1 4 8 1 1)36+ 
37-add_indirect_load_e2e_case(indirect_load_rank2_axisneg2_simd 2 -2 1 0 4 9 1 1 4 8 1 1)37+function(mark_indirect_load_wide_types test_name)
38-add_indirect_load_e2e_case(indirect_load_rank3_axis1_simt 3 1 0 1 3 5 8 1 2 7 6 1)38+ target_compile_definitions(${test_name}_codegen_v2 PRIVATE IL_DATA_FLOAT IL_INDEX_INT64)
39-add_indirect_load_e2e_case(indirect_load_rank3_axis1_simd 3 1 1 0 3 5 8 1 2 7 6 1)39+ target_compile_definitions(${test_name}_e2e_v2 PRIVATE IL_DATA_FLOAT IL_INDEX_INT64)
40-add_indirect_load_e2e_case(indirect_load_rank4_axis2_simt 4 2 0 1 3 3 5 8 2 2 5 6)40+endfunction()
41-add_indirect_load_e2e_case(indirect_load_rank4_axis2_simd 4 2 1 0 3 3 5 8 2 2 5 6)41+ 
42+function(mark_indirect_load_large_simt test_name)
43+ target_compile_definitions(${test_name}_e2e_v2 PRIVATE IL_EXPECT_SIMT_MULTI_ROUND)
44+endfunction()
45+ 
46+function(mark_indirect_load_input_pre_vf test_name expected)
47+ target_compile_definitions(${test_name}_codegen_v2 PRIVATE IL_EXPECT_INPUT_PRE_VF=${expected})
48+endfunction()
49+ 
50+function(mark_indirect_load_mixed_index_pre test_name)
51+ target_compile_definitions(${test_name}_codegen_v2 PRIVATE IL_MIXED_INDEX_PRE)
52+ target_compile_definitions(${test_name}_e2e_v2 PRIVATE IL_MIXED_INDEX_PRE)
53+endfunction()
54+ 
55+function(mark_indirect_load_output_post test_name output_post_type expect_simt)
56+ target_compile_definitions(${test_name}_codegen_v2 PRIVATE
57+ IL_OUTPUT_POST_TYPE=${output_post_type}
58+ IL_EXPECT_SIMT=${expect_simt})
59+ target_compile_definitions(${test_name}_e2e_v2 PRIVATE
60+ IL_OUTPUT_POST_TYPE=${output_post_type}
61+ IL_EXPECT_SIMT=${expect_simt})
62+endfunction()
63+ 
64+add_indirect_load_e2e_case(indirect_load_rank2_axis1_simt 2 1 1 0 1 5 8 1 1 4 6 1 1)
65+add_indirect_load_e2e_case(indirect_load_rank2_axis1_simd 2 1 1 1 0 5 8 1 1 4 6 1 1)
66+mark_indirect_load_static_shape(indirect_load_rank2_axis1_simd)
67+add_indirect_load_e2e_case(indirect_load_rank2_axisneg2_simt 2 -2 1 0 1 4 9 1 1 4 8 1 1)
68+add_indirect_load_e2e_case(indirect_load_rank2_axisneg2_simd 2 -2 1 1 0 4 9 1 1 4 8 1 1)
69+add_indirect_load_e2e_case(indirect_load_rank3_axis1_simt 3 1 1 0 1 3 5 8 1 2 7 6 1)
70+add_indirect_load_e2e_case(indirect_load_rank3_axis1_simd 3 1 1 1 0 3 5 8 1 2 7 6 1)
71+add_indirect_load_e2e_case(indirect_load_rank4_axis2_simt 4 2 1 0 1 3 3 5 8 2 2 5 6)
72+add_indirect_load_e2e_case(indirect_load_rank4_axis2_simd 4 2 1 1 0 3 3 5 8 2 2 5 6)
73+add_indirect_load_e2e_case(indirect_load_rank2_axis1_relu_exp2_pre_simd 2 1 3 0 0 5 8 1 1 4 6 1 1)
74+ 
75+# Static output expansion keeps both data and index preprocessing before IndirectLoad SIMD.
76+add_indirect_load_e2e_case(indirect_load_rank2_axis0_relu_pre_expanded_simd 2 0 1 1 0 4 9 1 1 6 8 1 1)
77+add_indirect_load_e2e_case(indirect_load_rank2_axis0_relu_relu_pre_expanded_simd 2 0 4 1 0 4 9 1 1 6 8 1 1)
78+add_indirect_load_e2e_case(indirect_load_rank3_axis2_exp2_pre_expanded_simd 3 2 2 0 0 4 5 6 1 3 4 11 1)
79+add_indirect_load_e2e_case(indirect_load_rank4_axis1_relu_exp2_pre_expanded_simd 4 1 3 0 0 2 3 4 5 2 7 4 3)
80+add_indirect_load_e2e_case(indirect_load_rank3_axis2_mixed_vf_pre_expanded_simd 3 2 3 0 0 4 5 6 1 3 4 11 1)
81+mark_indirect_load_static_shape(indirect_load_rank2_axis0_relu_pre_expanded_simd)
82+mark_indirect_load_static_shape(indirect_load_rank2_axis0_relu_relu_pre_expanded_simd)
83+mark_indirect_load_static_shape(indirect_load_rank3_axis2_exp2_pre_expanded_simd)
84+mark_indirect_load_static_shape(indirect_load_rank4_axis1_relu_exp2_pre_expanded_simd)
85+mark_indirect_load_static_shape(indirect_load_rank3_axis2_mixed_vf_pre_expanded_simd)
86+mark_indirect_load_mixed_index_pre(indirect_load_rank3_axis2_mixed_vf_pre_expanded_simd)
87+target_compile_definitions(indirect_load_rank2_axis0_relu_pre_expanded_simd_codegen_v2 PRIVATE IL_DATA_FLOAT)
88+target_compile_definitions(indirect_load_rank2_axis0_relu_pre_expanded_simd_e2e_v2 PRIVATE IL_DATA_FLOAT)
89+target_compile_definitions(indirect_load_rank2_axis0_relu_relu_pre_expanded_simd_codegen_v2 PRIVATE IL_DATA_FLOAT)
90+target_compile_definitions(indirect_load_rank2_axis0_relu_relu_pre_expanded_simd_e2e_v2 PRIVATE IL_DATA_FLOAT)
91+mark_indirect_load_input_pre_vf(indirect_load_rank2_axis0_relu_pre_expanded_simd 0)
92+mark_indirect_load_input_pre_vf(indirect_load_rank2_axis0_relu_relu_pre_expanded_simd 1)
93+ 
94+add_indirect_load_e2e_case(indirect_load_rank2_axis1_large_simt 2 1 1 0 1 2 64 1 1 2 32771 1 1)
95+mark_indirect_load_large_simt(indirect_load_rank2_axis1_large_simt)
96+add_indirect_load_e2e_case(indirect_load_rank2_axis1_float_int64_simd 2 1 0 1 0 4 8 1 1 3 6 1 1)
97+mark_indirect_load_wide_types(indirect_load_rank2_axis1_float_int64_simd)
98+ 
99+add_indirect_load_e2e_case(indirect_load_rank2_axis1_modified_bessel_k0_simd_only 2 1 0 0 0 5 8 1 1 4 6 1 1)
100+mark_indirect_load_output_post(indirect_load_rank2_axis1_modified_bessel_k0_simd_only 1 0)
101+target_compile_definitions(indirect_load_rank2_axis1_modified_bessel_k0_simd_only_codegen_v2 PRIVATE IL_DATA_FLOAT)
102+target_compile_definitions(indirect_load_rank2_axis1_modified_bessel_k0_simd_only_e2e_v2 PRIVATE IL_DATA_FLOAT)
103+target_include_directories(indirect_load_rank2_axis1_modified_bessel_k0_simd_only_e2e_v2 PRIVATE
104+ ${CODE_ROOT_DIR}/v35/ascendc/api_regbase)
105+ 
106+add_indirect_load_e2e_case(indirect_load_rank2_axis1_precision_cast_simt 2 1 5 0 1 5 8 1 1 4 6 1 1)
107+add_indirect_load_e2e_case(indirect_load_rank2_axis1_precision_cast_simd 2 1 5 1 0 5 8 1 1 4 6 1 1)
108+ 
109+add_indirect_load_e2e_case(indirect_load_rank2_axis1_bf16_int64_simt 2 1 0 0 1 5 8 1 1 4 6 1 1)
110+mark_indirect_load_output_post(indirect_load_rank2_axis1_bf16_int64_simt 2 1)
111+target_compile_definitions(indirect_load_rank2_axis1_bf16_int64_simt_codegen_v2 PRIVATE IL_DATA_BF16 IL_INDEX_INT64)
112+target_compile_definitions(indirect_load_rank2_axis1_bf16_int64_simt_e2e_v2 PRIVATE IL_DATA_BF16 IL_INDEX_INT64)
113+ 
114+add_indirect_load_e2e_case(indirect_load_rank2_axis1_uint32_int32_simd 2 1 0 1 0 5 8 1 1 4 6 1 1)
115+mark_indirect_load_output_post(indirect_load_rank2_axis1_uint32_int32_simd 3 1)
116+target_compile_definitions(indirect_load_rank2_axis1_uint32_int32_simd_codegen_v2 PRIVATE IL_DATA_UINT32)
117+target_compile_definitions(indirect_load_rank2_axis1_uint32_int32_simd_e2e_v2 PRIVATE IL_DATA_UINT32)
Mautofuse/tests/v35/st/backend_e2e_v2/indirect_load_store_test/indirect_load_store_backend_generator.cpp+135-13
@@ -29,8 +29,56 @@
29namespace {29namespace {
30constexpr size_t kRank = IL_RANK;30constexpr size_t kRank = IL_RANK;
31constexpr int64_t kAxis = IL_AXIS;31constexpr int64_t kAxis = IL_AXIS;
32-constexpr bool kHasInputPre = IL_HAS_INPUT_PRE;32+constexpr auto kInputPreType = static_cast<ascir::IndirectLoadInputPreType>(IL_INPUT_PRE_TYPE);
33+constexpr bool kHasInputPre = kInputPreType != ascir::IndirectLoadInputPreType::kNone;
33constexpr bool kUseExp2 = IL_USE_EXP2;34constexpr bool kUseExp2 = IL_USE_EXP2;
35+#ifdef IL_OUTPUT_POST_TYPE
36+constexpr auto kOutputPostType = static_cast<ascir::IndirectLoadOutputPostType>(IL_OUTPUT_POST_TYPE);
37+#else
38+constexpr auto kOutputPostType = ascir::IndirectLoadOutputPostType::kDefault;
39+#endif
40+#ifdef IL_EXPECT_SIMT
41+constexpr bool kExpectSimt = IL_EXPECT_SIMT;
42+#else
43+constexpr bool kExpectSimt = true;
44+#endif
45+#if defined(IL_DATA_BF16)
46+constexpr af::DataType kDataType = af::DT_BF16;
47+#elif defined(IL_DATA_UINT32)
48+constexpr af::DataType kDataType = af::DT_UINT32;
49+#elif defined(IL_DATA_FLOAT)
50+constexpr af::DataType kDataType = af::DT_FLOAT;
51+#else
52+constexpr af::DataType kDataType = af::DT_FLOAT16;
53+#endif
54+#ifdef IL_INDEX_INT64
55+constexpr af::DataType kIndexType = af::DT_INT64;
56+#else
57+constexpr af::DataType kIndexType = af::DT_INT32;
58+#endif
59+#ifdef IL_STATIC_SHAPE
60+constexpr bool kStaticShape = true;
61+#else
62+constexpr bool kStaticShape = false;
63+#endif
64+constexpr int64_t kInputNumel = IL_X_S0 * IL_X_S1 * IL_X_S2 * IL_X_S3;
65+constexpr int64_t kOutputNumel = IL_INDEX_S0 * IL_INDEX_S1 * IL_INDEX_S2 * IL_INDEX_S3;
66+constexpr bool kExpectInputPreMoved = kHasInputPre && (!kStaticShape || kOutputNumel <= kInputNumel);
67+#ifdef IL_MIXED_INDEX_PRE
68+constexpr bool kMixedIndexPre = true;
69+#else
70+constexpr bool kMixedIndexPre = false;
71+#endif
72+ 
73+std::vector<int64_t> GetStaticShape(bool input) {
74+ if (!kStaticShape) {
75+ return {};
76+ }
77+ std::vector<int64_t> shape = input ? std::vector<int64_t>{IL_X_S0, IL_X_S1, IL_X_S2, IL_X_S3}
78+ : std::vector<int64_t>{IL_INDEX_S0, IL_INDEX_S1, IL_INDEX_S2, IL_INDEX_S3};
79+ shape.resize(kRank);
80+ return shape;
81+}
34 82 
35const af::Axis *FindDerivedAxis(const af::AscGraph &graph, af::Axis::Type type, af::AxisId from) {83const af::Axis *FindDerivedAxis(const af::AscGraph &graph, af::Axis::Type type, af::AxisId from) {
36 for (const auto &axis : graph.GetAllAxis()) {84 for (const auto &axis : graph.GetAllAxis()) {
@@ -59,6 +107,31 @@ void ExpectAxisOrigins(af::AscGraph &graph, af::AxisId axis_id, const std::vecto
59 EXPECT_EQ(origins, expected);107 EXPECT_EQ(origins, expected);
60}108}
61 109 
110+void ExpectMixedInputPre(const af::AscNodePtr &indirect_load) {
111+ if (!kMixedIndexPre) {
112+ return;
113+ }
114+ const af::AscNodePtr data_post = ascgen_utils::indirect_load::GetInputProducer(indirect_load, 0UL);
115+ ASSERT_NE(data_post, nullptr);
116+ EXPECT_EQ(data_post->GetName(), "input_exp2");
117+ const af::AscNodePtr data_vf = ascgen_utils::indirect_load::GetInputProducer(data_post, 0UL);
118+ ASSERT_NE(data_vf, nullptr);
119+ EXPECT_TRUE(af::ops::IsOps<af::ascir_op::VectorFunc>(data_vf));
120+ 
121+ const af::AscNodePtr index_post = ascgen_utils::indirect_load::GetInputProducer(indirect_load, 1UL);
122+ ASSERT_NE(index_post, nullptr);
123+ EXPECT_EQ(index_post->GetName(), "index_floor_to_int");
124+ const af::AscNodePtr index_log2 = ascgen_utils::indirect_load::GetInputProducer(index_post, 0UL);
125+ ASSERT_NE(index_log2, nullptr);
126+ EXPECT_EQ(index_log2->GetName(), "index_log2");
127+ const af::AscNodePtr index_exp2 = ascgen_utils::indirect_load::GetInputProducer(index_log2, 0UL);
128+ ASSERT_NE(index_exp2, nullptr);
129+ EXPECT_EQ(index_exp2->GetName(), "index_exp2");
130+ const af::AscNodePtr index_vf = ascgen_utils::indirect_load::GetInputProducer(index_exp2, 0UL);
131+ ASSERT_NE(index_vf, nullptr);
132+ EXPECT_TRUE(af::ops::IsOps<af::ascir_op::VectorFunc>(index_vf));
133+}
134+ 
62void ExpectLoopFramework(af::AscGraph &graph, const af::AscNodePtr &indirect_load) {135void ExpectLoopFramework(af::AscGraph &graph, const af::AscNodePtr &indirect_load) {
63 ascgen_utils::indirect_load::TemplateAxes axes;136 ascgen_utils::indirect_load::TemplateAxes axes;
64 ascgen_utils::indirect_load::TemplateLogicalView logical_view;137 ascgen_utils::indirect_load::TemplateLogicalView logical_view;
@@ -107,9 +180,33 @@ void ExpectLoopFramework(af::AscGraph &graph, const af::AscNodePtr &indirect_loa
107 graph, axes.input_inner_axis,180 graph, axes.input_inner_axis,
108 std::vector<af::AxisId>(logical_view.data.axis_ids.begin() + split, logical_view.data.axis_ids.end()));181 std::vector<af::AxisId>(logical_view.data.axis_ids.begin() + split, logical_view.data.axis_ids.end()));
109 }182 }
183+ 
184+ const auto input_producer = ascgen_utils::indirect_load::GetInputProducer(indirect_load, 0UL);
185+ ASSERT_NE(input_producer, nullptr);
186+ if (kExpectInputPreMoved) {
187+ EXPECT_EQ(input_producer->GetName(), "input_load");
188+ } else if (kHasInputPre) {
189+#ifdef IL_EXPECT_INPUT_PRE_VF
190+ EXPECT_EQ(af::ops::IsOps<af::ascir_op::VectorFunc>(input_producer), IL_EXPECT_INPUT_PRE_VF);
191+#endif
192+ if (af::ops::IsOps<af::ascir_op::VectorFunc>(input_producer)) {
193+ EXPECT_EQ(ascgen_utils::indirect_load::GetTemplateRole(input_producer),
194+ ascgen_utils::indirect_load::TemplateRole::kSimdInputPre);
195+ } else {
196+ EXPECT_EQ(input_producer->GetName(),
197+ kInputPreType == ascir::IndirectLoadInputPreType::kRelu ? "input_relu" : "input_exp2");
198+ }
199+ }
200+ const auto index_producer = ascgen_utils::indirect_load::GetInputProducer(indirect_load, 1UL);
201+ ASSERT_NE(index_producer, nullptr);
202+ EXPECT_EQ(index_producer->GetName(), kMixedIndexPre ? "index_floor_to_int" : "index_abs");
203+ ExpectMixedInputPre(indirect_load);
204+#if IL_INPUT_PRE_TYPE == 5
205+ EXPECT_EQ(indirect_load->outputs()[0]->attr.dtype, af::DT_FLOAT16);
206+#endif
110}207}
111 208 
112-void CheckScheduledLoopFramework(ascir::FusedScheduledResult &result) {209+bool CheckScheduledLoopFramework(ascir::FusedScheduledResult &result) {
113 size_t simd_count = 0UL;210 size_t simd_count = 0UL;
114 size_t simt_count = 0UL;211 size_t simt_count = 0UL;
115 for (auto &candidates : result.node_idx_to_scheduled_results) {212 for (auto &candidates : result.node_idx_to_scheduled_results) {
@@ -131,12 +228,33 @@ void CheckScheduledLoopFramework(ascir::FusedScheduledResult &result) {
131 }228 }
132 }229 }
133 EXPECT_GT(simd_count, 0UL);230 EXPECT_GT(simd_count, 0UL);
134- EXPECT_GT(simt_count, 0UL);231+ if (kExpectSimt) {
232+ EXPECT_GT(simt_count, 0UL);
233+ return simd_count > 0UL && simt_count > 0UL;
234+ } else {
235+ EXPECT_EQ(simt_count, 0UL);
236+ return simd_count > 0UL && simt_count == 0UL;
237+ }
135}238}
136 239 
137void CheckGeneratedKernel(const std::string &kernel) {240void CheckGeneratedKernel(const std::string &kernel) {
138 EXPECT_NE(kernel.find("// IndirectLoad SIMD"), std::string::npos);241 EXPECT_NE(kernel.find("// IndirectLoad SIMD"), std::string::npos);
139- EXPECT_NE(kernel.find("// IndirectLoad SIMT"), std::string::npos);242+ EXPECT_NE(kernel.find("IndirectLoadSimd<"), std::string::npos);
243+ EXPECT_EQ(kernel.find("auto indirect_load_offset"), std::string::npos);
244+ if (kExpectSimt) {
245+ EXPECT_NE(kernel.find("// IndirectLoad SIMT"), std::string::npos);
246+ EXPECT_NE(kernel.find("IndirectLoadSimt<"), std::string::npos);
247+ EXPECT_EQ(kernel.find("IndirectLoadSimtKernel_"), std::string::npos);
248+ } else {
249+ EXPECT_EQ(kernel.find("// IndirectLoad SIMT"), std::string::npos);
250+ EXPECT_EQ(kernel.find("IndirectLoadSimt<"), std::string::npos);
251+ }
252+#ifdef IL_DATA_BF16
253+ EXPECT_NE(kernel.find("IndirectLoadSimt<bfloat16_t, int64_t, bfloat16_t"), std::string::npos);
254+#endif
255+#ifdef IL_DATA_UINT32
256+ EXPECT_NE(kernel.find("IndirectLoadSimd<uint32_t, int32_t"), std::string::npos);
257+#endif
140}258}
141} // namespace259} // namespace
142 260 
@@ -159,7 +277,9 @@ TEST_F(TestBackendIndirectLoadStoreE2e, IndirectLoadStoreCodegen) {
159#define REGISTER_TILING_DEFAULT(tiling)277#define REGISTER_TILING_DEFAULT(tiling)
160#define GET_TILING_DATA(t, tiling) AutofuseTilingData t = *(AutofuseTilingData*)tiling;278#define GET_TILING_DATA(t, tiling) AutofuseTilingData t = *(AutofuseTilingData*)tiling;
161)";279)";
162- auto graph = ascir::ShareGraph::IndirectLoadStoreFusedGraph(kRank, kAxis, af::DT_FLOAT16, kHasInputPre, kUseExp2);280+ auto graph = ascir::ShareGraph::IndirectLoadStoreFusedGraph(kRank, kAxis, kDataType, kIndexType, kInputPreType,
281+ kUseExp2, kOutputPostType, GetStaticShape(true),
282+ GetStaticShape(false), kMixedIndexPre);
163 ASSERT_NE(graph, nullptr);283 ASSERT_NE(graph, nullptr);
164 std::map<std::string, std::string> shape_info;284 std::map<std::string, std::string> shape_info;
165 for (size_t i = 0UL; i < 2UL * kRank; ++i) {285 for (size_t i = 0UL; i < 2UL * kRank; ++i) {
@@ -175,14 +295,16 @@ TEST_F(TestBackendIndirectLoadStoreE2e, IndirectLoadStoreCodegen) {
175 codegen::CodegenResult result;295 codegen::CodegenResult result;
176 testing::internal::CaptureStdout();296 testing::internal::CaptureStdout();
177 const auto optimize_status = optimizer.Optimize(graph, fused_schedule_result);297 const auto optimize_status = optimizer.Optimize(graph, fused_schedule_result);
178- if (optimize_status == 0) {298+ const std::string optimize_logs = testing::internal::GetCapturedStdout();
179- CheckScheduledLoopFramework(fused_schedule_result);299+ EXPECT_EQ(optimize_logs.find("[ERROR]"), std::string::npos) << optimize_logs;
180- }300+ ASSERT_EQ(optimize_status, 0) << optimize_logs;
181- const auto codegen_status = optimize_status == 0 ? codegen.Generate(shape_info, fused_schedule_result, result) : -1;301+ ASSERT_TRUE(CheckScheduledLoopFramework(fused_schedule_result)) << optimize_logs;
182- const std::string logs = testing::internal::GetCapturedStdout();302+ 
183- EXPECT_EQ(logs.find("[ERROR]"), std::string::npos) << logs;303+ testing::internal::CaptureStdout();
184- ASSERT_EQ(optimize_status, 0) << logs;304+ const auto codegen_status = codegen.Generate(shape_info, fused_schedule_result, result);
185- ASSERT_EQ(codegen_status, 0) << logs;305+ const std::string codegen_logs = testing::internal::GetCapturedStdout();
306+ EXPECT_EQ(codegen_logs.find("[ERROR]"), std::string::npos) << codegen_logs;
307+ ASSERT_EQ(codegen_status, 0) << codegen_logs;
186 CheckGeneratedKernel(result.kernel);308 CheckGeneratedKernel(result.kernel);
187 309 
188 std::fstream kernel_file(parts[0], std::ios::out);310 std::fstream kernel_file(parts[0], std::ios::out);
Mautofuse/tests/v35/st/backend_e2e_v2/indirect_load_store_test/test_e2e_indirect_load_store_kernel.cpp+91-18
@@ -22,29 +22,70 @@
22 22 
23extern "C" __global__ __aicore__ void indirect_load_store_test(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR workspace,23extern "C" __global__ __aicore__ void indirect_load_store_test(GM_ADDR x, GM_ADDR index, GM_ADDR y, GM_ADDR workspace,
24 GM_ADDR tiling);24 GM_ADDR tiling);
25+#ifdef IL_STATIC_SHAPE
26+extern "C" int64_t AutofuseTiling(AutofuseTilingData *, uint32_t *, uint32_t *, uint32_t, uint32_t);
27+#endif
25#if IL_RANK == 228#if IL_RANK == 2
29+#ifndef IL_STATIC_SHAPE
26extern "C" int64_t AutofuseTiling(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, AutofuseTilingData *, uint32_t *,30extern "C" int64_t AutofuseTiling(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, AutofuseTilingData *, uint32_t *,
27 uint32_t *, uint32_t, uint32_t);31 uint32_t *, uint32_t, uint32_t);
32+#endif
28constexpr std::array<int32_t, 2> kInputShape = {IL_X_S0, IL_X_S1};33constexpr std::array<int32_t, 2> kInputShape = {IL_X_S0, IL_X_S1};
29constexpr std::array<int32_t, 2> kIndexShape = {IL_INDEX_S0, IL_INDEX_S1};34constexpr std::array<int32_t, 2> kIndexShape = {IL_INDEX_S0, IL_INDEX_S1};
30#elif IL_RANK == 335#elif IL_RANK == 3
36+#ifndef IL_STATIC_SHAPE
31extern "C" int64_t AutofuseTiling(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, uint32_t s4, uint32_t s5,37extern "C" int64_t AutofuseTiling(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, uint32_t s4, uint32_t s5,
32 AutofuseTilingData *, uint32_t *, uint32_t *, uint32_t, uint32_t);38 AutofuseTilingData *, uint32_t *, uint32_t *, uint32_t, uint32_t);
39+#endif
33constexpr std::array<int32_t, 3> kInputShape = {IL_X_S0, IL_X_S1, IL_X_S2};40constexpr std::array<int32_t, 3> kInputShape = {IL_X_S0, IL_X_S1, IL_X_S2};
34constexpr std::array<int32_t, 3> kIndexShape = {IL_INDEX_S0, IL_INDEX_S1, IL_INDEX_S2};41constexpr std::array<int32_t, 3> kIndexShape = {IL_INDEX_S0, IL_INDEX_S1, IL_INDEX_S2};
35#elif IL_RANK == 442#elif IL_RANK == 4
43+#ifndef IL_STATIC_SHAPE
36extern "C" int64_t AutofuseTiling(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, uint32_t s4, uint32_t s5,44extern "C" int64_t AutofuseTiling(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t s3, uint32_t s4, uint32_t s5,
37 uint32_t s6, uint32_t s7, AutofuseTilingData *, uint32_t *, uint32_t *, uint32_t,45 uint32_t s6, uint32_t s7, AutofuseTilingData *, uint32_t *, uint32_t *, uint32_t,
38 uint32_t);46 uint32_t);
47+#endif
39constexpr std::array<int32_t, 4> kInputShape = {IL_X_S0, IL_X_S1, IL_X_S2, IL_X_S3};48constexpr std::array<int32_t, 4> kInputShape = {IL_X_S0, IL_X_S1, IL_X_S2, IL_X_S3};
40constexpr std::array<int32_t, 4> kIndexShape = {IL_INDEX_S0, IL_INDEX_S1, IL_INDEX_S2, IL_INDEX_S3};49constexpr std::array<int32_t, 4> kIndexShape = {IL_INDEX_S0, IL_INDEX_S1, IL_INDEX_S2, IL_INDEX_S3};
41#endif50#endif
42 51 
43namespace {52namespace {
44constexpr int32_t kAxis = IL_AXIS < 0 ? IL_AXIS + IL_RANK : IL_AXIS;53constexpr int32_t kAxis = IL_AXIS < 0 ? IL_AXIS + IL_RANK : IL_AXIS;
54+#if defined(IL_DATA_BF16)
55+using DataType = bfloat16_t;
56+#elif defined(IL_DATA_UINT32)
57+using DataType = uint32_t;
58+#elif defined(IL_DATA_FLOAT)
59+using DataType = float;
60+#else
61+using DataType = half;
62+#endif
63+#ifdef IL_INDEX_INT64
64+using IndexType = int64_t;
65+#else
66+using IndexType = int32_t;
67+#endif
68+ 
69+float BesselK0Reference(float x) {
70+ const float y = x * x / 4.0F;
71+ float series_term = 1.0F;
72+ float series_sum = 1.0F;
73+ for (int32_t order = 1; order <= 20; ++order) {
74+ const float scale = x * 0.5F / static_cast<float>(order);
75+ series_term *= scale * scale;
76+ series_sum += series_term;
77+ }
78+ const float poly =
79+ -0.57721566F +
80+ y * (0.42278420F +
81+ y * (0.23069756F + y * (0.03488590F + y * (0.00262698F + y * (0.00010750F + y * 0.00000740F)))));
82+ return -std::log(x * 0.5F) * series_sum + poly;
83+}
45 84 
46void RunTiling(AutofuseTilingData &tiling, uint32_t &workspace_size, uint32_t &block_dim) {85void RunTiling(AutofuseTilingData &tiling, uint32_t &workspace_size, uint32_t &block_dim) {
47-#if IL_RANK == 286+#ifdef IL_STATIC_SHAPE
87+ AutofuseTiling(&tiling, &workspace_size, &block_dim, 48, 192 * 1024);
88+#elif IL_RANK == 2
48 AutofuseTiling(kInputShape[0], kInputShape[1], kIndexShape[0], kIndexShape[1], &tiling, &workspace_size, &block_dim,89 AutofuseTiling(kInputShape[0], kInputShape[1], kIndexShape[0], kIndexShape[1], &tiling, &workspace_size, &block_dim,
49 48, 192 * 1024);90 48, 192 * 1024);
50#elif IL_RANK == 391#elif IL_RANK == 3
@@ -65,18 +106,24 @@ int32_t ElementCount(const std::array<int32_t, N> &shape) {
65 return count;106 return count;
66}107}
67 108 
68-void InitializeData(half *x, int32_t *index, half *expected, int32_t input_count, int32_t output_count) {109+void InitializeData(DataType *x, IndexType *index, DataType *expected, int32_t input_count, int32_t output_count) {
69 std::array<int32_t, IL_RANK> input_strides{};110 std::array<int32_t, IL_RANK> input_strides{};
70 input_strides.back() = 1;111 input_strides.back() = 1;
71 for (size_t i = input_strides.size() - 1UL; i > 0UL; --i) {112 for (size_t i = input_strides.size() - 1UL; i > 0UL; --i) {
72 input_strides[i - 1UL] = kInputShape[i] * input_strides[i];113 input_strides[i - 1UL] = kInputShape[i] * input_strides[i];
73 }114 }
74 for (int32_t i = 0; i < input_count; ++i) {115 for (int32_t i = 0; i < input_count; ++i) {
75- x[i] = static_cast<half>(static_cast<float>((i % 29) - 14) * 0.25F);116+#if defined(IL_DATA_UINT32)
117+ x[i] = static_cast<DataType>(i % 5);
118+#elif defined(IL_OUTPUT_POST_TYPE) && IL_OUTPUT_POST_TYPE == 1
119+ x[i] = static_cast<DataType>(static_cast<float>((i % 8) + 1) * 0.25F);
120+#else
121+ x[i] = static_cast<DataType>(static_cast<float>((i % 29) - 14) * 0.25F);
122+#endif
76 }123 }
77 for (int32_t i = 0; i < output_count; ++i) {124 for (int32_t i = 0; i < output_count; ++i) {
78 const int32_t index_value = (i * 3 + 1) % kInputShape[kAxis];125 const int32_t index_value = (i * 3 + 1) % kInputShape[kAxis];
79- index[i] = i % 2 == 0 ? index_value : -index_value;126+ index[i] = static_cast<IndexType>(i % 2 == 0 ? index_value : -index_value);
80 int32_t remaining = i;127 int32_t remaining = i;
81 int32_t input_offset = 0;128 int32_t input_offset = 0;
82 for (size_t dim = kIndexShape.size(); dim-- > 0UL;) {129 for (size_t dim = kIndexShape.size(); dim-- > 0UL;) {
@@ -85,13 +132,29 @@ void InitializeData(half *x, int32_t *index, half *expected, int32_t input_count
85 input_offset += (static_cast<int32_t>(dim) == kAxis ? index_value : coord) * input_strides[dim];132 input_offset += (static_cast<int32_t>(dim) == kAxis ? index_value : coord) * input_strides[dim];
86 }133 }
87 float value = static_cast<float>(x[input_offset]);134 float value = static_cast<float>(x[input_offset]);
88-#if IL_HAS_INPUT_PRE135+#if defined(IL_OUTPUT_POST_TYPE) && IL_OUTPUT_POST_TYPE == 1
89- value = std::max(value, 0.0F);136+ expected[i] = static_cast<DataType>(BesselK0Reference(value));
90-#endif137+#elif defined(IL_OUTPUT_POST_TYPE) && IL_OUTPUT_POST_TYPE == 2
91-#if IL_USE_EXP2138+ expected[i] = static_cast<DataType>(value);
92- expected[i] = static_cast<half>(-std::exp2(value));139+#elif defined(IL_OUTPUT_POST_TYPE) && IL_OUTPUT_POST_TYPE == 3
140+ expected[i] = static_cast<DataType>(std::exp2(value));
141+#elif IL_INPUT_PRE_TYPE == 5
142+ expected[i] = static_cast<DataType>(IL_USE_EXP2 ? std::exp2(value) : std::exp(value));
93#else143#else
94- expected[i] = static_cast<half>(-std::exp(value));144+#if IL_INPUT_PRE_TYPE == 1 || IL_INPUT_PRE_TYPE == 4
145+ value = std::max(value, 0.0F);
146+#elif IL_INPUT_PRE_TYPE == 3
147+ value = std::exp2(std::max(value, 0.0F));
148+#endif
149+#if IL_INPUT_PRE_TYPE == 2
150+ expected[i] = static_cast<DataType>(-std::exp2(value));
151+#elif IL_INPUT_PRE_TYPE == 3
152+ expected[i] = static_cast<DataType>(-value);
153+#elif IL_USE_EXP2
154+ expected[i] = static_cast<DataType>(-std::exp2(value));
155+#else
156+ expected[i] = static_cast<DataType>(-std::exp(value));
157+#endif
95#endif158#endif
96 }159 }
97}160}
@@ -101,32 +164,42 @@ TEST(E2EIndirectLoadStore, GeneratedKernelMatchesReference) {
101 const int32_t input_count = ElementCount(kInputShape);164 const int32_t input_count = ElementCount(kInputShape);
102 const int32_t output_count = ElementCount(kIndexShape);165 const int32_t output_count = ElementCount(kIndexShape);
103 const auto gm_free = [](void *ptr) { AscendC::GmFree(ptr); };166 const auto gm_free = [](void *ptr) { AscendC::GmFree(ptr); };
104- std::unique_ptr<half, decltype(gm_free)> x(reinterpret_cast<half *>(AscendC::GmAlloc(input_count * sizeof(half))),167+ std::unique_ptr<DataType, decltype(gm_free)> x(
105- gm_free);168+ reinterpret_cast<DataType *>(AscendC::GmAlloc(input_count * sizeof(DataType))), gm_free);
106- std::unique_ptr<int32_t, decltype(gm_free)> index(169+ std::unique_ptr<IndexType, decltype(gm_free)> index(
107- reinterpret_cast<int32_t *>(AscendC::GmAlloc(output_count * sizeof(int32_t))), gm_free);170+ reinterpret_cast<IndexType *>(AscendC::GmAlloc(output_count * sizeof(IndexType))), gm_free);
108- std::unique_ptr<half, decltype(gm_free)> output(171+ std::unique_ptr<DataType, decltype(gm_free)> output(
109- reinterpret_cast<half *>(AscendC::GmAlloc(output_count * sizeof(half))), gm_free);172+ reinterpret_cast<DataType *>(AscendC::GmAlloc(output_count * sizeof(DataType))), gm_free);
110 ASSERT_NE(x, nullptr);173 ASSERT_NE(x, nullptr);
111 ASSERT_NE(index, nullptr);174 ASSERT_NE(index, nullptr);
112 ASSERT_NE(output, nullptr);175 ASSERT_NE(output, nullptr);
113- std::vector<half> expected(static_cast<size_t>(output_count));176+ std::vector<DataType> expected(static_cast<size_t>(output_count));
114 InitializeData(x.get(), index.get(), expected.data(), input_count, output_count);177 InitializeData(x.get(), index.get(), expected.data(), input_count, output_count);
115- std::fill_n(output.get(), output_count, static_cast<half>(0.0F));178+ std::fill_n(output.get(), output_count, static_cast<DataType>(0.0F));
116 179 
117 AutofuseTilingData tiling_data{};180 AutofuseTilingData tiling_data{};
118 uint32_t workspace_size = 0U;181 uint32_t workspace_size = 0U;
119 uint32_t block_dim = 48U;182 uint32_t block_dim = 48U;
120 RunTiling(tiling_data, workspace_size, block_dim);183 RunTiling(tiling_data, workspace_size, block_dim);
184+#if !defined(IL_EXPECT_SIMT) || IL_EXPECT_SIMT
121 ASSERT_EQ(tiling_data.graph0_tiling_key, IL_TILING_KEY);185 ASSERT_EQ(tiling_data.graph0_tiling_key, IL_TILING_KEY);
186+#endif
122 ASSERT_GT(tiling_data.block_dim, 0U);187 ASSERT_GT(tiling_data.block_dim, 0U);
188+#ifdef IL_EXPECT_SIMT_MULTI_ROUND
189+ ASSERT_GT((static_cast<uint32_t>(output_count) + tiling_data.block_dim - 1U) / tiling_data.block_dim, 1024U);
190+#endif
123 191 
124 AscendC::SetKernelMode(KernelMode::AIV_MODE);192 AscendC::SetKernelMode(KernelMode::AIV_MODE);
125 ICPU_RUN_KF(indirect_load_store_test, tiling_data.block_dim, reinterpret_cast<uint8_t *>(x.get()),193 ICPU_RUN_KF(indirect_load_store_test, tiling_data.block_dim, reinterpret_cast<uint8_t *>(x.get()),
126 reinterpret_cast<uint8_t *>(index.get()), reinterpret_cast<uint8_t *>(output.get()), nullptr,194 reinterpret_cast<uint8_t *>(index.get()), reinterpret_cast<uint8_t *>(output.get()), nullptr,
127 reinterpret_cast<uint8_t *>(&tiling_data));195 reinterpret_cast<uint8_t *>(&tiling_data));
128 for (int32_t i = 0; i < output_count; ++i) {196 for (int32_t i = 0; i < output_count; ++i) {
197+#if defined(IL_DATA_BF16) || defined(IL_DATA_UINT32)
198+ EXPECT_EQ(static_cast<float>(output.get()[i]), static_cast<float>(expected[static_cast<size_t>(i)]))
199+ << "offset=" << i;
200+#else
129 EXPECT_NEAR(static_cast<float>(output.get()[i]), static_cast<float>(expected[static_cast<size_t>(i)]), 0.0625F)201 EXPECT_NEAR(static_cast<float>(output.get()[i]), static_cast<float>(expected[static_cast<size_t>(i)]), 0.0625F)
130 << "offset=" << i;202 << "offset=" << i;
203+#endif
131 }204 }
132}205}
Mautofuse/v35/ascendc/api_regbase/CMakeLists.txt+2-0
@@ -78,6 +78,8 @@ set(ascendc_api_regbase_extend_src
78 igammac_helper/asymptotic_series.h78 igammac_helper/asymptotic_series.h
79 igammac_helper/series_complement.h79 igammac_helper/series_complement.h
80 hypot.h80 hypot.h
81+ indirect_load_simd.h
82+ indirect_load_simt.h
81)83)
82 84 
83set(generated_regbase_files)85set(generated_regbase_files)
Aautofuse/v35/ascendc/api_regbase/indirect_load_simd.h+68-0
@@ -0,0 +1,68 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef AUTOFUSE_V35_ASCENDC_API_REGBASE_INDIRECT_LOAD_SIMD_H_
12+#define AUTOFUSE_V35_ASCENDC_API_REGBASE_INDIRECT_LOAD_SIMD_H_
13+ 
14+namespace AscendC {
15+namespace Internal {
16+template <int32_t Dim, int32_t Axis, int32_t Rank>
17+struct IndirectLoadSimdInnerOffset {
18+ __aicore__ inline static int64_t Call(int64_t inner, const int64_t *shape) {
19+ int64_t offset = inner % shape[Dim] * shape[Rank + Dim];
20+ if constexpr (Dim > Axis + 1) {
21+ offset += IndirectLoadSimdInnerOffset<Dim - 1, Axis, Rank>::Call(inner / shape[Dim], shape);
22+ }
23+ return offset;
24+ }
25+};
26+ 
27+template <int32_t Dim, int32_t Axis>
28+struct IndirectLoadSimdInnerSize {
29+ __aicore__ inline static int64_t Call(const int64_t *shape) {
30+ if constexpr (Dim == Axis) {
31+ return 1;
32+ } else {
33+ return shape[Dim] * IndirectLoadSimdInnerSize<Dim - 1, Axis>::Call(shape);
34+ }
35+ }
36+};
37+} // namespace Internal
38+ 
39+template <typename X, typename Index, int32_t Rank, int32_t Axis, typename... ShapeArgs>
40+__aicore__ inline void IndirectLoadSimd(const LocalTensor<X> &x, const LocalTensor<Index> &index,
41+ const LocalTensor<X> &y, const LocalTensor<uint8_t> &tmp, uint32_t actual_size,
42+ int64_t output_offset, int64_t input_axis, ShapeArgs... shape_args) {
43+ static_assert(Rank > 0 && Axis >= 0 && Axis < Rank, "IndirectLoad SIMD rank or axis is invalid.");
44+ static_assert(sizeof...(ShapeArgs) == static_cast<size_t>(2 * Rank), "IndirectLoad SIMD shape is invalid.");
45+ const int64_t shape[] = {static_cast<int64_t>(shape_args)...};
46+ const int64_t index_inner = Internal::IndirectLoadSimdInnerSize<Rank - 1, Axis>::Call(shape);
47+ const int64_t input_inner = shape[Rank + Axis];
48+ const int64_t input_slice_count = input_axis * input_inner;
49+ const int64_t output_slice_count = shape[Axis] * index_inner;
50+ const int64_t outer_begin = output_offset / output_slice_count;
51+ LocalTensor<uint32_t> offsets = tmp.template ReinterpretCast<uint32_t>();
52+ for (int64_t i = 0; i < actual_size; ++i) {
53+ const int64_t global_idx = output_offset + i;
54+ const int64_t outer_local = global_idx / output_slice_count - outer_begin;
55+ const int64_t index_value = static_cast<int64_t>(index.GetValue(i));
56+ int64_t input_inner_offset = 0;
57+ if constexpr (Axis + 1 < Rank) {
58+ input_inner_offset =
59+ Internal::IndirectLoadSimdInnerOffset<Rank - 1, Axis, Rank>::Call(global_idx % index_inner, shape);
60+ }
61+ const int64_t src_idx = outer_local * input_slice_count + index_value * input_inner + input_inner_offset;
62+ offsets.SetValue(i, static_cast<uint32_t>(src_idx * sizeof(X)));
63+ }
64+ Gather(y, x, offsets, static_cast<uint32_t>(0), actual_size);
65+}
66+} // namespace AscendC
67+ 
68+#endif // AUTOFUSE_V35_ASCENDC_API_REGBASE_INDIRECT_LOAD_SIMD_H_
Aautofuse/v35/ascendc/api_regbase/indirect_load_simt.h+72-0
@@ -0,0 +1,72 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef AUTOFUSE_V35_ASCENDC_API_REGBASE_INDIRECT_LOAD_SIMT_H_
12+#define AUTOFUSE_V35_ASCENDC_API_REGBASE_INDIRECT_LOAD_SIMT_H_
13+ 
14+namespace AscendC {
15+namespace Internal {
16+constexpr uint32_t kIndirectLoadSimtThreadNum = 1024U;
17+ 
18+template <int32_t Dim, int32_t Axis, int32_t Rank>
19+struct IndirectLoadSimtOffset {
20+ __simt_callee__ __aicore__ inline static int64_t Call(int64_t linear_index, int64_t indirect_index,
21+ const int64_t *shape) {
22+ int64_t input_offset = 0;
23+ if constexpr (Dim == Axis) {
24+ input_offset = indirect_index * shape[Rank + Dim];
25+ } else {
26+ input_offset = linear_index % shape[Dim] * shape[Rank + Dim];
27+ }
28+ if constexpr (Dim > 0) {
29+ input_offset +=
30+ IndirectLoadSimtOffset<Dim - 1, Axis, Rank>::Call(linear_index / shape[Dim], indirect_index, shape);
31+ }
32+ return input_offset;
33+ }
34+};
35+ 
36+template <typename X, typename Index, typename Y, int32_t Rank, int32_t Axis, typename IndexTransform,
37+ typename OutputTransform, typename... ShapeArgs>
38+__simt_vf__ __aicore__ LAUNCH_BOUND(kIndirectLoadSimtThreadNum) inline void IndirectLoadSimtKernel(
39+ __gm__ X *x, __gm__ Index *index, __gm__ Y *y, uint32_t actual_size, int64_t output_offset, int64_t x_axis_size,
40+ ShapeArgs... shape_args) {
41+ static_assert(Rank > 0 && Axis >= 0 && Axis < Rank, "IndirectLoad SIMT rank or axis is invalid.");
42+ static_assert(sizeof...(ShapeArgs) == static_cast<size_t>(2 * Rank), "IndirectLoad SIMT shape is invalid.");
43+ const int64_t shape[] = {static_cast<int64_t>(shape_args)...};
44+ for (uint32_t i = threadIdx.x; i < actual_size; i += blockDim.x) {
45+ const int64_t output_index = output_offset + i;
46+ const int64_t indirect_index = static_cast<int64_t>(IndexTransform::Call(index[output_index]));
47+ if (unlikely(indirect_index < 0 || indirect_index >= x_axis_size)) {
48+ y[output_index] = static_cast<Y>(0);
49+ continue;
50+ }
51+ const int64_t input_offset =
52+ IndirectLoadSimtOffset<Rank - 1, Axis, Rank>::Call(output_index, indirect_index, shape);
53+ y[output_index] = OutputTransform::Call(x[input_offset]);
54+ }
55+}
56+} // namespace Internal
57+ 
58+template <typename X, typename Index, typename Y, int32_t Rank, int32_t Axis, typename IndexTransform,
59+ typename OutputTransform, typename... ShapeArgs>
60+__aicore__ inline void IndirectLoadSimt(__gm__ X *x, __gm__ Index *index, __gm__ Y *y, uint32_t actual_size,
61+ int64_t output_offset, int64_t x_axis_size, ShapeArgs... shape_args) {
62+ Simt::VF_CALL<
63+ Internal::IndirectLoadSimtKernel<X, Index, Y, Rank, Axis, IndexTransform, OutputTransform, ShapeArgs...>>(
64+ Simt::Dim3(Internal::kIndirectLoadSimtThreadNum), x, index, y, actual_size, output_offset, x_axis_size,
65+ shape_args...);
66+ const int32_t event_id = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
67+ SetFlag<HardEvent::V_MTE3>(event_id);
68+ WaitFlag<HardEvent::V_MTE3>(event_id);
69+}
70+} // namespace AscendC
71+ 
72+#endif // AUTOFUSE_V35_ASCENDC_API_REGBASE_INDIRECT_LOAD_SIMT_H_
Mautofuse/v35/ascir/ascir_codegen_v2.h+13-0
@@ -30,6 +30,19 @@ class AscIrCodegenV2 : public AscIrCodegen {
30 virtual uint32_t GetMicroInstNum() const {30 virtual uint32_t GetMicroInstNum() const {
31 return 1U;31 return 1U;
32 }32 }
33+ 
34+ virtual bool IsSimtScalarSupported(const AscNode &node) const {
35+ (void)node;
36+ return false;
37+ }
38+ 
39+ virtual ge::graphStatus GenerateSimtScalarExpr(const AscNode &node, const std::vector<std::string> &inputs,
40+ std::string &expr) const {
41+ (void)node;
42+ (void)inputs;
43+ (void)expr;
44+ return af::FAILED;
45+ }
33};46};
34} // namespace ascir47} // namespace ascir
35} // namespace af48} // namespace af
Mautofuse/v35/codegen/ascendc_reg_base_api_register.cpp+8-0
@@ -242,6 +242,12 @@ Register::Register() {
242 };242 };
243 const std::string kAscendcShiftedChebyshevPolynomialWRegBaseStr = {243 const std::string kAscendcShiftedChebyshevPolynomialWRegBaseStr = {
244#include "shifted_chebyshev_polynomial_w_reg_base.h"244#include "shifted_chebyshev_polynomial_w_reg_base.h"
245+ };
246+ const std::string kAscendcIndirectLoadSimtRegBaseStr = {
247+#include "indirect_load_simt_reg_base.h"
248+ };
249+ const std::string kAscendcIndirectLoadSimdRegBaseStr = {
250+#include "indirect_load_simd_reg_base.h"
245 };251 };
246 std::unordered_map<std::string, std::string> api_to_file{252 std::unordered_map<std::string, std::string> api_to_file{
247 {"cast_reg_base.h", kAscendcCastRegStr},253 {"cast_reg_base.h", kAscendcCastRegStr},
@@ -313,6 +319,8 @@ Register::Register() {
313 {"shifted_chebyshev_polynomial_u_reg_base.h", kAscendcShiftedChebyshevPolynomialURegBaseStr},319 {"shifted_chebyshev_polynomial_u_reg_base.h", kAscendcShiftedChebyshevPolynomialURegBaseStr},
314 {"shifted_chebyshev_polynomial_v_reg_base.h", kAscendcShiftedChebyshevPolynomialVRegBaseStr},320 {"shifted_chebyshev_polynomial_v_reg_base.h", kAscendcShiftedChebyshevPolynomialVRegBaseStr},
315 {"shifted_chebyshev_polynomial_w_reg_base.h", kAscendcShiftedChebyshevPolynomialWRegBaseStr},321 {"shifted_chebyshev_polynomial_w_reg_base.h", kAscendcShiftedChebyshevPolynomialWRegBaseStr},
322+ {"indirect_load_simt_reg_base.h", kAscendcIndirectLoadSimtRegBaseStr},
323+ {"indirect_load_simd_reg_base.h", kAscendcIndirectLoadSimdRegBaseStr},
316 };324 };
317 325 
318 AscendCApiRegistry::GetInstance().RegisterApi(api_to_file);326 AscendCApiRegistry::GetInstance().RegisterApi(api_to_file);
Mautofuse/v35/codegen/reg_api_call/reg_indirect_load_api_call.cpp+40-276
@@ -10,15 +10,11 @@
10 10 
11#include "reg_indirect_load_api_call.h"11#include "reg_indirect_load_api_call.h"
12 12 
13-#include <map>
14#include <sstream>13#include <sstream>
15#include "api_call/utils/api_call_factory.h"14#include "api_call/utils/api_call_factory.h"
16-#include "asc_graph_utils.h"
17#include "ascir_ops.h"15#include "ascir_ops.h"
18#include "common/checker.h"16#include "common/checker.h"
19-#include "common_utils.h"
20#include "indirect_load_utils.h"17#include "indirect_load_utils.h"
21-#include "schedule_result.h"
22#include "v35/codegen/simt_scalar_call/simt_scalar_emitter.h"18#include "v35/codegen/simt_scalar_call/simt_scalar_emitter.h"
23 19 
24namespace codegen {20namespace codegen {
@@ -28,8 +24,6 @@ struct LogicalTensorInfo {
28 std::vector<ascir::SizeExpr> strides;24 std::vector<ascir::SizeExpr> strides;
29};25};
30 26 
31-constexpr uint32_t kIndirectLoadSimtThreadNum = 1024;
32- 
33bool IsAxisDerivedFrom(const TPipe &tpipe, ascir::AxisId axis_id, ascir::AxisId ancestor_axis_id) {27bool IsAxisDerivedFrom(const TPipe &tpipe, ascir::AxisId axis_id, ascir::AxisId ancestor_axis_id) {
34 if (axis_id == ancestor_axis_id) {28 if (axis_id == ancestor_axis_id) {
35 return true;29 return true;
@@ -100,37 +94,12 @@ af::Status CheckIndirectLoadShape(const LogicalTensorInfo &x_info, const Logical
100 return af::SUCCESS;94 return af::SUCCESS;
101}95}
102 96 
103-af::Status CollectVectorFuncNodes(const af::AscNodePtr &vf_node, std::vector<af::AscNodePtr> &nodes) {
104- if (vf_node == nullptr || !af::ops::IsOps<af::ascir_op::VectorFunc>(vf_node)) {
105- return af::SUCCESS;
106- }
107- const std::string *graph_name = af::AttrUtils::GetStr(vf_node->GetOpDescBarePtr(), "sub_graph_name");
108- GE_ASSERT_NOTNULL(graph_name, "Get sub graph name failed, vf node:%s", vf_node->GetNamePtr());
109- af::AscGraph owner_graph("owner_graph");
110- GE_ASSERT_SUCCESS(af::AscGraphUtils::FromComputeGraph(vf_node->GetOwnerComputeGraph(), owner_graph),
111- "Get owner graph failed, vf node:%s", vf_node->GetNamePtr());
112- af::AscGraph sub_graph("vf_sub_graph");
113- GE_ASSERT_SUCCESS(owner_graph.FindSubGraph(*graph_name, sub_graph), "Get sub graph failed, vf node:%s",
114- vf_node->GetNamePtr());
115- for (const af::AscNodePtr &sub_node : sub_graph.GetAllNodes()) {
116- if (af::ops::IsOps<af::ascir_op::Data>(sub_node) || af::ops::IsOps<af::ascir_op::Load>(sub_node) ||
117- af::ops::IsOps<af::ascir_op::Store>(sub_node) || af::ops::IsOps<af::ascir_op::Output>(sub_node)) {
118- continue;
119- }
120- nodes.emplace_back(sub_node);
121- }
122- return af::SUCCESS;
123-}
124- 
125af::Status CollectSimtInlineTransformNode(const af::AscNodePtr &node, std::vector<af::AscNodePtr> &nodes) {97af::Status CollectSimtInlineTransformNode(const af::AscNodePtr &node, std::vector<af::AscNodePtr> &nodes) {
126 GE_ASSERT_TRUE(ascgen_utils::indirect_load::GetTemplateRole(node) ==98 GE_ASSERT_TRUE(ascgen_utils::indirect_load::GetTemplateRole(node) ==
127 ascgen_utils::indirect_load::TemplateRole::kSimtInlineTransform,99 ascgen_utils::indirect_load::TemplateRole::kSimtInlineTransform,
128 "IndirectLoad SIMT transform node has no inline-transform role.");100 "IndirectLoad SIMT transform node has no inline-transform role.");
129- const size_t node_count = nodes.size();101+ GE_ASSERT_TRUE(!af::ops::IsOps<af::ascir_op::VectorFunc>(node),
130- GE_ASSERT_SUCCESS(CollectVectorFuncNodes(node, nodes));102+ "IndirectLoad SIMT transform must use scalar emission, node:%s", node->GetNamePtr());
131- if (nodes.size() != node_count) {
132- return af::SUCCESS;
133- }
134 GE_ASSERT_TRUE(node->GetInDataNodesSize() == 1UL, "IndirectLoad SIMT transform node must have exactly one input.");103 GE_ASSERT_TRUE(node->GetInDataNodesSize() == 1UL, "IndirectLoad SIMT transform node must have exactly one input.");
135 nodes.emplace_back(node);104 nodes.emplace_back(node);
136 return af::SUCCESS;105 return af::SUCCESS;
@@ -142,10 +111,6 @@ af::Status CollectSimtIndexPreNodes(const ascir::NodeView &node, std::vector<af:
142 auto peer_out_anchor = index_anchor->GetPeerOutAnchor();111 auto peer_out_anchor = index_anchor->GetPeerOutAnchor();
143 GE_ASSERT_NOTNULL(peer_out_anchor, "IndirectLoad index input anchor has no peer output.");112 GE_ASSERT_NOTNULL(peer_out_anchor, "IndirectLoad index input anchor has no peer output.");
144 auto current = std::dynamic_pointer_cast<af::AscNode>(peer_out_anchor->GetOwnerNode());113 auto current = std::dynamic_pointer_cast<af::AscNode>(peer_out_anchor->GetOwnerNode());
145- GE_ASSERT_SUCCESS(CollectVectorFuncNodes(current, nodes));
146- if (!nodes.empty()) {
147- return af::SUCCESS;
148- }
149 std::vector<af::AscNodePtr> reversed_nodes;114 std::vector<af::AscNodePtr> reversed_nodes;
150 while (current != nullptr && !af::ops::IsOps<af::ascir_op::Data>(current)) {115 while (current != nullptr && !af::ops::IsOps<af::ascir_op::Data>(current)) {
151 if (af::ops::IsOps<af::ascir_op::Load>(current)) {116 if (af::ops::IsOps<af::ascir_op::Load>(current)) {
@@ -175,99 +140,6 @@ af::Status CollectSimtOutputMetadata(const ascir::NodeView &node, std::vector<af
175 return af::FAILED;140 return af::FAILED;
176}141}
177 142 
178-struct SimdCodegenParams {
179- const Tensor &x;
180- const Tensor &index;
181- const Tensor &y;
182- af::AscNodePtr input_pre_vf;
183- size_t axis_pos;
184- std::string x_dtype_name;
185- std::string y_dtype_name;
186- std::string tmp_buf_name;
187- std::string output_offset_expr;
188- std::string index_axis_expr;
189- std::string input_axis_expr;
190- std::string index_inner_expr;
191- std::string input_inner_expr;
192- std::string index_outer_sizes_expr;
193- std::string input_outer_strides_expr;
194-};
195- 
196-void EmitSimdWindowSetup(const SimdCodegenParams &params, std::stringstream &ss) {
197- ss << " LocalTensor<uint8_t> indirect_load_tmp = " << params.tmp_buf_name << ";" << std::endl;
198- ss << " LocalTensor<" << params.x_dtype_name << "> input_used = indirect_load_tmp.template ReinterpretCast<"
199- << params.x_dtype_name << ">();" << std::endl;
200- ss << " int64_t output_offset = " << params.output_offset_expr << ";" << std::endl;
201- ss << " int64_t index_axis = " << params.index_axis_expr << ";" << std::endl;
202- ss << " int64_t input_axis = " << params.input_axis_expr << ";" << std::endl;
203- ss << " int64_t index_inner = " << params.index_inner_expr << ";" << std::endl;
204- ss << " int64_t input_inner = " << params.input_inner_expr << ";" << std::endl;
205- ss << " int64_t input_slice_count = input_axis * input_inner;" << std::endl;
206- ss << " int64_t output_slice_count = index_axis * index_inner;" << std::endl;
207- ss << " int64_t outer_begin = output_offset / output_slice_count;" << std::endl;
208- ss << " int64_t outer_end = (output_offset + " << params.y.actual_size
209- << " + output_slice_count - 1) / output_slice_count;" << std::endl;
210- ss << " int64_t outer_count = outer_end - outer_begin;" << std::endl;
211- ss << " int64_t input_window_count = outer_count * input_slice_count;" << std::endl;
212- ss << " constexpr int64_t indirect_load_block_bytes = 32;" << std::endl;
213- ss << " int64_t input_bytes = (input_window_count * sizeof(" << params.x_dtype_name
214- << ") + indirect_load_block_bytes - 1) / indirect_load_block_bytes * indirect_load_block_bytes;" << std::endl;
215- if (params.input_pre_vf != nullptr) {
216- ss << " int64_t input_pre_bytes = (input_window_count * sizeof(" << params.y_dtype_name
217- << ") + indirect_load_block_bytes - 1) / indirect_load_block_bytes * indirect_load_block_bytes;" << std::endl;
218- }
219-}
220- 
221-void EmitSimdInputWindow(const SimdCodegenParams &params, std::stringstream &ss) {
222- if (params.axis_pos == 0UL) {
223- ss << " DataCopyPadExtend<" << params.x_dtype_name << ", AscendC::PaddingMode::Normal>(input_used[0], " << params.x
224- << "[0], 1, input_window_count, 0, 0);" << std::endl;
225- } else {
226- ss << " int64_t index_sizes[] = {" << params.index_outer_sizes_expr << "};" << std::endl;
227- ss << " int64_t input_outer_strides[] = {" << params.input_outer_strides_expr << "};" << std::endl;
228- ss << " for (int64_t outer_local = 0; outer_local < outer_count; ++outer_local) {" << std::endl;
229- ss << " int64_t outer_linear = outer_begin + outer_local;" << std::endl;
230- ss << " int64_t input_window_offset = 0;" << std::endl;
231- ss << " for (int64_t axis = " << params.axis_pos << " - 1; axis >= 0; --axis) {" << std::endl;
232- ss << " int64_t coord = outer_linear % index_sizes[axis];" << std::endl;
233- ss << " outer_linear = outer_linear / index_sizes[axis];" << std::endl;
234- ss << " input_window_offset += coord * input_outer_strides[axis];" << std::endl;
235- ss << " }" << std::endl;
236- ss << " DataCopyPadExtend<" << params.x_dtype_name << ", AscendC::PaddingMode::Normal>("
237- << "input_used[outer_local * input_slice_count], " << params.x
238- << "[input_window_offset], 1, input_slice_count, 0, 0);" << std::endl;
239- ss << " }" << std::endl;
240- }
241- if (params.input_pre_vf == nullptr) {
242- ss << " LocalTensor<" << params.x_dtype_name << "> input_pre_used = input_used;" << std::endl;
243- return;
244- }
245- ss << " LocalTensor<uint8_t> input_pre_buf = indirect_load_tmp[input_bytes];" << std::endl;
246- ss << " LocalTensor<" << params.y_dtype_name << "> input_pre_used = input_pre_buf.template ReinterpretCast<"
247- << params.y_dtype_name << ">();" << std::endl;
248- ss << " VFCall" << params.input_pre_vf->GetName() << "((__local_mem__ " << params.y_dtype_name
249- << " *)input_pre_used[0].GetPhyAddr(), (__local_mem__ " << params.x_dtype_name
250- << " *)input_used[0].GetPhyAddr(), static_cast<uint32_t>(input_window_count));" << std::endl;
251-}
252- 
253-void EmitSimdGather(const SimdCodegenParams &params, std::stringstream &ss) {
254- const bool has_input_pre = params.input_pre_vf != nullptr;
255- ss << " LocalTensor<uint8_t> offset_buf = indirect_load_tmp["
256- << (has_input_pre ? "input_bytes + input_pre_bytes" : "input_bytes") << "];" << std::endl;
257- ss << " auto indirect_load_offset = offset_buf.template ReinterpretCast<uint32_t>();" << std::endl;
258- ss << " for (int64_t i = 0; i < " << params.y.actual_size << "; ++i) {" << std::endl;
259- ss << " int64_t global_idx = output_offset + i;" << std::endl;
260- ss << " int64_t outer_local = global_idx / output_slice_count - outer_begin;" << std::endl;
261- ss << " int64_t index_val = static_cast<int64_t>(" << params.index << ".GetValue(i));" << std::endl;
262- ss << " int64_t inner = global_idx % index_inner;" << std::endl;
263- ss << " int64_t src_idx = outer_local * input_slice_count + index_val * input_inner + inner;" << std::endl;
264- ss << " indirect_load_offset.SetValue(i, static_cast<uint32_t>(src_idx * sizeof("
265- << (has_input_pre ? params.y_dtype_name : params.x_dtype_name) << ")));" << std::endl;
266- ss << " }" << std::endl;
267- ss << " AscendC::Gather(" << params.y << ", input_pre_used, indirect_load_offset, static_cast<uint32_t>(0), "
268- << params.y.actual_size << ");" << std::endl;
269-}
270- 
271af::Status GenerateTypedTransform(const std::string &name, const std::string &input_dtype,143af::Status GenerateTypedTransform(const std::string &name, const std::string &input_dtype,
272 const std::vector<af::AscNodePtr> &nodes, std::stringstream &ss) {144 const std::vector<af::AscNodePtr> &nodes, std::stringstream &ss) {
273 std::string return_dtype = input_dtype;145 std::string return_dtype = input_dtype;
@@ -300,87 +172,6 @@ af::Status GenerateTypedTransform(const std::string &name, const std::string &in
300 ss << "};" << std::endl;172 ss << "};" << std::endl;
301 return af::SUCCESS;173 return af::SUCCESS;
302}174}
303- 
304-struct SimtKernelDefinitionParams {
305- const std::string &kernel_name;
306- const std::string &x_dtype_name;
307- const std::string &index_dtype_name;
308- const std::string &output_dtype_name;
309- const std::string &index_transform;
310- const std::string &output_transform;
311- size_t rank;
312- int64_t axis;
313-};
314- 
315-void EmitSimtKernelDefinition(const SimtKernelDefinitionParams &params, std::stringstream &ss) {
316- ss << "__simt_vf__ __aicore__ LAUNCH_BOUND(" << kIndirectLoadSimtThreadNum << ") inline void " << params.kernel_name
317- << "(__gm__ " << params.x_dtype_name << " *x, __gm__ " << params.index_dtype_name << " *index, __gm__ "
318- << params.output_dtype_name << " *y, uint32_t actual_size, int64_t output_offset, int64_t x_axis_size";
319- for (size_t i = 0; i < params.rank; ++i) {
320- ss << ", int64_t index_size" << i;
321- }
322- for (size_t i = 0; i < params.rank; ++i) {
323- ss << ", int64_t x_stride" << i;
324- }
325- ss << ") {" << std::endl;
326- ss << " for (uint32_t i = threadIdx.x; i < actual_size; i += blockDim.x) {" << std::endl;
327- ss << " int64_t output_index = output_offset + i;" << std::endl;
328- ss << " int64_t indirect_index = static_cast<int64_t>(" << params.index_transform
329- << "::Call(index[output_index]));" << std::endl;
330- ss << " if (unlikely(indirect_index < 0 || indirect_index >= x_axis_size)) {" << std::endl;
331- ss << " y[output_index] = static_cast<" << params.output_dtype_name << ">(0);" << std::endl;
332- ss << " continue;" << std::endl;
333- ss << " }" << std::endl;
334- ss << " int64_t linear_index = output_index;" << std::endl;
335- ss << " int64_t input_offset = 0;" << std::endl;
336- for (int64_t axis = static_cast<int64_t>(params.rank) - 1; axis >= 0; --axis) {
337- if (axis != params.axis) {
338- ss << " int64_t coord" << axis << " = linear_index % index_size" << axis << ";" << std::endl;
339- }
340- if (axis > 0) {
341- ss << " linear_index = linear_index / index_size" << axis << ";" << std::endl;
342- }
343- if (axis == params.axis) {
344- ss << " input_offset += indirect_index * x_stride" << axis << ";" << std::endl;
345- } else {
346- ss << " input_offset += coord" << axis << " * x_stride" << axis << ";" << std::endl;
347- }
348- }
349- ss << " y[output_index] = " << params.output_transform << "::Call(x[input_offset]);" << std::endl;
350- ss << " }" << std::endl;
351- ss << "}" << std::endl;
352-}
353- 
354-struct SimtCodegenParams {
355- const Tensor &x;
356- const Tensor &index;
357- std::string output_gm_expr;
358- std::string x_dtype_name;
359- std::string index_dtype_name;
360- std::string output_dtype_name;
361- std::string kernel_name;
362- std::string x_axis_size_expr;
363- std::string rank_sizes_expr;
364- std::string x_strides_expr;
365- std::string outer_tb_var;
366-};
367- 
368-void EmitSimtLaunch(const SimtCodegenParams &params, std::stringstream &ss) {
369- ss << " __gm__ " << params.x_dtype_name << " *x_ptr = (__gm__ " << params.x_dtype_name << " *)" << params.x
370- << ".GetPhyAddr();" << std::endl;
371- ss << " __gm__ " << params.index_dtype_name << " *index_ptr = (__gm__ " << params.index_dtype_name << " *)"
372- << params.index << ".GetPhyAddr();" << std::endl;
373- ss << " __gm__ " << params.output_dtype_name << " *y_ptr = (__gm__ " << params.output_dtype_name << " *)"
374- << params.output_gm_expr << ".GetPhyAddr();" << std::endl;
375- ss << " AscendC::Simt::VF_CALL<" << params.kernel_name << ">(AscendC::Simt::Dim3(" << kIndirectLoadSimtThreadNum
376- << "), x_ptr, index_ptr, y_ptr, static_cast<uint32_t>(" << params.outer_tb_var << "_loop_size), block_dim_offset, "
377- << params.x_axis_size_expr << params.rank_sizes_expr << params.x_strides_expr << ");" << std::endl;
378- ss << " int32_t event_id_v_to_mte3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::V_MTE3));"
379- << std::endl;
380- ss << " AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(event_id_v_to_mte3);" << std::endl;
381- ss << " AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(event_id_v_to_mte3);" << std::endl;
382-}
383- 
384} // namespace175} // namespace
385 176 
386Status IndirectLoadRegApiCall::ParseAttr(const ascir::NodeView &node) {177Status IndirectLoadRegApiCall::ParseAttr(const ascir::NodeView &node) {
@@ -390,20 +181,15 @@ Status IndirectLoadRegApiCall::ParseAttr(const ascir::NodeView &node) {
390 ascgen_utils::indirect_load::TemplateAxes template_axes;181 ascgen_utils::indirect_load::TemplateAxes template_axes;
391 GE_ASSERT_SUCCESS(ascgen_utils::indirect_load::GetTemplateAxes(node, template_axes));182 GE_ASSERT_SUCCESS(ascgen_utils::indirect_load::GetTemplateAxes(node, template_axes));
392 outer_axis_ = template_axes.outer_axis;183 outer_axis_ = template_axes.outer_axis;
393- template_id_ = ::ascir::GetTemplateIdOrDefault(*node, ascir::TemplateId::kIndirectLoadSimt);184+ template_id_ = ::ascir::GetTemplateIdOrDefault(*node);
185+ GE_ASSERT_TRUE(
186+ template_id_ == ascir::TemplateId::kIndirectLoadSimd || template_id_ == ascir::TemplateId::kIndirectLoadSimt,
187+ "IndirectLoad node[%s] has invalid template id[%d].", node->GetNamePtr(), static_cast<int32_t>(template_id_));
394 GE_ASSERT_SUCCESS(ascgen_utils::indirect_load::GetTemplateLogicalView(node, logical_view_));188 GE_ASSERT_SUCCESS(ascgen_utils::indirect_load::GetTemplateLogicalView(node, logical_view_));
395 const int64_t rank = static_cast<int64_t>(logical_view_.data.axis_ids.size());189 const int64_t rank = static_cast<int64_t>(logical_view_.data.axis_ids.size());
396 GE_ASSERT_TRUE(axis >= -rank && axis < rank, "IndirectLoad axis is out of range.");190 GE_ASSERT_TRUE(axis >= -rank && axis < rank, "IndirectLoad axis is out of range.");
397 axis_ = axis < 0L ? axis + rank : axis;191 axis_ = axis < 0L ? axis + rank : axis;
398- if (template_id_ == ascir::TemplateId::kIndirectLoadSimd) {192+ if (template_id_ == ascir::TemplateId::kIndirectLoadSimt) {
399- const auto input_producer = ascgen_utils::indirect_load::GetInputProducer(node, 0UL);
400- if (af::ops::IsOps<af::ascir_op::VectorFunc>(input_producer)) {
401- GE_ASSERT_TRUE(ascgen_utils::indirect_load::GetTemplateRole(input_producer) ==
402- ascgen_utils::indirect_load::TemplateRole::kSimdInputPre,
403- "IndirectLoad SIMD input VectorFunc has no input-pre role.");
404- input_pre_vf_ = input_producer;
405- }
406- } else {
407 GE_ASSERT_SUCCESS(CollectSimtIndexPreNodes(node, index_pre_nodes_));193 GE_ASSERT_SUCCESS(CollectSimtIndexPreNodes(node, index_pre_nodes_));
408 GE_ASSERT_SUCCESS(CollectSimtOutputMetadata(node, output_post_nodes_, output_gm_tensor_, output_dtype_));194 GE_ASSERT_SUCCESS(CollectSimtOutputMetadata(node, output_post_nodes_, output_gm_tensor_, output_dtype_));
409 GE_ASSERT_TRUE(!node->outputs().empty(), "IndirectLoad SIMT node has no output tensor.");195 GE_ASSERT_TRUE(!node->outputs().empty(), "IndirectLoad SIMT node has no output tensor.");
@@ -444,27 +230,18 @@ Status IndirectLoadRegApiCall::GenerateFuncDefinition(const TPipe &tpipe, const
444 "IndirectLoad SIMT input and index rank must be equal.");230 "IndirectLoad SIMT input and index rank must be equal.");
445 std::string x_dtype;231 std::string x_dtype;
446 std::string index_dtype;232 std::string index_dtype;
447- std::string output_dtype;
448 GE_ASSERT_SUCCESS(Tensor::DtypeName(x.dtype, x_dtype));233 GE_ASSERT_SUCCESS(Tensor::DtypeName(x.dtype, x_dtype));
449 GE_ASSERT_SUCCESS(Tensor::DtypeName(index.dtype, index_dtype));234 GE_ASSERT_SUCCESS(Tensor::DtypeName(index.dtype, index_dtype));
450- GE_ASSERT_SUCCESS(Tensor::DtypeName(output_dtype_, output_dtype));
451 const std::string valid_node_name = ascgen_utils::GenValidName(node_name);235 const std::string valid_node_name = ascgen_utils::GenValidName(node_name);
452 const std::string index_transform = "IndirectLoadIndexTransform_" + valid_node_name;236 const std::string index_transform = "IndirectLoadIndexTransform_" + valid_node_name;
453 const std::string output_transform = "IndirectLoadOutputTransform_" + valid_node_name;237 const std::string output_transform = "IndirectLoadOutputTransform_" + valid_node_name;
454- const std::string kernel_name =
455- "IndirectLoadSimtKernel_" + ascgen_utils::GenValidName(graph_name) + "_" + valid_node_name;
456- 
457 GELOGI(238 GELOGI(
458- "[IndirectLoad] Generate SIMT wrapper for node[%s], rank[%zu], axis[%ld], index_pre_nodes[%zu], "239+ "[IndirectLoad] Generate SIMT transforms for node[%s], rank[%zu], axis[%ld], index_pre_nodes[%zu], "
459- "output_post_nodes[%zu], kernel[%s].",240+ "output_post_nodes[%zu].",
460- node_name.c_str(), x_info.sizes.size(), axis_, index_pre_nodes_.size(), output_post_nodes_.size(),241+ node_name.c_str(), x_info.sizes.size(), axis_, index_pre_nodes_.size(), output_post_nodes_.size());
461- kernel_name.c_str());
462 242 
463 GE_ASSERT_SUCCESS(GenerateTypedTransform(index_transform, index_dtype, index_pre_nodes_, ss));243 GE_ASSERT_SUCCESS(GenerateTypedTransform(index_transform, index_dtype, index_pre_nodes_, ss));
464 GE_ASSERT_SUCCESS(GenerateTypedTransform(output_transform, x_dtype, output_post_nodes_, ss));244 GE_ASSERT_SUCCESS(GenerateTypedTransform(output_transform, x_dtype, output_post_nodes_, ss));
465- const SimtKernelDefinitionParams params{kernel_name, x_dtype, index_dtype, output_dtype,
466- index_transform, output_transform, x_info.sizes.size(), axis_};
467- EmitSimtKernelDefinition(params, ss);
468 return af::SUCCESS;245 return af::SUCCESS;
469}246}
470 247 
@@ -473,18 +250,15 @@ Status IndirectLoadRegApiCall::Generate(const TPipe &tpipe, const std::vector<as
473 const std::vector<std::reference_wrapper<const Tensor>> &outputs,250 const std::vector<std::reference_wrapper<const Tensor>> &outputs,
474 std::string &result) const {251 std::string &result) const {
475 GE_ASSERT_TRUE(inputs.size() == 2U && outputs.size() == 1U, "IndirectLoad expects 2 inputs and 1 output.");252 GE_ASSERT_TRUE(inputs.size() == 2U && outputs.size() == 1U, "IndirectLoad expects 2 inputs and 1 output.");
253+ GE_ASSERT_TRUE(template_id_ == ascir::TemplateId::kIndirectLoadSimd,
254+ "IndirectLoad tensor-based Generate only supports SIMD.");
476 const LogicalTensorInfo x_info = BuildLogicalTensorInfo(logical_view_.data, tpipe);255 const LogicalTensorInfo x_info = BuildLogicalTensorInfo(logical_view_.data, tpipe);
477 const LogicalTensorInfo index_info = BuildLogicalTensorInfo(logical_view_.index, tpipe);256 const LogicalTensorInfo index_info = BuildLogicalTensorInfo(logical_view_.index, tpipe);
478 const LogicalTensorInfo output_info = BuildLogicalTensorInfo(logical_view_.output, tpipe);257 const LogicalTensorInfo output_info = BuildLogicalTensorInfo(logical_view_.output, tpipe);
479 GE_ASSERT_SUCCESS(CheckIndirectLoadShape(x_info, index_info, output_info));258 GE_ASSERT_SUCCESS(CheckIndirectLoadShape(x_info, index_info, output_info));
480 (void)RegisterBasicDumpParam(this->api_name_, inputs, outputs);259 (void)RegisterBasicDumpParam(this->api_name_, inputs, outputs);
481- 260+ GELOGI("[IndirectLoad] Generate SIMD API body for node[%s].", node_name.c_str());
482- if (template_id_ == ascir::TemplateId::kIndirectLoadSimd) {261+ return GenerateSimd(tpipe, current_axis, inputs, outputs, result);
483- GELOGI("[IndirectLoad] Generate SIMD API body for node[%s].", node_name.c_str());
484- return GenerateSimd(tpipe, current_axis, inputs, outputs, result);
485- }
486- GELOGI("[IndirectLoad] Generate SIMT API body for node[%s].", node_name.c_str());
487- return GenerateSimt(tpipe, current_axis, inputs, result);
488}262}
489 263 
490Status IndirectLoadRegApiCall::Generate(const TPipe &tpipe, const std::vector<ascir::AxisId> &current_axis,264Status IndirectLoadRegApiCall::Generate(const TPipe &tpipe, const std::vector<ascir::AxisId> &current_axis,
@@ -515,37 +289,22 @@ Status IndirectLoadRegApiCall::GenerateSimd(const TPipe &tpipe, const std::vecto
515 const LogicalTensorInfo index_info = BuildLogicalTensorInfo(logical_view_.index, tpipe);289 const LogicalTensorInfo index_info = BuildLogicalTensorInfo(logical_view_.index, tpipe);
516 const size_t axis_pos = static_cast<size_t>(axis_);290 const size_t axis_pos = static_cast<size_t>(axis_);
517 std::string x_dtype_name;291 std::string x_dtype_name;
518- std::string y_dtype_name;292+ std::string index_dtype_name;
519- Tensor::DtypeName(x.dtype, x_dtype_name);293+ GE_ASSERT_SUCCESS(Tensor::DtypeName(x.dtype, x_dtype_name));
520- Tensor::DtypeName(y.dtype, y_dtype_name);294+ GE_ASSERT_SUCCESS(Tensor::DtypeName(index.dtype, index_dtype_name));
521- const bool has_input_pre = input_pre_vf_ != nullptr;295+ GE_ASSERT_TRUE(x.dtype == y.dtype, "IndirectLoad SIMD input/output dtype must match after input preprocess.");
522- GE_ASSERT_TRUE(has_input_pre || x.dtype == y.dtype,296+ GELOGD("[IndirectLoad] Generate SIMD body for node[%s], rank[%zu], axis[%zu].", node_name.c_str(),
523- "IndirectLoad SIMD without input preprocess requires matching input/output dtype.");297+ x_info.sizes.size(), axis_pos);
524- GELOGD("[IndirectLoad] Generate SIMD body for node[%s], rank[%zu], axis[%zu], has_input_pre[%d].", node_name.c_str(),
525- x_info.sizes.size(), axis_pos, static_cast<int32_t>(has_input_pre));
526 298 
527- const SimdCodegenParams params{
528- x,
529- index,
530- y,
531- input_pre_vf_,
532- axis_pos,
533- x_dtype_name,
534- y_dtype_name,
535- tpipe.tmp_buf.name + "_" + std::to_string(tmp_iter->second),
536- tpipe.tiler.Offset(current_axis, y.axis, y.axis_strides),
537- tpipe.tiler.Size(index_info.sizes[axis_pos]),
538- tpipe.tiler.Size(x_info.sizes[axis_pos]),
539- tpipe.tiler.Size(index_info.strides[axis_pos]),
540- tpipe.tiler.Size(x_info.strides[axis_pos]),
541- JoinSizeExprs(std::vector<ascir::SizeExpr>(index_info.sizes.begin(), index_info.sizes.begin() + axis_pos), tpipe),
542- JoinSizeExprs(std::vector<ascir::SizeExpr>(x_info.strides.begin(), x_info.strides.begin() + axis_pos), tpipe)};
543 std::stringstream ss;299 std::stringstream ss;
544 ss << "// IndirectLoad SIMD" << std::endl;300 ss << "// IndirectLoad SIMD" << std::endl;
545 ss << "{" << std::endl;301 ss << "{" << std::endl;
546- EmitSimdWindowSetup(params, ss);302+ ss << " AscendC::IndirectLoadSimd<" << x_dtype_name << ", " << index_dtype_name << ", " << x_info.sizes.size()
547- EmitSimdInputWindow(params, ss);303+ << ", " << axis_ << ">(" << std::endl;
548- EmitSimdGather(params, ss);304+ ss << " " << x << ", " << index << ", " << y << ", " << tpipe.tmp_buf.name << "_" << tmp_iter->second << ", "
305+ << y.actual_size << ", " << tpipe.tiler.Offset(current_axis, y.axis, y.axis_strides) << ", "
306+ << tpipe.tiler.Size(x_info.sizes[axis_pos]) << ", " << JoinSizeExprs(index_info.sizes, tpipe) << ", "
307+ << JoinSizeExprs(x_info.strides, tpipe) << ");" << std::endl;
549 ss << "}" << std::endl;308 ss << "}" << std::endl;
550 result = ss.str();309 result = ss.str();
551 return af::SUCCESS;310 return af::SUCCESS;
@@ -561,8 +320,6 @@ Status IndirectLoadRegApiCall::GenerateSimt(const TPipe &tpipe, const std::vecto
561 GELOGD("[IndirectLoad] Generate SIMT body for node[%s], output_gm[%s].", node_name.c_str(),320 GELOGD("[IndirectLoad] Generate SIMT body for node[%s], output_gm[%s].", node_name.c_str(),
562 output_gm_tensor_.c_str());321 output_gm_tensor_.c_str());
563 322 
564- const std::string x_axis_size = tpipe.tiler.Size(x_info.sizes[axis_]);
565- 
566 std::string outer_tb_var;323 std::string outer_tb_var;
567 const bool has_outer_tb =324 const bool has_outer_tb =
568 FindCurrentAxisVar(tpipe, current_axis, Axis::Type::kAxisTypeBlockInner, outer_axis_, outer_tb_var);325 FindCurrentAxisVar(tpipe, current_axis, Axis::Type::kAxisTypeBlockInner, outer_axis_, outer_tb_var);
@@ -571,16 +328,23 @@ Status IndirectLoadRegApiCall::GenerateSimt(const TPipe &tpipe, const std::vecto
571 GE_ASSERT_SUCCESS(Tensor::DtypeName(x.dtype, x_dtype));328 GE_ASSERT_SUCCESS(Tensor::DtypeName(x.dtype, x_dtype));
572 GE_ASSERT_SUCCESS(Tensor::DtypeName(index.dtype, index_dtype));329 GE_ASSERT_SUCCESS(Tensor::DtypeName(index.dtype, index_dtype));
573 GE_ASSERT_SUCCESS(Tensor::DtypeName(output_dtype_, output_dtype));330 GE_ASSERT_SUCCESS(Tensor::DtypeName(output_dtype_, output_dtype));
574- const std::string kernel_name =331+ const std::string valid_node_name = ascgen_utils::GenValidName(node_name);
575- "IndirectLoadSimtKernel_" + ascgen_utils::GenValidName(graph_name) + "_" + ascgen_utils::GenValidName(node_name);332+ const std::string index_transform = "IndirectLoadIndexTransform_" + valid_node_name;
576- const std::string rank_sizes = std::string(", ") + JoinSizeExprs(index_info.sizes, tpipe);333+ const std::string output_transform = "IndirectLoadOutputTransform_" + valid_node_name;
577- const std::string x_strides = std::string(", ") + JoinSizeExprs(x_info.strides, tpipe);
578- const SimtCodegenParams params{x, index, output_gm_tensor_, x_dtype, index_dtype, output_dtype,
579- kernel_name, x_axis_size, rank_sizes, x_strides, outer_tb_var};
580 std::stringstream ss;334 std::stringstream ss;
581 ss << "// IndirectLoad SIMT" << std::endl;335 ss << "// IndirectLoad SIMT" << std::endl;
582 ss << "{" << std::endl;336 ss << "{" << std::endl;
583- EmitSimtLaunch(params, ss);337+ ss << " __gm__ " << x_dtype << " *x_ptr = (__gm__ " << x_dtype << " *)" << x << ".GetPhyAddr();" << std::endl;
338+ ss << " __gm__ " << index_dtype << " *index_ptr = (__gm__ " << index_dtype << " *)" << index << ".GetPhyAddr();"
339+ << std::endl;
340+ ss << " __gm__ " << output_dtype << " *y_ptr = (__gm__ " << output_dtype << " *)" << output_gm_tensor_
341+ << ".GetPhyAddr();" << std::endl;
342+ ss << " AscendC::IndirectLoadSimt<" << x_dtype << ", " << index_dtype << ", " << output_dtype << ", "
343+ << x_info.sizes.size() << ", " << axis_ << ", " << index_transform << ", " << output_transform << ">("
344+ << std::endl;
345+ ss << " x_ptr, index_ptr, y_ptr, static_cast<uint32_t>(" << outer_tb_var << "_loop_size), block_dim_offset, "
346+ << tpipe.tiler.Size(x_info.sizes[axis_]) << ", " << JoinSizeExprs(index_info.sizes, tpipe) << ", "
347+ << JoinSizeExprs(x_info.strides, tpipe) << ");" << std::endl;
584 ss << "}" << std::endl;348 ss << "}" << std::endl;
585 result = ss.str();349 result = ss.str();
586 return af::SUCCESS;350 return af::SUCCESS;
Mautofuse/v35/codegen/reg_api_call/reg_indirect_load_api_call.h+1-4
@@ -15,12 +15,10 @@
15 15 
16#include "codegen_kernel.h"16#include "codegen_kernel.h"
17#include "indirect_load_utils.h"17#include "indirect_load_utils.h"
18-#include "schedule_result.h"
19 18 
20namespace codegen {19namespace codegen {
21class IndirectLoadRegApiCall final : public ApiCall {20class IndirectLoadRegApiCall final : public ApiCall {
22 public:21 public:
23- using ApiCall::Generate;
24 explicit IndirectLoadRegApiCall(const std::string &api_name) : ApiCall(api_name) {}22 explicit IndirectLoadRegApiCall(const std::string &api_name) : ApiCall(api_name) {}
25 ~IndirectLoadRegApiCall() final = default;23 ~IndirectLoadRegApiCall() final = default;
26 Status GenerateFuncDefinition(const TPipe &tpipe, const Tiler &tiler, std::stringstream &ss) const override;24 Status GenerateFuncDefinition(const TPipe &tpipe, const Tiler &tiler, std::stringstream &ss) const override;
@@ -42,9 +40,8 @@ class IndirectLoadRegApiCall final : public ApiCall {
42 40 
43 int64_t axis_ = 0;41 int64_t axis_ = 0;
44 ascir::AxisId outer_axis_ = af::kIdNone;42 ascir::AxisId outer_axis_ = af::kIdNone;
45- ascir::TemplateId template_id_{ascir::TemplateId::kIndirectLoadSimt};43+ ascir::TemplateId template_id_{ascir::TemplateId::kDefault};
46 ascgen_utils::indirect_load::TemplateLogicalView logical_view_;44 ascgen_utils::indirect_load::TemplateLogicalView logical_view_;
47- af::AscNodePtr input_pre_vf_;
48 std::vector<af::AscNodePtr> index_pre_nodes_;45 std::vector<af::AscNodePtr> index_pre_nodes_;
49 std::vector<af::AscNodePtr> output_post_nodes_;46 std::vector<af::AscNodePtr> output_post_nodes_;
50 std::string output_gm_tensor_;47 std::string output_gm_tensor_;
Mautofuse/v35/codegen/simt_scalar_call/simt_scalar_emitter.cpp+6-16
@@ -12,20 +12,11 @@
12 12 
13#include "ascir_ops.h"13#include "ascir_ops.h"
14#include "ascir_ops_utils.h"14#include "ascir_ops_utils.h"
15+#include "v35/ascir/ascir_codegen_v2.h"
15#include "common/checker.h"16#include "common/checker.h"
16#include "common_utils.h"17#include "common_utils.h"
17-#include "indirect_load_utils.h"
18-#include "schedule_result.h"
19 18 
20namespace codegen {19namespace codegen {
21-bool CanEmitSimtScalar(const ascir::NodeView &node) {
22- if (node == nullptr) {
23- return false;
24- }
25- const auto impl = ascgen_utils::GetAscIrCodegenImpl(node->GetType());
26- return impl != nullptr && impl->IsSimtScalarSupported(*node);
27-}
28- 
29af::Status EmitSimtScalarExpr(const ascir::NodeView &node, const std::vector<std::string> &inputs, std::string &expr) {20af::Status EmitSimtScalarExpr(const ascir::NodeView &node, const std::vector<std::string> &inputs, std::string &expr) {
30 GE_ASSERT_NOTNULL(node, "SIMT scalar node is null.");21 GE_ASSERT_NOTNULL(node, "SIMT scalar node is null.");
31 GE_ASSERT_TRUE(!inputs.empty() && inputs.size() == node->inputs.Size(),22 GE_ASSERT_TRUE(!inputs.empty() && inputs.size() == node->inputs.Size(),
@@ -34,12 +25,11 @@ af::Status EmitSimtScalarExpr(const ascir::NodeView &node, const std::vector<std
34 const auto impl = ascgen_utils::GetAscIrCodegenImpl(node->GetType());25 const auto impl = ascgen_utils::GetAscIrCodegenImpl(node->GetType());
35 GE_ASSERT_NOTNULL(impl, "SIMT scalar codegen is not registered for node %s[%s].", node->GetTypePtr(),26 GE_ASSERT_NOTNULL(impl, "SIMT scalar codegen is not registered for node %s[%s].", node->GetTypePtr(),
36 node->GetNamePtr());27 node->GetNamePtr());
37- GE_ASSERT_TRUE(impl->IsSimtScalarSupported(*node), "SIMT scalar codegen is not supported for node %s[%s].",28+ const auto *v2_impl = dynamic_cast<af::ascir::AscIrCodegenV2 *>(impl.get());
29+ GE_ASSERT_NOTNULL(v2_impl, "SIMT scalar codegen for node %s[%s] is not a V2 implementation.", node->GetTypePtr(),
30+ node->GetNamePtr());
31+ GE_ASSERT_TRUE(v2_impl->IsSimtScalarSupported(*node), "SIMT scalar codegen is not supported for node %s[%s].",
38 node->GetTypePtr(), node->GetNamePtr());32 node->GetTypePtr(), node->GetNamePtr());
39- return impl->GenerateSimtScalarExpr(*node, inputs, expr);33+ return v2_impl->GenerateSimtScalarExpr(*node, inputs, expr);
40-}
41- 
42-bool IsSkippedApiEmitProcessNode(const ascir::NodeView &node) {
43- return ascgen_utils::indirect_load::GetTemplateBehavior(std::dynamic_pointer_cast<af::AscNode>(node)).skips_api_emit;
44}34}
45} // namespace codegen35} // namespace codegen
Mautofuse/v35/codegen/simt_scalar_call/simt_scalar_emitter.h+0-2
@@ -16,9 +16,7 @@
16#include "ascir.h"16#include "ascir.h"
17 17 
18namespace codegen {18namespace codegen {
19-bool CanEmitSimtScalar(const ascir::NodeView &node);
20af::Status EmitSimtScalarExpr(const ascir::NodeView &node, const std::vector<std::string> &inputs, std::string &expr);19af::Status EmitSimtScalarExpr(const ascir::NodeView &node, const std::vector<std::string> &inputs, std::string &expr);
21-bool IsSkippedApiEmitProcessNode(const ascir::NodeView &node);
22} // namespace codegen20} // namespace codegen
23 21 
24#endif // __AUTOFUSE_SIMT_SCALAR_EMITTER_H__22#endif // __AUTOFUSE_SIMT_SCALAR_EMITTER_H__
Mautofuse/v35/codegen/vec_func_call/vec_func_call.cpp+0-16
@@ -320,22 +320,6 @@ Status VfCall::ParseSubGraph(const ascir::NodeView &vf_node, const ascir::ImplGr
320}320}
321 321 
322Status VfCall::ParseInputOutputInfo(const TPipe &tpipe) const {322Status VfCall::ParseInputOutputInfo(const TPipe &tpipe) const {
323- if (inputs.empty() && skip_api_emit) {
324- for (const auto *input : node->inputs()) {
325- GE_ASSERT_NOTNULL(input, "Skipped VfCall node %s has null input.", node->GetNamePtr());
326- auto tensor_ptr = tpipe.GetTensor(input->attr.mem.tensor_id);
327- GE_CHK_BOOL_RET_STATUS(tensor_ptr != nullptr, af::FAILED, "Check[Param] tensor_ptr is nullptr");
328- ub_inputs_.emplace_back(*tensor_ptr);
329- }
330- }
331- if (outputs.empty() && skip_api_emit) {
332- for (const auto *output : node->outputs()) {
333- GE_ASSERT_NOTNULL(output, "Skipped VfCall node %s has null output.", node->GetNamePtr());
334- auto tensor_ptr = tpipe.GetTensor(output->attr.mem.tensor_id);
335- GE_CHK_BOOL_RET_STATUS(tensor_ptr != nullptr, af::FAILED, "Check[Param] tensor_ptr is nullptr");
336- ub_outputs_.emplace_back(*tensor_ptr);
337- }
338- }
339 for (const auto &in : inputs) {323 for (const auto &in : inputs) {
340 auto tensor_ptr = tpipe.GetTensor(in->id);324 auto tensor_ptr = tpipe.GetTensor(in->id);
341 GE_CHK_BOOL_RET_STATUS(tensor_ptr != nullptr, af::FAILED, "Check[Param] tensor_ptr is nullptr");325 GE_CHK_BOOL_RET_STATUS(tensor_ptr != nullptr, af::FAILED, "Check[Param] tensor_ptr is nullptr");
Mautofuse/v35/optimize/partition/vector_func_partitioner.cpp+2-3
@@ -1169,9 +1169,8 @@ af::Status VectorFuncPartitioner::ModifySubgraphAttrs(af::AscGraph &vf_graph) {
1169af::Status VectorFuncPartitioner::BuildSubgraphs() {1169af::Status VectorFuncPartitioner::BuildSubgraphs() {
1170 for (const auto &cluster : cluster_dict_.GetAllClusters()) {1170 for (const auto &cluster : cluster_dict_.GetAllClusters()) {
1171 // 不包含隐式广播的单节点不融合1171 // 不包含隐式广播的单节点不融合
1172- const bool force_vf = ascgen_utils::indirect_load::ShouldApplyInputInnerVectorization(cluster->Nodes().front());1172+ if (!cluster->meta_data_.enable_vf ||
1173- if (!cluster->meta_data_.enable_vf || (cluster->Nodes().size() < kMinVfNodesNum && !force_vf &&1173+ (cluster->Nodes().size() < kMinVfNodesNum && !ascgen_utils::IsNodeContainsBrcInline(cluster->Nodes().back()))) {
1174- !ascgen_utils::IsNodeContainsBrcInline(cluster->Nodes().back()))) {
1175 continue;1174 continue;
1176 }1175 }
1177 // 1. create vf op && subgraph1176 // 1. create vf op && subgraph
Mautofuse/v35/optimize/template/nddma_template.cpp+1-1
@@ -207,7 +207,7 @@ af::Status NddmaTemplate::ProcessSliceToNddma(const af::AscNodePtr &node_load, b
207 GE_CHECK_NOTNULL(node_load->GetOpDesc());207 GE_CHECK_NOTNULL(node_load->GetOpDesc());
208 208 
209 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node_load);209 const auto indirect_load_behavior = ascgen_utils::indirect_load::GetTemplateBehavior(node_load);
210- if (indirect_load_behavior.uses_direct_gm_pipeline || indirect_load_behavior.skips_api_emit) {210+ if (indirect_load_behavior.uses_direct_gm_pipeline) {
211 return af::SUCCESS;211 return af::SUCCESS;
212 }212 }
213 213 
Mscripts/test/run_autofuse_test.sh+22-7
@@ -771,13 +771,28 @@ build_backend() {
771 load_gather_split_b_t_abs_store_test_e2e_v2 \771 load_gather_split_b_t_abs_store_test_e2e_v2 \
772 load_gather_tail_split_b_t_abs_store_test_e2e_v2 \772 load_gather_tail_split_b_t_abs_store_test_e2e_v2 \
773 load_gather_one_axis_split_b_t_abs_store_test_e2e_v2 \773 load_gather_one_axis_split_b_t_abs_store_test_e2e_v2 \
774- indirect_load_rank2_axis1_simt_e2e_v2 \774+ indirect_load_rank2_axis1_simt_e2e_v2 \
775- indirect_load_rank2_axis1_simd_e2e_v2 \775+ indirect_load_rank2_axis1_simd_e2e_v2 \
776- indirect_load_rank3_axis1_simt_e2e_v2 \776+ indirect_load_rank2_axisneg2_simt_e2e_v2 \
777- indirect_load_rank3_axis1_simd_e2e_v2 \777+ indirect_load_rank2_axisneg2_simd_e2e_v2 \
778- indirect_load_rank4_axis2_simt_e2e_v2 \778+ indirect_load_rank3_axis1_simt_e2e_v2 \
779- indirect_load_rank4_axis2_simd_e2e_v2 \779+ indirect_load_rank3_axis1_simd_e2e_v2 \
780- load_where_x2_x3_is_ubscalar_store_test_e2e_v2 \780+ indirect_load_rank4_axis2_simt_e2e_v2 \
781+ indirect_load_rank4_axis2_simd_e2e_v2 \
782+ indirect_load_rank2_axis1_relu_exp2_pre_simd_e2e_v2 \
783+ indirect_load_rank2_axis0_relu_pre_expanded_simd_e2e_v2 \
784+ indirect_load_rank2_axis0_relu_relu_pre_expanded_simd_e2e_v2 \
785+ indirect_load_rank3_axis2_exp2_pre_expanded_simd_e2e_v2 \
786+ indirect_load_rank4_axis1_relu_exp2_pre_expanded_simd_e2e_v2 \
787+ indirect_load_rank3_axis2_mixed_vf_pre_expanded_simd_e2e_v2 \
788+ indirect_load_rank2_axis1_large_simt_e2e_v2 \
789+ indirect_load_rank2_axis1_float_int64_simd_e2e_v2 \
790+ indirect_load_rank2_axis1_modified_bessel_k0_simd_only_e2e_v2 \
791+ indirect_load_rank2_axis1_precision_cast_simt_e2e_v2 \
792+ indirect_load_rank2_axis1_precision_cast_simd_e2e_v2 \
793+ indirect_load_rank2_axis1_bf16_int64_simt_e2e_v2 \
794+ indirect_load_rank2_axis1_uint32_int32_simd_e2e_v2 \
795+ load_where_x2_x3_is_ubscalar_store_test_e2e_v2 \
781 gather_reduce_store_test_e2e_v2 \796 gather_reduce_store_test_e2e_v2 \
782 load_where_store_test_e2e_v2 \797 load_where_store_test_e2e_v2 \
783 load_where_x2_is_ubscalar_store_test_e2e_v2 \798 load_where_x2_is_ubscalar_store_test_e2e_v2 \