from distutils.core import setup, Extension
import numpy as np

customcpu_backend = Extension('slo.backends._customcpu',
    sources = ['slo/backends/_customcpu.c'],
    include_dirs=[np.get_include()],
    extra_compile_args = ['-std=c11', '-fopenmp', '-m64', '-O3'],
    extra_link_args=['-fopenmp', '-mavx'],
)

customgpu_backend = Extension('slo.backends._customgpu',
    sources = ['slo/backends/_customgpu.c'],
    include_dirs=[np.get_include()],
    extra_compile_args = ['-std=c11'],
    extra_link_args=['slo/backends/libgpu.a', '-L/usr/local/cuda/lib64', '-lcudart', '-lcublas', '-lstdc++'],
)

setup(
    name = 'slo',
    packages = ['slo', 'slo.backends'],
    version = '1.0.0',
    ext_modules = [customcpu_backend, customgpu_backend],
)

@@ -1,8 +1,18 @@
+import os%0A
 from dis
@@ -67,16 +67,27 @@
  as np%0A%0A
+exts = %5B%5D%0A%0A
 customcp
@@ -328,17 +328,98 @@
 vx'%5D,%0A)%0A
-%0A
+exts.append(customcpu_backend)%0A%0Aif os.path.exists( 'slo/backends/libgpu.a' ):%0A    
 customgp
@@ -463,24 +463,28 @@
 customgpu',%0A
+    
     sources 
@@ -508,32 +508,36 @@
 _customgpu.c'%5D,%0A
+    
     include_dirs
@@ -553,32 +553,36 @@
 include()%5D,%0A    
+  
+  
 extra_compile_ar
@@ -596,24 +596,28 @@
 -std=c11'%5D,%0A
+    
     extra_li
@@ -714,16 +714,55 @@
 dc++'%5D,%0A
+    )%0A    exts.append(customgpu_backend
 )%0A%0Asetup
@@ -866,46 +866,12 @@
 s = 
-%5Bcustomcpu_backend, customgpu_backend%5D
+exts
 ,%0A)%0A

disable building of customgpu module if c library doesn't exist