Function initial_A

python/tvm/tirx/exec_context.py:196–204

Build A at PrimFunc device entry: all threads active, offsets all zero.

(*, lane_ext: int = 32, warp_ext: int, cta_ext: int = 1)

def initial_A(*, lane_ext: int = 32, warp_ext: int, cta_ext: int = 1) -> ActiveSet:
    """Build A at PrimFunc device entry: all threads active, offsets all zero."""
    return ActiveSet.from_axes(
        [
            ("laneid", AxisRange(lane_ext, 0)),
            ("warpid", AxisRange(warp_ext, 0)),
            ("cta_id", AxisRange(cta_ext, 0)),
        ]
    )


def filter_narrow(A: ActiveSet, binding: LaneBinding, lo: int, hi: int) -> ActiveSet: