diff --git a/lld/test/wasm/data-segments.ll b/lld/test/wasm/data-segments.ll
--- a/lld/test/wasm/data-segments.ll
+++ b/lld/test/wasm/data-segments.ll
@@ -84,8 +84,8 @@
 ; PASSIVE-NEXT:        Body:            0B
 ; PASSIVE-NEXT:      - Index:           2
 ; PASSIVE-NEXT:        Locals:          []
-; PASSIVE32-NEXT:        Body:            41B4D60041004101FE480200044041B4D6004101427FFE0102001A054180084100410DFC08000041900841004114FC08010041B4D6004102FE17020041B4D600417FFE0002001A0BFC0900FC09010B
-; PASSIVE64-NEXT:        Body:            42B4D60041004101FE480200044042B4D6004101427FFE0102001A054280084100410DFC08000042900841004114FC08010042B4D6004102FE17020042B4D600417FFE0002001A0BFC0900FC09010B
+; PASSIVE32-NEXT:        Body:             02400240024041B4D60041004101FE4802000E020001020B4180084100410DFC08000041900841004114FC08010041B4D6004102FE17020041B4D600417FFE0002001A0C010B41B4D6004101427FFE0102001A0BFC0900FC09010B
+; PASSIVE64-NEXT:        Body:            02400240024042B4D60041004101FE4802000E020001020B4280084100410DFC08000042900841004114FC08010042B4D6004102FE17020042B4D600417FFE0002001A0C010B42B4D6004101427FFE0102001A0BFC0900FC09010B
 ; PASSIVE-NEXT:  - Type:            DATA
 ; PASSIVE-NEXT:    Segments:
 ; PASSIVE-NEXT:      - SectionOffset:   3
@@ -121,8 +121,8 @@
 ; PASSIVE32-PIC-NEXT:          - Type:            I32
 ; PASSIVE64-PIC-NEXT:          - Type:            I64
 ; PASSIVE-PIC-NEXT:            Count:           1
-; PASSIVE32-PIC-NEXT:        Body:            230141B4CE006A2100200041004101FE480200044020004101427FFE0102001A05410023016A4100410DFC080000411023016A41004114FC08010020004102FE1702002000417FFE0002001A0BFC0900FC09010B
-; PASSIVE64-PIC-NEXT:        Body:            230142B4CE007C2100200041004101FE480200044020004101427FFE0102001A05420023017C4100410DFC080000421023017C41004114FC08010020004102FE1702002000417FFE0002001A0BFC0900FC09010B
+; PASSIVE32-PIC-NEXT:          Body:            230141B4CE006A2100024002400240200041004101FE4802000E020001020B410023016A4100410DFC080000411023016A41004114FC08010020004102FE1702002000417FFE0002001A0C010B20004101427FFE0102001A0BFC0900FC09010B
+; PASSIVE64-PIC-NEXT:          Body:            230142B4CE007C2100024002400240200041004101FE4802000E020001020B420023017C4100410DFC080000421023017C41004114FC08010020004102FE1702002000417FFE0002001A0C010B20004101427FFE0102001A0BFC0900FC09010B
 ; PASSIVE-PIC-NEXT:      - Index:           3
 ; PASSIVE-PIC-NEXT:        Locals:          []
 ; PASSIVE-PIC-NEXT:        Body:            0B
diff --git a/lld/wasm/Writer.cpp b/lld/wasm/Writer.cpp
--- a/lld/wasm/Writer.cpp
+++ b/lld/wasm/Writer.cpp
@@ -1022,22 +1022,17 @@
     // initialized. The generated code is as follows:
     //
     // (func $__wasm_init_memory
-    //  (if
-    //   (i32.atomic.rmw.cmpxchg align=2 offset=0
-    //    (i32.const $__init_memory_flag)
-    //    (i32.const 0)
-    //    (i32.const 1)
-    //   )
-    //   (then
-    //    (drop
-    //     (i32.atomic.wait align=2 offset=0
-    //      (i32.const $__init_memory_flag)
-    //      (i32.const 1)
-    //      (i32.const -1)
+    //  (block $drop
+    //   (block $wait
+    //    (block $init
+    //     (br_table $init $wait $drop
+    //      (i32.atomic.rmw.cmpxchg align=2 offset=0
+    //       (i32.const $__init_memory_flag)
+    //       (i32.const 0)
+    //       (i32.const 1)
+    //      )
     //     )
-    //    )
-    //   )
-    //   (else
+    //    ) ;; $init
     //    ( ... initialize data segments ... )
     //    (i32.atomic.store align=2 offset=0
     //     (i32.const $__init_memory_flag)
@@ -1049,8 +1044,16 @@
     //      (i32.const -1u)
     //     )
     //    )
+    //    (br $drop)
+    //   ) ;; $wait
+    //   (drop
+    //    (i32.atomic.wait align=2 offset=0
+    //     (i32.const $__init_memory_flag)
+    //     (i32.const 1)
+    //     (i32.const -1)
+    //    )
     //   )
