WIP: Switch WitnessTable to HOISTABLE

This commit deduplicates Witnessable by making it HOISTABLE.

Three problems were observed when changed to HOISTABLE, and they are
addressed in this commit.
 1. A HOISTABLE IRInst is immutable once created. `SetOperand()`
    functions became unavailable. There were code that was modifying
    WitnessTable after creating it. This commit changed the behavior so
    that all operands are prepared before creating a new WitnessTable,
    and all operands are given to its constructor.
 1. WitnessTable started to have duplicated decorations and children.
    Due to the nature of HOISTABLE, when create a new WitnessTable, a
    pre-existing WitnessTable is returned and reused. There isn't an
    easy way to tell if the WitnessTable is a brand new or reused at the
    moment. This commit assumes that it is a brand new only when it has
    no decorations and chilrend after its creation.
 1. In `SimplifyIR()`, the order of children were slightly changed as a
    result of the optimization. The behavior was little different when
    WitnessTable became HOISTABLE. This resulted in an error where
    WitnessTable has a WitnessEntry pointing to an incorrect
    `IRSpecialize`. In order for it to function properly, `IRSpecialize`
    had to appear before `IRWitnessTable`.

Author: jkwak-work

source/slang/slang-ir-autodiff.cpp

@@ -3177,14 +3177,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-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

@@ -2950,8 +2950,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.cpp

@@ -1694,6 +1694,21 @@ void addHoistableInst(IRBuilder* builder, IRInst* inst)
     while (insertBeforeInst && insertBeforeInst->getOp() == kIROp_Param)
         insertBeforeInst = insertBeforeInst->getNextInst();
 
+    if (inst->getOp() == kIROp_WitnessTable)
+    {
+        // WitnessTable may reference specialize inst-s from its WitnessEntry
+        // children. In this case, specialize insts must be cloned before the
+        // WitnessTable.
+        //
+        for (IRInst* iter = insertBeforeInst; iter; )
+        {
+            bool isSpecialize = (iter->getOp() == kIROp_Specialize);
+            iter = iter->getNextInst();
+            if (isSpecialize)
+                insertBeforeInst = iter;
+        }
+    }
+
     // For instructions that will be placed at module scope,
     // we don't care about relative ordering, but for everything
     // else, we want to ensure that an instruction comes after

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

@@ -8016,27 +8016,38 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
                         irSatisfyingWitnessTable = subBuilder->createWitnessTable(
                             irWitnessTableBaseType,
                             irWitnessTable->getConcreteType());
-                        auto mangledName = getMangledNameForConformanceWitness(
-                            subContext->astBuilder,
-                            astReqWitnessTable->witnessedType,
-                            astReqWitnessTable->baseType);
-                        subBuilder->addExportDecoration(
-                            irSatisfyingWitnessTable,
-                            mangledName.getUnownedSlice());
-                        if (isExportedType(astReqWitnessTable->witnessedType))
+
+                        // TODO: When WitnessTable became HOISTABLE, we needed a way to avoid
+                        // adding the same decorations and childrens redundantly.
+                        // There isn't an easy way to tell if the returned WitnessTable is a
+                        // brand new or a found one from the existing pool.
+                        // The code below assumes that when there are any decoration or child,
+                        // it is a pre-existed one.
+                        //
+                        if (irSatisfyingWitnessTable->getFirstDecorationOrChild() == nullptr)
                         {
-                            subBuilder->addHLSLExportDecoration(irSatisfyingWitnessTable);
-                            subBuilder->addKeepAliveDecoration(irSatisfyingWitnessTable);
-                        }
+                            auto mangledName = getMangledNameForConformanceWitness(
+                                subContext->astBuilder,
+                                astReqWitnessTable->witnessedType,
+                                astReqWitnessTable->baseType);
+                            subBuilder->addExportDecoration(
+                                irSatisfyingWitnessTable,
+                                mangledName.getUnownedSlice());
+                            if (isExportedType(astReqWitnessTable->witnessedType))
+                            {
+                                subBuilder->addHLSLExportDecoration(irSatisfyingWitnessTable);
+                                subBuilder->addKeepAliveDecoration(irSatisfyingWitnessTable);
+                            }
 
-                        // Recursively lower the sub-table.
-                        lowerWitnessTable(
-                            subContext,
-                            astReqWitnessTable,
-                            irSatisfyingWitnessTable,
-                            mapASTToIRWitnessTable);
+                            // Recursively lower the sub-table.
+                            lowerWitnessTable(
+                                subContext,
+                                astReqWitnessTable,
+                                irSatisfyingWitnessTable,
+                                mapASTToIRWitnessTable);
 
