masahi commented on a change in pull request #10561:
URL: https://github.com/apache/tvm/pull/10561#discussion_r824096015
##########
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:
Yeah, this function is not only used for ordinary lowering purposes, but
also abused for meta schedule task extraction or "apply history best" depending
on where it is called...
`Optional<ObjectRef> opt_mod_or_base_func =
meta_schedule::MetaScheduleContext::QueryInsideWithScope` is very opaque since
(1) we don't know if it returns anything and (2) we don't what concrete value
it returns.
After my task extraction refactor, `QueryInsideWithScope` can always return
a non-null, concrete `Schedule` object, since the `None` case is used only for
task extraction currently. cc @junrushao1994
--
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]
