# 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 def test_lower_warp_mem(): m = 128 A = tvm.placeholder((m,), name='A') B = tvm.compute((m,), lambda i: A[i] + 3, name='B') s = tvm.create_schedule(B.op) AA = s.cache_read(A, "warp", [B]) xo, xi = s[B].split(B.op.axis[0], 32) xi0, xi1 = s[B].split(xi, factor=16) tx = tvm.thread_axis("threadIdx.x") s[B].bind(xi1, tx) s[B].bind(xo, tvm.thread_axis("blockIdx.x")) s[AA].compute_at(s[B], xo) xo, xi = s[AA].split(s[AA].op.axis[0], 16) s[AA].bind(xi, tx) f = tvm.lower(s, [A, B]) fhost, fdevice = tvm.ir_pass.SplitHostDevice(f) fdevice = tvm.ir_pass.LowerWarpMemory(fdevice, 16) assert(fdevice.body.body.value.value == "local") assert(fdevice.body.body.body.extents[0].value == 2) if __name__ == "__main__": test_lower_warp_mem()