PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.04585v1cs.PLcs.CL

TL;DR

MLIR’s compile-time typing and C++ extension model fit Julia’s dynamically specialized, high-level programming model poorly. JLIR addresses this with a Julia-native, extensible SSA IR and automatic JACC kernel extraction, achieving strong GPU results without programmer annotations while retaining a documented array-transfer limitation.

  • Problem

    MLIR’s strong compile-time type requirements and C++ extension model are a poor match for Julia’s dynamically specialized language model.

  • Method

    JLIR provides a Julia-native multi-level SSA IR with dialects, extensible Julia-defined passes, JIT-compatible partial typing, and automatic JACCTransformPass kernel extraction.

  • Results

    Generated A100 kernels reached 96% of the hand-written Julia+CUDA Black–Scholes baseline, 3,023 GB/s on Jacobi, 85× GPU speedup for LLaMA-3 matmul vec, and 87% for untiled DGEMM without programmer annotation.

  • Takeaways & Limitations

    A plain serial Julia function can produce portable GPU-executable code through JLIR, with the reported pipeline latency below 1.5 ms.

  • Takeaways & Limitations

    JACCTransformPass creates JACC arrays for every input and output Julia array, adding data-transfer overhead and motivating future data-flow and array-reuse passes.

Abstract

from arXiv · show

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.

1 Introduction

JLIR brings MLIR-style, multi-level compiler infrastructure into Julia using Julia-native abstractions that remain compatible with dynamic specialization. It combines extensible dialects and passes with automatic extraction of GPU-portable JACC kernels from serial Julia.

  • The design preserves partially typed programs for transformation until JIT specialization resolves concrete types, addressing Julia’s type-agnostic semantics.
  • JLIR is a Julia-native, SSA-based intermediate representation framework that provides MLIR-style compiler infrastructure within Julia itself.
  • Its five built-in dialects represent arithmetic, control flow, memory operations, and function definitions without C++ code or a foreign build system.
  • The @dialect macro lets users define new operations in a few lines of Julia, with multiple dispatch supplying extensibility for operations and lowering rules.
  • JLIR’s composable pass infrastructure includes loop fusion, loop unrolling, resource estimation, and serialization, with passes inspectable and debuggable in the Julia REPL.
  • JACCTransformPass identifies 1D, 2D, mixed, and reduction loops in SSA form and rewrites them into GPU-portable JACC parallel constructs from serial Julia.
  • The demonstration defines custom operations and fusion passes, estimates arithmetic cost, lowers fused IR to executable Julia, and verifies correctness against references.
  • On an NVIDIA A100, generated kernels reached 96% of a hand-written Julia+CUDA Black–Scholes baseline, 3,023 GB/s for Jacobi, and 85× serial-Julia speedup for LLaMA-3 matmul.

2 Background

LLVM and MLIR provide reusable compiler infrastructure but expose abstraction and accessibility constraints for high-level dynamic languages. Julia offers homoiconic syntax, JIT specialization, and multiple-dispatch-based portability mechanisms that motivate JLIR’s design.

  • LLVM uses typed SSA IR and shared optimization backends, but its machine-oriented single level largely erases high-level structures such as loop nests and tensor shapes.
  • MLIR addresses this with extensible dialects, shared IR abstractions, and progressive lowering from domain-specific operations toward LLVM-compatible representations.
  • MLIR’s C++ toolchain and TableGen-based extension model reduce accessibility for communities working primarily in higher-level languages such as Python or Julia.
  • Julia combines dynamic typing with LLVM-based JIT compilation and multiple dispatch, producing type-specific LLVM IR at first execution.
  • Julia’s Expr trees are ordinary, manipulable Julia data structures that form an accessible stage before typed Julia IR and subsequent LLVM lowering.
  • JACC provides single-source parallel-for and parallel-reduce APIs, selecting CUDA, HIP, or multithreaded CPU execution according to the active backend.
  • JLIR’s compilation pipeline parses Expr trees into SSA-form JLIR, applies configurable passes, emits JACC Julia, and lowers remaining IR to executable Julia.

