jumerckx / thesis_code

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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)

Overview of the Repository

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.

Examples

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.

Installation Instructions

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:

Building MLIR

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

About

License:MIT License


Languages

Language:Julia 100.0%