[TTI][BPF] Ensure ArgumentPromotion Not Exceeding Target MaxArgs

[yonghong-song]

Summary

With LLVM patch https://reviews.llvm.org/D148269, we hit a linux kernel
bpf selftest compilation failure like below:

... 
progs/test_xdp_noinline.c:739:8: error: too many args to t8: i64 = GlobalAddress<ptr @encap_v4> 0, progs/test_xdp_noinline.c:739:8
            if (!encap_v4(xdp, cval, &pckt, dst, pkt_bytes))
                 ^   
... 
progs/test_xdp_noinline.c:321:6: error: defined with too many args
bool encap_v4(struct xdp_md *xdp, struct ctl_value *cval,
     ^   
...

Note that bpf selftests are compiled with -O2 which is
the recommended flag for bpf community.

The bpf backend calling convention is only allowing 5
parameters in registers and does not allow pass arguments
through stacks. In the above case, ArgumentPromotionPass
replaced parameter '&pckt' as two parameters, so the total
number of arguments after ArgumentPromotion pass becomes 6
and this caused later compilation failure during instruction
selection phase.

This patch added a TargetTransformInfo hook getMaxNumArgs()
which returns 5 for BPF and UINT_MAX for other targets.

[yonghong-song]

• add a test case

[aeubanks]

what about combining arguments in the backend (basically undoing this transformation)? that would generalize better to other passes and input IR

llvm/test/CodeGen/BPF/argumentpromotion.ll

1 ↗

(On Diff #514398)

this should only run argument promotion, not the entire O2 pipeline

also, this test case needs to be reduced a lot more, there's a lot of unrelated things in the IR here

[yonghong-song]

• for added test, do not use -O2, use -passes=argpromotion instead for opt. Also put the test under test/Transformations/ArgumentPromotion/BPF directory similar to ArgumentPromotion/X86.

[yonghong-song]

what about combining arguments in the backend (basically undoing this transformation)? that would generalize better to other passes and input IR

I would prefer to use a TTI hook as the hook itself essentially a profitability check. If the number of arguments is greater than the backend specified, the transformation is not profitable.
I suspect any other phases to do similar argument expanding transformation. If this does happen, we can use the same hook in those other transformations.

In our case, if input IR contains more than 5 arguments, the compiler will error out. This is okay and expected in bpf community since it is explicit in BPF specification that maximum 5 arguments are allowed. We do not need to add optimization to cope with such cases. We only need to deal with cases where user code is valid in terms of number of arguments but the compiler transformation makes it invalid.

llvm/test/CodeGen/BPF/argumentpromotion.ll

1 ↗

(On Diff #514398)

Good suggestion. Just remove O2 and run argument promotion only for the test. Also simplified the input IR a lot.

[aeubanks]

lgtm with test cleanup

llvm/test/Transforms/ArgumentPromotion/BPF/argpromotion.ll

25

a lot of this can still be removed, like dso_local, local_unnamed_addr, fastcc

[yonghong-song]

• remove function 'annotation' like dso_local, fastcc, local_unnamed_addr, unnamed_addr, as they won't affect argspromotion transformation.

llvm/test/Transforms/ArgumentPromotion/BPF/argpromotion.ll

25

done!

llvm/include/llvm/Analysis/TargetTransformInfo.h

   /// Used by the LowerTypeTests pass, which constructs an IR inline assembler
   /// node containing a jump table in a format suitable for the target, so it
   /// needs to know what format of jump table it can legally use.
   ///
   /// For non-Arm targets, this function isn't used. It defaults to returning
   /// false, but it shouldn't matter what it returns anyway.
   bool hasArmWideBranch(bool Thumb) const;
 
+  /// \return The maximum number of function arguments the target supports.
+  unsigned getMaxNumArgs() const;
+
   /// @}
 
 private:
   /// The abstract base class used to type erase specific TTI
   /// implementations.
   class Concept;
 
   /// The template model for the base class which wraps a concrete
public:
   virtual unsigned getMinTripCountTailFoldingThreshold() const = 0;
   virtual bool enableScalableVectorization() const = 0;
   virtual bool supportsScalableVectors() const = 0;
   virtual bool hasActiveVectorLength(unsigned Opcode, Type *DataType,
                                      Align Alignment) const = 0;
   virtual VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const = 0;
   virtual bool hasArmWideBranch(bool Thumb) const = 0;
+  virtual unsigned getMaxNumArgs() const = 0;
 };
 
 template <typename T>
 class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
   T Impl;
 
 public:
   Model(T Impl) : Impl(std::move(Impl)) {}
public:
   VPLegalization
   getVPLegalizationStrategy(const VPIntrinsic &PI) const override {
     return Impl.getVPLegalizationStrategy(PI);
   }
 
   bool hasArmWideBranch(bool Thumb) const override {
     return Impl.hasArmWideBranch(Thumb);
   }
+
+  unsigned getMaxNumArgs() const override {
+    return Impl.getMaxNumArgs();
+  }
 };
 
 template <typename T>
 TargetTransformInfo::TargetTransformInfo(T Impl)
     : TTIImpl(new Model<T>(Impl)) {}
 
 /// Analysis pass providing the \c TargetTransformInfo.
 ///

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

public:
   getVPLegalizationStrategy(const VPIntrinsic &PI) const {
     return TargetTransformInfo::VPLegalization(
         /* EVLParamStrategy */ TargetTransformInfo::VPLegalization::Discard,
         /* OperatorStrategy */ TargetTransformInfo::VPLegalization::Convert);
   }
 
   bool hasArmWideBranch(bool) const { return false; }
 
+  unsigned getMaxNumArgs() const { return UINT_MAX; }
+
 protected:
   // Obtain the minimum required size to hold the value (without the sign)
   // In case of a vector it returns the min required size for one element.
   unsigned minRequiredElementSize(const Value *Val, bool &isSigned) const {
     if (isa<ConstantDataVector>(Val) || isa<ConstantVector>(Val)) {
       const auto *VectorValue = cast<Constant>(Val);
 
       // In case of a vector need to pick the max between the min

llvm/lib/Analysis/TargetTransformInfo.cpp

 TargetTransformInfo::getVPLegalizationStrategy(const VPIntrinsic &VPI) const {
   return TTIImpl->getVPLegalizationStrategy(VPI);
 }
 
 bool TargetTransformInfo::hasArmWideBranch(bool Thumb) const {
   return TTIImpl->hasArmWideBranch(Thumb);
 }
 
+unsigned TargetTransformInfo::getMaxNumArgs() const {
+  return TTIImpl->getMaxNumArgs();
+}
+
 bool TargetTransformInfo::shouldExpandReduction(const IntrinsicInst *II) const {
   return TTIImpl->shouldExpandReduction(II);
 }
 
 unsigned TargetTransformInfo::getGISelRematGlobalCost() const {
   return TTIImpl->getGISelRematGlobalCost();
 }
 

llvm/lib/Target/BPF/BPFTargetTransformInfo.h

public:
   TTI::MemCmpExpansionOptions enableMemCmpExpansion(bool OptSize,
                                                     bool IsZeroCmp) const {
     TTI::MemCmpExpansionOptions Options;
     Options.LoadSizes = {8, 4, 2, 1};
     Options.MaxNumLoads = TLI->getMaxExpandSizeMemcmp(OptSize);
     return Options;
   }
 
+  unsigned getMaxNumArgs() const {
+    return 5;
+  }
+
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_BPF_BPFTARGETTRANSFORMINFO_H

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

static Function *promoteArguments(Function *F, FunctionAnalysisManager &FAM,
 
   const DataLayout &DL = F->getParent()->getDataLayout();
   auto &AAR = FAM.getResult<AAManager>(*F);
   const auto &TTI = FAM.getResult<TargetIRAnalysis>(*F);
 
   // Check to see which arguments are promotable.  If an argument is promotable,
   // add it to ArgsToPromote.
   DenseMap<Argument *, SmallVector<OffsetAndArgPart, 4>> ArgsToPromote;
+  unsigned NumArgsAfterPromote = F->getFunctionType()->getNumParams();
   for (Argument *PtrArg : PointerArgs) {
     // Replace sret attribute with noalias. This reduces register pressure by
     // avoiding a register copy.
     if (PtrArg->hasStructRetAttr()) {
       unsigned ArgNo = PtrArg->getArgNo();
       F->removeParamAttr(ArgNo, Attribute::StructRet);
       F->addParamAttr(ArgNo, Attribute::NoAlias);
       for (Use &U : F->uses()) {
         CallBase &CB = cast<CallBase>(*U.getUser());
         CB.removeParamAttr(ArgNo, Attribute::StructRet);
         CB.addParamAttr(ArgNo, Attribute::NoAlias);
       }
     }
 
     // If we can promote the pointer to its value.
     SmallVector<OffsetAndArgPart, 4> ArgParts;
 
     if (findArgParts(PtrArg, DL, AAR, MaxElements, IsRecursive, ArgParts)) {
       SmallVector<Type *, 4> Types;
       for (const auto &Pair : ArgParts)
         Types.push_back(Pair.second.Ty);
 
       if (areTypesABICompatible(Types, *F, TTI)) {
+        NumArgsAfterPromote += ArgParts.size() - 1;
         ArgsToPromote.insert({PtrArg, std::move(ArgParts)});
       }
     }
   }
 
   // No promotable pointer arguments.
   if (ArgsToPromote.empty())
     return nullptr;
 
+  if (NumArgsAfterPromote > TTI.getMaxNumArgs())
+    return nullptr;
+
   return doPromotion(F, FAM, ArgsToPromote);
 }
 
 PreservedAnalyses ArgumentPromotionPass::run(LazyCallGraph::SCC &C,
                                              CGSCCAnalysisManager &AM,
                                              LazyCallGraph &CG,
                                              CGSCCUpdateResult &UR) {
   bool Changed = false, LocalChange;

llvm/test/Transforms/ArgumentPromotion/BPF/argpromotion.ll

+; RUN: opt -passes=argpromotion -mtriple=bpf-pc-linux -S %s | FileCheck %s
+; Source:
+;  struct t {
+;    int a, b, c, d, e, f, g;
+;  };
+;  __attribute__((noinline)) static int foo1(struct t *p1, struct t *p2, struct t *p3) {
+;    return p1->a + p1->b + p2->c + p2->e + p3->f + p3->g;
+;  }
+;  __attribute__((noinline)) static int foo2(struct t *p1, struct t *p2, struct t *p3) {
+;    return p1->a + p1->b + p2->c + p2->e + p3->f;
+;  }
+;  void init(void *);
+;  int bar(void) {
+;    struct t v1, v2, v3;
+;    init(&v1); init(&v2); init(&v3);
+;    return foo1(&v1, &v2, &v3) + foo2(&v1, &v2, &v3);
+;  }
+; Compilation flag:
+;   clang -target bpf -O2 -S t.c -mllvm -print-before=argpromotion -mllvm -print-module-scope
+;   and then do some manual tailoring to remove some attributes/metadata which is not used
+;   by argpromotion pass.
+
+%struct.t = type { i32, i32, i32, i32, i32, i32, i32 }
+
+define dso_local i32 @bar() local_unnamed_addr {
+entry:
+  %v1 = alloca %struct.t, align 4
+  %v2 = alloca %struct.t, align 4
+  %v3 = alloca %struct.t, align 4
+  call void @init(ptr noundef nonnull %v1)
+  call void @init(ptr noundef nonnull %v2)
+  call void @init(ptr noundef nonnull %v3)
+  %call = call fastcc i32 @foo1(ptr noundef nonnull %v1, ptr noundef nonnull %v2, ptr noundef nonnull %v3)
+  %call1 = call fastcc i32 @foo2(ptr noundef nonnull %v1, ptr noundef nonnull %v2, ptr noundef nonnull %v3)
+  %add = add nsw i32 %call, %call1
+  ret i32 %add
+}
+
+declare dso_local void @init(ptr noundef) local_unnamed_addr
+
+define internal fastcc i32 @foo1(ptr nocapture noundef readonly %p1, ptr nocapture noundef readonly %p2, ptr nocapture noundef readonly %p3) unnamed_addr {
+entry:
+  %0 = load i32, ptr %p1, align 4
+  %b = getelementptr inbounds %struct.t, ptr %p1, i64 0, i32 1
+  %1 = load i32, ptr %b, align 4
+  %add = add nsw i32 %1, %0
+  %c = getelementptr inbounds %struct.t, ptr %p2, i64 0, i32 2
+  %2 = load i32, ptr %c, align 4
+  %add1 = add nsw i32 %add, %2
+  %e = getelementptr inbounds %struct.t, ptr %p2, i64 0, i32 4
+  %3 = load i32, ptr %e, align 4
+  %add2 = add nsw i32 %add1, %3
+  %f = getelementptr inbounds %struct.t, ptr %p3, i64 0, i32 5
+  %4 = load i32, ptr %f, align 4
+  %add3 = add nsw i32 %add2, %4
+  %g = getelementptr inbounds %struct.t, ptr %p3, i64 0, i32 6
+  %5 = load i32, ptr %g, align 4
+  %add4 = add nsw i32 %add3, %5
+  ret i32 %add4
+}
+
+; Without number-of-argument constraint, argpromotion will create a function signature with 6 arguments. Since
+; bpf target only supports maximum 5 arguments, so no argpromotion here.
+;
+; CHECK:  i32 @foo1(ptr nocapture noundef readonly %p1, ptr nocapture noundef readonly %p2, ptr nocapture noundef readonly %p3)
+
+define internal fastcc i32 @foo2(ptr noundef %p1, ptr noundef %p2, ptr noundef %p3) unnamed_addr {
+entry:
+  %0 = load i32, ptr %p1, align 4
+  %b = getelementptr inbounds %struct.t, ptr %p1, i64 0, i32 1
+  %1 = load i32, ptr %b, align 4
+  %add = add nsw i32 %0, %1
+  %c = getelementptr inbounds %struct.t, ptr %p2, i64 0, i32 2
+  %2 = load i32, ptr %c, align 4
+  %add1 = add nsw i32 %add, %2
+  %e = getelementptr inbounds %struct.t, ptr %p2, i64 0, i32 4
+  %3 = load i32, ptr %e, align 4
+  %add2 = add nsw i32 %add1, %3
+  %f = getelementptr inbounds %struct.t, ptr %p3, i64 0, i32 5
+  %4 = load i32, ptr %f, align 4
+  %add3 = add nsw i32 %add2, %4
+  ret i32 %add3
+}
+
+; Without number-of-argument constraint, argpromotion will create a function signature with 5 arguments, which equals
+; the maximum number of argument permitted by bpf backend, so argpromotion result code does work.
+;
+; CHECK:  i32 @foo2(i32 %p1.0.val, i32 %p1.4.val, i32 %p2.8.val, i32 %p2.16.val, i32 %p3.20.val)

llvm/test/Transforms/ArgumentPromotion/BPF/lit.local.cfg

+if not 'BPF' in config.root.targets:
+    config.unsupported = True