From 4b39362cefd996597df7bf192c65b8738f8e8e25 Mon Sep 17 00:00:00 2001
From: Vincent Tang <vtang@tenstorrent.com>
Date: Thu, 8 Feb 2024 18:26:31 +0000
Subject: [PATCH] #5222: switching hex8_to_hex32 from py to cpp

---
 tt_metal/hw/toolchain/hex8tohex32.py | 41 ---------------------
 tt_metal/jit_build/build.cpp         | 55 ++++++++++++++++++++++++----
 tt_metal/jit_build/build.hpp         |  1 -
 3 files changed, 48 insertions(+), 49 deletions(-)
 delete mode 100644 tt_metal/hw/toolchain/hex8tohex32.py

diff --git a/tt_metal/hw/toolchain/hex8tohex32.py b/tt_metal/hw/toolchain/hex8tohex32.py
deleted file mode 100644
index 8e0665f4fdc..00000000000
--- a/tt_metal/hw/toolchain/hex8tohex32.py
+++ /dev/null
@@ -1,41 +0,0 @@
-# SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
-
-# SPDX-License-Identifier: Apache-2.0
-
-#!/usr/bin/env python3
-
-import itertools
-import sys
-
-
-def write_data(outf, data, ptr):
-    if len(data) != 0:
-        outf.write("@%08x\n" % (ptr >> 2))
-        while len(data) % 4 != 0:
-            data.append(0)
-        for word_bytes in zip(*([iter(data)] * 4)):
-            outf.write("".join(["%02x" % b for b in reversed(word_bytes)]) + "\n")
-
-
-assert len(sys.argv) > 2
-in_lines = open(sys.argv[1]).readlines()
-data = []
-ptr = 0
-outf = open(sys.argv[2], "wt")
-for line in in_lines:
-    if line.startswith("@"):
-        addr = int(line[1:], 16)
-        if addr > ptr + 4:
-            write_data(outf, data, ptr)
-            ptr = addr
-            data = []
-            while ptr % 4 != 0:
-                data.append(0)
-                ptr -= 1
-        else:
-            while ptr + len(data) < addr:
-                data.append(0)
-    else:
-        data += [int(tok, 16) for tok in line.split()]
-
-write_data(outf, data, ptr)
diff --git a/tt_metal/jit_build/build.cpp b/tt_metal/jit_build/build.cpp
index dfb32fec39f..76a4016d92b 100644
--- a/tt_metal/jit_build/build.cpp
+++ b/tt_metal/jit_build/build.cpp
@@ -6,6 +6,10 @@
 #include <filesystem>
 #include <thread>
 #include <string>
+#include <iostream>
+#include <fstream>
+#include <sstream>
+#include <iomanip>
 
 #include "jit_build/build.hpp"
 #include "jit_build/genfiles.hpp"
@@ -46,7 +50,6 @@ void JitBuildEnv::init(uint32_t device_id, tt::ARCH arch)
     // Tools
     this->gpp_ = this->root_ + "tt_metal/third_party/sfpi/compiler/bin/riscv32-unknown-elf-g++ ";
     this->objcopy_ = this->root_ + "tt_metal/third_party/sfpi/compiler/bin/riscv32-unknown-elf-objcopy ";
-    this->hex8tohex32_ = string("python3 ") + this->root_ + "tt_metal/hw/toolchain/hex8tohex32.py ";
 
     // Flags
     string common_flags;
@@ -481,14 +484,52 @@ void JitBuildState::elf_to_hex8(const string& log_file, const string& out_dir) c
 
 void JitBuildState::hex8_to_hex32(const string& log_file, const string& out_dir) const
 {
-    string cmd;
-    cmd = "cd " + out_dir + " && ";
-    cmd += env_.hex8tohex32_ + this->target_name_ + ".hex.tmp " + this->target_name_ + ".hex";
+    ZoneScoped;
+    auto write_data = [](std::ofstream& outf, std::vector<uint64_t>& data, uint64_t& ptr){
+        if (!data.empty()) {
+            outf << "@" << std::setfill('0') << std::setw(8) << std::hex << (ptr >> 2) << "\n";
+            for (size_t i = 0; i < data.size(); i += 4) {
+                for (int j = 3; j >= 0; --j) {
+                    outf << std::setfill('0') << std::setw(2) << std::hex << data[i + j];
+                }
+                outf << "\n";
+            }
+        }
+        data.clear();
+    };
 
-    log_debug(tt::LogBuildKernels, "    hex8tohex32 cmd: {}", cmd);
-    if (!tt::utils::run_command(cmd, log_file, false)) {
-        build_failure(this->target_name_, "hex8tohex32.py", cmd, log_file);
+    auto pad_zeroes = [](std::vector<uint64_t>& data, uint32_t num){
+        for(unsigned int i = 0; i < num; i++){
+            data.push_back(0);
+        }
+    };
+
+    std::ifstream inf(out_dir + this->target_name_ + ".hex.tmp");
+    std::ofstream outf(out_dir + this->target_name_ + ".hex");
+    std::string line;
+    std::vector<uint64_t> data;
+    uint64_t ptr = 0;
+
+    while (std::getline(inf, line)) {
+        if (line[0] == '@') {
+            uint64_t addr = std::stol(line.substr(1), nullptr, 16);
+            if (addr > ptr + 4) {
+                write_data(outf, data, ptr);
+                ptr = addr;
+                pad_zeroes(data, (ptr % 4));
+                ptr -= ptr % 4;
+            } else {
+                pad_zeroes(data, (addr - ptr - data.size()));
+            }
+        } else {
+            std::istringstream iss(line);
+            std::string tok;
+            while (iss >> tok) {
+                data.push_back(std::stol(tok, nullptr, 16));
+            }
+        }
     }
+    write_data(outf, data, ptr);
 }
 
 // Given this elf (A) and a later elf (B):
diff --git a/tt_metal/jit_build/build.hpp b/tt_metal/jit_build/build.hpp
index 636e71c9f79..f31b60204d9 100644
--- a/tt_metal/jit_build/build.hpp
+++ b/tt_metal/jit_build/build.hpp
@@ -59,7 +59,6 @@ class JitBuildEnv {
     // Tools
     string gpp_;
     string objcopy_;
-    string hex8tohex32_;
 
     // Compilation options
     string cflags_;