BUILD_TYPE?=Release
BUILD_DIR?=./build
Python3_EXECUTABLE?=$(shell which python3)
BINDINGS_PYTHON_ENABLED=ON
DATAFLOW_EXECUTION_ENABLED=OFF
TIMING_ENABLED=OFF
CC_COMPILER=
CXX_COMPILER=
CUDA_SUPPORT?=OFF
CONCRETE_CORE_PATH?= $(shell pwd)/concrete-core
INSTALL_PREFIX?=$(abspath $(BUILD_DIR))/install
INSTALL_PATH=$(abspath $(INSTALL_PREFIX))/concretecompiler/
MAKEFILE_ROOT_DIR=$(shell pwd)

CONCRETE_OPTIMIZER_DIR ?= $(shell pwd)/concrete-optimizer

KEYSETCACHEDEV=/tmp/KeySetCache
KEYSETCACHECI ?= ../KeySetCache
KEYSETCACHENAME ?= KeySetCacheV4

HPX_VERSION?=1.7.1
HPX_URL=https://github.com/STEllAR-GROUP/hpx/archive/refs/tags/$(HPX_VERSION).tar.gz
HPX_TARBALL=$(shell pwd)/hpx-$(HPX_VERSION).tar.gz
HPX_LOCAL_DIR=$(shell pwd)/hpx-$(HPX_VERSION)
HPX_INSTALL_DIR?=$(HPX_LOCAL_DIR)/build

CONCRETE_CORE_FFI_VERSION?=0.2.0

ML_BENCH_SUBSET_ID=

# Find OS
OS=undefined
ifeq ($(shell uname), Linux)
	OS=linux
else ifeq ($(shell uname), Darwin)
	OS=darwin
endif

# Setup find arguments for MacOS
ifeq ($(OS), darwin)
FIND_EXECUTABLE_ARG=-perm +111
else
FIND_EXECUTABLE_ARG=-executable
endif

ARCHITECTURE=undefined
ifeq ($(shell uname -m), arm64)
	ARCHITECTURE=aarch64
else 
	ARCHITECTURE=amd64
endif

CONCRETE_CORE_FFI_TARBALL=concrete-core-ffi_$(CONCRETE_CORE_FFI_VERSION)_$(OS)_$(ARCHITECTURE).tar.gz

export PATH := $(abspath $(BUILD_DIR))/bin:$(PATH)

ifeq ($(shell which ccache),)
	CCACHE=OFF
else
	CCACHE=ON
endif

ifeq ($(CCACHE),ON)
		CMAKE_CCACHE_OPTIONS=-DCMAKE_CXX_COMPILER_LAUNCHER=ccache
else
		CMAKE_CCACHE_OPTIONS=
endif

ifneq ($(CC_COMPILER),)
		CC_COMPILER_OPTION=-DCMAKE_C_COMPILER=$(CC_COMPILER)
else
		CC_COMPILER_OPTION=
endif

ifneq ($(CXX_COMPILER),)
		CXX_COMPILER_OPTION=-DCMAKE_CXX_COMPILER=$(CXX_COMPILER)
else
		CXX_COMPILER_OPTION=
endif

# don't run parallel python tests if compiler doesn't support it
ifeq ($(DATAFLOW_EXECUTION_ENABLED),ON)
		PYTHON_TESTS_MARKER=""
else
		PYTHON_TESTS_MARKER="not parallel"
endif

all: concretecompiler python-bindings build-tests build-benchmarks build-mlbench doc

# concrete-core-ffi #######################################

CONCRETE_CORE_FFI_FOLDER = $(shell pwd)/concrete-core-ffi-$(CONCRETE_CORE_FFI_VERSION)

CONCRETE_CORE_FFI_ARTIFACTS=$(CONCRETE_CORE_FFI_FOLDER)/libconcrete_core_ffi.a $(CONCRETE_CORE_FFI_FOLDER)/concrete-core-ffi.h

concrete-core-ffi: $(CONCRETE_CORE_FFI_ARTIFACTS)

$(CONCRETE_CORE_FFI_ARTIFACTS): $(CONCRETE_CORE_FFI_FOLDER)

$(CONCRETE_CORE_FFI_FOLDER): $(CONCRETE_CORE_FFI_TARBALL)
	mkdir -p $(CONCRETE_CORE_FFI_FOLDER)
	tar -xvzf $(CONCRETE_CORE_FFI_TARBALL) --directory $(CONCRETE_CORE_FFI_FOLDER)

