Basic MLIR Code Generation into PyOP3

[SamSJackson]

PR: MLIR Code Generation for PyOP3

This PR presents the basic MLIR generation for operations such as array assignment.

MLIR is being considered as an alternative to Loopy as the multi-level compiler passes allow us to target CPUs and GPUs from a similar base IR. MLIR also provides compiler passes that would allow us to optimise for the respective GPU targets.

At this stage, this PR presents basic generation of MLIR for CPUs.

How does it work?

Within pyop3/compile/, the file core.py traverses the PyOP3 AST and passes buffer expressions to mlir.py.

The MLIRCodegenContext class parses the buffer expressions into SSAValues.
MLIR is an Static Single Assignment (SSA) compiler, which means that variables can be defined exactly once (i.e. x = x + 1 is not allowed).
MLIR is composed of dialects, where dialects correspond to components of an IR. These dialects are then lowered through compiler passes and optimisations, which allows us to have multiple targets.
In this work, the key dialects used:

• scf: Structured Control Flow serves as the iteration domain for our operations. Higher-level dialects such as linalg or affine could not be used due to the indirect memory accesses, which are not affine indexing operations.
• memref: This is our dialect for buffers. As an alternative, the higher-level tensor could be used but it does not make sense to mix dialect levels. Further, without linalg/affine, compiler optimisations for tensor cannot be utillised. Sources from xDSL also suggested that memref was the better choice here.
• arith: This represents constant values as well as arithmetic operations.

A key design choice in this implementation is that MLIR is produced as the PyOP3 AST is parsed by core.py.
Opposingly, Loopy collects all expressions into pymbolic expressions before summarily generating code.
Our approach was chosen because collecting the PyOP3 expressions would have required an additional IR for temporary storage of variables. The collection-generation approach was initially discussed and attempted but I concluded it was bloated.

As the MLIR FuncOp requires knowledge of all arguments before creation, MLIR is initially generated to a Block. Arguments can then be introduced procedurally. At finalisation, the arguments and operations are detached from the Block and a FuncOp is created, operations attached. Arguments are also re-ordered to match Loopy.

Iteration domains are produced eagerly and an internal symbol table is employed to ensure re-use of key SSA values (such as iteration variables and buffer args).

What happens next?

Foremost, there is some software engineering flaws with this implementation. Namely, the type inference.
There was a dtype property introduces into PyOP3 expressions but MLIR also uses an index type, an int required for load/storing memref/tensors. I have yet to produce a cleaner approach for dealing with this.
Also some more cleaning to do more generally, I am sure.

Regarding future work: this PR does not introduce compilation or execution for PyOP3.

Compilation requires an MLIR binary in PyOP3 as well as working on pyop3/cc.py. I have not viewed how difficult a task that this will be just yet.
For compilation, there must be a lowering pass (from mlir-opt), a translation pass (from mlir-translate) and a compilation to a shared library (from clang). Minimal passes for the lowering passes have been found, experimentation can be done in the future to find an optimal set of passes. The set of passes should be affected by target hardware.

In terms of execution, it is a case of wrapping numpy arrays in corresponding memref descriptors. This should happen in pyop3/insn/exec.py. It should not be a particularly invasive change.

[connorjward]

Super cool!

There was a dtype property introduces into PyOP3 expressions but MLIR also uses an index type, an int required for load/storing memref/tensors.

We do sort of know when we have indices vs integers. E.g. consider

idat[2*i+j]

We know that idat is the outermost thing and hence the dtype is an integer, but when we recurse to look at 2*i+j we now know that only index types make any sense.

Other than this I think my main points are:

1. The current scoping implementation is quite confusing. I much prefer context managers over try, finally.
2. MLIR is confusing. Code like

        # Move old block ops into the func entry block
        ops = list(self._entry_block.ops)
        for op in ops:
            op.detach()
        new_block.add_ops(ops)
        new_block.add_op(func.ReturnOp())

could really benefit from more comments as the API is utterly foreign.

[connorjward]

I'd turn this into another singledispatch function. We are redundantly turning pyop3.expr.Add into "add" just to case it. Just .register(pyop3.expr.Add) instead.

[connorjward]

To explain things a little, IntType, RealType and ScalarType all come from PETSc. They can change depending on how PETSc is configured. E.g. IntType is usually int32 but can also be int64.

[SamSJackson]

IntegerType requires an additional argument for bit-width. So, to work with the adaptable nature, I read the bit-width from np.dtype given as IntType and apply that to IntegerType

[connorjward]

I don't really understand this. Looks like it's just doing

self._entry_block.args = new_block.args

but in a strange way

[SamSJackson]

I think I will add lots of comments around this section.

This loop is re-ordering the function arguments (the SSA values).

When an argument is inserted into a block or operation, downstream SSA values are influenced.
The replace_by is addressing any downstream influences between the pair of SSA values (old_arg, new_arg).
i.e. we cannot simply re-order the function signature but also replace downstream mentions of old_arg.