[Bug] Possible issue with the "simplify pass" using the "propagate_knowns_to_simplify_expressions" flag

[sdalvi-quic]

I was trying to understand the use of T.assume in the simplification of expressions and I came across this example in the "test_tir_transform_simplify.py" file named as "TestSimplifyUsingPartiallyKnownBufferExpression". Here, I tried to modify the test and saw that, it was failing while using the "propagate_knowns_to_simplify_expressions" flag and T.assume statements. I have attached the test below.

class TestSimplifyUsingPartiallyKnownBufferExpression(BaseBeforeAfter):
    """An assumption about buffer contents may apply to only part of a buffer

    Like TestSimplifyUsingPartiallyKnownBufferConditional, but the
    conditional is expressed as part of T.assume, instead of in the
    control flow.
    """

    # propagate_knowns_to_prove_conditional = True
    propagate_knowns_to_simplify_expressions = True

    def before(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            T.evaluate(T.assume(i < 14 or A[i] == 0))

        for i in T.serial(16):
            if i < 14:
                if A[i] == 0:
                    A[i] = 42

    def expected(A: T.Buffer(16, "int32")):
        for i in T.serial(16):
            T.evaluate(T.assume(i < 14 or A[i] == 0))

        for i in T.serial(16):
            if i < 14:
                if A[i] == 0:
                    A[i] = 42

Here given the assumption statements T.assume(i < 14 or A[i] == 0), A[i] will have to be 0 if the range of i >= 14 but A[i] may or may not be 0 if i < 14. So, we should see the expected code as above where no statement is eliminated. But the IR returned after the pass eliminates the condition if A[i] == 0 and returns the IR as below :

def after(A: T.Buffer((16,), "int32")):
   T.func_attr({"global_symbol": "main"})
   for i in range(16):
       T.assume(i < 14 or A[i] == 0)
   for i in range(16):
      if i < 14:
           A[i] = 42

I see that the two flags "propagate_knowns_to_simplify_expressions" and "propagate_knowns_to_prove_conditional " are internally calling the function SimplifyInContext which is simplifying using the line :

expr = control_flow_block.known_at_block_start.SubstituteKnownBufferValues(
      std::move(expr), axis_var_lookup_, analyzer);

Can you please help to fix the issue with the test or point what modification needs to be done?

According to my understanding there is either an issue with the control flow graph generated when T.assume statement are present or there needs to be slight changes in the SimplifyInContext function which is simplifying the expression, please correct me if I am wrong.

@Lunderberg, it will be helpful if you could help with this issue.

[Lunderberg]

Thank you for the bug report, and that's definitely incorrect behavior. I think the problem derives from here. The additional_predicate isn't being appended to the buffer access, and so the T.assume(i < 14 or A[i]==0) is erroneously being treated as if it were T.assume(A[i]==0).

[sdalvi-quic]

Thank you Eric for pointing to the location which might be causing issue. Yes, the additional_predicate is not getting appended to the constraints.