mlir_codegen

MLIRBufferFType

Bases: BufferFType, MLIRArgumentFType, ABC

Abstract base class for the ftype of datastructures. The ftype defines how the data in an Buffer is organized and accessed.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRBufferFType(BufferFType, MLIRArgumentFType, ABC):
    """
    Abstract base class for the ftype of datastructures. The ftype defines how
    the data in an Buffer is organized and accessed.
    """

    @abstractmethod
    def mlir_length(self, ctx: MLIRContext, buffer):
        """
        Return MLIR code which loads a named buffer at the given index.
        """
        ...

    @abstractmethod
    def mlir_load(self, ctx: MLIRContext, buffer, index):
        """
        Return MLIR code which loads a named buffer at the given index.
        """
        ...

    @abstractmethod
    def mlir_store(self, ctx: MLIRContext, buffer, index, value=None):
        """
        Return MLIR code which stores a named buffer to the given index.
        """
        ...

    @abstractmethod
    def mlir_resize(self, ctx: MLIRContext, buffer, new_length):
        """Resize a buffer and store its updated descriptor in the slot box."""
        ...

mlir_length(ctx, buffer) abstractmethod

Return MLIR code which loads a named buffer at the given index.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_length(self, ctx: MLIRContext, buffer):
    """
    Return MLIR code which loads a named buffer at the given index.
    """
    ...

mlir_load(ctx, buffer, index) abstractmethod

Return MLIR code which loads a named buffer at the given index.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_load(self, ctx: MLIRContext, buffer, index):
    """
    Return MLIR code which loads a named buffer at the given index.
    """
    ...

mlir_resize(ctx, buffer, new_length) abstractmethod

Resize a buffer and store its updated descriptor in the slot box.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_resize(self, ctx: MLIRContext, buffer, new_length):
    """Resize a buffer and store its updated descriptor in the slot box."""
    ...

mlir_store(ctx, buffer, index, value=None) abstractmethod

Return MLIR code which stores a named buffer to the given index.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_store(self, ctx: MLIRContext, buffer, index, value=None):
    """
    Return MLIR code which stores a named buffer to the given index.
    """
    ...

MLIRCompiler

Bases: MLIRForm, AssemblyLoader

A class to compile and run FinchAssembly.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRCompiler(MLIRForm, asm.AssemblyLoader):
    """
    A class to compile and run FinchAssembly.

    """

    def __init__(self, ctx: MLIRLowerer | None = None):
        self.ctx: MLIRLowerer = MLIRGenerator() if ctx is None else ctx

    def lower(
        self, prgm: asm.Module, *, mode: CompilerMode | None = None
    ) -> MLIRLibrary:
        if prgm.head() != asm.Module:
            raise ValueError(
                "MLIRCompiler expects a Module as the head of the program, "
                f"got {type(prgm.head())}"
            )
        mlir_code = self.ctx(prgm, mode=mode).code
        logger.debug(f"Compiling MLIR code:\n{mlir_code}")
        context, module, engine = load_mlir_engine(mlir_code)
        kernels = {}
        for func in prgm.funcs:
            match func:
                case asm.Function(asm.Variable(func_name, func_type), _, _):
                    kernels[func_name] = MLIRKernel(
                        engine,
                        func_name,
                        func_type,
                    )
                case _:
                    raise NotImplementedError(
                        f"Unrecognized function type: {type(func)}"
                    )
        return MLIRLibrary((context, module, engine), kernels)

MLIRContext

Bases: Context

A context for generating MLIR code.

