197 lines
6.6 KiB
Python
197 lines
6.6 KiB
Python
"""
|
|
Comprehensive integration test for fp16/int16 bitwise operations.
|
|
Demonstrates real-world usage of the enhanced bitwise operators.
|
|
"""
|
|
|
|
import torch
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
|
|
|
|
|
def test_mixed_dtype_operations():
|
|
"""Test operations with mixed dtypes."""
|
|
print("=" * 70)
|
|
print("Testing Mixed Dtype Operations")
|
|
print("=" * 70)
|
|
|
|
# Create a dummy context for testing
|
|
class DummyContext:
|
|
class Start:
|
|
line = 0
|
|
column = 0
|
|
start = Start()
|
|
ctx = DummyContext()
|
|
|
|
# int8 operations
|
|
print("\n1. INT8 Bitwise Operations:")
|
|
a_int8 = torch.tensor([15, 7, 3], dtype=torch.int8)
|
|
visitor = UnifiedMathVisitor({"a": a_int8})
|
|
result = visitor._bitwise_not(a_int8)
|
|
print(f" Input: {a_int8} (dtype={a_int8.dtype})")
|
|
print(f" ~Input: {result} (dtype={result.dtype})")
|
|
assert result.dtype == torch.int8
|
|
|
|
# int16 operations
|
|
print("\n2. INT16 Bitwise Operations:")
|
|
a_int16 = torch.tensor([255, 127, 63], dtype=torch.int16)
|
|
b_int16 = torch.tensor([15, 31, 7], dtype=torch.int16)
|
|
visitor = UnifiedMathVisitor({"a": a_int16, "b": b_int16})
|
|
|
|
result_and = visitor._bitwise_op(a_int16, b_int16, torch.bitwise_and, lambda x, y: x & y, ctx)
|
|
result_or = visitor._bitwise_op(a_int16, b_int16, torch.bitwise_or, lambda x, y: x | y, ctx)
|
|
result_xor = visitor._bitwise_op(a_int16, b_int16, torch.bitwise_xor, lambda x, y: x ^ y, ctx)
|
|
|
|
print(f" a: {a_int16} (dtype={a_int16.dtype})")
|
|
print(f" b: {b_int16} (dtype={b_int16.dtype})")
|
|
print(f" a & b: {result_and}")
|
|
print(f" a | b: {result_or}")
|
|
print(f" a ^ b: {result_xor}")
|
|
assert all(t.dtype == torch.int16 for t in [result_and, result_or, result_xor])
|
|
|
|
# fp16 operations
|
|
print("\n3. FP16 Bitwise Operations (bit-level manipulation):")
|
|
a_fp16 = torch.tensor([1.0, -2.0, 3.5], dtype=torch.float16)
|
|
visitor = UnifiedMathVisitor({"a": a_fp16})
|
|
result = visitor._bitwise_not(a_fp16)
|
|
print(f" Input: {a_fp16} (dtype={a_fp16.dtype})")
|
|
print(f" Bit-flipped: {result} (dtype={result.dtype})")
|
|
assert result.dtype == torch.float16
|
|
|
|
# int32 operations (existing, should still work)
|
|
print("\n4. INT32 Bitwise Operations (backward compatibility):")
|
|
a_int32 = torch.tensor([65535, 32767, 16383], dtype=torch.int32)
|
|
visitor = UnifiedMathVisitor({"a": a_int32})
|
|
result = visitor._bitwise_not(a_int32)
|
|
print(f" Input: {a_int32} (dtype={a_int32.dtype})")
|
|
print(f" ~Input: {result} (dtype={result.dtype})")
|
|
assert result.dtype == torch.int32
|
|
|
|
print("\n" + "=" * 70)
|
|
print("✓ All mixed dtype operations completed successfully!")
|
|
print("=" * 70)
|
|
|
|
|
|
def test_scalar_list_tensor_combinations():
|
|
"""Test operations with different input types."""
|
|
print("\n" + "=" * 70)
|
|
print("Testing Scalar/List/Tensor Combinations")
|
|
print("=" * 70)
|
|
|
|
# Create a dummy context for testing
|
|
class DummyContext:
|
|
class Start:
|
|
line = 0
|
|
column = 0
|
|
start = Start()
|
|
ctx = DummyContext()
|
|
|
|
visitor = UnifiedMathVisitor({})
|
|
|
|
# Tensor & Tensor
|
|
print("\n1. Tensor & Tensor (int16):")
|
|
a = torch.tensor([7, 14, 21], dtype=torch.int16)
|
|
b = torch.tensor([3, 5, 7], dtype=torch.int16)
|
|
result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
|
|
print(f" {a} & {b} = {result}")
|
|
|
|
# Tensor & List
|
|
print("\n2. Tensor & List (mixed):")
|
|
a = torch.tensor([15, 14, 13], dtype=torch.int16)
|
|
b = [7, 3, 1]
|
|
result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
|
|
print(f" Tensor({a}) & List({b}) = {result}")
|
|
|
|
# Scalar & List
|
|
print("\n3. Scalar & List:")
|
|
a = 15
|
|
b = [7, 3, 1]
|
|
result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
|
|
print(f" {a} & {b} = {result}")
|
|
|
|
# List & List
|
|
print("\n4. List & List:")
|
|
a = [15, 14, 13]
|
|
b = [7, 3, 1]
|
|
result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
|
|
print(f" {a} & {b} = {result}")
|
|
|
|
print("\n" + "=" * 70)
|
|
print("✓ All scalar/list/tensor combinations work correctly!")
|
|
print("=" * 70)
|
|
|
|
|
|
def test_element_size_mapping():
|
|
"""Test the element size to dtype mapping for all supported types."""
|
|
print("\n" + "=" * 70)
|
|
print("Testing Element Size Mapping")
|
|
print("=" * 70)
|
|
|
|
visitor = UnifiedMathVisitor({})
|
|
|
|
test_dtypes = [
|
|
(torch.int8, "int8", 1),
|
|
(torch.int16, "int16", 2),
|
|
(torch.int32, "int32", 4),
|
|
(torch.int64, "int64", 8),
|
|
(torch.float16, "float16", 2),
|
|
(torch.float32, "float32", 4),
|
|
(torch.float64, "float64", 8),
|
|
]
|
|
|
|
print("\nDtype -> Element Size -> View Dtype Mapping:")
|
|
for dtype, name, expected_elem_size in test_dtypes:
|
|
tensor = torch.zeros(1, dtype=dtype)
|
|
elem_size = tensor.element_size()
|
|
view_dtype = visitor._get_bitwise_view_dtype(elem_size)
|
|
|
|
# Determine expected view dtype
|
|
if elem_size == 1:
|
|
expected_view = "int8"
|
|
elif elem_size == 2:
|
|
expected_view = "int16"
|
|
elif elem_size == 4:
|
|
expected_view = "int32"
|
|
else:
|
|
expected_view = "int64"
|
|
|
|
print(f" {name:12} -> {elem_size} byte(s) -> {str(view_dtype):18}")
|
|
assert elem_size == expected_elem_size
|
|
|
|
print("\n" + "=" * 70)
|
|
print("✓ Element size mapping is correct!")
|
|
print("=" * 70)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("\n")
|
|
print("#" * 70)
|
|
print("# FP16/INT16 Bitwise Operations - Comprehensive Integration Test")
|
|
print("#" * 70)
|
|
|
|
try:
|
|
test_element_size_mapping()
|
|
test_mixed_dtype_operations()
|
|
test_scalar_list_tensor_combinations()
|
|
|
|
print("\n")
|
|
print("#" * 70)
|
|
print("# ✓ ALL TESTS PASSED SUCCESSFULLY")
|
|
print("#" * 70)
|
|
print("\nSummary:")
|
|
print(" ✓ fp16 (float16) bitwise operations work correctly")
|
|
print(" ✓ int16 (int16) bitwise operations work correctly")
|
|
print(" ✓ int8, int32, int64 operations maintained")
|
|
print(" ✓ Float dtypes supported through bit-pattern manipulation")
|
|
print(" ✓ Mixed tensor/list/scalar operations supported")
|
|
print(" ✓ Backward compatibility preserved")
|
|
print("\n")
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ Test failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
sys.exit(1)
|