diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 65dd03e6f6a..5ac77e84020 100755 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -236,9 +236,9 @@ cc_library( deps = [ ":backend_configs", ":buffer_allocations", - ":cudnn_conv_runner", ":elemental_ir_emitter", ":gpu_constants", + ":gpu_conv_runner", ":gpu_executable", ":hlo_to_ir_bindings", ":ir_emission_utils", @@ -505,7 +505,7 @@ cc_library( ":backend_configs", ":buffer_allocations", ":cusolver_context", - ":cudnn_conv_runner", + ":gpu_conv_runner", ":gpu_debug_info_manager", ":gpu_types", ":hlo_execution_profiler", @@ -615,7 +615,7 @@ cc_library( deps = [ ":backend_configs", ":buffer_comparator", - ":cudnn_conv_runner", + ":gpu_conv_runner", ":gpu_executable", ":ir_emission_utils", ":stream_executor_util", @@ -637,14 +637,14 @@ cc_library( ) cc_library( - name = "cudnn_conv_algorithm_picker", - srcs = ["cudnn_conv_algorithm_picker.cc"], - hdrs = ["cudnn_conv_algorithm_picker.h"], + name = "gpu_conv_algorithm_picker", + srcs = ["gpu_conv_algorithm_picker.cc"], + hdrs = ["gpu_conv_algorithm_picker.h"], deps = [ ":backend_configs", ":buffer_comparator", - ":cudnn_conv_runner", ":gpu_autotuning_proto", + ":gpu_conv_runner", ":gpu_executable", ":hlo_algorithm_blacklist", ":ir_emission_utils", @@ -672,7 +672,7 @@ cc_library( ) cc_library( - name = "cudnn_conv_runner", + name = "gpu_conv_runner", srcs = ["gpu_conv_runner.cc"], hdrs = ["gpu_conv_runner.h"], deps = [ @@ -1057,6 +1057,7 @@ cc_library( ":cudnn_pad_for_convolutions", ":fusion_merger", ":gpu_constants", + ":gpu_conv_algorithm_picker", ":gpu_conv_padding_legalization", ":gpu_conv_rewriter", ":gpu_copy_insertion", @@ -1156,13 +1157,13 @@ cc_library( ], deps = [ ":cublas_gemm_pad_for_tensor_cores", - ":cudnn_conv_algorithm_picker", ":cudnn_fused_conv_rewriter", ":cudnn_pad_for_convolutions", ":cusolver_rewriter", ":gemm_algorithm_picker", ":gemm_rewriter", ":gpu_compiler", + ":gpu_conv_algorithm_picker", ":gpu_conv_padding_legalization", ":gpu_conv_rewriter", ":gpu_layout_assignment", diff --git a/tensorflow/compiler/xla/service/gpu/amdgpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/amdgpu_compiler.cc index ed49fcd584f..10dd4542612 100644 --- a/tensorflow/compiler/xla/service/gpu/amdgpu_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/amdgpu_compiler.cc @@ -16,9 +16,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/amdgpu_compiler.h" #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.h" #include "tensorflow/compiler/xla/service/gpu/gpu_conv_padding_legalization.h" #include "tensorflow/compiler/xla/service/gpu/gpu_conv_rewriter.h" -// TODO(whchung@gmail.com): Add gpu_conv_algorithm_picker after its PR merged. #include "tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h" #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h" #include "tensorflow/compiler/xla/service/gpu/target_constants.h" @@ -97,7 +97,7 @@ Status AMDGPUCompiler::OptimizeHloPostLayoutAssignment( options.set_is_layout_sensitive(true); pipeline.AddPass>(options); - // TODO(whchung@gmail.com): Add gpu_conv_algorithm_picker after its PR merged. + pipeline.AddPass(stream_exec, device_allocator); // Clean up new_tuple described above. pipeline.AddPass(); diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc b/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc similarity index 79% rename from tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc rename to tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc index f68f695b4d7..676982bfdbc 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc @@ -13,9 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.h" -#include "google/protobuf/any.pb.h" #include "absl/algorithm/container.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" @@ -48,6 +47,54 @@ using se::DeviceMemoryBase; using se::dnn::AlgorithmDesc; using tensorflow::AutotuneResult; +class ScratchAllocator : public se::ScratchAllocator { + public: + ScratchAllocator(int device_ordinal, + se::DeviceMemoryAllocator* memory_allocator) + : device_ordinal_(device_ordinal), memory_allocator_(memory_allocator) {} + + int64 GetMemoryLimitInBytes() override { + return 1LL << 32; // 4GB. TODO(jlebar): Tune this? + } + int64 TotalAllocatedBytes() { return total_allocated_bytes_; } + + StatusOr> AllocateBytes(int64 byte_size) override; + + template + StatusOr> Allocate(int64 num_elements) { + TF_ASSIGN_OR_RETURN(se::DeviceMemory bytes, + AllocateBytes(num_elements * sizeof(T))); + return se::DeviceMemory(bytes); + } + + private: + const int device_ordinal_; + se::DeviceMemoryAllocator* memory_allocator_; + std::vector allocated_buffers_; + int64 total_allocated_bytes_ = 0; +}; + +StatusOr> ScratchAllocator::AllocateBytes( + int64 byte_size) { + CHECK_GE(byte_size, 0) << "byte_size must be positive."; + if (byte_size > GetMemoryLimitInBytes()) { + return se::port::Status( + se::port::error::RESOURCE_EXHAUSTED, + absl::StrFormat( + "Allocating %d bytes exceeds the memory limit of %d bytes.", + byte_size, GetMemoryLimitInBytes())); + } + + TF_ASSIGN_OR_RETURN(se::OwningDeviceMemory allocated_buffer, + memory_allocator_->Allocate(device_ordinal_, byte_size, + /*retry_on_failure=*/false)); + total_allocated_bytes_ += byte_size; + + se::DeviceMemoryBase buffer_addr = *allocated_buffer; + allocated_buffers_.push_back(std::move(allocated_buffer)); + return se::DeviceMemory(buffer_addr); +} + std::vector GetAlgorithms(CudnnConvKind kind, se::StreamExecutor* stream_exec) { std::vector algorithms; @@ -198,7 +245,7 @@ auto& autotune_cache_stats GUARDED_BY(autotune_cache_lock) = *new ConvCacheStats(); } // anonymous namespace -StatusOr CudnnConvAlgorithmPicker::PickBestAlgorithm( +StatusOr GpuConvAlgorithmPicker::PickBestAlgorithm( const HloCustomCallInstruction* instr) { // Don't run this function concurrently on the same GPU. // @@ -226,22 +273,6 @@ StatusOr CudnnConvAlgorithmPicker::PickBestAlgorithm( autotune_cache_stats.cache_misses++; } - StatusOr result_or = PickBestAlgorithmNoCache(instr); - if (result_or.ok()) { - tensorflow::mutex_lock lock(autotune_cache_lock); - CHECK(autotune_cache.insert({key, result_or.ValueOrDie()}).second); - } - return result_or; -} - -StatusOr CudnnConvAlgorithmPicker::PickBestAlgorithmNoCache( - const HloCustomCallInstruction* instr) { - XLA_SCOPED_LOGGING_TIMER( - absl::StrCat("CudnnConvAlgorithmPicker::PickBestAlgorithmImpl for ", - instr->ToString())); - - const Shape& result_shape = instr->shape().tuple_shapes(0); - // Make sure any previous activity on this executor is done. We don't want to // interfere with programs that are still running on the GPU. if (!stream_exec_->SynchronizeAllActivity()) { @@ -269,6 +300,32 @@ StatusOr CudnnConvAlgorithmPicker::PickBestAlgorithmNoCache( return &stream_opt.value(); }(); + StatusOr result_or(InternalError("Unknown platform.")); + // Check StreamExecutor on which platform it is. ROCm and Cuda implementation + // have diverged. Secifically, we need to make sure redzone allocator related + // utilities are not used in ROCm routine + if (stream_exec_->platform_kind() == se::PlatformKind::kROCm) { + result_or = PickBestAlgorithmNoCacheRocm(instr, allocator, stream); + } else if (stream_exec_->platform_kind() == se::PlatformKind::kCuda) { + result_or = PickBestAlgorithmNoCacheCuda(instr, allocator, stream); + } + + if (result_or.ok()) { + tensorflow::mutex_lock lock(autotune_cache_lock); + CHECK(autotune_cache.insert({key, result_or.ValueOrDie()}).second); + } + return result_or; +} + +StatusOr +GpuConvAlgorithmPicker::PickBestAlgorithmNoCacheCuda( + const HloCustomCallInstruction* instr, se::DeviceMemoryAllocator* allocator, + se::Stream* stream) { + // Right now Redzone allocator is available in Cuda target only + XLA_SCOPED_LOGGING_TIMER(absl::StrCat( + "GpuConvAlgorithmPicker::PickBestAlgorithmImpl for ", instr->ToString())); + + const Shape& result_shape = instr->shape().tuple_shapes(0); int64 rng_state = 0; const auto initialize_buffer = [&stream, &rng_state]( @@ -526,8 +583,78 @@ StatusOr CudnnConvAlgorithmPicker::PickBestAlgorithmNoCache( instr->ToString()); } -StatusOr CudnnConvAlgorithmPicker::RunOnInstruction( - HloInstruction* instr) { +StatusOr +GpuConvAlgorithmPicker::PickBestAlgorithmNoCacheRocm( + const HloCustomCallInstruction* instr, se::DeviceMemoryAllocator* allocator, + se::Stream* stream) { + XLA_SCOPED_LOGGING_TIMER(absl::StrCat( + "GpuConvAlgorithmPicker::PickBestAlgorithmImpl for ", instr->ToString())); + + const auto device_ordinal = stream_exec_->device_ordinal(); + std::vector operand_buffers; + + ScratchAllocator input_output_allocator(device_ordinal, allocator); + const auto initialize_buffer = [stream](DeviceMemoryBase buffer) { + // Although we don't have evidence this matters, zero out the buffers + // before autotuning. It's conceivable that using uninitialized memory as + // the inputs might affect performance if e.g. the inputs contain + // denormals, and this is easy enough. + stream->ThenMemZero(&buffer, buffer.size()); + }; + + // Allocate space for the input, filter, and output of the convolution. We + // use a ScratchAllocator for this instead of calling allocator_ directly so + // that our allocations don't leak. + for (const auto* operand : instr->operands()) { + TF_ASSIGN_OR_RETURN(auto buffer, + input_output_allocator.AllocateBytes( + ShapeUtil::ByteSizeOf(operand->shape()))); + initialize_buffer(buffer); + operand_buffers.push_back(buffer); + } + + TF_ASSIGN_OR_RETURN( + auto result_buffer, + input_output_allocator.AllocateBytes( + ShapeUtil::ByteSizeOf(instr->shape().tuple_shapes(0)))); + initialize_buffer(result_buffer); + + ScratchAllocator scratch_allocator(device_ordinal, allocator); + se::dnn::ProfileResult profile_result; + VLOG(3) << "Auto-tuning for " << instr->ToString(); + RunConvOptions options; + options.profile_result = &profile_result; + + // ROCm: Set the overriding algorithm to empty to remind cudnn_conv_runner + // that the AlgorithmConfig in running convolution needs to be empty + options.algo_override = se::dnn::AlgorithmDesc(); + + bool launch_ok = + RunGpuConv(instr, absl::MakeSpan(operand_buffers), result_buffer, + &scratch_allocator, stream, options) + .ok(); + + AutotuneResult best_result; + if (launch_ok && profile_result.is_valid()) { + best_result.mutable_conv()->set_algorithm( + profile_result.algorithm().algo_id()); + best_result.mutable_conv()->set_tensor_ops_enabled( + profile_result.algorithm().tensor_ops_enabled()); + int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes(); + best_result.set_scratch_bytes(scratch_bytes_used); + *best_result.mutable_run_time() = tensorflow::proto_utils::ToDurationProto( + absl::Milliseconds(profile_result.elapsed_time_in_ms())); + + return best_result; + } + + return InternalError( + "All algorithms tried for convolution %s failed. Falling back to " + "default algorithm.", + instr->ToString()); +} + +StatusOr GpuConvAlgorithmPicker::RunOnInstruction(HloInstruction* instr) { CHECK(IsCustomCallToDnnConvolution(*instr)); StatusOr best_algo_or = @@ -577,7 +704,7 @@ StatusOr CudnnConvAlgorithmPicker::RunOnInstruction( return true; } -StatusOr CudnnConvAlgorithmPicker::RunOnComputation( +StatusOr GpuConvAlgorithmPicker::RunOnComputation( HloComputation* computation) { std::vector convs; for (auto* instr : computation->instructions()) { @@ -594,11 +721,11 @@ StatusOr CudnnConvAlgorithmPicker::RunOnComputation( return changed; } -StatusOr CudnnConvAlgorithmPicker::Run(HloModule* module) { - XLA_SCOPED_LOGGING_TIMER("CudnnConvAlgorithmPicker"); +StatusOr GpuConvAlgorithmPicker::Run(HloModule* module) { + XLA_SCOPED_LOGGING_TIMER("GpuConvAlgorithmPicker"); if (module->config().debug_options().xla_gpu_disable_autotune()) { - VLOG(2) << "Convolution auto-tuning disabled, CudnnConvAlgorithmPicker " + VLOG(2) << "Convolution auto-tuning disabled, GpuConvAlgorithmPicker " "returning early."; return false; } diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h b/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.h similarity index 69% rename from tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h rename to tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.h index 75a1087002b..dddbe2ddfdc 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONV_ALGORITHM_PICKER_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONV_ALGORITHM_PICKER_H_ +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_CONV_ALGORITHM_PICKER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_CONV_ALGORITHM_PICKER_H_ #include "absl/time/time.h" #include "absl/types/optional.h" @@ -32,17 +32,17 @@ namespace gpu { // Modifies CustomCalls to cudnn convolutions, choosing the best algorithm for // each and adding explicit scratch space to the CustomCalls. -class CudnnConvAlgorithmPicker : public HloModulePass { +class GpuConvAlgorithmPicker : public HloModulePass { public: // If the `allocator` parameter is not null, we will use it to allocate temp // memory while timing the various convolution algorithms. If it's null, // we'll use the default allocator on the StreamExecutor. - CudnnConvAlgorithmPicker(se::StreamExecutor* stream_exec, - se::DeviceMemoryAllocator* allocator) + GpuConvAlgorithmPicker(se::StreamExecutor* stream_exec, + se::DeviceMemoryAllocator* allocator) : stream_exec_(stream_exec), allocator_(allocator) {} absl::string_view name() const override { - return "cudnn-conv-algorithm-picker"; + return "gpu-conv-algorithm-picker"; } StatusOr Run(HloModule* module) override; @@ -52,14 +52,19 @@ class CudnnConvAlgorithmPicker : public HloModulePass { StatusOr RunOnInstruction(HloInstruction* instr); StatusOr PickBestAlgorithm( const HloCustomCallInstruction* instr); - StatusOr PickBestAlgorithmNoCache( - const HloCustomCallInstruction* instr); - se::StreamExecutor* stream_exec_; // never null - se::DeviceMemoryAllocator* allocator_; // may be null + StatusOr PickBestAlgorithmNoCacheCuda( + const HloCustomCallInstruction* instr, + se::DeviceMemoryAllocator* allocator, se::Stream* stream); + + StatusOr PickBestAlgorithmNoCacheRocm( + const HloCustomCallInstruction* instr, + se::DeviceMemoryAllocator* allocator, se::Stream* stream); + + se::StreamExecutor* stream_exec_; // never null + se::DeviceMemoryAllocator* allocator_; // may be null }; } // namespace gpu } // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUDNN_CONV_ALGORITHM_PICKER_H_ +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_GPU_CONV_ALGORITHM_PICKER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/gpu_conv_runner.cc b/tensorflow/compiler/xla/service/gpu/gpu_conv_runner.cc index c1db7916b3c..07b6c9108ae 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_conv_runner.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_conv_runner.cc @@ -223,7 +223,17 @@ Status RunGpuConvImpl(const GpuConvParams& params, auto output_buf = se::DeviceMemory(params.output_buf); AlgorithmConfig algorithm = params.algorithm; - if (options.algo_override) { + // in ROCm mode, the first call to run the convolution needs to trigger the + // code that calls miopenFind* API. That triggger is implicit, it is based + // on whether or not the AlgorithmConfig::algorithm is empty! So for the + // first call we need to ensure that the AlgorithmConfig::algorithm is + // empty. For all subsequent calls, we should use the value retrieved from + // the backend_config + if ((stream->parent()->platform_kind() == se::PlatformKind::kROCm) && + (options.algo_override.has_value()) && + (*options.algo_override == se::dnn::AlgorithmDesc())) { + algorithm = AlgorithmConfig(); + } else if (options.algo_override.has_value()) { algorithm = AlgorithmConfig(*options.algo_override); } @@ -261,8 +271,13 @@ StatusOr GetGpuConvParams( const Shape* filter_shape; const Shape* output_shape; - params.algorithm = se::dnn::AlgorithmConfig(se::dnn::AlgorithmDesc( - backend_config.algorithm(), backend_config.tensor_ops_enabled())); + // The third field is scratch size stored from conv_algorithm_picker + // The operand is added to the shape field of the conv instruction + // in GpuConvAlgorithmPicker::RunOnInstruction() call. + params.algorithm = se::dnn::AlgorithmConfig( + se::dnn::AlgorithmDesc(backend_config.algorithm(), + backend_config.tensor_ops_enabled()), + conv->shape().tuple_shapes(1).dimensions(0)); params.conv_result_scale = backend_config.conv_result_scale(); switch (params.kind) { diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc index 95a61089563..489cbd101e2 100755 --- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc @@ -22,12 +22,12 @@ limitations under the License. #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" #include "tensorflow/compiler/xla/service/dump.h" #include "tensorflow/compiler/xla/service/gpu/cublas_gemm_pad_for_tensor_cores.h" -#include "tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_fused_conv_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/cudnn_pad_for_convolutions.h" #include "tensorflow/compiler/xla/service/gpu/cusolver_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/gemm_algorithm_picker.h" #include "tensorflow/compiler/xla/service/gpu/gemm_rewriter.h" +#include "tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.h" #include "tensorflow/compiler/xla/service/gpu/gpu_conv_padding_legalization.h" #include "tensorflow/compiler/xla/service/gpu/gpu_conv_rewriter.h" #include "tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h" @@ -188,11 +188,11 @@ Status NVPTXCompiler::OptimizeHloPostLayoutAssignment( // The new tuple and gte instructions then be simplified away, because // nobody is expected to use the scratch value. // - // However, if we were to run CudnnConvAlgorithmPicker after fusion + // However, if we were to run GpuConvAlgorithmPicker after fusion // the gte(customcall, 0) would probably already be into a fusion node. We // can't simplify across HloComputation boundaries, so in this case we // wouldn't be able to simplify away the new_tuple bits. - pipeline.AddPass(stream_exec, device_allocator); + pipeline.AddPass(stream_exec, device_allocator); // Find the fastest algorithm for GEMMs. pipeline.AddPass(stream_exec, device_allocator);