# mypy: allow-untyped-defsr"""This package adds support for CUDA tensor types.It implements the same function as CPU tensors, but they utilizeGPUs for computation.It is lazily initialized, so you can always import it, and use:func:`is_available()` to determine if your system supports CUDA.:ref:`cuda-semantics` has more details about working with CUDA."""importimportlibimportosimportthreadingimporttracebackimportwarningsfromfunctoolsimportlru_cachefromtypingimportAny,Callable,cast,Optional,Unionimporttorchimporttorch._Cfromtorchimportdeviceas_devicefromtorch._utilsimport_dummy_type,_LazySeedTracker,classpropertyfromtorch.typesimportDevicefrom.importgdsfrom._utilsimport_get_device_indexfrom.graphsimport(CUDAGraph,graph,graph_pool_handle,is_current_stream_capturing,make_graphed_callables,)from.streamsimportEvent,ExternalStream,Streamtry:fromtorch._Cimport_cudart# type: ignore[attr-defined]exceptImportError:_cudart=None_initialized=False_tls=threading.local()_initialization_lock=threading.Lock()_queued_calls:list[tuple[Callable[[],None],list[str]]]=[]# don't invoke these until initialization occurs_is_in_bad_fork=getattr(torch._C,"_cuda_isInBadFork",lambda:False)_device_t=Union[_device,str,int,None]_HAS_PYNVML=False_PYNVML_ERR=Nonetry:fromtorchimportversionas_versiontry:ifnot_version.hip:importpynvml# type: ignore[import]else:importctypesfrompathlibimportPath# In ROCm (at least up through 6.3.2) there're 2 copies of libamd_smi.so:# - One at lib/libamd_smi.so# - One at share/amd_smi/amdsmi/libamd_smi.so## The amdsmi python module hardcodes loading the second one in share-# https://github.com/ROCm/amdsmi/blob/1d305dc9708e87080f64f668402887794cd46584/py-interface/amdsmi_wrapper.py#L174## See also https://github.com/ROCm/amdsmi/issues/72.## This creates an ODR violation if the copy of libamd_smi.so from lib# is also loaded (via `ld` linking, `LD_LIBRARY_PATH` or `rpath`).## In order to avoid the violation we hook CDLL and try using the# already loaded version of amdsmi, or any version in the processes# rpath/LD_LIBRARY_PATH first, so that we only load a single copy# of the .so.class_amdsmi_cdll_hook:def__init__(self)->None:self.original_CDLL=ctypes.CDLL# type: ignore[misc,assignment]paths=["libamd_smi.so"]ifrocm_home:=os.getenv("ROCM_HOME",os.getenv("ROCM_PATH")):paths=[os.path.join(rocm_home,"lib/libamd_smi.so")]+pathsself.paths:list[str]=pathsdefhooked_CDLL(self,name:Union[str,Path,None],*args:Any,**kwargs:Any)->ctypes.CDLL:ifnameandPath(name).name=="libamd_smi.so":forpathinself.paths:try:returnself.original_CDLL(path,*args,**kwargs)exceptOSError:passreturnself.original_CDLL(name,*args,**kwargs)# type: ignore[arg-type]def__enter__(self)->None:ctypes.CDLL=self.hooked_CDLL# type: ignore[misc,assignment]def__exit__(self,type:Any,value:Any,traceback:Any)->None:ctypes.CDLL=self.original_CDLL# type: ignore[misc]with_amdsmi_cdll_hook():importamdsmi# type: ignore[import]_HAS_PYNVML=TrueexceptModuleNotFoundError:passfinally:del_versionexceptImportErroraserr:_PYNVML_ERR=err# sometimes a lib is installed but the import fails for some other reason, so we log the error for later_lazy_seed_tracker=_LazySeedTracker()# Define dummy _CudaDeviceProperties type if PyTorch was compiled without CUDAifhasattr(torch._C,"_CudaDeviceProperties"):_CudaDeviceProperties=torch._C._CudaDevicePropertieselse:_CudaDeviceProperties=_dummy_type("_CudaDeviceProperties")# type: ignore[assignment, misc]ifhasattr(torch._C,"_cuda_exchangeDevice"):_exchange_device=torch._C._cuda_exchangeDeviceelse:def_exchange_device(device:int)->int:ifdevice<0:return-1raiseRuntimeError("PyTorch was compiled without CUDA support")ifhasattr(torch._C,"_cuda_maybeExchangeDevice"):_maybe_exchange_device=torch._C._cuda_maybeExchangeDeviceelse:def_maybe_exchange_device(device:int)->int:ifdevice<0:return-1raiseRuntimeError("PyTorch was compiled without CUDA support")has_half:bool=Truehas_magma:bool=torch._C._has_magmadefault_generators:tuple[torch._C.Generator]=()# type: ignore[assignment]def_is_compiled()->bool:r"""Return true if compile with CUDA support."""returnhasattr(torch._C,"_cuda_getDeviceCount")def_nvml_based_avail()->bool:returnos.getenv("PYTORCH_NVML_BASED_CUDA_CHECK")=="1"
[docs]defis_available()->bool:r"""Return a bool indicating if CUDA is currently available."""ifnot_is_compiled():returnFalseif_nvml_based_avail():# The user has set an env variable to request this availability check that attempts to avoid fork poisoning by# using NVML at the cost of a weaker CUDA availability assessment. Note that if NVML discovery/initialization# fails, this assessment falls back to the default CUDA Runtime API assessment (`cudaGetDeviceCount`)returndevice_count()>0else:# The default availability inspection never throws and returns 0 if the driver is missing or can't# be initialized. This uses the CUDA Runtime API `cudaGetDeviceCount` which in turn initializes the CUDA Driver# API via `cuInit`returntorch._C._cuda_getDeviceCount()>0
defis_bf16_supported(including_emulation:bool=True):r"""Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16."""# Check for ROCm, if true return true, no ROCM_VERSION check required,# since it is supported on AMD GPU archs.iftorch.version.hip:returnTrue# If CUDA is not available, than it does not support bf16 eitherifnotis_available():returnFalsedevice=torch.cuda.current_device()# Check for CUDA version and device compute capability.# This is a fast way to check for it.cuda_version=torch.version.cudaif(cuda_versionisnotNoneandint(cuda_version.split(".")[0])>=11andtorch.cuda.get_device_properties(device).major>=8):returnTrueifnotincluding_emulation:returnFalse# Finally try to create a bfloat16 device.return_check_bf16_tensor_supported(device)@lru_cache(maxsize=16)def_check_bf16_tensor_supported(device:_device_t):try:torch.tensor([1.0],dtype=torch.bfloat16,device=device)returnTrueexceptException:returnFalse
[docs]defis_tf32_supported()->bool:r"""Return a bool indicating if the current CUDA/ROCm device supports dtype tf32."""# Check for ROCm. If true, return false, since PyTorch does not currently support# tf32 on ROCm.iftorch.version.hip:returnFalse# Otherwise, tf32 is supported on CUDA platforms that natively (i.e. no emulation)# support bfloat16.returnis_bf16_supported(including_emulation=False)
def_sleep(cycles):torch._C._cuda_sleep(cycles)def_extract_arch_version(arch_string:str):"""Extracts the architecture string from a CUDA version"""base=arch_string.split("_")[1]base=base.removesuffix("a")returnint(base)def_check_capability():incorrect_binary_warn=""" Found GPU%d%s which requires CUDA_VERSION >= %d to work properly, but your PyTorch was compiled with CUDA_VERSION %d. Please install the correct PyTorch binary using instructions from https://pytorch.org """# noqa: F841old_gpu_warn=""" Found GPU%d%s which is of cuda capability %d.%d. PyTorch no longer supports this GPU because it is too old. The minimum cuda capability supported by this library is %d.%d. """iftorch.version.cudaisnotNone:# on ROCm we don't want this checkCUDA_VERSION=torch._C._cuda_getCompiledVersion()# noqa: F841fordinrange(device_count()):capability=get_device_capability(d)major=capability[0]minor=capability[1]name=get_device_name(d)current_arch=major*10+minormin_arch=min((_extract_arch_version(arch)forarchintorch.cuda.get_arch_list()),default=35,)ifcurrent_arch<min_arch:warnings.warn(old_gpu_warn%(d,name,major,minor,min_arch//10,min_arch%10))def_check_cubins():incompatible_device_warn="""{} with CUDA capability sm_{} is not compatible with the current PyTorch installation.The current PyTorch install supports CUDA capabilities {}.If you want to use the {} GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/"""iftorch.version.cudaisNone:# on ROCm we don't want this checkreturnarch_list=get_arch_list()iflen(arch_list)==0:returnsupported_sm=[_extract_arch_version(arch)forarchinarch_listif"sm_"inarch]foridxinrange(device_count()):cap_major,cap_minor=get_device_capability(idx)# NVIDIA GPU compute architectures are backward compatible within major versionsupported=any(sm//10==cap_majorforsminsupported_sm)ifnotsupported:device_name=get_device_name(idx)capability=cap_major*10+cap_minorwarnings.warn(incompatible_device_warn.format(device_name,capability," ".join(arch_list),device_name))
[docs]defis_initialized():r"""Return whether PyTorch's CUDA state has been initialized."""return_initializedandnot_is_in_bad_fork()
def_lazy_call(callable,**kwargs):with_initialization_lock:ifis_initialized():callable()else:# TODO(torch_deploy): this accesses linecache, which attempts to read the# file system to get traceback info. Patch linecache or do something# else here if this ends up being important.global_lazy_seed_trackerifkwargs.get("seed_all",False):_lazy_seed_tracker.queue_seed_all(callable,traceback.format_stack())elifkwargs.get("seed",False):_lazy_seed_tracker.queue_seed(callable,traceback.format_stack())else:# Don't store the actual traceback to avoid memory cycle_queued_calls.append((callable,traceback.format_stack()))_lazy_call(_check_capability)_lazy_call(_check_cubins)classDeferredCudaCallError(Exception):passOutOfMemoryError=torch._C.OutOfMemoryError
[docs]definit():r"""Initialize PyTorch's CUDA state. You may need to call this explicitly if you are interacting with PyTorch via its C API, as Python bindings for CUDA functionality will not be available until this initialization takes place. Ordinary users should not need this, as all of PyTorch's CUDA methods automatically initialize CUDA state on-demand. Does nothing if the CUDA state is already initialized. """_lazy_init()
def_lazy_init():global_initialized,_queued_callsifis_initialized()orhasattr(_tls,"is_initializing"):returnwith_initialization_lock:# We be double-checked locking, boys! This is OK because# the above test was GIL protected anyway. The inner test# is for when a thread blocked on some other thread which was# doing the initialization; when they get the lock, they will# find there is nothing left to do.ifis_initialized():return# It is important to prevent other threads from entering _lazy_init# immediately, while we are still guaranteed to have the GIL, because some# of the C calls we make below will release the GILif_is_in_bad_fork():raiseRuntimeError("Cannot re-initialize CUDA in forked subprocess. To use CUDA with ""multiprocessing, you must use the 'spawn' start method")ifnothasattr(torch._C,"_cuda_getDeviceCount"):raiseAssertionError("Torch not compiled with CUDA enabled")if_cudartisNone:raiseAssertionError("libcudart functions unavailable. It looks like you have a broken build?")# This function throws if there's a driver initialization error, no GPUs# are found or any other error occursif"CUDA_MODULE_LOADING"notinos.environ:os.environ["CUDA_MODULE_LOADING"]="LAZY"torch._C._cuda_init()# Some of the queued calls may reentrantly call _lazy_init();# we need to just return without initializing in that case.# However, we must not let any *other* threads in!_tls.is_initializing=True_queued_calls.extend(callsforcallsin_lazy_seed_tracker.get_calls()ifcalls)try:forqueued_call,orig_tracebackin_queued_calls:try:queued_call()exceptExceptionase:msg=(f"CUDA call failed lazily at initialization with error: {str(e)}\n\n"f"CUDA call was originally invoked at:\n\n{''.join(orig_traceback)}")raiseDeferredCudaCallError(msg)fromefinally:delattr(_tls,"is_initializing")_initialized=True
[docs]defcudart():r"""Retrieves the CUDA runtime API module. This function initializes the CUDA runtime environment if it is not already initialized and returns the CUDA runtime API module (_cudart). The CUDA runtime API module provides access to various CUDA runtime functions. Args: ``None`` Returns: module: The CUDA runtime API module (_cudart). Raises: RuntimeError: If CUDA cannot be re-initialized in a forked subprocess. AssertionError: If PyTorch is not compiled with CUDA support or if libcudart functions are unavailable. Example of CUDA operations with profiling: >>> import torch >>> from torch.cuda import cudart, check_error >>> import os >>> >>> os.environ['CUDA_PROFILE'] = '1' >>> >>> def perform_cuda_operations_with_streams(): >>> stream = torch.cuda.Stream() >>> with torch.cuda.stream(stream): >>> x = torch.randn(100, 100, device='cuda') >>> y = torch.randn(100, 100, device='cuda') >>> z = torch.mul(x, y) >>> return z >>> >>> torch.cuda.synchronize() >>> print("====== Start nsys profiling ======") >>> check_error(cudart().cudaProfilerStart()) >>> with torch.autograd.profiler.emit_nvtx(): >>> result = perform_cuda_operations_with_streams() >>> print("CUDA operations completed.") >>> check_error(torch.cuda.cudart().cudaProfilerStop()) >>> print("====== End nsys profiling ======") To run this example and save the profiling information, execute: >>> $ nvprof --profile-from-start off --csv --print-summary -o trace_name.prof -f -- python cudart_test.py This command profiles the CUDA operations in the provided script and saves the profiling information to a file named `trace_name.prof`. The `--profile-from-start off` option ensures that profiling starts only after the `cudaProfilerStart` call in the script. The `--csv` and `--print-summary` options format the profiling output as a CSV file and print a summary, respectively. The `-o` option specifies the output file name, and the `-f` option forces the overwrite of the output file if it already exists. """_lazy_init()return_cudart
[docs]classdevice:r"""Context-manager that changes the selected device. Args: device (torch.device or int): device index to select. It's a no-op if this argument is a negative integer or ``None``. """def__init__(self,device:Any):self.idx=_get_device_index(device,optional=True)self.prev_idx=-1def__enter__(self):self.prev_idx=torch.cuda._exchange_device(self.idx)def__exit__(self,type:Any,value:Any,traceback:Any):self.idx=torch.cuda._maybe_exchange_device(self.prev_idx)returnFalse
[docs]classdevice_of(device):r"""Context-manager that changes the current device to that of given object. You can use both tensors and storages as arguments. If a given object is not allocated on a GPU, this is a no-op. Args: obj (Tensor or Storage): object allocated on the selected device. """def__init__(self,obj):idx=obj.get_device()ifobj.is_cudaelse-1super().__init__(idx)
[docs]defset_device(device:_device_t)->None:r"""Set the current device. Usage of this function is discouraged in favor of :any:`device`. In most cases it's better to use ``CUDA_VISIBLE_DEVICES`` environmental variable. Args: device (torch.device or int): selected device. This function is a no-op if this argument is negative. """device=_get_device_index(device)ifdevice>=0:torch._C._cuda_setDevice(device)
[docs]defget_device_name(device:Optional[_device_t]=None)->str:r"""Get the name of a device. Args: device (torch.device or int or str, optional): device for which to return the name. This function is a no-op if this argument is a negative integer. It uses the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Returns: str: the name of the device """returnget_device_properties(device).name
[docs]defget_device_capability(device:Optional[_device_t]=None)->tuple[int,int]:r"""Get the cuda capability of a device. Args: device (torch.device or int or str, optional): device for which to return the device capability. This function is a no-op if this argument is a negative integer. It uses the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Returns: tuple(int, int): the major and minor cuda capability of the device """prop=get_device_properties(device)returnprop.major,prop.minor
[docs]defget_device_properties(device:Optional[_device_t]=None)->_CudaDeviceProperties:r"""Get the properties of a device. Args: device (torch.device or int or str, optional): device for which to return the properties of the device. It uses the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Returns: _CudaDeviceProperties: the properties of the device """_lazy_init()# will define _get_device_propertiesdevice=_get_device_index(device,optional=True)ifdevice<0ordevice>=device_count():raiseAssertionError("Invalid device id")return_get_device_properties(device)# type: ignore[name-defined]
[docs]defcan_device_access_peer(device:_device_t,peer_device:_device_t)->bool:r"""Check if peer access between two devices is possible."""_lazy_init()device=_get_device_index(device,optional=True)peer_device=_get_device_index(peer_device)ifdevice<0ordevice>=device_count():raiseAssertionError("Invalid device id")ifpeer_device<0orpeer_device>=device_count():raiseAssertionError("Invalid peer device id")returntorch._C._cuda_canDeviceAccessPeer(device,peer_device)
[docs]classStreamContext:r"""Context-manager that selects a given stream. All CUDA kernels queued within its context will be enqueued on a selected stream. Args: Stream (Stream): selected stream. This manager is a no-op if it's ``None``. .. note:: Streams are per-device. """cur_stream:Optional["torch.cuda.Stream"]def__init__(self,stream:Optional["torch.cuda.Stream"]):self.stream=streamself.idx=_get_device_index(None,True)ifnottorch.jit.is_scripting():ifself.idxisNone:self.idx=-1self.src_prev_stream=(Noneifnottorch.jit.is_scripting()elsetorch.cuda.default_stream(None))self.dst_prev_stream=(Noneifnottorch.jit.is_scripting()elsetorch.cuda.default_stream(None))def__enter__(self):# Local cur_stream variable for type refinementcur_stream=self.stream# Return if stream is None or CUDA device not availableifcur_streamisNoneorself.idx==-1:returnself.src_prev_stream=torch.cuda.current_stream(None)# If the stream is not on the current device, then# set the current stream on the deviceifself.src_prev_stream.device!=cur_stream.device:withdevice(cur_stream.device):self.dst_prev_stream=torch.cuda.current_stream(cur_stream.device)torch.cuda.set_stream(cur_stream)def__exit__(self,type:Any,value:Any,traceback:Any):# Local cur_stream variable for type refinementcur_stream=self.stream# If stream is None or no CUDA device available, returnifcur_streamisNoneorself.idx==-1:return# Reset the stream on the original device# and destination deviceifself.src_prev_stream.device!=cur_stream.device:# type: ignore[union-attr]torch.cuda.set_stream(self.dst_prev_stream)# type: ignore[arg-type]torch.cuda.set_stream(self.src_prev_stream)# type: ignore[arg-type]
[docs]defstream(stream:Optional["torch.cuda.Stream"])->StreamContext:r"""Wrap around the Context-manager StreamContext that selects a given stream. Arguments: stream (Stream): selected stream. This manager is a no-op if it's ``None``. .. note:: In eager mode stream is of type Stream class while in JIT it is an object of the custom class ``torch.classes.cuda.Stream``. """returnStreamContext(stream)
def_set_stream_by_id(stream_id,device_index,device_type):r"""set stream specified by the stream id, device index and device type Args: stream_id (int): stream id in stream pool device_index (int): device index in topo device_type (int): enum device type """torch._C._cuda_setStream(stream_id=stream_id,device_index=device_index,device_type=device_type,)
[docs]defset_stream(stream:Stream):r"""Set the current stream.This is a wrapper API to set the stream. Usage of this function is discouraged in favor of the ``stream`` context manager. Args: stream (Stream): selected stream. This function is a no-op if this argument is ``None``. """ifstreamisNone:return_set_stream_by_id(stream_id=stream.stream_id,device_index=stream.device_index,device_type=stream.device_type,)
def_parse_visible_devices()->Union[list[int],list[str]]:r"""Parse CUDA_VISIBLE_DEVICES environment variable."""var=os.getenv("CUDA_VISIBLE_DEVICES")iftorch.version.hip:hip_devices=os.getenv("HIP_VISIBLE_DEVICES")rocr_devices=os.getenv("ROCR_VISIBLE_DEVICES")# You must take care if both HIP and ROCR env vars are set as they have# different meanings. Both env vars accept either a list of ints or a# list of UUIDs. The ROCR env var is processed first which then reduces# the number of GPUs that HIP can select from.ifrocr_devicesisnotNone:rocr_count=len(rocr_devices.split(","))ifhip_devicesisnotNone:# sanity check if both env vars are setiflen(hip_devices.split(","))>rocr_count:raiseRuntimeError("HIP_VISIBLE_DEVICES contains more devices than ROCR_VISIBLE_DEVICES")# HIP_VISIBLE_DEVICES is preferred over ROCR_VISIBLE_DEVICESvar=hip_deviceselse:returnlist(range(rocr_count))elifhip_devicesisnotNone:var=hip_devicesifvarisNone:returnlist(range(64))def_strtoul(s:str)->int:"""Return -1 or positive integer sequence string starts with."""ifnots:return-1foridx,cinenumerate(s):ifnot(c.isdigit()or(idx==0andcin"+-")):breakifidx+1==len(s):idx+=1returnint(s[:idx])ifidx>0else-1defparse_list_with_prefix(lst:str,prefix:str)->list[str]:rcs:list[str]=[]foreleminlst.split(","):# Repeated id results in empty setifeleminrcs:returncast(list[str],[])# Anything other but prefix is ignoredifnotelem.startswith(prefix):breakrcs.append(elem)returnrcsifvar.startswith("GPU-"):returnparse_list_with_prefix(var,"GPU-")ifvar.startswith("MIG-"):returnparse_list_with_prefix(var,"MIG-")# CUDA_VISIBLE_DEVICES uses something like strtoul# which makes `1gpu2,2ampere` is equivalent to `1,2`rc:list[int]=[]foreleminvar.split(","):x=_strtoul(elem.strip())# Repeated ordinal results in empty setifxinrc:returncast(list[int],[])# Negative value aborts the sequenceifx<0:breakrc.append(x)returnrcdef_raw_device_count_amdsmi()->int:ifnot_HAS_PYNVML:# If amdsmi is not availablereturn-1try:amdsmi.amdsmi_init()exceptamdsmi.AmdSmiExceptionase:warnings.warn(f"Can't initialize amdsmi - Error code: {e.err_code}")return-1socket_handles=amdsmi.amdsmi_get_processor_handles()returnlen(socket_handles)def_raw_device_count_nvml()->int:r"""Return number of devices as reported by NVML or negative value if NVML discovery/initialization failed."""fromctypesimportbyref,c_int,CDLLnvml_h=CDLL("libnvidia-ml.so.1")rc=nvml_h.nvmlInit()ifrc!=0:warnings.warn("Can't initialize NVML")return-1dev_count=c_int(-1)rc=nvml_h.nvmlDeviceGetCount_v2(byref(dev_count))ifrc!=0:warnings.warn("Can't get nvml device count")return-1delnvml_hreturndev_count.valuedef_raw_device_uuid_amdsmi()->Optional[list[str]]:fromctypesimportbyref,c_int,c_void_p,CDLL,create_string_bufferifnot_HAS_PYNVML:# If amdsmi is not availablereturnNonetry:amdsmi.amdsmi_init()exceptamdsmi.AmdSmiException:warnings.warn("Can't initialize amdsmi")returnNonetry:socket_handles=amdsmi.amdsmi_get_processor_handles()dev_count=len(socket_handles)exceptamdsmi.AmdSmiException:warnings.warn("Can't get amdsmi device count")returnNoneuuids:list[str]=[]foridxinrange(dev_count):try:handler=amdsmi.amdsmi_get_processor_handles()[idx]exceptamdsmi.AmdSmiException:warnings.warn("Cannot get amd device handler")returnNonetry:uuid=amdsmi.amdsmi_get_gpu_asic_info(handler)["asic_serial"][2:]# Removes 0x prefix from serialexceptamdsmi.AmdSmiException:warnings.warn("Cannot get uuid for amd device")returnNoneuuids.append(str(uuid).lower())# Lower-case to match expected HIP_VISIBLE_DEVICES uuid inputreturnuuidsdef_raw_device_uuid_nvml()->Optional[list[str]]:r"""Return list of device UUID as reported by NVML or None if NVM discovery/initialization failed."""fromctypesimportbyref,c_int,c_void_p,CDLL,create_string_buffernvml_h=CDLL("libnvidia-ml.so.1")rc=nvml_h.nvmlInit()ifrc!=0:warnings.warn("Can't initialize NVML")returnNonedev_count=c_int(-1)rc=nvml_h.nvmlDeviceGetCount_v2(byref(dev_count))ifrc!=0:warnings.warn("Can't get nvml device count")returnNoneuuids:list[str]=[]foridxinrange(dev_count.value):dev_id=c_void_p()rc=nvml_h.nvmlDeviceGetHandleByIndex_v2(idx,byref(dev_id))ifrc!=0:warnings.warn("Can't get device handle")returnNonebuf_len=96buf=create_string_buffer(buf_len)rc=nvml_h.nvmlDeviceGetUUID(dev_id,buf,buf_len)ifrc!=0:warnings.warn("Can't get device UUID")returnNoneuuids.append(buf.raw.decode("ascii").strip("\0"))delnvml_hreturnuuidsdef_transform_uuid_to_ordinals(candidates:list[str],uuids:list[str])->list[int]:r"""Given the set of partial uuids and list of known uuids builds a set of ordinals excluding ambiguous partials IDs."""defuuid_to_ordinal(candidate:str,uuids:list[str])->int:best_match=-1foridx,uuidinenumerate(uuids):ifnotuuid.startswith(candidate):continue# Ambiguous candidateifbest_match!=-1:return-1best_match=idxreturnbest_matchrc:list[int]=[]forcandidateincandidates:iftorch.version.hip:candidate=candidate.replace("GPU-","",1)# Remove GPU-prefix to match amdsmi asic serialidx=uuid_to_ordinal(candidate,uuids)# First invalid ordinal stops parsingifidx<0:break# Duplicates result in empty setifidxinrc:returncast(list[int],[])rc.append(idx)returnrcdef_device_count_amdsmi()->int:visible_devices=_parse_visible_devices()ifnotvisible_devices:return0try:iftype(visible_devices[0])isstr:uuids=_raw_device_uuid_amdsmi()ifuuidsisNone:return-1# Create string version of visible devices to avoid mypy warningsvisible_device_str=cast(list[str],visible_devices)visible_devices=_transform_uuid_to_ordinals(visible_device_str,uuids)else:raw_cnt=_raw_device_count_amdsmi()ifraw_cnt<=0:returnraw_cnt# Trim the list up to a maximum available deviceforidx,valinenumerate(visible_devices):ifcast(int,val)>=raw_cnt:returnidxexceptOSError:return-1exceptAttributeError:return-1returnlen(visible_devices)def_device_count_nvml()->int:r"""Return number of devices as reported by NVML taking CUDA_VISIBLE_DEVICES into account. Negative value is returned if NVML discovery or initialization has failed. """visible_devices=_parse_visible_devices()ifnotvisible_devices:return0try:iftype(visible_devices[0])isstr:# Skip MIG parsingifvisible_devices[0].startswith("MIG-"):return-1uuids=_raw_device_uuid_nvml()ifuuidsisNone:return-1visible_devices=_transform_uuid_to_ordinals(cast(list[str],visible_devices),uuids)else:raw_cnt=_raw_device_count_nvml()ifraw_cnt<=0:returnraw_cnt# Trim the list up to a maximum available deviceforidx,valinenumerate(visible_devices):ifcast(int,val)>=raw_cnt:returnidxexceptOSError:return-1exceptAttributeError:return-1returnlen(visible_devices)def_get_nvml_device_index(device:Optional[Union[int,Device]])->int:r"""Return the NVML index of the device, taking CUDA_VISIBLE_DEVICES into account."""idx=_get_device_index(device,optional=True)visible_devices=_parse_visible_devices()iftype(visible_devices[0])isstr:uuids=_raw_device_uuid_nvml()ifuuidsisNone:raiseRuntimeError("Can't get device UUIDs")visible_devices=_transform_uuid_to_ordinals(cast(list[str],visible_devices),uuids)visible_devices=cast(list[int],visible_devices)ifidx<0oridx>=len(visible_devices):raiseRuntimeError(f"device {idx} is not visible (CUDA_VISIBLE_DEVICES={visible_devices})")returnvisible_devices[idx]_cached_device_count:Optional[int]=None
[docs]defdevice_count()->int:r"""Return the number of GPUs available."""global_cached_device_countifnot_is_compiled():return0if_cached_device_countisnotNone:return_cached_device_count# bypass _device_count_nvml() if rocm (not supported)nvml_count=_device_count_amdsmi()iftorch.version.hipelse_device_count_nvml()r=torch._C._cuda_getDeviceCount()ifnvml_count<0elsenvml_count# NB: Do not cache the device count prior to CUDA initialization, because# the number of devices can change due to changes to CUDA_VISIBLE_DEVICES# setting prior to CUDA initialization.if_initialized:_cached_device_count=rreturnr
[docs]defget_arch_list()->list[str]:r"""Return list CUDA architectures this library was compiled for."""ifnotis_available():return[]arch_flags=torch._C._cuda_getArchFlags()ifarch_flagsisNone:return[]returnarch_flags.split()
[docs]defget_gencode_flags()->str:r"""Return NVCC gencode flags this library was compiled with."""arch_list=get_arch_list()iflen(arch_list)==0:return""arch_list_=[arch.split("_")forarchinarch_list]return" ".join([f"-gencode compute=compute_{arch},code={kind}_{arch}"for(kind,arch)inarch_list_])
[docs]defcurrent_device()->int:r"""Return the index of a currently selected device."""_lazy_init()returntorch._C._cuda_getDevice()
[docs]defsynchronize(device:_device_t=None)->None:r"""Wait for all kernels in all streams on a CUDA device to complete. Args: device (torch.device or int, optional): device for which to synchronize. It uses the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). """_lazy_init()withtorch.cuda.device(device):returntorch._C._cuda_synchronize()
[docs]defipc_collect():r"""Force collects GPU memory after it has been released by CUDA IPC. .. note:: Checks if any sent CUDA tensors could be cleaned from the memory. Force closes shared memory file used for reference counting if there is no active counters. Useful when the producer process stopped actively sending tensors and want to release unused memory. """_lazy_init()returntorch._C._cuda_ipc_collect()
[docs]defcurrent_stream(device:Optional[_device_t]=None)->Stream:r"""Return the currently selected :class:`Stream` for a given device. Args: device (torch.device or int, optional): selected device. Returns the currently selected :class:`Stream` for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). """_lazy_init()streamdata=torch._C._cuda_getCurrentStream(_get_device_index(device,optional=True))returnStream(stream_id=streamdata[0],device_index=streamdata[1],device_type=streamdata[2])
[docs]defdefault_stream(device:Optional[_device_t]=None)->Stream:r"""Return the default :class:`Stream` for a given device. Args: device (torch.device or int, optional): selected device. Returns the default :class:`Stream` for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). """_lazy_init()streamdata=torch._C._cuda_getDefaultStream(_get_device_index(device,optional=True))returnStream(stream_id=streamdata[0],device_index=streamdata[1],device_type=streamdata[2])
[docs]defget_stream_from_external(data_ptr:int,device:Optional[_device_t]=None)->Stream:r"""Return a :class:`Stream` from an externally allocated CUDA stream. This function is used to wrap streams allocated in other libraries in order to facilitate data exchange and multi-library interactions. .. note:: This function doesn't manage the stream life-cycle, it is the user responsibility to keep the referenced stream alive while this returned stream is being used. Args: data_ptr(int): Integer representation of the `cudaStream_t` value that is allocated externally. device(torch.device or int, optional): the device where the stream was originally allocated. If device is specified incorrectly, subsequent launches using this stream may fail. """_lazy_init()streamdata=torch._C._cuda_getStreamFromExternal(data_ptr,_get_device_index(device,optional=True))returnStream(stream_id=streamdata[0],device_index=streamdata[1],device_type=streamdata[2])
[docs]defcurrent_blas_handle():r"""Return cublasHandle_t pointer to current cuBLAS handle"""_lazy_init()returntorch._C._cuda_getCurrentBlasHandle()
[docs]defset_sync_debug_mode(debug_mode:Union[int,str])->None:r"""Set the debug mode for cuda synchronizing operations. Args: debug_mode(str or int): if "default" or 0, don't error or warn on synchronizing operations, if "warn" or 1, warn on synchronizing operations, if "error" or 2, error out synchronizing operations. Warning: This is an experimental feature, and not all synchronizing operations will trigger warning or error. In particular, operations in torch.distributed and torch.sparse namespaces are not covered yet. """_lazy_init()ifisinstance(debug_mode,str):ifdebug_mode=="default":debug_mode=0elifdebug_mode=="warn":debug_mode=1elifdebug_mode=="error":debug_mode=2else:raiseRuntimeError("invalid value of debug_mode, expected one of `default`, `warn`, `error`")torch._C._cuda_set_sync_debug_mode(debug_mode)
[docs]defget_sync_debug_mode()->int:r"""Return current value of debug mode for cuda synchronizing operations."""_lazy_init()returntorch._C._cuda_get_sync_debug_mode()
def_get_pynvml_handler(device:Optional[Union[Device,int]]=None):ifnot_HAS_PYNVML:raiseModuleNotFoundError("pynvml does not seem to be installed or it can't be imported.")from_PYNVML_ERRfrompynvmlimportNVMLError_DriverNotLoadedtry:pynvml.nvmlInit()exceptNVMLError_DriverNotLoadedase:raiseRuntimeError("cuda driver can't be loaded, is cuda enabled?")fromedevice=_get_nvml_device_index(device)handle=pynvml.nvmlDeviceGetHandleByIndex(device)returnhandledef_get_amdsmi_handler(device:Optional[Union[Device,int]]=None):ifnot_HAS_PYNVML:raiseModuleNotFoundError("amdsmi does not seem to be installed or it can't be imported.")from_PYNVML_ERRtry:amdsmi.amdsmi_init()exceptamdsmi.AmdSmiExceptionase:raiseRuntimeError("amdsmi driver can't be loaded, requires >=ROCm5.6 installation")fromedevice=_get_amdsmi_device_index(device)handle=amdsmi.amdsmi_get_processor_handles()[device]returnhandledef_get_amdsmi_device_index(device:Optional[Union[int,Device]])->int:r"""Return the amdsmi index of the device, taking visible_devices into account."""idx=_get_device_index(device,optional=True)visible_devices=_parse_visible_devices()iftype(visible_devices[0])isstr:uuids=_raw_device_uuid_amdsmi()ifuuidsisNone:raiseRuntimeError("Can't get device UUIDs")visible_devices_str=cast(list[str],visible_devices)# Create str variable for mypyvisible_devices=_transform_uuid_to_ordinals(visible_devices_str,uuids)idx_map=dict(enumerate(cast(list[int],visible_devices)))ifidxnotinidx_map:raiseRuntimeError(f"device {idx} is not visible (HIP_VISIBLE_DEVICES={visible_devices})")returnidx_map[idx]def_get_amdsmi_device_memory_used(device:Optional[Union[Device,int]]=None)->int:handle=_get_amdsmi_handler()device=_get_amdsmi_device_index(device)# amdsmi_get_gpu_vram_usage returns mem usage in megabytesmem_mega_bytes=amdsmi.amdsmi_get_gpu_vram_usage(handle)["vram_used"]mem_bytes=mem_mega_bytes*1024*1024returnmem_bytesdef_get_amdsmi_memory_usage(device:Optional[Union[Device,int]]=None)->int:handle=_get_amdsmi_handler()device=_get_amdsmi_device_index(device)handle=amdsmi.amdsmi_get_processor_handles()[device]returnamdsmi.amdsmi_get_gpu_activity(handle)["umc_activity"]def_get_amdsmi_utilization(device:Optional[Union[Device,int]]=None)->int:handle=_get_amdsmi_handler()device=_get_amdsmi_device_index(device)handle=amdsmi.amdsmi_get_processor_handles()[device]returnamdsmi.amdsmi_get_gpu_activity(handle)["gfx_activity"]def_get_amdsmi_temperature(device:Optional[Union[Device,int]]=None)->int:handle=_get_amdsmi_handler(device)returnamdsmi.amdsmi_get_temp_metric(handle,amdsmi.AmdSmiTemperatureType.JUNCTION,amdsmi.AmdSmiTemperatureMetric.CURRENT,)def_get_amdsmi_power_draw(device:Optional[Union[Device,int]]=None)->int:handle=_get_amdsmi_handler(device)socket_power=amdsmi.amdsmi_get_power_info(handle)["average_socket_power"]ifsocket_power!="N/A":returnsocket_powerelse:returnamdsmi.amdsmi_get_power_info(handle)["current_socket_power"]def_get_amdsmi_clock_rate(device:Optional[Union[Device,int]]=None)->int:handle=_get_amdsmi_handler(device)clock_info=amdsmi.amdsmi_get_clock_info(handle,amdsmi.AmdSmiClkType.GFX)if"cur_clk"inclock_info:# ROCm 6.2 deprecationreturnclock_info["cur_clk"]else:returnclock_info["clk"]
[docs]defdevice_memory_used(device:Optional[Union[Device,int]]=None)->int:r"""Return used global (device) memory in bytes as given by `nvidia-smi` or `amd-smi`. Args: device (torch.device or int, optional): selected device. Returns statistic for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). """ifnottorch.version.hip:handle=_get_pynvml_handler()device=_get_nvml_device_index(device)handle=pynvml.nvmlDeviceGetHandleByIndex(device)returnpynvml.nvmlDeviceGetMemoryInfo(handle).usedelse:return_get_amdsmi_device_memory_used(device)
[docs]defmemory_usage(device:Optional[Union[Device,int]]=None)->int:r"""Return the percent of time over the past sample period during which global (device) memory was being read or written as given by `nvidia-smi`. Args: device (torch.device or int, optional): selected device. Returns statistic for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Warning: Each sample period may be between 1 second and 1/6 second, depending on the product being queried. """ifnottorch.version.hip:handle=_get_pynvml_handler()device=_get_nvml_device_index(device)handle=pynvml.nvmlDeviceGetHandleByIndex(device)returnpynvml.nvmlDeviceGetUtilizationRates(handle).memoryelse:return_get_amdsmi_memory_usage(device)
[docs]defutilization(device:Optional[Union[Device,int]]=None)->int:r"""Return the percent of time over the past sample period during which one or more kernels was executing on the GPU as given by `nvidia-smi`. Args: device (torch.device or int, optional): selected device. Returns statistic for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Warning: Each sample period may be between 1 second and 1/6 second, depending on the product being queried. """ifnottorch.version.hip:handle=_get_pynvml_handler(device)device=_get_nvml_device_index(device)handle=pynvml.nvmlDeviceGetHandleByIndex(device)returnpynvml.nvmlDeviceGetUtilizationRates(handle).gpuelse:return_get_amdsmi_utilization(device)
[docs]deftemperature(device:Optional[Union[Device,int]]=None)->int:r"""Return the average temperature of the GPU sensor in Degrees C (Centigrades). The average temperature is computed based on past sample period as given by `nvidia-smi`. Args: device (torch.device or int, optional): selected device. Returns statistic for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Warning: Each sample period may be between 1 second and 1/6 second, depending on the product being queried. """ifnottorch.version.hip:handle=_get_pynvml_handler(device)# 0 refers to the temperature sensor for the GPU die.returnpynvml.nvmlDeviceGetTemperature(handle,0)else:return_get_amdsmi_temperature(device)
[docs]defpower_draw(device:Optional[Union[Device,int]]=None)->int:r"""Return the average power draw of the GPU sensor in mW (MilliWatts) over the past sample period as given by `nvidia-smi` for Fermi or newer fully supported devices. Args: device (torch.device or int, optional): selected device. Returns statistic for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Warning: Each sample period may be between 1 second and 1/6 second, depending on the product being queried. """ifnottorch.version.hip:handle=_get_pynvml_handler(device)returnpynvml.nvmlDeviceGetPowerUsage(handle)else:return_get_amdsmi_power_draw(device)
[docs]defclock_rate(device:Optional[Union[Device,int]]=None)->int:r"""Return the clock speed of the GPU SM in MHz (megahertz) over the past sample period as given by `nvidia-smi`. Args: device (torch.device or int, optional): selected device. Returns statistic for the current device, given by :func:`~torch.cuda.current_device`, if :attr:`device` is ``None`` (default). Warning: Each sample period may be between 1 second and 1/6 second, depending on the product being queried. """ifnottorch.version.hip:handle=_get_pynvml_handler(device)returnpynvml.nvmlDeviceGetClockInfo(handle,1)else:return_get_amdsmi_clock_rate(device)
def_get_device(device:Union[int,str,torch.device])->torch.device:r"""Return the torch.device type object from the passed in device. Args: device (torch.device or int): selected device. """ifisinstance(device,str):device=torch.device(device)elifisinstance(device,int):device=torch.device("cuda",device)returndevicedef_get_generator(device:torch.device)->torch._C.Generator:r"""Return the CUDA Generator object for the given device. Args: device (torch.device): selected device. """idx=device.indexifidxisNone:idx=current_device()returntorch.cuda.default_generators[idx]def_set_rng_state_offset(offset:int,device:Union[int,str,torch.device]="cuda")->None:r"""Set the random number generator state offset of the specified GPU. Args: offset (int): The desired offset device (torch.device or int, optional): The device to set the RNG state. Default: ``'cuda'`` (i.e., ``torch.device('cuda')``, the current CUDA device). """final_device=_get_device(device)defcb():default_generator=_get_generator(final_device)default_generator.set_offset(offset)_lazy_call(cb)def_get_rng_state_offset(device:Union[int,str,torch.device]="cuda")->int:r"""Return the random number generator state offset of the specified GPU. Args: device (torch.device or int, optional): The device to return the RNG state offset of. Default: ``'cuda'`` (i.e., ``torch.device('cuda')``, the current CUDA device). .. warning:: This function eagerly initializes CUDA. """_lazy_init()final_device=_get_device(device)default_generator=_get_generator(final_device)returndefault_generator.get_offset()from.memoryimport*# noqa: F403from.randomimport*# noqa: F403################################################################################# Define Storage and Tensor classes################################################################################@staticmethod# type: ignore[misc]def_lazy_new(cls,*args,**kwargs):_lazy_init()# We may need to call lazy init again if we are a forked child# del _CudaBase.__new__returnsuper(_CudaBase,cls).__new__(cls,*args,**kwargs)class_CudaBase:is_cuda=Trueis_sparse=Falsedeftype(self,*args,**kwargs):# We could use a Protocol here to tell mypy that self has `get_device` method# but it is only available in the typing module on Python >= 3.8# or on typing_extensions module on Python >= 3.6withdevice(self.get_device()):# type: ignore[attr-defined]returnsuper().type(*args,**kwargs)# type: ignore[misc]__new__=_lazy_newfromtorch.storageimport_LegacyStorage,_warn_typed_storage_removalclass_CudaLegacyStorage(_LegacyStorage):@classmethoddeffrom_buffer(cls,*args,**kwargs):_warn_typed_storage_removal()raiseRuntimeError("from_buffer: Not available for CUDA storage")@classmethoddef_new_with_weak_ptr(cls,*args,**kwargs):raiseRuntimeError("_new_with_weak_ptr: Not available for CUDA storage")@classmethoddef_new_shared_filename(cls,manager,obj,size,*,device=None,dtype=None):raiseRuntimeError("_new_shared_filename: Not available for CUDA storage")classByteStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.uint8classDoubleStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.doubleclassFloatStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.floatclassHalfStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.halfclassLongStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.longclassIntStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.intclassShortStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.shortclassCharStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.int8classBoolStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.boolclassBFloat16Storage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.bfloat16classComplexDoubleStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.cdoubleclassComplexFloatStorage(_CudaLegacyStorage):@classpropertydefdtype(self):_warn_typed_storage_removal()returnself._dtype@classpropertydef_dtype(self):returntorch.cfloatdel_LegacyStoragedel_CudaLegacyStoragetorch._storage_classes.add(DoubleStorage)torch._storage_classes.add(FloatStorage)torch._storage_classes.add(LongStorage)torch._storage_classes.add(IntStorage)torch._storage_classes.add(ShortStorage)torch._storage_classes.add(CharStorage)torch._storage_classes.add(ByteStorage)torch._storage_classes.add(HalfStorage)torch._storage_classes.add(BoolStorage)torch._storage_classes.add(BFloat16Storage)torch._storage_classes.add(ComplexDoubleStorage)torch._storage_classes.add(ComplexFloatStorage)class_WrappedTritonKernel:"""Just a simple wrapper to store some metadata for testing purposes."""def__init__(self,kernel):self.kernel=kernelself.kernel_invoked=Falsedef__call__(self,*args,**kwargs):res=self.kernel(*args,**kwargs)self.kernel_invoked=Truereturnresdef_register_triton_kernels():iftorch._running_with_deploy():return@_WrappedTritonKerneldefkernel_impl(*args,**kwargs):fromtorch.sparse._triton_opsimportbsr_dense_mmreturnbsr_dense_mm(*args,skip_checks=True,**kwargs)@_WrappedTritonKerneldefaddmm_kernel_impl(*args,**kwargs):fromtorch.sparse._triton_opsimportbsr_dense_addmmreturnbsr_dense_addmm(*args,skip_checks=True,**kwargs)has_triton=importlib.util.find_spec("triton")isnotNoneifhas_triton:torch._TritonLibrary.registerOp("_triton_bsr_dense_mm_out","_triton_bsr_dense_mm_out(Tensor bsr, Tensor dense, *, Tensor(a!) out) -> Tensor(a!)",kernel_impl,"SparseCsrCUDA",)torch._TritonLibrary.registerOp("_triton_bsr_dense_addmm_out",("_triton_bsr_dense_addmm_out(Tensor input, Tensor bsr, Tensor dense,"" *, Scalar beta, Scalar alpha, Tensor(a!) out) -> Tensor(a!)"),addmm_kernel_impl,"SparseCsrCUDA",)_lazy_call(_register_triton_kernels)from.importamp,jiterator,nvtx,profiler,sparse,tunable__all__=[# Typed storage and tensors"BFloat16Storage","BFloat16Tensor","BoolStorage","BoolTensor","ByteStorage","ByteTensor","CharStorage","CharTensor","ComplexDoubleStorage","ComplexFloatStorage","DoubleStorage","DoubleTensor","FloatStorage","FloatTensor","HalfStorage","HalfTensor","IntStorage","IntTensor","LongStorage","LongTensor","ShortStorage","ShortTensor","CUDAGraph","CudaError","DeferredCudaCallError","Event","ExternalStream","Stream","StreamContext","amp","caching_allocator_alloc","caching_allocator_delete","caching_allocator_enable","can_device_access_peer","check_error","cudaStatus","cudart","current_blas_handle","current_device","current_stream","default_generators","default_stream","device","device_count","device_memory_used","device_of","empty_cache","get_allocator_backend","CUDAPluggableAllocator","change_current_allocator","get_arch_list","get_device_capability","get_device_name","get_device_properties","get_gencode_flags","get_per_process_memory_fraction","get_rng_state","get_rng_state_all","get_stream_from_external","get_sync_debug_mode","graph","graph_pool_handle","graphs","has_half","has_magma","init","initial_seed","ipc_collect","is_available","is_bf16_supported","is_current_stream_capturing","is_initialized","is_tf32_supported","jiterator","list_gpu_processes","make_graphed_callables","manual_seed","manual_seed_all","max_memory_allocated","max_memory_cached","max_memory_reserved","mem_get_info","memory","memory_allocated","memory_cached","memory_reserved","memory_snapshot","memory_stats","memory_stats_as_nested_dict","memory_summary","memory_usage","MemPool","MemPoolContext","use_mem_pool","temperature","power_draw","clock_rate","nccl","nvtx","profiler","random","reset_accumulated_memory_stats","reset_max_memory_allocated","reset_max_memory_cached","reset_peak_memory_stats","seed","seed_all","set_device","set_per_process_memory_fraction","set_rng_state","set_rng_state_all","set_stream","set_sync_debug_mode","sparse","stream","streams","synchronize","tunable","utilization",]
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.