JLIR: A Julia-Native MLIR-Inspired Intermediate Representation with Automatic JACC Kernel Extraction

Narasinga Rao Miniskar; Seyong Lee; Keita Teranishi; Jeffrey S Vetter

cs.PLSep 4, 2026

JLIR: A Julia-Native MLIR-Inspired Intermediate Representation with Automatic JACC Kernel Extraction

Authors: Narasinga Rao Miniskar, Seyong Lee, Keita Teranishi, Jeffrey S Vetter

Organizations: Oak Ridge National Laboratory, Oak Ridge, 37830, USA

Abstract

The Multi-Level Intermediate Representation (MLIR) has made reusable compiler infrastructure practical for domain-specific computation. However, MLIR's strong compile-time type requirements and low-level (C++) extension model can be a poor match for high-level, dynamically specialized languages such as Julia. MLIR has several drawbacks for dynamic programming languages in terms of the type system and level of abstraction. It is thus extremely challenging for non-compiler or scientific computing users to introduce new programming abstractions and express algorithm implementations in a form that remains both natural and optimizable. As a result, library interfaces for linear algebra, mesh processing, partial differential equations, and related domains often sit outside the compiler optimization path. We present JLIR (Julia-native Level Intermediate Representation), a Julia-native intermediate representation framework that brings the main benefits of MLIR-style multi-level, dialect-oriented compilation into the Julia ecosystem while remaining usable as ordinary Julia code. JLIR represents Julia programs before low-level lowering, supports extensible operations and transformation passes through Julia's language mechanisms, and allows partially typed programs to remain transformable until concrete types are known. The framework includes built-in dialects for arithmetic, control flow, functions, structured loops, and memory operations, and it also includes a lightweight mechanism for adding new domain operations without modifying the core system. To demonstrate JLIR's capabilities, we applied it to automatic Julia for Accelerators (JACC) kernel generation.

Figures & tables

Figure 1: Overview of JLIR’s capabilities: a Julia-native front-end, five built-in dialects, a JIT-compatible type system, a composable pass infrastructure, automatic GPU lowering via JACCTransformPass , code generation, and zero-friction packaging.

Figure 2: JLIR compilation pipeline. A serial Julia function ( Expr tree) is parsed into SSA-form JLIR across five dialects; user-defined operations may be injected via a .jld dialect file (dashed arrow). A configurable pass pipeline then applies transformations: JACCTransformPass emits GPU-portable JACC Julia (targeting CUDA/ROCm/CPU threads), and the Codegen / Julia Emitter pass lowers IR to plain executable Julia evaluated by Julia’s JIT. The ResourceEstimationPass produces an operation/memory/cycle report as a side output.

MLIR concept	MLIR mechanism	JLIR (Julia-native)
Operation type	C ++ struct + TableGen	Julia struct + @op DSL
Dialect definition	TableGen .td file	@dialect macro
Op dispatch	Virtual method table	Julia multiple dispatch
Pass interface	C ++ CRTP † abstract base class	Julia AbstractPass + dispatch
Front-end parser	Custom MLIR parser / C ++	Meta.parse + Expr walk
Type system	Strictly typed at IR build time	AnyType for JIT-resolved types

Table 1: Mapping of MLIR concepts to their Julia-native JLIR equivalents.

Pattern	Kernel	Julia Source	JLIR	Generated JACC
1D Parallel-for	rms_norm_scale	Listing 5	Listing 6	Listing 13
2D Parallel-for	jacobi2d	Listing 7	Listing 8	Listing 14
2D Parallel-for	gemm	Listing 9	Listing 10	Listing 15
(Mixed 2D+Sequential)
1D Parallel-reduce	rms_norm_reduce	Listing 11	Listing 12	Listing 16

Table 2: JACC Transformation Pass with four different pattern matches and examples of input Julia, JLIR, and output JACC transformation

Benchmark	Source LOC	Pipeline Time
GEMM	10	$<0.4$ ms
Jacobi 2D stencil	8	$<0.3$ ms
Black–Scholes	12	$<0.5$ ms
LLaMA-3 (3 kernels)	52	$<1.2$ ms

Table 3: JLIR pipeline compilation overhead per source function. LOC: Lines of Code

N	JACC CPU	cuBLAS	JACC GPU	Julia+CUDA
	[GF/s]	[TF/s]	[TF/s]	[TF/s]
