[XLA:GPU] Failed to determine best cudnn convolution algorithm

I got Failed to determine best cudnn convolution algorithm error when running facebook/wav2vec2-base-960h model using torch_xla on GPU in fp16 mode. This error only occurs when using fp16 and fp32 works fine.

Minimal HLO to reproduce:

HloModule Test

ENTRY main {
  x = f16[44,768,1,49]{3,2,1,0} parameter(0)
  y = f16[44,768,1,50]{3,2,1,0} parameter(1)
  ROOT %convolution.10022 = f16[1,128,48,768]{3,2,1,0} convolution(f16[44,768,1,49]{3,2,1,0} %x, f16[44,768,1,50]{3,2,1,0} %y), window={size=1x50 pad=0_0x64_64}, dim_labels=fb01_io01->01bf, batch_group_count=16
}

The following is the full log file.

(base) ubuntu@xla-p3-8x:~/src/tensorflow/xla_benchmark$ ../bazel-bin/tensorflow/compiler/xla/tools/replay_computation_gpu --use_fake_data=true --num_runs=1 --print_result=false conv.hlo 
2022-04-15 00:10:23.507808: I tensorflow/compiler/xla/service/platform_util.cc:69] platform Host present but no XLA compiler available: could not find registered compiler for platform Host -- check target linkage (hint: try adding tensorflow/compiler/jit:xla_cpu_jit as a dependency)
2022-04-15 00:10:24.852458: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:24.935392: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:24.944338: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:24.967499: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:24.983119: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:25.002219: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:25.025011: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:25.050733: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:983] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2022-04-15 00:10:25.052806: I tensorflow/compiler/xla/service/service.cc:174] XLA service 0x562e5c814b50 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2022-04-15 00:10:25.052832: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (0): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052839: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (1): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052844: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (2): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052849: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (3): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052854: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (4): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052859: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (5): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052863: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (6): Tesla V100-SXM2-32GB, Compute Capability 7.0
2022-04-15 00:10:25.052868: I tensorflow/compiler/xla/service/service.cc:182]   StreamExecutor device (7): Tesla V100-SXM2-32GB, Compute Capability 7.0
conv.hlo: is not HloSnapshot. Trying HloProto.
conv.hlo: is not HloProto. Trying HLO text.
2022-04-15 00:10:25.053972: I tensorflow/compiler/xla/tools/replay_computation.cc:470] Compiling 1 modules in parallel.
2022-04-15 00:10:25.622602: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] Loaded cuDNN version 8204
2022-04-15 00:10:26.108045: I tensorflow/compiler/xla/tools/replay_computation.cc:487] Done compiling; now running the modules.
2022-04-15 00:10:26.108867: E tensorflow/compiler/xla/tools/replay_computation.cc:491] Compilation failed: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv = (f16[1,128,48,768]{3,1,0,2}, u8[0]{0}) custom-call(f16[704,48,1,49]{0,3,2,1} %bitcast.2, f16[48,768,1,50]{0,3,2,1} %pad), window={size=1x50 pad=0_0x64_64}, dim_labels=fb01_io01->01bf, feature_group_count=16, custom_call_target="__cudnn$convForward", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: UNKNOWN: CUDNN_STATUS_BAD_PARAM
in tensorflow/stream_executor/cuda/cuda_dnn.cc(3520): 'op' CUDNN_BACKEND_OPERATION: cudnnFinalize Failed

To ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning.: hlo { hlo_module { name: "Test" entry_computation_name: "main" computations { name: "main" instructions { name: "x" opcode: "parameter" shape { element_type: F16 dimensions: 44 dimensions: 768 dimensions: 1 dimensions: 49 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } metadata { } frontend_attributes { } } instructions { name: "y" opcode: "parameter" shape { element_type: F16 dimensions: 44 dimensions: 768 dimensions: 1 dimensions: 50 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } metadata { } parameter_number: 1 id: 1 frontend_attributes { } } instructions { name: "convolution.10022" opcode: "convolution" shape { element_type: F16 dimensions: 1 dimensions: 128 dimensions: 48 dimensions: 768 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } metadata { } window { dimensions { size: 1 stride: 1 window_dilation: 1 base_dilation: 1 } dimensions { size: 50 stride: 1 padding_low: 64 padding_high: 64 window_dilation: 1 base_dilation: 1 } } convolution_dimension_numbers { kernel_output_feature_dimension: 1 kernel_spatial_dimensions: 2 kernel_spatial_dimensions: 3 input_batch_dimension: 1 output_batch_dimension: 2 output_feature_dimension: 3 input_spatial_dimensions: 2 input_spatial_dimensions: 3 output_spatial_dimensions: 0 output_spatial_dimensions: 1 } id: 2 operand_ids: 0 operand_ids: 1 feature_group_count: 1 precision_config { operand_precision: DEFAULT operand_precision: DEFAULT } batch_group_count: 16 frontend_attributes { } } program_shape { parameters { element_type: F16 dimensions: 44 dimensions: 768 dimensions: 1 dimensions: 49 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } parameters { element_type: F16 dimensions: 44 dimensions: 768 dimensions: 1 dimensions: 50 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } result { element_type: F16 dimensions: 1 dimensions: 128 dimensions: 48 dimensions: 768 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } parameter_names: "x" parameter_names: "y" } id: 2 root_id: 2 } host_program_shape { parameters { element_type: F16 dimensions: 44 dimensions: 768 dimensions: 1 dimensions: 49 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } parameters { element_type: F16 dimensions: 44 dimensions: 768 dimensions: 1 dimensions: 50 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } result { element_type: F16 dimensions: 1 dimensions: 128 dimensions: 48 dimensions: 768 layout { minor_to_major: 3 minor_to_major: 2 minor_to_major: 1 minor_to_major: 0 format: DENSE } is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false is_dynamic_dimension: false } parameter_names: "p0" parameter_names: "p1" } entry_computation_id: 2 input_output_alias { } dynamic_parameter_binding { } } }

Tested with tensorflow commit 75861c43005523e2552bb3f85b2f0defc16ea9cf, CUDA 11.4, CUDNN 8.2.