File size: 2,088 Bytes
2d8a802 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import os
import pathlib
import sys
import signal
import time
from torch.utils.cpp_extension import load
def timeout_handler(signum, frame):
print("Build timed out - this indicates a hanging issue")
sys.exit(1)
# Set up timeout
signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(180) # 3 minute timeout
repo = pathlib.Path(".").resolve()
os.environ.setdefault("TORCH_EXTENSIONS_DIR", str(repo / ".torch_extensions_debug"))
print("=== Testing with Single Source File ===")
try:
print("Building with just new_cumsum.cu...")
mod = load(
name="_megablocks_debug_single",
sources=["csrc/new_cumsum.cu"],
extra_include_paths=["csrc"],
extra_cflags=["-O3", "-std=c++17"],
extra_cuda_cflags=["-O3"],
verbose=True,
is_python_module=False,
)
print("✓ Single source build successful")
except Exception as e:
print(f"✗ Single source build failed: {e}")
print("\n=== Testing with Two Source Files ===")
try:
print("Building with new_cumsum.cu and new_histogram.cu...")
mod = load(
name="_megablocks_debug_double",
sources=["csrc/new_cumsum.cu", "csrc/new_histogram.cu"],
extra_include_paths=["csrc"],
extra_cflags=["-O3", "-std=c++17"],
extra_cuda_cflags=["-O3"],
verbose=True,
is_python_module=False,
)
print("✓ Double source build successful")
except Exception as e:
print(f"✗ Double source build failed: {e}")
print("\n=== Testing with grouped_gemm.cu Only ===")
try:
print("Building with just grouped_gemm.cu (most complex)...")
mod = load(
name="_megablocks_debug_gemm",
sources=["csrc/grouped_gemm/grouped_gemm.cu"],
extra_include_paths=["csrc"],
extra_cflags=["-O3", "-std=c++17"],
extra_cuda_cflags=["-O3"],
extra_ldflags=["-lhipblaslt"],
verbose=True,
is_python_module=False,
)
print("✓ grouped_gemm build successful")
except Exception as e:
print(f"✗ grouped_gemm build failed: {e}")
signal.alarm(0) # Cancel timeout
|