The context tracks SSA values, variable and slot bindings, indentation, nested blocks, constants, and operations emitted into the MLIR module.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRContext(Context):
    """
    A context for generating MLIR code.

    The context tracks SSA values, variable and slot bindings, indentation,
    nested blocks, constants, and operations emitted into the MLIR module.

    """

    def __init__(
        self,
        tab="  ",
        indent=1,
        bindings=None,
        slots=None,
        *,
        mode=None,
    ):
        if bindings is None:
            bindings = ScopedDict()
        if slots is None:
            slots = ScopedDict()

        super().__init__(mode=mode)
        self.tab = tab
        self.indent = indent
        self.bindings = bindings
        self.slots = slots

    @property
    def feed(self) -> str:
        return self.tab * self.indent

    def new_ssa(self) -> str:
        return "%" + self.freshen("v")

    def constant(self, value, type_: str) -> str:
        key = f".const:{value}:{type_}"
        if key in self.bindings:
            return self.bindings[key][0]
        ssa = self.new_ssa()
        self.exec(f"{self.feed}{ssa} = arith.constant {value} : {type_}")
        self.bindings[key] = (ssa, type_)
        return ssa

    def resolve(self, node):
        match node:
            case asm.Slot(var_n, _):
                if var_n in self.slots:
                    fields, _ = self.slots[var_n]
                    return fields
                raise KeyError(f"Slot {var_n} not found in context")
            case _:
                raise ValueError(f"Expected Slot, got: {type(node)}")

    def cache(self, name, val):
        if isinstance(val, asm.Literal | asm.Variable):
            return val
        var = asm.Variable(self.freshen(name), val.result_type)
        self(asm.Assign(var, val))
        return var

    def emit(self):
        return "\n".join([*self.preamble, *self.epilogue])

    def emit_global(self):
        body = self.emit()
        new_funcs = []

        for name, value in mlir_globals.items():
            if f"func.call @{name}(" in body:
                new_funcs.append(value)

        funcs = "\n".join(new_funcs)
        if funcs:
            funcs += "\n"

        return f"module {{\n{funcs}{body}\n}}\n"

    def block(self) -> MLIRContext:
        blk = super().block()
        blk.tab = self.tab
        blk.indent = self.indent
        blk.bindings = self.bindings
        blk.slots = self.slots
        return blk

    def subblock(self):
        blk = self.block()
        blk.indent = self.indent + 1
        blk.bindings = self.bindings.scope()
        blk.slots = self.slots.scope()
        return blk

    def loop_carried_names(self, body):
        names = []
        nodes = [body]

        while nodes:
            node = nodes.pop()
            match node:
                case asm.Assign(asm.Variable(name, _), _):
                    if name in self.bindings and name not in names:
                        names.append(name)
                case asm.SetAttr(obj, _, _):
                    while isinstance(obj, asm.GetAttr):
                        obj = obj.obj
                    if (
                        isinstance(obj, asm.Variable)
                        and obj.name in self.bindings
                        and obj.name not in names
                    ):
                        names.append(obj.name)
                case asm.Block(bodies):
                    nodes.extend(reversed(bodies))
                case asm.IfElse(_, if_body, else_body):
                    nodes.extend((else_body, if_body))
                case (
                    asm.If(_, nested)
                    | asm.ForLoop(_, _, _, nested)
                    | asm.BufferLoop(_, _, nested)
                    | asm.WhileLoop(_, nested)
                ):
                    nodes.append(nested)

        return names

    def __call__(self, prgm: asm.AssemblyNode):
        feed = self.feed
        match prgm:
            case asm.Assert(exp):
                condition = self(exp)
                message = json.dumps(f"Finch assertion failed: {exp}")
                self.exec(f"{feed}cf.assert {condition}, {message}")
                return None
            case asm.Literal(value):
                # A Scalar-wrapped literal (a sparse gap read) carries its dtype
                # in element_type; the Scalar's own ftype is a struct, which is
                # not a type `arith.constant` can be given.
                value_type = _promote_type_helper([prgm])
                value = getattr(value, "val", value)
                t = mlir_type(value_type)
                new = float(value) if MLIROperator.is_float(value_type) else int(value)
                return self.constant(new, t)

            case asm.Variable(name, _):
                if name not in self.bindings:
                    raise ValueError(f"Variable does not exist: {name!r}")
                value, type_ = self.bindings[name]
                if value.startswith("@"):
                    result = self.new_ssa()
                    self.exec(f"{feed}{result} = func.constant {value} : {type_}")
                    return result
                return value

            case asm.Assign(asm.Variable(var_n, var_t), val):
                v = self(val)
                self.bindings[var_n] = (v, mlir_type(var_t))
                return None

            case asm.Call(op, args):
                return mlir_function_call(op, self, *args)

            case asm.Length(buffer):
                buf_t = buffer.result_type
                if not isinstance(buf_t, MLIRBufferFType):
                    raise TypeError(f"Expected MLIR buffer type, got: {buf_t}")
                return buf_t.mlir_length(self, self.resolve(buffer))

            case asm.Load(buffer, index):
                buf_t = buffer.result_type
                if not isinstance(buf_t, MLIRBufferFType):
                    raise TypeError(f"Expected MLIR buffer type, got: {buf_t}")
                if self.mode.debug:
                    index = self.cache("index", index)
                    self(
                        asm.Assert(
                            asm.Call(
                                asm.Literal(ffuncs.and_),
                                (
                                    asm.Call(
                                        asm.Literal(ffuncs.ge),
                                        (index, asm.Literal(index.result_type(0))),
                                    ),
                                    asm.Call(
                                        asm.Literal(ffuncs.lt),
                                        (index, asm.Length(buffer)),
                                    ),
                                ),
                            )
                        )
                    )
                return buf_t.mlir_load(self, self.resolve(buffer), index)

            case asm.Store(buffer, index, value):
                buf_t = buffer.result_type
                if not isinstance(buf_t, MLIRBufferFType):
                    raise TypeError(f"Expected MLIR buffer type, got: {buf_t}")
                if self.mode.debug:
                    index = self.cache("index", index)
                    self(
                        asm.Assert(
                            asm.Call(
                                asm.Literal(ffuncs.and_),
                                (
                                    asm.Call(
                                        asm.Literal(ffuncs.ge),
                                        (index, asm.Literal(index.result_type(0))),
                                    ),
                                    asm.Call(
                                        asm.Literal(ffuncs.lt),
                                        (index, asm.Length(buffer)),
                                    ),
                                ),
                            )
                        )
                    )
                return buf_t.mlir_store(self, self.resolve(buffer), index, value)

            case asm.Resize(asm.Slot(var_n, _) as buffer, size):
                buf_t = buffer.result_type

                if not isinstance(buf_t, MLIRBufferFType):
                    raise TypeError(f"Expected MLIR buffer type, got {buf_t}")

                buf_t.mlir_resize(self, self.resolve(buffer), size)
                return None

            case asm.GetAttr(obj, attr):
                attrs = [attr.val]
                base = obj
                while True:
                    match base:
                        case asm.GetAttr(inner, a):
                            attrs.append(a.val)
                            base = inner
                        case _:
                            break
                attrs.reverse()
                obj_t = base.result_type
                if not isinstance(obj_t, StructFType):
                    raise TypeError(f"Expected struct type, got: {obj_t}")
                return mlir_getattr(obj_t, self, self(base), attrs)

            case asm.SetAttr(obj, asm.Literal(attr), value):
                match obj.result_type:
                    case MutableStructFType() as obj_t:
                        obj_value = self(obj)
                        val = self(value)
                        value_t = mlir_type(obj_t.struct_attrtype(attr))
                        stored_t = llvm_type(value_t)
                        if stored_t != value_t:
                            cast = self.new_ssa()
                            if value_t == "index":
                                self.exec(
                                    f"{feed}{cast} = arith.index_cast {val} "
                                    ": index to i64"
                                )
                            else:
                                self.exec(
                                    f"{feed}{cast} = "
                                    "builtin.unrealized_conversion_cast "
                                    f"{val} : {value_t} to {stored_t}"
                                )
                            val = cast
                        ptr = self.new_ssa()
                        self.exec(
                            f"{feed}{ptr} = llvm.getelementptr {obj_value}"
                            f"[0, {obj_t.struct_fieldnames.index(attr)}] "
                            f": (!llvm.ptr) -> !llvm.ptr, {mlir_struct_type(obj_t)}"
                        )
                        self.exec(
                            f"{feed}llvm.store {val}, {ptr} : {stored_t}, !llvm.ptr"
                        )
                        return None
                attrs = [attr]
                base = obj
                while isinstance(base, asm.GetAttr):
                    attrs.append(base.attr.val)
                    base = base.obj
                attrs.reverse()

                if not isinstance(base, asm.Variable):
                    raise TypeError(f"Expected variable-backed struct, got: {base}")

                name = base.name
                obj_t = base.result_type
                if not isinstance(obj_t, StructFType):
                    raise TypeError(f"Expected struct type, got: {obj_t}")

                indices = []
                field_t = obj_t
                for field in attrs:
                    if not isinstance(field_t, StructFType):
                        raise TypeError(f"Expected struct type, got: {field_t}")
                    indices.append(field_t.struct_fieldnames.index(field))
                    field_t = field_t.struct_attrtype(field)

                obj = self.bindings[name][0]
                val = self(value)
                value_t = mlir_type(field_t)
                stored_t = llvm_type(value_t)
                if stored_t != value_t:
                    cast = self.new_ssa()
                    if value_t == "index":
                        self.exec(
                            f"{feed}{cast} = arith.index_cast {val} : index to i64"
                        )
                    else:
                        self.exec(
                            f"{feed}{cast} = builtin.unrealized_conversion_cast "
                            f"{val} : {value_t} to {stored_t}"
                        )
                    val = cast

                result = self.new_ssa()

                self.exec(
                    f"{feed}{result} = llvm.insertvalue {val}, {obj}"
                    f"[{', '.join(map(str, indices))}] "
                    f": {mlir_type(obj_t)}"
                )

                self.bindings[name] = (result, mlir_type(obj_t))
                return None

            case asm.Unpack(asm.Slot(var_n, var_t), val):
                if val.result_type != var_t:
                    raise TypeError(f"Type mismatch: {val.result_type} != {var_t}")
                self.slots[var_n] = (var_t.mlir_unpack(self, var_n, val), var_t)
                return None

            case asm.Repack(asm.Slot(var_n, var_t)):
                obj, _ = self.slots[var_n]
                var_t.mlir_repack(self, var_n, obj)
                return None

            case asm.Block(bodies):
                ctx_2 = self.block()
                for body in bodies:
                    ctx_2(body)

                block = ctx_2.emit()
                if block:
                    self.exec(block)
                return None

            case asm.ForLoop(asm.Variable(var_n, var_t), start, end, body):
                lo, hi = self(start), self(end)
                for bound_t, name in (
                    (start.result_type, "lo"),
                    (end.result_type, "hi"),
                ):
                    if mlir_type(bound_t) != "index":
                        cast = self.new_ssa()
                        old = lo if name == "lo" else hi
                        self.exec(
                            f"{feed}{cast} = arith.index_cast {old} "
                            f": {mlir_type(bound_t)} to index"
                        )
                        if name == "lo":
                            lo = cast
                        else:
                            hi = cast
                step = self.constant(1, "index")
                iv = self.new_ssa()
                ctx_2 = self.subblock()
                loop_names = self.loop_carried_names(body)
                initial = [self.bindings[name] for name in loop_names]
                iter_args = [self.new_ssa() for _ in loop_names]
                for name, arg, (_, type_) in zip(
                    loop_names, iter_args, initial, strict=True
                ):
                    ctx_2.bindings[name] = (arg, type_)
                if mlir_type(var_t) == "index":
                    ctx_2.bindings[var_n] = (iv, "index")
                else:
                    iv_cast = ctx_2.new_ssa()
                    ctx_2.exec(
                        f"{ctx_2.feed}{iv_cast} = arith.index_cast {iv} "
                        f": index to {mlir_type(var_t)}"
                    )
                    ctx_2.bindings[var_n] = (iv_cast, mlir_type(var_t))
                ctx_2(body)

                if loop_names:
                    yield_val = [ctx_2.bindings[name][0] for name in loop_names]
                    yield_t = [type_ for _, type_ in initial]
                    ctx_2.exec(
                        f"{ctx_2.feed}scf.yield {', '.join(yield_val)} "
                        f": {', '.join(yield_t)}"
                    )
                    results = [self.new_ssa() for _ in loop_names]
                    result_sp = (
                        results[0]
                        if len(results) == 1
                        else f"{results[0]}:{len(results)}"
                    )
                    result_val = (
                        results
                        if len(results) == 1
                        else [f"{results[0]}#{i}" for i in range(len(results))]
                    )
                    iter_sp = [
                        f"{arg} = {value}"
                        for arg, (value, _) in zip(iter_args, initial, strict=True)
                    ]
                    self.exec(
                        f"{feed}{result_sp} = scf.for {iv} = {lo} to {hi} "
                        f"step {step} iter_args({', '.join(iter_sp)}) "
                        f"-> ({', '.join(yield_t)}) "
                        f"{{\n{ctx_2.emit()}\n{feed}}}"
                    )
                    for name, result, type_ in zip(
                        loop_names, result_val, yield_t, strict=True
                    ):
                        self.bindings[name] = (result, type_)
                    return None

                self.exec(
                    f"{feed}scf.for {iv} = {lo} to {hi} step {step} {{\n"
                    f"{ctx_2.emit()}\n"
                    f"{feed}}}"
                )
                return None

            case asm.BufferLoop(buffer, asm.Variable(var_n, var_t), body):
                buf_t = buffer.result_type
                if not isinstance(buf_t, MLIRBufferFType):
                    raise TypeError(f"Expected MLIR buffer type, got: {buf_t}")
                buf = self.resolve(buffer)
                length = buf_t.mlir_length(self, buf)

                zero = self.constant(0, "index")
                step = self.constant(1, "index")

                iv = self.new_ssa()
                ctx_2 = self.subblock()
                loop_names = self.loop_carried_names(body)
                initial = [self.bindings[name] for name in loop_names]
                iter_args = [self.new_ssa() for _ in loop_names]
                for name, arg, (_, type_) in zip(
                    loop_names, iter_args, initial, strict=True
                ):
                    ctx_2.bindings[name] = (arg, type_)

                idx_name = ctx_2.freshen(".buffer_index")
                ctx_2.bindings[idx_name] = (iv, "index")

                elem = buf_t.mlir_load(ctx_2, buf, asm.Variable(idx_name, algebra.intp))
                ctx_2.bindings[var_n] = (elem, mlir_type(var_t))
                ctx_2(body)

                if loop_names:
                    yield_val = [ctx_2.bindings[name][0] for name in loop_names]
                    yield_t = [type_ for _, type_ in initial]
                    ctx_2.exec(
                        f"{ctx_2.feed}scf.yield {', '.join(yield_val)} "
                        f": {', '.join(yield_t)}"
                    )
                    results = [self.new_ssa() for _ in loop_names]
                    result_sp = (
                        results[0]
                        if len(results) == 1
                        else f"{results[0]}:{len(results)}"
                    )
                    result_val = (
                        results
                        if len(results) == 1
                        else [f"{results[0]}#{i}" for i in range(len(results))]
                    )
                    iter_sp = [
                        f"{arg} = {value}"
                        for arg, (value, _) in zip(iter_args, initial, strict=True)
                    ]
                    self.exec(
                        f"{feed}{result_sp} = scf.for {iv} = {zero} to {length} "
                        f"step {step} iter_args({', '.join(iter_sp)}) "
                        f"-> ({', '.join(yield_t)}) "
                        f"{{\n{ctx_2.emit()}\n{feed}}}"
                    )
                    for name, result, type_ in zip(
                        loop_names, result_val, yield_t, strict=True
                    ):
                        self.bindings[name] = (result, type_)
                    return None

                self.exec(
                    f"{feed}scf.for {iv} = {zero} to {length} step {step} {{\n"
                    f"{ctx_2.emit()}\n"
                    f"{feed}}}"
                )
                return None

            case asm.If(condition, body):
                cond = self(condition)

                before = {}
                vars = self.bindings

                while vars is not None:
                    for i, j in vars.bindings.items():
                        if i not in before:
                            before[i] = j
                    vars = vars.parent

                new_ctx = self.subblock()
                new_ctx.bindings = ScopedDict(before.copy())
                new_ctx(body)
                body_bindings = new_ctx.bindings.bindings

                changed = {
                    name: body_bindings[name]
                    for name, old_binding in before.items()
                    if not name.startswith(".")
                    and body_bindings.get(name) != old_binding
                }
                if not changed:
                    self.exec(f"{feed}scf.if {cond} {{\n{new_ctx.emit()}\n{feed}}}")
                else:
                    names = list(changed)
                    new_result = [self.new_ssa() for _ in names]

                    new_vals = [changed[i][0] for i in names]
                    new_type = [changed[i][1] for i in names]
                    old_vals = [before[i][0] for i in names]

                    new_ctx.exec(
                        f"{new_ctx.feed}scf.yield "
                        f"{', '.join(new_vals)} : {', '.join(new_type)}"
                    )

                    self.exec(
                        f"{feed}{', '.join(new_result)} = scf.if {cond} "
                        f"-> ({', '.join(new_type)}) {{\n"
                        f"{new_ctx.emit()}\n"
                        f"{feed}}} else {{\n"
                        f"{new_ctx.feed}scf.yield {', '.join(old_vals)} : "
                        f"{', '.join(new_type)}\n"
                        f"{feed}}}"
                    )

                    for i, j, k in zip(
                        names,
                        new_result,
                        new_type,
                        strict=True,
                    ):
                        self.bindings[i] = (j, k)

                return None

            case asm.IfElse(condition, body, else_body):
                cond = self(condition)

                before = {}
                vars = self.bindings

                while vars is not None:
                    for i, j in vars.bindings.items():
                        if i not in before:
                            before[i] = j
                    vars = vars.parent

                new_ctx = self.subblock()
                new_ctx.bindings = ScopedDict(before.copy())
                new_ctx(body)
                body_bindings = new_ctx.bindings.bindings

                else_ctx = self.subblock()
                else_ctx.bindings = ScopedDict(before.copy())
                else_ctx(else_body)
                else_bindings = else_ctx.bindings.bindings

                names = sorted(
                    i
                    for i in set(body_bindings) & set(else_bindings)
                    if not i.startswith(".")
                    and (
                        body_bindings[i] != before.get(i)
                        or else_bindings[i] != before.get(i)
                    )
                )
                if not names:
                    self.exec(
                        f"{feed}scf.if {cond} {{\n"
                        f"{new_ctx.emit()}\n"
                        f"{feed}}} else {{\n"
                        f"{else_ctx.emit()}\n"
                        f"{feed}}}"
                    )
                else:
                    new_result = [self.new_ssa() for _ in names]

                    new_vals = [body_bindings[i][0] for i in names]
                    else_vals = [else_bindings[i][0] for i in names]
                    new_type = [body_bindings[i][1] for i in names]

                    new_ctx.exec(
                        f"{new_ctx.feed}scf.yield {', '.join(new_vals)} "
                        f": {', '.join(new_type)}"
                    )
                    else_ctx.exec(
                        f"{else_ctx.feed}scf.yield {', '.join(else_vals)} "
                        f": {', '.join(new_type)}"
                    )

                    self.exec(
                        f"{feed}{', '.join(new_result)} = "
                        f"scf.if {cond} -> ({', '.join(new_type)}) {{\n"
                        f"{new_ctx.emit()}\n"
                        f"{feed}}} else {{\n"
                        f"{else_ctx.emit()}\n"
                        f"{feed}}}"
                    )

                    for i, j, k in zip(
                        names,
                        new_result,
                        new_type,
                        strict=True,
                    ):
                        self.bindings[i] = (j, k)

                return None

            case asm.WhileLoop(condition, loop_body):
                loop_names = self.loop_carried_names(loop_body)

                if not loop_names:
                    condition_ctx = self.subblock()
                    cond = condition_ctx(condition)
                    condition_ctx.exec(f"{condition_ctx.feed}scf.condition({cond})")

                    body_ctx = self.subblock()
                    body_ctx(loop_body)
                    body_ctx.exec(f"{body_ctx.feed}scf.yield")

                    self.exec(
                        f"{feed}scf.while : () -> () {{\n"
                        f"{condition_ctx.emit()}\n"
                        f"{feed}}} do {{\n"
                        f"{body_ctx.emit()}\n"
                        f"{feed}}}"
                    )
                    return None

                before = {i: self.bindings[i] for i in loop_names}
                before_vals = [before[i][0] for i in loop_names]
                before_type = [before[i][1] for i in loop_names]
                loop_args = [self.new_ssa() for _ in loop_names]

                condition_ctx = self.subblock()
                for i, j, k in zip(
                    loop_names,
                    loop_args,
                    before_type,
                    strict=True,
                ):
                    condition_ctx.bindings.bindings[i] = (j, k)
                cond = condition_ctx(condition)
                condition_ctx.exec(
                    f"{condition_ctx.feed}scf.condition({cond}) "
                    f"{', '.join(loop_args)} : {', '.join(before_type)}"
                )

                body_ctx = self.subblock()
                for i, j, k in zip(
                    loop_names,
                    loop_args,
                    before_type,
                    strict=True,
                ):
                    body_ctx.bindings.bindings[i] = (j, k)
                body_ctx(loop_body)

                new_vals = [body_ctx.bindings[i][0] for i in loop_names]
                body_ctx.exec(
                    f"{body_ctx.feed}scf.yield "
                    f"{', '.join(new_vals)} : {', '.join(before_type)}"
                )

                arg_bindings = {
                    arg: (initial, type_)
                    for arg, initial, type_ in zip(
                        loop_args,
                        before_vals,
                        before_type,
                        strict=True,
                    )
                }

                new_ssa = ", ".join(
                    f"{arg} = {initial}" for arg, (initial, _) in arg_bindings.items()
                )
                body_args = ", ".join(
                    f"{arg}: {type_}" for arg, (_, type_) in arg_bindings.items()
                )
                result_types = ", ".join(type_ for _, type_ in arg_bindings.values())

                result = self.new_ssa()
                if len(arg_bindings) == 1:
                    result_name = result
                    result_vals = [result]
                else:
                    result_name = f"{result}:{len(arg_bindings)}"
                    result_vals = [f"{result}#{i}" for i in range(len(arg_bindings))]

                block = self.freshen("bb")
                self.exec(
                    f"{feed}{result_name} = scf.while ({new_ssa}) "
                    f": ({result_types}) -> ({result_types}) {{\n"
                    f"{condition_ctx.emit()}\n"
                    f"{feed}}} do {{\n"
                    f"{body_ctx.feed}^{block}({body_args}):\n"
                    f"{body_ctx.emit()}\n"
                    f"{feed}}}"
                )

                for i, j, (_, k) in zip(
                    loop_names,
                    result_vals[: len(loop_names)],
                    list(arg_bindings.values())[: len(loop_names)],
                    strict=True,
                ):
                    self.bindings[i] = (j, k)

                return None

            case asm.Function(
                asm.Variable(func_name, asm.AssemblyKernelFType(result_type=return_t)),
                args,
                body,
            ):
                ctx_2 = self.subblock()
                statement = []
                for arg in args:
                    match arg:
                        case asm.Variable(name, t):
                            ssa = "%" + name.replace("#", "_")
                            statement.append(f"{ssa}: {mlir_type(t)}")
                            ctx_2.bindings[name] = (ssa, mlir_type(t))
                        case _:
                            raise NotImplementedError(
                                f"Unrecognized argument type: {arg}"
                            )
                if isinstance(return_t, StructFType):
                    statement.append("%_ret: !llvm.ptr")
                    ctx_2.bindings[".ret_ptr"] = ("%_ret", "!llvm.ptr")
                    ret = ""
                else:
                    ret = (
                        ""
                        if return_t == algebra.none_
                        else f" -> {mlir_type(return_t)}"
                    )
                ctx_2(body)
                body_code = ctx_2.emit()
                feed = self.feed

                self.exec(
                    f"{feed}func.func @{func_name}({', '.join(statement)}){ret} "
                    f"attributes {{llvm.emit_c_interface}} {{\n"
                    f"{body_code}\n"
                    f"{feed}}}"
                )
                return None

            case asm.Return(value):
                if value.result_type == algebra.none_:
                    if not isinstance(value, asm.Literal):
                        self(value)
                    self.exec(f"{feed}func.return")
                elif isinstance(value.result_type, StructFType):
                    v = self(value)
                    ptr = self.bindings[".ret_ptr"][0]
                    self.exec(
                        f"{feed}llvm.store {v}, {ptr} "
                        f": {mlir_type(value.result_type)}, !llvm.ptr"
                    )
                    self.exec(f"{feed}func.return")
                else:
                    v = self(value)
                    self.exec(f"{feed}func.return {v} : {mlir_type(value.result_type)}")
                return None

            case asm.Module(funcs):
                for func in funcs:
                    self.bindings[func.name.name] = (
                        f"@{func.name.name}",
                        mlir_type(func.name.result_type),
                    )
                for func in funcs:
                    if not isinstance(func, asm.Function):
                        raise NotImplementedError(
                            f"Unrecognized function type: {type(func)}"
                        )
                    self(func)
                return None

            case node:
                raise NotImplementedError(f"Unrecognized node: {node}")

