wehub-resource-sync
309 lines
12 KiB
Python
309 lines
12 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, F811, F841
|
|
"""
|
|
Test DLPack integration between PyTorch and TVM.
|
|
|
|
This test verifies:
|
|
1. DLPack conversion from PyTorch to TVM
|
|
2. DLPack conversion from TVM to PyTorch
|
|
3. Data integrity preservation during conversion
|
|
4. Functionality equivalence between DLPack and numpy fallback
|
|
5. Error handling for unsupported data types
|
|
"""
|
|
|
|
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 TestDLPackIntegration:
|
|
def test_dlpack_pytorch_to_tvm_conversion(self):
|
|
pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)
|
|
|
|
tvm_tensor = tvm.runtime.from_dlpack(pytorch_tensor)
|
|
|
|
assert isinstance(tvm_tensor, tvm.runtime.Tensor)
|
|
assert tvm_tensor.shape == pytorch_tensor.shape
|
|
assert str(tvm_tensor.dtype) == str(pytorch_tensor.dtype).replace("torch.", "")
|
|
|
|
tvm_numpy = tvm_tensor.numpy()
|
|
pytorch_numpy = pytorch_tensor.numpy()
|
|
tvm.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)
|
|
|
|
def test_dlpack_pytorch_to_tvm_conversion_gpu(self):
|
|
if tvm.cuda().exist:
|
|
|
|
def run_and_check():
|
|
pytorch_tensor = torch.tensor(
|
|
[1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32, device="cuda"
|
|
)
|
|
|
|
tvm_tensor = tvm.runtime.from_dlpack(pytorch_tensor)
|
|
|
|
assert isinstance(tvm_tensor, tvm.runtime.Tensor)
|
|
assert tvm_tensor.shape == pytorch_tensor.shape
|
|
assert str(tvm_tensor.dtype) == str(pytorch_tensor.dtype).replace("torch.", "")
|
|
assert str(tvm_tensor.device) == "cuda:0"
|
|
|
|
# Move to CPU for numpy conversion
|
|
tvm_numpy = tvm_tensor.numpy()
|
|
pytorch_numpy = pytorch_tensor.cpu().numpy()
|
|
tvm.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
else:
|
|
pytest.skip("CUDA not available")
|
|
|
|
def test_dlpack_tvm_to_pytorch_conversion(self):
|
|
import numpy as np
|
|
|
|
data = np.array([1.0, 2.0, 3.0, 5.0], dtype="float32")
|
|
tvm_tensor = tvm.runtime.tensor(data)
|
|
|
|
pytorch_tensor = torch.from_dlpack(tvm_tensor)
|
|
|
|
assert isinstance(pytorch_tensor, torch.Tensor)
|
|
assert pytorch_tensor.shape == tvm_tensor.shape
|
|
assert pytorch_tensor.dtype == torch.float32
|
|
|
|
tvm_numpy = tvm_tensor.numpy()
|
|
pytorch_numpy = pytorch_tensor.numpy()
|
|
tvm.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)
|
|
|
|
def test_dlpack_tvm_to_pytorch_conversion_gpu(self):
|
|
if tvm.cuda().exist:
|
|
import numpy as np
|
|
|
|
data = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype="float32")
|
|
|
|
def run_and_check():
|
|
tvm_tensor = tvm.runtime.tensor(data, device=tvm.cuda(0))
|
|
|
|
pytorch_tensor = torch.from_dlpack(tvm_tensor)
|
|
|
|
assert isinstance(pytorch_tensor, torch.Tensor)
|
|
assert pytorch_tensor.shape == tvm_tensor.shape
|
|
assert pytorch_tensor.dtype == torch.float32
|
|
assert pytorch_tensor.device.type == "cuda"
|
|
|
|
tvm_numpy = tvm_tensor.numpy()
|
|
pytorch_numpy = pytorch_tensor.cpu().numpy()
|
|
tvm.testing.assert_allclose(tvm_numpy, pytorch_numpy, atol=1e-5)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
else:
|
|
pytest.skip("CUDA not available")
|
|
|
|
def test_dlpack_roundtrip_conversion(self):
|
|
"""Test roundtrip conversion: PyTorch -> TVM -> PyTorch."""
|
|
# Create PyTorch tensor
|
|
original_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)
|
|
|
|
# Convert to TVM
|
|
tvm_tensor = tvm.runtime.from_dlpack(original_tensor)
|
|
|
|
# Convert back to PyTorch
|
|
result_tensor = torch.from_dlpack(tvm_tensor)
|
|
|
|
# Verify roundtrip integrity
|
|
assert torch.allclose(original_tensor, result_tensor, atol=1e-5)
|
|
assert original_tensor.dtype == result_tensor.dtype
|
|
assert original_tensor.shape == result_tensor.shape
|
|
|
|
def test_dlpack_different_data_types(self):
|
|
"""Test DLPack conversion with different data types."""
|
|
test_types = [
|
|
(torch.float32, "float32"),
|
|
(torch.float64, "float64"),
|
|
(torch.int32, "int32"),
|
|
(torch.int64, "int64"),
|
|
]
|
|
|
|
for torch_dtype, tvm_dtype in test_types:
|
|
# Create PyTorch tensor
|
|
pytorch_tensor = torch.tensor([1, 2, 3], dtype=torch_dtype)
|
|
|
|
# Convert to TVM
|
|
tvm_tensor = tvm.runtime.from_dlpack(pytorch_tensor)
|
|
|
|
# Convert back to PyTorch
|
|
result_tensor = torch.from_dlpack(tvm_tensor)
|
|
|
|
# Verify conversion
|
|
assert torch.allclose(pytorch_tensor, result_tensor, atol=1e-5)
|
|
assert pytorch_tensor.dtype == result_tensor.dtype
|
|
|
|
def test_dlpack_different_shapes(self):
|
|
"""Test DLPack conversion with different tensor shapes."""
|
|
test_shapes = [
|
|
(1,),
|
|
(2, 3),
|
|
(4, 5, 6),
|
|
(1, 1, 1, 1),
|
|
]
|
|
|
|
for shape in test_shapes:
|
|
# Create PyTorch tensor
|
|
pytorch_tensor = torch.randn(shape, dtype=torch.float32)
|
|
|
|
# Convert to TVM
|
|
tvm_tensor = tvm.runtime.from_dlpack(pytorch_tensor)
|
|
|
|
# Convert back to PyTorch
|
|
result_tensor = torch.from_dlpack(tvm_tensor)
|
|
|
|
# Verify conversion
|
|
assert torch.allclose(pytorch_tensor, result_tensor, atol=1e-5)
|
|
assert pytorch_tensor.shape == result_tensor.shape
|
|
|
|
def test_dlpack_functionality_verification(self):
|
|
"""Test that DLPack and numpy conversions produce identical results."""
|
|
# Create large PyTorch tensor
|
|
size = 1000000
|
|
pytorch_tensor = torch.randn(size, dtype=torch.float32)
|
|
|
|
# Test DLPack conversion
|
|
tvm_tensor_dlpack = tvm.runtime.from_dlpack(pytorch_tensor)
|
|
|
|
# Test numpy conversion
|
|
numpy_array = pytorch_tensor.detach().cpu().numpy()
|
|
tvm_tensor_numpy = tvm.runtime.tensor(numpy_array)
|
|
|
|
# Verify both methods produce same result
|
|
result_dlpack = torch.from_dlpack(tvm_tensor_dlpack)
|
|
result_numpy = torch.from_numpy(tvm_tensor_numpy.numpy())
|
|
assert torch.allclose(result_dlpack, result_numpy, atol=1e-5)
|
|
|
|
# Verify data integrity
|
|
assert torch.allclose(result_dlpack, pytorch_tensor, atol=1e-5)
|
|
assert result_dlpack.shape == pytorch_tensor.shape
|
|
assert result_dlpack.dtype == pytorch_tensor.dtype
|
|
|
|
def test_dlpack_error_handling(self):
|
|
"""Test DLPack error handling for unsupported operations."""
|
|
# Test with non-contiguous tensor
|
|
pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)
|
|
non_contiguous = pytorch_tensor[::2] # Create non-contiguous view
|
|
|
|
# This should work (PyTorch handles non-contiguous tensors)
|
|
try:
|
|
tvm_tensor = tvm.runtime.from_dlpack(non_contiguous)
|
|
result_tensor = torch.from_dlpack(tvm_tensor)
|
|
assert torch.allclose(non_contiguous, result_tensor, atol=1e-5)
|
|
except Exception as e:
|
|
# If it fails, that's also acceptable
|
|
pass
|
|
|
|
def test_dlpack_with_base_py_module(self):
|
|
"""Test DLPack conversion within BasePyModule context."""
|
|
|
|
# Create a simple IRModule
|
|
@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)
|
|
|
|
# Create PyTorch tensor
|
|
input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
|
|
|
|
# Call TIR function (this will trigger DLPack conversion)
|
|
result = py_mod.call_tir(identity_func, [input_tensor], R.Tensor((3,), "float32"))
|
|
|
|
# Verify result
|
|
assert isinstance(result, torch.Tensor)
|
|
assert torch.allclose(result, input_tensor, atol=1e-5)
|
|
|
|
def test_dlpack_device_consistency(self):
|
|
"""Test DLPack conversion maintains device consistency."""
|
|
# Test CPU tensor
|
|
cpu_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
|
|
cpu_tvm = tvm.runtime.from_dlpack(cpu_tensor)
|
|
cpu_result = torch.from_dlpack(cpu_tvm)
|
|
|
|
assert cpu_result.device.type == "cpu"
|
|
assert torch.allclose(cpu_tensor, cpu_result, atol=1e-5)
|
|
|
|
# Note: GPU testing would require CUDA/OpenCL setup
|
|
# This is a basic test that CPU works correctly
|
|
|
|
def test_dlpack_memory_sharing(self):
|
|
"""Test that DLPack conversion shares memory when possible."""
|
|
# Create PyTorch tensor
|
|
pytorch_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype=torch.float32)
|
|
|
|
# Convert to TVM
|
|
tvm_tensor = tvm.runtime.from_dlpack(pytorch_tensor)
|
|
|
|
# Modify the original tensor
|
|
pytorch_tensor[0] = 10.0
|
|
|
|
# Convert back to PyTorch
|
|
result_tensor = torch.from_dlpack(tvm_tensor)
|
|
|
|
# The result should reflect the modification (memory sharing)
|
|
assert result_tensor[0] == 10.0
|
|
assert torch.allclose(pytorch_tensor, result_tensor, atol=1e-5)
|
|
|
|
def test_dlpack_batch_operations(self):
|
|
"""Test DLPack conversion with batch operations."""
|
|
# Create batch of tensors
|
|
batch_size = 10
|
|
pytorch_tensors = [torch.randn(5, dtype=torch.float32) for _ in range(batch_size)]
|
|
|
|
# Convert all to TVM
|
|
tvm_tensors = [tvm.runtime.from_dlpack(t) for t in pytorch_tensors]
|
|
|
|
# Convert all back to PyTorch
|
|
result_tensors = [torch.from_dlpack(t) for t in tvm_tensors]
|
|
|
|
# Verify all conversions
|
|
for i in range(batch_size):
|
|
assert torch.allclose(pytorch_tensors[i], result_tensors[i], atol=1e-5)
|
|
|
|
def test_dlpack_edge_cases(self):
|
|
"""Test DLPack conversion with edge cases."""
|
|
# Empty tensor
|
|
empty_tensor = torch.tensor([], dtype=torch.float32)
|
|
empty_tvm = tvm.runtime.from_dlpack(empty_tensor)
|
|
empty_result = torch.from_dlpack(empty_tvm)
|
|
|
|
assert empty_result.shape == empty_tensor.shape
|
|
assert empty_result.dtype == empty_tensor.dtype
|
|
|
|
# Single element tensor
|
|
single_tensor = torch.tensor([42.0], dtype=torch.float32)
|
|
single_tvm = tvm.runtime.from_dlpack(single_tensor)
|
|
single_result = torch.from_dlpack(single_tvm)
|
|
|
|
assert single_result.shape == single_tensor.shape
|
|
assert single_result[0] == 42.0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|