mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-12 12:47:21 +00:00
[Inference/Refactor] Refactor compilation mechanism and unified multi hw (#5613)
* refactor compilation mechanism and unified multi hw * fix file path bug * add init.py to make pybind a module to avoid relative path error caused by softlink * delete duplicated micros * fix micros bug in gcc
This commit is contained in:
60
extensions/csrc/kernel/cuda/utils/nvgpu_dev_info.h
Normal file
60
extensions/csrc/kernel/cuda/utils/nvgpu_dev_info.h
Normal file
@@ -0,0 +1,60 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <ostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "micros.h"
|
||||
|
||||
namespace colossalAI {
|
||||
namespace cuda {
|
||||
namespace utils {
|
||||
|
||||
class NVGPUDevInfo {
|
||||
public:
|
||||
explicit NVGPUDevInfo(int device_num) : device_num_(device_num) {
|
||||
CUDA_CHECK(cudaGetDeviceProperties(&prop_, device_num));
|
||||
}
|
||||
|
||||
std::array<int, 3> GetMaxGridDims() const {
|
||||
std::array<int, 3> ret;
|
||||
ret[0] = prop_.maxGridSize[0];
|
||||
ret[1] = prop_.maxGridSize[1];
|
||||
ret[2] = prop_.maxGridSize[2];
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::array<int, 3> GetMaxBlockDims() const {
|
||||
std::array<int, 3> ret;
|
||||
ret[0] = prop_.maxThreadsDim[0];
|
||||
ret[1] = prop_.maxThreadsDim[1];
|
||||
ret[2] = prop_.maxThreadsDim[2];
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::array<int, 2> GetCapability() const {
|
||||
std::array<int, 2> ret;
|
||||
ret[0] = prop_.major;
|
||||
ret[1] = prop_.minor;
|
||||
return ret;
|
||||
}
|
||||
|
||||
int GetMultiProcessorCount() const { return prop_.multiProcessorCount; }
|
||||
|
||||
int GetMaxThreadsPerMultiProcessor() const {
|
||||
return prop_.maxThreadsPerMultiProcessor;
|
||||
}
|
||||
|
||||
int GetMaxThreadsPerBlock() const { return prop_.maxThreadsPerBlock; }
|
||||
|
||||
private:
|
||||
int device_num_;
|
||||
cudaDeviceProp prop_;
|
||||
};
|
||||
|
||||
} // namespace utils
|
||||
} // namespace cuda
|
||||
} // namespace colossalAI
|
Reference in New Issue
Block a user