gnu: Add xnnpack-for-r-torch.

* gnu/packages/machine-learning.scm (xnnpack-for-r-torch): New variable.
(xnnpack)[inputs]: Use xnnpack-for-r-torch instead of xnnpack.

1 files changed, 70 insertions, 1 deletions

diff --git a/gnu/packages/machine-learning.scm b/gnu/packages/machine-learning.scm
index b72c442fc0..51ab46bb86 100644
--- a/gnu/packages/machine-learning.scm
+++ b/gnu/packages/machine-learning.scm
@@ -4251,6 +4251,74 @@ TensorFlow.js, PyTorch, and MediaPipe.")
          "i686-linux" "x86_64-linux"))
       (license license:bsd-3))))
 
+(define-public xnnpack-for-r-torch
+  (let ((version "0.0")
+        (commit "51a987591a6fc9f0fc0707077f53d763ac132cbf")
+        (revision "2"))
+    (package
+      (inherit xnnpack)
+      (version (git-version version revision commit))
+      (source
+       (origin
+         (method git-fetch)
+         (uri (git-reference (url (package-home-page xnnpack)) (commit commit)))
+         (file-name (git-file-name (package-name xnnpack) version))
+         (sha256
+          (base32
+           "1rzby82xq8d0rl1d148yz88jh9cpsw5c8b2yw7yg39mi7qmr55rm"))
+         (modules '((guix build utils)
+                    (ice-9 ftw)
+                    (ice-9 textual-ports)
+                    (srfi srfi-26)))
+         (snippet
+          '(begin
+             ;; Remove autogenerated files
+             (for-each
+              (lambda (dir)
+                (let ((gendir (string-append "src/" dir "/gen")))
+                  (when (file-exists? gendir)
+                    (delete-file-recursively gendir)
+                    ;; Needed for the scripts generating the files
+                    (mkdir gendir))))
+              (scandir "src" (negate (cut member <> '("." "..")))))
+             (delete-file-recursively "google3")
+             (delete-file "cmake/microkernels.cmake")
+             ;; Additional autogenerated files which contain the string
+             ;; "Auto-generated file"
+             (for-each
+              (lambda (dir)
+                (for-each
+                 (lambda (name)
+                   (let ((path (string-append dir "/" name)))
+                     (when (call-with-input-file path
+                             (lambda (port)
+                               (string-contains
+                                (get-string-all port)
+                                "Auto-generated file")))
+                       (delete-file path))))
+                 (scandir dir (negate (cut member <> '("." ".."))))))
+              '("test" "bench" "eval" "models" "src/enums" "src/xnnpack"))))))
+      (arguments
+       (substitute-keyword-arguments (package-arguments xnnpack)
+         ((#:phases phases)
+          #~(modify-phases #$phases
+              (replace 'generate-files
+                (lambda _
+                  (for-each
+                   (lambda (name)
+                     (when (and (string-prefix? "generate" name)
+                                (string-suffix? ".sh" name)
+                                (not (equal? "generate-amalgamation.sh" name)))
+                       (display (string-append name "\n"))
+                       (invoke "bash" (string-append "scripts/" name))))
+                   (scandir "scripts"))
+                  ;; These need to run after the above scripts
+                  (display "Remaining files\n")
+                  (invoke "python3" "tools/update-microkernels.py")
+                  (substitute* "tools/amalgamate-microkernels.py"
+                    (("BUILD") "BUILD.bazel"))
+                  (invoke "bash" "scripts/generate-amalgamation.sh"))))))))))
+
 ;; Warning: This package requires AVX2 or AVX-512 instructions.
 (define-public fbgemm
   (package
@@ -4851,7 +4919,8 @@ Note: currently this package does not provide GPU support.")
     (inputs
      (modify-inputs (package-inputs python-pytorch)
        (replace "qnnpack-pytorch" qnnpack-pytorch-for-r-torch)
-       (replace "oneapi-dnnl" oneapi-dnnl-for-r-torch)))
+       (replace "oneapi-dnnl" oneapi-dnnl-for-r-torch)
+       (replace "xnnpack" xnnpack-for-r-torch)))
     (propagated-inputs
      (modify-inputs (package-propagated-inputs python-pytorch)
        (append python-filelock