test_dot.py 2.39 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
Tianqi Chen committed
17 18 19 20 21 22 23 24 25
import tvm
import numpy as np

def lower(s, args, name="mydot"):
    binds = {}
    arg_list = []

    for x in args:
        assert isinstance(x, tvm.tensor.Tensor)
26
        buf = tvm.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
Tianqi Chen committed
27 28
        binds[x] = buf
        arg_list.append(buf)
29
    s = s.normalize()
Tianqi Chen committed
30 31
    bounds = tvm.schedule.InferBound(s)
    stmt = tvm.schedule.ScheduleOps(s, bounds)
32
    stmt = tvm.ir_pass.StorageFlatten(stmt, binds, 16)
Tianqi Chen committed
33 34
    stmt = tvm.ir_pass.CanonicalSimplify(stmt)
    stmt = tvm.ir_pass.Simplify(stmt)
35
    fapi = tvm.ir_pass.MakeAPI(stmt, name, arg_list, 0, True)
36
    fapi = tvm.ir_pass.LowerTVMBuiltin(fapi)
Tianqi Chen committed
37 38 39 40 41 42 43 44 45
    return fapi


def mybuild(fapi, target="llvm"):
    return


def test_dot():
    nn = 12
46
    n = tvm.convert(nn)
Tianqi Chen committed
47 48
    A = tvm.placeholder((n,), name='A')
    B = tvm.placeholder((n,), name='B')
49
    k = tvm.reduce_axis((0, n), 'k')
Tianqi Chen committed
50
    C = tvm.compute((1,), lambda _: tvm.sum(A[k] * B[k], axis=k), name='C')
51
    s = tvm.create_schedule(C.op)
Tianqi Chen committed
52 53 54
    fapi = lower(s, [A, B, C])

    def verify(target):
55
        if not tvm.module.enabled(target):
Tianqi Chen committed
56 57
            print("Target %s is not enabled" % target)
            return
58
        f = tvm.codegen.build_module(fapi, target)
Tianqi Chen committed
59 60 61 62 63 64
        # verify
        ctx = tvm.cpu(0)
        a = tvm.nd.array(np.random.uniform(size=(nn,)).astype(A.dtype), ctx)
        b = tvm.nd.array(np.random.uniform(size=(nn,)).astype(B.dtype), ctx)
        c  = tvm.nd.array(np.zeros((1,), dtype=C.dtype), ctx)
        f(a, b, c)
65
        tvm.testing.assert_allclose(
Tianqi Chen committed
66 67 68 69 70 71
            c.asnumpy(), np.dot(a.asnumpy(), b.asnumpy()), rtol=1e-4)

    verify("llvm")

if __name__ == "__main__":
    test_dot()