def simple_multistream_memory_planning(sorted_graph: List[Tensor]):
    """
    A specialized case for simple multi-stream execution.
    It uses more or slightly more GPU memory than greedy_by_size_memory_planner,
    depending on the input graph, but still significantly less
    than naive_memory_planning.
    """
    from aitemplate.utils.graph_utils import track_graph_timings

    # track the sequence
    time_stats = track_graph_timings(sorted_graph, {})

    # sort all operators by parallel execution order
    ops_by_order = defaultdict(list)
    for op, tracking in time_stats.op_parallel_trackers.items():
        ops_by_order[tracking.execution_order].append(op)

    # convert Dict[int, List[Operator]] into List[List[Operator]]
    max_parallel_ops = multistream_max_mem_parallel_ops()
    par_ops_seq = split_simple_multistream_parallel_ops(ops_by_order, max_parallel_ops)

    tensor_usage_records = _make_tensor_usage_records_simple_multistream(par_ops_seq)

    return _greedy_by_size_memory_planning(sorted_graph, tensor_usage_records)
A specialized case for simple multi-stream execution.
It uses more or slightly more GPU memory than greedy_by_size_memory_planner,
depending on the input graph, but still significantly less
than naive_memory_planning.
python/aitemplate/compiler/transform/memory_planning.py