60 lines
2.8 KiB
Go
60 lines
2.8 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
package core
|
|
|
|
import "sourcedock.dev/petrbalvin/tensor/internal/engine"
|
|
|
|
// CPU policy. The heavy loops of the library, matrix products,
|
|
// convolutions, element-wise maps, axis reductions, scans and the FFT,
|
|
// run in parallel across the worker count chosen here. The default is
|
|
// NumCPU, which is right for almost every machine; call SetNumCPU once
|
|
// at startup to override (e.g. for a container with a CPU quota, or a
|
|
// small embedded box).
|
|
|
|
// SetNumCPU sets the number of goroutines the parallel kernels may use
|
|
// and returns the previous value. Values below 1 reset to NumCPU. It
|
|
// is safe to call at any time; running kernels finish with their old
|
|
// worker count.
|
|
func SetNumCPU(n int) int { return engine.SetNumWorkers(n) }
|
|
|
|
// NumWorkers returns the current worker count.
|
|
func NumWorkers() int { return engine.NumWorkers() }
|
|
|
|
// workersFor bounds the worker count for a workload of n items.
|
|
func workersFor(n int) int { return engine.WorkersFor(n) }
|
|
|
|
// parallel splits the [0, n) range into chunks across workers.
|
|
func parallel(n int, fn func(start, end int)) { engine.Parallel(n, fn) }
|
|
|
|
// parallelMin splits the [0, n) range across workers only while each
|
|
// worker keeps at least minPerWorker elements; a workload below that
|
|
// floor runs on the calling goroutine.
|
|
func parallelMin(n, minPerWorker int, fn func(start, end int)) {
|
|
engine.ParallelMin(n, minPerWorker, fn)
|
|
}
|
|
|
|
// The spawn floors below are constraints on when a parallel kernel may
|
|
// spawn at all: a worker must carry enough elements for its own
|
|
// creation and scheduling to stay a small fraction of the work it runs,
|
|
// and below that floor the whole map is faster on the calling
|
|
// goroutine. The floor scales with the per-element cost of the kernel,
|
|
// so each kernel family carries its own measured value; the sweep in
|
|
// math_bench_extra_test.go guards the sizes around both crossovers.
|
|
|
|
// elementwiseMinPerWorker is the per-worker chunk floor for the
|
|
// element-independent arithmetic maps (Add, Sub, Mul, Div and the
|
|
// scalar maps). Their per-element work is a couple of machine
|
|
// instructions, so a spawned worker needs about a thousand elements
|
|
// before the chunk outweighs its own start-up; below roughly 32k
|
|
// elements the maps measure faster serial, above that parallel.
|
|
const elementwiseMinPerWorker = 1024
|
|
|
|
// mathFuncMinPerWorker is the per-worker chunk floor for the
|
|
// per-element math maps (realFunc, roundFunc and the specialised
|
|
// Sqrt). Their per-element work costs tens of cycles, so a worker
|
|
// amortises its start-up at far smaller chunks than the arithmetic
|
|
// maps; a floor in the thousands would keep workloads of a few thousand
|
|
// elements needlessly serial.
|
|
const mathFuncMinPerWorker = 64
|