Merge pull request #5755 from notoraptor/op-param-gpudnnsoftmax

doc/library/gof/params_type.txt

@@ -12,4 +12,5 @@ Reference
    :platform: Unix, Windows
    :synopsis: Wrapper class for op params
    :members:
-.. moduleauthor:: LISA
+   :member-order: bysource
\ No newline at end of file
+.. moduleauthor:: LISA

theano/gof/op.py

@@ -808,7 +808,7 @@ class Op(utils.object2, PureOp, CLinkerOp):
                 field = wrapper.fields[i]
                 _type = wrapper.types[i]
                 wrap_dict[field] = _type.filter(getattr(self, field), strict=False, allow_downcast=True)
-            return theano.gof.Params(wrapper, **wrap_dict)
+            return self.params_type.get_params(self)
         raise theano.gof.utils.MethodNotDefined('get_params')
 
     def prepare_node(self, node, storage_map, compute_map, impl):

theano/gof/tests/test_params_type.py

@@ -6,7 +6,7 @@ from theano.gof import Op, COp, Apply
 from theano import Generic
 from theano.scalar import Scalar
 from theano.tensor import TensorType
-from theano.gof import ParamsType, Params
+from theano.gof import ParamsType, Params, EnumList
 from theano import tensor
 from theano.tests import unittest_tools as utt
 
@@ -213,6 +213,35 @@ class TestParamsType(TestCase):
                     a3=2000.0 - 0.00000000000000001)
         assert w.values_eq_approx(o1, o3)
 
+    def test_params_type_with_enums(self):
+        # Test that we fail if we create a params type with common enum names inside different enum types.
+        try:
+            ParamsType(enum1=EnumList('A', 'B', 'C'), enum2=EnumList('A', 'B', 'F'))
+        except AttributeError:
+            pass
+        else:
+            raise Exception('ParamsType should fail with common enum names inside different enum types.')
+
+        # Test that we fail if we create a params type with common names in both aliases and constants.
+        try:
+            ParamsType(enum1=EnumList(('A', 'a'), ('B', 'b')), enum2=EnumList(('ONE', 'a'), ('TWO', 'two')))
+        except AttributeError:
+            ParamsType(enum1=EnumList(('A', 'a'), ('B', 'b')), enum2=EnumList(('ONE', 'one'), ('TWO', 'two')))
+        else:
+            raise Exception('ParamsType should fail when there are aliases with same names as some constants.')
+
+        # Test that we can access enum values through wrapper directly.
+        w = ParamsType(enum1=EnumList('A', ('B', 'beta'), 'C'), enum2=EnumList(('D', 'delta'), 'E', 'F'))
+        assert w.A == 0 and w.B == 1 and w.C == 2
+        assert w.D == 0 and w.E == 1 and w.F == 2
+        # Test constants access through aliases.
+        assert w.enum_from_alias('beta') == w.B
+        assert w.enum_from_alias('delta') == w.D
+        assert w.enum_from_alias('C') == w.C  # C is not an alias, so it should return a constant named C.
+        # Test that other regular wrapper attributes are still available.
+        assert len(w.fields) == len(w.types) == w.length
+        assert w.name
+
     def test_op_params(self):
         a, b, c = 2, 3, -7
         x = tensor.matrix(dtype='float64')

theano/gof/tests/test_types.py

@@ -81,18 +81,19 @@ def test_cdata():
 
 class MyOpEnumList(Op):
     __props__ = ('op_chosen',)
-    params_type = EnumList('ADD', 'SUB', 'MULTIPLY', 'DIVIDE', ctype='unsigned long long')
+    params_type = EnumList(('ADD', '+'), ('SUB', '-'), ('MULTIPLY', '*'), ('DIVIDE', '/'), ctype='unsigned long long')
 
     def __init__(self, choose_op):
         assert self.params_type.ADD == 0
         assert self.params_type.SUB == 1
         assert self.params_type.MULTIPLY == 2
         assert self.params_type.DIVIDE == 3
-        op_to_const = {'+': self.params_type.ADD,
-                       '-': self.params_type.SUB,
-                       '*': self.params_type.MULTIPLY,
-                       '/': self.params_type.DIVIDE}
-        self.op_chosen = op_to_const[choose_op]
+        assert self.params_type.fromalias('+') == self.params_type.ADD
+        assert self.params_type.fromalias('-') == self.params_type.SUB
+        assert self.params_type.fromalias('*') == self.params_type.MULTIPLY
+        assert self.params_type.fromalias('/') == self.params_type.DIVIDE
+        assert self.params_type.has_alias(choose_op)
+        self.op_chosen = choose_op
 
     def get_params(self, node):
         return self.op_chosen
@@ -204,7 +205,7 @@ class TestEnumTypes(TestCase):
         # Check that invalid enum value raises exception.
         try:
             EnumType(INVALID_VALUE='string is not allowed.')
-        except ValueError:
+        except TypeError:
             pass
         else:
             raise Exception('EnumType with invalid value should fail.')
@@ -218,6 +219,23 @@ class TestEnumTypes(TestCase):
         # Check access to attributes.
         assert len((e1.ctype, e1.C1, e1.C2, e1.C3, e1.C4, e1.C5, e1.C6)) == 7
 
+        # Check enum with aliases.
+        e1 = EnumType(A=('alpha', 0), B=('beta', 1), C=2)
+        e2 = EnumType(A=('alpha', 0), B=('beta', 1), C=2)
+        e3 = EnumType(A=('a', 0), B=('beta', 1), C=2)
+        assert e1 == e2
+        assert e1 != e3
+        assert e1.filter('beta') == e1.fromalias('beta') == e1.B == 1
+        assert e1.filter('C') == e1.fromalias('C') == e1.C == 2
+
+        # Check that invalid alias (same as a constant) raises exception.
+        try:
+            EnumList(('A', 'a'), ('B', 'B'))
+        except TypeError:
+            EnumList(('A', 'a'), ('B', 'b'))
+        else:
+            raise Exception('Enum with an alias name equal to a constant name should fail.')
+
     def test_op_with_enumlist(self):
         a = scalar.int32()
         b = scalar.int32()

