[WebAssembly] Support select and block for reference types

[aheejin]

Summary

This adds missing select instruction support and block return type
support for reference types. Also refactors WebAssemblyInstrRef.td and
rearranges tests in reference-types.s. Tests don't include exnref
types, because we currently don't support exnref for ref.null and
the type will be removed soon anyway.

[aheejin]

Add missing reference types in utility functions

[aheejin]

Fix a typo

[sbc100]

llvm/test/MC/WebAssembly/reference-types.s

8

Isn't it useful to continue to test externref and funcref as function return types? Or is that tested elsewhere?

[aheejin]

llvm/test/MC/WebAssembly/reference-types.s

8

llvm-mc cannot check the correctness of that. We can put any signature here and the test will pass anyway.

[sbc100]

llvm/test/MC/WebAssembly/reference-types.s

8

But it must be at least checked for valid type?

You can write (foo) -> (bar) can you? It might be good to check that these types are accepted in this position?

[aheejin]

Add tests to test reference types in params and results

[aheejin]

This comment was removed by aheejin.

[aheejin]

Fix signature

[aheejin]

Add missing COPY opcode handling in Peephole

[aheejin]

llvm/test/MC/WebAssembly/reference-types.s

8

I added two more tests to test reference type signatures.

[wingo]

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

343

See also https://reviews.llvm.org/D92323, which contains this. Happy for this to land though!

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

#include "WebAssemblyGenAsmMatcher.inc"
   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("funcref", WebAssembly::BlockType::Funcref)
+        .Case("externref", WebAssembly::BlockType::Externref)
         .Case("exnref", WebAssembly::BlockType::Exnref)
         .Case("void", WebAssembly::BlockType::Void)
         .Default(WebAssembly::BlockType::Invalid);
   }
 
   bool parseRegTypeList(SmallVectorImpl<wasm::ValType> &Types) {
     while (Lexer.is(AsmToken::Identifier)) {
       auto Type = parseType(Lexer.getTok().getString());

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

 enum class BlockType : unsigned {
   Invalid = 0x00,
   Void = 0x40,
   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),
+  Externref = unsigned(wasm::ValType::EXTERNREF),
+  Funcref = unsigned(wasm::ValType::FUNCREF),
   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,
 };
inline bool isArgument(unsigned Opc) {
   case WebAssembly::ARGUMENT_v4i32:
   case WebAssembly::ARGUMENT_v4i32_S:
   case WebAssembly::ARGUMENT_v2i64:
   case WebAssembly::ARGUMENT_v2i64_S:
   case WebAssembly::ARGUMENT_v4f32:
   case WebAssembly::ARGUMENT_v4f32_S:
   case WebAssembly::ARGUMENT_v2f64:
   case WebAssembly::ARGUMENT_v2f64_S:
+  case WebAssembly::ARGUMENT_funcref:
+  case WebAssembly::ARGUMENT_funcref_S:
+  case WebAssembly::ARGUMENT_externref:
+  case WebAssembly::ARGUMENT_externref_S:
   case WebAssembly::ARGUMENT_exnref:
   case WebAssembly::ARGUMENT_exnref_S:
     return true;
   default:
     return false;
   }
 }
 
 inline bool isCopy(unsigned Opc) {
   switch (Opc) {
   case WebAssembly::COPY_I32:
   case WebAssembly::COPY_I32_S:
   case WebAssembly::COPY_I64:
   case WebAssembly::COPY_I64_S:
   case WebAssembly::COPY_F32:
   case WebAssembly::COPY_F32_S:
   case WebAssembly::COPY_F64:
   case WebAssembly::COPY_F64_S:
   case WebAssembly::COPY_V128:
   case WebAssembly::COPY_V128_S:
+  case WebAssembly::COPY_FUNCREF:
+  case WebAssembly::COPY_FUNCREF_S:
+  case WebAssembly::COPY_EXTERNREF:
+  case WebAssembly::COPY_EXTERNREF_S:
   case WebAssembly::COPY_EXNREF:
   case WebAssembly::COPY_EXNREF_S:
     return true;
   default:
     return false;
   }
 }
 
 inline bool isTee(unsigned Opc) {
   switch (Opc) {
   case WebAssembly::TEE_I32:
   case WebAssembly::TEE_I32_S:
   case WebAssembly::TEE_I64:
   case WebAssembly::TEE_I64_S:
   case WebAssembly::TEE_F32:
   case WebAssembly::TEE_F32_S:
   case WebAssembly::TEE_F64:
   case WebAssembly::TEE_F64_S:
   case WebAssembly::TEE_V128:
   case WebAssembly::TEE_V128_S:
+  case WebAssembly::TEE_FUNCREF:
+  case WebAssembly::TEE_FUNCREF_S:
+  case WebAssembly::TEE_EXTERNREF:
+  case WebAssembly::TEE_EXTERNREF_S:
   case WebAssembly::TEE_EXNREF:
   case WebAssembly::TEE_EXNREF_S:
     return true;
   default:
     return false;
   }
 }
 

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

