release/22.x: [X86] Improve illegal return type handling in FastISel (#186723)

llvm/lib/Target/X86/X86FastISel.cpp

@@ -3265,11 +3265,27 @@ bool X86FastISel::fastLowerCall(CallLoweringInfo &CLI) {
   bool Is64Bit        = Subtarget->is64Bit();
   bool IsWin64        = Subtarget->isCallingConvWin64(CC);
 
-  // If the return type is illegal, don't bother to promote it, just fall back
-  // to DAG ISel.
-  MVT RetVT;
-  if (!isTypeLegal(CLI.RetTy, RetVT) && !CLI.RetTy->isVoidTy())
-    return false;
+  // If the return type is illegal, check if the ABI requires a type conversion
+  // that FastISel cannot handle. Fall back to DAG ISel in such cases.
+  // For example, bfloat is returned as f16 in XMM0, however FastISel would
+  // assign f32 register type and store it in FuncInfo.ValueMap. This would
+  // cause DAG incorrectly perform type conversion from f32 to bfloat after get
+  // the value from FuncInfo.ValueMap.
+  // However, i1 is promoted to i8 and return i8 defined by ABI, so FastISel can
+  // lower it without switching to DAGISel.
+  MVT RetVT = MVT::Other;
+  if (!isTypeLegal(CLI.RetTy, RetVT) && !CLI.RetTy->isVoidTy()) {
+    if (RetVT == MVT::Other)
+      return false; // Unknown type, let DAG ISel handle it.
+
+    // RetVT is not MVT::Other, it must be simple now. It is something rely on
+    // the logic of isTypeLegal().
+    MVT ABIVT = TLI.getRegisterTypeForCallingConv(CLI.RetTy->getContext(),
+                                                  CLI.CallConv, RetVT);
+    MVT RegVT = TLI.getRegisterType(CLI.RetTy->getContext(), RetVT);
+    if (ABIVT != RegVT)
+      return false;
+  }
 
   // Call / invoke instructions with NoCfCheck attribute require special
   // handling.

llvm/test/CodeGen/X86/i1-fast-isel.ll

@@ -0,0 +1,106 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc --fast-isel < %s -mtriple=x86_64-unknown-unknown | FileCheck %s
+
+define i8 @test_direct_call(ptr %f) nounwind {
+; CHECK-LABEL: test_direct_call:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    callq foo@PLT
+; CHECK-NEXT:    movzbl %al, %edi
+; CHECK-NEXT:    callq bar@PLT
+; CHECK-NEXT:    popq %rcx
+; CHECK-NEXT:    retq
+entry:
+  %call = call i1 @foo(ptr %f)
+  %call2 = call zeroext i8 @bar(i1 %call)
+  ret i8 %call2
+}
+
+define i8 @test_fast_direct_call(ptr %f) nounwind {
+; CHECK-LABEL: test_fast_direct_call:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    callq foo_fast@PLT
+; CHECK-NEXT:    movzbl %al, %edi
+; CHECK-NEXT:    callq bar@PLT
+; CHECK-NEXT:    popq %rcx
+; CHECK-NEXT:    retq
+entry:
+  %call = call fastcc i1 @foo_fast(ptr %f)
+  %call2 = call zeroext i8 @bar(i1 %call)
+  ret i8 %call2
+}
+
+define i8 @test_indirect_all(ptr %fptr, ptr %f) nounwind {
+; CHECK-LABEL: test_indirect_all:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    movq %rdi, %rbx
+; CHECK-NEXT:    movq %rsi, %rdi
+; CHECK-NEXT:    callq foo@PLT
+; CHECK-NEXT:    movzbl %al, %edi
+; CHECK-NEXT:    callq *%rbx
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    retq
+entry:
+  %call = call i1 @foo(ptr %f)
+  %call2 = call zeroext i8 %fptr(i1 %call)
+  ret i8 %call2
+}
+
+define i8 @test_indirect_all2(ptr %fptr, ptr %f, i1 %cond) nounwind {
+; CHECK-LABEL: test_indirect_all2:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rbp
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    movl %edx, %ebp
+; CHECK-NEXT:    movq %rdi, %rbx
+; CHECK-NEXT:    movq %rsi, %rdi
+; CHECK-NEXT:    callq foo@PLT
+; CHECK-NEXT:    testb $1, %bpl
+; CHECK-NEXT:    je .LBB3_2
+; CHECK-NEXT:  # %bb.1: # %exit
+; CHECK-NEXT:    movzbl %al, %edi
+; CHECK-NEXT:    callq *%rbx
+; CHECK-NEXT:    jmp .LBB3_3
+; CHECK-NEXT:  .LBB3_2: # %exit2
+; CHECK-NEXT:    movb $3, %al
+; CHECK-NEXT:  .LBB3_3: # %exit2
+; CHECK-NEXT:    addq $8, %rsp
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    retq
+entry:
+  %call = call i1 @foo(ptr %f)
+  br i1 %cond, label %exit, label %exit2
+
+exit:
+  %call2 = call zeroext i8 %fptr(i1 %call)
+  ret i8 %call2
+
+exit2:
+  ret i8 3
+}
+
+
+define i8 @test_fast_indirect_all(ptr %fptr, ptr %f) nounwind {
+; CHECK-LABEL: test_fast_indirect_all:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    pushq %rbx
+; CHECK-NEXT:    movq %rdi, %rbx
+; CHECK-NEXT:    movq %rsi, %rdi
+; CHECK-NEXT:    callq foo@PLT
+; CHECK-NEXT:    movzbl %al, %edi
+; CHECK-NEXT:    callq *%rbx
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    retq
+entry:
+  %call = call fastcc i1 @foo(ptr %f)
+  %call2 = call zeroext i8 %fptr(i1 %call)
+  ret i8 %call2
+}
+
+declare i1 @foo(ptr %f)
+declare zeroext i8 @bar(i1)
+declare fastcc i1 @foo_fast(ptr %f)