Bug 1279312 - Pass a TLS pointer hidden argument to WebAssembly functions. r=luke

WebAssembly functions now expect to be passed a hidden argument in WasmTlsReg
which is a pointer to a TlsData struct.

Temporarily allocate a TlsData instance in the wasm::Instance itself. When wasm
supports multithreading, we will need to allocate a TlsData instance per thread
per module instance.

This patch generates code to pass the TLS pointer to WebAssembly functions,
preserving it through intra-module calls. The pointer is not used for anything
yet, and the the TLS pointer register is not currently preserved across function
calls.

js/src/asmjs/WasmInstance.cpp

@@ -417,7 +417,8 @@ Instance::callImport_f64(int32_t funcImportIndex, int32_t argc, uint64_t* argv)
     return ToNumber(cx, rval, (double*)argv);
 }
 
-Instance::Instance(UniqueCodeSegment codeSegment,
+Instance::Instance(JSContext* cx,
+                   UniqueCodeSegment codeSegment,
                    const Metadata& metadata,
                    const ShareableBytes* maybeBytecode,
                    HandleWasmMemoryObject memory,
@@ -446,6 +447,8 @@ Instance::Instance(UniqueCodeSegment codeSegment,
 
     for (size_t i = 0; i < tables_.length(); i++)
         *addressOfTableBase(i) = tables_[i]->array();
+
+   updateStackLimit(cx);
 }
 
 bool
@@ -509,6 +512,13 @@ Instance::memoryLength() const
     return memory_->buffer().byteLength();
 }
 
+void
+Instance::updateStackLimit(JSContext* cx)
+{
+    // Capture the stack limit for cx's thread.
+    tlsData_.stackLimit = *(void**)cx->stackLimitAddressForJitCode(StackForUntrustedScript);
+}
+
 bool
 Instance::callExport(JSContext* cx, uint32_t funcIndex, CallArgs args)
 {
@@ -624,7 +634,7 @@ Instance::callExport(JSContext* cx, uint32_t funcIndex, CallArgs args)
 
         // Call the per-exported-function trampoline created by GenerateEntry.
         auto funcPtr = JS_DATA_TO_FUNC_PTR(ExportFuncPtr, codeSegment_->code() + func.entryOffset());
-        if (!CALL_GENERATED_2(funcPtr, exportArgs.begin(), codeSegment_->globalData()))
+        if (!CALL_GENERATED_3(funcPtr, exportArgs.begin(), codeSegment_->globalData(), tlsData()))
             return false;
     }
 

js/src/asmjs/WasmInstance.h

@@ -50,8 +50,14 @@ class Instance
     bool                                 profilingEnabled_;
     CacheableCharsVector                 funcLabels_;
 
+
     UniquePtr<GeneratedSourceMap>        maybeSourceMap_;
 
+    // Thread-local data for code running in this instance.
+    // When threading is supported, we need a TlsData object per thread per
+    // instance.
+    TlsData                              tlsData_;
+
     // Internal helpers:
     uint8_t** addressOfMemoryBase() const;
     void** addressOfTableBase(size_t tableIndex) const;
@@ -59,6 +65,9 @@ class Instance
     FuncImportExit& funcImportToExit(const FuncImport& fi);
     MOZ_MUST_USE bool toggleProfiling(JSContext* cx);
 
+    // Get this instance's TLS data pointer for the current thread.
+    TlsData* tlsData() { return &tlsData_; }
+
     // An instance keeps track of its innermost WasmActivation. A WasmActivation
     // is pushed for the duration of each call of an export.
     friend class js::WasmActivation;
@@ -74,7 +83,8 @@ class Instance
     static int32_t callImport_f64(int32_t importIndex, int32_t argc, uint64_t* argv);
 
   public:
-    Instance(UniqueCodeSegment codeSegment,
+    Instance(JSContext* cx,
+             UniqueCodeSegment codeSegment,
              const Metadata& metadata,
              const ShareableBytes* maybeBytecode,
              HandleWasmMemoryObject memory,
@@ -136,6 +146,9 @@ class Instance
     const MemoryAccess* lookupMemoryAccess(void* pc) const;
 #endif
 
+    // Update the instance's copy of the stack limit.
+    void updateStackLimit(JSContext*);
+
     // about:memory reporting:
 
     void addSizeOfMisc(MallocSizeOf mallocSizeOf,

js/src/asmjs/WasmIonCompile.cpp

@@ -96,6 +96,9 @@ class FunctionCompiler
 
     FuncCompileResults&        compileResults_;
 
+    // TLS pointer argument to the current function.
+    MAsmJSParameter*           tlsPointer_;
+
   public:
     FunctionCompiler(const ModuleGeneratorData& mg,
                      Decoder& decoder,
@@ -117,7 +120,8 @@ class FunctionCompiler
         maxStackArgBytes_(0),
         loopDepth_(0),
         blockDepth_(0),
-        compileResults_(compileResults)
+        compileResults_(compileResults),
+        tlsPointer_(nullptr)
     {}
 
     const ModuleGeneratorData& mg() const    { return mg_; }
@@ -145,6 +149,12 @@ class FunctionCompiler
                 return false;
         }
 
+        // Set up a parameter that receives the hidden TLS pointer argument.
+        tlsPointer_ = MAsmJSParameter::New(alloc(), ABIArg(WasmTlsReg), MIRType::Pointer);
+        curBlock_->add(tlsPointer_);
+        if (!mirGen_.ensureBallast())
+            return false;
+
         for (size_t i = args.length(); i < locals_.length(); i++) {
             MInstruction* ins = nullptr;
             switch (locals_[i]) {
@@ -813,6 +823,14 @@ class FunctionCompiler
         return args->stackArgs_.append(mir);
     }
 
+    // Add the hidden TLS pointer argument to CallArgs.
+    bool passTlsPointer(CallArgs* args)
+    {
+        if (inDeadCode())
+            return true;
+        return args->regArgs_.append(MAsmJSCall::Arg(AnyRegister(WasmTlsReg), tlsPointer_));
+    }
+
     void propagateMaxStackArgBytes(uint32_t stackBytes)
     {
         if (callStack_.empty()) {
@@ -1656,8 +1674,16 @@ EmitReturn(FunctionCompiler& f)
     return true;
 }
 
+// Is a callee within the same module instance?
+enum class IntraModule
+{
+    False,
+    True
+};
+
 static bool
-EmitCallArgs(FunctionCompiler& f, const Sig& sig, FunctionCompiler::CallArgs* args)
+EmitCallArgs(FunctionCompiler& f, const Sig& sig, IntraModule intraModule,
+             FunctionCompiler::CallArgs* args)
 {
     if (!f.startCallArgs(args))
         return false;
@@ -1676,6 +1702,11 @@ EmitCallArgs(FunctionCompiler& f, const Sig& sig, FunctionCompiler::CallArgs* ar
     if (!f.iter().readCallArgsEnd(numArgs))
         return false;
 
+    // Calls within the module pass the module's TLS pointer.
+    // Calls to other modules go through stubs that set up their TLS pointers.
+    if (intraModule == IntraModule::True)
+        f.passTlsPointer(args);
+
     f.finishCallArgs(args);
     return true;
 }
@@ -1693,7 +1724,7 @@ EmitCall(FunctionCompiler& f, uint32_t callOffset)
     const Sig& sig = *f.mg().funcSigs[calleeIndex];
 
     FunctionCompiler::CallArgs args(f, lineOrBytecode);
-    if (!EmitCallArgs(f, sig, &args))
+    if (!EmitCallArgs(f, sig, IntraModule::True, &args))
         return false;
 
     if (!f.iter().readCallReturn(sig.ret()))
@@ -1723,7 +1754,7 @@ EmitCallIndirect(FunctionCompiler& f, uint32_t callOffset)
     const Sig& sig = f.mg().sigs[sigIndex];
 
     FunctionCompiler::CallArgs args(f, lineOrBytecode);
-    if (!EmitCallArgs(f, sig, &args))
+    if (!EmitCallArgs(f, sig, IntraModule::True, &args))
         return false;
 
     MDefinition* callee;
@@ -1762,7 +1793,7 @@ EmitCallImport(FunctionCompiler& f, uint32_t callOffset)
     const Sig& sig = *funcImport.sig;
 
     FunctionCompiler::CallArgs args(f, lineOrBytecode);
-    if (!EmitCallArgs(f, sig, &args))
+    if (!EmitCallArgs(f, sig, IntraModule::False, &args))
         return false;
 
     if (!f.iter().readCallReturn(sig.ret()))

js/src/asmjs/WasmModule.cpp

@@ -574,7 +574,8 @@ Module::instantiate(JSContext* cx,
         if (!instanceObj)
             return false;
 
-        auto instance = cx->make_unique<Instance>(Move(codeSegment),
+        auto instance = cx->make_unique<Instance>(cx,
+                                                  Move(codeSegment),
                                                   *metadata_,
                                                   maybeBytecode,
                                                   memory,

js/src/asmjs/WasmStubs.cpp

@@ -91,10 +91,10 @@ static const unsigned FramePushedAfterSave = NonVolatileRegs.gprs().size() * siz
 #endif
 static const unsigned FramePushedForEntrySP = FramePushedAfterSave + sizeof(void*);
 
-// Generate a stub that enters wasm from a C++ caller via the native ABI.
-// The signature of the entry point is Module::CodePtr. The exported wasm
+// Generate a stub that enters wasm from a C++ caller via the native ABI. The
+// signature of the entry point is Module::ExportFuncPtr. The exported wasm
 // function has an ABI derived from its specific signature, so this function
-// must map from the ABI of CodePtr to the export's signature's ABI.
+// must map from the ABI of ExportFuncPtr to the export's signature's ABI.
 Offsets
 wasm::GenerateEntry(MacroAssembler& masm, const FuncExport& fe, bool usesHeap)
 {
@@ -132,6 +132,19 @@ wasm::GenerateEntry(MacroAssembler& masm, const FuncExport& fe, bool usesHeap)
     if (usesHeap)
         masm.loadAsmJSHeapRegisterFromGlobalData();
 
+    // Put the per-thread, per-module TLS pointer into WasmTlsReg.
+    // This is the third argument in the ExportFuncPtr prototype.
+#if defined(JS_CODEGEN_X86)
+    masm.loadPtr(
+      Address(masm.getStackPointer(), EntryFrameSize + masm.framePushed() + 2 * sizeof(void*)),
+      WasmTlsReg);
+#else
+    masm.movePtr(IntArgReg2, WasmTlsReg);
+#endif
+    // Make sure the TLS pointer is not clobbered by the following code.
+    MOZ_ASSERT(WasmTlsReg != ABINonArgReg0, "TLS pointer can't be scratch reg");
+    MOZ_ASSERT(WasmTlsReg != ABINonArgReg1, "TLS pointer can't be scratch reg");
+
     // Put the 'argv' argument into a non-argument/return register so that we
     // can use 'argv' while we fill in the arguments for the asm.js callee.
     // Also, save 'argv' on the stack so that we can recover it after the call.

js/src/asmjs/WasmTypes.h

@@ -900,7 +900,27 @@ struct ExportArg
     uint64_t hi;
 };
 
-typedef int32_t (*ExportFuncPtr)(ExportArg* args, uint8_t* global);
+// TLS data for a single module instance.
+//
+// Every WebAssembly function expects to be passed a hidden TLS pointer argument
+// in WasmTlsReg. The TLS pointer argument points to a TlsData struct.
+// Compiled functions expect that the TLS pointer does not change for the
+// lifetime of the thread.
+//
+// There is a TlsData per module instance per thread, so inter-module calls need
+// to pass the TLS pointer appropriate for the callee module.
+//
+// After the TlsData struct follows the module's declared TLS variables.
+//
+struct TlsData
+{
+    // Stack limit for the current thread. This limit is checked against the
+    // stack pointer in the prologue of functions that allocate stack space. See
+    // `CodeGenerator::generateWasm`.
+    void* stackLimit;
+};
+
+typedef int32_t (*ExportFuncPtr)(ExportArg* args, uint8_t* global, TlsData* tls);
 
 // Constants:
 

js/src/jit/JitCommon.h

@@ -33,6 +33,10 @@
     (js::jit::Simulator::Current()->call(                               \
         JS_FUNC_TO_DATA_PTR(uint8_t*, entry), 2, p0, p1))
 
+#define CALL_GENERATED_3(entry, p0, p1, p2)                             \
+    (js::jit::Simulator::Current()->call(                               \
+        JS_FUNC_TO_DATA_PTR(uint8_t*, entry), 3, p0, p1, p2))
+
 #else
 
 // Call into jitted code by following the ABI of the native architecture.
@@ -41,6 +45,7 @@
 
 #define CALL_GENERATED_1(entry, p0)         entry(p0)
 #define CALL_GENERATED_2(entry, p0, p1)     entry(p0, p1)
+#define CALL_GENERATED_3(entry, p0, p1, p2) entry(p0, p1, p2)
 
 #endif
 

js/src/jit/arm/Assembler-arm.h

@@ -74,6 +74,11 @@ static constexpr Register CallTempNonArgRegs[] = { r5, r6, r7, r8 };
 static const uint32_t NumCallTempNonArgRegs =
     mozilla::ArrayLength(CallTempNonArgRegs);
 
+// TLS pointer argument register for WebAssembly functions. This must not alias
+// any other register used for passing function arguments or return values.
+// Preserved by WebAssembly functions.
+static constexpr Register WasmTlsReg = r9;
+
 class ABIArgGenerator
 {
     unsigned intRegIndex_;

js/src/jit/arm64/Assembler-arm64.h

@@ -88,6 +88,11 @@ static constexpr Register GlobalReg =  { Registers::x20 };
 static constexpr Register HeapReg = { Registers::x21 };
 static constexpr Register HeapLenReg = { Registers::x22 };
 
+// TLS pointer argument register for WebAssembly functions. This must not alias
+// any other register used for passing function arguments or return values.
+// Preserved by WebAssembly functions.
+static constexpr Register WasmTlsReg = { Registers::x17 };
+
 // Define unsized Registers.
 #define DEFINE_UNSIZED_REGISTERS(N)  \
 static constexpr Register r##N = { Registers::x##N };

js/src/jit/mips-shared/Assembler-mips-shared.h

@@ -107,6 +107,11 @@ static constexpr Register ReturnReg = v0;
 static constexpr FloatRegister ReturnSimd128Reg = InvalidFloatReg;
 static constexpr FloatRegister ScratchSimd128Reg = InvalidFloatReg;
 
+// TLS pointer argument register for WebAssembly functions. This must not alias
+// any other register used for passing function arguments or return values.
+// Preserved by WebAssembly functions.
+static constexpr Register WasmTlsReg = s5;
+
 // A bias applied to the GlobalReg to allow the use of instructions with small
 // negative immediate offsets which doubles the range of global data that can be
 // accessed with a single instruction.

js/src/jit/x64/Assembler-x64.h

@@ -149,6 +149,11 @@ static constexpr uint32_t NumFloatArgRegs = 8;
 static constexpr FloatRegister FloatArgRegs[NumFloatArgRegs] = { xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7 };
 #endif
 
+// TLS pointer argument register for WebAssembly functions. This must not alias
+// any other register used for passing function arguments or return values.
+// Preserved by WebAssembly functions.
+static constexpr Register WasmTlsReg = r14;
+
 // Registers used in the GenerateFFIIonExit Enable Activation block.
 static constexpr Register AsmJSIonExitRegCallee = r10;
 static constexpr Register AsmJSIonExitRegE0 = rax;

js/src/jit/x86/Assembler-x86.h

@@ -52,6 +52,11 @@ static constexpr FloatRegister ScratchFloat32Reg = FloatRegister(X86Encoding::xm
 static constexpr FloatRegister ScratchDoubleReg = FloatRegister(X86Encoding::xmm7, FloatRegisters::Double);
 static constexpr FloatRegister ScratchSimd128Reg = FloatRegister(X86Encoding::xmm7, FloatRegisters::Simd128);
 
+// TLS pointer argument register for WebAssembly functions. This must not alias
+// any other register used for passing function arguments or return values.
+// Preserved by WebAssembly functions.
+static constexpr Register WasmTlsReg = esi;
+
 // Avoid ebp, which is the FramePointer, which is unavailable in some modes.
 static constexpr Register ArgumentsRectifierReg = esi;
 static constexpr Register CallTempReg0 = edi;