From 11158df5b1f6a799fbe659fbd8f2a39015812375 Mon Sep 17 00:00:00 2001
From: Serhiy Mytrovtsiy <mitrovtsiy@ukr.net>
Date: Tue, 4 Feb 2025 16:38:28 +0100
Subject: [PATCH] feat: added stress tests for CPU efficiency and performance
 cores and GPU

---
 Kit/helpers.swift             | 159 ++++++++++++++++++++++++++++++++++
 Stats/Views/AppSettings.swift |  51 ++++++++++-
 2 files changed, 209 insertions(+), 1 deletion(-)

diff --git a/Kit/helpers.swift b/Kit/helpers.swift
index 4a1ee6b439a..9dfbf289252 100644
--- a/Kit/helpers.swift
+++ b/Kit/helpers.swift
@@ -14,6 +14,7 @@ import Cocoa
 import ServiceManagement
 import UserNotifications
 import WebKit
+import Metal
 
 public struct LaunchAtLogin {
     private static let id = "\(Bundle.main.bundleIdentifier!).LaunchAtLogin"
@@ -1626,3 +1627,161 @@ public class VerticallyCenteredTextFieldCell: NSTextFieldCell {
         self.attributedStringValue.draw(in: titleRect)
     }
 }
+
+public class CPUeStressTest {
+    public var isRunning: Bool = false
+    
+    private var workers: [DispatchWorkItem] = []
+    private let queue = DispatchQueue.global(qos: .background)
+    
+    public init() {}
+    
+    public func start() {
+        guard !self.isRunning else { return }
+        self.isRunning = true
+        
+        let efficientCoreCount = ProcessInfo.processInfo.processorCount / 2
+        self.workers.removeAll()
+
+        for index in 0..<efficientCoreCount {
+            let worker = DispatchWorkItem { [weak self] in
+                self?.test(threadIndex: index)
+            }
+            self.workers.append(worker)
+            self.queue.async(execute: worker)
+        }
+    }
+    
+    public func stop() {
+        self.isRunning = false
+        self.workers.forEach { $0.cancel() }
+        self.workers.removeAll()
+    }
+    
+    private func test(threadIndex: Int) {
+        pthread_set_qos_class_self_np(QOS_CLASS_BACKGROUND, 0)
+        while isRunning {
+            _ = sqrt(987654.321)
+        }
+    }
+}
+
+public class CPUpStressTest {
+    public var isRunning = false
+    
+    private var workers: [DispatchWorkItem] = []
+    private let queue = DispatchQueue.global(qos: .userInteractive)
+
+    public init() {}
+
+    public func start() {
+        guard !self.isRunning else { return }
+        self.isRunning = true
+
+        let performanceCoreCount = ProcessInfo.processInfo.activeProcessorCount
+        self.workers.removeAll()
+
+        for index in 0..<performanceCoreCount {
+            let worker = DispatchWorkItem { [weak self] in
+                self?.test(threadIndex: index)
+            }
+            self.workers.append(worker)
+            self.queue.async(execute: worker)
+        }
+    }
+    
+    public func stop() {
+        self.isRunning = false
+        self.workers.forEach { $0.cancel() }
+        self.workers.removeAll()
+    }
+    
+    private func test(threadIndex: Int) {
+        pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0)
+        while isRunning {
+            _ = sqrt(987654.321)
+        }
+    }
+}
+
+public class GPUStressTest {
+    public var isRunning = false
+    
+    private let device: MTLDevice
+    private let commandQueue: MTLCommandQueue
+    private let pipeline: MTLComputePipelineState
+    private let dataSize = 50_000_000 // Large data size for GPU workload
+    private let bufferA: MTLBuffer
+    private let bufferB: MTLBuffer
+    private let bufferC: MTLBuffer
+    
+    public init?() {
+        guard let device = MTLCreateSystemDefaultDevice(), let queue = device.makeCommandQueue() else {
+            return nil
+        }
+        
+        self.device = device
+        self.commandQueue = queue
+        
+        let source = """
+        #include <metal_stdlib>
+        using namespace metal;
+        kernel void full_load_kernel(const device float* inA [[buffer(0)]],
+                                      const device float* inB [[buffer(1)]],
+                                      device float* outC [[buffer(2)]],
+                                      uint id [[thread_position_in_grid]]) {
+            outC[id] = (inA[id] * inB[id]) + sin(inA[id]) + cos(inB[id]) + tan(inA[id]) + log(inB[id]);
+        }
+        """
+        
+        do {
+            let library = try device.makeLibrary(source: source, options: nil)
+            let function = library.makeFunction(name: "full_load_kernel")!
+            self.pipeline = try device.makeComputePipelineState(function: function)
+        } catch {
+            return nil
+        }
+        
+        self.bufferA = device.makeBuffer(length: self.dataSize * MemoryLayout<Float>.size, options: .storageModeShared)!
+        self.bufferB = device.makeBuffer(length: self.dataSize * MemoryLayout<Float>.size, options: .storageModeShared)!
+        self.bufferC = device.makeBuffer(length: self.dataSize * MemoryLayout<Float>.size, options: .storageModeShared)!
+        
+        let dataA = [Float](repeating: 1.0, count: self.dataSize)
+        let dataB = [Float](repeating: 2.0, count: self.dataSize)
+        memcpy(self.bufferA.contents(), dataA, dataA.count * MemoryLayout<Float>.size)
+        memcpy(self.bufferB.contents(), dataB, dataB.count * MemoryLayout<Float>.size)
+    }
+    
+    public func start() {
+        guard !self.isRunning else { return }
+        self.isRunning = true
+        
+        DispatchQueue.global(qos: .userInitiated).async { [weak self] in
+            self?.test()
+        }
+    }
+    
+    public func stop() {
+        self.isRunning = false
+    }
+    
+    private func test() {
+        let threadGroupSize = MTLSize(width: 256, height: 1, depth: 1)
+        let gridSize = MTLSize(width: self.dataSize, height: 1, depth: 1)
+        
+        while self.isRunning {
+            guard let commandBuffer = self.commandQueue.makeCommandBuffer(),
+                  let commandEncoder = commandBuffer.makeComputeCommandEncoder() else {
+                break
+            }
+            
+            commandEncoder.setComputePipelineState(self.pipeline)
+            commandEncoder.setBuffer(self.bufferA, offset: 0, index: 0)
+            commandEncoder.setBuffer(self.bufferB, offset: 0, index: 1)
+            commandEncoder.setBuffer(self.bufferC, offset: 0, index: 2)
+            commandEncoder.dispatchThreads(gridSize, threadsPerThreadgroup: threadGroupSize)
+            commandEncoder.endEncoding()
+            commandBuffer.commit()
+        }
+    }
+}
diff --git a/Stats/Views/AppSettings.swift b/Stats/Views/AppSettings.swift
index 34a79a56aa0..6462a79cdc7 100644
--- a/Stats/Views/AppSettings.swift
+++ b/Stats/Views/AppSettings.swift
@@ -169,7 +169,24 @@ class ApplicationSettings: NSStackView {
         
         self.addArrangedSubview(scrollView)
         
-        NotificationCenter.default.addObserver(self, selector: #selector(toggleUninstallHelperButton), name: .fanHelperState, object: nil)
+        let CPUeButton = buttonView(#selector(self.toggleCPUeStressTest), text: localizedString("Run"))
+        let CPUpButton = buttonView(#selector(self.toggleCPUpStressTest), text: localizedString("Run"))
+        let GPUButton = buttonView(#selector(self.toggleGPUStressTest), text: localizedString("Run"))
+        
+        self.CPUeButton = CPUeButton
+        self.CPUpButton = CPUpButton
+        self.GPUButton = GPUButton
+        
+        var tests = [
+            PreferencesRow(localizedString("Efficiency cores"), component: CPUeButton),
+            PreferencesRow(localizedString("Performance cores"), component: CPUpButton)
+        ]
+        if self.GPUTest != nil {
+            tests.append(PreferencesRow(localizedString("GPU"), component: GPUButton))
+        }
+        scrollView.stackView.addArrangedSubview(PreferencesSection(label: localizedString("Stress tests"), tests))
+        
+        NotificationCenter.default.addObserver(self, selector: #selector(self.toggleUninstallHelperButton), name: .fanHelperState, object: nil)
     }
     
     required init?(coder: NSCoder) {
@@ -372,6 +389,38 @@ class ApplicationSettings: NSStackView {
     @objc private func uninstallHelper() {
         SMCHelper.shared.uninstall()
     }
+    
+    @objc private func toggleCPUeStressTest() {
+        if self.CPUeTest.isRunning {
+            self.CPUeTest.stop()
+            self.CPUeButton?.title = localizedString("Run")
+        } else {
+            self.CPUeTest.start()
+            self.CPUeButton?.title = localizedString("Stop")
+        }
+    }
+    
+    @objc private func toggleCPUpStressTest() {
+        if self.CPUpTest.isRunning {
+            self.CPUpTest.stop()
+            self.CPUpButton?.title = localizedString("Run")
+        } else {
+            self.CPUpTest.start()
+            self.CPUpButton?.title = localizedString("Stop")
+        }
+    }
+    
+    @objc private func toggleGPUStressTest() {
+        guard let test = self.GPUTest else { return }
+        
+        if test.isRunning {
+            test.stop()
+            self.GPUButton?.title = localizedString("Run")
+        } else {
+            test.start()
+            self.GPUButton?.title = localizedString("Stop")
+        }
+    }
 }
 
 private class ModuleSelectorView: NSStackView {