Tullio.jl Documentation

Tullio.jl is a flexible einsum macro for Julia that understands various array operations written in index notation. It supports matrix multiplication, permutations, convolutions, stencils, scatter/gather, and broadcasting.

Key features include:

• Automatic Loop Generation: Writes optimized nested loops based on index notation.
• Performance Optimizations: Uses multi-threading (via Threads.@spawn) and recursive tiling for large arrays.
• Package Integration:
  • Uses LoopVectorization.jl (@avx) for high-speed execution (disable with avx=false).
  • Uses KernelAbstractions.jl (@kernel) for GPU support (disable with cuda=false).
• Automatic Differentiation: Provides gradients for use with Tracker.jl or via ChainRules.jl (for Zygote.jl, Yota.jl, etc.).
  • Default: Symbolic derivative of the RHS (works for reductions over +, min, or max).
  • grad=Dual: Uses ForwardDiff.jl for more complex expressions (only for + reductions).

@tullio M[x,y,c] := N[x+i, y+j,c] * K[i,j]     # sum over i,j, and create M

@tullio S[x] = P[x,y] * log(Q[x,y] / R[y])     # sum over y, and write into S

@tullio A[i,j] += B[i,k,l] * C[l,j] * D[k,j]   # sum over k,l, and add to values in A

@tullio (*) Z[j] := X[ind[k],j] * exp(-Y[k])   # product over k

Tullio uses index notation to define how arrays are summed, broadcast, or concatenated.

Common patterns:

• Summation: Indices not appearing on the left-hand side are summed over.
• Broadcasting: If all indices on the RHS appear on the LHS, it performs broadcasting.
• Matrix Multiplication: Summing over a shared index between two arrays.
• Concatenation: Using nested loops to build higher-dimensional arrays.
• Reductions: Using specific functions like (max) or (*) to reduce indices.

using Tullio, Test
M = rand(1:20, 3, 7)

# Summation
@tullio S[1,c] := M[r,c]  # sum over r ∈ 1:3, for each c ∈ 1:7

# Broadcasting
@tullio Q[ρ,c] := M[ρ,c] + sqrt(S[1,c])

# Matrix Multiplication
mult(M,Q) = @tullio P[x,y] := M[x,c] * Q[y,c]

# Concatenation
R = [rand(Int8, 3, 4) for δ in 1:5]
@tullio T[j,i,δ] := R[δ][i,j] + 10im

# Reductions
@tullio (max) X[i] := abs2(T[j,i,δ])

# In-place update
dbl!(M, S) = @tullio M[r,c] = 2 * S[1,c]

Tullio supports complex indexing patterns for signal processing and stencils:

• Shifts: Use i+_ to indicate an automatic shift in the index range.
• Downsampling: The range of the index is the intersection of ranges allowed by both terms.
• Wrapping & Padding: Use functions like mod, clamp, or pad within the index to handle boundary conditions.
• Index by values: Indices can be determined by the values within an array (e.g., A[2K[j]+i]).

# Downsample
@tullio B[i] := (A[2i] + A[2i+1])/2

# Shifts
@tullio M[i,j] := A[i+j-1]  (j in 1:15)  # i in 1:7
@tullio M[i+_,j] := A[i+j]  (j in 1:15)  # i in 0:6, automatic shift "i+_"

# Wrapped & padded
@tullio M[i,j] := A[mod(i+j)]  (j in 1:15, i in 1:15)   # wraps around
@tullio M[i,j] := A[clamp(i+j)]  (j in 1:15, i in 1:15) # repeats edges
@tullio M[i+_,j] := A[pad(i+j, 3)]  (j in 1:15)         # fills with zeros

Tullio provides syntax for handling ranges, constants, and specialized indexing within the macro expression.

Specifying Ranges and Constants

• Manual Ranges: If ranges cannot be inferred, use A[i] := i^2 (i in 1:10).
• Fixing Indices: To prevent an index from being summed over, fix it to a constant using the $ syntax: A[i] := B[i, $col] - C[i, 2] (where col is a constant).
• Including Constants: Use the $ prefix for preferred constant inclusion: A[i] := $d * B[i]. Note that gradients are not calculated for d.

Specialized Indexing

