Python tvm.get_global_func函数代码示例

    def verify(target="llvm",
               algorithm=nnpack.ConvolutionAlgorithm.AUTO,
               with_bias=True):
        if not tvm.module.enabled(target):
            print("skip because %s is not enabled..." % target)
            return
        if not tvm.get_global_func("tvm.contrib.nnpack.fully_connected_inference", True):
            print("skip because extern function is not available")
            return
        if not nnpack.is_available():
            return

        ctx = tvm.cpu(0)
        transformed_kernel = nnpack.convolution_inference_weight_transform(
            kernel, algorithm=algorithm)
        output = nnpack.convolution_inference_without_weight_transform(
            data, transformed_kernel, bias if with_bias else None,
            [PAD, PAD, PAD, PAD], [STRIDE, STRIDE],
            algorithm=algorithm)

        s = tvm.create_schedule(output.op)

        f = tvm.build(s, [data, kernel, bias, output], target)

        na = np.random.uniform(size=dshape).astype(data.dtype)
        nb = np.random.uniform(size=kshape).astype(kernel.dtype)
        nc = np.random.uniform(size=bshape).astype(bias.dtype) if with_bias else np.zeros(bshape, dtype=bias.dtype)
        ta = tvm.nd.array(na, ctx)
        tb = tvm.nd.array(nb, ctx)
        tc = tvm.nd.array(nc, ctx)
        td = tvm.nd.array(np.zeros(oshape, dtype=output.dtype), ctx)
        f(ta, tb, tc, td)
        nd = np_conv(np.reshape(na, (BATCH, IC, IH, IW)), nb, PAD, STRIDE) + nc.reshape(1, bshape[0], 1, 1)
        tvm.testing.assert_allclose(
            td.asnumpy(), nd.reshape(BATCH, IC, IH, IW), rtol=1e-5)

 def check(factor):
     s = tvm.create_schedule(z.op)
     xo, xi = s[z].split(z.op.axis[0], factor=factor)
     vadd = intrin_vadd(factor)
     s[z].tensorize(xi, vadd)
     s = s.normalize()
     dom_map = tvm.schedule.InferBound(s)
     finfer = tvm.get_global_func("test.op.InferTensorizeRegion")
     out_dom, in_dom = finfer(s[z], dom_map)
     assert tvm.ir_pass.Equal(out_dom[z.op.axis[0]].extent, factor)
     assert tvm.ir_pass.Equal(out_dom[z.op.axis[0]].min, xo * factor)
     assert tvm.ir_pass.Equal(in_dom.items()[0][1][0].extent, factor)
     fmatch = tvm.get_global_func("test.op.MatchTensorizeBody")
     body = fmatch(s[z], out_dom, in_dom, vadd)
     assert tvm.ir_pass.Equal(tvm.ir_pass.CanonicalSimplify(body[0]),
                              tvm.ir_pass.CanonicalSimplify(vadd.op.body[0]))
     stmt = tvm.schedule.ScheduleOps(s, dom_map)
     tvm.lower(s, [x, y, z])

def test_conv2d():
    in_channel = 3
    out_channel = 64
    filter_h = 3
    filter_w = 3
    pad_h = 1
    pad_w = 1
    stride_h = 1
    stride_w = 1
    dilation_h = 1
    dilation_w = 1

    xshape = [1, in_channel, 128, 128]
    if not tvm.module.enabled("rocm"):
        print("skip because rocm is not enabled...")
        return
    if not tvm.get_global_func("tvm.contrib.miopen.conv2d.setup", True):
        print("skip because miopen is not enabled...")
        return
    wshape = (out_channel, in_channel, filter_h, filter_w)

    X = tvm.placeholder(xshape, name='X')
    W = tvm.placeholder(wshape, name='W')
    Y = miopen.conv2d_forward(X,
                              W,
                              stride_h,
                              stride_w,
                              pad_h,
                              pad_w,
                              dilation_h,
                              dilation_w,
                              conv_mode=0)

    yshape = [x.value for x in Y.shape]
    import topi
    with tvm.target.create("rocm -libs=miopen"):
        s = topi.generic.schedule_extern(Y)

    def verify():
        ctx = tvm.rocm(0)
        f = tvm.build(s, [X, W, Y], "rocm", target_host="llvm", name="conv2d")
        x = tvm.nd.array(np.random.uniform(-1, 1, xshape).astype(np.float32), ctx)
        w = tvm.nd.array(np.random.uniform(-1, 1, wshape).astype(np.float32), ctx)
        y = tvm.nd.array(np.random.uniform(-1, 1, yshape).astype(np.float32), ctx)
        f(x, w, y)

        Y_ref = topi.nn.conv2d_nchw(X, W, (stride_h, stride_w), (pad_h, pad_w))
        with tvm.target.rocm():
            s_ref = topi.generic.schedule_conv2d_nchw([Y_ref])
        f_ref = tvm.build(s_ref, [X, W, Y_ref], "rocm")
        y_ref = tvm.nd.array(np.random.uniform(-1, 1, yshape).astype(np.float32), ctx)
        f_ref(x, w, y_ref)
        print("Max abs diff:", np.max(np.abs(y.asnumpy() - y_ref.asnumpy())))
        tvm.testing.assert_allclose(y.asnumpy(), y_ref.asnumpy(), atol=1e-3)

    verify()

def stats():
    """Clear profiler statistics

    Returns
    -------
    stats : dict
        Current profiler statistics
    """
    x = tvm.get_global_func("vta.simulator.profiler_status")()
    return json.loads(x)

 def verify(A, B, C, target="llvm"):
     if not tvm.get_global_func("tvm.contrib.mps.conv2d", True):
         print("skip because extern function is not available")
         return
     ctx = tvm.metal(0)
     f = tvm.build(s1, [A, B, C], "metal")
     a = tvm.nd.array(np.random.uniform(size=(n, h, w, ci)).astype(A.dtype), ctx)
     b = tvm.nd.array(np.random.uniform(size=(co, kh, kw, ci)).astype(B.dtype), ctx)
     c = tvm.nd.array(np.zeros((n, h // stride, w // stride, co), dtype=C.dtype), ctx)
     f(a, b, c)

def test_get_global():
    targs = (10, 10.0, "hello")
    # register into global function table
    @tvm.register_func
    def my_packed_func(*args):
        assert(tuple(args) == targs)
        return 10
    # get it out from global function table
    f = tvm.get_global_func("my_packed_func")
    assert isinstance(f, tvm.Function)
    y = f(*targs)
    assert y == 10

def test_conv2d():
    in_channel = 3
    out_channel = 32
    filter_h = 3
    filter_w = 3
    pad_h = 1
    pad_w = 1
    stride_h = 1
    stride_w = 1
    dilation_h = 1
    dilation_w = 1

    xshape = [4, 3, 32, 32]
    if not tvm.module.enabled("cuda"):
        print("skip because cuda is not enabled...")
        return
    if not tvm.get_global_func("tvm.contrib.cudnn.conv2d.output_shape", True):
        print("skip because cudnn is not enabled...")
        return
    wshape = cudnn.conv2d_w_shape(in_channel,
                              out_channel,
                              filter_h,
                              filter_w)

    X = tvm.placeholder(xshape, name='X')
    W = tvm.placeholder(wshape, name='W')
    Y = cudnn.conv2d_forward(X,
                             W,
                             stride_h,
                             stride_w,
                             pad_h,
                             pad_w,
                             dilation_h,
                             dilation_w,
                             conv_mode=1,
                             tensor_format=0,
                             algo=1)
    yshape = [x.value for x in Y.shape]
    s =  tvm.create_schedule(Y.op)

    def verify():
        ctx = tvm.gpu(0)
        f = tvm.build(s, [X, W, Y], "cuda", target_host="llvm", name="conv2d")
        x = tvm.nd.array(np.random.uniform(-1, 1, xshape).astype(np.float32),
                         ctx)
        w = tvm.nd.array(np.random.uniform(-1, 1, wshape).astype(np.float32),
                         ctx)
        y = tvm.nd.array(np.random.uniform(-1, 1, yshape).astype(np.float32),
                         ctx)
        f(x, w, y)

    verify()

 def check_rfactor_no_reset_multi_reduction(factor, rfactor):
     s = tvm.create_schedule(C.op)
     x, y = C.op.axis
     rk = C.op.reduce_axis[0]
     yo, yi = s[C].split(y, factor=factor)
     ro, ri = s[C].split(rk, factor=rfactor)
     roo, roi = s[C].split(ro, factor=2)
     s[C].reorder(yo, roo, roi, yi, ri)
     gemv = intrin_gemv_no_reset(factor, rfactor)
     s[C].tensorize(yi, gemv)
     s = s.normalize()
     dom_map = tvm.schedule.InferBound(s)
     finfer = tvm.get_global_func("test.op.InferTensorizeRegion")
     out_dom, in_dom = finfer(s[C], dom_map)
     assert tvm.ir_pass.Equal(out_dom[x].extent, 1)
     assert tvm.ir_pass.Equal(out_dom[y].extent, factor)
     assert tvm.ir_pass.Equal(out_dom[y].min, yo * factor)
     fmatch = tvm.get_global_func("test.op.MatchTensorizeBody")
     body = fmatch(s[C], out_dom, in_dom, gemv)
     assert tvm.ir_pass.Equal(tvm.ir_pass.CanonicalSimplify(body[0]),
                              tvm.ir_pass.CanonicalSimplify(gemv.op.body[0]))
     stmt = tvm.schedule.ScheduleOps(s, dom_map)
     tvm.lower(s, [A, B, C])

 def verify(target="llvm"):
     if not tvm.module.enabled(target):
         print("skip because %s is not enabled..." % target)
         return
     if not tvm.get_global_func("tvm.contrib.random.normal", True):
         print("skip because extern function is not available")
         return
     ctx = tvm.cpu(0)
     f = tvm.build(s, [A], target)
     a = tvm.nd.array(np.zeros((m, n), dtype=A.dtype), ctx)
     f(a)
     na = a.asnumpy()
     assert abs(np.mean(na) - 3) < 1e-2
     assert abs(np.std(na) - 4) < 1e-2

 def verify(target="rocm"):
     if not tvm.module.enabled(target):
         print("skip because %s is not enabled..." % target)
         return
     if not tvm.get_global_func("tvm.contrib.rocblas.matmul", True):
         print("skip because extern function is not available")
         return
     ctx = tvm.rocm(0)
     f = tvm.build(s, [A, B, C], target)
     a = tvm.nd.array(np.random.uniform(size=(n, l)).astype(A.dtype), ctx)
     b = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), ctx)
     c = tvm.nd.array(np.zeros((n, m), dtype=C.dtype), ctx)
     f(a, b, c)
     tvm.testing.assert_allclose(
         c.asnumpy(), np.dot(a.asnumpy(), b.asnumpy()), rtol=1e-5)

 def verify(target="llvm"):
     if not tvm.module.enabled(target):
         print("skip because %s is not enabled..." % target)
         return
     if not tvm.get_global_func("tvm.contrib.cblas.matmul", True):
         print("skip because extern function is not avalable")
         return
     ctx = tvm.cpu(0)
     f = tvm.build(s, [A, B, D, bias], target)
     a = tvm.nd.array(np.random.uniform(size=(n, l)).astype(A.dtype), ctx)
     b = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), ctx)
     d = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), ctx)
     bb = 10.0
     f(a, b, d, bb)
     np.testing.assert_allclose(
         d.asnumpy(), np.dot(a.asnumpy(), b.asnumpy()) + bb, rtol=1e-5)

