Makefile

#  ParaGrad — Hardware-Agnostic Parallel Autograd Engine
#
#  make                            build everything (cpu)
#  make BACKEND=cuda               build with CUDA backend
#  make BACKEND=cublas             build with cuBLAS backend
#  make test                       unit + integration tests
#  make debug                      ASan build + unit tests
#  make bench_all BACKEND=cuda     full GPU benchmark suite
#  make info                       show resolved config
#  make help                       this list
#

#  Toolchain 
CXX       ?= g++
BACKEND   ?= cpu
CUDA_ARCH ?= sm_75

#  CUDA auto-detection (only needed for cuda/cublas) 
# Priority: CUDA_DIR > CUDA_HOME > CUDA_PATH > nvcc on PATH > common paths
ifneq ($(filter cuda cublas,$(BACKEND)),)

  ifndef CUDA_DIR
    ifdef CUDA_HOME
      CUDA_DIR := $(CUDA_HOME)
    else ifdef CUDA_PATH
      CUDA_DIR := $(CUDA_PATH)
    else ifneq ($(shell which nvcc 2>/dev/null),)
      CUDA_DIR := $(realpath $(dir $(shell which nvcc))/..)
    else ifneq ($(wildcard /usr/local/cuda),)
      CUDA_DIR := /usr/local/cuda
    else ifneq ($(wildcard /opt/cuda),)
      CUDA_DIR := /opt/cuda
    endif
  endif

  ifndef CUDA_DIR
    $(error CUDA not found. Install CUDA or set CUDA_DIR / CUDA_HOME)
  endif
  ifeq ($(wildcard $(CUDA_DIR)/include/cuda_runtime.h),)
    $(error CUDA_DIR=$(CUDA_DIR) invalid — missing include/cuda_runtime.h)
  endif

  # lib64 on most Linux, lib on some (Jetson, conda, macOS)
  ifneq ($(wildcard $(CUDA_DIR)/lib64/libcudart.so),)
    CUDA_LIB := $(CUDA_DIR)/lib64
  else ifneq ($(wildcard $(CUDA_DIR)/lib/x86_64-linux-gnu/libcudart.so),)
    CUDA_LIB := $(CUDA_DIR)/lib/x86_64-linux-gnu
  else ifneq ($(wildcard $(CUDA_DIR)/lib/x86_64-linux-gnu/libcudart.so),)
    CUDA_LIB := $(CUDA_DIR)/lib/x86_64-linux-gnu
  else ifneq ($(wildcard $(CUDA_DIR)/lib/libcudart.so),)
    CUDA_LIB := $(CUDA_DIR)/lib
  else ifneq ($(wildcard $(CUDA_DIR)/lib/libcudart.dylib),)
    CUDA_LIB := $(CUDA_DIR)/lib
  else
    $(error Cannot find libcudart under $(CUDA_DIR)/lib64 or $(CUDA_DIR)/lib)
  endif

  NVCC     := $(CUDA_DIR)/bin/nvcc
  CUDA_LIBS := -lcudart

  # cuda_fusion.cu JIT-compiles fused kernels via NVRTC + driver API
  ifneq ($(wildcard $(CUDA_LIB)/libnvrtc.so),)
    CUDA_LIBS += -lnvrtc -lcuda
  else ifneq ($(wildcard $(CUDA_LIB)/libnvrtc.dylib),)
    CUDA_LIBS += -lnvrtc -lcuda
  else
    $(warning libnvrtc not found in $(CUDA_LIB) — fused CUDA kernels will fail to JIT)
    CUDA_LIBS += -lcuda
  endif

endif  # cuda/cublas

#  Directories 
BUILD   = build/$(BACKEND)
BUILD_PARENT = build
CPU_DIR = src/backends/cpu
CUD_DIR = src/backends/cuda

