(define-module (epistemia packages machine-learning)
  #:use-module (guix packages)
  #:use-module (guix build-system go)
  #:use-module (guix download)
  #:use-module (guix utils)
  #:use-module (guix git-download)
  #:use-module (guix gexp)
  #:use-module ((guix licenses) #:prefix license:)
  #:use-module (gnu packages base)
  #:use-module (gnu packages golang-check)
  #:use-module (gnu packages golang-web)
  #:use-module (gnu packages golang-xyz)
  #:use-module (gnu packages llvm)
  #:use-module (gnu packages machine-learning)
  #:use-module (gnu packages rocm)
  #:use-module (gnu packages rocm-libs))

(define-public llama-cpp-rocm
  (let ((tag "b9515"))
    (package
      (inherit llama-cpp)
      (name "llama-cpp-rocm")
      (version (string-append "0.0.0-" tag))
      (source
       (origin
         (method git-fetch)
         (uri (git-reference
		(url "https://github.com/ggml-org/llama.cpp")
		(commit tag)))
         (file-name (git-file-name name tag))
         (sha256
          (base32 "1wfwan6x22dajpzl2hcavxy5q0wwrf142qrapy998cc0zkzn771c"))))
      (arguments
       (substitute-keyword-arguments (package-arguments llama-cpp)
	 ((#:configure-flags flags #~'())
	  #~(append #$flags
		    '(;; CMAKE_HIP_FLAGS is to workaround regression in ROCm 7.2
		      "-DCMAKE_HIP_FLAGS=\"-mllvm --amdgpu-unroll-threshold-local=600\""
		      #$(string-append "-DGPU_TARGETS=" (current-amd-gpu-targets-string))
		      ;; temporarily disable webui building until I figure out how to add second source for hf bucket
		      "-DLLAMA_USE_PREBUILT_UI=OFF" "-DLLAMA_BUILD_UI=OFF")))
	 ((#:phases phases)
	  #~(modify-phases #$phases
	      ;; mtp added test that also needs internet for downloading its model
	      (add-after 'fix-tests 'fix-tests-new
		(lambda _
		  (substitute* "tests/CMakeLists.txt"
		    (("llama_build_and_test\\(test-state-restore-fragmented\\.cpp .*")
		     "")
		    (("set_tests_properties\\(test-state-restore-fragmented PROPERTIES.*")
		     ""))
		  (substitute* "tests/CMakeLists.txt"
		    (("llama_build_and_test\\(test-recurrent-state-rollback\\.cpp .*")
		     "")
		    (("set_tests_properties\\(test-recurrent-state-rollback PROPERTIES.*")
		     ""))
		  (substitute* "tests/CMakeLists.txt"
		    (("llama_build_and_test\\(test-save-load-state\\.cpp .*")
		     "")
		    (("set_tests_properties\\(test-save-load-state PROPERTIES.*")
		     ""))))
	      (delete 'fix-python-shebang)))))
      (inputs
       (modify-inputs (package-inputs llama-cpp)
	 (replace "ggml" ggml-rocm))))))

(define-public ggml-rocm
  (package
    (inherit ggml)
    (name "ggml-rocm")
;;    (version "0.13.1")
    (version "9ca88e98e6d9e84cf77959931130d38c5b56c19b")
    (source
     (origin
       (method git-fetch)
       (uri (git-reference
              ;; (url "https://github.com/ggml-org/ggml")
	      (url "https://git.vineetk.net/ggml")
              ;; (commit (string-append "v" version))
	      (commit version)))
       (file-name (git-file-name name version))
       (sha256
        (base32 "0llmlnkcg8zrdj6v2wpl018iw6j0sp1fsg2if3w3zd91hv66jq5r"))))
    (arguments
     (substitute-keyword-arguments (package-arguments (@@ (gnu packages machine-learning) ggml))
       ((#:configure-flags flags #~'())
	#~(append #$flags (list (string-append "-DCMAKE_HIP_COMPILER_ROCM_ROOT=" #$rocm-hip-runtime)
				(string-append "-DCMAKE_HIP_FLAGS=--rocm-path=" #$rocm-hip-runtime
					       " --rocm-device-lib-path=" #$rocm-device-libs "/amdgcn/bitcode")
				#$(string-append "-DGPU_TARGETS=" (current-amd-gpu-targets-string))
				"-DGGML_CPU=ON"
				"-DGGML_HIP=ON"
				"-DGGML_HIP_ROCWMMA_FATTN=OFF"
				"-DGGML_CUDA_FA_ALL_QUANTS=ON"
				;; for stable-diffusion.cpp
				"-DGGML_CPU_ALL_VARIANTS=ON"
				"-DGGML_MAX_NAME=128")))))
    (inputs
     (modify-inputs (package-inputs ggml)
       (append clang-rocm
	       lld-rocm
	       rocm-cmake
	       rocm-device-libs
	       rocm-hipcc
	       rocm-hip-runtime
	       rocr-runtime
	       hipblas
	       hipblas-common
	       rocblas)))))