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, 509 insertions, 56 deletions

diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 78758d944..89e49816f 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -205,6 +205,24 @@ namespace Slang
         return nullptr;
     }
 
+    IROperandListBase IRInst::getAllAttrs()
+    {
+        // We assume as an invariant that all attributes appear at the end of the operand
+        // list, after all the non-attribute operands.
+        //
+        // We will therefore define a range that ends at the end of the operand list ...
+        //
+        IRUse* end = getOperands() + getOperandCount();
+        //
+        // ... and begins after the last non-attribute operand.
+        //
+        IRUse* cursor = getOperands();
+        while(cursor != end && !as<IRAttr>(cursor->get()))
+            cursor++;
+
+        return IROperandListBase(cursor, end);
+    }
+
     // IRConstant
 
     IRIntegerValue GetIntVal(IRInst* inst)
@@ -696,6 +714,371 @@ namespace Slang
     }
 
     //
+    // IRTypeLayout
+    //
+
+    IRTypeSizeAttr* IRTypeLayout::findSizeAttr(LayoutResourceKind kind)
+    {
+        // TODO: If we could assume the attributes were sorted
+        // by `kind`, then we could use a binary search here
+        // instead of linear.
+        //
+        // In practice, the number of entries will be very small,
+        // so the cost of the linear search should not be too bad.
+
+        for( auto sizeAttr : getSizeAttrs() )
+        {
+            if(sizeAttr->getResourceKind() == kind)
+                return sizeAttr;
+        }
+        return nullptr;
+    }
+
+    IRTypeLayout* IRTypeLayout::unwrapArray()
+    {
+        auto typeLayout = this;
+        while(auto arrayTypeLayout = as<IRArrayTypeLayout>(typeLayout))
+            typeLayout = arrayTypeLayout->getElementTypeLayout();
+        return typeLayout;
+    }
+
+    IRTypeLayout* IRTypeLayout::getPendingDataTypeLayout()
+    {
+        if(auto attr = findAttr<IRPendingLayoutAttr>())
+            return cast<IRTypeLayout>(attr->getLayout());
+        return nullptr;
+    }
+
+    IROperandList<IRTypeSizeAttr> IRTypeLayout::getSizeAttrs()
+    {
+        return findAttrs<IRTypeSizeAttr>();
+    }
+
+    IRTypeLayout::Builder::Builder(IRBuilder* irBuilder)
+        : m_irBuilder(irBuilder)
+    {}
+
+    void IRTypeLayout::Builder::addResourceUsage(
+        LayoutResourceKind  kind,
+        LayoutSize          size)
+    {
+        auto& resInfo = m_resInfos[Int(kind)];
+        resInfo.kind = kind;
+        resInfo.size += size;
+    }
+
+    void IRTypeLayout::Builder::addResourceUsage(IRTypeSizeAttr* sizeAttr)
+    {
+        addResourceUsage(
+            sizeAttr->getResourceKind(),
+            sizeAttr->getSize());
+    }
+
+    void IRTypeLayout::Builder::addResourceUsageFrom(IRTypeLayout* typeLayout)
+    {
+        for( auto sizeAttr : typeLayout->getSizeAttrs() )
+        {
+            addResourceUsage(sizeAttr);
+        }
+    }
+
+    IRTypeLayout* IRTypeLayout::Builder::build()
+    {
+        IRBuilder* irBuilder = getIRBuilder();
+
+        List<IRInst*> operands;
+
+        addOperands(operands);
+        addAttrs(operands);
+
+        return irBuilder->getTypeLayout(
+            getOp(),
+            operands);
+    }
+
+    void IRTypeLayout::Builder::addOperands(List<IRInst*>& operands)
+    {
+        addOperandsImpl(operands);
+    }
+
+    void IRTypeLayout::Builder::addAttrs(List<IRInst*>& operands)
+    {
+        auto irBuilder = getIRBuilder();
+
+        for(auto resInfo : m_resInfos)
+        {
+            if(resInfo.kind == LayoutResourceKind::None)
+                continue;
+
+            IRInst* sizeAttr = irBuilder->getTypeSizeAttr(
+                resInfo.kind,
+                resInfo.size);
+            operands.add(sizeAttr);
+        }
+
+        if( auto pendingTypeLayout = m_pendingTypeLayout )
+        {
+            operands.add(irBuilder->getPendingLayoutAttr(
+                pendingTypeLayout));
+        }
+
+        addAttrsImpl(operands);
+    }
+
+    //
+    // IRParameterGroupTypeLayout
+    //
+
+    void IRParameterGroupTypeLayout::Builder::addOperandsImpl(List<IRInst*>& ioOperands)
+    {
+        ioOperands.add(m_containerVarLayout);
+        ioOperands.add(m_elementVarLayout);
+        ioOperands.add(m_offsetElementTypeLayout);
+    }
+
+    IRParameterGroupTypeLayout* IRParameterGroupTypeLayout::Builder::build()
+    {
+        return cast<IRParameterGroupTypeLayout>(Super::Builder::build());
+    }
+
+    //
+    // IRStructTypeLayout
+    //
+
+    void IRStructTypeLayout::Builder::addAttrsImpl(List<IRInst*>& ioOperands)
+    {
+        auto irBuilder = getIRBuilder();
+        for(auto field : m_fields)
+        {
+            ioOperands.add(
+                irBuilder->getFieldLayoutAttr(field.key, field.layout));
+        }
+    }
+
+    //
+    // IRArrayTypeLayout
+    //
+
+    void IRArrayTypeLayout::Builder::addOperandsImpl(List<IRInst*>& ioOperands)
+    {
+        ioOperands.add(m_elementTypeLayout);
+    }
+
+    //
+    // IRStreamOutputTypeLayout
+    //
+
+    void IRStreamOutputTypeLayout::Builder::addOperandsImpl(List<IRInst*>& ioOperands)
+    {
+        ioOperands.add(m_elementTypeLayout);
+    }
+
+    //
+    // IRMatrixTypeLayout
+    //
+
+    IRMatrixTypeLayout::Builder::Builder(IRBuilder* irBuilder, MatrixLayoutMode mode)
+        : Super::Builder(irBuilder)
+    {
+        m_modeInst = irBuilder->getIntValue(irBuilder->getIntType(), IRIntegerValue(mode));
+    }
+
+    void IRMatrixTypeLayout::Builder::addOperandsImpl(List<IRInst*>& ioOperands)
+    {
+        ioOperands.add(m_modeInst);
+    }
+
+    //
+    // IRTaggedUnionTypeLayout
+    //
+
+    IRTaggedUnionTypeLayout::Builder::Builder(IRBuilder* irBuilder, LayoutSize tagOffset)
+        : Super::Builder(irBuilder)
+    {
+        m_tagOffset = irBuilder->getIntValue(irBuilder->getIntType(), tagOffset.raw);
+    }
+
+    void IRTaggedUnionTypeLayout::Builder::addCaseTypeLayout(IRTypeLayout* typeLayout)
+    {
+        m_caseTypeLayoutAttrs.add(getIRBuilder()->getCaseTypeLayoutAttr(typeLayout));
+    }
+
+    void IRTaggedUnionTypeLayout::Builder::addOperandsImpl(List<IRInst*>& ioOperands)
+    {
+        ioOperands.add(m_tagOffset);
+    }
+
+    void IRTaggedUnionTypeLayout::Builder::addAttrsImpl(List<IRInst*>& ioOperands)
+    {
+        for(auto attr : m_caseTypeLayoutAttrs)
+            ioOperands.add(attr);
+    }
+
+    //
+    // IRVarLayout
+    //
+
+    bool IRVarLayout::usesResourceKind(LayoutResourceKind kind)
+    {
+        // TODO: basing this check on whether or not the
+        // var layout has an entry for `kind` means that
+        // we can't just optimize away any entry where
+        // the offset is zero (which might be a small
+        // but nice optimization). We could consider shifting
+        // this test to use the entries on the type layout
+        // instead (since non-zero resource consumption
+        // should be an equivalent test).
+
+        return findOffsetAttr(kind) != nullptr;
+    }
+
+    IRSystemValueSemanticAttr* IRVarLayout::findSystemValueSemanticAttr()
+    {
+        return findAttr<IRSystemValueSemanticAttr>();
+    }
+
+    IRVarOffsetAttr* IRVarLayout::findOffsetAttr(LayoutResourceKind kind)
+    {
+        for( auto offsetAttr : getOffsetAttrs() )
+        {
+            if(offsetAttr->getResourceKind() == kind)
+                return offsetAttr;
+        }
+        return nullptr;
+    }
+
+    IROperandList<IRVarOffsetAttr> IRVarLayout::getOffsetAttrs()
+    {
+        return findAttrs<IRVarOffsetAttr>();
+    }
+
+    Stage IRVarLayout::getStage()
+    {
+        if(auto stageAttr = findAttr<IRStageAttr>())
+            return stageAttr->getStage();
+        return Stage::Unknown;
+    }
+
+    IRVarLayout* IRVarLayout::getPendingVarLayout()
+    {
+        if( auto pendingLayoutAttr = findAttr<IRPendingLayoutAttr>() )
+        {
+            return cast<IRVarLayout>(pendingLayoutAttr->getLayout());
+        }
+        return nullptr;
+    }
+
+    IRVarLayout::Builder::Builder(
+        IRBuilder*      irBuilder,
+        IRTypeLayout*   typeLayout)
+        : m_irBuilder(irBuilder)
+        , m_typeLayout(typeLayout)
+    {}
+
+    bool IRVarLayout::Builder::usesResourceKind(LayoutResourceKind kind)
+    {
+        return m_resInfos[Int(kind)].kind != LayoutResourceKind::None;
+    }
+
+    IRVarLayout::Builder::ResInfo* IRVarLayout::Builder::findOrAddResourceInfo(LayoutResourceKind kind)
+    {
+        auto& resInfo = m_resInfos[Int(kind)];
+        resInfo.kind = kind;
+        return &resInfo;
+    }
+
+    void IRVarLayout::Builder::setSystemValueSemantic(String const& name, UInt index)
+    {
+        m_systemValueSemantic = getIRBuilder()->getSystemValueSemanticAttr(name, index);
+    }
+
+    void IRVarLayout::Builder::setUserSemantic(String const& name, UInt index)
+    {
+        m_userSemantic = getIRBuilder()->getUserSemanticAttr(name, index);
+    }
+
+    void IRVarLayout::Builder::setStage(Stage stage)
+    {
+        m_stageAttr = getIRBuilder()->getStageAttr(stage);
+    }
+
+    void IRVarLayout::Builder::cloneEverythingButOffsetsFrom(
+        IRVarLayout* that)
+    {
+        if(auto systemValueSemantic = that->findAttr<IRSystemValueSemanticAttr>())
+            m_systemValueSemantic = systemValueSemantic;
+
+        if(auto userSemantic = that->findAttr<IRUserSemanticAttr>())
+            m_userSemantic = userSemantic;
+
+        if(auto stageAttr = that->findAttr<IRStageAttr>())
+            m_stageAttr = stageAttr;
+    }
+
+    IRVarLayout* IRVarLayout::Builder::build()
+    {
+        SLANG_ASSERT(m_typeLayout);
+
+        IRBuilder* irBuilder = getIRBuilder();
+
+        List<IRInst*> operands;
+
+        operands.add(m_typeLayout);
+
+        for(auto resInfo : m_resInfos)
+        {
+            if(resInfo.kind == LayoutResourceKind::None)
+                continue;
+
+            IRInst* varOffsetAttr = irBuilder->getVarOffsetAttr(
+                resInfo.kind,
+                resInfo.offset,
+                resInfo.space);
+            operands.add(varOffsetAttr);
+        }
+
+        if(auto semanticAttr = m_userSemantic)
+            operands.add(semanticAttr);
+
+        if(auto semanticAttr = m_systemValueSemantic)
+            operands.add(semanticAttr);
+
+        if(auto stageAttr = m_stageAttr)
+            operands.add(stageAttr);
+
+        if(auto pendingVarLayout = m_pendingVarLayout)
+        {
+            IRInst* pendingLayoutAttr = irBuilder->getPendingLayoutAttr(
+                pendingVarLayout);
+            operands.add(pendingLayoutAttr);
+        }
+
+        return irBuilder->getVarLayout(operands);
+    }
+
+    //
+    // IREntryPointLayout
+    //
+
+    IRStructTypeLayout* getScopeStructLayout(IREntryPointLayout* scopeLayout)
+    {
+        auto scopeTypeLayout = scopeLayout->getParamsLayout()->getTypeLayout();
+
+        if( auto constantBufferTypeLayout = as<IRParameterGroupTypeLayout>(scopeTypeLayout) )
+        {
+            scopeTypeLayout = constantBufferTypeLayout->getOffsetElementTypeLayout();
+        }
+
+        if( auto structTypeLayout = as<IRStructTypeLayout>(scopeTypeLayout) )
+        {
+            return structTypeLayout;
+        }
+
+        SLANG_UNEXPECTED("uhandled global-scope binding layout");
+        UNREACHABLE_RETURN(nullptr);
+    }
+
+    //
 
     IRBlock* IRBuilder::getBlock()
     {
@@ -3086,79 +3469,149 @@ namespace Slang
         addDecoration(inst, kIROp_HighLevelDeclDecoration, ptrConst);
     }
 
