3.0 KiB
3.0 KiB
CuTe Experimental APIs
Note: APIs in this module are experimental and subject to change.
This module serves as a staging area for new CuTe functionality that is still under active development. Performance, compile time, and interoperability with CuTe are works in progress. API signatures, behavior, and naming conventions may change without notice between releases.
Once these APIs are stabilized, they will be migrated to the main
cutesubmodules.Users are encouraged to experiment with these APIs but should be prepared to update their code as the interfaces evolve.
Core APIs (core.py)
elect_sync— Elects one thread within a warpget_mbarrier— Returns the mbarrier pointer for a given stage tokencreate_pipeline— Creates a circular buffer of synchronization primitives indexed by stage countcreate_pipeline_with_mask— Creates a pipeline with an arrival mask for cluster-scoped synchronizationpipeline_advance_iterator— Advances a pipeline iterator to the next stageproducer_acquire/producer_commit— Producer-side pipeline synchronizationconsumer_wait/consumer_release/consumer_tail— Consumer-side pipeline synchronizationget_pipeline_produce_stage/get_pipeline_consume_stage— Gets pipeline stage tokens
Memory APIs (memory.py)
allocate— Allocate a buffer with given type, layout, and address spacetma_load— Copy tensor from global memory to shared memory using TMAtma_load_multicast— Copy tensor from global memory to shared memory using TMA with multicasttma_store— Copy tensor from shared memory to global memory using TMAcopy— Copy tensor from src to dst using a given copy atom
Algorithm APIs (algorithm.py)
simt_auto_vec_copy— Copies a tensor between buffers with single thread (auto-vectorized)partition— Partition a buffer into a given layout and tilerpartition_and_copy— Combines partitioning and copying in a single operation
Math APIs (math.py)
dot— Computes a dot product of two tensors using an MMA atomdot_block_scaled— Computes a block-scaled dot product with scale factors
Pipeline Classes (pipeline.py)
GenericPipeline— Generic pipeline for any producer/consumer combinationTMAToUMMAPipeline— Pipeline for TMA load to UMMA consumptionTMAToAsyncPipeline— Pipeline for TMA load to async consumerAsyncToUMMAPipeline— Pipeline for async producer to UMMA consumptionUMMAtoAsyncPipeline— Pipeline for UMMA producer to async consumerTMAStorePipeline— Pipeline for SMEM producer to TMA store consumer
Utilities (utils.py)
get_cta_v_map_ab— Compute CTA-V map for A/B operandsget_cta_v_map_c— Compute CTA-V map for C operandmake_tmem_layout_acc— Derive TMEM accumulator buffer layout from a tiled MMAmake_tmem_layout_a— Derive TMEM A-operand buffer layout from a tiled MMAmake_t2r_rmem_layout— Derive per-thread RMEM buffer layout for the T2R epilogue copy