Extended graph lowering to use the stateless if/while ops when the if true/else body or while body funcs do not have stateful ops. #18509
Conversation
|
This PR can only be tested after the TF level code change is submitted, for adding the stateless op/while ops. I will add reviewers after that. |
| @@ -3131,7 +3133,10 @@ TFGraphFunctionLowering::lowerConditionalRegion(ConditionalSESERegion *r) { | |||
| // .Attr("Tin: list(type) >= 0") | |||
| // .Attr("Tout: list(type) >= 0") | |||
| auto opLocString = getUniqueName(loc, "op"); | |||
| auto *op = TF_NewOperation(graphFn.getGraph(), "If", opLocString.c_str()); | |||
| bool usesideEffectingIf = trueFnHasSideEffects || falseFnHasSideEffects; | |||
There was a problem hiding this comment.
Nit: camelcase useSideEffectingIf
… be used, when the if then/else body of If or the While body funcs do not have stateful ops. The are lowered to the same XLA ops. One use case is in the S4TF compiler: swiftlang/swift#18509 PiperOrigin-RevId: 207977126
… true/else body or while body funcs do not have stateful ops. This is a potential performance optimization, where stateless ops can be pruned away if TF-level graph compiler (e.g. grappler) can prove doing so is safe.
mhong
force-pushed
the
stateless_if_while
branch
from
ab240b3
to
1868f33
Compare
|
@swift-ci please test tensorflow |
mhong
changed the title
|
This PR is now ready for review. |
| @@ -3167,7 +3169,10 @@ TFGraphFunctionLowering::lowerConditionalRegion(ConditionalSESERegion *r) { | |||
| // .Attr("Tin: list(type) >= 0") | |||
| // .Attr("Tout: list(type) >= 0") | |||
| auto opLocString = getUniqueName(loc, "op"); | |||
| auto *op = TF_NewOperation(graphFn.getGraph(), "If", opLocString.c_str()); | |||
| bool useSideEffectingIf = trueFnHasSideEffects || falseFnHasSideEffects; | |||
There was a problem hiding this comment.
Both trueFnHasSideEffects and falseFnHasSideEffects are hard-coded false above. What would happen if there's a side-effecting operation in the body?
There was a problem hiding this comment.
They are output params of buildGraphFunction(). Does that clarify?
There was a problem hiding this comment.
Also, there are test cases that generate (the stateful) "If" and "While" ops
|
I'm not aware of specific optimizations targeting If/While, but in general Grappler can optimize stateless ops more (e.g. can prune away stateless ops when we run a graph function that produce no return values). So using these new stateless versions of the ops could trigger such optimization. |
|
Thanks all for the review! |
This is a potential performance optimization, where stateless ops can be pruned
away if TF-level graph compiler (e.g. grappler) can prove doing so is safe.
Updated TF repo commit hash to incorporate the new stateless if/while ops: tensorflow/tensorflow@8453e23.