MLIRForm

Bases: Form

Validates that a FinchAssembly program only uses constructs the MLIR backend can lower, raising a named error for each known limitation instead of failing deep inside code generation.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRForm(Form):
    """
    Validates that a FinchAssembly program only uses constructs the MLIR
    backend can lower, raising a named error for each known limitation
    instead of failing deep inside code generation.
    """

    @classmethod
    def validate_inputs(cls, prgm):
        match prgm:
            case asm.Module(funcs):
                for func in funcs:
                    cls.validate_function(func)
            case _:
                raise TypeError(f"MLIR backend expects an asm.Module, got {type(prgm)}")

    @classmethod
    def validate_function(cls, func):
        match func:
            case asm.Function(
                asm.Variable(
                    func_name, asm.AssemblyKernelFType(result_type=return_type)
                ),
                args,
                body,
            ):
                pass
            case _:
                raise TypeError(f"MLIR backend expects asm.Function, got {func}")
        defined = {}
        for arg in args:
            match arg:
                case asm.Variable(name, t):
                    cls.validate_type(func_name, f"argument {name!r}", t)
                    defined[name] = 0
                case _:
                    raise TypeError(
                        f"MLIR backend expects variable arguments in "
                        f"{func_name!r}, got {arg}"
                    )
        if return_type != algebra.none_:
            cls.validate_type(func_name, "return value", return_type)
        cls.validate_stmt(func_name, body, defined, depth=0)

    @classmethod
    def validate_type(cls, func_name, what, t):
        try:
            mlir_type(t)
        except NotImplementedError as e:
            raise NotImplementedError(
                f"MLIR backend cannot lower the type of {what} "
                f"in function {func_name!r}: {t}"
            ) from e

    @classmethod
    def validate_stmt(cls, func_name, node, defined, depth):
        match node:
            case asm.Block(bodies):
                for body in bodies:
                    cls.validate_stmt(func_name, body, defined, depth)
            case asm.Assign(asm.Variable(name, _), _):
                defined.setdefault(name, depth)
            case asm.ForLoop(asm.Variable(name, _), _, _, body):
                inner = dict(defined)
                inner[name] = depth + 1
                cls.validate_stmt(func_name, body, inner, depth + 1)
            case asm.BufferLoop(_, asm.Variable(name, _), body):
                inner = dict(defined)
                inner[name] = depth + 1
                cls.validate_stmt(func_name, body, inner, depth + 1)
            case asm.If(_, body):
                cls.validate_stmt(func_name, body, defined, depth)
            case asm.IfElse(_, body, else_body):
                cls.validate_stmt(func_name, body, defined, depth)
                cls.validate_stmt(func_name, else_body, defined, depth)
            case asm.WhileLoop(_, body):
                cls.validate_stmt(func_name, body, defined, depth)
            case (
                asm.Unpack(_, _)
                | asm.Repack(_)
                | asm.SetAttr(_, _, _)
                | asm.Store(_, _, _)
                | asm.Resize(_, _)
                | asm.Return(_)
                | asm.Assert(_)
            ):
                pass

            case _:
                raise NotImplementedError(
                    f"MLIR backend does not yet support "
                    f"{type(node).__name__} (in function {func_name!r})"
                )

