Special handling around return and liveness (#2234)

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

* Refactor Liveness pass, such that locations can be found independently of setting up ranges.

* Refactor around different stages of liveness span analysis.

* WIP Take into account PHI temporaries in liveness tracking.

* WIP First pass of PHI liveness refactor.

* Add BlockIndex.

* WIP Refactor phi liveness around inst runs.

* More improvements around liveness tracking.

* Bug fixes.
Special handling to not add multiple ends, at starts of blocks and after accesses.

* Fix test output.

* Use IRInsertLoc to track insertion point.

* Liveness markers don't have side effects.

* Fix typo in liveness test.

* Small improvements around setting SuccessorResult.

* Fix memory issue around reallocation and RAIIStackArray.
Update test output.

* Update test output for liveness.slang.

* Fix typo in SuccessorResult blockIndex.

* Small tidy up.

* Handle the root start block, correctly scoping the run.

* Split BlockInfo into 'Root' and 'Function'.
Store successors as BlockIndices.

* Tidy up around liveness tracking.

* Add head/tail support to ArrayViews.
Use Count where appropriate.
Use head/tail in liveness impl.

* Special handling if return is effectively a live variable.

* Update test output for improved return handling.

* Refactor how handling of return accesses.
Fix issue around liveness starts.

* Disable release warning for unused method.

* Some small improvements around liveness pass.

3 files changed, 169 insertions, 72 deletions

diff --git a/source/core/slang-array-view.h b/source/core/slang-array-view.h
index f6ac92004..99609ef69 100644
--- a/source/core/slang-array-view.h
+++ b/source/core/slang-array-view.h
@@ -72,6 +72,11 @@ namespace Slang
             return -1;
         }
 
