def _sch(decisions: List[List[int]]) -> Schedule: sch = Schedule(matmul, debug_mask="all") # pylint: disable=invalid-name d0, d1, d2 = decisions b0 = sch.get_block(name="C", func_name="main") root = sch.get_block(name="root", func_name="main") sch.get_consumers(block=b0) b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global") l2, l3, l4 = sch.get_loops(block=b0) v5, v6, v7, v8 = sch.sample_perfect_tile( loop=l2, n=4, max_innermost_factor=64, decision=d0, ) l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8]) v13, v14, v15, v16 = sch.sample_perfect_tile( loop=l3, n=4, max_innermost_factor=64, decision=d1, ) l17, l18, l19, l20 = sch.split(loop=l3, factors=[v13, v14, v15, v16]) v21, v22 = sch.sample_perfect_tile( loop=l4, n=2, max_innermost_factor=64, decision=d2, ) l23, l24 = sch.split(loop=l4, factors=[v21, v22]) sch.reorder(l9, l17, l10, l18, l23, l11, l19, l24, l12, l20) sch.reverse_compute_at(block=b1, loop=l18, preserve_unit_loops=True) v57 = sch.sample_categorical( candidates=[0, 16, 64, 512], probs=[0.25, 0.25, 0.25, 0.25], decision=0, ) sch.annotate(block_or_loop=root, ann_key="meta_schedule.unroll_explicit", ann_val=v57) # pylint: enable=invalid-name return sch
def _sch() -> Schedule: sch = Schedule(element_wise, debug_mask="all") # pylint: disable=invalid-name b0 = sch.get_block(name="C", func_name="main") l1, l2 = sch.get_loops(block=b0) l3 = sch.fuse(l1, l2) v4 = sch.sample_categorical( candidates=[32, 64, 128, 256, 512, 1024], probs=[ 0.16666666666666666, 0.16666666666666666, 0.16666666666666666, 0.16666666666666666, 0.16666666666666666, 0.16666666666666666, ], decision=3, ) l5, l6 = sch.split(loop=l3, factors=[None, v4]) sch.bind(loop=l5, thread_axis="blockIdx.x") sch.bind(loop=l6, thread_axis="threadIdx.x") # pylint: enable=invalid-name return sch