Add API to control interface specialization. (shader-slang#1925)

Author: csyonghe

slang.h

@@ -3050,6 +3050,7 @@ namespace slang
     typedef ISlangBlob IBlob;
 
     struct IComponentType;
+    struct ITypeConformance;
     struct IGlobalSession;
     struct IModule;
     struct ISession;
@@ -4023,6 +4024,32 @@ namespace slang
             */
         virtual SLANG_NO_THROW SlangResult SLANG_MCALL createCompileRequest(
             SlangCompileRequest**   outCompileRequest) = 0;
+
+        
+            /** Creates a `IComponentType` that represents a type's conformance to an interface.
+                The retrieved `ITypeConformance` objects can be included in a composite `IComponentType`
+                to explicitly specify which implementation types should be included in the final compiled
+                code. For example, if an module defines `IMaterial` interface and `AMaterial`,
+                `BMaterial`, `CMaterial` types that implements the interface, the user can exclude
+                `CMaterial` implementation from the resulting shader code by explcitly adding
+                `AMaterial:IMaterial` and `BMaterial:IMaterial` conformances to a composite
+                `IComponentType` and get entry point code from it. The resulting code will not have
+                anything related to `CMaterial` in the dynamic dispatch logic. If the user does not
+                explicitly include any `TypeConformances` to an interface type, all implementations to
+                that interface will be included by default. By linking a `ITypeConformance`, the user is
+                also given the opportunity to specify the dispatch ID of the implementation type. If
+                `conformanceIdOverride` is -1, there will be no override behavior and Slang will
+                automatically assign IDs to implementation types. The automatically assigned IDs can be
+                queried via `ISession::getTypeConformanceWitnessSequentialID`.
+
+                Returns SLANG_OK if succeeds, or SLANG_FAIL if `type` does not conform to `interfaceType`.
+            */
+        virtual SLANG_NO_THROW SlangResult SLANG_MCALL createTypeConformanceComponentType(
+            slang::TypeReflection* type,
+            slang::TypeReflection* interfaceType,
+            ITypeConformance** outConformance,
+            SlangInt conformanceIdOverride,
+            ISlangBlob** outDiagnostics) = 0;
     };
 
     #define SLANG_UUID_ISession ISession::getTypeGuid()
@@ -4204,6 +4231,12 @@ namespace slang
 
     #define SLANG_UUID_IEntryPoint IEntryPoint::getTypeGuid()
 
