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(N3) 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 ( ≈1,935 GB/s) at N≥2048 due to L2 cache reuse, reaching 3,023 GB/s at N=4096 . Julia+CUDA’s fixed 16×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), (4loads+1store)×8B/t . Values exceeding the A100 PCIe physical peak ( ≈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×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× 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× speedup; embed_lookup achieves 1.4× ; 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×
rms_norm_reduce
0.003
43.3
0.069
<1×
rms_norm_scale
0.0009
45.5
0.020
<1×
matmul_vec
96.7
38.0
1.137
85×
rope_rotate_q
0.044
52.9
0.048
≈1×
Table 7: LLaMA-3 8B micro-kernel latency at 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 μ s
0.04 μ s
12× faster
activation_sum
14.8 μ s
13.5 μ s
0.91×
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
∼ 60
DotElemFusionPass
17
∼ 40
Lowering + registration
auto
∼ 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 ( × 2 kernels)
JLIR ( μ s)
C ++ MLIR ( μ s)
Parse / Build IR
0 25
∼ 50
Fusion passes
∼ 10–50
∼ 520
Lower to arith
160
∼ 100
Total
∼ 185
∼ 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×
1 M
371
160
2.3×
10 M
557
489
1.1×
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× 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): (4loads+1store)×8B/t ; values exceeding the ≈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.
Multi-Level Intermediate Representation (MLIR) underlies modern ML compiler infrastructure (TensorFlow, JAX/StableHLO, PyTorch Inductor, IREE), yet appears only in trace amounts in code-LM pretraining corpora. MLIR is also extensible by design: new dialects ship per application domain, so a fine-tuned model per dialect does not scale. We ask whether inference-time priors derived mechanically from each dialect's Operation Definition Specification (ODS) can substitute for gradient-based adaptation. First, we release four natural-language-to-MLIR benchmarks across three dialects (MLIR-Spec-150, Linalg-Spec-30, StableHLO-Spec-30, StableHLO-Held-Out-200), totaling 410 in-scope NL-to-MLIR pairs, plus a 25-program out-of-grammar stress set, with functional references on 30 Spec prompts, under Apache-2.0 with Gebru datasheets and Croissant 1.1 metadata. Second, we build a three-layer schema-derived constraint stack: a CFG over op signatures (C1), type-domain splits from an ODS-extracted type lattice (C2), and an SSA-scope validator driving five-retry rejection sampling (C3). Porting it to StableHLO required no new constraint-layer code. On structurally dominated dialects, SmolLM2-1.7B with token-level constraints matches or exceeds 15B-34B code LMs running post-hoc rejection at a matched retry budget and 8-25x the per-generation speed: on linalg it reaches 80.0% verify-valid (three runs, n=125), 21 to 44 pp above CodeLlama-34B, Granite-Code-34B, and StarCoder2-15B with non-overlapping CIs. The comparison is between systems, not models: under the baselines' post-hoc protocol, SmolLM2 falls from 52.0% to 24.5% on arith+func. On arith+func and the templated StableHLO-Held-Out-200 the baselines match or beat the SLM; we scope these as non-win cells. We release benchmarks, decoder, every per-prompt generation, and a pinned reproducibility Dockerfile.
Large Language Models (LLMs) have become the dominant workload on modern AI accelerators, yet deploying them on specialized hardware still faces two core challenges: how to import a trained model into a compiler-friendly intermediate representation, and how to efficiently schedule the autoregressive inference loop under limited on-chip memory. This paper presents an MLIR (Multi-Level Intermediate Representation) based compilation method for large language models, illustrated using two dialects of operators, TopOp and TpuOp. TopOp serves as a high-level graph dialect that is independent of both the source framework and the target chip, and is responsible for expressing model semantics; TpuOp serves as the target hardware dialect, carrying chip-related decisions such as quantization, layer groups, and memory layout. A model is first represented as TopOp, then lowered layer by layer to TpuOp, and finally a deployable binary is generated. In addition, each Transformer layer is split into three stages for static compilation: prefill, prefill_kv (prefill with historical key-value cache), and decode, so as to accommodate the different computational characteristics of prompt-parallel processing and per-token generation. The method has been implemented in the TPU-MLIR compiler {https://github.com/sophgo/tpu-mlir} and the LLM-TPU deployment project {https://github.com/sophgo/LLM-TPU}, supporting a variety of generative models including the Qwen, Llama, InternVL, and MiniCPM-V series, as well as multiple quantization and deployment forms such as GPTQ, AWQ, and AutoRound.
We present the design and implementation of PolyBlocks, a modular and reusable MLIR-based compiler infrastructure for AI programming frameworks and AI chips. PolyBlocks is based on pass pipelines that compose transformations on loop nests and SSA, primarily relying on lightweight affine access analysis; the transformations are stitched together in specialized ways to realize high-performance code automatically by the use of analytical cost models and heuristics. The optimizations in these passes include multi-level tiling, fusion, on-chip scratchpad usage, mapping matmuls and convolutions to matrix units, fusing the attention layer, and several other transformations for parallelism and locality. They have been developed in a way that makes it easy to build PolyBlocks-based compilers to target new chips, reusing much of the infrastructure. PolyBlocks' design and architecture enable fully automatic code generation from high-level frameworks to low-level target-specific intrinsics. Experimental results from evaluating PolyBlocks-powered just-in-time compilation for PyTorch and JAX targeting NVIDIA GPUs show that it is able to match or outperform Torch Inductor and XLA in several cases, although the latter rely on a combination of vendor libraries and code generation. For individual operators like matmuls and convolutions, PolyBlocks-generated code is competitive with the best vendor-tuned libraries or hand-written kernels.