You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
74 lines
2.0 KiB
74 lines
2.0 KiB
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
|
*
|
|
* 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
|
|
*
|
|
* http://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. */
|
|
|
|
#pragma once
|
|
|
|
#if defined(__OSX__) || defined(__APPLE__)
|
|
#include <sys/sysctl.h>
|
|
#include <sys/types.h>
|
|
#endif
|
|
#include "unsupported/Eigen/CXX11/Tensor"
|
|
|
|
namespace paddle {
|
|
|
|
#if defined(__ANDROID__)
|
|
int GetCpuCount() {
|
|
FILE* fp = fopen("/sys/devices/system/cpu/possible", "r");
|
|
if (!fp) {
|
|
return 1;
|
|
}
|
|
int rank0, rank1;
|
|
int num = fscanf(fp, "%d-%d", &rank0, &rank1);
|
|
fclose(fp);
|
|
if (num < 2) return 1;
|
|
return rank1 + 1;
|
|
}
|
|
#elif defined(__OSX__) || defined(__APPLE__)
|
|
int GetCpuCount() {
|
|
int count = 0;
|
|
size_t len = sizeof(int);
|
|
sysctlbyname("hw.ncpu", &count, &len, NULL, 0);
|
|
return count > 0 ? count : 1;
|
|
}
|
|
#else
|
|
int GetCpuCount() { return 1; }
|
|
#endif
|
|
|
|
class EigenDeviceWarpper {
|
|
public: // NOLINT
|
|
#if EIGEN_USE_THREADS
|
|
static Eigen::ThreadPoolDevice* device() {
|
|
const int num_cpus = GetCpuCount();
|
|
const int num_threads = (num_cpus > 2) ? 2 : num_cpus;
|
|
static Eigen::ThreadPool tp(num_threads);
|
|
static Eigen::ThreadPoolDevice* device =
|
|
new Eigen::ThreadPoolDevice(&tp, num_threads);
|
|
return device;
|
|
}
|
|
|
|
static void free_device(Eigen::ThreadPoolDevice* device) {
|
|
// do nothing
|
|
}
|
|
#else
|
|
static Eigen::DefaultDevice* device() {
|
|
Eigen::DefaultDevice* device = new Eigen::DefaultDevice;
|
|
return device;
|
|
}
|
|
|
|
static void free_device(Eigen::DefaultDevice* device) { delete device; }
|
|
#endif
|
|
};
|
|
|
|
} // namespace paddle
|