Make IRWitnessTable HOISTABLE

This commit is to remove the duplications of IRWitnessTable by making it
HOISTABLE.

Author: jkwak-work

source/slang/slang-ir-autodiff.cpp

@@ -3179,14 +3179,27 @@ struct AutoDiffPass : public InstPassBase
                 List<IRInst*> args;
                 for (auto param : genType->getParams())
                     args.add(param);
-                as<IRWitnessTable>(innerResult.diffWitness)
-                    ->setConcreteType((IRType*)builder.emitSpecializeInst(
-                        builder.getTypeKind(),
-                        originalType,
-                        (UInt)args.getCount(),
-                        args.getBuffer()));
+
+                auto concreteType = as<IRType>(builder.emitSpecializeInst(
+                    builder.getTypeKind(),
+                    originalType,
+                    (UInt)args.getCount(),
+                    args.getBuffer()));
+
+                auto witnessTableType = innerResult.diffWitness->getFullType();
+                auto newWitnessTable = builder.createWitnessTable(witnessTableType, concreteType);
+
+                // Copy all entries from the old witness table to the new one
+                for (auto entry : as<IRWitnessTable>(innerResult.diffWitness)->getEntries())
+                {
+                    builder.createWitnessTableEntry(
+                        newWitnessTable,
+                        entry->getRequirementKey(),
+                        entry->getSatisfyingVal());
+                }
+
                 result.diffWitness =
-                    hoistValueFromGeneric(builder, innerResult.diffWitness, specInst, true);
+                    hoistValueFromGeneric(builder, newWitnessTable, specInst, true);
             }
             return result;
         }

source/slang/slang-ir-clone.cpp

@@ -309,6 +309,12 @@ void cloneDecoration(
     IRInst* newParent,
     IRModule* module)
 {
+    // If the parent is hoistable and is the same as the original decoration's parent,
+    // skip cloning to avoid duplicating decorations
+    if (getIROpInfo(newParent->getOp()).isHoistable() && 
+        oldDecoration->getParent() == newParent)
+        return;
+
     IRBuilder builder(module);
 
     if (auto first = newParent->getFirstDecorationOrChild())

source/slang/slang-ir-inst-defs.h

@@ -294,7 +294,7 @@ INST(GlobalConstant, globalConstant, 0, GLOBAL)
 
 INST(StructKey, key, 0, GLOBAL)
 INST(GlobalGenericParam, global_generic_param, 0, GLOBAL)
-INST(WitnessTable, witness_table, 0, 0)
+INST(WitnessTable, witness_table, 0, HOISTABLE)
 
 INST(IndexedFieldKey, indexedFieldKey, 2, HOISTABLE)
 

source/slang/slang-ir-insts.h

@@ -2934,8 +2934,6 @@ struct IRWitnessTable : IRInst
 
     IRType* getConcreteType() { return (IRType*)getOperand(0); }
 
-    void setConcreteType(IRType* t) { return setOperand(0, t); }
-
     IR_LEAF_ISA(WitnessTable)
 };
 

source/slang/slang-ir-link.cpp

@@ -224,7 +224,6 @@ IRInst* IRSpecContext::maybeCloneValue(IRInst* originalValue)
     case kIROp_StructKey:
     case kIROp_InterfaceRequirementEntry:
     case kIROp_GlobalGenericParam:
-    case kIROp_WitnessTable:
     case kIROp_InterfaceType:
         return cloneGlobalValue(this, originalValue);
 

source/slang/slang-lower-to-ir.cpp

@@ -8138,17 +8138,23 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         //
         auto irWitnessTableBaseType = lowerType(subContext, superType);
 
+        // Register a dummy value to avoid infinite recursions.
+        // Without this, the call to lowerType() can get into an infinite recursion.
+        //
+        context->setGlobalValue(
+            inheritanceDecl,
+            LoweredValInfo::simple(findOuterMostGeneric(subBuilder->getInsertLoc().getParent())));
+
+        auto irSubType = lowerType(subContext, subType);
+
         // Create the IR-level witness table
-        auto irWitnessTable = subBuilder->createWitnessTable(irWitnessTableBaseType, nullptr);
+        auto irWitnessTable = subBuilder->createWitnessTable(irWitnessTableBaseType, irSubType);
 
-        // Register the value now, rather than later, to avoid any possible infinite recursion.
+        // Override with the correct witness-table
         context->setGlobalValue(
             inheritanceDecl,
             LoweredValInfo::simple(findOuterMostGeneric(irWitnessTable)));
 
-        auto irSubType = lowerType(subContext, subType);
-        irWitnessTable->setConcreteType(irSubType);
-
         // TODO(JS):
         // Should the mangled name take part in obfuscation if enabled?
 

tests/autodiff/deduplicate-witness-table.slang

@@ -0,0 +1,33 @@
+//TEST:SIMPLE(filecheck=CHK):-stage compute -entry computeMain -target hlsl
+
+//CHK: struct DiffPair_1
+//CHK-NOT: struct DiffPair_2
+
+RWTexture2D<float> gOutputColor;
+
+struct ShadingFrame : IDifferentiable
+{
+    float3 T;
+}
+
+[Differentiable]
+float computeRay()
+{
+    float3 dir = 1.f;
+    return dot(dir, dir);
+}
+
+[Differentiable]
+float paramRay()
+{
+    DifferentialPair<float> dpDir = fwd_diff(computeRay)();
+    return dpDir.p;
+}
+
+[Shader("compute")]
+[NumThreads(1, 1, 1)]
+void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
+{
+    DifferentialPair<float> dpColor = fwd_diff(paramRay)();
+    gOutputColor[0] = dpColor.p;
+}