Output SPIR-V lifetimes (#2221)

* #include an absolute path didn't work - because paths were taken to always be relative.

* WIP tracking liveness.

* Skeleton around adding liveness instructions.

* Calling into liveness tracking logic.
Adds live start to var insts.

* Liveness macros have initial output.

* Looking at different initialization scenarios.

* Some discussion around liveness.

* WIP for working out liveness end.

* WIP Updated liveness using use lists.

* Is now adding liveness information

* Some small fixes.

* WIP around liveness.

* Seems to output liveness correctly for current scenario.

* Tidy up liveness code.

* Update comment arounds liveness to current status.

* Small fixes to liveness test.

* Add support for call in liveness analysis.

* Improve liveness example with array access.

* Small updates to comments.

* Disable liveness test because inconsistencies with output on CI system.

* First pass support for GLSL SPIR-V liveness support.

* Add the SPIRVOpDecoration.

* Fix signature for OpLivenessStop.

* Simplified by having a Kind type.

* Fix some issues brought up in PR.

* Rename liveness instructions.

* Merge with var-lifetime.
Small improvements.

* Improvements to the documentation/naming in GLSL liveness pass.
Add comment around possible improvements to the liveness pass.

9 files changed, 336 insertions, 6 deletions

diff --git a/build/visual-studio/slang/slang.vcxproj b/build/visual-studio/slang/slang.vcxproj
index 162419d0c..52c2db316 100644
--- a/build/visual-studio/slang/slang.vcxproj
+++ b/build/visual-studio/slang/slang.vcxproj
@@ -357,6 +357,7 @@ IF EXIST ..\..\..\external\slang-binaries\bin\windows-aarch64\slang-glslang.dll\
     <ClInclude Include="..\..\..\source\slang\slang-ir-extract-value-from-type.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-generics-lowering-context.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-glsl-legalize.h" />

+    <ClInclude Include="..\..\..\source\slang\slang-ir-glsl-liveness.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-hoist-local-types.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-inline.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-inst-defs.h" />

@@ -496,6 +497,7 @@ IF EXIST ..\..\..\external\slang-binaries\bin\windows-aarch64\slang-glslang.dll\
     <ClCompile Include="..\..\..\source\slang\slang-ir-extract-value-from-type.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-generics-lowering-context.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-glsl-legalize.cpp" />

+    <ClCompile Include="..\..\..\source\slang\slang-ir-glsl-liveness.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-hoist-local-types.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-inline.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-layout.cpp" />

diff --git a/build/visual-studio/slang/slang.vcxproj.filters b/build/visual-studio/slang/slang.vcxproj.filters
index e6f7e21e9..1337c0e55 100644
--- a/build/visual-studio/slang/slang.vcxproj.filters
+++ b/build/visual-studio/slang/slang.vcxproj.filters
@@ -168,6 +168,9 @@
     <ClInclude Include="..\..\..\source\slang\slang-ir-glsl-legalize.h">

       <Filter>Header Files</Filter>

     </ClInclude>

+    <ClInclude Include="..\..\..\source\slang\slang-ir-glsl-liveness.h">

+      <Filter>Header Files</Filter>

+    </ClInclude>

     <ClInclude Include="..\..\..\source\slang\slang-ir-hoist-local-types.h">

       <Filter>Header Files</Filter>

     </ClInclude>

@@ -581,6 +584,9 @@
     <ClCompile Include="..\..\..\source\slang\slang-ir-glsl-legalize.cpp">

       <Filter>Source Files</Filter>

     </ClCompile>

+    <ClCompile Include="..\..\..\source\slang\slang-ir-glsl-liveness.cpp">

+      <Filter>Source Files</Filter>

+    </ClCompile>

     <ClCompile Include="..\..\..\source\slang\slang-ir-hoist-local-types.cpp">

       <Filter>Source Files</Filter>

     </ClCompile>

diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index 3b39c1cb1..5b2ae6d19 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -1831,6 +1831,40 @@ void GLSLSourceEmitter::emitVectorTypeNameImpl(IRType* elementType, IRIntegerVal
     }
 }
 
