Port local_assert_no_cpu_op to the new back-end

theano/gpuarray/opt.py

@@ -2,10 +2,11 @@ from __future__ import absolute_import, print_function, division
 import copy
 import numpy
 import logging
+import pdb
 from six.moves import xrange
 
 import theano
-from theano import tensor, scalar, gof
+from theano import tensor, scalar, gof, config
 from theano.compile import optdb
 from theano.compile.ops import shape_i
 from theano.gof import (local_optimizer, EquilibriumDB, TopoOptimizer,
@@ -955,6 +956,28 @@ def local_gpu_elemwise_careduce(node):
                                 pre_scalar_op=scalar.basic.sqr)(inp)]
 
 
+@local_optimizer(None)
+def local_assert_no_cpu_op(node):
+    if (all([var.owner and isinstance(var.owner.op, HostFromGpu)
+             for var in node.inputs]) and
+        any([[c for c in var.clients if isinstance(c[0].op, GpuFromHost)]
+             for var in node.outputs])):
+
+            if config.assert_no_cpu_op == "warn":
+                _logger.warning(("CPU Op %s is detected in the computation "
+                                 "graph") % node)
+            elif config.assert_no_cpu_op == "raise":
+                raise AssertionError("The Op %s is on CPU." % node)
+            elif config.assert_no_cpu_op == "pdb":
+                pdb.set_trace()
+
+# Register the local_assert_no_cpu_op:
+assert_no_cpu_op = theano.tensor.opt.in2out(local_assert_no_cpu_op,
+                                            name='assert_no_cpu_op')
+# 49.2 is after device specialization & fusion optimizations for last transfers
+optdb.register('assert_no_cpu_op', assert_no_cpu_op, 49.2)
+
+
 def tensor_to_gpu(x, context_name):
     if isinstance(x.type, tensor.TensorType):
         y = GpuArrayType(broadcastable=x.type.broadcastable,

theano/gpuarray/tests/test_opt.py

 from __future__ import absolute_import, print_function, division
 import numpy
+from nose.tools import assert_raises
 
 import theano
 from theano import tensor
@@ -417,3 +418,33 @@ def test_local_lift_abstractconv_gpu_shape():
         theano.function([s, a, b], c, mode=mode_with_gpu)
     finally:
         theano.config.on_opt_error = prev
+
+
+def test_local_assert_no_cpu_op():
+    rng = numpy.random.RandomState(utt.fetch_seed())
+    m = rng.uniform(-1, 1, (10, 10)).astype("float32")
+    ms = gpuarray_shared_constructor(m, name="m_shared")
+    out = theano.tensor.tanh(ms).dot(ms.T)
+
+    mode_local_assert = mode_with_gpu.including("assert_no_cpu_op")
+    mode_local_assert = mode_local_assert.excluding("local_gpu_elemwise")
+
+    old = theano.config.assert_no_cpu_op
+    old2 = theano.config.on_opt_error
+    # If the flag is raise
+    try:
+        theano.config.assert_no_cpu_op = 'raise'
+        theano.config.on_opt_error = 'ignore'
+
+        assert_raises(AssertionError, theano.function,
+                      [], out, mode=mode_local_assert)
+    finally:
+        theano.config.assert_no_cpu_op = old
+        theano.config.on_opt_error = old2
+
+    # If the flag is ignore
+    try:
+        theano.config.assert_no_cpu_op = 'ignore'
+        theano.function([], out, mode=mode_local_assert)
+    finally:
+        theano.config.assert_no_cpu_op = old