wehub-resource-sync
9452 lines
348 KiB
Python
9452 lines
348 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: E501, F401, F841
|
|
import operator
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
from torch.export import export
|
|
from torch.nn import Module
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.relax.frontend.torch import from_exported_program
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
def verify_model(
|
|
torch_model,
|
|
example_args,
|
|
binding,
|
|
expected,
|
|
dynamic_shapes=None,
|
|
run_ep_decomposition=True,
|
|
keep_params_as_input=False,
|
|
unwrap_unit_return_tuple=False,
|
|
no_bind_return_tuple=False,
|
|
map_free_vars=False,
|
|
custom_convert_map=None,
|
|
):
|
|
exported_program = export(torch_model, args=example_args, dynamic_shapes=dynamic_shapes)
|
|
mod = from_exported_program(
|
|
exported_program,
|
|
run_ep_decomposition=run_ep_decomposition,
|
|
keep_params_as_input=keep_params_as_input,
|
|
unwrap_unit_return_tuple=unwrap_unit_return_tuple,
|
|
no_bind_return_tuple=no_bind_return_tuple,
|
|
custom_convert_map=custom_convert_map,
|
|
)
|
|
|
|
binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()}
|
|
expected = relax.transform.BindParams("main", binding)(expected)
|
|
tvm.ir.assert_structural_equal(mod, expected, map_free_vars=map_free_vars)
|
|
|
|
|
|
def verify_model_numerically(torch_model, example_args, rtol=1e-7, atol=1e-7):
|
|
"""Verify model by comparing numerical outputs between PyTorch and TVM."""
|
|
with torch.no_grad():
|
|
pytorch_output = torch_model(*example_args)
|
|
|
|
exported_program = export(torch_model, args=example_args)
|
|
mod = from_exported_program(exported_program)
|
|
target = tvm.target.Target("llvm")
|
|
ex = relax.build(mod, target)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
|
|
tvm_args = [tvm.runtime.tensor(arg.numpy()) for arg in example_args]
|
|
tvm_output = vm["main"](*tvm_args)
|
|
|
|
if hasattr(tvm_output, "numpy"):
|
|
tvm_output_np = tvm_output.numpy()
|
|
else:
|
|
tvm_output_np = tvm_output[0].numpy()
|
|
|
|
pytorch_output_np = (
|
|
pytorch_output.numpy()
|
|
if isinstance(pytorch_output, torch.Tensor)
|
|
else pytorch_output[0].numpy()
|
|
)
|
|
|
|
assert pytorch_output_np.shape == tvm_output_np.shape, (
|
|
f"Shape mismatch: PyTorch {pytorch_output_np.shape} vs TVM {tvm_output_np.shape}"
|
|
)
|
|
tvm.testing.assert_allclose(pytorch_output_np, tvm_output_np, rtol=rtol, atol=atol)
|
|
|
|
|
|
operator_basic_unary = [
|
|
(torch.abs, R.abs),
|
|
(torch.acos, R.acos),
|
|
(torch.acosh, R.acosh),
|
|
(torch.asin, R.asin),
|
|
(torch.asinh, R.asinh),
|
|
(torch.atan, R.atan),
|
|
(torch.atanh, R.atanh),
|
|
(torch.bitwise_not, R.bitwise_not),
|
|
(torch.ceil, R.ceil),
|
|
(torch.cos, R.cos),
|
|
(torch.cosh, R.cosh),
|
|
(torch.erf, R.erf),
|
|
(torch.exp, R.exp),
|
|
(torch.floor, R.floor),
|
|
(torch.ops.aten.gelu, R.nn.gelu),
|
|
(torch.log, R.log),
|
|
(torch.neg, R.negative),
|
|
(torch.relu, R.nn.relu),
|
|
(torch.round, R.round),
|
|
(torch.rsqrt, R.rsqrt),
|
|
(torch.sigmoid, R.sigmoid),
|
|
(torch.sin, R.sin),
|
|
(torch.sinh, R.sinh),
|
|
(torch.sign, R.sign),
|
|
(torch.sqrt, R.sqrt),
|
|
(torch.tan, R.tan),
|
|
(torch.tanh, R.tanh),
|
|
(torch.trunc, R.trunc),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("pytorch_op, relax_op", operator_basic_unary)
|
|
def test_basic_unary_ops(pytorch_op, relax_op):
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
class UnaryOp(Module):
|
|
def forward(self, input):
|
|
return pytorch_op(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(input_1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(UnaryOp(), example_args, {}, expected)
|
|
|
|
|
|
operator_bool_unary = [
|
|
(torch.isinf, R.isinf),
|
|
(torch.isnan, R.isnan),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("pytorch_op, relax_op", operator_bool_unary)
|
|
def test_bool_unary_ops(pytorch_op, relax_op):
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
class UnaryOp(Module):
|
|
def forward(self, input):
|
|
return pytorch_op(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="bool")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(input_1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(UnaryOp(), example_args, {}, expected)
|
|
|
|
|
|
def test_sqrt_integer_input():
|
|
"""Test that sqrt operation works with integer tensors by auto-converting to float."""
|
|
example_args = (torch.tensor([[4, 9, 16, 25]], dtype=torch.int64),)
|
|
|
|
class SqrtIntModel(Module):
|
|
def forward(self, input):
|
|
return torch.sqrt(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_int64:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 4), dtype="int64")) -> R.Tuple(
|
|
R.Tensor((1, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4), dtype="float32") = R.astype(input_1, dtype="float32")
|
|
lv1: R.Tensor((1, 4), dtype="float32") = R.sqrt(lv)
|
|
gv: R.Tuple(R.Tensor((1, 4), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(SqrtIntModel(), example_args, {}, expected_int64)
|
|
|
|
example_args_int32 = (torch.tensor([[1, 4, 9]], dtype=torch.int32),)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_int32:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3), dtype="int32")) -> R.Tuple(
|
|
R.Tensor((1, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3), dtype="float32") = R.astype(input_1, dtype="float32")
|
|
lv1: R.Tensor((1, 3), dtype="float32") = R.sqrt(lv)
|
|
gv: R.Tuple(R.Tensor((1, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(SqrtIntModel(), example_args_int32, {}, expected_int32)
|
|
|
|
|
|
def test_extended_unary_ops():
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
# celu
|
|
class Celu1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.celu = torch.nn.CELU()
|
|
|
|
def forward(self, input):
|
|
return self.celu(input)
|
|
|
|
class Celu2(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.celu(input)
|
|
|
|
# alpha * min(0, exp(x / alpha) - 1) + max(0, x)
|
|
@tvm.script.ir_module
|
|
class expected_celu:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
|
|
lv, R.const(1.0, "float32")
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater(
|
|
input, R.const(0.0, "float32")
|
|
)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv2, input, lv1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Celu1(), example_args, {}, expected_celu)
|
|
verify_model(Celu2(), example_args, {}, expected_celu)
|
|
|
|
# clamp
|
|
class Clamp(Module):
|
|
def forward(self, input):
|
|
return torch.clamp(input, min=0.1, max=0.5)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_clamp:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
input,
|
|
R.prim_value(T.float64(0.10000000000000001)),
|
|
R.prim_value(T.float64(0.5)),
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Clamp(), example_args, {}, expected_clamp)
|
|
|
|
class ClampMinOnly(Module):
|
|
def forward(self, input):
|
|
return torch.clamp(input, min=0.5, max=None)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_clamp_min_only:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
input, R.prim_value(T.float64(0.5)), R.prim_value(T.float64("inf"))
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(ClampMinOnly(), example_args, {}, expected_clamp_min_only)
|
|
|
|
class ClampTensors(Module):
|
|
def forward(self, input):
|
|
return torch.clamp(input, min=input, max=input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_clamp_tensors:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
|
|
input, R.shape([1, 3, 10, 10])
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.maximum(input, lv)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
|
|
input, R.shape([1, 3, 10, 10])
|
|
)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.minimum(lv1, lv2)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv3, R.prim_value(T.float64("-inf")), R.prim_value(T.float64("inf"))
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(ClampTensors(), example_args, {}, expected_clamp_tensors)
|
|
|
|
# dropout
|
|
|
|
class Dropout1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.dropout = torch.nn.Dropout(0.5)
|
|
|
|
def forward(self, input):
|
|
return self.dropout(input)
|
|
|
|
class Dropout2(Module):
|
|
def forward(self, input):
|
|
return torch.dropout(input, 0.5, train=True)
|
|
|
|
class Dropout3(Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.dropout_(input, 0.5, train=True)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_dropout_for_1_2:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (input,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_dropout_for_3:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.zeros(
|
|
R.shape([1, 3, 10, 10]), dtype="float32"
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv, R.const(0.5, "float32")
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(input, lv1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Dropout1(), example_args, {}, expected_dropout_for_1_2)
|
|
verify_model(Dropout2(), example_args, {}, expected_dropout_for_1_2)
|
|
verify_model(Dropout3(), example_args, {}, expected_dropout_for_3)
|
|
|
|
# elu
|
|
class Elu(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.elu = torch.nn.ELU()
|
|
|
|
def forward(self, input):
|
|
return self.elu(input)
|
|
|
|
class Elu2(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.elu(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_elu:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
|
|
R.const(1.0, "float32"), lv
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(lv1)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
|
|
R.const(-1.0, "float32"), lv2
|
|
)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input)
|
|
lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv3, lv4)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Elu(), example_args, {}, expected_elu)
|
|
verify_model(Elu2(), example_args, {}, expected_elu)
|
|
|
|
# hardsigmoid
|
|
class Hardsigmoid(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.hs = torch.nn.Hardsigmoid()
|
|
|
|
def forward(self, input):
|
|
return self.hs(input)
|
|
|
|
class Hardsigmoid2(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.hardsigmoid(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_hardsigmoid:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(
|
|
inp_0, R.const(3.0, "float32")
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv, R.prim_value(0), R.prim_value(T.float64("inf"))
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv1, R.prim_value(T.float64("-inf")), R.prim_value(6)
|
|
)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv2, R.const(6.0, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Hardsigmoid(), example_args, {}, expected_hardsigmoid)
|
|
verify_model(Hardsigmoid2(), example_args, {}, expected_hardsigmoid)
|
|
|
|
# hardwish
|
|
class Hardswish(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.hs = torch.nn.Hardswish()
|
|
|
|
def forward(self, input):
|
|
return self.hs(input)
|
|
|
|
class Hardswish2(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.hardswish(input)
|
|
|
|
class Hardswish3(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.hardswish_(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_hardswish_for_1_2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(
|
|
inp_0, R.const(3.0, "float32")
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv, R.prim_value(0), R.prim_value(T.float64("inf"))
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv1, R.prim_value(T.float64("-inf")), R.prim_value(6)
|
|
)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(inp_0, lv2)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv3, R.const(6.0, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_hardswish_for_3:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(
|
|
input, R.const(3.0, "float32")
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv, R.prim_value(0), R.prim_value(T.float64("inf"))
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
lv1, R.prim_value(T.float64("-inf")), R.prim_value(6)
|
|
)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(input, lv2)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv3, R.const(6.0, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Hardswish(), example_args, {}, expected_hardswish_for_1_2)
|
|
verify_model(Hardswish2(), example_args, {}, expected_hardswish_for_1_2)
|
|
verify_model(Hardswish3(), example_args, {}, expected_hardswish_for_3)
|
|
|
|
# isfinite
|
|
class IsFinite(Module):
|
|
def forward(self, input):
|
|
return torch.isfinite(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_isfinite:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="bool")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.abs(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.not_equal(
|
|
lv, R.const(float("inf"), "float32")
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.equal(input, input)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="bool") = R.multiply(lv2, lv1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(IsFinite(), example_args, {}, expected_isfinite)
|
|
|
|
# log2
|
|
class Log2(Module):
|
|
def forward(self, x):
|
|
return torch.log2(x)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected_log2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(inp_0)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv, R.const(0.69314718246459961, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Log2(), example_args, {}, Expected_log2)
|
|
|
|
# log10
|
|
class Log10(Module):
|
|
def forward(self, x):
|
|
return torch.log10(x)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected_log10:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(inp_0)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv, R.const(2.302585092994046, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Log10(), example_args, {}, Expected_log10)
|
|
|
|
# log1p
|
|
class Log1p(Module):
|
|
def forward(self, x):
|
|
return torch.log1p(x)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected_log1p:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(
|
|
R.add(inp_0, R.const(1, "float32"))
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Log1p(), example_args, {}, Expected_log1p)
|
|
|
|
# reciprocal
|
|
class Reciprocal(Module):
|
|
def forward(self, input):
|
|
return torch.reciprocal(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_reciprocal:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
R.const(1.0, "float32"), input_1
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Reciprocal(), example_args, {}, expected_reciprocal)
|
|
|
|
# Returns the maximum value of all elements in the input tensor.
|
|
class MaxModel(Module):
|
|
def forward(self, input):
|
|
return torch.max(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_max:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.max(input, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(MaxModel(), example_args, {}, expected_max)
|
|
|
|
# Returns the minimum value of all elements in the input tensor.
|
|
class MinModel(Module):
|
|
def forward(self, input):
|
|
return torch.min(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_min:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.min(input, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(MinModel(), example_args, {}, expected_min)
|
|
|
|
# relu6
|
|
class ReLU6_1(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.relu6 = torch.nn.ReLU6()
|
|
|
|
def forward(self, x):
|
|
return self.relu6(x)
|
|
|
|
class ReLU6_2(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.nn.functional.relu6(x)
|
|
|
|
class ReLU6_3(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.relu6_(x)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_relu6_1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
x, R.prim_value(T.float64(0.0)), R.prim_value(T.float64(6.0))
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_relu6_2:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu6(x)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_relu6_3:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
x, R.prim_value(0), R.prim_value(6)
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(ReLU6_1(), example_args, {}, expected_relu6_1)
|
|
verify_model(ReLU6_2(), example_args, {}, expected_relu6_2)
|
|
verify_model(ReLU6_3(), example_args, {}, expected_relu6_3)
|
|
|
|
# selu
|
|
class SELU(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.selu(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_selu:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
|
|
R.const(1.0, "float32"), lv
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(lv1)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
|
|
R.const(-1.6732631921768188, "float32"), lv2
|
|
)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input)
|
|
lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv3, lv4)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(SELU(), example_args, {}, expected_selu)
|
|
|
|
# silu
|
|
class SiLU(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.silu(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_silu:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sigmoid(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(input, lv)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(SiLU(), example_args, {}, expected_silu)
|
|
|
|
# silu_
|
|
class SiLU_(Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.silu_(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_silu_:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sigmoid(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(input, lv)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(SiLU_(), example_args, {}, expected_silu_)
|
|
|
|
# square
|
|
class Square(Module):
|
|
def forward(self, input):
|
|
return torch.square(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_square:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.power(
|
|
input, R.const(2.0, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Square(), example_args, {}, expected_square)
|
|
|
|
# relu_
|
|
class ReLU_(Module):
|
|
def forward(self, input):
|
|
return torch.relu_(input.clone())
|
|
|
|
@tvm.script.ir_module
|
|
class expected_relu_:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(ReLU_(), example_args, {}, expected_relu_)
|
|
|
|
|
|
def test_hardtanh():
|
|
class Hardtanh(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.ht = torch.nn.Hardtanh()
|
|
|
|
def forward(self, input):
|
|
return self.ht(input)
|
|
|
|
class Hardtanh2(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.hardtanh(input)
|
|
|
|
class Hardtanh3(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.hardtanh_(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_for_1_2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
|
|
inp_0, R.prim_value(T.float64(-1.0)), R.prim_value(T.float64(1.0))
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Hardtanh(), example_args, {}, expected_for_1_2)
|
|
verify_model(Hardtanh2(), example_args, {}, expected_for_1_2)
|
|
# In-place hardtanh_ yields the same program; mutation outputs are dropped.
|
|
verify_model(Hardtanh3(), example_args, {}, expected_for_1_2)
|
|
|
|
|
|
def test_softplus():
|
|
import torch
|
|
from torch.nn import Module
|
|
|
|
torch.set_grad_enabled(False)
|
|
|
|
class Softplus0(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.softplus = torch.nn.Softplus(1.0, 20.0)
|
|
|
|
def forward(self, x):
|
|
return self.softplus(x)
|
|
|
|
class Softplus1(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.softplus(input, 1.0, 20.0)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
|
|
x, R.const(1.0, "float32")
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(lv)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv1, R.const(1.0, "float32"))
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(lv2)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
|
|
lv3, R.const(1.0, "float32")
|
|
)
|
|
lv5: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater(
|
|
lv, R.const(20.0, "float32")
|
|
)
|
|
lv6: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv5, x, lv4)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv6,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Softplus0(), example_args, {}, expected)
|
|
verify_model(Softplus1(), example_args, {}, expected)
|
|
|
|
|
|
def test_leakyrelu():
|
|
import torch
|
|
from torch.nn import Module
|
|
|
|
torch.set_grad_enabled(False)
|
|
|
|
class LeakyReLU0(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.leakyrelu = torch.nn.LeakyReLU(0.02)
|
|
|
|
def forward(self, input):
|
|
return self.leakyrelu(input)
|
|
|
|
class LeakyReLU1(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.leaky_relu(input, 0.02)
|
|
|
|
class LeakyReLU2(Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.leaky_relu_(input, 0.02)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_for_1_2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.leakyrelu(input_1, alpha=0.02)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(LeakyReLU0(), example_args, {}, expected_for_1_2)
|
|
verify_model(LeakyReLU1(), example_args, {}, expected_for_1_2)
|
|
# In-place leaky_relu_ yields the same program; mutation outputs are dropped.
|
|
verify_model(LeakyReLU2(), example_args, {}, expected_for_1_2)
|
|
|
|
|
|
def test_logaddexp():
|
|
class LogAddExp(Module):
|
|
def forward(self, input1, input2):
|
|
return torch.logaddexp(input1, input2)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
input1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
input2: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater_equal(input1, input2)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv, input1, input2)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv, input2, input1)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.abs(input1)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="bool") = R.not_equal(
|
|
lv3, R.const(float("inf"), "float32")
|
|
)
|
|
lv5: R.Tensor((1, 3, 10, 10), dtype="bool") = R.equal(input1, input1)
|
|
lv6: R.Tensor((1, 3, 10, 10), dtype="bool") = R.multiply(lv5, lv4)
|
|
lv7: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_not(lv6)
|
|
lv8: R.Tensor((1, 3, 10, 10), dtype="bool") = R.equal(input1, input2)
|
|
lv9: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv7, lv8)
|
|
lv10: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(lv2, lv1)
|
|
lv11: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(lv10)
|
|
lv12: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(
|
|
lv11, R.const(1.0, "float32")
|
|
)
|
|
lv13: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(lv12)
|
|
lv14: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv1, lv13)
|
|
lv15: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv9, input1, lv14)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv15,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
)
|
|
verify_model(LogAddExp(), example_args, {}, expected)
|
|
|
|
|
|
def test_atan2():
|
|
class Atan2(Module):
|
|
def forward(self, lhs, rhs):
|
|
return torch.atan2(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.atan2(lhs, rhs)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
)
|
|
verify_model(Atan2(), example_args, {}, expected)
|
|
|
|
|
|
def test_logical_and():
|
|
class LogicalAnd(Module):
|
|
def forward(self, lhs, rhs):
|
|
return torch.logical_and(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
)
|
|
verify_model(LogicalAnd(), example_args, {}, expected)
|
|
|
|
|
|
def test_logical_not():
|
|
class LogicalNot(Module):
|
|
def forward(self, input):
|
|
return torch.logical_not(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="bool")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(input, dtype="bool")
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_not(lv)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(LogicalNot(), example_args, {}, expected)
|
|
|
|
|
|
def test_logical_or():
|
|
class LogicalOr(Module):
|
|
def forward(self, lhs, rhs):
|
|
return torch.logical_or(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
)
|
|
verify_model(LogicalOr(), example_args, {}, expected)
|
|
|
|
|
|
def test_logical_xor():
|
|
class LogicalXor(Module):
|
|
def forward(self, lhs, rhs):
|
|
return torch.logical_xor(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="bool")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
torch.randn(1, 3, 10, 10, dtype=torch.float32),
|
|
)
|
|
verify_model(LogicalXor(), example_args, {}, expected)
|
|
|
|
|
|
def test_pow_integer():
|
|
class Pow(Module):
|
|
def forward(self, input):
|
|
return input.pow(4)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(input: R.Tensor((4,), dtype="int64")) -> R.Tuple(R.Tensor((4,), dtype="int64")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((4,), dtype="int64") = R.multiply(input, input)
|
|
lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, input)
|
|
lv2: R.Tensor((4,), dtype="int64") = R.multiply(lv1, input)
|
|
gv: R.Tuple(R.Tensor((4,), dtype="int64")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.tensor([-1, 1, 2, 3], dtype=torch.int64),)
|
|
verify_model(Pow(), example_args, {}, expected)
|
|
|
|
|
|
def test_logsoftmax():
|
|
class LogSoftmax(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.lsm = torch.nn.LogSoftmax(dim=1)
|
|
|
|
def forward(self, input):
|
|
return self.lsm(input)
|
|
|
|
class LogSoftmax2(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.log_softmax(input, dim=1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.log_softmax(input_1, axis=1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(LogSoftmax(), example_args, {}, expected1)
|
|
verify_model(LogSoftmax2(), example_args, {}, expected1)
|
|
|
|
|
|
def test_prelu():
|
|
class Prelu1(Module):
|
|
def __init__(self, num_parameters=1, alpha=0.25):
|
|
super().__init__()
|
|
self.prelu = torch.nn.PReLU(num_parameters=num_parameters, init=alpha)
|
|
|
|
def forward(self, x):
|
|
return self.prelu(x)
|
|
|
|
class Prelu2(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.alpha = torch.nn.Parameter(torch.tensor([0.25]))
|
|
|
|
def forward(self, x):
|
|
return torch.nn.functional.prelu(x, self.alpha)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 1, 1, 1), dtype="float32") = R.reshape(
|
|
R.const([0.25], dtype="float32"), R.shape([1, 1, 1, 1])
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater(x, R.const(0.0, "float32"))
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(lv, x)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv1, x, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Prelu1(), example_args, {}, expected)
|
|
verify_model(Prelu2(), example_args, {}, expected)
|
|
|
|
|
|
def test_softmax():
|
|
class Softmax(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.sm = torch.nn.Softmax(dim=1)
|
|
|
|
def forward(self, input):
|
|
return self.sm(input)
|
|
|
|
class Softmax2(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.softmax(input, dim=1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.softmax(input_1, axis=1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Softmax(), example_args, {}, expected1)
|
|
verify_model(Softmax2(), example_args, {}, expected1)
|
|
|
|
|
|
def test_softsign():
|
|
class Softsign(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.ss = torch.nn.Softsign()
|
|
|
|
def forward(self, input):
|
|
return self.ss(input)
|
|
|
|
class Softsign2(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.softsign(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_softsign:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
abs_val = R.abs(input)
|
|
denom = R.add(abs_val, R.const(1.0, "float32"))
|
|
result = R.divide(input, denom)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (result,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Softsign(), example_args, {}, expected_softsign)
|
|
verify_model(Softsign2(), example_args, {}, expected_softsign)
|
|
|
|
|
|
def test_softshrink():
|
|
class Softshrink(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.softshrink = torch.nn.Softshrink(lambd=0.5)
|
|
|
|
def forward(self, input):
|
|
return self.softshrink(input)
|
|
|
|
class Softshrink2(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.softshrink(input, lambd=0.5)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_softshrink:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.abs(input)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater(lv, R.const(0.5, "float32"))
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sign(input)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
|
|
lv2, R.const(0.5, "float32")
|
|
)
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(input, lv3)
|
|
lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
|
|
input, R.const(0.0, "float32")
|
|
)
|
|
lv6: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv1, lv4, lv5)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv6,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Softshrink(), example_args, {}, expected_softshrink)
|
|
verify_model(Softshrink2(), example_args, {}, expected_softshrink)
|
|
|
|
|
|
def test_tril_triu():
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
|
|
class Tril(Module):
|
|
def forward(self, input):
|
|
return torch.tril(input, 1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_tril:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((10,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((1, 10), dtype="int64") = R.expand_dims(lv, axis=[-2])
|
|
lv2: R.Tensor((10,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv2, axis=[-1])
|
|
lv4: R.Tensor((10, 10), dtype="int64") = R.subtract(lv1, lv3)
|
|
lv5: R.Tensor((10, 10), dtype="bool") = R.less_equal(lv4, R.const(1, "int64"))
|
|
lv6: R.Tensor((), dtype="float32") = R.const(0.0, "float32")
|
|
lv7: R.Tensor((10, 10), dtype="float32") = R.where(lv5, input, lv6)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv7,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Tril(), example_args, {}, expected_tril)
|
|
|
|
class Triu(Module):
|
|
def forward(self, input):
|
|
return torch.triu(input, 1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_triu:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((10,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((1, 10), dtype="int64") = R.expand_dims(lv, axis=[-2])
|
|
lv2: R.Tensor((10,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv2, axis=[-1])
|
|
lv4: R.Tensor((10, 10), dtype="int64") = R.subtract(lv1, lv3)
|
|
lv5: R.Tensor((10, 10), dtype="bool") = R.greater_equal(lv4, R.const(1, "int64"))
|
|
lv6: R.Tensor((), dtype="float32") = R.const(0.0, "float32")
|
|
lv7: R.Tensor((10, 10), dtype="float32") = R.where(lv5, input, lv6)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv7,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Triu(), example_args, {}, expected_triu)
|
|
|
|
|
|
operator_binary_1 = [
|
|
(operator.add, R.add),
|
|
(torch.ops.aten.add_, R.add),
|
|
(torch.ops.aten.bitwise_or, R.bitwise_or),
|
|
(torch.ops.aten.bitwise_or_, R.bitwise_or),
|
|
(operator.sub, R.subtract),
|
|
(operator.mul, R.multiply),
|
|
(torch.ops.aten.mul_, R.multiply),
|
|
(operator.truediv, R.divide),
|
|
(operator.floordiv, R.floor_divide),
|
|
(torch.ops.aten.fmod, R.mod),
|
|
(operator.pow, R.power),
|
|
(operator.mod, R.floor_mod),
|
|
(operator.and_, R.bitwise_and),
|
|
(operator.or_, R.bitwise_or),
|
|
(operator.xor, R.bitwise_xor),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("op, relax_op", operator_binary_1)
|
|
def test_binary1(op, relax_op):
|
|
example_args1 = (
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
)
|
|
example_args2 = (torch.randn(10, 10, dtype=torch.float32),)
|
|
|
|
class Binary1(Module):
|
|
def __init__(self, op):
|
|
super().__init__()
|
|
self.op = op
|
|
|
|
def forward(self, lhs, rhs):
|
|
return self.op(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_binary1:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((10, 10), dtype="float32"),
|
|
rhs: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = relax_op(lhs, rhs)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Binary2(Module):
|
|
def __init__(self, op):
|
|
super().__init__()
|
|
self.op = op
|
|
|
|
def forward(self, lhs):
|
|
return self.op(lhs, 1.0)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_binary2:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = relax_op(lhs, R.const(1.0))
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# In-place ops (add_, mul_, ...) produce the same Relax program as their
|
|
# functional counterparts: mutation outputs are dropped by the importer.
|
|
verify_model(Binary1(op), example_args1, {}, expected_binary1)
|
|
verify_model(Binary2(op), example_args2, {}, expected_binary2)
|
|
|
|
|
|
operator_binary_scalar = [
|
|
(torch.ops.aten.add.Scalar, R.add),
|
|
(torch.ops.aten.bitwise_and.Scalar, R.bitwise_and),
|
|
(torch.ops.aten.bitwise_or.Scalar, R.bitwise_or),
|
|
(torch.ops.aten.bitwise_xor.Scalar, R.bitwise_xor),
|
|
(torch.ops.aten.div.Scalar, R.divide),
|
|
(torch.ops.aten.sub.Scalar, R.subtract),
|
|
(torch.ops.aten.mul.Scalar, R.multiply),
|
|
(torch.ops.aten.remainder.Scalar, R.floor_mod),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("op, relax_op", operator_binary_scalar)
|
|
def test_binary_scalar(op, relax_op):
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
class BinaryScalar(Module):
|
|
def __init__(self, op):
|
|
super().__init__()
|
|
self.op = op
|
|
|
|
def forward(self, lhs):
|
|
return self.op(lhs, 1.0)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_binary_scalar:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, R.const(1.0))
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(BinaryScalar(op), example_args, {}, expected_binary_scalar)
|
|
|
|
|
|
operator_binary_promote = [
|
|
(operator.add, R.add),
|
|
(operator.sub, R.subtract),
|
|
(operator.mul, R.multiply),
|
|
(operator.truediv, R.divide),
|
|
(operator.pow, R.power),
|
|
(operator.mod, R.floor_mod),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("op, relax_op", operator_binary_promote)
|
|
def test_binary_dtype_promotion(op, relax_op):
|
|
"""Ensure binary ops promote differing dtypes following PyTorch rules."""
|
|
|
|
class BinaryPromoteLHS(Module):
|
|
def forward(self, x):
|
|
arange_val = torch.arange(x.shape[1]) # int64 by default
|
|
return op(x, arange_val)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_promote_lhs:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(3), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((3,), dtype="float32") = R.astype(lv, dtype="float32")
|
|
lv2: R.Tensor((2, 3), dtype="float32") = relax_op(x, lv1)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class BinaryPromoteRHS(Module):
|
|
def forward(self, x):
|
|
arange_val = torch.arange(x.shape[1]) # int64 by default
|
|
return op(arange_val, x)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_promote_rhs:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(3), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((3,), dtype="float32") = R.astype(lv, dtype="float32")
|
|
lv2: R.Tensor((2, 3), dtype="float32") = relax_op(lv1, x)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 3, dtype=torch.float32),)
|
|
verify_model(BinaryPromoteLHS(), example_args, {}, expected_promote_lhs)
|
|
verify_model(BinaryPromoteRHS(), example_args, {}, expected_promote_rhs)
|
|
|
|
|
|
operator_binary_2 = [
|
|
(operator.eq, R.equal),
|
|
(operator.ne, R.not_equal),
|
|
(operator.lt, R.less),
|
|
(operator.le, R.less_equal),
|
|
(operator.gt, R.greater),
|
|
(operator.ge, R.greater_equal),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("op, relax_op", operator_binary_2)
|
|
def test_binary2(op, relax_op):
|
|
example_args1 = (
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
)
|
|
example_args2 = (torch.randn(10, 10, dtype=torch.float32),)
|
|
|
|
class Binary1(Module):
|
|
def __init__(self, op):
|
|
super().__init__()
|
|
self.op = op
|
|
|
|
def forward(self, lhs, rhs):
|
|
return self.op(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_binary1:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((10, 10), dtype="float32"),
|
|
rhs: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="bool")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="bool") = relax_op(lhs, rhs)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Binary2(Module):
|
|
def __init__(self, op):
|
|
super().__init__()
|
|
self.op = op
|
|
|
|
def forward(self, lhs):
|
|
return self.op(lhs, 1.0)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_binary2:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="bool")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="bool") = relax_op(lhs, R.const(1.0))
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Binary1(op), example_args1, {}, expected_binary1)
|
|
verify_model(Binary2(op), example_args2, {}, expected_binary2)
|
|
|
|
|
|
def test_binary3():
|
|
example_args1 = (
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
)
|
|
example_args2 = (torch.randn(10, 10, dtype=torch.float32),)
|
|
|
|
# Max
|
|
class Max1(Module):
|
|
def forward(self, x, y):
|
|
return torch.max(x, y)
|
|
|
|
@I.ir_module
|
|
class expected_max1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((10, 10), dtype="float32"),
|
|
inp_1: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.maximum(inp_0, inp_1)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Max1(), example_args1, {}, expected_max1)
|
|
|
|
# Min
|
|
class Min1(Module):
|
|
def forward(self, x, y):
|
|
return torch.min(x, y)
|
|
|
|
@I.ir_module
|
|
class expected_min1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((10, 10), dtype="float32"),
|
|
inp_1: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.minimum(inp_0, inp_1)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Min1(), example_args1, {}, expected_min1)
|
|
|
|
# RSub
|
|
class RSub1(Module):
|
|
def forward(self, x, y):
|
|
return torch.rsub(x, y)
|
|
|
|
class RSub2(Module):
|
|
def forward(self, x):
|
|
return torch.rsub(x, 5.0)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_rsub1:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 10), dtype="float32"), y: R.Tensor((10, 10), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.subtract(y, x)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_rsub2:
|
|
@R.function
|
|
def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.subtract(R.const(5.0, "float32"), x)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(RSub1(), example_args1, {}, expected_rsub1)
|
|
verify_model(RSub2(), example_args2, {}, expected_rsub2)
|
|
|
|
|
|
# IsIn
|
|
|
|
|
|
def test_isin():
|
|
class IsInModel(torch.nn.Module):
|
|
def forward(self, x, test_elements):
|
|
return torch.isin(x, test_elements)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((10, 10), dtype="float32"), test_elements: R.Tensor((8,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="bool")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10, 1), dtype="float32") = R.reshape(x, R.shape([10, 10, 1]))
|
|
lv1: R.Tensor((10, 10, 8), dtype="bool") = R.equal(lv, test_elements)
|
|
lv2: R.Tensor((10, 10, 8), dtype="int8") = R.astype(lv1, dtype="int8")
|
|
lv3: R.Tensor((10, 10), dtype="int8") = R.max(lv2, axis=[-1], keepdims=False)
|
|
lv4: R.Tensor((10, 10), dtype="bool") = R.astype(lv3, dtype="bool")
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(8, dtype=torch.float32),
|
|
)
|
|
verify_model(IsInModel(), example_args, {}, expected)
|
|
|
|
|
|
def test_div_mode():
|
|
# Case 1: Basic division (no rounding mode)
|
|
class DivModel(torch.nn.Module):
|
|
def forward(self, a, b):
|
|
return torch.div(a, b)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_div:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((64, 64), dtype="float32"), b: R.Tensor((64,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((64, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((64, 64), dtype="float32") = R.divide(a, b)
|
|
gv: R.Tuple(R.Tensor((64, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(64, 64, dtype=torch.float32),
|
|
torch.randn(64, dtype=torch.float32),
|
|
)
|
|
verify_model(DivModel(), example_args, {}, expected_div)
|
|
|
|
# Case 2: Division with trunc rounding
|
|
class DivTruncModel(torch.nn.Module):
|
|
def forward(self, a, b):
|
|
return torch.div(a, b, rounding_mode="trunc")
|
|
|
|
@tvm.script.ir_module
|
|
class expected_div_trunc:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((64, 64), dtype="float32"), b: R.Tensor((64,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((64, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((64, 64), dtype="float32") = R.divide(a, b)
|
|
lv1: R.Tensor((64, 64), dtype="float32") = R.trunc(lv)
|
|
gv: R.Tuple(R.Tensor((64, 64), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(DivTruncModel(), example_args, {}, expected_div_trunc)
|
|
|
|
# Case 3: Division with floor rounding
|
|
class DivFloorModel(torch.nn.Module):
|
|
def forward(self, a, b):
|
|
return torch.div(a, b, rounding_mode="floor")
|
|
|
|
@tvm.script.ir_module
|
|
class expected_div_floor:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((64, 64), dtype="float32"), b: R.Tensor((64,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((64, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((64, 64), dtype="float32") = R.floor_divide(a, b)
|
|
gv: R.Tuple(R.Tensor((64, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(DivFloorModel(), example_args, {}, expected_div_floor)
|
|
|
|
|
|
def test_batchnorm2d():
|
|
class BatchNorm2d1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.bn = torch.nn.BatchNorm2d(3)
|
|
|
|
def forward(self, input):
|
|
return self.bn(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((3,), dtype="float32"),
|
|
w2: R.Tensor((3,), dtype="float32"),
|
|
w3: R.Tensor((3,), dtype="float32"),
|
|
w4: R.Tensor((3,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
) = R.nn.batch_norm(
|
|
input_1,
|
|
w1,
|
|
w2,
|
|
w3,
|
|
w4,
|
|
axis=1,
|
|
epsilon=1e-05,
|
|
center=True,
|
|
scale=True,
|
|
momentum=0.1,
|
|
training=False,
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class BatchNorm2dCustom(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.bn = torch.nn.BatchNorm2d(3, eps=0.001, momentum=0.01)
|
|
|
|
def forward(self, input):
|
|
return self.bn(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((3,), dtype="float32"),
|
|
w2: R.Tensor((3,), dtype="float32"),
|
|
w3: R.Tensor((3,), dtype="float32"),
|
|
w4: R.Tensor((3,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
) = R.nn.batch_norm(
|
|
input_1,
|
|
w1,
|
|
w2,
|
|
w3,
|
|
w4,
|
|
axis=1,
|
|
epsilon=0.001,
|
|
center=True,
|
|
scale=True,
|
|
momentum=0.01,
|
|
training=False,
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model_1 = BatchNorm2d1().eval()
|
|
binding_1 = {
|
|
"w1": model_1.bn.weight.detach().numpy(),
|
|
"w2": model_1.bn.bias.detach().numpy(),
|
|
"w3": model_1.bn.running_mean.detach().numpy(),
|
|
"w4": model_1.bn.running_var.detach().numpy(),
|
|
}
|
|
verify_model(model_1, example_args, binding_1, expected1)
|
|
|
|
model_2 = BatchNorm2dCustom().eval()
|
|
binding_2 = {
|
|
"w1": model_2.bn.weight.detach().numpy(),
|
|
"w2": model_2.bn.bias.detach().numpy(),
|
|
"w3": model_2.bn.running_mean.detach().numpy(),
|
|
"w4": model_2.bn.running_var.detach().numpy(),
|
|
}
|
|
verify_model(model_2, example_args, binding_2, expected2)
|
|
|
|
class BatchNorm2dTraining(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.bn = torch.nn.BatchNorm2d(3, track_running_stats=True)
|
|
|
|
def forward(self, input):
|
|
return self.bn(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((2, 3, 4, 4), dtype="float32"),
|
|
w1: R.Tensor((3,), dtype="float32"),
|
|
w2: R.Tensor((3,), dtype="float32"),
|
|
w3: R.Tensor((3,), dtype="float32"),
|
|
w4: R.Tensor((3,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 3, 4, 4), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="int64") = R.add(R.const(0, "int64"), R.const(1, "int64"))
|
|
lv1: R.Tuple(
|
|
R.Tensor((2, 3, 4, 4), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
) = R.nn.batch_norm(
|
|
input_1,
|
|
w1,
|
|
w2,
|
|
w3,
|
|
w4,
|
|
axis=1,
|
|
epsilon=1e-5,
|
|
center=True,
|
|
scale=True,
|
|
momentum=0.1,
|
|
training=True,
|
|
)
|
|
lv2: R.Tensor((2, 3, 4, 4), dtype="float32") = lv1[0]
|
|
lv3: R.Tensor((3,), dtype="float32") = lv1[1]
|
|
lv4: R.Tensor((3,), dtype="float32") = R.zeros(R.shape([3]), dtype="float32")
|
|
lv5: R.Tuple(
|
|
R.Tensor((2, 3, 4, 4), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
R.Tensor((3,), dtype="float32"),
|
|
) = (lv2, lv3, lv4, lv4, lv4)
|
|
lv6: R.Tensor((2, 3, 4, 4), dtype="float32") = lv5[0]
|
|
lv7: R.Tensor((3,), dtype="float32") = lv5[3]
|
|
lv8: R.Tensor((3,), dtype="float32") = lv5[4]
|
|
gv: R.Tuple(R.Tensor((2, 3, 4, 4), dtype="float32")) = (lv6,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args_train = (torch.randn(2, 3, 4, 4, dtype=torch.float32),)
|
|
|
|
model_3 = BatchNorm2dTraining()
|
|
model_3.train() # Set to training mode
|
|
binding_3 = {
|
|
"w1": model_3.bn.weight.detach().numpy(),
|
|
"w2": model_3.bn.bias.detach().numpy(),
|
|
"w3": model_3.bn.running_mean.detach().numpy(),
|
|
"w4": model_3.bn.running_var.detach().numpy(),
|
|
}
|
|
verify_model(model_3, example_args_train, binding_3, expected3)
|
|
|
|
|
|
def test_adaptive_avgpool1d():
|
|
class AdaptiveAvgPool1d0(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AdaptiveAvgPool1d(output_size=5)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class AdaptiveAvgPool1d1(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.adaptive_avg_pool1d(input, output_size=5)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 5), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 10), dtype="float32") = R.expand_dims(input_1, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 5), dtype="float32") = R.nn.adaptive_avg_pool2d(
|
|
lv, output_size=[1, 5], layout="NCHW"
|
|
)
|
|
lv2: R.Tensor((1, 3, 5), dtype="float32") = R.squeeze(lv1, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 5), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, dtype=torch.float32),)
|
|
verify_model(AdaptiveAvgPool1d0(), example_args, {}, expected1)
|
|
verify_model(AdaptiveAvgPool1d1(), example_args, {}, expected1)
|
|
|
|
|
|
def test_adaptive_avgpool2d():
|
|
class AdaptiveAvgPool2d0(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AdaptiveAvgPool2d([10, 10])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class AdaptiveAvgPool2d1(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.adaptive_avg_pool2d(input, [10, 10])
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.adaptive_avg_pool2d(
|
|
input_1, output_size=[10, 10], layout="NCHW", out_layout="NCHW"
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(AdaptiveAvgPool2d0(), example_args, {}, expected1)
|
|
verify_model(AdaptiveAvgPool2d1(), example_args, {}, expected1)
|
|
|
|
|
|
def test_adaptive_avgpool3d():
|
|
class AdaptiveAvgPool3d0(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AdaptiveAvgPool3d([4, 4, 4])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class AdaptiveAvgPool3d1(torch.nn.Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.adaptive_avg_pool3d(input, [4, 4, 4])
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 4, 4, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = R.nn.adaptive_avg_pool3d(
|
|
input_1, output_size=[4, 4, 4], layout="NCDHW", out_layout="NCDHW"
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 4, 4, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 8, 8, 8, dtype=torch.float32),)
|
|
verify_model(AdaptiveAvgPool3d0(), example_args, {}, expected1)
|
|
verify_model(AdaptiveAvgPool3d1(), example_args, {}, expected1)
|
|
|
|
|
|
def test_addmm():
|
|
class Addmm1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x1, x2, x3):
|
|
return torch.addmm(x1, x2, x3)
|
|
|
|
class Addmm2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x1, x2, x3):
|
|
return torch.addmm(x1, x2, x3, beta=0.8, alpha=0.5)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
x1: R.Tensor((10, 10), dtype="float32"),
|
|
x2: R.Tensor((10, 10), dtype="float32"),
|
|
x3: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
|
|
lv1: R.Tensor((10, 10), dtype="float32") = R.add(x1, lv)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
x1: R.Tensor((10, 10), dtype="float32"),
|
|
x2: R.Tensor((10, 10), dtype="float32"),
|
|
x3: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
|
|
lv1: R.Tensor((10, 10), dtype="float32") = R.multiply(lv, R.const(0.5, "float32"))
|
|
lv2: R.Tensor((10, 10), dtype="float32") = R.multiply(x1, R.const(0.8, "float32"))
|
|
lv3: R.Tensor((10, 10), dtype="float32") = R.add(lv2, lv1)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
)
|
|
|
|
verify_model(Addmm1(), example_args, {}, expected1)
|
|
verify_model(Addmm2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_sparse_addmm():
|
|
class SparseAddmm1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x1, x2, x3):
|
|
return torch.sparse.addmm(x1, x2, x3)
|
|
|
|
class SparseAddmm2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x1, x2, x3):
|
|
return torch.sparse.addmm(x1, x2, x3, beta=0.8, alpha=0.5)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
x1: R.Tensor((10, 10), dtype="float32"),
|
|
x2: R.Tensor((10, 10), dtype="float32"),
|
|
x3: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
|
|
lv1: R.Tensor((10, 10), dtype="float32") = R.add(x1, lv)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
x1: R.Tensor((10, 10), dtype="float32"),
|
|
x2: R.Tensor((10, 10), dtype="float32"),
|
|
x3: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
|
|
lv1: R.Tensor((10, 10), dtype="float32") = R.multiply(lv, R.const(0.5, "float32"))
|
|
lv2: R.Tensor((10, 10), dtype="float32") = R.multiply(x1, R.const(0.8, "float32"))
|
|
lv3: R.Tensor((10, 10), dtype="float32") = R.add(lv2, lv1)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
torch.randn(10, 10, dtype=torch.float32),
|
|
)
|
|
|
|
verify_model(SparseAddmm1(), example_args, {}, expected1)
|
|
verify_model(SparseAddmm2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_avg_pool1d():
|
|
class AvgPool1d1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AvgPool1d(kernel_size=1)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 10), dtype="float32") = R.expand_dims(input, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 10), dtype="float32") = R.nn.avg_pool2d(
|
|
lv,
|
|
pool_size=[1, 1],
|
|
strides=[1, 1],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
ceil_mode=False,
|
|
count_include_pad=True,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv2: R.Tensor((1, 3, 10), dtype="float32") = R.squeeze(lv1, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 10), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool1d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AvgPool1d(kernel_size=3, stride=2, padding=1, ceil_mode=True)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class AvgPool1d3(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool1d(
|
|
input, kernel_size=3, stride=2, padding=1, ceil_mode=True
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 6), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 10), dtype="float32") = R.expand_dims(input, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 6), dtype="float32") = R.nn.avg_pool2d(
|
|
lv,
|
|
pool_size=[1, 3],
|
|
strides=[1, 2],
|
|
dilation=[1, 1],
|
|
padding=[0, 1, 0, 1],
|
|
ceil_mode=True,
|
|
count_include_pad=True,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv2: R.Tensor((1, 3, 6), dtype="float32") = R.squeeze(lv1, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 6), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool1d4(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool1d(input, kernel_size=2, stride=2, padding=0)
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 5), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 10), dtype="float32") = R.expand_dims(input, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 5), dtype="float32") = R.nn.avg_pool2d(
|
|
lv,
|
|
pool_size=[1, 2],
|
|
strides=[1, 2],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
ceil_mode=False,
|
|
count_include_pad=True,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv2: R.Tensor((1, 3, 5), dtype="float32") = R.squeeze(lv1, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 5), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, dtype=torch.float32),)
|
|
verify_model(AvgPool1d1(), example_args, {}, expected1)
|
|
verify_model(AvgPool1d2(), example_args, {}, expected2)
|
|
verify_model(AvgPool1d3(), example_args, {}, expected2)
|
|
verify_model(AvgPool1d4(), example_args, {}, expected3)
|
|
|
|
|
|
def test_avg_pool2d():
|
|
class AvgPool2d1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AvgPool2d(kernel_size=[1, 1])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.avg_pool2d(
|
|
input_1,
|
|
pool_size=[1, 1],
|
|
strides=[1, 1],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
count_include_pad=True,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool2d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AvgPool2d(kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class AvgPool2d3(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool2d(
|
|
input, kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv = R.nn.avg_pool2d(
|
|
input_1,
|
|
pool_size=[4, 4],
|
|
strides=[2, 2],
|
|
dilation=[1, 1],
|
|
padding=[2, 2, 2, 2],
|
|
ceil_mode=True,
|
|
count_include_pad=True,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
gv = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool2d4(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool2d(input, kernel_size=[2, 1], divisor_override=2)
|
|
|
|
@tvm.script.ir_module
|
|
class expected4:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv = R.nn.avg_pool2d(
|
|
input_1,
|
|
pool_size=[2, 1],
|
|
strides=[2, 1],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
ceil_mode=False,
|
|
count_include_pad=True,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
gv = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool2d5(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool2d(
|
|
input, kernel_size=[2, 1], divisor_override=2, count_include_pad=False
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected5:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv = R.nn.avg_pool2d(
|
|
input_1,
|
|
pool_size=[2, 1],
|
|
strides=[2, 1],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
ceil_mode=False,
|
|
count_include_pad=False,
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
gv = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(AvgPool2d1(), example_args, {}, expected1)
|
|
verify_model(AvgPool2d2(), example_args, {}, expected2)
|
|
verify_model(AvgPool2d3(), example_args, {}, expected2)
|
|
verify_model(AvgPool2d4(), example_args, {}, expected4)
|
|
verify_model(AvgPool2d5(), example_args, {}, expected5)
|
|
|
|
|
|
def test_avg_pool3d():
|
|
class AvgPool3d1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AvgPool3d(kernel_size=1)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 8, 8, 8), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = R.nn.avg_pool3d(
|
|
input_1,
|
|
pool_size=[1, 1, 1],
|
|
strides=[1, 1, 1],
|
|
dilation=[1, 1, 1],
|
|
padding=[0, 0, 0, 0, 0, 0],
|
|
ceil_mode=False,
|
|
count_include_pad=True,
|
|
layout="NCDHW",
|
|
out_layout="NCDHW",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 8, 8, 8), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool3d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.AvgPool3d(kernel_size=3, stride=2, padding=1, ceil_mode=True)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class AvgPool3d3(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool3d(
|
|
input, kernel_size=3, stride=2, padding=1, ceil_mode=True
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")):
|
|
with R.dataflow():
|
|
lv = R.nn.avg_pool3d(
|
|
input_1,
|
|
pool_size=[3, 3, 3],
|
|
strides=[2, 2, 2],
|
|
dilation=[1, 1, 1],
|
|
padding=[1, 1, 1, 1, 1, 1],
|
|
ceil_mode=True,
|
|
count_include_pad=True,
|
|
layout="NCDHW",
|
|
out_layout="NCDHW",
|
|
)
|
|
gv = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AvgPool3d4(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.avg_pool3d(input, kernel_size=[2, 1, 2], stride=[2, 1, 2])
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")):
|
|
with R.dataflow():
|
|
lv = R.nn.avg_pool3d(
|
|
input_1,
|
|
pool_size=[2, 1, 2],
|
|
strides=[2, 1, 2],
|
|
dilation=[1, 1, 1],
|
|
padding=[0, 0, 0, 0, 0, 0],
|
|
ceil_mode=False,
|
|
count_include_pad=True,
|
|
layout="NCDHW",
|
|
out_layout="NCDHW",
|
|
)
|
|
gv = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 8, 8, 8, dtype=torch.float32),)
|
|
verify_model(AvgPool3d1(), example_args, {}, expected1)
|
|
verify_model(AvgPool3d2(), example_args, {}, expected2)
|
|
verify_model(AvgPool3d3(), example_args, {}, expected2)
|
|
verify_model(AvgPool3d4(), example_args, {}, expected3)
|
|
|
|
|
|
def test_baddbmm():
|
|
class BAddBMM1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, c, x, y):
|
|
return torch.baddbmm(c, x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((4, 128, 512), dtype="float32"),
|
|
inp_1: R.Tensor((4, 128, 256), dtype="float32"),
|
|
inp_2: R.Tensor((4, 256, 512), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 128, 512), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(
|
|
inp_1, inp_2, out_dtype="float32"
|
|
)
|
|
lv1: R.Tensor((4, 128, 512), dtype="float32") = R.add(inp_0, lv)
|
|
gv: R.Tuple(R.Tensor((4, 128, 512), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class BAddBMM2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, c, x, y):
|
|
return torch.baddbmm(c, x, y, alpha=2, beta=0)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((4, 128, 512), dtype="float32"),
|
|
inp_1: R.Tensor((4, 128, 256), dtype="float32"),
|
|
inp_2: R.Tensor((4, 256, 512), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 128, 512), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(
|
|
inp_1, inp_2, out_dtype="float32"
|
|
)
|
|
lv1: R.Tensor((4, 128, 512), dtype="float32") = R.multiply(
|
|
lv, R.const(2, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((4, 128, 512), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class BAddBMM3(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, c, x, y):
|
|
return torch.baddbmm(c, x, y, alpha=2, beta=3)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((4, 128, 512), dtype="float32"),
|
|
inp_1: R.Tensor((4, 128, 256), dtype="float32"),
|
|
inp_2: R.Tensor((4, 256, 512), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 128, 512), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(
|
|
inp_1, inp_2, out_dtype="float32"
|
|
)
|
|
lv1: R.Tensor((4, 128, 512), dtype="float32") = R.multiply(
|
|
lv, R.const(2, "float32")
|
|
)
|
|
lv2: R.Tensor((4, 128, 512), dtype="float32") = R.multiply(
|
|
inp_0, R.const(3, "float32")
|
|
)
|
|
lv3: R.Tensor((4, 128, 512), dtype="float32") = R.add(lv2, lv1)
|
|
gv: R.Tuple(R.Tensor((4, 128, 512), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(4, 128, 512, dtype=torch.float32),
|
|
torch.randn(4, 128, 256, dtype=torch.float32),
|
|
torch.randn(4, 256, 512, dtype=torch.float32),
|
|
)
|
|
verify_model(
|
|
BAddBMM1(),
|
|
example_args,
|
|
{},
|
|
Expected1,
|
|
run_ep_decomposition=True,
|
|
)
|
|
|
|
verify_model(
|
|
BAddBMM2(),
|
|
example_args,
|
|
{},
|
|
Expected2,
|
|
run_ep_decomposition=True,
|
|
)
|
|
|
|
verify_model(
|
|
BAddBMM3(),
|
|
example_args,
|
|
{},
|
|
Expected3,
|
|
run_ep_decomposition=True,
|
|
)
|
|
|
|
|
|
def test_bmm():
|
|
class BMM(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x, y):
|
|
return torch.bmm(x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((4, 128, 256), dtype="float32"),
|
|
input_2: R.Tensor((4, 256, 512), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 128, 512), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(
|
|
input_1, input_2, out_dtype="float32"
|
|
)
|
|
gv: R.Tuple(R.Tensor((4, 128, 512), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(4, 128, 256, dtype=torch.float32),
|
|
torch.randn(4, 256, 512, dtype=torch.float32),
|
|
)
|
|
verify_model(
|
|
BMM(),
|
|
example_args,
|
|
{},
|
|
Expected,
|
|
run_ep_decomposition=True,
|
|
)
|
|
|
|
|
|
def test_conv_transpose1d():
|
|
class ConvTranspose1d1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.ConvTranspose1d(6, 6, 3, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
class ConvTranspose1d1Func(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.randn(size=[6, 6, 3])
|
|
self.bias = torch.randn(size=[6])
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.conv_transpose1d(input, self.weight, self.bias)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 6, 4), dtype="float32"),
|
|
w1: R.Tensor((6, 6, 3), dtype="float32"),
|
|
w2: R.Tensor((6,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 6), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 6), dtype="float32") = R.nn.conv1d_transpose(
|
|
input_1,
|
|
w1,
|
|
strides=[1],
|
|
padding=[0, 0],
|
|
output_padding=[0],
|
|
dilation=[1],
|
|
data_layout="NCW",
|
|
kernel_layout="IOW",
|
|
out_layout="NCW",
|
|
out_dtype="float32",
|
|
)
|
|
lv2: R.Tensor((1, 6, 1)) = R.reshape(w2, [1, 6, 1])
|
|
lv3: R.Tensor((1, 6, 6), dtype="float32") = R.add(lv1, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 6, 6), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class ConvTranspose1d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.ConvTranspose1d(6, 6, 3, bias=False)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 6, 4), dtype="float32"),
|
|
w1: R.Tensor((6, 6, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 6), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 6), dtype="float32") = R.nn.conv1d_transpose(
|
|
input_1,
|
|
w1,
|
|
strides=[1],
|
|
padding=[0, 0],
|
|
output_padding=[0],
|
|
dilation=[1],
|
|
data_layout="NCW",
|
|
kernel_layout="IOW",
|
|
out_layout="NCW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 6, 6), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 6, 4, dtype=torch.float32),)
|
|
|
|
model = ConvTranspose1d1()
|
|
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = ConvTranspose1d1Func()
|
|
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = ConvTranspose1d2()
|
|
binding = {"w1": model.conv.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected2)
|
|
|
|
|
|
def test_conv_transpose2d():
|
|
class ConvTranspose2d1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.ConvTranspose2d(3, 3, 7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
class ConvTranspose2d1Func(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.randn(size=[3, 3, 7, 7])
|
|
self.bias = torch.randn(size=[3])
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.conv_transpose2d(input, self.weight, self.bias)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((3, 3, 7, 7), dtype="float32"),
|
|
w2: R.Tensor((3,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 16, 16), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 3, 16, 16), dtype="float32") = R.nn.conv2d_transpose(
|
|
input_1,
|
|
w1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
output_padding=[0, 0],
|
|
dilation=[1, 1],
|
|
data_layout="NCHW",
|
|
kernel_layout="IOHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv2: R.Tensor((1, 3, 1, 1)) = R.reshape(w2, [1, 3, 1, 1])
|
|
lv3: R.Tensor((1, 3, 16, 16), dtype="float32") = R.add(lv1, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 3, 16, 16), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class ConvTranspose2d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.ConvTranspose2d(3, 3, 7, bias=False)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((3, 3, 7, 7), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 16, 16), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 3, 16, 16), dtype="float32") = R.nn.conv2d_transpose(
|
|
input_1,
|
|
w1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
output_padding=[0, 0],
|
|
dilation=[1, 1],
|
|
data_layout="NCHW",
|
|
kernel_layout="IOHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 16, 16), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model = ConvTranspose2d1()
|
|
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = ConvTranspose2d1Func()
|
|
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = ConvTranspose2d2()
|
|
binding = {"w1": model.conv.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected2)
|
|
|
|
|
|
def test_conv1d():
|
|
class Conv1D1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv1d(3, 6, 7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
class Conv1D1Func(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.randn(size=[6, 3, 7])
|
|
self.bias = torch.randn(size=[6])
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.conv1d(input, self.weight, self.bias)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
w1: R.Tensor((6, 3, 7), dtype="float32"),
|
|
w2: R.Tensor((6,), dtype="float32"),
|
|
input_1: R.Tensor((1, 3, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4), dtype="float32") = R.nn.conv1d(
|
|
input_1,
|
|
w1,
|
|
strides=[1],
|
|
padding=[0, 0],
|
|
dilation=[1],
|
|
data_layout="NCW",
|
|
kernel_layout="OIW",
|
|
out_layout="NCW",
|
|
out_dtype="float32",
|
|
)
|
|
lv2: R.Tensor((1, 6, 1), dtype="float32") = R.reshape(w2, [1, 6, 1])
|
|
lv3: R.Tensor((1, 6, 4), dtype="float32") = R.add(lv1, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Conv1D2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv1d(3, 6, 7, bias=False)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
w1: R.Tensor((6, 3, 7), dtype="float32"),
|
|
input_1: R.Tensor((1, 3, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4), dtype="float32") = R.nn.conv1d(
|
|
input_1,
|
|
w1,
|
|
strides=[1],
|
|
padding=[0, 0],
|
|
dilation=[1],
|
|
data_layout="NCW",
|
|
kernel_layout="OIW",
|
|
out_layout="NCW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, dtype=torch.float32),)
|
|
|
|
model = Conv1D1()
|
|
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Conv1D1Func()
|
|
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Conv1D2()
|
|
binding = {"w1": model.conv.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected2)
|
|
|
|
|
|
def test_conv2d():
|
|
class Conv2D1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
class Conv2D1Func(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.randn(size=[6, 3, 7, 7])
|
|
self.bias = torch.randn(size=[6])
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.conv2d(input, self.weight, self.bias)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
|
|
w2: R.Tensor((6,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4, 4), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
|
|
input_1,
|
|
w1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv2: R.Tensor((1, 6, 1, 1)) = R.reshape(w2, [1, 6, 1, 1])
|
|
lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4, 4), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Conv2D2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv2d(3, 6, 7, bias=False)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4, 4), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
|
|
input_1,
|
|
w1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4, 4), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model = Conv2D1()
|
|
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Conv2D1Func()
|
|
binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Conv2D2()
|
|
binding = {"w1": model.conv.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected2)
|
|
|
|
|
|
def test_conv3d():
|
|
class Conv3D1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv3d(3, 6, 7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
class Conv3D1Func(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.randn(size=[6, 3, 7, 7, 7])
|
|
self.bias = torch.randn(size=[6])
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.conv3d(input, self.weight, self.bias)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((6, 3, 7, 7, 7), dtype="float32"),
|
|
w2: R.Tensor((6,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4, 4, 4), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.nn.conv3d(
|
|
input_1,
|
|
w1,
|
|
strides=[1],
|
|
padding=[0, 0, 0],
|
|
dilation=[1],
|
|
data_layout="NCDHW",
|
|
kernel_layout="OIDHW",
|
|
out_layout="NCDHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv2: R.Tensor((1, 6, 1, 1, 1)) = R.reshape(w2, [1, 6, 1, 1, 1])
|
|
lv3: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.add(lv1, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4, 4, 4), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Conv3D2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv3d(3, 6, 7, bias=False)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((6, 3, 7, 7, 7), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4, 4, 4), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.nn.conv3d(
|
|
input_1,
|
|
w1,
|
|
strides=[1],
|
|
padding=[0, 0, 0],
|
|
dilation=[1],
|
|
data_layout="NCDHW",
|
|
kernel_layout="OIDHW",
|
|
out_layout="NCDHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4, 4, 4), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, 10, dtype=torch.float32),)
|
|
|
|
model = Conv3D1()
|
|
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Conv3D1Func()
|
|
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Conv3D2()
|
|
binding = {"w1": model.conv.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected2)
|
|
|
|
|
|
def test_pad():
|
|
class PadModel(torch.nn.Module):
|
|
def __init__(self, pad, mode="constant", value=0.0):
|
|
super().__init__()
|
|
self.pad = pad
|
|
self.mode = mode
|
|
self.value = value
|
|
|
|
def forward(self, x):
|
|
if self.mode == "constant":
|
|
return torch.nn.functional.pad(x, self.pad, mode=self.mode, value=self.value)
|
|
else:
|
|
return torch.nn.functional.pad(x, self.pad, mode=self.mode)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_constant:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 14, 12), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
|
|
x,
|
|
pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
|
|
pad_mode="constant",
|
|
pad_value=0.0,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 14, 12), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_reflect:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 14, 12), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((14,), dtype="int64") = R.arange(
|
|
R.prim_value(-2), R.prim_value(12), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((14,), dtype="int64") = R.abs(lv)
|
|
lv2: R.Tensor((14,), dtype="int64") = R.subtract(R.const(9, "int64"), lv1)
|
|
lv3: R.Tensor((14,), dtype="int64") = R.abs(lv2)
|
|
lv4: R.Tensor((14,), dtype="int64") = R.subtract(R.const(9, "int64"), lv3)
|
|
lv5: R.Tensor((1, 3, 14, 10), dtype="float32") = R.take(x, lv4, axis=2, mode="fast")
|
|
lv6: R.Tensor((12,), dtype="int64") = R.arange(
|
|
R.prim_value(-1), R.prim_value(11), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv7: R.Tensor((12,), dtype="int64") = R.abs(lv6)
|
|
lv8: R.Tensor((12,), dtype="int64") = R.subtract(R.const(9, "int64"), lv7)
|
|
lv9: R.Tensor((12,), dtype="int64") = R.abs(lv8)
|
|
lv10: R.Tensor((12,), dtype="int64") = R.subtract(R.const(9, "int64"), lv9)
|
|
lv11: R.Tensor((1, 3, 14, 12), dtype="float32") = R.take(
|
|
lv5, lv10, axis=3, mode="fast"
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 14, 12), dtype="float32")) = (lv11,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_replicate:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 14, 12), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((14,), dtype="int64") = R.arange(
|
|
R.prim_value(-2), R.prim_value(12), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((14,), dtype="int64") = R.clip(lv, R.prim_value(0), R.prim_value(9))
|
|
lv2: R.Tensor((1, 3, 14, 10), dtype="float32") = R.take(x, lv1, axis=2, mode="fast")
|
|
lv3: R.Tensor((12,), dtype="int64") = R.arange(
|
|
R.prim_value(-1), R.prim_value(11), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv4: R.Tensor((12,), dtype="int64") = R.clip(lv3, R.prim_value(0), R.prim_value(9))
|
|
lv5: R.Tensor((1, 3, 14, 12), dtype="float32") = R.take(
|
|
lv2, lv4, axis=3, mode="fast"
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 14, 12), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_circular:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 14, 12), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.zeros(
|
|
R.shape([1, 3, 14, 12]), dtype="float32"
|
|
)
|
|
|
|
lv1: R.Tensor((1, 3, 14, 10), dtype="float32") = R.strided_slice(
|
|
lv,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(11),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
x,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(10),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
lv1,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(12),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
lv2,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(10),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
|
|
lv4, R.shape([1, 3, 10, 10])
|
|
)
|
|
|
|
lv6: R.Tensor((1, 3, 14, 10), dtype="float32") = R.strided_slice(
|
|
lv,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(11),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv7: R.Tensor((1, 3, 14, 10), dtype="float32") = R.slice_scatter(
|
|
lv6, lv5, R.prim_value(2), R.prim_value(12), R.prim_value(1), axis=2
|
|
)
|
|
|
|
lv8: R.Tensor((1, 3, 14, 12), dtype="float32") = R.slice_scatter(
|
|
lv, lv7, R.prim_value(1), R.prim_value(11), R.prim_value(1), axis=3
|
|
)
|
|
|
|
lv9: R.Tensor((1, 3, 14, 1), dtype="float32") = R.strided_slice(
|
|
lv8,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv10: R.Tensor((1, 3, 14, 1), dtype="float32") = R.strided_slice(
|
|
lv8,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(10),),
|
|
(R.prim_value(11),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv11: R.Tensor((1, 3, 14, 1), dtype="float32") = R.broadcast_to(
|
|
lv10, R.shape([1, 3, 14, 1])
|
|
)
|
|
|
|
lv12: R.Tensor((1, 3, 14, 12), dtype="float32") = R.slice_scatter(
|
|
lv8, lv11, R.prim_value(0), R.prim_value(1), R.prim_value(1), axis=3
|
|
)
|
|
|
|
lv13: R.Tensor((1, 3, 14, 1), dtype="float32") = R.strided_slice(
|
|
lv12,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(11),),
|
|
(R.prim_value(12),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv14: R.Tensor((1, 3, 14, 1), dtype="float32") = R.strided_slice(
|
|
lv12,
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv15: R.Tensor((1, 3, 14, 1), dtype="float32") = R.broadcast_to(
|
|
lv14, R.shape([1, 3, 14, 1])
|
|
)
|
|
lv16: R.Tensor((1, 3, 14, 12), dtype="float32") = R.slice_scatter(
|
|
lv12, lv15, R.prim_value(11), R.prim_value(12), R.prim_value(1), axis=3
|
|
)
|
|
|
|
lv17: R.Tensor((1, 3, 2, 12), dtype="float32") = R.strided_slice(
|
|
lv16,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv18: R.Tensor((1, 3, 2, 12), dtype="float32") = R.strided_slice(
|
|
lv16,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(10),),
|
|
(R.prim_value(12),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv19: R.Tensor((1, 3, 2, 12), dtype="float32") = R.broadcast_to(
|
|
lv18, R.shape([1, 3, 2, 12])
|
|
)
|
|
|
|
lv20: R.Tensor((1, 3, 14, 12), dtype="float32") = R.slice_scatter(
|
|
lv16, lv19, R.prim_value(0), R.prim_value(2), R.prim_value(1), axis=2
|
|
)
|
|
lv21: R.Tensor((1, 3, 2, 12), dtype="float32") = R.strided_slice(
|
|
lv20,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(12),),
|
|
(R.prim_value(14),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv22: R.Tensor((1, 3, 2, 12), dtype="float32") = R.strided_slice(
|
|
lv20,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(4),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
|
|
lv23: R.Tensor((1, 3, 2, 12), dtype="float32") = R.broadcast_to(
|
|
lv22, R.shape([1, 3, 2, 12])
|
|
)
|
|
|
|
lv24: R.Tensor((1, 3, 14, 12), dtype="float32") = R.slice_scatter(
|
|
lv20, lv23, R.prim_value(12), R.prim_value(14), R.prim_value(1), axis=2
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 14, 12), dtype="float32")) = (lv24,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(PadModel(pad=[1, 1, 2, 2]), example_args, {}, expected_constant)
|
|
verify_model(
|
|
PadModel(pad=[1, 1, 2, 2], mode="reflect"),
|
|
example_args,
|
|
{},
|
|
expected_reflect,
|
|
run_ep_decomposition=True,
|
|
)
|
|
verify_model(
|
|
PadModel(pad=[1, 1, 2, 2], mode="replicate"),
|
|
example_args,
|
|
{},
|
|
expected_replicate,
|
|
run_ep_decomposition=True,
|
|
)
|
|
verify_model(
|
|
PadModel(pad=[1, 1, 2, 2], mode="circular"),
|
|
example_args,
|
|
{},
|
|
expected_circular,
|
|
run_ep_decomposition=True,
|
|
)
|
|
|
|
|
|
def test_pixel_shuffle():
|
|
class PixelShuffle1(torch.nn.Module):
|
|
def __init__(self, upscale_factor=2):
|
|
super().__init__()
|
|
self.pixel_shuffle = torch.nn.PixelShuffle(upscale_factor)
|
|
|
|
def forward(self, x):
|
|
return self.pixel_shuffle(x)
|
|
|
|
class PixelShuffle2(torch.nn.Module):
|
|
def __init__(self, upscale_factor=2):
|
|
super().__init__()
|
|
self.upscale_factor = upscale_factor
|
|
|
|
def forward(self, x):
|
|
return torch.nn.functional.pixel_shuffle(x, self.upscale_factor)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 8, 10, 15), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 20, 30), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 2, 2, 2, 10, 15), dtype="float32") = R.reshape(
|
|
x, R.shape([1, 2, 2, 2, 10, 15])
|
|
)
|
|
lv1: R.Tensor((1, 2, 10, 2, 15, 2), dtype="float32") = R.permute_dims(
|
|
lv, axes=[0, 1, 4, 2, 5, 3]
|
|
)
|
|
lv2: R.Tensor((1, 2, 20, 30), dtype="float32") = R.reshape(
|
|
lv1, R.shape([1, 2, 20, 30])
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 2, 20, 30), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 8, 10, 15, dtype=torch.float32),)
|
|
verify_model(PixelShuffle1(upscale_factor=2), example_args, {}, expected)
|
|
verify_model(PixelShuffle2(upscale_factor=2), example_args, {}, expected)
|
|
|
|
|
|
def test_einsum():
|
|
class Einsum1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return torch.einsum("ii", x)
|
|
|
|
class Einsum2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x, y):
|
|
return torch.einsum("i,j->ij", x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((4, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.einsum((inp_0,), subscripts="ii")
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((5,), dtype="float32"), inp_1: R.Tensor((4,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((5, 4), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 4), dtype="float32") = R.einsum(
|
|
(inp_0, inp_1), subscripts="i,j->ij"
|
|
)
|
|
gv: R.Tuple(R.Tensor((5, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(4, 4, dtype=torch.float32),)
|
|
verify_model(Einsum1(), example_args, {}, Expected1, run_ep_decomposition=False)
|
|
|
|
example_args = (torch.randn(5, dtype=torch.float32), torch.randn(4, dtype=torch.float32))
|
|
verify_model(Einsum2(), example_args, {}, Expected2, run_ep_decomposition=False)
|
|
|
|
|
|
def test_outer():
|
|
class Outer(torch.nn.Module):
|
|
def forward(self, x, y):
|
|
return torch.outer(x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(x: R.Tensor((3,), dtype="float32"), y: R.Tensor((4,), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1), dtype="float32") = R.reshape(x, R.shape([3, 1]))
|
|
lv1: R.Tensor((3, 4), dtype="float32") = R.multiply(lv, y)
|
|
gv: R.Tuple(R.Tensor((3, 4), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(3, dtype=torch.float32),
|
|
torch.randn(4, dtype=torch.float32),
|
|
)
|
|
verify_model(Outer(), example_args, {}, expected)
|
|
|
|
|
|
def test_embedding():
|
|
class Embedding(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.embedding = torch.nn.Embedding(10, 3)
|
|
|
|
def forward(self, input):
|
|
return self.embedding(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((4,), dtype="int64"), w1: R.Tensor((10, 3), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((4, 3), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((4,), dtype="int32") = R.astype(input_1, dtype="int32")
|
|
lv1: R.Tensor((4, 3), dtype="float32") = R.take(w1, lv, axis=0)
|
|
gv: R.Tuple(R.Tensor((4, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randint(low=-int(1e5), high=int(1e5), size=(4,), dtype=torch.int64),)
|
|
|
|
model = Embedding()
|
|
binding = {"w1": model.embedding.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
|
|
def test_groupnorm():
|
|
import torch
|
|
from torch.nn import Module
|
|
|
|
torch.set_grad_enabled(False)
|
|
torch.random.manual_seed(0)
|
|
|
|
class GroupNorm(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gn = torch.nn.GroupNorm(3, 3)
|
|
|
|
def forward(self, input):
|
|
return self.gn(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((3,), dtype="float32"),
|
|
w2: R.Tensor((3,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.group_norm(
|
|
input_1,
|
|
w1,
|
|
w2,
|
|
num_groups=3,
|
|
channel_axis=1,
|
|
axes=[2, 3],
|
|
epsilon=1.0000000000000001e-05,
|
|
center=True,
|
|
scale=True,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model = GroupNorm()
|
|
binding = {
|
|
"w1": model.gn.weight.detach().numpy(),
|
|
"w2": model.gn.bias.detach().numpy(),
|
|
}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
|
|
def test_instancenorm2d():
|
|
torch.set_grad_enabled(False)
|
|
torch.random.manual_seed(0)
|
|
|
|
class InstanceNorm2d(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gn = torch.nn.InstanceNorm2d(3)
|
|
|
|
def forward(self, input):
|
|
return self.gn(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((3,), dtype="float32"),
|
|
w2: R.Tensor((3,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.instance_norm(
|
|
input_1,
|
|
w1,
|
|
w2,
|
|
channel_axis=1,
|
|
axes=[0, 2, 3],
|
|
epsilon=1e-05,
|
|
center=True,
|
|
scale=True,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model = InstanceNorm2d()
|
|
binding = {
|
|
"w1": torch.ones(3).detach().numpy(),
|
|
"w2": torch.zeros(3).detach().numpy(),
|
|
}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
|
|
def test_layernorm():
|
|
class LayerNorm(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.ln = torch.nn.LayerNorm((10, 10))
|
|
|
|
def forward(self, input):
|
|
return self.ln(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
w1: R.Tensor((10, 10), dtype="float32"),
|
|
w2: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
|
|
input_1,
|
|
w1,
|
|
w2,
|
|
axes=[-2, -1],
|
|
epsilon=1e-05,
|
|
center=True,
|
|
scale=True,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model = LayerNorm()
|
|
binding = {
|
|
"w1": model.ln.weight.detach().numpy(),
|
|
"w2": model.ln.bias.detach().numpy(),
|
|
}
|
|
verify_model(LayerNorm(), example_args, binding, expected1)
|
|
|
|
|
|
def test_linear():
|
|
class Dense1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(10, 7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.linear(input)
|
|
|
|
class Dense1Func(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = torch.randn(size=[7, 10])
|
|
self.bias = torch.randn(size=[7])
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.linear(input, self.weight, self.bias)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
w1: R.Tensor((7, 10), dtype="float32"),
|
|
w2: R.Tensor((7,), dtype="float32"),
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 7), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((30, 10), dtype="float32") = R.reshape(input_1, R.shape([30, 10]))
|
|
lv1: R.Tensor((10, 7), dtype="float32") = R.permute_dims(w1, axes=[1, 0])
|
|
lv2: R.Tensor((30, 7), dtype="float32") = R.matmul(lv, lv1, out_dtype="float32")
|
|
lv3: R.Tensor((30, 7), dtype="float32") = R.add(w2, lv2)
|
|
lv4: R.Tensor((1, 3, 10, 7), dtype="float32") = R.reshape(
|
|
lv3, R.shape([1, 3, 10, 7])
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 7), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Dense2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(10, 7, bias=False)
|
|
|
|
def forward(self, input):
|
|
return self.linear(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
w1: R.Tensor((7, 10), dtype="float32"),
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 10, 7), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 7), dtype="float32") = R.permute_dims(w1, axes=[1, 0])
|
|
lv1: R.Tensor((30, 10), dtype="float32") = R.reshape(input_1, R.shape([30, 10]))
|
|
lv2: R.Tensor((30, 7), dtype="float32") = R.matmul(lv1, lv, out_dtype="float32")
|
|
lv3: R.Tensor((1, 3, 10, 7), dtype="float32") = R.reshape(
|
|
lv2, R.shape([1, 3, 10, 7])
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 7), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
model = Dense1()
|
|
binding = {"w1": model.linear.weight.detach().numpy(), "w2": model.linear.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Dense1Func()
|
|
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected1)
|
|
|
|
model = Dense2()
|
|
binding = {"w1": model.linear.weight.detach().numpy()}
|
|
verify_model(model, example_args, binding, expected2)
|
|
|
|
|
|
def test_maxpool1d():
|
|
class MaxPool1d(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool1d(kernel_size=2)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class MaxPool1d_functional(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.max_pool1d(input, kernel_size=2)
|
|
|
|
class MaxPool1d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool1d(kernel_size=3, stride=2)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 8), dtype="float32") = R.expand_dims(input_1, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 4), dtype="float32") = R.nn.max_pool2d(
|
|
lv,
|
|
pool_size=[1, 2],
|
|
strides=[1, 2],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv2: R.Tensor((1, 3, 1, 4), dtype="float32") = R.zeros_like(lv1)
|
|
lv3: R.Tuple(
|
|
R.Tensor((1, 3, 1, 4), dtype="float32"),
|
|
R.Tensor((1, 3, 1, 4), dtype="float32"),
|
|
) = (lv1, lv2)
|
|
lv4: R.Tensor((1, 3, 1, 4), dtype="float32") = lv3[0]
|
|
lv5: R.Tensor((1, 3, 4), dtype="float32") = R.squeeze(lv4, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 4), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 8), dtype="float32") = R.expand_dims(input_1, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 4), dtype="float32") = R.nn.max_pool2d(
|
|
lv,
|
|
pool_size=[1, 2],
|
|
strides=[1, 2],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv2: R.Tensor((1, 3, 1, 4), dtype="float32") = R.zeros_like(lv1)
|
|
lv3: R.Tuple(
|
|
R.Tensor((1, 3, 1, 4), dtype="float32"),
|
|
R.Tensor((1, 3, 1, 4), dtype="float32"),
|
|
) = (lv1, lv2)
|
|
lv4: R.Tensor((1, 3, 1, 4), dtype="float32") = lv3[0]
|
|
lv5: R.Tensor((1, 3, 4), dtype="float32") = R.squeeze(lv4, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 4), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 1, 10), dtype="float32") = R.expand_dims(input_1, axis=[-2])
|
|
lv1: R.Tensor((1, 3, 1, 4), dtype="float32") = R.nn.max_pool2d(
|
|
lv,
|
|
pool_size=[1, 3],
|
|
strides=[1, 2],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv2: R.Tensor((1, 3, 1, 4), dtype="float32") = R.zeros_like(lv1)
|
|
lv3: R.Tuple(
|
|
R.Tensor((1, 3, 1, 4), dtype="float32"),
|
|
R.Tensor((1, 3, 1, 4), dtype="float32"),
|
|
) = (lv1, lv2)
|
|
lv4: R.Tensor((1, 3, 1, 4), dtype="float32") = lv3[0]
|
|
lv5: R.Tensor((1, 3, 4), dtype="float32") = R.squeeze(lv4, axis=[-2])
|
|
gv: R.Tuple(R.Tensor((1, 3, 4), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Example inputs
|
|
example_args1 = (torch.randn(1, 3, 8, dtype=torch.float32),)
|
|
example_args2 = (torch.randn(1, 3, 8, dtype=torch.float32),)
|
|
example_args3 = (torch.randn(1, 3, 10, dtype=torch.float32),)
|
|
|
|
# Verify the models
|
|
verify_model(MaxPool1d(), example_args1, {}, expected1)
|
|
verify_model(MaxPool1d_functional(), example_args2, {}, expected2)
|
|
verify_model(MaxPool1d2(), example_args3, {}, expected3)
|
|
|
|
|
|
def test_maxpool2d():
|
|
class MaxPool2d(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool2d(kernel_size=[1, 1])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class MaxPool2d_functional(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1])
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.max_pool2d(
|
|
input_1,
|
|
pool_size=[1, 1],
|
|
strides=[1, 1],
|
|
dilation=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.zeros_like(lv)
|
|
lv2: R.Tuple(
|
|
R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
) = (lv, lv1)
|
|
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = lv2[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class MaxPool2d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool2d(kernel_size=[2, 2], dilation=[2, 3])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 4, 4), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 4, 4), dtype="float32") = R.nn.max_pool2d(
|
|
input_1,
|
|
pool_size=[2, 2],
|
|
strides=[2, 2],
|
|
dilation=[2, 3],
|
|
padding=[0, 0, 0, 0],
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv1: R.Tensor((1, 3, 4, 4), dtype="float32") = R.zeros_like(lv)
|
|
lv2: R.Tuple(
|
|
R.Tensor((1, 3, 4, 4), dtype="float32"), R.Tensor((1, 3, 4, 4), dtype="float32")
|
|
) = (lv, lv1)
|
|
lv3: R.Tensor((1, 3, 4, 4), dtype="float32") = lv2[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 4, 4), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class MaxPool2d3(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool2d(kernel_size=[4, 4], padding=2, stride=2)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 6, 6), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 6, 6), dtype="float32") = R.nn.max_pool2d(
|
|
input_1,
|
|
pool_size=[4, 4],
|
|
strides=[2, 2],
|
|
dilation=[1, 1],
|
|
padding=[2, 2, 2, 2],
|
|
layout="NCHW",
|
|
out_layout="NCHW",
|
|
)
|
|
lv1: R.Tensor((1, 3, 6, 6), dtype="float32") = R.zeros_like(lv)
|
|
lv2: R.Tuple(
|
|
R.Tensor((1, 3, 6, 6), dtype="float32"), R.Tensor((1, 3, 6, 6), dtype="float32")
|
|
) = (lv, lv1)
|
|
lv3: R.Tensor((1, 3, 6, 6), dtype="float32") = lv2[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 6, 6), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(MaxPool2d(), example_args, {}, expected1)
|
|
verify_model(MaxPool2d_functional(), example_args, {}, expected1)
|
|
verify_model(MaxPool2d2(), example_args, {}, expected2)
|
|
verify_model(MaxPool2d3(), example_args, {}, expected3)
|
|
|
|
|
|
def test_maxpool3d():
|
|
class MaxPool3d(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool3d(kernel_size=[1, 1, 1])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
class MaxPool3d_functional(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1])
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 4, 4, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 4, 4, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = R.nn.max_pool3d(
|
|
input_1,
|
|
pool_size=[1, 1, 1],
|
|
strides=[1, 1, 1],
|
|
dilation=[1, 1, 1],
|
|
padding=[0, 0, 0, 0, 0, 0],
|
|
layout="NCDHW",
|
|
out_layout="NCDHW",
|
|
)
|
|
lv1: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = R.zeros_like(lv)
|
|
lv2: R.Tuple(
|
|
R.Tensor((1, 3, 4, 4, 4), dtype="float32"),
|
|
R.Tensor((1, 3, 4, 4, 4), dtype="float32"),
|
|
) = (lv, lv1)
|
|
lv3: R.Tensor((1, 3, 4, 4, 4), dtype="float32") = lv2[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 4, 4, 4), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class MaxPool3d2(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[2, 2, 2])
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 3, 3, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 3, 3, 3), dtype="float32") = R.nn.max_pool3d(
|
|
input_1,
|
|
pool_size=[2, 2, 2],
|
|
strides=[2, 2, 2],
|
|
dilation=[2, 2, 2],
|
|
padding=[0, 0, 0, 0, 0, 0],
|
|
layout="NCDHW",
|
|
out_layout="NCDHW",
|
|
)
|
|
lv1: R.Tensor((1, 3, 3, 3, 3), dtype="float32") = R.zeros_like(lv)
|
|
lv2: R.Tuple(
|
|
R.Tensor((1, 3, 3, 3, 3), dtype="float32"),
|
|
R.Tensor((1, 3, 3, 3, 3), dtype="float32"),
|
|
) = (lv, lv1)
|
|
lv3: R.Tensor((1, 3, 3, 3, 3), dtype="float32") = lv2[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 3, 3, 3), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class MaxPool3d3(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pool = torch.nn.MaxPool3d(kernel_size=[3, 3, 3], padding=1, stride=2)
|
|
|
|
def forward(self, input):
|
|
return self.pool(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 5, 5, 5), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = R.nn.max_pool3d(
|
|
input_1,
|
|
pool_size=[3, 3, 3],
|
|
strides=[2, 2, 2],
|
|
dilation=[1, 1, 1],
|
|
padding=[1, 1, 1, 1, 1, 1],
|
|
layout="NCDHW",
|
|
out_layout="NCDHW",
|
|
)
|
|
lv1: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = R.zeros_like(lv)
|
|
lv2: R.Tuple(
|
|
R.Tensor((1, 3, 5, 5, 5), dtype="float32"),
|
|
R.Tensor((1, 3, 5, 5, 5), dtype="float32"),
|
|
) = (lv, lv1)
|
|
lv3: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = lv2[0]
|
|
gv: R.Tuple(R.Tensor((1, 3, 5, 5, 5), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Example input tensors
|
|
example_args1 = (torch.randn(1, 3, 4, 4, 4, dtype=torch.float32),)
|
|
example_args2 = (torch.randn(1, 3, 8, 8, 8, dtype=torch.float32),)
|
|
example_args3 = (torch.randn(1, 3, 10, 10, 10, dtype=torch.float32),)
|
|
|
|
# Verify the models with expected IR modules
|
|
verify_model(MaxPool3d(), example_args1, {}, expected1)
|
|
verify_model(MaxPool3d_functional(), example_args1, {}, expected1)
|
|
verify_model(MaxPool3d2(), example_args2, {}, expected2)
|
|
verify_model(MaxPool3d3(), example_args3, {}, expected3)
|
|
|
|
|
|
def test_scaled_dot_product_attention():
|
|
class Attention1(Module):
|
|
def forward(self, q, k, v):
|
|
return torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
q: R.Tensor((32, 8, 128, 64), dtype="float32"),
|
|
k: R.Tensor((32, 8, 128, 64), dtype="float32"),
|
|
v: R.Tensor((32, 8, 128, 64), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
|
|
q, axes=[0, 2, 1, 3]
|
|
)
|
|
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
|
|
k, axes=[0, 2, 1, 3]
|
|
)
|
|
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
|
|
v, axes=[0, 2, 1, 3]
|
|
)
|
|
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
|
|
lv, lv1, lv2, scale=None, causal_mask=None, window_size=None
|
|
)
|
|
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
|
|
lv3, axes=[0, 2, 1, 3]
|
|
)
|
|
gv: R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Attention2(Module):
|
|
def forward(self, q, k, v, mask):
|
|
return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask)
|
|
|
|
@I.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
q: R.Tensor((32, 8, 128, 64), dtype="float32"),
|
|
k: R.Tensor((32, 8, 128, 64), dtype="float32"),
|
|
v: R.Tensor((32, 8, 128, 64), dtype="float32"),
|
|
mask: R.Tensor((32, 8, 128, 128), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
|
|
q, axes=[0, 2, 1, 3]
|
|
)
|
|
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
|
|
k, axes=[0, 2, 1, 3]
|
|
)
|
|
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
|
|
v, axes=[0, 2, 1, 3]
|
|
)
|
|
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention_bias(
|
|
lv, lv1, lv2, mask, scale=None, causal_mask=None, window_size=None
|
|
)
|
|
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
|
|
lv3, axes=[0, 2, 1, 3]
|
|
)
|
|
gv: R.Tuple(R.Tensor((32, 8, 128, 64), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(
|
|
Attention1(),
|
|
(
|
|
torch.randn(32, 8, 128, 64, dtype=torch.float32),
|
|
torch.randn(32, 8, 128, 64, dtype=torch.float32),
|
|
torch.randn(32, 8, 128, 64, dtype=torch.float32),
|
|
),
|
|
{},
|
|
Expected1,
|
|
run_ep_decomposition=False,
|
|
)
|
|
|
|
verify_model(
|
|
Attention2(),
|
|
(
|
|
torch.randn(32, 8, 128, 64, dtype=torch.float32),
|
|
torch.randn(32, 8, 128, 64, dtype=torch.float32),
|
|
torch.randn(32, 8, 128, 64, dtype=torch.float32),
|
|
torch.randn(32, 8, 128, 128, dtype=torch.float32),
|
|
),
|
|
{},
|
|
Expected2,
|
|
run_ep_decomposition=False,
|
|
)
|
|
|
|
# Test 2D input (seq_len, head_dim) - bug fix for #18441
|
|
class Attention2D(Module):
|
|
def forward(self, x):
|
|
return torch.nn.functional.scaled_dot_product_attention(x, x, x, is_causal=False)
|
|
|
|
@I.ir_module
|
|
class Expected2D:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((8, 32), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((8, 32), dtype="float32")):
|
|
with R.dataflow():
|
|
# Expand to add batch dimension for query, key, value separately
|
|
# (8, 32) -> (1, 8, 32)
|
|
lv: R.Tensor((1, 8, 32), dtype="float32") = R.expand_dims(x, axis=[0])
|
|
lv1: R.Tensor((1, 8, 32), dtype="float32") = R.expand_dims(x, axis=[0])
|
|
lv2: R.Tensor((1, 8, 32), dtype="float32") = R.expand_dims(x, axis=[0])
|
|
# Expand to add num_heads dimension: (1, 8, 32) -> (1, 1, 8, 32)
|
|
lv3: R.Tensor((1, 1, 8, 32), dtype="float32") = R.expand_dims(lv, axis=[1])
|
|
lv4: R.Tensor((1, 1, 8, 32), dtype="float32") = R.expand_dims(lv1, axis=[1])
|
|
lv5: R.Tensor((1, 1, 8, 32), dtype="float32") = R.expand_dims(lv2, axis=[1])
|
|
# Attention operation: (1, 1, 8, 32) -> (1, 1, 8, 32)
|
|
lv6: R.Tensor((1, 1, 8, 32), dtype="float32") = R.nn.attention(
|
|
lv3, lv4, lv5, scale=None, causal_mask=None, window_size=None
|
|
)
|
|
# Squeeze batch and num_heads dimensions: (1, 1, 8, 32) -> (8, 32)
|
|
lv7: R.Tensor((8, 32), dtype="float32") = R.squeeze(lv6, axis=[0, 1])
|
|
gv: R.Tuple(R.Tensor((8, 32), dtype="float32")) = (lv7,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(
|
|
Attention2D(),
|
|
(torch.randn(8, 32, dtype=torch.float32),),
|
|
{},
|
|
Expected2D,
|
|
run_ep_decomposition=False,
|
|
)
|
|
|
|
|
|
def test_unbind():
|
|
class Unbind1(Module):
|
|
def forward(self, data):
|
|
return torch.unbind(data)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(data: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv3: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv, axis=[0])
|
|
lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[0])
|
|
lv5: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv2, axis=[0])
|
|
gv: R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
) = (lv3, lv4, lv5)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Unbind2(Module):
|
|
def forward(self, data):
|
|
return torch.unbind(data, dim=1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(data: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv1: R.Tensor((3, 1, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv2: R.Tensor((3, 1, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv3: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv, axis=[1])
|
|
lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[1])
|
|
lv5: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv2, axis=[1])
|
|
gv: R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
) = (lv3, lv4, lv5)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(data: R.Tensor((3, 1, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1, 3), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.squeeze(lv, axis=[1])
|
|
gv: R.Tuple(R.Tensor((3, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(3, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Unbind1(), example_args, {}, expected1)
|
|
verify_model(Unbind2(), example_args, {}, expected2)
|
|
single_dim_args = (torch.randn(3, 1, 3, dtype=torch.float32),)
|
|
verify_model(Unbind2(), single_dim_args, {}, expected3)
|
|
|
|
|
|
def test_interpolate():
|
|
class InterpolateBilinear(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.interpolate(input, (224, 224), mode="bilinear")
|
|
|
|
@tvm.script.ir_module
|
|
class expected_bilinear:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 112, 112), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 224, 224), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 224, 224), dtype="float32") = R.image.resize2d(
|
|
input,
|
|
R.shape([224, 224]),
|
|
roi=[T.float32(0.0), T.float32(0.0), T.float32(0.0), T.float32(0.0)],
|
|
layout="NCHW",
|
|
method="linear",
|
|
coordinate_transformation_mode="half_pixel",
|
|
rounding_method="round",
|
|
cubic_alpha=-0.75,
|
|
cubic_exclude=0,
|
|
extrapolation_value=0.0,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 224, 224), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class InterpolateNearest(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.interpolate(input, (224, 224), mode="nearest")
|
|
|
|
@tvm.script.ir_module
|
|
class expected_nearest:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 112, 112), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 224, 224), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 224, 224), dtype="float32") = R.image.resize2d(
|
|
input,
|
|
R.shape([224, 224]),
|
|
roi=[T.float32(0.0), T.float32(0.0), T.float32(0.0), T.float32(0.0)],
|
|
layout="NCHW",
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="half_pixel",
|
|
rounding_method="round",
|
|
cubic_alpha=-0.75,
|
|
cubic_exclude=0,
|
|
extrapolation_value=0.0,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 224, 224), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class InterpolateBicubic(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.interpolate(input, (224, 224), mode="bicubic")
|
|
|
|
@I.ir_module
|
|
class expected_bicubic:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 112, 112), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 224, 224), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((224,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(224), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((224,), dtype="float32") = R.astype(lv, dtype="float32")
|
|
lv2: R.Tensor((224,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(224), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv3: R.Tensor((224,), dtype="float32") = R.astype(lv2, dtype="float32")
|
|
lv4: R.Tensor((224,), dtype="float32") = R.add(lv3, R.const(0.5, "float32"))
|
|
lv5: R.Tensor((224,), dtype="float32") = R.multiply(lv4, R.const(0.5, "float32"))
|
|
lv6: R.Tensor((224,), dtype="float32") = R.subtract(lv5, R.const(0.5, "float32"))
|
|
lv7: R.Tensor((224,), dtype="float32") = R.add(lv1, R.const(0.5, "float32"))
|
|
lv8: R.Tensor((224,), dtype="float32") = R.multiply(lv7, R.const(0.5, "float32"))
|
|
lv9: R.Tensor((224,), dtype="float32") = R.subtract(lv8, R.const(0.5, "float32"))
|
|
lv10: R.Tensor((224, 1), dtype="float32") = R.expand_dims(lv9, axis=[-1])
|
|
lv11: R.Tensor((224,), dtype="float32") = R.floor(lv6)
|
|
lv12: R.Tensor((224, 1), dtype="float32") = R.floor(lv10)
|
|
lv13: R.Tensor((224, 1), dtype="float32") = R.subtract(lv10, lv12)
|
|
lv14: R.Tensor((224, 1), dtype="float32") = R.clip(
|
|
lv13, R.prim_value(T.float64(0.0)), R.prim_value(T.float64(1.0))
|
|
)
|
|
lv15: R.Tensor((224,), dtype="float32") = R.subtract(lv6, lv11)
|
|
lv16: R.Tensor((224,), dtype="float32") = R.clip(
|
|
lv15, R.prim_value(T.float64(0.0)), R.prim_value(T.float64(1.0))
|
|
)
|
|
lv17: R.Tensor((224,), dtype="int64") = R.astype(lv11, dtype="int64")
|
|
lv18: R.Tensor((224, 1), dtype="int64") = R.astype(lv12, dtype="int64")
|
|
lv19: R.Tensor((224, 1), dtype="int64") = R.subtract(lv18, R.const(1, "int64"))
|
|
lv20: R.Tensor((224, 1), dtype="int64") = R.add(lv18, R.const(1, "int64"))
|
|
lv21: R.Tensor((224, 1), dtype="int64") = R.add(lv18, R.const(2, "int64"))
|
|
lv22: R.Tensor((224,), dtype="int64") = R.subtract(lv17, R.const(1, "int64"))
|
|
lv23: R.Tensor((224,), dtype="int64") = R.add(lv17, R.const(1, "int64"))
|
|
lv24: R.Tensor((224,), dtype="int64") = R.add(lv17, R.const(2, "int64"))
|
|
lv25: R.Tensor((224,), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv16)
|
|
lv26: R.Tensor((448,), dtype="float32") = R.concat((lv16, lv25), axis=0)
|
|
lv27: R.Tensor((2, 224), dtype="float32") = R.reshape(lv26, R.shape([2, 224]))
|
|
lv28: R.Tensor((224,), dtype="float32") = R.add(lv16, R.const(1.0, "float32"))
|
|
lv29: R.Tensor((224,), dtype="float32") = R.subtract(R.const(2.0, "float32"), lv16)
|
|
lv30: R.Tensor((448,), dtype="float32") = R.concat((lv28, lv29), axis=0)
|
|
lv31: R.Tensor((2, 224), dtype="float32") = R.reshape(lv30, R.shape([2, 224]))
|
|
lv32: R.Tensor((2, 224), dtype="float32") = R.multiply(
|
|
lv31, R.const(-0.75, "float32")
|
|
)
|
|
lv33: R.Tensor((2, 224), dtype="float32") = R.subtract(
|
|
lv32, R.const(-3.75, "float32")
|
|
)
|
|
lv34: R.Tensor((2, 224), dtype="float32") = R.multiply(lv33, lv31)
|
|
lv35: R.Tensor((2, 224), dtype="float32") = R.add(lv34, R.const(-6.0, "float32"))
|
|
lv36: R.Tensor((2, 224), dtype="float32") = R.multiply(lv35, lv31)
|
|
lv37: R.Tensor((2, 224), dtype="float32") = R.subtract(
|
|
lv36, R.const(-3.0, "float32")
|
|
)
|
|
lv38: R.Tensor((2, 224), dtype="float32") = R.multiply(
|
|
lv27, R.const(1.25, "float32")
|
|
)
|
|
lv39: R.Tensor((2, 224), dtype="float32") = R.subtract(
|
|
lv38, R.const(2.25, "float32")
|
|
)
|
|
lv40: R.Tensor((2, 224), dtype="float32") = R.multiply(lv39, lv27)
|
|
lv41: R.Tensor((2, 224), dtype="float32") = R.multiply(lv40, lv27)
|
|
lv42: R.Tensor((2, 224), dtype="float32") = R.add(lv41, R.const(1.0, "float32"))
|
|
lv43: R.Tensor((1, 224), dtype="float32") = R.strided_slice(
|
|
lv37,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv44: R.Tensor((1, 224), dtype="float32") = R.strided_slice(
|
|
lv37,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv45: R.Tensor((224,), dtype="float32") = R.squeeze(lv43, axis=[0])
|
|
lv46: R.Tensor((224,), dtype="float32") = R.squeeze(lv44, axis=[0])
|
|
lv47: R.Tensor((1, 224), dtype="float32") = R.strided_slice(
|
|
lv42,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv48: R.Tensor((1, 224), dtype="float32") = R.strided_slice(
|
|
lv42,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv49: R.Tensor((224,), dtype="float32") = R.squeeze(lv47, axis=[0])
|
|
lv50: R.Tensor((224,), dtype="float32") = R.squeeze(lv48, axis=[0])
|
|
lv51: R.Tensor((224, 1), dtype="float32") = R.subtract(
|
|
R.const(1.0, "float32"), lv14
|
|
)
|
|
lv52: R.Tensor((448, 1), dtype="float32") = R.concat((lv14, lv51), axis=0)
|
|
lv53: R.Tensor((2, 224, 1), dtype="float32") = R.reshape(lv52, R.shape([2, 224, 1]))
|
|
lv54: R.Tensor((224, 1), dtype="float32") = R.add(lv14, R.const(1.0, "float32"))
|
|
lv55: R.Tensor((224, 1), dtype="float32") = R.subtract(
|
|
R.const(2.0, "float32"), lv14
|
|
)
|
|
lv56: R.Tensor((448, 1), dtype="float32") = R.concat((lv54, lv55), axis=0)
|
|
lv57: R.Tensor((2, 224, 1), dtype="float32") = R.reshape(lv56, R.shape([2, 224, 1]))
|
|
lv58: R.Tensor((2, 224, 1), dtype="float32") = R.multiply(
|
|
lv57, R.const(-0.75, "float32")
|
|
)
|
|
lv59: R.Tensor((2, 224, 1), dtype="float32") = R.subtract(
|
|
lv58, R.const(-3.75, "float32")
|
|
)
|
|
lv60: R.Tensor((2, 224, 1), dtype="float32") = R.multiply(lv59, lv57)
|
|
lv61: R.Tensor((2, 224, 1), dtype="float32") = R.add(lv60, R.const(-6.0, "float32"))
|
|
lv62: R.Tensor((2, 224, 1), dtype="float32") = R.multiply(lv61, lv57)
|
|
lv63: R.Tensor((2, 224, 1), dtype="float32") = R.subtract(
|
|
lv62, R.const(-3.0, "float32")
|
|
)
|
|
lv64: R.Tensor((2, 224, 1), dtype="float32") = R.multiply(
|
|
lv53, R.const(1.25, "float32")
|
|
)
|
|
lv65: R.Tensor((2, 224, 1), dtype="float32") = R.subtract(
|
|
lv64, R.const(2.25, "float32")
|
|
)
|
|
lv66: R.Tensor((2, 224, 1), dtype="float32") = R.multiply(lv65, lv53)
|
|
lv67: R.Tensor((2, 224, 1), dtype="float32") = R.multiply(lv66, lv53)
|
|
lv68: R.Tensor((2, 224, 1), dtype="float32") = R.add(lv67, R.const(1.0, "float32"))
|
|
lv69: R.Tensor((1, 224, 1), dtype="float32") = R.strided_slice(
|
|
lv63,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv70: R.Tensor((1, 224, 1), dtype="float32") = R.strided_slice(
|
|
lv63,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv71: R.Tensor((224, 1), dtype="float32") = R.squeeze(lv69, axis=[0])
|
|
lv72: R.Tensor((224, 1), dtype="float32") = R.squeeze(lv70, axis=[0])
|
|
lv73: R.Tensor((1, 224, 1), dtype="float32") = R.strided_slice(
|
|
lv68,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv74: R.Tensor((1, 224, 1), dtype="float32") = R.strided_slice(
|
|
lv68,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv75: R.Tensor((224, 1), dtype="float32") = R.squeeze(lv73, axis=[0])
|
|
lv76: R.Tensor((224, 1), dtype="float32") = R.squeeze(lv74, axis=[0])
|
|
lv77: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv19, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv78: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv22, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv79: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv78, axis=3, mode="fast"
|
|
)
|
|
lv80: R.Tensor((224,), dtype="int64") = R.squeeze(lv77, axis=None)
|
|
lv81: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv79, lv80, axis=2, mode="fast"
|
|
)
|
|
lv82: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv19, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv83: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv17, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv84: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv83, axis=3, mode="fast"
|
|
)
|
|
lv85: R.Tensor((224,), dtype="int64") = R.squeeze(lv82, axis=None)
|
|
lv86: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv84, lv85, axis=2, mode="fast"
|
|
)
|
|
lv87: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv19, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv88: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv23, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv89: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv88, axis=3, mode="fast"
|
|
)
|
|
lv90: R.Tensor((224,), dtype="int64") = R.squeeze(lv87, axis=None)
|
|
lv91: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv89, lv90, axis=2, mode="fast"
|
|
)
|
|
lv92: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv19, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv93: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv24, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv94: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv93, axis=3, mode="fast"
|
|
)
|
|
lv95: R.Tensor((224,), dtype="int64") = R.squeeze(lv92, axis=None)
|
|
lv96: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv94, lv95, axis=2, mode="fast"
|
|
)
|
|
lv97: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv81, lv45)
|
|
lv98: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv86, lv49)
|
|
lv99: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv97, lv98)
|
|
lv100: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv91, lv50)
|
|
lv101: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv99, lv100)
|
|
lv102: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv96, lv46)
|
|
lv103: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv101, lv102)
|
|
lv104: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv18, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv105: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv22, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv106: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv105, axis=3, mode="fast"
|
|
)
|
|
lv107: R.Tensor((224,), dtype="int64") = R.squeeze(lv104, axis=None)
|
|
lv108: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv106, lv107, axis=2, mode="fast"
|
|
)
|
|
lv109: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv18, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv110: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv17, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv111: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv110, axis=3, mode="fast"
|
|
)
|
|
lv112: R.Tensor((224,), dtype="int64") = R.squeeze(lv109, axis=None)
|
|
lv113: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv111, lv112, axis=2, mode="fast"
|
|
)
|
|
lv114: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv18, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv115: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv23, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv116: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv115, axis=3, mode="fast"
|
|
)
|
|
lv117: R.Tensor((224,), dtype="int64") = R.squeeze(lv114, axis=None)
|
|
lv118: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv116, lv117, axis=2, mode="fast"
|
|
)
|
|
lv119: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv18, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv120: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv24, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv121: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv120, axis=3, mode="fast"
|
|
)
|
|
lv122: R.Tensor((224,), dtype="int64") = R.squeeze(lv119, axis=None)
|
|
lv123: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv121, lv122, axis=2, mode="fast"
|
|
)
|
|
lv124: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv108, lv45)
|
|
lv125: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv113, lv49)
|
|
lv126: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv124, lv125)
|
|
lv127: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv118, lv50)
|
|
lv128: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv126, lv127)
|
|
lv129: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv123, lv46)
|
|
lv130: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv128, lv129)
|
|
lv131: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv20, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv132: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv22, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv133: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv132, axis=3, mode="fast"
|
|
)
|
|
lv134: R.Tensor((224,), dtype="int64") = R.squeeze(lv131, axis=None)
|
|
lv135: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv133, lv134, axis=2, mode="fast"
|
|
)
|
|
lv136: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv20, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv137: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv17, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv138: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv137, axis=3, mode="fast"
|
|
)
|
|
lv139: R.Tensor((224,), dtype="int64") = R.squeeze(lv136, axis=None)
|
|
lv140: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv138, lv139, axis=2, mode="fast"
|
|
)
|
|
lv141: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv20, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv142: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv23, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv143: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv142, axis=3, mode="fast"
|
|
)
|
|
lv144: R.Tensor((224,), dtype="int64") = R.squeeze(lv141, axis=None)
|
|
lv145: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv143, lv144, axis=2, mode="fast"
|
|
)
|
|
lv146: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv20, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv147: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv24, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv148: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv147, axis=3, mode="fast"
|
|
)
|
|
lv149: R.Tensor((224,), dtype="int64") = R.squeeze(lv146, axis=None)
|
|
lv150: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv148, lv149, axis=2, mode="fast"
|
|
)
|
|
lv151: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv135, lv45)
|
|
lv152: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv140, lv49)
|
|
lv153: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv151, lv152)
|
|
lv154: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv145, lv50)
|
|
lv155: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv153, lv154)
|
|
lv156: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv150, lv46)
|
|
lv157: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv155, lv156)
|
|
lv158: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv21, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv159: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv22, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv160: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv159, axis=3, mode="fast"
|
|
)
|
|
lv161: R.Tensor((224,), dtype="int64") = R.squeeze(lv158, axis=None)
|
|
lv162: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv160, lv161, axis=2, mode="fast"
|
|
)
|
|
lv163: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv21, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv164: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv17, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv165: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv164, axis=3, mode="fast"
|
|
)
|
|
lv166: R.Tensor((224,), dtype="int64") = R.squeeze(lv163, axis=None)
|
|
lv167: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv165, lv166, axis=2, mode="fast"
|
|
)
|
|
lv168: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv21, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv169: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv23, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv170: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv169, axis=3, mode="fast"
|
|
)
|
|
lv171: R.Tensor((224,), dtype="int64") = R.squeeze(lv168, axis=None)
|
|
lv172: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv170, lv171, axis=2, mode="fast"
|
|
)
|
|
lv173: R.Tensor((224, 1), dtype="int64") = R.clip(
|
|
lv21, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv174: R.Tensor((224,), dtype="int64") = R.clip(
|
|
lv24, R.prim_value(0), R.prim_value(111)
|
|
)
|
|
lv175: R.Tensor((1, 3, 112, 224), dtype="float32") = R.take(
|
|
input, lv174, axis=3, mode="fast"
|
|
)
|
|
lv176: R.Tensor((224,), dtype="int64") = R.squeeze(lv173, axis=None)
|
|
lv177: R.Tensor((1, 3, 224, 224), dtype="float32") = R.take(
|
|
lv175, lv176, axis=2, mode="fast"
|
|
)
|
|
lv178: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv162, lv45)
|
|
lv179: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv167, lv49)
|
|
lv180: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv178, lv179)
|
|
lv181: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv172, lv50)
|
|
lv182: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv180, lv181)
|
|
lv183: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv177, lv46)
|
|
lv184: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv182, lv183)
|
|
lv185: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv103, lv71)
|
|
lv186: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv130, lv75)
|
|
lv187: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv185, lv186)
|
|
lv188: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv157, lv76)
|
|
lv189: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv187, lv188)
|
|
lv190: R.Tensor((1, 3, 224, 224), dtype="float32") = R.multiply(lv184, lv72)
|
|
lv191: R.Tensor((1, 3, 224, 224), dtype="float32") = R.add(lv189, lv190)
|
|
gv: R.Tuple(R.Tensor((1, 3, 224, 224), dtype="float32")) = (lv191,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 112, 112, dtype=torch.float32),)
|
|
verify_model(InterpolateBilinear(), example_args, {}, expected_bilinear)
|
|
verify_model(InterpolateNearest(), example_args, {}, expected_nearest)
|
|
verify_model(InterpolateBicubic(), example_args, {}, expected_bicubic)
|
|
|
|
|
|
def test_interpolate_antialiased():
|
|
"""Test bilinear interpolation with antialiasing enabled."""
|
|
|
|
class InterpolateBilinearAA(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.interpolate(
|
|
input, size=(64, 64), mode="bilinear", align_corners=False, antialias=True
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_bilinear_aa:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 64, 64), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 64, 64), dtype="float32") = R.image.resize2d(
|
|
input,
|
|
R.shape([64, 64]),
|
|
roi=[T.float32(0.0), T.float32(0.0), T.float32(0.0), T.float32(0.0)],
|
|
layout="NCHW",
|
|
method="linear",
|
|
coordinate_transformation_mode="half_pixel",
|
|
rounding_method="round",
|
|
cubic_alpha=-0.75,
|
|
cubic_exclude=0,
|
|
extrapolation_value=0.0,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 64, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 32, 32, dtype=torch.float32),)
|
|
verify_model(InterpolateBilinearAA(), example_args, {}, expected_bilinear_aa)
|
|
|
|
|
|
def test_mean():
|
|
class Mean(Module):
|
|
def forward(self, input):
|
|
return input.mean(-1)
|
|
|
|
class MeanKeepDim(Module):
|
|
def forward(self, input: torch.Tensor):
|
|
return input.mean(-1, keepdim=True)
|
|
|
|
class MeanWithoutDim(Module):
|
|
def forward(self, input: torch.Tensor):
|
|
return input.mean()
|
|
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((256,), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((256,), dtype="float32") = R.mean(inp_0, axis=[-1], keepdims=False)
|
|
gv: R.Tuple(R.Tensor((256,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((256, 1), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((256, 1), dtype="float32") = R.mean(inp_0, axis=[-1], keepdims=True)
|
|
gv: R.Tuple(R.Tensor((256, 1), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.mean(inp_0, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(256, 256, dtype=torch.float32),)
|
|
verify_model(Mean(), example_args, {}, Expected1)
|
|
verify_model(MeanKeepDim(), example_args, {}, Expected2)
|
|
verify_model(MeanWithoutDim(), example_args, {}, Expected3)
|
|
|
|
|
|
def test_median():
|
|
class Median(Module):
|
|
def forward(self, input):
|
|
return input.median(-1)
|
|
|
|
class MedianKeepDim(Module):
|
|
def forward(self, input):
|
|
return input.median(-1, keepdim=True)
|
|
|
|
class MedianWithoutDim(Module):
|
|
def forward(self, input):
|
|
return input.median()
|
|
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((256,), dtype="float32"), R.Tensor((256,), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tuple(R.Tensor((256,), dtype="float32"), R.Tensor((256,), dtype="int64")) = (
|
|
R.median(inp_0, axis=[-1], keepdims=False)
|
|
)
|
|
lv1: R.Tensor((256,), dtype="float32") = lv[0]
|
|
lv2: R.Tensor((256,), dtype="int64") = lv[1]
|
|
gv: R.Tuple(R.Tensor((256,), dtype="float32"), R.Tensor((256,), dtype="int64")) = (
|
|
lv1,
|
|
lv2,
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((256, 1), dtype="float32"), R.Tensor((256, 1), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tuple(
|
|
R.Tensor((256, 1), dtype="float32"), R.Tensor((256, 1), dtype="int64")
|
|
) = R.median(inp_0, axis=[-1], keepdims=True)
|
|
lv1: R.Tensor((256, 1), dtype="float32") = lv[0]
|
|
lv2: R.Tensor((256, 1), dtype="int64") = lv[1]
|
|
gv: R.Tuple(
|
|
R.Tensor((256, 1), dtype="float32"), R.Tensor((256, 1), dtype="int64")
|
|
) = (lv1, lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.median(inp_0, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(256, 256, dtype=torch.float32),)
|
|
verify_model(Median(), example_args, {}, Expected1)
|
|
verify_model(MedianKeepDim(), example_args, {}, Expected2)
|
|
verify_model(MedianWithoutDim(), example_args, {}, Expected3)
|
|
|
|
|
|
def test_sum():
|
|
class Sum(Module):
|
|
def forward(self, x):
|
|
return torch.sum(x, (2, 1))
|
|
|
|
class SumKeepDim(Module):
|
|
def forward(self, x):
|
|
return torch.sum(x, (2, 1), keepdim=True)
|
|
|
|
class SumWithoutDim(Module):
|
|
def forward(self, x):
|
|
return torch.sum(x)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 4), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4), dtype="float32") = R.sum(inp_0, axis=[2, 1], keepdims=False)
|
|
gv: R.Tuple(R.Tensor((1, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 1, 1, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 1, 1, 4), dtype="float32") = R.sum(
|
|
inp_0, axis=[2, 1], keepdims=True
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 1, 1, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected3:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.sum(inp_0, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(Sum(), example_args, {}, expected1)
|
|
verify_model(SumKeepDim(), example_args, {}, expected2)
|
|
verify_model(SumWithoutDim(), example_args, {}, expected3)
|
|
|
|
|
|
def test_argmax_argmin():
|
|
example_args = (torch.randn(256, 256, dtype=torch.float32),)
|
|
|
|
class Argmax1(Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.argmax(input, dim=-1)
|
|
|
|
class Argmax2(Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.argmax(input, dim=-1, keepdim=True)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_argmax1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((256,), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((256,), dtype="int64") = R.argmax(inp_0, axis=-1, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((256,), dtype="int64")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_argmax2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((256, 1), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((256, 1), dtype="int64") = R.argmax(inp_0, axis=-1, keepdims=True)
|
|
gv: R.Tuple(R.Tensor((256, 1), dtype="int64")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Argmax1(), example_args, {}, expected_argmax1)
|
|
verify_model(Argmax2(), example_args, {}, expected_argmax2)
|
|
|
|
class Argmin1(Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.argmin(input)
|
|
|
|
class Argmin2(Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
def forward(self, input):
|
|
return torch.argmin(input, keepdim=True)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_argmin1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="int64") = R.argmin(inp_0, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="int64")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_argmin2:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 1), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 1), dtype="int64") = R.argmin(inp_0, axis=None, keepdims=True)
|
|
gv: R.Tuple(R.Tensor((1, 1), dtype="int64")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Argmin1(), example_args, {}, expected_argmin1)
|
|
verify_model(Argmin2(), example_args, {}, expected_argmin2)
|
|
|
|
|
|
def test_cat_concat():
|
|
class Cat0(Module):
|
|
def forward(self, x, y):
|
|
return torch.cat((x, y))
|
|
|
|
class Cat1(Module):
|
|
def forward(self, x, y):
|
|
return torch.cat((x, y), dim=1)
|
|
|
|
class Cat2(Module):
|
|
def forward(self, x, y):
|
|
return torch.cat((x, y), 1)
|
|
|
|
class Cat3(Module):
|
|
def forward(self, x, y):
|
|
return torch.concat((x, y), dim=0)
|
|
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
inp_1: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 3), dtype="float32") = R.concat((inp_0, inp_1), axis=0)
|
|
gv: R.Tuple(R.Tensor((4, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
inp_1: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 6), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 6), dtype="float32") = R.concat((inp_0, inp_1), axis=1)
|
|
gv: R.Tuple(R.Tensor((2, 6), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 3, dtype=torch.float32), torch.randn(2, 3, dtype=torch.float32))
|
|
verify_model(Cat0(), example_args, {}, Expected1)
|
|
verify_model(Cat1(), example_args, {}, Expected2)
|
|
verify_model(Cat2(), example_args, {}, Expected2)
|
|
verify_model(Cat3(), example_args, {}, Expected1)
|
|
|
|
|
|
def test_cumsum():
|
|
class Cumsum(Module):
|
|
def forward(self, input):
|
|
return torch.cumsum(input, dim=1, dtype=torch.int32)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3, 4), dtype="int32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 2, 3, 4), dtype="int32") = R.cumsum(input_1, axis=1, dtype="int32")
|
|
gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="int32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(Cumsum(), example_args, {}, expected1)
|
|
|
|
|
|
def test_expand():
|
|
class Expand1(Module):
|
|
def forward(self, x):
|
|
return x.expand(4, 2, 3, 4)
|
|
|
|
class Expand2(Module):
|
|
def forward(self, x):
|
|
return x.expand(4, -1, -1, 4)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 2, 3, 4), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 2, 3, 4), dtype="float32") = R.broadcast_to(x, (4, 2, 3, 4))
|
|
gv: R.Tuple(R.Tensor((4, 2, 3, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(Expand1(), example_args, {}, expected1)
|
|
verify_model(Expand2(), example_args, {}, expected1)
|
|
|
|
|
|
def test_flatten():
|
|
class Flatten(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.f = torch.nn.Flatten(2, -1)
|
|
|
|
def forward(self, input):
|
|
return self.f(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 100), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 100), dtype="float32") = R.reshape(input_1, (1, 3, 100))
|
|
gv: R.Tuple(R.Tensor((1, 3, 100), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Flatten(), example_args, {}, expected1)
|
|
|
|
|
|
def test_meshgrid():
|
|
class Meshgrid1(Module):
|
|
def forward(self, input1, input2):
|
|
return torch.meshgrid((input1, input2), indexing="ij")
|
|
|
|
class Meshgrid2(Module):
|
|
def forward(self, input1, input2):
|
|
return torch.meshgrid((input1, input2), indexing="xy")
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input1: R.Tensor((3,), dtype="float32"), input2: R.Tensor((3,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1), dtype="float32") = R.reshape(input1, R.shape([3, 1]))
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(lv, R.shape([3, 3]))
|
|
lv2: R.Tensor((1, 3), dtype="float32") = R.reshape(input2, R.shape([1, 3]))
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(lv2, R.shape([3, 3]))
|
|
gv: R.Tuple(
|
|
R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
|
|
) = (lv1, lv3)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
input1: R.Tensor((3,), dtype="float32"), input2: R.Tensor((3,), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1), dtype="float32") = R.reshape(input2, R.shape([3, 1]))
|
|
lv1: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(lv, R.shape([3, 3]))
|
|
lv2: R.Tensor((1, 3), dtype="float32") = R.reshape(input1, R.shape([1, 3]))
|
|
lv3: R.Tensor((3, 3), dtype="float32") = R.broadcast_to(lv2, R.shape([3, 3]))
|
|
gv: R.Tuple(
|
|
R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
|
|
) = (lv3, lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(3, dtype=torch.float32),
|
|
torch.randn(3, dtype=torch.float32),
|
|
)
|
|
verify_model(Meshgrid1(), example_args, {}, expected1)
|
|
verify_model(Meshgrid2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_permute():
|
|
class Permute1(Module):
|
|
def forward(self, x):
|
|
return x.permute(0, 3, 2, 1)
|
|
|
|
class Permute2(Module):
|
|
def forward(self, x):
|
|
return torch.permute(x, (0, 3, 2, 1))
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 4, 3, 2), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4, 3, 2), dtype="float32") = R.permute_dims(x, axes=[0, 3, 2, 1])
|
|
gv: R.Tuple(R.Tensor((1, 4, 3, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(Permute1(), example_args, {}, expected1)
|
|
verify_model(Permute2(), example_args, {}, expected1)
|
|
|
|
|
|
def test_repeat():
|
|
class Tile1(Module):
|
|
def forward(self, x: torch.Tensor):
|
|
return x.repeat(2)
|
|
|
|
class Tile2(Module):
|
|
def forward(self, x: torch.Tensor):
|
|
return x.repeat(4, 2)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((3,), dtype="float32")) -> R.Tuple(R.Tensor((6,), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((6,), dtype="float32") = R.tile(x, 2)
|
|
gv: R.Tuple(R.Tensor((6,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 6), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 6), dtype="float32") = R.tile(x, [4, 2])
|
|
gv: R.Tuple(R.Tensor((4, 6), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(3, dtype=torch.float32),)
|
|
verify_model(Tile1(), example_args, {}, expected1)
|
|
|
|
example_args = (torch.randn(1, 3, dtype=torch.float32),)
|
|
verify_model(Tile2(), example_args, {}, expected2)
|
|
|
|
example_args = (torch.randn(1, 3, dtype=torch.float32),)
|
|
verify_model(Tile2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_reshape():
|
|
class Reshape(Module):
|
|
def forward(self, x):
|
|
return x.reshape(2, 12)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 12), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
|
|
gv: R.Tuple(R.Tensor((2, 12), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(Reshape(), example_args, {}, expected1)
|
|
|
|
|
|
def test_reshape_as():
|
|
class ReshapeAs(Module):
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor):
|
|
return x.reshape_as(y)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 2, 3, 4), dtype="float32"),
|
|
y: R.Tensor((2, 12), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 12), dtype="float32")):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
|
|
gv: R.Tuple(R.Tensor((2, 12), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 2, 3, 4, dtype=torch.float32),
|
|
torch.randn(2, 12, dtype=torch.float32),
|
|
)
|
|
verify_model(ReshapeAs(), example_args, {}, expected1)
|
|
|
|
|
|
def test_roll():
|
|
class Roll1(Module):
|
|
def forward(self, x):
|
|
return torch.roll(x, 1)
|
|
|
|
class Roll2(Module):
|
|
def forward(self, x):
|
|
return torch.roll(x, -1, 0)
|
|
|
|
class Roll3(Module):
|
|
def forward(self, x):
|
|
return torch.roll(x, shifts=(2, 1), dims=(0, 1))
|
|
|
|
# Test case 1: torch.roll(x, 1)
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 2), dtype="int64")) -> R.Tuple(R.Tensor((4, 2), dtype="int64")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((8,), dtype="int64") = R.reshape(x, R.shape([8]))
|
|
lv1: R.Tensor((8,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(8), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv2: R.Tensor((8,), dtype="int64") = R.add(lv1, R.const(7, "int64"))
|
|
lv3: R.Tensor((8,), dtype="int64") = R.mod(lv2, R.const(8, "int64"))
|
|
lv4: R.Tensor((8,), dtype="int64") = R.take(lv, lv3, axis=0, mode="fast")
|
|
lv5: R.Tensor((4, 2), dtype="int64") = R.reshape(lv4, R.shape([4, 2]))
|
|
gv: R.Tuple(R.Tensor((4, 2), dtype="int64")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 2: torch.roll(x, -1, 0)
|
|
@I.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 2), dtype="int64")) -> R.Tuple(R.Tensor((4, 2), dtype="int64")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(4), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((4,), dtype="int64") = R.add(lv, R.const(1, "int64"))
|
|
lv2: R.Tensor((4,), dtype="int64") = R.mod(lv1, R.const(4, "int64"))
|
|
lv3: R.Tensor((4, 2), dtype="int64") = R.take(x, lv2, axis=0, mode="fast")
|
|
gv: R.Tuple(R.Tensor((4, 2), dtype="int64")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 3: torch.roll(x, shifts=(2,1), dims=(0,1))
|
|
@I.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 2), dtype="int64")) -> R.Tuple(R.Tensor((4, 2), dtype="int64")):
|
|
with R.dataflow():
|
|
# First roll along dim=0 with shift=2
|
|
lv: R.Tensor((4,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(4), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((4,), dtype="int64") = R.add(lv, R.const(2, "int64"))
|
|
lv2: R.Tensor((4,), dtype="int64") = R.mod(lv1, R.const(4, "int64"))
|
|
lv3: R.Tensor((4, 2), dtype="int64") = R.take(x, lv2, axis=0, mode="fast")
|
|
# Second roll along dim=1 with shift=1
|
|
lv4: R.Tensor((2,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(2), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv5: R.Tensor((2,), dtype="int64") = R.add(lv4, R.const(1, "int64"))
|
|
lv6: R.Tensor((2,), dtype="int64") = R.mod(lv5, R.const(2, "int64"))
|
|
lv7: R.Tensor((4, 2), dtype="int64") = R.take(lv3, lv6, axis=1, mode="fast")
|
|
gv: R.Tuple(R.Tensor((4, 2), dtype="int64")) = (lv7,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test inputs
|
|
example_input = torch.randint(0, 10, (4, 2), dtype=torch.int64)
|
|
|
|
# Run verification for each case
|
|
verify_model(Roll1(), (example_input,), {}, Expected1)
|
|
verify_model(Roll2(), (example_input,), {}, Expected2)
|
|
verify_model(Roll3(), (example_input,), {}, Expected3)
|
|
|
|
|
|
def test_select_slice():
|
|
class Slice1(Module):
|
|
def forward(self, x):
|
|
return x[0, 1::2, :, :3]
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 10, 3), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 10, 10), dtype="float32") = R.take(
|
|
x, R.const(0, "int64"), axis=0, mode="fast"
|
|
)
|
|
lv1: R.Tensor((1, 10, 10), dtype="float32") = R.strided_slice(
|
|
lv,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(9223372036854775807),),
|
|
(R.prim_value(2),),
|
|
assume_inbound=False,
|
|
)
|
|
lv2: R.Tensor((1, 10, 3), dtype="float32") = R.strided_slice(
|
|
lv1,
|
|
(R.prim_value(2),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 10, 3), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Slice2(Module):
|
|
def forward(self, x):
|
|
return x[:, None, None, :, None]
|
|
|
|
@I.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(x: R.Tensor((8, 16), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((8, 1, 1, 16, 1), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((8, 1, 16), dtype="float32") = R.expand_dims(x, axis=[1])
|
|
lv1: R.Tensor((8, 1, 1, 16), dtype="float32") = R.expand_dims(lv, axis=[2])
|
|
lv2: R.Tensor((8, 1, 1, 16, 1), dtype="float32") = R.expand_dims(lv1, axis=[4])
|
|
gv: R.Tuple(R.Tensor((8, 1, 1, 16, 1), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Slice1(), example_args, {}, expected1)
|
|
|
|
example_args = (torch.randn(8, 16, dtype=torch.float32),)
|
|
verify_model(Slice2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_slice_scatter():
|
|
class SliceScatter1(Module):
|
|
def forward(self, input, src):
|
|
return torch.slice_scatter(input, src, dim=1, start=1, end=7, step=2)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((8, 8, 10, 10), dtype="float32"),
|
|
b: R.Tensor((8, 3, 10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((8, 8, 10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((8, 8, 10, 10), dtype="float32") = R.slice_scatter(
|
|
a, b, R.prim_value(1), R.prim_value(7), R.prim_value(2), axis=1
|
|
)
|
|
gv: R.Tuple(R.Tensor((8, 8, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class SliceScatter2(Module):
|
|
def forward(self, input, src):
|
|
return torch.slice_scatter(input, src, dim=0, start=0, end=6, step=1)
|
|
|
|
@I.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((8, 16), dtype="float32"), b: R.Tensor((6, 16), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((8, 16), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((8, 16), dtype="float32") = R.slice_scatter(
|
|
a, b, R.prim_value(0), R.prim_value(6), R.prim_value(1), axis=0
|
|
)
|
|
gv: R.Tuple(R.Tensor((8, 16), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class SliceScatterNegative(Module):
|
|
def forward(self, input, src):
|
|
return torch.slice_scatter(input, src, dim=1, start=0, end=-2, step=1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_slice_scatter:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((2, 5), dtype="float32"), b: R.Tensor((2, 3), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((2, 5), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 5), dtype="float32") = R.slice_scatter(
|
|
a, b, R.prim_value(0), R.prim_value(3), R.prim_value(1), axis=1
|
|
)
|
|
gv: R.Tuple(R.Tensor((2, 5), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(8, 8, 10, 10, dtype=torch.float32), torch.randn(8, 3, 10, 10))
|
|
verify_model(SliceScatter1(), example_args, {}, expected1)
|
|
|
|
example_args = (torch.randn(8, 16, dtype=torch.float32), torch.randn(6, 16))
|
|
verify_model(SliceScatter2(), example_args, {}, expected2)
|
|
|
|
example_args = (torch.randn(2, 5, dtype=torch.float32), torch.randn(2, 3, dtype=torch.float32))
|
|
verify_model(SliceScatterNegative(), example_args, {}, expected_slice_scatter)
|
|
|
|
|
|
def test_slice_with_symbolic_end():
|
|
"""_slice correctly handles symbolic end values from dynamic shapes."""
|
|
|
|
class SliceIdentityModel(torch.nn.Module):
|
|
def forward(self, x):
|
|
# x[:, :x.size(1)] is an identity slice that torch.export emits
|
|
# as slice(x, 1, 0, sym_size_int(x, 1), 1) with dynamic shapes.
|
|
seq_len = x.size(1)
|
|
return x[:, :seq_len] + 0.0 # +0.0 to ensure output is a new tensor
|
|
|
|
# The identity slice is elided; only x + 0.0 remains.
|
|
@I.ir_module
|
|
class ExpectedIdentity:
|
|
@R.function
|
|
def main(x: R.Tensor(("s0", "s1", 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor(("s0", "s1", 4), dtype="float32")
|
|
):
|
|
s0 = T.int64()
|
|
s1 = T.int64()
|
|
R.func_attr({"tir_var_lower_bound": {"s27": 2, "s77": 2}})
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0, s1, 4), dtype="float32") = R.add(x, R.const(0.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((s0, s1, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 8, 4, dtype=torch.float32),)
|
|
batch = torch.export.Dim("batch", min=2)
|
|
seq = torch.export.Dim("seq", min=2)
|
|
dynamic_shapes = {"x": {0: batch, 1: seq}}
|
|
|
|
verify_model(
|
|
SliceIdentityModel(),
|
|
example_args,
|
|
{},
|
|
ExpectedIdentity,
|
|
dynamic_shapes=dynamic_shapes,
|
|
map_free_vars=True,
|
|
)
|
|
|
|
class SliceStaticModel(torch.nn.Module):
|
|
def forward(self, x):
|
|
# A non-identity static slice
|
|
return x[:, :3]
|
|
|
|
@tvm.script.ir_module
|
|
class ExpectedStatic:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 8, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 4), dtype="float32") = R.strided_slice(
|
|
x,
|
|
axes=[1],
|
|
begin=[0],
|
|
end=[3],
|
|
strides=[1],
|
|
)
|
|
gv: R.Tuple(R.Tensor((2, 3, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args_static = (torch.randn(2, 8, 4, dtype=torch.float32),)
|
|
verify_model(SliceStaticModel(), example_args_static, {}, ExpectedStatic)
|
|
|
|
|
|
def test_split():
|
|
class Chunk(Module):
|
|
def forward(self, input):
|
|
return torch.chunk(input, 3, dim=1)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tuple(
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
) = R.split(input, indices_or_sections=[1, 2], axis=1)
|
|
lv1: R.Tensor((1, 1, 10, 10), dtype="float32") = lv[0]
|
|
lv2: R.Tensor((1, 1, 10, 10), dtype="float32") = lv[1]
|
|
lv3: R.Tensor((1, 1, 10, 10), dtype="float32") = lv[2]
|
|
gv: R.Tuple(
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
R.Tensor((1, 1, 10, 10), dtype="float32"),
|
|
) = (lv1, lv2, lv3)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Unbind1(Module):
|
|
def forward(self, data):
|
|
return torch.unbind(data)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(data: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv3: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv, axis=[0])
|
|
lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[0])
|
|
lv5: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv2, axis=[0])
|
|
gv: R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
) = (lv3, lv4, lv5)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Unbind2(Module):
|
|
def forward(self, data):
|
|
return torch.unbind(data, dim=1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(data: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv1: R.Tensor((3, 1, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv2: R.Tensor((3, 1, 10, 10), dtype="float32") = R.strided_slice(
|
|
data,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(3),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
lv3: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv, axis=[1])
|
|
lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[1])
|
|
lv5: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv2, axis=[1])
|
|
gv: R.Tuple(
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
R.Tensor((3, 10, 10), dtype="float32"),
|
|
) = (lv3, lv4, lv5)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Chunk(), example_args, {}, Expected)
|
|
|
|
example_args = (torch.randn(3, 3, 10, 10, dtype=torch.float32),)
|
|
verify_model(Unbind1(), example_args, {}, expected1)
|
|
verify_model(Unbind2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_squeeze():
|
|
class Squeeze1(Module):
|
|
def forward(self, input):
|
|
return input.squeeze(1)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 4, 1), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 4, 1), dtype="float32") = R.squeeze(inp_0, axis=[1])
|
|
gv: R.Tuple(R.Tensor((3, 4, 1), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Squeeze2(Module):
|
|
def forward(self, input):
|
|
return input.squeeze()
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(input: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 4), dtype="float32") = R.squeeze(input, axis=[0, 1, 2, 3])
|
|
gv: R.Tuple(R.Tensor((3, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Squeeze3(Module):
|
|
def forward(self, input):
|
|
return input.squeeze(2)
|
|
|
|
@I.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 1, 4, 1), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 1, 4, 1), dtype="float32") = R.squeeze(inp_0, axis=[2])
|
|
gv: R.Tuple(R.Tensor((3, 1, 4, 1), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(3, 1, 4, 1, dtype=torch.float32),)
|
|
|
|
verify_model(Squeeze1(), example_args, {}, Expected1)
|
|
verify_model(Squeeze2(), example_args, {}, Expected2)
|
|
verify_model(Squeeze3(), example_args, {}, Expected3)
|
|
|
|
|
|
def test_stack():
|
|
class Stack0(Module):
|
|
def forward(self, x, y):
|
|
return torch.stack((x, y)) # default dim=0
|
|
|
|
class Stack1(Module):
|
|
def forward(self, x, y):
|
|
return torch.stack((x, y), dim=1)
|
|
|
|
class Stack2(Module):
|
|
def forward(self, x, y):
|
|
return torch.stack((x, y), 1) # positional dim
|
|
|
|
class Stack3(Module):
|
|
def forward(self, x, y):
|
|
return torch.stack((x, y), dim=-1) # negative dim
|
|
|
|
@I.ir_module
|
|
class Expected0:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 3), dtype="float32") = R.concat((x, y), axis=0)
|
|
lv1: R.Tensor((2, 2, 3), dtype="float32") = R.reshape(lv, R.shape([2, 2, 3]))
|
|
gv: R.Tuple(R.Tensor((2, 2, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 6), dtype="float32") = R.concat((x, y), axis=1)
|
|
lv1: R.Tensor((2, 2, 3), dtype="float32") = R.reshape(lv, R.shape([2, 2, 3]))
|
|
gv: R.Tuple(R.Tensor((2, 2, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 3, 2), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 1), dtype="float32") = R.expand_dims(x, axis=[2])
|
|
lv1: R.Tensor((2, 3, 1), dtype="float32") = R.expand_dims(y, axis=[2])
|
|
lv2: R.Tensor((2, 3, 2), dtype="float32") = R.concat((lv, lv1), axis=-1)
|
|
gv: R.Tuple(R.Tensor((2, 3, 2), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 3, dtype=torch.float32), torch.randn(2, 3, dtype=torch.float32))
|
|
|
|
verify_model(Stack0(), example_args, {}, Expected0)
|
|
verify_model(Stack1(), example_args, {}, Expected1)
|
|
verify_model(Stack2(), example_args, {}, Expected1)
|
|
verify_model(Stack3(), example_args, {}, Expected3)
|
|
|
|
|
|
def test_tile():
|
|
class Tile1(Module):
|
|
def forward(self, x):
|
|
return x.tile((2,))
|
|
|
|
class Tile2(Module):
|
|
def forward(self, x):
|
|
return x.tile(4, 2)
|
|
|
|
class Tile3(Module):
|
|
def forward(self, x):
|
|
return torch.tile(x, (4, 2))
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 6), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 6), dtype="float32") = R.tile(x, repeats=[1, 2])
|
|
gv: R.Tuple(R.Tensor((1, 6), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 6), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 6), dtype="float32") = R.tile(x, repeats=[4, 2])
|
|
gv: R.Tuple(R.Tensor((4, 6), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, dtype=torch.float32),)
|
|
verify_model(Tile1(), example_args, {}, expected1)
|
|
verify_model(Tile2(), example_args, {}, expected2)
|
|
verify_model(Tile3(), example_args, {}, expected2)
|
|
|
|
|
|
def test_transpose():
|
|
class Transpose(Module):
|
|
def forward(self, x):
|
|
return x.transpose(1, 3)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 4, 3, 2), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4, 3, 2), dtype="float32") = R.permute_dims(x, axes=[0, 3, 2, 1])
|
|
gv: R.Tuple(R.Tensor((1, 4, 3, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(Transpose(), example_args, {}, expected1)
|
|
|
|
|
|
def test_unsqueeze():
|
|
class Unsqueeze1(Module):
|
|
def forward(self, input):
|
|
return input.unsqueeze(1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 1, 3, 10, 10), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 1, 3, 10, 10), dtype="float32") = R.expand_dims(input_1, 1)
|
|
gv: R.Tuple(R.Tensor((1, 1, 3, 10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Unsqueeze2(Module):
|
|
def forward(self, input):
|
|
return input.unsqueeze(-1)
|
|
|
|
@tvm.script.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 10, 10, 1), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 10, 10, 1), dtype="float32") = R.expand_dims(input_1, -1)
|
|
gv: R.Tuple(R.Tensor((1, 3, 10, 10, 1), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
verify_model(Unsqueeze1(), example_args, {}, expected1)
|
|
verify_model(Unsqueeze2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_view():
|
|
class View(Module):
|
|
def forward(self, x):
|
|
return x.view(2, 12)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 12), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
|
|
gv: R.Tuple(R.Tensor((2, 12), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(View(), example_args, {}, expected1)
|
|
|
|
|
|
def test_as_strided():
|
|
class AsStrided(Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.as_strided.default(x, (3, 2, 2), (4, 2, 1))
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 2, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 2, 2), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 2, 2), dtype="float32") = R.reshape(x, (3, 2, 2))
|
|
gv: R.Tuple(R.Tensor((3, 2, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class AsStridedNonContiguous(Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.as_strided.default(x, (2, 2, 2), (6, 3, 1))
|
|
|
|
class AsStridedWithStorageOffset(Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.as_strided.default(x, (2, 2), (2, 1), 1)
|
|
|
|
example_args = (torch.randn(2, 2, 3, dtype=torch.float32),)
|
|
verify_model(AsStrided(), example_args, {}, Expected)
|
|
|
|
exported = export(AsStridedNonContiguous(), args=example_args)
|
|
with pytest.raises(AssertionError, match="non-contiguous stride"):
|
|
from_exported_program(exported)
|
|
|
|
example_args = (torch.randn(2, 2, dtype=torch.float32),)
|
|
exported = export(AsStridedWithStorageOffset(), args=example_args)
|
|
with pytest.raises(AssertionError, match="storage_offset"):
|
|
from_exported_program(exported)
|
|
|
|
|
|
def test_arange():
|
|
class Arange(Module):
|
|
def forward(self, input):
|
|
return torch.arange(0, 20, dtype=torch.int32)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((20,), dtype="int32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((20,), dtype="int32") = R.arange(0, 20, 1, dtype="int32")
|
|
gv: R.Tuple(R.Tensor((20,), dtype="int32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(Arange(), example_args, {}, Expected)
|
|
|
|
|
|
def test_hamming_window():
|
|
class HammingWindow(Module):
|
|
def forward(self, input):
|
|
return torch.hamming_window(20, True, dtype=torch.float32)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((20,), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((20,), dtype="float32") = R.hamming_window(
|
|
R.prim_value(20),
|
|
R.prim_value(True),
|
|
R.prim_value(T.float64(0.54000000000000004)),
|
|
R.prim_value(T.float64(0.46000000000000002)),
|
|
dtype="float32",
|
|
)
|
|
gv: R.Tuple(R.Tensor((20,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(HammingWindow(), example_args, {}, Expected)
|
|
|
|
|
|
def test_contiguous():
|
|
class Contiguous(Module):
|
|
def forward(self, input):
|
|
return input.contiguous()
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((10, 10), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (input,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(Contiguous(), example_args, {}, Expected)
|
|
|
|
|
|
def test_clone():
|
|
class Clone(Module):
|
|
def forward(self, input):
|
|
return torch.clone(input)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (input,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(Clone(), example_args, {}, Expected)
|
|
|
|
|
|
def test_empty():
|
|
class Empty(Module):
|
|
def forward(self, input):
|
|
return torch.empty((10, 10), dtype=torch.float32)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.zeros(
|
|
R.shape([10, 10]), dtype="float32"
|
|
)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(Empty(), example_args, {}, Expected)
|
|
|
|
|
|
def test_empty_without_dtype():
|
|
class EmptyWithoutDtype(Module):
|
|
def forward(self, input):
|
|
return torch.empty((5, 5))
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((5, 5), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 5), dtype="float32") = R.zeros(R.shape([5, 5]), dtype="float32")
|
|
gv: R.Tuple(R.Tensor((5, 5), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(EmptyWithoutDtype(), example_args, {}, Expected)
|
|
|
|
|
|
def test_fill():
|
|
class Fill(Module):
|
|
def forward(self, input: torch.Tensor):
|
|
return torch.fill(input, 1.5)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((10, 10), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype="float32") = R.full_like(
|
|
input, R.const(1.5, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(10, 10, dtype=torch.float32),)
|
|
verify_model(Fill(), example_args, {}, Expected)
|
|
|
|
|
|
def test_fill_inplace():
|
|
class FillInplace(Module):
|
|
def forward(self, input: torch.Tensor):
|
|
input.fill_(42.0)
|
|
return input
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((2, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.full_like(input, R.const(42.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 3, dtype=torch.float32),)
|
|
verify_model(FillInplace(), example_args, {}, Expected)
|
|
|
|
|
|
def test_masked_fill():
|
|
class Masked_Fill(Module):
|
|
def forward(self, input: torch.Tensor, mask: torch.Tensor):
|
|
return torch.masked_fill(input, mask, 0)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((128, 128), dtype="float32"), mask: R.Tensor((128, 128), dtype="bool")
|
|
) -> R.Tuple(R.Tensor((128, 128), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.const(0.0, "float32")
|
|
lv1: R.Tensor((128, 128), dtype="float32") = R.where(mask, lv, input)
|
|
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(128, 128, dtype=torch.float32),
|
|
torch.testing.make_tensor((128, 128), dtype=torch.bool, device="cpu"),
|
|
)
|
|
verify_model(Masked_Fill(), example_args, {}, Expected)
|
|
|
|
|
|
def test_masked_fill_inplace():
|
|
class Masked_Fill_Inplace(Module):
|
|
def forward(self, input: torch.Tensor, mask: torch.Tensor):
|
|
return input.masked_fill_(mask, 1.5)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((128, 128), dtype="float32"), mask: R.Tensor((128, 128), dtype="bool")
|
|
) -> R.Tuple(R.Tensor((128, 128), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.const(1.5, "float32")
|
|
lv1: R.Tensor((128, 128), dtype="float32") = R.where(mask, lv, input)
|
|
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(128, 128, dtype=torch.float32),
|
|
torch.testing.make_tensor((128, 128), dtype=torch.bool, device="cpu"),
|
|
)
|
|
verify_model(Masked_Fill_Inplace(), example_args, {}, Expected)
|
|
|
|
|
|
def test_masked_select():
|
|
class MaskedSelect(Module):
|
|
def forward(self, data: torch.Tensor, mask: torch.Tensor):
|
|
return torch.masked_select(data, mask)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((2, 3), dtype="float32"), mask: R.Tensor((2, 3), dtype="bool")
|
|
) -> R.Tuple(R.Tensor(dtype="float32", ndim=1)):
|
|
R.func_attr({"tir_var_lower_bound": {"u0": 0}, "tir_var_upper_bound": {"u0": 6}})
|
|
u0 = T.int64()
|
|
with R.dataflow():
|
|
lv: R.Tensor((6,), dtype="float32") = R.reshape(data, R.shape([6]))
|
|
lv1: R.Tensor((6,), dtype="bool") = R.reshape(mask, R.shape([6]))
|
|
lv2: R.Tensor(dtype="int64", ndim=2) = R.nonzero(lv1)
|
|
lv3: R.Tensor((1, u0), dtype="int64") = R.match_cast(
|
|
lv2, R.Tensor((1, u0), dtype="int64")
|
|
)
|
|
lv4: R.Tensor((u0,), dtype="int64") = R.squeeze(lv3, axis=[0])
|
|
lv5: R.Tensor((u0,), dtype="float32") = R.take(lv, lv4, axis=0, mode="fast")
|
|
lv6: R.Tensor((), dtype="bool") = R.const(True, "bool")
|
|
lv7: R.Tensor((), dtype="bool") = R.const(True, "bool")
|
|
gv: R.Tuple(R.Tensor((u0,), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(2, 3, dtype=torch.float32),
|
|
torch.tensor([[True, False, True], [False, True, False]]),
|
|
)
|
|
verify_model(MaskedSelect(), example_args, {}, Expected)
|
|
|
|
|
|
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled")
|
|
def test_masked_select_numerically():
|
|
class MaskedSelect(Module):
|
|
def forward(self, data: torch.Tensor, mask: torch.Tensor):
|
|
return torch.masked_select(data, mask)
|
|
|
|
example_args = (
|
|
torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32),
|
|
torch.tensor([[True, False, True], [False, True, False]]),
|
|
)
|
|
verify_model_numerically(MaskedSelect(), example_args)
|
|
|
|
|
|
def test_new_ones():
|
|
class NewOnes(Module):
|
|
def forward(self, x):
|
|
return x.new_ones(1, 2, 3)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 2, 3), dtype="float32") = R.full(
|
|
(1, 2, 3), R.const(1, "float32"), dtype="float32"
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 2, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, dtype=torch.float32),)
|
|
verify_model(NewOnes(), example_args, {}, expected1)
|
|
|
|
|
|
def test_new_zeros():
|
|
class NewZeros(torch.nn.Module):
|
|
def forward(self, x):
|
|
return x.new_zeros(1, 128, 128)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 128, 128), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 128, 128), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 128, 128), dtype="float32") = R.full(
|
|
R.shape([1, 128, 128]), R.const(0, "float32"), dtype="float32"
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 128, 128), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 128, 128, dtype=torch.float32),)
|
|
verify_model(NewZeros(), example_args, {}, expected1)
|
|
|
|
|
|
def test_copy():
|
|
class CopyBroadcast(Module):
|
|
def forward(self, x, src):
|
|
x.copy_(src)
|
|
return x
|
|
|
|
@tvm.script.ir_module
|
|
class expected_copy:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype="float32"), src: R.Tensor((), dtype="int64")) -> R.Tuple(
|
|
R.Tensor((2, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.astype(src, dtype="float32")
|
|
lv1: R.Tensor((2, 3), dtype="float32") = R.broadcast_to(lv, (2, 3))
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.zeros(2, 3, dtype=torch.float32), torch.tensor(1, dtype=torch.int64))
|
|
verify_model(CopyBroadcast(), example_args, {}, expected_copy)
|
|
|
|
|
|
def test_to_copy():
|
|
# float
|
|
class ToFloat(Module):
|
|
def forward(self, x):
|
|
return x.float()
|
|
|
|
@tvm.script.ir_module
|
|
class expected_float:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3, 4), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float32")) = (x,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# half
|
|
class ToHalf(Module):
|
|
def forward(self, x):
|
|
return x.half()
|
|
|
|
@tvm.script.ir_module
|
|
class expected_half:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3, 4), dtype="float16")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 2, 3, 4), dtype="float16") = R.astype(x, dtype="float16")
|
|
gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float16")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# type
|
|
class Type(Module):
|
|
def forward(self, x):
|
|
return x.type(torch.float32)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_type:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3, 4), dtype="float32")
|
|
):
|
|
# block 0
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float32")) = (x,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class To1(Module):
|
|
def forward(self, input):
|
|
return input.to(torch.float16)
|
|
|
|
@I.ir_module
|
|
class expected_to1:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3, 4), dtype="float16")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 2, 3, 4), dtype="float16") = R.astype(input, dtype="float16")
|
|
gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float16")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class To2(Module):
|
|
def forward(self, input):
|
|
return input.to("cpu")
|
|
|
|
@I.ir_module
|
|
class expected_to2:
|
|
@R.function
|
|
def main(input: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 2, 3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float32")) = (input,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
|
|
verify_model(ToFloat(), example_args, {}, expected_float)
|
|
verify_model(ToHalf(), example_args, {}, expected_half)
|
|
verify_model(Type(), example_args, {}, expected_type)
|
|
verify_model(To1(), example_args, {}, expected_to1)
|
|
verify_model(To2(), example_args, {}, expected_to2)
|
|
|
|
|
|
def test_keep_params():
|
|
class Conv2D1(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)
|
|
|
|
def forward(self, input):
|
|
return self.conv(input)
|
|
|
|
@tvm.script.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
|
|
conv_weight: R.Tensor((6, 3, 7, 7), dtype="float32"),
|
|
conv_bias: R.Tensor((6,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 6, 4, 4), dtype="float32")):
|
|
R.func_attr({"num_input": 1})
|
|
# block 0
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
|
|
input_1,
|
|
conv_weight,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(conv_bias, [1, 6, 1, 1])
|
|
lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
|
|
gv: R.Tuple(R.Tensor((1, 6, 4, 4), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
from tvm.relax.frontend import detach_params
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
model = Conv2D1()
|
|
|
|
exported_program = torch.export.export(model, example_args)
|
|
mod = from_exported_program(exported_program, keep_params_as_input=True)
|
|
mod, params = detach_params(mod)
|
|
tvm.ir.assert_structural_equal(mod, expected1)
|
|
func = mod["main"]
|
|
params = params["main"]
|
|
|
|
assert len(params) == len(func.params) - 1
|
|
for param_var, param_tensor in zip(func.params[1:], params):
|
|
assert tuple(x.value for x in param_var.ty.shape.values) == param_tensor.shape
|
|
assert param_var.ty.dtype == param_tensor.dtype
|
|
|
|
tvm.testing.assert_allclose(params[0].numpy(), model.conv.weight.detach().detach().numpy())
|
|
tvm.testing.assert_allclose(params[1].numpy(), model.conv.bias.detach().detach().numpy())
|
|
|
|
|
|
def test_unwrap_unit_return_tuple():
|
|
class Identity(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x):
|
|
return (x,)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
|
|
(256, 256), dtype="float32"
|
|
):
|
|
with R.dataflow():
|
|
gv: R.Tensor((256, 256), dtype="float32") = inp_0
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(256, 256, dtype=torch.float32),)
|
|
verify_model(Identity(), example_args, {}, Expected, unwrap_unit_return_tuple=True)
|
|
|
|
|
|
def test_no_bind_return_tuple():
|
|
class Identity(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x, y):
|
|
return (x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((256, 256), dtype="float32"),
|
|
inp_1: R.Tensor((256, 256), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((256, 256), dtype="float32"), R.Tensor((256, 256), dtype="float32")):
|
|
with R.dataflow():
|
|
gv: R.Tensor((256, 256), dtype="float32") = inp_0
|
|
gv1: R.Tensor((256, 256), dtype="float32") = inp_1
|
|
R.output(gv, gv1)
|
|
return (gv, gv1)
|
|
|
|
example_args = (
|
|
torch.randn(256, 256, dtype=torch.float32),
|
|
torch.randn(256, 256, dtype=torch.float32),
|
|
)
|
|
verify_model(Identity(), example_args, {}, Expected, no_bind_return_tuple=True)
|
|
|
|
|
|
def test_register_buffer():
|
|
class ModelWithBuffer(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.register_buffer("my_buffer", torch.randn(3, 4), persistent=False)
|
|
|
|
def forward(self, x):
|
|
return x + self.my_buffer
|
|
|
|
example_args = (torch.randn(2, 3, 4),)
|
|
ep = export(ModelWithBuffer(), args=example_args)
|
|
# Just verify that import works.
|
|
from_exported_program(ep)
|
|
|
|
|
|
def test_custom_op():
|
|
class AddOp(Module):
|
|
def forward(self, x, y):
|
|
return torch.ops.aten.add.Tensor(x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((5,), dtype="float32"),
|
|
y: R.Tensor((5,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((5,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5,), dtype="float32") = R.subtract(x, y)
|
|
gv: R.Tuple(R.Tensor((5,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
from tvm.relax.frontend.torch.exported_program_translator import (
|
|
ExportedProgramImporter,
|
|
)
|
|
|
|
def custom_add_converter(node: torch.fx.Node, self: ExportedProgramImporter) -> relax.Var:
|
|
x = self.env[node.args[0]]
|
|
y = self.env[node.args[1]]
|
|
|
|
return self.block_builder.emit(R.subtract(x, y))
|
|
|
|
example_args = (torch.randn(5, dtype=torch.float32), torch.randn(5, dtype=torch.float32))
|
|
verify_model(
|
|
AddOp(), example_args, {}, Expected, custom_convert_map={"add.Tensor": custom_add_converter}
|
|
)
|
|
|
|
|
|
def test_empty_like():
|
|
class EmptyLike(Module):
|
|
def forward(self, data):
|
|
return torch.empty_like(data)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((5,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((5,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5,), dtype="float32") = R.zeros(R.shape([5]), dtype="float32")
|
|
gv: R.Tuple(R.Tensor((5,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, dtype=torch.float32),)
|
|
|
|
verify_model(EmptyLike(), example_args, {}, Expected)
|
|
|
|
|
|
def test_one_hot():
|
|
class OneHot(Module):
|
|
def forward(self, indices):
|
|
return torch.nn.functional.one_hot(indices, num_classes=10)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
indices: R.Tensor((5,), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((5, 10), dtype="int64")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((5, 1), dtype="int64") = R.expand_dims(indices, axis=[-1])
|
|
lv2: R.Tensor((5, 10), dtype="bool") = R.equal(lv1, lv)
|
|
lv3: R.Tensor((5, 10), dtype="int64") = R.astype(lv2, dtype="int64")
|
|
gv: R.Tuple(R.Tensor((5, 10), dtype="int64")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randint(0, 10, (5,), dtype=torch.int64),)
|
|
|
|
verify_model(OneHot(), example_args, {}, Expected)
|
|
|
|
|
|
def test_ones_like():
|
|
class OnesLike(Module):
|
|
def forward(self, input):
|
|
return torch.ones_like(input)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((128, 128), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input, R.const(1, "int32"))
|
|
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.rand(128, 128, dtype=torch.float32),)
|
|
|
|
verify_model(OnesLike(), example_args, {}, Expected)
|
|
|
|
|
|
def test_zero_inplace():
|
|
class ZeroInplace(Module):
|
|
def forward(self, input):
|
|
return input.zero_()
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((128, 128), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input, R.const(0, "int32"))
|
|
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.rand(128, 128, dtype=torch.float32),)
|
|
|
|
verify_model(ZeroInplace(), example_args, {}, Expected)
|
|
|
|
|
|
def test_zeros():
|
|
class Zeros(Module):
|
|
def forward(self, input):
|
|
return torch.zeros(5, 2)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((5, 2), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 2), dtype="float32") = R.full(
|
|
R.shape([5, 2]), R.const(0.0, "float32"), dtype="float32"
|
|
)
|
|
gv: R.Tuple(R.Tensor((5, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.rand(128, 128, dtype=torch.float32),)
|
|
|
|
verify_model(Zeros(), example_args, {}, Expected)
|
|
|
|
|
|
def test_zeros_like():
|
|
class ZerosLike(Module):
|
|
def forward(self, input):
|
|
return torch.zeros_like(input)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((128, 128), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input, R.const(0, "int32"))
|
|
gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.rand(128, 128, dtype=torch.float32),)
|
|
verify_model(ZerosLike(), example_args, {}, Expected)
|
|
|
|
|
|
def test_randn():
|
|
class Randn(Module):
|
|
def forward(self, input):
|
|
return input + torch.randn(5, 3)
|
|
|
|
example_args = (torch.rand(5, 3, dtype=torch.float32),)
|
|
exported_program = export(Randn(), args=example_args)
|
|
mod = from_exported_program(exported_program)
|
|
func = mod["main"]
|
|
ret_ty = func.ret_ty
|
|
assert ret_ty.fields[0].shape[0] == 5
|
|
assert ret_ty.fields[0].shape[1] == 3
|
|
assert ret_ty.fields[0].dtype == "float32"
|
|
|
|
|
|
def test_randn_like():
|
|
class RandnLike(Module):
|
|
def forward(self, input):
|
|
return input + torch.randn_like(input)
|
|
|
|
example_args = (torch.rand(4, 6, dtype=torch.float32),)
|
|
exported_program = export(RandnLike(), args=example_args)
|
|
mod = from_exported_program(exported_program)
|
|
func = mod["main"]
|
|
ret_ty = func.ret_ty
|
|
assert ret_ty.fields[0].shape[0] == 4
|
|
assert ret_ty.fields[0].shape[1] == 6
|
|
assert ret_ty.fields[0].dtype == "float32"
|
|
|
|
|
|
def test_type_as():
|
|
class TypeAs(Module):
|
|
def forward(self, input, other):
|
|
return input.type_as(other)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((128, 128), dtype="float32"),
|
|
other: R.Tensor((128, 128), dtype="float16"),
|
|
) -> R.Tuple(R.Tensor((128, 128), dtype="float16")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((128, 128), dtype="float16") = R.astype(input, dtype="float16")
|
|
gv: R.Tuple(R.Tensor((128, 128), dtype="float16")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.rand(128, 128, dtype=torch.float32),
|
|
torch.rand(128, 128, dtype=torch.float16),
|
|
)
|
|
|
|
verify_model(TypeAs(), example_args, {}, Expected)
|
|
|
|
|
|
def test_select():
|
|
class Select(Module):
|
|
def forward(self, input):
|
|
return torch.select(input, 0, 1)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((3,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3,), dtype="float32") = R.take(inp_0, R.const(1, "int64"), axis=0)
|
|
gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 3, dtype=torch.float32),)
|
|
|
|
verify_model(Select(), example_args, {}, Expected)
|
|
|
|
|
|
def test_unflatten():
|
|
class Unflatten(Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.unflatten(input, 1, (3, 5))
|
|
|
|
class Unflatten1(Module):
|
|
def forward(self, input):
|
|
return torch.ops.aten.unflatten(input, -2, (3, 5))
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 15, 7), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 3, 5, 7), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 5, 7), dtype="float32") = R.reshape(inp_0, [2, 3, 5, 7])
|
|
gv: R.Tuple(R.Tensor((2, 3, 5, 7), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 15, 7, dtype=torch.float32),)
|
|
|
|
verify_model(Unflatten(), example_args, {}, Expected)
|
|
verify_model(Unflatten1(), example_args, {}, Expected)
|
|
|
|
|
|
def test_gather():
|
|
class Gather0(Module):
|
|
def forward(self, data, indices):
|
|
return torch.gather(data, 0, indices)
|
|
|
|
class Gather1(Module):
|
|
def forward(self, data, indices):
|
|
return torch.gather(data, 1, indices)
|
|
|
|
class Gather2(Module):
|
|
def forward(self, data, indices):
|
|
return torch.gather(data, -1, indices)
|
|
|
|
class Gather3(Module):
|
|
def forward(self, data, indices):
|
|
return torch.gather(data, -2, indices)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected0:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
inp_1: R.Tensor((2, 3), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=0)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
inp_1: R.Tensor((2, 3), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=1)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
inp_1: R.Tensor((2, 3), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-1)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
inp_1: R.Tensor((2, 3), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-2)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(2, 3, dtype=torch.float32),
|
|
torch.randint(0, 3, (2, 3), dtype=torch.int64),
|
|
)
|
|
|
|
verify_model(Gather0(), example_args, {}, Expected0)
|
|
verify_model(Gather1(), example_args, {}, Expected1)
|
|
verify_model(Gather2(), example_args, {}, Expected2)
|
|
verify_model(Gather3(), example_args, {}, Expected3)
|
|
|
|
|
|
def test_index_put():
|
|
# Test case 1: 1D input
|
|
class IndexPut1D(Module):
|
|
def forward(self, data, indices_0, values):
|
|
indices_tuple = (indices_0,)
|
|
return data.index_put_(indices_tuple, values, accumulate=False)
|
|
|
|
example_args_1d = (
|
|
torch.randn(64, dtype=torch.float32),
|
|
torch.randint(0, 64, (128,), dtype=torch.int64),
|
|
torch.randn(128, dtype=torch.float32),
|
|
)
|
|
|
|
@I.ir_module
|
|
class Expected1D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((64,), dtype="float32"),
|
|
indices_0: R.Tensor((128,), dtype="int64"),
|
|
values: R.Tensor((128,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((64,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((64,), dtype="float32") = R.index_put(
|
|
data, (indices_0,), values, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((64,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 2: 2D input
|
|
class IndexPut2D(Module):
|
|
def forward(self, data, indices_0, indices_1, values):
|
|
indices_tuple = (indices_0, indices_1)
|
|
return data.index_put_(indices_tuple, values, accumulate=False)
|
|
|
|
example_args_2d = (
|
|
torch.randn(32, 64, dtype=torch.float32),
|
|
torch.randint(0, 32, (128,), dtype=torch.int64),
|
|
torch.randint(0, 64, (128,), dtype=torch.int64),
|
|
torch.randn(128, dtype=torch.float32),
|
|
)
|
|
|
|
@I.ir_module
|
|
class Expected2D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((32, 64), dtype="float32"),
|
|
indices_0: R.Tensor((128,), dtype="int64"),
|
|
indices_1: R.Tensor((128,), dtype="int64"),
|
|
values: R.Tensor((128,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((32, 64), dtype="float32") = R.index_put(
|
|
data, (indices_0, indices_1), values, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((32, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 3: 3D input
|
|
class IndexPut3D(Module):
|
|
def forward(self, data, indices_0, indices_1, indices_2, values):
|
|
indices_tuple = (indices_0, indices_1, indices_2)
|
|
return data.index_put_(indices_tuple, values, accumulate=False)
|
|
|
|
example_args_3d = (
|
|
torch.randn(16, 32, 64, dtype=torch.float32),
|
|
torch.randint(0, 16, (128,), dtype=torch.int64),
|
|
torch.randint(0, 32, (128,), dtype=torch.int64),
|
|
torch.randint(0, 64, (128,), dtype=torch.int64),
|
|
torch.randn(128, dtype=torch.float32),
|
|
)
|
|
|
|
@I.ir_module
|
|
class Expected3D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((16, 32, 64), dtype="float32"),
|
|
indices_0: R.Tensor((128,), dtype="int64"),
|
|
indices_1: R.Tensor((128,), dtype="int64"),
|
|
indices_2: R.Tensor((128,), dtype="int64"),
|
|
values: R.Tensor((128,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((16, 32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put(
|
|
data, (indices_0, indices_1, indices_2), values, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((16, 32, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 4: 4D input
|
|
class IndexPut4D(Module):
|
|
def forward(self, data, indices_0, indices_1, indices_2, indices_3, values):
|
|
indices_tuple = (indices_0, indices_1, indices_2, indices_3)
|
|
return data.index_put_(indices_tuple, values, accumulate=False)
|
|
|
|
example_args_4d = (
|
|
torch.randn(8, 16, 32, 64, dtype=torch.float32),
|
|
torch.randint(0, 8, (128,), dtype=torch.int64),
|
|
torch.randint(0, 16, (128,), dtype=torch.int64),
|
|
torch.randint(0, 32, (128,), dtype=torch.int64),
|
|
torch.randint(0, 64, (128,), dtype=torch.int64),
|
|
torch.randn(128, dtype=torch.float32),
|
|
)
|
|
|
|
@I.ir_module
|
|
class Expected4D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((8, 16, 32, 64), dtype="float32"),
|
|
indices_0: R.Tensor((128,), dtype="int64"),
|
|
indices_1: R.Tensor((128,), dtype="int64"),
|
|
indices_2: R.Tensor((128,), dtype="int64"),
|
|
indices_3: R.Tensor((128,), dtype="int64"),
|
|
values: R.Tensor((128,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put(
|
|
data,
|
|
(indices_0, indices_1, indices_2, indices_3),
|
|
values,
|
|
accumulate=False,
|
|
)
|
|
gv: R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 5: 5D input
|
|
class IndexPut5D(Module):
|
|
def forward(self, data, indices_0, indices_1, indices_2, indices_3, indices_4, values):
|
|
indices_tuple = (indices_0, indices_1, indices_2, indices_3, indices_4)
|
|
return data.index_put_(indices_tuple, values, accumulate=False)
|
|
|
|
example_args_5d = (
|
|
torch.randn(4, 8, 16, 32, 64, dtype=torch.float32),
|
|
torch.randint(0, 4, (128,), dtype=torch.int64),
|
|
torch.randint(0, 8, (128,), dtype=torch.int64),
|
|
torch.randint(0, 16, (128,), dtype=torch.int64),
|
|
torch.randint(0, 32, (128,), dtype=torch.int64),
|
|
torch.randint(0, 64, (128,), dtype=torch.int64),
|
|
torch.randn(128, dtype=torch.float32),
|
|
)
|
|
|
|
@I.ir_module
|
|
class Expected5D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((4, 8, 16, 32, 64), dtype="float32"),
|
|
indices_0: R.Tensor((128,), dtype="int64"),
|
|
indices_1: R.Tensor((128,), dtype="int64"),
|
|
indices_2: R.Tensor((128,), dtype="int64"),
|
|
indices_3: R.Tensor((128,), dtype="int64"),
|
|
indices_4: R.Tensor((128,), dtype="int64"),
|
|
values: R.Tensor((128,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = R.index_put(
|
|
data,
|
|
(indices_0, indices_1, indices_2, indices_3, indices_4),
|
|
values,
|
|
accumulate=False,
|
|
)
|
|
gv: R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 6: 2D input with multi-dimensional index (broadcasting)
|
|
# This tests the multi-dimensional index support with broadcasting
|
|
class IndexPutBroadcast1D(Module):
|
|
def forward(self, data, indices_1):
|
|
indices_0 = torch.arange(data.shape[0]).unsqueeze(1)
|
|
values = torch.ones(data.shape[0], len(indices_1), dtype=data.dtype)
|
|
return data.index_put_((indices_0, indices_1), values, accumulate=False)
|
|
|
|
example_args_broadcast1 = (
|
|
torch.randn(32, 64, dtype=torch.float32),
|
|
torch.randint(0, 64, (10,), dtype=torch.int64),
|
|
)
|
|
|
|
@I.ir_module
|
|
class ExpectedBroadcast1D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((32, 64), dtype="float32"),
|
|
indices_1: R.Tensor((10,), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((32,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(32), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((32, 1), dtype="int64") = R.expand_dims(lv, axis=[1])
|
|
lv2: R.Tensor((32, 10), dtype="float32") = R.full(
|
|
R.shape([32, 10]), R.const(1.0, "float32"), dtype="float32"
|
|
)
|
|
lv3: R.Tensor((32, 64), dtype="float32") = R.index_put(
|
|
data, (lv1, indices_1), lv2, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((32, 64), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 7: 2D input with multi-dimensional index (second position)
|
|
class IndexPutBroadcast2D(Module):
|
|
def forward(self, data, indices_0):
|
|
indices_1 = torch.arange(data.shape[1]).unsqueeze(1)
|
|
values = torch.ones(len(indices_0), data.shape[1], dtype=data.dtype)
|
|
return data.index_put_((indices_0, indices_1), values, accumulate=False)
|
|
|
|
example_args_broadcast2 = (
|
|
torch.randn(32, 64, dtype=torch.float32),
|
|
torch.randint(0, 32, (10,), dtype=torch.int64),
|
|
)
|
|
|
|
@I.ir_module
|
|
class ExpectedBroadcast2D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((32, 64), dtype="float32"),
|
|
indices_0: R.Tensor((10,), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((64,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(64), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((64, 1), dtype="int64") = R.expand_dims(lv, axis=[1])
|
|
lv2: R.Tensor((10, 64), dtype="float32") = R.full(
|
|
R.shape([10, 64]), R.const(1.0, "float32"), dtype="float32"
|
|
)
|
|
lv3: R.Tensor((32, 64), dtype="float32") = R.index_put(
|
|
data, (indices_0, lv1), lv2, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((32, 64), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 8: 3D input with mixed 1D and 2D indices
|
|
class IndexPutBroadcast3D(Module):
|
|
def forward(self, data, indices_1):
|
|
indices_0 = torch.arange(data.shape[0]).unsqueeze(1)
|
|
indices_2 = torch.arange(data.shape[2]).unsqueeze(1)
|
|
values = torch.ones(data.shape[0], len(indices_1), data.shape[2], dtype=data.dtype)
|
|
return data.index_put_((indices_0, indices_1, indices_2), values, accumulate=False)
|
|
|
|
example_args_broadcast3d = (
|
|
torch.randn(16, 32, 64, dtype=torch.float32),
|
|
torch.randint(0, 32, (10,), dtype=torch.int64),
|
|
)
|
|
|
|
@I.ir_module
|
|
class ExpectedBroadcast3D:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((16, 32, 64), dtype="float32"),
|
|
indices_1: R.Tensor((10,), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((16, 32, 64), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((16,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(16), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((16, 1), dtype="int64") = R.expand_dims(lv, axis=[1])
|
|
lv2: R.Tensor((64,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(64), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv3: R.Tensor((64, 1), dtype="int64") = R.expand_dims(lv2, axis=[1])
|
|
lv4: R.Tensor((16, 10, 64), dtype="float32") = R.full(
|
|
R.shape([16, 10, 64]), R.const(1.0, "float32"), dtype="float32"
|
|
)
|
|
lv5: R.Tensor((16, 32, 64), dtype="float32") = R.index_put(
|
|
data, (lv1, indices_1, lv3), lv4, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((16, 32, 64), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Test case 9: batched indexing with slice (e.g., M[:, rows, cols] = x)
|
|
class IndexPutBatchedWithNone(Module):
|
|
def forward(self, x):
|
|
B = x.size(0)
|
|
M = torch.zeros(B, 11, 11)
|
|
rows = torch.arange(10)
|
|
cols = rows + 1
|
|
M[:, rows, cols] = x # Batched index assignment
|
|
return M
|
|
|
|
example_args_batched_none = (torch.randn(2, 10, dtype=torch.float32),)
|
|
|
|
@I.ir_module
|
|
class ExpectedBatchedWithNone:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((2, 11, 11), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 11, 11), dtype="float32") = R.full(
|
|
R.shape([2, 11, 11]), R.const(0.0, "float32"), dtype="float32"
|
|
)
|
|
lv1: R.Tensor((10,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv2: R.Tensor((10,), dtype="int64") = R.add(lv1, R.const(1, "int64"))
|
|
lv3: R.Tensor((2,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(2), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv4: R.Tensor((2, 1), dtype="int64") = R.reshape(lv3, R.shape([2, 1]))
|
|
lv5: R.Tensor((2, 11, 11), dtype="float32") = R.index_put(
|
|
lv, (lv4, lv1, lv2), x, accumulate=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((2, 11, 11), dtype="float32")) = (lv5,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# Run verification for each case
|
|
verify_model(IndexPut1D(), example_args_1d, {}, Expected1D)
|
|
verify_model(IndexPut2D(), example_args_2d, {}, Expected2D)
|
|
verify_model(IndexPut3D(), example_args_3d, {}, Expected3D)
|
|
verify_model(IndexPut4D(), example_args_4d, {}, Expected4D)
|
|
verify_model(IndexPut5D(), example_args_5d, {}, Expected5D)
|
|
verify_model(IndexPutBroadcast1D(), example_args_broadcast1, {}, ExpectedBroadcast1D)
|
|
verify_model(IndexPutBroadcast2D(), example_args_broadcast2, {}, ExpectedBroadcast2D)
|
|
verify_model(IndexPutBroadcast3D(), example_args_broadcast3d, {}, ExpectedBroadcast3D)
|
|
verify_model(IndexPutBatchedWithNone(), example_args_batched_none, {}, ExpectedBatchedWithNone)
|
|
|
|
|
|
def test_index_put_with_tuple_output():
|
|
class IndexPutTupleOutput(Module):
|
|
def forward(self, x, buf, idx):
|
|
values = x
|
|
buf[..., idx, idx] = values
|
|
return x[..., 1], buf
|
|
|
|
example_args = (
|
|
torch.ones(2, 3, 5, dtype=torch.float32),
|
|
torch.zeros(2, 3, 5, 5, dtype=torch.float32),
|
|
torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64),
|
|
)
|
|
|
|
exported_program = export(IndexPutTupleOutput(), args=example_args)
|
|
mod = from_exported_program(exported_program)
|
|
|
|
ret_ty = mod["main"].ret_ty
|
|
assert isinstance(ret_ty, relax.TupleType)
|
|
|
|
tensor_fields = [f for f in ret_ty.fields if isinstance(f, relax.TensorType)]
|
|
assert len(tensor_fields) >= 2
|
|
|
|
assert any(
|
|
len(f.shape) == 4 and int(f.shape[-2]) == 5 and int(f.shape[-1]) == 5 for f in tensor_fields
|
|
)
|
|
|
|
|
|
def test_m4d_diag_index_put_tuple_output_regression():
|
|
class M4D(Module):
|
|
def forward(self, x):
|
|
b, k, n = 2, 3, 5
|
|
buf = x.new_zeros(b, k, n, n)
|
|
idx = torch.arange(n, device=x.device)
|
|
|
|
diag = buf[..., idx, idx]
|
|
diag = torch.nn.functional.elu(diag) + 1.0 + 1e-8
|
|
buf[..., idx, idx] = diag
|
|
|
|
return x[..., :1], buf
|
|
|
|
ex_in = torch.zeros(2, 3, 5, dtype=torch.float32)
|
|
exported_program = export(M4D().eval(), args=(ex_in,))
|
|
|
|
exported_targets = [str(getattr(n, "target", "")) for n in exported_program.graph.nodes]
|
|
assert any("index_put" in target for target in exported_targets)
|
|
|
|
# Regression focus: importing this graph should not segfault at Tuple construction.
|
|
mod = from_exported_program(exported_program)
|
|
ret_ty = mod["main"].ret_ty
|
|
assert isinstance(ret_ty, relax.TupleType)
|
|
|
|
tensor_fields = [f for f in ret_ty.fields if isinstance(f, relax.TensorType)]
|
|
assert len(tensor_fields) >= 2
|
|
# x: (2, 3, 5) → x[..., :1]: (2, 3, 1)
|
|
assert any(len(f.shape) == 3 and int(f.shape[-1]) == 1 for f in tensor_fields)
|
|
# buf: (2, 3, 5, 5) → 4-D with spatial dims 5x5
|
|
assert any(
|
|
len(f.shape) == 4 and int(f.shape[-2]) == 5 and int(f.shape[-1]) == 5 for f in tensor_fields
|
|
)
|
|
|
|
|
|
def test_index_put_mutation_through_alias_regression():
|
|
class IndexPutAlias(Module):
|
|
def forward(self, x, idx, values):
|
|
y = torch.ops.aten.alias.default(x)
|
|
y[idx] = values
|
|
return x, y
|
|
|
|
example_args = (
|
|
torch.zeros(5, dtype=torch.float32),
|
|
torch.tensor([1, 3], dtype=torch.int64),
|
|
torch.tensor([2.0, 4.0], dtype=torch.float32),
|
|
)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((5,), dtype="float32"),
|
|
idx: R.Tensor((2,), dtype="int64"),
|
|
values: R.Tensor((2,), dtype="float32"),
|
|
) -> R.Tuple(
|
|
R.Tensor((5,), dtype="float32"),
|
|
R.Tensor((5,), dtype="float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5,), dtype="float32") = R.index_put(
|
|
x, (idx,), values, accumulate=False
|
|
)
|
|
# Mutation outputs introduced by functionalization are dropped;
|
|
# only the user outputs (x, y) remain.
|
|
gv: R.Tuple(
|
|
R.Tensor((5,), dtype="float32"),
|
|
R.Tensor((5,), dtype="float32"),
|
|
) = (
|
|
lv,
|
|
lv,
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(IndexPutAlias(), example_args, {}, Expected)
|
|
|
|
|
|
def test_flip():
|
|
class Flip0(Module):
|
|
def forward(self, data):
|
|
return torch.flip(data, [0])
|
|
|
|
class Flip1(Module):
|
|
def forward(self, data):
|
|
return torch.flip(data, [1])
|
|
|
|
@tvm.script.ir_module
|
|
class Expected0:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 2), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 2), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=0)
|
|
gv: R.Tuple(R.Tensor((2, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 2), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 2), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=1)
|
|
gv: R.Tuple(R.Tensor((2, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 2, dtype=torch.float32),)
|
|
|
|
verify_model(Flip0(), example_args, {}, Expected0)
|
|
verify_model(Flip1(), example_args, {}, Expected1)
|
|
|
|
|
|
def test_flip_multi_axis():
|
|
class FlipMulti(Module):
|
|
def forward(self, data):
|
|
return torch.flip(data, [0, 1])
|
|
|
|
class FlipNegMulti(Module):
|
|
def forward(self, data):
|
|
return torch.flip(data, dims=[-1, -2])
|
|
|
|
@tvm.script.ir_module
|
|
class ExpectedMulti:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.flip(inp_0, axis=0)
|
|
lv1: R.Tensor((2, 3), dtype="float32") = R.flip(lv, axis=1)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class ExpectedNegMulti:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="float32") = R.flip(inp_0, axis=-1)
|
|
lv1: R.Tensor((2, 3), dtype="float32") = R.flip(lv, axis=-2)
|
|
gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 3, dtype=torch.float32),)
|
|
|
|
verify_model(FlipMulti(), example_args, {}, ExpectedMulti)
|
|
verify_model(FlipNegMulti(), example_args, {}, ExpectedNegMulti)
|
|
|
|
|
|
def test_take():
|
|
class Take(Module):
|
|
def forward(self, data, indices):
|
|
return torch.take(data, indices)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
data: R.Tensor((5,), dtype="float32"),
|
|
indices: R.Tensor((3,), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((3,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5,), dtype="float32") = R.reshape(data, R.shape([5]))
|
|
lv1: R.Tensor((3,), dtype="float32") = R.take(lv, indices, axis=0, mode="fast")
|
|
gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(5, dtype=torch.float32),
|
|
torch.randint(0, 5, (3,), dtype=torch.int64),
|
|
)
|
|
|
|
verify_model(Take(), example_args, {}, Expected)
|
|
|
|
|
|
def test_any():
|
|
class AnyAten(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.any(x, dim=1)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3), dtype="bool"),
|
|
) -> R.Tuple(R.Tensor((2,), dtype="bool")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3), dtype="int8") = relax.op.astype(x, dtype="int8")
|
|
lv2: R.Tensor((2,), dtype="int8") = relax.op.max(lv, axis=1, keepdims=False)
|
|
lv3: R.Tensor((2,), dtype="bool") = relax.op.astype(lv2, dtype="bool")
|
|
gv: R.Tuple(R.Tensor((2,), dtype="bool")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.tensor([[0, 0, 0], [0, 1, 0]], dtype=torch.bool),)
|
|
verify_model(AnyAten(), example_args, {}, Expected)
|
|
|
|
|
|
def test_std():
|
|
# torch.std(x) defaults to correction=1 (Bessel); decomposes to var.correction + sqrt.
|
|
class Std(Module):
|
|
def forward(self, x):
|
|
return torch.std(x)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.variance(x, axis=None, keepdims=False)
|
|
lv1: R.Tensor((), dtype="float32") = R.multiply(lv, R.const(15.0 / 14.0, "float32"))
|
|
lv2: R.Tensor((), dtype="float32") = R.sqrt(lv1)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv2,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, 3, dtype=torch.float32),)
|
|
verify_model(Std(), example_args, {}, Expected)
|
|
|
|
|
|
def test_var():
|
|
# torch.var(x) defaults to correction=1 (Bessel).
|
|
class Var(Module):
|
|
def forward(self, x):
|
|
return torch.var(x)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.variance(x, axis=None, keepdims=False)
|
|
lv1: R.Tensor((), dtype="float32") = R.multiply(lv, R.const(15.0 / 14.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, 3, dtype=torch.float32),)
|
|
verify_model(Var(), example_args, {}, Expected)
|
|
|
|
|
|
def test_var_correction():
|
|
class VarCorrection2(Module):
|
|
def forward(self, x):
|
|
return torch.var(x, dim=-1, correction=2)
|
|
|
|
class VarCorrection0(Module):
|
|
def forward(self, x):
|
|
return torch.var(x, dim=1, correction=0)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 5), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2,), dtype="float32") = R.variance(x, axis=[-1], keepdims=False)
|
|
lv1: R.Tensor((2,), dtype="float32") = R.multiply(lv, R.const(5.0 / 3.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((2,), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected0:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 5), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2,), dtype="float32") = R.variance(x, axis=[1], keepdims=False)
|
|
gv: R.Tuple(R.Tensor((2,), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 5, dtype=torch.float32),)
|
|
verify_model(VarCorrection2(), example_args, {}, Expected2)
|
|
verify_model(VarCorrection0(), example_args, {}, Expected0)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"torch_dtype,relax_dtype",
|
|
[(torch.float32, "float32"), (torch.bool, "bool")],
|
|
)
|
|
def test_prod(torch_dtype, relax_dtype):
|
|
class Prod(Module):
|
|
def forward(self, x):
|
|
return torch.prod(x, dtype=torch_dtype)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((5, 3), dtype=relax_dtype),
|
|
) -> R.Tuple(R.Tensor((), dtype=relax_dtype)):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype=relax_dtype) = R.prod(x, axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype=relax_dtype)) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.ones(5, 3, dtype=torch_dtype),)
|
|
verify_model(Prod(), example_args, {}, Expected)
|
|
|
|
|
|
def test_cumprod():
|
|
class Cumprod(Module):
|
|
def forward(self, x):
|
|
return torch.cumprod(x, 0)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((5, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 3), dtype="float32") = R.cumprod(inp_0, axis=0, exclusive=False)
|
|
gv: R.Tuple(R.Tensor((5, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_input = torch.randn(5, 3, dtype=torch.float32)
|
|
verify_model(Cumprod(), (example_input,), {}, Expected)
|
|
|
|
|
|
def test_where():
|
|
class Where(Module):
|
|
def forward(self, condition, x, y):
|
|
return torch.where(condition, x, y)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((5, 3), dtype="bool"),
|
|
inp_1: R.Tensor((5, 3), dtype="float32"),
|
|
inp_2: R.Tensor((5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((5, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 3), dtype="float32") = R.where(inp_0, inp_1, inp_2)
|
|
gv: R.Tuple(R.Tensor((5, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
condition = torch.testing.make_tensor((5, 3), dtype=torch.bool, device="cpu")
|
|
x = torch.randn(5, 3, dtype=torch.float32)
|
|
y = torch.randn(5, 3, dtype=torch.float32)
|
|
|
|
verify_model(Where(), (condition, x, y), {}, Expected)
|
|
|
|
|
|
def test_bucketize():
|
|
class Bucketize(Module):
|
|
def forward(self, input_tensor, boundaries):
|
|
return torch.bucketize(input_tensor, boundaries)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
input: R.Tensor((20,), dtype="int64"), boundaries: R.Tensor((10,), dtype="int64")
|
|
) -> R.Tuple(R.Tensor((20,), dtype="int64")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((20,), dtype="int64") = R.bucketize(
|
|
input, boundaries, out_int32=False, right=False
|
|
)
|
|
gv: R.Tuple(R.Tensor((20,), dtype="int64")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
input_tensor = torch.arange(0, 20)
|
|
boundaries = torch.arange(0, 20, 2)
|
|
|
|
verify_model(Bucketize(), (input_tensor, boundaries), {}, Expected)
|
|
|
|
|
|
@pytest.mark.parametrize("right", [False, True])
|
|
@pytest.mark.parametrize("out_int32", [False, True])
|
|
def test_bucketize_numerically(right, out_int32):
|
|
class Bucketize(Module):
|
|
def forward(self, input_tensor, boundaries):
|
|
return torch.bucketize(input_tensor, boundaries, right=right, out_int32=out_int32)
|
|
|
|
input_tensor = torch.tensor([-0.5, 0.0, 0.5, 1.0, 2.0, 2.5], dtype=torch.float32)
|
|
boundaries = torch.tensor([0.0, 1.0, 2.0], dtype=torch.float32)
|
|
|
|
verify_model_numerically(Bucketize(), (input_tensor, boundaries))
|
|
|
|
|
|
def test_argsort():
|
|
class Argsort(Module):
|
|
def forward(self, x):
|
|
return torch.argsort(x, dim=1, descending=True)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((5, 3), dtype="float32")) -> R.Tuple(R.Tensor((5, 3), dtype="int32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 3), dtype="int32") = R.argsort(
|
|
x, axis=1, descending=True, dtype="int32"
|
|
)
|
|
lv1: R.Tensor((5, 3), dtype="float32") = R.gather_elements(x, lv, axis=1)
|
|
lv2: R.Tuple(R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")) = (
|
|
lv1,
|
|
lv,
|
|
)
|
|
lv3: R.Tensor((5, 3), dtype="int32") = lv2[1]
|
|
gv: R.Tuple(R.Tensor((5, 3), dtype="int32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, 3, dtype=torch.float32),)
|
|
verify_model(Argsort(), example_args, {}, Expected)
|
|
|
|
|
|
def test_topk():
|
|
class Topk(Module):
|
|
def forward(self, x):
|
|
return torch.topk(x, k=2, dim=1, largest=True, sorted=True)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((5, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")) = (
|
|
R.topk(x, k=2, axis=1, ret_type="both", largest=True, dtype="int64")
|
|
)
|
|
lv1: R.Tensor((5, 2), dtype="float32") = lv[0]
|
|
lv2: R.Tensor((5, 2), dtype="int64") = lv[1]
|
|
gv: R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")) = (
|
|
lv1,
|
|
lv2,
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, 3, dtype=torch.float32),)
|
|
verify_model(Topk(), example_args, {}, Expected)
|
|
|
|
|
|
def test_dynamic_shape():
|
|
class DynamicModel(torch.nn.Module):
|
|
def forward(self, x1, x2):
|
|
return torch.ops.aten.add.Tensor(x1, x2)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor(("s0", 4), dtype="float32"),
|
|
rhs: R.Tensor(("s0", 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor(("s0", 4), dtype="float32")):
|
|
s0 = T.int64()
|
|
R.func_attr({"tir_var_lower_bound": {"s24": 0}})
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0, 4), dtype="float32") = R.add(lhs, rhs)
|
|
gv: R.Tuple(R.Tensor((s0, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(2, 4), torch.randn(2, 4))
|
|
batch = torch.export.Dim("batch")
|
|
dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}
|
|
|
|
verify_model(
|
|
DynamicModel(),
|
|
example_args,
|
|
{},
|
|
Expected,
|
|
dynamic_shapes=dynamic_shapes,
|
|
run_ep_decomposition=True,
|
|
map_free_vars=True,
|
|
)
|
|
|
|
|
|
def test_broadcast_to():
|
|
class BroadcastTo(Module):
|
|
def forward(self, x):
|
|
return torch.broadcast_to(x, (5, 3))
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((5, 1), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((5, 3), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 3), dtype="float32") = R.broadcast_to(x, R.shape([5, 3]))
|
|
gv: R.Tuple(R.Tensor((5, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, 1, dtype=torch.float32),)
|
|
verify_model(BroadcastTo(), example_args, {}, Expected)
|
|
|
|
|
|
def test_narrow():
|
|
class Narrow(Module):
|
|
def forward(self, x):
|
|
return torch.narrow(x, 1, 0, 2)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((5, 3), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((5, 2), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5, 2), dtype="float32") = R.strided_slice(
|
|
x,
|
|
(R.prim_value(1),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(2),),
|
|
(R.prim_value(1),),
|
|
assume_inbound=False,
|
|
)
|
|
gv: R.Tuple(R.Tensor((5, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
|
|
return gv
|
|
|
|
example_args = (torch.randn(5, 3, dtype=torch.float32),)
|
|
verify_model(Narrow(), example_args, {}, Expected)
|
|
|
|
|
|
def test_item():
|
|
class Item(Module):
|
|
def forward(self, x):
|
|
return x.item()
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((1,), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.take(input, R.const(0, "int64"), axis=0)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, dtype=torch.float32),)
|
|
verify_model(Item(), example_args, {}, Expected)
|
|
|
|
|
|
def test_norm():
|
|
class Norm(Module):
|
|
def __init__(self, p, dim=None, keepdim=False):
|
|
super().__init__()
|
|
self.p = p
|
|
self.dim = dim
|
|
self.keepdim = keepdim
|
|
|
|
def forward(self, x):
|
|
return torch.norm(x, p=self.p, dim=self.dim, keepdim=self.keepdim)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.max(R.abs(inp_0), axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((), dtype="float32") = R.min(R.abs(inp_0), axis=None, keepdims=False)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected3:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
|
|
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(2, "float32"))
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
|
|
lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(0.5, "float32"))
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected4:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
|
|
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(1.0, "float32"))
|
|
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
|
|
lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(1.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected5:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 1, 1, 1), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
|
|
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(-4.0, "float32"))
|
|
lv2: R.Tensor((1, 1, 1, 1), dtype="float32") = R.sum(lv1, axis=None, keepdims=True)
|
|
lv3: R.Tensor((1, 1, 1, 1), dtype="float32") = R.power(
|
|
lv2, R.const(-0.25, "float32")
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 1, 1, 1), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected6:
|
|
@R.function
|
|
def main(
|
|
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 1, 1, 1), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
|
|
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(0.5, "float32"))
|
|
lv2: R.Tensor((1, 1, 1, 1), dtype="float32") = R.sum(lv1, axis=None, keepdims=True)
|
|
lv3: R.Tensor((1, 1, 1, 1), dtype="float32") = R.power(lv2, R.const(2.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((1, 1, 1, 1), dtype="float32")) = (lv3,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
norms = [
|
|
((float("inf"), None, False), Expected1),
|
|
((float("-inf"), None, False), Expected2),
|
|
((float(2), None, False), Expected3),
|
|
((1.0, None, False), Expected4),
|
|
((float(-4), None, True), Expected5),
|
|
((0.5, None, True), Expected6),
|
|
]
|
|
|
|
example_args = (torch.randn(1, 3, 5, 3, dtype=torch.float32),)
|
|
|
|
for (p, dim, keepdim), expected in norms:
|
|
verify_model(Norm(p, dim=dim, keepdim=keepdim), example_args, {}, expected)
|
|
|
|
|
|
def test_eye():
|
|
import pytest
|
|
|
|
class Eye1(Module):
|
|
def forward(self, input):
|
|
return torch.eye(3, 5, dtype=torch.float32)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(input: R.Tensor((3, 5), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 5), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3,), dtype="uint8") = R.arange(
|
|
R.prim_value(0), R.prim_value(3), R.prim_value(1), dtype="uint8"
|
|
)
|
|
lv1: R.Tensor((5,), dtype="uint8") = R.arange(
|
|
R.prim_value(0), R.prim_value(5), R.prim_value(1), dtype="uint8"
|
|
)
|
|
lv2: R.Tensor((3, 1), dtype="uint8") = R.expand_dims(lv, axis=[-1])
|
|
lv3: R.Tensor((3, 5), dtype="bool") = R.equal(lv2, lv1)
|
|
lv4: R.Tensor((3, 5), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
gv: R.Tuple(R.Tensor((3, 5), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
class Eye2(Module):
|
|
def forward(self, input):
|
|
return torch.eye(5, dtype=torch.float32)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(input: R.Tensor((5,), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((5, 5), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((5,), dtype="uint8") = R.arange(
|
|
R.prim_value(0), R.prim_value(5), R.prim_value(1), dtype="uint8"
|
|
)
|
|
lv1: R.Tensor((5,), dtype="uint8") = R.arange(
|
|
R.prim_value(0), R.prim_value(5), R.prim_value(1), dtype="uint8"
|
|
)
|
|
lv2: R.Tensor((5, 1), dtype="uint8") = R.expand_dims(lv, axis=[-1])
|
|
lv3: R.Tensor((5, 5), dtype="bool") = R.equal(lv2, lv1)
|
|
lv4: R.Tensor((5, 5), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
gv: R.Tuple(R.Tensor((5, 5), dtype="float32")) = (lv4,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args1 = (torch.randn(3, 5, dtype=torch.float32),)
|
|
verify_model(Eye1(), example_args1, {}, Expected1)
|
|
|
|
example_args2 = (torch.randn(5, dtype=torch.float32),)
|
|
verify_model(Eye2(), example_args2, {}, Expected2)
|
|
|
|
|
|
def test_cross_entropy():
|
|
class CrossEntropyModule(Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.criterion = nn.CrossEntropyLoss()
|
|
self.target = torch.tensor([0, 1, 2, 1])
|
|
|
|
def forward(self, x):
|
|
return self.criterion(x, self.target)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 3), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 3), dtype="float32") = R.nn.log_softmax(x, axis=1)
|
|
lv1: R.Tensor((4,), dtype="bool") = R.not_equal(
|
|
R.const([0, 1, 2, 1], dtype="int64"), R.const(-100, "int64")
|
|
)
|
|
lv2: R.Tensor((), dtype="int64") = R.const(0, "int64")
|
|
lv3: R.Tensor((4,), dtype="int64") = R.where(
|
|
lv1, R.const([0, 1, 2, 1], dtype="int64"), lv2
|
|
)
|
|
lv4: R.Tensor((4, 1), dtype="int64") = R.expand_dims(lv3, axis=[1])
|
|
lv5: R.Tensor((4, 1), dtype="float32") = R.gather_elements(lv, lv4, axis=1)
|
|
lv6: R.Tensor((4,), dtype="float32") = R.squeeze(lv5, axis=[1])
|
|
lv7: R.Tensor((4,), dtype="float32") = R.negative(lv6)
|
|
lv8: R.Tensor((4,), dtype="bool") = R.not_equal(
|
|
R.const([0, 1, 2, 1], dtype="int64"), R.const(-100, "int64")
|
|
)
|
|
lv9: R.Tensor((), dtype="float32") = R.const(0.0, "float32")
|
|
lv10: R.Tensor((4,), dtype="float32") = R.where(lv8, lv7, lv9)
|
|
lv11: R.Tensor((4,), dtype="bool") = R.not_equal(
|
|
R.const([0, 1, 2, 1], dtype="int64"), R.const(-100, "int64")
|
|
)
|
|
lv12: R.Tensor((4,), dtype="int64") = R.astype(lv11, dtype="int64")
|
|
lv13: R.Tensor((), dtype="int64") = R.sum(lv12, axis=None, keepdims=False)
|
|
lv14: R.Tensor((), dtype="float32") = R.astype(lv13, dtype="float32")
|
|
lv15: R.Tensor((), dtype="float32") = R.sum(lv10, axis=None, keepdims=False)
|
|
lv16: R.Tensor((), dtype="float32") = R.divide(lv15, lv14)
|
|
gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv16,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args1 = (torch.randn(4, 3, dtype=torch.float32),)
|
|
verify_model(CrossEntropyModule(), example_args1, {}, Expected1)
|
|
|
|
|
|
def test_linspace():
|
|
class Linspace(Module):
|
|
def forward(self, input):
|
|
return torch.linspace(0, 1, steps=9, dtype=torch.float32)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(input: R.Tensor((9, 9), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((9,), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((9,), dtype="int64") = R.arange(
|
|
R.prim_value(0), R.prim_value(9), R.prim_value(1), dtype="int64"
|
|
)
|
|
lv1: R.Tensor((9,), dtype="bool") = R.less(lv, R.const(4, "int64"))
|
|
lv2: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32")
|
|
lv3: R.Tensor((9,), dtype="float32") = R.multiply(lv2, R.const(0.125, "float32"))
|
|
lv4: R.Tensor((9,), dtype="float32") = R.add(lv3, R.const(0.0, "float32"))
|
|
lv5: R.Tensor((9,), dtype="int64") = R.subtract(R.const(8, "int64"), lv)
|
|
lv6: R.Tensor((9,), dtype="float32") = R.astype(lv5, dtype="float32")
|
|
lv7: R.Tensor((9,), dtype="float32") = R.multiply(lv6, R.const(0.125, "float32"))
|
|
lv8: R.Tensor((9,), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv7)
|
|
lv9: R.Tensor((9,), dtype="float32") = R.where(lv1, lv4, lv8)
|
|
gv: R.Tuple(R.Tensor((9,), dtype="float32")) = (lv9,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(9, 9, dtype=torch.float32),)
|
|
verify_model(Linspace(), example_args, {}, Expected)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"torch_dtype, relax_dtype",
|
|
[
|
|
(torch.float32, "float32"),
|
|
(torch.float16, "float16"),
|
|
(torch.bfloat16, "bfloat16"),
|
|
(torch.int64, "int64"),
|
|
(torch.int32, "int32"),
|
|
(torch.bool, "bool"),
|
|
],
|
|
)
|
|
def test_dtypes(torch_dtype, relax_dtype):
|
|
example_args = (
|
|
torch.testing.make_tensor((10, 10), dtype=torch_dtype, device="cpu", low=0, high=10),
|
|
torch.testing.make_tensor((10, 10), dtype=torch_dtype, device="cpu", low=0, high=10),
|
|
)
|
|
|
|
class Model(Module):
|
|
def forward(self, lhs: torch.Tensor, rhs: torch.Tensor):
|
|
return torch.ops.aten.add(lhs, rhs)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
lhs: R.Tensor((10, 10), dtype=relax_dtype),
|
|
rhs: R.Tensor((10, 10), dtype=relax_dtype),
|
|
) -> R.Tuple(R.Tensor((10, 10), dtype=relax_dtype)):
|
|
with R.dataflow():
|
|
lv: R.Tensor((10, 10), dtype=relax_dtype) = relax.op.add(lhs, rhs)
|
|
gv: R.Tuple(R.Tensor((10, 10), dtype=relax_dtype)) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(Model(), example_args, {}, Expected)
|
|
|
|
|
|
def test_mm():
|
|
class MatrixMultiply(Module):
|
|
def forward(self, a, b):
|
|
return torch.mm(a, b)
|
|
|
|
example_args = (
|
|
torch.randn(2, 3, dtype=torch.float32),
|
|
torch.randn(3, 4, dtype=torch.float32),
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor((2, 3), dtype="float32"),
|
|
b: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 4), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 4), dtype="float32") = R.matmul(a, b, out_dtype="float32")
|
|
gv: R.Tuple(R.Tensor((2, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(MatrixMultiply(), example_args, {}, Expected)
|
|
|
|
|
|
def test_sparse_mm():
|
|
class SparseMatrixMultiply(Module):
|
|
def forward(self, sparse_input, dense_input):
|
|
return torch.sparse.mm(sparse_input, dense_input)
|
|
|
|
indices = torch.tensor([[0, 1, 2], [2, 0, 1]])
|
|
values = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
|
|
sparse_input = torch.sparse_coo_tensor(indices, values, size=(3, 100))
|
|
dense_input = torch.randn(100, 50, dtype=torch.float32)
|
|
|
|
example_args = (sparse_input, dense_input)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
sparse_input: R.Tensor((3, 100), dtype="float32"),
|
|
dense_input: R.Tensor((100, 50), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((3, 50), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 50), dtype="float32") = R.full(
|
|
R.shape([3, 50]), R.const(0.0, "float32"), dtype="float32"
|
|
)
|
|
lv1: R.Tensor((3, 50), dtype="float32") = R.matmul(
|
|
sparse_input, dense_input, out_dtype="float32"
|
|
)
|
|
gv: R.Tuple(R.Tensor((3, 50), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(SparseMatrixMultiply(), example_args, {}, Expected)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_lstm():
|
|
class LSTM(nn.Module):
|
|
def __init__(self, input_size, hidden_size, batch_first, bidirectional):
|
|
super().__init__()
|
|
self.lstm = nn.LSTM(
|
|
input_size=input_size,
|
|
hidden_size=hidden_size,
|
|
num_layers=1,
|
|
batch_first=batch_first,
|
|
bidirectional=bidirectional,
|
|
)
|
|
|
|
def forward(self, x):
|
|
y, _ = self.lstm(x)
|
|
return y
|
|
|
|
# Unidirectional LSTM with batch_first=True
|
|
torch.manual_seed(42)
|
|
x = torch.randn(2, 3, 4, dtype=torch.float32)
|
|
verify_model_numerically(LSTM(4, 8, batch_first=True, bidirectional=False), (x,))
|
|
|
|
# Unidirectional LSTM with batch_first=False
|
|
torch.manual_seed(43)
|
|
x2 = torch.randn(4, 2, 3, dtype=torch.float32)
|
|
verify_model_numerically(LSTM(3, 6, batch_first=False, bidirectional=False), (x2,))
|
|
|
|
# Bidirectional LSTM with batch_first=True
|
|
torch.manual_seed(44)
|
|
x3 = torch.randn(2, 3, 4, dtype=torch.float32)
|
|
verify_model_numerically(LSTM(4, 8, batch_first=True, bidirectional=True), (x3,))
|
|
|
|
# Bidirectional LSTM with batch_first=False
|
|
torch.manual_seed(45)
|
|
x4 = torch.randn(4, 2, 3, dtype=torch.float32)
|
|
verify_model_numerically(LSTM(3, 6, batch_first=False, bidirectional=True), (x4,))
|
|
|
|
|
|
def test_tensor_none_tuple():
|
|
example_args = (torch.tensor([1.0, 2.0, 3.0]),)
|
|
|
|
class TensorNoneModel(Module):
|
|
def forward(self, x):
|
|
return x + 1, None
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((3,), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3,), dtype="float32"), R.Any
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3,), dtype="float32") = R.add(x, R.const(1.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((3,), dtype="float32"), R.Any) = (lv, R.null_value())
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(TensorNoneModel(), example_args, {}, Expected)
|
|
|
|
|
|
def test_gru():
|
|
class BasicGRU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gru = nn.GRU(
|
|
input_size=4,
|
|
hidden_size=8,
|
|
num_layers=1,
|
|
batch_first=True,
|
|
bidirectional=False,
|
|
)
|
|
|
|
def forward(self, x):
|
|
y, _ = self.gru(x)
|
|
return y
|
|
|
|
torch.manual_seed(42)
|
|
x = torch.randn(2, 3, 4, dtype=torch.float32)
|
|
model = BasicGRU()
|
|
with torch.no_grad():
|
|
pytorch_output = model(x)
|
|
exported_program = export(model, args=(x,))
|
|
mod = from_exported_program(exported_program)
|
|
target = tvm.target.Target("llvm")
|
|
ex = relax.build(mod, target)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
x_tvm = tvm.runtime.tensor(x.numpy())
|
|
tvm_output = vm["main"](x_tvm)
|
|
if hasattr(tvm_output, "numpy"):
|
|
tvm_output_np = tvm_output.numpy()
|
|
else:
|
|
tvm_output_np = tvm_output[0].numpy()
|
|
assert pytorch_output.shape == tvm_output_np.shape, (
|
|
f"Shape mismatch: PyTorch {pytorch_output.shape} vs TVM {tvm_output_np.shape}"
|
|
)
|
|
tvm.testing.assert_allclose(pytorch_output.numpy(), tvm_output_np, rtol=1e-4, atol=1e-5)
|
|
|
|
class SeqFirstGRU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gru = nn.GRU(
|
|
input_size=3,
|
|
hidden_size=6,
|
|
num_layers=1,
|
|
batch_first=False,
|
|
bidirectional=False,
|
|
)
|
|
|
|
def forward(self, x):
|
|
y, _ = self.gru(x)
|
|
return y
|
|
|
|
torch.manual_seed(43)
|
|
x2 = torch.randn(4, 2, 3, dtype=torch.float32)
|
|
model2 = SeqFirstGRU()
|
|
with torch.no_grad():
|
|
pytorch_output2 = model2(x2)
|
|
exported_program2 = export(model2, args=(x2,))
|
|
mod2 = from_exported_program(exported_program2)
|
|
ex2 = relax.build(mod2, target)
|
|
vm2 = relax.VirtualMachine(ex2, tvm.cpu())
|
|
x2_tvm = tvm.runtime.tensor(x2.numpy())
|
|
tvm_output2 = vm2["main"](x2_tvm)
|
|
if hasattr(tvm_output2, "numpy"):
|
|
tvm_output2_np = tvm_output2.numpy()
|
|
else:
|
|
tvm_output2_np = tvm_output2[0].numpy()
|
|
assert pytorch_output2.shape == tvm_output2_np.shape
|
|
tvm.testing.assert_allclose(pytorch_output2.numpy(), tvm_output2_np, rtol=1e-4, atol=1e-5)
|
|
|
|
# Test bidirectional GRU with batch_first=True
|
|
class BidirectionalGRU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gru = nn.GRU(
|
|
input_size=4,
|
|
hidden_size=5,
|
|
num_layers=1,
|
|
batch_first=True,
|
|
bidirectional=True,
|
|
)
|
|
|
|
def forward(self, x):
|
|
y, _ = self.gru(x)
|
|
return y
|
|
|
|
torch.manual_seed(44)
|
|
x3 = torch.randn(2, 3, 4, dtype=torch.float32)
|
|
model3 = BidirectionalGRU()
|
|
with torch.no_grad():
|
|
pytorch_output3 = model3(x3)
|
|
|
|
# Verify output shape is correct (hidden_size * 2 due to bidirectional)
|
|
assert pytorch_output3.shape == (
|
|
2,
|
|
3,
|
|
10,
|
|
), f"Expected shape (2, 3, 10), got {pytorch_output3.shape}"
|
|
|
|
exported_program3 = export(model3, args=(x3,))
|
|
mod3 = from_exported_program(exported_program3)
|
|
ex3 = relax.build(mod3, target)
|
|
vm3 = relax.VirtualMachine(ex3, tvm.cpu())
|
|
x3_tvm = tvm.runtime.tensor(x3.numpy())
|
|
tvm_output3 = vm3["main"](x3_tvm)
|
|
if hasattr(tvm_output3, "numpy"):
|
|
tvm_output3_np = tvm_output3.numpy()
|
|
else:
|
|
tvm_output3_np = tvm_output3[0].numpy()
|
|
assert pytorch_output3.shape == tvm_output3_np.shape, (
|
|
f"Shape mismatch: PyTorch {pytorch_output3.shape} vs TVM {tvm_output3_np.shape}"
|
|
)
|
|
tvm.testing.assert_allclose(pytorch_output3.numpy(), tvm_output3_np, rtol=1e-4, atol=1e-5)
|
|
|
|
# Test bidirectional GRU with batch_first=False
|
|
class SeqFirstBidirectionalGRU(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.gru = nn.GRU(
|
|
input_size=3,
|
|
hidden_size=4,
|
|
num_layers=1,
|
|
batch_first=False,
|
|
bidirectional=True,
|
|
)
|
|
|
|
def forward(self, x):
|
|
y, _ = self.gru(x)
|
|
return y
|
|
|
|
torch.manual_seed(45)
|
|
x4 = torch.randn(4, 2, 3, dtype=torch.float32) # (seq_len, batch, input_size)
|
|
model4 = SeqFirstBidirectionalGRU()
|
|
with torch.no_grad():
|
|
pytorch_output4 = model4(x4)
|
|
|
|
# Verify output shape (seq_len, batch, hidden_size * 2)
|
|
assert pytorch_output4.shape == (
|
|
4,
|
|
2,
|
|
8,
|
|
), f"Expected shape (4, 2, 8), got {pytorch_output4.shape}"
|
|
|
|
exported_program4 = export(model4, args=(x4,))
|
|
mod4 = from_exported_program(exported_program4)
|
|
ex4 = relax.build(mod4, target)
|
|
vm4 = relax.VirtualMachine(ex4, tvm.cpu())
|
|
x4_tvm = tvm.runtime.tensor(x4.numpy())
|
|
tvm_output4 = vm4["main"](x4_tvm)
|
|
if hasattr(tvm_output4, "numpy"):
|
|
tvm_output4_np = tvm_output4.numpy()
|
|
else:
|
|
tvm_output4_np = tvm_output4[0].numpy()
|
|
assert pytorch_output4.shape == tvm_output4_np.shape
|
|
tvm.testing.assert_allclose(pytorch_output4.numpy(), tvm_output4_np, rtol=1e-4, atol=1e-5)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_rnn_tanh():
|
|
target = tvm.target.Target("llvm")
|
|
|
|
def _check(rnn_kwargs, x_shape, seed):
|
|
class RNNWithState(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.rnn = nn.RNN(nonlinearity="tanh", num_layers=1, **rnn_kwargs)
|
|
|
|
def forward(self, x):
|
|
output, h_n = self.rnn(x)
|
|
return output, h_n
|
|
|
|
torch.manual_seed(seed)
|
|
x = torch.randn(*x_shape, dtype=torch.float32)
|
|
model = RNNWithState()
|
|
with torch.no_grad():
|
|
pt_out, pt_hn = model(x)
|
|
|
|
exported_program = export(model, args=(x,))
|
|
mod = from_exported_program(exported_program, run_ep_decomposition=False)
|
|
ex = relax.build(mod, target)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
tvm_outputs = vm["main"](tvm.runtime.tensor(x.numpy()))
|
|
tvm_out_np = tvm_outputs[0].numpy()
|
|
tvm_hn_np = tvm_outputs[1].numpy()
|
|
|
|
assert pt_out.shape == tvm_out_np.shape, (
|
|
f"output shape mismatch: PyTorch {tuple(pt_out.shape)} vs TVM {tvm_out_np.shape}"
|
|
)
|
|
assert pt_hn.shape == tvm_hn_np.shape, (
|
|
f"h_n shape mismatch: PyTorch {tuple(pt_hn.shape)} vs TVM {tvm_hn_np.shape}"
|
|
)
|
|
tvm.testing.assert_allclose(pt_out.numpy(), tvm_out_np, rtol=1e-4, atol=1e-5)
|
|
tvm.testing.assert_allclose(pt_hn.numpy(), tvm_hn_np, rtol=1e-4, atol=1e-5)
|
|
|
|
# batch_first, unidirectional
|
|
_check(
|
|
{"input_size": 4, "hidden_size": 8, "batch_first": True, "bidirectional": False},
|
|
(2, 3, 4),
|
|
seed=42,
|
|
)
|
|
# seq-first (batch_first=False), unidirectional
|
|
_check(
|
|
{"input_size": 3, "hidden_size": 6, "batch_first": False, "bidirectional": False},
|
|
(4, 2, 3),
|
|
seed=43,
|
|
)
|
|
# bidirectional, batch_first
|
|
_check(
|
|
{"input_size": 4, "hidden_size": 8, "batch_first": True, "bidirectional": True},
|
|
(2, 3, 4),
|
|
seed=44,
|
|
)
|
|
|
|
|
|
def test_dynamic_shape_with_range_constraints():
|
|
class DynamicModel(torch.nn.Module):
|
|
def forward(self, x1, x2):
|
|
return torch.ops.aten.add.Tensor(x1, x2)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x1: R.Tensor(("s0", 4), dtype="float32"), x2: R.Tensor(("s0", 4), dtype="float32")
|
|
) -> R.Tuple(R.Tensor(("s0", 4), dtype="float32")):
|
|
s0 = T.int64()
|
|
R.func_attr({"tir_var_lower_bound": {"s24": 1}, "tir_var_upper_bound": {"s24": 64}})
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0, 4), dtype="float32") = R.add(x1, x2)
|
|
gv: R.Tuple(R.Tensor((s0, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(8, 4), torch.randn(8, 4))
|
|
batch = torch.export.Dim("batch", min=1, max=64)
|
|
dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}}
|
|
|
|
verify_model(
|
|
DynamicModel(),
|
|
example_args,
|
|
{},
|
|
Expected,
|
|
dynamic_shapes=dynamic_shapes,
|
|
map_free_vars=True,
|
|
)
|
|
|
|
|
|
def test_dynamic_shape_with_addition_constraints():
|
|
class ConcatModel(torch.nn.Module):
|
|
def forward(self, x, y):
|
|
return torch.cat([x, y], dim=0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("s0", 4), dtype="float32"), y: R.Tensor(("s0___1", 4), dtype="float32")
|
|
) -> R.Tuple(R.Tensor(("s0 + s0___1", 4), dtype="float32")):
|
|
s0 = T.int64()
|
|
s0___1 = T.int64()
|
|
R.func_attr(
|
|
{
|
|
"tir_var_lower_bound": {"s77": 1, "s77___1": 2},
|
|
"tir_var_upper_bound": {"s77": 64, "s77___1": 65},
|
|
}
|
|
)
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0 + s0___1, 4), dtype="float32") = R.concat((x, y), axis=0)
|
|
gv: R.Tuple(R.Tensor((s0 + s0___1, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
batch = torch.export.Dim("batch", min=1, max=64)
|
|
example_args = (torch.randn(8, 4), torch.randn(9, 4))
|
|
dynamic_shapes = {"x": {0: batch}, "y": {0: batch + 1}}
|
|
|
|
verify_model(
|
|
ConcatModel(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes, map_free_vars=True
|
|
)
|
|
|
|
|
|
def test_dynamic_shape_with_subtraction_constraints():
|
|
class ConcatModel(torch.nn.Module):
|
|
def forward(self, x, y):
|
|
return torch.cat([x, y], dim=0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("s0___1", 4), dtype="float32"), y: R.Tensor(("s0", 4), dtype="float32")
|
|
) -> R.Tuple(R.Tensor(("s0___1 + s0", 4), dtype="float32")):
|
|
s0___1 = T.int64()
|
|
s0 = T.int64()
|
|
R.func_attr(
|
|
{
|
|
"tir_var_lower_bound": {"s17": 0, "s17___1": 1},
|
|
"tir_var_upper_bound": {"s17": 63, "s17___1": 64},
|
|
}
|
|
)
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0___1 + s0, 4), dtype="float32") = R.concat((x, y), axis=0)
|
|
gv: R.Tuple(R.Tensor((s0___1 + s0, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
batch = torch.export.Dim("batch", min=1, max=64)
|
|
example_args = (torch.randn(8, 4), torch.randn(7, 4))
|
|
dynamic_shapes = {"x": {0: batch}, "y": {0: batch - 1}}
|
|
|
|
verify_model(
|
|
ConcatModel(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes, map_free_vars=True
|
|
)
|
|
|
|
|
|
def test_dynamic_shape_with_multiplication_constraints():
|
|
class ConcatModel(torch.nn.Module):
|
|
def forward(self, x, y):
|
|
return torch.cat([x, y], dim=0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("s0", 4), dtype="float32"), y: R.Tensor(("s0_2", 4), dtype="float32")
|
|
) -> R.Tuple(R.Tensor(("s0 + s0_2", 4), dtype="float32")):
|
|
s0 = T.int64()
|
|
s0_2 = T.int64()
|
|
R.func_attr(
|
|
{
|
|
"tir_var_lower_bound": {"s77": 1, "s77_2": 2},
|
|
"tir_var_upper_bound": {"s77": 64, "s77_2": 128},
|
|
}
|
|
)
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0 + s0_2, 4), dtype="float32") = R.concat((x, y), axis=0)
|
|
gv: R.Tuple(R.Tensor((s0 + s0_2, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
batch = torch.export.Dim("batch", min=1, max=64)
|
|
example_args = (torch.randn(8, 4), torch.randn(16, 4))
|
|
dynamic_shapes = {"x": {0: batch}, "y": {0: batch * 2}}
|
|
|
|
verify_model(
|
|
ConcatModel(), example_args, {}, Expected, dynamic_shapes=dynamic_shapes, map_free_vars=True
|
|
)
|
|
|
|
|
|
def test_dynamic_shape_with_unbounded_constraints():
|
|
class DynamicModel(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.add.Tensor(x, x)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("s0", 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor(("s0", 4), dtype="float32")
|
|
):
|
|
s0 = T.int64()
|
|
R.func_attr({"tir_var_lower_bound": {"s77": 2}})
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0, 4), dtype="float32") = R.add(x, x)
|
|
gv: R.Tuple(R.Tensor((s0, 4), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(8, 4),)
|
|
batch = torch.export.Dim("batch", min=2)
|
|
dynamic_shapes = {"x": {0: batch}}
|
|
|
|
verify_model(
|
|
DynamicModel(),
|
|
example_args,
|
|
{},
|
|
Expected,
|
|
dynamic_shapes=dynamic_shapes,
|
|
map_free_vars=True,
|
|
)
|
|
|
|
|
|
def test_sym_size_int():
|
|
class SymSizeInt(Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x):
|
|
# TODO(@mshr-h): `torch.ops.aten.sym_size.int(x, self.dim)` would be ideal, but currently
|
|
# the ep frontend is not able to handle it.
|
|
return torch.add(x[0], torch.ops.aten.sym_size.int(x, self.dim))
|
|
|
|
@I.ir_module
|
|
class Expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((1, 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 4), dtype="float32") = R.take(
|
|
x, R.const(0, "int64"), axis=0, mode="fast"
|
|
)
|
|
lv1: R.Tensor((3, 4), dtype="float32") = R.add(lv, R.const(3.0, "float32"))
|
|
gv: R.Tuple(R.Tensor((3, 4), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args_1 = (torch.randn(1, 3, 4),)
|
|
verify_model(SymSizeInt(dim=1), example_args_1, {}, Expected1)
|
|
verify_model(SymSizeInt(dim=-2), example_args_1, {}, Expected1)
|
|
|
|
class SymSizeIntDynamic(Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.dim = dim
|
|
|
|
def forward(self, x):
|
|
shape_dim = torch.ops.aten.sym_size.int(x, self.dim)
|
|
return x.reshape(shape_dim, -1)
|
|
|
|
@I.ir_module
|
|
class Expected2:
|
|
@R.function
|
|
def main(x: R.Tensor(("s0", 3, 4), dtype="float32")) -> R.Tuple(
|
|
R.Tensor(("s0", 12), dtype="float32")
|
|
):
|
|
s0 = T.int64()
|
|
R.func_attr({"tir_var_lower_bound": {"s77": 0}})
|
|
with R.dataflow():
|
|
lv: R.Tensor((s0, 12), dtype="float32") = R.reshape(x, R.shape([s0, 12]))
|
|
gv: R.Tuple(R.Tensor((s0, 12), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args_2 = (torch.randn(2, 3, 4),)
|
|
dynamic_shapes = {"x": {0: torch.export.Dim("dim")}}
|
|
verify_model(
|
|
SymSizeIntDynamic(dim=0),
|
|
example_args_2,
|
|
{},
|
|
Expected2,
|
|
dynamic_shapes=dynamic_shapes,
|
|
map_free_vars=True,
|
|
)
|
|
|
|
|
|
def test_exponential():
|
|
class Exponential(Module):
|
|
def forward(self, x):
|
|
return x.exponential_()
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 8), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 8), dtype="float32") = R.zeros_like(x)
|
|
gv: R.Tuple(R.Tensor((4, 8), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(4, 8, dtype=torch.float32),)
|
|
verify_model(Exponential(), example_args, {}, Expected)
|
|
|
|
|
|
def test_max_dim():
|
|
class MaxDim1(Module):
|
|
def forward(self, x):
|
|
return torch.max(x, dim=1)
|
|
|
|
class MaxDim2(Module):
|
|
def forward(self, x):
|
|
return torch.max(x, dim=1, keepdim=True)
|
|
|
|
@I.ir_module
|
|
class expected1:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 8, 16), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 16), dtype="float32"), R.Tensor((4, 16), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tuple(
|
|
R.Tensor((4, 1, 16), dtype="float32"), R.Tensor((4, 1, 16), dtype="int64")
|
|
) = R.topk(x, k=1, axis=1, ret_type="both", largest=True, dtype="int64")
|
|
lv1: R.Tensor((4, 1, 16), dtype="float32") = lv[0]
|
|
lv2: R.Tensor((4, 16), dtype="float32") = R.squeeze(lv1, axis=[1])
|
|
lv3: R.Tensor((4, 1, 16), dtype="int64") = lv[1]
|
|
lv4: R.Tensor((4, 16), dtype="int64") = R.squeeze(lv3, axis=[1])
|
|
lv5: R.Tuple(
|
|
R.Tensor((4, 16), dtype="float32"), R.Tensor((4, 16), dtype="int64")
|
|
) = (lv2, lv4)
|
|
lv6: R.Tensor((4, 16), dtype="float32") = lv5[0]
|
|
lv7: R.Tensor((4, 16), dtype="int64") = lv5[1]
|
|
gv: R.Tuple(
|
|
R.Tensor((4, 16), dtype="float32"), R.Tensor((4, 16), dtype="int64")
|
|
) = (lv6, lv7)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class expected2:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 8, 16), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 1, 16), dtype="float32"), R.Tensor((4, 1, 16), dtype="int64")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tuple(
|
|
R.Tensor((4, 1, 16), dtype="float32"), R.Tensor((4, 1, 16), dtype="int64")
|
|
) = R.topk(x, k=1, axis=1, ret_type="both", largest=True, dtype="int64")
|
|
lv1: R.Tensor((4, 1, 16), dtype="float32") = lv[0]
|
|
lv2: R.Tensor((4, 1, 16), dtype="int64") = lv[1]
|
|
lv3: R.Tuple(
|
|
R.Tensor((4, 1, 16), dtype="float32"), R.Tensor((4, 1, 16), dtype="int64")
|
|
) = (lv1, lv2)
|
|
lv4: R.Tensor((4, 1, 16), dtype="float32") = lv3[0]
|
|
lv5: R.Tensor((4, 1, 16), dtype="int64") = lv3[1]
|
|
gv: R.Tuple(
|
|
R.Tensor((4, 1, 16), dtype="float32"), R.Tensor((4, 1, 16), dtype="int64")
|
|
) = (lv4, lv5)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(4, 8, 16, dtype=torch.float32),)
|
|
verify_model(MaxDim1(), example_args, {}, expected1)
|
|
verify_model(MaxDim2(), example_args, {}, expected2)
|
|
|
|
|
|
def test_alias():
|
|
class Alias(Module):
|
|
def forward(self, x):
|
|
return torch.ops.aten.alias(x)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((4, 8), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((4, 8), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
gv: R.Tuple(R.Tensor((4, 8), dtype="float32")) = (x,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(4, 8, dtype=torch.float32),)
|
|
verify_model(Alias(), example_args, {}, Expected)
|
|
|
|
|
|
def test_scatter_value():
|
|
class ScatterValue(Module):
|
|
def forward(self, x, index):
|
|
return x.scatter(1, index, 0.5)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((4, 8), dtype="float32"),
|
|
index: R.Tensor((4, 2), dtype="int64"),
|
|
) -> R.Tuple(R.Tensor((4, 8), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 2), dtype="float32") = R.broadcast_to(
|
|
R.const(0.5, "float32"), R.shape([4, 2])
|
|
)
|
|
lv1: R.Tensor((4, 8), dtype="float32") = R.scatter_elements(x, index, lv, axis=1)
|
|
gv: R.Tuple(R.Tensor((4, 8), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(4, 8, dtype=torch.float32),
|
|
torch.randint(0, 8, (4, 2), dtype=torch.int64),
|
|
)
|
|
verify_model(ScatterValue(), example_args, {}, Expected)
|
|
|
|
|
|
def test_grid_sample():
|
|
class GridSample(Module):
|
|
def forward(self, input, grid):
|
|
return torch.nn.functional.grid_sample(
|
|
input, grid, mode="bilinear", padding_mode="zeros", align_corners=True
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 4, 4), dtype="float32"),
|
|
grid: R.Tensor((1, 2, 2, 2), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 3, 2, 2), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 2, 2), dtype="float32") = R.image.grid_sample(
|
|
input_1,
|
|
grid,
|
|
method="bilinear",
|
|
layout="NCHW",
|
|
padding_mode="zeros",
|
|
align_corners=True,
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 2, 2), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 4, 4, dtype=torch.float32),
|
|
torch.randn(1, 2, 2, 2, dtype=torch.float32),
|
|
)
|
|
verify_model(GridSample(), example_args, {}, expected)
|
|
|
|
|
|
def test_torchvision_roi_align():
|
|
torchvision = pytest.importorskip("torchvision")
|
|
|
|
class ROIAlign(Module):
|
|
def forward(self, input, rois):
|
|
return torchvision.ops.roi_align(
|
|
input,
|
|
rois,
|
|
output_size=(3, 3),
|
|
spatial_scale=1.0,
|
|
sampling_ratio=2,
|
|
aligned=False,
|
|
)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
input_1: R.Tensor((1, 3, 8, 8), dtype="float32"),
|
|
rois: R.Tensor((2, 5), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 3, 3, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 3, 3), dtype="float32") = R.vision.roi_align(
|
|
input_1,
|
|
rois,
|
|
pooled_size=(3, 3),
|
|
spatial_scale=1.0,
|
|
sample_ratio=2,
|
|
layout="NCHW",
|
|
mode="avg",
|
|
)
|
|
gv: R.Tuple(R.Tensor((2, 3, 3, 3), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (
|
|
torch.randn(1, 3, 8, 8, dtype=torch.float32),
|
|
torch.tensor([[0.0, 1.0, 1.0, 6.0, 6.0], [0.0, 0.5, 0.5, 7.0, 7.0]], dtype=torch.float32),
|
|
)
|
|
verify_model(ROIAlign(), example_args, {}, expected)
|
|
|
|
|
|
def test_torchvision_roi_align_aligned():
|
|
torchvision = pytest.importorskip("torchvision")
|
|
|
|
class ROIAlign(Module):
|
|
def forward(self, input, rois):
|
|
return torchvision.ops.roi_align(
|
|
input,
|
|
rois,
|
|
output_size=(1, 1),
|
|
spatial_scale=1.0,
|
|
sampling_ratio=2,
|
|
aligned=True,
|
|
)
|
|
|
|
example_args = (
|
|
torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4),
|
|
torch.tensor([[0.0, 1.0, 1.0, 1.2, 1.2]], dtype=torch.float32),
|
|
)
|
|
verify_model_numerically(ROIAlign(), example_args, rtol=1e-5, atol=1e-5)
|
|
|
|
|
|
def test_upsample_nearest2d():
|
|
class UpsampleNearest2dScale(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.interpolate(input, scale_factor=2.0, mode="nearest")
|
|
|
|
class UpsampleNearest2dSize(Module):
|
|
def forward(self, input):
|
|
return torch.nn.functional.interpolate(input, size=(20, 20), mode="nearest")
|
|
|
|
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
|
|
|
|
@tvm.script.ir_module
|
|
class expected_scale:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 20, 20), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d(
|
|
input_1,
|
|
size=(20, 20),
|
|
layout="NCHW",
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="half_pixel",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 20, 20), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class expected_size:
|
|
@R.function
|
|
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
|
|
R.Tensor((1, 3, 20, 20), dtype="float32")
|
|
):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d(
|
|
input_1,
|
|
size=(20, 20),
|
|
layout="NCHW",
|
|
method="nearest_neighbor",
|
|
coordinate_transformation_mode="half_pixel",
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 3, 20, 20), dtype="float32")) = (lv,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify_model(UpsampleNearest2dScale(), example_args, {}, expected_scale)
|
|
verify_model(UpsampleNearest2dSize(), example_args, {}, expected_size)
|
|
|
|
|
|
def test_from_exported_program_sparse_csr_buffer():
|
|
class SparseCsrBufferModule(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
crow_indices = torch.tensor([0, 1, 2], dtype=torch.int64)
|
|
col_indices = torch.tensor([0, 1], dtype=torch.int64)
|
|
values = torch.tensor([1.0, 1.0], dtype=torch.float32, requires_grad=True)
|
|
csr_tensor = torch.sparse_csr_tensor(
|
|
crow_indices, col_indices, values, dtype=torch.float32
|
|
)
|
|
self.register_buffer("csr_tensor", csr_tensor)
|
|
self.csr_tensor.requires_grad_(True)
|
|
|
|
def forward(self, x):
|
|
csr2 = self.csr_tensor.to_sparse(layout=torch.sparse_csr)
|
|
y = torch.matmul(csr2, x)
|
|
return y.sum()
|
|
|
|
model = SparseCsrBufferModule().eval()
|
|
x = torch.ones((2, 1), dtype=torch.float32)
|
|
exported_program = export(model, (x,))
|
|
mod = from_exported_program(exported_program)
|
|
assert isinstance(mod, tvm.IRModule)
|
|
|
|
|
|
def test_cond_basic():
|
|
"""Basic data-dependent cond with runtime predicate."""
|
|
|
|
class CondModel(Module):
|
|
def forward(self, x):
|
|
def true_fn(x):
|
|
return x.cos()
|
|
|
|
def false_fn(x):
|
|
return x.sin()
|
|
|
|
return torch.cond(x.sum() > 0, true_fn, false_fn, (x,))
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def cond_true_branch_0(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tensor((3, 4), dtype="float32"):
|
|
gv: R.Tensor((3, 4), dtype="float32") = R.cos(x)
|
|
gv1: R.Tensor((3, 4), dtype="float32") = gv
|
|
return gv1
|
|
|
|
@R.function
|
|
def cond_false_branch_1(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tensor((3, 4), dtype="float32"):
|
|
gv: R.Tensor((3, 4), dtype="float32") = R.sin(x)
|
|
gv1: R.Tensor((3, 4), dtype="float32") = gv
|
|
return gv1
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((3, 4), dtype="float32")):
|
|
cls = expected
|
|
gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
gv1: R.Tensor((), dtype="bool") = R.greater(gv, R.const(0.0, "float32"))
|
|
if gv1:
|
|
gv2: R.Tensor((3, 4), dtype="float32") = cls.cond_true_branch_0(x)
|
|
cond_result: R.Tensor((3, 4), dtype="float32") = gv2
|
|
else:
|
|
gv3: R.Tensor((3, 4), dtype="float32") = cls.cond_false_branch_1(x)
|
|
cond_result: R.Tensor((3, 4), dtype="float32") = gv3
|
|
return (cond_result,)
|
|
|
|
verify_model(CondModel(), (torch.randn(3, 4),), {}, expected, map_free_vars=True)
|
|
|
|
|
|
def test_cond_shape_predicate():
|
|
"""Cond with a shape-derived predicate and dynamic shapes."""
|
|
|
|
class CondShapeModel(Module):
|
|
def forward(self, x):
|
|
def true_fn(x):
|
|
return x + 1.0
|
|
|
|
def false_fn(x):
|
|
return x - 1.0
|
|
|
|
return torch.cond(x.shape[0] > 4, true_fn, false_fn, (x,))
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def cond_true_branch_0(
|
|
x: R.Tensor(("s77", 4), dtype="float32"),
|
|
) -> R.Tensor(("s77", 4), dtype="float32"):
|
|
s77 = T.int64()
|
|
gv: R.Tensor((s77, 4), dtype="float32") = R.add(x, R.const(1.0, "float32"))
|
|
gv1: R.Tensor((s77, 4), dtype="float32") = gv
|
|
return gv1
|
|
|
|
@R.function
|
|
def cond_false_branch_1(
|
|
x: R.Tensor(("s77", 4), dtype="float32"),
|
|
) -> R.Tensor(("s77", 4), dtype="float32"):
|
|
s77 = T.int64()
|
|
gv: R.Tensor((s77, 4), dtype="float32") = R.subtract(x, R.const(1.0, "float32"))
|
|
gv1: R.Tensor((s77, 4), dtype="float32") = gv
|
|
return gv1
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(("s77", 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor(("s77", 4), dtype="float32")):
|
|
s77 = T.int64()
|
|
R.func_attr({"tir_var_lower_bound": {"s77": 1}})
|
|
cls = expected
|
|
gv: R.Tensor((), dtype="bool") = R.const(True, "bool")
|
|
if gv:
|
|
gv1: R.Tensor((s77, 4), dtype="float32") = cls.cond_true_branch_0(x)
|
|
cond_result: R.Tensor((s77, 4), dtype="float32") = gv1
|
|
else:
|
|
gv2: R.Tensor((s77, 4), dtype="float32") = cls.cond_false_branch_1(x)
|
|
cond_result: R.Tensor((s77, 4), dtype="float32") = gv2
|
|
return (cond_result,)
|
|
|
|
batch = torch.export.Dim("batch", min=1)
|
|
verify_model(
|
|
CondShapeModel(),
|
|
(torch.randn(3, 4),),
|
|
{},
|
|
expected,
|
|
dynamic_shapes={"x": {0: batch}},
|
|
map_free_vars=True,
|
|
)
|
|
|
|
|
|
def test_cond_tuple_output():
|
|
"""Cond where both branches return a tuple."""
|
|
|
|
class CondTupleModel(Module):
|
|
def forward(self, x):
|
|
def true_fn(x):
|
|
return (x.cos(), x.sin())
|
|
|
|
def false_fn(x):
|
|
return (x.sin(), x.cos())
|
|
|
|
return torch.cond(x.sum() > 0, true_fn, false_fn, (x,))
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def cond_true_branch_0(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((3, 4), dtype="float32"), R.Tensor((3, 4), dtype="float32")):
|
|
gv: R.Tensor((3, 4), dtype="float32") = R.cos(x)
|
|
gv1: R.Tensor((3, 4), dtype="float32") = R.sin(x)
|
|
gv2: R.Tensor((3, 4), dtype="float32") = gv
|
|
gv3: R.Tensor((3, 4), dtype="float32") = gv1
|
|
return (gv2, gv3)
|
|
|
|
@R.function
|
|
def cond_false_branch_1(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((3, 4), dtype="float32"), R.Tensor((3, 4), dtype="float32")):
|
|
gv: R.Tensor((3, 4), dtype="float32") = R.sin(x)
|
|
gv1: R.Tensor((3, 4), dtype="float32") = R.cos(x)
|
|
gv2: R.Tensor((3, 4), dtype="float32") = gv
|
|
gv3: R.Tensor((3, 4), dtype="float32") = gv1
|
|
return (gv2, gv3)
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((3, 4), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((3, 4), dtype="float32"), R.Tensor((3, 4), dtype="float32")):
|
|
cls = expected
|
|
gv: R.Tensor((), dtype="float32") = R.sum(x, axis=None, keepdims=False)
|
|
gv1: R.Tensor((), dtype="bool") = R.greater(gv, R.const(0.0, "float32"))
|
|
if gv1:
|
|
gv2: R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
) = cls.cond_true_branch_0(x)
|
|
cond_result: R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
) = gv2
|
|
else:
|
|
gv3: R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
) = cls.cond_false_branch_1(x)
|
|
cond_result: R.Tuple(
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
R.Tensor((3, 4), dtype="float32"),
|
|
) = gv3
|
|
gv4: R.Tensor((3, 4), dtype="float32") = cond_result[0]
|
|
gv5: R.Tensor((3, 4), dtype="float32") = cond_result[1]
|
|
return (gv4, gv5)
|
|
|
|
verify_model(CondTupleModel(), (torch.randn(3, 4),), {}, expected, map_free_vars=True)
|
|
|
|
|
|
def test_cond_nested():
|
|
"""Nested cond: a cond inside one of the branches."""
|
|
|
|
class CondNestedModel(Module):
|
|
def forward(self, x):
|
|
def true_fn(x):
|
|
def inner_true(x):
|
|
return x * 2.0
|
|
|
|
def inner_false(x):
|
|
return x * 3.0
|
|
|
|
return torch.cond(x.sum() > 1, inner_true, inner_false, (x,))
|
|
|
|
def false_fn(x):
|
|
return x - 1.0
|
|
|
|
return torch.cond(x.sum() > 0, true_fn, false_fn, (x,))
|
|
|
|
example_args = (torch.randn(3, 4),)
|
|
exported_program = export(CondNestedModel(), args=example_args)
|
|
mod = from_exported_program(exported_program)
|
|
|
|
assert isinstance(mod, tvm.IRModule)
|
|
|
|
# Should have at least 4 branch functions (2 outer + 2 inner) plus main
|
|
func_names = [gv.name_hint for gv in mod.get_global_vars()]
|
|
branch_funcs = [n for n in func_names if n != "main"]
|
|
assert len(branch_funcs) >= 4, (
|
|
f"Expected at least 4 branch functions for nested cond, got {branch_funcs}"
|
|
)
|
|
# Verify no duplicate function names
|
|
assert len(set(branch_funcs)) == len(branch_funcs), (
|
|
f"Duplicate branch function names: {branch_funcs}"
|
|
)
|
|
|
|
|
|
def test_affine_grid():
|
|
class AffineGrid(Module):
|
|
def forward(self, theta):
|
|
return torch.nn.functional.affine_grid(theta, [1, 3, 16, 16], align_corners=True)
|
|
|
|
@tvm.script.ir_module
|
|
class expected:
|
|
@R.function
|
|
def main(
|
|
theta: R.Tensor((1, 2, 3), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 16, 16, 2), dtype="float32")):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 2, 16, 16), dtype="float32") = R.image.affine_grid(
|
|
theta, size=(16, 16)
|
|
)
|
|
lv1: R.Tensor((1, 16, 16, 2), dtype="float32") = R.permute_dims(
|
|
lv, axes=[0, 2, 3, 1]
|
|
)
|
|
gv: R.Tuple(R.Tensor((1, 16, 16, 2), dtype="float32")) = (lv1,)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
example_args = (torch.randn(1, 2, 3, dtype=torch.float32),)
|
|
# Disable decomposition to keep aten.affine_grid_generator as a single op
|
|
verify_model(AffineGrid(), example_args, {}, expected, run_ep_decomposition=False)
|
|
|
|
|
|
def test_affine_grid_numerically():
|
|
"""Verify affine_grid numerical correctness: PyTorch vs TVM via our converter."""
|
|
|
|
class AffineGrid(Module):
|
|
def forward(self, theta):
|
|
return torch.nn.functional.affine_grid(theta, [2, 3, 8, 12], align_corners=True)
|
|
|
|
model = AffineGrid()
|
|
example_args = (torch.randn(2, 2, 3, dtype=torch.float32),)
|
|
|
|
with torch.no_grad():
|
|
pytorch_output = model(*example_args)
|
|
|
|
exported_program = export(model, args=example_args)
|
|
mod = from_exported_program(exported_program, run_ep_decomposition=False)
|
|
|
|
exe = tvm.compile(mod, target="llvm")
|
|
vm = relax.VirtualMachine(exe, tvm.cpu())
|
|
|
|
tvm_args = [tvm.runtime.tensor(arg.numpy()) for arg in example_args]
|
|
tvm_output = vm["main"](*tvm_args)
|
|
tvm_output_np = tvm_output[0].numpy()
|
|
|
|
tvm.testing.assert_allclose(tvm_output_np, pytorch_output.numpy(), rtol=1e-5, atol=1e-5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|