csullivan commented on code in PR #11692:
URL: https://github.com/apache/tvm/pull/11692#discussion_r896008247
##########
src/meta_schedule/extracted_task.cc:
##########
@@ -32,6 +38,42 @@ ExtractedTask::ExtractedTask(String task_name, IRModule mod,
Target target,
data_ = n;
}
+Optional<tir::PrimFunc> DefaultTaskFilter(const Array<te::Tensor>& args) {
+ using namespace ::tvm::te;
+ std::vector<Tensor> stack;
+ std::unordered_set<const TensorNode*> visited;
+ for (const Tensor& v : args) {
+ for (const PrimExpr& e : v->shape) {
+ // Dynamic shape is not supported for now
+ if (!e->IsInstance<IntImmNode>()) {
+ return NullOpt;
+ }
+ }
+ if (!visited.count(v.get())) {
+ visited.insert(v.get());
+ stack.push_back(v);
+ }
+ }
+ while (!stack.empty()) {
+ Tensor tensor = stack.back();
+ stack.pop_back();
+ if (tensor->op->IsInstance<PlaceholderOpNode>()) {
+ // do nothing
+ } else if (tensor->op->IsInstance<ComputeOpNode>()) {
Review Comment:
```suggestion
} else if (tensor->op->IsInstance<ComputeOpNode>() ||
tensor->op->IsInstance<ExternOpNode>()) {
```
You'll need to check for extern ops here as well
([ref](https://github.com/apache/tvm/pull/11692/files#diff-a3afd47a39285ff4100ee476114ca7a33f06b5a0929d8ecf6623c53c0de01ad8L55)).
--
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]
