test_hybrid_script.py 25.8 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.
17
import tvm, inspect, sys, traceback, numpy, pytest, types, os
18
from tvm.contrib import util
19
from tvm.hybrid import script
20
from tvm.hybrid.runtime import HYBRID_GLOBALS
21

22
@pytest.mark.skip
23
def run_and_check(func, args, var_dict={}, target='llvm', sch=None, outs=None):
24 25 26 27 28 29 30
    def tvm_val_2_py_val(val):
        val = tvm.ir_pass.Substitute(val, var_dict)
        val = tvm.ir_pass.Simplify(val)
        assert isinstance(val, (tvm.expr.IntImm, tvm.expr.UIntImm))
        return val.value

    ctx = tvm.context(target, 0)
31 32
    op = None

33 34 35 36 37 38 39 40
    if sch is None:
        outs = func(*tuple(tvm.convert(i) if isinstance(i, list) else i for i in args))
        op = outs[0].op if isinstance(outs, list) else outs.op
        sch = tvm.create_schedule(op)
    else:
        assert outs is not None
        assert isinstance(outs, list)
        op = outs[0].op
41 42 43 44 45 46

    emu_args = []
    nd_args = []
    for i in args:
        if isinstance(i, tvm.tensor.Tensor):
            shape = [tvm_val_2_py_val(j) for j in i.shape]
47 48
            emu_args.append(numpy.random.randn(*shape).astype(i.dtype))
            nd_args.append(tvm.nd.array(emu_args[-1], ctx))
49
        elif isinstance(i, tvm.expr.Var):
50 51
            emu_args.append(tvm_val_2_py_val(i))
            nd_args.append(emu_args[-1])
52 53 54
        else:
            assert isinstance(i, list)
            emu_args.append(numpy.array(i))
55

56 57
    compile_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))] + \
                   (outs if isinstance(outs, list) else [outs])
58
    module = tvm.build(sch,
59
                       compile_args,
60
                       target=target)
61
    assert module
62

63 64 65 66 67 68 69 70 71 72
    out_tensors = []
    for i in range(op.num_outputs):
        output = op.output(i)
        shape = [tvm_val_2_py_val(j) for j in output.shape]
        nd_args.append(tvm.nd.array(numpy.zeros(shape).astype(output.dtype), ctx))
        out_tensors.append(nd_args[-1])

    ref_data = func(*emu_args)
    if isinstance(ref_data, numpy.ndarray):
        ref_data = [ref_data]
73

74 75
    module(*nd_args)

76
    for nd, np in zip(out_tensors, ref_data):
77
        tvm.testing.assert_allclose(nd.asnumpy(), np, rtol=1e-5, atol=1e-5)
78

79 80 81 82 83
    module_args = [i for i in args if isinstance(i, (tvm.tensor.Tensor, tvm.expr.Var))]
    module_outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs
    h_module = tvm.hybrid.build(sch, module_args, module_outs)

    return h_module, module_args, module_outs
84

85
@script
86 87 88 89 90 91
def outer_product(n, m, a, b):
    """This is a simple outer product.
    Actually this function is not required to be documented.
    I write this docstring to test skipping docstring functionality.
    """
    c = output_tensor((n, m), a.dtype)
92 93
    for i in range(n):
        for j in range(m):
94
            assert i < n and j < m, "index out of range!"
95
            c[i, j] = a[i] * b[j]
96
    return c
97 98 99 100 101 102 103 104

#Test global function
#Test bridge between frontend and backend
def test_outer_product():
    n = tvm.var('n')
    m = tvm.var('m')
    a = tvm.placeholder((n, ), name='a')
    b = tvm.placeholder((m, ), name='b')
105 106 107 108 109 110 111 112

    try:
        c = outer_product(n, m, a, b)
        ir = c.op.body
    except IOError as err:
        assert sys.version_info[0] == 2 and str(err) == 'could not get source code'
        return

