[Attributor]Introduce call base context argument pathway.

[kuter]

Summary

This patch puts a constraint on the attributes that will receive call base context.

[jdoerfert]

llvm/include/llvm/Transforms/IPO/Attributor.h

1434

Nit: For better or worse we use F and Fn as a variable name for functions usually.

I read the comment and do not know what happens. We need to improve the wording, maybe elaborate.

llvm/lib/Transforms/IPO/Attributor.cpp

97

I'm confused by the description. Is the analysis invaliding values?

818

Nit: Assertions always with messages.

---

Can we have a explanation here or in the header what is happening here. The commit message should also contain more details.

826

Style: No braces. Use a range loop or part of the constructor (or use append).

[kuter]

fix wording. use pointer_iterator instead of for loop.
fix naming.

[kuter]

llvm/lib/Transforms/IPO/Attributor.cpp

826

That was my initial idea but since arg_begin() is a pointer the SmallVector was trying to converyt Argument to Value*. After investigating I found that there is something called pointer_iterator that would fix this issue.

[kuter]

add unit test.
fix bug (😄)
fix typo.

[kuter]

remove unused header.

[kuter]

llvm/unittests/Transforms/IPO/AttributorTestBase.h

53

contents BumpPtrAllocator will be freed after this block is dead.
This causes a use after free.

[jdoerfert]

llvm/unittests/Transforms/IPO/AttributorTestBase.h

53

Put the allocator in the same scope as the Attributor?

[kuter]

llvm/unittests/Transforms/IPO/AttributorTestBase.h

53

That's what I am planing to do. I just wanted to add a inline commend so I don't forget about this issue.

llvm/include/llvm/Transforms/IPO/Attributor.h

private:
   /// all call sites are known, hence the function has internal linkage.
   /// If true is returned, \p AllCallSitesKnown is set if all possible call
   /// sites of the function have been visited.
   bool checkForAllCallSites(function_ref<bool(AbstractCallSite)> Pred,
                             const Function &Fn, bool RequireAllCallSites,
                             const AbstractAttribute *QueryingAA,
                             bool &AllCallSitesKnown);
 
+  /// for function \p Fn Populate the cache for values that the call base
+  /// context is going to get propagated to.
+  void populateArgumentPathwayForFunction(Function *Fn);
+
   /// Determine if CallBase context in \p IRP should be propagated.
-  bool shouldPropagateCallBaseContext(const IRPosition &IRP);
+  /// \p IgnoreCommandLine is for ignoring the
+  /// `attributor-enable-call-site-specific-deduction` command line argument,
+  /// only used for testing.
+  bool shouldPropagateCallBaseContext(const IRPosition &IRP,
+                                      bool IgnoreCommandLine = false);
 
   /// Apply all requested function signature rewrites
   /// (\see registerFunctionSignatureRewrite) and return Changed if the module
   /// was altered.
   ChangeStatus
   rewriteFunctionSignatures(SmallPtrSetImpl<Function *> &ModifiedFns);
 
   /// Check if the Attribute \p AA should be seeded.
private:
 
   /// Functions, blocks, and instructions we delete after manifest is done.
   ///
   ///{
   SmallPtrSet<Function *, 8> ToBeDeletedFunctions;
   SmallPtrSet<BasicBlock *, 8> ToBeDeletedBlocks;
   SmallDenseSet<WeakVH, 8> ToBeDeletedInsts;
   ///}
+
+  /// Values per function that are relevant to the call base context.
+  DenseMap<Function *, SmallPtrSet<Value *, 32>> ArgumentPathway;
+
+  friend class AttributorTestBase;
 };
 
 /// An interface to query the internal state of an abstract attribute.
 ///
 /// The abstract state is a minimal interface that allows the Attributor to
 /// communicate with the abstract attributes about their internal state without
 /// enforcing or exposing implementation details, e.g., the (existence of an)
 /// underlying lattice.

llvm/lib/Transforms/IPO/Attributor.cpp

SeedAllowList("attributor-seed-allow-list", cl::Hidden,
                            "allowed to be seeded."),
                   cl::ZeroOrMore, cl::CommaSeparated);
 
 static cl::opt<bool> EnableCallSiteSpecific(
     "attributor-enable-call-site-specific-deduction", cl::Hidden,
     cl::desc("Allow the Attributor to do call site specific analysis"),
     cl::init(false));
 
