# This Makefile will be invoked by the python build system (e.g. via 'pip install'),
# but you can also build individual targets by invoking 'make' directly.
#
# Overrideable variables (either as env variables, or in a gitignored-file 'config.mk')
# NOTE: Conda users normally need none of these (found via $CONDA_PREFIX)
#
#   NVCC_OPT       e.g. '-O0 -g' (default is '-O3')
#   NVCC           e.g. '/usr/local/bin/nvcc' (the nvcc program; its flags are NVCCFLAGS)
#   NVCCFLAGS      e.g. '-std=c++20 ...' (all of nvcc's flags; usually tweak NVCC_OPT instead)
#   NVCC_ARCH      e.g. '-gencode arch=compute_90,code=sm_90' (default targets sm_80/86/89)
#   NVCC_DEPFLAGS  e.g. '-MMD -MP' (the default)
#   PYTHON         e.g. 'python3.11' (default is 'python3')

# Disable built-in rules and variables (must be first).
MAKEFLAGS += --no-builtin-rules
MAKEFLAGS += --no-builtin-variables

# Default target 'all' must be first target in Makefile.
# The 'bin' target builds a bunch of binaries in bin/...
# The 'lib' target builds the C++ library lib/libksgpu.so, and the python extension ksgpu/ksgpu_pybind11...so.
# The 'build_wheel' and 'build_sdist' targets are invoked by 'pip' (or 'make all').
all: bin lib build_wheel build_sdist

.PHONY: all bin lib build_wheel build_sdist clean


####################################################################################################
#
# Variables encoding configuration: PYTHON, NVCC, NVCCFLAGS, NVCC_OPT, NVCC_ARCH,
# NVCC_DEPFLAGS. These can be overridden from the environment or a gitignored
# 'config.mk' (see the overridable-variables list at the top of this file), so
# you normally don't need to edit the defaults below.
#
# FIXME some day I'll replace the env/config.mk mechanism with a proper configure script.

# Optional per-machine overrides (gitignored). Included before the '?=' defaults
# below so its assignments take effect.
-include config.mk

PYTHON ?= python3

# NVCC is the nvcc program alone; its flags live in NVCCFLAGS. Keeping them
# separate means you can point NVCC at a different toolkit (NVCC=/path/to/nvcc)
# without retyping the flags, and tweak the optimization level via NVCC_OPT
# (a debug build is just NVCC_OPT='-O0 -g') without retyping the whole command.
NVCC      ?= nvcc
NVCC_OPT  ?= -O3
NVCCFLAGS ?= -std=c++17 -m64 $(NVCC_OPT) --compiler-options -Wall,-fPIC

# Extra nvcc flags needed to build Makefile dependencies
#   -MMD create dep file, omitting "system" headers
#   -MP add phony target for each header in dep file (makes error reporting less confusing)
# Note: we don't need "-MT $@", since we use in-tree object filenames (x.cu -> x.o).
# Note: we don't need "-MT $*.d", since we use in-tree depfile names (x.cu -> x.d).
NVCC_DEPFLAGS ?= -MMD -MP

# NVIDIA archictecture.
DEFAULT_NVCC_ARCH = -gencode arch=compute_80,code=sm_80
DEFAULT_NVCC_ARCH += -gencode arch=compute_86,code=sm_86
DEFAULT_NVCC_ARCH += -gencode arch=compute_89,code=sm_89
# DEFAULT_ARCH += -gencode arch=compute_90,code=sm_90
NVCC_ARCH ?= $(DEFAULT_NVCC_ARCH)

# If building inside a conda env, add -L and an RPATH to $CONDA_PREFIX/lib on
# every link line. ksgpu currently has no conda-lib link deps, but keeping this
# scaffolding here matches pirate's Makefile and makes the shared libs
# self-locating if a future ksgpu source file grows a conda dep.
ifneq ($(CONDA_PREFIX),)
CONDA_LIBFLAGS = -L$(CONDA_PREFIX)/lib
CONDA_RPATHFLAGS = -Xcompiler '"-Wl,-rpath=$(CONDA_PREFIX)/lib"'
endif


####################################################################################################
#
# "Derived" config variables: PYTHON_INCDIR, NUMPY_INCDIR, PYBIND11_INCDIR, PYEXT_SUFFIX.
#
# These are autogenerated by makefile_helper.py, and cached in makefile_helper.out.
# PYEXT_SUFFIX is something like .cpython-312-x86_64-linux-gnu.so.


ifneq ($(MAKECMDGOALS),clean)
  include makefile_helper.out
endif

makefile_helper.out: makefile_helper.py Makefile
	$(PYTHON) makefile_helper.py


####################################################################################################


# The main output of the build process is these two libraries.
# Reminder: PYEXT_SUFFIX is something like .cpython-312-x86_64-linux-gnu.so.
KSGPU_LIB := lib/libksgpu.so
KSGPU_PYEXT = ksgpu/ksgpu_pybind11$(PYEXT_SUFFIX)