113 114 115 116 117 118 119 120 121 122 123 124
    #Check for i in (0, n)
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'i'
    assert ir.min.value == 0
    assert ir.extent.name == 'n'
    ibody = ir.body
    assert isinstance(ibody, tvm.stmt.For)
    #Check for j in (0, m)
    assert ibody.loop_var.name == 'j'
    assert ibody.min.value == 0
    assert ibody.extent.name == 'm'
    #Check loop body
125 126 127
    jblock = ibody.body
    assert isinstance(jblock, tvm.stmt.Block)
    jbody = jblock.first
128 129 130
    assert isinstance(jbody, tvm.stmt.AssertStmt)
    assert isinstance(jbody.message, tvm.expr.StringImm)
    assert jbody.message.value == "index out of range!"
131
    jbody = jblock.rest
132 133 134 135 136 137 138 139 140 141 142
    assert isinstance(jbody, tvm.stmt.Provide)
    assert jbody.func.name == 'c'
    assert len(jbody.args) == 2
    assert jbody.args[0].name == 'i'
    assert jbody.args[1].name == 'j'
    assert isinstance(jbody.value, tvm.expr.Mul)
    mul = jbody.value
    assert isinstance(mul.a, tvm.expr.Call)
    assert mul.a.name == 'a'
    assert mul.b.name == 'b'

143 144 145 146 147 148 149
    func, ins, outs = run_and_check(outer_product, [n, m, a, b], {n: 99, m: 101})
    temp = util.tempdir()
    path = temp.relpath('%s.py' % func.name)
    func.save(path)
    func_ = tvm.hybrid.HybridModule()
    func_.load(path)
    run_and_check(func_, ins, {n: 99, m: 101}, outs=outs)
150

151 152 153 154 155 156 157 158
    for key, _ in HYBRID_GLOBALS.items():
        assert key not in globals().keys()
        assert key not in outer_product.__globals__.keys()

#Test local function
#Test allocation of local variable
def test_fanout():
    @script
159
    def fanout(n, a):
160
        three = 3.0
161
        b = output_tensor((a.shape[0] - 3, ), a.dtype)
162 163 164
        for i in range(a.shape[0] - 3):
            sigma = 0.0
            for j in range(3):
165
                sigma += a[i + j]
166 167
            sigma = sigma / three
            b[i] = sigma
168
        return b
169 170

    n = tvm.var('n')
171
    a = tvm.placeholder((n, ), 'float32', name='a')
172 173 174 175 176 177
    try:
        b = fanout(n, a)
        ir = b.op.body
    except IOError as err:
        assert sys.version_info[0] == 2 and str(err) == 'could not get source code'
        return
178 179 180 181 182 183 184 185

    #Check for i in (0, n-3)
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'i'
    assert ir.min.value == 0
    assert tvm.ir_pass.Equal(ir.extent, n - 3)
    #Check loopbody
    ibody = ir.body
186 187 188 189 190 191
    assert isinstance(ibody, tvm.stmt.AttrStmt)
    abody = ibody.body
    assert isinstance(abody, tvm.stmt.Realize)
    assert abody.bounds[0].min.value == 0
    assert abody.bounds[0].extent.value == 1
    assert abody.func.name == 'sigma'
192
    #Check i loop body
193
    rbody = abody.body
194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233
    assert isinstance(rbody.first, tvm.stmt.Provide)
    assert rbody.first.func.name == 'sigma'
    assert len(rbody.first.args) == 1
    assert rbody.first.args[0].value == 0
    #Check fanout loop
    jloop = rbody.rest.first
    assert jloop.loop_var.name == 'j'
    assert jloop.min.value == 0
    assert jloop.extent.value == 3
    jbody = jloop.body
    assert isinstance(jbody, tvm.stmt.Provide)
    assert len(jbody.args) == 1
    assert jbody.args[0].value == 0
    assert jbody.func.name == 'sigma'
    assert isinstance(jbody.value, tvm.expr.Add)
    value = jbody.value
    assert isinstance(value.a, tvm.expr.Call)
    assert value.a.name == 'sigma'
    assert len(value.a.args) == 1
    assert value.a.args[0].value == 0
    assert value.b.name == 'a'
    assert len(value.b.args) == 1
    assert tvm.ir_pass.Equal(value.b.args[0], ir.loop_var + jloop.loop_var)
    divide= rbody.rest.rest.first
    assert isinstance(divide, tvm.stmt.Provide)
    assert len(divide.args) == 1
    assert divide.args[0].value == 0
    value = divide.value
    assert isinstance(value, tvm.expr.Mul)
    assert value.a.name == 'sigma'
    assert len(value.a.args) == 1
    assert value.a.args[0].value == 0
    assert abs(value.b.value - (1 / 3.0)) < 1e-5
    write = rbody.rest.rest.rest
    assert isinstance(write, tvm.stmt.Provide)
    assert write.func.name == 'b'
    assert write.value.name == 'sigma'
    assert len(write.value.args) == 1
    assert write.value.args[0].value == 0

