diff options
Diffstat (limited to 'tensorflow/compiler/xla/service/layout_assignment.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/layout_assignment.h | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index b75ecb311a..f9e8dbea2f 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -488,6 +488,9 @@ class LayoutAssignment : public HloPassInterface { } } + // Adds constraints related to host Send/Recv instructions. + Status BuildHostChannelConstraints(HloComputation* computation); + // Map containing the layouts of all computations assigned so // far. Computations are handled in a topological sort where computations are // handled before their caller instructions so the layouts of caller @@ -507,6 +510,10 @@ class LayoutAssignment : public HloPassInterface { // computations/instructions. ChannelLayoutConstraints channel_constraints_; + // Layout constraints for send/recv instructions which communicate with the + // host. + ChannelLayoutConstraints host_channel_constraints_; + // The set of HLO instructions which lacked any layout constraint, thus // receiving propagated default layouts. tensorflow::gtl::FlatSet<const HloInstruction*> |