Bug 1720514 - wasm: Add framework for intrinsics with basic i8vecmul.…

… r=jseward

This commit implements a framework for creating 'intrinsics' which are
natively implemented functions that are exposable as wasm functions
that can be called or linked against. A simple 8-bit vector product for
flexible vectors is implemented as a proof of concept.

The basic API is:
```
let module = wasmIntrinsicI8VecMul(); // WebAssembly.Module
let memory = new WebAssembly.Module({ initial: pageSize });
let instance = new WebAssembly.Instance(module,
  { "": { memory } });
instance.exports.i8vecmul(dest, src1, src2, len);
```

The implementation is mainly done through `CompileIntrisicModule` which
manually builds a ModuleEnvironment with an imported memory, and a single
exported function which is of the bytecode form:
```
(func (params ...)
  local.get 0
  ...
  local.get n
  private_intrinsic_opcode
)
```
The private_intrinsic_opcode is implemented as an instance call. An
additional heap base parameter is added which allows quick bounds
checking, similar to Instance::memory32Copy.

A followup will implement the intrinsic for the firefox
translations project.

Differential Revision: https://phabricator.services.mozilla.com/D119919

js/src/builtin/TestingFunctions.cpp

@@ -125,6 +125,7 @@
 #include "wasm/WasmBaselineCompile.h"
 #include "wasm/WasmCraneliftCompile.h"
 #include "wasm/WasmInstance.h"
+#include "wasm/WasmIntrinsic.h"
 #include "wasm/WasmIonCompile.h"
 #include "wasm/WasmJS.h"
 #include "wasm/WasmModule.h"
@@ -1946,6 +1947,23 @@ static bool WasmLoadedFromCache(JSContext* cx, unsigned argc, Value* vp) {
   return WasmReturnFlag(cx, argc, vp, Flag::Deserialized);
 }
 
+static bool WasmIntrinsicI8VecMul(JSContext* cx, unsigned argc, Value* vp) {
+  if (!wasm::HasSupport(cx)) {
+    JS_ReportErrorASCII(cx, "wasm support unavailable");
+    return false;
+  }
+
+  CallArgs args = CallArgsFromVp(argc, vp);
+
+  RootedWasmModuleObject module(cx);
+  if (!wasm::CompileIntrinsicModule(cx, wasm::IntrinsicOp::I8VecMul, &module)) {
+    ReportOutOfMemory(cx);
+    return false;
+  }
+  args.rval().set(ObjectValue(*module.get()));
+  return true;
+}
+
 static bool LargeArrayBufferEnabled(JSContext* cx, unsigned argc, Value* vp) {
   CallArgs args = CallArgsFromVp(argc, vp);
   args.rval().setBoolean(ArrayBufferObject::maxBufferByteLength() >
@@ -7948,6 +7966,10 @@ JS_FOR_WASM_FEATURES(WASM_FEATURE, WASM_FEATURE)
 "  Returns a boolean indicating whether a given module was deserialized directly from a\n"
 "  cache (as opposed to compiled from bytecode)."),
 
+    JS_FN_HELP("wasmIntrinsicI8VecMul", WasmIntrinsicI8VecMul, 0, 0,
+"wasmIntrinsicI8VecMul()",
+"  Returns a module that implements an i8 vector pairwise multiplication intrinsic."),
+
     JS_FN_HELP("largeArrayBufferEnabled", LargeArrayBufferEnabled, 0, 0,
 "largeArrayBufferEnabled()",
 "  Returns true if array buffers larger than 2GB can be allocated."),

js/src/jit-test/tests/wasm/intrinsics/directives.txt

@@ -0,0 +1 @@
+|jit-test| test-also=--wasm-compiler=optimizing; include:wasm.js

js/src/jit-test/tests/wasm/intrinsics/i8vecmul.js

@@ -0,0 +1,32 @@
+let memory = new WebAssembly.Memory({initial: 1});
+let bytes = new Uint8Array(memory.buffer);
+
+let module = wasmIntrinsicI8VecMul();
+let instance = new WebAssembly.Instance(module, {
+  "": {"memory": memory}
+});
+let {i8vecmul} = instance.exports;
+
+// Test basic vector pairwise product
+{
+  // [0, 1, 2, 3] . [0, 2, 4, 6] = [0, 2, 8, 18]
+  for (let i = 0; i < 4; i++) {
+    bytes[i] = i;
+    bytes[4 + i] = i * 2;
+  }
+  i8vecmul(
+    /* dest */ 8,
+    /* src1 */ 0,
+    /* src2 */ 4,
+    /* len */ 4);
+  for (let i = 0; i < 4; i++) {
+    assertEq(bytes[8 + i], i * i * 2);
+  }
+}
+
+// Test bounds checking
+{
+  assertErrorMessage(() => i8vecmul(PageSizeInBytes - 1, 0, 0, 2), WebAssembly.RuntimeError, /index out of bounds/);
+  assertErrorMessage(() => i8vecmul(0, PageSizeInBytes - 1, 0, 2), WebAssembly.RuntimeError, /index out of bounds/);
+  assertErrorMessage(() => i8vecmul(0, 0, PageSizeInBytes - 1, 2), WebAssembly.RuntimeError, /index out of bounds/);
+}

js/src/jit/IonTypes.h

@@ -869,6 +869,9 @@ enum ABIFunctionType : uint32_t {
   Args_Int32_GeneralInt32Int32Int32Int32Int32 = detail::MakeABIFunctionType(
       ArgType_Int32, {ArgType_General, ArgType_Int32, ArgType_Int32,
                       ArgType_Int32, ArgType_Int32, ArgType_Int32}),
+  Args_Int32_GeneralInt32Int32Int32Int32General = detail::MakeABIFunctionType(
+      ArgType_Int32, {ArgType_General, ArgType_Int32, ArgType_Int32,
+                      ArgType_Int32, ArgType_Int32, ArgType_General}),
   Args_Int32_GeneralInt32Int32Int32General = detail::MakeABIFunctionType(
       ArgType_Int32, {ArgType_General, ArgType_Int32, ArgType_Int32,
                       ArgType_Int32, ArgType_General}),

js/src/jit/arm/Simulator-arm.cpp

@@ -2374,6 +2374,8 @@ typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32Int32)(int32_t, int32_t,
                                                                int32_t);
 typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32Int32Int32)(
     int32_t, int32_t, int32_t, int32_t, int32_t, int32_t);
+typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32Int32General)(
+    int32_t, int32_t, int32_t, int32_t, int32_t, int32_t);
 typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32General)(
     int32_t, int32_t, int32_t, int32_t, int32_t);
 typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int64)(int32_t, int32_t,
@@ -2839,6 +2841,15 @@ void Simulator::softwareInterrupt(SimInstruction* instr) {
           setCallResult(result);
           break;
         }
+        case Args_Int32_GeneralInt32Int32Int32Int32General: {
+          Prototype_Int32_GeneralInt32Int32Int32Int32General target =
+              reinterpret_cast<
+                  Prototype_Int32_GeneralInt32Int32Int32Int32General>(external);
+          int64_t result = target(arg0, arg1, arg2, arg3, arg4, arg5);
+          scratchVolatileRegisters(/* scratchFloat = true */);
+          setCallResult(result);
+          break;
+        }
         case Args_Int32_GeneralInt32Int32Int32General: {
           Prototype_Int32_GeneralInt32Int32Int32General target =
               reinterpret_cast<Prototype_Int32_GeneralInt32Int32Int32General>(

js/src/jit/arm64/vixl/MozSimulator-vixl.cpp

@@ -541,6 +541,12 @@ typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32Int32Int32)(int64_t,
                                                                     int32_t,
                                                                     int32_t,
                                                                     int32_t);
+typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32Int32General)(int64_t,
+                                                                      int32_t,
+                                                                      int32_t,
+                                                                      int32_t,
+                                                                      int32_t,
+                                                                      int64_t);
 typedef int32_t (*Prototype_Int32_GeneralInt32Int32Int32General)(int64_t,
                                                                  int32_t,
                                                                  int32_t,
@@ -814,6 +820,13 @@ Simulator::VisitCallRedirection(const Instruction* instr)
       setGPR32Result(ret);
       break;
     }
+    case js::jit::Args_Int32_GeneralInt32Int32Int32Int32General: {
+      int32_t ret =
+          reinterpret_cast<Prototype_Int32_GeneralInt32Int32Int32Int32General>(
+              nativeFn)(x0, x1, x2, x3, x4, x5);
+      setGPR32Result(ret);
+      break;
+    }
     case js::jit::Args_Int32_GeneralInt32Int32Int32General: {
       int32_t ret =
           reinterpret_cast<Prototype_Int32_GeneralInt32Int32Int32General>(

js/src/wasm/WasmBaselineCompile.cpp

@@ -8676,6 +8676,8 @@ class BaseCompiler final : public BaseCompilerInterface {
   [[nodiscard]] bool emitVectorShiftRightI64x2();
 #  endif
 #endif
+
+  [[nodiscard]] bool emitIntrinsic(IntrinsicOp op);
 };
 
 // TODO: We want these to be inlined for sure; do we need an `inline` somewhere?
@@ -15492,6 +15494,26 @@ bool BaseCompiler::emitVectorShiftRightI64x2() {
 #  endif
 #endif  // ENABLE_WASM_SIMD
 
+bool BaseCompiler::emitIntrinsic(IntrinsicOp op) {
+  uint32_t lineOrBytecode = readCallSiteLineOrBytecode();
+  const Intrinsic& intrinsic = Intrinsic::getFromOp(op);
+
+  NothingVector params;
+  if (!iter_.readIntrinsic(intrinsic, &params)) {
+    return false;
+  }
+
+  if (deadCode_) {
+    return true;
+  }
+
+  // The final parameter of an intrinsic is implicitly the heap base
+  pushHeapBase();
+
+  // Call the intrinsic
+  return emitInstanceCall(lineOrBytecode, intrinsic.signature);
+}
+
 bool BaseCompiler::emitBody() {
   MOZ_ASSERT(stackMapGenerator_.framePushedAtEntryToBody.isSome());
 
@@ -17151,6 +17173,15 @@ bool BaseCompiler::emitBody() {
       case uint16_t(Op::MozPrefix):
         return iter_.unrecognizedOpcode(&op);
 
+      // private intrinsic operations
+      case uint16_t(Op::IntrinsicPrefix): {
+        if (!moduleEnv_.intrinsicsEnabled() ||
+            op.b1 >= uint32_t(IntrinsicOp::Limit)) {
+          return iter_.unrecognizedOpcode(&op);
+        }
+        CHECK_NEXT(emitIntrinsic(IntrinsicOp(op.b1)));
+      }
+
       default:
         return iter_.unrecognizedOpcode(&op);
     }

js/src/wasm/WasmBinary.h

@@ -70,6 +70,11 @@ class Opcode {
     static_assert(size_t(SimdOp::Limit) <= 0xFFFFFF, "fits");
     MOZ_ASSERT(size_t(op) < size_t(SimdOp::Limit));
   }
+  MOZ_IMPLICIT Opcode(IntrinsicOp op)
+      : bits_((uint32_t(op) << 8) | uint32_t(Op::IntrinsicPrefix)) {
+    static_assert(size_t(IntrinsicOp::Limit) <= 0xFFFFFF, "fits");
+    MOZ_ASSERT(size_t(op) < size_t(IntrinsicOp::Limit));
+  }
 
   bool isOp() const { return bits_ < uint32_t(Op::FirstPrefix); }
   bool isMisc() const { return (bits_ & 255) == uint32_t(Op::MiscPrefix); }

js/src/wasm/WasmBuiltins.cpp

@@ -252,6 +252,12 @@ const SymbolicAddressSignature SASigRefTest = {
     SymbolicAddress::RefTest, _I32, _Infallible, 3, {_PTR, _RoN, _RoN, _END}};
 const SymbolicAddressSignature SASigRttSub = {
     SymbolicAddress::RttSub, _RoN, _FailOnNullPtr, 3, {_PTR, _RoN, _RoN, _END}};
+const SymbolicAddressSignature SASigIntrI8VecMul = {
+    SymbolicAddress::IntrI8VecMul,
+    _VOID,
+    _FailOnNegI32,
+    6,
+    {_PTR, _I32, _I32, _I32, _I32, _PTR, _END}};
 
 }  // namespace wasm
 }  // namespace js
@@ -1271,6 +1277,12 @@ void* wasm::AddressOf(SymbolicAddress imm, ABIFunctionType* abiType) {
       *abiType = Args_General1;
       return FuncCast(PrintText, *abiType);
 #endif
+    case SymbolicAddress::IntrI8VecMul:
+      *abiType = MakeABIFunctionType(
+          ArgType_Int32, {ArgType_General, ArgType_Int32, ArgType_Int32,
+                          ArgType_Int32, ArgType_Int32, ArgType_General});
+      MOZ_ASSERT(*abiType == ToABIType(SASigIntrI8VecMul));
+      return FuncCast(Instance::intrI8VecMul, *abiType);
     case SymbolicAddress::Limit:
       break;
   }
@@ -1397,6 +1409,7 @@ bool wasm::NeedsBuiltinThunk(SymbolicAddress sym) {
     case SymbolicAddress::ArrayNew:
     case SymbolicAddress::RefTest:
     case SymbolicAddress::RttSub:
+    case SymbolicAddress::IntrI8VecMul:
       return true;
     case SymbolicAddress::Limit:
       break;

js/src/wasm/WasmBuiltins.h

@@ -123,6 +123,7 @@ enum class SymbolicAddress {
 #if defined(JS_CODEGEN_MIPS32)
   js_jit_gAtomic64Lock,
 #endif
+  IntrI8VecMul,
 #ifdef WASM_CODEGEN_DEBUG
   PrintI32,
   PrintPtr,
@@ -235,6 +236,7 @@ extern const SymbolicAddressSignature SASigPushRefIntoExn;
 extern const SymbolicAddressSignature SASigArrayNew;
 extern const SymbolicAddressSignature SASigRefTest;
 extern const SymbolicAddressSignature SASigRttSub;
+extern const SymbolicAddressSignature SASigIntrI8VecMul;
 
 bool IsRoundingFunction(SymbolicAddress callee, jit::RoundingMode* mode);
 

js/src/wasm/WasmCompile.cpp

@@ -100,6 +100,7 @@ FeatureArgs FeatureArgs::build(JSContext* cx, const FeatureOptions& options) {
   if (wormholeOverride) {
     features.v128 = true;
   }
+  features.intrinsics = options.intrinsics;
 
   return features;
 }

js/src/wasm/WasmCompileArgs.h

@@ -73,10 +73,12 @@ class Tiers {
 // available under prefs.)
 
 struct FeatureOptions {
-  FeatureOptions() : simdWormhole(false) {}
+  FeatureOptions() : simdWormhole(false), intrinsics(false) {}
 
   // May be set if javascript.options.wasm_simd_wormhole==true.
   bool simdWormhole;
+  // Enables intrinsic opcodes, only set in WasmIntrinsic.cpp.
+  bool intrinsics;
 };
 
 // Describes the features that control wasm compilation.
@@ -89,7 +91,8 @@ struct FeatureArgs {
 #undef WASM_FEATURE
             sharedMemory(Shareable::False),
         hugeMemory(false),
-        simdWormhole(false) {
+        simdWormhole(false),
+        intrinsics(false) {
   }
   FeatureArgs(const FeatureArgs&) = default;
   FeatureArgs& operator=(const FeatureArgs&) = default;
@@ -104,6 +107,7 @@ struct FeatureArgs {
   Shareable sharedMemory;
   bool hugeMemory;
   bool simdWormhole;
+  bool intrinsics;
 };
 
 // Describes the JS scripted caller of a request to compile a wasm module.

js/src/wasm/WasmConstants.h

@@ -450,7 +450,8 @@ enum class Op {
   // GC (experimental)
   RefEq = 0xd5,
 
-  FirstPrefix = 0xfb,
+  FirstPrefix = 0xfa,
+  IntrinsicPrefix = 0xfa,
   GcPrefix = 0xfb,
   MiscPrefix = 0xfc,
   SimdPrefix = 0xfd,
@@ -937,6 +938,21 @@ enum class ThreadOp {
   Limit
 };
 
+enum class IntrinsicOp {
+  // ------------------------------------------------------------------------
+  // These operators are emitted internally when compiling intrinsic modules
+  // and are rejected by wasm validation.  They are prefixed by
+  // IntrinsicPrefix.
+
+  // i8vecmul(dest: i32, src1: i32, src2: i32, len: i32)
+  //  Performs pairwise multiplication of two i8 vectors of 'len' specified at
+  //  'src1' and 'src2'. Output is written to 'dest'. This is used as a
+  //  basic self-test for intrinsics.
+  I8VecMul = 0x0,
+
+  Limit
+};
+
 enum class MozOp {
   // ------------------------------------------------------------------------
   // These operators are emitted internally when compiling asm.js and are

js/src/wasm/WasmFrameIter.cpp

@@ -1526,6 +1526,8 @@ static const char* ThunkedNativeToDescription(SymbolicAddress func) {
     case SymbolicAddress::js_jit_gAtomic64Lock:
       MOZ_CRASH();
 #endif
+    case SymbolicAddress::IntrI8VecMul:
+      return "call to native i8 vector multiplication intrinsic (in wasm)";
 #ifdef WASM_CODEGEN_DEBUG
     case SymbolicAddress::PrintI32:
     case SymbolicAddress::PrintPtr:

js/src/wasm/WasmInstance.cpp

@@ -1106,6 +1106,41 @@ bool Instance::initElems(uint32_t tableIndex, const ElemSegment& seg,
   return AnyRef::fromJSObject(subRtt.get()).forCompiledCode();
 }
 
+/* static */ int32_t Instance::intrI8VecMul(Instance* instance, uint32_t dest,
+                                            uint32_t src1, uint32_t src2,
+                                            uint32_t len, uint8_t* memBase) {
+  MOZ_ASSERT(SASigIntrI8VecMul.failureMode == FailureMode::FailOnNegI32);
+
+  const WasmArrayRawBuffer* rawBuf = WasmArrayRawBuffer::fromDataPtr(memBase);
+  size_t memLen = rawBuf->byteLength();
+
+  // Bounds check and deal with arithmetic overflow.
+  uint64_t destLimit = uint64_t(dest) + uint64_t(len);
+  uint64_t src1Limit = uint64_t(src1) + uint64_t(len);
+  uint64_t src2Limit = uint64_t(src2) + uint64_t(len);
+  if (destLimit > memLen || src1Limit > memLen || src2Limit > memLen) {
+    JSContext* cx = TlsContext.get();
+    JS_ReportErrorNumberASCII(cx, GetErrorMessage, nullptr,
+                              JSMSG_WASM_OUT_OF_BOUNDS);
+    return -1;
+  }
+
+  // Basic dot product
+  uint8_t* destPtr = &memBase[dest];
+  uint8_t* src1Ptr = &memBase[src1];
+  uint8_t* src2Ptr = &memBase[src2];
+  while (len > 0) {
+    *destPtr = (*src1Ptr) * (*src2Ptr);
+
+    destPtr++;
+    src1Ptr++;
+    src2Ptr++;
+    len--;
+  }
+
+  return 0;
+}
+
 // Note, dst must point into nonmoveable storage that is not in the nursery,
 // this matters for the write barriers.  Furthermore, for pointer types the
 // current value of *dst must be null so that only a post-barrier is required.

js/src/wasm/WasmInstance.h

@@ -247,6 +247,8 @@ class Instance {
   static int32_t refTest(Instance* instance, void* refPtr, void* rttPtr);
   static void* rttSub(Instance* instance, void* rttParentPtr,
                       void* rttSubCanonPtr);
+  static int32_t intrI8VecMul(Instance* instance, uint32_t dest, uint32_t src1,
+                              uint32_t src2, uint32_t len, uint8_t* memBase);
 };
 
 using UniqueInstance = UniquePtr<Instance>;