y2020/vision/sift/fast_gaussian_runner.py

from __future__ import print_function

import json
import sys
import subprocess
import os
import threading

from bazel_tools.tools.python.runfiles import runfiles

def main(params):
  r = runfiles.Create()
  generator = r.Rlocation('org_frc971/y2020/vision/sift/fast_gaussian_generator')

  ruledir = sys.argv[2]
  target_cpu = sys.argv[3]

  target = {
      'aarch64': 'arm-64-linux-no_asserts',
      'armv7': 'arm-32-linux-no_asserts',
      'k8': 'x86-64-linux-no_asserts',
  }[target_cpu]

  commands = []

  amd64_debian_sysroot = r.Rlocation('amd64_debian_sysroot/usr/lib/x86_64-linux-gnu/libc.so.6').rsplit('/', 4)[0]
  env = os.environ.copy()
  env['LD_LIBRARY_PATH'] = ':'.join([
      # TODO(brian): Figure this out again.  It is a bit aggressive.
      #amd64_debian_sysroot + '/lib/x86_64-linux-gnu',
      #amd64_debian_sysroot + '/lib',
      #amd64_debian_sysroot + '/usr/lib/x86_64-linux-gnu',
      #amd64_debian_sysroot + '/usr/lib',
  ])

  all_header = [
      '#ifndef Y2020_VISION_SIFT_FAST_GAUSSIAN_ALL_H_',
      '#define Y2020_VISION_SIFT_FAST_GAUSSIAN_ALL_H_',
      '#include "HalideBuffer.h"',
  ]

  for cols, rows in params['sizes']:
    for sigma, sigma_name, filter_width in params['sigmas']:
      name = "fast_gaussian_%dx%d_%s" % (cols, rows, sigma_name)

      commands.append([
          generator,
          '-g', 'gaussian_generator',
          '-o', ruledir,
          '-f', name,
          '-e', 'o,h,html',
          'target=%s-no_runtime' % target,
          'cols=%s' % cols,
          'rows=%s' % rows,
          'sigma=%s' % sigma,
          'filter_width=%s' % filter_width,
      ])
      all_header += [
          '#include "y2020/vision/sift/%s.h"' % name,
      ]

      name = "fast_gaussian_subtract_%dx%d_%s" % (cols, rows, sigma_name)

      commands.append([
          generator,
          '-g', 'gaussian_and_subtract_generator',
          '-o', ruledir,
          '-f', name,
          '-e', 'o,h,html',
          'target=%s-no_runtime' % target,
          'cols=%s' % cols,
          'rows=%s' % rows,
          'sigma=%s' % sigma,
          'filter_width=%s' % filter_width,
      ])
      all_header += [
          '#include "y2020/vision/sift/%s.h"' % name,
      ]

    name = 'fast_subtract_%dx%d' % (cols, rows)
    commands.append([
        generator,
        '-g', 'subtract_generator',
        '-o', ruledir,
        '-f', name,
        '-e', 'o,h,html',
        'target=%s-no_runtime' % target,
        'cols=%s' % cols,
        'rows=%s' % rows,
    ])
    all_header += [
        '#include "y2020/vision/sift/%s.h"' % name,
    ]
  commands.append([
      generator,
      '-r', 'fast_gaussian_runtime',
      '-o', ruledir,
      '-e', 'o',
      'target=%s' % target,
  ])

  all_header += [
      'namespace frc971 {',
      'namespace vision {',
      '// 0 is success. 1 is non-implemented size. Negative is a Halide error.',
      'inline int DoGeneratedFastGaussian(',
      '    Halide::Runtime::Buffer<const int16_t, 2> input,',
      '    Halide::Runtime::Buffer<int16_t, 2> output,',
      '    double sigma) {',
  ]

  for sigma, sigma_name, filter_width in params['sigmas']:
    for cols, rows in params['sizes']:
      name = "fast_gaussian_%dx%d_%s" % (cols, rows, sigma_name)
      all_header += [
          '  if (input.dim(0).extent() == %s' % cols,
          '      && input.dim(1).extent() == %s' % rows,
          '      && sigma == %s) {' % sigma,
          '    return %s(input, output);' % name,
          '  }',
      ]

  all_header += [
      '  return 1;',
      '}',
      'inline int DoGeneratedFastGaussianAndSubtract(',
      '    Halide::Runtime::Buffer<const int16_t, 2> input,',
      '    Halide::Runtime::Buffer<int16_t, 2> blurred,',
      '    Halide::Runtime::Buffer<int16_t, 2> difference,',
      '    double sigma) {',
  ]

  for sigma, sigma_name, filter_width in params['sigmas']:
    for cols, rows in params['sizes']:
      name = "fast_gaussian_subtract_%dx%d_%s" % (cols, rows, sigma_name)
      all_header += [
          '  if (input.dim(0).extent() == %s' % cols,
          '      && input.dim(1).extent() == %s' % rows,
          '      && sigma == %s) {' % sigma,
          '    return %s(input, blurred, difference);' % name,
          '  }',
      ]

  all_header += [
      '  return 1;',
      '}',
      'inline int DoGeneratedFastSubtract('
      '    Halide::Runtime::Buffer<const int16_t, 2> input_a,',
      '    Halide::Runtime::Buffer<const int16_t, 2> input_b,',
      '    Halide::Runtime::Buffer<int16_t, 2> output) {',
  ]
  for cols, rows in params['sizes']:
    name = 'fast_subtract_%dx%d' % (cols, rows)
    all_header += [
        '  if (input_a.dim(0).extent() == %s' % cols,
        '      && input_a.dim(1).extent() == %s) {' % rows,
        '    return %s(input_a, input_b, output);' % name,
        '  }',
    ]
  all_header += [
      '  return 1;',
      '}',
      '}  // namespace vision',
      '}  // namespace frc971',
      '#endif  // Y2020_VISION_SIFT_FAST_GAUSSIAN_ALL_H_',
  ]

  with open(os.path.join(ruledir, 'fast_gaussian_all.h'), 'w') as f:
    f.writelines([line + '\n' for line in all_header])

  commands_lock = threading.Lock()
  success = [True]

  def run_commands():
    while True:
      with commands_lock:
        if not commands:
          return
        if not success[0]:
          return
        command = commands.pop()
      try:
        subprocess.check_call(command, env=env)
      except:
        with commands_lock:
          success[0] = False
        raise
  threads = [threading.Thread(target=run_commands) for _ in range(4)]
  for thread in threads:
    thread.start()
  for thread in threads:
    thread.join()
  if not success[0]:
    sys.exit(1)

if __name__ == '__main__':
  main(json.loads(sys.argv[1]))