Lunderberg commented on code in PR #16296:
URL: https://github.com/apache/tvm/pull/16296#discussion_r1439521746
##########
src/arith/ir_mutator_with_analyzer.cc:
##########
@@ -96,6 +124,11 @@ Stmt IRMutatorWithAnalyzer::VisitStmt_(const LetStmtNode*
op) {
Stmt IRMutatorWithAnalyzer::VisitStmt_(const IfThenElseNode* op) {
PrimExpr condition = this->VisitExpr(op->condition);
PrimExpr real_condition = condition;
+
+ if (!IsPureCondition(real_condition, GetRef<Stmt>(op))) {
+ return StmtExprMutator::VisitStmt_(op);
Review Comment:
I think the early-bail-out is a bit too broad. If a conditional depends
both on pure and impure conditions, the early return would prevent the analyzer
from using the pure condition. These mixed conditionals could be produced when
lowering (e.g. merging nested conditional blocks) even if the initial kernel
doesn't contain any, so we probably should handle it.
```python
@T.prim_func
def func(A: T.Buffer(16, "int32"), B: T.Buffer(16,"int32"), C: T.Buffer(16,
"int32")):
for i in range(16):
# This conditional provides a pure constraint `i>1` along with
# an impure constraint `A[i]>1`.
if i>1 and A[i] > 1:
A[i] = A[i] + 1
# The `A[i]-1>0` conditional should not be simplified.
B[i] = T.if_then_else(A[i] - 1 > 0, 0, 1)
# However, the `i-1>0` should still be simplfied to `True`
C[i] = T.if_then_else(i - 1 > 0, 0, 1)
```
Instead of an early-return, I think we can get the desired effect by
changing the value provided in the `ConstraintContext`. There's an existing
utility,
[`ExtractConstraints`](https://github.com/apache/tvm/blob/main/src/arith/constraint_extract.h#L62)
that would break up the `cond1 and cond2` conditions, so they could be
inspected separately. I'm thinking something like the following (and
analogously for the constraint on the else case).
```c++
// Break up `pure_cond1 && pure_cond2 && impure_cond`, providing
// `pure_cond1 && pure_cond2` to the analyzer.
PrimExpr if_case_constraint = [&]() -> PrimExpr {
PrimExpr constraint = Bool(true);
for(const auto& partial_constraint: ExtractConstraints(real_condition)) {
if(IsPureCondition(partial_constraint, GetRef<Stmt>(op))) {
constraint = constraint && partial_constraint;
}
}
return constraint;
};
```
##########
src/arith/ir_mutator_with_analyzer.cc:
##########
@@ -20,13 +20,41 @@
/*!
* \file tvm/arith/ir_mutator_with_analyzer.cc
*/
+#include <vector>
+
#include "ir_mutator_with_analyzer.h"
#include <tvm/arith/iter_affine_map.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/op.h>
namespace tvm {
+namespace tir {
+
+bool IsPureCondition(const PrimExpr& condition, const Stmt& scope) {
+ CallEffectKind side_effect = SideEffect(condition);
+ if (side_effect > CallEffectKind::kReadState) {
+ return false;
+ } else if (side_effect == CallEffectKind::kReadState) {
+ // when the condition maybe stateful, enter into the slow path
+ // to visit the subtree one more time and ensure the state is not mutated.
+ std::vector<const VarNode*> read_vars;
+ UsesVar(condition, [&read_vars](const VarNode* v) {
+ read_vars.push_back(v);
+ return true;
Review Comment:
By returning `true`, the `VarTouchVisitor` will return early and won't visit
the rest of the TIR tree. As a result, any read variables after the first
would be missing from `read_vars`.
##########
src/tir/analysis/var_touch.cc:
##########
@@ -45,12 +47,16 @@ class VarTouchVisitor : public StmtExprVisitor {
void VisitExpr_(const VarNode* op) final { Handle(op); }
void VisitStmt_(const BufferStoreNode* op) final {
- Handle(op->buffer->data.get());
+ if (touch_write_) {
Review Comment:
This looks like it's a repetition of the logic in `VarUseDefAnalyzer`, which
already distinguishes between variable access, reads of a buffer, and writes to
a buffer. I'd recommend using it rather than duplicating similar logic here.
##########
src/arith/ir_mutator_with_analyzer.cc:
##########
@@ -96,6 +124,11 @@ Stmt IRMutatorWithAnalyzer::VisitStmt_(const LetStmtNode*
op) {
Stmt IRMutatorWithAnalyzer::VisitStmt_(const IfThenElseNode* op) {
PrimExpr condition = this->VisitExpr(op->condition);
PrimExpr real_condition = condition;
+
Review Comment:
Can we add a comment that the `IsPureCondition` is intended to avoid
introducing a constraint that may be invalidated? On first read, I thought
this was to avoid applying a simplification to a conditional that could remove
side-effects.
##########
src/arith/ir_mutator_with_analyzer.cc:
##########
@@ -20,13 +20,41 @@
/*!
* \file tvm/arith/ir_mutator_with_analyzer.cc
*/
+#include <vector>
+
#include "ir_mutator_with_analyzer.h"
#include <tvm/arith/iter_affine_map.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/op.h>
namespace tvm {
+namespace tir {
+
+bool IsPureCondition(const PrimExpr& condition, const Stmt& scope) {
Review Comment:
For any impure conditional, it looks like we'd end up visiting the body
three times (`UsesVar`, `WritesVar`, and then the actual mutation). For a
nested conditional, the nested body would be visited several times.
--
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]
