[XLA/GPU] Add CustomCallSchedule to give schedule hints to custom-calls. #48806

trentlo · 2021-04-28T17:57:52Z

Add schedule hints EARLY_AS_POSSIBLE and LATE_AS_POSSIBLE to custom-calls.
This supports a custom-call case, where a logical operation can be lowered into
two HLOs (e.g., PerformX and PerformXDone). We can utilize this mechanism to
either hide host latencies between the pair of the custom-calls or the two calls
can more accurately identify the def-use relationship (typically PerformX is
scheduled right after all of its producers have been scheduled and PerformXDone
is scheduled right before its first consumer.)

I need this change for implementing XLA Horovod ops. I have a working prototype internally within NVIDIA and it works well with this change for our tracked DL models, i.e., host overhead is well hidden and I do see some overlapping/parallelism between communication and computation.

Add schedule hints EARLY_AS_POSSIBLE and LATE_AS_POSSIBLE to custom-calls. This supports a custom-call case, where a logical operation can be lowered into two HLOs (e.g., PerformX and PerformXDone). We can utilize this mechanism to either hide host latencies between the pair of the custom-calls or the two calls can more accurately identify the def-use relationship (typically PerformX is scheduled right after all of its producers have been scheduled and PerformXDone is scheduled right before its first consumer.)

trentlo · 2021-04-28T18:20:11Z

@timshen91 could you help to review this PR? Thanks!

@cheshire @sanjoy FYI.

yunxing · 2021-04-29T12:00:12Z

tensorflow/compiler/xla/service/hlo_memory_scheduler.h

+// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing
+// function to MemorySchedulerAlgorithm to enforce certain hlo schedule
+// constraints desired for custom-calls.
+typedef std::function<HloInstructionSequence(const HloInstructionSequence&)>


Is it more composable if we do postprocessing as another pass after memory scheduling?

The scheduler tries to compare the mem usage of each schedule and pick the lowest one, so ideally postprocessing would happen in this file after the scheduling algorithm, but before the mem usage calculation.

Right, the main motivation to implement in the current way (vs. a separate pass) is to account for potential memory size changes caused by the postprocessors. In this way, the change is accounted into the memory usage calculation.

@yunxing @dimvar Further thoughts?

