​

Just-in-Time Compilation

​

How Codeflash Optimizes with JIT

1. Rewrites the code in a JIT-compatible format, which may involve breaking down complex functions into separate JIT-compiled components.
2. Generates appropriate tests that are compatible with JIT-compiled code, carefully handling data types since JIT compilers have stricter input type requirements.
3. Disables JIT compilation when running coverage and tracer. This ensures accurate coverage and trace data, since both rely on Python bytecode execution. JIT-compiled code bypasses Python bytecode, so it would prevent proper tracking.
4. Disables the Line Profiler for JIT compiled code. It could be possible to disable JIT compilation and run the line profiler, but that would lead to inaccurate information which could misguide the optimization process.

​

Configuration

​

When JIT Compilation Helps

• Numerical computations with loops that can’t be easily vectorized.
• Custom algorithms not covered by existing optimized libraries.
• Functions that are called repeatedly with consistent input types.
• Code that benefits from hardware-specific optimizations (SIMD acceleration).

​

Example

​

Function Definition

import torch
def complex_activation(x):
    """A custom activation with many small operations - compile makes a huge difference"""
    # Many sequential element-wise ops create kernel launch overhead
    x = torch.sin(x)
    x = x * torch.cos(x)
    x = x + torch.exp(-x.abs())
    x = x / (1 + x.pow(2))
    x = torch.tanh(x) * torch.sigmoid(x)
    x = x - 0.5 * x.pow(3)
    return x

​

Benchmarking Snippet (replace cuda with mps to run on your Mac)

import time
# Create compiled version
complex_activation_compiled = torch.compile(complex_activation)

# Benchmark
x = torch.randn(1000, 1000, device='cuda')

# Warmup steps are slower as the JIT compiler is understanding the function execution to compile into machine code
for _ in range(10):
    _ = complex_activation(x)
    _ = complex_activation_compiled(x)

# Time uncompiled
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    y = complex_activation(x)
torch.cuda.synchronize()
uncompiled_time = time.time() - start

# Time compiled
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    y = complex_activation_compiled(x)
torch.cuda.synchronize()
compiled_time = time.time() - start

print(f"Uncompiled: {uncompiled_time:.4f}s")
print(f"Compiled: {compiled_time:.4f}s")
print(f"Speedup: {uncompiled_time/compiled_time:.2f}x")

Uncompiled: 0.0176s
Compiled: 0.0063s
Speedup: 2.80x

1. Already vectorized - All operations are already PyTorch tensor ops.
2. Multiple Kernel Launches - Uncompiled code launches ~10 separate kernels. torch.compile fuses them into 1-2 kernels, eliminating kernel launch overhead.
3. No algorithmic improvement - The computation itself is already optimal.
4. Python overhead elimination - Removes Python interpreter overhead between operations.

​

When JIT Compilation May Not Help

• The code already uses highly optimized libraries (e.g., NumPy with MKL, cuBLAS, cuDNN).
• Functions have variable input types or shapes that prevent effective compilation.
• The compilation overhead exceeds the runtime savings for short-running functions.
• The code relies heavily on Python objects or dynamic features that JIT compilers can’t optimize.

​

Example

​

Function Definition

def adaptive_processing(x, threshold=0.5):
    """Function with data-dependent control flow - compile struggles here"""
    # Check how many values exceed threshold (data-dependent!)
    mask = x > threshold
    num_large = mask.sum().item()  # .item() causes graph break

    if num_large > x.numel() * 0.3:
        # Path 1: Many large values - use expensive operation
        result = torch.matmul(x, x.T)  # Already optimized by cuBLAS
        result = result.mean(dim=0)
    else:
        # Path 2: Few large values - use cheap operation
        result = x.mean(dim=1)

    return result

​

Benchmarking Snippet (replace cuda with mps to run on your Mac)

# Create compiled version
adaptive_processing_compiled = torch.compile(adaptive_processing)

# Test with data that causes branch variation
x = torch.randn(500, 500, device='cuda')

# Warmup steps are slower as the JIT compiler is understanding the function execution to compile into machine code
for _ in range(10):
    _ = adaptive_processing(x)
    _ = adaptive_processing_compiled(x)

# Benchmark with varying data (causes recompilation)
torch.cuda.synchronize()
start = time.time()
for i in range(100):
    # Vary the data to trigger different branches
    x_test = torch.randn(500, 500, device='cuda') + (i % 2)
    y = adaptive_processing(x_test)
torch.cuda.synchronize()
uncompiled_time = time.time() - start

torch.cuda.synchronize()
start = time.time()
for i in range(100):
    x_test = torch.randn(500, 500, device='cuda') + (i % 2)
    y = adaptive_processing_compiled(x_test)  # Recompiles frequently!
torch.cuda.synchronize()
compiled_time = time.time() - start

print(f"Uncompiled: {uncompiled_time:.4f}s")
print(f"Compiled: {compiled_time:.4f}s")
print(f"Slowdown: {compiled_time/uncompiled_time:.2f}x")

Uncompiled: 0.0296s
Compiled: 0.2847s
Slowdown: 9.63x

1. Graph breaks - .item() forces a graph break, negating compile benefits.
2. Recompilation overhead - Different branches cause expensive recompilation each time.
3. Dynamic control flow - Data-dependent conditionals can’t be optimized away.
4. Already optimized ops - matmul already uses cuBLAS; compile adds overhead without benefit.

​

Better Optimization Strategy

def optimized_version(x, threshold=0.5):
    """Remove data-dependent control flow - vectorize instead"""
    mask = (x > threshold).float()
    weight = (mask.mean() > 0.3).float()  # Keep on GPU

    # Compute both paths, blend based on weight (branchless)
    expensive = torch.matmul(x, x.T).mean(dim=0)
    cheap = x.mean(dim=1).squeeze()

    # Pad cheap result to match expensive dimensions
    cheap_padded = cheap.expand(expensive.shape[0])

    result = weight * expensive + (1 - weight) * cheap_padded
    return result

Optimized: 0.0277s
Speedup compared to Uncompiled: 1.57x

1. Eliminate .item() - Keep computation on GPU.
2. Branchless execution - Compute both paths, blend results.
3. Vectorization - Replace conditionals with masked operations.
4. Reduce Python overhead - Minimize host-device synchronization.

​

Supported JIT Frameworks

​

Numba (CPU Code)

• @jit - General-purpose JIT compilation with optional flags.
  • nopython=True - Compiles to machine code without falling back to the Python interpreter.
  • fastmath=True - Uses aggressive floating-point optimizations via LLVM’s fastmath flag.
  • cache=True - cache compiled function to disk which reduces future runtimes.
  • parallel=True - Parallelizes code inside loops.

​

PyTorch

• torch.compile() - Compiles a function or module for optimized execution.
  • mode - Controls the compilation strategy:
    • "default" - Balanced compilation with moderate optimization.
    • "reduce-overhead" - Minimizes Python overhead using CUDA graphs, ideal for small batches.
    • "max-autotune" - Spends more time auto-tuning to find the fastest kernels.
  • fullgraph=True - Requires the entire function to be captured as a single graph. Raises an error if graph breaks occur, useful for ensuring complete optimization.
  • dynamic=True - Enables dynamic shape support, allowing the compiled function to handle varying input sizes without recompilation.

​

TensorFlow

• @tf.function - Converts Python functions into TensorFlow graphs for optimized execution.
  • jit_compile=True - Enables XLA compilation, which performs whole-function optimization including operation fusion, memory layout optimization, and target-specific code generation.

​

JAX

• @jax.jit - JIT compiles functions using XLA with automatic operation fusion.