MCPcopy Create free account
hub / github.com/apache/tvm / CpAsyncBulkTensorNamespace

Class CpAsyncBulkTensorNamespace

python/tvm/backend/cuda/script.py:174–244  ·  view source on GitHub ↗

The CpAsyncBulkTensor instruction submodule.

Source from the content-addressed store, hash-verified

172
173
174class CpAsyncBulkTensorNamespace:
175 """The CpAsyncBulkTensor instruction submodule."""
176
177 def __init__(self):
178 self.g2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster)
179 self.g2c_tile_gather4 = _op_wrapper(
180 _cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster
181 )
182 self.s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global)
183 self.s2g_reduce = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global_reduce)
184 self.g2c_prefetch = _op_wrapper(
185 _cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch
186 )
187
188 @staticmethod
189 def g2c_bar_addr(
190 dim,
191 dst_ptr,
192 bar_addr,
193 tensormap_addr,
194 cta_mask,
195 cta_group,
196 cache_hint,
197 *coords,
198 cache_policy=None,
199 ):
200 _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP)
201 cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy)
202 return _tir_op.call_intrin(
203 "",
204 "tirx.ptx_cp_async_bulk_tensor_global_to_cluster",
205 dim,
206 dst_ptr,
207 bar_addr,
208 tensormap_addr,
209 cta_mask,
210 cta_group,
211 cache_policy,
212 int(has_cache_policy),
213 1,
214 *coords,
215 )
216
217 @staticmethod
218 def g2c_tile_gather4_bar_addr(
219 dim,
220 dst_ptr,
221 bar_addr,
222 tensormap_addr,
223 cta_mask,
224 cta_group,
225 cache_hint,
226 *coords,
227 cache_policy=None,
228 ):
229 _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP)
230 cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy)
231 return _tir_op.call_intrin(

Callers 1

__init__Method · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…