3 JLIR Design

JLIR maps MLIR’s structural abstractions onto Julia’s programming model through native parsing, SSA data structures, dialects, and multiple-dispatch extensibility. Its design supports both partially typed transformations and Julia-compatible code generation.

  • JLIR retains MLIR’s core abstractions while implementing them through Julia’s programming model rather than C++ infrastructure.
  • Because Julia exposes parsed Expr objects, JLIR’s front end recursively traverses the syntax tree without a tokenizer, grammar specification, or generated code.
  • A new JLIR operation can be added by defining a Julia struct and methods for its results, operands, and lowering, without altering existing code or recompiling the framework.
  • AnyType allows SSA values to remain unresolved until Julia’s JIT determines concrete types, while JLIR passes continue operating on those values.
  • SSA values may carry names, types including AnyType, and compile-time constants, enabling constant folding during construction and early loop-bound detection.
  • JLIR operations produce and consume SSA values, while regions and blocks represent structured control flow within the IR hierarchy.
  • The five core dialects cover arithmetic, control flow, functions, structured loops, and memory references, including Julia-compatible mutation and loop-carried values.
  • The @dialect macro generates boilerplate for new dialect operations from a concise inline specification.

4 Use Case: JACC Transform Pass

JACCTransformPass automatically classifies JLIR loop nests and rewrites suitable patterns into portable JACC kernels. Its structural detection uses loop-carried state, while free-variable analysis preserves kernel argument consistency; array creation overhead remains a limitation.

  • Transformation pipeline: JACCTransformPass traverses functions and converts suitable loop nests into JACC.parallel for or JACC.parallel reduce kernels.The pass generates GPU-ready Julia from serial sources without developer annotations.
  • Loop detection: An empty iter_args field marks parallel work, whereas loop-carried iter_args indicate sequential dependence.JLIR encodes this structural fact during parsing, allowing constant-time inspection during transformation.
  • Pattern detection: The pass recognizes 1D, 2D, mixed parallel-sequential, and reduction patterns through nested-loop and accumulator checks.Reduction detection requires one iter argument, an assignment update, a producing binary operation, and an operator in {+,*,min,max}.
  • Pattern detection: Reduction detection validates accumulator structure and extracts the per-element operand for parallel reduction generation.The supported reduction operators are +, *, min, and max.
  • Kernel extraction: Free-variable analysis collects external SSA values in first-use order and reuses that order for kernel parameters and JACC call arguments.This keeps extracted kernel signatures and call sites aligned without symbol-table lookup.
  • Limitation: The pass currently creates JACC arrays for every input and output array, adding data-transfer overhead and reducing performance efficiency.The paper identifies global-scope data-flow analysis and array reuse as future work.

5 Benchmarks

The benchmark suite evaluates automatic, annotation-free JLIR transformation across complementary scientific-computing and ML-inference workloads. It compares serial Julia, hand-written CUDA, generated JACC GPU, generated JACC CPU-thread, and, for GEMM, cuBLAS variants on an AMD EPYC/A100 system.

  • Benchmark design: Four complementary benchmarks preserve plain serial Julia sources that JLIR transforms automatically without manual annotations.The suite targets computational patterns spanning scientific computing and ML inference.
  • Execution variants: The evaluation compares serial Julia, hand-crafted CUDA, generated JACC GPU, generated JACC CPU-thread, and cuBLAS for GEMM on 128 CPU threads and an NVIDIA A100.GPU runtimes use synchronization barriers, while CPU runtimes use @elapsed.
  • GEMM: GEMM uses a three-level matrix-multiplication loop nest whose outer loops are parallel and inner accumulator loop is sequential.JLIR classifies this structure as mixed 2D plus sequential.
  • Jacobi 2D stencil: Jacobi 2D is classified as a trivially parallel nested loop and emitted as a JACC.parallel for kernel.Its effective bandwidth accounts for four loads and one store per output element.
  • Black–Scholes: Black–Scholes processes independent option contracts with a compute-intensive loop and precomputed scalar parameters passed as free-variable kernel arguments.The benchmark includes logarithm, exponential, and error-function evaluations.
  • LLaMA-3: LLaMA-3 evaluation uses token-embedding lookup, RMS normalization, and RoPE-NeOX micro-kernels extracted from attention and feed-forward layers.These workloads include gather, reduction, scaling, and in-place positional-encoding operations.

