added create_reduce_tensor_descriptor
@@ -450,6 +450,22 @@ def create_op_tensor_descriptor(op_type, dtype):
return desc
+def create_reduce_tensor_descriptor(reduce_type, dtype):
+ desc = Descriptor(cudnn.createReduceTensorDescriptor(),
+ py_cudnn.destroyReduceTensorDescriptor)
+ data_type = get_data_type(dtype)
+
+ if reduce_type in (cudnn.CUDNN_REDUCE_TENSOR_MIN, cudnn.CUDNN_REDUCE_TENSOR_MAX):
+ indicies = cudnn.CUDNN_REDUCE_TENSOR_FLATTENED_INDICES
+ else:
+ indicies = cudnn.CUDNN_REDUCE_TENSOR_NO_INDICES
+ cudnn.setReduceTensorDescriptor(desc.value, reduce_type, data_type
+ cudnn.CUDNN_NOT_PROPAGATE_NAN,
+ indicies,
+ cudnn.CUDNN_32BIT_INDICES)
+ return desc
+
+
cpdef bint is_tensor_core_available(dtype) except *:
return (_cudnn_version >= 7000 and
dtype == numpy.float16 and