+    struct ITypeConformance : public IComponentType
+    {
+        SLANG_COM_INTERFACE(0x73eb3147, 0xe544, 0x41b5, { 0xb8, 0xf0, 0xa2, 0x44, 0xdf, 0x21, 0x94, 0xb })
+    };
+    #define SLANG_UUID_ITypeConformance ITypeConformance::getTypeGuid()
+
         /** A module is the granularity of shader code compilation and loading.
 
         In most cases a module corresponds to a single compile "translation unit."

source/slang/slang-compiler.cpp

@@ -296,6 +296,92 @@ namespace Slang
         return empty;
     }
 
+    TypeConformance::TypeConformance(
+        Linkage* linkage,
+        SubtypeWitness* witness,
+        Int confomrmanceIdOverride,
+        DiagnosticSink* sink)
+        : ComponentType(linkage)
+        , m_subtypeWitness(witness)
+        , m_conformanceIdOverride(confomrmanceIdOverride)
+    {
+        addDepedencyFromWitness(witness);
+        m_irModule = generateIRForTypeConformance(this, m_conformanceIdOverride, sink);
+    }
+
+    void TypeConformance::addDepedencyFromWitness(SubtypeWitness* witness)
+    {
+        if (auto declaredWitness = as<DeclaredSubtypeWitness>(witness))
+        {
+            auto declModule = getModule(declaredWitness->declRef.getDecl());
+            m_moduleDependency.addDependency(declModule);
+            m_pathDependency.addDependency(declModule);
+            if (m_requirementSet.Add(declModule))
+            {
+                m_requirements.add(declModule);
+            }
+            // TODO: handle the specialization arguments in declaredWitness->declRef.substitutions.
+        }
+        else if (auto transitiveWitness = as<TransitiveSubtypeWitness>(witness))
+        {
+            addDepedencyFromWitness(transitiveWitness->midToSup);
+            addDepedencyFromWitness(transitiveWitness->subToMid);
+        }
+        else if (auto conjunctionWitness = as<ConjunctionSubtypeWitness>(witness))
+        {
+            auto left = as<SubtypeWitness>(conjunctionWitness->leftWitness);
+            if (left)
+                addDepedencyFromWitness(left);
+            auto right = as<SubtypeWitness>(conjunctionWitness->rightWitness);
+            if (right)
+                addDepedencyFromWitness(right);
+        }
+    }
+
+    ISlangUnknown* TypeConformance::getInterface(const Guid& guid)
+    {
+        if (guid == slang::ITypeConformance::getTypeGuid())
+            return static_cast<slang::ITypeConformance*>(this);
+
+        return Super::getInterface(guid);
+    }
+
+    List<Module*> const& TypeConformance::getModuleDependencies()
+    {
+        return m_moduleDependency.getModuleList();
+    }
+
+    List<String> const& TypeConformance::getFilePathDependencies()
+    {
+        return m_pathDependency.getFilePathList();
+    }
+
+    Index TypeConformance::getRequirementCount() { return m_requirements.getCount(); }
+
+    RefPtr<ComponentType> TypeConformance::getRequirement(Index index)
+    {
+        return m_requirements[index];
+    }
+
+    void TypeConformance::acceptVisitor(
+        ComponentTypeVisitor* visitor,
+        ComponentType::SpecializationInfo* specializationInfo)
+    {
+        SLANG_UNUSED(specializationInfo);
+        visitor->visitTypeConformance(this);
+    }
+
+    RefPtr<ComponentType::SpecializationInfo> TypeConformance::_validateSpecializationArgsImpl(
+        SpecializationArg const* args,
+        Index argCount,
+        DiagnosticSink* sink)
+    {
+        SLANG_UNUSED(args);
+        SLANG_UNUSED(argCount);
+        SLANG_UNUSED(sink);
+        return nullptr;
+    }
+
     //
 
     Profile Profile::lookUp(UnownedStringSlice const& name)

source/slang/slang-compiler.h

@@ -850,6 +850,125 @@ namespace Slang
         Profile m_profile;
     };
 
+    class TypeConformance
+        : public ComponentType
+        , public slang::ITypeConformance
+    {
+        typedef ComponentType Super;
+
+    public:
+        SLANG_REF_OBJECT_IUNKNOWN_ALL
+
+        ISlangUnknown* getInterface(const Guid& guid);
+
+        TypeConformance(
+            Linkage* linkage,
+            SubtypeWitness* witness,
+            Int confomrmanceIdOverride,
+            DiagnosticSink* sink);
+
+        // Forward `IComponentType` methods
+
+        SLANG_NO_THROW slang::ISession* SLANG_MCALL getSession() SLANG_OVERRIDE
+        {
+            return Super::getSession();
+        }
+
+        SLANG_NO_THROW slang::ProgramLayout* SLANG_MCALL
+            getLayout(SlangInt targetIndex, slang::IBlob** outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::getLayout(targetIndex, outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL getEntryPointCode(
+            SlangInt entryPointIndex,
+            SlangInt targetIndex,
+            slang::IBlob** outCode,
+            slang::IBlob** outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::getEntryPointCode(entryPointIndex, targetIndex, outCode, outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL specialize(
+            slang::SpecializationArg const* specializationArgs,
+            SlangInt specializationArgCount,
+            slang::IComponentType** outSpecializedComponentType,
+            ISlangBlob** outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::specialize(
+                specializationArgs,
+                specializationArgCount,
+                outSpecializedComponentType,
+                outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL link(
+            slang::IComponentType** outLinkedComponentType,
+            ISlangBlob** outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::link(outLinkedComponentType, outDiagnostics);
+        }
+
+        SLANG_NO_THROW SlangResult SLANG_MCALL getEntryPointHostCallable(
+            int entryPointIndex,
+            int targetIndex,
+            ISlangSharedLibrary** outSharedLibrary,
+            slang::IBlob** outDiagnostics) SLANG_OVERRIDE
+        {
+            return Super::getEntryPointHostCallable(
+                entryPointIndex, targetIndex, outSharedLibrary, outDiagnostics);
+        }
+
+        List<Module*> const& getModuleDependencies() SLANG_OVERRIDE;
+        List<String> const& getFilePathDependencies() SLANG_OVERRIDE;
+
+        SLANG_NO_THROW Index SLANG_MCALL getSpecializationParamCount() SLANG_OVERRIDE { return 0; }
+
+        /// Get the existential type parameter at `index`.
+        SpecializationParam const& getSpecializationParam(Index /*index*/) SLANG_OVERRIDE
+        {
+            static SpecializationParam emptyParam;
+            return emptyParam;
+        }
+
+        Index getRequirementCount() SLANG_OVERRIDE;
+        RefPtr<ComponentType> getRequirement(Index index) SLANG_OVERRIDE;
+        Index getEntryPointCount() SLANG_OVERRIDE { return 0; };
+        RefPtr<EntryPoint> getEntryPoint(Index index) SLANG_OVERRIDE
+        {
+            SLANG_UNUSED(index);
+            return nullptr;
+        }
+        String getEntryPointMangledName(Index /*index*/) SLANG_OVERRIDE { return ""; }
+
+        Index getShaderParamCount() SLANG_OVERRIDE { return 0; }
+        ShaderParamInfo getShaderParam(Index index) SLANG_OVERRIDE
+        {
+            SLANG_UNUSED(index);
+            return ShaderParamInfo();
+        }
+
+        SubtypeWitness* getSubtypeWitness() { return m_subtypeWitness; }
+        IRModule* getIRModule() { return m_irModule.Ptr(); }
+    protected:
+        void acceptVisitor(ComponentTypeVisitor* visitor, SpecializationInfo* specializationInfo)
+            SLANG_OVERRIDE;
+
+        RefPtr<SpecializationInfo> _validateSpecializationArgsImpl(
+            SpecializationArg const* args,
+            Index argCount,
+            DiagnosticSink* sink) SLANG_OVERRIDE;
+    private:
+        SubtypeWitness* m_subtypeWitness;
+        ModuleDependencyList m_moduleDependency;
+        FilePathDependencyList m_pathDependency;
+        List<RefPtr<Module>> m_requirements;
+        HashSet<Module*> m_requirementSet;
+        RefPtr<IRModule> m_irModule;
+        Int m_conformanceIdOverride;
+        void addDepedencyFromWitness(SubtypeWitness* witness);
+    };
+
     enum class PassThroughMode : SlangPassThroughIntegral
     {
         None = SLANG_PASS_THROUGH_NONE,	                    ///< don't pass through: use Slang compiler
@@ -1319,6 +1438,12 @@ namespace Slang
             slang::TypeReflection* type,
             slang::TypeReflection* interfaceType,
             uint32_t*              outId) override;
