Repository: pytorch
Issue: Come with a better strategy for TensorArg (error reporting)
Labels: module: performance,module: internals,triaged

TensorArg (in `aten/src/ATen/TensorUtils.h`) is a class I came up with when I was porting CuDNN convolutions from Python to C++. The original goals were as follows:

1. Create a set of utility functions to conveniently perform input checks (e.g., does it have the correct size; is it on the correct GPU, etc.)
2. Give good error messages when these tests fail (e.g., give the name and the argument number of the argument that failed the test)
3. Avoid copy-pasting error checking logic.

Goals (2) and (3) are in tension with each other, especially for convolutions. Consider the following example: `convolution_forward`, and `transposed_convolution_backward`. Algorithmically, these two functions are equivalent (and indeed, both functions backend to the same underlying function), but from a UI perspective the functions are very different (e.g., arguments names differ). If you want to not copy-paste error checking logic, you need to put this logic in the underlying backend function, but if you want good error messages you have to track extra information on a per-Tensor basis so that you can accurately describe what the name/argument position of an argument is.

TensorArg solves this problem by rewriting functions to pass `TensorArg` struct instead of `Tensor`, where `TensorArg` is augmented with extra information to track the name and position of an argument. Checking functions which make use of `TensorArg` can pull this information out directly when giving an error message.

However, TensorArg falls short of some other goals which I subsequently realized were quite important:

1. **Speed.** Boxing tensors into `TensorArg` structs, while fairly cheap, resulted in a perceptible, cumulative performance slowdown. Ugh! Some of the cost may be due to additional dynamic allocations. https://github.com/pytorch/pytorch/issues/3958 is the tracking issue.
2. **Brevity.** Suppose that I have a native function `f` which backends to `f_out`, which is also native. I cannot rewrite `f_out` to be `TensorArg`, because the calling convention for native functions requires `Tensor` arguments. So I have to write another backend function `_f_out`, and then have `f_out` and `f` call into it.
3. **Compositionality.** Essentially, the `TensorArg` strategy means we need a "shadow" hierarchy of TensorArg-consuming functions, which you call directly if you want to preserve error reporting. This means that if you ever need to go back in the direct Tensor interface, you cannot preserve error information.

Given this problem, it seems unlikely that we can recommend TensorArg for wider use in the ATen codebase, which is a shame because it is still very difficult to do good error reporting when writing code in ATen.

CC @goldsborough 

----

Proposal: Thread-local state error reporting information

Let's get rid of `TensorArg` so that we can achieve brevity and compositionality. In this case, how can we tell what the name of an argument is? We have thread-local error reporting information, which maintains a mapping (just an association list) from `TensorImpl*` to error information. Upon a failure, we perform a lookup in the mapping to resolve a `Tensor` into an actual argument name (or not, if we failed on an intermediate result; that's a bad outcome and means you didn't add enough error checking). The mapping is similar to how we handle input arguments in tracing.

One limitation of this approach is that it deals poorly when exactly the same tensor is passed as two arguments, e.g., `add(x, x)`. In this case, we cannot tell purely from the identity of the TensorImpl whether or not this was the first argument or the second argument. 

----

Proposal: Exception handler transformers

Let's get rid of `TensorArg` so that we can achieve brevity and compositionality. In this case, how can we tell what the name of an argument is? (Does this sound familiar? It's because it's the same as above). We don't: instead, we report whatever the local information of a function is and raise the exception.

At each function call site, we instrument an exception handler which is responsible for "rewriting" the error message so it makes sense in its locality. So for example, if we have:

```
void foo(Tensor x, Tensor y) {
  bar(y, x);
}
```

We must somehow make this equivalent to the following:

```
void foo(Tensor x, Tensor y) {
  try {
    bar(y, x);
  except (AtenError& e) {
    e.calling_pattern({{"x", 1}, {"y", 0}});
    throw;
  }
}
```

or something equivalent, where we can reconstruct the original argument pattern by traversing the exception handlers.

Limitation: It's not altogether clear how to automatically generate these wrappers without ... something like `TensorArg`, because the call pattern of `bar` matters, and that's not easily accessible from C++.

cc @ezyang @bhosmer @smessmer @ljk53 @bdhirsh @ailzhang @VitalyFedyunin @ngimel