-    //  )
+    //  ) ;; $drop
     //  ( ... drop data segments ... )
     // )
     //
@@ -1084,29 +1087,31 @@
       }
     };
 
-    // Atomically check whether this is the main thread.
+    // Set up destination blocks
+    writeU8(os, WASM_OPCODE_BLOCK, "block $drop");
+    writeU8(os, WASM_TYPE_NORESULT, "block type");
+    writeU8(os, WASM_OPCODE_BLOCK, "block $wait");
+    writeU8(os, WASM_TYPE_NORESULT, "block type");
+    writeU8(os, WASM_OPCODE_BLOCK, "block $init");
+    writeU8(os, WASM_TYPE_NORESULT, "block type");
+
+    // Atomically check whether we win the race.
     writeGetFlagAddress();
     writeI32Const(os, 0, "expected flag value");
-    writeI32Const(os, 1, "flag value");
+    writeI32Const(os, 1, "new flag value");
     writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
     writeUleb128(os, WASM_OPCODE_I32_RMW_CMPXCHG, "i32.atomic.rmw.cmpxchg");
     writeMemArg(os, 2, 0);
-    writeU8(os, WASM_OPCODE_IF, "IF");
-    writeU8(os, WASM_TYPE_NORESULT, "blocktype");
-
-    // Did not increment 0, so wait for main thread to initialize memory
-    writeGetFlagAddress();
-    writeI32Const(os, 1, "expected flag value");
-    writeI64Const(os, -1, "timeout");
-
-    writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
-    writeUleb128(os, WASM_OPCODE_I32_ATOMIC_WAIT, "i32.atomic.wait");
-    writeMemArg(os, 2, 0);
-    writeU8(os, WASM_OPCODE_DROP, "drop");
 
-    writeU8(os, WASM_OPCODE_ELSE, "ELSE");
+    // Based on the value, decide what to do next.
+    writeU8(os, WASM_OPCODE_BR_TABLE, "br_table");
+    writeUleb128(os, 2, "label vector length");
+    writeUleb128(os, 0, "label $init");
+    writeUleb128(os, 1, "label $wait");
+    writeUleb128(os, 2, "default label $drop");
 
-    // Did increment 0, so conditionally initialize passive data segments
+    // Initialize passive data segments
+    writeU8(os, WASM_OPCODE_END, "end $init");
     for (const OutputSegment *s : segments) {
       if (needsPassiveInitialization(s)) {
         // destination address
@@ -1145,9 +1150,23 @@
     writeMemArg(os, 2, 0);
     writeU8(os, WASM_OPCODE_DROP, "drop");
 
-    writeU8(os, WASM_OPCODE_END, "END");
+    // Branch to drop the segments
+    writeU8(os, WASM_OPCODE_BR, "br");
+    writeUleb128(os, 1, "label $drop");
+
+    // Wait for the winning thread to initialize memory
+    writeU8(os, WASM_OPCODE_END, "end $wait");
+    writeGetFlagAddress();
+    writeI32Const(os, 1, "expected flag value");
+    writeI64Const(os, -1, "timeout");
+
+    writeU8(os, WASM_OPCODE_ATOMICS_PREFIX, "atomics prefix");
+    writeUleb128(os, WASM_OPCODE_I32_ATOMIC_WAIT, "i32.atomic.wait");
+    writeMemArg(os, 2, 0);
+    writeU8(os, WASM_OPCODE_DROP, "drop");
 
     // Unconditionally drop passive data segments
+    writeU8(os, WASM_OPCODE_END, "end $drop");
     for (const OutputSegment *s : segments) {
       if (needsPassiveInitialization(s)) {
         // data.drop instruction
@@ -1156,6 +1175,8 @@
         writeUleb128(os, s->index, "segment index immediate");
       }
     }
+
+    // End the function
     writeU8(os, WASM_OPCODE_END, "END");
   }
 
diff --git a/llvm/include/llvm/BinaryFormat/Wasm.h b/llvm/include/llvm/BinaryFormat/Wasm.h
--- a/llvm/include/llvm/BinaryFormat/Wasm.h
+++ b/llvm/include/llvm/BinaryFormat/Wasm.h
@@ -284,8 +284,10 @@
 
 // Opcodes used in synthetic functions.
 enum : unsigned {
-  WASM_OPCODE_IF = 0x04,
-  WASM_OPCODE_ELSE = 0x05,
+  WASM_OPCODE_BLOCK = 0x02,
+  WASM_OPCODE_BR = 0x0c,
+  WASM_OPCODE_BR_TABLE = 0x0e,
+  WASM_OPCODE_RETURN = 0x0f,
   WASM_OPCODE_DROP = 0x1a,
   WASM_OPCODE_MISC_PREFIX = 0xfc,
   WASM_OPCODE_MEMORY_INIT = 0x08,