summaryrefslogtreecommitdiff
path: root/bitsandbytes/cextension.py
blob: f5b97fda6cf23ac4d1bbacacf8377efffb047998 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import ctypes as ct
import os
from warnings import warn

from bitsandbytes.cuda_setup.main import evaluate_cuda_setup


class CUDALibrary_Singleton(object):
    _instance = None

    def __init__(self):
        raise RuntimeError("Call get_instance() instead")

    def initialize(self):
        self.context = {}
        binary_name = evaluate_cuda_setup()
        if not os.path.exists(os.path.dirname(__file__) + f"/{binary_name}"):
            print(f"TODO: compile library for specific version: {binary_name}")
            print("defaulting to libbitsandbytes.so")
            self.lib = ct.cdll.LoadLibrary(
                os.path.dirname(__file__) + "/libbitsandbytes.so"
            )
        else:
            self.lib = ct.cdll.LoadLibrary(
                os.path.dirname(__file__) + f"/{binary_name}"
            )

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls.__new__(cls)
            cls._instance.initialize()
        return cls._instance


lib = CUDALibrary_Singleton.get_instance().lib
try:
    lib.cadam32bit_g32
    lib.get_context.restype = ct.c_void_p
    lib.get_cusparse.restype = ct.c_void_p
    COMPILED_WITH_CUDA = True
except AttributeError:
    warn(
        "The installed version of bitsandbytes was compiled without GPU support. "
        "8-bit optimizers and GPU quantization are unavailable."
    )
    COMPILED_WITH_CUDA = False