Method conv1d

tests/python/relax/test_transform_legalize_ops_nn.py:49–65

(A: T.Buffer((T.int64(2), T.int64(128), T.int64(28)), "float32"), B: T.Buffer((T.int64(64), T.int64(16), T.int64(3)), "float32"), group_conv1d_ncw: T.Buffer((T.int64(2), T.int64(64), T.int64(13)), "float32"))

        @T.prim_func(private=True, s_tir=True)
        def conv1d(A: T.Buffer((T.int64(2), T.int64(128), T.int64(28)), "float32"), B: T.Buffer((T.int64(64), T.int64(16), T.int64(3)), "float32"), group_conv1d_ncw: T.Buffer((T.int64(2), T.int64(64), T.int64(13)), "float32")):
            T.func_attr({"tirx.noalias": True})
            pad_temp = T.sblock_alloc_buffer((T.int64(2), T.int64(128), T.int64(30)))
            for i0, i1, i2 in T.grid(T.int64(2), T.int64(128), T.int64(30)):
                with T.sblock("pad_temp"):
                    v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
                    T.reads(A[v_i0, v_i1, v_i2 - T.int64(1)])
                    T.writes(pad_temp[v_i0, v_i1, v_i2])
                    pad_temp[v_i0, v_i1, v_i2] = T.if_then_else(T.int64(1) <= v_i2 and v_i2 < T.int64(29), A[v_i0, v_i1, v_i2 - T.int64(1)], T.float32(0))
            for nn, ff, yy, rc, ry in T.grid(T.int64(2), T.int64(64), T.int64(13), T.int64(16), T.int64(3)):
                with T.sblock("group_conv1d_ncw"):
                    v_nn, v_ff, v_yy, v_rc, v_ry = T.axis.remap("SSSRR", [nn, ff, yy, rc, ry])
                    T.reads(pad_temp[v_nn, v_ff // T.int64(8) * T.int64(16) + v_rc, v_yy * T.int64(2) + v_ry * T.int64(2)], B[v_ff, v_rc, v_ry])
                    T.writes(group_conv1d_ncw[v_nn, v_ff, v_yy])
                    with T.init():
                        group_conv1d_ncw[v_nn, v_ff, v_yy] = T.float32(0)
                    group_conv1d_ncw[v_nn, v_ff, v_yy] = group_conv1d_ncw[v_nn, v_ff, v_yy] + pad_temp[v_nn, v_ff // T.int64(8) * T.int64(16) + v_rc, v_yy * T.int64(2) + v_ry * T.int64(2)] * B[v_ff, v_rc, v_ry]
    # fmt: on

    mod = LegalizeOps()(Conv1d)