From 281781caab0228f909bf22c6a23567433b68daaf Mon Sep 17 00:00:00 2001 From: Lite Ye Date: Sun, 20 Nov 2022 00:31:10 -0500 Subject: [PATCH 1/2] Add option to disallow operator in TorchDynamo Add missing type annotation --- .../meta_schedule/testing/torchbench/run.py | 14 +++++ .../meta_schedule/testing/torchbench/utils.py | 63 ++++++++++++++++++- 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/python/tvm/meta_schedule/testing/torchbench/run.py b/python/tvm/meta_schedule/testing/torchbench/run.py index 65e1a1a59f3e..55ac02b0b743 100644 --- a/python/tvm/meta_schedule/testing/torchbench/run.py +++ b/python/tvm/meta_schedule/testing/torchbench/run.py @@ -110,6 +110,7 @@ from tvm._ffi import get_global_func from tvm.contrib.graph_executor import GraphModule from tvm.meta_schedule.testing.torchbench.utils import ( + DisallowedOperator, load_torchdynamo_benchmark_runner, same, timed, @@ -196,6 +197,12 @@ def parse_args(): default=5, help="The number of rounds to warmup before starting to measure the performance.", ) + args.add_argument( + "--disallowed-op", + type=str, + default="all", + help=DisallowedOperator.__doc__, + ) # Model selection args.add_argument( @@ -313,6 +320,12 @@ def parse_args(): parsed = args.parse_args() + if parsed.disallowed_op == "all": + disallowed_op = set(DisallowedOperator) + else: + disallowed_op = {DisallowedOperator(v) for v in parsed.disallowed_op.split(",")} + parsed.disallowed_op = disallowed_op + # Trim all args, otherwise it confuses the arg parser of timm_efficientdet sys.argv = sys.argv[:1] @@ -335,6 +348,7 @@ def parse_args(): IS_CUDA, cosine_similarity=ARGS.result_metric == ResultComparisonMetric.COSINE, float32=ARGS.float32, + disallowed_operators=ARGS.disallowed_op, ) diff --git a/python/tvm/meta_schedule/testing/torchbench/utils.py b/python/tvm/meta_schedule/testing/torchbench/utils.py index 8bd022a9cb18..d2e9555cab55 100644 --- a/python/tvm/meta_schedule/testing/torchbench/utils.py +++ b/python/tvm/meta_schedule/testing/torchbench/utils.py @@ -19,13 +19,33 @@ from TorchDynamo. """ +import functools import os import sys from dataclasses import dataclass +from enum import Enum +from typing import Set import torch # type: ignore +class DisallowedOperator(Enum): + """ + The operators to disallow in the fx graph produced by TorchDynamo. + This is to workaround the limitation in TVM's PyTorch frontend. + + - inplace_copy: aten::copy_ as inplace assign A[...] = ..., or method call A.copy_(...) + - einsum: torch.functional.einsum + - multihead_attention: torch.nn.MultiheadAttention + - as_stride: Tensor.as_stride + """ + + INPLACE_COPY = "inplace_copy" + EINSUM = "einsum" + MULTIHEAD_ATTENTION = "multihead_attention" + AS_STRIDE = "as_stride" + + def find_torchdynamo() -> str: """ Find the directory of TorchDynamo repo. @@ -57,14 +77,52 @@ def find_torchdynamo() -> str: sys.path.append(f"{DYNAMO_DIR}/benchmarks") # pylint: disable=wrong-import-position, unused-import +import torchdynamo # type: ignore from benchmarks.common import same, timed # type: ignore from torchbench import TorchBenchmarkRunner # type: ignore # pylint: disable=wrong-import-position, unused-import +def _disallow_operators(disallowed_ops: Set[DisallowedOperator]): + """ + Disallow certain operators in the fx graph produced by TorchDynamo. + There are two ways to disallow operator in TorchDynamo, + 1. Use the disallow_in_graph API, which only applies to free function call. + 2. Patch the TensorVariable class, which applies to method call on torch.Tensor. + """ + disallowed_tensor_methods: Set[str] = set() + + if DisallowedOperator.INPLACE_COPY in disallowed_ops: + torchdynamo.disallow_in_graph(torch.Tensor.copy_) + disallowed_tensor_methods.update({"copy_", "__setitem__"}) + + if DisallowedOperator.EINSUM in disallowed_ops: + torchdynamo.disallow_in_graph(torch.functional.einsum) + + if DisallowedOperator.MULTIHEAD_ATTENTION in disallowed_ops: + torchdynamo.disallow_in_graph(torch.nn.MultiheadAttention) + + if DisallowedOperator.AS_STRIDE in disallowed_ops: + disallowed_tensor_methods.add("as_stride") + + TensorVariable = torchdynamo.variables.tensor.TensorVariable + old_call_method = TensorVariable.call_method + + @functools.wraps(old_call_method) + def call_method(self, tx, name, args, kwargs): + if name in disallowed_tensor_methods: + raise torchdynamo.exc.Unsupported(f"Tensor.{name} not supported by TVM.") + return old_call_method(self, tx, name, args, kwargs) + + TensorVariable.call_method = call_method + + def load_torchdynamo_benchmark_runner( - is_cuda: bool, cosine_similarity: bool = False, float32: bool = False + is_cuda: bool, + cosine_similarity: bool = False, + float32: bool = False, + disallowed_operators: Set[DisallowedOperator] = None, ) -> TorchBenchmarkRunner: """ Load the benchmark runner from TorchDynamo. @@ -94,6 +152,9 @@ class RunnerArgs: runner.args = args runner.model_iter_fn = runner.forward_pass + if disallowed_operators: + _disallow_operators(disallowed_operators) + if is_cuda: # pylint: disable=import-outside-toplevel import benchmarks.common # type: ignore From 51803650106551d4eee0118caa8d54eb26b2886e Mon Sep 17 00:00:00 2001 From: Lite Ye Date: Mon, 21 Nov 2022 19:51:33 -0500 Subject: [PATCH 2/2] Fix lint --- python/tvm/meta_schedule/testing/torchbench/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/tvm/meta_schedule/testing/torchbench/utils.py b/python/tvm/meta_schedule/testing/torchbench/utils.py index d2e9555cab55..af81318d0006 100644 --- a/python/tvm/meta_schedule/testing/torchbench/utils.py +++ b/python/tvm/meta_schedule/testing/torchbench/utils.py @@ -106,16 +106,16 @@ def _disallow_operators(disallowed_ops: Set[DisallowedOperator]): if DisallowedOperator.AS_STRIDE in disallowed_ops: disallowed_tensor_methods.add("as_stride") - TensorVariable = torchdynamo.variables.tensor.TensorVariable - old_call_method = TensorVariable.call_method + tensor_variable_cls = torchdynamo.variables.tensor.TensorVariable + old_call_method = tensor_variable_cls.call_method @functools.wraps(old_call_method) - def call_method(self, tx, name, args, kwargs): + def call_method(self, translator, name, args, kwargs): if name in disallowed_tensor_methods: raise torchdynamo.exc.Unsupported(f"Tensor.{name} not supported by TVM.") - return old_call_method(self, tx, name, args, kwargs) + return old_call_method(self, translator, name, args, kwargs) - TensorVariable.call_method = call_method + tensor_variable_cls.call_method = call_method def load_torchdynamo_benchmark_runner(