From 7e7637b44648ad95b5cfd5d949527cfaf74f74fd Mon Sep 17 00:00:00 2001
From: Valentin Churavy <vchuravy@csail.mit.edu>
Date: Fri, 28 Jan 2022 15:01:17 -0500
Subject: [PATCH] Add build_local.jl script

---
 Brutus/Project.toml  |  3 ++-
 utils/Project.toml   |  6 +++++
 utils/build_local.jl | 64 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 utils/Project.toml
 create mode 100644 utils/build_local.jl

diff --git a/Brutus/Project.toml b/Brutus/Project.toml
index 9e5698d..a108348 100644
--- a/Brutus/Project.toml
+++ b/Brutus/Project.toml
@@ -5,6 +5,7 @@ version = "0.1.0"
 
 [deps]
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 
 [compat]
-julia = "1.5"
+julia = "1.6"
diff --git a/utils/Project.toml b/utils/Project.toml
new file mode 100644
index 0000000..a5ea13e
--- /dev/null
+++ b/utils/Project.toml
@@ -0,0 +1,6 @@
+[deps]
+CMake_jll = "3f4e10e2-61f2-5801-8945-23b9d642d0e6"
+Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+Scratch = "6c6a2e73-6563-6170-7368-637461726353"
diff --git a/utils/build_local.jl b/utils/build_local.jl
new file mode 100644
index 0000000..65bbccc
--- /dev/null
+++ b/utils/build_local.jl
@@ -0,0 +1,64 @@
+# Invoke with
+# `julia --project=deps deps/build_local.jl [build_dir]`
+
+Brutus = Base.UUID("61eb1bfa-7361-4325-ad38-22787b887f55")
+
+using Pkg, Scratch, Preferences, Libdl, CMake_jll
+
+# 1. Ensure that an appropriate LLVM_full_jll is installed
+Pkg.activate(; temp=true)
+llvm_assertions = try
+    cglobal((:_ZN4llvm24DisableABIBreakingChecksE, Base.libllvm_path()), Cvoid)
+    false
+catch
+    true
+end
+LLVM = if llvm_assertions
+    Pkg.add(name="LLVM_full_assert_jll", version=Base.libllvm_version)
+    using LLVM_full_assert_jll
+    LLVM_full_assert_jll
+else
+    Pkg.add(name="LLVM_full_jll", version=Base.libllvm_version)
+    using LLVM_full_jll
+    LLVM_full_jll
+end
+MLIR_DIR = joinpath(LLVM.artifact_dir, "lib", "cmake", "mlir")
+
+# 2. Get a scratch directory
+if length(ARGS) == 0
+    build_dir = get_scratch!(Brutus, "build")
+else
+    build_dir = only(ARGS)
+end
+isdir(build_dir) && rm(build_dir; recursive=true)
+source_dir = dirname(@__DIR__)
+julia=joinpath(Sys.BINDIR::String, Base.julia_exename())
+
+lit = joinpath(LLVM.artifact_dir, "tools", "lit", "lit.py")
+mlir_tblgen = joinpath(LLVM.artifact_dir, "tools", "mlir-tblgen")
+
+# Build!
+@info "Building" source_dir build_dir MLIR_DIR julia
+cmake() do cmake
+    run(`$cmake -DMLIR_DIR=$(MLIR_DIR)
+                -DJulia_EXECUTABLE=$(julia)
+                -DLLVM_EXTERNAL_LIT=$(lit)
+                -DMLIR_TABLEGEN_EXE=$(mlir_tblgen)
+                -B$(build_dir) -S$(source_dir)`)
+    run(`$cmake --build $(build_dir) --parallel $(Sys.CPU_THREADS)`)
+end
+
+# FIXME: Discover built libraries
+built_libs = filter(readdir(joinpath(scratch_dir, "Enzyme"))) do file
+    endswith(file, ".$(Libdl.dlext)") && startswith(file, "lib")
+end
+lib_path = joinpath(scratch_dir, "Enzyme", only(built_libs))
+isfile(lib_path) || error("Could not find library $lib_path in build directory")
+
+# Tell Enzyme_jll to load our library instead of the default artifact one
+set_preferences!(
+    joinpath(dirname(@__DIR__), "LocalPreferences.toml"),
+    "Enzyme_jll",
+    "libEnzyme_path" => lib_path;
+    force=true,
+)