(cl:defpackage #:avx2
  (:use :cl)
  (:export #:d4+ #:d4* #:d4ref #:d4set))

(cl:in-package #:avx2)

(sb-c:defknown (d4+ d4*) ((sb-ext:simd-pack-256 double-float)
                          (sb-ext:simd-pack-256 double-float))
    (sb-ext:simd-pack-256 double-float)
    (sb-c:movable sb-c:flushable sb-c:always-translatable)
  :overwrite-fndb-silently t)

(sb-c:defknown d4ref ((simple-array double-float (*)) sb-int:index)
    (sb-ext:simd-pack-256 double-float)
    (sb-c:movable sb-c:flushable sb-c:always-translatable)
  :overwrite-fndb-silently t)

(sb-c:defknown d4set ((simple-array double-float (*))
                      sb-int:index
                      (sb-ext:simd-pack-256 double-float))
    (sb-ext:simd-pack-256 double-float)
    (sb-c:movable sb-c:unsafely-flushable sb-c:always-translatable)
  :overwrite-fndb-silently t)

(cl:in-package #:sb-vm)

(define-vop (avx2:d4+)
  (:translate avx2:d4+)
  (:policy :fast-safe)
  (:args (x :scs (double-avx2-reg))
         (y :scs (double-avx2-reg)))
  (:arg-types simd-pack-256-double simd-pack-256-double)
  (:results (r :scs (double-avx2-reg)))
  (:result-types simd-pack-256-double)
  (:generator 4 (inst vaddpd r x y)))

(define-vop (avx2:d4*)
  (:translate avx2:d4*)
  (:policy :fast-safe)
  (:args (x :scs (double-avx2-reg))
         (y :scs (double-avx2-reg)))
  (:arg-types simd-pack-256-double simd-pack-256-double)
  (:results (r :scs (double-avx2-reg)))
  (:result-types simd-pack-256-double)
  (:generator 4 (inst vmulpd r x y)))

(define-vop (avx2:d4ref)
  (:translate avx2:d4ref)
  (:policy :fast-safe)
  (:args (array :scs (descriptor-reg))
         (index :scs (any-reg)))
  (:arg-types simple-array-double-float positive-fixnum)
  (:results (value :scs (double-avx2-reg)))
  (:result-types simd-pack-256-double)
  (:generator 5 (inst vmovupd value
                      (make-ea-for-float-ref
                       array index 0 32
                       :scale (ash 1 (- word-shift n-fixnum-tag-bits))))))

(define-vop (avx2:d4set)
  (:translate avx2:d4set)
  (:policy :fast-safe)
  (:args (array :scs (descriptor-reg))
         (index :scs (any-reg))
         (value :scs (double-avx2-reg) :target result))
  (:arg-types simple-array-double-float positive-fixnum simd-pack-256-double)
  (:results (result :scs (double-avx2-reg)))
  (:result-types simd-pack-256-double)
  (:generator 5 (inst vmovupd
                      (make-ea-for-float-ref
                       array index 0 32
                       :scale (ash 1 (- word-shift n-fixnum-tag-bits)))
                      value)
              (move result value)))

(cl:in-package #:avx2)

(defun d4+ (x y)
  (d4+ x y))

(defun d4* (x y)
  (d4* x y))

(defun d4ref (array index)
  (d4ref array index))

(defun d4set (array index value)
  (d4set array index value))

(defun jacobi-1d (src dst &optional (iterations 1))
  (declare (optimize (speed 3) (safety 0))
           (type (simple-array double-float (*)) src dst)
           (type (integer 0 #.most-positive-fixnum) iterations))
  (loop repeat iterations do
    (loop for index from 1 by 4 below (- (length src) 4) do
      (d4set dst index
             (d4* (d4+ (d4ref src (1- index))
                       (d4ref src (1+ index)))
                  (sb-ext:%make-simd-pack-256-double 0.5d0 0.5d0 0.5d0 0.5d0))))
    (rotatef src dst))
  src)

(defun jacobi-test (size iterations)
  (let ((a (make-array size :element-type 'double-float :initial-element 0.0d0))
        (b (make-array size :element-type 'double-float :initial-element 0.0d0)))
    (setf (aref a 0) 1.0d0)
    (setf (aref b 0) 1.0d0)
    (jacobi-1d a b iterations)))