Static Arguments as a Means of Increasing Performance.

[Jordan-Dennis]

Hi The jax Community,

Can the performance of a function be improved by making arguments static?

If the answer is yes then I am also curios about:

How much if at all, the target architecture matters?
On the top of you're head, are there any specific circumstances where this will, or will not work?
How likely is this to change as jax and mhlo develop?

I am interested in the subset of functions with static shapes, i.e. jax.jit(func) succeeds. This question is based on the discussion in this issue.

Below I have included a minimal(ish) working example based on our discussions:

import jax 

def hypotenuse(ccoords: float) -> float:
    x: float = jax.lax.index_in_dim(ccoords, 0, keepdims=False)
    y: float = jax.lax.index_in_dim(ccoords, 1, keepdims=False)
    x_sq: float = jax.lax.integer_pow(x, 2)
    y_sq: float = jax.lax.integer_pow(y, 2)
    return jax.lax.sqrt(x_sq + y_sq)

def coords(n: int, rad: float) -> float:
    arange: float = jax.lax.iota(float, n)
    max_: float = np.array(n - 1, dtype=float)
    grid: float = arange * 2. * rad / max_ - rad
    s: int = grid.size
    shape: tuple = (1, s, s) 
    x: float = jax.lax.broadcast_in_dim(grid, shape, (2,))
    y: float = jax.lax.broadcast_in_dim(grid, shape, (1,))
    return jax.lax.concatenate([x, y], 0)

def cartesian_to_polar(ccoords: float) -> float:
    x: float = jax.lax.index_in_dim(ccoords, 0)
    y: float = jax.lax.index_in_dim(ccoords, 1)
    x_sq: float = jax.lax.integer_pow(x, 2)
    y_sq: float = jax.lax.integer_pow(y, 2)
    hypot: float = jax.lax.sqrt(x_sq + y_sq)
    return jax.lax.concatenate([hypot, jax.lax.atan2(x, y)], 0)

def circ_ap_func(
        radius: float, 
        x: float, 
        y: float,
        rotation: float, 
        nsoft: float,
        pixel_scale: float) -> float:
    # Passing arguments to safe types. 
    centre: float = np.asarray([x, y]).astype(float)
    radius: float = np.asarray(radius).astype(float)
    rotation: float = np.asarray(rotation).astype(float)
    nsoft: float = np.asarray(nsoft).astype(float)
    
    # Organising coords
    ccoords: float = coords(npix, radius)
    
    # Translation 
    ccoords: float = ccoords - centre[:, None, None]
        
    # Rotation
    sin_alpha: float = jax.lax.sin(rotation)
    cos_alpha: float = jax.lax.cos(rotation)
    x: float = jax.lax.index_in_dim(ccoords, 0)
    y: float = jax.lax.index_in_dim(ccoords, 1)
    new_x: float = x * cos_alpha - y * sin_alpha
    new_y: float = x * sin_alpha + y * cos_alpha
    ccoords: float = jax.lax.concatenate([new_x, new_y], 0)        
        
    # Transformation 
    rho: float = hypotenuse(ccoords)
        
    # Linear softening
    distances: float = radius - rho
    lower: float = jax.lax.full_like(distances, 0., dtype=float)
    upper: float = jax.lax.full_like(distances, 1., dtype=float)
    inside: float = jax.lax.max(distances, lower)
    scaled: float = inside / nsoft / pixel_scale
    aperture: float = jax.lax.min(scaled, upper)
    return aperture


x: float = 0.
y: float = 0.
rotation: float = 0.
pixel_scale: float = 2. * radius / npix
wavelength: float = 550e-09
radius: float = 1.
npix: int = 128
nsoft: int = 3

dynamic: callable = jax.jit(circ_ap_func)
static: callable = jax.jit(circ_ap_func, inline=True, static_argnums=(1, 2, 3, 4))

# %%timeit
static(radius, x, y, rotation, nsoft, pixel_scale).block_until_ready()

# %%timeit
dynamic(radius, x, y, rotation, nsoft, pixel_scale).block_until_ready()

Regards

Jordan

[jakevdp]

Hi - thanks for the question. I think there's not enough information here to give an answer (I'm honestly not even sure what you're asking!) and the diuscussion you linked to is very long and it's difficult for me to understand which comments and code snippets are relevant to your question. Perhaps you could edit your question to more directly add a minimal example of the situation you're curious about?

[Jordan-Dennis]

Sorry for the lack of clarity @jakevdp. I'll put together a summary of the linked issue, but it may take some time on my end (this is not urgent). I've edited the question so that it (hopefully) makes more sense now.

[Jordan-Dennis]

I found the following functions useful trying to answer this question myself, perhaps they will also be useful for you.

import re
import difflib

def get_mhlo(
        func: callable, # Compiled 
        *args: object, 
        **kwargs: object
    ) -> str:
    mhlo: str = func.lower(*args, **kwargs).compile().as_text()
    return re.sub("metadata={.*}", "", mhlo)

def print_diff(mhlo: str, comp: str) -> list:
    mhlo_lines: list = mhlo.splitlines()
    comp_lines: list = comp.splitlines()
    diff: iter = difflib.unified_diff(mhlo_lines, comp_lines)
    for line in diff:
        print(line)