[SYCL] Rewrite aspect propagation function lookup using toposort

Signed-off-by: Dmitry Sidorov <dmitrii.s.sidorov@gmail.com>

Author: MrSidims

llvm/lib/SYCLLowerIR/SYCLPropagateAspectsUsage.cpp

@@ -511,7 +511,50 @@ void validateUsedAspectsForFunctions(const FunctionToAspectsMapTy &Map,
   }
 }
 
+/// Computes topological order of functions in the call graph.
+/// Returns functions in reverse topological order.
+/// This allows single-pass bottom-up propagation.
+std::vector<Function *> getTopologicalOrder(const CallGraphTy &CG,
+                                            const std::vector<Function *> &EntryPoints) {
+  std::vector<Function *> Result;
+  DenseMap<const Function *, unsigned> InDegree;
+
+  // Build reverse call graph and compute in-degrees.
+  DenseMap<Function *, SmallVector<Function *, 4>> ReverseCG;
+  for (const auto &[Caller, Callees] : CG) {
+    for (Function *Callee : Callees) {
+      ReverseCG[Callee].push_back(Caller);
+      InDegree[Caller]++;
+    }
+  }
+
+  // Start with functions that have no callees.
+  std::queue<Function *> Worklist;
+  for (const auto &[F, Callees] : CG) {
+    if (InDegree[F] == 0)
+      Worklist.push(F);
+  }
+
+  // Kahn's algorithm for topological sort.
+  while (!Worklist.empty()) {
+    Function *F = Worklist.front();
+    Worklist.pop();
+    Result.push_back(F);
+
+    auto It = ReverseCG.find(F);
+    if (It != ReverseCG.end()) {
+      for (Function *Caller : It->second) {
+        if (--InDegree[Caller] == 0)
+          Worklist.push(Caller);
+      }
+    }
+  }
+
+  return Result;
+}
+
 /// Propagates aspects from leaves up to the top of call graph.
+/// Uses topological sort for efficient single-pass propagation.
 /// NB! Call graph corresponds to call graph of SYCL code which
 /// can't contain recursive calls. So there can't be loops in
 /// a call graph. But there can be path's intersections.
@@ -534,6 +577,91 @@ void propagateAspectsThroughCG(Function *F, CallGraphTy &CG,
   AspectsMap[F].insert(LocalAspects.begin(), LocalAspects.end());
 }
 
+/// Processes each function exactly once in bottom-up order.
+void propagateAspectsThroughCGOptimized(const std::vector<Function *> &TopoOrder,
+                                        const CallGraphTy &CG,
+                                        FunctionToAspectsMapTy &AspectsMap) {
+  // Process in topological order.
+  for (Function *F : TopoOrder) {
+    auto It = CG.find(F);
+    if (It == CG.end())
+      continue;
+
+    // Merge aspects from all callees.
+    for (Function *Callee : It->second) {
+      const auto &CalleeAspects = AspectsMap[Callee];
+      AspectsMap[F].insert(CalleeAspects.begin(), CalleeAspects.end());
+    }
+  }
+}
+
+/// This reduces redundant type lookups when multiple instructions use the same type.
+AspectsSetTy getAspectsFromInstructions(Function &F,
+                                        TypeToAspectsMapTy &TypesWithAspects,
+                                        int FP64Aspect, bool FP64ConvEmu) {
+  AspectsSetTy Result;
+
+  // Collect unique types used across all instructions.
+  SmallDenseMap<const Type *, bool, 32> SeenTypes;
+
+  for (Instruction &I : instructions(F)) {
+    // Check instruction return type.
+    const Type *ReturnType = I.getType();
+    if (auto *AI = dyn_cast<AllocaInst>(&I))
+      ReturnType = AI->getAllocatedType();
+
+    bool IsFP64Conversion = FP64ConvEmu && isFP64ConversionInstruction(I);
+
+    // Only analyze this type once.
+    if (SeenTypes.try_emplace(ReturnType, IsFP64Conversion && hasDoubleType(ReturnType)).second) {
+      const AspectsSetTy &Aspects = getAspectsFromType(ReturnType, TypesWithAspects);
+      for (int Aspect : Aspects) {
+        if (!FP64ConvEmu || Aspect != FP64Aspect || !SeenTypes[ReturnType])
+          Result.insert(Aspect);
+      }
+    }
+
+    // Check operand types.
+    for (const auto &OperandIt : I.operands()) {
+      const Type *OpType = nullptr;
+      if (const auto *GV = dyn_cast<GlobalValue>(OperandIt->stripPointerCasts()))
+        OpType = GV->getValueType();
+      else
+        OpType = OperandIt->getType();
+
+      if (OpType && SeenTypes.try_emplace(OpType, IsFP64Conversion && hasDoubleType(OpType)).second) {
+        const AspectsSetTy &Aspects = getAspectsFromType(OpType, TypesWithAspects);
+        for (int Aspect : Aspects) {
+          if (!FP64ConvEmu || Aspect != FP64Aspect || !SeenTypes[OpType])
+            Result.insert(Aspect);
+        }
+      }
+    }
+
+    // Check GEP source type.
+    if (auto *GEPI = dyn_cast<GetElementPtrInst>(&I)) {
+      const Type *SourceType = GEPI->getSourceElementType();
+      if (SeenTypes.try_emplace(SourceType, IsFP64Conversion && hasDoubleType(SourceType)).second) {
+        const AspectsSetTy &Aspects = getAspectsFromType(SourceType, TypesWithAspects);
+        for (int Aspect : Aspects) {
+          if (!FP64ConvEmu || Aspect != FP64Aspect || !SeenTypes[SourceType])
+            Result.insert(Aspect);
+        }
+      }
+    }
+
+    // Check instruction-level metadata.
+    if (const MDNode *InstAspects = I.getMetadata("sycl_used_aspects")) {
+      for (const MDOperand &MDOp : InstAspects->operands()) {
+        const Constant *C = cast<ConstantAsMetadata>(MDOp)->getValue();
+        Result.insert(cast<ConstantInt>(C)->getSExtValue());
+      }
+    }
+  }
+
+  return Result;
+}
+
 /// Processes a function:
 ///  - checks if return and argument types are using any aspects
 ///  - checks if instructions are using any aspects
@@ -549,12 +677,14 @@ void processFunction(Function &F, FunctionToAspectsMapTy &FunctionToUsedAspects,
   assert(FP64AspectIt != AspectValues.end() &&
          "fp64 aspect was not found in the aspect values.");
   auto FP64Aspect = FP64AspectIt->second;
+
   const AspectsSetTy RetTyAspects =
       getAspectsFromType(F.getReturnType(), TypesWithAspects);
   for (const auto &Aspect : RetTyAspects)
     if (!FP64ConvEmu || (Aspect != FP64Aspect) ||
         !hasDoubleType(F.getReturnType()))
       FunctionToUsedAspects[&F].insert(Aspect);
+
   for (Argument &Arg : F.args()) {
     const AspectsSetTy ArgAspects =
         getAspectsFromType(Arg.getType(), TypesWithAspects);
@@ -564,19 +694,20 @@ void processFunction(Function &F, FunctionToAspectsMapTy &FunctionToUsedAspects,
         FunctionToUsedAspects[&F].insert(Aspect);
   }
 
+  // Optimized instruction analysis with type deduplication.
+  const AspectsSetTy InstrAspects =
+      getAspectsFromInstructions(F, TypesWithAspects, FP64Aspect, FP64ConvEmu);
+  FunctionToUsedAspects[&F].insert(InstrAspects.begin(), InstrAspects.end());
+
+  // Build call graph.
   for (Instruction &I : instructions(F)) {
-    const AspectsSetTy Aspects =
-        getAspectsUsedByInstruction(I, TypesWithAspects);
-    for (const auto &Aspect : Aspects)
-      if (!FP64ConvEmu || (Aspect != FP64Aspect) || !hasDoubleType(I) ||
-          !isFP64ConversionInstruction(I))
-        FunctionToUsedAspects[&F].insert(Aspect);
     if (const auto *CI = dyn_cast<CallInst>(&I)) {
       if (!CI->isIndirectCall() && CI->getCalledFunction())
         CG[&F].insert(CI->getCalledFunction());
     }
   }
 
+  // Collect aspects from metadata (combined to reduce lookups).
   auto CollectAspectsFromMD = [&F](const char* MDName, FunctionToAspectsMapTy &Map) {
     if (const MDNode *MD = F.getMetadata(MDName)) {
       AspectsSetTy Aspects;
@@ -696,23 +827,25 @@ buildFunctionsToAspectsMap(Module &M, TypeToAspectsMapTy &TypesWithAspects,
     collectVirtualFunctionSetInfo(F, VirtualFunctionSets);
   }
 
+  // Compute topological order once for both propagation passes.
+  std::vector<Function *> TopoOrder = getTopologicalOrder(CG, EntryPoints);
+
+  // Handle virtual function sets (still needs old recursive propagation)
   SmallPtrSet<const Function *, 16> Visited;
   for (Function *F : EntryPoints) {
-    propagateAspectsThroughCG(F, CG, FunctionToUsedAspects, Visited);
     processDeclaredVirtualFunctionSets(F, CG, FunctionToUsedAspects, Visited,
                                        VirtualFunctionSets);
   }
 
+  // Optimized single-pass propagation for used aspects.
+  propagateAspectsThroughCGOptimized(TopoOrder, CG, FunctionToUsedAspects);
+
   if (ValidateAspects)
     validateUsedAspectsForFunctions(FunctionToUsedAspects, AspectValues,
                                     EntryPoints, CG);
 
-  // The set of aspects from FunctionToDeclaredAspects should be merged to the
-  // set of FunctionToUsedAspects after validateUsedAspectsForFunctions call to
-  // avoid errors during validation.
-  Visited.clear();
-  for (Function *F : EntryPoints)
-    propagateAspectsThroughCG(F, CG, FunctionToDeclaredAspects, Visited);
+  // Optimized single-pass propagation for declared aspects.
+  propagateAspectsThroughCGOptimized(TopoOrder, CG, FunctionToDeclaredAspects);
 
   return {std::move(FunctionToUsedAspects),
           std::move(FunctionToDeclaredAspects)};