This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 3608afc3809bf012645693c8748bb56167992e6f Author: tqchen <[email protected]> AuthorDate: Fri Oct 16 02:33:56 2026 +0000 [DOCS] Add standalone TIRx Python API reference --- docs/.gitignore | 19 + docs/_templates/autosummary/callable.rst | 20 + docs/_templates/autosummary/class.rst | 32 + docs/conf.py | 115 ++++ docs/htaccess.in | 20 + docs/index.rst | 28 + docs/reference/python/backend-configuration.rst | 51 ++ docs/reference/python/cuda-script-namespaces.rst | 267 ++++++++ docs/reference/python/cuda.rst | 438 ++++++++++++ docs/reference/python/index.rst | 38 ++ docs/reference/python/ir.rst | 297 ++++++++ docs/reference/python/other-backends.rst | 271 ++++++++ .../python/runtime-and-top-level-utilities.rst | 117 ++++ docs/reference/python/script-and-printer.rst | 744 +++++++++++++++++++++ docs/reference/python/support.rst | 179 +++++ docs/reference/python/symbolic-analysis.rst | 132 ++++ .../python/target-and-code-generation.rst | 107 +++ docs/reference/python/testing.rst | 104 +++ docs/reference/python/tirx.rst | 456 +++++++++++++ docs/requirements.txt | 21 + .../tvm_tirx/backend/cuda/lang/tile_scheduler.py | 150 +++-- python/tvm_tirx/tirx/script/ir_builder/triton.py | 7 +- 22 files changed, 3540 insertions(+), 73 deletions(-) diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 0000000000..fabfa51909 --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1,19 @@ +# 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. + +_build/ +reference/python/generated/ diff --git a/docs/_templates/autosummary/callable.rst b/docs/_templates/autosummary/callable.rst new file mode 100644 index 0000000000..becfdba148 --- /dev/null +++ b/docs/_templates/autosummary/callable.rst @@ -0,0 +1,20 @@ +.. 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. + +{{ fullname | escape | underline }} + +.. autofunction:: {{ fullname }} diff --git a/docs/_templates/autosummary/class.rst b/docs/_templates/autosummary/class.rst new file mode 100644 index 0000000000..a77c1fc00f --- /dev/null +++ b/docs/_templates/autosummary/class.rst @@ -0,0 +1,32 @@ +.. 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. + +{{ fullname | escape | underline }} + +.. currentmodule:: {{ module }} + +{% if fullname == "tvm_tirx.ir.Var" %} +.. raw:: html + + <span id="tvm_tirx.tirx.Var"></span> +{% endif %} + +.. autoclass:: {{ objname }} + +{% if fullname == "tvm_tirx.runtime.DataType" %} +.. autoclass:: tvm_ffi.dtype +{% endif %} diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000000..423e399505 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,115 @@ +# 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. + +import os +import shutil +import sys +from pathlib import Path + +from sphinx import addnodes + +os.environ["TVM_FFI_BUILD_DOCS"] = "1" + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "python")) + +project = "tvm-tirx" +html_title = project +author = "Apache TVM Community" +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", + "sphinx_reredirects", +] +autosummary_generate = True +templates_path = ["_templates"] +exclude_patterns = ["_build"] +html_theme = "sphinx_book_theme" + +autodoc_default_options = { + "members": True, + "undoc-members": True, + "show-inheritance": True, + "inherited-members": False, + "member-order": "bysource", +} +autodoc_typehints = "description" +autodoc_preserve_defaults = True +napoleon_use_rtype = False +napoleon_use_ivar = True +napoleon_custom_sections = ["Usage"] + +# Match the FFI reference's optional PyTorch import convention. +autodoc_mock_imports = ["torch"] + +html_theme_options = { + "repository_url": "https://github.com/apache/tvm";, + "repository_branch": "tirx-standalone-root", + "path_to_docs": "docs", + "use_repository_button": True, + "show_toc_level": 2, + "max_navbar_depth": 2, + "collapse_navbar": True, +} + +# Case-insensitive filesystems need distinct class/function stub names. +autosummary_filename_map = { + "tvm_tirx.runtime.Device": "tvm_tirx.runtime.Device_class", + "tvm_tirx.runtime.device": "tvm_tirx.runtime.device_function", + "tvm_tirx.runtime.Tensor": "tvm_tirx.runtime.Tensor_class", + "tvm_tirx.runtime.tensor": "tvm_tirx.runtime.tensor_function", + "tvm_tirx.ir.prim.Div": "tvm_tirx.ir.prim.Div_class", + "tvm_tirx.ir.prim.div": "tvm_tirx.ir.prim.div_function", + "tvm_tirx.ir.prim.FloorDiv": "tvm_tirx.ir.prim.FloorDiv_class", + "tvm_tirx.ir.prim.floordiv": "tvm_tirx.ir.prim.floordiv_function", + "tvm_tirx.ir.prim.FloorMod": "tvm_tirx.ir.prim.FloorMod_class", + "tvm_tirx.ir.prim.floormod": "tvm_tirx.ir.prim.floormod_function", + "tvm_tirx.ir.prim.Max": "tvm_tirx.ir.prim.Max_class", + "tvm_tirx.ir.prim.max": "tvm_tirx.ir.prim.max_function", + "tvm_tirx.ir.prim.Min": "tvm_tirx.ir.prim.Min_class", + "tvm_tirx.ir.prim.min": "tvm_tirx.ir.prim.min_function", + "tvm_tirx.support.utils.FileLock": "tvm_tirx.support.utils.FileLock_class", + "tvm_tirx.support.utils.filelock": "tvm_tirx.support.utils.filelock_function", +} + + +redirects = { + "reference/python/generated/tvm_tirx.tirx.Var": "tvm_tirx.ir.Var.html", + "reference/python/generated/tvm_tirx.runtime.Tensor": "tvm_tirx.runtime.Tensor_class.html", +} + + +def _qualify_map_type_parameters(app, doctree): + """Keep generic Map parameters distinct from unrelated CUDA attributes.""" + for ref in doctree.findall(addnodes.pending_xref): + if ( + ref.get("py:module") == "tvm_tirx.ir" + and ref.get("py:class") == "Map" + and ref.get("reftarget") in {"K", "V"} + ): + ref["reftarget"] = "tvm_ffi.container." + ref["reftarget"] + + +def _copy_apache_config(app, exception): + if exception is None and app.builder.format == "html": + shutil.copyfile(Path(__file__).with_name("htaccess.in"), Path(app.outdir) / ".htaccess") + + +def setup(app): + app.connect("doctree-read", _qualify_map_type_parameters) + app.connect("build-finished", _copy_apache_config) diff --git a/docs/htaccess.in b/docs/htaccess.in new file mode 100644 index 0000000000..6400d8b820 --- /dev/null +++ b/docs/htaccess.in @@ -0,0 +1,20 @@ +# 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. + +<FilesMatch "\.Var\.(html|rst)$"> + SetHandler default-handler +</FilesMatch> diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 0000000000..4d675c39c2 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,28 @@ +.. 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. + +tvm-tirx +======== + +Table of Contents +----------------- + +.. toctree:: + :maxdepth: 1 + :caption: Reference + + reference/python/index.rst diff --git a/docs/reference/python/backend-configuration.rst b/docs/reference/python/backend-configuration.rst new file mode 100644 index 0000000000..032d1e23a0 --- /dev/null +++ b/docs/reference/python/backend-configuration.rst @@ -0,0 +1,51 @@ +.. 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 Configuration +===================== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.backend +---------------- + +.. currentmodule:: tvm_tirx.backend + +.. autosummary:: + :toctree: generated/ + + is_loaded + load + + +tvm_tirx.backend.config +----------------------- + +.. currentmodule:: tvm_tirx.backend.config + +.. autosummary:: + :toctree: generated/ + + add_backend_config_argument + argparse_backend_config + backend_config_json + copy_backend_config + merge_backend_configs + parse_backend_config + prepare_target diff --git a/docs/reference/python/cuda-script-namespaces.rst b/docs/reference/python/cuda-script-namespaces.rst new file mode 100644 index 0000000000..f459949ffc --- /dev/null +++ b/docs/reference/python/cuda-script-namespaces.rst @@ -0,0 +1,267 @@ +.. 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 Script Namespaces +====================== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.backend.cuda.script.iket +--------------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.script.iket + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + mark + official_event + range_end + range_pop + range_push + range_start + sentinel_token + + +tvm_tirx.backend.cuda.script.tcgen05 +------------------------------------ + +.. currentmodule:: tvm_tirx.backend.cuda.script.tcgen05 + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + encode_instr_descriptor + encode_instr_descriptor_block_scaled + encode_matrix_descriptor + + +tvm_tirx.backend.cuda.script.tile +--------------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.script.tile + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + apply + ewise + + +tvm_tirx.backend.cuda.script.wgmma +---------------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.script.wgmma + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + encode_matrix_descriptor + noop_barrier + + +tvm_tirx.tirx.script.nvshmem +---------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.nvshmem + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + barrier_all + fence + getmem_nbi + my_pe + n_pes + putmem_nbi + putmem_signal_nbi + quiet + signal_op + wait_until + + +tvm_tirx.tirx.script.nvshmem.getmem_nbi +--------------------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.nvshmem.getmem_nbi + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + block + warp + + +tvm_tirx.tirx.script.nvshmem.putmem_nbi +--------------------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.nvshmem.putmem_nbi + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + block + warp + + +tvm_tirx.tirx.script.nvshmem.putmem_signal_nbi +---------------------------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.nvshmem.putmem_signal_nbi + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + block + warp + + +tvm_tirx.tirx.script.ptx +------------------------ + +.. currentmodule:: tvm_tirx.tirx.script.ptx + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + abs + activemask + add + addr + and + applypriority + atom + bar + barrier + bfe + bfi + bfind + bmsk + brev + clmad + clusterlaunchcontrol + clz + cnot + copysign + cos + cp + createpolicy + cvt + cvt_pack + cvta + discard + div + dp2a + dp4a + elect_sync + ex2 + fabric + fence + fma + fns + getctarank + griddepcontrol + isspacep + ld + ldmatrix + ldu + lg2 + lop3 + mad + mad24 + mapa + match + max + mbarrier + min + mma + mov + movmatrix + mul + mul24 + multimem_cp + multimem_ld_reduce + multimem_red + multimem_red_async + multimem_st + multimem_st_async + neg + not + or + popc + pred + prefetch + prefetchu + prmt + rcp + red + red_async + redux_sync + rem + rsqrt + sad + selp + set + setmaxnreg + setp + shf + shfl_sync + shl + shr + sin + slct + spcompress + spdecompress + sqrt + st + st_async + st_bulk + stmatrix + sub + szext + tanh + tcgen05 + tensormap_cp_fenceproxy + tensormap_replace + testp + vote_sync + wgmma + xor + + +tvm_tirx.tirx.script.ptx_legacy +------------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.ptx_legacy + +.. autosummary:: + :toctree: generated/ + :template: autosummary/callable.rst + + ldmatrix + mma diff --git a/docs/reference/python/cuda.rst b/docs/reference/python/cuda.rst new file mode 100644 index 0000000000..53c63d4273 --- /dev/null +++ b/docs/reference/python/cuda.rst @@ -0,0 +1,438 @@ +.. 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 +==== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.backend.cuda +--------------------- + +.. currentmodule:: tvm_tirx.backend.cuda + +.. autosummary:: + :toctree: generated/ + + BackendConfig + export_cuda_host + register_backend + script_namespace + script_namespaces + + +tvm_tirx.backend.cuda.axis +-------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.axis + +.. autosummary:: + :toctree: generated/ + + register + + +tvm_tirx.backend.cuda.layout +---------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.layout + +.. autosummary:: + :toctree: generated/ + + SwizzleMode + mma_atom_layout + mma_atom_shape + mma_shared_layout + tmem_datapath_layout + tmem_mma_operand_layout + + +tvm_tirx.backend.cuda.thread_scope +---------------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.thread_scope + +.. autosummary:: + :toctree: generated/ + + ThreadScope + + +tvm_tirx.backend.cuda.compiler +------------------------------ + +.. currentmodule:: tvm_tirx.backend.cuda.compiler + +.. autosummary:: + :toctree: generated/ + + CompilationResult + compile_source + + +tvm_tirx.backend.cuda.iket +-------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.iket + +.. autosummary:: + :toctree: generated/ + + IketProfileError + IketProfileResult + IketProfiler + profile + run + validate_official_environment + + +tvm_tirx.backend.cuda.codegen +----------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.codegen + +.. autosummary:: + :toctree: generated/ + + PTXDataType + device_intrinsic + get_codegen + header_generator + register_codegen + + +tvm_tirx.backend.cuda.lang +-------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.lang + +.. autosummary:: + :toctree: generated/ + + BaseTileScheduler + ClusterPersistentScheduler2D + FlashAttentionLPTScheduler + FlashAttentionLinearScheduler + MBarrier + Pipeline + PipelineState + SMEMPool + TCGen05Bar + TMABar + + +tvm_tirx.backend.cuda.launch +---------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.launch + +.. autosummary:: + :toctree: generated/ + + AccessPolicyWindow + KernelAttributes + LaunchCompletionEvent + LaunchConfig + MemSyncDomainMap + ProgrammaticEvent + + +tvm_tirx.backend.cuda.op +------------------------ + +.. currentmodule:: tvm_tirx.backend.cuda.op + +.. autosummary:: + :toctree: generated/ + + TCGen05InstrDescriptorAttrs + TCGen05InstrDescriptorBlockScaledAttrs + TensorMapEncodeTiledAttr + any_sync + atomic_add + atomic_cas + ballot_sync + bfloat1622float2 + bfloat162float + block_dim + block_idx + bmma_sync + clock64 + cluster_cta_id + cluster_dim + cluster_id + cluster_sync + cta_max + cta_min + cta_pair_id + cta_reduce + cta_sum + cta_sync + cvta_generic_to_shared + decl_tmem + elect_sync + fadd2_rn + fdividef + ffs_u32 + float22bfloat162_rn + float22bfloat162_rn_from_float2 + float22half2 + float2_x + float2_y + float8tohalf8 + float_as_uint + fmul2_rn + fp8x4_e4m3_from_float4 + func_call + get_tmem_addr + grid_dim + grid_sync + half2float + half8tofloat8 + hmax2 + hmin2 + iket_mark + iket_official_event + iket_range_end + iket_range_pop + iket_range_push + iket_range_start + iket_sentinel_token + lane_id + linear_thread_id + make_float2 + mbarrier_wait + mbarrier_wait_acquire_cluster + mma_fill + mma_store + mov_sreg + nano_sleep + nvshmem_barrier_all + nvshmem_fence + nvshmem_getmem_nbi + nvshmem_getmem_nbi_block + nvshmem_getmem_nbi_warp + nvshmem_my_pe + nvshmem_n_pes + nvshmem_putmem_nbi + nvshmem_putmem_nbi_block + nvshmem_putmem_nbi_warp + nvshmem_putmem_signal_nbi + nvshmem_putmem_signal_nbi_block + nvshmem_putmem_signal_nbi_warp + nvshmem_quiet + nvshmem_signal_op + nvshmem_wait_until + printf + ptx_cp_async_legacy + ptx_legacy_ldmatrix + ptx_legacy_mma + reduce_add_sync_u32 + reduce_min_sync_u32 + runtime_instr_desc + sm100_2sm_leader_smem_addr + smem_addr_from_uint64 + syncthreads_and + syncthreads_or + tcgen05_encode_instr_descriptor + tcgen05_encode_instr_descriptor_block_scaled + tcgen05_encode_matrix_descriptor + tensormap_encode_tiled + thread_fence + thread_idx + thread_in_warpgroup + thread_rank + timer_end + timer_finalize + timer_init + timer_start + trap_when_assert_failed + uint_as_float + wait_until + warp_id + warp_in_warpgroup + warp_max + warp_min + warp_reduce + warp_sum + warp_sync + warpgroup_id + warpgroup_sync + wgmma_encode_matrix_descriptor + wgmma_noop_barrier + + +tvm_tirx.backend.cuda.ptx +------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.ptx + +.. autosummary:: + :toctree: generated/ + + PTXNamespace + instruction_forms + instruction_name + + +tvm_tirx.backend.cuda.script +---------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.script + +.. autosummary:: + :toctree: generated/ + + CudaTcgen05Namespace + CudaWgmmaNamespace + IketNamespace + NVSHMEMGetMemNBINamespace + NVSHMEMNamespace + NVSHMEMPutMemNBINamespace + NVSHMEMPutMemSignalNBINamespace + PTXLegacyNamespace + PTXNamespace + STIRNamespace + any_sync + ballot_sync + bfloat1622float2 + bfloat162float + block_dim + block_idx + clock64 + cluster_cta_id + cluster_dim + cluster_id + cluster_sync + cta_pair_id + cta_sync + cvta_generic_to_shared + decl_tmem + elect_sync + fadd2_rn + fdividef + ffs_u32 + float22bfloat162_rn + float22bfloat162_rn_from_float2 + float22half2 + float2_x + float2_y + float8tohalf8 + float_as_uint + fmul2_rn + fp8x4_e4m3_from_float4 + get_tmem_addr + grid_dim + grid_sync + half2float + half8tofloat8 + hmax2 + hmin2 + iket_mark + iket_official_event + iket_range_end + iket_range_pop + iket_range_push + iket_range_start + iket_sentinel_token + lane_id + linear_thread_id + make_float2 + mbarrier_wait + mbarrier_wait_acquire_cluster + mma_fill + mma_store + nano_sleep + printf + reduce_add_sync_u32 + reduce_min_sync_u32 + runtime_instr_desc + smem_addr_from_uint64 + syncthreads_and + syncthreads_or + tcgen05_encode_instr_descriptor + tcgen05_encode_instr_descriptor_block_scaled + tcgen05_encode_matrix_descriptor + tensormap_encode_tiled + thread_fence + thread_idx + thread_in_warpgroup + thread_rank + trap_when_assert_failed + uint_as_float + warp_id + warp_in_warpgroup + warp_sync + warpgroup_id + warpgroup_sync + wgmma_encode_matrix_descriptor + wgmma_noop_barrier + __shfl_sync + __shfl_up_sync + __shfl_down_sync + __shfl_xor_sync + __activemask + + +tvm_tirx.backend.cuda.tile.exec_context +--------------------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.tile.exec_context + +.. autosummary:: + :toctree: generated/ + + ActiveSet + AxisRange + ExecContext + ExecContextError + LaneBinding + Split + filter_modulo + filter_narrow + initial_A + scope_switch + + +tvm_tirx.backend.cuda.transforms +-------------------------------- + +.. currentmodule:: tvm_tirx.backend.cuda.transforms + +.. autosummary:: + :toctree: generated/ + + BindBackendConfig + LowerIket + SpecializeDeviceHelpers + SpecializeEntryHelpers + +tvm_tirx.backend.cuda.backend_config +------------------------------------ + +.. currentmodule:: tvm_tirx.backend.cuda.backend_config + +.. autosummary:: + :toctree: generated/ + + compiler_options + default_backend_config + resolve_backend_config + validate_backend_config diff --git a/docs/reference/python/index.rst b/docs/reference/python/index.rst new file mode 100644 index 0000000000..f348ca167a --- /dev/null +++ b/docs/reference/python/index.rst @@ -0,0 +1,38 @@ +.. 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. + +Python API +========== + +.. automodule:: tvm_tirx + :no-members: + +.. toctree:: + :maxdepth: 1 + + runtime-and-top-level-utilities + ir + symbolic-analysis + tirx + target-and-code-generation + script-and-printer + backend-configuration + cuda + other-backends + support + testing + cuda-script-namespaces diff --git a/docs/reference/python/ir.rst b/docs/reference/python/ir.rst new file mode 100644 index 0000000000..3912e3af29 --- /dev/null +++ b/docs/reference/python/ir.rst @@ -0,0 +1,297 @@ +.. 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. + +IR +== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.ir +----------- + +.. currentmodule:: tvm_tirx.ir + +.. autosummary:: + :toctree: generated/ + + AnyType + Array + AssertStmt + Attrs + BaseFunc + Bind + Break + Call + CallEffectKind + CallSiteLoc + CallingConv + Constant + Continue + DataTypeImm + DictAttrs + EnvFunc + Evaluate + Expr + ExprOperand + ExprWithOp + FloatImm + For + ForKind + FuncType + GenericConst + GlobalInfo + GlobalVar + IRModule + If + IntImm + LambdaExpr + Location + Map + MissingType + Node + Op + OpaqueExpr + OpaqueType + PrimType + PtrType + Range + RegionStmt + Return + Scriptable + SeqStmt + SourceLoc + SourceName + StagingExpr + Stmt + StringImm + StringType + TensorLoad + TensorRegion + TensorRegionType + TensorStore + Tuple + TupleGetItem + TupleType + Type + UNKNOWN_LOC + UnknownLoc + Var + While + assert_structural_equal + const + is_prim_expr + is_prim_var + load_json + make_node + register_op_attr + reinfer_type + save_json + stmt_list + stmt_seq + + +tvm_tirx.ir.base +---------------- + +.. currentmodule:: tvm_tirx.ir.base + +.. autosummary:: + :toctree: generated/ + + register_dialect_script_method + register_script_printer + +* :obj:`EnvFunc <tvm_tirx.ir.EnvFunc>` +* :obj:`Node <tvm_tirx.ir.Node>` +* :obj:`Scriptable <tvm_tirx.ir.Scriptable>` +* :obj:`assert_structural_equal <tvm_tirx.ir.assert_structural_equal>` +* :obj:`load_json <tvm_tirx.ir.load_json>` +* :obj:`save_json <tvm_tirx.ir.save_json>` + +tvm_tirx.ir.prim +---------------- + +.. currentmodule:: tvm_tirx.ir.prim + +.. autosummary:: + :toctree: generated/ + + Add + And + BinaryOpExpr + BitwiseAnd + BitwiseNot + BitwiseOr + BitwiseXor + Broadcast + Cast + CmpExpr + Div + EQ + FloorDiv + FloorMod + GE + GT + LE + LShift + LT + Let + LogicalExpr + Max + Min + Mod + Mul + NE + Not + Or + RShift + Ramp + Select + Shuffle + Sub + abs + acos + acosh + all + any + asin + asinh + assume + atan + atan2 + atanh + bitwise_and + bitwise_not + bitwise_or + bitwise_xor + ceil + ceildiv + clz + convert + copysign + cos + cosh + div + erf + exp + exp10 + exp2 + expr_deep_equal + floor + floordiv + floormod + fma + fmod + hypot + if_then_else + indexdiv + indexmod + infinity + isfinite + isinf + isnan + ldexp + likely + log + log10 + log1p + log2 + max + max_value + min + min_value + nearbyint + nextafter + popcount + pow + power + round + rsqrt + shift_left + shift_right + sigmoid + sin + sinh + sqrt + tan + tanh + trunc + truncdiv + truncmod + vscale + +* :obj:`Expr <tvm_tirx.ir.Expr>` +* :obj:`const <tvm_tirx.ir.const>` + +tvm_tirx.ir.transform +--------------------- + +.. currentmodule:: tvm_tirx.ir.transform + +.. autosummary:: + :toctree: generated/ + + ModulePass + Pass + PassContext + PassInfo + PassInstrument + PrintIR + Sequential + module_pass + pass_instrument + + +tvm_tirx.ir.transform.instrument +-------------------------------- + +.. currentmodule:: tvm_tirx.ir.transform.instrument + +.. autosummary:: + :toctree: generated/ + + DumpIR + PassPrintingInstrument + PassTimingInstrument + PrintAfterAll + PrintBeforeAll + + +tvm_tirx.ir.unique_name_supply +------------------------------ + +.. currentmodule:: tvm_tirx.ir.unique_name_supply + +.. autosummary:: + :toctree: generated/ + + UniqueNameSupply + + +tvm_tirx.ir.json_compact +------------------------ + +.. currentmodule:: tvm_tirx.ir.json_compact + +.. autosummary:: + :toctree: generated/ + + create_updater + get_version + upgrade_json diff --git a/docs/reference/python/other-backends.rst b/docs/reference/python/other-backends.rst new file mode 100644 index 0000000000..2a00fa2d48 --- /dev/null +++ b/docs/reference/python/other-backends.rst @@ -0,0 +1,271 @@ +.. 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. + +Other Backends +============== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.backend.adreno +----------------------- + +.. currentmodule:: tvm_tirx.backend.adreno + +.. autosummary:: + :toctree: generated/ + + register_backend + + +tvm_tirx.backend.hexagon +------------------------ + +.. currentmodule:: tvm_tirx.backend.hexagon + +.. autosummary:: + :toctree: generated/ + + register_backend + + +tvm_tirx.backend.hexagon.linker +------------------------------- + +.. currentmodule:: tvm_tirx.backend.hexagon.linker + +.. autosummary:: + :toctree: generated/ + + hexagon_link + link_shared + register_linker + + +tvm_tirx.backend.metal +---------------------- + +.. currentmodule:: tvm_tirx.backend.metal + +.. autosummary:: + :toctree: generated/ + + register_backend + script_namespace + script_namespaces + + +tvm_tirx.backend.metal.op +------------------------- + +.. currentmodule:: tvm_tirx.backend.metal.op + +.. autosummary:: + :toctree: generated/ + + cooperative_tensor_fill + cooperative_tensor_load + cooperative_tensor_multiply_accumulate + cooperative_tensor_store + make_filled_simdgroup_matrix + simd_shuffle + simd_shuffle_down + simd_shuffle_up + simdgroup_load + simdgroup_multiply_accumulate + simdgroup_store + + +tvm_tirx.backend.metal.script +----------------------------- + +.. currentmodule:: tvm_tirx.backend.metal.script + +.. autosummary:: + :toctree: generated/ + + MetalNamespace + + +tvm_tirx.backend.opencl +----------------------- + +.. currentmodule:: tvm_tirx.backend.opencl + +.. autosummary:: + :toctree: generated/ + + nd_mem_alloc_with_scope + register_backend + texture2d_load + texture2d_store + + +tvm_tirx.backend.opencl.script +------------------------------ + +.. currentmodule:: tvm_tirx.backend.opencl.script + +.. autosummary:: + :toctree: generated/ + + nd_mem_alloc_with_scope + texture2d_load + texture2d_store + + +tvm_tirx.backend.rocm +--------------------- + +.. currentmodule:: tvm_tirx.backend.rocm + +.. autosummary:: + :toctree: generated/ + + register_backend + + +tvm_tirx.backend.trn +-------------------- + +.. currentmodule:: tvm_tirx.backend.trn + +.. autosummary:: + :toctree: generated/ + + register_backend + script_namespace + script_namespaces + + +tvm_tirx.backend.trn.axis +------------------------- + +.. currentmodule:: tvm_tirx.backend.trn.axis + +.. autosummary:: + :toctree: generated/ + + register + + +tvm_tirx.backend.trn.layout +--------------------------- + +.. currentmodule:: tvm_tirx.backend.trn.layout + +.. autosummary:: + :toctree: generated/ + + is_trainium_layout + + +tvm_tirx.backend.trn.thread_scope +--------------------------------- + +.. currentmodule:: tvm_tirx.backend.trn.thread_scope + +.. autosummary:: + :toctree: generated/ + + ThreadScope + + +tvm_tirx.backend.trn.op +----------------------- + +.. currentmodule:: tvm_tirx.backend.trn.op + +.. autosummary:: + :toctree: generated/ + + 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 + + +tvm_tirx.backend.trn.pipeline +----------------------------- + +.. currentmodule:: tvm_tirx.backend.trn.pipeline + +.. autosummary:: + :toctree: generated/ + + finalize_device_passes_trn + trn_pipeline + + +tvm_tirx.backend.trn.script +--------------------------- + +.. currentmodule:: tvm_tirx.backend.trn.script + +.. autosummary:: + :toctree: generated/ + + NKINamespace + + +tvm_tirx.backend.trn.transform +------------------------------ + +.. currentmodule:: tvm_tirx.backend.trn.transform + +.. autosummary:: + :toctree: generated/ + + LowerTIRx + LowerTrainiumLayout + TrnNaiveAllocator + TrnPrivateTensorAlloc + + +tvm_tirx.backend.vulkan +----------------------- + +.. currentmodule:: tvm_tirx.backend.vulkan + +.. autosummary:: + :toctree: generated/ + + register_backend + + +tvm_tirx.backend.webgpu +----------------------- + +.. currentmodule:: tvm_tirx.backend.webgpu + +.. autosummary:: + :toctree: generated/ + + register_backend diff --git a/docs/reference/python/runtime-and-top-level-utilities.rst b/docs/reference/python/runtime-and-top-level-utilities.rst new file mode 100644 index 0000000000..f67e426e62 --- /dev/null +++ b/docs/reference/python/runtime-and-top-level-utilities.rst @@ -0,0 +1,117 @@ +.. 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. + +Runtime and Top-Level Utilities +=============================== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx +-------- + +.. currentmodule:: tvm_tirx + +.. autosummary:: + :toctree: generated/ + + get_global_func + hexagon + register_global_func + register_object + +* :obj:`DataType <tvm_tirx.runtime.DataType>` +* :obj:`DataTypeCode <tvm_tirx.runtime.DataTypeCode>` +* :obj:`IRModule <tvm_tirx.ir.IRModule>` +* :obj:`Object <tvm_tirx.runtime.Object>` +* :obj:`cpu <tvm_tirx.runtime.cpu>` +* :obj:`cuda <tvm_tirx.runtime.cuda>` +* :obj:`device <tvm_tirx.runtime.device>` +* :obj:`device_from_target <tvm_tirx.target.device_from_target>` +* :obj:`ext_dev <tvm_tirx.runtime.ext_dev>` +* :obj:`metal <tvm_tirx.runtime.metal>` +* :obj:`opencl <tvm_tirx.runtime.opencl>` +* :obj:`rocm <tvm_tirx.runtime.rocm>` +* :obj:`vpi <tvm_tirx.runtime.vpi>` +* :obj:`vulkan <tvm_tirx.runtime.vulkan>` + +tvm_tirx.runtime +---------------- + +.. currentmodule:: tvm_tirx.runtime + +.. autosummary:: + :toctree: generated/ + + DataType + DataTypeCode + Device + MemoryScope + Module + Object + ShapeTuple + Tensor + convert + cpu + cuda + device + empty + enabled + ext_dev + from_dlpack + load_module + metal + num_threads + opencl + rocm + system_lib + tensor + vpi + vulkan + + +tvm_tirx.runtime.memory_scope +----------------------------- + +.. currentmodule:: tvm_tirx.runtime.memory_scope + +* :obj:`MemoryScope <tvm_tirx.runtime.MemoryScope>` + +tvm_tirx.config +--------------- + +.. currentmodule:: tvm_tirx.config + +.. autosummary:: + :toctree: generated/ + + configuration + + +tvm_tirx.libinfo +---------------- + +.. currentmodule:: tvm_tirx.libinfo + +.. autosummary:: + :toctree: generated/ + + configuration + find_include_path + find_libtvm_runtime + package_lib_paths diff --git a/docs/reference/python/script-and-printer.rst b/docs/reference/python/script-and-printer.rst new file mode 100644 index 0000000000..920ef0e72d --- /dev/null +++ b/docs/reference/python/script-and-printer.rst @@ -0,0 +1,744 @@ +.. 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. + +Script and Printer +================== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.script +--------------- + +.. currentmodule:: tvm_tirx.script + +.. autosummary:: + :toctree: generated/ + + register_dialect + register_module_validator + + +tvm_tirx.script.parser +---------------------- + +.. currentmodule:: tvm_tirx.script.parser + +.. autosummary:: + :toctree: generated/ + + from_source + ir_module + make_decorator + make_macro_decorator + pyfunc + register_namespace + register_namespace_initializer + +* :obj:`parse <tvm_tirx.script.parser.from_source>` + +tvm_tirx.script.parser.protocol_registry +---------------------------------------- + +.. currentmodule:: tvm_tirx.script.parser.protocol_registry + +.. autosummary:: + :toctree: generated/ + + register_mutable_decl + register_scalar_annotation + + +tvm_tirx.script.ir_builder +-------------------------- + +.. currentmodule:: tvm_tirx.script.ir_builder + +.. autosummary:: + :toctree: generated/ + + AlreadyEmitted + IRBuilder + IRModuleFrame + Lambda + at_ + check_well_formed_ + constexpr + decl_function + dtype + dynamic + ir_module + meta_var + module_attrs + module_get_attr + module_global_infos + module_member_ + module_set_attr + resolve_global_info_ + resolve_global_info_args + with_at_group_ + +* :obj:`Any <tvm_tirx.ir.AnyType>` +* :obj:`Call <tvm_tirx.ir.Call>` +* :obj:`DataTypeImm <tvm_tirx.ir.DataTypeImm>` +* :obj:`Func <tvm_tirx.ir.FuncType>` +* :obj:`FuncType <tvm_tirx.ir.FuncType>` +* :obj:`GenericConst <tvm_tirx.ir.GenericConst>` +* :obj:`Missing <tvm_tirx.ir.MissingType>` +* :obj:`MissingType <tvm_tirx.ir.MissingType>` +* :obj:`Op <tvm_tirx.ir.Op>` +* :obj:`Opaque <tvm_tirx.ir.OpaqueType>` +* :obj:`Prim <tvm_tirx.ir.PrimType>` +* :obj:`PrimType <tvm_tirx.ir.PrimType>` +* :obj:`Ptr <tvm_tirx.ir.PtrType>` +* :obj:`Range <tvm_tirx.ir.Range>` +* :obj:`String <tvm_tirx.ir.StringType>` +* :obj:`StringImm <tvm_tirx.ir.StringImm>` +* :obj:`StringType <tvm_tirx.ir.StringType>` +* :obj:`TensorRegion <tvm_tirx.ir.TensorRegion>` +* :obj:`TensorRegionType <tvm_tirx.ir.TensorRegionType>` +* :obj:`Tuple <tvm_tirx.ir.Tuple>` +* :obj:`Type <tvm_tirx.ir.Type>` +* :obj:`Var <tvm_tirx.ir.Var>` +* :obj:`make_node <tvm_tirx.ir.make_node>` + +tvm_tirx.script.ir_builder.frame +-------------------------------- + +.. currentmodule:: tvm_tirx.script.ir_builder.frame + +.. autosummary:: + :toctree: generated/ + + AssertFrame + ElseFrame + ForFrame + IfFrame + RegionFrame + StmtFrame + ThenFrame + WhileFrame + +* :obj:`IRModuleFrame <tvm_tirx.script.ir_builder.IRModuleFrame>` + +tvm_tirx.script.ir_builder.ir +----------------------------- + +.. currentmodule:: tvm_tirx.script.ir_builder.ir + +* :obj:`Lambda <tvm_tirx.script.ir_builder.Lambda>` +* :obj:`constexpr <tvm_tirx.script.ir_builder.constexpr>` +* :obj:`dtype <tvm_tirx.script.ir_builder.dtype>` +* :obj:`dynamic <tvm_tirx.script.ir_builder.dynamic>` +* :obj:`meta_var <tvm_tirx.script.ir_builder.meta_var>` +* :obj:`module_attrs <tvm_tirx.script.ir_builder.module_attrs>` +* :obj:`module_get_attr <tvm_tirx.script.ir_builder.module_get_attr>` +* :obj:`module_global_infos <tvm_tirx.script.ir_builder.module_global_infos>` +* :obj:`module_set_attr <tvm_tirx.script.ir_builder.module_set_attr>` + +tvm_tirx.script.ir_builder.stmt +------------------------------- + +.. currentmodule:: tvm_tirx.script.ir_builder.stmt + +.. autosummary:: + :toctree: generated/ + + bind + +* :obj:`add_to_parent <tvm_tirx.tirx.script.add_to_parent>` +* :obj:`assert_ <tvm_tirx.tirx.script.assert_>` +* :obj:`break_ <tvm_tirx.tirx.script.break_>` +* :obj:`continue_ <tvm_tirx.tirx.script.continue_>` +* :obj:`else_ <tvm_tirx.tirx.script.else_>` +* :obj:`evaluate <tvm_tirx.tirx.script.evaluate>` +* :obj:`frame_scope <tvm_tirx.tirx.script.frame_scope>` +* :obj:`grid <tvm_tirx.tirx.script.grid>` +* :obj:`if_ <tvm_tirx.tirx.script.if_>` +* :obj:`region <tvm_tirx.tirx.script.region>` +* :obj:`return_ <tvm_tirx.tirx.script.return_>` +* :obj:`then_ <tvm_tirx.tirx.script.then_>` +* :obj:`unpack_ <tvm_tirx.tirx.script.unpack_>` +* :obj:`while_ <tvm_tirx.tirx.script.while_>` + +tvm_tirx.script.ir_builder.tracing +---------------------------------- + +.. currentmodule:: tvm_tirx.script.ir_builder.tracing + +.. autosummary:: + :toctree: generated/ + + LineTraceScope + + +tvm_tirx.script.printer +----------------------- + +.. currentmodule:: tvm_tirx.script.printer + +.. autosummary:: + :toctree: generated/ + + PrinterConfig + + +tvm_tirx.script.printer.doc +--------------------------- + +.. currentmodule:: tvm_tirx.script.printer.doc + +.. autosummary:: + :toctree: generated/ + + AssertDoc + AssignDoc + AttrAccessDoc + CallDoc + ClassDoc + CommentDoc + DictDoc + Doc + DocStringDoc + ExprDoc + ExprStmtDoc + ForDoc + FunctionDoc + IdDoc + IfDoc + IndexDoc + LambdaDoc + ListDoc + LiteralDoc + OperationDoc + OperationKind + ReturnDoc + ScopeDoc + SliceDoc + StmtBlockDoc + StmtDoc + TupleDoc + WhileDoc + + +tvm_tirx.tirx.script +-------------------- + +.. currentmodule:: tvm_tirx.tirx.script + +.. autosummary:: + :toctree: generated/ + + Let + LetAnnotation + Optional + Tuple + abi_field_get + abi_field_set + abs + acos + acosh + add_to_parent + address_of + alloc_local + alloc_scalar + alloc_shared + alloc_tensor + alloc_workspace + and_ + arg_ + asin + asinh + assert_ + assign + assume + assume_aligned + atan + atan2 + atanh + atomic_add + bf16 + bind + bind_ + bitwise_and + bitwise_not + bitwise_or + bitwise_xor + bool + break_ + call_cpacked + call_cpacked_lowered + call_extern + call_ffi_kernel + call_global_var_ + call_intrin + call_kernel + call_llvm_intrin + call_llvm_pure_intrin + call_packed + call_packed_lowered + call_pure_extern + call_spirv_pure_glsl450 + cast + ceil + ceildiv + check_well_formed_ + clz + compute_scope + continue_ + copysign + cos + cosh + cpu_parallel_barrier + decl_mutable_cell_ + decl_scalar + decl_tensor + device_context + device_entry + device_scope + else_ + emit + eq_ + erf + evaluate + exp + exp10 + exp2 + f16 + f32 + f64 + fabs + float16x16 + float16x2 + float16x32 + float16x4 + float16x64 + float16x8 + float32x16 + float32x2 + float32x32 + float32x4 + float32x64 + float32x8 + float4_e2m1fn + float4_e2m1fnx16 + float4_e2m1fnx2 + float4_e2m1fnx32 + float4_e2m1fnx4 + float4_e2m1fnx64 + float4_e2m1fnx8 + float64x16 + float64x2 + float64x32 + float64x4 + float64x64 + float64x8 + float6_e2m3fn + float6_e2m3fnx16 + float6_e2m3fnx2 + float6_e2m3fnx32 + float6_e2m3fnx4 + float6_e2m3fnx64 + float6_e2m3fnx8 + float6_e3m2fn + float6_e3m2fnx16 + float6_e3m2fnx2 + float6_e3m2fnx32 + float6_e3m2fnx4 + float6_e3m2fnx64 + float6_e3m2fnx8 + float8_e3m4 + float8_e3m4x16 + float8_e3m4x2 + float8_e3m4x32 + float8_e3m4x4 + float8_e3m4x64 + float8_e3m4x8 + float8_e4m3 + float8_e4m3b11fnuz + float8_e4m3b11fnuzx16 + float8_e4m3b11fnuzx2 + float8_e4m3b11fnuzx32 + float8_e4m3b11fnuzx4 + float8_e4m3b11fnuzx64 + float8_e4m3b11fnuzx8 + float8_e4m3fn + float8_e4m3fnuz + float8_e4m3fnuzx16 + float8_e4m3fnuzx2 + float8_e4m3fnuzx32 + float8_e4m3fnuzx4 + float8_e4m3fnuzx64 + float8_e4m3fnuzx8 + float8_e4m3fnx16 + float8_e4m3fnx2 + float8_e4m3fnx32 + float8_e4m3fnx4 + float8_e4m3fnx64 + float8_e4m3fnx8 + float8_e4m3x16 + float8_e4m3x2 + float8_e4m3x32 + float8_e4m3x4 + float8_e4m3x64 + float8_e4m3x8 + float8_e5m2 + float8_e5m2fnuz + float8_e5m2fnuzx16 + float8_e5m2fnuzx2 + float8_e5m2fnuzx32 + float8_e5m2fnuzx4 + float8_e5m2fnuzx64 + float8_e5m2fnuzx8 + float8_e5m2x16 + float8_e5m2x2 + float8_e5m2x32 + float8_e5m2x4 + float8_e5m2x64 + float8_e5m2x8 + float8_e8m0fnu + float8_e8m0fnux16 + float8_e8m0fnux2 + float8_e8m0fnux32 + float8_e8m0fnux4 + float8_e8m0fnux64 + float8_e8m0fnux8 + floor + floordiv + floormod + fma + fmod + for_ + frame_scope + free_workspace + func_attr + func_name_ + func_ret + function + function_ + ge_ + get_active_lane_mask + gpu_active_thread_selector + gpu_dp4a + gpu_fill_fragment + gpu_load_matrix_sync + gpu_mma_sync + gpu_storage_sync + gpu_store_matrix_sync + gpu_thread_allreduce + gpu_thread_filter + gpu_thread_invariant + gpu_thread_return + gpu_warp_activemask + gpu_warp_shuffle + gpu_warp_shuffle_down + gpu_warp_shuffle_up + gpu_warp_shuffle_xor + grid + gt_ + handle + handle_add_byte_offset + hypot + i16 + i32 + i64 + i8 + if_ + if_then_else + if_then_else_ + ignore_loop_partition + index_map + infinity + inline + int16x16 + int16x2 + int16x32 + int16x4 + int16x64 + int16x8 + int32x16 + int32x2 + int32x32 + int32x4 + int32x64 + int32x8 + int64x16 + int64x2 + int64x32 + int64x4 + int64x64 + int64x8 + int8x16 + int8x2 + int8x32 + int8x4 + int8x64 + int8x8 + isfinite + isinf + isnan + isnullptr + jit + kernel_replace_point + launch_thread + ldexp + le_ + likely + local_scalar + log + log10 + log1p + log2 + logical_and + logical_not + logical_or + lt_ + macro + masked_load + masked_store + max + max_value + meta_class + min + min_value + ne_ + nearbyint + nextafter + not_ + or_ + parallel + parallel_launch + popcount + pow + prefetch + ptr + ptr_byte_offset + range_ + region + register_script_namespace + reinterpret + resolve_global_info_ + return_ + round + rsqrt + select + serial + set_mutable_cell_ + setattr_ + setitem_ + shared_scalar + shift_left + shift_right + sigmoid + sin + sinh + sqrt + stack_alloca + stack_make_dltensor + stack_make_shape + static_assert + static_handle + tan + tanh + target + tensor_data_ptr + tensor_indices + tensor_store + then_ + throw_last_error + tmem + trunc + truncdiv + truncmod + u16 + u32 + u64 + u8 + uint16x16 + uint16x2 + uint16x32 + uint16x4 + uint16x64 + uint16x8 + uint32x16 + uint32x2 + uint32x32 + uint32x4 + uint32x64 + uint32x8 + uint64x16 + uint64x2 + uint64x32 + uint64x4 + uint64x64 + uint64x8 + uint8x16 + uint8x2 + uint8x32 + uint8x4 + uint8x64 + uint8x8 + undef + unpack_ + unroll + vector_combine + vector_high + vector_low + vectorized + void + vscale + while_ + +* :obj:`Add <tvm_tirx.ir.prim.Add>` +* :obj:`AlreadyEmitted <tvm_tirx.script.ir_builder.AlreadyEmitted>` +* :obj:`And <tvm_tirx.ir.prim.And>` +* :obj:`Any <tvm_tirx.ir.AnyType>` +* :obj:`Axis <tvm_tirx.tirx.layout.Axis>` +* :obj:`BitwiseAnd <tvm_tirx.ir.prim.BitwiseAnd>` +* :obj:`BitwiseNot <tvm_tirx.ir.prim.BitwiseNot>` +* :obj:`BitwiseOr <tvm_tirx.ir.prim.BitwiseOr>` +* :obj:`BitwiseXor <tvm_tirx.ir.prim.BitwiseXor>` +* :obj:`Broadcast <tvm_tirx.ir.prim.Broadcast>` +* :obj:`Call <tvm_tirx.ir.Call>` +* :obj:`CallEffectKind <tvm_tirx.ir.CallEffectKind>` +* :obj:`CallFFIKernelAttr <tvm_tirx.tirx.CallFFIKernelAttr>` +* :obj:`Cast <tvm_tirx.ir.prim.Cast>` +* :obj:`ComposeLayout <tvm_tirx.tirx.ComposeLayout>` +* :obj:`DataTypeImm <tvm_tirx.ir.DataTypeImm>` +* :obj:`Div <tvm_tirx.ir.prim.Div>` +* :obj:`EQ <tvm_tirx.ir.prim.EQ>` +* :obj:`FloatImm <tvm_tirx.ir.FloatImm>` +* :obj:`FloorDiv <tvm_tirx.ir.prim.FloorDiv>` +* :obj:`FloorMod <tvm_tirx.ir.prim.FloorMod>` +* :obj:`Func <tvm_tirx.ir.FuncType>` +* :obj:`FuncType <tvm_tirx.ir.FuncType>` +* :obj:`GE <tvm_tirx.ir.prim.GE>` +* :obj:`GT <tvm_tirx.ir.prim.GT>` +* :obj:`GenericConst <tvm_tirx.ir.GenericConst>` +* :obj:`IRBuilder <tvm_tirx.script.ir_builder.IRBuilder>` +* :obj:`IRModuleFrame <tvm_tirx.script.ir_builder.IRModuleFrame>` +* :obj:`IntImm <tvm_tirx.ir.IntImm>` +* :obj:`Iter <tvm_tirx.tirx.layout.Iter>` +* :obj:`LE <tvm_tirx.ir.prim.LE>` +* :obj:`LShift <tvm_tirx.ir.prim.LShift>` +* :obj:`LT <tvm_tirx.ir.prim.LT>` +* :obj:`Lambda <tvm_tirx.script.ir_builder.Lambda>` +* :obj:`Layout <tvm_tirx.tirx.Layout>` +* :obj:`Max <tvm_tirx.ir.prim.Max>` +* :obj:`MemoryScope <tvm_tirx.runtime.MemoryScope>` +* :obj:`Min <tvm_tirx.ir.prim.Min>` +* :obj:`Missing <tvm_tirx.ir.MissingType>` +* :obj:`MissingType <tvm_tirx.ir.MissingType>` +* :obj:`Mod <tvm_tirx.ir.prim.Mod>` +* :obj:`Mul <tvm_tirx.ir.prim.Mul>` +* :obj:`NE <tvm_tirx.ir.prim.NE>` +* :obj:`Not <tvm_tirx.ir.prim.Not>` +* :obj:`Op <tvm_tirx.ir.Op>` +* :obj:`Opaque <tvm_tirx.ir.OpaqueType>` +* :obj:`Or <tvm_tirx.ir.prim.Or>` +* :obj:`Prim <tvm_tirx.ir.PrimType>` +* :obj:`PrimType <tvm_tirx.ir.PrimType>` +* :obj:`Ptr <tvm_tirx.ir.PtrType>` +* :obj:`RShift <tvm_tirx.ir.prim.RShift>` +* :obj:`Ramp <tvm_tirx.ir.prim.Ramp>` +* :obj:`Range <tvm_tirx.ir.Range>` +* :obj:`Select <tvm_tirx.ir.prim.Select>` +* :obj:`Shuffle <tvm_tirx.ir.prim.Shuffle>` +* :obj:`String <tvm_tirx.ir.StringType>` +* :obj:`StringImm <tvm_tirx.ir.StringImm>` +* :obj:`StringType <tvm_tirx.ir.StringType>` +* :obj:`Sub <tvm_tirx.ir.prim.Sub>` +* :obj:`Tensor <tvm_tirx.tirx.TensorType>` +* :obj:`TensorLoad <tvm_tirx.tirx.TensorLoad>` +* :obj:`TensorMap <tvm_tirx.tirx.TensorMapType>` +* :obj:`TensorRegion <tvm_tirx.ir.TensorRegion>` +* :obj:`TensorRegionType <tvm_tirx.ir.TensorRegionType>` +* :obj:`ThreadScope <tvm_tirx.tirx.ThreadScope>` +* :obj:`TileLayout <tvm_tirx.tirx.TileLayout>` +* :obj:`Type <tvm_tirx.ir.Type>` +* :obj:`Var <tvm_tirx.ir.Var>` +* :obj:`at_ <tvm_tirx.script.ir_builder.at_>` +* :obj:`bfloat16 <tvm_tirx.tirx.script.bf16>` +* :obj:`boolean <tvm_tirx.tirx.script.bool>` +* :obj:`broadcast <tvm_tirx.ir.prim.Broadcast>` +* :obj:`constexpr <tvm_tirx.script.ir_builder.constexpr>` +* :obj:`decl_function <tvm_tirx.script.ir_builder.decl_function>` +* :obj:`dtype <tvm_tirx.script.ir_builder.dtype>` +* :obj:`dynamic <tvm_tirx.script.ir_builder.dynamic>` +* :obj:`emit_ <tvm_tirx.tirx.script.emit>` +* :obj:`float16 <tvm_tirx.tirx.script.f16>` +* :obj:`float32 <tvm_tirx.tirx.script.f32>` +* :obj:`float64 <tvm_tirx.tirx.script.f64>` +* :obj:`func_ret_type_ <tvm_tirx.tirx.script.func_ret>` +* :obj:`int16 <tvm_tirx.tirx.script.i16>` +* :obj:`int32 <tvm_tirx.tirx.script.i32>` +* :obj:`int64 <tvm_tirx.tirx.script.i64>` +* :obj:`int8 <tvm_tirx.tirx.script.i8>` +* :obj:`ir_module <tvm_tirx.script.parser.ir_module>` +* :obj:`llvm_lookup_intrinsic_id <tvm_tirx.codegen.llvm.lookup_intrinsic_id>` +* :obj:`make_node <tvm_tirx.ir.make_node>` +* :obj:`meta_var <tvm_tirx.script.ir_builder.meta_var>` +* :obj:`module_attrs <tvm_tirx.script.ir_builder.module_attrs>` +* :obj:`module_get_attr <tvm_tirx.script.ir_builder.module_get_attr>` +* :obj:`module_global_infos <tvm_tirx.script.ir_builder.module_global_infos>` +* :obj:`module_member_ <tvm_tirx.script.ir_builder.module_member_>` +* :obj:`module_set_attr <tvm_tirx.script.ir_builder.module_set_attr>` +* :obj:`pyfunc <tvm_tirx.script.parser.pyfunc>` +* :obj:`ramp <tvm_tirx.ir.prim.Ramp>` +* :obj:`resolve_global_info_args <tvm_tirx.script.ir_builder.resolve_global_info_args>` +* :obj:`smem <tvm_tirx.tirx.script.alloc_shared>` +* :obj:`uint16 <tvm_tirx.tirx.script.u16>` +* :obj:`uint32 <tvm_tirx.tirx.script.u32>` +* :obj:`uint64 <tvm_tirx.tirx.script.u64>` +* :obj:`uint8 <tvm_tirx.tirx.script.u8>` +* :obj:`with_at_group_ <tvm_tirx.script.ir_builder.with_at_group_>` + +tvm_tirx.tirx.script.jit +------------------------ + +.. currentmodule:: tvm_tirx.tirx.script.jit + +.. autosummary:: + :toctree: generated/ + + TIRJit + +* :obj:`OptionalAnnotation <tvm_tirx.tirx.script.Optional>` + +tvm_tirx.tirx.script.ir_builder.external_kernel +----------------------------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.ir_builder.external_kernel + +.. autosummary:: + :toctree: generated/ + + BaseKernel + SourceKernel + call_kernel + + +tvm_tirx.tirx.script.ir_builder.triton +-------------------------------------- + +.. currentmodule:: tvm_tirx.tirx.script.ir_builder.triton + +.. autosummary:: + :toctree: generated/ + + TritonKernel diff --git a/docs/reference/python/support.rst b/docs/reference/python/support.rst new file mode 100644 index 0000000000..712db8d6e9 --- /dev/null +++ b/docs/reference/python/support.rst @@ -0,0 +1,179 @@ +.. 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. + +Support +======= + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.support +---------------- + +.. currentmodule:: tvm_tirx.support + +.. autosummary:: + :toctree: generated/ + + describe + libinfo + + +tvm_tirx.support.cc +------------------- + +.. currentmodule:: tvm_tirx.support.cc + +.. autosummary:: + :toctree: generated/ + + create_executable + create_shared + create_staticlib + cross_compiler + get_cc + get_global_symbol_section_map + get_target_by_dump_machine + + +tvm_tirx.support.clang +---------------------- + +.. currentmodule:: tvm_tirx.support.clang + +.. autosummary:: + :toctree: generated/ + + create_llvm + find_clang + + +tvm_tirx.support.ndk +-------------------- + +.. currentmodule:: tvm_tirx.support.ndk + +.. autosummary:: + :toctree: generated/ + + create_shared + create_staticlib + get_global_symbol_section_map + + +tvm_tirx.support.nvcc +--------------------- + +.. currentmodule:: tvm_tirx.support.nvcc + +.. autosummary:: + :toctree: generated/ + + callback_libdevice_path + compile_cuda + find_cuda_path + find_libdevice_path + find_nvshmem_paths + get_cuda_version + get_target_compute_version + have_bf16 + have_cudagraph + have_fp16 + have_fp4 + have_fp8 + have_int8 + have_tensorcore + parse_compute_version + tvm_callback_cuda_compile + + +tvm_tirx.support.popen_pool +--------------------------- + +.. currentmodule:: tvm_tirx.support.popen_pool + +.. autosummary:: + :toctree: generated/ + + MapResult + PopenPoolExecutor + PopenWorker + StatusKind + kill_child_processes + + +tvm_tirx.support.rocm +--------------------- + +.. currentmodule:: tvm_tirx.support.rocm + +.. autosummary:: + :toctree: generated/ + + callback_rocm_bitcode_path + callback_rocm_link + find_lld + find_rocm_path + get_rocm_arch + have_matrixcore + parse_compute_version + rocm_link + + +tvm_tirx.support.tar +-------------------- + +.. currentmodule:: tvm_tirx.support.tar + +.. autosummary:: + :toctree: generated/ + + normalize_file_list_by_unpacking_tars + tar + untar + + +tvm_tirx.support.utils +---------------------- + +.. currentmodule:: tvm_tirx.support.utils + +.. autosummary:: + :toctree: generated/ + + DirectoryCreatedPastAtExit + FileLock + TempDirectory + filelock + is_source_path + tempdir + which + + +tvm_tirx.support.xcode +---------------------- + +.. currentmodule:: tvm_tirx.support.xcode + +.. autosummary:: + :toctree: generated/ + + compile_coreml + compile_metal + create_dylib + xcrun diff --git a/docs/reference/python/symbolic-analysis.rst b/docs/reference/python/symbolic-analysis.rst new file mode 100644 index 0000000000..10b6f7ef66 --- /dev/null +++ b/docs/reference/python/symbolic-analysis.rst @@ -0,0 +1,132 @@ +.. 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. + +Symbolic Analysis +================= + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.sym +------------ + +.. currentmodule:: tvm_tirx.sym + +.. autosummary:: + :toctree: generated/ + + Analyzer + CompareResult + ConstIntBound + IntSet + IntervalSet + IterMapExpr + IterMark + IterSplitExpr + IterSumExpr + ModularSet + ProofStrength + allow_uint_as_index + deduce_bound + detect_clip_bound + detect_iter_map + detect_linear_equation + estimate_region_lower_bound + estimate_region_strict_bound + estimate_region_upper_bound + inverse_affine_iter_map + iter_map_simplify + normalize_iter_map_to_expr + normalize_to_iter_sum + subspace_divide + + +tvm_tirx.sym.analyzer +--------------------- + +.. currentmodule:: tvm_tirx.sym.analyzer + +.. autosummary:: + :toctree: generated/ + + ConstraintScope + +* :obj:`Analyzer <tvm_tirx.sym.Analyzer>` +* :obj:`CompareResult <tvm_tirx.sym.CompareResult>` +* :obj:`ConstIntBound <tvm_tirx.sym.ConstIntBound>` +* :obj:`ModularSet <tvm_tirx.sym.ModularSet>` +* :obj:`ProofStrength <tvm_tirx.sym.ProofStrength>` +* :obj:`allow_uint_as_index <tvm_tirx.sym.allow_uint_as_index>` + +tvm_tirx.sym.int_set +-------------------- + +.. currentmodule:: tvm_tirx.sym.int_set + +.. autosummary:: + :toctree: generated/ + + neg_inf + pos_inf + union_lower_bound + +* :obj:`IntSet <tvm_tirx.sym.IntSet>` +* :obj:`IntervalSet <tvm_tirx.sym.IntervalSet>` +* :obj:`estimate_region_lower_bound <tvm_tirx.sym.estimate_region_lower_bound>` +* :obj:`estimate_region_strict_bound <tvm_tirx.sym.estimate_region_strict_bound>` +* :obj:`estimate_region_upper_bound <tvm_tirx.sym.estimate_region_upper_bound>` + +tvm_tirx.sym.iter_affine_map +---------------------------- + +.. currentmodule:: tvm_tirx.sym.iter_affine_map + +.. autosummary:: + :toctree: generated/ + + IterMapLevel + IterMapResult + +* :obj:`IterMapExpr <tvm_tirx.sym.IterMapExpr>` +* :obj:`IterMark <tvm_tirx.sym.IterMark>` +* :obj:`IterSplitExpr <tvm_tirx.sym.IterSplitExpr>` +* :obj:`IterSumExpr <tvm_tirx.sym.IterSumExpr>` +* :obj:`detect_iter_map <tvm_tirx.sym.detect_iter_map>` +* :obj:`inverse_affine_iter_map <tvm_tirx.sym.inverse_affine_iter_map>` +* :obj:`iter_map_simplify <tvm_tirx.sym.iter_map_simplify>` +* :obj:`normalize_iter_map_to_expr <tvm_tirx.sym.normalize_iter_map_to_expr>` +* :obj:`normalize_to_iter_sum <tvm_tirx.sym.normalize_to_iter_sum>` +* :obj:`subspace_divide <tvm_tirx.sym.subspace_divide>` + +tvm_tirx.sym.pattern +-------------------- + +.. currentmodule:: tvm_tirx.sym.pattern + +* :obj:`detect_clip_bound <tvm_tirx.sym.detect_clip_bound>` +* :obj:`detect_linear_equation <tvm_tirx.sym.detect_linear_equation>` + +tvm_tirx.sym.utils +------------------ + +.. currentmodule:: tvm_tirx.sym.utils + +.. autosummary:: + :toctree: generated/ + + get_const_int diff --git a/docs/reference/python/target-and-code-generation.rst b/docs/reference/python/target-and-code-generation.rst new file mode 100644 index 0000000000..626c311d19 --- /dev/null +++ b/docs/reference/python/target-and-code-generation.rst @@ -0,0 +1,107 @@ +.. 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. + +Target and Code Generation +========================== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.target +--------------- + +.. currentmodule:: tvm_tirx.target + +.. autosummary:: + :toctree: generated/ + + Target + TargetKind + device_from_target + list_tags + register_tag + + +tvm_tirx.target.target +---------------------- + +.. currentmodule:: tvm_tirx.target.target + +.. autosummary:: + :toctree: generated/ + + TargetFeatures + +* :obj:`Target <tvm_tirx.target.Target>` +* :obj:`TargetKind <tvm_tirx.target.TargetKind>` +* :obj:`device_from_target <tvm_tirx.target.device_from_target>` + +tvm_tirx.target.tag +------------------- + +.. currentmodule:: tvm_tirx.target.tag + +* :obj:`list_tags <tvm_tirx.target.list_tags>` +* :obj:`register_tag <tvm_tirx.target.register_tag>` + +tvm_tirx.target.detect_target +----------------------------- + +.. currentmodule:: tvm_tirx.target.detect_target + +.. autosummary:: + :toctree: generated/ + + detect_target_from_device + register_device_target_detector + + +tvm_tirx.codegen +---------------- + +.. currentmodule:: tvm_tirx.codegen + +.. autosummary:: + :toctree: generated/ + + Executable + build_module + load_static_library + target_has_features + + +tvm_tirx.codegen.llvm +--------------------- + +.. currentmodule:: tvm_tirx.codegen.llvm + +.. autosummary:: + :toctree: generated/ + + cpu_has_features + get_cpu_archlist + get_cpu_features + get_intrinsic_name + get_system_cpu + get_system_triple + get_system_x86_vendor + get_targets + get_vector_width + is_valid_cpu + lookup_intrinsic_id + version_major diff --git a/docs/reference/python/testing.rst b/docs/reference/python/testing.rst new file mode 100644 index 0000000000..d37fc0fc01 --- /dev/null +++ b/docs/reference/python/testing.rst @@ -0,0 +1,104 @@ +.. 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. + +Testing +======= + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.testing +---------------- + +.. currentmodule:: tvm_tirx.testing + +.. autosummary:: + :toctree: generated/ + + assert_allclose + device_enabled + echo + generate_random_array + identity_after + main + np_dtype_from_str + object_use_count + parameter + run_with_gpu_lock + terminate_self + test_raise_error + + +tvm_tirx.testing.env +-------------------- + +.. currentmodule:: tvm_tirx.testing.env + +.. autosummary:: + :toctree: generated/ + + build_flag_enabled + cuda_arch + has_adreno_opencl + has_cpu_feature + has_cublas + has_cuda + has_cuda_arch + has_cuda_compute + has_cudagraph + has_cudnn + has_gpu + has_hexagon + has_hexagon_toolchain + has_hipblas + has_llvm + has_llvm_min_version + has_matrixcore + has_metal + has_multi_gpu + has_nccl + has_nvcc_version + has_nvptx + has_nvshmem + has_opencl + has_rocm + has_vulkan + has_x86_avx512 + has_x86_vnni + + +tvm_tirx.testing.locking +------------------------ + +.. currentmodule:: tvm_tirx.testing.locking + +* :obj:`run_with_gpu_lock <tvm_tirx.testing.run_with_gpu_lock>` + +tvm_tirx.testing.utils +---------------------- + +.. currentmodule:: tvm_tirx.testing.utils + +* :obj:`assert_allclose <tvm_tirx.testing.assert_allclose>` +* :obj:`device_enabled <tvm_tirx.testing.device_enabled>` +* :obj:`generate_random_array <tvm_tirx.testing.generate_random_array>` +* :obj:`identity_after <tvm_tirx.testing.identity_after>` +* :obj:`main <tvm_tirx.testing.main>` +* :obj:`np_dtype_from_str <tvm_tirx.testing.np_dtype_from_str>` +* :obj:`parameter <tvm_tirx.testing.parameter>` +* :obj:`terminate_self <tvm_tirx.testing.terminate_self>` diff --git a/docs/reference/python/tirx.rst b/docs/reference/python/tirx.rst new file mode 100644 index 0000000000..0219f661b2 --- /dev/null +++ b/docs/reference/python/tirx.rst @@ -0,0 +1,456 @@ +.. 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. + +TIRx +==== + +.. contents:: Table of Contents + :local: + :depth: 1 + +tvm_tirx.tirx +------------- + +.. currentmodule:: tvm_tirx.tirx + +.. autosummary:: + :toctree: generated/ + + CallFFIKernelAttr + ComposeLayout + Function + IndexMap + Layout + TensorLoad + TensorMapType + TensorType + ThreadScope + TileLayout + abi_field_get + abi_field_set + address_of + alloc_workspace + assume_aligned + call_cpacked + call_cpacked_lowered + call_extern + call_ffi_kernel + call_intrin + call_llvm_intrin + call_llvm_pure_intrin + call_packed + call_packed_lowered + call_pure_extern + compile + cpu_parallel_barrier + decl_tensor + free_workspace + get_active_lane_mask + get_default_tir_pipeline + get_tir_pipeline + get_vscale_expr + gpu_dp4a + gpu_fill_fragment + gpu_load_matrix_sync + gpu_mma_sync + gpu_store_matrix_sync + gpu_thread_allreduce + handle_add_byte_offset + ignore_loop_partition + is_tensor_var + isnullptr + logaddexp + ptr_byte_offset + register_intrin_lowering + register_tir_pipeline + reinterpret + renew_def + stack_alloca + stack_make_dltensor + stack_make_shape + tensor_data_ptr + throw_last_error + undef + vector_combine + vector_high + vector_low + +* :obj:`Add <tvm_tirx.ir.prim.Add>` +* :obj:`And <tvm_tirx.ir.prim.And>` +* :obj:`BitwiseAnd <tvm_tirx.ir.prim.BitwiseAnd>` +* :obj:`BitwiseNot <tvm_tirx.ir.prim.BitwiseNot>` +* :obj:`BitwiseOr <tvm_tirx.ir.prim.BitwiseOr>` +* :obj:`BitwiseXor <tvm_tirx.ir.prim.BitwiseXor>` +* :obj:`Cast <tvm_tirx.ir.prim.Cast>` +* :obj:`Div <tvm_tirx.ir.prim.Div>` +* :obj:`EQ <tvm_tirx.ir.prim.EQ>` +* :obj:`FloorDiv <tvm_tirx.ir.prim.FloorDiv>` +* :obj:`FloorMod <tvm_tirx.ir.prim.FloorMod>` +* :obj:`GE <tvm_tirx.ir.prim.GE>` +* :obj:`GT <tvm_tirx.ir.prim.GT>` +* :obj:`LE <tvm_tirx.ir.prim.LE>` +* :obj:`LShift <tvm_tirx.ir.prim.LShift>` +* :obj:`LT <tvm_tirx.ir.prim.LT>` +* :obj:`Let <tvm_tirx.ir.prim.Let>` +* :obj:`Max <tvm_tirx.ir.prim.Max>` +* :obj:`Min <tvm_tirx.ir.prim.Min>` +* :obj:`Mod <tvm_tirx.ir.prim.Mod>` +* :obj:`Mul <tvm_tirx.ir.prim.Mul>` +* :obj:`NE <tvm_tirx.ir.prim.NE>` +* :obj:`Not <tvm_tirx.ir.prim.Not>` +* :obj:`Or <tvm_tirx.ir.prim.Or>` +* :obj:`RShift <tvm_tirx.ir.prim.RShift>` +* :obj:`Select <tvm_tirx.ir.prim.Select>` +* :obj:`Sub <tvm_tirx.ir.prim.Sub>` +* :obj:`Var <tvm_tirx.ir.Var>` +* :obj:`abs <tvm_tirx.ir.prim.abs>` +* :obj:`acos <tvm_tirx.ir.prim.acos>` +* :obj:`acosh <tvm_tirx.ir.prim.acosh>` +* :obj:`all <tvm_tirx.ir.prim.all>` +* :obj:`any <tvm_tirx.ir.prim.any>` +* :obj:`asin <tvm_tirx.ir.prim.asin>` +* :obj:`asinh <tvm_tirx.ir.prim.asinh>` +* :obj:`assume <tvm_tirx.ir.prim.assume>` +* :obj:`atan <tvm_tirx.ir.prim.atan>` +* :obj:`atan2 <tvm_tirx.ir.prim.atan2>` +* :obj:`atanh <tvm_tirx.ir.prim.atanh>` +* :obj:`bitwise_and <tvm_tirx.ir.prim.bitwise_and>` +* :obj:`bitwise_not <tvm_tirx.ir.prim.bitwise_not>` +* :obj:`bitwise_or <tvm_tirx.ir.prim.bitwise_or>` +* :obj:`bitwise_xor <tvm_tirx.ir.prim.bitwise_xor>` +* :obj:`ceil <tvm_tirx.ir.prim.ceil>` +* :obj:`ceildiv <tvm_tirx.ir.prim.ceildiv>` +* :obj:`clz <tvm_tirx.ir.prim.clz>` +* :obj:`convert <tvm_tirx.ir.prim.convert>` +* :obj:`copysign <tvm_tirx.ir.prim.copysign>` +* :obj:`cos <tvm_tirx.ir.prim.cos>` +* :obj:`cosh <tvm_tirx.ir.prim.cosh>` +* :obj:`div <tvm_tirx.ir.prim.div>` +* :obj:`erf <tvm_tirx.ir.prim.erf>` +* :obj:`exp <tvm_tirx.ir.prim.exp>` +* :obj:`exp10 <tvm_tirx.ir.prim.exp10>` +* :obj:`exp2 <tvm_tirx.ir.prim.exp2>` +* :obj:`floor <tvm_tirx.ir.prim.floor>` +* :obj:`floordiv <tvm_tirx.ir.prim.floordiv>` +* :obj:`floormod <tvm_tirx.ir.prim.floormod>` +* :obj:`fmod <tvm_tirx.ir.prim.fmod>` +* :obj:`hypot <tvm_tirx.ir.prim.hypot>` +* :obj:`if_then_else <tvm_tirx.ir.prim.if_then_else>` +* :obj:`indexdiv <tvm_tirx.ir.prim.indexdiv>` +* :obj:`indexmod <tvm_tirx.ir.prim.indexmod>` +* :obj:`infinity <tvm_tirx.ir.prim.infinity>` +* :obj:`isfinite <tvm_tirx.ir.prim.isfinite>` +* :obj:`isinf <tvm_tirx.ir.prim.isinf>` +* :obj:`isnan <tvm_tirx.ir.prim.isnan>` +* :obj:`ldexp <tvm_tirx.ir.prim.ldexp>` +* :obj:`likely <tvm_tirx.ir.prim.likely>` +* :obj:`log <tvm_tirx.ir.prim.log>` +* :obj:`log10 <tvm_tirx.ir.prim.log10>` +* :obj:`log1p <tvm_tirx.ir.prim.log1p>` +* :obj:`log2 <tvm_tirx.ir.prim.log2>` +* :obj:`max <tvm_tirx.ir.prim.max>` +* :obj:`max_value <tvm_tirx.ir.prim.max_value>` +* :obj:`min <tvm_tirx.ir.prim.min>` +* :obj:`min_value <tvm_tirx.ir.prim.min_value>` +* :obj:`nearbyint <tvm_tirx.ir.prim.nearbyint>` +* :obj:`nextafter <tvm_tirx.ir.prim.nextafter>` +* :obj:`popcount <tvm_tirx.ir.prim.popcount>` +* :obj:`pow <tvm_tirx.ir.prim.pow>` +* :obj:`power <tvm_tirx.ir.prim.power>` +* :obj:`round <tvm_tirx.ir.prim.round>` +* :obj:`rsqrt <tvm_tirx.ir.prim.rsqrt>` +* :obj:`shift_left <tvm_tirx.ir.prim.shift_left>` +* :obj:`shift_right <tvm_tirx.ir.prim.shift_right>` +* :obj:`sigmoid <tvm_tirx.ir.prim.sigmoid>` +* :obj:`sin <tvm_tirx.ir.prim.sin>` +* :obj:`sinh <tvm_tirx.ir.prim.sinh>` +* :obj:`sqrt <tvm_tirx.ir.prim.sqrt>` +* :obj:`tan <tvm_tirx.ir.prim.tan>` +* :obj:`tanh <tvm_tirx.ir.prim.tanh>` +* :obj:`trunc <tvm_tirx.ir.prim.trunc>` +* :obj:`truncdiv <tvm_tirx.ir.prim.truncdiv>` +* :obj:`truncmod <tvm_tirx.ir.prim.truncmod>` +* :obj:`vscale <tvm_tirx.ir.prim.vscale>` + +tvm_tirx.tirx.op +---------------- + +.. currentmodule:: tvm_tirx.tirx.op + +.. autosummary:: + :toctree: generated/ + + gpu_active_thread_selector + gpu_storage_sync + gpu_thread_filter + gpu_thread_invariant + gpu_thread_return + gpu_warp_activemask + gpu_warp_shuffle + gpu_warp_shuffle_down + gpu_warp_shuffle_up + gpu_warp_shuffle_xor + kernel_replace_point + masked_load + masked_store + +* :obj:`CallFFIKernelAttr <tvm_tirx.tirx.CallFFIKernelAttr>` +* :obj:`TensorLoad <tvm_tirx.ir.TensorLoad>` +* :obj:`TensorMapType <tvm_tirx.tirx.TensorMapType>` +* :obj:`abi_field_get <tvm_tirx.tirx.abi_field_get>` +* :obj:`abi_field_set <tvm_tirx.tirx.abi_field_set>` +* :obj:`address_of <tvm_tirx.tirx.address_of>` +* :obj:`alloc_workspace <tvm_tirx.tirx.alloc_workspace>` +* :obj:`assume_aligned <tvm_tirx.tirx.assume_aligned>` +* :obj:`call_cpacked <tvm_tirx.tirx.call_cpacked>` +* :obj:`call_cpacked_lowered <tvm_tirx.tirx.call_cpacked_lowered>` +* :obj:`call_extern <tvm_tirx.tirx.call_extern>` +* :obj:`call_ffi_kernel <tvm_tirx.tirx.call_ffi_kernel>` +* :obj:`call_intrin <tvm_tirx.tirx.call_intrin>` +* :obj:`call_llvm_intrin <tvm_tirx.tirx.call_llvm_intrin>` +* :obj:`call_llvm_pure_intrin <tvm_tirx.tirx.call_llvm_pure_intrin>` +* :obj:`call_packed <tvm_tirx.tirx.call_packed>` +* :obj:`call_packed_lowered <tvm_tirx.tirx.call_packed_lowered>` +* :obj:`call_pure_extern <tvm_tirx.tirx.call_pure_extern>` +* :obj:`cpu_parallel_barrier <tvm_tirx.tirx.cpu_parallel_barrier>` +* :obj:`free_workspace <tvm_tirx.tirx.free_workspace>` +* :obj:`get_active_lane_mask <tvm_tirx.tirx.get_active_lane_mask>` +* :obj:`get_vscale_expr <tvm_tirx.tirx.get_vscale_expr>` +* :obj:`gpu_dp4a <tvm_tirx.tirx.gpu_dp4a>` +* :obj:`gpu_fill_fragment <tvm_tirx.tirx.gpu_fill_fragment>` +* :obj:`gpu_load_matrix_sync <tvm_tirx.tirx.gpu_load_matrix_sync>` +* :obj:`gpu_mma_sync <tvm_tirx.tirx.gpu_mma_sync>` +* :obj:`gpu_store_matrix_sync <tvm_tirx.tirx.gpu_store_matrix_sync>` +* :obj:`gpu_thread_allreduce <tvm_tirx.tirx.gpu_thread_allreduce>` +* :obj:`handle_add_byte_offset <tvm_tirx.tirx.handle_add_byte_offset>` +* :obj:`ignore_loop_partition <tvm_tirx.tirx.ignore_loop_partition>` +* :obj:`is_tensor_var <tvm_tirx.tirx.is_tensor_var>` +* :obj:`isnullptr <tvm_tirx.tirx.isnullptr>` +* :obj:`logaddexp <tvm_tirx.tirx.logaddexp>` +* :obj:`ptr_byte_offset <tvm_tirx.tirx.ptr_byte_offset>` +* :obj:`register_intrin_lowering <tvm_tirx.tirx.register_intrin_lowering>` +* :obj:`reinterpret <tvm_tirx.tirx.reinterpret>` +* :obj:`stack_alloca <tvm_tirx.tirx.stack_alloca>` +* :obj:`stack_make_dltensor <tvm_tirx.tirx.stack_make_dltensor>` +* :obj:`stack_make_shape <tvm_tirx.tirx.stack_make_shape>` +* :obj:`tensor_data_ptr <tvm_tirx.tirx.tensor_data_ptr>` +* :obj:`throw_last_error <tvm_tirx.tirx.throw_last_error>` +* :obj:`undef <tvm_tirx.tirx.undef>` +* :obj:`vector_combine <tvm_tirx.tirx.vector_combine>` +* :obj:`vector_high <tvm_tirx.tirx.vector_high>` +* :obj:`vector_low <tvm_tirx.tirx.vector_low>` + +tvm_tirx.tirx.analysis +---------------------- + +.. currentmodule:: tvm_tirx.tirx.analysis + +.. autosummary:: + :toctree: generated/ + + undefined_vars + verify_memory + verify_well_formed + + +tvm_tirx.tirx.transform +----------------------- + +.. currentmodule:: tvm_tirx.tirx.transform + +.. autosummary:: + :toctree: generated/ + + AnnotateEntryFunc + Apply + BF16ComputeLegalize + BF16StorageLegalize + BindTarget + CommonSubexprElim + ConvertSSA + DecorateDeviceScope + FP8ComputeLegalize + FP8StorageLegalize + Filter + FlattenTensor + ForceNarrowIndexToInt32 + FunctionPass + HoistIf + LowerIntrin + LowerTIRx + LowerTVMBuiltin + LowerThreadAllreduce + MakePackedAPI + NarrowDataType + PointerValueTypeRewrite + RemoveNoOp + RemoveNoOpConfig + RenormalizeSplitPattern + SkipAssert + SplitHostDevice + StmtSimplify + StorageRewrite + TileDispatch + UnrollLoop + UnrollLoopConfig + VectorizeLoop + VerifyMemory + function_pass + + +tvm_tirx.tirx.layout +-------------------- + +.. currentmodule:: tvm_tirx.tirx.layout + +.. autosummary:: + :toctree: generated/ + + Axis + Iter + +* :obj:`ComposeLayout <tvm_tirx.tirx.ComposeLayout>` +* :obj:`Layout <tvm_tirx.tirx.Layout>` +* :obj:`TileLayout <tvm_tirx.tirx.TileLayout>` + +tvm_tirx.tirx.compilation_pipeline +---------------------------------- + +.. currentmodule:: tvm_tirx.tirx.compilation_pipeline + +.. autosummary:: + :toctree: generated/ + + default_tir_pipeline + finalize_device_passes + finalize_device_passes_tirx + finalize_host_passes + finalize_tir_pipeline + register_default_tir_pipeline_selector + tirx_pipeline + +* :obj:`get_default_tir_pipeline <tvm_tirx.tirx.get_default_tir_pipeline>` +* :obj:`get_tir_pipeline <tvm_tirx.tirx.get_tir_pipeline>` +* :obj:`register_tir_pipeline <tvm_tirx.tirx.register_tir_pipeline>` + +tvm_tirx.tirx.tile +------------------ + +.. currentmodule:: tvm_tirx.tirx.tile + +.. autosummary:: + :toctree: generated/ + + DispatchContext + Instruction + Operand + TensorCall + fail + + +tvm_tirx.tirx.tile.dispatch +--------------------------- + +.. currentmodule:: tvm_tirx.tirx.tile.dispatch + +.. autosummary:: + :toctree: generated/ + + DispatchFail + run_dispatch + +* :obj:`DispatchContext <tvm_tirx.tirx.tile.DispatchContext>` +* :obj:`fail <tvm_tirx.tirx.tile.fail>` + +tvm_tirx.tirx.tile.instruction +------------------------------ + +.. currentmodule:: tvm_tirx.tirx.tile.instruction + +.. autosummary:: + :toctree: generated/ + + namespace + +* :obj:`Instruction <tvm_tirx.tirx.tile.Instruction>` +* :obj:`Operand <tvm_tirx.tirx.tile.Operand>` +* :obj:`TensorCall <tvm_tirx.tirx.tile.TensorCall>` + +tvm_tirx.tirx.tile.views +------------------------ + +.. currentmodule:: tvm_tirx.tirx.tile.views + +.. autosummary:: + :toctree: generated/ + + Add + BinaryChain + BinaryOp + BinaryReduce + Cast + Copy + CopyAsync + Exp + Exp2 + Exp2WithScaleBias + ExpWithScaleBias + FDiv + FMA + Fill + Gemm + GemmAsync + Log2 + Log2WithScaleBias + Max + Maximum + Memset + Min + Minimum + Mul + Reciprocal + ReduceNegate + ReduceOp + Select + SiLU + Sqrt + SqrtWithScaleBias + Sub + Sum + UnaryOp + UnaryOpWithScaleBias + UnaryReduce + UnaryReduceWithScaleBias + Zero + + +tvm_tirx.tirx.bench +------------------- + +.. currentmodule:: tvm_tirx.tirx.bench + +.. autosummary:: + :toctree: generated/ + + CudaProfiler + DistributedBenchContext + EventType + bench + decode_tag + export_to_perfetto_trace + is_running_under_pytest + setup diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000000..1555dac5a6 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,21 @@ +# 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. + +sphinx>=7 +sphinx-book-theme +sphinx-reredirects +triton>=3.3 diff --git a/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py b/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py index dc83f44393..24eba47b80 100644 --- a/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py +++ b/python/tvm_tirx/backend/cuda/lang/tile_scheduler.py @@ -54,32 +54,35 @@ class BaseTileScheduler: class ClusterPersistentScheduler2D(BaseTileScheduler): - """ - Tile scheduler for cluster-based persistent kernels. + """Tile scheduler for cluster-based persistent kernels. Distributes a 2D tile grid across persistent clusters using group-major ordering for L2 cache locality. Each cluster starts at its cluster_id and strides by num_clusters to process tiles. Tile Ordering (group-major for L2 locality): + - Tiles are grouped into "L2 groups" of `l2_group_size` rows - Within a group, tiles are visited in column-major order within the group - Groups are processed in row-major order - Example with 4x4 tiles, l2_group_size=2: + Example with 4x4 tiles, l2_group_size=2:: + Group 0 (rows 0-1): 0 2 4 6 1 3 5 7 Group 1 (rows 2-3): 8 10 12 14 9 11 13 15 Serpentine Mode (serpentine=True): + - Uses CUTLASS-style 2D block swizzle with serpentine traversal - Grid is divided into swizzle_size x swizzle_size blocks - Within each block, tiles are visited in row-major order - Blocks are traversed in serpentine order (even block-rows forward, odd backward) - This provides better L2 locality by reusing both A and B tiles - Example with 4x4 tiles, swizzle_size=2, serpentine=True: + Example with 4x4 tiles, swizzle_size=2, serpentine=True:: + Block layout: Block(0,0) Block(0,1) Block(1,0) Block(1,1) @@ -127,77 +130,83 @@ class ClusterPersistentScheduler2D(BaseTileScheduler): Usage ----- - ```python - scheduler = ClusterPersistentScheduler2D( - "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES, - num_clusters=NUM_CLUSTERS, l2_group_size=8 - ) - scheduler.init(cluster_id) # cluster_id = cta_idx // CLUSTER_SIZE - - while scheduler.valid(): - m = T.meta_var(scheduler.m_idx) # current M tile - n = T.meta_var(scheduler.n_idx) # current N tile - # ... process tile (m, n) ... - scheduler.next_tile() - ``` + + .. code-block:: python + + scheduler = ClusterPersistentScheduler2D( + "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES, + num_clusters=NUM_CLUSTERS, l2_group_size=8 + ) + scheduler.init(cluster_id) # cluster_id = cta_idx // CLUSTER_SIZE + + while scheduler.valid(): + m = T.meta_var(scheduler.m_idx) # current M tile + n = T.meta_var(scheduler.n_idx) # current N tile + # ... process tile (m, n) ... + scheduler.next_tile() + Examples -------- Example 1: Basic persistent kernel - ``` - num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2 - cluster_m=1, cluster_n=1 (default, no tile subdivision) - - Group-major tile numbering (l2_group_size=2): - n=0 n=1 n=2 n=3 - m=0 0 2 4 6 ┐ L2 group 0 - m=1 1 3 5 7 ┘ - m=2 8 10 12 14 ┐ L2 group 1 - m=3 9 11 13 15 ┘ - - Work distribution (cluster starts at cluster_id, strides by num_clusters=3): - cluster 0: work_idx 0,3,6,9,12,15 -> tiles 0,3,6,9,12,15 - cluster 1: work_idx 1,4,7,10,13 -> tiles 1,4,7,10,13 - cluster 2: work_idx 2,5,8,11,14 -> tiles 2,5,8,11,14 - - Tile grid (which cluster handles each tile): - n=0 n=1 n=2 n=3 - m=0 C0 C2 C1 C0 ┐ L2 group 0 - m=1 C1 C0 C2 C1 ┘ - m=2 C2 C1 C0 C2 ┐ L2 group 1 - m=3 C0 C2 C1 C0 ┘ - - Tile sequence per cluster (in execution order): - cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3) - cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2) - cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3) - ``` + + .. code-block:: text + + num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2 + cluster_m=1, cluster_n=1 (default, no tile subdivision) + + Group-major tile numbering (l2_group_size=2): + n=0 n=1 n=2 n=3 + m=0 0 2 4 6 ┐ L2 group 0 + m=1 1 3 5 7 ┘ + m=2 8 10 12 14 ┐ L2 group 1 + m=3 9 11 13 15 ┘ + + Work distribution (cluster starts at cluster_id, strides by num_clusters=3): + cluster 0: work_idx 0,3,6,9,12,15 -> tiles 0,3,6,9,12,15 + cluster 1: work_idx 1,4,7,10,13 -> tiles 1,4,7,10,13 + cluster 2: work_idx 2,5,8,11,14 -> tiles 2,5,8,11,14 + + Tile grid (which cluster handles each tile): + n=0 n=1 n=2 n=3 + m=0 C0 C2 C1 C0 ┐ L2 group 0 + m=1 C1 C0 C2 C1 ┘ + m=2 C2 C1 C0 C2 ┐ L2 group 1 + m=3 C0 C2 C1 C0 ┘ + + Tile sequence per cluster (in execution order): + cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3) + cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2) + cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3) + Example 2: 2SM GEMM (typical B200 config) - ``` - M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1 - => M_TILES=8, N_TILES=4 - => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity) - - Scheduler params: - num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8 - cluster_m=1, cluster_n=1 - - Key: Scheduler outputs CLUSTER-level tiles. - All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler. - CTAs differentiate via cluster_rank (computed OUTSIDE scheduler): - cluster_rank = cta_idx % CLUSTER_SIZE - cb_m = cluster_rank % CLUSTER_M # 0 or 1 for 2SM - cb_n = cluster_rank // CLUSTER_M # 0 for 2SM - - Final CTA tile: - cta_m = m_idx * CLUSTER_M + cb_m - cta_n = n_idx * CLUSTER_N + cb_n - - Example: cluster 5 gets scheduler tile (1,2) - CTA rank=0 (cb_m=0): actual tile (2,2) - CTA rank=1 (cb_m=1): actual tile (3,2) - ``` + + .. code-block:: text + + M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1 + => M_TILES=8, N_TILES=4 + => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity) + + Scheduler params: + num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8 + cluster_m=1, cluster_n=1 + + Key: Scheduler outputs CLUSTER-level tiles. + All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler. + CTAs differentiate via cluster_rank (computed OUTSIDE scheduler): + cluster_rank = cta_idx % CLUSTER_SIZE + cb_m = cluster_rank % CLUSTER_M # 0 or 1 for 2SM + cb_n = cluster_rank // CLUSTER_M # 0 for 2SM + + Final CTA tile: + cta_m = m_idx * CLUSTER_M + cb_m + cta_n = n_idx * CLUSTER_N + cb_n + + Example: cluster 5 gets scheduler tile (1,2) + CTA rank=0 (cb_m=0): actual tile (2,2) + CTA rank=1 (cb_m=1): actual tile (3,2) + """ def __init__( @@ -469,7 +478,8 @@ class FlashAttentionLinearScheduler(BaseTileScheduler): """Linear 3D scheduler for flash attention (batch, head, m_block). Used for non-causal attention with simple linear decomposition. - Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using: + Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:: + batch = linear_idx // (num_heads * num_m_blocks) head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks m_block = linear_idx % num_m_blocks diff --git a/python/tvm_tirx/tirx/script/ir_builder/triton.py b/python/tvm_tirx/tirx/script/ir_builder/triton.py index 33ae1d87f4..603c9575c8 100644 --- a/python/tvm_tirx/tirx/script/ir_builder/triton.py +++ b/python/tvm_tirx/tirx/script/ir_builder/triton.py @@ -42,9 +42,10 @@ class TritonKernel(BaseKernel): This class bridges the Triton kernel with TVM runtime. The compilation includes the following steps: - - Deduce the kernel signature and generate the Triton kernel - - Embed the compiled kernel into the current IRModule as an external module - - Generate a call to the Triton kernel following its calling convention via call_packed. + + - Deduce the kernel signature and generate the Triton kernel + - Embed the compiled kernel into the current IRModule as an external module + - Generate a call to the Triton kernel following its calling convention via call_packed. """ def __init__(self, func):
