[stinkytofu] Implement legalizeImplicitSpecialRegisters function

for handling implicit special registers (SCC, VCC, EXEC) in instructions
And, use both in ToStinkyTofuUtils and StinkyBuildImplicitDependencyPass

Author: pschang-phy

shared/stinkytofu/include/stinkytofu/ir/asm/RegisterKey.hpp

@@ -56,6 +56,14 @@ inline RegKey toRegKey(const StinkyRegister& reg, unsigned offset = 0) {
     return {reg.reg.type, reg.reg.idx + offset};
 }
 
+/// Check if two registers are the same.
+/// @param reg1 The first register.
+/// @param reg2 The second register.
+/// @return True if the two registers are the same, false otherwise.
+inline bool isSameRegister(const StinkyRegister& reg1, const StinkyRegister& reg2) {
+    return toRegKey(reg1) == toRegKey(reg2);
+}
+
 /// Invoke fn(RegKey) for each DWORD in a register operand.
 /// Skips non-register operands (literals, immediates).
 template <typename Fn>

shared/stinkytofu/include/stinkytofu/transforms/asm/LegalizationUtils.hpp

@@ -116,6 +116,18 @@ STINKYTOFU_EXPORT Legalized legalizeDSStoreB192(StinkyInstruction* inst, AsmIRBu
 STINKYTOFU_EXPORT Legalized legalizeDSStoreB256(StinkyInstruction* inst, AsmIRBuilder& irBuilder,
                                                 GfxArchID archId, bool hasVgprMsb);
 
+// Legalize implicit special registers (SCC, VCC, EXEC) on an instruction.
+//
+// HW flags (Flags.def: IF_ImplicitRead/WriteSCC, IF_ImplicitReadVCC,
+// IF_ImplicitRead/WriteEXEC) declare implicit reads/writes that are not
+// encoded as explicit operands. This function inspects those flags and adds
+// the corresponding singleton register (sized by `wavefrontSize` for
+// VCC/EXEC) to the instruction's src/dest list — but only if the register
+// is not already present. The check matches by RegType and idx, which is
+// sufficient for SCC/VCC/EXEC since they are singletons.
+STINKYTOFU_EXPORT void legalizeImplicitSpecialRegisters(StinkyInstruction* inst,
+                                                        uint32_t wavefrontSize);
+
 }  // namespace stinkytofu
 
 #endif  // STINKYTOFU_LEGALIZATION_UTILS_HPP

shared/stinkytofu/include/stinkytofu/transforms/asm/StinkyBuildImplicitDependencyPass.hpp

