[FuncSpec] Consider constant struct arguments when specializing.

[labrinea]

Summary

Optionally enabled just like integer and floating point arguments.

[ChuanqiXu]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

880–881

This looks not the same with the original statements. For example, we won't do anything if V is the OldArg and ValueState doesn't contain the OldArg.

[labrinea]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

880–881

Looking at what handleCallArguments() does, makes me think that OldArg ought to be in either ValueState or StructValueState. That said I am wondering whether it's better to just call Solver.visit(CallSite) for the callsites of specializations right after we've updated them. This would essentially make markArgInFuncSpecialization() obsolete. Lemme try this and run a few tests.

[ChuanqiXu]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

880–881

Agreed. It looks cleaner at least.

[labrinea]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

880–881

I've raised D146158.

[labrinea]

rebased on parent

[ChuanqiXu]

LGTM now.

llvm/lib/Transforms/Utils/SCCPSolver.cpp

818

[labrinea]

This is basically the same revision as https://reviews.llvm.org/D145374?id=502612, where markArgInFuncSpecialization was adapted to support struct arguments. As @chill pointed out in D147132, the deprecated function has the potential to be more efficient if funcspec-for-literal-constant gets enabled by default, so I am putting it back in. To answer @ChuanqiXu's comment: arguments we are not specializing on, must be present either in ValueStateor StructValueState therefore the specialization can inherit their lattice value from the original function. The existing code does not handle structs, hence the check I suppose? (a struct won't be present in the standard value map)

[chill]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

668–673

Maybe delete this part?

[labrinea]

Changes from last revision:

• Renamed markArgInFuncSpecialization to handleSpecializationArguments as a better description.
• Rewrote the above function such that:
  • the lattice state of non constant arguments is copied from the original function without recomputing
  • the arguments are not pushed to the solver's worklists (there's no need for it)
  • Function::arg_iterator is used rather than Argument * directly
• Inlined isValidArgumentType and removed repeated check for int/float types

[labrinea]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

876

I'll replace this with ValueLatticeElement

[chill]

LGTM too, cheers!

llvm/lib/Transforms/Utils/SCCPSolver.cpp

843

handleXXX is very uninformative, it might be good in a contexts where we don't know
what kind of "handling" is done by the caller (like e.g. "visit"), but in this specific case
of setitng up/initialising lattice values for the parameters it's better called just like that, something like
initSpecializationParemetersLaticeValues, setupSpecArgsValueStates, etc.

[labrinea]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

843

Agree thanks. I was thinking of SetLatticeValueForSpecializationArguments but it was a bit long. We already have similar names: getLatticeValueFor, removeLatticeValueFor and resetLatticeValueFor (see D146158).

llvm/include/llvm/Transforms/Utils/SCCPSolver.h

public:
   // corresponding to elements of the structure are constants,
   // false otherwise.
   bool isStructLatticeConstant(Function *F, StructType *STy);
 
   /// Helper to return a Constant if \p LV is either a constant or a constant
   /// range with a single element.
   Constant *getConstant(const ValueLatticeElement &LV) const;
 
+  /// Return either a Constant or nullptr for a given Value.
+  Constant *getConstantOrNull(Value *V) const;
+
   /// Return a reference to the set of argument tracked functions.
   SmallPtrSetImpl<Function *> &getArgumentTrackedFunctions();
 
