onnxscript 0.1.0__py3-none-any.whl

onnxscript/tools/benchmark/benchmark_run.py

@@ -0,0 +1,140 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# pylint: disable=consider-using-with,import-outside-toplevel
+from __future__ import annotations
+
+import multiprocessing
+import os
+import platform
+import re
+import subprocess
+import sys
+
+
+class BenchmarkError(RuntimeError):
+    pass
+
+
+def get_machine() -> dict[str, str | int | float | tuple[int, int]]:
+    """Returns the machine specification."""
+    config: dict[str, str | int | float | tuple[int, int]] = dict(
+        machine=str(platform.machine()),
+        processor=str(platform.processor()),
+        version=str(sys.version),
+        config=int(multiprocessing.cpu_count()),
+        executable=str(sys.executable),
+    )
+    try:
+        import torch.cuda
+    except ImportError:
+        return config
+
+    config["has_cuda"] = bool(torch.cuda.is_available())
+    if config["has_cuda"]:
+        config["capability"] = torch.cuda.get_device_capability(0)
+        config["device_name"] = str(torch.cuda.get_device_name(0))
+    return config
+
+
+def _cmd_line(script_name: str, **kwargs: dict[str, str | int | float]) -> list[str]:
+    args = [sys.executable, "-m", script_name]
+    for k, v in kwargs.items():
+        args.append(f"--{k}")
+        args.append(str(v))
+    return args
+
+
+def _extract_metrics(text: str) -> dict[str, str]:
+    reg = re.compile(r":(.*?),(.*.?);")
+    res = reg.findall(text)
+    if len(res) == 0:
+        return {}
+    return dict(res)
+
+
+def _make_prefix(script_name: str, index: int) -> str:
+    name = os.path.splitext(script_name)[0]
+    return f"{name}_dort_c{index}_"
+
+
+def run_benchmark(
+    script_name: str,
+    configs: list[dict[str, str | int | float]],
+    verbose: int = 0,
+    stop_if_exception: bool = True,
+    dort_dump: bool = False,
+) -> list[dict[str, str | int | float | tuple[int, int]]]:
+    """
+    Runs a script multiple times and extract information from the output
+    following the pattern ``:<metric>,<value>;``.
+
+    :param script_name: python script to run
+    :param configs: list of execution to do
+    :param stop_if_exception: stop if one experiment failed, otherwise continue
+    :param verbose: use tqdm to follow the progress
+    :param dort_dump: dump onnx file if dort is used
+    :return: values
+    """
+    if verbose:
+        try:
+            from tqdm import tqdm
+
+            loop = tqdm(configs)
+        except ImportError:
+            loop = configs
+    else:
+        loop = configs
+
+    data: list[dict[str, str | int | float | tuple[int, int]]] = []
+    for i, config in enumerate(loop):
+        cmd = _cmd_line(script_name, **config)
+
+        if dort_dump:
+            os.environ["ONNXRT_DUMP_PATH"] = _make_prefix(script_name, i)
+        else:
+            os.environ["ONNXRT_DUMP_PATH"] = ""
+        if verbose > 3:
+            print(f"[run_benchmark] cmd={cmd if isinstance(cmd, str) else ' '.join(cmd)}")
+
+        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        try:
+            res = p.communicate(timeout=30)
+            out, err = res
+            serr = err.decode("utf-8", errors="ignore")
+        except subprocess.TimeoutExpired as e:
+            p.kill()
+            res = p.communicate()
+            out, err = res
+            serr = f"{e}\n:timeout,1;{err.decode('utf-8', errors='ignore')}"
+        sout = out.decode("utf-8", errors="ignore")
+
+        if "ONNXRuntimeError" in serr or "ONNXRuntimeError" in sout:
+            if stop_if_exception:  # pylint: disable=no-else-raise
+                raise RuntimeError(
+                    f"Unable to continue with config {config} due to the "
+                    f"following error\n{serr}"
+                    f"\n----OUTPUT--\n{sout}"
+                )
+
+        metrics = _extract_metrics(sout)
+        if len(metrics) == 0:
+            if stop_if_exception:  # pylint: disable=no-else-raise
+                raise BenchmarkError(
+                    f"Unable (2) to continue with config {config}, no metric was "
+                    f"collected.\n--ERROR--\n{serr}\n--OUTPUT--\n{sout}"
+                )
+            else:
+                metrics = {}
+        metrics.update(config)
+        metrics["ERROR"] = serr
+        metrics["OUTPUT"] = sout
+        metrics["CMD"] = f"[{' '.join(cmd)}]"
+        data.append(metrics)  # type: ignore[arg-type]
+        if verbose > 5:
+            print("--------------- ERROR")
+            print(serr)
+        if verbose >= 10:
+            print("--------------- OUTPUT")
+            print(sout)
+
+    return data

onnxscript/tools/benchmark/export_model.py

@@ -0,0 +1,207 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# pylint: disable=import-outside-toplevel
+
+import hashlib
+import pprint
+import textwrap
+import time
+from typing import Any
+
+
+def main(args=None):
+    import onnxscript.tools.benchmark
+
+    kwargs: dict[str, Any] = onnxscript.tools.benchmark.get_parsed_args(
+        "export_model",
+        description=textwrap.dedent(
+            """Measures the inference time for a particular model.
+            This script can be used to quickly evaluate the improvment made by a pattern optimization
+            for a particular model.
+
+            If one value contains ",", the script understand multiple commands
+            must be run. It computes all the possible configurations.
+            In that case, it produces a csv file (if output_data is not empty) with all the results.
+
+            Example with a large phi model::
+
+                python -m onnxscript.tools.benchmark.export_model --model phi --device cuda --config large --num_hidden_layers=6 --dtype=float32 --dynamic=0 --verbose=1 --exporter=dynamo
+
+            Example with a medium llama model::
+
+                python -m onnxscript.tools.benchmark.export_model --model llama --device cuda --config medium --num_hidden_layers=1 --dtype=float32 --dynamic=0 --verbose=1 --exporter=dynamo --optimization=rewrite/optimize/inline/llama0/onnxruntime
+            """
+        ),
+        repeat=(10, "number of inferences to measure"),
+        warmup=(5, "number of inferences to warm"),
+        model=("phi", "model to measure, llama, mistral, phi, ..."),
+        exporter=("dynamo", "script, dynamo"),
+        device=("cpu", "'cpu' or 'cuda'"),
+        target_opset=(18, "opset to convert into, use with backend=custom"),
+        config=("small", "default, medium, or small to test"),
+        verbose=(0, "verbosity"),
+        dump_folder=("", "if not empty, dump the model in that folder"),
+        dump_ort=(1, "produce the model optimized by onnxruntime"),
+        ort_optimize=(1, "enable or disable onnxruntime optimization"),
+        dtype=("default", "cast the model and the inputs into this type"),
+        dynamic=(0, "use dynamic shapes"),
+        num_hidden_layers=(1, "number of hidden layers"),
+        with_mask=(1, "with or without mask, dynamo may fail with a mask"),
+        optimization=(
+            "",
+            "optimization scenario, comma separated value, optimize, rewrite, "
+            "inline, set of patterns (default, onnxruntime, customops)",
+        ),
+        implementation=("eager", "eager or sdpa"),
+        memory_peak=(0, "measure the memory peak during conversion"),
+        output_data=(
+            "export_model.csv",
+            "produces a csv file with the data if multiple configurations are tested",
+        ),
+        new_args=args,
+    )
+    if onnxscript.tools.benchmark.multi_run(kwargs):
+        import onnxscript.tools.benchmark.benchmark_run
+
+        configs = onnxscript.tools.benchmark.make_configs(kwargs)
+        data = onnxscript.tools.benchmark.benchmark_run.run_benchmark(
+            "onnxscript.tools.benchmark.export_model",
+            configs,
+            kwargs["verbose"],
+            stop_if_exception=False,
+        )
+        if kwargs["verbose"] > 2:
+            pprint.pprint(data if kwargs["verbose"] > 3 else data[:2])
+        if kwargs["output_data"]:
+            df = onnxscript.tools.benchmark.make_dataframe_from_benchmark_data(data)
+            df.to_csv(kwargs["output_data"], index=False)
+            df.to_excel(kwargs["output_data"] + ".xlsx", index=False)
+            if kwargs["verbose"]:
+                print(df)
+    else:
+        print("-------------------")
+        print("[export_model]")
+        pprint.pprint(kwargs)
+        print("-------------------")
+
+        # Import is delayed so that help is being display faster (without having to import heavy packages).
+        import onnxscript.tools
+        import onnxscript.tools.memory_peak
+        import onnxscript.tools.transformers_models
+
+        print(
+            f"[export_model] create the model and inputs for {kwargs['model']!r} and config {kwargs['config']!r}"
+        )
+        begin = time.perf_counter()
+        model, example_inputs, dynamic_shapes = (
+            onnxscript.tools.transformers_models.get_model_and_inputs(
+                warmup=kwargs["warmup"],
+                repeat=kwargs["repeat"],
+                model=kwargs["model"],
+                config=kwargs["config"],
+                dynamic_shapes=kwargs["dynamic"],
+                device=kwargs["device"],
+                num_hidden_layers=kwargs["num_hidden_layers"],
+                with_mask=kwargs["with_mask"],
+                implementation=kwargs["implementation"],
+                dtype=kwargs["dtype"],
+            )
+        )
+        print(f"[export_model] model created in {time.perf_counter() - begin}")
+        if kwargs["dynamic"]:
+            print(f"[export_model] dynamic_shapes={dynamic_shapes}")
+        msg = [tuple(i.shape for i in inp) for inp in example_inputs]
+        print(f"[export_model] input_shapes={msg}")
+        conversion: dict[str, Any] = {}
+        memory_stats: dict[str, float] = {}
+
+        if kwargs["exporter"] == "eager":
+            print("[export_model] start benchmark")
+            begin = time.perf_counter()
+            result = onnxscript.tools.benchmark.run_inference(
+                model,
+                example_inputs,
+                warmup=kwargs["warmup"],
+                repeat=kwargs["repeat"],
+                verbose=kwargs["verbose"],
+            )
+            print(f"[export_model] benchmark done in {time.perf_counter() - begin}")
+        else:
+            print(
+                f"[export_model] export to onnx with exporter={kwargs['exporter']!r} "
+                f"and optimization={kwargs['optimization']!r}"
+            )
+            begin = time.perf_counter()
+            if kwargs["optimization"]:
+                m = hashlib.sha256()
+                m.update(kwargs["optimization"].encode())
+                so = m.hexdigest()[:5]
+            else:
+                so = ""
+            name = "_".join(
+                [
+                    kwargs["model"],
+                    kwargs["exporter"],
+                    "dynamic" if kwargs["dynamic"] else "static",
+                    kwargs["dtype"].replace("float", "fp"),
+                    kwargs["device"],
+                    kwargs["config"],
+                    f"h{kwargs['num_hidden_layers']}",
+                    so,
+                ],
+            )
+            filename = f"em_{name}.onnx"
+
+            memory_session = (
+                onnxscript.tools.memory_peak.start_spying_on(cuda=kwargs["device"] == "cuda")
+                if kwargs["memory_peak"]
+                else None
+            )
+            print(f"[export_model] start memory peak monitoring {memory_session}")
+            proto = onnxscript.tools.benchmark.common_export(
+                model=model,
+                inputs=example_inputs[0],
+                exporter=kwargs["exporter"],
+                target_opset=kwargs["target_opset"],
+                folder=kwargs["dump_folder"],
+                filename=filename,
+                dynamic_shapes=dynamic_shapes if kwargs["dynamic"] else None,
+                optimization=kwargs["optimization"],
+                verbose=kwargs["verbose"],
+                stats=conversion,
+            )
+            print(f"[export_model] export to onnx done in {time.perf_counter() - begin}")
+            if memory_session is not None:
+                memory_results = memory_session.stop()
+                print(f"[export_model] ends memory monitoring {memory_results}")
+                memory_stats = onnxscript.tools.memory_peak.flatten(
+                    memory_results, prefix="memory_"
+                )
+            else:
+                memory_stats = {}
+
+            result = onnxscript.tools.benchmark.run_onnx_inference(
+                proto,
+                example_inputs,
+                warmup=kwargs["warmup"],
+                repeat=kwargs["repeat"],
+                verbose=kwargs["verbose"],
+                ort_optimize=kwargs["ort_optimize"],
+                torch_model=model,
+            )
+
+        print("[export_model] end")
+        print("------------------------------")
+        for k, v in sorted(kwargs.items()):
+            print(f":{k},{v};")
+        for k, v in sorted(conversion.items()):
+            print(f":{k},{v};")
+        if memory_stats:
+            for k, v in memory_stats.items():
+                print(f":{k},{v};")
+        for k, v in sorted(result.items()):
+            print(f":{k},{v};")
+
+
+if __name__ == "__main__":
+    main()