-    void IRBuilder::addLayoutDecoration(IRInst* value, Layout* layout)
+    void IRBuilder::addLayoutDecoration(IRInst* value, IRLayout* layout)
+    {
+        addDecoration(value, kIROp_LayoutDecoration, layout);
+    }
+
+    IRTypeSizeAttr* IRBuilder::getTypeSizeAttr(
+        LayoutResourceKind kind,
+        LayoutSize size)
     {
-        IRLayout* irLayout = getLayout(layout);
-        addDecoration(value, kIROp_LayoutDecoration, irLayout);
+        auto kindInst = getIntValue(getIntType(), IRIntegerValue(kind));
+        auto sizeInst = getIntValue(getIntType(), IRIntegerValue(size.raw));
 
+        IRInst* operands[] = { kindInst, sizeInst };
 
+        return cast<IRTypeSizeAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_TypeSizeAttr,
+            SLANG_COUNT_OF(operands),
+            operands));
     }
 
-    IRLayout* IRBuilder::getLayout(Layout* astLayout)
+    IRVarOffsetAttr* IRBuilder::getVarOffsetAttr(
+        LayoutResourceKind  kind,
+        UInt                offset,
+        UInt                space)
     {
-        if (astLayout == nullptr)
-        {
-            return nullptr;
-        }
+        IRInst* operands[3];
+        UInt operandCount = 0;
 
-        IRLayout* irLayout = nullptr;
-        if(sharedBuilder->layoutMap.TryGetValue(astLayout, irLayout))
-        {
-            SLANG_ASSERT(irLayout->getASTLayout() == astLayout);
-            return irLayout;
-        }
+        auto kindInst = getIntValue(getIntType(), IRIntegerValue(kind));
+        operands[operandCount++] = kindInst;
+
+        auto offsetInst = getIntValue(getIntType(), IRIntegerValue(offset));
+        operands[operandCount++] = offsetInst;
 
-        if (EntryPointLayout* entryPointLayout = as<EntryPointLayout>(astLayout))
+        if(space)
         {
-            irLayout = createInst<IREntryPointLayout>(this, kIROp_EntryPointLayout, nullptr, getPtrValue(astLayout));
+            auto spaceInst = getIntValue(getIntType(), IRIntegerValue(space));
+            operands[operandCount++] = spaceInst;
         }
-        else if (VarLayout* varLayout = as<VarLayout>(astLayout))
-        {
-            UnownedStringSlice nameSlice;
-            if (varLayout->getVariable())
-            {
-                Name* name = varLayout->getName();
-                if (name)
-                {
-                    nameSlice = name->text.getUnownedSlice();
-                }
-            }
 
-            // Get the name as a literal.
-            // We use an empty length string, as we can't use a null inst ptr.
-            // If there was a 'null' instruction then it might make more sense to use that
-            IRStringLit* nameLit = getStringValue(nameSlice);
-            
-            // Layout, name, type layout, absolute layout
-            IRInst* args[4] = {
-                getPtrValue(astLayout),
-                nameLit, 
-                getLayout(varLayout->getTypeLayout()),
-                getLayout(varLayout->m_absoluteLayout)
-            };
+        return cast<IRVarOffsetAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_VarOffsetAttr,
+            operandCount,
+            operands));
+    }
 