-  /// Mark the constant arguments of a new function specialization. \p F points
-  /// to the cloned function and \p Args contains a list of constant arguments
-  /// represented as pairs of {formal,actual} values (the formal argument is
-  /// associated with the original function definition). All other arguments of
-  /// the specialization inherit the lattice state of their corresponding values
-  /// in the original function.
-  void markArgInFuncSpecialization(Function *F,
+  /// Set the Lattice Value for the arguments of a specialization \p F.
+  /// If an argument is Constant then its lattice value is marked with the
+  /// corresponding actual argument in \p Args. Otherwise, its lattice value
+  /// is inherited (copied) from the corresponding formal argument in \p Args.
+  void handleSpecializationArguments(Function *F,
                                      const SmallVectorImpl<ArgInfo> &Args);
 
   /// Mark all of the blocks in function \p F non-executable. Clients can used
   /// this method to erase a function from the module (e.g., if it has been
   /// completely specialized and is no longer needed).
   void markFunctionUnreachable(Function *F);
 
   void visit(Instruction *I);
   void visitCall(CallInst &I);

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 }
 
 Function *FunctionSpecializer::createSpecialization(Function *F, const SpecSig &S) {
   Function *Clone = cloneCandidateFunction(F);
 
   // Initialize the lattice state of the arguments of the function clone,
   // marking the argument on which we specialized the function constant
   // with the given value.
-  Solver.markArgInFuncSpecialization(Clone, S.Args);
+  Solver.handleSpecializationArguments(Clone, S.Args);
 
   Solver.addArgumentTrackedFunction(Clone);
   Solver.markBlockExecutable(&Clone->front());
 
   // Mark all the specialized functions
   Specializations.insert(Clone);
   ++NumSpecsCreated;
 

 }
 
 /// Determine if it is possible to specialise the function for constant values
 /// of the formal parameter \p A.
 bool FunctionSpecializer::isArgumentInteresting(Argument *A) {
   // No point in specialization if the argument is unused.
   if (A->user_empty())
     return false;
 
-  // For now, don't attempt to specialize functions based on the values of
-  // composite types.
-  Type *ArgTy = A->getType();
-  if (!ArgTy->isSingleValueType())
-    return false;
-
-  // Specialization of integer and floating point types needs to be explicitly
-  // enabled.
-  if (!SpecializeLiteralConstant &&
-      (ArgTy->isIntegerTy() || ArgTy->isFloatingPointTy()))
+  Type *Ty = A->getType();
+  if (!Ty->isPointerTy() && (!SpecializeLiteralConstant ||
+      (!Ty->isIntegerTy() && !Ty->isFloatingPointTy() && !Ty->isStructTy())))
     return false;
 
   // SCCP solver does not record an argument that will be constructed on
   // stack.
   if (A->hasByValAttr() && !A->getParent()->onlyReadsMemory())
     return false;
 
   // Check the lattice value and decide if we should attemt to specialize,
   // based on this argument. No point in specialization, if the lattice value
   // is already a constant.
-  const ValueLatticeElement &LV = Solver.getLatticeValueFor(A);
-  if (LV.isUnknownOrUndef() || LV.isConstant() ||
-      (LV.isConstantRange() && LV.getConstantRange().isSingleElement())) {
-    LLVM_DEBUG(dbgs() << "FnSpecialization: Nothing to do, parameter "
-                      << A->getNameOrAsOperand() << " is already constant\n");
-    return false;
-  }
-
-  LLVM_DEBUG(dbgs() << "FnSpecialization: Found interesting parameter "
-                    << A->getNameOrAsOperand() << "\n");
+  bool IsOverdefined = Ty->isStructTy()
+    ? any_of(Solver.getStructLatticeValueFor(A), SCCPSolver::isOverdefined)
+    : SCCPSolver::isOverdefined(Solver.getLatticeValueFor(A));
 
-  return true;
+  LLVM_DEBUG(
+    if (IsOverdefined)
+      dbgs() << "FnSpecialization: Found interesting parameter "
+             << A->getNameOrAsOperand() << "\n";
+    else
+      dbgs() << "FnSpecialization: Nothing to do, parameter "
+             << A->getNameOrAsOperand() << " is already constant\n";
+  );
+  return IsOverdefined;
 }
 