-                        irSatisfyingWitnessTable->moveToEnd();
+                            irSatisfyingWitnessTable->moveToEnd();
+                        }
                     }
                     irSatisfyingVal = irSatisfyingWitnessTable;
                 }
@@ -8145,58 +8156,74 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         //
         auto irWitnessTableBaseType = lowerType(subContext, superType);
 
-        // Create the IR-level witness table
-        auto irWitnessTable = subBuilder->createWitnessTable(irWitnessTableBaseType, nullptr);
-
-        // Register the value now, rather than later, to avoid any possible infinite recursion.
+        // 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(irWitnessTable)));
+            LoweredValInfo::simple(findOuterMostGeneric(subBuilder->getInsertLoc().getParent())));
 
         auto irSubType = lowerType(subContext, subType);
-        irWitnessTable->setConcreteType(irSubType);
 
-        // TODO(JS):
-        // Should the mangled name take part in obfuscation if enabled?
+        // Create the IR-level witness table
+        auto irWitnessTable = subBuilder->createWitnessTable(irWitnessTableBaseType, irSubType);
+
+        // TODO: When WitnessTable became HOISTABLE, we needed a way to avoid
+        // adding the same decorations and childrens redundantly.
+        // There isn't an easy way to tell if the returned WitnessTable is a
+        // brand new or a found one from the existing pool.
+        // The code below assumes that when there are any decoration or child,
+        // it is a pre-existed one.
+        //
+        if (irWitnessTable->getFirstDecorationOrChild() == nullptr)
+        {
+            // Override with the correct witness-table
+            context->setGlobalValue(
+                inheritanceDecl,
+                LoweredValInfo::simple(findOuterMostGeneric(irWitnessTable)));
 
-        addLinkageDecoration(
-            context,
-            irWitnessTable,
-            inheritanceDecl,
-            mangledName.getUnownedSlice());
+            // TODO(JS):
+            // Should the mangled name take part in obfuscation if enabled?
 
-        // If the witness table is for a COM interface, always keep it alive.
-        if (irWitnessTableBaseType->findDecoration<IRComInterfaceDecoration>())
-        {
-            subBuilder->addHLSLExportDecoration(irWitnessTable);
-        }
+            addLinkageDecoration(
+                context,
+                irWitnessTable,
+                inheritanceDecl,
+                mangledName.getUnownedSlice());
 
-        for (auto mod : parentDecl->modifiers)
-        {
-            if (as<HLSLExportModifier>(mod))
+            // If the witness table is for a COM interface, always keep it alive.
+            if (irWitnessTableBaseType->findDecoration<IRComInterfaceDecoration>())
             {
                 subBuilder->addHLSLExportDecoration(irWitnessTable);
-                subBuilder->addKeepAliveDecoration(irWitnessTable);
             }
-            else if (as<AutoDiffBuiltinAttribute>(mod))
+
+            for (auto mod : parentDecl->modifiers)
             {
-                subBuilder->addAutoDiffBuiltinDecoration(irWitnessTable);
+                if (as<HLSLExportModifier>(mod))
+                {
+                    subBuilder->addHLSLExportDecoration(irWitnessTable);
+                    subBuilder->addKeepAliveDecoration(irWitnessTable);
+                }
+                else if (as<AutoDiffBuiltinAttribute>(mod))
+                {
+                    subBuilder->addAutoDiffBuiltinDecoration(irWitnessTable);
+                }
             }
-        }
 
-        // Make sure that all the entries in the witness table have been filled in,
-        // including any cases where there are sub-witness-tables for conformances
-        bool isExplicitExtern = false;
-        if (!isImportedDecl(context, parentDecl, isExplicitExtern))
-        {
-            Dictionary<WitnessTable*, IRWitnessTable*> mapASTToIRWitnessTable;
-            lowerWitnessTable(
-                subContext,
-                inheritanceDecl->witnessTable,
-                irWitnessTable,
-                mapASTToIRWitnessTable);
+            // Make sure that all the entries in the witness table have been filled in,
+            // including any cases where there are sub-witness-tables for conformances
+            bool isExplicitExtern = false;
+            if (!isImportedDecl(context, parentDecl, isExplicitExtern))
+            {
+                Dictionary<WitnessTable*, IRWitnessTable*> mapASTToIRWitnessTable;
+                lowerWitnessTable(
+                    subContext,
+                    inheritanceDecl->witnessTable,
+                    irWitnessTable,
+                    mapASTToIRWitnessTable);
+            }
+            irWitnessTable->moveToEnd();
         }
-        irWitnessTable->moveToEnd();
 
         return LoweredValInfo::simple(
             finishOuterGenerics(subBuilder, irWitnessTable, outerGeneric));

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;
+}