MLIRKernel

Bases: AssemblyKernel

A class to represent MLIR kernel. It validates and serializes arguments, invokes the compiled function, and reconstructs the result.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRKernel(asm.AssemblyKernel):
    """
    A class to represent MLIR kernel. It validates and serializes arguments,
    invokes the compiled function, and reconstructs the result.

    """

    def __init__(self, engine, func_name, type_):
        super().__init__(type_)
        self.engine = engine
        self.func_name = func_name
        self.ret_type = type_.result_type
        self.argtypes = type_.arg_types

    def __call__(self, *args):
        if len(args) != len(self.argtypes):
            raise ValueError(
                f"Expected {len(self.argtypes)} arguments, got {len(args)}"
            )
        for argtype, arg in zip(self.argtypes, args, strict=False):
            if not fisinstance(arg, argtype):
                raise TypeError(f"Expected argument of type {argtype}, got {type(arg)}")
        serial_args = list(map(serialize_to_mlir, self.argtypes, args))

        packed = []
        for t, sa in zip(self.argtypes, serial_args, strict=False):
            if isinstance(t, BufferFType) and mlir_type(t) != "!llvm.ptr":
                packed.append(ctypes.pointer(ctypes.pointer(sa)))
            else:
                packed.append(ctypes.pointer(sa))

        if self.ret_type == algebra.none_:
            res = None
            self.engine.invoke(self.func_name, *packed)
        elif isinstance(self.ret_type, StructFType):
            res = mlir_ctype(self.ret_type)()
            self.engine.invoke(
                self.func_name, *packed, ctypes.pointer(ctypes.pointer(res))
            )
        else:
            res = mlir_ctype(self.ret_type)()
            self.engine.invoke(self.func_name, *packed, ctypes.pointer(res))

        for arg_type, arg, serial_arg in zip(
            self.argtypes, args, serial_args, strict=False
        ):
            deserialize_from_mlir(arg_type, arg, serial_arg)

        if self.ret_type is type(None):
            return None
        return construct_from_mlir(self.ret_type, res)