234 235
    func, ins, outs = run_and_check(fanout, [n, a], {n: 10})
    run_and_check(func, ins, {n: 10}, outs=outs)
236 237 238 239


def test_looptype():
    @script
240
    def looptype(a, b, c):
241 242 243 244
        d = output_tensor((16, ), 'int32')
        e = output_tensor((16, ), 'int32')
        f = output_tensor((16, ), 'int32')
        for i in parallel(16):
245
            d[i] = a[i]
246
        for j in vectorize(16):
247
            e[j] = b[j]
248
        for k in unroll(16):
249 250
            f[k] = c[k]
        return d, e, f
251

252 253 254
    a = tvm.placeholder((16, ), name='a', dtype='int32')
    b = tvm.placeholder((16, ), name='b', dtype='int32')
    c = tvm.placeholder((16, ), name='c', dtype='int32')
255 256 257 258 259
    try:
        d, e, f = looptype(a, b, c)
        ir = d.op.body
    except:
        return
260 261 262 263 264 265 266
    iloop = ir.first
    jloop = ir.rest.first
    kloop = ir.rest.rest
    assert iloop.for_type == tvm.stmt.For.Parallel
    assert jloop.for_type == tvm.stmt.For.Vectorized
    assert kloop.for_type == tvm.stmt.For.Unrolled

267 268
    func, ins, outs = run_and_check(looptype, [a, b, c])
    run_and_check(func, ins, outs=outs)
269 270


271 272
def test_if():
    @script
273 274 275
    def if_then_else(a):
        b = output_tensor((10, ), 'int32')
        c = output_tensor((10, ), 'int32')
276 277
        for i in range(10):
            if i % 2 == 0:
278
                c[i] = a[i]
279
            else:
280
                c[i] = b[i]
281 282
        for i in unroll(10):
            b[i] = -1 if i % 2 == 0 else 1
283
        return b, c
284 285 286

    a = tvm.placeholder((10, ), dtype='int32', name='a')

287 288
    func, ins, outs = run_and_check(if_then_else, [a])
    run_and_check(func, ins, outs=outs)
289

290 291 292 293 294 295 296 297 298 299
    @script
    def if_triple_condition(a):
        b = output_tensor((10, ), 'int32')
        for i in range(10):
            if 0 <= i < 5:
                b[i] = a[i]
            else:
                b[i] = a[i] + 1
        return b

300 301
    func, ins, outs = run_and_check(if_triple_condition, [a])
    run_and_check(func, ins, outs=outs)
302 303 304 305 306 307 308 309 310 311 312

    @script
    def if_and(a):
        b = output_tensor((10, ), 'int32')
        for i in range(10):
            if i >= 0 and i < 5:
                b[i] = a[i]
            else:
                b[i] = a[i] + 1
        return b

313 314
    func, ins, outs = run_and_check(if_and, [a])
    run_and_check(func, ins, outs=outs)
315

316 317 318

def test_bind():
    if not tvm.gpu(0).exist:
319
        print('[Warning] No GPU found! Skip bind test!')
320
        return
Jian Weng committed
321

322
    @script
323
    def vec_add(a, b):
324
        c = output_tensor((1000, ), 'float32')
