Commit e7c743ab authored by PanZezhong's avatar PanZezhong
Browse files

issue/1033 fix search python lib

parent 11d8d0bb
......@@ -150,7 +150,7 @@ target("flash-attn-nvidia")
local TORCH_DIR = os.iorunv("python", {"-c", "import torch, os; print(os.path.dirname(torch.__file__))"}):trim()
local PYTHON_INCLUDE = os.iorunv("python", {"-c", "import sysconfig; print(sysconfig.get_paths()['include'])"}):trim()
local PYTHON_LIB_DIR = os.iorunv("python", {"-c", "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))"}):trim()
local LIB_PYTHON = os.iorunv("python", {"-c", "import sysconfig, os; print(sysconfig.get_config_var('LDLIBRARY'))"}):trim()
local LIB_PYTHON = os.iorunv("python", {"-c", "import glob,sysconfig,os;print(glob.glob(os.path.join(sysconfig.get_config_var('LIBDIR'),'libpython*.so'))[0])"}):trim()
-- Include dirs (needed for both device and host)
target:add("includedirs", FLASH_ATTN_ROOT .. "/csrc/flash_attn/src", {public = false})
......@@ -160,15 +160,10 @@ target("flash-attn-nvidia")
target:add("includedirs", CUTLASS_ROOT .. "/include", {public = false})
target:add("includedirs", FLASH_ATTN_ROOT .. "/csrc/flash_attn", {public = false})
-- For device linking, only add CUDA-related link directories
target:add("linkdirs", TORCH_DIR .. "/lib", {public = false, force = true})
-- For host linking, we need to add these via link options
-- Use add_ldflags to pass library paths to the host linker only
target:add("ldflags", "-L" .. TORCH_DIR .. "/lib", {force = true})
target:add("ldflags", "-L" .. PYTHON_LIB_DIR, {force = true})
target:add("ldflags", "-l" .. LIB_PYTHON:gsub("%.so$", ""):gsub("^lib", ""), {force = true})
target:add("links", "torch", "torch_cuda", "torch_cpu", "c10", "c10_cuda", "torch_python")
-- Link libraries
target:add("linkdirs", TORCH_DIR .. "/lib", PYTHON_LIB_DIR)
target:add("links", "torch", "torch_cuda", "torch_cpu", "c10", "c10_cuda", "torch_python", LIB_PYTHON)
end
end)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment