Add strip assertions flag (iree-org#19014)

This change _temporarily_ adds `iree-opt-strip-assertions` to SDXL CI which should fix the regression discussed in iree-org#19002. Assertions within `linalg.generic` ops can mess with dispatch creation and lead to poorly performing dispatches, despite the fact that they get stripped in later pipelines. --------- Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
Groverkss · Nov 29, 2024 · 3835004 · 3835004
1 parent f51e1da
commit 3835004
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 10 deletions.
diff --git a/.github/workflows/pkgci_regression_test.yml b/.github/workflows/pkgci_regression_test.yml
@@ -221,7 +221,7 @@ jobs:
             --goldentime-rocm-clip-ms 18.5 \
             --goldentime-rocm-vae-ms 337.0 \
             --goldendispatch-rocm-unet 1531 \
-            --goldendispatch-rocm-clip 1141 \
+            --goldendispatch-rocm-clip 1139 \
             --goldendispatch-rocm-vae 246 \
             --goldensize-rocm-unet-bytes 2280000  \
             --goldensize-rocm-clip-bytes 860000 \
@@ -243,7 +243,7 @@ jobs:
             --goldentime-rocm-clip-ms 15.5 \
             --goldentime-rocm-vae-ms 80.0 \
             --goldendispatch-rocm-unet 1531 \
-            --goldendispatch-rocm-clip 1141 \
+            --goldendispatch-rocm-clip 1139 \
             --goldendispatch-rocm-vae 246 \
             --goldensize-rocm-unet-bytes 2270000 \
             --goldensize-rocm-clip-bytes 860000  \

diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -95,6 +95,8 @@ void buildGlobalOptimizationPassPipeline(
 
   // Preprocessing passes to get the program into a canonical state.
   FunctionLikeNest(mainPassManager)
+      .addPredicatedPass(transformOptions.options.stripAssertions,
+                         IREE::Util::createStripDebugOpsPass)
       .addPass(IREE::Util::createOptimizeIntArithmeticPass)
       .addPass(createLinalgQuantizedConvToConvPass)
       .addPass(createLinalgQuantizedMatmulToMatmulPass)
@@ -217,16 +219,10 @@ void buildGlobalOptimizationPassPipeline(
 
   FunctionLikeNest(mainPassManager)
       .addPass(IREE::Flow::createCanonicalizerPass)
-      .addPass(mlir::createCSEPass);
-
-  FunctionLikeNest(mainPassManager)
+      .addPass(mlir::createCSEPass)
       // After running const-eval to a fixed point and folding unit extent dims,
       // try any new raising opportunities.
-      .addPass(createRaiseSpecialOpsPass)
-      // Strip std.assert & co after we perform optimizations; prior to this we
-      // may use the assertions to derive information during analysis.
-      .addPredicatedPass(transformOptions.options.stripAssertions,
-                         IREE::Util::createStripDebugOpsPass);
+      .addPass(createRaiseSpecialOpsPass);
 
   // Export after const-eval. If the user wants to keep the input constants
   // as is in the final parameter archive, they will probably want to disable

diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_clip.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_clip.py
@@ -63,6 +63,7 @@
     "--iree-llvmcpu-fail-on-out-of-bounds-stack-allocation=false",
     "--iree-llvmcpu-distribution-size=32",
     "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
     "--iree-llvmcpu-enable-ukernels=all",
     "--iree-global-opt-enable-quantized-matmul-reassociation",
 ]
@@ -92,6 +93,7 @@ def SDXL_CLIP_COMMON_RUN_FLAGS(
     f"--iree-hip-target={rocm_chip}",
     "--iree-input-type=torch",
     "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
     "--iree-global-opt-propagate-transposes=true",
     "--iree-opt-outer-dim-concat=true",
     "--iree-hip-waves-per-eu=2",

diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_unet.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_unet.py
@@ -131,6 +131,7 @@
     "--iree-llvmcpu-fail-on-out-of-bounds-stack-allocation=false",
     "--iree-llvmcpu-distribution-size=32",
     "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
     "--iree-llvmcpu-enable-ukernels=all",
     "--iree-global-opt-enable-quantized-matmul-reassociation",
 ]
@@ -194,6 +195,7 @@ def SDXL_PUNET_INT8_FP8_OUT(
     "--iree-hal-target-backends=rocm",
     f"--iree-hip-target={rocm_chip}",
     "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
     "--iree-global-opt-propagate-transposes=true",
     "--iree-dispatch-creation-enable-fuse-horizontal-contractions=true",
     "--iree-dispatch-creation-enable-aggressive-fusion=true",

diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_vae.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_vae.py
@@ -43,6 +43,7 @@
     "--iree-llvmcpu-fail-on-out-of-bounds-stack-allocation=false",
     "--iree-llvmcpu-distribution-size=32",
     "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
     "--iree-llvmcpu-enable-ukernels=all",
     "--iree-global-opt-enable-quantized-matmul-reassociation",
 ]
@@ -63,6 +64,7 @@ def SDXL_VAE_COMMON_RUN_FLAGS(
     "--iree-hal-target-backends=rocm",
     f"--iree-hip-target={rocm_chip}",
     "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
     "--iree-global-opt-propagate-transposes=true",
     "--iree-opt-outer-dim-concat=true",
     "--iree-llvmgpu-enable-prefetch=true",