#  Compiler flags 
CXXFLAGS  = -std=c++17 -Wall -Wextra -O2 -Iinclude -I$(CPU_DIR) -Itests/unit -fopenmp
NVCCFLAGS = -std=c++17 -O2 -Iinclude -I$(CPU_DIR) -I$(CUD_DIR) \
            --expt-relaxed-constexpr -arch=$(CUDA_ARCH)
LDFLAGS   = -fopenmp

ifeq ($(BACKEND),cuda)
  CXXFLAGS  += -DPARAGRAD_USE_CUDA -I$(CUD_DIR) -I$(CUDA_DIR)/include
  NVCCFLAGS += -DPARAGRAD_USE_CUDA
  LDFLAGS   += -L$(CUDA_LIB) $(CUDA_LIBS) -Wl,-rpath,$(CUDA_LIB)
endif
ifeq ($(BACKEND),cublas)
  CXXFLAGS  += -DPARAGRAD_USE_CUDA -DPARAGRAD_USE_CUBLAS -I$(CUD_DIR) -I$(CUDA_DIR)/include
  NVCCFLAGS += -DPARAGRAD_USE_CUDA -DPARAGRAD_USE_CUBLAS
  LDFLAGS   += -L$(CUDA_LIB) $(CUDA_LIBS) -lcublas -Wl,-rpath,$(CUDA_LIB)
endif

#  Object files 
OBJS = $(BUILD)/tensor.o $(BUILD)/tape.o $(BUILD)/fusion.o $(BUILD)/cpu_backend.o \
       $(BUILD)/diagnostics.o
ifeq ($(BACKEND),cuda)
  OBJS += $(BUILD)/cuda_backend.o $(BUILD)/cuda_fusion.o
endif
ifeq ($(BACKEND),cublas)
  OBJS += $(BUILD)/cuda_backend.o $(BUILD)/cuda_fusion.o $(BUILD)/cublas_backend.o
endif

#  Headers (rebuild tracking) 
HDRS = include/paragrad/tensor.h include/paragrad/tape.h \
       include/paragrad/backend.h include/paragrad/fusion.h \
       include/paragrad/nn.h include/paragrad/log.h \
       include/paragrad/diagnostics.h \
       $(CPU_DIR)/cpu_backend.h tests/unit/test_helpers.h
ifneq ($(filter cuda cublas,$(BACKEND)),)
  HDRS += $(CUD_DIR)/cuda_backend.h
endif
ifeq ($(BACKEND),cublas)
  HDRS += $(CUD_DIR)/cublas_backend.h
endif

#  Target lists 
UNIT_NAMES   = test_tensor test_unary_ops test_binary_ops test_matmul \
               test_nn test_fusion test_edge_cases test_transformer_ops \
               test_fusion_regression test_memory_safety_regression \
               test_diagnostics
INTEG_NAMES  = test_training test_backend_consistency
UNIT_BINS    = $(addprefix $(BUILD)/, $(UNIT_NAMES))
INTEG_BINS   = $(addprefix $(BUILD)/, $(INTEG_NAMES))

EXAMPLES = $(BUILD)/train_lm $(BUILD)/generate_lm $(BUILD)/mnist_train
BENCHES  = $(BUILD)/bench_backends $(BUILD)/bench_fusion $(BUILD)/bench_training

#  Phony targets
.PHONY: all all_backends test test_unit test_integration examples benches debug clean purge \
        bench bench_fusion bench_training bench_all \
        update-benchmarks info help check

all: $(UNIT_BINS) $(INTEG_BINS) $(EXAMPLES) $(BENCHES)
	@echo ""
	@echo "Build complete  [BACKEND=$(BACKEND)]"
	@echo "  tests:    make test"
	@echo "  examples: make examples"
	@echo "  benches:  make benches"

#  Preflight check 
check:
	@echo "Checking build environment..."
	@$(CXX) --version | head -1
	@echo "  OpenMP: $$(echo 'int main(){return 0;}' | $(CXX) -fopenmp -x c++ - -o /dev/null 2>&1 && echo yes || echo no)"