def test_get_callback_with_node():
    x = tvm.convert(10)
    def test(y):
        assert y.handle != x.handle
        return y

    f2 = tvm.convert(test)
    # register into global function table
    @tvm.register_func
    def my_callback_with_node(y, f):
        assert y == x
        return f(y)

    # get it out from global function table
    f = tvm.get_global_func("my_callback_with_node")
    assert isinstance(f, tvm.Function)
    y = f(x, f2)
    assert(y.value == 10)

    def verify(target="llvm"):
        if not tvm.module.enabled(target):
            print("skip because %s is not enabled..." % target)
            return
        if not tvm.get_global_func("tvm.contrib.nnpack.fully_connected_inference", True):
            print("skip because extern function is not avalable")
            return
        ctx = tvm.cpu(0)
        f = tvm.build(s, [data, kernel, bias, output], target)

        na = np.random.uniform(size=dshape).astype(data.dtype)
        nb = np.random.uniform(size=kshape).astype(kernel.dtype)
        nc = np.zeros(bshape, dtype=bias.dtype)
        ta = tvm.nd.array(na, ctx)
        tb = tvm.nd.array(nb, ctx)
        tc = tvm.nd.array(nc, ctx)
        td = tvm.nd.array(np.zeros(oshape, dtype=output.dtype), ctx)
        f(ta, tb, tc, td)
        nd = np_conv(na, nb, PAD)
        np.testing.assert_allclose(
            td.asnumpy(), nd, rtol=1e-5)

def test_env_func():
    @tvm.register_func("test.env_func")
    def test(x):
        return x + 1

    f = tvm.get_global_func("test.env_func")
    x = tvm.get_env_func("test.env_func")
    assert x.name == "test.env_func"
    json_str = tvm.save_json([x])
    y = tvm.load_json(json_str)[0]
    assert y.name == x.name
    assert y(1) == 2
    assert y.func(1) == 2

    x = tvm.make.node("attrs.TestAttrs", name="xx", padding=(3,4), func=y)
    assert x.name == "xx"
    assert x.padding[0].value == 3
    assert x.padding[1].value == 4
    assert x.axis == 10
    x = tvm.load_json(tvm.save_json(x))
    assert isinstance(x.func, tvm.container.EnvFunc)
    assert x.func(10) == 11

def save_tensors(params):
    """Save parameter dictionary to binary bytes.

    The result binary bytes can be loaded by the
    GraphModule with API "load_params".

    Parameters
    ----------
    params : dict of str to NDArray
        The parameter dictionary.

    Returns
    -------
    param_bytes: bytearray
        Serialized parameters.
    """
    _save_tensors = tvm.get_global_func("_save_param_dict")

    args = []
    for k, v in params.items():
        args.append(k)
        args.append(tvm.nd.array(v))
    return _save_tensors(*args)