[XLA] Add simple HLO if conversion pass #22974

Keno · 2018-10-14T23:52:58Z

kConditional operations are currently generally disallowed in parallel contexts
(e.g. in mapped computations). The julia XLA frontend was running into this limitation
quite a bit, because existing julia code tends to use the terniary operator for select,
e.g. to describe the derivative of a max call (and thus a relu) - see the
definitions of the derivatives of max at
https://github.com/JuliaDiff/DiffRules.jl/blob/master/src/rules.jl#L94

To support these sorts of patterns, add a simple if conversion pass that converts
conditionals in parallel context by equivalent select operations (which are well supported),
i.e. a computation like:

if {
 %pif = () parameter(0)
 ROOT %cif = f32[] constant(0)
}

else {
 %pelse = () parameter(0)
 ROOT %celse = f32[] constant(1)
}

mapped {
 %a = f32[] parameter(0)
 %b = f32[] parameter(1)
 %lt = pred[] less-than(%a, %b)
 %t = () tuple()
 ROOT %conditional = f32[] conditional(%lt, %t, %t), true_computation=if, false_computation=else
}

ENTRY comp {
 %p1 = f32[1000]{0} parameter(0)
 %p2 = f32[1000]{0} parameter(1)
 ROOT %mapped = f32[1000]{0} map(%p1, %p2), dimensions={0}, to_apply=mapped
}

gets rewritten to

mapped {
 %a = f32[] parameter(0)
 %b = f32[] parameter(1)
 %cif = f32[] constant(0)
 %celse = f32[] constant(1)
 %lt = pred[] less-than(%a, %b)
 ROOT %select = f32[] select(%lt, %cif, %celse)
}

ENTRY comp {
 %p1 = f32[1000]{0} parameter(0)
 %p2 = f32[1000]{0} parameter(1)
 ROOT %mapped = f32[1000]{0} map(%p1, %p2) dimensions={0} to_apply=mapped
}

To keep things simple, this is accomplished by first rewriting the conditional
to two calls and a select and then inlining the individual calls. Naturally,
the transformation is only applied if the called computation do not
have side effects (which they generally don't if they're in parallel
context). In the future, it would be good to let MapInliner further
simplify this to an implicitly mapped select.

sanjoy · 2018-10-16T19:43:06Z

tensorflow/compiler/xla/service/conditional_to_select.cc

