Change workspace dir (#566)

In my development setup, I have different machines with different GPUs.
They share the home directory on a network filesystem. When I switch
between machines, since the JIT compilation flags change, I'll have to
recompile kernels every time.

One solution is to specify the same `TORCH_CUDA_ARCH_LIST` every time.
However, I keep forgetting that.

Another solution, as proposed in this PR, is to put different arch list
in different cache directory.

Author: abcdabcd987

python/flashinfer/jit/env.py

@@ -15,9 +15,20 @@
 """
 
 import pathlib
+import re
+
+from torch.utils.cpp_extension import _get_cuda_arch_flags
+
+
+def _get_workspace_dir_name() -> pathlib.Path:
+    flags = _get_cuda_arch_flags()
+    arch = "_".join(sorted(set(re.findall(r"compute_(\d+)", "".join(flags)))))
+    # e.g.: $HOME/.cache/flashinfer/75_80_89_90/
+    return pathlib.Path.home() / ".cache" / "flashinfer" / arch
+
 
 # use pathlib
-FLASHINFER_WORKSPACE_DIR = pathlib.Path.home() / ".flashinfer"
+FLASHINFER_WORKSPACE_DIR = _get_workspace_dir_name()
 FLASHINFER_JIT_DIR = FLASHINFER_WORKSPACE_DIR / "cached_ops"
 FLASHINFER_GEN_SRC_DIR = FLASHINFER_WORKSPACE_DIR / "generated"
 _project_root = pathlib.Path(__file__).resolve().parent.parent.parent