Example #1
0
def init_array(path):
    n = N.get()

    for i in range(n):
        for j in range(n):
            path[i, j] = datatype(i * j % 7 + 1)
            if (i + j) % 13 == 0 or (i + j) % 7 == 0 or (i + j) % 11 == 0:
                path[i, j] = datatype(999)


@dace.program(datatype[N, N])
def floyd_warshall(path):
    @dace.mapscope
    def k_map(k: _[0:N]):
        @dace.map
        def ij_map(i: _[0:N], j: _[0:N]):
            ik_dist << path[i, k]
            kj_dist << path[k, j]
            out >> path(1, lambda x, y: min(x, y))[i, j]
            out = ik_dist + kj_dist


if __name__ == '__main__':
    if polybench:
        polybench.main(sizes, args, [(0, 'path')], init_array, floyd_warshall)
    else:
        [k.set(v) for k, v in sizes[2].items()]
        init_array(*args)
        floyd_warshall(*args)
Example #2
0

@dace.program(datatype[M, N], datatype[N], datatype[N])
def atax(A, x, y):
    tmp = dace.define_local([M], dtype=datatype)

    @dace.map
    def reset_y(i: _[0:N]):
        out >> y[i]
        out = 0.0

    for i in range(M):

        @dace.map
        def compute_tmp(j: _[0:N]):
            inA << A[i, j]
            inx << x[j]
            out >> tmp(1, lambda a, b: a + b, 0)[i]
            out = inA * inx

        @dace.map
        def compute_y(j: _[0:N]):
            inA << A[i, j]
            intmp << tmp[i]
            outy >> y(1, lambda a, b: a + b)[j]
            outy = inA * intmp


if __name__ == '__main__':
    polybench.main(sizes, args, [(2, 'y')], init_array, atax)
Example #3
0
                        center << table[i, j]
                        swest << table[i + 1, j - 1]
                        out >> table[i, j]
                        out = max(center, swest)

            for k in range(i + 1, j, 1):

                @dace.tasklet
                def set_table_5():
                    center << table[i, j]
                    k_center << table[i, k]
                    k_south << table[k + 1, j]
                    out >> table[i, j]
                    out = max(center, k_center + k_south)


def print_result(filename, *args):
    with open(filename, 'w') as fp:
        fp.write("==BEGIN DUMP_ARRAYS==\n")
        fp.write("begin dump: %s\n" % 'table')
        for i in range(0, N.get()):
            for j in range(i, N.get()):
                fp.write("{} ".format(args[1][i, j]))
            fp.write("\n")
        fp.write("\nend   dump: %s\n" % 'table')
        fp.write("==END   DUMP_ARRAYS==\n")


if __name__ == '__main__':
    polybench.main(sizes, args, print_result, init_array, nussinov)
Example #4
0
    def mult_tmp(i: _[0:NI], j: _[0:NJ], k: _[0:NK]):
        in_a << A[i, k]
        in_b << B[k, j]
        in_alpha << alpha
        out >> tmp(1, lambda x, y: x + y)[i, j]
        out = in_alpha * in_a * in_b

    @dace.map
    def mult_d(i: _[0:NI], j: _[0:NL]):
        inp << D[i, j]
        in_beta << beta
        out >> D[i, j]

        out = inp * in_beta

    @dace.map
    def comp_d(i: _[0:NI], j: _[0:NL], k: _[0:NJ]):
        in_a << tmp[i, k]
        in_b << C[k, j]
        out >> D(1, lambda x, y: x + y)[i, j]
        out = in_a * in_b


if __name__ == '__main__':
    if polybench:
        polybench.main(sizes, args, [(3, 'D')], init_array, k2mm)
    else:
        [k.set(v) for k, v in sizes[2].items()]
        init_array(*args)
        k2mm(*args)
Example #5
0
File: trmm.py Project: zurvar/dace
def trmm(A, B, alpha):
    @dace.mapscope
    def compute(j: _[0:N]):
        @dace.mapscope
        def computecol(i: _[0:M]):
            tmp = dace.define_local_scalar(datatype)

            @dace.tasklet
            def reset_tmp():
                out >> tmp
                out = 0

            @dace.map
            def compute_elem(k: _[i + 1:M]):
                ia << A[k, i]
                ib << B[k, j]
                ob >> tmp(1, lambda a, b: a + b)
                ob = ia * ib

            @dace.tasklet
            def mult():
                ib << B[i, j]
                ialpha << alpha
                itmp << tmp
                ob >> B[i, j]
                ob = ialpha * (ib + itmp)


if __name__ == '__main__':
    polybench.main(sizes, args, outputs, init_array, trmm)
Example #6
0
            C[i, j] = datatype((i * j + 3) % n) / m


@dace.program(datatype[N, N], datatype[N, M], datatype[N, M], datatype[1],
              datatype[1])
def syr2k(C, A, B, alpha, beta):
    @dace.map
    def mult_c_rows(i: _[0:N]):
        @dace.map
        def mult_c_cols(j: _[0:i + 1]):
            ic << C[i, j]
            ib << beta
            oc >> C[i, j]
            oc = ic * ib

    @dace.map
    def compute(i: _[0:N], k: _[0:M]):
        @dace.map
        def compute_elem(j: _[0:i + 1]):
            ialpha << alpha
            ia << A[i, k]
            iat << A[j, k]
            ib << B[i, k]
            ibt << B[j, k]
            oc >> C(1, lambda a, b: a + b)[i, j]
            oc = ialpha * iat * ib + ialpha * ibt * ia


if __name__ == '__main__':
    polybench.main(sizes, args, outputs, init_array, syr2k)
Example #7
0
            L[i, j] = datatype(0)


@dace.program(datatype[N, N], datatype[N], datatype[N])
def trisolv(L, x, b):
    for i in range(0, N, 1):

        @dace.tasklet
        def init_x():
            in_b << b[i]
            out >> x[i]
            out = in_b

        @dace.map
        def set_x(j: _[0:i]):
            in_L << L[i, j]
            in_x << x[j]
            out >> x(1, lambda x, y: x + y)[i]
            out = -in_L * in_x

        @dace.tasklet
        def div():
            in_x << x[i]
            in_L << L[i, i]
            out >> x[i]
            out = in_x / in_L


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 'x')], init_array, trisolv)
Example #8
0
        @dace.map
        def update_ey(i: _[1:NX], j: _[0:NY]):
            eyin << ey[i, j]
            hz1 << hz[i, j]
            hz2 << hz[i - 1, j]
            eyout >> ey[i, j]
            eyout = eyin - datatype(0.5) * (hz1 - hz2)

        @dace.map
        def update_ex(i: _[0:NX], j: _[1:NY]):
            exin << ex[i, j]
            hz1 << hz[i, j]
            hz2 << hz[i, j - 1]
            exout >> ex[i, j]
            exout = exin - datatype(0.5) * (hz1 - hz2)

        @dace.map
        def update_hz(i: _[0:NX - 1], j: _[0:NY - 1]):
            hzin << hz[i, j]
            ex1 << ex[i, j + 1]
            ex2 << ex[i, j]
            ey1 << ey[i + 1, j]
            ey2 << ey[i, j]
            hzout >> hz[i, j]
            hzout = hzin - datatype(0.7) * (ex1 - ex2 + ey1 - ey2)


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'ex'), (1, 'ey'), (2, 'hz')], init_array,
                   fdtd2d)
Example #9
0
                out >> A[i, j]
                out = ij_in / jj_in

        @dace.map
        def k_loop2(k: _[0:i]):
            k_in << A[i, k]
            out >> A(1, lambda x, y: x + y)[i, i]
            out = -k_in * k_in

        @dace.tasklet
        def sqrt():
            inp << A[i, i]
            out >> A[i, i]
            out = math.sqrt(inp)


def print_result(filename, *args):
    with open(filename, 'w') as fp:
        fp.write("==BEGIN DUMP_ARRAYS==\n")
        fp.write("begin dump: %s\n" % 'A')
        for i in range(0, N.get()):
            for j in range(0, i + 1):
                fp.write("{:.7f} ".format(args[0][i, j]))
            fp.write("\n")
        fp.write("\nend   dump: %s\n" % 'A')
        fp.write("==END   DUMP_ARRAYS==\n")


if __name__ == '__main__':
    polybench.main(sizes, args, print_result, init_array, cholesky)