$(CONCRETE_CORE_FFI_TARBALL):
	curl -L https://github.com/zama-ai/concrete-core/releases/download/concrete-core-ffi-$(CONCRETE_CORE_FFI_VERSION)/$(CONCRETE_CORE_FFI_TARBALL) -o $(CONCRETE_CORE_FFI_TARBALL)

# concrete-optimizer ######################################

LIB_CONCRETE_OPTIMIZER_CPP = $(CONCRETE_OPTIMIZER_DIR)/target/libconcrete_optimizer_cpp.a

concrete-optimizer-lib:
	make -C $(CONCRETE_OPTIMIZER_DIR)/concrete-optimizer-cpp $(LIB_CONCRETE_OPTIMIZER_CPP)

# HPX #####################################################

install-hpx-from-source: $(HPX_LOCAL_DIR)
	mkdir -p $(HPX_LOCAL_DIR)/build
	cd $(HPX_LOCAL_DIR)/build && cmake \
		-DHPX_WITH_FETCH_ASIO=on \
		-DHPX_FILESYSTEM_WITH_BOOST_FILESYSTEM_COMPATIBILITY=ON \
		-DHPX_WITH_MALLOC=system ..
	cd $(HPX_LOCAL_DIR)/build && make -j2

$(HPX_TARBALL): 
	curl -L $(HPX_URL) -o $(HPX_TARBALL)

$(HPX_LOCAL_DIR): $(HPX_TARBALL)
	tar xzvf $(HPX_TARBALL)

$(BUILD_DIR)/configured.stamp:
	mkdir -p $(BUILD_DIR)
	cmake -B $(BUILD_DIR) -GNinja ../llvm-project/llvm/ \
	$(CMAKE_CCACHE_OPTIONS) \
	$(CC_COMPILER_OPTION) \
	$(CXX_COMPILER_OPTION) \
	-DLLVM_ENABLE_PROJECTS="mlir;clang;openmp" \
	-DLLVM_BUILD_EXAMPLES=OFF \
	-DLLVM_TARGETS_TO_BUILD="host" \
	-DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \
	-DLLVM_ENABLE_ASSERTIONS=ON \
	-DMLIR_ENABLE_BINDINGS_PYTHON=$(BINDINGS_PYTHON_ENABLED) \
	-DCONCRETELANG_BINDINGS_PYTHON_ENABLED=$(BINDINGS_PYTHON_ENABLED) \
	-DCONCRETELANG_DATAFLOW_EXECUTION_ENABLED=$(DATAFLOW_EXECUTION_ENABLED) \
	-DCONCRETELANG_TIMING_ENABLED=$(TIMING_ENABLED) \
	-DCONCRETE_FFI_RELEASE=$(CONCRETE_CORE_FFI_FOLDER) \
	-DHPX_DIR=${HPX_INSTALL_DIR}/lib/cmake/HPX \
	-DLLVM_EXTERNAL_PROJECTS=concretelang \
	-DLLVM_EXTERNAL_CONCRETELANG_SOURCE_DIR=. \
	-DPython3_EXECUTABLE=${Python3_EXECUTABLE} \
	-DCONCRETE_OPTIMIZER_DIR=${CONCRETE_OPTIMIZER_DIR} \
	-DCONCRETE_CORE_PATH=$(CONCRETE_CORE_PATH) \
	-DCONCRETELANG_CUDA_SUPPORT=${CUDA_SUPPORT}
	touch $@

build-initialized: concrete-optimizer-lib concrete-core-ffi $(BUILD_DIR)/configured.stamp

doc: build-initialized
	cmake --build $(BUILD_DIR) --target mlir-doc

concretecompiler: build-initialized
	cmake --build $(BUILD_DIR) --target concretecompiler

python-bindings: build-initialized
	cmake --build $(BUILD_DIR) --target ConcretelangMLIRPythonModules
	cmake --build $(BUILD_DIR) --target ConcretelangPythonModules

rust-bindings: install
	cd lib/Bindings/Rust && CONCRETE_COMPILER_INSTALL_DIR=$(INSTALL_PATH) cargo build --release

CAPI:
	cmake --build $(BUILD_DIR) --target CONCRETELANGCAPIFHE CONCRETELANGCAPIFHELINALG CONCRETELANGCAPISupport

clientlib: build-initialized
	cmake --build $(BUILD_DIR) --target ConcretelangClientLib

serverlib: build-initialized
	cmake --build $(BUILD_DIR) --target ConcretelangServerLib



