Merge pull request #4 from junrushao1994/feature/2022-08-15/group2-2

Extract leaf blocks only

Author: Kathryn-cat

src/meta_schedule/feature_extractor/per_block_feature.cc

@@ -26,7 +26,6 @@
 #include <unordered_set>
 #include <vector>
 
-#include "../utils.h"
 #include "./utils.h"
 
 namespace tvm {
@@ -618,7 +617,7 @@ void Feature::SubFeature::SetReuse(const LoopNest& loop_nest,
         extent = 1;
       }
       const IntVec& touch = buffer_touched_under_loop.at(loop_idx);
-      reuse_dis_iter = std::accumulate(touch.begin(), touch.end(), 1);
+      reuse_dis_iter = std::accumulate(touch.begin(), touch.end(), 0);
       reuse_dis_bytes = 0.0;
       int buffer_idx = -1;
       for (int64_t numel : touch) {
@@ -677,6 +676,7 @@ Feature::Feature(const BlockRealizeNode* realize, const LoopNest& loop_nest,
   for (SubFeature& feature : sub_features) {
     const BufferNode* buffer = feature.buffer_;
     int64_t numel = buffer_touched_under_loop.front().at(++buffer_idx);
+    LOG(INFO) << "feature: " << feature.buffer_->name << ", numel = " << numel;
     feature.SetFeature(loop_nest, cache_line_bytes, numel * buffer->dtype.bytes());
   }
   // Step 5. Calculate `for_touched_bytes`
@@ -999,15 +999,15 @@ class PerBlockFeatureCollector : private StmtVisitor {
     std::vector<Feature> result;
     result.reserve(collector.block_features_.size());
     for (const BlockRealizeNode* realize : collector.ordered_blocks_) {
+      LOG(INFO) << "Block: " << realize->block->name_hint;
       Feature& feature = collector.block_features_.at(realize);
-      if (feature.block_realize != nullptr) {
-        ICHECK(feature.group1);
-        ICHECK(feature.group2);
-        ICHECK(feature.group3);
-        ICHECK(feature.group4);
-        ICHECK(feature.group5);
-        result.push_back(std::move(feature));
-      }
+      ICHECK(feature.block_realize == realize);
+      ICHECK(feature.group1);
+      ICHECK(feature.group2);
+      ICHECK(feature.group3);
+      ICHECK(feature.group4);
+      ICHECK(feature.group5);
+      result.push_back(std::move(feature));
     }
     return result;
   }
@@ -1020,33 +1020,34 @@ class PerBlockFeatureCollector : private StmtVisitor {
         arith_intensity_curve_num_samples_(arith_intensity_curve_num_samples) {}
 
   void VisitStmt_(const BlockRealizeNode* realize) final {
-    if (!scopes_.empty()) {
-      ordered_blocks_.push_back(realize);
-    }
-    Feature& feature = block_features_[realize];
-    feature.block_realize = realize;
-    feature.group1 = std::make_unique<group1::Feature>(realize, loop_nest_, is_gpu_);
+    ordered_blocks_.push_back(realize);
+    int previous_num_blocks_visited = ++this->num_blocks_visited_;
     scopes_.push_back(realize);
     dfs_path_.push_back(realize);
     StmtVisitor::VisitStmt_(realize);
     dfs_path_.pop_back();
     scopes_.pop_back();
-    if (!scopes_.empty()) {
-      AddArithOpsToScope(&feature.group1->arith_ops);
-    }
-    IntVec for_touched_bytes;
-    feature.group2 = std::make_unique<group2::Feature>(realize, loop_nest_, cache_line_bytes_,
-                                                       scopes_.empty() ? nullptr : scopes_.back(),
-                                                       &for_touched_bytes, &analyzer_);
-    feature.group3 =
-        std::make_unique<group3::Feature>(arith_intensity_curve_num_samples_, loop_nest_,
-                                          for_touched_bytes, feature.group1->arith_ops);
-    feature.group4 = std::make_unique<group4::Feature>(loop_nest_, realize, &analyzer_,
-                                                       &alloc_buffer_outer_loops_);
-    feature.group5 = std::make_unique<group5::Feature>(loop_nest_);
-    // Erase the feature of the root block
-    if (scopes_.empty()) {
-      block_features_.erase(realize);
+    // only extract features for leaf blocks
+    if (previous_num_blocks_visited == this->num_blocks_visited_) {
+      IntVec for_touched_bytes;
+      Feature& feature = block_features_[realize];
+      feature.block_realize = realize;
+      if (feature.group1 == nullptr) {
+        feature.group1 = std::make_unique<group1::Feature>(realize, loop_nest_, is_gpu_);
+      }
+      feature.group2 = std::make_unique<group2::Feature>(realize, loop_nest_, cache_line_bytes_,
+                                                         scopes_.empty() ? nullptr : scopes_.back(),
+                                                         &for_touched_bytes, &analyzer_);
+      feature.group3 =
+          std::make_unique<group3::Feature>(arith_intensity_curve_num_samples_, loop_nest_,
+                                            for_touched_bytes, feature.group1->arith_ops);
+      feature.group4 = std::make_unique<group4::Feature>(loop_nest_, realize, &analyzer_,
+                                                         &alloc_buffer_outer_loops_);
+      feature.group5 = std::make_unique<group5::Feature>(loop_nest_);
+      block_features_.emplace(realize, Feature{});
+    } else {
+      ordered_blocks_.erase(
+          std::find(std::begin(ordered_blocks_), std::end(ordered_blocks_), realize));
     }
   }
 
@@ -1069,14 +1070,13 @@ class PerBlockFeatureCollector : private StmtVisitor {
     ICHECK(!scopes_.empty());
     group1::Feature::ArithOps arith_ops;
     arith_ops.AddExpr(store->value, loop_nest_.prod);
-    AddArithOpsToScope(&arith_ops);
-  }
-
-  void AddArithOpsToScope(group1::Feature::ArithOps* arith_ops) {
     const BlockRealizeNode* scope = scopes_.back();
-    // Add the arith_ops to the parent
-    group1::Feature::ArithOps& parent_arith_ops = block_features_[scope].group1->arith_ops;
-#define TVM_FEATURE_MATH_OP_ADD(Name) parent_arith_ops.Name += arith_ops->Name;
+    std::unique_ptr<group1::Feature>& feature = block_features_[scope].group1;
+    if (feature == nullptr) {
+      block_features_[scope].block_realize = scope;
+      feature = std::make_unique<group1::Feature>(scope, loop_nest_, is_gpu_);
+    }
+#define TVM_FEATURE_MATH_OP_ADD(Name) feature->arith_ops.Name += arith_ops.Name;
     TVM_FEATURE_MATH_OP_ADD(float_mad);
     TVM_FEATURE_MATH_OP_ADD(float_add_sub);
     TVM_FEATURE_MATH_OP_ADD(float_mul);
@@ -1097,6 +1097,7 @@ class PerBlockFeatureCollector : private StmtVisitor {
   }
 
   bool is_gpu_;
+  int num_blocks_visited_ = 0;
   int64_t cache_line_bytes_;
   int64_t arith_intensity_curve_num_samples_;
   arith::Analyzer analyzer_;
@@ -1134,6 +1135,7 @@ class PerBlockFeatureNode : public FeatureExtractorNode {
     using namespace tvm::tir::per_block_feature;
     static transform::Sequential passes = tir::transform::PassListForFeatureExtraction();
     mod = passes(std::move(mod));
+    LOG(INFO) << "mod =\n" << tir::AsTVMScript(mod);
     std::vector<Feature> features = PerBlockFeatureCollector::Collect(
         is_gpu, this->cache_line_bytes, this->arith_intensity_curve_num_samples, mod);
     int n_features = features.size();