[WebAssembly] Support WasmEHFuncInfo serialization

[aheejin]

Summary

This adds support for serialization of WasmEHFuncInfo, in the form of
<Source BB Number, Unwind destination BB number>. To make YAML mapping
work, we needed to make a copy of the existing SrcToUnwindDest map
within yaml::WebAssemblyMachineFunctionInfo.

It was hard to add EH MIR tests for CFGStackify because WasmEHFuncInfo
could not be read from test MIR files. This adds the serialization
support for that to make EH MIR tests easier.

[aheejin]

Fix comments

[aheejin]

Cosmetic changes in tests

[dschuff]

llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h

34

could this be const also?

[aheejin]

Add consts

llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h

 
 #ifndef LLVM_LIB_TARGET_WEBASSEMBLY_WEBASSEMBLYMACHINEFUNCTIONINFO_H
 #define LLVM_LIB_TARGET_WEBASSEMBLY_WEBASSEMBLYMACHINEFUNCTIONINFO_H
 
 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
 #include "llvm/BinaryFormat/Wasm.h"
 #include "llvm/CodeGen/MIRYamlMapping.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/WasmEHFuncInfo.h"
 #include "llvm/MC/MCSymbolWasm.h"
 
 namespace llvm {
 
 namespace yaml {
 struct WebAssemblyFunctionInfo;
 }
 
 /// This class is derived from MachineFunctionInfo and contains private
 /// WebAssembly-specific information for each MachineFunction.
 class WebAssemblyFunctionInfo final : public MachineFunctionInfo {
+  const MachineFunction &MF;
+
   std::vector<MVT> Params;
   std::vector<MVT> Results;
   std::vector<MVT> Locals;
 
   /// A mapping from CodeGen vreg index to WebAssembly register number.
   std::vector<unsigned> WARegs;
 
   /// A mapping from CodeGen vreg index to a boolean value indicating whether
class WebAssemblyFunctionInfo final : public MachineFunctionInfo {
   unsigned FrameBaseVreg = -1U;
   // The local holding the frame base. This is either FP or SP
   // after WebAssemblyExplicitLocals
   unsigned FrameBaseLocal = -1U;
 
   // Function properties.
   bool CFGStackified = false;
 
+  // Catchpad unwind destination info for wasm EH.
+  WasmEHFuncInfo *WasmEHInfo = nullptr;
+
 public:
-  explicit WebAssemblyFunctionInfo(MachineFunction &MF) {}
+  explicit WebAssemblyFunctionInfo(MachineFunction &MF)
+      : MF(MF), WasmEHInfo(MF.getWasmEHFuncInfo()) {}
   ~WebAssemblyFunctionInfo() override;
+
+  const MachineFunction &getMachineFunction() const { return MF; }
+
   void initializeBaseYamlFields(const yaml::WebAssemblyFunctionInfo &YamlMFI);
 
   void addParam(MVT VT) { Params.push_back(VT); }
   const std::vector<MVT> &getParams() const { return Params; }
 
   void addResult(MVT VT) { Results.push_back(VT); }
   const std::vector<MVT> &getResults() const { return Results; }
 
public:
   // For a given stackified WAReg, return the id number to print with push/pop.
   static unsigned getWARegStackId(unsigned Reg) {
     assert(Reg & INT32_MIN);
     return Reg & INT32_MAX;
   }
 
   bool isCFGStackified() const { return CFGStackified; }
   void setCFGStackified(bool Value = true) { CFGStackified = Value; }
+
+  WasmEHFuncInfo *getWasmEHFuncInfo() const { return WasmEHInfo; }
+  void setWasmEHFuncInfo(WasmEHFuncInfo *Info) { WasmEHInfo = Info; }
 };
 
 void computeLegalValueVTs(const Function &F, const TargetMachine &TM, Type *Ty,
                           SmallVectorImpl<MVT> &ValueVTs);
 
 // Compute the signature for a given FunctionType (Ty). Note that it's not the
 // signature for ContextFunc (ContextFunc is just used to get varous context)
 void computeSignatureVTs(const FunctionType *Ty, const Function *TargetFunc,
                          const Function &ContextFunc, const TargetMachine &TM,
                          SmallVectorImpl<MVT> &Params,
                          SmallVectorImpl<MVT> &Results);
 
 void valTypesFromMVTs(const ArrayRef<MVT> &In,
                       SmallVectorImpl<wasm::ValType> &Out);
 
 std::unique_ptr<wasm::WasmSignature>
 signatureFromMVTs(const SmallVectorImpl<MVT> &Results,
                   const SmallVectorImpl<MVT> &Params);
 
 namespace yaml {
 
+using BBNumberMap = DenseMap<int, int>;
+
 struct WebAssemblyFunctionInfo final : public yaml::MachineFunctionInfo {
   bool CFGStackified = false;
+  // The same as WasmEHFuncInfo's SrcToUnwindDest, but stored in the mapping of
+  // BB numbers
+  BBNumberMap SrcToUnwindDest;
 
   WebAssemblyFunctionInfo() = default;
   WebAssemblyFunctionInfo(const llvm::WebAssemblyFunctionInfo &MFI);
 
   void mappingImpl(yaml::IO &YamlIO) override;
   ~WebAssemblyFunctionInfo() = default;
 };
 
 template <> struct MappingTraits<WebAssemblyFunctionInfo> {
   static void mapping(IO &YamlIO, WebAssemblyFunctionInfo &MFI) {
     YamlIO.mapOptional("isCFGStackified", MFI.CFGStackified, false);
+    YamlIO.mapOptional("wasmEHFuncInfo", MFI.SrcToUnwindDest);
+  }
+};
+
+template <> struct CustomMappingTraits<BBNumberMap> {
+  static void inputOne(IO &YamlIO, StringRef Key,
+                       BBNumberMap &SrcToUnwindDest) {
+    YamlIO.mapRequired(Key.str().c_str(),
+                       SrcToUnwindDest[std::atoi(Key.str().c_str())]);
+  }
+
+  static void output(IO &YamlIO, BBNumberMap &SrcToUnwindDest) {
+    for (auto KV : SrcToUnwindDest)
+      YamlIO.mapRequired(std::to_string(KV.first).c_str(), KV.second);
   }
 };
 
 } // end namespace yaml
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.cpp

llvm::signatureFromMVTs(const SmallVectorImpl<MVT> &Results,
   auto Sig = std::make_unique<wasm::WasmSignature>();
   valTypesFromMVTs(Results, Sig->Returns);
   valTypesFromMVTs(Params, Sig->Params);
   return Sig;
 }
 
 yaml::WebAssemblyFunctionInfo::WebAssemblyFunctionInfo(
     const llvm::WebAssemblyFunctionInfo &MFI)
-    : CFGStackified(MFI.isCFGStackified()) {}
+    : CFGStackified(MFI.isCFGStackified()) {
+  auto *EHInfo = MFI.getWasmEHFuncInfo();
+  const llvm::MachineFunction &MF = MFI.getMachineFunction();
+  //  MFI.getWasmEHFuncInfo() is non-null only for functions with the
+  //  personality function.
+  if (EHInfo) {
+    // SrcToUnwindDest can contain stale mappings in case BBs are removed in
+    // optimizations, in case, for example, they are unreachable. We should not
+    // include their info.
+    SmallPtrSet<const MachineBasicBlock *, 16> MBBs;
+    for (const auto &MBB : MF)
+      MBBs.insert(&MBB);
+    for (auto KV : EHInfo->SrcToUnwindDest) {
+      auto *SrcBB = KV.first.get<MachineBasicBlock *>();
+      auto *DestBB = KV.second.get<MachineBasicBlock *>();
+      if (MBBs.count(SrcBB) && MBBs.count(DestBB))
+        SrcToUnwindDest[SrcBB->getNumber()] = DestBB->getNumber();
+    }
+  }
+}
 
 void yaml::WebAssemblyFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
   MappingTraits<WebAssemblyFunctionInfo>::mapping(YamlIO, *this);
 }
 
 void WebAssemblyFunctionInfo::initializeBaseYamlFields(
     const yaml::WebAssemblyFunctionInfo &YamlMFI) {
   CFGStackified = YamlMFI.CFGStackified;
+  if (WasmEHInfo) {
+    for (auto KV : YamlMFI.SrcToUnwindDest)
+      WasmEHInfo->setUnwindDest(MF.getBlockNumbered(KV.first),
+                                MF.getBlockNumbered(KV.second));
+  }
 }

llvm/test/CodeGen/WebAssembly/function-info.mir

-# RUN: llc -mtriple=wasm32-unknown-unknown -run-pass wasm-cfg-stackify %s -o - | FileCheck %s
+# RUN: llc -mtriple=wasm32-unknown-unknown -exception-model=wasm -mattr=+exception-handling -run-pass wasm-cfg-sort -run-pass wasm-cfg-stackify %s -o - | FileCheck %s
 
-# CHECK-LABEL: function_property_test
+--- |
+  target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+  target triple = "wasm32-unknown-unknown"
+
+  declare i32 @__gxx_wasm_personality_v0(...)
+  declare void @foo()
+  define void @function_property_test() {
+    ret void
+  }
+  define void @wasm_eh_info_test() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+    ret void
+  }
+...
+
+# CHECK-LABEL: name: function_property_test
 # CHECK: machineFunctionInfo:
 # CHECK:   isCFGStackified:   true
 name: function_property_test
 liveins:
   - { reg: '$arguments' }
 body: |
   bb.0:
     RETURN implicit-def dead $arguments
 ...
