[CI] Enable building StableHLO in CI (#4170)

Context: Due to the LLVM version upgrade, stablehlo CI tests have been
temporarily disabled.
Ref link of discussion: #4152

Author: penguin-wwy

build_tools/build_standalone.sh

@@ -20,7 +20,6 @@ cmake -GNinja -B"$build_dir" "$llvm_project_dir/llvm" \
   -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
   -DLLVM_ENABLE_PROJECTS=mlir \
   -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
-  -DTORCH_MLIR_ENABLE_STABLEHLO=OFF \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$project_dir" \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DLLVM_ENABLE_ASSERTIONS=ON \

build_tools/ci/build_posix.sh

@@ -51,7 +51,6 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -DLLVM_TARGETS_TO_BUILD=host \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DTORCH_MLIR_ENABLE_LTC=OFF \
-  -DTORCH_MLIR_ENABLE_STABLEHLO=OFF \
   -DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS=ON
 echo "::endgroup::"
 

build_tools/ci/test_posix.sh

@@ -28,9 +28,9 @@ case $torch_version in
 
     # TODO: Need to verify in the stable version
     # TODO: Enable for nightly once the Stabelhlo integration is done
-    # echo "::group::Run FxImporter2Stablehlo e2e integration tests"
-    # python3 -m e2e_testing.main --config=fx_importer_stablehlo -v
-    # echo "::endgroup::"
+    echo "::group::Run FxImporter2Stablehlo e2e integration tests"
+    python3 -m e2e_testing.main --config=fx_importer_stablehlo -v
+    echo "::endgroup::"
 
     echo "::group::Run FxImporter TOSA e2e integration tests"
     python3 -m e2e_testing.main --config=fx_importer_tosa -v

build_tools/e2eshark_build.sh

@@ -23,7 +23,6 @@ cmake -S "$repo_root/externals/llvm-project/llvm" -B "$build_dir" \
   -DPython3_FIND_VIRTUALENV=ONLY \
   -DLLVM_ENABLE_PROJECTS=mlir \
   -DLLVM_EXTERNAL_PROJECTS="torch-mlir" \
-  -DTORCH_MLIR_ENABLE_STABLEHLO=OFF \
   -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$PWD" \
   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
   -DLLVM_TARGETS_TO_BUILD=host

build_tools/python_deploy/build_linux_packages.sh

@@ -239,7 +239,6 @@ function build_in_tree() {
       -DLLVM_TARGETS_TO_BUILD=host \
       -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
       -DTORCH_MLIR_ENABLE_LTC=${enable_ltc} \
-      -DTORCH_MLIR_ENABLE_STABLEHLO=OFF \
       -DTORCH_MLIR_USE_INSTALLED_PYTORCH="$torch_from_bin" \
       -DTORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO} \
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
@@ -406,7 +405,6 @@ function build_out_of_tree() {
       -DMLIR_DIR="/main_checkout/torch-mlir/llvm-build/lib/cmake/mlir/" \
       -DMLIR_ENABLE_BINDINGS_PYTHON=OFF \
       -DTORCH_MLIR_ENABLE_LTC=${enable_ltc} \
-      -DTORCH_MLIR_ENABLE_STABLEHLO=OFF \
       -DTORCH_MLIR_USE_INSTALLED_PYTORCH="$torch_from_bin" \
       -DTORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO} \
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \

externals/stablehlo

@@ -33,6 +33,7 @@
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
 #include "stablehlo/conversions/linalg/transforms/Passes.h"
 #include "stablehlo/transforms/Passes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 #endif
 
 #ifdef TORCH_MLIR_ENABLE_TOSA

lib/InitAll.cpp

@@ -25,14 +25,14 @@ func.func @test_quantization_per_channel(%arg0: !torch.vtensor<[4,3,7,7],f32>) -
   %0 = torch.vtensor.literal(dense<[4.000000e-01, 1.000000e-01, 2.000000e-01, 3.000000e-01]> : tensor<4xf32>) : !torch.vtensor<[4],f32>
   %1 = torch.vtensor.literal(dense<[4, 1, 2, 3]> : tensor<4xsi8>) : !torch.vtensor<[4],si8>
   %int12 = torch.constant.int 12
-  %int1 = torch.constant.int 1
+  %zero = torch.constant.int 0
   // CHECK: %[[QUANT:.+]] = stablehlo.uniform_quantize %[[ARG0]]
-  // CHECK-SAME: (tensor<4x3x7x7xf32>) -> tensor<4x3x7x7x!quant.uniform<i8:f32:1, {0.4{{.*}}:4,0.1{{.*}}:1,0.2{{.*}}:2,0.3{{.*}}:3}>>
-  %2 = torch.aten.quantize_per_channel %arg0, %0, %1, %int1, %int12 : !torch.vtensor<[4,3,7,7],f32>, !torch.vtensor<[4],f32>, !torch.vtensor<[4],si8>, !torch.int, !torch.int -> !torch.vtensor<[4,3,7,7],!torch.qint8>
+  // CHECK-SAME: (tensor<4x3x7x7xf32>) -> tensor<4x3x7x7x!quant.uniform<i8:f32:0, {0.4{{.*}}:4,0.1{{.*}}:1,0.2{{.*}}:2,0.3{{.*}}:3}>>
+  %2 = torch.aten.quantize_per_channel %arg0, %0, %1, %zero, %int12 : !torch.vtensor<[4,3,7,7],f32>, !torch.vtensor<[4],f32>, !torch.vtensor<[4],si8>, !torch.int, !torch.int -> !torch.vtensor<[4,3,7,7],!torch.qint8>
   %3 = torch.aten.int_repr %2 : !torch.vtensor<[4,3,7,7],!torch.qint8> -> !torch.vtensor<[4,3,7,7],si8>
   // CHECK: %[[DEQ:.+]] = stablehlo.uniform_dequantize %[[QUANT]]
-  // CHECK-SAME: (tensor<4x3x7x7x!quant.uniform<i8:f32:1, {0.4{{.*}}:4,0.1{{.*}}:1,0.2{{.*}}:2,0.3{{.*}}:3}>>) -> tensor<4x3x7x7xf32>
-  %4 = torch.aten._make_per_channel_quantized_tensor %3, %0, %1, %int1 : !torch.vtensor<[4,3,7,7],si8>, !torch.vtensor<[4],f32>, !torch.vtensor<[4],si8>, !torch.int -> !torch.vtensor<[4,3,7,7],!torch.qint8>
+  // CHECK-SAME: (tensor<4x3x7x7x!quant.uniform<i8:f32:0, {0.4{{.*}}:4,0.1{{.*}}:1,0.2{{.*}}:2,0.3{{.*}}:3}>>) -> tensor<4x3x7x7xf32>
+  %4 = torch.aten._make_per_channel_quantized_tensor %3, %0, %1, %zero : !torch.vtensor<[4,3,7,7],si8>, !torch.vtensor<[4],f32>, !torch.vtensor<[4],si8>, !torch.int -> !torch.vtensor<[4,3,7,7],!torch.qint8>
   %5 = torch.aten.dequantize.self %4 : !torch.vtensor<[4,3,7,7],!torch.qint8> -> !torch.vtensor<[4,3,7,7],f32>
   return %5 : !torch.vtensor<[4,3,7,7],f32>
 }