def apply(self, sch: Schedule, block: BlockRV) -> List[Schedule]:
if _is_root(sch, block):
return [sch]
sch = sch.copy()
if sch.get(block).name_hint == "B":
sch.compute_inline(block)
return [sch]
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):
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 inline(sch: Schedule):
b125 = sch.get_block(name="A")
sch.compute_inline(block=b125)
b126 = sch.get_block(name="B")
sch.compute_inline(block=b126)
def apply(self, sch: Schedule, block: BlockRV):
if sch.get(block).name_hint == "root":
return [sch]
sch = sch.copy()
sch.compute_inline(block)
return [sch]
def inline(sch: Schedule):
b1 = sch.get_block(name="A")
b2 = sch.get_block(name="B")
sch.compute_inline(block=b1)
sch.compute_inline(block=b2)