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
