tilelang/examples/analyze

TVM IR Performance Analyzer

A performance analysis toolkit for TVM IR modules, Provides hardware-aware performance metrics including FLOPs, memory bandwidth utilization, and execution time estimation.

Features

• ​Operation Analysis: Supports arbitrary operations expressed in TVM IR (including GEMM and convolution)
• ​Memory Traffic Calculation: Tracks global memory transfers
• ​Architecture-aware Metrics: Pre-configured with NVIDIA GPU architectures (Ampere, Ada Lovelace)
• ​Performance Estimation: Predicts execution time using roofline model
• ​TVM Integration: Works with TVM IRModule and PrimFunc

Quick Start

GEMM Analysis Example

import tilelang.language as T
from tilelang.tools import Analyzer
from tilelang.carver.arch import CUDA

M = N = K = 1024

def kernel(block_M=128, block_N=128, block_K=32, num_stages=3, thread_num=128):
    @T.prim_func
    def main(A: T.Tensor((M, K), "float16"),
             B: T.Tensor((N, K), "float16"),
             C: T.Tensor((M, N), "float")):
        # ... (kernel definition)
    return main

cuda_device = CUDA("cuda")
result = Analyzer.analysis(kernel(), cuda_device)
print(result)

Convolution Analysis Example

import tilelang.language as T
from tilelang.tools import Analyzer
from tilelang.carver.arch import CUDA

def kernel(N=64, C=256, H=512, W=512, F=512, K=3, block_M=64, block_N=128):
    @T.prim_func
    def main(data: T.Tensor((N, H, W, C), "float16"),
             kernel: T.Tensor((K, K, C, F), "float16"),
             out: T.Tensor((N, (H-K+1), (W-K+1), F), "float")):
        # ... (convolution kernel definition)
    return main

cuda_device = CUDA("cuda")
result = Analyzer.analysis(kernel(), cuda_device)
print(result)

API Documentation

AnalysisResult Class

@dataclass(frozen=True)
class AnalysisResult:
    total_flops: int          # Total floating-point operations
    total_global_bytes: int   # Global memory traffic in bytes
    estimated_time: float     # Predicted execution time (seconds)
    tflops: float             # Achieved TFLOPS
    bandwidth_GBps: float     # Memory bandwidth utilization

Analyzer Class Methods

analysis(fn, device)

• ​Parameters:
  • fn: TVM IRModule or PrimFunc
  • device: Device configuration object
• Returns: AnalysisResult

Supported Architectures

# Extendable to custom hardware via: "compute_capability": (cores_per_SM, clock_GHz, flops_per_cycle, max_SM_count)
ARCH_CONFIGS = {
    "80": (128, 1.41, 2, 108),  # A100
    "86": (128, 1.70, 2, 84),   # RTX 3080
    "89": (128, 2.52, 2, 128)  # RTX 4090
}

Implementation Details

Performance Model

Uses roofline model with two constraints:

1. ​Compute Bound: Time = Total FLOPs / (SM Count × Cores/SM × Clock × FLOPs/Cycle)
2. ​Memory Bound: Time = Memory Bytes / (Bandwidth × Utilization)

IR Analysis Pass

1. ​Traversal: Walks through TVM IR using ir_transform
2. ​Operation Detection:
   • Counts FLOPs for all compute operations
   • Calculates memory traffic for all memory operations
3. ​Loop Handling:
   • Tracks nested loops for operation scaling
   • Accounts for block/grid dimensions

Key Metrics Calculation

Metric	Formula
FLOPs per GEMM	2 × M × N × K
Memory Traffic per Copy	elements × dtype_size × loop_product
Achieved TFLOPS	total_flops / estimated_time / 1e12
Memory Bandwidth	total_global_bytes / estimated_time

Limitations

1. Requires memory operations to be properly annotated in the IR
2. Assumes perfect memory coalescing and no bank conflicts

Supported Operations

Any operation expressed in TVM IR