junrushao1994 commented on a change in pull request #10686:
URL: https://github.com/apache/tvm/pull/10686#discussion_r830509216
##########
File path: src/auto_scheduler/feature.cc
##########
@@ -724,6 +731,15 @@ class PerStoreFeatureExtractor : public StmtExprVisitor {
ExtractAllocationFeature(node);
}
+ void VisitStmt_(const LetStmtNode* node) final {
+ // TODO(tkonolige): add arithmetic counts from this statement to counts of
inner stores.
+ ana_.Bind(node->var, node->value);
+ ICHECK(variable_definition_stack_.size() > 0)
+ << "Variable definition out size of a for loop is not handled by
feature extraction";
+ variable_definition_stack_.back().push_back(std::make_tuple(node->var,
node->value));
+ StmtExprVisitor::VisitStmt_(node);
Review comment:
Do we want to pop after visiting?
##########
File path: include/tvm/arith/analyzer.h
##########
@@ -407,11 +407,13 @@ class TVM_DLL Analyzer {
* \brief Notify all the sub-analyzers that var
* is created and binded to expr.
*
- * Each var can only be binded once.
+ * Each var can only be bound once.
*
* \param var The variable.
* \param expr The expression we bind to.
- * \param allow_override Whether we allow overriding an existing var's
expression.
+ * \param allow_override Whether we allow overriding an existing var's
+ * expression. This option should not be used if there is any
dependency
Review comment:
Thanks for clarifying this!
##########
File path: src/auto_scheduler/feature.cc
##########
@@ -233,14 +233,9 @@ AnnotationPosType GetAnnotationPosEncoding(const Var& var,
const Array<PrimExpr>
}
}
-// Return the extent of a for loop
-int64_t GetLoopExtent(const ForNode* node) {
- auto pint = node->extent.as<IntImmNode>();
- if (pint != nullptr) {
- return pint->value;
- } else {
- return 1;
- }
+// Return the maximum extent of a for loop
+int64_t GetLoopExtent(const ForNode* node, const Analyzer& ana) {
+ return ana.const_int_bound(node->extent)->max_value;
Review comment:
The reason that we didn't use const-int-bound in feature extraction is
that it becomes less helpful when encountering TIR of symbolic shapes.
##########
File path: tests/python/unittest/test_auto_scheduler_feature.py
##########
@@ -220,12 +220,15 @@ def tir_matmul(
def test_primfunc():
- features = auto_scheduler.feature.named_features_from_primfunc(tir_matmul)
Review comment:
+1 for TQ's idea
--
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]
