Add work-around for issue in cutensorReduction of cuTENSOR 1.2.1
@@ -72,6 +72,20 @@ cdef dict _dict_contraction_v10200 = {
_linalg.COMPUTE_TYPE_FP64: cutensor.COMPUTE_64F,
_linalg.COMPUTE_TYPE_FP32: cutensor.COMPUTE_32F},
}
+cdef dict _dict_compute_type = {
+ 'e': cutensor.R_MIN_16F,
+ 'f': cutensor.R_MIN_32F,
+ 'd': cutensor.R_MIN_64F,
+ 'F': cutensor.C_MIN_32F,
+ 'D': cutensor.C_MIN_64F,
+}
+cdef dict _dict_compute_type_v10200 = {
+ 'e': cutensor.COMPUTE_16F,
+ 'f': cutensor.COMPUTE_32F,
+ 'd': cutensor.COMPUTE_64F,
+ 'F': cutensor.COMPUTE_32F,
+ 'D': cutensor.COMPUTE_64F,
+}
cdef class Mode(object):
@@ -131,19 +145,16 @@ cdef int _get_cuda_dtype(numpy_dtype) except -1:
raise TypeError('Dtype {} is not supported'.format(numpy_dtype))
-cdef int _get_cutensor_dtype(numpy_dtype) except -1:
- if numpy_dtype == numpy.float16:
- return cutensor.R_MIN_16F
- elif numpy_dtype == numpy.float32:
- return cutensor.R_MIN_32F
- elif numpy_dtype == numpy.float64:
- return cutensor.R_MIN_64F
- elif numpy_dtype == numpy.complex64:
- return cutensor.C_MIN_32F
- elif numpy_dtype == numpy.complex128:
- return cutensor.C_MIN_64F
+cdef int _get_cutensor_compute_type(numpy_dtype) except -1:
+ if cutensor.get_version() >= 10200:
+ # version 1.2.0 or later
+ dict_compute_type = _dict_compute_type_v10200
else:
+ dict_compute_type = _dict_compute_type
+ key = numpy_dtype.char
+ if key not in dict_compute_type:
raise TypeError('Dtype {} is not supported'.format(numpy_dtype))
+ return dict_compute_type[key]
def create_mode(*mode):
@@ -668,7 +679,7 @@ def reduction(
A, desc_A, _auto_create_mode(A, mode_A),
_create_scalar(beta, compute_dtype),
C, desc_C, _auto_create_mode(C, mode_C),
- reduce_op, _get_cutensor_dtype(compute_dtype)
+ reduce_op, _get_cutensor_compute_type(compute_dtype)
)
@@ -783,6 +794,6 @@ def _try_reduction_routine(
out_arg,
desc_out,
_create_mode_with_cache(out_axis),
- reduce_op, _get_cutensor_dtype(compute_dtype))
+ reduce_op, _get_cutensor_compute_type(compute_dtype))
return out
