feat(rocjitsu): cdna5 OOB behavior implementation

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna5/addr_calc.cpp

@@ -4,7 +4,6 @@
 #include "rocjitsu/isa/arch/amdgpu/cdna5/addr_calc.h"
 #include "rocjitsu/isa/arch/amdgpu/generated/cdna5/operand.h"
 #include "rocjitsu/isa/arch/amdgpu/generated/cdna5/operand_types.h"
-#include "rocjitsu/isa/arch/amdgpu/shared/addr_calc_buffer.h"
 #include "rocjitsu/vm/amdgpu/compute_unit.h"
 #include "rocjitsu/vm/amdgpu/lds.h"
 #include "rocjitsu/vm/amdgpu/mem_state.h"
@@ -21,6 +20,8 @@
 namespace rocjitsu::cdna5 {
 namespace {
 
+constexpr uint64_t kBufferOffsetMask = (uint64_t{1} << 45) - 1;
+
 uint32_t scaled_vaddr_factor(const amdgpu::VectorMemState &d) {
   // LLVM folds gfx1250 scale_offset when the scale matches the full memory
   // access size. The encoded immediate offset remains a byte offset.
@@ -64,18 +65,6 @@ uint32_t resolved_vgpr_base(const amdgpu::Wavefront &wf, uint32_t operand,
          *Isa::resolved_vgpr_offset(wf, OperandType::OPR_VGPR, operand, role);
 }
 
-uint64_t buffer_base_addr(uint32_t srd0, uint32_t srd1) {
-  return (static_cast<uint64_t>(srd1 & 0x01FF'FFFFu) << 32) | srd0;
-}
-
-uint64_t buffer_num_records(uint32_t srd1, uint32_t srd2, uint32_t srd3) {
-  return ((static_cast<uint64_t>(srd3 & 0x3Fu) << 32) | srd2) << 7 | ((srd1 >> 25) & 0x7Fu);
-}
-
-uint32_t buffer_stride(uint32_t srd3) { return (srd3 >> 12) & 0xFFFFu; }
-
-bool buffer_oob_raw(uint32_t srd3) { return (srd3 & 0x8000'0000u) != 0; }
-
 void init_vector_mem_state(amdgpu::Wavefront &wf, amdgpu::VectorMemState &d) {
   uint64_t exec = wf.exec();
   d.lane_mask = exec;
@@ -86,6 +75,17 @@ void init_vector_mem_state(amdgpu::Wavefront &wf, amdgpu::VectorMemState &d) {
   d.cu_path = wf.cu().full_path();
 }
 
+int64_t logical_buffer_offset(uint32_t index, uint32_t stride, uint32_t voffset, int32_t ioffset) {
+  return static_cast<int64_t>(index) * stride + voffset + ioffset;
+}
+
+bool byte_range_exceeds(int64_t offset, uint32_t size, uint64_t bound) {
+  if (offset < 0)
+    return true;
+  uint64_t unsigned_offset = static_cast<uint64_t>(offset);
+  return unsigned_offset > bound || size > bound - unsigned_offset;
+}
+
 bool decode_flat_private_address(amdgpu::Wavefront &wf, uint64_t addr, uint64_t *translated) {
   uint32_t lane_stride = wf.scratch_lane_size();
   if (lane_stride == 0)
@@ -150,6 +150,21 @@ void flat_global_calculate_addresses(const Inst &inst, amdgpu::Wavefront &wf,
 
 } // namespace
 
+BufferResource decode_buffer_resource(uint32_t srd0, uint32_t srd1, uint32_t srd2, uint32_t srd3) {
+  static constexpr uint32_t stride_multipliers[] = {1, 4, 8, 32};
+  BufferResource resource{};
+  resource.base_address = (static_cast<uint64_t>(srd1 & 0x01FF'FFFFu) << 32) | srd0;
+  resource.num_records =
+      ((static_cast<uint64_t>(srd3 & 0x3Fu) << 32) | srd2) << 7 | ((srd1 >> 25) & 0x7Fu);
+  resource.raw_stride = (srd3 >> 12) & 0x3FFFu;
+  resource.stride_scale = static_cast<uint8_t>((srd3 >> 26) & 0x3u);
+  resource.stride = resource.raw_stride * stride_multipliers[resource.stride_scale];
+  resource.swizzle_enabled = ((srd3 >> 28) & 0x1u) != 0;
+  resource.oob_select = ((srd3 >> 29) & 0x1u) != 0;
+  resource.type = static_cast<uint8_t>((srd3 >> 30) & 0x3u);
+  return resource;
+}
+
 uint64_t smem_calculate_address(const SmemMachineInst &inst, amdgpu::Wavefront &wf,
                                 uint32_t access_size_bytes) {
   auto &cu = wf.cu();
@@ -233,12 +248,9 @@ void mubuf_calculate_addresses(const VbufferMachineInst &inst, amdgpu::Wavefront
   uint32_t srd1 = amdgpu::RegisterAccess(cu).read_sgpr(sb + 1);
   uint32_t srd2 = amdgpu::RegisterAccess(cu).read_sgpr(sb + 2);
   uint32_t srd3 = amdgpu::RegisterAccess(cu).read_sgpr(sb + 3);
-  uint64_t base_addr = buffer_base_addr(srd0, srd1);
-  uint64_t num_records = buffer_num_records(srd1, srd2, srd3);
-  uint32_t stride = buffer_stride(srd3);
-  bool oob_raw = buffer_oob_raw(srd3);
+  BufferResource resource = decode_buffer_resource(srd0, srd1, srd2, srd3);
   uint32_t soffset_val = has_smem_offset(inst.soffset) ? read_sreg_m0_operand(wf, inst.soffset) : 0;
-  int64_t ioff = static_cast<int64_t>(signed_ioffset(inst.ioffset));
+  int32_t ioff = signed_ioffset(inst.ioffset);
   uint32_t vbase = 0;
   if (inst.idxen || inst.offen)
     vbase = resolved_vgpr_base(wf, inst.vaddr, amdgpu::VgprMsbRole::Src0);
@@ -248,6 +260,18 @@ void mubuf_calculate_addresses(const VbufferMachineInst &inst, amdgpu::Wavefront
     uint32_t reg_count = (inst.idxen && inst.offen) ? 2 : 1;
     vaddr_region.emplace(regs.read_vgpr_region(vbase, reg_count, exec));
   }
+  assert(d.elem_size != 0 && d.num_elems != 0);
+  d.element_lane_masks.clear();
+  if (ioff < 0) {
+    // CDNA5 requires VBUFFER IOFFSET to be non-negative. Suppress illegal
+    // encodings defensively so their wrapped 45-bit address cannot reach L1.
+    d.element_lane_masks.assign(d.num_elems, 0);
+    d.lane_mask = 0;
+    return;
+  }
+  if (!resource.swizzle_enabled)
+    d.element_lane_masks.assign(d.num_elems, exec);
+  d.lane_mask = 0;
   for (uint32_t lane = 0; lane < wf.wf_size(); ++lane) {
     if (!(exec & (1ULL << lane)))
       continue;
@@ -261,21 +285,34 @@ void mubuf_calculate_addresses(const VbufferMachineInst &inst, amdgpu::Wavefront
     } else if (inst.offen) {
       voffset = vaddr_region->lane(0, lane);
     }
-    uint32_t offset_part = amdgpu::addr_calc::buffer_offset_part(voffset, ioff);
-    uint64_t total_offset =
-        amdgpu::addr_calc::buffer_total_offset(index, stride, offset_part, soffset_val);
-    bool oob;
-    if (!oob_raw && stride > 0) {
-      oob = index >= num_records;
-    } else {
-      oob = offset_part >= num_records;
+    int64_t logical_offset = logical_buffer_offset(index, resource.stride, voffset, ioff);
+    uint64_t buffer_offset = static_cast<uint64_t>(logical_offset) & kBufferOffsetMask;
+    d.per_lane_addr[lane] = resource.base_address + soffset_val + buffer_offset;
+
+    if (resource.swizzle_enabled) {
+      // Swizzled address generation is not modeled yet. Preserve the existing
+      // linear address behavior, but do not apply unswizzled bounds rules.
+      d.lane_mask |= uint64_t{1} << lane;
+      continue;
     }
-    if (oob) {
-      d.lane_mask &= ~(1ULL << lane);
-      d.per_lane_addr[lane] = 0;
-    } else {
-      d.per_lane_addr[lane] = (base_addr + total_offset) & 0xFFFFFFFFFFFFULL;
+
+    int64_t record_offset = static_cast<int64_t>(soffset_val) + voffset + ioff;
+    int64_t total_offset = static_cast<int64_t>(soffset_val) + logical_offset;
+    for (uint32_t elem = 0; elem < d.num_elems; ++elem) {
+      int64_t component_offset = static_cast<int64_t>(elem) * d.elem_size;
+      bool oob =
+          byte_range_exceeds(total_offset + component_offset, d.elem_size, resource.num_records);
+      if (resource.oob_select) {
+        oob = oob ||
+              byte_range_exceeds(record_offset + component_offset, d.elem_size, resource.stride);
+      }
+      if (oob)
+        d.element_lane_masks[elem] &= ~(uint64_t{1} << lane);
+      else
+        d.lane_mask |= uint64_t{1} << lane;
     }
+    if (!(d.lane_mask & (uint64_t{1} << lane)))
+      d.per_lane_addr[lane] = 0;
   }
 }
 

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna5/addr_calc.h

@@ -18,6 +18,20 @@ struct VectorMemState;
 
 namespace rocjitsu::cdna5 {
 
+/// @brief Decoded CDNA5 buffer resource descriptor fields.
+struct BufferResource {
+  uint64_t base_address = 0;
+  uint64_t num_records = 0;
+  uint32_t raw_stride = 0;
+  uint32_t stride = 0;
+  uint8_t stride_scale = 0;
+  bool swizzle_enabled = false;
+  bool oob_select = false;
+  uint8_t type = 0;
+};
+
+BufferResource decode_buffer_resource(uint32_t srd0, uint32_t srd1, uint32_t srd2, uint32_t srd3);
+
 /// @brief Sign-extend a CDNA5 24-bit IOFFSET field.
 inline int32_t signed_ioffset(uint32_t ioffset) { return static_cast<int32_t>(ioffset << 8) >> 8; }
 

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/l1_vector_cache.cpp

@@ -43,6 +43,23 @@ uint32_t for_each_coalesced_lane_run(const uint64_t *addrs, uint64_t lane_mask,
   return run_count;
 }
 
+bool all_elements_use_lane_mask(std::span<const uint64_t> element_lane_masks, uint64_t lane_mask,
+                                uint32_t num_elems) {
+  if (element_lane_masks.empty())
+    return true;
+  assert(element_lane_masks.size() == num_elems);
+  (void)num_elems;
+  return std::ranges::all_of(element_lane_masks,
+                             [lane_mask](uint64_t mask) { return mask == lane_mask; });
+}
+
+uint64_t fully_valid_lane_mask(std::span<const uint64_t> element_lane_masks, uint64_t lane_mask) {
+  uint64_t full_lane_mask = lane_mask;
+  for (uint64_t element_mask : element_lane_masks)
+    full_lane_mask &= element_mask;
+  return full_lane_mask;
+}
+
 } // namespace
 
 L1VectorCache::L1VectorCache(L2Cache *l2)
@@ -193,10 +210,30 @@ void L1VectorCache::write_bytes(uint64_t addr, const uint8_t *src, uint32_t size
 
 void L1VectorCache::load(const uint64_t *addrs, uint64_t lane_mask, uint32_t elem_size,
                          uint32_t num_elems, uint8_t *dst, Mtype mtype, bool non_temporal,
-                         bool request_l1_bypass, uint32_t wf_size, uint32_t vmid) {
+                         bool request_l1_bypass, uint32_t wf_size, uint32_t vmid,
+                         std::span<const uint64_t> element_lane_masks) {
   auto coherence_guard = DeviceCacheCoherence::instance().acquire_l1_access();
   synchronize_epoch_locked();
   uint32_t stride = num_elems * elem_size;
+  if (!all_elements_use_lane_mask(element_lane_masks, lane_mask, num_elems)) {
+    uint64_t full_lane_mask = fully_valid_lane_mask(element_lane_masks, lane_mask);
+    for_each_coalesced_lane_run(
+        addrs, full_lane_mask, wf_size, stride, [&](uint32_t first_lane, uint32_t run_lanes) {
+          read_bytes(addrs[first_lane], dst + first_lane * stride, run_lanes * stride, mtype,
+                     non_temporal, request_l1_bypass, vmid);
+        });
+    for (uint32_t elem = 0; elem < num_elems; ++elem) {
+      uint64_t mask = element_lane_masks[elem] & lane_mask & ~full_lane_mask;
+      while (mask) {
+        const uint32_t lane = std::countr_zero(mask);
+        mask &= ~(uint64_t{1} << lane);
+        read_bytes(addrs[lane] + static_cast<uint64_t>(elem) * elem_size,
+                   dst + lane * stride + elem * elem_size, elem_size, mtype, non_temporal,
+                   request_l1_bypass, vmid);
+      }
+    }
+    return;
+  }
   for_each_coalesced_lane_run(
       addrs, lane_mask, wf_size, stride, [&](uint32_t first_lane, uint32_t run_lanes) {
         read_bytes(addrs[first_lane], dst + first_lane * stride, run_lanes * stride, mtype,
@@ -206,14 +243,34 @@ void L1VectorCache::load(const uint64_t *addrs, uint64_t lane_mask, uint32_t ele
 
 void L1VectorCache::store(const uint64_t *addrs, uint64_t lane_mask, uint32_t elem_size,
                           uint32_t num_elems, const uint8_t *src, Mtype mtype, bool non_temporal,
-                          uint32_t wf_size, uint32_t vmid) {
+                          uint32_t wf_size, uint32_t vmid,
+                          std::span<const uint64_t> element_lane_masks) {
   auto coherence_guard = DeviceCacheCoherence::instance().acquire_l1_access();
   synchronize_epoch_locked();
   uint32_t stride = num_elems * elem_size;
   const uint32_t active_lanes = std::popcount(lane_mask);
   ++store_count_;
   if (active_lanes > 0)
     ++store_active_count_;
+  if (!all_elements_use_lane_mask(element_lane_masks, lane_mask, num_elems)) {
+    uint64_t full_lane_mask = fully_valid_lane_mask(element_lane_masks, lane_mask);
+    store_l2_writes_ += for_each_coalesced_lane_run(
+        addrs, full_lane_mask, wf_size, stride, [&](uint32_t first_lane, uint32_t run_lanes) {
+          write_bytes(addrs[first_lane], src + first_lane * stride, run_lanes * stride, mtype,
+                      non_temporal, vmid);
+        });
+    for (uint32_t elem = 0; elem < num_elems; ++elem) {
+      uint64_t mask = element_lane_masks[elem] & lane_mask & ~full_lane_mask;
+      store_l2_writes_ += std::popcount(mask);
+      while (mask) {
+        const uint32_t lane = std::countr_zero(mask);
+        mask &= ~(uint64_t{1} << lane);
+        write_bytes(addrs[lane] + static_cast<uint64_t>(elem) * elem_size,
+                    src + lane * stride + elem * elem_size, elem_size, mtype, non_temporal, vmid);
+      }
+    }
+    return;
+  }
   store_l2_writes_ += for_each_coalesced_lane_run(
       addrs, lane_mask, wf_size, stride, [&](uint32_t first_lane, uint32_t run_lanes) {
         write_bytes(addrs[first_lane], src + first_lane * stride, run_lanes * stride, mtype,

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/l1_vector_cache.h

@@ -8,6 +8,7 @@
 #include "simdojo/components/cache.h"
 
 #include <cstdint>
+#include <span>
 
 namespace rocjitsu {
 namespace amdgpu {
@@ -41,13 +42,20 @@ class L1VectorCache {
 
   void set_l2(L2Cache *l2);
   void set_memory(GpuMemory *mem);
+
+  /// @param element_lane_masks Empty when every element uses @p lane_mask;
+  /// otherwise contains exactly @p num_elems masks. In the latter form,
+  /// @p lane_mask is the union of lanes valid for at least one element.
   void load(const uint64_t *addrs, uint64_t lane_mask, uint32_t elem_size, uint32_t num_elems,
             uint8_t *dst, Mtype mtype, bool non_temporal, bool request_l1_bypass, uint32_t wf_size,
-            uint32_t vmid = 0);
+            uint32_t vmid = 0, std::span<const uint64_t> element_lane_masks = {});
 
+  /// @param element_lane_masks Empty when every element uses @p lane_mask;
+  /// otherwise contains exactly @p num_elems masks. In the latter form,
+  /// @p lane_mask is the union of lanes valid for at least one element.
   void store(const uint64_t *addrs, uint64_t lane_mask, uint32_t elem_size, uint32_t num_elems,
              const uint8_t *src, Mtype mtype, bool non_temporal, uint32_t wf_size,
-             uint32_t vmid = 0);
+             uint32_t vmid = 0, std::span<const uint64_t> element_lane_masks = {});
 
   void invalidate(uint64_t addr, uint32_t vmid = 0);
   void invalidate_all();

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/mem_state.h

@@ -91,6 +91,10 @@ struct VectorMemState : DynamicInstState {
   }
   std::array<uint64_t, 64> per_lane_addr = {};
   uint64_t lane_mask = 0;
+  /// Optional per-element lane validity. Empty means every element uses
+  /// lane_mask; otherwise the vector contains exactly num_elems masks and
+  /// lane_mask is their union.
+  std::vector<uint64_t> element_lane_masks;
   uint64_t exec_mask = 0; ///< EXEC mask at issue time. Set by addr calc functions.
                           ///< Writeback zeroes OOB lanes (exec_mask & ~lane_mask).
   uint32_t wf_size = 64;  ///< Wavefront width (set from wavefront's wf_size()).

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/memory_pipeline.cpp

@@ -493,12 +493,13 @@ void GlobalMemPipeline::initiate_access(Instruction &inst, Wavefront &wf) {
   }
 
   if (d.is_load) {
-    d.response_data.resize(d.wf_size * d.num_elems * d.elem_size);
+    d.response_data.assign(d.wf_size * d.num_elems * d.elem_size, 0);
     l1_->load(d.per_lane_addr.data(), d.lane_mask, d.elem_size, d.num_elems, d.response_data.data(),
-              d.mtype, d.non_temporal, d.request_force_l1_bypass, d.wf_size, wf.process_id());
+              d.mtype, d.non_temporal, d.request_force_l1_bypass, d.wf_size, wf.process_id(),
+              d.element_lane_masks);
   } else {
     l1_->store(d.per_lane_addr.data(), d.lane_mask, d.elem_size, d.num_elems, d.store_data.data(),
-               d.mtype, d.non_temporal, d.wf_size, wf.process_id());
+               d.mtype, d.non_temporal, d.wf_size, wf.process_id(), d.element_lane_masks);
   }
 }