ifneq ($(filter cuda cublas,$(BACKEND)),)
	@$(NVCC) --version | grep release
	@echo "  CUDA_DIR: $(CUDA_DIR)"
	@echo "  CUDA_LIB: $(CUDA_LIB)"
	@echo "  CUDA_ARCH: $(CUDA_ARCH)"
	@nvidia-smi --query-gpu=name,driver_version --format=csv,noheader 2>/dev/null || echo "  (nvidia-smi not available)"
endif
	@echo "Ready."

#  Build directory 
$(BUILD):
	mkdir -p $(BUILD)

#  Object compilation

# Core (C++)
$(BUILD)/tensor.o: src/tensor.cpp $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) -c $< -o $@

$(BUILD)/tape.o: src/tape.cpp $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) -c $< -o $@

$(BUILD)/fusion.o: src/fusion.cpp $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) -c $< -o $@

$(BUILD)/diagnostics.o: src/diagnostics.cpp $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) -c $< -o $@

$(BUILD)/cpu_backend.o: $(CPU_DIR)/cpu_backend.cpp $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) -c $< -o $@

# CUDA (.cu)
$(BUILD)/cuda_backend.o: $(CUD_DIR)/cuda_backend.cu $(HDRS) | $(BUILD)
	$(NVCC) $(NVCCFLAGS) -c $< -o $@

$(BUILD)/cuda_fusion.o: $(CUD_DIR)/cuda_fusion.cu $(HDRS) | $(BUILD)
	$(NVCC) $(NVCCFLAGS) -c $< -o $@

$(BUILD)/cublas_backend.o: $(CUD_DIR)/cublas_backend.cu $(HDRS) | $(BUILD)
	$(NVCC) $(NVCCFLAGS) -c $< -o $@

#  Tests

$(BUILD)/test_%: tests/unit/test_%.cpp $(OBJS) $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) $< $(OBJS) $(LDFLAGS) -o $@

$(BUILD)/test_training: tests/integration/test_training.cpp $(OBJS) $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) $< $(OBJS) $(LDFLAGS) -o $@

$(BUILD)/test_backend_consistency: tests/integration/test_backend_consistency.cpp $(OBJS) $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) $< $(OBJS) $(LDFLAGS) -o $@

define run_suite
	@failed=0; \
	for t in $(1); do \
		printf "\n── %s " "$$(basename $$t)"; \
		if $$t > /dev/null 2>&1; then printf "✓\n"; \
		else printf "✗\n"; $$t; failed=1; fi; \
	done; \
	echo ""; \
	if [ $$failed -eq 0 ]; then echo "  $(2): all passed"; \
	else echo "  $(2): FAILURES"; exit 1; fi
endef

test_unit: $(UNIT_BINS)
	@echo "\n════ unit tests [$(BACKEND)] ════"
	$(call run_suite,$(UNIT_BINS),unit)

test_integration: $(INTEG_BINS)
	@echo "\n════ integration tests [$(BACKEND)] ════"
	$(call run_suite,$(INTEG_BINS),integration)

test: test_unit test_integration

#  Debug / sanitizer build 
debug: CXXFLAGS += -DPARAGRAD_DEBUG -g -O0 -fsanitize=address
debug: LDFLAGS  += -fsanitize=address
debug: clean $(UNIT_BINS)
	@echo "\n════ debug tests [$(BACKEND)] ════"
	@for t in $(UNIT_BINS); do \
		echo "\n── $$(basename $$t)"; $$t || true; \
	done

#  Examples

$(BUILD)/train_lm: examples/train_lm.cpp $(OBJS) $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) $< $(OBJS) $(LDFLAGS) -o $@

$(BUILD)/generate_lm: examples/generate_lm.cpp $(OBJS) $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) $< $(OBJS) $(LDFLAGS) -o $@