325
        for tx in bind('threadIdx.x', 1000):
326
            c[tx] = a[tx] + b[tx]
327
        return c
328 329 330

    a = tvm.placeholder((1000, ), dtype='float32', name='a')
    b = tvm.placeholder((1000, ), dtype='float32', name='b')
331 332
    func, ins, outs = run_and_check(vec_add, [a, b], target='cuda')
    run_and_check(func, ins, outs=outs, target='cuda')
333

334 335 336 337 338 339 340 341 342 343 344
    @script
    def raw(a, b):
        c = output_tensor((1000, ), 'float32')
        for i in range(1000):
            c[i] = a[i] + b[i]
        return c

    c = raw(a, b)
    sch = tvm.create_schedule(c.op)
    x = tvm.thread_axis('threadIdx.x')
    sch[c].bind(c.op.axis[0], x)
345 346
    func, ins, outs = run_and_check(raw, [a, b], sch=sch, outs=[c], target='cuda')
    run_and_check(func, ins, outs=outs, target='cuda')
347

Jian Weng committed
348

349
    @tvm.hybrid.script
Jian Weng committed
350 351 352 353 354 355 356 357 358 359
    def foo(a):
        c = output_tensor((a.shape[0],), a.dtype)
        total = allocate((1,), a.dtype, 'local')
        len_i = a.shape[0]
        len_j = a.shape[1]
        for i in bind('threadIdx.x', len_i):
            total[0] = 0.
            for k in const_range(len_j):
                total[0] += a[i, k]
            c[i] = total[0]
360

361
        return c
362

Jian Weng committed
363 364 365 366 367 368 369 370
    a = tvm.placeholder((8, 4), 'float32')
    c = foo(a)
    s = tvm.create_schedule(c.op)
    ir = tvm.lower(s, [a, c], simple_mode=True)
    assert not isinstance(ir, tvm.stmt.AttrStmt)
    func, ins, outs = run_and_check(foo, [a], target='cuda')
    run_and_check(func, ins, outs=outs, target='cuda')

371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386
    @tvm.hybrid.script
    def max_threads(a):
        b = output_tensor(a.shape, a.dtype)
        n = a.shape[0]
        m = max_num_threads(True)
        for i in bind('threadIdx.x', m):
            for j in bind('blockIdx.x', ceil_div(n, m)):
                if i * m + j < n:
                    b[i * m + j] = a[i * m + j] + a[i * m + j]
        return b

    a = tvm.placeholder((10000, ), 'float32')
    with tvm.target.create('cuda'):
        func, ins, outs = run_and_check(max_threads, [a], target='cuda')
        run_and_check(func, ins, outs=outs, target='cuda')

387

388 389 390
def test_math_intrin():
    @script
    def intrin_real(a):
391 392 393 394 395 396 397 398 399 400
        b = output_tensor((8, ), 'float32')
        b[0] = sqrt(a[0])
        b[1] = log(a[1])
        b[2] = exp(a[2])
        b[3] = sigmoid(a[3])
        b[4] = power(a[4], a[5])
        b[5] = tanh(a[5])
        b[6] = min(a[4], a[5])
        b[7] = max(a[5], a[6])
        return b
401

402
    a8 = tvm.placeholder((8, ), dtype='float32', name='a')
403 404 405
    b8 = intrin_real(a8)
    sch = tvm.create_schedule(b8.op)
    func = tvm.build(sch, [a8, b8])
406
    assert func
407
    a = numpy.arange(2, 10).astype('float32')
408
    tvm_a = tvm.ndarray.array(a)
409 410 411 412
    tvm_b = tvm.ndarray.array(numpy.zeros((8, ), dtype='float32'))
    b = intrin_real(a)
    func(tvm_a, tvm_b)
    tvm.testing.assert_allclose(b, tvm_b.asnumpy(), rtol=1e-5)
413 414 415

    @script
    def intrin_int(a):
416 417 418
        b = output_tensor((1, ), 'int32')
        b[0] = popcount(a[0])
        return b
419 420

    a1 = tvm.placeholder((1, ), dtype='int32')
