Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/framework/gpu_util.py: 55%
11 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Contains GPU utility functions."""
17import collections
18import re
21# Matches the DeviceAttributes.physical_device_desc field.
22_PHYSICAL_DEVICE_DESCRIPTION_REGEX = re.compile(
23 r'name: ([^,]*), (?:.*compute capability: (\d+)\.(\d+))?')
26# compute_capability is a (major version, minor version) pair, or None if this
27# is not an Nvidia GPU.
28GpuInfo = collections.namedtuple('gpu_info', ['name', 'compute_capability'])
31def compute_capability_from_device_desc(device_attrs):
32 """Returns the GpuInfo given a DeviceAttributes proto.
34 Args:
35 device_attrs: A DeviceAttributes proto.
37 Returns
38 A gpu_info tuple. Both fields are None if `device_attrs` does not have a
39 valid physical_device_desc field.
40 """
41 # TODO(jingyue): The device description generator has to be in sync with
42 # this file. Another option is to put compute capability in
43 # DeviceAttributes, but I avoided that to keep DeviceAttributes
44 # target-independent. Reconsider this option when we have more things like
45 # this to keep in sync.
46 # LINT.IfChange
47 match = _PHYSICAL_DEVICE_DESCRIPTION_REGEX.search(
48 device_attrs.physical_device_desc)
49 # LINT.ThenChange(//tensorflow/core/common_runtime/gpu/gpu_device.cc)
50 if not match:
51 return GpuInfo(None, None)
52 cc = (int(match.group(2)), int(match.group(3))) if match.group(2) else None
53 return GpuInfo(match.group(1), cc)