From 127f40da7f774cba61bd371b5be265ed29cf84fb Mon Sep 17 00:00:00 2001
From: Artturin <Artturin@artturin.com>
Date: Wed, 2 Oct 2024 22:37:51 +0300
Subject: [PATCH] umd: Update

---
 pkgs/umd/default.nix           | 105 +++++++++++++++++++++++++++------
 pkgs/umd/fmt_mystery.patch     |  25 --------
 pkgs/umd/missing_headers.patch |  24 --------
 3 files changed, 87 insertions(+), 67 deletions(-)
 delete mode 100644 pkgs/umd/fmt_mystery.patch
 delete mode 100644 pkgs/umd/missing_headers.patch

diff --git a/pkgs/umd/default.nix b/pkgs/umd/default.nix
index 7f0cb5a..f76fe8e 100644
--- a/pkgs/umd/default.nix
+++ b/pkgs/umd/default.nix
@@ -1,24 +1,73 @@
 { pkgs }:
 
+let
+  depsDir = "deps";
+
+  version = "unstable-2024-10-01";
+
+  # Update outputHash in umd-deps too
+  src = pkgs.fetchFromGitHub {
+    owner = "tenstorrent";
+    repo = "tt-umd";
+    rev = "5293e508ade90e758b386068770a882667a535f0";
+    hash = "sha256-Pa6vFjsG30UnGwLKZc/OTmiSMkXs/0sAFwVGojE8IbY=";
+  };
+
+  umd-deps = pkgs.stdenv.mkDerivation {
+    name = "tt-umd-deps-${version}.tar.gz";
+
+    inherit src;
+
+    dontBuild = true;
+
+    outputHash = "sha256-/T7UJ1OCd4T68DXOxHjGLYiiLoEZIRt7PHbT9npT4uk=";
+    outputHashAlgo = "sha256";
+
+    nativeBuildInputs = with pkgs; [
+      cmake
+      git
+      cacert
+      python3
+    ];
+
+    ARCH_NAME = "wormhole_b0";
+    cmakeFlags = [
+      "-DCPM_DOWNLOAD_ALL=ON"
+      "-DCPM_SOURCE_CACHE=${depsDir}"
+      "-DTT_UMD_BUILD_TESTS=ON"
+    ];
+
+    postPatch = ''
+      cp ${pkgs.cpm-cmake}/share/cpm/CPM.cmake cmake/CPM.cmake
+    '';
+
+    installPhase = ''
+      runHook preInstall
+
+      # Prune the `.git` directories
+      find ${depsDir} -name .git -type d -prune -exec rm -rf {} \;;
+      # Build a reproducible tar, per instructions at https://reproducible-builds.org/docs/archives/
+      tar --owner=0 --group=0 --numeric-owner --format=gnu \
+          --sort=name --mtime="@$SOURCE_DATE_EPOCH" \
+          -czf $out \
+            ${depsDir} \
+
+      runHook postInstall
+    '';
+  };
+in
+
 pkgs.stdenv.mkDerivation {
   pname = "tt-umd";
-  version = "main-2024-02-01";
+  inherit version src;
 
-  src = pkgs.fetchFromGitHub {
-    owner = "tenstorrent-metal";
-    repo = "umd";
-    rev = "341f5b7b299f128faaf2ca446a03298cb781a645";
-    hash = "sha256-jMxhhFWnCjNZZvFiTCeuEHvxvE0+IoaP4NJkr/CDLy8=";
-  };
-
-  patches = [
-    ./fmt_mystery.patch
-    ./missing_headers.patch
-  ];
-
-  makeFlags = [
-    "DEVICE_CXX=${pkgs.stdenv.cc.targetPrefix}c++"
-    "ARCH_NAME=grayskull"
+  nativeBuildInputs = with pkgs; [
+    cmake
+    git
+    cacert
+    ninja
+    python3
+    removeReferencesTo
   ];
 
   buildInputs = with pkgs; [
@@ -28,8 +77,28 @@ pkgs.stdenv.mkDerivation {
     hwloc
   ];
 
+  ARCH_NAME = "wormhole_b0";
+  cmakeFlags = [
+    "-DCPM_SOURCE_CACHE=${depsDir}"
+    # libdevice.so
+    # RUNPATH              /build/source/build/_deps/nanomsg-build:/build/source/build/_deps/libuv-build:/nix/store/n1yy5f1754p2d6dhksvg6rwpayymw1fx-tt-umd-unstable-2024-10-01/lib/:...
+    # TODO: look in to fixing properly if there's a need.
+    "-DCMAKE_SKIP_BUILD_RPATH=ON"
+  ];
+
+  postUnpack = ''
+    mkdir -p $sourceRoot/build
+    tar -xf ${umd-deps} -C $sourceRoot/build
+  '';
+
+  postPatch = ''
+    cp ${pkgs.cpm-cmake}/share/cpm/CPM.cmake cmake/CPM.cmake
+  '';
+
   installPhase = ''
-    mkdir $out
-    mv build/lib $out
+    runHook preInstall
+    mkdir -p $out/lib
+    mv lib/libdevice.so $out/lib
+    runHook postInstall
   '';
 }
diff --git a/pkgs/umd/fmt_mystery.patch b/pkgs/umd/fmt_mystery.patch
deleted file mode 100644
index 0c03402..0000000
--- a/pkgs/umd/fmt_mystery.patch
+++ /dev/null
@@ -1,25 +0,0 @@
-diff --git a/device/cpuset_lib.cpp b/device/cpuset_lib.cpp
-index f8d5e03..93e89b7 100644
---- a/device/cpuset_lib.cpp
-+++ b/device/cpuset_lib.cpp
-@@ -9,6 +9,7 @@
- #include <thread>
- #include "device/device_api.h"
- #include <filesystem>
-+#include <fmt/std.h>
- namespace tt {
- 
- namespace fs = std::filesystem;
-diff --git a/device/tt_silicon_driver.cpp b/device/tt_silicon_driver.cpp
-index fb243ed..c8950ae 100644
---- a/device/tt_silicon_driver.cpp
-+++ b/device/tt_silicon_driver.cpp
-@@ -4342,7 +4342,7 @@ void tt_SiliconDevice::set_power_state(tt_DevicePowerState device_state) {
-             set_pcie_power_state(device_state);
-         } else {
-             int exit_code = set_remote_power_state(chip, device_state);
--            log_assert(exit_code == 0, "Failed to set power state to {} with exit code: {}", device_state, exit_code);
-+            log_assert(exit_code == 0, "Failed to set power state with exit code: {}", exit_code);
-         }
-     }
- }
diff --git a/pkgs/umd/missing_headers.patch b/pkgs/umd/missing_headers.patch
deleted file mode 100644
index fb4e0b9..0000000
--- a/pkgs/umd/missing_headers.patch
+++ /dev/null
@@ -1,24 +0,0 @@
-diff --git a/device/tt_cluster_descriptor.h b/device/tt_cluster_descriptor.h
-index c5c066e..84f17ce 100644
---- a/device/tt_cluster_descriptor.h
-+++ b/device/tt_cluster_descriptor.h
-@@ -11,6 +11,7 @@
- 
- #include <unordered_map>
- #include <unordered_set>
-+#include <cstdint>
- #include <set>
- #include <map>
- #include <tuple>
-diff --git a/device/tt_soc_descriptor.cpp b/device/tt_soc_descriptor.cpp
-index adf77c2..3813eab 100644
---- a/device/tt_soc_descriptor.cpp
-+++ b/device/tt_soc_descriptor.cpp
-@@ -5,6 +5,7 @@
- #include "tt_soc_descriptor.h"
- 
- #include <assert.h>
-+#include <cstdint>
- #include <fstream>
- #include <iostream>
- #include <regex>