content: """
  
  .. _low_cpu_memory_compilation:
  
  Low CPU Memory Compilation Example
  ==================================
  
  This example demonstrates compiling a model with a bounded CPU (host) memory
  budget using Torch-TensorRT Dynamo. Limiting host RAM use is helpful on
  memory-constrained machines or when compiling very large models.
  
  Key notes:
  - The toy model below has roughly 430 MB of parameters. We set the CPU
    memory budget to 2 GiB. At compile time, only about 900 MB of host RAM
    may remain available. We expect at most 403 * 4 = 1612 MB of memory to be used by the model.
    So the model is partitioned into two subgraphs to fit the memory budget.
  
  - Performance impact varies by model. When the number of TensorRT engines
    created is small, the impact is typically minimal.
  
  """
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch_tensorrt as torchtrt
  from torch_tensorrt.dynamo.conversion import CompilationSettings
  
  
  class net(nn.Module):
      def __init__(self):
          super().__init__()
          # Intentionally large layers to stress host memory during compilation.
          self.conv1 = nn.Conv2d(1024, 4096, 3, padding=1)
          self.bn1 = nn.BatchNorm2d(4096)
          self.conv2 = nn.Conv2d(4096, 1024, 3, padding=1)
          self.bn2 = nn.BatchNorm2d(1024)
          self.fc1 = nn.Linear(1024 * 56 * 56, 10)
  
      def forward(self, x):
          x = self.conv1(x)
          x = self.bn1(x)
          x = F.relu(x)
          x = F.max_pool2d(x, (2, 2))
          x = self.conv2(x)
          x = self.bn2(x)
          x = F.relu(x)
          x = F.max_pool2d(x, (2, 2))
          x = torch.flatten(x, 1)
          return self.fc1(x)
  
  
  model = net().eval()
  model.to("cuda")
  inputs = [torch.randn((1, 1024, 224, 224)).to("cuda")]
  
  enabled_precisions = {torch.float}
  use_python_runtime = False
  
  compilation_options = {
      "use_python_runtime": use_python_runtime,
      "enabled_precisions": enabled_precisions,
      "min_block_size": 1,
      "immutable_weights": True,
      "reuse_cached_engines": False,
      "enable_resource_partitioning": True,
      "cpu_memory_budget": 2 * 1024 * 1024 * 1024,  # 2 GiB in bytes
  }
  
  settings = CompilationSettings(**compilation_options)
  with torchtrt.dynamo.Debugger(
      log_level="debug",
      logging_dir="/home/profile/logging/moe",
      engine_builder_monitor=False,
  ):
  
      exp_program = torch.export.export(model, tuple(inputs))
      trt_gm = torchtrt.dynamo.compile(
          exp_program,
          inputs=inputs,
          **compilation_options,
      )
  
      # Expect two back-to-back TensorRT engines due to partitioning under the memory budget.
      print(trt_gm)
  
  
  """
  You should be able to see two back-to-back TensorRT engines in the graph
  
  Graph Structure:
  
     Inputs: List[Tensor: (1, 1024, 224, 224)@float32]
      ...
      TRT Engine #1 - Submodule name: _run_on_acc_0_resource_split_0
       Engine Inputs: List[Tensor: (1, 1024, 224, 224)@float32]
       Number of Operators in Engine: 9
       Engine Outputs: List[Tensor: (1, 1024, 112, 112)@float32]
      ...
      TRT Engine #2 - Submodule name: _run_on_acc_0_resource_split_1
       Engine Inputs: List[Tensor: (1, 1024, 112, 112)@float32]
       Number of Operators in Engine: 3
       Engine Outputs: List[Tensor: (1, 10)@float32]
      ...
     Outputs: List[Tensor: (1, 10)@float32]
  
    ------------------------- Aggregate Stats -------------------------
  
     Average Number of Operators per TRT Engine: 6.0
     Most Operators in a TRT Engine: 9
  
    ********** Recommendations **********
  
     - For minimal graph segmentation, select min_block_size=9 which would generate 1 TRT engine(s)
     - For moderate graph segmentation, select min_block_size=6 which would generate 1 TRT engine(s)
     - The current level of graph segmentation is equivalent to selecting min_block_size=3 which generates 2 TRT engine(s)
  GraphModule(
    (_run_on_acc_0_resource_split_0): TorchTensorRTModule()
    (_run_on_acc_0_resource_split_1): TorchTensorRTModule()
  )
  
  
  
  def forward(self, x):
      x, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
      _run_on_acc_0_resource_split_0 = self._run_on_acc_0_resource_split_0(x);  x = None
      _run_on_acc_0_resource_split_1 = self._run_on_acc_0_resource_split_1(_run_on_acc_0_resource_split_0);  _run_on_acc_0_resource_split_0 = None
      return pytree.tree_unflatten((_run_on_acc_0_resource_split_1,), self._out_spec)
  )
  """
  