xref: /aosp_15_r20/external/pytorch/binaries/inspect_gpu.cc (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 /**
2  * Copyright (c) 2016-present, Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include <cuda_runtime.h>
18 
19 #include <sstream>
20 #include <vector>
21 
22 #include "c10/util/Flags.h"
23 #include "caffe2/core/common_gpu.h"
24 #include "caffe2/core/init.h"
25 #include "caffe2/core/logging.h"
26 
27 using std::vector;
28 
29 C10_DECLARE_int(caffe2_log_level);
30 
main(int argc,char ** argv)31 int main(int argc, char** argv) {
32   caffe2::GlobalInit(&argc, &argv);
33   c10::SetUsageMessage(
34       "Inspects the GPUs on the current machine and prints out their details "
35       "provided by cuda.");
36 
37   int gpu_count;
38   CUDA_ENFORCE(cudaGetDeviceCount(&gpu_count));
39   for (int i = 0; i < gpu_count; ++i) {
40     LOG(INFO) << "Querying device ID = " << i;
41     caffe2::DeviceQuery(i);
42   }
43 
44   vector<vector<bool> > access_pattern;
45   CAFFE_ENFORCE(caffe2::GetCudaPeerAccessPattern(&access_pattern));
46 
47   std::stringstream sstream;
48   // Find topology
49   for (int i = 0; i < gpu_count; ++i) {
50     for (int j = 0; j < gpu_count; ++j) {
51       sstream << (access_pattern[i][j] ? "+" : "-") << " ";
52     }
53     sstream << std::endl;
54   }
55   LOG(INFO) << "Access pattern: " << std::endl << sstream.str();
56 
57   return 0;
58 }
59