$(BUILD)/mnist_train: examples/mnist_train.cpp $(OBJS) $(HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) $< $(OBJS) $(LDFLAGS) -o $@

examples: $(EXAMPLES)
	@echo "Built examples [$(BACKEND)]"
	@echo "  ./$(BUILD)/train_lm    <text_file> [options]"
	@echo "  ./$(BUILD)/generate_lm <checkpoint.pgr> [options]"
	@echo "  ./$(BUILD)/mnist_train"

#  Benchmarks
BENCH_HDRS = benchmarks/bench_utils.h $(HDRS)

$(BUILD)/bench_%: benchmarks/bench_%.cpp $(OBJS) $(BENCH_HDRS) | $(BUILD)
	$(CXX) $(CXXFLAGS) -Ibenchmarks $< $(OBJS) $(LDFLAGS) -o $@

benches: $(BENCHES)
	@echo "Built benchmarks [$(BACKEND)]"

bench:          $(BUILD)/bench_backends  ; ./$(BUILD)/bench_backends
bench_fusion:   $(BUILD)/bench_fusion    ; ./$(BUILD)/bench_fusion
bench_training: $(BUILD)/bench_training  ; ./$(BUILD)/bench_training
bench_all: $(BENCHES)
	@chmod +x benchmarks/run_all_benchmarks.sh
	./benchmarks/run_all_benchmarks.sh $(BACKEND)

#  Info / help

info:
	@echo "BACKEND   = $(BACKEND)"
	@echo "CXX       = $(CXX)"
	@echo "CXXFLAGS  = $(CXXFLAGS)"
ifneq ($(filter cuda cublas,$(BACKEND)),)
	@echo "CUDA_DIR  = $(CUDA_DIR)"
	@echo "CUDA_LIB  = $(CUDA_LIB)"
	@echo "CUDA_LIBS = $(CUDA_LIBS)"
	@echo "CUDA_ARCH = $(CUDA_ARCH)"
	@echo "NVCC      = $(NVCC)"
endif
	@echo "LDFLAGS   = $(LDFLAGS)"

help:
	@echo "ParaGrad build system"
	@echo ""
	@echo "  make [BACKEND=cpu|cuda|cublas]   build all targets"
	@echo "  make test                        run all tests"
	@echo "  make test_unit                   unit tests only"
	@echo "  make test_integration            integration tests only"
	@echo "  make debug                       ASan build + unit tests"
	@echo "  make examples                    build example programs"
	@echo "  make benches                     build benchmarks"
	@echo "  make bench                       backend comparison (transformer shapes)"
	@echo "  make bench_fusion                fusion speedup (chains + epilogue)"
	@echo "  make bench_training              end-to-end training throughput"
	@echo "  make bench_all                   build + run all benchmarks"
	@echo "  make update-benchmarks           run benchmarks + update README"
	@echo "  make check                       verify build environment"
	@echo "  make info                        show resolved config"
	@echo "  make clean                       remove build artifacts"
	@echo ""
	@echo "Override variables:"
	@echo "  BACKEND    cpu (default), cuda, cublas"
	@echo "  CUDA_DIR   path to CUDA toolkit (auto-detected)"
	@echo "  CUDA_ARCH  GPU architecture (default: sm_75)"
	@echo "  CXX        C++ compiler (default: g++)"

#  Benchmark + update README (run locally with GPU)
update-benchmarks:
	@chmod +x scripts/bench_and_update.sh
	./scripts/bench_and_update.sh --commit

clean:
	rm -rf $(BUILD_PARENT)/*

# Build all three backends sequentially.
all_backends:
	$(MAKE) all BACKEND=cpu
	$(MAKE) all BACKEND=cuda
	$(MAKE) all BACKEND=cublas
	@echo ""
	@echo "All backends built."
	@echo "  cpu    - build/cpu/"
	@echo "  cuda   - build/cuda/"
	@echo "  cublas - build/cublas/"