• Mapping/Clamping: Use A[mod(i), clamp(j)] to map indices to the valid axes of A. This also disables range inference from A.
• Padding: Use A[pad(i, 3)] to extend the range of i by inserting zeros. Use pad=NaN to use NaN as the padding value instead of zero.
• Shifted Output: To ensure shifted output indices start at 1, use an underscore on the left side: A[i+_] := ....

# Specifying a range
A[i] := i^2 (i in 1:10)

# Fixing an index to a constant
A[i] := B[i, $col]

# Using padding
A[pad(i, 3)]

For complex logic, @tullio expressions can be written using begin ... end blocks.

Key Behaviors:

• Explicit Ranges: If the macro cannot infer the output array's index ranges (e.g., when using complex logic or assignment), you must provide the index ranges explicitly in parentheses at the end of the expression.
• Index Assignment: Using the syntax xi = ... inside the block tells Tullio that xi is a local variable and should not be summed over.
• Manual @inbounds: If you use @inbounds inside a begin ... end block, Tullio will not automatically add @inbounds to the generated loops. It also assumes indices might go out of bounds.
• Symbolic Derivatives: Complex multi-line blocks may not support symbolic derivatives, but grad=Dual (Dual numbers) will work.

# Example: A convolution with cyclic indices
@tullio out[x,y,c] := begin
    xi = mod(x+i, axes(mat,1)) # 'xi =' prevents summing over xi
    @inbounds trunc(Int, mat[xi, mod(y+j), c] * kern[i,j])
end (x in 1:10, y in 1:10, i in -3:3, j in -3:3, c in 1:1)

Tullio supports automatic differentiation (AD) and GPU acceleration. You can use AD libraries like Tracker or Zygote to compute gradients of @tullio expressions. For GPU acceleration, use CUDA.jl or KernelAbstractions.jl by passing GPU arrays to the macro.

Important Limitations & Warnings:

• GPU Reductions: Complete reductions to a single scalar (e.g., sum) are currently not supported on the GPU and will be extremely slow.
• Scalar Gradients: Gradients are only calculated for arrays. Passing a scalar to a gradient function (e.g., gradient(a -> (@tullio _ := $a * A[i]), 3.14)) will result in zero.
• grad=Dual Performance: When using grad=Dual, the right-hand side of the expression is evaluated a second time during the backward pass. This saves memory but increases computation time for expensive functions.

using Tullio
using Tracker

# Define a matrix multiplication using Tullio
mul(A, B) = @tullio C[i,k] := A[i,j] * B[j,k]

A = rand(3,40); B = rand(40,500)

# Compute gradient with respect to A
ΔA = Tracker.gradient((A,B) -> sum(mul(A, B)), A, B)[1]

# GPU usage with CUDA
using CUDA
cu_A = cu(A); cu_B = cu(B)
cu_res = mul(cu_A, cu_B)
cu_ΔA = Tracker.gradient((A,B) -> sum(mul(A, B)), cu_A, cu_B)[1]

Tullio allows fine-grained control over performance and behavior via keyword arguments within the macro.

Performance Tuning

• threads: Set to true (default) for multi-threading, false to disable, or a specific threshold (e.g., threads=64^3) to control when work is divided.
• avx: Set to true (default) to use LoopVectorization.jl, false to disable, or an integer (e.g., avx=4) to specify unroll factors.
• tensor: Set to true (default) to use TensorOperations.jl, or false to disable.
• cuda: Controls GPU usage. Default is cuda=256 (passed to kernel(CUDA(), 256)). Requires KernelAbstractions.jl and CUDA.jl to be visible.

Gradient & Differentiation

• grad: Set to true (default) for gradient calculation, false to disable, or grad=Dual to use ForwardDiff.jl (requires ForwardDiff to be loaded).
• nograd: Disables gradient calculation for specific arrays. Use nograd=A or nograd=(A, B, C).

Output & Debugging

• verbose: Set to false (default), true to print inferred index ranges and gradients, or 2 to print all internal generated functions.
• init: Sets the initial value for reductions (e.g., init=0.0). While +, *, min, max, &, and | have sensible defaults, other reductions use zero by default.

Array Assignment

• A[i,j] := ...: Creates a new array.
• A[i,j] = ...: Writes into an existing array.
• A[i,j] += ...: Performs an in-place update on an existing array.

# Default settings:
@tullio threads=true fastmath=true avx=true tensor=true cuda=256 grad=Base verbose=false A[i,j] := ...