1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
//@HEADER
// ************************************************************************
//
// Kokkos v. 4.0
// Copyright (2022) National Technology & Engineering
// Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
// See https://kokkos.org/LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//@HEADER
#include <Kokkos_Core.hpp>
#include <iostream>
#include <string>
#include <thread>
int get_num_devices() {
int num_devices;
#if defined(KOKKOS_ENABLE_CUDA)
KOKKOS_IMPL_CUDA_SAFE_CALL(cudaGetDeviceCount(&num_devices));
#elif defined(KOKKOS_ENABLE_HIP)
KOKKOS_IMPL_HIP_SAFE_CALL(hipGetDeviceCount(&num_devices));
#elif defined(KOKKOS_ENABLE_OPENMPTARGET)
num_devices = omp_get_num_devices();
#elif defined(KOKKOS_ENABLE_OPENACC)
num_devices = acc_get_num_devices(acc_get_device_type());
#elif defined(KOKKOS_ENABLE_SYCL)
num_devices = sycl::device::get_devices(sycl::info::device_type::gpu).size();
#else
num_devices = -1;
#endif
KOKKOS_ASSERT(num_devices == Kokkos::num_devices());
return num_devices;
}
int get_device_id() {
int device_id;
#if defined(KOKKOS_ENABLE_CUDA)
KOKKOS_IMPL_CUDA_SAFE_CALL(cudaGetDevice(&device_id));
#elif defined(KOKKOS_ENABLE_HIP)
KOKKOS_IMPL_HIP_SAFE_CALL(hipGetDevice(&device_id));
#elif defined(KOKKOS_ENABLE_OPENMPTARGET)
device_id = omp_get_default_device();
#elif defined(KOKKOS_ENABLE_OPENACC)
device_id = acc_get_device_num(acc_get_device_type());
#elif defined(KOKKOS_ENABLE_SYCL)
// Not able to query the underlying runtime because there is no such thing as
// device currently being used with SYCL. We go through the Kokkos runtime
// which makes the assert below pointless but it still let us check that
// Kokkos selected the device we asked for from the Python tests.
device_id = Kokkos::device_id();
#else
device_id = -1;
#endif
KOKKOS_ASSERT(device_id == Kokkos::device_id());
return device_id;
}
int get_max_threads() {
#if defined(KOKKOS_ENABLE_OPENMP)
return omp_get_max_threads();
#elif defined(KOKKOS_ENABLE_THREADS)
return std::thread::hardware_concurrency();
#else
return 1;
#endif
}
int get_hwloc_enabled() {
#ifdef KOKKOS_ENABLE_HWLOC
return 1;
#else
return 0;
#endif
}
int get_num_threads() {
int const num_threads = Kokkos::DefaultHostExecutionSpace().concurrency();
KOKKOS_ASSERT(num_threads == Kokkos::num_threads());
return num_threads;
}
int get_disable_warnings() { return !Kokkos::show_warnings(); }
int get_tune_internals() { return Kokkos::tune_internals(); }
int print_flag(std::string const& flag) {
std::vector<std::string> valid_flags;
#define KOKKOS_TEST_PRINT_FLAG(NAME) \
if (flag == #NAME) { \
std::cout << get_##NAME() << '\n'; \
return EXIT_SUCCESS; \
} \
valid_flags.push_back(#NAME)
KOKKOS_TEST_PRINT_FLAG(num_threads);
KOKKOS_TEST_PRINT_FLAG(max_threads);
KOKKOS_TEST_PRINT_FLAG(device_id);
KOKKOS_TEST_PRINT_FLAG(num_devices);
KOKKOS_TEST_PRINT_FLAG(disable_warnings);
KOKKOS_TEST_PRINT_FLAG(tune_internals);
KOKKOS_TEST_PRINT_FLAG(hwloc_enabled);
#undef KOKKOS_TEST_PRINT_FLAG
std::cerr << "Invalid flag name " << flag << ". Valid names are ";
for (int i = 0; i < (int)valid_flags.size() - 1; ++i) {
std::cerr << valid_flags[i] << ", ";
}
std::cerr << "and " << valid_flags.back() << ".\n";
return EXIT_FAILURE;
}
int main(int argc, char* argv[]) { // NOLINT(bugprone-exception-escape)
Kokkos::ScopeGuard guard(argc, argv);
if (argc != 2) {
std::cerr << "Usage: <executable> NAME_OF_FLAG\n";
return EXIT_FAILURE;
}
int exit_code = print_flag(argv[1]);
return exit_code;
}
|
