sbcl-arrays.lisp

;;; -*- mode: Lisp; indent-tabs-mode: nil; -*-
;;;
;;; Copyright (c) 2010, Alexander Gavrilov (angavrilov@gmail.com)
;;;
;;; This file contains the groundwork for vectorized
;;; array access intrinsics.
;;;

(in-package #:SSE)

#|--------------------------------------|
 |  SSE ARRAY ELEMENT SIZE CALCULATION  |
 |--------------------------------------|#

(eval-when (:compile-toplevel :load-toplevel :execute)
  (defun sse-elt-shift-from-saetp (info)
    (and info
         (subtypep (saetp-specifier info) 'number)
         (not (saetp-fixnum-p info))
         (case (saetp-n-bits info)
           (8 0) (16 1) (32 2) (64 3) (128 4)))))

(defglobal %%size-shift-table%%
    (let ((arr (make-array (1+ widetag-mask) :initial-element nil)))
      (loop
         for info across *specialized-array-element-type-properties*
         for shift = (sse-elt-shift-from-saetp info)
         when shift
         do (setf (svref arr (saetp-typecode info)) shift))
      arr)
  "A table of element size shifts for supported SSE array types.")

(declaim (inline sse-elt-shift-of)
         (ftype (function (t) (integer 0 4)) sse-elt-shift-of))

(defun sse-elt-shift-of (obj)
  "Returns the SSE element size shift for the given object, or fails if it is not a valid SSE vector."
  (declare (optimize (safety 0)))
  (the (integer 0 4)
    (or (svref %%size-shift-table%%
               (if (sb-vm::%other-pointer-p obj)
                   (%other-pointer-widetag obj)
                   0))
        (error 'type-error
               :datum obj
               :expected-type 'sse-array))))

#|--------------------------------------|
 |     SSE-ARRAY TYPE AND ALLOCATION    |
 |--------------------------------------|#

(deftype sse-array (&optional (elt-type '* et-p) dims)
  "Type of arrays efficiently accessed by SSE aref intrinsics and returned by make-sse-array.
Should be assumed to be SIMPLE-ARRAY, except that displacing with MAKE-SSE-ARRAY is allowed."
  (if (eq elt-type '*)
      (progn
        (when et-p
          (error "SSE-ARRAY must have a specific element type."))
        `(simple-array * ,dims))
      (let* ((upgraded (upgraded-array-element-type elt-type))
             (shift (sse-elt-shift-from-saetp (find-saetp upgraded))))
        (when (null shift)
          (error "Invalid SSE-ARRAY element type: ~S" elt-type))
        (unless (subtypep upgraded elt-type)
          (warn "SSE-ARRAY element type ~S has been upgraded to ~S" elt-type upgraded))
        `(simple-array ,upgraded ,dims))))

(defun make-sse-array (dimensions &key (element-type '(unsigned-byte 8)) (initial-element nil ie-p) displaced-to (displaced-index-offset 0))
  "Allocates an SSE-ARRAY aligned to the 16-byte boundary. Flattens displacement chains for performance reasons."
  (let* ((upgraded (upgraded-array-element-type element-type))
         (shift (sse-elt-shift-from-saetp (find-saetp upgraded))))
    (when (null shift)
      (error "Invalid SSE-ARRAY element type: ~S" element-type))
    (if displaced-to
        ;; Fake displacement by allocating a simple-array header
        (let* ((dimensions (if (listp dimensions) dimensions (list dimensions)))
               (rank (length dimensions))
               (count (reduce #'* dimensions)))
          (unless (subtypep element-type (array-element-type displaced-to))
            (error "can't displace an array of type ~S into another of type ~S"
                   element-type (array-element-type displaced-to)))
          (with-array-data ((data displaced-to)
                            (start displaced-index-offset)
                            (end))
            (unless (= start 0)
              (error "SSE-ARRAY does not support displaced index offset."))
            (unless (<= count end)
              (array-bounding-indices-bad-error data start count))
            (if (= rank 1)
                (progn
                  (when (< count end)
                    (warn "SSE-ARRAY displaced size extended to the full length of the vector."))
                  data)
                (let ((new-array (make-array-header simple-array-widetag rank)))
                  (set-array-header new-array data count nil 0 dimensions nil t)))))
        ;; X86-64 vectors are already aligned to 16 bytes
        (apply #'make-array dimensions :element-type upgraded
               (if ie-p (list :initial-element initial-element))))))

#|--------------------------------------|
 |  AREF INTRINSIC DEFINITION HELPERS   |
 |--------------------------------------|#

(defconstant +vector-data-fixup+
  (- (* vector-data-offset n-word-bytes) other-pointer-lowtag)
  "Offset from a tagged vector pointer to its data")

(defmacro array-data-expr (array-var &optional is-vector)
  (ecase is-vector
    (:yes array-var)
    (:no `(%array-data-vector ,array-var))
    ((nil)
     `(if (array-header-p ,array-var)
          (%array-data-vector ,array-var)
          ,array-var))))

;; Depends on the vector-length field being in the same place
;; as the array fill pointer, which for simple-array is equal
;; to the total size.
;; The integer constant argument is the number of elements that
;; should be deducted from the size to account for SIMD access.
(defknown %sse-array-size (simple-array fixnum) array-total-size (flushable always-translatable dx-safe))

(define-vop (%sse-array-size/0)
  (:translate %sse-array-size)
  (:args (array :scs (descriptor-reg)))
  (:arg-types * (:constant (integer 0 0)))
  (:info gap)
  (:ignore gap)
  (:policy :fast-safe)
  (:results (result :scs (any-reg)))
  (:result-types tagged-num)
  (:generator 3
    (loadw result array vector-length-slot other-pointer-lowtag)))

(define-vop (%sse-array-size %sse-array-size/0)
  (:arg-types * (:constant (integer 1 16)))
  (:ignore)
  (:temporary (:sc any-reg) tmp)
  (:generator 8
    (loadw result array vector-length-slot other-pointer-lowtag)
    (inst mov tmp (fixnumize gap))
    (inst cmp result tmp)
    (inst cmov :ng tmp result)
    (inst sub result tmp)))

(defmacro with-sse-data (((sap-var data-var array) (offset-var index)) &body code)
  ;; Compute a SAP and offset for the specified array and index. Check bounds.
  (with-unique-names (data-index data-end elt-shift access-size)
    (once-only ((array array)
                (index index))
      `(locally
           (declare (optimize (insert-array-bounds-checks 0)))
         (with-array-data ((,data-var ,array)
                           (,data-index ,index)
                           (,data-end))
           (let* ((,sap-var (int-sap (get-lisp-obj-address ,data-var)))
                  (,elt-shift (sse-elt-shift-of ,data-var))
                  (,access-size (ash 16 (- ,elt-shift)))
                  (,offset-var (+ (ash ,data-index ,elt-shift) +vector-data-fixup+)))
             (declare (type system-area-pointer ,sap-var)
                      (type fixnum ,offset-var))
             (unless (<= 0 ,data-index (+ ,data-index ,access-size) ,data-end)
               (array-bounding-indices-bad-error ,array ,index (+ ,index ,access-size)))
             ,@code))))))

(defun sse-array-type-info-or-give-up (type)
  (cond ((and (array-type-p type)
              (not (array-type-complexp type)))
         (let* ((etype (array-type-specialized-element-type type))
                (shift (sse-elt-shift-from-saetp
                        (if (eq etype *wild-type*) nil
                            (find-saetp-by-ctype etype)))))
           (unless shift
             (give-up-ir1-transform "not a known SSE-compatible array element type: ~S"
                                    (type-specifier etype)))
           (values shift
                   (and (listp (array-type-dimensions type))
                        (if (null (cdr (array-type-dimensions type))) :yes :no)))))
        ((union-type-p type)
         ;; Support unions of array types with the same elt size
         (let (nonfirst rshift rdims)
           (dolist (subtype (union-type-types type))
             (multiple-value-bind (shift dims)
                 (sse-array-type-info-or-give-up subtype)
               (unless nonfirst
                 (setf nonfirst t
                       rshift shift
                       rdims dims))
               (unless (= rshift shift)
                 (give-up-ir1-transform
                  "union member types have different element sizes"))
               (unless (eq rdims dims)
                 (setf rdims nil))))
           (values rshift rdims)))
        (t
         (give-up-ir1-transform "not a simple array type"))))

(defun sse-array-info-or-give-up (lvar ref-size)
  ;; Look up the SSE element size and check if it is definitely a
  ;; vector
  (multiple-value-bind (shift dim-info)
      (sse-array-type-info-or-give-up (lvar-type lvar))
    (values (ash 1 shift)                 ; step
            (ash (1- ref-size) (- shift)) ; gap (size of SIMD overreach)
            dim-info)))

(defmacro def-aref-intrinsic (postfix rtype reader writer &key (ref-size 16) side-effect?)
  (let* ((rm-aref (symbolicate "ROW-MAJOR-AREF-" postfix))
         (rm-aset (if writer (symbolicate "ROW-MAJOR-ASET-" postfix)))
         (aref (symbolicate "AREF-" postfix))
         (aset (if writer (symbolicate "%ASET-" postfix)))
         (reader-vop (symbolicate "%" reader))
         (reader/ix-vop (symbolicate "%" reader "/IX"))
         (writer-vop (if writer (symbolicate "%" writer)))
         (writer/ix-vop (if writer (symbolicate "%" writer "/IX")))
         (rtype (or rtype '(values)))
         (known-flags (if side-effect?
                          '(dx-safe)
                          '(foldable flushable dx-safe)))
         (index-expression
          (if (= ref-size 0)
              ``(the signed-word index)
              ``(the signed-word (%check-bound array (%sse-array-size array ,gap) index)))))
    `(progn
       ;; ROW-MAJOR-AREF
       (export ',rm-aref)
       (defknown ,rm-aref (array index) ,rtype ,known-flags)
       (defun ,rm-aref (array index)
         (with-sse-data ((sap data array)
                         (offset index))
           (,reader-vop sap offset 1 0)))
       ;;
       (deftransform ,rm-aref ((array index) (simple-array t) * :important t)
         ,(format nil "open-code ~A" rm-aref)
         (multiple-value-bind (step gap is-vector)
             (sse-array-info-or-give-up array ,ref-size)
           (declare (ignorable gap))
           `(,',reader/ix-vop (array-data-expr array ,is-vector)
                              ,,index-expression
                              ,step ,+vector-data-fixup+)))
       ;; AREF
       (export ',aref)
       (defknown ,aref (array &rest index) ,rtype ,known-flags)
       (defun ,aref (array &rest indices)
         (declare (truly-dynamic-extent indices))
         (with-sse-data ((sap data array)
                         (offset (%array-row-major-index array indices)))
           (,reader-vop sap offset 1 0)))
       ;;
       (defoptimizer (,aref derive-type) ((array &rest indices) node)
         (assert-array-rank array (length indices))
         (values-specifier-type ',rtype))
       (deftransform ,aref ((array &rest indices) (simple-array &rest t) * :important t)
         ,(format nil "open-code ~A" aref)
         (multiple-value-bind (step gap is-vector)
             (sse-array-info-or-give-up array ,ref-size)
           (declare (ignorable gap))
           (let ((syms (make-gensym-list (length indices))))
             `(lambda (array ,@syms)
                (let ((index ,(if (eq is-vector :yes) (first syms)
                                  `(array-row-major-index array ,@syms))))
                  (,',reader/ix-vop (array-data-expr array ,is-vector)
                                    ,,index-expression
                                    ,step ,+vector-data-fixup+))))))
       ,@(if writer
             `(;; ROW-MAJOR-ASET
               (defknown ,rm-aset (array index sse-pack) ,rtype ())
               (defsetf ,rm-aref ,rm-aset)
               (defun ,rm-aset (array index new-value)
                 (with-sse-data ((sap data array)
                                 (offset index))
                   (,writer-vop sap offset 1 0 (the ,rtype new-value))
                   new-value))
               ;;
               (deftransform ,rm-aset ((array index value) (simple-array t t) * :important t)
                 ,(format nil "open-code ~A" rm-aset)
                 (multiple-value-bind (step gap is-vector)
                     (sse-array-info-or-give-up array ,ref-size)
                   (declare (ignorable gap))
                   `(progn
                      (,',writer/ix-vop (array-data-expr array ,is-vector)
                                        ,,index-expression
                                        ,step ,+vector-data-fixup+
                                        (the sse-pack value))
                      value)))
               ;; %ASET
               (defknown ,aset (array &rest t) ,rtype ())
               (defsetf ,aref ,aset)
               (defun ,aset (array &rest stuff)
                 (let ((new-value (car (last stuff))))
                   (with-sse-data ((sap data array)
                                   (offset (%array-row-major-index array (nbutlast stuff))))
                     (,writer-vop sap offset 1 0 (the ,rtype new-value))
                     new-value)))
               ;;
               (defoptimizer (,aset derive-type) ((array &rest stuff) node)
                 (assert-array-rank array (1- (length stuff)))
                 (assert-lvar-type (car (last stuff)) (specifier-type 'sse-pack)
                                   (lexenv-policy (node-lexenv node)))
                 (specifier-type ',rtype))
               (deftransform ,aset ((array &rest stuff) (simple-array &rest t) * :important t)
                 ,(format nil "open-code ~A" aset)
                 (multiple-value-bind (step gap is-vector)
                     (sse-array-info-or-give-up array ,ref-size)
                   (declare (ignorable gap))
                   (let ((syms (make-gensym-list (length stuff))))
                     `(lambda (array ,@syms)
                        (let ((index ,(if (eq is-vector :yes) (first syms)
                                          `(array-row-major-index array ,@(butlast syms)))))
                          (,',writer/ix-vop (array-data-expr array ,is-vector)
                                            ,,index-expression
                                            ,step ,+vector-data-fixup+
                                            (the sse-pack ,(car (last syms)))))
                        ,(car (last syms)))))))))))