Add type-flow analysis pass for specialized dynamic dispatch

[saipraveenb25]

Fixes: #7782
Unblocks: #6486

Type Flow Analysis for Specialized Dynamic Dispatch

This PR implements a transformation pass for dynamic dispatch in the Slang compiler, enabling the lowering of dynamic instructions (such as LookupWitnessMethod, MakeExistential, ...) into concrete instructions specialized to the set of possible elements that can be determined at compile time.

This PR also removes the old generics lowering approach that relied on detecting all available conformances in the global scope and generating dispatch functions based on this information.

Minor Note: This PR introduces new IR op-codes, but does not change the max version since these IR ops are strictly internal. They are created and destroyed within the lifetime of the typeflow specialization pass.

Overview

The implementation introduces an IR transformation pass that converts runtime polymorphic instructions into specialized dispatch functions and union types through 3 stages:
i. Analysis: determining sets of possible types/tables/functions/etc.. for each dynamic instruction through a data-flow analysis pass,
ii. Specialization: replacing dynamic instructions into specialized versions that operate on the known sets via "set" instructions (e.g. TypeSet, SetTagType, etc..), and
iii. Lowering: lowering sets into dispatch functions (for sets of functions) and union types (for sets of types).

Analysis: Propagating sets of types, funcs, tables & generics using static analysis.

The analysis step proceeds similar to a traditional inter-procedural, context-sensitive data-flow analysis step.
The primary difference is that the states being tracked for each instruction is a not a set of values, but of global definitions (e.g. types, functions, witness tables, generics, etc..)

We maintain a map of (context, instruction) to typeFlowInfo for each instruction.

• context is either an IRFunc* or an IRSpecialize* inst and the instruction is any child inst of the context. context must be a global inst with concrete operands. In case of IRSpecialize*, the operands dictate the bindings for the generic parameters (these may also be sets)

• typeFlowInfo is an "effective type" that we track for each instruction that may be dynamic (or have dynamic component). Here are the set of possible states:

  • FuncSet,TypeSet,GenericSet,WitnessTableSet: Represents a set of functions/types/generics/tables.
  • ElementOfSetType(set): Represents an element of the given set. Used for instructions whose value can be one of the elements in the set (e.g. the result of a LookupWitnessMethod, ExtractExistentialWitnessTable, ExtractExistentialType, etc..)
  • TaggedUnionType(tableSet : WitnessTableSet, typeSet : TypeSet): Represents a tagged union of index into tables in tableSet and a value whose type may be any of the types in typeSet. Typically used for instructions that create existentials.
  • UntaggedUnionType(typeSet : TypeSet): Represents a union of values whose type may be any of the types in typeSet. Typically used for instructions that extract values from existentials.
  • PtrTypeBase(typeFlowInfo)/ArrayType(typeFlowInfo, ...): Additionally, we allow type-flow info insts to be nested into pointer and array types. This enables us to analyze (and optimize) loads and stores from local vars, structs, arrays, etc.
  • Special Elements: We define a few extra op-codes to represent certain sentinel values that require special treatment.
    • NoneTypeElement, NoneWitnessTableElement: Represents a 'none' sentinel value for types and witness tables respectively, to be used to represent the none case for OptionalType
    • UninitializedTypeElement(interfaceType), UninitializedWitnessTableElement(interfaceType): Represents an uninitialized value of the given type. Used to represent cases where a value may be uninitialized (e.g. result of a LoadFromUninitializedMemory). This is primarily used to detect dynamic dispatch on uninitialized values and diagnose appropriate errors.
    • UnboundedTypeElement(interfaceType), UnboundedWitnessTableElement(interfaceType): Represents an unbounded set of types that conform to the given interface type. This is used to represent cases where we cannot determine a finite set of possible types/tables (e.g. when we need to compile a function for true dynamic dispatch by accepting v-table pointers). Note that the elements can be combined with known elements to represent a hybrid specialized + dynamic case.
    • UnboundedFuncElement(funcType), UnboundedGenericElement(genericType): Similar to above, but for functions and generics respectively, the funcType and genericType represent the type signature to be used to call into the element.

