[GitHub] [tvm] junrushao1994 commented on a change in pull request #10561: [TECompiler] Decouple TE compute and schedule lowering in ScheduleBuilder

[GitBox <gi...@apache.org>]

junrushao1994 commented on a change in pull request #10561:
URL: https://github.com/apache/tvm/pull/10561#discussion_r824284592



##########
File path: src/relay/backend/te_compiler_cache.cc
##########
@@ -45,6 +45,8 @@
 #include "../../te/operation/create_primfunc.h"
 #include "../op/memory/memory.h"
 #include "../transforms/pass_utils.h"
+#include "tvm/relay/op_strategy.h"
+#include "tvm/tir/function.h"

Review comment:
       nit: Maybe we can follow the convention using
   
   ```suggestion
   #include <tvm/relay/op_strategy.h>
   #include <tvm/tir/function.h>
   ```

##########
File path: src/relay/backend/te_compiler_cache.cc
##########
@@ -347,26 +270,127 @@ class ScheduleBuilder : public backend::MemoizedExprTranslator<Array<te::Tensor>
     return {tuple[op->index]};
   }
 
+ public:
+  // Additional outputs
+  Array<tvm::te::Tensor> fn_inputs_;
+  Array<te::Operation> scalars_;
+  std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors_;
+  std::string candidate_name_;
+  OpImplementation anchor_implementation_;
+
  private:
   tvm::Target target_;
-  Op anchor_op_;
-  Attrs anchor_attrs_;
-  int anchor_op_pattern_{0};
-  OpImplementation anchor_implementation_;
   std::ostringstream readable_name_stream_;
-  Array<te::Operation> scalars_;
-  std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors_;
-  bool use_auto_scheduler_;
-  bool use_meta_schedule_;
+  // Index of the global constants
+  static int const_index;
   // Cache device copy op for equivalence checking to reduce registry lookup
   // overhead for each invocation of call node when retrieving schedules.
   const Op& device_copy_op_;
-  bool create_schedule_;
-  // Index of the global constants
-  static int const_index;
 };
 
-int ScheduleBuilder::const_index = 0;
+int LowerToTECompute::const_index = 0;
+
+// Construct a schedule for a given Relay primitive function and target.
+class ScheduleBuilder : ExprVisitor {

Review comment:
       nit: always good to be explicit
   
   ```suggestion
   class ScheduleBuilder : public ExprVisitor {
   ```

##########
File path: src/relay/backend/te_compiler_cache.cc
##########
@@ -347,26 +270,127 @@ class ScheduleBuilder : public backend::MemoizedExprTranslator<Array<te::Tensor>
     return {tuple[op->index]};
   }
 
+ public:
+  // Additional outputs
+  Array<tvm::te::Tensor> fn_inputs_;
+  Array<te::Operation> scalars_;
+  std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors_;
+  std::string candidate_name_;
+  OpImplementation anchor_implementation_;
+
  private:
   tvm::Target target_;
-  Op anchor_op_;
-  Attrs anchor_attrs_;
-  int anchor_op_pattern_{0};
-  OpImplementation anchor_implementation_;
   std::ostringstream readable_name_stream_;
-  Array<te::Operation> scalars_;
-  std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors_;
-  bool use_auto_scheduler_;
-  bool use_meta_schedule_;
+  // Index of the global constants
+  static int const_index;
   // Cache device copy op for equivalence checking to reduce registry lookup
   // overhead for each invocation of call node when retrieving schedules.
   const Op& device_copy_op_;
-  bool create_schedule_;
-  // Index of the global constants
-  static int const_index;
 };
 
-int ScheduleBuilder::const_index = 0;
+int LowerToTECompute::const_index = 0;
+
+// Construct a schedule for a given Relay primitive function and target.
+class ScheduleBuilder : ExprVisitor {
+ public:
+  explicit ScheduleBuilder(Target target) : target_(target) {
+    // Whether to use auto_scheduler schedule.
+    use_auto_scheduler_ = backend::IsAutoSchedulerEnabled();
+  }
+
+  CachedFunc Create(const Function& relay_func, std::function<std::string(std::string)> renamer) {

Review comment:
       Per discussion with Masa, we decided to just clean up `meta_schedule/integration.cc` and make it clearer




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org