basic for if-control-flow

[PeiMu]

Hi @stellaraccident @silvasean @cathyzhyi and other contributors,

I'm not sure if it is appropriate to give a PR like this.

I've implemented a very basic demo for the if-control-flow. I'm not sure if the community has such a plan for supporting the control flow. To check out the demo, you can run the e2e_testing/torchscript/control_flow_test.py.

I've also implemented the for-control-flow, as shown in the e2e_testing/torchscript/control_flow_test.py. And the for-control-flow code runs successfully in this PR. But it seems we have something like unroll in --torch-reify-shape-calculations. Does this mean we will use these several passes to unroll the for-like control flow other than using the scf.for?

Any advice would appreciate a lot!

Best regards, Pei

[silvasean]

That pass should not be unrolling outside the shape calculation regions. Filed https://github.com/llvm/torch-mlir/issues/735 to track that.

I'm still not fully clear on the approach here. Can you show the IR before/after MaximizeValueSemantics for the follwing test cases?

def test_case_1(t):
  if t.sum() > 0:
    x = t + 1
  else:
    x = t - 1
  return x

def test_case_2(t):
  x = t
  if t.sum() > 0:
    x = t + 1
  return x

def test_case_3(t):
  if t.sum() > 0:
    if t.sum() == 42.:
      x = t + 2
    else:
      x = t + 3
  else:
    x = t - 1
  return x