tests/python/utils/test_embed_cubin.py - tvm-ffi - 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.
 """Tests for the embed_cubin utility."""

 from __future__ import annotations

 import subprocess
 import sys
 import tempfile
 from pathlib import Path

 import pytest
 from tvm_ffi.utils.embed_cubin import embed_cubin

 try:
     import torch
 except ImportError:
     torch = None  # ty: ignore[invalid-assignment]


 def _is_cuda_available() -> bool:
     """Check if CUDA is available for testing."""
     if torch is None:
         return False
     return torch.cuda.is_available()


 def _create_test_object_file(obj_path: Path) -> None:
     """Create a simple test object file with TVM_FFI_EMBED_CUBIN macro usage.

     We avoid including the full header to prevent CUDA dependency in tests.
     Instead, we manually declare the external symbols that TVM_FFI_EMBED_CUBIN would create.
     """
     cpp_code = """
     // Manually declare the external symbols that will be provided by embed_cubin
     // This is what TVM_FFI_EMBED_CUBIN(test_cubin) would do, but without CUDA headers
     extern "C" {
       extern const unsigned char __tvm_ffi__cubin_test_cubin[];
       extern const unsigned char __tvm_ffi__cubin_test_cubin_end[];
     }

     // Simple function that references the embedded CUBIN symbols
     const void* get_cubin_start() {
       return __tvm_ffi__cubin_test_cubin;
     }

     const void* get_cubin_end() {
       return __tvm_ffi__cubin_test_cubin_end;
     }
     """

     # Write C++ code to a temporary file
     cpp_file = obj_path.with_suffix(".cc")
     cpp_file.write_text(cpp_code)

     # Compile to object file
     # Try to find a C++ compiler
     compilers = ["g++", "clang++", "c++"]
     compiler = None
     for c in compilers:
         try:
             subprocess.run(
                 [c, "--version"],
                 check=True,
                 capture_output=True,
             )
             compiler = c
             break
         except (subprocess.CalledProcessError, FileNotFoundError):
             continue

     if compiler is None:
         pytest.skip("No C++ compiler found (tried g++, clang++, c++)")  # ty: ignore[invalid-argument-type, too-many-positional-arguments]

     assert isinstance(compiler, str), "Compiler is not a string"

     compile_cmd: list[str] = [compiler, "-c", "-fPIC", "-o", str(obj_path), str(cpp_file)]

     try:
         subprocess.run(compile_cmd, check=True, capture_output=True)
     except subprocess.CalledProcessError as e:
         pytest.skip(f"Failed to compile test object file: {e.stderr.decode('utf-8')}")  # ty: ignore[invalid-argument-type, too-many-positional-arguments]
     finally:
         # Clean up temporary C++ file
         cpp_file.unlink(missing_ok=True)


 def _create_test_cubin(cubin_path: Path) -> None:
     """Create a dummy CUBIN file for testing."""
     # Create a simple binary file with some recognizable content
     cubin_path.write_bytes(b"DUMMY_CUBIN_DATA_FOR_TESTING" * 10)


 def _check_symbols_in_object(obj_path: Path, expected_symbols: list[str]) -> bool:
     """Check if expected symbols exist in the object file using nm."""
     try:
         result = subprocess.run(
             ["nm", str(obj_path)],
             check=True,
             capture_output=True,
             text=True,
         )
         nm_output = result.stdout

         for symbol in expected_symbols:
             if symbol not in nm_output:
                 return False
         return True
     except (subprocess.CalledProcessError, FileNotFoundError):
         pytest.skip("nm tool not available")  # ty: ignore[invalid-argument-type, too-many-positional-arguments]
         return False


 @pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux")
 @pytest.mark.skipif(torch is None, reason="PyTorch not installed")
 @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available")
 def test_embed_cubin_basic() -> None:
     """Test basic embed_cubin functionality."""
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_path = Path(temp_dir)

         # Create test files
         input_obj = temp_path / "input.o"
         output_obj = temp_path / "output.o"
         cubin_file = temp_path / "test.cubin"

         _create_test_object_file(input_obj)
         _create_test_cubin(cubin_file)

         # Embed CUBIN into object file
         embed_cubin(
             cubin_path=cubin_file,
             input_obj_path=input_obj,
             output_obj_path=output_obj,
             name="test_cubin",
             verbose=False,
         )

         # Check that output file exists
         assert output_obj.exists(), "Output object file was not created"
         assert output_obj.stat().st_size > 0, "Output object file is empty"

         # Check that the output is larger than the input (it should contain both)
         assert output_obj.stat().st_size > input_obj.stat().st_size, (
             "Output object file should be larger than input"
         )

         # Check that expected symbols are present
         expected_symbols = [
             "__tvm_ffi__cubin_test_cubin",
             "__tvm_ffi__cubin_test_cubin_end",
         ]
         assert _check_symbols_in_object(output_obj, expected_symbols), (
             f"Expected symbols {expected_symbols} not found in output object file"
         )


 @pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux")
 @pytest.mark.skipif(torch is None, reason="PyTorch not installed")
 @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available")
 def test_embed_cubin_different_names() -> None:
     """Test embedding CUBIN with different names."""
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_path = Path(temp_dir)

         # Create test files
         input_obj = temp_path / "input.o"
         output_obj = temp_path / "output.o"
         cubin_file = temp_path / "test.cubin"

         _create_test_object_file(input_obj)
         _create_test_cubin(cubin_file)

         # Test with a different name
         custom_name = "my_custom_kernel"
         embed_cubin(
             cubin_path=cubin_file,
             input_obj_path=input_obj,
             output_obj_path=output_obj,
             name=custom_name,
             verbose=False,
         )

         # Check that symbols with custom name are present
         expected_symbols = [
             f"__tvm_ffi__cubin_{custom_name}",
             f"__tvm_ffi__cubin_{custom_name}_end",
         ]
         assert _check_symbols_in_object(output_obj, expected_symbols), (
             f"Expected symbols {expected_symbols} not found in output object file"
         )


 def test_embed_cubin_nonexistent_input() -> None:
     """Test that embed_cubin raises error for nonexistent input files."""
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_path = Path(temp_dir)

         input_obj = temp_path / "nonexistent_input.o"
         output_obj = temp_path / "output.o"
         cubin_file = temp_path / "nonexistent.cubin"

         # Test with nonexistent CUBIN file
         with pytest.raises(FileNotFoundError):
             embed_cubin(
                 cubin_path=cubin_file,
                 input_obj_path=input_obj,
                 output_obj_path=output_obj,
                 name="test",
                 verbose=False,
             )


 @pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux")
 @pytest.mark.skipif(torch is None, reason="PyTorch not installed")
 @pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available")
 def test_embed_cubin_verbose_mode() -> None:
     """Test that verbose mode doesn't crash."""
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_path = Path(temp_dir)

         # Create test files
         input_obj = temp_path / "input.o"
         output_obj = temp_path / "output.o"
         cubin_file = temp_path / "test.cubin"

         _create_test_object_file(input_obj)
         _create_test_cubin(cubin_file)

         # Run with verbose mode (should not crash)
         embed_cubin(
             cubin_path=cubin_file,
             input_obj_path=input_obj,
             output_obj_path=output_obj,
             name="test_cubin",
             verbose=True,  # Test verbose mode
         )

         assert output_obj.exists(), "Output object file was not created in verbose mode"
	# 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 embed_cubin utility."""

	from __future__ import annotations

	import subprocess
	import sys
	import tempfile
	from pathlib import Path

	import pytest
	from tvm_ffi.utils.embed_cubin import embed_cubin

	try:
	import torch
	except ImportError:
	torch = None # ty: ignore[invalid-assignment]


	def _is_cuda_available() -> bool:
	"""Check if CUDA is available for testing."""
	if torch is None:
	return False
	return torch.cuda.is_available()


	def _create_test_object_file(obj_path: Path) -> None:
	"""Create a simple test object file with TVM_FFI_EMBED_CUBIN macro usage.

	We avoid including the full header to prevent CUDA dependency in tests.
	Instead, we manually declare the external symbols that TVM_FFI_EMBED_CUBIN would create.
	"""
	cpp_code = """
	// Manually declare the external symbols that will be provided by embed_cubin
	// This is what TVM_FFI_EMBED_CUBIN(test_cubin) would do, but without CUDA headers
	extern "C" {
	extern const unsigned char __tvm_ffi__cubin_test_cubin[];
	extern const unsigned char __tvm_ffi__cubin_test_cubin_end[];
	}

	// Simple function that references the embedded CUBIN symbols
	const void* get_cubin_start() {
	return __tvm_ffi__cubin_test_cubin;
	}

	const void* get_cubin_end() {
	return __tvm_ffi__cubin_test_cubin_end;
	}
	"""

	# Write C++ code to a temporary file
	cpp_file = obj_path.with_suffix(".cc")
	cpp_file.write_text(cpp_code)

	# Compile to object file
	# Try to find a C++ compiler
	compilers = ["g++", "clang++", "c++"]
	compiler = None
	for c in compilers:
	try:
	subprocess.run(
	[c, "--version"],
	check=True,
	capture_output=True,
	)
	compiler = c
	break
	except (subprocess.CalledProcessError, FileNotFoundError):
	continue

	if compiler is None:
	pytest.skip("No C++ compiler found (tried g++, clang++, c++)") # ty: ignore[invalid-argument-type, too-many-positional-arguments]

	assert isinstance(compiler, str), "Compiler is not a string"

	compile_cmd: list[str] = [compiler, "-c", "-fPIC", "-o", str(obj_path), str(cpp_file)]

	try:
	subprocess.run(compile_cmd, check=True, capture_output=True)
	except subprocess.CalledProcessError as e:
	pytest.skip(f"Failed to compile test object file: {e.stderr.decode('utf-8')}") # ty: ignore[invalid-argument-type, too-many-positional-arguments]
	finally:
	# Clean up temporary C++ file
	cpp_file.unlink(missing_ok=True)


	def _create_test_cubin(cubin_path: Path) -> None:
	"""Create a dummy CUBIN file for testing."""
	# Create a simple binary file with some recognizable content
	cubin_path.write_bytes(b"DUMMY_CUBIN_DATA_FOR_TESTING" * 10)


	def _check_symbols_in_object(obj_path: Path, expected_symbols: list[str]) -> bool:
	"""Check if expected symbols exist in the object file using nm."""
	try:
	result = subprocess.run(
	["nm", str(obj_path)],
	check=True,
	capture_output=True,
	text=True,
	)
	nm_output = result.stdout

	for symbol in expected_symbols:
	if symbol not in nm_output:
	return False
	return True
	except (subprocess.CalledProcessError, FileNotFoundError):
	pytest.skip("nm tool not available") # ty: ignore[invalid-argument-type, too-many-positional-arguments]
	return False


	@pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux")
	@pytest.mark.skipif(torch is None, reason="PyTorch not installed")
	@pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available")
	def test_embed_cubin_basic() -> None:
	"""Test basic embed_cubin functionality."""
	with tempfile.TemporaryDirectory() as temp_dir:
	temp_path = Path(temp_dir)

	# Create test files
	input_obj = temp_path / "input.o"
	output_obj = temp_path / "output.o"
	cubin_file = temp_path / "test.cubin"

	_create_test_object_file(input_obj)
	_create_test_cubin(cubin_file)

	# Embed CUBIN into object file
	embed_cubin(
	cubin_path=cubin_file,
	input_obj_path=input_obj,
	output_obj_path=output_obj,
	name="test_cubin",
	verbose=False,
	)

	# Check that output file exists
	assert output_obj.exists(), "Output object file was not created"
	assert output_obj.stat().st_size > 0, "Output object file is empty"

	# Check that the output is larger than the input (it should contain both)
	assert output_obj.stat().st_size > input_obj.stat().st_size, (
	"Output object file should be larger than input"
	)

	# Check that expected symbols are present
	expected_symbols = [
	"__tvm_ffi__cubin_test_cubin",
	"__tvm_ffi__cubin_test_cubin_end",
	]
	assert _check_symbols_in_object(output_obj, expected_symbols), (
	f"Expected symbols {expected_symbols} not found in output object file"
	)


	@pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux")
	@pytest.mark.skipif(torch is None, reason="PyTorch not installed")
	@pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available")
	def test_embed_cubin_different_names() -> None:
	"""Test embedding CUBIN with different names."""
	with tempfile.TemporaryDirectory() as temp_dir:
	temp_path = Path(temp_dir)

	# Create test files
	input_obj = temp_path / "input.o"
	output_obj = temp_path / "output.o"
	cubin_file = temp_path / "test.cubin"

	_create_test_object_file(input_obj)
	_create_test_cubin(cubin_file)

	# Test with a different name
	custom_name = "my_custom_kernel"
	embed_cubin(
	cubin_path=cubin_file,
	input_obj_path=input_obj,
	output_obj_path=output_obj,
	name=custom_name,
	verbose=False,
	)

	# Check that symbols with custom name are present
	expected_symbols = [
	f"__tvm_ffi__cubin_{custom_name}",
	f"__tvm_ffi__cubin_{custom_name}_end",
	]
	assert _check_symbols_in_object(output_obj, expected_symbols), (
	f"Expected symbols {expected_symbols} not found in output object file"
	)


	def test_embed_cubin_nonexistent_input() -> None:
	"""Test that embed_cubin raises error for nonexistent input files."""
	with tempfile.TemporaryDirectory() as temp_dir:
	temp_path = Path(temp_dir)

	input_obj = temp_path / "nonexistent_input.o"
	output_obj = temp_path / "output.o"
	cubin_file = temp_path / "nonexistent.cubin"

	# Test with nonexistent CUBIN file
	with pytest.raises(FileNotFoundError):
	embed_cubin(
	cubin_path=cubin_file,
	input_obj_path=input_obj,
	output_obj_path=output_obj,
	name="test",
	verbose=False,
	)


	@pytest.mark.skipif(sys.platform != "linux", reason="embed_cubin only supported on Linux")
	@pytest.mark.skipif(torch is None, reason="PyTorch not installed")
	@pytest.mark.skipif(not _is_cuda_available(), reason="CUDA not available")
	def test_embed_cubin_verbose_mode() -> None:
	"""Test that verbose mode doesn't crash."""
	with tempfile.TemporaryDirectory() as temp_dir:
	temp_path = Path(temp_dir)

	# Create test files
	input_obj = temp_path / "input.o"
	output_obj = temp_path / "output.o"
	cubin_file = temp_path / "test.cubin"

	_create_test_object_file(input_obj)
	_create_test_cubin(cubin_file)

	# Run with verbose mode (should not crash)
	embed_cubin(
	cubin_path=cubin_file,
	input_obj_path=input_obj,
	output_obj_path=output_obj,
	name="test_cubin",
	verbose=True, # Test verbose mode
	)

	assert output_obj.exists(), "Output object file was not created in verbose mode"