# Copyright (c) Meta Platforms, Inc. and affiliates.

# All rights reserved.

#

# This source code is licensed under the BSD-style license found in the

# LICENSE file in the root directory of this source tree.

import argparse

import json

from typing import Any, Callable, Dict, List, Optional, Tuple

import torch

from executorch.devtools import parse_etrecord

from executorch.exir import ExportedProgram

from executorch.exir.backend.backend_api import LoweredBackendModule

def _get_tensor_data(node: torch.fx.Node, tensor: torch.Tensor) -> Dict[str, Any]:

return {

"name": node.name,

"numel": tensor.numel(),

"dtype": str(tensor.dtype)[6:], # Remove "torch." prefix

"element_size": tensor.element_size(),

"shape": list(tensor.shape),

"num_bytes": tensor.element_size() * tensor.numel(),

"nn_module_stack": (

str(node.meta["nn_module_stack"])

if "nn_module_stack" in node.meta

else None

),

}

def _get_delegate_blob_data(

node: torch.fx.Node,

lowered_backend_module: LoweredBackendModule,

delegate_deserializers: Optional[

Dict[str, Callable[[bytes], Dict[str, Any]]]

] = None,

) -> Dict[str, Any]:

delegate_blob_data = {

"name": node.name,

"backend_id": lowered_backend_module.backend_id,

"num_bytes": len(lowered_backend_module.processed_bytes),

}

if (

delegate_deserializers is not None

and lowered_backend_module.backend_id in delegate_deserializers

):

delegate_blob_data.update(

delegate_deserializers[lowered_backend_module.backend_id](

lowered_backend_module.processed_bytes

)

)

return delegate_blob_data

def _get_nested_model_data( # noqa: C901

graph_module: torch.fx.GraphModule,

delegate_deserializers: Optional[

Dict[str, Callable[[bytes], Dict[str, Any]]]

] = None,

tensor_data: Optional[List[Dict[str, Any]]] = None,

delegate_blob_data: Optional[List[Dict[str, Any]]] = None,

exported_program: Optional["ExportedProgram"] = None,

) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:

if tensor_data is None:

tensor_data = []

if delegate_blob_data is None:

delegate_blob_data = []

for node in graph_module.graph.nodes:

if node.op == "placeholder" and exported_program is not None:

sig = exported_program.graph_signature

fqn = None

if node.name in getattr(sig, "inputs_to_parameters", {}):

fqn = sig.inputs_to_parameters[node.name]

elif node.name in getattr(sig, "inputs_to_buffers", {}):

fqn = sig.inputs_to_buffers[node.name]

if fqn is not None:

tensor = exported_program.state_dict.get(fqn)

if isinstance(tensor, torch.Tensor):

tensor_data.append(_get_tensor_data(node, tensor))

elif node.op == "get_attr":

node_attr = getattr(node.graph.owning_module, node.target)

if isinstance(node_attr, torch.Tensor):

tensor_data.append(_get_tensor_data(node, node_attr))

elif isinstance(node_attr, torch.fx.GraphModule):

_get_nested_model_data(

node_attr, delegate_deserializers, tensor_data, delegate_blob_data

)

elif isinstance(node_attr, LoweredBackendModule):

delegate_blob_data.append(

_get_delegate_blob_data(node, node_attr, delegate_deserializers)

)

return (tensor_data, delegate_blob_data)

def generate_model_size_information(

model: ExportedProgram,

delegate_deserializers: Optional[

Dict[str, Callable[[bytes], Dict[str, Any]]]

] = None,

flatbuffer: Optional[bytes] = None,

) -> Dict[str, Any]:

"""

Generate a json-serializable Dict containing information about a model's

size. This includes data about individual tensors and delegate blobs.

Optionally:

- delegate_deserializers can be provided to manually specify additional

information to include for delegate blobs for specific backends.

- flatbuffer can be provided to include a comparison of total tensor data

size to overall model size

"""

tensor_and_delegate_blob_data = _get_nested_model_data(

model.graph_module, delegate_deserializers, exported_program=model

)

for data_list in tensor_and_delegate_blob_data:

data_list.sort(key=lambda data: data["num_bytes"], reverse=True)

(tensor_data, delegate_blob_data) = tensor_and_delegate_blob_data

total_tensor_data_size = sum(data["num_bytes"] for data in tensor_data)

total_delegate_blob_data_size = sum(

data["num_bytes"] for data in delegate_blob_data

)

overview = {

"total_tensor_data_size": total_tensor_data_size,

"total_delegate_blob_data_size": total_delegate_blob_data_size,

}

if flatbuffer is not None:

model_size = len(flatbuffer)

overview.update(

{

"serialization_metadata_size": (

model_size - total_tensor_data_size - total_delegate_blob_data_size

),

"model_size": model_size,

}

)

return {

"tensor_data": tensor_data,

"delegate_blob_data": delegate_blob_data,

"overview": overview,

}

def parse_args():

parser = argparse.ArgumentParser()

parser.add_argument(

"--etrecord_path",

required=True,

help="The path to the ETRecord for the model to generate size information for",

)

parser.add_argument(

"--output_path",

default="model_size_information.json",

help="The output path for the model size information as a json file",

)

args = parser.parse_args()

return args

def main():

args = parse_args()

etrecord = parse_etrecord(args.etrecord_path)

all_model_size_information = [

generate_model_size_information(

model=exported_program,

delegate_deserializers=None,

flatbuffer=None,

)

for (name, exported_program) in etrecord.graph_map.items()

]

with open(args.output_path, "w") as f:

f.write(json.dumps(all_model_size_information))

if __name__ == "__main__":

main()