Xilinx / mlir-air

MIT License
76 stars 26 forks source link

Implement python decorator to abstract module InsertionPoint #600

Closed hunhoffe closed 4 months ago

hunhoffe commented 4 months ago

This PR introduces a new python decorator, module_builder, which may be used to help create a module of code using the AIR python bindings. The purpose of this decorator is to help simplify the air python bindings for the benefit of the user.

As an example, I changed programming_examples/shim_dma_2d/shim_dma_2d.py to use this new decorator, as well as any tests that seem to create a module using python bindings.

It turns out there is already an @module decorator which is builtin with the MLIR Python bindings. However, I couldn't find much documentation on it. When I finally got it to work (and I'm not 100% sure this is how it's meant to be used, but it's my best attempt), the end result seemed a bit unintuitive so I think it's worth defining and using a new decorator.

Using new @module_builder decorator

@module_builder
def build_module():
    memrefTyInOut = MemRefType.get(IMAGE_SIZE, T.i32())

    # We will send the image worth of data in and out
    @FuncOp.from_py_func(memrefTyInOut, memrefTyInOut)
    def copy(arg0, arg1):
        ...

if __name__ == "__main__":
    module = build_module()
    print(module)

Using builtin @module decorator

def build_module(loc):
    @module(loc=loc)
    def my_module():
        memrefTyInOut = MemRefType.get(IMAGE_SIZE, T.i32())

        # We will send the image worth of data in and out
        @FuncOp.from_py_func(memrefTyInOut, memrefTyInOut)
        def copy(arg0, arg1):
            ...
    return my_module

if __name__ == "__main__":
    with Context() as ctx, Location.unknown() as loc:
        my_module = build_module(loc)
    print(my_module)

Two things I don't like about using the @module decorator are that: