[MLIR] How to add attribute and source location

[Kuree]

So I have the following magma code:

import magma as m 
    

class Foo(m.Circuit): 
    io = m.IO(I=m.In(m.Bits[8]), O=m.Out(m.Bits[8]))
    tmp = io.I + 42
    io.O @= tmp

m.compile("build/Foo", Foo, output="mlir")

Here is the generated MLIR:

module attributes {circt.loweringOptions = "locationInfoStyle=none"} {
    hw.module @Foo(%I: i8) -> (O: i8) {
        %0 = hw.constant 42 : i8
        %1 = comb.add %I, %0 : i8
        hw.output %1 : i8
    }
}   

What I want is something like this

module attributes {circt.loweringOptions = "locationInfoStyle=none"} {
    hw.module @Foo(%I: i8) -> (O: i8) attributes {hw.debug.name = ["io.I", "io.O"]} {
        %0 = hw.constant 42 : i8
        %1 = comb.add %I, %0 : i8
        hw.output %1 : i8
    }
}

My questions are:

1. How to add custom attributes to module and wires, e.g. hw.debug.name
2. How to add source locations from python? The generate RTL or MLIR doesn't contain any python location information.

Furthermore, if I do something like this:

import magma as m
    

class Foo(m.Circuit): 
    io = m.IO(I=m.In(m.Bits[8]), O=m.Out(m.Bits[8]))
    a = io.I + 42
    io.O @= a

m.compile("build/Foo", Foo, output="mlir")

The output doesn't change, but if magma can generate a sv.wire with a symbol name of a and annotate the source location, the users can then recognize the symbol in SystemVerilog and hgdb is able to compute proper breakpoint locations. From a implementation perspective this can be done via checking the local variable, so it shouldn't be that challenging. If C-extension is used, it should only introduce about 10% of overhead based on my experiments. From an end-user's perspective I think it's worth a try.

[leonardt]