6 Results

JLIR’s transformation pipeline adds negligible overhead while automatically generating GPU-portable kernels from serial Julia code. Performance depends on the generated kernel strategy: untiled GEMM trails cuBLAS, whereas Jacobi achieves high effective bandwidth through cache reuse.

  • Pipeline overhead: Every benchmark completes the full JLIR transformation in less than 1.5 ms, and this overhead is amortized by Julia’s first-invocation JIT compilation.The pipeline includes parsing, JACCTransformPass, and code emission.
  • GEMM performance: At peak, the JLIR-generated JACC GPU GEMM kernel reaches 87% of hand-written Julia+CUDA performance, or 1.89 versus 2.16 TF/s at N = 1024.The ratio declines to 69–70% at N ≥8192 because JACC wrapper register overhead reduces SM occupancy.
  • GEMM performance: cuBLAS reaches 13–17 TF/s across N ∈[1024, 16384], 7–15× faster than the untiled JACC GPU kernel because it uses FP64 tensor cores and shared-memory tiling.JLIR and Julia+CUDA use the same naive untiled three-level loop, so matching cuBLAS would require a shared-memory tiling pass.
  • Jacobi 2D stencil: Julia+CUDA reaches 3,814 GB/s at N = 4096, exceeding JACC’s 3,023 GB/s because its fixed 16×16 tiles provide higher reuse than JACC’s auto-selected 32×32 blocks.The comparison concerns effective bandwidth, not physical DRAM throughput.
  • Jacobi 2D stencil: The JACC GPU Jacobi kernel reaches 2,402 GB/s at N = 2048 and 3,023 GB/s at N = 4096 effective bandwidth through L2 cache reuse.These values are logical workload-level bandwidth and can exceed the A100 PCIe physical peak of approximately 1,935 GB/s because cache reuse reduces DRAM traffic.

6.4 Black–Scholes Option Pricing Performance

JLIR-generated JACC GPU kernels closely match hand-written Julia+CUDA on compute-bound Black–Scholes, while LLaMA-3 benefits only when kernels have enough work to amortize dispatch overhead.

  • Black–Scholes: 556 GF/s is 96% of the hand-written Julia+CUDA baseline’s 581 GF/s for Black–Scholes at n = 10 M options.Both implementations dispatch transcendental functions to identical CUDA math intrinsics.
  • Black–Scholes: At n = 100 K, Black–Scholes performance is limited because kernel launch overhead dominates the short computation.
  • LLaMA-3: 85× GPU speedup occurs for the LLaMA-3 matmul vec kernel, the dominant compute bottleneck at dim = 4096.Its serial latency is 96.7 ms versus 1.137 ms on the GPU.
  • LLaMA-3: The embed lookup kernel achieves only 1.4× GPU speedup because its serial runtime nearly matches kernel launch latency.
  • LLaMA-3: Three lightweight LLaMA-3 kernels receive no practical benefit from JACC because their runtimes are at or below GPU dispatch latency.Their serial runtimes are 0.0009–0.044 ms, compared with approximately 0.019–0.068 ms dispatch latency.
  • Implication: A profitability heuristic comparing estimated trip counts with empirical launch thresholds could suppress GPU dispatch for small kernels.

6.6 Dialect DSL Extensibility

JLIR’s Dialect DSL lets users define operations and fusion passes in Julia, integrate them with existing infrastructure, and preserve correctness through lowering. The approach reduces implementation effort and can improve generated-kernel performance when fusion removes allocations or redundant arithmetic.

  • Dialect definition: Four custom arithmetic operations were defined in a 24-line .jld file and exercised through a nine-phase compilation pipeline.
  • Dialect definition: Loading a dialect at runtime generates operation structures, builders, dispatch methods, printing support, and lowering integration without modifying JLIR internals.
  • User-defined fusion: ScaleAddFusionPass recognizes add(x, mul(s, y)) patterns and replaces them with ScaleAdd operations while removing dead intermediates.
  • Resource estimation: After fusion, axpy dot requires 2 arithmetic operations per iteration instead of 4, while memory reads remain at 3.
  • Correctness and performance: Both fused kernels match reference implementations within 10^-14 relative error, confirming semantics-preserving fusion and lowering.
  • Correctness and performance: The JLIR-lowered axpy dot is 12× faster than the reference because it avoids the reference’s temporary array allocation.
  • Comparison with MLIR: The C++ MLIR implementation requires 3.6× more code, while JLIR is 3× faster overall at runtime on the compared dialect workload.

6.7 Comparison with Reactant.jl

JLIR and Reactant offer different Julia GPU compilation trade-offs: Reactant can obtain cuBLAS-backed tiled GEMM performance, while JLIR has lower overhead and stronger Black–Scholes throughput at smaller workloads.

  • DGEMM: Reactant reaches 14.8 TF/s at N = 2048 and 17.3 TF/s at N = 4096 for DGEMM, matching cuBLAS within 1%.
  • DGEMM: JLIR’s untiled DGEMM reaches 1.2 TF/s at N = 4096, below Reactant’s cuBLAS-backed matmul path.
  • Compilation trade-offs: JLIR’s lightweight pass infrastructure adds less than 1.5 ms overhead, whereas Reactant’s compilation path takes several seconds.
  • Black–Scholes: At n = 10 M options, JACC GPU reaches 557 GF/s versus Reactant’s 489 GF/s, a 14% advantage for JLIR-generated code.
  • Black–Scholes: JLIR’s Black–Scholes advantage increases to 3.6× at n = 100 K because Reactant incurs higher per-kernel dispatch and JIT overhead.
  • Overall comparison: Across large workloads, JLIR-generated kernels provide substantial GPU acceleration without programmer annotations.

7 Related Work

JLIR re-hosts MLIR-style infrastructure in Julia, emphasizing low-friction extension and automatic extraction of parallelism from serial Julia code. Related systems instead target different language models, optimization dimensions, or domain-specific schedules.

  • MLIR: JLIR shares MLIR’s module–region–block–operation hierarchy, SSA model, and progressive multi-dialect lowering, but is implemented entirely in Julia.
  • MLIR: JLIR’s macro-based DSL defines dialect operations in a few lines without the C++ build environment required by MLIR’s TableGen workflow.
  • xDSL: xDSL similarly reimplements MLIR in Python, whereas JLIR is tailored to Julia’s JIT environment and supplies an end-to-end JACC GPU pipeline.
  • Equality saturation: JLIR and equality saturation address complementary optimization dimensions: JLIR extracts heterogeneous parallelism, while equality saturation performs general scalar optimization.
  • Julia GPU systems: KernelAbstractions.jl and CUDA.jl require explicit GPU-kernel definitions, while JLIR transforms serial Julia functions automatically through JACCTransformPass.
  • JACC: JACC provides runtime portability, while JLIR performs compile-time parallelism extraction from serial Julia loops.
  • Domain-specific systems: Halide and Triton specialize in image pipelines or tiled neural-network workloads, whereas JLIR targets arbitrary scientific loops with a modest peak-performance trade-off.
  • Reactant and JAX: Reactant inherits XLA backend optimizations but favors functional immutable array programs, while JLIR handles mutable in-place loops and Julia Expr trees.

8 Conclusion

