Kernel: gelu (Level 1, category: activation)
Input shapes: [[4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 1e-06)
Pallas runtime: 814.1361 ms; JAX XLA compiled runtime: 0.2041 ms; speedup vs compiled baseline: 0.0003

JAX reference implementation:
@jax.jit
def jax_gelu(x: jax.Array) -> jax.Array:
    return jax.nn.gelu(x)

Original Pallas kernel (TPU-oriented): """Level 1: Elementwise GELU via Pallas.

Demonstrates: transcendental functions (tanh, erf) inside kernels.
"""

from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L1/gelu", __doc__)

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def _gelu_kernel(x_ref, o_ref):
    x = x_ref[...]
    o_ref[...] = x * 0.5 * (1.0 + jnp.tanh(
        jnp.sqrt(2.0 / jnp.pi) * (x + 0.044715 * x ** 3)
    ))


def pallas_gelu(x: jax.Array) -> jax.Array:
    n = x.shape[0]
    block_size = min(1024, n)
    grid_size = n // block_size

    return pl.pallas_call(
        _gelu_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(grid_size,),
        in_specs=[pl.BlockSpec((block_size, *x.shape[1:]), lambda i: (i, *([0] * (x.ndim - 1))))],
        out_specs=pl.BlockSpec((block_size, *x.shape[1:]), lambda i: (i, *([0] * (x.ndim - 1)))),
    )(x)


pallas_kernel = pallas_gelu
task_name = "gelu"
input_shapes = [(4096, 4096)]
category = "activation"
level = 1

Fixed Pallas kernel (GPU-compatible): """Level 1: Elementwise GELU via Pallas.

Demonstrates: transcendental functions (tanh, erf) inside kernels.
"""

from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L1/gelu", __doc__)

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def _gelu_kernel(x_ref, o_ref):
    x = x_ref[...]
    o_ref[...] = x * 0.5 * (1.0 + jnp.tanh(
        jnp.sqrt(2.0 / jnp.pi) * (x + 0.044715 * x ** 3)
    ))


def pallas_gelu(x: jax.Array) -> jax.Array:
    bm = min(128, x.shape[0])
    bn = min(128, x.shape[1]) if x.ndim > 1 else x.shape[0]
    grid = (x.shape[0] // bm, x.shape[1] // bn) if x.ndim > 1 else (x.shape[0] // bm,)
    if x.ndim > 1:
        in_sp = [pl.BlockSpec((bm, bn), lambda i, j: (i, j))]
        out_sp = pl.BlockSpec((bm, bn), lambda i, j: (i, j))
    else:
        in_sp = [pl.BlockSpec((bm,), lambda i: (i,))]
        out_sp = pl.BlockSpec((bm,), lambda i: (i,))
    return pl.pallas_call(
        _gelu_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=grid,
        in_specs=in_sp,
        out_specs=out_sp,
    )(x)


pallas_kernel = pallas_gelu
task_name = "gelu"
input_shapes = [(4096, 4096)]
category = "activation"
level = 1

Diff of GPU fix: --- a/gelu.py
+++ b/gelu.py
@@ -20,16 +20,21 @@
 
 
 def pallas_gelu(x: jax.Array) -> jax.Array:
-    n = x.shape[0]
-    block_size = min(1024, n)
-    grid_size = n // block_size
-
+    bm = min(128, x.shape[0])
+    bn = min(128, x.shape[1]) if x.ndim > 1 else x.shape[0]
+    grid = (x.shape[0] // bm, x.shape[1] // bn) if x.ndim > 1 else (x.shape[0] // bm,)
+    if x.ndim > 1:
+        in_sp = [pl.BlockSpec((bm, bn), lambda i, j: (i, j))]
+        out_sp = pl.BlockSpec((bm, bn), lambda i, j: (i, j))
+    else:
+        in_sp = [pl.BlockSpec((bm,), lambda i: (i,))]
+        out_sp = pl.BlockSpec((bm,), lambda i: (i,))
     return pl.pallas_call(
         _gelu_kernel,
         out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
-        grid=(grid_size,),
-        in_specs=[pl.BlockSpec((block_size, *x.shape[1:]), lambda i: (i, *([0] * (x.ndim - 1))))],
-        out_specs=pl.BlockSpec((block_size, *x.shape[1:]), lambda i: (i, *([0] * (x.ndim - 1)))),
+        grid=grid,
+        in_specs=in_sp,
+        out_specs=out_sp,
     )(x)
 
 

Jaxpr IR: { lambda ; a:f32[4096,4096]. let
    b:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{f32[128,128]} d:Ref<default>{f32[128,128]}. let
          e:f32[128,128] <- c[...]
          f:f32[128,128] = mul e 0.5:f32[]
          g:f32[] = sqrt 0.6366197723675814:f32[]
          h:f32[128,128] = integer_pow[y=3] e
          i:f32[128,128] = mul 0.044714998453855515:f32[] h
          j:f32[128,128] = add e i
          k:f32[] = convert_element_type[new_dtype=float32 weak_type=False] g
          l:f32[128,128] = mul k j
          m:f32[128,128] = tanh l
          n:f32[128,128] = add 1.0:f32[] m
          o:f32[128,128] = mul f n
          p:f32[128,128], d[...] <- d[...], o
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a
  in (b,) }