[WebAssembly] Don't do SjLj transformation when there's only setjmp

[aheejin]

Summary

When there is a setjmp call in a function, we transform every callsite
of setjmp to record its information by calling saveSetjmp function,
and we also transform every callsite of a function that can longjmp to
to check if a longjmp occurred and if so jump to the corresponding
post-setjmp BB. Currently we are doing this for every function that
contains a call to setjmp, but if there is no other function call
within that function that can longjmp, this transformation of setjmp
callsite and all the preparation of setjmpTable in the entry of the
function are not necessary.

This checks if a setjmp-calling function has any other calls that can
longjmp, and if not, skips the function for the SjLj transformation.

[aheejin]

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

26

This is added because without this function call this function is not going to be a target of SjLj transformation with this CL.

[aheejin]

Make a function static

[dschuff]

llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

744

can this be collapsed into if (EnableEmSjLj && SetjmpF)?

749

[aheejin]

Address comments

llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp

class WebAssemblyLowerEmscriptenEHSjLj final : public ModulePass {
 
   Value *wrapInvoke(CallBase *CI);
   void wrapTestSetjmp(BasicBlock *BB, DebugLoc DL, Value *Threw,
                       Value *SetjmpTable, Value *SetjmpTableSize, Value *&Label,
                       Value *&LongjmpResult, BasicBlock *&EndBB);
   Function *getInvokeWrapper(CallBase *CI);
 
   bool areAllExceptionsAllowed() const { return EHAllowlistSet.empty(); }
-  bool canLongjmp(const Value *Callee) const;
-  bool isEmAsmCall(const Value *Callee) const;
   bool supportsException(const Function *F) const {
     return EnableEmEH && (areAllExceptionsAllowed() ||
                           EHAllowlistSet.count(std::string(F->getName())));
   }
 
   void rebuildSSA(Function &F);
 
