提交 74147ec4 authored 作者: nouiz's avatar nouiz

Merge pull request #967 from mrocklin/gpu-tags

add gpu tags to gpu optdbs
...@@ -43,7 +43,7 @@ gpu_optimizer = EquilibriumDB() ...@@ -43,7 +43,7 @@ gpu_optimizer = EquilibriumDB()
gpu_cut_copies = EquilibriumDB() gpu_cut_copies = EquilibriumDB()
gpu_seqopt = SequenceDB() gpu_seqopt = SequenceDB()
gpu_seqopt.register('gpu_local_optimizations', gpu_optimizer, 1, gpu_seqopt.register('gpu_local_optimizations', gpu_optimizer, 1,
'fast_run', 'inplace') 'fast_run', 'inplace', 'gpu')
gpu_seqopt.register('gpu_cut_transfers', gpu_cut_copies, 2, gpu_seqopt.register('gpu_cut_transfers', gpu_cut_copies, 2,
'fast_run', 'gpu') 'fast_run', 'gpu')
# DO NOT PUT fast_run in gpu_opt! This will ALWAYS enable the GPU! # DO NOT PUT fast_run in gpu_opt! This will ALWAYS enable the GPU!
...@@ -63,7 +63,7 @@ optdb.register('gpu_after_fusion', ...@@ -63,7 +63,7 @@ optdb.register('gpu_after_fusion',
def register_opt(*tags, **kwargs): def register_opt(*tags, **kwargs):
def f(local_opt): def f(local_opt):
name = (kwargs and kwargs.pop('name')) or local_opt.__name__ name = (kwargs and kwargs.pop('name')) or local_opt.__name__
gpu_optimizer.register(name, local_opt, 'fast_run', *tags) gpu_optimizer.register(name, local_opt, 'fast_run', 'gpu', *tags)
return local_opt return local_opt
return f return f
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论