IfStmt simplification did not handle AtomicOp correctly?

[k-ye]

Describe the bug

I think there is a bug when simplifying IfStmt, around how it handles atomics. See the screenshots below when running renderer.py for a difference.

If I disable the simplification when there is an AtomicOp inside IfStmt, the bug gets resolved.

I.e., change this line in simplify.cpp from

  if (is_global_write(clause[i].get()) ||

to

  if (clause[i]->is<GlobalStoreStmt>() ||

I'm running the CPU version on macOS Catalina 10.15.2.

Log/Screenshots

This is the result I got when running the renderer example on the head release 0.3.21 or the master branch 0726d49:

And if I patch in the fix, the rendered result will show the taichi logo

To Reproduce

python examples/renderer.py

[yuanming-hu]

Refactoring simplify.cpp to fix this and make the correctness more trivial...

[yuanming-hu]

The minimal script to reproduce this issue:

@ti.all_archs
def test_local_atomic_with_if():
  ret = ti.var(dt=ti.i32, shape=())
  ti.cfg.print_ir = True

  @ti.kernel
  def test():
    if True:
      x = 0
      x += 1
      ret[None] = x

  test()
  assert ret[None] == 1

[yuanming-hu]

Fixed for now by adding

taichi/taichi/transforms/simplify.cpp

Line 814 in b34fc92

op->instance_id == alloca->instance_id)) {

in is_atomic_value_used. Will take a deeper look into the whole simplify.cpp file tomorrow to figure out more systematic solutions. Releasing v0.3.22 with these improvements. Thanks!

[yuanming-hu]

Since no related issues have been raised about the simplified pass, I assume simplify.cpp works with no problem...