Repository: pytorch
Issue: Reduce degrees of freedom in strides so that they unambiguously specify layout permutation
Labels: triaged,enhancement

In the memory layout work in https://github.com/pytorch/pytorch/issues/19092 we have a problem where strides undetermine the layout of a tensor when size = 1 or 0. In particular, if I have a tensor with size `(1, 1, 1, 1)`, it is impossible to tell (from the strides) if it is an NHWC or NCHW tensor (or indeed, any other permutation of the layout). This is because when size = 1 or 0, strides are undetermined, The strides are undetermined because the indexing formula `stride * index` always has `index = 0`, so it doesn't matter what you pick for strides. So you may very well find that your tensor has strides `(1, 1, 1, 1)` which doesn't tell you anything about its memory layout.

We argue that we should fully specify the strides of this tensor from a set of sentinel strides, resolving the ambiguity of whether or not you have an NHWC or NCHW tensor. In the example above, your strides would either be `(D0, D1, D2, D3)` (NCHW) or `(D0, D3, D1, D2)` (a NHWC physical, logically viewed as NCHW via the permutation (0, 3, 1, 2); notice that the layout permutation is **reverse** of the view permutation). When strides are uniquely determined, the layout also is uniquely determined, solving the information loss problem.

## The scheme

Intuitively, the layout permutation of a tensor is the `dims` that when run on `permute(*dims)` that would have taken the physical layout of a tensor to its logical layout. So for example, a NHWC tensor viewed as a NCHW tensor has the layout permutation `(0, 3, 1, 2)`, because `nhwc.permute(0, 3, 1, 2)` gives you channels-first.  The idea is to add enough information to strides so that we can always recover the layout permutation from it.

We reserve a set of distinguished stride values `D0`, `D1`, `D2` ... which compare unequal with natural strides. The mnemonic for these distinguished stride values is that `D0` is `stride(0)` of a contiguous tensor, if the stride would otherwise be undetermined (e.g., its corresponding size is 1, or the tensor is empty). Let `undetermined_stride(0) = D0` and so forth. We define a total order on undetermined strides `D0 > D1 > D2 > ...` (the order is in reverse of the natural order on the indices, for symmetry with the fact that with a contiguous tensor, the largest strides are associated with the lowest dimension). 

The strides for a contiguous tensor with the identity layout permutation (e.g., NCHW) of size `sizes` is computed in the following way:

```
def identity_permutation_contiguous_strides(sizes):
  stride = 1
  dim = len(sizes)
  strides = [0] * dim
  is_empty = any(s == 0 for s in sizes);

  for i in range(dim - 1, -1, -1): # dim - 1 to 0 inclusive
    if sizes[i] == 1 or is_empty:
      strides[i] = undetermined_stride(i)
    else:
      strides[i] = stride
    stride *= size[i]
```

Examples:
* sizes = (2, 2, 2), strides = (4, 2, 1)
* sizes = (2, 1, 2), strides = (2, D1, 1)
* sizes = (2, 0, 2), strides = (D0, D1, D2)

Invariants:
* Every stride in strides is unique.
* `is_undetermined_stride(stride[i])` if and only if `size[i] == 1` or `numel == 0`
* A size of a tensor *uniquely* determines what strides it has (but not vice versa; consider an empty tensor)

With these undetermined strides, we can now unambiguously identify if a tensor is NCHW contiguous or NHWC contiguous, by doing a direct equality comparison with the appropriately permuted strides of the contiguous strides for that size.

```
def is_nchw_contiguous(t):
  return t.dim() == 4 and identity_permutation_contiguous_strides(t.size()) == t.stride()

def permute(xs, permutation):
  return (xs[i] for i in permutation)

def is_nhwc_contiguous(t):
  return t.dim() == 4 and permute(identity_permutation_contiguous_strides(t.size()), (0, 3, 1, 2)) == t.stride()
```

With more work, you can also determine the permutation based on the strides, even if they're not contiguous.

