Metadsl integration

Hello all,

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.

There isn’t “built in” control-flow operations in MLIR. The cond_br is part of the “standard dialect” (despite its name is just a dialect like others).