# These get compiled into lib/libksgpu.so
# Note: some files are .cpp (compiled via nvcc forwarding to host compiler)
# and some are .cu (require nvcc's CUDA frontend for device code).
LIB_CU_SRCFILES := \
  src_lib/memcpy_kernels.cu \
  src_lib/test_utils.cu

LIB_CPP_SRCFILES := \
  src_lib/Array.cpp \
  src_lib/CpuThreadPool.cpp \
  src_lib/Dtype.cpp \
  src_lib/assert_arrays_equal.cpp \
  src_lib/cuda_utils.cpp \
  src_lib/mem_utils.cpp \
  src_lib/rand_utils.cpp \
  src_lib/string_utils.cpp

LIB_SRCFILES := $(LIB_CU_SRCFILES) $(LIB_CPP_SRCFILES)

# These get compiled into ksgpu/ksgpu_pybind11....so
PYEXT_CPP_SRCFILES := \
  src_pybind11/ksgpu_pybind11.cpp \
  src_pybind11/pybind11_utils.cpp

PYEXT_SRCFILES := $(PYEXT_CPP_SRCFILES)

# These are in 1-1 corresponding with executables in bin/
# For example, 'src_bin/time-atomic-add.cu' gets compiled to 'bin/time-atomic-add'.
BIN_SRCFILES := \
  src_bin/time-atomic-add.cu \
  src_bin/time-fma.cu \
  src_bin/time-global-memory.cu \
  src_bin/time-l2-cache.cu \
  src_bin/time-local-transpose.cu \
  src_bin/time-memcpy-kernels.cu \
  src_bin/time-shared-memory.cu \
  src_bin/time-tensor-cores.cu \
  src_bin/time-warp-shuffle.cu \
  src_bin/scratch.cu \
  src_bin/reverse-engineer-mma.cu \
  src_bin/test-array.cu \
  src_bin/test-device-transpose-kernels.cu \
  src_bin/test-memcpy-kernels.cu \
  src_bin/test-sparse-mma.cu \
  src_bin/show-devices.cu

# Must list all python source files here.
# (Otherwise they won't show up in 'pip install' or pypi.)
PYFILES := \
  ksgpu/__init__.py \
  ksgpu/CudaStreamWrapper.py \
  ksgpu/pybind11_injections.py \
  ksgpu/tests.py \
  ksgpu/utils.py

# Must list all header files here.
# (Otherwise they won't show up in 'pip install' or pypi.)
HFILES := \
  include/ksgpu.hpp \
  include/ksgpu/Array.hpp \
  include/ksgpu/CpuThreadPool.hpp \
  include/ksgpu/Dtype.hpp \
  include/ksgpu/KernelTimer.hpp \
  include/ksgpu/constexpr_functions.hpp \
  include/ksgpu/cuda_utils.hpp \
  include/ksgpu/device_fp16.hpp \
  include/ksgpu/device_mma.hpp \
  include/ksgpu/device_transposes.hpp \
  include/ksgpu/mem_utils.hpp \
  include/ksgpu/memcpy_kernels.hpp \
  include/ksgpu/rand_utils.hpp \
  include/ksgpu/string_utils.hpp \
  include/ksgpu/test_utils.hpp \
  include/ksgpu/time_utils.hpp \
  include/ksgpu/xassert.hpp \
  include/ksgpu/dlpack.h \
  include/ksgpu/pybind11.hpp \
  include/ksgpu/pybind11_utils.hpp

# 'make clean' deletes {*~, *.o, *.d, *.so, *.pyc} from these dirs.
CLEAN_DIRS := . include include/ksgpu lib src_bin src_lib src_pybind11 ksgpu ksgpu/__pycache__

# Extra files to be deleted by 'make clean'.
# Note that 'ksgpu/include' and 'ksgpu/lib' are symlinks, so we put them in CLEAN_FILES, not CLEAN_RMDIRS
CLEAN_FILES := sdist_files.txt wheel_files.txt makefile_helper.out ksgpu/include ksgpu/lib

# Directories that should be empty at the end of 'make clean', and can be deleted.
CLEAN_RMDIRS := bin lib ksgpu/__pycache__


####################################################################################################


LIB_OFILES := $(LIB_CU_SRCFILES:%.cu=%.o) $(LIB_CPP_SRCFILES:%.cpp=%.o)
PYEXT_OFILES := $(PYEXT_CPP_SRCFILES:%.cpp=%.o)
BIN_XFILES := $(BIN_SRCFILES:src_bin/%.cu=bin/%)

# Must include all .d files, or build will break!
ALL_SRCFILES := $(LIB_SRCFILES) $(PYEXT_SRCFILES) $(BIN_SRCFILES)
DEPFILES := $(LIB_CU_SRCFILES:%.cu=%.d) $(LIB_CPP_SRCFILES:%.cpp=%.d)
DEPFILES += $(PYEXT_CPP_SRCFILES:%.cpp=%.d) $(BIN_SRCFILES:%.cu=%.d)

