def __init__(self,
device_id=None,
enable_winograd=True,
deterministic=True,
scratch_size=0):
drv.init()
self.device_id = device_id if device_id is not None else 0
# check compute capability
self.compute_capability = drv.Device(
self.device_id).compute_capability()
if self.compute_capability[0] < 3:
raise RuntimeError("Unsupported GPU")
# context
self.ctx = drv.Device(self.device_id).make_context()
# attributes
self.stream = None
self.warmup = False
self.scratch_size = scratch_size
self.scratch_offset = 0
self.sm_count = _get_sm_count()
# store GPU memory size in bytes
self.gpu_memory_size = drv.mem_get_info()[1]
# Fall back to CUDA C kernels on older (pre-Maxwell) GPU generations
if self.compute_capability[0] < 5:
# TODO: this is not fully supported in graph yet
self.use_cudac_kernels = True
else:
self.use_cudac_kernels = False
# TODO
# self.cublas_handle = cublas.cublasCreate()
self.pcg = rng_mrg()
self.enable_winograd = enable_winograd
self.deterministic = deterministic
self.cache_dir = get_cache_dir()
def pcg(self):
return rng_mrg()