3535 * 4. This pass currently works for op_pattern kElemWise and kBroadcast.
3636 */
3737
38+ #include < tvm/node/attr_registry_map.h>
3839#include < tvm/relax/expr.h>
3940#include < tvm/relay/op_attr_types.h>
4041#include < tvm/tir/builtin.h>
4142#include < tvm/tir/function.h>
4243#include < tvm/tir/op.h>
4344#include < tvm/tir/stmt_functor.h>
4445#include < tvm/tir/transform.h>
45- # include < tvm/node/attr_registry_map.h >
46+
4647#include < optional>
4748
4849#include " ../../arith/constraint_extract.h"
@@ -78,15 +79,17 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer {
7879 /* This class analyzes the complete primfunc.
7980 It parses the buffer assumptions and eliminates the redundant branch
8081 introduced due to layout specific padding by leveraging from buffer assumptions.
81- On eliminating the branch there are more opportunities to vectorize the code and improve performance.
82+ On eliminating the branch there are more opportunities to vectorize the code
83+ and improve performance.
8284
8385 Example:
8486 -------------
8587 Prim Func Before :
8688 for (...)
8789 T.assume( assume_condition or A[i] == 0 )
8890 for (...)
89- out = T.if_then_else(if_then_else_condition, 0, function(A)) # here function(A) is some function on Var A
91+ out = T.if_then_else(if_then_else_condition, 0, function(A))
92+ # here function(A) is some function on Var A
9093
9194 Prim Func After :
9295 for (...)
@@ -96,12 +99,13 @@ class ParseAssumeAndOvercompute : public IRMutatorWithAnalyzer {
9699 --------------
97100 # High-level implementation details :
98101 1. The pass parses the assume statement and stores the relevant information.
99- 2. The pass tries to evaluate the then_clause and else_clause in then_condition_context and else_condition_context.
102+ 2. The pass tries to evaluate the then_clause and else_clause in then_condition_context
103+ and else_condition_context.
100104 It checks if the context of the assume statement (for condition indices and
101- assume_condition) is same as the context of the if_then_else statement (for condition indices and
102- if_then_else condition). If context is same and the expression inside if_then_else statement is a function of the
103- buffer assumption (eg A in above example), then the pass substitutes the value from the buffer assumption and
104- simplifies the expression .
105+ assume_condition) is same as the context of the if_then_else statement (for condition indices
106+ and if_then_else condition). If context is same and the expression inside if_then_else statement
107+ is a function of the buffer assumption (eg A in above example),
108+ then the pass substitutes the value from the buffer assumption and simplifies the expression.
105109 3. The pass then checks if then_clause and else_clause evaluate to same value.
106110 If yes, then return the else_clause if we are in the then_condition_context (since then_clause
107111 will be true in this context and if else_clause is also evaluating to true then we can directly
0 commit comments