theano/gpuarray/dnn.py

@@ -12,7 +12,7 @@ from theano import Op, Apply, tensor, config, Variable
 from theano.scalar import as_scalar, constant, Log, get_scalar_type
 from theano.tensor import as_tensor_variable
 from theano.gradient import DisconnectedType, grad_not_implemented
-from theano.gof import Optimizer, local_optimizer, COp
+from theano.gof import Optimizer, local_optimizer, COp, ParamsType, CEnumType
 from theano.gof.cmodule import GCC_compiler
 from theano.gof.type import CDataType, Generic
 from theano.compile import optdb
@@ -234,6 +234,11 @@ class DnnBase(COp):
             ptr = ctx.cudnn_handle.value
             res = handle_type.make_value(ptr)
             ctx.cudnn_handle_param = res
+        if isinstance(self.params_type, ParamsType):
+            if not self.params_type.has_type(handle_type):
+                raise TypeError('DnnBase: params_type must take into account the cuDNN handle type.')
+            handle_field = self.params_type.get_field(handle_type)
+            return self.params_type.get_params(self, **{handle_field: ctx.cudnn_handle_param})
         return ctx.cudnn_handle_param
 
     def __init__(self, files=None, c_func=None):
@@ -1504,6 +1509,18 @@ class GpuDnnSoftmaxBase(DnnBase):
     """
 
     __props__ = ('mode', 'algo')
+    # Neither inputs nor output types properties are used
+    # neither in dnn_base.c nor in dnn_softmax*.c,
+    # so we can disable input checking.
+    check_input = False
+    params_type = ParamsType(algo=CEnumType(('CUDNN_SOFTMAX_FAST', 'fast'),
+                                            ('CUDNN_SOFTMAX_LOG', 'log'),
+                                            ('CUDNN_SOFTMAX_ACCURATE', 'accurate'),
+                                            ctype='cudnnSoftmaxAlgorithm_t'),
+                             mode=CEnumType(('CUDNN_SOFTMAX_MODE_INSTANCE', 'instance'),
+                                            ('CUDNN_SOFTMAX_MODE_CHANNEL', 'channel'),
+                                            ctype='cudnnSoftmaxMode_t'),
+                             handle=handle_type)
 
     def __init__(self, algo, mode):
         DnnBase.__init__(self, [self.file], self.c_func)
@@ -1520,21 +1537,6 @@ class GpuDnnSoftmaxBase(DnnBase):
         else:
             return [shape[1]]
 
-    def get_op_params(self):
-        if self.mode == 'instance':
-            mode = "CUDNN_SOFTMAX_MODE_INSTANCE"
-        else:
-            mode = "CUDNN_SOFTMAX_MODE_CHANNEL"
-
-        if self.algo == 'fast':
-            algo = "CUDNN_SOFTMAX_FAST"
-        elif self.algo == 'log':
-            algo = "CUDNN_SOFTMAX_LOG"
-        else:
-            algo = "CUDNN_SOFTMAX_ACCURATE"
-
-        return [("SOFTMAX_MODE", mode), ("SOFTMAX_ALGO", algo)]
-
 
 class GpuDnnSoftmax(GpuDnnSoftmaxBase):
 

theano/gpuarray/dnn_softmax.c

@@ -35,7 +35,7 @@ if (APPLY_SPECIFIC(output) != NULL)
 
 int APPLY_SPECIFIC(softmax)(PyGpuArrayObject *x,
                             PyGpuArrayObject **out,
-                            cudnnHandle_t _handle) {
+                            PARAMS_TYPE* wrapper) {
   PyGpuContextObject *c = x->context;
   cudnnStatus_t err;
 
@@ -83,9 +83,9 @@ int APPLY_SPECIFIC(softmax)(PyGpuArrayObject *x,
     cuda_wait((*out)->ga.data, GPUARRAY_CUDA_WAIT_WRITE);
 
     err = cudnnSoftmaxForward(
-      _handle,
-      SOFTMAX_ALGO,
-      SOFTMAX_MODE,
+      wrapper->handle,
+      wrapper->algo,
+      wrapper->mode,
       alpha,
       APPLY_SPECIFIC(input),
       PyGpuArray_DEV_DATA(x),

theano/gpuarray/dnn_softmax_grad.c

@@ -46,7 +46,7 @@ if (APPLY_SPECIFIC(dx) != NULL)
 int APPLY_SPECIFIC(softmax_grad)(PyGpuArrayObject *dy,
                                  PyGpuArrayObject *sm,
                                  PyGpuArrayObject **dx,
-                                 cudnnHandle_t _handle) {
+                                 PARAMS_TYPE* wrapper) {
   PyGpuContextObject *c = dy->context;
   cudnnStatus_t err;
 
@@ -97,9 +97,9 @@ int APPLY_SPECIFIC(softmax_grad)(PyGpuArrayObject *dy,
     cuda_wait((*dx)->ga.data, GPUARRAY_CUDA_WAIT_WRITE);
 
     err = cudnnSoftmaxBackward(
-      _handle,
-      SOFTMAX_ALGO,
-      SOFTMAX_MODE,
+      wrapper->handle,
+      wrapper->algo,
+      wrapper->mode,
       alpha,
       APPLY_SPECIFIC(sm),
       PyGpuArray_DEV_DATA(sm),