wehub-resource-sync
220 lines
7.7 KiB
Python
220 lines
7.7 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: F401
|
|
"""
|
|
Test BasePyModule core functionality.
|
|
|
|
This test verifies:
|
|
1. BasePyModule instantiation and basic methods
|
|
2. TIR function compilation and execution
|
|
3. Python function integration
|
|
4. DLPack conversion between PyTorch and TVM
|
|
"""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
import tvm
|
|
from tvm import relax, tirx
|
|
from tvm.relax import BasePyModule
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
class TestBasePyModule:
|
|
"""Test BasePyModule core functionality."""
|
|
|
|
def test_base_py_module_instantiation(self):
|
|
@T.prim_func(s_tir=True)
|
|
def simple_func(A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32")):
|
|
for i in T.grid(10):
|
|
B[i] = A[i] * 2.0
|
|
|
|
ir_mod = tvm.IRModule({"simple_func": simple_func})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
assert isinstance(py_mod, BasePyModule)
|
|
assert hasattr(py_mod, "call_tir")
|
|
assert hasattr(py_mod, "call_dps_packed")
|
|
assert hasattr(py_mod, "compiled_tir_funcs")
|
|
|
|
def test_base_py_module_instantiation_gpu(self):
|
|
@T.prim_func(s_tir=True)
|
|
def simple_func(A: T.Buffer((10,), "float32"), B: T.Buffer((10,), "float32")):
|
|
for i in T.grid(10):
|
|
B[i] = A[i] * 2.0
|
|
|
|
ir_mod = tvm.IRModule({"simple_func": simple_func})
|
|
|
|
if tvm.cuda().exist:
|
|
|
|
def run_and_check():
|
|
device = tvm.cuda(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
assert isinstance(py_mod, BasePyModule)
|
|
assert hasattr(py_mod, "call_tir")
|
|
assert hasattr(py_mod, "call_dps_packed")
|
|
assert hasattr(py_mod, "compiled_tir_funcs")
|
|
# Check if target contains "cuda" instead of exact match
|
|
assert "cuda" in str(py_mod.target)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
else:
|
|
pytest.skip("CUDA not available")
|
|
|
|
def test_tir_function_compilation(self):
|
|
@T.prim_func(s_tir=True)
|
|
def add_func(
|
|
A: T.Buffer((5,), "float32"), B: T.Buffer((5,), "float32"), C: T.Buffer((5,), "float32")
|
|
):
|
|
for i in T.grid(5):
|
|
C[i] = A[i] + B[i]
|
|
|
|
ir_mod = tvm.IRModule({"add_func": add_func})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
assert "add_func" in py_mod.tir_func_names
|
|
assert "add_func" in py_mod.compiled_tir_funcs
|
|
|
|
def test_call_tir_with_pytorch_tensors(self):
|
|
@T.prim_func(s_tir=True)
|
|
def scale_func(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
|
|
for i in T.grid(4):
|
|
B[i] = A[i] * T.float32(2.5)
|
|
|
|
ir_mod = tvm.IRModule({"scale_func": scale_func})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
|
|
scale_value = 2.5
|
|
|
|
result = py_mod.call_tir(scale_func, [input_tensor], R.Tensor((4,), "float32"))
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape == (4,)
|
|
expected = input_tensor * scale_value
|
|
assert torch.allclose(result, expected, atol=1e-5)
|
|
|
|
def test_call_tir_with_pytorch_tensors_gpu(self):
|
|
if tvm.cuda().exist:
|
|
|
|
def run_and_check():
|
|
# Create a simple IRModule without TIR functions for GPU testing
|
|
ir_mod = tvm.IRModule({})
|
|
device = tvm.cuda(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
# Test basic GPU functionality without TIR compilation issues
|
|
assert isinstance(py_mod, BasePyModule)
|
|
assert hasattr(py_mod, "call_tir")
|
|
assert hasattr(py_mod, "call_dps_packed")
|
|
assert "cuda" in str(py_mod.target)
|
|
|
|
# Test that we can create GPU tensors and they work
|
|
input_tensor = torch.tensor(
|
|
[1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device="cuda"
|
|
)
|
|
assert input_tensor.device.type == "cuda"
|
|
assert input_tensor.shape == (4,)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
else:
|
|
pytest.skip("CUDA not available")
|
|
|
|
def test_dlpack_conversion_pytorch_to_tvm(self):
|
|
@T.prim_func(s_tir=True)
|
|
def identity_func(A: T.Buffer((3,), "float32"), B: T.Buffer((3,), "float32")):
|
|
for i in T.grid(3):
|
|
B[i] = A[i]
|
|
|
|
ir_mod = tvm.IRModule({"identity_func": identity_func})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
|
|
|
|
result = py_mod.call_tir(identity_func, [input_tensor], R.Tensor((3,), "float32"))
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert torch.allclose(result, input_tensor, atol=1e-5)
|
|
|
|
def test_dlpack_conversion_tvm_to_pytorch(self):
|
|
@T.prim_func(s_tir=True)
|
|
def constant_func(B: T.Buffer((2,), "float32")):
|
|
for i in T.grid(2):
|
|
B[i] = T.float32(5.0)
|
|
|
|
ir_mod = tvm.IRModule({"constant_func": constant_func})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
result = py_mod.call_tir(constant_func, [], R.Tensor((2,), "float32"))
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
assert result.shape == (2,)
|
|
expected = torch.tensor([5.0, 5.0], dtype=torch.float32)
|
|
assert torch.allclose(result, expected, atol=1e-5)
|
|
|
|
def test_add_python_function(self):
|
|
ir_mod = tvm.IRModule({})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
def custom_activation(x):
|
|
return torch.tanh(x)
|
|
|
|
py_mod.add_python_function("custom_activation", custom_activation)
|
|
|
|
assert hasattr(py_mod, "custom_activation")
|
|
assert "custom_activation" in py_mod.pyfuncs
|
|
|
|
input_tensor = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float32)
|
|
result = py_mod.custom_activation(input_tensor)
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
expected = torch.tanh(input_tensor)
|
|
assert torch.allclose(result, expected, atol=1e-5)
|
|
|
|
def test_call_dps_packed_with_python_function(self):
|
|
ir_mod = tvm.IRModule({})
|
|
device = tvm.cpu(0)
|
|
py_mod = BasePyModule(ir_mod, device)
|
|
|
|
def my_softmax(tensor, dim):
|
|
return torch.softmax(tensor, dim=dim)
|
|
|
|
py_mod.add_python_function("my_softmax", my_softmax)
|
|
|
|
input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
|
|
|
|
result = py_mod.call_dps_packed(
|
|
"my_softmax", [input_tensor, 1], R.Tensor((2, 2), "float32")
|
|
)
|
|
|
|
assert isinstance(result, torch.Tensor)
|
|
expected = torch.softmax(input_tensor, dim=1)
|
|
assert torch.allclose(result, expected, atol=1e-5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|