cornell-zhang / heterocl

HeteroCL: A Multi-Paradigm Programming Infrastructure for Software-Defined Heterogeneous Computing
https://cornell-zhang.github.io/heterocl/
Apache License 2.0
322 stars 92 forks source link

[API] Add support for user-defined modules #68

Closed seanlatias closed 5 years ago

seanlatias commented 5 years ago

This PR allows users to defined a HeteroCL module. As the name indicates, a HeteroCL module will be synthesized as a hardware module during synthesis. The advantage of using a module is that by reusing the same module, there will be only one instantiation.

To define a HeteroCL module, we first need to define a Python function. Then, we can use a decorator@hcl.module to register the defined Python function. Following is an example.

# define a HeteroCL module, we need to provide the shape of the argument
@hcl.module([(10,)])
def foo(A):
    A[0] = 1

# call the module exactly the same as how you call a Python function
foo(A)

You can also have return values within a module. To return a value, we need to use the API hcl.return_, which is a part of HeteroCL DSL.

@hcl.module([(10,)])
def foo(A):
    hcl.return_(A[0])

with hcl.Stage():
    B[0] = foo(A)

For now, we only allow the module to return an expression. Namely, you cannot return a Tensor. However, as shown in the first example, you can update an input Tensor.

There are more advanced usages, please check the tutorial for more examples and details.

Other updates

Users can now pack/unpack with a specified axis. For more details, please check the test file.