GITHUB_URL=https://api.github.com/repos/zama-ai/concrete-compiler-internal
GITHUB_URL_LIST_ARTIFACTS="${GITHUB_URL}/actions/artifacts?name=${KEYSETCACHENAME}&per_page=1"
CURL=curl -H"Accept: application/vnd.github.v3+json" -H"authorization: Bearer ${GITHUB_TOKEN}"
keysetcache.zip: REDIRECT_URL = $(shell ${CURL} -s ${GITHUB_URL_LIST_ARTIFACTS} | grep archive_download_url | grep -o 'http[^"]\+')
keysetcache.zip:
	${CURL} --location -o keysetcache.zip ${REDIRECT_URL}
	du -h keysetcache.zip

keysetcache_ci_populated: keysetcache.zip
	unzip keysetcache.zip -d ${KEYSETCACHECI}
	du -sh ${KEYSETCACHECI}
	rm keysetcache.zip

keysetcache_populated: keysetcache.zip
	unzip keysetcache.zip -d ${KEYSETCACHEDEV}
	du -sh ${KEYSETCACHEDEV}
	rm keysetcache.zip


# test

build-tests: build-unit-tests build-end-to-end-tests

run-tests: run-check-tests run-unit-tests run-end-to-end-tests run-python-tests

## check-tests

run-check-tests: concretecompiler file-check not
	$(BUILD_DIR)/bin/llvm-lit -v tests/check_tests

## unit-tests

build-unit-tests: build-initialized
	cmake --build $(BUILD_DIR) --target ConcretelangUnitTests

run-unit-tests: build-unit-tests
	find $(BUILD_DIR)/tools/concretelang/tests/unit_tests -name unit_tests_concretelang* $(FIND_EXECUTABLE_ARG) -type f | xargs -n1 ./run_test_bin.sh

## python-tests

run-python-tests: python-bindings concretecompiler
	PYTHONPATH=${PYTHONPATH}:$(BUILD_DIR)/tools/concretelang/python_packages/concretelang_core LD_PRELOAD=$(BUILD_DIR)/lib/libConcretelangRuntime.so pytest -vs -m $(PYTHON_TESTS_MARKER) tests/python

test-compiler-file-output: concretecompiler
	pytest -vs tests/test_compiler_file_output

## rust-tests
run-rust-tests: rust-bindings
	cd lib/Bindings/Rust && CONCRETE_COMPILER_INSTALL_DIR=$(INSTALL_PATH) LD_LIBRARY_PATH=$(INSTALL_PATH)/lib cargo test --release

## end-to-end-tests

build-end-to-end-jit-test: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_jit_test

build-end-to-end-test: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_test

build-end-to-end-jit-encrypted-tensor: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_jit_encrypted_tensor

build-end-to-end-jit-fhelinalg: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_jit_fhelinalg

build-end-to-end-jit-lambda: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_jit_lambda

build-end-to-end-tests: build-end-to-end-jit-test build-end-to-end-test build-end-to-end-jit-encrypted-tensor build-end-to-end-jit-fhelinalg build-end-to-end-jit-lambda

### end-to-end-tests CPU

FIXTURE_CPU_DIR=tests/end_to_end_fixture/tests_cpu

$(FIXTURE_CPU_DIR)/%.yaml: tests/end_to_end_fixture/%_gen.py
	mkdir -p $(FIXTURE_CPU_DIR)
	$(Python3_EXECUTABLE) $< > $@

$(FIXTURE_CPU_DIR)/bug_report.yaml:
	unzip -o $(FIXTURE_CPU_DIR)/bug_report.zip -d $(FIXTURE_CPU_DIR)

generate-cpu-tests: $(FIXTURE_CPU_DIR)/end_to_end_leveled.yaml $(FIXTURE_CPU_DIR)/end_to_end_apply_lookup_table.yaml $(FIXTURE_CPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml $(FIXTURE_CPU_DIR)/bug_report.yaml

run-end-to-end-tests: build-end-to-end-tests generate-cpu-tests
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_jit_test
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_jit_encrypted_tensor
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_jit_fhelinalg
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_jit_lambda
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_test \
	--backend=cpu --jit $(FIXTURE_CPU_DIR)/*.yaml

### end-to-end-tests GPU

FIXTURE_GPU_DIR=tests/end_to_end_fixture/tests_gpu

$(FIXTURE_GPU_DIR):
	mkdir -p $(FIXTURE_GPU_DIR)

$(FIXTURE_GPU_DIR)/end_to_end_apply_lookup_table.yaml: tests/end_to_end_fixture/end_to_end_apply_lookup_table_gen.py
	$(Python3_EXECUTABLE) $< --bitwidth 1 2 3 4 5 6 7 > $@

$(FIXTURE_GPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml: tests/end_to_end_fixture/end_to_end_linalg_apply_lookup_table_gen.py
	$(Python3_EXECUTABLE) $< --bitwidth 1 2 3 4 5 6 7 > $@

generate-gpu-tests: $(FIXTURE_GPU_DIR) $(FIXTURE_GPU_DIR)/end_to_end_apply_lookup_table.yaml $(FIXTURE_GPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml

run-end-to-end-tests-gpu: build-end-to-end-test generate-gpu-tests
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_test \
	--backend=gpu --library /tmp/concrete_compiler/gpu_tests/ \
	$(FIXTURE_GPU_DIR)/*.yaml

## end-to-end-dataflow-tests

build-end-to-end-dataflow-tests: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_jit_auto_parallelization
	cmake --build $(BUILD_DIR) --target end_to_end_jit_distributed
	cmake --build $(BUILD_DIR) --target end_to_end_jit_aes_short

run-end-to-end-dataflow-tests: build-end-to-end-dataflow-tests
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_jit_auto_parallelization
	$(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_jit_distributed

# benchmark

build-benchmarks: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_benchmark

## benchmark CPU

BENCHMARK_CPU_DIR=tests/end_to_end_fixture/benchmarks_cpu

$(BENCHMARK_CPU_DIR):
	mkdir -p $@

$(BENCHMARK_CPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml: tests/end_to_end_fixture/end_to_end_linalg_apply_lookup_table_gen.py
	$(Python3_EXECUTABLE) $< --n-ct 64 128 1024 > $@

$(BENCHMARK_CPU_DIR)/%.yaml: tests/end_to_end_fixture/%_gen.py
	$(Python3_EXECUTABLE) $< > $@

generate-cpu-benchmarks: $(BENCHMARK_CPU_DIR) $(BENCHMARK_CPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml $(BENCHMARK_CPU_DIR)/end_to_end_apply_lookup_table.yaml $(BENCHMARK_CPU_DIR)/end_to_end_leveled.yaml

run-cpu-benchmarks: build-benchmarks generate-cpu-benchmarks
	$(BUILD_DIR)/bin/end_to_end_benchmark \
	--backend=cpu \
	--benchmark_out=benchmarks_results.json --benchmark_out_format=json \
	$(BENCHMARK_CPU_DIR)/*.yaml

## benchmark GPU

BENCHMARK_GPU_DIR=tests/end_to_end_fixture/benchmarks_gpu

$(BENCHMARK_GPU_DIR):
	mkdir -p $@

$(BENCHMARK_GPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml: tests/end_to_end_fixture/end_to_end_linalg_apply_lookup_table_gen.py
	$(Python3_EXECUTABLE) $< \
	--bitwidth 1 2 3 4 5 6 7 --n-ct 1 128 1024 2048 8192


generate-gpu-benchmarks: $(BENCHMARK_GPU_DIR) $(BENCHMARK_GPU_DIR)/end_to_end_linalg_apply_lookup_table.yaml

run-gpu-benchmarks: build-benchmarks generate-cpu-benchmarks
	$(BUILD_DIR)/bin/end_to_end_benchmark \
	--backend=gpu \
	--benchmark_out=benchmarks_results.json --benchmark_out_format=json \
	$(BENCHMARK_CPU_DIR)/*.yaml



build-mlbench: build-initialized
	cmake --build $(BUILD_DIR) --target end_to_end_mlbench

generate-mlbench:
	mkdir -p tests/end_to_end_benchmarks/mlbench
	rm -rf tests/end_to_end_benchmarks/mlbench/*
	unzip tests/end_to_end_benchmarks/mlbench.zip -d tests/end_to_end_benchmarks/mlbench
	rm -f tests/end_to_end_benchmarks/mlbench/**/*\=*
	find tests/end_to_end_benchmarks/mlbench -name "*.mlir" -exec sed -e '1d' -e 's/ func / func.func /g' -e 's/ linalg.tensor_/ tensor./g' -e '$$d' -i {} \;
	$(Python3_EXECUTABLE) tests/end_to_end_benchmarks/generate_bench_yaml.py tests/end_to_end_benchmarks/mlbench tests/end_to_end_benchmarks/mlbench/end_to_end_mlbench

run-mlbench: build-mlbench generate-mlbench
	tests/end_to_end_benchmarks/end_to_end_mlbench.sh tests/end_to_end_benchmarks/mlbench/ $(BUILD_DIR)/bin/end_to_end_mlbench

run-mlbench-subset: build-mlbench generate-mlbench
	@[ "${ML_BENCH_SUBSET_ID}" ] || ( echo "ML_BENCH_SUBSET_ID is not set"; exit 1 )
	tests/end_to_end_benchmarks/end_to_end_mlbench.sh tests/end_to_end_benchmarks/mlbench/end_to_end_mlbench_$(ML_BENCH_SUBSET_ID).yaml $(BUILD_DIR)/bin/end_to_end_mlbench

show-stress-tests-summary:
	@echo '------ Stress tests summary ------'
	@echo
	@echo 'Rates:'
	@cd tests/stress_tests/trace && grep success_rate -R
	@echo
	@echo 'Parameters issues:'
	@cd tests/stress_tests/trace && grep BAD -R || echo 'No issues'

stress-tests: concretecompiler
	pytest -vs tests/stress_tests

# useful for faster cache generation, need pytest-parallel
stress-tests-fast-cache: concretecompiler
	pytest --workers auto -vs tests/stress_tests

# LLVM/MLIR dependencies

all-deps: file-check not

file-check: build-initialized
	cmake --build $(BUILD_DIR) --target FileCheck

not: build-initialized
	cmake --build $(BUILD_DIR) --target not

mlir-cpu-runner: build-initialized
	cmake --build $(BUILD_DIR) --target mlir-cpu-runner

opt: build-initialized
	cmake --build $(BUILD_DIR) --target opt

mlir-opt: build-initialized
	cmake --build $(BUILD_DIR) --target mlir-opt

mlir-translate: build-initialized
	cmake --build $(BUILD_DIR) --target mlir-translate

update-python-version:
	echo "__version__ = \"`git describe --tags --abbrev=0 | grep -e '[0-9].*' -o`\"" > lib/Bindings/Python/version.txt

check-python-format:
	black --check tests/python/ lib/Bindings/Python/concrete/

python-format:
	black tests/python/ lib/Bindings/Python/concrete/

python-lint:
	pylint --rcfile=../pylintrc lib/Bindings/Python/concrete/compiler

check-rust-format:
	cd lib/Bindings/Rust && cargo fmt  --check

rust-format:
	cd lib/Bindings/Rust && cargo fmt 

# libraries we want to have in the installation that aren't already a deps of other targets
install-deps:
	cmake --build $(BUILD_DIR) --target MLIRCAPIRegistration

install: concretecompiler concrete-optimizer-lib CAPI install-deps
	$(info Install prefix set to $(INSTALL_PREFIX))
	$(info Installing under $(INSTALL_PATH))
	mkdir -p $(INSTALL_PATH)/include
	cp -R $(abspath $(BUILD_DIR))/bin $(INSTALL_PATH)
	cp -R $(abspath $(BUILD_DIR))/lib $(INSTALL_PATH)
	cp $(LIB_CONCRETE_OPTIMIZER_CPP) $(INSTALL_PATH)/lib/
	cp $(CONCRETE_OPTIMIZER_DIR)/concrete-optimizer-cpp/src/cpp/concrete-optimizer.hpp $(INSTALL_PATH)/include
	cd $(MAKEFILE_ROOT_DIR)/include && find . -iregex '^.*\.\(h\|hpp\|td\)$$' -exec cp --parents {} $(INSTALL_PATH)/include \;
	cd $(MAKEFILE_ROOT_DIR)/../llvm-project/llvm/include && find . -iregex '^.*\.\(h\|hpp\|td\)$$' -exec cp --parents {} $(INSTALL_PATH)/include \;
	cd $(MAKEFILE_ROOT_DIR)/../llvm-project/mlir/include && find . -iregex '^.*\.\(h\|hpp\|td\)$$' -exec cp --parents {} $(INSTALL_PATH)/include \;
	cd $(abspath $(BUILD_DIR))/include && find . -iname '*.inc' -exec cp --parents {} $(INSTALL_PATH)/include \;
	cd $(abspath $(BUILD_DIR))/tools/concretelang/include && find . -iname '*.inc' -exec cp --parents {} $(INSTALL_PATH)/include \;
	cd $(abspath $(BUILD_DIR))/tools/mlir/include && find . -iname '*.inc' -exec cp --parents {} $(INSTALL_PATH)/include \;

.PHONY: build-initialized \
	build-end-to-end-jit \
	concretecompiler \
	python-bindings \
	add-deps \
	file-check \
	not \
	update-python-version \
	python-lint \
	python-format \
	check-python-format \
	concrete-optimizer-lib \
	concrete-core-ffi \
	build-tests \
	run-tests \
	run-check-tests \
	build-unit-tests \
	run-unit-tests \
	run-python-tests \
	build-end-to-end-tests \
	build-end-to-end-dataflow-tests \
	run-end-to-end-dataflow-tests \
	concrete-core-ffi \
	opt \
	mlir-opt \
	mlir-cpu-runner \
	mlir-translate