# 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. import tvm from tvm import te import os def test_unroll_loop(): ib = tvm.tir.ir_builder.create() dtype = 'int64' n = te.size_var('n') Ab = tvm.tir.decl_buffer((n, ), dtype) Aptr = ib.buffer_ptr(Ab) # for i in 0 to n-1: with ib.for_range(n, n + 2, name="i") as i: with ib.for_range(0, 8, name="i", for_type="unroll") as j: Aptr[j + 1] = Aptr[i] + 1 stmt = ib.get() assert isinstance(stmt, tvm.tir.For) ret = tvm.tir.ir_pass.UnrollLoop(stmt, 16, 8, 0, True) assert not isinstance(ret, tvm.tir.For) ret = tvm.tir.ir_pass.UnrollLoop(stmt, 15, 8, 0, True) assert isinstance(ret, tvm.tir.For) ret = tvm.tir.ir_pass.UnrollLoop(stmt, 16, 8, 0, False) assert isinstance(ret, tvm.tir.For) assert ret.for_type == tvm.tir.For.Unrolled ib = tvm.tir.ir_builder.create() ib.scope_attr(tvm.tir.const(0, "int32"), "pragma_auto_unroll_max_step", 16) ib.emit(stmt) wrapped = ib.get() wrapped = tvm.tir.SeqStmt([wrapped, stmt]) assert isinstance(ret, tvm.tir.For) ret = tvm.tir.ir_pass.UnrollLoop(wrapped, 0, 8, 0, False) assert isinstance(ret[0], tvm.tir.For) assert ret[0].for_type == tvm.tir.For.Unrolled assert isinstance(ret[1], tvm.tir.For) assert ret[1].for_type != tvm.tir.For.Unrolled def test_unroll_fake_loop(): ib = tvm.tir.ir_builder.create() dtype = 'int32' n = te.size_var('n') Ab = tvm.tir.decl_buffer((n, ), dtype) Aptr = ib.buffer_ptr(Ab) # for i in 0 to n-1: with ib.for_range(0, 1, name="i") as i: Aptr[i*2] = 3 with ib.for_range(0, 10, name="j") as j: Aptr[j + 1] = Aptr[i] + 1 stmt = ib.get() ret = tvm.tir.ir_pass.UnrollLoop(stmt, 8, 0, 1, True) assert isinstance(ret[0], tvm.tir.Store) def test_unroll_single_count_loops(): n = te.size_var('n') A = te.placeholder((n,), name='A') B = te.compute((n,), lambda *i: A(*i), name='B') s = te.create_schedule(B.op) s = s.normalize() dom_map = tvm.te.schedule.InferBound(s) stmt = tvm.te.schedule.ScheduleOps(s, dom_map) # all parameters to UnrolLoops are default values except for # auto_unroll_max_extent which has been set to 1 (default:0) after_unroll_stmt = tvm.tir.ir_pass.UnrollLoop(stmt, 0, 8, 1, True) assert after_unroll_stmt == stmt if __name__ == "__main__": test_unroll_loop() test_unroll_fake_loop() test_unroll_single_count_loops()