Example #10
0
                a

        @dace.map
        def a(i: _[1:N - 1], j: _[1:N - 1], k: _[1:N - 1]):
            a11 << B[i + 1, j, k]
            a12 << B[i - 1, j, k]
            a21 << B[i, j + 1, k]
            a22 << B[i, j - 1, k]
            a31 << B[i, j, k + 1]
            a32 << B[i, j, k - 1]
            a << B[i, j, k]
            b >> A[i, j, k]

            b = 0.125 * (a11 - datatype(2.0) * a + a12) +\
                0.125 * (a21 - datatype(2.0) * a + a22) +\
                0.125 * (a31 - datatype(2.0) * a + a32) +\
                a


def init_array(A, B):  #, N, tsteps):
    n = N.get()
    for i in range(n):
        for j in range(n):
            for k in range(n):
                A[i, j, k] = datatype((i + j + (n - k)) * 10) / n
                B[i, j, k] = datatype((i + j + (n - k)) * 10) / n


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'A')], init_array, heat3d)
Example #11
0
def lu(A):
    for i in range(0, N, 1):
        for j in range(0, i, 1):

            @dace.map
            def k_loop1(k: _[0:j]):
                i_in << A[i, k]
                j_in << A[k, j]
                out >> A(1, lambda x, y: x + y)[i, j]
                out = -i_in * j_in

            @dace.tasklet
            def div():
                ij_in << A[i, j]
                jj_in << A[j, j]
                out >> A[i, j]
                out = ij_in / jj_in

        for j in range(i, N, 1):

            @dace.map
            def k_loop2(k: _[0:i]):
                i_in << A[i, k]
                j_in << A[k, j]
                out >> A(1, lambda x, y: x + y)[i, j]
                out = -i_in * j_in


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'A')], init_array, lu)
Example #12
0
            inA << A[r, q, s]
            inC4 << C4[s, p]
            s >> sum(1, lambda a, b: a + b, 0)[r, q, p]
            s = inA * inC4

        @dace.map
        def compute_A(p: _[0:NP]):
            insum << sum[r, q, p]
            out >> A[r, q, p]
            out = insum

        # @dace.program(datatype[NP], datatype[NP, NP], datatype[NP],
        #               datatype[NP])
        # def internal(nA, nC4, nSum, nAout):
        #     @dace.map
        #     def compute_sum(p: _[0:NP], s: _[0:NP]):
        #         inA << nA[s]
        #         inC4 << nC4[s, p]
        #         s >> nSum(1, lambda a, b: a + b, 0)[p]
        #         s = inA * inC4

        #     @dace.map
        #     def compute_A(p: _[0:NP]):
        #         insum << nSum[p]
        #         out >> nAout[p]
        #         out = insum


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'A')], init_array, doitgen)
Example #13
0
        def set_y3():
            in_w << w
            out >> y[i]
            out = in_w

    for i in range(N - 1, -1, -1):

        @dace.tasklet
        def init_w4():
            in_y << y[i]
            out >> w
            out = in_y

        @dace.map
        def set_w4(j: _[i + 1:N]):
            in_A << A[i, j]
            in_x << x[j]
            out >> w(1, lambda x, y: x + y)
            out = -in_A * in_x

        @dace.tasklet
        def set_x4():
            in_w << w
            in_A << A[i, i]
            out >> x[i]
            out = in_w / in_A


if __name__ == '__main__':
    polybench.main(sizes, args, [(2, 'x')], init_array, ludcmp)
