( # pylint: disable=too-many-locals,missing-docstring
self,
func: tirx.PrimFunc,
target: Target,
_: bool,
)
| 51 | """ |
| 52 | |
| 53 | def apply( # pylint: disable=too-many-locals,missing-docstring |
| 54 | self, |
| 55 | func: tirx.PrimFunc, |
| 56 | target: Target, |
| 57 | _: bool, |
| 58 | ) -> s_tir.Schedule: |
| 59 | if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target): |
| 60 | return None |
| 61 | max_threads_per_block = base.max_threads_per_block(target) |
| 62 | |
| 63 | sch = s_tir.Schedule(func) |
| 64 | block_infos = normalize_prim_func(sch) |
| 65 | |
| 66 | if block_infos is None: |
| 67 | return None |
| 68 | |
| 69 | block_infos = try_inline(sch, block_infos) |
| 70 | reduction_blocks: list[tuple[s_tir.schedule.SBlockRV, s_tir.schedule.LoopRV]] = [] |
| 71 | for block in block_infos: |
| 72 | s_loops: list[s_tir.schedule.LoopRV] = [] |
| 73 | r_loops: list[s_tir.schedule.LoopRV] = [] |
| 74 | o_loops: list[s_tir.schedule.LoopRV] = [] |
| 75 | dom_kind = block.dom_kind() |
| 76 | block = block.block_rv |
| 77 | |
| 78 | if any( |
| 79 | [sch.get(loop_rv).thread_binding is not None for loop_rv in sch.get_loops(block)] |
| 80 | ): |
| 81 | continue |
| 82 | |
| 83 | if len(sch.get_loops(block)) == 0 and _has_internal_thread_env(sch.get(block).body): |
| 84 | # The block (e.g. an opaque sort kernel) launches its own |
| 85 | # threads; binding an outer loop would conflict with them. |
| 86 | continue |
| 87 | |
| 88 | for loop, iter_type in zip(sch.get_loops(block), dom_kind): |
| 89 | {"S": s_loops, "R": r_loops, "O": o_loops}[iter_type].append(loop) |
| 90 | |
| 91 | if not s_loops: |
| 92 | s_loops.append(sch.add_unit_loop(block)) |
| 93 | sch.reorder(*s_loops, *r_loops, *o_loops) |
| 94 | bx, tx = sch.split( # pylint: disable=invalid-name |
| 95 | sch.fuse(*s_loops), |
| 96 | factors=[None, max_threads_per_block], |
| 97 | ) |
| 98 | sch.bind(bx, "blockIdx.x") |
| 99 | sch.bind(tx, "threadIdx.x") |
| 100 | |
| 101 | if len(r_loops) > 0: |
| 102 | reduction_blocks.append((block, r_loops[0])) |
| 103 | |
| 104 | for block, r_loop in reduction_blocks: |
| 105 | sch.decompose_reduction(block, r_loop) |
| 106 | |
| 107 | return sch |
nothing calls this directly
no test coverage detected