class_docstr: This is the core class for AutoChunk.
  
  It defines the framework of the strategy of AutoChunk.
  Chunks will be selected one by one until search stops.
  
  The chunk search is as follows:
  1. find the peak memory node
  2. find the max chunk region according to the peak memory node
  3. find all possible chunk regions in the max chunk region
  4. find the best chunk region for current status
  5. goto 1
  
  Attributes:
      gm: graph model
      print_mem (bool): print estimated memory
      trace_index: trace the flow of every dim of every node to find all free dims
      trace_flow: determine the region chunk strategy
      reorder_graph: reorder nodes to improve chunk efficiency
      estimate_memory: estimate memory with chunk
      select_chunk: select the best chunk region
  
  Args:
      gm: graph model
      max_memory (int): max memory in MB
      print_mem (bool): print estimated memory
class_docstr_tok: ["This","is","the","core","class","for","AutoChunk",".","It","defines","the","framework","of","the","strategy","of","AutoChunk",".","Chunks","will","be","selected","one","by","one","until","search","stops",".","The","chunk","search","is","as","follows",":","1.","find","the","peak","memory","node","2.","find","the","max","chunk","region","according","to","the","peak","memory","node","3.","find","all","possible","chunk","regions","in","the","max","chunk","region","4.","find","the","best","chunk","region","for","current","status","5.","goto","1","Attributes",":","gm",":","graph","model","print_mem","(","bool",")",":","print","estimated","memory","trace_index",":","trace","the","flow","of","every","dim","of","every","node","to","find","all","free","dims","trace_flow",":","determine","the","region","chunk","strategy","reorder_graph",":","reorder","nodes","to","improve","chunk","efficiency","estimate_memory",":","estimate","memory","with","chunk","select_chunk",":","select","the","best","chunk","region","Args",":","gm",":","graph","model","max_memory","(","int",")",":","max","memory","in","MB","print_mem","(","bool",")",":","print","estimated","memory"]
docstr: Search every possible region within the max chunk region.
  
  Args:
      max_chunk_region (Tuple)
      peak_node (Node): peak memory node
  
  Returns:
      possible_chunk_region (List)
docstr_tok: ["Search","every","possible","region","within","the","max","chunk","region",".","Args",":","max_chunk_region","(","Tuple",")","peak_node","(","Node",")",":","peak","memory","node","Returns",":","possible_chunk_region","(","List",")"]
code: def _search_possible_chunk_regions(
      self, max_chunk_region: Tuple, peak_region: Node
  ) -> List:
      """
      Search every possible region within the max chunk region.
  
      Args:
          max_chunk_region (Tuple)
          peak_node (Node): peak memory node
  
      Returns:
          possible_chunk_region (List)
      """
      possible_chunk_region = []
      output_trace = copy.deepcopy(self.trace_indice.indice_trace_list)
      input_trace = []  # trace of a node's input nodes
      for _, n in enumerate(self.node_mgr.get_node_list()):
          cur_trace = {}
          for arg in n.args:
              if type(arg) == type(
                  n
              ) and not is_non_compute_node_except_placeholder(arg):
                  cur_trace[
                      arg
                  ] = self.trace_indice._find_trace_from_node(arg)
          input_trace.append(cur_trace)
  
      for start_idx in range(max_chunk_region[0], peak_region[0] + 1):
          for end_idx in range(peak_region[1], max_chunk_region[1] + 1):
              # skip non compute nodes
              if is_non_compute_node(
                  self.node_mgr.get_node_by_idx(start_idx)
              ) or is_non_compute_node(
                  self.node_mgr.get_node_by_idx(end_idx)
              ):
                  continue
              # select free dim
              chunk_info = self._find_chunk_info(
                  input_trace, output_trace, start_idx, end_idx
              )
              if len(chunk_info) > 0:
                  possible_chunk_region.extend(chunk_info)
      return possible_chunk_region
  
file_path: colossalai/colossalai/autochunk/search_chunk.py
imports: import copy
  from typing import Dict, List, Tuple
  from torch.fx.node import Node
  from .estimate_memory import EstimateMemory
  from .reorder_graph import ReorderGraph
  from .select_chunk import SelectChunk
  from .trace_flow import TraceFlow
  from .trace_indice import TraceIndice
  from .utils import NodeMgr, get_logger, is_non_compute_node, is_non_compute_node_except_placeholder
label_desc: Use image node_id 6 for calling the SearchChunk obj's underlying member method code with example usage: obj._search_possible_chunk_regions(max_chunk_region, peak_region) and returns: possible_chunk_region