while loop simplifier missing about mlir-hlo

[Mogball]

I recently changed scf.while to be simplifiable by SCCP, after fixing a few bugs with getRegionSuccessors and how that function was getting used. The same can be done for mhlo.while.

[lipracer]

Thanks a lot, I will make a reference and then implement the simplification of mhlo::while.

[joker-eph]

Isn't a folder for the WhileOp when the return of the condition is false enough here?

[lipracer]

Yeah.First we need to propagate the constant to cond's body, then try to fold it's.

while(%cst0, %cst1) {
  cond (%arg0, %arg1) {
    return compare(%arg0, %arg1)  
  }
}

and the mhlo-sink-constants-to-control-flow just clone constantOp into body, didn't clone arguments.

[joker-eph]

I have a patch for the folder, but I also looked into the RegionBranchOpInterface and it seems like there is a limitation to the interface that prevents us from using it in MHLO.

In particular the getMutableSuccessorOperands for RegionBranchTerminatorOpInterface can't be used with our while op because the condition terminator does not take as operands the values to pass to the body...

[lipracer]

Can we apply the framework of data flow analysis, implemented by constant propagation? Including other control flow simplifications, I'm not sure if scf dailect does these mhlo's as well, but we need these for completeness.

[joker-eph]

This relies on RegionBranchOpInterface, which is difficult to apply here (what I'm describing above)

[Mogball]

In particular the getMutableSuccessorOperands for RegionBranchTerminatorOpInterface can't be used with our while op because the condition terminator does not take as operands the values to pass to the body...

I've run into this limitation before but just added as operands the forwarded values. I will check whether the function really needs to return a MutableOperandRange or whether it can just return a ValueRange. Then, the terminator could just return the block arguments.

[joker-eph]

Thanks!

[lipracer]

Thanks.I need to know about RegionBranchOpInterface. But xla service already has while_loop_simplifier.

[Mogball]

I will check whether the function really needs to return a MutableOperandRange or whether it can just return a ValueRange

The only upstream use of the MutableOperandRange return type is in buffer deallocation.

This signature matches with the interface for BranchOpInterface in that it allows transformations to reason and modify control-flow constructs. I'm not sure it can go away.