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# ==============================================================================
16"""Retrieves CUDA compute capability from NVIDIA webpage and creates a `.csv`.
17
18This module is mainly written to supplement for `../config_detector.py`
19which retrieves CUDA compute capability from existing golden file.
20
21The golden file resides inside `./golden` directory.
22
23Usage:
24  python cuda_compute_capability.py
25
26Output:
27  Creates `compute_capability.csv` file in the same directory by default. If
28  the file already exists, then it overwrites the file.
29
30  In order to use the new `.csv` as the golden, then it should replace the
31  original golden file (`./golden/compute_capability_golden.csv`) with the
32  same file name and path.
33"""
34import collections
35import difflib
36import os
37import re
38import urllib.request as urllib
39
40from absl import app
41from absl import flags
42
43
44FLAGS = flags.FLAGS
45PATH_TO_DIR = "tensorflow/tools/tensorflow_builder/config_detector"
46CUDA_CC_GOLDEN_DIR = PATH_TO_DIR + "/data/golden/compute_capability_golden.csv"
47
48
49def retrieve_from_web(generate_csv=False):
50  """Retrieves list of all CUDA compute capability from NVIDIA webpage.
51
52  Args:
53    generate_csv: Boolean for generating an output file containing
54                  the results.
55
56  Returns:
57    OrderedDict that is a list of all CUDA compute capability listed on the
58    NVIDIA page. Order goes from top to bottom of the webpage content (.html).
59  """
60  url = "https://developer.nvidia.com/cuda-gpus"
61  source = urllib.request.urlopen(url)
62  matches = []
63  while True:
64    line = source.readline()
65    if "</html>" in line:
66      break
67    else:
68      gpu = re.search(r"<a href=.*>([\w\S\s\d\[\]\,]+[^*])</a>(<a href=.*)?.*",
69                      line)
70      capability = re.search(
71          r"([\d]+).([\d]+)(/)?([\d]+)?(.)?([\d]+)?.*</td>.*", line)
72      if gpu:
73        matches.append(gpu.group(1))
74      elif capability:
75        if capability.group(3):
76          capability_str = capability.group(4) + "." + capability.group(6)
77        else:
78          capability_str = capability.group(1) + "." + capability.group(2)
79        matches.append(capability_str)
80
81  return create_gpu_capa_map(matches, generate_csv)
82
83
84def retrieve_from_golden():
85  """Retrieves list of all CUDA compute capability from a golden file.
86
87  The following file is set as default:
88    `./golden/compute_capability_golden.csv`
89
90  Returns:
91    Dictionary that lists of all CUDA compute capability in the following
92    format:
93      {'<GPU name>': ['<version major>.<version minor>', ...], ...}
94
95    If there are multiple versions available for a given GPU, then it
96    appends all supported versions in the value list (in the key-value
97    pair.)
98  """
99  out_dict = dict()
100  with open(CUDA_CC_GOLDEN_DIR) as g_file:
101    for line in g_file:
102      line_items = line.split(",")
103      val_list = []
104      for item in line_items[1:]:
105        val_list.append(item.strip("\n"))
106      out_dict[line_items[0]] = val_list
107
108  return out_dict
109
110
111def create_gpu_capa_map(match_list,
112                        generate_csv=False,
113                        filename="compute_capability"):
114  """Generates a map between GPU types and corresponding compute capability.
115
116  This method is used for retrieving CUDA compute capability from the web only.
117
118  Args:
119    match_list: List of all CUDA compute capability detected from the webpage.
120    generate_csv: Boolean for creating csv file to store results.
121    filename: String that is the name of the csv file (without `.csv` ending).
122
123  Returns:
124    OrderedDict that lists in the incoming order of all CUDA compute capability
125    provided as `match_list`.
126  """
127  gpu_capa = collections.OrderedDict()
128  include = False
129  gpu = ""
130  cnt = 0
131  mismatch_cnt = 0
132  for match in match_list:
133    if "Products" in match:
134      if not include:
135        include = True
136
137      continue
138    elif "www" in match:
139      include = False
140      break
141
142    if include:
143      if gpu:
144        if gpu in gpu_capa:
145          gpu_capa[gpu].append(match)
146        else:
147          gpu_capa[gpu] = [match]
148
149        gpu = ""
150        cnt += 1
151        if len(list(gpu_capa.keys())) < cnt:
152          mismatch_cnt += 1
153          cnt = len(list(gpu_capa.keys()))
154
155      else:
156        gpu = match
157
158  if generate_csv:
159    f_name = filename + ".csv"
160    write_csv_from_dict(f_name, gpu_capa)
161
162  return gpu_capa
163
164
165def write_csv_from_dict(filename, input_dict):
166  """Writes out a `.csv` file from an input dictionary.
167
168  After writing out the file, it checks the new list against the golden
169  to make sure golden file is up-to-date.
170
171  Args:
172    filename: String that is the output file name.
173    input_dict: Dictionary that is to be written out to a `.csv` file.
174  """
175  f = open(PATH_TO_DIR + "/data/" + filename, "w")
176  for k, v in input_dict.items():
177    line = k
178    for item in v:
179      line += "," + item
180
181    f.write(line + "\n")
182
183  f.flush()
184  print("Wrote to file %s" % filename)
185  check_with_golden(filename)
186
187
188def check_with_golden(filename):
189  """Checks the newly created CUDA compute capability file with the golden.
190
191  If differences are found, then it prints a list of all mismatches as
192  a `WARNING`.
193
194  Golden file must reside in `golden/` directory.
195
196  Args:
197    filename: String that is the name of the newly created file.
198  """
199  path_to_file = PATH_TO_DIR + "/data/" + filename
200  if os.path.isfile(path_to_file) and os.path.isfile(CUDA_CC_GOLDEN_DIR):
201    with open(path_to_file, "r") as f_new:
202      with open(CUDA_CC_GOLDEN_DIR, "r") as f_golden:
203        diff = difflib.unified_diff(
204            f_new.readlines(),
205            f_golden.readlines(),
206            fromfile=path_to_file,
207            tofile=CUDA_CC_GOLDEN_DIR
208        )
209        diff_list = []
210        for line in diff:
211          diff_list.append(line)
212
213        if diff_list:
214          print("WARNING: difference(s) found between new csv and golden csv.")
215          print(diff_list)
216        else:
217          print("No difference found between new csv and golen csv.")
218
219
220def print_dict(py_dict):
221  """Prints dictionary with formatting (2 column table).
222
223  Args:
224    py_dict: Dictionary that is to be printed out in a table format.
225  """
226  for gpu, cc in py_dict.items():
227    print("{:<25}{:<25}".format(gpu, cc))
228
229
230def main(argv):
231  if len(argv) > 2:
232    raise app.UsageError("Too many command-line arguments.")
233
234  retrieve_from_web(generate_csv=True)
235
236
237if __name__ == "__main__":
238  app.run(main)
239