MLIRLibrary

Bases: AssemblyLibrary

A class to represent MLIR module.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRLibrary(asm.AssemblyLibrary):
    """
    A class to represent MLIR module.

    """

    def __init__(self, mlir_module, kernels):
        self.mlir_module = mlir_module
        self.kernels = kernels

    def __getattr__(self, name):
        if name in self.kernels:
            return self.kernels[name]
        raise AttributeError(
            f"{type(self).__name__!r} object has no attribute {name!r}"
        )

MLIRLowerer

Bases: Stage

Source code in src/finch/codegen/mlir_codegen/stages.py

class MLIRLowerer(Stage):
    @abstractmethod
    def lower(self, prgm: asm.Module, *, mode: CompilerMode | None = None) -> MLIRCode:
        """
        Lower the given assembly program to textual MLIR.

        """

lower(prgm, *, mode=None) abstractmethod

Lower the given assembly program to textual MLIR.

Source code in src/finch/codegen/mlir_codegen/stages.py

@abstractmethod
def lower(self, prgm: asm.Module, *, mode: CompilerMode | None = None) -> MLIRCode:
    """
    Lower the given assembly program to textual MLIR.

    """

MLIRUnpackableFType

Bases: ABC

Abstract base class for unpackable formats in MLIR. Unpackable formats must also support other functions with unpacked inputs in addition to variable ones.

