cuda.py

#!/usr/bin/env python
# encoding: utf-8
# Thomas Nagy, 2010

"cuda"

import os
from waflib import Task
from waflib.TaskGen import extension
from waflib.Tools import ccroot, c_preproc
from waflib.Configure import conf

class cuda(Task.Task):
    run_str = '${NVCC} ${CUDAFLAGS} ${CXXFLAGS} ${FRAMEWORKPATH_ST:FRAMEWORKPATH} ${CPPPATH_ST:INCPATHS} ${DEFINES_ST:DEFINES} ${CXX_SRC_F}${SRC} ${CXX_TGT_F}${TGT}'
    color   = 'GREEN'
    ext_in  = ['.h']
    vars    = ['CCDEPS']
    scan    = c_preproc.scan
    shell   = False

@extension('.cu', '.cuda')
def c_hook(self, node):
    return self.create_compiled_task('cuda', node)

def configure(conf):
    conf.find_program('nvcc', var='NVCC')
    conf.options.cuda_dir=conf.root.find_node(conf.env.NVCC).parent.parent.abspath()
    conf.find_cuda_libs()

@conf
def find_cuda_libs(self):
    if not self.env.NVCC:
        self.fatal('check for nvcc first')

    d = self.root.find_node(self.env.NVCC).parent.parent

    node = d.find_node('include')
    _includes = node and node.abspath() or ''

    _libpath=[]
    for x in ('lib64', 'lib'):
        try:
            _libpath.append(d.find_node(x).abspath())
        except:
            pass
    # this should not raise any error
        # self.check_cxx(header='cuda.h', lib='cuda', libpath=_libpath, includes=_includes)
    self.check_cxx(header='cuda.h', lib='cudart', libpath=_libpath, includes=_includes)
    self.check_cxx(header='cufft.h', lib='cufft', libpath=_libpath, includes=_includes)