def test_bound(): A = te.placeholder((7, 7), name="A") B = te.compute((7, 7), lambda i, j: A[i, j], name="B") C = te.compute((3, 3), lambda i, j: B[i + 1, j * 2], name="C") s = te.create_schedule(C.op) KgeN.lower(s, [A, C]) assert s[B].op.axis[0].range.end.val == 3 assert s[B].op.axis[1].range.end.val == 5
def test_fuse(): A = te.placeholder((8, 8), name="A") B = te.compute((8, 8), lambda i, j: A[i, j], name="B") s = te.create_schedule(B.op) x, y = s[B].op.axis fused = s[B].fuse(x, y) KgeN.lower(s, [A, B]) assert fused.range.end.val == 64
def test_split1(): A = te.placeholder((9, 9), name="A") B = te.compute((9, 9), lambda i, j: A[i, j], name="B") s = te.create_schedule(B.op) x, _ = s[B].op.axis x_o, x_i = s[B].split(x, 2) KgeN.lower(s, [A, B]) assert x_o.range.end.val == 5 assert x_i.range.end.val == 2
def test_if_then_else(): A = te.placeholder((8, 8), name="A") B = te.compute((8, 8), lambda i, j: A[i, j], name="B") C = te.compute( (3, 3), lambda i, j: te.if_then_else(te.all(i < 2, j < 1), B[i, j] + 1, 0), name="C") s = te.create_schedule(C.op) KgeN.lower(s, [A, C]) x, y = s[B].op.axis assert x.range.end.val == 2 assert y.range.end.val == 1
import KgeN from KgeN import te m = 256 A = te.placeholder((m, m), name="A") B = te.compute((m, m), lambda i, j: 2 + A[i, j], name="B") C = te.compute((m, m), lambda i, j: B[i + j, j] * 2, name="C") # schedule s = te.create_schedule(C.op) outer, inner = s[C].split(s[C].op.axis[0], 32) B_outer, B_inner = s[B].split(s[B].op.axis[0], 32) # s[C].reorder(inner, outer) # fused = s[C].fuse(outer, inner) s[B].compute_at(s[C], s[C].op.axis[1]) # lower func = KgeN.lower(s, [A, C]) print(KgeN.build(func))
import KgeN from KgeN import te M = 128 A = te.compute((M, M), lambda i, j: 1, name="A") s = te.create_schedule(A.op) func = KgeN.lower(s, [A]) print(KgeN.build(func))
import KgeN from KgeN import te M = 128 A = te.placeholder((M, M), name="A") B = te.compute((M - 2, M - 2), lambda i, j: (A[i - 1, j] + A[i, j] + A[i + 1, j]) / 3, name="B") s = te.create_schedule(B.op) func = KgeN.lower(s, [A, B]) print(KgeN.build(func))
import KgeN from KgeN import te C = te.compute((5, 16), lambda ci, cj: 5, name='C') D = te.compute((5, 16), lambda di, dj: C[di, dj] * 2, name='D') E = te.compute((5, 16), lambda ei, ej: D[ei, ej] * 4, name='E') s = te.create_schedule(E.op) s[C].compute_at(s[D], s[D].op.axis[1]) s[D].compute_at(s[E], s[E].op.axis[1]) func = KgeN.lower(s, [E]) print(KgeN.build(func))
import KgeN from KgeN import te C = te.compute((5, 16), lambda ci, cj: ci * cj, name='C') D = te.compute((5, 16), lambda di, dj: C[di, dj] * 2, name='D') s = te.create_schedule(D.op) s[C].compute_inline() func = KgeN.lower(s, [D]) print(KgeN.build(func))