From dd50382b3237897d336f2c8f7b7f555b3c851abb Mon Sep 17 00:00:00 2001
From: Tim Dettmers <tim.dettmers@gmail.com>
Date: Sun, 31 Jul 2022 17:47:44 -0700
Subject: Full evaluate_cuda setup with integration test.

---
 bitsandbytes/cextension.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

(limited to 'bitsandbytes/cextension.py')

diff --git a/bitsandbytes/cextension.py b/bitsandbytes/cextension.py
index 2374c35..4bc7bf7 100644
--- a/bitsandbytes/cextension.py
+++ b/bitsandbytes/cextension.py
@@ -1,9 +1,34 @@
 import ctypes as ct
 import os
 from warnings import warn
+from bitsandbytes.cuda_setup import evaluate_cuda_setup
 
-lib = ct.cdll.LoadLibrary(os.path.dirname(__file__) + '/libbitsandbytes.so')
 
+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
-- 
cgit v1.2.3