提交 b94282cb authored 作者: Dustin Webb's avatar Dustin Webb 提交者: Frederic

Added a new decorator and tagged a few opts for optimization when using the gpu…

Added a new decorator and tagged a few opts for optimization when using the gpu or optimizer=fast_compile.
上级 67c0dfea
...@@ -315,6 +315,12 @@ def register_canonicalize(lopt, *tags, **kwargs): ...@@ -315,6 +315,12 @@ def register_canonicalize(lopt, *tags, **kwargs):
return lopt return lopt
def register_canonicalize_tags(*tags, **kwargs):
def register(lopt):
return register_canonicalize(lopt, *tags, **kwargs)
return register
def register_stabilize(lopt, *tags, **kwargs): def register_stabilize(lopt, *tags, **kwargs):
name = (kwargs and kwargs.pop('name')) or lopt.__name__ name = (kwargs and kwargs.pop('name')) or lopt.__name__
compile.optdb['stabilize'].register(name, lopt, 'fast_run', *tags) compile.optdb['stabilize'].register(name, lopt, 'fast_run', *tags)
...@@ -1304,7 +1310,7 @@ def local_track_shape_i(node): ...@@ -1304,7 +1310,7 @@ def local_track_shape_i(node):
@register_specialize @register_specialize
@register_canonicalize @register_canonicalize_tags('gpu')
@gof.local_optimizer([Subtensor]) @gof.local_optimizer([Subtensor])
def local_subtensor_make_vector(node): def local_subtensor_make_vector(node):
# replace all subtensor(make_vector) like: # replace all subtensor(make_vector) like:
...@@ -1354,8 +1360,7 @@ def local_subtensor_make_vector(node): ...@@ -1354,8 +1360,7 @@ def local_subtensor_make_vector(node):
#TODO: the other optimization for and, or, xor, le and ge see ticket #496. #TODO: the other optimization for and, or, xor, le and ge see ticket #496.
@register_canonicalize_tags('fast_compile')
@register_canonicalize
@register_specialize @register_specialize
@gof.local_optimizer([T.Elemwise]) @gof.local_optimizer([T.Elemwise])
def local_useless_elemwise(node): def local_useless_elemwise(node):
...@@ -3508,7 +3513,7 @@ def local_reduce_join(node): ...@@ -3508,7 +3513,7 @@ def local_reduce_join(node):
#else the reduction do something about the dtype. #else the reduction do something about the dtype.
@register_canonicalize @register_canonicalize_tags('fast_compile')
@gof.local_optimizer(ALL_REDUCE) @gof.local_optimizer(ALL_REDUCE)
def local_cut_useless_reduce(node): def local_cut_useless_reduce(node):
"""Sum(a, axis=[]) -> a """ """Sum(a, axis=[]) -> a """
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论