chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,289 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-docstring,
|
||||
from collections.abc import Sequence
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from tvm_ffi import Shape
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import dlight as dl
|
||||
from tvm.script import tirx as T
|
||||
from tvm.testing import env
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
|
||||
np_zero = np.full((16, 16), 0.0, "float16")
|
||||
np_one = np.full((32, 32), 1.0, "float32")
|
||||
np_two = np.full((16, 16), 2.0, "float16")
|
||||
np_three = np.full((32, 32), 3.0, "float32")
|
||||
|
||||
reserved_nseq = 4
|
||||
max_history = 4
|
||||
num_layers = 1
|
||||
# Note that kernels in this test file cannot support 1-dim states.
|
||||
states = [((16, 16), "float16"), ((32, 32), "float32")]
|
||||
|
||||
f_clear = None
|
||||
f_add_sequence = None
|
||||
f_remove_sequence = None
|
||||
f_fork_sequence = None
|
||||
f_popn = None
|
||||
f_begin_forward = None
|
||||
f_end_forward = None
|
||||
f_get = None
|
||||
f_set = None
|
||||
f_debug_get = None
|
||||
|
||||
f_tir_gets = []
|
||||
f_tir_sets = []
|
||||
|
||||
# pylint: enable=invalid-name
|
||||
|
||||
|
||||
def set_global_func():
|
||||
global f_clear, f_add_sequence, f_remove_sequence, f_fork_sequence, f_popn
|
||||
global f_begin_forward, f_end_forward, f_get, f_set, f_debug_get
|
||||
global f_tir_gets, f_tir_sets
|
||||
|
||||
f_clear = tvm.get_global_func("vm.builtin.kv_state_clear")
|
||||
f_add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
|
||||
f_remove_sequence = tvm.get_global_func("vm.builtin.kv_state_remove_sequence")
|
||||
f_fork_sequence = tvm.get_global_func("vm.builtin.kv_state_fork_sequence")
|
||||
f_popn = tvm.get_global_func("vm.builtin.kv_state_popn")
|
||||
f_begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
|
||||
f_end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward")
|
||||
f_get = tvm.get_global_func("vm.builtin.rnn_state_get")
|
||||
f_set = tvm.get_global_func("vm.builtin.rnn_state_set")
|
||||
f_debug_get = tvm.get_global_func("vm.builtin.rnn_state_debug_get")
|
||||
|
||||
target = tvm.target.Target("cuda")
|
||||
|
||||
def _build(tir_func):
|
||||
mod = tvm.IRModule({"main": tir_func})
|
||||
with target:
|
||||
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod) # pylint: disable=not-callable
|
||||
f = tvm.tirx.build(mod["main"], target=target)
|
||||
return f.main
|
||||
|
||||
_f_tir_gets, _f_tir_sets = [], []
|
||||
for state in states:
|
||||
shape, dtype = state
|
||||
_f_tir_gets.append(_build(rnn_state_get(shape, dtype)))
|
||||
_f_tir_sets.append(_build(rnn_state_set(shape, dtype)))
|
||||
|
||||
f_tir_gets = _f_tir_gets
|
||||
f_tir_sets = _f_tir_sets
|
||||
|
||||
|
||||
def create_rnn_state(device):
|
||||
f_create = tvm.get_global_func("vm.builtin.rnn_state_create")
|
||||
init_values = [
|
||||
tvm.runtime.tensor(np_zero, device=device),
|
||||
tvm.runtime.tensor(np_one, device=device),
|
||||
]
|
||||
return f_create(num_layers, reserved_nseq, max_history, f_tir_gets, f_tir_sets, init_values)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rnn_state():
|
||||
set_global_func()
|
||||
return create_rnn_state
|
||||
|
||||
|
||||
def verify_state(state, seq_ids, expected_values):
|
||||
layer_id = 0
|
||||
for seq_id in seq_ids:
|
||||
for state_id, expected_value in enumerate(expected_values[seq_id]):
|
||||
state_value = f_debug_get(state, layer_id, state_id, seq_id)
|
||||
tvm.testing.assert_allclose(state_value.numpy(), expected_value)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_rnn_state_get(rnn_state): # pylint: disable=redefined-outer-name
|
||||
def run_and_check():
|
||||
device = tvm.cuda()
|
||||
state = rnn_state(device)
|
||||
f_clear(state)
|
||||
f_add_sequence(state, 0)
|
||||
f_begin_forward(state, Shape([0]), Shape([1]))
|
||||
tvm_nd_0 = tvm.runtime.tensor(np.empty((1, 16, 16), "float16"), device=device)
|
||||
tvm_nd_1 = tvm.runtime.tensor(np.empty((1, 32, 32), "float32"), device=device)
|
||||
f_get(state, 0, 0, tvm_nd_0)
|
||||
f_get(state, 0, 1, tvm_nd_1)
|
||||
f_end_forward(state)
|
||||
tvm.testing.assert_allclose(tvm_nd_0.numpy(), np.zeros((1, 16, 16), "float16"))
|
||||
tvm.testing.assert_allclose(tvm_nd_1.numpy(), np.ones((1, 32, 32), "float32"))
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_rnn_state_set(rnn_state): # pylint: disable=redefined-outer-name
|
||||
def run_and_check():
|
||||
device = tvm.cuda()
|
||||
state = rnn_state(device)
|
||||
f_clear(state)
|
||||
for seq_id in range(3):
|
||||
f_add_sequence(state, seq_id)
|
||||
f_begin_forward(state, Shape([0, 2]), Shape([1, 1]))
|
||||
f_set(
|
||||
state,
|
||||
0,
|
||||
0,
|
||||
tvm.runtime.tensor(np.full((2, 16, 16), 2.0, "float16"), device=device),
|
||||
)
|
||||
f_set(
|
||||
state,
|
||||
0,
|
||||
1,
|
||||
tvm.runtime.tensor(np.full((2, 32, 32), 3.0, "float32"), device=device),
|
||||
)
|
||||
f_end_forward(state)
|
||||
expected_values = [[np_two, np_three], [np_zero, np_one], [np_two, np_three]]
|
||||
verify_state(state, [0, 1, 2], expected_values)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_rnn_state_popn(rnn_state): # pylint: disable=redefined-outer-name
|
||||
def run_and_check():
|
||||
device = tvm.cuda()
|
||||
state = rnn_state(device)
|
||||
f_clear(state)
|
||||
f_add_sequence(state, 0)
|
||||
f_begin_forward(state, Shape([0]), Shape([1]))
|
||||
f_set(state, 0, 0, tvm.runtime.tensor(np_two.reshape(1, 16, 16), device=device))
|
||||
f_set(state, 0, 1, tvm.runtime.tensor(np_three.reshape(1, 32, 32), device=device))
|
||||
f_end_forward(state)
|
||||
verify_state(state, [0], [[np_two, np_three]])
|
||||
f_popn(state, 0, 1)
|
||||
verify_state(state, [0], [[np_zero, np_one]])
|
||||
with pytest.raises(RuntimeError):
|
||||
f_popn(state, 0, 1)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_rnn_state_fork_sequence(rnn_state): # pylint: disable=redefined-outer-name
|
||||
def run_and_check():
|
||||
device = tvm.cuda()
|
||||
state = rnn_state(device)
|
||||
f_clear(state)
|
||||
f_add_sequence(state, 0)
|
||||
f_begin_forward(state, Shape([0]), Shape([1]))
|
||||
f_set(state, 0, 0, tvm.runtime.tensor(np_two.reshape(1, 16, 16), device=device))
|
||||
f_set(state, 0, 1, tvm.runtime.tensor(np_three.reshape(1, 32, 32), device=device))
|
||||
f_end_forward(state)
|
||||
f_fork_sequence(state, 0, 1, -1)
|
||||
verify_state(state, [0, 1], [[np_two, np_three], [np_two, np_three]])
|
||||
f_popn(state, 1, 1)
|
||||
verify_state(state, [0, 1], [[np_two, np_three], [np_zero, np_one]])
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def rnn_state_get(
|
||||
shape: Sequence[int],
|
||||
dtype: str,
|
||||
):
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def _rnn_state_get(
|
||||
var_storage: T.handle,
|
||||
var_seq_slot_ids: T.handle,
|
||||
var_history_slot_ids: T.handle,
|
||||
var_output: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
|
||||
storage = T.match_buffer(var_storage, (reserved_nseq, max_history, *shape), dtype)
|
||||
seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32")
|
||||
history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32")
|
||||
output = T.match_buffer(var_output, (batch_size, *shape), dtype)
|
||||
|
||||
for i in range(batch_size):
|
||||
for s in T.grid(*shape):
|
||||
with T.sblock("copy"):
|
||||
vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s])
|
||||
seq_id: T.let[T.int32] = seq_slot_ids[vi]
|
||||
history_id: T.let[T.int32] = history_slot_ids[vi]
|
||||
# The following line is equivalent to:
|
||||
# `output[vi, *vs] = storage[seq_id, history_id, *vs]`
|
||||
# However, unpacking operator in subscript requires Python 3.11 or newer
|
||||
T.buffer_store(
|
||||
output, T.BufferLoad(storage, [seq_id, history_id, *vs]), [vi, *vs]
|
||||
)
|
||||
# fmt: on
|
||||
return _rnn_state_get
|
||||
|
||||
|
||||
def rnn_state_set(
|
||||
shape: Sequence[int | tirx.Var],
|
||||
dtype: str,
|
||||
):
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def _rnn_state_set(
|
||||
var_storage: T.handle,
|
||||
var_seq_slot_ids: T.handle,
|
||||
var_history_slot_ids: T.handle,
|
||||
var_data: T.handle,
|
||||
):
|
||||
batch_size = T.int32()
|
||||
|
||||
storage = T.match_buffer(var_storage, (reserved_nseq, max_history, *shape), dtype)
|
||||
seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32")
|
||||
history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32")
|
||||
data = T.match_buffer(var_data, (batch_size, *shape), dtype)
|
||||
|
||||
for i in range(batch_size):
|
||||
for s in T.grid(*shape):
|
||||
with T.sblock("copy"):
|
||||
vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s])
|
||||
seq_id: T.let[T.int32] = seq_slot_ids[vi]
|
||||
history_id: T.let[T.int32] = (history_slot_ids[vi] + 1) % T.cast(
|
||||
max_history, "int32"
|
||||
)
|
||||
# The following line is equivalent to:
|
||||
# `storage[seq_id, history_id, *vs] = data[vi, *vs]`
|
||||
# However, unpacking operator in subscript requires Python 3.11 or newer
|
||||
T.buffer_store(
|
||||
storage, T.BufferLoad(data, [vi, *vs]), [seq_id, history_id, *vs]
|
||||
)
|
||||
|
||||
# fmt: on
|
||||
|
||||
return _rnn_state_set
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
set_global_func()
|
||||
rnn_state = create_rnn_state
|
||||
test_rnn_state_get(rnn_state)
|
||||
test_rnn_state_set(rnn_state)
|
||||
test_rnn_state_popn(rnn_state)
|
||||
test_rnn_state_fork_sequence(rnn_state)
|
||||
Reference in New Issue
Block a user