From 6162d21a5bce7bc31b9a6748304e8f954fdf37f2 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sun, 26 Dec 2021 21:11:14 +0800
Subject: [PATCH] LLVM PGO support

---
 flake.nix            | 61 ++++++++++++++++++++++++++++++++++++++++++++
 llvm/default.nix     |  3 ++-
 nac3artiq/Cargo.toml |  3 +++
 nac3artiq/src/lib.rs | 10 ++++++++
 4 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/flake.nix b/flake.nix
index 7c860146f..1e37b47d4 100644
--- a/flake.nix
+++ b/flake.nix
@@ -66,6 +66,67 @@
               '';
           }
         );
+
+        # LLVM PGO support
+        llvm-nac3-instrumented = pkgs.callPackage "${self}/llvm" {
+          stdenv = pkgs.llvmPackages_13.stdenv;
+          extraCmakeFlags = [ "-DLLVM_BUILD_INSTRUMENTED=IR" ];
+        };
+        nac3artiq-instrumented = pkgs.python3Packages.toPythonModule (
+          pkgs.rustPlatform.buildRustPackage {
+            name = "nac3artiq-instrumented";
+            src = self;
+            cargoLock = { lockFile = ./Cargo.lock; };
+            nativeBuildInputs = [ pkgs.python3 llvm-nac3-instrumented ];
+            buildInputs = [ pkgs.python3 llvm-nac3-instrumented ];
+            cargoBuildFlags = [ "--package" "nac3artiq" "--features" "init-llvm-profile" ];
+            doCheck = false;
+            configurePhase =
+              ''
+              export CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUSTFLAGS="-C link-arg=-L${pkgs.llvmPackages_13.compiler-rt}/lib/linux -C link-arg=-lclang_rt.profile-x86_64"
+              '';
+            installPhase =
+              ''
+              TARGET_DIR=$out/${pkgs.python3Packages.python.sitePackages}
+              mkdir -p $TARGET_DIR
+              cp target/x86_64-unknown-linux-gnu/release/libnac3artiq.so $TARGET_DIR/nac3artiq.so
+              '';
+          }
+        );
+        nac3artiq-profile = pkgs.stdenvNoCC.mkDerivation {
+          name = "nac3artiq-profile";
+          src = self;
+          buildInputs = [ (pkgs.python3.withPackages(ps: [ ps.numpy nac3artiq-instrumented ])) pkgs.lld_13 pkgs.llvmPackages_13.libllvm ];
+          phases = [ "buildPhase" "installPhase" ];
+          # TODO: get more representative code.
+          buildPhase = "python $src/nac3artiq/demo/demo.py";
+          installPhase =
+            ''
+            mkdir $out
+            llvm-profdata merge -o $out/llvm.profdata /build/llvm/build/profiles/*
+            '';
+        };
+        llvm-nac3-pgo = pkgs.callPackage "${self}/llvm" {
+          stdenv = pkgs.llvmPackages_13.stdenv;
+          extraCmakeFlags = [ "-DLLVM_PROFDATA_FILE=${nac3artiq-profile}/llvm.profdata" ];
+        };
+        nac3artiq-pgo = pkgs.python3Packages.toPythonModule (
+          pkgs.rustPlatform.buildRustPackage {
+            name = "nac3artiq-pgo";
+            src = self;
+            cargoLock = { lockFile = ./Cargo.lock; };
+            nativeBuildInputs = [ pkgs.python3 llvm-nac3-pgo ];
+            buildInputs = [ pkgs.python3 llvm-nac3-pgo ];
+            cargoBuildFlags = [ "--package" "nac3artiq" ];
+            cargoTestFlags = [ "--package" "nac3ast" "--package" "nac3parser" "--package" "nac3core" "--package" "nac3artiq" ];
+            installPhase =
+              ''
+              TARGET_DIR=$out/${pkgs.python3Packages.python.sitePackages}
+              mkdir -p $TARGET_DIR
+              cp target/x86_64-unknown-linux-gnu/release/libnac3artiq.so $TARGET_DIR/nac3artiq.so
+              '';
+          }
+        );
       };
 
       packages.x86_64-w64-mingw32 = rec {
diff --git a/llvm/default.nix b/llvm/default.nix
index e49506cd3..36ed6ec20 100644
--- a/llvm/default.nix
+++ b/llvm/default.nix
@@ -12,6 +12,7 @@
 , debugVersion ? false
 , enableManpages ? false
 , enableSharedLibraries ? false
+, extraCmakeFlags ? []
 }:
 
 let
@@ -149,7 +150,7 @@ in stdenv.mkDerivation (rec {
         ];
       in "-DCROSS_TOOLCHAIN_FLAGS_NATIVE:list=${lib.concatStringsSep ";" nativeToolchainFlags}"
     )
-  ];
+  ] ++ extraCmakeFlags;
 
   postBuild = ''
     rm -fR $out
diff --git a/nac3artiq/Cargo.toml b/nac3artiq/Cargo.toml
index 09098f8d2..602062c94 100644
--- a/nac3artiq/Cargo.toml
+++ b/nac3artiq/Cargo.toml
@@ -19,3 +19,6 @@ nac3core = { path = "../nac3core" }
 version = "0.1.0-beta.4"
 default-features = false
 features = ["llvm13-0", "target-x86", "target-arm", "target-riscv", "no-libffi-linking"]
+
+[features]
+init-llvm-profile = []
diff --git a/nac3artiq/src/lib.rs b/nac3artiq/src/lib.rs
index 91d766b51..c085ba8ca 100644
--- a/nac3artiq/src/lib.rs
+++ b/nac3artiq/src/lib.rs
@@ -685,8 +685,18 @@ impl Nac3 {
     }
 }
 
+#[cfg(feature = "init-llvm-profile")]
+extern "C" {
+    fn __llvm_profile_initialize();
+}
+
 #[pymodule]
 fn nac3artiq(_py: Python, m: &PyModule) -> PyResult<()> {
+    #[cfg(feature = "init-llvm-profile")]
+    unsafe {
+        __llvm_profile_initialize();
+    }
+
     Target::initialize_all(&InitializationConfig::default());
     m.add_class::<Nac3>()?;
     Ok(())