Example #14
0
File: adi.py Project: orausch/dace
            u0i >> u(-1)[ir, 0]
            pi0 >> p(-1)[ir, 0]
            qi0 >> q(-1)[ir, 0]

            pout >> p[ir, :]
            qout >> q[ir, :]
            uout >> u[ir, :]

            uNi >> u(-1)[ir, N - 1]

            u0i = datatype(1.0)
            pi0 = datatype(0.0)
            qi0 = datatype(1.0)

            for j in range(1, N - 1):
                pout[j] = -in_f / (in_d * pin[j - 1] + in_e)
                qout[j] = (-in_a * vin_prev[j] +
                           (datatype(1.0) + datatype(2.0) * in_a) * vin[j] -
                           in_c * vin_next[j] -
                           in_d * qin[j - 1]) / (in_d * pin[j - 1] + in_e)

            uNi = datatype(1.0)

            for j in range(N - 2, 0, -1):
                uout[j] = pin[j] * uin[j + 1] + qin[j]


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'u')], init_array, adi)
Example #15
0
    for t in range(tsteps):
        for i in range(1, N - 1):
            for j in range(1, N - 1):

                @dace.tasklet
                def a():
                    a1 << A[i - 1, j - 1]
                    a2 << A[i - 1, j]
                    a3 << A[i - 1, j + 1]
                    a4 << A[i, j - 1]
                    a5 << A[i, j]
                    a6 << A[i, j + 1]
                    a7 << A[i + 1, j - 1]
                    a8 << A[i + 1, j]
                    a9 << A[i + 1, j + 1]
                    out >> A[i, j]

                    out = (a1 + a2 + a3 + a4 + a5 + a6 + a7 + a8 +
                           a9) / datatype(9.0)


def init_array(A, tsteps):
    n = N.get()
    for i in range(n):
        for j in range(n):
            A[i, j] = datatype(i * (j + 2) + 2) / n


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'A')], init_array, seidel2d)
Example #16
0
        for i in range(W - 1, -1, -1):

            @dace.tasklet
            def comp_y22():
                in_img << imgOut[i, j]
                in_tp1 << tp1
                in_tp2 << tp2
                in_yp1 << yp1
                in_yp2 << yp2
                out_y2 >> y2[i, j]
                out_tp1 >> tp1
                out_tp2 >> tp2
                out_yp1 >> yp1
                out_yp2 >> yp2
                out_y2 = a7 * in_tp1 + a8 * in_tp2 + b1 * in_yp1 + b2 * in_yp2
                out_tp2 = in_tp1
                out_tp1 = in_img
                out_yp2 = in_yp1
                out_yp1 = out_y2

    @dace.map
    def comp_iout2(i: _[0:W], j: _[0:H]):
        in_y1 << y1[i, j]
        in_y2 << y2[i, j]
        out_img >> imgOut[i, j]
        out_img = c1 * (in_y1 + in_y2)


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 'imgOut')], init_array, deriche)
Example #17
0
        oud >> data[i, j]
        oud = (ind - m) / (math.sqrt(datatype(N)) * sd)

    @dace.map
    def comp_corr_diag(i: _[0:M]):
        corrout >> corr[i, i]
        corrout = 1.0

    @dace.mapscope
    def comp_corr_row(i: _[0:M - 1]):
        @dace.mapscope
        def comp_corr_col(j: _[i + 1:M]):
            @dace.map
            def comp_cov_k(k: _[0:N]):
                indi << data[k, i]
                indj << data[k, j]
                cov_ij >> corr(1, lambda x, y: x + y, 0)[i, j]
                cov_ij = (indi * indj)

    @dace.mapscope
    def symmetrize(i: _[0:M - 1]):
        @dace.map
        def symmetrize_col(j: _[i + 1:M]):
            corrin << corr[i, j]
            corrout >> corr[j, i]
            corrout = corrin


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 'corr')], init_array, correlation)
Example #18
0
            in_R << R[k, k]
            out_Q >> Q[i, k]
            out_Q = in_A / in_R

        @dace.mapscope
        def set_rna(j: _[k + 1:N]):
            # for j in range(k+1, N, 1):

            @dace.tasklet
            def init_r():
                out_R >> R[k, j]
                out_R = datatype(0)

            @dace.map
            def set_r(i: _[0:M]):
                in_A << A[i, j]
                in_Q << Q[i, k]
                out_R >> R(1, lambda x, y: x + y)[k, j]
                out_R = in_A * in_Q

            @dace.map
            def set_a(i: _[0:M]):
                in_R << R[k, j]
                in_Q << Q[i, k]
                out_A >> A(1, lambda x, y: x + y)[i, j]
                out_A = -in_R * in_Q


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 'R'), (2, 'Q')], init_array, gramschmidt)
Example #19
0
        for j in range(n):
            A[i, j] = datatype((i * j + 1) % n) / n
            B[i, j] = datatype((i * j + 2) % n) / n