Kariddi · 2021-04-29T15:14:34Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+        if (CustomCallWithSchedule(user,
+                                   CustomCallSchedule::EARLY_AS_POSSIBLE) &&
+            absl::c_all_of(user->operands(), [&](const HloInstruction* opnd) {
+              return scheduled.contains(opnd);


Should control-dependencies also be checked?

Good question. I was also wondering if checking control dependencies is necessary in this particular use case of custom-calls. I think it might not be necessary as in the current code base, i.e., I don't find any codes adding control dependencies to custom-calls. However, let's be conservative and add the check for control dependencies for future proof.

Kariddi · 2021-04-29T15:20:28Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+        if (CustomCallWithSchedule(opnd,
+                                   CustomCallSchedule::LATE_AS_POSSIBLE) &&
+            absl::c_all_of(opnd->users(), [&](const HloInstruction* u) {
+              return scheduled.contains(u);


Same about control-dependencies. I think you might need to check it.

Kariddi · 2021-04-29T15:23:29Z

tensorflow/compiler/xla/service/hlo_memory_scheduler.h

+// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing
+// function to MemorySchedulerAlgorithm to enforce certain hlo schedule
+// constraints desired for custom-calls.
+typedef std::function<HloInstructionSequence(const HloInstructionSequence&)>


Kariddi · 2021-04-29T15:24:09Z

tensorflow/compiler/xla/service/hlo_parser_test.cc

+
+ENTRY %CustomCall () -> f32[1,2,3] {
+  %constant = f32[1]{0} constant({12345})
+  ROOT %custom-call = f32[1,2,3]{0,2,1} custom-call(f32[1]{0} %constant), custom_call_target="foo\"bar", schedule=LATE_AS_POSSIBLE


Could you add the tests also for EARLY ?

Lee165 · 2021-04-29T15:35:06Z

Please don't send me any more documents

…

---Original--- From: "Marcello ***@***.***> Date: Thu, Apr 29, 2021 23:32 PM To: ***@***.***>; Cc: ***@***.***>; Subject: Re: [tensorflow/tensorflow] [XLA/GPU] Add CustomCallSchedule to give schedule hints to custom-calls. (#48806) @Kariddi commented on this pull request. In tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc: > + const std::vector<HloInstruction*>& instrs = input.instructions(); + for (HloInstruction* instr : instrs) { + if (scheduled.contains(instr)) { + continue; + } + + early_as_possible_sched.push_back(instr); + scheduled.insert(instr); + + for (HloInstruction* user : instr->users()) { + // Schedule any user who has the attribute `early_as_possible` and all + // of its producers have been scheduled. + if (CustomCallWithSchedule(user, + CustomCallSchedule::EARLY_AS_POSSIBLE) && + absl::c_all_of(user->operands(), [&](const HloInstruction* opnd) { + return scheduled.contains(opnd); Should control-dependencies also be checked? In tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc: > + for (auto it = early_as_possible_sched.rbegin(); + it != early_as_possible_sched.rend(); it++) { + if (scheduled.contains(*it)) { + continue; + } + + late_as_possible_sched.push_front(*it); + scheduled.insert(*it); + + for (HloInstruction* opnd : (*it)->unique_operands()) { + // Schedule any opnd who has the attribute `late_as_possible` if all of + // its users have been scheduled. + if (CustomCallWithSchedule(opnd, + CustomCallSchedule::LATE_AS_POSSIBLE) && + absl::c_all_of(opnd->users(), [&](const HloInstruction* u) { + return scheduled.contains(u); Same about control-dependencies. I think you might need to check it. In tensorflow/compiler/xla/service/hlo_parser_test.cc: > @@ -1079,6 +1079,18 @@ ENTRY %CustomCallWithAliasing (p0: (f32[2,2], f32[42,2,3]), p1: f32[123,4]) -> ( ROOT %custom-call = (f32[123,4]{0,1}, f32[2,2]{0,1}, f32[1,2,3]{0,1,2}) custom-call((f32[2,2]{0,1}, f32[42,2,3]{0,1,2}) %p0, f32[123,4]{0,1} %p1), custom_call_target="baz", output_to_operand_aliasing={{0}: (1, {}), {1}: (0, {0})} } +)" +}, +// CustomCall with schedule. +{ +"CustomCallWithSchedule", +R"(HloModule custom_call + +ENTRY %CustomCall () -> f32[1,2,3] { + %constant = f32[1]{0} constant({12345}) + ROOT %custom-call = f32[1,2,3]{0,2,1} custom-call(f32[1]{0} %constant), custom_call_target="foo\"bar", schedule=LATE_AS_POSSIBLE Could you add the tests also for EARLY ? In tensorflow/compiler/xla/service/hlo_memory_scheduler.h: > @@ -32,6 +32,12 @@ limitations under the License. namespace xla { +// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing +// function to MemorySchedulerAlgorithm to enforce certain hlo schedule +// constraints desired for custom-calls. +typedef std::function<HloInstructionSequence(const HloInstructionSequence&)> +1 — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub, or unsubscribe.

trentlo

Thanks for taking a look. Some response inlined. Will update the codes soon.

trentlo · 2021-04-29T17:47:56Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+        if (CustomCallWithSchedule(user,
+                                   CustomCallSchedule::EARLY_AS_POSSIBLE) &&
+            absl::c_all_of(user->operands(), [&](const HloInstruction* opnd) {
+              return scheduled.contains(opnd);


Good question. I was also wondering if checking control dependencies is necessary in this particular use case of custom-calls. I think it might not be necessary as in the current code base, i.e., I don't find any codes adding control dependencies to custom-calls. However, let's be conservative and add the check for control dependencies for future proof.

trentlo · 2021-04-29T17:53:20Z

tensorflow/compiler/xla/service/hlo_memory_scheduler.h

+// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing
+// function to MemorySchedulerAlgorithm to enforce certain hlo schedule
+// constraints desired for custom-calls.
+typedef std::function<HloInstructionSequence(const HloInstructionSequence&)>


Right, the main motivation to implement in the current way (vs. a separate pass) is to account for potential memory size changes caused by the postprocessors. In this way, the change is accounted into the memory usage calculation.

trentlo · 2021-04-29T17:53:42Z

tensorflow/compiler/xla/service/hlo_parser_test.cc

+
+ENTRY %CustomCall () -> f32[1,2,3] {
+  %constant = f32[1]{0} constant({12345})
+  ROOT %custom-call = f32[1,2,3]{0,2,1} custom-call(f32[1]{0} %constant), custom_call_target="foo\"bar", schedule=LATE_AS_POSSIBLE


…dencies.

trentlo · 2021-04-29T19:36:38Z

Codes updated. Please help to take another look. Thanks!

trentlo · 2021-05-03T20:26:01Z

@Kariddi A test failure is because I missed the initialization of custom_call_schedule_ in one of the HloCustomCallInstruction constructors. (EDIT: just checked. It was missing because that constructor is new and not in NV's forked branch.) I push a commit to fix it. I don't think other test/build failures in the CI are related to this PR.

Could you please approve the PR again? Thanks!

trentlo · 2021-05-04T20:57:16Z

@Kariddi A test failure is because I missed the initialization of custom_call_schedule_ in one of the HloCustomCallInstruction constructors. (EDIT: just checked. It was missing because that constructor is new and not in NV's forked branch.) I push a commit to fix it. I don't think other test/build failures in the CI are related to this PR.

Could you please approve the PR again? Thanks!

ping~

sanjoy

I have a couple of nitpicky comments.

sanjoy · 2021-05-05T07:10:45Z

tensorflow/compiler/xla/service/hlo_parser_test.cc

+
+ENTRY %CustomCall () -> f32[1,2,3] {
+  %constant = f32[1]{0} constant({12345})
+  %custom-call.0 = f32[1,2,3]{0,2,1} custom-call(f32[1]{0} %constant), custom_call_target="foo", schedule=EARLY_AS_POSSIBLE


How about s/EARLY_AS_POSSIBLE/EARLIEST and s/LATE_AS_POSSIBLE/LATEST?

Make sense. Updated.

sanjoy · 2021-05-05T07:12:50Z

tensorflow/compiler/xla/service/hlo_instructions.h

@@ -1562,6 +1568,8 @@ class HloCustomCallInstruction : public HloInstruction {
  std::vector<std::pair<ShapeIndex, std::pair<int64, ShapeIndex>>>
      output_to_operand_aliasing_;
  absl::optional<Literal> literal_;
+  // A custom-call schedule hint.
+  CustomCallSchedule custom_call_schedule_;


You could just add a default initializer (to CustomCallSchedule::NONE) here instead of doing it in every constructor.

I know. I just wanted to conform to the original coding style in the file. Let me know if you still want me to switch to use default initializers.

trentlo

Updated. Please help to take another look. Thanks!

trentlo · 2021-05-05T17:42:19Z

tensorflow/compiler/xla/service/hlo_parser_test.cc

+
+ENTRY %CustomCall () -> f32[1,2,3] {
+  %constant = f32[1]{0} constant({12345})
+  %custom-call.0 = f32[1,2,3]{0,2,1} custom-call(f32[1]{0} %constant), custom_call_target="foo", schedule=EARLY_AS_POSSIBLE


Make sense. Updated.

trentlo · 2021-05-05T17:43:59Z

tensorflow/compiler/xla/service/hlo_instructions.h

@@ -1562,6 +1568,8 @@ class HloCustomCallInstruction : public HloInstruction {
  std::vector<std::pair<ShapeIndex, std::pair<int64, ShapeIndex>>>
      output_to_operand_aliasing_;
  absl::optional<Literal> literal_;
+  // A custom-call schedule hint.
+  CustomCallSchedule custom_call_schedule_;


I know. I just wanted to conform to the original coding style in the file. Let me know if you still want me to switch to use default initializers.

EARLY_AS_POSSIBLE -> EARLIEST. LATE_AS_POSSIBLE -> LATEST.

trentlo · 2021-05-07T19:11:01Z

@sanjoy are you still on this?

cheshire · 2021-05-12T21:49:06Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+// relationship of the two calls (typically PerformX is scheduled right after
+// all of its producers have been scheduled and PerformXDone is scheduled right
+// before its first consumer.)
+HloInstructionSequence postprocessor_to_custom_schedule(


Conceptually this is a bit dodgy: we are exposing the schedule to all clients, but only the GPU backend uses it.

Should we explicitly error out on other backends? Or ignore it?

The current behavior is that if some clients add the schedule attributes to backends other than GPU, the attributes are simply ignored as the their postprocessors are empty.

An alternative is to make a default postprocessor to give warnings. E.g., "the schedule is set but ignored on the backend." How do you think?

Note that the current implementation in effect ignores the schedule attributes if they are used in other backends. This is a safe and correct behavior by all means.

After some thought, it is good for me to leave this in the current way (i.e., no postprocessors in other backends and effectively ignore the schedule attributes). Let me know if you have other thoughts or I don't address your concern.

cheshire · 2021-05-12T21:50:07Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+    };
+    const std::vector<HloInstruction*>& instrs = input.instructions();
+    for (HloInstruction* instr : instrs) {
+      if (scheduled.contains(instr)) {


Haven't we just defined the lambda checking precisely this? Or more concretely: do we need the lambda at all then?

The lambda is needed because the absl::c_all_of() below requires a function form.

absl::c_all_of(user->operands(), is_scheduled) && absl::c_all_of(user->control_predecessors(), is_scheduled)) {

It is cleaner to also use the lambda here. Will update it.

cheshire · 2021-05-12T21:51:02Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+      return scheduled.contains(instr);
+    };
+    const std::vector<HloInstruction*>& instrs = input.instructions();
+    for (HloInstruction* instr : instrs) {


Optional nitpick: maybe just use input.instructions() inline?

cheshire · 2021-05-12T21:51:35Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+// relationship of the two calls (typically PerformX is scheduled right after
+// all of its producers have been scheduled and PerformXDone is scheduled right
+// before its first consumer.)
+HloInstructionSequence postprocessor_to_custom_schedule(


Style guide says functions are UpperCase

Right. I don't know why I made it lower case. Will revise it.

cheshire · 2021-05-12T21:53:28Z

tensorflow/compiler/xla/service/hlo.proto

@@ -237,6 +243,9 @@ message HloInstructionProto {
  repeated xla.CustomCallOutputOperandAliasing
      custom_call_output_operand_aliasing = 74;

+  // Specifies the desired schedule for the custom-call.


Specify that the field is only present for custom calls?

Will add comments to make it clear.

cheshire · 2021-05-12T21:53:50Z

tensorflow/compiler/xla/service/hlo_instruction.cc

@@ -3823,6 +3828,24 @@ StatusOr<PrecisionConfig::Precision> StringToPrecision(const string& name) {
  return found->second;
 }

+StatusOr<CustomCallSchedule> StringToCustomCallSchedule(const string& name) {


Nit: absl::string_view for param

cheshire · 2021-05-12T21:54:09Z

tensorflow/compiler/xla/service/hlo_instruction.cc

@@ -3823,6 +3828,24 @@ StatusOr<PrecisionConfig::Precision> StringToPrecision(const string& name) {
  return found->second;
 }

+StatusOr<CustomCallSchedule> StringToCustomCallSchedule(const string& name) {
+  static std::unordered_map<string, CustomCallSchedule>* map = [] {


Can we use absl::flat_hash_map?

Yes, updated.

tensorflow/compiler/xla/service/hlo_instruction.cc

cheshire · 2021-05-12T21:55:39Z

tensorflow/compiler/xla/service/hlo_memory_scheduler.h

+// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing
+// function to MemorySchedulerAlgorithm to enforce certain hlo schedule
+// constraints desired for custom-calls.
+typedef std::function<HloInstructionSequence(const HloInstructionSequence&)>


Nit: "using" is preferred to typedef.

cheshire · 2021-05-12T21:56:01Z

tensorflow/compiler/xla/service/hlo_memory_scheduler.h

+// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing
+// function to MemorySchedulerAlgorithm to enforce certain hlo schedule
+// constraints desired for custom-calls.
+typedef std::function<HloInstructionSequence(const HloInstructionSequence&)>


@yunxing @dimvar Further thoughts?

trentlo

Thanks George for taking a look. I will update the codes accordingly soon.

trentlo · 2021-05-13T19:03:29Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+// relationship of the two calls (typically PerformX is scheduled right after
+// all of its producers have been scheduled and PerformXDone is scheduled right
+// before its first consumer.)
+HloInstructionSequence postprocessor_to_custom_schedule(


The current behavior is that if some clients add the schedule attributes to backends other than GPU, the attributes are simply ignored as the their postprocessors are empty.

An alternative is to make a default postprocessor to give warnings. E.g., "the schedule is set but ignored on the backend." How do you think?

trentlo · 2021-05-13T19:07:55Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+// relationship of the two calls (typically PerformX is scheduled right after
+// all of its producers have been scheduled and PerformXDone is scheduled right
+// before its first consumer.)
+HloInstructionSequence postprocessor_to_custom_schedule(


Right. I don't know why I made it lower case. Will revise it.

trentlo · 2021-05-13T19:12:15Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+    };
+    const std::vector<HloInstruction*>& instrs = input.instructions();
+    for (HloInstruction* instr : instrs) {
+      if (scheduled.contains(instr)) {


The lambda is needed because the absl::c_all_of() below requires a function form.

absl::c_all_of(user->operands(), is_scheduled) && absl::c_all_of(user->control_predecessors(), is_scheduled)) {

It is cleaner to also use the lambda here. Will update it.

trentlo · 2021-05-13T19:13:49Z

tensorflow/compiler/xla/service/hlo.proto

@@ -237,6 +243,9 @@ message HloInstructionProto {
  repeated xla.CustomCallOutputOperandAliasing
      custom_call_output_operand_aliasing = 74;

+  // Specifies the desired schedule for the custom-call.


Will add comments to make it clear.

tensorflow/compiler/xla/service/hlo_instruction.cc

trentlo

Codes updated. Please help to take a look again and let me know if all the previous comments are addressed. Thanks!

trentlo · 2021-05-14T00:07:18Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+// relationship of the two calls (typically PerformX is scheduled right after
+// all of its producers have been scheduled and PerformXDone is scheduled right
+// before its first consumer.)
+HloInstructionSequence postprocessor_to_custom_schedule(


trentlo · 2021-05-14T00:07:31Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+      return scheduled.contains(instr);
+    };
+    const std::vector<HloInstruction*>& instrs = input.instructions();
+    for (HloInstruction* instr : instrs) {


trentlo · 2021-05-14T00:07:41Z

tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc

+    };
+    const std::vector<HloInstruction*>& instrs = input.instructions();
+    for (HloInstruction* instr : instrs) {
+      if (scheduled.contains(instr)) {


trentlo · 2021-05-14T00:07:51Z

tensorflow/compiler/xla/service/hlo.proto

@@ -237,6 +243,9 @@ message HloInstructionProto {
  repeated xla.CustomCallOutputOperandAliasing
      custom_call_output_operand_aliasing = 74;

+  // Specifies the desired schedule for the custom-call.


trentlo · 2021-05-14T00:08:03Z

tensorflow/compiler/xla/service/hlo_instruction.cc

@@ -3823,6 +3828,24 @@ StatusOr<PrecisionConfig::Precision> StringToPrecision(const string& name) {
  return found->second;
 }

+StatusOr<CustomCallSchedule> StringToCustomCallSchedule(const string& name) {


trentlo · 2021-05-14T00:08:13Z

tensorflow/compiler/xla/service/hlo_instruction.cc

@@ -3823,6 +3828,24 @@ StatusOr<PrecisionConfig::Precision> StringToPrecision(const string& name) {
  return found->second;
 }

+StatusOr<CustomCallSchedule> StringToCustomCallSchedule(const string& name) {
+  static std::unordered_map<string, CustomCallSchedule>* map = [] {


Yes, updated.

trentlo · 2021-05-14T00:08:30Z

tensorflow/compiler/xla/service/hlo_memory_scheduler.h

+// Postprocessor of the HloInstructionSequence. This is an opt-in postprocessing
+// function to MemorySchedulerAlgorithm to enforce certain hlo schedule
+// constraints desired for custom-calls.
+typedef std::function<HloInstructionSequence(const HloInstructionSequence&)>


cheshire · 2021-05-14T22:59:09Z

Warning: Unhelpful next ID: saw 76 but field "custom_call_schedule" already uses it. Use 77 instead? [misleading_next_id]

trentlo · 2021-05-15T07:00:00Z

Will do. I have some internet issue. Will do when my internet is up.

…

On Fri, May 14, 2021 at 3:59 PM George Karpenkov ***@***.***> wrote: Warning: Unhelpful next ID: saw 76 but field "custom_call_schedule" already uses it. Use 77 instead? [misleading_next_id] — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#48806 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AC3VUT5PTKKLR6B7O2KQGRLTNWTM7ANCNFSM43XR3SYA> .

trentlo · 2021-05-17T22:37:32Z

Warning: Unhelpful next ID: saw 76 but field "custom_call_schedule" already uses it. Use 77 instead? [misleading_next_id]

Updated. Please help to take a look again.

PiperOrigin-RevId: 375939575 Change-Id: If25885fd0f0e48e961a5596ff76a83fdf61c36b3

google-ml-butler bot added the size:L CL Change Size: Large label Apr 28, 2021

google-ml-butler bot requested review from joker-eph and sanjoy April 28, 2021 17:57

google-cla bot added the cla: yes label Apr 28, 2021

trentlo force-pushed the xla_custom_call_for_hvd_042721 branch from f4f3a55 to 7de009c Compare April 28, 2021 18:12

gbaned self-assigned this Apr 29, 2021

gbaned added this to Assigned Reviewer in PR Queue via automation Apr 29, 2021

yunxing reviewed Apr 29, 2021

View reviewed changes

Kariddi reviewed Apr 29, 2021

View reviewed changes

trentlo commented Apr 29, 2021

View reviewed changes

trentlo added 2 commits April 29, 2021 12:24

[XLA/GPU] Make postprocessor_to_custom_schedule respect control depen…

f2227ba

…dencies.

[XLA/GPU] Test CustomCallSchedule::EARLY_AS_POSSIBLE in hlo_parser_test.

8ac6a19

trentlo requested review from yunxing and Kariddi April 29, 2021 19:36

Kariddi approved these changes Apr 30, 2021

View reviewed changes

PR Queue automation moved this from Assigned Reviewer to Approved by Reviewer Apr 30, 2021

google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Apr 30, 2021

kokoro-team removed the kokoro:force-run Tests on submitted change label Apr 30, 2021

google-ml-butler bot removed the ready to pull PR ready for merge process label May 3, 2021

[XLA] Initialize custom_call_schedule_ in HloCustomCallInstruction.

ba45314

trentlo requested a review from Kariddi May 3, 2021 20:26

sanjoy suggested changes May 5, 2021

View reviewed changes

PR Queue automation moved this from Approved by Reviewer to Reviewer Requested Changes May 5, 2021

trentlo commented May 5, 2021

View reviewed changes

[XLA/GPU] Rename enums in CustomCallSchedule

d9e163d

EARLY_AS_POSSIBLE -> EARLIEST. LATE_AS_POSSIBLE -> LATEST.

trentlo requested a review from sanjoy May 6, 2021 18:33

cheshire requested changes May 12, 2021

View reviewed changes

trentlo commented May 13, 2021

View reviewed changes

trentlo commented May 14, 2021

View reviewed changes

trentlo requested a review from cheshire May 14, 2021 00:09

[XLA/GPU] Code style polishing.

9b741a9

cheshire approved these changes May 14, 2021

View reviewed changes

google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels May 14, 2021

kokoro-team removed the kokoro:force-run Tests on submitted change label May 14, 2021

sanjoy removed their request for review May 14, 2021 06:59

gbaned removed the ready to pull PR ready for merge process label May 17, 2021

trentlo requested a review from cheshire May 17, 2021 22:37

[XLA] Update Next id in hlo.proto.

6571b86

cheshire approved these changes May 17, 2021

View reviewed changes

google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels May 17, 2021

kokoro-team removed the kokoro:force-run Tests on submitted change label May 17, 2021

akuegel added the kokoro:force-run Tests on submitted change label May 19, 2021

kokoro-team removed the kokoro:force-run Tests on submitted change label May 19, 2021

copybara-service bot merged commit ef862e4 into tensorflow:master May 19, 2021

PR Queue automation moved this from Reviewer Requested Changes to Merged May 19, 2021

copybara-service bot pushed a commit that referenced this pull request May 26, 2021

Add custom call scheduling from #48806 to Python clients.

5201b6e

PiperOrigin-RevId: 375939575 Change-Id: If25885fd0f0e48e961a5596ff76a83fdf61c36b3

trentlo mentioned this pull request Jul 23, 2021

XLA HorovodAllreduce for tf.function(jit_compile=True) horovod/horovod#3053

Merged

4 tasks

[XLA/GPU] Add CustomCallSchedule to give schedule hints to custom-calls. #48806

[XLA/GPU] Add CustomCallSchedule to give schedule hints to custom-calls. #48806

Conversation

trentlo commented Apr 28, 2021

trentlo commented Apr 28, 2021

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

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

Choose a reason for hiding this comment

Lee165 commented Apr 29, 2021 via email

trentlo 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

trentlo commented Apr 29, 2021

trentlo commented May 3, 2021 • edited

trentlo commented May 4, 2021

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

trentlo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

trentlo commented May 7, 2021

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

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

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

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

Choose a reason for hiding this comment

trentlo 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

trentlo 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cheshire commented May 14, 2021

trentlo commented May 15, 2021 via email

trentlo commented May 17, 2021

trentlo commented May 3, 2021 •

edited