wehub-resource-sync
264 lines
8.0 KiB
Python
264 lines
8.0 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.
|
|
"""Tests for the Executable class."""
|
|
|
|
import os
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.runtime import Executable
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
@tvm.script.ir_module
|
|
class MyModule:
|
|
@T.prim_func(s_tir=True)
|
|
def add(
|
|
A: T.Buffer((10,), "float32"),
|
|
B: T.Buffer((10,), "float32"),
|
|
C: T.Buffer((10,), "float32"),
|
|
):
|
|
for i in range(10):
|
|
C[i] = A[i] + B[i]
|
|
|
|
|
|
def test_executable_init():
|
|
"""Test initialization of Executable class."""
|
|
lib = tvm.tirx.build(MyModule, target="llvm")
|
|
executable = Executable(lib)
|
|
|
|
assert executable.mod is lib
|
|
assert executable._jitted_mod is None
|
|
|
|
|
|
def test_executable_getitem():
|
|
"""Test __getitem__ method of Executable class."""
|
|
lib = tvm.tirx.build(MyModule, target="llvm")
|
|
executable = Executable(lib)
|
|
|
|
# Jit the module first
|
|
executable.jit()
|
|
|
|
# Test __getitem__
|
|
add_func = executable["add"]
|
|
|
|
# Verify the function works
|
|
a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
|
|
b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
|
|
c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
|
|
add_func(a, b, c)
|
|
|
|
# Check results
|
|
tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
|
|
|
|
def test_executable_jit_already_jitted():
|
|
"""Test jit method when module is already jitted."""
|
|
lib = tvm.tirx.build(MyModule, target="llvm")
|
|
executable = Executable(lib)
|
|
|
|
# First jit call
|
|
jitted_mod1 = executable.jit()
|
|
|
|
# Second jit call should return the cached jitted module
|
|
jitted_mod2 = executable.jit()
|
|
assert jitted_mod2 is jitted_mod1
|
|
|
|
# Test with force_recompile
|
|
jitted_mod3 = executable.jit(force_recompile=True)
|
|
# The module might be different after force recompilation
|
|
|
|
# Verify both modules work correctly
|
|
a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
|
|
b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
|
|
c1 = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
c2 = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
|
|
jitted_mod1["add"](a, b, c1)
|
|
jitted_mod3["add"](a, b, c2)
|
|
|
|
tvm.testing.assert_allclose(c1.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
tvm.testing.assert_allclose(c2.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
|
|
|
|
def test_executable_export_library():
|
|
"""Test export_library method."""
|
|
lib = tvm.tirx.build(MyModule, target="llvm")
|
|
executable = Executable(lib)
|
|
|
|
# Create a temporary directory for the library
|
|
temp_dir = tempfile.mkdtemp()
|
|
try:
|
|
lib_path = os.path.join(temp_dir, "test_lib.so")
|
|
executable.export_library(lib_path)
|
|
|
|
# Verify the library was created
|
|
assert os.path.exists(lib_path)
|
|
|
|
# Load the library back
|
|
loaded_mod = tvm.runtime.load_module(lib_path)
|
|
assert loaded_mod is not None
|
|
|
|
# Test the loaded module
|
|
a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
|
|
b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
|
|
c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
|
|
loaded_mod["add"](a, b, c)
|
|
|
|
# Check results
|
|
tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
finally:
|
|
# Clean up
|
|
if os.path.exists(temp_dir):
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir)
|
|
|
|
|
|
def test_executable_export_library_with_workspace():
|
|
"""Test export_library method with workspace_dir."""
|
|
lib = tvm.tirx.build(MyModule, target="llvm")
|
|
executable = Executable(lib)
|
|
|
|
# Create temporary directories
|
|
temp_dir = tempfile.mkdtemp()
|
|
workspace_dir = tempfile.mkdtemp()
|
|
|
|
try:
|
|
lib_path = os.path.join(temp_dir, "test_lib.so")
|
|
executable.export_library(lib_path, workspace_dir=workspace_dir)
|
|
|
|
# Verify the library was created
|
|
assert os.path.exists(lib_path)
|
|
|
|
# Load the library back
|
|
loaded_mod = tvm.runtime.load_module(lib_path)
|
|
assert loaded_mod is not None
|
|
|
|
# Test the loaded module
|
|
a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
|
|
b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
|
|
c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
|
|
loaded_mod["add"](a, b, c)
|
|
|
|
# Check results
|
|
tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
finally:
|
|
# Clean up
|
|
for directory in [temp_dir, workspace_dir]:
|
|
if os.path.exists(directory):
|
|
import shutil
|
|
|
|
shutil.rmtree(directory)
|
|
|
|
|
|
def test_executable_integration():
|
|
"""Integration test for Executable with a simple TVM module."""
|
|
# Create target and build
|
|
target = tvm.target.Target("llvm")
|
|
lib = tvm.tirx.build(MyModule, target=target)
|
|
|
|
# Create an executable
|
|
executable = Executable(lib)
|
|
|
|
# Test jit
|
|
jitted_mod = executable.jit()
|
|
assert jitted_mod is not None
|
|
|
|
# Test __getitem__
|
|
add_func = executable["add"]
|
|
assert add_func is not None
|
|
|
|
# Test the function works
|
|
a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
|
|
b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
|
|
c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
|
|
add_func(a, b, c)
|
|
|
|
# Check results
|
|
tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
|
|
# Test export_library
|
|
temp_dir = tempfile.mkdtemp()
|
|
try:
|
|
lib_path = os.path.join(temp_dir, "test_lib.so")
|
|
executable.export_library(lib_path)
|
|
|
|
# Verify the library was created
|
|
assert os.path.exists(lib_path)
|
|
|
|
# Load the library back
|
|
loaded_mod = tvm.runtime.load_module(lib_path)
|
|
assert loaded_mod is not None
|
|
|
|
# Test the loaded module
|
|
loaded_add = loaded_mod["add"]
|
|
c_loaded = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
loaded_add(a, b, c_loaded)
|
|
|
|
# Check results
|
|
tvm.testing.assert_allclose(c_loaded.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
|
|
finally:
|
|
# Clean up
|
|
if os.path.exists(temp_dir):
|
|
import shutil
|
|
|
|
shutil.rmtree(temp_dir)
|
|
|
|
|
|
def test_executable_jit_force_recompile():
|
|
"""Test jit method with force_recompile=True."""
|
|
# Create target and build
|
|
target = tvm.target.Target("c")
|
|
lib = tvm.tirx.build(MyModule, target=target)
|
|
|
|
# Create an executable
|
|
executable = Executable(lib)
|
|
|
|
# First jit call
|
|
jitted_mod1 = executable.jit()
|
|
|
|
# Second jit call without force_recompile should return the same module
|
|
jitted_mod2 = executable.jit()
|
|
assert jitted_mod1 is jitted_mod2
|
|
|
|
# Third jit call with force_recompile should return a new module
|
|
jitted_mod3 = executable.jit(force_recompile=True)
|
|
assert jitted_mod3 is not jitted_mod1
|
|
|
|
# Test the function works
|
|
a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
|
|
b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
|
|
c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))
|
|
|
|
jitted_mod3["add"](a, b, c)
|
|
|
|
# Check results
|
|
tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|