Fix DCE on mutable calls in a loop. (#2943)

* Fix DCE on mutable calls in a loop.

* More accurate in-loop test.

* code review  fixes.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

source/slang/slang-ir-dce.cpp

@@ -104,6 +104,9 @@ struct DeadCodeEliminationContext
     bool processInst(IRInst* root)
     {
         bool result = false;
+
+        module->invalidateAllAnalysis();
+
         for (;;)
         {
             liveInsts.clear();
@@ -185,6 +188,12 @@ struct DeadCodeEliminationContext
                 // decision of whether a child (or decoration)
                 // should be live when its parent is to a subroutine.
                 //
+
+                if (auto func = as<IRGlobalValueWithCode>(inst))
+                {
+                    module->findOrCreateDominatorTree(func);
+                }
+
                 for (auto child : inst->getDecorationsAndChildren())
                 {
                     if (shouldInstBeLiveIfParentIsLive(child))
@@ -208,6 +217,8 @@ struct DeadCodeEliminationContext
             //
             phiRemoved = false;
             result |= eliminateDeadInstsRec(root);
+
+
             if (!phiRemoved)
                 break;
         }
@@ -271,6 +282,12 @@ struct DeadCodeEliminationContext
             {
                 changed |= eliminateDeadInstsRec(child);
             }
+            if (changed)
+            {
+                // If the function body is changed, invalidate its dominator tree.
+                if (auto func = as<IRGlobalValueWithCode>(inst))
+                    module->invalidateAnalysisForInst(func);
+            }
         }
         return changed;
     }

source/slang/slang-ir-propagate-func-properties.cpp

@@ -14,6 +14,26 @@ class FuncPropertyPropagationContext
     virtual bool propagate(IRBuilder& builder, IRFunc* func) = 0;
 };
 
