From 43661754b3182daa7fa86e7f1cea2ca9244dbd3d Mon Sep 17 00:00:00 2001
From: Parker Williams <pawilliams@hrl.com>
Date: Thu, 29 Sep 2022 16:48:16 -0700
Subject: [PATCH] check for diagonal matrices in find-diag-in-e-basis

It seems that LAPACK becomes unstable on ARM when attempting to
find the eigenvectors and eigenvalues of a diagonal matrix. This
instability is triggered in find-diagonalizer-in-e-basis. The
operation becomes trivial when run on a diagonal matrix, so we
check for this and return the identity to avoid making a
potentially unstable LAPACK call.

Fixes #842.
---
 src/compilers/approx.lisp  | 68 +++++++++++++++++++++-----------------
 src/matrix-operations.lisp |  7 ++++
 2 files changed, 45 insertions(+), 30 deletions(-)

diff --git a/src/compilers/approx.lisp b/src/compilers/approx.lisp
index 4610ef2e4..860354015 100644
--- a/src/compilers/approx.lisp
+++ b/src/compilers/approx.lisp
@@ -199,36 +199,44 @@ This function should be completely deterministic (i.e., produce the same value f
   (check-type m magicl:matrix)
   (let* ((u (magicl:@ +edag-basis+ m +e-basis+))
          (gammag (magicl:@ u (magicl:transpose u))))
-    (loop :for attempt :from 1 :to num-attempts :do
-      (let* ((coeff (diagonalizer-number-generator attempt))
-             (matrix (magicl:map (lambda (z)
-                                   (+ (* coeff       (realpart z))
-                                      (* (- 1 coeff) (imagpart z))))
-                                 gammag))
-             (evecs (ensure-positive-determinant
-                     (orthonormalize-matrix!
-                      (nth-value 1 (magicl:eig matrix)))))
-             (evals (magicl:diag
-                     (magicl:@ (magicl:transpose evecs)
-                               gammag
-                               evecs)))
-             (v (magicl:@ evecs
-                          (from-diag evals)
-                          (magicl:transpose evecs))))
-        (when (and (double= 1.0d0 (magicl:det evecs))
-                   (magicl:every #'double= gammag v))
-          (when *check-math*
-            (assert (magicl:every #'double~
-                                  (eye 4 :type 'double-float)
-                                  (magicl:@ (magicl:transpose evecs)
-                                            evecs))
-                (evecs)
-                "The calculated eigenvectors were not found to be orthonormal. ~
-               EE^T =~%~A"
-                (magicl:@ (magicl:transpose evecs)
-                          evecs)))
-          (return-from find-diagonalizer-in-e-basis evecs)))))
-  (error 'diagonalizer-not-found :matrix m :attempts num-attempts))
+    (cond
+      ;; XXX: It seems that LAPACK distributions on ARM do not work or
+      ;; are not stable when calculating eigenvectors of a diagonal
+      ;; matrix. We have not closely investigated why, but checking
+      ;; for diagonal matrices as a special case seems to work.
+      ((diagonal-matrix-p gammag)
+       (magicl:eye (magicl:shape m) :type (magicl:element-type m)))
+      (t
+       (loop :for attempt :from 1 :to num-attempts :do
+         (let* ((coeff (diagonalizer-number-generator attempt))
+                (matrix (magicl:map (lambda (z)
+                                      (+ (* coeff       (realpart z))
+                                         (* (- 1 coeff) (imagpart z))))
+                                    gammag))
+                (evecs (ensure-positive-determinant
+                        (orthonormalize-matrix!
+                         (nth-value 1 (magicl:eig matrix)))))
+                (evals (magicl:diag
+                        (magicl:@ (magicl:transpose evecs)
+                                  gammag
+                                  evecs)))
+                (v (magicl:@ evecs
+                             (from-diag evals)
+                             (magicl:transpose evecs))))
+           (when (and (double= 1.0d0 (magicl:det evecs))
+                      (magicl:every #'double= gammag v))
+             (when *check-math*
+               (assert (magicl:every #'double~
+                                     (eye 4 :type 'double-float)
+                                     (magicl:@ (magicl:transpose evecs)
+                                               evecs))
+                       (evecs)
+                       "The calculated eigenvectors were not found to be orthonormal. ~
+                        EE^T =~%~A"
+                       (magicl:@ (magicl:transpose evecs)
+                                 evecs)))
+             (return-from find-diagonalizer-in-e-basis evecs))))
+       (error 'diagonalizer-not-found :matrix m :attempts num-attempts)))))
 
 (defun diagonalizer-in-e-basis (m)
   "For M in SU(4), compute an SO(4) column matrix of eigenvectors of E^* M E (E^* M E)^T.
diff --git a/src/matrix-operations.lisp b/src/matrix-operations.lisp
index 10a21c481..c5cb7908c 100644
--- a/src/matrix-operations.lisp
+++ b/src/matrix-operations.lisp
@@ -8,6 +8,13 @@
 
 (in-package #:cl-quil)
 
+(defun diagonal-matrix-p (m)
+  "Is M diagonal?"
+  (dotimes (row (magicl:nrows m) t)
+    (dotimes (col (magicl:ncols m))
+      (when (and (/= row col) (not (double~ 0 (magicl:tref m row col))))
+        (return-from diagonal-matrix-p nil)))))
+
 (defun matrix-first-column-equality (x y)
   (check-type x magicl:matrix)
   (check-type y magicl:matrix)