Python create_scheduleの例

コード例 #1

ファイル: test_bound.py プロジェクト: Cjkkkk/KgeN

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

コード例 #2

ファイル: test_bound.py プロジェクト: Cjkkkk/KgeN

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

コード例 #3

ファイル: test_bound.py プロジェクト: Cjkkkk/KgeN

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

コード例 #4

ファイル: test_bound.py プロジェクト: Cjkkkk/KgeN

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

コード例 #5

ファイル: schedule.py プロジェクト: Cjkkkk/KgeN

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))

コード例 #6

ファイル: const.py プロジェクト: Cjkkkk/KgeN

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))

コード例 #7

ファイル: blur.py プロジェクト: Cjkkkk/KgeN

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))

コード例 #8

ファイル: nest_attach.py プロジェクト: Cjkkkk/KgeN

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))

コード例 #9

ファイル: compute_inline.py プロジェクト: Cjkkkk/KgeN

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))