third_party/waf/waflib/extras/halide.py

   1 #!/usr/bin/python
   2 # -*- coding: utf-8 -*-
   3 # Halide code generation tool
   4
   5 __author__ = __maintainer__ = "Jérôme Carretero <cJ-waf@zougloub.eu>"
   6 __copyright__ = "Jérôme Carretero, 2014"
   7
   8 """
   9
  10 Tool to run `Halide <http://halide-lang.org>`_ code generators.
  11
  12 Usage::
  13
  14    bld(
  15     name='pipeline',
  16      # ^ Reference this in use="..." for things using the generated code
  17     #target=['pipeline.o', 'pipeline.h']
  18      # ^ by default, name.{o,h} is added, but you can set the outputs here
  19     features='halide',
  20     halide_env="HL_TRACE=1 HL_TARGET=host-opencl-gpu_debug",
  21      # ^ Environment passed to the generator,
  22      # can be a dict, k/v list, or string.
  23     args=[],
  24      # ^ Command-line arguments to the generator (optional),
  25      # eg. to give parameters to the scheduling
  26     source='pipeline_gen',
  27      # ^ Name of the source executable
  28    )
  29
  30
  31 Known issues:
  32
  33
  34 - Currently only supports Linux (no ".exe")
  35
  36 - Doesn't rerun on input modification when input is part of a build
  37   chain, and has been modified externally.
  38
  39 """
  40
  41 import os
  42 from waflib import Task, Utils, Options, TaskGen, Errors
  43
  44 class run_halide_gen(Task.Task):
  45         color = 'CYAN'
  46         vars = ['HALIDE_ENV', 'HALIDE_ARGS']
  47         run_str = "${SRC[0].abspath()} ${HALIDE_ARGS}"
  48         def __str__(self):
  49                 stuff = "halide"
  50                 stuff += ("[%s]" % (",".join(
  51                  ('%s=%s' % (k,v)) for k, v in sorted(self.env.env.items()))))
  52                 return Task.Task.__str__(self).replace(self.__class__.__name__,
  53                  stuff)
  54
  55 @TaskGen.feature('halide')
  56 @TaskGen.before_method('process_source')
  57 def halide(self):
  58         Utils.def_attrs(self,
  59          args=[],
  60          halide_env={},
  61         )
  62
  63         bld = self.bld
  64
  65         env = self.halide_env
  66         try:
  67                 if isinstance(env, str):
  68                         env = dict(x.split('=') for x in env.split())
  69                 elif isinstance(env, list):
  70                         env = dict(x.split('=') for x in env)
  71                 assert isinstance(env, dict)
  72         except Exception as e:
  73                 if not isinstance(e, ValueError) \
  74                  and not isinstance(e, AssertionError):
  75                         raise
  76                 raise Errors.WafError(
  77                  "halide_env must be under the form" \
  78                  " {'HL_x':'a', 'HL_y':'b'}" \
  79                  " or ['HL_x=y', 'HL_y=b']" \
  80                  " or 'HL_x=y HL_y=b'")
  81
  82         src = self.to_nodes(self.source)
  83         assert len(src) == 1, "Only one source expected"
  84         src = src[0]
  85
  86         args = Utils.to_list(self.args)
  87
  88         def change_ext(src, ext):
  89                 # Return a node with a new extension, in an appropriate folder
  90                 name = src.name
  91                 xpos = src.name.rfind('.')
  92                 if xpos == -1:
  93                         xpos = len(src.name)
  94                 newname = name[:xpos] + ext
  95                 if src.is_child_of(bld.bldnode):
  96                         node = src.get_src().parent.find_or_declare(newname)
  97                 else:
  98                         node = bld.bldnode.find_or_declare(newname)
  99                 return node
 100
 101         def to_nodes(self, lst, path=None):
 102                 tmp = []
 103                 path = path or self.path
 104                 find = path.find_or_declare
 105
 106                 if isinstance(lst, self.path.__class__):
 107                         lst = [lst]
 108
 109                 for x in Utils.to_list(lst):
 110                         if isinstance(x, str):
 111                                 node = find(x)
 112                         else:
 113                                 node = x
 114                         tmp.append(node)
 115                 return tmp
 116
 117         tgt = to_nodes(self, self.target)
 118         if not tgt:
 119                 tgt = [change_ext(src, '.o'), change_ext(src, '.h')]
 120         cwd = tgt[0].parent.abspath()
 121         task = self.create_task('run_halide_gen', src, tgt, cwd=cwd)
 122         task.env.append_unique('HALIDE_ARGS', args)
 123         if task.env.env == []:
 124                 task.env.env = {}
 125         task.env.env.update(env)
 126         task.env.HALIDE_ENV = " ".join(("%s=%s" % (k,v)) for (k,v) in sorted(env.items()))
 127         task.env.HALIDE_ARGS = args
 128
 129         try:
 130                 self.compiled_tasks.append(task)
 131         except AttributeError:
 132                 self.compiled_tasks = [task]
 133         self.source = []
 134
 135 def configure(conf):
 136         if Options.options.halide_root is None:
 137                 conf.check_cfg(package='Halide', args='--cflags --libs')
 138         else:
 139                 halide_root = Options.options.halide_root
 140                 conf.env.INCLUDES_HALIDE = [ os.path.join(halide_root, "include") ]
 141                 conf.env.LIBPATH_HALIDE = [ os.path.join(halide_root, "lib") ]
 142                 conf.env.LIB_HALIDE = ["Halide"]
 143
 144                 # You might want to add this, while upstream doesn't fix it
 145                 #conf.env.LIB_HALIDE += ['ncurses', 'dl', 'pthread']
 146
 147 def options(opt):
 148         opt.add_option('--halide-root',
 149          help="path to Halide include and lib files",
 150         )
 151