diff --git a/src/exo/core/LoopIR_pprint.py b/src/exo/core/LoopIR_pprint.py index 4464976e3..ac97c8b95 100644 --- a/src/exo/core/LoopIR_pprint.py +++ b/src/exo/core/LoopIR_pprint.py @@ -12,6 +12,20 @@ from .internal_cursors import Node, Gap, Block, Cursor, InvalidCursorError, GapType from .prelude import * + +def _format_code(code): + # See the following file for customization options: + # https://github.com/google/yapf/blob/main/yapf/yapflib/style.py + code, _ = FormatCode( + code, + style_config={ + "based_on_style": "pep8", + "column_limit": 160, + }, + ) + return code.rstrip("\n") + + # --------------------------------------------------------------------------- # # --------------------------------------------------------------------------- # # Notes on Layout Schemes... @@ -110,10 +124,7 @@ def str(self): assert len(self._lines) == 1 return self._lines[0] - fmtstr, linted = FormatCode("\n".join(self._lines)) - if isinstance(self._node, LoopIR.proc): - assert linted, "generated unlinted code..." - return fmtstr + return _format_code("\n".join(self._lines)) def push(self, only=None): if only is None: @@ -321,10 +332,6 @@ def ptype(self, t): # LoopIR Pretty Printing -def _format_code(code): - return FormatCode(code)[0].rstrip("\n") - - @extclass(LoopIR.proc) def __str__(self): return _format_code("\n".join(_print_proc(self, PrintEnv(), ""))) diff --git a/tests/golden/test_config/test_ld.txt b/tests/golden/test_config/test_ld.txt index 0fe2c60f3..6d4099bed 100644 --- a/tests/golden/test_config/test_ld.txt +++ b/tests/golden/test_config/test_ld.txt @@ -3,8 +3,7 @@ def config_ld_i8(scale: f32 @ DRAM, src_stride: stride @ DRAM): # ConfigLoad.scale = scale ConfigLoad.src_stride = src_stride -def ld_i8(n: size, m: size, scale: f32 @ DRAM, src: i8[n, m] @ DRAM, - dst: i8[n, 16] @ GEMM_SCRATCH): +def ld_i8(n: size, m: size, scale: f32 @ DRAM, src: i8[n, m] @ DRAM, dst: i8[n, 16] @ GEMM_SCRATCH): assert n <= 16 assert m <= 16 assert stride(src, 1) == 1 diff --git a/tests/golden/test_cursors/test_basic_forwarding2.txt b/tests/golden/test_cursors/test_basic_forwarding2.txt index 1c3e62cfe..84b7030b1 100644 --- a/tests/golden/test_cursors/test_basic_forwarding2.txt +++ b/tests/golden/test_cursors/test_basic_forwarding2.txt @@ -1,5 +1,4 @@ -def filter1D(ow: size, kw: size, x: f32[ow + kw - 1] @ DRAM, y: f32[ow] @ DRAM, - w: f32[kw] @ DRAM): +def filter1D(ow: size, kw: size, x: f32[ow + kw - 1] @ DRAM, y: f32[ow] @ DRAM, w: f32[kw] @ DRAM): for outXo in seq(0, ow / 4): sum: f32[4] @ DRAM # <-- NODE for outXi in seq(0, 4): diff --git a/tests/golden/test_cursors/test_basic_forwarding3.txt b/tests/golden/test_cursors/test_basic_forwarding3.txt index 1c3e62cfe..84b7030b1 100644 --- a/tests/golden/test_cursors/test_basic_forwarding3.txt +++ b/tests/golden/test_cursors/test_basic_forwarding3.txt @@ -1,5 +1,4 @@ -def filter1D(ow: size, kw: size, x: f32[ow + kw - 1] @ DRAM, y: f32[ow] @ DRAM, - w: f32[kw] @ DRAM): +def filter1D(ow: size, kw: size, x: f32[ow + kw - 1] @ DRAM, y: f32[ow] @ DRAM, w: f32[kw] @ DRAM): for outXo in seq(0, ow / 4): sum: f32[4] @ DRAM # <-- NODE for outXi in seq(0, 4): diff --git a/tests/golden/test_halide_ops/test_compute_at_with_prologue.txt b/tests/golden/test_halide_ops/test_compute_at_with_prologue.txt index cc12975d6..370061777 100644 --- a/tests/golden/test_halide_ops/test_compute_at_with_prologue.txt +++ b/tests/golden/test_halide_ops/test_compute_at_with_prologue.txt @@ -9,5 +9,4 @@ def foo(): producer[1 + y, x] = 1.0 for x in seq(0, 10): producer[1 + y, 1 + x] = 1.0 - consumer[y, x] = producer[y, x] + producer[y, 1 + x] + producer[ - 1 + y, x] + producer[1 + y, 1 + x] \ No newline at end of file + consumer[y, x] = producer[y, x] + producer[y, 1 + x] + producer[1 + y, x] + producer[1 + y, 1 + x] \ No newline at end of file diff --git a/tests/golden/test_halide_ops/test_schedule_blur1d.txt b/tests/golden/test_halide_ops/test_schedule_blur1d.txt index e268dce51..6e50b7368 100644 --- a/tests/golden/test_halide_ops/test_schedule_blur1d.txt +++ b/tests/golden/test_halide_ops/test_schedule_blur1d.txt @@ -1,25 +1,18 @@ -def blur1d_compute_at_store_root(n: size, consumer: i8[n] @ DRAM, - inp: i8[n + 6] @ DRAM): +def blur1d_compute_at_store_root(n: size, consumer: i8[n] @ DRAM, inp: i8[n + 6] @ DRAM): producer: i8[1 + n] @ DRAM for i in seq(0, n): for ii in seq(0, 2): - producer[i + ii] = (inp[i + ii] + inp[1 + i + ii] + - inp[2 + i + ii] + inp[3 + i + ii] + - inp[4 + i + ii] + inp[5 + i + ii]) / 6.0 + producer[i + ii] = (inp[i + ii] + inp[1 + i + ii] + inp[2 + i + ii] + inp[3 + i + ii] + inp[4 + i + ii] + inp[5 + i + ii]) / 6.0 consumer[i] = (producer[i] + producer[1 + i]) / 2.0 def blur1d_compute_at(n: size, consumer: i8[n] @ DRAM, inp: i8[n + 6] @ DRAM): for i in seq(0, n): producer: i8[2] @ DRAM for ii in seq(0, 2): - producer[ii] = (inp[i + ii] + inp[1 + i + ii] + inp[2 + i + ii] + - inp[3 + i + ii] + inp[4 + i + ii] + - inp[5 + i + ii]) / 6.0 + producer[ii] = (inp[i + ii] + inp[1 + i + ii] + inp[2 + i + ii] + inp[3 + i + ii] + inp[4 + i + ii] + inp[5 + i + ii]) / 6.0 consumer[i] = (producer[0] + producer[1]) / 2.0 def blur1d_inline(n: size, consumer: i8[n] @ DRAM, inp: i8[n + 6] @ DRAM): for i in seq(0, n): - consumer[i] = ((inp[i] + inp[1 + i] + inp[2 + i] + inp[3 + i] + - inp[4 + i] + inp[5 + i]) / 6.0 + - (inp[1 + i] + inp[2 + i] + inp[3 + i] + inp[4 + i] + - inp[5 + i] + inp[6 + i]) / 6.0) / 2.0 \ No newline at end of file + consumer[i] = ((inp[i] + inp[1 + i] + inp[2 + i] + inp[3 + i] + inp[4 + i] + inp[5 + i]) / 6.0 + + (inp[1 + i] + inp[2 + i] + inp[3 + i] + inp[4 + i] + inp[5 + i] + inp[6 + i]) / 6.0) / 2.0 \ No newline at end of file diff --git a/tests/golden/test_halide_ops/test_schedule_blur2d.txt b/tests/golden/test_halide_ops/test_schedule_blur2d.txt index d38bb0fb5..ef44f1a49 100644 --- a/tests/golden/test_halide_ops/test_schedule_blur2d.txt +++ b/tests/golden/test_halide_ops/test_schedule_blur2d.txt @@ -1,5 +1,4 @@ -def blur2d_compute_at_i_store_root(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_compute_at_i_store_root(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[1 + n, 1 + n] @ DRAM for i in seq(0, n): @@ -7,12 +6,9 @@ def blur2d_compute_at_i_store_root(n: size, consumer: i8[n, n] @ DRAM, for j in seq(0, 1 + n): producer[i + ii, j] = sin[i + ii, j] for j in seq(0, n): - consumer[i, - j] = (producer[i, j] + producer[i, 1 + j] + - producer[1 + i, j] + producer[1 + i, 1 + j]) / 4.0 + consumer[i, j] = (producer[i, j] + producer[i, 1 + j] + producer[1 + i, j] + producer[1 + i, 1 + j]) / 4.0 -def blur2d_compute_at_j_store_root(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_compute_at_j_store_root(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[1 + n, 1 + n] @ DRAM for i in seq(0, n): @@ -20,12 +16,9 @@ def blur2d_compute_at_j_store_root(n: size, consumer: i8[n, n] @ DRAM, for ji in seq(0, 2): for ii in seq(0, 2): producer[i + ii, j + ji] = sin[i + ii, j + ji] - consumer[i, - j] = (producer[i, j] + producer[i, 1 + j] + - producer[1 + i, j] + producer[1 + i, 1 + j]) / 4.0 + consumer[i, j] = (producer[i, j] + producer[i, 1 + j] + producer[1 + i, j] + producer[1 + i, 1 + j]) / 4.0 -def blur2d_compute_at_i(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_compute_at_i(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 for i in seq(0, n): producer: i8[2, 1 + n] @ DRAM @@ -33,11 +26,9 @@ def blur2d_compute_at_i(n: size, consumer: i8[n, n] @ DRAM, for j in seq(0, 1 + n): producer[ii, j] = sin[i + ii, j] for j in seq(0, n): - consumer[i, j] = (producer[0, j] + producer[0, 1 + j] + - producer[1, j] + producer[1, 1 + j]) / 4.0 + consumer[i, j] = (producer[0, j] + producer[0, 1 + j] + producer[1, j] + producer[1, 1 + j]) / 4.0 -def blur2d_compute_at_j_store_at_i(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_compute_at_j_store_at_i(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 for i in seq(0, n): producer: i8[2, 1 + n] @ DRAM @@ -45,13 +36,10 @@ def blur2d_compute_at_j_store_at_i(n: size, consumer: i8[n, n] @ DRAM, for ji in seq(0, 2): for ii in seq(0, 2): producer[ii, j + ji] = sin[i + ii, j + ji] - consumer[i, j] = (producer[0, j] + producer[0, 1 + j] + - producer[1, j] + producer[1, 1 + j]) / 4.0 + consumer[i, j] = (producer[0, j] + producer[0, 1 + j] + producer[1, j] + producer[1, 1 + j]) / 4.0 -def blur2d_inline(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_inline(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 for i in seq(0, n): for j in seq(0, n): - consumer[i, j] = (sin[i + 0, j + 0] + sin[i + 0, j + 1] + - sin[i + 1, j + 0] + sin[i + 1, j + 1]) / 4.0 \ No newline at end of file + consumer[i, j] = (sin[i + 0, j + 0] + sin[i + 0, j + 1] + sin[i + 1, j + 0] + sin[i + 1, j + 1]) / 4.0 \ No newline at end of file diff --git a/tests/golden/test_halide_ops/test_schedule_tiled_blur2d.txt b/tests/golden/test_halide_ops/test_schedule_tiled_blur2d.txt index d0b86272c..187df0397 100644 --- a/tests/golden/test_halide_ops/test_schedule_tiled_blur2d.txt +++ b/tests/golden/test_halide_ops/test_schedule_tiled_blur2d.txt @@ -1,5 +1,4 @@ -def blur2d_tiled(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_tiled(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[n + 1, n + 1] @ DRAM for i in seq(0, n + 1): @@ -9,15 +8,10 @@ def blur2d_tiled(n: size, consumer: i8[n, n] @ DRAM, for j in seq(0, n / 4): for ii in seq(0, 4): for ji in seq(0, 4): - consumer[ - 4 * i + ii, 4 * j + - ji] = (producer[4 * i + ii, 4 * j + ji] + - producer[4 * i + ii, 4 * j + ji + 1] + - producer[4 * i + ii + 1, 4 * j + ji] + - producer[4 * i + ii + 1, 4 * j + ji + 1]) / 4.0 + consumer[4 * i + ii, 4 * j + ji] = (producer[4 * i + ii, 4 * j + ji] + producer[4 * i + ii, 4 * j + ji + 1] + + producer[4 * i + ii + 1, 4 * j + ji] + producer[4 * i + ii + 1, 4 * j + ji + 1]) / 4.0 -def blur2d_tiled_compute_at_i(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_tiled_compute_at_i(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[1 + n, 1 + n] @ DRAM for i in seq(0, n / 4): @@ -27,32 +21,23 @@ def blur2d_tiled_compute_at_i(n: size, consumer: i8[n, n] @ DRAM, for j in seq(0, n / 4): for ii in seq(0, 4): for ji in seq(0, 4): - consumer[ii + 4 * i, ji + 4 * j] = ( - producer[ii + 4 * i, ji + 4 * j] + - producer[ii + 4 * i, 1 + ji + 4 * j] + - producer[1 + ii + 4 * i, ji + 4 * j] + - producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 + consumer[ii + 4 * i, ji + 4 * j] = (producer[ii + 4 * i, ji + 4 * j] + producer[ii + 4 * i, 1 + ji + 4 * j] + + producer[1 + ii + 4 * i, ji + 4 * j] + producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 -def blur2d_tiled_compute_at_j(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_tiled_compute_at_j(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[1 + n, 1 + n] @ DRAM for i in seq(0, n / 4): for j in seq(0, n / 4): for ji in seq(0, 5): for ii in seq(0, 5): - producer[ii + 4 * i, ji + 4 * j] = sin[ii + 4 * i, - ji + 4 * j] + producer[ii + 4 * i, ji + 4 * j] = sin[ii + 4 * i, ji + 4 * j] for ii in seq(0, 4): for ji in seq(0, 4): - consumer[ii + 4 * i, ji + 4 * j] = ( - producer[ii + 4 * i, ji + 4 * j] + - producer[ii + 4 * i, 1 + ji + 4 * j] + - producer[1 + ii + 4 * i, ji + 4 * j] + - producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 + consumer[ii + 4 * i, ji + 4 * j] = (producer[ii + 4 * i, ji + 4 * j] + producer[ii + 4 * i, 1 + ji + 4 * j] + + producer[1 + ii + 4 * i, ji + 4 * j] + producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 -def blur2d_tiled_compute_at_ii(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_tiled_compute_at_ii(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[1 + n, 1 + n] @ DRAM for i in seq(0, n / 4): @@ -60,18 +45,12 @@ def blur2d_tiled_compute_at_ii(n: size, consumer: i8[n, n] @ DRAM, for ii in seq(0, 4): for iii in seq(0, 2): for ji in seq(0, 5): - producer[ii + iii + 4 * i, - ji + 4 * j] = sin[ii + iii + 4 * i, - ji + 4 * j] + producer[ii + iii + 4 * i, ji + 4 * j] = sin[ii + iii + 4 * i, ji + 4 * j] for ji in seq(0, 4): - consumer[ii + 4 * i, ji + 4 * j] = ( - producer[ii + 4 * i, ji + 4 * j] + - producer[ii + 4 * i, 1 + ji + 4 * j] + - producer[1 + ii + 4 * i, ji + 4 * j] + - producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 + consumer[ii + 4 * i, ji + 4 * j] = (producer[ii + 4 * i, ji + 4 * j] + producer[ii + 4 * i, 1 + ji + 4 * j] + + producer[1 + ii + 4 * i, ji + 4 * j] + producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 -def blur2d_tiled_compute_at_ji(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_tiled_compute_at_ji(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 producer: i8[1 + n, 1 + n] @ DRAM for i in seq(0, n / 4): @@ -80,17 +59,11 @@ def blur2d_tiled_compute_at_ji(n: size, consumer: i8[n, n] @ DRAM, for ji in seq(0, 4): for jii in seq(0, 2): for iii in seq(0, 2): - producer[ii + iii + 4 * i, - ji + jii + 4 * j] = sin[ii + iii + 4 * i, - ji + jii + 4 * j] - consumer[ii + 4 * i, ji + 4 * j] = ( - producer[ii + 4 * i, ji + 4 * j] + - producer[ii + 4 * i, 1 + ji + 4 * j] + - producer[1 + ii + 4 * i, ji + 4 * j] + - producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 + producer[ii + iii + 4 * i, ji + jii + 4 * j] = sin[ii + iii + 4 * i, ji + jii + 4 * j] + consumer[ii + 4 * i, ji + 4 * j] = (producer[ii + 4 * i, ji + 4 * j] + producer[ii + 4 * i, 1 + ji + 4 * j] + + producer[1 + ii + 4 * i, ji + 4 * j] + producer[1 + ii + 4 * i, 1 + ji + 4 * j]) / 4.0 -def blur2d_tiled_compute_at_and_store_at_ji(n: size, consumer: i8[n, n] @ DRAM, - sin: i8[n + 1, n + 1] @ DRAM): +def blur2d_tiled_compute_at_and_store_at_ji(n: size, consumer: i8[n, n] @ DRAM, sin: i8[n + 1, n + 1] @ DRAM): assert n % 4 == 0 for i in seq(0, n / 4): for j in seq(0, n / 4): @@ -99,8 +72,5 @@ def blur2d_tiled_compute_at_and_store_at_ji(n: size, consumer: i8[n, n] @ DRAM, producer: i8[2, 2] @ DRAM for jii in seq(0, 2): for iii in seq(0, 2): - producer[iii, jii] = sin[ii + iii + 4 * i, - ji + jii + 4 * j] - consumer[ii + 4 * i, ji + - 4 * j] = (producer[0, 0] + producer[0, 1] + - producer[1, 0] + producer[1, 1]) / 4.0 \ No newline at end of file + producer[iii, jii] = sin[ii + iii + 4 * i, ji + jii + 4 * j] + consumer[ii + 4 * i, ji + 4 * j] = (producer[0, 0] + producer[0, 1] + producer[1, 0] + producer[1, 1]) / 4.0 \ No newline at end of file diff --git a/tests/golden/test_im2col/test_im2col.txt b/tests/golden/test_im2col/test_im2col.txt index 3febf84c0..f36cbb338 100644 --- a/tests/golden/test_im2col/test_im2col.txt +++ b/tests/golden/test_im2col/test_im2col.txt @@ -1,20 +1,17 @@ -def im2col(C: size, W: size, R: size, x: R[C, W] @ DRAM, - y: R[C + 1, R + 1, W + 1] @ DRAM): +def im2col(C: size, W: size, R: size, x: R[C, W] @ DRAM, y: R[C + 1, R + 1, W + 1] @ DRAM): for c in seq(0, C): for r in seq(0, R): for i in seq(0, W): if 0 <= i - r: y[c, r, i] = x[c, i - r] -def matmul(K: size, C: size, W: size, R: size, w: R[K, C, R] @ DRAM, - res: R[K, W] @ DRAM, y: R[C + 1, R + 1, W + 1] @ DRAM): +def matmul(K: size, C: size, W: size, R: size, w: R[K, C, R] @ DRAM, res: R[K, W] @ DRAM, y: R[C + 1, R + 1, W + 1] @ DRAM): for k in seq(0, K): for c in seq(0, C): for r in seq(0, R): for i in seq(0, W): if 0 <= i - r: res[k, i] += w[k, c, r] * y[c, r, i] -def im2col_conv(K: size, C: size, W: size, R: size, w: R[K, C, R] @ DRAM, - x: R[C, W] @ DRAM, res: R[K, W] @ DRAM): +def im2col_conv(K: size, C: size, W: size, R: size, w: R[K, C, R] @ DRAM, x: R[C, W] @ DRAM, res: R[K, W] @ DRAM): for k_init in seq(0, K): for i_init in seq(0, W): res[k_init, i_init] = 0.0 diff --git a/tests/golden/test_neon/test_gen_neon_simple_math_scheduling.txt b/tests/golden/test_neon/test_gen_neon_simple_math_scheduling.txt index b1d9e2f3a..843e2f726 100644 --- a/tests/golden/test_neon/test_gen_neon_simple_math_scheduling.txt +++ b/tests/golden/test_neon/test_gen_neon_simple_math_scheduling.txt @@ -10,5 +10,4 @@ def simple_math_neon_sched(n: size, x: R[n] @ DRAM, y: R[n] @ DRAM): neon_vst_4xf32(x[4 * io:4 + 4 * io], xVec[0:4]) if n % 4 > 0: for ii in seq(0, n % 4): - x[ii + n / 4 * - 4] = x[ii + n / 4 * 4] * y[ii + n / 4 * 4] * y[ii + n / 4 * 4] \ No newline at end of file + x[ii + n / 4 * 4] = x[ii + n / 4 * 4] * y[ii + n / 4 * 4] * y[ii + n / 4 * 4] \ No newline at end of file diff --git a/tests/golden/test_new_eff/test_reorder_loops_4pt_stencil_succeed.txt b/tests/golden/test_new_eff/test_reorder_loops_4pt_stencil_succeed.txt index 27598dfc9..a7636e42e 100644 --- a/tests/golden/test_new_eff/test_reorder_loops_4pt_stencil_succeed.txt +++ b/tests/golden/test_new_eff/test_reorder_loops_4pt_stencil_succeed.txt @@ -2,5 +2,4 @@ def foo(N: size, x: R[N, N] @ DRAM): for j in seq(0, N): for i in seq(0, N): if 0 < i and i < N - 1 and (0 < j and j < N - 1): - x[i, j] += -1.0 / 4.0 * (x[i - 1, j] + x[i + 1, j] + - x[i, j - 1] + x[i, j + 1]) \ No newline at end of file + x[i, j] += -1.0 / 4.0 * (x[i - 1, j] + x[i + 1, j] + x[i, j - 1] + x[i, j + 1]) \ No newline at end of file diff --git a/tests/golden/test_new_eff/test_reorder_loops_requiring_seq.txt b/tests/golden/test_new_eff/test_reorder_loops_requiring_seq.txt index eaa00f8b2..7e30b21ef 100644 --- a/tests/golden/test_new_eff/test_reorder_loops_requiring_seq.txt +++ b/tests/golden/test_new_eff/test_reorder_loops_requiring_seq.txt @@ -2,5 +2,4 @@ def foo(N: size, x: R[N, N] @ DRAM): for j in seq(0, N): for i in seq(0, N): if i > 0 and j > 0: - x[i, j] += -1.0 / 3.0 * (x[i - 1, j] + x[i - 1, j - 1] + - x[i, j - 1]) \ No newline at end of file + x[i, j] += -1.0 / 3.0 * (x[i - 1, j] + x[i - 1, j - 1] + x[i, j - 1]) \ No newline at end of file diff --git a/tests/golden/test_schedules/test_cut_loop_syrk.txt b/tests/golden/test_schedules/test_cut_loop_syrk.txt index 46f7192ea..8ab6c725a 100644 --- a/tests/golden/test_schedules/test_cut_loop_syrk.txt +++ b/tests/golden/test_schedules/test_cut_loop_syrk.txt @@ -1,5 +1,4 @@ -def SYRK(M: size, K: size, A: [f32][M, K] @ DRAM, A_t: [f32][M, K] @ DRAM, - C: [f32][M, M] @ DRAM): +def SYRK(M: size, K: size, A: [f32][M, K] @ DRAM, A_t: [f32][M, K] @ DRAM, C: [f32][M, M] @ DRAM): assert M >= 1 assert K >= 1 assert stride(A, 1) == 1 diff --git a/tests/golden/test_schedules/test_extract_subproc7.txt b/tests/golden/test_schedules/test_extract_subproc7.txt index 6b951277b..c3756fa3c 100644 --- a/tests/golden/test_schedules/test_extract_subproc7.txt +++ b/tests/golden/test_schedules/test_extract_subproc7.txt @@ -1,12 +1,10 @@ -def gemv(m: size, n: size, alpha: R @ DRAM, beta: R @ DRAM, - A: [R][m, n] @ DRAM, x: [R][n] @ DRAM, y: [R][m] @ DRAM): +def gemv(m: size, n: size, alpha: R @ DRAM, beta: R @ DRAM, A: [R][m, n] @ DRAM, x: [R][n] @ DRAM, y: [R][m] @ DRAM): assert stride(A, 1) == 1 for i in seq(0, m): y[i] = y[i] * beta for j in seq(0, n): fooooo(m, n, alpha, A, x, y, j) -def fooooo(m: size, n: size, alpha: R @ DRAM, A: [R][m, n] @ DRAM, - x: [R][n] @ DRAM, y: [R][m] @ DRAM, j: index): +def fooooo(m: size, n: size, alpha: R @ DRAM, A: [R][m, n] @ DRAM, x: [R][n] @ DRAM, y: [R][m] @ DRAM, j: index): assert stride(A, 1) == 1 assert 0 <= j assert j < n diff --git a/tests/golden/test_schedules/test_fold_buffer_blur.txt b/tests/golden/test_schedules/test_fold_buffer_blur.txt index 4a19b323f..60a2fe48e 100644 --- a/tests/golden/test_schedules/test_fold_buffer_blur.txt +++ b/tests/golden/test_schedules/test_fold_buffer_blur.txt @@ -5,15 +5,9 @@ def blur(H: size, W: size, inp: i8[H + 2, W] @ DRAM, out: i8[H, W] @ DRAM): blur_x: i8[3, W] @ DRAM for ii in seq(0, 2): for j in seq(0, -2 + W): - blur_x[ii, j] = inp[ii + 32 * io, - j] + inp[ii + 32 * io, - 1 + j] + inp[ii + 32 * io, 2 + j] + blur_x[ii, j] = inp[ii + 32 * io, j] + inp[ii + 32 * io, 1 + j] + inp[ii + 32 * io, 2 + j] for ii in seq(0, 32): for j in seq(0, -2 + W): - blur_x[(2 + ii) % 3, - j] = inp[ii + 32 * io, - j] + inp[ii + 32 * io, - 1 + j] + inp[ii + 32 * io, 2 + j] + blur_x[(2 + ii) % 3, j] = inp[ii + 32 * io, j] + inp[ii + 32 * io, 1 + j] + inp[ii + 32 * io, 2 + j] for j in seq(0, -2 + W): - out[ii + 32 * io, j] = blur_x[ii % 3, j] + blur_x[ - (1 + ii) % 3, j] + blur_x[(2 + ii) % 3, j] \ No newline at end of file + out[ii + 32 * io, j] = blur_x[ii % 3, j] + blur_x[(1 + ii) % 3, j] + blur_x[(2 + ii) % 3, j] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_fold_buffer_unsharp.txt b/tests/golden/test_schedules/test_fold_buffer_unsharp.txt index 81d1ef3d1..3c6a26435 100644 --- a/tests/golden/test_schedules/test_fold_buffer_unsharp.txt +++ b/tests/golden/test_schedules/test_fold_buffer_unsharp.txt @@ -1,5 +1,4 @@ -def exo_unsharp_base(W: size, H: size, output: f32[3, H, W] @ DRAM, - input: f32[3, H + 6, W + 6] @ DRAM): +def exo_unsharp_base(W: size, H: size, output: f32[3, H, W] @ DRAM, input: f32[3, H + 6, W + 6] @ DRAM): assert H % 32 == 0 for y in par(0, H / 32): gray: f32[8, 6 + W] @ DRAM @@ -7,27 +6,16 @@ def exo_unsharp_base(W: size, H: size, output: f32[3, H, W] @ DRAM, blur_y: f32[1, 6 + W] @ DRAM for yi in seq(0, 6): for x in seq(0, 6 + W): - gray[yi, x] = input[0, yi + 32 * y, - x] + input[1, yi + 32 * y, - x] + input[2, yi + 32 * y, x] + gray[yi, x] = input[0, yi + 32 * y, x] + input[1, yi + 32 * y, x] + input[2, yi + 32 * y, x] for y_i in seq(0, 32): for x in seq(0, 6 + W): - gray[(6 + y_i) % 8, - x] = input[0, 6 + y_i + 32 * y, - x] + input[1, 6 + y_i + 32 * y, - x] + input[2, 6 + y_i + 32 * y, x] + gray[(6 + y_i) % 8, x] = input[0, 6 + y_i + 32 * y, x] + input[1, 6 + y_i + 32 * y, x] + input[2, 6 + y_i + 32 * y, x] for x in seq(0, 6 + W): - blur_y[0, x] = gray[(3 + y_i) % 8, x] + gray[ - (2 + y_i) % 8, x] + gray[(4 + y_i) % 8, x] + gray[ - (1 + y_i) % 8, x] + gray[(5 + y_i) % 8, x] + gray[ - y_i % 8, x] + gray[(6 + y_i) % 8, x] + blur_y[0, x] = gray[(3 + y_i) % 8, x] + gray[(2 + y_i) % 8, x] + gray[(4 + y_i) % 8, x] + gray[(1 + y_i) % 8, x] + gray[ + (5 + y_i) % 8, x] + gray[y_i % 8, x] + gray[(6 + y_i) % 8, x] for x in seq(0, W): - ratio[0, x] = ( - gray[(3 + y_i) % 8, 3 + x] - - (blur_y[0, 3 + x] + blur_y[0, 2 + x] + blur_y[0, 4 + x] + - blur_y[0, 1 + x] + blur_y[0, 5 + x] + blur_y[0, x] + - blur_y[0, 6 + x])) / gray[(3 + y_i) % 8, 3 + x] + ratio[0, x] = (gray[(3 + y_i) % 8, 3 + x] - (blur_y[0, 3 + x] + blur_y[0, 2 + x] + blur_y[0, 4 + x] + blur_y[0, 1 + x] + blur_y[0, 5 + x] + + blur_y[0, x] + blur_y[0, 6 + x])) / gray[(3 + y_i) % 8, 3 + x] for c in seq(0, 3): for x in seq(0, W): - output[c, y_i + 32 * y, - x] = ratio[0, x] * input[c, 3 + y_i + 32 * y, 3 + x] \ No newline at end of file + output[c, y_i + 32 * y, x] = ratio[0, x] * input[c, 3 + y_i + 32 * y, 3 + x] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_simplify2.txt b/tests/golden/test_schedules/test_simplify2.txt index a0f2ad74e..c4dcbf407 100644 --- a/tests/golden/test_schedules/test_simplify2.txt +++ b/tests/golden/test_schedules/test_simplify2.txt @@ -1,5 +1,4 @@ -def foo(A: i8[32, 64] @ DRAM, B: i8[16, 128] @ DRAM, C: i32[32, 32] @ DRAM, - ko: size, ji_unroll: size, ii_unroll: size): +def foo(A: i8[32, 64] @ DRAM, B: i8[16, 128] @ DRAM, C: i32[32, 32] @ DRAM, ko: size, ji_unroll: size, ii_unroll: size): for io in seq(0, 1): for jo in seq(0, 1): Btile1: i8[16, 64] @ DRAM diff --git a/tests/golden/test_schedules/test_specialize_sizes.txt b/tests/golden/test_schedules/test_specialize_sizes.txt index 3e5f30a35..0af67100e 100644 --- a/tests/golden/test_schedules/test_specialize_sizes.txt +++ b/tests/golden/test_schedules/test_specialize_sizes.txt @@ -1,5 +1,4 @@ -def gemm(M: size, N: size, K: size, C: f32[M, N] @ DRAM, A: f32[M, K] @ DRAM, - B: f32[K, N] @ DRAM, alpha: f32 @ DRAM): +def gemm(M: size, N: size, K: size, C: f32[M, N] @ DRAM, A: f32[M, K] @ DRAM, B: f32[K, N] @ DRAM, alpha: f32 @ DRAM): if N <= 64: for i in seq(0, M): for j in seq(0, N): diff --git a/tests/golden/test_schedules/test_stage_mem.txt b/tests/golden/test_schedules/test_stage_mem.txt index ecca52186..44e6f2505 100644 --- a/tests/golden/test_schedules/test_stage_mem.txt +++ b/tests/golden/test_schedules/test_stage_mem.txt @@ -10,9 +10,7 @@ def sqmat(n: size, A: R[n, n] @ DRAM, B: R[n, n] @ DRAM): for ii in seq(0, 4): for jj in seq(0, 4): for kk in seq(0, 4): - Atile[ii, - jj] += B[ii + 4 * i, kk + - 4 * k] * B[kk + 4 * k, jj + 4 * j] + Atile[ii, jj] += B[ii + 4 * i, kk + 4 * k] * B[kk + 4 * k, jj + 4 * j] for i0 in seq(0, 4): for i1 in seq(0, 4): A[i0 + 4 * i, i1 + 4 * j] = Atile[i0, i1] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_stage_mem_accum.txt b/tests/golden/test_schedules/test_stage_mem_accum.txt index 9f2a18df2..980f28fc1 100644 --- a/tests/golden/test_schedules/test_stage_mem_accum.txt +++ b/tests/golden/test_schedules/test_stage_mem_accum.txt @@ -10,9 +10,7 @@ def sqmat(n: size, A: R[n, n] @ DRAM, B: R[n, n] @ DRAM): for ii in seq(0, 4): for jj in seq(0, 4): for kk in seq(0, 4): - Atile[ii, - jj] += B[ii + 4 * i, kk + - 4 * k] * B[kk + 4 * k, jj + 4 * j] + Atile[ii, jj] += B[ii + 4 * i, kk + 4 * k] * B[kk + 4 * k, jj + 4 * j] for i0 in seq(0, 4): for i1 in seq(0, 4): A[i0 + 4 * i, i1 + 4 * j] += Atile[i0, i1] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_stage_mem_recursive.txt b/tests/golden/test_schedules/test_stage_mem_recursive.txt index 0c631a409..440a568d0 100644 --- a/tests/golden/test_schedules/test_stage_mem_recursive.txt +++ b/tests/golden/test_schedules/test_stage_mem_recursive.txt @@ -4,14 +4,9 @@ def recursive(n: size, y: R[n] @ DRAM, x: R[n] @ DRAM): for io in seq(0, (-2 + n) / 4): y_tmp: R[4] @ DRAM y_tmp[0] = y[1 + 4 * io] + y[4 * io] + x[4 * io] - y_tmp[1] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[ - 1 + 4 * io] - y_tmp[2] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[ - 1 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io]) + x[2 + - 4 * io] - y_tmp[3] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[ - 1 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io]) + x[ - 2 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io] + - y[1 + 4 * io] + x[1 + 4 * io]) + x[3 + 4 * io] + y_tmp[1] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[1 + 4 * io] + y_tmp[2] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[1 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io]) + x[2 + 4 * io] + y_tmp[3] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[1 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io]) + x[2 + 4 * io] + ( + y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[1 + 4 * io]) + x[3 + 4 * io] for i0 in seq(0, 4): y[2 + i0 + 4 * io] = y_tmp[i0] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_stage_mem_twice.txt b/tests/golden/test_schedules/test_stage_mem_twice.txt index b826202a7..3a4c1272e 100644 --- a/tests/golden/test_schedules/test_stage_mem_twice.txt +++ b/tests/golden/test_schedules/test_stage_mem_twice.txt @@ -14,5 +14,4 @@ def sqmat(n: size, A: R[n, n] @ DRAM, B: R[n, n] @ DRAM): for ii in seq(0, 4): for jj in seq(0, 4): for kk in seq(0, 4): - A[ii + 4 * i, - jj + 4 * j] += B1[ii, kk] * B2[kk, jj] \ No newline at end of file + A[ii + 4 * i, jj + 4 * j] += B1[ii, kk] * B2[kk, jj] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_unify3.txt b/tests/golden/test_schedules/test_unify3.txt index 4ce30848f..b81452f15 100644 --- a/tests/golden/test_schedules/test_unify3.txt +++ b/tests/golden/test_schedules/test_unify3.txt @@ -1,5 +1,4 @@ def foo(n: size, z: R[n] @ DRAM, x: R[n] @ DRAM, y: R[n] @ DRAM): assert n % 4 == 0 for i in seq(0, n / 4): - simd_add4(z[4 * i + 0:4 * i + 4], x[4 * i + 0:4 * i + 4], - y[4 * i + 0:4 * i + 4]) \ No newline at end of file + simd_add4(z[4 * i + 0:4 * i + 4], x[4 * i + 0:4 * i + 4], y[4 * i + 0:4 * i + 4]) \ No newline at end of file diff --git a/tests/golden/test_uast/test_alloc_nest.txt b/tests/golden/test_uast/test_alloc_nest.txt index f82ca4eab..3293dcd63 100644 --- a/tests/golden/test_uast/test_alloc_nest.txt +++ b/tests/golden/test_uast/test_alloc_nest.txt @@ -1,5 +1,4 @@ -def alloc_nest(n: size, m: size, x: R[n, m], y: R[n, m] @ DRAM, - res: R[n, m] @ DRAM): +def alloc_nest(n: size, m: size, x: R[n, m], y: R[n, m] @ DRAM, res: R[n, m] @ DRAM): # @instr TEST for i in seq(0, n): rloc: R[m] @ DRAM @@ -12,4 +11,4 @@ def alloc_nest(n: size, m: size, x: R[n, m], y: R[n, m] @ DRAM, for j in seq(0, m): rloc[j] = xloc[j] + yloc[j] for j in seq(0, m): - res[i, j] = rloc[j] + res[i, j] = rloc[j] \ No newline at end of file diff --git a/tests/golden/test_uast/test_conv1d.txt b/tests/golden/test_uast/test_conv1d.txt index 3544924d2..83ef97d2b 100644 --- a/tests/golden/test_uast/test_conv1d.txt +++ b/tests/golden/test_uast/test_conv1d.txt @@ -5,4 +5,4 @@ def conv1d(n: size, m: size, r: size, x: R[n], w: R[m], res: R[r]): for i in seq(0, r): for j in seq(0, n): if i <= j and j < i + m: - res[i] += x[j] * w[i - j + m - 1] + res[i] += x[j] * w[i - j + m - 1] \ No newline at end of file diff --git a/tests/golden/test_uast/test_unary_neg.txt b/tests/golden/test_uast/test_unary_neg.txt index c6231e18a..e6f7f0509 100644 --- a/tests/golden/test_uast/test_unary_neg.txt +++ b/tests/golden/test_uast/test_unary_neg.txt @@ -1,4 +1,4 @@ def negate_array(n: size, x: R[n], res: R[n] @ DRAM): # @instr TEST for i in seq(0, n): - res[i] = -x[i] + -x[i] - -(x[i] + 0.0) + res[i] = -x[i] + -x[i] - -(x[i] + 0.0) \ No newline at end of file diff --git a/tests/golden/test_x86/test_gen_avx2_simple_math_scheduling.txt b/tests/golden/test_x86/test_gen_avx2_simple_math_scheduling.txt index eb8e75158..8c91eef8f 100644 --- a/tests/golden/test_x86/test_gen_avx2_simple_math_scheduling.txt +++ b/tests/golden/test_x86/test_gen_avx2_simple_math_scheduling.txt @@ -10,5 +10,4 @@ def simple_math_avx2_sched(n: size, x: R[n] @ DRAM, y: R[n] @ DRAM): mm256_storeu_ps(x[8 * io:8 + 8 * io], xVec[0:8]) if n % 8 > 0: for ii in seq(0, n % 8): - x[ii + n / 8 * - 8] = x[ii + n / 8 * 8] * y[ii + n / 8 * 8] * y[ii + n / 8 * 8] \ No newline at end of file + x[ii + n / 8 * 8] = x[ii + n / 8 * 8] * y[ii + n / 8 * 8] * y[ii + n / 8 * 8] \ No newline at end of file