tests/python/relax/test_tvmscript_pyfunc.py - tvm - Git at Google

 # 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.
 """
 Test TVMScript @I.pyfunc decorator functionality.

 This test verifies:
 1. @I.pyfunc decorator works correctly
 2. Python functions are properly integrated into IRModule
 3. BasePyModule inheritance is handled correctly
 4. ExternFunc nodes are created for Python functions
 """

 import pytest
 import torch
 import tvm
 from tvm import relax
 from tvm.script import ir as I, relax as R, tir as T
 from tvm.relax import BasePyModule
 import numpy as np


 @I.ir_module
 class TestPyFuncModule(BasePyModule):
     """Test module with Python functions using @I.pyfunc decorator."""

     @I.pyfunc
     def pytorch_processor(x: torch.Tensor) -> torch.Tensor:
         """Python function that processes PyTorch tensors."""
         return torch.nn.functional.relu(x) * 2.0

     @I.pyfunc
     def pytorch_adder(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """Python function that adds two PyTorch tensors."""
         return x + y

     @I.pyfunc
     def pytorch_complex_ops(x: torch.Tensor) -> torch.Tensor:
         """Complex PyTorch operations."""
         result = torch.nn.functional.softmax(x, dim=0)
         result = torch.nn.functional.dropout(result, p=0.1, training=False)
         return result * 10.0

     @T.prim_func
     def simple_tir_func(
         var_A: T.handle,
         var_B: T.handle,
     ):
         T.func_attr({"tir.noalias": True})
         n = T.int32()
         A = T.match_buffer(var_A, (n,), "float32")
         B = T.match_buffer(var_B, (n,), "float32")

         for i in T.grid(n):
             with T.block("copy"):
                 vi = T.axis.remap("S", [i])
                 B[vi] = A[vi]


 class TestTVMScriptPyFunc:
     def test_pyfunc_decorator_creates_pyfuncs_attribute(self):
         module = TestPyFuncModule

         assert hasattr(module, "pyfuncs"), "Module should have pyfuncs attribute"

         pyfuncs = module.pyfuncs
         assert isinstance(pyfuncs, dict), "pyfuncs should be a dictionary"

         expected_functions = ["pytorch_processor", "pytorch_adder", "pytorch_complex_ops"]
         for func_name in expected_functions:
             assert func_name in pyfuncs, f"Function {func_name} should be in pyfuncs"

     def test_pyfunc_functions_are_callable(self):
         """Test that Python functions in pyfuncs are callable."""
         module = TestPyFuncModule
         pyfuncs = module.pyfuncs

         # Test pytorch_processor
         processor_func = pyfuncs["pytorch_processor"]
         assert callable(processor_func), "pytorch_processor should be callable"

         # Test pytorch_adder
         adder_func = pyfuncs["pytorch_adder"]
         assert callable(adder_func), "pytorch_adder should be callable"

         # Test pytorch_complex_ops
         complex_func = pyfuncs["pytorch_complex_ops"]
         assert callable(complex_func), "pytorch_complex_ops should be callable"

     def test_pyfunc_functions_execute_correctly(self):
         """Test that Python functions execute correctly."""
         module = TestPyFuncModule
         pyfuncs = module.pyfuncs

         # Create test data
         x = torch.tensor([1.0, -2.0, 3.0, -4.0, 5.0], dtype=torch.float32)
         y = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5], dtype=torch.float32)

         # Test pytorch_processor
         processor_func = pyfuncs["pytorch_processor"]
         processor_result = processor_func(x)

         assert isinstance(processor_result, torch.Tensor)
         expected = torch.nn.functional.relu(x) * 2.0
         assert torch.allclose(processor_result, expected, atol=1e-5)

         # Test pytorch_adder
         adder_func = pyfuncs["pytorch_adder"]
         adder_result = adder_func(x, y)

         assert isinstance(adder_result, torch.Tensor)
         expected = x + y
         assert torch.allclose(adder_result, expected, atol=1e-5)

         # Test pytorch_complex_ops
         complex_func = pyfuncs["pytorch_complex_ops"]
         complex_result = complex_func(x)

         assert isinstance(complex_result, torch.Tensor)
         # Note: dropout is non-deterministic, so we just check shape and type
         assert complex_result.shape == x.shape
         assert complex_result.dtype == x.dtype

     def test_pyfunc_module_has_functions_attribute(self):
         """Test that the module has functions attribute for IRModule operations."""
         module = TestPyFuncModule

         # Check if functions attribute exists
         assert hasattr(module, "functions"), "Module should have functions attribute"

         functions = module.functions
         # TVM IRModule.functions is not a standard dict, but has dict-like behavior
         assert hasattr(functions, "__getitem__"), "functions should support dict-like access"
         assert hasattr(functions, "__iter__"), "functions should be iterable"

     def test_pyfunc_module_script_method(self):
         """Test that the module has script() method for TVMScript output."""
         module = TestPyFuncModule

         # Check if script method exists
         assert hasattr(module, "script"), "Module should have script method"

         # Test script method execution
         script_output = module.script()
         assert isinstance(script_output, str), "script() should return a string"
         assert len(script_output) > 0, "script() should return non-empty string"

     def test_pyfunc_module_inheritance_flag(self):
         """Test that the module has BasePyModule inheritance flag."""
         module = TestPyFuncModule

         # Check if inheritance flag exists (this might not be set in all implementations)
         if hasattr(module, "_base_py_module_inherited"):
             assert module._base_py_module_inherited, "Inheritance flag should be True"
         else:
             # Alternative: check if the module supports Python functions
             assert hasattr(module, "pyfuncs"), "Module should support Python functions"

         # Check if original class is preserved (this might not be set in all implementations)
         if hasattr(module, "_original_class"):
             assert module._original_class is not None, "Original class should be preserved"
         else:
             # Alternative: check if module is callable (ModuleFactory)
             assert hasattr(module, "__call__"), "Module should be callable (ModuleFactory)"

     def test_pyfunc_module_creation_and_execution(self):
         module = TestPyFuncModule

         assert hasattr(module, "__call__"), "Module should be callable"

         device = tvm.cpu(0)
         instance = module(device)

         assert isinstance(instance, BasePyModule), "Instance should be BasePyModule"
         assert hasattr(instance, "pyfuncs"), "Instance should have pyfuncs"

         x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
         result = instance.pytorch_processor(x)

         assert isinstance(result, torch.Tensor)
         expected = torch.nn.functional.relu(x) * 2.0
         assert torch.allclose(result, expected, atol=1e-5)

     def test_pyfunc_module_creation_and_execution_gpu(self):
         module = TestPyFuncModule

         if tvm.cuda().exist:
             device = tvm.cuda(0)
             instance = module(device)

             assert isinstance(instance, BasePyModule), "Instance should be BasePyModule"
             assert hasattr(instance, "pyfuncs"), "Instance should have pyfuncs"

             x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device="cuda")
             result = instance.pytorch_processor(x)

             assert isinstance(result, torch.Tensor)
             assert result.device.type == "cuda"
             expected = torch.nn.functional.relu(x) * 2.0
             assert torch.allclose(result, expected, atol=1e-5)
         else:
             pytest.skip("CUDA not available")

     def test_pyfunc_with_tir_integration(self):
         """Test that Python functions can work with TIR functions."""
         module = TestPyFuncModule

         # Create instance
         device = tvm.cpu(0)
         instance = module(device)

         # Test TIR function execution
         n = 5
         input_tensor = torch.randn(n, dtype=torch.float32)

         # Call TIR function - it needs 3 arguments: input, output, and size
         # But call_tir handles the output buffer creation, so we only pass input and size
         # Note: TIR functions expect TVM types, not Python types
         result = instance.call_tir(
             instance.simple_tir_func,
             [input_tensor],  # Only pass input tensor, let call_tir handle the rest
             R.Tensor((n,), "float32"),
         )

         # Verify result
         assert isinstance(result, torch.Tensor)
         assert result.shape == (n,)
         assert torch.allclose(result, input_tensor, atol=1e-5)

     def test_pyfunc_decorator_preserves_function_signatures(self):
         """Test that @I.pyfunc decorator preserves function signatures."""
         module = TestPyFuncModule
         pyfuncs = module.pyfuncs

         # Check function signatures
         import inspect

         # pytorch_processor signature
         processor_func = pyfuncs["pytorch_processor"]
         sig = inspect.signature(processor_func)
         params = list(sig.parameters.keys())
         assert len(params) == 1, "pytorch_processor should have 1 parameter"
         assert params[0] == "x", "First parameter should be 'x'"

         # pytorch_adder signature
         adder_func = pyfuncs["pytorch_adder"]
         sig = inspect.signature(adder_func)
         params = list(sig.parameters.keys())
         assert len(params) == 2, "pytorch_adder should have 2 parameters"
         assert params[0] == "x", "First parameter should be 'x'"
         assert params[1] == "y", "Second parameter should be 'y'"


 if __name__ == "__main__":
     pytest.main([__file__])
	# 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.
	"""
	Test TVMScript @I.pyfunc decorator functionality.

	This test verifies:
	1. @I.pyfunc decorator works correctly
	2. Python functions are properly integrated into IRModule
	3. BasePyModule inheritance is handled correctly
	4. ExternFunc nodes are created for Python functions
	"""

	import pytest
	import torch
	import tvm
	from tvm import relax
	from tvm.script import ir as I, relax as R, tir as T
	from tvm.relax import BasePyModule
	import numpy as np


	@I.ir_module
	class TestPyFuncModule(BasePyModule):
	"""Test module with Python functions using @I.pyfunc decorator."""

	@I.pyfunc
	def pytorch_processor(x: torch.Tensor) -> torch.Tensor:
	"""Python function that processes PyTorch tensors."""
	return torch.nn.functional.relu(x) * 2.0

	@I.pyfunc
	def pytorch_adder(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
	"""Python function that adds two PyTorch tensors."""
	return x + y

	@I.pyfunc
	def pytorch_complex_ops(x: torch.Tensor) -> torch.Tensor:
	"""Complex PyTorch operations."""
	result = torch.nn.functional.softmax(x, dim=0)
	result = torch.nn.functional.dropout(result, p=0.1, training=False)
	return result * 10.0

	@T.prim_func
	def simple_tir_func(
	var_A: T.handle,
	var_B: T.handle,
	):
	T.func_attr({"tir.noalias": True})
	n = T.int32()
	A = T.match_buffer(var_A, (n,), "float32")
	B = T.match_buffer(var_B, (n,), "float32")

	for i in T.grid(n):
	with T.block("copy"):
	vi = T.axis.remap("S", [i])
	B[vi] = A[vi]


	class TestTVMScriptPyFunc:
	def test_pyfunc_decorator_creates_pyfuncs_attribute(self):
	module = TestPyFuncModule

	assert hasattr(module, "pyfuncs"), "Module should have pyfuncs attribute"

	pyfuncs = module.pyfuncs
	assert isinstance(pyfuncs, dict), "pyfuncs should be a dictionary"

	expected_functions = ["pytorch_processor", "pytorch_adder", "pytorch_complex_ops"]
	for func_name in expected_functions:
	assert func_name in pyfuncs, f"Function {func_name} should be in pyfuncs"

	def test_pyfunc_functions_are_callable(self):
	"""Test that Python functions in pyfuncs are callable."""
	module = TestPyFuncModule
	pyfuncs = module.pyfuncs

	# Test pytorch_processor
	processor_func = pyfuncs["pytorch_processor"]
	assert callable(processor_func), "pytorch_processor should be callable"

	# Test pytorch_adder
	adder_func = pyfuncs["pytorch_adder"]
	assert callable(adder_func), "pytorch_adder should be callable"

	# Test pytorch_complex_ops
	complex_func = pyfuncs["pytorch_complex_ops"]
	assert callable(complex_func), "pytorch_complex_ops should be callable"

	def test_pyfunc_functions_execute_correctly(self):
	"""Test that Python functions execute correctly."""
	module = TestPyFuncModule
	pyfuncs = module.pyfuncs

	# Create test data
	x = torch.tensor([1.0, -2.0, 3.0, -4.0, 5.0], dtype=torch.float32)
	y = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5], dtype=torch.float32)

	# Test pytorch_processor
	processor_func = pyfuncs["pytorch_processor"]
	processor_result = processor_func(x)

	assert isinstance(processor_result, torch.Tensor)
	expected = torch.nn.functional.relu(x) * 2.0
	assert torch.allclose(processor_result, expected, atol=1e-5)

	# Test pytorch_adder
	adder_func = pyfuncs["pytorch_adder"]
	adder_result = adder_func(x, y)

	assert isinstance(adder_result, torch.Tensor)
	expected = x + y
	assert torch.allclose(adder_result, expected, atol=1e-5)

	# Test pytorch_complex_ops
	complex_func = pyfuncs["pytorch_complex_ops"]
	complex_result = complex_func(x)

	assert isinstance(complex_result, torch.Tensor)
	# Note: dropout is non-deterministic, so we just check shape and type
	assert complex_result.shape == x.shape
	assert complex_result.dtype == x.dtype

	def test_pyfunc_module_has_functions_attribute(self):
	"""Test that the module has functions attribute for IRModule operations."""
	module = TestPyFuncModule

	# Check if functions attribute exists
	assert hasattr(module, "functions"), "Module should have functions attribute"

	functions = module.functions
	# TVM IRModule.functions is not a standard dict, but has dict-like behavior
	assert hasattr(functions, "__getitem__"), "functions should support dict-like access"
	assert hasattr(functions, "__iter__"), "functions should be iterable"

	def test_pyfunc_module_script_method(self):
	"""Test that the module has script() method for TVMScript output."""
	module = TestPyFuncModule

	# Check if script method exists
	assert hasattr(module, "script"), "Module should have script method"

	# Test script method execution
	script_output = module.script()
	assert isinstance(script_output, str), "script() should return a string"
	assert len(script_output) > 0, "script() should return non-empty string"

	def test_pyfunc_module_inheritance_flag(self):
	"""Test that the module has BasePyModule inheritance flag."""
	module = TestPyFuncModule

	# Check if inheritance flag exists (this might not be set in all implementations)
	if hasattr(module, "_base_py_module_inherited"):
	assert module._base_py_module_inherited, "Inheritance flag should be True"
	else:
	# Alternative: check if the module supports Python functions
	assert hasattr(module, "pyfuncs"), "Module should support Python functions"

	# Check if original class is preserved (this might not be set in all implementations)
	if hasattr(module, "_original_class"):
	assert module._original_class is not None, "Original class should be preserved"
	else:
	# Alternative: check if module is callable (ModuleFactory)
	assert hasattr(module, "__call__"), "Module should be callable (ModuleFactory)"

	def test_pyfunc_module_creation_and_execution(self):
	module = TestPyFuncModule

	assert hasattr(module, "__call__"), "Module should be callable"

	device = tvm.cpu(0)
	instance = module(device)

	assert isinstance(instance, BasePyModule), "Instance should be BasePyModule"
	assert hasattr(instance, "pyfuncs"), "Instance should have pyfuncs"

	x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
	result = instance.pytorch_processor(x)

	assert isinstance(result, torch.Tensor)
	expected = torch.nn.functional.relu(x) * 2.0
	assert torch.allclose(result, expected, atol=1e-5)

	def test_pyfunc_module_creation_and_execution_gpu(self):
	module = TestPyFuncModule

	if tvm.cuda().exist:
	device = tvm.cuda(0)
	instance = module(device)

	assert isinstance(instance, BasePyModule), "Instance should be BasePyModule"
	assert hasattr(instance, "pyfuncs"), "Instance should have pyfuncs"

	x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device="cuda")
	result = instance.pytorch_processor(x)

	assert isinstance(result, torch.Tensor)
	assert result.device.type == "cuda"
	expected = torch.nn.functional.relu(x) * 2.0
	assert torch.allclose(result, expected, atol=1e-5)
	else:
	pytest.skip("CUDA not available")

	def test_pyfunc_with_tir_integration(self):
	"""Test that Python functions can work with TIR functions."""
	module = TestPyFuncModule

	# Create instance
	device = tvm.cpu(0)
	instance = module(device)

	# Test TIR function execution
	n = 5
	input_tensor = torch.randn(n, dtype=torch.float32)

	# Call TIR function - it needs 3 arguments: input, output, and size
	# But call_tir handles the output buffer creation, so we only pass input and size
	# Note: TIR functions expect TVM types, not Python types
	result = instance.call_tir(
	instance.simple_tir_func,
	[input_tensor], # Only pass input tensor, let call_tir handle the rest
	R.Tensor((n,), "float32"),
	)

	# Verify result
	assert isinstance(result, torch.Tensor)
	assert result.shape == (n,)
	assert torch.allclose(result, input_tensor, atol=1e-5)

	def test_pyfunc_decorator_preserves_function_signatures(self):
	"""Test that @I.pyfunc decorator preserves function signatures."""
	module = TestPyFuncModule
	pyfuncs = module.pyfuncs

	# Check function signatures
	import inspect

	# pytorch_processor signature
	processor_func = pyfuncs["pytorch_processor"]
	sig = inspect.signature(processor_func)
	params = list(sig.parameters.keys())
	assert len(params) == 1, "pytorch_processor should have 1 parameter"
	assert params[0] == "x", "First parameter should be 'x'"

	# pytorch_adder signature
	adder_func = pyfuncs["pytorch_adder"]
	sig = inspect.signature(adder_func)
	params = list(sig.parameters.keys())
	assert len(params) == 2, "pytorch_adder should have 2 parameters"
	assert params[0] == "x", "First parameter should be 'x'"
	assert params[1] == "y", "Second parameter should be 'y'"


	if __name__ == "__main__":
	pytest.main([__file__])