[WebAssembly] Support WasmEHFuncInfo serialization

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.

Reviewed By: dschuff

Differential Revision: https://reviews.llvm.org/D97174

llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.cpp

@@ -104,7 +104,26 @@ llvm::signatureFromMVTs(const SmallVectorImpl<MVT> &Results,
 
 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);
@@ -113,4 +132,9 @@ void yaml::WebAssemblyFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
 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/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h

@@ -19,6 +19,7 @@
 #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 {
@@ -30,6 +31,8 @@ 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;
@@ -63,9 +66,16 @@ class WebAssemblyFunctionInfo final : public MachineFunctionInfo {
   // 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); }
@@ -151,6 +161,9 @@ public:
 
   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,
@@ -172,8 +185,13 @@ signatureFromMVTs(const SmallVectorImpl<MVT> &Results,
 
 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);
@@ -185,6 +203,20 @@ struct WebAssemblyFunctionInfo final : public yaml::MachineFunctionInfo {
 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);
   }
 };
 

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

@@ -1,6 +1,20 @@
-# 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
@@ -10,3 +24,34 @@ 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
+...