@@ -31,16 +31,21 @@ namespace stinkytofu {
 class Pass;
 
 /**
- * @brief Attaches LDS pseudo-registers to instructions for implicit dependency tracking.
+ * @brief Attaches implicit registers to instructions for dependency tracking.
  *
- * Adds RegType::LDS pseudo-registers (keyed by MemTokenData token IDs) to LDS-related
- * instructions so that the def-use chain enforces barrier ordering:
+ * Two kinds of implicit dependencies are materialized as registers so that the
+ * def-use chain builder can see them:
  *
- *   LDS writers (tensor_load, ds_write) — token to dest (defines)
- *   LDS readers (ds_read)               — token to src  (uses)
- *   Barriers                            — token to both src and dest
+ * 1) Special registers (SCC, VCC, EXEC) declared via HW flags
+ *    (Flags.def: IF_ImplicitRead/WriteSCC, IF_ImplicitReadVCC,
+ *     IF_ImplicitRead/WriteEXEC). The corresponding singleton register is
+ *    added to src/dest if not already present.
  *
- * This creates the dependency chain:  writer → barrier → reader.
+ * 2) RegType::LDS pseudo-registers keyed by MemTokenData token IDs:
+ *      LDS writers (tensor_load, ds_write) — token to dest (defines)
+ *      LDS readers (ds_read)               — token to src  (uses)
+ *      Barriers                            — token to both src and dest
+ *    This creates the dependency chain:  writer → barrier → reader.
  */
 STINKYTOFU_EXPORT std::unique_ptr<Pass> createStinkyBuildImplicitDependencyPass();
 

shared/stinkytofu/src/conversion/rocisa/ToStinkyTofuUtils.cpp

@@ -210,6 +210,10 @@ Legalized legalizeInstruction(StinkyInstruction* inst, rocisa::Instruction* roci
                               AsmIRBuilder& irBuilder, GfxArchID archId,
                               const std::map<std::string, int>& asmCaps,
                               const std::map<std::string, int>& archCaps, bool hasVgprMsb) {
+    // Attach implicit special registers (SCC/VCC/`EXEC) declared by HW flags
+    // (Flags.def) to the instruction.
+    legalizeImplicitSpecialRegisters(inst, getWaveFrontSize(archId));
+
     if (isBranch(*inst)) {
         // Handle branch instructions
         rocisa::BranchInstruction* branchInst =
@@ -317,19 +321,6 @@ void addRegistersToInstruction(StinkyInstruction* stinkyInst, const rocisa::Inst
         }
     }
 
-    // Add implicit special registers driven by HW flags (Flags.def).
-    if (stinkyInst->is(IF_ImplicitReadSCC)) stinkyInst->addSrcReg(StinkyRegister::getSCCRegister());
-    if (stinkyInst->is(IF_ImplicitWriteSCC))
-        stinkyInst->addDestReg(StinkyRegister::getSCCRegister());
-
-    uint32_t wfs = getWaveFrontSize(archId);
-    if (stinkyInst->is(IF_ImplicitReadVCC))
-        stinkyInst->addSrcReg(StinkyRegister::getVCCRegister(wfs));
-    if (stinkyInst->is(IF_ImplicitReadEXEC))
-        stinkyInst->addSrcReg(StinkyRegister::getEXECRegister(wfs));
-    if (stinkyInst->is(IF_ImplicitWriteEXEC))
-        stinkyInst->addDestReg(StinkyRegister::getEXECRegister(wfs));
-
 #ifndef NDEBUG
     // Verify: read-write operands must exist in both destRegs and srcRegs.
     {

shared/stinkytofu/src/transforms/asm/LegalizationUtils.cpp

@@ -28,6 +28,7 @@
 
 #include "stinkytofu/hardware/ArchHelper.hpp"
 #include "stinkytofu/hardware/GfxIsa.hpp"
+#include "stinkytofu/ir/asm/RegisterKey.hpp"
 #include "stinkytofu/ir/asm/StinkyAsmIR.hpp"
 #include "stinkytofu/ir/asm/StinkyModifiers.hpp"
 
@@ -595,4 +596,36 @@ Legalized legalizeDSStoreB256(StinkyInstruction* inst, AsmIRBuilder& irBuilder,
     return {store1, store2};
 }
 
+namespace {
+// Add `reg` to the dest list of `inst`, unless a register with the same
+// RegType/idx is already present. SCC/VCC/EXEC are singletons, so RegType+idx
+// is sufficient to detect duplicates introduced by an upstream stage or by an
+// instruction that already encodes the register as an explicit operand.
+void addUniqueSpecialDest(StinkyInstruction* inst, const StinkyRegister& reg) {
+    for (const StinkyRegister& d : inst->getDestRegs())
+        if (isSameRegister(d, reg)) return;
+    inst->addDestReg(reg);
+}
+
+void addUniqueSpecialSrc(StinkyInstruction* inst, const StinkyRegister& reg) {
+    for (const StinkyRegister& s : inst->getSrcRegs())
+        if (isSameRegister(s, reg)) return;
+    inst->addSrcReg(reg);
+}
+}  // namespace
+
+void legalizeImplicitSpecialRegisters(StinkyInstruction* inst, uint32_t wavefrontSize) {
+    if (inst == nullptr) return;
+
+    if (inst->is(IF_ImplicitReadSCC)) addUniqueSpecialSrc(inst, StinkyRegister::getSCCRegister());
+    if (inst->is(IF_ImplicitWriteSCC)) addUniqueSpecialDest(inst, StinkyRegister::getSCCRegister());
+
+    if (inst->is(IF_ImplicitReadVCC))
+        addUniqueSpecialSrc(inst, StinkyRegister::getVCCRegister(wavefrontSize));
+    if (inst->is(IF_ImplicitReadEXEC))
+        addUniqueSpecialSrc(inst, StinkyRegister::getEXECRegister(wavefrontSize));
+    if (inst->is(IF_ImplicitWriteEXEC))
+        addUniqueSpecialDest(inst, StinkyRegister::getEXECRegister(wavefrontSize));
+}
+
 }  // namespace stinkytofu

shared/stinkytofu/src/transforms/asm/StinkyBuildImplicitDependencyPass.cpp

@@ -32,21 +32,31 @@
 #include "stinkytofu/core/PassManager.hpp"
 #include "stinkytofu/ir/asm/StinkyAsmIR.hpp"
 #include "stinkytofu/support/ErrorHandling.hpp"
+#include "stinkytofu/transforms/asm/LegalizationUtils.hpp"
 
 #define DEBUG_TYPE "StinkyBuildImplicitDependencyPass"
 
-// Implicit LDS dependency pass
-// ============================
-// Attaches RegType::LDS pseudo-registers to instructions based on their MemTokenData
-// token IDs.  The instruction type determines src vs dest placement:
+// Implicit dependency pass
+// ========================
+// Attaches implicit registers to instructions so that the def-use chain builder
+// can see dependencies that are not encoded as explicit operands. Two kinds of
+// implicit dependencies are handled:
 //
-//   tensor_load / ds_write  →  LDS token to dest  (LDS producer)
-//   ds_read                 →  LDS token to src   (LDS consumer)
-//   barrier / signal / wait →  LDS token to both  (synchronization point)
+// 1) Implicit special registers (SCC, VCC, EXEC) driven by HW flags
+//    (Flags.def: IF_ImplicitRead/WriteSCC, IF_ImplicitReadVCC,
+//     IF_ImplicitRead/WriteEXEC). The corresponding special register is added
+//    to src/dest if not already present.
 //
-// The def-use chain builder then sees:
-//   producer(def LDS[t]) → barrier(use+def LDS[t]) → consumer(use LDS[t])
-// which forces the scheduler to respect: producers → barrier → consumers.
+// 2) RegType::LDS pseudo-registers (keyed by MemTokenData token IDs). The
+//    instruction type determines src vs dest placement:
+//
+//      tensor_load / ds_write  →  LDS token to dest  (LDS producer)
+//      ds_read                 →  LDS token to src   (LDS consumer)
+//      barrier / signal / wait →  LDS token to both  (synchronization point)
+//
+//    The def-use chain builder then sees:
+//      producer(def LDS[t]) → barrier(use+def LDS[t]) → consumer(use LDS[t])
+//    which forces the scheduler to respect: producers → barrier → consumers.
 
 namespace {
 using namespace stinkytofu;
@@ -186,18 +196,27 @@ static void checkConsistentMemTokens(const BasicBlock& bb) {
 
 void setPseudoRegistersInBlock(BasicBlock& bb, PassContext& passCtx,
                                const std::unordered_set<const BasicBlock*>& checkBlocks) {
+    bool doLdsTokenHandling = true;
     if (!passCtx.getPassFeatureConfig().barrierConfig.unrollMovableBarrier) {
-        PASS_DEBUG(std::cerr << "[BuildImplicitDep] skip BB label=\"" << bb.getLabel()
-                             << "\" (unrollMovableBarrier=false)\n");
-        return;
+        PASS_DEBUG(std::cerr << "[BuildImplicitDep] skip LDS-token handling BB label=\""
+                             << bb.getLabel() << "\" (unrollMovableBarrier=false)\n");
+        doLdsTokenHandling = false;
     }
 
-    checkConsistentMemTokens(bb);
+    if (doLdsTokenHandling) {
+        checkConsistentMemTokens(bb);
+    }
 
+    const uint32_t wavefrontSize = passCtx.getWavefrontSize();
     for (auto it = bb.begin(); it != bb.end(); ++it) {
         auto* inst = dyn_cast<StinkyInstruction>(it.getNodePtr());
         if (!inst) continue;
 
+        // Always attach implicit special registers (SCC/VCC/EXEC) declared by HW flags
+        legalizeImplicitSpecialRegisters(inst, wavefrontSize);
+
+        if (!doLdsTokenHandling) continue;
+
         const MemTokenData* mt = inst->getModifier<MemTokenData>();
         if (!mt) continue;
         assert(!mt->tokens.empty() && "MemTokenData with empty tokens");