code: def emit_code_with_activation_checkpoint(
      body, ckpt_func, nodes, emit_node_func, delete_unused_value_func
  ):
      # find the activation checkpoint regions
      ckpt_regions = _find_ckpt_regions(nodes)
      start_idx = [item[0] for item in ckpt_regions]
      end_idx = [item[1] for item in ckpt_regions]
      input_vars = []
      output_vars = []
      within_ckpt_region = False
  
      # find the offload regions
      offload_regions, offload_labels = _find_offload_regions(nodes)
      offload_starts = [item[0] for item in offload_regions]
      offload_ends = [item[1] for item in offload_regions]
      offload_inputs = []
      offload_outputs = []
      within_offload_region = False
  
      node_list = list(nodes)
  
      # use this variable to avoid inserting hook functions
      # to ckpt_func repeatedly
      is_hook_inserted = False
  
      # find the input and output var names for each region
      for idx, (start, end) in enumerate(ckpt_regions):
          ckpt_node_list = node_list[start : end + 1]
          inputs, outputs = _find_input_and_output_nodes(ckpt_node_list)
          input_vars.append(inputs)
          output_vars.append(outputs)
  
      # find the input and output var names for each offload region
      for idx, (start, end) in enumerate(offload_regions):
          offload_node_list = node_list[start : end + 1]
          inputs, outputs = _find_input_and_output_nodes(
              offload_node_list
          )
          offload_inputs.append(inputs)
          offload_outputs.append(outputs)
  
      # append code text to body
      for idx, node in enumerate(node_list):
          # if this is the first node of the ckpt region
          # append the ckpt function definition
          if idx in start_idx:
              label = start_idx.index(idx)
              ckpt_fn_def = _gen_ckpt_fn_def(label, input_vars[label])
              ckpt_func.append(f"{ckpt_fn_def}\n")
              within_ckpt_region = True
  
          if idx in offload_starts:
              offload_label = offload_labels[offload_starts.index(idx)]
              _, offload_input, offload_bar = offload_label
              within_offload_region = True
  
              # insert hook functions if needed
              if not is_hook_inserted:
                  pack_hook, unpack_hook = _gen_saved_tensors_hooks()
                  ckpt_func.insert(
                      0, "\n".join([pack_hook, unpack_hook]) + "\n"
                  )
                  is_hook_inserted = True
  
              if offload_input and offload_bar:
                  body.append(_gen_save_on_cpu_context())
  
              elif offload_input:
                  for par in offload_inputs[offload_label[0]]:
                      body.append(f"setattr({par}, 'offload', True)\n")
                  body.append(
                      _gen_save_tensors_hooks_context(
                          offload_input=True
                      )
                  )
  
              else:
                  for par in offload_inputs[offload_label[0]]:
                      body.append(f"setattr({par}, 'offload', False)\n")
                  body.append(
                      _gen_save_tensors_hooks_context(
                          offload_input=False
                      )
                  )
  
          # NOTE: emit_node does not emit a string with newline. It depends
          # on delete_unused_values to append one
          # NOTE: currently we separate body and ckpt_func definition
          if within_ckpt_region:
              emit_node_func(node, ckpt_func)
              ckpt_func[-1] = "    " + ckpt_func[-1]
              delete_unused_value_func(node, ckpt_func)
  
          elif within_offload_region:
              emit_node_func(node, body)
              body[-1] = "    " + body[-1]
              delete_unused_value_func(node, body)
  
          else:
              emit_node_func(node, body)
              delete_unused_value_func(node, body)
  
          if idx in end_idx:
              # if this is the last node of the ckpt region
              # generate return statement
              label = end_idx.index(idx)
              return_statement = _gen_ckpt_output(output_vars[label])
              return_statement = f"    {return_statement}\n\n"
              ckpt_func.append(return_statement)
  
              # we need to check if the checkpoint need to offload the input
              start_node_idx = start_idx[label]
              if "activation_offload" in node_list[start_node_idx].meta:
                  activation_offload = node_list[start_node_idx].meta[
                      "activation_offload"
                  ]
              else:
                  activation_offload = False
  
              # we need to check if the checkpoint need use_reentrant=False
              use_reentrant = True
              non_leaf_input = 0
              for var in input_vars[label]:
                  input_node = next(
                      item for item in node_list if item.name == var
                  )
                  if input_node.op != "placeholder":
                      non_leaf_input = 1
                  for user in input_node.users:
                      if "activation_checkpoint" in user.meta:
                          if (
                              user.meta["activation_checkpoint"]
                              == label
                          ):
                              if user.op == "call_module":
                                  if hasattr(
                                      user.graph.owning_module.get_submodule(
                                          user.target
                                      ),
                                      "inplace",
                                  ):
                                      use_reentrant = not user.graph.owning_module.get_submodule(
                                          user.target
                                      ).inplace
  
                              elif user.op == "call_function":
                                  if "inplace" in user.kwargs:
                                      use_reentrant = not user.kwargs[
                                          "inplace"
                                      ]
  
              # if all the inputs are leaf nodes, we need to set use_reentrant = False
              if not non_leaf_input:
                  use_reentrant = False
  
              # generate checkpoint function call in a new line
              usage = _gen_ckpt_usage(
                  label,
                  activation_offload,
                  input_vars[label],
                  output_vars[label],
                  use_reentrant,
              )
              usage += "\n"
              body.append(usage)
              within_ckpt_region = False
  
          if idx in offload_ends:
              within_offload_region = False
  
file_path: colossalai/colossalai/fx/codegen/activation_checkpoint_codegen.py
imports: from typing import Any, Dict, Iterable, List, Tuple
  import torch
  import colossalai
label_desc: Use image node_id 14 for calling a global function with example usage: emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func, delete_unused_value_func) without return types