421 422 423
    b1 = intrin_int(a1)
    sch = tvm.create_schedule(b1.op)
    func = tvm.build(sch, [a1, b1])
424
    assert func
425
    a = numpy.array([114514]).astype('int32')
426
    tvm_a = tvm.ndarray.array(a)
427 428 429 430
    tvm_b = tvm.ndarray.array(numpy.array([0]).astype('int32'))
    b = intrin_int(a)
    func(tvm_a, tvm_b)
    assert tvm_b.asnumpy()[0] == b[0]
431

432
# test non caconical loops
433 434
def test_non_zero():
    @tvm.hybrid.script
435 436
    def blur(a):
        b = output_tensor((30, 30), 'float32')
437 438 439 440 441
        for i in range(2, 32):
            for j in range(2, 32):
                s = 0.0
                for di in range(3):
                    for dj in range(3):
442
                        s += a[i-di, j-dj]
443
                b[i-2, j-2] = s / 9.0
444 445 446
        return b

    a = tvm.placeholder((32, 32), 'float32', 'a')
447 448
    func, ins, outs = run_and_check(blur, [a])
    run_and_check(func, ins, outs=outs)
449 450

    @tvm.hybrid.script
451 452
    def triangle(a, b):
        c = output_tensor((10, 10), dtype='float32')
453 454 455
        for i in range(10):
            for j in range(i, 10):
                c[i, j] = a[i] * b[j]
456
        return c
457 458 459 460

    a = tvm.placeholder((10, ), dtype='float32', name='a')
    b = tvm.placeholder((10, ), dtype='float32', name='b')

461 462
    func, ins, outs = run_and_check(triangle, [a, b])
    run_and_check(func, ins, outs=outs)
463 464 465

def test_allocate():
    @tvm.hybrid.script
466 467
    def blur2d(a):
        b = output_tensor((30, 30), 'float32')
468 469 470 471 472 473 474
        for i in range(30):
            ha = allocate((3, 30), 'float32')
            for j in range(3):
                for k in range(30):
                    ha[j, k] = a[i+j, k] + a[i+j, k+1] + a[i+j, k+2]
            for j in range(30):
                b[i, j] = (ha[0, j] + ha[1, j] + ha[2, j]) / 9.0
475
        return b
476 477

    a = tvm.placeholder((32, 32), 'float32', 'a')
478 479 480 481
    b = blur2d(a)
    sch = tvm.create_schedule(b.op)
    func, ins, outs = run_and_check(blur2d, [a])
    run_and_check(func, ins, outs=outs)
482 483 484

    if tvm.gpu().exist:
        @tvm.hybrid.script
485 486
        def share_vec_add(a, b):
            c = output_tensor((256, ), 'float32')
487 488 489 490 491 492 493 494
            shared = allocate((256, ), 'float32', 'shared')
            for i in bind("threadIdx.x", 256):
                shared[i] = a[i]
            local = allocate((256, ), 'float32', 'local')
            for i in bind("threadIdx.x", 256):
                local[i] = b[i]
            for i in bind("threadIdx.x", 256):
                c[i] = shared[i] + local[i]
495
            return c
496 497 498

        a = tvm.placeholder((256, ), dtype='float32', name='a')
        b = tvm.placeholder((256, ), dtype='float32', name='b')
Jian Weng committed
499
        c = share_vec_add(a, b)
500 501
        func, ins, outs = run_and_check(share_vec_add, [a, b], target='cuda')
        run_and_check(func, ins, outs=outs, target='cuda')
502 503
    else:
        print('[Warning] No GPU found! Skip shared mem test!')
504

