[WebAssembly] Make stack pointer args inhibit tail calls

[tlively]

Summary

Also make return calls terminator instructions so epilogues are
inserted before them rather than after them. Together, these changes
make WebAssembly's tail call optimization more stack-safe.

[merge_guards_bot]

Unit tests: fail. 62433 tests passed, 1 failed and 845 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_timedmutex_requirements/thread_timedmutex_class/try_lock_for.pass.cpp

clang-tidy: fail. clang-tidy found 0 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[dschuff]

It doesn't look like there's anything wasm-specific here. Surely this also inhibits tail calling in other backends too? Are frontends suppost to avoid putting 'tail call' in the IR in this case?

[tlively]

In D73943#1857542, @dschuff wrote:

It doesn't look like there's anything wasm-specific here. Surely this also inhibits tail calling in other backends too? Are frontends suppost to avoid putting 'tail call' in the IR in this case?

I took a look at some of the other backends and their logic is much more complicated than what we have here. I didn't see any code that did exactly what I'm doing here, but I assume that they are doing something else to the same effect that is possibly more conservative. It's also possible that they do some magic to move stack arguments before tail calling. Experimentation shows that clang does not add tail to the call when a stack argument is passed anyway, so perhaps other frontends do that work themselves as well.

[dschuff]

The langref says (https://llvm.org/docs/LangRef.html#call-instruction) that the tail call marker implies that the callee does not access allocas from the caller. So it seems like it *should* mean that the backend can depend on this property (that you're checking for here). It also means that the frontend should guarantee it as best it knows how, and optimizations should not introduce it (or remove the attribute if they do?). There are lots of ways to sneak pointers into places (aliasing, going through memory, etc etc) so I'd expect the check in this CL to be brittle.

But if other backends are checking for this kind of property anyway that would be fishy, so now I'm a bit confused. Are they just checking for particular target-specific (or calling-convention-specific) properties?

[aheejin]

I also did a quick search for what other backends are doing, and they are all kind of different. Many targets seem to implement some function called [[ https://github.com/llvm/llvm-project/search?q=isEligibleForTailCallOptimization&unscoped_q=isEligibleForTailCallOptimization | isEligibleForTailCallOptimization ]], even though the superclass TargetLowering does not have that method. Anyway, I think it would be fine if we add things we need on an as-needed basis.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

691

I think there can be other cases other than geps and aliases; looking at this function might give some ideas. Skimming at the code, it looks [[ https://github.com/llvm/llvm-project/blob/0c3b2986ac6b71abc649811c3ec9cb0bf073c7d8/llvm/lib/IR/Value.cpp#L605-L607 | stripPointerCastsAndOffsets ]] does the most comprehensive checking we need.

[tlively]

In D73943#1858326, @dschuff wrote:

The langref says (https://llvm.org/docs/LangRef.html#call-instruction) that the tail call marker implies that the callee does not access allocas from the caller. So it seems like it *should* mean that the backend can depend on this property (that you're checking for here). It also means that the frontend should guarantee it as best it knows how, and optimizations should not introduce it (or remove the attribute if they do?)

I agree with your interpretation, but there doesn't seem to be any validation that the callee doesn't access allocas of the caller, so this change makes the backend more robust in the presence of incorrect IR that nonetheless passes IR validation. That seems like a useful but non-essential property to have, so I would be ok not merging this except that I don't think there are any downsides.

There are lots of ways to sneak pointers into places (aliasing, going through memory, etc etc) so I'd expect the check in this CL to be brittle.

The check does as much as it can, but you're right that it can't look through operations that lose type information. That means it is possible to generate IR that hides its reference to caller allocas in its tail call, but it's ok if we sometimes generate bad code in that case because the LangRef said you weren't supposed to do that.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

691

stripPointerCastsAndOffsets doesn't look through aliases, so it doesn't replace stripPointerCastsAndAliases. I could use both in the loop, but I think manually fetching the GEP pointer is clearer. I could probably be convinced otherwise, though.

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

691

Ah, you're right that it doesn't look through aliases. But the current implementation here does not look through bitcasts or call return args that stripPointerCastsAndOffsets. How about something like this?

Value *Val = Arg.get(), *NewVal = nullptr;
while (Val != NewVal) {
  NewVal = Val->stripPoinsterCastsAndOffsets()->stripPointerCastsAndAliases();
}

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

691

('checks' was missing at the end of first sentence)

[aheejin]

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

691

Talked in person, I was incorrect that this misses bitcasts and such; it includes checks for all that because this manually goes through geps.

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

WebAssemblyTargetLowering::LowerCall(CallLoweringInfo &CLI,
   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");
 
   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;
+    auto NoTail = [&](const char *Msg) {
+      if (CLI.CS && CLI.CS.isMustTailCall())
+        fail(DL, DAG, Msg);
+      CLI.IsTailCall = false;
+    };
+
+    if (!Subtarget->hasTailCall())
+      NoTail("WebAssembly 'tail-call' feature not enabled");
+
+    // Varargs calls cannot be tail calls because the buffer is on the stack
+    if (CLI.IsVarArg)
+      NoTail("WebAssembly does not support varargs tail calls");
+
+    if (!TypesMatch)
+      NoTail("WebAssembly tail call requires caller and callee return types to "
+             "match");
+
+    // If pointers to local stack values are passed, we cannot tail call
+    if (CLI.CS) {
+      for (auto &Arg : CLI.CS.args()) {
+        Value *Val = Arg.get();
+        // Trace the value back through pointer operations
+        while (true) {
+          Value *Src = Val->stripPointerCastsAndAliases();
+          if (auto *GEP = dyn_cast<GetElementPtrInst>(Src))
+            Src = GEP->getPointerOperand();
+          if (Val == Src)
+            break;
+          Val = Src;
         }
-    } 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");
+        if (isa<AllocaInst>(Val)) {
+          NoTail(
+              "WebAssembly does not support tail calling with stack arguments");
+          break;
         }
       }
     }
   }
 
   SmallVectorImpl<ISD::InputArg> &Ins = CLI.Ins;
   if (Ins.size() > 1)
     fail(DL, DAG, "WebAssembly doesn't support more than 1 returned value yet");

llvm/lib/Target/WebAssembly/WebAssemblyInstrCall.td

 
 let IsCanonical = 1 in {
 defm CALL_VOID :
   I<(outs), (ins function32_op:$callee, variable_ops),
     (outs), (ins function32_op:$callee),
     [(WebAssemblycall0 (i32 imm:$callee))],
     "call    \t$callee", "call\t$callee", 0x10>;
 
-let isReturn = 1 in
+let isReturn = 1, isTerminator = 1, hasCtrlDep = 1, isBarrier = 1 in
 defm RET_CALL :
   I<(outs), (ins function32_op:$callee, variable_ops),
     (outs), (ins function32_op:$callee),
     [(WebAssemblyretcall (i32 imm:$callee))],
     "return_call    \t$callee", "return_call\t$callee", 0x12>,
   Requires<[HasTailCall]>;
 
 let isCodeGenOnly = 1 in

llvm/test/CodeGen/WebAssembly/tailcall.ll

 ; 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
 }
 
+; Stack-allocated arguments inhibit tail calls
 
+; CHECK-LABEL: stack_arg:
+; CHECK: i32.call
+define i32 @stack_arg(i32* %x) {
+  %a = alloca i32
+  %v = tail call i32 @stack_arg(i32* %a)
+  ret i32 %v
+}
+
+; CHECK-LABEL: stack_arg_gep:
+; CHECK: i32.call
+define i32 @stack_arg_gep(i32* %x) {
+  %a = alloca { i32, i32 }
+  %p = getelementptr { i32, i32 }, { i32, i32 }* %a, i32 0, i32 1
+  %v = tail call i32 @stack_arg_gep(i32* %p)
+  ret i32 %v
+}
+
+; CHECK-LABEL: stack_arg_cast:
+; CHECK: global.get $push{{[0-9]+}}=, __stack_pointer
+; CHECK: global.set __stack_pointer, $pop{{[0-9]+}}
+; FAST: i32.call ${{[0-9]+}}=, stack_arg_cast, $pop{{[0-9]+}}
+; CHECK: global.set __stack_pointer, $pop{{[0-9]+}}
+; SLOW: return_call stack_arg_cast, ${{[0-9]+}}
+define i32 @stack_arg_cast(i32 %x) {
+  %a = alloca [64 x i32]
+  %i = ptrtoint [64 x i32]* %a to i32
+  %v = tail call i32 @stack_arg_cast(i32 %i)
+  ret i32 %v
+}
 
 ; Check that the signatures generated for external indirectly
 ; return-called functions include the proper return types
 
 ; YAML-LABEL: - Index:           8
 ; YAML-NEXT:    ParamTypes:
 ; YAML-NEXT:      - I32
 ; YAML-NEXT:      - F32