提交 e486864f authored 作者: Frederic Bastien's avatar Frederic Bastien

Make opt GpuAlloc memset0 work in the new back-end

上级 cef892bb
...@@ -74,7 +74,8 @@ def register_opt(*tags, **kwargs): ...@@ -74,7 +74,8 @@ def register_opt(*tags, **kwargs):
return f return f
register_opt('fast_compile')(theano.tensor.opt.local_track_shape_i) register_opt('fast_compile')(theano.tensor.opt.local_track_shape_i)
register_opt(final_opt=True, name='gpua_constant_folding')(
tensor.opt.constant_folding)
gpu_optimizer.register('local_remove_all_assert', gpu_optimizer.register('local_remove_all_assert',
theano.tensor.opt.local_remove_all_assert, theano.tensor.opt.local_remove_all_assert,
'unsafe') 'unsafe')
......
...@@ -125,6 +125,14 @@ def test_local_gpualloc_memset_0(): ...@@ -125,6 +125,14 @@ def test_local_gpualloc_memset_0():
o = numpy.ones((1,), dtype='float32') o = numpy.ones((1,), dtype='float32')
ones = numpy.ones((2,), dtype='float32') ones = numpy.ones((2,), dtype='float32')
# Test with 0 from CPU op.
a = tensor.alloc(z, i)
f = theano.function([i], a, mode=mode_with_gpu)
topo = f.maker.fgraph.toposort()
assert len(topo) == 2
assert isinstance(topo[0].op, GpuAlloc) and topo[0].op.memset_0
assert (numpy.asarray(f(6)) == 0).all()
# Test with 0 # Test with 0
a = GpuAlloc(test_ctx_name)(z, i) a = GpuAlloc(test_ctx_name)(z, i)
f = theano.function([i], a, mode=mode_with_gpu) f = theano.function([i], a, mode=mode_with_gpu)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论