Access Matrix

include/tvm/auto_scheduler/compute_dag.h

@@ -262,6 +262,13 @@ class ComputeDAG : public ObjectRef {
    */
   String PrintStepsAsPython(const Array<Step>& transform_steps) const;
 
+  /*!
+   * \brief Compute the access matrix of compute DAG.
+   * \param enabled Enables feature extraction.
+   * \return The flattened access matrix.
+   */
+  std::vector<int> ComputeAccessMatrix(bool enabled = true) const;
+
   /*!
    * \brief Print the compute DAG to a string. This is also used to generate the ComputeDAG hash.
    * \param simple_mode Simple mode will only include the op names and brief compute.

include/tvm/auto_scheduler/feature.h

@@ -46,16 +46,17 @@ namespace auto_scheduler {
  * \param cache_line_size The size of cache line in bytes
  * \param max_n_bufs The maximum number of extracted buffers for one statement
  * \param ret The returned feature vector
+ * \param access_matrix Access matrix feature vector
  */
 void GetPerStoreFeature(const Stmt& stmt, int cache_line_size, int max_n_bufs,
-                        std::vector<float>* ret);
+                        std::vector<float>* ret, std::vector<int>* access_matrix, bool access_matrix_enabled);
 
 /*
  * \brief Get the names of elements in the feature vector. Use this for debug and inspection.
  * \param max_n_bufs The maximum number of extracted buffers for one statement
  * \param ret The returned names.
  */
