[MetaSchedule] TorchBench tuning script: add option to disallow operators in sub graph

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,
 )
 
 

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")
+
+    tensor_variable_cls = torchdynamo.variables.tensor.TensorVariable
+    old_call_method = tensor_variable_cls.call_method
+
+    @functools.wraps(old_call_method)
+    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, translator, name, args, kwargs)
+
+    tensor_variable_cls.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