Support in-place operations on Variable views

[colesbury]

Currently, we don't allow in-place operations on views if there are any other live Variables sharing the same data. We should allow in-place operations and compute the gradient correctly. There are a few changes necessary to support this.

Viewing operations include view, narrow, select, expand, [un]squeeze, and unfold. By "a view" I mean a Variable that is a view on another Variable.

1. Add a field "base" to Variable. Every view has a pointer to a single base Variable. (The base is never a view)
2. In-place operations on views change the grad_fn of the base, not of the view.
3. The grad_fn on a view may become stale. So views also store an expected_version

In-place operations on views

Every in-place operation on a view inserts three Functions to base.grad_fn. For example:

x = base[2]
x.mul_(y)

becomes:

Stride(MulBackward(Unstride(<old base.grad_fn>), <y.grad_fn>))

Assume, for now, that base is contiguous. The Stride op takes the grad_output, makes it contiguous, and applies the stride, sizes, and storageOffset of the view (x). This performs the identical viewing operation on grad_output that was performed on base to get x. For the above example, this is equivalent to:

grad_input = grad_output[2]

Unstride performs the reverse operation. It creates a grad_input with the same shape as base, applies the striding of the view x to grad_input, and copies grad_output into that view. For the above example, this is equivalent to:

grad_input = zeros_like(base)
grad_input[2] = grad_output

Stride and Unstride are both implemented in terms of a differentiable restride operation that generalizes all of our viewing operations:

Tensor restride(Tensor self, IntList sizes, IntList strides, int offset);

Stale grad_fn on views

An in-place operation on a view makes the grad_fn of all other views stale. The compute_next_functions needs to be updated to handle views: it should re-create the Function for a view if it's stale before returning it.

Non-contiguous base Variables

If base is not contiguous, then we also need the strides/sizes of the base in restride:

Tensor restride(Tensor self, IntList sizes, IntList strides, IntList base_sizes, IntList base_strides, int offset) {
  Tensor copy = self.type().tensor(base_sizes, base_strides);
  copy.copy_(self);
  return restride(copy, sizes, strides, offset);
}

[colesbury]

cc @apaszke, @soumith, @ezyang, @gchanan

This is important for fully implementing indexing on Variables. It's very difficult to combine basic slicing and advanced indexing in __setitem__ without this, and even basic slicing in __setitem__ becomes more complicated.

[soumith]

this seems like a pretty reasonable proposal. i think you should go for it.

[apaszke]

A few comments:

1. Having stale state is terrible. It's not only compute_next_functions that will have to be updated. We'll have to go over a lot of code and make sure we never actually access state until we know it's up to date. I can see a lot of easy to miss bugs created because of this. However, I can see that this might be the only option to have reasonable perf in this setting...
2. The cost of backward for in-place operations on views might become unbearable. People will think that their code will be blazing fast because they're only working in-place, but this will mean a ton of Unstrides in the backward, and each has to allocate full-sized tensors and zero them. That's the problem that makes backward of pack_padded_sequence slow. It has a lot of narrows inside, and just writing out backward manually would make it a loooot faster. In general I think it's possible to overcome this problem, but this requires structuring the backward graph a bit differently (it would have to be a single atomic autograd op, that does slicing, performs another call, and copies the result back into the original thing. add clones where necessary)
3. You're talking about modifications to views. What about modifications to base Variable? Does it just make all the views stale and that's it?
4. I don't understand the part about non-contiguous base Variables. If you were to create a new non-contiguous tensor, you would still need to allocate a large contiguous chunk of memory for it, but the copy_ would only initialize the referenced parts, leaving garbage between them. Not sure if it would ever be used again. If yes, it's an obvious bug, otherwise it might be possible to avoid allocating this huge chunk and compress the operation

[apaszke]

Also sorry if my comment sounds too negative. Thanks a lot for writing up this proposal. It looks very reasonable and I'm looking forward to adding this to autograd. However, it might require some polishing.

[colesbury]

One way to think about this is that views don't store grad_fn. When they're used in a computation, we dynamically create a view Function that does the inverse view based on the sizes and strides of the base and the view. But for performance reasons, it's convenient to cache the grad_fn along with a saved version like in SavedVariable to avoid unnecessarily creating duplicate functions when a view is used multiple times.

There's a single accessor Variable.grad_fn() that everything goes through. I'm going to make VariableImpl.grad_fn private or at least rename it to _grad_fn.

I don't think this is specific to in-place operations. Narrowing operations are free in the forward-pass, but expensive in the backward pass because you have to pad the gradient with zeros.

In-place ops on views have the cost of the basic backwards op (e.g. MulBackward) plus the view backwards (e.g. NarrowBackward).

(If the base is not contiguous, the cost might be ~2x that)

It'd be good to address this in general, perhaps with sparse gradients.

Yes, it makes the grad_fn on the views stale implicitly by increment the shared version counter. To be clear, the graph of Function nodes is still immutable, only Variable.grad_fn can change. And the base Variable does not have a reference to its views. The reference is one-way from views to base.

This is a bit tricky to explain. The viewing operation is completely described by the sizes, strides, and relative storageOffset of the base and view variables. We want to perform an equivalent viewing operation on the incoming grad as the one that created the view. (Note: grad is the same size as base, but not necessarily the same strides).

If the base and grad are both contiguous, then we apply the view.sizes(), view.strides(), and relative storage offset to the grad. (This is basically free).

If it's not the case, then we can do the restride op from above. You're right that if base has "holes" then it'll create more memory then necessary. These holes will never be accessed, because the view is guaranteed to only contain a subset (inclusive) of the elements in base.

There's an optimization which avoids creating the wholes in most cases: divide the strides by the GCD.

[ezyang]

For the record, I've talked to @colesbury about this in person and I'm really convinced about the technical feasibility of the idea.

In general I think it's possible to overcome this problem, but this requires structuring the backward graph a bit differently (it would have to be a single atomic autograd op, that does slicing, performs another call, and copies the result back into the original thing. add clones where necessary)

Sounds like a job for the optimizer :) (Though, we're losing all of our inplace information already in the JIT. Win some, lose some.)