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:
loc is not a positional argument, so I need to be able to say loc=<some_var> at the time of function definition. To ensure there is a loc variable in scope, I had to keep an extra layer of function definitions (build_module() AND my_module(), instead of just one)
At the end of the build_module() function, you must return my_module (no explicit invocation) instead of return my_module(). This seems unintuitive.
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
decoratorUsing builtin
@module
decoratorTwo things I don't like about using the
@module
decorator are that:loc
is not a positional argument, so I need to be able to sayloc=<some_var>
at the time of function definition. To ensure there is aloc
variable in scope, I had to keep an extra layer of function definitions (build_module()
ANDmy_module()
, instead of just one)build_module()
function, you mustreturn my_module
(no explicit invocation) instead ofreturn my_module()
. This seems unintuitive.