• The fixed-point analysis iterates by adding the users of insts whose type-data has changed in each iteration until no more changes are made.

  • If arguments to a branch-inst are modified, the control-flow edge is added to the work queue.
  • If arguments to a 'call' inst are modified, the callee is added as an interprocedural edge to propagate information to the callee's parameters.
  • If a return value is modified, the 'back-edge' to all call sites is added to the work queue.
  • All of these 'work items' carry a context instruction. Inter-procedural edges carry both a source and destination context.

Specialization: Convert dynamic lookups to use tag operations on known sets

Transforms dynamic dispatch instructions into tag operations and lowers certain insts (e.g call) to use these tags where necessary. This pass also rewrites the types of instructions to use the set-based types.

We also introduce a new type:

• SetTagType(set): Represents a run-time identifier that can be used to select an element in set.

Any insts that were assigned an ElementOfSetType(set) during analysis will be specialized into insts that produce a SetTagType(set) to represent the run-time identifier, and any insts that were using the element will be specialized to accept the tag instead.
In some cases (like Call), this involves some re-writes to convert the tag into something callable.

Tag Operations:

• GetTagForSuperSet(tag), GetTagForSubSet(tag): Maps a tag for a sub-set to a tag for a super-set or vice-versa (the sets are encoded in the types of the get-tag inst and the operand)
• GetTagForMappedSet(tag, key): Maps a tag for a set to another set under a table lookup operation.
• GetTagForSpecializedSet(tag, specArgs...): Maps a tag for a generic set to a specialized set under specialization with the given arguments.
• GetTagFromSequentialID(id, interfaceType): Converts witness-table sequential IDs (which are global & may be externally defined) to set tags.
• GetSequentialIDFromTag(tag, interfaceType): Converts witness-table sequential IDs (which are global & may be externally defined) to set tags.
• GetTagForElementInSet(value, set): Given a concrete value and a set, returns the corresponding tag for that value in the set.
• GetElementFromTag(tag : SetTagType(set)): Given a tag and a set, returns the corresponding concrete value in the set.
• CastInterfaceToTaggedUnionPtr(ptr): Convert a pointer to an interface-typed location into a pointer to a tagged union of known types. This is required to avoid changing the layout of an interface type that is used globally (e.g. const buffers, structured buffers, etc.), while still loading into a known tagged union type.

Tagged Union Operations:

• MakeTaggedUnion(tag : SetTagType(tableSet), value : UntaggedUnionType(typeSet)): Create a tagged-union value from a tag and a value.
• GetTagFromTaggedUnion(taggedUnionValue : TaggedUnionType(typeSet, tableSet)): Translate a tagged-union value to its corresponding tag for the witness table represented in the tagged-union's set.
• GetTypeTagFromTaggedUnion(taggedUnionValue : TaggedUnionType(typeSet, tableSet)): Translate a tagged-union value to its corresponding tag for the type in the tagged-union's set. Note that the run-time information about the type is not actually used anywhere in the current compiler, so this inst turns into an undefined (poison) value currently.
• GetValueFromTaggedUnion(taggedUnionValue : TaggedUnionType(typeSet, tableSet)): Translate a tagged-union value to its corresponding value in the tagged-union's set.

Logically all tagged union operations can be thought of as operating on a tuple of (tag, value), where the tag is an identifier for the set of witness tables and the value may be of any of the types in the set of types. These are lowered into MakeTuple/GetTupleElement operations.

Dispatch Operations:

• GetDispatcher(witnessTableSet : WitnessTableSet, key : StructKey): Given a set of witness tables and a key, returns a callable function that dispatches to the correct implementation based on the tag (note that this uses the table tag and not the func tag, to skip the extra mapping step)
• GetSpecializedDispatcher(witnessTableSet : WitnessTableSet, key : StructKey, specArgs...): Given a set of witness tables, a key, and specialization arguments, returns a callable function that dispatches to the correct implementation based on the tag and the specialization arguments.

Specialization Rules For Dynamic Insts

In general, every inst with a non-trivial type-flow-data has its type replaced with the type extracted from the type-flow-data.

For certain instructions, we have additional lowering logic:

1. Call

   • Multiple possible callees become a FuncSet
   • Call sites use tag to identify the specific function
   • If the callee is a Specialize inst, any specialization arguments that are themselves
     set-tags are also passed in as arguments to the call.

2. LookupWitnessMethod

   • Lowered into GetTagForMappedSet(tag, key)
   • Types of the operand and inst encode the source and dest sets

3. Existential Insts

   • ExtractExistentialValue and ExtractExistentialWitnessTable become extraction/reinterpretation operations. ExtractExistentialType is replaced by its inferred TypeSet.

4. Load/Store

   • If an IRStore is writing a value of set type that is different from the pointer's lowered type, then we cast the result (by reinterpreting the value and mapping the tag) before writing.
   • Similarly, if an IRLoad is reading a value of set type that is different from the pointer's lowered type, we cast the loaded value before using.

Specialization of Generics With Set Operands:

The specialization pass can create specialized versions of generics that use sets instead of concrete elements.
This PR introduces an alternative approach to specialization specializeGenericWithSetArgs that handles such IRSpecialize insts.

Unlike the existing specializeGeneric, specializeGenericWithSetArgs introduces additional tag variables for each witness table set operand to keep track of the index during execution.
Example:

// Generic function
%1 = Generic {
%2 = Block
{ 
    %3 : TypeType = Param() // T
    %4 : WitnessTableType(%ifoo) = Param() // T : IFoo
    ...
    %4 : Func {
    %5 = Block 
    {
        %6 : %firstParamType = Param()
        %7 : %secondParamType = Param()
        ...
    }
    }
}
}

// Specialization inst
%8 = Specialize(%1, UntaggedUnionType(TypeSet(A, B)), WitnessTableSet(A_Table, B_Table))

// Specialization result
%9 = Func {
    %10 = Block {
        // Generic parameters that need a run-time tag (i.e. index) get 
        // added to the front as function parameters.
        // Currently, only table sets need tags.
        //
        %11 : SetTagType(WitnessTableSet(A_Table, B_Table)) = Param()
        // Function parameters are at the end.
        %12 : %firstParamType = Param()
        %13 : %secondParamType = Param()
        ...
    }
}

This process of adding tag parameters is only applied for generic functions. For generic struct types, the usual specialization is applied, which simply substitutes the set itself (since the contents of the specialized struct type cannot depend on anything truly dynamic).

Lowering: Generate integer maps, dispatch functions and any-value-types

The Analysis and Specialization parts are run entirely within the specialization loop through
specializeDynamicInsts().

Lowering, on the other hand, happens after the main specialization loop is complete. Most instructions are lowered just before the re-interpret lowering pass. The sequential ID mapping instructions are lowered with the generics lowering step since the sequential IDs need to be assigned first.

Key insts lowered here are:

• GetDispatcher & GetSpecializedDispatcher: Dispatch function with switch statements + calls to witness table wrappers
• UntaggedUnionType: Convert to AnyValueType using the max size of the types in the type set
• GetTagForSuperSet/GetTagForSubSet/GetTagForMappedSet/GetTagForSpecializedSet: These insts are lowered into integer mappings
  by using a mapping from set elements to integers. The current logic assigns the same integer ID for an element, independent of which set they appear in, so GetTagForSuperSet and GetTagForSubSet become no-ops, while GetTagForMappedSet becomes an switch-statement from source set' integers to destination set's integer. GetTagForSpecializedSet should be consumed during the specialization step and turn into a GetSpecializedDispatcher instead, so we do not expect to see it during lowering.
