[WebAssembly] Allow multivalue types in block signature operands

Summary:
Renames `ExprType` to the more apt `BlockType` and adds a variant for
multivalue blocks. Currently non-void blocks are only generated at the
end of functions where the block return type needs to agree with the
function return type, and that remains true for multivalue
blocks. That invariant means that the actual signature does not need
to be stored in the block signature `MachineOperand` because it can be
inferred by `WebAssemblyMCInstLower` from the return type of the
parent function. `WebAssemblyMCInstLower` continues to lower block
signature operands to immediates when possible but lowers multivalue
signatures to function type symbols. The AsmParser and Disassembler
are updated to handle multivalue block types as well.

Reviewers: aheejin, dschuff, aardappel

Subscribers: sbc100, jgravelle-google, hiraditya, sunfish, llvm-commits

Tags: #llvm

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

llvm-svn: 374933

Author: tlively

llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp

@@ -313,16 +313,17 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     return Optional<wasm::ValType>();
   }
 
-  WebAssembly::ExprType parseBlockType(StringRef ID) {
-    return StringSwitch<WebAssembly::ExprType>(ID)
-        .Case("i32", WebAssembly::ExprType::I32)
-        .Case("i64", WebAssembly::ExprType::I64)
-        .Case("f32", WebAssembly::ExprType::F32)
-        .Case("f64", WebAssembly::ExprType::F64)
-        .Case("v128", WebAssembly::ExprType::V128)
-        .Case("exnref", WebAssembly::ExprType::Exnref)
-        .Case("void", WebAssembly::ExprType::Void)
-        .Default(WebAssembly::ExprType::Invalid);
+  WebAssembly::BlockType parseBlockType(StringRef ID) {
+    // Multivalue block types are handled separately in parseSignature
+    return StringSwitch<WebAssembly::BlockType>(ID)
+        .Case("i32", WebAssembly::BlockType::I32)
+        .Case("i64", WebAssembly::BlockType::I64)
+        .Case("f32", WebAssembly::BlockType::F32)
+        .Case("f64", WebAssembly::BlockType::F64)
+        .Case("v128", WebAssembly::BlockType::V128)
+        .Case("exnref", WebAssembly::BlockType::Exnref)
+        .Case("void", WebAssembly::BlockType::Void)
+        .Default(WebAssembly::BlockType::Invalid);
   }
 
   bool parseRegTypeList(SmallVectorImpl<wasm::ValType> &Types) {
@@ -416,7 +417,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
   }
 
   void addBlockTypeOperand(OperandVector &Operands, SMLoc NameLoc,
-                           WebAssembly::ExprType BT) {
+                           WebAssembly::BlockType BT) {
     Operands.push_back(std::make_unique<WebAssemblyOperand>(
         WebAssemblyOperand::Integer, NameLoc, NameLoc,
         WebAssemblyOperand::IntOp{static_cast<int64_t>(BT)}));
@@ -456,6 +457,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     // If this instruction is part of a control flow structure, ensure
     // proper nesting.
     bool ExpectBlockType = false;
+    bool ExpectFuncType = false;
     if (Name == "block") {
       push(Block);
       ExpectBlockType = true;
@@ -494,6 +496,10 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
       if (pop(Name, Function) || ensureEmptyNestingStack())
         return true;
     } else if (Name == "call_indirect" || Name == "return_call_indirect") {
+      ExpectFuncType = true;
+    }
+
+    if (ExpectFuncType || (ExpectBlockType && Lexer.is(AsmToken::LParen))) {
       // This has a special TYPEINDEX operand which in text we
       // represent as a signature, such that we can re-build this signature,
       // attach it to an anonymous symbol, which is what WasmObjectWriter
@@ -502,6 +508,8 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
       auto Signature = std::make_unique<wasm::WasmSignature>();
       if (parseSignature(Signature.get()))
         return true;
+      // Got signature as block type, don't need more
+      ExpectBlockType = false;
       auto &Ctx = getStreamer().getContext();
       // The "true" here will cause this to be a nameless symbol.
       MCSymbol *Sym = Ctx.createTempSymbol("typeindex", true);
@@ -526,7 +534,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
         if (ExpectBlockType) {
           // Assume this identifier is a block_type.
           auto BT = parseBlockType(Id.getString());
-          if (BT == WebAssembly::ExprType::Invalid)
+          if (BT == WebAssembly::BlockType::Invalid)
             return error("Unknown block type: ", Id);
           addBlockTypeOperand(Operands, NameLoc, BT);
           Parser.Lex();
@@ -594,7 +602,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     }
     if (ExpectBlockType && Operands.size() == 1) {
       // Support blocks with no operands as default to void.
-      addBlockTypeOperand(Operands, NameLoc, WebAssembly::ExprType::Void);
+      addBlockTypeOperand(Operands, NameLoc, WebAssembly::BlockType::Void);
     }
     Parser.Lex();
     return false;

llvm/lib/Target/WebAssembly/Disassembler/LLVMBuild.txt

@@ -18,5 +18,5 @@
 type = Library
 name = WebAssemblyDisassembler
 parent = WebAssembly
-required_libraries = WebAssemblyDesc MCDisassembler WebAssemblyInfo Support
+required_libraries = WebAssemblyDesc MCDisassembler WebAssemblyInfo Support MC
 add_to_library_groups = WebAssembly

llvm/lib/Target/WebAssembly/Disassembler/WebAssemblyDisassembler.cpp

@@ -24,6 +24,7 @@
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCSymbolWasm.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/TargetRegistry.h"
@@ -213,10 +214,29 @@ MCDisassembler::DecodeStatus WebAssemblyDisassembler::getInstruction(
         return MCDisassembler::Fail;
       break;
     }
-    // block_type operands (uint8_t).
+    // block_type operands:
     case WebAssembly::OPERAND_SIGNATURE: {
-      if (!parseImmediate<uint8_t>(MI, Size, Bytes))
+      int64_t Val;
+      uint64_t PrevSize = Size;
+      if (!nextLEB(Val, Bytes, Size, true))
         return MCDisassembler::Fail;
+      if (Val < 0) {
+        // Negative values are single septet value types or empty types
+        if (Size != PrevSize + 1) {
+          MI.addOperand(
+              MCOperand::createImm(int64_t(WebAssembly::BlockType::Invalid)));
+        } else {
+          MI.addOperand(MCOperand::createImm(Val & 0x7f));
+        }
+      } else {
+        // We don't have access to the signature, so create a symbol without one
+        MCSymbol *Sym = getContext().createTempSymbol("typeindex", true);
+        auto *WasmSym = cast<MCSymbolWasm>(Sym);
+        WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
+        const MCExpr *Expr = MCSymbolRefExpr::create(
+            WasmSym, MCSymbolRefExpr::VK_WASM_TYPEINDEX, getContext());
+        MI.addOperand(MCOperand::createExpr(Expr));
+      }
       break;
     }
     // FP operands.

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp

@@ -272,9 +272,21 @@ void WebAssemblyInstPrinter::printWebAssemblyP2AlignOperand(const MCInst *MI,
 void WebAssemblyInstPrinter::printWebAssemblySignatureOperand(const MCInst *MI,
                                                               unsigned OpNo,
                                                               raw_ostream &O) {
-  auto Imm = static_cast<unsigned>(MI->getOperand(OpNo).getImm());
-  if (Imm != wasm::WASM_TYPE_NORESULT)
-    O << WebAssembly::anyTypeToString(Imm);
+  const MCOperand &Op = MI->getOperand(OpNo);
+  if (Op.isImm()) {
+    auto Imm = static_cast<unsigned>(Op.getImm());
+    if (Imm != wasm::WASM_TYPE_NORESULT)
+      O << WebAssembly::anyTypeToString(Imm);
+  } else {
+    auto Expr = cast<MCSymbolRefExpr>(Op.getExpr());
+    auto *Sym = cast<MCSymbolWasm>(&Expr->getSymbol());
+    if (Sym->getSignature()) {
+      O << WebAssembly::signatureToString(Sym->getSignature());
+    } else {
+      // Disassembler does not currently produce a signature
+      O << "unknown_type";
+    }
+  }
 }
 
 // We have various enums representing a subset of these types, use this

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCCodeEmitter.cpp

@@ -152,6 +152,7 @@ void WebAssemblyMCCodeEmitter::encodeInstruction(
         break;
       case WebAssembly::OPERAND_FUNCTION32:
       case WebAssembly::OPERAND_OFFSET32:
+      case WebAssembly::OPERAND_SIGNATURE:
       case WebAssembly::OPERAND_TYPEINDEX:
       case WebAssembly::OPERAND_GLOBAL:
       case WebAssembly::OPERAND_EVENT:

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyMCTargetDesc.h

@@ -122,16 +122,22 @@ enum TOF {
 namespace llvm {
 namespace WebAssembly {
 
-/// This is used to indicate block signatures.
-enum class ExprType : unsigned {
+/// Used as immediate MachineOperands for block signatures
+enum class BlockType : unsigned {
+  Invalid = 0x00,
   Void = 0x40,
-  I32 = 0x7F,
-  I64 = 0x7E,
-  F32 = 0x7D,
-  F64 = 0x7C,
-  V128 = 0x7B,
-  Exnref = 0x68,
-  Invalid = 0x00
+  I32 = unsigned(wasm::ValType::I32),
+  I64 = unsigned(wasm::ValType::I64),
+  F32 = unsigned(wasm::ValType::F32),
+  F64 = unsigned(wasm::ValType::F64),
+  V128 = unsigned(wasm::ValType::V128),
+  Exnref = unsigned(wasm::ValType::EXNREF),
+  // Multivalue blocks (and other non-void blocks) are only emitted when the
+  // blocks will never be exited and are at the ends of functions (see
+  // WebAssemblyCFGStackify::fixEndsAtEndOfFunction). They also are never made
+  // to pop values off the stack, so the exact multivalue signature can always
+  // be inferred from the return type of the parent function in MCInstLower.
+  Multivalue = 0xffff,
 };
 
 /// Instruction opcodes emitted via means other than CodeGen.

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

@@ -315,12 +315,12 @@ void WebAssemblyCFGStackify::placeBlockMarker(MachineBasicBlock &MBB) {
   //   br_on_exn 0, $__cpp_exception
   //   rethrow
   // end_block
-  WebAssembly::ExprType ReturnType = WebAssembly::ExprType::Void;
+  WebAssembly::BlockType ReturnType = WebAssembly::BlockType::Void;
   if (IsBrOnExn) {
     const char *TagName = BrOnExn->getOperand(1).getSymbolName();
     if (std::strcmp(TagName, "__cpp_exception") != 0)
       llvm_unreachable("Only C++ exception is supported");
-    ReturnType = WebAssembly::ExprType::I32;
+    ReturnType = WebAssembly::BlockType::I32;
   }
 
   auto InsertPos = getLatestInsertPos(Header, BeforeSet, AfterSet);
@@ -406,7 +406,7 @@ void WebAssemblyCFGStackify::placeLoopMarker(MachineBasicBlock &MBB) {
   auto InsertPos = getEarliestInsertPos(&MBB, BeforeSet, AfterSet);
   MachineInstr *Begin = BuildMI(MBB, InsertPos, MBB.findDebugLoc(InsertPos),
                                 TII.get(WebAssembly::LOOP))
-                            .addImm(int64_t(WebAssembly::ExprType::Void));
+                            .addImm(int64_t(WebAssembly::BlockType::Void));
 
   // Decide where in Header to put the END_LOOP.
   BeforeSet.clear();
@@ -575,7 +575,7 @@ void WebAssemblyCFGStackify::placeTryMarker(MachineBasicBlock &MBB) {
   MachineInstr *Begin =
       BuildMI(*Header, InsertPos, Header->findDebugLoc(InsertPos),
               TII.get(WebAssembly::TRY))
-          .addImm(int64_t(WebAssembly::ExprType::Void));
+          .addImm(int64_t(WebAssembly::BlockType::Void));
 
   // Decide where in Header to put the END_TRY.
   BeforeSet.clear();
@@ -1129,7 +1129,7 @@ bool WebAssemblyCFGStackify::fixUnwindMismatches(MachineFunction &MF) {
       MachineInstr *NestedTry =
           BuildMI(*MBB, *RangeBegin, RangeBegin->getDebugLoc(),
                   TII.get(WebAssembly::TRY))
-              .addImm(int64_t(WebAssembly::ExprType::Void));
+              .addImm(int64_t(WebAssembly::BlockType::Void));
 
       // Create the nested EH pad and fill instructions in.
       MachineBasicBlock *NestedEHPad = MF.CreateMachineBasicBlock();
@@ -1231,54 +1231,28 @@ void WebAssemblyCFGStackify::fixEndsAtEndOfFunction(MachineFunction &MF) {
   if (MFI.getResults().empty())
     return;
 
-  // TODO: Generalize from value types to function types for multivalue
-  WebAssembly::ExprType RetType;
-  switch (MFI.getResults().front().SimpleTy) {
-  case MVT::i32:
-    RetType = WebAssembly::ExprType::I32;
-    break;
-  case MVT::i64:
-    RetType = WebAssembly::ExprType::I64;
-    break;
-  case MVT::f32:
-    RetType = WebAssembly::ExprType::F32;
-    break;
-  case MVT::f64:
-    RetType = WebAssembly::ExprType::F64;
-    break;
-  case MVT::v16i8:
-  case MVT::v8i16:
-  case MVT::v4i32:
-  case MVT::v2i64:
-  case MVT::v4f32:
-  case MVT::v2f64:
-    RetType = WebAssembly::ExprType::V128;
-    break;
-  case MVT::exnref:
-    RetType = WebAssembly::ExprType::Exnref;
-    break;
-  default:
-    llvm_unreachable("unexpected return type");
-  }
+  // MCInstLower will add the proper types to multivalue signatures based on the
+  // function return type
+  WebAssembly::BlockType RetType =
+      MFI.getResults().size() > 1
+          ? WebAssembly::BlockType::Multivalue
+          : WebAssembly::BlockType(
+                WebAssembly::toValType(MFI.getResults().front()));
 
   for (MachineBasicBlock &MBB : reverse(MF)) {
     for (MachineInstr &MI : reverse(MBB)) {
       if (MI.isPosition() || MI.isDebugInstr())
         continue;
-      if (MI.getOpcode() == WebAssembly::END_BLOCK) {
-        if (MFI.getResults().size() > 1)
-          report_fatal_error("Multivalue block signatures not implemented yet");
-        EndToBegin[&MI]->getOperand(0).setImm(int32_t(RetType));
-        continue;
-      }
-      if (MI.getOpcode() == WebAssembly::END_LOOP) {
-        if (MFI.getResults().size() > 1)
-          report_fatal_error("Multivalue loop signatures not implemented yet");
+      switch (MI.getOpcode()) {
+      case WebAssembly::END_BLOCK:
+      case WebAssembly::END_LOOP:
+      case WebAssembly::END_TRY:
         EndToBegin[&MI]->getOperand(0).setImm(int32_t(RetType));
         continue;
+      default:
+        // Something other than an `end`. We're done.
+        return;
       }
-      // Something other than an `end`. We're done.
-      return;
     }
   }
 }

llvm/lib/Target/WebAssembly/WebAssemblyMCInstLower.cpp

@@ -163,6 +163,21 @@ MCOperand WebAssemblyMCInstLower::lowerSymbolOperand(const MachineOperand &MO,
   return MCOperand::createExpr(Expr);
 }
 
+MCOperand WebAssemblyMCInstLower::lowerTypeIndexOperand(
+    SmallVector<wasm::ValType, 1> &&Returns,
+    SmallVector<wasm::ValType, 4> &&Params) const {
+  auto Signature = std::make_unique<wasm::WasmSignature>(std::move(Returns),
+                                                         std::move(Params));
+  MCSymbol *Sym = Printer.createTempSymbol("typeindex");
+  auto *WasmSym = cast<MCSymbolWasm>(Sym);
+  WasmSym->setSignature(Signature.get());
+  Printer.addSignature(std::move(Signature));
+  WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
+  const MCExpr *Expr =
+      MCSymbolRefExpr::create(WasmSym, MCSymbolRefExpr::VK_WASM_TYPEINDEX, Ctx);
+  return MCOperand::createExpr(Expr);
+}
+
 // Return the WebAssembly type associated with the given register class.
 static wasm::ValType getType(const TargetRegisterClass *RC) {
   if (RC == &WebAssembly::I32RegClass)
@@ -178,6 +193,16 @@ static wasm::ValType getType(const TargetRegisterClass *RC) {
   llvm_unreachable("Unexpected register class");
 }
 
+static void getFunctionReturns(const MachineInstr *MI,
+                               SmallVectorImpl<wasm::ValType> &Returns) {
+  const Function &F = MI->getMF()->getFunction();
+  const TargetMachine &TM = MI->getMF()->getTarget();
+  Type *RetTy = F.getReturnType();
+  SmallVector<MVT, 4> CallerRetTys;
+  computeLegalValueVTs(F, TM, RetTy, CallerRetTys);
+  valTypesFromMVTs(CallerRetTys, Returns);
+}
+
 void WebAssemblyMCInstLower::lower(const MachineInstr *MI,
                                    MCInst &OutMI) const {
   OutMI.setOpcode(MI->getOpcode());
@@ -208,8 +233,6 @@ void WebAssemblyMCInstLower::lower(const MachineInstr *MI,
       if (I < Desc.NumOperands) {
         const MCOperandInfo &Info = Desc.OpInfo[I];
         if (Info.OperandType == WebAssembly::OPERAND_TYPEINDEX) {
-          MCSymbol *Sym = Printer.createTempSymbol("typeindex");
-
           SmallVector<wasm::ValType, 4> Returns;
           SmallVector<wasm::ValType, 4> Params;
 
@@ -228,26 +251,21 @@ void WebAssemblyMCInstLower::lower(const MachineInstr *MI,
 
           // return_call_indirect instructions have the return type of the
           // caller
-          if (MI->getOpcode() == WebAssembly::RET_CALL_INDIRECT) {
-            const Function &F = MI->getMF()->getFunction();
-            const TargetMachine &TM = MI->getMF()->getTarget();
-            Type *RetTy = F.getReturnType();
-            SmallVector<MVT, 4> CallerRetTys;
-            computeLegalValueVTs(F, TM, RetTy, CallerRetTys);
-            valTypesFromMVTs(CallerRetTys, Returns);
-          }
+          if (MI->getOpcode() == WebAssembly::RET_CALL_INDIRECT)
+            getFunctionReturns(MI, Returns);
 
-          auto *WasmSym = cast<MCSymbolWasm>(Sym);
-          auto Signature = std::make_unique<wasm::WasmSignature>(std::move(Returns),
-                                                            std::move(Params));
-          WasmSym->setSignature(Signature.get());
-          Printer.addSignature(std::move(Signature));
-          WasmSym->setType(wasm::WASM_SYMBOL_TYPE_FUNCTION);
-
-          const MCExpr *Expr = MCSymbolRefExpr::create(
-              WasmSym, MCSymbolRefExpr::VK_WASM_TYPEINDEX, Ctx);
-          MCOp = MCOperand::createExpr(Expr);
+          MCOp = lowerTypeIndexOperand(std::move(Returns), std::move(Params));
           break;
+        } else if (Info.OperandType == WebAssembly::OPERAND_SIGNATURE) {
+          auto BT = static_cast<WebAssembly::BlockType>(MO.getImm());
+          assert(BT != WebAssembly::BlockType::Invalid);
+          if (BT == WebAssembly::BlockType::Multivalue) {
+            SmallVector<wasm::ValType, 1> Returns;
+            getFunctionReturns(MI, Returns);
+            MCOp = lowerTypeIndexOperand(std::move(Returns),
+                                         SmallVector<wasm::ValType, 4>());
+            break;
+          }
         }
       }
       MCOp = MCOperand::createImm(MO.getImm());