def input_tile_data_pad(sch: Schedule):
b115 = sch.get_block(name="input_tile")
(b116,) = sch.get_consumers(block=b115)
_, _, _, l120, _, _, _, _ = sch.get_loops(block=b116)
sch.compute_at(block=b115, loop=l120, preserve_unit_loops=True)
sch.set_scope(block=b115, buffer_index=0, storage_scope="local")
b127 = sch.get_block(name="data_pad")
sch.compute_inline(block=b127)
b3 = sch.get_block(name="data_pack")
l25, l26, l27, l28, _, _, _, _ = sch.get_loops(block=b3)
l33 = sch.fuse(l25, l26, l27, l28)
v34 = sch.sample_categorical(
candidates=[32, 64, 128, 256, 512, 1024],
probs=[
0.16666666666666666,
0.16666666666666666,
0.16666666666666666,
0.16666666666666666,
0.16666666666666666,
0.16666666666666666,
],
decision=2,
)
l35, l36 = sch.split(loop=l33, factors=[None, v34])
sch.bind(loop=l35, thread_axis="blockIdx.x")
sch.bind(loop=l36, thread_axis="threadIdx.x")
def input_tile_data_pad(sch: Schedule):
b78 = sch.get_block(name="input_tile")
l80 = sch.sample_compute_location(block=b78, decision=4)
sch.compute_at(block=b78, loop=l80, preserve_unit_loops=True)
b81 = sch.get_block(name="data_pad")
l83 = sch.sample_compute_location(block=b81, decision=-2)
sch.compute_at(block=b81, loop=l83, preserve_unit_loops=True)
def input_tile_data_pad(sch: Schedule):
b115 = sch.get_block(name="input_tile")
(b116, ) = sch.get_consumers(block=b115)
_, _, _, l120, _, _, _, _ = sch.get_loops(block=b116)
sch.compute_at(block=b115, loop=l120, preserve_unit_loops=True)
sch.set_scope(block=b115, buffer_index=0, storage_scope="local")
b127 = sch.get_block(name="data_pad")
sch.compute_inline(block=b127)
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)