• GetTagFromSequentialID/GetSequentialIDFromTag: This inst is lowered after the witness-table sequential IDs are assigned. Creates a key-mapping function that maps the sequential ID in the table's decoration to a local index into the set.
• Tagged union related insts are lowered into corresponding MakeTuple and GetTupleElement operations.

Examples:

Here is an end-to-end example of analysis, specialization and lowering applied to a function that uses dynamic insts (most types are inlined for readability)

%IFoo : TypeKind = InterfaceType(/* entries */)

%A : StructType = StructType(/* fields */)
%A_calc : FuncType(FloatType(), FloatType()) = Func { /* ... */ }
%AisIFoo : WitnessTableType(%IFoo) = WitnessTable
{
    WitnessTableEntry("calc", %A_calc)
    /* witness table entries */
}

%B : StructType = StructType(/* fields */)
%B_calc : FuncType(FloatType(), FloatType()) = Func { /* ... */ }
%BisIFoo : WitnessTableType(%IFoo) = WitnessTable
{
    WitnessTableEntry("calc", %B_calc)
    /* witness table entries */
}

// compute : uint, float -> float
%compute : FuncType(FloatType(), UintType(), FloatType()) = Func
{
    %1 : BasicBlockType() = Block
    {
        %id : UIntType() = Param()
        %x : FloatType() = Param()
        %11 : BoolType() = Eq(%id, 0)
        ifElse(%11, %id, %x, %4)
    }
    %2 : BasicBlockType() = Block
    {
        %11 : %A = /* Make value of type A */
        %12 : %IFoo = MakeExistential(%11, %A, %AisIFoo)
        branch(%4, %12)
    }
    %3 : BasicBlockType() = Block
    {
        %11 : %B = /* Make value of type B */
        %12 : %IFoo = MakeExistential(%11, %B, %BisIFoo)
        branch(%4, %12)
    }
    %4 : BasicBlockType() = Block
    {
        %41 : %IFoo = Param()
        %42 : TypeKind = ExtractExistentialType(%41)
        %43 : %42 = ExtractExistentialValue(%12)
        %44 : WitnessTableType(%IFoo) = ExtractExistentialWitnessTable(%12)
        %45 : FuncType(FloatType(), FloatType()) = LookupWitnessMethod(%44, "calc")
        %46 : Call(%45, %x)
        %47 : Return(%46)
    }
}

After the Analysis step, these are the assignments that we will hold on to in the
dictionary. Note: nothing in the function is modified at this point, though set
insts may be created in the global scope.

// compute : uint, float -> float
%compute : FuncType(FloatType(), UintType(), FloatType()) = Func
{
    %1 : BasicBlockType() = Block
    {
        %id : UIntType() = Param() |= none 
        %x : FloatType() = Param() |= none
        %11 : BoolType() = Eq(%id, 0) |= none
        IfElse(%11, %id, %x, %4) 
    }
    %2 : BasicBlockType() = Block
    {
        %11 : %A = /* Make value of type A */ |= none
        %12 : %IFoo = MakeExistential(%11, %A, %AisIFoo) |= TaggedUnionType(
                                                                TypeSet(%A),
                                                                WitnessTableSet(%AisIFoo))
        UnconditionalBranch(%4, %12)
    }
    %3 : BasicBlockType() = Block
    {
        %11 : %B = /* Make value of type B */
        %12 : %IFoo = MakeExistential(%11, %B, %BisIFoo) |= TaggedUnionType(
                                                                TypeSet(%B),
                                                                WitnessTableSet(%BisIFoo))
        UnconditionalBranch(%4, %12)
    }
    %4 : BasicBlockType() = Block
    {
        %41 : %IFoo = Param() |= TaggedUnionType(
                                        TypeSet(%A, %B),
                                        WitnessTableSet(%AisIFoo, %BisIFoo))
        %42 : TypeKind = ExtractExistentialType(%41) 
                                |= ElementOfSetType(
                                        TypeSet(%A, %B))
        %43 : %42 = ExtractExistentialValue(%12) 
                                |= UntaggedUnionType(
                                        TypeSet(%A, %B))
        %44 : WitnessTableType(%IFoo) = ExtractExistentialWitnessTable(%12) 
                                |= ElementOfSetType(
                                        WitnessTableSet(%AisIFoo, %BisIFoo))
        %45 : FuncType(FloatType(), FloatType()) = LookupWitnessMethod(%44, "calc")
                                |= ElementOfSetType(
                                        FuncSet(%A_calc, %B_calc))
        %46 : FloatType() = Call(%45, %x) |= none
        Return(%46) |= none
    }
}