+static cl::opt<unsigned>
+    CallSiteSpecificDepth("attributor-call-site-specifc-deduction-depth",
+                          cl::Hidden,
+                          cl::desc("Maximum depth of values that call site "
+                                   "specific analysis context will reach."),
+                          cl::init(6));
+
 /// Logic operators for the change status enum class.
 ///
 ///{
 ChangeStatus llvm::operator|(ChangeStatus l, ChangeStatus r) {
   return l == ChangeStatus::CHANGED ? l : r;
 }
 ChangeStatus llvm::operator&(ChangeStatus l, ChangeStatus r) {
   return l == ChangeStatus::UNCHANGED ? l : r;
for (unsigned u = 0; u < Uses.size(); ++u) {
     LLVM_DEBUG(dbgs() << "[Attributor] Call site callback failed for "
                       << *ACS.getInstruction() << "\n");
     return false;
   }
 
   return true;
 }
 
-bool Attributor::shouldPropagateCallBaseContext(const IRPosition &IRP) {
-  // TODO: Maintain a cache of Values that are
-  // on the pathway from a Argument to a Instruction that would effect the
-  // liveness/return state etc.
-  return EnableCallSiteSpecific;
+void Attributor::populateArgumentPathwayForFunction(Function *Fn) {
+  /// Efficiently compute the intersections between Values that can be
+  /// reached forwards from the arguments and the Values that can be
+  /// reached backwards from the instructions that would benefit.
+  /// from call site specific information.
+  assert(Fn != nullptr && "There must be a function!");
+
+  SmallVector<Value *, 32> SuccessfullValues;
+  SmallVector<Value *, 32> Worklist(
+      pointer_iterator<Argument *>(Fn->arg_begin()),
+      pointer_iterator<Argument *>(Fn->arg_end()));
+  SmallPtrSet<Value *, 32> ForwardsVisited;
+
+  for (unsigned I = 0; I < CallSiteSpecificDepth && !Worklist.empty(); I++) {
+    SmallVector<Value *, 32> NewWorklist;
+    while (!Worklist.empty()) {
+      Value *Val = Worklist.pop_back_val();
+
+      // Avoid infinite loops.
+      if (!ForwardsVisited.insert(Val).second)
+        continue;
+
+      if (auto *Instr = dyn_cast<Instruction>(Val)) {
+        // Keep track of important instructions.
+        if (Instr->isTerminator())
+          SuccessfullValues.push_back(Instr);
+      }
+      for (User *U : Val->users())
+        NewWorklist.push_back(U);
+    }
+    Worklist.swap(NewWorklist);
+  }
+
+  Worklist.assign(SuccessfullValues.begin(), SuccessfullValues.end());
+
+  SmallSet<Value *, 32> BackwardsVisited;
+  for (unsigned I = 0; I < CallSiteSpecificDepth && !Worklist.empty(); I++) {
+    SmallVector<Value *, 32> NewWorklist;
+    while (!Worklist.empty()) {
+      Value *Val = Worklist.pop_back_val();
+
+      // Avoid infinite loops.
+      if (!BackwardsVisited.insert(Val).second)
+        continue;
+
+      if (auto Inst = dyn_cast<Instruction>(Val)) {
+        for (Use &U : Inst->operands()) {
+
+          // Avoid visiting nodes that can not be reached by the arguments.
+          // this creates the intersection.
+          if (!ForwardsVisited.count(U.get()))
+            continue;
+          NewWorklist.push_back(U.get());
+        }
+      }
+    }
+    Worklist.swap(NewWorklist);
+  }
+
+  LLVM_DEBUG(dbgs() << "[Attributor] Argument pathway for function: "
+                    << Fn->getName() << "\n");
+  for (Value *V : BackwardsVisited) {
+    LLVM_DEBUG(dbgs() << "[Attributor]" << *V << " is in Argument Pathway."
+                      << "\n");
+  }
+  ArgumentPathway[Fn] = std::move(BackwardsVisited);
+}
+
+bool Attributor::shouldPropagateCallBaseContext(const IRPosition &IRP,
+                                                bool IgnoreCommandLine) {
+  if (!IgnoreCommandLine && !EnableCallSiteSpecific)
+    return false;
+
+  // Certain positions must always propogate call base context.
+  switch (IRP.getPositionKind()) {
+  case IRPosition::IRP_RETURNED:
+  case IRPosition::IRP_ARGUMENT:
+  case IRPosition::IRP_FUNCTION:
+    return true;
+  case IRPosition::IRP_FLOAT:
+    break;
+  default:
+    return false;
+  }
+  Function *Fn = IRP.getAssociatedFunction();
+  if (!Fn)
+    return false;
+
+  // Look up the value to see if it is one of the values that will recieve
+  // call base context.
+  auto ArgPathwayIt = ArgumentPathway.find(Fn);
+  if (ArgPathwayIt == ArgumentPathway.end()) {
+    // No cached information, create the `patway`.
+    populateArgumentPathwayForFunction(Fn);
+    ArgPathwayIt = ArgumentPathway.find(Fn);
+  }
+
+  return ((*ArgPathwayIt).second).count(&IRP.getAnchorValue());
 }
 
 bool Attributor::checkForAllReturnedValuesAndReturnInsts(
     function_ref<bool(Value &, const SmallSetVector<ReturnInst *, 4> &)> Pred,
     const AbstractAttribute &QueryingAA) {
 
   const IRPosition &IRP = QueryingAA.getIRPosition();
   // Since we need to provide return instructions we have to have an exact

llvm/unittests/Transforms/IPO/AttributorTest.cpp

 //===- AttributorTest.cpp - Attributor unit tests ------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===---------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "AttributorTestBase.h"
 #include "llvm/ADT/StringRef.h"
-#include "llvm/Analysis/CGSCCPassManager.h"
-#include "llvm/Analysis/CallGraphSCCPass.h"
-#include "llvm/Analysis/LoopAnalysisManager.h"
-#include "llvm/AsmParser/Parser.h"
-#include "llvm/Support/Allocator.h"
-#include "llvm/Testing/Support/Error.h"
-#include "llvm/Transforms/Utils/CallGraphUpdater.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "gtest/gtest.h"
 #include <memory>
 
 namespace llvm {
 
 TEST_F(AttributorTestBase, IRPPositionCallBaseContext) {
   const char *ModuleString = R"(
     define i32 @foo(i32 %a) {
     entry:
       ret i32 %a
     }
   )";
 
   parseModule(ModuleString);
 
   Function *F = M->getFunction("foo");
   IRPosition Pos = IRPosition::function(*F, (const llvm::CallBase *)0xDEADBEEF);
   EXPECT_TRUE(Pos.hasCallBaseContext());
   EXPECT_FALSE(Pos.stripCallBaseContext().hasCallBaseContext());
 }
 
+TEST_F(AttributorTestBase, CallBaseContextPathway) {
+  const char *ModuleString = R"(
+  define i32 @foo(i32 %a, i32 %b) #0 {
+    ; Distance to return is 4 (should have cb context)
+    %a1 = add nsw i32 %a, 1
+    %a2 = add nsw i32 %a1, 1
+    %a3 = add nsw i32 %a2, 1
+    
+    ; Distance to return is 7 (should not have cb context)
+    %b1 = add nsw i32 %b, 1
+    %b2 = add nsw i32 %b1, 1
+    %b3 = add nsw i32 %b2, 1
+    %b4 = add nsw i32 %b3, 1
+    %b5 = add nsw i32 %b4, 1
+    %b6 = add nsw i32 %b5, 1
+    %b7 = add nsw i32 %b6, 1
+
+    %c = add nsw i32 %b7, %a3
+    ret i32 %c
+  })";
+
+  auto Instr = [](Function *F, StringRef Name) -> Instruction * {
+    for (inst_iterator I = inst_begin(F); I != inst_end(F); I++) {
+      if (I->getName() == Name)
+        return &*I;
+    }
+    llvm_unreachable("Instruction name not found!");
+    return nullptr;
+  };
+
+  Module &M = parseModule(ModuleString);
+  Attributor &A = createAttributor(M);
+  Function *F = M.getFunction("foo");
+
+#define EXPECT_PROPAGATE(n)                                                    \
+  EXPECT_TRUE(doesPropagateCBContext(A, *Instr(F, (n))))
+#define EXPECT_NOT_PROPAGATE(n)                                                \
+  EXPECT_FALSE(doesPropagateCBContext(A, *Instr(F, (n))))
+
+  EXPECT_PROPAGATE("a1");
+  EXPECT_PROPAGATE("a2");
+  EXPECT_PROPAGATE("a3");
+
+  EXPECT_NOT_PROPAGATE("b1");
+  EXPECT_NOT_PROPAGATE("b2");
+  EXPECT_NOT_PROPAGATE("b3");
+  EXPECT_NOT_PROPAGATE("b4");
+  EXPECT_NOT_PROPAGATE("b5");
+  EXPECT_NOT_PROPAGATE("b6");
+  EXPECT_NOT_PROPAGATE("b7");
+
+  EXPECT_PROPAGATE("c");
+
+#undef EXPECT_PROPAGATE
+#undef EXPECT_NOT_PROPAGATE
+}
+
 TEST_F(AttributorTestBase, TestCast) {
   const char *ModuleString = R"(
     define i32 @foo(i32 %a, i32 %b) {
     entry:
       %c = add i32 %a, %b
       ret i32 %c
     }
   )";
 
   Module &M = parseModule(ModuleString);
