[llvm] [DirectX] Propagate shader flags mask of callees to callers (PR #118306)

[Justin Bogner via llvm-commits]

================
@@ -46,21 +49,63 @@ static void updateFunctionFlags(ComputedShaderFlags &CSF,
   }
 }
 
-void ModuleShaderFlags::initialize(const Module &M) {
-  // Collect shader flags for each of the functions
-  for (const auto &F : M.getFunctionList()) {
-    if (F.isDeclaration())
-      continue;
-    ComputedShaderFlags CSF;
-    for (const auto &BB : F)
-      for (const auto &I : BB)
-        updateFunctionFlags(CSF, I);
-    // Insert shader flag mask for function F
-    FunctionFlags.push_back({&F, CSF});
-    // Update combined shader flags mask
-    CombinedSFMask.merge(CSF);
+void ModuleShaderFlags::initialize(Module &M) {
+  CallGraph CG(M);
+
+  // Compute Shader Flags Mask for all functions using post-order visit of SCC
+  // of the call graph.
+  for (scc_iterator<CallGraph *> SCCI = scc_begin(&CG); !SCCI.isAtEnd();
+       ++SCCI) {
+    const std::vector<CallGraphNode *> &CurSCC = *SCCI;
+
+    // Union of shader masks of all functions in CurSCC
+    ComputedShaderFlags SCCSF;
+    for (CallGraphNode *CGN : CurSCC) {
+      Function *F = CGN->getFunction();
+      if (!F)
+        continue;
+
+      if (F->isDeclaration())
+        continue;
+
+      ComputedShaderFlags CSF;
+      for (const auto &BB : *F)
+        for (const auto &I : BB)
+          updateFunctionFlags(CSF, I);
+      // Update combined shader flags mask for all functions in this SCC
+      SCCSF.merge(CSF);
+    }
+
+    // Update combined shader flags mask for all functions of the module
+    CombinedSFMask.merge(SCCSF);
+
+    // Shader flags mask of each of the functions in an SCC of the call graph is
+    // the union of all functions in the SCC. Update shader flags masks of
+    // functions in SCC accordingly. This is trivially true if SCC contains one
+    // function.
+    for (CallGraphNode *CGN : CurSCC) {
+      Function *F = CGN->getFunction();
+      if (!F)
+        continue;
+      // If F already has a shader flag mask associated as result of
+      // any of its callee's flags being propagated, merge SCCSF with
+      // existing flags. Else set its mask to SCCSF.
+      if (FunctionFlags.contains(F))
+        FunctionFlags[F].merge(SCCSF);
+      else
+        FunctionFlags[F] = SCCSF;
+      // Propagate Shader Flag Masks to callers of F
+      for (const auto User : F->users()) {
+        if (const CallInst *CI = dyn_cast<CallInst>(User)) {
+          const Function *Caller = CI->getParent()->getParent();
+          if (FunctionFlags.contains(Caller))
+            FunctionFlags[Caller].merge(SCCSF);
+          else
+            FunctionFlags[Caller] = SCCSF;
+        }
+      }
----------------
bogner wrote:

We shouldn't need a loop to propagate flags to callers. Instead, we should handle this in `updateFunctionFlags` with something like:
```c++
  if (auto *CI = dyn_cast<CallInst>(&I))
    CSF.merge(FunctionFlags[CI->getCalledFunction()]);
```

https://github.com/llvm/llvm-project/pull/118306