Extend DCE to handle pure function calls that returns via `inout` parameters. (#2700)

2 files changed, 45 insertions, 19 deletions

diff --git a/source/slang/slang-ir-autodiff-unzip.cpp b/source/slang/slang-ir-autodiff-unzip.cpp
index 16862bb19..7e01fde28 100644
--- a/source/slang/slang-ir-autodiff-unzip.cpp
+++ b/source/slang/slang-ir-autodiff-unzip.cpp
@@ -581,20 +581,6 @@ IRFunc* DiffUnzipPass::extractPrimalFunc(
                 }
                 instsToRemove.add(inst);
             }
-            else if (auto primalCtx = inst->findDecoration<IRBackwardDerivativePrimalContextDecoration>())
-            {
-                if (inst->getOp() == kIROp_Call)
-                {
-                    // The primal calls should be marked as no side effect so they can be DCE'd if possible.
-                    // We can only do so if the intermediate context of the callee is stored.
-                    //
-                    if (primalCtx->getBackwardDerivativePrimalContextVar()
-                            ->findDecoration<IRPrimalValueStructKeyDecoration>())
-                    {
-                        builder.addSimpleDecoration<IRNoSideEffectDecoration>(inst);
-                    }
-                }
-            }
         }
     }
 
diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp
index de03a1661..a27afee8a 100644
--- a/source/slang/slang-ir-util.cpp
+++ b/source/slang/slang-ir-util.cpp
@@ -609,17 +609,57 @@ bool isPureFunctionalCall(IRCall* call)
     {
         // If the function has no side effect and is not writing to any outputs,
         // we can safely treat the call as a normal inst.
-        bool hasOutArg = false;
+        IRFunc* parentFunc = nullptr;
         for (UInt i = 0; i < call->getArgCount(); i++)
         {
-            if (isValueType(call->getArg(i)->getDataType()))
+            auto arg = call->getArg(i);
+            if (isValueType(arg->getDataType()))
                 continue;
+
             // If the argument type is not a known value type,
             // assume it is a pointer or handle through which side effect can take place.
-            hasOutArg = true;
-            break;
+            if (!parentFunc)
+            {
+                parentFunc = getParentFunc(call);
+                if (!parentFunc)
+                    return false;
+            }
+
+            if (arg->getOp() == kIROp_Var && getParentFunc(arg) == parentFunc)
+            {
+                // If the pointer argument is a local variable (thus can't alias with other addresses)
+                // and it is never read from in the function, we can safely treat the call as having
+                // no side-effect.
+                // This is a conservative test, but is sufficient to detect the most common case where
+                // a temporary variable is used as the inout argument and the result stored in the temp
+                // variable isn't being used elsewhere in the parent func.
+                // 
+                // A more aggresive test can check all other address uses reachable from the call site
+                // and see if any of them are aliasing with the argument.
+                for (auto use = arg->firstUse; use; use = use->nextUse)
+                {
+                    if (as<IRDecoration>(use->getUser()))
+                        continue;
+                    switch (use->getUser()->getOp())
+                    {
+                    case kIROp_Store:
+                        // We are fine with stores into the variable, since store operations
+                        // are not dependent on whatever we do in the call here.
+                        continue;
+                    default:
+                        // 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.
+                        return false;
+                    }
+                }
+            }
+            else
+            {
+                return false;
+            }
         }
-        return !hasOutArg;
+        return true;
     }
     return false;
 }