def split_simple_multistream_parallel_ops(ops_by_order, max_parallel_ops: int):
    """
    Make sure that no more than max_parallel_ops operators are run in parallel.

    Say, on the first step op1, op2 and op3 can be executed in parallel.
    On the second one, it is op4 and op5.
    On the third one it is op6, op7, op8, op9.
    Then, ops_by_order is something like
      { 1: [op1, op2, op3], 2: [op4, op5], 3: [op6, op7, op8, op9] }
    Given max_parallel_ops=2, the output will be:
      [[op1, op2], [op3], [op4, op5], [op6, op7], [op8, op9]]

    Parameters
    ----------
    ops_by_order : Dict[int, List[Operator]]
        A dictionary, its keys represent the execution order
        and its values represent operators that are executed in parallel.
    max_parallel_ops : int
        Number of operators that are allowed to be run in parallel

    Output : List[List[Operator]]
        transformed sequence of operators to execute.

    """
    assert max_parallel_ops > 0

    # todo: a better splitting algorithm can be implemented,
    # the one that splits operators into max_parallel_ops buckets
    # so that the amount of needed memory is about the same.
    # use priority_queue for this and iteratively add to the
    # bucket that has the lowest 'assigned' memory.

    output = []

    execution_orders = sorted(ops_by_order.keys())
    for execution_order in execution_orders:
        ops = ops_by_order[execution_order]

        ops_parallel = []
        for op in ops:
            ops_parallel.append(op)
            if len(ops_parallel) >= max_parallel_ops:
                output.append(ops_parallel)
                ops_parallel = []

        if len(ops_parallel) > 0:
            output.append(ops_parallel)

    # done
    return output
Make sure that no more than max_parallel_ops operators are run in parallel.

Say, on the first step op1, op2 and op3 can be executed in parallel.
On the second one, it is op4 and op5.
On the third one it is op6, op7, op8, op9.
Then, ops_by_order is something like
{ 1: [op1, op2, op3], 2: [op4, op5], 3: [op6, op7, op8, op9] }
Given max_parallel_ops=2, the output will be:
[[op1, op2], [op3], [op4, op5], [op6, op7], [op8, op9]]

Parameters
----------
ops_by_order : Dict[int, List[Operator]]
A dictionary, its keys represent the execution order
and its values represent operators that are executed in parallel.
max_parallel_ops : int
Number of operators that are allowed to be run in parallel

Output : List[List[Operator]]
transformed sequence of operators to execute.
python/aitemplate/utils/graph_utils.py