add __torch_function__ API override mechanism

benchmarks/README.md

@@ -27,3 +27,63 @@ Please refer to each subfolder to discover each benchmark suite
 
 * [Fast RNNs benchmarks](fastrnns/README.md)
 
+
+
+## PyTorch ASV benchmarks
+
+Benchmarking PyTorch with Airspeed Velocity.
+
+
+Usage
+-----
+
+Airspeed Velocity manages building and Python virtualenvs or conda envs by
+itself, unless told otherwise (e.g. with `--python=same`).
+To run the benchmarks, you do not need to install a development version of
+PyTorch to your current Python environment.
+TODO: check that the isolated build feature works, so far just used
+`--python=same`.
+
+Run a benchmark against currently installed PyTorch version (don't
+record the result)::
+
+    asv run --python=same
+
+Compare change in benchmark results to another version::
+
+    TODO
+
+Run ASV commands (record results and generate HTML)::
+
+    cd benchmarks
+    asv run --skip-existing-commits --steps 10 ALL
+    asv publish
+    asv preview
+
+More on how to use ``asv`` can be found in the
+[ASV documentation](https://asv.readthedocs.io).
+Command-line help is available as usual via `asv --help` and `asv run --help`.
+
+
+
+Writing benchmarks
+------------------
+
+See the ASV documentation for basics on how to write benchmarks.
+
+Some things to consider:
+
+- The benchmark suite should be importable with any PyTorch version.
+
+- The benchmark parameters etc. should not depend on which PyTorch version
+  is installed.
+
+- Try to keep the runtime of the benchmark reasonable.
+
+- Prefer ASV's `time_` methods for benchmarking times rather than cooking up
+  time measurements via `time.clock`, even if it requires some juggling when
+  writing the benchmark.
+
+- Preparing input tensors etc. should generally be put in the `setup` method
+  rather than the `time_` methods, to avoid counting preparation time together
+  with the time of the benchmarked operation.

benchmarks/asv.conf.json

@@ -0,0 +1,85 @@
+{
+    // The version of the config file format.  Do not change, unless
+    // you know what you are doing.
+    "version": 1,
+
+    // The name of the project being benchmarked
+    "project": "PyTorch",
+
+    // The project's homepage
+    "project_url": "https://pytorch.org/",
+
+    // The URL or local path of the source code repository for the
+    // project being benchmarked
+    "repo": "..",
+
+    // List of branches to benchmark. If not provided, defaults to "master"
+    // (for git) or "tip" (for mercurial).
+    "branches": ["master"],
+
+    // The DVCS being used.  If not set, it will be automatically
+    // determined from "repo" by looking at the protocol in the URL
+    // (if remote), or by looking for special directories, such as
+    // ".git" (if local).
+    "dvcs": "git",
+
+    // The tool to use to create environments.  May be "conda",
+    // "virtualenv" or other value depending on the plugins in use.
+    // If missing or the empty string, the tool will be automatically
+    // determined by looking for tools on the PATH environment
+    // variable.
+    "environment_type": "conda",
+
+    // the base URL to show a commit for the project.
+    "show_commit_url": "https://github.com/pytorch/pytorch/commit/",
+
+    // The Pythons you'd like to test against.  If not provided, defaults
+    // to the current version of Python used to run `asv`.
+    "pythons": ["3.6"],
+
+    // The matrix of dependencies to test.  Each key is the name of a
+    // package (in PyPI) and the values are version numbers.  An empty
+    // list indicates to just test against the default (latest)
+    // version.
+    "matrix": {
+        "six": [],
+    },
+
+    // The directory (relative to the current directory) that benchmarks are
+    // stored in.  If not provided, defaults to "benchmarks"
+    "benchmark_dir": "benchmarks",
+
+    // The directory (relative to the current directory) to cache the Python
+    // environments in.  If not provided, defaults to "env"
+    "env_dir": "env",
+
+
+    // The directory (relative to the current directory) that raw benchmark
+    // results are stored in.  If not provided, defaults to "results".
+    "results_dir": "results",
+
+    // The directory (relative to the current directory) that the html tree
+    // should be written to.  If not provided, defaults to "html".
+    "html_dir": "html",
+
+    // The number of characters to retain in the commit hashes.
+    // "hash_length": 8,
+
+    // `asv` will cache wheels of the recent builds in each
+    // environment, making them faster to install next time.  This is
+    // number of builds to keep, per environment.
+    "build_cache_size": 2,
+
+    // The commits after which the regression search in `asv publish`
+    // should start looking for regressions. Dictionary whose keys are
+    // regexps matching to benchmark names, and values corresponding to
+    // the commit (exclusive) after which to start looking for
+    // regressions.  The default is to start from the first commit
+    // with results. If the commit is `null`, regression detection is
+    // skipped for the matching benchmark.
+    //
+    // "regressions_first_commits": {
+    //    "some_benchmark": "352cdf",  // Consider regressions only after this commit
+    //    "another_benchmark": null,   // Skip regression detection altogether
+    // }
+}

benchmarks/benchmarks/__init__.py

@@ -0,0 +1 @@
+from __future__ import absolute_import, division, print_function

benchmarks/benchmarks/bench_overrides.py

@@ -0,0 +1,58 @@
+from .common import Benchmark
+
+import torch
+from torch import Tensor
+
+class DuckTensor(object):
+    def __torch_function__(self, func, args=(), kwargs=None):
+        pass
+
+HANDLED_FUNCTIONS = {}
+
+def implements(torch_function):
+    "Register an implementation of a torch function for a Tensor-like object."
+    def decorator(func):
+        HANDLED_FUNCTIONS[torch_function.__name__] = func
+        return func
+    return decorator
+
+class SubTensor(Tensor):
+    def __torch_function__(self, func, args=(), kwargs=None):
+        if(kwargs is None):
+            kwargs = {}
+
+        if func not in HANDLED_FUNCTIONS:
+            return NotImplemented
+        # Note: this allows subclasses that don't override
+        # __torch_function__ to handle DiagonalTensor objects.
+        return HANDLED_FUNCTIONS[func](*args, **kwargs)
+
+@implements(torch.add)
+def add(mat1, mat2):
+    "Implementation of torch.mm for DiagonalTensor objects"
+    return 0
+
+@implements(torch.mm)
+def mm(mat1, mat2):
+    "Implementation of torch.mm for DiagonalTensor objects"
+    return 1
+
+class TorchFunction(Benchmark):
+
+    def setup(self):
+        self.t1 = torch.ones(2, 2, dtype=torch.float32)
+        self.t2 = torch.zeros(2, 2, dtype=torch.float32)
+        self.t3 = SubTensor([[1, 1], [1, 1.]])
+        self.t4 = SubTensor([[0, 0], [0, 0.]])
+
+    def time_add(self):
+        torch.add(self.t1, self.t2)
+
+    def time_matmul(self):
+        torch.mm(self.t1, self.t2)
+
+    def time_subtensor_add(self):
+        torch.add(self.t3, self.t4)
+
+    def time_subtensor_multipy(self):
+        torch.mm(self.t3, self.t4)

benchmarks/benchmarks/common.py

@@ -0,0 +1,18 @@
+from __future__ import absolute_import, division, print_function
+
+import random
+import torch
+import numpy
+
+
+# Better not to use random numbers for benchmarks, but just in case,
+# seed everything
+random.seed(123123459)
+torch.manual_seed(123123459)
+numpy.random.seed(123123459)
+
+
+class Benchmark(object):
+    # asv auto-selects number of iterations, so benchmarks runs in between
+    # goal_time/10 and goal_time (seconds)
+    goal_time = 0.25

test/run_test.py

@@ -60,6 +60,7 @@
     'type_promotion',
     'jit_disabled',
     'function_schema',
+    'overrides',
 ]
 
 # skip < 3.3 because mock is added in 3.3 and is used in rpc_fork and rpc_spawn