[Inline][X86] Avoid inlining if it would create ABI-incompatible calls (PR52660)

[nikic]

Summary

X86 allows inlining functions if the callee target features are a subset of the caller target features. This ensures that we don't inline something into a caller that does not support it.

However, this does not account for possible call ABI mismatches as a result of inlining. If a call passing a vector argument was originally in -avx function, calling another -avx function, the vector is passed in xmm. If we now inline it into an +avx function, then it will be passed in ymm, even though the callee expects it in xmm.

Fix this by scanning over all calls in the function and checking whether ABI incompatibility is possible. Calls that only pass scalar types are excluded, as I believe those always use the same ABI independent of target features (right?)

Fixes https://github.com/llvm/llvm-project/issues/52660.

[pengfei]

The fix seems good. Just two concerns here:

• I wonder what's the algorithm when inlining. Given we have a calling chain A->B->C. If we can inline both B and C into A, is B inlined firstly to A, then C to AB, or C inlined firstly to B, then BC to A? I'd worry if we are pessimistic for the former case.
• Another concern is the compiling time. Scanning seems inevitable, but can we avoid to repeat scanning if the inlined function is called in many place?

[nikic]

In D116036#3206048, @pengfei wrote:

The fix seems good. Just two concerns here:

• I wonder what's the algorithm when inlining. Given we have a calling chain A->B->C. If we can inline both B and C into A, is B inlined firstly to A, then C to AB, or C inlined firstly to B, then BC to A? I'd worry if we are pessimistic for the former case.

With few exceptions, inlining happens bottom-up, so B->C will be inlined first and then A->BC. So calls in the callee will already be eliminated, unless they can't be inlined for some reason (noinline, cost etc).

• Another concern is the compiling time. Scanning seems inevitable, but can we avoid to repeat scanning if the inlined function is called in many place?

I think this will not have much practical impact mainly because this kind of target feature mismatch is rare. Usually target features on all functions are the same. Inlining already does multiple scans of the callee, so at least this doesn't change anything asymptotically.

[pengfei]

Thanks for the explanation! LGTM, but please leave some time for others comments.

llvm/lib/Target/X86/X86TargetTransformInfo.cpp

 /// code size, latency and uop count.
 //===----------------------------------------------------------------------===//
 
 #include "X86TargetTransformInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/BasicTTIImpl.h"
 #include "llvm/CodeGen/CostTable.h"
 #include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/Support/Debug.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "x86tti"
 
 //===----------------------------------------------------------------------===//
bool X86TTIImpl::areInlineCompatible(const Function *Caller,
   const TargetMachine &TM = getTLI()->getTargetMachine();
 
   // Work this as a subsetting of subtarget features.
   const FeatureBitset &CallerBits =
       TM.getSubtargetImpl(*Caller)->getFeatureBits();
   const FeatureBitset &CalleeBits =
       TM.getSubtargetImpl(*Callee)->getFeatureBits();
 
+  // Check whether features are the same (apart from the ignore list).
   FeatureBitset RealCallerBits = CallerBits & ~InlineFeatureIgnoreList;
   FeatureBitset RealCalleeBits = CalleeBits & ~InlineFeatureIgnoreList;
-  return (RealCallerBits & RealCalleeBits) == RealCalleeBits;
+  if (RealCallerBits == RealCalleeBits)
+    return true;
+
+  // If the features are a subset, we need to additionally check for calls
+  // that may become ABI-incompatible as a result of inlining.
+  if ((RealCallerBits & RealCalleeBits) != RealCalleeBits)
+    return false;
+
+  for (const Instruction &I : instructions(Callee)) {
+    if (const auto *CB = dyn_cast<CallBase>(&I)) {
+      SmallVector<Type *, 8> Types;
+      for (Value *Arg : CB->args())
+        Types.push_back(Arg->getType());
+      if (!CB->getType()->isVoidTy())
+        Types.push_back(CB->getType());
+
+      // Simple types are always ABI compatible.
+      auto IsSimpleTy = [](Type *Ty) {
+        return !Ty->isVectorTy() && !Ty->isAggregateType();
+      };
+      if (all_of(Types, IsSimpleTy))
+        continue;
+
+      if (Function *NestedCallee = CB->getCalledFunction()) {
+        // Assume that intrinsics are always ABI compatible.
+        if (NestedCallee->isIntrinsic())
+          continue;
+
+        // Do a precise compatibility check.
+        if (!areTypesABICompatible(Caller, NestedCallee, Types))
+          return false;
+      } else {
+        // We don't know the target features of the callee,
+        // assume it is incompatible.
+        return false;
+      }
+    }
+  }
+  return true;
 }
 
 bool X86TTIImpl::areTypesABICompatible(const Function *Caller,
                                        const Function *Callee,
                                        const ArrayRef<Type *> &Types) const {
   if (!BaseT::areTypesABICompatible(Caller, Callee, Types))
     return false;
 

llvm/test/Transforms/Inline/X86/call-abi-compatibility.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
 ; RUN: opt < %s -mtriple=x86_64-unknown-linux-gnu -S -inline | FileCheck %s
 
 ; Test for PR52660.
 
 ; This call should not get inlined, because it would make the callee_not_avx
 ; call ABI incompatible.
-; TODO: Currently gets inlined.
 define void @caller_avx() "target-features"="+avx" {
 ; CHECK-LABEL: define {{[^@]+}}@caller_avx
 ; CHECK-SAME: () #[[ATTR0:[0-9]+]] {
-; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @callee_not_avx(<4 x i64> <i64 0, i64 1, i64 2, i64 3>)
+; CHECK-NEXT:    call void @caller_not_avx()
 ; CHECK-NEXT:    ret void
 ;
   call void @caller_not_avx()
   ret void
 }
 
 define internal void @caller_not_avx() {
+; CHECK-LABEL: define {{[^@]+}}@caller_not_avx() {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @callee_not_avx(<4 x i64> <i64 0, i64 1, i64 2, i64 3>)
+; CHECK-NEXT:    ret void
+;
   call i64 @callee_not_avx(<4 x i64> <i64 0, i64 1, i64 2, i64 3>)
   ret void
 }
 
 define i64 @callee_not_avx(<4 x i64> %arg) noinline {
 ; CHECK-LABEL: define {{[^@]+}}@callee_not_avx
 ; CHECK-SAME: (<4 x i64> [[ARG:%.*]]) #[[ATTR1:[0-9]+]] {
 ; CHECK-NEXT:    [[V:%.*]] = extractelement <4 x i64> [[ARG]], i64 2
 ; CHECK-NEXT:    ret i64 [[V]]
 ;
   %v = extractelement <4 x i64> %arg, i64 2
   ret i64 %v
 }
 
 ; This call also shouldn't be inlined, as we don't know whether callee_unknown
 ; is ABI compatible or not.
-; TODO: Currently gets inlined.
 define void @caller_avx2() "target-features"="+avx" {
 ; CHECK-LABEL: define {{[^@]+}}@caller_avx2
 ; CHECK-SAME: () #[[ATTR0]] {
-; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @callee_unknown(<4 x i64> <i64 0, i64 1, i64 2, i64 3>)
+; CHECK-NEXT:    call void @caller_not_avx2()
 ; CHECK-NEXT:    ret void
 ;
   call void @caller_not_avx2()
   ret void
 }
 
 define internal void @caller_not_avx2() {
+; CHECK-LABEL: define {{[^@]+}}@caller_not_avx2() {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @callee_unknown(<4 x i64> <i64 0, i64 1, i64 2, i64 3>)
+; CHECK-NEXT:    ret void
+;
   call i64 @callee_unknown(<4 x i64> <i64 0, i64 1, i64 2, i64 3>)
   ret void
 }
 
 declare i64 @callee_unknown(<4 x i64>)
 
 ; This call should get inlined, because we assume that intrinsics are always
 ; ABI compatible.