#!/usr/bin/env python

import os
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext
import sys

# CUDA specific config
# nvcc is assumed to be in users PATH and LD_LIBRARY_PATH is assumed to contain
# location of -lcublas and -lcudart
nvcc_compile_args = ['-O', '--ptxas-options=-v', '--compiler-options',
                     "'-fPIC'"]
cuda_libs = ['cublas', 'cudart']
try:
    cuda_lib_dirs = os.environ.get("LD_LIBRARY_PATH").split(":")
except AttributeError:
    # LD_LIBRARY_PATH not set
    cuda_lib_dirs = []

cudamat_ext = Extension('cudamat.libcudamat',
                        sources=['cudamat/cudamat.cu',
                                 'cudamat/cudamat_kernels.cu'],
                        libraries=cuda_libs,
                        library_dirs=cuda_lib_dirs,
                        extra_compile_args=nvcc_compile_args)
cudalearn_ext = Extension('cudamat.libcudalearn',
                          sources=['cudamat/learn.cu',
                                   'cudamat/learn_kernels.cu'],
                          libraries=cuda_libs,
                          library_dirs=cuda_lib_dirs,
                          extra_compile_args=nvcc_compile_args)


class CUDA_build_ext(build_ext):
    """
    Custom build_ext command that compiles CUDA files.
    Note that all extension source files will be processed with this compiler.
    """
    def build_extensions(self):
        self.compiler.src_extensions.append('.cu')
        self.compiler.set_executable('compiler_so', 'nvcc')
        build_ext.build_extensions(self)

setup(name="cudamat",
      version="0.3",
      description="Performs linear algebra computation on the GPU via CUDA",
      ext_modules=[cudamat_ext, cudalearn_ext],
      packages=find_packages(exclude=['examples', 'test']),
      include_package_data=True,
      package_data={'cudamat': ['rnd_multipliers_32bit.txt']},
      author="Volodymyr Mnih",
      url="https://github.com/cudamat/cudamat",
      cmdclass={'build_ext': CUDA_build_ext})

@@ -202,92 +202,15 @@
 user
+'
 s PATH
- and LD_LIBRARY_PATH is assumed to contain%0A# location of -lcublas and -lcudart
 %0Anvc
@@ -334,164 +334,8 @@
 las'
-, 'cudart'%5D%0Atry:%0A    cuda_lib_dirs = os.environ.get(%22LD_LIBRARY_PATH%22).split(%22:%22)%0Aexcept AttributeError:%0A    # LD_LIBRARY_PATH not set%0A    cuda_lib_dirs = %5B
 %5D%0A%0Ac
@@ -547,60 +547,8 @@
 bs,%0A
-                        library_dirs=cuda_lib_dirs,%0A
     
@@ -825,62 +825,8 @@
 bs,%0A
-                          library_dirs=cuda_lib_dirs,%0A
     
@@ -1153,16 +1153,16 @@
 ('.cu')%0A
-
         
@@ -1213,16 +1213,83 @@
 'nvcc')%0A
+        self.compiler.set_executable('linker_so', 'nvcc --shared')%0A
         

Use nvcc for linking to avoid CUDA library lookup