refactor rewrite_tensorize

Author: masahi

src/meta_schedule/postproc/rewrite_tensorize.cc

@@ -29,12 +29,7 @@ using tir::LoopRV;
 
 void ApplyTensorization(const tir::Schedule& sch, const String& func_name,
                         const tir::PrimFuncNode* func, bool vectorize_init_loop) {
-  struct RewriteJob {
-    std::string block_name;
-    std::function<void(tir::BlockRV)> postproc_fn;
-  };
-
-  std::vector<RewriteJob> jobs;
+  std::vector<std::pair<std::string, std::function<void(tir::BlockRV)>>> jobs;
 
   tir::PostOrderVisit(func->body, [=, &jobs](const ObjectRef& obj) -> bool {
     if (const auto* block = obj.as<tir::BlockNode>()) {
@@ -43,36 +38,34 @@ void ApplyTensorization(const tir::Schedule& sch, const String& func_name,
               tir::GetAnn<String>(block_sref, tir::attr::meta_schedule_auto_tensorize)) {
         std::string block_name = block_sref->StmtAs<tir::BlockNode>()->name_hint;
         if (block_name.find("init") == std::string::npos) {
-          jobs.push_back({block_name, [sch, intrin_name](tir::BlockRV block) {
-                            sch->Tensorize(block, intrin_name.value());
-                          }});
+          jobs.emplace_back(block_name, [sch, intrin_name](tir::BlockRV block) {
+            sch->Tensorize(block, intrin_name.value());
+          });
         } else if (vectorize_init_loop) {
-          jobs.push_back({block_name, [sch](tir::BlockRV block) {
-                            Array<BlockRV> child_blocks = sch->GetChildBlocks(block);
-                            ICHECK(child_blocks.size() == 1);
-                            Array<LoopRV> init_loops = sch->GetLoops(child_blocks[0]);
-                            ICHECK(init_loops.size() == 1);
-                            sch->Vectorize(init_loops[0]);
-                          }});
+          jobs.emplace_back(block_name, [sch](tir::BlockRV block) {
+            Array<BlockRV> child_blocks = sch->GetChildBlocks(block);
+            ICHECK(child_blocks.size() == 1);
+            Array<LoopRV> init_loops = sch->GetLoops(child_blocks[0]);
+            ICHECK(init_loops.size() == 1);
+            sch->Vectorize(init_loops[0]);
+          });
         }
       }
     }
     return true;
   });
 
-  for (auto job : jobs) {
-    tir::BlockRV block = sch->GetBlock(job.block_name, func_name);
+  for (auto kv : jobs) {
+    tir::BlockRV block = sch->GetBlock(kv.first, func_name);
     sch->Unannotate(block, tir::attr::meta_schedule_auto_tensorize);
-    job.postproc_fn(block);
+    kv.second(block);
   }
 }
 
 class RewriteTensorizeNode : public PostprocNode {
  public:
-  // Inherited from PostprocNode
   void InitializeWithTuneContext(const TuneContext& context) final {}
 
-  // Inherited from PostprocNode
   bool Apply(const tir::Schedule& sch) final;
 
   void VisitAttrs(tvm::AttrVisitor* v) {}