cornell-zhang / heterocl

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

[IRBuilder] `buffer_at` not built in IR #477

Closed zzzDavid closed 1 year ago

zzzDavid commented 1 year ago

hcl.buffer_at was not connected correctly from the frontend to IR. A test example:

def test_bconv_nhwc_buffer_at():
    bs = 4
    ic, oc = 6, 16
    ih, iw = 8, 8
    kh, kw = 3, 3
    oh, ow = ih - kh + 1, iw - kw + 1

    hcl.init(hcl.UInt(1))
    A = hcl.placeholder((bs, ih, iw, ic))
    F = hcl.placeholder((oc, kh, kw, ic))

    def conv(A, F):
        rc = hcl.reduce_axis(0, ic)
        rh = hcl.reduce_axis(0, kh)
        rw = hcl.reduce_axis(0, kw)
        L = ic * kh * kw
        B = hcl.compute(
            (bs, oh, ow, oc),
            lambda n, h, w, c: L
            - (
                hcl.sum(
                    A[n, h + rh, w + rw, rc] ^ F[c, rh, rw, rc],
                    axis=[rh, rw, rc],
                    dtype=hcl.Int(32),
                )
                << 1
            ),
            name="B",
            dtype=hcl.Int(32),
        )
        return B

    s = hcl.create_schedule([A, F], conv)
    B = conv.B
    buf = s.buffer_at(B, s[B], B.axis[2])
    LB = s.reuse_at(A, s[B], B.axis[1])
    WB = s.reuse_at(LB, s[B], B.axis[2])
    f = hcl.build(s)

    np_A = np.random.randint(0, 2, size=(bs, ih, iw, ic))
    np_B = np.random.randint(0, 2, size=(oc, kh, kw, ic))
    np_C = np.zeros((bs, oh, ow, oc), dtype="int")

    for n in range(0, bs):
        for y in range(0, oh):
            for x in range(0, ow):
                for c in range(0, oc):
                    for rc in range(0, ic):
                        for rh in range(0, kh):
                            for rw in range(0, kw):
                                np_C[n][y][x][c] += 1 - 2 * (
                                    np_A[n][y + rh][x + rw][rc] ^ np_B[c][rh][rw][rc]
                                )

    hcl_A = hcl.asarray(np_A, dtype=hcl.UInt(1))
    hcl_B = hcl.asarray(np_B, dtype=hcl.UInt(1))
    hcl_C = hcl.asarray(np_C, dtype=hcl.Int(32))

    f(hcl_A, hcl_B, hcl_C)

    assert np.array_equal(np_C, hcl_C.asnumpy())

Error message:

.BufferAt is not an attribute in hcl_d