+static bool isKnownOpCodeWithSideEffect(IROp op)
+{
+    switch (op)
+    {
+    case kIROp_ifElse:
+    case kIROp_unconditionalBranch:
+    case kIROp_Switch:
+    case kIROp_Return:
+    case kIROp_loop:
+    case kIROp_Call:
+    case kIROp_Param:
+    case kIROp_Unreachable:
+    case kIROp_Store:
+    case kIROp_SwizzledStore:
+        return true;
+    default:
+        return false;
+    }
+}
+
 class ReadNoneFuncPropertyPropagationContext : public FuncPropertyPropagationContext
 {
 public:
@@ -40,22 +60,10 @@ class ReadNoneFuncPropertyPropagationContext : public FuncPropertyPropagationCon
             for (auto inst : block->getChildren())
             {
                 // Is this inst known to not have global side effect/analyzable?
-                if (inst->mightHaveSideEffects())
+                if (!isKnownOpCodeWithSideEffect(inst->getOp()))
                 {
-                    switch (inst->getOp())
+                    if (inst->mightHaveSideEffects())
                     {
-                    case kIROp_ifElse:
-                    case kIROp_unconditionalBranch:
-                    case kIROp_Switch:
-                    case kIROp_Return:
-                    case kIROp_loop:
-                    case kIROp_Call:
-                    case kIROp_Param:
-                    case kIROp_Unreachable:
-                    case kIROp_Store:
-                    case kIROp_SwizzledStore:
-                        break;
-                    default:
                         // We have a inst that has side effect and is not understood by this method.
                         // e.g. bufferStore, discard, etc.
                         hasSideEffectCall = true;
@@ -238,33 +246,21 @@ class NoSideEffectFuncPropertyPropagationContext : public FuncPropertyPropagatio
         {
             for (auto inst : block->getChildren())
             {
-                // Is this inst known to not have global side effect/analyzable?
-                if (inst->mightHaveSideEffects())
+                if (!isKnownOpCodeWithSideEffect(inst->getOp()))
                 {
-                    switch (inst->getOp())
+                    // Is this inst known to not have global side effect/analyzable?
+                    if (inst->mightHaveSideEffects())
                     {
-                    case kIROp_ifElse:
-                    case kIROp_unconditionalBranch:
-                    case kIROp_Switch:
-                    case kIROp_Return:
-                    case kIROp_loop:
-                    case kIROp_Call:
-                    case kIROp_Param:
-                    case kIROp_Unreachable:
-                    case kIROp_Store:
-                    case kIROp_SwizzledStore:
-                        break;
-                    default:
                         // We have a inst that has side effect and is not understood by this method.
                         // e.g. bufferStore, discard, etc.
                         hasSideEffectCall = true;
                         break;
                     }
-                }
-                else
-                {
-                    // A side effect free inst can't generate side effects for the function.
-                    continue;
+                    else
+                    {
+                        // A side effect free inst can't generate side effects for the function.
+                        continue;
+                    }
                 }
 
                 if (auto call = as<IRCall>(inst))

source/slang/slang-ir-simplify-cfg.cpp

@@ -742,6 +742,12 @@ static bool processFunc(IRGlobalValueWithCode* func)
         if (!blocksRemoved)
             break;
     }
+    if (changed)
+    {
+        auto module = func->getModule();
+        if (module)
+            module->invalidateAnalysisForInst(func);
+    }
     return changed;
 }
 

source/slang/slang-ir-util.cpp

@@ -665,6 +665,11 @@ bool areCallArgumentsSideEffectFree(IRCall* call)
                 return false;
         }
 
+        auto module = parentFunc->getModule();
+        if (!module)
+            return false;
+        auto dom = module->findDominatorTree(parentFunc);
+
         if (arg->getOp() == kIROp_Var && getParentFunc(arg) == parentFunc)
         {
             // If the pointer argument is a local variable (thus can't alias with other addresses)
@@ -688,9 +693,49 @@ bool areCallArgumentsSideEffectFree(IRCall* call)
                     // are not dependent on whatever we do in the call here.
                     continue;
                 default:
-                    // Skip the call itself, since we are checking if the call has side effect.
+                    // Skip the call itself if the var is used as an argument to an out parameter
+                    // since we are checking if the call has side effect.
+                    // We can't treat the call as side effect free if var is used as an inout parameter,
+                    // because if the call is inside a loop there will be a visible side effect after
+                    // the call.
                     if (use->getUser() == call)
+                    {
+                        auto funcType = as<IRFuncType>(call->getCallee()->getDataType());
+                        if (!funcType)
+                            return false;
+                        if (funcType->getParamCount() > i && as<IROutType>(funcType->getParamType(i)))
+                            continue;
+
+                        // We are an argument to an inout parameter.
+                        // We can only treat the call as side effect free if the call is not inside a loop.
+                        // 
+                        // If we don't have the loop information here, we will conservatively return false.
+                        //
+                        if (!dom)
+                            return false;
+
+                        // If we have dominator tree available, use it to check if the call is inside a loop.
+                        auto callBlock = as<IRBlock>(call->getParent());
+                        if (!callBlock) return false;
+                        auto varBlock = as<IRBlock>(arg->getParent());
+                        if (!varBlock) return false;
+                        auto idom = callBlock;
+                        while (idom != varBlock)
+                        {
+                            idom = dom->getImmediateDominator(idom);
+                            if (!idom)
+                                return false; // If we are here, var does not dominate the call, which should never happen.
+                            if (auto loop = as<IRLoop>(idom->getTerminator()))
+                            {
+                                if (!dom->dominates(loop->getBreakBlock(), callBlock))
+                                    return false; // The var is used in a loop, must return false.
+                            }
+                        }
+                        // If we reach here, the var is used as an inout parameter for the call, but the call
+                        // is not nested in a loop at an higher nesting level than where the var is defined,
+                        // so we can treat the use as DCE-able.
                         continue;
+                    }
                     // We have some other unknown use of the variable address, they can
                     // be loads, or calls using addresses derived from the variable,
                     // we will treat the call as having side effect to be safe.
@@ -718,6 +763,11 @@ bool isPureFunctionalCall(IRCall* call)
 
 bool isSideEffectFreeFunctionalCall(IRCall* call)
 {
+    // If the call has been marked as no-side-effect, we
+    // will treat it so, by-passing all other checks.
+    if (call->findDecoration<IRNoSideEffectDecoration>())
+        return false;
+
     if (!doesCalleeHaveSideEffect(call->getCallee()))
     {
         return areCallArgumentsSideEffectFree(call);

source/slang/slang-ir.cpp

@@ -5,6 +5,8 @@
 
 #include "../core/slang-basic.h"
 
+#include "slang-ir-dominators.h"
+
 #include "slang-mangle.h"
 
 namespace Slang
@@ -4067,6 +4069,20 @@ namespace Slang
         return module;
     }
 
+    IRDominatorTree* IRModule::findOrCreateDominatorTree(IRGlobalValueWithCode* func)
+    {
+        IRAnalysis* analysis = m_mapInstToAnalysis.tryGetValue(func);
+        if (analysis)
+            return analysis->getDominatorTree();
+        else
+        {
+            m_mapInstToAnalysis[func] = IRAnalysis();
+            analysis = m_mapInstToAnalysis.tryGetValue(func);
+        }
+        analysis->domTree = computeDominatorTree(func);
+        return analysis->getDominatorTree();
+    }
+
     void addGlobalValue(
         IRBuilder*  builder,
         IRInst*     value)
@@ -7082,6 +7098,8 @@ namespace Slang
                 module->getDeduplicationContext()->getConstantMap().remove(IRConstantKey{ constInst });
             }
             module->getDeduplicationContext()->getInstReplacementMap().remove(this);
+            if (auto func = as<IRGlobalValueWithCode>(this))
+                module->invalidateAnalysisForInst(func);
         }
         removeArguments();
         removeFromParent();
@@ -7153,10 +7171,6 @@ namespace Slang
                 // common subexpression elimination, etc.
                 //
                 auto call = cast<IRCall>(this);
-                // If the call has been marked as no-side-effect, we
-                // will treat it so, by-passing all other checks.
-                if (call->findDecoration<IRNoSideEffectDecoration>())
-                    return false;
                 return !isSideEffectFreeFunctionalCall(call);
             }
             break;
@@ -7610,6 +7624,11 @@ namespace Slang
         return inst;
     }
 
+    IRDominatorTree* IRAnalysis::getDominatorTree()
+    {
+        return static_cast<IRDominatorTree*>(domTree.get());
+    }
+
 } // namespace Slang
 
 #if SLANG_VC

