@@ -653,8 +653,8 @@ struct Promoter {
653653 Value resolveSlot (Value projectionOrSlot);
654654
655655 void captureAcrossWait ();
656- void captureAcrossWait (ProbeOp probeOp , ArrayRef<WaitOp> waitOps,
657- Liveness &liveness, DominanceInfo &dominance);
656+ void captureValueAcrossWait (Value value , ArrayRef<WaitOp> waitOps,
657+ Liveness &liveness, DominanceInfo &dominance);
658658
659659 void constructLattice ();
660660 void propagateBackward ();
@@ -847,7 +847,9 @@ Value Promoter::resolveSlot(Value projectionOrSlot) {
847847// / Explicitly capture any probes that are live across an `llhd.wait` as block
848848// / arguments and destination operand of that wait. This ensures that replacing
849849// / the probe with a reaching definition later on will capture the value of the
850- // / reaching definition before the wait.
850+ // / reaching definition before the wait. Also capture any extract ops of probes
851+ // / that are live across wait, to allow the extract result (e.g., i1) to be
852+ // / captured instead of the full probe result (e.g., i8).
851853void Promoter::captureAcrossWait () {
852854 if (region.hasOneBlock ())
853855 return ;
@@ -861,53 +863,82 @@ void Promoter::captureAcrossWait() {
861863 Liveness liveness (region.getParentOp ());
862864
863865 SmallVector<WaitOp> crossingWaitOps;
866+
867+ // Capture probes that are live across wait ops.
864868 for (auto &block : region) {
865869 for (auto probeOp : block.getOps <ProbeOp>()) {
866870 for (auto waitOp : waitOps)
867871 if (liveness.getLiveness (waitOp->getBlock ())->isLiveOut (probeOp))
868872 crossingWaitOps.push_back (waitOp);
869873 if (!crossingWaitOps.empty ()) {
870- captureAcrossWait (probeOp, crossingWaitOps, liveness, dominance);
874+ captureValueAcrossWait (probeOp, crossingWaitOps, liveness, dominance);
875+ crossingWaitOps.clear ();
876+ }
877+ }
878+ }
879+
880+ // Capture extract ops of probes that are live across wait ops. This handles
881+ // cases like `@(posedge data[0])` where the extract result (i1) should be
882+ // captured instead of the probe result (i8), enabling the probe to be
883+ // hoisted.
884+ for (auto &block : region) {
885+ for (auto extractOp : block.getOps <comb::ExtractOp>()) {
886+ // Only handle extracts of probe results.
887+ if (!extractOp.getInput ().getDefiningOp <ProbeOp>())
888+ continue ;
889+ for (auto waitOp : waitOps)
890+ if (liveness.getLiveness (waitOp->getBlock ())->isLiveOut (extractOp))
891+ crossingWaitOps.push_back (waitOp);
892+ if (!crossingWaitOps.empty ()) {
893+ captureValueAcrossWait (extractOp, crossingWaitOps, liveness, dominance);
871894 crossingWaitOps.clear ();
872895 }
873896 }
874897 }
875898}
876899
877- // / Add a probe as block argument to a list of wait ops and update uses of the
878- // / probe to use the added block arguments as appropriate. This may insert
879- // / additional block arguments in case the probe and added block arguments both
900+ // / Add a value as block argument to a list of wait ops and update uses of the
901+ // / value to use the added block arguments as appropriate. This may insert
902+ // / additional block arguments in case the value and added block arguments both
880903// / reach the same block.
881- void Promoter::captureAcrossWait (ProbeOp probeOp, ArrayRef<WaitOp> waitOps,
882- Liveness &liveness, DominanceInfo &dominance) {
904+ void Promoter::captureValueAcrossWait (Value value, ArrayRef<WaitOp> waitOps,
905+ Liveness &liveness,
906+ DominanceInfo &dominance) {
883907 LLVM_DEBUG ({
884- llvm::dbgs () << " Capture " << probeOp << " \n " ;
908+ llvm::dbgs () << " Capture " << value << " \n " ;
885909 for (auto waitOp : waitOps)
886910 llvm::dbgs () << " - Across " << waitOp << " \n " ;
887911 });
888912
889- // Calculate the merge points for this probe once it gets promoted to block
913+ // Get the block where the value is defined.
914+ Block *defBlock = nullptr ;
915+ if (auto *defOp = value.getDefiningOp ())
916+ defBlock = defOp->getBlock ();
917+ else
918+ defBlock = cast<BlockArgument>(value).getOwner ();
919+
920+ // Calculate the merge points for this value once it gets promoted to block
890921 // arguments across the wait ops.
891922 auto &domTree = dominance.getDomTree (®ion);
892923 llvm::IDFCalculatorBase<Block, false > idfCalculator (domTree);
893924
894- // Calculate the set of blocks which will define this probe as a distinct
925+ // Calculate the set of blocks which will define this value as a distinct
895926 // value.
896927 SmallPtrSet<Block *, 4 > definingBlocks;
897- definingBlocks.insert (probeOp-> getBlock () );
928+ definingBlocks.insert (defBlock );
898929 for (auto waitOp : waitOps)
899930 definingBlocks.insert (waitOp.getDest ());
900931 idfCalculator.setDefiningBlocks (definingBlocks);
901932
902- // Calculate where the probe is live.
933+ // Calculate where the value is live.
903934 SmallPtrSet<Block *, 16 > liveInBlocks;
904935 for (auto &block : region)
905- if (liveness.getLiveness (&block)->isLiveIn (probeOp ))
936+ if (liveness.getLiveness (&block)->isLiveIn (value ))
906937 liveInBlocks.insert (&block);
907938 idfCalculator.setLiveInBlocks (liveInBlocks);
908939
909940 // Calculate the merge points where we will have to insert block arguments for
910- // this probe .
941+ // this value .
911942 SmallVector<Block *> mergePointsVec;
912943 idfCalculator.calculate (mergePointsVec);
913944 SmallPtrSet<Block *, 16 > mergePoints (mergePointsVec.begin (),
@@ -916,29 +947,28 @@ void Promoter::captureAcrossWait(ProbeOp probeOp, ArrayRef<WaitOp> waitOps,
916947 mergePoints.insert (waitOp.getDest ());
917948 LLVM_DEBUG (llvm::dbgs () << " - " << mergePoints.size () << " merge points\n " );
918949
919- // Perform a depth-first search starting at the block containing the probe ,
950+ // Perform a depth-first search starting at the block containing the value ,
920951 // which dominates all its uses. When we encounter a block that is a merge
921952 // point, insert a block argument.
922953 struct WorklistItem {
923954 DominanceInfoNode *domNode;
924955 Value reachingDef;
925956 };
926957 SmallVector<WorklistItem> worklist;
927- worklist.push_back ({domTree.getNode (probeOp-> getBlock ()), probeOp });
958+ worklist.push_back ({domTree.getNode (defBlock), value });
928959
929960 while (!worklist.empty ()) {
930961 auto item = worklist.pop_back_val ();
931962 auto *block = item.domNode ->getBlock ();
932963
933- // If this block is a merge point, insert a block argument for the probe .
964+ // If this block is a merge point, insert a block argument for the value .
934965 if (mergePoints.contains (block))
935- item.reachingDef =
936- block->addArgument (probeOp.getType (), probeOp.getLoc ());
966+ item.reachingDef = block->addArgument (value.getType (), value.getLoc ());
937967
938- // Replace any uses of the probe in this block with the current reaching
968+ // Replace any uses of the value in this block with the current reaching
939969 // definition.
940970 for (auto &op : *block)
941- op.replaceUsesOfWith (probeOp , item.reachingDef );
971+ op.replaceUsesOfWith (value , item.reachingDef );
942972
943973 // If the terminator of this block branches to a merge point, add the
944974 // current reaching definition as a destination operand.
0 commit comments