[Tracking Issue] TIR-level CSE follow up

[masahi]

Follow up items after #9482

• Extend the notion of expression equality to take advantage of the purity annotation of TIR intrinsics. See the discussion in Implementation of Common Subexpression Elimination for TIR #9482 (comment) and below for the motivation.
• Rewrite the test case using TVMScript. Implementation of Common Subexpression Elimination for TIR #9482 (comment)
• Explore possibilities of refactoring existing TIR utility functions as a special case of the new functionality in the PR. Implementation of Common Subexpression Elimination for TIR #9482 (comment)

@FranckQC @jroesch, also cc @mbs-octoml @mikepapadim @electriclilies

cc @Hzfengsy @junrushao

[FranckQC]

Thanks @masahi for the summary.
I'll be more than happy to look into these (in particular item 1 on commoning pure functions and item 3 on refactoring some existing functions as special cases of some news functions from the CSE) as soon as I have some time!
I know item 1 would make the CSE pass useful for you @masahi, so I'll do that as soon as possible.

Kind regards,
Franck

[yuanfz98]

Hello,

For semantic extension, I listed some items:

• commutativity (identifying for instance (x+y) with (y+x))
• associativity (identifying for instance (x+y)+z with x+(y+z))
• distributivity (x*(y+z) and x*y+x*z)

We may call tvm/arith library to do the job.

bool EquivalentTerms(const PrimExpr& a, const PrimExpr& b) {
  // For now, we just check the syntactic equality, but that could later become a semantic test,
  // for instance identifying computations modulo commutativity (like x+y and y+x), or modulo
  // associativity (like (x+y)+z and x+(y+z)), etc.
  arith::Analyzer analyser;
  PrimExpr a_simplified = analyser.Simplify(a);
  PrimExpr b_simplified = analyser.Simplify(b);
  return EqualTerms(a_simplified, b_simplified);
}

An example test :

def test_semantic_equiv_distributivity():
    i1 = te.var("i1")
    i2 = te.var("i2")
    x = te.var("x")
    y = te.var("y")
    z = te.var("z")
    dtype = "int32"
    buffer = tvm.tir.decl_buffer((50,), dtype)
    body = tvm.tir.SeqStmt(
        [
            tvm.tir.Store(buffer.data, x*(y+z), i1),
            tvm.tir.Store(buffer.data, x*y+x*z, i2),
        ]
    )

    mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([i1, i2, x, y, z], body))
    body = tvm.tir.transform.CommonSubexprElimTIR()(mod)

    tvm.transform.PrintIR()(body)

And I got:

[15:10:03] /home/yuan/Coding/compiler/repos/tvm/src/ir/transform.cc:566: PrintIR():
#[version = "0.0.5"]
@main = primfn(i1: int32, i2: int32, x: int32, y: int32, z: int32) -> () {
  let cse_var_1: int32 = (x*(y + z))
   {
    buffer: Pointer(int32)[i1] = cse_var_1
    buffer[i2] = cse_var_1
  }
}

However I found that the arith system may not be able to do the commutativity (x+y != y+x and x*y != y*x).

[FranckQC]

Hi @yuanfz98

Thanks for your interest in the pass!
That's great, I didn't know that there was an arith::Analyzer which was already able to algebraic simplifications of terms!

When I worked on the CSE pass, I designed it with the idea that it should be easy to extend it for working with terms modulo some equations (by just replacing the EquivalentTerms() function as you did), but I though that we would have to write the normalization procedures on terms for doing that (which would simplify terms within an algebraic structure). And writing such normalization functions is quite a bit annoying to do (it implies deciding on some order terms, sorting them, etc). So I left it for a potential future work. I wasn't even sure that the efforts to implement such normalization procedures would be worth it, because I didn't know how often semantically equivalent computations appear in a TIR program. But as we already have this Simplify() method from arith::Analyzer that does simplification work, we can indeed reuse it without even thinking about it. That's great!

Too bad that arith::Analyzer can't deal with commutativity. I had a quick look at how RewriteSimplifier (which is used by Simplify()) is implemented, and it works by stating rewriting rules that could be applied. They are tried in the order in which they appear with the TVM_TRY_REWRITE() directive.
So instead of computing a normal form in an algebraic structure as I described above, it seems that they try to apply rewriting rules (in a specific order), and they keep applying them, hoping that it will converge to the simplified term. I guess it often does some interesting simplifications, but it is not guaranteed to lead to the simplified term if it exists (unlike a normalization procedure simplifying a term in an algrebraic structure like a Monoid, a Group or a Ring). And without limiting the number of rewriting steps, it would not even be guaranteed to terminate.

The reason why they can't define a commutativity rule is therefore likely because it would create infinite loops of rewriting:
With the rule x+y --> y+x,
the term a+b can be rewritten to b+a, which can again be rewritten to a+b, etc.
Of course, we could limit the rewriting to a fixed number of steps (and they already do it : Simplify() calls the RewriteSimplifier a fixed amount of times, by default 2) to avoid the infinite rewriting, but then this rule for commutativity would be useless. In order to deal with commutativity, one really has to decide an order on subterms and to re-order terms using this order.

In short, for rewriting rules that can be applied again and which lead back to the original term (i.e, A --> B --> A), this strategy of just stating rewrite rules by priority won't work unfortunately.
I guess it's just a limitation of the approach they have used for implementing Simplify() (which was a very good idea for TVM, as they deal with much more than the axioms of algebraic structures!).

Simplify() will do the amount of work that it can do and it will be better than nothing.
Very well spotted @yuanfz98 !

[FranckQC]

Would you like to make a PR for this @yuanfz98 ?
It's great that we can have comparisons modulo associativity and modulo distributivity. It's probably pretty rare in practice, but as we can have it for free, it's probably good to take.

Many thanks again for your interest in the pass!

[yuanfz98]

Would you like to make a PR for this @yuanfz98 ? It's great that we can have comparisons modulo associativity and modulo distributivity. It's probably pretty rare in practice, but as we can have it for free, it's probably good to take.

Many thanks again for your interest in the pass!

Sure, I will make a PR this week !

[driazati]

@masahi is this still in progress? the related PRs have merged so I'm closing this tracking issue but please re-open if it's not totally done

[masahi]

Yeah none of the items I raised have been done actually