p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[MSLITE][Develop] add hardware judge for tensor rt delegate

This commit is contained in:
yangruoqi713 2021-07-02 17:14:46 +08:00
parent e1d25946bc
commit a530c86e31
1 changed files with 16 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
mindspore/lite/src/delegate/tensorrt/tensorrt_delegate.cc
View File

@ -15,7 +15,10 @@
*/
#include "src/delegate/tensorrt/tensorrt_delegate.h"
#include <cuda_runtime.h>
#include <vector>
#include <fstream>
#include <string>
#include "src/delegate/delegate_utils.h"
#include "src/delegate/tensorrt/op/activation_tensorrt.h"
#include "src/delegate/tensorrt/op/shape_tensorrt.h"
@ -31,7 +34,20 @@
#include "src/delegate/tensorrt/op/scale_tensorrt.h"
namespace mindspore::lite {
bool IsHardwareSupport() {
int driver_version = 0;
cudaDriverGetVersion(&driver_version);
if (driver_version == 0) {
MS_LOG(WARNING) << "No nvidia GPU driver.";
return false;
}
return true;
}
int TensorRTDelegate::Init() {
if (!IsHardwareSupport()) {
return RET_NOT_SUPPORT;
}
op_func_lists_.clear();
op_func_lists_ = {
{schema::PrimitiveType_Activation, GetTensorRTOp<ActivationTensorRT>},