extconf.rb

require 'mkmf'
require 'rbconfig'

# Logging level
$VERBOSE = true

host_os   = RbConfig::CONFIG['host_os']
on_macos  = host_os =~ /darwin/
on_linux  = host_os =~ /linux/

# Brief dependency summary for users on all platforms
puts "\n== tensor native extension build =="
puts "Host OS: #{host_os}"
puts "This extension can use (optionally):"
puts "  - OpenMP (libomp/libgomp) for multi-threading"
puts "  - OpenBLAS (openblas) for BLAS optimizations"
puts "Environment overrides:"
puts "  - OPENBLAS_DIR  : root with include/ and lib/ for OpenBLAS"
puts "  - OMP_INCLUDE_DIR / OMP_LIB_DIR: OpenMP headers/libs"
puts "  - PKG_CONFIG     : path to pkg-config\n"

# Debug environment variables
puts "PKG_CONFIG_PATH: #{ENV['PKG_CONFIG_PATH']}"
puts "PATH:           #{ENV['PATH']}"
puts "pkg-config:     #{`which pkg-config`.chomp}"

# Set absolute path to pkg-config if needed (primarily for Homebrew on macOS)
if on_macos && ENV['PKG_CONFIG'].to_s.empty?
  ENV['PKG_CONFIG'] = '/opt/homebrew/bin/pkg-config'
end
puts "Using pkg-config at: #{ENV['PKG_CONFIG']}"

# Check for required headers
puts "Checking for ruby.h..."
unless have_header('ruby.h')
  abort "ruby.h not found. Please ensure Ruby development headers are installed."
end

# Detect architecture
is_arm = RbConfig::CONFIG['host_cpu'] =~ /arm|aarch64/
is_x86 = RbConfig::CONFIG['host_cpu'] =~ /x86|x64|i386/

# Detect if we're using Apple's clang
is_apple_clang = RbConfig::CONFIG['CC'] =~ /^clang/ && RUBY_PLATFORM =~ /darwin/

# OpenMP support (optional)
omp_include_dir = ENV['OMP_INCLUDE_DIR']
omp_lib_dir     = ENV['OMP_LIB_DIR']

if omp_include_dir && !omp_include_dir.empty?
  $INCFLAGS << " -I#{omp_include_dir}"
end
if omp_lib_dir && !omp_lib_dir.empty?
  $LDFLAGS << " -L#{omp_lib_dir}"
end

openmp_enabled = false

if is_apple_clang
  puts "Detected Apple's clang, trying LLVM OpenMP..."
  libomp_prefix = "/opt/homebrew/opt/libomp"
  if Dir.exist?(libomp_prefix)
    $INCFLAGS << " -I#{libomp_prefix}/include"
    $LDFLAGS << " -L#{libomp_prefix}/lib"
    $CFLAGS << " -Xpreprocessor -fopenmp"
    $LDFLAGS << " -lomp"
    openmp_enabled = have_library('omp')
    puts "LLVM OpenMP support #{openmp_enabled ? 'enabled' : 'not found'}"
    unless openmp_enabled
      puts "To enable: brew install libomp"
    end
  else
    puts "LLVM libomp not found under /opt/homebrew/opt/libomp."
    puts "To enable OpenMP on macOS: brew install libomp"
  end
else
  # Standard OpenMP support for non-Apple compilers
  if have_library('omp') || have_library('gomp')
    $CFLAGS << " -fopenmp"
    $LDFLAGS << " -fopenmp"
    openmp_enabled = true
    puts "OpenMP support enabled via -fopenmp"
  else
    puts "OpenMP library (libomp/libgomp) not found."
    if on_linux
      puts "On Debian/Ubuntu: sudo apt-get install libomp-dev (or ensure gcc's libgomp is available)."
      puts "On Fedora/RHEL:   sudo dnf install libgomp"
    end
  end
end

# Optional OpenBLAS support
openblas_include_path = nil
openblas_lib_path     = nil

