chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,842 @@
|
||||
# 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
|
||||
"""CLML integration operator tests."""
|
||||
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax, rpc
|
||||
from tvm.relax.backend.adreno import clml
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import relax as relax_builder
|
||||
|
||||
|
||||
def get_relax_conv2d_mod(
|
||||
data_shape,
|
||||
weight_shape,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
weight_layout="OIHW",
|
||||
groups=1,
|
||||
dtype="float32",
|
||||
has_bias=False,
|
||||
has_bn=False,
|
||||
has_activation=False,
|
||||
has_pad=False,
|
||||
is_depthwise=False,
|
||||
):
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
if has_pad:
|
||||
p = (0, 0, 0, 0, padding[0], padding[0], padding[1], padding[1])
|
||||
orig_data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
data = R.nn.pad(orig_data, pad_width=p, pad_value=0.0)
|
||||
padding = (0, 0, 0, 0)
|
||||
else:
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
weight = R.arg("weight", R.Tensor(weight_shape, dtype))
|
||||
if has_bias:
|
||||
bias = R.arg("bias", R.Tensor((1, weight_shape[0], 1, 1), dtype))
|
||||
|
||||
is_depthwise = data_shape[1] == weight_shape[0] == groups
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(
|
||||
R.nn.conv2d(
|
||||
data,
|
||||
weight,
|
||||
out_dtype=dtype,
|
||||
strides=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
data_layout="NCHW",
|
||||
kernel_layout=weight_layout,
|
||||
groups=groups,
|
||||
)
|
||||
)
|
||||
if has_bias:
|
||||
output = R.emit(output + bias)
|
||||
if has_bn:
|
||||
gamma = R.arg("gamma", R.Tensor((weight_shape[0],), dtype))
|
||||
beta = R.arg("beta", R.Tensor((weight_shape[0],), dtype))
|
||||
mean = R.arg("mean", R.Tensor((weight_shape[0],), dtype))
|
||||
variance = R.arg("variance", R.Tensor((weight_shape[0],), dtype))
|
||||
output = R.emit(
|
||||
R.nn.batch_norm(output, gamma, beta, mean, variance, axis=1, epsilon=1e-5)[
|
||||
0
|
||||
]
|
||||
)
|
||||
if has_activation:
|
||||
output = R.emit(R.nn.relu(output))
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_clml_conv2d_codegen(
|
||||
data_shape,
|
||||
weight_shape,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
weight_layout="OIHW",
|
||||
groups=1,
|
||||
dtype="float32",
|
||||
has_bias=False,
|
||||
has_bn=False,
|
||||
has_activation=False,
|
||||
has_pad=False,
|
||||
is_depthwise=False,
|
||||
):
|
||||
kernel_h, kernel_w = weight_shape[2], weight_shape[3]
|
||||
channels = weight_shape[0]
|
||||
if len(padding) == 2:
|
||||
padding = (padding[0], padding[1], padding[0], padding[1])
|
||||
output_height = ((data_shape[2] - kernel_h + padding[0] + padding[2]) / stride[0]) + 1
|
||||
output_width = ((data_shape[3] - kernel_w + padding[1] + padding[3]) / stride[1]) + 1
|
||||
output_shape = (1, channels, int(output_height), int(output_width))
|
||||
out_dtype = dtype
|
||||
is_depthwise = data_shape[1] == channels == groups
|
||||
|
||||
weight_layout = "IOHW" if is_depthwise else "OIHW"
|
||||
if weight_layout == "OIHW":
|
||||
weight_shape = (channels, data_shape[1] // groups, kernel_h, kernel_w)
|
||||
else:
|
||||
weight_shape = (data_shape[1] // groups, channels, kernel_h, kernel_w)
|
||||
|
||||
if is_depthwise:
|
||||
name = "openclml.nn.depthwise_conv2d"
|
||||
else:
|
||||
name = "openclml.nn.conv2d"
|
||||
|
||||
node = {
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [],
|
||||
"attrs": {
|
||||
"groups": groups,
|
||||
"num_outputs": 1,
|
||||
"data_layout": "NCHW",
|
||||
"kernel_layout": weight_layout,
|
||||
"dilation": dilation,
|
||||
"out_layout": "NCHW",
|
||||
"out_dtype": out_dtype,
|
||||
"shape": [list(output_shape)],
|
||||
"dtype": [dtype],
|
||||
"padding": padding,
|
||||
"strides": stride,
|
||||
},
|
||||
}
|
||||
|
||||
if has_activation:
|
||||
node["attrs"]["activation_type"] = "relu"
|
||||
|
||||
nodes = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(data_shape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
]
|
||||
|
||||
nodes.append(
|
||||
{
|
||||
"op": "const",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(weight_shape)], "dtype": [str(dtype)]},
|
||||
}
|
||||
)
|
||||
|
||||
if has_bias:
|
||||
bias_dtype = dtype
|
||||
nodes.append(
|
||||
{
|
||||
"op": "const",
|
||||
"name": "",
|
||||
"attrs": {
|
||||
"shape": [[1, weight_shape[1] if is_depthwise else weight_shape[0], 1, 1]],
|
||||
"dtype": [bias_dtype],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
if has_bn:
|
||||
bn_shape = [1, weight_shape[0], 1, 1]
|
||||
# conv2d + bn --> conv2d + Add due to OptimizeBatchNorm transformation Pass
|
||||
nodes.append(
|
||||
{
|
||||
"name": "",
|
||||
"op": "const",
|
||||
"attrs": {"dtype": [dtype], "shape": [[1, weight_shape[0], 1, 1]]},
|
||||
},
|
||||
)
|
||||
|
||||
input_idx = 0
|
||||
for _ in range(len(nodes)):
|
||||
node["inputs"].append([input_idx, 0, 0])
|
||||
input_idx += 1
|
||||
node["attrs"]["num_inputs"] = len(nodes)
|
||||
nodes.append(node)
|
||||
return nodes
|
||||
|
||||
|
||||
def get_relax_conv2d_transpose_mod(
|
||||
data_shape,
|
||||
weight_shape,
|
||||
channels,
|
||||
stride,
|
||||
padding,
|
||||
dtype="float32",
|
||||
):
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
weight = R.arg("weight", R.Tensor(weight_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(
|
||||
R.nn.conv2d_transpose(
|
||||
data,
|
||||
weight,
|
||||
groups=1,
|
||||
strides=stride,
|
||||
padding=padding,
|
||||
kernel_layout="OIHW",
|
||||
data_layout="NCHW",
|
||||
)
|
||||
)
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_conv2d_transpose_expected_codegen(
|
||||
dshape, kshape, channels, kernel_size, strides, padding, dilation, dtype, output_shape
|
||||
):
|
||||
attrs = {
|
||||
"data_layout": "NCHW",
|
||||
"kernel_layout": "OIHW",
|
||||
"groups": 1,
|
||||
"dilation": dilation,
|
||||
"num_inputs": 2,
|
||||
"num_outputs": 1,
|
||||
"padding": padding,
|
||||
"shape": [list(output_shape)],
|
||||
"dtype": [dtype],
|
||||
"strides": strides,
|
||||
"out_dtype": "",
|
||||
"out_layout": "NCHW",
|
||||
"output_padding": [0, 0],
|
||||
}
|
||||
|
||||
exp_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(dshape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
{
|
||||
"op": "const",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(kshape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
{
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [[0, 0, 0], [1, 0, 0]],
|
||||
"attrs": attrs,
|
||||
},
|
||||
]
|
||||
return exp_codegen
|
||||
|
||||
|
||||
def get_batchnorm_mod(data_shape, channels, axis, epsilon, dtype):
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
gamma = R.arg("gamma", R.Tensor((channels,), dtype))
|
||||
beta = R.arg("beta", R.Tensor((channels,), dtype))
|
||||
mean = R.arg("moving_mean", R.Tensor((channels,), dtype))
|
||||
variance = R.arg("moving_var", R.Tensor((channels,), dtype))
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(
|
||||
R.nn.batch_norm(data, gamma, beta, mean, variance, axis, epsilon)[0]
|
||||
)
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_binary_op_mod(a_shape, b_shape, op, dtype):
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
a = R.arg("a", R.Tensor(a_shape, dtype))
|
||||
b = R.arg("b", R.Tensor(b_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(op(a, b))
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
|
||||
low, high = 0, 1
|
||||
a_data = np.random.uniform(low, high, size=(a_shape)).astype(dtype)
|
||||
b_data = np.random.uniform(low, high, size=(b_shape)).astype(dtype)
|
||||
|
||||
return (tvm.IRModule({"main": func}), (a_data, b_data))
|
||||
|
||||
|
||||
def get_unary_op_mod(a_shape, op, dtype):
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
a = R.arg("a", R.Tensor(a_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(op(a))
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
|
||||
low, high = 0, 1
|
||||
a_data = np.random.uniform(low, high, size=(a_shape)).astype(dtype)
|
||||
|
||||
return (tvm.IRModule({"main": func}), (a_data,))
|
||||
|
||||
|
||||
def get_relax_maxpool_mod(
|
||||
data_shape, dtype, pool_size, stride=None, dilation=(1, 1), padding=(0, 0), has_pad=False
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
data_shape (tuple): Input tensor shape
|
||||
pool_size (tuple): Pooling window size (height, width)
|
||||
stride (tuple, optional): Stride of pooling operation. Defaults to pool_size.
|
||||
dilation (tuple, optional): Dilation rate. Defaults to (1, 1).
|
||||
padding (tuple, optional): Padding for the input tensor. Defaults to (0, 0).
|
||||
dtype (str, optional): Data type. Defaults to "float32".
|
||||
has_pad (bool, optional): Whether to apply explicit padding. Defaults to False.
|
||||
|
||||
Returns:
|
||||
tvm.IRModule: Relax MaxPool module
|
||||
"""
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
|
||||
if has_pad:
|
||||
p = (0, 0, 0, 0, padding[0], padding[1], padding[0], padding[1])
|
||||
orig_data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
data = R.nn.pad(orig_data, pad_width=p, pad_value=float("-inf"))
|
||||
padding = (0, 0)
|
||||
else:
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(
|
||||
R.nn.max_pool2d(
|
||||
data,
|
||||
pool_size=pool_size,
|
||||
strides=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
layout="NCHW",
|
||||
)
|
||||
)
|
||||
R.output(output)
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_maxpool_expected_codegen(input_shape, pool_size, stride, padding, pool_type, dtype):
|
||||
import math
|
||||
|
||||
adjusted_input_shape = [
|
||||
input_shape[0],
|
||||
input_shape[1],
|
||||
input_shape[2] + padding[0] + padding[1],
|
||||
input_shape[3] + padding[2] + padding[3],
|
||||
]
|
||||
|
||||
pool_height = math.floor(((adjusted_input_shape[2] - pool_size[0]) / stride[0]) + 1)
|
||||
pool_width = math.floor(((adjusted_input_shape[3] - pool_size[1]) / stride[1]) + 1)
|
||||
output_shape = [adjusted_input_shape[0], adjusted_input_shape[1], pool_height, pool_width]
|
||||
|
||||
attrs = {
|
||||
"ceil_mode": 0,
|
||||
"dilation": [1, 1],
|
||||
"layout": "NCHW",
|
||||
"num_inputs": 1,
|
||||
"num_outputs": 1,
|
||||
"out_layout": "NCHW",
|
||||
"padding": list(padding),
|
||||
"pool_size": pool_size,
|
||||
"shape": [list(output_shape)],
|
||||
"dtype": [dtype],
|
||||
"strides": stride,
|
||||
"count_include_pad": 0,
|
||||
}
|
||||
if sum(padding):
|
||||
attrs["count_include_pad"] = 0
|
||||
|
||||
exp_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(adjusted_input_shape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
{
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [[0, 0, 0]],
|
||||
"attrs": attrs,
|
||||
},
|
||||
]
|
||||
return exp_codegen
|
||||
|
||||
|
||||
def get_relax_avgpool_mod(data_shape, dtype, pool_size, stride, dilation, padding, has_pad):
|
||||
"""
|
||||
Args:
|
||||
data_shape (tuple): Input tensor shape
|
||||
pool_size (tuple): Pooling window size (height, width)
|
||||
stride (tuple, optional): Stride of pooling operation. Defaults to pool_size.
|
||||
dilation (tuple, optional): Dilation rate. Defaults to (1, 1).
|
||||
padding (tuple, optional): Padding for the input tensor. Defaults to (0, 0).
|
||||
dtype (str, optional): Data type. Defaults to "float32".
|
||||
has_pad (bool, optional): Whether to apply explicit padding. Defaults to False.
|
||||
count_include_pad (bool, optional): Whether to include padding in averaging. Defaults to True.
|
||||
|
||||
Returns:
|
||||
tvm.IRModule: Relax AvgPool module
|
||||
"""
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
|
||||
if has_pad:
|
||||
p = (0, 0, 0, 0, padding[0], padding[1], padding[0], padding[1])
|
||||
orig_data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
data = R.nn.pad(orig_data, pad_width=p, pad_value=0.0)
|
||||
padding = (0, 0)
|
||||
else:
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(
|
||||
R.nn.avg_pool2d(
|
||||
data,
|
||||
pool_size=pool_size,
|
||||
strides=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
layout="NCHW",
|
||||
)
|
||||
)
|
||||
R.output(output)
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_avgpool_expected_codegen(input_shape, pool_size, stride, padding, pool_type, dtype):
|
||||
import math
|
||||
|
||||
adjusted_input_shape = [
|
||||
input_shape[0],
|
||||
input_shape[1],
|
||||
input_shape[2] + padding[0] + padding[1],
|
||||
input_shape[3] + padding[2] + padding[3],
|
||||
]
|
||||
|
||||
pool_height = math.floor(((adjusted_input_shape[2] - pool_size[0]) / stride[0]) + 1)
|
||||
pool_width = math.floor(((adjusted_input_shape[3] - pool_size[1]) / stride[1]) + 1)
|
||||
output_shape = [adjusted_input_shape[0], adjusted_input_shape[1], pool_height, pool_width]
|
||||
|
||||
attrs = {
|
||||
"ceil_mode": 0,
|
||||
"dilation": [1, 1],
|
||||
"layout": "NCHW",
|
||||
"num_inputs": 1,
|
||||
"num_outputs": 1,
|
||||
"out_layout": "NCHW",
|
||||
"padding": list(padding),
|
||||
"pool_size": pool_size,
|
||||
"shape": [list(output_shape)],
|
||||
"dtype": [dtype],
|
||||
"strides": stride,
|
||||
"count_include_pad": 0,
|
||||
}
|
||||
if sum(padding):
|
||||
attrs["count_include_pad"] = 0
|
||||
|
||||
exp_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(adjusted_input_shape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
{
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [[0, 0, 0]],
|
||||
"attrs": attrs,
|
||||
},
|
||||
]
|
||||
return exp_codegen
|
||||
|
||||
|
||||
def get_relax_reshape_mod(input_shape, output_shape, dtype):
|
||||
"""
|
||||
Args:
|
||||
input_shape (tuple): Input tensor shape
|
||||
output_shape (tuple): Desired output tensor shape
|
||||
dtype (str, optional): Data type. Defaults to "float32".
|
||||
|
||||
Returns:
|
||||
tvm.IRModule: Relax Reshape module
|
||||
"""
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
data = R.arg("data", R.Tensor(input_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(R.reshape(data, output_shape))
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_relax_reshape_codegen(input_shape, output_shape, dtype):
|
||||
def compute_output_shape(input_shape, output_shape):
|
||||
input_elements = np.prod(input_shape)
|
||||
specified_elements = np.prod([dim for dim in output_shape if dim != -1])
|
||||
missing_dim = input_elements // specified_elements
|
||||
return [int(dim) if dim != -1 else int(missing_dim) for dim in output_shape]
|
||||
|
||||
expected_output_shape = compute_output_shape(input_shape, output_shape)
|
||||
|
||||
expected_codegen_str = [
|
||||
{
|
||||
"attrs": {
|
||||
"dtype": [dtype],
|
||||
"shape": [list(input_shape)],
|
||||
},
|
||||
"name": "",
|
||||
"op": "input",
|
||||
},
|
||||
{
|
||||
"attrs": {
|
||||
"dtype": [dtype],
|
||||
"num_inputs": 1,
|
||||
"num_outputs": 1,
|
||||
"shape": [expected_output_shape],
|
||||
},
|
||||
"inputs": [[0, 0, 0]],
|
||||
"name": "",
|
||||
"op": "kernel",
|
||||
},
|
||||
]
|
||||
return expected_codegen_str
|
||||
|
||||
|
||||
def get_relax_global_avgpool_mod(data_shape, keepdims, dtype):
|
||||
"""
|
||||
Create a Relax module for Global Average Pooling (GAP).
|
||||
|
||||
Args:
|
||||
data_shape (tuple): Input tensor shape (N, C, H, W)
|
||||
dtype (str): Data type
|
||||
|
||||
Returns:
|
||||
tvm.IRModule: Relax GAP module
|
||||
"""
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(R.mean(data, axis=[2, 3], keepdims=keepdims))
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_global_avgpool_expected_codegen(input_shape, keep_dims, dtype):
|
||||
"""
|
||||
Generate expected codegen for Global Average Pooling.
|
||||
|
||||
Args:
|
||||
input_shape (tuple): Input shape (N, C, H, W)
|
||||
dtype (str): Data type
|
||||
|
||||
Returns:
|
||||
dict: Expected codegen output
|
||||
"""
|
||||
output_shape = (
|
||||
[input_shape[0], input_shape[1]]
|
||||
if not keep_dims
|
||||
else [input_shape[0], input_shape[1], 1, 1]
|
||||
)
|
||||
attrs = {
|
||||
"num_inputs": 1,
|
||||
"num_outputs": 1,
|
||||
"shape": [list(output_shape)],
|
||||
"dtype": [dtype],
|
||||
"axis": [2, 3],
|
||||
"keepdims": 1 if keep_dims else 0,
|
||||
}
|
||||
|
||||
exp_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(input_shape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
{"op": "kernel", "name": "", "inputs": [[0, 0, 0]], "attrs": attrs},
|
||||
]
|
||||
return exp_codegen
|
||||
|
||||
|
||||
def get_relax_global_maxpool_mod(data_shape, keepdims, dtype):
|
||||
"""
|
||||
Create a Relax module for Global Average Pooling (GAP).
|
||||
|
||||
Args:
|
||||
data_shape (tuple): Input tensor shape (N, C, H, W)
|
||||
dtype (str): Data type
|
||||
|
||||
Returns:
|
||||
tvm.IRModule: Relax GAP module
|
||||
"""
|
||||
N, C, H, W = data_shape
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
data = R.arg("data", R.Tensor(data_shape, dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
output = R.emit(
|
||||
R.nn.max_pool2d(
|
||||
data, pool_size=(H, W), strides=(1, 1), padding=(0, 0), layout="NCHW"
|
||||
)
|
||||
)
|
||||
R.output(output)
|
||||
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_global_maxpool_expected_codegen(input_shape, pool_size, stride, padding, pool_type, dtype):
|
||||
import math
|
||||
|
||||
adjusted_input_shape = [
|
||||
input_shape[0],
|
||||
input_shape[1],
|
||||
input_shape[2] + padding[0] + padding[1],
|
||||
input_shape[3] + padding[2] + padding[3],
|
||||
]
|
||||
|
||||
output_shape = [adjusted_input_shape[0], adjusted_input_shape[1], 1, 1]
|
||||
|
||||
attrs = {
|
||||
"ceil_mode": 0,
|
||||
"dilation": [1, 1],
|
||||
"layout": "NCHW",
|
||||
"num_inputs": 1,
|
||||
"num_outputs": 1,
|
||||
"out_layout": "NCHW",
|
||||
"padding": padding,
|
||||
"pool_size": pool_size,
|
||||
"shape": [list(output_shape)],
|
||||
"dtype": [dtype],
|
||||
"strides": stride,
|
||||
"count_include_pad": 0,
|
||||
}
|
||||
if sum(padding):
|
||||
attrs["count_include_pad"] = 0
|
||||
|
||||
exp_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [list(adjusted_input_shape)], "dtype": [str(dtype)]},
|
||||
},
|
||||
{
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [[0, 0, 0]],
|
||||
"attrs": attrs,
|
||||
},
|
||||
]
|
||||
return exp_codegen
|
||||
|
||||
|
||||
def get_dequant_matmul_module(K, N):
|
||||
@I.ir_module(s_tir=True)
|
||||
class DequantMatmul:
|
||||
@R.function
|
||||
def main(
|
||||
input: R.Tensor((1, "seq_len", K), dtype="float16"),
|
||||
weight: R.Tensor((K // 8, N), dtype="uint32"),
|
||||
scale: R.Tensor((K // 32, N), dtype="float16"),
|
||||
):
|
||||
seq_len = T.int64()
|
||||
cls = DequantMatmul
|
||||
with R.dataflow():
|
||||
lv2 = relax.call_tir(
|
||||
cls.dequantize,
|
||||
(weight, scale),
|
||||
out_ty=R.Tensor((K, N), dtype="float16"),
|
||||
)
|
||||
gv: R.Tensor((1, seq_len, N), dtype="float16") = relax.op.matmul(
|
||||
input, lv2, out_dtype="float16"
|
||||
)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def dequantize(weight: T.handle, scale: T.handle, var_dequantize: T.handle):
|
||||
T.func_attr({"tirx.noalias": T.bool(True)})
|
||||
lm_head_q_weight1 = T.match_buffer(weight, (T.int64(K // 8), T.int64(N)), "uint32")
|
||||
lm_head_q_scale1 = T.match_buffer(scale, (T.int64(K // 32), T.int64(N)), "float16")
|
||||
dequantize = T.match_buffer(var_dequantize, (T.int64(K), T.int64(N)), "float16")
|
||||
# with T.sblock("root"):
|
||||
compute = T.alloc_buffer((T.int64(K), T.int64(N)), "float16")
|
||||
for i0, i1 in T.grid(T.int64(K), T.int64(N)):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(lm_head_q_weight1[v_i0 // T.int64(8), v_i1])
|
||||
T.writes(compute[v_i0, v_i1])
|
||||
compute[v_i0, v_i1] = T.Cast(
|
||||
"float16",
|
||||
T.bitwise_and(
|
||||
T.shift_right(
|
||||
lm_head_q_weight1[v_i0 // T.int64(8), v_i1],
|
||||
T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4)),
|
||||
),
|
||||
T.uint32(15),
|
||||
),
|
||||
)
|
||||
for i0, i1 in T.grid(T.int64(K), T.int64(N)):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(compute[v_i0, v_i1], lm_head_q_scale1[v_i0 // T.int64(32), v_i1])
|
||||
T.writes(dequantize[v_i0, v_i1])
|
||||
dequantize[v_i0, v_i1] = (
|
||||
compute[v_i0, v_i1] - T.float16(7.0)
|
||||
) * lm_head_q_scale1[v_i0 // T.int64(32), v_i1]
|
||||
|
||||
return DequantMatmul
|
||||
|
||||
|
||||
def get_dequant_vec_matmul_module(K, N):
|
||||
@I.ir_module(s_tir=True)
|
||||
class DequantVecMatmul:
|
||||
@R.function
|
||||
def main(
|
||||
input: R.Tensor((1, 1, K), dtype="float16"),
|
||||
weight: R.Tensor((K // 8, "vocab_size"), dtype="uint32"),
|
||||
scale: R.Tensor((K // 32, "vocab_size"), dtype="float16"),
|
||||
):
|
||||
vocab_size = T.int64()
|
||||
cls = DequantVecMatmul
|
||||
with R.dataflow():
|
||||
lv2 = relax.call_tir(
|
||||
cls.dequantize,
|
||||
(weight, scale),
|
||||
out_ty=R.Tensor((K, vocab_size), dtype="float16"),
|
||||
)
|
||||
gv: R.Tensor((1, 1, vocab_size), dtype="float16") = relax.op.matmul(
|
||||
input, lv2, out_dtype="float16"
|
||||
)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def dequantize(weight: T.handle, scale: T.handle, var_dequantize: T.handle):
|
||||
T.func_attr({"tirx.noalias": T.bool(True)})
|
||||
vocab_size = T.int64()
|
||||
lm_head_q_weight1 = T.match_buffer(weight, (T.int64(K // 8), vocab_size), "uint32")
|
||||
lm_head_q_scale1 = T.match_buffer(scale, (T.int64(K // 32), vocab_size), "float16")
|
||||
dequantize = T.match_buffer(var_dequantize, (T.int64(K), vocab_size), "float16")
|
||||
# with T.sblock("root"):
|
||||
compute = T.alloc_buffer((T.int64(K), vocab_size), "float16")
|
||||
for i0, i1 in T.grid(T.int64(K), vocab_size):
|
||||
with T.sblock("compute"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(lm_head_q_weight1[v_i0 // T.int64(8), v_i1])
|
||||
T.writes(compute[v_i0, v_i1])
|
||||
compute[v_i0, v_i1] = T.Cast(
|
||||
"float16",
|
||||
T.bitwise_and(
|
||||
T.shift_right(
|
||||
lm_head_q_weight1[v_i0 // T.int64(8), v_i1],
|
||||
T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4)),
|
||||
),
|
||||
T.uint32(15),
|
||||
),
|
||||
)
|
||||
for i0, i1 in T.grid(T.int64(K), vocab_size):
|
||||
with T.sblock("dequantize"):
|
||||
v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(compute[v_i0, v_i1], lm_head_q_scale1[v_i0 // T.int64(32), v_i1])
|
||||
T.writes(dequantize[v_i0, v_i1])
|
||||
dequantize[v_i0, v_i1] = (
|
||||
compute[v_i0, v_i1] - T.float16(7.0)
|
||||
) * lm_head_q_scale1[v_i0 // T.int64(32), v_i1]
|
||||
|
||||
return DequantVecMatmul
|
||||
@@ -0,0 +1,675 @@
|
||||
# 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
|
||||
"""CLML integration operator tests."""
|
||||
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from mod_utils import (
|
||||
get_avgpool_expected_codegen,
|
||||
get_batchnorm_mod,
|
||||
get_binary_op_mod,
|
||||
get_clml_conv2d_codegen,
|
||||
get_conv2d_transpose_expected_codegen,
|
||||
get_dequant_matmul_module,
|
||||
get_dequant_vec_matmul_module,
|
||||
get_global_avgpool_expected_codegen,
|
||||
get_global_maxpool_expected_codegen,
|
||||
get_maxpool_expected_codegen,
|
||||
get_relax_avgpool_mod,
|
||||
get_relax_conv2d_mod,
|
||||
get_relax_conv2d_transpose_mod,
|
||||
get_relax_global_avgpool_mod,
|
||||
get_relax_global_maxpool_mod,
|
||||
get_relax_maxpool_mod,
|
||||
get_relax_reshape_codegen,
|
||||
get_relax_reshape_mod,
|
||||
get_unary_op_mod,
|
||||
)
|
||||
from utils import skip_unless_adreno_clml, verify_results
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax, rpc
|
||||
from tvm.relax.backend.adreno import clml
|
||||
from tvm.relax.backend.adreno.clml import OpenCLMLOffLoad, OpenCLMLOffLoadForLLM
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import relax as relax_builder
|
||||
|
||||
CLML_VERSION = clml.clml_sdk_version()
|
||||
TARGET_CLML_VERSION = int(os.environ.get("ADRENO_TARGET_CLML_VERSION", 4))
|
||||
clml_target = tvm.target.Target("qcom/adreno-opencl-clml")
|
||||
ref_target = tvm.target.Target("opencl")
|
||||
|
||||
|
||||
def verify_clml_codegen(clml_mod, clml_codegen):
|
||||
clml_mod = OpenCLMLOffLoadForLLM(clml_target)(clml_mod)
|
||||
clml_mod = OpenCLMLOffLoad()(clml_mod)
|
||||
|
||||
source = clml_mod.attrs["external_mods"][0].inspect_source()
|
||||
codegen = json.loads(source)["nodes"]
|
||||
for node in range(len(codegen)):
|
||||
if codegen[node]["op"] == "input" or codegen[node]["op"] == "const":
|
||||
codegen[node]["name"] = ""
|
||||
if codegen[node]["op"] == "kernel":
|
||||
codegen[node]["name"] = ""
|
||||
|
||||
codegen_str = json.dumps(codegen, sort_keys=True, indent=2)
|
||||
known_good_codegen_str = json.dumps(clml_codegen, sort_keys=True, indent=2)
|
||||
assert codegen_str == known_good_codegen_str, (
|
||||
f"The JSON produced by codegen does not match the expected result. \n"
|
||||
f"Actual={codegen_str} \n"
|
||||
f"Expected={known_good_codegen_str}"
|
||||
)
|
||||
|
||||
|
||||
def verify(
|
||||
mod, clml_codegen, inputs_np, params_np, target_minimum_clml_version=None, target_test=True
|
||||
):
|
||||
mod = tvm.relax.transform.BindParams("main", params_np)(mod)
|
||||
codegen_mod, clml_mod = mod.clone(), mod
|
||||
verify_clml_codegen(codegen_mod, clml_codegen)
|
||||
|
||||
if (
|
||||
target_minimum_clml_version is not None
|
||||
and TARGET_CLML_VERSION < target_minimum_clml_version
|
||||
):
|
||||
print(f"Skipped Eval Tests for {inspect.stack()[1].function} function", flush=True)
|
||||
return
|
||||
|
||||
if "ADRENO_TARGET" not in os.environ:
|
||||
return
|
||||
|
||||
if target_test:
|
||||
verify_results(clml_mod, target=clml_target, ref_target=ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"kernel_h, kernel_w, padding, stride, dilation, out_channels, shape, has_bias, has_bn, has_activation, has_pad, is_depthwise",
|
||||
[
|
||||
(3, 3, (1, 1), (1, 1), (1, 1), 64, (3, 224, 224), False, True, False, True, False),
|
||||
(3, 3, (1, 1), (1, 1), (1, 1), 64, (3, 224, 224), False, True, False, False, False),
|
||||
# (5, 5, (2, 2), (1, 1), (1, 1), 16, (16, 64, 64), False, True, True, False, False),
|
||||
# (7, 7, (3, 3), (2, 2), (1, 1), 32, (3, 224, 224), True, False, True, True, False),
|
||||
(3, 3, (0, 0), (1, 1), (1, 1), 512, (256, 14, 14), True, False, True, False, False),
|
||||
(1, 1, (0, 0), (1, 1), (1, 1), 1024, (512, 7, 7), True, False, True, False, False),
|
||||
(1, 3, (0, 0), (1, 1), (1, 1), 64, (64, 7, 7), True, False, True, False, False),
|
||||
(3, 1, (0, 0), (1, 1), (1, 1), 64, (64, 7, 7), False, True, True, True, False),
|
||||
],
|
||||
)
|
||||
def test_conv2d_offload(
|
||||
kernel_h,
|
||||
kernel_w,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_channels,
|
||||
shape,
|
||||
has_bias,
|
||||
has_bn,
|
||||
has_activation,
|
||||
has_pad,
|
||||
is_depthwise,
|
||||
dtype,
|
||||
):
|
||||
low, high = -0.01, 0.01
|
||||
rtol, atol = 1e-3, 1e-3
|
||||
if CLML_VERSION > 3:
|
||||
rtol, atol = 1e-2, 1e-2 # @clml precision
|
||||
|
||||
data_shape = (1, *shape)
|
||||
if is_depthwise:
|
||||
groups = data_shape[1] // out_channels
|
||||
else:
|
||||
groups = 1
|
||||
padding = (padding[0], padding[1], padding[0], padding[1])
|
||||
|
||||
weight_format = "IOHW" if is_depthwise else "OIHW"
|
||||
weight_shape = (out_channels, data_shape[1] // groups, kernel_h, kernel_w)
|
||||
|
||||
data = np.random.uniform(low, high, size=data_shape).astype(dtype)
|
||||
weight = np.random.uniform(low, high, size=weight_shape).astype(dtype)
|
||||
bias = np.random.uniform(low, high, size=(1, weight_shape[0], 1, 1)).astype(dtype)
|
||||
|
||||
gamma = np.random.uniform(low, high, size=(weight_shape[0],)).astype(dtype)
|
||||
beta = np.random.uniform(low, high, size=(weight_shape[0],)).astype(dtype)
|
||||
mean = np.random.uniform(low, high, size=(weight_shape[0],)).astype(dtype)
|
||||
variance = np.random.uniform(low, high, size=(weight_shape[0],)).astype(dtype)
|
||||
|
||||
inputs_np = [data]
|
||||
params_np = {"weight": weight}
|
||||
if has_bias:
|
||||
params_np["bias"] = bias
|
||||
if has_bn:
|
||||
params_np.update({"gamma": gamma, "beta": beta, "mean": mean, "variance": variance})
|
||||
|
||||
mod = get_relax_conv2d_mod(
|
||||
data_shape,
|
||||
weight_shape,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
weight_layout=weight_format,
|
||||
groups=groups,
|
||||
dtype=dtype,
|
||||
has_bias=has_bias,
|
||||
has_bn=has_bn,
|
||||
has_activation=has_activation,
|
||||
has_pad=has_pad,
|
||||
is_depthwise=is_depthwise,
|
||||
)
|
||||
clml_codegen = get_clml_conv2d_codegen(
|
||||
data_shape,
|
||||
weight_shape,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
padding=padding,
|
||||
weight_layout=weight_format,
|
||||
groups=groups,
|
||||
dtype=dtype,
|
||||
has_bias=has_bias,
|
||||
has_bn=has_bn,
|
||||
has_activation=has_activation,
|
||||
has_pad=has_pad,
|
||||
is_depthwise=is_depthwise,
|
||||
)
|
||||
verify(mod, clml_codegen, inputs_np, params_np)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"dshape, kshape, channels, kernel_size, strides, padding, out_shape",
|
||||
[
|
||||
((1, 256, 100, 100), (64, 256, 4, 4), 64, (4, 4), (2, 2), (0, 0, 0, 0), (1, 64, 202, 202)),
|
||||
((1, 64, 200, 200), (64, 64, 4, 4), 64, (4, 4), (2, 2), (1, 1, 1, 1), (1, 64, 400, 400)),
|
||||
((1, 64, 200, 200), (64, 64, 4, 4), 64, (4, 4), (2, 2), (1, 1, 1, 1), (1, 64, 400, 400)),
|
||||
((1, 64, 400, 400), (16, 64, 4, 4), 16, (4, 4), (2, 2), (1, 1, 1, 1), (1, 16, 800, 800)),
|
||||
],
|
||||
)
|
||||
def test_conv2d_transpose(
|
||||
dshape, kshape, channels, kernel_size, strides, padding, dtype, out_shape
|
||||
):
|
||||
low, high = -1, 1
|
||||
|
||||
data = np.random.uniform(low, high, size=dshape).astype(dtype)
|
||||
weight = np.random.uniform(low, high, size=kshape).astype(dtype)
|
||||
|
||||
inputs_np = [data]
|
||||
params_np = {"weight": weight}
|
||||
|
||||
mod = get_relax_conv2d_transpose_mod(
|
||||
dshape,
|
||||
kshape,
|
||||
channels=channels,
|
||||
stride=strides,
|
||||
padding=padding,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
clml_codegen = get_conv2d_transpose_expected_codegen(
|
||||
dshape=dshape,
|
||||
kshape=kshape,
|
||||
channels=channels,
|
||||
kernel_size=kernel_size,
|
||||
strides=strides,
|
||||
padding=padding,
|
||||
dilation=(1, 1),
|
||||
dtype=dtype,
|
||||
output_shape=out_shape,
|
||||
)
|
||||
verify(mod, clml_codegen, inputs_np, params_np, target_test=False)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.skipif(
|
||||
CLML_VERSION < 3,
|
||||
reason="Requires compiler supporting CLML v5 or above",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"trials",
|
||||
[
|
||||
[(1, 64, 14, 14), 1, 3e-4],
|
||||
[(1, 14, 256, 256), 1, 3e-4],
|
||||
[(1, 14, 256, 256), 1, 3e-4],
|
||||
[(1, 256, 1, 1), 1, 3e-4],
|
||||
],
|
||||
)
|
||||
def test_batchnorm(dtype, trials):
|
||||
low, high = 0, 1
|
||||
(input_shape, axis, epsilon) = trials
|
||||
channels = input_shape[axis]
|
||||
|
||||
def _get_axis_tuple(axis):
|
||||
if axis == 0:
|
||||
return (1, 2, 3)
|
||||
elif axis == 1:
|
||||
return (0, 2, 3)
|
||||
elif axis == 2:
|
||||
return (0, 1, 3)
|
||||
else:
|
||||
return (0, 1, 2)
|
||||
|
||||
data = np.random.uniform(low, high, size=(input_shape)).astype(dtype)
|
||||
gamma = np.random.uniform(low, high, size=(channels)).astype(dtype)
|
||||
beta = np.random.uniform(low, high, size=(channels)).astype(dtype)
|
||||
mean = np.mean(data, _get_axis_tuple(axis), keepdims=False)
|
||||
variance = np.var(data, _get_axis_tuple(axis), keepdims=False)
|
||||
|
||||
inputs_np = [data]
|
||||
params_np = {"gamma": gamma, "beta": beta, "moving_mean": mean, "moving_var": variance}
|
||||
mod = get_batchnorm_mod(input_shape, channels, axis, epsilon, dtype)
|
||||
clml_codegen = [
|
||||
{
|
||||
"attrs": {"dtype": [dtype], "shape": [input_shape]},
|
||||
"name": "",
|
||||
"op": "input",
|
||||
},
|
||||
{"attrs": {"dtype": [dtype], "shape": [[channels]]}, "name": "", "op": "const"},
|
||||
{"attrs": {"dtype": [dtype], "shape": [[channels]]}, "name": "", "op": "const"},
|
||||
{"attrs": {"dtype": [dtype], "shape": [[channels]]}, "name": "", "op": "const"},
|
||||
{"attrs": {"dtype": [dtype], "shape": [[channels]]}, "name": "", "op": "const"},
|
||||
{
|
||||
"attrs": {
|
||||
"axis": axis,
|
||||
"center": 1,
|
||||
"dtype": [dtype],
|
||||
"momentum": 0.10000000000000001,
|
||||
"epsilon": 0.00029999999999999997,
|
||||
"num_inputs": 5,
|
||||
"num_outputs": 1,
|
||||
"scale": 1,
|
||||
"training": 1,
|
||||
"shape": [input_shape],
|
||||
},
|
||||
"inputs": [[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0], [4, 0, 0]],
|
||||
"name": "",
|
||||
"op": "kernel",
|
||||
},
|
||||
]
|
||||
verify(mod, clml_codegen, inputs_np, params_np)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"a_shape, b_shape, op",
|
||||
[
|
||||
((1, 64, 14, 14), (1, 64, 14, 14), R.add),
|
||||
((1, 256), (1, 256), R.add),
|
||||
((1, 64, 14, 14), (1, 64, 14, 14), R.subtract),
|
||||
((1, 256), (1, 256), R.subtract),
|
||||
((1, 64, 14, 14), (1, 64, 14, 14), R.multiply),
|
||||
((1, 256), (1, 256), R.multiply),
|
||||
((1, 64, 14, 14), (1, 64, 14, 14), R.divide),
|
||||
((1, 256), (1, 256), R.divide),
|
||||
((1, 64, 14, 14), (1, 64, 14, 14), R.minimum),
|
||||
((1, 256), (1, 256), R.minimum),
|
||||
((1, 64, 14, 14), (1, 64, 14, 14), R.maximum),
|
||||
((1, 256), (1, 256), R.maximum),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
def test_binary_ops(a_shape, b_shape, op, dtype):
|
||||
(mod, inputs_np) = get_binary_op_mod(a_shape, b_shape, op, dtype)
|
||||
clml_codegen = [
|
||||
{
|
||||
"attrs": {
|
||||
"dtype": [dtype],
|
||||
"shape": [a_shape],
|
||||
},
|
||||
"name": "",
|
||||
"op": "input",
|
||||
},
|
||||
{
|
||||
"attrs": {
|
||||
"dtype": [dtype],
|
||||
"shape": [b_shape],
|
||||
},
|
||||
"name": "",
|
||||
"op": "input",
|
||||
},
|
||||
{
|
||||
"attrs": {
|
||||
"dtype": [dtype],
|
||||
"num_inputs": 2,
|
||||
"num_outputs": 1,
|
||||
"shape": [a_shape],
|
||||
},
|
||||
"inputs": [[0, 0, 0], [1, 0, 0]],
|
||||
"name": "",
|
||||
"op": "kernel",
|
||||
},
|
||||
]
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
"float32",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"a_shape, op",
|
||||
[
|
||||
((1, 64, 14, 14), R.nn.relu),
|
||||
((1, 256, 1, 1), R.nn.relu),
|
||||
((1, 14, 256, 256), R.nn.relu),
|
||||
((1, 14, 14, 256), R.nn.relu),
|
||||
],
|
||||
)
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
def test_unary_ops(a_shape, op, dtype):
|
||||
(mod, inputs_np) = get_unary_op_mod(a_shape, op, dtype)
|
||||
clml_codegen = [
|
||||
{
|
||||
"attrs": {
|
||||
"dtype": [dtype],
|
||||
"shape": [a_shape],
|
||||
},
|
||||
"name": "",
|
||||
"op": "input",
|
||||
},
|
||||
{
|
||||
"attrs": {
|
||||
"activation_type": "relu",
|
||||
"dtype": [dtype],
|
||||
"num_inputs": 1,
|
||||
"num_outputs": 1,
|
||||
"shape": [a_shape],
|
||||
},
|
||||
"inputs": [[0, 0, 0]],
|
||||
"name": "",
|
||||
"op": "kernel",
|
||||
},
|
||||
]
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"trials",
|
||||
[
|
||||
[(1, 64, 147, 147), (3, 3), (2, 2), (1, 1), (0, 0, 0, 0), False],
|
||||
[(1, 256, 17, 17), (3, 3), (1, 1), (1, 1), (0, 0, 0, 0), False],
|
||||
[(1, 1024, 14, 14), (3, 3), (1, 1), (1, 1), (0, 0, 0, 0), False],
|
||||
# With padding is realized as nn.pad + pool
|
||||
# [(1, 32, 256, 256), (3, 3), (2, 2), (1, 1), (1, 1, 1, 1), True],
|
||||
# [(1, 32, 256, 256), (3, 3), (2, 2), (1, 1), (0, 1, 0, 1), True],
|
||||
# [(1, 32, 256, 256), (2, 2), (2, 2), (1, 1), (1, 1, 1, 1), True],
|
||||
# [(1, 32, 256, 256), (2, 2), (2, 2), (1, 1), (1, 0, 1, 0), True],
|
||||
],
|
||||
)
|
||||
def test_max_pool(dtype, trials):
|
||||
low, high = -1, 1
|
||||
(input_shape, pool_size, stride, dilation, padding, has_pad) = trials
|
||||
mod = get_relax_maxpool_mod(input_shape, dtype, pool_size, stride, dilation, padding, has_pad)
|
||||
clml_codegen = get_maxpool_expected_codegen(
|
||||
input_shape, pool_size, stride, padding, "maxpool2d", dtype
|
||||
)
|
||||
|
||||
inputs_np = [np.random.uniform(low, high, size=input_shape).astype(dtype)]
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"trials",
|
||||
[
|
||||
[(1, 64, 147, 147), (3, 3), (2, 2), (1, 1), (0, 0, 0, 0), False],
|
||||
[(1, 256, 17, 17), (3, 3), (1, 1), (1, 1), (0, 0, 0, 0), False],
|
||||
[(1, 1024, 14, 14), (3, 3), (1, 1), (1, 1), (0, 0, 0, 0), False],
|
||||
# With padding is realized as nn.pad + pool
|
||||
# [(1, 32, 256, 256), (3, 3), (2, 2), (1, 1), (1, 1, 1, 1), True],
|
||||
# [(1, 32, 256, 256), (3, 3), (2, 2), (1, 1), (0, 1, 0, 1), True],
|
||||
# [(1, 32, 256, 256), (2, 2), (2, 2), (1, 1), (1, 1, 1, 1), True],
|
||||
# [(1, 32, 256, 256), (2, 2), (2, 2), (1, 1), (1, 0, 1, 0), True],
|
||||
],
|
||||
)
|
||||
def test_avg_pool(dtype, trials):
|
||||
low, high = -1, 1
|
||||
(input_shape, pool_size, stride, dilation, padding, has_pad) = trials
|
||||
mod = get_relax_avgpool_mod(input_shape, dtype, pool_size, stride, dilation, padding, has_pad)
|
||||
clml_codegen = get_avgpool_expected_codegen(
|
||||
input_shape, pool_size, stride, padding, "avg_pool2d", dtype
|
||||
)
|
||||
|
||||
inputs_np = [np.random.uniform(low, high, size=input_shape).astype(dtype)]
|
||||
params_np = {}
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"trials",
|
||||
[
|
||||
[(1, 3, 32, 32), (1, 4, -1, 32)],
|
||||
[(1, 4, 8, 32), (1, 4, -1, 16)],
|
||||
[(1, 64, 3, 3), (1, 32, 3, -1)],
|
||||
],
|
||||
)
|
||||
def test_reshape(dtype, trials):
|
||||
low, high = -1, 1
|
||||
(input_shape, output_shape) = trials
|
||||
mod = get_relax_reshape_mod(input_shape, output_shape, dtype)
|
||||
clml_codegen = get_relax_reshape_codegen(input_shape, output_shape, dtype)
|
||||
|
||||
inputs_np = [np.random.uniform(low, high, size=input_shape).astype(dtype)]
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Codegen Comparision Failing")
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"trials",
|
||||
[
|
||||
[(1, 64, 147, 147), True],
|
||||
[(1, 256, 17, 17), False],
|
||||
[(1, 1024, 14, 14), True],
|
||||
[(1, 32, 256, 256), False],
|
||||
],
|
||||
)
|
||||
def test_global_avg_pool(dtype, trials):
|
||||
"""Test function for global average pooling."""
|
||||
low, high = -1, 1
|
||||
(input_shape, keep_dims) = trials
|
||||
N, C, H, W = input_shape
|
||||
pool_size, stride, padding = (H, W), (1, 1), (0, 0, 0, 0)
|
||||
mod = get_relax_global_avgpool_mod(input_shape, keep_dims, dtype)
|
||||
clml_codegen = get_global_maxpool_expected_codegen(
|
||||
input_shape, pool_size, stride, padding, "global_max", dtype
|
||||
)
|
||||
|
||||
inputs_np = [np.random.uniform(low, high, size=input_shape).astype(dtype)]
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_clml
|
||||
@pytest.mark.parametrize("dtype", ["float32"])
|
||||
@pytest.mark.parametrize(
|
||||
"trials",
|
||||
[
|
||||
[(1, 64, 147, 147), True],
|
||||
[(1, 256, 17, 17), False],
|
||||
[(1, 1024, 14, 14), True],
|
||||
[(1, 32, 256, 256), False],
|
||||
],
|
||||
)
|
||||
def test_global_max_pool(dtype, trials):
|
||||
"""Test function for global average pooling."""
|
||||
low, high = -1, 1
|
||||
(input_shape, keep_dims) = trials
|
||||
N, C, H, W = input_shape
|
||||
pool_size, stride, padding = (H, W), (1, 1), (0, 0, 0, 0)
|
||||
mod = get_relax_global_maxpool_mod(input_shape, keep_dims, dtype)
|
||||
clml_codegen = get_global_maxpool_expected_codegen(
|
||||
input_shape, pool_size, stride, padding, "global_max", dtype
|
||||
)
|
||||
|
||||
inputs_np = [np.random.uniform(low, high, size=input_shape).astype(dtype)]
|
||||
verify(mod, clml_codegen, inputs_np, {})
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
CLML_VERSION < 5,
|
||||
reason="Requires target device with CLML v5 or above",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"K, N, M",
|
||||
[
|
||||
(4096, 11008, 256),
|
||||
(2048, 32768, 128),
|
||||
(4096, 4096, 512),
|
||||
(4096, 22016, 64),
|
||||
(16384, 2048, 128),
|
||||
(2048, 2560, 1024),
|
||||
(3072, 9216, 256),
|
||||
(14336, 4096, 128),
|
||||
(1536, 17920, 128),
|
||||
(8960, 1536, 1024),
|
||||
],
|
||||
)
|
||||
def test_dequant_matmul(K, N, M):
|
||||
x_data = np.random.uniform(-0.1, 0.1, size=(1, M, K)).astype("float16")
|
||||
weight = np.random.randint(0, 100, size=(K // 8, N)).astype("uint32")
|
||||
scale = np.random.uniform(-0.1, 0.1, size=(K // 32, N)).astype("float16")
|
||||
|
||||
mod = get_dequant_matmul_module(K, N)
|
||||
clml_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [[[K // 8, N]]], "dtype": [["uint32"]]},
|
||||
},
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [[[K // 32, N]]], "dtype": [["float16"]]},
|
||||
},
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [[[1, -1, K]]], "dtype": [["float16"]]},
|
||||
},
|
||||
{
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [[0, 0, 0], [1, 0, 0], [2, 0, 0]],
|
||||
"attrs": {
|
||||
"dtype": ["float16"],
|
||||
"num_inputs": 3,
|
||||
"num_outputs": 1,
|
||||
"out_dtype": ["float16"],
|
||||
"shape": [[1, -1, N]],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
inputs_np = [x_data, weight, scale]
|
||||
verify(mod, clml_codegen, inputs_np, {}, target_minimum_clml_version=5)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
CLML_VERSION < 5,
|
||||
reason="Requires compiler supporting CLML v5 or above",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"K, N",
|
||||
[
|
||||
(4096, 11008),
|
||||
(2048, 32768),
|
||||
(4096, 4096),
|
||||
(4096, 22016),
|
||||
(16384, 2048),
|
||||
(2048, 2560),
|
||||
(3072, 9216),
|
||||
(4096, 28672),
|
||||
(14336, 4096),
|
||||
(1536, 17920),
|
||||
(8960, 1536),
|
||||
],
|
||||
)
|
||||
def test_dequant_vec_matmul(K, N):
|
||||
x_data = np.random.uniform(-0.1, 0.1, size=(1, 1, K)).astype("float16")
|
||||
weight = np.random.randint(0, 100, size=(K // 8, N)).astype("uint32")
|
||||
scale = np.random.uniform(-0.1, 0.1, size=(K // 32, N)).astype("float16")
|
||||
|
||||
mod = get_dequant_vec_matmul_module(K, N)
|
||||
clml_codegen = [
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [[[K // 8, -1]]], "dtype": [["uint32"]]},
|
||||
},
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [[[K // 32, -1]]], "dtype": [["float16"]]},
|
||||
},
|
||||
{
|
||||
"op": "input",
|
||||
"name": "",
|
||||
"attrs": {"shape": [[[1, 1, K]]], "dtype": [["float16"]]},
|
||||
},
|
||||
{
|
||||
"op": "kernel",
|
||||
"name": "",
|
||||
"inputs": [[0, 0, 0], [1, 0, 0], [2, 0, 0]],
|
||||
"attrs": {
|
||||
"dtype": ["float16"],
|
||||
"num_inputs": 3,
|
||||
"num_outputs": 1,
|
||||
"out_dtype": ["float16"],
|
||||
"shape": [[1, 1, -1]],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
inputs_np = (x_data, weight, scale)
|
||||
verify(mod, clml_codegen, inputs_np, {}, target_minimum_clml_version=5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,789 @@
|
||||
# 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: F401, F841
|
||||
|
||||
import copy
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("onnx")
|
||||
|
||||
import onnx
|
||||
from utils import verify_results
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.relax.frontend.onnx import from_onnx
|
||||
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import relax as relax_builder
|
||||
|
||||
TARGETS = [tvm.target.Target("qcom/adreno-opencl-texture")]
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_network_resnet():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Resnet:
|
||||
@R.function
|
||||
def main(
|
||||
data: R.Tensor((1, 3, 224, 224), dtype="float32"),
|
||||
resnetv22_batchnorm0_gamma: R.Tensor((3,), dtype="float32"),
|
||||
resnetv22_batchnorm0_beta: R.Tensor((3,), dtype="float32"),
|
||||
resnetv22_batchnorm0_running_mea: R.Tensor((3,), dtype="float32"),
|
||||
resnetv22_batchnorm0_running_var: R.Tensor((3,), dtype="float32"),
|
||||
resnetv22_conv0_weight: R.Tensor((64, 3, 7, 7), dtype="float32"),
|
||||
resnetv22_batchnorm1_gamma: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_batchnorm1_beta: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_batchnorm1_running_mea: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_batchnorm1_running_var: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm0_gamma: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm0_beta: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm0_running_mea: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm0_running_var: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_conv0_weight: R.Tensor((64, 64, 3, 3), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm1_gamma: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm1_beta: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm1_running_mea: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm1_running_var: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_conv1_weight: R.Tensor((64, 64, 3, 3), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm2_gamma: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm2_beta: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm2_running_mea: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm2_running_var: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_conv2_weight: R.Tensor((64, 64, 3, 3), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm3_gamma: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm3_beta: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm3_running_mea: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_batchnorm3_running_var: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage1_conv3_weight: R.Tensor((64, 64, 3, 3), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm0_gamma: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm0_beta: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm0_running_mea: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm0_running_var: R.Tensor((64,), dtype="float32"),
|
||||
resnetv22_stage2_conv0_weight: R.Tensor((128, 64, 3, 3), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm1_gamma: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm1_beta: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm1_running_mea: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm1_running_var: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_conv1_weight: R.Tensor((128, 128, 3, 3), dtype="float32"),
|
||||
resnetv22_stage2_conv2_weight: R.Tensor((128, 64, 1, 1), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm2_gamma: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm2_beta: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm2_running_mea: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm2_running_var: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_conv3_weight: R.Tensor((128, 128, 3, 3), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm3_gamma: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm3_beta: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm3_running_mea: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_batchnorm3_running_var: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage2_conv4_weight: R.Tensor((128, 128, 3, 3), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm0_gamma: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm0_beta: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm0_running_mea: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm0_running_var: R.Tensor((128,), dtype="float32"),
|
||||
resnetv22_stage3_conv0_weight: R.Tensor((256, 128, 3, 3), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm1_gamma: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm1_beta: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm1_running_mea: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm1_running_var: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_conv1_weight: R.Tensor((256, 256, 3, 3), dtype="float32"),
|
||||
resnetv22_stage3_conv2_weight: R.Tensor((256, 128, 1, 1), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm2_gamma: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm2_beta: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm2_running_mea: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm2_running_var: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_conv3_weight: R.Tensor((256, 256, 3, 3), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm3_gamma: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm3_beta: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm3_running_mea: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_batchnorm3_running_var: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage3_conv4_weight: R.Tensor((256, 256, 3, 3), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm0_gamma: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm0_beta: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm0_running_mea: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm0_running_var: R.Tensor((256,), dtype="float32"),
|
||||
resnetv22_stage4_conv0_weight: R.Tensor((512, 256, 3, 3), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm1_gamma: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm1_beta: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm1_running_mea: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm1_running_var: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_conv1_weight: R.Tensor((512, 512, 3, 3), dtype="float32"),
|
||||
resnetv22_stage4_conv2_weight: R.Tensor((512, 256, 1, 1), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm2_gamma: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm2_beta: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm2_running_mea: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm2_running_var: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_conv3_weight: R.Tensor((512, 512, 3, 3), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm3_gamma: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm3_beta: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm3_running_mea: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_batchnorm3_running_var: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_stage4_conv4_weight: R.Tensor((512, 512, 3, 3), dtype="float32"),
|
||||
resnetv22_batchnorm2_gamma: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_batchnorm2_beta: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_batchnorm2_running_mea: R.Tensor((512,), dtype="float32"),
|
||||
resnetv22_batchnorm2_running_var: R.Tensor((512,), dtype="float32"),
|
||||
reshape_attr_tensor164: R.Tensor((2,), dtype="int64"),
|
||||
resnetv22_dense0_weight: R.Tensor((1000, 512), dtype="float32"),
|
||||
resnetv22_dense0_bias: R.Tensor((1000,), dtype="float32"),
|
||||
) -> R.Tensor((1, 1000), dtype="float32"):
|
||||
with R.dataflow():
|
||||
lv: R.Tuple(
|
||||
R.Tensor((1, 3, 224, 224), dtype="float32"),
|
||||
R.Tensor((3,), dtype="float32"),
|
||||
R.Tensor((3,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
data,
|
||||
resnetv22_batchnorm0_gamma,
|
||||
resnetv22_batchnorm0_beta,
|
||||
resnetv22_batchnorm0_running_mea,
|
||||
resnetv22_batchnorm0_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv1: R.Tensor((1, 3, 224, 224), dtype="float32") = lv[0]
|
||||
lv2: R.Tensor((3,), dtype="float32") = lv[1]
|
||||
lv3: R.Tensor((3,), dtype="float32") = lv[2]
|
||||
lv4: R.Tensor((1, 64, 112, 112), dtype="float32") = R.nn.conv2d(
|
||||
lv1,
|
||||
resnetv22_conv0_weight,
|
||||
strides=[2, 2],
|
||||
padding=[3, 3, 3, 3],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv5: R.Tuple(
|
||||
R.Tensor((1, 64, 112, 112), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv4,
|
||||
resnetv22_batchnorm1_gamma,
|
||||
resnetv22_batchnorm1_beta,
|
||||
resnetv22_batchnorm1_running_mea,
|
||||
resnetv22_batchnorm1_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv6: R.Tensor((1, 64, 112, 112), dtype="float32") = lv5[0]
|
||||
lv7: R.Tensor((64,), dtype="float32") = lv5[1]
|
||||
lv8: R.Tensor((64,), dtype="float32") = lv5[2]
|
||||
lv9: R.Tensor((1, 64, 112, 112), dtype="float32") = R.nn.relu(lv6)
|
||||
lv10: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.max_pool2d(
|
||||
lv9,
|
||||
pool_size=[3, 3],
|
||||
strides=[2, 2],
|
||||
dilation=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
ceil_mode=False,
|
||||
count_include_pad=False,
|
||||
layout="NCHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv11: R.Tuple(
|
||||
R.Tensor((1, 64, 56, 56), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv10,
|
||||
resnetv22_stage1_batchnorm0_gamma,
|
||||
resnetv22_stage1_batchnorm0_beta,
|
||||
resnetv22_stage1_batchnorm0_running_mea,
|
||||
resnetv22_stage1_batchnorm0_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv12: R.Tensor((1, 64, 56, 56), dtype="float32") = lv11[0]
|
||||
lv13: R.Tensor((64,), dtype="float32") = lv11[1]
|
||||
lv14: R.Tensor((64,), dtype="float32") = lv11[2]
|
||||
lv15: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv12)
|
||||
lv16: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d(
|
||||
lv15,
|
||||
resnetv22_stage1_conv0_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv17: R.Tuple(
|
||||
R.Tensor((1, 64, 56, 56), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv16,
|
||||
resnetv22_stage1_batchnorm1_gamma,
|
||||
resnetv22_stage1_batchnorm1_beta,
|
||||
resnetv22_stage1_batchnorm1_running_mea,
|
||||
resnetv22_stage1_batchnorm1_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv18: R.Tensor((1, 64, 56, 56), dtype="float32") = lv17[0]
|
||||
lv19: R.Tensor((64,), dtype="float32") = lv17[1]
|
||||
lv20: R.Tensor((64,), dtype="float32") = lv17[2]
|
||||
lv21: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv18)
|
||||
lv22: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d(
|
||||
lv21,
|
||||
resnetv22_stage1_conv1_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv23: R.Tensor((1, 64, 56, 56), dtype="float32") = R.add(lv22, lv10)
|
||||
lv24: R.Tuple(
|
||||
R.Tensor((1, 64, 56, 56), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv23,
|
||||
resnetv22_stage1_batchnorm2_gamma,
|
||||
resnetv22_stage1_batchnorm2_beta,
|
||||
resnetv22_stage1_batchnorm2_running_mea,
|
||||
resnetv22_stage1_batchnorm2_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv25: R.Tensor((1, 64, 56, 56), dtype="float32") = lv24[0]
|
||||
lv26: R.Tensor((64,), dtype="float32") = lv24[1]
|
||||
lv27: R.Tensor((64,), dtype="float32") = lv24[2]
|
||||
lv28: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv25)
|
||||
lv29: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d(
|
||||
lv28,
|
||||
resnetv22_stage1_conv2_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv30: R.Tuple(
|
||||
R.Tensor((1, 64, 56, 56), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv29,
|
||||
resnetv22_stage1_batchnorm3_gamma,
|
||||
resnetv22_stage1_batchnorm3_beta,
|
||||
resnetv22_stage1_batchnorm3_running_mea,
|
||||
resnetv22_stage1_batchnorm3_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv31: R.Tensor((1, 64, 56, 56), dtype="float32") = lv30[0]
|
||||
lv32: R.Tensor((64,), dtype="float32") = lv30[1]
|
||||
lv33: R.Tensor((64,), dtype="float32") = lv30[2]
|
||||
lv34: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv31)
|
||||
lv35: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d(
|
||||
lv34,
|
||||
resnetv22_stage1_conv3_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv36: R.Tensor((1, 64, 56, 56), dtype="float32") = R.add(lv35, lv23)
|
||||
lv37: R.Tuple(
|
||||
R.Tensor((1, 64, 56, 56), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
R.Tensor((64,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv36,
|
||||
resnetv22_stage2_batchnorm0_gamma,
|
||||
resnetv22_stage2_batchnorm0_beta,
|
||||
resnetv22_stage2_batchnorm0_running_mea,
|
||||
resnetv22_stage2_batchnorm0_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv38: R.Tensor((1, 64, 56, 56), dtype="float32") = lv37[0]
|
||||
lv39: R.Tensor((64,), dtype="float32") = lv37[1]
|
||||
lv40: R.Tensor((64,), dtype="float32") = lv37[2]
|
||||
lv41: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv38)
|
||||
lv42: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.conv2d(
|
||||
lv41,
|
||||
resnetv22_stage2_conv0_weight,
|
||||
strides=[2, 2],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv43: R.Tuple(
|
||||
R.Tensor((1, 128, 28, 28), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv42,
|
||||
resnetv22_stage2_batchnorm1_gamma,
|
||||
resnetv22_stage2_batchnorm1_beta,
|
||||
resnetv22_stage2_batchnorm1_running_mea,
|
||||
resnetv22_stage2_batchnorm1_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv44: R.Tensor((1, 128, 28, 28), dtype="float32") = lv43[0]
|
||||
lv45: R.Tensor((128,), dtype="float32") = lv43[1]
|
||||
lv46: R.Tensor((128,), dtype="float32") = lv43[2]
|
||||
lv47: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.relu(lv44)
|
||||
lv48: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.conv2d(
|
||||
lv47,
|
||||
resnetv22_stage2_conv1_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv49: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.conv2d(
|
||||
lv41,
|
||||
resnetv22_stage2_conv2_weight,
|
||||
strides=[2, 2],
|
||||
padding=[0, 0, 0, 0],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv50: R.Tensor((1, 128, 28, 28), dtype="float32") = R.add(lv48, lv49)
|
||||
lv51: R.Tuple(
|
||||
R.Tensor((1, 128, 28, 28), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv50,
|
||||
resnetv22_stage2_batchnorm2_gamma,
|
||||
resnetv22_stage2_batchnorm2_beta,
|
||||
resnetv22_stage2_batchnorm2_running_mea,
|
||||
resnetv22_stage2_batchnorm2_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv52: R.Tensor((1, 128, 28, 28), dtype="float32") = lv51[0]
|
||||
lv53: R.Tensor((128,), dtype="float32") = lv51[1]
|
||||
lv54: R.Tensor((128,), dtype="float32") = lv51[2]
|
||||
lv55: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.relu(lv52)
|
||||
lv56: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.conv2d(
|
||||
lv55,
|
||||
resnetv22_stage2_conv3_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv57: R.Tuple(
|
||||
R.Tensor((1, 128, 28, 28), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv56,
|
||||
resnetv22_stage2_batchnorm3_gamma,
|
||||
resnetv22_stage2_batchnorm3_beta,
|
||||
resnetv22_stage2_batchnorm3_running_mea,
|
||||
resnetv22_stage2_batchnorm3_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv58: R.Tensor((1, 128, 28, 28), dtype="float32") = lv57[0]
|
||||
lv59: R.Tensor((128,), dtype="float32") = lv57[1]
|
||||
lv60: R.Tensor((128,), dtype="float32") = lv57[2]
|
||||
lv61: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.relu(lv58)
|
||||
lv62: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.conv2d(
|
||||
lv61,
|
||||
resnetv22_stage2_conv4_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv63: R.Tensor((1, 128, 28, 28), dtype="float32") = R.add(lv62, lv50)
|
||||
lv64: R.Tuple(
|
||||
R.Tensor((1, 128, 28, 28), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
R.Tensor((128,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv63,
|
||||
resnetv22_stage3_batchnorm0_gamma,
|
||||
resnetv22_stage3_batchnorm0_beta,
|
||||
resnetv22_stage3_batchnorm0_running_mea,
|
||||
resnetv22_stage3_batchnorm0_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv65: R.Tensor((1, 128, 28, 28), dtype="float32") = lv64[0]
|
||||
lv66: R.Tensor((128,), dtype="float32") = lv64[1]
|
||||
lv67: R.Tensor((128,), dtype="float32") = lv64[2]
|
||||
lv68: R.Tensor((1, 128, 28, 28), dtype="float32") = R.nn.relu(lv65)
|
||||
lv69: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.conv2d(
|
||||
lv68,
|
||||
resnetv22_stage3_conv0_weight,
|
||||
strides=[2, 2],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv70: R.Tuple(
|
||||
R.Tensor((1, 256, 14, 14), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv69,
|
||||
resnetv22_stage3_batchnorm1_gamma,
|
||||
resnetv22_stage3_batchnorm1_beta,
|
||||
resnetv22_stage3_batchnorm1_running_mea,
|
||||
resnetv22_stage3_batchnorm1_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv71: R.Tensor((1, 256, 14, 14), dtype="float32") = lv70[0]
|
||||
lv72: R.Tensor((256,), dtype="float32") = lv70[1]
|
||||
lv73: R.Tensor((256,), dtype="float32") = lv70[2]
|
||||
lv74: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.relu(lv71)
|
||||
lv75: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.conv2d(
|
||||
lv74,
|
||||
resnetv22_stage3_conv1_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv76: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.conv2d(
|
||||
lv68,
|
||||
resnetv22_stage3_conv2_weight,
|
||||
strides=[2, 2],
|
||||
padding=[0, 0, 0, 0],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv77: R.Tensor((1, 256, 14, 14), dtype="float32") = R.add(lv75, lv76)
|
||||
lv78: R.Tuple(
|
||||
R.Tensor((1, 256, 14, 14), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv77,
|
||||
resnetv22_stage3_batchnorm2_gamma,
|
||||
resnetv22_stage3_batchnorm2_beta,
|
||||
resnetv22_stage3_batchnorm2_running_mea,
|
||||
resnetv22_stage3_batchnorm2_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv79: R.Tensor((1, 256, 14, 14), dtype="float32") = lv78[0]
|
||||
lv80: R.Tensor((256,), dtype="float32") = lv78[1]
|
||||
lv81: R.Tensor((256,), dtype="float32") = lv78[2]
|
||||
lv82: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.relu(lv79)
|
||||
lv83: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.conv2d(
|
||||
lv82,
|
||||
resnetv22_stage3_conv3_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv84: R.Tuple(
|
||||
R.Tensor((1, 256, 14, 14), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv83,
|
||||
resnetv22_stage3_batchnorm3_gamma,
|
||||
resnetv22_stage3_batchnorm3_beta,
|
||||
resnetv22_stage3_batchnorm3_running_mea,
|
||||
resnetv22_stage3_batchnorm3_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv85: R.Tensor((1, 256, 14, 14), dtype="float32") = lv84[0]
|
||||
lv86: R.Tensor((256,), dtype="float32") = lv84[1]
|
||||
lv87: R.Tensor((256,), dtype="float32") = lv84[2]
|
||||
lv88: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.relu(lv85)
|
||||
lv89: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.conv2d(
|
||||
lv88,
|
||||
resnetv22_stage3_conv4_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv90: R.Tensor((1, 256, 14, 14), dtype="float32") = R.add(lv89, lv77)
|
||||
lv91: R.Tuple(
|
||||
R.Tensor((1, 256, 14, 14), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
R.Tensor((256,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv90,
|
||||
resnetv22_stage4_batchnorm0_gamma,
|
||||
resnetv22_stage4_batchnorm0_beta,
|
||||
resnetv22_stage4_batchnorm0_running_mea,
|
||||
resnetv22_stage4_batchnorm0_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv92: R.Tensor((1, 256, 14, 14), dtype="float32") = lv91[0]
|
||||
lv93: R.Tensor((256,), dtype="float32") = lv91[1]
|
||||
lv94: R.Tensor((256,), dtype="float32") = lv91[2]
|
||||
lv95: R.Tensor((1, 256, 14, 14), dtype="float32") = R.nn.relu(lv92)
|
||||
lv96: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.conv2d(
|
||||
lv95,
|
||||
resnetv22_stage4_conv0_weight,
|
||||
strides=[2, 2],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv97: R.Tuple(
|
||||
R.Tensor((1, 512, 7, 7), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv96,
|
||||
resnetv22_stage4_batchnorm1_gamma,
|
||||
resnetv22_stage4_batchnorm1_beta,
|
||||
resnetv22_stage4_batchnorm1_running_mea,
|
||||
resnetv22_stage4_batchnorm1_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv98: R.Tensor((1, 512, 7, 7), dtype="float32") = lv97[0]
|
||||
lv99: R.Tensor((512,), dtype="float32") = lv97[1]
|
||||
lv100: R.Tensor((512,), dtype="float32") = lv97[2]
|
||||
lv101: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.relu(lv98)
|
||||
lv102: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.conv2d(
|
||||
lv101,
|
||||
resnetv22_stage4_conv1_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv103: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.conv2d(
|
||||
lv95,
|
||||
resnetv22_stage4_conv2_weight,
|
||||
strides=[2, 2],
|
||||
padding=[0, 0, 0, 0],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv104: R.Tensor((1, 512, 7, 7), dtype="float32") = R.add(lv102, lv103)
|
||||
lv105: R.Tuple(
|
||||
R.Tensor((1, 512, 7, 7), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv104,
|
||||
resnetv22_stage4_batchnorm2_gamma,
|
||||
resnetv22_stage4_batchnorm2_beta,
|
||||
resnetv22_stage4_batchnorm2_running_mea,
|
||||
resnetv22_stage4_batchnorm2_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv106: R.Tensor((1, 512, 7, 7), dtype="float32") = lv105[0]
|
||||
lv107: R.Tensor((512,), dtype="float32") = lv105[1]
|
||||
lv108: R.Tensor((512,), dtype="float32") = lv105[2]
|
||||
lv109: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.relu(lv106)
|
||||
lv110: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.conv2d(
|
||||
lv109,
|
||||
resnetv22_stage4_conv3_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv111: R.Tuple(
|
||||
R.Tensor((1, 512, 7, 7), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv110,
|
||||
resnetv22_stage4_batchnorm3_gamma,
|
||||
resnetv22_stage4_batchnorm3_beta,
|
||||
resnetv22_stage4_batchnorm3_running_mea,
|
||||
resnetv22_stage4_batchnorm3_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv112: R.Tensor((1, 512, 7, 7), dtype="float32") = lv111[0]
|
||||
lv113: R.Tensor((512,), dtype="float32") = lv111[1]
|
||||
lv114: R.Tensor((512,), dtype="float32") = lv111[2]
|
||||
lv115: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.relu(lv112)
|
||||
lv116: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.conv2d(
|
||||
lv115,
|
||||
resnetv22_stage4_conv4_weight,
|
||||
strides=[1, 1],
|
||||
padding=[1, 1, 1, 1],
|
||||
dilation=[1, 1],
|
||||
groups=1,
|
||||
data_layout="NCHW",
|
||||
kernel_layout="OIHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
lv117: R.Tensor((1, 512, 7, 7), dtype="float32") = R.add(lv116, lv104)
|
||||
lv118: R.Tuple(
|
||||
R.Tensor((1, 512, 7, 7), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
R.Tensor((512,), dtype="float32"),
|
||||
) = R.nn.batch_norm(
|
||||
lv117,
|
||||
resnetv22_batchnorm2_gamma,
|
||||
resnetv22_batchnorm2_beta,
|
||||
resnetv22_batchnorm2_running_mea,
|
||||
resnetv22_batchnorm2_running_var,
|
||||
axis=1,
|
||||
epsilon=9.9999997473787516e-06,
|
||||
center=True,
|
||||
scale=True,
|
||||
momentum=0.10000000000000001,
|
||||
)
|
||||
lv119: R.Tensor((1, 512, 7, 7), dtype="float32") = lv118[0]
|
||||
lv120: R.Tensor((512,), dtype="float32") = lv118[1]
|
||||
lv121: R.Tensor((512,), dtype="float32") = lv118[2]
|
||||
lv122: R.Tensor((1, 512, 7, 7), dtype="float32") = R.nn.relu(lv119)
|
||||
lv123: R.Tensor((1, 512, 1, 1), dtype="float32") = R.mean(
|
||||
lv122, axis=[2, 3], keepdims=True
|
||||
)
|
||||
lv124: R.Tensor((1, 512), dtype="float32") = R.reshape(lv123, R.shape([1, 512]))
|
||||
lv125: R.Tensor((512, 1000), dtype="float32") = R.permute_dims(
|
||||
resnetv22_dense0_weight, axes=[1, 0]
|
||||
)
|
||||
lv126: R.Tensor((1, 1000), dtype="float32") = R.matmul(lv124, lv125)
|
||||
gv: R.Tensor((1, 1000), dtype="float32") = R.add(lv126, resnetv22_dense0_bias)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
verify_results(Resnet, target, tvm.target.Target("llvm"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,792 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import pytest
|
||||
from utils import skip_unless_adreno_opencl_vulkan, verify_results
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script.parser import ir as I
|
||||
from tvm.script.parser import relax as R
|
||||
|
||||
TARGETS = [
|
||||
tvm.target.Target("qcom/adreno-opencl-texture"),
|
||||
# tvm.target.Target("qcom/adreno-vulkan-texture"),
|
||||
]
|
||||
ref_target = tvm.target.Target("llvm")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 64, 56, 56), "float32"), w: R.Tensor((32, 64, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 32, 54, 54), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_relu():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_relu_conv2d_relu():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
x0: R.Tensor((2, 16, 28, 28), "float32") = R.nn.relu(x)
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x0, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_relu_tanh():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
||||
gv3: R.Tensor((2, 4, 26, 26), "float32") = R.tanh(gv2)
|
||||
R.output(gv3)
|
||||
return gv3
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_add():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"),
|
||||
w: R.Tensor((4, 16, 3, 3), "float32"),
|
||||
bias: R.Tensor((2, 4, 26, 26), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.add(gv, bias)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_sum():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=2):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4), "float32") = R.sum(gv, axis=[2, 3])
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_sum_keepdims():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=2):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 1, 1), "float32") = R.sum(gv, axis=[2, 3], keepdims=True)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_sum_reduce():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=2):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 26), "float32") = R.sum(gv, axis=[1, 2])
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_transpose():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((26, 26, 4, 2), "float32") = R.permute_dims(gv, axes=[3, 2, 1, 0])
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_expand_dims():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=6):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 1, 4, 1, 26, 26), "float32") = R.expand_dims(gv, axis=(-3, 1))
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_squeeze():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((1, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=3):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((1, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((4, 26, 26), "float32") = R.squeeze(gv, axis=[0])
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_strided_slice():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 2, 9, 7), dtype="float32") = R.strided_slice(
|
||||
gv, begin=[0, 0, 0], end=[4, 26, 26], strides=[2, 3, 4], axes=[1, 2, 3]
|
||||
)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_relu_concat():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
||||
gv3: R.Tensor((2, 8, 26, 26), "float32") = R.concat((gv, gv2), axis=1)
|
||||
R.output(gv3)
|
||||
return gv3
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_relu_concat_split():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
||||
gv3: R.Tensor((2, 8, 26, 26), "float32") = R.concat((gv, gv2), axis=1)
|
||||
gv4 = R.split(gv3, indices_or_sections=2, axis=1)
|
||||
# TODO @Siva: Multi value return have an issue at runtime.
|
||||
gv5 = gv4[0]
|
||||
R.output(gv5)
|
||||
return gv5
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_relu_concat_split_transpose_concat():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
||||
gv3: R.Tensor((2, 8, 26, 26), "float32") = R.concat((gv, gv2), axis=1)
|
||||
gv4 = R.split(gv3, indices_or_sections=2, axis=1)
|
||||
gv5: R.Tensor((26, 26, 4, 2), "float32") = R.permute_dims(gv4[0], axes=[3, 2, 1, 0])
|
||||
gv6: R.Tensor((26, 26, 4, 2), "float32") = R.permute_dims(gv4[1], axes=[3, 2, 1, 0])
|
||||
gv7: R.Tensor((26, 26, 8, 2), "float32") = R.concat((gv5, gv6), axis=2)
|
||||
R.output(gv7)
|
||||
return gv7
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Known failure: numerical mismatch in texture lowering")
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_maxpool2d():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2 = R.nn.max_pool2d(
|
||||
gv,
|
||||
pool_size=[2, 2],
|
||||
strides=[2, 2],
|
||||
padding=[0, 0],
|
||||
layout="NCHW",
|
||||
out_layout="NCHW",
|
||||
)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Known failure: numerical mismatch in texture lowering")
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_avgpool2d():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2 = R.nn.adaptive_avg_pool2d(gv, output_size=[13, 13], layout="NCHW")
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_softmax():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2 = R.nn.softmax(gv, axis=1)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_layernorm():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"),
|
||||
w: R.Tensor((4, 16, 3, 3), "float32"),
|
||||
gamma: R.Tensor((26, 26), dtype="float32"),
|
||||
beta: R.Tensor((26, 26), dtype="float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.layer_norm(
|
||||
gv, gamma, beta, axes=[-2, -1]
|
||||
)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_binary_broadcast():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"),
|
||||
w: R.Tensor((4, 16, 3, 3), "float32"),
|
||||
bias: R.Tensor((26, 26), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.add(gv, bias)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_binary_ewise_scalar():
|
||||
target = TARGETS[0]
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 16, 28, 28), "float32"), w: R.Tensor((4, 16, 3, 3), "float32")
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
||||
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.add(gv, R.const(1, "float32"))
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_residual_block():
|
||||
target = TARGETS[0]
|
||||
r"""
|
||||
- some kind of residual block followed by convolution to have texture after residual block
|
||||
- scalar data type verification which should be mapped to global memory scope
|
||||
layout_transform (NCHW->NCHW4c)
|
||||
| <- buffer
|
||||
conv2d (1) <- to get textures as output
|
||||
/ \
|
||||
conv2d (2) |
|
||||
\ /
|
||||
add <- add should be fused into conv2d (2)
|
||||
multiply to scalar <- buffer to the input of multiply scalar value
|
||||
relu
|
||||
| <- texture in intermediate tensor
|
||||
conv2d (3)
|
||||
relu
|
||||
| <- buffer
|
||||
layout_transform (NCHW4c->NCHW)
|
||||
"""
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 32, 40, 40), "float32"),
|
||||
w1: R.Tensor((32, 32, 2, 2), "float32"),
|
||||
w2: R.Tensor((32, 32, 1, 1), "float32"),
|
||||
w3: R.Tensor((32, 32, 2, 2), "float32"),
|
||||
bias: R.Tensor((1, 32, 1, 1), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv = R.nn.conv2d(x, w1, strides=[2, 2], out_dtype="float32")
|
||||
gv1 = R.add(gv, bias)
|
||||
gv2 = R.nn.relu(gv1)
|
||||
gv3 = R.nn.conv2d(gv2, w2, strides=[1, 1], out_dtype="float32")
|
||||
bias_1 = R.multiply(bias, R.const(0.15, "float32"))
|
||||
gv4 = R.add(gv3, bias_1)
|
||||
gv5 = R.nn.relu(gv4)
|
||||
gv6 = R.nn.conv2d(gv5, w3, strides=[2, 2], out_dtype="float32")
|
||||
gv7 = R.nn.relu(gv6)
|
||||
R.output(gv7)
|
||||
return gv7
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_conv2d_fallback_to_buffer_conv2d():
|
||||
target = TARGETS[0]
|
||||
r"""
|
||||
layout_transform (NCHW->NCHW4c)
|
||||
| <- texture
|
||||
conv2d (1) <- textures as output
|
||||
/ \
|
||||
conv2d (2) conv2d (3) <- conv2d (2) emits texture, conv2d (3) emits buffer
|
||||
\ / <- concat shouldn't support textures here
|
||||
concatenation
|
||||
| <- buffer
|
||||
layout_transform (NCHW4c->NCHW)
|
||||
"""
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 32, 40, 40), "float32"),
|
||||
w1: R.Tensor((96, 32, 2, 2), "float32"),
|
||||
w2: R.Tensor((32, 96, 2, 2), "float32"),
|
||||
w3: R.Tensor((5, 96, 2, 2), "float32"),
|
||||
bias1: R.Tensor((1, 96, 1, 1), "float32"),
|
||||
bias2: R.Tensor((1, 32, 1, 1), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv = R.nn.conv2d(x, w1, strides=[2, 2], out_dtype="float32")
|
||||
gv1 = R.add(gv, bias1)
|
||||
gv2 = R.nn.relu(gv1)
|
||||
gv3 = R.nn.conv2d(gv2, w2, strides=[2, 2], out_dtype="float32")
|
||||
gv6 = R.nn.conv2d(gv2, w3, strides=[2, 2], out_dtype="float32")
|
||||
gv7 = R.concat((gv3, gv6), axis=1)
|
||||
R.output(gv7)
|
||||
return gv7
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_conv2d_conv2d_conv2d_concat():
|
||||
target = TARGETS[0]
|
||||
r"""
|
||||
layout_transform (NCHW->NCHW4c)
|
||||
| <- texture
|
||||
conv2d (1) <- textures as output
|
||||
/ \
|
||||
conv2d (2) conv2d (3)
|
||||
\ / <- concat does support textures here
|
||||
concatenation
|
||||
| <- buffer
|
||||
layout_transform (NCHW4c->NCHW)
|
||||
"""
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 32, 40, 40), "float32"),
|
||||
w1: R.Tensor((96, 32, 2, 2), "float32"),
|
||||
w2: R.Tensor((32, 96, 2, 2), "float32"),
|
||||
w3: R.Tensor((8, 96, 2, 2), "float32"),
|
||||
bias1: R.Tensor((1, 96, 1, 1), "float32"),
|
||||
bias2: R.Tensor((1, 32, 1, 1), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv = R.nn.conv2d(x, w1, strides=[2, 2], out_dtype="float32")
|
||||
gv1 = R.add(gv, bias1)
|
||||
gv2 = R.nn.relu(gv1)
|
||||
gv3 = R.nn.conv2d(gv2, w2, strides=[2, 2], out_dtype="float32")
|
||||
gv6 = R.nn.conv2d(gv2, w3, strides=[2, 2], out_dtype="float32")
|
||||
gv7 = R.concat((gv3, gv6), axis=1)
|
||||
R.output(gv7)
|
||||
return gv7
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Known failure: numerical mismatch in texture lowering")
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_pooling_branching_texture_params():
|
||||
target = TARGETS[0]
|
||||
r"""
|
||||
Verification of the pooling and many branches having textures
|
||||
layout_transform (NCHW->NCHW4c)
|
||||
| <- texture
|
||||
conv2d (0) <- to get textures
|
||||
| <- textures
|
||||
pooling
|
||||
/ \ \ <- textures
|
||||
conv2d (1) conv2d (2) conv2d (3)
|
||||
\ / |
|
||||
add | <- to have the only one output, will be fused
|
||||
\ /
|
||||
add <- to have the only one output, will be fused
|
||||
| <- buffer
|
||||
layout_transform (NCHW4c->NCHW)
|
||||
"""
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 32, 40, 40), "float32"),
|
||||
w1: R.Tensor((32, 32, 1, 1), "float32"),
|
||||
w2: R.Tensor((32, 32, 2, 2), "float32"),
|
||||
w3: R.Tensor((32, 32, 1, 1), "float32"),
|
||||
w4: R.Tensor((32, 32, 2, 2), "float32"),
|
||||
bias1: R.Tensor((1, 32, 1, 1), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
gv = R.nn.conv2d(x, w1, strides=[1, 1], out_dtype="float32")
|
||||
gv1 = R.nn.max_pool2d(gv, pool_size=[2, 2], strides=[2, 2])
|
||||
gv2 = R.nn.conv2d(
|
||||
gv1, w2, padding=[0, 0, 1, 1], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
gv5 = R.nn.conv2d(
|
||||
gv1, w3, padding=[0, 0, 0, 0], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
gv6 = R.nn.conv2d(
|
||||
gv1, w4, padding=[0, 1, 1, 0], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
gv8 = R.add(gv2, gv5)
|
||||
gv9 = R.add(gv8, gv6)
|
||||
R.output(gv9)
|
||||
return gv9
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_injective_inputs1():
|
||||
target = TARGETS[0]
|
||||
r"""
|
||||
Input
|
||||
/ \
|
||||
/ |
|
||||
| /
|
||||
conv2d (1) /
|
||||
| /
|
||||
conv2d (2) mean
|
||||
/ \ /
|
||||
| | \ /
|
||||
| | (3) add
|
||||
| | |
|
||||
| \ /
|
||||
\ mul
|
||||
\ /
|
||||
add
|
||||
|
||||
"""
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((1, 4, 40, 40), "float32"),
|
||||
w1: R.Tensor((4, 4, 3, 3), "float32"),
|
||||
w2: R.Tensor((4, 4, 3, 3), "float32"),
|
||||
w3: R.Tensor((4, 4, 3, 3), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
mean = R.mean(x, axis=1, keepdims=True)
|
||||
conv1 = R.nn.conv2d(
|
||||
x, w1, padding=[1, 1, 1, 1], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
conv2 = R.nn.conv2d(
|
||||
conv1, w2, padding=[1, 1, 1, 1], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
ad3 = R.add(conv1, conv2)
|
||||
ad1 = R.add(mean, conv1)
|
||||
ad2 = R.multiply(ad1, conv2)
|
||||
gv = R.add(ad3, ad2)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@skip_unless_adreno_opencl_vulkan
|
||||
@pytest.mark.skipif(not tvm.testing.device_enabled(TARGETS[0]), reason="opencl not enabled")
|
||||
def test_injective_nwo_inputs2():
|
||||
target = TARGETS[0]
|
||||
r"""
|
||||
Input
|
||||
/ \
|
||||
| \
|
||||
conv2d \
|
||||
| /
|
||||
conv2d mean /
|
||||
/ \ /
|
||||
add | \ |
|
||||
| | \ |
|
||||
| | \ /
|
||||
| | (3) add
|
||||
| | |
|
||||
| \ /
|
||||
| \ /
|
||||
\ mul
|
||||
\ /
|
||||
add
|
||||
|
||||
"""
|
||||
|
||||
@I.ir_module
|
||||
class Input:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((1, 4, 40, 40), "float32"),
|
||||
w1: R.Tensor((4, 4, 3, 3), "float32"),
|
||||
w2: R.Tensor((4, 4, 3, 3), "float32"),
|
||||
w3: R.Tensor((4, 4, 3, 3), "float32"),
|
||||
) -> R.Tensor(None, "float32", ndim=4):
|
||||
with R.dataflow():
|
||||
mean = R.mean(x, axis=1, keepdims=True)
|
||||
conv1 = R.nn.conv2d(
|
||||
x, w1, padding=[1, 1, 1, 1], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
conv2 = R.nn.conv2d(
|
||||
conv1, w2, padding=[1, 1, 1, 1], strides=[1, 1], out_dtype="float32"
|
||||
)
|
||||
ad3 = R.add(conv1, conv2)
|
||||
ad1 = R.add(mean, conv1)
|
||||
ad2 = R.multiply(ad1, conv2)
|
||||
gv = R.add(ad2, ad3)
|
||||
R.output(gv)
|
||||
return gv
|
||||
|
||||
verify_results(Input, target, ref_target)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,283 @@
|
||||
# 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: F401
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.ir.module import IRModule
|
||||
from tvm.script.parser import ir as I
|
||||
from tvm.script.parser import relax as R
|
||||
from tvm.script.parser import tirx as T
|
||||
|
||||
|
||||
def verify(input, expected):
|
||||
mod = tvm.relax.backend.adreno.transform.FoldVDeviceScopeChange()(input)
|
||||
tvm.ir.assert_structural_equal(mod, expected)
|
||||
|
||||
|
||||
def test_maxpool2d_scope_folding():
|
||||
@I.ir_module(s_tir=True)
|
||||
class Input:
|
||||
I.module_global_infos(
|
||||
{
|
||||
"vdevice": [
|
||||
I.vdevice({"device": "adreno", "kind": "opencl"}, 0, "global.texture-weight"),
|
||||
I.vdevice({"device": "adreno", "kind": "opencl"}, 0, "global"),
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def max_pool2d_opencl(
|
||||
gv: T.Buffer((T.int64(2), T.int64(1), T.int64(26), T.int64(26), T.int64(4)), "float32"),
|
||||
pool_max: T.Buffer(
|
||||
(T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4)), "float32"
|
||||
),
|
||||
):
|
||||
# with T.sblock("root"):
|
||||
for ax0, ax1, ax2, ax3, ax4, rv0, rv1 in T.grid(
|
||||
T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4), T.int64(2), T.int64(2)
|
||||
):
|
||||
with T.sblock("pool_max"):
|
||||
v_ax0, v_ax1, v_ax2, v_ax3, v_ax4, v_rv0, v_rv1 = T.axis.remap(
|
||||
"SSSSSRR", [ax0, ax1, ax2, ax3, ax4, rv0, rv1]
|
||||
)
|
||||
T.reads(
|
||||
gv[
|
||||
v_ax0,
|
||||
v_ax1,
|
||||
v_ax2 * T.int64(2) + v_rv0,
|
||||
v_ax3 * T.int64(2) + v_rv1,
|
||||
v_ax4,
|
||||
]
|
||||
)
|
||||
T.writes(pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4])
|
||||
T.sblock_attr({"schedule_rule": "meta_schedule.pool_max"})
|
||||
with T.init():
|
||||
pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = T.float32(
|
||||
-340282346638528859811704183484516925440.0
|
||||
)
|
||||
pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = T.max(
|
||||
pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4],
|
||||
gv[
|
||||
v_ax0,
|
||||
v_ax1,
|
||||
v_ax2 * T.int64(2) + v_rv0,
|
||||
v_ax3 * T.int64(2) + v_rv1,
|
||||
v_ax4,
|
||||
],
|
||||
)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def te_layout_transform(
|
||||
x: T.Buffer((T.int64(2), T.int64(4), T.int64(26), T.int64(26)), "float32"),
|
||||
te_layout_transform: T.Buffer(
|
||||
(T.int64(2), T.int64(1), T.int64(26), T.int64(26), T.int64(4)), "float32"
|
||||
),
|
||||
):
|
||||
# with T.sblock("root"):
|
||||
for self, i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(26), T.int64(26)):
|
||||
with T.sblock("te_layout_transform"):
|
||||
v_self, v_i0, v_i1, v_i2 = T.axis.remap("SSSS", [self, i0, i1, i2])
|
||||
T.reads(x[v_self, v_i0, v_i1, v_i2])
|
||||
T.writes(
|
||||
te_layout_transform[
|
||||
v_self, v_i0 // T.int64(4), v_i1, v_i2, v_i0 % T.int64(4)
|
||||
]
|
||||
)
|
||||
te_layout_transform[
|
||||
v_self, v_i0 // T.int64(4), v_i1, v_i2, v_i0 % T.int64(4)
|
||||
] = x[v_self, v_i0, v_i1, v_i2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def te_layout_transform2(
|
||||
lv2: T.Buffer(
|
||||
(T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4)), "float32"
|
||||
),
|
||||
te_layout_transform: T.Buffer(
|
||||
(T.int64(2), T.int64(4), T.int64(13), T.int64(13)), "float32"
|
||||
),
|
||||
):
|
||||
# with T.sblock("root"):
|
||||
for self, i0, i1, i2, i3 in T.grid(
|
||||
T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4)
|
||||
):
|
||||
with T.sblock("te_layout_transform"):
|
||||
v_self, v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSSS", [self, i0, i1, i2, i3])
|
||||
T.reads(lv2[v_self, v_i0, v_i1, v_i2, v_i3])
|
||||
T.writes(te_layout_transform[v_self, v_i3, v_i1, v_i2])
|
||||
te_layout_transform[v_self, v_i3, v_i1, v_i2] = lv2[
|
||||
v_self, v_i0, v_i1, v_i2, v_i3
|
||||
]
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 4, 26, 26), dtype="float32", vdevice="opencl:1:global"),
|
||||
) -> R.Tensor((2, 4, 13, 13), dtype="float32", vdevice="opencl:1:global"):
|
||||
cls = Input
|
||||
with R.dataflow():
|
||||
lv = R.call_tir(
|
||||
cls.te_layout_transform,
|
||||
(x,),
|
||||
out_ty=R.Tensor(
|
||||
(2, 1, 26, 26, 4), dtype="float32", vdevice="opencl:0:global.texture-weight"
|
||||
),
|
||||
)
|
||||
lv2 = R.call_tir(
|
||||
cls.max_pool2d_opencl,
|
||||
(lv,),
|
||||
out_ty=R.Tensor(
|
||||
(2, 1, 13, 13, 4), dtype="float32", vdevice="opencl:0:global.texture-weight"
|
||||
),
|
||||
)
|
||||
lv5: R.Tensor((2, 1, 13, 13, 4), dtype="float32", vdevice="opencl:1:global") = (
|
||||
R.to_vdevice(lv2, dst_vdevice="opencl:1:global")
|
||||
)
|
||||
gv2 = R.call_tir(
|
||||
cls.te_layout_transform2,
|
||||
(lv5,),
|
||||
out_ty=R.Tensor((2, 4, 13, 13), dtype="float32", vdevice="opencl:1:global"),
|
||||
)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class Expected:
|
||||
I.module_global_infos(
|
||||
{
|
||||
"vdevice": [
|
||||
I.vdevice({"device": "adreno", "kind": "opencl"}, 0, "global.texture-weight"),
|
||||
I.vdevice({"device": "adreno", "kind": "opencl"}, 0, "global"),
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def max_pool2d_opencl(
|
||||
gv: T.Buffer((T.int64(2), T.int64(1), T.int64(26), T.int64(26), T.int64(4)), "float32"),
|
||||
pool_max: T.Buffer(
|
||||
(T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4)), "float32"
|
||||
),
|
||||
):
|
||||
# with T.sblock("root"):
|
||||
for ax0, ax1, ax2, ax3, ax4, rv0, rv1 in T.grid(
|
||||
T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4), T.int64(2), T.int64(2)
|
||||
):
|
||||
with T.sblock("pool_max"):
|
||||
v_ax0, v_ax1, v_ax2, v_ax3, v_ax4, v_rv0, v_rv1 = T.axis.remap(
|
||||
"SSSSSRR", [ax0, ax1, ax2, ax3, ax4, rv0, rv1]
|
||||
)
|
||||
T.reads(
|
||||
gv[
|
||||
v_ax0,
|
||||
v_ax1,
|
||||
v_ax2 * T.int64(2) + v_rv0,
|
||||
v_ax3 * T.int64(2) + v_rv1,
|
||||
v_ax4,
|
||||
]
|
||||
)
|
||||
T.writes(pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4])
|
||||
T.sblock_attr({"schedule_rule": "meta_schedule.pool_max"})
|
||||
with T.init():
|
||||
pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = T.float32(
|
||||
-340282346638528859811704183484516925440.0
|
||||
)
|
||||
pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = T.max(
|
||||
pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4],
|
||||
gv[
|
||||
v_ax0,
|
||||
v_ax1,
|
||||
v_ax2 * T.int64(2) + v_rv0,
|
||||
v_ax3 * T.int64(2) + v_rv1,
|
||||
v_ax4,
|
||||
],
|
||||
)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def te_layout_transform(
|
||||
x: T.Buffer((T.int64(2), T.int64(4), T.int64(26), T.int64(26)), "float32"),
|
||||
te_layout_transform: T.Buffer(
|
||||
(T.int64(2), T.int64(1), T.int64(26), T.int64(26), T.int64(4)), "float32"
|
||||
),
|
||||
):
|
||||
# with T.sblock("root"):
|
||||
for self, i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(26), T.int64(26)):
|
||||
with T.sblock("te_layout_transform"):
|
||||
v_self, v_i0, v_i1, v_i2 = T.axis.remap("SSSS", [self, i0, i1, i2])
|
||||
T.reads(x[v_self, v_i0, v_i1, v_i2])
|
||||
T.writes(
|
||||
te_layout_transform[
|
||||
v_self, v_i0 // T.int64(4), v_i1, v_i2, v_i0 % T.int64(4)
|
||||
]
|
||||
)
|
||||
te_layout_transform[
|
||||
v_self, v_i0 // T.int64(4), v_i1, v_i2, v_i0 % T.int64(4)
|
||||
] = x[v_self, v_i0, v_i1, v_i2]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def te_layout_transform2(
|
||||
lv2: T.Buffer(
|
||||
(T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4)), "float32"
|
||||
),
|
||||
te_layout_transform: T.Buffer(
|
||||
(T.int64(2), T.int64(4), T.int64(13), T.int64(13)), "float32"
|
||||
),
|
||||
):
|
||||
# with T.sblock("root"):
|
||||
for self, i0, i1, i2, i3 in T.grid(
|
||||
T.int64(2), T.int64(1), T.int64(13), T.int64(13), T.int64(4)
|
||||
):
|
||||
with T.sblock("te_layout_transform"):
|
||||
v_self, v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSSS", [self, i0, i1, i2, i3])
|
||||
T.reads(lv2[v_self, v_i0, v_i1, v_i2, v_i3])
|
||||
T.writes(te_layout_transform[v_self, v_i3, v_i1, v_i2])
|
||||
te_layout_transform[v_self, v_i3, v_i1, v_i2] = lv2[
|
||||
v_self, v_i0, v_i1, v_i2, v_i3
|
||||
]
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((2, 4, 26, 26), dtype="float32", vdevice="opencl:1:global"),
|
||||
) -> R.Tensor((2, 4, 13, 13), dtype="float32", vdevice="opencl:1:global"):
|
||||
cls = Expected
|
||||
with R.dataflow():
|
||||
lv = R.call_tir(
|
||||
cls.te_layout_transform,
|
||||
(x,),
|
||||
out_ty=R.Tensor(
|
||||
(2, 1, 26, 26, 4), dtype="float32", vdevice="opencl:0:global.texture-weight"
|
||||
),
|
||||
)
|
||||
lv5 = R.call_tir(
|
||||
cls.max_pool2d_opencl,
|
||||
(lv,),
|
||||
out_ty=R.Tensor((2, 1, 13, 13, 4), dtype="float32", vdevice="opencl:1:global"),
|
||||
)
|
||||
gv2 = R.call_tir(
|
||||
cls.te_layout_transform2,
|
||||
(lv5,),
|
||||
out_ty=R.Tensor((2, 4, 13, 13), dtype="float32", vdevice="opencl:1:global"),
|
||||
)
|
||||
R.output(gv2)
|
||||
return gv2
|
||||
|
||||
verify(Input, Expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,218 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.support import ndk
|
||||
|
||||
# Test Infra
|
||||
|
||||
|
||||
class run_time_check:
|
||||
def __init__(self, device):
|
||||
self.device = device
|
||||
|
||||
def check(self):
|
||||
# Ensure adreno specific tests
|
||||
if self.device == "real":
|
||||
return "ADRENO_TARGET" in os.environ
|
||||
|
||||
# Adreno CI
|
||||
if "ADRENO_TARGET" in os.environ:
|
||||
return True
|
||||
|
||||
# Tests that can run on generic targets too
|
||||
elif self.device == "opencl":
|
||||
return tvm.opencl().exist
|
||||
elif self.device == "vulkan":
|
||||
return tvm.vulkan().exist
|
||||
elif self.device == "any":
|
||||
return tvm.opencl().exist or tvm.vulkan().exist
|
||||
else:
|
||||
return False
|
||||
|
||||
def __call__(self):
|
||||
return self.check
|
||||
|
||||
|
||||
# Eager skips for Adreno GPU tests, resolved at import time. Pair each with
|
||||
# ``@pytest.mark.gpu`` at the test site so CI's ``-m gpu`` filter selects it.
|
||||
|
||||
# OpenCL or Vulkan
|
||||
skip_unless_adreno_opencl_vulkan = pytest.mark.skipif(
|
||||
not run_time_check("any").check(),
|
||||
reason="need adreno opencl or vulkan",
|
||||
)
|
||||
|
||||
# CLML Codegen
|
||||
skip_unless_adreno_clml = pytest.mark.skipif(
|
||||
tvm.get_global_func("relax.is_openclml_runtime_enabled", allow_missing=True) is None,
|
||||
reason="need adreno openclml",
|
||||
)
|
||||
|
||||
|
||||
def is_target_available(target):
|
||||
if "clml" in target.attrs.get("keys", []) and "ADRENO_TARGET" not in os.environ:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class SessionManager:
|
||||
def __init__(self):
|
||||
self.is_remote = SessionManager.is_target_rpc()
|
||||
|
||||
def __enter__(self):
|
||||
if self.is_remote:
|
||||
self.RPC_TRACKER_HOST = os.getenv("TVM_TRACKER_HOST", "localhost")
|
||||
self.RPC_TRACKER_PORT = int(os.getenv("TVM_TRACKER_PORT", 7979))
|
||||
self.RPC_DEVICE_KEY = os.getenv("RPC_DEVICE_KEY", "android")
|
||||
|
||||
self.tracker = tvm.rpc.connect_tracker(self.RPC_TRACKER_HOST, self.RPC_TRACKER_PORT)
|
||||
self.rpc = self.tracker.request(self.RPC_DEVICE_KEY, priority=0, session_timeout=600)
|
||||
else:
|
||||
self.rpc = tvm.rpc.LocalSession()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.rpc.get_function("CloseRPCConnection")()
|
||||
|
||||
def load_module(self, ex: relax.VMExecutable):
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
file_name = "vm_library.so"
|
||||
file_path = os.path.join(tempdir, file_name)
|
||||
if self.is_remote:
|
||||
ex.export_library(
|
||||
file_path, fcompile=ndk.create_shared, options=["-shared", "-fPIC", "-lm"]
|
||||
)
|
||||
else:
|
||||
ex.export_library(file_path)
|
||||
|
||||
self.rpc.upload(file_path)
|
||||
rexec = self.rpc.load_module(file_name)
|
||||
return rexec
|
||||
|
||||
def device(self, device: str):
|
||||
return self.rpc.device(device)
|
||||
|
||||
@staticmethod
|
||||
def is_target_rpc():
|
||||
"""
|
||||
Checks if the target is a remote device.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool: True if RPC_TARGET is set, False otherwise
|
||||
"""
|
||||
return os.environ.get("ADRENO_TARGET") == "adreno"
|
||||
|
||||
|
||||
def run_local(mod, inputs, target):
|
||||
"""
|
||||
Run the Relax module on the local CPU for verification.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : tvm.IRModule
|
||||
The Relax IRModule to execute.
|
||||
inputs : list of numpy.ndarray
|
||||
The input data for the module.
|
||||
save_lib : bool, optional
|
||||
Whether to save the compiled library. Default is False.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tvm.runtime.NDArray or tuple of tvm.runtime.NDArray
|
||||
The output from the module execution.
|
||||
"""
|
||||
ex = relax.build(mod, target)
|
||||
dev = tvm.cpu()
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs]
|
||||
vm.set_input("main", *inputs)
|
||||
vm.invoke_stateful("main")
|
||||
tvm_output = vm.get_outputs("main")
|
||||
if isinstance(tvm_output, tuple):
|
||||
tvm_output = tuple(out.numpy() for out in tvm_output)
|
||||
else:
|
||||
tvm_output = (tvm_output.numpy(),)
|
||||
return tvm_output
|
||||
|
||||
|
||||
def build_and_run(mod, inputs, tgt):
|
||||
if SessionManager.is_target_rpc():
|
||||
tgt = tvm.target.Target(tgt, host={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
|
||||
else:
|
||||
tgt = tvm.target.Target(tgt, host={"kind": "llvm"})
|
||||
|
||||
relax_pipeline = relax.pipeline.get_default_pipeline(tgt)
|
||||
tir_pipeline = tvm.tirx.get_default_tir_pipeline(tgt)
|
||||
mod = relax_pipeline(mod)
|
||||
|
||||
ex = tvm.compile(mod, tgt, tir_pipeline=tir_pipeline)
|
||||
|
||||
def run_and_check():
|
||||
with SessionManager() as sess:
|
||||
rexec = sess.load_module(ex)
|
||||
dev = sess.device(tgt.kind.name)
|
||||
|
||||
if "vdevice" in mod.global_infos:
|
||||
device_arr = [dev for _ in range(len(mod.global_infos["vdevice"]))]
|
||||
else:
|
||||
device_arr = [dev]
|
||||
vm = relax.VirtualMachine(rexec, device_arr)
|
||||
device_inputs = [tvm.runtime.tensor(ip, dev) for ip in inputs]
|
||||
vm.set_input("main", *device_inputs)
|
||||
vm.invoke_stateful("main")
|
||||
tvm_output = vm.get_outputs("main")
|
||||
if isinstance(tvm_output, tuple):
|
||||
return tuple(out.numpy() for out in tvm_output)
|
||||
return (tvm_output.numpy(),)
|
||||
|
||||
if SessionManager.is_target_rpc():
|
||||
return run_and_check()
|
||||
return tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
def verify_results(mod, target, ref_target):
|
||||
if not is_target_available(target):
|
||||
print("Skipping Eval Tests", flush=True)
|
||||
return
|
||||
|
||||
inputs = []
|
||||
for arg in mod["main"].params:
|
||||
shape = tuple(shape_val.value for shape_val in arg.ty.shape.values)
|
||||
inputs.append(np.random.uniform(0, 1, size=shape).astype(arg.ty.dtype))
|
||||
|
||||
mod_org, mod_ref = mod, mod.clone()
|
||||
|
||||
mod_ref = tvm.relax.transform.DecomposeOpsForInference()(mod_ref)
|
||||
if ref_target.kind.name == "llvm":
|
||||
rs_ref = run_local(mod_ref, inputs, ref_target)
|
||||
else:
|
||||
rs_ref = build_and_run(mod_ref, inputs, ref_target)
|
||||
|
||||
rs_org = build_and_run(mod_org, inputs, target)
|
||||
|
||||
for vl_org, vl_ref in zip(rs_org, rs_ref):
|
||||
tvm.testing.assert_allclose(vl_org, vl_ref, rtol=1e-3, atol=1e-3)
|
||||
Reference in New Issue
Block a user