Do you have an example of how to use a C-extension to get the local variable name? Would you be doing it during the io.O @= a call (i.e. in the imatmul impl)? Or the a = io.I + 42 line? (we can't override the assignment logic to lookup the variable at that time, but we could do it when the output of the add instance is created).

We are also working on a preprocessor that adds in symbol information without having to do it at runtime, so that may be a viable longterm option

[leonardt]

@rsetaluri do you think we can add a generic API for passing through attributes to MLIR for module/wire objects? We could simply expose a dictionary on the objects, or add an "add_attribute" method.

[leonardt]

Both source locations and variable names are inferred if you use m.config.set_debug_mode(True), but I don't remember if we added support for this in MLIR (it was added originally for the CoreIR backend). @rsetaluri did we ever test this with the new backend? The downside to debug mode is that it's slow, which might be fine for these initial experiments, but for larger designs we are looking into the preprocessor approach (unless there's a C-extension way to do this as well?)

[rsetaluri]

@rsetaluri do you think we can add a generic API for passing through attributes to MLIR for module/wire objects? We could simply expose a dictionary on the objects, or add an "add_attribute" method.

Yes, see here:

magma/magma/backend/mlir/mlir.py

Line 145 in 1d7b3b9

attr_dict: Mapping = default_field(dict, init=False)

Now, not all Ops print the attrs because annoyingly, the formatting can be op-specific. So we would just need to add a few liens for the requisite ops (e.g. sv.Wire) to print the attr dict as well.

[Kuree]

You would do it inside the io.O @= a call. Here is some code I use to compute the source location and local variables. Copied to here:

std::optional<std::pair<std::string, uint32_t>> get_fn_ln(uint32_t num_frame_back) {
    using namespace kratos;
    // get caller frame info
    PyFrameObject *frame = PyThreadState_Get()->frame;
    uint32_t i = 0;
    while (frame->f_back && (++i) < num_frame_back) {
        frame = frame->f_back;
    }
    if (frame) {
        uint32_t line_num = PyFrame_GetLineNumber(frame);
        struct py::detail::string_caster<std::string> repr;
        py::handle handle(frame->f_code->co_filename);
        repr.load(handle, true);
        if (repr) {
            // resolve full path
            std::string filename = fs::abspath(repr);
            return std::make_pair(filename, line_num);
        }
    }
    return std::nullopt;
}

py::dict get_frame_local(uint32_t num_frame_back) {
    PyFrameObject *frame = PyThreadState_Get()->frame;
    uint32_t i = 0;
    while (frame->f_back && (++i) < num_frame_back) {
        frame = frame->f_back;
    }
    if (frame) {
        // implementation copied from
        // https://github.com/python/cpython/blob/master/Python/ceval.c
        // PyEval_GetLocals(void)
        PyFrame_FastToLocals(frame);
        auto local = frame->f_locals;
        if (local) {
            py::handle obj(local);
            return obj.cast<py::dict>();
        }
    }
    return py::dict();
}

The function above returns the entire local variables, but you can pick test out the names first to avoid copying between C++ and Python.

[rsetaluri]

Both source locations and variable names are inferred if you use m.config.set_debug_mode(True), but I don't remember if we added support for this in MLIR (it was added originally for the CoreIR backend). @rsetaluri did we ever test this with the new backend? The downside to debug mode is that it's slow, which might be fine for these initial experiments, but for larger designs we are looking into the preprocessor approach (unless there's a C-extension way to do this as well?)

Right now we don't actually include any source info in the generated IR. I can work towards this.

[leonardt]

You would do it inside the io.O @= a call. Here is some code I use to compute the source location and local variables. Copied to here:

std::optional<std::pair<std::string, uint32_t>> get_fn_ln(uint32_t num_frame_back) {
    using namespace kratos;
    // get caller frame info
    PyFrameObject *frame = PyThreadState_Get()->frame;
    uint32_t i = 0;
    while (frame->f_back && (++i) < num_frame_back) {
        frame = frame->f_back;
    }
    if (frame) {
        uint32_t line_num = PyFrame_GetLineNumber(frame);
        struct py::detail::string_caster<std::string> repr;
        py::handle handle(frame->f_code->co_filename);
        repr.load(handle, true);
        if (repr) {
            // resolve full path
            std::string filename = fs::abspath(repr);
            return std::make_pair(filename, line_num);
        }
    }
    return std::nullopt;
}

py::dict get_frame_local(uint32_t num_frame_back) {
    PyFrameObject *frame = PyThreadState_Get()->frame;
    uint32_t i = 0;
    while (frame->f_back && (++i) < num_frame_back) {
        frame = frame->f_back;
    }
    if (frame) {
        // implementation copied from
        // https://github.com/python/cpython/blob/master/Python/ceval.c
        // PyEval_GetLocals(void)
        PyFrame_FastToLocals(frame);
        auto local = frame->f_locals;
        if (local) {
            py::handle obj(local);
            return obj.cast<py::dict>();
        }
    }
    return py::dict();
}

The function above returns the entire local variables, but you can pick test out the names first to avoid copying between C++ and Python.

Cool thanks, I can try this out. Perhaps we should consider making a micro pip package that provides a Python API for this info, might be useful in general to people.

[Kuree]

Perhaps we should consider making a micro pip package that provides a Python API for this info, might be useful in general to people.

@leonardt That sounds like a good idea.

Most of the performance slowsdown with inspect package comes from the fact that Python has to populate all the fields in a frame object. These C-extension functions bypass that.

[Kuree]

Right now we don't actually include any source info in the generated IR. I can work towards this.

@rsetaluri Do you have some ETA on this feature?

[leonardt]

For get_frame_local, is the idea that you get the locals dict then for each variable, you "name" it based on the key in the locals dict?

[leonardt]

Hmm, now that I think about it, during something like @=, we could lookup the value being wired in locals to get it's name (I think this is what you meant by test)

[Kuree]

For get_frame_local, is the idea that you get the locals dict then for each variable, you "name" it based on the key in the locals dict?

I get the whole local variables then iterative them to filter out any integer values or custom type value. Here here is an example:

a = 1
b = net()
d = "foo"
local_vars = get_frame_local()
res = {}
for name, value in local_vars.items():
    if isinstance(value, (int, net)):
        res[name] = value

Then you can compute the symbol table based on this information. This filtering can also be done in C-extension for performance reasons.

Hmm, now that I think about it, during something like @=, we could lookup the value being wired in locals to get it's name (I think this is what you meant by test)

This should also work, depending on what kind of information you wan to store. The issue is if the left hand side is a new variable, it won't show up in the locals.

[mahorowitz]

Any progress on this issue? I have been pushing Keyi to get hgdb to work with Magma code before he graduates, and getting source information into the IR is a necessary first step.

[rsetaluri]

I am working towards a DAC submission this week so quite blocked on that. I will take this up next week if that is ok.

[leonardt]

Prototype C extension almost done, see draft here: #1189

I'll need to check whether this catches all the wiring cases properly, the main issue is making sure the source info is being recorded at the appropriate callsite (e.g. not internally in magma). Also will take some time to setup the package distribution for linux/mac.

I think we can add names using the metaclass access to the class namespace, so we don't need a C API extension for this, but we will have the issue of names defined inside loops being overwritten (so only the last value will get named).

[rsetaluri]

My questions are:

1. How to add custom attributes to module and wires, e.g. hw.debug.name
2. How to add source locations from python? The generate RTL or MLIR doesn't contain any python location information.

@Kuree Can you provide an exact spec for what you need? Right now we can emit loc("myfile.py":10:20) style location attributes. For hw.module ops we can currently emit attributes, we just need to figure out what to emit (and where).

[Kuree]

I need:

1. loc() for every statement that has source-level correspondence.
2. Inserting attribute to the module op.
3. Obtain the Python-level symbol name for each net or reg symbol and then annotate it as an attribute for the AssignOp.

[Kuree]

Any updates? Using latest magma with m.config.set_debug_mode(True) still doesn't produce any location info in the generated MLIR file.

[rsetaluri]

Try setting location_info_style=True in m.compile()

[Kuree]

Try setting location_info_style=True in m.compile()

I got loc(unknown) instead. Only the module declaration has proper location, but assignments do not.

[rsetaluri]

Yes we are still adding more locations for the other ops.

[rsetaluri]

You can use the mlir-source-loc-additions branch which is WIP which adds (possibly incorect...) annotations to most ops.