Unify per-lane address decoding

gpu-simulator/trace-parser/trace_parser.cc

@@ -97,47 +97,6 @@ kernel_trace_t::kernel_trace_t(const std::string &filePath) {
 
 kernel_trace_t::~kernel_trace_t() { delete tracez_reader; }
 
-void inst_memadd_info_t::base_stride_decompress(
-    unsigned long long base_address, int stride,
-    const std::bitset<WARP_SIZE> &mask) {
-  bool first_bit1_found = false;
-  bool last_bit1_found = false;
-  unsigned long long addra = base_address;
-  for (int s = 0; s < WARP_SIZE; s++) {
-    if (mask.test(s) && !first_bit1_found) {
-      first_bit1_found = true;
-      addrs[s] = base_address;
-    } else if (first_bit1_found && !last_bit1_found) {
-      if (mask.test(s)) {
-        addra += stride;
-        addrs[s] = addra;
-      } else
-        last_bit1_found = true;
-    } else
-      addrs[s] = 0;
-  }
-}
-
-void inst_memadd_info_t::base_delta_decompress(
-    unsigned long long base_address, const std::vector<long long> &deltas,
-    const std::bitset<WARP_SIZE> &mask) {
-  bool first_bit1_found = false;
-  long long last_address = 0;
-  unsigned delta_index = 0;
-  for (int s = 0; s < 32; s++) {
-    if (mask.test(s) && !first_bit1_found) {
-      addrs[s] = base_address;
-      first_bit1_found = true;
-      last_address = base_address;
-    } else if (mask.test(s) && first_bit1_found) {
-      assert(delta_index < deltas.size());
-      addrs[s] = last_address + deltas[delta_index++];
-      last_address = addrs[s];
-    } else
-      addrs[s] = 0;
-  }
-}
-
 void tma_inst_memaddr_info_t::base_delta_decompress(
     unsigned long long base_address, const std::vector<long long> &deltas,
     const std::bitset<WARP_SIZE> &mask) {
@@ -284,46 +243,8 @@ bool inst_trace_t::parse_from_string(std::string trace, unsigned trace_version,
     memadd_info->width = get_datawidth_from_opcode(opcode_tokens);
 
     ss >> std::dec >> address_mode;
-    if (address_mode == address_format::list_all) {
-      // read addresses one by one from the file
-      for (int s = 0; s < WARP_SIZE; s++) {
-        if (mask_bits.test(s))
-          ss >> std::hex >> memadd_info->addrs[s];
-        else
-          memadd_info->addrs[s] = 0;
-      }
-    } else if (address_mode == address_format::base_stride) {
-      // read addresses as base address and stride
-      unsigned long long base_address = 0;
-      int stride = 0;
-      ss >> std::hex >> base_address;
-      ss >> std::dec >> stride;
-      memadd_info->base_stride_decompress(base_address, stride, mask_bits);
-    } else if (address_mode == address_format::base_delta) {
-      unsigned long long base_address = 0;
-      std::vector<long long> deltas;
-      // read addresses as base address and deltas
-      ss >> std::hex >> base_address;
-      // No delta will be pushed if there is only one bit set in the mask
-      // Find the first bit set
-      int first_bit_set = -1;
-      for (int s = 0; s < WARP_SIZE; s++) {
-        if (mask_bits.test(s)) {
-          first_bit_set = s;
-          break;
-        }
-      }
-      for (int s = first_bit_set + 1; s < WARP_SIZE; s++) {
-        if (mask_bits.test(s)) {
-          long long delta = 0;
-          ss >> std::dec >> delta;
-          deltas.push_back(delta);
-        }
-      }
-      memadd_info->base_delta_decompress(base_address, deltas, mask_bits);
-    } else {
-      assert(0 && "Unsupported address mode");
-    }
+    std::vector<uint64_t> addrs = decode(address_mode, ss, mask_bits);
+    std::memcpy(memadd_info->addrs, addrs.data(), WARP_SIZE * sizeof(uint64_t));
   }
 
   ss >> imm;

gpu-simulator/trace-parser/trace_parser.h

@@ -13,7 +13,8 @@
 #include <string>
 #include <vector>
 
-#define WARP_SIZE 32
+#include "../../util/tracer_nvbit/tracer_tool/decode.h"
+
 #define MAX_DST 1
 #define MAX_SRC 4
 
@@ -31,15 +32,6 @@ enum address_scope {
   SYS_MEM,
 };
 
-// TODO Fix this shared enum with tracer tool
-enum address_format {
-  list_all = 0,
-  base_stride = 1,
-  base_delta = 2,
-  tma_list_all = 3,
-  tma_base_delta = 4
-};
-
 struct trace_command {
   std::string command_string;
   command_type m_type;
@@ -48,12 +40,6 @@ struct trace_command {
 struct inst_memadd_info_t {
   uint64_t addrs[WARP_SIZE];
   int32_t width;
-
-  void base_stride_decompress(unsigned long long base_address, int stride,
-                              const std::bitset<WARP_SIZE> &mask);
-  void base_delta_decompress(unsigned long long base_address,
-                             const std::vector<long long> &deltas,
-                             const std::bitset<WARP_SIZE> &mask);
 };
 
 struct tma_inst_memaddr_info_t {

util/tracer_nvbit/tracer_tool/decode.h

@@ -0,0 +1,117 @@
+#pragma once
+
+#include <bitset>
+#include <cassert>
+#include <cstdint>
+#include <sstream>
+#include <stdexcept>
+#include <vector>
+
+constexpr int WARP_SIZE = 32;
+
+enum address_format {
+  list_all = 0,
+  base_stride = 1,
+  base_delta = 2,
+  tma_list_all = 3,
+  tma_base_delta = 4
+};
+
+inline std::vector<uint64_t> decode_list_all(std::stringstream &ss,
+                                             std::bitset<WARP_SIZE> mask_bits) {
+  std::vector<uint64_t> addrs;
+  uint64_t addr = 0;
+
+  // read addresses one by one from the file
+  for (int s = 0; s < WARP_SIZE; s++) {
+    if (mask_bits.test(s)) {
+      ss >> std::hex >> addr;
+      if (ss.fail())
+        throw std::runtime_error("String stream input failed prematurely");
+      addrs.push_back(addr);
+    } else
+      addrs.push_back(0);
+  }
+  assert(addrs.size() == WARP_SIZE);
+  return addrs;
+}
+
+inline std::vector<uint64_t>
+decode_base_stride(std::stringstream &ss, std::bitset<WARP_SIZE> mask_bits) {
+  std::vector<uint64_t> addrs;
+  uint64_t addr = 0;
+
+  // read addresses as base address and stride
+  int stride = 0;
+  ss >> std::hex >> addr;
+  ss >> std::dec >> stride;
+  if (ss.fail())
+    throw std::runtime_error("String stream input failed prematurely");
+  bool first_bit1_found = false;
+  bool last_bit1_found = false;
+
+  for (int s = 0; s < WARP_SIZE; s++) {
+    if (!first_bit1_found && mask_bits.test(s)) {
+      first_bit1_found = true;
+      addrs.push_back(addr);
+    } else if (first_bit1_found && !last_bit1_found) {
+      if (mask_bits.test(s)) {
+        addr += stride;
+        addrs.push_back(addr);
+      } else {
+        last_bit1_found = true;
+        addrs.push_back(0);
+      }
+    } else
+      addrs.push_back(0);
+  }
+  assert(addrs.size() == WARP_SIZE);
+  return addrs;
+}
+
+inline std::vector<uint64_t>
+decode_base_delta(std::stringstream &ss, std::bitset<WARP_SIZE> mask_bits) {
+  std::vector<uint64_t> addrs;
+  uint64_t addr = 0;
+  int64_t delta = 0;
+
+  // read addresses as base address and deltas
+  ss >> std::hex >> addr;
+  if (ss.fail())
+    throw std::runtime_error("String stream input failed prematurely");
+  bool first_bit1_found = false;
+  for (int s = 0; s < WARP_SIZE; s++) {
+    if (!first_bit1_found && mask_bits.test(s)) {
+      first_bit1_found = true;
+      addrs.push_back(addr);
+    } else if (first_bit1_found && mask_bits.test(s)) {
+      ss >> std::dec >> delta;
+      if (ss.fail())
+        throw std::runtime_error("String stream input failed prematurely");
+      addr += delta;
+      addrs.push_back(addr);
+    } else
+      addrs.push_back(0);
+  }
+  assert(addrs.size() == WARP_SIZE);
+  return addrs;
+}
+
+inline std::vector<uint64_t> decode(unsigned address_mode,
+                                    std::stringstream &ss,
+                                    std::bitset<WARP_SIZE> mask_bits) {
+
+  if (address_mode == 0)
+    return decode_list_all(ss, mask_bits);
+  if (address_mode == 1)
+    return decode_base_stride(ss, mask_bits);
+  if (address_mode == 2)
+    return decode_base_delta(ss, mask_bits);
+  if (address_mode == 3 || address_mode == 4)
+    throw std::runtime_error(
+        "decode() does not support TMA addressing modes as they are written "
+        "usage-specific; reference usage-specific code at other callsites (for "
+        "instance, trace_parser.cc or post-traces-processing.cpp) and do not "
+        "invoke decode() here");
+  throw std::runtime_error("Unsupported address mode");
+}

util/tracer_nvbit/tracer_tool/tracer_tool.cu

@@ -38,6 +38,7 @@
 #include "watchdog.h"
 
 /* for general utils functions */
+#include "decode.h"
 #include "utils.h"
 
 #define TRACER_VERSION "5"
@@ -252,14 +253,6 @@ bool should_trace_kernel(uint64_t kernel_id, const std::string &kernel_name) {
   return false;
 }
 
-enum address_format {
-  list_all = 0,
-  base_stride = 1,
-  base_delta = 2,
-  tma_list_all = 3,
-  tma_base_delta = 4
-};
-
 /* File pointers for the kernels, and stats files */
 static FILE *kernelsFile = NULL;
 static FILE *statsFile = NULL;

util/tracer_nvbit/tracer_tool/traces-processing/post-traces-processing.cpp

@@ -22,6 +22,7 @@
 #include <zdict.h>
 #include <zstd.h>
 
+#include "../decode.h"
 #include "../tracez_format.h"
 
 using namespace std;
@@ -600,43 +601,11 @@ static mbar_inst_info classify_mbar_inst(const string &line,
   } else {
     unsigned addr_mode;
     ss >> dec >> addr_mode;
-    int first_lane = -1;
-    for (int s = 0; s < 32 && first_lane < 0; s++)
-      if ((mask >> s) & 1u)
-        first_lane = s;
-    if (first_lane < 0)
-      return r;
-    if (addr_mode == 0) { // list_all: one address per active lane
-      for (int s = first_lane; s < 32; s++) {
-        if (!((mask >> s) & 1u))
-          continue;
-        uint64_t addr = 0;
-        ss >> hex >> addr;
-        addrs.push_back(addr);
-      }
-    } else if (addr_mode == 1) { // base_stride: one contiguous run of lanes
-      uint64_t base = 0;
-      long long stride = 0;
-      ss >> hex >> base;
-      ss >> dec >> stride;
-      uint64_t addr = base;
-      for (int s = first_lane; s < 32 && ((mask >> s) & 1u); s++) {
-        addrs.push_back(addr);
-        addr += stride;
-      }
-    } else if (addr_mode == 2) { // base_delta: delta from the previous lane
-      uint64_t addr = 0;
-      ss >> hex >> addr;
-      addrs.push_back(addr);
-      for (int s = first_lane + 1; s < 32; s++) {
-        if (!((mask >> s) & 1u))
-          continue;
-        long long delta = 0;
-        ss >> dec >> delta;
-        addr += delta;
-        addrs.push_back(addr);
-      }
-    } else {
+    try {
+      addrs = decode(addr_mode, ss, mask);
+    } catch (const std::exception &e) {
+      fprintf(stderr, "Warning: %s (mode %u), classifying as OTHER\n", e.what(),
+              addr_mode);
       return r;
     }
   }