chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,301 @@
|
||||
# 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.
|
||||
"""Configure pytest"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
import tvm.topi.testing
|
||||
from tvm import te
|
||||
from tvm.contrib import cblas, dnnl, mkl
|
||||
|
||||
|
||||
def verify_matmul_add(
|
||||
matrix_m, matrix_l, matrix_n, lib, transa=False, transb=False, dtype="float32"
|
||||
):
|
||||
"""Tests matmul+add op"""
|
||||
bias = te.var("bias", dtype=dtype)
|
||||
ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
|
||||
bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
|
||||
input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
|
||||
input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
|
||||
matmul_result = lib.matmul(input1_data, input2_data, transa, transb)
|
||||
final_result = te.compute(
|
||||
matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
|
||||
)
|
||||
|
||||
def get_numpy(a, b, matrix_bias, transa, transb):
|
||||
if transa:
|
||||
a = a.transpose()
|
||||
if transb:
|
||||
b = b.transpose()
|
||||
return np.dot(a, b) + matrix_bias
|
||||
|
||||
def compiling(f, name="test_matmul_add", ext=".so"):
|
||||
path = name + ext
|
||||
f.export_library(path)
|
||||
mod = tvm.runtime.load_module(path)
|
||||
f = mod[name]
|
||||
return f
|
||||
|
||||
def verify(target="llvm"):
|
||||
if not tvm.testing.device_enabled(target):
|
||||
print(f"skip because {target} is not enabled...")
|
||||
return
|
||||
if not tvm.get_global_func(lib.__name__ + ".matmul", True):
|
||||
print("skip because extern function is not available")
|
||||
return
|
||||
dev = tvm.cpu(0)
|
||||
name = "test_matmul_add"
|
||||
f = tvm.compile(
|
||||
te.create_prim_func([input1_data, input2_data, final_result, bias]).with_attr(
|
||||
"global_symbol", name
|
||||
),
|
||||
target=target,
|
||||
)
|
||||
if target == "c":
|
||||
f = compiling(f, name)
|
||||
matrix_input1 = tvm.runtime.tensor(
|
||||
np.random.uniform(size=ashape).astype(input1_data.dtype), dev
|
||||
)
|
||||
matrix_input2 = tvm.runtime.tensor(
|
||||
np.random.uniform(size=bshape).astype(input2_data.dtype), dev
|
||||
)
|
||||
matrix_result = tvm.runtime.tensor(
|
||||
np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
|
||||
)
|
||||
matrix_bias = 10.0
|
||||
f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
|
||||
tvm.testing.assert_allclose(
|
||||
matrix_result.numpy(),
|
||||
get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), matrix_bias, transa, transb),
|
||||
rtol=1e-5,
|
||||
)
|
||||
|
||||
verify("llvm")
|
||||
verify("c")
|
||||
|
||||
|
||||
def test_matmul_add():
|
||||
"""Tests of matmul+add op"""
|
||||
verify_matmul_add(235, 128, 1024, cblas)
|
||||
verify_matmul_add(235, 128, 1024, cblas, True, False)
|
||||
verify_matmul_add(235, 128, 1024, cblas, False, True)
|
||||
verify_matmul_add(235, 128, 1024, cblas, True, True)
|
||||
verify_matmul_add(235, 128, 1024, mkl)
|
||||
verify_matmul_add(235, 128, 1024, mkl, True, False)
|
||||
verify_matmul_add(235, 128, 1024, mkl, False, True)
|
||||
verify_matmul_add(235, 128, 1024, mkl, True, True)
|
||||
verify_matmul_add(235, 128, 1024, dnnl)
|
||||
verify_matmul_add(235, 128, 1024, dnnl, True, False)
|
||||
verify_matmul_add(235, 128, 1024, dnnl, False, True)
|
||||
verify_matmul_add(235, 128, 1024, dnnl, True, True)
|
||||
verify_matmul_add(1, 16, 4, cblas)
|
||||
verify_matmul_add(1, 16, 3, cblas, True, False)
|
||||
verify_matmul_add(1, 16, 3, cblas, False, False)
|
||||
verify_matmul_add(1, 16, 3, cblas, True, True)
|
||||
verify_matmul_add(1, 16, 4, mkl)
|
||||
verify_matmul_add(1, 16, 3, mkl, True, False)
|
||||
verify_matmul_add(1, 16, 3, mkl, False, False)
|
||||
verify_matmul_add(1, 16, 3, mkl, True, True)
|
||||
verify_matmul_add(1, 16, 4, dnnl)
|
||||
verify_matmul_add(1, 16, 3, dnnl, True, False)
|
||||
verify_matmul_add(1, 16, 3, dnnl, False, False)
|
||||
verify_matmul_add(1, 16, 3, dnnl, True, True)
|
||||
|
||||
|
||||
def verify_quantized_matmul_add(matrix_m, matrix_l, matrix_n, transa=False, transb=False):
|
||||
"""Tests quantized matmul+add op"""
|
||||
if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
|
||||
pytest.skip("Quantized dense is supported only for MKL. TVM GPU CI uses openblas")
|
||||
data_dtype = "uint8"
|
||||
kernel_dtype = "int8"
|
||||
out_dtype = "int32"
|
||||
bias = te.var("bias", dtype=out_dtype)
|
||||
ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
|
||||
bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
|
||||
input1_data = te.placeholder(ashape, name="input1_data", dtype=data_dtype)
|
||||
input2_data = te.placeholder(bshape, name="input2_data", dtype=kernel_dtype)
|
||||
matmul_result = mkl.matmul_u8s8s32(input1_data, input2_data, transa, transb, dtype=out_dtype)
|
||||
final_result = te.compute(
|
||||
matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
|
||||
)
|
||||
|
||||
def get_numpy(a, b, matrix_bias, transa, transb):
|
||||
if transa:
|
||||
a = a.transpose()
|
||||
if transb:
|
||||
b = b.transpose()
|
||||
return np.dot(a, b) + matrix_bias
|
||||
|
||||
def verify(target="llvm"):
|
||||
if not tvm.testing.device_enabled(target):
|
||||
print(f"skip because {target} is not enabled...")
|
||||
return
|
||||
if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
|
||||
print("skip because extern function is not available")
|
||||
return
|
||||
dev = tvm.cpu(0)
|
||||
f = tvm.compile(
|
||||
te.create_prim_func([input1_data, input2_data, final_result, bias]), target=target
|
||||
)
|
||||
matrix_input1 = tvm.runtime.tensor(
|
||||
np.random.randint(low=0, high=50, size=ashape).astype(input1_data.dtype), dev
|
||||
)
|
||||
matrix_input2 = tvm.runtime.tensor(
|
||||
np.random.randint(low=0, high=50, size=bshape).astype(input2_data.dtype), dev
|
||||
)
|
||||
matrix_result = tvm.runtime.tensor(
|
||||
np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
|
||||
)
|
||||
matrix_bias = 10
|
||||
f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
|
||||
tvm.testing.assert_allclose(
|
||||
matrix_result.numpy(),
|
||||
get_numpy(
|
||||
matrix_input1.numpy().astype("int32"),
|
||||
matrix_input2.numpy().astype("int32"),
|
||||
matrix_bias,
|
||||
transa,
|
||||
transb,
|
||||
),
|
||||
rtol=1e-5,
|
||||
)
|
||||
|
||||
verify()
|
||||
|
||||
|
||||
def test_quantized_matmul_add():
|
||||
"""Tests of quantized matmul+add op"""
|
||||
verify_quantized_matmul_add(235, 128, 1024)
|
||||
verify_quantized_matmul_add(235, 128, 1024, True, False)
|
||||
verify_quantized_matmul_add(235, 128, 1024, False, True)
|
||||
verify_quantized_matmul_add(235, 128, 1024, True, True)
|
||||
verify_quantized_matmul_add(1, 16, 4)
|
||||
verify_quantized_matmul_add(1, 16, 3, True, False)
|
||||
verify_quantized_matmul_add(1, 16, 3, False, True)
|
||||
verify_quantized_matmul_add(1, 16, 3, True, True)
|
||||
|
||||
|
||||
def verify_batch_matmul(
|
||||
batch_a,
|
||||
batch_b,
|
||||
matrix_m,
|
||||
matrix_l,
|
||||
matrix_n,
|
||||
lib,
|
||||
transa=False,
|
||||
transb=False,
|
||||
dtype="float32",
|
||||
):
|
||||
"""Tests matmul op where matrices are in batch"""
|
||||
batch = max(batch_a, batch_b)
|
||||
ashape = (batch_a, matrix_l, matrix_n) if transa else (batch_a, matrix_n, matrix_l)
|
||||
bshape = (batch_b, matrix_m, matrix_l) if transb else (batch_b, matrix_l, matrix_m)
|
||||
input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
|
||||
input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
|
||||
matmul_result = lib.batch_matmul(input1_data, input2_data, transa, transb)
|
||||
final_result = te.compute(
|
||||
matmul_result.shape, lambda k, i, j: matmul_result[k, i, j], name="final_result"
|
||||
)
|
||||
|
||||
def get_numpy(a, b, transa, transb):
|
||||
if transa:
|
||||
a = a.transpose(0, 2, 1)
|
||||
if not transb:
|
||||
b = b.transpose(0, 2, 1)
|
||||
return tvm.topi.testing.batch_matmul(a, b)
|
||||
|
||||
def compiling(f, name="test_batch_matmul", ext=".so"):
|
||||
path = name + ext
|
||||
f.export_library(path)
|
||||
mod = tvm.runtime.load_module(path)
|
||||
f = mod[name]
|
||||
return f
|
||||
|
||||
def verify(target="llvm"):
|
||||
if not tvm.testing.device_enabled(target):
|
||||
print(f"skip because {target} is not enabled...")
|
||||
return
|
||||
if not tvm.get_global_func(lib.__name__ + ".matmul", True):
|
||||
print("skip because extern function is not available")
|
||||
return
|
||||
dev = tvm.cpu(0)
|
||||
name = "test_batch_matmul"
|
||||
f = tvm.compile(
|
||||
te.create_prim_func([input1_data, input2_data, final_result]), target=target
|
||||
)
|
||||
if target == "c":
|
||||
f = compiling(f, name)
|
||||
matrix_input1 = tvm.runtime.tensor(
|
||||
np.random.uniform(size=ashape).astype(input1_data.dtype), dev
|
||||
)
|
||||
matrix_input2 = tvm.runtime.tensor(
|
||||
np.random.uniform(size=bshape).astype(input2_data.dtype), dev
|
||||
)
|
||||
matrix_result = tvm.runtime.tensor(
|
||||
np.zeros((batch, matrix_n, matrix_m), dtype=final_result.dtype), dev
|
||||
)
|
||||
f(matrix_input1, matrix_input2, matrix_result)
|
||||
tvm.testing.assert_allclose(
|
||||
matrix_result.numpy(),
|
||||
get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), transa, transb),
|
||||
rtol=1e-5,
|
||||
)
|
||||
|
||||
verify("llvm")
|
||||
verify("c")
|
||||
|
||||
|
||||
def test_batch_matmul():
|
||||
"""Tests of matmul op where matrices are in batch"""
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, cblas)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, False)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, cblas, False, True)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, True)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, mkl)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, False)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, mkl, False, True)
|
||||
verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, True)
|
||||
verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
|
||||
verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
|
||||
verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
|
||||
verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
|
||||
verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
|
||||
verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
|
||||
verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
|
||||
verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, cblas)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, False)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, cblas, False, False)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, True)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, cblas)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, mkl)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, False)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, mkl, False, False)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, True)
|
||||
verify_batch_matmul(1, 1, 1, 16, 3, mkl)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_matmul_add()
|
||||
test_quantized_matmul_add()
|
||||
test_batch_matmul()
|
||||
Reference in New Issue
Block a user