505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541
def test_upstream():
    @tvm.hybrid.script
    def upstream(a):
        b = output_tensor((20, ), 'float32')
        for i in range(20):
            b[i] = a[i] * i
        return b

    a = tvm.placeholder((20, ), 'float32')
    b = tvm.placeholder((20, ), 'float32')
    c = tvm.compute((20, ), lambda x: a[x] + b[x])
    d = upstream(c)
    sch = tvm.create_schedule([c.op, d.op])
    ir = tvm.lower(sch, [a, b, d], simple_mode=True)
    func = tvm.build(sch, [a, b, d])
    assert(func)

    a = numpy.random.randn(20).astype('float32')
    b = numpy.random.randn(20).astype('float32')
    ref = numpy.zeros((20, ), 'float32')
    for i in range(20):
        ref[i] = (a[i] + b[i]) * i

    tvm_a = tvm.nd.array(a)
    tvm_b = tvm.nd.array(b)
    tvm_d = tvm.nd.array(numpy.zeros((20, )).astype('float32'))

    func(tvm_a, tvm_b, tvm_d)
    tvm.testing.assert_allclose(tvm_d.asnumpy(), ref, 1e-5, 1e-5)

def test_downstream():
    @tvm.hybrid.script
    def downstream(a):
        b = output_tensor((20, ), 'float32')
        for i in range(20):
            b[i] = a[i] * i
        return b
542

543

544 545 546
    a = tvm.placeholder((20, ), 'float32')
    b = downstream(a)
    c = tvm.compute((20, ), lambda x: b[x] + 1.0)
547

548 549 550 551 552 553 554 555 556 557 558 559 560 561
    sch = tvm.create_schedule(c.op)
    module = tvm.build(sch, [a, c])
    assert module

    a = numpy.random.randn(20).astype('float32')
    ref = numpy.zeros((20, )).astype('float32')
    for i in range(20):
        ref[i] = (a[i] * i) + 1.0

    tvm_a = tvm.nd.array(a)
    tvm_c = tvm.nd.array(numpy.zeros((20, )).astype('float32'))
    module(tvm_a, tvm_c)
    tvm.testing.assert_allclose(tvm_c.asnumpy(), ref, 1e-5, 1e-5)

562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587
def test_const_param():
    @tvm.hybrid.script
    def add_something(a, b):
        c = output_tensor((11, ), 'int32')
        for i in range(11):
            c[i] = a[i] + b
        return c

    a = tvm.placeholder((11, ), dtype='int32', name='a')
    b = tvm.const(11, 'int32')
    c = add_something(a, b)
    sch = tvm.create_schedule(c.op)
    module = tvm.build(sch, [a, c], 'llvm')
    assert(module)

    np_a = numpy.arange(11).astype('int32')
    np_b = 11
    np_c = numpy.zeros((11, )).astype('int32')

    nd_a = tvm.ndarray.array(np_a)
    nd_c = tvm.ndarray.array(numpy.zeros((11, )).astype('int32'))
    module(nd_a, nd_c)
    ref = add_something(np_a, 11)

    tvm.testing.assert_allclose(nd_c.asnumpy(), ref, 1e-5, 1e-5)

