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][Pass][Schedule][Backend] Add array partitioning #93

Closed seanlatias closed 5 years ago

seanlatias commented 5 years ago

In this PR, the following is done.

  1. Add a new IR node Partition for partitioning an array
  2. Add a memory primitive partition
  3. Extend the Allocate IR node so that it can have multiple attributes, such as Partition
  4. Add an IR pass that simplifies the partition operation
  5. Add HLS codegen for array partitioning
  6. For multi-dimensional arrays, HLS codegen will unflatten 1D index back to multi-dimensional
seanlatias commented 5 years ago

Detailed Usage

The API takes in four arguments.

s.partition(Tensor, partition_type, dim, factor)

The first one is the Tensor to be partitioned. It can be a placeholder, a compute result, or an intermediate variable. The second one is the type. We provide three choices, which is the same as Vivado HLS. The three types are Complete, Block, and Cycle. By default, it is Complete. We can use it by using Partition.{type}. The third one is the dimension. Similar to Vivado HLS, if dim=0, we partition all dimensions. Otherwise, only the specified dimension will be partitioned. Note that this argument can only be applied to tensors with more-than-one dimension. Finally, the fourth argument is the partition factor. If Complete is specified, this field would be ignored. Following is an example.

A = hcl.placeholder((10, 10))
B = hcl.compute(A.shape, lambda x, y: A[x, y])
s = hcl.create_schedule([A, B])
s.partition(A, Partition.Block, dim=1, factor=5)
s.partition(B, Partition.Cycle, dim=0, factor=2)