@dace.program(datatype[N, N], datatype[N, N], datatype[N], datatype[N], datatype[N], datatype[1], datatype[1])
def gesummv(A, B, tmp, x, y, alpha, beta):
    @dace.map
    def compute_ty(i: _[0:N], j: _[0:N]):
        ia << A[i, j]
        ib << B[i, j]
        ix << x[j]
        ot >> tmp(1, lambda a, b: a + b, 0)[i]
        oy >> y(1, lambda a, b: a + b, 0)[i]

        ot = ia * ix
        oy = ib * ix

    @dace.map
    def update_y(i: _[0:N]):
        iy << y[i]
        ialpha << alpha
        ibeta << beta
        it << tmp[i]
        oy >> y[i]
        oy = ialpha * it + ibeta * iy


if __name__ == '__main__':
    polybench.main(sizes, args, outputs, init_array, gesummv)
Example #20
0
        def set_alpha():
            in_r << r[k]
            in_sum << sum
            in_b << beta
            out_a >> alpha
            out_a = -(in_r + in_sum) / in_b

        @dace.map
        def set_zeta(i: _[0:k]):
            in_y << y[i]
            kin_y << y[k - i - 1]
            in_a << alpha
            out_z >> z[i]
            out_z = in_y + in_a * kin_y

        @dace.map
        def set_y1(i: _[0:k]):
            in_z << z[i]
            out_y >> y[i]
            out_y = in_z

        @dace.tasklet
        def set_y2():
            in_a << alpha
            out_y >> y[k]
            out_y = in_a


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 'y')], init_array, durbin)
Example #21
0
            B[i, j] = datatype(i * (j + 2) % nj) / nj


@dace.program(datatype[NI, NJ], datatype[NI, NK], datatype[NK, NJ],
              datatype[1], datatype[1])
def gemm(C, A, B, alpha, beta):
    @dace.map
    def mult_c(i: _[0:NI], j: _[0:NJ]):
        inp << C[i, j]
        in_beta << beta
        out >> C[i, j]

        out = inp * in_beta

    @dace.map
    def comp(i: _[0:NI], k: _[0:NK], j: _[0:NJ]):
        in_a << A[i, k]
        in_b << B[k, j]
        in_alpha << alpha
        out >> C(1, lambda x, y: x + y)[i, j]
        out = in_alpha * in_a * in_b


if __name__ == '__main__':
    if polybench:
        polybench.main(sizes, args, [(0, 'C')], init_array, gemm)
    else:
        [k.set(v) for k, v in sizes[2].items()]
        init_array(*args)
        gemm(*args)
Example #22
0
File: 3mm.py Project: cpenny42/dace
    @dace.map
    def mult_E(i: _[0:NI], j: _[0:NJ], k: _[0:NK]):
        in_a << A[i, k]
        in_b << B[k, j]
        out >> E(1, lambda x, y: x + y, 0)[i, j]
        out = in_a * in_b

    @dace.map
    def mult_F(i: _[0:NJ], j: _[0:NL], k: _[0:NM]):
        in_a << C[i, k]
        in_b << D[k, j]
        out >> F(1, lambda x, y: x + y, 0)[i, j]
        out = in_a * in_b

    @dace.map
    def mult_G(i: _[0:NI], j: _[0:NL], k: _[0:NJ]):
        in_a << E[i, k]
        in_b << F[k, j]
        out >> G(1, lambda x, y: x + y, 0)[i, j]
        out = in_a * in_b


if __name__ == '__main__':
    if polybench:
        polybench.main(sizes, args, [(4, 'G')], init_array, k3mm)
    else:
        [k.set(v) for k, v in sizes[2].items()]
        init_array(*args)
        k3mm(*args)
Example #23
0
        oa = ia + iu1 * iv1 + iu2 * iv2

    @dace.map
    def comp_y(i: _[0:N], j: _[0:N]):
        ib << beta
        ia << A[j, i]
        iy << y[j]
        ox >> x(1, lambda a, b: a + b)[i]

        ox = ib * ia * iy

    @dace.map
    def comp_xz(i: _[0:N]):
        ix << x[i]
        iz << z[i]
        ox >> x[i]
        ox = ix + iz

    @dace.map
    def comp_w(i: _[0:N], j: _[0:N]):
        ialpha << alpha
        ia << A[i, j]
        ix << x[j]
        ow >> w(1, lambda a, b: a + b)[i]
        ow = ialpha * ia * ix


