From ea6f79938f2a7ebfb0af2a77a5a80f415a02dbb5 Mon Sep 17 00:00:00 2001 From: George Karpenkov Date: Tue, 23 Feb 2021 18:09:05 -0800 Subject: [PATCH] [XLA:GPU] Allow using ChannelLayoutConstraints for XLA:GPU PiperOrigin-RevId: 359178668 Change-Id: Iac7816ca263e6f0283f59f02d4f0844a76691f77 --- tensorflow/compiler/xla/service/gpu/gpu_compiler.cc | 4 +++- .../xla/service/gpu/gpu_layout_assignment.h | 13 ++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc index 23e4908268a..3fad7380164 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc @@ -296,9 +296,11 @@ Status GpuCompiler::OptimizeHloModule( // Layout assignment uses alias analysis, which requires the call graph to // be flattened. pipeline.AddPass(); + ChannelLayoutConstraints layout_constraints; pipeline.AddPass( hlo_module->mutable_entry_computation_layout(), - LayoutAssignment::InstructionCanChangeLayout, stream_exec); + LayoutAssignment::InstructionCanChangeLayout, stream_exec, + &layout_constraints); TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status()); } diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h index 6a48e55fd2e..1e6a18077a0 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h @@ -29,12 +29,15 @@ namespace gpu { // layout constraints for operands and results of library calls. class GpuLayoutAssignment : public LayoutAssignment { public: - explicit GpuLayoutAssignment(ComputationLayout* entry_computation_layout, - std::function - instruction_can_change_layout_func, - se::StreamExecutor* stream_executor) + explicit GpuLayoutAssignment( + ComputationLayout* entry_computation_layout, + std::function + instruction_can_change_layout_func, + se::StreamExecutor* stream_executor, + ChannelLayoutConstraints* channel_constraints = nullptr) : LayoutAssignment(entry_computation_layout, - std::move(instruction_can_change_layout_func)), + std::move(instruction_can_change_layout_func), + channel_constraints), stream_executor_(stream_executor) {} ~GpuLayoutAssignment() override {}