-            irLayout = createInst<IRVarLayout>(this, kIROp_VarLayout, nullptr, 4, args);
-        }
-        else if (TypeLayout* typeLayout = as<TypeLayout>(astLayout))
-        {
-            irLayout = createInst<IRTypeLayout>(this, kIROp_TypeLayout, nullptr, getPtrValue(astLayout));
-        }
-        else
-        {
-            SLANG_UNEXPECTED("Unknown layout type");
-        }
+    IRPendingLayoutAttr* IRBuilder::getPendingLayoutAttr(
+        IRLayout* pendingLayout)
+    {
+        IRInst* operands[] = { pendingLayout };
+
+        return cast<IRPendingLayoutAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_PendingLayoutAttr,
+            SLANG_COUNT_OF(operands),
+            operands));
+    }
+
+    IRStructFieldLayoutAttr* IRBuilder::getFieldLayoutAttr(
+        IRStructKey*    key,
+        IRVarLayout*    layout)
+    {
+        IRInst* operands[] = { key, layout };
+
+        return cast<IRStructFieldLayoutAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_StructFieldLayoutAttr,
+            SLANG_COUNT_OF(operands),
+            operands));
+    }
+
+    IRCaseTypeLayoutAttr* IRBuilder::getCaseTypeLayoutAttr(
+        IRTypeLayout*   layout)
+    {
+        IRInst* operands[] = { layout };
+
+        return cast<IRCaseTypeLayoutAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_CaseTypeLayoutAttr,
+            SLANG_COUNT_OF(operands),
+            operands));
+    }
+
+    IRSemanticAttr* IRBuilder::getSemanticAttr(
+        IROp            op,
+        String const&   name,
+        UInt            index)
+    {
+        auto nameInst = getStringValue(name.getUnownedSlice());
+        auto indexInst = getIntValue(getIntType(), index);
+
+        IRInst* operands[] = { nameInst, indexInst };
+
+        return cast<IRSemanticAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            op,
+            SLANG_COUNT_OF(operands),
+            operands));
+    }
+
+    IRStageAttr* IRBuilder::getStageAttr(Stage stage)
+    {
+        auto stageInst = getIntValue(getIntType(), IRIntegerValue(stage));
+        IRInst* operands[] = { stageInst };
+        return cast<IRStageAttr>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_StageAttr,
+            SLANG_COUNT_OF(operands),
+            operands));
+    }
 
