I had the pleasure of meeting a few MLIR folks at the recent C4ML conference.
I am working on a very similar project, called
metadsl. I am coming from a different background and targeting a different language for users of the framework (Python), so it looks a bit different (more functional dataflow than CFG, like this topic).
One of my chief questions after leaving the conference is how I can build off and integrate with MLIR, due to all the activity happening in this space.
So a key requirement is to be able to support the Python community in building their own DSLs in Python using it’s type annotations to define the specifications and making the execution happen through replacement patterns, much like the translation in MLIR between dialects.
I know there are variety of software engineering issues to tackle in terms of calling out to MLIR as a shared library from a Python C extension, but those don’t seem like dealbreakers. The bigger question is how to map the underlying semantics of MLIR to expose them in Python.
One nice tool I have been building for myself is a debugger I can use in JupyterLab to step through the replacements on my graph and see it move (source):
Another is the ability to write replacement rules as if you are were dealing just regular python expressions, by defining a function that takes a list of wildcards and return the source and targets. For example, here is the replacement rule for evaluating outer product as a functional array operation:
T = typing.TypeVar("T") Shape = Vec[Integer] Indices = Vec[Integer] class Array(Expression, typing.Generic[T]): @expression @classmethod def create(cls, shape: Shape, index_fn: Abstraction[Indices, T]) -> Array[T]: ... @expression def __matmul__(self, other: Array[T]) -> Array[T]: ... @register @rule def _outer_product( l_shape: Shape, r_shape: Shape, l_idx: Abstraction[Indices, Integer], r_idx: Abstraction[Indices, Integer], ): l_size = l_shape.length @Abstraction.from_fn def index_fn(indices: Indices) -> Integer: return l_idx(indices.take(l_size)) * r_idx(indices.drop(l_size)) return ( Array.create(l_shape, l_idx) @ Array.create(r_shape, r_idx), Array.create(l_shape + r_shape, index_fn), )
Ideally, I would be able to take forms like this and translate them to MLIR dialects and rules at runtime and we could then use MLIR to translate them to other dialects.
One question I had is about control flow in MLIR. Are there any built in operations? In the spec I see some references to
cond_br. Is this simply part of another dialect or is it somehow intrinsic to MLIR?
Another question is it possible to export a declarative specification for any dialect? I.e. some kind of format that tells me what types it provides and what operators? I would expect no, because it seems like many of these things seem like they are dynamically computed from C++ methods? The idea being that if we could do this then we could create deeply embedded versions of these dialects in another language, like Python, mapping MLIR types to Python type annotations, as possible.