DNS - Do not hoist encoding ops if the producer is not a flow region op.

Signed-off-by: hanhanW <[email protected]>
iree-org · Feb 13, 2025 · 6ebfcaa · 6ebfcaa
1 parent beae744
commit 6ebfcaa
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/compiler/src/iree/compiler/DispatchCreation/FuseEncodingOpsIntoDispatchRegions.cpp b/compiler/src/iree/compiler/DispatchCreation/FuseEncodingOpsIntoDispatchRegions.cpp
@@ -42,6 +42,9 @@ struct FuseEncodingOpsIntoDispatchRegionsPass
 
     SmallVector<IREE::Encoding::SetEncodingOp> encodingOps;
     funcOp->walk([&](IREE::Encoding::SetEncodingOp encodingOp) {
+      if (encodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) {
+        return;
+      }
       encodingOps.push_back(encodingOp);
     });
 

diff --git a/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp b/compiler/src/iree/compiler/DispatchCreation/HoistEncodingOps.cpp
@@ -195,7 +195,9 @@ void HoistEncodingOpsPass::runOnOperation() {
 
   SmallVector<IREE::Encoding::SetEncodingOp> candidates;
   funcOp->walk([&](IREE::Encoding::SetEncodingOp setEncodingOp) {
-    if (setEncodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>()) {
+    if (setEncodingOp->getParentOfType<IREE::Flow::DispatchRegionOp>() &&
+        setEncodingOp.getSource()
+            .getDefiningOp<IREE::Flow::DispatchRegionOp>()) {
       candidates.push_back(setEncodingOp);
     }
   });