kernel-fusion-blog

kernel-fusion.py

import torch
import triton
import triton.language as tl

# ==========================================
# Pointwise Fusion Example (PyTorch Source)
# ==========================================

def pointwise_example(x, w, b):
    """
    Original PyTorch implementation of a 
    standard neural network layer pattern.
    """
    # Multiple element-wise operations
    tmp = x * w         # multiply
    tmp = tmp + b       # add
    tmp = tmp.sigmoid() # sigmoid activation
    return tmp


# ==========================================
# Unfused Triton Kernels
# ==========================================

@triton.jit
def mul_kernel(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    """Kernel 1: Multiply"""
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.load(in_ptr1 + x0, xmask)
    tmp2 = tmp0 * tmp1
    tl.store(out_ptr0 + x0, tmp2, xmask)

@triton.jit
def add_kernel(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    """Kernel 2: Add"""
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.load(in_ptr1 + x0, xmask)
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + x0, tmp2, xmask)

@triton.jit
def sigmoid_kernel(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    """Kernel 3: Sigmoid"""
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.sigmoid(tmp0)
    tl.store(out_ptr0 + x0, tmp1, xmask)


# ==========================================
# Fused Triton Kernel
# ==========================================

@triton.jit
def triton_poi_fused_add_mul_sigmoid_0(in_ptr0, in_ptr1, in_ptr2, 
                                        out_ptr0, xnumel, XBLOCK: tl.constexpr):
    """Fused Kernel: Multiply -> Add -> Sigmoid"""
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex

    # Load all inputs once from Global Memory
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = tl.load(in_ptr1 + (x0), xmask)
    tmp3 = tl.load(in_ptr2 + (x0), xmask)

    # Fused pointwise operations happen in GPU Registers
    tmp2 = tmp0 * tmp1
    tmp4 = tmp2 + tmp3
    tmp5 = tl.sigmoid(tmp4)

    # Store final result only to Global Memory
    tl.store(out_ptr0 + (x0), tmp5, xmask)

# ==========================================
# Expected Reduciton Kernel from Script
# ==========================================
    
@triton.jit
def triton_per_fused_add_mul_sum_0(
    in_out_ptr0,
    in_ptr0,
    xnumel,
    r0_numel,
    XBLOCK: tl.constexpr,
):
    xnumel = 1024
    r0_numel = 1024
    R0_BLOCK: tl.constexpr = 1024

    rnumel = r0_numel
    RBLOCK: tl.constexpr = R0_BLOCK

    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel

    r0_index = tl.arange(0, R0_BLOCK)[None, :]
    r0_offset = 0
    r0_mask = tl.full([R0_BLOCK], True, tl.int1)[None, :]

    roffset = r0_offset
    rindex = r0_index

    r0_1 = r0_index
    x0 = xindex

    tmp0 = tl.load(in_ptr0 + (r0_1 + 1024 * x0), xmask, other=0.0)
    tmp1 = tl.full([1, 1], 2.0, tl.float32)
    tmp2 = tmp0 * tmp1

    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, R0_BLOCK])
    tmp5 = tl.where(xmask, tmp3, 0)

    tmp6 = tl.sum(tmp5, 1)[:, None].to(tl.float32)

    tmp7 = tl.full([1, 1], 1.0, tl.float32)
    tmp8 = tmp6 + tmp7

    tl.store(in_out_ptr0 + (x0), tmp8, xmask)

# ==========================================
# Reduction Fusion Example Script
# ==========================================

def reduction_example(x):
    """Example of pointwise + reduction fusion"""
    tmp = x * 2.0
    result = tmp.sum(dim=-1)
    result = result + 1.0
    return result

if __name__ == "__main__":
    # Create test input on GPU
    if torch.cuda.is_available():
        x = torch.randn(1024, 1024, device='cuda')

        # Compile the function
        compiled_fn = torch.compile(reduction_example)
        
        # Run to trigger compilation and see logs if TORCH_LOGS="output_code"
        result_fused = compiled_fn(x)
        print("Compilation successful and kernel executed.")
    else:
        print("CUDA not available. Run on a GPU-enabled machine to see Triton kernels.")