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

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 a -avx function, calling another -avx function, the
vector is passed in xmm. If we now inline it into a +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.

Fixes #52660.

Differential Revision: https://reviews.llvm.org/D116036

Author: nikic

Diff for: llvm/lib/Target/X86/X86TargetTransformInfo.cpp

@@ -43,6 +43,7 @@
 #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"
 
@@ -5187,9 +5188,48 @@ bool X86TTIImpl::areInlineCompatible(const Function *Caller,
   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,

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

@@ -5,18 +5,21 @@
 
 ; 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
 }
@@ -33,18 +36,21 @@ define i64 @callee_not_avx(<4 x i64> %arg) noinline {
 
 ; 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
 }