[IR] Add CallSiteOnly flag to CallBase's attribute getters

[aqjune]

Summary

I found that a few methods in CallBase return the value of attributes
attached to the call instruction without looking into the attributes
in the called function's declaration.
This causes sub-optimal results in some places.
For example, Value::getPointerDereferenceableBytes is calling
CallBase::getDereferenceableBytes which does not look into
dereferenceable attributes at the function declaration.
If the declaration had larger dereferenceable attribute value, this
results in a sub-optimal result.

To address this, I added 'CallSiteOnly' flag and updated such users to set this flag to false to get better result.

[aqjune]

Well, I found a way to get around this issue without changing these APIs. :)

llvm/include/llvm/IR/InstrTypes.h

 //===- llvm/InstrTypes.h - Important Instruction subclasses -----*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
public:
   LLVM_ATTRIBUTE_DEPRECATED(unsigned getRetAlignment() const,
                             "Use getRetAlign() instead") {
     if (const auto MA = Attrs.getRetAlignment())
       return MA->value();
     return 0;
   }
 
   /// Extract the alignment of the return value.
-  MaybeAlign getRetAlign() const { return Attrs.getRetAlignment(); }
+  /// If CallSiteOnly is true, look into the attribute that is explicitly
+  /// attached to this call site only. Otherwise, look into the called
+  /// function's declaration/definition as well and return the maximum value.
+  MaybeAlign getRetAlign(bool CallSiteOnly = true) const {
+    MaybeAlign A = Attrs.getRetAlignment();
+    auto *Fn = getCalledFunction();
+    if (!CallSiteOnly && Fn) {
+      if (Fn->hasAttribute(AttributeList::ReturnIndex, Attribute::Alignment)) {
+        A = max(A, Fn->getAttribute(AttributeList::ReturnIndex,
+                                    llvm::Attribute::Alignment)
+                       .getAlignment()
+                       .getValue());
+      }
+    }
+    return A;
+  }
 
   /// Extract the alignment for a call or parameter (0=unknown).
   LLVM_ATTRIBUTE_DEPRECATED(unsigned getParamAlignment(unsigned ArgNo) const,
                             "Use getParamAlign() instead") {
     if (const auto MA = Attrs.getParamAlignment(ArgNo))
       return MA->value();
     return 0;
   }
 
   /// Extract the alignment for a call or parameter (0=unknown).
