[WebAssembly] Fix conflict between ret legalization and sjlj

[loladiro]

Summary

When the WebAssembly backend encounters a return type that doesn't
fit within i32, SelectionDAG performs sret demotion, adding an
additional argument to the start of the function that contains
a pointer to an sret buffer to use instead. However, this conflicts
with the emscripten sjlj lowering pass. There we translate calls like:

	call {i32, i32} @foo()

into (in pseudo-llvm

	%addr = @foo
	call {i32, i32} @__invoke_{i32,i32}(%addr)

i.e. we perform an indirect call through an extra function.
However, the sret transform now transforms this into
the equivalent of

%addr = @foo
%sret = alloca {i32, i32}
call {i32, i32} @__invoke_{i32,i32}(%sret, %addr)

(while simultaneously translation the implementation of @foo as well).
Unfortunately, this doesn't work out. The __invoke_ ABI expected
the function address to be the first argument, causing crashes.

There is several possible ways to fix this:

1. Implementing the sret rewrite at the IR level as well and performing it as part of lowering to __invoke
2. Fixing the wasm backend to recognize that __invoke has a special ABI
3. A change to the binaryen/emscripten ABI to recognize this situation

This revision implements the middle option, teaching the backend to
treat __invoke_ functions specially in sret lowering. This is achieved
by

1. Introducing a new CallingConv ID for invoke functions
2. Adding a hook to CanLowerReturn to override which argument is used for the sret pointer.

This revision isn't quite complete and definitely needs a test
(and I'm considering dropping the hook in 2) in favor of just
swizzeling the arguments in lowerCall), but I have verified
that it indeed fixes the observed crashes, so it should be
sufficient to discuss whether this is the correct approach.

[aheejin]

Thank you for fixing this! This is a long-known bug we haven't fixed so far, because wasm C++ frontend lowers struct returns into sret arguments in the frontend, so we don't usually encounter this bug. Just out of curiosity, how did you discover this bug? (The reporter of this bug found it while using the Rust frontend, I heard)

I like this approach. I prefer the option 2 among the three options you described: 1 is redundantly doing what ISel is already doing, and 3 makes LLVM-backend generated code incorrect so we have to rely on the other tool to fix it, which seems not very desirable. As you said, having to add an extra argument to the interface of CanLowerReturn is a bit sad but I don't really mind that if others are OK with it.

@dschuff @tlively Any more opinions?

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

818 ↗

(On Diff #212390)

Real nit: In this file it seems we mostly don't add braces to one-line ifs. The same for the if clause below.

[tlively]

In D65463#1607742, @aheejin wrote:

I like this approach. I prefer the option 2 among the three options you described: 1 is redundantly doing what ISel is already doing, and 3 makes LLVM-backend generated code incorrect so we have to rely on the other tool to fix it, which seems not very desirable. As you said, having to add an extra argument to the interface of CanLowerReturn is a bit sad but I don't really mind that if others are OK with it.

@dschuff @tlively Any more opinions?

I agree that this is the best direction to go in. It avoids duplicated effort and magic.

[loladiro]

Updated to add a test and confine the changes to the wasm backend code.

[tlively]

LGTM with a small test format change.

llvm/test/CodeGen/WebAssembly/lower-em-sjlj-sret.ll

25 ↗

(On Diff #213216)

You can pass -wasm-keep-registers to llc to keep the output in SSA form, which usually is more readable in tests.

[aheejin]

LGTM with nits (+ -wasm-keep-registers thing as @tlively said). Thank you!

llvm/test/CodeGen/WebAssembly/lower-em-exceptions.ll

19 ↗

(On Diff #213216)

Nit: How about adding call cc{{.*}} here too show that invokes have been converted to calls, which is one of the main things this pass does? The same thing for similar lines in this file and lower-em-sjlj.ll too.

llvm/test/CodeGen/WebAssembly/lower-em-sjlj-sret.ll

13 ↗

(On Diff #213216)

Real nit: we can maybe remove hidden

[loladiro]

llvm/test/CodeGen/WebAssembly/lower-em-exceptions.ll

19 ↗

(On Diff #213216)

Ok, will do prior to commit.

llvm/test/CodeGen/WebAssembly/lower-em-sjlj-sret.ll

13 ↗

(On Diff #213216)

sure

[loladiro]

Just out of curiosity, how did you discover this bug?

Realized I missed this question on my first read through. This came up compiling the julia system image. We do generally also lower things to sret. However, for things like Union{RegEx, Int64}, we emit that as a tagged union {jl_value_t *, i8}, which doesn't get sret'ed, because ideally we'd like it in two registers.

[aheejin]

Oh I see. Thanks!

llvm/trunk/include/llvm/IR/CallingConv.h

enum {
     AMDGPU_ES = 96,
 
     // Calling convention between AArch64 Advanced SIMD functions
     AArch64_VectorCall = 97,
 
     /// Calling convention between AArch64 SVE functions
     AArch64_SVE_VectorCall = 98,
 
+    /// Calling convention for emscripten __invoke_* functions. The first
+    /// argument is required to be the function ptr being indirectly called.
+    /// The remainder matches the regular calling convention.
+    WASM_EmscriptenInvoke = 99,
+
     /// The highest possible calling convention ID. Must be some 2^k - 1.
     MaxID = 1023
   };
 
 } // end namespace CallingConv
 
 } // end namespace llvm
 
 #endif // LLVM_IR_CALLINGCONV_H

llvm/trunk/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

static bool callingConvSupported(CallingConv::ID CallConv) {
   // We currently support the language-independent target-independent
   // conventions. We don't yet have a way to annotate calls with properties like
   // "cold", and we don't have any call-clobbered registers, so these are mostly
   // all handled the same.
   return CallConv == CallingConv::C || CallConv == CallingConv::Fast ||
          CallConv == CallingConv::Cold ||
          CallConv == CallingConv::PreserveMost ||
          CallConv == CallingConv::PreserveAll ||
-         CallConv == CallingConv::CXX_FAST_TLS;
+         CallConv == CallingConv::CXX_FAST_TLS ||
+         CallConv == CallingConv::WASM_EmscriptenInvoke;
 }
 
 SDValue
 WebAssemblyTargetLowering::LowerCall(CallLoweringInfo &CLI,
                                      SmallVectorImpl<SDValue> &InVals) const {
   SelectionDAG &DAG = CLI.DAG;
   SDLoc DL = CLI.DL;
   SDValue Chain = CLI.Chain;
WebAssemblyTargetLowering::LowerCall(CallLoweringInfo &CLI,
   }
 
   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;
+
+  // The generic code may have added an sret argument. If we're lowering an
+  // invoke function, the ABI requires that the function pointer be the first
+  // argument, so we may have to swap the arguments.
+  if (CallConv == CallingConv::WASM_EmscriptenInvoke && Outs.size() >= 2 &&
+      Outs[0].Flags.isSRet()) {
+    std::swap(Outs[0], Outs[1]);
+    std::swap(OutVals[0], OutVals[1]);
+  }
+
   unsigned NumFixedArgs = 0;
   for (unsigned I = 0; I < Outs.size(); ++I) {
     const ISD::OutputArg &Out = Outs[I];
     SDValue &OutVal = OutVals[I];
     if (Out.Flags.isNest())
       fail(DL, DAG, "WebAssembly hasn't implemented nest arguments");
     if (Out.Flags.isInAlloca())
       fail(DL, DAG, "WebAssembly hasn't implemented inalloca arguments");

llvm/trunk/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

Value *WebAssemblyLowerEmscriptenEHSjLj::wrapInvoke(CallOrInvoke *CI) {
   // Invoke function wrapper in JavaScript
   SmallVector<Value *, 16> Args;
   // Put the pointer to the callee as first argument, so it can be called
   // within the invoke wrapper later
   Args.push_back(CI->getCalledValue());
   Args.append(CI->arg_begin(), CI->arg_end());
   CallInst *NewCall = IRB.CreateCall(getInvokeWrapper(CI), Args);
   NewCall->takeName(CI);
-  NewCall->setCallingConv(CI->getCallingConv());
+  NewCall->setCallingConv(CallingConv::WASM_EmscriptenInvoke);
   NewCall->setDebugLoc(CI->getDebugLoc());
 
   // Because we added the pointer to the callee as first argument, all
   // argument attribute indices have to be incremented by one.
   SmallVector<AttributeSet, 8> ArgAttributes;
   const AttributeList &InvokeAL = CI->getAttributes();
 
   // No attributes for the callee pointer.

llvm/trunk/test/CodeGen/WebAssembly/lower-em-exceptions-whitelist.ll

 
 define void @do_catch() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 ; CHECK-LABEL: @do_catch(
 entry:
   invoke void @foo()
           to label %invoke.cont unwind label %lpad
 ; CHECK: entry:
 ; CHECK-NEXT: store i32 0, i32*
-; CHECK-NEXT: call void @__invoke_void(void ()* @foo)
+; CHECK-NEXT: call cc{{.*}} void @__invoke_void(void ()* @foo)
 
 invoke.cont:                                      ; preds = %entry
   br label %try.cont
 
 lpad:                                             ; preds = %entry
   %0 = landingpad { i8*, i32 }
           catch i8* null
   %1 = extractvalue { i8*, i32 } %0, 0

llvm/trunk/test/CodeGen/WebAssembly/lower-em-exceptions.ll

 ; Test invoke instruction with clauses (try-catch block)
 define void @clause() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 ; CHECK-LABEL: @clause(
 entry:
   invoke void @foo(i32 3)
           to label %invoke.cont unwind label %lpad
 ; CHECK: entry:
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
-; CHECK-NEXT: call void @__invoke_void_i32(void (i32)* @foo, i32 3)
+; CHECK-NEXT: call cc{{.*}} void @__invoke_void_i32(void (i32)* @foo, i32 3)
 ; CHECK-NEXT: %[[__THREW__VAL:.*]] = load i32, i32* @__THREW__
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
 ; CHECK-NEXT: %cmp = icmp eq i32 %[[__THREW__VAL]], 1
 ; CHECK-NEXT: br i1 %cmp, label %lpad, label %invoke.cont
 
 invoke.cont:                                      ; preds = %entry
   br label %try.cont
 

 ; Test invoke instruction with filters (functions with throw(...) declaration)
 define void @filter() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 ; CHECK-LABEL: @filter(
 entry:
   invoke void @foo(i32 3)
           to label %invoke.cont unwind label %lpad
 ; CHECK: entry:
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
-; CHECK-NEXT: call void @__invoke_void_i32(void (i32)* @foo, i32 3)
+; CHECK-NEXT: call cc{{.*}} void @__invoke_void_i32(void (i32)* @foo, i32 3)
 ; CHECK-NEXT: %[[__THREW__VAL:.*]] = load i32, i32* @__THREW__
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
 ; CHECK-NEXT: %cmp = icmp eq i32 %[[__THREW__VAL]], 1
 ; CHECK-NEXT: br i1 %cmp, label %lpad, label %invoke.cont
 
 invoke.cont:                                      ; preds = %entry
   ret void
 

 ; Test if argument attributes indices in newly created call instructions are correct
 define void @arg_attributes() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 ; CHECK-LABEL: @arg_attributes(
 entry:
   %0 = invoke noalias i8* @bar(i8 signext 1, i8 zeroext 2)
           to label %invoke.cont unwind label %lpad
 ; CHECK: entry:
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
-; CHECK-NEXT: %0 = call noalias i8* @"__invoke_i8*_i8_i8"(i8* (i8, i8)* @bar, i8 signext 1, i8 zeroext 2)
+; CHECK-NEXT: %0 = call cc{{.*}} noalias i8* @"__invoke_i8*_i8_i8"(i8* (i8, i8)* @bar, i8 signext 1, i8 zeroext 2)
 
 invoke.cont:                                      ; preds = %entry
   br label %try.cont
 
 lpad:                                             ; preds = %entry
   %1 = landingpad { i8*, i32 }
           catch i8* bitcast (i8** @_ZTIi to i8*)
           catch i8* null

llvm/trunk/test/CodeGen/WebAssembly/lower-em-sjlj-sret.ll

+; RUN: llc < %s -asm-verbose=false -enable-emscripten-sjlj -wasm-keep-registers | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+%struct.__jmp_buf_tag = type { [6 x i32], i32, [32 x i32] }
+
+declare i32 @setjmp(%struct.__jmp_buf_tag*) #0
+declare {i32, i32} @returns_struct()
+
+; Test the combination of backend legalization of large return types and the
+; Emscripten sjlj transformation
+define {i32, i32} @legalized_to_sret() {
+entry:
+  %env = alloca [1 x %struct.__jmp_buf_tag], align 16
+  %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %env, i32 0, i32 0
+  %call = call i32 @setjmp(%struct.__jmp_buf_tag* %arraydecay) #0
+; This is the function pointer to pass to invoke.
+; It needs to be the first argument (that's what we're testing here)
+; CHECK:  i32.const $push[[FPTR:[0-9]+]]=, returns_struct
+; This is the sret stack region (as an offset from the stack pointer local)
+; CHECK:  call  "__invoke_{i32.i32}", $pop[[FPTR]]
+  %ret = call {i32, i32} @returns_struct()
+  ret {i32, i32} %ret
+}
+
+attributes #0 = { returns_twice }

llvm/trunk/test/CodeGen/WebAssembly/lower-em-sjlj.ll

 ; CHECK-NEXT: %[[SETJMP_TABLE1:.*]] = call i32* @saveSetjmp(%struct.__jmp_buf_tag* %[[ARRAYDECAY]], i32 1, i32* %[[SETJMP_TABLE]], i32 %[[SETJMP_TABLE_SIZE]])
 ; CHECK-NEXT: %[[SETJMP_TABLE_SIZE1:.*]] = call i32 @getTempRet0()
 ; CHECK-NEXT: br label %entry.split
 
 ; CHECK: entry.split:
 ; CHECK-NEXT: phi i32 [ 0, %entry ], [ %[[LONGJMP_RESULT:.*]], %if.end ]
 ; CHECK-NEXT: %[[ARRAYDECAY1:.*]] = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %[[BUF]], i32 0, i32 0
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
-; CHECK-NEXT: call void @"__invoke_void_%struct.__jmp_buf_tag*_i32"(void (%struct.__jmp_buf_tag*, i32)* @emscripten_longjmp_jmpbuf, %struct.__jmp_buf_tag* %[[ARRAYDECAY1]], i32 1)
+; CHECK-NEXT: call cc{{.*}} void @"__invoke_void_%struct.__jmp_buf_tag*_i32"(void (%struct.__jmp_buf_tag*, i32)* @emscripten_longjmp_jmpbuf, %struct.__jmp_buf_tag* %[[ARRAYDECAY1]], i32 1)
 ; CHECK-NEXT: %[[__THREW__VAL:.*]] = load i32, i32* @__THREW__
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
 ; CHECK-NEXT: %[[CMP0:.*]] = icmp ne i32 %__THREW__.val, 0
 ; CHECK-NEXT: %[[THREWVALUE_VAL:.*]] = load i32, i32* @__threwValue
 ; CHECK-NEXT: %[[CMP1:.*]] = icmp ne i32 %[[THREWVALUE_VAL]], 0
 ; CHECK-NEXT: %[[CMP:.*]] = and i1 %[[CMP0]], %[[CMP1]]
 ; CHECK-NEXT: br i1 %[[CMP]], label %if.then1, label %if.else1
 
entry:
   %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %buf, i32 0, i32 0
   %call = call i32 @setjmp(%struct.__jmp_buf_tag* %arraydecay) #0
   call void @foo()
   ret void
 ; CHECK: entry:
 ; CHECK: %[[SETJMP_TABLE:.*]] = call i32* @saveSetjmp(
 
 ; CHECK: entry.split:
-; CHECK: call void @__invoke_void(void ()* @foo)
+; CHECK: @__invoke_void(void ()* @foo)
 
 ; CHECK: entry.split.split:
 ; CHECK-NEXT: %[[BUF:.*]] = bitcast i32* %[[SETJMP_TABLE]] to i8*
 ; CHECK-NEXT: tail call void @free(i8* %[[BUF]])
 ; CHECK-NEXT: ret void
 }
 
 ; Test a case when a function call is within try-catch, after a setjmp
 define hidden void @exception_and_longjmp() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 ; CHECK-LABEL: @exception_and_longjmp
 entry:
   %buf = alloca [1 x %struct.__jmp_buf_tag], align 16
   %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %buf, i32 0, i32 0
   %call = call i32 @setjmp(%struct.__jmp_buf_tag* %arraydecay) #0
   invoke void @foo()
           to label %try.cont unwind label %lpad
 
 ; CHECK: entry.split:
 ; CHECK: store i32 0, i32* @__THREW__
-; CHECK-NEXT: call void @__invoke_void(void ()* @foo)
+; CHECK-NEXT: call cc{{.*}} void @__invoke_void(void ()* @foo)
 ; CHECK-NEXT: %[[__THREW__VAL:.*]] = load i32, i32* @__THREW__
 ; CHECK-NEXT: store i32 0, i32* @__THREW__
 ; CHECK-NEXT: %[[CMP0:.*]] = icmp ne i32 %[[__THREW__VAL]], 0
 ; CHECK-NEXT: %[[THREWVALUE_VAL:.*]] = load i32, i32* @__threwValue
 ; CHECK-NEXT: %[[CMP1:.*]] = icmp ne i32 %[[THREWVALUE_VAL]], 0
 ; CHECK-NEXT: %[[CMP:.*]] = and i1 %[[CMP0]], %[[CMP1]]
 ; CHECK-NEXT: br i1 %[[CMP]], label %if.then1, label %if.else1
 

 ; Test that the allocsize attribute is being transformed properly
 declare i8 *@allocator(i32, %struct.__jmp_buf_tag*) #3
 define hidden i8 *@allocsize() {
 ; CHECK-LABEL: @allocsize
 entry:
   %buf = alloca [1 x %struct.__jmp_buf_tag], align 16
   %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %buf, i32 0, i32 0
   %call = call i32 @setjmp(%struct.__jmp_buf_tag* %arraydecay) #0
-; CHECK: i8* @"__invoke_i8*_i32_%struct.__jmp_buf_tag*"([[ARGS:.*]]) #[[ALLOCSIZE_ATTR:[0-9]+]]
+; CHECK: call cc{{.*}} i8* @"__invoke_i8*_i32_%struct.__jmp_buf_tag*"([[ARGS:.*]]) #[[ALLOCSIZE_ATTR:[0-9]+]]
   %alloc = call i8* @allocator(i32 20, %struct.__jmp_buf_tag* %arraydecay) #3
   ret i8 *%alloc
 }
 
 declare void @foo()
 ; Function Attrs: returns_twice
 declare i32 @setjmp(%struct.__jmp_buf_tag*) #0
 ; Function Attrs: noreturn