Source code in src/finch/codegen/mlir_codegen/mlir.py

class MLIRUnpackableFType(ABC):
    """
    Abstract base class for unpackable formats in MLIR. Unpackable formats must also
    support other functions with unpacked inputs in addition to variable ones.
    """

    @abstractmethod
    def mlir_unpack(self, ctx, lhs, rhs):
        """
        Convert a value to an unpacked representation in MLIR. Returns a NamedTuple
        of unpacked variable names, etc. The `lhs` is the variable namespace to
        assign to.
        """
        ...

    @abstractmethod
    def mlir_repack(self, ctx, lhs, rhs):
        """
        Update an object based on an unpacked representation. The `rhs` is the
        unpacked representation to update from, and `lhs` is a variable name referring
        to the original object to update.
        """
        ...

mlir_repack(ctx, lhs, rhs) abstractmethod

Update an object based on an unpacked representation. The rhs is the unpacked representation to update from, and lhs is a variable name referring to the original object to update.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_repack(self, ctx, lhs, rhs):
    """
    Update an object based on an unpacked representation. The `rhs` is the
    unpacked representation to update from, and `lhs` is a variable name referring
    to the original object to update.
    """
    ...

mlir_unpack(ctx, lhs, rhs) abstractmethod

Convert a value to an unpacked representation in MLIR. Returns a NamedTuple of unpacked variable names, etc. The lhs is the variable namespace to assign to.

