chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,263 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user