After Specialization, the insts have their types rewritten, and any dynamic uses (e.g. calls, loads, stores, etc.)
are rewritten accordingly

%g1 : FuncType(FloatType(), SetTagType(WitnessTableSet(%AisIFoo, %BisIFoo)), FloatType()) 
    = GetDispatcher(WitnessTableSet(%AisIFoo, %BisIFoo), "calc")
// compute : uint, float -> float
%compute : FuncType(FloatType(), UintType(), FloatType()) = Func
{
    %1 : BasicBlockType() = Block
    {
        %id : UIntType() = Param() 
        %x : FloatType() = Param() 
        %11 : BoolType() = Eq(%id, 0) 
        IfElse(%11, %id, %x, %4) 
    }
    %2 : BasicBlockType() = Block
    {
        %11 : %A = /* Make value of type A */ 
        %12 : SetTagType(WitnessTableSet(%A)) = GetTagForElementInSet(%A, WitnessTableSet(%A))
        %13 : TaggedUnionType(TypeSet(%A), WitnessTableSet(%AisIFoo)) 
            = MakeTaggedUnion(%12, %11) // tag ID is always '0' for single element sets, 
                                    // since the IDs are local
        %14 : SetTagType(WitnessTableSet(%AisIFoo)) = GetTagFromTaggedUnion(%13)
        %15 : %A = GetValueFromTaggedUnion(%13)
        // Since the type of the arg (%12) doesn't match the parameter (%41), we need to 
        // 'upcast' it.
        %16 : SetTagType(WitnessTableSet(%AisIFoo, %BisIFoo)) = 
            GetTagForSuperSet(%14)
        %17 : UntaggedUnionType(TypeSet(%A, %B)) = Reinterpret(%15)
        %18 : TaggedUnionType(
                    TypeSet(%A, %B),
                    WitnessTableSet(%AisIFoo, %BisIFoo)) 
            = MakeTaggedUnion(%15, %16)
        UnconditionalBranch(%4, %18)
    }
    %3 : BasicBlockType() = Block
    {
        // Same as above, but for %B
        ...
        UnconditionalBranch(%4, /*..*/)
    }
    %4 : BasicBlockType() = Block
    {
        %41 : TaggedUnionType(
                    TypeSet(%A, %B),
                    WitnessTableSet(%AisIFoo, %BisIFoo)) = Param() 
        
        // Replaced entirely by it's type-set 
        %42 : SetTagType(TypeSet(%A, %B)) = GetTypeTagFromTaggedUnion(%41) 
        
        %43 : UntaggedUnionType(TypeSet(%A, %B)) = ExtractExistentialValue(%41)
        %44 : SetTagType(WitnessTableSet(%AisIFoo, %BisIFoo)) = 
            GetTagFromTaggedUnion(%41, 0)
        %45 : SetTagType(FuncSet(%A_calc, %B_calc)) = 
            
        
        // Callee is replaced with a dispatcher (see %g1 above)
        // which represents a wrapper function that dispatches to 
        // the appropriate implementation based on the witness table tag.
        //
        // The resulting function takes in the tag as the first argument
        //
        // If the callee (or any one in the set) is a dynamic specialization, 
        // any necessary tags are also appended to the front.
        //
        %46 : FloatType() = Call(%g1, %45, %x)

        Return(%46)
    }
}

