在 cupy 中启用 cub 的正确方法是什么?

What is the proper way to enable cub in cupy?

我正在尝试找出在 cupy 中启用 cub 的正确方法,但到目前为止还没有成功。我查看了文档,但找不到任何东西。目前我启用 cub 是这样的:

import cupy.core._accelerator as _acc
_acc.set_routine_accelerators(['cub'])
_acc.set_reduction_accelerators(['cub'])

在执行上述代码之前,cub被禁用。我通过 运行:

确认
cupy.core.get_reduction_accelerators()
cupy.core.get_routine_accelerators()

其中 return 一个空列表 ([])。在 运行 第一个片段中的代码之后,上面的函数 return [1] (无论那是什么意思)。此外,我可以注意到 cupy.nansum.

等函数的巨大性能差异

如您所见,函数 cupy.set_routine_acceleratorscupy.set_reduction_accelerators 属于私有 API (cupy.core._accelerator),这意味着我不应该调用它们。

我正在使用 Python 3.7.6cupy 8.1.0

谢谢

记录在案的方法是通过 CUPY_ACCELERATORS 环境变量。导出 CUPY_ACCELERATORS=cub。 docs.cupy.dev/en/stable/reference/environment.html 如您所见,set_*_accelerators 主要是我们用于测试的私有 API。他们返回 1 的原因是我们使用了枚举,我同意这令人困惑……也许我们可以更改它:)。