+
+---
+# CHECK-LABEL: name: wasm_eh_info_test
+# CHECK: machineFunctionInfo:
+name: wasm_eh_info_test
+liveins:
+  - { reg: '$arguments' }
+# CHECK:  wasmEHFuncInfo:
+# bb.2 becomes bb.1 and bb.3 becomes bb.2 after CFGSort.
+# CHECK-NEXT:    1:               2
+machineFunctionInfo:
+  wasmEHFuncInfo:
+    2:               3
+body:             |
+  bb.0:
+    successors: %bb.1, %bb.2
+    CALL @foo, implicit-def dead $arguments, implicit $sp32, implicit $sp64, implicit-def dead $arguments, implicit $sp32, implicit $sp64
+
+  bb.1:
+    RETURN implicit-def dead $arguments
+
+  bb.2 (landing-pad):
+    successors: %bb.1, %bb.3
+    %0:i32 = CATCH &__cpp_exception, implicit-def dead $arguments
+    CALL @foo, implicit-def dead $arguments, implicit $sp32, implicit $sp64, implicit-def dead $arguments, implicit $sp32, implicit $sp64
+    BR %bb.1, implicit-def $arguments
+
+  bb.3 (landing-pad):
+    CATCH_ALL implicit-def $arguments
+    RETHROW 0, implicit-def $arguments
+...