[WebAssembly] Do not emit tail calls with return type mismatch

[tlively]

Summary

return_call and return_call_indirect are only valid if the return
types of the callee and caller match. We were previously not enforcing
that, which was producing invalid modules.

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

649

It's preexisting, but GuaranteedTailCallOpt's description says:

When the flag is on, participating targets will perform tail call optimization on all calls which use the fastcc calling convention and which satisfy certain target-independent criteria (being at the end of a function, having the same return type as their parent function, etc.), using an alternate ABI if necessary.

Does this mean, should we at least try to do tail calls even if CLI.IsTailCall is not true, i.e., there is no tailcall property in the instruction of .ll file? If so, it seems the check for this condition should be outside of if (CLI.IsTailCall). Let me know if I'm mistaken.

664

We don't call fail in this case?

677

I'm unsure how the user interface should work. When should we error out and when should we silently switch to non-tail calls? I might be mistaken, but it looks to me that unless CLI.CS.isMustTailCall() is true we don't need to error out, no? But currently MustTail contains much broader conditions, so I'm wondering what the UI should look like.

[tlively]

• Record proper return types for indirect tail calls as well

[tlively]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

649

I don't believe this is necessary. I checked and neither RISCV nor AArch64 try to do tail calls unless CLI.IsTailCall is true. There is some pass that runs that marks all possible tail calls as tail, so I guess in the common case at least it would be duplicating effort to check again so late in the backend.

664

It should be an LLVM IR-level validation error if the prototypes of callers and their musttail callees do no match, so we should never observe this condition. I will add a comment.

677

RISCV and AArch only error out when CLI.CS.isMustTailCall() is true, as you say. I will update this code to match that behavior.

[tlively]

• Address comments, add a few more tests

[tlively]

• Do not allow varargs tail calls

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

649

Would we need a test case for vararg too?

677

Should we document about these restrictions (such as no vararg) here too, like other platforms?

llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp

238

Why should we do this? Is the return type of the tail-called function different from that of the caller? If so, wouldn't we have already disabled tail calling in ISelLowering as you did above?

llvm/test/CodeGen/WebAssembly/tailcall.ll

215

I'm not sure what this test is trying to test..?

[tlively]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

649

There already is one that was updated for this change.

677

Yes, I'll add documentation.

llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp

238

Up above the return types for the current machine instruction (Returns) are determined from the instruction's defs. But return_calls do not have defs, so we need this extra code to record their return types by examining the caller's return type.

llvm/test/CodeGen/WebAssembly/tailcall.ll

215

This tests that we get the ReturnType: I32 correctly. Without the change you asked about above we incorrectly determine that the return type is none.

[tlively]

• Update documentation

[aheejin]

LGTM!

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

649

Oh I didn't see it, sorry.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

