Reviewing MLIRBuilder for Impero->MLIR by SamSJackson · Pull Request #5552 · firedrakeproject/firedrake

SamSJackson · 2026-10-23T13:30:55Z

MLIRBuilder - Class to build MLIR from an Impero AST

Principal idea is to receive an Impero AST and eagerly generate MLIR.

Most translation happens within def process. If an object is already processed, the previously-generated SSA value will be used as opposed to re-processing. Additional functions are helper functions for a _process method.

We use memref.Globals for temporary buffers as opposed to memref.alloc (heap memory) or memref.alloca (stack memory) because it generalises to GPU easier and we assign memory hierarchy on GPU (in per-thread, shared, or global memory). This assumes that the temporary buffers have static sizing.

There is also a constant-cache - to avoid a bloated MLIR, even if it is cleaned in the compiler passes.

Core assumptions being made

We are always working with 1D/flat buffers.
The only dynamic loop-extent is the cell-level loop.
Impero.Block objects can be ignored and simply process the children.

To improve

I am confident that we can receive arguments in advance.
- This would allow us to generate FuncOp from beginning and avoid argument sorting.
Need to improve docstrings
Not handling possible math functions

Ongoing bugs

Indexing is not correct for accessing temporary buffers.
Probably more than have not appeared yet

- Requires parsing GEM and Impero types - Function arguments can be known in advance if we parse Impero AST - Need to determine if Parallel indices are determined at GEM->Impero or Impero->MLIR pass. - GEM->Impero happens in FIAT, no changes from `fiat/connorjward/mlir` branch (but there are some changes, such as __str__ definitions that I have locally)

- Some major bugs to solve: - Evaluate does not work for tensor expressions (literally everything from GEM) - Temporary buffer memory is allocated on alloc/alloca - it should be alloc for runtime & non-read-only vs global for compile-time readonly - gem.Variable improperly assumed to be function arguments. Function arguments should be parsed ahead of time. - Need way to determine between ForOp and ParallelOp. Ideally comes from GEM -> Impero parsing step.

- Evaluate expressions return to temporary variables - Need to determine a way to solve their mapping

Regretfully, the MLIR is invalid Bugs to fix: - Not appropriately handling multiindex/VariableIndex - imp.Return is not aligning the processed values with the expected memref store type (i.e. actual: memref.store val, buf[indices] : f64 expected: memref.store val, buf[indices] : memref<Nxf64> for some value N or dynamic index )

Getting segfaults on execution. Debugging in process. cc.py is still goofy with implementation but just prototype for now.

Not flattening indices correctly

Removed segfault and traversing tree in reverse. Only issue now is that the final scatter is not accumulating

Assemble now works for one-form. Barring that the results are not correct. Implementing ReturnAccumulate would fix this but upstream changes will make that pointless. (ReturnAccumulate being replaced with Assignment)

Getting rid of overlapping functions Pretty much finished now for CPU implementations. Waiting for upstream changes to validate. Possible improvements: - Pass kernel buffers to GemExecutable from GemCodegenContext: - Allows us to create FuncOp from getgo and remove Block/detaching and reattaching ops. - Also avoids sorting of arguments

Added a mapping for buffer->device based on desired location, given generation - A device is placed on CPU if it is used for runtime-variable domain extent Need to make temporary arrays use memref.Global as opposed to memref.Alloc because it enables easier shared memory workgroups on GPU and we can easily initialise to 0 on CPU and GPU

- GlobalOps are used instead of `alloc` because not re-initialised per thread. Memory scope can be assigned (per-thread, per-threadblock, global memory) - Outermost function developed for entering outermost parallel (when we enter cell loop) - Additional cleaning and refactoring

connorjward · 2026-10-23T13:57:14Z

+    rank = 1 if arg.shape is None else len(arg.shape)
+    return ctypes.POINTER(Memref.ctype(rank, np.dtype(arg.dtype)))
+
+class Memref:


Suggested change

class Memref:

class MemrefBuilder:

connorjward · 2026-10-23T13:58:44Z

+        return f"<{self.name}, dtype: {self.dtype}, shape: {self.shape or '?'}>"
+
+def get_mlir_type(dtype):
+    # NOTE: Default to f64 if dtype is None, possibly wrong.


At this point we should know all of the dtypes with the possible exception of literals. I think a default here will just cause confusion down the line.

connorjward · 2026-10-23T14:03:09Z

+
+    dt = np.dtype(dtype)
+    if dt.kind == "f":
+        return f64 if dt.itemsize == 8 else f32


Single precision etc? I think we should be explicit about catching everything and raising errors if we don't know what to do

connorjward · 2026-10-23T14:05:15Z

+        self.process(root)
+        n = len(self._arguments)
+
+        perm = sorted(range(n), key=lambda i: self._arguments[i].name)


I am happy to have an arbitrary argument order in the generated code. In which case, can this go?

connorjward · 2026-10-23T14:08:01Z

+        if len(multiindex) > 1: 
+            multiindex = multiindex[1:] 
+
+        return self._load(aggregate, multiindex, [1])


As your comments suggest, this does look quite ugly and will need cleanup

connorjward · 2026-10-23T14:08:35Z

+        return self.process(leaf.expression)
+
+    @_process.register(gem.Indexed)
+    @_process.register(gem.Gather)


Suggested change

@_process.register(gem.Gather)

not using any more

connorjward · 2026-10-23T14:10:49Z

+    _process.register(gem.FloorDiv)(_binary_handler(None, arith.FloorDivSIOp))
+    _process.register(gem.Remainder)(_binary_handler(None, arith.RemSIOp))
+    _process.register(gem.LogicalAnd)(_binary_handler(None, arith.AndIOp))
+    _process.register(gem.LogicalOr)(_binary_handler(None, arith.OrIOp))


_handler_mapping = { gem.Sum: (AddfOp, AddiOp), ... } def _process(op: gem.BinaryOp): val1, val2 = _handler_mapping[type(op)] return _binary_handler(val1, val2)

(assuming something like BinaryOp exists, it should)

also fix the bad variable names

connorjward · 2026-10-23T14:11:17Z

+    _process.register(gem.LogicalAnd)(_binary_handler(None, arith.AndIOp))
+    _process.register(gem.LogicalOr)(_binary_handler(None, arith.OrIOp))
+
+    _CMPF = {"<": "olt", "<=": "ole", ">": "ogt", ">=": "oge", "==": "oeq", "!=": "one"}


certainly needs a comment or clearer name

connorjward · 2026-10-23T14:12:38Z

            print("Not removing cached libraries")
+
+
+@contextmanager


I already have this in pyop3.debug.timer!

connorjward · 2026-10-23T14:13:31Z

                os.close(descriptor)

-                if not compiler.ld:
+                if extension == "mlir":


I think this needs a refactor, but I can do that because I know the code better

SamSJackson added 20 commits October 3, 2026 10:16

Solving Evaluate parsing

60a4bf3

- Evaluate expressions return to temporary variables - Need to determine a way to solve their mapping

Impero->MLIR Generating and Compiling

01962fd

Getting segfaults on execution. Debugging in process. cc.py is still goofy with implementation but just prototype for now.

Allocating 0 values but indices still broke

66e9956

Not flattening indices correctly

Still debugging this indices access problem

2391bd6

Debugging multiindex problem

6dd9a5c

Removed segfault and traversing tree in reverse. Only issue now is that the final scatter is not accumulating

Hiding arg print

b14ae2a

Merge branch 'connorjward/mlir' into SamSJackson/mlir

b1ee716

Compiling and executing (no segfaults!)

29458ce

Assemble now works for one-form. Barring that the results are not correct. Implementing ReturnAccumulate would fix this but upstream changes will make that pointless. (ReturnAccumulate being replaced with Assignment)

Cleaning code

c046b3f

WILO note for where i left off

5f7de55

Merging connor's change

9ae274c

MLIR changes pre-Connor merge

d5cddf7

Merging Connor's change for GemCompileR

6d7c50e

MLIRBuilder for Impero->MLIR

12fd14f

SamSJackson added the base:main Run this PR using a main (dev) build label Oct 23, 2026

SamSJackson requested a review from connorjward October 23, 2026 13:32

SamSJackson changed the title ~~Sam s jackson/mlir~~ Reviewing MLIRBuilder for Impero->MLIR Oct 23, 2026

connorjward requested changes Oct 23, 2026

View reviewed changes

Navigation Menu

Repository navigation

Reviewing MLIRBuilder for Impero->MLIR - #5552

Reviewing MLIRBuilder for Impero->MLIR#5552
SamSJackson wants to merge 20 commits into
connorjward/mlirfrom
SamSJackson/mlir

SamSJackson commented Oct 23, 2026 •

edited

Loading

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

connorjward Oct 23, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

SamSJackson commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

MLIRBuilder - Class to build MLIR from an Impero AST

Core assumptions being made

To improve

Ongoing bugs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

SamSJackson commented Oct 23, 2026 •

edited

Loading