source/slang/slang-ir.h

@@ -2053,6 +2053,14 @@ struct IRDeduplicationContext
     ConstantMap m_constantMap;
 };
 
+struct IRDominatorTree;
+
+struct IRAnalysis
+{
+    RefPtr<RefObject> domTree;
+    IRDominatorTree* getDominatorTree();
+};
+
 struct IRModule : RefObject
 {
 public:
@@ -2072,6 +2080,17 @@ struct IRModule : RefObject
 
     IRDeduplicationContext* getDeduplicationContext() const { return &m_deduplicationContext; }
 
+    IRDominatorTree* findDominatorTree(IRGlobalValueWithCode* func)
+    {
+        IRAnalysis* analysis = m_mapInstToAnalysis.tryGetValue(func);
+        if (analysis)
+            return analysis->getDominatorTree();
+        return nullptr;
+    }
+    IRDominatorTree* findOrCreateDominatorTree(IRGlobalValueWithCode* func);
+    void invalidateAnalysisForInst(IRGlobalValueWithCode* func) { m_mapInstToAnalysis.remove(func); }
+    void invalidateAllAnalysis() { m_mapInstToAnalysis.clear(); }
+
     IRInstListBase getGlobalInsts() const { return getModuleInst()->getChildren(); }
 
         /// Create an empty instruction with the `op` opcode and space for
@@ -2139,6 +2158,8 @@ struct IRModule : RefObject
 
         /// Holds the obfuscated source map for this module if applicable
     ComPtr<IBoxValue<SourceMap>> m_obfuscatedSourceMap;
+
+    Dictionary<IRInst*, IRAnalysis> m_mapInstToAnalysis;
 };
 
 struct IRSpecializationDictionaryItem : public IRInst

tests/bugs/mutating/mutating-call-in-loop-dce.slang

@@ -0,0 +1,39 @@
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+
+// Confirm that a SideEffectFree mutable method does not get DCE'd when
+// it is called from within a loop.
+
+struct C
+{
+    int a;
+    [mutating]
+    int add()
+    {
+        a++;
+        return a;
+    }
+    [mutating]
+    void init()
+    {
+        a = 0;
+    }
+}
+int doSomething(C d)
+{
+    int rs = 0;
+    for (int i = 0; i < 10; i++)
+        rs = d.add();
+    return rs;
+}
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out, name outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(1, 1, 1)]
+void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
+{
+    int tid = dispatchThreadID.x;
+    C c;
+    c.init();
+    outputBuffer[tid] = doSomething(c);
+}

tests/bugs/mutating/mutating-call-in-loop-dce.slang.expected.txt

@@ -0,0 +1,2 @@
+type: int32_t
+10