halide.py 3.9 KB

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