[AMDGPU] Don't emit llvm.amdgcn.permlane64 on CDNA

quadrants/runtime/llvm/llvm_context.cpp

@@ -552,43 +552,66 @@ std::unique_ptr<llvm::Module> QuadrantsLLVMContext::module_from_file(const std::
       }
       patch_intrinsic("amdgpu_clock_i64", llvm::Intrinsic::amdgcn_s_memtime);
       patch_intrinsic("amdgpu_ds_bpermute", llvm::Intrinsic::amdgcn_ds_bpermute);
-      // ``llvm.amdgcn.permlane64`` exchanges a 32-bit value between lanes ``i`` and ``i ^ 32`` in a single instruction.
-      // We use it to extend the SIMD32-scoped ``ds_bpermute`` (every shuffle op lowers to that) into a wave64-aware
-      // cross-half shuffle on RDNA: ``ds_bpermute`` reads within the lane's own 32-lane SIMD cluster, ``permlane64``
-      // brings the other SIMD's value to this lane, and we select between the two based on which half the target lane
-      // sits in. See ``amdgpu_cross_half_shuffle_i32`` in runtime.cpp. The instruction is gfx940+ (CDNA3) and gfx11+
-      // (RDNA3+) only -- on earlier wave64-capable targets (gfx9xx CDNA1/2, gfx10.x RDNA1/2) the AMDGPU LLVM backend
-      // hits "Cannot select" while lowering the intrinsic, so we have to provide a software emulation.
+      // The wave64 cross-half subgroup shuffle (``amdgpu_cross_half_shuffle_i32`` in runtime.cpp) is built from
+      // ``ds_bpermute`` plus, on some targets, ``permlane64``. How those behave -- and therefore how we patch them --
+      // depends on the architecture family:
       //
-      // The emulation is a wave-local LDS roundtrip: each lane writes its ``value`` to ``lds[wave_base + lane]``,
-      // a wavefront-scope acquire-release fence lowers to ``s_waitcnt lgkmcnt(0)`` (drains outstanding LDS writes),
-      // and each lane then reads ``lds[wave_base + (lane ^ 32)]``. On RDNA wave64-emulation the two SIMD32 halves of
-      // the wave issue store / load in two passes apiece, but the waitcnt between them guarantees both halves' stores
-      // are committed to LDS before either half's loads issue, so the cross-half routing is correct.  ``wave_base``
-      // is ``(workitem.id.x >> 6) << 6``, scoping the LDS slot to a single wave so multi-wave workgroups don't
-      // collide.  The LDS buffer is a 1024-entry per-workgroup global (4 KiB) -- enough for the AMDGPU 1024-thread
-      // workgroup max at wave64.  The buffer is only materialised on this code path, so kernels on permlane64-capable
-      // hardware (the common case) pay zero LDS for cross-half shuffles.
+      //   * GCN / CDNA (gfx9xx: gfx900/906 Vega, gfx908/gfx90a CDNA1/2, gfx940/gfx941/gfx942 CDNA3, gfx950 CDNA4):
+      //     ``ds_bpermute`` addresses the full wave64 directly, so the cross-half shuffle is a single wide
+      //     ``ds_bpermute`` (lane mask 63) and ``permlane64`` is unnecessary. Critically, ``v_permlane64_b32`` does
+      //     not exist on CDNA -- emitting ``llvm.amdgcn.permlane64`` makes the backend select the ``V_PERMLANE64_B32``
+      //     pseudo, which has no valid MC opcode for CDNA, and then crash with a bare SIGSEGV inside
+      //     ``SIInstrInfo::getInstSizeInBytes`` during branch relaxation (genesis-world #2962). So on CDNA we patch
+      //     ``amdgpu_permlane64`` to the identity, which neutralises the helper's (RDNA-shaped) cross-SIMD branch
+      //     without ever emitting the intrinsic.
+      //   * RDNA3/4 (gfx11 / gfx12): ``ds_bpermute`` is SIMD32-scoped (lane mask 31), so the top half is reached via
+      //     the native single-instruction ``v_permlane64_b32``.
+      //   * RDNA1/2 (gfx10.x): ``ds_bpermute`` is SIMD32-scoped (lane mask 31), but ``v_permlane64_b32`` does not
+      //     exist yet, so we emulate the lane ``i`` <-> ``i ^ 32`` swap with an LDS roundtrip (below).
       //
-      // The intrinsic is overloaded on its element type (signature ``T -> T`` for any 32-bit-or-smaller ``T``), so we
-      // have to pass the explicit ``i32`` type alongside the ID -- otherwise ``CreateIntrinsic`` segfaults inside
-      // ``getDeclaration()`` while resolving the mangled name.
+      // The LDS emulation writes each lane's ``value`` to ``lds[wave_base + lane]``, issues a wavefront-scope
+      // acquire-release fence (lowers to ``s_waitcnt lgkmcnt(0)``: drains outstanding LDS writes without the
+      // cross-wave ``s_barrier`` a workgroup-scope fence would emit, which would deadlock if only some waves reach
+      // this point), then reads back ``lds[wave_base + (lane ^ 32)]``. ``wave_base`` is ``(workitem.id.x >> 6) << 6``,
+      // scoping the slot to a single wave so multi-wave workgroups don't collide. The buffer is a 1024-entry
+      // per-workgroup global (4 KiB, the AMDGPU 1024-thread wave64 max), materialised only on this path, so kernels
+      // on the other two paths pay zero LDS for cross-half shuffles.
+      //
+      // ``patch_intrinsic`` for permlane64 passes the explicit ``i32`` type alongside the ID because the intrinsic is
+      // overloaded on its element type (signature ``T -> T`` for any 32-bit-or-smaller ``T``); otherwise
+      // ``CreateIntrinsic`` segfaults inside ``getDeclaration()`` while resolving the mangled name.
       auto mcpu_str = AMDGPUContext::get_instance().get_mcpu();
-      bool has_permlane64 = (mcpu_str == "gfx940" || mcpu_str == "gfx941" || mcpu_str == "gfx942" ||
-                             mcpu_str.substr(0, 5) == "gfx11" || mcpu_str.substr(0, 5) == "gfx12");
-      // Escape hatch for validating the LDS software emulation on hardware that natively supports
-      // ``v_permlane64_b32``: setting ``QD_AMDGPU_FORCE_PERMLANE64_FALLBACK=1`` forces the JIT to take the LDS path
-      // even on gfx11+ / gfx940+, so we can exercise the fallback on a working AMD box (gfx1100 / gfx942) without
-      // needing a gfx10.x runner.  Has no effect on non-AMDGPU backends.
-      if (const char *force_fallback = std::getenv("QD_AMDGPU_FORCE_PERMLANE64_FALLBACK")) {
-        if (force_fallback[0] == '1') {
-          has_permlane64 = false;
-        }
+      bool is_gcn_cdna = (mcpu_str.substr(0, 4) == "gfx9");
+      bool has_native_permlane64 = (mcpu_str.substr(0, 5) == "gfx11" || mcpu_str.substr(0, 5) == "gfx12");
+
+      // Patch the ds_bpermute lane mask used by ``amdgpu_cross_half_shuffle_i32``: 63 where ``ds_bpermute`` is
+      // wave64-wide (GCN/CDNA), 31 where it is SIMD32-scoped (RDNA, paired with the permlane64 swap above).
+      if (auto mask_func = module->getFunction("amdgpu_ds_bpermute_lane_mask")) {
+        mask_func->deleteBody();
+        auto bb = llvm::BasicBlock::Create(*ctx, "entry", mask_func);
+        IRBuilder<> builder(*ctx);
+        builder.SetInsertPoint(bb);
+        builder.CreateRet(llvm::ConstantInt::get(llvm::Type::getInt32Ty(*ctx), is_gcn_cdna ? 63 : 31));
+        QuadrantsLLVMContext::mark_inline(mask_func);
       }
-      if (has_permlane64) {
+
+      if (has_native_permlane64) {
         patch_intrinsic("amdgpu_permlane64", llvm::Intrinsic::amdgcn_permlane64, true, {llvm::Type::getInt32Ty(*ctx)});
+      } else if (is_gcn_cdna) {
+        // CDNA: the wide ds_bpermute already reaches all 64 lanes, so permlane64 is unnecessary -- and emitting it
+        // crashes the backend. Patch it to the identity so the helper's cross-SIMD branch returns the same value as
+        // its same-SIMD branch, making the per-lane select a true no-op.
+        if (auto permlane64_func = module->getFunction("amdgpu_permlane64")) {
+          permlane64_func->deleteBody();
+          auto bb = llvm::BasicBlock::Create(*ctx, "entry", permlane64_func);
+          IRBuilder<> builder(*ctx);
+          builder.SetInsertPoint(bb);
+          builder.CreateRet(&*permlane64_func->arg_begin());
+          QuadrantsLLVMContext::mark_inline(permlane64_func);
+        }
       } else if (auto permlane64_func = module->getFunction("amdgpu_permlane64")) {
-        // LDS-based software emulation.  Layout: ``[1024 x i32] addrspace(3)`` indexed by ``wave_base + lane``.
+        // gfx10.x RDNA1/2: LDS-based software emulation. Layout: ``[1024 x i32] addrspace(3)`` indexed by ``wave_base +
+        // lane``.
         auto i32_ty = llvm::Type::getInt32Ty(*ctx);
         auto buf_ty = llvm::ArrayType::get(i32_ty, 1024);
         auto lds_global = llvm::cast_or_null<llvm::GlobalVariable>(module->getNamedValue("__amdgpu_permlane64_lds"));

quadrants/runtime/llvm/runtime_module/runtime.cpp

@@ -1015,13 +1015,13 @@ i32 amdgpu_ds_bpermute(i32 byte_index, i32 value) {
 }
 
 // Exchanges a 32-bit value between lanes ``i`` and ``i ^ 32`` in a single instruction. The native instruction
-// ``v_permlane64_b32`` is only available on gfx940+ (CDNA3) and gfx11+ (RDNA3+); ``llvm_context.cpp`` detects the
-// target at JIT time and patches this stub to either the ``llvm.amdgcn.permlane64`` intrinsic (on supported
-// hardware) or an LDS-roundtrip software emulation (on gfx9xx CDNA1/2 and gfx10.x RDNA1/2). The emulation has higher
-// latency (LDS store + ``s_waitcnt`` + LDS load -- roughly tens of cycles per call vs. a few for the native swap),
-// but produces correct cross-half results on RDNA wave64 emulation hardware. Used by
-// ``amdgpu_cross_half_shuffle_i32`` below to repair the cross-half story for ``ds_bpermute``, which is SIMD32-scoped
-// on RDNA.
+// ``v_permlane64_b32`` is RDNA-only -- it exists on gfx11 (RDNA3) and gfx12 (RDNA4), but on NO CDNA part (the AMD
+// assembler rejects it on gfx908/gfx90a/gfx940/gfx941/gfx942/gfx950). ``llvm_context.cpp`` detects the target at JIT
+// time and patches this stub to either the ``llvm.amdgcn.permlane64`` intrinsic (on gfx11+/gfx12 only) or an
+// LDS-roundtrip software emulation (on every CDNA part and on gfx10.x RDNA1/2). The emulation has higher latency (LDS
+// store + ``s_waitcnt`` + LDS load -- roughly tens of cycles per call vs. a few for the native swap), but produces
+// correct cross-half results on wave64 hardware. Used by ``amdgpu_cross_half_shuffle_i32`` below to repair the
+// cross-half story for ``ds_bpermute``, which is SIMD32-scoped on RDNA.
 i32 amdgpu_permlane64(i32 value) {
   __builtin_trap();
   return 0;
@@ -1059,6 +1059,15 @@ i32 amdgpu_lane_id() {
   return amdgpu_mbcnt_hi(-1, amdgpu_mbcnt_lo(-1, 0));
 }
 
+// Lane-index mask applied to ``ds_bpermute``'s target in ``amdgpu_cross_half_shuffle_i32``. ``ds_bpermute``'s reach
+// differs by ISA, so ``llvm_context.cpp`` patches this at JIT load time: 63 on GCN/CDNA (gfx9xx), where
+// ``ds_bpermute`` addresses the full wave64 directly, and 31 on RDNA (gfx10/11/12), where it is SIMD32-scoped and the
+// top half is reached via the ``permlane64`` swap instead. Defaults to 31 so an unpatched build stays correct on
+// RDNA. (runtime.cpp is compiled at -O0 to bitcode, so this stays a real call until the JIT patches + inlines it.)
+i32 amdgpu_ds_bpermute_lane_mask() {
+  return 31;
+}
+
 // Wave64-aware "read ``value`` from lane ``target_lane``" gather for AMDGPU. Shared by every i32 shuffle variant
 // (``shuffle`` / ``shuffle_down`` / ``shuffle_up``); the f32 / i64 / f64 wrappers below decompose into i32 calls and
 // therefore inherit the wave64 fix for free.
@@ -1075,13 +1084,16 @@ i32 amdgpu_lane_id() {
 // so ``ds_bpermute(byte, permlane64(value))`` effectively reads from lanes 32-63. We always compute both reads and
 // select between them branchlessly based on the high bit of ``target_lane``: bit 5 picks the half.
 //
-// Note this is correct on every AMDGPU target we run on. On CDNA (gfx9xx, gfx940/942) ``ds_bpermute`` could in
-// principle directly address all 64 lanes, but because we always mask the byte argument to ``(target_lane & 31) * 4``
-// we never test that path -- on both ISAs the byte index is in [0, 128) and only addresses the bottom half. The
-// ``permlane64`` swap then supplies the top-half data: on hardware with the native instruction (gfx940+ CDNA3 /
-// gfx11+ RDNA3+) this is a single ``v_permlane64_b32``; on older wave64-capable targets (gfx9xx CDNA1/2, gfx10.x
-// RDNA1/2) the JIT patches ``amdgpu_permlane64`` to an LDS roundtrip that produces the same result at higher latency
-// (see the patching logic in ``llvm_context.cpp``).
+// The lane-index mask (``amdgpu_ds_bpermute_lane_mask()``) and ``permlane64`` are both JIT-patched per architecture
+// (see ``llvm_context.cpp``), because ``ds_bpermute``'s reach differs by ISA:
+//   * GCN / CDNA (gfx9xx, incl. gfx940/gfx942/gfx950): ``ds_bpermute`` addresses the full wave64, so the mask is 63
+//     and a single ``ds_bpermute((target_lane & 63) * 4, value)`` already returns lane ``target_lane`` for the whole
+//     wave. ``permlane64`` is patched to the identity there -- ``v_permlane64_b32`` does not exist on CDNA and
+//     emitting it crashes the backend (genesis-world #2962) -- so ``from_other_half`` equals ``from_self_half`` and
+//     the select below is a true no-op.
+//   * RDNA (gfx10/11/12): ``ds_bpermute`` is SIMD32-scoped, so the mask is 31 and ``ds_bpermute`` only reaches the
+//     issuing lane's own 32-lane half. The top half is supplied by the ``permlane64`` swap -- a single
+//     ``v_permlane64_b32`` on gfx11/gfx12, or an LDS-roundtrip emulation on gfx10.x.
 //
 // OOR target lanes (``target_lane < 0`` or ``target_lane >= 64``): we mask to ``target_lane & 31`` for the byte and
 // ``& 32`` for the half-bit. The behaviour for OOR targets is implementation-defined on every backend (CUDA's
@@ -1093,12 +1105,12 @@ i32 amdgpu_cross_half_shuffle_i32(i32 target_lane, i32 value) {
   // wave -- lifting it above the select keeps the AMDGPU backend happy and lets it issue exactly one
   // ``v_permlane64_b32``. ``ds_bpermute`` on RDNA wave64 is SIMD32-scoped with a 5-bit address (top half of the wave
   // is unreachable directly), so ``from_self_half`` handles the same-SIMD case and ``from_other_half`` handles the
-  // cross-SIMD case via the ``swapped`` payload. On CDNA the wave is one SIMD64 so both reads return the same value
-  // and the select is a no-op; we don't try to optimize that out because the dead read is cheap (LLVM CSE may fold
-  // it anyway).
+  // cross-SIMD case via the ``swapped`` payload. On CDNA the lane mask is 63 (``ds_bpermute`` is wave64-wide) and
+  // ``permlane64`` is patched to the identity, so ``swapped == value``, both reads return lane ``target_lane``, and
+  // the select is a true no-op; we don't optimize that out because the dead read is cheap (LLVM CSE may fold it).
   i32 self_lane = amdgpu_lane_id();
   i32 swapped = amdgpu_permlane64(value);
-  i32 byte = (target_lane & 31) * 4;
+  i32 byte = (target_lane & amdgpu_ds_bpermute_lane_mask()) * 4;
   // ``llvm.amdgcn.ds.bpermute`` is the real hardware ``ds_bpermute_b32`` -- but if LLVM's uniformity analysis decides
   // ``byte`` is uniform across the wave (e.g. ``target_lane`` is a compile-time constant), it sometimes lowers to a
   // ``v_readlane_b32``-style instruction that addresses lanes 0..31 wave-globally rather than SIMD32-locally. On

tests/python/test_simt.py

@@ -3845,8 +3845,12 @@ def test_subgroup_exclusive_add_tiled_log2_size_6():
 # to the ``permlane64``-based cross-half helper a lane in the bottom half could not read the top half (and vice
 # versa).  All five tests are gated to ``log2_group_size() == 6`` so they only assert anything on real wave64
 # hardware -- CUDA and SPIR-V backends with wave32 skip the absolute-correctness check (the cross-half partner is
-# out of range there, which is implementation-defined).  CDNA (gfx9xx, MI300X) already had a wave64-wide
-# ``ds_bpermute`` so its behaviour is unchanged by the fix; the new helper is observably a no-op on that path.
+# out of range there, which is implementation-defined).  Both CDNA and RDNA wave64 run these, via different lowerings
+# of ``amdgpu_cross_half_shuffle_i32``: on CDNA (gfx9xx, incl. gfx942 MI300X) ``ds_bpermute`` is wave64-wide, so the
+# cross-half read is a single wide ``ds_bpermute`` (lane mask 63) and ``permlane64`` is unused; on RDNA ``ds_bpermute``
+# is SIMD32-scoped (lane mask 31) and the top half comes from ``amdgpu_permlane64`` -- the native ``v_permlane64_b32``
+# on gfx11+/gfx12, or an LDS-roundtrip emulation on gfx10.x.  ``v_permlane64_b32`` does not exist on CDNA and emitting
+# it there crashes the AMDGPU backend (genesis-world issue #2962), so the JIT must never emit it on gfx9xx.
 # --------------------------------------------------------------------------------------------------------------------