Refactor lib_generator to set library and source paths

microsoft · Jul 11, 2024 · fb78244 · fb78244
1 parent f5b9999
commit fb78244
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/bitblas/builder/lib_generator/__init__.py b/bitblas/builder/lib_generator/__init__.py
@@ -70,3 +70,9 @@ def get_source_path(self):
 
     def get_lib_path(self):
         return self.libpath
+
+    def set_lib_path(self, libpath):
+        self.libpath = libpath
+
+    def set_src_path(self, srcpath):
+        self.srcpath = srcpath
diff --git a/bitblas/ops/operator.py b/bitblas/ops/operator.py
@@ -302,9 +302,11 @@ def update_runtime_module(self, rt_mod, srcpath=None, libpath=None):
         self.function_handle = rt_mod.get_function(rt_mod.entry_name).handle
         self.torch_func = to_pytorch_func(rt_mod)
         if srcpath is not None:
-            self.srcpath = srcpath
+            assert self.lib_generator is not None, "lib_generator is not initialized"
+            self.lib_generator.set_src_path(srcpath)
         if libpath is not None:
-            self.libpath = libpath
+            assert self.lib_generator is not None, "lib_generator is not initialized"
+            self.lib_generator.set_lib_path(libpath)
             self.lib = ctypes.CDLL(libpath)
             self.lib.init()