added support for all rk35xx models to NPU core runtime and cpu affinity core selection

cpuaffinity.go

@@ -2,6 +2,7 @@ package rknnlite
 
 import (
 	"fmt"
+	"strings"
 	"syscall"
 	"unsafe"
 )
@@ -11,17 +12,76 @@ const (
 	RK3588FastCores = uintptr(0b11110000)
 	// RK3588SlowCores is the cpu affinity mask of the efficient cortex A55 cores 0-3
 	RK3588SlowCores = uintptr(0b00001111)
-	// RK3588Allcores is the cpu affinity mask for all cortext A76 and A55 cores 0-7
+	// RK3588Allcores is the cpu affinity mask for all cortex A76 and A55 cores 0-7
 	RK3588AllCores = uintptr(0b11111111)
 
 	// RK3582FastCores is the cpu affinity mask of the fast cortex A76 cores 4-5
 	RK3582FastCores = uintptr(0b00110000)
 	// RK3582SlowCores is the cpu affinity mask of the efficient cortex A55 cores 0-3
 	RK3582SlowCores = uintptr(0b00001111)
-	// RK3582Allcores is the cpu affinity mask for all cortext A76 and A55 cores 0-5
+	// RK3582Allcores is the cpu affinity mask for all cortex A76 and A55 cores 0-5
 	RK3582AllCores = uintptr(0b00111111)
+
+	// RK3576FastCores is the cpu affinity mask of the fast cortex A72 cores 4-7
+	RK3576FastCores = uintptr(0b11110000)
+	// RK3576SlowCores is the cpu affinity mask of the efficient cortex A53 cores 0-3
+	RK3576SlowCores = uintptr(0b00001111)
+	// RK3576Allcores is the cpu affinity mask for all cortex A72 and A53 cores 0-7
+	RK3576AllCores = uintptr(0b11111111)
+
+	// RK3568AllCores is the cpu affinity mask of all cortex A55 (2Ghz) cores 0-3
+	RK3568AllCores = uintptr(0b00001111)
+
+	// RK3566AllCores is the cpu affinity mask of all cortex A55 (1.6Ghz) cores 0-3
+	RK3566AllCores = uintptr(0b00001111)
+
+	// RK3562AllCores is the cpu affinity mask of all cortex A53 cores 0-3
+	RK3562AllCores = uintptr(0b00001111)
 )
 
+// CoreType specifies the CPU core type
+type CoreType int
+
+const (
+	FastCores CoreType = 0
+	SlowCores CoreType = 1
+	AllCores  CoreType = 2
+)
+
+// coreMaskList defines a list of CPU core masks for lookup by key
+var coreMaskList = map[string]map[CoreType]uintptr{
+	"rk3562": {
+		SlowCores: RK3562AllCores,
+		FastCores: RK3562AllCores,
+		AllCores:  RK3562AllCores,
+	},
+	"rk3566": {
+		SlowCores: RK3566AllCores,
+		FastCores: RK3566AllCores,
+		AllCores:  RK3566AllCores,
+	},
+	"rk3568": {
+		SlowCores: RK3568AllCores,
+		FastCores: RK3568AllCores,
+		AllCores:  RK3568AllCores,
+	},
+	"rk3576": {
+		SlowCores: RK3576SlowCores,
+		FastCores: RK3576FastCores,
+		AllCores:  RK3576AllCores,
+	},
+	"rk3582": {
+		SlowCores: RK3582SlowCores,
+		FastCores: RK3582FastCores,
+		AllCores:  RK3582AllCores,
+	},
+	"rk3588": {
+		SlowCores: RK3588SlowCores,
+		FastCores: RK3588FastCores,
+		AllCores:  RK3588AllCores,
+	},
+}
+
 // SetCPUAffinity sets the CPU Affinity mask of the program to run on the specified
 // cores
 func SetCPUAffinity(mask uintptr) error {
@@ -63,3 +123,28 @@ func CPUCoreMask(cores []int) uintptr {
 
 	return mask
 }
+
+// SetCPUAffinityByPlatform sets the CPU Affinity mask of the program to run
+// on the specified CPU cores based on the given platform string of
+// rk3562|rk3566|rk3568|rk3576|rk3582|rk3582|rk3588
+func SetCPUAffinityByPlatform(platform string, ct CoreType) error {
+
+	var useCores uintptr
+	matched := false
+
+	platform = strings.TrimSpace(platform)
+	platform = strings.ToLower(platform)
+
+	if platform, ok := coreMaskList[platform]; ok {
+		if coretype, ok := platform[ct]; ok {
+			useCores = coretype
+			matched = true
+		}
+	}
+
+	if !matched {
+		return fmt.Errorf("unknown platform: %s", platform)
+	}
+
+	return SetCPUAffinity(useCores)
+}

runtime.go

@@ -8,6 +8,7 @@ import "C"
 import (
 	"fmt"
 	"os"
+	"strings"
 	"unsafe"
 )
 
@@ -37,9 +38,11 @@ var (
 	RK3588 = []CoreMask{NPUCore0, NPUCore1, NPUCore2}
 	RK3582 = []CoreMask{NPUCore0, NPUCore1, NPUCore2}
 	RK3576 = []CoreMask{NPUCore0, NPUCore1}
-	RK3568 = []CoreMask{NPUCore0}
-	RK3566 = []CoreMask{NPUCore0}
-	RK3562 = []CoreMask{NPUCore0}
+	// RK356x is a single NPU core and does not support the setCoreMask() call
+	// so we skip it
+	RK3568 = []CoreMask{NPUSkipSetCore}
+	RK3566 = []CoreMask{NPUSkipSetCore}
+	RK3562 = []CoreMask{NPUSkipSetCore}
 )
 
 // ErrorCodes
@@ -280,3 +283,28 @@ func (r *Runtime) InputAttrs() []TensorAttr {
 func (r *Runtime) OutputAttrs() []TensorAttr {
 	return r.outputAttrs
 }
+
+// NewRuntimeByPlatform returns a RKNN run time instance and automatically
+// selects the NPU cores to run on the given platform string of
+// rk3562|rk3566|rk3568|rk3576|rk3582|rk3582|rk3588
+// Provide the full path and filename of the RKNN compiled model file to run
+func NewRuntimeByPlatform(platform string, modelFile string) (*Runtime, error) {
+
+	platform = strings.TrimSpace(platform)
+	platform = strings.ToLower(platform)
+
+	var useCore CoreMask
+
+	switch platform {
+	case "rk3562", "rk3566", "rk3568":
+		useCore = NPUSkipSetCore
+
+	case "rk3576", "rk3582", "rk3588":
+		useCore = NPUCoreAuto
+
+	default:
+		return nil, fmt.Errorf("unknown platform: %s", platform)
+	}
+
+	return NewRuntime(modelFile, useCore)
+}