github.com/grailbio/base@v0.0.11/simd/float_amd64.go (about)

     1  // Copyright 2021 GRAIL, Inc.  All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  // +build amd64,!appengine
     6  
     7  package simd
     8  
     9  import (
    10  	"math"
    11  	"reflect"
    12  	"unsafe"
    13  
    14  	"golang.org/x/sys/cpu"
    15  )
    16  
    17  //go:noescape
    18  func findNaNOrInf64SSSE3Asm(data unsafe.Pointer, nElem int) int
    19  
    20  //go:noescape
    21  func findNaNOrInf64AVX2Asm(data unsafe.Pointer, nElem int) int
    22  
    23  var avx2Available bool
    24  
    25  func init() {
    26  	avx2Available = cpu.X86.HasAVX2
    27  	// possible todo: detect FMA and/or AVX512DQ.
    28  }
    29  
    30  // FindNaNOrInf64 returns the position of the first NaN/inf value if one is
    31  // present, and -1 otherwise.
    32  func FindNaNOrInf64(data []float64) int {
    33  	nElem := len(data)
    34  	if nElem < 16 {
    35  		for i, x := range data {
    36  			if (math.Float64bits(x) & (0x7ff << 52)) == (0x7ff << 52) {
    37  				return i
    38  			}
    39  		}
    40  		return -1
    41  	}
    42  	dataHeader := (*reflect.SliceHeader)(unsafe.Pointer(&data))
    43  	if avx2Available {
    44  		return findNaNOrInf64AVX2Asm(unsafe.Pointer(dataHeader.Data), nElem)
    45  	} else {
    46  		return findNaNOrInf64SSSE3Asm(unsafe.Pointer(dataHeader.Data), nElem)
    47  	}
    48  }