wehub-resource-sync
210 lines
6.9 KiB
Python
210 lines
6.9 KiB
Python
# 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
|
|
import tvm.testing
|
|
from tvm.script import tirx as T
|
|
from tvm.tirx.buffer import Buffer
|
|
from tvm.tirx.function import PrimFunc
|
|
from tvm.tirx.stmt import SBlock
|
|
|
|
|
|
def _check_func_signature_remap(lhs: PrimFunc, rhs: PrimFunc):
|
|
assert lhs != rhs
|
|
for x, y in zip(lhs.params, rhs.params):
|
|
assert x != y
|
|
assert lhs.buffer_map[x] != rhs.buffer_map[y]
|
|
|
|
|
|
def _check_buffer_decl(lhs: Buffer, rhs: Buffer):
|
|
assert lhs != rhs
|
|
assert lhs.data != rhs.data
|
|
|
|
|
|
def _check_block_signature_remap(lhs: SBlock, rhs: SBlock):
|
|
assert lhs != rhs
|
|
for x, y in zip(lhs.iter_vars, rhs.iter_vars):
|
|
assert x != y
|
|
assert x.var != y.var
|
|
for x, y in zip(lhs.alloc_buffers, rhs.alloc_buffers):
|
|
_check_buffer_decl(x, y)
|
|
for x, y in zip(lhs.match_buffers, rhs.match_buffers):
|
|
assert x != y
|
|
_check_buffer_decl(x.buffer, y.buffer)
|
|
|
|
|
|
def test_simple():
|
|
@T.prim_func(s_tir=True)
|
|
# Buffer A should be remapped
|
|
def elementwise(A: T.Buffer((128, 128), "float32")):
|
|
# Buffer B should be remapped
|
|
B = T.sblock_alloc_buffer((128, 128), "float32")
|
|
# i, j should be remapped
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
# vi, vj should be remapped
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
T.reads(A[vi, vj])
|
|
T.writes(B[vi, vj])
|
|
B[vi, vj] = A[vi, vj] * 2.0
|
|
|
|
f1 = elementwise
|
|
f2 = tvm.s_tir.renew_defs(f1)
|
|
tvm.ir.assert_structural_equal(f1, f2)
|
|
|
|
_check_func_signature_remap(f1, f2)
|
|
# check root block
|
|
_check_block_signature_remap(f1.body.block, f2.body.block)
|
|
# check remap of i
|
|
assert f1.body.block.body.loop_var != f2.body.block.body.loop_var
|
|
# check remap of j
|
|
assert f1.body.block.body.body.loop_var != f2.body.block.body.body.loop_var
|
|
|
|
# check inner block
|
|
def _get_sblock(f):
|
|
return f.body.block.body.body.body.block
|
|
|
|
_check_block_signature_remap(_get_sblock(f1), _get_sblock(f2))
|
|
|
|
|
|
def test_match_buffer():
|
|
# well-formed checker complains about multiple definitions for variable A0_s1,
|
|
# likely stemming from strides=[s, s]
|
|
@T.prim_func(check_well_formed=False, s_tir=True)
|
|
# A and B should be remapped
|
|
def func_match_buffer(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")):
|
|
with T.sblock("root"):
|
|
s = T.int32()
|
|
e = T.int32()
|
|
# A0 should be remapped
|
|
A0 = T.match_buffer(
|
|
A[0:128, 0:128],
|
|
shape=(128, 128),
|
|
dtype="float32",
|
|
# s and e should be remapped
|
|
strides=[s, s],
|
|
elem_offset=e,
|
|
)
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A0[vi, vj] * 2.0
|
|
|
|
f1 = func_match_buffer
|
|
f2 = tvm.s_tir.renew_defs(f1)
|
|
tvm.ir.assert_structural_equal(f1, f2)
|
|
|
|
_check_func_signature_remap(f1, f2)
|
|
_check_block_signature_remap(f1.body.block, f2.body.block)
|
|
assert f1.body.block.body.loop_var != f2.body.block.body.loop_var
|
|
|
|
def _get_sblock(f):
|
|
return f.body.block
|
|
|
|
block1 = _get_sblock(f1)
|
|
block2 = _get_sblock(f2)
|
|
_check_block_signature_remap(block1, block2)
|
|
|
|
matched_buffer1 = block1.match_buffers[0].buffer
|
|
matched_buffer2 = block2.match_buffers[0].buffer
|
|
# Stride var s should be remapped
|
|
assert matched_buffer1.strides[0] != matched_buffer2.strides[0]
|
|
assert matched_buffer1.strides[1] != matched_buffer2.strides[1]
|
|
# s should be only remapped once
|
|
assert matched_buffer1.strides[0] == matched_buffer1.strides[1]
|
|
assert matched_buffer2.strides[0] == matched_buffer2.strides[1]
|
|
# Element-offset var e should be remapped
|
|
assert matched_buffer1.elem_offset != matched_buffer2.elem_offset
|
|
|
|
|
|
def test_undefined_buffer():
|
|
@T.prim_func(s_tir=True)
|
|
def access_alloc():
|
|
# Buffer A should be remapped
|
|
A = T.alloc_buffer((128,), "float16")
|
|
T.evaluate(A.data)
|
|
for i in range(128):
|
|
A[i] = A[i] + T.float16(1.0)
|
|
|
|
f1 = access_alloc
|
|
f2 = tvm.s_tir.renew_defs(f1)
|
|
tvm.ir.assert_structural_equal(f1, f2)
|
|
|
|
# AllocBuffer is now a flat statement in SeqStmt
|
|
assert f1.body.seq[0].buffer.data != f2.body.seq[0].buffer.data
|
|
|
|
def _get_buffer_store_buffer(f):
|
|
# SeqStmt: [AllocBuffer, Evaluate, For]; For body has the BufferStore
|
|
return f.body.seq[2].body.buffer
|
|
|
|
_check_buffer_decl(_get_buffer_store_buffer(f1), _get_buffer_store_buffer(f2))
|
|
|
|
|
|
def test_symbolic_func():
|
|
@T.prim_func(s_tir=True)
|
|
def symbolic_func(a: T.handle, b: T.handle, n: T.int32):
|
|
m = T.int32()
|
|
A = T.match_buffer(a, (n, m))
|
|
B = T.match_buffer(b, (n, m * 2))
|
|
for i, j in T.grid(n, m):
|
|
B[i, j * 2] = A[i, j]
|
|
B[i, j * 2 + 1] = A[i, j]
|
|
|
|
f1 = symbolic_func
|
|
f2 = tvm.s_tir.renew_defs(f1)
|
|
tvm.ir.assert_structural_equal(f1, f2)
|
|
|
|
|
|
def test_buffer_map():
|
|
@T.prim_func(s_tir=True)
|
|
def main(a: T.handle, b: T.handle):
|
|
m = T.int64()
|
|
A = T.match_buffer(a, (m * 2,))
|
|
B = T.match_buffer(b, (m, 2))
|
|
for i, j in T.grid(m, 2):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi * 2 + vj]
|
|
|
|
f1 = main
|
|
f2 = tvm.s_tir.renew_defs(main)
|
|
tvm.ir.assert_structural_equal(f1, f2)
|
|
assert f1.buffer_map[f1.params[1]].shape[0] != f2.buffer_map[f2.params[1]].shape[0]
|
|
|
|
|
|
def test_gather():
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def take(
|
|
A: T.Buffer((4096, 4096), "float16"),
|
|
B: T.Buffer((1,), "int32"),
|
|
T_take: T.Buffer((1, 4096), "float16"),
|
|
):
|
|
for ax0, ax1 in T.grid(1, 4096):
|
|
with T.sblock("T_take"):
|
|
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
|
T.reads(A[B[v_ax0], v_ax1], B[v_ax0])
|
|
T.writes(T_take[v_ax0, v_ax1])
|
|
T_take[v_ax0, v_ax1] = A[B[v_ax0], v_ax1]
|
|
|
|
f1 = take
|
|
f2 = tvm.s_tir.renew_defs(take)
|
|
tvm.ir.assert_structural_equal(f1, f2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|