Upgrade to PyTorch 2.3. (#546)

As discussed on Discord, this is a significant upgrade because it is the
first stable release that has a fully functional `torch.export.export`
with the preferred dynamic shapes support. It is also just prior to
nightlies that completely remove support for the old constraints based
API, so is therefore a good point to stop for a moment and support both
styles.

This patch makes a number of API changes:

* Issues deprecation warnings if the `constraints=` keyword for jittable
is used, otherwise not passing it to PyTorch. This should make jittable
not immediately incompatible with later nightlies unless if that feature
is used.
* Adds the ability for a `CompiledModule` to directly have an attribute
of a `torch.export.ExportedProgram`, allowing the user to pre-export
with Torch and then construct a compiled module from that (vs the
`jittable` approach where the `CompiledModule` API was directly invoking
Torch internals to do so). This defaults to exporting as `public` if
given a name not starting with an underscore and private otherwise.
Private ExportedPrograms can be called from procedures just as with
`jittable`.
* `shark_turbine.aot.export()` now accepts either an `CompiledModule`,
`nn.Module`, a or a `torch.export.ExportedProgram`. For the last two, a
new `external_params=` bool is available to control whether parameters
are inlined or externalized. For an `nn.Module` arguments corresponding
to `torch.export.export` are added. Internally, for an `nn.Module`, it
simply calls `torch.export.export`. `jittable` is no longer used
internally.

Some attempt has been made to be backwards compatible with Torch 2.1.0.
New features will not work, but we should be able to support a short
buffer window where older pinned systems are not completely broken. The
repository prior to this patch will be branched to `torch_2.1`.

Breaking changes:

* ops.iree.trace_tensors (plural) had to be removed because the PyTorch
auto functionalization thing has a TODO around lists of tensors. We can
add a wrapper that takes a list and invokves trace_tensors multiple
times and/or ass a `functional_trace_tensors` which works a bit better
with the infra.
* stateless_llama_test.py::test_rerotated_torch_comparison marked as
expectedFailure. Filed #560

Author: stellaraccident

.github/workflows/test.yml

@@ -39,12 +39,11 @@ jobs:
           # Note: We install in three steps in order to satisfy requirements
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
-          pip install --index-url https://download.pytorch.org/whl/cpu \
-            -r core/pytorch-cpu-requirements.txt \
-            -r core/torchvision-requirements.txt
+          pip install -r core/pytorch-cpu-requirements.txt
           pip install --upgrade \
             -r core/requirements.txt \
-            -r mypy-requirements.txt
+            -r mypy-requirements.txt \
+            -r serving/requirements.txt
           pip install -e core[testing] -e serving[testing]
 
       - name: Run core tests

.github/workflows/test_models.yml

@@ -38,12 +38,10 @@ jobs:
           # Note: We install in three steps in order to satisfy requirements
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
-          pip install --index-url https://download.pytorch.org/whl/cpu \
-            -r core/pytorch-cpu-requirements.txt \
-            -r core/torchvision-requirements.txt          
-          pip install --upgrade -r core/requirements.txt
-          pip install -e core[testing]
-          pip install -e models
+          pip install -r core/pytorch-cpu-requirements.txt
+          pip install --pre --upgrade -r core/requirements.txt
+          pip install --pre -e core[testing]
+          pip install --pre -e models
       
       - name: Show current free memory
         run: |

.github/workflows/test_sdxl.yml

@@ -31,8 +31,7 @@ jobs:
           # from non default locations first. Installing the PyTorch CPU
           # wheels saves multiple minutes and a lot of bandwidth on runner setup.
           pip install --index-url https://download.pytorch.org/whl/cpu \
-            -r core/pytorch-cpu-requirements.txt \
-            -r core/torchvision-requirements.txt          
+            -r core/pytorch-cpu-requirements.txt
           pip install --upgrade -r core/requirements.txt
           pip install -e core[testing,torch-cpu-nightly]
           pip install --upgrade -r models/requirements.txt

MANIFEST.in

@@ -1,5 +1,4 @@
 include README.md
 include requirements.txt
 include pytorch-cpu-requirements.txt
-include torchvision-requirements.txt
 include version_info.json

README.md

@@ -45,9 +45,7 @@ pip install shark-turbine
 The above does install some unecessary cuda/cudnn packages for cpu use. To avoid this you
 can specify pytorch-cpu and install via:
 ```
-pip install --index-url https://download.pytorch.org/whl/cpu \
-    -r core/pytorch-cpu-requirements.txt \
-    -r core/torchvision-requirements.txt
+pip install -r core/pytorch-cpu-requirements.txt
 pip install shark-turbine
 ```
 

core/examples/aot_mlp/mlp_export_dynamic.py

@@ -49,7 +49,12 @@ def main(self, x=aot.AbstractTensor(None, 97, 8, dtype=torch.float32)):
         )
 
 
-exported = aot.export(CompiledMLP)
+batch = torch.export.Dim("batch")
+exported = aot.export(
+    model,
+    args=(torch.empty([2, 97, 8], dtype=torch.float32),),
+    dynamic_shapes={"x": {0: batch}},
+)
 # Note that dynamic Torch IR is created below.
 exported.print_readable()
 

core/iree-requirements.txt

@@ -1,2 +1,2 @@
-iree-compiler==20240311.828
-iree-runtime==20240311.828
+iree-compiler==20240327.844
+iree-runtime==20240327.844

core/pytorch-cpu-requirements.txt

@@ -1,3 +1,3 @@
 --pre
-torch==2.1.0
-mpmath==1.3.0
+--index-url https://download.pytorch.org/whl/test/cpu
+-r pytorch-requirements.txt

core/pytorch-requirements.txt

@@ -0,0 +1,3 @@
+torch==2.3.0
+torchaudio
+torchvision

core/requirements.txt

@@ -4,6 +4,5 @@
 # versions, not specific).
 -f https://openxla.github.io/iree/pip-release-links.html
 
--r pytorch-cpu-requirements.txt
--r torchvision-requirements.txt
+-r pytorch-requirements.txt
 -r iree-requirements.txt