Apache TVM

TIRx is an open-source, hardware-native Domain Specific Language (DSL) and compiler designed for machine-learning kernels. It is built on top of the Apache TVM compiler infrastructure and targets GPUs and specialized AI accelerators.

Unlike high-level DSLs that hide hardware details, TIRx uses a lower and more explicit boundary. It allows developers to keep orchestration (pipeline state, synchronization, roles, and backend intrinsics) in hardware-native source code, while exposing recurring tile-level structures to the compiler to enable optimization and reuse.

Apache TVM is a machine learning compilation framework designed for Python-first development and universal deployment. It allows developers to take pre-trained machine learning models, compile them, and generate deployable modules that can be embedded and run across diverse environments (mobile, edge, bare metal, etc.).

Key characteristics include:

• Python-first: The optimization process is fully customizable in Python, allowing for pipeline changes without recompiling the TVM stack.
• Composable: Optimization passes, libraries, and code generation (codegen) can be easily composed into the existing pipeline.
• Universal Deployment: The TVM runtime is designed to work in non-Python environments and supports zero-copy data exchange with ecosystems like PyTorch, TensorFlow, and TensorRT via DLPack support.

The TVM Vulkan runtime is composed of three primary architectural components that manage device lifecycle, thread-local state, and kernel execution:

1. VulkanDeviceAPI: Implements the standard TVM DeviceAPI. It manages core Vulkan data structures, initializes the Vulkan instance and devices, and queries for available extensions.
2. VulkanThreadEntry: Manages thread-local state. It maintains a staging buffer for data copies and a VulkanStream for each device.
3. VulkanWrappedFunc: Handles kernel launches. It retrieves a VulkanPipeline instance from the VulkanModuleNode and executes the kernel on the active VulkanStream using either immediate or deferred modes.

TVM uses several specialized modules for different stages of the compilation and optimization pipeline:

• tvm/relax: High-level IR for computational graphs. Uses relax.transform for optimizations.
• tvm/tirx: Core IR for TensorIR (low-level tensor functions). Includes tirx::PrimFunc, analysis, and transformation passes.
• tvm/s_tir: Schedulable TIR. Provides schedule primitives (tiling, vectorization) and auto-tuning tools like MetaSchedule and DLight.
• tvm/arith: Tools for analyzing integer arithmetic properties (bounds, positiveness) used by TensorIR passes.
• tvm/te: Tensor Expression DSL. Use te.create_prim_func to convert expressions into tirx::PrimFuncs.
• tvm/topi: Tensor Operator Inventory. A collection of pre-defined operators (e.g., numpy-like) for common deep learning workloads.

The tvm.relax.backend module provides specialized backend support for compiling and running Relax programs on various hardware targets. Depending on your target hardware, you should use the corresponding sub-module to access backend-specific optimizations and compilation logic.

Available backend sub-modules include:

• tvm.relax.backend.cuda: For NVIDIA CUDA devices.
• tvm.relax.backend.rocm: For AMD ROCm devices.
• tvm.relax.backend.metal: For Apple Metal devices.
• tvm.relax.backend.adreno: For Qualcomm Adreno GPUs.
• tvm.relax.backend.gpu_generic: For generic GPU acceleration.
• tvm.relax.backend.cpu_generic: For generic CPU execution.
• tvm.relax.backend.contrib: For additional contribution-based backend utilities.

The tvm.contrib package provides a collection of utility modules for various hardware accelerators, math libraries, and data handling tasks. These modules allow users to interface with low-level libraries (like BLAS, cuDNN, or MKL) and perform specialized operations within the TVM ecosystem.

Key submodules include:

• Linear Algebra & BLAS: cblas, cublas, cublaslt, hipblas, mkl, and thrust for high-performance matrix operations.
• Deep Learning Accelerators: cudnn (NVIDIA), nnpack, dnnl (Intel), and hexagon (Qualcomm).
• Data & Interop: dlpack for zero-copy tensor sharing and download for fetching remote resources.
• Specialized Utilities: coreml_runtime for Apple CoreML, random for random number generation, and pickle_memoize for caching.

Apache TVM follows a structured flow to transform high-level model descriptions into deployable modules. The architecture is primarily composed of two major components: TensorIR and Relax.

To understand how to use TVM, you should follow this mental model:

1. Overall Flow: The sequence of steps from high-level model input to a deployable artifact.
2. Runtime-based View: How components interact during the compilation and execution process.
3. Logical Modules: The static organization of the codebase and how different modules relate to one another.

For deep dives into the specific IR (Intermediate Representation) layers, refer to the TensorIR Deep Dive and Relax Deep Dive documentation.

The tvm.s_tir.dlight module provides tools for applying automated Tensor IR (TIR) optimizations. It is organized into submodules targeting specific hardware backends and analysis tasks:

• tvm.s_tir.dlight.gpu: Optimization routines for GPU backends.
• tvm.s_tir.dlight.cpu: Optimization routines for CPU backends.
• tvm.s_tir.dlight.adreno: Optimization routines specifically for Adreno GPUs.
• tvm.s_tir.dlight.analysis: Tools for analyzing TIR to guide optimizations.
• tvm.s_tir.dlight.base: Base classes and fundamental utilities for the DLight framework.
• tvm.s_tir.dlight.benchmark: Utilities for benchmarking the performance of optimized TIR schedules.