Bug 1609057 - Add support for stack results in Ion wasm compiler r=lth

This patch extends MIR with a MWasmStackResultArea which contains the
locations to which stack result values will be written.  The intention
is that the MWasmStackResultArea will get passed as an additional
argument to calls with stack results, and the results will get extracted
from the area via instances of MWasmStackResult.

Next steps will be to wire up WasmIonCompile and implement the back-end,
extending the register allocator to handle stack result areas specially.

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

--HG--
extra : moz-landing-system : lando

js/src/jit/IonAnalysis.cpp

@@ -2951,6 +2951,7 @@ static bool IsResumableMIRType(MIRType type) {
     case MIRType::Doublex2:  // NYI, see also RSimdBox::recover
     case MIRType::Int64:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       return false;
   }
   MOZ_CRASH("Unknown MIRType.");

js/src/jit/IonTypes.h

@@ -448,13 +448,14 @@ enum class MIRType : uint8_t {
   // Types above are specialized.
   Value,
   ObjectOrNull,
-  None,         // Invalid, used as a placeholder.
-  Slots,        // A slots vector
-  Elements,     // An elements vector
-  Pointer,      // An opaque pointer that receives no special treatment
-  RefOrNull,    // Wasm Ref/AnyRef/NullRef: a raw JSObject* or a raw (void*)0
-  Shape,        // A Shape pointer.
-  ObjectGroup,  // An ObjectGroup pointer.
+  None,          // Invalid, used as a placeholder.
+  Slots,         // A slots vector
+  Elements,      // An elements vector
+  Pointer,       // An opaque pointer that receives no special treatment
+  RefOrNull,     // Wasm Ref/AnyRef/NullRef: a raw JSObject* or a raw (void*)0
+  StackResults,  // Wasm multi-value stack result area, which may contain refs
+  Shape,         // A Shape pointer.
+  ObjectGroup,   // An ObjectGroup pointer.
   Last = ObjectGroup,
   // Representing both SIMD.IntBxN and SIMD.UintBxN.
   Int8x16 = Int32 | (4 << VECTOR_SCALE_SHIFT),
@@ -601,6 +602,8 @@ static inline const char* StringFromMIRType(MIRType type) {
       return "Pointer";
     case MIRType::RefOrNull:
       return "RefOrNull";
+    case MIRType::StackResults:
+      return "StackResults";
     case MIRType::Shape:
       return "Shape";
     case MIRType::ObjectGroup:

js/src/jit/LIR.cpp

@@ -367,6 +367,8 @@ static const char* DefTypeName(LDefinition::Type type) {
       return "simd128int";
     case LDefinition::SIMD128FLOAT:
       return "simd128float";
+    case LDefinition::STACKRESULTS:
+      return "stackresults";
 #  ifdef JS_NUNBOX32
     case LDefinition::TYPE:
       return "t";

js/src/jit/LIR.h

@@ -401,7 +401,6 @@ class LDefinition {
     MUST_REUSE_INPUT
   };
 
-  // This should be kept in sync with LIR.cpp's TypeChars.
   enum Type {
     GENERAL,     // Generic, integer or pointer-width data (GPR).
     INT32,       // int32 data (GPR).
@@ -411,6 +410,7 @@ class LDefinition {
     DOUBLE,      // 64-bit floating-point value (FPU).
     SIMD128INT,  // 128-bit SIMD integer vector (FPU).
     SIMD128FLOAT,  // 128-bit SIMD floating point vector (FPU).
+    STACKRESULTS,  // A variable-size stack allocation that may contain objects.
 #ifdef JS_NUNBOX32
     // A type virtual register must be followed by a payload virtual
     // register, as both will be tracked as a single gcthing.
@@ -547,6 +547,8 @@ class LDefinition {
       case MIRType::Int64:
         return LDefinition::GENERAL;
 #endif
+      case MIRType::StackResults:
+        return LDefinition::STACKRESULTS;
       case MIRType::Int8x16:
       case MIRType::Int16x8:
       case MIRType::Int32x4:

js/src/jit/Lowering.cpp

@@ -4436,6 +4436,14 @@ void LIRGenerator::visitWasmRegister64Result(MWasmRegister64Result* ins) {
   add(lir, ins);
 }
 
+void LIRGenerator::visitWasmStackResultArea(MWasmStackResultArea* ins) {
+  MOZ_CRASH("unimplemented");
+}
+
+void LIRGenerator::visitWasmStackResult(MWasmStackResult* ins) {
+  MOZ_CRASH("unimplemented");
+}
+
 void LIRGenerator::visitWasmCall(MWasmCall* ins) {
   bool needsBoundsCheck = true;
   if (ins->callee().isTable()) {

js/src/jit/MIR.h

@@ -11862,6 +11862,91 @@ class MWasmRegister64Result : public MWasmResultBase<Register64> {
   TRIVIAL_NEW_WRAPPERS
 };
 
+class MWasmStackResultArea : public MNullaryInstruction {
+ public:
+  class StackResult {
+    // Offset in bytes from lowest address of stack result area.
+    uint32_t offset_;
+    MIRType type_;
+
+   public:
+    StackResult() : type_(MIRType::Undefined) {}
+    StackResult(uint32_t offset, MIRType type) : offset_(offset), type_(type) {}
+
+    bool initialized() const { return type_ != MIRType::Undefined; }
+    uint32_t offset() const {
+      MOZ_ASSERT(initialized());
+      return offset_;
+    }
+    MIRType type() const {
+      MOZ_ASSERT(initialized());
+      return type_;
+    }
+    uint32_t endOffset() const { return offset() + MIRTypeToSize(type()); }
+  };
+
+ private:
+  FixedList<StackResult> results_;
+
+  explicit MWasmStackResultArea() : MNullaryInstruction(classOpcode) {
+    setResultType(MIRType::StackResults);
+  }
+
+  void assertInitialized() const {
+    MOZ_ASSERT(results_.length() != 0);
+#ifdef DEBUG
+    for (size_t i = 0; i < results_.length(); i++) {
+      MOZ_ASSERT(results_[i].initialized());
+    }
+#endif
+  }
+
+ public:
+  INSTRUCTION_HEADER(WasmStackResultArea)
+  TRIVIAL_NEW_WRAPPERS
+
+  MOZ_MUST_USE bool init(TempAllocator& alloc, size_t stackResultCount) {
+    MOZ_ASSERT(results_.length() == 0);
+    MOZ_ASSERT(stackResultCount > 0);
+    return results_.init(alloc, stackResultCount);
+  }
+
+  size_t resultCount() const { return results_.length(); }
+  const StackResult& result(size_t n) const {
+    MOZ_ASSERT(results_[n].initialized());
+    return results_[n];
+  }
+  void initResult(size_t n, const StackResult& loc) {
+    MOZ_ASSERT(!results_[n].initialized());
+    MOZ_ASSERT((n == 0) == (loc.offset() == 0));
+    MOZ_ASSERT_IF(n > 0, loc.offset() >= result(n - 1).endOffset());
+    results_[n] = loc;
+  }
+
+  uint32_t byteSize() const {
+    assertInitialized();
+    return result(resultCount() - 1).endOffset();
+  }
+};
+
+class MWasmStackResult : public MUnaryInstruction, public NoTypePolicy::Data {
+  uint32_t resultIdx_;
+
+  const MWasmStackResultArea::StackResult& result() const {
+    return resultArea()->toWasmStackResultArea()->result(resultIdx_);
+  }
+
+  MWasmStackResult(MWasmStackResultArea* resultArea, size_t idx)
+      : MUnaryInstruction(classOpcode, resultArea), resultIdx_(idx) {
+    setResultType(result().type());
+  }
+
+ public:
+  INSTRUCTION_HEADER(WasmStackResult)
+  TRIVIAL_NEW_WRAPPERS
+  NAMED_OPERANDS((0, resultArea))
+};
+
 class MWasmCall final : public MVariadicInstruction, public NoTypePolicy::Data {
   wasm::CallSiteDesc desc_;
   wasm::CalleeDesc callee_;

js/src/jit/StackSlotAllocator.h

@@ -91,6 +91,8 @@ class StackSlotAllocator {
       case LDefinition::SIMD128INT:
       case LDefinition::SIMD128FLOAT:
         return 16;
+      case LDefinition::STACKRESULTS:
+        MOZ_CRASH("Stack results area must be allocated manually");
     }
     MOZ_CRASH("Unknown slot type");
   }

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

@@ -48,6 +48,7 @@ ABIArg ABIArgGenerator::softNext(MIRType type) {
     case MIRType::Int32:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       if (intRegIndex_ == NumIntArgRegs) {
         current_ = ABIArg(stackOffset_);
         stackOffset_ += sizeof(uint32_t);
@@ -109,6 +110,7 @@ ABIArg ABIArgGenerator::hardNext(MIRType type) {
     case MIRType::Int32:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       if (intRegIndex_ == NumIntArgRegs) {
         current_ = ABIArg(stackOffset_);
         stackOffset_ += sizeof(uint32_t);

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

@@ -35,6 +35,7 @@ ABIArg ABIArgGenerator::next(MIRType type) {
     case MIRType::Int64:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       if (intRegIndex_ == NumIntArgRegs) {
         current_ = ABIArg(stackOffset_);
         stackOffset_ += sizeof(uintptr_t);

js/src/jit/mips32/Assembler-mips32.cpp

@@ -26,6 +26,7 @@ ABIArg ABIArgGenerator::next(MIRType type) {
     case MIRType::Int32:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       if (GetIntArgReg(usedArgSlots_, &destReg)) {
         current_ = ABIArg(destReg);
       } else {

js/src/jit/mips64/Assembler-mips64.cpp

@@ -22,7 +22,8 @@ ABIArg ABIArgGenerator::next(MIRType type) {
     case MIRType::Int32:
     case MIRType::Int64:
     case MIRType::Pointer:
-    case MIRType::RefOrNull: {
+    case MIRType::RefOrNull:
+    case MIRType::StackResults: {
       Register destReg;
       if (GetIntArgReg(usedArgSlots_, &destReg)) {
         current_ = ABIArg(destReg);

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

@@ -47,6 +47,7 @@ ABIArg ABIArgGenerator::next(MIRType type) {
     case MIRType::Int64:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       current_ = ABIArg(IntArgRegs[regIndex_++]);
       break;
     case MIRType::Float32:
@@ -77,6 +78,7 @@ ABIArg ABIArgGenerator::next(MIRType type) {
     case MIRType::Int64:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       if (intRegIndex_ == NumIntArgRegs) {
         current_ = ABIArg(stackOffset_);
         stackOffset_ += sizeof(uint64_t);

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

@@ -20,6 +20,7 @@ ABIArg ABIArgGenerator::next(MIRType type) {
     case MIRType::Float32:
     case MIRType::Pointer:
     case MIRType::RefOrNull:
+    case MIRType::StackResults:
       current_ = ABIArg(stackOffset_);
       stackOffset_ += sizeof(uint32_t);
       break;

js/src/wasm/WasmBaselineCompile.cpp

@@ -1076,12 +1076,14 @@ int32_t BaseLocalIter::pushLocal(size_t nbytes) {
 void BaseLocalIter::settle() {
   if (!argsIter_.done()) {
     mirType_ = argsIter_.mirType();
+    MIRType concreteType = mirType_;
     switch (mirType_) {
-      case MIRType::Pointer:
+      case MIRType::StackResults:
         // The pointer to stack results is handled like any other argument:
         // either addressed in place if it is passed on the stack, or we spill
         // it in the frame if it's in a register.
         MOZ_ASSERT(args_.isSyntheticStackResultPointerArg(index_));
+        concreteType = MIRType::Pointer;
         [[fallthrough]];
       case MIRType::Int32:
       case MIRType::Int64:
@@ -1089,7 +1091,7 @@ void BaseLocalIter::settle() {
       case MIRType::Float32:
       case MIRType::RefOrNull:
         if (argsIter_->argInRegister()) {
-          frameOffset_ = pushLocal(MIRTypeToSize(mirType_));
+          frameOffset_ = pushLocal(MIRTypeToSize(concreteType));
         } else {
           frameOffset_ = -(argsIter_->offsetFromArgBase() + sizeof(Frame));
         }
@@ -1097,7 +1099,7 @@ void BaseLocalIter::settle() {
       default:
         MOZ_CRASH("Argument type");
     }
-    if (mirType_ == MIRType::Pointer) {
+    if (mirType_ == MIRType::StackResults) {
       stackResultPointerOffset_ = frameOffset();
       // Advance past the synthetic stack result pointer argument and fall
       // through to the next case.

js/src/wasm/WasmGC.cpp

@@ -184,8 +184,6 @@ bool GenerateStackmapEntriesForTrapExit(const ArgTypeVector& args,
   }
 
   for (ABIArgIter i(args); !i.done(); i++) {
-    MOZ_ASSERT((i.mirType() == MIRType::Pointer) ==
-               args.isSyntheticStackResultPointerArg(i.index()));
     if (!i->argInRegister() || i.mirType() != MIRType::RefOrNull) {
       continue;
     }

js/src/wasm/WasmStubs.cpp

@@ -315,7 +315,7 @@ static void SetupABIArguments(MacroAssembler& masm, const FuncExport& fe,
           masm.load64(src, iter->gpr64());
         } else if (type == MIRType::RefOrNull) {
           masm.loadPtr(src, iter->gpr());
-        } else if (type == MIRType::Pointer) {
+        } else if (type == MIRType::StackResults) {
           MOZ_ASSERT(args.isSyntheticStackResultPointerArg(iter.index()));
           MOZ_CRASH("multiple function results not yet implemented");
         } else {
@@ -378,7 +378,7 @@ static void SetupABIArguments(MacroAssembler& masm, const FuncExport& fe,
                                                  iter->offsetFromArgBase()));
             break;
           }
-          case MIRType::Pointer: {
+          case MIRType::StackResults: {
             MOZ_ASSERT(args.isSyntheticStackResultPointerArg(iter.index()));
             MOZ_CRASH("multiple function results not yet implemented");
             break;
@@ -1379,7 +1379,7 @@ void wasm::GenerateDirectCallFromJit(MacroAssembler& masm, const FuncExport& fe,
         case MIRType::RefOrNull:
           GenPrintPtr(DebugChannel::Function, masm, iter->gpr());
           break;
-        case MIRType::Pointer:
+        case MIRType::StackResults:
           MOZ_ASSERT(args.isSyntheticStackResultPointerArg(iter.index()));
           GenPrintPtr(DebugChannel::Function, masm, iter->gpr());
           break;
@@ -1449,7 +1449,7 @@ void wasm::GenerateDirectCallFromJit(MacroAssembler& masm, const FuncExport& fe,
             masm.storePtr(scratch, dst);
             break;
           }
-          case MIRType::Pointer: {
+          case MIRType::StackResults: {
             MOZ_CRASH("multi-value in ion to wasm fast path unimplemented");
           }
           default: {

js/src/wasm/WasmTypes.h

@@ -1212,7 +1212,7 @@ class ArgTypeVector {
   jit::MIRType operator[](size_t i) const {
     MOZ_ASSERT(i < length());
     if (isSyntheticStackResultPointerArg(i)) {
-      return jit::MIRType::Pointer;
+      return jit::MIRType::StackResults;
     }
     return ToMIRType(args_[naturalIndex(i)]);
   }