def find_lib_path()

in python/tvm/_ffi/libinfo.py [0:0]

• 73 lines of code
• 34 McCabe index (conditional complexity)

def find_lib_path(name=None, search_path=None, optional=False):
    """Find dynamic library files.

    Parameters
    ----------
    name : list of str
        List of names to be found.

    Returns
    -------
    lib_path : list(string)
        List of all found path to the libraries
    """
    use_runtime = os.environ.get("TVM_USE_RUNTIME_LIB", False)
    dll_path = get_dll_directories()

    if search_path is not None:
        if isinstance(search_path, list):
            dll_path = dll_path + search_path
        else:
            dll_path.append(search_path)

    if name is not None:
        if isinstance(name, list):
            lib_dll_path = []
            for n in name:
                lib_dll_path += [os.path.join(p, n) for p in dll_path]
        else:
            lib_dll_path = [os.path.join(p, name) for p in dll_path]
        runtime_dll_path = []
        ext_lib_dll_path = []
    else:
        if sys.platform.startswith("win32"):
            lib_dll_names = ["libtvm.dll", "tvm.dll"]
            runtime_dll_names = ["libtvm_runtime.dll", "tvm_runtime.dll"]
            ext_lib_dll_names = [
                "3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.dll",
                "3rdparty/libflash_attn/src/libflash_attn.dll",
            ]
        elif sys.platform.startswith("darwin"):
            lib_dll_names = ["libtvm.dylib"]
            runtime_dll_names = ["libtvm_runtime.dylib"]
            ext_lib_dll_names = [
                "3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.dylib",
                "3rdparty/libflash_attn/src/libflash_attn.dylib",
            ]
        else:
            lib_dll_names = ["libtvm.so"]
            runtime_dll_names = ["libtvm_runtime.so"]
            ext_lib_dll_names = [
                "3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.so",
                "3rdparty/libflash_attn/src/libflash_attn.so",
            ]

        name = lib_dll_names + runtime_dll_names + ext_lib_dll_names
        lib_dll_path = [
            os.path.join(p, name)
            for name in lib_dll_names
            for p in dll_path
            if not p.endswith("python/tvm")
        ]
        runtime_dll_path = [
            os.path.join(p, name)
            for name in runtime_dll_names
            for p in dll_path
            if not p.endswith("python/tvm")
        ]
        ext_lib_dll_path = [os.path.join(p, name) for name in ext_lib_dll_names for p in dll_path]
    if not use_runtime:
        # try to find lib_dll_path
        lib_found = [p for p in lib_dll_path if os.path.exists(p) and os.path.isfile(p)]
        lib_found += [p for p in runtime_dll_path if os.path.exists(p) and os.path.isfile(p)]
        lib_found += [p for p in ext_lib_dll_path if os.path.exists(p) and os.path.isfile(p)]
    else:
        # try to find runtime_dll_path
        use_runtime = True
        lib_found = [p for p in runtime_dll_path if os.path.exists(p) and os.path.isfile(p)]

    if not lib_found:
        if not optional:
            message = (
                f"Cannot find libraries: {name}\n"
                + "List of candidates:\n"
                + "\n".join(lib_dll_path + runtime_dll_path)
            )
            raise RuntimeError(message)
        return None

    if use_runtime:
        sys.stderr.write("Loading runtime library %s... exec only\n" % lib_found[0])
        sys.stderr.flush()
    return lib_found