Bug 1559965 - Wasm: Implement new atomic.fence instruction. r=bbouvier

This commit implements the 'atomic.fence' Wasm instruction.

Issue: https://github.com/WebAssembly/threads/issues/140
Overview: https://github.com/WebAssembly/threads/pull/141

The instruction is encoded as, 0xFE 0x03, with a reserved byte trailing for a future
memory order immediate. The instruction is implemented by emitting a memoryBarrier
through the macro assembler.

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

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

js/src/jit-test/tests/wasm/binary.js

@@ -256,10 +256,10 @@ function checkIllegalPrefixed(prefix, opcode) {
 //
 // June 2017 threads draft:
 //
-//  0x00 .. 0x02 are wait/wake ops
+//  0x00 .. 0x03 are wait/wake/fence ops
 //  0x10 .. 0x4f are primitive atomic ops
 
-for (let i = 3; i < 0x10; i++)
+for (let i = 0x4; i < 0x10; i++)
     checkIllegalPrefixed(ThreadPrefix, i);
 
 for (let i = 0x4f; i < 0x100; i++)

js/src/jit-test/tests/wasm/fence.js

@@ -0,0 +1,22 @@
+// Test that `atomic.fence` is a valid instruction of type `[] -> []`
+wasmFullPass(
+`(module
+	(func atomic.fence)
+	(export "run" 0)
+)`);
+
+// Test that `atomic.fence` works with non-shared memory
+wasmFullPass(
+`(module
+	(memory 1)
+	(func atomic.fence)
+	(export "run" 0)
+)`);
+
+// Test that `atomic.fence` works with shared memory
+wasmFullPass(
+`(module
+	(memory 1 1 shared)
+	(func atomic.fence)
+	(export "run" 0)
+)`);

js/src/jit/CodeGenerator.cpp

@@ -13980,6 +13980,11 @@ void CodeGenerator::visitWasmCompareAndSelect(LWasmCompareAndSelect* ins) {
   MOZ_CRASH("in CodeGenerator::visitWasmCompareAndSelect: unexpected types");
 }
 
+void CodeGenerator::visitWasmFence(LWasmFence* lir) {
+  MOZ_ASSERT(gen->compilingWasm());
+  masm.memoryBarrier(MembarFull);
+}
+
 static_assert(!std::is_polymorphic<CodeGenerator>::value,
               "CodeGenerator should not have any virtual methods");
 

js/src/jit/Lowering.cpp

@@ -5052,5 +5052,9 @@ void LIRGenerator::visitWasmSelect(MWasmSelect* ins) {
   defineReuseInput(lir, ins, LWasmSelect::TrueExprIndex);
 }
 
+void LIRGenerator::visitWasmFence(MWasmFence* ins) {
+  add(new (alloc()) LWasmFence, ins);
+}
+
 static_assert(!std::is_polymorphic<LIRGenerator>::value,
               "LIRGenerator should not have any virtual methods");

js/src/jit/MIR.h

@@ -11436,6 +11436,19 @@ class MAsmJSStoreHeap
   }
 };
 
+class MWasmFence : public MNullaryInstruction {
+ protected:
+  MWasmFence() : MNullaryInstruction(classOpcode) { setGuard(); }
+
+ public:
+  INSTRUCTION_HEADER(WasmFence)
+  TRIVIAL_NEW_WRAPPERS
+
+  AliasSet getAliasSet() const override { return AliasSet::None(); }
+
+  ALLOW_CLONE(MWasmFence)
+};
+
 class MWasmCompareExchangeHeap : public MVariadicInstruction,
                                  public NoTypePolicy::Data {
   wasm::MemoryAccessDesc access_;

js/src/jit/shared/LIR-shared.h

@@ -6405,6 +6405,12 @@ class LWasmCompareExchangeHeap : public LInstructionHelper<1, 4, 4> {
   }
 };
 
+class LWasmFence : public LInstructionHelper<0, 0, 0> {
+ public:
+  LIR_HEADER(WasmFence);
+  explicit LWasmFence() : LInstructionHelper(classOpcode) {}
+};
+
 class LWasmAtomicExchangeHeap : public LInstructionHelper<1, 3, 4> {
  public:
   LIR_HEADER(WasmAtomicExchangeHeap);

js/src/wasm/WasmAST.h

@@ -440,7 +440,8 @@ enum class AstExprKind {
   UnaryOperator,
   Unreachable,
   Wait,
-  Wake
+  Wake,
+  Fence
 };
 
 class AstExpr : public AstNode {
@@ -817,6 +818,11 @@ class AstWake : public AstExpr {
   AstExpr& count() const { return *count_; }
 };
 
+struct AstFence : AstExpr {
+  static const AstExprKind Kind = AstExprKind::Fence;
+  AstFence() : AstExpr(AstExprKind::Fence, ExprType::Void) {}
+};
+
 class AstMemOrTableCopy : public AstExpr {
   bool isMem_;
   AstRef destTable_;

js/src/wasm/WasmBaselineCompile.cpp

@@ -6852,6 +6852,7 @@ class BaseCompiler final : public BaseCompilerInterface {
   MOZ_MUST_USE bool emitAtomicStore(ValType type, Scalar::Type viewType);
   MOZ_MUST_USE bool emitWait(ValType type, uint32_t byteSize);
   MOZ_MUST_USE bool emitWake();
+  MOZ_MUST_USE bool emitFence();
   MOZ_MUST_USE bool emitAtomicXchg(ValType type, Scalar::Type viewType);
   void emitAtomicXchg64(MemoryAccessDesc* access, ValType type,
                         WantResult wantResult);
@@ -10199,6 +10200,18 @@ bool BaseCompiler::emitWake() {
   return emitInstanceCall(lineOrBytecode, SASigWake);
 }
 
+bool BaseCompiler::emitFence() {
+  if (!iter_.readFence()) {
+    return false;
+  }
+  if (deadCode_) {
+    return true;
+  }
+
+  masm.memoryBarrier(MembarFull);
+  return true;
+}
+
 // Bulk memory must be available if shared memory is enabled.
 bool BaseCompiler::bulkmemOpsEnabled() {
 #ifndef ENABLE_WASM_BULKMEM_OPS
@@ -11602,6 +11615,8 @@ bool BaseCompiler::emitBody() {
             CHECK_NEXT(emitWait(ValType::I32, 4));
           case uint32_t(ThreadOp::I64Wait):
             CHECK_NEXT(emitWait(ValType::I64, 8));
+          case uint32_t(ThreadOp::Fence):
+            CHECK_NEXT(emitFence());
 
           case uint32_t(ThreadOp::I32AtomicLoad):
             CHECK_NEXT(emitAtomicLoad(ValType::I32, Scalar::Int32));

js/src/wasm/WasmConstants.h

@@ -402,6 +402,7 @@ enum class ThreadOp {
   Wake = 0x00,
   I32Wait = 0x01,
   I64Wait = 0x02,
+  Fence = 0x03,
 
   // Load and store
   I32AtomicLoad = 0x10,

js/src/wasm/WasmIonCompile.cpp

@@ -279,6 +279,14 @@ class FunctionCompiler {
     return constant;
   }
 
+  void fence() {
+    if (inDeadCode()) {
+      return;
+    }
+    MWasmFence* ins = MWasmFence::New(alloc());
+    curBlock_->add(ins);
+  }
+
   template <class T>
   MDefinition* unary(MDefinition* op) {
     if (inDeadCode()) {
@@ -2800,6 +2808,15 @@ static bool EmitWait(FunctionCompiler& f, ValType type, uint32_t byteSize) {
   return true;
 }
 
+static bool EmitFence(FunctionCompiler& f) {
+  if (!f.iter().readFence()) {
+    return false;
+  }
+
+  f.fence();
+  return true;
+}
+
 static bool EmitWake(FunctionCompiler& f) {
   uint32_t lineOrBytecode = f.readCallSiteLineOrBytecode();
 
@@ -3854,6 +3871,8 @@ static bool EmitBodyExprs(FunctionCompiler& f) {
             CHECK(EmitWait(f, ValType::I32, 4));
           case uint32_t(ThreadOp::I64Wait):
             CHECK(EmitWait(f, ValType::I64, 8));
+          case uint32_t(ThreadOp::Fence):
+            CHECK(EmitFence(f));
 
           case uint32_t(ThreadOp::I32AtomicLoad):
             CHECK(EmitAtomicLoad(f, ValType::I32, Scalar::Int32));

js/src/wasm/WasmOpIter.cpp

@@ -318,6 +318,8 @@ OpKind wasm::Classify(OpBytes op) {
         case ThreadOp::I32Wait:
         case ThreadOp::I64Wait:
           return OpKind::Wait;
+        case ThreadOp::Fence:
+          return OpKind::Fence;
         case ThreadOp::I32AtomicLoad:
         case ThreadOp::I64AtomicLoad:
         case ThreadOp::I32AtomicLoad8U:

js/src/wasm/WasmOpIter.h

@@ -144,6 +144,7 @@ enum class OpKind {
   End,
   Wait,
   Wake,
+  Fence,
   AtomicLoad,
   AtomicStore,
   AtomicBinOp,
@@ -442,6 +443,7 @@ class MOZ_STACK_CLASS OpIter : private Policy {
   MOZ_MUST_USE bool readWait(LinearMemoryAddress<Value>* addr,
                              ValType resultType, uint32_t byteSize,
                              Value* value, Value* timeout);
+  MOZ_MUST_USE bool readFence();
   MOZ_MUST_USE bool readAtomicLoad(LinearMemoryAddress<Value>* addr,
                                    ValType resultType, uint32_t byteSize);
   MOZ_MUST_USE bool readAtomicStore(LinearMemoryAddress<Value>* addr,
@@ -1702,6 +1704,19 @@ inline bool OpIter<Policy>::readWait(LinearMemoryAddress<Value>* addr,
   return true;
 }
 
+template <typename Policy>
+inline bool OpIter<Policy>::readFence() {
+  MOZ_ASSERT(Classify(op_) == OpKind::Fence);
+  uint8_t flags;
+  if (!readFixedU8(&flags)) {
+    return fail("expected memory order after fence");
+  }
+  if (flags != 0) {
+    return fail("non-zero memory order not supported yet");
+  }
+  return true;
+}
+
 template <typename Policy>
 inline bool OpIter<Policy>::readAtomicLoad(LinearMemoryAddress<Value>* addr,
                                            ValType resultType,

js/src/wasm/WasmTextToBinary.cpp

@@ -85,6 +85,7 @@ class WasmToken {
     Error,
     Export,
     ExtraConversionOpcode,
+    Fence,
     Field,
     Float,
     Func,
@@ -230,7 +231,7 @@ class WasmToken {
     MOZ_ASSERT(begin != end);
     MOZ_ASSERT(kind_ == AtomicCmpXchg || kind_ == AtomicLoad ||
                kind_ == AtomicRMW || kind_ == AtomicStore || kind_ == Wait ||
-               kind_ == Wake);
+               kind_ == Wake || kind_ == Fence);
     u.threadOp_ = op;
   }
   explicit WasmToken(const char16_t* begin)
@@ -305,6 +306,7 @@ class WasmToken {
       case DataDrop:
       case Drop:
       case ElemDrop:
+      case Fence:
       case GetGlobal:
       case GetLocal:
       case If:
@@ -941,6 +943,9 @@ WasmToken WasmTokenStream::next() {
         if (consume(u"wake") || consume(u"notify")) {
           return WasmToken(WasmToken::Wake, ThreadOp::Wake, begin, cur_);
         }
+        if (consume(u"fence")) {
+          return WasmToken(WasmToken::Fence, ThreadOp::Fence, begin, cur_);
+        }
         break;
       }
       break;
@@ -3987,6 +3992,8 @@ static AstExpr* ParseExprBody(WasmParseContext& c, WasmToken token,
       return ParseWait(c, token.threadOp(), inParens);
     case WasmToken::Wake:
       return ParseWake(c, inParens);
+    case WasmToken::Fence:
+      return new (c.lifo) AstFence();
     case WasmToken::BinaryOpcode:
       return ParseBinaryOperator(c, token.op(), inParens);
     case WasmToken::Block:
@@ -5819,6 +5826,8 @@ static bool ResolveExpr(Resolver& r, AstExpr& expr) {
       return ResolveWait(r, expr.as<AstWait>());
     case AstExprKind::Wake:
       return ResolveWake(r, expr.as<AstWake>());
+    case AstExprKind::Fence:
+      return true;
     case AstExprKind::MemOrTableCopy:
       return ResolveMemOrTableCopy(r, expr.as<AstMemOrTableCopy>());
     case AstExprKind::DataOrElemDrop:
@@ -6423,6 +6432,10 @@ static bool EncodeWake(Encoder& e, AstWake& s) {
          e.writeOp(ThreadOp::Wake) && EncodeLoadStoreFlags(e, s.address());
 }
 
+static bool EncodeFence(Encoder& e, AstFence& s) {
+  return e.writeOp(ThreadOp::Fence) && e.writeFixedU8(0);
+}
+
 static bool EncodeMemOrTableCopy(Encoder& e, AstMemOrTableCopy& s) {
   return EncodeExpr(e, s.dest()) && EncodeExpr(e, s.src()) &&
          EncodeExpr(e, s.len()) &&
@@ -6622,6 +6635,8 @@ static bool EncodeExpr(Encoder& e, AstExpr& expr) {
       return EncodeWait(e, expr.as<AstWait>());
     case AstExprKind::Wake:
       return EncodeWake(e, expr.as<AstWake>());
+    case AstExprKind::Fence:
+      return EncodeFence(e, expr.as<AstFence>());
     case AstExprKind::MemOrTableCopy:
       return EncodeMemOrTableCopy(e, expr.as<AstMemOrTableCopy>());
     case AstExprKind::DataOrElemDrop:

js/src/wasm/WasmValidate.cpp

@@ -994,6 +994,9 @@ static bool DecodeFunctionBodyExprs(const ModuleEnvironment& env,
             LinearMemoryAddress<Nothing> addr;
             CHECK(iter.readWait(&addr, ValType::I64, 8, &nothing, &nothing));
           }
+          case uint32_t(ThreadOp::Fence): {
+            CHECK(iter.readFence());
+          }
           case uint32_t(ThreadOp::I32AtomicLoad): {
             LinearMemoryAddress<Nothing> addr;
             CHECK(iter.readAtomicLoad(&addr, ValType::I32, 4));