Add block_filter attribute in ScheduleState to enable scheduling a su…

…bset of blocks. (apache#58)

include/tvm/tir/schedule/state.h

@@ -102,6 +102,10 @@ class ScheduleStateNode : public Object {
    * \sa ScheduleDebugMask
    */
   int debug_mask;
+  /*!
+   * \brief The filter to apply on blocks.
+   */
+  Optional<String> block_filter = NullOpt;
 
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("mod", &mod);
@@ -150,6 +154,14 @@ class ScheduleStateNode : public Object {
    * have block vars, since the affine flag depends on the outer scope of stmt.
    */
   TVM_DLL void UpdateScopeBlockInfo(const Stmt& stmt);
+  /*!
+   * \brief Set block filter.
+   */
+  TVM_DLL void SetBlockFilter(String name);
+  /*!
+   * \brief Unset block filter.
+   */
+  TVM_DLL void UnsetBlockFilter();
   /*!
    * \brief Get the BlockScope correpsonding to the sref of scope root block
    * \param scope_root The block sref to be retrieved

python/tvm/tir/schedule/schedule.py

@@ -198,6 +198,12 @@ def trace(self) -> Optional[Trace]:
         """Returns the internally maintained trace of scheduling program execution"""
         return _ffi_api.ScheduleGetTrace(self)  # type: ignore # pylint: disable=no-member
 
+    def set_block_filter(self, name: str) -> None:
+        _ffi_api.ScheduleSetBlockFilter(self.state, name)  # type: ignore # pylint: disable=no-member
+
+    def unset_block_filter(self) -> None:
+        _ffi_api.ScheduleUnsetBlockFilter(self.state)  # type: ignore # pylint: disable=no-member
+
     def copy(self) -> "Schedule":
         """Returns a copy of the schedule, including both the state and the symbol table,
         * guaranteeing that