def _make_tensor_usage_records_simple_multistream(
    par_ops_seq: List[List[Operator]],
) -> List[TensorUsageRecord]:
    """
    Generalized version of _make_tensor_usage_records() which
    assumes that several ops may be executed on every step.

    Simple multistream algo iteratively tracks sets of operators
    that can be run in parallel independently on each iteration.

    par_ops_seq contains lists of operators that can be run
    in parallel on every algorithm iteration.

    Technically, the regular _make_tensor_usage_records() version
    is similar to the following one:

       def _make_tensor_usage_records(sorted_ops):
         par_ops_seq = [sorted_ops]
         return _make_tensor_usage_records_simple_multistream(par_ops_seq)

    This version is kept as a separate one, because multistreaming
    feature is still somewhat experimental.
    """

    num_of_ops = len(par_ops_seq)
    tensor_records = defaultdict(
        lambda: TensorUsageRecord(
            tensor=None, first_op_idx=num_of_ops, last_op_idx=-1, size=None
        )
    )

    for op_idx, par_ops in enumerate(par_ops_seq):
        for op in par_ops:
            for tensor in op._attrs["inputs"] + op._attrs["outputs"]:
                # Skip weights and inputs since we don't overwrite them.
                # Note that it might be OK to overwrite inputs, but let's be
                # consertative for now and not surprise users. We could always
                # make a flag to do that later if it's needed.
                if tensor._attrs["is_param"]:
                    continue
                name = tensor._attrs["name"]
                this_tensor = tensor_records[name].tensor
                if this_tensor is None:
                    tensor_records[name].tensor = tensor
                else:
                    # make sure we didn't screw up anything
                    assert (
                        tensor == this_tensor
                    ), f"existing tensor: {this_tensor}, new tensor: {tensor}, op: {op}"

                first_op_idx = tensor_records[name].first_op_idx
                last_op_idx = tensor_records[name].last_op_idx
                tensor_records[name].first_op_idx = min(first_op_idx, op_idx)
                tensor_records[name].last_op_idx = max(last_op_idx, op_idx)
                # An output tensor's lifetime extends to the last op.
                if tensor._attrs["is_output"]:
                    tensor_records[name].last_op_idx = num_of_ops - 1

                size = tensor_records[name].size
                tensor_size = tensor.size_bytes(alignment=64)
                if size is None:
                    tensor_records[name].size = tensor_size
                else:
                    # make sure we didn't screw up anything
                    assert size == tensor_size

    # tensor views extend the lifetime of the original tensors
    tensor_views = []
    for name, tensor_record in tensor_records.items():
        this_tensor = tensor_record.tensor
        if this_tensor._attrs["is_view_of"]:
            orig_tensor = _find_original_tensor(this_tensor)
            # view of input
            if orig_tensor._attrs["is_param"]:
                continue
            orig_tensor_name = orig_tensor._attrs["name"]
            assert orig_tensor_name in tensor_records
            tensor_records[orig_tensor_name].last_op_idx = max(
                tensor_records[orig_tensor_name].last_op_idx, tensor_record.last_op_idx
            )
            tensor_views.append(name)

    # remove tensor views from tensor_records
    for name in tensor_views:
        del tensor_records[name]

    # sanity checks
    # make sure we have valid indices and sizes
    records = tensor_records.values()
    for tensor, first_op_idx, last_op_idx, size in records:
        assert tensor is not None
        assert 0 <= first_op_idx < num_of_ops
        assert 0 <= last_op_idx < num_of_ops
        assert first_op_idx <= last_op_idx
        assert size is not None

    return list(records)
Generalized version of _make_tensor_usage_records() which
assumes that several ops may be executed on every step.

Simple multistream algo iteratively tracks sets of operators
that can be run in parallel independently on each iteration.

par_ops_seq contains lists of operators that can be run
in parallel on every algorithm iteration.

Technically, the regular _make_tensor_usage_records() version
is similar to the following one:

def _make_tensor_usage_records(sorted_ops):
par_ops_seq = [sorted_ops]
return _make_tensor_usage_records_simple_multistream(par_ops_seq)

This version is kept as a separate one, because multistreaming
feature is still somewhat experimental.
python/aitemplate/compiler/transform/memory_planning.py