wehub-resource-sync
381 lines
10 KiB
Python
381 lines
10 KiB
Python
import pytest
|
|
|
|
pytest.importorskip("jaxlib", reason="jaxlib not available")
|
|
pytest.importorskip("jax", reason="jax not available")
|
|
|
|
# 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, F841
|
|
|
|
# pylint: disable=c-extension-no-member
|
|
|
|
import functools
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.relax.frontend.stablehlo import from_stablehlo
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
def generate_np_inputs(
|
|
input_shapes: tuple | list[tuple], dtype: str = "float32"
|
|
) -> np.ndarray | list[np.ndarray]:
|
|
"""Generate numpy data as the inputs of model
|
|
|
|
Parameters
|
|
----------
|
|
input_shapes: Union[Tuple, List[Tuple]]
|
|
shapes for inputs
|
|
dtype: str
|
|
the data type of inputs
|
|
|
|
Results
|
|
-------
|
|
out: List[np.ndarray]
|
|
numpy input data
|
|
"""
|
|
if not isinstance(input_shapes[0], list | tuple):
|
|
return [np.random.uniform(size=input_shapes).astype(dtype)]
|
|
out = []
|
|
for input_shape in input_shapes:
|
|
out.append(np.random.uniform(size=input_shape).astype(dtype))
|
|
return out
|
|
|
|
|
|
def np2jnp(inputs_np: np.ndarray | list[np.ndarray]):
|
|
"""Convert data from numpy to jax.numpy
|
|
|
|
Parameters
|
|
----------
|
|
inputs_np: Union[np.ndarray, List[np.ndarray]]
|
|
numpy input data
|
|
|
|
Results
|
|
-------
|
|
out: Union[jnp.ndarray, List[jnp.ndarray]]
|
|
jax numpy data
|
|
"""
|
|
import jax.numpy as jnp
|
|
|
|
# Use jnp.asarray to avoid unnecessary memory copies
|
|
inputs_jnp = []
|
|
if isinstance(inputs_np, tuple | list):
|
|
for input_np in inputs_np:
|
|
inputs_jnp.append(jnp.asarray(input_np))
|
|
return inputs_jnp
|
|
return jnp.asarray(inputs_np)
|
|
|
|
|
|
def check_correctness(
|
|
jax_jit_mod,
|
|
input_shapes: tuple | list[tuple],
|
|
dtype: str = "float32",
|
|
) -> None:
|
|
"""Run a jax model and the translated TVM IRModule,
|
|
verify the inference accuracy.
|
|
|
|
Parameters
|
|
----------
|
|
jax_jit_mod: jaxlib.xla_extension.CompiledFunction
|
|
The input jax jitted model
|
|
input_shapes: Union[Tuple, List[Tuple]]
|
|
shapes for inputs
|
|
dtype: str
|
|
the data type of inputs
|
|
"""
|
|
# Generate numpy inputs
|
|
inputs_np = generate_np_inputs(input_shapes, dtype)
|
|
# Get the jax numpy data
|
|
inputs_jnp = np2jnp(inputs_np)
|
|
|
|
# lower the jitted function to StableHLO
|
|
lowered = jax_jit_mod.lower(*inputs_np)
|
|
|
|
# lowered.as_text(dialect="stablehlo") generates text format
|
|
# compiler_ir generates the related jaxlib.mlir.Module
|
|
stablehlo_module = lowered.compiler_ir(dialect="stablehlo")
|
|
|
|
# Convert the StableHLO IR to Relax
|
|
ir_mod = from_stablehlo(stablehlo_module)
|
|
|
|
# Run the jax jitted model with the input jax numpy data
|
|
jax_output = jax_jit_mod(*inputs_jnp)
|
|
|
|
# TODO (yongwww): support multiple targets,
|
|
# "llvm" should be good for this check
|
|
target = tvm.target.Target("llvm", host="llvm")
|
|
# Compile and run
|
|
ex = tvm.compile(ir_mod, target)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
vm.set_input("main", *inputs_np)
|
|
vm.invoke_stateful("main")
|
|
tvm_output = vm.get_outputs("main")
|
|
|
|
# Single ouput
|
|
if isinstance(tvm_output, tvm.runtime.Tensor):
|
|
tvm.testing.assert_allclose(tvm_output.numpy(), jax_output, rtol=1e-5, atol=1e-5)
|
|
return
|
|
|
|
# Multiple ouputs
|
|
assert len(tvm_output) == len(jax_output), "numbers of outputs mismatch"
|
|
for tvm_out, jax_out in zip(tvm_output, jax_output):
|
|
tvm.testing.assert_allclose(tvm_out.numpy(), jax_out, rtol=1e-5, atol=1e-5)
|
|
|
|
|
|
def get_vm_res(
|
|
ir_mod: tvm.IRModule, weights: np.ndarray | list[np.ndarray]
|
|
) -> tvm.runtime.Tensor | list[tvm.runtime.Tensor]:
|
|
"""Compile and run an ir_module on Relax VM
|
|
|
|
Parameters
|
|
----------
|
|
ir_mod: tvm.IRModule
|
|
input ir module
|
|
|
|
weights: Union[np.ndarray, List[np.ndarray]]
|
|
input weights
|
|
|
|
Results
|
|
-------
|
|
out: Union[tvm.runtime.Tensor, List[tvm.runtime.Tensor]]
|
|
inference result
|
|
"""
|
|
target = tvm.target.Target("llvm", host="llvm")
|
|
# Compile and run
|
|
ex = tvm.compile(ir_mod, target)
|
|
vm = relax.VirtualMachine(ex, tvm.cpu())
|
|
vm.set_input("main", *weights)
|
|
vm.invoke_stateful("main")
|
|
tvm_output = vm.get_outputs("main")
|
|
return tvm_output
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_add_dynamic():
|
|
add_dyn = """
|
|
func.func @test(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%1 = stablehlo.add %arg0, %arg1 : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
func.return %1 : tensor<?x?xf32>
|
|
}
|
|
"""
|
|
|
|
mod = from_stablehlo(add_dyn)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
arg0: R.Tensor(("n_0", "n_1"), dtype="float32"),
|
|
arg1: R.Tensor(("n_2", "n_3"), dtype="float32"),
|
|
) -> R.Tensor(dtype="float32", ndim=2):
|
|
n_0 = T.int64()
|
|
n_1 = T.int64()
|
|
n_2 = T.int64()
|
|
n_3 = T.int64()
|
|
with R.dataflow():
|
|
lv: R.Tensor(dtype="float32", ndim=2) = R.add(arg0, arg1)
|
|
gv: R.Tensor(dtype="float32", ndim=2) = lv
|
|
R.output(gv)
|
|
return gv
|
|
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_unary():
|
|
import jax
|
|
|
|
def _rsqrt(x):
|
|
return jax.lax.rsqrt(x)
|
|
|
|
def _sqrt(x):
|
|
return jax.lax.sqrt(x)
|
|
|
|
def _sin(x):
|
|
return jax.lax.sin(x)
|
|
|
|
def _sinh(x):
|
|
return jax.lax.sinh(x)
|
|
|
|
def _cos(x):
|
|
return jax.lax.cos(x)
|
|
|
|
def _cosh(x):
|
|
return jax.lax.cos(x)
|
|
|
|
def _exp(x):
|
|
return jax.lax.exp(x)
|
|
|
|
def _round(x):
|
|
return jax.lax.round(x)
|
|
|
|
input_shapes = (2, 3, 4)
|
|
for fn in [_rsqrt, _sqrt, _sin, _cos, _cosh, _exp, _round]:
|
|
check_correctness(jax.jit(fn), input_shapes)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_binary():
|
|
import jax
|
|
|
|
def fn(x, y):
|
|
r1 = x + y
|
|
r2 = r1 * r1
|
|
r3 = r2 / r1
|
|
r = r2 - r3
|
|
return r
|
|
|
|
input_shape = (1, 2, 3)
|
|
input_shapes = (input_shape, input_shape)
|
|
|
|
# jit the function
|
|
jit_fn = jax.jit(fn)
|
|
|
|
# verify inference accuracy
|
|
check_correctness(jit_fn, input_shapes)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_const():
|
|
import jax
|
|
|
|
def fn(x):
|
|
return x + 1
|
|
|
|
check_correctness(jax.jit(fn), (2,))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_maximum():
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
def fn(x, y):
|
|
return jnp.maximum(x, y)
|
|
|
|
check_correctness(jax.jit(fn), ((2, 3), (2, 3)))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_minimum():
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
def fn(x, y):
|
|
return jnp.minimum(x, y)
|
|
|
|
check_correctness(jax.jit(fn), ((2, 3), (2, 3)))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
@pytest.mark.skip(
|
|
reason="jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed."
|
|
)
|
|
# TODO(mshr-h): may be fixed by upgrading jax to >=0.4.33
|
|
def test_reduce():
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
def fn(x):
|
|
return jnp.mean(x, axis=(1, 2))
|
|
|
|
check_correctness(jax.jit(fn), (2, 3, 4, 5))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
@pytest.mark.skip(
|
|
reason="jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed."
|
|
)
|
|
# TODO(mshr-h): may be fixed by upgrading jax to >=0.4.33
|
|
def test_reduce_window():
|
|
import jax
|
|
from flax import linen as nn
|
|
|
|
def fn(x):
|
|
return nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")
|
|
|
|
check_correctness(jax.jit(fn), (2, 3, 4))
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
def test_dot_general():
|
|
import jax
|
|
|
|
def fn(x, y):
|
|
return jax.lax.dot_general(x, y, (([1], [0]), ([], [])))
|
|
|
|
input_shapes = ((1, 512), (512, 2))
|
|
check_correctness(jax.jit(fn), input_shapes)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
|
|
@pytest.mark.skip(
|
|
reason="jaxlib.xla_extension.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed."
|
|
)
|
|
# TODO(yongwww): fix flaky error of "invalid device ordinal"
|
|
def test_conv():
|
|
import jax
|
|
import jax.random as jrandom
|
|
from flax import linen as nn
|
|
|
|
conv = nn.Conv(64, (7, 7), (2, 2), padding=[(3, 3), (3, 3)], name="conv_init")
|
|
input_shape = (7, 7, 5, 64)
|
|
|
|
input_np = generate_np_inputs(input_shape)[0]
|
|
input_jnp = np2jnp(input_np)
|
|
# initialize the conv
|
|
weights = conv.init(jrandom.PRNGKey(0), input_jnp)
|
|
# get jax inference output
|
|
jax_output = conv.apply(weights, input_jnp)
|
|
|
|
# assemble numpy data using weights generated above
|
|
kernel_np = np.asarray(weights["params"]["kernel"])
|
|
bias_np = np.asarray(weights["params"]["bias"])
|
|
inputs_np = [bias_np, kernel_np, input_np]
|
|
|
|
# jit and lower to StableHLO
|
|
apply = functools.partial(conv.apply)
|
|
stablehlo_module = jax.jit(apply).lower(weights, input_jnp).compiler_ir(dialect="stablehlo")
|
|
|
|
# convert in Relax
|
|
ir_mod = from_stablehlo(stablehlo_module)
|
|
# compile and run
|
|
tvm_output = get_vm_res(ir_mod, inputs_np)
|
|
# verify accuracy
|
|
tvm.testing.assert_allclose(tvm_output.numpy(), jax_output, rtol=1e-5, atol=1e-5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|