我如何知道安装了 numba 或 tensorflow 的 python 代码中每个块的最大线程数?

How do I know the maximum number of threads per block in python code with either numba or tensorflow installed?

python 中是否有安装了 numba 或 tensorflow 的代码? 例如,如果我想知道 GPU 内存信息,我可以简单地使用:

from numba import cuda
gpus = cuda.gpus.lst
for gpu in gpus:
    with gpu:
        meminfo = cuda.current_context().get_memory_info()
        print("%s, free: %s bytes, total, %s bytes" % (gpu, meminfo[0], meminfo[1]))

在 numba。 但是我找不到任何代码可以为我提供每个块信息的最大线程数。 我希望代码检测每个块的最大线程数,并进一步计算每个方向的指定块数。

Is there any code in python with either numba or tensorflow installed?

我找不到。 numba 设备 class 似乎具有检索设备属性的功能:

In [9]: ddd=numba.cuda.get_current_device()

In [10]: print(ddd)
<CUDA device 0 'b'GeForce GTX 970''>

In [11]: print(ddd.attributes)
{}

但至少在我使用的 numba 版本 (0.31.0) 中,字典似乎没有被填充。否则,在这个阶段,numba 似乎不会公开任何用于检索设备或编译函数属性的常规驱动程序 API 功能。

from numba import cuda
gpu = cuda.get_current_device()
print("name = %s" % gpu.name)
print("maxThreadsPerBlock = %s" % str(gpu.MAX_THREADS_PER_BLOCK))
print("maxBlockDimX = %s" % str(gpu.MAX_BLOCK_DIM_X))
print("maxBlockDimY = %s" % str(gpu.MAX_BLOCK_DIM_Y))
print("maxBlockDimZ = %s" % str(gpu.MAX_BLOCK_DIM_Z))
print("maxGridDimX = %s" % str(gpu.MAX_GRID_DIM_X))
print("maxGridDimY = %s" % str(gpu.MAX_GRID_DIM_Y))
print("maxGridDimZ = %s" % str(gpu.MAX_GRID_DIM_Z))
print("maxSharedMemoryPerBlock = %s" % str(gpu.MAX_SHARED_MEMORY_PER_BLOCK))
print("asyncEngineCount = %s" % str(gpu.ASYNC_ENGINE_COUNT))
print("canMapHostMemory = %s" % str(gpu.CAN_MAP_HOST_MEMORY))
print("multiProcessorCount = %s" % str(gpu.MULTIPROCESSOR_COUNT))
print("warpSize = %s" % str(gpu.WARP_SIZE))
print("unifiedAddressing = %s" % str(gpu.UNIFIED_ADDRESSING))
print("pciBusID = %s" % str(gpu.PCI_BUS_ID))
print("pciDeviceID = %s" % str(gpu.PCI_DEVICE_ID))

这些似乎是当前支持的所有属性。我找到了列表 here,它与 CUDA 文档中的枚举值相匹配,因此扩展它非常简单。例如,我添加了 CU_DEVICE_ATTRIBUTE_TOTAL_CONSTANT_MEMORY = 9,现在可以正常工作了。

如果这个周末我有时间,我会尝试将它们四舍五入,更新文档,并提交 PR。