SDIST_FILES := pyproject.toml Makefile makefile_helper.py
SDIST_FILES += $(PYFILES) $(ALL_SRCFILES) $(HFILES)

# Some symlinks for the wheel:
#  - header file include/%.hpp gets symlinked to ksgpu/include/%.hpp
#  - library lib/libksgpu.so gets symlinked to ksgpu/lib/libksgpu.so
#  - python extension ksgpu/ksgpu_pybind11...so does not need to be symlinked/renamed.
WHEEL_FILES := $(PYFILES) $(KSGPU_PYEXT) ksgpu/$(KSGPU_LIB)
WHEEL_FILES += $(HFILES:%=ksgpu/%)

# Phony targets. The special targets 'build_wheel' and 'build_sdist' are needed by pip/pipmake.
lib: $(KSGPU_LIB) $(KSGPU_PYEXT)
bin: $(BIN_XFILES)
build_wheel: wheel_files.txt $(KSGPU_LIB) $(KSGPU_PYEXT)
build_sdist: sdist_files.txt

# Symlink {include,lib} into python directory 'ksgpu'.
ksgpu/include:
	ln -s ../include $@
ksgpu/lib:
	ln -s ../lib $@

# Build object files in src_lib/ or src_bin/ from .cu files
%.o: %.cu %.d
	$(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) $(NVCC_DEPFLAGS) -c -o $@ $<

# Build object files in src_lib/ from .cpp files
# Note: nvcc forwards .cpp files to the host compiler (no CUDA frontend processing).
%.o: %.cpp %.d
	$(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) $(NVCC_DEPFLAGS) -c -o $@ $<

# Build object files in src_pybind11/ from .cpp files, with special flags.
src_pybind11/%.o: src_pybind11/%.cpp src_pybind11/%.d
	$(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) $(NVCC_DEPFLAGS) -I$(PYTHON_INCDIR) -I$(NUMPY_INCDIR) -I$(PYBIND11_INCDIR) -c -o $@ $<

# Build the C++ library (lib/libksgpu.so)
$(KSGPU_LIB): $(LIB_OFILES)
	@mkdir -p lib
	$(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) -shared -o $@ $^ $(CONDA_LIBFLAGS) $(CONDA_RPATHFLAGS)

# Build binaries (bin/*)
bin/%: src_bin/%.o $(KSGPU_LIB)
	@mkdir -p bin/
	$(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) -o $@ $^

# Build the python extension (ksgpu/ksgpu_pybind11...so)
# We want it to automatically pull in the C++ library ksgpu/lib/libksgpu.so.
#
# The python extension has been built correctly if 'objdump -x' shows the following:
#   NEEDED   libksgpu.so
#   RUNPATH  $ORIGIN/lib
#
# The quoting can be understood by working backwards as follows:
#  - g++ command line should look like:   g++ -Wl,-rpath="\$ORIGIN/lib"
#  - nvcc command line should look like:  nvcc -Xcompiler '"-Wl,-rpath=\\$ORIGIN/lib"'
#  - Makefile line should look like:      nvcc -Xcompiler '"-Wl,-rpath=\\$$ORIGIN/lib"'

$(KSGPU_PYEXT): $(PYEXT_OFILES) $(KSGPU_LIB) ksgpu/lib
	$(NVCC) $(NVCCFLAGS) $(NVCC_ARCH) -shared -o $@ $(PYEXT_OFILES) -lksgpu -Lksgpu/lib $(CONDA_LIBFLAGS) -Xcompiler '"-Wl,-rpath=\\$$ORIGIN/lib"' $(CONDA_RPATHFLAGS)

# Needed by pip/pipmake: list of all files that go into the (non-editable) wheel.
wheel_files.txt: Makefile ksgpu/include ksgpu/lib
	rm -f $@
	for f in $(WHEEL_FILES); do echo $$f; done >>$@

# Needed by pip/pipmake: list of all files that go into the sdist.
sdist_files.txt: Makefile
	rm -f $@
	for f in $(SDIST_FILES); do echo $$f; done >>$@

clean:
	@for f in $(foreach d,$(CLEAN_DIRS),$(wildcard $d/*~ $d/*.o $d/*.d $d/*.so $d/*.pyc)); do echo rm $$f; rm $$f; done
	@for f in $(wildcard $(CLEAN_FILES) $(BIN_XFILES)); do echo rm $$f; rm $$f; done
	@for d in $(wildcard $(CLEAN_RMDIRS)); do echo rmdir $$d; rmdir $$d; done

# Specifying .SECONDARY with no prerequisites disables auto-deletion of intermediate files.
.SECONDARY:

# If a depfile is absent, build can still proceed.
$(DEPFILES):

# Include any depfiles which are present.
include $(wildcard $(DEPFILES))
