wehub-resource-sync
843 lines
27 KiB
Python
843 lines
27 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, F401, F841
|
|
"""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
|