1 /*
2  * Copyright 2018 The Android Open Source Project
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 "Thread.h"
18 
19 #include <sched.h>
20 #include <unistd.h>
21 
22 namespace samples {
23 
getNumCpus()24 int32_t getNumCpus() {
25     static int32_t sNumCpus = []() {
26         pid_t pid = gettid();
27         cpu_set_t cpuSet;
28         CPU_ZERO(&cpuSet);
29         sched_getaffinity(pid, sizeof(cpuSet), &cpuSet);
30 
31         int32_t numCpus = 0;
32         while (CPU_ISSET(numCpus, &cpuSet)) {
33             ++numCpus;
34         }
35 
36         return numCpus;
37     }();
38 
39     return sNumCpus;
40 }
41 
setAffinity(int32_t cpu)42 void setAffinity(int32_t cpu) {
43     cpu_set_t cpuSet;
44     CPU_ZERO(&cpuSet);
45     CPU_SET(cpu, &cpuSet);
46     sched_setaffinity(gettid(), sizeof(cpuSet), &cpuSet);
47 }
48 
setAffinity(Affinity affinity)49 void setAffinity(Affinity affinity) {
50     const int32_t numCpus = getNumCpus();
51 
52     cpu_set_t cpuSet;
53     CPU_ZERO(&cpuSet);
54     for (int32_t cpu = 0; cpu < numCpus; ++cpu) {
55         switch (affinity) {
56             case Affinity::None:
57                 CPU_SET(cpu, &cpuSet);
58                 break;
59             case Affinity::Even:
60                 if (cpu % 2 == 0) CPU_SET(cpu, &cpuSet);
61                 break;
62             case Affinity::Odd:
63                 if (cpu % 2 == 1) CPU_SET(cpu, &cpuSet);
64                 break;
65         }
66     }
67 
68     sched_setaffinity(gettid(), sizeof(cpuSet), &cpuSet);
69 }
70 
71 } // namespace samples {
72