588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620
def test_value_index():
    @tvm.hybrid.script
    def kernel_a(a):
        b = output_tensor((16, ), 'int32')
        c = output_tensor((4, 4), 'int32')
        for i in range(16):
            b[i] = a[i] + 2
            c[i // 4, i % 4] = a[i] + 1
        return b, c

    @tvm.hybrid.script
    def kernel_b(b, a):
        c = output_tensor((4, 4), 'int32')
        for i in range(4):
            for j in range(4):
                c[i, j] = a[i * 4 + j] * b[i, j]
        return c

    a = tvm.placeholder((16, ), 'int32')
    b, c = kernel_a(a)
    d = kernel_b(c, b)
    sch = tvm.create_schedule(d.op)
    module = tvm.build(sch, [a, d])
    assert module

    np_a = numpy.arange(16).astype('int32')
    np_b, np_c = kernel_a(np_a)
    ref = kernel_b(np_c, np_b)

    res = tvm.ndarray.array(numpy.zeros((4, 4)).astype('int32'))
    module(tvm.ndarray.array(np_a), res)
    tvm.testing.assert_allclose(res.asnumpy(), ref)

621 622 623
def test_func_call():
    @tvm.hybrid.script
    def foo(a, b):
624
        for i in range(len(a)):
625
            a[i] = i + 1.0
626
        for i in range(len(a)):
627 628 629 630 631 632 633
            b[i] = i + 1.0
        c = outer_product(10, 10, a, b)
        d = output_tensor(c.shape, c.dtype)
        for i in range(10):
            for j in range(10):
                d[i, j] = c[i, j] + i * j
        return d
634

635 636
    a = tvm.placeholder((10, ), name='a')
    b = tvm.placeholder((10, ), name='b')
637 638
    func, ins, outs = run_and_check(foo, [a, b])
    run_and_check(func, ins, outs=outs)
639 640 641 642 643 644 645 646 647 648 649 650 651

def test_bool():
    @tvm.hybrid.script
    def foo(a):
        b = output_tensor(a.shape, a.dtype)
        b[0] = 1.2
        for i in range(1, a.shape[0] - 1):
            if a[i] * a[i - 1] < a[i] or a[i] * a[i - 1] < a[i - 1] or i * a[i] == a[i]:
                b[i] = a[i]
            else:
                b[i] = 0.0
        return b
    a = tvm.placeholder((10, ), name='a')
652 653
    func, ins, outs = run_and_check(foo, [a])
    run_and_check(func, ins, outs=outs)
654

655 656 657 658
def test_const_range():
    @tvm.hybrid.script
    def foo(a, b):
        c = output_tensor(a.shape, a.dtype)
659
        d = output_tensor(a.shape, 'int32')
660 661 662

        for i in const_range(2):
            for j in const_range(5):
663
                c[i, j] = float32(int32(a[i, j]) + b[i, j])
664 665 666

        for i in const_range(len(b)):
            for j in const_range(len(b[0])):
667
                d[i, j] = int32(a[i, j] + b[i, j])
668 669 670

        return c, d

671
    a = tvm.placeholder((2, 5), name='a', dtype='float32')
672
    b = [[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]]
673 674
    func, ins, outs = run_and_check(foo, [a, b])
    run_and_check(func, ins, outs=outs)
675

676 677 678 679 680 681 682 683 684 685 686 687 688 689
    @tvm.hybrid.script
    def goo(a, b):
        c = output_tensor(a.shape, a.dtype)
        len_b = len(b)
        for i in const_range(len_b * 2):
            if i < len_b:
                c[i] = a[i] + b[i]
            else:
                c[i - len_b] = a[i - len_b] + b[i - len_b]
        return c
    a = tvm.placeholder((5, ), name='a', dtype='int32')
    b = [1, 2, 3, 4, 5]
    c = goo(a, tvm.convert(b))
    sch = tvm.create_schedule(c.op)
690 691
    func, ins, outs = run_and_check(goo, [a, b])
    run_and_check(func, ins, outs=outs)
692 693 694 695 696 697 698 699 700 701 702 703 704

    @tvm.hybrid.script
    def hoo(a, b):
        c = output_tensor(a.shape, a.dtype)
        len_b = len(b)
        for i in range(a.shape[0]):
            for j in const_range(len(b)):
                d = a[i] * b[j]
                d += a[i] + b[j]
                c[i] = d
        return c
    a = tvm.placeholder((5, ), name='a', dtype='int32')
    b = [1, 2, 3, 4, 5]
705 706
    func, ins, outs = run_and_check(hoo, [a, b])
    run_and_check(func, ins, outs=outs)
707

708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747
def test_schedule():
    @script
    def outer_product(a, b):
        c = output_tensor((64, 64), a.dtype)
        for i in range(64):
            for j in range(64):
                c[i, j] = a[i] * b[j]
        return c
    a = tvm.placeholder((64,), name='a', dtype='float32')
    b = tvm.placeholder((64,), name='b', dtype='float32')
    c = outer_product(a, b)

    # Test perfect loop split
    # Test loop reorder
    # Test loop annotation
    sch = tvm.create_schedule(c.op)
    i, j = c.op.axis
    io, ii = sch[c].split(i, 4)
    sch[c].parallel(ii)
    jo, ji = sch[c].split(j, 4)
    joo, joi = sch[c].split(jo, 4)
    sch[c].vectorize(ji)
    sch[c].reorder(ii, io, joo, joi, ji)
    ir = tvm.lower(sch, [a, b, c], simple_mode=True)
    assert isinstance(ir, tvm.stmt.ProducerConsumer)
    ir = ir.body
    assert isinstance(ir, tvm.stmt.AttrStmt)
    ir = ir.body
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'i.inner'
    ir = ir.body
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'i.outer'
    ir = ir.body
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'j.outer.outer'
    ir = ir.body
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'j.outer.inner'
    ir = ir.body
748 749
    func, ins, outs = run_and_check(outer_product, [a, b], sch=sch, outs=[c])
    run_and_check(func, ins, outs=outs)
750 751 752 753 754 755 756 757 758 759 760

    # Test fuse
    sch = tvm.create_schedule(c.op)
    sch[c].fuse(c.op.axis[0], c.op.axis[1])
    ir = tvm.lower(sch, [a, b, c], simple_mode=True)
    assert isinstance(ir, tvm.stmt.ProducerConsumer)
    ir = ir.body
    assert isinstance(ir, tvm.stmt.AttrStmt)
    ir = ir.body
    assert isinstance(ir, tvm.stmt.For)
    assert ir.loop_var.name == 'i.j.fused'
761 762
    func, ins, outs = run_and_check(outer_product, [a, b], sch=sch, outs=[c])
    run_and_check(func, ins, outs=outs)
763 764 765 766 767

    # Test imperfect loop split
    sch = tvm.create_schedule(c.op)
    sch[c].split(c.op.axis[0], 3)
    ir = tvm.lower(sch, [a, b, c], simple_mode=True)
768 769
    func, ins, outs = run_and_check(outer_product, [a, b], sch=sch, outs=[c])
    run_and_check(func, ins, outs=outs)
770 771 772

    # Test loop binds

773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790
def test_capture():
    n = 8

    constant_tuple = (10, n)
    constant_list = [[1, 2], [3, n]]
    const_value = 1

    @tvm.hybrid.script
    def add_something(a):
        c = output_tensor((constant_tuple[1],), 'int32')
        for i in range(constant_tuple[1]):
            c[i] = a[i] + constant_list[1][const_value]
        return c

    a = tvm.placeholder((n, ), dtype='int32', name='a')

    func, ins, outs = run_and_check(add_something, [a])
    run_and_check(func, ins, outs=outs)
791

792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822
def test_array_inputs():
    @script
    def sum_array(inputs):
        out = output_tensor((10,), inputs[0].dtype)
        n = len(inputs)
        for i in range(10):
            for j in const_range(n):
                out[i] += inputs[j][i]
        return out
    n = 5
    inputs = []
    for i in range(n):
        inputs.append(tvm.placeholder((10,), name='t%s' % i, dtype='float32'))
    
    out = sum_array(tvm.convert(inputs))
    assert len(out.op.inputs) == n

    sch = tvm.create_schedule(out.op)
    mod = tvm.build(sch, inputs + [out], target='llvm')
    assert mod

    input_nd = []
    out_ref = numpy.zeros((10,))
    for _ in range(n):
        arr = numpy.random.uniform(size=(10,)).astype('float32')
        input_nd.append(tvm.nd.array(arr))
        out_ref += arr
    out_nd = tvm.nd.array(numpy.zeros((10,), 'float32'))
    mod(*input_nd, out_nd)
    tvm.testing.assert_allclose(out_nd.asnumpy(), out_ref)

823 824 825 826 827 828 829
if __name__ == "__main__":
    test_outer_product()
    test_fanout()
    test_looptype()
    test_if()
    test_bind()
    test_math_intrin()
830 831
    test_non_zero()
    test_allocate()
832 833
    test_upstream()
    test_downstream()
834
    test_const_param()
835
    test_value_index()
836 837
    test_func_call()
    test_bool()
838
    test_const_range()
839
    test_schedule()
840
    test_capture()
841
    test_array_inputs()
842 843
    # TODO:
    # test_inplace()