Initial support for dynamic dispatch using "tagged union" types (shader-slang#772)

* Initial support for dynamic dispatch using "tagged union" types

Suppose a user declares some generic shader code, like the following:

```hlsl
interface IFrobnicator { ... }
type_param T : IFrobincator;
ParameterBlock<T : IFrobnicator> gFrobnicator;
...
gFrobincator.frobnicate(value);
```

and then they have some concrete implementations of the required interface:

```hlsl
struct A : IFrobnicator { ... }
struct B : IFrobnicator { ... }
```

The current Slang compiler allows them to generate  distinct compiled kernels for the case of `T=A` and the case of `T=B`. This means that the decision of which implementation to use must be made at or before the time when a shader gets bound in the application.

This change adds a new ability where the Slang compiler can generate code to handle the case where `T` might be *either* `A` or `B`, and which case it is will be determined dynamically at runtime. This means a single compiled kernel can handle both cases, and the decision about which code path to run can be made any time before the shader executes.

This new option is supported by defining a *tagged union* type. Via the API, the user specifies that `T` should be specialized to `__TaggedUnion(A,B)` (the double underscore indicates that this is an experimental and unsupported feature at present). We refer to the types `A` and `B` here as the "case" types of the tagged union. Conceptually, the compiler synthesizes a type something like:

```hlsl
struct TU { union { A a; B b; } payload; uint tag; }
```

The user can then allocate a constant buffer to hold their tagged union type, and when they pick a concrete type to use (say `B`), they fill in the first `sizeof(B)` bytes of their buffer with data describing a `B` instance, and then set the `tag` field to the appopriate 0-based index of the case type they chose (in this case the `B` case gets the tag value `1`).

Actually implementing tagged unions takes a few main steps:

* Type parsing was extended to special-case `__TaggedUnion` as a contextual keyword. This is really only intended to be used when parsing types from the API or command-line, and Bad Things are likely to happen if a user ever puts it directly in their code. Eventually construction of tagged unions should be an API feature and not part of the language syntax.

* Semantic checking was extended to recognize that a tagged union like `__TaggedUnion(A,B)` shoud support an interface like `IFrobnicator` whenever all of the case types suport it, as long as the interface is "safe" for use with tagged unions (which means it doesn't use a few of the advancd langauge features like associated types).

* The IR was extended with instructions to represent tagged union types and to extract their tag and the payload for the different cases as needed.

* IR generation was extended to synthesize implementations of interface methods for any interface that a tagged union needs to support. Right now the implementation is simplistic and only handles simple method requirements, which it does by emitting a `switch` instruction to pick between the different cases.

* A new IR pass was introduced to "desugar" any tagged union types used in the code. The downstream HLSL and GLSL compilers don't support `union`s, so we have to instead emit a tagged union as a "bag of bits" and implement loading the data for particular cases from it manually.

* Final code emit mostly Just Works after the above steps, but we had to introduce an explicit IR instruction for bit-casting to handle the output of the desugaring pass.

There are a bunch of gaps and caveats in this implementation, but that seems reasonable for something that is an experimental feature. The various `TODO` comments and assertion failures in unimplemented cases are intended, so that this work can be checked in even if it isn't feature-complete.

* fixup: missing files

* fixup: typos

Author: Tim Foley

source/slang/check.cpp

@@ -9,6 +9,8 @@
 
 namespace Slang
 {
+    RefPtr<TypeType> getTypeType(
+        Type* type);
 
         /// Should the given `decl` nested in `parentDecl` be treated as a static rather than instance declaration?
     bool isEffectivelyStatic(
@@ -677,7 +679,7 @@ namespace Slang
                     auto baseExprType = baseExpr->type.type;
                     RefPtr<SharedTypeExpr> baseTypeExpr = new SharedTypeExpr();
                     baseTypeExpr->base.type = baseExprType;
-                    baseTypeExpr->type = new TypeType(baseExprType);
+                    baseTypeExpr->type.type = getTypeType(baseExprType);
 
                     auto expr = new StaticMemberExpr();
                     expr->loc = loc;
@@ -2071,9 +2073,7 @@ namespace Slang
         {
             RefPtr<TypeCastExpr> castExpr = createImplicitCastExpr();
 
-            auto typeType = new TypeType();
-            typeType->setSession(getSession());
-            typeType->type = toType;
+            auto typeType = getTypeType(toType);
 
             auto typeExpr = new SharedTypeExpr();
             typeExpr->type.type = typeType;
@@ -5547,6 +5547,60 @@ namespace Slang
             return witness;
         }
 
+            /// Is the given interface one that a tagged-union type can conform to?
+            ///
+            /// If a tagged union type `__TaggedUnion(A,B)` is going to be
+            /// plugged in for a type parameter `T : IFoo` then we need to
+            /// be sure that the interface `IFoo` doesn't have anything
+            /// that could lead to unsafe/unsound behavior. This function
+            /// checks that all the requirements on the interfaceare safe ones.
+            ///
+        bool isInterfaceSafeForTaggedUnion(
+            DeclRef<InterfaceDecl> interfaceDeclRef)
+        {
+            for( auto memberDeclRef : getMembers(interfaceDeclRef) )
+            {
+                if(!isInterfaceRequirementSafeForTaggedUnion(interfaceDeclRef, memberDeclRef))
+                    return false;
+            }
+
+            return true;
+        }
+
+            /// Is the given interface requirement one that a tagged-union type can satisfy?
+            ///
+            /// Unsafe requirements include any `static` requirements,
+            /// any associated types, and also any requirements that make
+            /// use of the `This` type (once we support it).
+            ///
+        bool isInterfaceRequirementSafeForTaggedUnion(
+            DeclRef<InterfaceDecl>  interfaceDeclRef,
+            DeclRef<Decl>           requirementDeclRef)
+        {
+            if(auto callableDeclRef = requirementDeclRef.As<CallableDecl>())
+            {
+                // A `static` method requirement can't be satisfied by a
+                // tagged union, because there is no tag to dispatch on.
+                //
+                if(requirementDeclRef.getDecl()->HasModifier<HLSLStaticModifier>())
+                    return false;
+
+                // TODO: We will eventually want to check that any callable
+                // requirements do not use the `This` type or any associated
+                // types in ways that could lead to errors.
+                //
+                // For now we are disallowing interfaces that have associated
+                // types completely, and we haven't implemented the `This`
+                // type, so we should be safe.
+
+                return true;
+            }
+            else
+            {
+                return false;
+            }
+        }
+
         bool doesTypeConformToInterfaceImpl(
             RefPtr<Type>            originalType,
             RefPtr<Type>            type,
@@ -5661,6 +5715,69 @@ namespace Slang
                     }
                 }
             }
+            else if(auto taggedUnionType = type->As<TaggedUnionType>())
+            {
+                // A tagged union type conforms to an interface if all of
+                // the constituent types in the tagged union conform.
+                //
+                // We will iterate over the "case" types in the tagged
+                // union, and check if they conform to the interface.
+                // Along the way we will collect the conformance witness
+                // values *if* we are being asked to produce a witness
+                // value for the tagged union itself (that is, if
+                // `outWitness` is non-null).
+                //
+                List<RefPtr<Val>> caseWitnesses;
+                for(auto caseType : taggedUnionType->caseTypes)
+                {
+                    RefPtr<Val> caseWitness;
+
+                    if(!doesTypeConformToInterfaceImpl(
+                        caseType,
+                        caseType,
+                        interfaceDeclRef,
+                        outWitness ? &caseWitness : nullptr,
+                        nullptr))
+                    {
+                        return false;
+                    }
+
+                    if(outWitness)
+                    {
+                        caseWitnesses.Add(caseWitness);
+                    }
+                }
+
+                // We also need to validate the requirements on
+                // the interface to make sure that they are suitable for
+                // use with a tagged-union type.
+                //
+                // For example, if the interface includes a `static` method
+                // (which can therefore be called without a particular instance),
+                // then we wouldn't know what implementation of that method
+                // to use because there is no tag value to dispatch on.
+                //
+                // We will start out being conservative about what we accept
+                // here, just to keep things simple.
+                //
+                if(!isInterfaceSafeForTaggedUnion(interfaceDeclRef))
+                    return false;
+
+                // If we reach this point then we have a concrete
+                // witness for each of the case types, and that is
+                // enough to build a witness for the tagged union.
+                //
+                if(outWitness)
+                {
+                    RefPtr<TaggedUnionSubtypeWitness> taggedUnionWitness = new TaggedUnionSubtypeWitness();
+                    taggedUnionWitness->sub = taggedUnionType;
+                    taggedUnionWitness->sup = DeclRefType::Create(getSession(), interfaceDeclRef);
+                    taggedUnionWitness->caseWitnesses.SwapWith(caseWitnesses);
+
+                    *outWitness = taggedUnionWitness;
+                }
+                return true;
+            }
 
             // default is failure
             return false;
@@ -8090,6 +8207,23 @@ namespace Slang
             return expr;
         }
 
+        RefPtr<Expr> visitTaggedUnionTypeExpr(TaggedUnionTypeExpr* expr)
+        {
+            // We have an expression of the form `__TaggedUnion(A, B, ...)`
+            // which will evaluate to a tagged-union type over `A`, `B`, etc.
+            //
+            RefPtr<TaggedUnionType> type = new TaggedUnionType();
+            expr->type = QualType(getTypeType(type));
+
+            for( auto& caseTypeExpr : expr->caseTypes )
+            {
+                caseTypeExpr = CheckProperType(caseTypeExpr);
+                type->caseTypes.Add(caseTypeExpr.type);
+            }
+
+            return expr;
+        }
+
 
 
 
@@ -9039,7 +9173,7 @@ namespace Slang
             scopesToTry.Add(module->moduleDecl->scope);
 
         List<RefPtr<Type>> globalGenericArgs;
-        for (auto name : entryPoint->genericParameterTypeNames)
+        for (auto name : entryPoint->genericArgStrings)
         {
             // parse type name
             RefPtr<Type> type;
@@ -9059,6 +9193,25 @@ namespace Slang
                 return;
             }
 
+            // The following is a bit of a hack.
+            //
+            // Back-end code generation relies on us having computed layouts for all tagged
+            // unions that end up being used in the code, which means we need a way to find
+            // all such types that get used in a module (and the stuff it imports).
+            //
+            // The Right Way to handle this would probably be to have each `ModuleDecl` track
+            // any tagged union types that get created in the context of that module, and
+            // then combine those lists later.
+            //
+            // For now we are assuming a tagged union type only comes into existence
+            // as a (top-level) argument for a generic type parameter, so that we
+            // can check for them here and cache them on the entry point request.
+            //
+            if( auto taggedUnionType = type->As<TaggedUnionType>() )
+            {
+                entryPoint->taggedUnionTypes.Add(taggedUnionType);
+            }
+
             globalGenericArgs.Add(type);
         }
 

