xref: /aosp_15_r20/external/tensorflow/third_party/gpus/find_rocm_config.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2020 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"""Prints ROCm library and header directories and versions found on the system.
16
17The script searches for ROCm library and header files on the system, inspects
18them to determine their version and prints the configuration to stdout.
19The path to inspect is specified through an environment variable (ROCM_PATH).
20If no valid configuration is found, the script prints to stderr and
21returns an error code.
22
23The script takes the directory specified by the ROCM_PATH environment variable.
24The script looks for headers and library files in a hard-coded set of
25subdirectories from base path of the specified directory. If ROCM_PATH is not
26specified, then "/opt/rocm" is used as it default value
27
28"""
29
30import io
31import os
32import re
33import sys
34
35
36class ConfigError(Exception):
37  pass
38
39
40def _get_default_rocm_path():
41  return "/opt/rocm"
42
43
44def _get_rocm_install_path():
45  """Determines and returns the ROCm installation path."""
46  rocm_install_path = _get_default_rocm_path()
47  if "ROCM_PATH" in os.environ:
48    rocm_install_path = os.environ["ROCM_PATH"]
49  # rocm_install_path = os.path.realpath(rocm_install_path)
50  return rocm_install_path
51
52
53def _get_composite_version_number(major, minor, patch):
54  return 10000 * major + 100 * minor + patch
55
56
57def _get_header_version(path, name):
58  """Returns preprocessor defines in C header file."""
59  for line in io.open(path, "r", encoding="utf-8"):
60    match = re.match(r"#define %s +(\d+)" % name, line)
61    if match:
62      value = match.group(1)
63      return int(value)
64
65  raise ConfigError('#define "{}" is either\n'.format(name) +
66                    "  not present in file {} OR\n".format(path) +
67                    "  its value is not an integer literal")
68
69
70def _find_rocm_config(rocm_install_path):
71
72  def rocm_version_numbers_pre_rocm50(path, prior_err):
73    version_file = os.path.join(path, ".info/version-dev")
74    if not os.path.exists(version_file):
75      raise ConfigError("{} ROCm version file ".format(prior_err) +
76                        '"{}" not found either.'.format(version_file))
77    version_numbers = []
78    with open(version_file) as f:
79      version_string = f.read().strip()
80      version_numbers = version_string.split(".")
81    major = int(version_numbers[0])
82    minor = int(version_numbers[1])
83    patch = int(version_numbers[2].split("-")[0])
84    return major, minor, patch
85
86  def rocm_version_numbers_post_rocm50(path):
87    version_file = os.path.join(path, "include/rocm_version.h")
88    if not os.path.exists(version_file):
89      return False, 'ROCm version file "{}" not found.'.format(version_file) +\
90        " Trying an alternate approach to determine the ROCm version.", 0,0,0
91    major = _get_header_version(version_file, "ROCM_VERSION_MAJOR")
92    minor = _get_header_version(version_file, "ROCM_VERSION_MINOR")
93    patch = _get_header_version(version_file, "ROCM_VERSION_PATCH")
94    return True, "", major, minor, patch
95
96  status, error_msg, major, minor, patch = \
97    rocm_version_numbers_post_rocm50(rocm_install_path)
98
99  if not status:
100    major, minor, patch = \
101      rocm_version_numbers_pre_rocm50(rocm_install_path, error_msg)
102
103  rocm_config = {
104      "rocm_version_number": _get_composite_version_number(major, minor, patch)
105  }
106
107  return rocm_config
108
109
110def _find_hipruntime_config(rocm_install_path):
111
112  def hipruntime_version_number(path):
113    version_file = os.path.join(path, "hip/include/hip/hip_version.h")
114    if not os.path.exists(version_file):
115      raise ConfigError(
116          'HIP Runtime version file "{}" not found'.format(version_file))
117    # This header file has an explicit #define for HIP_VERSION, whose value
118    # is (HIP_VERSION_MAJOR * 100 + HIP_VERSION_MINOR)
119    # Retreive the major + minor and re-calculate here, since we do not
120    # want get into the business of parsing arith exprs
121    major = _get_header_version(version_file, "HIP_VERSION_MAJOR")
122    minor = _get_header_version(version_file, "HIP_VERSION_MINOR")
123    return 100 * major + minor
124
125  hipruntime_config = {
126      "hipruntime_version_number": hipruntime_version_number(rocm_install_path)
127  }
128
129  return hipruntime_config
130
131
132def _find_miopen_config(rocm_install_path):
133
134  def miopen_version_numbers(path):
135    version_file = os.path.join(path, "miopen/include/miopen/version.h")
136    if not os.path.exists(version_file):
137      raise ConfigError(
138          'MIOpen version file "{}" not found'.format(version_file))
139    major = _get_header_version(version_file, "MIOPEN_VERSION_MAJOR")
140    minor = _get_header_version(version_file, "MIOPEN_VERSION_MINOR")
141    patch = _get_header_version(version_file, "MIOPEN_VERSION_PATCH")
142    return major, minor, patch
143
144  major, minor, patch = miopen_version_numbers(rocm_install_path)
145
146  miopen_config = {
147      "miopen_version_number":
148          _get_composite_version_number(major, minor, patch)
149  }
150
151  return miopen_config
152
153
154def _find_rocblas_config(rocm_install_path):
155
156  def rocblas_version_numbers(path):
157    possible_version_files = [
158        "rocblas/include/rocblas-version.h",  # ROCm 3.7 and prior
159        "rocblas/include/internal/rocblas-version.h",  # ROCm 3.8
160    ]
161    version_file = None
162    for f in possible_version_files:
163      version_file_path = os.path.join(path, f)
164      if os.path.exists(version_file_path):
165        version_file = version_file_path
166        break
167    if not version_file:
168      raise ConfigError(
169          "rocblas version file not found in {}".format(
170              possible_version_files))
171    major = _get_header_version(version_file, "ROCBLAS_VERSION_MAJOR")
172    minor = _get_header_version(version_file, "ROCBLAS_VERSION_MINOR")
173    patch = _get_header_version(version_file, "ROCBLAS_VERSION_PATCH")
174    return major, minor, patch
175
176  major, minor, patch = rocblas_version_numbers(rocm_install_path)
177
178  rocblas_config = {
179      "rocblas_version_number":
180          _get_composite_version_number(major, minor, patch)
181  }
182
183  return rocblas_config
184
185
186def _find_rocrand_config(rocm_install_path):
187
188  def rocrand_version_number(path):
189    possible_version_files = [
190        "rocrand/include/rocrand_version.h",  # ROCm 5.0 and prior
191        "include/rocrand/rocrand_version.h",  # ROCm 5.1
192    ]
193    version_file = None
194    for f in possible_version_files:
195      version_file_path = os.path.join(path, f)
196      if os.path.exists(version_file_path):
197        version_file = version_file_path
198        break
199    if not version_file:
200      raise ConfigError(
201          "rocrand version file not found in {}".format(possible_version_files))
202    version_number = _get_header_version(version_file, "ROCRAND_VERSION")
203    return version_number
204
205  rocrand_config = {
206      "rocrand_version_number": rocrand_version_number(rocm_install_path)
207  }
208
209  return rocrand_config
210
211
212def _find_rocfft_config(rocm_install_path):
213
214  def rocfft_version_numbers(path):
215    version_file = os.path.join(path, "rocfft/include/rocfft-version.h")
216    if not os.path.exists(version_file):
217      raise ConfigError(
218          'rocfft version file "{}" not found'.format(version_file))
219    major = _get_header_version(version_file, "rocfft_version_major")
220    minor = _get_header_version(version_file, "rocfft_version_minor")
221    patch = _get_header_version(version_file, "rocfft_version_patch")
222    return major, minor, patch
223
224  major, minor, patch = rocfft_version_numbers(rocm_install_path)
225
226  rocfft_config = {
227      "rocfft_version_number":
228          _get_composite_version_number(major, minor, patch)
229  }
230
231  return rocfft_config
232
233
234def _find_hipfft_config(rocm_install_path):
235
236  def hipfft_version_numbers(path):
237    version_file = os.path.join(path, "hipfft/include/hipfft-version.h")
238    if not os.path.exists(version_file):
239      raise ConfigError(
240          'hipfft version file "{}" not found'.format(version_file))
241    major = _get_header_version(version_file, "hipfftVersionMajor")
242    minor = _get_header_version(version_file, "hipfftVersionMinor")
243    patch = _get_header_version(version_file, "hipfftVersionPatch")
244    return major, minor, patch
245
246  major, minor, patch = hipfft_version_numbers(rocm_install_path)
247
248  hipfft_config = {
249      "hipfft_version_number":
250          _get_composite_version_number(major, minor, patch)
251  }
252
253  return hipfft_config
254
255
256def _find_roctracer_config(rocm_install_path):
257
258  def roctracer_version_numbers(path):
259    version_file = os.path.join(path, "roctracer/include/roctracer.h")
260    if not os.path.exists(version_file):
261      raise ConfigError(
262          'roctracer version file "{}" not found'.format(version_file))
263    major = _get_header_version(version_file, "ROCTRACER_VERSION_MAJOR")
264    minor = _get_header_version(version_file, "ROCTRACER_VERSION_MINOR")
265    # roctracer header does not have a patch version number
266    patch = 0
267    return major, minor, patch
268
269  major, minor, patch = roctracer_version_numbers(rocm_install_path)
270
271  roctracer_config = {
272      "roctracer_version_number":
273          _get_composite_version_number(major, minor, patch)
274  }
275
276  return roctracer_config
277
278
279def _find_hipsparse_config(rocm_install_path):
280
281  def hipsparse_version_numbers(path):
282    version_file = os.path.join(path, "hipsparse/include/hipsparse-version.h")
283    if not os.path.exists(version_file):
284      raise ConfigError(
285          'hipsparse version file "{}" not found'.format(version_file))
286    major = _get_header_version(version_file, "hipsparseVersionMajor")
287    minor = _get_header_version(version_file, "hipsparseVersionMinor")
288    patch = _get_header_version(version_file, "hipsparseVersionPatch")
289    return major, minor, patch
290
291  major, minor, patch = hipsparse_version_numbers(rocm_install_path)
292
293  hipsparse_config = {
294      "hipsparse_version_number":
295          _get_composite_version_number(major, minor, patch)
296  }
297
298  return hipsparse_config
299
300def _find_hipsolver_config(rocm_install_path):
301
302  def hipsolver_version_numbers(path):
303    possible_version_files = [
304        "hipsolver/include/hipsolver-version.h",  # ROCm 5.0 and prior
305        "hipsolver/include/internal/hipsolver-version.h",  # ROCm 5.1
306    ]
307    version_file = None
308    for f in possible_version_files:
309      version_file_path = os.path.join(path, f)
310      if os.path.exists(version_file_path):
311        version_file = version_file_path
312        break
313    if not version_file:
314      raise ConfigError("hipsolver version file not found in {}".format(
315          possible_version_files))
316    major = _get_header_version(version_file, "hipsolverVersionMajor")
317    minor = _get_header_version(version_file, "hipsolverVersionMinor")
318    patch = _get_header_version(version_file, "hipsolverVersionPatch")
319    return major, minor, patch
320
321  major, minor, patch = hipsolver_version_numbers(rocm_install_path)
322
323  hipsolver_config = {
324      "hipsolver_version_number":
325          _get_composite_version_number(major, minor, patch)
326  }
327
328  return hipsolver_config
329
330
331def _find_rocsolver_config(rocm_install_path):
332
333  def rocsolver_version_numbers(path):
334    version_file = os.path.join(path, "rocsolver/include/rocsolver-version.h")
335    if not os.path.exists(version_file):
336      raise ConfigError(
337          'rocsolver version file "{}" not found'.format(version_file))
338    major = _get_header_version(version_file, "ROCSOLVER_VERSION_MAJOR")
339    minor = _get_header_version(version_file, "ROCSOLVER_VERSION_MINOR")
340    patch = _get_header_version(version_file, "ROCSOLVER_VERSION_PATCH")
341    return major, minor, patch
342
343  major, minor, patch = rocsolver_version_numbers(rocm_install_path)
344
345  rocsolver_config = {
346      "rocsolver_version_number":
347          _get_composite_version_number(major, minor, patch)
348  }
349
350  return rocsolver_config
351
352
353def find_rocm_config():
354  """Returns a dictionary of ROCm components config info."""
355  rocm_install_path = _get_rocm_install_path()
356  if not os.path.exists(rocm_install_path):
357    raise ConfigError(
358        'Specified ROCM_PATH "{}" does not exist'.format(rocm_install_path))
359
360  result = {}
361
362  result["rocm_toolkit_path"] = rocm_install_path
363  result.update(_find_rocm_config(rocm_install_path))
364  result.update(_find_hipruntime_config(rocm_install_path))
365  result.update(_find_miopen_config(rocm_install_path))
366  result.update(_find_rocblas_config(rocm_install_path))
367  result.update(_find_rocrand_config(rocm_install_path))
368  result.update(_find_rocfft_config(rocm_install_path))
369  if result["rocm_version_number"] >= 40100:
370    result.update(_find_hipfft_config(rocm_install_path))
371  result.update(_find_roctracer_config(rocm_install_path))
372  result.update(_find_hipsparse_config(rocm_install_path))
373  if result["rocm_version_number"] >= 40500:
374    result.update(_find_hipsolver_config(rocm_install_path))
375  result.update(_find_rocsolver_config(rocm_install_path))
376
377  return result
378
379
380def main():
381  try:
382    for key, value in sorted(find_rocm_config().items()):
383      print("%s: %s" % (key, value))
384  except ConfigError as e:
385    sys.stderr.write("\nERROR: {}\n\n".format(str(e)))
386    sys.exit(1)
387
388
389if __name__ == "__main__":
390  main()
391