+        bool containsMemory(const ThisType& rhs) const
+        {
+            return rhs.getBuffer() >= getBuffer() && rhs.end() <= end(); 
+        }
+
         bool operator==(const ThisType& rhs) const
         {
             if (&rhs == this)
diff --git a/source/slang/slang-ir-liveness.cpp b/source/slang/slang-ir-liveness.cpp
index 1e88763fd..05fccbe43 100644
--- a/source/slang/slang-ir-liveness.cpp
+++ b/source/slang/slang-ir-liveness.cpp
@@ -183,6 +183,9 @@ struct LivenessContext
     LivenessContext(IRModule* module):
         m_module(module)
     {
+        // Disable warning if not used
+        SLANG_UNUSED(&LivenessContext::_isAnyRunInst);
+        
         m_sharedBuilder.init(module);
         m_builder.init(m_sharedBuilder);
     }
@@ -228,7 +231,10 @@ struct LivenessContext
     void _addInst(IRInst* inst);
 
         /// True if it's an instruction of interest and so will go within a run for a block
-    bool _isRunInst(IRInst* inst);
+    bool _isNormalRunInst(IRInst* inst);
+
+        /// Returns true if is a normal run inst, or if is a return that accesses
+    bool _isAnyRunInst(IRInst* inst);
 
         // Returns the index in the run of a start for the current root, else -1
     Index _indexOfRootStart(const ConstArrayView<IRInst*>& run);
@@ -241,8 +247,11 @@ struct LivenessContext
 
     void _maybeAddEndBeforeInst(IRInst* inst);
 
+        /// Maybe insert an end after the instruction
+    void _maybeAddEndAfterRunIndex(BlockIndex blockIndex, const ConstArrayView<IRInst*>& run, Index runIndex);
+
         // Add a live end instruction at the start of block, referencing the root
-    void _maybeAddEndAtBlockStart(IRBlock* block);
+    void _maybeAddEndAtBlockStart(BlockIndex blockIndex);
 
         /// Look from inst for an LiveEndRange to the root.
     IRInst* _findRootEnd(IRInst* inst);
@@ -256,6 +265,10 @@ struct LivenessContext
         /// Get the block from the index
     IRBlock* _getBlock(BlockIndex blockIndex) const { return m_functionBlockInfos[Index(blockIndex)].block; }
 
+        /// True if the terminator can be considered an access
+        /// This allows us to elide a scope end if the root is returned 
+    bool _isAccessTerminator(IRTerminatorInst* terminator);
+
         /// Gets the instructions of interest for this info, in the order they appear within the block
     ConstArrayView<IRInst*> _getRun(const BlockInfo* info) 
     {
@@ -297,8 +310,10 @@ struct LivenessContext
     IRBuilder m_builder;
 };
 
-void LivenessContext::_maybeAddEndAtBlockStart(IRBlock* block)
+void LivenessContext::_maybeAddEndAtBlockStart(BlockIndex blockIndex)
 {
+    auto block = _getBlock(blockIndex);
+
     // Insert before the first ordinary inst
     auto inst = block->getFirstOrdinaryInst();
     // A block has to end with a terminator... so must always be an ordinary inst, if there is a function body
@@ -403,11 +418,12 @@ Index LivenessContext::_findLastAccess(const ConstArrayView<IRInst*>& run)
 {
     for (Index i = run.getCount() - 1; i >= 0; --i)
     {
+        auto inst = run[i];
         // If it's not a live start, it must be an access
-        if (as<IRLiveRangeStart>(run[i]) == nullptr)
+        if (as<IRLiveRangeStart>(inst) == nullptr)
         {
             // Check it really is an access inst..
-            SLANG_ASSERT(_isRunInst(run[i]));
+            SLANG_ASSERT(_isAnyRunInst(inst));
             return i;
         }
     }
@@ -435,9 +451,18 @@ IRInst* LivenessContext::_findRootEnd(IRInst* inst)
     return nullptr;
 }
 
+void LivenessContext::_maybeAddEndAfterRunIndex(BlockIndex blockIndex, const ConstArrayView<IRInst*>& run, Index runIndex)
+{
+    SLANG_UNUSED(blockIndex);
+    return _maybeAddEndAfterInst(run[runIndex]);
+}
+
 void LivenessContext::_maybeAddEndAfterInst(IRInst* inst)
 {
-    if (!_findRootEnd(inst->getNextInst()))
+    // We can't add after the inst, if it's a terminator
+    // or if we find an end. 
+    if (as<IRTerminatorInst>(inst) == nullptr && 
+        !_findRootEnd(inst->getNextInst()))
     {
         // Just add end of scope after the inst 
         m_builder.setInsertLoc(IRInsertLoc::after(inst));
@@ -468,11 +493,8 @@ LivenessContext::BlockResult LivenessContext::_completeBlock(BlockIndex blockInd
     // If we found one, that is the end of the range
     if (lastAccessIndex >= 0)
     {
-        IRInst* lastAccessInst = run[lastAccessIndex];
-
-        // Insert an end after the last access (if not one already)
-        _maybeAddEndAfterInst(lastAccessInst);
-
+        _maybeAddEndAfterRunIndex(blockIndex, run, lastAccessIndex);
+        
         // Add the result
         return _addBlockResult(blockIndex, BlockResult::Found);
     }
@@ -560,7 +582,7 @@ LivenessContext::BlockResult LivenessContext::_processBlock(BlockIndex blockInde
             if (successorResult == BlockResult::NotFound)
             {
                 const auto successorBlockIndex = successors[i];
-                _maybeAddEndAtBlockStart(_getBlock(successorBlockIndex));
+                _maybeAddEndAtBlockStart(successorBlockIndex);
                 _addBlockResult(successorBlockIndex, BlockResult::Found);
             }
         }
@@ -703,6 +725,7 @@ void LivenessContext::_findAliasesAndAccesses(IRInst* root)
                     // These will never take place on the var which is accessed through a pointer, so can be ignored
                     break;
                 }
+                
                 default: break;
             }
 
@@ -741,9 +764,6 @@ void LivenessContext::_findAndEmitRangeEnd(IRLiveRangeStart* liveRangeStart)
     m_rootLiveStart = liveRangeStart;
     m_rootLiveStartBlock = as<IRBlock>(liveRangeStart->getParent());
 
-    m_root = liveRangeStart->getReferenced();
-    m_rootBlock = as<IRBlock>(m_root->parent);
-
     // If either of these asserts fail it probably means there hasn't been a call
     // to `_findAliasesAndAccesses` which is required before this function can be called.
     // 
@@ -794,23 +814,22 @@ void LivenessContext::_findAndEmitRangeEnd(IRLiveRangeStart* liveRangeStart)
 
     // Set back to nullptr for safety
     m_rootLiveStart = nullptr;
-    m_root = nullptr;
-    m_rootBlock = nullptr;
     m_rootLiveStartBlock = nullptr;
 }
 
-bool LivenessContext::_isRunInst(IRInst* inst)
+bool LivenessContext::_isNormalRunInst(IRInst* inst)
 {
     const auto op = inst->getOp();
 
-    // If it's a live start then we need to track
+    // Detect if it's a range start for the root.
     if (op == kIROp_LiveRangeStart)
     {
-        return true;
+        auto start = as<IRLiveRangeStart>(inst);
+        return start->getReferenced() == m_root;
     }
 
     // NOTE!
-    // These are the only ops *currently* that indicate an access.
+    // The ops in th elist above are the only ops *currently* that indicate an access.
     // Has to be consistent with `_findAliasesAndAccesses`
     if (op == kIROp_Call || 
         op == kIROp_Load || 
@@ -824,26 +843,81 @@ bool LivenessContext::_isRunInst(IRInst* inst)
     return false;
 }
 
+bool LivenessContext::_isAccessTerminator(IRTerminatorInst* terminator)
+{
+    // This is to special case when a return, returns a root or an alias
+    // 
+    // We need to detect if the return value accesses the root
+
+    if (terminator->getOp() == kIROp_ReturnVal)
+    {
+        // We are going to special case return if it hits an alias
+        auto returnVal = static_cast<IRReturnVal*>(terminator);
+
+        auto val = returnVal->getVal();
+        
+        // TODO(JS): This is perhaps somewhat argable, but it means if 
+        // we have a cast between uint/int (for example) that isn't a problem
+
+        // Strip construct
+        if (val->getOp() == kIROp_Construct && val->getOperandCount() == 1)
+        {
+            val = val->getOperand(0);
+        }
+
+        // If it *is* the root it's an access
+        if (val == m_root)
+        {
+            return true;
+        }
+
+        // If it's a load, see what is being loaded from an alias to the root
+        if (auto load = as<IRLoad>(val))
+        {
+            const auto valPtr = load->getPtr();
+            return m_aliases.contains(valPtr);
+        }
+    }
+
+    return false;
+}
+
+bool LivenessContext::_isAnyRunInst(IRInst* inst)
+{
+    if (auto terminator = as<IRTerminatorInst>(inst))
+    {
+        return _isAccessTerminator(terminator);
+    }
+    return _isNormalRunInst(inst);
+}
+
 void LivenessContext::_findInstRunsForBlocks()
 {
     const Count count = m_blockInfos.getCount();
     for (Index i = 0; i < count; ++i)
     {
-        auto& blockInfo = m_blockInfos[i];
+        const auto blockIndex = BlockIndex(i);
+
+        // Get the block
+        auto block = _getBlock(blockIndex);
+
+        // Get the block info
+        auto* blockInfo = _getBlockInfo(blockIndex);
 
-        if (blockInfo.instCount == 0)
+        const auto start = m_instRuns.getCount();
+        blockInfo->runStart = start;
+
+        if (blockInfo->instCount == 0)
         {
             // Nothing to do if it's empty
-            SLANG_ASSERT(blockInfo.runCount == 0);
+            SLANG_ASSERT(blockInfo->runCount == 0);
         }
-        else if (blockInfo.instCount == 1)
+        else if (blockInfo->instCount == 1)
         {
-            // This is the easy case, since we don't need to determine the order of the instructions
-            blockInfo.runStart = m_instRuns.getCount();
-            blockInfo.runCount = 1;
-            SLANG_ASSERT(blockInfo.lastInst);
-            m_instRuns.add(blockInfo.lastInst);
-            continue;
+            // This is the easy case, since we don't need to determine the order of the instructions   
+            SLANG_ASSERT(blockInfo->lastInst);
+            m_instRuns.add(blockInfo->lastInst);
+            blockInfo->runCount = 1;
         }
         else
         {
@@ -853,19 +927,15 @@ void LivenessContext::_findInstRunsForBlocks()
             //
             // For now we just add them all.
 
-            const auto start = m_instRuns.getCount();
-            blockInfo.runStart = start;
-            blockInfo.runCount = blockInfo.instCount;
+            blockInfo->runCount = blockInfo->instCount;
 
-            m_instRuns.setCount(start + blockInfo.instCount);
+            m_instRuns.setCount(start + blockInfo->instCount);
             IRInst** dst = m_instRuns.getBuffer() + start;
 
             // Find all of the instructions of interest in order
-            auto block = _getBlock(BlockIndex(i));
-
             for (auto inst : block->getChildren())
             {
-                if (_isRunInst(inst))
+                if (_isNormalRunInst(inst))
                 {
                     *dst++ = inst;
                     if (dst == m_instRuns.end())
@@ -875,10 +945,30 @@ void LivenessContext::_findInstRunsForBlocks()
                 }
             }
             SLANG_ASSERT(dst == m_instRuns.end());
+        }   
+
+        SLANG_ASSERT(blockInfo->runCount == blockInfo->instCount);
+
+        // Special case the terminator - we allow a return that accesses the root 
+        // to be added to the run. 
+        //
+        // TODO(JS): We might want this behavior to be switchable with an option.
+        // If we don't add the terminator, everything else will behave correctly with regard
+        // adding live range end markers.
+        {
+            auto terminator = block->getTerminator();
+            if (_isAccessTerminator(terminator))
+            {
+                m_instRuns.add(terminator);
+                blockInfo->runCount++;
+            }
         }
 
-        // The run count and the inst count must match at this point
-        SLANG_ASSERT(blockInfo.runCount == blockInfo.instCount);
+        SLANG_ASSERT(blockInfo->runStart + blockInfo->runCount == m_instRuns.getCount());
+
+        // The run count must be at least as many as the found instCount
+        // There can be more instructions as we allow some special cases (for example around return)
+        SLANG_ASSERT(blockInfo->runCount >= blockInfo->instCount);
     }
 }
 
@@ -898,6 +988,10 @@ void LivenessContext::_processRoot(const Location* locations, Count locationsCou
     
     auto root = locations[0].root;
 
+    // Set the root
+    m_root = root;
+    m_rootBlock = as<IRBlock>(m_root->parent);
+
     // Add all the live starts
     m_liveRangeStarts.setCount(locationsCount);
     for (Index i = 0; i < locationsCount; ++i)
@@ -928,6 +1022,10 @@ void LivenessContext::_processRoot(const Location* locations, Count locationsCou
         // We want to process this RangeStart for the root, to find all of the ends
         _findAndEmitRangeEnd(liveStart);
     }
+
+    // No root is active in processing
+    m_root = nullptr;
+    m_rootBlock = nullptr;
 }
 
 void LivenessContext::_processLocationsInFunction(const Location* locations, Count count)
diff --git a/tests/experimental/liveness/liveness.slang.expected b/tests/experimental/liveness/liveness.slang.expected
index 77f14714e..7301d88bf 100644
--- a/tests/experimental/liveness/liveness.slang.expected
+++ b/tests/experimental/liveness/liveness.slang.expected
@@ -102,15 +102,9 @@ int someSlowFunc_0(int a_0)
 #line 27
     livenessEnd_1(i_0, 0);
 
-#line 27
-    uint _S6 = v_0;
-
-#line 27
-    livenessEnd_0(v_0, 0);
-
 
 
-    return int(_S6);
+    return int(v_0);
 }
 
 
@@ -126,16 +120,16 @@ struct SomeStruct_0
 #line 17
 SomeStruct_0 makeSomeStruct_0()
 {
-    const int  _S7[100] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
+    const int  _S6[100] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
 
 #line 19
-    SomeStruct_0 s_0 = { 0, 0, _S7 };
+    SomeStruct_0 s_0 = { 0, 0, _S6 };
     return s_0;
 }
 
 
 #line 15
-layout(std430, binding = 1) buffer _S8 {
+layout(std430, binding = 1) buffer _S7 {
     int _data[];
 } anotherBuffer_0;
 
@@ -155,7 +149,7 @@ int somethingElse_0(inout SomeStruct_0 s_2)
 
 
 #line 4
-layout(std430, binding = 0) buffer _S9 {
+layout(std430, binding = 0) buffer _S8 {
     int _data[];
 } outputBuffer_0;
 
@@ -221,14 +215,14 @@ void main()
         livenessStart_2(t_0, 0);
 
 #line 57
-        SomeStruct_0 _S10 = makeSomeStruct_0();
+        SomeStruct_0 _S9 = makeSomeStruct_0();
 
 #line 57
-        t_0 = _S10;
-        const int  _S11[100] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
+        t_0 = _S9;
+        const int  _S10[100] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
 
 #line 58
-        SomeStruct_0 u_0 = { 0, 0, _S11 };
+        SomeStruct_0 u_0 = { 0, 0, _S10 };
 
 #line 58
         SomeStruct_0 u_1;
@@ -258,7 +252,7 @@ void main()
 #line 68
             x_1 = u_0;
             x_1.x_0 = ((anotherBuffer_0)._data[(uint(v_1 & 3))]) + 1;
-            SomeStruct_0 _S12 = x_1;
+            SomeStruct_0 _S11 = x_1;
 
 #line 70
             livenessEnd_2(x_1, 0);
@@ -267,7 +261,7 @@ void main()
             livenessStart_2(u_1, 0);
 
 #line 70
-            u_1 = _S12;
+            u_1 = _S11;
 
 #line 70
         }
@@ -276,58 +270,58 @@ void main()
 
         s_3.c_0[index_0 & 7] = s_3.c_0[index_0 & 7] + 1;
 
-        int _S13 = s_3.x_0 + t_0.x_0;
+        int _S12 = s_3.x_0 + t_0.x_0;
 
 #line 76
-        SomeStruct_0 _S14 = u_1;
+        SomeStruct_0 _S13 = u_1;
 
 #line 76
         livenessEnd_2(u_1, 0);
 
 #line 76
-        int _S15 = _S13 + _S14.x_0;
+        int _S14 = _S12 + _S13.x_0;
 
 #line 76
-        int _S16 = doThing_0(t_0);
+        int _S15 = doThing_0(t_0);
 
 #line 76
-        int _S17 = _S15 + _S16;
+        int _S16 = _S14 + _S15;
 
 #line 76
-        int _S18 = somethingElse_0(t_0);
+        int _S17 = somethingElse_0(t_0);
 
 #line 76
         livenessEnd_2(t_0, 0);
 
 #line 76
-        int _S19 = _S17 + _S18;
+        int _S18 = _S16 + _S17;
 
 #line 76
-        int _S20 = s_3.c_0[2];
+        int _S19 = s_3.c_0[2];
 
 #line 76
         livenessEnd_2(s_3, 0);
 
 #line 76
-        int _S21 = _S19 + _S20;
+        int _S20 = _S18 + _S19;
 
 #line 76
-        int _S22 = res_0;
+        int _S21 = res_0;
 
 #line 76
         livenessEnd_1(res_0, 0);
 
 #line 76
-        int res_1 = _S22 + _S21;
+        int res_1 = _S21 + _S20;
 
 #line 52
-        int _S23 = i_2;
+        int _S22 = i_2;
 
 #line 52
         livenessEnd_1(i_2, 0);
 
 #line 52
-        int i_3 = _S23 + 1;
+        int i_3 = _S22 + 1;
 
 #line 52
         livenessStart_0(i_2, 0);
@@ -348,13 +342,13 @@ void main()
     livenessEnd_1(i_2, 0);
 
 #line 52
-    int _S24 = res_0;
+    int _S23 = res_0;
 
 #line 52
     livenessEnd_1(res_0, 0);
 
 #line 79
-    ((outputBuffer_0)._data[(uint(index_0))]) = _S24;
+    ((outputBuffer_0)._data[(uint(index_0))]) = _S23;
     return;
 }