-/// Check if the valuy \p V  (an actual argument) is a constant or can only
+/// Check if the value \p V  (an actual argument) is a constant or can only
 /// have a constant value. Return that constant.
 Constant *FunctionSpecializer::getCandidateConstant(Value *V) {
   if (isa<PoisonValue>(V))
     return nullptr;
 
   // TrackValueOfGlobalVariable only tracks scalar global variables.
   if (auto *GV = dyn_cast<GlobalVariable>(V)) {
     // Check if we want to specialize on the address of non-constant
     // global values.
     if (!GV->isConstant() && !SpecializeOnAddress)
       return nullptr;
 
     if (!GV->getValueType()->isSingleValueType())
       return nullptr;
   }
 
   // Select for possible specialisation values that are constants or
   // are deduced to be constants or constant ranges with a single element.
   Constant *C = dyn_cast<Constant>(V);
-  if (!C) {
-    const ValueLatticeElement &LV = Solver.getLatticeValueFor(V);
-    if (LV.isConstant())
-      C = LV.getConstant();
-    else if (LV.isConstantRange() && LV.getConstantRange().isSingleElement()) {
-      assert(V->getType()->isIntegerTy() && "Non-integral constant range");
-      C = Constant::getIntegerValue(V->getType(),
-                                    *LV.getConstantRange().getSingleElement());
-    } else
-      return nullptr;
-  }
-
+  if (!C)
+    C = Solver.getConstantOrNull(V);
   return C;
 }
 
 void FunctionSpecializer::updateCallSites(Function *F, const Spec *Begin,
                                           const Spec *End) {
   // Collect the call sites that need updating.
   SmallVector<CallBase *> ToUpdate;
   for (User *U : F->users())

llvm/lib/Transforms/Utils/SCCPSolver.cpp

static bool canRemoveInstruction(Instruction *I) {
   // those cases by falling through to here.
   // TODO: Mark globals as being constant earlier, so
   // TODO: wouldInstructionBeTriviallyDead() knows that atomic loads
   // TODO: are safe to remove.
   return isa<LoadInst>(I);
 }
 
 bool SCCPSolver::tryToReplaceWithConstant(Value *V) {
-  Constant *Const = nullptr;
-  if (V->getType()->isStructTy()) {
-    std::vector<ValueLatticeElement> IVs = getStructLatticeValueFor(V);
-    if (llvm::any_of(IVs, isOverdefined))
-      return false;
-    std::vector<Constant *> ConstVals;
-    auto *ST = cast<StructType>(V->getType());
-    for (unsigned i = 0, e = ST->getNumElements(); i != e; ++i) {
-      ValueLatticeElement V = IVs[i];
-      ConstVals.push_back(SCCPSolver::isConstant(V)
-                              ? getConstant(V)
-                              : UndefValue::get(ST->getElementType(i)));
-    }
-    Const = ConstantStruct::get(ST, ConstVals);
-  } else {
-    const ValueLatticeElement &IV = getLatticeValueFor(V);
-    if (isOverdefined(IV))
-      return false;
-
-    Const = SCCPSolver::isConstant(IV) ? getConstant(IV)
-                                       : UndefValue::get(V->getType());
-  }
-  assert(Const && "Constant is nullptr here!");
-
+  Constant *Const = getConstantOrNull(V);
+  if (!Const)
   // Replacing `musttail` instructions with constant breaks `musttail` invariant
   // unless the call itself can be removed.
   // Calls with "clang.arc.attachedcall" implicitly use the return value and
   // those uses cannot be updated with a constant.
   CallBase *CB = dyn_cast<CallBase>(V);
   if (CB && ((CB->isMustTailCall() &&
               !canRemoveInstruction(CB)) ||
              CB->getOperandBundle(LLVMContext::OB_clang_arc_attachedcall))) {
void markOverdefined(Value *V) {
     else
       markOverdefined(ValueState[V], V);
   }
 
   bool isStructLatticeConstant(Function *F, StructType *STy);
 
   Constant *getConstant(const ValueLatticeElement &LV) const;
 
+  Constant *getConstantOrNull(Value *V) const;
+
   SmallPtrSetImpl<Function *> &getArgumentTrackedFunctions() {
     return TrackingIncomingArguments;
   }
 
-  void markArgInFuncSpecialization(Function *F,
-                                   const SmallVectorImpl<ArgInfo> &Args);
+  void handleSpecializationArguments(Function *F,
 
   void markFunctionUnreachable(Function *F) {
     for (auto &BB : *F)
       BBExecutable.erase(&BB);
   }
 
   void solveWhileResolvedUndefsIn(Module &M) {
     bool ResolvedUndefs = true;
Constant *SCCPInstVisitor::getConstant(const ValueLatticeElement &LV) const {
   if (LV.isConstantRange()) {
     const auto &CR = LV.getConstantRange();
     if (CR.getSingleElement())
       return ConstantInt::get(Ctx, *CR.getSingleElement());
   }
   return nullptr;
 }
 
-void SCCPInstVisitor::markArgInFuncSpecialization(
-    Function *F, const SmallVectorImpl<ArgInfo> &Args) {
+Constant *SCCPInstVisitor::getConstantOrNull(Value *V) const {
+  Constant *Const = nullptr;
+  if (V->getType()->isStructTy()) {
+    std::vector<ValueLatticeElement> LVs = getStructLatticeValueFor(V);
+    if (any_of(LVs, SCCPSolver::isOverdefined))
+      return nullptr;
+    std::vector<Constant *> ConstVals;
+    auto *ST = cast<StructType>(V->getType());
+    for (unsigned I = 0, E = ST->getNumElements(); I != E; ++I) {
+      ValueLatticeElement LV = LVs[I];
+      ConstVals.push_back(SCCPSolver::isConstant(LV)
+                              ? getConstant(LV)
+                              : UndefValue::get(ST->getElementType(I)));
+    }
+    Const = ConstantStruct::get(ST, ConstVals);
+  } else {
+    const ValueLatticeElement &LV = getLatticeValueFor(V);
+    if (SCCPSolver::isOverdefined(LV))
+      return nullptr;
+    Const = SCCPSolver::isConstant(LV) ? getConstant(LV)
+                                       : UndefValue::get(V->getType());
+  }
+  assert(Const && "Constant is nullptr here!");
+  return Const;
+}
+
+void SCCPInstVisitor::handleSpecializationArguments(Function *F,
+                          const SmallVectorImpl<ArgInfo> &Args) {
   assert(!Args.empty() && "Specialization without arguments");
   assert(F->arg_size() == Args[0].Formal->getParent()->arg_size() &&
          "Functions should have the same number of arguments");
 
   auto Iter = Args.begin();
-  Argument *NewArg = F->arg_begin();
-  Argument *OldArg = Args[0].Formal->getParent()->arg_begin();
+  Function::arg_iterator NewArg = F->arg_begin();
+  Function::arg_iterator OldArg = Args[0].Formal->getParent()->arg_begin();
   for (auto End = F->arg_end(); NewArg != End; ++NewArg, ++OldArg) {
 
     LLVM_DEBUG(dbgs() << "SCCP: Marking argument "
                       << NewArg->getNameOrAsOperand() << "\n");
 
-    if (Iter != Args.end() && OldArg == Iter->Formal) {
-      // Mark the argument constants in the new function.
-      markConstant(NewArg, Iter->Actual);
+    // Mark the argument constants in the new function
+    // or copy the lattice state over from the old function.
+    if (Iter != Args.end() && Iter->Formal == &*OldArg) {
+      if (auto *STy = dyn_cast<StructType>(NewArg->getType())) {
+        for (unsigned I = 0, E = STy->getNumElements(); I != E; ++I) {
+          ValueLatticeElement &NewValue = StructValueState[{&*NewArg, I}];
+          NewValue.markConstant(Iter->Actual->getAggregateElement(I));
+        }
+      } else {
+        ValueState[&*NewArg].markConstant(Iter->Actual);
+      }
       ++Iter;
-    } else if (ValueState.count(OldArg)) {
-      // For the remaining arguments in the new function, copy the lattice state
-      // over from the old function.
-      //
-      // Note: This previously looked like this:
-      // ValueState[NewArg] = ValueState[OldArg];
-      // This is incorrect because the DenseMap class may resize the underlying
-      // memory when inserting `NewArg`, which will invalidate the reference to
-      // `OldArg`. Instead, we make sure `NewArg` exists before setting it.
-      auto &NewValue = ValueState[NewArg];
-      NewValue = ValueState[OldArg];
-      pushToWorkList(NewValue, NewArg);
+    } else {
+      if (auto *STy = dyn_cast<StructType>(NewArg->getType())) {
+        for (unsigned I = 0, E = STy->getNumElements(); I != E; ++I) {
+          ValueLatticeElement &NewValue = StructValueState[{&*NewArg, I}];
+          NewValue = StructValueState[{&*OldArg, I}];
+        }
+      } else {
+        auto &NewValue = ValueState[&*NewArg];
+        NewValue = ValueState[&*OldArg];
+      }
     }
   }
 }
 
 void SCCPInstVisitor::visitInstruction(Instruction &I) {
   // All the instructions we don't do any special handling for just
   // go to overdefined.
   LLVM_DEBUG(dbgs() << "SCCP: Don't know how to handle: " << I << '\n');
   markOverdefined(&I);
 }
 

 bool SCCPSolver::isStructLatticeConstant(Function *F, StructType *STy) {
   return Visitor->isStructLatticeConstant(F, STy);
 }
 
 Constant *SCCPSolver::getConstant(const ValueLatticeElement &LV) const {
   return Visitor->getConstant(LV);
 }
 
+Constant *SCCPSolver::getConstantOrNull(Value *V) const {
+  return Visitor->getConstantOrNull(V);
+}
+
 SmallPtrSetImpl<Function *> &SCCPSolver::getArgumentTrackedFunctions() {
   return Visitor->getArgumentTrackedFunctions();
 }
 
-void SCCPSolver::markArgInFuncSpecialization(
-    Function *F, const SmallVectorImpl<ArgInfo> &Args) {
-  Visitor->markArgInFuncSpecialization(F, Args);
+void SCCPSolver::handleSpecializationArguments(Function *F,
+                     const SmallVectorImpl<ArgInfo> &Args) {
+  Visitor->handleSpecializationArguments(F, Args);
 }
 
 void SCCPSolver::markFunctionUnreachable(Function *F) {
   Visitor->markFunctionUnreachable(F);
 }
 
 void SCCPSolver::visit(Instruction *I) { Visitor->visit(I); }
 
 void SCCPSolver::visitCall(CallInst &I) { Visitor->visitCall(I); }

llvm/test/Transforms/FunctionSpecialization/constant-struct.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+
+; RUN: opt -passes="ipsccp<func-spec>" -force-specialization \
+; RUN:     -funcspec-for-literal-constant -S < %s | FileCheck %s
+
+define i32 @foo(i32 %y0, i32 %y1) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[Y:%.*]] = insertvalue { i32, i32 } undef, i32 [[Y0:%.*]], 0
+; CHECK-NEXT:    [[YY:%.*]] = insertvalue { i32, i32 } [[Y]], i32 [[Y1:%.*]], 1
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i32 @add.1({ i32, i32 } { i32 2, i32 3 }, { i32, i32 } [[YY]])
+; CHECK-NEXT:    ret i32 [[CALL]]
+;
+entry:
+  %y = insertvalue { i32, i32 } undef, i32 %y0, 0
+  %yy = insertvalue { i32, i32 } %y, i32 %y1, 1
+  %call = tail call i32 @add({i32, i32} {i32 2, i32 3}, {i32, i32} %yy)
+  ret i32 %call
+}
+
+define i32 @bar(i32 %x0, i32 %x1) {
+; CHECK-LABEL: @bar(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[X:%.*]] = insertvalue { i32, i32 } undef, i32 [[X0:%.*]], 0
+; CHECK-NEXT:    [[XX:%.*]] = insertvalue { i32, i32 } [[X]], i32 [[X1:%.*]], 1
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i32 @add.2({ i32, i32 } [[XX]], { i32, i32 } { i32 3, i32 2 })
+; CHECK-NEXT:    ret i32 [[CALL]]
+;
+entry:
+  %x = insertvalue { i32, i32 } undef, i32 %x0, 0
+  %xx = insertvalue { i32, i32 } %x, i32 %x1, 1
+  %call = tail call i32 @add({i32, i32} %xx, {i32, i32} {i32 3, i32 2})
+  ret i32 %call
+}
+
+define internal i32 @add({i32, i32} %x, {i32, i32} %y) {
+entry:
+  %x0 = extractvalue {i32, i32} %x, 0
+  %y0 = extractvalue {i32, i32} %y, 0
+  %add0 = add nsw i32 %x0, %y0
+  %x1 = extractvalue {i32, i32} %x, 1
+  %y1 = extractvalue {i32, i32} %y, 1
+  %add1 = add nsw i32 %x1, %y1
+  %mul = mul i32 %add0, %add1
+  ret i32 %mul
+}