remove ensure_compile_time_eval because hlo shows it's already optimized

[spraharsh]

Addresses #738

The main issue comes up in the evaluation of

with jax.ensure_compile_time_eval():
    return is_successful(result) | is_event(result)

This also shows up in jax-ml/jax#31461 where jax.ensure_compile_time_eval doesn't play well with sharding.

In this particular case however ensure_compile_time_eval() makes no difference to optimized code. to reproduce that, run,

# check_hlo.py
import functools
import re
import sys
from contextlib import nullcontext
from pathlib import Path

import diffrax as dfx
import jax
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

label = sys.argv[1]  # before or after
assert label in ("before", "after")

mesh = Mesh(np.array(jax.devices("cpu")[:2]), ("i",))
spec = P("i")


@jax.jit
@functools.partial(
    jax.shard_map,
    mesh=mesh,
    in_specs=spec,
    out_specs=spec,
    check_vma=False,
)
@jax.vmap
def run(y0):
    return dfx.diffeqsolve(
        dfx.ODETerm(lambda t, y, args: -y),
        dfx.Tsit5(),
        t0=0,
        t1=1,
        dt0=0.1,
        y0=y0,
    ).ys


# The original fails with the ambient context, so use its working baseline.
context = jax.set_mesh(mesh) if label == "after" else nullcontext()
with context:
    y0 = jax.device_put(
        np.array([10, 20], dtype=np.float32),
        NamedSharding(mesh, spec),
    )
    compiled = run.lower(y0).compile()
    hlo = compiled.as_text()

Path(f"{label}.raw-hlo.txt").write_text(hlo)

# Retain all computation bodies; exclude header/source-location tables.
start = re.search(r"^(?:%|ENTRY %)", hlo, re.MULTILINE).start()
instructions = hlo[start:]

# Remove instruction debug metadata; retain all other attributes.
instructions = re.sub(r", metadata=\{[^\n]*?\}", "", instructions)
Path(f"{label}.instructions.txt").write_text(instructions)

with

JAX_PLATFORMS=cpu JAX_NUM_CPU_DEVICES=2 PYTHONPATH=. python check_hlo.py before

before the patch is applied (Note that before does not have a mesh set, setting a mesh leads to the compilation error described). Then run

JAX_PLATFORMS=cpu JAX_NUM_CPU_DEVICES=2 PYTHONPATH=. python check_hlo.py after diff -u before.instructions.txt after.instructions.txt

after the patch is applied. The instructions show no difference

[spraharsh]

The original reason ensure compile_time_eval was added was to optimize code for "benchmarks/against_scan.py" in this commit . I double checked that removing it does not change emitted HLO

Run

import re
import sys
from pathlib import Path

import jax

label = sys.argv[1]  # before or after

# Load the benchmark definitions without running its speed tests.
path = Path("benchmarks/against_scan.py")
source = "\n".join(
    line for line in path.read_text().splitlines()
    if not line.startswith("speedtest(")
)
benchmark = {}
exec(compile(source, str(path), "exec"), benchmark)

compiled = jax.jit(benchmark["dfx_fn"]).lower(
    benchmark["fields"], benchmark["ts"]
).compile()

hlo = compiled.as_text()
Path(f"{label}.raw-hlo.txt").write_text(hlo)

# Exclude module header/source tables and instruction debug metadata.
start = re.search(r"^(?:%|ENTRY %)", hlo, re.MULTILINE).start()
instructions = re.sub(r", metadata=\{[^\n]*?\}", "", hlo[start:])
Path(f"{label}.instructions.txt").write_text(instructions)

before the patch with

JAX_PLATFORMS=cuda XLA_PYTHON_CLIENT_PREALLOCATE=false PYTHONPATH=. \
python check_benchmark_hlo.py before

. Then run

JAX_PLATFORMS=cuda XLA_PYTHON_CLIENT_PREALLOCATE=false PYTHONPATH=. \
python check_benchmark_hlo.py after

diff -u before.instructions.txt after.instructions.txt

There should be no difference in instructions