-  MaybeAlign getParamAlign(unsigned ArgNo) const {
-    return Attrs.getParamAlignment(ArgNo);
+  /// If CallSiteOnly is true, look into the attribute that is explicitly
+  /// attached to this call site only. Otherwise, look into the called
+  /// function's declaration/definition as well and return the maximum value.
+  MaybeAlign getParamAlign(unsigned ArgNo, bool CallSiteOnly = true) const {
+    MaybeAlign A = Attrs.getParamAlignment(ArgNo);
+    auto *Fn = getCalledFunction();
+    if (!CallSiteOnly && Fn && ArgNo < Fn->arg_size() &&
+        Fn->hasParamAttribute(ArgNo, Attribute::Alignment)) {
+      A = max(A, Fn->getParamAlign(ArgNo).getValue());
+    }
+    return A;
   }
 
   /// Extract the byval type for a call or parameter.
   Type *getParamByValType(unsigned ArgNo) const {
     Type *Ty = Attrs.getParamByValType(ArgNo);
     return Ty ? Ty : getArgOperand(ArgNo)->getType()->getPointerElementType();
   }
 
   /// Extract the preallocated type for a call or parameter.
   Type *getParamPreallocatedType(unsigned ArgNo) const {
     Type *Ty = Attrs.getParamPreallocatedType(ArgNo);
     return Ty ? Ty : getArgOperand(ArgNo)->getType()->getPointerElementType();
   }
 
   /// Extract the number of dereferenceable bytes for a call or
   /// parameter (0=unknown).
-  uint64_t getDereferenceableBytes(unsigned i) const {
-    return Attrs.getDereferenceableBytes(i);
+  /// If CallSiteOnly is true, look into the attribute that is explicitly
+  /// attached to this call site only. Otherwise, look into the called
+  /// function's declaration/definition as well and return the maximum value.
+  uint64_t getDereferenceableBytes(unsigned i, bool CallSiteOnly = true) const {
+    uint64_t Bytes = Attrs.getDereferenceableBytes(i);
+    auto *Fn = getCalledFunction();
+    if (!CallSiteOnly && Fn) {
+      bool Flag = i == AttributeList::ReturnIndex ||
+                  (i - AttributeList::FirstArgIndex) < Fn->arg_size();
+      if (Flag)
+        Bytes = std::max(Bytes, Fn->getDereferenceableBytes(i));
+    }
+    return Bytes;
   }
 
   /// Extract the number of dereferenceable_or_null bytes for a call or
   /// parameter (0=unknown).
-  uint64_t getDereferenceableOrNullBytes(unsigned i) const {
-    return Attrs.getDereferenceableOrNullBytes(i);
+  /// If CallSiteOnly is true, look into the attribute that is explicitly
+  /// attached to this call site only. Otherwise, look into the called
+  /// function's declaration/definition as well and return the maximum value.
+  uint64_t getDereferenceableOrNullBytes(unsigned i,
+                                         bool CallSiteOnly = true) const {
+    uint64_t Bytes = Attrs.getDereferenceableOrNullBytes(i);
+    auto *Fn = getCalledFunction();
+    if (!CallSiteOnly && Fn) {
+      bool Flag = i == AttributeList::ReturnIndex ||
+                  (i - AttributeList::FirstArgIndex) < Fn->arg_size();
+      if (Flag)
+        Bytes = std::max(Bytes, Fn->getDereferenceableOrNullBytes(i));
+    }
+    return Bytes;
   }
 
   /// Return true if the return value is known to be not null.
   /// This may be because it has the nonnull attribute, or because at least
   /// one byte is dereferenceable and the pointer is in addrspace(0).
   bool isReturnNonNull() const;
 
   /// Determine if the return value is marked with NoAlias attribute.

llvm/include/llvm/IR/IntrinsicInst.h

 //===-- llvm/IntrinsicInst.h - Intrinsic Instruction Wrappers ---*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
public:
 
   unsigned getDestAddressSpace() const {
     return cast<PointerType>(getRawDest()->getType())->getAddressSpace();
   }
 
   /// FIXME: Remove this function once transition to Align is over.
   /// Use getDestAlign() instead.
   unsigned getDestAlignment() const {
-    if (auto MA = getParamAlign(ARG_DEST))
+    if (auto MA = getParamAlign(ARG_DEST, /*CallSiteOnly=*/false))
       return MA->value();
     return 0;
   }
-  MaybeAlign getDestAlign() const { return getParamAlign(ARG_DEST); }
+  MaybeAlign getDestAlign() const {
+    return getParamAlign(ARG_DEST, /*CallSiteOnly=*/false);
+  }
 
   /// Set the specified arguments of the instruction.
   void setDest(Value *Ptr) {
     assert(getRawDest()->getType() == Ptr->getType() &&
            "setDest called with pointer of wrong type!");
     setArgOperand(ARG_DEST, Ptr);
   }
 
public:
 
   unsigned getSourceAddressSpace() const {
     return cast<PointerType>(getRawSource()->getType())->getAddressSpace();
   }
 
   /// FIXME: Remove this function once transition to Align is over.
   /// Use getSourceAlign() instead.
   unsigned getSourceAlignment() const {
-    if (auto MA = BaseCL::getParamAlign(ARG_SOURCE))
+    if (auto MA = BaseCL::getParamAlign(ARG_SOURCE, /*CallSiteOnly=*/false))
       return MA->value();
     return 0;
   }
 
   MaybeAlign getSourceAlign() const {
-    return BaseCL::getParamAlign(ARG_SOURCE);
+    return BaseCL::getParamAlign(ARG_SOURCE, /*CallSiteOnly=*/false);
   }
 
   void setSource(Value *Ptr) {
     assert(getRawSource()->getType() == Ptr->getType() &&
            "setSource called with pointer of wrong type!");
     BaseCL::setArgOperand(ARG_SOURCE, Ptr);
   }
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

 //===- SelectionDAGBuilder.cpp - Selection-DAG building -------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
void SelectionDAGBuilder::visitTargetIntrinsic(const CallInst &I,
 
   if (!I.getType()->isVoidTy()) {
     if (VectorType *PTy = dyn_cast<VectorType>(I.getType())) {
       EVT VT = TLI.getValueType(DAG.getDataLayout(), PTy);
       Result = DAG.getNode(ISD::BITCAST, getCurSDLoc(), VT, Result);
     } else
       Result = lowerRangeToAssertZExt(DAG, I, Result);
 
-    MaybeAlign Alignment = I.getRetAlign();
-    if (!Alignment)
-      Alignment = F->getAttributes().getRetAlignment();
+    MaybeAlign Alignment = I.getRetAlign(false);
     // Insert `assertalign` node if there's an alignment.
     if (InsertAssertAlign && Alignment) {
       Result =
           DAG.getAssertAlign(getCurSDLoc(), Result, Alignment.valueOrOne());
     }
 
     setValue(&I, Result);
   }

llvm/lib/IR/Instructions.cpp

 //===- Instructions.cpp - Implement the LLVM instructions -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
if (auto *F = getCalledFunction())
     return F->getIntrinsicID();
   return Intrinsic::not_intrinsic;
 }
 
 bool CallBase::isReturnNonNull() const {
   if (hasRetAttr(Attribute::NonNull))
     return true;
 
-  if (getDereferenceableBytes(AttributeList::ReturnIndex) > 0 &&
-           !NullPointerIsDefined(getCaller(),
-                                 getType()->getPointerAddressSpace()))
+  if (getDereferenceableBytes(AttributeList::ReturnIndex,
+                              /*CallSiteOnly=*/false) > 0 &&
+      !NullPointerIsDefined(getCaller(), getType()->getPointerAddressSpace()))
     return true;
 
   return false;
 }
 
 Value *CallBase::getReturnedArgOperand() const {
   unsigned Index;
 

llvm/lib/IR/Value.cpp

 //===-- Value.cpp - Implement the Value class -----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
if (DerefBytes == 0) {
       }
     }
 
     if (DerefBytes == 0) {
       DerefBytes = A->getDereferenceableOrNullBytes();
       CanBeNull = true;
     }
   } else if (const auto *Call = dyn_cast<CallBase>(this)) {
-    DerefBytes = Call->getDereferenceableBytes(AttributeList::ReturnIndex);
+    DerefBytes = Call->getDereferenceableBytes(AttributeList::ReturnIndex,
+                                               /*CallSiteOnly=*/false);
     if (DerefBytes == 0) {
-      DerefBytes =
-          Call->getDereferenceableOrNullBytes(AttributeList::ReturnIndex);
+      DerefBytes = Call->getDereferenceableOrNullBytes(
+          AttributeList::ReturnIndex, /*CallSiteOnly=*/false);
       CanBeNull = true;
     }
   } else if (const LoadInst *LI = dyn_cast<LoadInst>(this)) {
     if (MDNode *MD = LI->getMetadata(LLVMContext::MD_dereferenceable)) {
       ConstantInt *CI = mdconst::extract<ConstantInt>(MD->getOperand(0));
       DerefBytes = CI->getLimitedValue();
     }
     if (DerefBytes == 0) {

llvm/unittests/IR/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   Support
   Passes
   TransformUtils
   ScalarOpts
   )
 
 add_llvm_unittest(IRTests
   AbstractCallSiteTest.cpp
   AsmWriterTest.cpp
   AttributesTest.cpp
   BasicBlockTest.cpp
   CFGBuilder.cpp
+  CallBaseTest.cpp
   ConstantRangeTest.cpp
   ConstantsTest.cpp
   DataLayoutTest.cpp
   DebugInfoTest.cpp
   DebugTypeODRUniquingTest.cpp
   DemandedBitsTest.cpp
   DominatorTreeTest.cpp
   DominatorTreeBatchUpdatesTest.cpp

llvm/unittests/IR/CallBaseTest.cpp

+//===--------------- CallBaseTest.cpp - CallBase Unittests ----------------===//
+//
+// 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/AsmParser/Parser.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("CallBaseTests", errs());
+  return Mod;
+}
+
+TEST(CallBase, Call) {
+  LLVMContext C;
+
+  const char *IR =
+      "declare align 4 dereferenceable(100) dereferenceable_or_null(200) "
+      "    i8* @g(i8* align(16) "
+      "               dereferenceable(960) "
+      "               dereferenceable_or_null(9600) %p)\n"
+      "define void @f1() {\n"
+      "  call align 2 dereferenceable(60) dereferenceable_or_null(80) "
+      "     i8* @g(i8* align(8) "
+      "                dereferenceable(480) "
+      "                dereferenceable_or_null(4800) null)\n"
+      "  ret void\n"
+      "}\n"
+      "define void @f2() {\n"
+      "  call align 8 dereferenceable(112) dereferenceable_or_null(208) "
+      "     i8* @g(i8* align(32) "
+      "                dereferenceable(1920) "
+      "                dereferenceable_or_null(19200) null)\n"
+      "  ret void\n"
+      "}\n";
+
+  std::unique_ptr<Module> M = parseIR(C, IR);
+  ASSERT_TRUE(M);
+
+  auto RetIdx = AttributeList::ReturnIndex;
+  auto FstIdx = AttributeList::FirstArgIndex;
+
+  {
+    Function *F = M->getFunction("f1");
+    ASSERT_NE(F, nullptr);
+    const CallBase *I = dyn_cast<CallBase>(&*F->begin()->begin());
+    MaybeAlign A;
+
+    // I->getXX(_, true) < I->getXX(_, false) because the callsite attributes
+    // in f2 has smaller values than the attributes in the declaration of g.
+
+    ASSERT_EQ(I->getDereferenceableBytes(RetIdx, true), 60u);
+    ASSERT_EQ(I->getDereferenceableBytes(RetIdx, false), 100u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(RetIdx, true), 80u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(RetIdx, false), 200u);
+
+    A = I->getRetAlign(true);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 2u);
+    A = I->getRetAlign(false);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 4u);
+
+    ASSERT_EQ(I->getDereferenceableBytes(FstIdx, true), 480u);
+    ASSERT_EQ(I->getDereferenceableBytes(FstIdx, false), 960u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(FstIdx, true), 4800u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(FstIdx, false), 9600u);
+
+    A = I->getParamAlign(0, true);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 8u);
+    A = I->getParamAlign(0, false);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 16u);
+  }
+
+  {
+    Function *F = M->getFunction("f2");
+    ASSERT_NE(F, nullptr);
+    const CallBase *I = dyn_cast<CallBase>(&*F->begin()->begin());
+    MaybeAlign A;
+
+    // I->getXX(_, true) == I->getXX(_, false) because the callsite attributes
+    // in f1 has larger values than the attributes in the declaration of g.
+
+    ASSERT_EQ(I->getDereferenceableBytes(RetIdx, true), 112u);
+    ASSERT_EQ(I->getDereferenceableBytes(RetIdx, false), 112u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(RetIdx, true), 208u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(RetIdx, false), 208u);
+
+    A = I->getRetAlign(true);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 8u);
+    A = I->getRetAlign(false);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 8u);
+
+    ASSERT_EQ(I->getDereferenceableBytes(FstIdx, true), 1920u);
+    ASSERT_EQ(I->getDereferenceableBytes(FstIdx, false), 1920u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(FstIdx, true), 19200u);
+    ASSERT_EQ(I->getDereferenceableOrNullBytes(FstIdx, false), 19200u);
+
+    A = I->getParamAlign(0, true);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 32u);
+    A = I->getParamAlign(0, false);
+    ASSERT_TRUE(A.hasValue());
+    ASSERT_EQ(A.getValue().value(), 32u);
+  }
+}