Functionalization integration

[bdhirsh]

Summary:
This is the consolidated branch of all the functionalization changes. For individual changes, please visit the corresponding PRs for details.

Here is a brief summary of the highlights: Funtionalization is a dispatcher pass introduced in upstream to remove views and mutations from a PyTorch program to produce functional graphs which is better for backend compilers to do their optimizations. We have a in-house infrastructure to do similar tricks as the HLOs we generated are always functional. The benefits of adopting the upstream Funtionalization pass is that we can get rid of the in-house view infrastructure that we have struggled to maintain for years, and let upstream do all those heavy lifting for us.

Implementation details:

1. To enable Funtionalization, we just need to wrap our newly created at::Tensor that holds a XLATensor in a FunctionalWrapper and then return the wrapper to Python. Then for any consecutive ops will then first go to the Funtionalization pass before reaching us. Correspondingly, we then have to unwrap the FunctionalWrapper before getting the XLATensor. Basically a thin layer called FunctionalWrapper is added to the whole aten-xla bridge process. FunctionalWrapper <=> at::Tensor <=> XLATensor.
2. To support the new view ops, for each view op, we have to implement at most two variants: view_copy: the view op but returns a copy instead; view_scatter: sometimes we need extra logic to reapply the updated value to the base view.
3. For in-place ops, we have a new op called _propagate_xla_data to keep any in-place ops optimization we had before active.

Design Doc

Test Plan:
CI

[JackCaoG]

@bdhirsh build is failing, is that expected?

[bdhirsh]

whoops - I think this is now fixed, but I'll confirm

[wonjoolee95]

Okay, fixed some build failures that I had after rebasing the upstream PyTorch PR and this PR. The CI failure on this CI is not up to date, it is being shown because the pinned upstream PR is out of date.

Now I can see that bunch of PyTorch's index_reduce tests failing. For a specific example:

======================================================================
ERROR: test_index_reduce_reduce_mean_xla_bfloat16 (__main__.TestTorchDeviceTypeXLA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_internal/common_device_type.py", line 378, in instantiated_test
    result = test(self, **param_kwargs)
  File "test/test_torch.py", line 3011, in test_index_reduce
    dest.index_reduce_(dim, idx, src, reduce, include_self=include_self)
IndexError: select(): index 0 out of range for tensor of size [3, 4, 0] at dimension 2

----------------------------------------------------------------------
Ran 1 test in 13.533s

This test failure is caused at line https://github.com/pytorch/pytorch/blob/master/test/test_torch.py#L3007. Putting some debugging statements, I can see that the parameters (dim, idx, src) that cause the failures are:

dim=2
idx=tensor([], device='xla:0', dtype=torch.int64)
src=tensor([], device='xla:0', size=(3, 4, 0), dtype=torch.bfloat16)

The line dest.index_reduce_(dim, idx, src, reduce, include_self=include_self) doesn't seem to trigger XLA's select ops either. @bdhirsh, any quick ideas why this line might be give such error?

[wonjoolee95]

Opened a separate PR on PyTorch side, since I didn't have access to update/rebase the original one -- pytorch/pytorch#88787.

With the latest rebase, the new sym int copy ops (79cbdb6) caused some failures that was addressed by updating at::native::narrow_copy_dense_symint to at::functionalization::functionalize_aten_op_symint<ATEN_OP(narrow_copy)>. Need to double check if the dense can be safely dropped.

[wonjoolee95]

Was able to fix some of the errors seen before, but now I'm also seeing an odd failure for a very simple test test_bool_tensor_value_change at https://github.com/pytorch/pytorch/blob/master/test/test_torch.py#L2760:

test/test_torch.py:2764: UserWarning: 0The operator aten::as_strided appears to be a view operator, but it has no implementation for the backend "xla:0". View operators don't support falling back to run on the CPU, since the tensor's storage cannot be shared across devices. (Triggered internally at /workspace/pytorch/aten/src/ATen/native/CPUFallback.cpp:188.)
  x[0] = False
E
======================================================================
ERROR: test_bool_tensor_value_change_xla (__main__.TestTorchDeviceTypeXLA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_internal/common_device_type.py", line 391, in instantiated_test
    raise rte
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_internal/common_device_type.py", line 378, in instantiated_test
    result = test(self, **param_kwargs)
  File "test/test_torch.py", line 2764, in test_bool_tensor_value_change
    x[0] = False
RuntimeError: self__storage_saved.value().is_alias_of(result.storage()) INTERNAL ASSERT FAILED at "/workspace/pytorch/torch/csrc/autograd/generated/VariableType_3.cpp":11489, please report a bug to PyTorch. 

----------------------------------------------------------------------
Ran 1 test in 0.081s

FAILED (errors=1)

According to the error message, it seems like the op as_strided doesn't seem to get correctly dispatched into XLA hence getting dispatched to CPU and failing with the error above. If I test without the changes in this PR, I can confirm that this op gets dispatched to XLA. I'm syncing offline with @bdhirsh to fix this.

Meanwhile, I'll get started with a draft PR to remove view/aliasing infrastructure in XLA.

[alanwaketan]

@wonjoolee95 Can you use TF_VLOG instead of std::cout for your own logging? Such that we can optionally turn them on or off by envs.

[wonjoolee95]

Seems like the upstream PR might need a rebase. @alanwaketan, do we need the commit [Revert "Use C10_AS_INTARRAYREF_SLOW]?

[alanwaketan]

My bad, yea, the upstream PR will need a rebase as well.

[alanwaketan]

I'm preparing to land the branch finally. @yeounoh Please review the SPMD part and make sure it's correct. @vanbasten23 Please review the DS part and make sure it's correct.

[JackCaoG]

Thanks and great work @alanwaketan @wonjoolee95 @vanbasten23 @bdhirsh !

[alanwaketan]

Thanks Jack and Xiongfei for reviewing.