+void GLSLSourceEmitter::emitTypeImpl(IRType* type, const StringSliceLoc* nameAndLoc)
+{
+    if (auto refType = as<IRRefType>(type))
+    {
+        m_writer->emit("spirv_by_reference ");
+        type = refType->getValueType();
+    }
+    return Super::emitTypeImpl(type, nameAndLoc);
+}
+
+void GLSLSourceEmitter::emitParamTypeImpl(IRType* type, String const& name)
+{
+    if (auto refType = as<IRRefType>(type))
+    {
+        // 
+        m_writer->emit("spirv_by_reference  ");
+        type = refType->getValueType();
+    }
+    Super::emitParamTypeImpl(type, name);
+}
+
+void GLSLSourceEmitter::emitFuncDecorationImpl(IRDecoration* decoration)
+{
+    if (decoration->getOp() == kIROp_SPIRVOpDecoration)
+    {
+        m_writer->emit("spirv_instruction(id = ");
+        emitSimpleValue(decoration->getOperand(0));
+        m_writer->emit(")\n");
+    }
+    else
+    {
+        Super::emitFuncDecorationImpl(decoration);
+    }
+}
 
 void GLSLSourceEmitter::emitSimpleTypeImpl(IRType* type)
 {
diff --git a/source/slang/slang-emit-glsl.h b/source/slang/slang-emit-glsl.h
index c770cce34..481efed63 100644
--- a/source/slang/slang-emit-glsl.h
+++ b/source/slang/slang-emit-glsl.h
@@ -37,6 +37,9 @@ protected:
     virtual void emitVectorTypeNameImpl(IRType* elementType, IRIntegerValue elementCount) SLANG_OVERRIDE;
     virtual void emitVarDecorationsImpl(IRInst* varDecl) SLANG_OVERRIDE;
     virtual void emitMatrixLayoutModifiersImpl(IRVarLayout* layout) SLANG_OVERRIDE;
+    virtual void emitTypeImpl(IRType* type, const StringSliceLoc* nameAndLoc) SLANG_OVERRIDE;
+    virtual void emitParamTypeImpl(IRType* type, String const& name) SLANG_OVERRIDE;
+    virtual void emitFuncDecorationImpl(IRDecoration* decoration) SLANG_OVERRIDE;
 
     virtual void handleRequiredCapabilitiesImpl(IRInst* inst) SLANG_OVERRIDE;
 
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 0f43762c1..1b94c588b 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -39,6 +39,7 @@
 #include "slang-ir-validate.h"
 #include "slang-ir-wrap-structured-buffers.h"
 #include "slang-ir-liveness.h"
+#include "slang-ir-glsl-liveness.h"
 
 #include "slang-legalize-types.h"
 #include "slang-lower-to-ir.h"
@@ -762,6 +763,11 @@ Result linkAndOptimizeIR(
         addLivenessTrackingToModule(irModule);
 
         dumpIRIfEnabled(codeGenContext, irModule, "LIVENESS");
+
+        if (isKhronosTarget(targetRequest))
+        {
+            applyGLSLLiveness(irModule);
+        }
     }
 
     // We include one final step to (optionally) dump the IR and validate
diff --git a/source/slang/slang-ir-glsl-liveness.cpp b/source/slang/slang-ir-glsl-liveness.cpp
new file mode 100644
index 000000000..d7a3e0395
--- /dev/null
+++ b/source/slang/slang-ir-glsl-liveness.cpp
@@ -0,0 +1,240 @@
+#include "slang-ir-glsl-liveness.h"
+
+#include "slang-ir-insts.h"
+#include "slang-ir.h"
+
+#include "slang-ir-dominators.h"
+
+namespace Slang
+{
+
+namespace { // anonymous
+
+struct GLSLLivenessContext
+{
+    enum class Kind
+    {
+        Start,
+        End,
+        CountOf,
+    };
+
+        /// Process the module
+    void processModule();
+
+    GLSLLivenessContext(IRModule* module):
+        m_module(module)
+    {
+        m_sharedBuilder.init(module);
+        m_builder.init(m_sharedBuilder);
+    }
+
+    void _replaceMarker(IRLiveRangeMarker* liveMarker);
+    void _addDecorations(Kind kind, IRFunc* func);
+
+        /// Process a function in the module
+    void _processFunction(IRFunc* funcInst);
+
+    IRType* _getReferencedType(IRInst* referenced);
+
+    Kind getKind(IROp op)
+    {
+        switch (op)
+        {
+            case kIROp_LiveRangeStart: return Kind::Start;
+            case kIROp_LiveRangeEnd:   return Kind::End;
+            default: break;
+        }
+        SLANG_UNREACHABLE("Invalid op");
+    }
+
+    // Entry holds information about each of the function kinds
+    struct Entry
+    {
+        Dictionary<IRType*, IRFunc*> m_funcs;       ///< Map of the parameter type to functions implementing (for the kind)
+        IRStringLit* m_nameHintLiteral = nullptr;   ///< Name hint string literal of the function 
+        IRInst* m_spirvOpLiteral = nullptr;         ///< The SPIR-V opcode for the kind
+    };
+
+    List<IRLiveRangeMarker*> m_markerInsts;         ///< All of the liveness marker instrucitons found
+
+    Entry m_entries[Index(Kind::CountOf)];          /// Entry for each kind of function
+
+    IRStringLit* m_extensionStringLiteral = nullptr;        ///< The string literal holding the SPIR-V extension needed 
+    IRInst* m_zeroIntLiteral = nullptr;                     ///< Zero value literal
+
+    IRModule* m_module;
+    SharedIRBuilder m_sharedBuilder;
+    IRBuilder m_builder;
+};
+
+void GLSLLivenessContext::_processFunction(IRFunc* funcInst)
+{
+    // Iterate through blocks in the function, looking for variables to live track
+    for (auto block = funcInst->getFirstBlock(); block; block = block->getNextBlock())
+    {
+        for (auto inst = block->getFirstChild(); inst; inst = inst->getNextInst())
+        {
+            IRLiveRangeMarker* marker = as<IRLiveRangeMarker>(inst);
+            if (marker)
+            {
+                m_markerInsts.add(marker);
+            }
+        }
+    }
+}
+
+void GLSLLivenessContext::_addDecorations(Kind kind, IRFunc* func)
+{
+    // We might(?) want to add a decoration saying this is GLSL specific, but at this point
+    // we can only be in GLSL dependent IR.
+    //
+    // m_builder.addTargetDecoration();
+
+    // We need the spirv extension
+    m_builder.addDecoration(func, kIROp_RequireGLSLExtensionDecoration, m_extensionStringLiteral);
+
+    const auto& entry = m_entries[Index(kind)];
+    if (entry.m_nameHintLiteral)
+    {
+        m_builder.addNameHintDecoration(func, entry.m_nameHintLiteral);
+    }
+
+    m_builder.addDecoration(func, kIROp_SPIRVOpDecoration, entry.m_spirvOpLiteral);
+}
+
+IRType* GLSLLivenessContext::_getReferencedType(IRInst* referenced)
+{
+    auto type = referenced->getDataType();
+
+    if (type->getOp() == kIROp_PtrType)
+    {
+        type = static_cast<IRPtrType*>(type)->getValueType();
+    }
+
+    return type;
+}
+
+void GLSLLivenessContext::_replaceMarker(IRLiveRangeMarker* markerInst)
+{
+    const auto kind = getKind(markerInst->getOp());
+    auto& entry = m_entries[Index(kind)];
+
+    IRInst* referenced = markerInst->getReferenced();
+    IRType* referencedType = _getReferencedType(referenced);
+
+    IRFunc* func = nullptr;
+    
+    if (IRFunc** funcPtr = entry.m_funcs.TryGetValue(referencedType))
+    {
+        func = *funcPtr;
+    }
+    else
+    {
+        // We didn't find a function for the type, so lets create one. It has a signature of
+        //
+        // void func(Ref<ReferencedType> target, int sizeInBytes)
+
+        IRType* paramTypes[] = 
+        {
+            m_builder.getRefType(referencedType),       ///< Use a reference to the referenced type
+            m_builder.getIntType(),                     ///< The size type
+        };
+
+        func = m_builder.createFunc();
+
+        auto funcType = m_builder.getFuncType(SLANG_COUNT_OF(paramTypes), paramTypes, m_builder.getVoidType());
+        m_builder.setDataType(func, funcType);
+
+        // Add any decorations to the new function
+        _addDecorations(kind, func);
+
+        // Add to the map
+        entry.m_funcs.Add(referencedType, func);
+    }
+    SLANG_ASSERT(func);
+
+    // Create a call to the function in the form of...
+    // func(referencedItem, 0);
+
+    // As per the SPIR-V documentation around the OpLifetimeStart/OpLifetimeEnd
+    // https://www.khronos.org/registry/SPIR-V/specs/unified1/SPIRV.html#OpLifetimeStart
+    // 
+    // If the type is known the size should be passed as 0
+    IRInst* args[] = 
+    {
+        referenced,
+        m_zeroIntLiteral,
+    };
+
+    // Set the location at the marker to add the call
+    m_builder.setInsertLoc(IRInsertLoc::after(markerInst));
+    m_builder.emitCallInst(m_builder.getVoidType(), func, SLANG_COUNT_OF(args), args);
+
+    // We don't need the marker anymore
+    markerInst->removeAndDeallocate();
+}
+
+void GLSLLivenessContext::processModule()
+{
+    // Find all of the liveness marker insts
+    //
+    // This is done prior to processing, so we don't need to worry about traversal when
+    // instructions are replaced.
+
+    IRModuleInst* moduleInst = m_module->getModuleInst();
+    for (IRInst* child : moduleInst->getChildren())
+    {
+        // We want to find all of the functions, and process them
+        if (auto funcInst = as<IRFunc>(child))
+        {
+            // Then we want to look through their definition
+            // inserting instructions that mark the liveness start/end
+            _processFunction(funcInst);
+        }
+    }
+    
+    // If we didn't find any liveness marker instructions then we are done
+    if (!m_markerInsts.getCount())
+    {
+        return;
+    }
+
+    // Zero value literal
+    m_zeroIntLiteral = m_builder.getIntValue(m_builder.getIntType(), 0);
+
+    // Set up some values that will be needed on instructions
+    m_extensionStringLiteral = m_builder.getStringValue(UnownedStringSlice::fromLiteral("GL_EXT_spirv_intrinsics"));
+
+    // The op values are from the SPIR-V spec
+    // https://www.khronos.org/registry/SPIR-V/specs/unified1/SPIRV.html#OpLifetimeStart
+
+    {
+        auto& entry = m_entries[Index(Kind::Start)];
+        entry.m_nameHintLiteral = m_builder.getStringValue(UnownedStringSlice::fromLiteral("livenessStart"));
+        entry.m_spirvOpLiteral = m_builder.getIntValue(m_builder.getIntType(), 256);
+    }
+    {
+        auto& entry = m_entries[Index(Kind::End)];
+        entry.m_nameHintLiteral = m_builder.getStringValue(UnownedStringSlice::fromLiteral("livenessEnd"));
+        entry.m_spirvOpLiteral = m_builder.getIntValue(m_builder.getIntType(), 257);
+    }
+
+    // Iterate across instructions, replacing with a call to a generated function (one that just is a declaration defining the SPIR-V op)
+    for (auto markerInst : m_markerInsts)
+    {
+        _replaceMarker(markerInst);
+    }
+}
+
+} // anonymous
+
+void applyGLSLLiveness(IRModule* module)
+{
+    GLSLLivenessContext context(module);
+
+    context.processModule();
+}
+
+} // namespace Slang
+    
\ No newline at end of file
diff --git a/source/slang/slang-ir-glsl-liveness.h b/source/slang/slang-ir-glsl-liveness.h
new file mode 100644
index 000000000..401df57fc
--- /dev/null
+++ b/source/slang/slang-ir-glsl-liveness.h
@@ -0,0 +1,20 @@
+// slang-ir-liveness.h
+#ifndef SLANG_IR_GLSL_LIVENESS_H
+#define SLANG_IR_GLSL_LIVENESS_H
+
+namespace Slang
+{
+
+struct IRModule;
+
+	/// Converts liveness marker instructions in a module into SPIR-V lifetime ops.
+	/// 
+	/// It does this by using the GL_EXT_spirv_intrinsics extension. 
+	///
+	/// The transformation takes place at the IR level, inserting new functions as needed for 
+	/// the types referenced via liveness markers.
+void applyGLSLLiveness(IRModule* module);
+
+}
+
+#endif // SLANG_IR_LIVENESS_H
\ No newline at end of file
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index d5cb49046..fa977fd89 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/source/slang/slang-ir-inst-defs.h
@@ -638,6 +638,8 @@ INST(HighLevelDeclDecoration,               highLevelDecl,          1, 0)
 
     INST(SemanticDecoration, semantic, 2, 0)
 
+    INST(SPIRVOpDecoration, spirvOpDecoration, 1, 0)
+
         /// Marks a struct type as being used as a structured buffer block.
         /// Recognized by SPIRV-emit pass so we can emit a SPIRV `BufferBlock` decoration.
     INST(SPIRVBufferBlockDecoration, spvBufferBlock, 0, 0)
diff --git a/source/slang/slang-ir-liveness.cpp b/source/slang/slang-ir-liveness.cpp
index 3d1626691..7667a33f0 100644
--- a/source/slang/slang-ir-liveness.cpp
+++ b/source/slang/slang-ir-liveness.cpp
@@ -83,6 +83,23 @@ is created that is then just copied into the variable. That temporary being some
 could perhaps have liveness issues.
 */
 
+/* This implementation could potentially be improved in a few ways:
+
+## Use Dominator tree block indexing
+
+The dominator tree has a mapping from block insts to integer values. These integer value could be used to allow information about blocks to 
+be stored in an array. A downside, would be exposing that aspect of the implementation to the dominator tree.
+
+## Store if a block has an access instruction
+
+As it stands, when traversing the tree to find the last access, the implementation searches for each block backwards to find the last access
+by looking if it's in the m_accessSet. 
+
+This searching could be avoided, by when accesses are added the block they are in is marked as having an access. If a block had no accesses
+this would remove the linear search for the first access instruction, if the block indicates it doesn't have any. The downside being that every 
+access would also have to mark the block it is in.
+*/
+
 namespace { // anonymous
 
 struct LivenessContext
@@ -118,9 +135,6 @@ struct LivenessContext
         /// Process a 'root'. A variable that has liveness tracking
     void processRoot(const RootInfo& rootInfo);
 
-        /// Process a function in the module
-    void processFunction(IRFunc* funcInst);
-
         /// Process the module
     void processModule();
 
@@ -131,6 +145,9 @@ struct LivenessContext
         m_builder.init(m_sharedBuilder);
     }
 
+        /// Process a function in the module
+    void _processFunction(IRFunc* funcInst);
+
         // Add a live end instruction at the start of block, referencing the root 'root'.
     void _addLiveRangeEndAtBlockStart(IRBlock* block, IRInst* root);
 
@@ -463,7 +480,7 @@ void LivenessContext::processRoot(const RootInfo& rootInfo)
     }
 }
 
-void LivenessContext::processFunction(IRFunc* funcInst)
+void LivenessContext::_processFunction(IRFunc* funcInst)
 {
     List<RootInfo> rootInfos;
 
@@ -508,14 +525,14 @@ void LivenessContext::processModule()
 
     IRModuleInst* moduleInst = m_module->getModuleInst();
 
-    for (IRInst* child = moduleInst->getFirstDecorationOrChild(); child; child = child->getNextInst())
+    for (IRInst* child : moduleInst->getChildren())
     {
         // We want to find all of the functions, and process them
         if (auto funcInst = as<IRFunc>(child))
         {
             // Then we want to look through their definition
             // inserting instructions that mark the liveness start/end
-            processFunction(funcInst);
+            _processFunction(funcInst);
         }
     }
 }