+
+namespace xla {
+
+StatusOr<bool> DoConditionalToSelect(HloInstruction* conditional) {


Please either make this static or put in an anonymous namespace.

sanjoy · 2018-10-16T19:44:50Z

tensorflow/compiler/xla/service/conditional_to_select.cc

+StatusOr<bool> DoConditionalToSelect(HloInstruction* conditional) {
+  // Only allow conditional to select if the called computations
+  // do not have side effects.
+  for (HloComputation* computation : conditional->called_computations()) {


Up to you, but I think conditional->true_computation()->HasSideEffects() || conditional->false_computation()->HasSideEffects() looks a bit cleaner.

sanjoy · 2018-10-16T19:48:23Z

tensorflow/compiler/xla/service/conditional_to_select.cc

+          conditional->false_computation()));
+  HloInstruction* select_op =
+      computation->AddInstruction(HloInstruction::CreateTernary(
+          conditional->shape(), HloOpcode::kSelect,


I suspect this may need to be a kTupleSelect if conditional->shape() is a tuple.

Also, I think we need to broadcast the condition when creating the select if conditional->shape() is not a scalar.

sanjoy · 2018-10-16T19:49:23Z

tensorflow/compiler/xla/service/conditional_to_select.cc

+  TF_RETURN_IF_ERROR(
+      call_graph->VisitNodes([&](const CallGraphNode& node) -> Status {
+        std::vector<HloInstruction*> ToInline;
+        if (node.context() != CallContext::kParallel) return Status::OK();


We prefer using braces even for single statement if and for.

Keno · 2018-10-16T22:25:31Z

Feedback addressed.

sanjoy · 2018-10-16T23:13:03Z

tensorflow/compiler/xla/service/conditional_to_select.h

+namespace xla {
+
+// A pass which transforms conditionals to selects in places where conditionals
+// are not allowed to appear (e.g. mapped computation)


Just to be clear, this is a backend constraint right (that we can't codegen this) not a general constraint (i.e. the un-transformed IR still passes the verifier)? If yes, would be nice to note explicitly.

It passes the verifier, but fails in buffer assignment. The backend folks would have to say what the long term plan here is.

Ok, can you instead s/not allowed/legal but not supported/ or something like that? I don't want folks to get misled by this comment about what HLO is allowed.

Keno · 2018-10-24T20:27:03Z

Looks like there's some whitespace issues with my BUILD changes? Should I update the PR to fix that or do you want to do that as part of the merge?

sanjoy · 2018-10-24T23:40:13Z

Hi Keno,

It'd be great if you could fix it on the PR itself; I believe the error message has a diff? If you run into any unforeseen issues let me know and I'll send you a buildified BUILD file.

Keno · 2018-10-25T17:23:04Z

Alright. Let's try that. I assume you'll have to re-trigger CI since that doesn't seem to run automatically.

Keno · 2018-11-05T03:50:32Z

bump

… TPUs Will be fixed once tensorflow/tensorflow#22974 is deployed to TPUs.

sanjoy

lgtm with comments addressed but not actually lgtm'ing since I'm not sure if the code will be auto merged automatically without giving you time to address the comments.

sanjoy · 2018-11-05T04:07:38Z

tensorflow/compiler/xla/service/conditional_to_select.h

+namespace xla {
+
+// A pass which transforms conditionals to selects in places where conditionals
+// are not allowed to appear (e.g. mapped computation)


Ok, can you instead s/not allowed/legal but not supported/ or something like that? I don't want folks to get misled by this comment about what HLO is allowed.

sanjoy · 2018-11-05T04:11:09Z

tensorflow/compiler/xla/service/conditional_to_select.cc

+  // which expects the condition to be a scalar.
+  if (!ShapeUtil::IsScalar(conditional->shape()) &&
+      !ShapeUtil::IsTuple(conditional->shape())) {
+    condition = computation->AddInstruction(HloInstruction::CreateBroadcast(


This can be a new utility in tensorflow/compiler/xla/service/hlo_creation_utils.h

I see that it already has MakeSelectHlo. In principle that could be expanded do do this whole rigamarole. Should I do that? Or should I create a separate function?

Making the existing switch between select and triple-select sgtm.

Keno · 2018-11-05T19:39:41Z

I've addressed those two comments. @davidel had some thoughts on the split of work here between HLO and the backend, but I think the conclusion ended up being that while we should fix this in the backend also, the HLO-level pass in nevertheless useful.

sanjoy · 2019-02-13T22:21:59Z

@ymodak I'm happy to manually pull in the CL to Google and check it in (and resolve the merge conflicts internally). However it is missing the [ready to pull] label -- can I just add that label directly?

yifeif · 2019-02-27T00:12:28Z

@Keno do you mind resolving the conflicts? Thank you!

Keno · 2019-02-27T18:05:12Z

I was under the impression @sanjoy would do that on import

sanjoy · 2019-02-27T18:42:17Z

I was under the impression @sanjoy would do that on import

I tried to do that but it seems like we don't have a way to pull in a change with conflicts into piper (so that I can later resolve the conflicts). @yifeif - is that accurate?

rthadur · 2019-03-13T21:06:06Z

@Keno @sanjoy @yifeif any updates on resolving conflicts ?

Keno · 2019-03-13T21:11:38Z

Let me just do it. Sounds like that's easiest.

rthadur · 2019-03-13T21:28:33Z

Let me just do it. Sounds like that's easiest.

thank you

PiperOrigin-RevId: 239060519

Keno · 2019-03-20T01:56:27Z

Thanks, @sanjoy.

googlebot added the cla: yes label Oct 14, 2018

drpngx requested a review from sanjoy October 15, 2018 17:32

drpngx added the awaiting review Pull request awaiting review label Oct 15, 2018

sanjoy suggested changes Oct 16, 2018

View reviewed changes

Keno force-pushed the kf/ifconversion branch from 735945f to 97214c7 Compare October 16, 2018 22:25

sanjoy previously approved these changes Oct 16, 2018

View reviewed changes

tensorflowbutler removed the awaiting review Pull request awaiting review label Oct 17, 2018

ymodak self-assigned this Oct 18, 2018

ymodak added the kokoro:force-run Tests on submitted change label Oct 23, 2018

kokoro-team removed the kokoro:force-run Tests on submitted change label Oct 23, 2018

ymodak added awaiting testing (then merge) kokoro:force-run Tests on submitted change labels Oct 23, 2018

kokoro-team removed the kokoro:force-run Tests on submitted change label Oct 23, 2018

Keno dismissed sanjoy’s stale review via c162bf4 October 25, 2018 17:22

Keno force-pushed the kf/ifconversion branch from 97214c7 to c162bf4 Compare October 25, 2018 17:22

sanjoy added the kokoro:force-run Tests on submitted change label Oct 25, 2018

kokoro-team removed the kokoro:force-run Tests on submitted change label Oct 25, 2018

Keno added a commit to Keno/DiffRules.jl that referenced this pull request Nov 5, 2018

Add a hack for work around control flow not working in map context on…

83ea447

… TPUs Will be fixed once tensorflow/tensorflow#22974 is deployed to TPUs.

sanjoy added the kokoro:force-run Tests on submitted change label Nov 5, 2018

kokoro-team removed the kokoro:force-run Tests on submitted change label Nov 5, 2018

sanjoy suggested changes Nov 5, 2018

View reviewed changes

Keno force-pushed the kf/ifconversion branch from c162bf4 to fc816c9 Compare November 5, 2018 19:36

sanjoy previously approved these changes Nov 20, 2018

View reviewed changes

ymodak added the kokoro:force-run Tests on submitted change label Nov 20, 2018

ymodak added the kokoro:force-run Tests on submitted change label Feb 13, 2019

ymodak added ready to pull PR ready for merge process and removed ready to pull PR ready for merge process labels Feb 13, 2019

tensorflowbutler removed the stat:awaiting response Status - Awaiting response from author label Feb 15, 2019

ymodak added ready to pull PR ready for merge process and removed ready to pull PR ready for merge process kokoro:force-run Tests on submitted change labels Feb 19, 2019

Merge branch 'master' into kf/ifconversion

7026001

sanjoy dismissed their stale review via 7026001 March 16, 2019 17:53

PR Queue automation moved this from Approved by Reviewer to Reviewer Requested Changes Mar 16, 2019

PR Queue automation moved this from Reviewer Requested Changes to Approved by Reviewer Mar 16, 2019

sanjoy approved these changes Mar 16, 2019

View reviewed changes

rthadur assigned rthadur and unassigned ymodak Mar 18, 2019

rthadur added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process and removed ready to pull PR ready for merge process labels Mar 18, 2019

kokoro-team removed the kokoro:force-run Tests on submitted change label Mar 18, 2019

tensorflow-copybara merged commit 7026001 into tensorflow:master Mar 18, 2019

PR Queue automation moved this from Approved by Reviewer to Merged Mar 18, 2019

tensorflow-copybara pushed a commit that referenced this pull request Mar 18, 2019

Merge pull request #22974 from JuliaComputing:kf/ifconversion

ca6e8a1

PiperOrigin-RevId: 239060519

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[XLA] Add simple HLO if conversion pass #22974

[XLA] Add simple HLO if conversion pass #22974

Keno commented Oct 14, 2018

sanjoy Oct 16, 2018

sanjoy Oct 16, 2018

sanjoy Oct 16, 2018

sanjoy Oct 16, 2018

sanjoy Oct 16, 2018

Keno commented Oct 16, 2018

sanjoy Oct 16, 2018

Keno Oct 16, 2018

sanjoy Nov 5, 2018

Keno commented Oct 24, 2018

sanjoy commented Oct 24, 2018

Keno commented Oct 25, 2018

Keno commented Nov 5, 2018

sanjoy left a comment

sanjoy Nov 5, 2018

sanjoy Nov 5, 2018

Keno Nov 5, 2018

sanjoy Nov 5, 2018

Keno commented Nov 5, 2018

sanjoy commented Feb 13, 2019

yifeif commented Feb 27, 2019

Keno commented Feb 27, 2019

sanjoy commented Feb 27, 2019

rthadur commented Mar 13, 2019

Keno commented Mar 13, 2019

rthadur commented Mar 13, 2019

Keno commented Mar 20, 2019


		namespace xla {

		StatusOr<bool> DoConditionalToSelect(HloInstruction* conditional) {

[XLA] Add simple HLO if conversion pass #22974

[XLA] Add simple HLO if conversion pass #22974

Conversation

Keno commented Oct 14, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Keno commented Oct 16, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Keno commented Oct 24, 2018

sanjoy commented Oct 24, 2018

Keno commented Oct 25, 2018

Keno commented Nov 5, 2018

sanjoy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Keno commented Nov 5, 2018

sanjoy commented Feb 13, 2019

yifeif commented Feb 27, 2019

Keno commented Feb 27, 2019

sanjoy commented Feb 27, 2019

rthadur commented Mar 13, 2019

Keno commented Mar 13, 2019

rthadur commented Mar 13, 2019

Keno commented Mar 20, 2019