Don't hoist bit-extend op as a leaf (#19871)

There are a ton bit extend ops that are getting hoisted which simply
convert the weights from f16 to f32 (these ops are fairly small so they
don't trigger the max size increase threshold i.e. 1024 elems). Instead,
we want these ops to be fused with their consumers to prevent
materializing the high bit-width tensors.

---------

Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>

.github/workflows/pkgci_regression_test.yml

@@ -131,9 +131,9 @@ jobs:
             --goldentime-rocm-unet-ms 255.0 \
             --goldentime-rocm-clip-ms 14.5 \
             --goldentime-rocm-vae-ms 310.0 \
-            --goldendispatch-rocm-unet 1602 \
-            --goldendispatch-rocm-clip 1139 \
-            --goldendispatch-rocm-vae 246 \
+            --goldendispatch-rocm-unet 1236 \
+            --goldendispatch-rocm-clip 967 \
+            --goldendispatch-rocm-vae 208 \
             --goldensize-rocm-unet-bytes 2280000  \
             --goldensize-rocm-clip-bytes 860000 \
             --goldensize-rocm-vae-bytes 840000 \
@@ -156,16 +156,16 @@ jobs:
             --goldentime-rocm-unet-ms 80.0 \
             --goldentime-rocm-clip-ms 15.0 \
             --goldentime-rocm-vae-ms 75.0 \
-            --goldendispatch-rocm-unet 1602 \
-            --goldendispatch-rocm-clip 1139 \
-            --goldendispatch-rocm-vae 246 \
+            --goldendispatch-rocm-unet 1236 \
+            --goldendispatch-rocm-clip 967 \
+            --goldendispatch-rocm-vae 208 \
             --goldensize-rocm-unet-bytes 2270000 \
             --goldensize-rocm-clip-bytes 860000  \
             --goldensize-rocm-vae-bytes 840000 \
             --goldentime-rocm-punet-int8-fp16-ms 50.0 \
             --goldentime-rocm-punet-int8-fp8-ms 52.0 \
-            --goldendispatch-rocm-punet-int8-fp16 1424 \
-            --goldendispatch-rocm-punet-int8-fp8 1704 \
+            --goldendispatch-rocm-punet-int8-fp16 1419 \
+            --goldendispatch-rocm-punet-int8-fp8 1699 \
             --goldensize-rocm-punet-int8-fp8-bytes 2800000 \
             --goldensize-rocm-punet-int8-fp16-bytes 2560000 \
             --rocm-chip gfx942 \
@@ -185,16 +185,16 @@ jobs:
             --goldentime-rocm-unet-ms 195.0 \
             --goldentime-rocm-clip-ms 15.0 \
             --goldentime-rocm-vae-ms 190.0 \
-            --goldendispatch-rocm-unet 1602 \
-            --goldendispatch-rocm-clip 1139 \
-            --goldendispatch-rocm-vae 246 \
+            --goldendispatch-rocm-unet 1236 \
+            --goldendispatch-rocm-clip 967 \
+            --goldendispatch-rocm-vae 208 \
             --goldensize-rocm-unet-bytes 2270000 \
             --goldensize-rocm-clip-bytes 860000 \
             --goldensize-rocm-vae-bytes 840000 \
             --goldentime-rocm-punet-int8-fp16-ms 140.0 \
             --goldentime-rocm-punet-int8-fp8-ms 150 \
-            --goldendispatch-rocm-punet-int8-fp16 1424 \
-            --goldendispatch-rocm-punet-int8-fp8 1704 \
+            --goldendispatch-rocm-punet-int8-fp16 1419 \
+            --goldendispatch-rocm-punet-int8-fp8 1699 \
             --goldensize-rocm-punet-int8-fp8-bytes 2800000 \
             --goldensize-rocm-punet-int8-fp16-bytes 2560000 \
             --rocm-chip gfx942 \

compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel

@@ -36,6 +36,7 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/Flow/IR",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
+        "//compiler/src/iree/compiler/Dialect/LinalgExt/Utils",
         "//compiler/src/iree/compiler/Dialect/Stream/IR",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "@llvm-project//mlir:ArithDialect",

compiler/src/iree/compiler/ExternalInterfaces/CMakeLists.txt

@@ -42,6 +42,7 @@ iree_cc_library(
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::LinalgExt::IR
+    iree::compiler::Dialect::LinalgExt::Utils
     iree::compiler::Dialect::Stream::IR
     iree::compiler::Dialect::Util::IR
   PUBLIC

compiler/src/iree/compiler/ExternalInterfaces/UtilExternalModels.cpp

@@ -12,6 +12,7 @@
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
@@ -334,26 +335,33 @@ struct HoistableLinalgOpInterface
           HoistableLinalgOpInterface<OpTy>, OpTy> {
   bool isHoistableOp(Operation *) const { return true; }
 
-  /// FillOp and broadcasting ops are not hoistableLeaf ops, since it is
-  /// typically better to fuse them with their consumers.
+  // Determines if a linalg op is a hoistable leaf, based on heuristics.
   bool isHoistableLeafOp(Operation *op) const {
+    // Don't hoist bit extend ops because fusing them with their
+    // consumers prevents materializing the high bit-width tensor and they
+    // preform very little real computation.
+    if (IREE::LinalgExt::isBitExtendOp(op)) {
+      return false;
+    }
+
+    // Hoist all non-generic linalg ops except for fill ops which should be
+    // fused with their consumers.
     auto genericOp = llvm::dyn_cast<linalg::GenericOp>(op);
-    if (!genericOp)
+    if (!genericOp) {
       return !isa<linalg::FillOp>(op);
-    // Generally, we prefer to not hoist broadcasts.
-    // Detect op that only broadcast input as fusing them makes the new
+    }
+    if (linalg::isaFillOpInterface(genericOp).has_value()) {
+      return false;
+    }
+
+    // Don't hoist broadcast-like ops because fusing them makes the new
     // op cheaper.
-    if (genericOp.getNumParallelLoops() == genericOp.getNumLoops() &&
-        isa<linalg::YieldOp>(genericOp.getBody()->front())) {
-      for (OpOperand *opOperand : genericOp.getDpsInputOperands()) {
-        AffineMap indexingMap = genericOp.getMatchingIndexingMap(opOperand);
-        if (indexingMap.isProjectedPermutation() &&
-            indexingMap.getNumDims() != indexingMap.getNumResults()) {
-          return false;
-        }
-      }
+    if (linalg::isaBroadcastOpInterface(genericOp).has_value()) {
+      return false;
     }
-    return !linalg::isaFillOpInterface(genericOp).has_value();
+
+    // Hoist all other ops.
+    return true;
   }
   bool isAtomicallyHoistableOp(Operation *) const { return true; }
   bool isOperandHoistable(Operation *, OpOperand *) const { return true; }