-
-  SetVector<Function *> Functions;
-  AnalysisGetter AG;
-  for (Function &F : M)
-    Functions.insert(&F);
-
-  CallGraphUpdater CGUpdater;
-  BumpPtrAllocator Allocator;
-  InformationCache InfoCache(M, AG, Allocator, nullptr);
-  Attributor A(Functions, InfoCache, CGUpdater);
+  Attributor &A = createAttributor(M);
 
   Function *F = M.getFunction("foo");
 
   AbstractAttribute *AA = (AbstractAttribute *)&(
       A.getOrCreateAAFor<AAIsDead>(IRPosition::function(*F)));
 
   EXPECT_TRUE(AA);
 
   const auto *SFail = dyn_cast<AAAlign>(AA);
   const auto *SSucc = dyn_cast<AAIsDead>(AA);
 
   ASSERT_EQ(SFail, nullptr);
   ASSERT_TRUE(SSucc);
 }
 
-} // namespace llvm
+} // namespace llvm
 

llvm/unittests/Transforms/IPO/AttributorTestBase.h

 //===- llvm/unittest/Transforms/IPO/AttributorTestBase.h -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This file defines an AttributorTestBase class, which provides helpers to
 /// parse a LLVM IR string and create Attributor
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_UNITTESTS_TRANSFORMS_ATTRIBUTOR_TESTBASE_H
 #define LLVM_UNITTESTS_TRANSFORMS_ATTRIBUTOR_TESTBASE_H
 
+#include "llvm/ADT/iterator.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Testing/Support/Error.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 #include "gtest/gtest.h"
 #include <memory>
 
 namespace llvm {
 
 /// Helper class to create a module from assembly string and an Attributor
 class AttributorTestBase : public testing::Test {
 protected:
   std::unique_ptr<LLVMContext> Ctx;
   std::unique_ptr<Module> M;
+  std::unique_ptr<Attributor> A;
 
   AttributorTestBase() : Ctx(new LLVMContext) {}
 
   Module &parseModule(const char *ModuleString) {
     SMDiagnostic Err;
     M = parseAssemblyString(ModuleString, Err, *Ctx);
     EXPECT_TRUE(M);
     return *M;
   }
+
+  Attributor &createAttributor(Module &M) {
+    SetVector<Function *> Functions(
+        pointer_iterator<Module::iterator>(M.begin()),
+        pointer_iterator<Module::iterator>(M.end()));
+
+    AnalysisGetter AG;
+    CallGraphUpdater CGUpdater;
+    BumpPtrAllocator Allocator;
+    InformationCache InfoCache(M, AG, Allocator, nullptr);
+    A = std::make_unique<Attributor>(Functions, InfoCache, CGUpdater);
+    return *A;
+  }
+
+  bool doesPropagateCBContext(Attributor &A, Instruction &Inst) {
+    return A.shouldPropagateCallBaseContext(IRPosition::value(Inst),
+                                            true /* IgnoreCommandLine */);
+  }
 };
 
 } // namespace llvm
 
-#endif
+#endif