if __name__ == '__main__':
    polybench.main(sizes, args, outputs, init_array, gemver)
Example #24
0
def init_array(x1, x2, y_1, y_2, A):
    n = N.get()

    for i in range(n):
        x1[i] = datatype(i % n) / n
        x2[i] = datatype((i + 1) % n) / n
        y_1[i] = datatype((i + 3) % n) / n
        y_2[i] = datatype((i + 4) % n) / n
        for j in range(n):
            A[i, j] = datatype(i * j % n) / n


@dace.program(datatype[N], datatype[N], datatype[N], datatype[N], datatype[N,
                                                                           N])
def mvt(x1, x2, y_1, y_2, A):
    @dace.map
    def compute(i: _[0:N], j: _[0:N]):
        in_A1 << A[i, j]
        in_A2 << A[j, i]
        iny1 << y_1[j]
        iny2 << y_2[j]
        out1 >> x1(1, lambda a, b: a + b)[i]
        out2 >> x2(1, lambda a, b: a + b)[i]
        out1 = in_A1 * iny1
        out2 = in_A2 * iny2


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'x1'), (1, 'x2')], init_array, mvt)
Example #25
0
File: bicg.py Project: tbennun/dace
    for i in range(m):
        p[i] = datatype(i % m) / m
    for i in range(n):
        r[i] = datatype(i % n) / n
        for j in range(m):
            A[i, j] = datatype(i * (j + 1) % n) / n


@dace.program(datatype[N, M], datatype[M], datatype[N], datatype[M],
              datatype[N])
def bicg(A, s, q, p, r):
    @dace.map
    def reset_s(i: _[0:M]):
        out >> s[i]
        out = 0.0

    @dace.map
    def compute(i: _[0:N], j: _[0:M]):
        inA << A[i, j]
        inr << r[i]
        inp << p[j]
        outs >> s(1, lambda a, b: a + b)[j]
        outq >> q(1, lambda a, b: a + b)[i]
        outs = inr * inA
        outq = inA * inp


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 's'), (2, 'q')], init_array, bicg)
Example #26
0
    @dace.map
    def sub_mean(i: _[0:N], j: _[0:M]):
        ind << data[i, j]
        m << mean[j]
        oud >> data[i, j]
        oud = ind - m

    @dace.mapscope
    def comp_cov_row(i: _[0:M]):
        @dace.mapscope
        def comp_cov_col(j: _[i:M]):
            @dace.map
            def comp_cov_k(k: _[0:N]):
                indi << data[k, i]
                indj << data[k, j]
                cov_ij >> cov(1, lambda x, y: x + y, 0)[i, j]
                cov_ij = (indi * indj)

    @dace.mapscope
    def symmetrize(i: _[0:M]):
        @dace.map
        def symmetrize_col(j: _[i:M]):
            cov_ij << cov[i, j]
            covout >> cov(2)[:, :]
            covout[i, j] = cov_ij / (N - 1)
            covout[j, i] = cov_ij / (N - 1)


if __name__ == '__main__':
    polybench.main(sizes, args, [(1, 'cov')], init_array, covariance)
Example #27
0
@dace.program(datatype[N], datatype[N])  #, dace.int32, dace.int32)
def jacobi1d(A, B):  #, N, tsteps):
    for t in range(tsteps):

        @dace.map
        def a(i: _[1:N - 1]):
            a1 << A[i - 1]
            a2 << A[i]
            a3 << A[i + 1]
            b >> B[i]
            b = 0.33333 * (a1 + a2 + a3)

        @dace.map
        def b(i: _[1:N - 1]):
            a1 << B[i - 1]
            a2 << B[i]
            a3 << B[i + 1]
            b >> A[i]
            b = 0.33333 * (a1 + a2 + a3)


def init_array(A, B):  #, N, tsteps):
    n = N.get()
    for i in range(n):
        A[i] = datatype(i + 2) / n
        B[i] = datatype(i + 3) / n


if __name__ == '__main__':
    polybench.main(sizes, args, [(0, 'A')], init_array, jacobi1d)