Creating a type graph from a function body with macros

General Usage
,
1

I have a package I am working on where users will write code that looks like (using the package’s @circuit macro):

@circuit Foo begin
    parameters : [
        a => 0.5,
        b => 10
    ]

    circuit : (dut::Foo)(x::SomeType) -> begin
        y = x * dut.a + x * dut.b
        z = y + dut.b

        return z
    end
end

From this code, I want return an expression that does a few things. First, it defines a new struct that’s callable:

@kwdef struct Foo
    a::Float64 = 0.5
    b::Int = 10
end

function (dut::Foo)(x::SomeType)
    y = x * dut.a + x * dut.b
    z = y + dut.b

    return z
end

It also ultimately returns a Module object which stores info about the “circuit” being built. The whole point of this macro is to be syntactic sugar for defining the stuff above and populating the Module object. Most importantly, one of the fields of the Module is a DFG (data flow graph) stored as a MetaDiGraph (from MetaGraphs.jl). Each node in the graph is an operation in the AST of the function body above, and each node has three metadata fields:

  1. The operator as a Symbol which I can get easily enough in the macro
  2. A list of inputs as a vector of Tuple{Symbol, Symbol}. The first element in the tuple is the name of each input (e.g. :y), and the second element is the type (e.g. SomeType).
  3. A list of outputs as a vector of Tuple{Symbol, Symbol}. Same as 2 above.

The issue I am facing is extracting the type information. A constraint placed on the user is that the function signature must contain fully specified type information. Ideally, I am hoping that there is some metaprogramming magic (IRTools.jl maybe?) that allows me to get the return type of each node in the AST (e.g. x * dut.a).

I have my doubts that this is the case, since I understand that macros work on syntax not typed trees. I am happy to leave the types in the Module above be Any then have a function execute at runtime to fill in the type information. I am thinking something like multiple dispatch. As I type this it is becoming apparent that this runtime function might look something like @adjoint from Zygote.jl. Does this approach seem like the best one?

This question is more open-ended that I set out for it to be, but any insight from more experienced package developers would be appreciated. Thanks!

1 Like
2

I don’t understand this part. Where is it specified? The user of what, your macro?

In general, “don’t explicitly rely on type inference” is the standard advice. There are some tricks around it. Eg. [foo(x) for x in 1:10] where foo(x) = x==5 ? "hey" : 2 will start with an array of typeof(foo(1)), then when foo(5) is hit, it will reallocate the whole array as a Vector{Any}

Sounds about right? Can you build the DFG after evaluating x * dut.a for some inputs x and dut?

3

Yeah, I met the user of the macro would need to type (dut::Foo)(x::SomeType). This is just an informally specified constraint via the docs.

Not exactly, since the DFG is built within the macro, and I can’t call the function until after the macro creates it. But I am thinking of basically creating a function in the macro called _extracttypes(dut::Foo, x::SomeType) which is the exact same as the normal function evaluation except every function call is replaced with _extracttypes(m::Module, f, args...) that does what you suggest. It evaluates f(args) and inspects the return type, then modifies the DFG in m accordingly. Since there is only one function that the user will use on the module, generate(m::Module), I should be able to call _extracttypes at the start of that function and that will give me what I need.

Thanks for your help!