Generate MLIR from Julia code.
Jojo.generate(f, Tuple{i64, i64}) do a,b
a>b ? a*b : a+b
end
func.func @"#14"(%arg0: i64, %arg1: i64) -> i64 attributes {llvm.emit_c_interface} {
%0 = arith.cmpi sgt, %arg0, %arg1 : i64
cf.cond_br %0, ^bb1, ^bb2
^bb1: // pred: ^bb0
%1 = arith.muli %arg0, %arg1 : i64
return %1 : i64
^bb2: // pred: ^bb0
%2 = arith.addi %arg0, %arg1 : i64
return %2 : i64
}
(see /examples/simple.jl
for complete example)
generate.jl
contains the main MLIR code generation interpreter loop.
codegenContext.jl
contains the AbstractCodegenContext
and default CodegenContext
.
Inlining policy changes and Boolean conversions are handled by the MLIRInterpreter defined in abstract.jl
.
src2src.jl
is an experiment to generate MLIR code by transforming the Julia IR code and executing it, instead manual abstract interpretation over the regular Julia SSA IR.
Filename | Description |
---|---|
simple.jl |
Illustrates the basic MLIR code generation process + execution on a simple function with control flow. |
regions.jl |
Example of a custom CodegenContext and intrinsic functions that generate operations with regions through the use of higher-order functions. |
einsum.jl |
An example of a DSL for generating linalg.generic operations given an einsum description. |
gpu_vadd.jl |
vadd kernel similar to CUDA.jl example. |
gpu_wmma.jl |
Vendor-neutral WMMA operations. |
gpu_mma_from_linalg.jl |
Example of transforming a m16n8k16 linalg.matmul into hardware-specific NVPTX (mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 ). |
transform.jl |
Port of the MLIR transform tutorial that reproduces a Halide schedule. |
src2src_example.jl |
Demonstration of the approach of generating MLIR code by transforming and executing Julia SSA IR. |
The code depends on Julia compiler internals and therefore won't work with many different Julia versions.
Development and testing was done with 1.11.0-beta1
.
Julia Version 1.11.0-beta1
Commit 08e1fc0abb9 (2024-04-10 08:40 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 12 × Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
WORD_SIZE: 64
LLVM: libLLVM-16.0.6 (ORCJIT, skylake)
Instantiate the packages in Manifest.toml
:
] instantiate
This will install forks with small changes that haven't yet been upstreamed for:
To run the examples, a path to the MLIR C API library libMLIR-C.so
needs to be provided.
This can be done by creating LocalPreferences.toml
file in the root directory containing:
[MLIR_jll]
mlir_c_path = "[LLVM install directory]/lib/libMLIR-C.so"
A sufficiently recent version of LLVM/MLIR might work, but development was done on commit 8a237ab7d9022d24441544ba25be480f0c944f5a.
Alternatively, my fork of LLVM includes a few additional commits that work around a build error, as well as support for extracting generated PTX to run with CUDA.jl, and some commits from a stale pull-request upstream that allows lowering WMMA operations for AMD: https://github.com/jumerckx/llvm-project/
LLVM has to be built with the following flags:
LLVM_INSTALL_UTILS=ON
LLVM_BUILD_LLVM_DYLIB=ON
LLVM_EXTERNAL_MLIR_SOURCE_DIR=/home/jumerckx/llvm-project/mlir
LLVM_TARGETS_TO_BUILD=host;NVPTX;AMDGPU # NVPTX and AMDGPU are only needed for NVIDIA and AMD GPUs respectively.
LLVM_TOOL_MLIR_BUILD=ON
MLIR_BUILD_MLIR_C_DYLIB=ON
MLIR_ENABLE_CUDA_RUNNER=ON