def gauss(input, k, rdom, name): blur_x = hl.Func(name + "_x") output = hl.Func(name) x, y, c, xi, yi = hl.Var("x"), hl.Var("y"), hl.Var("c"), hl.Var("xi"), hl.Var("yi") val = hl.Expr("val") if input.dimensions() == 2: blur_x[x, y] = hl.sum(input[x + rdom, y] * k[rdom]) val = hl.sum(blur_x[x, y + rdom] * k[rdom]) if input.output_types()[0] == hl.UInt(16): val = hl.u16(val) output[x, y] = val else: blur_x[x, y, c] = hl.sum(input[x + rdom, y, c] * k[rdom]) val = hl.sum(blur_x[x, y + rdom, c] * k[rdom]) if input.output_types()[0] == hl.UInt(16): val = hl.u16(val) output[x, y, c] = val blur_x.compute_at(output, x).vectorize(x, 16) output.compute_root().tile(x, y, xi, yi, 256, 128).vectorize(xi, 16).parallel(y) return output
def merge_temporal(images, alignment): weight = hl.Func("merge_temporal_weights") total_weight = hl.Func("merge_temporal_total_weights") output = hl.Func("merge_temporal_output") ix, iy, tx, ty, n = hl.Var('ix'), hl.Var('iy'), hl.Var('tx'), hl.Var('ty'), hl.Var('n') rdom0 = hl.RDom([(0, 16), (0, 16)]) rdom1 = hl.RDom([(1, images.dim(2).extent() - 1)]) imgs_mirror = hl.BoundaryConditions.mirror_interior(images, [(0, images.width()), (0, images.height())]) layer = box_down2(imgs_mirror, "merge_layer") offset = Point(alignment[tx, ty, n]).clamp(Point(MINIMUM_OFFSET, MINIMUM_OFFSET), Point(MAXIMUM_OFFSET, MAXIMUM_OFFSET)) al_x = idx_layer(tx, rdom0.x) + offset.x / 2 al_y = idx_layer(ty, rdom0.y) + offset.y / 2 ref_val = layer[idx_layer(tx, rdom0.x), idx_layer(ty, rdom0.y), 0] alt_val = layer[al_x, al_y, n] factor = 8.0 min_distance = 10 max_distance = 300 # max L1 distance, otherwise the value is not used distance = hl.sum(hl.abs(hl.cast(hl.Int(32), ref_val) - hl.cast(hl.Int(32), alt_val))) / 256 normal_distance = hl.max(1, hl.cast(hl.Int(32), distance) / factor - min_distance / factor) # Weight for the alternate frame weight[tx, ty, n] = hl.select(normal_distance > (max_distance - min_distance), 0.0, 1.0 / normal_distance) total_weight[tx, ty] = hl.sum(weight[tx, ty, rdom1]) + 1 offset = Point(alignment[tx, ty, rdom1]) al_x = idx_im(tx, ix) + offset.x al_y = idx_im(ty, iy) + offset.y ref_val = imgs_mirror[idx_im(tx, ix), idx_im(ty, iy), 0] alt_val = imgs_mirror[al_x, al_y, rdom1] # Sum all values according to their weight, and divide by total weight to obtain average output[ix, iy, tx, ty] = hl.sum(weight[tx, ty, rdom1] * alt_val / total_weight[tx, ty]) + ref_val / total_weight[ tx, ty] weight.compute_root().parallel(ty).vectorize(tx, 16) total_weight.compute_root().parallel(ty).vectorize(tx, 16) output.compute_root().parallel(ty).vectorize(ix, 32) return output
def box_down2(input, name): output = hl.Func(name) x, y, n = hl.Var("x"), hl.Var("y"), hl.Var('n') rdom = hl.RDom([(0, 2), (0, 2)]) output[x, y, n] = hl.cast( hl.UInt(16), hl.sum(hl.cast(hl.UInt(32), input[2 * x + rdom.x, 2 * y + rdom.y, n])) / 4) output.compute_root().parallel(y).vectorize(x, 16) return output
def tone_map(input, width, height, compression, gain): print(f'Compression: {compression}, gain: {gain}') normal_dist = hl.Func("luma_weight_distribution") grayscale = hl.Func("grayscale") output = hl.Func("tone_map_output") x, y, c, v = hl.Var("x"), hl.Var("y"), hl.Var("c"), hl.Var("v") rdom = hl.RDom([(0, 3)]) normal_dist[v] = hl.f32(hl.exp(-12.5 * hl.pow(hl.f32(v) / 65535 - 0.5, 2))) grayscale[x, y] = hl.u16(hl.sum(hl.u32(input[x, y, rdom])) / 3) dark = grayscale comp_const = 1 gain_const = 1 comp_slope = (compression - comp_const) / (TONE_MAP_PASSES) gain_slope = (gain - gain_const) / (TONE_MAP_PASSES) for i in range(TONE_MAP_PASSES): print(' pass', i) norm_comp = i * comp_slope + comp_const norm_gain = i * gain_slope + gain_const bright = brighten(dark, norm_comp) dark_gamma = gamma_correct(dark) bright_gamma = gamma_correct(bright) dark_gamma = combine2(dark_gamma, bright_gamma, width, height, normal_dist) dark = brighten(gamma_inverse(dark_gamma), norm_gain) output[x, y, c] = hl.u16_sat(hl.u32(input[x, y, c]) * hl.u32(dark[x, y]) / hl.u32(hl.max(1, grayscale[x, y]))) grayscale.compute_root().parallel(y).vectorize(x, 16) normal_dist.compute_root().vectorize(v, 16) return output
def gaussian_down4(input, name): output = hl.Func(name) k = hl.Func(name + "_filter") x, y, n = hl.Var("x"), hl.Var("y"), hl.Var('n') rdom = hl.RDom([(-2, 5), (-2, 5)]) k[x, y] = 0 k[-2, -2] = 2 k[-1, -2] = 4 k[0, -2] = 5 k[1, -2] = 4 k[2, -2] = 2 k[-2, -1] = 4 k[-1, -1] = 9 k[0, -1] = 12 k[1, -1] = 9 k[2, -1] = 4 k[-2, 0] = 5 k[-1, 0] = 12 k[0, 0] = 15 k[1, 0] = 12 k[2, 0] = 5 k[-2, 1] = 4 k[-1, 1] = 9 k[0, 1] = 12 k[1, 1] = 9 k[2, 1] = 4 k[-2, 2] = 2 k[-1, 2] = 4 k[0, 2] = 5 k[1, 2] = 4 k[2, 2] = 2 output[x, y, n] = hl.cast( hl.UInt(16), hl.sum( hl.cast( hl.UInt(32), input[4 * x + rdom.x, 4 * y + rdom.y, n] * k[rdom.x, rdom.y])) / 159) k.compute_root().parallel(y).parallel(x) output.compute_root().parallel(y).vectorize(x, 16) return output
def srgb(input, ccm): srgb_matrix = hl.Func("srgb_matrix") output = hl.Func("srgb_output") x, y, c = hl.Var("x"), hl.Var("y"), hl.Var("c") rdom = hl.RDom([(0, 3)]) srgb_matrix[x, y] = hl.f32(0) srgb_matrix[0, 0] = hl.f32(ccm[0][0]) srgb_matrix[1, 0] = hl.f32(ccm[0][1]) srgb_matrix[2, 0] = hl.f32(ccm[0][2]) srgb_matrix[0, 1] = hl.f32(ccm[1][0]) srgb_matrix[1, 1] = hl.f32(ccm[1][1]) srgb_matrix[2, 1] = hl.f32(ccm[1][2]) srgb_matrix[0, 2] = hl.f32(ccm[2][0]) srgb_matrix[1, 2] = hl.f32(ccm[2][1]) srgb_matrix[2, 2] = hl.f32(ccm[2][2]) output[x, y, c] = hl.u16_sat(hl.sum(srgb_matrix[rdom, c] * input[x, y, rdom])) return output
def align_layer(layer, prev_alignment, prev_min, prev_max): scores = hl.Func(layer.name() + "_scores") alignment = hl.Func(layer.name() + "_alignment") xi, yi, tx, ty, n = hl.Var("xi"), hl.Var("yi"), hl.Var('tx'), hl.Var( 'ty'), hl.Var('n') rdom0 = hl.RDom([(0, 16), (0, 16)]) rdom1 = hl.RDom([(-4, 8), (-4, 8)]) # Alignment of the previous (more coarse) layer scaled to this (finer) layer prev_offset = DOWNSAMPLE_RATE * Point( prev_alignment[prev_tile(tx), prev_tile(ty), n]).clamp( prev_min, prev_max) x0 = idx_layer(tx, rdom0.x) y0 = idx_layer(ty, rdom0.y) # (x,y) coordinates in the search region relative to the offset obtained from the alignment of the previous layer x = x0 + prev_offset.x + xi y = y0 + prev_offset.y + yi ref_val = layer[x0, y0, 0] # Value of reference frame (the first frame) alt_val = layer[x, y, n] # alternate frame value # L1 distance between reference frame and alternate frame d = hl.abs(hl.cast(hl.Int(32), ref_val) - hl.cast(hl.Int(32), alt_val)) scores[xi, yi, tx, ty, n] = hl.sum(d) # Alignment for each tile, where L1 distances are minimum alignment[tx, ty, n] = Point(hl.argmin(scores[rdom1.x, rdom1.y, tx, ty, n])) + prev_offset scores.compute_at(alignment, tx).vectorize(xi, 8) alignment.compute_root().parallel(ty).vectorize(tx, 16) return alignment
def demosaic(input, width, height): print(f'width: {width}, height: {height}') f0 = hl.Buffer(hl.Int(32), [5, 5], "demosaic_f0") f1 = hl.Buffer(hl.Int(32), [5, 5], "demosaic_f1") f2 = hl.Buffer(hl.Int(32), [5, 5], "demosaic_f2") f3 = hl.Buffer(hl.Int(32), [5, 5], "demosaic_f3") f0.translate([-2, -2]) f1.translate([-2, -2]) f2.translate([-2, -2]) f3.translate([-2, -2]) d0 = hl.Func("demosaic_0") d1 = hl.Func("demosaic_1") d2 = hl.Func("demosaic_2") d3 = hl.Func("demosaic_3") output = hl.Func("demosaic_output") x, y, c = hl.Var("x"), hl.Var("y"), hl.Var("c") rdom0 = hl.RDom([(-2, 5), (-2, 5)]) # rdom1 = hl.RDom([(0, width / 2), (0, height / 2)]) input_mirror = hl.BoundaryConditions.mirror_interior(input, [(0, width), (0, height)]) f0.fill(0) f1.fill(0) f2.fill(0) f3.fill(0) f0_sum = 8 f1_sum = 16 f2_sum = 16 f3_sum = 16 f0[0, -2] = -1 f0[0, -1] = 2 f0[-2, 0] = -1 f0[-1, 0] = 2 f0[0, 0] = 4 f0[1, 0] = 2 f0[2, 0] = -1 f0[0, 1] = 2 f0[0, 2] = -1 f1[0, -2] = 1 f1[-1, -1] = -2 f1[1, -1] = -2 f1[-2, 0] = -2 f1[-1, 0] = 8 f1[0, 0] = 10 f1[1, 0] = 8 f1[2, 0] = -2 f1[-1, 1] = -2 f1[1, 1] = -2 f1[0, 2] = 1 f2[0, -2] = -2 f2[-1, -1] = -2 f2[0, -1] = 8 f2[1, -1] = -2 f2[-2, 0] = 1 f2[0, 0] = 10 f2[2, 0] = 1 f2[-1, 1] = -2 f2[0, 1] = 8 f2[1, 1] = -2 f2[0, 2] = -2 f3[0, -2] = -3 f3[-1, -1] = 4 f3[1, -1] = 4 f3[-2, 0] = -3 f3[0, 0] = 12 f3[2, 0] = -3 f3[-1, 1] = 4 f3[1, 1] = 4 f3[0, 2] = -3 d0[x, y] = hl.u16_sat(hl.sum(hl.i32(input_mirror[x + rdom0.x, y + rdom0.y]) * f0[rdom0.x, rdom0.y]) / f0_sum) d1[x, y] = hl.u16_sat(hl.sum(hl.i32(input_mirror[x + rdom0.x, y + rdom0.y]) * f1[rdom0.x, rdom0.y]) / f1_sum) d2[x, y] = hl.u16_sat(hl.sum(hl.i32(input_mirror[x + rdom0.x, y + rdom0.y]) * f2[rdom0.x, rdom0.y]) / f2_sum) d3[x, y] = hl.u16_sat(hl.sum(hl.i32(input_mirror[x + rdom0.x, y + rdom0.y]) * f3[rdom0.x, rdom0.y]) / f3_sum) R_row = y % 2 == 0 B_row = y % 2 != 0 R_col = x % 2 == 0 B_col = x % 2 != 0 at_R = c == 0 at_G = c == 1 at_B = c == 2 output[x, y, c] = hl.select(at_R & R_row & B_col, d1[x, y], at_R & B_row & R_col, d2[x, y], at_R & B_row & B_col, d3[x, y], at_G & R_row & R_col, d0[x, y], at_G & B_row & B_col, d0[x, y], at_B & B_row & R_col, d1[x, y], at_B & R_row & B_col, d2[x, y], at_B & R_row & R_col, d3[x, y], input[x, y]) d0.compute_root().parallel(y).vectorize(x, 16) d1.compute_root().parallel(y).vectorize(x, 16) d2.compute_root().parallel(y).vectorize(x, 16) d3.compute_root().parallel(y).vectorize(x, 16) output.compute_root().parallel(y).align_bounds(x, 2).unroll(x, 2).align_bounds(y, 2).unroll(y, 2).vectorize(x, 16) return output
def bilateral_filter(input, width, height): print(' bilateral_filter') k = hl.Buffer(hl.Float(32), [7, 7], "gauss_kernel") k.translate([-3, -3]) weights = hl.Func("bilateral_weights") total_weights = hl.Func("bilateral_total_weights") bilateral = hl.Func("bilateral") output = hl.Func("bilateral_filter_output") x, y, dx, dy, c = hl.Var("x"), hl.Var("y"), hl.Var("dx"), hl.Var("dy"), hl.Var("c") rdom = hl.RDom([(-3, 7), (-3, 7)]) k.fill(0) k[-3, -3] = 0.000690 k[-2, -3] = 0.002646 k[-1, -3] = 0.005923 k[0, -3] = 0.007748 k[1, -3] = 0.005923 k[2, -3] = 0.002646 k[3, -3] = 0.000690 k[-3, -2] = 0.002646 k[-2, -2] = 0.010149 k[-1, -2] = 0.022718 k[0, -2] = 0.029715 k[1, -2] = 0.022718 k[2, -2] = 0.010149 k[3, -2] = 0.002646 k[-3, -1] = 0.005923 k[-2, -1] = 0.022718 k[-1, -1] = 0.050855 k[0, -1] = 0.066517 k[1, -1] = 0.050855 k[2, -1] = 0.022718 k[3, -1] = 0.005923 k[-3, 0] = 0.007748 k[-2, 0] = 0.029715 k[-1, 0] = 0.066517 k[0, 0] = 0.087001 k[1, 0] = 0.066517 k[2, 0] = 0.029715 k[3, 0] = 0.007748 k[-3, 1] = 0.005923 k[-2, 1] = 0.022718 k[-1, 1] = 0.050855 k[0, 1] = 0.066517 k[1, 1] = 0.050855 k[2, 1] = 0.022718 k[3, 1] = 0.005923 k[-3, 2] = 0.002646 k[-2, 2] = 0.010149 k[-1, 2] = 0.022718 k[0, 2] = 0.029715 k[1, 2] = 0.022718 k[2, 2] = 0.010149 k[3, 2] = 0.002646 k[-3, 3] = 0.000690 k[-2, 3] = 0.002646 k[-1, 3] = 0.005923 k[0, 3] = 0.007748 k[1, 3] = 0.005923 k[2, 3] = 0.002646 k[3, 3] = 0.000690 input_mirror = hl.BoundaryConditions.mirror_interior(input, [(0, width), (0, height)]) dist = hl.cast(hl.Float(32), hl.cast(hl.Int(32), input_mirror[x, y, c]) - hl.cast(hl.Int(32), input_mirror[x + dx, y + dy, c])) sig2 = 100 threshold = 25000 score = hl.select(hl.abs(input_mirror[x + dx, y + dy, c]) > threshold, 0, hl.exp(-dist * dist / sig2)) weights[dx, dy, x, y, c] = k[dx, dy] * score total_weights[x, y, c] = hl.sum(weights[rdom.x, rdom.y, x, y, c]) bilateral[x, y, c] = hl.sum(input_mirror[x + rdom.x, y + rdom.y, c] * weights[rdom.x, rdom.y, x, y, c]) / \ total_weights[x, y, c] output[x, y, c] = hl.cast(hl.Float(32), input[x, y, c]) output[x, y, 1] = bilateral[x, y, 1] output[x, y, 2] = bilateral[x, y, 2] weights.compute_at(output, y).vectorize(x, 16) output.compute_root().parallel(y).vectorize(x, 16) output.update(0).parallel(y).vectorize(x, 16) output.update(1).parallel(y).vectorize(x, 16) return output
def main(): # Declare some Vars to use below. x, y = hl.Var("x"), hl.Var("y") # Load a grayscale image to use as an input. image_path = os.path.join(os.path.dirname(__file__), "../../tutorial/images/gray.png") input_data = imageio.imread(image_path) if True: # making the image smaller to go faster input_data = input_data[:160, :150] assert input_data.dtype == np.uint8 input = hl.Buffer(input_data) # You can define a hl.Func in multiple passes. Let's see a toy # example first. if True: # The first definition must be one like we have seen already # - a mapping from Vars to an hl.Expr: f = hl.Func("f") f[x, y] = x + y # We call this first definition the "pure" definition. # But the later definitions can include computed expressions on # both sides. The simplest example is modifying a single point: f[3, 7] = 42 # We call these extra definitions "update" definitions, or # "reduction" definitions. A reduction definition is an # update definition that recursively refers back to the # function's current value at the same site: if False: e = f[x, y] + 17 print("f[x, y] + 17", e) print("(f[x, y] + 17).type()", e.type()) print("(f[x, y]).type()", f[x, y].type()) f[x, y] = f[x, y] + 17 # If we confine our update to a single row, we can # recursively refer to values in the same column: f[x, 3] = f[x, 0] * f[x, 10] # Similarly, if we confine our update to a single column, we # can recursively refer to other values in the same row. f[0, y] = f[0, y] / f[3, y] # The general rule is: Each hl.Var used in an update definition # must appear unadorned in the same position as in the pure # definition in all references to the function on the left- # and right-hand sides. So the following definitions are # legal updates: # x is used, so all uses of f must have x as the first argument. f[x, 17] = x + 8 # y is used, so all uses of f must have y as the second argument. f[0, y] = y * 8 f[x, x + 1] = x + 8 f[y / 2, y] = f[0, y] * 17 # But these ones would cause an error: # f[x, 0) = f[x + 1, 0) <- First argument to f on the right-hand-side must be 'x', not 'x + 1'. # f[y, y + 1) = y + 8 <- Second argument to f on the left-hand-side must be 'y', not 'y + 1'. # f[y, x) = y - x <- Arguments to f on the left-hand-side are in the wrong places. # f[3, 4) = x + y <- Free variables appear on the right-hand-side # but not the left-hand-side. # We'll realize this one just to make sure it compiles. The # second-to-last definition forces us to realize over a # domain that is taller than it is wide. f.realize(100, 101) # For each realization of f, each step runs in its entirety # before the next one begins. Let's trace the loads and # stores for a simpler example: g = hl.Func("g") g[x, y] = x + y # Pure definition g[2, 1] = 42 # First update definition g[x, 0] = g[x, 1] # Second update definition g.trace_loads() g.trace_stores() g.realize(4, 4) # Reading the log, we see that each pass is applied in turn. The # equivalent Python is: result = np.empty((4, 4), dtype=np.int) # Pure definition for yy in range(4): for xx in range(4): result[yy][xx] = xx + yy # First update definition result[1][2] = 42 # Second update definition for xx in range(4): result[0][xx] = result[1][xx] # end of section # Putting update passes inside loops. if True: # Starting with this pure definition: f = hl.Func("f") f[x, y] = x + y # Say we want an update that squares the first fifty rows. We # could do this by adding 50 update definitions: # f[x, 0) = f[x, 0) * f[x, 0) # f[x, 1) = f[x, 1) * f[x, 1) # f[x, 2) = f[x, 2) * f[x, 2) # ... # f[x, 49) = f[x, 49) * f[x, 49) # Or equivalently using a compile-time loop in our C++: # for (int i = 0 i < 50 i++) { # f[x, i) = f[x, i) * f[x, i) # # But it's more manageable and more flexible to put the loop # in the generated code. We do this by defining a "reduction # domain" and using it inside an update definition: r = hl.RDom([(0, 50)]) f[x, r] = f[x, r] * f[x, r] halide_result = f.realize(100, 100) # The equivalent Python is: py_result = np.empty((100, 100), dtype=np.int) for yy in range(100): for xx in range(100): py_result[yy][xx] = xx + yy for xx in range(100): for rr in range(50): # The loop over the reduction domain occurs inside of # the loop over any pure variables used in the update # step: py_result[rr][xx] = py_result[rr][xx] * py_result[rr][xx] # Check the results match: for yy in range(100): for xx in range(100): assert halide_result[xx, yy] == py_result[yy][xx], \ "halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result[xx, yy], py_result[yy][xx]) # Now we'll examine a real-world use for an update definition: # computing a histogram. if True: # Some operations on images can't be cleanly expressed as a pure # function from the output coordinates to the value stored # there. The classic example is computing a histogram. The # natural way to do it is to iterate over the input image, # updating histogram buckets. Here's how you do that in Halide: histogram = hl.Func("histogram") # Histogram buckets start as zero. histogram[x] = 0 # Define a multi-dimensional reduction domain over the input image: r = hl.RDom([(0, input.width()), (0, input.height())]) # For every point in the reduction domain, increment the # histogram bucket corresponding to the intensity of the # input image at that point. histogram[input[r.x, r.y]] += 1 halide_result = histogram.realize(256) # The equivalent Python is: py_result = np.empty((256), dtype=np.int) for xx in range(256): py_result[xx] = 0 for r_y in range(input.height()): for r_x in range(input.width()): py_result[input_data[r_x, r_y]] += 1 # Check the answers agree: for xx in range(256): assert py_result[xx] == halide_result[xx], \ "halide_result(%d) = %d instead of %d" % (xx, halide_result[xx], py_result[xx]) # Scheduling update steps if True: # The pure variables in an update step and can be # parallelized, vectorized, split, etc as usual. # Vectorizing, splitting, or parallelize the variables that # are part of the reduction domain is trickier. We'll cover # that in a later lesson. # Consider the definition: f = hl.Func("x") f[x, y] = x * y # Set the second row to equal the first row. f[x, 1] = f[x, 0] # Set the second column to equal the first column plus 2. f[1, y] = f[0, y] + 2 # The pure variables in each stage can be scheduled # independently. To control the pure definition, we schedule # as we have done in the past. The following code vectorizes # and parallelizes the pure definition only. f.vectorize(x, 4).parallel(y) # We use hl.Func::update(int) to get a handle to an update step # for the purposes of scheduling. The following line # vectorizes the first update step across x. We can't do # anything with y for this update step, because it doesn't # use y. f.update(0).vectorize(x, 4) # Now we parallelize the second update step in chunks of size # 4. yo, yi = hl.Var("yo"), hl.Var("yi") f.update(1).split(y, yo, yi, 4).parallel(yo) halide_result = f.realize(16, 16) # Here's the equivalent (serial) C: py_result = np.empty((16, 16), dtype=np.int) # Pure step. Vectorized in x and parallelized in y. for yy in range(16): # Should be a parallel for loop for x_vec in range(4): xx = [x_vec * 4, x_vec * 4 + 1, x_vec * 4 + 2, x_vec * 4 + 3] py_result[yy][xx[0]] = xx[0] * yy py_result[yy][xx[1]] = xx[1] * yy py_result[yy][xx[2]] = xx[2] * yy py_result[yy][xx[3]] = xx[3] * yy # First update. Vectorized in x. for x_vec in range(4): xx = [x_vec * 4, x_vec * 4 + 1, x_vec * 4 + 2, x_vec * 4 + 3] py_result[1][xx[0]] = py_result[0][xx[0]] py_result[1][xx[1]] = py_result[0][xx[1]] py_result[1][xx[2]] = py_result[0][xx[2]] py_result[1][xx[3]] = py_result[0][xx[3]] # Second update. Parallelized in chunks of size 4 in y. for yo in range(4): # Should be a parallel for loop for yi in range(4): yy = yo * 4 + yi py_result[yy][1] = py_result[yy][0] + 2 # Check the C and Halide results match: for yy in range(16): for xx in range(16): assert halide_result[xx, yy] == py_result[yy][xx], \ "halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result[xx, yy], py_result[yy][xx]) # That covers how to schedule the variables within a hl.Func that # uses update steps, but what about producer-consumer # relationships that involve compute_at and store_at? Let's # examine a reduction as a producer, in a producer-consumer pair. if True: # Because an update does multiple passes over a stored array, # it's not meaningful to inline them. So the default schedule # for them does the closest thing possible. It computes them # in the innermost loop of their consumer. Consider this # trivial example: producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 producer[x] += 1 consumer[x] = 2 * producer[x] halide_result = consumer.realize(10) # The equivalent Python is: py_result = np.empty((10), dtype=np.int) for xx in range(10): producer_storage = np.empty((1), dtype=np.int) # Pure step for producer producer_storage[0] = xx * 17 # Update step for producer producer_storage[0] = producer_storage[0] + 1 # Pure step for consumer py_result[xx] = 2 * producer_storage[0] # Check the results match for xx in range(10): assert halide_result[xx] == py_result[xx], \ "halide_result(%d) = %d instead of %d" % (xx, halide_result[xx], py_result[xx]) # For all other compute_at/store_at options, the reduction # gets placed where you would expect, somewhere in the loop # nest of the consumer. # Now let's consider a reduction as a consumer in a # producer-consumer pair. This is a little more involved. if True: if True: # Case 1: The consumer references the producer in the pure step # only. producer, consumer = hl.Func("producer"), hl.Func("consumer") # The producer is pure. producer[x] = x * 17 consumer[x] = 2 * producer[x] consumer[x] += 1 # The valid schedules for the producer in this case are # the default schedule - inlined, and also: # # 1) producer.compute_at(x), which places the computation of # the producer inside the loop over x in the pure step of the # consumer. # # 2) producer.compute_root(), which computes all of the # producer ahead of time. # # 3) producer.store_root().compute_at(x), which allocates # space for the consumer outside the loop over x, but fills # it in as needed inside the loop. # # Let's use option 1. producer.compute_at(consumer, x) halide_result = consumer.realize(10) # The equivalent Python is: py_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range(10): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 py_result[xx] = 2 * producer_storage[0] # Update step for the consumer for xx in range(10): py_result[xx] += 1 # All of the pure step is evaluated before any of the # update step, so there are two separate loops over x. # Check the results match for xx in range(10): assert halide_result[xx] == py_result[xx], \ "halide_result(%d) = %d instead of %d" % (xx, halide_result[xx], py_result[xx]) if True: # Case 2: The consumer references the producer in the update step # only producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 consumer[x] = x consumer[x] += producer[x] # Again we compute the producer per x coordinate of the # consumer. This places producer code inside the update # step of the producer, because that's the only step that # uses the producer. producer.compute_at(consumer, x) # Note however, that we didn't say: # # producer.compute_at(consumer.update(0), x). # # Scheduling is done with respect to Vars of a hl.Func, and # the Vars of a hl.Func are shared across the pure and # update steps. halide_result = consumer.realize(10) # The equivalent Python is: py_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range(10): py_result[xx] = xx # Update step for the consumer for xx in range(10): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 py_result[xx] += producer_storage[0] # Check the results match for xx in range(10): assert halide_result[xx] == py_result[xx], \ "halide_result(%d) = %d instead of %d" % (xx, halide_result[xx], py_result[xx]) if True: # Case 3: The consumer references the producer in # multiple steps that share common variables producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 consumer[x] = producer[x] * x consumer[x] += producer[x] # Again we compute the producer per x coordinate of the # consumer. This places producer code inside both the # pure and the update step of the producer. So there ends # up being two separate realizations of the producer, and # redundant work occurs. producer.compute_at(consumer, x) halide_result = consumer.realize(10) # The equivalent Python is: py_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range(10): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 py_result[xx] = producer_storage[0] * xx # Update step for the consumer for xx in range(10): # Another copy of the pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 py_result[xx] += producer_storage[0] # Check the results match for xx in range(10): assert halide_result[xx] == py_result[xx], \ "halide_result(%d) = %d instead of %d" % (xx, halide_result[xx], py_result[xx]) if True: # Case 4: The consumer references the producer in # multiple steps that do not share common variables producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x, y] = x * y consumer[x, y] = x + y consumer[x, 0] = producer[x, x - 1] consumer[0, y] = producer[y, y - 1] # In this case neither producer.compute_at(consumer, x) # nor producer.compute_at(consumer, y) will work, because # either one fails to cover one of the uses of the # producer. So we'd have to inline producer, or use # producer.compute_root(). # Let's say we really really want producer to be # compute_at the inner loops of both consumer update # steps. Halide doesn't allow multiple different # schedules for a single hl.Func, but we can work around it # by making two wrappers around producer, and scheduling # those instead: # Attempt 2: producer_wrapper_1, producer_wrapper_2, consumer_2 = hl.Func(), hl.Func(), hl.Func() producer_wrapper_1[x, y] = producer[x, y] producer_wrapper_2[x, y] = producer[x, y] consumer_2[x, y] = x + y consumer_2[x, 0] += producer_wrapper_1[x, x - 1] consumer_2[0, y] += producer_wrapper_2[y, y - 1] # The wrapper functions give us two separate handles on # the producer, so we can schedule them differently. producer_wrapper_1.compute_at(consumer_2, x) producer_wrapper_2.compute_at(consumer_2, y) halide_result = consumer_2.realize(10, 10) # The equivalent Python is: py_result = np.empty((10, 10), dtype=np.int) # Pure step for the consumer for yy in range(10): for xx in range(10): py_result[yy][xx] = xx + yy # First update step for consumer for xx in range(10): producer_wrapper_1_storage = np.empty((1), dtype=np.int) producer_wrapper_1_storage[0] = xx * (xx - 1) py_result[0][xx] += producer_wrapper_1_storage[0] # Second update step for consumer for yy in range(10): producer_wrapper_2_storage = np.empty((1), dtype=np.int) producer_wrapper_2_storage[0] = yy * (yy - 1) py_result[yy][0] += producer_wrapper_2_storage[0] # Check the results match for yy in range(10): for xx in range(10): assert halide_result[xx, yy] == py_result[yy][xx], \ "halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result[xx, yy], py_result[yy][xx]) if True: # Case 5: Scheduling a producer under a reduction domain # variable of the consumer. # We are not just restricted to scheduling producers at # the loops over the pure variables of the consumer. If a # producer is only used within a loop over a reduction # domain (hl.RDom) variable, we can also schedule the # producer there. producer, consumer = hl.Func("producer"), hl.Func("consumer") r = hl.RDom([(0, 5)]) producer[x] = x * 17 consumer[x] = x + 10 consumer[x] += r + producer[x + r] producer.compute_at(consumer, r) halide_result = consumer.realize(10) # The equivalent Python is: py_result = np.empty((10), dtype=np.int) # Pure step for the consumer. for xx in range(10): py_result[xx] = xx + 10 # Update step for the consumer. for xx in range(10): # The loop over the reduction domain is always the inner loop. for rr in range(5): # We've schedule the storage and computation of # the producer here. We just need a single value. producer_storage = np.empty((1), dtype=np.int) # Pure step of the producer. producer_storage[0] = (xx + rr) * 17 # Now use it in the update step of the consumer. py_result[xx] += rr + producer_storage[0] # Check the results match for xx in range(10): assert halide_result[xx] == py_result[xx], \ "halide_result(%d) = %d instead of %d" % (xx, halide_result[xx], py_result[xx]) # A real-world example of a reduction inside a producer-consumer chain. if True: # The default schedule for a reduction is a good one for # convolution-like operations. For example, the following # computes a 5x5 box-blur of our grayscale test image with a # hl.clamp-to-edge boundary condition: # First add the boundary condition. clamped = hl.BoundaryConditions.repeat_edge(input) # Define a 5x5 box that starts at (-2, -2) r = hl.RDom([(-2, 5), (-2, 5)]) # Compute the 5x5 sum around each pixel. local_sum = hl.Func("local_sum") local_sum[x, y] = 0 # Compute the sum as a 32-bit integer local_sum[x, y] += clamped[x + r.x, y + r.y] # Divide the sum by 25 to make it an average blurry = hl.Func("blurry") blurry[x, y] = hl.cast(hl.UInt(8), local_sum[x, y] / 25) halide_result = blurry.realize(input.width(), input.height()) # The default schedule will inline 'clamped' into the update # step of 'local_sum', because clamped only has a pure # definition, and so its default schedule is fully-inlined. # We will then compute local_sum per x coordinate of blurry, # because the default schedule for reductions is # compute-innermost. Here's the equivalent Python: #cast_to_uint8 = lambda x_: np.array([x_], dtype=np.uint8)[0] local_sum = np.empty((1), dtype=np.int32) py_result = hl.Buffer(hl.UInt(8), [input.width(), input.height()]) for yy in range(input.height()): for xx in range(input.width()): # FIXME this loop is quite slow # Pure step of local_sum local_sum[0] = 0 # Update step of local_sum for r_y in range(-2, 2 + 1): for r_x in range(-2, 2 + 1): # The clamping has been inlined into the update step. clamped_x = min(max(xx + r_x, 0), input.width() - 1) clamped_y = min(max(yy + r_y, 0), input.height() - 1) local_sum[0] += input[clamped_x, clamped_y] # Pure step of blurry # py_result(x, y) = (uint8_t)(local_sum[0] / 25) #py_result[xx, yy] = cast_to_uint8(local_sum[0] / 25) # hl.cast done internally py_result[xx, yy] = int(local_sum[0] / 25) # Check the results match for yy in range(input.height()): for xx in range(input.width()): assert halide_result[xx, yy] == py_result[xx, yy], \ "halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result[xx, yy], py_result[xx, yy]) # Reduction helpers. if True: # There are several reduction helper functions provided in # Halide.h, which compute small reductions and schedule them # innermost into their consumer. The most useful one is # "sum". f1 = hl.Func("f1") r = hl.RDom([(0, 100)]) f1[x] = hl.sum(r + x) * 7 # Sum creates a small anonymous hl.Func to do the reduction. It's # equivalent to: f2, anon = hl.Func("f2"), hl.Func("anon") anon[x] = 0 anon[x] += r + x f2[x] = anon[x] * 7 # So even though f1 references a reduction domain, it is a # pure function. The reduction domain has been swallowed to # define the inner anonymous reduction. halide_result_1 = f1.realize(10) halide_result_2 = f2.realize(10) # The equivalent Python is: py_result = np.empty((10), dtype=np.int) for xx in range(10): anon = np.empty((1), dtype=np.int) anon[0] = 0 for rr in range(100): anon[0] += rr + xx py_result[xx] = anon[0] * 7 # Check they all match. for xx in range(10): assert halide_result_1[xx] == py_result[xx], \ "halide_result_1(%d) = %d instead of %d" % (xx, halide_result_1[xx], py_result[xx]) assert halide_result_2[xx] == py_result[xx], \ "halide_result_2(%d) = %d instead of %d" % (xx, halide_result_2[xx], py_result[xx]) print("Success!") return 0
def main(): # Declare some Vars to use below. x, y = hl.Var("x"), hl.Var("y") # Load a grayscale image to use as an input. image_path = os.path.join(os.path.dirname(__file__), "../../tutorial/images/gray.png") input_data = imread(image_path) if True: # making the image smaller to go faster input_data = input_data[:160, :150] assert input_data.dtype == np.uint8 input = hl.Buffer(input_data) # You can define a hl.Func in multiple passes. Let's see a toy # example first. if True: # The first definition must be one like we have seen already # - a mapping from Vars to an hl.Expr: f = hl.Func("f") f[x, y] = x + y # We call this first definition the "pure" definition. # But the later definitions can include computed expressions on # both sides. The simplest example is modifying a single point: f[3, 7] = 42 # We call these extra definitions "update" definitions, or # "reduction" definitions. A reduction definition is an # update definition that recursively refers back to the # function's current value at the same site: if False: e = f[x, y] + 17 print("f[x, y] + 17", e) print("(f[x, y] + 17).type()", e.type()) print("(f[x, y]).type()", f[x, y].type()) f[x, y] = f[x, y] + 17 # If we confine our update to a single row, we can # recursively refer to values in the same column: f[x, 3] = f[x, 0] * f[x, 10] # Similarly, if we confine our update to a single column, we # can recursively refer to other values in the same row. f[0, y] = f[0, y] / f[3, y] # The general rule is: Each hl.Var used in an update definition # must appear unadorned in the same position as in the pure # definition in all references to the function on the left- # and right-hand sides. So the following definitions are # legal updates: f[x, 17] = x + 8 # x is used, so all uses of f must have x as the first argument. f[0, y] = y * 8 # y is used, so all uses of f must have y as the second argument. f[x, x + 1] = x + 8 f[y / 2, y] = f[0, y] * 17 # But these ones would cause an error: # f[x, 0) = f[x + 1, 0) <- First argument to f on the right-hand-side must be 'x', not 'x + 1'. # f[y, y + 1) = y + 8 <- Second argument to f on the left-hand-side must be 'y', not 'y + 1'. # f[y, x) = y - x <- Arguments to f on the left-hand-side are in the wrong places. # f[3, 4) = x + y <- Free variables appear on the right-hand-side but not the left-hand-side. # We'll realize this one just to make sure it compiles. The # second-to-last definition forces us to realize over a # domain that is taller than it is wide. f.realize(100, 101) # For each realization of f, each step runs in its entirety # before the next one begins. Let's trace the loads and # stores for a simpler example: g = hl.Func("g") g[x, y] = x + y # Pure definition g[2, 1] = 42 # First update definition g[x, 0] = g[x, 1] # Second update definition g.trace_loads() g.trace_stores() g.realize(4, 4) # Reading the log, we see that each pass is applied in turn. The equivalent C is: result = np.empty((4, 4), dtype=np.int) # Pure definition for yy in range(4): for xx in range(4): result[yy][xx] = xx + yy # First update definition result[1][2] = 42 # Second update definition for xx in range(4): result[0][xx] = result[1][xx] # end of section # Putting update passes inside loops. if True: # Starting with this pure definition: f = hl.Func("f") f[x, y] = x + y # Say we want an update that squares the first fifty rows. We # could do this by adding 50 update definitions: # f[x, 0) = f[x, 0) * f[x, 0) # f[x, 1) = f[x, 1) * f[x, 1) # f[x, 2) = f[x, 2) * f[x, 2) # ... # f[x, 49) = f[x, 49) * f[x, 49) # Or equivalently using a compile-time loop in our C++: # for (int i = 0 i < 50 i++) { # f[x, i) = f[x, i) * f[x, i) # # But it's more manageable and more flexible to put the loop # in the generated code. We do this by defining a "reduction # domain" and using it inside an update definition: r = hl.RDom(0, 50) f[x, r] = f[x, r] * f[x, r] halide_result = f.realize(100, 100) # The equivalent C is: c_result = np.empty((100, 100), dtype=np.int) for yy in range(100): for xx in range(100): c_result[yy][xx] = xx + yy for xx in range(100): for rr in range(50): # The loop over the reduction domain occurs inside of # the loop over any pure variables used in the update # step: c_result[rr][xx] = c_result[rr][xx] * c_result[rr][xx] # Check the results match: for yy in range(100): for xx in range(100): if halide_result(xx, yy) != c_result[yy][xx]: raise Exception( "halide_result(%d, %d) = %d instead of %d" % (xx, yy, halide_result(xx, yy), c_result[yy][xx])) return -1 # Now we'll examine a real-world use for an update definition: # computing a histogram. if True: # Some operations on images can't be cleanly expressed as a pure # function from the output coordinates to the value stored # there. The classic example is computing a histogram. The # natural way to do it is to iterate over the input image, # updating histogram buckets. Here's how you do that in Halide: histogram = hl.Func("histogram") # Histogram buckets start as zero. histogram[x] = 0 # Define a multi-dimensional reduction domain over the input image: r = hl.RDom(0, input.width(), 0, input.height()) # For every point in the reduction domain, increment the # histogram bucket corresponding to the intensity of the # input image at that point. histogram[input[r.x, r.y]] += 1 halide_result = histogram.realize(256) # The equivalent C is: c_result = np.empty((256), dtype=np.int) for xx in range(256): c_result[xx] = 0 for r_y in range(input.height()): for r_x in range(input.width()): c_result[input_data[r_x, r_y]] += 1 # Check the answers agree: for xx in range(256): if c_result[xx] != halide_result(xx): raise Exception("halide_result(%d) = %d instead of %d" % (xx, halide_result(xx), c_result[xx])) return -1 # Scheduling update steps if True: # The pure variables in an update step and can be # parallelized, vectorized, split, etc as usual. # Vectorizing, splitting, or parallelize the variables that # are part of the reduction domain is trickier. We'll cover # that in a later lesson. # Consider the definition: f = hl.Func("x") f[x, y] = x * y # Set the second row to equal the first row. f[x, 1] = f[x, 0] # Set the second column to equal the first column plus 2. f[1, y] = f[0, y] + 2 # The pure variables in each stage can be scheduled # independently. To control the pure definition, we schedule # as we have done in the past. The following code vectorizes # and parallelizes the pure definition only. f.vectorize(x, 4).parallel(y) # We use hl.Func::update(int) to get a handle to an update step # for the purposes of scheduling. The following line # vectorizes the first update step across x. We can't do # anything with y for this update step, because it doesn't # use y. f.update(0).vectorize(x, 4) # Now we parallelize the second update step in chunks of size # 4. yo, yi = hl.Var("yo"), hl.Var("yi") f.update(1).split(y, yo, yi, 4).parallel(yo) halide_result = f.realize(16, 16) # Here's the equivalent (serial) C: c_result = np.empty((16, 16), dtype=np.int) # Pure step. Vectorized in x and parallelized in y. for yy in range(16): # Should be a parallel for loop for x_vec in range(4): xx = [x_vec * 4, x_vec * 4 + 1, x_vec * 4 + 2, x_vec * 4 + 3] c_result[yy][xx[0]] = xx[0] * yy c_result[yy][xx[1]] = xx[1] * yy c_result[yy][xx[2]] = xx[2] * yy c_result[yy][xx[3]] = xx[3] * yy # First update. Vectorized in x. for x_vec in range(4): xx = [x_vec * 4, x_vec * 4 + 1, x_vec * 4 + 2, x_vec * 4 + 3] c_result[1][xx[0]] = c_result[0][xx[0]] c_result[1][xx[1]] = c_result[0][xx[1]] c_result[1][xx[2]] = c_result[0][xx[2]] c_result[1][xx[3]] = c_result[0][xx[3]] # Second update. Parallelized in chunks of size 4 in y. for yo in range(4): # Should be a parallel for loop for yi in range(4): yy = yo * 4 + yi c_result[yy][1] = c_result[yy][0] + 2 # Check the C and Halide results match: for yy in range(16): for xx in range(16): if halide_result(xx, yy) != c_result[yy][xx]: raise Exception( "halide_result(%d, %d) = %d instead of %d" % (xx, yy, halide_result(xx, yy), c_result[yy][xx])) return -1 # That covers how to schedule the variables within a hl.Func that # uses update steps, but what about producer-consumer # relationships that involve compute_at and store_at? Let's # examine a reduction as a producer, in a producer-consumer pair. if True: # Because an update does multiple passes over a stored array, # it's not meaningful to inline them. So the default schedule # for them does the closest thing possible. It computes them # in the innermost loop of their consumer. Consider this # trivial example: producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 producer[x] += 1 consumer[x] = 2 * producer[x] halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) for xx in range(10): producer_storage = np.empty((1), dtype=np.int) # Pure step for producer producer_storage[0] = xx * 17 # Update step for producer producer_storage[0] = producer_storage[0] + 1 # Pure step for consumer c_result[xx] = 2 * producer_storage[0] # Check the results match for xx in range(10): if halide_result(xx) != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % (xx, halide_result(xx), c_result[xx])) return -1 # For all other compute_at/store_at options, the reduction # gets placed where you would expect, somewhere in the loop # nest of the consumer. # Now let's consider a reduction as a consumer in a # producer-consumer pair. This is a little more involved. if True: if True: # Case 1: The consumer references the producer in the pure step only. producer, consumer = hl.Func("producer"), hl.Func("consumer") # The producer is pure. producer[x] = x * 17 consumer[x] = 2 * producer[x] consumer[x] += 1 # The valid schedules for the producer in this case are # the default schedule - inlined, and also: # # 1) producer.compute_at(x), which places the computation of # the producer inside the loop over x in the pure step of the # consumer. # # 2) producer.compute_root(), which computes all of the # producer ahead of time. # # 3) producer.store_root().compute_at(x), which allocates # space for the consumer outside the loop over x, but fills # it in as needed inside the loop. # # Let's use option 1. producer.compute_at(consumer, x) halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range(10): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] = 2 * producer_storage[0] # Update step for the consumer for xx in range(10): c_result[xx] += 1 # All of the pure step is evaluated before any of the # update step, so there are two separate loops over x. # Check the results match for xx in range(10): if halide_result(xx) != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % (xx, halide_result(xx), c_result[xx])) return -1 if True: # Case 2: The consumer references the producer in the update step only producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 consumer[x] = x consumer[x] += producer[x] # Again we compute the producer per x coordinate of the # consumer. This places producer code inside the update # step of the producer, because that's the only step that # uses the producer. producer.compute_at(consumer, x) # Note however, that we didn't say: # # producer.compute_at(consumer.update(0), x). # # Scheduling is done with respect to Vars of a hl.Func, and # the Vars of a hl.Func are shared across the pure and # update steps. halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range(10): c_result[xx] = xx # Update step for the consumer for xx in range(10): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] += producer_storage[0] # Check the results match for xx in range(10): if halide_result(xx) != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % (xx, halide_result(xx), c_result[xx])) return -1 if True: # Case 3: The consumer references the producer in # multiple steps that share common variables producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 consumer[x] = producer[x] * x consumer[x] += producer[x] # Again we compute the producer per x coordinate of the # consumer. This places producer code inside both the # pure and the update step of the producer. So there ends # up being two separate realizations of the producer, and # redundant work occurs. producer.compute_at(consumer, x) halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range(10): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] = producer_storage[0] * xx # Update step for the consumer for xx in range(10): # Another copy of the pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] += producer_storage[0] # Check the results match for xx in range(10): if halide_result(xx) != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % (xx, halide_result(xx), c_result[xx])) return -1 if True: # Case 4: The consumer references the producer in # multiple steps that do not share common variables producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x, y] = x * y consumer[x, y] = x + y consumer[x, 0] = producer[x, x - 1] consumer[0, y] = producer[y, y - 1] # In this case neither producer.compute_at(consumer, x) # nor producer.compute_at(consumer, y) will work, because # either one fails to cover one of the uses of the # producer. So we'd have to inline producer, or use # producer.compute_root(). # Let's say we really really want producer to be # compute_at the inner loops of both consumer update # steps. Halide doesn't allow multiple different # schedules for a single hl.Func, but we can work around it # by making two wrappers around producer, and scheduling # those instead: # Attempt 2: producer_wrapper_1, producer_wrapper_2, consumer_2 = hl.Func( ), hl.Func(), hl.Func() producer_wrapper_1[x, y] = producer[x, y] producer_wrapper_2[x, y] = producer[x, y] consumer_2[x, y] = x + y consumer_2[x, 0] += producer_wrapper_1[x, x - 1] consumer_2[0, y] += producer_wrapper_2[y, y - 1] # The wrapper functions give us two separate handles on # the producer, so we can schedule them differently. producer_wrapper_1.compute_at(consumer_2, x) producer_wrapper_2.compute_at(consumer_2, y) halide_result = consumer_2.realize(10, 10) # The equivalent C is: c_result = np.empty((10, 10), dtype=np.int) # Pure step for the consumer for yy in range(10): for xx in range(10): c_result[yy][xx] = xx + yy # First update step for consumer for xx in range(10): producer_wrapper_1_storage = np.empty((1), dtype=np.int) producer_wrapper_1_storage[0] = xx * (xx - 1) c_result[0][xx] += producer_wrapper_1_storage[0] # Second update step for consumer for yy in range(10): producer_wrapper_2_storage = np.empty((1), dtype=np.int) producer_wrapper_2_storage[0] = yy * (yy - 1) c_result[yy][0] += producer_wrapper_2_storage[0] # Check the results match for yy in range(10): for xx in range(10): if halide_result(xx, yy) != c_result[yy][xx]: print("halide_result(%d, %d) = %d instead of %d", xx, yy, halide_result(xx, yy), c_result[yy][xx]) return -1 if True: # Case 5: Scheduling a producer under a reduction domain # variable of the consumer. # We are not just restricted to scheduling producers at # the loops over the pure variables of the consumer. If a # producer is only used within a loop over a reduction # domain (hl.RDom) variable, we can also schedule the # producer there. producer, consumer = hl.Func("producer"), hl.Func("consumer") r = hl.RDom(0, 5) producer[x] = x * 17 consumer[x] = x + 10 consumer[x] += r + producer[x + r] producer.compute_at(consumer, r) halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer. for xx in range(10): c_result[xx] = xx + 10 # Update step for the consumer. for xx in range(10): for rr in range( 5 ): # The loop over the reduction domain is always the inner loop. # We've schedule the storage and computation of # the producer here. We just need a single value. producer_storage = np.empty((1), dtype=np.int) # Pure step of the producer. producer_storage[0] = (xx + rr) * 17 # Now use it in the update step of the consumer. c_result[xx] += rr + producer_storage[0] # Check the results match for xx in range(10): if halide_result(xx) != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % (xx, halide_result(xx), c_result[xx])) return -1 # A real-world example of a reduction inside a producer-consumer chain. if True: # The default schedule for a reduction is a good one for # convolution-like operations. For example, the following # computes a 5x5 box-blur of our grayscale test image with a # hl.clamp-to-edge boundary condition: # First add the boundary condition. clamped = hl.repeat_edge(input) # Define a 5x5 box that starts at (-2, -2) r = hl.RDom(-2, 5, -2, 5) # Compute the 5x5 sum around each pixel. local_sum = hl.Func("local_sum") local_sum[x, y] = 0 # Compute the sum as a 32-bit integer local_sum[x, y] += clamped[x + r.x, y + r.y] # Divide the sum by 25 to make it an average blurry = hl.Func("blurry") blurry[x, y] = hl.cast(hl.UInt(8), local_sum[x, y] / 25) halide_result = blurry.realize(input.width(), input.height()) # The default schedule will inline 'clamped' into the update # step of 'local_sum', because clamped only has a pure # definition, and so its default schedule is fully-inlined. # We will then compute local_sum per x coordinate of blurry, # because the default schedule for reductions is # compute-innermost. Here's the equivalent C: #cast_to_uint8 = lambda x_: np.array([x_], dtype=np.uint8)[0] local_sum = np.empty((1), dtype=np.int32) c_result = hl.Buffer(hl.UInt(8), input.width(), input.height()) for yy in range(input.height()): for xx in range(input.width()): # FIXME this loop is quite slow # Pure step of local_sum local_sum[0] = 0 # Update step of local_sum for r_y in range(-2, 2 + 1): for r_x in range(-2, 2 + 1): # The clamping has been inlined into the update step. clamped_x = min(max(xx + r_x, 0), input.width() - 1) clamped_y = min(max(yy + r_y, 0), input.height() - 1) local_sum[0] += input(clamped_x, clamped_y) # Pure step of blurry #c_result(x, y) = (uint8_t)(local_sum[0] / 25) #c_result[xx, yy] = cast_to_uint8(local_sum[0] / 25) c_result[xx, yy] = int(local_sum[0] / 25) # hl.cast done internally # Check the results match for yy in range(input.height()): for xx in range(input.width()): if halide_result(xx, yy) != c_result(xx, yy): raise Exception( "halide_result(%d, %d) = %d instead of %d" % (xx, yy, halide_result(xx, yy), c_result(xx, yy))) return -1 # Reduction helpers. if True: # There are several reduction helper functions provided in # Halide.h, which compute small reductions and schedule them # innermost into their consumer. The most useful one is # "sum". f1 = hl.Func("f1") r = hl.RDom(0, 100) f1[x] = hl.sum(r + x) * 7 # Sum creates a small anonymous hl.Func to do the reduction. It's equivalent to: f2, anon = hl.Func("f2"), hl.Func("anon") anon[x] = 0 anon[x] += r + x f2[x] = anon[x] * 7 # So even though f1 references a reduction domain, it is a # pure function. The reduction domain has been swallowed to # define the inner anonymous reduction. halide_result_1 = f1.realize(10) halide_result_2 = f2.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) for xx in range(10): anon = np.empty((1), dtype=np.int) anon[0] = 0 for rr in range(100): anon[0] += rr + xx c_result[xx] = anon[0] * 7 # Check they all match. for xx in range(10): if halide_result_1(xx) != c_result[xx]: print("halide_result_1(%d) = %d instead of %d", x, halide_result_1(x), c_result[x]) return -1 if halide_result_2(xx) != c_result[xx]: print("halide_result_2(%d) = %d instead of %d", x, halide_result_2(x), c_result[x]) return -1 # A complex example that uses reduction helpers. if False: # non-sense to port SSE code to python, skipping this test # Other reduction helpers include "product", "minimum", # "maximum", "hl.argmin", and "argmax". Using hl.argmin and argmax # requires understanding tuples, which come in a later # lesson. Let's use minimum and maximum to compute the local # spread of our grayscale image. # First, add a boundary condition to the input. clamped = hl.Func("clamped") x_clamped = hl.clamp(x, 0, input.width() - 1) y_clamped = hl.clamp(y, 0, input.height() - 1) clamped[x, y] = input[x_clamped, y_clamped] box = hl.RDom(-2, 5, -2, 5) # Compute the local maximum minus the local minimum: spread = hl.Func("spread") spread[x, y] = (maximum(clamped(x + box.x, y + box.y)) - minimum(clamped(x + box.x, y + box.y))) # Compute the result in strips of 32 scanlines yo, yi = hl.Var("yo"), hl.Var("yi") spread.split(y, yo, yi, 32).parallel(yo) # Vectorize across x within the strips. This implicitly # vectorizes stuff that is computed within the loop over x in # spread, which includes our minimum and maximum helpers, so # they get vectorized too. spread.vectorize(x, 16) # We'll apply the boundary condition by padding each scanline # as we need it in a circular buffer (see lesson 08). clamped.store_at(spread, yo).compute_at(spread, yi) halide_result = spread.realize(input.width(), input.height()) # The C equivalent is almost too horrible to contemplate (and # took me a long time to debug). This time I want to time # both the Halide version and the C version, so I'll use sse # intrinsics for the vectorization, and openmp to do the # parallel for loop (you'll need to compile with -fopenmp or # similar to get correct timing). #ifdef __SSE2__ # Don't include the time required to allocate the output buffer. c_result = hl.Buffer(hl.UInt(8), input.width(), input.height()) #ifdef _OPENMP t1 = datetime.now() #endif # Run this one hundred times so we can average the timing results. for iters in range(100): pass # #pragma omp parallel for # for yo in range((input.height() + 31)/32): # y_base = hl.min(yo * 32, input.height() - 32) # # # Compute clamped in a circular buffer of size 8 # # (smallest power of two greater than 5). Each thread # # needs its own allocation, so it must occur here. # # clamped_width = input.width() + 4 # clamped_storage = np.empty((clamped_width * 8), dtype=np.uint8) # # for yi in range(32): # y = y_base + yi # # uint8_t *output_row = &c_result(0, y) # # # Compute clamped for this scanline, skipping rows # # already computed within this slice. # int min_y_clamped = (yi == 0) ? (y - 2) : (y + 2) # int max_y_clamped = (y + 2) # for (int cy = min_y_clamped cy <= max_y_clamped cy++) { # # Figure out which row of the circular buffer # # we're filling in using bitmasking: # uint8_t *clamped_row = clamped_storage + (cy & 7) * clamped_width # # # Figure out which row of the input we're reading # # from by clamping the y coordinate: # int clamped_y = std::hl.min(std::hl.max(cy, 0), input.height()-1) # uint8_t *input_row = &input(0, clamped_y) # # # Fill it in with the padding. # for (int x = -2 x < input.width() + 2 ): # int clamped_x = std::hl.min(std::hl.max(x, 0), input.width()-1) # *clamped_row++ = input_row[clamped_x] # # # # # Now iterate over vectors of x for the pure step of the output. # for (int x_vec = 0 x_vec < (input.width() + 15)/16 x_vec++) { # int x_base = std::hl.min(x_vec * 16, input.width() - 16) # # # Allocate storage for the minimum and maximum # # helpers. One vector is enough. # __m128i minimum_storage, maximum_storage # # # The pure step for the maximum is a vector of zeros # maximum_storage = (__m128i)_mm_setzero_ps() # # # The update step for maximum # for (int max_y = y - 2 max_y <= y + 2 max_y++) { # uint8_t *clamped_row = clamped_storage + (max_y & 7) * clamped_width # for (int max_x = x_base - 2 max_x <= x_base + 2 max_): # __m128i v = _mm_loadu_si128((__m128i const *)(clamped_row + max_x + 2)) # maximum_storage = _mm_max_epu8(maximum_storage, v) # # # # # The pure step for the minimum is a vector of # # ones. Create it by comparing something to # # itself. # minimum_storage = (__m128i)_mm_cmpeq_ps(_mm_setzero_ps(), # _mm_setzero_ps()) # # # The update step for minimum. # for (int min_y = y - 2 min_y <= y + 2 min_y++) { # uint8_t *clamped_row = clamped_storage + (min_y & 7) * clamped_width # for (int min_x = x_base - 2 min_x <= x_base + 2 min_): # __m128i v = _mm_loadu_si128((__m128i const *)(clamped_row + min_x + 2)) # minimum_storage = _mm_min_epu8(minimum_storage, v) # # # # # Now compute the spread. # __m128i spread = _mm_sub_epi8(maximum_storage, minimum_storage) # # # Store it. # _mm_storeu_si128((__m128i *)(output_row + x_base), spread) # # # # del clamped_storage # # end of hundred iterations # Skip the timing comparison if we don't have openmp # enabled. Otherwise it's unfair to C. #ifdef _OPENMP t2 = datetime.now() # Now run the Halide version again without the # jit-compilation overhead. Also run it one hundred times. for iters in range(100): spread.realize(halide_result) t3 = datetime.now() # Report the timings. On my machine they both take about 3ms # for the 4-megapixel input (fast!), which makes sense, # because they're using the same vectorization and # parallelization strategy. However I find the Halide easier # to read, write, debug, modify, and port. print("Halide spread took %f ms. C equivalent took %f ms" % ((t3 - t2).total_seconds() * 1000, (t2 - t1).total_seconds() * 1000)) #endif # _OPENMP # Check the results match: for yy in range(input.height()): for xx in range(input.width()): if halide_result(xx, yy) != c_result(xx, yy): raise Exception( "halide_result(%d, %d) = %d instead of %d" % (xx, yy, halide_result(xx, yy), c_result(xx, yy))) return -1 #endif # __SSE2__ else: print("(Skipped the SSE2 section of the code, " "since non-sense in python world.)") print("Success!") return 0
def findStereoCorrespondence(left, right, SADWindowSize, minDisparity, numDisparities, xmin, xmax, ymin, ymax, x_tile_size=32, y_tile_size=32, test=False, uniquenessRatio=0.15, disp12MaxDiff=1): """ Returns Func (left: Func, right: Func) """ x, y, c, d = Var("x"), Var("y"), Var("c"), Var("d") diff = Func("diff") diff[d, x, y] = h.cast(UInt(16), h.abs(left[x, y] - right[x - d, y])) win2 = SADWindowSize / 2 diff_T = Func("diff_T") xi, xo, yi, yo = Var("xi"), Var("xo"), Var("yi"), Var("yo") diff_T[d, xi, yi, xo, yo] = diff[d, xi + xo * x_tile_size + xmin, yi + yo * y_tile_size + ymin] cSAD, vsum = Func("cSAD"), Func("vsum") rk = RDom(-win2, SADWindowSize, "rk") rxi, ryi = RDom(1, x_tile_size - 1, "rxi"), RDom(1, y_tile_size - 1, "ryi") if test: vsum[d, xi, yi, xo, yo] = h.sum(diff_T[d, xi, yi + rk, xo, yo]) cSAD[d, xi, yi, xo, yo] = h.sum(vsum[d, xi + rk, yi, xo, yo]) else: vsum[d, xi, yi, xo, yo] = h.select(yi != 0, h.cast(UInt(16), 0), h.sum(diff_T[d, xi, rk, xo, yo])) vsum[d, xi, ryi, xo, yo] = vsum[d, xi, ryi - 1, xo, yo] + diff_T[ d, xi, ryi + win2, xo, yo] - diff_T[d, xi, ryi - win2 - 1, xo, yo] cSAD[d, xi, yi, xo, yo] = h.select(xi != 0, h.cast(UInt(16), 0), h.sum(vsum[d, rk, yi, xo, yo])) cSAD[d, rxi, yi, xo, yo] = cSAD[d, rxi - 1, yi, xo, yo] + vsum[d, rxi + win2, yi, xo, yo] - vsum[d, rxi - win2 - 1, yi, xo, yo] rd = RDom(minDisparity, numDisparities) disp_left = Func("disp_left") disp_left[xi, yi, xo, yo] = h.Tuple(h.cast(UInt(16), minDisparity), h.cast(UInt(16), (2 << 16) - 1)) disp_left[xi, yi, xo, yo] = h.tuple_select( cSAD[rd, xi, yi, xo, yo] < disp_left[xi, yi, xo, yo][1], h.Tuple(h.cast(UInt(16), rd), cSAD[rd, xi, yi, xo, yo]), h.Tuple(disp_left[xi, yi, xo, yo])) FILTERED = -16 disp = Func("disp") disp[x, y] = h.select( # x > xmax-xmin or y > ymax-ymin, x < xmax, h.cast( UInt(16), disp_left[x % x_tile_size, y % y_tile_size, x / x_tile_size, y / y_tile_size][0]), h.cast(UInt(16), FILTERED)) # Schedule vector_width = 8 disp.compute_root() \ .tile(x, y, xo, yo, xi, yi, x_tile_size, y_tile_size).reorder(xi, yi, xo, yo) \ .vectorize(xi, vector_width).parallel(xo).parallel(yo) # reorder storage disp_left.reorder_storage(xi, yi, xo, yo) diff_T.reorder_storage(xi, yi, xo, yo, d) vsum.reorder_storage(xi, yi, xo, yo, d) cSAD.reorder_storage(xi, yi, xo, yo, d) disp_left.compute_at(disp, xo).reorder(xi, yi, xo, yo) \ .vectorize(xi, vector_width) \ .update() \ .reorder(xi, yi, rd, xo, yo).vectorize(xi, vector_width) if test: cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) else: cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \ .update() \ .reorder(yi, rxi, xo, yo, d).vectorize(yi, vector_width) vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \ .update() \ .reorder(xi, ryi, xo, yo, d).vectorize(xi, vector_width) return disp
def main(): # Declare some Vars to use below. x, y = hl.Var ("x"), hl.Var ("y") # Load a grayscale image to use as an input. image_path = os.path.join(os.path.dirname(__file__), "../../tutorial/images/gray.png") input_data = imread(image_path) if True: # making the image smaller to go faster input_data = input_data[:160, :150] assert input_data.dtype == np.uint8 input = hl.Buffer(input_data) # You can define a hl.Func in multiple passes. Let's see a toy # example first. if True: # The first definition must be one like we have seen already # - a mapping from Vars to an hl.Expr: f = hl.Func("f") f[x, y] = x + y # We call this first definition the "pure" definition. # But the later definitions can include computed expressions on # both sides. The simplest example is modifying a single point: f[3, 7] = 42 # We call these extra definitions "update" definitions, or # "reduction" definitions. A reduction definition is an # update definition that recursively refers back to the # function's current value at the same site: if False: e = f[x, y] + 17 print("f[x, y] + 17", e) print("(f[x, y] + 17).type()", e.type()) print("(f[x, y]).type()", f[x,y].type()) f[x, y] = f[x, y] + 17 # If we confine our update to a single row, we can # recursively refer to values in the same column: f[x, 3] = f[x, 0] * f[x, 10] # Similarly, if we confine our update to a single column, we # can recursively refer to other values in the same row. f[0, y] = f[0, y] / f[3, y] # The general rule is: Each hl.Var used in an update definition # must appear unadorned in the same position as in the pure # definition in all references to the function on the left- # and right-hand sides. So the following definitions are # legal updates: f[x, 17] = x + 8 # x is used, so all uses of f must have x as the first argument. f[0, y] = y * 8 # y is used, so all uses of f must have y as the second argument. f[x, x + 1] = x + 8 f[y/2, y] = f[0, y] * 17 # But these ones would cause an error: # f[x, 0) = f[x + 1, 0) <- First argument to f on the right-hand-side must be 'x', not 'x + 1'. # f[y, y + 1) = y + 8 <- Second argument to f on the left-hand-side must be 'y', not 'y + 1'. # f[y, x) = y - x <- Arguments to f on the left-hand-side are in the wrong places. # f[3, 4) = x + y <- Free variables appear on the right-hand-side but not the left-hand-side. # We'll realize this one just to make sure it compiles. The # second-to-last definition forces us to realize over a # domain that is taller than it is wide. f.realize(100, 101) # For each realization of f, each step runs in its entirety # before the next one begins. Let's trace the loads and # stores for a simpler example: g = hl.Func("g") g[x, y] = x + y # Pure definition g[2, 1] = 42 # First update definition g[x, 0] = g[x, 1] # Second update definition g.trace_loads() g.trace_stores() g.realize(4, 4) # Reading the log, we see that each pass is applied in turn. The equivalent C is: result = np.empty( (4,4), dtype=np.int) # Pure definition for yy in range(4): for xx in range(4): result[yy][xx] = xx + yy # First update definition result[1][2] = 42 # Second update definition for xx in range(4): result[0][xx] = result[1][xx] # end of section # Putting update passes inside loops. if True: # Starting with this pure definition: f = hl.Func("f") f[x, y] = x + y # Say we want an update that squares the first fifty rows. We # could do this by adding 50 update definitions: # f[x, 0) = f[x, 0) * f[x, 0) # f[x, 1) = f[x, 1) * f[x, 1) # f[x, 2) = f[x, 2) * f[x, 2) # ... # f[x, 49) = f[x, 49) * f[x, 49) # Or equivalently using a compile-time loop in our C++: # for (int i = 0 i < 50 i++) { # f[x, i) = f[x, i) * f[x, i) # # But it's more manageable and more flexible to put the loop # in the generated code. We do this by defining a "reduction # domain" and using it inside an update definition: r = hl.RDom([(0, 50)]) f[x, r] = f[x, r] * f[x, r] halide_result = f.realize(100, 100) # The equivalent C is: c_result = np.empty((100, 100), dtype=np.int) for yy in range(100): for xx in range(100): c_result[yy][xx] = xx + yy for xx in range(100): for rr in range(50): # The loop over the reduction domain occurs inside of # the loop over any pure variables used in the update # step: c_result[rr][xx] = c_result[rr][xx] * c_result[rr][xx] # Check the results match: for yy in range(100): for xx in range(100): if halide_result[xx, yy] != c_result[yy][xx]: raise Exception("halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result[xx, yy], c_result[yy][xx])) return -1 # Now we'll examine a real-world use for an update definition: # computing a histogram. if True: # Some operations on images can't be cleanly expressed as a pure # function from the output coordinates to the value stored # there. The classic example is computing a histogram. The # natural way to do it is to iterate over the input image, # updating histogram buckets. Here's how you do that in Halide: histogram = hl.Func("histogram") # Histogram buckets start as zero. histogram[x] = 0 # Define a multi-dimensional reduction domain over the input image: r = hl.RDom([(0, input.width()), (0, input.height())]) # For every point in the reduction domain, increment the # histogram bucket corresponding to the intensity of the # input image at that point. histogram[input[r.x, r.y]] += 1 halide_result = histogram.realize(256) # The equivalent C is: c_result = np.empty((256), dtype=np.int) for xx in range(256): c_result[xx] = 0 for r_y in range(input.height()): for r_x in range(input.width()): c_result[input_data[r_x, r_y]] += 1 # Check the answers agree: for xx in range(256): if c_result[xx] != halide_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % ( xx, halide_result[xx], c_result[xx])) return -1 # Scheduling update steps if True: # The pure variables in an update step and can be # parallelized, vectorized, split, etc as usual. # Vectorizing, splitting, or parallelize the variables that # are part of the reduction domain is trickier. We'll cover # that in a later lesson. # Consider the definition: f = hl.Func("x") f[x, y] = x*y # Set the second row to equal the first row. f[x, 1] = f[x, 0] # Set the second column to equal the first column plus 2. f[1, y] = f[0, y] + 2 # The pure variables in each stage can be scheduled # independently. To control the pure definition, we schedule # as we have done in the past. The following code vectorizes # and parallelizes the pure definition only. f.vectorize(x, 4).parallel(y) # We use hl.Func::update(int) to get a handle to an update step # for the purposes of scheduling. The following line # vectorizes the first update step across x. We can't do # anything with y for this update step, because it doesn't # use y. f.update(0).vectorize(x, 4) # Now we parallelize the second update step in chunks of size # 4. yo, yi = hl.Var("yo"), hl.Var("yi") f.update(1).split(y, yo, yi, 4).parallel(yo) halide_result = f.realize(16, 16) # Here's the equivalent (serial) C: c_result = np.empty((16, 16), dtype=np.int) # Pure step. Vectorized in x and parallelized in y. for yy in range( 16): # Should be a parallel for loop for x_vec in range(4): xx = [x_vec*4, x_vec*4+1, x_vec*4+2, x_vec*4+3] c_result[yy][xx[0]] = xx[0] * yy c_result[yy][xx[1]] = xx[1] * yy c_result[yy][xx[2]] = xx[2] * yy c_result[yy][xx[3]] = xx[3] * yy # First update. Vectorized in x. for x_vec in range(4): xx = [x_vec*4, x_vec*4+1, x_vec*4+2, x_vec*4+3] c_result[1][xx[0]] = c_result[0][xx[0]] c_result[1][xx[1]] = c_result[0][xx[1]] c_result[1][xx[2]] = c_result[0][xx[2]] c_result[1][xx[3]] = c_result[0][xx[3]] # Second update. Parallelized in chunks of size 4 in y. for yo in range(4): # Should be a parallel for loop for yi in range(4): yy = yo*4 + yi c_result[yy][1] = c_result[yy][0] + 2 # Check the C and Halide results match: for yy in range( 16): for xx in range( 16 ): if halide_result[xx, yy] != c_result[yy][xx]: raise Exception("halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result[xx, yy], c_result[yy][xx])) return -1 # That covers how to schedule the variables within a hl.Func that # uses update steps, but what about producer-consumer # relationships that involve compute_at and store_at? Let's # examine a reduction as a producer, in a producer-consumer pair. if True: # Because an update does multiple passes over a stored array, # it's not meaningful to inline them. So the default schedule # for them does the closest thing possible. It computes them # in the innermost loop of their consumer. Consider this # trivial example: producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x*17 producer[x] += 1 consumer[x] = 2 * producer[x] halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) for xx in range(10): producer_storage = np.empty((1), dtype=np.int) # Pure step for producer producer_storage[0] = xx * 17 # Update step for producer producer_storage[0] = producer_storage[0] + 1 # Pure step for consumer c_result[xx] = 2 * producer_storage[0] # Check the results match for xx in range( 10 ): if halide_result[xx] != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % ( xx, halide_result[xx], c_result[xx])) return -1 # For all other compute_at/store_at options, the reduction # gets placed where you would expect, somewhere in the loop # nest of the consumer. # Now let's consider a reduction as a consumer in a # producer-consumer pair. This is a little more involved. if True: if True: # Case 1: The consumer references the producer in the pure step only. producer, consumer = hl.Func("producer"), hl.Func("consumer") # The producer is pure. producer[x] = x*17 consumer[x] = 2 * producer[x] consumer[x] += 1 # The valid schedules for the producer in this case are # the default schedule - inlined, and also: # # 1) producer.compute_at(x), which places the computation of # the producer inside the loop over x in the pure step of the # consumer. # # 2) producer.compute_root(), which computes all of the # producer ahead of time. # # 3) producer.store_root().compute_at(x), which allocates # space for the consumer outside the loop over x, but fills # it in as needed inside the loop. # # Let's use option 1. producer.compute_at(consumer, x) halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range( 10 ): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] = 2 * producer_storage[0] # Update step for the consumer for xx in range( 10 ): c_result[xx] += 1 # All of the pure step is evaluated before any of the # update step, so there are two separate loops over x. # Check the results match for xx in range( 10 ): if halide_result[xx] != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % ( xx, halide_result[xx], c_result[xx])) return -1 if True: # Case 2: The consumer references the producer in the update step only producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 consumer[x] = x consumer[x] += producer[x] # Again we compute the producer per x coordinate of the # consumer. This places producer code inside the update # step of the producer, because that's the only step that # uses the producer. producer.compute_at(consumer, x) # Note however, that we didn't say: # # producer.compute_at(consumer.update(0), x). # # Scheduling is done with respect to Vars of a hl.Func, and # the Vars of a hl.Func are shared across the pure and # update steps. halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range( 10 ): c_result[xx] = xx # Update step for the consumer for xx in range( 10 ): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] += producer_storage[0] # Check the results match for xx in range( 10 ): if halide_result[xx] != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % ( xx, halide_result[xx], c_result[xx])) return -1 if True: # Case 3: The consumer references the producer in # multiple steps that share common variables producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x] = x * 17 consumer[x] = producer[x] * x consumer[x] += producer[x] # Again we compute the producer per x coordinate of the # consumer. This places producer code inside both the # pure and the update step of the producer. So there ends # up being two separate realizations of the producer, and # redundant work occurs. producer.compute_at(consumer, x) halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer for xx in range( 10 ): # Pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] = producer_storage[0] * xx # Update step for the consumer for xx in range( 10 ): # Another copy of the pure step for producer producer_storage = np.empty((1), dtype=np.int) producer_storage[0] = xx * 17 c_result[xx] += producer_storage[0] # Check the results match for xx in range( 10 ): if halide_result[xx] != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % ( xx, halide_result[xx], c_result[xx])) return -1 if True: # Case 4: The consumer references the producer in # multiple steps that do not share common variables producer, consumer = hl.Func("producer"), hl.Func("consumer") producer[x, y] = x*y consumer[x, y] = x + y consumer[x, 0] = producer[x, x-1] consumer[0, y] = producer[y, y-1] # In this case neither producer.compute_at(consumer, x) # nor producer.compute_at(consumer, y) will work, because # either one fails to cover one of the uses of the # producer. So we'd have to inline producer, or use # producer.compute_root(). # Let's say we really really want producer to be # compute_at the inner loops of both consumer update # steps. Halide doesn't allow multiple different # schedules for a single hl.Func, but we can work around it # by making two wrappers around producer, and scheduling # those instead: # Attempt 2: producer_wrapper_1, producer_wrapper_2, consumer_2 = hl.Func(), hl.Func(), hl.Func() producer_wrapper_1[x, y] = producer[x, y] producer_wrapper_2[x, y] = producer[x, y] consumer_2[x, y] = x + y consumer_2[x, 0] += producer_wrapper_1[x, x-1] consumer_2[0, y] += producer_wrapper_2[y, y-1] # The wrapper functions give us two separate handles on # the producer, so we can schedule them differently. producer_wrapper_1.compute_at(consumer_2, x) producer_wrapper_2.compute_at(consumer_2, y) halide_result = consumer_2.realize(10, 10) # The equivalent C is: c_result = np.empty((10, 10), dtype=np.int) # Pure step for the consumer for yy in range( 10): for xx in range( 10 ): c_result[yy][xx] = xx + yy # First update step for consumer for xx in range( 10 ): producer_wrapper_1_storage = np.empty((1), dtype=np.int) producer_wrapper_1_storage[0] = xx * (xx-1) c_result[0][xx] += producer_wrapper_1_storage[0] # Second update step for consumer for yy in range( 10): producer_wrapper_2_storage = np.empty((1), dtype=np.int) producer_wrapper_2_storage[0] = yy * (yy-1) c_result[yy][0] += producer_wrapper_2_storage[0] # Check the results match for yy in range( 10): for xx in range( 10 ): if halide_result[xx, yy] != c_result[yy][xx]: print("halide_result(%d, %d) = %d instead of %d", xx, yy, halide_result[xx, yy], c_result[yy][xx]) return -1 if True: # Case 5: Scheduling a producer under a reduction domain # variable of the consumer. # We are not just restricted to scheduling producers at # the loops over the pure variables of the consumer. If a # producer is only used within a loop over a reduction # domain (hl.RDom) variable, we can also schedule the # producer there. producer, consumer = hl.Func("producer"), hl.Func("consumer") r = hl.RDom([(0, 5)]) producer[x] = x * 17 consumer[x] = x + 10 consumer[x] += r + producer[x + r] producer.compute_at(consumer, r) halide_result = consumer.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) # Pure step for the consumer. for xx in range(10): c_result[xx] = xx + 10 # Update step for the consumer. for xx in range( 10 ): for rr in range(5): # The loop over the reduction domain is always the inner loop. # We've schedule the storage and computation of # the producer here. We just need a single value. producer_storage = np.empty((1), dtype=np.int) # Pure step of the producer. producer_storage[0] = (xx + rr) * 17 # Now use it in the update step of the consumer. c_result[xx] += rr + producer_storage[0] # Check the results match for xx in range( 10 ): if halide_result[xx] != c_result[xx]: raise Exception("halide_result(%d) = %d instead of %d" % ( xx, halide_result[xx], c_result[xx])) return -1 # A real-world example of a reduction inside a producer-consumer chain. if True: # The default schedule for a reduction is a good one for # convolution-like operations. For example, the following # computes a 5x5 box-blur of our grayscale test image with a # hl.clamp-to-edge boundary condition: # First add the boundary condition. clamped = hl.BoundaryConditions.repeat_edge(input) # Define a 5x5 box that starts at (-2, -2) r = hl.RDom([(-2, 5), (-2, 5)]) # Compute the 5x5 sum around each pixel. local_sum = hl.Func("local_sum") local_sum[x, y] = 0 # Compute the sum as a 32-bit integer local_sum[x, y] += clamped[x + r.x, y + r.y] # Divide the sum by 25 to make it an average blurry = hl.Func("blurry") blurry[x, y] = hl.cast(hl.UInt(8), local_sum[x, y] / 25) halide_result = blurry.realize(input.width(), input.height()) # The default schedule will inline 'clamped' into the update # step of 'local_sum', because clamped only has a pure # definition, and so its default schedule is fully-inlined. # We will then compute local_sum per x coordinate of blurry, # because the default schedule for reductions is # compute-innermost. Here's the equivalent C: #cast_to_uint8 = lambda x_: np.array([x_], dtype=np.uint8)[0] local_sum = np.empty((1), dtype=np.int32) c_result = hl.Buffer(hl.UInt(8), [input.width(), input.height()]) for yy in range(input.height()): for xx in range(input.width()): # FIXME this loop is quite slow # Pure step of local_sum local_sum[0] = 0 # Update step of local_sum for r_y in range(-2, 2+1): for r_x in range(-2, 2+1): # The clamping has been inlined into the update step. clamped_x = min(max(xx + r_x, 0), input.width()-1) clamped_y = min(max(yy + r_y, 0), input.height()-1) local_sum[0] += input[clamped_x, clamped_y] # Pure step of blurry #c_result(x, y) = (uint8_t)(local_sum[0] / 25) #c_result[xx, yy] = cast_to_uint8(local_sum[0] / 25) c_result[xx, yy] = int(local_sum[0] / 25) # hl.cast done internally # Check the results match for yy in range(input.height()): for xx in range(input.width()): if halide_result[xx, yy] != c_result[xx, yy]: raise Exception("halide_result(%d, %d) = %d instead of %d" % (xx, yy, halide_result[xx, yy], c_result[xx, yy])) return -1 # Reduction helpers. if True: # There are several reduction helper functions provided in # Halide.h, which compute small reductions and schedule them # innermost into their consumer. The most useful one is # "sum". f1 = hl.Func ("f1") r = hl.RDom([(0, 100)]) f1[x] = hl.sum(r + x) * 7 # Sum creates a small anonymous hl.Func to do the reduction. It's equivalent to: f2, anon = hl.Func("f2"), hl.Func("anon") anon[x] = 0 anon[x] += r + x f2[x] = anon[x] * 7 # So even though f1 references a reduction domain, it is a # pure function. The reduction domain has been swallowed to # define the inner anonymous reduction. halide_result_1 = f1.realize(10) halide_result_2 = f2.realize(10) # The equivalent C is: c_result = np.empty((10), dtype=np.int) for xx in range( 10 ): anon = np.empty((1), dtype=np.int) anon[0] = 0 for rr in range(100): anon[0] += rr + xx c_result[xx] = anon[0] * 7 # Check they all match. for xx in range( 10 ): if halide_result_1[xx] != c_result[xx]: print("halide_result_1(%d) = %d instead of %d", xx, halide_result_1[xx], c_result[xx]) return -1 if halide_result_2[xx] != c_result[xx]: print("halide_result_2(%d) = %d instead of %d", xx, halide_result_2[xx], c_result[xx]) return -1 # A complex example that uses reduction helpers. if False: # non-sense to port SSE code to python, skipping this test # Other reduction helpers include "product", "minimum", # "maximum", "hl.argmin", and "argmax". Using hl.argmin and argmax # requires understanding tuples, which come in a later # lesson. Let's use minimum and maximum to compute the local # spread of our grayscale image. # First, add a boundary condition to the input. clamped = hl.Func("clamped") x_clamped = hl.clamp(x, 0, input.width()-1) y_clamped = hl.clamp(y, 0, input.height()-1) clamped[x, y] = input[x_clamped, y_clamped] box = hl.RDom([(-2, 5), (-2, 5)]) # Compute the local maximum minus the local minimum: spread = hl.Func("spread") spread[x, y] = (maximum(clamped(x + box.x, y + box.y)) - minimum(clamped(x + box.x, y + box.y))) # Compute the result in strips of 32 scanlines yo, yi = hl.Var("yo"), hl.Var("yi") spread.split(y, yo, yi, 32).parallel(yo) # Vectorize across x within the strips. This implicitly # vectorizes stuff that is computed within the loop over x in # spread, which includes our minimum and maximum helpers, so # they get vectorized too. spread.vectorize(x, 16) # We'll apply the boundary condition by padding each scanline # as we need it in a circular buffer (see lesson 08). clamped.store_at(spread, yo).compute_at(spread, yi) halide_result = spread.realize(input.width(), input.height()) # The C equivalent is almost too horrible to contemplate (and # took me a long time to debug). This time I want to time # both the Halide version and the C version, so I'll use sse # intrinsics for the vectorization, and openmp to do the # parallel for loop (you'll need to compile with -fopenmp or # similar to get correct timing). #ifdef __SSE2__ # Don't include the time required to allocate the output buffer. c_result = hl.Buffer(hl.UInt(8), input.width(), input.height()) #ifdef _OPENMP t1 = datetime.now() #endif # Run this one hundred times so we can average the timing results. for iters in range(100): pass # #pragma omp parallel for # for yo in range((input.height() + 31)/32): # y_base = hl.min(yo * 32, input.height() - 32) # # # Compute clamped in a circular buffer of size 8 # # (smallest power of two greater than 5). Each thread # # needs its own allocation, so it must occur here. # # clamped_width = input.width() + 4 # clamped_storage = np.empty((clamped_width * 8), dtype=np.uint8) # # for yi in range(32): # y = y_base + yi # # uint8_t *output_row = &c_result(0, y) # # # Compute clamped for this scanline, skipping rows # # already computed within this slice. # int min_y_clamped = (yi == 0) ? (y - 2) : (y + 2) # int max_y_clamped = (y + 2) # for (int cy = min_y_clamped cy <= max_y_clamped cy++) { # # Figure out which row of the circular buffer # # we're filling in using bitmasking: # uint8_t *clamped_row = clamped_storage + (cy & 7) * clamped_width # # # Figure out which row of the input we're reading # # from by clamping the y coordinate: # int clamped_y = std::hl.min(std::hl.max(cy, 0), input.height()-1) # uint8_t *input_row = &input(0, clamped_y) # # # Fill it in with the padding. # for (int x = -2 x < input.width() + 2 ): # int clamped_x = std::hl.min(std::hl.max(x, 0), input.width()-1) # *clamped_row++ = input_row[clamped_x] # # # # # Now iterate over vectors of x for the pure step of the output. # for (int x_vec = 0 x_vec < (input.width() + 15)/16 x_vec++) { # int x_base = std::hl.min(x_vec * 16, input.width() - 16) # # # Allocate storage for the minimum and maximum # # helpers. One vector is enough. # __m128i minimum_storage, maximum_storage # # # The pure step for the maximum is a vector of zeros # maximum_storage = (__m128i)_mm_setzero_ps() # # # The update step for maximum # for (int max_y = y - 2 max_y <= y + 2 max_y++) { # uint8_t *clamped_row = clamped_storage + (max_y & 7) * clamped_width # for (int max_x = x_base - 2 max_x <= x_base + 2 max_): # __m128i v = _mm_loadu_si128((__m128i const *)(clamped_row + max_x + 2)) # maximum_storage = _mm_max_epu8(maximum_storage, v) # # # # # The pure step for the minimum is a vector of # # ones. Create it by comparing something to # # itself. # minimum_storage = (__m128i)_mm_cmpeq_ps(_mm_setzero_ps(), # _mm_setzero_ps()) # # # The update step for minimum. # for (int min_y = y - 2 min_y <= y + 2 min_y++) { # uint8_t *clamped_row = clamped_storage + (min_y & 7) * clamped_width # for (int min_x = x_base - 2 min_x <= x_base + 2 min_): # __m128i v = _mm_loadu_si128((__m128i const *)(clamped_row + min_x + 2)) # minimum_storage = _mm_min_epu8(minimum_storage, v) # # # # # Now compute the spread. # __m128i spread = _mm_sub_epi8(maximum_storage, minimum_storage) # # # Store it. # _mm_storeu_si128((__m128i *)(output_row + x_base), spread) # # # # del clamped_storage # # end of hundred iterations # Skip the timing comparison if we don't have openmp # enabled. Otherwise it's unfair to C. #ifdef _OPENMP t2 = datetime.now() # Now run the Halide version again without the # jit-compilation overhead. Also run it one hundred times. for iters in range(100): spread.realize(halide_result) t3 = datetime.now() # Report the timings. On my machine they both take about 3ms # for the 4-megapixel input (fast!), which makes sense, # because they're using the same vectorization and # parallelization strategy. However I find the Halide easier # to read, write, debug, modify, and port. print("Halide spread took %f ms. C equivalent took %f ms" % ( (t3 - t2).total_seconds() * 1000, (t2 - t1).total_seconds() * 1000)) #endif # _OPENMP # Check the results match: for yy in range(input.height()): for xx in range(input.width()): if halide_result(xx, yy) != c_result(xx, yy): raise Exception("halide_result(%d, %d) = %d instead of %d" % ( xx, yy, halide_result(xx, yy), c_result(xx, yy))) return -1 #endif # __SSE2__ else: print("(Skipped the SSE2 section of the code, " "since non-sense in python world.)") print("Success!") return 0
def findStereoCorrespondence(left, right, SADWindowSize, minDisparity, numDisparities, xmin, xmax, ymin, ymax, x_tile_size=32, y_tile_size=32, test=False, uniquenessRatio=0.15, disp12MaxDiff=1): """ Returns Func (left: Func, right: Func) """ x, y, c, d = Var("x"), Var("y"), Var("c"), Var("d") diff = Func("diff") diff[d, x, y] = h.cast(UInt(16), h.abs(left[x, y] - right[x-d, y])) win2 = SADWindowSize/2 diff_T = Func("diff_T") xi, xo, yi, yo = Var("xi"), Var("xo"), Var("yi"), Var("yo") diff_T[d, xi, yi, xo, yo] = diff[d, xi + xo * x_tile_size + xmin, yi + yo * y_tile_size + ymin] cSAD, vsum = Func("cSAD"), Func("vsum") rk = RDom(-win2, SADWindowSize, "rk") rxi, ryi = RDom(1, x_tile_size - 1, "rxi"), RDom(1, y_tile_size - 1, "ryi") if test: vsum[d, xi, yi, xo, yo] = h.sum(diff_T[d, xi, yi+rk, xo, yo]) cSAD[d, xi, yi, xo, yo] = h.sum(vsum[d, xi+rk, yi, xo, yo]) else: vsum[d, xi, yi, xo, yo] = h.select(yi != 0, h.cast(UInt(16), 0), h.sum(diff_T[d, xi, rk, xo, yo])) vsum[d, xi, ryi, xo, yo] = vsum[d, xi, ryi-1, xo, yo] + diff_T[d, xi, ryi+win2, xo, yo] - diff_T[d, xi, ryi-win2-1, xo, yo] cSAD[d, xi, yi, xo, yo] = h.select(xi != 0, h.cast(UInt(16), 0), h.sum(vsum[d, rk, yi, xo, yo])) cSAD[d, rxi, yi, xo, yo] = cSAD[d, rxi-1, yi, xo, yo] + vsum[d, rxi+win2, yi, xo, yo] - vsum[d, rxi-win2-1, yi, xo, yo] rd = RDom(minDisparity, numDisparities) disp_left = Func("disp_left") disp_left[xi, yi, xo, yo] = h.Tuple(h.cast(UInt(16), minDisparity), h.cast(UInt(16), (2<<16)-1)) disp_left[xi, yi, xo, yo] = h.tuple_select( cSAD[rd, xi, yi, xo, yo] < disp_left[xi, yi, xo, yo][1], h.Tuple(h.cast(UInt(16), rd), cSAD[rd, xi, yi, xo, yo]), h.Tuple(disp_left[xi, yi, xo, yo])) FILTERED = -16 disp = Func("disp") disp[x, y] = h.select( # x > xmax-xmin or y > ymax-ymin, x < xmax, h.cast(UInt(16), disp_left[x % x_tile_size, y % y_tile_size, x / x_tile_size, y / y_tile_size][0]), h.cast(UInt(16), FILTERED)) # Schedule vector_width = 8 disp.compute_root() \ .tile(x, y, xo, yo, xi, yi, x_tile_size, y_tile_size).reorder(xi, yi, xo, yo) \ .vectorize(xi, vector_width).parallel(xo).parallel(yo) # reorder storage disp_left.reorder_storage(xi, yi, xo, yo) diff_T .reorder_storage(xi, yi, xo, yo, d) vsum .reorder_storage(xi, yi, xo, yo, d) cSAD .reorder_storage(xi, yi, xo, yo, d) disp_left.compute_at(disp, xo).reorder(xi, yi, xo, yo) \ .vectorize(xi, vector_width) \ .update() \ .reorder(xi, yi, rd, xo, yo).vectorize(xi, vector_width) if test: cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) else: cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \ .update() \ .reorder(yi, rxi, xo, yo, d).vectorize(yi, vector_width) vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \ .update() \ .reorder(xi, ryi, xo, yo, d).vectorize(xi, vector_width) return disp