[Arith][Refactor] Extract And/Or/Not handling from RewriteSimplifier

[Lunderberg]

Previously, rewrites of boolean expressions were handled directly within RewriteSimplifier visitors. This commit extracts the handling of And, Or, and Not nodes into an independent internal utility. The RewriteBooleanOperators utility can simplifies these nodes, assuming that the arguments to these expressions are already simplified. (e.g. Finding the negation of (a < b) without needing to recurse into a or b.) This will be used as part of additional simplifications required by buffer layout
padding.

Because no external functionality should be added or changed by this refactor, this commit does not introduce any additional unit tests.

[tqchen]

Thanks @Lunderberg . RewriteSimplifier is in the heart of a lot of things, so changes to this module should be done with extra amount of care. For example, in some of our cases we have recursive calls into simplify that further triggers other rewrite simplifications of the expressions within the boolean expr (TVM_TRY_RECURSIVE_REWRITE), so refactoring booleaning simplfication into a separate function may not retain that deep recursive behavior.

We can see the values of introducing boolean simplification as being described. For extra caution, it might be helpful to simply introduce a bit of duplication of rules and have an independent implementation here. Since SimplifyBooleanOps are used in more limited cases.

Additionally, is it OK to simply invoke the full simplification instead? It might be useful to discuss the usecase scenarios to see if we can also retrofit to some of existing simplifiers

[Lunderberg]

@tqchen Thank you, and I definitely agree that any changes to RewriteSimplifier should be done with care, because of how foundational it is, but also any improvements to it may yield improvements across the codebase. I didn't do so for this PR, because it isn't intended to alter any functionality, but any that do are going to be behind feature flags to avoid incidental breakage or unnecessary CPU usage.

For example, in some of our cases we have recursive calls into simplify that further triggers other rewrite simplifications of the expressions within the boolean expr (TVM_TRY_RECURSIVE_REWRITE), so refactoring booleaning simplfication into a separate function may not retain that deep recursive behavior.

Regarding the TVM_TRY_RECURSIVE_REWRITE functionality, the behavior was considered and should be preserved by this refactor. The recursive rewrites used for boolean expressions are to move a Not to be internal to Or/And, regardless of the depth of the or/and chain contained in the Not. The refactor allows a Not to be recursively propagated into the chain, using TVM_TRY_RECURSIVE_REWRITE, recursively visiting the And/Or as needed.

This should also improve the performance. On the main branch, the expression !(a && (b || c)), the expressions a is visited twice, and both b and c are visited three times, even though I think only the first visit is necessary. I wrote out and validated the sequence of visits to convince myself, shown below.

Sequence of visits on main branch

1. Visit `!(a && (b || c))` 1. Visit `a && (b || c)` 1. Visit `a` 2. Visit `b || c` 1. Visit `b` 2. Visit `c` 3. Return `b || c` 3. Return `a && (b || c)` 2. Recursive rewrite, `!(a && (b || c))` becomes `(!a) || (!(b||c))` 3. Visit `(!a) || (!(b||c))` 1. Visit of `!a` 1. Visit `a` 2. Return `!a` 2. Visit `(!(b||c))` 1. Visit `b||c` 1. Visit `b` 2. Visit `c` 3. Return `b||c` 2. Recursive rewrite, `!(b||c)` becomes `(!b) && (!c)` 3. Visit `(!b) && (!c)` 1. Visit `!b` 1. Visit `b` 2. Return `!b` 2. Visit `!c` 1. Visit `c` 2. Return `!c` 3. Return `(!b) && (!c)` 4. Return `(!b) && (!c)` 3. Return `(!a) || ((!b) && (!c))` 4. Return `(!a) || ((!b) && (!c))`

Sequence of visits after PR#12942

1. RW, Visit `!(a && (b || c))` 1. RW, Visit `a && (b || c)` 1. RW, Visit `a` 2. RW, Visit `b || c` 1. RW, Visit `b` 2. RW, Visit `c` 3. Return `b || c` 3. Return `a && (b || c)` 2. Delegate to BoolRW 3. Recursive rewrite, `!(a && (b || c))` becomes `(!a) || (!(b||c))` 4. BoolRW, Visit `(!a) || (!(b||c))` 1. BoolRW, Visit `!a` 2. BoolRW, Visit `(!(b||c))` 1. BoolRW, Visit `b||c` 2. Recursive rewrite, `!(b||c)` becomes `(!b) && (!c)` 3. BoolRW, Visit `(!b) && (!c)` 1. Visit `!b` 2. Visit `!c` 3. Return `(!b) && (!c)` 4. Return `(!b) && (!c)` 3. Return `(!a) || ((!b) && (!c))` 5. Return `(!a) || ((!b) && (!c))`

For extra caution, it might be helpful to simply introduce a bit of duplication of rules and have an independent implementation here.

In the short-term, I agree that a copy/paste of the relevant portions would be the safest, and would minimize the scope of any changes. In the long-term, I'd worry about having multiple duplicate implementations, whether improvements/bugfixes in one get propagated over to the other, or the impact of requiring a developer to know which of multiple related implementations should be called.

Additionally, is it OK to simply invoke the full simplification instead?

For the test cases I've run, it is okay to use the full simplification rather than limiting it to boolean expressions. That said, the test cases were designed to test the functionality being added, and so I'm not sure the final cost of the additional visits.

[tqchen]

Thanks @Lunderberg for detailed analysis. Get it. Given full simplification is also applicable. It also depends on how many times simplify is being called in the related places. If they are constant, perhaps we could use full simplification instead for now, then circle back if we have places where we clearly see there is a need.

[Lunderberg]

That makes sense, and agreed as a path forward. Closing this PR until and unless it becomes necessary or the performance difference becomes much larger.

[tqchen]

Thanks @Lunderberg for constructive discussions

[Lunderberg]

Absolutely, and thank you for the same, @tqchen!