def root_anno(sch: Schedule):
b8 = sch.get_block(name="root", func_name="main")
v140 = sch.sample_categorical(
candidates=[0, 16, 64, 512, 1024],
probs=[
0.20000000000000001,
0.20000000000000001,
0.20000000000000001,
0.20000000000000001,
0.20000000000000001,
],
decision=2,
)
sch.annotate(block_or_loop=b8, ann_key="meta_schedule.unroll_explicit", ann_val=v140)
def bgemm(sch: Schedule):
b31 = sch.get_block(name="bgemm")
sch.annotate(
block_or_loop=b31,
ann_key="meta_schedule.tiling_structure",
ann_val="SSSRRSRS",
)
b32 = sch.cache_write(block=b31,
write_buffer_index=0,
storage_scope="local")
b31, b32 = b32, b31
l33, l34, l35, l36, l37 = sch.get_loops(block=b32)
v38, v39, v40, v41, v42 = sch.sample_perfect_tile(
n=5,
loop=l33,
max_innermost_factor=64,
decision=[1, 1, 1, 1, 6],
)
l43, l44, l45, l46, l47 = sch.split(loop=l33,
factors=[v38, v39, v40, v41, v42])
v48, v49, v50, v51, v52 = sch.sample_perfect_tile(
n=5,
loop=l34,
max_innermost_factor=64,
decision=[1, 1, 1, 3, 2],
)
l53, l54, l55, l56, l57 = sch.split(loop=l34,
factors=[v48, v49, v50, v51, v52])
v58, v59, v60, v61, v62 = sch.sample_perfect_tile(
n=5,
loop=l35,
max_innermost_factor=64,
decision=[3, 1, 1, 1, 3],
)
l63, l64, l65, l66, l67 = sch.split(loop=l35,
factors=[v58, v59, v60, v61, v62])
v68, v69, v70, v71, v72 = sch.sample_perfect_tile(
n=5,
loop=l36,
max_innermost_factor=64,
decision=[4, 2, 1, 4, 4],
)
l73, l74, l75, l76, l77 = sch.split(loop=l36,
factors=[v68, v69, v70, v71, v72])
v78, v79, v80 = sch.sample_perfect_tile(
n=3,
loop=l37,
max_innermost_factor=64,
decision=[32, 1, 4],
)
l81, l82, l83 = sch.split(loop=l37, factors=[v78, v79, v80])
sch.reorder(
# fmt: off
l43,
l53,
l63,
l73,
l44,
l54,
l64,
l74,
l45,
l55,
l65,
l75,
l81,
l82,
l46,
l56,
l66,
l76,
l83,
l47,
l57,
l67,
l77,
# fmt: on
)
l84 = sch.fuse(l43, l53, l63, l73)
sch.bind(loop=l84, thread_axis="blockIdx.x")
l85 = sch.fuse(l44, l54, l64, l74)
sch.bind(loop=l85, thread_axis="vthread.x")
l86 = sch.fuse(l45, l55, l65, l75)
sch.bind(loop=l86, thread_axis="threadIdx.x")
b87 = sch.cache_read(block=b32,
read_buffer_index=1,
storage_scope="shared")
sch.compute_at(block=b87, loop=l81, preserve_unit_loops=True)
_, _, _, _, l92, l93, l94, l95 = sch.get_loops(block=b87)
sch.fuse(l92, l93, l94, l95)
v97 = sch.sample_categorical(
candidates=[1, 2, 3, 4],
probs=[0.25, 0.25, 0.25, 0.25],
decision=1,
)
sch.annotate(
block_or_loop=b87,
ann_key="meta_schedule.cooperative_fetch",
ann_val=v97,
)
b101 = sch.cache_read(block=b32,
read_buffer_index=2,
storage_scope="shared")
sch.compute_at(block=b101, loop=l81, preserve_unit_loops=True)
_, _, _, _, l106, l107, l108, l109 = sch.get_loops(block=b101)
sch.fuse(l106, l107, l108, l109)
v110 = sch.sample_categorical(
candidates=[1, 2, 3, 4],
probs=[0.25, 0.25, 0.25, 0.25],
decision=1,
)
sch.annotate(
block_or_loop=b101,
ann_key="meta_schedule.cooperative_fetch",
ann_val=v110,
)
sch.reverse_compute_at(block=b31, loop=l86, preserve_unit_loops=True)
def bgemm(sch: Schedule):
bgemm = sch.get_block(name="bgemm")
write_cache = sch.cache_write(
block=bgemm,
write_buffer_index=0,
storage_scope="global",
)
sch.annotate(
block_or_loop=bgemm,
ann_key="meta_schedule.tiling_structure",
ann_val="SSRSRS",
)
# b33, b34 = b34, b33
l35, l36, l37, l38, l39 = sch.get_loops(block=bgemm)
v40, v41, v42, v43 = sch.sample_perfect_tile(
n=4,
loop=l35,
max_innermost_factor=64,
decision=[1, 2, 3, 1],
)
l44, l45, l46, l47 = sch.split(loop=l35, factors=[v40, v41, v42, v43])
v48, v49, v50, v51 = sch.sample_perfect_tile(
n=4,
loop=l36,
max_innermost_factor=64,
decision=[1, 1, 1, 6],
)
l52, l53, l54, l55 = sch.split(loop=l36, factors=[v48, v49, v50, v51])
v56, v57, v58, v59 = sch.sample_perfect_tile(
n=4,
loop=l37,
max_innermost_factor=64,
decision=[1, 1, 1, 9],
)
l60, l61, l62, l63 = sch.split(loop=l37, factors=[v56, v57, v58, v59])
v64, v65, v66, v67 = sch.sample_perfect_tile(
n=4,
loop=l38,
max_innermost_factor=64,
decision=[2, 1, 16, 4],
)
l68, l69, l70, l71 = sch.split(loop=l38, factors=[v64, v65, v66, v67])
v72, v73 = sch.sample_perfect_tile(
n=2,
loop=l39,
max_innermost_factor=64,
decision=[16, 8],
)
l74, l75 = sch.split(loop=l39, factors=[v72, v73])
sch.reorder(
# fmt: off
l44,
l52,
l60,
l68,
l45,
l53,
l61,
l69,
l74,
l46,
l54,
l62,
l70,
l75,
l47,
l55,
l63,
l71,
# fmt: on
)
sch.reverse_compute_at(block=write_cache,
loop=l69,
preserve_unit_loops=True)