256	0.8	0.818	0.54	0.91
512	5.9	4.265	1.35	1.65
1,024	22.0	13.17	1.89	2.16
2,048	24.0	16.83	1.28	1.96
4,096	—	17.27	1.20	1.88

Table 4: DGEMM throughput (NVIDIA A100 PCIe, Float64 ). JACC CPU is omitted (—) for

N>2048

. JACC GPU and Julia+CUDA both use the same untiled

O(N^{3})

loop; cuBLAS uses double-precision (FP64) tensor cores with shared-memory tiling. All GPU values in TF/s; JACC CPU in GF/s.

Figure 3: Jacobi 2D stencil: effective bandwidth (GB/s). JACC GPU exceeds the A100 PCIe physical DRAM peak (

{\approx}1{,}935

GB/s) at

N\!\geq\!2048

due to L2 cache reuse, reaching 3,023 GB/s at

N=4096

. Julia+CUDA’s fixed

16{\times}16

tiles achieve higher reuse (3,814 GB/s).

N	Serial	CPU thr.	Julia+CUDA	JACC GPU
512	17.4	0.2	758.3	499.8
1024	16.2	0.9	2226.9	1353.2
2048	6.0	4.5	3227.9	2401.7
4096	—	17.7	3814.0	3023.4

Table 5: Jacobi 2D stencil: effective (logical) memory bandwidth (GB/s),

(4\,\text{loads}+1\,\text{store})\times 8\,\text{B}/t

. Values exceeding the A100 PCIe physical peak (

{\approx}1{,}935

GB/s) reflect L2 cache hit reuse in the five-point stencil, not super-physical DRAM throughput. “—” denotes serial runs skipped as too slow. CPU thr. = 128-thread Threads.@threads baseline.

Figure 4: Black–Scholes throughput (GF/s). At

n=10

M JACC GPU reaches 96% of Julia+CUDA (556 vs 581 GF/s); both kernels invoke the same CUDA transcendental device intrinsics. The gap at

n=100

K reflects kernel launch overhead dominating a short computation.

N	Serial	CPU thr.	Julia+CUDA	JACC GPU
100 K	0.4	0.04	118.1	77.0
1 M	0.4	0.40	434.6	351.0
10 M	—	4.93	581.3	555.5

Table 6: Black–Scholes throughput (GF/s, NVIDIA A100 PCIe, Float64 ). Julia+CUDA is a hand-written @cuda kernel with the identical algorithm. Both GPU paths dispatch the transcendentals ( log , exp ,

2\times\texttt{erf}

) to the same CUDA math device intrinsics. CPU Thr. = 128-thread Threads.@threads baseline; “—” denotes serial runs skipped as too slow.

Figure 5: LLaMA-3 8B kernel latency (ms, log scale) for serial CPU, JACC CPU (128 threads), and JACC GPU across five kernels at dim=4096. matmul_vec achieves

85\times

GPU speedup; launch overhead dominates the lighter kernels.

Figure 6: LLaMA-3 8B speedup of JACC GPU over the serial CPU baseline at dim=4096. matmul_vec achieves

85\times

speedup; embed_lookup achieves

1.4\times

; the remaining kernels are launch overhead dominated.

Kernel	Serial	JACC CPU	JACC GPU	GPU
	[ms]	[ms]	[ms]	Speedup
embed_lookup	0.028	38.8	0.020	$1.4\times$
rms_norm_reduce	0.003	43.3	0.069	$<1\times$
rms_norm_scale	0.0009	45.5	0.020	$<1\times$
matmul_vec	96.7	38.0	1.137	$85\times$
rope_rotate_q	0.044	52.9	0.048	${\approx}1\times$

Table 7: LLaMA-3 8B micro-kernel latency at

\mathit{dim}=4096

, 32 heads. JACC CPU = 128-thread Threads.@threads ; serial ref = single-threaded Julia. GPU speedup = serial / JACC GPU.

Metric	Before fusion	After fusion
Ops in loop body	4 BinOp	2 ScaleAddOp
Arithmetic / iter	$4n$	$2n$
Memory reads / iter	$3n$	$3n$

Table 8: Resource estimation for axpy_dot before and after ScaleAddFusionPass .

n

denotes the loop trip count.

Kernel	Ref.	JLIR	Ratio
axpy_dot	0.40 $\mu$ s	0.04 $\mu$ s	$12\times$ faster
activation_sum	14.8 $\mu$ s	13.5 $\mu$ s	$0.91\times$