JLIR implements an MLIR-style, Julia-native compilation framework with extensible dialects, passes, and automatic JACC kernel extraction. On NVIDIA A100 benchmarks, generated kernels achieved strong performance without programmer annotations, while current transformations remain limited for cache-sensitive and structurally complex programs.

  • Conclusion: JLIR provides a structured SSA IR, five native dialects, a pass infrastructure, and direct parsing of Julia Expr trees without C++ components.The framework adopts MLIR’s multi-level approach within Julia’s language and tooling ecosystem.
  • Conclusion: JACCTransformPass automatically detects common parallel and reduction loop nests and rewrites them into JACC parallel constructs.Supported patterns include 1D, 2D, mixed parallel-sequential, and associative reduction loops.
  • Conclusion: 96% of the hand-written Julia+CUDA baseline was achieved for Black–Scholes, alongside 3,023 GB/s Jacobi bandwidth, an 85× LLaMA-3 speedup, and 87% of untiled DGEMM.These NVIDIA A100 results required no programmer annotation.
  • Conclusion: The current pass lacks shared-memory tiling and register blocking, declines data-dependent or dependency-heavy nests, and defaults unannotated types to AnyType.These constraints limit cache-sensitive kernel performance and type-directed optimization.
  • Conclusion: JLIR’s planned extensions include tiling and JACC-array reuse to improve bandwidth utilization, arithmetic intensity, and data-transfer efficiency.The proposed tiling pass would operate on the JLIR Memref dialect before JACCTransformPass.
  • Conclusion: JLIR lowers the barrier to high-performance, portable GPU computing by keeping compiler infrastructure within Julia’s familiar language and tooling ecosystem.The supported scope is Julia-native compilation following the MLIR paradigm.

Appendix A: Julia source code and corresponding JLIR code

The appendix presents Julia examples and their corresponding JLIR SSA representations for parallel-for, mixed parallel-sequential, and parallel-reduction patterns. Empty loop iterated arguments distinguish independent loops from loops carrying accumulators.

  • 1D parallel-for: A 1D parallel-for loop has independent iterations and is represented by a JLIR ForOp with empty iter args.The RMS normalization scaling example loads x and weight, multiplies by rms_inv, and stores the result.
  • 2D parallel-for: A 2D parallel-for nests two independent loops, represented by two nested ForOps with empty iter args.The Jacobi stencil updates each output element from neighboring input values.
  • Mixed parallel-sequential: Mixed parallel-sequential GEMM uses empty iter args for the outer loops and one iter arg for the inner k-loop accumulator.The accumulator is initialized, updated through multiplication and addition, yielded, and stored to C.
  • 1D parallel-reduce: A 1D parallel-reduce loop accumulates a scalar with an associative operator, represented by one ForOp carrying the reduction value.The RMS norm example squares each element and adds it to ss.

Appendix B: JACC Generated Output using the JLIR to JACC Transform Pass

The generated JACC output extracts loop bodies into kernels and dispatches them through parallel-for or parallel-reduce operations. Multidimensional launches pass dimension tuples, while mixed GEMM preserves the inner accumulation loop sequentially.

  • 1D parallel-for: A 1D parallel-for extraction makes free variables kernel parameters and dispatches the extracted body through JACC.parallel_for.The generated kernel receives the index and arrays needed by the original loop body.
  • 2D parallel-for: A 2D parallel-for passes tuple(N, M) to JACC.parallel_for, whose kernel receives j and i coordinates for the CUDA thread grid.The inner loop body becomes the kernel and preserves the Jacobi stencil computation.
  • Mixed parallel-sequential: Mixed GEMM parallelizes the outer i and j loops while retaining the k-loop sequentially inside the generated kernel.This preserves correctness of the dot-product accumulation.
  • 1D parallel-reduce: Parallel reduction removes the accumulation and assignment operations from the kernel body, returning per-element contributions to JACC.parallel_reduce.The reduction operator and initial value are passed as keyword arguments; JACC selects a backend-specific reduction strategy.
Loading 2609.04585v1…