[TTI][NFCI] Introduce two new target transform hooks

[jdoerfert]

Summary

TargetTransformInfo::willResultInPoisonOnCallEdge and
TargetTransformInfo::isValidCallBaseForCallee allow the middle end to
prune potential call targets if such a call would be undefined.
The default implementation does not try to be very agressive, TODO's are
placed for improvments, or target specific hooks.

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

822–833

Move this into Attributes somewhere to check for ABI compatible attributes?

823

Should also check inreg matches

823

Also careful of arbitrary string attributes?

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

843

signext and zeroext probably need to match, at least for the return

[jdoerfert]

Check for sized and scalable

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

823

Maybe skip address space pairs that aren't isNoopAddrSpaceCast

[jdoerfert]

Address comments.

[nikic]

Can you please explain what the motivation for this is? I mean, I can see the benefit in theory, but I'm not clear on where this would help in practice.

This seems like a very dangerous change to me. If we want to do this, I'd recommend to go for a default implementation of "always valid" and then carve out very specific exceptions based on real-world motivation.

Two immediate issues I see here are:

• On x86_32 at least, clang will pass arguments by value, even though they really have indirect byval ABI. The reasoning is that LLVM will actually lowering those arguments with the correct ABI, and passing values as SSA values is of course much more amenable to optimization. However, Rust will instead generate the "correct" byval ABI. These calls are incompatible at the IR level (you will have a byval ptr argument on one side an something like i32 on the other -- or worse, a non-byval ptr that has a completely different meaning!) but become compatible after lowering. Clang is doing something extremely fishy here, but we still need to make sure that this does not get miscompiled (as your current code will do, on multiple levels).
• There are recurring problems with callee vs caller ABI attribute mismatches. IIRC there were some attempts to clean this up as part of the opaque pointer migration, by making lowering only consider callee ABI attributes. However, this ran into a whole series of issues with attribute mismatches in practice, and ended up getting reverted. Maybe @aeubanks has more context on that. Missing signext/zeroext attributes are also a recurring problem for targets like s390x. And this also plays into the previous point: A ptr with and without byval are not necessarily an ABI mismatch after lowering.

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

846

I think you missed the base call site CC matches function CC

[jdoerfert]

I'll move @nikic's comment to https://discourse.llvm.org/t/ub-or-not-ub-whats-a-legal-call/

[jdoerfert]

Make it less aggressive. Mostly based on noundef and no more attribute mismatch identification.

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

855

I didn't know any mismatches were allowed?

856

Should just fix these underlying queries being non const

886

Typo arugments

890–893

Should move this to AttributeFuncs? Is there any useful case where this differs from areInlineCompatible?

898

Should at least check the ABI attributes

[arsenm]

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

815

I'd assume no

816

I'd assume no. Do we even allow unsized parameters?

llvm/include/llvm/Analysis/TargetTransformInfo.h

 
   /// \returns True if the caller and callee agree on how \p Types will be
   /// passed to or returned from the callee.
   /// to the callee.
   /// \param Types List of types to check.
   bool areTypesABICompatible(const Function *Caller, const Function *Callee,
                              const ArrayRef<Type *> &Types) const;
 
+  /// \returns True if \p ParamTy at a call base can be passed to \p ArgTy at a
+  /// callee.
+  bool willResultInPoisonOnCallEdge(Type *ParamTy, Type *ArgTy) const;
+
+  /// \returns True if \p CB may legally call \p Callee via an indirect call.
+  bool isValidCallBaseForCallee(const CallBase *CB,
+                                const Function *Callee) const;
+
   /// The type of load/store indexing.
   enum MemIndexedMode {
     MIM_Unindexed, ///< No indexing.
     MIM_PreInc,    ///< Pre-incrementing.
     MIM_PreDec,    ///< Pre-decrementing.
     MIM_PostInc,   ///< Post-incrementing.
     MIM_PostDec    ///< Post-decrementing.
   };
virtual void getMemcpyLoopResidualLoweringType(
       unsigned RemainingBytes, unsigned SrcAddrSpace, unsigned DestAddrSpace,
       unsigned SrcAlign, unsigned DestAlign,
       std::optional<uint32_t> AtomicCpySize) const = 0;
   virtual bool areInlineCompatible(const Function *Caller,
                                    const Function *Callee) const = 0;
   virtual bool areTypesABICompatible(const Function *Caller,
                                      const Function *Callee,
                                      const ArrayRef<Type *> &Types) const = 0;