+        SLANG_NO_THROW SlangResult SLANG_MCALL createTypeConformanceComponentType(
+            slang::TypeReflection* type,
+            slang::TypeReflection* interfaceType,
+            slang::ITypeConformance** outConformance,
+            SlangInt conformanceIdOverride,
+            ISlangBlob** outDiagnostics) override;
         SLANG_NO_THROW SlangResult SLANG_MCALL createCompileRequest(
             SlangCompileRequest**   outCompileRequest) override;
 
@@ -1756,6 +1881,7 @@ namespace Slang
         virtual void visitModule(Module* module, Module::ModuleSpecializationInfo* specializationInfo) = 0;
         virtual void visitComposite(CompositeComponentType* composite, CompositeComponentType::CompositeSpecializationInfo* specializationInfo) = 0;
         virtual void visitSpecialized(SpecializedComponentType* specialized) = 0;
+        virtual void visitTypeConformance(TypeConformance* conformance) = 0;
 
     protected:
         // These helpers can be used to recurse into the logical children of a

source/slang/slang-ir-link.cpp

@@ -274,7 +274,7 @@ IRInst* IRSpecContext::maybeCloneValue(IRInst* originalValue)
 
     default:
         {
-            // In the deafult case, assume that we have some sort of "hoistable"
+            // In the default case, assume that we have some sort of "hoistable"
             // instruction that requires us to create a clone of it.
 
             UInt argCount = originalValue->getOperandCount();
@@ -439,6 +439,8 @@ static void cloneExtraDecorations(
 
             case kIROp_BindExistentialSlotsDecoration:
             case kIROp_LayoutDecoration:
+            case kIROp_PublicDecoration:
+            case kIROp_SequentialIDDecoration:
                 if(!clonedInst->findDecorationImpl(decoration->getOp()))
                 {
                     cloneInst(context, builder, decoration);

source/slang/slang-ir-specialize.cpp

@@ -2158,6 +2158,11 @@ IRInst* specializeGenericImpl(
             if( auto returnValInst = as<IRReturnVal>(ii) )
             {
                 auto specializedVal = findCloneForOperand(&env, returnValInst->getVal());
+
+                // Clone decorations on the orignal `specialize` inst over to the newly specialized
+                // value.
+                cloneInstDecorationsAndChildren(
+                    &env, &sharedBuilderStorage, specializeInst, specializedVal);
                 return specializedVal;
             }