pytorch
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
27using std::vector;28
29C10_DECLARE_int(caffe2_log_level);30
31int main(int argc, char** argv) {32caffe2::GlobalInit(&argc, &argv);33c10::SetUsageMessage(34"Inspects the GPUs on the current machine and prints out their details "35"provided by cuda.");36
37int gpu_count;38CUDA_ENFORCE(cudaGetDeviceCount(&gpu_count));39for (int i = 0; i < gpu_count; ++i) {40LOG(INFO) << "Querying device ID = " << i;41caffe2::DeviceQuery(i);42}43
44vector<vector<bool> > access_pattern;45CAFFE_ENFORCE(caffe2::GetCudaPeerAccessPattern(&access_pattern));46
47std::stringstream sstream;48// Find topology49for (int i = 0; i < gpu_count; ++i) {50for (int j = 0; j < gpu_count; ++j) {51sstream << (access_pattern[i][j] ? "+" : "-") << " ";52}53sstream << std::endl;54}55LOG(INFO) << "Access pattern: " << std::endl << sstream.str();56
57return 0;58}
59