-        SLANG_ASSERT(irLayout);
-        SLANG_ASSERT(irLayout->getASTLayout() == astLayout);
 
-        sharedBuilder->layoutMap[astLayout] = irLayout;
+    IRTypeLayout* IRBuilder::getTypeLayout(IROp op, List<IRInst*> const& operands)
+    {
+        return cast<IRTypeLayout>(findOrEmitHoistableInst(
+            getVoidType(),
+            op,
+            operands.getCount(),
+            operands.getBuffer()));
+    }
 
-        addGlobalValue(this, irLayout);
+    IRVarLayout* IRBuilder::getVarLayout(List<IRInst*> const& operands)
+    {
+        return cast<IRVarLayout>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_VarLayout,
+            operands.getCount(),
+            operands.getBuffer()));
+    }
 
-        // need to keep in scope 
-        addRefObjectToFree(astLayout);
+    IREntryPointLayout* IRBuilder::getEntryPointLayout(
+        IRVarLayout* paramsLayout,
+        IRVarLayout* resultLayout)
+    {
+        IRInst* operands[] = { paramsLayout, resultLayout };
 
-        return irLayout;
+        return cast<IREntryPointLayout>(findOrEmitHoistableInst(
+            getVoidType(),
+            kIROp_EntryPointLayout,
+            SLANG_COUNT_OF(operands),
+            operands));
     }
 
     //