source/slang/compiler.h

@@ -122,9 +122,8 @@ namespace Slang
         // The name of the entry point function (e.g., `main`)
         Name* name;
 
-        // The type names we want to substitute into the
-        // global generic type parameters
-        List<String> genericParameterTypeNames;
+            /// Source code for the generic arguments to use for the generic parameters of the entry point.
+        List<String> genericArgStrings;
 
         // The profile that the entry point will be compiled for
         // (this is a combination of the target stage, and also
@@ -156,6 +155,9 @@ namespace Slang
         RefPtr<FuncDecl> decl;
 
         RefPtr<Substitutions> globalGenericSubst;
+
+            /// Any tagged union types that were referenced by the generic arguments of the entry point.
+        List<RefPtr<TaggedUnionType>> taggedUnionTypes;
     };
 
     enum class PassThroughMode : SlangPassThrough

source/slang/diagnostic-defs.h

@@ -347,7 +347,7 @@ DIAGNOSTIC(38102, Error, initializerNotInsideType, "an 'init' declaration is onl
 DIAGNOSTIC(38102, Error, accessorMustBeInsideSubscriptOrProperty, "an accessor declaration is only allowed inside a subscript or property declaration")
 
 DIAGNOSTIC(38020, Error, mismatchEntryPointTypeArgument, "expecting $0 entry-point type arguments, provided $1.")
-DIAGNOSTIC(38021, Error, typeArgumentDoesNotConformToInterface, "type argument `$1` for generic parameter `$0` does not conform to interface `$1`.")
+DIAGNOSTIC(38021, Error, typeArgumentDoesNotConformToInterface, "type argument `$1` for generic parameter `$0` does not conform to interface `$2`.")
 
 DIAGNOSTIC(38022, Error, cannotSpecializeGlobalGenericToItself, "the global type parameter '$0' cannot be specialized to itself")
 DIAGNOSTIC(38023, Error, cannotSpecializeGlobalGenericToAnotherGenericParam, "the global type parameter '$0' cannot be specialized using another global type parameter ('$1')")

source/slang/emit.cpp

@@ -12,6 +12,7 @@
 #include "ir-specialize.h"
 #include "ir-specialize-resources.h"
 #include "ir-ssa.h"
+#include "ir-union.h"
 #include "ir-validate.h"
 #include "legalize-types.h"
 #include "lower-to-ir.h"
@@ -3922,13 +3923,100 @@ struct EmitVisitor
             }
             break;
 