After lowering

// `GetDispatcher` becomes a dispatch function
// g1 : uint, float -> float
%g1 : FuncType(FloatType(), UIntType(), FloatType()) = Func
{
    // Switch-case based dispatch function.
    ...
}

// compute : uint, float -> float
%compute : FuncType(FloatType(), UintType(), FloatType()) = Func
{
    %1 : BasicBlockType() = Block
    {
        %id : UIntType() = Param() 
        %x : FloatType() = Param() 
        %11 : BoolType() = Eq(%id, 0) 
        IfElse(%11, %id, %x, %4) 
    }
    %2 : BasicBlockType() = Block
    {
        %11 : %A = /* Make value of type A */ 
        %12 : TupleType(UIntType(), %A) 
            = MakeTuple(0, %11) // %A has a unique global index of 0
        %14 : %A = GetTupleElement(%12, 1)
        
        // Reinterpret the value part to fit into the larger AnyValueType
        // and pass the tag through as-is. Tag IDs are globally consistent.
        // 
        %16 : AnyValueType(max(sizeof(%A), sizeof(%B))) = Reinterpret(%14)
        %17 : TupleType(
                    UIntType(),
                    AnyValueType(max(sizeof(%A), sizeof(%B))))
            = MakeTuple(%12, %16)
        UnconditionalBranch(%4, %17)
    }
    %3 : BasicBlockType() = Block
    {
        %21 : %B = /* Make value of type B */ 
        %22 : TupleType(UIntType(), %B) 
            = MakeTuple(1, %21) // %B has a unique global index of 1
        %24 : %B = GetTupleElement(%22, 1)

        %26 : AnyValueType(max(sizeof(%A), sizeof(%B))) = Reinterpret(%24)
        %27 : TupleType(
                    UIntType(),
                    AnyValueType(max(sizeof(%A), sizeof(%B))))
            = MakeTuple(%22, %26)
        UnconditionalBranch(%4, %27)
    }
    %4 : BasicBlockType() = Block
    {
        %41 : TupleType(UIntType(), AnyValueType(max(sizeof(%A), sizeof(%B)))) 
            = Param() 
        
        %43 : AnyValueType(max(sizeof(%A), sizeof(%B))) = GetTupleElement(%41, 1)
        %44 : UIntType() = 
            GetTupleElement(%41, 0) 
    
        %46 : Call(%g1, %44, %x) 

        Return(%46)
    }
}

Breaking Changes

While this PR does not change the language or spec, it enforces some of the rules around existentials more strictly, which can cause old code to fail with diagnostic errors rather than silently generating code that might have undefined behavior at runtime.

Summary of patterns to look out for:

• Initializing interface-type objects with default constructor {}. E.g. IFoo foo = {};

  • This will cause an error, since an interface-type object has no default value.
  • Fix Option 1: IFoo foo = MyDefaultFooImpl();. Create your own default concrete type that conforms to IFoo
  • Fix Option 2: Use Optional<IFoo> foo = {}; which initializes to none. Uses of foo need to change to foo.value and should check for foo.hasValue first.

• Potentially uninitialized interface-type objects:

  IFoo foo;
  if(someCondition)
      foo = MyFooImpl();
  foo.calc(); // Error 50101: Cannot dynamically dispatch on potentially uninitialized interface-type object 'IFoo'

  • Fix Option 1: Ensure that all code paths initialize the interface-type object before use.
  • Fix Option 2: Use Optional<IFoo> foo = {}; when creating the object.

[jhelferty-nv]

Reapproving based on Tess's previous approval, and Sai's assurance that changes were minor to address branch drift.