Skip to content

Commit de93114

Browse files
committed
Added a potential solution for windows
1 parent d99f183 commit de93114

File tree

1 file changed

+28
-0
lines changed

1 file changed

+28
-0
lines changed

py/torch_tensorrt/dynamo/utils.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -881,3 +881,31 @@ def release_memory() -> None:
881881
logger.warning("Failed to release CPU memory.")
882882
except Exception:
883883
logger.warning("Failed to release CPU memory.")
884+
885+
elif platform.system() == "Windows":
886+
try:
887+
from ctypes import wintypes as wt
888+
889+
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
890+
psapi = ctypes.WinDLL("psapi", use_last_error=True)
891+
892+
GetCurrentProcess = kernel32.GetCurrentProcess
893+
GetCurrentProcess.restype = wt.HANDLE
894+
hProcess = GetCurrentProcess()
895+
896+
EmptyWorkingSet = psapi.EmptyWorkingSet
897+
EmptyWorkingSet.argtypes = [wt.HANDLE]
898+
EmptyWorkingSet.restype = wt.BOOL
899+
EmptyWorkingSet(hProcess)
900+
901+
SetProcessWorkingSetSize = kernel32.SetProcessWorkingSetSize
902+
SetProcessWorkingSetSize.argtypes = [
903+
wt.HANDLE,
904+
ctypes.c_size_t,
905+
ctypes.c_size_t,
906+
]
907+
SetProcessWorkingSetSize.restype = wt.BOOL
908+
SIZE_T_MAX = ctypes.c_size_t(-1).value # 0xFFFFFFFFFFFFFFFF on x64
909+
SetProcessWorkingSetSize(hProcess, SIZE_T_MAX, SIZE_T_MAX)
910+
except Exception:
911+
logger.warning("Failed to release CPU memory.")

0 commit comments

Comments
 (0)