The CpAsync instruction submodule.
| 121 | |
| 122 | |
| 123 | class CpAsyncNamespace: |
| 124 | """The CpAsync instruction submodule.""" |
| 125 | |
| 126 | def __init__(self): |
| 127 | self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_commit_group) |
| 128 | self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_wait_group) |
| 129 | # Legacy variant: takes (dst_ptr, dst_offset, src_ptr, src_offset, |
| 130 | # cp_size). Offsets are folded into the pointers; coexists with |
| 131 | # the fork-native ``__call__`` form. |
| 132 | self.legacy = _dtype_forward(_cuda_op.ptx_cp_async_legacy) |
| 133 | self.bulk = CpAsyncBulkNamespace() |
| 134 | self.mbarrier = CpAsyncMbarrierNamespace() |
| 135 | |
| 136 | def __call__(self, *args, **kwds): |
| 137 | # Accept the legacy 6-arg form ``(elem_dtype, dst, dst_off, src, |
| 138 | # src_off, cp_size)`` that the printer round-trips for the raw |
| 139 | # ``tirx.ptx_cp_async`` Call emitted by |
| 140 | # ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``. The pass-emitted |
| 141 | # Call has 5 args (no ``tvm_access_ptr`` fold) and a |
| 142 | # per-element-dtype Call.dtype, so build it directly. |
| 143 | if len(args) == 6 and isinstance(args[0], str) and "dtype" not in kwds: |
| 144 | import tvm |
| 145 | |
| 146 | elem_dtype, dst, dst_off, src, src_off, cp_size = args |
| 147 | return tvm.tirx.Call( |
| 148 | tvm.DataType(elem_dtype), |
| 149 | tvm.ir.Op.get("tirx.ptx_cp_async"), |
| 150 | [dst, dst_off, src, src_off, cp_size], |
| 151 | ) |
| 152 | return _dtype_forward(_cuda_op.ptx_cp_async)(*args, **kwds) |
| 153 | |
| 154 | # __call__ corresponds to ptx_cp_async |
| 155 | __tir_call_op_name__ = "ptx_cp_async" |
| 156 | |
| 157 | |
| 158 | class CpAsyncBulkNamespace: |
no outgoing calls
no test coverage detected
searching dependent graphs…