Partial specialization was not checking the callsite to make sure it was using the...

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@105528 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/lib/Transforms/IPO/PartialSpecialization.cpp b/lib/Transforms/IPO/PartialSpecialization.cpp
index 084b94e53566132bcfab5a3255c4abc9654b3fb4..6a288fef64eeb4eed671bf63aa967c2c5269308e 100644 (file)
--- a/lib/Transforms/IPO/PartialSpecialization.cpp
+++ b/lib/Transforms/IPO/PartialSpecialization.cpp
@@ -32,6 +32,10 @@
 using namespace llvm;
 
 STATISTIC(numSpecialized, "Number of specialized functions created");
+STATISTIC(numReplaced, "Number of callers replaced by specialization");
+
+// Maximum number of arguments markable interested
+static const int MaxInterests = 6;
 
 // Call must be used at least occasionally
 static const int CallsMin = 5;
@@ -40,8 +44,9 @@ static const int CallsMin = 5;
 static const double ConstValPercent = .1;
 
 namespace {
+  typedef SmallVector<int, MaxInterests> InterestingArgVector;
   class PartSpec : public ModulePass {
-    void scanForInterest(Function&, SmallVector<int, 6>&);
+    void scanForInterest(Function&, InterestingArgVector&);
     int scanDistribution(Function&, int, std::map<Constant*, int>&);
   public :
     static char ID; // Pass identification, replacement for typeid
@@ -61,11 +66,13 @@ static Function*
 SpecializeFunction(Function* F, 
                    DenseMap<const Value*, Value*>& replacements) {
   // arg numbers of deleted arguments
-  DenseSet<unsigned> deleted;
+  DenseMap<unsigned, const Argument*> deleted;
   for (DenseMap<const Value*, Value*>::iterator 
          repb = replacements.begin(), repe = replacements.end();
-       repb != repe; ++repb)
-    deleted.insert(cast<Argument>(repb->first)->getArgNo());
+       repb != repe; ++repb) {
+    Argument const *arg = cast<const Argument>(repb->first);
+    deleted[arg->getArgNo()] = arg;
+  }
 
   Function* NF = CloneFunction(F, replacements);
   NF->setLinkage(GlobalValue::InternalLinkage);
@@ -80,9 +87,23 @@ SpecializeFunction(Function* F,
       if (CS.getCalledFunction() == F) {
         
         SmallVector<Value*, 6> args;
-        for (unsigned x = 0; x < CS.arg_size(); ++x)
-          if (!deleted.count(x))
-            args.push_back(CS.getArgument(x));
+        // Assemble the non-specialized arguments for the updated callsite.
+        // In the process, make sure that the specialized arguments are
+        // constant and match the specialization.  If that's not the case,
+        // this callsite needs to call the original or some other
+        // specialization; don't change it here.
+        CallSite::arg_iterator as = CS.arg_begin(), ae = CS.arg_end();
+        for (CallSite::arg_iterator ai = as; ai != ae; ++ai) {
+          DenseMap<unsigned, const Argument*>::iterator delit = deleted.find(
+            std::distance(as, ai));
+          if (delit == deleted.end())
+            args.push_back(cast<Value>(ai));
+          else {
+            Constant *ci = dyn_cast<Constant>(ai);
+            if (!(ci && ci == replacements[delit->second]))
+              goto next_use;
+          }
+        }
         Value* NCall;
         if (CallInst *CI = dyn_cast<CallInst>(i)) {
           NCall = CallInst::Create(NF, args.begin(), args.end(), 
@@ -99,8 +120,11 @@ SpecializeFunction(Function* F,
         }
         CS.getInstruction()->replaceAllUsesWith(NCall);
         CS.getInstruction()->eraseFromParent();
+        ++numReplaced;
       }
     }
+  next_use:
+    ;
   }
   return NF;
 }
@@ -111,7 +135,7 @@ bool PartSpec::runOnModule(Module &M) {
   for (Module::iterator I = M.begin(); I != M.end(); ++I) {
     Function &F = *I;
     if (F.isDeclaration() || F.mayBeOverridden()) continue;
-    SmallVector<int, 6> interestingArgs;
+    InterestingArgVector interestingArgs;
     scanForInterest(F, interestingArgs);
 
     // Find the first interesting Argument that we can specialize on
@@ -143,7 +167,7 @@ bool PartSpec::runOnModule(Module &M) {
 
 /// scanForInterest - This function decides which arguments would be worth
 /// specializing on.
-void PartSpec::scanForInterest(Function& F, SmallVector<int, 6>& args) {
+void PartSpec::scanForInterest(Function& F, InterestingArgVector& args) {
   for(Function::arg_iterator ii = F.arg_begin(), ee = F.arg_end();
       ii != ee; ++ii) {
     for(Value::use_iterator ui = ii->use_begin(), ue = ii->use_end();

diff --git a/test/Transforms/PartialSpecialize/dg.exp b/test/Transforms/PartialSpecialize/dg.exp
new file mode 100644 (file)
index 0000000..f200589
--- /dev/null
+++ b/test/Transforms/PartialSpecialize/dg.exp
@@ -0,0 +1,3 @@
+load_lib llvm.exp
+
+RunLLVMTests [lsort [glob -nocomplain $srcdir/$subdir/*.{ll,c,cpp}]]

diff --git a/test/Transforms/PartialSpecialize/two-specializations.ll b/test/Transforms/PartialSpecialize/two-specializations.ll
new file mode 100644 (file)
index 0000000..9a371a5
--- /dev/null
+++ b/test/Transforms/PartialSpecialize/two-specializations.ll
@@ -0,0 +1,32 @@
+; If there are two specializations of a function, make sure each callsite
+; calls the right one.
+;
+; RUN: opt -S -partialspecialization %s | FileCheck %s
+declare void @callback1()
+declare void @callback2()
+
+define internal void @UseCallback(void()* %pCallback) {
+  call void %pCallback()
+  ret void
+}
+
+define void @foo(void()* %pNonConstCallback)
+{
+Entry:
+; CHECK: Entry
+; CHECK-NEXT: call void @UseCallback1()
+; CHECK-NEXT: call void @UseCallback1()
+; CHECK-NEXT: call void @UseCallback2()
+; CHECK-NEXT: call void @UseCallback(void ()* %pNonConstCallback)
+; CHECK-NEXT: call void @UseCallback1()
+; CHECK-NEXT: call void @UseCallback2()
+; CHECK-NEXT: call void @UseCallback2()
+  call void @UseCallback(void()* @callback1)
+  call void @UseCallback(void()* @callback1)
+  call void @UseCallback(void()* @callback2)
+  call void @UseCallback(void()* %pNonConstCallback)
+  call void @UseCallback(void()* @callback1)
+  call void @UseCallback(void()* @callback2)
+  call void @UseCallback(void()* @callback2)
+  ret void
+}