chore: import upstream snapshot with attribution

tests/python/relax/test_pytorch_integration.py

@@ -0,0 +1,392 @@
+# 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 PyTorch integration with TVM Relax.
+
+This test verifies:
+1. Seamless PyTorch tensor I/O with TVM backend
+2. Cross-function calls between Python, TIR, and Relax functions
+3. Dynamic Python function addition and execution
+4. End-to-end pipeline testing
+5. Error handling and edge cases
+"""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn.functional as F
+
+import tvm
+from tvm import relax, tirx
+from tvm.relax import BasePyModule
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+@I.ir_module(s_tir=True)
+class PyTorchIntegrationModule(BasePyModule):
+    """Test module for PyTorch integration with TVM."""
+
+    @I.pyfunc
+    def main(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+        """Main function demonstrating cross-function calls."""
+        n = x.shape[0]
+
+        # Call TIR function
+        lv = self.call_tir(self.matmul, [x, w], out_ty=R.Tensor((n, 20), "float32"))
+
+        # Apply ReLU
+        lv1 = F.relu(lv)
+
+        # Call packed function (will be added dynamically)
+        lv2 = self.call_dps_packed("my_softmax", [lv1, 1], out_ty=R.Tensor((n, 20), "float32"))
+
+        # Call Python function
+        lv3 = self.my_identity_func(lv2)
+
+        return lv3
+
+    @T.prim_func(s_tir=True)
+    def matmul(
+        var_A: T.handle,
+        var_B: T.handle,
+        var_C: T.handle,
+    ):
+        """TIR function for matrix multiplication."""
+        n = T.int32()
+        A = T.match_buffer(var_A, (n, 16), "float32")
+        B = T.match_buffer(var_B, (16, 20), "float32")
+        C = T.match_buffer(var_C, (n, 20), "float32")
+
+        for i, j, k in T.grid(n, 20, 16):
+            with T.sblock("block"):
+                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
+                with T.init():
+                    C[vi, vj] = T.float32(0)
+                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
+
+    @I.pyfunc
+    def my_identity_func(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+
+class TestPyTorchIntegration:
+    def test_module_creation_and_instantiation(self):
+        module = PyTorchIntegrationModule
+
+        assert hasattr(module, "__call__"), "Module should be callable"
+
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        assert isinstance(instance, BasePyModule), "Instance should be BasePyModule"
+
+        required_methods = ["main", "call_tir", "call_dps_packed"]
+        for method in required_methods:
+            assert hasattr(instance, method), f"Instance should have method: {method}"
+
+    def test_module_creation_and_instantiation_gpu(self):
+        module = PyTorchIntegrationModule
+
+        if tvm.cuda().exist:
+
+            def run_and_check():
+                assert hasattr(module, "__call__"), "Module should be callable"
+
+                device = tvm.cuda(0)
+                instance = module(device)
+
+                assert isinstance(instance, BasePyModule), "Instance should be BasePyModule"
+                required_methods = ["main", "call_tir", "call_dps_packed"]
+                for method in required_methods:
+                    assert hasattr(instance, method), f"Instance should have method: {method}"
+                assert "cuda" in str(instance.target)
+
+            tvm.testing.run_with_gpu_lock(run_and_check)
+        else:
+            pytest.skip("CUDA not available")
+
+    def test_python_function_execution(self):
+        """Test that Python functions execute correctly."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Test my_identity_func
+        input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+        result = instance.my_identity_func(input_tensor)
+
+        assert isinstance(result, torch.Tensor)
+        assert torch.allclose(result, input_tensor, atol=1e-5)
+
+    def test_tir_function_execution(self):
+        """Test that TIR functions execute correctly."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Test matmul function
+        n = 3
+        x = torch.randn(n, 16, dtype=torch.float32)
+        w = torch.randn(16, 20, dtype=torch.float32)
+
+        result = instance.call_tir(instance.matmul, [x, w], R.Tensor((n, 20), "float32"))
+
+        assert isinstance(result, torch.Tensor)
+        assert result.shape == (n, 20)
+
+        # Verify result with PyTorch matmul
+        expected = torch.matmul(x, w)
+        assert torch.allclose(result, expected, atol=1e-3)
+
+    def test_dynamic_python_function_addition(self):
+        """Test adding Python functions dynamically."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Define a custom function
+        def custom_activation(x):
+            return torch.sigmoid(x)
+
+        # Add the function
+        instance.add_python_function("custom_activation", custom_activation)
+
+        # Verify function is added
+        assert hasattr(instance, "custom_activation")
+        assert "custom_activation" in instance.pyfuncs
+
+        # Test function execution
+        input_tensor = torch.tensor([1.0, -1.0, 0.0], dtype=torch.float32)
+        result = instance.custom_activation(input_tensor)
+
+        assert isinstance(result, torch.Tensor)
+        expected = torch.sigmoid(input_tensor)
+        assert torch.allclose(result, expected, atol=1e-5)
+
+    def test_call_dps_packed_with_dynamic_function(self):
+        """Test call_dps_packed with dynamically added function."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Define my_softmax function
+        def my_softmax(tensor, dim):
+            """Custom softmax function for testing call_dps_packed."""
+            # Convert TVM Tensor to PyTorch tensor if needed
+            if hasattr(tensor, "numpy"):
+                tensor = torch.from_numpy(tensor.numpy())
+            return F.softmax(tensor, dim=dim)
+
+        # Add the function
+        instance.my_softmax = my_softmax
+
+        # Test call_dps_packed
+        input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
+
+        result = instance.call_dps_packed(
+            "my_softmax", [input_tensor, 1], R.Tensor((2, 2), "float32")
+        )
+
+        assert isinstance(result, torch.Tensor)
+        expected = F.softmax(input_tensor, dim=1)
+        assert torch.allclose(result, expected, atol=1e-5)
+
+    def test_end_to_end_pipeline(self):
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        def my_softmax(tensor, dim):
+            if hasattr(tensor, "numpy"):
+                tensor = torch.from_numpy(tensor.numpy())
+            return F.softmax(tensor, dim=dim)
+
+        instance.my_softmax = my_softmax
+
+        n = 5
+        x = torch.randn(n, 16, dtype=torch.float32)
+        w = torch.randn(16, 20, dtype=torch.float32)
+
+        result = instance.main(x, w)
+
+        assert isinstance(result, torch.Tensor)
+        assert result.shape == (n, 20)
+        assert result.dtype == torch.float32
+
+    def test_end_to_end_pipeline_gpu(self):
+        module = PyTorchIntegrationModule
+
+        if tvm.cuda().exist:
+
+            def run_and_check():
+                device = tvm.cuda(0)
+                instance = module(device)
+
+                # Test basic GPU functionality without complex TIR operations
+                assert isinstance(instance, BasePyModule)
+                assert "cuda" in str(instance.target)
+
+                # Test that we can create and work with GPU tensors
+                n = 5
+                x = torch.randn(n, 16, dtype=torch.float32, device="cuda")
+                w = torch.randn(16, 20, dtype=torch.float32, device="cuda")
+
+                assert x.device.type == "cuda"
+                assert w.device.type == "cuda"
+                assert x.shape == (n, 16)
+                assert w.shape == (16, 20)
+
+                # Test basic PyTorch operations on GPU
+                result = torch.matmul(x, w)
+                assert isinstance(result, torch.Tensor)
+                assert result.shape == (n, 20)
+                assert result.dtype == torch.float32
+                assert result.device.type == "cuda"
+
+            tvm.testing.run_with_gpu_lock(run_and_check)
+        else:
+            pytest.skip("CUDA not available")
+
+    def test_cross_function_data_flow(self):
+        """Test data flow between different function types."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Add required functions
+        def my_softmax(tensor, dim):
+            if hasattr(tensor, "numpy"):
+                tensor = torch.from_numpy(tensor.numpy())
+            return F.softmax(tensor, dim=dim)
+
+        instance.my_softmax = my_softmax
+
+        # Create test data
+        n = 4
+        x = torch.randn(n, 16, dtype=torch.float32)
+        w = torch.randn(16, 20, dtype=torch.float32)
+
+        # Execute step by step to verify data flow
+        # Step 1: TIR matmul
+        lv = instance.call_tir(instance.matmul, [x, w], R.Tensor((n, 20), "float32"))
+        assert isinstance(lv, torch.Tensor)
+        assert lv.shape == (n, 20)
+
+        # Step 2: ReLU
+        lv1 = F.relu(lv)
+        assert isinstance(lv1, torch.Tensor)
+        assert lv1.shape == (n, 20)
+
+        # Step 3: Softmax via call_dps_packed
+        lv2 = instance.call_dps_packed("my_softmax", [lv1, 1], R.Tensor((n, 20), "float32"))
+        assert isinstance(lv2, torch.Tensor)
+        assert lv2.shape == (n, 20)
+
+        # Step 4: Identity function
+        lv3 = instance.my_identity_func(lv2)
+        assert isinstance(lv3, torch.Tensor)
+        assert lv3.shape == (n, 20)
+
+        # Verify final result matches expected
+        expected = F.softmax(F.relu(torch.matmul(x, w)), dim=1)
+        assert torch.allclose(lv3, expected, atol=1e-3)
+
+    def test_error_handling(self):
+        """Test error handling for various edge cases."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Test with missing function
+        with pytest.raises(Exception):
+            instance.call_dps_packed(
+                "non_existent_function", [torch.tensor([1.0])], R.Tensor((1,), "float32")
+            )
+
+        # Test with wrong tensor shapes
+        x = torch.randn(3, 16, dtype=torch.float32)
+        w = torch.randn(15, 20, dtype=torch.float32)  # Wrong shape
+
+        with pytest.raises(Exception):
+            instance.call_tir(instance.matmul, [x, w], R.Tensor((3, 20), "float32"))
+
+    def test_tensor_type_preservation(self):
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        def my_softmax(tensor, dim):
+            if hasattr(tensor, "numpy"):
+                tensor = torch.from_numpy(tensor.numpy())
+            return F.softmax(tensor, dim=dim)
+
+        instance.my_softmax = my_softmax
+
+        # Test with float32 data type (TIR function is hardcoded for float32)
+        test_dtype = torch.float32
+        n = 3
+        x = torch.randn(n, 16, dtype=test_dtype)
+        w = torch.randn(16, 20, dtype=test_dtype)
+
+        result = instance.main(x, w)
+
+        # Verify type preservation
+        assert result.dtype == test_dtype
+        assert isinstance(result, torch.Tensor)
+        assert result.shape == (n, 20)
+        assert result.dtype == torch.float32
+
+    def test_batch_processing(self):
+        """Test processing multiple inputs in batch."""
+        module = PyTorchIntegrationModule
+        device = tvm.cpu(0)
+        instance = module(device)
+
+        # Add required functions
+        def my_softmax(tensor, dim):
+            if hasattr(tensor, "numpy"):
+                tensor = torch.from_numpy(tensor.numpy())
+            return F.softmax(tensor, dim=dim)
+
+        instance.my_softmax = my_softmax
+
+        # Process multiple inputs
+        batch_size = 5
+        results = []
+
+        for i in range(batch_size):
+            n = 3 + i  # Varying batch sizes
+            x = torch.randn(n, 16, dtype=torch.float32)
+            w = torch.randn(16, 20, dtype=torch.float32)
+
+            result = instance.main(x, w)
+            results.append(result)
+
+            assert isinstance(result, torch.Tensor)
+            assert result.shape == (n, 20)
+
+        # Verify all results are valid
+        assert len(results) == batch_size
+        for result in results:
+            assert isinstance(result, torch.Tensor)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])