 public:
Function *WebAssemblyLowerEmscriptenEHSjLj::getInvokeWrapper(CallBase *CI) {
 
   FunctionType *FTy = FunctionType::get(CalleeFTy->getReturnType(), ArgTys,
                                         CalleeFTy->isVarArg());
   Function *F = getEmscriptenFunction(FTy, "__invoke_" + Sig, M);
   InvokeWrappers[Sig] = F;
   return F;
 }
 
-bool WebAssemblyLowerEmscriptenEHSjLj::canLongjmp(const Value *Callee) const {
+static bool canLongjmp(const Value *Callee) {
   if (auto *CalleeF = dyn_cast<Function>(Callee))
     if (CalleeF->isIntrinsic())
       return false;
 
   // Attempting to transform inline assembly will result in something like:
   //     call void @__invoke_void(void ()* asm ...)
   // which is invalid because inline assembly blocks do not have addresses
   // and can't be passed by pointer. The result is a crash with illegal IR.
if (CalleeName == "__cxa_begin_catch" || CalleeName == "__cxa_end_catch" ||
       CalleeName == "__cxa_allocate_exception" || CalleeName == "__cxa_throw" ||
       CalleeName == "__clang_call_terminate")
     return false;
 
   // Otherwise we don't know
   return true;
 }
 
-bool WebAssemblyLowerEmscriptenEHSjLj::isEmAsmCall(const Value *Callee) const {
+static bool isEmAsmCall(const Value *Callee) {
   StringRef CalleeName = Callee->getName();
   // This is an exhaustive list from Emscripten's <emscripten/em_asm.h>.
   return CalleeName == "emscripten_asm_const_int" ||
          CalleeName == "emscripten_asm_const_double" ||
          CalleeName == "emscripten_asm_const_int_sync_on_main_thread" ||
          CalleeName == "emscripten_asm_const_double_sync_on_main_thread" ||
          CalleeName == "emscripten_asm_const_async_on_main_thread";
 }
static void replaceLongjmpWithEmscriptenLongjmp(Function *LongjmpF,
   // with (int(*)(jmp_buf*, int))emscripten_longjmp.
   if (!LongjmpF->uses().empty()) {
     Value *EmLongjmp =
         IRB.CreateBitCast(EmLongjmpF, LongjmpF->getType(), "em_longjmp");
     LongjmpF->replaceAllUsesWith(EmLongjmp);
   }
 }
 
+static bool containsLongjmpableCalls(const Function *F) {
+  for (const auto &BB : *F)
+    for (const auto &I : BB)
+      if (const auto *CB = dyn_cast<CallBase>(&I))
+        if (canLongjmp(CB->getCalledOperand()))
+          return true;
+  return false;
+}
+
 bool WebAssemblyLowerEmscriptenEHSjLj::runOnModule(Module &M) {
   LLVM_DEBUG(dbgs() << "********** Lower Emscripten EH & SjLj **********\n");
 
   LLVMContext &C = M.getContext();
   IRBuilder<> IRB(C);
 
   Function *SetjmpF = M.getFunction("setjmp");
   Function *LongjmpF = M.getFunction("longjmp");
-  bool SetjmpUsed = SetjmpF && !SetjmpF->use_empty();
-  bool LongjmpUsed = LongjmpF && !LongjmpF->use_empty();
-  DoSjLj = EnableEmSjLj && (SetjmpUsed || LongjmpUsed);
 
   auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
   assert(TPC && "Expected a TargetPassConfig");
   auto &TM = TPC->getTM<WebAssemblyTargetMachine>();
 
   if (EnableEmEH && TM.Options.ExceptionModel == ExceptionHandling::Wasm)
     report_fatal_error("-exception-model=wasm not allowed with "
                        "-enable-emscripten-cxx-exceptions");
if (EnableEmEH) {
     ResumeF = getEmscriptenFunction(ResumeFTy, "__resumeException", &M);
 
     // Register llvm_eh_typeid_for function
     FunctionType *EHTypeIDTy =
         FunctionType::get(IRB.getInt32Ty(), IRB.getInt8PtrTy(), false);
     EHTypeIDF = getEmscriptenFunction(EHTypeIDTy, "llvm_eh_typeid_for", &M);
   }
 
+  if (EnableEmSjLj && SetjmpF) {
+    // Precompute setjmp users
+    for (User *U : SetjmpF->users()) {
+      Function *UserF = cast<Instruction>(U)->getFunction();
+      // If a function that calls setjmp does not contain any other calls that
+      // can longjmp, we don't need to do any transformation on that function,
+      // so can ignore it
+      if (containsLongjmpableCalls(UserF))
+        SetjmpUsers.insert(UserF);
+    }
+  }
+
+  bool SetjmpUsed = SetjmpF && !SetjmpUsers.empty();
+  bool LongjmpUsed = LongjmpF && !LongjmpF->use_empty();
+  DoSjLj = EnableEmSjLj && (SetjmpUsed || LongjmpUsed);
+
   // Function registration and data pre-gathering for setjmp/longjmp handling
   if (DoSjLj) {
     // Register emscripten_longjmp function
     FunctionType *FTy = FunctionType::get(
         IRB.getVoidTy(), {getAddrIntType(&M), IRB.getInt32Ty()}, false);
     EmLongjmpF = getEmscriptenFunction(FTy, "emscripten_longjmp", &M);
 
     if (SetjmpF) {
       // Register saveSetjmp function
       FunctionType *SetjmpFTy = SetjmpF->getFunctionType();
       FTy = FunctionType::get(Type::getInt32PtrTy(C),
                               {SetjmpFTy->getParamType(0), IRB.getInt32Ty(),
                                Type::getInt32PtrTy(C), IRB.getInt32Ty()},
                               false);
       SaveSetjmpF = getEmscriptenFunction(FTy, "saveSetjmp", &M);
 
       // Register testSetjmp function
       FTy = FunctionType::get(
           IRB.getInt32Ty(),
           {getAddrIntType(&M), Type::getInt32PtrTy(C), IRB.getInt32Ty()},
           false);
       TestSetjmpF = getEmscriptenFunction(FTy, "testSetjmp", &M);
-
-      // Precompute setjmp users
-      for (User *U : SetjmpF->users()) {
-        auto *UI = cast<Instruction>(U);
-        SetjmpUsers.insert(UI->getFunction());
-      }
     }
   }
 
   // Exception handling transformation
   if (EnableEmEH) {
     for (Function &F : M) {
       if (F.isDeclaration())
         continue;

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

; CHECK-LABEL: entry
   ; CHECK: call i8* @malloc
   ; CHECK: call i32* @saveSetjmp
   ; CHECK: call i32 @getTempRet0
   %retval = alloca i32, align 4
   %jmp = alloca [1 x %struct.__jmp_buf_tag], align 16
   store i32 0, i32* %retval, align 4
   %arraydecay = getelementptr inbounds [1 x %struct.__jmp_buf_tag], [1 x %struct.__jmp_buf_tag]* %jmp, i32 0, i32 0
   %call = call i32 @setjmp(%struct.__jmp_buf_tag* %arraydecay) #0
+  call void @foo()
   ret void
 
 ; CHECK-LABEL: entry.split
-  ; CHECK: call void @free
-  ; CHECK: ret void
 }
 
 ; This is a dummy dlmalloc implemenation only to make compiler pass, because an
 ; alias (malloc) has to point an actual definition.
 define i8* @dlmalloc(i32) {
   %p = inttoptr i32 0 to i8*
   ret i8* %p
 }
 
+declare void @foo()
 ; Function Attrs: returns_twice
 declare i32 @setjmp(%struct.__jmp_buf_tag*) #0
 
 attributes #0 = { returns_twice }

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

 ; 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 where a function has a setjmp call but no other calls that can
+; longjmp. We don't need to do any transformation in this case.
+define void @setjmp_only(i8* %ptr) {
+; CHECK-LABEL: @setjmp_only
+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
+  ; free cannot longjmp
+  call void @free(i8* %ptr)
+  ret void
+; CHECK-NOT: @malloc
+; CHECK-NOT: %setjmpTable
+; CHECK-NOT: @saveSetjmp
+; CHECK-NOT: @testSetjmp
+}
+
 ; Test SSA validity
 define void @ssa(i32 %n) {
 ; CHECK-LABEL: @ssa
 entry:
   %buf = alloca [1 x %struct.__jmp_buf_tag], align 16
   %cmp = icmp sgt i32 %n, 5
   br i1 %cmp, label %if.then, label %if.end
 ; CHECK: entry: