[WebAssembly] Implement addrspacecast to funcref (llvm#166820)

Adds lowering of `addrspacecast [0 -> 20]` to allow easy conversion of
function pointers to Wasm `funcref`

When given a constant function pointer, it lowers to a direct
`ref.func`. Otherwise it lowers to a `table.get` from
`__indirect_function_table` using the provided pointer as the index.

Author: QuantumSegfault

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.cpp

@@ -411,6 +411,10 @@ WebAssemblyTargetLowering::WebAssemblyTargetLowering(
   setOperationAction(ISD::INTRINSIC_W_CHAIN, MVT::Other, Custom);
   setOperationAction(ISD::INTRINSIC_VOID, MVT::Other, Custom);
 
+  // Allow converting function ptrs in address space 0 to Wasm funcref (address
+  // space 20)
+  setOperationAction(ISD::ADDRSPACECAST, MVT::funcref, Custom);
+
   setMaxAtomicSizeInBitsSupported(64);
 
   // Always convert switches to br_tables unless there is only one case, which
@@ -1756,6 +1760,8 @@ SDValue WebAssemblyTargetLowering::LowerOperation(SDValue Op,
     return LowerMUL_LOHI(Op, DAG);
   case ISD::UADDO:
     return LowerUADDO(Op, DAG);
+  case ISD::ADDRSPACECAST:
+    return LowerADDRSPACECAST(Op, DAG);
   }
 }
 
@@ -1899,6 +1905,58 @@ SDValue WebAssemblyTargetLowering::LowerUADDO(SDValue Op,
   return DAG.getMergeValues(Ops, DL);
 }
 
+SDValue WebAssemblyTargetLowering::LowerADDRSPACECAST(SDValue Op,
+                                                      SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+
+  AddrSpaceCastSDNode *ACN = cast<AddrSpaceCastSDNode>(Op.getNode());
+
+  if (ACN->getSrcAddressSpace() !=
+          WebAssembly::WasmAddressSpace::WASM_ADDRESS_SPACE_DEFAULT ||
+      ACN->getDestAddressSpace() !=
+          WebAssembly::WasmAddressSpace::WASM_ADDRESS_SPACE_FUNCREF)
+    return SDValue();
+
+  if (ACN->getValueType(0) != MVT::funcref) {
+    reportFatalInternalError("Cannot addrspacecast to funcref addrspace with "
+                             "results other than MVT::funcref");
+  }
+
+  SDValue Src = ACN->getOperand(0);
+
+  // Lower addrspacecasts of direct/constant function ptrs to ref.func
+  if (auto *GA = dyn_cast<GlobalAddressSDNode>(
+          Src->getOpcode() == WebAssemblyISD::Wrapper ? Src->getOperand(0)
+                                                      : Src)) {
+    auto *GV = GA->getGlobal();
+
+    if (const Function *F = dyn_cast<Function>(GV)) {
+      SDValue FnAddress = DAG.getTargetGlobalAddress(F, DL, MVT::i32);
+
+      SDValue RefFuncNode =
+          DAG.getNode(WebAssemblyISD::REF_FUNC, DL, MVT::funcref, FnAddress);
+      return RefFuncNode;
+    }
+  }
+
+  // Lower everything else to a table.get from the indirect function table
+  const MachineFunction &MF = DAG.getMachineFunction();
+
+  MVT PtrVT = getPointerTy(MF.getDataLayout());
+
+  MCSymbolWasm *Table =
+      WebAssembly::getOrCreateFunctionTableSymbol(MF.getContext(), Subtarget);
+  SDValue TableSym = DAG.getMCSymbol(Table, PtrVT);
+
+  SDValue TableSlot = Op.getOperand(0);
+
+  SDValue Result(DAG.getMachineNode(WebAssembly::TABLE_GET_FUNCREF, DL,
+                                    MVT::funcref, TableSym, TableSlot),
+                 0);
+
+  return Result;
+}
+
 SDValue WebAssemblyTargetLowering::Replace128Op(SDNode *N,
                                                 SelectionDAG &DAG) const {
   assert(Subtarget->hasWideArithmetic());

llvm/lib/Target/WebAssembly/WebAssemblyISelLowering.h

@@ -121,6 +121,7 @@ class WebAssemblyTargetLowering final : public TargetLowering {
   SDValue LowerMUL_LOHI(SDValue Op, SelectionDAG &DAG) const;
   SDValue Replace128Op(SDNode *N, SelectionDAG &DAG) const;
   SDValue LowerUADDO(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerADDRSPACECAST(SDValue Op, SelectionDAG &DAG) const;
 
   // Custom DAG combine hooks
   SDValue

llvm/lib/Target/WebAssembly/WebAssemblyInstrRef.td

@@ -11,6 +11,11 @@
 ///
 //===----------------------------------------------------------------------===//
 
+def WebAssemblyRefFunc_t : SDTypeProfile<1, 1, [SDTCisVT<0, funcref>, SDTCisPtrTy<1>]>;
+def WebAssemblyRefFunc :
+    SDNode<"WebAssemblyISD::REF_FUNC", WebAssemblyRefFunc_t,
+           []>;
+
 multiclass REF_I<WebAssemblyRegClass rc, ValueType vt, string ht> {
   defm REF_NULL_#rc : I<(outs rc:$dst), (ins),
                         (outs), (ins),
@@ -42,7 +47,7 @@ defm REF_TEST_FUNCREF : I<(outs I32:$res), (ins TypeIndex:$type, FUNCREF:$ref),
                         Requires<[HasGC]>;
 
 defm REF_FUNC : I<(outs FUNCREF:$res), (ins function32_op:$func),
-                    (outs), (ins function32_op:$func), [],
+                    (outs), (ins function32_op:$func), [(set FUNCREF:$res, (WebAssemblyRefFunc tglobaladdr:$func))],
                     "ref.func\t$func", "ref.func $func", 0xd2>,
                 Requires<[HasReferenceTypes]>;
 

llvm/test/CodeGen/WebAssembly/addrspacecast-funcref.ll

@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=wasm32-unknown-unknown -mattr=+reference-types < %s | FileCheck -check-prefixes=CHECK,WASM32 %s
+; RUN: llc -mtriple=wasm64-unknown-unknown -mattr=+reference-types < %s | FileCheck -check-prefixes=CHECK,WASM64 %s
+
+%funcref = type ptr addrspace(20) ;; addrspace 20 is nonintegral
+
+declare void @foo();
+
+@global_var = local_unnamed_addr global i32 0
+
+define %funcref @cast_const_funcptr() {
+; CHECK-LABEL: cast_const_funcptr:
+; CHECK:         .functype cast_const_funcptr () -> (funcref)
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    ref.func foo
+; CHECK-NEXT:    # fallthrough-return
+  %result = addrspacecast ptr @foo to %funcref
+  ret %funcref %result
+}
+
+define %funcref @cast_const_not_funcptr() {
+; WASM32-LABEL: cast_const_not_funcptr:
+; WASM32:         .functype cast_const_not_funcptr () -> (funcref)
+; WASM32-NEXT:  # %bb.0:
+; WASM32-NEXT:    i32.const global_var
+; WASM32-NEXT:    table.get __indirect_function_table
+; WASM32-NEXT:    # fallthrough-return
+;
+; WASM64-LABEL: cast_const_not_funcptr:
+; WASM64:         .functype cast_const_not_funcptr () -> (funcref)
+; WASM64-NEXT:  # %bb.0:
+; WASM64-NEXT:    i64.const global_var
+; WASM64-NEXT:    table.get __indirect_function_table
+; WASM64-NEXT:    # fallthrough-return
+  %result = addrspacecast ptr @global_var to %funcref
+  ret %funcref %result
+}
+
+define %funcref @cast_param_funcptr(ptr %funcptr) {
+; WASM32-LABEL: cast_param_funcptr:
+; WASM32:         .functype cast_param_funcptr (i32) -> (funcref)
+; WASM32-NEXT:  # %bb.0:
+; WASM32-NEXT:    local.get 0
+; WASM32-NEXT:    table.get __indirect_function_table
+; WASM32-NEXT:    # fallthrough-return
+;
+; WASM64-LABEL: cast_param_funcptr:
+; WASM64:         .functype cast_param_funcptr (i64) -> (funcref)
+; WASM64-NEXT:  # %bb.0:
+; WASM64-NEXT:    local.get 0
+; WASM64-NEXT:    table.get __indirect_function_table
+; WASM64-NEXT:    # fallthrough-return
+  %result = addrspacecast ptr %funcptr to %funcref
+  ret %funcref %result
+}