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][Refactor] Fix all tests and complete all compute primitives #81

Closed seanlatias closed 5 years ago

seanlatias commented 5 years ago

In this PR, the following things are done.

  1. Complete basic implementation of compute_at. More complex cases will be handled in the near future. For example, merging two loop nests with different bounds.

Now after the user uses compute_at, an attach_scope attribute statements will be inserted. Then we can continue applying other scheduling functions (e.g., reorder, split, fuse). You can check the test cases for more examples. One thing to note is that we need to perform corresponding substitutions for the producer stages that are computed at the current (consumer) stage. Thus, we support both directions for applying a sequence of optimizations. For example,

A = hcl.placeholer((10, 10))
B = hcl.compute(A.shape, lambda x1, y1: A[x1, y1] + 1)
C = hcl.compute(B.shape, lambda x2, y2: B[x2, y2] + 1)

# first direction: compute_at then reorder
s = hcl.create_schedule([A, C])
s[B].compute_at(s[C], C.axis[1])
s[C].reorder(C.axis[1], C.axis[0])

# second direction: reorder then compute_at
# note that we need to reorder both producer and consumer
# otherwise, it does not make sense
s = hcl.create_schedule([A, C])
s[B].reorder(B.axis[1], B.axis[0])
s[C].reorder(C.axis[1], C.axis[0])
# note that we also need to compute at a different axis
s[B].compute_at(s[C], C.axis[0])

# both scheduling ends up with the following IR
for (y2, 0, 10)
  for (x2, 0, 10)
    allocate B[1]
    B[0] = A[x2, y2] + 1
    C[x2, y2] = B[0] + 1

The attribute statements will be replaced during lowering. More specifically, it happens at the stage of ScheduleOp. When MakePipeline, the body of the producer will be placed in the body of the consumer.

  1. Fix all test cases. Some of the original test cases do not make sense and thus are updated.

  2. Remove all scan related functions and files.