提交 00f184d5 authored 作者: abergeron's avatar abergeron

Merge pull request #3308 from ivdorelian/patch-1

Update sandbox/cuda/dnn_base.c to work on Windows
......@@ -5,31 +5,44 @@ static cudnnHandle_t _handle = NULL;
static int
c_set_tensorNd(CudaNdarray *var, cudnnTensorDescriptor_t desc) {
int dim = CudaNdarray_NDIM(var);
int strides[dim];
int *strides = (int *)malloc(dim * sizeof(int));
int default_str = 1;
for (int i = dim-1; i >= 0; i--)
{
if (CudaNdarray_HOST_STRIDES(var)[i])
strides[i] = CudaNdarray_HOST_STRIDES(var)[i];
else
strides[i] = default_str;
default_str *= CudaNdarray_HOST_DIMS(var)[i];
int return_value = 0;
if (strides != NULL) {
for (int i = dim-1; i >= 0; i--)
{
if (CudaNdarray_HOST_STRIDES(var)[i])
strides[i] = CudaNdarray_HOST_STRIDES(var)[i];
else
strides[i] = default_str;
default_str *= CudaNdarray_HOST_DIMS(var)[i];
}
cudnnStatus_t err = cudnnSetTensorNdDescriptor(desc, CUDNN_DATA_FLOAT, dim,
CudaNdarray_HOST_DIMS(var),
strides);
if (err != CUDNN_STATUS_SUCCESS) {
PyErr_Format(PyExc_RuntimeError,
"Could not set tensorNd descriptor: %s"
"dim=%d",
cudnnGetErrorString(err), dim);
return_value = -1;
}
} else {
PyErr_Format(PyExc_MemoryError,
"Could not allocate memory for strides array of size %d.",
dim);
return_value = -1;
}
cudnnStatus_t err = cudnnSetTensorNdDescriptor(desc, CUDNN_DATA_FLOAT, dim,
CudaNdarray_HOST_DIMS(var),
strides);
if (err != CUDNN_STATUS_SUCCESS) {
PyErr_Format(PyExc_RuntimeError,
"Could not set tensorNd descriptor: %s"
"dim=%d",
cudnnGetErrorString(err), dim);
return -1;
}
return 0;
free(strides);
return return_value;
}
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论