dispatch.go raw

   1  package native
   2  
   3  import (
   4  	"unsafe"
   5      
   6      `github.com/klauspost/cpuid/v2`
   7  	"github.com/cloudwego/base64x/internal/rt"
   8  	"github.com/cloudwego/base64x/internal/native/avx2"
   9  	"github.com/cloudwego/base64x/internal/native/sse"
  10  )
  11  
  12  var (
  13      hasAVX2 = cpuid.CPU.Has(cpuid.AVX2)
  14      hasSSE = cpuid.CPU.Has(cpuid.SSE)
  15  )
  16  
  17  var (
  18  	S_b64decode uintptr
  19  	S_b64encode uintptr
  20  )
  21  
  22  var (
  23  	F_b64decode func(out unsafe.Pointer, src unsafe.Pointer, len int, mod int) (ret int)
  24  	F_b64encode func(out unsafe.Pointer, src unsafe.Pointer, mod int)
  25  )
  26  
  27  func useAVX2() {
  28  	avx2.Use()
  29  	S_b64decode = avx2.S_b64decode
  30  	S_b64encode = avx2.S_b64encode
  31  
  32  	F_b64decode = avx2.F_b64decode
  33  	F_b64encode = avx2.F_b64encode
  34  }
  35  
  36  func useSSE() {
  37  	sse.Use()
  38  	S_b64decode = sse.S_b64decode
  39  	S_b64encode = sse.S_b64encode
  40  
  41  	F_b64decode = sse.F_b64decode
  42  	F_b64encode = sse.F_b64encode
  43  }
  44  
  45  //go:nosplit
  46  func B64Decode(out *[]byte, src unsafe.Pointer, len int, mod int) (ret int) {
  47      return F_b64decode(rt.NoEscape(unsafe.Pointer(out)), rt.NoEscape(unsafe.Pointer(src)), len, mod)
  48  }
  49  
  50  //go:nosplit
  51  func B64Encode(out *[]byte, src *[]byte, mod int) {
  52  	F_b64encode(rt.NoEscape(unsafe.Pointer(out)), rt.NoEscape(unsafe.Pointer(src)), mod)
  53  }
  54  
  55  func init() {
  56  	if hasAVX2 {
  57  		useAVX2()
  58  	} else if hasSSE {
  59  		useSSE()
  60  	} else {
  61  		panic("Unsupported CPU, lacks of AVX2 or SSE CPUID Flag. maybe it's too old to run Sonic.")
  62  	}
  63  }
  64