[XLA:GPU] Make GpuBfloat16Support compatible with AOT compilation

anlunx · tensorflower-gardener · commit d6f5d7891537 · 2022-12-15T12:18:18.000-08:00
GpuBfloat16Support should work without stream executor.

PiperOrigin-RevId: 495661838
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
@@ -192,7 +192,14 @@ class GpuBfloat16Support : public BFloat16Support {
   explicit GpuBfloat16Support(bool supports_matrix_multiplication,
                               se::StreamExecutor* stream_exec)
       : supports_matrix_multiplication_(supports_matrix_multiplication),
-        stream_exec_(stream_exec) {}
+        is_conv_bf16_supported_(IsConvBf16Supported(stream_exec)) {}
+
+  explicit GpuBfloat16Support(bool supports_matrix_multiplication,
+                              se::dnn::VersionInfo cudnn_version,
+                              se::CudaComputeCapability cuda_compute_capability)
+      : supports_matrix_multiplication_(supports_matrix_multiplication),
+        is_conv_bf16_supported_(
+            IsConvBf16Supported(cudnn_version, cuda_compute_capability)) {}
 
   bool SupportsBF16Operand(const HloInstruction& hlo,
                            int64_t operand_index) const override {
@@ -235,30 +242,41 @@ class GpuBfloat16Support : public BFloat16Support {
       case HloOpcode::kBitcast:
         return true;
       case HloOpcode::kConvolution:
-        return IsConvBF16Supported();
+        return is_conv_bf16_supported_;
       default:
         return supports_matrix_multiplication_ &&
                gpu::IsMatrixMultiplication(hlo);
     }
   }
 
-  bool IsConvBF16Supported() const {
-    if (se::dnn::DnnSupport* dnn = stream_exec_->AsDnn()) {
+  static bool IsConvBf16Supported(se::StreamExecutor* stream_exec) {
+    if (se::dnn::DnnSupport* dnn = stream_exec->AsDnn()) {
       se::port::StatusOr<se::dnn::VersionInfo> cudnn_version =
           dnn->GetVersion();
-      return cudnn_version.ok() &&
-             (cudnn_version->major_version() > 8 ||
-              (cudnn_version->major_version() == 8 &&
-               cudnn_version->minor_version() >= 2)) &&
-             stream_exec_->GetDeviceDescription()
-                 .cuda_compute_capability()
-                 .IsAtLeast(se::CudaComputeCapability::AMPERE);
+      if (cudnn_version.ok()) {
+        auto cuda_compute_capability =
+            stream_exec->GetDeviceDescription().cuda_compute_capability();
+        return (cudnn_version->major_version() > 8 ||
+                (cudnn_version->major_version() == 8 &&
+                 cudnn_version->minor_version() >= 2)) &&
+               cuda_compute_capability.IsAtLeast(
+                   se::CudaComputeCapability::AMPERE);
+      }
     }
     return false;
   }
 
+  static bool IsConvBf16Supported(
+      se::dnn::VersionInfo cudnn_version,
+      se::CudaComputeCapability cuda_compute_capability) {
+    return (cudnn_version.major_version() > 8 ||
+            (cudnn_version.major_version() == 8 &&
+             cudnn_version.minor_version() >= 2)) &&
+           cuda_compute_capability.IsAtLeast(se::CudaComputeCapability::AMPERE);
+  }
+
   bool supports_matrix_multiplication_;
-  se::StreamExecutor* stream_exec_;
+  bool is_conv_bf16_supported_;
 };
 
 int64_t GetSizeOfShape(const Shape& shape, int pointer_size) {
