edg-l/edgarluque.com
1
0
Fork 0
mirror of https://github.com/edg-l/edgarluque.com.git synced 2024-11-09 09:38:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

blog

Browse source
This commit is contained in:
Edgar 2023-11-27 11:52:19 +01:00
parent b0daa09a54
commit 6dec10da55
No known key found for this signature in database
GPG key ID: 70ADAE8F35904387
1 changed files with 142 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

148
content/blog/mlir-with-rust.md
View file

@ -1,8 +1,7 @@
+++
title = "Intro to LLVM and MLIR with Rust and Melior"
date = 2023-11-26
date = 2023-11-27
description = "Learning MLIR with too many dialects."
draft = true
[taxonomies]
categories = ["rust", "MLIR", "LLVM"]
+++
@ -34,7 +33,7 @@ It has a simple type system, you can work with integers of any bit size (i1, i32
Probably the biggest wall you will hit is learning about GEPs (Get Element Ptr), it's often misunderstood how it works, so they even have a entire documentation page for it: <https://llvm.org/docs/GetElementPtr.html>.
Another thing that may need attention are [PHI nodes](https://stackoverflow.com/questions/11485531/what-exactly-phi-instruction-does-and-how-to-use-it-in-llvm), which are how LLVM selects a value that comes from control flow branches due to the nature of SSA.
The API to build such IR have the following structures:
The API to build such IR have the following structure:
- Context: Opaquely owns and manages the global data of the LLVM infrastructure, including types, uniquing tables, etc.
- Module: The top level container of all other IR objects, it is akin to a compile unit, holds a list of global variables, functions, libraries or other modules it depends on, a symbol table, and target data such as the layout.
@ -46,9 +45,9 @@ If you want to use LLVM with Rust in a type safe manner, I recommend the really
So what is MLIR? It goes a level above, in that LLVM IR itself is one of it's dialects.
MLIR is kind of a IR of IRs, and it supports many of them using "dialects". For example, you may have heard of NVVM IR (CUDA), MLIR supports it through the [NVVM](https://mlir.llvm.org/docs/Dialects/NVVMDialect/) dialect, but there is also a more generic and higher level [GPU](https://mlir.llvm.org/docs/Dialects/GPU/) dialect.
MLIR is kind of a IR of IRs, and it supports many of them using "dialects". For example, you may have heard of NVVM IR (CUDA), MLIR supports modeling it through the [NVVM](https://mlir.llvm.org/docs/Dialects/NVVMDialect/) dialect (or [ROCDL](https://mlir.llvm.org/docs/Dialects/ROCDLDialect/) for AMD), but there is also a more generic and higher level [GPU](https://mlir.llvm.org/docs/Dialects/GPU/) dialect.
Those dialects define *conversion* [passes](https://mlir.llvm.org/docs/Passes/) between them, meaning you can convert IR code using the GPU dialect to the NVVM dialect.
Those dialects define *conversion* [passes](https://mlir.llvm.org/docs/Passes/) between them, meaning for example, you can convert IR code using the GPU dialect to the NVVM dialect.
They also may define dialect passes, for example the `-gpu-map-parallel-loops` which greedily maps loops to GPU hardware dimensions.
@ -56,6 +55,7 @@ Some notable dialects:
- [Builtin](https://mlir.llvm.org/docs/Dialects/Builtin/): The builtin dialect contains a core set of Attributes, Operations, and Types that have wide applicability across a very large number of domains and abstractions. Many of the components of this dialect are also instrumental in the implementation of the core IR.
- [Affine](https://mlir.llvm.org/docs/Dialects/Affine/): This dialect provides a powerful abstraction for affine operations and analyses.
- [Async](https://mlir.llvm.org/docs/Dialects/AsyncDialect/): Types and operations for async dialect This dialect contains operations for modeling asynchronous execution.
- [SCF](https://mlir.llvm.org/docs/Dialects/SCFDialect/): The scf (structured control flow) dialect contains operations that represent control flow constructs such as if and for. Being structured means that the control flow has a structure unlike, for example, gotos or asserts.
- [CF](https://mlir.llvm.org/docs/Dialects/ControlFlowDialect/): This dialect contains low-level, i.e. non-region based, control flow constructs. These constructs generally represent control flow directly on SSA blocks of a control flow graph.
- [LLVM](https://mlir.llvm.org/docs/Dialects/LLVM/): This dialect maps LLVM IR into MLIR by defining the corresponding operations and types. LLVM IR metadata is usually represented as MLIR attributes, which offer additional structure verification.
@ -64,9 +64,9 @@ Some notable dialects:
- [TOSA](https://mlir.llvm.org/docs/Dialects/TOSA/): TOSA was developed after parallel efforts to rationalize the top-down picture from multiple high-level frameworks, as well as a bottom-up view of different hardware target concerns (CPU, GPU and NPU), and reflects a set of choices that attempt to manage both sets of requirements.
- [Func](https://mlir.llvm.org/docs/Dialects/Func/): This dialect contains operations surrounding high order function abstractions, such as calls.
You can also [make your own dialect](https://mlir.llvm.org/docs/Tutorials/CreatingADialect/), useful to make a domain specific language, in this dialect you can define transformations to other dialects, passes, etc.
You can also [make your own dialect](https://mlir.llvm.org/docs/Tutorials/CreatingADialect/), useful to make a domain specific language for example, in this dialect you can define transformations to other dialects, passes, etc.
All these dialects can exist in your code at the same time, but at the end, you want to execute your code, for this there are Targets, one is LLVM IR itself. In this case, you would need to use passes to convert all dialects to the LLVM dialect, and then you can make the [translation from MLIR to LLVM IR](https://mlir.llvm.org/docs/TargetLLVMIR/).
All these dialects can exist in your MLIR code at the same time, but at the end, you want to execute your code, for this there are Targets, one is LLVM IR itself. In this case, you would need to use passes to convert all dialects to the LLVM dialect, and then you can make the [translation from MLIR to LLVM IR](https://mlir.llvm.org/docs/TargetLLVMIR/).
The structure of MLIR is recursive as follows:
@ -78,6 +78,24 @@ The top level module is also a operation, which holds a single region with a sin
A region can have 1 or more blocks, each block can have one or more operations, a operation can use 1 or more regions.
## Operations
These provides the functionality, and what make up the bulk of MLIR.
A operation has the following properties:
- Name: The name of the operation, a unique identified within MLIR. Operations live within a dialect, so they are refered to using `dialect.operation`,
for example `arith.add`
- Traits: A operation can have a set of traits that affect the syntax or semantics, for example, whether it has side effects,
whether it's a block terminator, etc.
- Constraints: These are used when verifying a operation is correct, for example whether the operands have matching shape and types.
- Arguments: There are 2 types of arguments: operands, which are runtime values, and attributes, which are compile time constant values.
- Regions: A operation can accept regions, which contain blocks within.
- Results: The results of the operation, which can be more than 1. For example `arith.add` has 1 result, defined by the type of it's arguments.
- Successors: When a operation is a terminator it needs successors, for example `cf.br` which is a unconditional jump and thus a branching operation, accepts a single successor (block).
You can read more about operations in the [Operation Definition Specification](https://mlir.llvm.org/docs/DefiningDialects/Operations/).
To use MLIR with Rust, I recommend [melior](https://github.com/raviqqe/melior), here is a snippet making a function that adds 2 numbers:
```rust
@ -90,6 +108,7 @@ use melior::{
// We need a registry to hold all the dialects
let registry = DialectRegistry::new();
// Register all dialects that come with MLIR.
register_all_dialects(&registry);
// The MLIR context, like the LLVM one.
@ -105,6 +124,7 @@ let location = Location::unknown(&context);
let module = Module::new(location);
// A integer-like type with platform dependent bit width. (like size_t or usize)
// This is a type defined in the Builtin dialect.
let index_type = Type::index(&context);
// Append a `func::func` operation to the body (a block) of the module.
@ -117,7 +137,9 @@ let index_type = Type::index(&context);
// These blocks each can have more operations.
module.body().append_operation(func::func(
&context,
// accepts a StringAttribute which is the function name.
StringAttribute::new(&context, "add"),
// A type attribute, defining the function signature.
TypeAttribute::new(
FunctionType::new(&context, &[index_type, index_type], &[index_type]).into()
),
@ -141,6 +163,9 @@ module.body().append_operation(func::func(
func::r#return( &[sum.result(0).unwrap().into()], location)
);
// The Func operation requires a region,
// we add the block we created to the region and return it,
// which is passed as an argument to the `func::func` function.
let region = Region::new();
region.append_block(block);
region
@ -151,3 +176,112 @@ module.body().append_operation(func::func(
assert!(module.as_operation().verify());
```
Here is a more complex function, using the `SCF` dialect, which allows us to use a `while` loop:
```rust
let context = Context::new();
load_all_dialects(&context);
let location = Location::unknown(&context);
let module = Module::new(location);
let index_type = Type::index(&context);
let float_type = Type::float64(&context);
module.body().append_operation(func::func(
&context,
StringAttribute::new(&context, "foo"),
TypeAttribute::new(FunctionType::new(&context, &[], &[]).into()),
{
let block = Block::new(&[]);
let initial = block.append_operation(arith::constant(
&context,
IntegerAttribute::new(0, index_type).into(),
location,
));
block.append_operation(r#while(
&[initial.result(0).unwrap().into()],
&[float_type],
{
let block = Block::new(&[(index_type, location)]);
let condition = block.append_operation(arith::constant(
&context,
IntegerAttribute::new(0, IntegerType::new(&context, 1).into())
.into(),
location,
));
let result = block.append_operation(arith::constant(
&context,
FloatAttribute::new(&context, 42.0, float_type).into(),
location,
));
block.append_operation(super::condition(
condition.result(0).unwrap().into(),
&[result.result(0).unwrap().into()],
location,
));
let region = Region::new();
region.append_block(block);
region
},
{
let block = Block::new(&[(float_type, location)]);
let result = block.append_operation(arith::constant(
&context,
IntegerAttribute::new(42, Type::index(&context)).into(),
location,
));
block.append_operation(r#yield(
&[result.result(0).unwrap().into()],
location,
));
let region = Region::new();
region.append_block(block);
region
},
location,
));
block.append_operation(func::r#return(&[], location));
let region = Region::new();
region.append_block(block);
region
},
&[],
location,
));
assert!(module.as_operation().verify());
```
This code generates the following MLIR IR:
```mlir
module {
func.func @foo() {
%c0 = arith.constant 0 : index
%0 = scf.while (%arg0 = %c0) : (index) -> f64 {
%false = arith.constant false
%cst = arith.constant 4.200000e+01 : f64
scf.condition(%false) %cst : f64
} do {
^bb0(%arg0: f64):
%c42 = arith.constant 42 : index
scf.yield %c42 : index
}
return
}
}
```
There is way more to MLIR, but this is meant to be a small introduction.