Use scratchData on `IRInst` to replace HashSets. (#2978)

* Use scratchData on `IRInst` to replace HashSets.

* Update test results.

* Initialize scratchData.

* Update autodiff documentation.

* Use enum instead of bool.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 32 insertions, 9 deletions

diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index a88d43db3..bf5893db7 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -18,6 +18,7 @@
 #include "slang-ir-insts.h"
 #include "slang-legalize-types.h"
 #include "slang-mangle.h"
+#include "slang-ir-util.h"
 
 namespace Slang
 {
@@ -3466,6 +3467,9 @@ static LegalVal legalizeGlobalParam(
     }
 }
 
+static constexpr int kHasBeenAddedOrProcessedScratchBitIndex = 0;
+static constexpr int kHasBeenAddedScratchBitIndex = 1;
+
 struct IRTypeLegalizationPass
 {
     IRTypeLegalizationContext* context;
@@ -3480,13 +3484,29 @@ struct IRTypeLegalizationPass
     // instructions have ever been added to the work list.
 
     List<IRInst*> workList;
-    HashSet<IRInst*> hasBeenAddedOrProcessedSet;
-    HashSet<IRInst*> addedToWorkListSet;
+
+    bool hasBeenAddedOrProcessed(IRInst* inst)
+    {
+        if (!inst) return true;
+        return (inst->scratchData & (1 << kHasBeenAddedOrProcessedScratchBitIndex)) != 0;
+    }
+    void setHasBeenAddedOrProcessed(IRInst* inst)
+    {
+        inst->scratchData |= (1 << kHasBeenAddedOrProcessedScratchBitIndex);
+    }
+    bool addedToWorkList(IRInst* inst)
+    {
+        return (inst->scratchData & (1 << kHasBeenAddedScratchBitIndex)) != 0;
+    }
+    void setAddedToWorkList(IRInst* inst)
+    {
+        inst->scratchData |= (1 << kHasBeenAddedScratchBitIndex);
+    }
 
     bool hasBeenAddedToWorkListOrProcessed(IRInst* inst)
     {
-        if (hasBeenAddedToWorkList(inst)) return true;
-        return hasBeenAddedOrProcessedSet.contains(inst);
+        if (!inst) return true;
+        return (inst->scratchData != 0);
     }
 
     // We will add a simple query to check whether an instruciton
@@ -3526,7 +3546,7 @@ struct IRTypeLegalizationPass
         //
         if(inst->getOp() == kIROp_InterfaceRequirementEntry) return true;
 
-        return addedToWorkListSet.contains(inst);
+        return addedToWorkList(inst);
     }
 
     // Next we define a convenience routine for adding something to the work list.
@@ -3535,15 +3555,17 @@ struct IRTypeLegalizationPass
     {
         // We want to avoid adding anything we've already added or processed.
         //
-        if(addedToWorkListSet.contains(inst))
+        if(addedToWorkList(inst))
             return;
         workList.add(inst);
-        addedToWorkListSet.add(inst);
-        hasBeenAddedOrProcessedSet.add(inst);
+        setAddedToWorkList(inst);
+        setHasBeenAddedOrProcessed(inst);
     }
 
     void processModule(IRModule* module)
     {
+        initializeScratchData(module->getModuleInst());
+
         // In order to process an entire module, we start by adding the
         // root module insturction to our work list, and then we will
         // proceed to process instructions until the work list goes dry.
@@ -3568,7 +3590,8 @@ struct IRTypeLegalizationPass
                 //
                 List<IRInst*> workListCopy;
                 Swap(workListCopy, workList);
-                addedToWorkListSet.clear();
+
+                resetScratchDataBit(module->getModuleInst(), kHasBeenAddedScratchBitIndex);
 
                 // Now we simply process each instruction on the copy of
                 // the work list, knowing that `processInst` may add additional