text: #!/usr/bin/python3
  
  # Halide tutorial lesson 12.
  
  # This lesson demonstrates how to use Halide to run code on a GPU.
  
  # This lesson can be built by invoking the command:
  #    make tutorial_lesson_12_using_the_gpu
  # in a shell with the current directory at the top of the halide source tree.
  # Otherwise, see the platform-specific compiler invocations below.
  
  # On linux, you can compile and run it like so:
  # g++ lesson_12*.cpp -g -std=c++11 -I ../include -L ../bin -lHalide `libpng-config --cflags --ldflags` -lpthread -ldl -o lesson_12
  # LD_LIBRARY_PATH=../bin ./lesson_12
  
  # On os x:
  # g++ lesson_12*.cpp -g -std=c++11 -I ../include -L ../bin -lHalide `libpng-config --cflags --ldflags` -o lesson_12
  # DYLD_LIBRARY_PATH=../bin ./lesson_12
  
  #include "Halide.h"
  #include <stdio.h>
  #using namespace Halide
  import halide as hl
  
  # Include some support code for loading pngs.
  #include "image_io.h"
  from imageio import imread
  import os.path
  
  # Include a clock to do performance testing.
  #include "clock.h"
  from datetime import datetime
  
  
  # Define some Vars to use.
  x, y, c, i, ii, xi, yi = hl.Var("x"), hl.Var("y"), hl.Var("c"), hl.Var("i"), hl.Var("ii"), hl.Var("xi"), hl.Var("yi")
  
  # We're going to want to schedule a pipeline in several ways, so we
  # define the pipeline in a class so that we can recreate it several
  # times with different schedules.
  class MyPipeline:
  
      def __init__(self, input):
  
          assert input.type() == hl.UInt(8)
  
          self.lut = hl.Func("lut")
          self.padded = hl.Func("padded")
          self.padded16 = hl.Func("padded16")
          self.sharpen = hl.Func("sharpen")
          self.curved = hl.Func("curved")
          self.input = input
  
  
          # For this lesson, we'll use a two-stage pipeline that sharpens
          # and then applies a look-up-table (LUT).
  
          # First we'll define the LUT. It will be a gamma curve.
          self.lut[i] = hl.cast(hl.UInt(8), hl.clamp(pow(i / 255.0, 1.2) * 255.0, 0, 255))
  
          # Augment the input with a boundary condition.
          self.padded[x, y, c] = input[hl.clamp(x, 0, input.width()-1),
                                  hl.clamp(y, 0, input.height()-1), c]
  
          # Cast it to 16-bit to do the math.
          self.padded16[x, y, c] = hl.cast(hl.UInt(16), self.padded[x, y, c])
  
          # Next we sharpen it with a five-tap filter.
          self.sharpen[x, y, c] = (self.padded16[x, y, c] * 2-
                              (self.padded16[x - 1, y, c] +
                               self.padded16[x, y - 1, c] +
                               self.padded16[x + 1, y, c] +
                               self.padded16[x, y + 1, c]) / 4)
  
          # Then apply the LUT.
          self.curved[x, y, c] = self.lut[self.sharpen[x, y, c]]
  
  
      # Now we define methods that give our pipeline several different
      # schedules.
      def schedule_for_cpu(self):
          # Compute the look-up-table ahead of time.
          self.lut.compute_root()
  
          # Compute color channels innermost. Promise that there will
          # be three of them and unroll across them.
          self.curved.reorder(c, x, y) \
                .bound(c, 0, 3) \
                .unroll(c)
  
          # Look-up-tables don't vectorize well, so just parallelize
          # curved in slices of 16 scanlines.
          yo, yi = hl.Var("yo"), hl.Var("yi")
          self.curved.split(y, yo, yi, 16) \
                .parallel(yo)
  
          # Compute sharpen as needed per scanline of curved, reusing
          # previous values computed within the same strip of 16
          # scanlines.
          self.sharpen.store_at(self.curved, yo) \
                      .compute_at(self.curved, yi)
  
          # Vectorize the sharpen. It's 16-bit so we'll vectorize it 8-wide.
          self.sharpen.vectorize(x, 8)
  
          # Compute the padded input at the same granularity as the
          # sharpen. We'll leave the hl.cast to 16-bit inlined into
          # sharpen.
          self.padded.store_at(self.curved, yo) \
              .compute_at(self.curved, yi)
  
          # Also vectorize the padding. It's 8-bit, so we'll vectorize
          # 16-wide.
          self.padded.vectorize(x, 16)
  
          # JIT-compile the pipeline for the CPU.
          self.curved.compile_jit()
  
          return
  
  
      # Now a schedule that uses CUDA or OpenCL.
      def schedule_for_gpu(self):
          # We make the decision about whether to use the GPU for each
          # hl.Func independently. If you have one hl.Func computed on the
          # CPU, and the next computed on the GPU, Halide will do the
          # copy-to-gpu under the hood. For this pipeline, there's no
          # reason to use the CPU for any of the stages. Halide will
          # copy the input image to the GPU the first time we run the
          # pipeline, and leave it there to reuse on subsequent runs.
  
          # As before, we'll compute the LUT once at the start of the
          # pipeline.
          self.lut.compute_root()
  
          # Let's compute the look-up-table using the GPU in 16-wide
          # one-dimensional thread blocks. First we split the index
          # into blocks of size 16:
          block, thread = hl.Var("block"), hl.Var("thread")
          self.lut.split(i, block, thread, 16)
          # Then we tell cuda that our Vars 'block' and 'thread'
          # correspond to CUDA's notions of blocks and threads, or
          # OpenCL's notions of thread groups and threads.
          self.lut.gpu_blocks(block) \
                  .gpu_threads(thread)
  
          # This is a very common scheduling pattern on the GPU, so
          # there's a shorthand for it:
  
          # lut.gpu_tile(i, ii, 16)
  
          # hl.Func::gpu_tile method is similar to hl.Func::tile, except that
          # it also specifies that the tile coordinates correspond to
          # GPU blocks, and the coordinates within each tile correspond
          # to GPU threads.
  
          # Compute color channels innermost. Promise that there will
          # be three of them and unroll across them.
          self.curved.reorder(c, x, y) \
                     .bound(c, 0, 3) \
                     .unroll(c)
  
          # Compute curved in 2D 8x8 tiles using the GPU.
          self.curved.gpu_tile(x, y, xi, yi, 8, 8)
  
          # This is equivalent to:
          # curved.tile(x, y, xo, yo, xi, yi, 8, 8)
          #       .gpu_blocks(xo, yo)
          #       .gpu_threads(xi, yi)
  
          # We'll leave sharpen as inlined into curved.
  
          # Compute the padded input as needed per GPU block, storing the
          # intermediate result in shared memory. hl.Var::gpu_blocks, and
          # hl.Var::gpu_threads exist to help you schedule producers within
          # GPU threads and blocks.
          self.padded.compute_at(self.curved, x)
  
          # Use the GPU threads for the x and y coordinates of the
          # padded input.
          self.padded.gpu_threads(x, y)
  
          # JIT-compile the pipeline for the GPU. CUDA or OpenCL are
          # not enabled by default. We have to construct a hl.Target
          # object, enable one of them, and then pass that target
          # object to compile_jit. Otherwise your CPU will very slowly
          # pretend it's a GPU, and use one thread per output pixel.
  
          # Start with a target suitable for the machine you're running
          # this on.
          target = hl.get_host_target()
  
          # Then enable OpenCL or CUDA.
          #use_opencl = False
          use_opencl = True
          if use_opencl:
              # We'll enable OpenCL here, because it tends to give better
              # performance than CUDA, even with NVidia's drivers, because
              # NVidia's open source LLVM backend doesn't seem to do all
              # the same optimizations their proprietary compiler does.
              target.set_feature(hl.TargetFeature.OpenCL)
              print("(Using OpenCL)")
          else:
              # Uncomment the next line and comment out the line above to
              # try CUDA instead.
              target.set_feature(hl.TargetFeature.CUDA)
              print("(Using CUDA)")
  
          # If you want to see all of the OpenCL or CUDA API calls done
          # by the pipeline, you can also enable the Debug
          # flag. This is helpful for figuring out which stages are
          # slow, or when CPU -> GPU copies happen. It hurts
          # performance though, so we'll leave it commented out.
          # target.set_feature(hl.TargetFeature.Debug)
  
          self.curved.compile_jit(target)
  
      def test_performance(self):
          # Test the performance of the scheduled MyPipeline.
  
          output = hl.Buffer(hl.UInt(8),
                          [self.input.width(), self.input.height(), self.input.channels()])
  
          # Run the filter once to initialize any GPU runtime state.
          self.curved.realize(output)
  
          # Now take the best of 3 runs for timing.
          best_time = float("inf")
          for i in range(3):
  
              t1 = datetime.now()
  
              # Run the filter 100 times.
              for j in range(100):
                  self.curved.realize(output)
  
              # Force any GPU code to finish by copying the buffer back to the CPU.
              output.copy_to_host()
  
              t2 = datetime.now()
  
              elapsed = (t2 - t1).total_seconds()
              if elapsed < best_time:
                  best_time = elapsed
          # end of "best of three times"
  
          print("%1.4f milliseconds" % (best_time * 1000))
  
      def test_correctness(self, reference_output):
  
          assert reference_output.type() == hl.UInt(8)
          output = self.curved.realize(self.input.width(),
                                       self.input.height(),
                                       self.input.channels())
          assert output.type() == hl.UInt(8)
  
          # Check against the reference output.
          for c in range(self.input.channels()):
              for y in range(self.input.height()):
                  for x in range(self.input.width()):
                      if output[x, y, c] != reference_output[x, y, c]:
                          print(
                              "Mismatch between output (%d) and "
                              "reference output (%d) at %d, %d, %d" % (
                                  output[x, y, c],
                                  reference_output[x, y, c],
                                  x, y, c))
                          return
  
          print("CPU and GPU outputs are consistent.")
  
  
  def main():
      # Load an input image.
      image_path = os.path.join(os.path.dirname(__file__), "../../tutorial/images/rgb.png")
      input = hl.Buffer(imread(image_path))
  
      # Allocated an image that will store the correct output
      reference_output = hl.Buffer(hl.UInt(8), [input.width(), input.height(), input.channels()])
  
      print("Testing performance on CPU:")
      p1 = MyPipeline(input)
      p1.schedule_for_cpu()
      p1.test_performance()
      p1.curved.realize(reference_output)
  
      if have_opencl():
          print("Testing performance on GPU:")
          p2 = MyPipeline(input)
          p2.schedule_for_gpu()
          p2.test_performance()
          p2.test_correctness(reference_output)
      else:
          print("Not testing performance on GPU, "
                "because I can't find the opencl library")
  
      return 0
  
  
  
  def have_opencl():
      """
      A helper function to check if OpenCL seems to exist on this machine.
      :return: bool
      """
  
      import ctypes
      import platform
  
      try:
          if platform.system() == "Windows":
              ret = ctypes.windll.LoadLibrary("OpenCL.dll") != None
          elif platform.system() == "Darwin": # apple
              ret =  ctypes.cdll.LoadLibrary("/System/Library/Frameworks/OpenCL.framework/Versions/Current/OpenCL") != None
          elif platform.system() == "Linux":
              ret = ctypes.cdll.LoadLibrary("libOpenCL.so") != None
          else:
              raise Exception("Cannot check for opencl presence "
                              "on unknown system '%s'" % platform.system())
      except OSError:
              ret = False
  
      return ret
  
  
  if __name__ == "__main__":
      main()
  