Source code in src/finch/codegen/mlir_codegen/mlir.py

@abstractmethod
def mlir_unpack(self, ctx, lhs, rhs):
    """
    Convert a value to an unpacked representation in MLIR. Returns a NamedTuple
    of unpacked variable names, etc. The `lhs` is the variable namespace to
    assign to.
    """
    ...

mlir_ctype(s)

Return the ctypes equivalent of an FType or MLIR type string, so Python can pass arguments to and read results from a compiled MLIR kernel.

Source code in src/finch/codegen/mlir_codegen/mlir.py

def mlir_ctype(s: FType | str):
    """
    Return the ctypes equivalent of an FType or MLIR type string, so Python
    can pass arguments to and read results from a compiled MLIR kernel.
    """
    match s:
        case MutableStructFType():
            return ctypes.POINTER(mlir_struct_ctype(s))
        case StructFType():
            return mlir_struct_ctype(s)
        case FType():
            return mlir_ctype(mlir_type(s))

    if s.startswith("memref<") and s.endswith(">"):
        res = mlir_memrefs.get(s)
        if res is None:
            from mlir.runtime import make_nd_memref_descriptor

            t = s[len("memref<") : -1].split("x")
            rank = 0
            while rank < len(t) - 1 and (t[rank] == "?" or t[rank].isdigit()):
                rank += 1
            elem = "x".join(t[rank:])
            res = make_nd_memref_descriptor(rank, mlir_ctype(elem))
            mlir_memrefs[s] = res
        return res

    match s:
        case "!llvm.ptr":
            return ctypes.c_void_p
        case "i1":
            return ctypes.c_bool
        case "i8":
            return ctypes.c_int8
        case "i16":
            return ctypes.c_int16
        case "i32":
            return ctypes.c_int32
        case "index" | "i64":
            return ctypes.c_int64
        case "f32":
            return ctypes.c_float
        case "f64":
            return ctypes.c_double
        case _:
            raise NotImplementedError(f"No ctypes mapping for MLIR type {s}")

