/
decide.cpp
104 lines (86 loc) · 2.72 KB
/
decide.cpp
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
// Copyright 2019 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <pybind11/pybind11.h>
namespace py = pybind11;
#ifdef _WIN32
// Windows
#include <intrin.h>
#define cpuid(info, x) __cpuidex(info, x, 0)
#else
// GCC Intrinsics
#include <cpuid.h>
void cpuid(int info[4], int infoType){
__cpuid_count(infoType, 0, info[0], info[1], info[2], info[3]);
}
#endif
enum Instructions { AVX512F = 0, AVX2 = 1, SSE4_1 = 2, BASIC = 3};
int detect_instructions() {
Instructions instr = BASIC;
int info[4];
cpuid(info, 0);
int nIds = info[0];
if (nIds >= 1) {
cpuid(info, 1);
if ((info[2] & (1 << 19)) != 0) {
instr = SSE4_1;
}
}
if (nIds >= 7) {
cpuid(info, 7);
if ((info[1] & (1 << 5))!= 0) {
instr = AVX2;
}
if ((info[1] & (1 << 16)) != 0) {
instr = AVX512F;
}
}
return static_cast<int>(instr);
}
enum GPUCapabilities {
CUDA = 0, CUSTATEVEC = 1, NO_GPU = 10, NO_CUSTATEVEC = 11 };
// For now, GPU detection is performed at compile time, as our wheels are
// generated on Github Actions runners which do not have GPU support.
//
// Users wishing to use qsim with GPU will need to compile locally on a device
// which has the necessary CUDA toolkit.
int detect_gpu() {
#ifdef __NVCC__
GPUCapabilities gpu = CUDA;
#else
GPUCapabilities gpu = NO_GPU;
#endif
return gpu;
}
// For now, cuStateVec detection is performed at compile time, as our wheels
// are generated on Github Actions runners which do not have GPU support.
//
// Users wishing to use qsim with cuStateVec will need to compile locally on
// a device which has the necessary CUDA toolkit and cuStateVec library.
int detect_custatevec() {
#if defined(__NVCC__) && defined(__CUSTATEVEC__)
GPUCapabilities gpu = CUSTATEVEC;
#else
GPUCapabilities gpu = NO_CUSTATEVEC;
#endif
return gpu;
}
PYBIND11_MODULE(qsim_decide, m) {
m.doc() = "pybind11 plugin"; // optional module docstring
// Methods for returning amplitudes
m.def("detect_instructions", &detect_instructions, "Detect SIMD");
// Detect available GPUs.
m.def("detect_gpu", &detect_gpu, "Detect GPU");
// Detect cuStateVec.
m.def("detect_custatevec", &detect_custatevec, "Detect cuStateVec");
}