Table 9: Wall-clock performance of JLIR-lowered vs. reference Julia for the DSL benchmark kernels (

N=1{,}024

, minimum of 20 repetitions).

Component	JLIR	C ++ MLIR
Dialect op definitions	26 (.jld)	41 (.td + header + glue)
ScaleAddFusionPass	24	$\sim$ 60
DotElemFusionPass	17	$\sim$ 40
Lowering + registration	auto	$\sim$ 67
Build system	none (Julia pkg)	CMake + mlir-tblgen + ninja
Total (code lines)	67	243

Table 10: Source lines of code (comments and blank lines excluded) for the myops dialect. “auto” = generated by the JLIR DSL with no user-written code.

Phase ( $\times$ 2 kernels)	JLIR ( $\mu$ s)	C ++ MLIR ( $\mu$ s)
Parse / Build IR	0 25	$\sim$ 50
Fusion passes	$\sim$ 10–50	$\sim$ 520
Lower to arith	160	$\sim$ 100
Total	$\sim$ 185	$\sim$ 620

Table 11: Pipeline latency for the myops dialect on two kernels (warm context; minimum of 20–50 repetitions).

N	JACC GPU	Reactant GPU	cuBLAS
256	0.58	0.37	0.82
512	1.38	2.81	4.27
1,024	1.89	9.43	13.17
2,048	1.29	14.76	16.83
4,096	1.20	17.31	17.27

Table 12: DGEMM throughput comparison: JACC GPU (JLIR-generated), Reactant.jl (XLA/cuBLAS), and cuBLAS (NVIDIA A100 PCIe, Float64 ). JACC GPU and Reactant GPU timed in separate processes; cuBLAS from Table 4 . All values in TF/s.

n	JACC GPU	Reactant GPU	JACC/React.
100 K	87	25	$3.5\times$
1 M	371	160	$2.3\times$
10 M	557	489	$1.1\times$

Table 13: Black–Scholes throughput (GF/s): JACC GPU (JLIR-generated) vs. Reactant.jl (XLA), NVIDIA A100 PCIe, Float64 . Timed in separate processes; JACC values from Table 6 .

Benchmark	JACC GPU peak	Julia+CUDA peak	cuBLAS peak	JACC/Julia+CUDA
Jacobi 2D (eff. BW ‡ )	3,023 GB/s	3,814 GB/s	—	79%
Black–Scholes	556 GF/s	581 GF/s	—	96%
GEMM ( $N=1024$ ) †	1,889 GF/s	2,158 GF/s	13,168 GF/s	87%
matmul_vec (LLaMA-3)	$85\times$ over serial	—	—	—

Table 14: Peak performance summary of JLIR-generated JACC GPU kernels (NVIDIA A100 PCIe). Julia+CUDA is a hand‑written @cuda kernel with the identical algorithm. ‡ Jacobi bandwidth is effective (logical):

(4\,\text{loads}+1\,\text{store})\times 8\,\text{B}/t

; values exceeding the

{\approx}1{,}935

GB/s physical peak reflect L2 cache hit reuse. † cuBLAS uses FP64 tensor cores with shared-memory tiling (16–17 TF/s); JACC GPU and Julia+CUDA both use the same untiled algorithm. The 87% JACC/Julia+CUDA ratio is preserved relative to cuBLAS.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 21

Figure 22

Figure 23

Figure 24

Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR
May 14, 2026Plawan Kumar RathConstrained Decoding Code Generation
Meta
An MLIR-Based Compilation Method for Large Language Models
Jul 17, 2026Pengchao Hu, Zhibin Xin, Yifan Chen +3LLM Serving LLM Inference Scheduling
Sophgo Inc.
PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks
Mar 6, 2026Uday Bondhugula, Akshay Baviskar, Navdeep Katel +3Compiler Optimization GPU Kernel Generation
Polymage Labs and Indian Institute of Science, Bengaluru, Karnataka, India · Polymage Labs, Bengaluru, India

May 14, 2026J/K move · Enter open · S save

Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR

Plawan Kumar Rath

Abstract

Explore similar work

Cross-Dialect Generalization Without Retraining: Benchmarks and Evaluation of Schema-Derived Constrained Decoding for MLIR

An MLIR-Based Compilation Method for Large Language Models

PolyBlocks: A Compiler Infrastructure for AI Chips and Programming Frameworks