update

bytedance · Aug 22, 2024 · a45ec01 · a45ec01
1 parent 9bd20d8
commit a45ec01
Show file tree

Hide file tree

Showing 11 changed files with 56 additions and 25 deletions.
diff --git a/.github/workflows/torch-frontend-ci.yaml b/.github/workflows/torch-frontend-ci.yaml
@@ -41,5 +41,6 @@ jobs:
       - name: Checkout byteir repo
         uses: actions/checkout@v3
       - name: Build and test TorchFrontend
+        run: python3 -m pip install -r ./frontends/torch-frontend/torch-cuda-requirements.txt
         run: ./frontends/torch-frontend/scripts/build_and_test.sh
         shell: bash
diff --git a/frontends/README.md b/frontends/README.md
@@ -2,21 +2,21 @@
 
 ByteIR Frontends includes Tensorflow, PyTorch, and ONNX.
 
-Each of them can generates mhlo dialects from the corresponding frontend.
+Each of them can generates stablehlo dialects from the corresponding frontend.
 
 Each frontend can be built independently with the corresponding requirement and dependencies. 
 Note it may or may not be guaranteed using the same version of dependencies, e.g. LLVM, with the other frontend, due to convenience of development.
 
-But each frontend will be guaranteed to generate compatible mhlo format with the ByteIR compiler.
+But each frontend will be guaranteed to generate compatible stablehlo format with the ByteIR compiler.
 
 ## [TensorFlow](tf-frontend/README.md)
-tf graph --> tf dialect --> mhlo dialect pipeline
+tf graph --> tf dialect --> stablehlo dialect pipeline
 
 ## [PyTorch](torch-frontend/README.md)
-PyTorch --> torch dialect --> mhlo dialect pipeline
+PyTorch --> torch dialect --> stablehlo dialect pipeline
 
 ## [ONNX](onnx-frontend/README.md)
-onnx graph --> onnx dialect --> mhlo dialect
+onnx graph --> onnx dialect --> stablehlo dialect
 
 
 

diff --git a/frontends/torch-frontend/README.md b/frontends/torch-frontend/README.md
@@ -1,5 +1,5 @@
 # Torch Frontend
-torch-frontend is a project to build customized torch model --> torch dialect --> mhlo dialect pipeline, where we could add extended dialect and passes.
+torch-frontend is a project to build customized torch model --> torch dialect --> stablehlo dialect pipeline, where we could add extended dialect and passes.
 
 
 ## Quick Start

diff --git a/frontends/torch-frontend/requirements.txt → ...nds/torch-frontend/build-requirements.txt b/frontends/torch-frontend/requirements.txt → ...nds/torch-frontend/build-requirements.txt
@@ -5,22 +5,19 @@
 # torchvision==0.16.0+cpu
 
 # cuda torch and torchvision
---extra-index-url https://download.pytorch.org/whl/cu118
---pre
-torch==2.1.0+cu118
-torchvision==0.16.0+cu118
+# --extra-index-url https://download.pytorch.org/whl/cu118
+# --pre
+# torch==2.1.0+cu118
+# torchvision==0.16.0+cu118
 
 # cuda torch and torchvision nightly
 # --extra-index-url https://download.pytorch.org/whl/nightly/cu118
 # --pre
 # torch==2.1.0.dev20230820+cu118
 # torchvision==0.16.0.dev20230820+cu118
 
-transformers==4.29.2
-
 
 # The following copied from torch-mlir
-numpy
 
 # Build requirements.
 pybind11
@@ -31,9 +28,3 @@ ninja
 pyyaml
 packaging
 
-# Test Requirements
-pillow
-pytest==8.1.0
-dill
-multiprocess
-expecttest
diff --git a/frontends/torch-frontend/scripts/build.sh b/frontends/torch-frontend/scripts/build.sh
@@ -3,6 +3,26 @@
 set -e
 set -x
 
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --disable-jit-ir)
+      TORCH_FRONTEND_ENABLE_JIT_IR_IMPORTER=OFF
+      shift
+      ;;
+    --no-test)
+      TORCH_FRONTEND_TEST=OFF
+      shift
+      ;;
+    *)
+      echo "Invalid option: $1"
+      exit 1
+      ;;
+  esac
+done
+
+TORCH_FRONTEND_ENABLE_JIT_IR_IMPORTER=${TORCH_FRONTEND_ENABLE_JIT_IR_IMPORTER:-ON}
+TORCH_FRONTEND_TEST=${TORCH_FRONTEND_TEST:-ON}
+
 # path to script
 CUR_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
 # path to byteir root
@@ -21,11 +41,16 @@ cmake -S . \
       -DCMAKE_BUILD_TYPE=Release \
       -DCMAKE_C_COMPILER=gcc \
       -DCMAKE_CXX_COMPILER=g++ \
+      -DTORCH_FRONTEND_ENABLE_JIT_IR_IMPORTER=${TORCH_FRONTEND_ENABLE_JIT_IR_IMPORTER} \
       -DCMAKE_CXX_FLAGS="-Wno-unused-but-set-parameter -Wno-unused-but-set-variable" \
       -DPython3_EXECUTABLE=$(which python3)
 
 cmake --build ./build --target all
 
-PYTHONPATH=build/python_packages/:build/torch_mlir_build/python_packages/torch_mlir python3 -m pytest torch-frontend/python/test
+if [[ $TORCH_FRONTEND_TEST == "ON" ]]; then
+  python3 -m pip install -r test-requirements.txt
+  install_mhlo_tools
+  PYTHONPATH=build/python_packages/:build/torch_mlir_build/python_packages/torch_mlir python3 -m pytest torch-frontend/python/test
+fi
 
 popd
diff --git a/frontends/torch-frontend/scripts/build_and_test.sh b/frontends/torch-frontend/scripts/build_and_test.sh
@@ -49,6 +49,8 @@ cmake -S . \
 cmake --build ./build --target all
 
 if [[ $TORCH_FRONTEND_TEST == "ON" ]]; then
+  python3 -m pip install -r test-requirements.txt
+  install_mhlo_tools
   PYTHONPATH=build/python_packages/:build/torch_mlir_build/python_packages/torch_mlir TORCH_DISABLE_NATIVE_FUNCOL=1 python3 -m pytest torch-frontend/python/test
 fi
 

diff --git a/frontends/torch-frontend/scripts/envsetup.sh b/frontends/torch-frontend/scripts/envsetup.sh
@@ -32,9 +32,8 @@ function apply_patches() {
 function prepare_for_build_with_prebuilt() {
   pushd ${PROJ_DIR}
   # install requirements
-  python3 -m pip install -r requirements.txt
+  python3 -m pip install -r build-requirements.txt
   # python3 -m pip install --no-cache-dir torch==2.1.0+cu118 torchvision==0.16.0+cu118 -f https://download.pytorch.org/whl/torch_stable.html
-  install_mhlo_tools
 
   # initialize submodule
   git submodule update --init -f $TORCH_MLIR_ROOT
@@ -49,9 +48,8 @@ function prepare_for_build_with_prebuilt() {
 function prepare_for_build() {
   pushd ${PROJ_DIR}
   # install requirements
-  python3 -m pip install -r requirements.txt
+  python3 -m pip install -r build-requirements.txt
   # python3 -m pip install --no-cache-dir torch==2.1.0+cu118 torchvision==0.16.0+cu118 -f https://download.pytorch.org/whl/torch_stable.html
-  install_mhlo_tools
 
   # initialize submodule
   git submodule update --init --recursive -f $TORCH_MLIR_ROOT

diff --git a/frontends/torch-frontend/test-requirements.txt b/frontends/torch-frontend/test-requirements.txt
@@ -0,0 +1,9 @@
+numpy
+transformers==4.29.2
+
+# Test Requirements
+pillow
+pytest==8.1.0
+dill
+multiprocess
+expecttest
diff --git a/frontends/torch-frontend/torch-cpu-requirements.txt b/frontends/torch-frontend/torch-cpu-requirements.txt
@@ -1,3 +1,4 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 --pre
 torch==2.1.0+cpu
+torchvision==0.16.0+cpu
diff --git a/frontends/torch-frontend/torch-cuda-requirements.txt b/frontends/torch-frontend/torch-cuda-requirements.txt
@@ -0,0 +1,4 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+--pre
+torch==2.1.0+cu118
+torchvision==0.16.0+cu118
diff --git a/tests/build_and_test_e2e.sh b/tests/build_and_test_e2e.sh
@@ -13,13 +13,13 @@ bash scripts/compiler/build_and_test.sh --no-test
 # build runtime
 bash scripts/runtime/build_and_test.sh --cuda --python --no-test
 # build torch_frontend
+pip3 install -r $ROOT_PROJ_DIR/frontends/torch-frontend/torch-cuda-requirements.txt
 bash frontends/torch-frontend/scripts/build_and_test.sh --no-test
 
 pip3 install $ROOT_PROJ_DIR/external/AITemplate/python/dist/*.whl
 pip3 install $ROOT_PROJ_DIR/compiler/build/python/dist/*.whl
 pip3 install $ROOT_PROJ_DIR/runtime/python/dist/*.whl
 pip3 install $ROOT_PROJ_DIR/frontends/torch-frontend/build/torch-frontend/python/dist/*.whl
-pip3 install -r $ROOT_PROJ_DIR/frontends/torch-frontend/torch-requirements.txt
 pip3 install flash_attn==2.5.3
 source scripts/prepare.sh
 install_mhlo_tools