WebAssemblyTargetLowering::LowerCall(CallLoweringInfo &CLI,
   CallingConv::ID CallConv = CLI.CallConv;
   if (!callingConvSupported(CallConv))
     fail(DL, DAG,
          "WebAssembly doesn't support language-specific or target-specific "
          "calling conventions yet");
   if (CLI.IsPatchPoint)
     fail(DL, DAG, "WebAssembly doesn't support patch point yet");
 
-  // Fail if tail calls are required but not enabled
-  if (!Subtarget->hasTailCall()) {
-    if ((CallConv == CallingConv::Fast && CLI.IsTailCall &&
-         MF.getTarget().Options.GuaranteedTailCallOpt) ||
-        (CLI.CS && CLI.CS.isMustTailCall()))
-      fail(DL, DAG, "WebAssembly 'tail-call' feature not enabled");
+  if (CLI.IsTailCall) {
+    bool MustTail = CLI.CS && CLI.CS.isMustTailCall();
+    if (Subtarget->hasTailCall() && !CLI.IsVarArg) {
+      // Do not tail call unless caller and callee return types match
+      const Function &F = MF.getFunction();
+      const TargetMachine &TM = getTargetMachine();
+      Type *RetTy = F.getReturnType();
+      SmallVector<MVT, 4> CallerRetTys;
+      SmallVector<MVT, 4> CalleeRetTys;
+      computeLegalValueVTs(F, TM, RetTy, CallerRetTys);
+      computeLegalValueVTs(F, TM, CLI.RetTy, CalleeRetTys);
+      bool TypesMatch = CallerRetTys.size() == CalleeRetTys.size() &&
+                        std::equal(CallerRetTys.begin(), CallerRetTys.end(),
+                                   CalleeRetTys.begin());
+      if (!TypesMatch) {
+        // musttail in this case would be an LLVM IR validation failure
+        assert(!MustTail);
+        CLI.IsTailCall = false;
+      }
+    } else {
       CLI.IsTailCall = false;
+      if (MustTail) {
+        if (CLI.IsVarArg) {
+          // The return would pop the argument buffer
+          fail(DL, DAG, "WebAssembly does not support varargs tail calls");
+        } else {
+          fail(DL, DAG, "WebAssembly 'tail-call' feature not enabled");
+        }
+      }
+    }
   }
 
   SmallVectorImpl<ISD::InputArg> &Ins = CLI.Ins;
   if (Ins.size() > 1)
     fail(DL, DAG, "WebAssembly doesn't support more than 1 returned value yet");
 
   SmallVectorImpl<ISD::OutputArg> &Outs = CLI.Outs;
   SmallVectorImpl<SDValue> &OutVals = CLI.OutVals;
   unsigned NumFixedArgs = 0;

llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp

case MachineOperand::MO_Immediate:
             if (MO.isReg())
               Params.push_back(getType(MRI.getRegClass(MO.getReg())));
 
           // call_indirect instructions have a callee operand at the end which
           // doesn't count as a param.
           if (WebAssembly::isCallIndirect(MI->getOpcode()))
             Params.pop_back();
 
+          // return_call_indirect instructions have the return type of the
+          // caller
+          if (MI->getOpcode() == WebAssembly::RET_CALL_INDIRECT) {
+            const Function &F = MI->getMF()->getFunction();
+            const TargetMachine &TM = MI->getMF()->getTarget();
+            Type *RetTy = F.getReturnType();
+            SmallVector<MVT, 4> CallerRetTys;
+            computeLegalValueVTs(F, TM, RetTy, CallerRetTys);
+            valTypesFromMVTs(CallerRetTys, Returns);
+          }
+
           auto *WasmSym = cast<MCSymbolWasm>(Sym);
           auto Signature = make_unique<wasm::WasmSignature>(std::move(Returns),
                                                             std::move(Params));
           WasmSym->setSignature(Signature.get());
           Printer.addSignature(std::move(Signature));
           WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
 
           const MCExpr *Expr = MCSymbolRefExpr::create(

llvm/test/CodeGen/WebAssembly/tailcall.ll

 ; RUN: llc < %s -asm-verbose=false -verify-machineinstrs -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -mattr=+tail-call | FileCheck --check-prefixes=CHECK,SLOW %s
 ; RUN: llc < %s -asm-verbose=false -verify-machineinstrs -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -fast-isel -mattr=+tail-call | FileCheck --check-prefixes=CHECK,FAST %s
+; RUN: llc < %s --filetype=obj -mattr=+tail-call | obj2yaml | FileCheck --check-prefix=YAML %s
 
-; Test that the tail-call attribute is accepted
+; Test that the tail calls lower correctly
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
 %fn = type <{i32 (%fn, i32, i32)*}>
 declare i1 @foo(i1)
 declare i1 @bar(i1)
 

 ; FAST: i32.call $push[[L:[0-9]+]]=, baz, $pop{{[0-9]+}}, $pop{{[0-9]+}}, $pop{{[0-9]+}}{{$}}
 ; FAST: return $pop[[L]]{{$}}
 declare i32 @baz(i32, i32, i32)
 define i32 @mismatched_prototypes() {
   %v = tail call i32 @baz(i32 0, i32 42, i32 6)
   ret i32 %v
 }
 
+; CHECK-LABEL: mismatched_return_void:
+; CHECK: i32.call $drop=, baz, $pop{{[0-9]+}}, $pop{{[0-9]+}}, $pop{{[0-9]+}}{{$}}
+; CHECK: return{{$}}
+define void @mismatched_return_void() {
+  %v = tail call i32 @baz(i32 0, i32 42, i32 6)
+  ret void
+}
+
+; CHECK-LABEL: mismatched_return_f32:
+; CHECK: i32.call $push[[L:[0-9]+]]=, baz, $pop{{[0-9]+}}, $pop{{[0-9]+}}, $pop{{[0-9]+}}{{$}}
+; CHECK: f32.reinterpret_i32 $push[[L1:[0-9]+]]=, $pop[[L]]{{$}}
+; CHECK: return $pop[[L1]]{{$}}
+define float @mismatched_return_f32() {
+  %v = tail call i32 @baz(i32 0, i32 42, i32 6)
+  %u = bitcast i32 %v to float
+  ret float %u
+}
+
+; CHECK-LABEL: mismatched_indirect_void:
+; CHECK: i32.call_indirect $drop=, $0, $1, $2, $0{{$}}
+; CHECK: return{{$}}
+define void @mismatched_indirect_void(%fn %f, i32 %x, i32 %y) {
+  %p = extractvalue %fn %f, 0
+  %v = tail call i32 %p(%fn %f, i32 %x, i32 %y)
+  ret void
+}
+
+; CHECK-LABEL: mismatched_indirect_f32:
+; CHECK: i32.call_indirect $push[[L:[0-9]+]]=, $0, $1, $2, $0{{$}}
+; CHECK: f32.reinterpret_i32 $push[[L1:[0-9]+]]=, $pop[[L]]{{$}}
+; CHECK: return $pop[[L1]]{{$}}
+define float @mismatched_indirect_f32(%fn %f, i32 %x, i32 %y) {
+  %p = extractvalue %fn %f, 0
+  %v = tail call i32 %p(%fn %f, i32 %x, i32 %y)
+  %u = bitcast i32 %v to float
+  ret float %u
+}
+
 ; CHECK-LABEL: mismatched_byval:
 ; CHECK: i32.store
 ; CHECK: return_call quux, $pop{{[0-9]+}}{{$}}
 declare i32 @quux(i32* byval)
 define i32 @mismatched_byval(i32* %x) {
   %v = tail call i32 @quux(i32* byval %x)
   ret i32 %v
 }
 
 ; CHECK-LABEL: varargs:
 ; CHECK: i32.store
-; CHECK: return_call var, $1{{$}}
+; CHECK: i32.call $0=, var, $1{{$}}
+; CHECK: return $0{{$}}
 declare i32 @var(...)
 define i32 @varargs(i32 %x) {
   %v = tail call i32 (...) @var(i32 %x)
   ret i32 %v
 }
 
+; Type transformations inhibit tail calls, even when they are nops
+
+; CHECK-LABEL: mismatched_return_zext:
+; CHECK: i32.call
+define i32 @mismatched_return_zext() {
+  %v = tail call i1 @foo(i1 1)
+  %u = zext i1 %v to i32
+  ret i32 %u
+}
+
+; CHECK-LABEL: mismatched_return_sext:
+; CHECK: i32.call
+define i32 @mismatched_return_sext() {
+  %v = tail call i1 @foo(i1 1)
+  %u = sext i1 %v to i32
+  ret i32 %u
+}
+
+; CHECK-LABEL: mismatched_return_trunc:
+; CHECK: i32.call
+declare i32 @int()
+define i1 @mismatched_return_trunc() {
+  %v = tail call i32 @int()
+  %u = trunc i32 %v to i1
+  ret i1 %u
+}
+
+
+
+; Check that the signatures generated for external indirectly
+; return-called functions include the proper return types
+
+; YAML-LABEL: - Index:           8
+; YAML-NEXT:    ReturnType:      I32
+; YAML-NEXT:    ParamTypes:
+; YAML-NEXT:      - I32
+; YAML-NEXT:      - F32
+; YAML-NEXT:      - I64
+; YAML-NEXT:      - F64
+define i32 @unique_caller(i32 (i32, float, i64, double)** %p) {
+  %f = load i32 (i32, float, i64, double)*, i32 (i32, float, i64, double)** %p
+  %v = tail call i32 %f(i32 0, float 0., i64 0, double 0.)
+  ret i32 %v
+}
+
 ; CHECK-LABEL: .section .custom_section.target_features
 ; CHECK-NEXT: .int8 1
 ; CHECK-NEXT: .int8 43
 ; CHECK-NEXT: .int8 9
 ; CHECK-NEXT: .ascii "tail-call"