You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/06/14 01:12:46 UTC
(tvm) branch main updated: [TIRx] Post-bringup follow-ups: op-dispatch, namespaces, launch bounds, gemm-async, backend reorg (#19757)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new bb6f8aec55 [TIRx] Post-bringup follow-ups: op-dispatch, namespaces, launch bounds, gemm-async, backend reorg (#19757)
bb6f8aec55 is described below
commit bb6f8aec55d85f221492aee51380356b66767dca
Author: Bohan Hou <bo...@andrew.cmu.edu>
AuthorDate: Sat Jun 13 18:12:40 2026 -0700
[TIRx] Post-bringup follow-ups: op-dispatch, namespaces, launch bounds, gemm-async, backend reorg (#19757)
This PR batches several post-bringup TIRx follow-ups, rebased onto
current `main`.
### Changes
- **op-dispatch**: per-call exec scope via `Tx.<scope>.op`; remove
`ExecScopeStmt`
- **namespaces**: split TIRx op namespaces; remove tile-primitive kind
attrs
- **codegen**: support explicit CUDA launch bounds
- **gemm-async**: support contiguous-axis (K-major) operand slicing
- **backend reorg**: move in-tree GPU backends out of core into
`src/backend/<target>/` and `python/tvm/backend/<target>/`
(codegen/runtime/op), with the corresponding `CMakeLists.txt` /
`cmake/modules` and include-path updates
### Testing
- Builds with `USE_CUDA=ON` / `USE_LLVM=ON`
- The TIRx Python test suite (`tests/python/tirx/`) passes locally
---
CMakeLists.txt | 22 +-
cmake/modules/CUDA.cmake | 4 +-
cmake/modules/Hexagon.cmake | 63 +-
cmake/modules/LLVM.cmake | 6 +-
cmake/modules/Metal.cmake | 4 +-
cmake/modules/OpenCL.cmake | 4 +-
cmake/modules/ROCM.cmake | 2 +-
cmake/modules/Vulkan.cmake | 14 +-
docs/reference/api/python/tirx/backend.rst | 6 -
include/tvm/tirx/builtin.h | 270 -
include/tvm/tirx/op.h | 2 -
include/tvm/tirx/target_builtin/cuda.h | 745 ---
include/tvm/tirx/target_builtin/trn.h | 156 -
python/tvm/__init__.py | 8 +-
python/tvm/backend/__init__.py | 209 +
.../default.py => backend/adreno/__init__.py} | 30 +-
.../adreno.py => backend/adreno/target_tags.py} | 2 +-
python/tvm/backend/cuda/__init__.py | 71 +
python/tvm/backend/cuda/lang/__init__.py | 70 +
.../tvm/{tirx => backend/cuda}/lang/alloc_pool.py | 10 +-
python/tvm/{tirx => backend/cuda}/lang/pipeline.py | 0
.../tvm/{tirx => backend/cuda}/lang/smem_desc.py | 2 +-
.../{tirx => backend/cuda}/lang/tile_scheduler.py | 0
.../tvm/{tirx => backend/cuda}/lang/warp_role.py | 0
python/tvm/{tirx => backend/cuda}/op.py | 4071 +--------------
.../cuda/operator}/__init__.py | 3 +-
.../cuda/operator/intrinsics}/__init__.py | 0
.../cuda}/operator/intrinsics/_schema.py | 4 +-
.../cuda/operator/intrinsics}/cp_async.py | 7 +-
.../cuda/operator/intrinsics}/header.py | 0
.../cuda/operator/intrinsics}/math.py | 4 +-
.../cuda/operator/intrinsics}/memory.py | 4 +-
.../cuda/operator/intrinsics}/misc.py | 4 +-
.../cuda/operator/intrinsics}/mma.py | 2 +-
.../cuda/operator/intrinsics}/nvshmem.py | 2 +-
.../cuda/operator/intrinsics}/registry.py | 0
.../cuda/operator/intrinsics}/sync.py | 10 +-
.../cuda/operator/intrinsics}/tcgen05.py | 2 +-
.../cuda/operator/intrinsics}/types.py | 0
.../cuda/operator/intrinsics}/utils.py | 0
.../cuda/operator/intrinsics}/wgmma.py | 2 +-
.../cuda/operator/tile_primitive}/__init__.py | 4 +
.../cuda/operator/tile_primitive}/common.py | 0
.../cuda/operator/tile_primitive}/copy/__init__.py | 0
.../cuda/operator/tile_primitive}/copy/_common.py | 0
.../operator/tile_primitive}/copy/_swizzle_iter.py | 0
.../cuda/operator/tile_primitive}/copy/fallback.py | 0
.../operator/tile_primitive}/copy/gmem_smem.py | 2 +-
.../operator/tile_primitive}/copy/ld_stmatrix.py | 2 +-
.../cuda/operator/tile_primitive}/copy/reg.py | 2 +-
.../cuda/operator/tile_primitive}/copy/utils.py | 0
.../tile_primitive}/copy_async/__init__.py | 0
.../operator/tile_primitive}/copy_async/dsmem.py | 0
.../operator/tile_primitive}/copy_async/ldgsts.py | 2 +-
.../tile_primitive}/copy_async/tcgen05_cp.py | 2 +-
.../tile_primitive}/copy_async/tcgen05_ldst.py | 0
.../operator/tile_primitive}/copy_async/tma.py | 0
.../operator/tile_primitive}/copy_async/utils.py | 0
.../tile_primitive}/elementwise/__init__.py | 2 +-
.../tile_primitive}/elementwise/_common.py | 0
.../tile_primitive}/elementwise/ops/__init__.py | 0
.../tile_primitive}/elementwise/ops/binary.py | 0
.../tile_primitive}/elementwise/ops/cast.py | 0
.../tile_primitive}/elementwise/ops/fma.py | 0
.../tile_primitive}/elementwise/ops/unary.py | 0
.../operator/tile_primitive}/elementwise/reg.py | 2 +-
.../tile_primitive}/elementwise/register.py | 0
.../operator/tile_primitive}/elementwise/smem.py | 2 +-
.../elementwise/vec_emit/__init__.py | 0
.../elementwise/vec_emit/binary_f32x2.py | 0
.../elementwise/vec_emit/cast_vec2.py | 0
.../elementwise/vec_emit/fma_f32x2.py | 0
.../operator/tile_primitive}/exec_scope_utils.py | 0
.../cuda/operator/tile_primitive}/gemm/__init__.py | 0
.../operator/tile_primitive}/gemm/mma_m16n8k_.py | 0
.../tile_primitive}/gemm_async/__init__.py | 0
.../operator/tile_primitive}/gemm_async/tcgen05.py | 211 +-
.../cuda/operator/tile_primitive}/gemm_utils.py | 0
.../cuda/operator/tile_primitive}/layout_utils.py | 0
.../tile_primitive}/permute_layout/__init__.py | 0
.../permute_layout/warp_xor_swizzle.py | 0
.../operator/tile_primitive}/reduction/__init__.py | 0
.../operator/tile_primitive}/reduction/local.py | 4 +-
.../operator/tile_primitive}/reduction/shared.py | 2 +-
.../tile_primitive}/reduction/sm100_packed.py | 2 +-
.../operator/tile_primitive}/reduction/utils.py | 9 +-
.../cuda/operator/tile_primitive}/tma_utils.py | 0
python/tvm/backend/cuda/script.py | 571 +++
.../cuda.py => backend/cuda/target_tags.py} | 2 +-
.../default.py => backend/hexagon/__init__.py} | 26 +-
.../hexagon.py => backend/hexagon/target_tags.py} | 2 +-
python/tvm/backend/metal/__init__.py | 58 +
python/tvm/backend/metal/op.py | 84 +
python/tvm/backend/metal/script.py | 55 +
.../metal.py => backend/metal/target_tags.py} | 4 +-
.../{tirx/backend => backend/opencl}/__init__.py | 10 +-
.../tvm/{tirx/backend => backend/rocm}/__init__.py | 10 +-
python/tvm/backend/trn/__init__.py | 68 +
python/tvm/backend/trn/layout.py | 123 +
python/tvm/backend/trn/op.py | 153 +
.../backend => backend/trn/operator}/__init__.py | 7 +-
.../trn/operator/tile_primitive}/__init__.py | 0
.../operator/tile_primitive}/binary/__init__.py | 0
.../trn/operator/tile_primitive}/binary/default.py | 4 +-
.../trn/operator/tile_primitive}/binary/utils.py | 7 +-
.../trn/operator/tile_primitive}/common.py | 0
.../tile_primitive}/compose_op/__init__.py | 0
.../tile_primitive}/compose_op/binary_chain.py | 0
.../tile_primitive}/compose_op/binary_reduce.py | 0
.../tile_primitive}/compose_op/compose_op.py | 0
.../tile_primitive}/compose_op/reduce_negate.py | 0
.../tile_primitive}/compose_op/unary_reduce.py | 0
.../operator/tile_primitive}/compose_op/utils.py | 3 +-
.../trn/operator/tile_primitive}/copy/__init__.py | 0
.../trn/operator/tile_primitive}/copy/default.py | 11 +-
.../trn/operator/tile_primitive}/dim_utils.py | 0
.../trn/operator/tile_primitive}/gemm/__init__.py | 0
.../trn/operator/tile_primitive}/gemm/default.py | 11 +-
.../tile_primitive}/instruction_generator.py | 11 +-
.../trn/operator/tile_primitive}/private_alloc.py | 6 +-
.../operator/tile_primitive}/reduction/__init__.py | 0
.../operator/tile_primitive}/reduction/default.py | 2 +-
.../operator/tile_primitive}/reduction/utils.py | 11 +-
.../operator/tile_primitive}/select/__init__.py | 0
.../trn/operator/tile_primitive}/select/default.py | 5 +-
.../trn/operator/tile_primitive}/unary/__init__.py | 0
.../trn/operator/tile_primitive}/unary/default.py | 4 +-
.../trn/operator/tile_primitive}/unary/utils.py | 9 +-
.../tile_primitive}/unary/with_bias_scale.py | 4 +-
.../operator/tile_primitive}/workspace_utils.py | 0
python/tvm/backend/trn/pipeline.py | 58 +
python/tvm/backend/trn/script.py | 58 +
.../tvm/backend/trn/target_tags.py | 25 +-
.../trn => backend/trn/transform}/__init__.py | 27 +-
.../trn/transform}/naive_allocator.py | 3 +-
.../trn/transform}/private_buffer_alloc.py | 0
.../{tirx/backend => backend/vulkan}/__init__.py | 10 +-
.../{tirx/backend => backend/webgpu}/__init__.py | 10 +-
python/tvm/s_tir/tensor_intrin/cuda.py | 4 +-
python/tvm/s_tir/tensor_intrin/metal.py | 8 +-
python/tvm/target/tag_registry/__init__.py | 4 -
python/tvm/tirx/__init__.py | 17 +-
python/tvm/tirx/backend/__init__.py | 4 +-
python/tvm/tirx/backend/adreno/__init__.py | 17 -
python/tvm/tirx/bench.py | 8 +-
python/tvm/tirx/compilation_pipeline.py | 39 +-
python/tvm/tirx/lang/alloc_pool.py | 527 +-
python/tvm/tirx/lang/pipeline.py | 242 +-
python/tvm/tirx/lang/smem_desc.py | 53 +-
python/tvm/tirx/lang/tile_scheduler.py | 814 +--
python/tvm/tirx/lang/warp_role.py | 142 +-
python/tvm/tirx/op.py | 5285 +-------------------
python/tvm/tirx/operator/intrinsics/_common.py | 2 +-
.../tvm/tirx/operator/tile_primitive/__init__.py | 11 +-
.../operator/tile_primitive/dispatch_context.py | 4 +
python/tvm/tirx/script/builder/ir.py | 656 +--
python/tvm/tirx/transform/__init__.py | 1 -
python/tvm/tirx/transform/transform.py | 11 +
src/arith/const_int_bound.cc | 7 +-
src/arith/ir_mutator_with_analyzer.cc | 9 +-
src/arith/ir_visitor_with_analyzer.cc | 5 +-
src/arith/rewrite_simplify.cc | 11 +-
.../cuda => backend/cuda/codegen}/codegen_cuda.cc | 74 +-
.../cuda => backend/cuda/codegen}/codegen_cuda.h | 2 +-
.../cuda/codegen}/cuda_fallback_module.cc | 2 +-
.../cuda/codegen}/cuda_fallback_module.h | 4 +-
.../cuda/codegen}/intrin_rule_cuda.cc | 49 +-
.../cuda/codegen}/literal/cuda_half_t.h | 0
.../cuda/codegen}/literal/cuda_int8_t.h | 0
.../cuda/codegen}/llvm/codegen_nvptx.cc | 14 +-
src/{target/cuda => backend/cuda/codegen}/ptx.cc | 28 +-
src/{target/cuda => backend/cuda/codegen}/ptx.h | 0
src/backend/cuda/codegen/register.cc | 143 +
.../prof_utils.h => backend/cuda/op/register.cc} | 19 +-
.../cuda.cc => backend/cuda/op/target_builtin.cc} | 69 +-
.../cuda => backend/cuda/runtime}/cuda_common.h | 2 +-
.../cuda/runtime}/cuda_device_api.cc | 0
.../cuda => backend/cuda/runtime}/cuda_module.cc | 8 +-
.../cuda/runtime}/l2_cache_flush.cc | 2 +-
.../cuda/runtime/vm}/cuda_graph_builtin.cc | 4 +-
.../hexagon/codegen}/hexagon_fallback_module.cc | 4 +-
.../hexagon/codegen}/hexagon_fallback_module.h | 4 +-
.../hexagon/codegen}/llvm/codegen_hexagon.cc | 16 +-
.../hexagon/codegen}/llvm/intrin_rule_hexagon.cc | 35 +-
src/backend/hexagon/codegen/register.cc | 64 +
.../hexagon => backend/hexagon/runtime}/README.md | 0
.../hexagon/runtime}/hexagon_buffer.cc | 0
.../hexagon/runtime}/hexagon_buffer.h | 0
.../hexagon/runtime}/hexagon_buffer_manager.h | 0
.../hexagon/runtime}/hexagon_common.cc | 0
.../hexagon/runtime}/hexagon_common.h | 0
.../hexagon/runtime}/hexagon_device_api.cc | 2 +-
.../hexagon/runtime}/hexagon_device_api.h | 0
.../hexagon/runtime}/hexagon_htp.cc | 0
.../hexagon/runtime}/hexagon_htp.h | 0
.../hexagon/runtime}/hexagon_hvx.cc | 0
.../hexagon/runtime}/hexagon_hvx.h | 0
.../hexagon/runtime}/hexagon_module.cc | 10 +-
.../hexagon/runtime}/hexagon_power_manager.cc | 0
.../hexagon/runtime}/hexagon_power_manager.h | 0
.../hexagon/runtime}/hexagon_thread_manager.cc | 0
.../hexagon/runtime}/hexagon_thread_manager.h | 0
.../hexagon/runtime}/hexagon_user_dma.cc | 0
.../hexagon/runtime}/hexagon_user_dma.h | 0
.../runtime}/hexagon_user_dma_descriptors.h | 0
.../runtime}/hexagon_user_dma_instructions.h | 0
.../hexagon/runtime}/hexagon_user_dma_registers.h | 0
.../hexagon/runtime}/hexagon_vtcm_pool.cc | 0
.../hexagon/runtime}/hexagon_vtcm_pool.h | 0
.../hexagon/runtime}/ops/conv2d.h | 0
.../hexagon/runtime}/ops/conv2d_fp16_hvx.cc | 0
.../hexagon/runtime}/ops/conv2d_quant_hvx.cc | 0
.../hexagon/runtime}/ops/conv_utils.cc | 0
.../hexagon/runtime}/profiler/README.md | 0
.../hexagon/runtime}/profiler/lwp_handler.S | 0
.../hexagon/runtime}/profiler/prof_utils.cc | 0
.../hexagon/runtime}/profiler/prof_utils.h | 0
.../hexagon/runtime}/qhl/qhl_wrapper.cc | 0
.../hexagon/runtime}/ring_buffer.h | 0
.../hexagon/runtime}/rpc/android/session.cc | 6 +-
.../hexagon/runtime}/rpc/android_bash.sh.template | 0
.../hexagon/runtime}/rpc/hexagon/rpc_server.cc | 6 +-
.../hexagon/runtime}/rpc/hexagon_rpc.idl | 0
.../runtime}/rpc/simulator/hexagon_sim_proto.h | 0
.../hexagon/runtime}/rpc/simulator/rpc_server.cc | 2 +-
.../hexagon/runtime}/rpc/simulator/session.cc | 6 +-
.../metal/codegen}/codegen_metal.cc | 24 +-
.../metal/codegen}/codegen_metal.h | 2 +-
.../metal/codegen}/intrin_rule_metal.cc | 22 +-
.../metal/codegen}/metal_fallback_module.cc | 2 +-
.../metal/codegen}/metal_fallback_module.h | 4 +-
src/backend/metal/codegen/register.cc | 63 +
.../prof_utils.h => backend/metal/op/register.cc} | 19 +-
src/backend/metal/op/target_builtin.cc | 62 +
.../metal => backend/metal/runtime}/metal_common.h | 2 +-
.../metal/runtime}/metal_device_api.mm | 0
.../metal/runtime}/metal_module.mm | 10 +-
.../opencl/codegen}/codegen_opencl.cc | 18 +-
.../opencl/codegen}/codegen_opencl.h | 2 +-
.../opencl/codegen}/intrin_rule_opencl.cc | 38 +-
.../opencl/codegen}/opencl_fallback_module.cc | 4 +-
.../opencl/codegen}/opencl_fallback_module.h | 4 +-
src/backend/opencl/codegen/register.cc | 67 +
.../opencl/runtime}/opencl_common.h | 8 +-
.../opencl/runtime}/opencl_device_api.cc | 4 +-
.../opencl/runtime}/opencl_module.cc | 13 +-
.../opencl/runtime}/opencl_wrapper/README.md | 0
.../runtime}/opencl_wrapper/opencl_wrapper.cc | 0
.../opencl/runtime}/source_utils.h | 0
.../opencl => backend/opencl/runtime}/texture.h | 0
.../rocm/codegen}/llvm/codegen_amdgpu.cc | 14 +-
.../rocm/codegen}/llvm/intrin_rule_rocm.cc | 8 +-
src/backend/rocm/codegen/register.cc | 148 +
.../rocm/codegen}/rocm_fallback_module.cc | 2 +-
.../rocm/codegen}/rocm_fallback_module.h | 4 +-
.../rocm => backend/rocm/runtime}/rocm_common.h | 2 +-
.../rocm/runtime}/rocm_device_api.cc | 0
.../rocm => backend/rocm/runtime}/rocm_module.cc | 8 +-
.../source => backend/trn/codegen}/codegen_trn.cc | 59 +-
.../source => backend/trn/codegen}/codegen_trn.h | 2 +-
.../trn/codegen/register.cc} | 55 +-
.../prof_utils.h => backend/trn/op/register.cc} | 19 +-
.../trn.cc => backend/trn/op/target_builtin.cc} | 34 +-
src/backend/trn/transform/lower_trainium_layout.cc | 361 ++
.../vulkan/codegen}/build_vulkan.cc | 14 +-
.../vulkan/codegen}/codegen_spirv.cc | 21 +-
.../vulkan/codegen}/codegen_spirv.h | 4 +-
.../vulkan/codegen}/intrin_rule_spirv.cc | 26 +-
.../vulkan/codegen}/ir_builder.cc | 0
.../vulkan => backend/vulkan/codegen}/ir_builder.h | 0
src/backend/vulkan/codegen/register.cc | 93 +
.../vulkan/codegen}/spirv_support.cc | 0
.../vulkan/codegen}/spirv_support.h | 0
.../vulkan/codegen}/spirv_utils.cc | 4 +-
.../vulkan/codegen}/spirv_utils.h | 2 +-
.../vulkan/codegen}/vulkan_fallback_module.cc | 8 +-
.../vulkan/codegen}/vulkan_fallback_module.h | 4 +-
.../vulkan => backend/vulkan/runtime}/README.md | 0
.../vulkan/runtime}/spirv_shader.h | 0
.../vulkan => backend/vulkan/runtime}/thread_map.h | 0
.../vulkan/runtime}/vulkan_amdrgp.cc | 0
.../vulkan/runtime}/vulkan_amdrgp.h | 0
.../vulkan/runtime}/vulkan_buffer.cc | 0
.../vulkan/runtime}/vulkan_buffer.h | 0
.../vulkan/runtime}/vulkan_common.cc | 0
.../vulkan/runtime}/vulkan_common.h | 0
.../vulkan/runtime}/vulkan_device.cc | 2 +-
.../vulkan/runtime}/vulkan_device.h | 0
.../vulkan/runtime}/vulkan_device_api.cc | 0
.../vulkan/runtime}/vulkan_device_api.h | 2 +-
.../vulkan/runtime}/vulkan_instance.cc | 2 +-
.../vulkan/runtime}/vulkan_instance.h | 0
.../vulkan/runtime}/vulkan_module.cc | 8 +-
.../vulkan/runtime}/vulkan_stream.cc | 2 +-
.../vulkan/runtime}/vulkan_stream.h | 0
.../vulkan/runtime}/vulkan_wrapped_func.cc | 4 +-
.../vulkan/runtime}/vulkan_wrapped_func.h | 6 +-
.../webgpu/codegen}/codegen_webgpu.cc | 20 +-
.../webgpu/codegen}/codegen_webgpu.h | 2 +-
.../webgpu/codegen}/intrin_rule_webgpu.cc | 20 +-
src/backend/webgpu/codegen/register.cc | 80 +
.../webgpu/codegen}/webgpu_fallback_module.cc | 2 +-
.../webgpu/codegen}/webgpu_fallback_module.h | 4 +-
src/relax/transform/static_plan_block_memory.cc | 2 +-
src/runtime/extra/contrib/clml/clml_runtime.h | 2 +-
.../extra/contrib/cublas/cublas_json_runtime.cc | 2 +-
src/runtime/extra/contrib/cublas/cublas_utils.cc | 2 +-
.../contrib/cudnn/cudnn_frontend/attention.cc | 2 +-
src/runtime/extra/contrib/cudnn/cudnn_utils.h | 2 +-
src/runtime/extra/contrib/curand/curand.cc | 2 +-
.../cutlass/fp16_group_gemm_runner_sm100.cuh | 2 +-
.../cutlass/fp16_group_gemm_runner_sm90.cuh | 2 +-
.../fp8_groupwise_scaled_gemm_runner_sm100.cuh | 2 +-
.../fp8_groupwise_scaled_gemm_runner_sm90.cuh | 2 +-
...p8_groupwise_scaled_group_gemm_runner_sm100.cuh | 2 +-
src/runtime/extra/contrib/cutlass/gemm_runner.cuh | 2 +-
.../extra/contrib/hipblas/hipblas_json_runtime.cc | 2 +-
src/runtime/extra/contrib/hipblas/hipblas_utils.cc | 2 +-
src/runtime/extra/contrib/nvshmem/dist_gemm.cu | 2 +-
src/runtime/extra/contrib/nvshmem/init.cc | 2 +-
.../extra/contrib/nvshmem/memory_allocator.cc | 2 +-
.../extra/contrib/tensorrt/tensorrt_calibrator.h | 2 +-
src/runtime/extra/contrib/thrust/thrust.cu | 2 +-
.../extra/disco/cuda_ipc/cuda_ipc_memory.cc | 2 +-
src/runtime/extra/disco/nccl/nccl_context.h | 4 +-
src/runtime/vm/attn_utils.h | 2 +-
src/runtime/vm/hexagon/builtin.cc | 2 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 2 +-
src/s_tir/backend/adreno/texture_flatten.cc | 2 +-
.../postproc/rewrite_cooperative_fetch.cc | 7 +-
src/s_tir/schedule/analysis/analysis.cc | 5 +-
src/s_tir/transform/inject_permuted_layout.cc | 29 +-
src/s_tir/transform/inject_ptx_async_copy.cc | 10 +-
src/s_tir/transform/inject_ptx_ldg32.cc | 4 +-
src/s_tir/transform/inject_software_pipeline.cc | 24 +-
src/s_tir/transform/memhammer_lower_auto_copy.cc | 6 +-
.../transform/memhammer_tensorcore_rewrite.cc | 7 +-
.../transform/merge_shared_memory_allocations.cc | 4 +-
src/s_tir/transform/tensorcore_infer_fragment.cc | 13 +-
src/target/llvm/intrin_rule_nvptx.cc | 3 +-
src/target/tag.cc | 13 -
src/target/target_kind.cc | 318 --
src/tirx/analysis/filter_canonical.cc | 5 +-
src/tirx/ir/data_type_rewriter.cc | 9 +-
src/tirx/ir/exec_scope.cc | 4 +-
src/tirx/ir/stmt.cc | 5 +-
src/tirx/op/builtin.cc | 153 -
src/tirx/op/op.cc | 47 +-
src/tirx/script/builder/frame.cc | 3 +-
src/tirx/transform/lower_intrin.cc | 4 +-
src/tirx/transform/lower_tvm_builtin.cc | 24 +-
src/tirx/transform/lower_warp_memory.cc | 51 +-
src/tirx/transform/remove_no_op.cc | 11 +-
src/tirx/transform/tile_primitive_dispatch.cc | 5 +-
.../s_tir/dlight/test_gpu_matmul_tensorize.py | 20 +-
tests/python/tirx-base/test_tir_op_types.py | 15 +-
.../tile_primitive/cuda/copy/test_fallback.py | 4 +-
.../tile_primitive/cuda/copy/test_gmem_smem.py | 2 +-
.../tile_primitive/cuda/copy/test_swizzle_iter.py | 10 +-
.../tile_primitive/cuda/copy_async/test_dsmem.py | 2 +-
.../cuda/copy_async/test_smem_tmem.py | 2 +-
.../tile_primitive/cuda/copy_async/test_tma.py | 8 +-
.../tile_primitive/cuda/elementwise/test_unary.py | 10 +-
.../cuda/gemm_async/test_gemm_async.py | 92 +-
.../cuda/permute_layout/test_permute_layout.py | 4 +-
.../tile_primitive/trn/test_compose_op_trn.py | 14 +-
.../operator/tile_primitive/trn/test_copy_trn.py | 8 +-
.../operator/tile_primitive/trn/test_gemm_trn.py | 4 +-
.../tile_primitive/trn/test_private_alloc_trn.py | 2 +-
.../tile_primitive/trn/test_reduction_trn.py | 4 +-
.../operator/tile_primitive/trn/test_unary_trn.py | 2 +-
tests/python/tirx/test_alloc_pool.py | 6 +-
tests/python/tirx/test_layout.py | 10 +-
tests/python/tirx/test_op_namespace_cleanup.py | 121 +
tests/python/tirx/test_printer_tir_namespaces.py | 198 +-
.../transform/test_transform_naive_allocator.py | 2 +-
376 files changed, 4536 insertions(+), 14207 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6e25f10e7f..ad99c4c6ac 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -339,13 +339,19 @@ tvm_file_glob(GLOB CODEGEN_SRCS
src/target/source/*.cc
src/target/canonicalizer/*.cc
src/target/canonicalizer/llvm/*.cc
- src/target/cuda/*.cc
- src/target/rocm/*.cc
- src/target/hexagon/*.cc
- src/target/metal/*.cc
- src/target/opencl/*.cc
- src/target/vulkan/vulkan_fallback_module.cc
- src/target/webgpu/*.cc
+ src/backend/cuda/codegen/*.cc
+ src/backend/cuda/op/*.cc
+ src/backend/hexagon/codegen/*.cc
+ src/backend/metal/codegen/*.cc
+ src/backend/metal/op/*.cc
+ src/backend/opencl/codegen/*.cc
+ src/backend/rocm/codegen/*.cc
+ src/backend/trn/codegen/*.cc
+ src/backend/trn/op/*.cc
+ src/backend/trn/transform/*.cc
+ src/backend/vulkan/codegen/register.cc
+ src/backend/vulkan/codegen/vulkan_fallback_module.cc
+ src/backend/webgpu/codegen/*.cc
)
list(APPEND COMPILER_SRCS ${CODEGEN_SRCS})
@@ -361,7 +367,7 @@ tvm_file_glob(GLOB RUNTIME_SRCS
src/runtime/rpc/minrpc/*.cc
)
# Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra.
-# Note: src/runtime/{cuda,vulkan,opencl,metal,rocm,hexagon}/* move to per-backend DSOs.
+# Note: src/backend/*/runtime sources move to per-backend DSOs.
set(TVM_RUNTIME_EXT_OBJS "")
if(BUILD_FOR_HEXAGON)
diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake
index 0028e04fcc..53f4bf8d86 100644
--- a/cmake/modules/CUDA.cmake
+++ b/cmake/modules/CUDA.cmake
@@ -62,8 +62,8 @@ endif(USE_CUDA)
if(USE_CUDA)
message(STATUS "Build cuda device runtime")
- tvm_file_glob(GLOB RUNTIME_CUDA_SRCS src/runtime/cuda/*.cc)
- tvm_file_glob(GLOB VM_CUDA_BUILTIN_SRC_CC src/runtime/vm/cuda/*.cc)
+ tvm_file_glob(GLOB RUNTIME_CUDA_SRCS src/backend/cuda/runtime/*.cc)
+ tvm_file_glob(GLOB VM_CUDA_BUILTIN_SRC_CC src/backend/cuda/runtime/vm/*.cc)
add_library(tvm_runtime_cuda_objs OBJECT ${RUNTIME_CUDA_SRCS} ${VM_CUDA_BUILTIN_SRC_CC})
target_link_libraries(tvm_runtime_cuda_objs PUBLIC tvm_ffi_header)
diff --git a/cmake/modules/Hexagon.cmake b/cmake/modules/Hexagon.cmake
index c92fc70799..512a893c73 100644
--- a/cmake/modules/Hexagon.cmake
+++ b/cmake/modules/Hexagon.cmake
@@ -50,7 +50,8 @@ macro(file_glob_append _output_list)
set(${_output_list} ${_tmp1})
endmacro()
-set(TVMRT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src/runtime")
+set(TVMRT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src/backend/hexagon/runtime")
+set(TVM_CORE_RUNTIME_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src/runtime")
if(DEFINED USE_HEXAGON_DEVICE)
message(WARNING "USE_HEXAGON_DEVICE is deprecated, use USE_HEXAGON instead")
@@ -79,7 +80,7 @@ endif()
if(NOT USE_HEXAGON)
# USE_HEXAGON=OFF: codegen still works through the per-backend fallback
- # module (src/target/hexagon/hexagon_fallback_module.cc), which is always
+ # module (src/backend/hexagon/codegen/hexagon_fallback_module.cc), which is always
# compiled into libtvm via CODEGEN_SRCS. No opt-stub registration is
# needed.
return()
@@ -122,14 +123,14 @@ if(BUILD_FOR_HEXAGON)
# When building FOR Hexagon (the DSP itself), all runtime sources go into
# the single libtvm_runtime (static or shared). No per-backend DSO split.
file_glob_append(RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/*.cc"
+ "${TVMRT_SOURCE_DIR}/*.cc"
)
# Add builtins to RelaxVM
tvm_file_glob(GLOB VM_BUILTIN_SRC_CC src/runtime/vm/hexagon/*.cc)
list(APPEND RUNTIME_SRCS ${VM_BUILTIN_SRC_CC})
else()
file_glob_append(RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/hexagon_module.cc"
+ "${TVMRT_SOURCE_DIR}/hexagon_module.cc"
)
endif()
@@ -137,7 +138,7 @@ set(htp_supported_archs "v68" "v69" "v73" "v75")
list(FIND htp_supported_archs "${USE_HEXAGON_ARCH}" supported_arch_index)
if(${supported_arch_index} EQUAL -1)
# Exclude User DMA files when building for archs below v68
- list(REMOVE_ITEM RUNTIME_HEXAGON_SRCS "${TVMRT_SOURCE_DIR}/hexagon/hexagon_user_dma.cc")
+ list(REMOVE_ITEM RUNTIME_HEXAGON_SRCS "${TVMRT_SOURCE_DIR}/hexagon_user_dma.cc")
endif()
if(BUILD_FOR_HEXAGON)
@@ -160,7 +161,7 @@ if(BUILD_FOR_HEXAGON)
# QHL support.
if(USE_HEXAGON_QHL)
file_glob_append(TVM_QHL_WRAPPER_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/qhl/*.cc"
+ "${TVMRT_SOURCE_DIR}/qhl/*.cc"
)
include_directories(
@@ -181,20 +182,20 @@ if(BUILD_FOR_HEXAGON)
if(${supported_arch_index} GREATER -1)
# Hand-written ops
file_glob_append(RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/ops/*.cc"
+ "${TVMRT_SOURCE_DIR}/ops/*.cc"
)
include_directories(
- "${TVMRT_SOURCE_DIR}/hexagon/ops"
+ "${TVMRT_SOURCE_DIR}/ops"
)
set_source_files_properties(
- "${TVMRT_SOURCE_DIR}/hexagon/ops/conv2d_quant_hvx.cc"
+ "${TVMRT_SOURCE_DIR}/ops/conv2d_quant_hvx.cc"
PROPERTIES COMPILE_FLAGS "-mhvx"
)
set_source_files_properties(
- "${TVMRT_SOURCE_DIR}/hexagon/ops/conv2d_fp16_hvx.cc"
+ "${TVMRT_SOURCE_DIR}/ops/conv2d_fp16_hvx.cc"
PROPERTIES COMPILE_FLAGS "-mhvx"
)
endif()
@@ -249,21 +250,21 @@ if(USE_HEXAGON_RPC)
add_custom_command(
OUTPUT
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.h"
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_skel.c"
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc.h"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc_skel.c"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc_stub.c"
COMMAND
${QAIC_EXE_PATH} ${QAIC_FLAGS}
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.idl"
- -o "${TVMRT_SOURCE_DIR}/hexagon/rpc"
- MAIN_DEPENDENCY "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.idl"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc.idl"
+ -o "${TVMRT_SOURCE_DIR}/rpc"
+ MAIN_DEPENDENCY "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc.idl"
)
if("${CMAKE_C_COMPILER_ID}" STREQUAL "GNU" OR "${CMAKE_C_COMPILER_ID}" STREQUAL "Clang")
# We can't easily fix this at the source-code level, because the .c file is generated
# by the qaic program. But it should be safe to ignore the warning:
# https://stackoverflow.com/questions/13905200/is-it-wise-to-ignore-gcc-clangs-wmissing-braces-warning
- set_source_files_properties("${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c"
+ set_source_files_properties("${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc_stub.c"
PROPERTY COMPILE_FLAGS "-Wno-missing-braces")
endif()
endfunction()
@@ -273,12 +274,12 @@ if(USE_HEXAGON_RPC)
add_android_paths()
build_rpc_idl()
file_glob_append(RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/android/*.cc"
+ "${TVMRT_SOURCE_DIR}/rpc/android/*.cc"
)
# Add this file separately, because it's auto-generated, and glob won't
# find it during cmake-time.
list(APPEND RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc_stub.c"
)
list(APPEND TVM_RUNTIME_LINKER_LIBS cdsprpc)
@@ -289,30 +290,30 @@ if(USE_HEXAGON_RPC)
# Include the generic RPC code into the TVM runtime.
list(APPEND RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/rpc/minrpc/minrpc_server.h"
- "${TVMRT_SOURCE_DIR}/rpc/minrpc/rpc_reference.h"
- "${TVMRT_SOURCE_DIR}/rpc/rpc_module.cc"
- "${TVMRT_SOURCE_DIR}/rpc/rpc_endpoint.cc"
- "${TVMRT_SOURCE_DIR}/rpc/rpc_session.cc"
+ "${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/minrpc/minrpc_server.h"
+ "${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/minrpc/rpc_reference.h"
+ "${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/rpc_module.cc"
+ "${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/rpc_endpoint.cc"
+ "${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/rpc_session.cc"
# TODO(masahi): Remove rpc_local_session.cc after verifying that things work without it
- "${TVMRT_SOURCE_DIR}/rpc/rpc_local_session.cc"
+ "${TVM_CORE_RUNTIME_SOURCE_DIR}/rpc/rpc_local_session.cc"
)
- set(HEXAGON_PROFILER_DIR "${TVMRT_SOURCE_DIR}/hexagon/profiler")
+ set(HEXAGON_PROFILER_DIR "${TVMRT_SOURCE_DIR}/profiler")
# Add the hardware-specific RPC code into the skel library.
set_property(SOURCE ${HEXAGON_PROFILER_DIR}/lwp_handler.S PROPERTY LANGUAGE C)
add_library(hexagon_rpc_skel SHARED
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon/rpc_server.cc"
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_skel.c"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon/rpc_server.cc"
+ "${TVMRT_SOURCE_DIR}/rpc/hexagon_rpc_skel.c"
"${HEXAGON_PROFILER_DIR}/prof_utils.cc"
"${HEXAGON_PROFILER_DIR}/lwp_handler.S"
)
target_include_directories(hexagon_rpc_skel
- SYSTEM PRIVATE "${TVMRT_SOURCE_DIR}/hexagon/rpc"
+ SYSTEM PRIVATE "${TVMRT_SOURCE_DIR}/rpc"
)
# Add the simulator-specific RPC code into a shared library to be
# executed via run_main_on_sim.
add_library(hexagon_rpc_sim SHARED
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/simulator/rpc_server.cc"
+ "${TVMRT_SOURCE_DIR}/rpc/simulator/rpc_server.cc"
"${HEXAGON_PROFILER_DIR}/prof_utils.cc"
"${HEXAGON_PROFILER_DIR}/lwp_handler.S"
)
@@ -324,7 +325,7 @@ if(USE_HEXAGON_RPC)
find_hexagon_toolchain()
add_hexagon_wrapper_paths()
file_glob_append(RUNTIME_HEXAGON_SRCS
- "${TVMRT_SOURCE_DIR}/hexagon/rpc/simulator/session.cc"
+ "${TVMRT_SOURCE_DIR}/rpc/simulator/session.cc"
)
list(APPEND TVM_RUNTIME_LINKER_LIBS "-lwrapper")
endif()
diff --git a/cmake/modules/LLVM.cmake b/cmake/modules/LLVM.cmake
index f944b41304..8ad8bfc53c 100644
--- a/cmake/modules/LLVM.cmake
+++ b/cmake/modules/LLVM.cmake
@@ -44,9 +44,9 @@ if(NOT ${USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
add_definitions(-DTVM_LLVM_HAS_AARCH64_TARGET=${TVM_LLVM_HAS_AARCH64_TARGET})
tvm_file_glob(GLOB COMPILER_LLVM_SRCS
src/target/llvm/*.cc
- src/target/cuda/llvm/*.cc
- src/target/rocm/llvm/*.cc
- src/target/hexagon/llvm/*.cc
+ src/backend/cuda/codegen/llvm/*.cc
+ src/backend/rocm/codegen/llvm/*.cc
+ src/backend/hexagon/codegen/llvm/*.cc
)
list(APPEND TVM_LINKER_LIBS ${LLVM_LIBS})
list(APPEND COMPILER_SRCS ${COMPILER_LLVM_SRCS})
diff --git a/cmake/modules/Metal.cmake b/cmake/modules/Metal.cmake
index c593d0d420..02526b4c9c 100644
--- a/cmake/modules/Metal.cmake
+++ b/cmake/modules/Metal.cmake
@@ -19,7 +19,7 @@ if(USE_METAL)
message(STATUS "Build metal device runtime")
find_library(METAL_LIB Metal)
find_library(FOUNDATION_LIB Foundation)
- tvm_file_glob(GLOB RUNTIME_METAL_SRCS src/runtime/metal/*.mm)
+ tvm_file_glob(GLOB RUNTIME_METAL_SRCS src/backend/metal/runtime/*.mm)
add_library(tvm_runtime_metal_objs OBJECT ${RUNTIME_METAL_SRCS})
target_link_libraries(tvm_runtime_metal_objs PUBLIC tvm_ffi_header)
@@ -33,5 +33,5 @@ if(USE_METAL)
tvm_configure_target_library(tvm_runtime_metal RUNTIME_MODULE)
endif(USE_METAL)
# When USE_METAL=OFF the codegen-side fallback in
-# src/target/metal/metal_fallback_module.cc handles construction; no opt
+# src/backend/metal/codegen/metal_fallback_module.cc handles construction; no opt
# stub is needed (it is always compiled via CODEGEN_SRCS in CMakeLists.txt).
diff --git a/cmake/modules/OpenCL.cmake b/cmake/modules/OpenCL.cmake
index f833832d4c..392457da7a 100644
--- a/cmake/modules/OpenCL.cmake
+++ b/cmake/modules/OpenCL.cmake
@@ -16,13 +16,13 @@
# under the License.
if(USE_OPENCL)
- tvm_file_glob(GLOB RUNTIME_OPENCL_SRCS src/runtime/opencl/*.cc)
+ tvm_file_glob(GLOB RUNTIME_OPENCL_SRCS src/backend/opencl/runtime/*.cc)
set(_opencl_libs "")
if(${USE_OPENCL} MATCHES ${IS_TRUE_PATTERN})
message(STATUS "Enabled runtime search for OpenCL library location")
file_glob_append(RUNTIME_OPENCL_SRCS
- "src/runtime/opencl/opencl_wrapper/opencl_wrapper.cc"
+ "src/backend/opencl/runtime/opencl_wrapper/opencl_wrapper.cc"
)
include_directories(SYSTEM "3rdparty/OpenCL-Headers")
else()
diff --git a/cmake/modules/ROCM.cmake b/cmake/modules/ROCM.cmake
index a2d1516558..ce9ad1414f 100644
--- a/cmake/modules/ROCM.cmake
+++ b/cmake/modules/ROCM.cmake
@@ -32,7 +32,7 @@ if(USE_ROCM)
endif()
message(STATUS "Build rocm device runtime")
- tvm_file_glob(GLOB RUNTIME_ROCM_SRCS src/runtime/rocm/*.cc)
+ tvm_file_glob(GLOB RUNTIME_ROCM_SRCS src/backend/rocm/runtime/*.cc)
set(_rocm_libs ${ROCM_HIPHCC_LIBRARY})
if(ROCM_HSA_LIBRARY)
diff --git a/cmake/modules/Vulkan.cmake b/cmake/modules/Vulkan.cmake
index ba51e4b842..2f34b63240 100644
--- a/cmake/modules/Vulkan.cmake
+++ b/cmake/modules/Vulkan.cmake
@@ -29,12 +29,12 @@ if(USE_VULKAN)
include_directories(SYSTEM ${Vulkan_INCLUDE_DIRS})
message(STATUS "Build with Vulkan support")
tvm_file_glob(GLOB COMPILER_VULKAN_SRCS
- src/target/vulkan/build_vulkan.cc
- src/target/vulkan/codegen_spirv.cc
- src/target/vulkan/intrin_rule_spirv.cc
- src/target/vulkan/ir_builder.cc
- src/target/vulkan/spirv_support.cc
- src/target/vulkan/spirv_utils.cc
+ src/backend/vulkan/codegen/build_vulkan.cc
+ src/backend/vulkan/codegen/codegen_spirv.cc
+ src/backend/vulkan/codegen/intrin_rule_spirv.cc
+ src/backend/vulkan/codegen/ir_builder.cc
+ src/backend/vulkan/codegen/spirv_support.cc
+ src/backend/vulkan/codegen/spirv_utils.cc
)
list(APPEND COMPILER_SRCS ${COMPILER_VULKAN_SRCS})
list(APPEND TVM_LINKER_LIBS ${Vulkan_SPIRV_TOOLS_LIBRARY})
@@ -44,7 +44,7 @@ endif(USE_VULKAN)
if(USE_VULKAN)
message(STATUS "Build vulkan device runtime")
- tvm_file_glob(GLOB RUNTIME_VULKAN_SRCS src/runtime/vulkan/*.cc)
+ tvm_file_glob(GLOB RUNTIME_VULKAN_SRCS src/backend/vulkan/runtime/*.cc)
add_library(tvm_runtime_vulkan_objs OBJECT ${RUNTIME_VULKAN_SRCS})
target_link_libraries(tvm_runtime_vulkan_objs PUBLIC tvm_ffi_header)
diff --git a/docs/reference/api/python/tirx/backend.rst b/docs/reference/api/python/tirx/backend.rst
index c4e6689bfe..35d6433802 100644
--- a/docs/reference/api/python/tirx/backend.rst
+++ b/docs/reference/api/python/tirx/backend.rst
@@ -23,9 +23,3 @@ tvm.tirx.backend
.. automodule:: tvm.tirx.backend
:members:
:imported-members:
-
-tvm.tirx.backend.adreno
-***********************
-.. automodule:: tvm.tirx.backend.adreno
- :members:
- :imported-members:
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index f25e48daf3..953ee8daa1 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -560,49 +560,6 @@ TVM_DLL const Op& tvm_global_barrier_kinit();
* }
*/
TVM_DLL const Op& tvm_thread_allreduce();
-
-// Metal SimdGroup matrix intrinsics
-
-/*!
- * \brief tvm intrinsic for initializing and simdgroup with given value.
- * \note only 8x8 shape is supported by Metal Spec and TVM, but we still keep shape as params,
- * keeping the similar interface with Metal Spec.
- *
- * void make_filled_simdgroup_matrix(Var d, PrimExpr index, PrimExpr value,
- * int col = 8, int row = 8);
- */
-TVM_DLL const Op& make_filled_simdgroup_matrix();
-
-/*!
- * \brief tvm intrinsic for loading data from device memory or threadgroup memory to simdgroup.
- * \note only 8x8 shape is supported by Metal Spec and TVM, but we still keep shape as params,
- * keeping the similar interface with Metal Spec.
- *
- * void simdgroup_load(Var d, PrimExpr index, PrimExpr ptr, PrimExpr stride,
- int col = 8, int row = 8, bool transpose_matrix = false);
- */
-TVM_DLL const Op& simdgroup_load();
-
-/*!
- * \brief tvm intrinsic for storing data from simdgroup to device memory or threadgroup memory.
- * \note only 8x8 shape is supported by Metal Spec and TVM, but we still keep shape as params,
- * keeping the similar interface with Metal Spec.
- *
- * void simdgroup_store(Var d, PrimExpr index, PrimExpr ptr, PrimExpr stride,
- * int col = 8, int row = 8, bool transpose_matrix = false);
- */
-TVM_DLL const Op& simdgroup_store();
-
-/*!
- * \brief tvm intrinsic for multiply and accumulate two matrices in simdgroup
- * \note only 8x8 shape is supported by Metal Spec and TVM, but we still keep shape as params,
- * keeping the similar interface with Metal Spec.
- *
- * void simdgroup_mma(Var d, PrimExpr index_d, Var a, PrimExpr index_a,
- * Var b, PrimExpr index_b, Var c, PrimExpr index_c);
- */
-TVM_DLL const Op& simdgroup_multiply_accumulate();
-
// Metal cooperative_tensor intrinsics (MetalPerformancePrimitives / Metal 4)
/*!
@@ -856,233 +813,6 @@ enum TVMStructFieldKind : int {
* \brief Print the content of a buffer during runtime.
*/
TVM_DLL const Op& print_buffer();
-
-/*!
- * \brief tvm intrinsic for initializing the CUDA profiler, and store profiling result in a buffer.
- *
- * void timer_init_cuda(Var profiler_buffer, Var profiler_tag, Var profiler_write_offset, int
- * num_groups, Expr group_id) {
- * // initialize the tag and write to pos 0 in the buffer
- * // initialize write offset for every leader thread in warp group across all blocks
- * }
- */
-TVM_DLL const Op& timer_init_cuda();
-
-/*!
- * \brief tvm intrinsic for starting the timer for profiling a specific event,
- * and storing profiling result in a buffer.
- *
- * void timer_start_cuda(IntImm event_type, Var profiler_buffer, Var profiler_tag,
- * Var profiler_write_offset, IntImm profiler_write_stride, Expr leader_cond)
- * {
- * // each leader thread in warp group gets the time stamp and event type, combine with the tag
- * // and write to corresponding offset in buffer
- * // each leader thread advance offset by stride
- * }
- */
-TVM_DLL const Op& timer_start_cuda();
-
-/*!
- * \brief tvm intrinsic for ending the timer for profiling a specific event,
- * and storing profiling result in a buffer.
- *
- * void timer_end_cuda(IntImm event_type, Var profiler_buffer, Var profiler_tag,
- * Var profiler_write_offset, IntImm profiler_write_stride, Expr leader_cond) {
- * // each leader thread in warp group gets the time stamp and event type, combine with the tag
- * // and write to corresponding offset in buffer
- * // each leader thread advance offset by stride
- * }
- */
-TVM_DLL const Op& timer_end_cuda();
-
-/*!
- * \brief tvm intrinsic for finalize the timer for profiling,
- * and storing profiling result in a buffer.
- *
- * void timer_finalize_cuda(Var profiler_buffer, Var profiler_tag, Var profiler_write_offset,
- * IntImm profiler_write_stride, Expr leader_cond) {
- * // each leader thread in warp group gets the time stamp and end signal, combine with the tag
- * // and write to corresponding offset in buffer
- * // each leader thread advance offset by stride
- * }
- */
-TVM_DLL const Op& timer_finalize_cuda();
-
-/*!
- * \brief tvm intrinsic for cuda atomic add instruction
- */
-TVM_DLL const Op& cuda_atomic_add();
-
-/*!
- * \brief tvm intrinsic for cuda thread fence instruction
- */
-TVM_DLL const Op& cuda_thread_fence();
-
-/*!
- * \brief tvm intrinsic for cuda warpgroup sync instruction
- */
-TVM_DLL const Op& cuda_warpgroup_sync();
-
-/*!
- * \brief Warp-level butterfly shuffle-XOR reduction.
- *
- * cuda_warp_reduce(value, op, width) reduces value across width adjacent
- * lanes using the specified operation ("sum", "max", "min").
- */
-TVM_DLL const Op& cuda_warp_reduce();
-
-/*!
- * \brief CTA-wide reduction via warp shuffle + shared memory.
- *
- * cuda_cta_reduce(value, op, num_warps, scratch) reduces value across
- * the entire CTA using the specified operation ("sum", "max", "min").
- */
-TVM_DLL const Op& cuda_cta_reduce();
-
-/*!
- * \brief Typed load/store copy of num_bytes bytes.
- *
- * cuda_copy_bytes(dst, src, num_bytes) copies num_bytes bytes from src to dst
- * using a single typed load/store (uint4, uint2, unsigned int, etc.).
- * num_bytes must be one of {1, 2, 4, 8, 16}.
- */
-TVM_DLL const Op& cuda_copy_bytes();
-
-/*!
- * \brief tvm intrinsic for cuda warp sync instruction
- */
-TVM_DLL const Op& cuda_warp_sync();
-
-/*!
- * \brief tvm intrinsic for cuda block-wide sync (syncthreads)
- */
-TVM_DLL const Op& cuda_cta_sync();
-
-/*!
- * \brief tvm intrinsic for cuda grid-wide sync (cooperative groups)
- */
-TVM_DLL const Op& cuda_grid_sync();
-
-/*!
- * \brief tvm intrinsic for cuda cluster-wide sync instruction
- */
-TVM_DLL const Op& cuda_cluster_sync();
-
-/*!
- * \brief tvm intrinsic that returns ``cooperative_groups::thread_rank()``
- * for the enclosing CTA (linear thread index within the block).
- */
-TVM_DLL const Op& cuda_thread_rank();
-
-/*!
- * \brief tvm intrinsic for cuda half to float conversion
- */
-TVM_DLL const Op& cuda_half2float();
-
-/*!
- * \brief tvm intrinsic for cuda bfloat16 to float conversion
- */
-TVM_DLL const Op& cuda_bfloat162float();
-
-/*!
- * \brief tvm intrinsic for a helper converting float2 to half2 with rounding
- */
-TVM_DLL const Op& cuda_float22half2();
-
-/*!
- * \brief tvm intrinsic to trap when an assertion failed (cond == false)
- */
-TVM_DLL const Op& cuda_trap_when_assert_failed();
-
-/*!
- * \brief tvm intrinsic to modify runtime instruction descriptor
- */
-TVM_DLL const Op& cuda_runtime_instr_desc();
-
-/*!
- * \brief tvm intrinsic to convert 8 half2 lanes to 8 float2 lanes
- */
-TVM_DLL const Op& cuda_half8tofloat8();
-
-/*!
- * \brief tvm intrinsic to convert 8 float2 lanes to 8 half2 lanes with rounding
- */
-TVM_DLL const Op& cuda_float8tohalf8();
-
-/*!
- * \brief tvm intrinsic for cuda syncthreads_and instruction
- */
-TVM_DLL const Op& cuda_syncthreads_and();
-
-/*!
- * \brief tvm intrinsic for cuda syncthreads_or instruction
- */
-TVM_DLL const Op& cuda_syncthreads_or();
-
-/*!
- * \brief tvm intrinsic for cuda nano sleep instruction
- */
-TVM_DLL const Op& cuda_nano_sleep();
-
-/*!
- * \brief tvm intrinsic for cuda atomic compare and swap instruction
- */
-TVM_DLL const Op& cuda_atomic_cas();
-
-/*!
- * \brief tvm intrinsic for cuda printf instruction
- */
-TVM_DLL const Op& cuda_printf();
-
-/*!
- * \brief tvm intrinsic for cuda ldg instruction
- */
-TVM_DLL const Op& cuda_ldg();
-
-/*!
- * \brief tvm intrinsic for cuda tmem address calculation
- */
-TVM_DLL const Op& cuda_get_tmem_addr();
-
-/*!
- * \brief tvm intrinsic for PTX fast exp2 approximation (ex2.approx.ftz.f32)
- */
-TVM_DLL const Op& ptx_exp2();
-
-/*!
- * \brief tvm intrinsic for PTX fast reciprocal approximation (rcp.approx.ftz.f32)
- */
-TVM_DLL const Op& ptx_rcp();
-
-/*!
- * \brief tvm intrinsic for PTX warp-wide any predicate (__any_sync)
- */
-TVM_DLL const Op& ptx_any_sync();
-
-/*!
- * \brief tvm intrinsic for PTX 3-input max instruction (sm_100a+)
- */
-TVM_DLL const Op& ptx_reduce3_max_f32();
-
-/*!
- * \brief tvm intrinsic for PTX 3-input min instruction (sm_100a+)
- */
-TVM_DLL const Op& ptx_reduce3_min_f32();
-
-TVM_DLL const Op& ptx_add_f32();
-TVM_DLL const Op& ptx_add_f32x2();
-TVM_DLL const Op& ptx_add_f64();
-TVM_DLL const Op& ptx_sub_f32();
-TVM_DLL const Op& ptx_sub_f32x2();
-TVM_DLL const Op& ptx_sub_f64();
-TVM_DLL const Op& ptx_mul_f32();
-TVM_DLL const Op& ptx_mul_f32x2();
-TVM_DLL const Op& ptx_mul_f64();
-TVM_DLL const Op& ptx_fma_f32();
-TVM_DLL const Op& ptx_fma_f32x2();
-TVM_DLL const Op& ptx_fma_f64();
-TVM_DLL const Op& ptx_max_f32();
-
} // namespace builtin
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op.h b/include/tvm/tirx/op.h
index 60b292bbb2..2027665712 100644
--- a/include/tvm/tirx/op.h
+++ b/include/tvm/tirx/op.h
@@ -35,8 +35,6 @@
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/target_builtin/cuda.h>
-#include <tvm/tirx/target_builtin/trn.h>
#include <algorithm>
#include <limits>
diff --git a/include/tvm/tirx/target_builtin/cuda.h b/include/tvm/tirx/target_builtin/cuda.h
deleted file mode 100644
index ff10ee0b43..0000000000
--- a/include/tvm/tirx/target_builtin/cuda.h
+++ /dev/null
@@ -1,745 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/tir/target_builtin/cuda.h
- * \brief TIR builtin intrinsics specific to CUDA target.
- */
-#ifndef TVM_TIRX_TARGET_BUILTIN_CUDA_H_
-#define TVM_TIRX_TARGET_BUILTIN_CUDA_H_
-
-#include <tvm/tirx/expr.h>
-#include <tvm/tirx/op.h>
-
-namespace tvm {
-namespace tirx {
-namespace builtin {
-
-// TODO(tvm-team) TensorCore specific intrinsics should be directly registered under
-// cuda. namespace and used through op.
-/*!
- * \brief tvm intrinsic for tensor core load operators.
- *
- * void tvm_load_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
- * Expr index, Expr buffer_ptr, Expr stride,
- * StringImm layout) {
- * // m, n, k are the shape of wmma fragment.
- * // Determine fragment layout(column-major or row major) by layout.
- * // fragments must be in 'wmma.matrix_a' or 'wmma.matrix_b' scope.
- * nvcuda::wmma::load_matrix_sync(fragment[index], buffer_ptr, stride);
- * }
- */
-TVM_DLL const Op& tvm_load_matrix_sync();
-
-/*!
- * \brief tvm intrinsic for tensor core mma_sync operators.
- *
- * void tvm_mma_sync(Var fragment_d, Expr index_d,
- * Var fragment_a, Expr index_a,
- * Var fragment_b, Expr index_b,
- * Var fragment_c, Expr index_c) {
- * nvcuda::wmma::mma_sync(fragment_d[index_d], fragment_a[index_a],
- * fragment_b[index_b], fragment_c[index_c]);
- * }
- */
-TVM_DLL const Op& tvm_mma_sync();
-
-/*!
- * \brief tvm intrinsic for tensor core bmma_sync operators.
- *
- * void tvm_bmma_sync(Var fragment_d, Expr index_d,
- * Var fragment_a, Expr index_a,
- * Var fragment_b, Expr index_b,
- * Var fragment_c, Expr index_c) {
- * nvcuda::wmma::bmma_sync(fragment_d[index_d], fragment_a[index_a],
- * fragment_b[index_b], fragment_c[index_c]);
- * }
- */
-TVM_DLL const Op& tvm_bmma_sync();
-
-/*!
- * \brief tvm intrinsic for tensor core fill_fragment operators.
- *
- * void tvm_fill_fragment(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
- * Expr index, Expr value) {
- * // m, n, k are the shape of wmma fragment
- * // fragments must be in 'wmma.accumulator' scope.
- * nvcuda::wmma::fill_fragment(fragment[index], value);
- * }
- */
-TVM_DLL const Op& tvm_fill_fragment();
-
-/*!
- * \brief tvm intrinsic for tensor core store operators.
- *
- * void tvm_store_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
- * Expr index, Expr buffer_ptr, Expr stride,
- * StringImm layout) {
- * // m, n, k are the shape of wmma fragment
- * // fragments must be in 'wmma.accumulator' scope.
- * nvcuda::wmma::store_matrix_sync(fragment[index], buffer_ptr, stride, layout);
- * }
- */
-TVM_DLL const Op& tvm_store_matrix_sync();
-
-/*!
- * \brief tvm intrinsic for ptx tensor core mma instructions.
- *
- * void ptx_mma(StringImm shape, StringImm A_layout, StringImm B_layout,
- * StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
- * Var multiplicand_a, Expr a_index,
- * Var multiplicand_b, Expr b_index,
- * Var accumulator, Expr c_index, bool saturate);
- */
-TVM_DLL const Op& ptx_mma();
-
-/*!
- * \brief ptx mma / ldmatrix / mma_store / mma_fill variants that take
- * ``(ptr_var, offset)`` pairs (not a folded access_ptr Call). Codegen
- * emits ``ptr + offset`` C pointer arithmetic; ``lower_warp_memory``
- * rewrites the offset's group component to its thread-local index.
- */
-TVM_DLL const Op& ptx_mma_legacy();
-TVM_DLL const Op& ptx_ldmatrix_legacy();
-TVM_DLL const Op& mma_store_legacy();
-TVM_DLL const Op& mma_fill_legacy();
-
-/*!
- * \brief tvm intrinsic for ptx predicate load with 32-bit data type.
- *
- */
-TVM_DLL const Op& ptx_ldg32();
-
-/*!
- * \brief tvm intrinsic for sparse tensor core ptx instructions.
- *
- * void ptx_mma_sp(StringImm shape, StringImm A_layout, StringImm B_layout,
- * StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
- * Var multiplicand_a, Expr a_index,
- * Var multiplicand_b, Expr b_index,
- * Var accumulator, Expr c_index,
- * Var metadata, Expr meta_index,
- * Var sparse_selector, bool saturate);
- */
-TVM_DLL const Op& ptx_mma_sp();
-
-/*!
- * \brief tvm intrinsic for ptx load matrix from shared memory.
- *
- * void ptx_ldmatrix(Bool trans, IntImm num, StringImm type,
- * Var local_ptr, Expr local_offset,
- * Var smem_ptr, Expr smem_offset);
- */
-TVM_DLL const Op& ptx_ldmatrix();
-
-/*!
- * \brief tvm intrinsics for ptx async copy from global to shared memory using cp.async
- *
- * void ptx_cp_async(Var shared_ptr,
- * Expr shared_offset,
- * Var global_ptr,
- * Expr global_offset,
- * size_t bytes);
- */
-TVM_DLL const Op& ptx_cp_async();
-
-/*!
- * \brief tvm intrinsics for ptx async copy from global to shared memory using cp.async.bulk
- *
- * void ptx_cp_async_bulk(Var shared_ptr,
- * Expr shared_offset,
- * Var global_ptr,
- * Expr global_offset,
- * size_t bytes,
- * int barrier_arr_id,
- * int barrier_id);
- */
-TVM_DLL const Op& ptx_cp_async_bulk();
-
-/*!
- * \brief tvm intrinsics for ptx async bulk copy from shared::cta to shared::cluster
- *
- * void ptx_cp_async_bulk_shared_to_cluster(Expr dst_ptr,
- * Expr src_ptr,
- * Expr size,
- * Expr mbar);
- */
-TVM_DLL const Op& ptx_cp_async_bulk_shared_to_cluster();
-
-/*!
- * \brief tvm intrinsics for ptx async copy commit and wait.
- *
- * void ptx_cp_async_commit_group();
- * void ptx_cp_async_wait_group(int num);
- *
- */
-TVM_DLL const Op& ptx_cp_async_commit_group();
-TVM_DLL const Op& ptx_cp_async_wait_group();
-
-/*!
- * \brief tvm intrinsics for ptx async copy barrier using cp.async.mbarrier.arrive
- *
- * ptx_cp_async_mbarrier_arrive(int barrier_arr_id, int barrier_id)
- *
- */
-TVM_DLL const Op& ptx_cp_async_mbarrier_arrive();
-
-/*!
- * \brief PTX fence instruction: fence.{sem}.{scope}
- *
- * ptx_fence(StringImm sem, StringImm scope)
- */
-TVM_DLL const Op& ptx_fence();
-
-/*!
- * \brief PTX fence.proxy.async instruction: fence.proxy.async[.{space}]
- *
- * ptx_fence_proxy_async(StringImm space)
- */
-TVM_DLL const Op& ptx_fence_proxy_async();
-
-/*!
- * \brief tvm instrinsics to call mbarrier.init.shared::cta.b64
- *
- * ptx_mbarrier_init(uint64_t* bar_ptr, int thread_count)
- */
-TVM_DLL const Op& ptx_mbarrier_init();
-
-/*!
- * \brief tvm instrinsics to call
- * mbarrier.arrive.shared::cta.b64
- * or
- * @p mapa.shared::cluster.u32
- * @p mbarrier.arrive.shared::cluster.b64
- */
-TVM_DLL const Op& ptx_mbarrier_arrive();
-
-/*!
- * \brief tvm instrinsics to call
- * mbarrier.arrive.expect_tx.shared.b64
- * or
- * @p mapa.shared::cluster.u32
- * @p mbarrier.arrive.expect_tx.shared.b64
- *
- * ptx_mbarrier_arrive_expect_tx(uint64_t* bar_ptr, int byte_count)
- */
-TVM_DLL const Op& ptx_mbarrier_arrive_expect_tx();
-
-/*!
- * \brief tvm instrinsics to call mbarrier.try_wait.parity repeatedly until it returns true
- *
- * ptx_mbarrier_try_wait(uint64_t* bar_ptr, int phase)
- */
-TVM_DLL const Op& ptx_mbarrier_try_wait();
-
-/*!
- * \brief tvm instrinsics to call bar.arrive a, b
- *
- * bar_arrive(int name_bar_id, int thread_count)
- */
-TVM_DLL const Op& ptx_bar_arrive();
-
-/*!
- * \brief tvm instrinsics to call bar.sync a, {b}
- *
- * bar_sync(int name_bar_id, int thread_count)
- */
-TVM_DLL const Op& ptx_bar_sync();
-
-/*!
- * \brief tvm instrinsics to call
- * cp.async.bulk.tensor.dim.shared::cluster.global.tile.mbarrier::complete_tx::bytes
- *
- * TMA alignment requirement:
- * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
- *
- * ptx_cp_async_bulk_tensor_global_to_cluster(int dim, PrimExpr dst_ptr, PrimExpr bar_ptr,
- * PrimExpr tensormap_addr, int...coords, int cta_mask, int cta_group, string cache_hint)
- */
-TVM_DLL const Op& ptx_cp_async_bulk_tensor_global_to_cluster();
-
-/*!
- * \brief tvm intrinsic to call
- * cp.async.bulk.tensor.dim.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes
- *
- * TMA alignment requirement:
- * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
- *
- * ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster(int dim, PrimExpr dst_ptr, PrimExpr
- * bar_ptr, PrimExpr tensormap_addr, int...coords, int cta_mask, int cta_group, string cache_hint)
- */
-TVM_DLL const Op& ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster();
-
-/*!
- * \brief tvm instrinsics to call
- * cp.async.bulk.tensor.dim.global.shared::cta.tile。bulk_group
- *
- * TMA alignment requirement:
- * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
- *
- * ptx_cp_async_bulk_tensor_shared_to_global(int dim, PrimExpr src_ptr, PrimExpr tensormap_addr,
- * int...coords, string cache_hint)
- */
-TVM_DLL const Op& ptx_cp_async_bulk_tensor_shared_to_global();
-
-/*!
- * \brief tvm instrinsics to call
- * cp.async.bulk.prefetch.tensor.dim.L2.global.tile
- *
- * TMA alignment requirement:
- * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
- *
- * ptx_cp_async_bulk_tensor_global_to_cluster_prefetch(int dim, PrimExpr tensormap_addr,
- * int...coords, string cache_hint)
- */
-TVM_DLL const Op& ptx_cp_async_bulk_tensor_global_to_cluster_prefetch();
-
-/*!
- * \brief tvm instrinsics to call
- * cp.reduce.async.bulk.tensor.dim.dst.src.redOp
- *
- * ptx_cp_async_bulk_tensor_shared_to_global_reduce(int dim, PrimExpr src_ptr, PrimExpr
- * tensormap_addr, int...coords, string cache_hint)
- */
-TVM_DLL const Op& ptx_cp_async_bulk_tensor_shared_to_global_reduce();
-
-/*!
- * \brief tvm instrinsics to call cp.async.bulk.commit_group
- *
- * ptx_cp_async_bulk_commit_group()
- */
-TVM_DLL const Op& ptx_cp_async_bulk_commit_group();
-
-/*!
- * \brief tvm instrinsics to call cp.async.bulk.wait_group{.read} N
- *
- * ptx_cp_async_bulk_wait_group(int N, bool read)
- */
-TVM_DLL const Op& ptx_cp_async_bulk_wait_group();
-
-/*!
- * \brief tvm instrinsics to call barrier.cluster.arrive{.sem}{.aligned}
- *
- * ptx_barrier_cluster_arrive(string sem, bool aligned)
- */
-TVM_DLL const Op& ptx_barrier_cluster_arrive();
-
-/*!
- * \brief tvm instrinsics to call barrier.cluster.wait.{acquire}{.aligned}
- *
- * ptx_barrier_cluster_wait(bool acquire, bool aligned)
- */
-TVM_DLL const Op& ptx_barrier_cluster_wait();
-
-/*!
- * \brief tvm instrinsics to call elect.sync _|p, membermask and return the predicate
- *
- * elect_sync(membermask)
- */
-TVM_DLL const Op& ptx_elect_sync();
-
-/*!
- * \brief PTX fence.mbarrier_init.release.cluster instruction
- *
- * ptx_fence_mbarrier_init()
- */
-TVM_DLL const Op& ptx_fence_mbarrier_init();
-
-/*!
- * \brief tvm instrinsics to fetch PTX pre-defined registers
- *
- * ptx_fetch_register(int bits, string reg_name)
- */
-TVM_DLL const Op& ptx_fetch_register();
-
-/*!
- * \brief PTX programmatic dependent launch synchronization.
- */
-TVM_DLL const Op& ptx_griddepcontrol_wait();
-TVM_DLL const Op& ptx_griddepcontrol_launch_dependents();
-
-/*!
- * \brief tvm intrinsic for storing the result of PTX MMA into a destination pointer.
- * For example, if each thread in a warp of size 32 has 4 elements from the result of
- * m16xn8xk16 MMA in its registers, this intrinsic can be used to store the result in a
- * 16x8 region in shared or global memory.
- *
- * There is no real PTX instruction that does that, but we want to hide details of
- * complex index manipulation behind this intrinsic to simplify TIR lowering passes (e.g.
- * LowerWarpMemory).
- *
- * void mma_store(IntImm m, IntImm n, Var dst_ptr, Var src_ptr, Expr src_offset, Var dst_stride);
- */
-TVM_DLL const Op& mma_store();
-
-/*!
- * \brief tvm intrinsic for zero-initializing an MMA accumulation register.
- * For example, if each thread in a warp of size 32 has 8 elements from the A matrix in
- * m16xn8xk16 MMA in its registers, this intrinsic can be used to zero-initialize its
- * 4 accumulation registers.
- *
- * There is no real PTX instruction that does that, but we introduce this intrinsic for the
- * same reason as mma_store above.
- *
- * void mma_fill(IntImm local_size, Var local_ptr, Expr offset);
- */
-TVM_DLL const Op& mma_fill();
-
-/*!
- * \brief tvm intrinsic to encode matrix descriptor for wgmma instructions.
- *
- * ptx_wgmma_encode_matrix_descriptor(PrimExpr ptr, PrimExpr ldo, PrimExpr sdo, int swizzle)
- */
-TVM_DLL const Op& ptx_wgmma_encode_matrix_descriptor();
-
-/*!
- * \brief tvm intrinsic to call "" : "+r"(reg) :: "memory"
- *
- * ptx_wgmma_noop_barrier()
- */
-TVM_DLL const Op& ptx_wgmma_noop_barrier();
-
-/*!
- * \brief tvm intrinsic to call wgmma.mma_async.sync.aligned.shape.dtype.atype.btype
- * where both A and B are in shared memory.
- *
- * ptx_wgmma_mma_async_ss()
- */
-TVM_DLL const Op& ptx_wgmma_mma_async_ss();
-
-/*!
- * \brief tvm intrinsic to call wgmma.mma_async.sync.aligned.shape.dtype.atype.btype
- * where A is in register and B is in shared memory.
- *
- * ptx_wgmma_mma_async_rs()
- */
-TVM_DLL const Op& ptx_wgmma_mma_async_rs();
-
-/*!
- * \brief tvm intrinsic to call wgmma.fence.sync.aligned;
- *
- * ptx_wgmma_fence()
- */
-TVM_DLL const Op& ptx_wgmma_fence();
-
-/*!
- * \brief tvm intrinsic to call wgmma.commit_group.sync.aligned;
- *
- * ptx_wgmma_commit_group()
- */
-TVM_DLL const Op& ptx_wgmma_commit_group();
-
-/*!
- * \brief tvm intrinsic to call wgmma.wait_group.sync.aligned;
- *
- * ptx_wgmma_wait_group(int N)
- */
-TVM_DLL const Op& ptx_wgmma_wait_group();
-
-/*!
- * \brief tvm intrinsic to call stmatrix.sync.aligned.m8n8.num{.trans}.shared.b16 [p], r;
- *
- * ptx_stmatrix(int num, bool trans, PrimExpr ptr, PrimExpr... vars)
- */
-TVM_DLL const Op& ptx_stmatrix();
-
-/*!
- * \brief tvm intrinsic to call setmaxnreg.action.sync.aligned.u32 imm-reg-count
- */
-TVM_DLL const Op& ptx_setmaxnreg();
-
-/*!
- * \brief tvm intrinsic to call ld.global.acquire.gpu.b32
- *
- * ptx_ld_global_acquire()
- */
-TVM_DLL const Op& ptx_ld_global_acquire();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.alloc.cta_group.sync.aligned;
- *
- * ptx_tcgen05_alloc(Var dst_ptr, int n_cols, int cta_group)
- */
-TVM_DLL const Op& ptx_tcgen05_alloc();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.dealloc.cta_group.sync.aligned;
- *
- * ptx_tcgen05_dealloc(uint32_t taddr, int n_cols, int cta_group)
- */
-TVM_DLL const Op& ptx_tcgen05_dealloc();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.relinquish_alloc_permit.cta_group.sync.aligned;
- *
- * ptx_tcgen05_relinquish_alloc_permit(int cta_group)
- */
-TVM_DLL const Op& ptx_tcgen05_relinquish_alloc_permit();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.fence::before_thread_sync;
- *
- * ptx_tcgen05_fence_before_thread_sync()
- */
-TVM_DLL const Op& ptx_tcgen05_fence_before_thread_sync();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.fence::after_thread_sync;
- *
- * ptx_tcgen05_fence_after_thread_sync()
- */
-TVM_DLL const Op& ptx_tcgen05_fence_after_thread_sync();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.ld.sync.aligned;
- *
- * ptx_tcgen05_ld()
- */
-TVM_DLL const Op& ptx_tcgen05_ld();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.st.sync.aligned;
- *
- * ptx_tcgen05_st()
- */
-TVM_DLL const Op& ptx_tcgen05_st();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.wait::ld.sync.aligned;
- *
- * ptx_tcgen05_wait_ld()
- */
-TVM_DLL const Op& ptx_tcgen05_wait_ld();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.wait::st.sync.aligned;
- *
- * ptx_tcgen05_wait_st()
- */
-TVM_DLL const Op& ptx_tcgen05_wait_st();
-
-/*!
- * \brief tvm intrinsic to encode matrix descriptor for tcgen05 instructions.
- *
- * ptx_tcgen05_encode_matrix_descriptor(PrimExpr ptr, PrimExpr ldo, PrimExpr sdo, int swizzle)
- */
-TVM_DLL const Op& ptx_tcgen05_encode_matrix_descriptor();
-
-/*!
- * \brief tvm intrinsic to encode instruction descriptor for tcgen05 MMA.
- *
- * ptx_tcgen05_encode_instr_descriptor(PrimExpr desc, string d_dtype, string a_dtype, string
- * b_dtype, int M, int N, int K, bool trans_a, bool trans_b, int n_cta_groups, bool neg_a, bool
- * neg_b, bool sat_d, bool is_sparse)
- */
-TVM_DLL const Op& ptx_tcgen05_encode_instr_descriptor();
-
-/*!
- * \brief tvm intrinsic to encode instruction descriptor for tcgen05 MMA block scaled.
- *
- * ptx_tcgen05_encode_instr_descriptor_block_scaled(PrimExpr desc, string d_dtype,
- * string a_dtype, string b_dtype, string sfa_dtype, string stb_dtype,
- * int M, int N, int K, bool trans_a, bool trans_b,
- * int n_cta_groups, bool neg_a, bool neg_b, bool is_sparse)
- */
-TVM_DLL const Op& ptx_tcgen05_encode_instr_descriptor_block_scaled();
-
-/*!
- * \brief tvm intrinsic to call tcgen05.mma.cta_group.kind without block scaling.
- *
- * ptx_tcgen05_mma()
- */
-TVM_DLL const Op& ptx_tcgen05_mma();
-
-/*!
- * \brief tvm intrinsic to call tcgen05.mma.cta_group.kind.block_scale{.scale_vec_size}
- *
- * ptx_tcgen05_mma_block_scale()
- */
-TVM_DLL const Op& ptx_tcgen05_mma_block_scale();
-
-/*!
- * \brief tvm intrinsic to call tcgen05.mma.sp.cta_group.kind without block scaling.
- *
- * ptx_tcgen05_mma_sp()
- */
-TVM_DLL const Op& ptx_tcgen05_mma_sp();
-
-/*!
- * \brief tvm intrinsic to call tcgen05.mma.sp.cta_group.kind.block_scale{.scale_vec_size}
- *
- * ptx_tcgen05_mma_sp_block_scale()
- */
-TVM_DLL const Op& ptx_tcgen05_mma_sp_block_scale();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.commit.cta_group
- *
- * ptx_tcgen05_commit()
- */
-TVM_DLL const Op& ptx_tcgen05_commit();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.cp.cta_group
- *
- * ptx_tcgen05_cp()
- */
-TVM_DLL const Op& ptx_tcgen05_cp();
-
-/*!
- * \brief tvm instrinsics to call tcgen05.shift.cta_group.down
- *
- * ptx_tcgen05_shift()
- */
-TVM_DLL const Op& ptx_tcgen05_shift();
-
-/*!
- * \brief tvm instrinsics to call map_shared_rank
- *
- * ptx_map_shared_rank(PrimExpr ptr, int rank)
- */
-TVM_DLL const Op& ptx_map_shared_rank();
-
-/*!
- * \brief tvm instrinsics to call a CUDA function. Source code is provided as a string.
- *
- * cuda_func_call(String func_name, PrimExpr... args, String source_code)
- */
-TVM_DLL const Op& cuda_func_call();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_my_pe() operation.
- *
- * int nvshmem_my_pe()
- */
-TVM_DLL const Op& nvshmem_my_pe();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_n_pes() operation.
- *
- * int nvshmem_n_pes()
- */
-TVM_DLL const Op& nvshmem_n_pes();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_getmem_nbi() operation.
- *
- * void nvshmem_getmem_nbi(void *dest, const void *source, size_t nelems, int pe)
- */
-TVM_DLL const Op& nvshmem_getmem_nbi();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_putmem_nbi() operation.
- *
- * void nvshmem_putmem_nbi(void *dest, const void *source, size_t nelems, int pe)
- */
-TVM_DLL const Op& nvshmem_putmem_nbi();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_getmem_nbi_warp() operation.
- *
- * void nvshmemx_getmem_nbi_warp(void *dest, const void *source, size_t nelems, int pe)
- */
-TVM_DLL const Op& nvshmem_getmem_nbi_warp();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_putmem_nbi_warp() operation.
- *
- * void nvshmemx_putmem_nbi_warp(void *dest, const void *source, size_t nelems, int pe)
- */
-TVM_DLL const Op& nvshmem_putmem_nbi_warp();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_getmem_nbi_block() operation.
- *
- * void nvshmemx_getmem_nbi_block(void *dest, const void *source, size_t nelems, int pe)
- */
-TVM_DLL const Op& nvshmem_getmem_nbi_block();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_putmem_nbi_block() operation.
- *
- * void nvshmemx_putmem_nbi_block(void *dest, const void *source, size_t nelems, int pe)
- */
-TVM_DLL const Op& nvshmem_putmem_nbi_block();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_signal_op() operation.
- *
- * void nvshmemx_signal_op(uint64_t *sig_addr, uint64_t signal, int sig_op, int pe)
- */
-TVM_DLL const Op& nvshmem_signal_op();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_FuncParam{TYPENAME}_wait_until() operation.
- *
- * void nvshmem_FuncParam{TYPENAME}_wait_until(TYPE *ivar, int cmp, TYPE cmp_value)
- */
-TVM_DLL const Op& nvshmem_wait_until();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_quiet() operation.
- *
- * void nvshmem_quiet()
- */
-TVM_DLL const Op& nvshmem_quiet();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_putmem_signal_nbi() operation.
- *
- * void nvshmemx_putmem_signal_nbi(void *dest, const void *source, size_t nelems, uint64_t
- * *sig_addr, uint64_t signal, int sig_op, int pe)
- */
-TVM_DLL const Op& nvshmem_putmem_signal_nbi();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_putmem_signal_nbi_warp() operation.
- *
- * void nvshmemx_putmem_signal_nbi_warp(void *dest, const void *source, size_t nelems, uint64_t
- * *sig_addr, uint64_t signal, int sig_op, int pe)
- */
-TVM_DLL const Op& nvshmem_putmem_signal_nbi_warp();
-
-/*!
- * \brief nvshmem intrinsics for nvshmemx_putmem_signal_nbi_block() operation.
- *
- * void nvshmemx_putmem_signal_nbi_block(void *dest, const void *source, size_t nelems,
- * uint64_t *sig_addr, uint64_t signal, int sig_op, int pe)
- */
-TVM_DLL const Op& nvshmem_putmem_signal_nbi_block();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_fence() operation.
- *
- * void nvshmem_fence()
- */
-TVM_DLL const Op& nvshmem_fence();
-
-/*!
- * \brief nvshmem intrinsics for nvshmem_barrier_all() operation.
- *
- * void nvshmem_barrier_all()
- */
-TVM_DLL const Op& nvshmem_barrier_all();
-
-} // namespace builtin
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIRX_TARGET_BUILTIN_CUDA_H_
diff --git a/include/tvm/tirx/target_builtin/trn.h b/include/tvm/tirx/target_builtin/trn.h
deleted file mode 100644
index 556156bc13..0000000000
--- a/include/tvm/tirx/target_builtin/trn.h
+++ /dev/null
@@ -1,156 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/tir/target_builtin/trn.h
- * \brief TIR builtin intrinsics specific to Trainium target.
- */
-#ifndef TVM_TIRX_TARGET_BUILTIN_TRN_H_
-#define TVM_TIRX_TARGET_BUILTIN_TRN_H_
-
-#include <tvm/tirx/expr.h>
-#include <tvm/tirx/op.h>
-
-namespace tvm {
-namespace tirx {
-namespace builtin {
-
-/*!
- * \brief nki intrinsics for load operation.
- *
- * nki_load(result, data)
- */
-TVM_DLL const Op& nki_load();
-/*!
- * \brief nki intrinsics for store operation.
- *
- * nki_store(result, data)
- */
-TVM_DLL const Op& nki_store();
-/*!
- * \brief nki intrinsics for tensor_copy operation.
- *
- * nki_tensor_copy(result, data)
- */
-TVM_DLL const Op& nki_tensor_copy();
-/*!
- * \brief nki intrinsics for matmul operation.
- *
- * nki_matmul(C, A, B, accum)
- *
- * equivalent to C += A.T @ B (if accum is true), or C = A.T @ B (if accum is false)
- */
-TVM_DLL const Op& nki_matmul();
-
-/*!
- * \brief nki intrinsics for activation operation.
- *
- * nki_activation(result, data, opcode, bias, scale)
- */
-TVM_DLL const Op& nki_activation();
-
-/*!
- * \brief nki intrinsics for reciprocal operation.
- *
- * nki_reciprocal(result, data)
- */
-TVM_DLL const Op& nki_reciprocal();
-
-/*!
- * \brief nki intrinsics for tensortensor operation.
- *
- * nki_tensortensor(result, operand0, operand1, opcode)
- */
-TVM_DLL const Op& nki_tensortensor();
-
-/*!
- * \brief nki intrinsics for tensorscalar operation.
- *
- * nki_tensorscalar(result, operand0, operand1, opcode, reverse)
- */
-TVM_DLL const Op& nki_tensorscalar();
-
-/*!
- * \brief nki intrinsics for tensorreduce operation.
- *
- * nki_tensorreduce(result, data, opcode, negate, axes)
- */
-TVM_DLL const Op& nki_tensorreduce();
-
-/*!
- * \brief nki intrinsics for memset operation.
- *
- * nki_memset(result, value)
- */
-TVM_DLL const Op& nki_memset();
-
-/*!
- * \brief nki intrinsics for activation reduce operation.
- *
- * nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias, scale)
- */
-TVM_DLL const Op& nki_activation_reduce();
-
-/*!
- * \brief nki intrinsics for tensorscalar reduce operation.
- *
- * nki_tensorscalar_reduce(reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode,
- * reverse)
- */
-TVM_DLL const Op& nki_tensorscalar_reduce();
-
-/*!
- * \brief nki intrinsics for initializing identity tensor.
- *
- * nki_identity(result, size)
- */
-TVM_DLL const Op& nki_identity();
-
-/*!
- * \brief nki intrinsics for scalar tensor tensor operation.
- *
- * (data op1 operand1) op2 (operand2) where op1 is tensor-scalar and op2 is tensor-tensor
- *
- * nki_scalar_tensor_tensor(result, data, operand0, operand1, opcode0, opcode1, reverse0, reverse1)
- *
- */
-TVM_DLL const Op& nki_scalar_tensor_tensor();
-
-/*!
- * \brief nki intrinsics for scalar tensor scalar operation.
- *
- * (data op1 operand1) op2 (operand2) where op1 and op2 are tensor-scalar
- *
- * nki_scalar_tensor_scalar(result, data, operand0, operand1, opcode0, opcode1, reverse0, reverse1)
- *
- */
-TVM_DLL const Op& nki_scalar_tensor_scalar();
-
-/*!
- * \brief nki intrinsics for affine_select operation.
- *
- * nki_affine_select(result, pred, true_value, false_value)
- */
-TVM_DLL const Op& nki_affine_select();
-
-} // namespace builtin
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIRX_TARGET_BUILTIN_TRN_H_
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index bb171170fa..b38e86a0ae 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -51,6 +51,9 @@ from . import script
# tvm.tirx — registers itself via tvm.script.register_dialect in its __init__
from . import tirx
+# tvm.backend — owns backend Python load hooks
+from . import backend
+
# tvm.target
from . import target
@@ -72,10 +75,7 @@ from .support import rocm as _rocm, nvcc as _nvcc
# Relax contain modules that are only available in compiler package
# Do not import them if TVM is built with runtime only
if not _RUNTIME_ONLY:
- # tile_primitive imports both Python Op class declarations (Zero, Add, ...)
- # and per-target dispatch schedule registrations. Must run before relax so
- # any relax pass that looks up a schedule sees them.
- from .tirx.operator import tile_primitive
+ backend.load_all()
# tvm.relax — registers itself via tvm.script.register_dialect in its __init__
from . import relax
diff --git a/python/tvm/backend/__init__.py b/python/tvm/backend/__init__.py
new file mode 100644
index 0000000000..2243f5e034
--- /dev/null
+++ b/python/tvm/backend/__init__.py
@@ -0,0 +1,209 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Backend-owned Python modules and load hooks."""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import sys
+import types
+from pkgutil import extend_path
+from typing import Any
+
+__path__ = extend_path(__path__, __name__) # type: ignore[name-defined]
+
+_BUILTIN_BACKENDS = (
+ "cuda",
+ "metal",
+ "rocm",
+ "trn",
+ "opencl",
+ "vulkan",
+ "webgpu",
+ "hexagon",
+ "adreno",
+)
+_LOADED_BACKENDS: dict[str, Any] = {}
+
+
+class _AliasModule(types.ModuleType):
+ """Module object that exposes a backend module under a public alias."""
+
+ def __init__(self, fullname: str, module):
+ super().__init__(fullname, getattr(module, "__doc__", None))
+ self.__dict__["__tvm_backend_module__"] = module
+ self.__dict__["__package__"] = fullname.rpartition(".")[0]
+ if hasattr(module, "__all__"):
+ self.__dict__["__all__"] = module.__all__
+ if hasattr(module, "__path__"):
+ self.__dict__["__path__"] = []
+
+ def __getattr__(self, name: str):
+ return getattr(self.__dict__["__tvm_backend_module__"], name)
+
+ def __setattr__(self, name: str, value):
+ setattr(self.__dict__["__tvm_backend_module__"], name, value)
+
+ def __delattr__(self, name: str):
+ delattr(self.__dict__["__tvm_backend_module__"], name)
+
+ def __dir__(self):
+ return sorted(set(super().__dir__()) | set(dir(self.__dict__["__tvm_backend_module__"])))
+
+
+class _AliasLoader:
+ """Loader that returns an already-resolved module for an alias spec."""
+
+ def __init__(self, fullname: str, module):
+ self._fullname = fullname
+ self._module = module
+
+ def create_module(self, spec):
+ return _get_alias_module(self._fullname, self._module)
+
+ def exec_module(self, module):
+ _set_module_alias(self._fullname, self._module)
+ return None
+
+ def is_package(self, fullname):
+ return hasattr(self._module, "__path__")
+
+
+def _redirect_tirx_backend_alias(fullname: str) -> str | None:
+ prefix = "tvm.tirx."
+ if not fullname.startswith(prefix):
+ return None
+ rest = fullname[len(prefix) :]
+ backend_name, sep, tail = rest.partition(".")
+ if not sep or backend_name not in _LOADED_BACKENDS:
+ return None
+ return f"tvm.backend.{backend_name}.{tail}"
+
+
+class _BackendAliasFinder:
+ """Redirect ``tvm.tirx.<backend>.*`` imports to ``tvm.backend.<backend>.*``."""
+
+ @classmethod
+ def find_spec(cls, fullname, path, target=None):
+ redirected = _redirect_tirx_backend_alias(fullname)
+ if redirected is None:
+ return None
+ module = importlib.import_module(redirected)
+ _set_module_alias(fullname, module)
+ loader = _AliasLoader(fullname, module)
+ spec = importlib.util.spec_from_loader(
+ fullname, loader, is_package=hasattr(module, "__path__")
+ )
+ if spec is not None and hasattr(module, "__path__"):
+ spec.submodule_search_locations = []
+ return spec
+
+
+if not any(isinstance(finder, _BackendAliasFinder) for finder in sys.meta_path):
+ sys.meta_path.insert(0, _BackendAliasFinder())
+
+
+def _get_alias_module(alias: str, module):
+ existing = sys.modules.get(alias)
+ if (
+ isinstance(existing, _AliasModule)
+ and existing.__dict__.get("__tvm_backend_module__") is module
+ ):
+ return existing
+ return _AliasModule(alias, module)
+
+
+def _set_module_alias(alias: str, module, *, direct: bool = False) -> None:
+ alias_module = module if direct else _get_alias_module(alias, module)
+ sys.modules[alias] = alias_module
+ parent_name, _, child_name = alias.rpartition(".")
+ parent = sys.modules.get(parent_name)
+ if parent is not None:
+ setattr(parent, child_name, alias_module)
+
+
+def _alias_loaded_backend_modules(name: str) -> None:
+ backend_prefix = f"tvm.backend.{name}"
+ public_prefix = f"tvm.tirx.{name}"
+ for module_name, module in sorted(list(sys.modules.items())):
+ if module_name == backend_prefix or module_name.startswith(f"{backend_prefix}."):
+ public_name = f"{public_prefix}{module_name[len(backend_prefix) :]}"
+ _set_module_alias(public_name, module, direct=module_name == backend_prefix)
+
+
+def _import_backend(name: str):
+ module_name = f"tvm.backend.{name}"
+ try:
+ return importlib.import_module(module_name)
+ except ModuleNotFoundError as err:
+ if err.name == module_name:
+ raise ImportError(
+ f"Cannot load TVM backend {name!r}: expected Python package {module_name!r}. "
+ "Install the backend package or check the backend name."
+ ) from err
+ raise
+
+
+def load(name: str) -> None:
+ """Load a backend's Python registration hooks.
+
+ Loading is idempotent. A backend package must live at ``tvm.backend.<name>``
+ and expose ``register_backend()``.
+ """
+
+ if name in _LOADED_BACKENDS:
+ return None
+
+ module = _import_backend(name)
+ register_backend = getattr(module, "register_backend", None)
+ if register_backend is None:
+ raise AttributeError(f"Backend package 'tvm.backend.{name}' has no register_backend()")
+
+ import tvm.tirx as tirx # pylint: disable=import-outside-toplevel
+
+ setattr(tirx, name, module)
+ sys.modules[f"tvm.tirx.{name}"] = module
+ _LOADED_BACKENDS[name] = module
+ try:
+ register_backend()
+ _alias_loaded_backend_modules(name)
+ except Exception:
+ _LOADED_BACKENDS.pop(name, None)
+ if getattr(tirx, name, None) is module:
+ delattr(tirx, name)
+ if sys.modules.get(f"tvm.tirx.{name}") is module:
+ sys.modules.pop(f"tvm.tirx.{name}", None)
+ raise
+ return None
+
+
+def load_all() -> None:
+ """Load all in-tree backend Python hooks."""
+
+ for name in _BUILTIN_BACKENDS:
+ load(name)
+ return None
+
+
+def is_loaded(name: str) -> bool:
+ """Return whether a backend has been loaded."""
+
+ return name in _LOADED_BACKENDS
+
+
+__all__ = ["is_loaded", "load", "load_all"]
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py b/python/tvm/backend/adreno/__init__.py
similarity index 59%
copy from python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py
copy to python/tvm/backend/adreno/__init__.py
index f7a7b886d0..8f79af80d1 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py
+++ b/python/tvm/backend/adreno/__init__.py
@@ -14,20 +14,26 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+"""Adreno-owned backend hooks."""
-"""Reduction dispatch variant registrations."""
+from importlib import import_module
-from tvm.tirx.operator.tile_primitive import register_dispatch
+_LAZY_SUBMODULES = {"target_tags"}
-from ...common import ReduceOpType
-from .utils import reduction_trn
-for _op_name, _op_type in {
- "sum": ReduceOpType.SUM,
- "max": ReduceOpType.MAX,
- "min": ReduceOpType.MIN,
-}.items():
+def register_backend():
+ """Register Adreno-owned Python semantics."""
+ import tvm.backend as backend # pylint: disable=import-outside-toplevel
- @register_dispatch(_op_name, "trn", variant="reduction", priority=0)
- def _reduction_dispatch(op, sctx, _ty=_op_type):
- return reduction_trn(op, _ty, sctx)
+ backend.load("opencl")
+ backend.load("vulkan")
+ import_module(f"{__name__}.target_tags")
+
+
+def __getattr__(name: str):
+ if name in _LAZY_SUBMODULES:
+ return import_module(f"{__name__}.{name}")
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["register_backend", "target_tags"]
diff --git a/python/tvm/target/tag_registry/adreno.py b/python/tvm/backend/adreno/target_tags.py
similarity index 97%
rename from python/tvm/target/tag_registry/adreno.py
rename to python/tvm/backend/adreno/target_tags.py
index 388253c3bd..daba9acf9d 100644
--- a/python/tvm/target/tag_registry/adreno.py
+++ b/python/tvm/backend/adreno/target_tags.py
@@ -16,7 +16,7 @@
# under the License.
"""Qualcomm Adreno GPU target tags."""
-from .registry import register_tag
+from tvm.target import register_tag
register_tag(
"qcom/adreno-opencl",
diff --git a/python/tvm/backend/cuda/__init__.py b/python/tvm/backend/cuda/__init__.py
new file mode 100644
index 0000000000..34592d2de3
--- /dev/null
+++ b/python/tvm/backend/cuda/__init__.py
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA-owned TIRx modules."""
+
+from importlib import import_module
+
+_LAZY_SUBMODULES = {"lang", "op", "operator", "script", "target_tags"}
+
+
+def register_backend():
+ """Register CUDA-owned Python semantics."""
+ from tvm.tirx.script.builder import ir as builder_ir # pylint: disable=import-outside-toplevel
+
+ for name, namespace in script_namespaces().items():
+ builder_ir.register_script_namespace(name, namespace)
+
+ import_module(f"{__name__}.operator.intrinsics")
+ import_module(f"{__name__}.operator.tile_primitive")
+ import_module(f"{__name__}.target_tags")
+
+
+def script_namespaces(**_):
+ """Return CUDA-owned TVMScript namespaces."""
+ from .script import ( # pylint: disable=import-outside-toplevel
+ CUDANamespace,
+ NVSHMEMNamespace,
+ PTXNamespace,
+ )
+
+ return {
+ "cuda": CUDANamespace(),
+ "nvshmem": NVSHMEMNamespace(),
+ "ptx": PTXNamespace(),
+ }
+
+
+def script_namespace(**kwargs):
+ """Return the CUDA TVMScript namespace object."""
+ return script_namespaces(**kwargs)["cuda"]
+
+
+def __getattr__(name: str):
+ if name in _LAZY_SUBMODULES:
+ return import_module(f"{__name__}.{name}")
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+ "lang",
+ "op",
+ "operator",
+ "register_backend",
+ "script",
+ "script_namespace",
+ "script_namespaces",
+ "target_tags",
+]
diff --git a/python/tvm/backend/cuda/lang/__init__.py b/python/tvm/backend/cuda/lang/__init__.py
new file mode 100644
index 0000000000..a70ce50196
--- /dev/null
+++ b/python/tvm/backend/cuda/lang/__init__.py
@@ -0,0 +1,70 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA-specific TIRx language helpers."""
+
+from importlib import import_module
+
+__all__ = [
+ "BaseTileScheduler",
+ "ClusterPersistentScheduler2D",
+ "FlashAttentionLPTScheduler",
+ "FlashAttentionLinearScheduler",
+ "GroupMajor3D",
+ "IndexedTripleTileScheduler",
+ "MBarrier",
+ "Pipeline",
+ "PipelineState",
+ "RankAwareGroupMajorTileScheduler",
+ "SMEMPool",
+ "SmemDescriptor",
+ "TCGen05Bar",
+ "TMABar",
+ "TMEMPool",
+ "TMEMStages",
+ "WarpRole",
+ "WarpgroupRole",
+]
+
+_HELPER_MODULES = {
+ "MBarrier": ".pipeline",
+ "Pipeline": ".pipeline",
+ "PipelineState": ".pipeline",
+ "BaseTileScheduler": ".tile_scheduler",
+ "ClusterPersistentScheduler2D": ".tile_scheduler",
+ "FlashAttentionLPTScheduler": ".tile_scheduler",
+ "FlashAttentionLinearScheduler": ".tile_scheduler",
+ "GroupMajor3D": ".tile_scheduler",
+ "IndexedTripleTileScheduler": ".tile_scheduler",
+ "RankAwareGroupMajorTileScheduler": ".tile_scheduler",
+ "SMEMPool": ".alloc_pool",
+ "TCGen05Bar": ".pipeline",
+ "TMABar": ".pipeline",
+ "TMEMPool": ".alloc_pool",
+ "TMEMStages": ".alloc_pool",
+ "WarpRole": ".warp_role",
+ "WarpgroupRole": ".warp_role",
+ "SmemDescriptor": ".smem_desc",
+}
+
+
+def __getattr__(name: str):
+ module_name = _HELPER_MODULES.get(name)
+ if module_name is None:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+ value = getattr(import_module(module_name, __name__), name)
+ globals()[name] = value
+ return value
diff --git a/python/tvm/tirx/lang/alloc_pool.py b/python/tvm/backend/cuda/lang/alloc_pool.py
similarity index 97%
copy from python/tvm/tirx/lang/alloc_pool.py
copy to python/tvm/backend/cuda/lang/alloc_pool.py
index 48bb9929c6..725264a35b 100644
--- a/python/tvm/tirx/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -68,7 +68,7 @@ def _shape_product(shape):
def _auto_swizzle_mode(dtype):
"""Select the default MMA swizzle mode for a shared-memory allocation."""
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import SwizzleMode
+ from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
del dtype
return SwizzleMode.SWIZZLE_128B_ATOM
@@ -76,7 +76,7 @@ def _auto_swizzle_mode(dtype):
def _swizzle_atom_bytes(swizzle_mode):
"""Return the row width (in bytes) of one swizzle atom for *swizzle_mode*."""
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import SwizzleMode
+ from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
return {
SwizzleMode.SWIZZLE_NONE: 0,
@@ -111,7 +111,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
"""
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import SwizzleMode
+ from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
return
@@ -354,7 +354,7 @@ class TMEMPool:
what gemm dispatch iterates over). When ``shape`` has 3 dims, the first
is treated as a pipe-depth outer.
"""
- from tvm.tirx.operator.tile_primitive.cuda.gemm_async.tcgen05 import sf_tmem_layout
+ from tvm.backend.cuda.operator.tile_primitive.gemm_async.tcgen05 import sf_tmem_layout
if len(shape) == 2:
pipe_depth, rows, last = None, shape[0], shape[1]
@@ -477,7 +477,7 @@ class SMEMPool:
def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024):
"""Allocate MMA-compatible shared memory with an inferred swizzle layout."""
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import (
+ from tvm.backend.cuda.operator.tile_primitive.tma_utils import (
SwizzleMode,
mma_shared_layout,
)
diff --git a/python/tvm/tirx/lang/pipeline.py b/python/tvm/backend/cuda/lang/pipeline.py
similarity index 100%
copy from python/tvm/tirx/lang/pipeline.py
copy to python/tvm/backend/cuda/lang/pipeline.py
diff --git a/python/tvm/tirx/lang/smem_desc.py b/python/tvm/backend/cuda/lang/smem_desc.py
similarity index 95%
copy from python/tvm/tirx/lang/smem_desc.py
copy to python/tvm/backend/cuda/lang/smem_desc.py
index c858cb7069..d2561d17dc 100644
--- a/python/tvm/tirx/lang/smem_desc.py
+++ b/python/tvm/backend/cuda/lang/smem_desc.py
@@ -17,8 +17,8 @@
"""SMEM matrix descriptor helper for tcgen05 / wgmma."""
+from tvm.backend.cuda.operator.tile_primitive.common import smem_desc_add_16B_offset
from tvm.script import tirx as T
-from tvm.tirx.operator.tile_primitive.cuda.common import smem_desc_add_16B_offset
@T.meta_class
diff --git a/python/tvm/tirx/lang/tile_scheduler.py b/python/tvm/backend/cuda/lang/tile_scheduler.py
similarity index 100%
copy from python/tvm/tirx/lang/tile_scheduler.py
copy to python/tvm/backend/cuda/lang/tile_scheduler.py
diff --git a/python/tvm/tirx/lang/warp_role.py b/python/tvm/backend/cuda/lang/warp_role.py
similarity index 100%
copy from python/tvm/tirx/lang/warp_role.py
copy to python/tvm/backend/cuda/lang/warp_role.py
diff --git a/python/tvm/tirx/op.py b/python/tvm/backend/cuda/op.py
similarity index 56%
copy from python/tvm/tirx/op.py
copy to python/tvm/backend/cuda/op.py
index 4d80ac378e..e76d5fbe24 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -14,3150 +14,39 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=redefined-builtin, invalid-name, too-many-arguments
-"""Operators used in TIR expression."""
+# pylint: disable=invalid-name, too-many-arguments
+"""CUDA, PTX, and NVSHMEM TIR intrinsic builders."""
-from typing import Any
-
-import tvm_ffi
-from tvm_ffi import Array
-
-import tvm
-from tvm import tirx
-from tvm.ir import Op, PointerType, PrimExpr
-from tvm.ir.base import Span
-from tvm.ir.type import TensorMapType
-from tvm.runtime import const
-
-from . import _ffi_api
-from .buffer import Buffer
-from .expr import BufferLoad, Call, CommReducer, IntImm, PrimExprWithOp, Var
-
-# Choice / IntAttr value tables — single source of truth in
-# tvm.tirx.operator.intrinsics._common. Re-exported here under their
-# underscored names so the existing _choice(name, value, _FOO) call sites
-# below keep working without changes.
-from .operator.intrinsics._common import CLUSTER_BARRIER_SEM as _CLUSTER_BARRIER_SEM
-from .operator.intrinsics._common import CP_ASYNC_BULK_CACHE_HINT as _CP_ASYNC_BULK_CACHE_HINT
-from .operator.intrinsics._common import CP_ASYNC_BULK_RED_OP as _CP_ASYNC_BULK_RED_OP
-from .operator.intrinsics._common import CP_ASYNC_CACHE_HINT as _CP_ASYNC_CACHE_HINT
-from .operator.intrinsics._common import CP_ASYNC_FILL_MODE as _CP_ASYNC_FILL_MODE
-from .operator.intrinsics._common import CP_ASYNC_PREFETCH_SIZE as _CP_ASYNC_PREFETCH_SIZE
-from .operator.intrinsics._common import F32X2_ROUND as _F32X2_ROUND
-from .operator.intrinsics._common import FENCE_PROXY_ASYNC_SPACE as _FENCE_PROXY_ASYNC_SPACE
-from .operator.intrinsics._common import FENCE_SCOPE as _FENCE_SCOPE
-from .operator.intrinsics._common import FENCE_SEM as _FENCE_SEM
-from .operator.intrinsics._common import LDMATRIX_DTYPE as _LDMATRIX_DTYPE
-from .operator.intrinsics._common import LDMATRIX_NUM as _LDMATRIX_NUM
-from .operator.intrinsics._common import NVSHMEM_CMP as _NVSHMEM_CMP
-from .operator.intrinsics._common import NVSHMEM_SIG_OP as _NVSHMEM_SIG_OP
-from .operator.intrinsics._common import TCGEN05_CP_DECOMPRESS as _TCGEN05_CP_DECOMPRESS
-from .operator.intrinsics._common import TCGEN05_CP_MULTICAST as _TCGEN05_CP_MULTICAST
-from .operator.intrinsics._common import TCGEN05_CP_SHAPES as _TCGEN05_CP_SHAPES
-from .operator.intrinsics._common import TCGEN05_CTA_GROUP as _TCGEN05_CTA_GROUP
-from .operator.intrinsics._common import TCGEN05_LDST_SHAPES as _TCGEN05_LDST_SHAPES
-
-tir = tirx # alias for backward compat with upstream tir.convert() calls
-
-_DEVICE_INTRIN_PREFIX_TO_NAMESPACE = {
- "cuda_": "cuda",
- "ptx_": "ptx",
- "nvshmem_": "nvshmem",
- "nki_": "nki",
-}
-
-
-def _canonical_device_intrin_name(func_name: str) -> str:
- """Return the canonical registry name for statically registered device intrinsics."""
-
- if not isinstance(func_name, str) or not func_name.startswith("tirx."):
- return func_name
- basename = func_name[len("tirx.") :]
- if "." in basename:
- return func_name
- for prefix, namespace in _DEVICE_INTRIN_PREFIX_TO_NAMESPACE.items():
- if basename.startswith(prefix):
- return f"tirx.{namespace}.{basename[len(prefix) :]}"
- return func_name
-
-
-def _pack_buffer(buf, span=None):
- """Build intrinsics that packs the buffer."""
- shape = Call("handle", "tirx.tvm_stack_make_shape", buf.shape, span=span)
- strides = (
- Call("handle", "tirx.tvm_stack_make_shape", buf.strides, span=span) if buf.strides else 0
- )
- pack_args = [
- buf.data,
- shape,
- strides,
- len(buf.shape),
- const(0, dtype=buf.dtype),
- buf.elem_offset,
- ]
- return Call("handle", Op.get("tirx.tvm_stack_make_array"), pack_args, span=span)
-
-
-def call_packed_lowered(*args, span=None):
- """Lowered version of call packed.
- The argument to packed function can be Expr or Buffer.
- The argument is the corresponding POD type when Expr is presented.
- When the argument is Buffer, the corresponding PackedFunc
- will receive an TVMArrayHandle whose content is valid during the callback period.
- If the PackedFunc is a python callback, then the corresponding argument is Tensor.
-
- Parameters
- ----------
- args : list of Expr or Buffer.
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
-
- See Also
- --------
- te.extern : Create tensor with extern function call.
- """
- call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
- return Call("int32", Op.get("tirx.tvm_call_packed_lowered"), call_args, span=span)
-
-
-def call_cpacked_lowered(*args, span=None):
- """Lowered version of call c-packed.
- Same as call_packed, except that the first argument is the function name
- (as in call_extern), and the last argument is the resource handle.
-
- Parameters
- ----------
- args : list of Expr or Buffer.
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
-
- See Also
- --------
- te.extern : Create tensor with extern function call.
- """
- call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
- return Call("int32", Op.get("tirx.tvm_call_cpacked_lowered"), call_args, span=span)
-
-
-def call_packed(*args, span=None):
- """Build expression by call an external packed function.
-
- The argument to packed function can be Expr or Buffer.
- The argument is the corresponding POD type when Expr is presented.
-
- When the argument is Buffer, the corresponding PackedFunc
- will receive an TVMArrayHandle whose content is valid during the callback period.
- If the PackedFunc is a python callback, then the corresponding argument is Tensor.
-
- Parameters
- ----------
- args : list of Expr or Buffer.
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
-
- See Also
- --------
- te.extern : Create tensor with extern function call.
- """
- call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
- return Call("int32", Op.get("tirx.tvm_call_packed"), call_args, span=span)
-
-
-def call_cpacked(*args, span=None):
- """Build expression by call an external packed function.
-
- Same as call_packed, except that the first argument is the function name
- (as in call_extern), and the last argument is the resource handle.
-
- Parameters
- ----------
- args : list of Expr or Buffer.
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
-
- See Also
- --------
- te.extern : Create tensor with extern function call.
- """
- call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
- return Call("int32", Op.get("tirx.tvm_call_cpacked"), call_args, span=span)
-
-
-def call_intrin(dtype, func_name, *args, attrs=None, span=None):
- """Build expression by calling an intrinsic function.
-
- Intrinsics can be overloaded with multiple data types via
- the intrinsic translation rule.
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- func_name: str
- The intrinsic function name.
-
- args : list
- Positional arguments.
-
- attrs : Optional[tvm.ir.Attrs or Dict[str, Object]]
- Additional attributes for the call.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- if isinstance(func_name, str):
- func_name = _canonical_device_intrin_name(func_name)
- return Call(dtype, func_name, args, attrs=attrs, span=span)
-
-
-def call_pure_extern(dtype, func_name, *args, span=None):
- """Build expression by calling a pure extern function.
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- func_name: str
- The extern function name.
-
- args : list
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return Call(dtype, Op.get("tirx.call_pure_extern"), [func_name, *args], span=span)
-
-
-def call_extern(dtype, func_name, *args, span=None):
- """Build expression by calling a extern function.
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- func_name: str
- The extern function name.
-
- args : list
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return Call(dtype, Op.get("tirx.call_extern"), [func_name, *args], span=span)
-
-
-def _require_float_arg(op_name, x):
- x = tirx.convert(x)
- if "float" not in x.dtype and "bfloat" not in x.dtype:
- raise TypeError(f"tirx.{op_name} only supports floating-point inputs, but got {x.dtype}")
- return x
-
-
-def call_llvm_intrin(dtype, name, *args, span=None):
- """Build expression by calling a llvm intrinsic function
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- name : str
- The name of the llvm intrinsic function.
-
- args : list
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- # pylint: disable=import-outside-toplevel
- from tvm.target import codegen
-
- if isinstance(name, str):
- llvm_id = codegen.llvm_lookup_intrinsic_id(name)
- elif isinstance(name, IntImm):
- llvm_id = name.value
- else:
- llvm_id = name
- if llvm_id == 0:
- raise ValueError(f"Unknown llvm intrinsic function {name}")
- return call_intrin(
- dtype,
- Op.get("tirx.call_llvm_intrin"),
- tvm.tirx.const(llvm_id, "uint32"),
- *args,
- span=span,
- )
-
-
-def call_llvm_pure_intrin(dtype, name, *args, span=None):
- """Build expression by calling a pure llvm intrinsic function
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- name : str
- The name of the llvm intrinsic function.
-
- args : list
- Positional arguments.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- # pylint: disable=import-outside-toplevel
- from tvm.target import codegen
-
- if isinstance(name, str):
- llvm_id = codegen.llvm_lookup_intrinsic_id(name)
- elif isinstance(name, IntImm):
- llvm_id = name.value
- else:
- llvm_id = name
- if llvm_id == 0:
- raise ValueError(f"Unknown llvm intrinsic function {name}")
- return call_intrin(
- dtype,
- Op.get("tirx.call_llvm_pure_intrin"),
- tvm.tirx.const(llvm_id, "uint32"),
- *args,
- span=span,
- )
-
-
-def tvm_stack_alloca(dtype_str, num):
- """Return new on stack dtype[num]
-
- Parameters
- ----------
- dtype_str : str
- The data type of array.
-
- num : int
- The size of array.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.tvm_stack_alloca", dtype_str, num)
-
-
-def tvm_stack_make_shape(*args):
- """Allocate a shape tuple on stack, return the handle
-
- Parameters
- ----------
- args : int
- The tuple shape.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.tvm_stack_make_shape", *args)
-
-
-def tvm_stack_make_array(data, shape, strides, ndim, arr_dtype, elem_offset):
- """Allocate a Tensor(DLTensor) on stack, return the handle
-
- Parameters
- ----------
- data : Expr
- The data of array.
-
- shape : Expr
- The shape of array.
-
- strides : Expr
- The strides of array.
-
- ndim : Expr
- The dimensions of array.
-
- arr_dtype : Expr
- The data type of array.
-
- elem_offse : Expr
- The element offset of array.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.tvm_stack_make_array",
- data,
- shape,
- strides,
- ndim,
- arr_dtype,
- elem_offset,
- )
-
-
-def assume(cond=None):
- """Provide a true statement that can be used for simplifications
-
- Parameters
- ----------
- cond : Expr
- The constraint condition.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("bool", "tirx.assume", cond)
-
-
-def undef():
- """Returns an initialized but arbitrary value
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("int32", "tirx.undef")
-
-
-def call_tir(global_var: tvm.ir.GlobalVar, *args):
- """Performs a call into another PrimFunc in the same IRModule
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- assert isinstance(global_var, tvm.ir.GlobalVar)
-
- dtype = "void"
- if global_var.struct_info is not None:
- ret_sinfo = global_var.struct_info.ret
- if hasattr(ret_sinfo, "dtype"):
- dtype = ret_sinfo.dtype
-
- return Call(dtype=dtype, op=global_var, args=args)
-
-
-def start_profile_intrinsic(id):
- """Start profile intrinsic.
- Parameters
- ----------
- id : int
- The intrinsic id.
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.start_profile_intrinsic", id)
-
-
-def end_profile_intrinsic(id):
- """End profile intrinsic.
- Parameters
- ----------
- id : int
- The intrinsic id.
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.end_profile_intrinsic", id)
-
-
-def tvm_tuple(*value):
- """Create a tuple structure in value field of AttrStmt
-
- Parameters
- ----------
- value : Expr
- The value in tuple.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.tvm_tuple", *value)
-
-
-def handle_add_byte_offset(handle, offset):
- """Add offset to handle
-
- Parameters
- ----------
- handle : Expr
- The handle.
-
- offset : int
- The offset.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.handle_add_byte_offset", handle, offset)
-
-
-def tvm_struct_get(arr, index, field, dtype):
- """Get struct field value in array
-
- Parameters
- ----------
- dtype : str
- The date type of the result.
-
- arr : StructType*
- The array of struct.
-
- index : int
- The index of struct.
-
- field : int
- The field of struct.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.tvm_struct_get", arr, index, field)
-
-
-def tvm_struct_set(arr, index, field, value):
- """Set value in struct field in array
-
- Parameters
- ----------
- arr : StructType*
- The array of struct.
-
- index : int
- The index of struct.
-
- field : int
- The field of struct.
-
- value : Expr
- The value to be set in field.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("int32", "tirx.tvm_struct_set", arr, index, field, value)
-
-
-def _is_tensormap_var(obj: Var) -> bool:
- type_annotation = obj.type_annotation
- return isinstance(type_annotation, PointerType) and isinstance(
- type_annotation.element_type, TensorMapType
- )
-
-
-def address_of(obj: Buffer | BufferLoad | Var, span: Span | None = None) -> PrimExpr:
- """Returns the address of a buffer element or addressable variable.
-
- Parameters
- ----------
- obj: Union[Buffer, BufferLoad, Var]
- The buffer, buffer load, or addressable variable.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- if isinstance(obj, Buffer):
- n_dim = len(obj.shape)
- buffer_load = BufferLoad(obj, [0] * n_dim)
- return call_intrin("handle", "tirx.address_of", buffer_load, span=span)
- elif isinstance(obj, Var):
- dtype = "uint64" if _is_tensormap_var(obj) else "handle"
- return call_intrin(dtype, "tirx.address_of", obj, span=span)
- elif isinstance(obj, BufferLoad):
- return call_intrin("handle", "tirx.address_of", obj, span=span)
- else:
- raise ValueError(f"Invalid object type: {type(obj)}")
-
-
-def lookup_param(param_name, span=None):
- """Returns the param by name
-
- Parameters
- ----------
- param_name : str
- The name of param.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.lookup_param", param_name, span=span)
-
-
-def tvm_thread_allreduce(*freduce_args):
- """Perform allreduce inside threadblock.
-
- Parameters
- ----------
- freduce_args : Expr
- The args.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.tvm_thread_allreduce", *freduce_args)
-
-
-def tvm_thread_invariant(cond):
- """Mark condition as thread invariant.
-
- Parameters
- ----------
- cond : Expr
- The condition.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- assert isinstance(cond, PrimExpr)
- return call_intrin(cond.dtype, "tirx.tvm_thread_invariant", cond)
-
-
-def tvm_storage_sync(storage_scope, is_load=False, num_blocks=-1):
- """Perform synchronization in specified scope.
-
- Parameters
- ----------
- storage_scope : str
- The storage scope to perform synchronization.
-
- is_load : bool
- Whether to perform load synchronization. (for global sync only)
-
- num_blocks : int
- The number of blocks to synchronize. (for global sync only)
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("void", "tirx.tvm_storage_sync", storage_scope, is_load, num_blocks)
-
-
-def tvm_kernel_replace_point():
- """Mark where a transform should replace generated kernel initialization."""
- return call_intrin("void", "tirx.tvm_kernel_replace_point")
-
-
-def tvm_global_barrier_kinit():
- """Initialize the global barrier.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("void", "tirx.tvm_global_barrier_kinit")
-
-
-def tvm_warp_shuffle(mask, value, warp_id, width, warp_size):
- """Exchange value between threads inside a warp.
-
- Parameters
- ----------
- mask : PrimExpr
- The warp mask indicates active threads inside warp.
- value : PrimExpr
- The value to exchange.
- warp_id : PrimExpr
- The source lane index to fetch value.
- width : PrimExpr
- The width of sub-sections to perform warp shuffle.
- warp_size : PrimExpr
- The warp size.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(value.dtype, "tirx.tvm_warp_shuffle", mask, value, warp_id, width, warp_size)
-
-
-def tvm_warp_shuffle_up(mask, value, offset, width, warp_size):
- """Copy value from a lane with lower (by offset) index relative to caller.
-
- Parameters
- ----------
- mask : PrimExpr
- The warp mask indicates active threads inside warp.
- value : PrimExpr
- The value to exchange.
- offset : PrimExpr
- The difference between source lane index and destination lane index:
- `offset = dst_lane_idx - src_lane_idx`
- width : PrimExpr
- The width of sub-sections to perform warp shuffle.
- warp_size : PrimExpr
- The warp size.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- value.dtype, "tirx.tvm_warp_shuffle_up", mask, value, offset, width, warp_size
- )
-
-
-def tvm_warp_shuffle_down(mask, value, offset, width, warp_size):
- """Copy value from a lane with higher (by offset) index relative to caller.
-
- Parameters
- ----------
- mask : PrimExpr
- The warp mask indicates active threads inside warp.
- value : PrimExpr
- The value to exchange.
- offset : PrimExpr
- The difference between source lane index and destination lane index:
- `offset = src_lane_idx - dst_lane_idx`
- width : PrimExpr
- The width of sub-sections to perform warp shuffle.
- warp_size : PrimExpr
- The warp size.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- value.dtype, "tirx.tvm_warp_shuffle_down", mask, value, offset, width, warp_size
- )
-
-
-def tvm_warp_shuffle_xor(mask, value, lane_mask, width, warp_size):
- """Copy value from a lane with index computed by `src_lane_idx ^ lane_mask`.
-
- Parameters
- ----------
- mask : PrimExpr
- The warp mask indicates active threads inside warp.
- value : PrimExpr
- The value to exchange.
- lane_mask : PrimExpr
- The mask to compute source lane index:
- width : PrimExpr
- The width of sub-sections to perform warp shuffle.
- warp_size : PrimExpr
- The warp size.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- value.dtype, "tirx.tvm_warp_shuffle_xor", mask, value, lane_mask, width, warp_size
- )
-
-
-def tvm_warp_activemask():
- """Return a 32-bit mask indicates currently active threads in a calling warp.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("uint32", "tirx.tvm_warp_activemask")
-
-
-def type_annotation(dtype):
- """Create a type annotation expression
-
- Parameters
- ----------
- dtype : Expr
- The data type.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.type_annotation")
-
-
-def tvm_access_ptr(ptype, data, offset, extent, rw_mask):
- """Get head access address with memory access pattern info
-
- Parameters
- ----------
- ptype : Expr or str
- The data type of pointer. If a ``str``, it is wrapped via
- :func:`type_annotation` so that the lowering rule (which reads
- ``args[0].dtype()`` for the cast type) sees the intended dtype
- instead of ``void`` from a raw StringImm.
-
- data : DType*
- The data of pointer.
-
- offset : int
- The offset of pointer.
-
- extent : int
- The extent of pointer.
-
- rw_mask : int
- The read write mask.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- if isinstance(ptype, str):
- ptype = type_annotation(ptype)
- return call_intrin("handle", "tirx.tvm_access_ptr", ptype, data, offset, extent, rw_mask)
-
-
-def ptr_byte_offset(data, byte_offset, dtype):
- """Cast ``data + byte_offset`` to ``dtype*``.
-
- ``byte_offset`` is always in bytes. Use this when the source CUDA shape
- needs an explicitly typed local pointer derived from a byte-addressed base.
- """
- if isinstance(dtype, str):
- dtype = type_annotation(dtype)
- return call_intrin("handle", "tirx.ptr_byte_offset", data, byte_offset, dtype)
-
-
-def tvm_throw_last_error():
- """Throw TVMGetLastError()
-
- Returns
- -------
- ret : PrimExpr
- The return expression
- """
- return call_intrin("handle", "tirx.tvm_throw_last_error")
-
-
-def make_filled_simdgroup_matrix(
- d: Var,
- index: PrimExpr,
- value: PrimExpr,
- col: int = 8,
- row: int = 8,
-):
- """Create a filled SIMDGroup matrix
-
- Parameters
- ----------
- d : var
- The simdgroup var
-
- index : PrimExpr
- The index of the matrix.
-
- value : PrimExpr
- The value to fill.
-
- col : int
- The number of columns.
-
- row : int
- The number of rows.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.make_filled_simdgroup_matrix", d, index, value, col, row)
-
-
-def simdgroup_load(
- d: Var,
- index: PrimExpr,
- ptr: PrimExpr,
- stride: PrimExpr,
- col: int = 8,
- row: int = 8,
- transpose_matrix: bool = False,
-):
- """Load data from device memory or threadgroup memory to simdgroup
-
- Parameters
- ----------
- d : var
- The simdgroup var
-
- index : PrimExpr
- The index of the matrix.
-
- ptr : PrimExpr
- The pointer.
-
- stride : PrimExpr
- The stride.
-
- col : int
- The number of columns.
-
- row : int
- The number of rows.
-
- transpose_matrix : bool
- Whether to transpose the matrix.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.simdgroup_load",
- d,
- index,
- ptr,
- stride,
- col,
- row,
- transpose_matrix,
- )
-
-
-def simdgroup_store(
- d: PrimExpr,
- index: PrimExpr,
- ptr: PrimExpr,
- stride: PrimExpr,
- col: int = 8,
- row: int = 8,
- transpose_matrix: bool = False,
-):
- """Store data from simdgroup to device memory or threadgroup memory
-
- Parameters
- ----------
- d : PrimExpr
- The SIMDGroup.
-
- index : PrimExpr
- The index of the matrix.
-
- ptr : PrimExpr
- The pointer.
-
- stride : PrimExpr
- The stride.
-
- col : int
- The number of columns.
-
- row : int
- The number of rows.
-
-
- transpose_matrix : bool
- Whether to transpose the matrix.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.simdgroup_store",
- d,
- index,
- ptr,
- stride,
- col,
- row,
- transpose_matrix,
- )
-
-
-def simdgroup_multiply_accumulate(
- d: Var,
- index_d: PrimExpr,
- a: Var,
- index_a: PrimExpr,
- b: Var,
- index_b: PrimExpr,
- c: Var,
- index_c: PrimExpr,
-):
- """Multiply and accumulate two matrices in simdgroup
- i.e. d = a * b + c
-
- Parameters
- ----------
- d : Var
- The destination matrix.
-
- index_d : PrimExpr
- The index of the destination matrix.
-
- a : Var
- The first matrix.
-
- index_a : PrimExpr
- The index of the first matrix.
-
- b : Var
- The second matrix.
-
- index_b : PrimExpr
- The index of the second matrix.
-
- c : Var
- The third matrix.
-
- index_c : PrimExpr
- The index of the third matrix.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.simdgroup_multiply_accumulate",
- d,
- index_d,
- a,
- index_a,
- b,
- index_b,
- c,
- index_c,
- )
-
-
-def cooperative_tensor_fill(
- d: Var,
- index: PrimExpr,
- value: PrimExpr,
- rows: int,
- cols: int,
-):
- return call_intrin("handle", "tirx.cooperative_tensor_fill", d, index, value, rows, cols)
-
-
-def cooperative_tensor_load(
- d: Var,
- index: PrimExpr,
- ptr: PrimExpr,
- stride: PrimExpr,
- rows: int,
- cols: int,
- transpose_matrix: bool = False,
- mma_M: int = 0,
- mma_N: int = 0,
- mma_K: int = 0,
- operand_role: int = 0,
-):
- return call_intrin(
- "handle",
- "tirx.cooperative_tensor_load",
- d,
- index,
- ptr,
- stride,
- rows,
- cols,
- transpose_matrix,
- mma_M,
- mma_N,
- mma_K,
- operand_role,
- )
-
-
-def cooperative_tensor_store(
- d: PrimExpr,
- index: PrimExpr,
- ptr: PrimExpr,
- stride: PrimExpr,
- rows: int,
- cols: int,
- transpose_matrix: bool = False,
- mma_M: int = 0,
- mma_N: int = 0,
- mma_K: int = 0,
- operand_role: int = 0,
-):
- return call_intrin(
- "handle",
- "tirx.cooperative_tensor_store",
- d,
- index,
- ptr,
- stride,
- rows,
- cols,
- transpose_matrix,
- mma_M,
- mma_N,
- mma_K,
- operand_role,
- )
-
-
-def cooperative_tensor_multiply_accumulate(
- d: Var,
- index_d: PrimExpr,
- a: Var,
- index_a: PrimExpr,
- b: Var,
- index_b: PrimExpr,
- c: Var,
- index_c: PrimExpr,
- M: int,
- N: int,
- K: int,
- transpose_a: bool = False,
- transpose_b: bool = False,
-):
- return call_intrin(
- "handle",
- "tirx.cooperative_tensor_multiply_accumulate",
- d,
- index_d,
- a,
- index_a,
- b,
- index_b,
- c,
- index_c,
- M,
- N,
- K,
- transpose_a,
- transpose_b,
- )
-
-
-def vectorlow(dtype, vec):
- """Get the low level half of the vector
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- vec : list
- The input vector.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.vectorlow", vec)
-
-
-def vectorhigh(dtype, vec):
- """Get the high level half of the vector
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- vec : list
- The input vector.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.vectorhigh", vec)
-
-
-def vectorcombine(dtype, vec1, vec2):
- """Concat two vectors
-
- Parameters
- ----------
- vec1 : list
- The input vector.
-
- vec2 : list
- The input vector.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.vectorcombine", vec1, vec2)
-
-
-def dp4a(vec1, vec2, acc=0):
- """Dot product of two int8x4 vectors and add an optional accumulator
-
- Parameters
- ----------
- vec1 : int8x4
- The input vector.
-
- vec2 : int8x4
- The input vector.
-
- acc : int32
- The accumulator.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("int32", "tirx.dp4a", vec1, vec2, acc)
-
-
-def ret(val, span=None):
- """Create a tir return expression
-
- Parameters
- ----------
- val : Expr
- The returned tir expression, whose data type is int, float or void pointer.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- ret : PrimExpr
- The return expression
- """
-
- return _ffi_api.ret(val, span)
-
-
-def any(*args, span=None):
- """Create a new experssion of the union of all conditions in the arguments
-
- Parameters
- ----------
- args : list
- List of symbolic boolean expressions
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- expr: Expr
- Expression
- """
- if not args:
- raise ValueError("Any must take at least 1 argument")
- if len(args) == 1:
- return args[0]
- val = _ffi_api._OpOr(args[0], args[1], span) # type: ignore
- for i in range(2, len(args)):
- val = _ffi_api._OpOr(val, args[i], span) # type: ignore
- return val
-
-
-def all(*args, span=None):
- """Create a new expression of the intersection of all conditions in the
- arguments
-
- Parameters
- ----------
- args : list
- List of symbolic boolean expressions
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- expr: Expr
- Expression
- """
- if not args:
- raise ValueError("Any must take at least 1 argument")
- if len(args) == 1:
- return args[0]
- val = _ffi_api._OpAnd(args[0], args[1], span) # type: ignore
- for i in range(2, len(args)):
- val = _ffi_api._OpAnd(val, args[i], span) # type: ignore
- return val
-
-
-@tvm_ffi.register_global_func("tvm.default_trace_action")
-def _tvm_default_trace_action(*args):
- print(list(args))
-
-
-def trace(args, trace_action="tvm.default_trace_action"):
- """Trace tensor data at the runtime.
-
- The trace function allows to trace specific tensor at the
- runtime. The tracing value should come as last argument.
- The trace action should be specified, by default
- tvm.default_trace_action is used.
-
- Parameters
- ----------
- args : list of Expr or Buffers.
- Positional arguments.
-
- trace_action : str.
- The name of the trace action.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
-
- See Also
- --------
- tvm.tirx.call_packed : Creates packed function.
- """
- if not isinstance(args, list):
- raise Exception("tvm.tirx.trace consumes the args as list type")
- call_args = [_pack_buffer(x) if isinstance(x, Buffer) else x for x in args]
- call_args.insert(0, trace_action)
- return tvm.tirx.Call(args[-1].dtype, Op.get("tirx.tvm_call_trace_packed"), call_args)
-
-
-def min_value(dtype, span=None):
- """minimum value of dtype
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- value : tvm.Expr
- The minimum value of dtype.
- """
- return _ffi_api.min_value(dtype, span) # type: ignore
-
-
-def max_value(dtype: str, span: Span | None = None) -> Any:
- """maximum value of dtype
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- value : tvm.Expr
- The maximum value of dtype.
- """
- return _ffi_api.max_value(dtype, span) # type: ignore
-
-
-def infinity(dtype: str, span: Span | None = None) -> Any:
- """infinity value of dtype
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- value : tvm.Expr
- The infinity value of dtype.
- """
- return _ffi_api.infinity(dtype, span) # type: ignore
-
-
-def reinterpret(dtype, value, span: Span | None = None) -> Any:
- """infinity value of dtype
-
- Parameters
- ----------
- dtype : str
- The data type.
-
- value : PrimExpr
- The input value.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- value : tvm.Expr
- The reinterpret cast value of dtype.
- """
- return _ffi_api.reinterpret(dtype, value, span) # type: ignore
-
-
-def exp(x):
- """Take exponential of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.exp", x)
-
-
-def exp2(x):
- """Calculate 2**x
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.exp2", x)
-
-
-def exp10(x):
- """Calculate 10**x
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.exp10", x)
-
-
-def fma(x, y, z):
- """Take fused multiply-add of input x, y, z.
-
- Parameters
- ----------
- x : PrimExpr
- First input argument.
-
- y : PrimExpr
- Second input argument.
-
- z : PrimExpr
- Third input argument.
-
- Returns
- -------
- out : PrimExpr
- The result of x * y + z.
- """
- x = tir.convert(x)
- y = tir.convert(y)
- z = tir.convert(z)
- return call_intrin(x.dtype, "tirx.fma", x, y, z)
-
-
-def erf(x):
- """Take gauss error function of the input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.erf", x)
-
-
-def tanh(x):
- """Take hyperbolic tanh of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.tanh", x)
-
-
-def sigmoid(x):
- """Quick function to get sigmoid
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.sigmoid", x)
-
-
-def log(x):
- """Take log of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.log", x)
-
-
-def log2(x):
- """Take log2 of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.log2", x)
-
-
-def log10(x):
- """Take log10 of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.log10", x)
-
-
-def log1p(x):
- """Take log(x + 1) with respect to input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.log1p", x)
-
-
-def tan(x):
- """Take tan of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = _require_float_arg("tan", x)
- return call_intrin(x.dtype, "tirx.tan", x)
-
-
-def cos(x):
- """Take cos of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = _require_float_arg("cos", x)
- return call_intrin(x.dtype, "tirx.cos", x)
-
-
-def cosh(x):
- """Take cosh of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.cosh", x)
-
-
-def acos(x):
- """Take acos of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.acos", x)
-
-
-def acosh(x):
- """Take acos of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.acosh", x)
-
-
-def sin(x):
- """Take sin of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = _require_float_arg("sin", x)
- return call_intrin(x.dtype, "tirx.sin", x)
-
-
-def sinh(x):
- """Take sinh of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.sinh", x)
-
-
-def asin(x):
- """Take asin of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.asin", x)
-
-
-def asinh(x):
- """Take asinh of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.asinh", x)
-
-
-def atan(x):
- """Take atan of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.atan", x)
-
-
-def atanh(x):
- """Take atanh of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.atanh", x)
-
-
-def atan2(x1, x2):
- """Take arctan2(x1, x2).
-
- Parameters
- ----------
- x1 : PrimExpr
- Input argument.
-
- x2 : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x1 = tir.convert(x1)
- x2 = tir.convert(x2)
- return call_intrin(x1.dtype, "tirx.atan2", x1, x2)
-
-
-def sqrt(x):
- """Take square root of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.sqrt", x)
-
-
-def rsqrt(x):
- """Take reciprocal of square root of input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.rsqrt", x)
-
-
-def clz(x):
- """Count leading zero bits of an integer x.
-
- Parameters
- ----------
- x : PrimExpr
- Input 32 or 64 bit integer.
- The result is undefined if the input is 0.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return call_intrin("int32", "tirx.clz", x)
-
-
-def floor(x: PrimExprWithOp, span=None):
- """Take floor of float input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.floor(x, span) # type: ignore
-
-
-def ceil(x, span=None):
- """Take ceil of float input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.ceil(x, span) # type: ignore
-
-
-def trunc(x, span=None):
- """Get truncated value of the input.
-
- The truncated value of the scalar x is the
- nearest integer i which is closer to zero than x is.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.trunc(x, span) # type: ignore
-
-
-def abs(x, span=None):
- """Get absolute value of the input element-wise.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.abs(x, span) # type: ignore
-
-
-def bitwise_and(x, y, span=None):
- """Take bitwise and of two values
-
- Parameters
- ----------
- x : PrimExpr
- Left operand
-
- y : PrimExpr
- Right operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : PrimExpr
- The result.
- """
- return _ffi_api.bitwise_and(x, y, span)
-
-
-def bitwise_not(x, span=None):
- """Take bitwise not of input value
-
- Parameters
- ----------
- x : PrimExpr
- Input operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : PrimExpr
- The result.
- """
- return _ffi_api.bitwise_not(x, span)
-
-
-def bitwise_or(x, y, span=None):
- """Take bitwise or of two values
-
- Parameters
- ----------
- x : PrimExpr
- Left operand
-
- y : PrimExpr
- Right operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : PrimExpr
- The result.
- """
- return _ffi_api.bitwise_or(x, y, span)
-
-
-def bitwise_xor(x, y, span=None):
- """Take bitwise xor of two values
-
- Parameters
- ----------
- x : PrimExpr
- Left operand
-
- y : PrimExpr
- Right operand
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- res : PrimExpr
- The result.
- """
- return _ffi_api.bitwise_xor(x, y, span)
-
-
-def round(x, span=None):
- """Round elements of the array to the nearest integer.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.round(x, span) # type: ignore
-
-
-def nearbyint(x, span=None):
- """Round elements of the array to the nearest integer.
- This intrinsic uses llvm.nearbyint instead of llvm.round
- which is faster but will results different from te.round.
- Notably nearbyint rounds according to the rounding mode,
- whereas te.round (llvm.round) ignores that.
- For differences between the two see:
- https://en.cppreference.com/w/cpp/numeric/math/round
- https://en.cppreference.com/w/cpp/numeric/math/nearbyint
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.nearbyint(x, span) # type: ignore
-
-
-def nextafter(x1, x2):
- """Return the next floating-point value after x1 towards x2.
-
- Parameters
- ----------
- x1 : PrimExpr
- Input argument.
-
- x2 : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x1 = tir.convert(x1)
- x2 = tir.convert(x2)
- return call_intrin(x1.dtype, "tirx.nextafter", x1, x2) # type: ignore
-
-
-def hypot(x1, x2):
- """Equivalent to sqrt(x1**2 + x2**2), element-wise.
-
- Parameters
- ----------
- x1 : PrimExpr
- Input argument.
-
- x2 : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x1 = tir.convert(x1)
- x2 = tir.convert(x2)
- return call_intrin(x1.dtype, "tirx.hypot", x1, x2) # type: ignore
-
-
-def copysign(x1, x2):
- """Change the sign of x1 to that of x2, element-wise.
-
- Parameters
- ----------
- x1 : PrimExpr
- Input argument.
-
- x2 : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x1 = tir.convert(x1)
- x2 = tir.convert(x2)
- return call_intrin(x1.dtype, "tirx.copysign", x1, x2) # type: ignore
-
-
-def ldexp(x1, x2):
- """Returns x1 * (2 ** x2).
-
- Parameters
- ----------
- x1 : PrimExpr
- Input argument.
-
- x2 : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x1 = tir.convert(x1)
- x2 = tir.convert(x2)
- return call_intrin(x1.dtype, "tirx.ldexp", x1, x2) # type: ignore
-
-
-def likely(cond, span=None):
- """Mark condition as likely.
-
- Parameters
- ----------
-
- cond : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The marked expression.
- """
- return _ffi_api.likely(cond, span) # type: ignore
-
-
-def filter(var, pred, *, span=None): # pylint: disable=redefined-builtin
- """Thread-set filter escape hatch.
-
- Use this wrapper only when the predicate is *not* in the canonical
- thread-filter grammar (see ``src/tirx/analysis/filter_canonical.h``).
- Canonical predicates -- pure conjunctions of ``scopeid_var <op> const``
- comparisons plus bare ``T.ptx.elect_sync()`` calls -- are recognized by
- the lowering pass directly from ``if cond:``, so the wrapper is redundant
- for them.
-
- When wrapped: ``var`` (a ``ScopeIdDef``-declared scope identifier) tells
- the compiler which active-set axis to collapse to a singleton when the
- opaque predicate evaluates true; ``pred`` is preserved verbatim and
- evaluated at runtime.
-
- The legacy three-argument range form ``filter(var, lo, hi)`` has been
- removed -- write ``lo <= var and var < hi`` (or ``var == lo`` when
- ``hi == lo + 1``) at the call site instead.
- """
- return call_intrin("bool", "tirx.filter", var, pred, span=span)
-
-
-def selector(var, pred, span=None):
- """Analysis-only active-thread selector.
-
- ``selector(var, pred)`` denotes the unique value of ``var`` in the current
- active domain for which ``pred`` is true. It is intended for compiler
- metadata and should not survive to executable codegen.
- """
- return call_intrin(var.dtype, "tirx.selector", var, pred, span=span)
-
-
-def isnan(x, span=None):
- """Check if input value is Nan.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.isnan(x, span) # type: ignore
-
-
-def isnullptr(x, span=None):
- """Check if input value is nullptr.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return call_intrin("bool", "tirx.isnullptr", x, span=span) # type: ignore
-
-
-def isfinite(x, span=None):
- """Check if input value is finite.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.isfinite(x, span) # type: ignore
-
-
-def isinf(x, span=None):
- """Check if input value is infinite.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return _ffi_api.isinf(x, span) # type: ignore
-
-
-def power(x, y, span=None):
- """x power y
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- y : PrimExpr
- The exponent
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- z : PrimExpr
- The result.
- """
- return _ffi_api._OpPow(x, y, span) # type: ignore
-
-
-def pow(x, y, span=None):
- """x power y
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- y : PrimExpr
- The exponent
-
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- z : PrimExpr
- The result.
- """
- return _ffi_api._OpPow(x, y, span) # type: ignore
-
-
-def popcount(x):
- """Count the number of set bits in input x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- x = tir.convert(x)
- return call_intrin(x.dtype, "tirx.popcount", x)
-
-
-def q_multiply_shift(x, y, q, s):
- """Execute a multiplication between two Q-numbers x and y
- followed by a right shift s. The mathematical expression is:
-
- out = round(x*y*2^-s)
-
- More about Q-numbers here: https://en.wikipedia.org/wiki/Q_(number_format)
- The rounding rule is to the nearest value, rounding half up
- (i.e., round(x.1) = x and round (x.5) = x+1)
-
- Parameters
- ----------
- x : PrimExpr
- First Q-number
- y : PrimExpr
- Second Q-number
- q : PrimExpr
- Number of fractional bits in x and y. Needs to be > 0
- s : PrimExpr
- Integer shift
-
- Returns
- -------
- y : PrimExpr
- The result.
- """
- return call_intrin("int32", "tirx.q_multiply_shift", x, y, q, s)
-
-
-def q_multiply_shift_per_axis(
- x: PrimExpr,
- y: PrimExpr,
- ls: PrimExpr,
- rs: PrimExpr,
- q: IntImm,
- is_lshift_required: IntImm,
- is_rshift_required: IntImm,
-):
- """Execute a multiplication between two Q-numbers x and y
-
- Parameters
- ----------
- x : PrimExpr
- First Q-number.
- y : PrimExpr
- Second Q-number.
- ls : PrimExpr
- Integer left shift.
- rs : PrimExpr
- Integer right shift.
- q : IntImm
- Number of fractional bits in x and y. Needs to be > 0.
- is_lshift_required : IntImm
- Whether we need to do left shift or not.
- is_rshift_required : IntImm
- Whether we need to do right shift or not.
-
- Returns
- -------
- z : PrimExpr
- The result.
- """
- return call_intrin(
- "int32",
- "tirx.q_multiply_shift_per_axis",
- x,
- y,
- ls,
- rs,
- q,
- is_lshift_required,
- is_rshift_required,
- )
-
-
-def shift_left(x, y, span=None):
- """Return the result of x left shifted by y bits.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- y : PrimExpr
- Input argument.
-
- Returns
- -------
- z : PrimExpr
- The result.
- """
- return _ffi_api.left_shift(x, y, span)
-
-
-def shift_right(x, y, span=None):
- """Return the result of x right shifted by y bits.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
-
- y : PrimExpr
- Input argument.
-
- Returns
- -------
- z : PrimExpr
- The result.
- """
- return _ffi_api.right_shift(x, y, span)
-
-
-def fmod(x, y):
- """Return the remainder of x divided by y with the same sign as x.
-
- Parameters
- ----------
- x : PrimExpr
- Input argument.
- y : PrimExpr
- Input argument.
-
- Returns
- -------
- z : PrimExpr
- The result.
- """
- x = tir.convert(x)
- y = tir.convert(y)
- return call_intrin(x.dtype, "tirx.fmod", x, y)
-
-
-def if_then_else(cond, t, f, span=None):
- """Conditional selection expression.
-
- Parameters
- ----------
- cond : PrimExpr
- The condition
-
- t : PrimExpr
- The result expression if cond is true.
-
- f : PrimExpr
- The result expression if cond is false.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- result : Node
- The result of conditional expression.
-
- Note
- ----
- Unlike Select, if_then_else will not execute
- the branch that does not satisfy the condition.
- You can use it to guard against out of bound access.
- Unlike Select, if_then_else cannot be vectorized
- if some lanes in the vector have different conditions.
- """
- return _ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore
-
-
-def div(a, b, span=None):
- """Compute a / b as in C/C++ semantics.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand, known to be non-negative.
-
- b : PrimExpr
- The right hand operand, known to be non-negative.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
- Note
- ----
- When operands are integers, returns truncdiv(a, b, span).
- """
- return _ffi_api._OpDiv(a, b, span) # type: ignore
-
-
-def indexdiv(a, b, span=None):
- """Compute floor(a / b) where a and b are non-negative.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand, known to be non-negative.
-
- b : PrimExpr
- The right hand operand, known to be non-negative.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
-
- Note
- ----
- Use this function to split non-negative indices.
- This function may take advantage of operands'
- non-negativeness.
- """
- return _ffi_api._OpIndexDiv(a, b, span) # type: ignore
-
-
-def indexmod(a, b, span=None):
- """Compute the remainder of indexdiv. a and b are non-negative.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand, known to be non-negative.
-
- b : PrimExpr
- The right hand operand, known to be non-negative.
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
-
- Note
- ----
- Use this function to split non-negative indices.
- This function may take advantage of operands'
- non-negativeness.
- """
- return _ffi_api._OpIndexMod(a, b, span) # type: ignore
-
-
-def truncdiv(a, b, span=None):
- """Compute the truncdiv of two expressions.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand
-
- b : PrimExpr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
-
- Note
- ----
- This is the default integer division behavior in C.
- """
- return _ffi_api._OpTruncDiv(a, b, span) # type: ignore
-
-
-def truncmod(a, b, span=None):
- """Compute the truncmod of two expressions.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand
-
- b : PrimExpr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
-
- Note
- ----
- This is the default integer division behavior in C.
- """
- return _ffi_api._OpTruncMod(a, b, span) # type: ignore
-
-
-def floordiv(a, b, span=None):
- """Compute the floordiv of two expressions.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand
-
- b : PrimExpr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
- """
- return _ffi_api._OpFloorDiv(a, b, span) # type: ignore
-
-
-def logaddexp(a, b, span=None):
- """Compute the logaddexp of two expressions.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand
-
- b : PrimExpr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
- """
- return _ffi_api._OpLogAddExp(a, b, span) # type: ignore
-
-
-def floormod(a, b, span=None):
- """Compute the floormod of two expressions.
-
- Parameters
- ----------
- a : PrimExpr
- The left hand operand
-
- b : PrimExpr
- The right hand operand
-
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- res : PrimExpr
- The result expression.
- """
- return _ffi_api._OpFloorMod(a, b, span) # type: ignore
-
-
-def ceildiv(lhs, rhs, span=None):
- """Generic ceildiv operator.
-
- Parameters
- ----------
- lhs : object
- The left operand.
- rhs : object
- The right operand.
- span : Optional[Span]
- The location of this operator in the source.
-
- Returns
- -------
- op : tvm.Expr
- The result Expr of ceildiv operaton.
- """
- return _ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore
-
-
-def comm_reducer(fcombine, fidentity, name="reduce"):
- """Create a commutative reducer for reduction.
-
- Parameters
- ----------
- fcombine : function(Expr -> Expr -> Expr)
- A binary function which takes two Expr as input to return a Expr.
-
- fidentity : function(str -> Expr)
- A function which takes a type string as input to return a const Expr.
-
- Returns
- -------
- reducer : function
- A function which creates a reduce expression over axis.
- There are two ways to use it:
-
- 1. accept (expr, axis, where) to produce an Reduce Expr on
- specified axis;
- 2. simply use it with multiple Exprs.
-
- Example
- -------
- .. code-block:: python
-
- n = te.var("n")
- m = te.var("m")
- mysum = te.comm_reducer(lambda x, y: x+y,
- lambda t: tvm.tirx.const(0, dtype=t), name="mysum")
- A = te.placeholder((n, m), name="A")
- k = te.reduce_axis((0, m), name="k")
- B = te.compute((n,), lambda i: mysum(A[i, k], axis=k), name="B")
- """
-
- def _reduce_directly(*args):
- num = len(args)
- # process `where` is None
- if num == 3 and args[2] is None:
- num = 2
- res = args[0]
- for i in range(num - 1):
- res = fcombine(res, args[i + 1])
- return res
-
- def _make_reduce(expr, axis, where=None, init=None):
- code = fcombine.__code__
- assert fcombine.__code__.co_argcount == 2
- expr = tir.convert(expr)
- if init is not None:
- init = tir.convert(init)
- if isinstance(expr, Array):
- size = len(expr)
- lhs = []
- rhs = []
- dtypes = []
- for i in range(size):
- dtype = expr[i].dtype
- dtypes.append(dtype)
- lname = code.co_varnames[0] + "_" + str(i)
- lhs.append(Var(lname, dtype))
- rname = code.co_varnames[1] + "_" + str(i)
- rhs.append(Var(rname, dtype))
- if init is None:
- init = []
- result = fcombine(lhs, rhs)
- id_elem = fidentity(*dtypes)
- else:
- assert isinstance(expr, tvm.ir.PrimExpr)
- size = 1
- dtype = expr.dtype
- lvar = Var(code.co_varnames[0], dtype)
- rvar = Var(code.co_varnames[1], dtype)
- result = [fcombine(lvar, rvar)]
- id_elem = [fidentity(dtype)]
- lhs = [lvar]
- rhs = [rvar]
- expr = [expr]
- if init is not None:
- init = [init]
- combiner = CommReducer(lhs, rhs, result, id_elem)
- if not isinstance(axis, list | tuple | tvm.ir.Array):
- axis = [axis]
- if where is None:
- where = tir.convert(True)
- if init is None:
- outputs = tuple(
- tvm.tirx.Reduce(combiner, expr, axis, where, i, []) for i in range(size)
- )
- else:
- outputs = tuple(
- tvm.tirx.Reduce(combiner, expr, axis, where, i, init) for i in range(size)
- )
- return outputs[0] if size == 1 else outputs
-
- # pylint: disable=keyword-arg-before-vararg
- def reducer(expr, axis, where=None, init=None, *args):
- if isinstance(axis, tvm.tirx.IterVar | list | tuple):
- assert not args
- return _make_reduce(expr, axis, where, init)
-
- if where is None:
- assert not args
- assert init is None
- return _reduce_directly(expr, axis)
- elif init is None:
- assert not args
- return _reduce_directly(expr, axis, where)
- else:
- return _reduce_directly(expr, axis, where, init, *args)
-
- doc_str = """Create a {0} expression over axis.
-
- Parameters
- ----------
- expr : PrimExpr
- The source expression.
- axis : IterVar
- The reduction IterVar axis
- where : optional, Expr
- Filtering predicate of the reduction.
- Returns
- -------
- value : PrimExpr
- The result value.
-
- Example
- -------
- .. code-block:: python
-
- m = te.var("m")
- n = te.var("n")
- A = te.placeholder((m, n), name="A")
- k = te.reduce_axis((0, n), name="k")
-
- # there are two way to use this {0} reducer:
- # mode 1, accept (expr, axis, where) to produce an Reduce Expr
- # tvm.{0} represents tvm.te.{0} or tvm.tirx.{0}.
- B = te.compute((m,), lambda i: tvm.{0}(A[i, k], axis=k), name="B")
-
- # mode 2, simply use it with multiple Exprs:
- {0}_res = tvm.{0}(m, n)
- """
- reducer.__doc__ = doc_str.format(name)
- return reducer
-
-
-def TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint):
- """Backend function to allocate temporal workspace
-
- Parameters
- ----------
- device_type : int
- The device type which the space will be allocated.
-
- device_id : int
- The device id which the space will be allocated.
-
- nbytes : int
- The size of the space requested.
-
- dtype_code_hint : int
- The type code of the array elements. Only used in certain backends such as OpenGL.
-
- dtype_bits_hint : int
- The type bits of the array elements. Only used in certain backends such as OpenGL.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.TVMBackendAllocWorkspace",
- device_type,
- device_id,
- nbytes,
- dtype_code_hint,
- dtype_bits_hint,
- )
-
-
-def TVMBackendFreeWorkspace(device_type, device_id, ptr):
- """Backend function to free temporal workspace.
-
- Parameters
- ----------
- device_type : int
- The device type which the space will be allocated.
-
- device_id : int
- The device id which the space will be allocated.
-
- ptr : Var
- The result allocated space pointer.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("int32", "tirx.TVMBackendFreeWorkspace", device_type, device_id, ptr)
-
-
-def anylist_getitem(list_handle, index):
- """Returns an item from any list.
- list_handle: Var
- The handle to anylist
- index : int
- The index
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.anylist_getitem", list_handle, index)
-
-
-def anylist_resetitem(list_handle, index):
- """Reset an item from any list.
- list_handle: Var
- The handle to anylist
- index : int
- The index
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("int", "tirx.anylist_resetitem", list_handle, index)
-
-
-def anylist_setitem_call_packed(list_handle, index, func_name, *args):
- """Set anylist item by result of packed call.
- list_handle: Var
- The handle to anylist
- index : int
- The index
- func_name: str
- The name of the function to be called.
- args:
- Extra arguments
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "int", "tirx.anylist_setitem_call_packed", list_handle, index, func_name, *args
- )
-
-
-def anylist_setitem_call_cpacked(list_handle, index, func_name, *args):
- """Set anylist item by result of packed call.
- list_handle: Var
- The handle to anylist
- index : int
- The index
- func_name: str
- The name of the function to be called.
- args:
- Extra arguments
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "int", "tirx.anylist_setitem_call_cpacked", list_handle, index, func_name, *args
- )
-
-
-def vscale():
- """Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
- (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
- Returns
- -------
- call : PrimExpr
- Call to the vscale intrinsic
- """
- return call_intrin("int32", "tirx.vscale")
-
-
-def get_active_lane_mask(dtype, base, limit):
- """
- Calculate a predicate mask given an upper bound (limit) and a current value (base).
-
- It will be lowered to the llvm.get.active.lane.mask intrinsic.
- (https://llvm.org/docs/LangRef.html#llvm-get-active-lane-mask-intrinsics)
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- base : PrimExpr
- An expression reprsenting the base.
-
- limit : PrimExpr
- An expression representing the limit.
- """
- return call_intrin(dtype, "tirx.get_active_lane_mask", base, limit)
-
-
-def get_vscale_expr(dtype: str | tvm_ffi.dtype, min_size: int = 128) -> PrimExpr:
- """
- Create a datatype dependent scalable expression.
-
- Parameters
- ----------
- dtype : Union[str, tvm_ffi.DataType]
- Element data type.
- min_size : int
- The minimum size of the scalable vector in bits.
- """
- if isinstance(dtype, str):
- dtype = tvm_ffi.dtype(dtype)
- return min_size // dtype.bits * vscale()
-
-
-def ignore_loop_partition(predicate) -> PrimExpr:
- """
- Annotate a predicate not be considered as target condition of loop partition.
-
- Parameters
- ----------
- predicate : PrimExpr
- The annotated predicate expression.
- """
- return call_intrin("bool", "tirx.ignore_loop_partition", predicate)
-
-
-# pylint: disable=unnecessary-lambda
-sum = comm_reducer(lambda x, y: x + y, lambda t: const(0, dtype=t), name="sum")
-min = comm_reducer(lambda x, y: _ffi_api._OpMin(x, y, None), max_value, name="min") # type: ignore
-max = comm_reducer(lambda x, y: _ffi_api._OpMax(x, y, None), min_value, name="max") # type: ignore
+from __future__ import annotations
+from tvm import tirx
+from tvm.ir import Op, PrimExpr
+from tvm.runtime import const
+from tvm.tirx.expr import Call
+from tvm.tirx.op import bitwise_and, call_intrin, tvm_access_ptr
+from tvm.tirx.operator.intrinsics._common import CLUSTER_BARRIER_SEM as _CLUSTER_BARRIER_SEM
+from tvm.tirx.operator.intrinsics._common import (
+ CP_ASYNC_BULK_CACHE_HINT as _CP_ASYNC_BULK_CACHE_HINT,
+)
+from tvm.tirx.operator.intrinsics._common import CP_ASYNC_BULK_RED_OP as _CP_ASYNC_BULK_RED_OP
+from tvm.tirx.operator.intrinsics._common import CP_ASYNC_CACHE_HINT as _CP_ASYNC_CACHE_HINT
+from tvm.tirx.operator.intrinsics._common import CP_ASYNC_FILL_MODE as _CP_ASYNC_FILL_MODE
+from tvm.tirx.operator.intrinsics._common import CP_ASYNC_PREFETCH_SIZE as _CP_ASYNC_PREFETCH_SIZE
+from tvm.tirx.operator.intrinsics._common import F32X2_ROUND as _F32X2_ROUND
+from tvm.tirx.operator.intrinsics._common import FENCE_PROXY_ASYNC_SPACE as _FENCE_PROXY_ASYNC_SPACE
+from tvm.tirx.operator.intrinsics._common import FENCE_SCOPE as _FENCE_SCOPE
+from tvm.tirx.operator.intrinsics._common import FENCE_SEM as _FENCE_SEM
+from tvm.tirx.operator.intrinsics._common import LDMATRIX_DTYPE as _LDMATRIX_DTYPE
+from tvm.tirx.operator.intrinsics._common import LDMATRIX_NUM as _LDMATRIX_NUM
+from tvm.tirx.operator.intrinsics._common import NVSHMEM_CMP as _NVSHMEM_CMP
+from tvm.tirx.operator.intrinsics._common import NVSHMEM_SIG_OP as _NVSHMEM_SIG_OP
+from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_DECOMPRESS as _TCGEN05_CP_DECOMPRESS
+from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_MULTICAST as _TCGEN05_CP_MULTICAST
+from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_SHAPES as _TCGEN05_CP_SHAPES
+from tvm.tirx.operator.intrinsics._common import TCGEN05_CTA_GROUP as _TCGEN05_CTA_GROUP
+from tvm.tirx.operator.intrinsics._common import TCGEN05_LDST_SHAPES as _TCGEN05_LDST_SHAPES
+
+tir = tirx
########################################################
# CUDA native builtins
@@ -3509,215 +398,6 @@ def cuda_float8tohalf8(src_addr, dst_addr):
return call_intrin("", "tirx.cuda_float8tohalf8", src_addr, dst_addr)
-def tvm_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout):
- """TVM intrinsic for tensor core load operators
-
- Parameters
- ----------
- fragment : Var
- The wmma fragment.
-
- m : UIntImm
- The shape of wmma fragment.
-
- n : UIntImm
- The shape of wmma fragment.
-
- k : UIntImm
- The shape of wmma fragment.
-
- index : Expr
- The fragment index.
-
- buffer_ptr : Expr
- The fragment buffer pointer.
-
- stride : Expr
- The fragment stride.
-
- layout : Literal["row_major", "column_major"]
- The fragment layout.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle", "tirx.tvm_load_matrix_sync", fragment, m, n, k, index, buffer_ptr, stride, layout
- )
-
-
-def tvm_mma_sync(
- fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c
-):
- """TVM intrinsic for tensor core mma_sync operators
-
- Parameters
- ----------
- fragment_d : Var
- The wmma fragment_d.
-
- index_d : Expr
- The fragment_d index.
-
- fragment_a : Var
- The wmma fragment_a.
-
- index_a : Expr
- The fragment_a index.
-
- fragment_b : Var
- The wmma fragment_b.
-
- index_b : Expr
- The fragment_b index.
-
- fragment_c : Var
- The wmma fragment_c.
-
- index_c : Expr
- The fragment_c index.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.tvm_mma_sync",
- fragment_d,
- index_d,
- fragment_a,
- index_a,
- fragment_b,
- index_b,
- fragment_c,
- index_c,
- )
-
-
-def tvm_bmma_sync(
- fragment_d, index_d, fragment_a, index_a, fragment_b, index_b, fragment_c, index_c
-):
- """TVM intrinsic for tensor core bmma_sync operators
-
- Parameters
- ----------
- fragment_d : Var
- The bwmma fragment_d.
-
- index_d : Expr
- The fragment_d index.
-
- fragment_a : Var
- The bwmma fragment_a.
-
- index_a : Expr
- The fragment_a index.
-
- fragment_b : Var
- The bwmma fragment_b.
-
- index_b : Expr
- The fragment_b index.
-
- fragment_c : Var
- The bwmma fragment_c.
-
- index_c : Expr
- The fragment_c index.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle",
- "tirx.tvm_bmma_sync",
- fragment_d,
- index_d,
- fragment_a,
- index_a,
- fragment_b,
- index_b,
- fragment_c,
- index_c,
- )
-
-
-def tvm_fill_fragment(fragment, m, n, k, index, value):
- """TVM intrinsic for tensor core fill_fragment operators
-
- Parameters
- ----------
- fragment : Var
- The wmma fragment
-
- m : UIntImm
- The shape of wmma fragment.
-
- n : UIntImm
- The shape of wmma fragment.
-
- k : UIntImm
- The shape of wmma fragment.
-
- index : Expr
- The fragment index.
-
- value : Expr
- The value to be filled in fragment.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("handle", "tirx.tvm_fill_fragment", fragment, m, n, k, index, value)
-
-
-def tvm_store_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout):
- """TVM intrinsic for tensor core store operators
-
- Parameters
- ----------
- fragment : Var
- The wmma fragment.
-
- m : UIntImm
- The shape of wmma fragment.
-
- n : UIntImm
- The shape of wmma fragment.
-
- k : UIntImm
- The shape of wmma fragment.
-
- index : Expr
- The fragment index.
-
- buffer_ptr : Expr
- The fragment buffer pointer.
-
- stride : Expr
- The fragment stride.
-
- layout : Literal["row_major", "column_major"]
- The fragment layout.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "handle", "tirx.tvm_store_matrix_sync", fragment, m, n, k, index, buffer_ptr, stride, layout
- )
-
-
def ptx_mma_sp(
dtype,
shape,
@@ -3820,92 +500,6 @@ def ptx_mma_sp(
)
-def mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
- """TVM intrinsic for storing the result of PTX MMA into a destination pointer
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- m : IntImm
- The shape of mma fragment.
-
- n : IntImm
- The shape of mma fragment.
-
- dst_ptr : Var
- The destination pointer variable.
-
- src_ptr : Var
- The source pointer variable.
-
- src_offset : Expr
- The source offset.
-
- dst_stride : Var
- The destination stride.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.mma_store", m, n, dst_ptr, src_ptr, src_offset, dst_stride)
-
-
-def mma_store_legacy(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
- """mma_store with apache-style signature.
-
- ``dst_ptr`` is typically a ``tvm_access_ptr`` Call (so the caller can
- encode the destination's element dtype + base offset), and
- ``src_ptr + src_offset`` is the raw warp accumulator + element offset.
- Codegen does ``ptr + offset`` C pointer arithmetic; lower_warp_memory
- rewrites src_offset's group component to a thread-local index."""
- return call_intrin(
- dtype,
- "tirx.mma_store_legacy",
- m,
- n,
- dst_ptr,
- src_ptr,
- src_offset,
- dst_stride,
- )
-
-
-def mma_fill(dtype, local_size, local_ptr, offset):
- """TVM intrinsic for zero-initalizing an MMA accumulation registor
-
- Parameters
- ----------
- dtype : str
- The data type of the result.
-
- local_size : IntImm
- The number of elements.
-
- local_ptr : Var
- The destination pointer variable.
-
- offset : Expr
- The destination offset.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(dtype, "tirx.mma_fill", local_size, local_ptr, offset)
-
-
-def mma_fill_legacy(dtype, local_size, local_ptr, offset):
- """mma_fill with (ptr_var, offset). Codegen emits ``ptr + offset``
- C pointer arithmetic; lower_warp_memory rewrites the offset's group
- component to a thread-local index."""
- return call_intrin(dtype, "tirx.mma_fill_legacy", local_size, local_ptr, offset)
-
-
def ptx_cp_async_bulk(
dtype, shared_ptr, shared_offset, global_ptr, global_offset, bytes, barrier_id
):
@@ -4949,11 +1543,44 @@ def ptx_mma_sp_legacy(*all_args):
)
+def mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
+ """Store the result of PTX MMA into a destination pointer."""
+
+ return call_intrin(dtype, "tirx.mma_store", m, n, dst_ptr, src_ptr, src_offset, dst_stride)
+
+
+def mma_store_legacy(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
+ """mma_store with apache-style pointer/offset semantics."""
+
+ return call_intrin(
+ dtype,
+ "tirx.mma_store_legacy",
+ m,
+ n,
+ dst_ptr,
+ src_ptr,
+ src_offset,
+ dst_stride,
+ )
+
+
+def mma_fill(dtype, local_size, local_ptr, offset):
+ """Zero-initialize an MMA accumulation register."""
+
+ return call_intrin(dtype, "tirx.mma_fill", local_size, local_ptr, offset)
+
+
+def mma_fill_legacy(dtype, local_size, local_ptr, offset):
+ """mma_fill with apache-style pointer/offset semantics."""
+
+ return call_intrin(dtype, "tirx.mma_fill_legacy", local_size, local_ptr, offset)
+
+
def ptx_ldmatrix(trans, num, dtype, smem_ptr, *dst_handles):
"""TVM intrinsic for ldmatrix.sync.aligned.m8n8.x{num}{.trans}.shared.{dtype}.
Mirrors the PTX ISA destination form: each output register is a separate
- operand. Pass ``T.address_of(buf[idx])`` (or ``buf.ptr_to([idx])``) for
+ operand. Pass ``T.address_of(buf[idx])`` (or ``buf.ptr_to([idx])``) for
each destination — the slots may be non-contiguous.
Parameters
@@ -5023,7 +1650,6 @@ def _wrap_or_fold_access_ptr(ptr, offset, elem_dtype):
passes that assume a raw buffer var. Detect that case and fold the
outer offset into the inner one.
"""
- from tvm.ir import Op # local import to avoid cycles
is_access_ptr_call = (
isinstance(ptr, Call) and isinstance(ptr.op, Op) and ptr.op.name == "tirx.tvm_access_ptr"
@@ -6211,44 +2837,6 @@ def ptx_tcgen05_commit(bar, cta_group=1, cta_mask=0, *, pred=None):
return call_intrin("", "tirx.ptx_tcgen05_commit", *args)
-def print_buffer(buffer_var, dtype, is_string, is_scalar, dim_num, *shape):
- """Print out buffer memory (tensor, string, or scalar) during runtime on cuda.
- This print function allows printing out buffer in tvm during runtime without
- dumping all the cuda code.
- Parameters
- ----------
- buffer_var : Var
- The data pointer of the buffer that needs to be printed out.
- dtype : DataType
- The data type of the buffer.
- is_string: Bool
- Whether the buffer is a string (dtype is Int8 by default in the backend).
- is_scalar: Bool
- Whether the buffer is a scalar.
- dim_num : Int
- The number of dimensions of the buffer
- *shape : Tuple
- The dimensions of the buffer in order.
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- final_shape_args = []
- if len(shape) == 1 and isinstance(shape[0], tuple | list | tvm.ir.Array):
- # Case 1: Called as print_buffer(..., dim, (s1, s2, ...))
- # The user provided a tuple/list as the single shape argument.
- final_shape_args = list(shape[0])
- else:
- # Case 2: Called as print_buffer(..., dim, s1, s2, ...)
- # This is how TVMScript parser will call it.
- final_shape_args = list(shape)
-
- return _ffi_api.print_buffer(
- buffer_var, dtype, is_string, is_scalar, dim_num, *final_shape_args
- )
-
-
def timer_init_cuda(profiler_buffer, profiler_tag, profiler_write_offset, num_groups, group_id):
"""TVM intrinsic for initializing the CUDA profiler, and store profiling result in a buffer.
@@ -7297,61 +3885,16 @@ def cuda_atomic_cas(ptr, old_val, new_val):
old_val: PrimExpr
The old value.
- new_val: PrimExpr
- The new value.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- old_val = tir.convert(old_val)
- return call_intrin(old_val.dtype, "tirx.cuda_atomic_cas", ptr, old_val, new_val)
-
-
-def thread_return():
- """TVM intrinsic to call thread_return()
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.thread_return")
-
-
-def continue_loop(span=None):
- """Create a tir intrinsic call to represent continue expression
-
- Parameters
- ----------
- span : Optional[Span]
- The location of this operator in the source code.
-
- Returns
- -------
- ret : PrimExpr
- The continue expression
- """
-
- return _ffi_api.continue_loop(span)
-
-
-def break_loop(span=None):
- """Create a tir intrinsic call to represent break expression
-
- Parameters
- ----------
- span : Optional[Span]
- The location of this operator in the source code.
+ new_val: PrimExpr
+ The new value.
Returns
-------
- ret : PrimExpr
- The break expression
+ call : PrimExpr
+ The call expression.
"""
-
- return _ffi_api.break_loop(span)
+ old_val = tir.convert(old_val)
+ return call_intrin(old_val.dtype, "tirx.cuda_atomic_cas", ptr, old_val, new_val)
########################################################
@@ -7738,475 +4281,3 @@ def nvshmem_barrier_all():
"""
return call_intrin("", "tirx.nvshmem_barrier_all")
-
-
-########################################################
-# NKI builtins
-########################################################
-
-
-def nki_load(res, data):
- """TVM intrinsic to call nki load instruction
-
- Parameters
- ----------
- res : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_load", res, data)
-
-
-def nki_store(res, data):
- """TVM intrinsic to call nki store instruction
-
- Parameters
- ----------
- res : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_store", res, data)
-
-
-def nki_tensor_copy(res, data):
- """TVM intrinsic to call nki tensor copy instruction
-
- Parameters
- ----------
- res : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_tensor_copy", res, data)
-
-
-def nki_matmul(res, lhs, rhs, accum=True):
- """TVM intrinsic to call nki matmul instruction
-
- Parameters
- ----------
- res : BufferLoad
- The result buffer.
-
- lhs: BufferLoad
- The left hand side buffer.
-
- rhs: BufferLoad
- The right hand side buffer.
-
- accum: bool
- Whether to accumulate the result.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_matmul", res, lhs, rhs, accum)
-
-
-def nki_activation(result, data, opcode, bias=0.0, scale=1.0):
- """TVM intrinsic to call nki activation instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- opcode: str
- The opcode.
-
- bias: PrimExpr
- The bias.
-
- scale: PrimExpr
- The scale.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_activation", result, data, opcode, bias, scale)
-
-
-def nki_reciprocal(result, data):
- """TVM intrinsic to call nki reciprocal instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_reciprocal", result, data)
-
-
-def nki_tensorreduce(result, data, opcode, negate, *axes):
- """TVM intrinsic to call nki tensorreduce instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- opcode: str
- The opcode.
-
- negate: bool
- Whether to negate the result.
-
- axes: Tuple[int]
- The axes to reduce over.
-
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_tensorreduce", result, data, opcode, negate, *axes)
-
-
-def nki_tensortensor(result, operand0, operand1, opcode):
- """TVM intrinsic to call nki tensortensor instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- operand0: BufferLoad
- The first operand buffer.
-
- operand1: BufferLoad
- The second operand buffer.
-
- opcode: str
- The opcode.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_tensortensor", result, operand0, operand1, opcode)
-
-
-def nki_tensorscalar(result, operand0, operand1, opcode, reverse=False):
- """TVM intrinsic to call nki tensorscalar instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- operand0: BufferLoad
- The first operand buffer.
-
- operand1: PrimExpr
- The second operand scalar.
-
- opcode: str
- The opcode.
-
- reverse: bool
- Whether to reverse the operands.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_tensorscalar", result, operand0, operand1, opcode, reverse)
-
-
-def nki_memset(result, value):
- """TVM intrinsic to call nki memset instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- value: PrimExpr
- The value to set.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_memset", result, value)
-
-
-def nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias=0.0, scale=1.0):
- """TVM intrinsic to call nki activation reduce instruction
-
- act_res = act_op(data * scale + bias)
- reduce_res = reduce_op(act_res)
-
- Parameters
- ----------
- reduce_res : BufferLoad
- The result buffer of reduction.
-
- act_res : BufferLoad
- The result buffer of activation.
-
- data: BufferLoad
- The data buffer.
-
- opcode: str
- The opcode.
-
- reduce_opcode: str
- The reduce opcode.
-
- bias: PrimExpr
- The bias.
-
- scale: PrimExpr
- The scale.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "",
- "tirx.nki_activation_reduce",
- reduce_res,
- act_res,
- data,
- opcode,
- reduce_opcode,
- bias,
- scale,
- )
-
-
-def nki_tensorscalar_reduce(
- reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse=False
-):
- """TVM intrinsic to call nki tensorscalar reduce instruction
-
- tensorscalar_res = tensorscalar_op(operand0, operand1)
- reduce_res = reduce_op(tensorscalar_res)
-
- Parameters
- ----------
- reduce_res : BufferLoad
- The result buffer of reduction.
-
- tensorscalar_res : BufferLoad
- The result buffer of tensorscalar.
-
- operand0: BufferLoad
- The first operand buffer.
-
- operand1: PrimExpr
- The second operand scalar.
-
- opcode: str
- The opcode.
-
- reduce_opcode: str
- The reduce opcode.
-
- reverse: bool
- Whether to reverse the operands of tensorscalar.
- """
- return call_intrin(
- "",
- "tirx.nki_tensorscalar_reduce",
- reduce_res,
- tensorscalar_res,
- operand0,
- operand1,
- opcode,
- reduce_opcode,
- reverse,
- )
-
-
-def nki_identity(result, size):
- """TVM intrinsic to call nki identity instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- size: PrimExpr
- The size of the identity tensor.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_identity", result, size)
-
-
-def nki_scalar_tensor_tensor(
- result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
-):
- """TVM intrinsic to call nki scalar tensor tensor instruction
- (data op0 operand0) op1 (operand1) , where op0 is tensor-scalar and op1 is tensor-tensor
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- operand0: PrimExpr
- The first operand scalar.
-
- operand1: BufferLoad
- The second operand buffer.
-
- opcode0: str
- The first opcode.
-
- opcode1: str
- The second opcode.
-
- reverse0: bool
- Whether to reverse the first operand.
-
- reverse1: bool
- Whether to reverse the second operand.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "",
- "tirx.nki_scalar_tensor_tensor",
- result,
- data,
- operand0,
- operand1,
- opcode0,
- opcode1,
- reverse0,
- reverse1,
- )
-
-
-def nki_scalar_tensor_scalar(
- result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
-):
- """TVM intrinsic to call nki scalar tensor scalar instruction
- (data op0 operand0) op1 (operand1) , where op0 and op1 are tensor-scalar
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- data: BufferLoad
- The data buffer.
-
- operand0: PrimExpr
- The first operand scalar.
-
- operand1: PrimExpr
- The second operand scalar.
-
- opcode0: str
- The first opcode.
-
- opcode1: str
- The second opcode.
-
- reverse0: bool
- Whether to reverse the first operand.
-
- reverse1: bool
- Whether to reverse the second operand.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin(
- "",
- "tirx.nki_scalar_tensor_scalar",
- result,
- data,
- operand0,
- operand1,
- opcode0,
- opcode1,
- reverse0,
- reverse1,
- )
-
-
-def nki_affine_select(result, pred, true_value, false_value):
- """TVM intrinsic to call nki affine select instruction
-
- Parameters
- ----------
- result : BufferLoad
- The result buffer.
-
- pred: PrimExpr
- The predicate.
-
- true_value: PrimExpr
- The true value.
-
- false_value: PrimExpr
- The false value.
-
- Returns
- -------
- call : PrimExpr
- The call expression.
- """
- return call_intrin("", "tirx.nki_affine_select", result, pred, true_value, false_value)
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/__init__.py b/python/tvm/backend/cuda/operator/__init__.py
similarity index 88%
copy from python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/__init__.py
copy to python/tvm/backend/cuda/operator/__init__.py
index e406e9c3fd..744c9bafc9 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/__init__.py
+++ b/python/tvm/backend/cuda/operator/__init__.py
@@ -14,5 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+"""CUDA backend operator registrations and helpers."""
-from .warp_xor_swizzle import *
+__all__ = ["intrinsics", "tile_primitive"]
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/__init__.py b/python/tvm/backend/cuda/operator/intrinsics/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/intrinsics/cuda/__init__.py
rename to python/tvm/backend/cuda/operator/intrinsics/__init__.py
diff --git a/python/tvm/tirx/operator/intrinsics/_schema.py b/python/tvm/backend/cuda/operator/intrinsics/_schema.py
similarity index 97%
rename from python/tvm/tirx/operator/intrinsics/_schema.py
rename to python/tvm/backend/cuda/operator/intrinsics/_schema.py
index 57e409e955..67c8b8cebd 100644
--- a/python/tvm/tirx/operator/intrinsics/_schema.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/_schema.py
@@ -43,8 +43,8 @@ from __future__ import annotations
from collections.abc import Callable
-from tvm.tirx.op import cuda_func_call
-from tvm.tirx.operator.intrinsics.cuda.registry import register_codegen
+from tvm.backend.cuda.op import cuda_func_call
+from tvm.backend.cuda.operator.intrinsics.registry import register_codegen
# C primitive type → TVM dtype string. Used when the caller specifies a
# non-void ``return_type`` but no explicit ``tvm_return_type`` — the helper
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/cp_async.py b/python/tvm/backend/cuda/operator/intrinsics/cp_async.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/cp_async.py
rename to python/tvm/backend/cuda/operator/intrinsics/cp_async.py
index 2eeb0821d6..3e6bc015e8 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/cp_async.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/cp_async.py
@@ -27,9 +27,9 @@ compatibility helpers.
"""
import tvm
-from tvm.tirx.op import cuda_func_call
+from tvm.backend.cuda.op import cuda_func_call
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .utils import parse_str
@@ -200,7 +200,8 @@ def codegen_ptx_cp_async(*args):
Accepts three call shapes (sorted by arity):
* 5 args ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size)`` —
- the legacy form emitted by ``s_tir/transform/InjectPTXAsyncCopy``.
+ the legacy form emitted by
+ ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``.
Offsets are folded into the pointers via ``tvm_access_ptr`` (in
bytes; offsets are pre-scaled by the pass) and the call is
forwarded with default cache / predicate / fill_mode.
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/header.py b/python/tvm/backend/cuda/operator/intrinsics/header.py
similarity index 100%
rename from python/tvm/tirx/operator/intrinsics/cuda/header.py
rename to python/tvm/backend/cuda/operator/intrinsics/header.py
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/math.py b/python/tvm/backend/cuda/operator/intrinsics/math.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/math.py
rename to python/tvm/backend/cuda/operator/intrinsics/math.py
index 37cd57d871..d93caa8145 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/math.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/math.py
@@ -26,9 +26,9 @@ CUDA side:
* warp / CTA reductions (templated butterfly shuffle-XOR).
"""
-from tvm.tirx.op import cuda_func_call
+from tvm.backend.cuda.op import cuda_func_call
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import register_codegen
from .utils import parse_str, validate_power_of_two_range
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/memory.py b/python/tvm/backend/cuda/operator/intrinsics/memory.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/memory.py
rename to python/tvm/backend/cuda/operator/intrinsics/memory.py
index 152e1434ca..b91a7945f2 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/memory.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/memory.py
@@ -33,9 +33,9 @@ CUDA side:
"""
from tvm import DataType
-from tvm.tirx.op import cuda_func_call
+from tvm.backend.cuda.op import cuda_func_call
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .utils import parse_str
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/misc.py b/python/tvm/backend/cuda/operator/intrinsics/misc.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/misc.py
rename to python/tvm/backend/cuda/operator/intrinsics/misc.py
index 0cca2cd194..2ee21ba57d 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/misc.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/misc.py
@@ -31,9 +31,9 @@ import hashlib
import json
import tvm
-from tvm.tirx.op import cuda_func_call
+from tvm.backend.cuda.op import cuda_func_call
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .utils import parse_str
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/mma.py b/python/tvm/backend/cuda/operator/intrinsics/mma.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/mma.py
rename to python/tvm/backend/cuda/operator/intrinsics/mma.py
index 7c59987368..5d6bccdcb4 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/mma.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/mma.py
@@ -32,7 +32,7 @@ from dataclasses import dataclass
from tvm import DataType
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .types import PTXDataType
from .utils import parse_str
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/nvshmem.py b/python/tvm/backend/cuda/operator/intrinsics/nvshmem.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/nvshmem.py
rename to python/tvm/backend/cuda/operator/intrinsics/nvshmem.py
index af7fa4c990..dc922a0e12 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/nvshmem.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/nvshmem.py
@@ -17,7 +17,7 @@
# pylint: disable=redefined-builtin, invalid-name
"""NVSHMEM intrinsics. Each backend call is one ``device_intrinsic(...)``."""
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
_NVSHMEM = ("nvshmem",)
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/registry.py b/python/tvm/backend/cuda/operator/intrinsics/registry.py
similarity index 100%
rename from python/tvm/tirx/operator/intrinsics/cuda/registry.py
rename to python/tvm/backend/cuda/operator/intrinsics/registry.py
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/sync.py b/python/tvm/backend/cuda/operator/intrinsics/sync.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/sync.py
rename to python/tvm/backend/cuda/operator/intrinsics/sync.py
index 4386336660..0fcdb31a46 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/sync.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/sync.py
@@ -32,8 +32,14 @@ CUDA-side helpers:
* warpgroup sync (``bar.sync``)
"""
-from .._common import CLUSTER_BARRIER_SEM, FENCE_PROXY_ASYNC_SPACE, FENCE_SCOPE, FENCE_SEM
-from .._schema import device_intrinsic
+from tvm.tirx.operator.intrinsics._common import (
+ CLUSTER_BARRIER_SEM,
+ FENCE_PROXY_ASYNC_SPACE,
+ FENCE_SCOPE,
+ FENCE_SEM,
+)
+
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .utils import parse_str
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/tcgen05.py b/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/tcgen05.py
rename to python/tvm/backend/cuda/operator/intrinsics/tcgen05.py
index ef30a85d0f..029a6bf87c 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/tcgen05.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py
@@ -25,7 +25,7 @@ shared parts callable.
import tvm
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .types import PTXDataType
from .utils import parse_str, validate_cta_group, validate_power_of_two_range
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/types.py b/python/tvm/backend/cuda/operator/intrinsics/types.py
similarity index 100%
rename from python/tvm/tirx/operator/intrinsics/cuda/types.py
rename to python/tvm/backend/cuda/operator/intrinsics/types.py
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/utils.py b/python/tvm/backend/cuda/operator/intrinsics/utils.py
similarity index 100%
rename from python/tvm/tirx/operator/intrinsics/cuda/utils.py
rename to python/tvm/backend/cuda/operator/intrinsics/utils.py
diff --git a/python/tvm/tirx/operator/intrinsics/cuda/wgmma.py b/python/tvm/backend/cuda/operator/intrinsics/wgmma.py
similarity index 99%
rename from python/tvm/tirx/operator/intrinsics/cuda/wgmma.py
rename to python/tvm/backend/cuda/operator/intrinsics/wgmma.py
index 87666db581..1ec52a5e4e 100644
--- a/python/tvm/tirx/operator/intrinsics/cuda/wgmma.py
+++ b/python/tvm/backend/cuda/operator/intrinsics/wgmma.py
@@ -25,7 +25,7 @@ device_intrinsic with an attr-driven ``c_signature`` callable.
import tvm
-from .._schema import device_intrinsic
+from ._schema import device_intrinsic
from .registry import CODEGEN_REGISTRY, register_codegen
from .types import PTXDataType
from .utils import parse_str
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/__init__.py
similarity index 89%
rename from python/tvm/tirx/operator/tile_primitive/cuda/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/__init__.py
index cea930c362..c4f37b0253 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/__init__.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/__init__.py
@@ -16,5 +16,9 @@
# under the License.
from .copy import *
+from .copy_async import *
from .elementwise import *
+from .gemm import *
+from .gemm_async import *
+from .permute_layout import *
from .reduction import *
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/common.py b/python/tvm/backend/cuda/operator/tile_primitive/common.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/common.py
rename to python/tvm/backend/cuda/operator/tile_primitive/common.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/_common.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/_common.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/_swizzle_iter.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/_swizzle_iter.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/fallback.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/fallback.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/gmem_smem.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/gmem_smem.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py
index aee24c62e4..727f7a18df 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/copy/gmem_smem.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py
@@ -87,7 +87,7 @@ def _divides_thread_cnt(
def _is_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_cuda():
+ if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
return False, f"unsupported exec_scope {sctx.scope_kind}"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/ld_stmatrix.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/ld_stmatrix.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py
index 75b0c9015b..96a4545c2b 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/copy/ld_stmatrix.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py
@@ -61,7 +61,7 @@ def _compute_r_perm(r):
def _is_ldstmatrix(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_cuda():
+ if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("warp", "warpgroup", "cta"):
return False, f"unsupported exec_scope {sctx.scope_kind} (need warp, warpgroup, or cta)"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/reg.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/reg.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py
index 5ea8d40e93..3f30a83c03 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/copy/reg.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py
@@ -177,7 +177,7 @@ def _s_side_slice_ok(op_call: TilePrimitiveCall) -> tuple[bool, str | None]:
def _is_reg_copy(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_cuda():
+ if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
return False, f"unsupported exec_scope {sctx.scope_kind}"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy/utils.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy/utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/dsmem.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/dsmem.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/ldgsts.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/ldgsts.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py
index 8c86f75ac8..e57228db6b 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/ldgsts.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py
@@ -89,7 +89,7 @@ def _divides_thread_cnt_ldgsts(
def _is_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_cuda():
+ if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
return False, f"unsupported exec_scope {sctx.scope_kind}"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tcgen05_cp.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py
index 3a9d819478..5cccc307f5 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py
@@ -408,7 +408,7 @@ def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> Pr
desc_buf = _get_or_create_desc(sctx, s_buf, LDO_field, SDO_field, sw)
t_addr = t_buf.allocated_addr
- from tvm.tirx.operator.tile_primitive.cuda.common import smem_desc_add_16B_offset
+ from tvm.backend.cuda.operator.tile_primitive.common import smem_desc_add_16B_offset
# Flatten the N-D middle iteration into a single T.unroll. Each iteration's
# per-dim index is (flat // stride) % extent, summed into the t/s offsets.
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tcgen05_ldst.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tcgen05_ldst.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tma.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/tma.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/copy_async/utils.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/copy_async/utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py
similarity index 95%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py
index 872cad0867..576dc437b9 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/__init__.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py
@@ -30,7 +30,7 @@ PTX/CUDA intrinsics live in ``vec_emit/`` (``binary_f32x2`` / ``cast_vec2``
from .register import *
# Suppress submodule-attribute leakage. Without an explicit ``__all__`` here,
-# ``from tvm.tirx.operator.tile_primitive.cuda.elementwise import *`` (run by
+# ``from tvm.backend.cuda.operator.tile_primitive.elementwise import *`` (run by
# tile_primitive/__init__.py) re-exports the implicit submodule attributes
# (``ops``, ``reg``, ``smem``, ``vec_emit``) — and ``ops`` in particular
# shadows the top-level ``tile_primitive/ops.py`` (BinaryReduce / UnaryReduce
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/_common.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/_common.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/binary.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/binary.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/cast.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/cast.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/cast.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/cast.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/fma.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/fma.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/fma.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/fma.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/unary.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/ops/unary.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/reg.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/reg.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py
index 063e4c3979..eddf9f3d8e 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/reg.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py
@@ -110,7 +110,7 @@ def is_reg_ewise(spec):
"""Predicate factory: dispatch accepted iff all operands in ``local`` scope."""
def check(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_cuda:
+ if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
return False, f"unsupported scope {sctx.scope_kind}"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/register.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/register.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/register.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/register.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/smem.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/smem.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py
index 3ac7405d77..c2cb5469be 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/smem.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py
@@ -60,7 +60,7 @@ def is_smem_ewise(spec):
"""Predicate factory: dispatch accepted iff all operands in ``shared*``."""
def check(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
- if not sctx.is_cuda:
+ if not sctx.is_target("cuda"):
return False, "non-cuda target"
if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
return False, f"unsupported scope {sctx.scope_kind}"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/binary_f32x2.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/binary_f32x2.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/cast_vec2.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/cast_vec2.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/fma_f32x2.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/elementwise/vec_emit/fma_f32x2.py
rename to python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/exec_scope_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/exec_scope_utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/gemm/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/gemm/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/gemm/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/gemm/mma_m16n8k_.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/gemm/mma_m16n8k_.py
rename to python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/gemm_async/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/gemm_async/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/gemm_async/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/gemm_async/tcgen05.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py
similarity index 83%
rename from python/tvm/tirx/operator/tile_primitive/cuda/gemm_async/tcgen05.py
rename to python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py
index c19bcda622..0afa042b5e 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py
@@ -46,7 +46,12 @@ from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt, TilePrimitiveCall
from ..common import get_st_extent, smem_desc_add_16B_offset
from ..exec_scope_utils import single_thread
-from ..tma_utils import SwizzleMode, mma_atom_layout, mma_atom_shape
+from ..tma_utils import (
+ SwizzleMode,
+ get_swizzle_mode_from_layout,
+ mma_atom_layout,
+ mma_atom_shape,
+)
# Mirror of ``format_map`` in the dense ``encode_instr_descriptor`` codegen
# (``python/tvm/tirx/operator/intrinsics/cuda/tcgen05.py``). Used to fold the
@@ -516,80 +521,152 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -
Tuple of (swizzle_mode, ldo, sdo, is_mn_major).
"""
region = list(buf_region.region)
- slice_layout = buf.layout.slice(buf.shape, region)
+
+ def _match(slice_layout, shape_2d):
+ """Match ``slice_layout`` (of ``shape_2d``) against the swizzle atoms.
+
+ Returns ``(swizzle_mode, ldo, sdo, is_mn_major)`` or ``None``.
+ """
+
+ def _try_atom(atom, atom_shape):
+ if any(s % a != 0 for s, a in zip(shape_2d, atom_shape)):
+ return None
+ atom_size = functools.reduce(operator.mul, atom_shape, 1)
+ tiler = atom.is_tile_inner(slice_layout, shape_2d, atom_shape)
+ if tiler is None:
+ return None
+ tiler_shape = [s // a for s, a in zip(shape_2d, atom_shape)]
+ tiler_grouped, seps = tiler.canonicalize().group(tiler_shape)
+ elem_per_128b = 128 // tvm.DataType(dtype).bits
+ ldo = (tiler_grouped.shard[-1].stride * atom_size) // elem_per_128b
+ sdo = (tiler_grouped.shard[-2].stride * atom_size) // elem_per_128b
+ return mode, ldo, sdo
+
+ for mode in (
+ SwizzleMode.SWIZZLE_128B_ATOM,
+ SwizzleMode.SWIZZLE_64B_ATOM,
+ SwizzleMode.SWIZZLE_32B_ATOM,
+ ):
+ swizzle_atom = mma_atom_layout(dtype, mode)
+ base_shape = mma_atom_shape(dtype, mode) # [8, T*s]
+ swapped_shape = [base_shape[1], base_shape[0]] # [T*s, 8]
+
+ # MN-major atom: compose SwizzleLayout with stride-reversed TileLayout
+ # so the first dim (T*s) is contiguous instead of the second.
+ # Needed when the penultimate dim is physically contiguous.
+ mn_tile = TileLayout(S[tuple(swapped_shape) : (1, swapped_shape[0])])
+ mn_atom = ComposeLayout(swizzle_atom, mn_tile)
+
+ # Determine K-major vs MN-major based on which dim is contiguous.
+ # K-major: K dim contiguous (last dim for [MN,K], first dim for [K,MN])
+ # MN-major: MN dim contiguous
+ #
+ # The plain swizzle_atom has last dim contiguous.
+ # The mn_atom has first dim contiguous.
+ #
+ # For non-transposed [MN, K]: K is last dim
+ # - K-major = swizzle_atom with [8, T*s] (K contiguous in last dim)
+ # - MN-major = mn_atom with [T*s, 8] (MN contiguous in first dim)
+ # For transposed [K, MN]: MN is last dim
+ # - K-major = mn_atom with [T*s, 8] (K contiguous in first dim)
+ # - MN-major = swizzle_atom with [8, T*s] (MN contiguous in last dim)
+ if is_transposed:
+ candidates = [
+ (False, mn_atom, swapped_shape), # K-major: K in first dim
+ (True, swizzle_atom, base_shape), # MN-major: MN in last dim
+ ]
+ else:
+ candidates = [
+ (False, swizzle_atom, base_shape), # K-major: K in last dim
+ (True, mn_atom, swapped_shape), # MN-major: MN in first dim
+ ]
+
+ for is_mn_major, atom, atom_shape in candidates:
+ result = _try_atom(atom, atom_shape)
+ if result is not None:
+ sw, ldo_val, sdo_val = result
+ # shard[-1] = last-dim groups, shard[-2] = first-dim groups.
+ # LBO strides MN-groups for MN-major, K-groups for K-major.
+ # Non-transposed [MN,K]: last=K, first=MN → swap for MN-major
+ # Transposed [K,MN]: last=MN, first=K → swap for K-major
+ if is_mn_major != is_transposed:
+ ldo_val, sdo_val = sdo_val, ldo_val
+ return sw, ldo_val, sdo_val, is_mn_major
+ return None
+
+ # The MMA SMEM descriptor describes the buffer's *physical* swizzle, which
+ # spans whole atoms and is a property of the buffer -- not of any sub-tile.
+ # Round the contiguous (innermost) axis up to a swizzle-atom multiple before
+ # matching, so the descriptor is always derived from a whole number of atoms:
+ # * for an already atom-aligned region (full-K, stride-axis P@V slices,
+ # non-swizzled buffers) this is a no-op -- desc_region == region, so the
+ # derived (swizzle, ldo, sdo, is_mn_major) are identical to matching the
+ # region directly;
+ # * for a sub-atom contiguous slice (fine K-major split-K, e.g.
+ # Asmem[..., :, lo:hi]) it rounds up to the smallest atom count covering
+ # the slice and describes it from the buffer origin.
+ # Either way the actual [lo:hi] range is addressed by K_iters + the per-MMA
+ # -tile 16B offset (from the *sliced* region in _a_operand / _b_desc_val),
+ # not by this descriptor. Verified hardware-correct for every MMA_K-aligned
+ # contiguous slice -- this is what enables fine K-major split-K.
+ cax = len(region) - 1 # innermost (contiguous) axis
+ elem_per_16b = 128 // DataType(dtype).bits
+ phys_mode = get_swizzle_mode_from_layout(buf.layout)
+ desc_region = list(region)
+ if phys_mode in (
+ SwizzleMode.SWIZZLE_128B_ATOM,
+ SwizzleMode.SWIZZLE_64B_ATOM,
+ SwizzleMode.SWIZZLE_32B_ATOM,
+ ):
+ atom_inner = mma_atom_shape(dtype, phys_mode)[-1]
+ contig = int(region[cax].extent)
+ rounded = ((contig + atom_inner - 1) // atom_inner) * atom_inner
+ if rounded != contig:
+ # Sub-atom contiguous slice. The per-tile 16B offset that locates the
+ # slice start must be exact, so the start has to sit on a 16B
+ # (= elem_per_16b element) boundary; otherwise reject rather than
+ # silently mis-address.
+ if not analyzer.can_prove_equal(
+ tvm.tirx.floormod(region[cax].min, elem_per_16b), 0
+ ):
+ raise ValueError(
+ f"gemm_async: contiguous-axis slice start {region[cax].min} is not "
+ f"16B-aligned (={elem_per_16b} elements, dtype {dtype}). A sub-atom "
+ f"contiguous slice must start on a 16B boundary; otherwise keep that "
+ f"axis full and split on a stride/outer axis instead (lay the operand "
+ f"out MN-major)."
+ )
+ desc_region[cax] = tvm.ir.Range.from_min_extent(0, rounded)
+
+ slice_layout = buf.layout.slice(buf.shape, desc_region)
# Strip unit dims to get the 2D matrix shape.
- shape_2d = [int(r.extent) for r in region if int(r.extent) != 1]
+ shape_2d = [int(r.extent) for r in desc_region if int(r.extent) != 1]
assert len(shape_2d) == 2, (
- f"Expected exactly 2 non-unit dims in region {[int(r.extent) for r in region]}"
+ f"Expected exactly 2 non-unit dims in region {[int(r.extent) for r in desc_region]}"
)
-
- def _try_atom(atom, atom_shape):
- if any(s % a != 0 for s, a in zip(shape_2d, atom_shape)):
- return None
- atom_size = functools.reduce(operator.mul, atom_shape, 1)
- tiler = atom.is_tile_inner(slice_layout, shape_2d, atom_shape)
- if tiler is None:
- return None
- tiler_shape = [s // a for s, a in zip(shape_2d, atom_shape)]
- tiler_grouped, seps = tiler.canonicalize().group(tiler_shape)
- elem_per_128b = 128 // tvm.DataType(dtype).bits
- ldo = (tiler_grouped.shard[-1].stride * atom_size) // elem_per_128b
- sdo = (tiler_grouped.shard[-2].stride * atom_size) // elem_per_128b
- return mode, ldo, sdo
-
- for mode in (
+ result = _match(slice_layout, shape_2d)
+ if result is not None:
+ return result
+
+ # Genuinely unsupported: the layout doesn't tile any swizzle atom even at
+ # full atom width. Actionable error (the old generic "no swizzle mode"
+ # message read like a hard limit and was mistaken for one).
+ hint = ""
+ if phys_mode in (
SwizzleMode.SWIZZLE_128B_ATOM,
SwizzleMode.SWIZZLE_64B_ATOM,
SwizzleMode.SWIZZLE_32B_ATOM,
):
- swizzle_atom = mma_atom_layout(dtype, mode)
- base_shape = mma_atom_shape(dtype, mode) # [8, T*s]
- swapped_shape = [base_shape[1], base_shape[0]] # [T*s, 8]
-
- # MN-major atom: compose SwizzleLayout with stride-reversed TileLayout
- # so the first dim (T*s) is contiguous instead of the second.
- # Needed when the penultimate dim is physically contiguous.
- mn_tile = TileLayout(S[tuple(swapped_shape) : (1, swapped_shape[0])])
- mn_atom = ComposeLayout(swizzle_atom, mn_tile)
-
- # Determine K-major vs MN-major based on which dim is contiguous.
- # K-major: K dim contiguous (last dim for [MN,K], first dim for [K,MN])
- # MN-major: MN dim contiguous
- #
- # The plain swizzle_atom has last dim contiguous.
- # The mn_atom has first dim contiguous.
- #
- # For non-transposed [MN, K]: K is last dim
- # - K-major = swizzle_atom with [8, T*s] (K contiguous in last dim)
- # - MN-major = mn_atom with [T*s, 8] (MN contiguous in first dim)
- # For transposed [K, MN]: MN is last dim
- # - K-major = mn_atom with [T*s, 8] (K contiguous in first dim)
- # - MN-major = swizzle_atom with [8, T*s] (MN contiguous in last dim)
- if is_transposed:
- candidates = [
- (False, mn_atom, swapped_shape), # K-major: K in first dim
- (True, swizzle_atom, base_shape), # MN-major: MN in last dim
- ]
- else:
- candidates = [
- (False, swizzle_atom, base_shape), # K-major: K in last dim
- (True, mn_atom, swapped_shape), # MN-major: MN in first dim
- ]
-
- for is_mn_major, atom, atom_shape in candidates:
- result = _try_atom(atom, atom_shape)
- if result is not None:
- sw, ldo_val, sdo_val = result
- # shard[-1] = last-dim groups, shard[-2] = first-dim groups.
- # LBO strides MN-groups for MN-major, K-groups for K-major.
- # Non-transposed [MN,K]: last=K, first=MN → swap for MN-major
- # Transposed [K,MN]: last=MN, first=K → swap for K-major
- if is_mn_major != is_transposed:
- ldo_val, sdo_val = sdo_val, ldo_val
- return sw, ldo_val, sdo_val, is_mn_major
-
+ atom_inner = mma_atom_shape(dtype, phys_mode)[-1]
+ atom_bytes = atom_inner * DataType(dtype).bits // 8
+ hint = (
+ f" The buffer is physically {phys_mode.name}-swizzled (contiguous atom "
+ f"= {atom_inner} elements / {atom_bytes} B); the region's layout does not "
+ f"tile it."
+ )
raise ValueError(
- f"No compatible swizzle mode found for dtype {dtype} with region shape {shape_2d}"
+ f"gemm_async: no MMA SMEM descriptor matches region shape {shape_2d} "
+ f"for dtype {dtype}.{hint}"
)
if a_is_tmem:
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/gemm_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/gemm_utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/layout_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/layout_utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/permute_layout/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/warp_xor_swizzle.py b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/permute_layout/warp_xor_swizzle.py
rename to python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/reduction/__init__.py
rename to python/tvm/backend/cuda/operator/tile_primitive/reduction/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/local.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/reduction/local.py
rename to python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py
index b05618f153..3316aaf2a1 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/local.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py
@@ -66,10 +66,10 @@ from tvm.script import tirx as T
from tvm.tirx import BufferRegion, PrimFunc
from tvm.tirx.layout import TileLayout, laneid
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import ReduceOpType
from ..common import get_indices, get_st_extent
from ..layout_utils import get_local_region, get_sublayout_from_region
from .utils import (
@@ -156,7 +156,7 @@ def validate_reduction_local(
dst_br, src_br = op.output, op.input
dst, src = dst_br.buffer, src_br.buffer
- if not (src.scope() == "local" and dst.scope() == "local" and sctx.is_cuda()):
+ if not (src.scope() == "local" and dst.scope() == "local" and sctx.is_target("cuda")):
return False, "expected local scope and CUDA target"
if src.dtype != dst.dtype:
return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/shared.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/reduction/shared.py
rename to python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py
index 8bee09ecc3..38986df86b 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/shared.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py
@@ -61,10 +61,10 @@ from tvm.arith.analyzer import Analyzer
from tvm.script import tirx as T
from tvm.tirx import BufferRegion, PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import ReduceOpType
from ..common import get_indices, get_st_extent, next_power_of_2
from .utils import (
_analyze_axes,
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/sm100_packed.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py
similarity index 99%
rename from python/tvm/tirx/operator/tile_primitive/cuda/reduction/sm100_packed.py
rename to python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py
index 5b8540ecac..fd506fa14b 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/sm100_packed.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py
@@ -48,10 +48,10 @@ import operator
from tvm.script import tirx as T
from tvm.tirx import BufferRegion, PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import ReduceOpType
from ..common import sm_version_ok
from ..exec_scope_utils import exec_scope_ok
from .utils import (
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/utils.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py
similarity index 97%
rename from python/tvm/tirx/operator/tile_primitive/cuda/reduction/utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py
index b53b5d1810..6248cddea1 100644
--- a/python/tvm/tirx/operator/tile_primitive/cuda/reduction/utils.py
+++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py
@@ -25,9 +25,9 @@ from tvm.arith.analyzer import Analyzer
from tvm.script import tirx as T
from tvm.tirx import BufferRegion
from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import ReduceOpType
from ..common import match_scope
reduce_op_table = {
@@ -250,7 +250,12 @@ def _local_scope_match(op: TilePrimitiveCall, sctx: DispatchContext):
op = TilePrimitiveCall.downcast(op)
src, dst = op.input.buffer, op.output.buffer
ok = all(
- [src.scope() == "local", dst.scope() == "local", src.dtype == dst.dtype, sctx.is_cuda()]
+ [
+ src.scope() == "local",
+ dst.scope() == "local",
+ src.dtype == dst.dtype,
+ sctx.is_target("cuda"),
+ ]
)
if not ok:
return (False, "src/dst must be local scope with matching dtype on CUDA")
diff --git a/python/tvm/tirx/operator/tile_primitive/cuda/tma_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/cuda/tma_utils.py
rename to python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py
diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py
new file mode 100644
index 0000000000..a1148f9b67
--- /dev/null
+++ b/python/tvm/backend/cuda/script.py
@@ -0,0 +1,571 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA TVMScript namespaces."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from tvm.backend.cuda import op as _cuda_op
+from tvm.tirx import Buffer
+from tvm.tirx import op as _tir_op
+from tvm.tirx.script.builder.ir import _dtype_forward, _op_wrapper
+
+# pylint: disable=protected-access
+
+
+def _ptx_ldg32(reg, guard, addr, local_addr):
+ if isinstance(addr, Buffer):
+ addr = addr[0]
+ return _tir_op.call_intrin(reg.dtype, "tirx.ptx.ldg32", reg, guard, addr, local_addr)
+
+
+_ptx_ldg32.__tir_op_name__ = "ptx.ldg32"
+
+
+class PTXNamespace:
+ """The PTX instruction submodule."""
+
+ def __init__(self):
+ self.ldg32 = _ptx_ldg32
+ self.ldmatrix = _dtype_forward(_cuda_op.ptx_ldmatrix)
+ # Apache-compatible variant. Same lowered intrinsic as
+ # ``ldmatrix`` but accepts the historical ``(trans, num, dtype,
+ # local_ptr, local_offset, smem_ptr, smem_offset)`` form. Coexists
+ # with the fork-native version so upstream-derived tests keep
+ # working without rewriting their tirx code.
+ self.ldmatrix_legacy = _dtype_forward(_cuda_op.ptx_ldmatrix_legacy)
+ self.stmatrix = _op_wrapper(_cuda_op.ptx_stmatrix)
+ self.setmaxnreg: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_setmaxnreg)
+ self.elect_sync: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_elect_sync)
+ self.fetch_register: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_fetch_register)
+ self.ld = _op_wrapper(_cuda_op.ptx_ld)
+ self.ld_acquire = _op_wrapper(_cuda_op.ptx_ld_acquire)
+ self.ld_volatile = _op_wrapper(_cuda_op.ptx_ld_volatile)
+ self.ld_global_acquire = _op_wrapper(_cuda_op.ptx_ld_global_acquire)
+ self.red_scalar = _op_wrapper(_cuda_op.ptx_red_scalar)
+ self.atom_scalar = _op_wrapper(_cuda_op.ptx_atom_scalar)
+ self.prefetch_tensormap = _op_wrapper(_cuda_op.ptx_prefetch_tensormap)
+ self.mbarrier_test_wait_parity = _op_wrapper(_cuda_op.ptx_mbarrier_test_wait_parity)
+ self.cp_async_bulk_g2s_cta = _op_wrapper(_cuda_op.ptx_cp_async_bulk_g2s_cta)
+ self.cp_async_bulk_g2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_g2s_cluster)
+ self.cp_async_bulk_s2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2s_cluster)
+ self.cp_async_bulk_s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2g)
+ self.st = _op_wrapper(_cuda_op.ptx_st)
+ self.st_bulk = _op_wrapper(_cuda_op.ptx_st_bulk)
+ self.fns_b32 = _op_wrapper(_cuda_op.ptx_fns_b32)
+ self.add_rn_f32_bf16 = _op_wrapper(_cuda_op.ptx_add_rn_f32_bf16)
+ self.mapa = _op_wrapper(_cuda_op.ptx_mapa)
+ self.map_shared_rank = _op_wrapper(_cuda_op.ptx_map_shared_rank)
+ self.any_sync = _op_wrapper(_cuda_op.ptx_any_sync)
+ # Math operations
+ self.exp2 = _op_wrapper(_cuda_op.ptx_exp2)
+ self.rcp = _op_wrapper(_cuda_op.ptx_rcp)
+ self.reduce3_min_f32 = _op_wrapper(_cuda_op.ptx_reduce3_min_f32)
+ self.reduce3_max_f32 = _op_wrapper(_cuda_op.ptx_reduce3_max_f32)
+ # add/sub/mul/fma DPS form: (d_addr, a, b[, c], *, rounding, ftz[, sat])
+ self.add_f32 = _op_wrapper(_cuda_op.ptx_add_f32)
+ self.add_f32x2 = _op_wrapper(_cuda_op.ptx_add_f32x2)
+ self.add_f64 = _op_wrapper(_cuda_op.ptx_add_f64)
+ self.sub_f32 = _op_wrapper(_cuda_op.ptx_sub_f32)
+ self.sub_f32x2 = _op_wrapper(_cuda_op.ptx_sub_f32x2)
+ self.sub_f64 = _op_wrapper(_cuda_op.ptx_sub_f64)
+ self.mul_f32 = _op_wrapper(_cuda_op.ptx_mul_f32)
+ self.mul_f32x2 = _op_wrapper(_cuda_op.ptx_mul_f32x2)
+ self.mul_f64 = _op_wrapper(_cuda_op.ptx_mul_f64)
+ self.fma_f32 = _op_wrapper(_cuda_op.ptx_fma_f32)
+ self.fma_f32x2 = _op_wrapper(_cuda_op.ptx_fma_f32x2)
+ self.fma_f64 = _op_wrapper(_cuda_op.ptx_fma_f64)
+ self.max_f32 = _op_wrapper(_cuda_op.ptx_max_f32)
+ self.mma = MmaNamespace()
+ self.cp_async = CpAsyncNamespace()
+ self.wgmma = WgmmaNamespace()
+ self.mbarrier = MbarrierNamespace()
+ self.tcgen05 = Tcgen05Namespace()
+ self.bar = BarNamespace()
+ self.barrier = BarrierNamespace()
+ self.fence = FenceNamespace()
+ self.griddepcontrol = GriddepcontrolNamespace()
+
+
+class MmaNamespace:
+ """The MMA instruction submodule."""
+
+ def __init__(self):
+ self.sp = _dtype_forward(_cuda_op.ptx_mma_sp)
+ # Apache-compatible variant of ptx_mma. Coexists with the
+ # fork-native ``__call__`` form (``T.ptx.mma(...)``).
+ self.legacy = _dtype_forward(_cuda_op.ptx_mma_legacy)
+ # __call__ corresponds to ptx_mma
+ self.__tir_call_op_name__ = "ptx_mma"
+
+ def __call__(self, *args, **kwds):
+ return _dtype_forward(_cuda_op.ptx_mma)(*args, **kwds)
+
+
+class CpAsyncNamespace:
+ """The CpAsync instruction submodule."""
+
+ def __init__(self):
+ self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_commit_group)
+ self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_wait_group)
+ # Legacy variant: takes (dst_ptr, dst_offset, src_ptr, src_offset,
+ # cp_size). Offsets are folded into the pointers; coexists with
+ # the fork-native ``__call__`` form.
+ self.legacy = _dtype_forward(_cuda_op.ptx_cp_async_legacy)
+ self.bulk = CpAsyncBulkNamespace()
+ self.mbarrier = CpAsyncMbarrierNamespace()
+
+ def __call__(self, *args, **kwds):
+ # Accept the legacy 6-arg form ``(elem_dtype, dst, dst_off, src,
+ # src_off, cp_size)`` that the printer round-trips for the raw
+ # ``tirx.ptx_cp_async`` Call emitted by
+ # ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``. The pass-emitted
+ # Call has 5 args (no ``tvm_access_ptr`` fold) and a
+ # per-element-dtype Call.dtype, so build it directly.
+ if len(args) == 6 and isinstance(args[0], str) and "dtype" not in kwds:
+ import tvm
+
+ elem_dtype, dst, dst_off, src, src_off, cp_size = args
+ return tvm.tirx.Call(
+ tvm.DataType(elem_dtype),
+ tvm.ir.Op.get("tirx.ptx_cp_async"),
+ [dst, dst_off, src, src_off, cp_size],
+ )
+ return _dtype_forward(_cuda_op.ptx_cp_async)(*args, **kwds)
+
+ # __call__ corresponds to ptx_cp_async
+ __tir_call_op_name__ = "ptx_cp_async"
+
+
+class CpAsyncBulkNamespace:
+ """The CpAsyncBulk instruction submodule."""
+
+ def __init__(self):
+ self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_bulk_commit_group)
+ self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_bulk_wait_group)
+ self.tensor = CpAsyncBulkTensorNamespace()
+ self.s2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_shared_to_cluster)
+
+ def __call__(self, *args, **kwds):
+ return _dtype_forward(_cuda_op.ptx_cp_async_bulk)(*args, **kwds)
+
+ # __call__ corresponds to ptx_cp_async_bulk
+ __tir_call_op_name__ = "ptx_cp_async_bulk"
+
+
+class CpAsyncBulkTensorNamespace:
+ """The CpAsyncBulkTensor instruction submodule."""
+
+ def __init__(self):
+ self.g2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster)
+ self.g2c_tile_gather4 = _op_wrapper(
+ _cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster
+ )
+ self.s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global)
+ self.s2g_reduce = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global_reduce)
+ self.g2c_prefetch = _op_wrapper(
+ _cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch
+ )
+
+ @staticmethod
+ def g2c_bar_addr(
+ dim,
+ dst_ptr,
+ bar_addr,
+ tensormap_addr,
+ cta_mask,
+ cta_group,
+ cache_hint,
+ *coords,
+ cache_policy=None,
+ ):
+ _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP)
+ cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy)
+ return _tir_op.call_intrin(
+ "",
+ "tirx.ptx_cp_async_bulk_tensor_global_to_cluster",
+ dim,
+ dst_ptr,
+ bar_addr,
+ tensormap_addr,
+ cta_mask,
+ cta_group,
+ cache_policy,
+ int(has_cache_policy),
+ 1,
+ *coords,
+ )
+
+ @staticmethod
+ def g2c_tile_gather4_bar_addr(
+ dim,
+ dst_ptr,
+ bar_addr,
+ tensormap_addr,
+ cta_mask,
+ cta_group,
+ cache_hint,
+ *coords,
+ cache_policy=None,
+ ):
+ _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP)
+ cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy)
+ return _tir_op.call_intrin(
+ "",
+ "tirx.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster",
+ dim,
+ dst_ptr,
+ bar_addr,
+ tensormap_addr,
+ cta_mask,
+ cta_group,
+ cache_policy,
+ int(has_cache_policy),
+ 1,
+ *coords,
+ )
+
+
+class CpAsyncMbarrierNamespace:
+ """The CpAsyncMbarrier instruction submodule."""
+
+ def __init__(self):
+ self.arrive = _op_wrapper(_cuda_op.ptx_cp_async_mbarrier_arrive)
+
+
+class WgmmaNamespace:
+ """The WGMMA instruction submodule."""
+
+ def __init__(self):
+ self.fence: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_wgmma_fence)
+ self.commit_group = _op_wrapper(_cuda_op.ptx_wgmma_commit_group)
+ self.wait_group = _op_wrapper(_cuda_op.ptx_wgmma_wait_group)
+ self.noop_barrier = _op_wrapper(_cuda_op.ptx_wgmma_noop_barrier)
+ self.mma_async = WgmmaMmaAsyncNamespace()
+ self.encode_matrix_descriptor = _op_wrapper(_cuda_op.ptx_wgmma_encode_matrix_descriptor)
+
+
+class WgmmaMmaAsyncNamespace:
+ """The WGMMA MMAAsync instruction submodule."""
+
+ def __init__(self):
+ self.ss = _op_wrapper(_cuda_op.ptx_wgmma_mma_async_ss)
+ self.rs = _op_wrapper(_cuda_op.ptx_wgmma_mma_async_rs)
+
+
+class MbarrierNamespace:
+ """The Mbarrier instruction submodule."""
+
+ def __init__(self):
+ self.init = _op_wrapper(_cuda_op.ptx_mbarrier_init)
+ self.try_wait = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait)
+ self.try_wait_once = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_once)
+ self.arrive = MbarrierArriveNamespace()
+
+
+class MbarrierArriveNamespace:
+ """The Mbarrier Arrive instruction submodule."""
+
+ def __init__(self):
+ self.expect_tx = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_expect_tx)
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.ptx_mbarrier_arrive)(*args, **kwds)
+
+ # __call__ corresponds to ptx_mbarrier_arrive
+ __tir_call_op_name__ = "ptx_mbarrier_arrive"
+
+
+class Tcgen05Namespace:
+ """The Tcgen05 instruction submodule."""
+
+ def __init__(self):
+ self.alloc = _op_wrapper(_cuda_op.ptx_tcgen05_alloc)
+ self.dealloc = _op_wrapper(_cuda_op.ptx_tcgen05_dealloc)
+ self.relinquish_alloc_permit = _op_wrapper(_cuda_op.ptx_tcgen05_relinquish_alloc_permit)
+ self.encode_matrix_descriptor = _op_wrapper(_cuda_op.ptx_tcgen05_encode_matrix_descriptor)
+ self.encode_instr_descriptor = _op_wrapper(_cuda_op.ptx_tcgen05_encode_instr_descriptor)
+ self.encode_instr_descriptor_block_scaled = _op_wrapper(
+ _cuda_op.ptx_tcgen05_encode_instr_descriptor_block_scaled
+ )
+ self.ld = _op_wrapper(_cuda_op.ptx_tcgen05_ld)
+ self.st = _op_wrapper(_cuda_op.ptx_tcgen05_st)
+ self.cp = _op_wrapper(_cuda_op.ptx_tcgen05_cp)
+ self.shift = _op_wrapper(_cuda_op.ptx_tcgen05_shift)
+ self.commit = _op_wrapper(_cuda_op.ptx_tcgen05_commit)
+ self.wait = Tcgen05WaitNamespace()
+ self.mma = Tcgen05MmaNamespace()
+ self.fence = Tcgen05FenceNamespace()
+
+
+class Tcgen05FenceNamespace:
+ """The Tcgen05 Fence instruction submodule."""
+
+ def __init__(self):
+ self.before_thread_sync = _op_wrapper(_cuda_op.ptx_tcgen05_fence_before_thread_sync)
+ self.after_thread_sync = _op_wrapper(_cuda_op.ptx_tcgen05_fence_after_thread_sync)
+
+
+class Tcgen05MmaNamespace:
+ """The Tcgen05 MMA instruction submodule."""
+
+ def __init__(self):
+ self.block_scale = _op_wrapper(_cuda_op.ptx_tcgen05_mma_block_scale)
+ self.sp = Tcgen05MmaSpNamespace()
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.ptx_tcgen05_mma)(*args, **kwds)
+
+ # __call__ corresponds to ptx_tcgen05_mma
+ __tir_call_op_name__ = "ptx_tcgen05_mma"
+
+
+class Tcgen05MmaSpNamespace:
+ """Tcgen05 Sparse MMA instruction submodule."""
+
+ def __init__(self):
+ self.block_scale = _op_wrapper(_cuda_op.ptx_tcgen05_mma_sp_block_scale)
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.ptx_tcgen05_mma_sp)(*args, **kwds)
+
+ # __call__ corresponds to ptx_tcgen05_mma_sp
+ __tir_call_op_name__ = "ptx_tcgen05_mma_sp"
+
+
+class Tcgen05WaitNamespace:
+ """The Tcgen05 Wait instruction submodule."""
+
+ def __init__(self):
+ self.ld = _op_wrapper(_cuda_op.ptx_tcgen05_wait_ld)
+ self.st = _op_wrapper(_cuda_op.ptx_tcgen05_wait_st)
+
+
+class BarNamespace:
+ """The Bar instruction submodule."""
+
+ def __init__(self):
+ self.arrive = _op_wrapper(_cuda_op.ptx_bar_arrive)
+ self.sync = _op_wrapper(_cuda_op.ptx_bar_sync)
+
+
+class BarrierNamespace:
+ """The Barrier instruction submodule."""
+
+ def __init__(self):
+ self.cluster = BarrierClusterNamespace()
+
+
+class BarrierClusterNamespace:
+ """The BarrierCluster instruction submodule."""
+
+ def __init__(self):
+ self.arrive = _op_wrapper(_cuda_op.ptx_barrier_cluster_arrive)
+ self.wait = _op_wrapper(_cuda_op.ptx_barrier_cluster_wait)
+
+
+class FenceNamespace:
+ """PTX fence instruction submodule."""
+
+ def __init__(self):
+ self.proxy_async = _op_wrapper(_cuda_op.ptx_fence_proxy_async)
+ self.mbarrier_init = _op_wrapper(_cuda_op.ptx_fence_mbarrier_init)
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.ptx_fence)(*args, **kwds)
+
+ __tir_call_op_name__ = "ptx_fence"
+
+
+class GriddepcontrolNamespace:
+ """PTX griddepcontrol instruction submodule (sm_90+)."""
+
+ def __init__(self):
+ self.wait = _op_wrapper(_cuda_op.ptx_griddepcontrol_wait)
+ self.launch_dependents = _op_wrapper(_cuda_op.ptx_griddepcontrol_launch_dependents)
+
+
+class CUDANamespace:
+ """The CUDA intrinsics submodule."""
+
+ def __init__(self):
+ self.atomic_add = _op_wrapper(_cuda_op.cuda_atomic_add)
+ self.thread_fence = _op_wrapper(_cuda_op.cuda_thread_fence)
+ self.warpgroup_sync = _op_wrapper(_cuda_op.cuda_warpgroup_sync)
+ self.warp_sync = _op_wrapper(_cuda_op.cuda_warp_sync)
+ self.warp_reduce = _op_wrapper(_cuda_op.cuda_warp_reduce)
+ self.warp_sum = _op_wrapper(_cuda_op.cuda_warp_sum)
+ self.warp_max = _op_wrapper(_cuda_op.cuda_warp_max)
+ self.warp_min = _op_wrapper(_cuda_op.cuda_warp_min)
+ self.cta_reduce = _op_wrapper(_cuda_op.cuda_cta_reduce)
+ self.cta_sum = _op_wrapper(_cuda_op.cuda_cta_sum)
+ self.cta_max = _op_wrapper(_cuda_op.cuda_cta_max)
+ self.cta_min = _op_wrapper(_cuda_op.cuda_cta_min)
+ self.copy_bytes = _op_wrapper(_cuda_op.cuda_copy_bytes)
+ self.copy_128b = _op_wrapper(_cuda_op.cuda_copy_128b)
+ self.copy_64b = _op_wrapper(_cuda_op.cuda_copy_64b)
+ self.copy_32b = _op_wrapper(_cuda_op.cuda_copy_32b)
+ self.copy_16b = _op_wrapper(_cuda_op.cuda_copy_16b)
+ self.copy_8b = _op_wrapper(_cuda_op.cuda_copy_8b)
+ self.cta_sync = _op_wrapper(_cuda_op.cuda_cta_sync)
+ self.grid_sync = _op_wrapper(_cuda_op.cuda_grid_sync)
+ self.cluster_sync = _op_wrapper(_cuda_op.cuda_cluster_sync)
+ self.thread_rank = _op_wrapper(_cuda_op.cuda_thread_rank)
+ self.trap_when_assert_failed = _op_wrapper(_cuda_op.cuda_trap_when_assert_failed)
+ self.runtime_instr_desc = _op_wrapper(_cuda_op.cuda_runtime_instr_desc)
+ self.half2float = _op_wrapper(_cuda_op.cuda_half2float)
+ self.bfloat162float = _op_wrapper(_cuda_op.cuda_bfloat162float)
+ self.float22half2 = _op_wrapper(_cuda_op.cuda_float22half2)
+ self.half8tofloat8 = _op_wrapper(_cuda_op.cuda_half8tofloat8)
+ self.float8tohalf8 = _op_wrapper(_cuda_op.cuda_float8tohalf8)
+ self.syncthreads_and = _op_wrapper(_cuda_op.cuda_syncthreads_and)
+ self.syncthreads_or = _op_wrapper(_cuda_op.cuda_syncthreads_or)
+ self.nano_sleep = _op_wrapper(_cuda_op.cuda_nano_sleep)
+ self.atomic_cas = _op_wrapper(_cuda_op.cuda_atomic_cas)
+ self.func_call = _op_wrapper(_cuda_op.cuda_func_call)
+ self.printf = _op_wrapper(_cuda_op.cuda_printf)
+ self.ldg = _op_wrapper(_cuda_op.cuda_ldg)
+ self.get_tmem_addr = _op_wrapper(_cuda_op.cuda_get_tmem_addr)
+ self.cvta_generic_to_shared = _op_wrapper(_cuda_op.cuda_cvta_generic_to_shared)
+ self.smem_addr_from_uint64 = _op_wrapper(_cuda_op.cuda_smem_addr_from_uint64)
+ self.sm100_tma_2sm_mbarrier_addr = _op_wrapper(_cuda_op.cuda_sm100_tma_2sm_mbarrier_addr)
+ self.uint_as_float = _op_wrapper(_cuda_op.cuda_uint_as_float)
+ self.float_as_uint = _op_wrapper(_cuda_op.cuda_float_as_uint)
+ self.ballot_sync = _op_wrapper(_cuda_op.cuda_ballot_sync)
+ self.ffs_u32 = _op_wrapper(_cuda_op.cuda_ffs_u32)
+ self.reduce_add_sync_u32 = _op_wrapper(_cuda_op.cuda_reduce_add_sync_u32)
+ self.reduce_min_sync_u32 = _op_wrapper(_cuda_op.cuda_reduce_min_sync_u32)
+ self.clock64 = _op_wrapper(_cuda_op.cuda_clock64)
+ self.make_float2 = _op_wrapper(_cuda_op.cuda_make_float2)
+ self.float2_x = _op_wrapper(_cuda_op.cuda_float2_x)
+ self.float2_y = _op_wrapper(_cuda_op.cuda_float2_y)
+ self.fmul2_rn = _op_wrapper(_cuda_op.cuda_fmul2_rn)
+ self.fadd2_rn = _op_wrapper(_cuda_op.cuda_fadd2_rn)
+ self.float22bfloat162_rn = _op_wrapper(_cuda_op.cuda_float22bfloat162_rn)
+ self.float22bfloat162_rn_from_float2 = _op_wrapper(
+ _cuda_op.cuda_float22bfloat162_rn_from_float2
+ )
+ self.bfloat1622float2 = _op_wrapper(_cuda_op.cuda_bfloat1622float2)
+ self.hmin2 = _op_wrapper(_cuda_op.cuda_hmin2)
+ self.hmax2 = _op_wrapper(_cuda_op.cuda_hmax2)
+ self.fp8x4_e4m3_from_float4 = _op_wrapper(_cuda_op.cuda_fp8x4_e4m3_from_float4)
+ self.timer_init = _op_wrapper(_cuda_op.timer_init_cuda)
+ self.timer_start = _op_wrapper(_cuda_op.timer_start_cuda)
+ self.timer_end = _op_wrapper(_cuda_op.timer_end_cuda)
+ self.timer_finalize = _op_wrapper(_cuda_op.timer_finalize_cuda)
+ self.mma_store = _dtype_forward(_cuda_op.mma_store)
+ self.mma_fill = _dtype_forward(_cuda_op.mma_fill)
+ self.mma_store_legacy = _dtype_forward(_cuda_op.mma_store_legacy)
+ self.mma_fill_legacy = _dtype_forward(_cuda_op.mma_fill_legacy)
+ setattr(self, "__shfl_sync", self._shfl_sync)
+ setattr(self, "__shfl_up_sync", self._shfl_up_sync)
+ setattr(self, "__shfl_down_sync", self._shfl_down_sync)
+ setattr(self, "__shfl_xor_sync", self._shfl_xor_sync)
+ setattr(self, "__activemask", self._activemask)
+
+ @staticmethod
+ def _shfl_sync(mask, var, lane, width):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(var.dtype, "tirx.cuda.__shfl_sync", mask, var, lane, width)
+
+ @staticmethod
+ def _shfl_up_sync(mask, var, delta, width):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(var.dtype, "tirx.cuda.__shfl_up_sync", mask, var, delta, width)
+
+ @staticmethod
+ def _shfl_down_sync(mask, var, delta, width):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(var.dtype, "tirx.cuda.__shfl_down_sync", mask, var, delta, width)
+
+ @staticmethod
+ def _shfl_xor_sync(mask, var, lane_mask, width):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(
+ var.dtype, "tirx.cuda.__shfl_xor_sync", mask, var, lane_mask, width
+ )
+
+ @staticmethod
+ def _activemask():
+ return _tir_op.call_intrin("uint32", "tirx.cuda.__activemask")
+
+
+class NVSHMEMNamespace:
+ """The NVSHMEM intrinsics submodule."""
+
+ def __init__(self):
+ self.my_pe = _op_wrapper(_cuda_op.nvshmem_my_pe)
+ self.n_pes = _op_wrapper(_cuda_op.nvshmem_n_pes)
+ self.signal_op = _op_wrapper(_cuda_op.nvshmem_signal_op)
+ self.wait_until = _op_wrapper(_cuda_op.nvshmem_wait_until)
+ self.quiet = _op_wrapper(_cuda_op.nvshmem_quiet)
+ self.fence = _op_wrapper(_cuda_op.nvshmem_fence)
+ self.barrier_all = _op_wrapper(_cuda_op.nvshmem_barrier_all)
+ self.getmem_nbi = NVSHMEMGetMemNBINamespace()
+ self.putmem_nbi = NVSHMEMPutMemNBINamespace()
+ self.putmem_signal_nbi = NVSHMEMPutMemSignalNBINamespace()
+
+
+class NVSHMEMGetMemNBINamespace:
+ """The NVSHMEM GetMemNBI intrinsics submodule."""
+
+ def __init__(self):
+ self.warp = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_warp)
+ self.block = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_block)
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.nvshmem_getmem_nbi)(*args, **kwds)
+
+ # __call__ corresponds to nvshmem_getmem_nbi
+ __tir_call_op_name__ = "nvshmem_getmem_nbi"
+
+
+class NVSHMEMPutMemNBINamespace:
+ """The NVSHMEM PutMemNBI intrinsics submodule."""
+
+ def __init__(self):
+ self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_warp)
+ self.block = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_block)
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.nvshmem_putmem_nbi)(*args, **kwds)
+
+ # __call__ corresponds to nvshmem_putmem_nbi
+ __tir_call_op_name__ = "nvshmem_putmem_nbi"
+
+
+class NVSHMEMPutMemSignalNBINamespace:
+ """The NVSHMEM PutMemSignalNBI intrinsics submodule."""
+
+ def __init__(self):
+ self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_warp)
+ self.block = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_block)
+
+ def __call__(self, *args, **kwds):
+ return _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi)(*args, **kwds)
+
+ # __call__ corresponds to nvshmem_putmem_signal_nbi
+ __tir_call_op_name__ = "nvshmem_putmem_signal_nbi"
+
+
+__all__ = ["CUDANamespace", "NVSHMEMNamespace", "PTXNamespace"]
diff --git a/python/tvm/target/tag_registry/cuda.py b/python/tvm/backend/cuda/target_tags.py
similarity index 99%
rename from python/tvm/target/tag_registry/cuda.py
rename to python/tvm/backend/cuda/target_tags.py
index 6b1bd9e8a8..87c83b378f 100644
--- a/python/tvm/target/tag_registry/cuda.py
+++ b/python/tvm/backend/cuda/target_tags.py
@@ -16,7 +16,7 @@
# under the License.
"""NVIDIA CUDA target tags."""
-from .registry import register_tag
+from tvm.target import register_tag
def _register_cuda_tag(name, arch, shared_mem=49152, regs=65536, **extra):
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py b/python/tvm/backend/hexagon/__init__.py
similarity index 65%
copy from python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py
copy to python/tvm/backend/hexagon/__init__.py
index f7a7b886d0..3852e36ccc 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py
+++ b/python/tvm/backend/hexagon/__init__.py
@@ -14,20 +14,22 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+"""Hexagon-owned backend hooks."""
-"""Reduction dispatch variant registrations."""
+from importlib import import_module
-from tvm.tirx.operator.tile_primitive import register_dispatch
+_LAZY_SUBMODULES = {"target_tags"}
-from ...common import ReduceOpType
-from .utils import reduction_trn
-for _op_name, _op_type in {
- "sum": ReduceOpType.SUM,
- "max": ReduceOpType.MAX,
- "min": ReduceOpType.MIN,
-}.items():
+def register_backend():
+ """Register Hexagon-owned Python semantics."""
+ import_module(f"{__name__}.target_tags")
- @register_dispatch(_op_name, "trn", variant="reduction", priority=0)
- def _reduction_dispatch(op, sctx, _ty=_op_type):
- return reduction_trn(op, _ty, sctx)
+
+def __getattr__(name: str):
+ if name in _LAZY_SUBMODULES:
+ return import_module(f"{__name__}.{name}")
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["register_backend", "target_tags"]
diff --git a/python/tvm/target/tag_registry/hexagon.py b/python/tvm/backend/hexagon/target_tags.py
similarity index 98%
rename from python/tvm/target/tag_registry/hexagon.py
rename to python/tvm/backend/hexagon/target_tags.py
index a43c3fbd93..645567cabc 100644
--- a/python/tvm/target/tag_registry/hexagon.py
+++ b/python/tvm/backend/hexagon/target_tags.py
@@ -16,7 +16,7 @@
# under the License.
"""Qualcomm Hexagon target tags."""
-from .registry import register_tag
+from tvm.target import register_tag
_ONE_MB = 2**20
diff --git a/python/tvm/backend/metal/__init__.py b/python/tvm/backend/metal/__init__.py
new file mode 100644
index 0000000000..d42806433f
--- /dev/null
+++ b/python/tvm/backend/metal/__init__.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal-owned TIRx modules."""
+
+from importlib import import_module
+
+_LAZY_SUBMODULES = {"op", "script", "target_tags"}
+
+
+def register_backend():
+ """Register Metal-owned Python semantics."""
+ from tvm.tirx.script.builder import ir as builder_ir # pylint: disable=import-outside-toplevel
+
+ for name, namespace in script_namespaces().items():
+ builder_ir.register_script_namespace(name, namespace)
+ import_module(f"{__name__}.target_tags")
+
+
+def script_namespaces(**_):
+ """Return Metal-owned TVMScript namespaces."""
+ from .script import MetalNamespace # pylint: disable=import-outside-toplevel
+
+ return {"metal": MetalNamespace()}
+
+
+def script_namespace(**kwargs):
+ """Return the Metal TVMScript namespace object."""
+ return script_namespaces(**kwargs)["metal"]
+
+
+def __getattr__(name: str):
+ if name in _LAZY_SUBMODULES:
+ return import_module(f"{__name__}.{name}")
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+ "op",
+ "register_backend",
+ "script",
+ "script_namespace",
+ "script_namespaces",
+ "target_tags",
+]
diff --git a/python/tvm/backend/metal/op.py b/python/tvm/backend/metal/op.py
new file mode 100644
index 0000000000..f948576020
--- /dev/null
+++ b/python/tvm/backend/metal/op.py
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal-owned TIR intrinsic builders."""
+
+from __future__ import annotations
+
+from tvm.tirx.op import call_intrin
+
+
+def make_filled_simdgroup_matrix(d, index, value, col=8, row=8):
+ """Create a filled SIMDGroup matrix."""
+
+ return call_intrin("handle", "tirx.make_filled_simdgroup_matrix", d, index, value, col, row)
+
+
+def simdgroup_load(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+ """Load data from device or threadgroup memory to simdgroup."""
+
+ return call_intrin(
+ "handle",
+ "tirx.simdgroup_load",
+ d,
+ index,
+ ptr,
+ stride,
+ col,
+ row,
+ transpose_matrix,
+ )
+
+
+def simdgroup_store(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+ """Store data from simdgroup to device or threadgroup memory."""
+
+ return call_intrin(
+ "handle",
+ "tirx.simdgroup_store",
+ d,
+ index,
+ ptr,
+ stride,
+ col,
+ row,
+ transpose_matrix,
+ )
+
+
+def simdgroup_multiply_accumulate(d, index_d, a, index_a, b, index_b, c, index_c):
+ """Multiply and accumulate two matrices in simdgroup."""
+
+ return call_intrin(
+ "handle",
+ "tirx.simdgroup_multiply_accumulate",
+ d,
+ index_d,
+ a,
+ index_a,
+ b,
+ index_b,
+ c,
+ index_c,
+ )
+
+
+__all__ = [
+ "make_filled_simdgroup_matrix",
+ "simdgroup_load",
+ "simdgroup_multiply_accumulate",
+ "simdgroup_store",
+]
diff --git a/python/tvm/backend/metal/script.py b/python/tvm/backend/metal/script.py
new file mode 100644
index 0000000000..7c5d45564a
--- /dev/null
+++ b/python/tvm/backend/metal/script.py
@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal TVMScript namespace."""
+
+from __future__ import annotations
+
+from tvm.backend.metal import op as _metal_op
+from tvm.tirx import Buffer
+from tvm.tirx import op as _tir_op
+from tvm.tirx.script.builder.ir import _op_wrapper
+
+
+class MetalNamespace:
+ """The Metal intrinsics submodule."""
+
+ def __init__(self):
+ self.make_filled_simdgroup_matrix = _op_wrapper(_metal_op.make_filled_simdgroup_matrix)
+ self.simdgroup_load = _op_wrapper(_metal_op.simdgroup_load)
+ self.simdgroup_store = _op_wrapper(_metal_op.simdgroup_store)
+ self.simdgroup_multiply_accumulate = _op_wrapper(_metal_op.simdgroup_multiply_accumulate)
+
+ @staticmethod
+ def simd_shuffle(var, lane):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(var.dtype, "tirx.metal.simd_shuffle", var, lane)
+
+ @staticmethod
+ def simd_shuffle_up(var, delta):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(var.dtype, "tirx.metal.simd_shuffle_up", var, delta)
+
+ @staticmethod
+ def simd_shuffle_down(var, delta):
+ if isinstance(var, Buffer):
+ var = var[0]
+ return _tir_op.call_intrin(var.dtype, "tirx.metal.simd_shuffle_down", var, delta)
+
+
+__all__ = ["MetalNamespace"]
diff --git a/python/tvm/target/tag_registry/metal.py b/python/tvm/backend/metal/target_tags.py
similarity index 94%
rename from python/tvm/target/tag_registry/metal.py
rename to python/tvm/backend/metal/target_tags.py
index 6727db7c30..6ee758216d 100644
--- a/python/tvm/target/tag_registry/metal.py
+++ b/python/tvm/backend/metal/target_tags.py
@@ -16,14 +16,14 @@
# under the License.
"""Apple Metal GPU target tags."""
-from .registry import register_tag
+from tvm.target import register_tag
_METAL_HOST_TRIPLE = "arm64-apple-macos"
def _register_metal_tag(name, max_threads, shared_mem, warp_size, mcpu):
try:
- from ..codegen import llvm_is_valid_cpu
+ from tvm.target.codegen import llvm_is_valid_cpu
if not llvm_is_valid_cpu(mcpu, _METAL_HOST_TRIPLE):
return
diff --git a/python/tvm/tirx/backend/__init__.py b/python/tvm/backend/opencl/__init__.py
similarity index 83%
copy from python/tvm/tirx/backend/__init__.py
copy to python/tvm/backend/opencl/__init__.py
index 862bed83d2..a80696e590 100644
--- a/python/tvm/tirx/backend/__init__.py
+++ b/python/tvm/backend/opencl/__init__.py
@@ -14,6 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The TIR Adreno backend passes"""
+"""OpenCL-owned backend hooks."""
-from . import adreno
+
+def register_backend():
+ """Register OpenCL-owned Python semantics."""
+ return None
+
+
+__all__ = ["register_backend"]
diff --git a/python/tvm/tirx/backend/__init__.py b/python/tvm/backend/rocm/__init__.py
similarity index 83%
copy from python/tvm/tirx/backend/__init__.py
copy to python/tvm/backend/rocm/__init__.py
index 862bed83d2..d7574e974a 100644
--- a/python/tvm/tirx/backend/__init__.py
+++ b/python/tvm/backend/rocm/__init__.py
@@ -14,6 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The TIR Adreno backend passes"""
+"""ROCm-owned TIRx modules."""
-from . import adreno
+
+def register_backend():
+ """Register ROCm-owned Python semantics."""
+ return None
+
+
+__all__ = ["register_backend"]
diff --git a/python/tvm/backend/trn/__init__.py b/python/tvm/backend/trn/__init__.py
new file mode 100644
index 0000000000..7650ea87d5
--- /dev/null
+++ b/python/tvm/backend/trn/__init__.py
@@ -0,0 +1,68 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trainium-owned TIRx modules."""
+
+from importlib import import_module
+
+_LAZY_SUBMODULES = {"layout", "op", "operator", "pipeline", "script", "target_tags", "transform"}
+
+
+def register_backend():
+ """Register Trainium-owned Python semantics."""
+ from tvm.tirx import compilation_pipeline # pylint: disable=import-outside-toplevel
+ from tvm.tirx.script.builder import ir as builder_ir # pylint: disable=import-outside-toplevel
+
+ for name, namespace in script_namespaces().items():
+ builder_ir.register_script_namespace(name, namespace)
+
+ import_module(f"{__name__}.operator.tile_primitive")
+ trn_pipeline = import_module(f"{__name__}.pipeline")
+ import_module(f"{__name__}.target_tags")
+ import_module(f"{__name__}.transform")
+ compilation_pipeline.register_tir_pipeline("trn", trn_pipeline.trn_pipeline)
+
+
+def script_namespace(op_wrapper=None):
+ """Return the Trainium TVMScript namespace object."""
+ from .script import NKINamespace # pylint: disable=import-outside-toplevel
+
+ return NKINamespace(op_wrapper)
+
+
+def script_namespaces(op_wrapper=None, **_):
+ """Return Trainium-owned TVMScript namespaces."""
+ return {"nki": script_namespace(op_wrapper)}
+
+
+def __getattr__(name: str):
+ if name in _LAZY_SUBMODULES:
+ return import_module(f"{__name__}.{name}")
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+ "layout",
+ "op",
+ "operator",
+ "pipeline",
+ "register_backend",
+ "script",
+ "script_namespace",
+ "script_namespaces",
+ "target_tags",
+ "transform",
+]
diff --git a/python/tvm/backend/trn/layout.py b/python/tvm/backend/trn/layout.py
new file mode 100644
index 0000000000..5d5f089591
--- /dev/null
+++ b/python/tvm/backend/trn/layout.py
@@ -0,0 +1,123 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trainium-specific TIRx layout helpers."""
+
+from __future__ import annotations
+
+import functools
+import operator
+import re
+
+import tvm
+from tvm.tirx.expr import PrimExpr
+from tvm.tirx.layout import Axis, Iter, Layout, S, TileLayout
+
+_TRN_MEMORY_AXES = {"F", "P", "Bank"}
+_PSUM_MAX_ELEM_PER_BANK = 512
+
+
+def is_trainium_layout(layout: Layout | None) -> bool:
+ """Return whether a layout uses only Trainium memory axes."""
+ if not isinstance(layout, TileLayout):
+ return False
+ return not any(
+ iter.axis.is_memory() and iter.axis.name not in _TRN_MEMORY_AXES for iter in layout.shard
+ )
+
+
+def trainium_layout(annotation: str, shape: tuple[PrimExpr], is_psum: bool = False) -> TileLayout:
+ """Create a Trainium tile layout from a PF annotation string and logical shape."""
+ analyzer = tvm.arith.Analyzer()
+ assert re.fullmatch(r"[PF]*", annotation), (
+ f"annotation {annotation} must be a string of 'P' and 'F'"
+ )
+ assert len(annotation) == len(shape), (
+ f"annotation {annotation} and shape {shape} must have the same length"
+ )
+ num_p_dim = annotation.count("P")
+ if num_p_dim == 1:
+ p_idx = annotation.index("P")
+ p_dim = shape[p_idx]
+ assert analyzer.can_prove(p_dim <= 128 or p_dim % 128 == 0), (
+ f"There is only 1 P in the annotation. Partition size {p_dim} must be less than "
+ "or equal to 128 or a multiple of 128"
+ )
+ if analyzer.can_prove(p_dim > 128):
+ annotation = "F" + annotation
+ shape = (p_dim // 128, *shape[:p_idx], 128, *shape[p_idx + 1 :])
+ elif num_p_dim > 1:
+ p_dim_prod = functools.reduce(
+ operator.mul, [s for s, c in zip(shape, annotation) if c == "P"]
+ )
+ assert analyzer.can_prove(p_dim_prod <= 128), (
+ f"There are {num_p_dim} Ps in the annotation. Partition size {p_dim_prod} must be "
+ "less than or equal to 128"
+ )
+
+ f_shape = [s for s, c in zip(shape, annotation) if c == "F"]
+ p_shape = [s for s, c in zip(shape, annotation) if c == "P"]
+ f_strides = Layout._get_default_strides(f_shape, 1) # pylint: disable=protected-access
+ p_strides = Layout._get_default_strides(p_shape, 1) # pylint: disable=protected-access
+ f_tile_layout = TileLayout(S[tuple(f_shape) : tuple(s @ Axis.F for s in f_strides)])
+ p_tile_layout = TileLayout(S[tuple(p_shape) : tuple(s @ Axis.P for s in p_strides)])
+ result = []
+ f_index = p_index = 0
+
+ for char in annotation:
+ if char == "F":
+ result.append(f_tile_layout.shard[f_index])
+ f_index += 1
+ else:
+ result.append(p_tile_layout.shard[p_index])
+ p_index += 1
+ if num_p_dim == 1 and analyzer.can_prove(p_dim > 128):
+ higher_p = result[0]
+ result = result[1:]
+ result = [*result[:p_idx], higher_p, *result[p_idx:]]
+
+ res = TileLayout.from_iters(result, [], {})
+ if is_psum:
+ res = to_psum_layout(res)
+ return res
+
+
+def to_psum_layout(layout: TileLayout) -> TileLayout:
+ """Convert a Trainium sbuf layout to its psum physical-bank layout."""
+ analyzer = tvm.arith.Analyzer()
+ shard = []
+ for iter in layout.shard:
+ if iter.axis.name == "F":
+ if analyzer.can_prove(iter.stride % _PSUM_MAX_ELEM_PER_BANK == 0):
+ stride = analyzer.simplify(iter.stride // _PSUM_MAX_ELEM_PER_BANK)
+ shard.append(Iter(iter.extent, stride, Axis.get("Bank")))
+ elif analyzer.can_prove(_PSUM_MAX_ELEM_PER_BANK % iter.stride == 0):
+ c = analyzer.simplify(_PSUM_MAX_ELEM_PER_BANK // iter.stride)
+ if analyzer.can_prove(iter.extent < c):
+ shard.append(iter)
+ elif analyzer.can_prove(iter.extent % c == 0):
+ shard.append(Iter(analyzer.simplify(iter.extent // c), 1, Axis.get("Bank")))
+ shard.append(Iter(c, iter.stride, Axis.get("F")))
+ else:
+ raise ValueError(f"layout {layout} can not be converted to psum layout")
+ else:
+ raise ValueError(f"layout {layout} can not be converted to psum layout")
+ else:
+ shard.append(iter)
+ return TileLayout.from_iters(shard, [], {})
+
+
+__all__ = ["is_trainium_layout", "to_psum_layout", "trainium_layout"]
diff --git a/python/tvm/backend/trn/op.py b/python/tvm/backend/trn/op.py
new file mode 100644
index 0000000000..d919e4fb85
--- /dev/null
+++ b/python/tvm/backend/trn/op.py
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trainium-owned NKI intrinsic Python wrappers."""
+
+from __future__ import annotations
+
+from tvm.tirx.op import call_intrin
+
+
+def nki_load(res, data):
+ return call_intrin("", "tirx.nki_load", res, data)
+
+
+def nki_store(res, data):
+ return call_intrin("", "tirx.nki_store", res, data)
+
+
+def nki_tensor_copy(res, data):
+ return call_intrin("", "tirx.nki_tensor_copy", res, data)
+
+
+def nki_matmul(res, lhs, rhs, accum=True):
+ return call_intrin("", "tirx.nki_matmul", res, lhs, rhs, accum)
+
+
+def nki_activation(result, data, opcode, bias=0.0, scale=1.0):
+ return call_intrin("", "tirx.nki_activation", result, data, opcode, bias, scale)
+
+
+def nki_reciprocal(result, data):
+ return call_intrin("", "tirx.nki_reciprocal", result, data)
+
+
+def nki_tensorreduce(result, data, opcode, negate, *axes):
+ return call_intrin("", "tirx.nki_tensorreduce", result, data, opcode, negate, *axes)
+
+
+def nki_tensortensor(result, operand0, operand1, opcode):
+ return call_intrin("", "tirx.nki_tensortensor", result, operand0, operand1, opcode)
+
+
+def nki_tensorscalar(result, operand0, operand1, opcode, reverse=False):
+ return call_intrin("", "tirx.nki_tensorscalar", result, operand0, operand1, opcode, reverse)
+
+
+def nki_memset(result, value):
+ return call_intrin("", "tirx.nki_memset", result, value)
+
+
+def nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias=0.0, scale=1.0):
+ return call_intrin(
+ "",
+ "tirx.nki_activation_reduce",
+ reduce_res,
+ act_res,
+ data,
+ opcode,
+ reduce_opcode,
+ bias,
+ scale,
+ )
+
+
+def nki_tensorscalar_reduce(
+ reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse=False
+):
+ return call_intrin(
+ "",
+ "tirx.nki_tensorscalar_reduce",
+ reduce_res,
+ tensorscalar_res,
+ operand0,
+ operand1,
+ opcode,
+ reduce_opcode,
+ reverse,
+ )
+
+
+def nki_identity(result, size):
+ return call_intrin("", "tirx.nki_identity", result, size)
+
+
+def nki_scalar_tensor_tensor(
+ result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+):
+ return call_intrin(
+ "",
+ "tirx.nki_scalar_tensor_tensor",
+ result,
+ data,
+ operand0,
+ operand1,
+ opcode0,
+ opcode1,
+ reverse0,
+ reverse1,
+ )
+
+
+def nki_scalar_tensor_scalar(
+ result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+):
+ return call_intrin(
+ "",
+ "tirx.nki_scalar_tensor_scalar",
+ result,
+ data,
+ operand0,
+ operand1,
+ opcode0,
+ opcode1,
+ reverse0,
+ reverse1,
+ )
+
+
+def nki_affine_select(result, pred, true_value, false_value):
+ return call_intrin("", "tirx.nki_affine_select", result, pred, true_value, false_value)
+
+
+__all__ = [
+ "nki_activation",
+ "nki_activation_reduce",
+ "nki_affine_select",
+ "nki_identity",
+ "nki_load",
+ "nki_matmul",
+ "nki_memset",
+ "nki_reciprocal",
+ "nki_scalar_tensor_scalar",
+ "nki_scalar_tensor_tensor",
+ "nki_store",
+ "nki_tensor_copy",
+ "nki_tensorreduce",
+ "nki_tensorscalar",
+ "nki_tensorscalar_reduce",
+ "nki_tensortensor",
+]
diff --git a/python/tvm/tirx/backend/__init__.py b/python/tvm/backend/trn/operator/__init__.py
similarity index 86%
copy from python/tvm/tirx/backend/__init__.py
copy to python/tvm/backend/trn/operator/__init__.py
index 862bed83d2..2d9b93d998 100644
--- a/python/tvm/tirx/backend/__init__.py
+++ b/python/tvm/backend/trn/operator/__init__.py
@@ -14,6 +14,9 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The TIR Adreno backend passes"""
+"""Trainium backend operator package.
-from . import adreno
+Loaded by the Trainium backend registration hook.
+"""
+
+__all__ = ["tile_primitive"]
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/binary/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/binary/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/binary/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/binary/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/binary/default.py b/python/tvm/backend/trn/operator/tile_primitive/binary/default.py
similarity index 97%
rename from python/tvm/tirx/operator/tile_primitive/trn/binary/default.py
rename to python/tvm/backend/trn/operator/tile_primitive/binary/default.py
index 3fa565b1f4..9546c2f6e4 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/binary/default.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/binary/default.py
@@ -20,9 +20,9 @@
from tvm.script import tirx as T
from tvm.tirx import FloatImm, PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import MapOpType
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import MapOpType
from ..common import init_analyzer, nki_dim
from ..instruction_generator import InstructionGenerator
from .utils import InstType, binary_map_ops, try_find_inst_nary
@@ -32,7 +32,7 @@ def binary_trn(
op: TilePrimitiveCall, binary_op: MapOpType, sctx: DispatchContext
) -> PrimFunc | None:
"""Generate a binary operation schedule for Trainium."""
- if not (sctx.is_trn() and sctx.scope_kind == "thread"):
+ if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
assert binary_op in binary_map_ops, f"Unsupported binary operation {binary_op}"
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/binary/utils.py b/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py
similarity index 97%
rename from python/tvm/tirx/operator/tile_primitive/trn/binary/utils.py
rename to python/tvm/backend/trn/operator/tile_primitive/binary/utils.py
index 0f0c0e053f..891c5e9c48 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/binary/utils.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py
@@ -20,9 +20,10 @@
from enum import Enum
from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.tirx import BufferRegion, FloatImm
+from tvm.tirx.operator.tile_primitive.common import MapOpType
-from ...common import MapOpType
from ..dim_utils import get_ewise_dim_map
from ..instruction_generator import InstructionGenerator
@@ -72,8 +73,8 @@ def try_find_inst_nary(
valid_buffers = all(
[
dst.layout and all(src.layout for src in srcs if src is not None),
- dst.layout.is_trainium(),
- all(src.layout.is_trainium() for src in srcs if src is not None),
+ is_trainium_layout(dst.layout),
+ all(is_trainium_layout(src.layout) for src in srcs if src is not None),
dst.scope() == "trn.sbuf",
all(src.scope() in ["trn.sbuf", "trn.psum"] for src in srcs if src is not None),
]
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/common.py b/python/tvm/backend/trn/operator/tile_primitive/common.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/common.py
rename to python/tvm/backend/trn/operator/tile_primitive/common.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/binary_chain.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/binary_chain.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/binary_reduce.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/binary_reduce.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/compose_op.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/compose_op.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/compose_op.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/compose_op.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/reduce_negate.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/reduce_negate.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/reduce_negate.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/reduce_negate.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/unary_reduce.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/unary_reduce.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/utils.py b/python/tvm/backend/trn/operator/tile_primitive/compose_op/utils.py
similarity index 95%
rename from python/tvm/tirx/operator/tile_primitive/trn/compose_op/utils.py
rename to python/tvm/backend/trn/operator/tile_primitive/compose_op/utils.py
index 9fbaa524fd..1509482be8 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/compose_op/utils.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/compose_op/utils.py
@@ -18,8 +18,7 @@
"""Shared helpers for compose operator dispatches."""
from tvm.ir import Op
-
-from ...common import ReduceOpType
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
# Operation code mappings
opcode_table = {
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/copy/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/copy/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/copy/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/copy/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/copy/default.py b/python/tvm/backend/trn/operator/tile_primitive/copy/default.py
similarity index 96%
rename from python/tvm/tirx/operator/tile_primitive/trn/copy/default.py
rename to python/tvm/backend/trn/operator/tile_primitive/copy/default.py
index b1a0b20786..123698b621 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/copy/default.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/copy/default.py
@@ -17,6 +17,7 @@
"""Implementation of copy operator dispatchs."""
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.script import tirx as T
from tvm.tirx import PrimFunc
from tvm.tirx.operator.tile_primitive import (
@@ -86,7 +87,7 @@ def transpose_schedule(
if "identity" not in op.workspace:
assert sctx.alloc_only, (
- "Identity tensor must be specified in workspace. Run tvm.tirx.transform.trn.TrnPrivateBufferAlloc first." # noqa: E501
+ "Identity tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
)
identity_tensor = T.buffer(
(p_size, rhs_f_size), src_region.buffer.dtype, scope="trn.sbuf", buffer_name="identity"
@@ -143,7 +144,7 @@ def transpose_schedule(
if "acc_psum" not in op.workspace:
assert sctx.alloc_only, (
- "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.transform.trn.TrnPrivateBufferAlloc first." # noqa: E501
+ "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
)
acc_psum = T.buffer(
(max_psum_banks, p_size, largest_psum_per_bank),
@@ -202,9 +203,9 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
dst.scope() in ["global", "trn.sbuf", "trn.psum"],
src.scope() != "global" or dst.scope() != "global",
(src.scope() == "global" and isinstance(src.layout, T.TileLayout))
- or (src.scope() in ["trn.sbuf", "trn.psum"] and src.layout.is_trainium()),
+ or (src.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(src.layout)),
(dst.scope() == "global" and isinstance(dst.layout, T.TileLayout))
- or (dst.scope() in ["trn.sbuf", "trn.psum"] and dst.layout.is_trainium()),
+ or (dst.scope() in ["trn.sbuf", "trn.psum"] and is_trainium_layout(dst.layout)),
]
)
@@ -232,7 +233,7 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
if not inst_gen.check_partition_dim_match(src_region, dst_region):
return transpose_schedule(op, inst_gen, sctx)
- if src.layout.is_trainium():
+ if is_trainium_layout(src.layout):
inst = inst_gen.find_max_inst_size_from_one_region(src_region)
inst = inst_gen.fit_inst_tile_to_region(inst, dst_region)
src_to_dst = True
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/dim_utils.py b/python/tvm/backend/trn/operator/tile_primitive/dim_utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/dim_utils.py
rename to python/tvm/backend/trn/operator/tile_primitive/dim_utils.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/gemm/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/gemm/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/gemm/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/gemm/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/gemm/default.py b/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py
similarity index 97%
rename from python/tvm/tirx/operator/tile_primitive/trn/gemm/default.py
rename to python/tvm/backend/trn/operator/tile_primitive/gemm/default.py
index ca572ba781..eba660ab6d 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/gemm/default.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py
@@ -21,6 +21,7 @@ import functools
import operator
from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.ir import assert_structural_equal
from tvm.script import tirx as T
from tvm.tirx import BufferRegion, PrimFunc
@@ -110,7 +111,7 @@ def get_pf_dim_from_buffer_region(
def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
"""Schedule GEMM operation on Trainium."""
# Basic validation checks
- if not (sctx.is_trn() and sctx.scope_kind == "thread"):
+ if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
# Extract arguments
@@ -147,9 +148,9 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
A.dtype == B.dtype,
A.scope() == "trn.sbuf" and B.scope() == "trn.sbuf",
C.scope() == "trn.psum" or C.scope() == "trn.sbuf",
- A.layout.is_trainium(),
- B.layout.is_trainium(),
- C.layout.is_trainium(),
+ is_trainium_layout(A.layout),
+ is_trainium_layout(B.layout),
+ is_trainium_layout(C.layout),
A.layout.size("P") == B.layout.size("P"),
]
), "Invalid buffer layout and scope"
@@ -252,7 +253,7 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
acc_psum_shape = (max_psum_banks, p_size, largest_psum_per_bank)
if "acc_psum" not in op.workspace:
- assert sctx.alloc_only, "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.transform.trn.TrnPrivateBufferAlloc first." # noqa: E501
+ assert sctx.alloc_only, "Accumulation psum buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
acc_psum = T.buffer(
acc_psum_shape,
"float32",
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/instruction_generator.py b/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py
similarity index 98%
rename from python/tvm/tirx/operator/tile_primitive/trn/instruction_generator.py
rename to python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py
index 58163d4b14..d501aab826 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/instruction_generator.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py
@@ -25,6 +25,7 @@ from operator import mul
import tvm
from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.ir import Range
from tvm.script import tirx as T
from tvm.tirx import BufferRegion, PrimExpr, Var
@@ -379,7 +380,7 @@ class InstructionGenerator:
for i in reversed(dims):
for j in reversed(range(seps[i], seps[i + 1])):
it = shards[j]
- is_partition = it.axis.name == "P" if layout.is_trainium() else False
+ is_partition = it.axis.name == "P" if is_trainium_layout(layout) else False
logical_iter_dims = bind_iters[i][j - seps[i]]
for d in range(-1, len(logical_iter_dims)):
next_logical_stride = (
@@ -479,7 +480,9 @@ class InstructionGenerator:
is_free_dim: bool = True,
) -> LogicalIterList:
layout = self.split_layout_views[buffer_region]
- assert layout.is_trainium(), " Cannot propagate instruction information from HBM tensor"
+ assert is_trainium_layout(layout), (
+ " Cannot propagate instruction information from HBM tensor"
+ )
iters = layout.shard
seps = self.seps[buffer_region]
ret = [[[] for _ in range(seps[i], seps[i + 1])] for i in range(len(buffer_region.region))]
@@ -633,7 +636,7 @@ class InstructionGenerator:
mapped_data_iter = to_layout.shard[data_iter_map[i]]
if inst_stride_from is None:
inst_stride_from = data_iter.stride
- if not to_layout.is_trainium() and mapped_data_iter.stride != 1:
+ if not is_trainium_layout(to_layout) and mapped_data_iter.stride != 1:
# dma copy must be contiguous on hbm
break
inst_stride_to = mapped_data_iter.stride
@@ -649,7 +652,7 @@ class InstructionGenerator:
dim_map = self.dim_mapper.get_dim_map(buffer_region_1, buffer_region_2)
layout_1 = self.split_layout_views[buffer_region_1]
layout_2 = self.split_layout_views[buffer_region_2]
- if not layout_1.is_trainium() or not layout_2.is_trainium():
+ if not is_trainium_layout(layout_1) or not is_trainium_layout(layout_2):
return True
seps_1 = self.seps[buffer_region_1]
seps_2 = self.seps[buffer_region_2]
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/private_alloc.py b/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py
similarity index 96%
rename from python/tvm/tirx/operator/tile_primitive/trn/private_alloc.py
rename to python/tvm/backend/trn/operator/tile_primitive/private_alloc.py
index fe3f0a54bb..830ab45f2d 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/private_alloc.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py
@@ -17,6 +17,9 @@
from typing import Any
+from tvm.backend.trn.operator.tile_primitive.common import init_analyzer, nki_dim
+from tvm.backend.trn.operator.tile_primitive.dim_utils import get_ewise_dim_map
+from tvm.backend.trn.operator.tile_primitive.instruction_generator import InstructionGenerator
from tvm.script import tirx as T
from tvm.tirx import Buffer, FloatImm, Stmt
from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
@@ -29,9 +32,6 @@ from tvm.tirx.operator.tile_primitive.ops import (
UnaryReduce,
)
from tvm.tirx.operator.tile_primitive.registry import f_op_dispatcher
-from tvm.tirx.operator.tile_primitive.trn.common import init_analyzer, nki_dim
-from tvm.tirx.operator.tile_primitive.trn.dim_utils import get_ewise_dim_map
-from tvm.tirx.operator.tile_primitive.trn.instruction_generator import InstructionGenerator
from tvm.tirx.stmt import TilePrimitiveCall
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/reduction/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/reduction/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/reduction/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/reduction/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py b/python/tvm/backend/trn/operator/tile_primitive/reduction/default.py
similarity index 94%
rename from python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py
rename to python/tvm/backend/trn/operator/tile_primitive/reduction/default.py
index f7a7b886d0..0a1d8b6cbb 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/reduction/default.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/reduction/default.py
@@ -18,8 +18,8 @@
"""Reduction dispatch variant registrations."""
from tvm.tirx.operator.tile_primitive import register_dispatch
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
-from ...common import ReduceOpType
from .utils import reduction_trn
for _op_name, _op_type in {
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/reduction/utils.py b/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py
similarity index 95%
rename from python/tvm/tirx/operator/tile_primitive/trn/reduction/utils.py
rename to python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py
index 1d9840e5d6..43bd92d56d 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/reduction/utils.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py
@@ -17,12 +17,13 @@
"""Shared helpers for reduction schedules."""
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.script import tirx as T
from tvm.tirx import PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import ReduceOpType
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_reduction_dim_map
from ..instruction_generator import InstructionGenerator
@@ -46,7 +47,7 @@ def generate_intermediate_buffer(
check_workspace_buffer(intermediate_buffer, intermediate_shape, "trn.sbuf")
else:
assert sctx.alloc_only, (
- "Partial reduce buffer must be specified in workspace. Run tvm.tirx.transform.trn.TrnPrivateBufferAlloc first." # noqa: E501
+ "Partial reduce buffer must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
)
intermediate_buffer = T.buffer(
intermediate_shape,
@@ -73,7 +74,7 @@ def reduction_trn(
Returns:
Optional[PrimFunc]: The scheduled function, or None if not applicable.
"""
- if not (sctx.is_trn() and sctx.scope_kind == "thread"):
+ if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
dst_buffer_region, src_buffer_region, axes, accum = op.args[:4]
@@ -93,8 +94,8 @@ def reduction_trn(
src.layout and dst.layout,
src.scope() == "trn.sbuf" or src.scope() == "trn.psum",
dst.scope() == "trn.sbuf",
- src.layout.is_trainium(),
- dst.layout.is_trainium(),
+ is_trainium_layout(src.layout),
+ is_trainium_layout(dst.layout),
src.layout.size("P") == dst.layout.size("P"),
]
), "Invalid layout"
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/select/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/select/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/select/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/select/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/select/default.py b/python/tvm/backend/trn/operator/tile_primitive/select/default.py
similarity index 97%
rename from python/tvm/tirx/operator/tile_primitive/trn/select/default.py
rename to python/tvm/backend/trn/operator/tile_primitive/select/default.py
index 27136a3ac3..07486e6ae7 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/select/default.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/select/default.py
@@ -17,6 +17,7 @@
"""Implementation of select schedules."""
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.script import tirx as T
from tvm.tirx import BufferRegion, FloatImm, PrimFunc, TilePrimitiveCall
from tvm.tirx.operator.tile_primitive import (
@@ -63,8 +64,8 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
buffer_conditions = [
dst.buffer.layout and true_value.buffer.layout,
dst.buffer.scope() == "trn.sbuf" and true_value.buffer.scope() == "trn.sbuf",
- true_value.buffer.layout.is_trainium(),
- dst.buffer.layout.is_trainium(),
+ is_trainium_layout(true_value.buffer.layout),
+ is_trainium_layout(dst.buffer.layout),
]
if not all(buffer_conditions):
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/unary/__init__.py b/python/tvm/backend/trn/operator/tile_primitive/unary/__init__.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/unary/__init__.py
rename to python/tvm/backend/trn/operator/tile_primitive/unary/__init__.py
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/unary/default.py b/python/tvm/backend/trn/operator/tile_primitive/unary/default.py
similarity index 96%
rename from python/tvm/tirx/operator/tile_primitive/trn/unary/default.py
rename to python/tvm/backend/trn/operator/tile_primitive/unary/default.py
index e336daa717..69facea7f9 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/unary/default.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/unary/default.py
@@ -19,9 +19,9 @@
from tvm.tirx import FloatImm, PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import MapOpType
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import MapOpType
from ..common import init_analyzer
from ..instruction_generator import InstructionGenerator
from .utils import (
@@ -35,7 +35,7 @@ from .utils import (
def unary_trn(op: TilePrimitiveCall, unary_op: MapOpType, sctx: DispatchContext) -> PrimFunc | None:
"""Schedule unary operation on Trainium."""
# Check execution environment
- if not (sctx.is_trn() and sctx.scope_kind == "thread"):
+ if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
# Extract operation arguments
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/unary/utils.py b/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py
similarity index 96%
rename from python/tvm/tirx/operator/tile_primitive/trn/unary/utils.py
rename to python/tvm/backend/trn/operator/tile_primitive/unary/utils.py
index 7a757609b0..dc58862bb4 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/unary/utils.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py
@@ -18,10 +18,11 @@
"""Shared helpers, op tables, and validation functions for unary operator dispatches."""
from tvm.arith.analyzer import Analyzer
+from tvm.backend.trn.layout import is_trainium_layout
from tvm.script import tirx as T
from tvm.tirx import BufferRegion, FloatImm
+from tvm.tirx.operator.tile_primitive.common import MapOpType
-from ...common import MapOpType
from ..common import nki_dim
from ..dim_utils import get_ewise_dim_map
from ..instruction_generator import InstructionGenerator
@@ -56,8 +57,8 @@ def try_find_inst_unary(
src.layout and dst.layout,
src.scope() in ("trn.sbuf", "trn.psum"),
dst.scope() == "trn.sbuf",
- src.layout.is_trainium(),
- dst.layout.is_trainium(),
+ is_trainium_layout(src.layout),
+ is_trainium_layout(dst.layout),
]
)
@@ -97,7 +98,7 @@ def get_const_bias_tensor(bias, shape, dtype, workspace, sctx):
"""Create or retrieve a constant bias tensor."""
if "const_bias" not in workspace:
assert sctx.alloc_only, (
- "Constant bias tensor must be specified in workspace. Run tvm.tirx.transform.trn.TrnPrivateBufferAlloc first." # noqa: E501
+ "Constant bias tensor must be specified in workspace. Run tvm.tirx.trn.transform.TrnPrivateBufferAlloc first." # noqa: E501
)
# Create new bias buffer
bias_buffer = T.buffer(shape, dtype, scope="trn.sbuf", buffer_name="const_bias")
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/unary/with_bias_scale.py b/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py
similarity index 96%
rename from python/tvm/tirx/operator/tile_primitive/trn/unary/with_bias_scale.py
rename to python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py
index 399d8cfa6d..26fb567027 100644
--- a/python/tvm/tirx/operator/tile_primitive/trn/unary/with_bias_scale.py
+++ b/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py
@@ -19,9 +19,9 @@
from tvm.tirx import BufferRegion, PrimFunc
from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import MapOpType
from tvm.tirx.stmt import TilePrimitiveCall
-from ...common import MapOpType
from ..binary import try_find_inst_nary
from ..common import init_analyzer
from ..instruction_generator import InstructionGenerator
@@ -33,7 +33,7 @@ def unary_with_bias_scale_trn(
) -> PrimFunc | None:
"""Schedule unary operation with bias and scale on Trainium."""
# Check execution environment
- if not (sctx.is_trn() and sctx.scope_kind == "thread"):
+ if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
# Extract operation arguments with defaults
diff --git a/python/tvm/tirx/operator/tile_primitive/trn/workspace_utils.py b/python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py
similarity index 100%
rename from python/tvm/tirx/operator/tile_primitive/trn/workspace_utils.py
rename to python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py
diff --git a/python/tvm/backend/trn/pipeline.py b/python/tvm/backend/trn/pipeline.py
new file mode 100644
index 0000000000..d0663218e6
--- /dev/null
+++ b/python/tvm/backend/trn/pipeline.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trainium TIRX pipeline entrypoints."""
+
+import tvm
+from tvm import tirx
+from tvm.tirx.compilation_pipeline import finalize_host_passes
+
+from . import transform as trn_transform
+
+
+def trn_pipeline():
+ """The Trainium pipeline used in tvm.tirx.build."""
+
+ @tvm.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+ """Lower TIRx for the Trainium backend."""
+ passes = [
+ trn_transform.TrnPrivateBufferAlloc(),
+ trn_transform.TrnNaiveAllocator(),
+ tirx.transform.TilePrimitiveDispatch(),
+ trn_transform.LowerTrainiumLayout(),
+ tvm.s_tir.transform.DecorateDeviceScope(),
+ tirx.transform.StmtSimplify(),
+ tirx.transform.LowerTIRxOpaque(),
+ tvm.s_tir.transform.LoopPartition(),
+ tvm.s_tir.transform.HoistIfThenElse(),
+ tirx.transform.StmtSimplify(),
+ tirx.transform.RemoveNoOp(),
+ tirx.transform.AnnotateEntryFunc(),
+ tirx.transform.SplitHostDevice(),
+ tirx.transform.MakePackedAPI(),
+ ]
+ return tvm.ir.transform.Sequential(passes)(mod)
+
+ return _pipeline, finalize_host_passes, finalize_device_passes_trn
+
+
+def finalize_device_passes_trn(): # pylint: disable=unused-argument
+ """The finalization passes for the Trainium backend."""
+ return tvm.ir.transform.Sequential([tirx.transform.StmtSimplify()])
+
+
+__all__ = ["finalize_device_passes_trn", "trn_pipeline"]
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
new file mode 100644
index 0000000000..5d507f462b
--- /dev/null
+++ b/python/tvm/backend/trn/script.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trainium TVMScript namespaces."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from . import op as _trn_op
+
+OpWrapper = Callable[[Callable[..., Any]], Callable[..., Any]]
+
+
+def _default_op_wrapper() -> OpWrapper:
+ from tvm.tirx.script.builder.ir import _op_wrapper # pylint: disable=import-outside-toplevel
+
+ return _op_wrapper
+
+
+class NKINamespace:
+ """The NKI instructions submodule."""
+
+ def __init__(self, op_wrapper: OpWrapper | None = None):
+ wrap = op_wrapper or _default_op_wrapper()
+ self.load = wrap(_trn_op.nki_load)
+ self.store = wrap(_trn_op.nki_store)
+ self.tensor_copy = wrap(_trn_op.nki_tensor_copy)
+ self.matmul = wrap(_trn_op.nki_matmul)
+ self.activation = wrap(_trn_op.nki_activation)
+ self.activation_reduce = wrap(_trn_op.nki_activation_reduce)
+ self.reciprocal = wrap(_trn_op.nki_reciprocal)
+ self.tensorreduce = wrap(_trn_op.nki_tensorreduce)
+ self.tensortensor = wrap(_trn_op.nki_tensortensor)
+ self.tensorscalar = wrap(_trn_op.nki_tensorscalar)
+ self.tensorscalar_reduce = wrap(_trn_op.nki_tensorscalar_reduce)
+ self.scalar_tensor_tensor = wrap(_trn_op.nki_scalar_tensor_tensor)
+ self.scalar_tensor_scalar = wrap(_trn_op.nki_scalar_tensor_scalar)
+ self.memset = wrap(_trn_op.nki_memset)
+ self.identity = wrap(_trn_op.nki_identity)
+ self.affine_select = wrap(_trn_op.nki_affine_select)
+
+
+__all__ = ["NKINamespace", "OpWrapper"]
diff --git a/src/runtime/hexagon/rpc/android_bash.sh.template b/python/tvm/backend/trn/target_tags.py
similarity index 62%
copy from src/runtime/hexagon/rpc/android_bash.sh.template
copy to python/tvm/backend/trn/target_tags.py
index c45b03818f..54c63d7226 100644
--- a/src/runtime/hexagon/rpc/android_bash.sh.template
+++ b/python/tvm/backend/trn/target_tags.py
@@ -1,4 +1,3 @@
-#!/bin/sh
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
@@ -15,17 +14,23 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+"""AWS Trainium target tags."""
-export LD_LIBRARY_PATH=.
+from tvm.target import register_tag
-# Enable FARF-based logging for Hexagon code invoked by 'tvm_rpc_android_server'.
-export ADSP_LIBRARY_PATH=`pwd`
-echo <FARF_CONFIG> > tvm_rpc_android.farf
+def _register_aws_trn1_tag(name, cores):
+ register_tag(
+ name,
+ {
+ "kind": "trn",
+ "num-cores": cores,
+ "partition_size": 128,
+ "max_sbuf_size_per_partition": 196608,
+ "max_psum_size_per_partition": 16384,
+ },
+ )
-./tvm_rpc_android server --port=<RPC_SERVER_PORT> --tracker=<RPC_TRACKER_HOST>:<RPC_TRACKER_PORT> --key=<HEXAGON_REMOTE_DEVICE_KEY> >${PWD}/tvm_rpc_android.log 2>&1 &
-rpc_pid=$!
-
-rm -f rpc_pid.txt
-echo $rpc_pid >> rpc_pid.txt
+_register_aws_trn1_tag("aws/trn1/trn1.2xlarge", 2)
+_register_aws_trn1_tag("aws/trn1/trn1.32xlarge", 32)
diff --git a/python/tvm/tirx/transform/trn/__init__.py b/python/tvm/backend/trn/transform/__init__.py
similarity index 68%
rename from python/tvm/tirx/transform/trn/__init__.py
rename to python/tvm/backend/trn/transform/__init__.py
index 0aaf3062c8..84dac89cec 100644
--- a/python/tvm/tirx/transform/trn/__init__.py
+++ b/python/tvm/backend/trn/transform/__init__.py
@@ -1,4 +1,3 @@
-# isort: skip_file
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
@@ -18,17 +17,30 @@
"""Trainium-specific TIRX transformations."""
# pylint: disable=invalid-name
-# Fork-only TIRX-specific passes. They decorate their pass body with
-# `@prim_func_pass(...)` at module-load time, which triggers an FFI call to
-# construct PassInfo -- not runtime-safe. Loading them lazily preserves
-# apache's discipline that `import tvm.tirx.transform.trn` performs no
-# compiler-side FFI calls (required for `TVM_USE_RUNTIME_LIB=1`).
+from tvm_ffi import get_global_func
+
+import tvm
+from tvm import tirx
+
_LAZY_TRANSFORMS = {
"TrnNaiveAllocator": ".naive_allocator",
"TrnPrivateBufferAlloc": ".private_buffer_alloc",
}
+def LowerTrainiumLayout():
+ """Lower Trainium layouts to backend physical buffer shapes and indices."""
+ return get_global_func("tirx.backend.trn.transform.LowerTrainiumLayout")()
+
+
+def LowerTIRx():
+ """Lower TIRx tile primitive calls for the Trainium backend."""
+ return tvm.ir.transform.Sequential(
+ [tirx.transform.TilePrimitiveDispatch(), LowerTrainiumLayout()],
+ name="tirx.backend.trn.LowerTIRx",
+ )
+
+
def __getattr__(name):
target = _LAZY_TRANSFORMS.get(name)
if target is None:
@@ -36,3 +48,6 @@ def __getattr__(name):
from importlib import import_module # pylint: disable=import-outside-toplevel
return getattr(import_module(target, __name__), name)
+
+
+__all__ = ["LowerTIRx", "LowerTrainiumLayout", "TrnNaiveAllocator", "TrnPrivateBufferAlloc"]
diff --git a/python/tvm/tirx/transform/trn/naive_allocator.py b/python/tvm/backend/trn/transform/naive_allocator.py
similarity index 98%
rename from python/tvm/tirx/transform/trn/naive_allocator.py
rename to python/tvm/backend/trn/transform/naive_allocator.py
index 1720a32d69..ccab7add39 100644
--- a/python/tvm/tirx/transform/trn/naive_allocator.py
+++ b/python/tvm/backend/trn/transform/naive_allocator.py
@@ -21,10 +21,9 @@ from tvm import DataType
from tvm.tirx import AllocBuffer, IntImm
from tvm.tirx.buffer import Buffer
from tvm.tirx.stmt_functor import StmtVisitor
+from tvm.tirx.transform.common import BufferReplacer
from tvm.tirx.transform.function_pass import prim_func_pass
-from ..common import BufferReplacer
-
def is_const_shape(buffer: Buffer) -> bool:
for i in buffer.shape:
diff --git a/python/tvm/tirx/transform/trn/private_buffer_alloc.py b/python/tvm/backend/trn/transform/private_buffer_alloc.py
similarity index 100%
rename from python/tvm/tirx/transform/trn/private_buffer_alloc.py
rename to python/tvm/backend/trn/transform/private_buffer_alloc.py
diff --git a/python/tvm/tirx/backend/__init__.py b/python/tvm/backend/vulkan/__init__.py
similarity index 83%
copy from python/tvm/tirx/backend/__init__.py
copy to python/tvm/backend/vulkan/__init__.py
index 862bed83d2..343875aa8a 100644
--- a/python/tvm/tirx/backend/__init__.py
+++ b/python/tvm/backend/vulkan/__init__.py
@@ -14,6 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The TIR Adreno backend passes"""
+"""Vulkan-owned backend hooks."""
-from . import adreno
+
+def register_backend():
+ """Register Vulkan-owned Python semantics."""
+ return None
+
+
+__all__ = ["register_backend"]
diff --git a/python/tvm/tirx/backend/__init__.py b/python/tvm/backend/webgpu/__init__.py
similarity index 83%
copy from python/tvm/tirx/backend/__init__.py
copy to python/tvm/backend/webgpu/__init__.py
index 862bed83d2..cb682a56d2 100644
--- a/python/tvm/tirx/backend/__init__.py
+++ b/python/tvm/backend/webgpu/__init__.py
@@ -14,6 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The TIR Adreno backend passes"""
+"""WebGPU-owned backend hooks."""
-from . import adreno
+
+def register_backend():
+ """Register WebGPU-owned Python semantics."""
+ return None
+
+
+__all__ = ["register_backend"]
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
index 0e2047af32..fcd88fdae3 100644
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ b/python/tvm/s_tir/tensor_intrin/cuda.py
@@ -580,7 +580,7 @@ def get_mma_fill_intrin(dtype, local_size):
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
- T.mma_fill_legacy(local_size, C_warp.data, C_warp.elem_offset, dtype=dtype)
+ T.cuda.mma_fill_legacy(local_size, C_warp.data, C_warp.elem_offset, dtype=dtype)
)
return mma_fill_desc, mma_fill_impl
@@ -637,7 +637,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
- T.mma_store_legacy(
+ T.cuda.mma_store_legacy(
M_DIM,
N_DIM,
C.access_ptr("w"),
diff --git a/python/tvm/s_tir/tensor_intrin/metal.py b/python/tvm/s_tir/tensor_intrin/metal.py
index a789581d4b..894aeea656 100644
--- a/python/tvm/s_tir/tensor_intrin/metal.py
+++ b/python/tvm/s_tir/tensor_intrin/metal.py
@@ -60,7 +60,7 @@ def get_make_filled_simdgroup_matrix_intrin(
with T.sblock("root"):
T.reads()
T.writes(A[0:col, 0:row])
- T.make_filled_simdgroup_matrix(
+ T.metal.make_filled_simdgroup_matrix(
A.data,
index=get_simdgroup_index(A, d1, col, row),
value=T.float32(0),
@@ -122,7 +122,7 @@ def get_simdgroup_load_intrin(
with T.sblock("root"):
T.reads(A[0:col, 0:row])
T.writes(C[0:col, 0:row])
- T.simdgroup_load(
+ T.metal.simdgroup_load(
C.data,
index=get_simdgroup_index(C, d1, col, row),
ptr=A.access_ptr("r"),
@@ -179,7 +179,7 @@ def get_simdgroup_store_intrin(
with T.sblock("root"):
T.reads(A[0:col, 0:row])
T.writes(C[0:col, 0:row])
- T.simdgroup_store(
+ T.metal.simdgroup_store(
A.data,
index=get_simdgroup_index(A, s1, col, row),
ptr=C.access_ptr("w"),
@@ -223,7 +223,7 @@ def get_simdgroup_multiply_accumulate_intrin(
with T.sblock("root"):
T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:k_dim, 0:n_dim])
T.writes(C[0:m_dim, 0:n_dim])
- T.simdgroup_multiply_accumulate(
+ T.metal.simdgroup_multiply_accumulate(
C.data,
get_simdgroup_index(C, c1, m_dim, n_dim),
A.data,
diff --git a/python/tvm/target/tag_registry/__init__.py b/python/tvm/target/tag_registry/__init__.py
index 5ace43847e..114edc448c 100644
--- a/python/tvm/target/tag_registry/__init__.py
+++ b/python/tvm/target/tag_registry/__init__.py
@@ -21,13 +21,9 @@ Importing this package registers all Python-defined target tags.
"""
from . import registry
-from . import cuda
from . import arm_cpu
from . import riscv_cpu
from . import aws_cpu
-from . import metal
-from . import hexagon
-from . import adreno
# Validate all tags at import time
registry.list_tags()
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 4378a9dfbe..5e8a2184bd 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -64,16 +64,6 @@ from .op import (
tvm_bmma_sync,
tvm_fill_fragment,
)
-from .op import ptx_mma, ptx_mma_sp, mma_store, mma_fill
-from .op import ptx_mma_legacy, ptx_mma_sp_legacy, mma_store_legacy, mma_fill_legacy
-from .op import ptx_ldmatrix, ptx_cp_async, ptx_cp_async_bulk, ptx_cp_async_bulk_shared_to_cluster
-from .op import ptx_ldmatrix_legacy, ptx_cp_async_legacy
-from .op import (
- make_filled_simdgroup_matrix,
- simdgroup_load,
- simdgroup_multiply_accumulate,
- simdgroup_store,
-)
from .op import vectorlow, vectorhigh, vectorcombine
from .op import infinity, reinterpret
from .op import exp, exp2, exp10, log, log2, log10, log1p, ldexp, clz
@@ -113,16 +103,11 @@ from .functor import PyStmtExprVisitor, PyStmtExprMutator
from tvm.base import _RUNTIME_ONLY as _RUNTIME_ONLY_TIRX # pylint: disable=wrong-import-position
if not _RUNTIME_ONLY_TIRX:
- # CUDA codegen registration. Each family module registers codegen via
- # @register_codegen (hand-written ops) and ptx_intrinsic /
- # cuda_helper_intrinsic (schema-declared ops); the schema declarations
- # also inject Python wrappers into `tvm.tirx.op`. Must come before
- # anything downstream that looks up wrappers or the codegen registry.
- from .operator.intrinsics import cuda as _intrinsics_cuda
from .build import build
from .compilation_pipeline import (
get_tir_pipeline,
get_default_tir_pipeline,
+ register_tir_pipeline,
)
import tvm.script
diff --git a/python/tvm/tirx/backend/__init__.py b/python/tvm/tirx/backend/__init__.py
index 862bed83d2..6f5333629f 100644
--- a/python/tvm/tirx/backend/__init__.py
+++ b/python/tvm/tirx/backend/__init__.py
@@ -14,6 +14,4 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The TIR Adreno backend passes"""
-
-from . import adreno
+"""TIRx backend compatibility package."""
diff --git a/python/tvm/tirx/backend/adreno/__init__.py b/python/tvm/tirx/backend/adreno/__init__.py
deleted file mode 100644
index 06f85d091d..0000000000
--- a/python/tvm/tirx/backend/adreno/__init__.py
+++ /dev/null
@@ -1,17 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""The TIR Adreno backend passes"""
diff --git a/python/tvm/tirx/bench.py b/python/tvm/tirx/bench.py
index d12ff2e3d0..aa06763759 100644
--- a/python/tvm/tirx/bench.py
+++ b/python/tvm/tirx/bench.py
@@ -613,7 +613,7 @@ class CudaProfiler:
@T.inline
def init(self, group_id: tvm.tirx.PrimExpr):
if self.profiler_enabled:
- T.timer_init_cuda(
+ T.cuda.timer_init(
self.buffer.data,
self.profiler_tag.data,
self.profiler_write_offset.data,
@@ -624,7 +624,7 @@ class CudaProfiler:
@T.inline
def start(self, event_type: Enum, leader: None | tvm.tirx.PrimExpr | bool = None):
if self.profiler_enabled:
- T.timer_start_cuda(
+ T.cuda.timer_start(
event_type,
self.buffer.data,
self.profiler_tag.data,
@@ -636,7 +636,7 @@ class CudaProfiler:
@T.inline
def end(self, event_type: Enum, leader: None | tvm.tirx.PrimExpr | bool = None):
if self.profiler_enabled:
- T.timer_end_cuda(
+ T.cuda.timer_end(
event_type,
self.buffer.data,
self.profiler_tag.data,
@@ -648,7 +648,7 @@ class CudaProfiler:
@T.inline
def finalize(self, leader: None | tvm.tirx.PrimExpr | bool = None):
if self.profiler_enabled:
- T.timer_finalize_cuda(
+ T.cuda.timer_finalize(
self.buffer.data,
self.profiler_tag.data,
self.profiler_write_offset.data,
diff --git a/python/tvm/tirx/compilation_pipeline.py b/python/tvm/tirx/compilation_pipeline.py
index f79af3493f..d2847332b4 100644
--- a/python/tvm/tirx/compilation_pipeline.py
+++ b/python/tvm/tirx/compilation_pipeline.py
@@ -99,33 +99,6 @@ def tirx_pipeline():
return _pipeline, finalize_host_passes, finalize_device_passes
-def trn_pipeline():
- """The Trainium pipeline used in tvm.tirx.build"""
-
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
- """The default lowering passes for TRN backend."""
- tvm.transform.PassContext.current()
- passes = [
- tirx.transform.trn.TrnPrivateBufferAlloc(),
- tirx.transform.trn.TrnNaiveAllocator(),
- tirx.transform.LowerTIRx(),
- tvm.s_tir.transform.DecorateDeviceScope(),
- tirx.transform.StmtSimplify(),
- tirx.transform.LowerTIRxOpaque(),
- tvm.s_tir.transform.LoopPartition(),
- tvm.s_tir.transform.HoistIfThenElse(),
- tirx.transform.StmtSimplify(),
- tirx.transform.RemoveNoOp(),
- tirx.transform.AnnotateEntryFunc(),
- tirx.transform.SplitHostDevice(),
- tirx.transform.MakePackedAPI(),
- ]
- return tvm.ir.transform.Sequential(passes)(mod)
-
- return _pipeline, finalize_host_passes, finalize_device_passes_trn
-
-
def finalize_host_passes(): # pylint: disable=unused-argument
"""The default finalization passes for TIR backend."""
host_pass_list = [
@@ -153,14 +126,14 @@ def finalize_device_passes_tirx(): # pylint: disable=unused-argument
return tvm.ir.transform.Sequential(device_pass_list)
-def finalize_device_passes_trn(): # pylint: disable=unused-argument
- """The default finalization passes for TRN backend."""
- device_pass_list = [tirx.transform.StmtSimplify()]
- return tvm.ir.transform.Sequential(device_pass_list)
+# global map of pre-built pipelines
+PIPELINE_MAP = {"default": default_tir_pipeline, "tirx": tirx_pipeline}
-# global map of pre-built pipelines
-PIPELINE_MAP = {"default": default_tir_pipeline, "tirx": tirx_pipeline, "trn": trn_pipeline}
+def register_tir_pipeline(name: str, pipeline_factory) -> None:
+ """Register a named TIR pipeline factory."""
+
+ PIPELINE_MAP[name] = pipeline_factory
def get_tir_pipeline(name: str | None = None, **kwargs) -> tvm.transform.Pass:
diff --git a/python/tvm/tirx/lang/alloc_pool.py b/python/tvm/tirx/lang/alloc_pool.py
index 48bb9929c6..b45bcef60c 100644
--- a/python/tvm/tirx/lang/alloc_pool.py
+++ b/python/tvm/tirx/lang/alloc_pool.py
@@ -2,528 +2,7 @@
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""SMEM and TMEM bump-allocator pools for TIRX kernels."""
+# to you under the Apache License, Version 2.0.
+"""Compatibility redirect for CUDA allocation pool helpers."""
-from __future__ import annotations
-
-import functools
-import operator
-
-from tvm import DataType
-from tvm.tirx.layout import S, TCol, TileLayout, TLane
-
-# ---------------------------------------------------------------------------
-# ir_builder helpers — imported lazily to avoid circular deps at module level
-# ---------------------------------------------------------------------------
-
-_ir = None
-
-
-def _get_ir():
- global _ir
- if _ir is None:
- from tvm.tirx.script.builder import ir as _mod
-
- _ir = _mod
- return _ir
-
-
-def _get_frame():
- from tvm.tirx.script.builder import frame
-
- return frame
-
-
-# ---------------------------------------------------------------------------
-# Shared utilities
-# ---------------------------------------------------------------------------
-
-_POOL_UNSET = object()
-
-
-def _default_tmem_layout(rows, cols):
- return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)])
-
-
-def _emit_stmt(expr):
- ir = _get_ir()
- ir.add_to_parent(ir.evaluate(expr))
-
-
-def _shape_product(shape):
- return functools.reduce(operator.mul, shape, 1)
-
-
-def _auto_swizzle_mode(dtype):
- """Select the default MMA swizzle mode for a shared-memory allocation."""
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import SwizzleMode
-
- del dtype
- return SwizzleMode.SWIZZLE_128B_ATOM
-
-
-def _swizzle_atom_bytes(swizzle_mode):
- """Return the row width (in bytes) of one swizzle atom for *swizzle_mode*."""
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import SwizzleMode
-
- return {
- SwizzleMode.SWIZZLE_NONE: 0,
- SwizzleMode.SWIZZLE_32B_ATOM: 32,
- SwizzleMode.SWIZZLE_64B_ATOM: 64,
- SwizzleMode.SWIZZLE_128B_ATOM: 128,
- }[swizzle_mode]
-
-
-def _suggest_swizzle_for_row_bytes(row_bytes):
- """Pick the largest valid swizzle mode whose atom row fits within *row_bytes*."""
-
- for atom_bytes, mode in (
- (128, "SWIZZLE_128B_ATOM"),
- (64, "SWIZZLE_64B_ATOM"),
- (32, "SWIZZLE_32B_ATOM"),
- ):
- if row_bytes >= atom_bytes and row_bytes % atom_bytes == 0:
- return mode
- return "SWIZZLE_NONE"
-
-
-def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
- """Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible.
-
- ``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]``
- over the last two logical dimensions of *shape*. If the row width or row count of
- the request is smaller than (or not a multiple of) the atom, the underlying
- ``Layout.tile_to`` lowers to a ``floordiv``/``floormod`` by zero and raises an
- opaque internal "Divide by zero" diagnostic from ``tile_tile_ops.cc``. Catch the
- misconfiguration here so callers see *what* is wrong and *how* to fix it.
-
- Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
- """
- from tvm.tirx.operator.tile_primitive.cuda.tma_utils import SwizzleMode
-
- if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
- return
-
- if len(shape) < 2:
- raise ValueError(
- f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; "
- f"swizzled MMA layouts tile over the last two dims (rows, cols). "
- f"Use swizzle_mode='none' for 1-D allocations."
- )
-
- # Only validate concrete int dims; symbolic dims fall through (the analyzer
- # in C++ will still ICHECK on them, but at least we don't false-positive).
- rows = shape[-2]
- cols = shape[-1]
- if not (isinstance(rows, int) and isinstance(cols, int)):
- return
-
- dtype_bytes = DataType(dtype).bits // 8
- if dtype_bytes == 0:
- # Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so
- # use a fractional check expressed via bits.
- col_bits = cols * DataType(dtype).bits
- atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8
- if col_bits < atom_bits or col_bits % atom_bits != 0:
... 15282 lines suppressed ...