message: Perform redzone checking in conv_ops_3d
diff: @@ -36,7 +36,12 @@ limitations under the License.
  #include "tensorflow/core/protobuf/autotuning.pb.h"
  #include "tensorflow/core/util/proto/proto_utils.h"
  using stream_executor::dnn::DimIndex;
  -#endif
  +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  +#if GOOGLE_CUDA
  +#include "tensorflow/stream_executor/cuda/ptxas_utils.h"
  +#include "tensorflow/stream_executor/cuda/redzone_allocator.h"
  +#include "tensorflow/stream_executor/tf_allocator_adapter.h"
  +#endif // GOOGLE_CUDA
  namespace tensorflow {
  @@ -399,6 +404,10 @@ struct LaunchConvOp<GPUDevice, T> {
  {{out_planes, out_rows, out_cols}}, out_depth),
  &transformed_output));
  + se::TfAllocatorAdapter tf_allocator_adapter(
  + stream->parent()->platform(), ctx->device()->GetAllocator({}));
  + se::cuda::RedzoneAllocator rz_allocator(stream, &tf_allocator_adapter,
  + se::cuda::PtxCompilationOptions());
  auto input_ptr = AsDeviceMemory(input.template flat<T>().data(),
  input.template flat<T>().size());
  auto filter_ptr =
  @@ -407,6 +416,8 @@ struct LaunchConvOp<GPUDevice, T> {
  auto output_ptr =
  AsDeviceMemory(transformed_output.template flat<T>().data(),
  transformed_output.template flat<T>().size());
  + se::DeviceMemory<T> output_ptr_rz(
  + WrapRedzoneBestEffort(&rz_allocator, output_ptr));
  static int64 ConvolveScratchSize = GetDnnWorkspaceLimit(
  "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32); // 4GB by default
  @@ -452,12 +463,19 @@ struct LaunchConvOp<GPUDevice, T> {
  // TODO(zhengxq): profile each algorithm multiple times to better
  // accuracy.
  DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx);
  + se::cuda::RedzoneAllocator rz_scratch_allocator(
  + stream, &tf_allocator_adapter, se::cuda::PtxCompilationOptions(),
  + /*memory_limit=*/ConvolveScratchSize);
  + se::ScratchAllocator* allocator_used =
  + !RedzoneCheckDisabled()
  + ? static_cast<se::ScratchAllocator*>(&rz_scratch_allocator)
  + : static_cast<se::ScratchAllocator*>(&scratch_allocator);
  ProfileResult profile_result;
  bool cudnn_launch_status =
  stream
  ->ThenConvolveWithAlgorithm(
  input_desc, input_ptr, filter_desc, filter_ptr, conv_desc,
  - output_desc, &output_ptr, &scratch_allocator,
  + output_desc, &output_ptr_rz, allocator_used,
  AlgorithmConfig(profile_algorithm), &profile_result)
  .ok();
  if (cudnn_launch_status) {
  @@ -467,9 +485,15 @@ struct LaunchConvOp<GPUDevice, T> {
  result.mutable_conv()->set_algorithm(profile_algorithm.algo_id());
  result.mutable_conv()->set_tensor_ops_enabled(
  profile_algorithm.tensor_ops_enabled());
  - result.set_scratch_bytes(scratch_allocator.TotalByteSize());
  + result.set_scratch_bytes(
  + !RedzoneCheckDisabled()
  + ? rz_scratch_allocator
  + .TotalAllocatedBytesExcludingRedzones()
  + : scratch_allocator.TotalByteSize());
  *result.mutable_run_time() = proto_utils::ToDurationProto(
  absl::Milliseconds(profile_result.elapsed_time_in_ms()));
  + CheckRedzones(rz_scratch_allocator, &result);
  + CheckRedzones(rz_allocator, &result);
  }
  }
  }
  