User IR-based layout for all IR steps (#1084)

This change builds on previous work that moves toward a more IR-based representation of layout.
Those steps added some instructions for representing layout in the IR (initially just proxies for the AST layout objects), and an explicit lowering pass that could build a target-specific IR module that binds parameters and entry points to layout information.

This change aims to complete that work, in the sense that the IR representation of layout is now self-contained and does not rely on having pointers back into the AST-level representation.
Achieving this requires two main kinds of work:

1. Update any code that used layout information derived from the IR (most notably all the `slang-emit-*` code) to use the new IR representation and its accessors.

2. Update any code that *constructs* layouts using information derived from the IR to construct IR layouts instead.

The biggest new infrastructure feature in this change is support for "attributes" in the IR (I'd welcome feedback on the naming).
An attribute can either be thought of like key/value arguments that can be added to certain instructions to encode optional data, or alternatively like a decoration that is referenced as an operand instead of a child.
The value of attributes over decorations is that they can affect the hash/identity of an instruction (which decorations can't), while the advantage of decorations is that they can easily be added/removed over the lifetime of an instruction (which attributes can't).
We mostly use them here to represent operands that are logically optional.

Once attributes are available, the encoding of layout information into the IR is mostly straightforward:

* An `IRVarLayout` has a fixed operand for its type layout, and can accept a few different attributes
  * Zero or more `IRVarOffsetAttr`s that specify the offset of the variable for a given resource kind. These are equivalent to the `VarLayout::ResourceInfo`s at the AST level.
  * An optional `IRUserSemanticAttr` and `IRSystemValueSemanticAttr` to represent the (possibly derived) semantic of a varying input/output parameter.
  * An option `IRStageAttr` to represent the known stage for a parameter.

* An `IREntryPointLayout` has a var layout for the entry point parameters (logically grouped in to a struct) and another var layout for the result parameter.

* There is a small type hierarchy rooted at `IRTypeLayout` where each subtype can add fixed operands and attributes that are expected to appear. It also supports `IRTypeSizeAttr`s that serve a similar role to the `IRVarOffsetAttr`s.

* Structure types maintain the mapping of fields to their var layouts using `IRStructFieldLayoutAttr`s.

With the encoding in place, most of the changes in category (1) (code that just *uses* rather than *creates* layouts) was straightforward. The biggest different beyond name changes was that everything needs to be fetched using accessors instead of bare fields. It would have been possible to stage this commit and make the diffs smaller by first introducing mandatory acessors to the AST layout types.

The changes in category (2) were more involved. There were a lot of places in the existing code where a `TypeLayout` or `VarLayout` would be created, and then initialized piecemeal over several lines of code (and sometimes even across functions). Because of the way that layouts need to support many optional properties, it did not seem practical to just have monolithic factory functions that took all the options as arguments, so I instead opted for a builder approach.

The builders for `IRVarLayout` and `IREntryPointLayout` are both straightforward, and honestly there is no realy need for a builder for entry point layouts right now, but I was trying to future-proof in case we decidd to add some optional attributes to them.

The builders for type layouts are more involved because of the inheritance hierarchy. Each concrete sub-type of type layout needs to define its own builder type that customizes the opcode, operands, and attributes of the final instruction.

The refactoring that had to go into this change was a nice excuse to clean up a few ugly warts in the AST layout code that were largely there to support IR use cases. While this change adds a lot of new infrastructure code to the IR, most of the client code has stayed the same or gotten simpler.

One annoying wart that remains with this change is the notion of an "offset element type layout" for parameter group types. That idea was added to deal with a legacy feature in the reflection API that we realized was a mistake, but unfortunately having that "offset" layout handy made writing a few other pieces of code simpler so that there are use cases of the feature even in the IR. Removing those uses is do-able, but requires careful refactoring so it is best left to a follow-on change.

Another thing that could be considered for a follow-on change is how much information should be specified when constructing a `Builder` for an IR type layout, and how much should be allowed to be specified statefully/piecemeal. It would be nice to force all the required operands to be specified up front, but `IRParameterGroupTypeLayout::Builder` doesn't currently work that way because so much of the client code that needs it involved a lot of stateful setting and would need to be refactored heavily to provide the necessary information up front.

1 files changed, 283 insertions, 4 deletions

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index cd398f4e7..113dcbfad 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -6714,6 +6714,278 @@ RefPtr<IRModule> TargetProgram::getOrCreateIRModuleForLayout(DiagnosticSink* sin
     return m_irModuleForLayout;
 }
 
+    /// Specialized IR generation context for when generating IR for layouts.
+struct IRLayoutGenContext : IRGenContext
+{
+    IRLayoutGenContext(SharedIRGenContext* shared)
+        : IRGenContext(shared)
+    {}
+
+        /// Cache for custom key instructions used for entry-point parameter layout information.
+    Dictionary<ParamDecl*, IRStructKey*> mapEntryPointParamToKey;
+};
+
+    /// Lower an AST-level type layout to an IR-level type layout.
+IRTypeLayout* lowerTypeLayout(
+    IRLayoutGenContext* context,
+    TypeLayout*         typeLayout);
+
+    /// Lower an AST-level variable layout to an IR-level variable layout.
+IRVarLayout* lowerVarLayout(
+    IRLayoutGenContext* context,
+    VarLayout*          varLayout);
+
+    /// Shared code for most `lowerTypeLayout` cases.
+    ///
+    /// Handles copying of resource usage and pending data type layout
+    /// from the AST `typeLayout` to the specified `builder`.
+    ///
+static IRTypeLayout* _lowerTypeLayoutCommon(
+    IRLayoutGenContext*     context,
+    IRTypeLayout::Builder*  builder,
+    TypeLayout*             typeLayout)
+{
+    for( auto resInfo : typeLayout->resourceInfos )
+    {
+        builder->addResourceUsage(resInfo.kind, resInfo.count);
+    }
+
+    if( auto pendingTypeLayout = typeLayout->pendingDataTypeLayout )
+    {
+        builder->setPendingTypeLayout(
+            lowerTypeLayout(context, pendingTypeLayout));
+    }
+
+    return builder->build();
+}
+
+IRTypeLayout* lowerTypeLayout(
+    IRLayoutGenContext* context,
+    TypeLayout*         typeLayout)
+{
+    // TODO: We chould consider caching the layouts we create based on `typeLayout`
+    // and re-using them. This isn't strictly necessary because we emit the
+    // instructions as "hoistable" which should give us de-duplication, and it wouldn't
+    // help much until/unless the AST level gets less wasteful about how it computes layout.
+
+    // We will use casting to detect if `typeLayout` is
+    // one of the cases that requires a dedicated sub-type
+    // of IR type layout.
+    //
+    if( auto paramGroupTypeLayout = as<ParameterGroupTypeLayout>(typeLayout) )
+    {
+        IRParameterGroupTypeLayout::Builder builder(context->irBuilder);
+
+        builder.setContainerVarLayout(
+            lowerVarLayout(context, paramGroupTypeLayout->containerVarLayout));
+        builder.setElementVarLayout(
+            lowerVarLayout(context, paramGroupTypeLayout->elementVarLayout));
+        builder.setOffsetElementTypeLayout(
+            lowerTypeLayout(context, paramGroupTypeLayout->offsetElementTypeLayout));
+
+        return _lowerTypeLayoutCommon(context, &builder, paramGroupTypeLayout);
+    }
+    else if( auto structTypeLayout = as<StructTypeLayout>(typeLayout) )
+    {
+        IRStructTypeLayout::Builder builder(context->irBuilder);
+
+        for( auto fieldLayout : structTypeLayout->fields )
+        {
+            auto fieldDecl = fieldLayout->varDecl;
+
+            IRStructKey* irFieldKey = nullptr;
+            if(auto paramDecl = as<ParamDecl>(fieldDecl) )
+            {
+                // There is a subtle special case here.
+                //
+                // A `StructTypeLayout` might be used to represent
+                // the parameters of an entry point, and this is the
+                // one and only case where the "fields" being used
+                // might actually be `ParamDecl`s.
+                //
+                // The IR encoding of structure type layouts relies
+                // on using field "key" instructions to identify
+                // the fields, but these don't exist (by default)
+                // for function parameters.
+                //
+                // To get around this problem we will create key
+                // instructions to stand in for the entry-point parameters
+                // as needed when generating layout.
+                //
+                // We need to cache the generated keys on the context,
+                // so that if we run into another type layout for the
+                // same entry point we will re-use the same keys.
+                //
+                if( !context->mapEntryPointParamToKey.TryGetValue(paramDecl, irFieldKey) )
+                {
+                    irFieldKey = context->irBuilder->createStructKey();
+
+                    // TODO: It might eventually be a good idea to attach a mangled
+                    // name to the key we just generated (derived from the entry point
+                    // and parameter name), even though parameters don't usually have
+                    // linkage.
+                    //
+                    // Doing so would ensure that if we ever combined partial layout
+                    // information from different modules they would agree on the key
+                    // to use for entry-point parameters.
+                    //
+                    // For now this is a non-issue because both the creation and use
+                    // of these keys will be local to a single `IREntryPointLayout`,
+                    // and we don't support combination at a finer granularity than that.
+
+                    context->mapEntryPointParamToKey.Add(paramDecl, irFieldKey);
+                }
+            }
+            else
+            {
+                IRInst* irFieldKeyInst = getSimpleVal(context,
+                    ensureDecl(context, fieldDecl));
+                irFieldKey = as<IRStructKey>(irFieldKeyInst);
+            }
+            SLANG_ASSERT(irFieldKey);
+
+            auto irFieldLayout = lowerVarLayout(context, fieldLayout);
+            builder.addField(irFieldKey, irFieldLayout);
+        }
+
+        return _lowerTypeLayoutCommon(context, &builder, structTypeLayout);
+    }
+    else if( auto arrayTypeLayout = as<ArrayTypeLayout>(typeLayout) )
+    {
+        auto irElementTypeLayout = lowerTypeLayout(context, arrayTypeLayout->elementTypeLayout);
+        IRArrayTypeLayout::Builder builder(context->irBuilder, irElementTypeLayout);
+        return _lowerTypeLayoutCommon(context, &builder, arrayTypeLayout);
+    }
+    else if( auto taggedUnionTypeLayout = as<TaggedUnionTypeLayout>(typeLayout) )
+    {
+        IRTaggedUnionTypeLayout::Builder builder(context->irBuilder, taggedUnionTypeLayout->tagOffset);
+
+        for( auto caseTypeLayout : taggedUnionTypeLayout->caseTypeLayouts )
+        {
+            builder.addCaseTypeLayout(
+                lowerTypeLayout(
+                    context,
+                    caseTypeLayout));
+        }
+
+        return _lowerTypeLayoutCommon(context, &builder, taggedUnionTypeLayout);
+    }
+    else if( auto streamOutputTypeLayout = as<StreamOutputTypeLayout>(typeLayout) )
+    {
+        auto irElementTypeLayout = lowerTypeLayout(context, streamOutputTypeLayout->elementTypeLayout);
+
+        IRStreamOutputTypeLayout::Builder builder(context->irBuilder, irElementTypeLayout);
+        return _lowerTypeLayoutCommon(context, &builder, streamOutputTypeLayout);
+    }
+    else if( auto matrixTypeLayout = as<MatrixTypeLayout>(typeLayout) )
+    {
+        // TODO: Our support for explicit layouts on matrix types is minimal, so whether
+        // or not we even include `IRMatrixTypeLayout` doesn't impact any behavior we
+        // currently test.
+        //
+        // Our handling of matrix types and their layout needs a complete overhaul, but
+        // that isn't something we can get to right away, so we'll just try to pass
+        // along this data as best we can for now.
+
+        IRMatrixTypeLayout::Builder builder(context->irBuilder, matrixTypeLayout->mode);
+        return _lowerTypeLayoutCommon(context, &builder, matrixTypeLayout);
+    }
+    else
+    {
+        // If no special case applies we will build a generic `IRTypeLayout`.
+        //
+        IRTypeLayout::Builder builder(context->irBuilder);
+        return _lowerTypeLayoutCommon(context, &builder, typeLayout);
+    }
+}
+
+IRVarLayout* lowerVarLayout(
+    IRLayoutGenContext* context,
+    VarLayout*          varLayout)
+{
+    auto irTypeLayout = lowerTypeLayout(context, varLayout->typeLayout);
+    IRVarLayout::Builder irLayoutBuilder(context->irBuilder, irTypeLayout);
+
+    for( auto resInfo : varLayout->resourceInfos )
+    {
+        auto irResInfo = irLayoutBuilder.findOrAddResourceInfo(resInfo.kind);
+        irResInfo->offset = resInfo.index;
+        irResInfo->space = resInfo.space;
+    }
+
+    if( auto pendingVarLayout = varLayout->pendingVarLayout )
+    {
+        irLayoutBuilder.setPendingVarLayout(
+            lowerVarLayout(context, pendingVarLayout));
+    }
+
+    // We will only generate layout information with *either* a system-value
+    // semantic or a user-defined semantic, and we will always check for
+    // the system-value semantic first because the AST-level representation
+    // seems to encode both when a system-value semantic is present.
+    //
+    if( varLayout->systemValueSemantic.getLength() )
+    {
+        irLayoutBuilder.setSystemValueSemantic(
+            varLayout->systemValueSemantic,
+            varLayout->systemValueSemanticIndex);
+    }
+    else if( varLayout->semanticName.getLength() )
+    {
+        irLayoutBuilder.setUserSemantic(
+            varLayout->semanticName,
+            varLayout->semanticIndex);
+    }
+
+    if( varLayout->stage != Stage::Unknown )
+    {
+        irLayoutBuilder.setStage(varLayout->stage);
+    }
+
+    return irLayoutBuilder.build();
+}
+
+    /// Handle the lowering of an entry-point result layout to the IR
+IRVarLayout* lowerEntryPointResultLayout(
+    IRLayoutGenContext* context,
+    VarLayout*          layout)
+{
+    // The easy case is when there is a non-null `layout`, because we
+    // can handle it like any other var layout.
+    //
+    if(layout)
+        return lowerVarLayout(context, layout);
+
+    // Right now the AST-level layout logic will leave a null layout
+    // for the result when an entry point has a `void` result type.
+    //
+    // TODO: We should fix this at the AST level instead of the IR,
+    // but doing so would impact reflection, where clients could
+    // be using a null check to test for a `void` result.
+    //
+    // As a workaround, we will create an empty type layout and
+    // an empty var layout that represents it, consistent with the
+    // way that a `void` value consumes no resources.
+    //
+    IRTypeLayout::Builder typeLayoutBuilder(context->irBuilder);
+    auto irTypeLayout = typeLayoutBuilder.build();
+    IRVarLayout::Builder varLayoutBuilder(context->irBuilder, irTypeLayout);
+    return varLayoutBuilder.build();
+}
+
+    /// Lower AST-level layout information for an entry point to the IR
+IREntryPointLayout* lowerEntryPointLayout(
+    IRLayoutGenContext* context,
+    EntryPointLayout*   entryPointLayout)
+{
+    auto irParamsLayout = lowerVarLayout(context, entryPointLayout->parametersLayout);
+    auto irResultLayout = lowerEntryPointResultLayout(context, entryPointLayout->resultLayout);
+
+    return context->irBuilder->getEntryPointLayout(
+        irParamsLayout,
+        irResultLayout);
+}
+
 RefPtr<IRModule> TargetProgram::createIRModuleForLayout(DiagnosticSink* sink)
 {
     if(m_irModuleForLayout)
@@ -6735,7 +7007,7 @@ RefPtr<IRModule> TargetProgram::createIRModuleForLayout(DiagnosticSink* sink)
         sink);
     auto sharedContext = &sharedContextStorage;
 
-    IRGenContext contextStorage(sharedContext);
+    IRLayoutGenContext contextStorage(sharedContext);
     auto context = &contextStorage;
 
     SharedIRBuilder sharedBuilderStorage;
@@ -6768,9 +7040,11 @@ RefPtr<IRModule> TargetProgram::createIRModuleForLayout(DiagnosticSink* sink)
         //
         auto irVar = getSimpleVal(context, ensureDecl(context, varDecl));
 
+        auto irLayout = lowerVarLayout(context, varLayout);
+
         // Now attach the decoration to the variable.
         //
-        builder->addLayoutDecoration(irVar, varLayout);
+        builder->addLayoutDecoration(irVar, irLayout);
     }
 
     for( auto entryPointLayout : programLayout->entryPoints )
@@ -6785,14 +7059,19 @@ RefPtr<IRModule> TargetProgram::createIRModuleForLayout(DiagnosticSink* sink)
             builder->addImportDecoration(irFunc, getMangledName(funcDeclRef).getUnownedSlice());
         }
 
-        builder->addLayoutDecoration(irFunc, entryPointLayout);
+        auto irEntryPointLayout = lowerEntryPointLayout(context, entryPointLayout);
+
+        builder->addLayoutDecoration(irFunc, irEntryPointLayout);
     }
 
     for( auto taggedUnionTypeLayout : programLayout->taggedUnionTypeLayouts )
     {
         auto taggedUnionType = taggedUnionTypeLayout->getType();
         auto irType = lowerType(context, taggedUnionType);
-        builder->addLayoutDecoration(irType, taggedUnionTypeLayout);
+
+        auto irTypeLayout = lowerTypeLayout(context, taggedUnionTypeLayout);
+
+        builder->addLayoutDecoration(irType, irTypeLayout);
     }
 
     m_irModuleForLayout = irModule;