wehub-resource-sync
247 lines
8.8 KiB
Python
247 lines
8.8 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.
|
|
# ruff: noqa: E741
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.script.parser import ir as I
|
|
from tvm.script.parser import relax as R
|
|
from tvm.testing import assert_allclose
|
|
from tvm.testing.utils import check_numerical_grads
|
|
|
|
|
|
def rand(dtype, *shape):
|
|
return tvm.runtime.tensor(np.random.rand(*shape).astype(dtype))
|
|
|
|
|
|
def _legalize_and_build(mod, target, dev):
|
|
ex = tvm.compile(mod, target)
|
|
vm = relax.VirtualMachine(ex, dev)
|
|
return vm
|
|
|
|
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_manual_gradient():
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
|
|
# The expression computed is sum((2x - 2y) * (y + z))
|
|
# the gradient of x is broadcast_to(2y + 2z, x.shape)
|
|
# the gradient of y is collapse_sum_to((2x - 4y - 2z), y.shape)
|
|
# the gradient of z is collapse_sum_to((2x - 2y), z.shape)
|
|
# the gradient of u is 0
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 5), "float32"),
|
|
y: R.Tensor((5,), "float32"),
|
|
z: R.Tensor((5,), "float32"),
|
|
u: R.Tensor((5,), "float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, x)
|
|
lv2 = R.subtract(lv1, y)
|
|
lv3 = R.subtract(lv2, y)
|
|
lv4 = R.add(y, z)
|
|
lv5 = R.multiply(lv3, lv4)
|
|
lv6 = R.sum(lv5)
|
|
R.output(lv6)
|
|
return lv6
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
|
|
args = [rand("float32", 3, 5), rand("float32", 5), rand("float32", 5), rand("float32", 5)]
|
|
args_np = [x.numpy() for x in args]
|
|
|
|
vm = _legalize_and_build(After, target, dev)
|
|
output, grads = vm["main_adjoint"](*args)
|
|
output_np = np.sum((2 * args_np[0] - 2 * args_np[1]) * (args_np[1] + args_np[2]))
|
|
assert_allclose(output.numpy(), output_np, atol=1e-4)
|
|
|
|
expected_grads_nd = [
|
|
(2 * args_np[1] + 2 * args_np[2]) * np.ones_like(args_np[0]),
|
|
np.sum((2 * args_np[0] - 4 * args_np[1] - 2 * args_np[2]), axis=0),
|
|
np.sum((2 * args_np[0] - 2 * args_np[1]), axis=0),
|
|
np.zeros_like(args_np[3]),
|
|
]
|
|
for i, j in zip(grads, expected_grads_nd):
|
|
assert_allclose(i.numpy(), j, atol=1e-4)
|
|
|
|
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_mlp_blockbuilder():
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
layers, in_size, out_size, hidden_size, batch_size = 3, 5, 5, 5, 4
|
|
|
|
input_list = [relax.Var("x", R.Tensor((batch_size, in_size), "float32"))]
|
|
w_list = (
|
|
[relax.Var("w_0", R.Tensor((in_size, hidden_size), "float32"))]
|
|
+ [
|
|
relax.Var("w_" + str(i + 1), R.Tensor((hidden_size, hidden_size), "float32"))
|
|
for i in range(layers - 2)
|
|
]
|
|
+ [relax.Var("w_" + str(layers - 1), R.Tensor((hidden_size, out_size), "float32"))]
|
|
)
|
|
b_list = [
|
|
relax.Var("b_" + str(i), R.Tensor((hidden_size,), "float32")) for i in range(layers - 1)
|
|
] + [relax.Var("b_" + str(layers - 1), R.Tensor((out_size,), "float32"))]
|
|
label_list = [relax.Var("y", R.Tensor((batch_size,), "int64"))]
|
|
args_list = input_list + w_list + b_list + label_list
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("MLP", args_list):
|
|
with bb.dataflow():
|
|
current = input_list[0]
|
|
for i in range(layers):
|
|
lv0 = bb.emit(R.matmul(current, w_list[i]))
|
|
lv1 = bb.emit(R.add(lv0, b_list[i]))
|
|
current = bb.emit(R.nn.relu(lv1) if i < layers - 1 else lv1)
|
|
logits = R.nn.log_softmax(current)
|
|
loss = bb.emit(R.nn.nll_loss(logits, label_list[0]))
|
|
gv0 = bb.emit_output(loss)
|
|
bb.emit_func_output(gv0)
|
|
|
|
Before = bb.get()
|
|
After = relax.transform.Gradient("MLP", w_list + b_list)(Before)
|
|
# Check numerical gradients equal
|
|
args = []
|
|
for arg in After["MLP_adjoint"].params:
|
|
shape = [int(l) for l in arg.ty.shape]
|
|
if arg.ty.dtype == "int64":
|
|
args.append(
|
|
tvm.runtime.tensor(np.random.randint(0, out_size, size=shape).astype(np.int64))
|
|
)
|
|
else: # float32
|
|
args.append(rand("float32", *shape))
|
|
|
|
vm_before = _legalize_and_build(Before, target, dev)
|
|
vm_after = _legalize_and_build(After, target, dev)
|
|
_, grad = vm_after["MLP_adjoint"](*args)
|
|
|
|
def func(*inputs):
|
|
loss = vm_before["MLP"](args[0], *[tvm.runtime.tensor(i) for i in inputs], args[-1])
|
|
return loss.numpy()
|
|
|
|
check_numerical_grads(func, [i.numpy() for i in args[1:-1]], [i.numpy() for i in grad])
|
|
|
|
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_complex():
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
cst = relax.const(np.ones((6,)), dtype="float32")
|
|
cst1 = relax.const(np.array(3), dtype="int64")
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((6,), "float32"), y: R.Tensor((6, 3, 4), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.split(x, 2)
|
|
lv2 = lv1[0]
|
|
lv3 = lv1[1]
|
|
lv4 = lv2 + lv3
|
|
lv5 = (lv4, lv3)
|
|
lv6 = R.concat(lv5)
|
|
lv7 = (x, x)
|
|
lv8 = R.concat(lv7)
|
|
lv9 = R.concat(lv7)
|
|
lv10 = R.add(lv8, lv9)
|
|
lv11 = R.split(lv10, 2)
|
|
lv12 = R.add(lv6, lv11[0])
|
|
lv13 = cst
|
|
lv14 = R.add(lv12, lv13)
|
|
lv15 = R.subtract(lv13, lv14)
|
|
lv16 = R.multiply(lv14, lv15)
|
|
lv17 = R.multiply(lv15, lv16)
|
|
lv18 = R.tanh(lv17)
|
|
lv19 = R.sigmoid(lv18)
|
|
lv20 = R.permute_dims(y, axes=[0, 2, 1])
|
|
lv21 = R.sigmoid(lv20)
|
|
lv22 = R.matmul(y, lv21)
|
|
lv23 = R.sum(lv22, axis=[1, 2])
|
|
lv24 = R.add(lv19, lv23)
|
|
lv25 = R.nn.log_softmax(lv24)
|
|
gv = R.nn.nll_loss(lv25, cst1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
args = []
|
|
for arg in After["main_adjoint"].params:
|
|
shape = [int(l) for l in arg.ty.shape]
|
|
args.append(rand("float32", *shape))
|
|
|
|
vm_before = _legalize_and_build(Before, target, dev)
|
|
vm_after = _legalize_and_build(After, target, dev)
|
|
_, grad = vm_after["main_adjoint"](*args)
|
|
|
|
def func(*inputs):
|
|
loss = vm_before["main"](*[tvm.runtime.tensor(i) for i in inputs])
|
|
return loss.numpy()
|
|
|
|
check_numerical_grads(func, [i.numpy() for i in args], [i.numpy() for i in grad])
|
|
|
|
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_matmul():
|
|
target = "llvm"
|
|
dev = tvm.device(target)
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), "float32"), y: R.Tensor((3, 3), "float32")):
|
|
with R.dataflow():
|
|
lv1 = R.matmul(x, y)
|
|
lv2 = R.permute_dims(x)
|
|
lv3 = R.matmul(lv2, y)
|
|
lv4 = R.permute_dims(y)
|
|
lv5 = R.matmul(x, lv4)
|
|
lv6 = R.permute_dims(x)
|
|
lv7 = R.permute_dims(y)
|
|
lv8 = R.matmul(lv6, lv7)
|
|
lv9 = lv1 + lv3 + lv5 + lv8
|
|
gv = R.sum(lv9)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
After = relax.transform.Gradient("main")(Before)
|
|
args = []
|
|
for arg in After["main_adjoint"].params:
|
|
shape = [int(l) for l in arg.ty.shape]
|
|
args.append(rand("float32", *shape))
|
|
|
|
vm_before = _legalize_and_build(Before, target, dev)
|
|
vm_after = _legalize_and_build(After, target, dev)
|
|
_, grad = vm_after["main_adjoint"](*args)
|
|
|
|
def func(*inputs):
|
|
loss = vm_before["main"](*[tvm.runtime.tensor(i) for i in inputs])
|
|
return loss.numpy()
|
|
|
|
check_numerical_grads(func, [i.numpy() for i in args], [i.numpy() for i in grad])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|