Revert windows solution. Not working

Author: cehongwang

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -596,7 +596,7 @@ def _save_weight_mapping(self) -> None:
         torch.cuda.empty_cache()
 
     @needs_refit  # type: ignore[misc]
-    def _insert_engine_to_cache(self, hash_val: str, engine: bytes) -> None:
+    def _insert_engine_to_cache(self, hash_val: str, engine: trt.ICudaEngine) -> None:
         serialized_engine = engine.serialize()
         # TODO: @Evan is waiting for TRT's feature to cache the weight-stripped engine
         # if not self.compilation_settings.strip_engine_weights:
@@ -735,7 +735,7 @@ def run(
                         return interpreter_result  # type: ignore[no-any-return]
 
         self._construct_trt_network_def()
-        _LOGGER.info(
+        _LOGGER.debug(
             f"CPU memory usage after network construction: {get_cpu_memory_usage()} MB"
         )
 

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -2,18 +2,15 @@
 
 import io
 import logging
-from typing import Any, List, Optional, Sequence
+from typing import Any, List, NamedTuple, Optional, Sequence
 
 import torch
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._features import ENABLED_FEATURES
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
 from torch_tensorrt.dynamo._settings import CompilationSettings
-from torch_tensorrt.dynamo.conversion._TRTInterpreter import (
-    TRTInterpreter,
-    TRTInterpreterResult,
-)
+from torch_tensorrt.dynamo.conversion._TRTInterpreter import TRTInterpreter
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
 from torch_tensorrt.dynamo.utils import (
     get_cpu_memory_usage,
@@ -24,6 +21,14 @@
 logger = logging.getLogger(__name__)
 
 
+class SerializedInterpreterResult(NamedTuple):
+    serialized_engine: bytes
+    input_names: Sequence[str]
+    output_names: Sequence[str]
+    weight_name_map: Optional[dict[Any, Any]]
+    requires_output_allocator: bool
+
+
 def infer_module_output_dtypes(
     module: torch.fx.GraphModule,
     truncate_double: bool = False,
@@ -34,7 +39,7 @@ def infer_module_output_dtypes(
     """
     outputs = [node for node in module.graph.nodes if node.op == "output"]
     outputs = outputs[0].args
-    return get_output_dtypes(outputs, truncate_double)
+    return get_output_dtypes(outputs, truncate_double)  # type: ignore
 
 
 def interpret_module_to_result(
@@ -44,7 +49,7 @@ def interpret_module_to_result(
     arg_inputs: Optional[Sequence[Input]] = None,
     kwarg_inputs: Optional[dict[str, Any]] = None,
     engine_cache: Optional[BaseEngineCache] = None,
-) -> TRTInterpreterResult:
+) -> SerializedInterpreterResult:
     """Interpret an FX module to a TRTInterpreterResult
     Args:
         module: FX GraphModule to interpret
@@ -84,16 +89,18 @@ def interpret_module_to_result(
     with io.BytesIO() as engine_bytes:
         engine_bytes.write(serialized_engine)
         serialized_engine = engine_bytes.getvalue()
-
-    interpreter_result = TRTInterpreterResult(
-        engine=serialized_engine,
+        logger.debug(
+            f"CPU memory usage after serializing engine: {get_cpu_memory_usage()} MB"
+        )
+    serialized_interpreter_result = SerializedInterpreterResult(
+        serialized_engine=serialized_engine,
         input_names=interpreter_result.input_names,
         output_names=interpreter_result.output_names,
         weight_name_map=interpreter_result.weight_name_map,
         requires_output_allocator=interpreter_result.requires_output_allocator,
     )
 
-    return interpreter_result
+    return serialized_interpreter_result
 
 
 def convert_module(
@@ -132,7 +139,7 @@ def convert_module(
         )
 
     return rt_cls(
-        serialized_engine=interpreter_result.engine,
+        serialized_engine=interpreter_result.serialized_engine,
         input_binding_names=list(interpreter_result.input_names),
         output_binding_names=list(interpreter_result.output_names),
         name=name,

py/torch_tensorrt/dynamo/utils.py

@@ -882,32 +882,3 @@ def release_memory() -> None:
                 logger.warning("Failed to release CPU memory.")
         except Exception:
             logger.warning("Failed to release CPU memory.")
-
-    elif platform.system() == "Windows":
-        from ctypes import wintypes
-
-        kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
-        psapi = ctypes.WinDLL("psapi", use_last_error=True)
-
-        GetCurrentProcess = kernel32.GetCurrentProcess
-        GetCurrentProcess.restype = wintypes.HANDLE
-        hproc = GetCurrentProcess()
-
-        HeapSetInformation = kernel32.HeapSetInformation
-        HeapSetInformation.argtypes = [
-            wintypes.HANDLE,
-            ctypes.c_int,
-            ctypes.c_void_p,
-            ctypes.c_size_t,
-        ]
-        HeapSetInformation.restype = wintypes.BOOL
-        GetProcessHeap = kernel32.GetProcessHeap
-        GetProcessHeap.restype = wintypes.HANDLE
-        ok = False
-        try:
-            HeapOptimizeResources = 3
-            hheap = GetProcessHeap()
-            if HeapSetInformation(hheap, HeapOptimizeResources, None, 0):
-                ok = True
-        except Exception:
-            logger.warning("Failed to release CPU memory.")