Repository: shader-slang/slang
File: source/slang/slang-ir-specialize.cpp (C++)

Code before:
                    // Recursively evaluate the value
                    return tryEvaluateLinkTimeConstant(val, builder);
                }
            }
            break;
        }

        // Can't evaluate, return original
        return inst;
    }

    // Now that we know when we can specialize a generic, and how
    // to do it, we can write a subroutine that takes a
    // `specialize(g, a, b, c, ...)` instruction and performs
    // specialization if it is possible.
    //
    bool maybeSpecializeGeneric(IRSpecialize* specInst)
    {
        // First, try to evaluate any link-time constant expressions in the arguments.
        // This allows extern const overrides to be applied before specialization.
        IRBuilder evalBuilderStorage(module);
        auto evalBuilder = &evalBuilderStorage;
        for (UInt i = 1; i < specInst->getOperandCount(); ++i) // Skip operand 0 (the base)
        {
            auto arg = specInst->getOperand(i);
            auto evaluatedArg = tryEvaluateLinkTimeConstant(arg, evalBuilder);
            if (evaluatedArg != arg)
            {
                specInst->setOperand(i, evaluatedArg);
            }
        }

        // We will only attempt to specialize when all of the
        // operands to the `speicalize(...)` instruction are
        // themselves fully specialized.
        //
        if (!areAllOperandsFullySpecialized(specInst))
            return false;

        // The invariant that the arguments are fully specialized
        // should mean that `a, b, c, ...` are in a form that
        // we can work with, but it does *not* guarantee
        // that the `g` operand is something we can work with.
        //
        // We can only perform specialization in the case where
        // the base `g` is a known `generic` instruction.
        //
        auto baseVal = specInst->getBase();
        auto genericVal = as<IRGeneric>(baseVal);
        if (!genericVal)
            return false;

Review comment (bug):
This feels the wrong place to add constant folding. In `specializeGeneric`, there is already a call to `tryConstantFoldInst`, which should be enough. I don't think we need all these `tryEvaluateLinkTimeConstant` here.

Code after:
                        }

                        auto specDiffFunc = as<IRSpecialize>(decor->getOperand(0));

                        // If the base is specialized, the JVP version must be also be a specialized
                        // generic.
                        //
                        SLANG_RELEASE_ASSERT(specDiffFunc);

                        // Build specialization arguments from specInst.
                        // Note that if we've reached this point, we can safely assume
                        // that our args are fully specialized/concrete.
                        //
                        UCount argCount = specInst->getArgCount();
                        ShortList<IRInst*> args;
                        for (UIndex ii = 0; ii < argCount; ii++)
                            args.add(specInst->getArg(ii));

                        IRBuilder builder(module);

                        // Specialize the custom derivative function type with the original
                        // arguments.
                        builder.setInsertInto(module);
                        auto newDiffFuncType = builder.emitSpecializeInst(
                            builder.getTypeKind(),
                            specDiffFunc->getBase()->getDataType(),
                            argCount,
                            args.getArrayView().getBuffer());

                        // Specialize the custom derivative function with the original arguments.
                        builder.setInsertBefore(specInst);
                        auto newDiffFunc = builder.emitSpecializeInst(
                            (IRType*)newDiffFuncType,
                            specDiffFunc->getBase(),
                            argCount,
                            args.getArrayView().getBuffer());

                        // Add the new spec insts to the list so they get specialized with
                        // the usual logic.
                        //
                        addToWorkList(newDiffFuncType);
                        addToWorkList(newDiffFunc);

                        builder.addDecoration(specInst, decor->getOp(), newDiffFunc);

                        specialized = true;
                    }
                    if (specialized)
                        return true;
                }
            }