jroesch commented on a change in pull request #10561:
URL: https://github.com/apache/tvm/pull/10561#discussion_r824079170
##########
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:
While you are refactoring this code it might be good to better annotate
it with comments describing the pieces here as now this code (due to me and
Mark refactoring) has quite a few branches.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