+  virtual bool willResultInPoisonOnCallEdge(Type *ParamTy, Type *ArgTy) const = 0;
+  virtual bool isValidCallBaseForCallee(const CallBase *CB,
+                                        const Function *Callee) const = 0;
   virtual bool isIndexedLoadLegal(MemIndexedMode Mode, Type *Ty) const = 0;
   virtual bool isIndexedStoreLegal(MemIndexedMode Mode, Type *Ty) const = 0;
   virtual unsigned getLoadStoreVecRegBitWidth(unsigned AddrSpace) const = 0;
   virtual bool isLegalToVectorizeLoad(LoadInst *LI) const = 0;
   virtual bool isLegalToVectorizeStore(StoreInst *SI) const = 0;
   virtual bool isLegalToVectorizeLoadChain(unsigned ChainSizeInBytes,
                                            Align Alignment,
                                            unsigned AddrSpace) const = 0;
public:
   bool areInlineCompatible(const Function *Caller,
                            const Function *Callee) const override {
     return Impl.areInlineCompatible(Caller, Callee);
   }
   bool areTypesABICompatible(const Function *Caller, const Function *Callee,
                              const ArrayRef<Type *> &Types) const override {
     return Impl.areTypesABICompatible(Caller, Callee, Types);
   }
+  bool willResultInPoisonOnCallEdge(Type *ParamTy, Type *ArgTy) const override {
+    return Impl.willResultInPoisonOnCallEdge(ParamTy, ArgTy);
+  }
+  bool isValidCallBaseForCallee(const CallBase *CB,
+                                const Function *Callee) const override {
+    return Impl.isValidCallBaseForCallee(CB, Callee);
+  }
   bool isIndexedLoadLegal(MemIndexedMode Mode, Type *Ty) const override {
     return Impl.isIndexedLoadLegal(Mode, Ty, getDataLayout());
   }
   bool isIndexedStoreLegal(MemIndexedMode Mode, Type *Ty) const override {
     return Impl.isIndexedStoreLegal(Mode, Ty, getDataLayout());
   }
   unsigned getLoadStoreVecRegBitWidth(unsigned AddrSpace) const override {
     return Impl.getLoadStoreVecRegBitWidth(AddrSpace);

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

 /// a TargetTransformInfo-conforming class.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_TARGETTRANSFORMINFOIMPL_H
 #define LLVM_ANALYSIS_TARGETTRANSFORMINFOIMPL_H
 
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/VectorUtils.h"
+#include "llvm/IR/Attributes.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PatternMatch.h"
 #include <optional>
 #include <utility>
 
public:
   bool areTypesABICompatible(const Function *Caller, const Function *Callee,
                              const ArrayRef<Type *> &Types) const {
     return (Caller->getFnAttribute("target-cpu") ==
             Callee->getFnAttribute("target-cpu")) &&
            (Caller->getFnAttribute("target-features") ==
             Callee->getFnAttribute("target-features"));
   }
 
+  bool willResultInPoisonOnCallEdge(Type *ParamTy, Type *ArgTy) const {
+    if (ParamTy == ArgTy)
+      return false;
+
+    // TODO: Is scalable vs non-scalable OK?
+    // TODO: Is sized vs non-sized OK?
+    // TODO: For sized and non-scalable, we could check the sizes.
+
+    // Floating point vs non-floating point, not OK.
+    if (ParamTy->isFloatingPointTy() != ArgTy->isFloatingPointTy())
+      return true;
+
+    // Non-vector: ptr/int vs non-ptr/int, not OK.
+    if (!ArgTy->isVectorTy() && !ParamTy->isVectorTy())
+      if (ArgTy->isIntOrPtrTy() != ParamTy->isIntOrPtrTy())
+        return true;
+
+    if (ArgTy->isPointerTy() && ParamTy->isPointerTy()) {
+      // Implicit non-trivial AS cast, not OK.
+      if (!isNoopAddrSpaceCast(ParamTy->getPointerAddressSpace(),
+                               ArgTy->getPointerAddressSpace()))
+        return true;
+    }
+
+    // Give up.
+    return false;
+  }
+
+  bool isValidCallBaseForCallee(const CallBase *CB,
+                                const Function *Callee) const {
+    // Given a call and callee definition, if their calling conventions don't
+    // match, and neither one of the calling conventions is compatible with C
+    // calling convention this call must be unreachable, as the call is
+    // undefined. Declarations are excluded because a prototype may
+    // not actually end up matching the implementation's calling conv for a
+    // variety of reasons (e.g. it may be written in assembly).
+    // Note: Copied from InstCombineCalls, apparently TLI can be used there but
+    // not TTI (due to policy).
+    if (!Callee->isDeclaration() &&
+        (Callee->getCallingConv() != CB->getCallingConv() &&
+         !(Callee->getCallingConv() == llvm::CallingConv::C &&
+           TargetLibraryInfoImpl::isCallingConvCCompatible(
+               const_cast<CallBase *>(CB))) &&
+         !(CB->getCallingConv() == llvm::CallingConv::C &&
+           TargetLibraryInfoImpl::isCallingConvCCompatible(
+               const_cast<Function *>(Callee)))))
+      return false;
+
+    // TODO: Check the additional musttail rules here.
+
+    Type *FnRetTy = Callee->getReturnType();
+    Type *CBRetTy = CB->getType();
+    if (!CBRetTy->isVoidTy() && CB->hasRetAttr(Attribute::NoUndef)) {
+      // Void function would "return poison", which would be UB.
+      if (FnRetTy->isVoidTy())
+        return false;
+      // Non-void function that would "return poison", which would be UB.
+      if (!FnRetTy->isVoidTy() &&
+          willResultInPoisonOnCallEdge(CBRetTy, FnRetTy))
+        return false;
+    }
+
+    // TODO: Depending on the target, call sites might not be allowed void if
+    // the callee is not returning void.
+
+    unsigned FnNumArgs = Callee->arg_size();
+    unsigned CBNumArgs = CB->arg_size();
+
+    // Check arguments that are poison due to the lack of call site
+    // parameters.
+    for (unsigned ArgNo = CBNumArgs; ArgNo < FnNumArgs; ++ArgNo)
+      if (Callee->getArg(ArgNo)->hasAttribute(Attribute::NoUndef))
+        return false;
+
+    // TODO: Depending on the source and maybe target we could disallow too
+    // many arugments. For now, we allow it.
+
+    // Check if we will violate noundef for arguments.
+    for (unsigned ArgNo = 0; ArgNo < std::min(CBNumArgs, FnNumArgs); ++ArgNo)
+      if (Callee->hasParamAttribute(ArgNo, Attribute::NoUndef) &&
+          willResultInPoisonOnCallEdge(CB->getArgOperand(ArgNo)->getType(),
+                                       Callee->getArg(ArgNo)->getType()))
+        return false;
+
+    // TODO: Check FnNumArgs till CBNumArgs for OK variadic argument
+    // types/attributes.
+
+    // TODO: Check attributes.
+
+    // Seems valid enough.
+    return true;
+  }
+
   bool isIndexedLoadLegal(TTI::MemIndexedMode Mode, Type *Ty,
                           const DataLayout &DL) const {
     return false;
   }
 
   bool isIndexedStoreLegal(TTI::MemIndexedMode Mode, Type *Ty,
                            const DataLayout &DL) const {
     return false;

llvm/include/llvm/IR/Attributes.h

 /// include attributes like nocapture, which constrain the function
 /// implementation rather than the passed value.
 AttributeMask getUBImplyingAttributes();
 
 /// \returns Return true if the two functions have compatible target-independent
 /// attributes for inlining purposes.
 bool areInlineCompatible(const Function &Caller, const Function &Callee);
 
-
 /// Checks  if there are any incompatible function attributes between
 /// \p A and \p B.
 ///
 /// \param [in] A - The first function to be compared with.
 /// \param [in] B - The second function to be compared with.
 /// \returns true if the functions have compatible attributes.
 bool areOutlineCompatible(const Function &A, const Function &B);
 

llvm/lib/Analysis/TargetTransformInfo.cpp

 }
 
 bool TargetTransformInfo::areTypesABICompatible(
     const Function *Caller, const Function *Callee,
     const ArrayRef<Type *> &Types) const {
   return TTIImpl->areTypesABICompatible(Caller, Callee, Types);
 }
 
+bool TargetTransformInfo::willResultInPoisonOnCallEdge(Type *ParamTy,
+                                                     Type *ArgTy) const {
+  return TTIImpl->willResultInPoisonOnCallEdge(ParamTy, ArgTy);
+}
+
+bool TargetTransformInfo::isValidCallBaseForCallee(
+    const CallBase *CB, const Function *Callee) const {
+  return TTIImpl->isValidCallBaseForCallee(CB, Callee);
+}
+
 bool TargetTransformInfo::isIndexedLoadLegal(MemIndexedMode Mode,
                                              Type *Ty) const {
   return TTIImpl->isIndexedLoadLegal(Mode, Ty);
 }
 
 bool TargetTransformInfo::isIndexedStoreLegal(MemIndexedMode Mode,
                                               Type *Ty) const {
   return TTIImpl->isIndexedStoreLegal(Mode, Ty);