+        case kIROp_BitCast:
+            {
+                // TODO: we can simplify the logic for arbitrary bitcasts
+                // by always bitcasting the source to a `uint*` type (if it
+                // isn't already) and then bitcasting that to the destination
+                // type (if it isn't already `uint*`.
+                //
+                // For now we are assuming the source type is *already*
+                // a `uint*` type of the appropriate size.
+                //
+//                auto fromType = extractBaseType(inst->getOperand(0)->getDataType());
+                auto toType = extractBaseType(inst->getDataType());
+                switch(getTarget(ctx))
+                {
+                case CodeGenTarget::GLSL:
+                    switch(toType)
+                    {
+                    default:
+                        emit("/* unhandled */");
+                        break;
+
+                    case BaseType::UInt:
+                        break;
+
+                    case BaseType::Int:
+                        emitIRType(ctx, inst->getDataType());
+                        break;
+
+                    case BaseType::Float:
+                        emit("uintBitsToFloat(");
+                        break;
+                    }
+                    break;
+
+                case CodeGenTarget::HLSL:
+                    switch(toType)
+                    {
+                    default:
+                        emit("/* unhandled */");
+                        break;
+
+                    case BaseType::UInt:
+                        break;
+                    case BaseType::Int:
+                        emit("(");
+                        emitIRType(ctx, inst->getDataType());
+                        emit(")");
+                        break;
+                    case BaseType::Float:
+                        emit("asfloat");
+                        break;
+                    }
+                    break;
+
+
+                default:
+                    SLANG_UNEXPECTED("unhandled codegen target");
+                    break;
+                }
+
+                emit("(");
+                emitIROperand(ctx, inst->getOperand(0), mode, kEOp_General);
+                emit(")");
+            }
+            break;
+
         default:
             emit("/* unhandled */");
             break;
         }
         maybeCloseParens(needClose);
     }
 