mlir_type(t)

Convert an FType into the MLIR type string

Source code in src/finch/codegen/mlir_codegen/mlir.py

def mlir_type(t: FType):
    """
    Convert an FType into the MLIR type string
    """
    match t:
        case asm.AssemblyKernelFType():
            args = [mlir_type(arg) for arg in t.arg_types]
            if isinstance(t.result_type, StructFType):
                args.append("!llvm.ptr")
                ret = "()"
            else:
                ret = (
                    "()" if t.result_type == algebra.none_ else mlir_type(t.result_type)
                )
            return f"({', '.join(args)}) -> {ret}"
        case MLIRArgumentFType():
            return t.mlir_type()
        case algebra.bool_:
            return "i1"
        case algebra.intp | algebra.int_:
            # A Python `int` is a 64-bit integer, which this backend represents
            # as `index` like every other one. Mapping it to `i64` instead made
            # the emitted type depend on whether a literal or a variable came
            # first in a commutative call.
            return "index"
        case algebra.float_:
            return "f64"
        case algebra.ftypes.FDTypeNumpy():
            dt = np_dtype(t)
            match dt.kind:
                case "b":
                    return "i1"
                case "i" | "u":
                    return f"i{dt.itemsize * 8}"
                case "f":
                    return f"f{dt.itemsize * 8}"
                case _:
                    raise NotImplementedError(f"No MLIR type for numpy dtype {dt}")
        case MutableStructFType():
            return "!llvm.ptr"
        case StructFType():
            return mlir_struct_type(t)
        case _:
            raise NotImplementedError(f"No MLIR type mapping for {t}")