Merge main into dtype-functions-staging (#1935)

Signed-off-by: Gaurav Shukla <[email protected]> Signed-off-by: Prateek Gupta <[email protected]> Co-authored-by: Jiahao Li <[email protected]> Co-authored-by: Yuanqiang Liu <[email protected]> Co-authored-by: Vivek Khandelwal <[email protected]> Co-authored-by: Chi_Liu <[email protected]> Co-authored-by: Victor Guerra <[email protected]> Co-authored-by: Victor Guerra <[email protected]> Co-authored-by: powderluv <[email protected]> Co-authored-by: Ashay Rane <[email protected]> Co-authored-by: Eric Kunze <[email protected]> Co-authored-by: Gleb Kazantaev <[email protected]> Co-authored-by: Matthias Gehre <[email protected]> Co-authored-by: Yi Wang <[email protected]> Co-authored-by: Sean Silva <[email protected]> Co-authored-by: Zachary Cetinic <[email protected]> Co-authored-by: Tanyo Kwok <[email protected]> Co-authored-by: Zachary Cetinic <[email protected]> Co-authored-by: Kunwar Grover <[email protected]> Co-authored-by: Ziheng Jiang <[email protected]> Co-authored-by: Ziheng Jiang <[email protected]> Co-authored-by: Maksim Levental <[email protected]> Co-authored-by: Gaurav Shukla <[email protected]> Co-authored-by: Prateek Gupta <[email protected]> Co-authored-by: nvda <[email protected]> Co-authored-by: Ahmed S. Taei <[email protected]> Co-authored-by: Priya Savithiri <[email protected]> Co-authored-by: Zhekun Zhang <[email protected]> Co-authored-by: zhekun.zhang <[email protected]> Co-authored-by: Kan Chen <[email protected]> Co-authored-by: gpetters94 <[email protected]>
llvm · Mar 15, 2023 · 042d58b · 042d58b
1 parent ce7abf4
commit 042d58b
Show file tree

Hide file tree

Showing 165 changed files with 5,801 additions and 1,845 deletions.
diff --git a/.github/actions/setup-build/action.yml b/.github/actions/setup-build/action.yml
@@ -17,7 +17,7 @@ runs:
   - name: Set up Python
     uses: actions/setup-python@v4
     with:
-      python-version: '3.10'
+      python-version: '3.11'
 
   - name: Install MLIR Python depends
     run: |
@@ -26,7 +26,8 @@ runs:
 
   - name: Install PyTorch nightly depends
     run: |
-      python -m pip install -r requirements.txt
+      python -m pip install -r pytorch-requirements.txt
+      python -m pip install -r build-requirements.txt
     shell: bash
 
   - name: Install prerequisites (Linux)

diff --git a/.github/workflows/RollPyTorch.yml b/.github/workflows/RollPyTorch.yml
@@ -8,12 +8,19 @@ on:
 jobs:
   build_linux:
     name: Manylinux Build
-    runs-on: ubuntu-latest
+    runs-on: a100
 
     # Don't run this in everyone's forks.
     if: github.repository == 'llvm/torch-mlir'
 
     steps:
+
+    - name: Prepare workspace
+      run: |
+        # Clear the workspace directory so that we don't run into errors about
+        # existing lock files.
+        sudo rm -rf $GITHUB_WORKSPACE/*
+
     - name: Get torch-mlir
       uses: actions/checkout@v3
       with:
@@ -31,6 +38,7 @@ jobs:
 
         cd ${GITHUB_WORKSPACE}
         python -m pip install wheel
+        sudo apt-get install unzip
 
         # Fetch the most recent nightly torchvision release
         VISION_RELEASE=$(python -m pip index versions -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre torchvision | grep "Available versions" | tr ' ' '\n' | grep "^[0-9]" | sort --version-sort --reverse | head -n1 | tr -d ',' | sed 's/\([^+]*\).*/\1/')
@@ -44,7 +52,8 @@ jobs:
         # Read the version from the downloaded whl file without extracting it
         PT_RELEASE=$(unzip -p torch-*.whl 'torch-*/METADATA' | grep "^Version:" | awk '{ print $2 }' | sed 's/\([^+]*\).*/\1/')
         echo "Found torch release ${PT_RELEASE}"
-        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorch==%s\ntorchvision==%s\n" "${PT_RELEASE}" "${VISION_RELEASE}" > pytorch-requirements.txt
+        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorch==%s\n" "${PT_RELEASE}" > pytorch-requirements.txt
+        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorchvision==%s\n" "${VISION_RELEASE}" > torchvision-requirements.txt
 
         # Read the commit hash from the downloaded whl file without extracting it
         PT_HASH=$(unzip -p torch-"${PT_RELEASE}"*.whl torch/version.py | grep git_version | awk '{ print $3 }' | tr -d "'")
@@ -96,7 +105,7 @@ jobs:
         git fetch --recurse-submodules=no
         git checkout main
         git pull origin main
-        git add pytorch-hash.txt pytorch-requirements.txt lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
+        git add pytorch-hash.txt pytorch-requirements.txt torchvision-requirements.txt lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
         git diff --cached --exit-code || (git commit -m "update PyTorch version to ${{ env.PT_RELEASE }}" && git push --set-upstream origin main)
 
     - name: Update PyTorch Build Cache (if running on main branch)

diff --git a/.github/workflows/bazelBuildAndTest.yml b/.github/workflows/bazelBuildAndTest.yml
@@ -20,6 +20,12 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
+    - name: Prepare workspace
+      run: |
+        # Clear the workspace directory so that we don't run into errors about
+        # existing lock files.
+        sudo rm -rf $GITHUB_WORKSPACE/*
+
     - name: Checkout torch-mlir
       uses: actions/checkout@v3
       with:

diff --git a/.github/workflows/buildAndTest.yml b/.github/workflows/buildAndTest.yml
@@ -51,6 +51,14 @@ jobs:
     runs-on: ${{ matrix.os }}
 
     steps:
+
+    - name: Prepare workspace
+      if: ${{ matrix.os-arch == 'ubuntu-x86_64' }}
+      run: |
+        # Clear the workspace directory so that we don't run into errors about
+        # existing lock files.
+        sudo rm -rf $GITHUB_WORKSPACE/*
+    
     - name: Checkout torch-mlir
       uses: actions/checkout@v3
       with:
@@ -113,7 +121,7 @@ jobs:
           -DLLVM_USE_HOST_TOOLS=ON \
           -DLLVM_ENABLE_ZSTD=OFF \
           -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-          -DTORCH_MLIR_ENABLE_MHLO=OFF \
+          -DTORCH_MLIR_ENABLE_STABLEHLO=OFF \
           -DTORCH_MLIR_ENABLE_LTC=OFF \
           -DTORCH_MLIR_USE_INSTALLED_PYTORCH="${{ matrix.torch-binary }}" \
           -DMACOSX_DEPLOYMENT_TARGET=12.0 \

diff --git a/.github/workflows/buildRelease.yml b/.github/workflows/buildRelease.yml
@@ -13,8 +13,25 @@ on:
 jobs:
   build_linux:
     name: Manylinux Build
-    runs-on: ubuntu-latest
+    runs-on: a100
+    strategy:
+      matrix:
+        package: [ torch-mlir, torch-mlir-core ]
+        py_version: [ cp38-cp38, cp310-cp310, cp311-cp311 ]
+        exclude:
+          - package: torch-mlir-core
+            py_version: cp38-cp38
+          - package: torch-mlir-core
+            py_version: cp310-cp310
+
     steps:
+
+    - name: Prepare workspace
+      run: |
+        # Clear the workspace directory so that we don't run into errors about
+        # existing lock files.
+        sudo rm -rf $GITHUB_WORKSPACE/*
+
     - name: Get torch-mlir
       uses: actions/checkout@v3
       with:
@@ -28,7 +45,7 @@ jobs:
         python -m pip install wheel
         TM_PACKAGE_VERSION=${{ github.event.inputs.python_package_version }}
         printf "TORCH_MLIR_PYTHON_PACKAGE_VERSION=%s\n" $TM_PACKAGE_VERSION > ./torch_mlir_package_version
-        ./build_tools/python_deploy/build_linux_packages.sh
+        TM_PYTHON_VERSIONS=${{ matrix.py_version }} TM_PACKAGES=${{ matrix.package }} ./build_tools/python_deploy/build_linux_packages.sh
               
     # If we were given a release_id, then upload the package we just built
     # to the github releases page.
@@ -56,7 +73,7 @@ jobs:
       run: mkdir dist
     - name: Copy releases to publish to dist directory
       if: github.event.inputs.release_id != ''
-      run: cp build_tools/python_deploy/wheelhouse/torch_mlir-*.whl dist/
+      run: cp build_tools/python_deploy/wheelhouse/torch_mlir*.whl dist/
 
     # Wheels must be published from a linux environment.
     #
@@ -70,6 +87,9 @@ jobs:
   build_macos:
     name: MacOS Build
     runs-on: macos-latest
+    strategy:
+      matrix:
+        package: [ torch-mlir, torch-mlir-core ]
     steps:
     - name: Get torch-mlir
       uses: actions/checkout@v3
@@ -85,7 +105,7 @@ jobs:
         TM_PACKAGE_VERSION=${{ github.event.inputs.python_package_version }}
         printf "TORCH_MLIR_PYTHON_PACKAGE_VERSION=%s\n" $TM_PACKAGE_VERSION > ./torch_mlir_package_version
         sudo ./build_tools/python_deploy/install_macos_deps.sh
-        TORCH_MLIR_PYTHON_VERSIONS="3.10" ./build_tools/python_deploy/build_macos_packages.sh
+        packages=${{ matrix.package }} TORCH_MLIR_PYTHON_VERSIONS="3.11" ./build_tools/python_deploy/build_macos_packages.sh
 
     # If we were given a release_id, then upload the package we just built
     # to the github releases page.
@@ -113,7 +133,7 @@ jobs:
       run: mkdir dist
     - name: Copy releases to publish to dist directory
       if: github.event.inputs.release_id != ''
-      run: cp build_tools/python_deploy/wheelhouse/torch_mlir-*.whl dist/
+      run: cp build_tools/python_deploy/wheelhouse/torch_mlir*.whl dist/
 
     # Wheels must be published from a linux environment.
     #
@@ -127,6 +147,9 @@ jobs:
   build_windows:
     name: Windows Build
     runs-on: windows-latest
+    strategy:
+      matrix:
+        package: [ torch-mlir, torch-mlir-core ]
     steps:
     - name: Get torch-mlir
       uses: actions/checkout@v3
@@ -142,6 +165,14 @@ jobs:
     - name: Build Python wheels and smoke test.
       shell: pwsh
       run: |
+        if ( "${{ matrix.package }}" -eq "torch-mlir-core" )
+        {
+          $env:TORCH_MLIR_ENABLE_JIT_IR_IMPORTER='0'
+          $env:TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS='1'
+        } else {
+          $env:TORCH_MLIR_ENABLE_JIT_IR_IMPORTER='1'
+          $env:TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS='0'
+        }
         $env:TORCH_MLIR_PYTHON_PACKAGE_VERSION = '${{ github.event.inputs.python_package_version }}'
         ./build_tools/python_deploy/build_windows.ps1
 
@@ -172,7 +203,7 @@ jobs:
       continue-on-error: true
     - name: Copy releases to publish to dist directory
       if: github.event.inputs.release_id != ''
-      run: cp ./wheelhouse/torch_mlir-*.whl dist/
+      run: cp ./wheelhouse/torch_mlir*.whl dist/
 
     # Wheels must be published from a linux environment.
     #
@@ -216,4 +247,4 @@ jobs:
     #  if: github.event.inputs.release_id != ''
     #  uses: pypa/[email protected]
     #  with:
-    #    password: ${{ secrets.PYPI_API_TOKEN }}
+    #    password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.github/workflows/gh-pages-releases.yml b/.github/workflows/gh-pages-releases.yml
@@ -13,8 +13,13 @@ jobs:
     if: github.repository == 'llvm/torch-mlir'
 
     steps:
+      - name: Prepare workspace
+        run: |
+          # Clear the workspace directory so that we don't run into errors about
+          # existing lock files.
+          sudo rm -rf $GITHUB_WORKSPACE/*
       - name: Checking out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
       - name: Run scrape releases script

diff --git a/.github/workflows/oneshotSnapshotPackage.yml b/.github/workflows/oneshotSnapshotPackage.yml
@@ -10,8 +10,14 @@ jobs:
     # Don't run this in everyone's forks.
     if: github.repository == 'llvm/torch-mlir'
     steps:
+      - name: Prepare workspace
+        run: |
+          # Clear the workspace directory so that we don't run into errors about
+          # existing lock files.
+          sudo rm -rf $GITHUB_WORKSPACE/*
+
       - name: Checking out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
 

diff --git a/.github/workflows/releaseSnapshotPackage.yml b/.github/workflows/releaseSnapshotPackage.yml
@@ -13,8 +13,15 @@ jobs:
     # Don't run this in everyone's forks.
     if: github.repository == 'llvm/torch-mlir'
     steps:
+
+      - name: Prepare workspace
+        run: |
+          # Clear the workspace directory so that we don't run into errors about
+          # existing lock files.
+          sudo rm -rf $GITHUB_WORKSPACE/*
+
       - name: Checking out repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}
 

diff --git a/.gitignore b/.gitignore
@@ -32,3 +32,6 @@ bazel-*
 build_oot/
 docker_venv/
 llvm-build/
+
+# C++ build artifacts
+compile_commands.json
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -36,12 +36,18 @@ macro(torch_mlir_add_llvm_external_project name identifier location)
   set(LLVM_EXTERNAL_PROJECTS ${LLVM_EXTERNAL_PROJECTS} CACHE STRING "" FORCE)
 endmacro()
 
-option(TORCH_MLIR_ENABLE_MHLO "Add mhlo dialect" ON)
-if(TORCH_MLIR_ENABLE_MHLO)
-  add_definitions(-DTORCH_MLIR_ENABLE_MHLO)
+option(TORCH_MLIR_ENABLE_STABLEHLO "Add stablehlo dialect" ON)
+if(TORCH_MLIR_ENABLE_STABLEHLO)
+  add_definitions(-DTORCH_MLIR_ENABLE_STABLEHLO)
 endif()
 
+option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
 option(TORCH_MLIR_ENABLE_LTC "Enables LTC backend" OFF)
+option(TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS "Build Torch dialect MLIR Python bindings but neither JIT IR Importer nor LTC backend" OFF)
+if(TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS)
+  set(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER OFF)
+  set(TORCH_MLIR_ENABLE_LTC OFF)
+endif()
 
 if(TORCH_MLIR_ENABLE_LTC)
     set(ENV{TORCH_MLIR_ENABLE_LTC} 1)
@@ -109,7 +115,6 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR OR TORCH_MLIR_OUT_OF_TREE_
   # Don't try to compile the python extensions at the moment. We need
   # to import lots of dependencies from AddMLIRPython to make this work.
   set(MLIR_ENABLE_BINDINGS_PYTHON 1)
-  option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
 
   set(TORCH-MLIR_BUILT_STANDALONE 1)
   set(BACKEND_PACKAGE_STRING "LLVM ${LLVM_PACKAGE_VERSION}")
@@ -119,7 +124,6 @@ else()
   # In-tree build with LLVM_EXTERNAL_PROJECTS=torch-mlir
 
   option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF)
-  option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
 
   # TODO: Fix this upstream so that global include directories are not needed.
   set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir)
@@ -128,8 +132,8 @@ else()
   set(MLIR_INCLUDE_DIRS "${MLIR_INCLUDE_DIR};${MLIR_GENERATED_INCLUDE_DIR}")
 endif()
 
-if (TORCH_MLIR_ENABLE_MHLO)
-  set(MHLO_BUILD_EMBEDDED ON)
+if (TORCH_MLIR_ENABLE_STABLEHLO)
+  set(STABLEHLO_BUILD_EMBEDDED ON)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/externals/mlir-hlo
     ${CMAKE_CURRENT_BINARY_DIR}/mlir-hlo
     EXCLUDE_FROM_ALL)

diff --git a/README.md b/README.md
@@ -8,13 +8,12 @@ necessarily a reflection of the completeness or stability of the code, it
 does indicate that the project is not yet endorsed as a component of LLVM.
 
 [PyTorch](https://pytorch.org)
-An open source machine learning framework that accelerates the path from research prototyping to production deployment.
+PyTorch is an open source machine learning framework that facilitates the seamless transition from research and prototyping to production-level deployment. 
 
 [MLIR](https://mlir.llvm.org)
-The MLIR project is a novel approach to building reusable and extensible compiler infrastructure. MLIR aims to address software fragmentation, improve compilation for heterogeneous hardware, significantly reduce the cost of building domain specific compilers, and aid in connecting existing compilers together.
-
+The MLIR project offers a novel approach for building extensible and reusable compiler architectures, which address the issue of software fragmentation, reduce the cost of developing domain-specific compilers, improve compilation for heterogeneous hardware, and promote compatibility between existing compilers.
 [Torch-MLIR](https://github.com/llvm/torch-mlir)
-Multiple Vendors use MLIR as the middle layer, mapping from platform frameworks like PyTorch, JAX, and TensorFlow into MLIR and then progressively lowering down to their target hardware. We have seen half a dozen custom lowerings from PyTorch to MLIR. Having canonical lowerings from the PyTorch ecosystem to the MLIR ecosystem provides much needed relief to hardware vendors to focus on their unique value rather than implementing yet another PyTorch frontend for MLIR. The goal is to be similar to current hardware vendors adding LLVM target support instead of each one also implementing Clang / a C++ frontend.
+Several vendors have adopted MLIR as the middle layer in their systems, enabling them to map frameworks such as PyTorch, JAX, and TensorFlow into MLIR and subsequently lower them to their target hardware. We have observed half a dozen custom lowerings from PyTorch to MLIR, making it easier for hardware vendors to focus on their unique value, rather than needing to implement yet another PyTorch frontend for MLIR. The ultimate aim is to be similar to the current hardware vendors adding LLVM target support, rather than each one implementing Clang or a C++ frontend.
 
 [![Release Build](https://github.com/llvm/torch-mlir/actions/workflows/buildRelease.yml/badge.svg)](https://github.com/llvm/torch-mlir/actions/workflows/buildRelease.yml)
 
@@ -43,15 +42,26 @@ We have few paths to lower down to the Torch MLIR Dialect.
 
 ## Install torch-mlir snapshot
 
-This installs a pre-built snapshot of torch-mlir for Python 3.7/3.8/3.9/3.10 on Linux and macOS.
+At the time of writing, we release pre-built snapshot of torch-mlir for Python 3.10 on Linux and macOS.
 
+If you have Python 3.10, the following commands initialize a virtual environment.
 ```shell
-python -m venv mlir_venv
+python3.10 -m venv mlir_venv
 source mlir_venv/bin/activate
-# Some older pip installs may not be able to handle the recent PyTorch deps
+```
+
+Or, if you want to switch over multiple versions of Python using conda, you can create a conda environment with Python 3.10.
+```shell
+conda create -n torch-mlir python=3.10
+conda activate torch-mlir
 python -m pip install --upgrade pip
-pip install --pre torch-mlir torchvision -f https://llvm.github.io/torch-mlir/package-index/ --extra-index-url https://download.pytorch.org/whl/nightly/cpu
-# This will install the corresponding torch and torchvision nightlies
+```
+
+Then, we can install torch-mlir with the corresponding torch and torchvision nightlies.
+```
+pip install --pre torch-mlir torchvision \
+  -f https://llvm.github.io/torch-mlir/package-index/
+  --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 ```
 
 ## Demos

diff --git a/build-requirements.txt b/build-requirements.txt
@@ -1,5 +1,3 @@
--r pytorch-requirements.txt
-
 numpy
 pybind11
 wheel