-void GetPerStoreFeatureName(int max_n_bufs, std::vector<std::string>* ret);
+void GetPerStoreFeatureName(int max_n_bufs, std::vector<std::string>* ret, bool access_matrix);
 
 /*!
  * \brief Get per-store feature from states of the same task
@@ -68,7 +69,7 @@ void GetPerStoreFeatureName(int max_n_bufs, std::vector<std::string>* ret);
  */
 void GetPerStoreFeaturesFromStates(const Array<State>& states, const SearchTask& task,
                                    int skip_first_n_feature_extraction, int max_n_bufs,
-                                   std::vector<std::vector<float> >* features);
+                                   std::vector<std::vector<float> >* features, bool access_matrix);
 
 /*!
  * \brief Get per-store feature from states of different tasks
@@ -81,7 +82,7 @@ void GetPerStoreFeaturesFromStates(const Array<State>& states, const SearchTask&
  */
 void GetPerStoreFeaturesFromStates(const Array<State>& states, const std::vector<SearchTask>& tasks,
                                    int skip_first_n_feature_extraction, int max_n_bufs,
-                                   std::vector<std::vector<float> >* features);
+                                   std::vector<std::vector<float> >* features, bool access_matrix);
 
 /*!
  * \brief Get per-store features from a log file
@@ -96,7 +97,7 @@ void GetPerStoreFeaturesFromStates(const Array<State>& states, const std::vector
 void GetPerStoreFeaturesFromFile(const std::string& filename, int max_lines, int max_n_bufs,
                                  std::vector<std::vector<float> >* features,
                                  std::vector<float>* normalized_throughputs,
-                                 std::vector<int>* task_ids);
+                                 std::vector<int>* task_ids, bool access_matrix);
 
 /*!
  * \brief Get per-store features from measurement input/result pairs
@@ -114,7 +115,7 @@ void GetPerStoreFeaturesFromMeasurePairs(const Array<MeasureInput>& inputs,
                                          int skip_first_n_feature_extraction, int max_n_bufs,
                                          std::vector<std::vector<float> >* features,
                                          std::vector<float>* normalized_throughputs,
-                                         std::vector<int>* task_ids);
+                                         std::vector<int>* task_ids, bool access_matrix);
 
 }  // namespace auto_scheduler
 }  // namespace tvm

python/tvm/auto_scheduler/cost_model/lgbm_model.py

@@ -113,7 +113,8 @@ def __init__(
         use_gpu=False,
         few_shot_learning="base_only",
         verbose_eval=25,
-        seed=None):
+        seed=None,
+        access_matrix=True):
 
         global lgbm
         try:
@@ -135,18 +136,21 @@ def __init__(
         self.few_shot_learning = few_shot_learning
         self.verbose_eval = verbose_eval
         self.workload_embed_dict = dict()
+        self.access_matrix = access_matrix
 
         # lgbm params
         if params is None:
             self.lgbm_params = {
                 'boosting_type': 'gbdt',
-                'num_leaves': 72,
-                'learning_rate': 0.1632095,
-                'feature_fraction': 0.84375,
-                'bagging_fraction': 0.89435,
-                'bagging_freq': 4,
+                'num_leaves': 31,
+                'learning_rate': 0.05,
+                'feature_fraction': 0.9,
+                'bagging_fraction': 0.8,
+                'bagging_freq': 5,
                 'verbose': 0,
-                'min_sum_hessian_in_leaf': 4,
+                'min_child_weight': 2,
+                'in_sum_in_hessian': 0,
+                'min_data_in_leaf': 0
             }
         else:
             self.lgbm_params = params 
@@ -189,7 +193,7 @@ def fit_local(self, train_set, valid_set=None):
         elif self.few_shot_learning == "plus_per_task":
             base_preds = self._predict_a_dataset(self.base_model, train_set)
             for task in train_set.tasks():
-                diff_train_set = Dataset()
+                diff_train_set = Dataset(self.access_matrix)
                 diff_train_set.load_task_data(
                     task,
                     train_set.features[task],
@@ -249,7 +253,7 @@ def _fit_a_model(self, train_set, valid_set=None, valid_train_set=None):
             verbose_eval=self.verbose_eval  
         )
 
-        feature_names = list(get_per_store_feature_names()) + ['max', 'min', 'add', 
+        feature_names = list(get_per_store_feature_names(self.access_matrix)) + ['max', 'min', 'add', 
             'Conv2dOutput', 'conv2d_winograd', 'DepthwiseConv2d',
             'dense', 'softmax', 'compute(b, i, j)']
         feature_importances = bst.feature_importance()
@@ -285,7 +289,7 @@ def register_new_task(self, task):
 
     def make_diff_set(self, base_model, dataset):
         base_preds = self._predict_a_dataset(base_model, dataset)
-        diff_set = Dataset()
+        diff_set = Dataset(self.access_matrix)
         for task in dataset.tasks():
             diff_set.load_task_data(
                 task,
@@ -365,15 +369,16 @@ def save(self, filename):
 class LGBModel(PythonBasedModel):
     """The wrapper of LGBModelInternal. So we can use it in end-to-end search."""
     def __init__(self, few_shot_learning="base_only", verbose_eval=25,
-                 num_warmup_sample=100, seed=None, disable_update=False):
+                 num_warmup_sample=100, seed=None, disable_update=False, access_matrix=True):
         super().__init__()
 
         self.num_warmup_sample = num_warmup_sample
         self.disable_update = disable_update
         self.model = LGBModelInternal(few_shot_learning=few_shot_learning,
                                       verbose_eval=verbose_eval,
-                                      seed=seed)
-        self.dataset = Dataset()
+                                      seed=seed, access_matrix=access_matrix)
+        self.dataset = Dataset(self.access_matrix)
+        self.access_matrix = access_matrix
 
     def update(self, inputs, results):
         if self.disable_update or len(inputs) <= 0:
@@ -384,7 +389,7 @@ def update(self, inputs, results):
         logger.info("LGBModel Training time: %.2f s", time.time() - tic)
 
     def predict(self, task, states):
-        features = get_per_store_features_from_states(states, task)
+        features = get_per_store_features_from_states(states, task, access_matrix=self.access_matrix)
         if self.model is not None and len(self.dataset) > self.num_warmup_sample:
             learning_task = LearningTask(task.workload_key, str(task.target))
             eval_dataset = Dataset.create_one_task(learning_task, features, None)

python/tvm/auto_scheduler/cost_model/mlp_model.py

@@ -320,7 +320,7 @@ def moving_average(average, update):
 
 class MLPModelInternal:
     def __init__(self, device=None, few_shot_learning="base_only", use_workload_embedding=True, use_target_embedding=False,
-                 loss_type='lambdaRankLoss'):
+                 loss_type='lambdaRankLoss', access_matrix=True):
         if device is None:
             if torch.cuda.device_count():
                 device = 'cuda:0'
@@ -330,7 +330,7 @@ def __init__(self, device=None, few_shot_learning="base_only", use_workload_embe
         # Common parameters
         self.net_params = {
             "type": "SegmentSumMLP",
-            "in_dim": 164 + (10 if use_workload_embedding else 0),  
+            "in_dim": 164 + (250 if access_matrix else 0) + (10 if use_workload_embedding else 0),  
             "hidden_dim": 256,
             "out_dim": 1,
         }
@@ -429,7 +429,7 @@ def fit_local(self, train_set, valid_set=None):
             self.loss_func = torch.nn.MSELoss()
             self.net_params['add_sigmoid'] = True
             base_preds = self._predict_a_dataset(self.base_model, train_set)
-            diff_train_set = Dataset()
+            diff_train_set = Dataset(self.access_matrix)
             for task in train_set.tasks():
                 diff_train_set.load_task_data(
                     task,
@@ -439,7 +439,7 @@ def fit_local(self, train_set, valid_set=None):
 
             if valid_set:
                 base_preds = self._predict_a_dataset(self.base_model, valid_set)
-                diff_valid_set = Dataset()
+                diff_valid_set = Dataset(self.access_matrix)
                 for task in valid_set.tasks():
                     diff_valid_set.load_task_data(
                         task,
@@ -456,7 +456,7 @@ def fit_local(self, train_set, valid_set=None):
         elif self.few_shot_learning == "plus_per_task":
             base_preds = self._predict_a_dataset(self.base_model, train_set)
             for task in train_set.tasks():
-                diff_train_set = Dataset()
+                diff_train_set = Dataset(self.access_matrix)
                 diff_train_set.load_task_data(
                     task,
                     train_set.features[task],
@@ -785,12 +785,13 @@ def find_class(self, module, name):
 class MLPModel(PythonBasedModel):
     """The wrapper of MLPModelInternal. So we can use it in end-to-end search."""
 
-    def __init__(self, few_shot_learning="base_only", disable_update=False):
+    def __init__(self, few_shot_learning="base_only", disable_update=False, access_matrix=True):
         super().__init__()
 
         self.disable_update = disable_update
-        self.model = MLPModelInternal(few_shot_learning=few_shot_learning)
-        self.dataset = Dataset()
+        self.model = MLPModelInternal(few_shot_learning=few_shot_learning, access_matrix=access_matrix)
+        self.dataset = Dataset(self.access_matrix)
+        self.access_matrix = access_matrix
 
     def update(self, inputs, results):
         if self.disable_update or len(inputs) <= 0:
@@ -801,7 +802,7 @@ def update(self, inputs, results):
         logger.info("MLPModel Training time: %.2f s", time.time() - tic)
 
     def predict(self, task, states):
-        features = get_per_store_features_from_states(states, task)
+        features = get_per_store_features_from_states(states, task, access_matrix=self.access_matrix)
         if self.model is not None:
             learning_task = LearningTask(task.workload_key, str(task.target))
             eval_dataset = Dataset.create_one_task(learning_task, features, None)

python/tvm/auto_scheduler/cost_model/tabnet_model.py

@@ -784,7 +784,7 @@ def moving_average(average, update):
 
 class TabNetModelInternal:
     def __init__(self, use_gpu=True, device=None, few_shot_learning="base_only", use_workload_embedding=True, use_target_embedding=False,
-                 loss_type='lambdaRankLoss'):
+                 loss_type='lambdaRankLoss', access_matrix=True):
         print('tabnet')
         if device is None:
             if torch.cuda.device_count() and use_gpu:
@@ -795,7 +795,7 @@ def __init__(self, use_gpu=True, device=None, few_shot_learning="base_only", use
         # Common parameters
         self.net_params = {
             "type": "SegmentSumMLP",
-            "in_dim": 164 + (10 if use_workload_embedding else 0),
+            "in_dim": 164 + (250 if access_matrix else 0) + (10 if use_workload_embedding else 0),
             "hidden_dim": 256,
             "out_dim": 1,
         }
@@ -873,7 +873,7 @@ def fit_local(self, train_set, valid_set=None):
             self.loss_func = torch.nn.MSELoss()
             self.net_params['add_sigmoid'] = True
             base_preds = self._predict_a_dataset(self.base_model, train_set)
-            diff_train_set = Dataset()
+            diff_train_set = Dataset(self.access_matrix)
             for task in train_set.tasks():
                 diff_train_set.load_task_data(
                     task,
@@ -883,7 +883,7 @@ def fit_local(self, train_set, valid_set=None):
 
             if valid_set:
                 base_preds = self._predict_a_dataset(self.base_model, valid_set)
-                diff_valid_set = Dataset()
+                diff_valid_set = Dataset(self.access_matrix)
                 for task in valid_set.tasks():
                     diff_valid_set.load_task_data(
                         task,
@@ -900,7 +900,7 @@ def fit_local(self, train_set, valid_set=None):
         elif self.few_shot_learning == "plus_per_task":
             base_preds = self._predict_a_dataset(self.base_model, train_set)
             for task in train_set.tasks():
-                diff_train_set = Dataset()
+                diff_train_set = Dataset(self.access_matrix)
                 diff_train_set.load_task_data(
                     task,
                     train_set.features[task],
@@ -1125,12 +1125,13 @@ def find_class(self, module, name):
 class TabNetModel(PythonBasedModel):
     """The wrapper of TabNetModelInternal. So we can use it in end-to-end search."""
 
-    def __init__(self, few_shot_learning="base_only", disable_update=False):
+    def __init__(self, few_shot_learning="base_only", disable_update=False, access_matrix=True):
         super().__init__()
 
         self.disable_update = disable_update
-        self.model = TabNetModelInternal(few_shot_learning=few_shot_learning)
-        self.dataset = Dataset()
+        self.model = TabNetModelInternal(few_shot_learning=few_shot_learning, access_matrix=access_matrix)
+        self.dataset = Dataset(self.access_matrix)
+        self.access_matrix = access_matrix
 
     def update(self, inputs, results):
         if self.disable_update or len(inputs) <= 0:
@@ -1141,7 +1142,7 @@ def update(self, inputs, results):
         logger.info("TabNetModel Training time: %.2f s", time.time() - tic)
 
     def predict(self, task, states):
-        features = get_per_store_features_from_states(states, task)
+        features = get_per_store_features_from_states(states, task, access_matrix=self.access_matrix)
         if self.model is not None:
             learning_task = LearningTask(task.workload_key, str(task.target))
             eval_dataset = Dataset.create_one_task(learning_task, features, None)