From c223b72b786e702bf3b6541aad8415d3822f42c6 Mon Sep 17 00:00:00 2001
From: Zihao Ye <expye@outlook.com>
Date: Mon, 29 Jan 2024 00:10:16 +0800
Subject: [PATCH] add build meta

---
 python/setup.py    | 26 ++++++++++++++++++++++++++
 python/version.txt |  1 -
 2 files changed, 26 insertions(+), 1 deletion(-)
 delete mode 120000 python/version.txt

diff --git a/python/setup.py b/python/setup.py
index ef73ee27..d928c6a0 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -14,10 +14,14 @@
 limitations under the License.
 """
 import pathlib
+import os
+import re
 import datetime
 import subprocess
+import platform
 
 import setuptools
+import torch
 import torch.utils.cpp_extension as torch_cpp_ext
 
 root = pathlib.Path(__name__).parent.resolve().parent
@@ -44,6 +48,28 @@ def get_version():
     return version
 
 
+def get_cuda_version() -> tuple[int, int]:
+    if torch_cpp_ext.CUDA_HOME is None:
+        nvcc = "nvcc"
+    else:
+        nvcc = os.path.join(torch_cpp_ext.CUDA_HOME, "bin/nvcc")
+    txt = subprocess.check_output([nvcc, "--version"], text=True)
+    major, minor = map(int, re.findall(r"release (\d+)\.(\d+),", txt)[0])
+    return major, minor
+
+
+def generate_build_meta() -> None:
+    d = {}
+    version = get_version()
+    d["cuda_major"], d["cuda_minor"] = get_cuda_version()
+    d["torch"] = torch.__version__
+    d["python"] = platform.python_version()
+    d["TORCH_CUDA_ARCH_LIST"] = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
+    with open(root / "python/flashinfer/_build_meta.py", "w") as f:
+        f.write(f"__version__ = {version!r}\n")
+        f.write(f"build_meta = {d!r}")
+
+
 def remove_unwanted_pytorch_nvcc_flags():
     REMOVE_NVCC_FLAGS = [
         "-D__CUDA_NO_HALF_OPERATORS__",
diff --git a/python/version.txt b/python/version.txt
deleted file mode 120000
index aa4e5bec..00000000
--- a/python/version.txt
+++ /dev/null
@@ -1 +0,0 @@
-../version.txt
\ No newline at end of file