Function _factor_warpid

python/tvm/tirx/exec_context.py:295–310

(warp: AxisRange)

def _factor_warpid(warp: AxisRange) -> tuple[AxisRange, AxisRange] | None:
    if warp.stride != 1:
        return None
    off = warp.offset
    ext = warp.extent
    wid_off = off % WG_SIZE
    wgid_off = off // WG_SIZE

    if wid_off == 0 and ext % WG_SIZE == 0:
        return (
            AxisRange(extent=WG_SIZE, offset=0),
            AxisRange(extent=ext // WG_SIZE, offset=wgid_off),
        )
    if ext <= WG_SIZE - wid_off:
        return (AxisRange(extent=ext, offset=wid_off), AxisRange(extent=1, offset=wgid_off))
    return None


def _flat_product_range(