onnxscript/tools/benchmark/export_model_batch.py

@@ -0,0 +1,146 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# pylint: disable=import-outside-toplevel
+
+from __future__ import annotations
+
+import pprint
+import textwrap
+from typing import Any
+
+import onnxscript.tools.benchmark
+
+
+def main(args: list[str] | None = None):
+    kwargs: dict[str, Any] = onnxscript.tools.benchmark.get_parsed_args(
+        "export_model",
+        description=textwrap.dedent(
+            """Measures the inference time for a particular model.
+            It runs export_model to compare several optimization settings.
+
+            Example::
+
+                python -m onnxscript.tools.benchmark.export_model_batch --model phi --device cuda --config medium --num_hidden_layers=1 --dtype=float32 --dynamic=0 --verbose=1
+            """
+        ),
+        repeat=(10, "number of inferences to measure"),
+        warmup=(5, "number of inferences to warm"),
+        model=("phi", "model to measure, llama, mistral, phi, ..."),
+        device=("cpu", "'cpu' or 'cuda'"),
+        target_opset=(18, "opset to convert into, use with backend=custom"),
+        config=("small", "default, medium, or small to test"),
+        verbose=(0, "verbosity"),
+        dtype=("default", "cast the model and the inputs into this type"),
+        dynamic=(0, "use dynamic shapes"),
+        num_hidden_layers=(1, "number of hidden layers"),
+        with_mask=(1, "with or without mask, dynamo may fail with a mask"),
+        implementation=("eager", "eager or sdpa"),
+        new_args=args,
+    )
+
+    print("-------------------")
+    print("[export_model]")
+    pprint.pprint(kwargs)
+    print("-------------------")
+
+    import pandas
+
+    try:
+        import openpyxl
+    except ImportError:
+        openpyxl = None
+
+    from onnxscript.tools.benchmark.benchmark_helpers import (
+        BenchmarkError,
+        run_benchmark,
+    )
+
+    script_name = "onnxscript.tools.benchmark.export_model"
+
+    configs: list[dict[str, Any]] = [
+        dict(exporter="eager"),
+        dict(ort_optimize=1, exporter="script"),
+        dict(ort_optimize=1, optimization="optimize/rewrite/inline", exporter="script"),
+        dict(ort_optimize=0, optimization="optimize/rewrite/inline", exporter="script"),
+        dict(ort_optimize=1, optimization="", exporter="dynamo"),
+        dict(ort_optimize=1, optimization="optimize/rewrite/inline", exporter="dynamo"),
+        dict(ort_optimize=0, optimization="optimize/rewrite/inline", exporter="dynamo"),
+    ]
+    common_kwargs: dict[str, Any] = kwargs.copy()
+    common_kwargs["verbose"] = max(common_kwargs["verbose"] - 1, 0)
+    for c in configs:
+        c.update(common_kwargs)
+
+    if kwargs["verbose"]:
+        for i, cf in enumerate(configs):
+            print(f"[export_common_batch] config {i + 1}: {cf}")
+
+    ################################
+    # Running configuration.
+
+    try:
+        data = run_benchmark(
+            script_name,
+            configs,
+            verbose=kwargs["verbose"],
+            stop_if_exception=False,
+        )
+        data_collected = True
+    except BenchmarkError as e:
+        if kwargs["verbose"]:
+            print(e)
+        data_collected = False
+
+    prefix = "_".join(
+        [
+            "emb_",
+            kwargs["model"],
+            "dynamic" if kwargs["dynamic"] else "static",
+            kwargs["dtype"].replace("float", "fp"),
+            kwargs["device"],
+            kwargs["config"],
+            f"h{kwargs['num_hidden_layers']}",
+        ],
+    )
+
+    if data_collected:
+        df = pandas.DataFrame(data)
+        df = df.drop(["OUTPUT", "ERROR"], axis=1)
+        df["repeat_time"] = df["repeat_time"].astype(float)
+        df_eager = df[(df["implementation"] == "eager") & (df["exporter"] == "eager")][
+            "repeat_time"
+        ].dropna()
+        if df_eager.shape[0] > 0:
+            min_eager = df_eager.min()
+            df["increase"] = df["repeat_time"] / min_eager - 1
+        filename = f"{prefix}_with_cmd.csv"
+        df.to_csv(filename, index=False)
+
+        df = df.drop(["CMD"], axis=1)
+        filename = f"{prefix}.csv"
+        df.to_csv(filename, index=False)
+        df = pandas.read_csv(filename)  # to cast type
+        print(df)
+
+        # summary
+        cs = [
+            c
+            for c in ["exporter", "optimization", "warmup_time", "repeat_time", "increase"]
+            if c in df.columns
+        ]
+        dfs = df[cs]
+        if openpyxl:
+            filename = f"{prefix}_summary.xlsx"
+            dfs.to_excel(filename, index=False)
+        filename = f"{prefix}_summary.csv"
+        dfs.to_csv(filename, index=False)
+        print(dfs)
+
+    ########################
+    # First lines.
+
+    print(df.head(2).T)
+
+
+if __name__ == "__main__":
+    main()