+    BaseType extractBaseType(IRType* inType)
+    {
+        auto type = inType;
+        for(;;)
+        {
+            if(auto irBaseType = as<IRBasicType>(type))
+            {
+                return irBaseType->getBaseType();
+            }
+            else if(auto vecType = as<IRVectorType>(type))
+            {
+                type = vecType->getElementType();
+                continue;
+            }
+            else
+            {
+                return BaseType::Void;
+            }
+        }
+    }
+
     void emitIRInst(
         EmitContext*    ctx,
         IRInst*         inst,
@@ -6565,6 +6653,14 @@ String emitEntryPoint(
 #endif
         validateIRModuleIfEnabled(compileRequest, irModule);
 
+        // Desguar any union types, since these will be illegal on
+        // various targets.
+        //
+        desugarUnionTypes(irModule);
+#if 0
+        dumpIRIfEnabled(compileRequest, irModule, "UNIONS DESUGARED");
+#endif
+        validateIRModuleIfEnabled(compileRequest, irModule);
 
 
         // Any code that makes use of existential (interface) types
@@ -6595,8 +6691,6 @@ String emitEntryPoint(
         //
         specializeGenerics(irModule);
 
-
-
         // Debugging code for IR transformations...
 #if 0
         dumpIRIfEnabled(compileRequest, irModule, "SPECIALIZED");

source/slang/expr-defs.h

@@ -193,3 +193,14 @@ RAW(
     DeclRef<VarDeclBase> declRef;
 )
 END_SYNTAX_CLASS()
+
+    /// A type expression of the form `__TaggedUnion(A, ...)`.
+    ///
+    /// An expression of this form will resolve to a `TaggedUnionType`
+    /// when checked.
+    ///
+SYNTAX_CLASS(TaggedUnionTypeExpr, Expr)
+RAW(
+    List<TypeExp> caseTypes;
+)
+END_SYNTAX_CLASS()