explicit_openblas = ENV['OPENBLAS_DIR']
if explicit_openblas && Dir.exist?(explicit_openblas)
  openblas_root = explicit_openblas
elsif on_macos
  openblas_versions = Dir.glob("/opt/homebrew/Cellar/openblas/*").sort
  openblas_root = openblas_versions.last
else
  openblas_root = nil
end

if openblas_root && Dir.exist?(openblas_root)
  openblas_include_path = File.join(openblas_root, "include")
  openblas_lib_path     = File.join(openblas_root, "lib")
  $INCFLAGS << " -I#{openblas_include_path}"
  $LDFLAGS << " -L#{openblas_lib_path}"
  $LIBPATH << openblas_lib_path
  puts "Using OpenBLAS from #{openblas_root}"
elsif explicit_openblas
  puts "OPENBLAS_DIR=#{explicit_openblas} was set but does not exist; ignoring."
end

# Check for BLAS support using pkg-config (cross-platform)
puts "Checking for BLAS (OpenBLAS) support..."
if pkg_config("openblas")
  puts "BLAS support found via pkg-config! Enabling BLAS optimizations."
elsif openblas_lib_path && Dir.exist?(openblas_lib_path)
  puts "pkg-config failed to find OpenBLAS. Using manual configuration from #{openblas_lib_path}."
  $LDFLAGS << " -lopenblas"
else
  puts "OpenBLAS not found. Continuing without BLAS-specific optimizations."
  if on_linux
    puts "On Debian/Ubuntu: sudo apt-get install libopenblas-dev"
    puts "On Fedora/RHEL:   sudo dnf install openblas-devel"
  end
  puts "Or set OPENBLAS_DIR to your OpenBLAS installation root (with include/ and lib/)."
end

# Clean up incompatible OpenMP flags from pkg-config when using Apple's clang
if is_apple_clang
  $CFLAGS = $CFLAGS.to_s.gsub(/\s-fopenmp\b/, "")
end

# Check for math library
have_library('m')

# Check for SIMD headers based on architecture
puts "Checking for SIMD headers..."
if is_x86 && have_header('immintrin.h')
  $CFLAGS << " -mavx2"
  $defs << "-DHAVE_IMMINTRIN_H"
  puts "AVX2 support found! Enabling AVX2 optimizations."
elsif is_arm && have_header('arm_neon.h')
  if RbConfig::CONFIG['host_os'] =~ /darwin/
    $CFLAGS << " -DHAVE_ARM_NEON_H"
    puts "ARM NEON support found on macOS! Enabling NEON optimizations."
  else
    $CFLAGS << " -DHAVE_ARM_NEON_H -mfpu=neon"
    puts "ARM NEON support found! Enabling NEON optimizations."
  end
  $defs << "-DHAVE_ARM_NEON_H"
else
  puts "No SIMD headers found. SIMD optimizations will not be available."
end

# Add optimization flags
$CFLAGS << " -O3 -march=native"
$CFLAGS << " -fno-common -pipe"
$CFLAGS << " -D_XOPEN_SOURCE -D_DARWIN_C_SOURCE -D_DARWIN_UNLIMITED_SELECT -D_REENTRANT"

# Additional include and library paths
if on_macos
  $INCFLAGS << " -I/opt/homebrew/include"
  $LDFLAGS << " -L/opt/homebrew/lib"
end
$INCFLAGS << " -I/usr/local/include"
$LDFLAGS << " -L/usr/local/lib"

# Ensure proper linking for neural network operations
$defs << "-DUSE_NEURAL_OPS"
$defs << "-DUSE_PARALLEL_OPS"

# Print final configuration
puts "\nFinal Configuration:"
puts "Compiler: #{RbConfig::CONFIG['CC']}"
puts "CFLAGS: #{$CFLAGS}"
puts "LDFLAGS: #{$LDFLAGS}"
puts "INCFLAGS: #{$INCFLAGS}"
puts "Definitions: #{$defs.join(' ')}"

# Create the makefile for the native extension (tensor_ext)
create_makefile('tensor_ext')