def _compute_workspace(sorted_graph: List[Tensor]) -> Workspace:
    """
    Compute the workspace for the model, which can be used as scratch memory by ops.
    This pass examines two attributes on every function in the graph:
    - workspace: The amount of memory in bytes to be used as shared scratch memory.
      Here, "shared" means that other ops are allowed to write to this memory.
    - unique_workspace: The amount of memory in bytes to be used as exclusive scratch memory.
      If set, this pass will assign the op a "unique_workspace_offset". This can be used at
      codegen time to set a pointer to the region of exclusive shared memory.

    The returned Workspace has two attributes:
    - shared_size: The total memory needed for all op's shared scratch memory (i.e. the maximum
    of all workspace attributes)
    - unique_size: The total memory needed for all unique scratch memory (i.e. the sum of
    all unique_workspace attributes)

    During codegen, the workspace gets set up like this:
    [--unique 1--][--unique 2--]...[--unique N--][--shared--]
    """
    unique_workspace_size = 0
    max_workspace = 0
    for node in sorted_graph:
        for func in node._attrs["src_ops"]:
            if "workspace" in func._attrs:
                max_workspace = max(max_workspace, func._attrs["workspace"])
            if (
                "unique_workspace" in func._attrs
                and "unique_workspace_offset" not in func._attrs
            ):
                func._attrs["unique_workspace_offset"] = unique_workspace_size
                unique_workspace_size += func._attrs["unique_workspace"]
    return Workspace(max_workspace, unique_workspace_size)

    Compute the workspace for the model, which can be used as scratch memory by ops.
    This pass examines two attributes on every function in the graph:
    - workspace: The amount of memory in bytes to be used as shared scratch memory.
      Here, "shared" means that other ops are allowed to write to this memory.
    - unique_workspace: The amount of memory in bytes to be used as exclusive scratch memory.
      If set, this pass will assign the op a "unique_workspace_offset". This can be used at
      codegen time to set a pointer to the region of exclusive shared memory.

    The returned Workspace has two attributes:
    - shared_size: The total memory needed for all op's shared scratch memory (i.e. the maximum
    of all workspace attributes)
    - unique_size: The total memory needed for all unique scratch memory (i.e. the sum of
    all unique_workspace attributes)

    During codegen, the workspace gets set up like this:
    [--unique 1--][--unique 2--]...[--unique N--][--shared--]
    
python/aitemplate/compiler/transform/memory_planning.py