提交 7a85fa42 authored 作者: Frederic's avatar Frederic

Detect if we are able to init cuDNN handle before auto enable it.

上级 959bfc79
......@@ -1851,7 +1851,8 @@ class GCC_compiler(object):
return (compilation_ok, run_ok, out, err)
@staticmethod
def try_flags(flag_list, preambule="", try_run=False, output=False):
def try_flags(flag_list, preambule="", body="",
try_run=False, output=False):
'''
Try to compile a dummy file with these flags.
......@@ -1865,6 +1866,7 @@ class GCC_compiler(object):
%(preambule)s
int main(int argc, char** argv)
{
%(body)s
return 0;
}
""" % locals())
......
......@@ -25,8 +25,26 @@ def dnn_available():
dnn_available.msg = "Device not supported by cuDNN"
dnn_available.avail = False
else:
preambule = """
#include <cudnn.h>
#include <stdio.h>
#include <cuda.h>
#include <cudnn_helper.h>
"""
body = """
cudnnHandle_t _handle = NULL;
cudnnStatus_t err;
if ((err = cudnnCreate(&_handle)) != CUDNN_STATUS_SUCCESS) {
fprintf(stderr, "could not create cuDNN handle: %s",
cudnnGetErrorString(err));
return 1;
}
"""
comp, run, out, err = gof.cmodule.GCC_compiler.try_flags(
["-l", "cudnn"], preambule="#include <cudnn.h>",
["-l", "cudnn", "-I" + os.path.dirname(__file__)],
preambule=preambule, body=body,
try_run=True, output=True)
dnn_available.avail = comp and run
......@@ -90,7 +108,7 @@ cudnnHandle_t _handle = NULL;
return ["""{
cudnnStatus_t err;
if ((err = cudnnCreate(&_handle)) != CUDNN_STATUS_SUCCESS) {
PyErr_Format(PyExc_RuntimeError, "could not create cudnn handle: %%s",
PyErr_Format(PyExc_RuntimeError, "could not create cuDNN handle: %%s",
cudnnGetErrorString(err));
return %s;
}
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论