case MVT::f64:
     return wasm::ValType::F64;
   case MVT::v16i8:
   case MVT::v8i16:
   case MVT::v4i32:
   case MVT::v2i64:
   case MVT::v4f32:
   case MVT::v2f64:
     return wasm::ValType::V128;
+  case MVT::funcref:
+    return wasm::ValType::FUNCREF;
+  case MVT::externref:
+    return wasm::ValType::EXTERNREF;
   case MVT::exnref:
     return wasm::ValType::EXNREF;
   default:
     llvm_unreachable("unexpected type");
   }
 }

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

static unsigned getCopyOpcode(const TargetRegisterClass *RC) {
   if (RC == &WebAssembly::I64RegClass)
     return WebAssembly::COPY_I64;
   if (RC == &WebAssembly::F32RegClass)
     return WebAssembly::COPY_F32;
   if (RC == &WebAssembly::F64RegClass)
     return WebAssembly::COPY_F64;
   if (RC == &WebAssembly::V128RegClass)
     return WebAssembly::COPY_V128;
+  if (RC == &WebAssembly::FUNCREFRegClass)
+    return WebAssembly::COPY_FUNCREF;
+  if (RC == &WebAssembly::EXTERNREFRegClass)
+    return WebAssembly::COPY_EXTERNREF;
   if (RC == &WebAssembly::EXNREFRegClass)
     return WebAssembly::COPY_EXNREF;
   llvm_unreachable("Unexpected register class");
 }
 
 // When MBB is split into MBB and Split, we should unstackify defs in MBB that
 // have their uses in Split.
 static void unstackifyVRegsUsedInSplitBB(MachineBasicBlock &MBB,

llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp

static unsigned getDropOpcode(const TargetRegisterClass *RC) {
   if (RC == &WebAssembly::I64RegClass)
     return WebAssembly::DROP_I64;
   if (RC == &WebAssembly::F32RegClass)
     return WebAssembly::DROP_F32;
   if (RC == &WebAssembly::F64RegClass)
     return WebAssembly::DROP_F64;
   if (RC == &WebAssembly::V128RegClass)
     return WebAssembly::DROP_V128;
+  if (RC == &WebAssembly::FUNCREFRegClass)
+    return WebAssembly::DROP_FUNCREF;
+  if (RC == &WebAssembly::EXTERNREFRegClass)
+    return WebAssembly::DROP_EXTERNREF;
   if (RC == &WebAssembly::EXNREFRegClass)
     return WebAssembly::DROP_EXNREF;
   llvm_unreachable("Unexpected register class");
 }
 
 /// Get the appropriate local.get opcode for the given register class.
 static unsigned getLocalGetOpcode(const TargetRegisterClass *RC) {
   if (RC == &WebAssembly::I32RegClass)

llvm/lib/Target/WebAssembly/WebAssemblyFastISel.cpp

for (auto const &Arg : F->args()) {
     case MVT::v4f32:
       Opc = WebAssembly::ARGUMENT_v4f32;
       RC = &WebAssembly::V128RegClass;
       break;
     case MVT::v2f64:
       Opc = WebAssembly::ARGUMENT_v2f64;
       RC = &WebAssembly::V128RegClass;
       break;
+    case MVT::funcref:
+      Opc = WebAssembly::ARGUMENT_funcref;
+      RC = &WebAssembly::FUNCREFRegClass;
+      break;
+    case MVT::externref:
+      Opc = WebAssembly::ARGUMENT_externref;
+      RC = &WebAssembly::EXTERNREFRegClass;
+      break;
     case MVT::exnref:
       Opc = WebAssembly::ARGUMENT_exnref;
       RC = &WebAssembly::EXNREFRegClass;
       break;
     default:
       return false;
     }
     unsigned ResultReg = createResultReg(RC);
bool WebAssemblyFastISel::selectSelect(const Instruction *I) {
   case MVT::f64:
     Opc = WebAssembly::SELECT_F64;
     RC = &WebAssembly::F64RegClass;
     break;
   case MVT::exnref:
     Opc = WebAssembly::SELECT_EXNREF;
     RC = &WebAssembly::EXNREFRegClass;
     break;
+  case MVT::funcref:
+    Opc = WebAssembly::SELECT_FUNCREF;
+    RC = &WebAssembly::FUNCREFRegClass;
+    break;
+  case MVT::externref:
+    Opc = WebAssembly::SELECT_EXTERNREF;
+    RC = &WebAssembly::EXTERNREFRegClass;
+    break;
   default:
     return false;
   }
 
   unsigned ResultReg = createResultReg(RC);
   BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc, TII.get(Opc), ResultReg)
       .addReg(TrueReg)
       .addReg(FalseReg)
bool WebAssemblyFastISel::selectRet(const Instruction *I) {
   case MVT::f32:
   case MVT::f64:
   case MVT::v16i8:
   case MVT::v8i16:
   case MVT::v4i32:
   case MVT::v2i64:
   case MVT::v4f32:
   case MVT::v2f64:
+  case MVT::funcref:
+  case MVT::externref:
   case MVT::exnref:
     break;
   default:
     return false;
   }
 
   unsigned Reg;
   if (FuncInfo.Fn->getAttributes().hasAttribute(0, Attribute::SExt))

llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.cpp

void WebAssemblyInstrInfo::copyPhysReg(MachineBasicBlock &MBB,
   else if (RC == &WebAssembly::I64RegClass)
     CopyOpcode = WebAssembly::COPY_I64;
   else if (RC == &WebAssembly::F32RegClass)
     CopyOpcode = WebAssembly::COPY_F32;
   else if (RC == &WebAssembly::F64RegClass)
     CopyOpcode = WebAssembly::COPY_F64;
   else if (RC == &WebAssembly::V128RegClass)
     CopyOpcode = WebAssembly::COPY_V128;
+  else if (RC == &WebAssembly::FUNCREFRegClass)
+    CopyOpcode = WebAssembly::COPY_FUNCREF;
+  else if (RC == &WebAssembly::EXTERNREFRegClass)
+    CopyOpcode = WebAssembly::COPY_EXTERNREF;
   else if (RC == &WebAssembly::EXNREFRegClass)
     CopyOpcode = WebAssembly::COPY_EXNREF;
   else
     llvm_unreachable("Unexpected register class");
 
   BuildMI(MBB, I, DL, get(CopyOpcode), DestReg)
       .addReg(SrcReg, KillSrc ? RegState::Kill : 0);
 }

llvm/lib/Target/WebAssembly/WebAssemblyInstrRef.td

 // WebAssemblyInstrRef.td - WebAssembly reference type codegen --*- tablegen -*-
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// WebAssembly reference type operand codegen constructs.
 ///
 //===----------------------------------------------------------------------===//
 
-defm SELECT_EXNREF : I<(outs EXNREF:$dst),
-                       (ins EXNREF:$lhs, EXNREF:$rhs, I32:$cond),
-                       (outs), (ins),
-                       [(set EXNREF:$dst,
-                         (select I32:$cond, EXNREF:$lhs, EXNREF:$rhs))],
-                       "exnref.select\t$dst, $lhs, $rhs, $cond",
-                       "exnref.select", 0x1b>;
-
-def : Pat<(select (i32 (setne I32:$cond, 0)), EXNREF:$lhs, EXNREF:$rhs),
-          (SELECT_EXNREF EXNREF:$lhs, EXNREF:$rhs, I32:$cond)>;
-def : Pat<(select (i32 (seteq I32:$cond, 0)), EXNREF:$lhs, EXNREF:$rhs),
-          (SELECT_EXNREF EXNREF:$rhs, EXNREF:$lhs, I32:$cond)>;
-
-multiclass REF<WebAssemblyRegClass rt> {
-  defm REF_NULL_#rt : I<(outs rt:$res), (ins HeapType:$heaptype),
+multiclass REF_I<WebAssemblyRegClass reg, ValueType vt> {
+  defm REF_NULL_#reg : I<(outs reg:$res), (ins HeapType:$heaptype),
                          (outs), (ins HeapType:$heaptype),
                          [],
                          "ref.null\t$res, $heaptype",
                          "ref.null\t$heaptype",
-                        0xd0>;
+                         0xd0>,
+                       Requires<[HasReferenceTypes]>;
+  defm SELECT_#reg: I<(outs reg:$dst), (ins reg:$lhs, reg:$rhs, I32:$cond),
+                      (outs), (ins),
+                      [(set reg:$dst,
+                        (select I32:$cond, reg:$lhs, reg:$rhs))],
+                      vt#".select\t$dst, $lhs, $rhs, $cond",
+                      vt#".select", 0x1b>,
+                    Requires<[HasReferenceTypes]>;
 }
 
-defm "" : REF<FUNCREF>, Requires<[HasReferenceTypes]>;
-defm "" : REF<EXTERNREF>, Requires<[HasReferenceTypes]>;
+defm "" : REF_I<FUNCREF, funcref>;
+defm "" : REF_I<EXTERNREF, externref>;
+defm "" : REF_I<EXNREF, exnref>;
+
+foreach reg = [FUNCREF, EXTERNREF, EXNREF] in {
+def : Pat<(select (i32 (setne I32:$cond, 0)), reg:$lhs, reg:$rhs),
+          (!cast<Instruction>("SELECT_"#reg) reg:$lhs, reg:$rhs, I32:$cond)>;
+def : Pat<(select (i32 (seteq I32:$cond, 0)), reg:$lhs, reg:$rhs),
+          (!cast<Instruction>("SELECT_"#reg) reg:$rhs, reg:$lhs, I32:$cond)>;
+}

llvm/lib/Target/WebAssembly/WebAssemblyPeephole.cpp

if (!MFI.isVRegStackified(Reg)) {
         CopyLocalOpc = WebAssembly::COPY_F32;
         break;
       case WebAssembly::F64RegClassID:
         CopyLocalOpc = WebAssembly::COPY_F64;
         break;
       case WebAssembly::V128RegClassID:
         CopyLocalOpc = WebAssembly::COPY_V128;
         break;
+      case WebAssembly::FUNCREFRegClassID:
+        CopyLocalOpc = WebAssembly::COPY_FUNCREF;
+        break;
+      case WebAssembly::EXTERNREFRegClassID:
+        CopyLocalOpc = WebAssembly::COPY_EXTERNREF;
+        break;
       case WebAssembly::EXNREFRegClassID:
         CopyLocalOpc = WebAssembly::COPY_EXNREF;
         break;
       default:
         llvm_unreachable("Unexpected register class for return operand");
       }
       Register NewReg = MRI.createVirtualRegister(RegClass);
       BuildMI(MBB, MI, MI.getDebugLoc(), TII.get(CopyLocalOpc), NewReg)

llvm/test/MC/WebAssembly/reference-types.s

 # RUN: llvm-mc -show-encoding -triple=wasm32-unknown-unknown -mattr=+reference-types < %s | FileCheck %s
 # RUN: llvm-mc -show-encoding -triple=wasm64-unknown-unknown -mattr=+reference-types < %s | FileCheck %s
 
-#      CHECK: ref_null_externref:
-# CHECK-NEXT:         .functype ref_null_externref () -> (externref)
-#      CHECK: ref.null extern # encoding: [0xd0,0x6f]
-# CHECK-NEXT: end_function
-ref_null_externref:
-  .functype ref_null_externref () -> (externref)
+# CHECK-LABEL: ref_null_test:
+# CHECK: ref.null func   # encoding: [0xd0,0x70]
+# CHECK: ref.null extern # encoding: [0xd0,0x6f]
+ref_null_test:
+  .functype ref_null_test () -> ()
+  ref.null func
+  drop
   ref.null extern
+  drop
   end_function
 
-#      CHECK: ref_null_funcref:
-# CHECK-NEXT:         .functype ref_null_funcref () -> (funcref)
-#      CHECK: ref.null func # encoding: [0xd0,0x70]
-# CHECK-NEXT: end_function
-ref_null_funcref:
-  .functype ref_null_funcref () -> (funcref)
+# CHECK-LABEL: ref_sig_test_funcref:
+# CHECK-NEXT: .functype ref_sig_test_funcref (funcref) -> (funcref)
+ref_sig_test_funcref:
+  .functype ref_sig_test_funcref (funcref) -> (funcref)
+  end_function
+
+# CHECK-LABEL: ref_sig_test_externref:
+# CHECK-NEXT: .functype ref_sig_test_externref (externref) -> (externref)
+ref_sig_test_externref:
+  .functype ref_sig_test_externref (externref) -> (externref)
+  end_function
+
+# CHECK-LABEL: ref_select_test:
+# CHECK: funcref.select   # encoding: [0x1b]
+# CHECK: externref.select # encoding: [0x1b]
+ref_select_test:
+  .functype ref_select_test () -> ()
+  ref.null func
   ref.null func
+  i32.const 0
+  funcref.select
+  drop
+  ref.null extern
+  ref.null extern
+  i32.const 0
+  externref.select
+  drop
+  end_function
+
+# CHECK-LABEL: ref_block_test:
+# CHECK: block funcref
+# CHECK: block externref
+ref_block_test:
+  .functype ref_block_test () -> ()
+  block funcref
+  block externref
+  end_block
+  end_block
   end_function