cpufeatures.go raw

   1  //go:build amd64
   2  
   3  package p256k1
   4  
   5  import (
   6  	"sync"
   7  	"sync/atomic"
   8  
   9  	"github.com/klauspost/cpuid/v2"
  10  )
  11  
  12  // CPU feature flags
  13  var (
  14  	// hasAVX2CPU indicates whether the CPU supports AVX2 instructions.
  15  	// This is detected at startup and never changes.
  16  	hasAVX2CPU bool
  17  
  18  	// hasBMI2CPU indicates whether the CPU supports BMI2 instructions.
  19  	// BMI2 provides MULX, ADCX, ADOX for efficient carry-chain arithmetic.
  20  	hasBMI2CPU bool
  21  
  22  	// hasADXCPU indicates whether the CPU supports ADX instructions.
  23  	// ADX provides ADCX/ADOX for parallel carry chains.
  24  	hasADXCPU bool
  25  
  26  	// avx2Disabled allows runtime disabling of AVX2 for testing/debugging.
  27  	// Uses atomic operations for thread-safety without locks on the fast path.
  28  	avx2Disabled atomic.Bool
  29  
  30  	// bmi2Disabled allows runtime disabling of BMI2 for testing/debugging.
  31  	bmi2Disabled atomic.Bool
  32  
  33  	// initOnce ensures CPU detection runs exactly once
  34  	initOnce sync.Once
  35  )
  36  
  37  func init() {
  38  	initOnce.Do(detectCPUFeatures)
  39  }
  40  
  41  // detectCPUFeatures detects CPU capabilities at startup
  42  func detectCPUFeatures() {
  43  	hasAVX2CPU = cpuid.CPU.Has(cpuid.AVX2)
  44  	hasBMI2CPU = cpuid.CPU.Has(cpuid.BMI2)
  45  	hasADXCPU = cpuid.CPU.Has(cpuid.ADX)
  46  }
  47  
  48  // HasAVX2 returns true if AVX2 is available and enabled.
  49  // This is the function that should be called in hot paths to decide
  50  // whether to use AVX2-optimized code paths.
  51  func HasAVX2() bool {
  52  	return hasAVX2CPU && !avx2Disabled.Load()
  53  }
  54  
  55  // HasAVX2CPU returns true if the CPU supports AVX2, regardless of whether
  56  // it's been disabled via SetAVX2Enabled.
  57  func HasAVX2CPU() bool {
  58  	return hasAVX2CPU
  59  }
  60  
  61  // SetAVX2Enabled enables or disables the use of AVX2 instructions.
  62  // This is useful for benchmarking to compare AVX2 vs non-AVX2 performance,
  63  // or for debugging. Pass true to enable AVX2 (default), false to disable.
  64  // This function is thread-safe.
  65  func SetAVX2Enabled(enabled bool) {
  66  	avx2Disabled.Store(!enabled)
  67  }
  68  
  69  // IsAVX2Enabled returns whether AVX2 is currently enabled.
  70  // Returns true if AVX2 is both available on the CPU and not disabled.
  71  func IsAVX2Enabled() bool {
  72  	return HasAVX2()
  73  }
  74  
  75  // HasBMI2 returns true if BMI2 is available and enabled.
  76  // BMI2 provides MULX for efficient multiplication without affecting flags,
  77  // enabling parallel carry chains with ADCX/ADOX.
  78  func HasBMI2() bool {
  79  	return hasBMI2CPU && hasADXCPU && !bmi2Disabled.Load()
  80  }
  81  
  82  // HasBMI2CPU returns true if the CPU supports BMI2, regardless of whether
  83  // it's been disabled via SetBMI2Enabled.
  84  func HasBMI2CPU() bool {
  85  	return hasBMI2CPU
  86  }
  87  
  88  // HasADXCPU returns true if the CPU supports ADX (ADCX/ADOX instructions).
  89  func HasADXCPU() bool {
  90  	return hasADXCPU
  91  }
  92  
  93  // SetBMI2Enabled enables or disables the use of BMI2 instructions.
  94  // This is useful for benchmarking to compare BMI2 vs non-BMI2 performance.
  95  // Pass true to enable BMI2 (default), false to disable.
  96  // This function is thread-safe.
  97  func SetBMI2Enabled(enabled bool) {
  98  	bmi2Disabled.Store(!enabled)
  99  }
 100  
 101  // IsBMI2Enabled returns whether BMI2 is currently enabled.
 102  // Returns true if BMI2+ADX are both available on the CPU and not disabled.
 103  func IsBMI2Enabled() bool {
 104  	return HasBMI2()
 105  }
 106