```
def permutation_of(strides):
    # Invariant:
    #     if strides[i] == undetermined_stride(j),
    #     then permutation[i] = j
    #
    # Invariant:
    #     if permutation[i] > permutation[j]
    #         and not is_undetermined_stride(strides[i])
    #         and not is_undetermined_stride(strides[j]),
    #     then strides[i] < strides[j]
    #     (note that the comparison here is inverted
    #     since strides are ordered in decreasing order)

    # Partition strides into determined/undetermined strides
    determined_stride_indices = []
    undetermined_stride_indices = []
    for i, s in enumerate(strides):
        if is_undetermined_stride(s):
            undetermined_stride_indices.append(i)
        else:
            determined_stride_indices.append(i)

    # Sort
    determined_stride_indices.sort(key=lambda i: strides[i], reverse=True)
    undetermined_stride_indices.sort(key=lambda i: strides[i], reverse=True)

    # Reassemble permutation
    permutation = [-1] * len(strides)
    permutation_i = 0
    determined_i = 0
    undetermined_i = 0
    for i in range(0, len(strides)):
        if undetermined_i < len(undetermined_stride_indices):
            j = undetermined_stride_indices[undetermined_i]
            if strides[j] == undetermined_stride(i):
                permutation[j] = permutation_i
                permutation_i += 1
                undetermined_i += 1
                continue
        j = determined_stride_indices[determined_i]
        permutation[j] = permutation_i
        permutation_i += 1
        determined_i += 1

    return permutation
```

## Discussion

* Zero strides are canonically placed at the end of the dimension order. This is a little arbitrary; they could be put anywhere in the order, but we put them at the end so we can use the natural ordering between 0 < 1
 * Layout permutations account for “weird strides”, e.g., rolling windows. The tie-break when strides are equal prefers ordering larger windows before smaller windows; this is also an arbitrary choice.
 * The result of contiguous() operation should result in strides with the identity permutation. This diverges with our behavior in some cases today `torch.randn(2,3,1,4).permute(0,1,3,2).contiguous()`; we should change this behavior. (h/t Natalia for this example)

## Propagation rules

* Layout permutation propagation is largely determined by strides. If an operation produces a tensor whose strides uniquely determine a layout permutation, we don’t have to give any further rules for layout permutation propagation.
* Operators should behave *consistently* with respect to permutations: in other words, if an operator produces a layout permutation given some input layout permutation on a specific size/stride, it should produce that layout permutation for *all* sizes and strides consistent with the input layout permutation.
    * Example: Interpolate on a 2x2 tensor preserves permutation in the output. Thus, it should preserve permutation for 1x1 tensors.
    * In general, operations either preserve permutation, or explicitly change the permutation (transposes).

## Physical representation

We specified the existence of undetermined strides; how should we actually represent them? We observe:

* Value of stride doesn’t matter when size = 1, or when you have an empty tensor
* We represent strides as 64-bit integers, but we never actually use all of the bits in a 64-bit integer (since you can’t possibly have that much memory)

Proposal: instead of storing strides, we store strides in a packed format:

```

+-------------+-----------------------------------+
| tag (2 bit) | stride (62 bit)                   |
+-------------+-----------------------------------+

tag meanings:
  00 - regular stride
  11 - reserved for negative strides
  01 - reserved
  10 - permutation sentinel 
 
```

We require the stride for every dimension with size = 1 to be assigned a unique permutation sentinel (but we do not prescribe any other requirements on the sentinel). So for example, for a tensor of size (3, 3, 1, 1), the strides may be (3, 1, Sentinel 0, Sentinel 1). If a tensor is empty, we require every stride to be a unique permutation sentinel. With sentinels, the induced layout permutation is guaranteed to be unique and can be derived directly from the strides.

As a further optimization, well known permutations (such as the identity permutation, or NHWC) could be recorded in a separate field for quick checking. Memory format names like “channels_first” or “channels_last” correspond to families of permutations over arbitrary dimensionality. (NCHW would represent a permutation for fixed dim=4).